#!/usr/bin/env python
# -*- coding: utf-8 -*-
# This file is part of the 'astrophysix' Python package.
#
# Copyright © Commissariat a l'Energie Atomique et aux Energies Alternatives (CEA)
#
#  FREE SOFTWARE LICENCING
#  -----------------------
# This software is governed by the CeCILL license under French law and abiding by the rules of distribution of free
# software. You can use, modify and/or redistribute the software under the terms of the CeCILL license as circulated by
# CEA, CNRS and INRIA at the following URL: "http://www.cecill.info". As a counterpart to the access to the source code
# and rights to copy, modify and redistribute granted by the license, users are provided only with a limited warranty
# and the software's author, the holder of the economic rights, and the successive licensors have only limited
# liability. In this respect, the user's attention is drawn to the risks associated with loading, using, modifying
# and/or developing or reproducing the software by the user in light of its specific status of free software, that may
# mean that it is complicated to manipulate, and that also therefore means that it is reserved for developers and
# experienced professionals having in-depth computer knowledge. Users are therefore encouraged to load and test the
# software's suitability as regards their requirements in conditions enabling the security of their systems and/or data
# to be ensured and, more generally, to use and operate it in the same conditions as regards security. The fact that
# you are presently reading this means that you have had knowledge of the CeCILL license and that you accept its terms.
#
#
# COMMERCIAL SOFTWARE LICENCING
# -----------------------------
# You can obtain this software from CEA under other licencing terms for commercial purposes. For this you will need to
# negotiate a specific contract with a legal representative of CEA.
#
from __future__ import print_function, unicode_literals

import datetime
import os

import h5py
import numpy as N
from astrophysix import units as U
from astrophysix.simdm import Project, ProjectCategory, SimulationStudy
from astrophysix.simdm.datafiles import Datafile, PlotInfo, PlotType
from astrophysix.simdm.experiment import (
    AppliedAlgorithm,
    ParameterSetting,
    ParameterVisibility,
    ResolvedPhysicalProcess,
    Simulation,
)
from astrophysix.simdm.protocol import (
    Algorithm,
    AlgoType,
    InputParameter,
    PhysicalProcess,
    Physics,
    SimulationCode,
)
from astrophysix.simdm.results import GenericResult, Snapshot
from astrophysix.utils.file import FileType
from matplotlib import gridspec
from mpl_toolkits.axes_grid1 import AxesGrid, Grid

from plotter import *

P.rcParams["text.usetex"] = False


pp_params = default_params()
pp_params.input.nml_filename = "disk.nml"
# pp_params.out.interactive = True
pp_params.pymses.map_size = 2048
pp_params.pymses.zoom = 4
pp_params.pymses.filter = False
pp_params.pymses.variables = ["rho", "vel", "P", "g"]
pp_params.pymses.multiprocessing = True
pp_params.process.verbose = True
pp_params.disk.enable = True
pp_params.disk.nb_bin = 100
pp_params.pdf.nb_bin = 100
pp_params.process.num_process = 10
in_dir = "/drf/projets/alfven-data/nbrucy/simus/fragdisk"
out_dir = "/dsm/anais/storageA/nbrucy/visus/fragdisk/mnras"
nml_key = "cloud_params/beta_cool"

# --- Runs  -----

pp_params.astrophysix.simu_fmt = "beta{nml[cloud_params/beta_cool]:g}_{tag:.8}"
pp_params.astrophysix.descr_fmt = (
    "Group {tag:.8}, $\\beta$ = {nml[cloud_params/beta_cool]}"
)

pl_orp = Plotter(
    in_dir,
    filter_name="104_beta4_jr13",
    in_nums="last",
    path_out=out_dir,
    pp_params=pp_params,
)
orp = cst.Unit.create_unit("ORP", base_unit=pl_orp.comp.info["unit_time"] * 0.79)

# JR13_TIC
runs = "*_jr13"
pl_jr13 = Plotter(
    in_dir,
    filter_name=runs,
    in_nums="all",
    sort_run_by=nml_key,
    path_out=out_dir,
    tag="jr13_tic_mnras",
    pp_params=pp_params,
    unit_time=orp,
)
print("JR13_TIC defined")

# JR12_TIC
runs_12 = "0[0-9][0-9]_beta*_jr12"
pl_jr12_tic = Plotter(
    in_dir,
    filter_name=runs_12,
    in_nums="all",
    sort_run_by=nml_key,
    filter_nml=("cloud_params", "!=", 7),
    path_out=out_dir,
    tag="jr12_tic_mnras",
    pp_params=pp_params,
    unit_time=orp,
)

pp_params.astrophysix.simu_fmt = "beta{nml[cloud_params/beta_cool]:g}_{tag:.4}"
pp_params.astrophysix.descr_fmt = (
    "Group {tag:.4}, $\\beta$ = {nml[cloud_params/beta_cool]}"
)
print("JR12_TIC defined")

# JR12
in_dir_conv = "/drf/projets/alfven-data/nbrucy/simus/conv_disk"
out_dir_conv = "/dsm/anais/storageA/nbrucy/visus/conv_disk"
runs = "[7-8][0-9]_beta*_j*"
pl_jr12 = Plotter(
    in_dir_conv,
    filter_name=runs,
    in_nums="all",
    sort_run_by=nml_key,
    path_out=out_dir,
    tag="jr12_mnras",
    pp_params=pp_params,
    unit_time=orp,
)
print("JR12 defined")

