# models/forecast.py
# PyTorch-based Time Series Forecasting Models

import pandas as pd
import numpy as np
from typing import Literal, Tuple, Optional
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import warnings
warnings.filterwarnings('ignore')

ForecastModelType = Literal['lstm', 'bilstm', 'gru', 'elm', 'transformer', 'tcn', 'arima', 'prophet', 'linear']


class LSTMModel(nn.Module):
    """LSTM for time series forecasting"""
    def __init__(self, input_size: int = 1, hidden_size: int = 64, num_layers: int = 2, output_size: int = 1):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=0.2)
        self.fc = nn.Sequential(
            nn.Linear(hidden_size, 32),
            nn.ReLU(),
            nn.Linear(32, output_size)
        )
    
    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        out = self.fc(lstm_out[:, -1, :])
        return out


class BiLSTMModel(nn.Module):
    """Bidirectional LSTM for time series forecasting"""
    def __init__(self, input_size: int = 1, hidden_size: int = 64, num_layers: int = 2, output_size: int = 1):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, 
                          dropout=0.2, bidirectional=True)
        self.fc = nn.Sequential(
            nn.Linear(hidden_size * 2, 32),
            nn.ReLU(),
            nn.Linear(32, output_size)
        )
    
    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        out = self.fc(lstm_out[:, -1, :])
        return out


class GRUModel(nn.Module):
    """GRU for time series forecasting"""
    def __init__(self, input_size: int = 1, hidden_size: int = 64, num_layers: int = 2, output_size: int = 1):
        super().__init__()
        self.gru = nn.GRU(input_size, hidden_size, num_layers, batch_first=True, dropout=0.2)
        self.fc = nn.Sequential(
            nn.Linear(hidden_size, 32),
            nn.ReLU(),
            nn.Linear(32, output_size)
        )
    
    def forward(self, x):
        gru_out, _ = self.gru(x)
        out = self.fc(gru_out[:, -1, :])
        return out


class ELMModel(nn.Module):
    """Extreme Learning Machine - single hidden layer with random weights"""
    def __init__(self, input_size: int = 10, hidden_size: int = 128, output_size: int = 1):
        super().__init__()
        # Random weights (not trained)
        self.W = nn.Parameter(torch.randn(input_size, hidden_size), requires_grad=False)
        self.b = nn.Parameter(torch.randn(hidden_size), requires_grad=False)
        # Output weights (trained)
        self.beta = nn.Parameter(torch.randn(hidden_size, output_size))
        
    def forward(self, x):
        # Flatten input
        batch_size = x.shape[0]
        x_flat = x.reshape(batch_size, -1)
        # Hidden layer
        H = torch.relu(torch.matmul(x_flat, self.W) + self.b)
        # Output
        out = torch.matmul(H, self.beta)
        return out


class TransformerModel(nn.Module):
    """Transformer-based model for time series forecasting"""
    def __init__(self, input_size: int = 1, d_model: int = 64, nhead: int = 4, 
                 num_layers: int = 2, output_size: int = 1):
        super().__init__()
        self.embedding = nn.Linear(input_size, d_model)
        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward=256, 
                                                   batch_first=True, dropout=0.2)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
        self.fc = nn.Sequential(
            nn.Linear(d_model, 32),
            nn.ReLU(),
            nn.Linear(32, output_size)
        )
    
    def forward(self, x):
        x = self.embedding(x)
        x = self.transformer(x)
        out = self.fc(x[:, -1, :])
        return out


class TCNBlock(nn.Module):
    """Temporal Convolutional Network Block"""
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3, dilation: int = 1):
        super().__init__()
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, 
                             padding=(kernel_size-1)*dilation, dilation=dilation)
        self.norm = nn.BatchNorm1d(out_channels)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.2)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.norm(x)
        x = self.relu(x)
        x = self.dropout(x)
        return x


