
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
deep learning RUL models core class definition

Author: rameshbk
Last modified: Shweta, Nov 18th, 2025

Configurable RUL Estimation Models:
- BiLSTM (PyTorch)
- TCM (Temporal Convolutional Model)
- RandomForest Regressor
- DecisionTree Regressor

"""

import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.tree import DecisionTreeRegressor
from IntelliMaint.utils import Utils

# Optional PyTorch
try:
    import torch
    from torch import nn
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    print("[INFO] PyTorch is not installed. Please install PyTorch to use BiLSTMModel and TCMModel.")
    

# =====================================================================
# 1. BiLSTM MODEL (Configurable)
# =====================================================================
if TORCH_AVAILABLE:

    class BiLSTMModel:
        def __init__(self, config=None, verbose=False):
            """
            Bi-directional LSTM model for RUL sequence prediction.

            Args:
                config : dict, optional
                    Dictionary containing configuration parameters.

                Configurable Parameters:
                    input_size (int): Number of input features per timestep.
                        Default: 1

                    hidden_size (int): LSTM hidden dimension.
                        Default: 32

                    num_layers (int): Number of stacked LSTM layers.
                        Default: 1

                verbose : bool
                    Enables detailed debug logs.


            Returns:
                None
            """
            self.verbose = verbose
            ut = Utils(verbose=verbose)
            self.logger = ut.get_logger(self.__class__.__name__)
            self._print_debug = ut.print_debug

            default_config = {
                "input_size": 1,
                "hidden_size": 32,
                "num_layers": 1,
            }

            self.config = {**default_config, **(config or {})}

            if self.verbose:
                print(
                    f"[DEBUG] BiLSTMModel initialized with "
                    f"input_size={self.config['input_size']}, "
                    f"hidden_size={self.config['hidden_size']}, "
                    f"num_layers={self.config['num_layers']}"
                )

            self.logger.info(
                f"Initialized BiLSTMModel with config: "
                f"input_size={self.config['input_size']}, "
                f"hidden_size={self.config['hidden_size']}, "
                f"num_layers={self.config['num_layers']}"
            )

            self.model = self._build_model()

        def _build_model(self):
            """
            Builds the internal BiLSTM neural network architecture.

            This method constructs:
                - A bidirectional LSTM layer
                - A fully connected output layer

            Returns:
                nn.Module: The configured BiLSTM network.
            """
            input_size = self.config["input_size"]
            hidden = self.config["hidden_size"]
            layers = self.config["num_layers"]

            class _BiLSTM(nn.Module):
                def __init__(self):
                    super().__init__()
                    self.lstm = nn.LSTM(
                        input_size=input_size,
                        hidden_size=hidden,
                        num_layers=layers,
                        batch_first=True,
                        bidirectional=True,
                    )
                    self.fc = nn.Linear(hidden * 2, 1)

                def forward(self, x):
                    out, _ = self.lstm(x)
                    last = out[:, -1, :]
                    return self.fc(last)

            self.logger.info("BiLSTM model initialized.")
            return _BiLSTM()

        def predict(self, x):
            """
            Predicts output values using the BiLSTM model.

            Args:
                x (array-like):
                    Input array of shape (batch, sequence_length, input_size).

            Returns:
                np.ndarray: Model predictions as a NumPy array.
            """
            self.model.eval()
            with torch.no_grad():
                t = torch.tensor(x, dtype=torch.float32)
                preds = self.model(t).cpu().numpy()

            # Debug output
            self.logger.debug(f"BiLSTM predictions sample: {preds[:5]}")
            
            return preds
        
        def forecast_autoregressive(self, predict_fn, last_window, steps):
            """
            Multi-step autoregressive forecasting helper.

            Args:
                predict_fn : callable
                    Function that takes X_window and returns prediction.

                last_window : array-like
                    Latest known window of HI values.

                steps : int
                    Number of future predictions to generate.

            Returns:
                np.ndarray : Forecasted HI values.
            """
            window = last_window.astype(float).copy()
            preds = []

            for _ in range(steps):
                X = window.reshape(1, -1)
                y = float(predict_fn(X))
                preds.append(y)

                window = np.roll(window, -1)
                window[-1] = y

            return np.array(preds)


        def compute_rul_steps(self, last_window, forecast, threshold):
            """
            Computes RUL (in steps) until HI reaches a threshold.

            Args:
                last_window (array-like): Last known sequence.
                forecast (array-like): Predicted future values.
                threshold (float): Failure threshold.

            Returns:
                float : RUL steps (or inf if never reached).
            """
            hi_full = np.concatenate([last_window, forecast])
            idx = np.where(hi_full >= threshold)[0]

            if len(idx) == 0:
                return np.inf

            first = idx[0]
            return max(0.0, float(first - (len(last_window) - 1)))
                


# =====================================================================
# 2. TCM MODEL (Configurable)
# =====================================================================
if TORCH_AVAILABLE:

    class TCMModel:
        def __init__(self, config=None, verbose=False):
            """
            Temporal Convolution Model (TCM) for sequence prediction.

            Args:
                config : dict, optional
                    Configurable parameters include:

                Configurable Parameters:
                    channels (int): Number of convolutional channels.
                        Default: 32

                    kernel_size (int): Convolution kernel size.
                        Default: 3

                    num_layers (int): Number of dilated conv layers.
                        Default: 3

                verbose : bool
                    Enables detailed debug logs.

            Returns:
                None
            """
            self.verbose = verbose
            ut = Utils(verbose=verbose)
            self.logger = ut.get_logger(self.__class__.__name__)
            self._print_debug = ut.print_debug

            default_config = {
                "channels": 32,
                "kernel_size": 3,
                "num_layers": 3,
            }

            self.config = {**default_config, **(config or {})}

            if self.verbose:
                print(
                    f"[DEBUG] TCMModel initialized with "
                    f"channels={self.config['channels']}, "
                    f"kernel_size={self.config['kernel_size']}, "
                    f"num_layers={self.config['num_layers']}"
                )

            self.logger.info(
                f"Initialized TCMModel with config: "
                f"channels={self.config['channels']}, "
                f"kernel_size={self.config['kernel_size']}, "
                f"num_layers={self.config['num_layers']}"
            )

            self.model = self._build_model()

        def _build_model(self):
            """
            Builds the TCM (Temporal Convolutional Model) architecture.

            The constructed model includes:
                - A stack of dilated convolutional layers
                - ReLU activations
                - Adaptive average pooling
                - A fully connected output layer

            Returns:
                nn.Module: The configured TCM network.
            """
            channels = self.config["channels"]
            kernel = self.config["kernel_size"]
            layers_n = self.config["num_layers"]

            class _TCM(nn.Module):
                def __init__(self):
                    super().__init__()
                    layers = []
                    in_ch = 1
                    dilations = [1, 2, 4, 8, 16][:layers_n]

                    for d in dilations:
                        layers.append(
                            nn.Conv1d(
                                in_channels=in_ch,
                                out_channels=channels,
                                kernel_size=kernel,
                                dilation=d,
                                padding=d * (kernel - 1) // 2,
                            )
                        )
                        layers.append(nn.ReLU())
                        in_ch = channels

                    self.net = nn.Sequential(*layers)
                    self.pool = nn.AdaptiveAvgPool1d(1)
                    self.fc = nn.Linear(channels, 1)

                def forward(self, x):
                    out = self.net(x)
                    out = self.pool(out)[:, :, 0]
                    return self.fc(out)

            self.logger.info("TCM model initialized.")
            return _TCM()

        def predict(self, x):
            """
            Predicts output values using the TCM model.

            Args:
                x (array-like):
                    Input tensor with shape (batch, channels=1, sequence_length).

            Returns:
                np.ndarray: Predicted values as a NumPy array.
            """
            self.model.eval()
            with torch.no_grad():
                t = torch.tensor(x, dtype=torch.float32)
                preds = self.model(t).cpu().numpy()

            self.logger.debug(f"TCM predictions sample: {preds[:5]}")

            return preds
        
        def forecast_autoregressive(self, predict_fn, last_window, steps):
            """
            Multi-step autoregressive forecasting helper.

            Args:
                predict_fn : callable
                    Function that takes X_window and returns prediction.

                last_window : array-like
                    Latest known window of HI values.

                steps : int
                    Number of future predictions to generate.

            Returns:
                np.ndarray : Forecasted HI values.
            """
            window = last_window.astype(float).copy()
            preds = []

            for _ in range(steps):
                X = window.reshape(1, -1)
                y = float(predict_fn(X))
                preds.append(y)

                window = np.roll(window, -1)
                window[-1] = y

            return np.array(preds)


        def compute_rul_steps(self, last_window, forecast, threshold):
            """
            Computes RUL (in steps) until HI reaches a threshold.

            Args:
                last_window (array-like): Last known sequence.
                forecast (array-like): Predicted future values.
                threshold (float): Failure threshold.

            Returns:
                float : RUL steps (or inf if never reached).
            """
            hi_full = np.concatenate([last_window, forecast])
            idx = np.where(hi_full >= threshold)[0]

            if len(idx) == 0:
                return np.inf

            first = idx[0]
            return max(0.0, float(first - (len(last_window) - 1)))

# =====================================================================
# 3. SKLEARN MODELS (RF + DT)
# =====================================================================

class RandomForestRegressorModel:
    def __init__(self, config=None, verbose=False):
        """
        Random Forest regression model (scikit-learn).

        Args:
            config : dict, optional
                Configurable parameters:
                n_estimators (int): Number of trees. Default: 300  
                max_depth (int or None): Tree depth. Default: None

            verbose : bool
                If True, prints debug logs.

        Returns:
            None
        """
        self.verbose = verbose
        ut = Utils(verbose=verbose)
        self.logger = ut.get_logger(self.__class__.__name__)
        self._print_debug = ut.print_debug

        default_config = {
            "n_estimators": 300,
            "max_depth": None,
        }

        self.config = {**default_config, **(config or {})}

        if self.verbose:
            print(
                f"[DEBUG] RandomForestRegressorModel initialized with "
                f"n_estimators={self.config['n_estimators']}, "
                f"max_depth={self.config['max_depth']}"
            )

        self.logger.info(
            f"Initialized RandomForestRegressorModel with config: "
            f"n_estimators={self.config['n_estimators']}, "
            f"max_depth={self.config['max_depth']}"
        )

        self.model = RandomForestRegressor(
            n_estimators=self.config["n_estimators"],
            max_depth=self.config["max_depth"],
            random_state=0
        )

        self.logger.info("RandomForestRegressor initialized successfully.")

    def train(self, X, y):
        """
        Trains the RandomForestRegressor on the provided dataset.

        Args:
            X (array-like): Training input features.
            y (array-like): Target output values.

        Returns:
            None
        """
        self.model.fit(X, y)
        self.logger.info("RandomForestRegressor training completed.")

    def predict(self, X):
        """
        Generates predictions using the trained Random Forest model.

        Args:
            X (array-like): Input features.

        Returns:
            np.ndarray: Predicted output values.
        """
        preds = self.model.predict(X)
        self.logger.debug(f"RF predictions sample: {preds[:5]}")
        
        return preds
    
    def forecast_autoregressive(self, predict_fn, last_window, steps):
        """
        Multi-step autoregressive forecasting helper.

        Args:
            predict_fn : callable
                Function that takes X_window and returns prediction.

            last_window : array-like
                Latest known window of HI values.

            steps : int
                Number of future predictions to generate.

        Returns:
            np.ndarray : Forecasted HI values.
        """
        window = last_window.astype(float).copy()
        preds = []

        for _ in range(steps):
            X = window.reshape(1, -1)
            y = float(predict_fn(X))
            preds.append(y)

            window = np.roll(window, -1)
            window[-1] = y

        return np.array(preds)


    def compute_rul_steps(self, last_window, forecast, threshold):
        """
        Computes RUL (in steps) until HI reaches a threshold.

        Args:
            last_window (array-like): Last known sequence.
            forecast (array-like): Predicted future values.
            threshold (float): Failure threshold.

        Returns:
            float : RUL steps (or inf if never reached).
        """
        hi_full = np.concatenate([last_window, forecast])
        idx = np.where(hi_full >= threshold)[0]

        if len(idx) == 0:
            return np.inf

        first = idx[0]
        return max(0.0, float(first - (len(last_window) - 1)))


class DecisionTreeRegressorModel:
    def __init__(self, config=None, verbose=False):
        """
        Decision Tree regression model (scikit-learn).

        Args:
            config : dict, optional
                Configurable Parameters:
                    max_depth (int or None)
                        Default: None

            verbose : bool
                Enable detailed logs.

        Returns:
            None
        """
        self.verbose = verbose
        ut = Utils(verbose=verbose)
        self.logger = ut.get_logger(self.__class__.__name__)
        self._print_debug = ut.print_debug

        default_config = {
            "max_depth": None
        }

        self.config = {**default_config, **(config or {})}

        if self.verbose:
            print(
                f"[DEBUG] DecisionTreeRegressorModel initialized with "
                f"max_depth={self.config['max_depth']}"
            )

        self.logger.info(
            f"Initialized DecisionTreeRegressorModel with config: "
            f"max_depth={self.config['max_depth']}"
        )

        self.model = DecisionTreeRegressor(
            max_depth=self.config["max_depth"],
            random_state=0
        )

        self.logger.info("DecisionTreeRegressor initialized successfully.")

    def train(self, X, y):
        """
        Fits the DecisionTreeRegressor to the training dataset.

        Args:
            X (array-like): Training input features.
            y (array-like): Training target values.

        Returns:
            None
        """
        self.model.fit(X, y)
        self.logger.info("DecisionTreeRegressor training completed.")

    def predict(self, X):
        """
        Predicts outputs using the trained Decision Tree regressor.

        Args:
            X (array-like): Input feature matrix.

        Returns:
            np.ndarray: Predicted values.
        """
        preds = self.model.predict(X)
        self.logger.debug(f"DT predictions sample: {preds[:5]}")

        return preds
    
    def forecast_autoregressive(self, predict_fn, last_window, steps):
        """
        Multi-step autoregressive forecasting helper.

        Args:
            predict_fn : callable
                Function that takes X_window and returns prediction.

            last_window : array-like
                Latest known window of HI values.

            steps : int
                Number of future predictions to generate.

        Returns:
            np.ndarray : Forecasted HI values.
        """
        window = last_window.astype(float).copy()
        preds = []

        for _ in range(steps):
            X = window.reshape(1, -1)
            y = float(predict_fn(X))
            preds.append(y)

            window = np.roll(window, -1)
            window[-1] = y

        return np.array(preds)


    def compute_rul_steps(self, last_window, forecast, threshold):
        """
        Computes RUL (in steps) until HI reaches a threshold.

        Args:
            last_window (array-like): Last known sequence.
            forecast (array-like): Predicted future values.
            threshold (float): Failure threshold.

        Returns:
            float : RUL steps (or inf if never reached).
        """
        hi_full = np.concatenate([last_window, forecast])
        idx = np.where(hi_full >= threshold)[0]

        if len(idx) == 0:
            return np.inf

        first = idx[0]
        return max(0.0, float(first - (len(last_window) - 1)))



