# 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 gzip
import pathlib
import shutil

import numpy as np
import pandas as pd

from emerge_res_integration_tool import config
from emerge_res_integration_tool import constants
from emerge_res_integration_tool import utils

from metrix_launcher.core.metrix_simulation import MetrixSimulation

from emerge_res_integration_tool.network_model import NetworkModel, BatteryFleet
from emerge_res_integration_tool.timeseries_manager import TimeSeriesManager
from emerge_res_integration_tool.knitro_optimizer import KnitroBatteryOptimizer


class OptimizationEngine:
    """
    High-level orchestration of the iterative optimization, tying together
    the network model, batteries, time series and the Metrix simulator.
    """

    def __init__(
        self,
        folder_path,
        substation_ids,
        work_dir_name: str = constants.WORK_DIR_NAME,
        battery_costs: tuple[float, float] | None = None,  # (redisp_down, redisp_up)
        battery_size: tuple[float, float] | None = None,   # (emax_mwh, pmax_mw)
    ):
        """
        User-facing entrypoint.

        Args:
            folder_path: Base folder containing an `input/` directory compatible with
                Metrix naming conventions.
            substation_ids: List of substation IDs where batteries will be added.
            work_dir_name: Name of the working directory created under folder_path.
                Inputs are copied there so original user inputs are never modified.
            force_recreate_workdir: If True, delete and recreate the working directory.
        """
        self.folder_path = pathlib.Path(folder_path)
        self.substation_ids = list(substation_ids)

        self.work_dir = self.folder_path / work_dir_name
        self.work_input_dir = self.work_dir / constants.INPUT_DATA_FOLDER
        self.work_output_dir = self.folder_path / constants.OUTPUT_DATA_FOLDER
        self.work_output_dir.mkdir(parents=True, exist_ok=True)

        self._prepare_working_directory()

        # Core objects are built from the working directory
        self.network_model, self.timeseries_manager, self.metrix_simulator = (
            self._build_components_from_workdir()
        )

        # Network
        self.network = self.network_model.load()

        # Asset identifiers
        self.battery_ids = self.network_model.get_battery_ids()
        self.hydro_ids = self.network_model.get_hydro_ids()
        self.thermal_ids = self.network_model.get_thermal_ids()
        self.solar_ids = self.network_model.get_solar_ids()
        self.wind_ids = self.network_model.get_wind_ids()

        # Batteries
        self.battery_size = battery_size  # (Emax_MWh, Pmax_MW) or None
        
        if self.battery_ids:
            emax, pmax = battery_size if battery_size is not None else (constants.BAT_EMAX, constants.BAT_PMAX)
            self.battery_fleet = BatteryFleet.from_ids_with_defaults(
                self.battery_ids,
                default_emax=emax,
                default_pmax=pmax,
            )
            if battery_costs is not None:
                self.battery_fleet.redispatch_costs = battery_costs
        else:
            self.battery_fleet = BatteryFleet([])

        # Knitro optimiser for battery setpoints (lazy initialization)
        self.knitro_optimizer = None

    def _ensure_knitro_optimizer(self):
        """
        Lazily build the Knitro optimizer once batteries are present on the network.
        """
        if not self.battery_fleet.batteries:
            raise ValueError("No batteries detected on the network for Knitro optimisation.")
        if self.knitro_optimizer is None:
            self.knitro_optimizer = KnitroBatteryOptimizer(
                network_model=self.network_model,
                metrix_simulator=self.metrix_simulator,
                timeseries_manager=self.timeseries_manager,
                battery_fleet=self.battery_fleet,
                output_dir=self.work_output_dir,
            )

    def _prepare_working_directory(self):
        """Delegate preprocessing directory setup to `utils.py`."""
        utils.prepare_working_directory(self.folder_path, self.work_dir)

    def _build_components_from_workdir(self):
        """Delegate component building to `utils.py`."""
        return utils.build_components_from_workdir(self.work_dir)

    def _ensure_metrix_inputs(self, metrix_simulator: MetrixSimulation):
        """Delegate Metrix input bootstrap to `utils.py`."""
        utils.ensure_metrix_inputs(metrix_simulator)

    def _ensure_battery_timeseries_columns(self):
        """Delegate battery timeseries bootstrap to `utils.py`."""
        utils.ensure_battery_timeseries_columns(
            work_input_dir=self.work_input_dir,
            battery_fleet=self.battery_fleet,
        )
        self.timeseries_manager.reload_initial()

    def run_metrix_on_current_network(self) -> pd.DataFrame:
        """
        Run Metrix once on the current network/timeseries and return the CSV results.
        """
        config.logger.info("Starting baseline PowSyBl-Metrix run (no optimisation)")
        return_code, std_err = self.metrix_simulator.run_metrix()
        config.logger.info(f"Baseline Metrix returned with exit code {return_code}")
        config.logger.debug(std_err)

        results_path = pathlib.Path(self.metrix_simulator.csv_results_file_path)
        gz_path = pathlib.Path(str(results_path) + ".gz")

        if not results_path.exists():
            # If only the gzipped file exists, extract it the same way as cost helpers
            if gz_path.exists():
                config.logger.info("Extracting gzipped Metrix results for baseline run")
                with gzip.open(gz_path, "rb") as f_in, open(results_path, "wb") as f_out:
                    shutil.copyfileobj(f_in, f_out)
            else:
                raise FileNotFoundError(
                    f"Metrix results file not found at {results_path} (or {gz_path})"
                )

        run_metrix_normal = pd.read_csv(results_path, sep=";")
        return run_metrix_normal
    
    def run_with_knitro(self, enable_multistart: bool = False):
        """
        Solve the battery setpoints with Knitro in a single optimisation pass.
        """
        self._ensure_knitro_optimizer()
        config.logger.info("Launching Knitro-based optimisation")
        self.timeseries_manager.reset_for_iteration(0)

        solution = self.knitro_optimizer.solve(enable_multistart=enable_multistart)
        setpoints = solution["setpoints"]

        # Persist the chosen dispatch for record-keeping
        self.timeseries_manager.apply_battery_setpoints(setpoints)
        self.timeseries_manager.save_iteration_inputs(self.network_model)

        return {"solution": solution, "setpoints": setpoints,}

    def run_with_batteries_and_knitro(
        self,
        substation_ids=None,
        enable_multistart: bool = False,
        battery_size: tuple[float, float] | None = None,
    ):
        """
        End-to-end workflow:
          1. Save a copy of the original network.xiidm.
          2. Run Metrix once on the original network (no batteries).
          3. Add one battery per requested substation and overwrite network.xiidm.
          4. Rebuild battery-related structures and run Knitro optimisation.
          5. Restore the original network.xiidm at the end.
        """
        if substation_ids is None:
            substation_ids = self.substation_ids

        if not substation_ids:
            raise ValueError("run_with_batteries_and_knitro requires at least one substation ID")

        original_path = pathlib.Path(self.network_model.network_filepath)
        saved_path = original_path.with_name(constants.SAVED_NETWORK_FILE)

        # Save original network once
        if not saved_path.exists():
            shutil.copy2(original_path, saved_path)
            config.logger.info(f"Saved original network to {saved_path}")

        try:
            # 1) Baseline Metrix run on original network
            run_metrix_normal = self.run_metrix_on_current_network()

            # 2) Add batteries on a duplicate XIIDM network file
            config.logger.info(
                f"Adding batteries on substations: {list(substation_ids)}"
            )
            network_with_batteries_path = self.network_model.add_batteries_to_network(
                substation_ids=substation_ids
            )
            config.logger.info(
                f"Network with batteries saved to {network_with_batteries_path}"
            )

            # Overwrite network.xiidm used by Metrix with the battery version
            shutil.copy2(network_with_batteries_path, original_path)
            self.network_model.network_filepath = str(original_path)

            # Reload network and refresh battery fleet on the modified XIIDM
            self.network = self.network_model.load()
            self.battery_ids = self.network_model.get_battery_ids()
            effective_size = battery_size if battery_size is not None else self.battery_size
            emax, pmax = effective_size if effective_size is not None else (constants.BAT_EMAX, constants.BAT_PMAX)
            self.battery_fleet = BatteryFleet.from_ids_with_defaults(
                self.battery_ids,
                default_emax=emax,
                default_pmax=pmax,
            )
            self._ensure_battery_timeseries_columns()
            self.knitro_optimizer = None
            self._ensure_knitro_optimizer()
            self.knitro_optimizer.set_baseline_run(run_metrix_normal)

            # 3) Optimisation of battery profiles with Knitro on the modified network
            optimisation_results = self.run_with_knitro(
                enable_multistart=enable_multistart
            )

            return {
                "run_metrix_normal": run_metrix_normal,
                **optimisation_results,
            }
        finally:
            # 4) Restore original network.xiidm
            if saved_path.exists():
                shutil.copy2(saved_path, original_path)
                self.network_model.network_filepath = str(original_path)
                config.logger.info(f"Restored original network from {saved_path}")

    def run_findiff_sweep(self, stepsize_values=None):
        """Run multiple Knitro optimizations with different FINDIFF_RELSTEPSIZE
        values and consolidate results."""
        if stepsize_values is None:
            stepsize_values = constants.STEPSIZE_VALUES

        config.logger.info(
            f"Starting FINDIFF_RELSTEPSIZE sweep with {len(stepsize_values)} values"
        )

        all_rows = []

        for stepsize in stepsize_values:
            config.logger.info(
                f"Running optimization with FINDIFF_RELSTEPSIZE = {stepsize}"
            )

            # Reset timeseries
            self.timeseries_manager.reset_for_iteration(0)

            # Ensure optimizer is ready for this configuration
            self._ensure_knitro_optimizer()

            # Run optimization
            solution = self.knitro_optimizer.solve(
                enable_multistart=False, findiff_relstepsize=stepsize
            )

            for t in range(self.knitro_optimizer.horizon):
                row = {
                    "stepsize": stepsize,
                    "timestep": t,
                    "total_cost": solution["objective"],
                    "adequacy_cost": solution["adequacy_costs"][t],
                    "redispatching_cost": solution["redispatching_costs"][t],
                }
                
                # E0 values
                for bat_id, e0_val in solution["e0_values"].items():
                    row[f"E0_{bat_id}"] = e0_val
                
                # Setpoints
                for bat_id, profile in solution["setpoints"].items():
                    row[f"setpoint_{bat_id}"] = profile[t]
                
                all_rows.append(row)

        # DataFrame final
        combined_df = pd.DataFrame(all_rows)
        output_dir = self.work_output_dir
        
        combined_df.to_csv(
            output_dir / constants.SENSITIVITY_FINDIFF_FILE,
            sep=";", 
            index=False
        )
        
        config.logger.info("Saved combined results")
        return combined_df
    
    def run_battery_params_sweep(self, param_combinations=None):
        """
        Run multiple Knitro optimizations with different battery Emax/Pmax configurations
        """
        if param_combinations is None:
            param_combinations = [
                {"Emax": 500, "Pmax": 50},
                {"Emax": 200, "Pmax": 100},
                {"Emax": 50, "Pmax": 100},
            ]
        
        config.logger.info(f"Starting battery params sweep with {len(param_combinations)} configurations")
        
        all_rows = []
        
        for params in param_combinations:
            emax = params["Emax"]
            pmax = params["Pmax"]
            
            config.logger.info(f"Running optimization: Emax={emax} MWh, Pmax={pmax} MW")
            
            # Recreate batteries with new parameters
            self.battery_fleet = BatteryFleet.from_ids_with_defaults(
                self.battery_ids, 
                default_emax=emax, 
                default_pmax=pmax
            )
            
            self.knitro_optimizer = KnitroBatteryOptimizer(
                network_model=self.network_model,
                metrix_simulator=self.metrix_simulator,
                timeseries_manager=self.timeseries_manager,
                battery_fleet=self.battery_fleet,
                output_dir=self.work_output_dir,
            )
            
            # Reset
            self.timeseries_manager.reset_for_iteration(0)
            
            # Run optimization
            solution = self.knitro_optimizer.solve(
                enable_multistart=False,
                findiff_relstepsize=1e-2
            )
            
            # Write output file
            for t in range(self.knitro_optimizer.horizon):
                row = {
                    "Emax": emax,
                    "Pmax": pmax,
                    "timestep": t,
                    "total_cost": solution["objective"],
                    "adequacy_cost": solution["adequacy_costs"][t],
                    "redispatching_cost": solution["redispatching_costs"][t],
                }
                
                # E0 values
                for bat_id, e0_val in solution["e0_values"].items():
                    row[f"E0_{bat_id}"] = e0_val
                
                # Setpoints
                for bat_id, profile in solution["setpoints"].items():
                    row[f"setpoint_{bat_id}"] = profile[t]
                
                all_rows.append(row)
        
        # Un seul DataFrame
        combined_df = pd.DataFrame(all_rows)
        output_dir = self.work_output_dir
        
        combined_df.to_csv(
            output_dir / constants.SENSITIVITY_PARAMS_FILE,
            sep=";", 
            index=False
        )
        
        config.logger.info("Saved combined results")
        return combined_df
    
    def run_initial_guess_sweep(self, initial_guess_configs):
        """
        Run multiple Knitro optimizations with different initial guesses
        """
        config.logger.info(f"Starting initial guess sweep with {len(initial_guess_configs)} configurations")
        
        all_rows = []
        
        for guess_config in initial_guess_configs:
            name = guess_config["name"]
            
            config.logger.info(f"Running optimization with initial guess: {name}")
            
            # Reset
            self.timeseries_manager.reset_for_iteration(0)
            
            # Run optimization avec les paramètres custom
            self._ensure_knitro_optimizer()
            solution = self.knitro_optimizer.solve(
                enable_multistart=False,
                findiff_relstepsize=1e-2,
                use_custom_guess=guess_config["use_custom_guess"],
                custom_setpoints_mw=guess_config.get("custom_setpoints_mw"),
                custom_e0_mwh=guess_config.get("custom_e0_mwh"),
            )
            
            # Pour chaque timestep
            for t in range(self.knitro_optimizer.horizon):
                row = {
                    "initial_guess_name": name,
                    "timestep": t,
                    "total_cost": solution["objective"],
                    "adequacy_cost": solution["adequacy_costs"][t],
                    "redispatching_cost": solution["redispatching_costs"][t],
                }
                
                # E0 values
                for bat_id, e0_val in solution["e0_values"].items():
                    row[f"E0_{bat_id}"] = e0_val
                
                # Setpoints
                for bat_id, profile in solution["setpoints"].items():
                    row[f"setpoint_{bat_id}"] = profile[t]
                
                all_rows.append(row)
        
        # DataFrame final
        combined_df = pd.DataFrame(all_rows)
        output_dir = self.work_output_dir
        
        combined_df.to_csv(
            output_dir / constants.SENSITIVITY_INITIAL_GUESS_FILE,
            sep=";", 
            index=False
        )
        
        config.logger.info("Saved combined results")
        return combined_df