from serenitas.utils.db2 import dbconn
from lz4.block import decompress
import datetime
from construct import Enum, CString, Struct, Int32ul, Int32sl, Int64ul, Int64sl, Float64l, this

TRatePt = Struct("fDate" / Int64sl,
                 "fRate" / Float64l)
TCurve = Struct("fNumItems" / Int32sl,
                "fArray" / TRatePt[this.fNumItems],
                "fBaseDate" / Int64sl,
                "fBasis" / Float64l,
                "fDayCountConv" / Int64sl)
YieldCurve = Struct("curve" / TCurve,
                    "n_dates" / Int64ul,
                    "dates" / Int64sl[this.n_dates])
CurveName = Struct("Seniority" / Enum(Int32ul, Senior=0, Subordinated=1),
                   "DocClause" / Enum(Int32ul, XR14=0, MR14=1, MM14=2, CR14=3),
                   "ticker" / CString("utf8"))

SurvivalCurve = Struct("curve" / TCurve,
                       "recovery" / Float64l[this.curve.fNumItems],
                       "defaulted" / Int64sl,
                       "name" / CurveName)

NewTCurve = Struct("fNumItems" / Int32sl,
                   "fDayCountConv" / Int32sl,
                   "fBaseDate" / Int64sl,
                   "fBasis" / Float64l,
                   "fArray" / TRatePt[this.fNumItems])

serenitasdb = dbconn("serenitasdb")
with serenitasdb.cursor(binary=True) as c:
    c.execute('SELECT curve FROM rate_curves WHERE effective_date=%s AND curve_type=%s',
              (datetime.date(2023, 1, 9), 532))
    (state,) = c.fetchone()
    curve = decompress(state, 500)
    yc = YieldCurve.parse(curve)
    c.execute('SELECT curve FROM cds_curves WHERE date=%s AND redcode=%s',
              (datetime.date(2023, 1, 9), 'EKFHAH'))
    (state,) = c.fetchone()
    curve = decompress(state, 600)
    sc = SurvivalCurve.parse(curve)

new_yc = NewTCurve.build({"fNumItems":yc.curve.fNumItems, "fDayCountConv":yc.curve.fDayCountConv, "fBaseDate":yc.curve.fBaseDate,
                          "fBasis":yc.curve.fBasis, "fArray":yc.curve.fArray})