"""
Merge strategies — different algorithms for combining weight tensors from
multiple models into one.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Sequence

import numpy as np


class MergeStrategy(ABC):
    """Abstract base for weight merging strategies."""

    name: str = "base"

    @abstractmethod
    def merge_tensors(
        self,
        tensors: List[np.ndarray],
        weights: Optional[List[float]] = None,
    ) -> np.ndarray:
        """Merge multiple weight tensors into one.

        Args:
            tensors: List of weight tensors (same shape) from different models.
            weights: Optional per-model importance weights (sum to 1).
                If None, equal weighting is used.

        Returns:
            Merged weight tensor.
        """


class WeightedAverageStrategy(MergeStrategy):
    """Simple weighted average of weight tensors.

    merged = w1*A + w2*B + w3*C + ...

    Fast, simple, and surprisingly effective. Good default choice.
    """

    name = "weighted_average"

    def merge_tensors(
        self,
        tensors: List[np.ndarray],
        weights: Optional[List[float]] = None,
    ) -> np.ndarray:
        n = len(tensors)
        if weights is None:
            weights = [1.0 / n] * n

        # Normalize weights
        total = sum(weights)
        weights = [w / total for w in weights]

        result = np.zeros_like(tensors[0], dtype=np.float64)
        for tensor, w in zip(tensors, weights):
            result += w * tensor.astype(np.float64)

        return result.astype(tensors[0].dtype)


class SLERPStrategy(MergeStrategy):
    """Spherical Linear Interpolation (SLERP) between weight tensors.

    Interpolates along the surface of a hypersphere rather than linearly.
    Produces smoother blending that preserves weight norms better.

    For >2 models, performs pairwise SLERP in a tournament bracket.
    """

    name = "slerp"

    def merge_tensors(
        self,
        tensors: List[np.ndarray],
        weights: Optional[List[float]] = None,
    ) -> np.ndarray:
        n = len(tensors)
        if n == 1:
            return tensors[0].copy()

        if weights is None:
            weights = [1.0 / n] * n

        # Normalize weights
        total = sum(weights)
        weights = [w / total for w in weights]

        # For 2 tensors: direct SLERP
        if n == 2:
            return self._slerp(tensors[0], tensors[1], weights[1])

        # For >2: iterative pairwise SLERP weighted by importance
        # Start with the first tensor, then SLERP each subsequent one in
        result = tensors[0].astype(np.float64)
        cumulative_weight = weights[0]

        for i in range(1, n):
            # The interpolation factor is the proportion of the new model's
            # weight relative to the cumulative total
            cumulative_weight += weights[i]
            t = weights[i] / cumulative_weight
            result = self._slerp(result, tensors[i], t)

        return result.astype(tensors[0].dtype)

    @staticmethod
    def _slerp(a: np.ndarray, b: np.ndarray, t: float) -> np.ndarray:
        """SLERP between two tensors.

        Args:
            a, b: Tensors of the same shape.
            t: Interpolation factor (0 = a, 1 = b).
        """
        a_flat = a.flatten().astype(np.float64)
        b_flat = b.flatten().astype(np.float64)

        # Compute norms
        a_norm = np.linalg.norm(a_flat)
        b_norm = np.linalg.norm(b_flat)

        if a_norm < 1e-10 or b_norm < 1e-10:
            # Degenerate case: fall back to linear interpolation
            return ((1 - t) * a + t * b).astype(a.dtype)

        # Normalize
        a_unit = a_flat / a_norm
        b_unit = b_flat / b_norm

        # Angle between vectors
        dot = np.clip(np.dot(a_unit, b_unit), -1.0, 1.0)
        omega = np.arccos(dot)

        if omega < 1e-10:
            # Vectors are nearly parallel: fall back to linear
            result = ((1 - t) * a_flat + t * b_flat)
        else:
            sin_omega = np.sin(omega)
            result = (
                np.sin((1 - t) * omega) / sin_omega * a_flat
                + np.sin(t * omega) / sin_omega * b_flat
            )

        # Interpolate the norms linearly
        merged_norm = (1 - t) * a_norm + t * b_norm
        result_norm = np.linalg.norm(result)
        if result_norm > 1e-10:
            result = result * (merged_norm / result_norm)

        return result.reshape(a.shape).astype(a.dtype)


class TIESStrategy(MergeStrategy):
    """TIES-Merging: Trim, Elect Sign, Merge.

    Reference: Yadav et al. — "Resolving Interference When Merging Models" (2023)

    Steps:
    1. Compute task vectors (delta from a base/average)
    2. Trim small-magnitude deltas (keep top-k%)
    3. Elect dominant sign per parameter (majority vote)
    4. Average only the deltas that agree with the elected sign
    5. Apply merged delta back to base

    Best for merging many models where parameter conflicts are common.
    """

    name = "ties"

    def __init__(self, density: float = 0.5) -> None:
        """
        Args:
            density: Fraction of parameters to keep (0.0–1.0).
                0.5 = keep top 50% of deltas by magnitude.
        """
        self.density = density

    def merge_tensors(
        self,
        tensors: List[np.ndarray],
        weights: Optional[List[float]] = None,
    ) -> np.ndarray:
        n = len(tensors)
        if n == 1:
            return tensors[0].copy()

        if weights is None:
            weights = [1.0 / n] * n

        # Normalize weights
        total = sum(weights)
        weights = [w / total for w in weights]

        # Step 0: Compute base (weighted average as reference)
        base = np.zeros_like(tensors[0], dtype=np.float64)
        for tensor, w in zip(tensors, weights):
            base += w * tensor.astype(np.float64)

        # Step 1: Compute task vectors (deltas from base)
        deltas = [
            (tensor.astype(np.float64) - base) for tensor in tensors
        ]

        # Step 2: Trim — zero out small-magnitude deltas
        trimmed = []
        for delta in deltas:
            flat = np.abs(delta.flatten())
            if len(flat) == 0:
                trimmed.append(delta)
                continue
            threshold = np.quantile(flat, 1.0 - self.density)
            mask = np.abs(delta) >= threshold
            trimmed.append(delta * mask)

        # Step 3: Elect sign — majority vote per parameter
        sign_votes = np.zeros_like(base)
        for delta, w in zip(trimmed, weights):
            sign_votes += w * np.sign(delta)
        elected_sign = np.sign(sign_votes)

        # Step 4: Average deltas that agree with elected sign
        merged_delta = np.zeros_like(base)
        count = np.zeros_like(base)
        for delta, w in zip(trimmed, weights):
            agree_mask = np.sign(delta) == elected_sign
            merged_delta += w * delta * agree_mask
            count += agree_mask * w

        # Avoid division by zero
        safe_count = np.where(count > 0, count, 1.0)
        merged_delta = merged_delta / safe_count

        # Step 5: Apply merged delta to base
        result = base + merged_delta

        return result.astype(tensors[0].dtype)


class DAREStrategy(MergeStrategy):
    """DARE — Drop And REscale Merging.

    Reference: Yu et al. — "Language Models are Super Mario" (2023)

    Steps:
    1. Compute task vectors (delta from base average)
    2. Randomly drop a fraction of delta parameters (set to zero)
    3. Rescale surviving parameters by 1/(1 - drop_rate) to preserve magnitude
    4. Sum rescaled deltas and apply back to base

    Reduces inter-model interference by sparsifying the merged delta.
    """

    name = "dare"

    def __init__(self, drop_rate: float = 0.5, seed: int = 42) -> None:
        """
        Args:
            drop_rate: Fraction of delta parameters to drop (0.0–1.0).
                0.5 = drop 50% of each task vector's parameters.
            seed: Random seed for reproducible drops.
        """
        self.drop_rate = drop_rate
        self.seed = seed

    def merge_tensors(
        self,
        tensors: List[np.ndarray],
        weights: Optional[List[float]] = None,
    ) -> np.ndarray:
        n = len(tensors)
        if n == 1:
            return tensors[0].copy()

        if weights is None:
            weights = [1.0 / n] * n

        # Normalize weights
        total = sum(weights)
        weights = [w / total for w in weights]

        rng = np.random.default_rng(self.seed)

        # Step 0: Compute base (weighted average)
        base = np.zeros_like(tensors[0], dtype=np.float64)
        for tensor, w in zip(tensors, weights):
            base += w * tensor.astype(np.float64)

        # Step 1: Compute task vectors
        deltas = [(tensor.astype(np.float64) - base) for tensor in tensors]

        # Step 2-3: Drop and rescale each delta
        rescale_factor = 1.0 / max(1.0 - self.drop_rate, 1e-8)
        dropped_deltas = []
        for delta in deltas:
            mask = rng.random(delta.shape) >= self.drop_rate
            dropped_delta = delta * mask * rescale_factor
            dropped_deltas.append(dropped_delta)

        # Step 4: Weighted sum of rescaled deltas
        merged_delta = np.zeros_like(base)
        for delta, w in zip(dropped_deltas, weights):
            merged_delta += w * delta

        result = base + merged_delta
        return result.astype(tensors[0].dtype)


class BreadcrumbsStrategy(MergeStrategy):
    """Model Breadcrumbs — Magnitude-based sparse merging.

    Keeps only the top-k% of each task vector by absolute magnitude,
    then averages the surviving parameters. Simpler than TIES — no sign
    election step. Works well when models are fine-tuned from the same base.

    Steps:
    1. Compute task vectors (delta from base average)
    2. Keep only top density% parameters by magnitude per model
    3. Average all surviving deltas (weighted)
    4. Apply merged delta back to base
    """

    name = "breadcrumbs"

    def __init__(self, density: float = 0.3) -> None:
        """
        Args:
            density: Fraction of top-magnitude parameters to keep (0.0–1.0).
                0.3 = keep top 30% of each task vector.
        """
        self.density = density

    def merge_tensors(
        self,
        tensors: List[np.ndarray],
        weights: Optional[List[float]] = None,
    ) -> np.ndarray:
        n = len(tensors)
        if n == 1:
            return tensors[0].copy()

        if weights is None:
            weights = [1.0 / n] * n

        total = sum(weights)
        weights = [w / total for w in weights]

        # Base: weighted average
        base = np.zeros_like(tensors[0], dtype=np.float64)
        for tensor, w in zip(tensors, weights):
            base += w * tensor.astype(np.float64)

        # Task vectors
        deltas = [(tensor.astype(np.float64) - base) for tensor in tensors]

        # Keep top-k% by magnitude per delta
        sparse_deltas = []
        for delta in deltas:
            flat = np.abs(delta.flatten())
            if len(flat) == 0:
                sparse_deltas.append(delta)
                continue
            threshold = np.quantile(flat, 1.0 - self.density)
            mask = np.abs(delta) >= threshold
            sparse_deltas.append(delta * mask)

        # Weighted average of sparse deltas
        merged_delta = np.zeros_like(base)
        for delta, w in zip(sparse_deltas, weights):
            merged_delta += w * delta

        result = base + merged_delta
        return result.astype(tensors[0].dtype)


class ConsensusStrategy(MergeStrategy):
    """Consensus Merge — Only keep parameters where models agree.

    Ultra-conservative strategy: only merges parameters where all (or a
    quorum of) models agree on the sign/direction. Conflicting parameters
    are zeroed. Produces very clean merges with minimal interference.

    Steps:
    1. Compute task vectors from base
    2. For each parameter, count how many models have positive vs negative delta
    3. Keep only parameters where agreement ratio >= consensus_threshold
    4. Average the agreeing deltas
    5. Apply back to base
    """

    name = "consensus"

    def __init__(self, consensus_threshold: float = 0.7) -> None:
        """
        Args:
            consensus_threshold: Minimum fraction of models that must agree
                on a parameter's sign (0.5–1.0). 1.0 = unanimous.
        """
        self.consensus_threshold = max(0.5, min(1.0, consensus_threshold))

    def merge_tensors(
        self,
        tensors: List[np.ndarray],
        weights: Optional[List[float]] = None,
    ) -> np.ndarray:
        n = len(tensors)
        if n == 1:
            return tensors[0].copy()

        if weights is None:
            weights = [1.0 / n] * n

        total = sum(weights)
        weights = [w / total for w in weights]

        # Base: weighted average
        base = np.zeros_like(tensors[0], dtype=np.float64)
        for tensor, w in zip(tensors, weights):
            base += w * tensor.astype(np.float64)

        # Task vectors
        deltas = [(tensor.astype(np.float64) - base) for tensor in tensors]

        # Count sign agreement per parameter
        signs = np.stack([np.sign(d) for d in deltas], axis=0)  # [n, ...]
        positive_count = np.sum(signs > 0, axis=0)
        negative_count = np.sum(signs < 0, axis=0)
        max_agreement = np.maximum(positive_count, negative_count)
        nonzero_count = np.sum(signs != 0, axis=0)

        # Compute agreement ratio (where nonzero)
        safe_count = np.where(nonzero_count > 0, nonzero_count, 1)
        agreement_ratio = max_agreement / safe_count

        # Consensus mask: keep only where agreement exceeds threshold
        consensus_mask = agreement_ratio >= self.consensus_threshold

        # Determine elected sign (majority)
        elected_sign = np.where(positive_count >= negative_count, 1.0, -1.0)
        elected_sign = np.where(nonzero_count == 0, 0.0, elected_sign)

        # Average deltas that agree with elected sign, masked by consensus
        merged_delta = np.zeros_like(base)
        weight_sum = np.zeros_like(base)
        for delta, w in zip(deltas, weights):
            agree = np.sign(delta) == elected_sign
            merged_delta += w * delta * agree * consensus_mask
            weight_sum += w * agree * consensus_mask

        safe_weight = np.where(weight_sum > 0, weight_sum, 1.0)
        merged_delta = merged_delta / safe_weight

        result = base + merged_delta
        return result.astype(tensors[0].dtype)


# Strategy registry for lookup by name
STRATEGY_REGISTRY: Dict[str, type] = {
    "weighted_average": WeightedAverageStrategy,
    "slerp": SLERPStrategy,
    "ties": TIESStrategy,
    "dare": DAREStrategy,
    "breadcrumbs": BreadcrumbsStrategy,
    "consensus": ConsensusStrategy,
}


def get_strategy(name: str, **kwargs) -> MergeStrategy:
    """Get a merge strategy by name.

    Args:
        name: One of 'weighted_average', 'slerp', 'ties', 'dare',
            'breadcrumbs', 'consensus'.
        **kwargs: Strategy-specific arguments.

    Returns:
        Instantiated strategy.
    """
    if name not in STRATEGY_REGISTRY:
        raise ValueError(
            f"Unknown strategy '{name}'. Available: {list(STRATEGY_REGISTRY.keys())}"
        )
    return STRATEGY_REGISTRY[name](**kwargs)
