# To add a new cell, type '# %%'
# To add a new markdown cell, type '# %% [markdown]'
# %% 
using HDF5
using Plots
using GR
# using QuadGK
using Printf
using LaTeXStrings
using Statistics
ENV["GKSwstype"] = "100"
# %%

rootdirectory = "/home/oku/SimulationData/isogal"
tstart = 100
time = 100 # number of snapshots
H = 0.25
rlim = 10 # limit of radius (active only plane mode)
binwidthr = 0.25
binwidthrho = 0.2
rhomin = -30
rhomax = -18
binwidthsigma = 0.2
sigmamin = -3
sigmamax = 4
nSF = 10
TSF = 1e4
UnitLength_in_cm = 3.085678e21
UnitMass_in_g = 1.989e43
UnitDensity_in_cgs = UnitMass_in_g / UnitLength_in_cm^3
HydrogenMassFrac = 0.76
ProtonMass = 1.6726e-24

struct Param
  H::Float32
  rlim::Float32
  binwidthr::Float32
  binwidthrho::Float32
  binwidthsigma::Float32
  rhomin::Float32
  rhomax::Float32
  sigmamin::Float32
  sigmamax::Float32
  nSF::Float32
  TSF::Float32
  UnitDensity_in_cgs::Float32
  HydrogenMassFrac::Float32
  ProtonMass::Float32
end

param = Param(H, rlim, binwidthr, binwidthrho, binwidthsigma, rhomin, rhomax, sigmamin, sigmamax, nSF, TSF, UnitDensity_in_cgs, HydrogenMassFrac, ProtonMass)

models = [
    ["ss_model/ver10081/fiducial", "Fiducial", "green", "solid"],
    ["ss_model/ver10081/nonstochastic", "Non-stochastic", "blue", "dash"],
    ["ss_model/ver10081/sniikinetic", "SNII kinetic", "purple", "dot"],
    ["ss_model/ver10081/sniithermal", "SNII thermal", "red", "dashdot"],
    ["ss_model/ver10081/nofb", "No FB", "black", "dashdotdot"],
    # ["ss_model/ver10081/med", "High-reso", "red", "dash"],
    # ["ss_model/ver08261/med_T5", "08261 T5", "blue", "solid"],
    # ["ss_model/ver08261/med_r2", "08261 med_r2", "blue", "solid"],
    # ["ss_model/ver08261/med_p2", "08261 med_p2", "purple", "solid"],
    # ["original/Halov3", "Osaka", "black", "dashdotdot"],
        ]
# %%

function getcenter(sph)
    x = y = z = mtot = 0
    for i in 1:length(sph["mass"])
      x += sph["rho"][i] * sph["mass"][i] * sph["pos"][1, i]
      y += sph["rho"][i] * sph["mass"][i] * sph["pos"][2, i]
      z += sph["rho"][i] * sph["mass"][i] * sph["pos"][3, i]
      mtot += sph["rho"][i] * sph["mass"][i]
    end
    x /= mtot
    y /= mtot
    z /= mtot
    return [x, y, z]
end
function getvelocity(sph)
    x = y = z = mtot = 0
    for i in 1:length(sph["mass"])
      x += sph["mass"][i] * sph["vel"][1, i]
      y += sph["mass"][i] * sph["vel"][2, i]
      z += sph["mass"][i] * sph["vel"][3, i]
      mtot += sph["mass"][i]
    end
    x /= mtot
    y /= mtot
    z /= mtot
    return [x, y, z]
end
function padding(array)
    ret = []
    for i in 1:length(array)
        value = array[i] > 0 ? array[i] : 1e-10
        push!(ret, value)
    end
    return ret
end

function calcsigmaz(sph, galaxy)
    numbinr = param.rlim/param.binwidthr
    vz_r = zeros(Int(numbinr))
    m_r = zeros(Int(numbinr))
    for i in 1:length(sph["mass"])
        x = sph["pos"][1, i] - galaxy["pos"][1, 1]
        y = sph["pos"][2, i] - galaxy["pos"][2, 1]
        z = sph["pos"][3, i] - galaxy["pos"][3, 1]
        vz = sph["vel"][3, i] - galaxy["vel"][3, 1]
        m = sph["mass"][i]
        rho = sph["rho"][i]
        
        # nH = rho * param.UnitDensity_in_cgs * param.HydrogenMassFrac / param.ProtonMass
        # if nH > param.nSF && sph["temp"][i] < param.TSF
        
            r = sqrt(x^2 + y^2)
            if r < param.rlim && abs(z) < param.H
                ibinr = Int(ceil(r/param.binwidthr))
                vz_r[ibinr] += m * (vz)^2
                # vz_r[ibinr] += m * z^2
                m_r[ibinr] += m
            end
        # end
    end
    vz_r = vz_r./m_r
    vz_r = sqrt.(vz_r)

    return vz_r
