from libc.stdlib cimport malloc, free
from curve cimport (JpmcdsBuildIRZeroCurve, JpmcdsZeroPrice, JpmcdsMakeTCurve,
                    Basis, CONTINUOUS, JpmcdsCleanSpreadCurve)
from date cimport (JpmcdsStringToDateInterval, pydate_to_TDate, dcc,
                   JpmcdsDateIntervalToFreq, JpmcdsDateFwdThenAdjust, TDate_to_pydate,
                   JpmcdsDateFromBusDaysOffset)
from date import dcc_tostring
from cdsone cimport JpmcdsStringToStubMethod, TStubMethod

cdef int SUCCESS = 0

cpdef public enum BadDay:
    FOLLOW = <long>'F'
    PREVIOUS = <long>'P'
    NONE = <long>'N'
    MODIFIED = <long>'M'

cdef class Curve:
   def __dealloc__(self):
       if self._thisptr is not NULL:
           JpmcdsFreeTCurve(self._thisptr)

   def inspect(self):
       return {'base_date': TDate_to_pydate(self._thisptr.fBaseDate),
               'basis': self._thisptr.fBasis,
               'day_count_convention': dcc_tostring(self._thisptr.fDayCountConv),
               'data': fArray_to_list(self._thisptr.fArray, self._thisptr.fNumItems)}

cdef fArray_to_list(TRatePt* fArray, int fNumItems):
    cdef size_t i
    cdef list l = []
    for i in range(fNumItems):
        l.append((TDate_to_pydate(fArray[i].fDate), fArray[i].fRate))
    return l

cdef class ZeroCurve(Curve):

    def __init__(self, date, str types,
                 list periods, double[:] rates,
                 str mm_dcc, str fixed_swap_period, str float_swap_period,
                 str fixed_swap_dcc, str float_swap_dcc, BadDay bad_day_conv):
        """ Initialize a zero coupon curve

        instruments need to be sorted by tenor
        """
        cdef:
           double fixed_freq
           double float_freq
           TDateInterval ivl
           char* routine = 'zerocurve'
           TDate value_date = pydate_to_TDate(date)

        self._dates = <TDate*>malloc(len(periods) * sizeof(TDate))
        self._ninstr = len(periods)

        cdef TDate settle_date
        if JpmcdsDateFromBusDaysOffset(value_date, 2, "None", &settle_date)!= SUCCESS:
            raise ValueError

        cdef TDateInterval tmp
        for i, p in enumerate(periods):
            period_bytes = p.encode('utf-8')
            if JpmcdsStringToDateInterval(period_bytes, routine, &tmp) != SUCCESS:
                raise ValueError
            if JpmcdsDateFwdThenAdjust(settle_date, &tmp, MODIFIED,
                                       "None", &self._dates[i]) != SUCCESS:
                raise ValueError('Invalid interval')

        fixed_bytes = fixed_swap_period.encode('utf-8')
        float_bytes = float_swap_period.encode('utf-8')
        types_bytes = types.encode('utf-8')

        if JpmcdsStringToDateInterval(fixed_bytes, routine, &ivl) != SUCCESS:
            raise ValueError
        if JpmcdsDateIntervalToFreq(&ivl, &fixed_freq) != SUCCESS:
            raise ValueError
        if JpmcdsStringToDateInterval(float_bytes, routine, &ivl) != SUCCESS:
            raise ValueError
        if JpmcdsDateIntervalToFreq(&ivl, &float_freq) != SUCCESS:
            raise ValueError

        self._thisptr = JpmcdsBuildIRZeroCurve(
            value_date, types_bytes, self._dates,
            &rates[0], len(periods), dcc(mm_dcc), <long> fixed_freq,
            <long> float_freq, dcc(fixed_swap_dcc), dcc(float_swap_dcc),
            bad_day_conv, b"None"
        )

    def __dealloc__(self):
        if self._dates is not NULL:
            free(self._dates)

    def discount_factor(self, date):
        if self._thisptr is NULL:
            raise ValueError('curve is empty')
        cdef TDate discount_date = pydate_to_TDate(date)
        return JpmcdsZeroPrice(self._thisptr, discount_date)

    def list_dates(self):
        cdef size_t i
        return [TDate_to_pydate(self._dates[i]) for i in range(self._ninstr)]

cdef class SpreadCurve(Curve):

    def __init__(self, today, ZeroCurve zc, start_date, step_in_date,
                 cash_settle_date, list end_dates, double[:] coupon_rates,
                 double recovery_rate, int pay_accrued_on_default):

        cdef TDate today_c = pydate_to_TDate(today)
        cdef TDate step_in_date_c = pydate_to_TDate(step_in_date)
        cdef TDate cash_settle_date_c = pydate_to_TDate(cash_settle_date)
        cdef TDate start_date_c = pydate_to_TDate(start_date)
        cdef TDate* end_dates_c = <TDate*>malloc(len(end_dates) * sizeof(TDate))
        self._thisptr = NULL
        cdef size_t i

        for i, d in enumerate(end_dates):
            end_dates_c[i] = pydate_to_TDate(d)
        cdef TStubMethod stub_type

        if JpmcdsStringToStubMethod(b"f/s", &stub_type) != 0:
            raise ValueError("can't convert stub")

        self._thisptr = JpmcdsCleanSpreadCurve(today_c,
                                               zc._thisptr,
                                               start_date_c,
                                               step_in_date_c,
                                               cash_settle_date_c,
                                               len(end_dates),
                                               end_dates_c,
                                               &coupon_rates[0],
                                               NULL,
                                               recovery_rate,
                                               pay_accrued_on_default,
                                               NULL,
                                               dcc('ACT/360'),
                                               &stub_type,
                                               <long>'M',
                                               b'NONE')

    @classmethod
    def from_flat_hazard(cls, base_date, double rate, Basis basis = CONTINUOUS,
                         str day_count_conv = 'Actual/365F'):
        cdef TDate base_date_c = pydate_to_TDate(base_date)
        cdef SpreadCurve sc = cls.__new__(cls)
        cdef TDate max_date = 200000 # can go higher but this should be more than enough
        cdef TDate* dates = <TDate*>malloc(sizeof(TDate))

        cdef double* rates = <double*>malloc(sizeof(double))
        dates[0] = max_date
        rates[0] = rate

        sc._thisptr = JpmcdsMakeTCurve(base_date_c, dates, rates, 1,
                                       <double>basis, dcc(day_count_conv))
        return sc