import numpy as np
from scipy import sparse


def make_model3(p, r, i, s, gps,prm) : #function M = make_model(p, r, i, s, gps,prm)  % Kenya

    r = r.copy()
    r['progression'] = np.array(r['progression'])
    r['reactivation'] = np.array(r['reactivation'])
    r['r_sym'] = np.array(r['sym'])

        #% --- Get the linear rates ------------------------------------------------
    m  = np.zeros((i['nstates'],i['nstates'])) #m  = np.zeros(i.nstates);
    m2 = np.zeros((i['nstates'],i['nstates'])) #m2 = np.zeros(i.nstates);
    m4 = np.zeros((i['nstates'],i['nstates']))

    for ia in range(len(gps['age'])): #ia = 1:length(gps.age)
        age  = gps['age'][ia] #age  = gps.age{ia};

        for istr in range(len(gps['strains'])): #istr = 1:length(gps.strains)
            strain = gps['strains'][istr] #strain = gps.strains{istr};
    
            for ih in range(len(gps['hiv'])): #ih = 1:length(gps.hiv)
                hiv = gps['hiv'][ih] #hiv = gps.hiv{ih};
        
                #getst = @(st) i.(st).(age).(hiv).(strain) #getst = @(st) i.(st).(age).(hiv).(strain);
                Isc  = i['Isc'][age][hiv][strain] #Isc  = getst('Isc');     % Subclinical 
                I    = i['I'][age][hiv][strain] #I    = getst('I');       % clinical symptom 
                E    = i['E'][age][hiv][strain] #E    = getst('E');
                Dxs  = i['Dx'][age][hiv][strain] #Dxs  = getst('Dx');
                Txs  = i['Tx'][age][hiv][strain] #Txs  = getst('Tx');
                Tx2s = i['Tx2'][age][hiv][strain] #Tx2s = getst('Tx2');
                Rlo  = i['Rlo'][age][hiv][strain] #Rlo  = getst('Rlo');
                Rhi  = i['Rhi'][age][hiv][strain] #Rhi  = getst('Rhi');
                R    = i['R'][age][hiv][strain] #R    = getst('R');
      
                #% Group Selectors
                ismdr   = strain == 'mdr' #ismdr   = strcmp(strain, 'mdr');
        
                for iv in range(len(gps['vacc'])): #iv = 1:length(gps.vacc)
                    vacc = gps['vacc'][iv] #vacc = gps.vacc{iv};
                    #getst2 = @(st) i.(st).(age).(hiv).(vacc).(strain) #getst2 = @(st) i.(st).(age).(hiv).(vacc).(strain);
                    Lf = i['Lf'][age][hiv][vacc][strain]  #Lf = getst2('Lf') ;
                    Ls = i['Ls'][age][hiv][vacc][strain] #Ls = getst2('Ls');
            
                    #% --- Fast progression and LTBI stabilisation
                    source  = Lf #source  = Lf;
                    destins = np.concatenate([Isc,                                                   Ls]) #destins = [Isc,                                                   Ls];
                    rates   = [r['progression'][ia,ih]*p['PT_PLHIV'][ih]*(1-(iv==1)*p['VE'][1]), r['LTBI_stabil'][ih]] #rates   = [r.progression(ia,ih).*p.PT_PLHIV(ih)*(1-(iv==2)*p.VE(2)), r.LTBI_stabil(ih)];     % p.PT_PLHIV = [1 1 (1-0.6)] during PT among PLHIV
                    m4[destins, source] = m4[destins, source] + rates #m(destins, source) = m(destins, source) + rates';
            
                    #% --- Reactivation
                    source = Ls; destin = Isc; rate = r['reactivation'][ia,ih]*p['PT_PLHIV'][ih]*(1-(iv==1)*p['VE'][1]) #source = Ls; destin = Isc; rate = r.reactivation(ia,ih).*p.PT_PLHIV(ih)*(1-(iv==2)*p.VE(2));  % p.PT_PLHIV = [1 1 (1-0.6)] during PT among PLHIV
                    m4[destins, source] = m4[destins, source] + rate #m(destin, source) = m(destin, source) + rate;
           
                #end
        
                #% --- Sub-clinical to clinical Infection
                source = Isc #source = Isc;
                destin = I #               destin = I;
                rate   = r['r_sym']#rate   = r.sym;                                 % rate of symptom development (estimate this value)
                m[destin, source] = m[destin, source] + rate #m(destin, source) = m(destin, source) + rate;
        
                #% Introduce public/private sector
                #% --- Primary careseeking, including access to public sector care
                source = I; destin = Dxs['pu']; rate = r['cs'][ia]*p['pu'] + r['access'] #source = I; destin = Dxs.pu; rate = r.cs(ia)*p.pu + r.access;
                m2[destin, source] = m2[destin, source] + rate #m2(destin, source) = m2(destin, source) + rate;
        
                source = I; destin = Dxs['pr']; rate = r['cs'][ia]*(1-p['pu']) #source = I; destin = Dxs.pr; rate = r.cs(ia)*(1-p.pu);
                m2[destin, source] = m2[destin, source] + rate #m2(destin, source) = m2(destin, source) + rate;
        
                #% --- Secondary careseeking
                source = E; destin = Dxs['pu']; rate = r['cs2'][ia]*p['pu'] #source = E; destin = Dxs.pu; rate = r.cs2(ia)*p.pu;
                m2[destin, source] = m2[destin, source] + rate #m2(destin, source) = m2(destin, source) + rate;
        
                source = E; destin = Dxs['pr']; rate = r['cs2'][ia]*(1-p['pu']) #source = E; destin = Dxs.pr; rate = r.cs2(ia)*(1-p.pu);
                m2[destin, source] = m2[destin, source] + rate #m2(destin, source) = m2(destin, source) + rate;
        
                #% --- Case finding from subclinical TB
                source = Isc; destin = Dxs['pu']; rate = r['cs3'][ia]*p['pu'] #source = Isc; destin = Dxs.pu; rate = r.cs3(ia)*p.pu ;
                m2[destin, source] = m2[destin, source] + rate #m2(destin, source) = m2(destin, source) + rate;
        
                source = Isc; destin = Dxs['pr']; rate = r['cs3'][ia]*(1-p['pu']) #source = Isc; destin = Dxs.pr; rate = r.cs3(ia)*(1-p.pu);
                m2[destin, source] = m2[destin, source] + rate #m2(destin, source) = m2(destin, source) + rate;
  
        
                for ip in range(len(gps['provs'])): #ip = 1:length(gps.provs)
                    prov = gps['provs'][ip] #prov = gps.provs{ip};
            
                    Dx = Dxs[prov]; Tx = Txs[prov]; Tx2 = Tx2s[prov] #Dx = Dxs.(prov); Tx = Txs.(prov); Tx2 = Tx2s.(prov);
            
                    try:
                        pFLinit = p['Dx'][ip]*p['Tx_init'][ip]*(1 - ismdr*p['MDR_rec'][ip])#pFLinit = p.Dx(ip)*p.Tx_init(ip)*(1 - ismdr*p.MDR_rec(ip)); %*p.Tx_init2(ip)
                    except:
                        KeyboardInterrupt
                    #end
                    pSLinit = p['Dx'][ip]*p['Tx_init2'][ip]*ismdr*p['MDR_rec'][ip] #pSLinit = p.Dx(ip)*p.Tx_init2(ip)*ismdr*p.MDR_rec(ip);
                    p_ltfu  = 1-(pFLinit+pSLinit) #; %p.Dx(ip)*p.Tx_init(ip) #p_ltfu  = 1-(pFLinit+pSLinit); %p.Dx(ip)*p.Tx_init(ip);
            
                    source  = Dx #source  = Dx;
                    destins = np.concatenate([Tx,       Tx2,      E]) #destins =          [Tx,       Tx2,      E];
                    rates   = r['Dx']*np.array([pFLinit,  pSLinit,  p_ltfu]) #rates   = r.Dx*[pFLinit,  pSLinit,  p_ltfu];
                    if(np.shape(rates)==(3,1)):
                        rates = rates.T[0]
                    m[destins, source] = m[destins, source] + rates #m(destins, source) = m(destins, source) + rates';
            
                    #% --- FL Treatment
                    pFLcure = p['cure'][ip]*(1-ismdr) #pFLcure = p.cure(ip)*(1-ismdr);
                    pSLtran = p['SL_trans'][ip]*ismdr #pSLtran = p.SL_trans(ip)*ismdr;
                    rMDRacq = r['MDR_acqu']*(1-ismdr)#                    rMDRacq = r.MDR_acqu*(1-ismdr);
                    if(np.shape(rMDRacq)==(1,)):
                        rMDRacq = rMDRacq[0]
            
                    source  = Tx #source  = Tx;
                    destins = np.concatenate([Rlo,            Tx2,                        E,                              Rhi,             i['Tx'][age][hiv]['mdr'][prov]]) #destins = [Rlo            Tx2                        E                              Rhi,             i.Tx.(age).(hiv).mdr.(prov)];
                    rates   = [r['Tx']*pFLcure,  r['Tx']*(1-pFLcure)*pSLtran,  r['Tx']*(1-pFLcure)*(1-pSLtran),  r['default'][ip],   rMDRacq] #rates   = [r.Tx*pFLcure,  r.Tx*(1-pFLcure)*pSLtran,  r.Tx*(1-pFLcure)*(1-pSLtran),  r.default(ip),   rMDRacq];
                    m[destins, source] = m[destins, source] + rates #m(destins, source) = m(destins, source) + rates';
            
                    #% --- SL Treatment
                    source  = Tx2 #source  = Tx2;
                    destins = np.concatenate([Rlo,                 E]) #destins = [Rlo                 E];          %Rhi];
                    rates   = [r['Tx2']*p['cure2'][ip],  r['Tx2']*(1-p['cure2'][ip])+ r['default2'][ip]] #rates   = [r.Tx2*p.cure2(ip),  r.Tx2*(1-p.cure2(ip))+ r.default2(ip)];
                    m[destins, source] = m[destins, source] + rates#m(destins, source) = m(destins, source) + rates';
            
            
                #end  % provider loop
        
                #% --- Relapse
                sources = np.concatenate([Rlo, Rhi, R]) #sources = [Rlo, Rhi, R];
                destin  = Isc #destin  = Isc;
                rates   = r['relapse'] #rates   = r.relapse;
                m[destin, sources] = m[destin, sources] + rates
        
                sources = [Rlo, Rhi] #sources = [Rlo, Rhi];
                destin  = R#destin  = R;
                rates   = 0.5
                m[destin, sources] = m[destin, sources] + rates #m(destin, sources) = m(destin, sources) + rates;
        
                #% --- Self cure
                sources = np.intersect1d(np.intersect1d(np.intersect1d(s['infectious'],s[hiv]), s[strain]),s[age]) #sources = intersect(intersect(intersect(s.infectious,s.(hiv)), s.(strain)),s.(age)); 
                destin  = Rhi #destin  = Rhi;
                rates   = r['self_cure'][ih] #rates   = r.self_cure(ih);
                m[destin, sources] = m[destin, sources] + rates #m(destin, sources) = m(destin, sources) + rates;
        
            #end  % hiv loop
        #end % MDR loop

        for ih in range(len(gps['hiv'])) :#ih = 1:length(gps.hiv)
            hiv = gps['hiv'][ih] #hiv = gps.hiv{ih};
    
            source = i['U'][age][hiv]['v0']#source = i.U.(age).(hiv).v0;
            destin = i['U'][age][hiv]['v1'] #destin = i.U.(age).(hiv).v1;
            rate   = r['vacc'][ia] #rate   = r.vacc(ia);
            m[destin,source] = m[destin,source] + rate #m(destin,source) = m(destin,source) + rate;
    
            source = i['U'][age][hiv]['v1'] #source = i.U.(age).(hiv).v1;
            destin = i['U'][age][hiv]['v2'] #destin = i.U.(age).(hiv).v2;
            rate   = r['waning'] #rate   = r.waning;
            m[destin,source] = m[destin,source] + rate#m(destin,source) = m(destin,source) + rate;
    
            for istr in range(len(gps['strains'])): #istr = 1:length(gps.strains)
                strain = gps['strains'][istr] #strain = gps.strains{istr};
        
                source = i['Lf'][age][hiv]['v0'][strain]#source = i.Lf.(age).(hiv).v0.(strain);
                destin = i['Lf'][age][hiv]['v1'][strain] #destin = i.Lf.(age).(hiv).v1.(strain);
                rate   = r['vacc'][ia] #rate   = r.vacc(ia);
                m[destin,source] = m[destin,source] + rate #m(destin,source) = m(destin,source) + rate;
        
                source = i['Lf'][age][hiv]['v1'][strain] #source = i.Lf.(age).(hiv).v1.(strain);
                destin = i['Lf'][age][hiv]['v2'][strain] #destin = i.Lf.(age).(hiv).v2.(strain);
                rate   = r['waning'] #rate   = r.waning;
                m[destin,source] = m[destin,source] + rate #m(destin,source) = m(destin,source) + rate;
            #end
        #end
    #end

