# coding: utf-8

import numpy as np
from plotter import U


def get_dispersion(dset, name):
    """
    Compute dispersion from dset["name"]
    """
    vel = dset[name]
    mass = dset["mass_kg"]
    mass_tot = np.sum(mass)
    mean = np.sum(mass * vel) / mass_tot
    return np.sqrt(np.sum(mass * (vel - mean) ** 2) / mass_tot)


def get_polar_sigma(dset):
    """
    Get speed dispersion in polar coordinates
    """
    return {
        velname: get_dispersion(dset, velname) for velname in ["velphi", "velr", "velz"]
    }


def get_gas_dm_stars(pp):
    # Load  arrays
    pp.load_parts()
    pp.load_cells()
    cells = pp.cells
    parts = pp.parts

    # Compute extra fields and convert units
    for dset in (cells, parts):
        dset["pos_kpc"] = dset["pos"] - np.array([0.5, 0.5, 0.5])
        dset["pos_kpc"] *= pp.info["unit_length"].express(U.kpc)
        dset["r"] = np.sqrt(np.sum(dset["pos_kpc"][:, :2] ** 2, axis=1))
        dset["phi"] = np.angle(dset["pos_kpc"][:, 0] + dset["pos_kpc"][:, 1] * 1j)
        dset["velphi"] = pp.getter_vect_phi(dset, "vel") * pp.info[
            "unit_velocity"
        ].express(U.km_s)
        dset["velr"] = pp.getter_vect_r(dset, "vel") * pp.info["unit_velocity"].express(
            U.km_s
        )
        dset["velz"] = dset["vel"][:, 2] * pp.info["unit_velocity"].express(U.km_s)

    cells["mass_kg"] = (
        cells["rho"]
        * cells["dx"] ** 3
        * (pp.info["unit_density"] * pp.info["unit_length"] ** 3).express(U.kg)
    )

    parts["mass_kg"] = parts["mass"] * pp.info["unit_mass"].express(U.kg)

    for dset in (cells, parts):
        dset["ek"] = dset["mass_kg"] * np.sum(dset["vel"] ** 2, axis=1)
        dset["ek"] *= (U.kg * pp.info["unit_velocity"] ** 2).express(U.J)

    # Separate DM from stars
    mass_dm = np.max(parts["mass_kg"])
    mask_dm = parts["mass_kg"] == mass_dm
    mask_star = parts["mass_kg"] < mass_dm

    # Create separated arrays
    gas = cells
    dm = {key: parts[key][mask_dm] for key in parts}
    stars = {key: parts[key][mask_star] for key in parts}

    # Store arrays and return them
    pp.gas = gas
    pp.dm = dm
    pp.stars = stars
    return gas, dm, stars


def extract_polar_region(dset, r=4, dr=0.5, phi=0, dphi=0.125, z=0, dz=0.5):
    """
    Returns a mask for dset of the polar volume within polar coordinates
    [r - dr, r + dr] x [phi - dphi, phi + dphi] x [z -dz, z + dz]
    """
    mask_box = (
        (np.abs(dset["r"] - r) < dr)
        & (np.abs(dset["phi"] - phi) < dphi)
        & (np.abs(dset["pos"][:, 2] - z) < dz)
    )
    return {key: dset[key][mask_box] for key in dset}


def sector_analysis(pp, r=4, dr=0.5, phi=0, dphi=0.125, z=0, dz=0.5):
    """
    Analyze box at given coordinates
    """
    gds = [pp.gas, pp.dm, pp.stars]
    gds_box = [extract_polar_region(dset, r, dr, phi, dphi, z, dz) for dset in gds]
    result = {}
    result["ek"] = np.array([np.sum(dset["ek"]) for dset in gds_box])  # J
    result["mass"] = np.array([np.sum(dset["mass_kg"]) for dset in gds_box])  # kg
    result["ek_spe"] = result["ek"] / result["mass"]  # J.kg^-1
    for dset in gds_box:
        sigma = get_polar_sigma(dset)
        for key in sigma:
            keyres = "sigma_" + key
            if keyres in result:
                result[keyres].append(sigma[key])
            else:
                result[keyres] = [sigma[key]]
    return result


def ring_analysis(pp, r=4, dr=0.5, dphi=0.125, z=0, dz=0.5, do_mean=True):
    """
    Compute the average at a given of quantities computed in polar sectors.
    """

    phi_sectors = np.arange(-np.pi + dphi, np.pi - dphi, 2 * dphi)
    data_sec = [sector_analysis(pp, r, dr, phi, dphi, z, dz) for phi in phi_sectors]

    data = {}
    mwavg = {}
    tot_mass = [np.sum([d["mass"][i] for d in data_sec]) for i in range(3)]
    for key in data_sec[0]:
        mwavg[key] = []
        for i in range(3):
            mwavg[key].append(
                np.sum([d["mass"][i] * d[key][i] for d in data_sec]) / tot_mass[i]
            )

        data[key] = np.array([d[key] for d in data_sec])

    return mwavg, phi_sectors, data


def get_time_from_relax(pp):
    pp.load_parts()
    epoch = pp.parts["epoch"].copy()
    epoch *= pp.info["unit_time"].express(U.Myr)
    trelax = np.min(epoch[epoch > 0])
    tfromrelax = np.max(epoch - trelax)
    return tfromrelax


def get_last_sfr(pp):
    pp.load_parts()
    epoch = pp.parts["epoch"].copy()
    epoch *= pp.info["unit_time"].express(U.year)
    mass = pp.parts["mass"].copy()
    mass *= pp.info["unit_mass"].express(U.Msun)
    mask = epoch > 0
    masstot, time = np.histogram(epoch[mask], weights=mass[mask], bins=100)
    dtime = np.diff(time)
    sfr = masstot[-1] / dtime[-1]
    return sfr


def colmean4kpc(pp):
    pp.coldens("z")
    pp.rr("z")
    col = pp.get_value("/maps/coldens_z", unit=U.coldens)
    rr = pp.get_value("/maps/rr_z", unit=U.kpc)
    colmean = np.mean(col[np.abs(rr - 4) < 0.5])
    return colmean


def allinone(pp):

    get_gas_dm_stars(pp)
    res = {}
    res["run"] = pp.run
    res["num"] = pp.num
    res["coldens"] = colmean4kpc(pp)
    res["sfr"] = get_last_sfr(pp)
    res["time"] = get_time_from_relax(pp)
    ring = ring_analysis(pp)[0]
    for key in ring:
        for i, fluid in enumerate(["gas", "dm", "stars"]):
            res[f"{key}_{fluid}"] = ring[key][i]
    return res