from for_pencil import varmax
import pencil as pc
from misc import hfluct2d, deriv
import pylab as P
import numpy as N

P.rc('text', usetex=False)
P.rc("lines", linewidth=2)

index = pc.read_index()
dim = pc.read_dim()
param = pc.read_param(quiet=True)
par2 = pc.read_param(param2=True, quiet=True)
vmax=varmax()

nvar=5
var1=vmax-nvar+1

frad = N.empty((dim.nz, nvar), dtype='Float64')
fconv = N.empty_like(frad)
fkin = N.empty_like(frad)
Rstar = param.cp*(param.gamma-1.)/param.gamma
for k, ivar in enumerate(range(var1, vmax+1)):
    print 'read VAR%i'%ivar
    var = pc.read_var(ivar=ivar, run2D=True, quiet=True,
          param=param, dim=dim, index=index, magic=['TT','rho'], )
    ux = var.ux[var.n1:var.n2, var.l1:var.l2]
    uz = var.uz[var.n1:var.n2, var.l1:var.l2]
    rho = N.exp(var.lnrho[var.n1:var.n2, var.l1:var.l2])
    temp = var.TT[var.n1:var.n2, var.l1:var.l2]
    # radiative flux (only temp is recorded here)
    frad[:, k] = temp.mean(axis=1)
    # kinetic flux
    a = hfluct2d(uz)*(rho*(ux**2+uz**2))
    fkin[:, k] = 0.5*a.mean(axis=1)
    # convective flux
    #a = uz*hfluct2d(rho*temp*Rstar)
    a = rho*uz*hfluct2d(temp)
    fconv[:, k] = a.mean(axis=1)
z = var.z[var.n1:var.n2]

# temporal means
frad = frad.mean(axis=1)
frad = -par2.hcond0*deriv(z, frad)
fconv = fconv.mean(axis=1)
fkin = fkin.mean(axis=1)
ftot = frad+fconv+fkin

P.ioff()
P.plot(z, frad, 'b-', label='F_rad')
P.plot(z, fconv, 'g--', label='F_conv')
P.plot(z, fkin, 'k:', label='F_kin')
P.plot(z, ftot, 'r-', label='F_tot')
P.ylabel('Fluxes', fontsize=16)
P.legend(loc=0)
P.xlabel('Height', fontsize=16)
#P.axvline(param.z1, c='k', lw=2)
#P.axvline(param.z2, c='k', lw=2)
P.xlim(xmin=param.xyz0[2], xmax=param.xyz1[2])
P.subplots_adjust(right=0.95, top=0.95)
P.show()
P.ion()