import os
import torch
import torch.nn as nn
from tqdm.auto import tqdm
from copy import deepcopy

class Trainer:
    """
    Generic trainer for PyTorch models with support for early stopping and visualization.
    
    This class handles the training loop, validation, and evaluation of models
    for various time series tasks.
    
    Parameters:
    -----------
    model : torch.nn.Module
        The model to train.
    train_loader : torch.utils.data.DataLoader
        DataLoader for training data.
    val_loader : torch.utils.data.DataLoader, optional
        DataLoader for validation data.
    test_loader : torch.utils.data.DataLoader, optional
        DataLoader for test data.
    lr : float, default=1e-3
        Learning rate for the optimizer.
    weight_decay : float, default=1e-4
        Weight decay (L2 penalty) for the optimizer.
    epochs : int, default=1000
        Maximum number of training epochs.
    accum_steps : int, default=1
        Number of steps to accumulate gradients before updating weights.
    patience : int, default=5
        Number of epochs with no improvement after which training will be stopped.
    min_delta : float, default=1e-4
        Minimum change in validation loss to qualify as an improvement.
    max_grad_norm : float, default=5.0
        Maximum norm of gradients for gradient clipping.
    device : str, optional
        Device to use for training ('cuda' or 'cpu'). If None, will use CUDA if available.
    save_dir : str, default="./checkpoints"
        Directory to save model checkpoints.
    callbacks : list, optional
        List of callback functions to call after each epoch.
    criterion : callable, optional
        Loss function. If None, MSELoss will be used.
    optimizer : torch.optim.Optimizer, optional
        Optimizer. If None, AdamW will be used.
    scheduler : torch.optim.lr_scheduler._LRScheduler, optional
        Learning rate scheduler. If None, ReduceLROnPlateau will be used.
    """
    def __init__(
        self,
        model,
        train_loader,
        val_loader=None,
        test_loader=None,
        lr=1e-3,
        weight_decay=1e-4,
        epochs=1000,
        accum_steps=1,
        patience=5,
        min_delta=1e-4,
        max_grad_norm=5.0,
        device=None,
        save_dir="./checkpoints",
        callbacks=None,
        criterion=None,
        optimizer=None,
        scheduler=None,
    ):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.model = model.to(self.device)
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.test_loader = test_loader

        self.epochs = epochs
        self.accum_steps = accum_steps
        self.patience = patience
        self.min_delta = min_delta
        self.max_grad_norm = max_grad_norm
        self.save_dir = save_dir
        os.makedirs(save_dir, exist_ok=True)

        # --- Loss function ---
        self.criterion = criterion or nn.MSELoss()

        # --- Optimizer ---
        self.optimizer = optimizer or torch.optim.AdamW(
            self.model.parameters(), lr=lr, weight_decay=weight_decay
        )

        # --- Scheduler ---
        self.scheduler = scheduler or torch.optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer,
            mode='min',
            factor=0.1,
            patience=5,
            threshold=1e-4,
            threshold_mode='rel',
            cooldown=0,
            min_lr=1e-6
        ) 

        self.train_losses = []
        self.val_losses = []
        self.callbacks = callbacks or []
        self._viz_snapshots = []


    def fit(self):
        """
        Train the model for the specified number of epochs.
        
        This method handles the training loop, validation, early stopping,
        and callback execution.
        
        Returns:
        --------
        self
            The trainer instance.
        """
        best_val = float("inf")
        best_state = None
        no_improve = 0

        epoch_bar = tqdm(range(self.epochs), desc="Training", position=0)
        for epoch in epoch_bar:
            train_loss = self._run_epoch(self.train_loader, train=True, accum_steps=self.accum_steps)
            val_loss = self._run_epoch(self.val_loader, train=False) if self.val_loader else None

            if self.scheduler:
                if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                # ReduceLROnPlateau requires the validation loss
                    self.scheduler.step(val_loss)
                else:
                    self.scheduler.step()


            self.train_losses.append(train_loss)
            if val_loss is not None:
                self.val_losses.append(val_loss)
            
            # --- Update tqdm display ---
            postfix = {"Train Loss": f"{train_loss:.6f}"}
            if val_loss is not None:
                postfix["Val Loss"] = f"{val_loss:.6f}"
            epoch_bar.set_postfix(postfix)
            epoch_bar.set_description(f"Epoch {epoch+1}/{self.epochs}")
            
            # Callbacks (for logging, visualization, etc.)
            self._current_epoch = epoch
            for cb in self.callbacks:
                cb(epoch=epoch, trainer=self)

            # --- Early stopping ---
            if self.val_loader is not None:
                if best_val - val_loss > self.min_delta:
                    best_val = val_loss
                    best_state = deepcopy(self.model.state_dict())
                    no_improve = 0
                else:
                    no_improve += 1
                    if no_improve >= self.patience:
                        print("Early stopping triggered.")
                        break

        if best_state:
            self.model.load_state_dict(best_state)
        
        return self

    def _run_epoch(self, loader, train=True, accum_steps=1):
        """
        Run a single epoch of training or validation.
        
        Parameters:
        -----------
        loader : torch.utils.data.DataLoader
            DataLoader for the data.
        train : bool, default=True
            If True, perform training (backward pass and optimization).
            If False, perform validation only.
        accum_steps : int, default=1
            Number of steps to accumulate gradients before updating weights.
            
        Returns:
        --------
        float
            Average loss for the epoch.
        """
        if loader is None:
            return None

        self.model.train(train)
        total_loss = 0.0

        with torch.set_grad_enabled(train):
            for step, (data, target) in enumerate(loader):
                target = target.to(self.device)
                preds = self.model(data)
                loss = self.criterion(preds, target) / accum_steps

                if not torch.isfinite(loss):
                    print("⚠️ Non-finite loss, skipping batch")
                    self.optimizer.zero_grad(set_to_none=True)
                    continue

                if train:
                    loss.backward()

                    if (step + 1) % accum_steps == 0:
                        torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
                        self.optimizer.step()
                        self.optimizer.zero_grad()

                total_loss += loss.item()

                # --- Record visualization snapshot during validation ---
                if not train:
                    if hasattr(self, "_viz_snapshots"):
                        # Keep only last 5 for memory efficiency
                        if len(self._viz_snapshots) > 5:
                            self._viz_snapshots.pop(0)
        
                        self._viz_snapshots.append({
                            "epoch": getattr(self, "_current_epoch", -1),
                            "target": target.detach().cpu(),
                            "pred": preds.detach().cpu(),
                            "input": data['ts_array'].detach().cpu(),
                        })

        return total_loss / len(loader)

    def evaluate(self, loader):
        """
        Evaluate the model on a dataset.
        
        Parameters:
        -----------
        loader : torch.utils.data.DataLoader
            DataLoader for the evaluation data.
            
        Returns:
        --------
        float
            Average loss on the evaluation data.
        """
        self.model.eval()
        total_loss = self._run_epoch(loader, train=False)
        print(f"Evaluation Loss: {total_loss:.6f}")
        return total_loss