from .get_addresses import get_addresses 
import numpy as np
from scipy import sparse
import pickle

#clear all;
#for name in vars().keys():
#  del name

#gc.collect()

gps0 = {}
gps0['age']  = ['a1','a2','a3','a4','a5'] #gps0['age']  = {'a1','a2','a3','a4','a5'}
gps0['vacc'] = ['v0','v1','v2'] #gps0.vacc = {'v0','v1','v2'}; % After model calibration we can add these stages
gps0['hiv']  = ['h0','h1','hart'] #gps0.hiv  = {'h0','h1','hart'};
gps0['strains'] = ['ds','mdr'] #gps0.strains = {'ds','mdr'};
gps0['provs'] = ['pu','pr'] #gps0.provs = {'pu','pr'};

gps = {}
gps['age']  = ['a1','a2','a3','a4','a5'] #gps.age  = {'a1','a2','a3','a4','a5'};
gps['vacc'] = ['v0'] #gps.vacc = {'v0'};%,'v1','v2'}; % After model calibration we can add these stages
gps['hiv']  = ['h0','h1','hart'] #gps.hiv  = {'h0','h1','hart'};
gps['strains'] = ['ds','mdr'] #gps.strains = {'ds','mdr'};
gps['provs'] = ['pu','pr'] #gps.provs = {'pu','pr'};

gps1 = {}
gps1['age']  = ['a1','a2','a3','a4','a5'] #gps1.age  = {'a1','a2','a3','a4','a5'}; 
gps1['vacc'] = ['v0'] #gps1.vacc = {'v0'}; 
gps1['hiv']  = ['h0'] #gps1.hiv  = {'h0'}; 
gps1['strains'] = ['ds'] #gps1.strains = {'ds'}; 
gps1['provs'] = ['pr'] #gps1.provs = {'pr'};  

gps2 = {}
gps2['age']  = ['a1','a2','a3','a4','a5'] #gps2.age  = {'a1','a2','a3','a4','a5'}; 
gps2['vacc'] = ['v0']#gps2.vacc = {'v0'};  
gps2['hiv']  = ['h0'] #gps2.hiv  = {'h0'}; 
gps2['strains'] = ['ds','mdr'] #gps2.strains = {'ds','mdr'};
gps2['provs'] = ['pr'] #gps2.provs = {'pr'}; 

states0 = ['U'] #states0 = {'U'};                              % structured by hiv and vaccine status 
states1 = ['Lf','Ls'] #states1 = {'Lf','Ls'};                        % structured by hiv, vaccine and strain 
states2 = ['Isc','I','E','Rlo','Rhi','R'] #states2 = {'Isc','I','E','Rlo','Rhi','R'};    % structured by hiv and strain 
states3 = ['Dx','Tx','Tx2'] #states3 = {'Dx','Tx','Tx2'};                  % structured by hiv, strain and provider type

#% For the full model with vaccination
[i0, s0, d0, lim0] = get_addresses([states0, gps0['age'], gps0['hiv'], gps0['vacc']], {}, {}, {}, 0) #[i0, s0, d0, lim0] = get_addresses({states0, gps0.age, gps0.hiv, gps0.vacc}, [], [], [], 0);
[i0, s0, d0, lim0] = get_addresses([states1, gps0['age'], gps0['hiv'], gps0['vacc'], gps['strains']], i0, s0, d0, lim0)#[i0, s0, d0, lim0] = get_addresses({states1, gps0.age, gps0.hiv, gps0.vacc, gps.strains}, i0, s0, d0, lim0)
[i0, s0, d0, lim0] = get_addresses([states2, gps0['age'], gps0['hiv'], gps0['strains']], i0, s0, d0, lim0)#[i0, s0, d0, lim0] = get_addresses({states2, gps0.age, gps0.hiv, gps0.strains}, i0, s0, d0, lim0)
[i0, s0, d0, lim0] = get_addresses([states3, gps0['age'], gps0['hiv'], gps0['strains'], gps['provs']], i0, s0, d0, lim0)#[i0, s0, d0, lim0] = get_addresses({states3, gps0.age, gps0.hiv, gps0.strains, gps.provs}, i0, s0, d0, lim0)
d0 = [k for r, k in d0.items()] #d0 = chr(d0) #d0 = char(d0);

