import numpy as np
#import os
#import MapTPstoPINs
#import ReadTargetPopAndIntvnData
#import CalcResourceNeeds
# import InterventionInit

from .RunGlobalEpiModelb import RunGlobalEpiModelb
import pickle
import json
import math
from .get_distribution_fns import get_distribution_fns
from scipy import interpolate

#from InterventionScenario import InterventionScenario
from .InterventionInit import InterventionInit

from .get_objective import get_objective as f_obj
from loguru import logger

import os

def run_dynamical_model(ISO3, prm_modvar, xs_modvar, pop_data_modvar, Hi_scn, Hi_ref, TxTSFL_scn, TxTSFL_ref, Hd_scn, Hd_ref, firstImpactYr, finalImpactYr, finalYr, final_model_path=''):
    logger.debug('Starting run_dynamical_model')
    Model_setup_file_name = f"{final_model_path}/data/SetupModel/Model_setup.pkl"
    init0 = []
    Model_init0_file_name = f"{final_model_path}/data/json_init/" + ISO3 + "_init0.json"
    if os.path.exists(Model_init0_file_name):
         if os.path.isfile(Model_init0_file_name):
              with open(Model_init0_file_name , "r") as fp:#with open(f"{final_model_path}\\data\\CtriesDefaults\\TB.JSON", "r") as fp:
                init0 = json.load(fp)['init0']
                 
    with open(Model_setup_file_name, 'rb') as fp:
        setup_data = pickle.load(fp)
    logger.debug(f'Hi_scn {Hi_scn}')
    logger.debug(f'Hi_ref {Hi_ref}')
    logger.debug(f'TxTSFL_scn {TxTSFL_scn}')
    logger.debug(f'TxTSFL_ref {TxTSFL_ref}')
    logger.debug(f'Hd_scn {Hd_scn}')
    logger.debug(f'Hd_ref {Hd_ref}')


    ScnPARS = [ISO3]
    #ScnPARS{1}=ISO3;

    #close all;%figures
    #clear all;%variables 

    #%Location of Dynamical model 
    #cd C:\work\GlobalModel\NGA_Final

    #%these two lines can be commented out in the initial Python translation
    #load calibres1.mat xsto2 outsto2 xi obj; 
    #load Model_setup.mat i0 s0 gps0 sel0 agg0 r p prm d s;
    #######################%%%%mat_file_name = f"{final_model_path}/data/Calibres/" + ISO3 + ".mat"
    # jso_file_name = f"{final_model_path}/data/CtriesDefaults/" + ISO3 + "_calibres.json"

    #################mat_contents = scipy.io.loadmat(mat_file_name) #mat_contents = scipy.io.loadmat('calibres1.mat')
    mat_contents = {}

    # if data_glob_input==None:
    #     f_glob_input = open('TBExample.JSON')
    #     data_glob_input = json.load(f_glob_input)


    
    #Loading NGA specific data
    # if prm_modvar:
    data = pop_data_modvar.copy()
    prm = prm_modvar.copy()

    if len(prm['r']['ageing'])==5 :
        prm['r']['ageing'] = prm['r']['ageing'][1:]

    r = prm['r'].copy()
    p = prm['p'].copy()
    xs = xs_modvar
    # else:
    #     f_data_ctr = open(jso_file_name)
    #     data_ctr = json.load(f_data_ctr)
    #     data = data_ctr['data']
    #     prm = data_ctr['prm']
    #     r = data_ctr['r']
    #     p = data_ctr['p']
    #     xs = data_ctr['xs']
    
    #cs_global = r['cs']
    #cs2_global = r['cs2']
    #pu_global = p['pu']

    #% Age-group of Population in each compartment [0-4,5-9,10-14,15-65,65+]   
    #prm['N'] = np.array([0.2, 0.2, 0.2, 0.2, 0.2]) #prm.N = [0.2 0.2 0.2 0.2 0.2]; % Initial population (assumption)
    #prm['cm'] = data['contact_mat'] #prm.cm = data.contact_mat;     % Contact matrix correcting with population weight
    #prm['gth'] = np.array([0, 0.03, 0.035, 0.044]) #prm.gth = [0 0.03 0.035 0.044]; %[0,data.gth];        % population growth rate
    #r['mort'] = data['age_mort'] #r.mort = data.age_mort';        % Age specific mortality rate
    #r['ageing'] = np.array([0.1774, 0.1696, 0.1594, 0.0050]) #r.ageing = [0.1774 0.1696 0.1594 0.0050]; %data.ageing';   % Ageing rate

    # self cure is missing from r. I am adding it manually. Carel should make sure that it is included in the country specific JSON file0.1667        0   0.1667 
    #r['self_cure'] = np.array([0.1667,        0,   0.1667])


    #% --- Get calibration targets ---------------------------------------------
    #%data.pop00_22 = data.pop00_22;
    popn          = data['pop00_22'][-1] #popn          = data.pop00_22(end); % 2022
    logger.debug(f"popn {popn}")
    logger.debug(f"HIV_prev before {np.array(data['HIV_prev'])}")
    data['inc2015']  = np.array(data['inc_all'])[np.array(data['inc_all']).shape[0]-8,:] #data.inc2015  = data.inc_all(end-7,:); 
    data['inc2022']  = np.array(data['inc_all'])[np.array(data['inc_all']).shape[0]-1,:] #data.inc2022  = data.inc_all(end,:); 
    data['inc2000'] = np.array(data['inc_all'])[np.array(data['inc_all']).shape[0]-23,:]
    #% data.inc_h1  = data.inc_h1;
    #% data.mdr2015  = data.mdr2015; 
    #% data.mdr2022  = data.mdr2022;
    #data['mort_H0']  = np.array(data['mort_H0'])[np.array(data['mort_H0']).shape[0]-1,:] #data.mort_H0  = data.mort_H0(end,:);
    data['mort_H02000'] = np.array(data['mort_H0'])[np.array(data['mort_H0']).shape[0]-23,:]
    data['mort_H02022'] = np.array(data['mort_H0'])[np.array(data['mort_H0']).shape[0]-1,:]
    #%data.mort_H1  = data.mort_H1;
    data['mort_H12022'] = np.array(data['mort_H1'])
    data['noti'] = np.array([0.8, 1, 1.2])*np.array(data['noti'])[-1] #data.noti     = [0.8 1 1.2]*data.noti;
    data['mdriniTX'] = np.array([0.8, 1, 1.2])*np.array(data['mdriniTX'])[-1]  #data.mdriniTX = [0.8 1 1.2]*data.mdriniTX; 
    #data['sym']      = np.array([0.36, 0.50, 0.80])#data.sym      = [0.36 0.50 0.80];        % proportion of prevalent TB that has symptoms

    #% Get the HIV data                                                         
    prm['ART_start'] = data['ART_start'] #prm.ART_start = data.ART_start;                     %ART_start;
    data['ART_covg'] =  np.array([0.85, 1, 1.15])*data['ART_covg'] #data.ART_covg =  [0.85 1 1.15]*data.ART_covg;       %ARTcovg_2022;
    data['HIV_prev'] = np.array([0.85, 1, 1.15])*(np.array(data['HIV_prev'])/popn) #data.HIV_prev = [0.85 1 1.15]*(data.HIV_prev./popn);%HIVprev_2022 (proportion of population)
    HIV_incd = data['HIV_incd'] #HIV_incd = data.HIV_incd;                           %HIV incidence (proportion of population) 1980-2022

    logger.debug(f"HIV_prev {data['HIV_prev']} ")
    logger.debug(f"ART_covg {data['ART_covg']} ")
    logger.debug(f"ART_start {data['ART_start']} ")
    logger.debug(f"HIV_incd {data['HIV_incd']} ")

    logger.debug(f"noti {data['noti']}")
    logger.debug(f"mort_H12022 {data['mort_H12022']}")

    #% Extrapolate HIV incidence by 20 years
    ys1 = HIV_incd #;   %1980-2022
    xs1 = np.array(range(len(ys1))) #xs1 = 1:length(ys1);
    xs2 = np.array(range(len(ys1)))+20#xs2 = 1:length(ys1)+20;    % 2022-2042
    #ys2 = interp1(xs1,ys1,xs2,'linear','extrap');  % linier trend for the next 21 years
    fnys2 = interpolate.interp1d(xs1,ys1,kind='linear',fill_value='extrapolate')
    ys2 = [fnys2(o) for o in xs2]


    #GUY: I am commenting this out because Matlab uses something different 
    #prm['rHIV'] = ys2 #prm.rHIV = ys2;
    #prm['r'] = r #@Carel: is this OK? The r in prm does not have the expected structure
    #prm['r']['self_cure'] = np.array([0.1667,        0,   0.1667])
    #prm['N'] = np.array([0.2, 0.2, 0.2, 0.2, 0.2]) #prm.N = [0.2 0.2 0.2 0.2 0.2]; % Initial population (assumption)
    #prm['cm'] = data['contact_mat'] #prm.cm = data.contact_mat;     % Contact matrix correcting with population weight
    #prm['gth'] = np.array([0, 0.03, 0.035, 0.044]) #prm.gth = [0 0.03 0.035 0.044]; %[0,data.gth];        % population growth rate
    #prm['r']['ageing'] = np.array([0.1774, 0.1696, 0.1594, 0.0050]) #r.ageing = [0.1774 0.1696 0.1594 0.0050]; %data.ageing';   % Ageing rate

    #prm['r']['cs'] =  cs_global
    #prm['r']['cs2'] =  cs2_global
    #prm['p']['pu'] = pu_global
    #r = prm['r']
    #END GUY: I am commenting this out because Matlab uses something different 

    show = 0
    f0 = get_distribution_fns(data['sym'],   'lognorm', show) #f0 = get_distribution_fns(data.sym,         'lognorm', show);#f0 = get_distribution_fns(data.inc2015,   'lognorm', show)
    
    #I add this here so that the code can run, Carel should update 'data' in the JSON file 
    #data['inc2000'] = data['inc2022'] 
    #data['mort_H02000'] = data['mort_H0']
    #data['mort_H02022'] = data['mort_H0']
    #data['noti_cu'] = data['noti']
    data['noti2022'] = data['noti']*np.array([0.8, 1, 1.2])
    data['noti_cu'] = np.sum(np.array(data['c_newinc'][-1]))*np.array([0.8, 1, 1.2]) #data['noti_cu'] = sum(data['c_newinc'])*np.array([0.8, 1, 1.2])
    logger.debug(f"noti2022 {data['noti2022']}")
    logger.debug(f"noti_cu {data['noti_cu']}")
    #data['mort_H12022'] = data['mort_H1']

    f1 = get_distribution_fns(data['inc2000'],     'lognorm', show) #f1 = get_distribution_fns(data.inc2000,     'lognorm', show);  
    f2 = get_distribution_fns(data['inc2022'],     'lognorm', show) #f2 = get_distribution_fns(data.inc2022,     'lognorm', show);
    f3 = get_distribution_fns(data['mort_H02000'], 'lognorm', show) #f3 = get_distribution_fns(data.mort_H02000, 'lognorm', show);
    f4 = get_distribution_fns(data['mort_H02022'], 'lognorm', show)#f4 = get_distribution_fns(data.mort_H02022, 'lognorm', show);
    f5 = get_distribution_fns(data['noti_cu'],     'lognorm', show)#f5 = get_distribution_fns(data.noti_cu,     'lognorm', show);
    f6 = get_distribution_fns(data['noti2022'],    'lognorm', show)#f6 = get_distribution_fns(data.noti2022,    'lognorm', show);
    f7 = get_distribution_fns(data['mdr2015'],     'lognorm', show)#f7 = get_distribution_fns(data.mdr2015,     'lognorm', show);
    f8 = get_distribution_fns(data['mdr2022'],     'lognorm', show)#f8 = get_distribution_fns(data.mdr2022,     'lognorm', show);
    f9 = get_distribution_fns(data['mdriniTX'],    'lognorm', show)#f9 = get_distribution_fns(data.mdriniTX,    'lognorm', show); 
    f10 = get_distribution_fns(data['HIV_prev'],   'lognorm', show)#f10 = get_distribution_fns(data.HIV_prev,   'lognorm', show);
    f11 = get_distribution_fns(data['inc_h1'],     'lognorm', show) #f11 = get_distribution_fns(data.inc_h1,     'lognorm', show);
    f12 = get_distribution_fns(data['ART_covg'],   'lognorm', show)#f12 = get_distribution_fns(data.ART_covg,   'lognorm', show);
    f13 = get_distribution_fns(data['mort_H12022'],'lognorm', show)#f13 = get_distribution_fns(data.mort_H12022,'lognorm', show);

    logger.debug("after get_distribution_fns")
    #lhd.fn  = @(sym,inc2000, inc2022, mort_H02000,mort_H02022,noti_cu, noti2022, mdr2015, mdr2022, mdriniTX,HIV_prev,inc_h1,ART_covg, mort_H12022)f0(sym)+f1(inc2000)+ 10*f2(inc2022)+f3(mort_H02000)+ 10*f4(mort_H02022)+ f5(noti_cu)+f6(noti2022)+f7(mdr2015)+f8(mdr2022)+f9(mdriniTX)+f10(HIV_prev)+f11(inc_h1)+f12(ART_covg)+f13(mort_H12022);
    def fn(sym,inc2000, inc2022, mort_H02000,mort_H02022,noti_cu, noti2022, mdr2015, mdr2022, mdriniTX,HIV_prev,inc_h1,ART_covg, mort_H12022):   
            res_0 = f0[0](sym)
            res_1 = f1[0](inc2000) 
            res_2 = 10*f2[0](inc2022)
            res_3 = f3[0](mort_H02000) 
            res_4 = 10*f4[0](mort_H02022) 
            res_5 = f5[0](noti_cu)
            res_6 = f6[0](noti2022)
            res_7 = f7[0](mdr2015)
            res_8 = f8[0](mdr2022)
            res_9 = f9[0](mdriniTX)
            res_10 = f10[0](HIV_prev)
            res_11 = f11[0](inc_h1)
            res_12 = f12[0](ART_covg)
            res_13 = f13[0](mort_H12022)
            result = res_0 + res_1 + res_2 + res_3 + res_4 + res_5 + res_6 + res_7 + res_8 + res_9 + res_10 + res_11 + res_12 + res_13
            return result


    lhd = {}
    lhd['fn']  = fn  
    lhd['sgn'] = -math.inf

    logger.debug("defining obj")
    def obj(): #def obj(x): # The signature was changed on June 09 2024 after Carel decided that we stop passing xs as a parameter
        #[r_out, r_aux] = f_obj(x, prm.copy(), setup_data['ref'], setup_data['sel'].copy(), setup_data['agg'].copy(), setup_data['gps'].copy(), setup_data['gps1'].copy(), setup_data['gps2'].copy(),lhd)
        [r_out, r_aux] = f_obj(prm.copy(), setup_data['ref'], setup_data['sel'].copy(), setup_data['agg'].copy(), setup_data['gps'].copy(), setup_data['gps1'].copy(), setup_data['gps2'].copy(),lhd)
        return([r_out,r_aux])

    #mat_contents['xsto2'] = np.array([xs])
    #mat_contents['outsto2'] = np.array([0])
    #PARS = [mat_contents['xsto2'],mat_contents['outsto2'],setup_data['xi'],obj,setup_data['i0'],setup_data['s0'],setup_data['gps0'],
    #        setup_data['sel0'],setup_data['agg0'],r.copy(),data_glob_input['<TB_DynamicalModelPRM_V1>']['p'],
    #    prm]
    PARS = [[0],[0],setup_data['xi'],obj,setup_data['i0'],setup_data['s0'],setup_data['gps0'],
            setup_data['sel0'],setup_data['agg0'],r.copy(),p.copy(), prm.copy(), init0.copy() ]


    #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    #%Empty intervention object
    EpiIntvn=InterventionInit()
    
    EpiIntvn['year1'] = np.array([firstImpactYr])
    EpiIntvn['year2'] = np.array([finalImpactYr])
    EpiIntvn['year3'] = np.array([finalYr+1])
    

    #[epi_outputs_base]=RunGlobalEpiModelb(PARS,EpiIntvn,PathToModel,PathToScenarioModel)
    
    # EpiIntvn empty and generates no impact
    logger.debug(f'EpiIntvn base {EpiIntvn}')
    epi_outputs_base =RunGlobalEpiModelb(PARS, EpiIntvn, True) #The paths are not parameters in the version of the Matlab procedures I have access to 

    logger.debug(f'epi_outputs_base {epi_outputs_base}')
    # with open('epi_outputs_base.JSON', 'r') as fP:
    #     tmp = ujson.load(fP)
    #     epi_outputs_base = [np.array(base) for base in tmp]
    inct_base=epi_outputs_base[0] #inct_base=epi_outputs_base{1};%.*DPProjBase(4:12)'/10^5;
    inch1_base=epi_outputs_base[1] #inch1_base=epi_outputs_base{2};%.*DPProjBase(4:12)'/10^5;
    mdr_base=epi_outputs_base[2] #mdr_base=epi_outputs_base{3};%.*DPProjBase(4:12)'/10^5;
    morth0_base=epi_outputs_base[3] #morth0_base=epi_outputs_base{4};%.*DPProjBase(4:12)'/10^5;
    morth1_base=epi_outputs_base[4] #morth1_base=epi_outputs_base{5};%.*DPProjBase(4:12)'/10^5;
    mort_base=morth0_base+morth1_base#mort_base=morth0_base+morth1_base;
    noti_base=epi_outputs_base[5] #noti_base=epi_outputs_base{6};%.*DPProjBase(4:12)'/10^5;
    notimdr_base=epi_outputs_base[6] #notimdr_base=epi_outputs_base{7};%.*DPProjBase(4:12)'/10^5;
    prevtb_base=epi_outputs_base[7]

    # EpiIntvn['year1'] = np.array([2023])
    # EpiIntvn['year2'] = np.array([2026])
    # EpiIntvn['year3'] = np.array([2031])

    EpiIntvn['year1_vac'] = np.array([firstImpactYr])
    EpiIntvn['year2_vac'] = np.array([finalImpactYr])
    EpiIntvn['year3_vac'] = np.array([finalYr+1])

    #%TB detection
    EpiIntvn['r_cs'] = np.ones(5) ## EpiIntvn.r_cs=ones(1,5);          
    EpiIntvn['r_cs2'] = np.ones(5)# EpiIntvn.r_cs2=ones(1,5);         
    #%EpiIntvn.r_cs3=ones(1,5);

    #%Impact on detection
    Hd_ratio = Hd_scn/Hd_ref #Hd_ratio=Hd_scn./Hd_ref; 
    logger.debug(f'Hd_ratio {Hd_ratio}')
    

    for a in range(5) : #for a=1:5
        EpiIntvn['r_cs'][a] = Hd_ratio  
        EpiIntvn['r_cs2'][a] = Hd_ratio#EiIntvn.r_cs2(a)=Hd_ratio;
    # end
   
    #%Prevention
    EpiIntvn['r_progression'] = np.zeros((5,3)) #IntvnScn.r_progression=zeros(5,3);
    EpiIntvn['r_reactivation'] = np.zeros((5,3)) #IntvnScn.r_reactivation=zeros(5,3);

    #%Impact on infection
    Hi_ratio=(1-Hi_scn)/(1-Hi_ref) #;
    logger.debug(f'Hi_ratio {Hi_ratio}')

    EpiIntvn['r_progression'][:,0] = 1-Hi_ratio #IntvnScn.r_progression(:,1)=1-Hi_ratio;	
    EpiIntvn['r_reactivation'][:,0]= 1-Hi_ratio #IntvnScn.r_reactivation(:,1)=1-Hi_ratio;
   
    #%Treatment Success
    #Impact on treatment success
    TxTSFL_ratio=(1-TxTSFL_scn)/(1-TxTSFL_ref)#;
    logger.debug(f'TxTSFL_ratio {TxTSFL_ratio}')

    EpiIntvn['p_cure'] = np.array([TxTSFL_ratio])
    EpiIntvn['p_cure2'] = np.array([TxTSFL_ratio])

    #%SL treatment duration, only if treatment duration is set
    if len(EpiIntvn['p_cure2'])!=0:
        EpiIntvn['r_Tx2'] = np.array([12/9]) #IntvnScn.r_Tx2 = 12/9; % Shorter Tx2 regimen (9 months)

    #[epi_outputs_scn]=RunGlobalEpiModelb(PARS,EpiIntvn,PathToModel,PathToScenarioModel)
    # EpiIntvn not empty and generates impact
    epi_outputs_scn =RunGlobalEpiModelb(PARS,EpiIntvn,False) #The paths are not parameters in the version of the Matlab procedures I have access to 
    logger.debug(f'epi_outputs_scn {epi_outputs_scn}')
    # with open('epi_outputs_scn.JSON', 'r') as fP:
    #     tmp = ujson.load(fP)
    #     epi_outputs_scn = [np.array(scn) for scn in tmp]
    inct_scn=epi_outputs_scn[0] #inct_base=epi_outputs_base{1};%.*DPProjBase(4:12)'/10^5;
    inch1_scn=epi_outputs_scn[1] #inch1_base=epi_outputs_base{2};%.*DPProjBase(4:12)'/10^5;
    mdr_scn=epi_outputs_scn[2] #mdr_base=epi_outputs_base{3};%.*DPProjBase(4:12)'/10^5;
    morth0_scn=epi_outputs_scn[3] #morth0_base=epi_outputs_base{4};%.*DPProjBase(4:12)'/10^5;
    morth1_scn=epi_outputs_scn[4] #morth1_base=epi_outputs_base{5};%.*DPProjBase(4:12)'/10^5;
    mort_scn=morth0_scn+morth1_scn#mort_base=morth0_base+morth1_base;
    noti_scn=epi_outputs_scn[5] #noti_base=epi_outputs_base{6};%.*DPProjBase(4:12)'/10^5;
    notimdr_scn=epi_outputs_scn[6] #notimdr_base=epi_outputs_base{7};%.*DPProjBase(4:12)'/10^5;
    prevtb_scn=epi_outputs_scn[7]

    Outputs_Scn = []
    Outputs_Scn.append(epi_outputs_base) 
    Outputs_Scn.append(epi_outputs_scn) 
    return Outputs_Scn 
    
     
    #end


