#!/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 os

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

# ----------------------------------------------- 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.StarFormation,
    project_title="Ismfeed",
    alias="ismfeed",
    short_description="Impact of feedback and turbulence on star formation",
    general_description="""Impact of feedback and turbulence on star formation. The simulation presented here are described in Brucy et al. 2020, ApJL, L38""",
    data_description="The data available in this project...",
    directory_path="~nbrucy/simus/ismfeed",
)
print(proj)
# -------------------------------------------------------------------------------------------------------------------- #


# --------------------------------------- Simulation code definition ------------------------------------------------- #
ramses = SimulationCode(
    name="Ramses 3 (MHD)",
    code_name="Ramses",
    code_version="3.10.1",
    alias="RAMSES_3",
    url="https://www.ics.uzh.ch/~teyssier/ramses/RAMSES.html",
    description="Ramses MHD code",
)
# => Add algorithms : available algorithm types are :
# - AlgoType.AdaptiveMeshRefinement
# - AlgoType.VoronoiMovingMesh
# - AlgoType.SmoothParticleHydrodynamics
# - AlgoType.Godunov
# - AlgoType.PoissonMultigrid
# - AlgoType.PoissonConjugateGradient
# - AlgoType.ParticleMesh
# - AlgoType.FriendOfFriend
# - AlgoType.HLLCRiemann
# - AlgoType.RayTracer
# amr = ramses.algorithms.add(Algorithm(algo_type=AlgoType.AdaptiveMeshRefinement, description="AMR"))
ramses.algorithms.add(
    Algorithm(algo_type=AlgoType.Godunov, description="Godunov scheme")
)
ramses.algorithms.add(
    Algorithm(algo_type=AlgoType.HLLCRiemann, description="HLLC Riemann solver")
)
ramses.algorithms.add(
    Algorithm(
        algo_type=AlgoType.PoissonMultigrid, description="Multigrid Poisson solver"
    )
)
ramses.algorithms.add(
    Algorithm(algo_type=AlgoType.ParticleMesh, description="PM solver")
)

# => Add input parameters
ramses.input_parameters.add(
    InputParameter(
        key="amr_params/levelmin",
        name="Lmin",
        description="min. level of AMR refinement",
    )
)
ramses.input_parameters.add(
    InputParameter(
        key="amr_params/levelmax",
        name="Lmax",
        description="max. level of AMR refinement",
    )
)

ramses.input_parameters.add(
    InputParameter(
        key="turb_params/turb_rms", name="f_rms", description="Amplitude of the driving"
    )
)
ramses.input_parameters.add(
    InputParameter(
        key="cloud_params/dens0", name="n_0", description="Midplane density in cm**-3"
    )
)
ramses.input_parameters.add(
    InputParameter(
        key="cloud_params/bx_bound",
        name="bx_bound",
        description="imposed magnetic field at the x boundary (1 implies that the magnetic pressure is equal to thermal pressure in WNM, which corresponds to about 5muG)",
    )
)


# => Add physical processes : available physics are :
# - Physics.SelfGravity
# - Physics.Hydrodynamics
# - Physics.MHD
# - Physics.StarFormation
# - Physics.SupernovaeFeedback
# - Physics.AGNFeedback
# - Physics.MolecularCooling
ramses.physical_processes.add(
    PhysicalProcess(
        physics=Physics.StarFormation,
        description="Star Formation is triggered when density overpass",
    )
)
ramses.physical_processes.add(
    PhysicalProcess(
        physics=Physics.MHD, description="Magneto-hydrodynamical equations are solved"
    )
)
ramses.physical_processes.add(
    PhysicalProcess(physics=Physics.SelfGravity, description="Self-Gravity is applied.")
)
ramses.physical_processes.add(
    PhysicalProcess(physics=Physics.SupernovaeFeedback, description="SN feedback")
)
# -------------------------------------------------------------------------------------------------------------------- #


# -------------------------------------------- Simulation setup ------------------------------------------------------ #
simu = Simulation(
    simu_code=ramses,
    name="Noturb_1",
    alias="noturb_1",
    description="Simulation without turbulence",
    directory_path="~nbrucy/simus/turb/",
)
proj.simulations.add(simu)

# Add applied algorithms implementation details. Warning : corresponding algorithms must have been added in the 'ramses'
# simulation code.
# simu.applied_algorithms.add(AppliedAlgorithm(algorithm=amr, details="My AMR implementation [Teyssier 2002]"))
# simu.applied_algorithms.add(AppliedAlgorithm(algorithm=ramses.algorithms[AlgoType.HLLCRiemann.name],
#                                             details="My Riemann solver implementation [Teyssier 2002]"))