#% For the full model without vaccination
[i, s, d, lim] = get_addresses([states0, gps['age'], gps['hiv'], gps['vacc']], {}, {}, {}, 0)#[i, s, d, lim] = get_addresses({states0, gps.age, gps.hiv, gps.vacc}, [], [], [], 0);
[i, s, d, lim] = get_addresses([states1, gps['age'], gps['hiv'], gps['vacc'], gps['strains']], i, s, d, lim) #[i, s, d, lim] = get_addresses({states1, gps.age, gps.hiv, gps.vacc, gps.strains}, i, s, d, lim);
[i, s, d, lim] = get_addresses([states2, gps['age'], gps['hiv'], gps['strains']], i, s, d, lim) #[i, s, d, lim] = get_addresses({states2, gps.age, gps.hiv, gps.strains}, i, s, d, lim);
[i, s, d, lim] = get_addresses([states3, gps['age'], gps['hiv'], gps['strains'], gps['provs']], i, s, d, lim) #[i, s, d, lim] = get_addresses({states3, gps.age, gps.hiv, gps.strains, gps.provs}, i, s, d, lim);
d = [k for r, k in d.items()] #d = chr(d) #d = char(d);

#% For equilibrium model without mdr/hiv/vaccine/public sector 
[i1, s1, d1, lim1] = get_addresses([states0, gps1['age'], gps1['hiv'], gps1['vacc']], {}, {}, {}, 0) #[i1, s1, d1, lim1] = get_addresses({states0, gps1.age, gps1.hiv, gps1.vacc}, [], [], [], 0);
[i1, s1, d1, lim1] = get_addresses([states1, gps1['age'], gps1['hiv'], gps1['vacc'], gps1['strains']], i1, s1, d1, lim1)#[i1, s1, d1, lim1] = get_addresses({states1, gps1.age, gps1.hiv, gps1.vacc, gps1.strains}, i1, s1, d1, lim1);
[i1, s1, d1, lim1] = get_addresses([states2, gps1['age'], gps1['hiv'], gps1['strains']], i1, s1, d1, lim1)#[i1, s1, d1, lim1] = get_addresses({states2, gps1.age, gps1.hiv, gps1.strains}, i1, s1, d1, lim1);
[i1, s1, d1, lim1] = get_addresses([states3, gps1['age'], gps1['hiv'], gps1['strains'], gps1['provs']], i1, s1, d1, lim1)#[i1, s1, d1, lim1] = get_addresses({states3, gps1.age, gps1.hiv, gps1.strains, gps1.provs}, i1, s1, d1, lim1);
d1 = [k for r, k in d1.items()] #d1 = chr(d1)#d1 = char(d1);

#% Indtroducing mdr with the previous small model
[i2, s2, d2, lim2] = get_addresses([states0, gps2['age'], gps2['hiv'], gps2['vacc']], {}, {}, {}, 0) #[i2, s2, d2, lim2] = get_addresses({states0, gps2.age, gps2.hiv, gps2.vacc}, [], [], [], 0);
[i2, s2, d2, lim2] = get_addresses([states1, gps2['age'], gps2['hiv'], gps2['vacc'], gps2['strains']], i2, s2, d2, lim2)#[i2, s2, d2, lim2] = get_addresses({states1, gps2.age, gps2.hiv, gps2.vacc, gps2.strains}, i2, s2, d2, lim2);
[i2, s2, d2, lim2] = get_addresses([states2, gps2['age'], gps2['hiv'], gps2['strains']], i2, s2, d2, lim2)#[i2, s2, d2, lim2] = get_addresses({states2, gps2.age, gps2.hiv, gps2.strains}, i2, s2, d2, lim2);
[i2, s2, d2, lim2] = get_addresses([states3, gps2['age'], gps2['hiv'], gps2['strains'], gps2['provs']], i2, s2, d2, lim2) #[i2, s2, d2, lim2] = get_addresses({states3, gps2.age, gps2.hiv, gps2.strains, gps2.provs}, i2, s2, d2, lim2);
d2 = [k for r, k in d2.items()] #d2 = chr(d2) #d2 = char(d2);

s['infectious']      = np.concatenate([s['Isc'], s['I'], s['E'], s['Dx'], np.intersect1d(s['Tx'], s['mdr'])]) #s.infectious      = [s.Isc, s.I, s.E, s.Dx, intersect(s.Tx, s.mdr)];
s['infectious_wosc'] = np.concatenate([s['I'], s['E'], s['Dx'], np.intersect1d(s['Tx'], s['mdr'])]) #s.infectious_wosc = [s.I, s.E, s.Dx, intersect(s.Tx, s.mdr)];    % Infectious compartment without subclinical (contributing on TB death)
s['prevalent']       = np.unique(np.concatenate([s['infectious'], np.concatenate([s['Tx'],s['Tx2']])])) #s.prevalent       = unique([s.infectious, [s.Tx,s.Tx2]]);

