import numpy as np
from .alloc_parameters import alloc_parameters
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

def RunGlobalEpiModel(PARS,IntvnScn): #function [epi_outputs]=RunGlobalEpiModel(PARS,IntvnScn)

    epi_outputs=[] #epi_outputs=[];
    #%clear all; %load calibres2; xsto = xsto4;
    #%load calibres1; 

    xsto2 = PARS[0] #xsto2=PARS{1};
    outsto2 = PARS[1] #outsto2=PARS{2};

    xsto = xsto2 #xsto = xsto2;
    [M,I] = max(outsto2) #[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};


    #%load Model_setup;

    xs = xsto[I,:] #xs = xsto(I,:);  % Most probable parameter set 

    ix0 = 1e4 
    nx = 2 
    dx = round((np.size(xsto,1)-ix0)/nx) #ix0 = 1e4; nx = 2; dx = round((size(xsto,1)-ix0)/nx); %2e4 nx = 150
    xs_all = xsto[ix0::dx,:,0] #xs_all = xsto(ix0:dx:end,:,1);

    pct_xs= np.percenttile(xs_all,[2.5,50,97.5]) #pct_xs= prctile(xs_all,[2.5,50,97.5])


    #% % --- 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; p0 = p #;
    inct = [] #;
    noti = [] #;
    inch1= []
    mdr = []
    morth0= []
    morth1= []
    notimdr= []
    

    #% 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,:] #init = aux.soln(end,1:end);
    
    #% 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 = [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 = [s0['Dx'],s0['Tx'],s0['Tx2']] #inds4 = [s0.Dx,s0.Tx,s0.Tx2];
    inds5 =   s0['Tx2'][-1] + np.array(range(0,8)) #inds5 =  s0.Tx2(end)+[1:9] #inds5 =  s0.Tx2(end)+[1:9];  % Auxiliary terms
    ind = [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[0,i0['nx']] #init0 = zeros(1,i0.nx);
    init0[ind] = init #init0(ind) = init;   

    tref = np.arrange(2022, tend2 + 1, 1) #tref = [2022:1:tend2];
     
    #% --- Now simulate without intervention (baseline)
    M0 = make_model3(p, r, i0, s0, gps0, prm)# M0 = make_model3(p, r, i0, s0, gps0, prm); 
    prm['growth'] = prm['gth'][2] #= prm.gth(3);
    #geq = @(t,in) goveqs_scaleup(t, in, 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, create_steps(2022, tend2, 1/12), args=(M0, M0, [2022, 2023],prm, sel0, p, agg0))

    soln  = soln0
    t     = t0
    #soln0 = interp1(t, soln, tref) #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, 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, create_steps(2022, tend1, 1/12), args=(M0, M0, [2022, 2023],prm, sel0, p, agg0))
    initb = solna[-1,:]
    
    #% --- Next simulate with interventions
    #%r1 = r; p1 = p;
    #%Set intervention to model parameters
    [p1,r1] = InterventionScenario(p,r,IntvnScn) #[p1,r1]=SetInterventionScenario(p,r,IntvnScn);
    
  
    #%     p1.pu  = p.pu + (1- p.pu)*0.35; 
    #%     p1.Dx(1) = 0.95;          % [0.95 0.95];  
    #%     
    #%     r1.cs(4)  = 1.2*r1.cs(4);       %choose specific age group
    #%     r1.cs2(4) = 1.2*r1.cs2(4);      %choose specific age group
    #%     r1.cs3(4) = 1.2*r1.cs(4);       %choose specific age group
    #%     
    #%     qq = 0.05;                                % Preventive therapy among household contacts independent of hiv status
    #%     r1.progression(4,:)  = r.progression(4,:)*(1-qq);    % Replace by this  --> p.PT_PLHIV = [1 1 (1-0.6)]
    #%     r1.reactivation(4,:) = r.reactivation(4,:)*(1-qq);
    #%     
    #%     p.PT_PLHIV = [1 1 (1-0.6)];                % PT among PLHIV reduces rate by 60%
    #%     
    #%     p1.MDR_rec = [0.5,0]; %0.7; %0.8;
    #%     r1.Tx2 = 9/12;                             % Shorter Tx2 regimn (9 months)
    #%     p1.cure2 = [0.7, 0]; %0.8; %0.9;
    #%     r1.default(1) = r.default(1)*(1-0.05);     % reduction of default rate(pu) (by inccreasing Tx completion and reducing LTFU)     
    #%  
    #%     r1.ART_init = xs(16)*(1+0.5);
    #%     
    
    M1 = make_model3(p1, r1, i0, s0, gps0, prm)
      
    #geq = @(t,in) goveqs_scaleup(t, in, 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/12), args=(M0, M1, tend1 + np.array([0, 3]), 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)#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);
    #initv = solnv(end,:);
    tv, solnv = ode15s(goveqs_scaleup, initb, create_steps(tend1, tend1b, 1/12), args=(M0, M1, tend1 + np.array([0, 3]), i0, prm, sel0, agg0))
    initv = solnv[-1,:]

    #% Choose suitable rate of vaccination and vaccine efficacy
    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, 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/12), args=(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 = interp1(t, soln, tref);
    soln9 = np.zeros((len(tref),np.shape(soln)[1])) #soln8 = interp1(t, soln, tref)#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 = 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 #; % hiv-negative mortality
    morth1[:,:] = np.squeeze(dallmat[:,i0['aux']['mort'](2),:])*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
    
    epi_outputs[0] = inct
    epi_outputs[1] = inch1
    epi_outputs[2] = mdr
    epi_outputs[3] = morth0
    epi_outputs[4] = morth1
    epi_outputs[5] = noti
    epi_outputs[6] = notimdr
    
    
    return(epi_outputs)
    
    
  
#%   prvt = squeeze(sum(allmat(:,s0.prevalent,:),2))*1e5;       % Prevalence
     
#% end
#% fprintf('\n');

#% inc_pct  = permute(prctile(inct,[2.5,50,97.5],3)*1e5,[2,1,3]);
#% hiv_pct  = permute(prctile(inch1,[2.5,50,97.5],3)*1e5,[2,1,3]);
#% mdr_pct  = permute(prctile(mdr,[2.5,50,97.5],3)*1e5,[2,1,3]);
#% mrth0_pct  = permute(prctile(morth0,[2.5,50,97.5],3)*1e5,[2,1,3]);
#% mrth1_pct  = permute(prctile(morth1,[2.5,50,97.5],3)*1e5,[2,1,3]);
#% noti_pct = permute(prctile(noti,[2.5,50,97.5],3)*1e5,[2,1,3]);   
#% notimdr_pct = permute(prctile(notimdr,[2.5,50,97.5],3)*1e5,[2,1,3]);        

#% save res_forward; %with disruption   
#%  
#%  Figure_final
 
#figure; lw = 1.5;
#subplot(2,2,1)
#plot(2022:tend2-1,inct,'linewidth',lw)
#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'); %

#end
 