# coding: utf-8
import sys
import numpy as np
import os
import pymses

import matplotlib as mpl

if os.environ.get("DISPLAY", "") == "":
    print("No display found. Using non-interactive Agg backend")
    mpl.use("Agg")

import pylab as P
import glob as glob
import pickle as pickle

from pymses.sources.ramses import output
from pymses.analysis import Camera, raytracing, slicing
from pymses.filters import CellsToPoints
from pymses.analysis import ScalarOperator, FractionOperator, MaxLevelOperator

P.rcParams["image.cmap"] = "plasma"
P.rcParams["savefig.dpi"] = 800


def make_image_disk(
    path,
    num,
    path_out=None,
    order="<",
    force=False,
    tag="",
    vel_red=20,
    map_size=512,
    put_title=True,
    cpuamr=False,
):
    """
    Make several useful image of an output of a simulation

    Parameters
    ----------
    path      path of the Ramses output
    num       Ramses output number
    path_out  path of the pipeline output
    order     '<' or '>' TODO
    force     if set, erase any existing pipeline output files
    tag       string to add to the output name
    vel_red   number of point where velocity should be plot in the slices
    map_size  size of the map
    cpuamr    plot also levels and cpus at each step
    """

    ro = pymses.RamsesOutput(path, num, order=order)
    amr = ro.amr_source(["rho", "vel", "P", "Bl", "Br"])
    rad = 0.5
    center = [0.5, 0.5, 0.5]

    make_image_aux(
        amr,
        ro,
        center,
        rad,
        num,
        path,
        force=force,
        path_out=path_out,
        map_size=map_size,
        vel_red=vel_red,
        tag=tag,
        cpuamr=cpuamr,
        put_title=put_title,
    )


