from scipy.io import savemat
import numpy as np
import pandas as pd
#import os
import scipy.io
import openpyxl

import .InterventionInit
import matplotlib.pyplot as plt

from .alloc_parameters import alloc_parameters
from .make_model3 import make_model3
import pickle
import scipy
import numpy as np
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 .ode15s import ode15s, create_steps
from .goveqs_scaleup import goveqs_scaleup

import time

def run_simulate_forward(ISO3= 'NGA', data_glob_input=None):
   
    Model_setup_file_name = "data/SetupModel/Model_setup.pkl"
    with open(Model_setup_file_name, 'rb') as fp:
        setup_data = pickle.load(fp)

    xi = setup_data['xi'].copy()
    i0 = setup_data['i0']
    s0 = setup_data['s0']
    gps0 = setup_data['gps0']
    sel0 = setup_data['sel0']
    agg0 = setup_data['agg0'] 
    d0 = setup_data['d0']
    ref = setup_data['ref']  
    agg = setup_data['agg']
    gps = setup_data['gps']
    gps1 = setup_data['gps1']
    gps2 = setup_data['gps2']
    sel = setup_data['sel']

    ScnPARS = [ISO3]
    jso_file_name = "data/CtriesDefaults/" + ISO3 + "_calibres.json"

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

    expl_prm = data_glob_input['<TB_DynamicalModelPRM_V1>']
    expl_p = data_glob_input['<TB_DynamicalModelPRM_V1>']['p']
    expl_r = data_glob_input['<TB_DynamicalModelPRM_V1>']['r']
    
    cs_global = expl_r['cs']
    cs2_global = expl_r['cs2']
    pu_global = expl_p['pu']

    #Loading ctr specific data
    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 = np.array(data_ctr['xs'])
    prm['r']['self_cure'] = np.array([0.1667,        0,   0.1667])
    
    prm['r']['cs'] =  cs_global
    prm['r']['cs2'] =  cs2_global
    prm['p']['pu'] = pu_global
    r = prm['r']

    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     
    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['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_H12022'] = np.array(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
    
    data['noti2022'] = data['noti']*np.array([0.8, 1, 1.2])
    data['noti_cu'] = sum(data['c_newinc'])*np.array([0.8, 1, 1.2])

    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);

    #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
 
    def obj(x):
        [r_out, r_aux] = f_obj(x, prm.copy(), ref, sel, agg, gps, gps1, gps2,lhd)
        return([r_out,r_aux]) 


    #clear all; %load calibres2; xsto = xsto4;    
    #load calibres1; xsto = xsto2; %(2e4:end,:);
    #[M,I] = max(outsto2); %(2e4:end,:));
    # load Model_setup;
    #xs = xsto(I,:);  % Most probable parameter set 
    #ix0 = 2e4; nx = 150; dx = round((size(xsto,1)-ix0)/nx); %2e4 nx = 150
    #xs_all = xsto(ix0:dx:end,:,1);
    #pct_xs= prctile(xs_all,[2.5,50,97.5])


    #% % --- Get the Regional notifications
    tend1   = 2024+1 #;  % Intervention starts except vaccination  
    tend1b  = 2028#;  % Vaccination starts
    tend2   = 2035+1#;  % End date for simulation

    #% --- Do the simulations 
    #opts = odeset('NonNegative',[1:i0.nstates],'Refine',64,'AbsTol',1e-10,'RelTol',1e-10);

    r0 = r.copy(); p0 = p.copy()
    inct = []
    noti = []
    inch1 = []
    mdr = []
    morth0 = []
    morth1 = []
    noti = []
    notimdr = []
    artcovg = []
    prvty = []
    prvt_mdry = []
    prvt_hivy = []
    ARTI = []
    incty = []
    inc_hivy = []
    inc_mdry = []
    mrty = []
    mrt_hivy = []
    notiy = []
    notimdry = []
    LTBI_recent = []
    LTBI_remote = []
    LTBI = []
    pop = []
    Sus1 = []
    Sus2 = []

    #% mk = round(size(xs,1)/25);
    #% for ii = 1:size(xs,1)
    #%     if mod(ii,mk) == 0; fprintf('%0.5g ',ii/mk); end 
    
    r, p = alloc_parameters(xs,prm['r'],prm['p'], xi)
    #% Get the initial conditions
    [out, aux] = obj(xs)
    init = aux['soln'][-1,:]
    
    #% Extend initial condition for the full model with vaccine compartments  
    inds1 = np.intersect1d(s0['U'], s0['v0']) #inds1 = intersect(s0.U, s0.v0);
    inds2 = np.intersect1d(np.concatenate([s0['Lf'], s0['Ls']]), s0['v0'])#inds2 = intersect([s0.Lf, s0.Ls], s0.v0);
    inds3 = np.concatenate([s0['Isc'], s0['I'], s0['E'], s0['Rlo'], s0['Rhi'], s0['R']])#inds3 = [s0.Isc, s0.I, s0.E, s0.Rlo, s0.Rhi, s0.R];
    inds4 = np.concatenate([s0['Dx'],s0['Tx'],s0['Tx2']]) #inds4 = [s0.Dx,s0.Tx,s0.Tx2];
    inds5 =  s0['Tx2'][-1]+np.array(range(9))#;  #% Auxiliary terms
    ind = np.concatenate([inds1, inds2, inds3,inds4,inds5])
   
    #% To allocate results from a model without vaccine to input vector for full model with vaccine
    init0 = np.zeros((1,i0['nx']))#init0 = zeros(1,i0.nx);
    init0[0][ind] = init #init0(ind) = init;   

    tref = range(2022,tend2+1,1) #tref = [2022:1:tend2];
     
    #% --- Now simulate without intervention (baseline)
    p['pu'] = p['pu_2022']; r['cs'] = r['cs_2022']; r['cs2'] = r['cs2_2022']# p.pu = p.pu_2022; r.cs = r.cs_2022; r.cs2 = r.cs2_2022;
    prm['growth'] = prm['gth'][3]#prm.growth = prm.gth(4);
    M0 = make_model3(p, r, i0, s0, gps0, prm); 
        
    #% trend=(1-((data.inc2015(2)-data.inc2022(2))/(data.inc2015(2)*(2022-2015))));
    #%r.betared = 0.987; %trend; % Additional reduction of beta to capture current incidence trend.
    #geq = @(t,in) goveqs_scaleup(t, in, r, M0, M0, [2022 2023], i0, prm, sel0, agg0);
    #[t0, soln0] = ode15s(geq, [2022:tend2], init0, odeset('NonNegative',[1:i0.nstates]));
    t0, soln0 = ode15s(goveqs_scaleup, init0[0], create_steps(2022, tend2, 1/12), args=(r, M0, M0, [2022, 2023],i0,prm, sel0, agg0))
       
    soln  = soln0
    t     = t0
    #soln0 = interp1(t, soln, tref);
    soln0 = np.zeros((len(tref),np.shape(soln)[1]))
    for jj in range(np.shape(soln)[1]):#range(i['nx']):
        temp_fn_sol = interpolate.interp1d(t, soln[:,jj],fill_value="extrapolate") 
        soln0[:,jj] = [temp_fn_sol(ti) for ti in tref]

     
    #% Find the initial condition at the starting point for the interventions
    #geq = @(t,in) goveqs_scaleup(t, in, r, M0, M0, [2022 2023], i0, prm, sel0, agg0);
    #[ta, solna] = ode15s(geq, [2022:tend1], init0, odeset('NonNegative',[1:i0.nstates]));
    #initb = solna(end,:);
    ta, solna = ode15s(goveqs_scaleup, init0[0], create_steps(2022, tend1, 1), args=(r, M0, M0, [2022, 2023],i0,prm, sel0,  agg0))
    initb = solna[-1,:]
   
    #% --- Next simulate with interventions
    r1 = r.copy(); p1 = p.copy()
    p1['pu']  = p['pu'] + (1- p['pu'])*0.35; 
    p1['Dx'][0] = 0.95;          #% [0.95 0.95];  
    
    #%     r1.cs(4)  = 0.8*r1.cs(4);  %1.2     %choose specific age group
    #%     r1.cs2(4) = 0.8*r1.cs2(4); %1.2     %choose specific age group
    #%     r1.cs3(4) = 0.8*r1.cs(4);  %1.2     %choose specific age group
    
    age_gp = 3 #age_gp = 4 #;  % choose specific age group
    csfac = 1.2 #; % Choose any positive number to increase or decrease of care-seeking rate 
    r1['cs'][age_gp]  = max(r['cs_1997'][age_gp],r['cs'][age_gp]*csfac)#r1.cs(age_gp)  = max(r.cs_1997(age_gp),r.cs(age_gp)*csfac);
    r1['cs2'][age_gp] = max(r['cs2_1997'][age_gp],r['cs2'][age_gp]*csfac) #r1.cs2(age_gp) = max(r.cs2_1997(age_gp),r.cs2(age_gp)*csfac);
    r1['cs3'][age_gp] = max(r['cs_1997'][age_gp],r['cs'][age_gp]*csfac) #r1.cs3(age_gp) = max(r.cs_1997(age_gp),r.cs(age_gp)*csfac);
    
    qq = 0.05;                                          # % Preventive therapy among household contacts independent of hiv status
    r1['progression'][3,:]  = r['progression'][3,:]*(1-qq) #r1.progression(4,:)  = r.progression(4,:)*(1-qq);    % Replace by this  --> p.PT_PLHIV = [1 1 (1-0.6)]
    r1['reactivation'][3,:] = r['reactivation'][3,:]*(1-qq)#r1.reactivation(4,:) = r.reactivation(4,:)*(1-qq);    
    p['PT_PLHIV'] = np.array([1, 1, (1-0.6)]) #p.PT_PLHIV = [1 1 (1-0.6)];                           % PT among PLHIV reduces rate by 60%
    
    p1['MDR_rec'] = np.array([0.5,0]) #p1.MDR_rec = [0.5,0]; %0.7; %0.8;
    r1['Tx2'] = 9/12 #r1.Tx2 = 9/12;                             % Shorter Tx2 regimn (9 months)
    p1['cure2'] = np.array([0.7, 0]) #p1.cure2 = [0.7, 0]; %0.8; %0.9;
    r1['default'][0] = r['default'][0]*(1-0.05) #r1.default(1) = r.default(1)*(1-0.05);     % reduction of default rate(pu) (by inccreasing Tx completion and reducing LTFU)     
 
    r1['ART_init'] = r['ART_init']*(1+0.5) #r1.ART_init = r.ART_init*(1+0.5);          % Increase of ART coverage
    M1 = make_model3(p1, r1, i0, s0, gps0,prm)
      
    #geq = @(t,in) goveqs_scaleup(t, in, r, M0, M1, tend1 + [0  3], i0, prm, sel0, agg0);
    #[tb, solnb] = ode15s(geq, [tend1 tend2], initb, opts); 
    tb, solnb = ode15s(goveqs_scaleup, initb, create_steps(tend1, tend2, 1/2), args=(r, M0, M1, tend1 + np.array([0,  3]),i0,prm, sel0, agg0))
   
    soln  = np.concatenate([solna, solnb[1:,:]]) #soln  = [solna; solnb(2:end,:)]; 
    t     = np.concatenate([ta,    tb[1:]]) #t     = [ta;    tb(2:end)];
     
    soln8 = np.zeros((len(tref),np.shape(soln)[1])) #soln8 = interp1(t, soln, tref); 
    for jj in range(np.shape(soln)[1]):#range(i['nx']):
        temp_fn_sol = interpolate.interp1d(t, soln[:,jj],fill_value="extrapolate") 
        soln8[:,jj] = [temp_fn_sol(ti) for ti in tref]
    
    #% --- Next simulate with vaccination
    #% Find vaccination starting point
    #[tv, solnv] = ode15s(geq, [tend1 tend1b], initb, opts);
    tv, solnv = ode15s(goveqs_scaleup, initb, create_steps(tend1, tend1b, 1/2), args=(r, M0, M1, tend1 + np.array([0,  3]),i0,prm, sel0, agg0))
    initv = solnv[-1,:] #initv = solnv(end,:);
   
    #% Choose suitable rate of vaccination and vaccine efficacy
    r2 = r1.copy(); p2 = p1.copy() #r2 = r1; p2 = p1;
    p2['VE'] = np.array([0,  0.7]) #p2.VE = [0  0.7];                    % Efficacy--> [infection prevention, disease prevention]
    r2['vacc'] = np.array([0, 0, 5, 5, 5]) #r2.vacc = [0 0 5 5 5];               % Vaccination rate (chosen)
    M2 = make_model3(p2, r2, i0, s0, gps0, prm) #M2 = make_model3(p2, r2, i0, s0, gps0, prm); % With vaccine
      
    #geq = @(t,in) goveqs_scaleup(t, in,r, M1, M2, tend1b + [0  6], i0, prm, sel0, agg0); 
    #[tc, solnc] = ode15s(geq, [tend1b tend2], initv, opts);
    tc, solnc = ode15s(goveqs_scaleup, initv, create_steps(tend1b, tend2, 1/2), args=(r, M1, M2, tend1b + np.array([0,  6]),i0,prm, sel0, agg0))
    
    soln  = np.concatenate([solna, solnv[1:,:], solnc[1:,:]]) #soln  = [solna; solnv(2:end,:); solnc(2:end,:)];
    t     = np.concatenate([ta,    tv[1:], tc[1:]]) #t     = [ta;    tv(2:end);      tc(2:end)];
     
    soln9 = np.zeros((len(tref),np.shape(soln)[1])) #soln8 = interp1(t, soln, tref); 
    for jj in range(np.shape(soln)[1]):#range(i['nx']):
        temp_fn_sol = interpolate.interp1d(t, soln[:,jj],fill_value="extrapolate") 
        soln9[:,jj] = [temp_fn_sol(ti) for ti in tref]
    
    #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%     
    #% --- Bring them all together (baseline, interventions other than vaccine, interventions including vaccine)
    #allmat = cat(3, soln0,soln8,soln9); 
    allmat = np.dstack((soln0, soln8, soln9)) #allmat = cat(3, soln0,soln8,soln9); 
    dallmat = np.diff(allmat,n=1,axis=0) #dallmat = np.diff(allmat,[],1)
    
    inct = np.squeeze(dallmat[:,i0['aux']['inc'][0],:])*1e5 #inct(:,:) = squeeze(dallmat(:,i0.aux.inc(1),:))*1e5;    % Total incidence
    inch1 = np.squeeze(dallmat[:,i0['aux']['inc'][1],:])*1e5 #inch1(:,:) = squeeze(dallmat(:,i0.aux.inc(2),:))*1e5;   % HIV incidence
    mdr = np.squeeze(dallmat[:,i0['aux']['inc'][2],:])*1e5 #mdr(:,:) = squeeze(dallmat(:,i0.aux.inc(3),:))*1e5;     % MDR incidence
    morth0 = np.squeeze(dallmat[:,i0['aux']['mort'][0],:])*1e5 #morth0(:,:) = squeeze(dallmat(:,i0.aux.mort(1),:))*1e5; % hiv-negative mortality
    morth1 = np.squeeze(dallmat[:,i0['aux']['mort'][1],:])*1e5 #morth1(:,:) = squeeze(dallmat(:,i0.aux.mort(2),:))*1e5; % hiv-positive mortality
    noti = np.squeeze(np.sum(dallmat[:,i0['aux']['noti'],:],1))*1e5 #noti(:,:) = squeeze(sum(dallmat(:,i0.aux.noti,:),2))*1e5; % Total notification
    notimdr = np.squeeze(dallmat[:,i0['aux']['noti2'],:])*1e5 #notimdr(:,:) = squeeze(dallmat(:,i0.aux.noti2,:))*1e5;  % MDR Treatment initiation
   
    artcovg = np.squeeze(np.sum(allmat[:,s0['hart'],:],1)/np.sum(allmat[:,np.concatenate([s0['h1'],s0['hart']]),:],1))#artcovg(:,:) = squeeze(sum(allmat(:,s.hart,:),2)./sum(allmat(:,[s.h1,s.hart],:),2));
    
    #%   prvt = squeeze(sum(allmat(:,s0.prevalent,:),2))*1e5;       % Prevalence
    prvty = np.squeeze(np.sum(allmat[1:,s0['prevalent'],:],1)) #prvty(:,:) = squeeze(sum(allmat(2:end,s0.prevalent,:),2));                              %yearly prevalence
    prvt_mdry = np.squeeze(np.sum(allmat[1:,np.intersect1d(s0['prevalent'],s0['mdr']),:],1)) #prvt_mdry(:,:) = squeeze(sum(allmat(2:end,intersect(s0.prevalent,s0.mdr),:),2));         %yearly mdr prevalence
    prvt_hivy = np.squeeze(np.sum(allmat[1:,np.intersect1d(s0['prevalent'],[s0['h1'],s0['hart']]),:],1)) #prvt_hivy(:,:) = squeeze(sum(allmat(2:end,intersect(s0.prevalent,[s0.h1,s0.hart]),:),2)); %yearly hiv prevalence
    ARTI = np.squeeze(xs[xi['r_beta']]*np.sum(allmat[1:,np.intersect1d(s0['infectious'],s0['h0']),:],1)/np.sum(allmat[1:,np.array(range(len(d0))),:],1)) #ARTI(:,:) = squeeze(xs(xi.r_beta)*sum(allmat(2:end,intersect(s0.infectious,s.h0),:),2)./sum(allmat(2:end,1:size(d0,1),:),2));  % Yearly

    incty = np.squeeze(np.diff(allmat[:,i0['aux']['inc'][0],:],0)) #incty(:,:) = squeeze(diff(allmat(:,i0.aux.inc(1),:),[],1));      %yearly all incidence
    inc_hivy = np.squeeze(np.diff(allmat[:,i0['aux']['inc'][1],:],0)) #inc_hivy(:,:) = squeeze(diff(allmat(:,i0.aux.inc(2),:),[],1));   %yearly hiv-tb incidence
    inc_mdry = np.squeeze(np.diff(allmat[:,i0['aux']['inc'][2],:],0)) #inc_mdry(:,:) = squeeze(diff(allmat(:,i0.aux.inc(3),:),[],1));   %yearly mdr incidence
    mrty = np.squeeze(np.diff(allmat[:,i0['aux']['mort'][0],:],0)) #mrty(:,:)  = squeeze(diff(allmat(:,i0.aux.mort(1),:),[],1));     %yearly hiv-negative TB mortality
    mrt_hivy  = np.squeeze(np.diff(allmat[:,i0['aux']['mort'][1],:],0)) #mrt_hivy(:,:)  = squeeze(diff(allmat(:,i0.aux.mort(2),:),[],1)); %yearly hiv-positive TB mortality
    notiy = np.squeeze(np.sum(np.diff(allmat[:,i0['aux']['noti'],:],0),1)) #notiy(:,:) = squeeze(sum(diff(allmat(:,i0.aux.noti,:),[],1),2)); %yearly notification 
    notimdry = np.squeeze(np.diff(allmat[:,i0['aux']['noti2'],:],0)) #notimdry(:,:) = squeeze(diff(allmat(:,i0.aux.noti2,:),[],1));    %yearly MDR Tx initiation
   
    LTBI_recent = np.squeeze(np.sum(allmat[:,np.intersect1d(s0['Lf'],s0['h0']),:],1)) #LTBI_recent(:,:) = squeeze(sum(allmat(:,intersect(s0.Lf,s0.h0),:),2));   %yearly LTBI (only hiv-ve) 
    LTBI_remote = np.squeeze(np.sum(allmat[:,np.intersect1d(s0['Ls'],s0['h0']),:],1)) #LTBI_remote(:,:) = squeeze(sum(allmat(:,intersect(s0.Ls,s0.h0),:),2));   %yearly LTBI (only hiv-ve)

    LTBI = np.squeeze(np.sum(allmat[:,np.concatenate([s0['Lf'],s0['Ls']]),:],1))  #LTBI(:,:) = squeeze(sum(allmat(:,[s0.Lf,s0.Ls],:),2));         %yearly total LTBI
    pop = np.squeeze(np.sum(allmat[:,np.concatenate([s0['U'],s0['Lf'],s0['Ls'],s0['prevalent'],s0['Rlo'],s0['Rhi'],s0['R']]),:],1)) #pop(:,:) = squeeze(sum(allmat(:,[s0.U,s0.Lf,s0.Ls,s0.prevalent,s0.Rlo,s0.Rhi,s0.R],:),2)); 
    #% TB_woTx(:,:,ii) = squeeze(sum(allmat(2:12:end,s.infectious,:),2));         %yearly Active TB cases without Tx 
   
    Sus1 = np.squeeze(np.sum(allmat[:,np.intersect1d(np.concatenate([s0['U'],s0['Lf'],s0['Ls'], s0['Rlo'], s0['Rhi'], s0['R']]),s0['h0']),:],1)) #Sus1(:,:) = squeeze(sum(allmat(:,intersect([s0.U,s0.Lf,s0.Ls, s0.Rlo, s0.Rhi, s0.R],s0.h0),:),2)); % Susceptiable population including U
    Sus2 = np.squeeze(np.sum(allmat[:,np.intersect1d(np.concatenate([s0['Lf'],s0['Ls'], s0['Rlo'], s0['Rhi'], s0['R']]),s0['h0']),:],1))#Sus2(:,:) = squeeze(sum(allmat(:,intersect([s0.Lf,s0.Ls, s0.Rlo, s0.Rhi, s0.R],s0.h0),:),2));     % Susceptiable population excluding U
        
    tab = {}
    tab['pop'] = np.concatenate([np.array(data['pop00_22'])[np.array(range(len(data['pop00_22'])-3,len(data['pop00_22'])-1))], np.array([43326489, 43044540,	42760880,	42476294,	42190115,	41901694,	41610116,	41314544,	41014863])]) #tab.pop       = [data.pop00_22(end-3:end-1), 43326489, 43044540	42760880	42476294	42190115	41901694	41610116	41314544	41014863]; % pop_19_30;                                        %pop from 2019-2029
    tab['prvty_pct'] = prvty*1e5 #np.transpose(prvty*1e5,[1,0,2]) #tab.prvty_pct = permute(prvty*1e5,[2,1,3]);       %yearly prevalence
    tab['prvtmdry_pct'] = prvt_mdry*1e5 #np.transpose(prvt_mdry*1e5,[1,0,2])#tab.prvtmdry_pct = permute(prvt_mdry*1e5,[2,1,3]);%yearly mdr-TB-prevalence
    tab['prvthivy_pct'] = prvt_hivy*1e5 #np.transpose(prvt_hivy*1e5,[1,0,2]) #tab.prvthivy_pct = permute(prvt_hivy*1e5,[2,1,3]);%yearly hiv-TB-prevalence
    tab['ARTI'] = ARTI*1e2 #np.transpose(ARTI*1e2,[1,0,2]) #tab.ARTI = permute(ARTI*1e2,[2,1,3]);            % In percentage    

    tab['incy_pct']  = incty*1e5 #np.transpose(incty*1e5,[1,0,2]) #tab.incy_pct  = permute(incty*1e5,[2,1,3]);       %yearly incidence
    tab['inchivy_pct']  = inc_hivy*1e5 #np.transpose(inc_hivy*1e5,[1,0,2]) #tab.inchivy_pct  = permute(inc_hivy*1e5,[2,1,3]); %yearly hiv-TB incidence
    tab['incmdry_pct']  = inc_mdry*1e5 #np.transpose(inc_mdry*1e5,[1,0,2]) #tab.incmdry_pct  = permute(inc_mdry*1e5,[2,1,3]); %yearly mdr-incidence

    tab['mrty_pct']  = mrty*1e5 #np.transpose(mrty*1e5,[1,0,2]) #tab.mrty_pct  = permute(mrty*1e5,[2,1,3]);        %yearly hiv-ve mortality
    tab['mrthivy_pct']  = mrt_hivy*1e5 #np.transpose(mrt_hivy*1e5,[1,0,2]) #tab.mrthivy_pct  = permute(mrt_hivy*1e5,[2,1,3]); %yearly hiv-TB mortality

    tab['LTBIrecent_pct'] = LTBI_recent*1e5 #np.transpose(LTBI_recent*1e5,[1,0,2])#tab.LTBIrecent_pct = permute(LTBI_recent*1e5,[2,1,3]);  %yearly LTBI recent(fast)
    tab['LTBIremote_pct'] = LTBI_remote*1e5 #np.transpose(LTBI_remote*1e5,[1,0,2]) #tab.LTBIremote_pct = permute(LTBI_remote*1e5,[2,1,3]);  %yearly LTBI remote(slow)

    tab['LTBI_pct'] = LTBI*1e5 #np.transpose(LTBI*1e5,[1,0,2])#tab.LTBI_pct = permute(LTBI*1e5,[2,1,3]);         %yearly LTBI remote(slow)
    tab['population_pct'] = pop*1e5 #np.transpose(pop*1e5,[1,0,2])#tab.population_pct = permute(pop*1e5,[2,1,3]);    %yearly LTBI remote(slow)

    tab['noti_pct'] = notiy*1e5 #np.transpose(notiy*1e5,[1,0,2]) #tab.noti_pct = permute(notiy*1e5,[2,1,3]);        % Yearly notification
    tab['notimdry_pct'] = notimdry*1e5 #np.transpose(notimdry*1e5,[1,0,2]) #tab.notimdry_pct = permute(notimdry*1e5,[2,1,3]); %yearly mdr Tx initiation
    
    return tend1, tend1b, tend2, inct, mdr, noti, notimdr, tab
    #save res_forward; %with disruption   


