# Contains commonly used functions related to plotting results
# of the Strömgren Sphere Examples and other gas related functions.

import copy
import os

import numpy as np
import unyt

# species masses in atomic mass units
mamu = {"e": 0.0, "HI": 1.0, "HII": 1.0, "HeI": 4.0, "HeII": 4.0, "HeIII": 4.0}


def get_number_densities(Temp, XH, XHe):
    """
    Compute the number densities of all species at a given
    temperature following Katz, Hernquist, and Weinberg 1996

    Temp: temperature [unyt quantity]
    XH: total mass fraction of all hydrogen species (HI + HII)
    XHe: total mass fraction of all helium species (HeI + HeII + HeIII)
    """

    # n_H = X_H * rho_gas / m_H
    # then m_H = rho_gas * X_H / n_H
    # n_He = X_He * rho_gas / m_He =
    #      = (1 - X_H) * rho_gas / (4 * m_H)
    #      = (1 - X_H) * rho_gas / (4 * rho_gas * x_H / n_H)
    #      = (1 - X_H) / (4 * X_H) * n_H
    #      =  X_He / 4(1 - X_He) * n_H = y * n_H

    if XH == 0:
        nH = 0.0
        nHe = 1.0
    else:
        nH = XH
        nHe = XHe / 4

    # NOTE: This is not the ionization threshold!
    if Temp < 5000 * unyt.K:
        nH0 = nH
        nHp = 0.0
        nHe0 = nHe
        nHep = 0.0
        nHepp = 0.0

    else:

        Temp.convert_to_cgs()
        T = Temp.v
        # Recombination rate for H+ in units of cm^3 s^-1
        A_Hp = (
            8.40e-11
            / np.sqrt(T)
            * (T * 1e-3) ** (-0.2)
            * 1.0
            / (1.0 + (T * 1e-6) ** 0.7)
        )

        # Dielectronic recombination rate for He+ in units of cm^3 s^-1
        A_d = (
            1.9e-3
            / T ** 1.5
            * np.exp(-470000.0 / T)
            * (1.0 + 0.3 * np.exp(-94000.0 / T))
        )
        # Recombination rate for He+ in units of cm^3 s^-1
        A_Hep = 1.5e-10 / T ** 0.6353
        # Recombination rate for He++ in units of cm^3 s^-1
        A_Hepp = (
            3.36e-10
            / np.sqrt(T)
            * (T * 1e-3) ** (-0.2)
            * 1.0
            / (1.0 + (T * 1e-6) ** 0.7)
        )
        # collisional ionization rate for H0 in units of cm^3 s^-1
        #  G_H0 = 1.17e-10 * np.sqrt(T) * np.exp(-157809.1 / T) * 1. / (1. + np.sqrt(T*1e-5))
        G_H0 = (
            5.85e-11
            * np.sqrt(T)
            * np.exp(-157809.1 / T)
            * 1.0
            / (1.0 + np.sqrt(T * 1e-5))
        )
        # collisional ionization rate for He0 in units of cm^3 s^-1
        G_He0 = (
            2.38e-11
            * np.sqrt(T)
            * np.exp(-285335.4 / T)
            * 1.0
            / (1.0 + np.sqrt(T * 1e-5))
        )
        # collisional ionization rate for He+ in units of cm^3 s^-1
        G_Hep = (
            5.68e-12
            * np.sqrt(T)
            * np.exp(-631515.0 / T)
            * 1.0
            / (1.0 + np.sqrt(T * 1e-5))
        )

        # Katz et al. 1996 eq. 33 - 38
        # Note: We assume all photoionization rates to be zero.
        # Also, here we don't care about the actual number density, i.e.
        # about the volume, since it'll cancel out later when we compute
        # the mass fractions.

        nH0 = nH * A_Hp / (A_Hp + G_H0)
        nHp = nH - nH0
        nHep = nHe / (1.0 + (A_Hep + A_d) / G_He0 + G_Hep / A_Hepp)
        nHe0 = nHep * (A_Hep + A_d) / G_He0
        nHepp = nHep * G_Hep / A_Hepp

    # electron density
    ne = nHp + nHep + 2.0 * nHepp

    return nH0, nHp, nHe0, nHep, nHepp, ne


