"""
ModalityBridge — merge models that operate on DIFFERENT data types.

Solves the cross-modal problem: how to merge a text model with a
protein model, or a vision model with an audio model.

The key insight: instead of feeding the SAME data to all models,
each model gets data in ITS OWN format, but PAIRED together.
The bridge learns to align their outputs in a common space.

Requires PAIRED data — examples where both modalities are represented:
- Text description + protein sequence
- Image + caption
- Audio + transcript

Usage::

    from model_osmosis.merging.modality_bridge import ModalityBridge

    bridge = ModalityBridge()
    merged = bridge.merge(
        models=[text_model, protein_model],
        paired_data=[text_samples, protein_samples],
        # text_samples[i] and protein_samples[i] describe the same thing
    )
"""

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


class ModalityEncoder(nn.Module):
    """Encodes one modality's output into the common space."""

    def __init__(self, in_dim: int, common_dim: int) -> None:
        super().__init__()
        hidden = max(128, common_dim * 2)
        self.encoder = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.ReLU(),
            nn.LayerNorm(hidden),
            nn.Linear(hidden, common_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        encoded = self.encoder(x)
        # L2 normalize for cosine similarity alignment
        return F.normalize(encoded, p=2, dim=-1)


class BridgedModel(nn.Module):
    """The result of cross-modal merging.

    Accepts input from ANY modality and produces output in the common
    space that contains knowledge from ALL modalities.

    Args:
        encoders: Per-modality encoders (input → common space).
        decoder: Common space → output.
        common_dim: Dimension of the common space.
        modality_names: Names for each modality.
    """

    def __init__(
        self,
        encoders: nn.ModuleList,
        decoder: nn.Module,
        common_dim: int,
        modality_names: List[str],
    ) -> None:
        super().__init__()
        self.encoders = encoders
        self.decoder = decoder
        self.common_dim = common_dim
        self.modality_names = modality_names

    def encode(self, x: torch.Tensor, modality: int = 0) -> torch.Tensor:
        """Encode input from a specific modality to common space.

        Args:
            x: Input tensor from one modality.
            modality: Index of the modality (0, 1, 2, ...).

        Returns:
            Common-space representation.
        """
        return self.encoders[modality](x)

    def forward(
        self, x: torch.Tensor, modality: int = 0
    ) -> torch.Tensor:
        """Full forward: encode modality → common space → decode.

        Args:
            x: Input from one modality.
            modality: Which modality this input belongs to.

        Returns:
            Decoded output.
        """
        common = self.encode(x, modality)
        return self.decoder(common)

    def translate(
        self,
        x: torch.Tensor,
        from_modality: int,
        to_modality: int,
    ) -> torch.Tensor:
        """Translate between modalities via the common space.

        Encode from one modality, decode in another's style.

        Args:
            x: Input from source modality.
            from_modality: Source modality index.
            to_modality: Target modality index (for future per-modality decoders).

        Returns:
            Common-space representation (usable by target modality).
        """
        return self.encode(x, from_modality)

    def summary(self) -> str:
        lines = [f"BridgedModel ({len(self.modality_names)} modalities)"]
        for i, name in enumerate(self.modality_names):
            enc_params = sum(p.numel() for p in self.encoders[i].parameters())
            lines.append(f"  [{i}] {name}: encoder = {enc_params:,} params")
        dec_params = sum(p.numel() for p in self.decoder.parameters())
        lines.append(f"  Decoder: {dec_params:,} params")
        lines.append(f"  Common space: {self.common_dim}-dim")
        total = sum(p.numel() for p in self.parameters())
        lines.append(f"  Total: {total:,} params")
        return "\n".join(lines)


class ModalityBridge:
    """Merge models that operate on DIFFERENT data types.

    Uses paired data to learn a common representation space where
    different modalities can be aligned and merged.

    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],
        paired_data: List[np.ndarray],
        modality_names: Optional[List[str]] = None,
        output_dim: Optional[int] = None,
        decoder_hidden: int = 512,
        epochs: int = 100,
        lr: float = 1e-3,
        batch_size: int = 32,
        temperature: float = 0.07,
    ) -> BridgedModel:
        """Merge models from different modalities.

        Args:
            models: List of models (one per modality).
            paired_data: List of numpy arrays, one per modality.
                paired_data[i][j] and paired_data[k][j] describe
                the SAME concept in different modalities.
            modality_names: Names for each modality.
            output_dim: Output dimension of the decoder.
            decoder_hidden: Hidden size for the decoder.
            epochs: Training epochs for alignment.
            lr: Learning rate.
            batch_size: Training batch size.
            temperature: Temperature for contrastive alignment loss.

        Returns:
            A BridgedModel with aligned modality encoders.
        """
        n = len(models)
        if n < 2:
            raise ValueError("Need at least 2 models from different modalities")
        if len(paired_data) != n:
            raise ValueError(
                f"Need paired data for each model: got {len(paired_data)} "
                f"data arrays for {n} models"
            )

        # Validate paired data has same number of samples
        num_samples = len(paired_data[0])
        for i, data in enumerate(paired_data):
            if len(data) != num_samples:
                raise ValueError(
                    f"All paired data must have same number of samples. "
                    f"Modality 0 has {num_samples}, modality {i} has {len(data)}"
                )

        modality_names = modality_names or [f"modality_{i}" for i in range(n)]

        if self.verbose:
            print(f"[ModalityBridge] Bridging {n} modalities: {modality_names}")
            print(f"  Paired samples: {num_samples}")

        # Step 1: Collect outputs from each model on its own data
        if self.verbose:
            print("[ModalityBridge] Step 1/3: Collecting modality outputs...")

        model_outputs = []
        output_dims = []
        for i, (model, data) in enumerate(zip(models, paired_data)):
            model.eval()
            x = torch.from_numpy(data).float()
            with torch.no_grad():
                out = model(x)
            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"  {modality_names[i]}: output dim = {output_dims[-1]}")

        # Step 2: Train modality encoders with contrastive alignment
        if self.verbose:
            print("[ModalityBridge] Step 2/3: Learning cross-modal alignment...")

        encoders = nn.ModuleList([
            ModalityEncoder(odim, self.common_dim) for odim in output_dims
        ])

        optimizer = torch.optim.Adam(encoders.parameters(), lr=lr)

        for epoch in range(epochs):
            # Shuffle indices for mini-batching
            indices = torch.randperm(num_samples)
            total_loss = 0.0
            num_batches = 0

            for start in range(0, num_samples, batch_size):
                end = min(start + batch_size, num_samples)
                batch_idx = indices[start:end]
                bs = len(batch_idx)

                optimizer.zero_grad()

                # Encode each modality's batch
                encoded = [
                    encoders[i](model_outputs[i][batch_idx])
                    for i in range(n)
                ]

                # Contrastive alignment loss (CLIP-style)
                # Paired samples should be close, non-paired should be far
                loss = torch.tensor(0.0)
                pairs = 0
                for i in range(n):
                    for j in range(i + 1, n):
                        # Cosine similarity matrix
                        sim = torch.mm(encoded[i], encoded[j].t()) / temperature

                        # Diagonal = paired samples (should be high)
                        labels = torch.arange(bs)
                        loss_ij = F.cross_entropy(sim, labels)
                        loss_ji = F.cross_entropy(sim.t(), labels)
                        loss += (loss_ij + loss_ji) / 2
                        pairs += 1

                if pairs > 0:
                    loss = loss / pairs

                loss.backward()
                optimizer.step()
                total_loss += loss.item()
                num_batches += 1

            if self.verbose and (epoch + 1) % 20 == 0:
                avg = total_loss / max(num_batches, 1)
                print(f"  Epoch {epoch+1}/{epochs}: alignment loss = {avg:.4f}")

        # Step 3: Build decoder from common space
        if self.verbose:
            print("[ModalityBridge] Step 3/3: Building decoder...")

        if output_dim is None:
            output_dim = max(output_dims)

        decoder = nn.Sequential(
            nn.Linear(self.common_dim, decoder_hidden),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(decoder_hidden, decoder_hidden // 2),
            nn.ReLU(),
            nn.Linear(decoder_hidden // 2, output_dim),
        )

        # Train decoder on merged common-space representations
        with torch.no_grad():
            # Get all encodings
            all_encoded = [encoders[i](model_outputs[i]) for i in range(n)]
            # Average all modalities in common space = merged representation
            merged_common = torch.stack(all_encoded).mean(dim=0)

            # Target: average of padded original outputs
            padded = []
            for out in model_outputs:
                if out.shape[-1] < output_dim:
                    padded.append(F.pad(out, (0, output_dim - out.shape[-1])))
                elif out.shape[-1] > output_dim:
                    padded.append(out[:, :output_dim])
                else:
                    padded.append(out)
            merged_target = torch.stack(padded).mean(dim=0)

        dec_optimizer = torch.optim.Adam(decoder.parameters(), lr=lr)
        dataset = torch.utils.data.TensorDataset(merged_common, merged_target)
        loader = torch.utils.data.DataLoader(
            dataset, batch_size=batch_size, shuffle=True
        )

        best_loss = float("inf")
        best_state = None

        decoder.train()
        for epoch in range(epochs):
            total_loss = 0.0
            for x_batch, y_batch in loader:
                dec_optimizer.zero_grad()
                pred = decoder(x_batch)
                loss = F.mse_loss(pred, y_batch)
                loss.backward()
                dec_optimizer.step()
                total_loss += loss.item()

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

        if best_state is not None:
            decoder.load_state_dict(best_state)
        decoder.eval()

        # Build the bridged model
        bridged = BridgedModel(
            encoders=encoders,
            decoder=decoder,
            common_dim=self.common_dim,
            modality_names=modality_names,
        )

        if self.verbose:
            print(f"[ModalityBridge] ✓ Bridge complete!")
            print(bridged.summary())

        return bridged

    def unsupervised_merge(
        self,
        models: List[nn.Module],
        unpaired_data: List[np.ndarray],
        modality_names: Optional[List[str]] = None,
        output_dim: Optional[int] = None,
        decoder_hidden: int = 512,
        epochs: int = 100,
        lr: float = 1e-3,
        batch_size: int = 32,
    ) -> BridgedModel:
        """Merge models from different modalities WITHOUT paired data.

        Uses adversarial alignment: a discriminator tries to tell which
        modality an encoding came from. The encoders learn to fool it.
        When the discriminator can't tell the difference, the spaces
        are aligned — no dictionary needed.

        Args:
            models: List of models (one per modality).
            unpaired_data: List of numpy arrays, one per modality.
                These do NOT need to be paired — each model gets
                its own independent data.
            modality_names: Names for each modality.
            output_dim: Output dimension of the decoder.
            decoder_hidden: Hidden size for the decoder.
            epochs: Training epochs for alignment.
            lr: Learning rate.
            batch_size: Training batch size.

        Returns:
            A BridgedModel with adversarially aligned modality encoders.
        """
        n = len(models)
        if n < 2:
            raise ValueError("Need at least 2 models from different modalities")
        if len(unpaired_data) != n:
            raise ValueError(
                f"Need data for each model: got {len(unpaired_data)} "
                f"data arrays for {n} models"
            )

        modality_names = modality_names or [f"modality_{i}" for i in range(n)]

        if self.verbose:
            print(f"[ModalityBridge] Unsupervised bridging of {n} modalities")
            print(f"  No paired data needed!")

        # Step 1: Collect outputs
        if self.verbose:
            print("[ModalityBridge] Step 1/3: Collecting modality outputs...")

        model_outputs = []
        output_dims = []
        for i, (model, data) in enumerate(zip(models, unpaired_data)):
            model.eval()
            x = torch.from_numpy(data).float()
            with torch.no_grad():
                out = model(x)
            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"  {modality_names[i]}: {len(data)} samples, "
                      f"output dim = {output_dims[-1]}")

        # Step 2: Adversarial alignment
        if self.verbose:
            print("[ModalityBridge] Step 2/3: Adversarial alignment "
                  "(no pairs needed)...")

        encoders = nn.ModuleList([
            ModalityEncoder(odim, self.common_dim) for odim in output_dims
        ])

        # Discriminator: tries to guess which modality an encoding came from
        discriminator = nn.Sequential(
            nn.Linear(self.common_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, n),  # n classes = n modalities
        )

        enc_optimizer = torch.optim.Adam(encoders.parameters(), lr=lr)
        disc_optimizer = torch.optim.Adam(discriminator.parameters(), lr=lr)

        min_samples = min(len(d) for d in model_outputs)

        for epoch in range(epochs):
            # --- Train discriminator: learn to tell modalities apart ---
            disc_optimizer.zero_grad()

            disc_loss = torch.tensor(0.0)
            for i in range(n):
                idx = torch.randint(0, len(model_outputs[i]), (batch_size,))
                encoded = encoders[i](model_outputs[i][idx]).detach()
                pred = discriminator(encoded)
                labels = torch.full((batch_size,), i, dtype=torch.long)
                disc_loss += F.cross_entropy(pred, labels)
            disc_loss = disc_loss / n

            disc_loss.backward()
            disc_optimizer.step()

            # --- Train encoders: fool the discriminator ---
            enc_optimizer.zero_grad()

            fool_loss = torch.tensor(0.0)
            for i in range(n):
                idx = torch.randint(0, len(model_outputs[i]), (batch_size,))
                encoded = encoders[i](model_outputs[i][idx])
                pred = discriminator(encoded)
                # Fool: make discriminator think ALL encodings are from
                # a uniform distribution (equal probability for each modality)
                uniform = torch.ones(batch_size, n) / n
                fool_loss += -(uniform * F.log_softmax(pred, dim=-1)).sum(dim=-1).mean()

            # Distribution matching: all encoders should produce
            # similar statistical distributions
            means = []
            for i in range(n):
                idx = torch.randint(0, len(model_outputs[i]), (batch_size,))
                encoded = encoders[i](model_outputs[i][idx])
                means.append(encoded.mean(dim=0))

            dist_loss = torch.tensor(0.0)
            count = 0
            for i in range(n):
                for j in range(i + 1, n):
                    dist_loss += F.mse_loss(means[i], means[j])
                    count += 1
            if count > 0:
                dist_loss = dist_loss / count

            total_enc_loss = fool_loss / n + 0.5 * dist_loss
            total_enc_loss.backward()
            enc_optimizer.step()

            if self.verbose and (epoch + 1) % 20 == 0:
                # Check discriminator accuracy (lower = better alignment)
                with torch.no_grad():
                    correct = 0
                    total = 0
                    for i in range(n):
                        idx = torch.randint(0, len(model_outputs[i]),
                                            (min(50, len(model_outputs[i])),))
                        encoded = encoders[i](model_outputs[i][idx])
                        pred = discriminator(encoded).argmax(dim=-1)
                        correct += (pred == i).sum().item()
                        total += len(idx)
                    acc = correct / total
                print(f"  Epoch {epoch+1}/{epochs}: "
                      f"disc_acc={acc:.2%} (lower=better aligned)")

        # Step 3: Build decoder
        if self.verbose:
            print("[ModalityBridge] Step 3/3: Building decoder...")

        if output_dim is None:
            output_dim = max(output_dims)

        decoder = nn.Sequential(
            nn.Linear(self.common_dim, decoder_hidden),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(decoder_hidden, decoder_hidden // 2),
            nn.ReLU(),
            nn.Linear(decoder_hidden // 2, output_dim),
        )

        # Train decoder on merged common-space representations
        with torch.no_grad():
            all_encoded = [encoders[i](model_outputs[i]) for i in range(n)]
            # Use all data from all modalities
            merged_parts = []
            target_parts = []
            for i in range(n):
                merged_parts.append(all_encoded[i])
                out = model_outputs[i]
                if out.shape[-1] < output_dim:
                    out = F.pad(out, (0, output_dim - out.shape[-1]))
                elif out.shape[-1] > output_dim:
                    out = out[:, :output_dim]
                target_parts.append(out)

            merged_common = torch.cat(merged_parts, dim=0)
            merged_target = torch.cat(target_parts, dim=0)

        dec_optimizer = torch.optim.Adam(decoder.parameters(), lr=lr)
        dataset = torch.utils.data.TensorDataset(merged_common, merged_target)
        loader = torch.utils.data.DataLoader(
            dataset, batch_size=batch_size, shuffle=True
        )

        best_loss = float("inf")
        best_state = None

        decoder.train()
        for epoch in range(epochs):
            total_loss = 0.0
            for x_batch, y_batch in loader:
                dec_optimizer.zero_grad()
                pred = decoder(x_batch)
                loss = F.mse_loss(pred, y_batch)
                loss.backward()
                dec_optimizer.step()
                total_loss += loss.item()

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

        if best_state is not None:
            decoder.load_state_dict(best_state)
        decoder.eval()

        bridged = BridgedModel(
            encoders=encoders,
            decoder=decoder,
            common_dim=self.common_dim,
            modality_names=modality_names,
        )

        if self.verbose:
            print(f"[ModalityBridge] ✓ Unsupervised bridge complete!")
            print(bridged.summary())

        return bridged

