import numpy as np
from .make_model3 import make_model3

#from make_model import make_model
from .goveqs_scaleup import goveqs_scaleup
from .ode15s import ode15s, create_steps

from .InterventionScenario import InterventionScenario

from scipy import interpolate
import collections.abc


def RunGlobalEpiModelb(PARS,IntvnScn, bDoBaseline=False): # function [epi_outputs]=RunGlobalEpiModelb(PARS,IntvnScn)

    epi_outputs=[] #epi_outputs=[];

    #xsto2 = PARS[0] #xsto2=PARS{1};
    #outsto2 = PARS[1] #outsto2=PARS{2};
    
    #xsto = xsto2
    #M , I = outsto2.max(), outsto2.argmax() #[M,I] = max(outsto2);

    xi = PARS[2] #xi=PARS{3};
    obj= PARS[3] #obj=PARS{4};

    i0 = PARS[4] #i0=PARS{5};
    s0 = PARS[5] #s0=PARS{6};
    gps0=PARS[6] #gps0=PARS{7};
    sel0=PARS[7] #sel0=PARS{8};
    agg0=PARS[8] #agg0=PARS{9};
    r = PARS[9] #r=PARS{10};
    p = PARS[10] #p=PARS{11};
    prm=PARS[11] #prm=PARS{12};
    init0 = PARS[12] #init0

    #%load Model_setup;
    #xs = xsto[I,:] #xs = xsto(I,:);  % Most probable parameter set #Commented out on June 09 2024 after Carel suggested we stop passing xs as a parameter

    #% % --- Get the Regional notifications
    tend1   = IntvnScn['year1'] #tend1   = IntvnScn.year1;%2024;  % Intervention starts except vaccination  
    tend1b  = IntvnScn['year2'] #tend1b  = IntvnScn.year2;%2028;  % Vaccination starts
    tend2   = IntvnScn['year3'] #tend2   = IntvnScn.year3;%2036;  % End date for simulation

    #% --- Do the simulations 
    opts = {'NonNegative': list(range(1, i0['nstates'] + 1)), 'rtol': 1e-10, 'atol': 1e-10, 'max_step': 1/64}#opts = odeset('NonNegative',[1:i0.nstates],'Refine',64,'AbsTol',1e-10,'RelTol',1e-10);

    r0 = r.copy(); p0 = p.copy()

    inct   = []
    inch1  = []
    mdr    = []
    morth0 = []
    morth1 = []
    noti   = []
    notimdr= []
    prevtb = []


    #[r,p] = alloc_parameters(xs,prm['r'],prm['p'],xi) #[r,p] = alloc_parameters(xs,prm.r,prm.p,xi);
    #% Get the initial conditions
    init = init0.copy()
    if init == []: #Check if init0 was provided
        [out, aux] = obj() #[out, aux] = obj(xs);
        init = aux['soln'][-1,:] #init = aux.soln(end,1:end);    

    #init_matlab = np.loadtxt("C:\\Users\\GuyMahiane\\Documents\\GM_work\\Carel\\20231010_TB_IHT\\2024-07-11\\NGA_final_vaccine2f\\init0.txt")
    #init = init_matlab

    #% 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([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(1,10)) #inds5 =  s0.Tx2(end)+[1:9] #inds5 =  s0.Tx2(end)+[1:9];  % Auxiliary terms
    ind = np.concatenate([inds1, inds2, inds3,inds4,inds5]) #ind = [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 = np.arange(2022, tend2 + 1, 1) #tref = [2022:1:tend2];

    # % --- Now simulate without intervention (baseline)
    p['VE'] = np.array([0.7, 0.7]) #Carel asked to add this on 2024-10-18
    #p['VE'] = np.array([0.0, 0.0]) 
    M0 = make_model3(p, r, i0, s0, gps0, prm)#M0 = make_model3(p, r, i0, s0, gps0, prm); 
    #**** Apparently, this was changed at some point. Guy is making a change on 2024-07-24
    #prm['growth'] = prm['gth'][2] #prm.growth = prm.gth(3);
    prm['growth'] = prm['gth'][3] #prm.growth = prm.gth(4);
    #****End 

    #% Find the initial condition at the starting point for the interventions
    #geq = @(t,in) goveqs_scaleup(t, in, 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[0], 1/2), args=(r, M0, M0, [2022, 2023], i0, prm, sel0, agg0))
    initb = solna[-1,:]

    #% --- Next simulate with interventions
    #%r1 = r; p1 = p;
    #%Set intervention to model parameters
    r1 = r.copy()
    p1 = p.copy()
    
    if not bDoBaseline:
        [p1,r1] = InterventionScenario(p,r,IntvnScn) #[p1,r1]=SetInterventionScenario(p,r,IntvnScn);

    if isinstance(r1['Tx2'], (collections.abc.Sequence, np.ndarray)) :
        if len(r1['Tx2'])==1:
            r1['Tx2'] = r1['Tx2'][0]

    for cur_x in ['cure', 'cure2']:
        if cur_x in list(p1):
            if len(p1[cur_x])>=1:
                for ii in range(len(p1[cur_x])):
                    if isinstance(p1[cur_x][ii], (collections.abc.Sequence, np.ndarray)):
                        p1[cur_x][ii] = p1[cur_x][ii][0]  

    M1 = make_model3(p1, r1, i0, s0, gps0, prm) #M1 = make_model3(p1, r1, i0, s0, gps0, prm);

    #% p1
    #% r1
    #% pause

    #geq = @(t,in) goveqs_scaleup(t, in, M0, M1, [tend1 tend1b], i0, prm, sel0, agg0);
    #[tb, solnb] = ode15s(geq, [tend1 tend2], initb, opts);

    tb, solnb = ode15s(goveqs_scaleup, initb, create_steps(tend1[0], tend2[0], 1/2), args=(r1, M0, M1, np.array([tend1[0], tend1b[0]]), i0, prm, sel0, agg0))
    initb = solnb[-1,:]

    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]


    allmat = soln8 
    dallmat = np.diff(allmat,n=1,axis=0) #dallmat = diff(allmat,[],1);

    #pop = np.squeeze(np.sum(dallmat[:,np.concatenate([s0['U'],s0['Lf'],s0['Ls'],s0['prevalent'],s0['Rlo'],s0['Rhi'],s0['R']])],1))#pop(:,:) = squeeze(sum(allmat(1:end-1,[s0.U,s0.Lf,s0.Ls,s0.prevalent,s0.Rlo,s0.Rhi,s0.R],:),2)); %Total population
    #pop = np.squeeze(np.sum(allmat[:,:],1))[:-1]#pop(:,:) = squeeze(sum(allmat(1:end-1,[s0.U,s0.Lf,s0.Ls,s0.prevalent,s0.Rlo,s0.Rhi,s0.R],:),2)); %Total population
    pop = np.sum(soln[0:(np.shape(allmat)[0]-1),0:(i0['nstates']-1)],1)

    inct = np.squeeze(dallmat[:,i0['aux']['inc'][0]])/pop*1e5 #inct(:,:) = squeeze(dallmat(:,i0.aux.inc(1),:))*1e5;    % Total incidence
    inch1 = np.squeeze(dallmat[:,i0['aux']['inc'][1]])/pop*1e5 #inch1(:,:) = squeeze(dallmat(:,i0.aux.inc(2),:))*1e5;   % HIV incidence
    mdr = np.squeeze(dallmat[:,i0['aux']['inc'][2]])/pop*1e5 #mdr(:,:) = squeeze(dallmat(:,i0.aux.inc(3),:))*1e5;     % MDR incidence
    morth0 = np.squeeze(dallmat[:,i0['aux']['mort'][0]])/pop*1e5 #; % hiv-negative mortality
    morth1 = np.squeeze(dallmat[:,i0['aux']['mort'][1]])/pop*1e5 #morth1(:,:) = squeeze(dallmat(:,i0.aux.mort(2),:))*1e5; % hiv-positive mortality
    noti = np.squeeze(np.sum(dallmat[:,i0['aux']['noti']],1))/pop*1e5 #noti(:,:) = squeeze(sum(dallmat(:,i0.aux.noti,:),2))*1e5; % Total notification
    notimdr = np.squeeze(dallmat[:,i0['aux']['noti2']])/pop*1e5 #notimdr(:,:) = squeeze(dallmat(:,i0.aux.noti2,:))*1e5;  % MDR Treatment initiation
    prevtb = np.squeeze(np.sum(allmat[1:,s0['prevalent']],1))*1e5 #;%TB prevalence #prevtb(:,: ) = squeeze(sum(allmat(2:end,s0.prevalent,:),2))*1e5;%TB prevalence
  
    
    #pdb.set_trace() 

    epi_outputs = []
    epi_outputs.append(inct)#epi_outputs[0] = inct
    epi_outputs.append(inch1) #epi_outputs[1] = inch1
    epi_outputs.append(mdr) #epi_outputs[2] = mdr
    epi_outputs.append(morth0) #epi_outputs[3] = morth0
    epi_outputs.append(morth1) #epi_outputs[4] = morth1
    epi_outputs.append(noti) #epi_outputs[5] = noti
    epi_outputs.append(notimdr) #epi_outputs[6] = notimdr
    epi_outputs.append(prevtb)
    
    return(epi_outputs)

#end

