# 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