def make_image_aux(
    amr,
    ro,
    center,
    radius,
    num,
    path,
    force=False,
    path_out=None,
    map_size=512,
    vel_red=20,
    tag="",
    cpuamr=False,
    put_title=True,
):
    """
    Make several useful image of an output of a simulation, auxillary function

    Parameters
    ----------
    amr
    ro        pymses.RamsesOutput object

    center    3D array for coordinates center
    num       output number
    path_out  path of the pipeline output
    force     if set, erase any existing pipeline output files
    tag       string to add to the output name
    vel_red   number of point where velocity should be plot in the slices
    map_size  size of the map
    cpuamr    plot also levels and cpus at each step
    """

    lbox = ro.info["boxlen"]  # boxlen in codeunits
    lbox_units = lbox

    G = 1.0  # Gravitational constant

    ntick = 6
    title_ax = {"x": "x (code)", "y": "y (code)", "z": "z (code)"}

    time = ro.info["time"]  # time in codeunits
    title = "t=" + str(time)[0:5] + " (code)"

    if path_out is not None:
        directory = path_out
    else:
        directory = path

    name = directory + "/coldens_z" + "_" + tag + "_" + format(num, "05") + ".jpeg"
    if len(glob.glob(name)) == 1 and not force:
        return

    rho_op = ScalarOperator(lambda dset: dset["rho"], ro.info["unit_density"])
    rt = raytracing.RayTracer(amr, ro.info, rho_op)

    axes_los = ["x", "y", "z"]  # Line of sight axes
    axes_h = ["y", "x", "x"]  # Horizontal axes
    axes_v = ["z", "z", "y"]  # Vertical axes

    ax_nb = {"x": 0, "y": 1, "z": 2}

    image_names = ["coldens", "rho", "T"]

    for i, ax_los in enumerate(axes_los):
        ax_h = axes_h[i]
        ax_v = axes_v[i]

        cam = Camera(
            center=center,
            line_of_sight_axis=ax_los,
            region_size=[2.0 * radius, 2.0 * radius],
            distance=radius,
            far_cut_depth=radius,
            up_vector=ax_v,
            map_max_size=map_size,
        )

        datamap = rt.process(cam, surf_qty=True)

        # Column density
        dmap_col = datamap.map.T * lbox
        map_col = np.log10(dmap_col)

        P.close()
        im = P.imshow(
            map_col,
            extent=[
                (-radius + center[0]) * lbox_units,
                (radius + center[0]) * lbox_units,
                (-radius + center[1]) * lbox_units,
                (radius + center[1]) * lbox_units,
            ],
            origin="lower",
        )

        P.locator_params(axis=ax_h, nbins=ntick)
        P.locator_params(axis=ax_v, nbins=ntick)

        if put_title:
            P.title(title)

        P.xlabel(title_ax[ax_h])
        P.ylabel(title_ax[ax_v])

        cbar = P.colorbar(im)
        cbar.set_label(r"$log(N)$ (code)")
        name = directory + "/coldens_" + ax_los + "_" + tag + "_" + format(num, "05")
        name_im = name + ".jpeg"
        P.savefig(name_im)

        # Rho slice
        dmap_rho = slicing.SliceMap(amr, cam, rho_op, z=0.0)
        map_rho = np.log10(dmap_rho.map)
        map_rho = map_rho.T

        vh_op = ScalarOperator(
            lambda dset: dset["vel"][:, ax_nb[ax_h]], ro.info["unit_velocity"]
        )
        dmap_vh = slicing.SliceMap(amr, cam, vh_op, z=0.0)
        map_vh_red = dmap_vh.map[
            ::vel_red, ::vel_red
        ]  # take only a subset of velocities

        map_vh_red = map_vh_red.T

        vv_op = ScalarOperator(
            lambda dset: dset["vel"][:, ax_nb[ax_v]], ro.info["unit_velocity"]
        )
        dmap_vv = slicing.SliceMap(amr, cam, vv_op, z=0.0)
        map_vv_red = dmap_vv.map[::vel_red, ::vel_red]
        map_vv_red = map_vv_red.T

        P.close()
        im = P.imshow(
            map_rho,
            extent=[
                (-radius + center[0]) * lbox_units,
                (radius + center[0]) * lbox_units,
                (-radius + center[1]) * lbox_units,
                (radius + center[1]) * lbox_units,
            ],
            origin="lower",
        )
        P.locator_params(axis=ax_h, nbins=ntick)
        P.locator_params(axis=ax_v, nbins=ntick)
        nh = map_vh_red.shape[0]
        nv = map_vv_red.shape[1]
        vec_h = (
            np.arange(nh) * 2.0 / nh * radius - radius + center[0] + radius / nh
        ) * lbox_units
        vec_v = (
            np.arange(nv) * 2.0 / nv * radius - radius + center[1] + radius / nv
        ) * lbox_units
        hh, vv = np.meshgrid(vec_h, vec_v)
        max_v = np.max(np.sqrt(map_vh_red ** 2 + map_vv_red ** 2))
        Q = P.quiver(hh, vv, map_vh_red, map_vv_red, units="width")
        P.quiverkey(
            Q,
            0.7,
            0.95,
            max_v,
            r"$" + str(max_v)[0:4] + "$ (code)",
            labelpos="E",
            coordinates="figure",
        )

        if put_title:
            P.title(title)
        P.xlabel(title_ax[ax_h])
        P.ylabel(title_ax[ax_v])
        cbar = P.colorbar(im)
        cbar.set_label(r"$log(n)$ (code)")

        name = directory + "/rho_" + ax_los + "_" + tag + "_" + format(num, "05")
        name_im = name + ".jpeg"
        P.savefig(name_im)

        P_op = ScalarOperator(lambda dset: dset["P"], ro.info["unit_pressure"])
        dmap_P = slicing.SliceMap(amr, cam, P_op, z=0.0)
        P.close()
        dmap_T = dmap_P.map.T / dmap_rho.map.T
        map_T = np.log10(dmap_T)

        im = P.imshow(
            map_T,
            extent=[
                (-radius + center[0]) * lbox_units,
                (radius + center[0]) * lbox_units,
                (-radius + center[1]) * lbox_units,
                (radius + center[1]) * lbox_units,
            ],
            origin="lower",
        )

        P.locator_params(axis="x", nbins=ntick)
        P.locator_params(axis="y", nbins=ntick)

        if put_title:
            P.title(title)
        P.xlabel(title_ax[ax_h])
        P.ylabel(title_ax[ax_v])
        cbar = P.colorbar(im)
        cbar.set_label(r"$log(T) \, (K)$")

        name = directory + "/T_" + ax_los + "_" + tag + "_" + format(num, "05")
        name_im = name + ".jpeg"
        P.savefig(name_im)

        # Toomre parameter

        if False and ax_los == "z":
            map_Q = np.log10(np.sqrt(dmap_P.map.T / dmap_rho.map.T)) / (
                np.pi * map_col * G
            )

            im = P.imshow(
                map_T,
                extent=[
                    (-radius + center[0]) * lbox_units,
                    (radius + center[0]) * lbox_units,
                    (-radius + center[1]) * lbox_units,
                    (radius + center[1]) * lbox_units,
                ],
                origin="lower",
            )

            P.locator_params(axis="x", nbins=ntick)
            P.locator_params(axis="y", nbins=ntick)

            if put_title:
                P.title(title)
                P.xlabel(title_ax[ax_h])
                P.ylabel(title_ax[ax_v])
                cbar = P.colorbar(im)
                cbar.set_label(r"$log(T) \, (K)$")

                name = directory + "/Q_" + ax_los + "_" + tag + "_" + format(num, "05")
                name_im = name + ".jpeg"
                P.savefig(name_im)

        if cpuamr:
            level_op = MaxLevelOperator()
            amr.set_read_levelmax(20)
            rt = raytracing.RayTracer(amr, ro.info, level_op)
            datamap = rt.process(cam, surf_qty=True)
            map_level = datamap.map.T

            P.close()
            im = P.imshow(
                map_level,
                extent=[
                    (-radius + center[0]) * lbox_units,
                    (radius + center[0]) * lbox_units,
                    (-radius + center[1]) * lbox_units,
                    (radius + center[1]) * lbox_units,
                ],
                origin="lower",
            )

            P.locator_params(axis="x", nbins=ntick)
            P.locator_params(axis="y", nbins=ntick)

            if put_title:
                P.title(title)
            P.xlabel(title_ax[ax_h])
            P.ylabel(title_ax[ax_v])
            cbar = P.colorbar(im)
            cbar.set_label(r"level")

            name = directory + "/level_" + ax_los + "_" + tag + "_" + format(num, "05")
            name_im = name + ".jpeg"
            P.savefig(name_im)

            cpu_op = ScalarOperator(
                lambda dset: dset.icpu * (np.ones(dset["P"].shape)),
                ro.info["unit_pressure"],
            )
            rt = raytracing.RayTracer(amr, ro.info, cpu_op)
            datamap = rt.process(cam, surf_qty=True)
            map_cpu = datamap.map.T

            P.close()
            im = P.imshow(
                map_cpu,
                extent=[
                    (-radius + center[0]) * lbox_units,
                    (radius + center[0]) * lbox_units,
                    (-radius + center[1]) * lbox_units,
                    (radius + center[1]) * lbox_units,
                ],
                origin="lower",
            )
            P.locator_params(axis="x", nbins=ntick)
            P.locator_params(axis="y", nbins=ntick)

            if put_title:
                P.title(title)

            P.xlabel(title_ax[ax_h])
            P.ylabel(title_ax[ax_v])
            cbar = P.colorbar(im)
            cbar.set_label(r"cpu")

            name = directory + "/cpu_" + ax_los + "_" + tag + "_" + format(num, "05")
            name_im = name + ".jpeg"
            P.savefig(name_im)


