import contextlib
import logging
import os
import string
import time
import warnings
from copy import deepcopy
from datetime import datetime
from pathlib import Path
from typing import Any

import matplotlib
import pandas as pd
import torch
from PIL import Image
from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.callbacks import Callback
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.utilities import rank_zero_only
from torch import nn

from venturi.config import Config, instantiate

try:
    import wandb
except ImportError:
    _has_wandb = False
else:
    _has_wandb = True

if os.environ.get("DISPLAY", "") == "":
    # Use non-interactive Agg backend
    matplotlib.use("Agg")
import matplotlib.pyplot as plt


class LossCollection(nn.Module):
    """Container class to combine multiple loss functions.

    The constructor takes a Config object with the format:

    loss1_name:
        instance:
            _target_: path.to.LossClass
            arg1: value1
            ...
        loss_weight: float
    loss2_name:
        ...

    Each call to the forward method generates the total loss and, if return_logs is True, a
    dictionary with the individual loss values with the format {loss_name: loss_value}.
    """

    def __init__(self, cfg_loss: Config, return_logs: bool = True, prefix: str = ""):
        """Args:
        cfg_loss: Configuration dictionary for the loss functions.
        return_logs: Whether to return individual loss values as logs.
        prefix: Prefix to add to the loss names in the logs.
        """
        super().__init__()

        # If there is a single loss without weight, set its weight to 1.0
        if len(cfg_loss) == 1:
            name, config = next(iter(cfg_loss.items()))
            if "loss_weight" not in config:
                cfg_loss[name]["loss_weight"] = 1.0

        self.loss_map = nn.ModuleDict()
        self.cfg_loss = cfg_loss
        self.return_logs = return_logs
        self.weights = {}

        # Register components
        for name, config in cfg_loss.items():
            self.loss_map[f"{prefix}{name}"] = instantiate(config["instance"])
            self.weights[f"{prefix}{name}"] = config["loss_weight"]

    def forward(self, input, target):
        """Passes the input and target to every child loss.

        Args:
            input: Output from the network.
            target: Ground truth data.
        """
        total_loss = 0.0
        logs = {}

        for name, loss_fn in self.loss_map.items():
            weight = self.weights[name]

            val = loss_fn(input, target)

            total_loss += weight * val
            logs[name] = val.detach()

        output = (total_loss, logs) if self.return_logs else total_loss

        return output

    def clone(self, prefix: str = "") -> "LossCollection":
        """Make a copy of the class."""
        return self.__class__(deepcopy(self.cfg_loss), self.return_logs, prefix)


class PlottingCallback(Callback):
    """Reads the CSV log generated by Lightning and creates a static plot (plots.png) at the
    end of every training epoch.
    """

    @rank_zero_only
    def on_train_epoch_end(self, trainer, pl_module):
        if trainer.log_dir is None:
            raise ValueError("Trainer log_dir is None; cannot save plots.")
        logger_dir = Path(trainer.log_dir)
        metrics_file = logger_dir / "metrics.csv"

        if not metrics_file.exists():
            raise ValueError(
                f"Metrics file not found: {metrics_file}. Maybe you forgot to enable CSVLogger?"
            )

        try:
            try:
                df = pd.read_csv(metrics_file)
            except Exception:
                # There might be racing conditions in multi-GPU training. If reading fails,
                # we skip plotting for this epoch.
                return

            fig, ax = plt.subplots(1, 2, figsize=(15, 5))

            # Loss Plot
            args_p = pl_module.args.logging.plot.left_plot
            has_line = False
            for key in args_p.metrics:  
                key_p = key
                if key_p == "train/loss":
                    # Lightning automatically adds 'step' or 'epoch' suffixes to train loss
                    key_p = "train/loss_epoch"
                if key_p in df.columns:
                    data = df[["epoch", key_p]].dropna()
                    if len(data) > 1:
                        ax[0].plot(data["epoch"], data[key_p], label=key)
                        has_line = True

            ax[0].set_xlabel("epoch")
            if has_line:
                ax[0].legend()
            ax[0].set_ylim(args_p.ylim.min, args_p.ylim.max)

            # Metrics Plot
            args_p = pl_module.args.logging.plot.right_plot
            has_line = False
            for key in args_p.metrics:  
                if key in df.columns:
                    data = df[["epoch", key]].dropna()
                    if len(data) > 1:
                        ax[1].plot(data["epoch"], data[key], label=key)
                        has_line = True

            ax[1].set_xlabel("epoch")
            if has_line:
                ax[1].legend()
            ax[1].set_ylim(args_p.ylim.min, args_p.ylim.max)
            # Save
            plt.tight_layout()
            plt.savefig(logger_dir / "plots.png")
            plt.close(fig)

        except Exception as e:
            print(f"Error generating plot: {e}")