end

function calcsigma(sph, galaxy)
    numbinr = param.rlim/param.binwidthr
    vr_r = zeros(Int(numbinr))
    m_r = zeros(Int(numbinr))
    for i in 1:length(sph["mass"])
        x = sph["pos"][1, i] - galaxy["pos"][1, 1]
        y = sph["pos"][2, i] - galaxy["pos"][2, 1]
        z = sph["pos"][3, i] - galaxy["pos"][3, 1]
        vx = sph["vel"][1, i] - galaxy["vel"][1, 1]
        vy = sph["vel"][2, i] - galaxy["vel"][2, 1]
        vz = sph["vel"][3, i] - galaxy["vel"][3, 1]
        m = sph["mass"][i]
        
        r = sqrt(x^2 + y^2)
        vr = (-y*vx + x*vy)/r

        if r < param.rlim && abs(z) < param.H
            ibinr = Int(ceil(r/param.binwidthr))
            vr_r[ibinr] += m*vr
            m_r[ibinr] += m
        end
    end
    vr_r = vr_r./m_r

    sph["sigma"] = zeros(length(sph["mass"]))
    for i in 1:length(sph["mass"])
        x = sph["pos"][1, i] - galaxy["pos"][1, 1]
        y = sph["pos"][2, i] - galaxy["pos"][2, 1]
        z = sph["pos"][3, i] - galaxy["pos"][3, 1]
        vx = sph["vel"][1, i] - galaxy["vel"][1, 1]
        vy = sph["vel"][2, i] - galaxy["vel"][2, 1]
        vz = sph["vel"][3, i] - galaxy["vel"][3, 1]
        m = sph["mass"][i]
        
        r = sqrt(x^2 + y^2)

        if r < param.rlim && abs(z) < param.H
            ibinr = Int(ceil(r/param.binwidthr))
            vr = vr_r[ibinr]
            vrx = vr * -y / r
            vry = vr * x / r

            sigma = (vx - vrx)^2 + (vy - vry)^2 + vz^2
            sph["sigma"][i] = sqrt(sigma)
        end
    end
    return vr_r
end

function main(sph, galaxy)
    ret = Dict()
    numbinrho = (param.rhomax - param.rhomin)/param.binwidthrho
    numbinsigma = (param.sigmamax - param.sigmamin)/param.binwidthsigma

    ret["dm/dlogrho"] = zeros(Int(numbinrho))
    ret["dm/dlogsigma"] = zeros(Int(numbinsigma))
    ret["dm/dlogsigmaSF"] = zeros(Int(numbinsigma))
    ret["ddm/dlogcsdlogsigma"] = zeros(Int(numbinsigma), Int(numbinsigma))
    ret["ddm/dlogcsdlogsigmaSF"] = zeros(Int(numbinsigma), Int(numbinsigma))
    Mtot = 0 # total gas mass
    MSF = 0 # total starforming gas mass 
    for i in 1:length(sph["mass"])
        x = sph["pos"][1, i] - galaxy["pos"][1, 1]
        y = sph["pos"][2, i] - galaxy["pos"][2, 1]
        z = sph["pos"][3, i] - galaxy["pos"][3, 1]
        vx = sph["vel"][1, i] - galaxy["vel"][1, 1]
        vy = sph["vel"][2, i] - galaxy["vel"][2, 1]
        vz = sph["vel"][3, i] - galaxy["vel"][3, 1]
        m = sph["mass"][i]
        rho = sph["rho"][i]
        sigma = sph["sigma"][i]
        temp = sph["temp"][i]
        cs = sqrt(5/3 * sph["press"][i] / sph["dens"][i])
        Mach = sigma / cs

        H = abs(z)
        R = sqrt(x^2 + y^2)
        if H < param.H && R < param.rlim
            Mtot += m

            ibinrho = Int(ceil((log10(rho*param.UnitDensity_in_cgs) - param.rhomin)/ param.binwidthrho))
            ret["dm/dlogrho"][ibinrho] += m

            ibinMach = Int(ceil((log10(Mach) - param.sigmamin)/param.binwidthsigma))
            ret["dm/dlogsigma"][ibinMach] += m

            ibinsigma = Int(ceil((log10(sigma) - param.sigmamin)/param.binwidthsigma))
            ibincs = Int(ceil((log10(cs) - param.sigmamin)/param.binwidthsigma))

            ret["ddm/dlogcsdlogsigma"][ibinsigma, ibincs] += m

            nH = rho * param.UnitDensity_in_cgs * param.HydrogenMassFrac / param.ProtonMass
            if nH > param.nSF && temp < param.TSF
                MSF += m

                ret["dm/dlogsigmaSF"][ibinMach] += m

                ret["ddm/dlogcsdlogsigmaSF"][ibinsigma, ibincs] += m

            end
        end
    end
    
    # ret["dm/dlogrho"] = ret["dm/dlogrho"]./Mtot 
    # ret["dm/dlogsigma"] = ret["dm/dlogsigma"]./Mtot
    # ret["dm/dlogsigmaSF"] = ret["dm/dlogsigmaSF"]./MSF
    ret["dm/dlogrho"] = ret["dm/dlogrho"]./param.binwidthrho
    ret["dm/dlogsigma"] = ret["dm/dlogsigma"]./param.binwidthsigma
    ret["dm/dlogsigmaSF"] = ret["dm/dlogsigmaSF"]./param.binwidthsigma
    ret["ddm/dlogcsdlogsigma"] = ret["ddm/dlogcsdlogsigma"]./param.binwidthsigma^2
    ret["ddm/dlogcsdlogsigmaSF"] = ret["ddm/dlogcsdlogsigmaSF"]./param.binwidthsigma^2
    return ret
