# Copyright 2021 IRT Saint Exupéry, https://www.irt-saintexupery.com
#
# This work is licensed under a BSD 0-Clause License.
#
# Permission to use, copy, modify, and/or distribute this software
# for any purpose with or without fee is hereby granted.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
# WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL
# THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT,
# OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING
# FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
# NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
# WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

"""
Usage of the model calibration based on curve outputs
======================================================

Calibrate a model based on curve outputs.
"""

# %%
from __future__ import annotations

import logging

from gemseo.algos.opt.nlopt.settings.nlopt_cobyla_settings import NLOPT_COBYLA_Settings
from gemseo_calibration.calibrator import CalibrationMetricSettings
from gemseo_calibration.measures.integrated_measure import CurveScaling
from numpy import atleast_1d

from vimseo import EXAMPLE_RUNS_DIR
from vimseo.api import activate_logger
from vimseo.api import create_model
from vimseo.core.model_settings import IntegratedModelSettings
from vimseo.io.space_io import SpaceToolFileIO
from vimseo.problems.mock.mock_curves.mock_curves import MockCurves
from vimseo.storage_management.base_storage_manager import PersistencyPolicy
from vimseo.tools.calibration.calibration_step import CalibrationStep
from vimseo.tools.calibration.calibration_step import CalibrationStepInputs
from vimseo.tools.calibration.calibration_step import CalibrationStepSettings
from vimseo.tools.calibration.input_data import CALIBRATION_INPUT_DATA
from vimseo.utilities.generate_validation_reference import (
    generate_reference_from_parameter_space,
)

# %%
# We first define the logger level:
activate_logger(level=logging.INFO)

# %%
# We want to calibrate an analytical model that takes inputs :math:`x, x_1` and
# returns the curve :math:`(y_{axis}, y)` where :math:`\mathbf{y_{axis}} \in [0,1]` and
# :math:`\mathbf{y}=x \times \mathbf{y_{axis}} + x_1`.
# The objective is to find the best $x$ such that the simulated
# and reference $y$ match.

# %%
# Then, we need to create reference data.
# They are generated from the model to calibrate, which is biased by imposing a
# modified $x$.
# Several samples are generated by varying $x_1$.
X_TARGET = 1.5
MockCurves.CURVE_NB_POINTS = 10
model_name = "MockCurves"
load_case = "Dummy"
reference_mock_curves = create_model(
    model_name,
    load_case,
    model_options=IntegratedModelSettings(
        directory_archive_persistency=PersistencyPolicy.DELETE_ALWAYS,
        directory_scratch_persistency=PersistencyPolicy.DELETE_ALWAYS,
    ),
)
reference_mock_curves.default_input_data["x"] = atleast_1d(X_TARGET)
reference_mock_curves.cache = None
reference_data = generate_reference_from_parameter_space(
    reference_mock_curves,
    SpaceToolFileIO()
    .read(CALIBRATION_INPUT_DATA / "experimental_space_mock_curves.json")
    .parameter_space,
    n_samples=2,
    as_dataset=True,
)

# %%
# We now define the model used for the calibration:
model = create_model(
    model_name,
    load_case,
    model_options=IntegratedModelSettings(
        directory_archive_root=EXAMPLE_RUNS_DIR / "archive/calibration_curves",
        directory_scratch_root=EXAMPLE_RUNS_DIR / "scratch/calibration_curves",
        cache_file_path=EXAMPLE_RUNS_DIR
        / f"caches/calibration_curves/{model_name}_{load_case}_cache.hdf",
    ),
)

# %%
# Then, a step of calibration is defined.
# It uses the :class:``SBPISE`` metric, which computes the area
# between the reference and the simulated curves.
# The curves are scaled to zero-mean and unitary standard deviation
# by setting the argument ``scaling`` to ``CurveScaling.XYRange``,
# which is mandatory when calibrating for:
#  - several metrics in the same step
#  - several load cases in the same step
output_name = "y"
step = CalibrationStep(working_directory="curves")
step.execute(
    inputs=CalibrationStepInputs(
        reference_data={
            "Dummy": reference_data,
        },
    ),
    settings=CalibrationStepSettings(
        name_to_models={"Dummy": model},
        control_outputs={
            output_name: CalibrationMetricSettings(
                measure="SBPISE",
                mesh="y_axis",
                scaling=CurveScaling.XYRange,
            ).model_dump()
        },
        input_names=[
            "x_1",
        ],
        parameter_names=["x"],
        optimizer_settings=NLOPT_COBYLA_Settings(max_iter=50),
    ),
)
step.save_results()

# %%
# We can show the prior parameters, i.e. the optimizer starting point:
step.result.prior_parameters

# %%
# The outputs can be compared to the reference data, before and after calibration:
figs = step.plot_results(step.result, show=True, save=False)
figs["Dummy"]["simulated_versus_reference_curve_y_versus_y_axis"]

# %%
# The curves that have been defined as ``control_outputs`` can be retrieved as
# Pandas DataFrame:
step.result.curve_data