from copy import deepcopy
from os import environ
from requests import post, get
from loguru import logger

from AvenirCommon.Util import GBRange
from AvenirCommon.Database import GB_get_db_json

from SpectrumCommon.Const.TB import TB_Fort_Final_Year

from SpectrumCommon.Const.GB.GBVersions import TB_FORT_INPUT_VERSION, TB_FORT_OUTPUT_VERSION

def get_tb_fort_outputs(ISO3, vsn=TB_FORT_OUTPUT_VERSION):
    fort_outputs = GB_get_db_json(environ['AVENIR_SW_DEFAULT_DATA_CONNECTION'], 'tuberculosis', f'fort/outputs/{ISO3}_{vsn}.JSON')

    # Rescaling notification to the reported 2021 values
    # idx   = reported_year-TB_FirstHist_Year
    # scale = reported_notif/fort_outputs['nMid'][idx]
    # fort_outputs['nMid'][idx:] = scale*np.array(fort_outputs['nMid'][idx:])
    # fort_outputs['nLo'][idx:]  = scale*np.array(fort_outputs['nLo'][idx:])
    # fort_outputs['nHi'][idx:]  = scale*np.array(fort_outputs['nHi'][idx:])
    # fort_outputs['nSd'][idx:]  = scale*np.array(fort_outputs['nSd'][idx:])
 
    return fort_outputs


def get_fort_version():
    fort_url = environ.get("AVENIR_SW_FORT_URL", "https://tbbetafort-server.azurewebsites.net")
    response = get(f'{fort_url}/version')
    return response.json()['version']

def call_tb_fort(ISO3, fort_inputs, final_year, tb_slice, fort_slice, HRd = None, HRi = None , ORt = None):
    # fort_inputs = update_tb_fort_inputs(fort_inputs, tb_slice, fort_slice, HRd, HRi, ORt)

    fort_inputs["hRd"][fort_slice] = HRd[tb_slice] if HRd is not None else [1.0]*len(fort_inputs["hRd"][fort_slice])  
    fort_inputs["hRi"][fort_slice] = HRi[tb_slice] if HRi is not None else [1.0]*len(fort_inputs["hRi"][fort_slice])  
    fort_inputs["oRt"][fort_slice] = ORt[tb_slice] if ORt is not None else [1.0]*len(fort_inputs["oRt"][fort_slice])  

    if (ISO3 in ['IDN', 'VNM']):
         fort_inputs['modelType'] = 'failsafe' 
    else: 
         fort_inputs['modelType'] = 'IPn2'


    logger.debug(f"Calling TB Fort API version {get_fort_version()} with model type {fort_inputs['modelType']} for {ISO3} with final year {final_year} and years {fort_inputs['year'][0]}-{fort_inputs['year'][-1]}")

    fort_url = environ.get("AVENIR_SW_FORT_URL", "https://tbbetafort-server.azurewebsites.net")
    # fort_url = "https://tbfort-server.azurewebsites.net"
    # with open('PeteDoddsModel_Inputs.JSON', 'w+') as fp:
        # ujson.dump(fort_inputs, fp)   
    # if type(HRd)!=type(None):
    #     with open('Live_Inputs.JSON', 'r') as fp:
    #         fort_inputs =  ujson.load(fp)   
    response = post(f'{fort_url}/projection', json=fort_inputs)
    # response = post('http://localhost:9004/projection', json=fort_inputs)
    fort_outputs = response.json()


    # with open('PeteDoddsModel_'+ISO3+'_Outputs.JSON', 'w+') as fp:
    #     ujson.dump(fort_outputs, fp)   
    for key in fort_outputs:
        for i,e in enumerate(fort_outputs[key]):
            if e=='NA':
               fort_outputs[key][i] = 0

    # if 'error' in fort_outputs:
        # raise Exception('TB Fort API called failed')
               
    return fort_outputs


