"""
CrossMerger — merge models of ANY architecture into one.

Unlike weight-merging (which requires same architecture), CrossMerger
merges what models KNOW, not what they weigh.

How it works:
1. Run all source models on sample data → collect their outputs
2. Project all outputs to a common dimension
3. Blend/merge in the common space
4. Train a unified student model on the merged knowledge

Result: a single new model with combined knowledge from all sources,
regardless of their original architectures.

Usage::

    from model_osmosis.merging.cross_merger import CrossMerger

    merger = CrossMerger()
    merged = merger.merge(
        models=[resnet, transformer, cnn],
        sample_data=X,
        output_dim=256,
        student_hidden=512,
    )
    # merged is a single nn.Module with knowledge from all 3
"""

from __future__ import annotations

import copy
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from model_osmosis.utils.determinism import make_torch_generator, set_global_seed


class CommonProjector(nn.Module):
    """Projects any model's output to a common dimension."""

    def __init__(self, in_dim: int, common_dim: int) -> None:
        super().__init__()
        self.proj = nn.Sequential(
            nn.Linear(in_dim, max(64, common_dim)),
            nn.ReLU(),
            nn.Linear(max(64, common_dim), common_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.proj(x)


class MergedStudent(nn.Module):
    """The student model that learns merged knowledge.

    This is the final output — a clean nn.Module that contains
    knowledge from all source models, regardless of architecture.
    """

    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        output_dim: int,
        num_layers: int = 3,
    ) -> None:
        super().__init__()
        layers = []
        prev = input_dim
        for i in range(num_layers - 1):
            h = hidden_dim if i == 0 else hidden_dim // 2
            layers.extend([nn.Linear(prev, h), nn.ReLU(), nn.Dropout(0.1)])
            prev = h
        layers.append(nn.Linear(prev, output_dim))
        self.net = nn.Sequential(*layers)
        self.input_dim = input_dim
        self.output_dim = output_dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class CrossMerger:
    """Merge models of different architectures into one.

    Works by merging FUNCTIONAL knowledge (what models output)
    rather than structural knowledge (weights), so any architecture
    combination works.

    Args:
        common_dim: Dimension of the common representation space.
        verbose: Print progress.
    """

    def __init__(
        self,
        common_dim: int = 256,
        verbose: bool = True,
    ) -> None:
        self.common_dim = common_dim
        self.verbose = verbose

    def merge(
        self,
        models: List[nn.Module],
        sample_data: np.ndarray,
        output_dim: Optional[int] = None,
        student_hidden: int = 512,
        student_layers: int = 3,
        weights: Optional[List[float]] = None,
        epochs: int = 100,
        lr: float = 1e-3,
        batch_size: int = 32,
        projection_epochs: Optional[int] = None,
        common_dim: Optional[int] = None,
        seed: int = 42,
    ) -> MergedStudent:
        """Merge models of any architecture into a single model.

        Args:
            models: List of nn.Module models (ANY architecture).
            sample_data: Sample input data for knowledge extraction.
            output_dim: Output dimension of the merged model.
                Auto-detected if not provided.
            student_hidden: Hidden size of the student model.
            student_layers: Number of layers in the student model.
            weights: Per-model importance weights (default: equal).
            epochs: Training epochs for the student.
            lr: Learning rate.
            batch_size: Training batch size.

        Returns:
            A MergedStudent model containing merged knowledge.
        """
        n = len(models)
        if n == 0:
            raise ValueError("Need at least one model to merge")

        if self.verbose:
            print(f"[CrossMerger] Merging {n} models of different architectures")

        set_global_seed(seed)
        x_tensor = torch.from_numpy(sample_data).float()
        input_dim = sample_data.shape[-1]

        # Step 1: Collect outputs from all models
        if self.verbose:
            print("[CrossMerger] Step 1/4: Collecting model outputs...")

        model_outputs = []
        output_dims = []
        for i, model in enumerate(models):
            model.eval()
            with torch.no_grad():
                out = model(x_tensor)
            if out.dim() > 2:
                out = out.view(out.size(0), -1)
            model_outputs.append(out)
            output_dims.append(out.shape[-1])
            if self.verbose:
                print(f"  Model {i}: output dim = {output_dims[-1]}")

        # Auto-detect output dim
        if output_dim is None:
            output_dim = max(output_dims)
            if self.verbose:
                print(f"  Output dim (auto): {output_dim}")

        # Step 2: Learn projections to common space
        if self.verbose:
            print("[CrossMerger] Step 2/4: Learning common projections...")

        effective_common_dim = int(common_dim or self.common_dim)
        projectors = nn.ModuleList([
            CommonProjector(odim, effective_common_dim) for odim in output_dims
        ])

        # Train projectors to align representations
        proj_optimizer = torch.optim.Adam(projectors.parameters(), lr=lr)

        proj_epochs = (
            int(projection_epochs)
            if projection_epochs is not None
            else min(int(epochs), 50)
        )
        proj_epochs = max(1, proj_epochs)
        for epoch in range(proj_epochs):
            proj_optimizer.zero_grad()

            # Project all outputs to common space
            projected = [
                projectors[i](model_outputs[i]) for i in range(n)
            ]

            # Loss: projections should be close to each other (alignment)
            alignment_loss = torch.tensor(0.0)
            count = 0
            for i in range(n):
                for j in range(i + 1, n):
                    alignment_loss += F.mse_loss(projected[i], projected[j])
                    count += 1
            if count > 0:
                alignment_loss = alignment_loss / count

            # Regularization: projections should preserve information
            # (not collapse to zero)
            diversity_loss = torch.tensor(0.0)
            for p in projected:
                diversity_loss -= torch.var(p)
            diversity_loss = diversity_loss / n

            loss = alignment_loss + 0.1 * diversity_loss
            loss.backward()
            proj_optimizer.step()

        if self.verbose:
            print(f"  Alignment loss: {alignment_loss.item():.6f}")

        # Step 3: Create merged targets
        if self.verbose:
            print("[CrossMerger] Step 3/4: Merging knowledge in common space...")

        # Set weights
        if weights is None:
            weights = [1.0 / n] * n
        else:
            total = sum(weights)
            weights = [w / total for w in weights]

        # Merge in common space: weighted average of projected outputs
        with torch.no_grad():
            projected = [projectors[i](model_outputs[i]) for i in range(n)]
            merged_knowledge = torch.zeros_like(projected[0])
            for i in range(n):
                merged_knowledge += weights[i] * projected[i]

        # Also create merged OUTPUT targets (in original output space)
        # by weighted averaging the original model outputs (padded to same dim)
        padded_outputs = []
        for out in model_outputs:
            if out.shape[-1] < output_dim:
                padded = F.pad(out, (0, output_dim - out.shape[-1]))
            elif out.shape[-1] > output_dim:
                padded = out[:, :output_dim]
            else:
                padded = out
            padded_outputs.append(padded)

        with torch.no_grad():
            merged_targets = torch.zeros_like(padded_outputs[0])
            for i in range(n):
                merged_targets += weights[i] * padded_outputs[i]

        # Step 4: Train student on merged knowledge
        if self.verbose:
            print("[CrossMerger] Step 4/4: Training merged student model...")

        student = MergedStudent(
            input_dim=input_dim,
            hidden_dim=student_hidden,
            output_dim=output_dim,
            num_layers=student_layers,
        )

        optimizer = torch.optim.Adam(student.parameters(), lr=lr)
        dataset = torch.utils.data.TensorDataset(x_tensor, merged_targets)
        loader = torch.utils.data.DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=True,
            generator=make_torch_generator(seed),
        )

        student.train()
        best_loss = float("inf")
        best_state = None

        for epoch in range(epochs):
            total_loss = 0.0
            for x_batch, target_batch in loader:
                optimizer.zero_grad()
                pred = student(x_batch)
                loss = F.mse_loss(pred, target_batch)
                loss.backward()
                optimizer.step()
                total_loss += loss.item()

            avg_loss = total_loss / len(loader)
            if avg_loss < best_loss:
                best_loss = avg_loss
                best_state = copy.deepcopy(student.state_dict())

            if self.verbose and (epoch + 1) % 20 == 0:
                print(f"  Epoch {epoch+1}/{epochs}: loss = {avg_loss:.6f}")

        # Load best weights
        if best_state is not None:
            student.load_state_dict(best_state)

        student.eval()

        if self.verbose:
            total_params = sum(p.numel() for p in student.parameters())
            print(f"[CrossMerger] ✓ Merged model: {total_params:,} params, "
                  f"input={input_dim}, output={output_dim}")
            print(f"[CrossMerger]   Best loss: {best_loss:.6f}")

        return student