# Add parameter setting. Warning : corresponding input parameter must have been added in the 'ramses' simulation code.
# Available parameter visibility options are :
# - ParameterVisibility.NOT_DISPLAYED
# - ParameterVisibility.ADVANCED_DISPLAY
# - ParameterVisibility.BASIC_DISPLAY
simu.parameter_settings.add(
    ParameterSetting(
        input_param=ramses.input_parameters["levelmin"],
        value=8,
        visibility=ParameterVisibility.BASIC_DISPLAY,
    )
)
simu.parameter_settings.add(
    ParameterSetting(
        input_param=lmax, value=12, visibility=ParameterVisibility.BASIC_DISPLAY
    )
)

# Add resolved physical process implementation details. Warning : corresponding physical process must have been added to
# the 'ramses' simulation code
simu.resolved_physics.add(
    ResolvedPhysicalProcess(
        physics=ramses.physical_processes[Physics.StarFormation.name],
        details="Star formation specific implementation",
    )
)
simu.resolved_physics.add(
    ResolvedPhysicalProcess(
        physics=grav, details="self-gravity specific implementation"
    )
)
# -------------------------------------------------------------------------------------------------------------------- #


# -------------------------------------- Simulation generic result and snapshots ------------------------------------- #
# Generic result
gres = GenericResult(
    name="Key result 1 !",
    description="My description",
    directory_path="/my/path/to/result",
)
simu.generic_results.add(gres)

# Simulation snapshot
# In one-line
sn = simu.snapshots.add(
    Snapshot(
        name="My best snapshot !",
        description="My first snapshot description",
        time=(125, U.kyr),
        physical_size=(250.0, U.kpc),
        directory_path="/path/to/snapshot1",
        data_reference="OUTPUT_00056",
    )
)
# Or create snapshot, then add it to the simulation
sn2 = Snapshot(
    name="My second best snapshot !",
    description="My second snapshot description",
    time=(0.26, U.Myr),
    physical_size=(0.25, U.Mpc),
    directory_path="/path/to/snapshot2",
    data_reference="OUTPUT_00158",
)
simu.snapshots.add(sn2)
# -------------------------------------------------------------------------------------------------------------------- #


# ---------------------------------------------------- Result datafiles ---------------------------------------------- #
# Datafile creation
imf_df = sn.datafiles.add(
    Datafile(
        name="Initial mass function plot",
        description="This is my plot detailed description",
    )
)

# Add attached files to a datafile (1 per file type). Available file types are :
# - FileType.HDF5_FILE
# - FileType.PNG_FILE
# - FileType.JPEG_FILE
# - FileType.FITS_FILE
# - FileType.TARGZ_FILE
# - FileType.PICKLE_FILE
# - FileType.JSON_FILE
# - FileType.CSV_FILE
# - FileType.ASCII_FILE
imf_df[FileType.PNG_FILE] = os.path.join(
    "/data", "io", "datafiles", "plot_image_IMF.png"
)
imf_df[FileType.JPEG_FILE] = os.path.join(
    "/data", "io", "datafiles", "plot_with_legend.jpg"
)
imf_df[FileType.FITS_FILE] = os.path.join(
    "/data", "io", "datafiles", "cassiopea_A_0.5-1.5keV.fits"
)
imf_df[FileType.TARGZ_FILE] = os.path.join("/data", "io", "datafiles", "archive.tar.gz")
imf_df[FileType.JSON_FILE] = os.path.join(
    "/data", "io", "datafiles", "test_header_249.json"
)
imf_df[FileType.ASCII_FILE] = os.path.join("/data", "io", "datafiles", "abstract.txt")
imf_df[FileType.HDF5_FILE] = os.path.join("/data", "io", "HDF5", "study.h5")
imf_df[FileType.PICKLE_FILE] = os.path.join(
    "/data", "io", "datafiles", "dict_saved.pkl"
)

# Datafile plot information (for plot future updates and online interactive visualisation on Galactica web pages).
# Available plot types are :
# - LINE_PLOT
# - SCATTER_PLOT
# - HISTOGRAM
# - HISTOGRAM_2D
# - IMAGE
# - MAP_2D
imf_df.plot_info = PlotInfo(
    plot_type=PlotType.LINE_PLOT,
    xaxis_values=N.array([10.0, 20.0, 30.0, 40.0, 50.0]),
    yaxis_values=N.array([1.256, 2.456, 3.921, 4.327, 5.159]),
    xaxis_log_scale=False,
    yaxis_log_scale=False,
    xaxis_label="Mass",
    yaxis_label="Probability",
    xaxis_unit=U.Msun,
    plot_title="Initial mass function",
    yaxis_unit=U.Mpc,
)
# -------------------------------------------------------------------------------------------------------------------- #


# Save study in HDF5 file
study = SimulationStudy(project=proj)
study.save_HDF5("./frig_study.h5")

# Eventually reload it from HDF5 file to edit its content
# study = SimulationStudy.load_HDF5("./frig_study.h5")