from scipy.integrate import solve_ivp
from plotter import U
from snapshotprocessor import mean_by_bins
from utils import snapshotselector as select_snapshot

import numpy as np
import os
import matplotlib as mpl

import matplotlib.ticker as tick
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter

from utils.units import convert_exp

ssfr_sun = 2.5e-3

mp = 1.4 * 1.66 * 10 ** (-24) * U.g
z0 = 150 * U.pc


def convert_coldens_s(n0, z0=z0):
    return (np.sqrt(2 * np.pi) * mp * z0 * (n0 * U.cm ** (-3))).express(U.coldens)


convert_coldens = np.vectorize(convert_coldens_s)


def get_sinks(self, overwrite=False, stellar=True, convert_units=False, sk=True):

    self.sinks_from_log(overwrite=overwrite)
    self.coarse_step_from_log(overwrite=overwrite)
    self.sinks = self.get_value("/series/sinks_from_log")

    if stellar:
        self.stellar_from_log(overwrite=overwrite)
        self.stellar = self.get_value("/dataset/stellar_from_log")
        self.stellar["sn_time"] = {}
        for run in self.runs:
            self.stellar["sn_time"][run] = (
                self.stellar["time"][run] + self.stellar["lifetime"][run]
            ) * self.info["unit_time"].express(U.year)
            ind_sort = self.stellar["sn_time"][run].argsort()
            for key in self.stellar:
                self.stellar[key][run] = self.stellar[key][run][ind_sort]
        self.stellar["cum_mass"] = {}
        sn = {}
        sn["time"] = {}
        sn["cum_mass"] = {}

    self.sinks["cum_mass"] = {}

    for run in self.runs:

        if convert_units:
            self.sinks["time"][run] *= self.info["unit_time"].express(U.year)  # year

        self.sinks["cum_mass"][run] = self.sinks["mass_sink"][run].copy()  # Msun

        if stellar:
            self.stellar["cum_mass"][run] = np.cumsum(
                self.stellar["mass"][run]
            ) * self.info["unit_mass"].express(U.Msun)
            sn["time"][run], idx, count = np.unique(
                self.stellar["sn_time"][run], return_index=True, return_counts=True
            )
            idx += count - 1
            sn["cum_mass"][run] = self.stellar["cum_mass"][run][idx]
            ind_sn = np.searchsorted(self.sinks["time"][run], sn["time"][run])
            for i, indi in enumerate(ind_sn[ind_sn < ind_sn[-1]]):
                self.sinks["cum_mass"][run][indi : ind_sn[i + 1]] += sn["cum_mass"][
                    run
                ][i]

    if sk:
        time_gas, total_mass, mass_gas = self.total_mass()  # year, Msun, Msun

        sk_ssfr = {}
        for run in self.runs:
            coldens = mass_gas[run] / 1e6  # in Msun/pc²
            time_gas[run] *= self.info["unit_time"].express(U.Myr)  # Myr
            sk_ssfr[run] = ssfr_sun * (coldens / 10) ** 1.4  # Msun/kpc²/yr-1

        def dermass(t, sm, run, fact_sfr=1):
            ind = min(np.searchsorted(time_gas[run], t), sk_ssfr[run].size - 1)
            return sk_ssfr[run][ind] * fact_sfr * 1e6  # Msun/Myr-1

        tspan = np.linspace(0, 200, 100)  # Myr

        self.time_esm = {}
        self.expected_stellar_mass = {}
        self.expected_stellar_mass_max = {}
        self.expected_stellar_mass_min = {}

        for run in self.runs:
            sol = solve_ivp(
                dermass, (tspan[0], tspan[-1]), [0], t_eval=tspan, args=[run]
            )
            self.time_esm[run], self.expected_stellar_mass[run] = (
                sol["t"],
                sol["y"][0],
            )  # Myr, Msun
            sol_max = solve_ivp(
                dermass, (tspan[0], tspan[-1]), [0], t_eval=tspan, args=[run, 3]
            )
            self.expected_stellar_mass_max[run] = sol_max["y"][0]  # Myr, Msun
            sol_min = solve_ivp(
                dermass, (tspan[0], tspan[-1]), [0], t_eval=tspan, args=[run, 1.0 / 3]
            )
            self.expected_stellar_mass_min[run] = sol_min["y"][0]  # Myr, Msun


