import pandas as pd
from pandas.tseries.offsets import CustomBusinessDay, Day, QuarterBegin
from pandas.tseries.holiday import get_calendar, HolidayCalendarFactory, GoodFriday
import unittest

fed_cal = get_calendar('USFederalHolidayCalendar')
bond_cal = HolidayCalendarFactory('BondCalendar', fed_cal, GoodFriday)
bus_day = CustomBusinessDay(calendar=bond_cal())

def prev_immdate(valdate):
    dates = [pd.datetime(valdate.year-1, 12, 20)] + \
            [pd.datetime(valdate.year, m, 20) for m in [3, 6, 9, 12]]
    dates = [bus_day.rollforward(d) for d in dates]
    for i, d in enumerate(dates[1:], 1):
        if d >= valdate:
            break
    else:
        return dates[-1]
    return d if d==valdate else dates[i-1]

def imm_dates(start_date, end_date):
    start_date = bus_day.rollback(start_date)-19*Day()
    start_date = QuarterBegin(startingMonth=3).rollback(start_date)
    #should be close=left I think
    return (pd.date_range(start_date, end_date, freq='QS-MAR', closed=None).
            shift(19, freq='D').
            shift(0, bus_day))

def days_accrued(tradedate):
    start_protection = pd.Timestamp(tradedate) + Day()
    delta = start_protection-prev_immdate(start_protection)
    return delta.days

def isleapyear(date):
    return ((date.year % 4) == 0 & (~(date.year % 100 == 0) | date.year % 400 == 0))

def yearfrac(date1, date2, daycount):
    if daycount == '30/360':
        d1 = date1.dt.day.copy()
        d2 = date2.dt.day.copy()
        d2[(d2==31) & (d1.isin([30,31]))] = 30
        d1.replace({31:30})
        y1 = date1.dt.year
        y2 = date2.dt.year
        m1 = date1.dt.month
        m2 = date2.dt.month
        return (360*(y2-y1)+30*(m2-m1)+d2-d1)/360
    days_accrued = date2 - date1
    if daycount == 'ACT/365':
        return days_accrued.dt.days/365
    if daycount == 'ACT/360':
        return days_accrued.dt.days/360

class TestDaysAccrued(unittest.TestCase):
    def test(self):
        dates = [('2015-01-07', 17),
                 ('2015-06-19', 92),
                 ('2015-03-19', 0),
                 ('2015-06-20', 93),
                 ('2015-06-21', 0),
                 ('2015-12-21', 1),
                 ('2015-12-28', 8)]
        for date, days in dates:
            self.assertEqual(days_accrued(date), days)

if __name__=="__main__":
    unittest.main()