s0['infectious']      = np.concatenate([s0['Isc'], s0['I'], s0['E'], s0['Dx'], np.intersect1d(s0['Tx'], s0['mdr'])]) #s0.infectious      = [s0.Isc, s0.I, s0.E, s0.Dx, intersect(s0.Tx, s0.mdr)];
s0['infectious_wosc'] = np.concatenate([s0['I'], s0['E'], s0['Dx'], np.intersect1d(s0['Tx'], s0['mdr'])]) #s0.infectious_wosc = [s0.I, s0.E, s0.Dx, intersect(s0.Tx, s0.mdr)];    % Infectious compartment without subclinical (contributing on TB death)
s0['prevalent']       = np.unique(np.concatenate([s0['infectious'], np.concatenate([s0['Tx'],s0['Tx2']])])) #s0.prevalent       = unique([s0.infectious, [s0.Tx,s0.Tx2]]);

s1['infectious']      = np.concatenate([s1['Isc'], s1['I'], s1['E'], s1['Dx']]) #s1.infectious      = [s1.Isc, s1.I, s1.E, s1.Dx];
s1['infectious_wosc'] = np.concatenate([s1['I'], s1['E'], s1['Dx']]) #s1.infectious_wosc = [s1.I, s1.E, s1.Dx];    % Infectious compartment without subclinical (contributing on TB death)
s1['prevalent']       = np.unique(np.concatenate([s1['infectious'], np.concatenate([s1['Tx'],s1['Tx2']])])) #s1.prevalent       = unique([s1.infectious, [s1.Tx,s1.Tx2]]);

s2['infectious']      = np.concatenate([s2['Isc'], s2['I'], s2['E'], s2['Dx'], np.intersect1d(s2['Tx'], s2['mdr'])]) #s2.infectious      = [s2.Isc, s2.I, s2.E, s2.Dx, intersect(s2.Tx, s2.mdr)];
s2['infectious_wosc'] = np.concatenate([s2['I'], s2['E'], s2['Dx'], np.intersect1d(s2['Tx'], s2['mdr'])]) #s2.infectious_wosc = [s2.I, s2.E, s2.Dx, intersect(s2.Tx, s2.mdr)];    % Infectious compartment without subclinical (contributing on TB death)
s2['prevalent']       = np.unique(np.concatenate([s2['infectious'], np.concatenate([s2['Tx'],s2['Tx2']])])) #s2.prevalent       = unique([s2.infectious, [s2.Tx,s2.Tx2]]);


#% --- Include the auxiliaries ---------------------------------------------
names = ['inc','noti','noti2','mort']
lgths =     [3,     3,      1,     2]

if not np.isin('aux',list(i)):
    i['aux'] = {}

for ii in range(len(names)): #for ii = 1:length(names)
    inds = lim + np.array(range(lgths[ii])) #inds = lim + [1:lgths(ii)];
    i['aux'][list(names)[ii]] = inds #i.aux.(names{ii}) = inds;
    lim = inds[-1]+1 #lim = inds(end);
#end
  
if not np.isin('aux',list(i0)):
    i0['aux'] = {}  
for ii in range(len(names)): #for ii = 1:length(names)
    inds0 = lim0 + np.array(range(lgths[ii])) #inds0 = lim0 + [1:lgths(ii)];
    i0['aux'][list(names)[ii]] = inds0 #i0.aux.(names{ii}) = inds0;
    lim0 = inds0[-1]+1 #lim0 = inds0(end);
#end

i['nx']  = lim
i0['nx']= lim0
i1['nx'] = lim1
i2['nx'] = lim2
#% --- Make aggregators and selectors --------------------------------------

#% Selectors for the incidence
agg = {}
tmp = np.zeros((3,i['nstates'])) #tmp = zeros(3,i.nstates);
tmp[0,s['Isc']] = 1 #tmp(1,s.Isc) = 1;                                            % All
tmp[1,[np.intersect1d(s['Isc'],s['h1']),np.intersect1d(s['Isc'],s['hart'])]] = 1#tmp(2,[intersect(s.Isc,s.h1),intersect(s.Isc,s.hart)]) = 1;  % hiv-TB
tmp[2,np.intersect1d(s['Isc'],s['mdr'])] = 1 #tmp(3,intersect(s.Isc,s.mdr)) = 1;                           % mdr-TB 
agg['inc'] = sparse.csc_matrix(tmp) #agg.inc = sparse(tmp);

