import sys
import numpy as np
import os
import pymses
import pylab as P
import glob as glob
import pandas as pd
import pickle as pickle


from pymses.sources.ramses import output
from pymses.analysis import Camera, raytracing, slicing


def make_image_temp(
    amr,
    ro,
    center,
    radius,
    num,
    path,
    i_im=0,
    force=False,
    path_out=None,
    col_dens_only=False,
    map_size=512,
    vel_red=20,
    im_pos=0.0,
    tag="",
    cpuamr=False,
    put_title=True,
    mass_norm=1.0,
    AU_units=False,
):
    """
    Make several useful image of an output of a simulation

    Parameters
    ----------
    amr
    ro     ramses output object

    center 3D array for coordinates center
    num    output number
    """

    lbox = ro.info["boxlen"]  # boxlen in codeunits (=>pc)

    (
        AU,
        pc,
        Ms,
        Myr,
        scale_n,
        scale_d,
        scale_t,
        scale_l,
        scale_v,
        scale_T2,
        scale_ener,
        scale_mag,
        microG,
        km_s,
        Cwnm,
        scale_mass,
        unit_col,
        lbox_pc,
    ) = normalisation(ro)

    ntick = 6
    if AU_units:
        units = pc / AU
        lbox_units = lbox * units
        titre_x = "x (AU)"
        titre_y = "y (AU)"
        titre_z = "z (AU)"
    else:
        units = 1.0
        lbox_units = lbox
        titre_x = "x (pc)"
        titre_y = "y (pc)"
        titre_z = "z (pc)"

    lbox_cm = lbox * pc  # lbox in cm

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

    titre = "t=" + str(time)[0:5] + " (Myr)"

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

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

    ax_names = ["x", "y", "z"]
    up_axes = ["z", "z", "x"]
    other_axes = ["z", "x", "y"]

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

    for i, ax_name in enumerate(ax_names):
        cam = Camera(
            center=center,
            line_of_sight_axis=ax_name,
            region_size=[2.0 * radius, 2.0 * radius],
            distance=radius,
            far_cut_depth=radius,
            up_vector=up_axes[i],
            map_max_size=map_size,
        )

        datamap = rt.process(cam, surf_qty=True)
        map_col = np.log10(datamap.map.T * lbox_cm)

        P.clf()
        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=other_axes[i], nbins=ntick)
        P.locator_params(axis=up_axes[i], nbins=ntick)

        if put_title:
            P.title(titre)

        P.xlabel(titre_x)
        P.ylabel(titre_y)
        cbar = P.colorbar(im)
        cbar.set_label(r"$log(N) \, cm^{-2}$")

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

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

        vx_op = ScalarOperator(lambda dset: dset["vel"][:, 0], ro.info["unit_velocity"])
        dmap_vx = slicing.SliceMap(amr, cam_z, vx_op, z=0.0)
        map_vx_red = dmap_vx.map[::vel_red, ::vel_red]

        map_vx_red = map_vx_red.T

        vy_op = ScalarOperator(lambda dset: dset["vel"][:, 1], ro.info["unit_velocity"])
        dmap_vy = slicing.SliceMap(amr, cam, vy_op, z=0.0)
        map_vy_red = dmap_vy.map[::vel_red, ::vel_red]

        map_vy_red = map_vy_red.T

        P.clf()
        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=other_axes[i], nbins=ntick)
        P.locator_params(axis=up_axes[i], nbins=ntick)

        nx = map_vx_red.shape[0]
        ny = map_vx_red.shape[1]
        vec_x = (
            np.arange(nx) * 2.0 / nx * radius - radius + center[0] + radius / nx
        ) * lbox_units
        vec_y = (
            np.arange(ny) * 2.0 / ny * radius - radius + center[1] + radius / nx
        ) * lbox_units
        xx, yy = np.meshgrid(vec_x, vec_y)

        map_vx_red = map_vx_red * scale_v / km_s
        map_vy_red = map_vy_red * scale_v / km_s
        max_v = np.max(np.sqrt(map_vx_red ** 2 + map_vy_red ** 2))

        Q = P.quiver(xx, yy, map_vx_red, map_vy_red, units="width")
        P.quiverkey(
            Q,
            0.7,
            0.95,
            max_v,
            r"$" + str(max_v)[0:4] + "\, km \, s^{-1}$",
            labelpos="E",
            coordinates="figure",
        )

        if put_title:
            P.title(titre)
        P.xlabel(titre_x)
        P.ylabel(titre_y)
        cbar = P.colorbar(im)
        cbar.set_label(r"$log(n) \, (cm^{-3})$")

        name = directory + "/rho_z" + "_" + tag + str(i_im) + "_" + 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_z, P_op, z=0.0)
        P.clf()
        map_T = np.log10(dmap_P.map.T / dmap_rho.map.T * scale_T2)

        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(titre)

        P.xlabel(titre_x)
        P.ylabel(titre_y)
        cbar = P.colorbar(im)
        cbar.set_label(r"$log(T) \, (K)$")

        name = directory + "/T_z" + "_" + tag + str(i_im) + "_" + 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_z, surf_qty=True)
            map_level = datamap.map.T

            P.clf()
            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(titre)
            P.xlabel(titre_x)
            P.ylabel(titre_y)
            cbar = P.colorbar(im)
            cbar.set_label(r"level")

            name = (
                directory + "/level_z" + "_" + tag + str(i_im) + "_" + format(num, "05")
            )
            name_im = name + ".jpeg"
            P.savefig(name_im)
            if ps:
                P.savefig(name + ".ps")

            if descrip is not None:
                dd = {}
                dd.update(
                    {
                        "name": "mean AMR level in the z-direction"
                        + " ("
                        + tag
                        + str(i_im)
                        + ")"
                    }
                )
                dd.update({"type": "image"})
                dd.update({"description": "Mean AMR level in the z-direction."})
                dd.update({"display": 0})
                dd.update({"item": i_im})
                dd.update({"tag": tag})
                data = {"jpeg": name_im}
                dd.update({"data": data})
                list_im.append(dd)

            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_z, surf_qty=True)
            map_cpu = datamap.map.T

            P.clf()
            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(titre)

            P.xlabel(titre_x)
            P.ylabel(titre_y)
            cbar = P.colorbar(im)
            cbar.set_label(r"level")

            name = (
                directory + "/cpu_z" + "_" + tag + str(i_im) + "_" + 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=100.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
    if (
        not force
        and len(glob.glob(directory_out + "/prop_disk_" + str(num).zfill(5) + ".save"))
        != 0
    ):
        return

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

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

    (
        AU,
        pc,
        Ms,
        Myr,
        scale_n,
        scale_d,
        scale_t,
        scale_l,
        scale_v,
        scale_T2,
        scale_ener,
        scale_mag,
        microG,
        km_s,
        Cwnm,
        scale_mass,
        unit_col,
        lbox_pc,
    ) = me.normalisation(ro)

    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

    # 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"]
    # 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
    # TODO Check units
    G = 6.8e-8
    cs = np.sqrt(cells["P"] / cells["rho"]) * scale_v  # 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 or 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]

    # TODO Check what do that does
    nzoom = 9
    eps = 0.5 ** nzoom
    map_coldens, map_w13, xedges, yedges = me.make_hierarch_map(
        pos_disk_x,
        pos_disk_y,
        pos_disk_z,
        dx_disk,
        rho_disk,
        rho_disk,
        eps,
        center=[0.0, 0.0, 0.0],
        make_image=do_plot,
        path_out=directory_out,
        tag="xy_" + str(num).zfill(5),
    )

    map_coldens, map_w13, xedges, yedges = me.make_hierarch_map(
        pos_disk_z,
        pos_disk_x,
        pos_disk_y,
        dx_disk,
        rho_disk,
        rho_disk,
        eps,
        center=[0.0, 0.0, 0.0],
        make_image=do_plot,
        path_out=directory_out,
        tag="xz_" + str(num).zfill(5),
    )

    # Initialize binned quantities
    norm_rad = lbox * scale_l / AU  # radius in AU
    rdisk_AU = rad_disk * norm_rad

    cs_rad = np.zeros(nbin - 1)
    temp_rad = np.zeros(nbin - 1)
    press_rad = np.zeros(nbin - 1)
    rho_rad = np.zeros(nbin - 1)
    coldens_rad = np.zeros(nbin - 1)
    v_az_rad = np.zeros(nbin - 1)
    v_rad_rad = np.zeros(nbin - 1)

    for i in range(nbin - 1):
        mask_bin = (rdisk_AU > rad[i]) or (rdisk_AU < rad[i + 1])

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

        # TODO vérifier unités
        # 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] * dvol_disk[mask])
            * (lbox * pc) ** 3
            / ((rad[i + 1] - rad[i]) * (rad[i + 1] + rad[i]) * np.pi * AU ** 2)
        )

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

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

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

        prop_disk = {
            "time": time,
            "rad_AU": rad[0 : nbin - 1],
            "center": pos_mass,
            "alpha_rey": alpha_rey_rad,
            "v_rad": v_rad_rad,
            "v_az": v_az_rad,
            "coldens": coldens_rad,
            "rho": rho_rad,
            "press": press_rad,
            "temp": temp_rad,
            "cs": cs_rad,
        }

    # 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

    # store the results
    name_save = directory_out + "/prop_disk_" + str(num).zfill(5) + ".save"
    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:
        print("no pickle file for disk properties. Run single_z_disc_prop")
        return
    f = open(name_save, "r")
    prop_disk = pickle.load(f)
    f.close()

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

    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.clf()
    P.plot(
        np.log10(prop_disc["rad_AU"]),
        np.log10(prop_disc["temp"]),
        color="k",
        linewidth=2,
    )
    P.ylabel(r"$\log(T) \, (K)$")
    P.xlabel("disc radius (AU)")

    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.clf()

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

    P.plot((prop_disc["rad_AU"]), ((prop_disc["v_rad"])), color="k", linewidth=2)
    P.plot((prop_disc["rad_AU"]), (abs(prop_disc["v_az"])), color="r", linewidth=2)
    P.plot((prop_disc["rad_AU"]), ((prop_disc["cs"])), color="c", linewidth=2)

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

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

    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.clf()
    P.plot(
        np.log10(prop_disc["rad_AU"]),
        np.log10(prop_disc["coldens"]),
        color="k",
        linewidth=2,
    )
    P.ylabel(r"$\log(N) \, (cm^{-2})$")
    P.xlabel("disc radius (AU)")

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

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

    P.plot(prop_disc["rad_AU"], prop_disc["alpha_rey"], color="b", linewidth=2)

    P.legend((r"$\alpha _{Rey}$"), loc="upper right")

    P.ylabel(r"$\alpha}$")
    P.xlabel("disc radius (AU)")

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