"""
OsmoFormat Operations - Merge, stack, and transform OsmoModels.

Operations work directly on the universal OsmoFormat, enabling:
- Merging models from ANY framework together
- Stacking models as Mixture of Experts
- Knowledge transfer between architectures
- Layer-wise composition
"""

from dataclasses import replace
from datetime import datetime
from typing import Callable, Dict, List, Optional, Tuple, Union

import numpy as np

from model_osmosis.format.osmo_format import (
    OsmoModel,
    OsmoManifest,
    ArchitectureSpec,
    LayerSpec,
    LayerType,
    ModelModality,
    FrameworkOrigin,
    OSMO_FORMAT_VERSION,
)


def _normalize_weights(weights: List[float]) -> List[float]:
    """Normalize weights to sum to 1."""
    total = sum(weights)
    if total == 0:
        return [1.0 / len(weights)] * len(weights)
    return [w / total for w in weights]


def _slerp(v0: np.ndarray, v1: np.ndarray, t: float) -> np.ndarray:
    """Spherical linear interpolation between two vectors/tensors."""
    # Flatten for interpolation
    shape = v0.shape
    v0_flat = v0.flatten().astype(np.float64)
    v1_flat = v1.flatten().astype(np.float64)

    # Normalize
    v0_norm = np.linalg.norm(v0_flat)
    v1_norm = np.linalg.norm(v1_flat)

    if v0_norm == 0 or v1_norm == 0:
        return ((1 - t) * v0 + t * v1).astype(v0.dtype)

    v0_unit = v0_flat / v0_norm
    v1_unit = v1_flat / v1_norm

    # Compute angle
    dot = np.clip(np.dot(v0_unit, v1_unit), -1.0, 1.0)
    theta = np.arccos(dot)

    if abs(theta) < 1e-6:
        return ((1 - t) * v0 + t * v1).astype(v0.dtype)

    # SLERP
    sin_theta = np.sin(theta)
    s0 = np.sin((1 - t) * theta) / sin_theta
    s1 = np.sin(t * theta) / sin_theta

    # Interpolate magnitudes
    result_norm = (1 - t) * v0_norm + t * v1_norm
    interp = s0 * v0_flat + s1 * v1_flat
    interp_norm = np.linalg.norm(interp)
    if interp_norm < 1e-12:
        return ((1 - t) * v0 + t * v1).astype(v0.dtype)
    result = interp * result_norm / interp_norm

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


def _ties_merge(
    tensors: List[np.ndarray],
    weights: List[float],
    density: float = 0.5,
) -> np.ndarray:
    """
    TIES merging - Task Inference for Expert Sampling.

    Resolves sign conflicts and trims small values.
    """
    if len(tensors) == 1:
        return tensors[0]

    # Compute deltas from mean
    mean_tensor = sum(t * w for t, w in zip(tensors, weights))
    deltas = [t - mean_tensor for t in tensors]

    # Determine sign by majority vote
    signs = np.sign(sum(np.sign(d) * w for d, w in zip(deltas, weights)))

    # Trim based on density (keep top k% of values by magnitude)
    merged = np.zeros_like(tensors[0])

    for delta, weight in zip(deltas, weights):
        # Trim small values
        threshold = np.percentile(np.abs(delta), (1 - density) * 100)
        mask = np.abs(delta) >= threshold

        # Apply sign alignment
        aligned_delta = np.where(np.sign(delta) == signs, delta, 0)

        merged += aligned_delta * weight * mask

    return (mean_tensor + merged).astype(tensors[0].dtype)


