"""
Diffusion model adapter — load and merge Stable Diffusion, SDXL components.

Supports:
- Loading UNet, VAE, and text encoder components
- Merging UNets from different models using ModelOsmosis strategies
- Loading from HuggingFace Hub or local paths

Usage::

    from model_osmosis.adapters.diffusion import DiffusionAdapter

    adapter = DiffusionAdapter()

    # Load UNet as expert
    unet = adapter.load_unet("stabilityai/stable-diffusion-2-1")

    # Load text encoder
    text_enc = adapter.load_text_encoder("stabilityai/stable-diffusion-2-1")

    # Merge UNets from multiple models
    merged = adapter.merge_unets(
        ["model-1", "model-2"],
        weights=[0.6, 0.4],
        strategy="slerp",
    )
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Union

import torch
import torch.nn as nn
import numpy as np


class DiffusionUNetExpert(nn.Module):
    """Wrapped diffusion UNet as stackable expert.

    The UNet is the core of diffusion models — it predicts noise
    to be removed at each denoising step.

    Args:
        unet: The UNet2DConditionModel.
        name: Human-readable name.
        cross_attention_dim: Dimension of cross-attention (text conditioning).
        in_channels: Number of input channels (latent space).
    """

    def __init__(
        self,
        unet: nn.Module,
        name: str = "unet_expert",
        cross_attention_dim: int = 768,
        in_channels: int = 4,
    ) -> None:
        super().__init__()
        self.unet = unet
        self.name = name
        self.cross_attention_dim = cross_attention_dim
        self.in_channels = in_channels
        self.feature_dim = cross_attention_dim
        self.hidden_size = cross_attention_dim

        # Freeze weights
        for param in self.unet.parameters():
            param.requires_grad = False

    def forward(
        self,
        latent: torch.Tensor,
        timestep: Union[torch.Tensor, int],
        encoder_hidden_states: torch.Tensor,
        **kwargs,
    ) -> torch.Tensor:
        """Forward pass through UNet.

        Args:
            latent: Noisy latent (batch, channels, height, width).
            timestep: Denoising timestep (batch,) or int.
            encoder_hidden_states: Text embeddings (batch, seq_len, hidden).

        Returns:
            Predicted noise (batch, channels, height, width).
        """
        with torch.no_grad():
            output = self.unet(
                latent,
                timestep,
                encoder_hidden_states=encoder_hidden_states,
                **kwargs,
            )
        return output.sample

    def extra_repr(self) -> str:
        return (
            f"name={self.name}, cross_attention_dim={self.cross_attention_dim}, "
            f"in_channels={self.in_channels}"
        )


class TextEncoderExpert(nn.Module):
    """Wrapped text encoder as stackable expert.

    Args:
        encoder: The text encoder model (CLIP or similar).
        name: Human-readable name.
        hidden_size: Output hidden dimension.
    """

    def __init__(
        self,
        encoder: nn.Module,
        name: str = "text_encoder",
        hidden_size: int = 768,
    ) -> None:
        super().__init__()
        self.encoder = encoder
        self.name = name
        self.hidden_size = hidden_size
        self.feature_dim = hidden_size

        for param in self.encoder.parameters():
            param.requires_grad = False

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Encode text to embeddings.

        Args:
            input_ids: Token IDs (batch, seq_len).
            attention_mask: Optional attention mask.

        Returns:
            Text embeddings (batch, seq_len, hidden_size).
        """
        with torch.no_grad():
            outputs = self.encoder(
                input_ids=input_ids,
                attention_mask=attention_mask,
            )
        if hasattr(outputs, "last_hidden_state"):
            return outputs.last_hidden_state
        return outputs[0]

    def extra_repr(self) -> str:
        return f"name={self.name}, hidden_size={self.hidden_size}"