class ImageSaveCallback(Callback):
    """Saves validation predictions to disk."""

    def __init__(
        self,
        run_path: str | Path,
        val_img_indices: list[int],
        log_disk: bool = True,
        mean: torch.Tensor | float | None = None,
        std: torch.Tensor | float | None = None,
        log_wandb: bool = False,
        class_labels: list[str] | None = None,
    ):
        """Args:
        run_path: Path to the run folder where images will be saved.
        val_img_indices: List of image indices to log.
        log_disk: Whether to save images to disk.
        mean: Mean used for normalization (for denormalizing images).
        std: Std used for normalization (for denormalizing images).
        log_wandb: Whether to also log images to wandb.
        class_labels: List of class labels for segmentation masks. Only used by wandb logging.
        """
        super().__init__()

        run_path = Path(run_path)

        self.run_path = run_path
        self.val_img_indices = set(val_img_indices)
        self.log_disk = log_disk
        self.log_wandb = log_wandb
        self.mean = mean
        self.std = std
        self.class_labels = class_labels

    @rank_zero_only
    def on_fit_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
        if self.log_disk:
            # Create folders
            for idx in self.val_img_indices:
                save_dir = self.run_path / "images" / f"image_{idx}"
                save_dir.mkdir(parents=True, exist_ok=True)

    # Only the main process logs images in multi-GPU training, which will miss some
    # indices but avoids complications.
    @rank_zero_only
    def on_validation_batch_end(
        self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0 
    ):
        inputs, targets = batch
        batch_size = inputs.size(0)
        start_idx = batch_idx * batch_size
        end_idx = start_idx + batch_size

        # Check if any target index is in this batch
        current_indices = set(range(start_idx, end_idx))
        intersect = self.val_img_indices.intersection(current_indices)

        if not intersect:
            return

        with torch.no_grad():
            logits = outputs["logits"]

            # Apply argmax for multiclass or sigmoid>0.5 for binary
            if logits.shape[1] == 1:
                preds = (torch.sigmoid(logits) > 0.5).float()
            else:
                preds = torch.argmax(logits, dim=1).unsqueeze(1).float()

        for global_idx in intersect:
            local_idx = global_idx - start_idx

            img_tensor = inputs[local_idx].cpu()
            target_tensor = targets[local_idx].cpu()
            pred_tensor = preds[local_idx].cpu()

            img_prep = img_tensor

            # Denormalize image if necessary
            if self.mean is not None and self.std is not None:
                mean = torch.as_tensor(self.mean, device=img_prep.device)
                std = torch.as_tensor(self.std, device=img_prep.device)
                if mean.ndim == 0:
                    mean = mean.view(1, 1, 1)
                else:
                    mean = mean.view(-1, 1, 1)

                if std.ndim == 0:
                    std = std.view(1, 1, 1)
                else:
                    std = std.view(-1, 1, 1)
                if img_prep.shape[0] == 1 and (mean.shape[0] > 1 or std.shape[0] > 1):
                    raise ValueError("Image has 1 channel but mean and std have multiple values.")

                img_prep = img_prep * std + mean

            # Convert to HWC and scale to [0,255]
            if img_prep.ndim == 3 and img_prep.shape[0] in [1, 3, 4]:
                img_prep = img_prep.permute(1, 2, 0).squeeze()  # C,H,W -> H,W,C

            if img_prep.min() >= 0.0 and img_prep.max() <= 1.0:
                img_prep = img_prep * 255.0
            elif img_prep.min() >= -1.0 and img_prep.max() <= 1.0:
                img_prep = (img_prep + 1.0) * 127.5
            elif img_prep.min() < 0.0 or img_prep.max() > 255:
                min = img_prep.min()
                img_prep = (img_prep - min) / (img_prep.max() - min) * 255.0
            img_prep = img_prep.clamp(0, 255)
            target_tensor = target_tensor.squeeze()
            pred_tensor = pred_tensor.squeeze()

            if self.log_disk:
                self._save_to_disk(trainer, global_idx, img_prep, target_tensor, pred_tensor)
            if self.log_wandb:
                self._log_to_wandb(trainer, global_idx, img_prep, target_tensor, pred_tensor)

    def _save_to_disk(self, trainer, idx, img, target, pred):
        # Scale target and pred to [0,255]
        if target.max() > 0:
            target = target.float()
            target = target / target.max() * 255
        if pred.max() > 0:
            pred = pred.float()
            pred = pred / pred.max() * 255
        target = target.byte().numpy()
        pred = pred.byte().numpy()
        img = img.byte().numpy()

        img_pil = Image.fromarray(img)
        target_pil = Image.fromarray(target)
        pred_pil = Image.fromarray(pred)
        # Create image with 3 panels: input, target, prediction
        w, h = img_pil.size
        combined = Image.new("RGB", (w * 3, h))
        combined.paste(img_pil.convert("RGB"), (0, 0))
        combined.paste(target_pil.convert("RGB"), (w, 0))
        combined.paste(pred_pil.convert("RGB"), (w * 2, 0))

        filename = f"epoch_{trainer.current_epoch}.png"
        combined.save(self.run_path / "images" / f"image_{idx}" / filename)

    def _log_to_wandb(self, trainer, idx, img, target, pred):
        """Logs an interactive segmentation mask."""

        if not _has_wandb:
            raise ImportError("wandb is not installed. Cannot log images to wandb.")

        masks = {
            "predictions": {"mask_data": pred.byte().numpy()},
            "ground_truth": {"mask_data": target.byte().numpy()},
        }
        class_labels = self.class_labels
        if class_labels is not None:
            class_labels_dict = dict(enumerate(class_labels))
            masks["predictions"]["class_labels"] = class_labels_dict
            masks["ground_truth"]["class_labels"] = class_labels_dict

        # Create the WandB Image object
        wandb_image = wandb.Image(
            img.byte().numpy(),
            masks=masks,
            caption=f"Image {idx}",
        )

        for logger in trainer.loggers:
            if isinstance(logger, WandbLogger):
                logger.experiment.log(
                    {f"val_predictions/image_{idx}": [wandb_image]}, step=trainer.global_step
                )


