import os
import os.path
import datetime
import csv
from csv import reader
import json
import re
import psycopg2
from common import root
from db import query_db, with_connection
import sys
import yaml

pattern1 = re.compile("REINVEST\[\w+::REINV_TBA(\d)\]\[DEAL,(\d+)\]=.*$")
# reinv end date
pattern2 = re.compile("(STANDARD_VAR\[)\w+(::#REINV_END,\d+\]=)(\d.*$)")
# reinvprices 1
pattern3 = re.compile("STANDARD_VAR\[\w+::#PRICE100_TBA(\d),(\d+)\]=")
pattern5 = re.compile("STANDARD_VAR\[\w+::#REINVPCT_TBA(\d),(\d+)\]=")
pattern7 = re.compile("KINGS3")
pattern8 = re.compile("(#COLLATREINV_REINV_PCT_EXT\[)\w+(::\*\]\[DEAL,\d+\])=100")
pattern9 = re.compile("(?P<a>SEVERITY\[\w+,\d+\]=)mkt\(70\)")

# we use 84 so that it's both dividable by 2 and 3
global_reinvfloatpercentage = 84
global_reinvfixedpercentage = 16

@with_connection
def dealname_from_cusip(conn, cusips):
    with conn.cursor() as c:
        c.callproc("dealname_from_cusip", params = cusip)
        dealnames = [d[0] for d in c.fetchall()]
    return dealnames

def get_reinv_assets(dealname, workdate):
    sqlstr = 'select * from et_historicaldealinfo(%s, %s) where ReinvFlag IS TRUE'
    reinvassets = query_db(sqlstr, params = (dealname, workdate), one = False)
    d = {}
    for line in reinvassets:
        d[line[3]] = line[22]
    return d

def get_recovery(dealname, workdate, defaultrecovery = 50):
    """ compute average price of defaulted assets
    """
    sqlstr = "select sum(coalesce(price, %s) * currentbalance)/sum(currentbalance) " + \
             "from et_aggdealinfo_historical(%s, %s) where defaultedflag is True"
    recovery = query_db(sqlstr, params = (defaultrecovery, dealname, workdate))

    try:
        recovery = recovery[0]
    except TypeError:
        recovery = defaultrecovery

    if not recovery:
        recovery = defaultrecovery
    return float(recovery)

def get_reinvenddate(dealname):
    sqlstr = 'SELECT \"Reinv End Date\" from latest_clo_universe where dealname=%s'
    reinvenddate = query_db(sqlstr, params = (dealname,))[0]
    if reinvenddate:
        reinvenddate = reinvenddate.strftime("%Y%m%d")
        return reinvenddate
    else:
        raise Exception("missing reinvestment end date")

