# 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 json
import shutil
import urllib3

import pandas as pd
import requests

from emerge_res_integration_tool import config, config_test, constants, utils
from emerge_res_integration_tool.optimization_engine import OptimizationEngine

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


class QuickStartIntegrationEmerge:

    @classmethod
    def setUpClass(cls):
        """Call the CIRCE API once for the whole test class and store the response."""
        payload = {
            "battery_power": constants.BATTERY_POWER,
            "scenario_id": constants.SCENARIO_ID,
        }
        headers = {
            "accept": "application/json",
            "Content-Type": "application/json",
        }

        config.logger.info(f"Sending request to {constants.URL} ...")
        config.logger.info(f"  battery_power : {len(constants.BATTERY_POWER)} values")
        config.logger.info(f"  scenario_id   : {constants.SCENARIO_ID}")

        response = requests.post(
            constants.URL,
            headers=headers,
            data=json.dumps(payload),
            verify=False,
        )
        response.raise_for_status()

        cls.data = response.json()
        cls.optimal_buses = cls.data.get("optimal_buses", [])

        config.logger.info(f"API response status : {cls.data.get('status')}")
        config.logger.info(f"Optimal buses       : {cls.optimal_buses}")

    def test_iterative_emerge_loop(self):
        """
        Iterative EMERGE loop between CIRCE and Artelys tools until convergence or max iterations.

        At each iteration:
        1. Call OSP tool with the current battery profile
        2. Place a battery on the first substation returned
        3. Run Artelys-EMERGE optimization to get a new battery profile
        4. Repeat until the first substation returned by OSP tool matches the one
            used to generate the input battery profile, or until 10 iterations.

        Results are exported as a CSV file with one column per iteration.
        """
        MAX_ITERATIONS = 10

        test_dir = config_test.TEST_DATA_DIR / constants.OPTIM_MOZAMBIQUE_FOLDER
        network_path = (
            test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_NETWORK_FILE
        )
        timeseries_path = (
            test_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_TIMESERIES
        )

        headers = {
            "accept": "application/json",
            "Content-Type": "application/json",
        }

        #  Iteration 0 – use the baseline battery profile from constants
        current_battery_power = constants.BATTERY_POWER  # baseline profile

        results_columns = {}  # will hold one entry per iteration
        iteration = 0
        previous_substation = None  # substation used to generate current profile

        while iteration < MAX_ITERATIONS:
            iter_label = f"Iteration {iteration}"
            config.logger.info(f"--- {iter_label} ---")

            # 1. Store the battery profile used at this iteration
            results_columns[f"Profil battery iteration {iteration}"] = (
                current_battery_power
            )

            # 2. Call CIRCE tool with the current battery profile
            payload = {
                "battery_power": current_battery_power,
                "scenario_id": constants.SCENARIO_ID,
            }
            config.logger.info(f"Calling OSP API (iteration {iteration}) ...")
            response = requests.post(
                constants.URL,
                headers=headers,
                data=json.dumps(payload),
                verify=False,
            )
            response.raise_for_status()
            circe_data = response.json()

            optimal_buses = circe_data.get("optimal_buses", [])

            config.logger.info(
                f"OSP tool optimal buses (iteration {iteration}): {optimal_buses}"
            )

            # 3. Store the substation list returned by CIRCE for this iteration
            results_columns[f"List of substations OSP tool - iteration {iteration}"] = (
                optimal_buses
            )

            # 4. Check convergence: the first substation CIRCE recommends must be
            #    the same one that was used to generate the current input profile.
            first_bus = optimal_buses[0]
            first_substation = utils.find_substation_id(network_path, first_bus)

            if (
                previous_substation is not None
                and first_substation == previous_substation
            ):
                config.logger.info(
                    f"Convergence reached at iteration {iteration}: "
                    f"first substation '{first_substation}' matches the one used "
                    f"to generate the current battery profile."
                )
                break

            # 5. Map bus names to substation IDs and run EMERGE optimisation
            substation_ids = []
            for bus in optimal_buses[: constants.NB_BATTERIES]:
                substation = utils.find_substation_id(network_path, bus)
                substation_ids.append(substation)

            # Restore the initial timeseries before each optimisation run
            initial_timeseries_path = timeseries_path.parent / (
                "initial_" + timeseries_path.name
            )
            if initial_timeseries_path.exists():
                shutil.copyfile(initial_timeseries_path, timeseries_path)

            config.logger.info(
                f"Running EMERGE optimisation with batteries on: {substation_ids}"
            )
            engine = OptimizationEngine(
                folder_path=test_dir,
                substation_ids=substation_ids,
                battery_size=(100, 10),
            )
            results_dict = engine.run_with_batteries_and_knitro()
            config.logger.info(f"EMERGE optimisation finished (iteration {iteration})")

            # 6. Extract the new battery profile from the optimisation results
            new_battery_power = results_dict["setpoints"]["BAT"]
            if hasattr(new_battery_power, "tolist"):
                new_battery_power = new_battery_power.tolist()

            # Prepare next iteration
            previous_substation = first_substation
            current_battery_power = new_battery_power
            iteration += 1

        else:
            config.logger.warning(
                f"Maximum iterations ({MAX_ITERATIONS}) reached without convergence."
            )

        #  Export results to CSV
        output_csv_path = (
            test_dir / constants.OUTPUT_DATA_FOLDER / constants.INTEGRATION_EMERGE_FILE
        )

        # Pad all columns to the same length before building the DataFrame
        max_len = max(len(v) for v in results_columns.values())
        padded = {
            col: list(values) + [None] * (max_len - len(list(values)))
            for col, values in results_columns.items()
        }

        df = pd.DataFrame(padded)
        df.to_csv(output_csv_path, index=False)
        config.logger.info(f"Results exported to {output_csv_path}")


if __name__ == "__main__":
    runner = QuickStartIntegrationEmerge()
    runner.test_iterative_emerge_loop()
