"""Week-1 Track-2 20% calibrated mixed-bit compression entry points.

The artifact stores exact GPTQ codes and scales for the W3/W4 body, W8 full
attention, 30,000 selected tied embedding/output rows, and a small deterministic
token-string predictor for omitted rows.  Zlib wraps that tensor payload
losslessly for checkpoint storage.  Restoration produces an ordinary BF16
Hugging Face checkpoint.
"""

from __future__ import annotations

import json
import importlib.util
import os
import sys
import sysconfig
from pathlib import Path

# This file name is required by the course interface, but ``code`` is also a
# Python standard-library module used by ``pdb`` during PyTorch import.  When a
# user runs a wrapper from this directory, Python can resolve this file for both
# names.  Publish the stdlib API before importing torch so that the recursive
# ``pdb -> code`` import remains valid.
if __name__ == "code":
    _stdlib_code_path = Path(sysconfig.get_path("stdlib")) / "code.py"
    _stdlib_code_spec = importlib.util.spec_from_file_location(
        "_cs6013_stdlib_code", _stdlib_code_path
    )
    if _stdlib_code_spec is None or _stdlib_code_spec.loader is None:
        raise ImportError(f"could not load Python stdlib code module: {_stdlib_code_path}")
    _stdlib_code = importlib.util.module_from_spec(_stdlib_code_spec)
    _stdlib_code_spec.loader.exec_module(_stdlib_code)
    for _stdlib_name in (
        "InteractiveInterpreter",
        "InteractiveConsole",
        "interact",
        "compile_command",
    ):
        globals()[_stdlib_name] = getattr(_stdlib_code, _stdlib_name)

import torch

LOCAL_SRC = Path(__file__).resolve().parent / "src"
if LOCAL_SRC.is_dir() and str(LOCAL_SRC) not in sys.path:
    sys.path.insert(0, str(LOCAL_SRC))

from eaimath.artifact import (
    dequantize_gptq_model,
    pack_model_state,
    read_trace_corpus,
    restore_artifact,
    save_artifact,
)
from eaimath.buckets import BY_NAME
from eaimath.embedding_predictor import fit_token_predictor
from eaimath.model import load_tokenizer
from eaimath.pack import state_dict_bytes
from eaimath.vocab import build_keep_set, token_frequencies

SUBMISSION_HF_REPO = "safffrron/25M2111-Week01-Track2-20-Submission01"
EXPECTED_BITS = {"mlp": 3, "linear_attn": 4, "full_attn": 8, "embed": 8}


def _validated_inputs() -> tuple[Path, Path, dict]:
    source = Path(
        os.environ.get(
            "EAIMATH_B20_GPTQ_SOURCE", "checkpoints/r13_short_gptq_m3l4a8e8"
        )
    )
    corpus = Path(os.environ.get("EAIMATH_VOCAB_CORPUS", "data/traces.jsonl"))
    if not source.is_dir() or not corpus.is_file():
        raise FileNotFoundError(
            "Set EAIMATH_B20_GPTQ_SOURCE to the reproduced m3l4a8e8 GPTQ "
            "checkpoint and EAIMATH_VOCAB_CORPUS to the verified trace JSONL."
        )
    config_path = source / "experiment_config.json"
    if not config_path.is_file():
        raise FileNotFoundError(f"GPTQ experiment config is missing: {config_path}")
    config = json.loads(config_path.read_text())
    if config.get("bits") != EXPECTED_BITS:
        raise ValueError(f"expected m3l4a8e8, found bits={config.get('bits')}")
    if set(config.get("gptq_components", [])) != {"mlp", "linear_attn"}:
        raise ValueError("GPTQ source must calibrate both MLP and linear attention")
    if config.get("activation_order") != "gar":
        raise ValueError("GPTQ source must use the selected GAR activation order")
    return source, corpus, config


def convert_from_hf_checkpoint(
    model_name: str,
    output_path: str,
    sparsity: float = 0.5,
) -> None:
    """Pack the reproduced calibrated source into the physical 20% artifact.

    Run ``training/reproduce_source.sh`` and
    ``compression/reproduce_gptq.sh`` first. ``sparsity`` is accepted for the
    supplied course interface but is not used by this dense quantization method.
    """
    _ = sparsity
    source, corpus, source_config = _validated_inputs()
    problems, completions, corpus_stats = read_trace_corpus([corpus])
    tokenizer = load_tokenizer(str(source))
    counts = token_frequencies(problems + completions, tokenizer)
    priority = token_frequencies(problems, tokenizer)
    keep = build_keep_set(counts, tokenizer, 30_000, priority_counts=priority)

    loaded, model, replaced, exact_weights, capture_error = dequantize_gptq_model(
        source
    )
    if capture_error != 0.0:
        raise RuntimeError(f"GPTQ code capture error: {capture_error}")
    payload, _ = pack_model_state(
        model.state_dict(),
        keep["keep_ids"],
        group_size=128,
        exact_entries=exact_weights,
        recipe_bits=EXPECTED_BITS,
        recipe_name="m3l4a8e8",
    )
    payload["source"] = {
        "base_model": model_name,
        "gptq_source": str(source),
        "experiment_config": source_config,
        "vocab_corpus": str(corpus),
        "vocab_strategy": "problem-first",
        "corpus_stats": corpus_stats,
        "keep_set": {key: value for key, value in keep.items() if key != "keep_ids"},
        "gptq_linears": replaced,
    }
    payload["report"]["source_gptq_linears"] = replaced
    payload["report"]["gptq_code_capture_error"] = capture_error
    payload["report"]["training_token_coverage"] = keep.get("token_coverage")

    row_names = [
        name
        for name, entry in payload["state_dict"].items()
        if entry["kind"] in {"quant_rows", "raw_rows"}
    ]
    if len(row_names) != 1:
        raise RuntimeError(f"expected one canonical tied embedding, found {row_names}")
    predictor, predictor_report = fit_token_predictor(
        model.state_dict()[row_names[0]].detach().cpu(),
        tokenizer,
        sample_size=65_536,
        device=os.environ.get("EAIMATH_PREDICTOR_DEVICE", "cpu"),
    )
    payload["embedding_predictor"] = predictor
    payload["report"]["embedding_predictor"] = predictor_report
    payload["report"]["tensor_bytes"] += state_dict_bytes(
        {"embedding_predictor": predictor["basis"]}
    )

    keep_path = Path(output_path).with_suffix(Path(output_path).suffix + ".keep_ids.json")
    keep_path.parent.mkdir(parents=True, exist_ok=True)
    keep_path.write_text(json.dumps(keep["keep_ids"]))
    final = save_artifact(payload, Path(output_path), lossless_codec="zlib", zlib_level=6)
    ceiling = int(BY_NAME["B"].high_gb * 1e9)
    if int(final["disk_bytes"]) >= ceiling:
        raise RuntimeError(
            f"artifact is outside the 20% ceiling: {final['disk_bytes']:,} >= {ceiling:,}"
        )
    del model, loaded
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def convert_to_hf_checkpoint(
    model_name: str,
    checkpoint_path: str,
    output_path: str,
) -> None:
    """Restore the self-contained artifact to a full BF16 HF checkpoint."""
    report = restore_artifact(
        model_name,
        checkpoint_path,
        output_path,
        embedding_fill="token",
    )
    Path(output_path, "submission_report.json").write_text(json.dumps(report, indent=2))