class DiffusionAdapter:
    """Load and merge diffusion model components.

    Args:
        device: Device to load models to.
        torch_dtype: Data type (torch.float16 recommended for diffusion).
        verbose: Print loading progress.
    """

    def __init__(
        self,
        device: str = "cpu",
        torch_dtype: Optional[torch.dtype] = None,
        verbose: bool = True,
    ) -> None:
        self.device = device
        self.torch_dtype = torch_dtype or torch.float16
        self.verbose = verbose

        # Check diffusers availability
        try:
            import diffusers
            self._diffusers_available = True
        except ImportError:
            self._diffusers_available = False

    def _check_diffusers(self) -> None:
        """Check if diffusers is available."""
        if not self._diffusers_available:
            raise ImportError(
                "Diffusion adapter requires `diffusers`. "
                "Install with: pip install diffusers"
            )

    def load_unet(
        self,
        model_name: str,
        name: Optional[str] = None,
        subfolder: str = "unet",
        variant: Optional[str] = None,
    ) -> DiffusionUNetExpert:
        """Load a UNet component from a diffusion model.

        Args:
            model_name: HuggingFace model ID or local path.
            name: Human-readable name.
            subfolder: Subfolder containing UNet (usually "unet").
            variant: Model variant (e.g., "fp16").

        Returns:
            DiffusionUNetExpert ready for use.
        """
        self._check_diffusers()

        from diffusers import UNet2DConditionModel

        if self.verbose:
            print(f"[Diffusion] Loading UNet from '{model_name}'...")

        load_kwargs = {
            "torch_dtype": self.torch_dtype,
        }
        if variant:
            load_kwargs["variant"] = variant

        unet = UNet2DConditionModel.from_pretrained(
            model_name,
            subfolder=subfolder,
            **load_kwargs,
        )
        unet = unet.to(self.device)
        unet.eval()

        if self.verbose:
            total_params = sum(p.numel() for p in unet.parameters())
            print(f"[Diffusion] Loaded UNet: {total_params:,} params")

        return DiffusionUNetExpert(
            unet=unet,
            name=name or f"{model_name.split('/')[-1]}_unet",
            cross_attention_dim=unet.config.cross_attention_dim,
            in_channels=unet.config.in_channels,
        )

    def load_text_encoder(
        self,
        model_name: str,
        name: Optional[str] = None,
        subfolder: str = "text_encoder",
    ) -> TextEncoderExpert:
        """Load a text encoder from a diffusion model.

        Args:
            model_name: HuggingFace model ID or local path.
            name: Human-readable name.
            subfolder: Subfolder containing text encoder.

        Returns:
            TextEncoderExpert ready for use.
        """
        self._check_diffusers()

        from transformers import CLIPTextModel

        if self.verbose:
            print(f"[Diffusion] Loading text encoder from '{model_name}'...")

        encoder = CLIPTextModel.from_pretrained(
            model_name,
            subfolder=subfolder,
            torch_dtype=self.torch_dtype,
        )
        encoder = encoder.to(self.device)
        encoder.eval()

        if self.verbose:
            total_params = sum(p.numel() for p in encoder.parameters())
            print(f"[Diffusion] Loaded text encoder: {total_params:,} params")

        return TextEncoderExpert(
            encoder=encoder,
            name=name or f"{model_name.split('/')[-1]}_text_encoder",
            hidden_size=encoder.config.hidden_size,
        )

    def load_vae(
        self,
        model_name: str,
        name: Optional[str] = None,
        subfolder: str = "vae",
    ) -> nn.Module:
        """Load a VAE from a diffusion model.

        Args:
            model_name: HuggingFace model ID or local path.
            name: Human-readable name.
            subfolder: Subfolder containing VAE.

        Returns:
            AutoencoderKL model.
        """
        self._check_diffusers()

        from diffusers import AutoencoderKL

        if self.verbose:
            print(f"[Diffusion] Loading VAE from '{model_name}'...")

        vae = AutoencoderKL.from_pretrained(
            model_name,
            subfolder=subfolder,
            torch_dtype=self.torch_dtype,
        )
        vae = vae.to(self.device)
        vae.eval()

        if self.verbose:
            total_params = sum(p.numel() for p in vae.parameters())
            print(f"[Diffusion] Loaded VAE: {total_params:,} params")

        return vae

    def merge_unets(
        self,
        model_names: List[str],
        weights: Optional[List[float]] = None,
        strategy: str = "weighted_average",
        name: str = "merged_unet",
    ) -> DiffusionUNetExpert:
        """Merge UNets from multiple diffusion models.

        Args:
            model_names: List of HuggingFace model IDs.
            weights: Per-model weights (default: equal).
            strategy: Merge strategy ('weighted_average', 'slerp', 'ties').
            name: Name for the merged model.

        Returns:
            DiffusionUNetExpert with merged weights.
        """
        self._check_diffusers()

        from diffusers import UNet2DConditionModel

        if weights is None:
            weights = [1.0 / len(model_names)] * len(model_names)
        else:
            total = sum(weights)
            weights = [w / total for w in weights]

        if self.verbose:
            print(f"[Diffusion] Merging {len(model_names)} UNets with {strategy}...")
            for m, w in zip(model_names, weights):
                print(f"  - {m}: {w:.3f}")

        # Load all UNets
        unets = []
        for model_name in model_names:
            unet = UNet2DConditionModel.from_pretrained(
                model_name,
                subfolder="unet",
                torch_dtype=torch.float32,  # Use float32 for merging
            )
            unets.append(unet)

        # Get base config
        base_config = unets[0].config

        # Merge weights
        merged_state = {}
        base_state = unets[0].state_dict()

        for key in base_state.keys():
            tensors = [u.state_dict()[key].float() for u in unets]

            if strategy == "weighted_average":
                merged_state[key] = sum(w * t for w, t in zip(weights, tensors))

            elif strategy == "slerp":
                # Pairwise SLERP
                result = tensors[0]
                cumulative_weight = weights[0]
                for i in range(1, len(tensors)):
                    t = cumulative_weight / (cumulative_weight + weights[i])
                    result = self._slerp_tensor(result, tensors[i], t)
                    cumulative_weight += weights[i]
                merged_state[key] = result

            elif strategy == "ties":
                # TIES: Trim, Elect sign, Merge
                stacked = torch.stack(tensors)
                # Trim small values
                abs_vals = stacked.abs()
                threshold = abs_vals.quantile(0.5)
                trimmed = stacked.clone()
                trimmed[abs_vals < threshold] = 0
                # Elect sign
                signs = torch.sign(trimmed)
                elected_sign = torch.sign(signs.sum(dim=0))
                elected_sign[elected_sign == 0] = 1
                # Merge agreeing values
                mask = signs == elected_sign.unsqueeze(0)
                masked = trimmed * mask.float()
                weight_tensor = torch.tensor(weights).view(-1, *([1] * (stacked.dim() - 1)))
                weighted = masked * weight_tensor
                count = mask.float().sum(dim=0).clamp(min=1)
                merged_state[key] = weighted.sum(dim=0) / count

            else:
                # Default to weighted average
                merged_state[key] = sum(w * t for w, t in zip(weights, tensors))

            # Convert back to target dtype
            merged_state[key] = merged_state[key].to(self.torch_dtype)

        # Create merged UNet
        merged_unet = UNet2DConditionModel.from_config(base_config)
        merged_unet.load_state_dict(merged_state)
        merged_unet = merged_unet.to(self.device)
        merged_unet.eval()

        # Clean up
        del unets

        if self.verbose:
            total_params = sum(p.numel() for p in merged_unet.parameters())
            print(f"[Diffusion] Created merged UNet: {total_params:,} params")

        return DiffusionUNetExpert(
            unet=merged_unet,
            name=name,
            cross_attention_dim=base_config.cross_attention_dim,
            in_channels=base_config.in_channels,
        )

    def _slerp_tensor(
        self, a: torch.Tensor, b: torch.Tensor, t: float
    ) -> torch.Tensor:
        """Spherical linear interpolation."""
        a_flat = a.flatten().float()
        b_flat = b.flatten().float()

        a_norm = a_flat / (a_flat.norm() + 1e-8)
        b_norm = b_flat / (b_flat.norm() + 1e-8)

        dot = torch.clamp(torch.dot(a_norm, b_norm), -1.0, 1.0)
        omega = torch.acos(dot)

        if omega.abs() < 1e-6:
            return (1 - t) * a + t * b

        sin_omega = torch.sin(omega)
        result_flat = (
            torch.sin((1 - t) * omega) / sin_omega * a_flat +
            torch.sin(t * omega) / sin_omega * b_flat
        )

        result_norm = (1 - t) * a_flat.norm() + t * b_flat.norm()
        result_flat = result_flat / (result_flat.norm() + 1e-8) * result_norm

        return result_flat.view(a.shape)
