models.py

#  MIT License
#
#  Copyright (c) 2025 Honda Research Institute Europe GmbH
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to deal
#  in the Software without restriction, including without limitation the rights
#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#  copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#
#  The above copyright notice and this permission notice shall be included in all
#  copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#  SOFTWARE.
#
#    SPDX-License-Identifier: MIT
#

import torch
import logging
from pathlib import Path
from enum import Enum

import lightning as L
import numpy as np
import pandas as pd
from typing import Any, Callable
from sklearn.linear_model import LinearRegression
from sklearn.neighbors import KNeighborsRegressor
from xgboost import XGBRegressor
from torch import optim, nn, Tensor
from torch.utils.data import DataLoader, TensorDataset
from lightning.pytorch.utilities.types import OptimizerLRScheduler, STEP_OUTPUT
from lightning.pytorch.callbacks.early_stopping import EarlyStopping
from sklearn.preprocessing import StandardScaler

from error_compensation.interfaces import MultishotCompensator, SingleshotCompensator, ModelRegressor


class Features:
    base_features: list = ["P_dem", "theta_air_0"]
    ambient_temp: list = ["theta_air_0"]
    interior_temp: list = [f"theta_{i}" for i in range(1, 10)]
    pv_feature: list = ["P_ren"]
    past_temp_features: list = ["theta_air_1", "theta_air_2"]
    categorical_tod_features: list = ["time_of_day"]
    categorical_dow_features: list = ["day_of_week"]
    categorical_doy_features: list = ["day_of_year"]
    cyclic_tod_features: list = ["sin_time_of_day", "cos_time_of_day"]
    cyclic_dow_features: list = ["sin_day_of_week", "cos_day_of_week"]
    cyclic_doy_features: list = ["sin_day_of_year", "cos_day_of_year"]
    server_load: list = ["P_server_o4", "P_server_cis"]


class XGBMultishotCompensator(MultishotCompensator):
    def __init__(self, name: str, features: list[str], random_states: list[int] = None):
        if random_states is None:
            random_states = [42] * 9

        models = [XGBRegressor(objective="reg:squarederror", random_state=random_states[i]) for i in range(9)]
        super().__init__(name, features, models=models)


class LinearMultishotCompensator(MultishotCompensator):
    def __init__(self, name: str, features: list[str]):
        models = [LinearRegression() for _ in range(9)]
        super().__init__(name, features, models=models)


class LinearSingleshotCompensator(SingleshotCompensator):
    def __init__(self, name: str, features: list[str]):
        super().__init__(name, features, model=LinearRegression())


class XGBSingleshotCompensator(SingleshotCompensator):
    def __init__(self, name: str, features: list[str], random_state: int = None):
        model = XGBRegressor(random_state=random_state)
        super().__init__(name, features, model)


class KNNMultishotCompensator(MultishotCompensator):
    requires_scaling = True

    def __init__(self, name: str, features: list[str]):
        models = [KNeighborsRegressor() for _ in range(9)]
        super().__init__(name, features, models=models)


class KNNSingleshotCompensator(SingleshotCompensator):
    requires_scaling = True

    def __init__(self, name: str, features: list[str]):
        model = KNeighborsRegressor()
        super().__init__(name, features, model=model)


class TwoLayerReluMLP(L.LightningModule):
    loss_fn: Callable = None

    def __init__(self, inputs: list[str], n_outputs: int = 9, n_first_layer: int = 128, n_second_layer: int = 128, loss_fn: Callable = nn.functional.mse_loss):
        super().__init__()

        n_inputs = len(inputs)

        self.linear_relu_stack = nn.Sequential(
            nn.Linear(n_inputs, n_first_layer),
            nn.ReLU(),
            nn.Linear(n_first_layer, n_second_layer),
            nn.ReLU(),
            nn.Linear(n_second_layer, n_outputs)
        )

        self.loss_fn = loss_fn

    def _calculate_loss(self, batch, batch_idx) -> Tensor:
        x, y = batch

        pred = self.linear_relu_stack(x)
        loss = self.loss_fn(pred, y)

        return loss

    def training_step(self, batch, batch_idx) -> Tensor:
        loss = self._calculate_loss(batch, batch_idx)
        self.log("train_loss", loss)

        return loss

    def validation_step(self, batch, batch_idx) -> STEP_OUTPUT:
        loss = self._calculate_loss(batch, batch_idx)
        self.log("val_loss", loss)

        return loss

    def test_step(self, batch, batch_idx) -> STEP_OUTPUT:
        loss = self._calculate_loss(batch, batch_idx)
        self.log("test_loss", loss)

        return loss

    def forward(self, x) -> Any:
        return self.linear_relu_stack.forward(x)

    def configure_optimizers(self) -> OptimizerLRScheduler:
        optimizer = optim.Adam(self.parameters(), lr=1e-3, weight_decay=1e-4, eps=1e-6)
        scheduler = optim.lr_scheduler.CyclicLR(optimizer, base_lr=1e-6, max_lr=1e-3, cycle_momentum=False)

        # scheduler = torch.optim.lr_scheduler.OneCycleLR(
        #     optimizer, max_lr=1e-3, total_steps=self.trainer.estimated_stepping_batches
        # )

        return [optimizer], [scheduler]