class TimeLoggerCallback(Callback):
    """Logs the time taken for each training and validation epoch."""

    @rank_zero_only
    def on_train_epoch_start(self, trainer, pl_module):
        self.train_epoch_start = time.time()

    @rank_zero_only
    def on_train_epoch_end(self, trainer, pl_module):
        epoch_time = time.time() - self.train_epoch_start
        pl_module.log("train/epoch_time", epoch_time, on_epoch=True)

    @rank_zero_only
    def on_validation_epoch_start(self, trainer, pl_module):
        self.val_epoch_start = time.time()

    @rank_zero_only
    def on_validation_epoch_end(self, trainer, pl_module):
        if trainer.sanity_checking:
            return
        epoch_time = time.time() - self.val_epoch_start
        pl_module.log("val/epoch_time", epoch_time, on_epoch=True)


class TrainingTimeLoggerCallback(Callback):
    """Logs the time taken for a complete training run."""

    def __init__(self, run_path: str | Path):
        super().__init__()
        self.run_path = Path(run_path)

    @rank_zero_only
    def on_train_start(self, trainer, pl_module):
        now = datetime.now()
        dt_string = now.strftime("%Y-%m-%d %H:%M:%S")
        with open(self.run_path / "training_time.txt", "w") as f:
            f.write(f"Training started at: {dt_string}\n")

    @rank_zero_only
    def on_train_end(self, trainer, pl_module):
        now = datetime.now()
        dt_string = now.strftime("%Y-%m-%d %H:%M:%S")
        with open(self.run_path / "training_time.txt", "a") as f:
            f.write(f"Training ended at: {dt_string}\n")