def compute_sfr(self, target_start=0.05, target_end=0.3):
    mass = self.sinks["cum_mass"]  # Msun
    time = self.sinks["time"]  # year
    sfr = {}
    sfr_err = {}
    tend = {}
    tstart = {}
    for run in self.runs:
        tend[run] = (
            select_snapshot.time_mcons(self, run, target=target_end) * 1e6
        )  # year
        tstart[run] = (
            select_snapshot.time_mcons(self, run, target=target_start) * 1e6
        )  # year
        if tend[run] < tstart[run]:
            tend[run] = time[run][-1]

    for run in self.runs:
        idx1 = time[run].searchsorted(tstart[run])
        idx2 = min(time[run].searchsorted(tend[run]), len(time[run]) - 1)

        sfr[run] = (mass[run][idx2] - mass[run][idx1]) / (
            time[run][idx2] - time[run][idx1]
        )  # Msun/year
        sfr_other = []
        for i in range(idx1, idx2 - 10):
            sfr_other.append(
                (mass[run][i + 10 : idx2 + 1] - mass[run][i])
                / (time[run][i + 10 : idx2 + 1] - time[run][i])
            )  # Msun/year

        sfr_err[run] = np.std(np.concatenate(sfr_other))  # Msun/year
    self.sfr = np.array(list(sfr.values()))  # Msun/year
    self.sfr_err = np.array(list(sfr_err.values()))  # Msun/year
    return self.sfr, self.sfr_err


def get_nml_array(
    pl,
    xkey,
    cmap=None,
    logcmap=False,
    put_colorbar=True,
    cbar_fmt=tick.FormatStrFormatter("%.2g"),
):
    pl.study.nml(xkey, overwrite=True)
    x = pl.study.get_value(f"/comp/nml_{xkey}")
    xname = os.path.basename(xkey)

    if xname in pl.value_convert:
        x = np.array([pl.value_convert[xname](x_v) for x_v in x])
    if xname in pl.label_convert:
        xlabel = pl.label_convert[xname]
    else:
        xlabel = xname

    colors = None

    if cmap is not None:
        if logcmap:
            sm = plt.cm.ScalarMappable(
                cmap=cmap, norm=mpl.colors.LogNorm(vmin=min(x), vmax=max(x))
            )
        else:
            sm = plt.cm.ScalarMappable(
                cmap=cmap, norm=mpl.colors.Normalize(vmin=min(x), vmax=max(x))
            )
        print(x)
        if put_colorbar:
            cb = plt.colorbar(sm)
            cb.set_label(xlabel)
            cb.ax.yaxis.set_major_formatter(cbar_fmt)

        def colors(xi):
            return sm.cmap(sm.norm(xi))

    return x, xname, xlabel, colors


def histo_speed(pli, redo=False):
    fig, axes = plt.subplots(1, 4, figsize=(16, 5), sharey=True)
    fig2, axes2 = plt.subplots(1, 4, figsize=(16, 5), sharey=True)
    fig3, axes3 = plt.subplots(1, 4, figsize=(16, 5), sharey=True)

    direction = ["x", "y", "z"]

    for pp, ax, ax2, ax3 in zip(
        pli.get_snap_list(
            select={"filter_nml": (pli.nml_key[0], "in", [1.5, 3, 6, 12])}
        ),
        axes,
        axes2,
        axes3,
    ):

        pp.load_cells()
        vel = pp.cells["vel"][:] * pp.info["unit_velocity"].express(U.km_s)
        pos = pp.cells["pos"][:] * pp.info["unit_length"].express(U.pc)

        if redo:
            pp.vmean_z = []
            pp.vmean_x = []

        mass = (
            pp.cells["rho"]
            * pp.info["unit_density"].express(U.Msun / U.pc**3)
            * (pp.cells["dx"] * pp.info["unit_length"].express(U.pc)) ** 3
        )
        for i in range(3):
            ax.hist(
                vel[:, i],
                range=[-100, 100],
                bins=100,
                histtype="step",
                density=True,
                ls="-",
                label=f"$v_{direction[i]}$",
                lw=2,
            )
            ax.set_yscale("log")
            ax.set_ylim(1e-3, 2e-1)
            ax.set_xlabel("$v$ [km/s]")
            ax.set_title(pli.get_label_run(nml_key=pli.nml_key[:2], run=pp.run))
            ax2.set_title(pli.get_label_run(nml_key=pli.nml_key[:2], run=pp.run))
            ax3.set_title(pli.get_label_run(nml_key=pli.nml_key[:2], run=pp.run))

            if redo:
                zbin, vmean_z = mean_by_bins(pos[:, 2], vel[:, i], weights=mass)
                pp.zbin = zbin
                pp.vmean_z.append(vmean_z)

                xbin, vmean_x = mean_by_bins(pos[:, 0], vel[:, i], weights=mass)
                pp.xbin = zbin
                pp.vmean_x.append(vmean_x)
            else:
                zbin = pp.zbin
                xbin = pp.xbin

                vmean_x = pp.vmean_x[i]
                vmean_z = pp.vmean_z[i]

            ax2.plot(zbin, vmean_z, label=f"$v_{direction[i]}$", lw=2)
            ax2.set_xlabel("z [pc]")

            ax3.plot(xbin, vmean_x, label=f"$v_{direction[i]}$", lw=2)
            ax3.set_xlabel("x [pc]")

    fig.suptitle(pli.name)
    fig2.suptitle(pli.name)
    fig3.suptitle(pli.name)
    axes[0].set_ylabel("Mass weighted PDF")
    axes[0].legend()
    axes2[0].set_ylabel("v [km/s]")
    axes2[0].legend()
    axes3[0].set_ylabel("v [km/s]")
    axes3[0].legend()