def disk_prop(
    path_in,
    num,
    path_out=None,
    force=False,
    nb_bin=20,
    rad_ext=1.0,
    mass_star=1.0,
    pos_star=np.array([1.0, 1.0, 1.0]),
):
    """
    Compute properties of a disk in the plane (0,x,y)
    with a protostar at the center of the box

    The region of the disk is defined by its scale height

    Parameters
    ----------

    path_in   path of the input data files (output of ramses)
    num       id of the output
    path_out  optional path to the output files
    force     if set, redo ouptut even if already done

    nb_bin    Number of radial bins
    rad_ext   Outer radius of the disk

    pos_star  position of the central protostar
    mass_star mass of the central protostar
    """

    # Set th output directory
    if path_out is not None:
        directory_out = path_out
    else:
        directory_out = path_in

    # Check if the output file exists, and exit if it is the case
    name_save = directory_out + "/prop_disk_" + str(num).zfill(5) + ".save"
    if not force and len(glob.glob(name_save)) != 0:
        return

    # Compute the bins array
    lrad = np.log10(rad_ext)
    rad = np.logspace(lrad - 2.0, lrad, num=nb_bin)

    # Get Ramses data
    ro = pymses.RamsesOutput(path_in, num)
    lbox = ro.info["boxlen"]  # boxlen in codeunits (=>pc)

    time = ro.info["time"]  # * scale_t / Myr

    # Get array of cell positions
    amr = ro.amr_source(["rho", "vel", "Br", "Bl", "P"])
    cell_source = CellsToPoints(amr)
    cells = cell_source.flatten()
    dx = cells.get_sizes()
    pos = cells.points * lbox

    # Get positions in the frame of the protostar
    pos = pos - pos_star
    # Get cylindrical radius
    rc = np.sqrt(pos[:, 0] ** 2 + pos[:, 1] ** 2)
    # Get velocities
    vel = cells["vel"] * lbox
    # Get radial component of velocity
    norm_pos = rc
    norm_pos[rc == 0] = 1.0e-10  # Avoid division per 0
    v_rad = (pos[:, 0] * vel[:, 0] + pos[:, 1] * vel[:, 1]) / norm_pos
    # Get azimuthal component of velocity
    v_az = (pos[:, 0] * vel[:, 1] - pos[:, 1] * vel[:, 0]) / norm_pos

    # Select cells that are actually in the disk, ie within the scale height
    G = 1.0
    cs = np.sqrt(cells["P"] / cells["rho"])  # sound velocity
    height = cs * np.sqrt(rc ** 3 / (G * mass_star))
    mask_pos = np.abs(pos[:, 2]) < height  # condition on position
    mask_dens = cells["rho"] > 1.0e6  # condition on density
    mask = mask_pos  # & mask_dens
    print("Number of selected cells ", np.sum(mask))

    pos_disk = pos[mask]
    rc_disk = rc[mask]
    vel_disk = vel[mask]
    rho_disk = cells["rho"][mask]  # density
    press_disk = cells["P"][mask]  # pressure
    dx_disk = dx[mask]
    dvol_disk = dx_disk ** 3
    v_rad_disk = v_rad[mask]
    v_az_disk = v_az[mask]
    v_kepl = np.sqrt(mass_star * G / rc_disk)

    # Initialize binned quantities
    cs_rad = np.zeros(nb_bin - 1)
    temp_rad = np.zeros(nb_bin - 1)
    press_rad = np.zeros(nb_bin - 1)
    rho_rad = np.zeros(nb_bin - 1)
    coldens_rad = np.zeros(nb_bin - 1)
    v_az_rad = np.zeros(nb_bin - 1)
    v_kepl_rad = np.zeros(nb_bin - 1)
    v_rad_rad = np.zeros(nb_bin - 1)
    alpha_rey_rad = np.zeros(nb_bin - 1)
    Q_kepl_rad = np.zeros(nb_bin - 1)

    for i in range(nb_bin - 1):
        mask_bin = (rc_disk > rad[i]) & (rc_disk < rad[i + 1])

        print(
            "Bin #{} : {} cells between {} and {}".format(
                i, np.sum(mask_bin), rad[i], rad[i + 1]
            )
        )

        press_rad[i] = np.sum(press_disk[mask_bin] * dvol_disk[mask_bin]) / np.sum(
            dvol_disk[mask_bin]
        )
        rho_rad[i] = np.sum(rho_disk[mask_bin] * dvol_disk[mask_bin]) / np.sum(
            dvol_disk[mask_bin]
        )
        temp_rad[i] = press_rad[i] / rho_rad[i]

        # TODO verifier unites
        # Surface of a bin :  S = dr * 2 * pi * r with
        #  dr = rad[i + 1] - rad[i] and r = (rad[i + 1] + rad[i]) / 2.
        coldens_rad[i] = (
            np.sum(rho_disk[mask_bin] * dvol_disk[mask_bin])
            * (lbox) ** 3
            / ((rad[i + 1] - rad[i]) * (rad[i + 1] + rad[i]) * np.pi)
        )

        v_az_rad[i] = np.sum(
            v_az_disk[mask_bin] * rho_disk[mask_bin] * dvol_disk[mask_bin]
        ) / np.sum(rho_disk[mask_bin] * dvol_disk[mask_bin])

        v_rad_rad[i] = np.sum(
            v_rad_disk[mask_bin] * rho_disk[mask_bin] * dvol_disk[mask_bin]
        ) / np.sum(rho_disk[mask_bin] * dvol_disk[mask_bin])

        alpha_rey_rad[i] = (
            (
                np.sum(
                    v_az_disk[mask_bin]
                    * v_rad_disk[mask_bin]
                    * rho_disk[mask_bin]
                    * dvol_disk[mask_bin]
                )
                / np.sum(dvol_disk[mask_bin] * press_disk[mask_bin])
                - v_az_rad[i] * v_rad_rad[i] * rho_rad[i] / press_rad[i]
            )
            * v_az_rad[i]
            / abs(v_az_rad[i])
        )

        v_kepl_rad[i] = np.sum(
            v_kepl[mask_bin] * rho_disk[mask_bin] * dvol_disk[mask_bin]
        ) / np.sum(rho_disk[mask_bin] * dvol_disk[mask_bin])

    # Convert to good units (TODO check)
    cs_rad = np.sqrt(temp_rad)  # *scale_v / km_s
    temp_rad = temp_rad  # * scale_T2
    press_rad = press_rad  # * scale_v**2 * scale_d

    v_az_rad = v_az_rad  # * scale_v / km_s
    v_rad_rad = v_rad_rad  # * scale_v / km_s
    v_kepl_rad = v_kepl_rad

    Q_kepl_rad = cs_rad * v_az_rad / (np.pi * G * coldens_rad * rad[0 : nb_bin - 1])

    prop_disk = {
        "time": time,
        "rad": rad[0 : nb_bin - 1],
        "center": pos_star,
        "alpha_rey": alpha_rey_rad,
        "v_rad": v_rad_rad,
        "v_az": v_az_rad,
        "v_kepl": v_kepl_rad,
        "coldens": coldens_rad,
        "rho": rho_rad,
        "press": press_rad,
        "temp": temp_rad,
        "cs": cs_rad,
        "Q_kepl": Q_kepl_rad,
    }

    # store the results
    f = open(name_save, "w")
    pickle.dump(prop_disk, f)
    f.close()