tend1, tend1b, tend2, inct, mdr, noti, notimdr, tab = run_simulate_forward('NGA',None)

fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(nrows=2, ncols=2)

ax1.plot(np.array(range(2022,tend2)),inct)
ax1.set_xlabel('Time (year)')
ax1.set_ylabel('Incidence')

ax2.plot(np.array(range(2022,tend2)),mdr)
ax2.set_xlabel('Time (year)')
ax2.set_ylabel('MDR incidence')

ax3.plot(np.array(range(2022,tend2)),noti)
ax3.set_xlabel('Time (year)')
ax3.set_ylabel('Notifications')

ax4.plot(np.array(range(2022,tend2)),notimdr)
ax4.set_xlabel('Time (year)')
ax4.set_ylabel('MDR treatment initiation')

fig.show()
a=0
#%  Figure_final
 
#figure; lw = 1.5;
#subplot(2,2,1)
#plot(2022:tend2-1,inct,'linewidth',lw) %inct
#ylabel('All Incidence')
#subplot(2,2,2)
#plot(2022:tend2-1,mdr,'linewidth',lw)
#ylabel('MDR incidence')
#subplot(2,2,3)
#plot(2022:tend2-1,noti,'linewidth',lw)
#ylabel('Notification')
#subplot(2,2,4)
#pl1 = plot(2022:tend2-1,notimdr,'linewidth',lw);
#ylabel('MDR treatment initiation')
#%legend(pl1,'Baseline','+PPM','+ Improved diagnostics, routine TB services','+ Upstream case-finding','+ case-finding among subclinical','+ Preventive therapy (risk group)','+ Upfront DST','+ Shorter duration of 2ndline Tx','+ Improved secondline Tx success','+ Vaccination','location','SouthEast'); %
#legend(pl1,'Baseline','Interventions without vaccine','+ Vaccination','location','SouthEast'); %
 