class OsmoMerger:
    """
    Merge multiple OsmoModels into one.

    Supports various merge strategies:
    - weighted_average: Simple weighted average of parameters
    - slerp: Spherical interpolation (better for maintaining norms)
    - ties: Task-specific merging with conflict resolution
    - layer_wise: Different strategy per layer type
    """

    def __init__(
        self,
        strategy: str = "weighted_average",
        ties_density: float = 0.5,
    ):
        """
        Args:
            strategy: Merge strategy ("weighted_average", "slerp", "ties")
            ties_density: For TIES, fraction of parameters to keep
        """
        self.strategy = strategy
        self.ties_density = ties_density

    def merge(
        self,
        models: List[OsmoModel],
        weights: Optional[List[float]] = None,
        layer_strategies: Optional[Dict[LayerType, str]] = None,
    ) -> OsmoModel:
        """
        Merge multiple OsmoModels.

        Args:
            models: List of OsmoModels to merge
            weights: Optional weights (default: equal weights)
            layer_strategies: Optional per-layer-type strategies

        Returns:
            Merged OsmoModel
        """
        if len(models) == 0:
            raise ValueError("Need at least one model to merge")

        if len(models) == 1:
            return models[0]

        # Default equal weights
        if weights is None:
            weights = [1.0] * len(models)
        weights = _normalize_weights(weights)

        # Find common parameters
        all_param_names = set(models[0].parameters.keys())
        for model in models[1:]:
            all_param_names &= set(model.parameters.keys())

        if not all_param_names:
            raise ValueError("Models have no common parameters to merge")

        # Create layer type lookup
        layer_type_lookup = {}
        for model in models:
            for layer in model.architecture.layers:
                if layer.name not in layer_type_lookup:
                    layer_type_lookup[layer.name] = layer.layer_type

        # Merge parameters
        merged_params = {}

        for name in all_param_names:
            tensors = [m.parameters[name] for m in models]

            # Check shapes match
            shapes = [t.shape for t in tensors]
            if not all(s == shapes[0] for s in shapes):
                # Skip incompatible parameters
                continue

            # Determine strategy for this parameter
            layer_type = layer_type_lookup.get(name, LayerType.CUSTOM)
            strategy = self.strategy
            if layer_strategies and layer_type in layer_strategies:
                strategy = layer_strategies[layer_type]

            # Apply merge strategy
            if strategy == "slerp" and len(tensors) == 2:
                merged = _slerp(tensors[0], tensors[1], weights[1])
            elif strategy == "ties":
                merged = _ties_merge(tensors, weights, self.ties_density)
            else:
                # Weighted average
                merged = sum(t * w for t, w in zip(tensors, weights)).astype(tensors[0].dtype)

            merged_params[name] = merged

        # Build merged architecture (use first model as template)
        merged_layers = []
        for layer in models[0].architecture.layers:
            if layer.name in merged_params:
                merged_layers.append(LayerSpec(
                    name=layer.name,
                    layer_type=layer.layer_type,
                    shape=list(merged_params[layer.name].shape),
                    dtype=str(merged_params[layer.name].dtype),
                    attributes=layer.attributes,
                ))

        merged_arch = ArchitectureSpec(
            layers=merged_layers,
            hidden_size=models[0].architecture.hidden_size,
            num_layers=models[0].architecture.num_layers,
            num_attention_heads=models[0].architecture.num_attention_heads,
            vocab_size=models[0].architecture.vocab_size,
            max_position=models[0].architecture.max_position,
            intermediate_size=models[0].architecture.intermediate_size,
            architecture_type=models[0].architecture.architecture_type,
        )

        # Build manifest
        source_names = [m.manifest.name for m in models]
        weight_strs = [f"{m.manifest.name}:{w:.2f}" for m, w in zip(models, weights)]

        merged_manifest = OsmoManifest(
            version=OSMO_FORMAT_VERSION,
            name=f"merged-{'-'.join(source_names)[:50]}",
            description=f"Merged from: {', '.join(weight_strs)} using {self.strategy}",
            modality=models[0].manifest.modality,
            framework_origin=FrameworkOrigin.UNKNOWN,  # Framework-agnostic now
            parameter_count=sum(p.size for p in merged_params.values()),
            created_at=datetime.now().isoformat(),
            tags=["merged", self.strategy],
        )

        return OsmoModel(
            manifest=merged_manifest,
            architecture=merged_arch,
            parameters=merged_params,
        )


