"""
modeling.py - Model Definitions for LAI Prediction Pipeline

This module defines PyTorch-based deep learning components and data handling classes
used for predicting current year LAI from future LAI and static yearly variables.

Included Components:
- BackwardLAIPredictor: A sequence-to-sequence RNN architecture (optionally with attention)
  for backward-in-time LAI prediction.
- LAIDataset: A PyTorch Dataset for organizing LAI time series and static feature inputs.
- train_model: Model training loop with optional teacher forcing and early stopping.

These components are central to model training and evaluation.
"""

import torch
import time
import logging
import torch.nn as nn
from torch.utils.data import Dataset
from utils import evaluate


class BackwardLAIPredictor(nn.Module):
    """
    A sequence-to-sequence RNN model for predicting current-year LAI values
    from the following year's LAI (12 months) and static yearly variables.

    Supports optional attention mechanism and output constraints.

    Args:
        hidden_size (int): Size of the hidden state in RNN layers.
        rnn_cell_type (str): Type of RNN cell ('rnn', 'gru', 'lstm').
        use_attention (bool): Whether to apply attention over the encoder outputs.
        dropout (float): Dropout probability after decoder RNN.
        num_yearly_variables (int): Number of static yearly variables (default: 27).
        constrain_output (bool): Whether to apply sigmoid scaling to constrain predictions to (0, 7).
    """

    def __init__(self, hidden_size=64, rnn_cell_type="rnn", use_attention=False, dropout=0.0, num_yearly_variables=27, constrain_output = False):
        super().__init__()

        self.hidden_size = hidden_size
        self.rnn_cell_type = rnn_cell_type.lower()
        self.use_attention = use_attention
        self.constrain_output = constrain_output

        # Choose dropout layer or identity depending on dropout value
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() # nn.Identity() is a no-op layer, faster than dropout = 0

        # Map string to appropriate RNN class
        if self.rnn_cell_type == "gru":
            RNNCell = nn.GRU
        elif self.rnn_cell_type == "lstm":
            RNNCell = nn.LSTM
        elif self.rnn_cell_type == "rnn":
            RNNCell = nn.RNN
        else:
            raise ValueError(f"Unsupported RNN cell type: {rnn_cell_type}. Choose from 'rnn', 'gru', or 'lstm'.")

        # ----- Encoder -----
        # RNN encoder to process the 12 months of next year's LAI (1 feature per timestep)
        self.encoder_rnn = RNNCell(
            input_size=1,               # Only LAI value per timestep
            hidden_size=hidden_size,    # Size of the hidden state
            batch_first=True            # Input/output format: (batch, seq, feature)
        )

        # ----- Decoder -----
        # RNN decoder to predict the current year's LAI values, one at a time, autoregressively
        decoder_input_size = 1 + num_yearly_variables  # Previous LAI prediction + 27 static yearly variables
        if use_attention:
            decoder_input_size += hidden_size  # Add context vector if attention is used

        self.decoder_rnn = RNNCell(
            input_size=decoder_input_size,
            hidden_size=hidden_size,
            batch_first=True
        )

        # Fully connected layer to project RNN output (hidden state) to a single LAI value
        self.fc = nn.Linear(hidden_size, 1)

    def forward(self, lai_next_year, yearly_vars, targets=None, teacher_forcing_ratio=0.0):
        """
        Forward pass through the model. Predicts current year LAI values month-by-month.

        Args:
            lai_next_year (Tensor): Input LAI for the next year (shape: B x 12).
            yearly_vars (Tensor): Static yearly features (shape: B x num_yearly_variables).
            targets (Tensor, optional): Ground truth current-year LAI values (B x 12).
            teacher_forcing_ratio (float): Probability of using true targets as decoder input during training.

        Returns:
            Tensor: Predicted current-year LAI values (B x 12).
        """
        # ----- Encoder -----
        # Encode the 12 months of next year's LAI values (shape: B x 12 x 1)
        # The encoder internally loops through the sequence and returns the final hidden state
        encoder_input = lai_next_year.unsqueeze(-1)  # (B, 12) → (B, 12, 1)
        encoder_outputs, hidden = self.encoder_rnn(encoder_input)  # encoder_outputs: (B, 12, H), hidden: (1, B, H)

        # ----- Decoder -----
        # Use January of next year (last input to encoder, due to Dec→Jan flip) as initial decoder input
        dec_input = encoder_input[:, -1:, :].detach()

        # Expand static yearly variables to match decoder input format
        # Shape becomes (B, 1, 27) to be concatenated at each time step
        static_context = yearly_vars.unsqueeze(1)  # (B, 1, 27)

        outputs = []

        # Unpack hidden state properly for LSTM
        if self.rnn_cell_type == "lstm" and isinstance(hidden, tuple):
            h_dec, c_dec = hidden  # hidden: tuple (h_n, c_n)
        else:
            h_dec = hidden
            c_dec = None

        # Autoregressively predict each of the 12 months of the current year (backwards)
        for t in range(12):
            if self.use_attention:
                # Dot-product attention: context = softmax(q ⋅ k^T) ⋅ v
                # Query: current hidden state (B, 1, H)
                # Keys/Values: encoder outputs (B, 12, H)

                if self.rnn_cell_type == "lstm":
                    query = h_dec[-1].unsqueeze(1)  # (B, 1, H)
                else:
                    query = h_dec[-1].unsqueeze(1) if h_dec.dim() == 3 else h_dec.unsqueeze(1)  # (B, 1, H)

                # Compute attention scores using dot product
                # (B, 1, H) × (B, H, 12) → (B, 1, 12)
                scores = torch.bmm(query, encoder_outputs.transpose(1, 2)) # (B, 1, 12)

                # Normalize scores to get attention weights
                attention_weights = torch.softmax(scores, dim=-1)  # (B, 1, 12)

                # Compute context vector as weighted sum of encoder outputs
                # (B, 1, 12) × (B, 12, H) → (B, 1, H)
                context_vector = torch.bmm(attention_weights, encoder_outputs) # (B, 1, H)

                # Combine decoder input, static context, and attention context
                rnn_input = torch.cat([dec_input, static_context, context_vector], dim=-1)  # (B, 1, 1+27+H)
            else:
                # Combine decoder input and static context only
                rnn_input = torch.cat([dec_input, static_context], dim=-1)  # (B, 1, 1+27)

            # Run one step of the decoder RNN using previous hidden state
            if self.rnn_cell_type == "lstm":
                out, (h_dec, c_dec) = self.decoder_rnn(rnn_input, (h_dec, c_dec))
            else:
                out, h_dec = self.decoder_rnn(rnn_input, h_dec)

            # Apply dropout to RNN output (just identity if dropout is not enabled)
            out = self.dropout(out)

            # Project RNN output to a predicted LAI value
            lai_pred = self.fc(out)  # (B, 1, 1)

            # Constrain output to range (0,7) if enabled
            if self.constrain_output:
                lai_pred = 7.0 * torch.sigmoid(lai_pred)

            # Store prediction
            outputs.append(lai_pred)

            # Decide whether to use teacher forcing
            if self.training and targets is not None and torch.rand(1).item() < teacher_forcing_ratio:
                # Use the actual ground truth as next input
                dec_input = targets[:, t].unsqueeze(1).unsqueeze(2).detach()  # (B, 1, 1)
            else:
                # Use the model's prediction
                dec_input = lai_pred.detach()

        # Concatenate predictions into final output after squeeze: (B, 12, 1) → (B, 12)
        return torch.cat(outputs, dim=1).squeeze(-1) # (B, 12)