sel ={}
tmp = np.zeros((i['nstates'],i['nstates'])) #tmp = zeros(i.nstates);
tmp[s['Isc'],:] = 1 #tmp(s.Isc,:) = 1;
tmp[s['h1'],s['h0']] = 0; tmp[s['hart'],s['h1']] = 0; tmp[s['h0'],s['h1']] = 0; tmp[s['h1'],s['hart']] = 0 #tmp(s.h1,s.h0) = 0; tmp(s.hart,s.h1) = 0; tmp(s.h0,s.h1) = 0; tmp(s.h1,s.hart) = 0;
tmp[s['ds'],s['mdr']]= 0; tmp[s['mdr'],s['ds']] = 0 #tmp(s.ds,s.mdr)= 0; tmp(s.mdr,s.ds) = 0;   % ADD
sel['inc'] = tmp - np.diag(np.diag(tmp)) #sel.inc = tmp - diag(diag(tmp));

tmp = np.zeros((i['nstates'],i['nstates'])) #tmp = zeros(i.nstates);
tmp[np.intersect1d(s['Tx'],s['mdr']),np.intersect1d(s['Tx'],s['ds'])] = 1 #tmp(intersect(s.Tx,s.mdr),intersect(s.Tx,s.ds)) = 1;
sel['acqu'] = sparse.csc_matrix(tmp - np.diag(np.diag(tmp))) #sel.acqu = sparse(tmp - diag(diag(tmp)));

#% Selectors for public sector notifications
tmp = np.zeros((3,i['nstates']))#tmp = zeros(3,i.nstates);
tmp[0,np.intersect1d(np.intersect1d([s['Tx'],s['Tx2']],s['h0']),s['pu'])] = 1 #tmp(1,intersect(intersect([s.Tx,s.Tx2],s.h0),s.pu)) = 1;  % public sector notification
tmp[1,np.intersect1d(np.intersect1d([s['Tx'],s['Tx2']],s['h1']),s['pu'])] = 1 #tmp(2,intersect(intersect([s.Tx,s.Tx2],s.h1),s.pu)) = 1;
tmp[2,np.intersect1d(np.intersect1d([s['Tx'],s['Tx2']],s['hart']),s['pu'])] = 1 #tmp(3,intersect(intersect([s.Tx,s.Tx2],s.hart),s.pu)) = 1;
agg['noti'] = sparse.csc_matrix(tmp) #agg.noti = sparse(tmp);

tmp = np.zeros((i['nstates'],i['nstates'])) #tmp = zeros(i.nstates);
tmp[[s['Tx'],s['Tx2']],:] = 1 #tmp([s.Tx,s.Tx2],:) = 1;
tmp[s['h1'],s['h0']] = 0; tmp[s['hart'],s['h1']] = 0; tmp[s['h0'],s['h1']] = 0; tmp[s['h1'],s['hart']] = 0 #tmp(s.h1,s.h0) = 0; tmp(s.hart,s.h1) = 0; tmp(s.h0,s.h1) = 0; tmp(s.h1,s.hart) = 0;
tmp[s['ds'],s['mdr']]= 0; tmp[s['mdr'],s['ds']] = 0
sel['noti'] = tmp - np.diag(np.diag(tmp))

#% Selectors for second-line notifications
tmp = np.zeros((i['nstates'],i['nstates']))
tmp[s['Tx2'],:] = 1
tmp[s['h1'],s['h0']] = 0; tmp[s['hart'],s['h1']] = 0; tmp[s['h0'],s['h1']] = 0; tmp[s['h1'],s['hart']] = 0
tmp[s['ds'],s['mdr']]= 0; tmp[s['mdr'],s['ds']] = 0
sel['noti2'] = tmp - np.diag(np.diag(tmp))


#% Selectors for the incidence from the full model with vaccine
agg0={}
tmp = np.zeros((3,i0['nstates']))
tmp[0,s0['Isc']] = 1  #                                          % All
tmp[1,np.concatenate([np.intersect1d(s0['Isc'],s0['h1']),np.intersect1d(s0['Isc'],s0['hart'])])] = 1 #;  % hiv-TB
tmp[2,np.intersect1d(s0['Isc'],s0['mdr'])] = 1                           #% mdr-TB 
agg0['inc'] = sparse.csc_matrix(tmp)

