"""
Merge #6: Klear-Reasoner-8B → existing 5-model merge.

Klear-Reasoner-8B is based on Qwen3-8B-Base (SAME architecture as our
Qwen3-VL-8B target). This means:
    - Direct 1:1 layer matching (36 → 36)
    - Neurons are aligned (identity Q matrices)
    - No permutation needed
    - Transport plans take under 1 minute
    - Total merge: ~10-15 minutes

Think of this like adding more reasoning spice to an already-cooked dish.
The base model already has reasoning from DeepSeek, MiMo, and Falcon.
Klear-Reasoner adds EVEN MORE focused reasoning (90.5% AIME 2024).

Alpha = 0.3 (conservative — we don't want to overwrite what's already
merged). The existing model is the dominant one; Klear adds a boost.

Usage:
    python -m td_fuse.merge_klear

Expects:
    - Existing merge at: td_fuse_outputs/healed_final/ (or after_falcon_5model/)
    - Will download Klear-Reasoner-8B from HuggingFace
    - Output: td_fuse_outputs/6model_klear/
"""

import os
import sys
import gc
import time
import torch
import numpy as np
from pathlib import Path

# ============================================================================
# CONFIGURATION
# ============================================================================

# Where is our existing 5-model merge?
EXISTING_MERGE_PATHS = [
    "td_fuse_outputs/healed_final",
    "td_fuse_checkpoints/after_falcon",       # backup location
    "td_fuse_checkpoints/after_falcon_5model", # HF backup
]

# Klear-Reasoner-8B config
KLEAR_HF_ID = "Kwai-Klear/Klear-Reasoner-8B"
KLEAR_ALPHA = 0.3     # Conservative: keep 70% existing, add 30% Klear reasoning
KLEAR_NAME = "Klear-Reasoner-8B"

# Output
OUTPUT_DIR = "td_fuse_outputs/6model_klear"

# Vision encoder prefixes — NEVER merge into these
VISION_SKIP = ["visual", "merger", "model.visual", "model.merger"]

# Thinking tokens to protect
THINK_TOKEN_IDS = [151667, 151668]


def find_existing_merge():
    """Find the existing 5-model merge checkpoint."""
    for path in EXISTING_MERGE_PATHS:
        safetensors = list(Path(path).glob("*.safetensors")) if Path(path).exists() else []
        if safetensors:
            total_size = sum(f.stat().st_size for f in safetensors)
            if total_size > 1_000_000_000:  # > 1GB
                print(f"[merge6] Found existing merge at: {path} ({total_size/1e9:.1f} GB)")
                return path
    return None


def load_model(path, dtype=torch.bfloat16):
    """Load model — auto-detects Qwen3-VL."""
    from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig

    print(f"[merge6] Loading model from: {path}")
    t = time.time()

    try:
        config = AutoConfig.from_pretrained(path, trust_remote_code=True)
        model_type = getattr(config, 'model_type', '')
        config_class = type(config).__name__.lower()
        if 'qwen3_vl' in model_type or 'qwen3vl' in config_class:
            from transformers import Qwen3VLForConditionalGeneration
            print(f"[merge6] Loading as Qwen3-VL model")
            model = Qwen3VLForConditionalGeneration.from_pretrained(
                path, device_map="auto", torch_dtype=dtype
            )
        else:
            model = AutoModelForCausalLM.from_pretrained(
                path, device_map="auto", torch_dtype=dtype
            )
    except Exception as e:
        print(f"[merge6] Auto-detect failed ({e}), using AutoModelForCausalLM")
        model = AutoModelForCausalLM.from_pretrained(
            path, device_map="auto", torch_dtype=dtype
        )

    tokenizer = AutoTokenizer.from_pretrained(path)
    print(f"[merge6] Loaded in {time.time()-t:.0f}s")
    return model, tokenizer


def should_skip(key):
    """Check if a parameter should be skipped during merge."""
    # Skip vision encoder
    for prefix in VISION_SKIP:
        if key.startswith(prefix):
            return True
    return False