end

## main part

# %%
using Plots
gr()

Plots.scalefontsizes(1.2)

for j in tstart:time
    Results = Dict()
    
    Results["dm/dlogrho"] = [[] for i in 1:length(models)]
    Results["dm/dlogsigma"] = [[] for i in 1:length(models)]
    Results["dm/dlogsigmaSF"] = [[] for i in 1:length(models)]
    Results["ddm/dlogcsdlogsigma"] = [zeros(1, 1) for i in 1:length(models)]
    Results["ddm/dlogcsdlogsigmaSF"] = [zeros(1, 1) for i in 1:length(models)]
    Results["rotationalvelocity"] = [[] for i in 1:length(models)]
    Results["sigmaz"] = [[] for i in 1:length(models)]
    
    @time for i in 1:length(models)
        ## load data
        path2snapshot = @sprintf("%s/%s/snapshot_%03d/snapshot_%03d.hdf5", rootdirectory, models[i][1], j, j) 

        sph = Dict()
        galaxy = Dict()
        h5open(path2snapshot, "r") do fp
            sph["pos"] = read(fp,"PartType0/Coordinates")
            sph["vel"] = read(fp,"PartType0/Velocities")
            sph["hsml"] = read(fp, "PartType0/SmoothingLength")
            sph["mass"] = read(fp, "PartType0/Masses").*1e10
            sph["rho"] = read(fp, "PartType0/Density")
            sph["temp"] = read(fp, "PartType0/Temperature")
            sph["u"] = read(fp, "PartType0/InternalEnergy")
            sph["press"] = read(fp, "PartType0/Pressure")
            sph["dens"] = read(fp, "PartType0/EnergyWeightedDensity")
        end
        galaxy["pos"] = getcenter(sph)
        galaxy["vel"] = getvelocity(sph)
        
        Results["rotationalvelocity"][i] = calcsigma(sph, galaxy)

        Results["sigmaz"][i] = calcsigmaz(sph, galaxy)
        
        ret = main(sph, galaxy)

        ## save 
        for l in collect(keys(ret))
            Results[l][i]= ret[l]
        end
    end

    RHO = 0:(param.rhomax - param.rhomin)/param.binwidthrho - 1
    RHO = RHO .* param.binwidthrho
    RHO = RHO .+ (param.rhomin + 0.5*param.binwidthrho)
    Plots.plot(ylim=(1e5, 1e10), xlim=(-26, -20), yscale=:log10, xlabel=latexstring("\$\\log\\,(\\rho\\,[\\mathrm{g cm}^{-3}])\$"), ylabel=latexstring("\$dm/d\\log \\rho\\,[M_\\odot]\$"))
    for i in 1:length(models)
        Plots.plot!(RHO, Results["dm/dlogrho"][i], label=models[i][2], color=Symbol(models[i][3]), linestyle=Symbol(models[i][4]))
    end
    figurename = @sprintf("movie/dmdlogrho%03d.png", j)
    Plots.savefig(figurename)
    
    SIGMA = 0:(param.sigmamax - param.sigmamin)/param.binwidthsigma - 1
    SIGMA = SIGMA .* param.binwidthsigma
    SIGMA = SIGMA .+ (param.sigmamin + 0.5*param.binwidthsigma)
    Plots.plot(ylim=(1e5, 1e10), yscale=:log10, xlabel=latexstring("\$\\log\\,(\\sigma / c_s)\$"), ylabel=latexstring("\$dm/d\\log \\sigma\\,[M_\\odot]\$"))
    # Plots.plot(ylim=(1e-10, 1e0), yscale=:log10, xlabel=latexstring("\$\\log\\,(\\sigma\\,[\\mathrm{km s}^{-1}])\$"), ylabel=latexstring("\$dm/d\\log \\sigma\$"))
    for i in 1:length(models)
        Plots.plot!(SIGMA, Results["dm/dlogsigma"][i], label=models[i][2], color=Symbol(models[i][3]), linestyle=Symbol(models[i][4]))
    end
    figurename = @sprintf("movie/dmdlogsigma%03d.png", j)
    Plots.savefig(figurename)
    
    Plots.plot(ylim=(1e5, 1e10), yscale=:log10, xlabel=latexstring("\$\\log\\,(\\sigma / c_s)\$"), ylabel=latexstring("\$dm/d\\log \\sigma\\,[M_\\odot]\$"))
    # Plots.plot(ylim=(1e-10, 1e0), yscale=:log10, xlabel=latexstring("\$\\log\\,(\\sigma\\,[\\mathrm{km s}^{-1}])\$"), ylabel=latexstring("\$dm/d\\log \\sigma\$"))
    for i in 1:length(models)
        Plots.plot!(SIGMA, Results["dm/dlogsigmaSF"][i], label=models[i][2], color=Symbol(models[i][3]), linestyle=Symbol(models[i][4]))
    end
    figurename = @sprintf("movie/dmdlogsigmaSF%03d.png", j)
    Plots.savefig(figurename)
    
    Plots.contour(xlim=(0, 3), ylim=(0, 3), xlabel=latexstring("\$\\log\\,c_s\$"), ylabel=latexstring("\$\\log\\,v\$"), levels=[6, 7, 8, 9])
    # Plots.plot(ylim=(1e-10, 1e0), yscale=:log10, xlabel=latexstring("\$\\log\\,(\\sigma\\,[\\mathrm{km s}^{-1}])\$"), ylabel=latexstring("\$dm/d\\log \\sigma\$"))
    for i in 1:length(models)
        Plots.contour!(SIGMA, SIGMA, log10.(Results["ddm/dlogcsdlogsigmaSF"][i]), label=models[i][2], color=Symbol(models[i][3]), linestyle=Symbol(models[i][4]), levels=[6, 7, 8, 9])
    end
    figurename = @sprintf("movie/ddmdlogcsdlogsigmaSF%03d.png", j)
    Plots.savefig(figurename)
    
    
    R = 0:param.rlim/param.binwidthr - 1
    R = R .* param.binwidthr
    R = R .+ 0.5*param.binwidthr
    Plots.plot(ylim=(0,300), xlim=(0, param.rlim), xlabel=latexstring("\$R\\,[\\mathrm{kpc}]\$"), ylabel=latexstring("\$ V_r\\,[\\mathrm{km s}^{-1}]\$"))
    for i in 1:length(models)
        Plots.plot!(R, Results["rotationalvelocity"][i], label=models[i][2], color=Symbol(models[i][3]), linestyle=Symbol(models[i][4]))
    end
    figurename = @sprintf("movie/rotationalvelocity%03d.png", j)
    Plots.savefig(figurename)

    Plots.plot(xlim=(0, param.rlim), xlabel=latexstring("\$R\\,[\\mathrm{kpc}]\$"), ylabel=latexstring("\$\\sigma_{z}\\,[\\mathrm{km s^{-1}}]\$"))
    for i in 1:length(models)
        Plots.plot!(R, Results["sigmaz"][i], label=models[i][2], color=Symbol(models[i][3]), linestyle=Symbol(models[i][4]))
    end
    figurename = @sprintf("movie/sigmaz%03d.png", j)
    Plots.savefig(figurename)
end

