# coding: utf-8

import numpy as np
import pandas as pd
from plotter import U
import snapshotprocessor

mp = 1.4 * 1.66 * 10 ** (-24) * U.g
z0 = 150 * U.pc
sink_header = [
    "Id",
    "M",
    "dmf",
    "x",
    "y",
    "z",
    "vx",
    "vy",
    "vz",
    "rot_period",
    "lx",
    "ly",
    "lz",
    "acc_rate",
    "acc_lum",
    "age",
    "int_lum",
    "Teff",
]


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


convert_coldens = np.vectorize(convert_coldens_s)


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

def get_sinks(pp):
    csv_name = f"{pp.path}/output_{pp.num:05}/sink_{pp.num:05}.csv"
    return pd.read_csv(csv_name, header=None, names=sink_header)


def analyze_box(pp):
    pp.cells["mass"] = snapshotprocessor.mass_func(pp.cells)
    coldens = pp.get_value("/maps/coldens_z", unit=U.coldens)
    sinks = get_sinks(pp)
    sinks["mass"] = sinks.M
    res = {}
    dirs = ["x", "y", "z"]
    res["time"] = pp.info["time"] * pp.info["unit_time"].express(U.Myr)
    for i, dir in enumerate(dirs):
        pp.cells[f"v{dir}"] = pp.cells["vel"][:, i]
        res[f"sigma_{dir}"] = get_dispersion(pp.cells, f"v{dir}") * pp.info[
            "unit_velocity"
        ].express(U.km_s)
        res[f"sigma_sinks_{dir}"] = get_dispersion(sinks, f"v{dir}") * pp.info[
            "unit_velocity"
        ].express(U.km_s)
    res["coldens_mean"] = np.mean(coldens)
    res["n0"] = pp.get_nml("cloud_params/dens0")
    res["mass"] = np.sum(
        pp.cells["mass"]
        * (pp.info["unit_density"] * pp.info["unit_length"] ** 3).express(U.Msun)
    )
    res["coldens_initial"] = convert_coldens_s(res["n0"])
    res["mass_initial"] = res["coldens_initial"] * 1e6
    res["coldens_mean"] = np.mean(coldens)
    res["coldens_beam"] = res["mass"] / (pp.info["unit_length"].express(U.pc)) ** 2
    res["nsink"] = sinks.M.count()
    res["mass_sink"] = np.sum(sinks.M)

    return res


def load_wrapper(pp, fun):
    """
    Wrapper to load_unload data and map function
    """
    pp.load_cells(keys=["pos", "vel", "dx", "rho"])
    pp.coldens("z")
    res = fun(pp)
    pp.unload_cells()
    return res


def allinone(pp):
    return load_wrapper(pp, analyze_box)