#Derive the initial series expansions for CAMB
#AL, December 02. 
#Assume neutrinos highly relativistic, baryons cold,
#Lowest order in tight coupling

restart;
read "scal_eqs.txt";

#include terms up to t^Ord
Ord:=3;
DoQuint = "F";

Order:=Ord+5;

makeseries:=proc(v,x,ord)
 local Res,i;
 Res:=0;
 for i from 0 to ord do
   Res:=Res + cat('s',v)[i]*x^i;
 end do;
 v(x)=series(Res,x,ord); 
end;
    

expandfully:=proc(x)
 eval(x);
 %;
 %;
end;

assign(no_numassive);
assign(varsubs);


Phi(t):=-phi(t)-1/2*S(t)^2*kappa*rhopi(t)/k^2;
Psi(t):=phi(t)-1/2*S(t)^2*kappa*rhopi(t)/k^2;


Psi_eq:=k*(diff(Phi(t),t)-H(t)*Psi(t))+1/2*kappa*S(t)^2*((rho_t+p_t)*sigma(t)+rhoq(t));


Rb:=1-Rc;
v(t):=3/4*qg(t);

omb:=Rb*omm;
omc:=Rc*omm;
omr:=Rv*omm^2*H0^2/omega^2;
omg:=omm^2*H0^2/omega^2-omr;

pb(t):=0;c2(t):=0;
pig(t):=0;

etabar(t) := eta(t) - 2*H(t)*rhoq(t)/(rho_t+p_t)/k;
chi(t):=-etabar(t)/2/Kfac;
#etabar2 should be same as etabar(t)
etabar2:=2*Kfac*(Phi(t)+2/3/(1+K/H(t)^2)*(diff(Phi(t),t)/H(t) - Psi(t))/(1+p_t/rho_t));

rhog(t):=3*omg*H0^2/kappa/S(t)^4;
rhor(t):=3*omr*H0^2/kappa/S(t)^4;
rhob(t):=3*omb*H0^2/kappa/S(t)^3;
rhoc(t):=3*omc*H0^2/kappa/S(t)^3;
rhov(t):=3*omv*H0^2/kappa;
H(t):=H_t;

assign(subtots);

if (DoQuint = "T") then
#Quint
assign(subtotsquint);
quint_vars:={clv};
makeseries(V,p,Order);
subs(p=psi(t),convert(%,polynom));
V(t):=rhs(%);
makeseries(psi,t,Order);
assign(%);
spsi[0]:=0;
V(t):=series(simplify(V(t)),t,Order):
V1(t):=diff(V(t),t)/diff(psi(t),t);
V2(t):=diff(V1(t),t)/diff(psi(t),t);

else
assign(no_quint);
end if;


K:=solve(Kfac=Kfac_sub,K);
Kfac:=beta[2];



#Start series solving
dsolve({expandfully(Friedmann),S(0)=0,D(S)(0)=H0^2*omm/omega},S(t),type=series);
assign(%);

# S(t):=omm*H0^2/omega^2*(omega*t+(omega*t)^2/4-K/6*omega*t^3-K/48*omega^2*t^4);


Rb:=1-Rc;
v(t):=3/4*qg(t);

omb:=Rb*omm;
omc:=Rc*omm;
omr:=Rv*omm^2*H0^2/omega^2;
omg:=omm^2*H0^2/omega^2-omr;

pb(t):=0;c2(t):=0;
pig(t):=0;

#Use tight coupling result from Anthony's thesis
drag(t):=-rhob(t)/(3*rhob(t)+4*rhog(t))*(k*clxg(t)+4*H(t)*v(t));

#Some frame invariant variables

hclxc(t):=clxc(t) + 3/2*eta(t)/Kfac;
hclxb(t):=clxb(t) + 3/2*eta(t)/Kfac;
hclxr(t):=clxr(t) + 2*eta(t)/Kfac;
hclxg(t):=clxg(t) + 2*eta(t)/Kfac;
hvb(t):=v(t)+sigma(t);
hqr(t):=qr(t)+ 4/3*sigma(t);
hqg(t):=qg(t)+ 4/3*sigma(t);

hatvars:={hclxc,hclxb,hclxr,hclxg,hvb,hqr,hqg};


solvevars:= {qr,qg,clxg,clxc,clxb,clxr,pir,G_3,G_4,eta};
auxvars:={chi,z,phi,Phi,Psi,sigma} union hatvars;

