import numpy as np
from scipy import interpolate

def goveqs_basis2(input, t,  r, M, i, prm, sel, agg): # % <--- Arguments changed to accommodate prm
    out = np.zeros(i['nx'])#len(input)
    invec = input[0:i['nstates']]

    if t<1980 :
        rHIV=0
    else:
        tempfn = interpolate.interp1d(np.array(range(len(prm['rHIV']))), prm['rHIV'], fill_value="extrapolate")
        ti = t-1980
        rHIV = tempfn(ti) 

    #% r.progression = r.progression.*repelem(((r.betared).^max((t-2003),0)),5)';   % for 5 age groups
    #fac = pow(r['betared'],max(t-2000, 0))#fac = pow(r['betared'][0],max(t-2000, 0))#fac = pow(r['betared'][0],max(t-2003, 0)) #fac = (r.betared)^max(t-2003, 0); % trend on activation and progression rate
    fac_t  = t #fac_t  = min(t,2022); #Carel changed this on 2024-10-18 
    beta_red = r['betared']
    if ((fac_t>=2030) & (fac_t<=2050)):
       #beta_rem = np.linspace(beta_red,1,70)
       #tempfnbta = interpolate.interp1d(np.array(range(len(beta_rem))), beta_rem, fill_value="extrapolate")#beta_red=interp1([2030:2100],[beta_rem(1),beta_rem],fac_t);
       #beta_red = tempfnbta(fac_t-1980)
       beta_rem = np.linspace(beta_red,min(1,beta_red),50)
       tempfnbta = interpolate.interp1d(np.array(range(len(beta_rem))), beta_rem, fill_value="extrapolate")#beta_red=interp1([2030:2100],[beta_rem(1),beta_rem],fac_t);
       beta_red = tempfnbta(fac_t-2030)

    ##fac    = pow(r['betared'],max(fac_t-2000, 0)) #fac    = (r.betared)^max(fac_t-2000, 0); % trend on activation and progression rate
    fac    = pow(beta_red,max(fac_t-2000, 0))

    try :
      #% Normalise by populations
      lam = M['lambda']*invec/sum(invec) #lam = M.lambda*invec/sum(invec);
      nlin_ds = (lam[0]*M['nlin']['a1']['ds']+ lam[1]*M['nlin']['a2']['ds']+ lam[2]*M['nlin']['a3']['ds']+lam[3]*M['nlin']['a4']['ds']+ lam[4]*M['nlin']['a5']['ds'])#nlin_ds = (lam(1)*M.nlin.a1.ds+ lam(2)*M.nlin.a2.ds+ lam(3)*M.nlin.a3.ds+lam(4)*M.nlin.a4.ds+ lam(5)*M.nlin.a5.ds)
      nlin_mdr = lam[5]*M['nlin']['a1']['mdr']+ lam[6]*M['nlin']['a2']['mdr']+ lam[7]*M['nlin']['a3']['mdr'] +lam[8]*M['nlin']['a4']['mdr']+ lam[9]*M['nlin']['a5']['mdr'] #nlin_mdr = lam(6)*M.nlin.a1.mdr+ lam(7)*M.nlin.a2.mdr+ lam(8)*M.nlin.a3.mdr +lam(9)*M.nlin.a4.mdr+ lam(10)*M.nlin.a5.mdr;
      allmat = M['lin'] + M['Dxlin'] + rHIV*M['linHIV'] + fac*M['linLTBI'] + nlin_ds + nlin_mdr # allmat = M.lin + M.Dxlin + rHIV*M.linHIV + fac*M.linLTBI + nlin_ds + nlin_mdr 
      allmat = allmat.toarray()#out = allmat*invec;
      out1 = allmat@invec#np.matmul(allmat, invec_T)
    except:
       KeyboardInterrupt 
            

    # Implement deaths
    mortvec_sum = np.sum(M['mortvec'], 1) #morts = sum(M.mortvec,2).*invec;
    morts = mortvec_sum*invec#morts = np.asmatrix(mortvec_sum*invec).T#[ mortvec_sum[idx]*inv  for idx, inv in enumerate(invec)]
    out1 = out1 - morts

    out[0:i['nstates']] = out1

    # Implement births
    births = sum(morts) + prm['growth'] #births = sum(morts) + prm.growth; #births = sum(morts)
    out[i['U']['a1']['h0']['v0']] = out[i['U']['a1']['h0']['v0']]+births #out(i.U.a1.h0.v0) = out(i.U.a1.h0.v0)+births#out[i['U']['h0']['v0']] = out[i['U']['h0']['v0']]+births

    # Get the auxiliaries
    #out = [o.sum() for o in out]
    out[i['aux']['inc']]  = agg['inc']@(sel['inc']*allmat)@invec #out(i.aux.inc)     = agg.inc*(sel.inc.*allmat)*invec;
    out[i['aux']['noti']] = agg['noti']@(sel['noti']*allmat)@invec #out(i.aux.noti)    = agg.noti*(sel.noti.*allmat)*invec;
    out[i['aux']['noti2']] = sum((sel['noti2']*allmat)@invec) #out(i.aux.noti2)   = sum((sel.noti2.*allmat)*invec);
    out[i['aux']['mort'][0]] = sum(M['mortvec'][:,1]*invec) #out(i.aux.mort(1)) = sum(M.mortvec(:,2).*invec);
    out[i['aux']['mort'][1]] = sum(M['mortvec'][:,2]*invec) #out(i.aux.mort(2)) = sum(M.mortvec(:,3).*invec);

    return out

