import numpy as np
from scipy import sparse

def make_model(p, r, i, s, gps):
    
    # --- Get the linear rates ------------------------------------------------
    size = [i['nstates']]*2
    m    = np.zeros(size)
    m2 = np.zeros(size)
    #m2pr = np.zeros(size)

    for istr in gps['strains']:#for ip in range(len(gps['strains'])):
        strain = istr#gps['strains'][ip]
        for ih in gps['hiv']:#for ih in range(len(gps['hiv'])):
            hiv = ih#gps['hiv'][ih]
            idx_ih = (ih=='h1') + 2*(ih=='hart') 
            Isc = i['Isc'][hiv][strain]#i['st']['hiv']['strain']['Isc']
            I = i['I'][hiv][strain]#i['st']['hiv']['strain']['I']
            E = i['E'][hiv][strain]
            Dxs = i['Dx'][hiv][strain]
            Txs = i['Tx'][hiv][strain]
            Tx2s = i['Tx2'][hiv][strain]
            Rlo = i['Rlo'][hiv][strain]
            Rhi = i['Rhi'][hiv][strain]
            R = i['R'][hiv][strain]

            #% Group Selectors
            ismdr = strain=='mdr'
            for iv in gps['vacc']:
                #vacc = gps['vacc'][iv]
                idx_iv = (iv=='v1') + 2*(iv=='v2') 
                Lf = i['Lf'][hiv][iv][strain]
                Ls = i['Ls'][hiv][iv][strain]

                #% --- Fast progression and LTBI stabilisation
                m[Isc][Lf] = m[Isc][Lf] + r['progression'][idx_ih]*p['PT_PLHIV'][idx_ih]*(1-(idx_iv==1)*p['VE'][1])
                m[Ls][Lf] = m[Ls][Lf] + r['LTBI_stabil'][idx_ih]
    
                # --- Reactivation
                m[Isc][Ls] = m[Isc][Ls] + r['reactivation'][idx_ih]*p['PT_PLHIV'][idx_ih]*(1-(idx_iv==1)*p['VE'][1])

            #% --- Sub-clinical to clinical Infection    
            m2[I][Isc] = m2[I][Isc] + r['r_sym']

            # % Introduce public/private sector
            # % --- Primary careseeking, including access to public sector care
            m2[Dxs['pu']][I] = m2[Dxs['pu']][I] + r['cs']*p['pu'] + r['access']
            m2[Dxs['pr']][I] = m2[Dxs['pr']][I] + r['cs']*(1-p['pu'])

            #% --- Secondary careseeking
            m2[Dxs['pu']][E] = m2[Dxs['pu']][E] + r['cs2']*p['pu']
            m2[Dxs['pr']][E] = m2[Dxs['pr']][E] + r['cs2']*(1-p['pu'])

            #% --- Case finding from subclinical TB
            m2[Dxs['pu']][Isc] = m2[Dxs['pu']][Isc] + r['cs3']*p['pu']
            m2[Dxs['pr']][Isc] = m2[Dxs['pr']][Isc] + r['cs3']*(1-p['pu'])

            for ip in gps['provs']:#for ip in range(len(gps['provs'])):
                prov = ip
                idx_ip = 0+(ip=='pr')
                Dx = Dxs[prov]
                Tx = Txs[prov]
                Tx2 = Tx2s[prov]

                try: 
                  pFLinit = p['Dx'][idx_ip]*p['Tx_init'][idx_ip]*(1-ismdr*p['MDR_rec'][idx_ip])
                except:
                    KeyboardInterrupt

                pSLinit = p['Dx'][idx_ip]*p['Tx_init2'][idx_ip]*ismdr*p['MDR_rec'][idx_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);
                m[E][Dx] = m[E][Dx] + r['Dx']*p_ltfu
                m[Tx2][Dx] = m[Tx2][Dx] + r['Dx']*pSLinit
                m[Tx][Dx] = m[Tx][Dx] + r['Dx']*pFLinit
                #% --- FL Treatment
                pFLcure = p['cure'][idx_ip]*(1-ismdr) #            pFLcure = p.cure(ip)*(1-ismdr);
                pSLtran = p['SL_trans'][idx_ip]*ismdr #pSLtran = p.SL_trans(ip)*ismdr;
                rMDRacq = r['MDR_acqu']*(1-ismdr) #

                #source  = Tx;
                #destins = [Rlo            Tx2                        E                              Rhi,             i.Tx.(hiv).mdr.(prov)];
                #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[i['Tx'][hiv]['mdr'][prov]][Tx] = m[i['Tx'][hiv]['mdr'][prov]][Tx] + rMDRacq
                m[Rhi][Tx] = m[Rhi][Tx] + r['default'][idx_ip]
                m[E][Tx] = m[E][Tx] + r['Tx']*(1-pFLcure)*(1-pSLtran)
                m[Tx2][Tx] = m[Tx2][Tx] + r['Tx']*(1-pFLcure)*pSLtran
                m[Rlo][Tx] = m[Rlo][Tx] + r['Tx']*pFLcure

                #% --- SL Treatment
                #source  = Tx2;
                #destins = [Rlo                 E];          %Rhi];
                #rates   = [r.Tx2*p.cure2(ip),  r.Tx2*(1-p.cure2(ip))+ r.default2(ip)];
                #m(destins, source) = m(destins, source) + rates';
                m[E][Tx2] = m[E][Tx2] + r['Tx2']*(1-p['cure2'][idx_ip])+ r['default2'][idx_ip]
                m[Rlo][Tx2] = m[Rlo][Tx2] + r['Tx2']*p['cure2'][idx_ip]
                #End Provider Loop

            #% --- Relapse
            #sources = [Rlo, Rhi, R];
            #destin  = Isc;
            #rates   = r.relapse;
            #m(destin, sources) = m(destin, sources) + rates;
            #m[Isc][R] = m[Isc][R] + r['relapse']
            #m[Isc][Rhi] = m[Isc][Rhi] + r['relapse']
            #m[Isc][Rlo] = m[Isc][Rlo] + r['relapse']
             
            m[Isc][[Rlo,Rhi,R]] = m[Isc][[Rlo,Rhi,R]] + r['relapse']
            
            #sources = [Rlo, Rhi];
            #destin  = R;
            #rates   = 0.5;
            #m(destin, sources) = m(destin, sources) + rates;
            #m[R][Rhi] = m[R][Rhi] + 0.5
            #m[R][Rlo] = m[R][Rlo] + 0.5
            m[R][[Rhi,Rlo]] = m[R][[Rhi,Rlo]] + 0.5
        
            #% --- Self cure
            #sources = intersect(intersect(s.infectious,s.(hiv)), s.(strain));
            #destin  = Rhi;
            #rates   = r.self_cure(ih);
            #m(destin, sources) = m(destin, sources) + rates; 
            #for source in s['infectious'].intersection(s['hiv']).intersection(s['strain']): 
            #    m[Rhi][source] = m[Rhi][source] + r['self_cure'][ih]  
            source = np.intersect1d(np.intersect1d(s['infectious'],s[hiv]),s[strain])    
            m[Rhi][source] = m[Rhi][source] + r['self_cure'][idx_ih] 
        #End hiv loop
    #End MDR loop        

    for ih in gps['hiv']:#for ih in range(len(gps['hiv'])):
        hiv = ih#gps['hiv'][ih]
        source = i['U'][hiv]['v0']
        destin = i['U'][hiv]['v1']
        m[destin][source] = m[destin][source] + r['vacc']

        source = i['U'][hiv]['v1']
        destin = i['U'][hiv]['v2']
        m[destin][source] = m[destin][source] + r['waning']

        for istr in gps['strains']:#for istr in range(len(gps['strains'])):
            strain = istr#gps['strains'][istr]
            source = i['Lf'][hiv]['v0'][strain]
            destin = i['Lf'][hiv]['v1'][strain]
            m[destin][source] = m[destin][source] + r['vacc']

            source = i['Lf'][hiv]['v1'][strain]
            destin = i['Lf'][hiv]['v2'][strain]
            m[destin][source] = m[destin][source] + r['waning']

    #% --- 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.sub2ind(size(m), destins, sources) #;#inds    = sub2ind(size(m), destins, sources);
    inds = np.ravel_multi_index([destins, sources], dims=np.shape(m3), order='F')
    #rates   = 1;
    #m3[inds] = m3[inds] + 1#m3(inds) = m3(inds) + rates;
    #m3.ravel()[inds] = m3.ravel()[inds] + 1#m3(inds) = m3(inds) + rates;
    m3[destins,sources] = m3[destins,sources] + 1

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

    #% --- 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['nlin'] = {}
    #% --- Get the nonlinear rates ---------------------------------------------
    #% --- Allocating transitions
    for istr in gps['strains']:#for istr in range(len(gps['hiv'])):
        strain = istr

        m = np.zeros((i['nstates'],i['nstates']))
        for ih in gps['hiv']:#for istr in range(len(gps['hiv'])):
            hiv = ih
            idx_ih = (ih=='h1') + 2*(ih=='hart')
            for iv in gps['vacc']:#for iv in range(len(gps['vacc'])):
               vacc = iv
               idx_iv = (iv=='v1')+2*(iv=='v2')
               #sources = [i['U'][hiv][vacc], (([s['Lf'], s['Ls'], s['Rlo'], s['Rhi'], s['R']].intersecion(s[hiv])).intersection(s[strain])).intersection(s[vacc])]
               sources = np.concatenate([[i['U'][hiv][vacc]], np.intersect1d(np.intersect1d(np.intersect1d(np.concatenate([s['Lf'], s['Ls'], s['Rlo'], s['Rhi'], s['R']]),s[hiv]),s[strain]),s[vacc])])
               destin = i['Lf'][hiv][vacc][strain] #destin  = i.Lf.(hiv).(vacc).(strain);
               m[destin][sources] = m[destin][sources] + (1-(idx_iv==1)*p['VE'][0])
               #%M.nlin.(hiv).(strain) = sparse(m - diag(sum(m,1)));
        

            #% Adjust for any immune protection
            cols = np.intersect1d(np.concatenate([s['Lf'], s['Ls'], s['Rlo'], s['Rhi'], s['R']]),s[hiv])
            m[:,cols] = m[:,cols]*(1-p['imm'][idx_ih])#m(:,cols) = m(:,cols)*(1-p.imm(ih));

        M['nlin'][strain] = sparse.csc_matrix(m - np.diag(np.sum(m,0)))#M.nlin.(strain) = sparse(m - diag(sum(m,1)));

    #% --- Getting force-of-infection
    m = np.zeros((2,i['nstates'])) #m = zeros(2,i.nstates);
    m[0,np.intersect1d(s['infectious'],s['ds'])]  = r['beta'] #m(1,intersect(s.infectious, s.ds))  = r.beta;
    m[1,np.intersect1d(s['infectious'],s['mdr'])] = r['beta']*p['rfbeta_mdr'] #m(2,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'],3))
    m[:,0]            = r['mort']#m(:,1)            = r.mort;
    m[s['h1'],0]      = 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(np.concatenate([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
