from dateutil import rrule
import datetime

"""
    forked from
    https://gist.github.com/jckantor/d100a028027c5a6b8340

    modified for backtesting up to 1986, the original rules are valid for time from now on
"""

# Generate ruleset for holiday observances on the NYSE

def NYSE_holidays(start_date, end_date):
    """
        rules for holidays on the NYSE from 1986,
        modified start date for some of the holidays

        for day in NYSE_holidays(datetime.datetime(2017,1,1),datetime.datetime(2017,12,31)):
            print day.strftime('%b %d %Y')

        print( list(NYSE_holidays(datetime.datetime(2017,1,1),datetime.datetime(2017,12,31))) )


        Args:
            start and end date format - datetime.datetime(2017,1,1)

         return:
             NYSE holidays
    """
    rs = rrule.rruleset()

    # Include all potential holiday observances
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start_date, until=end_date, bymonth=12, bymonthday=31, byweekday=rrule.FR)) # New Years Day
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start_date, until=end_date, bymonth= 1, bymonthday= 1))                     # New Years Day
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start_date, until=end_date, bymonth= 1, bymonthday= 2, byweekday=rrule.MO)) # New Years Day
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart= datetime.datetime(1998, 1, 1), until=end_date, bymonth= 1, byweekday= rrule.MO(3)))            # Martin Luther King Day
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=datetime.datetime(1971, 1, 1), until=end_date, bymonth= 2, byweekday= rrule.MO(3)))            # Washington's Birthday
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start_date, until=end_date, byeaster= -2))                                  # Good Friday
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=datetime.datetime(1971, 1, 1), until=end_date, bymonth= 5, byweekday= rrule.MO(-1)))           # Memorial Day
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start_date, until=end_date, bymonth= 7, bymonthday= 3, byweekday=rrule.FR)) # Independence Day
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start_date, until=end_date, bymonth= 7, bymonthday= 4))                     # Independence Day
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start_date, until=end_date, bymonth= 7, bymonthday= 5, byweekday=rrule.MO)) # Independence Day
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start_date, until=end_date, bymonth= 9, byweekday= rrule.MO(1)))            # Labor Day
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start_date, until=end_date, bymonth=11, byweekday= rrule.TH(4)))            # Thanksgiving Day
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start_date, until=end_date, bymonth=12, bymonthday=24, byweekday=rrule.FR)) # Christmas
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start_date, until=end_date, bymonth=12, bymonthday=25))                     # Christmas
    rs.rrule(rrule.rrule(rrule.YEARLY, dtstart=start_date, until=end_date, bymonth=12, bymonthday=26, byweekday=rrule.MO)) # Christmas

    # Exclude potential holidays that fall on weekends
    rs.exrule(rrule.rrule(rrule.WEEKLY, dtstart=start_date, until=end_date, byweekday=(rrule.SA,rrule.SU)))

    # those dates are not holidays but for some reason they are treated as ones
    # Not Exclude,
    rs.exdate(datetime.datetime(2010, 12, 31))
    rs.exdate(datetime.datetime(2004, 12, 31))
    rs.exdate(datetime.datetime(1999, 12, 31))

    rs.exdate(datetime.datetime(1993, 12, 31))
    rs.exdate(datetime.datetime(1982, 12, 31))
    rs.exdate(datetime.datetime(1976, 12, 31))
    rs.exdate(datetime.datetime(1971, 12, 31))
    rs.exdate(datetime.datetime(1965, 12, 31))

    return rs

# Generate ruleset for NYSE trading days

def NYSE_trading_days(start_date, end_date):
    """
        NYSE trading days from 1986
        in addition to the rules added add hoc events:
        9-11, hurricanes, presidential funerals

        tdays = len(list(NYSE_trading_days(datetime.datetime(2016,1,26),datetime.datetime(2016,02,26))))

        Args:
            start and end date

         return:
             NYSE trading days start and end date included
    """
    rs = rrule.rruleset()

    rs.rrule(rrule.rrule(rrule.DAILY, dtstart=start_date, until=end_date))

    # Exclude weekends and NYSE holidays
    rs.exrule(rrule.rrule(rrule.WEEKLY, dtstart=start_date, byweekday=(rrule.SA,rrule.SU)))
    rs.exrule(NYSE_holidays(start_date,end_date))

    # Exclude a day of mourning for President George H.W. Bush.
    rs.exdate(datetime.datetime(2018, 12, 5))

    # Exclude Hurricane Sandy
    rs.exdate(datetime.datetime(2012, 10, 29))
    rs.exdate(datetime.datetime(2012, 10, 30))

    # Exclude a day of mourning for President Ford.
    rs.exdate(datetime.datetime(2007, 1, 2))

    # Exclude a day of mourning for Ronald W. Reagan, who died on Saturday, June 5, 2004.
    rs.exdate(datetime.datetime(2004, 6, 11))

    # Exclude PresidentialFuneral-RichardNixon
    rs.exdate(datetime.datetime(1994, 4, 27))

    # Exclude 9-11  (2001, 9, 11) - (2001, 9, 14)
    rs.exdate(datetime.datetime(2001, 9, 11))
    rs.exdate(datetime.datetime(2001, 9, 12))
    rs.exdate(datetime.datetime(2001, 9, 13))
    rs.exdate(datetime.datetime(2001, 9, 14))

    return rs


if __name__ == '__main__':
    # Examples

    # List all NYSE holiday observances for the coming year
    print("NYSE Holidays\n")
    for day in NYSE_holidays(datetime.datetime(2016,1,1),datetime.datetime(2016,12,31)):
        print(day.strftime('%b %d %Y'))

    # Count NYSE trading days in next 5 years
    print("\n\nTrading Days\n")
    for yr in range(2015,2020):
        tdays = len(list(NYSE_trading_days(datetime.datetime(yr,1,1),datetime.datetime(yr,12,31))))
        print("{0}  {1}".format(yr,tdays))