import numpy as np 
import matplotlib.pyplot as plt
import h5py
import astropy.units as u, astropy.constants as c

from load_data_arepo import load_fields_arepo
from galsec import Galsec

def sfr_history(data, tmin=None, tmax=None, average_time=1, **kwargs):
    """History of SFR. Time in Myr.
    SFR in Msun/year
    Parameters
    ----------
    data : _type_
        _description_
    tmin : _type_
        _description_
    tmax : _type_
        _description_
    average_time : _type_
        _description_
    """
    plt.figure(constrained_layout=True, figsize=(5,4))
    if tmin is None:
        tmin = np.floor(np.min(data["stars"]["birth_time"]))
    if tmax is None:
        tmax = np.ceil(np.max(data["stars"]["birth_time"]))
    bins = int(np.ceil((tmax - tmin) / average_time))
    tmax = tmin +  bins * average_time
    plt.hist(data["stars"]["birth_time"],
             weights=data["stars"]["mass"] / (average_time*1e6), 
             bins=bins, range=[tmin, tmax], 
             histtype="step", **kwargs)
    plt.ylabel("SFR [M$_\odot$/yr]")
    plt.xlabel("Time [Myr]")
    plt.savefig("sfr_history.png")

def ring_stuff(galsec, delta_r=1):
    galsec.ring_analysis(delta_r * u.kpc, rmax=12 * u.kpc)
    r_stars = galsec.rings['stars']["r"]
    r_gas = galsec.rings['gas']["r"]
    surface = np.pi * (delta_r * u.kpc) * r_gas
    
    plt.figure(constrained_layout=True, figsize=(5,4))
    sfr = galsec.rings['stars']["sfr"]
    ssfr = sfr / surface
    plt.plot(r_stars, ssfr)
    plt.ylabel("SSFR [M$_\odot$/yr/kpc$^2$]")
    plt.xlabel("R [kpc]")
    plt.savefig("sfr_profile.png")

    
    plt.figure(constrained_layout=True, figsize=(5,4))
    velphi = - galsec.rings['gas']["velphi"]
    plt.plot(r_gas, velphi)
    plt.ylabel(r"$v_\varphi$ [km/s]")
    plt.xlabel("R [kpc]")
    plt.savefig("rotation_curve.png")
    
    
    plt.figure(constrained_layout=True, figsize=(5,4))
    sigma_velphi = galsec.rings['gas']["sigma_velphi"]
    sigma_velr = galsec.rings['gas']["sigma_velr"]
    sigma_velz = galsec.rings['gas']["sigma_velz"]
    plt.plot(r_gas, sigma_velphi, label=r"$\sigma_{v,\varphi}$")
    plt.plot(r_gas, sigma_velr, label=r"$\sigma_{v,r}$")
    plt.plot(r_gas, sigma_velz, label=r"$\sigma_{v,z}$")
    plt.ylabel(r"$\sigma$ [km/s]")
    plt.xlabel("R [kpc]")
    plt.legend()
    plt.savefig("dispersion_profile.png")

    

def run_pipeline(path):
    
    # Galsec analysis
    data = load_fields_arepo(path)
    galsec = Galsec(data)
    
    sfr_history(data)
    ring_stuff(galsec)
    
    del data
    del galsec
    
    

if __name__ == "__main__":
    run_pipeline("/home/nbrucy/Travail/Postdoc/Ecogal/MW/junia2_0.25kpc/OUTPUT_SN/snap_400.hdf5")