# 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 numpy as np
import pandas as pd
import pathlib
import shutil

from emerge_res_integration_tool import constants
from emerge_res_integration_tool.network_model import NetworkModel


class TimeSeriesManager:
    """
    Manages input and derived timeseries for the optimisation loop.
    """

    def __init__(
        self,
        timeseries_path: str,
    ):
        # Main timeseries.csv path (will be reused for all Metrix computations)
        self.timeseries_path = timeseries_path
        # Initial timeseries path (a file copy is done when creating this instance)
        initial_timeseries_filename = ("initial_" +
        str(pathlib.Path(self.timeseries_path).name))
        self.initial_timeseries_path = pathlib.Path(self.timeseries_path).parent / initial_timeseries_filename
        shutil.copyfile(self.timeseries_path, self.initial_timeseries_path)
        # Path for output penultimate timeseries (saved for record keeping)
        penultimate_timeseries_filename = ("penultimate_" +
        str(pathlib.Path(self.timeseries_path).name))
        self.penultimate_timeseries_path = pathlib.Path(self.timeseries_path).parent.parent / constants.OUTPUT_DATA_FOLDER / penultimate_timeseries_filename

        # Create Dataframe structures for timeseries
        self.initial_df = pd.read_csv(self.initial_timeseries_path, sep=";")
        self.timeseries_df = self.initial_df.copy()
        self.previous_timeseries_df = self.initial_df.copy()
        self.horizon = len(self.initial_df)

    def get_initial_setpoints(self, battery_ids):
        """
        Extract initial setpoints per battery from the initial dataframe to
        seed the Knitro starting point.
        """
        profiles = {}
        for bat_id in battery_ids:
            target_col = f"{bat_id}_targetP"
            if target_col in self.initial_df.columns:
                profiles[bat_id] = self.initial_df[target_col].to_numpy()
            elif bat_id in self.initial_df.columns:
                profiles[bat_id] = self.initial_df[bat_id].to_numpy()
            else:
                profiles[bat_id] = np.zeros(self.horizon, dtype=float)
        return profiles

    def get_initial_soc(self, bat_id: str):
        soc_col = f"{bat_id}_SOC"
        if soc_col in self.initial_df.columns:
            return float(self.initial_df[soc_col].iloc[0])
        return 0.0
    
    def reload_initial(self):
        """
        Reload initial_df from disk after external modifications (e.g. battery
        column injection). Must be called before any reset_for_iteration(0).
        """
        self.initial_df = pd.read_csv(self.initial_timeseries_path, sep=";")
        self.timeseries_df = self.initial_df.copy()
        self.previous_timeseries_df = self.initial_df.copy()

    def apply_battery_setpoints(self, battery_setpoints: dict):
        """
        Write battery setpoints.
        """
        for bat_id, profile in battery_setpoints.items():
            if len(profile) < self.horizon:
                raise ValueError(f"Profile for {bat_id} should have {self.horizon} values")
            profile = np.asarray(profile[:self.horizon], dtype=float)
            target_col = f"{bat_id}_targetP"
            min_col = f"{bat_id}_minP"
            max_col = f"{bat_id}_maxP"
            self.timeseries_df[target_col] = profile
            self.timeseries_df[min_col] = profile
            self.timeseries_df[max_col] = profile
        self.timeseries_df.to_csv(self.timeseries_path, sep=";", index=False)

    def reset_for_iteration(self, iteration_index: int):
        """
        Prepares the timeseries for a new iteration by keeping a copy of the
        previous one and (for iteration 0) ensuring we start from the initial
        data.
        """
        if iteration_index > 0:
            self.previous_timeseries_df = self.timeseries_df.copy()
        else:
            self.previous_timeseries_df = self.initial_df.copy()
            self.timeseries_df = self.initial_df.copy()

    def save_iteration_inputs(self, network_model: NetworkModel):
        self.timeseries_df.to_csv(self.timeseries_path, sep=";", index=False)
        self.previous_timeseries_df.to_csv(
            self.penultimate_timeseries_path, sep=";", index=False
        )
        network_model.save_temporary(
            pathlib.Path(self.timeseries_path).parent.parent / constants.OUTPUT_DATA_FOLDER / constants.TMP_NETWORK_FILE
        )
