# Copyright (c) 2026 EarthBridge Team.
# Credits: Built on open-source libraries and papers acknowledged in README.md citations.

"""Local Diffusion inference pipeline with branching and fusion.

Provides :class:`LocalDiffusionPipeline` for running hallucination-aware
inference.  The pipeline supports two modes:

1. **Standard mode** – standard DDPM/DDIM reverse-process sampling.
2. **Branch-and-fuse mode** – the input anomaly mask partitions the
   conditioning image into in-distribution (IND) and out-of-distribution
   (OOD) regions.  Two denoising branches run in parallel; at a fusion
   timestep they merge and a single unified denoising continues.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import torch
from PIL import Image

from src.models.local_diffusion import LocalDiffusionUNet
from src.schedulers.local_diffusion import LocalDiffusionScheduler


@dataclass
class LocalDiffusionPipelineOutput:
    """Output class for Local Diffusion pipeline.

    Attributes
    ----------
    images : list of PIL.Image.Image, np.ndarray, or torch.Tensor
        Generated images after translation.
    nfe : int
        Number of function evaluations (model forward passes).
    """

    images: Union[List[Image.Image], np.ndarray, torch.Tensor]
    nfe: int


class LocalDiffusionPipeline:
    """Inference pipeline for Local Diffusion with optional branching and fusion.

    The pipeline wraps a :class:`LocalDiffusionUNet` and a
    :class:`LocalDiffusionScheduler` with a consistent API for loading,
    preprocessing, and postprocessing.

    Example
    -------
    ::

        from src.models.local_diffusion import create_unet
        from src.schedulers.local_diffusion import LocalDiffusionScheduler
        from src.pipelines.local_diffusion import LocalDiffusionPipeline

        unet = create_unet(dim=32, channels=1)
        scheduler = LocalDiffusionScheduler(num_train_timesteps=250)
        pipeline = LocalDiffusionPipeline(unet=unet, scheduler=scheduler)
        output = pipeline(cond_image=my_tensor, output_type="pt")
    """

    def __init__(
        self,
        unet: LocalDiffusionUNet,
        scheduler: LocalDiffusionScheduler,
    ) -> None:
        self.unet = unet
        self.scheduler = scheduler

    @property
    def device(self) -> torch.device:
        return next(self.unet.parameters()).device

    @property
    def dtype(self) -> torch.dtype:
        return next(self.unet.parameters()).dtype

    def to(self, device: torch.device | str) -> "LocalDiffusionPipeline":
        self.unet = self.unet.to(device)
        return self

    def prepare_inputs(
        self,
        image: Union[torch.Tensor, Image.Image, List[Image.Image]],
        device: torch.device,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        """Prepare input images for the pipeline.

        Converts PIL images or numpy arrays to normalised tensors in
        ``[-1, 1]`` range.
        """
        if isinstance(image, Image.Image):
            image = [image]

        if isinstance(image, list) and isinstance(image[0], Image.Image):
            images = []
            for img in image:
                img_array = np.array(img, dtype=np.float32)
                if img_array.max() > 1.0:
                    img_array = img_array / 255.0
                if img_array.ndim == 2:
                    img_array = img_array[:, :, np.newaxis]
                img_tensor = torch.from_numpy(img_array).permute(2, 0, 1)
                images.append(img_tensor)
            image = torch.stack(images)

        if isinstance(image, np.ndarray):
            image = torch.from_numpy(image)

        if image.min() >= 0 and image.max() <= 1.0:
            image = image * 2 - 1
        elif image.max() > 1.0:
            image = image / 255.0 * 2 - 1

        return image.to(device=device, dtype=dtype)

    @torch.no_grad()
    def __call__(
        self,
        cond_image: Union[torch.Tensor, Image.Image, List[Image.Image]],
        anomaly_mask: Optional[torch.Tensor] = None,
        num_inference_steps: Optional[int] = None,
        use_ddim: bool = False,
        ddim_steps: Optional[int] = None,
        branch_out: bool = False,
        fusion_timestep: int = 2,
        clip_min: float = -1.0,
        clip_max: float = 1.0,
        output_type: str = "pil",
        return_dict: bool = True,
    ) -> Union[LocalDiffusionPipelineOutput, tuple]:
        """Generate translated images via diffusion sampling.

        Parameters
        ----------
        cond_image : torch.Tensor or PIL.Image.Image or list of PIL.Image.Image
            Conditioning image (e.g. degraded, noisy, or corrupted input).
        anomaly_mask : torch.Tensor, optional
            Binary mask indicating OOD regions (1 = anomalous, 0 = normal).
            Required when ``branch_out=True``.
        num_inference_steps : int, optional
            Number of DDPM steps.  Defaults to the scheduler's training steps.
        use_ddim : bool
            Use DDIM sampling instead of DDPM.
        ddim_steps : int, optional
            Number of DDIM steps (ignored when ``use_ddim=False``).
        branch_out : bool
            Enable branch-and-fuse mode for hallucination suppression.
        fusion_timestep : int
            Timestep at which to fuse the OOD and IND branches.
        clip_min, clip_max : float
            Clipping bounds for predicted x_0.
        output_type : str
            ``"pil"``, ``"np"``, or ``"pt"`` (default ``"pil"``).
        return_dict : bool
            If ``True``, return a :class:`LocalDiffusionPipelineOutput`.

        Returns
        -------
        LocalDiffusionPipelineOutput or tuple
            Translated images and the number of function evaluations.
        """
        device = self.device
        dtype = self.dtype

        cond = self.prepare_inputs(cond_image, device, dtype)
        bs, channels, h, w = cond.shape
        T = num_inference_steps if num_inference_steps is not None else self.scheduler.num_train_timesteps

        self.unet.eval()
        nfe = 0

        if use_ddim:
            images, nfe = self._ddim_sample(
                cond, anomaly_mask, ddim_steps or T, branch_out,
                fusion_timestep, clip_min, clip_max, device, dtype,
            )
        else:
            images, nfe = self._ddpm_sample(
                cond, anomaly_mask, T, branch_out,
                fusion_timestep, clip_min, clip_max, device, dtype,
            )

        images = images.clamp(clip_min, clip_max)

        if output_type == "pil":
            images = self._convert_to_pil(images)
        elif output_type == "np":
            images = self._convert_to_numpy(images)

        if not return_dict:
            return (images, nfe)
        return LocalDiffusionPipelineOutput(images=images, nfe=nfe)

    # ------------------------------------------------------------------
    # DDPM sampling loop
    # ------------------------------------------------------------------

    def _ddpm_sample(
        self,
        cond: torch.Tensor,
        anomaly_mask: Optional[torch.Tensor],
        T: int,
        branch_out: bool,
        fusion_timestep: int,
        clip_min: float,
        clip_max: float,
        device: torch.device,
        dtype: torch.dtype,
    ) -> tuple:
        bs, channels, h, w = cond.shape
        img = torch.randn(bs, channels, h, w, device=device, dtype=dtype)
        nfe = 0
        branching = branch_out and anomaly_mask is not None

        if branching:
            mask = anomaly_mask.to(device).float()
            mask_binary = (mask >= 1.0).float()
            img_out = torch.randn_like(img)
            img_in = torch.randn_like(img)

        for t_idx in reversed(range(0, T)):
            t = torch.full((bs,), t_idx, device=device, dtype=torch.long)

            if branching and t_idx > fusion_timestep:
                # Branch mode: separate OOD and IND denoising
                cond_out = cond * mask_binary
                cond_in = cond * (1.0 - mask_binary)

                out_pred = self.unet(img_out, cond_out, t)
                in_pred = self.unet(img_in, cond_in, t)
                nfe += 2

                step_out = self.scheduler.step(out_pred, t, img_out, clip_min, clip_max)
                step_in = self.scheduler.step(in_pred, t, img_in, clip_min, clip_max)

                img_out = step_out.prev_sample
                img_in = step_in.prev_sample

                if t_idx == fusion_timestep + 1:
                    # Fuse at next step
                    x_start_fused = (
                        step_out.pred_x_start * mask_binary
                        + step_in.pred_x_start * (1.0 - mask_binary)
                    )
                    img = torch.where(mask_binary.bool(), img_out, img_in)
                    branching = False
            else:
                # Standard denoising
                model_output = self.unet(img, cond, t)
                nfe += 1
                step_result = self.scheduler.step(model_output, t, img, clip_min, clip_max)
                img = step_result.prev_sample

        return img, nfe

    # ------------------------------------------------------------------
    # DDIM sampling loop
    # ------------------------------------------------------------------

    def _ddim_sample(
        self,
        cond: torch.Tensor,
        anomaly_mask: Optional[torch.Tensor],
        ddim_steps: int,
        branch_out: bool,
        fusion_timestep: int,
        clip_min: float,
        clip_max: float,
        device: torch.device,
        dtype: torch.dtype,
    ) -> tuple:
        bs, channels, h, w = cond.shape
        T = self.scheduler.num_train_timesteps

        # Build DDIM time schedule
        times = torch.linspace(-1, T - 1, steps=ddim_steps + 1)
        times = list(reversed(times.int().tolist()))
        time_pairs = list(zip(times[:-1], times[1:]))

        img = torch.randn(bs, channels, h, w, device=device, dtype=dtype)
        nfe = 0
        branching = branch_out and anomaly_mask is not None

        if branching:
            mask = anomaly_mask.to(device).float()
            mask_binary = (mask >= 1.0).float()
            img_out = torch.randn_like(img)
            img_in = torch.randn_like(img)

        for time_now, time_next in time_pairs:
            t = torch.full((bs,), time_now, device=device, dtype=torch.long)
            t_next = torch.full((bs,), time_next, device=device, dtype=torch.long)

            if branching and time_now > fusion_timestep:
                cond_out = cond * mask_binary
                cond_in = cond * (1.0 - mask_binary)

                out_pred = self.unet(img_out, cond_out, t)
                in_pred = self.unet(img_in, cond_in, t)
                nfe += 2

                step_out = self.scheduler.ddim_step(out_pred, t, t_next, img_out, clip_min, clip_max)
                step_in = self.scheduler.ddim_step(in_pred, t, t_next, img_in, clip_min, clip_max)

                img_out = step_out.prev_sample
                img_in = step_in.prev_sample

                if time_next <= fusion_timestep:
                    img = torch.where(mask_binary.bool(), img_out, img_in)
                    branching = False
            else:
                model_output = self.unet(img, cond, t)
                nfe += 1
                step_result = self.scheduler.ddim_step(model_output, t, t_next, img, clip_min, clip_max)
                img = step_result.prev_sample

        return img, nfe

    # ------------------------------------------------------------------
    # Output conversion helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _convert_to_pil(images: torch.Tensor) -> List[Image.Image]:
        """Convert tensor in [-1, 1] to PIL images."""
        images = (images + 1) / 2
        images = images.clamp(0, 1)
        images = images.cpu().permute(0, 2, 3, 1).numpy()
        images = (images * 255).round().astype(np.uint8)
        pil_images = []
        for img in images:
            if img.shape[2] == 1:
                pil_images.append(Image.fromarray(img.squeeze(2), mode="L"))
            else:
                pil_images.append(Image.fromarray(img))
        return pil_images

    @staticmethod
    def _convert_to_numpy(images: torch.Tensor) -> np.ndarray:
        """Convert tensor in [-1, 1] to numpy array in [0, 1]."""
        images = (images + 1) / 2
        images = images.clamp(0, 1)
        return images.cpu().permute(0, 2, 3, 1).numpy()