def get_tb_fort_inputs(ISO3, final_year, tb_slice, fort_slice, HRd, HRi, ORt,  fort_vsn=TB_FORT_INPUT_VERSION):
    fort_inputs = GB_get_db_json(environ['AVENIR_SW_DEFAULT_DATA_CONNECTION'], 'tuberculosis', f'fort/inputs/{ISO3}_{fort_vsn}.JSON') 
    # fort_inputs = {}
    # null = None
    # fort_inputs["year"] = [2000,2001,2002,2003,2004,2005,2006,2007,2008,2009,2010,2011,2012,2013,2014,2015,2016,2017,2018,2019,2020,2021,2022,2023,2024,2025,2026,2027,2028,2029,2030]
    # fort_inputs["iHat"] = [282000,282000,280000,275000,270000,264000,259000,254000,250000,245000,239000,233000,225000,218000,206000,197000,187000,177000,168000,160000,151000,143000,null,null,null,null,null,null,null,null,null]
    # fort_inputs["sEi" ] = [165306,160714,153316,144643,135204,126020,118112,111224,104337,97449,90561,65306,76020,69133,60459,38776,39031,33929,29847,26276,25765,26276,null,null,null,null,null,null,null,null,null]
    # fort_inputs["nHat"] = [91101,94957,110289,117600,123127,124262,122198,128844,141157,148936,154694,156539,145323,131677,119592,137960,127407,117705,114233,112597,108714,104854,null,null,null,null,null,null,null,null,null]
    # fort_inputs["sEn" ] = [6972,7267,8440,9000,9423,9510,9352,9861,10803,11398,11839,11980,11122,10077,9152,10558,9751,9008,8742,8617,8320,8025,null,null,null,null,null,null,null,null,null]
    # fort_inputs["mHat"] = [108000,103000,92000,85000,78000,75000,74000,68000,60000,53000,47000,42000,41000,44000,44000,32000,31000,31000,27000,25000,23000,21000,null,null,null,null,null,null,null,null,null]
    # fort_inputs["sEm" ] = [26276,24235,21173,19388,17857,17092,16837,15561,13520,11990,10459,6633,7908,8418,8673,5357,5612,5867,5102,4847,4592,4337,null,null,null,null,null,null,null,null,null]
    # fort_inputs["pHat"] = [null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null]
    # fort_inputs["sEp" ] = [null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null]
    # fort_inputs["TXf" ] = [0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035,0.035]
    # fort_inputs["modelType"] = "IP"
    # #estimate TXf
    #     PropHIVp         = notif_pct[TB_HIVpCases]
    #     PropHIVn         = 1 - PropHIVp
    #     PropHIVponART    = notif_pct[TB_OnARTCases]
    #     PropHIVpNoART    = 1- PropHIVponART
    #     PropHIVponART6m  = 0.1 * PropHIVponART
    #     PropHIVponART12m = 0.9 * PropHIVponART
    #     # check that     PropHIVn+PropHIVpNoART+PropHIVponART6m+PropHIVponART12m=1
    #     TXf              = PropHIVn*(3/100)+PropHIVpNoART*(9/100)+PropHIVponART6m*(6/100)+PropHIVponART12m*(4/100)
    #     # check value for Ethiopia for which TXf ~ 0.035. Check value for Eswatini should closer to the HIV values
    if fort_vsn in ['V1', 'V2']:
        fort_inputs["sEp"] = []
        fort_inputs["tXf"] = []
        fort_inputs["pHat"] = []
        for yr in GBRange(fort_inputs["year"][0], fort_inputs["year"][-1]):
            fort_inputs["sEp"].append(None)
            fort_inputs["pHat"].append(None)
            fort_inputs["tXf"].append(0.035)

        for yr in GBRange(fort_inputs["year"][-1]+1, TB_Fort_Final_Year):
            fort_inputs["year"].append(yr)    
            for key in ["iHat", "sEi", "nHat", "sEn", "mHat", "sEm", "pHat", "sEp"]:
                fort_inputs[key].append(None)
            fort_inputs["tXf"].append(0.035)
        fort_inputs["hRd"] = [1]*len(fort_inputs["year"]) 
        fort_inputs["hRi"] = [1]*len(fort_inputs["year"]) 
        fort_inputs["oRt"] = [1]*len(fort_inputs["year"]) 


    if type(HRd)!=type(None):
        fort_inputs["hRd"][fort_slice] = HRd[tb_slice]  

    if type(HRi)!=type(None):
        fort_inputs["hRi"][fort_slice] = HRi[tb_slice]

    if type(ORt)!=type(None):
        fort_inputs["oRt"][fort_slice] = ORt[tb_slice]
     
    return fort_inputs