sel0 = {}
tmp = np.zeros((i0['nstates'],i0['nstates']))
tmp[s0['Isc'],:] = 1
tmp[s0['h1'],s0['h0']] = 0; tmp[s0['hart'],s0['h1']] = 0; tmp[s0['h0'],s0['h1']] = 0; tmp[s0['h1'],s0['hart']] = 0
tmp[s0['ds'],s0['mdr']]= 0; tmp[s0['mdr'],s0['ds']] = 0 #;   % ADD
sel0['inc'] = tmp - np.diag(np.diag(tmp))

tmp = np.zeros((i0['nstates'],i0['nstates']))
tmp[np.intersect1d(s0['Tx'],s0['mdr']),np.intersect1d(s0['Tx'],s0['ds'])] = 1
sel0['acqu'] = sparse.csc_matrix(tmp - np.diag(np.diag(tmp)))

#% Selectors for public sector notifications
tmp = np.zeros((3,i0['nstates']))
tmp[0,np.intersect1d(np.intersect1d([s0['Tx'],s0['Tx2']],s0['h0']),s0['pu'])] = 1 #tmp(1,intersect(intersect([s0.Tx,s0.Tx2],s0.h0),s0.pu)) = 1;  % public sector notification
tmp[1,np.intersect1d(np.intersect1d([s0['Tx'],s0['Tx2']],s0['h1']),s0['pu'])] = 1 #tmp(2,intersect(intersect([s0.Tx,s0.Tx2],s0.h1),s0.pu)) = 1;
tmp[2,np.intersect1d(np.intersect1d([s0['Tx'],s0['Tx2']],s0['hart']),s0['pu'])] = 1 #tmp(3,intersect(intersect([s0.Tx,s0.Tx2],s0.hart),s0.pu)) = 1;
agg0['noti'] = sparse.csc_matrix(tmp)

tmp = np.zeros((i0['nstates'],i0['nstates']))
tmp[[s0['Tx'],s0['Tx2']],:] = 1
tmp[s0['h1'],s0['h0']] = 0; tmp[s0['hart'],s0['h1']] = 0; tmp[s0['h0'],s0['h1']] = 0; tmp[s0['h1'],s0['hart']] = 0
tmp[s0['ds'],s0['mdr']]= 0; tmp[s0['mdr'],s0['ds']] = 0
sel0['noti'] = tmp - np.diag(np.diag(tmp))

#% Selectors for second-line notifications
tmp = np.zeros((i0['nstates'],i0['nstates']))
tmp[s0['Tx2'],:] = 1
tmp[s0['h1'],s0['h0']] = 0; tmp[s0['hart'],s0['h1']] = 0; tmp[s0['h0'],s0['h1']] = 0; tmp[s0['h1'],s0['hart']] = 0
tmp[s0['ds'],s0['mdr']]= 0; tmp[s0['mdr'],s0['ds']] = 0
sel0['noti2'] = tmp - np.diag(np.diag(tmp))

#% --- Define variables and ranges -----------------------------------------
#Change here
#names = ['r_beta','rfbeta_mdr','rfbeta_hiv','r_sym','p_pu','r_cs','r_cs2','r_mort_TB','p_Dx','p_Tx_complete','p_MDRrec2019', 'r_MDR_acqu','r_ART_init','r_HIV_mort','r_self_cure','p_HIV_relrate'];  # % <--- Updated
#lgths =        [1,           1,           1,     1,     1,     1,      1,          2,     2,              2,             1,            1,          1,           1,            1,            1]
#names = ['r_beta','rfbeta_mdr','rfbeta_hiv','r_sym','p_pu','r_cs','r_cs2','r_mort_TB','p_Dx','p_Tx_complete','p_MDRrec2022', 'r_MDR_acqu','r_ART_init','r_HIV_mort','r_self_cure','p_HIV_relrate', 'r_betared'];  # % <--- Updated
#lgths =        [1,           1,           1,     1,     1,     1,      1,          2,     2,              2,             1,            1,          1,           1,            1,            1, 1]

names = ['r_beta','rfbeta_mdr','rfbeta_hiv','r_sym','p_pu','r_cs','rf_cs2','r_mort_TB','p_Dx','p_Tx_complete','p_MDRrec2022', 'r_MDR_acqu','r_ART_init','r_HIV_mort','r_self_cure','p_HIV_relrate', 'rf_p_imm', 'rf_cs2022', 'r_betared'] #;   %'rf_LTBI_cure', <--- Updated
lgths =        [1,           1,           1,     1,     1,     1,      1,          2,     2,              2,             1,            1,            1,           1,            1,               1,           1,          1,         1]


