#!/usr/bin/env python3
"""Establish the lineage of Piko-9b by direct tensor comparison.

The script answers three questions with evidence rather than assertion:

1. Where does every tensor in the published checkpoint come from?
2. Are the Piko QLoRA adapters actually baked into the published weights?
3. Is the vision tower trained, or copied verbatim from its source?

It works on safetensors headers plus a sampled set of full tensors, so it does
not need to read all ~18 GB of every checkpoint.

Example
-------
python scripts/analyze_lineage.py \
    --candidate  .../merged_model/piko9-phase15-native-vision-base \
    --language   .../merged_model/wraithfast-phase15-150k-full-ft \
    --vision     .../models/qwen3.5-9b-vision \
    --adapter    .../adapters/piko9-phase5-native-vision-qlora \
    --adapter    .../adapters/piko9-120k-qlora \
    --output     reports/lineage_analysis.json
"""

from __future__ import annotations

import argparse
import hashlib
import json
import random
import sys
from pathlib import Path
from typing import Any

try:
    import torch
    from safetensors import safe_open
except ImportError as exc:  # pragma: no cover - environment problem, not logic
    sys.exit(f"Missing dependency: {exc}. Install torch and safetensors.")


# --------------------------------------------------------------------------- #
# checkpoint access
# --------------------------------------------------------------------------- #


class Checkpoint:
    """Read-only view over a sharded (or single-file) safetensors checkpoint."""

    def __init__(self, root: Path) -> None:
        self.root = root
        self.weight_map = self._load_weight_map()
        self._handles: dict[str, Any] = {}

    def _load_weight_map(self) -> dict[str, str]:
        index = self.root / "model.safetensors.index.json"
        if index.is_file():
            return json.loads(index.read_text(encoding="utf-8"))["weight_map"]

        singles = sorted(self.root.glob("*.safetensors"))
        if not singles:
            raise FileNotFoundError(f"No safetensors found under {self.root}")
        mapping: dict[str, str] = {}
        for shard in singles:
            with safe_open(shard, framework="pt", device="cpu") as handle:
                for key in handle.keys():
                    mapping[key] = shard.name
        return mapping

    @property
    def keys(self) -> set[str]:
        return set(self.weight_map)

    def _handle(self, filename: str) -> Any:
        if filename not in self._handles:
            self._handles[filename] = safe_open(self.root / filename, framework="pt", device="cpu")
        return self._handles[filename]

    def meta(self, key: str) -> dict[str, Any]:
        view = self._handle(self.weight_map[key]).get_slice(key)
        return {"dtype": str(view.get_dtype()), "shape": [int(d) for d in view.get_shape()]}

    def tensor(self, key: str) -> torch.Tensor:
        return self._handle(self.weight_map[key]).get_tensor(key)


def tensor_digest(tensor: torch.Tensor) -> str:
    """Content hash that is stable across dtypes and layouts."""
    contiguous = tensor.contiguous()
    return hashlib.sha256(
        contiguous.view(torch.uint8).numpy().tobytes()
        if contiguous.dtype not in (torch.bfloat16,)
        else contiguous.view(torch.int16).numpy().tobytes()
    ).hexdigest()


def tensor_stats(tensor: torch.Tensor) -> dict[str, float]:
    values = tensor.to(torch.float32)
    return {
        "mean": float(values.mean()),
        "std": float(values.std()),
        "absmax": float(values.abs().max()),
        "numel": int(values.numel()),
    }


# --------------------------------------------------------------------------- #
# comparisons
# --------------------------------------------------------------------------- #


def compare_key_sets(candidate: Checkpoint, sources: dict[str, Checkpoint]) -> dict[str, Any]:
    result: dict[str, Any] = {
        "candidate_tensor_count": len(candidate.keys),
        "sources": {},
    }
    covered: set[str] = set()
    for name, source in sources.items():
        shared = candidate.keys & source.keys
        covered |= shared
        result["sources"][name] = {
            "tensor_count": len(source.keys),
            "shared_with_candidate": len(shared),
            "only_in_source": sorted(source.keys - candidate.keys)[:20],
        }
    result["candidate_keys_unexplained"] = sorted(candidate.keys - covered)
    return result