eqs:={simplify(dclxg-diff(clxg(t),t)),simplify(dclxr-diff(clxr(t),t)),simplify(dclxb-diff(clxb(t),t)),simplify(dclxc-diff(clxc(t),t)),simplify(dqr-diff(qr(t),t)),simplify(dqg-diff(qg(t),t)),simplify(diff(pir(t),t)-dpir),simplify(diff(G_3(t),t)-dG_3),simplify(deta-diff(eta(t),t)),simplify(diff(G_4(t),t)-dG_4)};


knownvars:={H0,omm,omv,k,beta[2],Kf[2],Kf[3],Kf[4],omega,Rv,Rc,sclxb[0],sclxc[0],sclxr[0],seta[0],sqr[0]};

if (DoQuint = "T") then
#quint
eqs:=eqs union {diff(diff(psi(t),t),t)-dpsidot,diff(diff(clv(t),t),t)-dclvdot};
solvevars:=solvevars union {clv};
DeltaQ(t):=(diff(psi(t),t)*diff(clv(t),t)+S(t)^2*clv(t)*V1(t))/(diff(psi(t),t)^2/2+S(t)^2*V(t));
auxvars:=auxvars union {DeltaQ}:
knownvars:= knownvars union {sclv[0],sV[0],sV[1],sV[2],sV[3],sV[4],sV[5]};
end if;

series_subs:=map(makeseries,solvevars,t,Ord+3);
assign(%);




spir[0]:=0;   
sG_3[0]:=0;
sG_3[1]:=0;
sG_4[0]:=0;
sG_4[1]:=0;
sG_4[2]:=0;
G_5(t):=sG_5[4]*t^4+sG_5*t^5;

series_eqs:=map(expandfully,eqs);

map(series,series_eqs,t):
map(simplify,%):
series_eqs:=%;

dosolve:=proc(ineqs)
 local tmp,vars;
 map(series,ineqs,t);
 map(simplify,%);
 map(convert,%,polynom);
 map(tcoeff,%,t):
 tmp:=map(simplify,%);
 vars:=indets(tmp) minus knownvars;
 print(vars);
 solve(tmp,vars);
 assign(%);
end;

for i from 1 to Ord+1 do
 dosolve(series_eqs);
end do;


series('leadterm'(expandfully(etabar(t))),t);
#should be seta0


doresult:=proc(v,sub)
  local tmp;
  tmp:=series(simplify(series(subs(sub,expandfully(v(t))),t,Ord+3)),t,Ord+1);
  convert(%,polynom);
  while indets(lcoeff(%,t)) minus knownvars <> {} do 
   tmp:=series(tmp,t,order(tmp)-1);
   convert(%,polynom);
  end do;
  'v'(t) = tmp;
end;


getresult:=proc(sub)
  map(doresult,[op(solvevars),op(auxvars)],sub);
end;


ad_sub:={seta[0]=2*Kfac,sclxb[0]=0,sclxr[0]=0,sqr[0]=0,sclxc[0]=0,sclv[0]=0};
CDM_sub:={seta[0]=0,sclxb[0]=0,sclxr[0]=0,sqr[0]=0,sclxc[0]=1,sclv[0]=0};
baryon_sub:={seta[0]=0,sclxb[0]=1,sclxr[0]=0,sqr[0]=0,sclxc[0]=0,Rc=1-R_b,sclv[0]=0};
nu_sub:={seta[0]=0,sclxb[0]=0,sclxr[0]=1,sqr[0]=0,sclxc[0]=0,sclv[0]=0};
nu_vel_sub:={seta[0]=0,sclxb[0]=0,sclxr[0]=0,sqr[0]=1,sclxc[0]=0,sclv[0]=0};
nu_pi_sub:={seta[0]=0,sclxb[0]=0,sclxr[0]=0,sqr[0]=0,sclxc[0]=0,spir[0]=1,sclv[0]=0};
null_sub:={seta[0]=0,sclxb[0]=-Rc/Rb*sclxc[0],sclxr[0]=0,sqr[0]=0,sclxc[0]=0,sclv[0]=0};
Q_sub:={seta[0]=0,sclxb[0]=0,sclxr[0]=0,sqr[0]=0,sclxc[0]=0,sclv[0]=1};

ad_subs:=getresult(ad_sub);
CDM_subs:=getresult(CDM_sub);
nu_subs:=getresult(nu_sub);
nu_vel_subs:=getresult(nu_vel_sub);
Q_subs:=getresult(Q_sub);

unassignt:=proc(x)
 unassign('x'(t));
end;

map(unassignt,solvevars);
map(unassignt,auxvars);

#Can now access the solutions ad_subs, CDM_subs, etc..
#E.g. to get clxg(t) (=Delta_\gamma) in the CDM mode use
#subs(CDM_subs,clxg(t));