class LAIDataset(Dataset):
    """
    Custom PyTorch Dataset for LAI prediction. Stores next year's LAI,
    static yearly variables, and current year's target LAI values.

    Args:
        lai_next_year (np.ndarray or Tensor): Input LAI for the next year (N x 12).
        yearly_vars (np.ndarray or Tensor): Static yearly features (N x 27).
        lai_target (np.ndarray or Tensor): Target LAI for the current year (N x 12).
    """

    def __init__(self, lai_next_year, yearly_vars, lai_target):
        """
        All inputs should be flattened to shape (N, *) where N = num_locations * num_years
        """
        self.lai_next_year = lai_next_year
        self.yearly_vars = yearly_vars
        self.lai_target = lai_target

    def __len__(self):
        return self.lai_next_year.shape[0]

    def __getitem__(self, idx):
        return (
            self.lai_next_year[idx],  # (12,)
            self.yearly_vars[idx],    # (27,)
            self.lai_target[idx],     # (12,)
        )

def train_model(
    model, train_loader, val_loader, device, criterion, optimizer, num_epochs,
    teacher_forcing_ratio=0.0, use_early_stopping=False, early_stopping_patience=5, logger=logging
):
    """
    Trains the given model using a provided DataLoader and training settings.

    Supports teacher forcing and early stopping.

    Args:
        model (nn.Module): The model to train.
        train_loader (DataLoader): Training data loader.
        val_loader (DataLoader): Validation data loader.
        device (torch.device): Device to train on ('cpu' or 'cuda').
        criterion (loss function): Loss function to minimize.
        optimizer (torch.optim.Optimizer): Optimizer for parameter updates.
        num_epochs (int): Number of training epochs.
        teacher_forcing_ratio (float): Probability of using ground truth at each decoding step.
        use_early_stopping (bool): Whether to apply early stopping.
        early_stopping_patience (int): Number of epochs to wait for improvement before stopping.
        logger (logging.Logger): Logger for training progress.

    Returns:
        tuple:
            model (nn.Module): Trained model (best model if early stopping enabled).
            train_losses (list of float): Training loss values for each epoch.
            val_losses (list of float): Validation loss values for each epoch.
            best_epoch (int): Epoch with lowest validation loss.
            final_epoch (int): Final completed epoch.
            training_duration (float): Total training time in seconds.
    """

    start_time = time.time()
    train_losses = []
    val_losses = []
    best_val_loss = float('inf')
    patience_counter = 0
    best_model_state = None

    for epoch in range(num_epochs):
        model.train() # Set model to training mode (enables dropout)
        total_loss = 0.0

        for lai_next_year, yearly_vars, lai_target in train_loader:
            lai_next_year = lai_next_year.to(device)
            yearly_vars = yearly_vars.to(device)
            lai_target = lai_target.to(device)

            optimizer.zero_grad() # Clear accumulated gradients from previous step

            # output: (B, 12), same as lai_target
            output = model(lai_next_year, yearly_vars, targets=lai_target, teacher_forcing_ratio=teacher_forcing_ratio)
            loss = criterion(output, lai_target) # MSE over (B, 12), using the whole 12 months
            loss.backward() # Compute gradients w.r.t. model parameters
            optimizer.step() # Update model parameters using gradients
            total_loss += loss.item()

        # Average training and validation loss for this epoch
        avg_train_loss = total_loss / len(train_loader)
        avg_val_loss = evaluate(model, val_loader, device, criterion)

        train_losses.append(avg_train_loss)
        val_losses.append(avg_val_loss)

        logger.info(f"Epoch [{epoch+1}/{num_epochs}] | Train Loss: {avg_train_loss:.6f} | Val Loss: {avg_val_loss:.6f}")

        if use_early_stopping:
            # If validation loss improves, save best model state
            if avg_val_loss < best_val_loss:
                best_val_loss = avg_val_loss
                patience_counter = 0
                best_model_state = model.state_dict()
                logger.info("Validation loss improved. Resetting patience.")

            # If not, increment patience; stop if patience exceeds threshold
            else:
                patience_counter += 1
                logger.info(f"No improvement. Early stopping patience: {patience_counter}/{early_stopping_patience}")
                if patience_counter >= early_stopping_patience:
                    logger.info(f"Early stopping triggered at epoch {epoch+1}")
                    break

    training_duration = time.time() - start_time
    logger.info(f"Training completed in {training_duration:.2f} seconds.")

    # If early stopping was enabled and best model found, reload that state
    if use_early_stopping and best_model_state is not None:
        model.load_state_dict(best_model_state)
        logger.info("Loaded best model from early stopping checkpoint.")

    best_epoch = val_losses.index(min(val_losses)) + 1 if use_early_stopping else len(train_losses)
    
    final_epoch = len(train_losses)

    return model, train_losses, val_losses, best_epoch, final_epoch, training_duration
