import matplotlib.pyplot as plt
import scienceplots

plt.style.use(["science", "no-latex", "bright"])


class PlotCallback:
    """
    Callback for visualizing model training progress and predictions.
    
    This class provides visualization capabilities for different types of time series
    tasks including reconstruction, forecasting, and classification.
    
    Parameters:
    -----------
    every : int, default=5
        Plot every n epochs.
    show_recon : bool, default=False
        If True, visualize reconstructions.
    show_forecast : bool, default=False
        If True, visualize forecasting predictions.
    show_class : bool, default=False
        If True, visualize classification results (bar plot per class).
    n_samples : int, default=3
        Number of samples to plot.
    channels : int or list, default=0
        Channel(s) to visualize.
    """
    def __init__(
        self,
        every=5,
        show_recon=False,
        show_forecast=False,
        show_class=False,
        n_samples=3,
        channels=0,
    ):
        """
        Args:
            every: plot every n epochs
            show_recon: if True, visualize reconstructions
            show_forecast: if True, visualize forecasting predictions
            show_class: if True, visualize classification results (bar plot per class)
            n_samples: number of samples to plot
            channels: int or list of channels to visualize
        """
        self.every = every
        self.show_recon = show_recon
        self.show_forecast = show_forecast
        self.show_class = show_class
        self.n_samples = n_samples
        self.channels = channels if isinstance(channels, list) else [channels]

    def __call__(self, epoch, trainer):
        """
        Execute the callback at the end of an epoch.
        
        Parameters:
        -----------
        epoch : int
            Current epoch number.
        trainer : Trainer
            Trainer instance containing model and visualization data.
        """
        if epoch % self.every != 0:
            return

        self.plot_losses(trainer)

        # Visualization depending on model type
        if self.show_recon and trainer._viz_snapshots:
            self.plot_reconstructions(trainer)

        if self.show_forecast and trainer._viz_snapshots:
            self.plot_forecasts(trainer)

        if self.show_class and trainer._viz_snapshots:
            self.plot_classification(trainer)

    # ------------------------------------------------------------------ #
    # Loss curves
    # ------------------------------------------------------------------ #
    def plot_losses(self, trainer):
        """
        Plot training and validation loss curves.
        
        Parameters:
        -----------
        trainer : Trainer
            Trainer instance containing loss history.
        """
        plt.figure(figsize=(6, 4))
        plt.plot(trainer.train_losses, label="Train")
        if trainer.val_losses:
            plt.plot(trainer.val_losses, label="Val")
        plt.xlabel("Epoch")
        plt.ylabel("Loss")
        plt.title("Training Curves")
        plt.legend()
        plt.grid(alpha=0.3)
        plt.tight_layout()
        plt.show()

    # ------------------------------------------------------------------ #
    # Reconstruction visualization
    # ------------------------------------------------------------------ #
    def plot_reconstructions(self, trainer):
        """
        Plot original time series and their reconstructions.
        
        Parameters:
        -----------
        trainer : Trainer
            Trainer instance containing visualization snapshots.
        """
        snap = trainer._viz_snapshots[-1]
        tgt = snap["target"][: self.n_samples]
        rec = snap["pred"][: self.n_samples]
        T = tgt.size(1)
        C = len(self.channels)

        fig, axes = plt.subplots(C, self.n_samples, figsize=(5 * self.n_samples, 3 * C), squeeze=False)
        for i in range(self.n_samples):
            for j, ch in enumerate(self.channels):
                ax = axes[j, i]
                ax.plot(range(T), tgt[i, :, ch], label="Target")
                ax.plot(range(T), rec[i, :, ch], linestyle="--", label="Recon")
                ax.set_title(f"Epoch {snap['epoch']} — Sample {i} — ch {ch}")
                ax.legend()
                ax.grid(alpha=0.3)
        plt.tight_layout()
        plt.show()

    # ------------------------------------------------------------------ #
    # Forecast visualization
    # ------------------------------------------------------------------ #
    def plot_forecasts(self, trainer):
        """
        Plot input time series and their forecasted continuations.
        
        Parameters:
        -----------
        trainer : Trainer
            Trainer instance containing visualization snapshots.
        """
        snap = trainer._viz_snapshots[-1]
        tgt = snap["target"][: self.n_samples]   # [B, horizon, C]
        pred = snap["pred"][: self.n_samples]    # [B, horizon, C]
        x_in = snap["input"][: self.n_samples]    # [B, T, C]
        T_in = x_in.size(1)
        H = tgt.size(1)
        C = len(self.channels)

        fig, axes = plt.subplots(C, self.n_samples, figsize=(5 * self.n_samples, 3 * C), squeeze=False)
        for i in range(self.n_samples):
            for j, ch in enumerate(self.channels):
                ax = axes[j, i]
                ax.plot(range(-T_in, 0), x_in[i, :, ch], label="Input", color="gray")
                ax.plot(range(H), tgt[i, :, ch], label="Target")
                ax.plot(range(H), pred[i, :, ch], linestyle="--", label="Pred")
                ax.set_title(f"Epoch {snap['epoch']} — Sample {i} — ch {ch}")
                ax.legend()
                ax.grid(alpha=0.3)
        plt.tight_layout()
        plt.show()

    # ------------------------------------------------------------------ #
    # Classification visualization
    # ------------------------------------------------------------------ #
    def plot_classification(self, trainer):
        """
        Plot classification results as bar charts of class probabilities.
        
        Parameters:
        -----------
        trainer : Trainer
            Trainer instance containing visualization snapshots.
        """
        snap = trainer._viz_snapshots[-1]
        logits = snap["pred"][: self.n_samples]   # [B, num_classes]
        labels = snap["target"][: self.n_samples]   # [B]

        fig, axes = plt.subplots(1, self.n_samples, figsize=(4 * self.n_samples, 4), squeeze=False)
        for i in range(self.n_samples):
            ax = axes[0, i]
            ax.bar(range(logits.size(1)), logits[i].detach().cpu().numpy(), alpha=0.7)
            ax.axhline(0, color="k", linewidth=0.5)
            ax.set_title(f"Sample {i} — True: {labels[i].item()}")
            ax.set_xlabel("Class")
            if i == 0:
                ax.set_ylabel("logits") 
        plt.tight_layout()
        plt.show()