class MLPSingleShotCompensator(SingleshotCompensator):
    requires_scaling = True
    requires_validation = True
    model: TwoLayerReluMLP
    training_max_epochs: int
    output_scaler: StandardScaler = None
    _loss_fn: Callable = None

    def __init__(
        self,
        name: str,
        features: list[str],
        input_scaler: StandardScaler = None,
        output_scaler: StandardScaler = None,
        loss_fn: Callable = nn.functional.mse_loss,
        training_max_epochs: int = 50,
        initialize_model: bool = False,
    ):
        self._loss_fn = loss_fn
        if initialize_model:
            model = TwoLayerReluMLP(features, loss_fn=loss_fn)
        else:
            model = None

        self.training_max_epochs = training_max_epochs

        if input_scaler:
            self.input_scaler = input_scaler
        if output_scaler:
            self.output_scaler = output_scaler

        super().__init__(name, features, model)

    def dump(self, filepath: Path):
        self.model.eval()
        super().dump(filepath)

    def load(self, filepath: Path):
        pass

    def _infer(self, x: torch.Tensor) -> torch.Tensor:
        return self.model(x)

    def predict(self, X: dict[str, float or np.array] or pd.DataFrame) -> np.array:
        x = self._prepare_input(X)

        if not self.input_scaler:
            raise ValueError("Scaler not set yet, not trained yet? set_scaler() not called?")

        x = torch.tensor(x, dtype=torch.float32)

        if not self.model:
            raise ValueError("model not set yet, not trained yet? initialize_model = False?")

        self.model.eval()

        y_hat = self._infer(x)
        y_hat = self.output_scaler.inverse_transform(y_hat.detach().numpy())
        # y_hat = y_hat.detach().numpy()

        if isinstance(y_hat, np.ndarray) and y_hat.shape[0] == 1:
            y_hat = y_hat[0]

        return y_hat

    def train(self, X_train: pd.DataFrame, y_train: pd.DataFrame, **kwargs):
        if "X_val" not in kwargs or "y_val" not in kwargs:
            raise ValueError("MLP training requires validation set for early stopping")

        n_outputs = len(y_train.columns)
        if not self.model:
            self.model = TwoLayerReluMLP(self.features, n_outputs=n_outputs, loss_fn=self._loss_fn)

        X_val = kwargs["X_val"]
        y_val = kwargs["y_val"]

        # filter training features
        X_train = X_train[self.features]
        X_val = X_val[self.features]

        self.input_scaler = StandardScaler()
        self.input_scaler.fit(X_train.values)

        self.output_scaler = StandardScaler()
        self.output_scaler.fit(y_train.values)

        X_train = self.input_scaler.transform(X_train.values)
        y_train = self.output_scaler.transform(y_train.values)
        # y_train = y_train.values

        X_val = self.input_scaler.transform(X_val.values)
        y_val = self.output_scaler.transform(y_val.values)
        # y_val = y_val.values

        X_train = torch.tensor(X_train, dtype=torch.float32)
        y_train = torch.tensor(y_train, dtype=torch.float32)

        X_val = torch.tensor(X_val, dtype=torch.float32)
        y_val = torch.tensor(y_val, dtype=torch.float32)

        train_set = TensorDataset(torch.tensor(X_train), torch.tensor(y_train))
        val_set = TensorDataset(torch.tensor(X_val), torch.tensor(y_val))

        # Create data loaders.
        batch_size = 64
        train_dataloader = DataLoader(train_set, batch_size=batch_size)
        val_dataloader = DataLoader(val_set, batch_size=batch_size)

        trainer = L.Trainer(max_epochs=self.training_max_epochs, callbacks=[EarlyStopping(monitor="val_loss", mode="min", min_delta=1e-3)])
        trainer.fit(model=self.model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)

    def test(self, X_test: pd.DataFrame, y_test: pd.DataFrame):
        X_test = X_test[self.features]

        if not self.input_scaler:
            raise ValueError("Scaler not set yet, not trained yet? set_scaler() not called?")

        X_test = self.input_scaler.transform(X_test)
        X_test = torch.tensor(X_test, dtype=torch.float32)
        y_test = self.output_scaler.transform(y_test)

        test_set = TensorDataset(torch.tensor(X_test), torch.tensor(y_test))
        test_dataloader = DataLoader(test_set, batch_size=len(test_set))

        trainer = L.Trainer(max_epochs=50, callbacks=[EarlyStopping(monitor="val_loss", mode="min")])
        result = trainer.test(model=self.model, dataloaders=test_dataloader)

        return result

    def set_scaler(self, input_scaler: StandardScaler, output_scaler: StandardScaler):
        self.input_scaler = input_scaler
        self.output_scaler = output_scaler


class PCNNCompensator(MLPSingleShotCompensator):
    def __init__(
            self,
            name: str,
            features: list[str],
            input_scaler: StandardScaler,
            output_scaler: StandardScaler,
            trained_pcnn_model,
    ):

        super().__init__(
            name=name,
            features=features,
            input_scaler=input_scaler,
            output_scaler=output_scaler,
            initialize_model=False
        )

        self.model = trained_pcnn_model.mlp

    def _prepare_input(self, X: dict[str, float or np.array] or pd.DataFrame):
        # overwritten to enable named scalers, as in PCNN pipeline, scalers are all working on column names

        if isinstance(X, pd.DataFrame):
            x = X[[feature for feature in self.features]]
        else:
            x = pd.DataFrame(data={feature: [X[feature]] for feature in self.features})

        x = self.input_scaler.transform(x)

        return x

    def train(self, X_train: pd.DataFrame, y_train: pd.DataFrame, **kwargs):
        logging.warning("Pretrained PCNN compensator cannot be trained, nothing will be done")