
from scipy import interpolate
import numpy as np

def goveqs_basis_disruption(input, t, M, notif_fac, i, prm, sel, agg):
     
    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) 
        #if ti < 0:
        #    rHIV = tempfn(0)
        #elif ti> len(prm['rHIV'])-1:
        #    rHIV = tempfn(len(prm['rHIV'])-1)
        #else:
        #    rHIV = tempfn(ti) 

    # Normalise by populations
    lam = M['lambda']*invec/sum(invec)
    #fac = interpolate.interp1(2019+(0:length(notif_fac))/12, [[1 1];notif_fac], t)
    
    yrs = [y/12+2019 for y in range(len(notif_fac))]
    fnfac = interpolate.interp1d(2019+np.array(range(len(notif_fac)+1))/12, np.concatenate([[1], np.concatenate(notif_fac)]), kind='next', fill_value="extrapolate")
    
    fac=fnfac(t) 

    #if t<2019:
    #    fac = fnfac(2019)
    #else:
    #    fac=fnfac(t) 

    # Pull them together
    allmat = M['lin'] + M['Dxlin']*fac  + rHIV*M['linHIV'] + lam[0]*M['nlin']['ds'] + lam[1]*M['nlin']['mdr']


    allmat = allmat.toarray()

    #invec_T = np.asmatrix(invec).T
    #out1 = np.matmul(allmat, invec_T)
    out1 = allmat@invec

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

    #for idx in range(i['nstates']):
    #    out[idx] = out1[idx]
    #Changing the for loop since it is not efficient
    out[0:i['nstates']] = out1

    # Implement births
    births = sum(morts)
    out[i['U']['h0']['v0']] = out[i['U']['h0']['v0']]+births

    # Get the auxiliaries
    #sel_inci = (sel['inc']*allmat)
    #sel_noti = (sel['noti']*allmat)
    ##sel_mort = (sel['mort']*allmat)
    #sel_noti2 = (sel['noti2']*allmat)

    #sel_invec_inci = np.matmul(sel_inci, invec_T)
    #sel_invec_noti = np.matmul(sel_noti, invec_T)
    #sel_invec_noti2 = np.matmul(sel_noti2, invec_T) 

    #temp_arr = agg['inc']*sel_invec_inci
    #out[i['aux']['inc']]  = temp_arr.flatten()#np.matmul(agg['inc'], sel_invec_inci)
    #out[i['aux']['noti'][0]] = agg['noti'][0]*sel_invec_noti#np.matmul(agg['noti'][0], sel_invec_noti)
    #out[i['aux']['noti'][1]] = agg['noti'][1]*sel_invec_noti#np.matmul(agg['noti'][1], sel_invec_noti)
    
    #out[i['aux']['noti2']] = sum(sel['noti2']*sel_invec_noti)#sum(np.matmul(agg['noti2'], sel_invec_noti))

    #out[i['aux']['mort'][0]] = sum(M['mortvec'][:,1]*invec_T)
    #out[i['aux']['mort'][1]] = sum(M['mortvec'][:,2]*invec_T)
    
    #out = [o.sum() for o in out]

    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'][0]] = sum(M['mortvec'][:,1]*invec)
    out[i['aux']['mort'][1]] = sum(M['mortvec'][:,2]*invec)
            
    return out