def direct_merge(target_model, source_state, alpha=0.3):
    """
    Direct same-architecture merge: blend source weights into target.

    Since Klear-Reasoner-8B has the same Qwen3-8B architecture:
    - All parameter names match exactly
    - No layer mapping needed
    - No neuron permutation needed
    - Just simple alpha blending: W = alpha * source + (1-alpha) * target

    We skip:
    - Vision encoder weights (visual.*, merger.*)
    - Embedding/lm_head IF vocab sizes don't match
    """
    target_state = target_model.state_dict()

    # Check vocab sizes — try multiple key patterns
    target_vocab = None
    source_vocab = None
    for key in target_state:
        if "embed_tokens" in key:
            target_vocab = target_state[key].shape[0]
            break
    for key in source_state:
        if "embed_tokens" in key:
            source_vocab = source_state[key].shape[0]
            break
    print(f"[merge6] Target vocab: {target_vocab}, Source vocab: {source_vocab}")
    skip_embeddings = (target_vocab != source_vocab) if (target_vocab and source_vocab) else False
    if skip_embeddings:
        print(f"[merge6] Vocab mismatch ({target_vocab} vs {source_vocab}) — skipping embeddings")
    else:
        print(f"[merge6] Vocab sizes match — merging embeddings too")

    fused = 0
    skipped = 0
    total = len(target_state)

    # Debug: show first 5 keys from each to verify naming
    target_keys = list(target_state.keys())
    source_keys = list(source_state.keys())
    print(f"[merge6] Target has {len(target_keys)} params, Source has {len(source_keys)} params")
    print(f"[merge6] Target first 5 keys: {target_keys[:5]}")
    print(f"[merge6] Source first 5 keys: {source_keys[:5]}")

    # Save original think token embeddings for restoration
    think_embeds = {}
    for key in target_state:
        if "embed_tokens" in key:
            for tid in THINK_TOKEN_IDS:
                if tid < target_state[key].shape[0]:
                    think_embeds[(key, tid)] = target_state[key][tid].clone()

    for i, target_key in enumerate(target_state):
        if should_skip(target_key):
            skipped += 1
            continue

        if skip_embeddings and ("embed_tokens" in target_key or "lm_head" in target_key):
            skipped += 1
            continue

        # Find source key — handle Qwen3-VL prefix differences
        # Qwen3-VL keys: model.language_model.layers.0... or model.model.layers.0...
        # Klear keys:    model.layers.0...
        source_key = None
        candidates = [
            target_key,                                              # exact match
            target_key.replace("language_model.", ""),               # strip language_model.
            target_key.replace("model.model.", "model."),            # model.model.X -> model.X
            target_key.replace("model.", "", 1),                     # strip first model.
            f"model.{target_key}",                                   # add model. prefix
        ]
        for candidate in candidates:
            if candidate in source_state:
                source_key = candidate
                break

        if source_key is None:
            skipped += 1
            if skipped <= 10:
                print(f"  [skip] No source match: {target_key}")
            continue

        target_w = target_state[target_key]
        source_w = source_state[source_key]

        # Shape check
        if target_w.shape != source_w.shape:
            skipped += 1
            if skipped <= 5:
                print(f"  [skip] Shape mismatch: {target_key} ({target_w.shape} vs {source_w.shape})")
            continue

        # Blend: W_final = alpha * source + (1-alpha) * target
        fused_w = alpha * source_w.to(target_w.device, dtype=target_w.dtype) + (1 - alpha) * target_w
        target_state[target_key] = fused_w
        fused += 1

        if fused % 50 == 0:
            print(f"  Fused {fused}/{total} params...")
            sys.stdout.flush()

    # Restore think token embeddings
    for (key, tid), embed in think_embeds.items():
        if key in target_state and tid < target_state[key].shape[0]:
            target_state[key][tid] = embed
            print(f"[merge6] Protected think token {tid}")

    # Load fused state
    missing, unexpected = target_model.load_state_dict(target_state, strict=False)
    if missing:
        print(f"[merge6] {len(missing)} missing keys (likely vision quant params — safe)")
    print(f"[merge6] Fused {fused} params, skipped {skipped}")
    return target_model


def run_merge():
    """Main merge pipeline for Klear-Reasoner-8B."""
    merge_start = time.time()

    print("\n" + "=" * 60)
    print("MERGE #6: Klear-Reasoner-8B → 5-model merge")
    print(f"Alpha: {KLEAR_ALPHA} (30% Klear, 70% existing)")
    print(f"Started at: {time.strftime('%H:%M:%S')}")
    print("=" * 60)
    sys.stdout.flush()

    # --- Step 1: Find existing merge ---
    existing_path = find_existing_merge()
    if not existing_path:
        print("[merge6] ERROR: Cannot find existing 5-model merge!")
        print(f"[merge6] Looked in: {EXISTING_MERGE_PATHS}")
        sys.exit(1)

    # --- Step 2: Load existing merged model ---
    print("\n[merge6] Step 1/4: Loading existing 5-model merge...")
    sys.stdout.flush()
    target_model, target_tokenizer = load_model(existing_path)

    # --- Step 3: Download & load Klear-Reasoner-8B ---
    print("\n[merge6] Step 2/4: Loading Klear-Reasoner-8B (will download if needed)...")
    sys.stdout.flush()
    t = time.time()

    from transformers import AutoModelForCausalLM as AMLM
    source_model = AMLM.from_pretrained(
        KLEAR_HF_ID,
        device_map="cpu",  # Keep on CPU to save GPU VRAM
        torch_dtype=torch.bfloat16,
    )
    source_state = source_model.state_dict()
    print(f"[merge6] Klear loaded in {time.time()-t:.0f}s ({len(source_state)} params)")
    sys.stdout.flush()

    # Free source model object (keep state dict)
    del source_model
    gc.collect()

    # --- Step 4: Direct merge ---
    print("\n[merge6] Step 3/4: Merging weights (direct same-arch blend)...")
    sys.stdout.flush()
    t = time.time()
    target_model = direct_merge(target_model, source_state, alpha=KLEAR_ALPHA)
    del source_state
    gc.collect()
    torch.cuda.empty_cache()
    print(f"[merge6] Merge done in {time.time()-t:.0f}s")

    # --- Step 5: Save ---
    print(f"\n[merge6] Step 4/4: Saving to {OUTPUT_DIR}...")
    sys.stdout.flush()
    t = time.time()
    output_path = Path(OUTPUT_DIR)
    output_path.mkdir(parents=True, exist_ok=True)

    target_model.save_pretrained(str(output_path), safe_serialization=True)
    target_tokenizer.save_pretrained(str(output_path))

    # Verify save
    saved_files = list(output_path.glob("*.safetensors"))
    total_size = sum(f.stat().st_size for f in saved_files) if saved_files else 0
    if not saved_files or total_size < 1_000_000_000:
        print(f"[merge6] WARNING: Save may have failed! Only {total_size/1e9:.1f} GB")
    else:
        print(f"[merge6] Saved: {len(saved_files)} files, {total_size/1e9:.1f} GB")
    print(f"[merge6] Save done in {time.time()-t:.0f}s")

    total_min = (time.time() - merge_start) / 60
    print(f"\n{'=' * 60}")
    print(f"MERGE #6 COMPLETE: {total_min:.1f} min total")
    print(f"Output: {OUTPUT_DIR}")
    print(f"{'=' * 60}")
    sys.stdout.flush()

    return str(output_path)


if __name__ == "__main__":
    run_merge()