def get_number_densities_array(Temp, XH, XHe):
    """
    Compute the number densities of all species at a given
    temperature following Katz, Hernquist, and Weinberg 1996

    Temp: temperature [unyt array]
    XH: total mass fraction of all hydrogen species (HI + HII)
    XHe: total mass fraction of all helium species (HeI + HeII + HeIII)
    """

    # n_H = X_H * rho_gas / m_H
    # then m_H = rho_gas * X_H / n_H
    # n_He = X_He * rho_gas / m_He =
    #      = (1 - X_H) * rho_gas / (4 * m_H)
    #      = (1 - X_H) * rho_gas / (4 * rho_gas * x_H / n_H)
    #      = (1 - X_H) / (4 * X_H) * n_H
    #      =  X_He / 4(1 - X_He) * n_H = y * n_H

    nH = np.zeros(XH.shape, dtype=np.float64)
    nHe = np.zeros(XH.shape, dtype=np.float64)

    mask = XH == 0
    nH[mask] = 0.0
    nHe[mask] = 1.0

    inv_mask = np.logical_not(mask)
    nH[inv_mask] = XH[inv_mask]
    nHe[inv_mask] = 0.25 * XHe[inv_mask]

    nH0 = np.zeros(XH.shape, dtype=np.float64)
    nHp = np.zeros(XH.shape, dtype=np.float64)
    nHe0 = np.zeros(XH.shape, dtype=np.float64)
    nHep = np.zeros(XH.shape, dtype=np.float64)
    nHepp = np.zeros(XH.shape, dtype=np.float64)

    # NOTE: This is not the ionization threshold!
    neutral = Temp < 5000 * unyt.K

    nH0[neutral] = nH[neutral]
    nHp[neutral] = 0.0
    nHe0[neutral] = nHe[neutral]
    nHep[neutral] = 0.0
    nHepp[neutral] = 0.0

    Temp.convert_to_cgs()
    T = Temp.v
    ionized = np.logical_not(neutral)

    # Recombination rate for H+ in units of cm^3 s^-1
    A_Hp = (
        8.40e-11 / np.sqrt(T) * (T * 1e-3) ** (-0.2) * 1.0 / (1.0 + (T * 1e-6) ** 0.7)
    )
    # Dielectronic recombination rate for He+ in units of cm^3 s^-1
    A_d = 1.9e-3 / T ** 1.5 * np.exp(-470000.0 / T) * (1.0 + 0.3 * np.exp(-94000.0 / T))
    # Recombination rate for He+ in units of cm^3 s^-1
    A_Hep = 1.5e-10 / T ** 0.6353
    # Recombination rate for He++ in units of cm^3 s^-1
    A_Hepp = (
        3.36e-10 / np.sqrt(T) * (T * 1e-3) ** (-0.2) * 1.0 / (1.0 + (T * 1e-6) ** 0.7)
    )
    # collisional ionization rate for H0 in units of cm^3 s^-1
    G_H0 = (
        5.85e-11 * np.sqrt(T) * np.exp(-157809.1 / T) * 1.0 / (1.0 + np.sqrt(T * 1e-5))
    )
    # collisional ionization rate for He0 in units of cm^3 s^-1
    G_He0 = (
        2.38e-11 * np.sqrt(T) * np.exp(-285335.4 / T) * 1.0 / (1.0 + np.sqrt(T * 1e-5))
    )
    # collisional ionization rate for He+ in units of cm^3 s^-1
    G_Hep = (
        5.68e-12 * np.sqrt(T) * np.exp(-631515.0 / T) * 1.0 / (1.0 + np.sqrt(T * 1e-5))
    )

    # Katz et al. 1996 eq. 33 - 38
    # Note: We assume all photoionization rates to be zero.
    # Also, here we don't care about the actual number density, i.e.
    # about the volume, since it'll cancel out later when we compute
    # the mass fractions.

    nH0[ionized] = nH[ionized] * A_Hp[ionized] / (A_Hp[ionized] + G_H0[ionized])
    nHp[ionized] = nH[ionized] - nH0[ionized]
    nHep[ionized] = nHe[ionized] / (
        1.0
        + (A_Hep[ionized] + A_d[ionized]) / G_He0[ionized]
        + G_Hep[ionized] / A_Hepp[ionized]
    )
    nHe0[ionized] = nHep[ionized] * (A_Hep[ionized] + A_d[ionized]) / G_He0[ionized]
    nHepp[ionized] = nHep[ionized] * G_Hep[ionized] / A_Hepp[ionized]

    # electron density
    ne = nHp + nHep + 2.0 * nHepp

    return nH0, nHp, nHe0, nHep, nHepp, ne


