"""Find a (quant_method, format) pair that lets transformers parse this
modelopt MIXED_PRECISION checkpoint, and check that per-layer schemes resolve
to the correct precision for BOTH the FP8 and the NVFP4 tensors."""

import gc
import json
import os
import shutil
import traceback

SRC = os.environ.get("DWCT_NEMOTRON", "/models/nemotron")
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(HERE, "..", ".qprobe"))

print("=" * 72)
print("environment")
print("=" * 72)
try:
    import transformers
    print("transformers", transformers.__version__)
except Exception as exc:
    print("transformers ERR", exc)

try:
    import compressed_tensors as ct
    print("compressed_tensors", getattr(ct, "__version__", "?"))
except Exception as exc:
    print("compressed_tensors IMPORT FAILED:", type(exc).__name__, exc)

try:
    from compressed_tensors.config import QuantizationFormat
    print("QuantizationFormat values:", [f.value for f in QuantizationFormat])
except Exception as exc:
    print("QuantizationFormat ERR", exc)

for mod in ("quantizer_compressed_tensors", "quantizer_nvfp4", "quantizer_finegrained_fp8"):
    try:
        m = __import__("transformers.quantizers." + mod, fromlist=["*"])
        print(mod, "->", [n for n in dir(m) if not n.startswith("_") and ("Quant" in n or "Scheme" in n)])
    except Exception as exc:
        print(mod, "ERR", exc)

from transformers import PreTrainedModel
if not isinstance(getattr(PreTrainedModel, "all_tied_weights_keys", None), property):
    def _g(self):
        return getattr(self, "_all_tied_weights_keys", None) or getattr(self, "_tied_weights_keys", None) or []
    def _s(self, v):
        self._all_tied_weights_keys = v
    PreTrainedModel.all_tied_weights_keys = property(_g, _s)

PROBE_LAYERS = [
    "language_model.backbone.layers.1.mixer.experts.0.up_proj",       # NVFP4
    "language_model.backbone.layers.1.mixer.shared_experts.up_proj",  # FP8
    "language_model.backbone.layers.0.mixer.in_proj",                 # FP8
    "language_model.backbone.embeddings",                             # dense
]

CANDIDATES = [
    ("compressed-tensors", "nvfp4-pack-quantized"),
    ("compressed-tensors", "mixed-precision"),
    ("nvfp4", "nvfp4-pack-quantized"),
    ("nvfp4", "mixed-precision"),
    ("compressed-tensors", "float-quantized"),
]


def make_overlay(tag, method, fmt):
    dst = os.path.join(ROOT, tag)
    os.makedirs(dst, exist_ok=True)
    for name in os.listdir(SRC):
        s, d = os.path.join(SRC, name), os.path.join(dst, name)
        if os.path.lexists(d):
            continue
        try:
            os.symlink(s, d)
        except OSError:
            shutil.copytree(s, d) if os.path.isdir(s) else shutil.copy2(s, d)
    cfg = json.load(open(os.path.join(SRC, "config.json")))
    q = cfg["quantization_config"]
    q["quant_method"] = method
    q["format"] = fmt
    q["quantization_status"] = "compressed"
    cfg["quantization_config"] = q
    p = os.path.join(dst, "config.json")
    if os.path.islink(p):
        os.unlink(p)
    json.dump(cfg, open(p, "w"), indent=2)
    return dst


def scheme_of(qz, name):
    for obj in (qz, getattr(qz, "quantization_config", None)):
        fn = getattr(obj, "get_scheme", None)
        if fn is None:
            continue
        try:
            return fn(name)
        except Exception as exc:
            return "ERR " + repr(exc)[:140]
    return "no get_scheme"


def try_candidate(method, fmt):
    tag = ("%s__%s" % (method, fmt)).replace("-", "_")
    print("-" * 72)
    print("CANDIDATE  quant_method=%r  format=%r" % (method, fmt))
    try:
        d = make_overlay(tag, method, fmt)
        from transformers import AutoConfig
        from transformers.quantizers import auto as qauto
        cfg = AutoConfig.from_pretrained(d, trust_remote_code=True)
        qc = getattr(cfg, "quantization_config", None)
        print("  parsed config class:", type(qc).__name__)
        get_q = getattr(qauto, "get_quantizer", None)
        qz = get_q(qc) if get_q else qauto.AutoQuantizer.from_config(qc)
        print("  quantizer:", type(qz).__name__)
        for name in PROBE_LAYERS:
            sch = scheme_of(qz, name)
            label = type(sch).__name__ if not isinstance(sch, str) else sch
            print("  scheme(%-52s) -> %s" % (name.split("backbone.")[-1], label))
            if isinstance(sch, str) and sch.startswith("ERR"):
                print("  RESULT: scheme resolution failed")
                return False
        print("  building on meta ...")
        import torch
        from transformers import AutoModelForCausalLM
        with torch.device("meta"):
            model = AutoModelForCausalLM.from_pretrained(
                d, trust_remote_code=True, dtype=torch.bfloat16,
                attn_implementation="eager",
            )
        try:
            mod = model.language_model.backbone.layers[1].mixer.experts[0].up_proj
            print("  expert module:", type(mod).__name__)
        except Exception as exc:
            print("  expert module lookup ERR", exc)
        del model
        gc.collect()
        print("  RESULT: OK")
        return True
    except Exception as exc:
        print("  RESULT: FAIL", type(exc).__name__, repr(exc)[:300])
        traceback.print_exc(limit=2)
        return False


winners = []
for method, fmt in CANDIDATES:
    try:
        if try_candidate(method, fmt):
            winners.append((method, fmt))
    except Exception:
        traceback.print_exc(limit=3)

print("=" * 72)
print("WORKING COMBINATIONS:", winners)