class OsmoStacker:
    """
    Stack OsmoModels as Mixture of Experts.

    Creates a gating mechanism that routes inputs to different experts.
    """

    def __init__(
        self,
        gating_type: str = "linear",
        top_k: int = 2,
    ):
        """
        Args:
            gating_type: Type of gating ("linear", "mlp")
            top_k: Number of experts to use per input
        """
        self.gating_type = gating_type
        self.top_k = top_k

    def stack(
        self,
        models: List[OsmoModel],
        expert_names: Optional[List[str]] = None,
    ) -> OsmoModel:
        """
        Stack models as MoE.

        Args:
            models: List of expert OsmoModels
            expert_names: Optional names for each expert

        Returns:
            Stacked OsmoModel with gating parameters
        """
        if len(models) < 2:
            raise ValueError("Need at least 2 models for stacking")

        if expert_names is None:
            expert_names = [f"expert_{i}" for i in range(len(models))]

        # Collect all parameters with expert prefix
        stacked_params = {}
        num_experts = len(models)

        for i, (model, name) in enumerate(zip(models, expert_names)):
            for param_name, arr in model.parameters.items():
                stacked_params[f"{name}/{param_name}"] = arr

        # Create gating parameters
        hidden_size = models[0].architecture.hidden_size

        if self.gating_type == "linear":
            # Simple linear gate: input_dim -> num_experts
            gate_weight = np.random.randn(hidden_size, num_experts).astype(np.float32) * 0.02
            stacked_params["gate/weight"] = gate_weight
        else:
            # MLP gate
            gate_hidden = hidden_size // 4
            stacked_params["gate/fc1/weight"] = np.random.randn(hidden_size, gate_hidden).astype(np.float32) * 0.02
            stacked_params["gate/fc1/bias"] = np.zeros(gate_hidden, dtype=np.float32)
            stacked_params["gate/fc2/weight"] = np.random.randn(gate_hidden, num_experts).astype(np.float32) * 0.02
            stacked_params["gate/fc2/bias"] = np.zeros(num_experts, dtype=np.float32)

        # Build architecture
        all_layers = []

        # Add gate layers
        all_layers.append(LayerSpec(
            name="gate/weight" if self.gating_type == "linear" else "gate/fc1/weight",
            layer_type=LayerType.GATE,
            shape=list(stacked_params["gate/weight" if self.gating_type == "linear" else "gate/fc1/weight"].shape),
            dtype="float32",
        ))

        # Add expert layers
        for model, name in zip(models, expert_names):
            for layer in model.architecture.layers:
                prefixed_name = f"{name}/{layer.name}"
                if prefixed_name in stacked_params:
                    all_layers.append(LayerSpec(
                        name=prefixed_name,
                        layer_type=layer.layer_type,
                        shape=list(stacked_params[prefixed_name].shape),
                        dtype=str(stacked_params[prefixed_name].dtype),
                        attributes={**layer.attributes, "expert": name},
                    ))

        stacked_arch = ArchitectureSpec(
            layers=all_layers,
            hidden_size=hidden_size,
            num_layers=models[0].architecture.num_layers,
            num_attention_heads=models[0].architecture.num_attention_heads,
            architecture_type="mixture_of_experts",
            extra={
                "num_experts": num_experts,
                "top_k": self.top_k,
                "gating_type": self.gating_type,
                "expert_names": expert_names,
            },
        )

        stacked_manifest = OsmoManifest(
            version=OSMO_FORMAT_VERSION,
            name=f"moe-{num_experts}experts",
            description=f"Stacked MoE with {num_experts} experts, top-{self.top_k} routing",
            modality=models[0].manifest.modality,
            framework_origin=FrameworkOrigin.UNKNOWN,
            parameter_count=sum(p.size for p in stacked_params.values()),
            created_at=datetime.now().isoformat(),
            tags=["moe", "stacked", f"top-{self.top_k}"],
        )

        return OsmoModel(
            manifest=stacked_manifest,
            architecture=stacked_arch,
            parameters=stacked_params,
        )


class OsmoComposer:
    """
    Compose new models by combining layers from different OsmoModels.

    Enables fine-grained control over which layers come from which model.
    """

    def compose(
        self,
        models: Dict[str, OsmoModel],
        layer_mapping: Dict[str, Tuple[str, str]],
    ) -> OsmoModel:
        """
        Compose a new model from layers of different models.

        Args:
            models: Dict of name -> OsmoModel
            layer_mapping: Dict of output_layer_name -> (source_model_name, source_layer_name)

        Returns:
            Composed OsmoModel
        """
        composed_params = {}
        composed_layers = []

        for out_name, (model_name, src_name) in layer_mapping.items():
            if model_name not in models:
                raise ValueError(f"Unknown model: {model_name}")

            model = models[model_name]
            if src_name not in model.parameters:
                raise ValueError(f"Layer {src_name} not found in {model_name}")

            composed_params[out_name] = model.parameters[src_name].copy()

            # Find layer spec
            for layer in model.architecture.layers:
                if layer.name == src_name:
                    composed_layers.append(LayerSpec(
                        name=out_name,
                        layer_type=layer.layer_type,
                        shape=layer.shape,
                        dtype=layer.dtype,
                        attributes={**layer.attributes, "source": f"{model_name}/{src_name}"},
                    ))
                    break

        # Use first model as architecture template
        first_model = next(iter(models.values()))

        composed_arch = ArchitectureSpec(
            layers=composed_layers,
            hidden_size=first_model.architecture.hidden_size,
            num_layers=len([l for l in composed_layers if l.layer_type == LayerType.ATTENTION]),
            architecture_type="composed",
        )

        composed_manifest = OsmoManifest(
            version=OSMO_FORMAT_VERSION,
            name="composed-model",
            description=f"Composed from: {', '.join(models.keys())}",
            modality=first_model.manifest.modality,
            framework_origin=FrameworkOrigin.UNKNOWN,
            parameter_count=sum(p.size for p in composed_params.values()),
            created_at=datetime.now().isoformat(),
            tags=["composed"],
        )

        return OsmoModel(
            manifest=composed_manifest,
            architecture=composed_arch,
            parameters=composed_params,
        )


def merge_osmo(
    models: List[OsmoModel],
    weights: Optional[List[float]] = None,
    strategy: str = "weighted_average",
) -> OsmoModel:
    """Convenience function to merge OsmoModels."""
    merger = OsmoMerger(strategy=strategy)
    return merger.merge(models, weights=weights)


def stack_osmo(
    models: List[OsmoModel],
    top_k: int = 2,
    gating_type: str = "linear",
) -> OsmoModel:
    """Convenience function to stack OsmoModels as MoE."""
    stacker = OsmoStacker(gating_type=gating_type, top_k=top_k)
    return stacker.stack(models)
