# Copyright (c) 2026, Artelys (https://www.artelys.com)
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
# SPDX-License-Identifier: MPL-2.0
#
#
# @author Léo Rossignol <leo.rossignol at artelys.com>
# @author Alexis Godefroy <alexis.godefroy at artelys.com>
#
import pathlib
import unittest
import pytest
import shutil

from emerge_res_integration_tool import config, config_test, constants

import numpy as np
import pandas as pd
import warnings

from metrix_launcher.core.metrix_simulation import MetrixSimulation

from emerge_res_integration_tool.optimization_engine import OptimizationEngine
from emerge_res_integration_tool.network_model import NetworkModel
from emerge_res_integration_tool.timeseries_manager import TimeSeriesManager


warnings.filterwarnings("ignore", category=pd.errors.PerformanceWarning)


class TestKnitroMetrixRuns(unittest.TestCase):

    @pytest.mark.skip(reason="OK but can be long to run")
    def test_1_simple_knitro_metrix_run(self):
        """Run a simple optim of battery dispatch plan based on 3A use case from powsybl-tutorials.
        The optimization is performed only on 6 timesteps"""

        # Set the working directory for Knitro x Metrix optimization process
        test_dir = (
            config_test.TEST_DATA_DIR / constants.OPTIM_TEST_FOLDER
        )

        # Create NetworkModel instance
        network_model = NetworkModel(
            test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_NETWORK_FILE
        )
        config.logger.info("NetworkModel instance created")

        # Re-write initial timeseries if needed (in case the test is re-run,
        # to avoid starting from the solution)
        timeseries_path = test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_TIMESERIES
        initial_timeseries_filename = "initial_" + str(pathlib.Path(timeseries_path).name)
        initial_timeseries_path = pathlib.Path(timeseries_path).parent / initial_timeseries_filename
        if initial_timeseries_path.exists():
            shutil.copyfile(initial_timeseries_path, timeseries_path)
        
        # Create TimeSeriesManager instance
        timeseries_manager = TimeSeriesManager(
            timeseries_path=timeseries_path
        )

        # Create Metrix simulation instance
        metrix_simulator = MetrixSimulation(folder=test_dir, variant_count=str(timeseries_manager.horizon))

        # Create and execute the Knitro x Metrix optimization process
        config.logger.info("Starting optimization using OptimizationEngine")
        engine = OptimizationEngine(network_model, timeseries_manager, metrix_simulator)

        results_dict = engine.run_with_knitro()

        # Check if it was successful: results_dict contains correct results
        config.logger.info(f"Results (setpoints): {results_dict['setpoints']}")
        self.assertTrue(results_dict["setpoints"].get("BAT") is not None)
        self.assertTrue(len(results_dict["setpoints"]["BAT"]) == 6)

    @pytest.mark.skip(reason="OK but can be long to run")
    def test_1b_end_to_end_baseline_plus_battery_solution_csv(self):
        """
        Run baseline Metrix, add one battery on a chosen substation, run Knitro,
        and check that the exported final solution CSV has the expected columns.
        """
        test_dir = (
            config_test.TEST_DATA_DIR / constants.OPTIM_TEST_FOLDER
        )

        timeseries_path = test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_TIMESERIES
        initial_timeseries_filename = "initial_" + str(pathlib.Path(timeseries_path).name)
        initial_timeseries_path = pathlib.Path(timeseries_path).parent / initial_timeseries_filename
        if initial_timeseries_path.exists():
            shutil.copyfile(initial_timeseries_path, timeseries_path)

        engine = OptimizationEngine(
            folder_path=test_dir,
            substation_ids=["NO"],
        )

        results_dict = engine.run_with_batteries_and_knitro()

        self.assertTrue("run_metrix_normal" in results_dict)
        self.assertTrue(results_dict["setpoints"].get("BAT") is not None)

        exported_path = (
            test_dir / constants.OUTPUT_DATA_FOLDER / constants.BEST_RESULTS_PATH
        )
        self.assertTrue(exported_path.exists())

        df = pd.read_csv(exported_path, sep=";")
        expected_cols = [
            "Timestep",
            "BAT_E0 (MWh)",
            "BAT_profile (MW)",
            "BAT_SOC",
            constants.OUTPUT_CUM_COST,
            constants.OUTPUT_ADEQUACY_COST,
            constants.OUTPUT_REDISP_COST,
            constants.OUTPUT_CUM_SAVED_COST,
            constants.OUTPUT_CUM_CURTAILMENT,
        ]
        self.assertEqual(expected_cols, list(df.columns))
    
    @pytest.mark.skip(reason="OK but can be long to run")    
    def test_1c_several_batteries(self):
        test_dir = (
            config_test.TEST_DATA_DIR / constants.OPTIM_TEST_FOLDER
        )

        timeseries_path = test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_TIMESERIES
        initial_timeseries_filename = "initial_" + str(pathlib.Path(timeseries_path).name)
        initial_timeseries_path = pathlib.Path(timeseries_path).parent / initial_timeseries_filename
        if initial_timeseries_path.exists():
            shutil.copyfile(initial_timeseries_path, timeseries_path)

        engine = OptimizationEngine(
            folder_path=test_dir,
            substation_ids=["NO", "SE"],
        )

        results_dict = engine.run_with_batteries_and_knitro()

        self.assertTrue("run_metrix_normal" in results_dict)

        exported_path = (
            test_dir / constants.OUTPUT_DATA_FOLDER / constants.BEST_RESULTS_PATH
        )
        self.assertTrue(exported_path.exists())
        
    def test_1d_knitro_metrix_run_mozambique(self):
        """
        Run optimization on Mozambique network.
        """
        test_dir = (
            config_test.TEST_DATA_DIR / constants.OPTIM_MOZAMBIQUE_FOLDER
        )

        timeseries_path = test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_TIMESERIES
        initial_timeseries_filename = "initial_" + str(pathlib.Path(timeseries_path).name)
        initial_timeseries_path = pathlib.Path(timeseries_path).parent / initial_timeseries_filename
        if initial_timeseries_path.exists():
            shutil.copyfile(initial_timeseries_path, timeseries_path)

        engine = OptimizationEngine(
            folder_path=test_dir,
            substation_ids=["S15"],
        )

        results_dict = engine.run_with_batteries_and_knitro()

        self.assertTrue("run_metrix_normal" in results_dict)
        self.assertTrue(results_dict["setpoints"].get("BAT") is not None)

        exported_path = (
            test_dir / constants.OUTPUT_DATA_FOLDER / constants.BEST_RESULTS_PATH
        )
        self.assertTrue(exported_path.exists())

        df = pd.read_csv(exported_path, sep=";")
        expected_cols = [
            "Timestep",
            "BAT_E0 (MWh)",
            "BAT_profile (MW)",
            "BAT_SOC",
            constants.OUTPUT_CUM_COST,
            constants.OUTPUT_ADEQUACY_COST,
            constants.OUTPUT_REDISP_COST,
            constants.OUTPUT_CUM_SAVED_COST,
            constants.OUTPUT_CUM_CURTAILMENT,
        ]
        self.assertEqual(expected_cols, list(df.columns))

    @pytest.mark.skip(reason="OK but can be long to run")
    def test_2_sensitivity_analysis_knitro_stepsize_param(self):
        """Run a sensitivity analysis on Knitro parameter FINDIFF_RELSTEPSIZE"""

        # Set the working directory for Knitro x Metrix optimization process
        test_dir = (
            config_test.TEST_DATA_DIR
            / constants.SENSITIVITY_TEST_FOLDER
        )

        # Create NetworkModel instance
        network_model = NetworkModel(
            test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_NETWORK_FILE
        )

        # Re-write initial timeseries if needed (in case the test is re-run,
        # to avoid starting from the solution)
        timeseries_path = test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_TIMESERIES
        initial_timeseries_filename = "initial_" + str(pathlib.Path(timeseries_path).name)
        initial_timeseries_path = pathlib.Path(timeseries_path).parent / initial_timeseries_filename
        if initial_timeseries_path.exists():
            shutil.copyfile(initial_timeseries_path, timeseries_path)
        
        # Create TimeSeriesManager instance
        timeseries_manager = TimeSeriesManager(
            timeseries_path=timeseries_path
        )

        # Create Metrix simulation instance
        metrix_simulator = MetrixSimulation(folder=test_dir, variant_count=str(timeseries_manager.horizon))

        # Create and execute the Knitro x Metrix optimization process
        config.logger.info("Starting optimization using OptimizationEngine")
        engine = OptimizationEngine(network_model, timeseries_manager, metrix_simulator)

        config.logger.info("Running FINDIFF_RELSTEPSIZE sweep")
        engine.run_findiff_sweep([1e-2, 1e-5])
        config.logger.info("FINDIFF_RELSTEPSIZE sweep done")
        
    @pytest.mark.skip(reason="OK but can be long to run")
    def test_3_sensitivity_analysis_battery_params(self):
        """Run a sensitivity analysis on battery parameters (Emax, Pmax)"""

        # Set the working directory for Knitro x Metrix optimization process
        test_dir = (
            config_test.TEST_DATA_DIR
            / constants.SENSITIVITY_TEST_FOLDER
        )

        # Create NetworkModel instance
        network_model = NetworkModel(
            test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_NETWORK_FILE
        )

        # Re-write initial timeseries if needed (in case the test is re-run,
        # to avoid starting from the solution)
        timeseries_path = test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_TIMESERIES
        initial_timeseries_filename = "initial_" + str(pathlib.Path(timeseries_path).name)
        initial_timeseries_path = pathlib.Path(timeseries_path).parent / initial_timeseries_filename
        if initial_timeseries_path.exists():
            shutil.copyfile(initial_timeseries_path, timeseries_path)
        
        # Create TimeSeriesManager instance
        timeseries_manager = TimeSeriesManager(
            timeseries_path=timeseries_path
        )

        # Create Metrix simulation instance
        metrix_simulator = MetrixSimulation(folder=test_dir, variant_count=str(timeseries_manager.horizon))

        # Create and execute the Knitro x Metrix optimization process
        config.logger.info("Starting optimization using OptimizationEngine")
        engine = OptimizationEngine(network_model, timeseries_manager, metrix_simulator)

        # Define battery parameter combinations to test
        param_combinations = [
            {"Emax": 500, "Pmax": 50},
            {"Emax": 200, "Pmax": 100},
        ]

        config.logger.info("Running battery params sweep")
        engine.run_battery_params_sweep(param_combinations)
        config.logger.info("Battery params sweep done")
        
    @pytest.mark.skip(reason="OK but can be long to run")    
    def test_4_sensitivity_analysis_initial_guess(self):
        """Run a sensitivity analysis on initial guess (to consolidate confidence in the results)"""

        # Set the working directory for Knitro x Metrix optimization process
        test_dir = (
            config_test.TEST_DATA_DIR
            / constants.SENSITIVITY_TEST_FOLDER
        )

        # Create NetworkModel instance
        network_model = NetworkModel(
            test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_NETWORK_FILE
        )

        # Re-write initial timeseries if needed (in case the test is re-run,
        # to avoid starting from the solution)
        timeseries_path = test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_TIMESERIES
        initial_timeseries_filename = "initial_" + str(pathlib.Path(timeseries_path).name)
        initial_timeseries_path = pathlib.Path(timeseries_path).parent / initial_timeseries_filename
        if initial_timeseries_path.exists():
            shutil.copyfile(initial_timeseries_path, timeseries_path)
        
        # Create TimeSeriesManager instance
        timeseries_manager = TimeSeriesManager(
            timeseries_path=timeseries_path
        )

        # Create Metrix simulation instance
        metrix_simulator = MetrixSimulation(
            folder=test_dir, 
            variant_count=str(timeseries_manager.horizon)
        )

        # Create OptimizationEngine
        config.logger.info("Starting optimization using OptimizationEngine")
        engine = OptimizationEngine(network_model, timeseries_manager, metrix_simulator)
        
        # Get battery parameters for computing actual values
        # Assuming all batteries have same Pmax/Emax (adjust if needed)
        pmax = engine.battery_fleet.batteries[0].pmax
        emax = engine.battery_fleet.batteries[0].emax

        # Define different initial guess configurations to test
        initial_guess_configs = [
            {
                "name": "heuristic_cold_start",
                "use_custom_guess": False,
                "custom_setpoints_mw": None,
                "custom_e0_mwh": None,
            },
            {
                "name": "warm_start_near_optimal",
                "use_custom_guess": True,
                # Example: quasi-optimal solution from a previous run
                "custom_setpoints_mw": [
                    pmax, pmax, -pmax, -pmax, pmax, pmax
                ],
                "custom_e0_mwh": emax,  # Start at 60% SOC (if Emax=500)
            },
            {
                "name": "aggressive_discharge",
                "use_custom_guess": True,
                "custom_setpoints_mw": [
                    pmax, pmax, pmax, pmax, pmax, pmax
                ],
                "custom_e0_mwh": emax,  # Start fully charged
            },
            {
                "name": "aggressive_charge",
                "use_custom_guess": True,
                "custom_setpoints_mw": [
                    -pmax, -pmax, -pmax, -pmax, -pmax, -pmax
                ],
                "custom_e0_mwh": 0,  # Start fully discharged
            },
            {
                "name": "nothing_planned",
                "use_custom_guess": True,
                "custom_setpoints_mw": [
                    0, 0, 0, 0, 0, 0
                ],
                "custom_e0_mwh": 0.5 * emax,  # Start at 50% SOC
            },
            {
                "name": "random_perturbation",
                "use_custom_guess": True,
                # Random profile for comparison
                "custom_setpoints_mw": [
                    0.3*pmax, -0.8*pmax, 0.5*pmax, -0.2*pmax, -0.1*pmax, 0.7*pmax
                ],
                "custom_e0_mwh": 0.4*emax,
            },
        ]

        config.logger.info(f"Running initial guess sweep with {len(initial_guess_configs)} configurations")
        engine.run_initial_guess_sweep(initial_guess_configs)
        config.logger.info("Initial guess sweep done")
        
    @pytest.mark.skip(reason="OK but can be long to run")
    def test_5_battery_size_parameter(self):
        """
        Verify that the battery_size parameter is correctly propagated to the
        BatteryFleet and used during optimisation.
        """
        test_dir = (
            config_test.TEST_DATA_DIR / constants.OPTIM_MOZAMBIQUE_FOLDER
        )

        timeseries_path = test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_TIMESERIES
        initial_timeseries_path = (
            timeseries_path.parent / ("initial_" + timeseries_path.name)
        )
        if initial_timeseries_path.exists():
            shutil.copyfile(initial_timeseries_path, timeseries_path)

        custom_emax = 300.0
        custom_pmax = 75.0

        engine = OptimizationEngine(
            folder_path=test_dir,
            substation_ids=["S15"],
            battery_size=(custom_emax, custom_pmax),
        )

        results_dict = engine.run_with_batteries_and_knitro()

        # Battery fleet should reflect the custom size
        self.assertEqual(len(engine.battery_fleet.batteries), 1)
        bat = engine.battery_fleet.batteries[0]
        self.assertAlmostEqual(bat.emax, custom_emax)
        self.assertAlmostEqual(bat.pmax, custom_pmax)

        # Optimisation should still produce valid setpoints
        self.assertIn("BAT", results_dict["setpoints"])
        profile = results_dict["setpoints"]["BAT"]
        self.assertEqual(len(profile), engine.knitro_optimizer.horizon)
        