class LightningLogFilter(logging.Filter):
    """Filter for annoying Lightning messages."""

    def filter(self, record):
        msg = str(record.msg)

        forbidden_phrases = [
            "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES",
        ]

        # Verify if any forbidden phrase is in the message
        return not any(phrase in msg for phrase in forbidden_phrases)


def is_rank_zero():
    """Check if the current process is rank zero."""
    return int(os.environ.get("RANK", "0")) == 0


def silence_lightning():
    """Silence annoying Lightning messages."""

    warnings.filterwarnings("ignore", module="lightning")
    warnings.filterwarnings("ignore", message=".*does not have many workers.*")
    warnings.filterwarnings("ignore", ".*The number of training batches.*")
    warnings.filterwarnings(
        "ignore", message=".*isinstance.treespec, LeafSpec. is deprecated.*", category=FutureWarning
    )

    lightning_filter = LightningLogFilter()

    # Apply the filter to the root of Lightning
    logging.getLogger("lightning.pytorch").addFilter(lightning_filter)
    logging.getLogger("lightning.fabric").addFilter(lightning_filter)

    # The above is not enough, so we iterate over all active loggers
    for logger_name in logging.Logger.manager.loggerDict:
        if "lightning" in logger_name:
            logger = logging.getLogger(logger_name)
            if isinstance(logger, logging.Logger):
                logger.addFilter(lightning_filter)
                logger.setLevel(logging.ERROR)


def find_key_recursive(data: Config, target_key: str) -> Any:
    """Recursively searches for a key in a nested dictionary.
    Returns the value of the first match found.
    """
    if target_key in data:
        return data[target_key]

    for value in data.values():
        if isinstance(value, Config):
            result = find_key_recursive(value, target_key)
            if result is not None:
                return result

    return None


def generate_name_from_config(config: Config, template: str) -> str:
    """Fills a template string using values found recursively in a Config.

    Args:
        config: The configuration object.
        template: A string like "model_{arch}_lr{lr:.4f}".

    Returns:
        The formatted string.

    Raises:
        KeyError: If a key required by the template is not found in the config.
    """
    # Extract all variable names from the template (e.g., "lr", "bs")
    # Formatter.parse returns tuples: (literal_text, field_name, format_spec, conversion)
    needed_keys = {
        field_name for _, field_name, _, _ in string.Formatter().parse(template) if field_name
    }

    # Find values for these keys
    context = {}
    for key in needed_keys:
        value = find_key_recursive(config, key)

        if value is None:
            raise KeyError(f"Template requires key '{key}', but it was not found in the config.")

        context[key] = value

    # Format the string using the found values
    return template.format(**context)


def delete_wandb_run(project_name: str, run_path: str):
    """Delete a wandb run."""

    api = wandb.Api()
    runs = api.runs(project_name, {"config.logging.run_path": run_path})
    with contextlib.suppress(ValueError, TypeError):
        # Get number of runs to verify if there are any
        len(runs)

    for run in runs:
        run.delete()


def get_next_name(path: Path) -> Path:
    """If 'path' exists, appends _2, _3, etc. until a free name is found."""

    directory = path.parent
    base_name = path.stem
    if not path.exists():
        return path

    counter = 2
    while True:
        candidate = directory / f"{base_name}_{counter}{path.suffix}"
        if not candidate.exists():
            return candidate
        counter += 1