def attribute_tensors(
    candidate: Checkpoint,
    sources: dict[str, Checkpoint],
    sample_size: int,
    seed: int,
) -> dict[str, Any]:
    """For a random sample of tensors, decide which source they came from."""
    rng = random.Random(seed)
    keys = sorted(candidate.keys)

    # Stratify so both the language and the vision half are represented.
    language = [k for k in keys if ".visual." not in k]
    vision = [k for k in keys if ".visual." in k]
    per_half = max(1, sample_size // 2)
    sampled = rng.sample(language, min(per_half, len(language)))
    sampled += rng.sample(vision, min(per_half, len(vision)))

    records: list[dict[str, Any]] = []
    tally: dict[str, int] = {}
    for key in sorted(sampled):
        candidate_tensor = candidate.tensor(key)
        candidate_hash = tensor_digest(candidate_tensor)
        record: dict[str, Any] = {
            "key": key,
            "meta": candidate.meta(key),
            "matches": [],
            "differs_from": [],
            "absent_from": [],
        }
        for name, source in sources.items():
            if key not in source.keys:
                record["absent_from"].append(name)
                continue
            source_tensor = source.tensor(key)
            if source_tensor.shape != candidate_tensor.shape:
                record["differs_from"].append({"source": name, "reason": "shape"})
                continue
            if tensor_digest(source_tensor) == candidate_hash:
                record["matches"].append(name)
            else:
                delta = (
                    (source_tensor.to(torch.float32) - candidate_tensor.to(torch.float32))
                    .abs()
                    .max()
                )
                record["differs_from"].append(
                    {"source": name, "reason": "value", "max_abs_delta": float(delta)}
                )

        verdict = "+".join(record["matches"]) if record["matches"] else "no_exact_match"
        record["verdict"] = verdict
        tally[verdict] = tally.get(verdict, 0) + 1
        records.append(record)

    return {"sampled": len(records), "tally": tally, "records": records}


def adapter_effect(
    candidate: Checkpoint,
    language: Checkpoint,
    adapter_dir: Path,
    max_modules: int,
) -> dict[str, Any]:
    """Check whether a LoRA adapter's delta is present in the candidate weights.

    For each targeted module we reconstruct ``scaling * B @ A`` and test the
    hypothesis ``candidate == language + delta``.  If the candidate instead
    equals ``language`` exactly, the adapter was never merged.
    """
    config_path = adapter_dir / "adapter_config.json"
    weights_path = adapter_dir / "adapter_model.safetensors"
    if not (config_path.is_file() and weights_path.is_file()):
        return {"adapter": adapter_dir.name, "status": "not_found"}

    config = json.loads(config_path.read_text(encoding="utf-8"))
    rank = config.get("r")
    alpha = config.get("lora_alpha")
    scaling = (alpha / rank) if rank else None

    with safe_open(weights_path, framework="pt", device="cpu") as handle:
        adapter_keys = list(handle.keys())
        pairs: dict[str, dict[str, str]] = {}
        for key in adapter_keys:
            if ".lora_A." in key:
                pairs.setdefault(key.split(".lora_A.")[0], {})["A"] = key
            elif ".lora_B." in key:
                pairs.setdefault(key.split(".lora_B.")[0], {})["B"] = key

        checks: list[dict[str, Any]] = []
        merged_votes = {"merged": 0, "not_merged": 0, "inconclusive": 0}

        for module, parts in sorted(pairs.items()):
            if len(checks) >= max_modules:
                break
            if "A" not in parts or "B" not in parts:
                continue

            target = _resolve_target_key(module, candidate.keys)
            if target is None:
                continue

            lora_a = handle.get_tensor(parts["A"]).to(torch.float32)
            lora_b = handle.get_tensor(parts["B"]).to(torch.float32)
            delta = (lora_b @ lora_a) * (scaling or 1.0)
            delta_norm = float(delta.abs().max())

            candidate_tensor = candidate.tensor(target).to(torch.float32)
            if target not in language.keys:
                continue
            language_tensor = language.tensor(target).to(torch.float32)

            if delta.shape != language_tensor.shape:
                continue

            gap_unmerged = float((candidate_tensor - language_tensor).abs().max())
            gap_merged = float((candidate_tensor - (language_tensor + delta)).abs().max())

            if delta_norm == 0.0:
                verdict = "inconclusive"  # zero-valued adapter proves nothing
            elif gap_unmerged == 0.0 and gap_merged > 0.0:
                verdict = "not_merged"
            elif gap_merged < gap_unmerged:
                verdict = "merged"
            else:
                verdict = "inconclusive"

            merged_votes[verdict] += 1
            checks.append(
                {
                    "module": module,
                    "candidate_tensor": target,
                    "lora_delta_absmax": delta_norm,
                    "absmax_candidate_minus_language": gap_unmerged,
                    "absmax_candidate_minus_language_plus_delta": gap_merged,
                    "verdict": verdict,
                }
            )

    if merged_votes["not_merged"] and not merged_votes["merged"]:
        conclusion = "adapter_is_NOT_present_in_published_weights"
    elif merged_votes["merged"] and not merged_votes["not_merged"]:
        conclusion = "adapter_IS_present_in_published_weights"
    elif not checks:
        conclusion = "no_comparable_modules"
    else:
        conclusion = "inconclusive"

    return {
        "adapter": adapter_dir.name,
        "status": "checked",
        "rank": rank,
        "lora_alpha": alpha,
        "scaling": scaling,
        "base_model_name_or_path": config.get("base_model_name_or_path"),
        "lora_module_count": len(pairs),
        "modules_checked": len(checks),
        "votes": merged_votes,
        "conclusion": conclusion,
        "checks": checks,
    }


def _resolve_target_key(module: str, available: set[str]) -> str | None:
    """Map a PEFT module path onto a checkpoint tensor key."""
    stem = module
    for prefix in ("base_model.model.", "base_model."):
        stem = stem.removeprefix(prefix)
    candidates = [
        f"{stem}.weight",
        f"model.{stem}.weight",
        f"model.language_model.{stem}.weight",
        f"{stem.replace('model.layers.', 'model.language_model.layers.')}.weight",
    ]
    for candidate in candidates:
        if candidate in available:
            return candidate
    return None


# --------------------------------------------------------------------------- #
# entry point
# --------------------------------------------------------------------------- #


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--candidate", required=True, type=Path)
    parser.add_argument("--language", required=True, type=Path)
    parser.add_argument("--vision", required=True, type=Path)
    parser.add_argument("--adapter", action="append", default=[], type=Path)
    parser.add_argument("--sample-size", type=int, default=24)
    parser.add_argument("--adapter-modules", type=int, default=6)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--output", type=Path, default=Path("reports/lineage_analysis.json"))
    args = parser.parse_args()

    for label, path in (
        ("candidate", args.candidate),
        ("language", args.language),
        ("vision", args.vision),
    ):
        if not path.is_dir():
            sys.exit(f"{label} checkpoint not found: {path}")

    print(f"Opening candidate  : {args.candidate}", flush=True)
    candidate = Checkpoint(args.candidate)
    print(f"Opening language   : {args.language}", flush=True)
    language = Checkpoint(args.language)
    print(f"Opening vision     : {args.vision}", flush=True)
    vision = Checkpoint(args.vision)

    sources = {"language": language, "vision": vision}

    report: dict[str, Any] = {
        "inputs": {
            "candidate": str(args.candidate),
            "language": str(args.language),
            "vision": str(args.vision),
            "adapters": [str(a) for a in args.adapter],
        },
        "seed": args.seed,
    }

    print("Comparing tensor key sets...", flush=True)
    report["key_sets"] = compare_key_sets(candidate, sources)

    print(f"Attributing a sample of {args.sample_size} tensors...", flush=True)
    report["attribution"] = attribute_tensors(candidate, sources, args.sample_size, args.seed)

    report["adapters"] = []
    for adapter_dir in args.adapter:
        print(f"Testing adapter {adapter_dir.name}...", flush=True)
        report["adapters"].append(
            adapter_effect(candidate, language, adapter_dir, args.adapter_modules)
        )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(f"\nWrote {args.output}")

    print("\n--- summary ---")
    print("tensor attribution:", report["attribution"]["tally"])
    for entry in report["adapters"]:
        print(f"adapter {entry['adapter']}: {entry.get('conclusion', entry['status'])}")


if __name__ == "__main__":
    main()