xi = {}; lim = 0
for ii in range(len(names)): #for ii = 1:length(names)
    inds = lim + np.array(range(lgths[ii]))
    xi[names[ii]] = inds.astype(int) #xi.(names{ii}) = inds;
    #if names[ii]=='r_mort_TB':
    #    xxx = 0
    lim = lim + lgths[ii] #lim = inds(end);
#end
#%xi.calib = xi.r_HIV_mort;
xi['nx'] = lim

bds = np.zeros((xi['nx'],2))
prm = {}
bds[xi['r_beta'],:]       = np.array([0, 40]) #; % ds
bds[xi['rfbeta_mdr'],:]   = np.array([0, 1])
bds[xi['rfbeta_hiv'],:]   = np.array([0, 1])
bds[xi['r_sym'],:]        = np.array([0.1, 100])
bds[xi['p_pu'],:]         = np.array([0, 1])
bds[xi['r_cs'],:]         = np.array([0.1, 10]) #bds[xi['r_cs'],:]         = np.array([0.1, 100])
bds[xi['rf_cs2'],:]        = np.array([1, 40])#bds(xi.rf_cs2,:)        = [1 40]; ; #bds[xi['r_cs2'],:]        = np.array([1, 24])
#%bds(xi.r_Tx_init,:]   = [0.1 10];
bds[xi['r_mort_TB'],:]    = 1/6*np.array([[0, 2], [0, 100]]) #1/6*[0 2; 0 100];
bds[xi['p_Dx'],:]         = [np.array([0.75, 0.9]), np.array([0.1,  0.7])] #[0.75 0.9; 0.1  0.3]; % 0.4 0.7];  
bds[xi['p_Tx_complete'],:]= [np.array([0.75, 0.95]), np.array([0.4, 0.8])] #[0.75 0.95; 0.4 0.8];
bds[xi['p_MDRrec2022'],:] = np.array([0,  1])#bds[xi['p_MDRrec2019'],:] = np.array([0,  1])
bds[xi['r_MDR_acqu'],:]   = np.array([0,  0.06])
#% bds(xi.r_HIV_acqu,:)  = [0 10];
bds[xi['r_ART_init'],:]   = np.array([0, 10])
bds[xi['r_HIV_mort'],:]   = np.array([0, 10])
bds[xi['r_self_cure'],:]  = 1/6*np.array([0.85, 1.15])
bds[xi['p_HIV_relrate'],:]= np.array([1, 100])

bds[xi['rf_p_imm'],:]      = np.array([0.75, 1.25]) #bds(xi.rf_p_imm,:)      = [0.75, 1.25]; %[0.1 0.75];
bds[xi['rf_cs2022'],:]     = np.array([1, 10]) #bds(xi.rf_cs2022,:)     = [1 10];
#%bds(xi.rf_LTBI_cure,:)  = [0.75, 1.25];

bds[xi['r_betared'],:]     = np.array([0.8, 1]) #bds(xi.r_betared,:)     = [0.8 1];

prm['bounds'] = bds.T #prm.bounds = bds';
#% --- Define baseline parameter values ------------------------------------


#% Inter-age contact matrix 
#%prm.contact = [[1.3710,1.4317,0.0500];[2.5175,2.9007,0.0993];[0.2789,0.3350,0.0160]];% %Age group(0-24,25-60,60+) 
#%prm.contact = ones(3,3);% %Age group(0-24,25-60,60+) 
#% Age-group of Population in each compartment [0-4,5-10,11-15,16-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]; %[0.39 0.31 0.24   0.021  0.032]; %1970

#% Contact matrix correcting with population weight
prm['cm'] = [[0.9591,    0.4910,    0.1820,    0.6734,    0.0088],
            [0.4908,    2.1476,    0.7891,    0.8469,    0.0191],
            [0.2415,    0.5639,    2.4042,    1.3988,    0.0219],
            [1.2417,    1.0995,    1.1480,   10.0745,    0.0586],
            [0.0289,    0.0264,    0.0266,    0.0517,    0.0063]]
#%prm.cm = prm.contact./repmat(prm.N,size(prm.N,2),1); % 5 age groups
#%prm.gth = [0 0.02 0.02 0.02]; %[0.056, 0.037, 0.034, 0.02] ; %prm.growth = 0.02;