def get_mass_fractions(T, XH, XHe):
    """
    Compute the mass fractions of all species at a
    given temperature

    T: temperature
    XH: total mass fraction of all hydrogen species (HI + HII)
    XHe: total mass fraction of all helium species (HeI + HeII + HeIII)
    """

    # first get number densities
    if isinstance(XH, np.ndarray):
        nH0, nHp, nHe0, nHep, nHepp, ne = get_number_densities_array(T, XH, XHe)
    else:
        nH0, nHp, nHe0, nHep, nHepp, ne = get_number_densities(T, XH, XHe)

    # now get mass denities in units of atomic mass units
    mH0 = nH0
    mHp = nHp
    mHe0 = 4.0 * nHe0
    mHep = 4.0 * nHep
    mHepp = 4.0 * nHepp
    # neglect electron mass contributions

    mtot = mH0 + mHp + mHe0 + mHep + mHepp  # + me

    XH0 = mH0 / mtot
    XHp = mHp / mtot
    XHe0 = mHe0 / mtot
    XHep = mHep / mtot
    XHepp = mHepp / mtot
    #  Xe = me / mtot

    return XH0, XHp, XHe0, XHep, XHepp


def internal_energy(T, mu, gamma):
    """
    Compute the internal energy of the gas for a given
    temperature and mean molecular weight
    """
    # Using u = 1 / (gamma - 1) * p / rho
    #   and p = N/V * kT = rho / (mu * m_u) * kT

    u = unyt.boltzmann_constant * T / (gamma - 1) / (mu * unyt.atomic_mass_unit)
    return u


def get_soundspeed_from_internal_energy(data):
    """
    Compute the local sound speed for all particles.
    data: swiftsimio.load() object.
    """

    u = data.gas.internal_energies
    gamma = data.metadata.gas_gamma

    return np.sqrt(u * gamma / (gamma - 1))


def get_soundspeed_from_density_pressure(data):
    """
    Compute the local sound speed for all particles.
    data: swiftsimio.load() object.
    """

    gamma = data.metadata.gas_gamma
    P = data.gas.pressures
    rho = data.gas.densities

    return np.sqrt(gamma * P / rho)


def mean_molecular_weight(XH0, XHp, XHe0, XHep, XHepp):
    """
    Determines the mean molecular weight for given 
    mass fractions of
        hydrogen:   XH0
        H+:         XHp
        He:         XHe0
        He+:        XHep
        He++:       XHepp

    returns:
        mu: mean molecular weight [in atomic mass units]
        NOTE: to get the actual mean mass, you still need
        to multiply it by m_u, as is tradition in the formulae
    """

    # 1/mu = sum_j X_j / A_j * (1 + E_j)
    # A_H    = 1, E_H    = 0
    # A_Hp   = 1, E_Hp   = 1
    # A_He   = 4, E_He   = 0
    # A_Hep  = 4, E_Hep  = 1
    # A_Hepp = 4, E_Hepp = 2
    one_over_mu = XH0 + 2 * XHp + 0.25 * XHe0 + 0.5 * XHep + 0.75 * XHepp

    return 1.0 / one_over_mu


