"""
LLMCrossMerger — merge LLMs of ANY architecture into one.

Unlike weight-merging (which requires same architecture), LLMCrossMerger
merges what models KNOW by distilling knowledge from multiple teacher LLMs
into a single student model.

How it works:
1. Load each teacher one at a time → extract hidden state representations
2. Project all representations to a common dimension
3. Merge via weighted average in the common space
4. Fine-tune a student LLM to match the merged representations

Result: a single HuggingFace-compatible model with combined knowledge from
all source models, regardless of their original architectures.

Usage::

    from model_osmosis.merging.llm_cross_merger import LLMCrossMerger

    merger = LLMCrossMerger()
    model, tokenizer = merger.merge(
        teacher_ids=["Qwen/Qwen2.5-Coder-1.5B", "deepseek-ai/deepseek-coder-1.3b-base"],
        sample_texts=["def fibonacci(n):", "class Node:", ...],
        weights=[0.6, 0.4],
    )
    # model is a full HuggingFace PreTrainedModel
"""

from __future__ import annotations

import gc
import math
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple, Union

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


@dataclass
class LLMCrossMergeConfig:
    """Configuration for LLM cross-architecture merging."""

    student_model_id: Optional[str] = None
    """HF repo ID for the student base model. None = use first teacher."""

    common_dim: int = 256
    """Dimension of the common projection space."""

    distillation_epochs: int = 3
    """Number of training epochs over the sample texts."""

    learning_rate: float = 2e-5
    """AdamW learning rate for student fine-tuning."""

    batch_size: int = 4
    """Micro-batch size for distillation."""

    max_seq_len: int = 512
    """Maximum token sequence length."""

    gradient_accumulation_steps: int = 8
    """Effective batch = batch_size * gradient_accumulation_steps."""

    warmup_ratio: float = 0.1
    """Fraction of total steps for LR warmup."""

    lm_loss_weight: float = 0.1
    """Weight for language modeling loss (prevents catastrophic forgetting)."""

    projection_epochs: int = 50
    """Epochs to train the alignment projectors."""

    device: str = "auto"
    """Device: 'cpu', 'cuda', 'auto'."""

    torch_dtype: str = "float16"
    """Precision for teacher loading: 'float16', 'bfloat16', 'float32'."""

    use_lora: bool = True
    """Whether to use LoRA for distillation (saves VRAM)."""

    lora_r: int = 8
    """LoRA rank."""

    lora_alpha: int = 32
    """LoRA alpha."""

    verbose: bool = True