def plot_disk_prop(path, num, force=False, pdf=False, tag=""):
    """
    Plot properties of a disk

    num       id of the ramses output
    path      path to the properties file
    force     if set, redo plots even if already done

    pdf       if set, do output in pdf as well

    """

    # Load property file
    name_save = path + "/prop_disk_" + str(num).zfill(5) + ".save"

    # Check if the properties file exists
    if len(glob.glob(name_save)) == 0:
        throw("no pickle file for disk properties. Run single_z_disk_prop")
    f = open(name_save, "r")
    prop_disk = pickle.load(f)
    f.close()

    # Check if the output file exists, and exit if it is the case
    name_save = path + "/rho_disk_r_" + str(num).zfill(5) + ".jpeg"
    if not force and len(glob.glob(name_save)) != 0:
        return

    P.close()
    P.plot(
        np.log10(prop_disk["rad"]), np.log10(prop_disk["rho"]), color="k", linewidth=2
    )
    P.ylabel(r"$\log(n) \, (cm^{-3})$")
    P.xlabel("disk radius")

    if pdf:
        P.savefig(path + "/rho_disk_r_" + str(num).zfill(5) + ".pdf")
    P.savefig(path + "/rho_disk_r_" + str(num).zfill(5) + ".jpeg")

    P.close()
    P.plot(
        np.log10(prop_disk["rad"]), np.log10(prop_disk["temp"]), color="k", linewidth=2
    )
    P.ylabel(r"$\log(T) \, (K)$")
    P.xlabel("disk radius")

    if pdf:
        P.savefig(path + "/T_disk_r_" + str(num).zfill(5) + ".pdf")
    P.savefig(path + "/T_disk_r_" + str(num).zfill(5) + ".jpeg")

    P.close()

    P.xscale("log")
    P.yscale("symlog", linthreshy=0.01)

    P.plot((prop_disk["rad"]), ((prop_disk["v_rad"])), color="k", linewidth=2)
    P.plot((prop_disk["rad"]), ((prop_disk["v_kepl"])), color="b", linewidth=2)
    P.plot((prop_disk["rad"]), (abs(prop_disk["v_az"])), color="r", linewidth=2)
    P.plot((prop_disk["rad"]), ((prop_disk["cs"])), color="c", linewidth=2)

    P.legend((r"$v_r$", r"$v_{kepl}$", r"$v_\phi$", r"$c_s$"), loc="upper right")

    P.ylabel(r"$V \, (km s^{-1})$")
    P.xlabel("disk radius")

    if pdf:
        P.savefig(path + "/V_disk_r_" + str(num).zfill(5) + ".pdf")
    P.savefig(path + "/V_disk_r_" + str(num).zfill(5) + ".jpeg")

    P.close()
    P.plot(
        np.log10(prop_disk["rad"]),
        np.log10(prop_disk["coldens"]),
        color="k",
        linewidth=2,
    )
    P.ylabel(r"$\log(N) \, (cm^{-2})$")
    P.xlabel("disk radius ")

    if pdf:
        P.savefig(path + "/coldens_disk_r_" + str(num).zfill(5) + ".pdf")
    P.savefig(path + "/coldens_disk_r_" + str(num).zfill(5) + ".jpeg")

    # Alpha
    P.close()
    P.xscale("log")
    P.yscale("log")
    P.ylim([1e-5, 1.0])

    P.plot(prop_disk["rad"], abs(prop_disk["alpha_rey"]), color="b", linewidth=2)

    P.ylabel(r"$\alpha$")
    P.xlabel("disk radius ")

    if pdf:
        P.savefig(path + "/alpha_disk_r_" + str(num).zfill(5) + ".pdf")
    P.savefig(path + "/alpha_disk_r_" + str(num).zfill(5) + ".jpeg")

    # Q

    P.close()
    P.xscale("log")
    P.yscale("log", linthreshy=0.001)

    P.plot(prop_disk["rad"], abs(prop_disk["Q_kepl"]), color="b", linewidth=2)

    P.ylabel(r"$Q$")
    P.xlabel("disk radius ")

    if pdf:
        P.savefig(path + "/Q_r_" + str(num).zfill(5) + ".pdf")
    P.savefig(path + "/Q_r_" + str(num).zfill(5) + ".jpeg")