# def make_clump_hop(self, threshold_density=10):
#     """
#     Apply HOP algorithm

#     Args:
#         threshold_density (float): select only cells over threshold
#     """

#     # Selection of cells
#     mask = (
#         self.cells["rho"] * self.info["unit_density"].express(U.H_cc)
#         > threshold_density
#     )
#     ncells = np.sum(mask)

#     # fill the matrice with ID, x,y,z and masses of particles
#     cells_group = np.zeros((ncells, 14))
#     cells_group[:, 0] = np.arange(ncells)  # index
#     position = self.cells["pos"][mask] * self.info["unit_length"].express(U.pc)
#     cells_group[:, 1:4] = position  # position
#     density = self.cells["rho"][mask] * self.info["unit_density"].express(U.H_cc)
#     size = self.cells["dx"][mask] * self.info["unit_length"].express(U.pc)
#     cells_group[:, 4] = (
#         density * size ** 3 * (U.H_cc * U.pc ** 3).express(U.Msun)
#     )  # mass
#     cells_group[:, 6] = density
#     velocity = self.cells["vel"][mask] * self.info["unit_velocity"].express(U.km_s)
#     pressure = self.cells["P"][mask] * self.info["unit_pressure"].express(U.bar)
#     temperature = (pressure / density) * (
#         (U.bar / U.H_cc) * (1.4 * U.mH) / U.kB
#     ).express(U.K)
#     cells_group[:, 7] = temperature
#     cells_group[:, 8] = pressure
#     cells_group[:, 9:12] = velocity
#     cells_group[:, 12] = self.cells["phi"][mask]
#     cells_group[:, 13] = size ** 3

#     # save file.txt
#     head = str(ncells)
#     np.savetxt(
#         self.filename[:-3] + "_hop.txt",
#         cells_group[:, :5],
#         fmt="%10d %.10e %.10e %.10e %.10e",
#         header=head,
#         delimiter=" ",
#         comments=" ",
#     )

#     # save file.den
#     f = open(self.filename[:-3] + "_hop.den", "wb")
#     f.write(pack("I", ncells))
#     self.cells["rho"][mask].astype("f").tofile(f)
#     f.close()

#     # exec HOP algo
#     h = HOP(self.filename[:-3] + "_hop.txt", os.path.dirname(self.filename))
#     h.process_hop(force=True)

#     # get the igroup array
#     group_ids = h.get_group_ids()

#     # sort it and apply the sorting to the coordinates
#     # this means that the particules of group 1 are written first then of group 2 etc...
#     ind_sort = np.argsort(group_ids)
#     cells_group = cells_group[ind_sort]
#     cells_group[:, 5] = group_ids[ind_sort]

#     # Make it a pandas' DataFrame
#     cells_group = pd.DataFrame(
#         cells_group,
#         columns=[
#             "id",
#             "x",
#             "y",
#             "z",
#             "mass",
#             "group",
#             "density",
#             "temperature",
#             "pressure",
#             "vx",
#             "vy",
#             "vz",
#             "phi",
#             "volume",
#         ],
#     )

#     self.clumps = cells_group

#     return cells_group


