"""
LoRA adapter — load and merge PEFT/LoRA adapters.

Supports:
- Loading a base model with a single LoRA adapter
- Merging multiple LoRA adapters into one model
- Using ModelOsmosis merge strategies (weighted_average, slerp, ties) on LoRA deltas

Usage::

    from model_osmosis.adapters.lora import LoRAAdapter

    adapter = LoRAAdapter()

    # Load base model with single LoRA
    model = adapter.load_with_lora(
        base_model="meta-llama/Llama-2-7b-hf",
        lora_path="username/my-lora-adapter",
        merge_lora=True,  # Merge LoRA weights into base
    )

    # Merge multiple LoRAs
    model = adapter.load_with_multiple_loras(
        base_model="meta-llama/Llama-2-7b-hf",
        lora_paths=["lora1", "lora2", "lora3"],
        lora_weights=[0.5, 0.3, 0.2],
        merge_strategy="ties",
    )
"""

from __future__ import annotations

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

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


class LoRAExpert(nn.Module):
    """Wrapped LoRA-adapted model as stackable expert.

    Args:
        model: The merged model (base + LoRA).
        hidden_size: Model's hidden dimension.
        name: Human-readable name.
        base_model_name: Original base model ID.
        lora_paths: LoRA adapters that were applied.
    """

    def __init__(
        self,
        model: nn.Module,
        hidden_size: int,
        name: str = "lora_expert",
        base_model_name: Optional[str] = None,
        lora_paths: Optional[List[str]] = None,
    ) -> None:
        super().__init__()
        self.model = model
        self.hidden_size = hidden_size
        self.name = name
        self.base_model_name = base_model_name
        self.lora_paths = lora_paths or []

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass."""
        if x.dim() == 2:
            x = x.unsqueeze(1)

        with torch.no_grad():
            outputs = self.model(inputs_embeds=x)

        if hasattr(outputs, "last_hidden_state"):
            hidden = outputs.last_hidden_state
        elif hasattr(outputs, "hidden_states") and outputs.hidden_states:
            hidden = outputs.hidden_states[-1]
        else:
            hidden = outputs[0] if isinstance(outputs, tuple) else outputs

        if hidden.dim() == 3:
            hidden = hidden[:, -1, :]

        return hidden

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


class LoRAAdapter:
    """Load and merge PEFT/LoRA adapters.

    Args:
        device: Device to load models to.
        torch_dtype: Data type for model weights.
        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
        self.verbose = verbose

        # Check PEFT availability
        try:
            from peft import PeftModel, PeftConfig
            self._peft_available = True
        except ImportError:
            self._peft_available = False

    def _check_peft(self) -> None:
        """Check if PEFT is available."""
        if not self._peft_available:
            raise ImportError(
                "PEFT library required for LoRA support. "
                "Install with: pip install peft"
            )

    def load_with_lora(
        self,
        base_model: str,
        lora_path: str,
        merge_lora: bool = True,
        name: Optional[str] = None,
    ) -> LoRAExpert:
        """Load a base model with a LoRA adapter.

        Args:
            base_model: HuggingFace model ID for base model.
            lora_path: Path to LoRA adapter (HF Hub or local).
            merge_lora: If True, merge LoRA weights into base model.
            name: Optional name for the expert.

        Returns:
            LoRAExpert ready for stacking.
        """
        self._check_peft()

        from peft import PeftModel
        from transformers import AutoModel, AutoConfig

        name = name or f"{base_model.split('/')[-1]}+lora"

        if self.verbose:
            print(f"[LoRA] Loading base model: {base_model}")
            print(f"[LoRA] Loading LoRA adapter: {lora_path}")

        # Get config for hidden size
        config = AutoConfig.from_pretrained(base_model)
        hidden_size = getattr(config, "hidden_size", None)
        if hidden_size is None:
            hidden_size = getattr(config, "d_model", 768)

        # Load base model
        load_kwargs = {}
        if self.torch_dtype:
            load_kwargs["torch_dtype"] = self.torch_dtype

        base = AutoModel.from_pretrained(base_model, **load_kwargs)

        # Load and apply LoRA
        model = PeftModel.from_pretrained(base, lora_path)

        if merge_lora:
            if self.verbose:
                print("[LoRA] Merging LoRA weights into base model...")
            model = model.merge_and_unload()

        model = model.to(self.device)
        model.eval()

        if self.verbose:
            total_params = sum(p.numel() for p in model.parameters())
            print(f"[LoRA] Loaded '{name}': {total_params:,} params")

        return LoRAExpert(
            model=model,
            hidden_size=hidden_size,
            name=name,
            base_model_name=base_model,
            lora_paths=[lora_path],
        )

    def load_with_multiple_loras(
        self,
        base_model: str,
        lora_paths: List[str],
        lora_weights: Optional[List[float]] = None,
        merge_strategy: str = "weighted_average",
        name: Optional[str] = None,
    ) -> LoRAExpert:
        """Load base model with multiple LoRA adapters merged.

        Args:
            base_model: HuggingFace model ID.
            lora_paths: List of LoRA adapter paths.
            lora_weights: Weights for each LoRA (default: equal).
            merge_strategy: How to combine LoRAs ('weighted_average', 'ties', 'slerp').
            name: Optional name for the expert.

        Returns:
            LoRAExpert with all LoRAs merged.
        """
        self._check_peft()

        from peft import PeftModel
        from transformers import AutoModel, AutoConfig

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

        name = name or f"{base_model.split('/')[-1]}+{len(lora_paths)}loras"

        if self.verbose:
            print(f"[LoRA] Loading base model: {base_model}")
            print(f"[LoRA] Merging {len(lora_paths)} LoRA adapters with {merge_strategy}")

        # Get config
        config = AutoConfig.from_pretrained(base_model)
        hidden_size = getattr(config, "hidden_size", None)
        if hidden_size is None:
            hidden_size = getattr(config, "d_model", 768)

        # Load base model
        load_kwargs = {}
        if self.torch_dtype:
            load_kwargs["torch_dtype"] = self.torch_dtype

        base = AutoModel.from_pretrained(base_model, **load_kwargs)
        base_state = {k: v.clone() for k, v in base.state_dict().items()}

        # Extract deltas from each LoRA
        lora_deltas = []
        for i, lora_path in enumerate(lora_paths):
            if self.verbose:
                print(f"[LoRA] Extracting delta from: {lora_path}")

            # Load LoRA
            peft_model = PeftModel.from_pretrained(
                AutoModel.from_pretrained(base_model, **load_kwargs),
                lora_path,
            )

            # Merge to get effective weights
            merged = peft_model.merge_and_unload()
            merged_state = merged.state_dict()

            # Compute delta: merged - base
            delta = {}
            for key in base_state:
                if key in merged_state:
                    delta[key] = merged_state[key] - base_state[key]

            lora_deltas.append(delta)

            # Clean up
            del peft_model, merged

        # Merge deltas using specified strategy
        if self.verbose:
            print(f"[LoRA] Merging deltas with strategy: {merge_strategy}")

        merged_delta = self._merge_deltas(lora_deltas, lora_weights, merge_strategy)

        # Apply merged delta to base
        final_state = {}
        for key in base_state:
            if key in merged_delta:
                final_state[key] = base_state[key] + merged_delta[key]
            else:
                final_state[key] = base_state[key]

        base.load_state_dict(final_state)
        base = base.to(self.device)
        base.eval()

        if self.verbose:
            total_params = sum(p.numel() for p in base.parameters())
            print(f"[LoRA] Created '{name}': {total_params:,} params")

        return LoRAExpert(
            model=base,
            hidden_size=hidden_size,
            name=name,
            base_model_name=base_model,
            lora_paths=lora_paths,
        )

    def _merge_deltas(
        self,
        deltas: List[Dict[str, torch.Tensor]],
        weights: List[float],
        strategy: str,
    ) -> Dict[str, torch.Tensor]:
        """Merge multiple weight deltas using specified strategy.

        Args:
            deltas: List of state dict deltas.
            weights: Per-delta weights.
            strategy: Merge strategy name.

        Returns:
            Merged delta dict.
        """
        # Get all keys
        all_keys = set()
        for d in deltas:
            all_keys.update(d.keys())

        merged = {}

        if strategy == "weighted_average":
            for key in all_keys:
                tensors = []
                key_weights = []
                for i, delta in enumerate(deltas):
                    if key in delta and delta[key].abs().sum() > 0:
                        tensors.append(delta[key])
                        key_weights.append(weights[i])

                if tensors:
                    # Normalize weights for this key
                    w_sum = sum(key_weights)
                    if w_sum > 0:
                        merged[key] = sum(
                            w / w_sum * t for w, t in zip(key_weights, tensors)
                        )

        elif strategy == "slerp":
            # SLERP for 2 deltas, pairwise for more
            for key in all_keys:
                tensors = [d.get(key) for d in deltas if key in d]
                tensors = [t for t in tensors if t is not None and t.abs().sum() > 0]

                if len(tensors) == 0:
                    continue
                elif len(tensors) == 1:
                    merged[key] = tensors[0] * weights[0]
                else:
                    # 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[key] = result

        elif strategy == "ties":
            # TIES: Trim, Elect sign, Merge
            for key in all_keys:
                tensors = []
                for delta in deltas:
                    if key in delta:
                        tensors.append(delta[key])

                if not tensors:
                    continue

                # Stack tensors
                stacked = torch.stack(tensors)

                # Trim: zero out small values (keep top 50%)
                abs_vals = stacked.abs()
                threshold = abs_vals.quantile(0.5)
                trimmed = stacked.clone()
                trimmed[abs_vals < threshold] = 0

                # Elect sign: majority vote
                signs = torch.sign(trimmed)
                sign_sum = signs.sum(dim=0)
                elected_sign = torch.sign(sign_sum)
                elected_sign[elected_sign == 0] = 1

                # Merge: average only values that agree with elected sign
                mask = signs == elected_sign.unsqueeze(0)
                masked = trimmed * mask.float()

                # Weighted average of agreeing values
                weight_tensor = torch.tensor(weights, device=stacked.device).view(-1, *([1] * (stacked.dim() - 1)))
                weighted = masked * weight_tensor
                count = mask.float().sum(dim=0).clamp(min=1)
                merged[key] = weighted.sum(dim=0) / count

        else:
            # Default to weighted average
            return self._merge_deltas(deltas, weights, "weighted_average")

        return merged

    def _slerp_tensor(
        self, a: torch.Tensor, b: torch.Tensor, t: float
    ) -> torch.Tensor:
        """Spherical linear interpolation between two tensors.

        Args:
            a: First tensor.
            b: Second tensor.
            t: Interpolation factor (0 = a, 1 = b).

        Returns:
            Interpolated tensor.
        """
        # Flatten for dot product
        a_flat = a.flatten().float()
        b_flat = b.flatten().float()

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

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

        # Handle parallel vectors
        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
        )

        # Interpolate norms
        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).to(a.dtype)

    def get_lora_info(self, lora_path: str) -> Dict[str, Any]:
        """Get information about a LoRA adapter without loading it.

        Args:
            lora_path: Path to LoRA adapter.

        Returns:
            Dict with LoRA configuration info.
        """
        self._check_peft()

        from peft import PeftConfig

        config = PeftConfig.from_pretrained(lora_path)

        return {
            "base_model": config.base_model_name_or_path,
            "task_type": str(config.task_type),
            "r": getattr(config, "r", None),
            "lora_alpha": getattr(config, "lora_alpha", None),
            "lora_dropout": getattr(config, "lora_dropout", None),
            "target_modules": getattr(config, "target_modules", None),
        }