prm['gth'] = [0, 0.039, 0.039, 0.035] #prm.gth = [0 0.039 0.039 0.035]; %[0.056, 0.037, 0.034, 0.02] ; %prm.growth = 0.02;

r={}
r['mort'] = [0.012, 0.0012, 0.0009, 0.0048, 0.0622] #r.mort = [0.012, 0.0012, 0.0009, 0.0048, 0.0622];  % Average 1997-2019
r['ageing'] = [0.186,  0.184,  0.162,  0.0047] #r.ageing = [0.186  0.184  0.162  0.0047]; % Estimated using simple model run

#% Natural history (h0, h1, hart)
r['progression']  = 0.0826*np.array([1, 10, 10*0.4])*np.array([1, 1, 0.5]) #r.progression  = 0.0826*[1 10 10*0.4].*[1 1 0.5];
r['LTBI_stabil']  = 0.872*np.array([1, 0, 1]) #r.LTBI_stabil  = 0.872*[1 0 1];
r['reactivation'] = 0.0006*np.array([1, 100, 100*0.4])*np.array([1, 1, 0.5]) #r.reactivation = 0.0006*[1 100 100*0.4].*[1 1 0.5];
#% r.self_cure    = 1/6*[1 0 1];
#% r.mort_TB      = 1/6;
r['relapse']      = np.array([0.032, 0.14, 0.0015]) #r.relapse      = [0.032 0.14 0.0015];
p = {}
p['imm']          = np.array([0.8, 0, 0.8]) #p.imm          = [0.8 0 0.8];

#% Diagnosis stage
#% r.cs           = 1;
#% p.pu           = 0.5;
#% p.Dx           = 0.65; % 0.85; 
r['Dx']             = 52 #r.Dx             = 52;
#% r.cs2          = 2;

p['MDR_rec2015'] = np.array([0.001, 0]) #p.MDR_rec2015 = [0.001 0]; % p.MDR_rec2012  = [0.036 0]; % %CHECK for each country [pu pr]
p['Tx_init2']    = np.array([0.88, 0]) #p.Tx_init2    = [0.88 0] #p.Tx_init2    = [0.88 0]; %CHECK for each country [pu pr]
p['SL_trans']    = np.array([0.88, 0]) #; %[0.2 0]; [pu pr] #p.SL_trans    = [0.88 0]; %[0.2 0]; [pu pr]

#% Treatment stage
p['Tx_init']     = np.array([1, 1])#p.Tx_init     = [1 1];
r['Tx']      = 2 #r.Tx      = 2;
r['Tx2']     = 0.5 #r.Tx2     = 0.5;   
#% p.default = 1 - 0.71;
#% r.default = r.Tx*p.default./(1-p.default);
p['cure']    = np.array([1, 1]) #p.cure    = [1 1];

p['tsrsl']     = np.array([0.48, 1e-6]) #p.tsrsl     = [0.48 1e-6];                % Second-line treatment success rate in pu and pr
r['default2']  = r['Tx2']*(1-p['tsrsl'])/p['tsrsl'] #r.default2  = r.Tx2*(1-p.tsrsl)./p.tsrsl;
p['cure2']     = np.array([0.5, 0]) #p.cure2     = [0.5 0];                 % Should I consider a different value for those with HIV?

#% Intervention parameter
#% r.pt          = 0;
p['PT_PLHIV']    = np.array([1, 1, 1]) #p.PT_PLHIV    = [1 1 1]; % With PT intervenion among PLHIv [1 1 (1-0.6)]
r['access']      = 0#r.access      = 0;
r['cs3']         = 0 #r.cs3         = 0;          % case finding rate among subclinical TB
p['VE']          = np.array([0, 0]) #p.VE          = [0 0];
r['vacc']        = np.array([0, 0, 0, 0, 0]) #r.vacc        = [0 0 0 0 0]; 
r['waning']      = 0 #r.waning      = 0;

#% Bring them all together
prm['p'] = p; prm['r'] = r #prm.p = p; prm.r = r;
ref = {}
ref['i'] = i;   ref['s'] = s;   ref['d'] = d;  ref['xi'] = xi #ref.i = i;   ref.s = s;   ref.d = d;  ref.xi = xi;
#%ref.i0 = i0; ref.s0 = s0; ref.d0 = d0; 
ref['i1'] = i1; ref['s1'] = s1; ref['d1'] = d1 #ref.i1 = i1; ref.s1 = s1; ref.d1 = d1; 
ref['i2'] = i2; ref['s2'] = s2; ref['d2'] = d2 #ref.i2 = i2; ref.s2 = s2; ref.d2 = d2; 