# JR11
runs = "*beta*_jr11"
pl_l11 = Plotter(
    in_dir_conv,
    filter_name=runs,
    sort_run_by=nml_key,
    filter_nml=("cloud_params/beta_cool", ">", 3),
    path_out=out_dir_conv,
    tag="jr11_mnras",
    pp_params=pp_params,
    unit_time=orp,
)

print("JR11 defined")


pls = [pl_l11, pl_jr12, pl_jr12_tic, pl_jr13]


# ----------------------------------------------- Project creation --------------------------------------------------- #
# Available project categories are :
# - ProjectCategory.SolarMHD
# - ProjectCategory.PlanetaryAtmospheres
# - ProjectCategory.StarPlanetInteractions
# - ProjectCategory.StarFormation
# - ProjectCategory.Supernovae
# - ProjectCategory.GalaxyFormation
# - ProjectCategory.GalaxyMergers
# - ProjectCategory.Cosmology
proj = Project(
    category=ProjectCategory.StarPlanetInteractions,
    project_title="Fragdisk",
    alias="FRAGDISK",
    short_description="Fragmentation of self-gravitating disks",
    general_description="""
    Study of the fragmentation of self-gravitating disks. See Brucy & Hennebelle 2021 (submitted) for more details.

    This database is currently being completed.

    Abstract:

    Self-gravitating disks are believed to play an important role in astrophysics in particular regarding the star and planet formation process.
    In this context, disks subject to an idealized cooling process, characterized by a cooling timescale β expressed in unit of orbital timescale,
    have been extensively studied. We take advantage of the Riemann solver and the 3D Godunov scheme implemented in the code Ramses
    to perform high resolution simulations, complementing previous studies that have used Smoothed Particle Hydrodynamics (SPH) or 2D grid codes.
    """,
    data_description="""The data available for this project is the underlying data of the article Brucy & Hennebelle 2021.
    The data is not already fully uploaded. 3D datacube extraction on demand is planned""",
    directory_path="~nbrucy/simus/fragdisk",
)
print(proj)
# -------------------------------------------------------------------------------------------------------------------- #

# -------------------------------------------------------------------------------------------------------------------- #
redo = True
for pl in pls:
    pl.pp_params.process.verbose = True
    pl.comp.pp_params.process.verbose = True
    for run in pl.runs:
        simu = pl.simulations[run]
        proj.simulations.add(simu)


# -------------------------------------------------------------------------------------------------------------------- #


for pl in pls:
    select = {"time": 4.5}
    pl.coldens(
        "z",
        overwrite=redo,
        overwrite_dep=False,
        unit_space=cst.cm,
        unit_time=orp,
        nml_key="cloud_params/beta_cool",
        vmin=1e-2,
        vmax=1e2,
        put_units=False,
        select=select,
        label=r"$\Sigma$",
    )

    pl.coldens(
        "y",
        overwrite=redo,
        overwrite_dep=False,
        unit_space=cst.cm,
        unit_time=orp,
        nml_key="cloud_params/beta_cool",
        vmin=1e-2,
        vmax=1e2,
        put_units=False,
        select=select,
        label=r"$\Sigma$",
    )

    pl.slice_rho(
        "z",
        overwrite=redo,
        overwrite_dep=False,
        unit_space=cst.cm,
        unit_time=orp,
        nml_key="cloud_params/beta_cool",
        put_units=False,
        select=select,
        label=r"$\rho$",
    )

    pl.slice_rho(
        "y",
        overwrite=redo,
        overwrite_dep=False,
        unit_space=cst.cm,
        unit_time=orp,
        nml_key="cloud_params/beta_cool",
        put_units=False,
        select=select,
        label=r"$\rho$",
    )

    pl.slice_P(
        "z",
        overwrite=redo,
        overwrite_dep=False,
        unit_space=cst.cm,
        unit_time=orp,
        nml_key="cloud_params/beta_cool",
        put_units=False,
        select=select,
        label=r"$P$",
    )

    pl.pdf_coldens(
        "z",
        overwrite=redo,
        overwrite_dep=False,
        unit_time=orp,
        nml_key="cloud_params/beta_cool",
        label=r"$\log(\sigma)$",
        kind="step",
        color="k",
        select=select,
    )


# -------------------------------------------------------------------------------------------------------------------- #


# Create HDF5 files
for simu in proj.simulations:
    for snap in simu.snapshots:
        for df in snap.datafiles:
            name = df[FileType.JPEG_FILE].filename
            name = os.path.splitext(name)[0] + ".h5"
            h5 = h5py.File(out_dir + "/" + name, "w")
            p = h5.create_group("plot")
            df.plot_info.hsp_save_to_h5(p)
            h5.close()
            df[FileType.HDF5_FILE] = out_dir + "/" + name

for param in ramses.input_parameters:
    param.key = os.path.basename(param.key)

study = SimulationStudy(project=proj)

for sim in study.project.simulations:
    for snap in sim.snapshots:
        snap.time = (snap.time[0], cst.year)

proj.galactica_validity_check()
study.save_HDF5(out_dir + "/fragdisk_study.h5")