import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.preprocessing import StandardScaler

class PriceLSTM(nn.Module):
    """
    PyTorch LSTM model for day-ahead electricity price mean forecasting.
    """
    def __init__(self, input_dim: int, hidden_dim: int, num_layers: int, output_dim: int = 1):
        super().__init__()
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)
        self.linear = nn.Linear(hidden_dim, output_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        lstm_out, _ = self.lstm(x)
        return self.linear(lstm_out[:, -1, :])


class LSTMSequenceModel:
    """
    Wrapper class managing standard-scaling, sequencing, PyTorch training, 
    and out-of-sample prediction.
    """
    def __init__(self, seq_length: int = 24, hidden_dim: int = 64, num_layers: int = 2, 
                 lr: float = 0.001, epochs: int = 15, batch_size: int = 64):
        self.seq_length = seq_length
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.lr = lr
        self.epochs = epochs
        self.batch_size = batch_size
        
        self.scaler_X = StandardScaler()
        self.scaler_y = StandardScaler()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = None

    def _create_sequences(self, X: np.ndarray, y: np.ndarray = None):
        xs, ys = [], []
        limit = len(X) - self.seq_length
        for i in range(limit):
            xs.append(X[i:(i + self.seq_length)])
            if y is not None:
                ys.append(y[i + self.seq_length])
        
        if y is not None:
            return np.array(xs), np.array(ys)
        return np.array(xs)

    def fit(self, X: pd.DataFrame, y: pd.Series):
        """Fits the LSTM model on X and y by scaling data and creating sequences."""
        # Standard scale
        X_scaled = self.scaler_X.fit_transform(X)
        y_scaled = self.scaler_y.fit_transform(y.values.reshape(-1, 1))

        # Create input sequences
        X_seq, y_seq = self._create_sequences(X_scaled, y_scaled)

        # Build DataLoader
        dataset = TensorDataset(torch.Tensor(X_seq), torch.Tensor(y_seq))
        loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True)

        # Initialize network
        self.model = PriceLSTM(
            input_dim=X.shape[1], 
            hidden_dim=self.hidden_dim, 
            num_layers=self.num_layers
        ).to(self.device)

        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
        criterion = nn.MSELoss()

        # Training loop
        self.model.train()
        for epoch in range(self.epochs):
            for batch_X, batch_y in loader:
                batch_X = batch_X.to(self.device)
                batch_y = batch_y.to(self.device)
                
                optimizer.zero_grad()
                outputs = self.model(batch_X)
                loss = criterion(outputs, batch_y)
                loss.backward()
                optimizer.step()

        return self

    def predict(self, X: pd.DataFrame, y: pd.Series = None):
        """
        Generates predictions for the given features X.
        If y is provided, aligns and returns (y_pred, y_aligned) for evaluation metrics.
        If y is None, returns y_pred alone.
        """
        if self.model is None:
            raise ValueError("Model is not fitted yet. Call fit() first.")
        
        # Scale X using trained scaler
        X_scaled = self.scaler_X.transform(X)

        self.model.eval()
        with torch.no_grad():
            if y is not None:
                y_scaled = self.scaler_y.transform(y.values.reshape(-1, 1))
                X_seq, y_seq = self._create_sequences(X_scaled, y_scaled)
                X_tensor = torch.Tensor(X_seq).to(self.device)
                
                preds_scaled = self.model(X_tensor).cpu().numpy()
                y_pred = self.scaler_y.inverse_transform(preds_scaled).flatten()
                y_aligned = self.scaler_y.inverse_transform(y_seq).flatten()
                return y_pred, y_aligned
            else:
                X_seq = self._create_sequences(X_scaled)
                X_tensor = torch.Tensor(X_seq).to(self.device)
                
                preds_scaled = self.model(X_tensor).cpu().numpy()
                y_pred = self.scaler_y.inverse_transform(preds_scaled).flatten()
                return y_pred