class TCNModel(nn.Module):
    """Temporal Convolutional Network for time series forecasting"""
    def __init__(self, input_size: int = 1, channels: list = None, output_size: int = 1):
        super().__init__()
        if channels is None:
            channels = [32, 64, 64]
        
        layers = []
        in_ch = input_size
        for i, out_ch in enumerate(channels):
            layers.append(TCNBlock(in_ch, out_ch, dilation=2**i))
            in_ch = out_ch
        
        self.network = nn.Sequential(*layers)
        self.fc = nn.Sequential(
            nn.AdaptiveAvgPool1d(1),
            nn.Flatten(),
            nn.Linear(channels[-1], 32),
            nn.ReLU(),
            nn.Linear(32, output_size)
        )
    
    def forward(self, x):
        # TCN expects (batch, channels, length)
        x = x.transpose(1, 2)
        x = self.network(x)
        out = self.fc(x)
        return out


class TimeSeriesForecaster:
    """PyTorch-based Time Series Forecaster"""
    def __init__(self, model_type: ForecastModelType = 'lstm', horizon: int = 7, 
                 hidden_size: int = 64, learning_rate: float = 0.001, epochs: int = 50,
                 input_lags: int = 10):
        self.model_type = model_type.lower()
        self.horizon = horizon
        self.input_lags = input_lags  # Number of past days to look at
        self.hidden_size = hidden_size
        self.learning_rate = learning_rate
        self.epochs = epochs
        self.model = None
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.ts_data = None
        self.value_col = None
        self.scaler_min = None
        self.scaler_max = None

    def _create_sequences(self, data: np.ndarray, window_size: int = None) -> Tuple[np.ndarray, np.ndarray]:
        """Create sliding window sequences using input_lags as look-back window"""
        if window_size is None:
            window_size = self.input_lags
        X, y = [], []
        for i in range(len(data) - window_size):
            X.append(data[i:i+window_size])
            y.append(data[i+window_size])
        return np.array(X), np.array(y)

    def _normalize(self, data: np.ndarray) -> np.ndarray:
        """Min-max normalization"""
        self.scaler_min = np.min(data)
        self.scaler_max = np.max(data)
        return (data - self.scaler_min) / (self.scaler_max - self.scaler_min + 1e-8)

    def _denormalize(self, data: np.ndarray) -> np.ndarray:
        """Reverse normalization"""
        return data * (self.scaler_max - self.scaler_min) + self.scaler_min

    def fit(self, df: pd.DataFrame, value_col: str = 'mw_max'):
        """Train the model"""
        self.ts_data = df[value_col].values.astype(np.float32)
        self.value_col = value_col
        
        # For statistical models
        if self.model_type in ['arima', 'prophet']:
            if self.model_type == 'arima':
                from statsmodels.tsa.arima.model import ARIMA
                self.model = ARIMA(self.ts_data, order=(7, 1, 0)).fit()
            else:  # prophet
                from prophet import Prophet
                prophet_df = pd.DataFrame({
                    'ds': df.index if isinstance(df.index, pd.DatetimeIndex) else pd.date_range(start='2024-01-01', periods=len(df)),
                    'y': self.ts_data
                })
                self.model = Prophet(yearly_seasonality=True, daily_seasonality=False)
                self.model.fit(prophet_df)
            return
        
        # For linear model
        if self.model_type == 'linear':
            from sklearn.linear_model import LinearRegression
            X = np.arange(len(self.ts_data)).reshape(-1, 1)
            self.model = LinearRegression()
            self.model.fit(X, self.ts_data)
            return
        
        # For deep learning models
        if self.model_type not in ['lstm', 'bilstm', 'gru', 'elm', 'transformer', 'tcn']:
            raise ValueError(f"Unsupported model type: {self.model_type}")
        
        # Normalize data
        normalized_data = self._normalize(self.ts_data)
        
        # Create sequences using input_lags
        X, y = self._create_sequences(normalized_data, window_size=self.input_lags)
        X = X.reshape(X.shape[0], X.shape[1], 1)  # (samples, window, features)
        
        # Convert to tensors
        X_tensor = torch.FloatTensor(X).to(self.device)
        y_tensor = torch.FloatTensor(y).reshape(-1, 1).to(self.device)
        
        # Create data loader
        dataset = TensorDataset(X_tensor, y_tensor)
        dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
        
        # Create model
        if self.model_type == 'lstm':
            self.model = LSTMModel(input_size=1, hidden_size=self.hidden_size, output_size=1).to(self.device)
        elif self.model_type == 'bilstm':
            self.model = BiLSTMModel(input_size=1, hidden_size=self.hidden_size, output_size=1).to(self.device)
        elif self.model_type == 'gru':
            self.model = GRUModel(input_size=1, hidden_size=self.hidden_size, output_size=1).to(self.device)
        elif self.model_type == 'elm':
            self.model = ELMModel(input_size=self.horizon, hidden_size=128, output_size=1).to(self.device)
        elif self.model_type == 'transformer':
            self.model = TransformerModel(input_size=1, d_model=64, output_size=1).to(self.device)
        elif self.model_type == 'tcn':
            self.model = TCNModel(input_size=1, output_size=1).to(self.device)
        
        # Train
        optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
        criterion = nn.MSELoss()
        
        for epoch in range(self.epochs):
            total_loss = 0
            for X_batch, y_batch in dataloader:
                optimizer.zero_grad()
                outputs = self.model(X_batch)
                loss = criterion(outputs, y_batch)
                loss.backward()
                optimizer.step()
                total_loss += loss.item()
            
            if (epoch + 1) % max(1, self.epochs // 5) == 0:
                print(f"Epoch {epoch+1}/{self.epochs}, Loss: {total_loss/len(dataloader):.4f}")
        
        print(f"✅ {self.model_type.upper()} model trained")

    def predict(self, future_steps: int = None) -> np.ndarray:
        """Generate forecast"""
        if future_steps is None:
            future_steps = self.horizon
        
        if self.model_type in ['arima', 'prophet']:
            if self.model_type == 'arima':
                forecast = self.model.forecast(steps=future_steps)
                return np.maximum(forecast.values if hasattr(forecast, 'values') else forecast, 0)
            else:  # prophet
                future = self.model.make_future_dataframe(periods=future_steps)
                forecast = self.model.predict(future)
                return np.maximum(forecast['yhat'][-future_steps:].values, 0)
        
        if self.model_type == 'linear':
            X = np.arange(len(self.ts_data)).reshape(-1, 1)
            X_future = np.arange(len(self.ts_data), len(self.ts_data) + future_steps).reshape(-1, 1)
            pred = self.model.predict(X_future)
            return np.maximum(pred.flatten(), 0)
        
        if self.model_type in ['lstm', 'bilstm', 'gru', 'elm', 'transformer', 'tcn']:
            self.model.eval()
            predictions = []
            normalized_data = self._normalize(self.ts_data)
            
            with torch.no_grad():
                # Iteratively predict future steps
                current_window = normalized_data[-self.input_lags:].copy()
                
                for step in range(future_steps):
                    window_reshaped = current_window.reshape(1, self.input_lags, 1)
                    
                    if self.model_type == 'elm':
                        window_reshaped = window_reshaped.reshape(1, -1)
                    
                    X_tensor = torch.FloatTensor(window_reshaped).to(self.device)
                    pred_norm = self.model(X_tensor).cpu().numpy().flatten()[0]
                    predictions.append(pred_norm)
                    
                    # Update window for next prediction
                    current_window = np.append(current_window[1:], pred_norm)
            
            # Denormalize predictions
            pred = np.array(predictions)
            pred = self._denormalize(pred)
            return np.maximum(pred, 0)
        
        raise ValueError("Model not trained")

    def get_confidence_interval(self, future_steps: int = None) -> Tuple[np.ndarray, np.ndarray]:
        """Get 95% confidence interval"""
        if future_steps is None:
            future_steps = self.horizon
        
        forecast = self.predict(future_steps)
        std = np.std(self.ts_data)
        margin = 1.96 * std
        return (forecast - margin, forecast + margin)