import datetime
import pandas as pd

from dateutil.relativedelta import relativedelta, WE
from pandas.tseries.offsets import CustomBusinessDay, Day, QuarterBegin
from pandas.tseries.holiday import get_calendar, HolidayCalendarFactory, GoodFriday

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


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 previous_twentieth(d):
    r = datetime.date(d.year, d.month, 20)
    if r > d:
        r -= relativedelta(months=1)
    mod = r.month % 3
    if mod != 0:
        r -= relativedelta(months=mod)
    return r


def imm_date(d):
    if isinstance(d, str):
        d = datetime.date.fromisoformat(d)
    r = d + relativedelta(day=1, weekday=WE(3))
    if r <= d:
        return imm_date(r + relativedelta(months=1, day=1))
    return r


def days_accrued(tradedate):
    tradedate = pd.Timestamp(tradedate)
    start_protection = tradedate + Day()
    # we want the the largest twentieth imm date rolled forward such that
    # it's less than start_protection
    effective_date = pd.Timestamp(
        previous_twentieth(bus_day.rollback(start_protection).date())
    )
    effective_date = bus_day.rollforward(effective_date)
    delta = start_protection - effective_date
    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