#% --- Introduction of Ageing rate                                           
    sources = s['a1'] #sources = s.a1;
    destins = s['a2'] #destins = s.a2;
    
    inds = np.ravel_multi_index((destins, sources), m.shape, order='C') ##inds = np.ravel_multi_index((destins, sources), m.shape, order='F') #inds    = sub2ind(size(m), destins, sources);
    #inds = np.ravel_multi_index([destins, sources], dims=np.shape(m), order='F')

    rate   = r['ageing'][0] #rate   = r.ageing(1);
    m.ravel()[inds] = m.ravel()[inds] + rate

    sources = s['a2'] #sources = s.a2;
    destins = s['a3'] #destins = s.a3;
    inds = np.ravel_multi_index((destins, sources), m.shape, order='C') ##inds = np.ravel_multi_index((destins, sources), m.shape, order='F') # inds    = sub2ind(size(m), destins, sources);
    rate   = r['ageing'][1] #rate   = r.ageing(2);
    m.ravel()[inds] = m.ravel()[inds] + rate #m(inds) = m(inds) + rate;

    sources = s['a3'] #sources = s.a3;
    destins = s['a4'] #destins = s.a4;
    inds = np.ravel_multi_index((destins, sources), m.shape, order='C') #inds = np.ravel_multi_index((destins, sources), m.shape, order='F') #inds    = sub2ind(size(m), destins, sources);
    rate   = r['ageing'][2] #rate   = r.ageing(3);
    m.ravel()[inds] = m.ravel()[inds] + rate #m(inds) = m(inds) + rate;

    sources = s['a4']#sources = s.a4;
    destins = s['a5'] #destins = s.a5;
    inds = np.ravel_multi_index((destins, sources), m.shape, order='C') ##inds = np.ravel_multi_index((destins, sources), m.shape, order='F') #inds    = sub2ind(size(m), destins, sources);
    rate   = r['ageing'][3] #rate   = r.ageing(4);
    m.ravel()[inds] = m.ravel()[inds] + rate #m(inds) = m(inds) + rate;

    #% --- HIV acquisition                                                      % <--- NEW
    m3 = np.zeros((i['nstates'],i['nstates'])) #m3 = zeros(i.nstates);
    sources = s['h0'] #sources = s.h0;
    destins = s['h1'] #destins = s.h1;
    inds = np.ravel_multi_index((destins, sources), m.shape, order='C') #inds = np.ravel_multi_index((destins, sources), m.shape, order='F') #inds    = sub2ind(size(m), destins, sources);
    rates   = 1 #rates   = 1;
    m3.ravel()[inds] = m3.ravel()[inds] + rates #   m3(inds) = m3(inds) + rates;

    #% --- ART initiation
    sources = s['h1'] #sources = s.h1;
    destins = s['hart'] #destins = s.hart;
    inds = np.ravel_multi_index((destins, sources), m.shape, order='C') #inds = np.ravel_multi_index((destins, sources), m.shape, order='F') ##inds = np.ravel_multi_index((destins, sources), m.shape, order='F') #inds    = sub2ind(size(m), destins, sources);
    rates   = r['ART_init'] #rates   = r.ART_init;
    m.ravel()[inds] = m.ravel()[inds] + rates #m(inds) = m(inds) + rates;

    #% --- Bring them together
    M = {}
    M['lin']   = sparse.csc_matrix(m - np.diag(np.sum(m,0))) #M.lin    = sparse(m - diag(sum(m,1)));
    M['Dxlin']  = sparse.csc_matrix(m2 - np.diag(np.sum(m2,0))) #M.Dxlin  = sparse(m2 - diag(sum(m2,1)));
    M['linHIV'] = sparse.csc_matrix(m3 - np.diag(np.sum(m3,0))) #M.linHIV = sparse(m3 - diag(sum(m3,1)));                                   % <--- NEW, now HIV acquisition is separate matrix made in lines 62-67
    M['linLTBI'] = sparse.csc_matrix(m4 - np.diag(np.sum(m4,0))) #M.linLTBI = sparse(m4 - diag(sum(m4,1)));
    #% --- Get the nonlinear rates ---------------------------------------------
    M['nlin'] = {}
    for ia in range(len(gps['age'])):
        age  = gps['age'][ia]
        M['nlin'][age] = {}

    #% --- Allocating transitions
    for istr in range(len(gps['strains'])): #istr = 1:length(gps.strains)
        strain = gps['strains'][istr]
      
        for ia in range(len(gps['age'])): #ia = 1:length(gps.age)
            age  = gps['age'][ia]
    
            m = np.zeros((i['nstates'],i['nstates'])) #m = zeros(i.nstates);
            for ih in range(len(gps['hiv'])) : #ih = 1:length(gps.hiv)
                hiv = gps['hiv'][ih] #hiv = gps.hiv{ih};
        
                for iv in range(len(gps['vacc'])): #iv = 1:length(gps.vacc)
                    vacc = gps['vacc'][iv] #vacc = gps.vacc{iv};
            
                    tidx0 = np.intersect1d(np.intersect1d(np.intersect1d(np.intersect1d(np.concatenate([s['Lf'], s['Ls'], s['Rlo'], s['Rhi'], s['R']]),s[age]),s[hiv]),s[strain]),s[vacc])
                    sources = [i['U'][age][hiv][vacc], tidx0] #sources = [i.U.(age).(hiv).(vacc), intersect(intersect(intersect(intersect([s.Lf, s.Ls, s.Rlo, s.Rhi, s.R],s.(age)),s.(hiv)),s.(strain)),s.(vacc))];
                    if len(tidx0) == 0:
                        sources = [i['U'][age][hiv][vacc]]
                    else:
                        sources = np.concatenate(sources)      
                    destin  = i['Lf'][age][hiv][vacc][strain] #destin  = i.Lf.(age).(hiv).(vacc).(strain);
                    m[destin, sources] = m[destin, sources] + (1-(iv==1)*p['VE'][0])#m(destin, sources) = m(destin, sources) + (1-(iv==2)*p.VE(1));
            
                    #%M.nlin.(hiv).(strain) = sparse(m - diag(sum(m,1)));
                #end
        
                #% Adjust for any immune protection
                cols = np.intersect1d(np.concatenate([s['Lf'], s['Ls'], s['Rlo'], s['Rhi'], s['R']]), s[hiv]) #cols = intersect([s.Lf, s.Ls, s.Rlo, s.Rhi, s.R], s.(hiv));
                m[:,cols] = m[:,cols]*(1-p['imm'][ih]) #m(:,cols) = m(:,cols)*(1-p.imm(ih));
            #end
            M['nlin'][age][strain] = sparse.csc_matrix(m - np.diag(np.sum(m,0))) #M.nlin.(age).(strain) = sparse(m - diag(sum(m,1)));
        #end
    #end

    #% construct basic building block of repeating contact matrix, as per five-age category
    tmp = np.zeros((len(gps['age']),i['nstates'])) #tmp = zeros(length(gps.age),i.nstates);
    for ia_sus in range(len(gps['age'])): #ia_sus = 1:length(gps.age)
        for ia_inf in range(len(gps['age'])): #ia_inf = 1:length(gps.age)
            cols = np.intersect1d(s['infectious'],s[gps['age'][ia_inf]]) #cols = intersect(s.infectious,s.(gps.age{ia_inf}));
            tmp[ia_sus,cols] = np.array(prm['cm'])[ia_sus,ia_inf] #tmp(ia_sus,cols) = prm.cm(ia_sus,ia_inf);
        #end
    #end    

    #% --- Getting force-of-infection
    m = np.zeros((10,i['nstates'])) #m = zeros(10,i.nstates);            % DS/mdr and 5 age-groups
    m[0:5,np.intersect1d(s['infectious'], s['ds'])]   = r['beta']*tmp[:, np.intersect1d(s['infectious'], s['ds'])] #m(1:5,intersect(s.infectious, s.ds))   = r.beta*tmp(:,intersect(s.infectious, s.ds));
    m[5:10,np.intersect1d(s['infectious'], s['mdr'])]  = r['beta']*p['rfbeta_mdr']*tmp[:,np.intersect1d(s['infectious'], s['mdr'])] #%r.beta*p.rfbeta_mdr #m(6:10,intersect(s.infectious, s.mdr))  = r.beta*p.rfbeta_mdr*tmp(:,intersect(s.infectious, s.mdr)); %r.beta*p.rfbeta_mdr;
    m[:,[s['h1'], s['hart']]] = m[:,[s['h1'], s['hart']]]*p['rfbeta_hiv'] #m(:,[s.h1, s.hart]) = m(:,[s.h1, s.hart])*p.rfbeta_hiv;
    m[:,s['Isc']] = m[:,s['Isc']]*0.75 #m(:,s.Isc) = m(:,s.Isc)*0.75;
    M['lambda'] = sparse.csc_matrix(m) #M.lambda = sparse(m);

 
    #% --- Get the mortality rates
    m = np.zeros((i['nstates'],5)) #m = zeros(i.nstates,5);
    #%m(:,1)            = r.mort;
    m[s['a1'],0]         = r['mort'][0] #m(s.a1,1)         = r.mort(1);
    m[s['a2'],0]         = r['mort'][1] #m(s.a2,1)         = r.mort(2);
    m[s['a3'],0]         = r['mort'][2] #m(s.a3,1)         = r.mort(3);
    m[s['a4'],0]         = r['mort'][3]#m(s.a4,1)         = r.mort(4);
    m[s['a5'],0]         = r['mort'][4] #m(s.a5,1)         = r.mort(5);

    m[s['h1'],0]         = r['HIV_mort'] #m(s.h1,1)         = r.HIV_mort;                                            % HIV+ve, no TB
    #% HIV-ve TB
    inds = np.intersect1d([s['h0']], s['infectious_wosc']) #inds = intersect([s.h0], s.infectious_wosc);
    m[inds,1] = r['mort_TB'][0] #m(inds,2) = r.mort_TB(1);
    #% HIV/TB coinfection
    inds = np.intersect1d([s['h1'],s['hart']], s['infectious_wosc']) #inds = intersect([s.h1,s.hart], s.infectious_wosc);
    m[inds,2] = r['mort_TB'][1] #m(inds,3) = r.mort_TB(2);
    #%keyboard

    M['mortvec'] = m #M.mortvec = m;

    return(M)