import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from prophet import Prophet
from pmdarima import auto_arima
import logging

logger = logging.getLogger(__name__)

class LSTMForecaster(nn.Module):
    """LSTM model for time series forecasting."""
    
    def __init__(self, input_size: int, hidden_size: int, num_layers: int, 
                 output_size: int, dropout: float = 0.2):
        super(LSTMForecaster, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, 
                           batch_first=True, dropout=dropout)
        self.dropout = nn.Dropout(dropout)
        self.linear = nn.Linear(hidden_size, output_size)
    
    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        lstm_out = self.dropout(lstm_out[:, -1, :])  # Take the last output
        out = self.linear(lstm_out)
        return out

class AdvancedModelTrainer:
    """Trainer for advanced forecasting models."""
    
    def __init__(self, config: dict):
        self.config = config
        
    def train_lstm(self, X_train: np.ndarray, y_train: np.ndarray,
                  X_val: np.ndarray = None, 
                  y_val: np.ndarray = None) -> nn.Module:
        """Train LSTM model."""
        model_config = self.config['lstm']
        
        # Convert to PyTorch tensors
        train_dataset = TensorDataset(
            torch.FloatTensor(X_train), 
            torch.FloatTensor(y_train)
        )
        train_loader = DataLoader(train_dataset, batch_size=model_config['batch_size'], shuffle=True)
        
        # Initialize model
        model = LSTMForecaster(
            input_size=X_train.shape[2],
            hidden_size=model_config['hidden_size'],
            num_layers=model_config['num_layers'],
            output_size=y_train.shape[1],
            dropout=model_config['dropout']
        )
        
        # Training setup
        criterion = nn.MSELoss()
        optimizer = torch.optim.Adam(model.parameters(), lr=model_config['learning_rate'])
        
        # Training loop
        for epoch in range(model_config['epochs']):
            model.train()
            epoch_loss = 0
            
            for batch_X, batch_y in train_loader:
                optimizer.zero_grad()
                predictions = model(batch_X)
                loss = criterion(predictions, batch_y)
                loss.backward()
                
                # Gradient clipping
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                
                optimizer.step()
                epoch_loss += loss.item()
            
            if epoch % 10 == 0:
                logger.info(f'Epoch {epoch}, Loss: {epoch_loss/len(train_loader):.4f}')
        
        return model
    
    def train_prophet(self, df: pd.DataFrame, 
                     date_col: str, 
                     value_col: str) -> Prophet:
        """Train Facebook Prophet model."""
        prophet_df = df[[date_col, value_col]].rename(
            columns={date_col: 'ds', value_col: 'y'}
        )
        
        model = Prophet(
            changepoint_prior_scale=self.config['prophet'].get('changepoint_prior_scale', 0.05),
            seasonality_prior_scale=self.config['prophet'].get('seasonality_prior_scale', 10),
            yearly_seasonality=self.config['prophet'].get('yearly_seasonality', True),
            weekly_seasonality=self.config['prophet'].get('weekly_seasonality', True),
            daily_seasonality=self.config['prophet'].get('daily_seasonality', False)
        )
        
        model.fit(prophet_df)
        return model
    
    def train_auto_arima(self, series: pd.Series) -> object:
        """Train auto ARIMA model."""
        model = auto_arima(
            series,
            start_p=1,
            start_q=1,
            max_p=3,
            max_q=3,
            seasonal=True,
            m=7,
            stepwise=True,
            suppress_warnings=True,
            error_action='ignore'
        )
        return model