#% --- Get calibration targets ---------------------------------------------
#% 2019

##LOOKS LIKE DATA ARE MISSING. SHOULD BE FIXED BEFORE Setup_Model can be run
#load('../Collect_data/NGA.mat')

# #Loading NGA specific data from JSON File
# f_data_ctr = open('JSON_data.JSON')
# data_ctr = json.load(f_data_ctr)
# data = data_ctr['data']
# prm = data_ctr['prm']
# r = data_ctr['r']
# r = data_ctr['p']
# xs = data_ctr['xs']

# #% 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

# #% --- Get calibration targets ---------------------------------------------
# #%data.pop00_22 = data.pop00_22;
# popn          = data['pop00_22'][-1] #popn          = data.pop00_22(end); % 2022
# 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.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_H1  = data.mort_H1;
# data['noti']     = np.array([0.8, 1, 1.2])*np.array(data['noti']) #data.noti     = [0.8 1 1.2]*data.noti;
# data['mdriniTX'] = np.array([0.8, 1, 1.2])*np.array(data['mdriniTX']) #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])*(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


# #% 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]

# prm['rHIV'] = ys2 #prm.rHIV = ys2;


# #% num = [data.inc_all; data.inc_h1; data.noti; data.ART_covg; data.HIV_prev; data.mort_H0; data.mort_H1;data.mdr2015;data.mdr2019;data.mdriniTX];
# show = 0
# f0 = get_distribution_fns(data['inc2015'],   'lognorm', show) #f0 = get_distribution_fns(data.inc2015,   'lognorm', show)
# f1 = get_distribution_fns(data['inc2022'],   'lognorm', show) #f1 = get_distribution_fns(data['inc_all'],   'lognorm', show), changed
# f2 = get_distribution_fns(data['inc_h1'],    'lognorm', show)
# f3 = get_distribution_fns(data['noti'],      'lognorm', show)
# f4 = get_distribution_fns(data['ART_covg'],  'lognorm', show)
# f5 = get_distribution_fns(data['HIV_prev'],  'lognorm', show)
# f6 = get_distribution_fns(data['mort_H0'],   'lognorm', show)
# f7 = get_distribution_fns(data['mort_H1'],   'lognorm', show)
# f8 = get_distribution_fns(data['mdr2015'],   'lognorm', show)
# f9 = get_distribution_fns(data['mdr2022'],   'lognorm', show) #f9 = get_distribution_fns(data['mdr2019'],   'lognorm', show), changed
# f10 = get_distribution_fns(data['mdriniTX'], 'lognorm', show); #%show ->0
# f11 = get_distribution_fns(data['sym'],      'lognorm', show)
      
# def fn(inc2015, inc2022, inc_h1, noti, ART_covg, HIV_prev, mort_H0, mort_H1, mdr2015, mdr2019, mdriniTX,sym):
#         res_0 = 10*f1[0](inc2015)
#         res_1 = 10*f1[0](inc2022)
#         res_2 = 10*f2[0](inc_h1) 
#         res_3 = f3[0](noti) 
#         res_4 = f4[0](ART_covg)
#         res_5 = f5[0](HIV_prev) 
#         res_6 = f6[0](mort_H0) 
#         res_7 = f7[0](mort_H1)
#         res_9 = 10*f9[0](mdr2019) 
#         res_10 = f10[0](mdriniTX)
#         res_11 = f11[0](sym) #; %+ f8(mdr2015)
#         result = res_0 + res_1 + res_2 + res_3 + res_4 + res_5 + res_6 + res_7 + res_9 + res_10 + res_11
#         return result


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

#save Model_setup;
#* Get all local variables as a dictionary
variables_dict = locals()

#* Filter out unwanted variables (like built-in ones)
variables_dict = {k: v for k, v in variables_dict.items() if not (k.startswith('__') or  k in {'gc', 'gc', 'get_addresses',
'np', 'scipy', 'sparse', 'interpolate', 'get_distribution_fns', 'math', 'pickle','json', 'show', 'variables_dict', 'f_data_ctr', 'f0',
'f1', 'f2', 'f3', 'f4', 'f5', 'f6', 'f7', 'f8', 'f9', 'f10', 'f11','fn','lhd'})}


#* Save all variables to a file named 'Model_setup.pkl'
with open('data/SetupModel/Model_setup.pkl', 'wb') as file:
    pickle.dump(variables_dict, file)