def plot_mass_sfr(
    pl,
    xkey="turb_params/comp_frac",
    start=0.01,
    end=0.4,
    cmap=None,
    logcmap=False,
    redo=False,
    fig_height=4,
    ax_sfr=None,
    logsfr=True,
    marker="o",
    color_sfr="k",
    scale=0.9,
    do_sfr=True,
    plot_dir=".",
):
    plt.figure(figsize=np.array([5, fig_height]) * scale)
    get_sinks(pl.study, overwrite=redo)

    pl.value_convert["Bx"] = lambda x: x * 7.6189439
    pl.label_convert["Bx"] = r"B$_0$ [$\mu$G]"
    pl.value_convert["bx_bound"] = lambda x: x * 7.6189439
    pl.label_convert["bx_bound"] = r"B$_0$ [$\mu$G]"
    pl.label_convert["comp_frac"] = r"$\chi$"

    x_fmt = FuncFormatter(lambda x, p: f"{convert_exp(x,2)}")
    x, xname, xlabel, colors = get_nml_array(pl, xkey, cmap, logcmap, cbar_fmt=x_fmt)

    try:
        n0 = pl.study.get_nml("galbox_params/dens0", pl.runs[0])
    except KeyError():
        n0 = pl.study.get_nml("cloud_params/dens0", pl.runs[0])

    mass = pl.study.sinks["cum_mass"]
    time = pl.study.sinks["time"]
    tend = {}
    tstart = {}

    for i, run in enumerate(pl.runs):
        tend[run] = (select_snapshot.time_mcons(pl.study, run, target=end) * 1e6,)
        tstart[run] = select_snapshot.time_mcons(pl.study, run, target=start) * 1e6
        if tend[run] < tstart[run]:
            tend[run] = time[run][-1]

        idx1 = time[run].searchsorted(tend[run])
        idx2 = time[run].searchsorted(tstart[run])
        mask = time[run] < tend[run]

        if cmap is None:
            (p,) = plt.plot(
                time[run][mask] * 1e-6,
                mass[run][mask] * 1e-6,
                label=f"{xlabel} = {x[i]:.2f}",
            )
        else:
            color = colors(x[i])
            (p,) = plt.plot(time[run][mask] * 1e-6, mass[run][mask] * 1e-6, color=color)
        plt.scatter(
            [time[run][idx1] * 1e-6],
            [mass[run][idx1] * 1e-6],
            marker="<",
            color=p.get_color(),
        )
        plt.scatter(
            [time[run][idx2] * 1e-6],
            [mass[run][idx2] * 1e-6],
            marker=">",
            color=p.get_color(),
        )

    plt.title(f"$\Sigma_0$ = {convert_coldens(n0):.1f}" + " M$_\odot$.pc$^{-1}$")
    plt.xlabel("Time [Myr]")
    plt.ylabel("Mass accreted [$10^6$ M$_\odot$]")
    plt.savefig(f"{plot_dir}/{xname}_n{n0}_mass.pdf")

    ##################################################
    if do_sfr:
        label = f"$\Sigma_0$ = {convert_coldens(n0):.1f}" + " M$_\odot$.pc$^{-2}$"

        if ax_sfr is None:
            plt.figure(figsize=(5, fig_height))
        else:
            plt.sca(ax_sfr)
        if redo or not hasattr(pl.study, "sfr"):
            sfr, sfr_err = compute_sfr(pl.study, target_start=start, target_end=end)
        else:
            sfr, sfr_err = pl.study.sfr, pl.study.sfr_err
        if logsfr:
            plt.errorbar(
                x=x,
                y=np.log10(sfr),
                yerr=sfr_err / sfr,
                color=color_sfr,
                marker=marker,
                ls=":",
                lw=0.5,
                label=label,
                elinewidth=1.5,
            )
            plt.hlines(
                np.log10(
                    ssfr_sun * (convert_coldens(n0) * (1 - end / 2.0) / 10) ** 1.4
                ),
                xmin=min(x),
                xmax=max(x),
                color=color_sfr,
                ls="--",
            )
            plt.ylabel("log$(\Sigma_{\mathrm{SFR}})$ " + rf"[${U.ssfrK.latex}$]")
        else:
            plt.errorbar(
                x=x,
                y=sfr,
                yerr=sfr_err,
                color=color_sfr,
                marker=marker,
                ls=":",
                lw=0.5,
                label=label,
            )
            plt.hlines(
                ssfr_sun * (convert_coldens(n0) * (1 - end / 2.0) / 10) ** 1.4,
                xmin=min(x),
                xmax=max(x),
                color=color_sfr,
                ls="--",
            )
            plt.ylabel("$\Sigma_{\mathrm{SFR}}$ " + rf"[${U.ssfrK.latex}$]")

        plt.xlabel(f"{xlabel}")
        plt.gca().xaxis.set_major_formatter(x_fmt)
        plt.legend()

        plt.savefig(f"{plot_dir}/{xname}_n{n0}_logsfr.pdf")