def gas_temperature(u, mu, gamma):
    """
    Compute the gas temperature given the specific internal 
    energy u and the mean molecular weight mu
    """

    # Using u = 1 / (gamma - 1) * p / rho
    #   and p = N/V * kT = rho / (mu * m_u) * kT

    T = u * (gamma - 1) * mu * unyt.atomic_mass_unit / unyt.boltzmann_constant

    return T.to("K")


def get_snapshot_list(snapshot_basename="output", plot_all=True, snapnr=0):
    """
    Find the snapshot(s) that are to be plotted 
    and return their names as list
    """

    snaplist = []

    if plot_all:
        dirlist = os.listdir()
        for f in dirlist:
            if f.startswith(snapshot_basename) and f.endswith("hdf5"):
                snaplist.append(f)

        snaplist = sorted(snaplist)

    else:
        fname = snapshot_basename + "_" + str(snapnr).zfill(4) + ".hdf5"
        if not os.path.exists(fname):
            print("Didn't find file", fname)
            quit(1)
        snaplist.append(fname)

    return snaplist


def get_imf(scheme, data):
    """
    Get the ion mass fraction (imf) according to the scheme.
    return a class with ion mass fraction for species X, 
    including HI, HII, HeI, HeII, HeIII:
    The ion mass function can be accessed through: imf.X
    The unit is in m_X/m_tot, where m_X is the mass in species X
    and m_tot is the total gas mass.
    """
    if scheme.startswith("GEAR M1closure"):
        imf = data.gas.ion_mass_fractions
    elif scheme.startswith("SPH M1closure"):
        # atomic mass
        mass_fraction_hydrogen = data.gas.rt_element_mass_fractions.hydrogen
        imf = copy.deepcopy(data.gas.rt_species_abundances)
        named_columns = data.gas.rt_species_abundances.named_columns
        for column in named_columns:
            # abundance is in n_X/n_H unit. We convert it to mass fraction by multipling mass fraction of H
            mass_fraction = (
                getattr(data.gas.rt_species_abundances, column)
                * mass_fraction_hydrogen
                * mamu[column]
            )
            setattr(imf, column, mass_fraction)
    else:
        raise ValueError("Unknown scheme", scheme)
    return imf


def get_abundances(scheme, data):
    """
    Get the species abundance according to the scheme
    return a class with normalized number densities for abunance X, 
    including HI, HII, HeI, HeII, HeIII:
    The abundances can be accessed through: sA.X
    The unit is in n_X/n_H, where n_X is the number density of species X
    and n_H is the number density of hydrogen.
    """
    if scheme.startswith("GEAR M1closure"):
        # atomic mass
        sA = copy.deepcopy(data.gas.ion_mass_fractions)
        mass_fraction_hydrogen = (
            data.gas.ion_mass_fractions.HI + data.gas.ion_mass_fractions.HII
        )
        # abundance is in n_X/n_H unit. We convert mass fraction to abundance by dividing mass fraction of H
        sA.HI = data.gas.ion_mass_fractions.HI / mass_fraction_hydrogen / mamu["HI"]
        sA.HII = data.gas.ion_mass_fractions.HII / mass_fraction_hydrogen / mamu["HII"]
        sA.HeI = data.gas.ion_mass_fractions.HeI / mass_fraction_hydrogen / mamu["HeI"]
        sA.HeII = (
            data.gas.ion_mass_fractions.HeII / mass_fraction_hydrogen / mamu["HeII"]
        )
        sA.HeIII = (
            data.gas.ion_mass_fractions.HeIII / mass_fraction_hydrogen / mamu["HeIII"]
        )
    elif scheme.startswith("SPH M1closure"):
        sA = data.gas.rt_species_abundances
    else:
        raise ValueError("Unknown scheme", scheme)
    return sA


def trim_paramstr(paramstr):
    """
    clean up strings in the form [x,y,z,...]
    and return an array: array([x,y,z,...])
    """
    paramstr = paramstr.strip()
    if paramstr.startswith("["):
        paramstr = paramstr[1:]
    if paramstr.endswith("]"):
        paramstr = paramstr[:-1]

    params = paramstr.split(",")
    paramtrimmed = []
    for er in params:
        paramtrimmed.append(float(er))
    return paramtrimmed