def generate_scenarios(workdate, dealname):
    prometheus = os.path.join(root, "Scenarios", "prometheus.sss")
    n_scenarios = 100
    basedir = os.path.join(root, "Scenarios", "Intex curves_" + workdate)
    defaultedprice = get_recovery(dealname, workdate)
    replace = "\g<a>{0:.3f}".format(defaultedprice)
    with open(os.path.join(basedir, "csv", dealname + ".config")) as fh:
        config = yaml.load(fh)
    reinvflag = config['reinvflag']
    if reinvflag:
        reinvenddate = get_reinvenddate(dealname)
        reinv_assets = get_reinv_assets(dealname, workdate)
        n_float_assets = len([v for v in reinv_assets.values() if v == 'FLOAT'])
        n_fixed_assets = len([v for v in reinv_assets.values() if v == 'FIXED'])
        rollingmat = config['rollingmat']
        if n_fixed_assets == 0:
            reinvfixedpercentage = 0
        else:
            reinvfixedpercentage = global_reinvfixedpercentage / n_fixed_assets
        if n_float_assets > 0:
            reinvfloatpercentage = (100 - n_fixed_assets * reinvfixedpercentage)/n_float_assets

    try:
        with open(os.path.join(basedir, "csv", dealname + "-reinvprices.csv"), "r") as fh:
            dr = csv.DictReader(fh)
            reinvprices = {f: [] for f in dr.fieldnames}
            for line in dr:
                for f in dr.fieldnames:
                    reinvprices[f].append("{0:.3f}".format(float(line[f])))
    except IOError:
        reinvflag = False

    sssfile = os.path.join(basedir, "sss", dealname + ".sss")
    if not os.path.exists(os.path.join(basedir, "sss")):
        os.makedirs(os.path.join(basedir, "sss"))
    cdrscenarios = os.path.join(basedir, "csv", dealname + "-cdr.csv")
    recoveryscenarios = os.path.join(basedir, "csv", dealname + "-recovery.csv")
    fhsss = open(sssfile, "w")
    fhcdr = open(cdrscenarios, "r")
    fhrecovery = open(recoveryscenarios, "r")
    csvcdr = reader(fhcdr)
    csvrecovery = reader(fhrecovery)
    cdrline = next(csvcdr)
    cdrline = "\t".join(["{0:.3f}".format(float(cdr)) for cdr in cdrline])
    recoveryline = next(csvrecovery)
    recoveryline = "\t".join(["{0:.3f}".format(float(recovery)) for recovery in recoveryline])

    i=1
    with open(prometheus) as fh:
        for line in fh:
            line = line.rstrip()

            if "DEAL_NAME" in line:
                newline = "DEAL_NAME=" +  dealname.upper() + "\r\n"
                fhsss.write(newline)
                continue
            if not reinvflag and pattern8.match(line):
                line = re.sub(pattern8, r"\1{0}\2=0".format(dealname.upper()), line)
                fhsss.write(line + "\r\n")
                continue
            if not reinvflag and "DO_REINV" in line:
                fhsss.write("DO_REINV=0" + "\r\n")
                continue
            m = pattern1.match(line)
            if reinvflag and m:
                reinv_number, scen_number = m.groups()
                key = "REINV_TBA" + reinv_number
                if key in reinv_assets:
                    if reinv_assets[key] == "FLOAT":
                        coupon = 4
                    elif reinv_assets[key] == "FIXED":
                        coupon = 7
                    line = "REINVEST[{0}::REINV_TBA{1}][DEAL,{2}]=".format(dealname.upper(),
                                                                           reinv_number,
                                                                           scen_number)
                    line += "COUP_SPR={0}|AMORT=Bullet|USE_REINVEST_PIP=1|MAT_DATE={1}|".format(coupon,
                                                                                                rollingmat)
                    fhsss.write(line + "\r\n")
                continue
            if reinvflag and pattern2.match(line):
                line = re.sub(pattern2, r"\1{0}\2{1}", line).format(dealname.upper(), reinvenddate)
                fhsss.write(line + "\r\n")
                continue
            m = pattern3.match(line)
            if reinvflag and m:
                reinv_number, scen_number = m.groups()
                if dealname=="litpt3" and reinv_number=="1":
                    line = "STANDARD_VAR[LITPT3::#REINVLOANP100,{0}]".format(scen_number) + \
                           "={0}".format(" ".join(reinvprices["REINV_TBA1"]))
                    fhsss.write(line + "\r\n")
                    continue
                reinv_name = "REINV_TBA" + reinv_number
                if reinv_name in reinvprices:
                    line = \
                    "STANDARD_VAR[{0}::#PRICE100_TBA{1},{2}]={3}".format(dealname.upper(),
                                                                         reinv_number,
                                                                         scen_number,
                                                                         " ".join(reinvprices[reinv_name]))
                    fhsss.write(line + "\r\n")
                continue
            m = pattern5.match(line)
            if reinvflag and m:
                reinv_number, scen_number = m.groups()
                reinv_name = "REINV_TBA" + reinv_number
                if reinv_number=="1":
                    if dealname=="litpt3":
                        line = "STANDARD_VAR[LITPT3::#LOANREINVPCT,{0}]=100".format(scen_number)
                        fhsss.write(line + "\r\n")
                        continue
                    if dealname=="flags4":
                        line = "STANDARD_VAR[FLAGS4::#PCT100_TBA1,{0}]=100".format(scen_number)
                        fhsss.write(line + "\r\n")
                        continue
                if reinv_name in reinv_assets:
                    if reinv_assets[reinv_name] == 'FIXED':
                        line = "STANDARD_VAR[{0}::#REINVPCT_TBA{1},{2}]={3}".format(dealname.upper(),
                                                                                    reinv_number,
                                                                                    scen_number,
                                                                                    reinvfixedpercentage)
                    elif reinv_assets[reinv_name] == 'FLOAT':
                        line =  "STANDARD_VAR[{0}::#REINVPCT_TBA{1},{2}]={3}".format(dealname.upper(),
                                                                                     reinv_number,
                                                                                     scen_number,
                                                                                     reinvfloatpercentage)

                    fhsss.write(line + "\r\n")
                continue
            if pattern7.search(line):
                line = re.sub(pattern7, dealname.upper(), line)
                fhsss.write(line + "\r\n")
                continue
            if "LOSS_RATE[DEAL,{0}]".format(i) in line:
                newcdrline = "LOSS_RATE[DEAL,{0}]=".format(i) + cdrline
                fhsss.write(newcdrline + "\r\n")
                continue
            if "LOSS_SEVERITY[DEAL,{0}]".format(i) in line:
                newrecoveryline = "LOSS_SEVERITY[DEAL,{0}]=".format(i) + recoveryline
                fhsss.write(newrecoveryline + "\r\n")
                i = i + 1
                if i <= n_scenarios:
                    cdrline = next(csvcdr)
                    cdrline = "\t".join(["{0:.3f}".format(float(cdr)) for cdr in cdrline]) + "\r\n"
                    recoveryline = next(csvrecovery)
                    recoveryline = "\t".join(["{0:.3f}".format(float(recovery)) \
                                                  for recovery in recoveryline]) + "\r\n"
                continue
            if "LOSS_NONPERF_SEVERITY" in line:
                line = re.sub(pattern9, replace, line)
            fhsss.write(line + "\r\n")
    fhsss.close()
    fhrecovery.close()
    fhcdr.close()
    print("generated scenarios for: {0}".format(dealname))

if __name__ == "__main__":
    if len(sys.argv) > 1:
        workdate = sys.argv[1]
    else:
        workdate = str(datetime.date.today())
    if len(sys.argv) > 2:
        dealnames = sys.argv[2:]
    else:
        dealnames = [d.split(".")[0] for d in
                    os.listdir(os.path.join(root,  "Scenarios",
                                            "Intex curves_" + workdate, "csv"))
                    if "RData" in d]
    for dealname in dealnames:
        generate_scenarios(workdate, dealname)