class HiddenStateCollector:
    """Load a teacher model one at a time, extract hidden states, free memory.

    This ensures only ONE large model is on GPU/CPU at a time, keeping
    memory usage manageable for consumer hardware.
    """

    def __init__(
        self,
        device: torch.device,
        torch_dtype: torch.dtype,
        max_seq_len: int,
        verbose: bool = True,
    ) -> None:
        self.device = device
        self.torch_dtype = torch_dtype
        self.max_seq_len = max_seq_len
        self.verbose = verbose

    def collect(
        self,
        model_id: str,
        texts: List[str],
        batch_size: int = 4,
    ) -> Tuple[torch.Tensor, int]:
        """Extract hidden states from a teacher model.

        Loads the model, runs inference on all texts, extracts the last
        hidden state (mean-pooled over sequence), then unloads the model.

        Args:
            model_id: HuggingFace model ID or local path.
            texts: Sample texts to run through the model.
            batch_size: Inference batch size.

        Returns:
            hidden_states: (N, hidden_size) tensor on CPU.
            hidden_size: The model's hidden dimension.
        """
        from transformers import AutoModelForCausalLM, AutoTokenizer

        if self.verbose:
            print(f"  Loading teacher: {model_id}")

        # Load tokenizer
        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        # Load model
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=self.torch_dtype,
            trust_remote_code=True,
            device_map=str(self.device) if self.device.type == "cuda" else None,
        )
        if self.device.type != "cuda":
            model = model.to(self.device)
        model.eval()

        hidden_size = model.config.hidden_size
        if self.verbose:
            total_params = sum(p.numel() for p in model.parameters())
            print(f"    Parameters: {total_params:,}, hidden_size: {hidden_size}")

        # Collect hidden states batch by batch
        all_hidden = []
        num_batches = math.ceil(len(texts) / batch_size)

        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i : i + batch_size]

            inputs = tokenizer(
                batch_texts,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=self.max_seq_len,
            ).to(self.device)

            with torch.no_grad():
                outputs = model(
                    **inputs,
                    output_hidden_states=True,
                )

            # Get last hidden state: (batch, seq_len, hidden_size)
            last_hidden = outputs.hidden_states[-1]

            # Mean-pool over sequence (masked)
            mask = inputs["attention_mask"].unsqueeze(-1).float()
            pooled = (last_hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)

            all_hidden.append(pooled.cpu().float())

            if self.verbose and (i // batch_size + 1) % 10 == 0:
                print(f"    Batch {i // batch_size + 1}/{num_batches}")

        # Concatenate all hidden states
        hidden_states = torch.cat(all_hidden, dim=0)  # (N, hidden_size)

        if self.verbose:
            print(f"    Collected: {hidden_states.shape}")

        # Free memory
        del model, tokenizer, inputs, outputs
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()

        if self.verbose:
            print(f"    Teacher unloaded.")

        return hidden_states, hidden_size


class LLMCrossMerger:
    """Merge LLMs of different architectures into one.

    Works by extracting hidden state representations from multiple
    teacher models, projecting them to a common space, merging,
    and then distilling into a student LLM.

    The student is a real HuggingFace transformer model that can be
    used with model.generate(), saved with save_pretrained(), and
    works with the obfuscation pipeline.

    Args:
        config: Configuration for the merger.
    """

    def __init__(self, config: Optional[LLMCrossMergeConfig] = None) -> None:
        self.config = config or LLMCrossMergeConfig()

    def merge(
        self,
        teacher_ids: List[str],
        sample_texts: List[str],
        weights: Optional[List[float]] = None,
        student_model_id: Optional[str] = None,
    ) -> Tuple[Any, Any]:
        """Merge multiple LLMs of any architecture into one.

        Args:
            teacher_ids: List of HuggingFace model IDs.
            sample_texts: Sample text corpus for knowledge extraction.
                More diverse samples = better merge quality.
            weights: Per-teacher importance weights (default: equal).
            student_model_id: HF model ID for the student base.
                Defaults to the first teacher.

        Returns:
            (model, tokenizer) — a full HuggingFace-compatible model
            and its tokenizer.
        """
        cfg = self.config
        n = len(teacher_ids)

        if n == 0:
            raise ValueError("Need at least one teacher model")
        if len(sample_texts) == 0:
            raise ValueError("Need sample texts for knowledge extraction")

        # Resolve device
        device = self._resolve_device(cfg.device)
        dtype = self._resolve_dtype(cfg.torch_dtype)

        if cfg.verbose:
            print(f"[LLMCrossMerger] Merging {n} LLMs")
            print(f"  Teachers: {teacher_ids}")
            print(f"  Sample texts: {len(sample_texts)}")
            print(f"  Device: {device}, dtype: {dtype}")

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

        # Step 1: Collect hidden states from each teacher
        if cfg.verbose:
            print(f"\n[Step 1/{4}] Collecting hidden states from teachers...")

        collector = HiddenStateCollector(
            device=device,
            torch_dtype=dtype,
            max_seq_len=cfg.max_seq_len,
            verbose=cfg.verbose,
        )

        teacher_hiddens = []  # List of (N, hidden_size_i)
        teacher_dims = []  # List of hidden_size_i

        for i, teacher_id in enumerate(teacher_ids):
            if cfg.verbose:
                print(f"\n  Teacher {i + 1}/{n}:")
            hidden_states, hidden_size = collector.collect(
                model_id=teacher_id,
                texts=sample_texts,
                batch_size=cfg.batch_size,
            )
            teacher_hiddens.append(hidden_states)
            teacher_dims.append(hidden_size)

        # Step 2: Project to common space and merge
        if cfg.verbose:
            print(f"\n[Step 2/{4}] Projecting to common space and merging...")
            print(f"  Teacher dimensions: {teacher_dims}")
            print(f"  Common dimension: {cfg.common_dim}")

        merged_targets = self._project_and_merge(
            teacher_hiddens, teacher_dims, weights
        )

        if cfg.verbose:
            print(f"  Merged targets: {merged_targets.shape}")

        # Step 3: Initialize student
        student_id = student_model_id or cfg.student_model_id or teacher_ids[0]
        if cfg.verbose:
            print(f"\n[Step 3/{4}] Initializing student from: {student_id}")

        student, student_tokenizer = self._initialize_student(student_id, device, dtype)

        # Step 4: Distill merged knowledge into student
        if cfg.verbose:
            print(f"\n[Step 4/{4}] Distilling merged knowledge into student...")

        student = self._distill(
            student=student,
            tokenizer=student_tokenizer,
            sample_texts=sample_texts,
            merged_targets=merged_targets,
            device=device,
        )

        student.eval()
        if cfg.verbose:
            total_params = sum(p.numel() for p in student.parameters())
            print(f"\n[LLMCrossMerger] Done! Student model: {total_params:,} params")

        return student, student_tokenizer

    def _project_and_merge(
        self,
        teacher_hiddens: List[torch.Tensor],
        teacher_dims: List[int],
        weights: List[float],
    ) -> torch.Tensor:
        """Project teacher hidden states to common space and merge.

        Trains linear projectors to align representations, then
        creates weighted-average merged targets.

        Returns:
            merged: (N, common_dim) tensor.
        """
        cfg = self.config
        n = len(teacher_hiddens)
        common_dim = cfg.common_dim

        # Create projectors
        projectors = nn.ModuleList([
            nn.Linear(dim, common_dim) for dim in teacher_dims
        ])

        # Train projectors to align representations
        optimizer = torch.optim.Adam(projectors.parameters(), lr=1e-3)

        for epoch in range(cfg.projection_epochs):
            optimizer.zero_grad()

            projected = [
                projectors[i](teacher_hiddens[i]) for i in range(n)
            ]

            # Alignment loss: projections should be close
            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

            # Diversity: prevent collapse
            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()
            optimizer.step()

        if cfg.verbose:
            print(f"  Projection alignment loss: {alignment_loss.item():.6f}")

        # Create merged targets
        with torch.no_grad():
            projected = [projectors[i](teacher_hiddens[i]) for i in range(n)]
            merged = torch.zeros_like(projected[0])
            for i in range(n):
                merged += weights[i] * projected[i]

        return merged.detach()

    def _initialize_student(
        self,
        model_id: str,
        device: torch.device,
        dtype: torch.dtype,
    ) -> Tuple[Any, Any]:
        """Load the student model and tokenizer."""
        from transformers import AutoModelForCausalLM, AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=dtype,
            trust_remote_code=True,
        )

        if self.config.verbose:
            total = sum(p.numel() for p in model.parameters())
            print(f"  Student loaded: {total:,} params, hidden_size: {model.config.hidden_size}")

        return model, tokenizer

    def _distill(
        self,
        student: Any,
        tokenizer: Any,
        sample_texts: List[str],
        merged_targets: torch.Tensor,
        device: torch.device,
    ) -> Any:
        """Fine-tune student to match merged hidden state targets.

        Uses MSE loss on pooled hidden states + optional LM loss
        to prevent catastrophic forgetting.
        """
        cfg = self.config

        student_hidden_size = student.config.hidden_size
        common_dim = cfg.common_dim

        # Temporary projection head for distillation
        projector = nn.Linear(student_hidden_size, common_dim)
        projector = projector.to(device)

        # Move student to device
        student = student.to(device)

        # Apply LoRA if requested and on GPU
        if cfg.use_lora and device.type == "cuda":
            try:
                from peft import LoraConfig, get_peft_model
                lora_config = LoraConfig(
                    r=cfg.lora_r,
                    lora_alpha=cfg.lora_alpha,
                    target_modules="all-linear",
                    lora_dropout=0.05,
                    bias="none",
                    task_type="CAUSAL_LM",
                )
                if cfg.verbose:
                    print(f"  [Pipeline] Applying LoRA (r={cfg.lora_r})...")
                student = get_peft_model(student, lora_config)
                student.print_trainable_parameters()
            except ImportError:
                if cfg.verbose:
                    print("  [Warning] peft not installed — skipping LoRA")

        # Enable gradient checkpointing to save memory
        if device.type == "cuda":
            try:
                student.gradient_checkpointing_enable()
                if cfg.verbose:
                    print("  [Pipeline] Gradient checkpointing enabled")
            except Exception as e:
                if cfg.verbose:
                    print(f"  [Pipeline] Gradient checkpointing not available: {e}")

        student.train()

        # Tokenize all sample texts
        encodings = tokenizer(
            sample_texts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=cfg.max_seq_len,
        )

        input_ids = encodings["input_ids"]
        attention_mask = encodings["attention_mask"]

        # Create dataset
        dataset = torch.utils.data.TensorDataset(
            input_ids,
            attention_mask,
            merged_targets,
        )
        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=cfg.batch_size,
            shuffle=True,
        )

        # Optimizer
        all_params = list(student.parameters()) + list(projector.parameters())
        trainable_params = [p for p in all_params if p.requires_grad]

        if not trainable_params:
            if cfg.verbose:
                print("  [Warning] No trainable parameters found in student. Only training projection head.")
            trainable_params = list(projector.parameters())

        optimizer = torch.optim.AdamW(trainable_params, lr=cfg.learning_rate, weight_decay=0.01)

        # Learning rate scheduler with warmup
        total_steps = len(dataloader) * cfg.distillation_epochs
        warmup_steps = int(total_steps * cfg.warmup_ratio)

        def lr_lambda(step):
            if step < warmup_steps:
                return float(step) / float(max(1, warmup_steps))
            progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))
            return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))

        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

        # Training loop
        global_step = 0
        accum_steps = cfg.gradient_accumulation_steps

        for epoch in range(cfg.distillation_epochs):
            epoch_loss = 0.0
            epoch_mse = 0.0
            epoch_lm = 0.0
            num_batches = 0

            for batch_idx, (ids, mask, targets) in enumerate(dataloader):
                ids = ids.to(device)
                mask = mask.to(device)
                targets = targets.to(device)

                # Forward pass with hidden states
                outputs = student(
                    input_ids=ids,
                    attention_mask=mask,
                    labels=ids,  # For LM loss
                    output_hidden_states=True,
                )

                # Get last hidden state and mean-pool
                last_hidden = outputs.hidden_states[-1]  # (B, seq, hidden)
                mask_expanded = mask.unsqueeze(-1).float()
                pooled = (last_hidden * mask_expanded).sum(dim=1) / mask_expanded.sum(dim=1).clamp(min=1)

                # Project to common dim
                projected = projector(pooled)  # (B, common_dim)

                # MSE loss on merged targets
                mse_loss = F.mse_loss(projected, targets)

                # LM loss (prevents forgetting)
                lm_loss = outputs.loss if outputs.loss is not None else torch.tensor(0.0, device=device)

                # Combined loss
                loss = mse_loss + cfg.lm_loss_weight * lm_loss
                loss = loss / accum_steps

                loss.backward()

                if (batch_idx + 1) % accum_steps == 0 or (batch_idx + 1) == len(dataloader):
                    torch.nn.utils.clip_grad_norm_(all_params, 1.0)
                    optimizer.step()
                    scheduler.step()
                    optimizer.zero_grad()
                    global_step += 1

                epoch_loss += loss.item() * accum_steps
                epoch_mse += mse_loss.item()
                epoch_lm += lm_loss.item()
                num_batches += 1

            avg_loss = epoch_loss / num_batches
            avg_mse = epoch_mse / num_batches
            avg_lm = epoch_lm / num_batches

            if cfg.verbose:
                print(
                    f"  Epoch {epoch + 1}/{cfg.distillation_epochs}: "
                    f"loss={avg_loss:.4f} (mse={avg_mse:.4f}, lm={avg_lm:.4f})"
                )

        # Clean up: remove projector, move student back to half precision
        del projector
        
        # Merge LoRA weights if applicable
        if hasattr(student, "merge_and_unload"):
            if cfg.verbose:
                print("  [Pipeline] Merging LoRA weights...")
            student = student.merge_and_unload()
            
        try:
            student = student.half()
        except (ValueError, TypeError):
            pass
        student.eval()

        # Move to CPU to save memory
        student = student.cpu()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()

        return student

    @staticmethod
    def _resolve_device(device_str: str) -> torch.device:
        """Resolve device string to torch.device."""
        if device_str == "auto":
            if torch.cuda.is_available():
                return torch.device("cuda")
            return torch.device("cpu")
        return torch.device(device_str)

    @staticmethod
    def _resolve_dtype(dtype_str: str) -> torch.dtype:
        """Resolve dtype string to torch.dtype."""
        dtypes = {
            "float16": torch.float16,
            "bfloat16": torch.bfloat16,
            "float32": torch.float32,
        }
        return dtypes.get(dtype_str, torch.float16)
