import numpy as np

def goveqs_basissmall(f_in, t, r, M, i, prm): #function [out, lam] = goveqs_basissmall(t, in, M, i, prm)             % <--- Arguments changed to accommodate prm

    invec = f_in[0:i['nstates']] #invec = in(1:i.nstates);
    #r['betared'] = np.array([r['betared']])
    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)); 
    
    #%try
    #% Normalise by populations
    lam = M['lambda']*invec/sum(invec)

    #What if t is less than 2015?
    #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_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'])#*pow(r['betared'][0],t-2015)#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;
    if i['nx']>60 : #i.nx>60  %length(gps.strains)>1
        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; 
    else:
        nlin_mdr = 0   
    #end
    #%allmat = M.lin + M.Dxlin + rHIV*M.linHIV + nlin_ds + nlin_mdr;
    #allmat = M['lin'] + M['Dxlin'] +  nlin_ds + nlin_mdr #allmat = M.lin + M.Dxlin +  nlin_ds + nlin_mdr;
    #%allmat = M.lin + M.Dxlin + M.linLTBI +rHIV*M.linHIV + nlin_ds + nlin_mdr;
    #New, 2024-04-02
    allmat = M['lin'] + M['Dxlin'] + fac*M['linLTBI'] +  nlin_ds + nlin_mdr

    out = allmat@invec #out = allmat*invec;
    #% catch
    #%    keyboard; 
    #% end

    #% Implement deaths
    morts = np.sum(M['mortvec'],1)*invec#    morts = sum(M.mortvec,2).*invec;
    out = out - morts #out = out - morts;

    #% Implement births
    births = sum(morts) + prm['growth'] #births = sum(morts) + prm.growth;
    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;

    #% % Get the auxiliaries
    #% out(i.aux.inc)     = agg.inc*(sel.inc.*allmat)*invec;
    #% out(i.aux.noti)    = agg.noti*(sel.noti.*allmat)*invec;
    #% out(i.aux.noti2)   = sum((sel.noti2.*allmat)*invec);
    #% out(i.aux.mort(1)) = sum(M.mortvec(:,2).*invec);
    #% out(i.aux.mort(2)) = sum(M.mortvec(:,3).*invec);
    #n_out = np.concatenate([out, lam])
    #Python does not allow returning a vector that does not have the same length as the initial condition. The solution here would be either 
    #to pass a vector as parameter and write lam in it or increase the size of the initial condition...
    return(out) 