"""Run frozen MatrixFormer-Impute or MatrixFormer+ on a numerical matrix."""

import argparse
import hashlib
import json
from pathlib import Path

import numpy as np
from huggingface_hub import hf_hub_download

from tabimpute.routing import TwoExpertRouter
from tabimpute.streamed_inference import StreamedImputePFN


REPO_ID = "Tabimpute/TabImpute"


def artifact_path(filename, local_dir, revision):
    if local_dir is not None:
        path = local_dir / filename
        if not path.is_file():
            raise FileNotFoundError(f"Missing local artifact: {path}")
        return path
    return Path(hf_hub_download(REPO_ID, filename, revision=revision))


def verify_sha256(path, expected):
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    if digest.hexdigest() != expected:
        raise ValueError(f"SHA-256 mismatch: {path}")


def load_expert(name, manifest, args):
    spec = manifest["experts"][name]
    config_path = artifact_path(spec["config"], args.local_dir, args.revision)
    checkpoint_path = artifact_path(spec["checkpoint"], args.local_dir, args.revision)
    verify_sha256(config_path, spec["config_sha256"])
    verify_sha256(checkpoint_path, spec["checkpoint_sha256"])
    config = json.loads(config_path.read_text(encoding="utf-8"))
    return StreamedImputePFN(
        checkpoint_path=str(checkpoint_path),
        json_config={"config": config["model"]},
        device=args.device,
        prediction_stat="mean",
        precision_mode=args.precision,
        max_num_rows=args.max_rows or None,
        max_num_cols=args.max_cols or None,
    )


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", choices=("impute", "plus"), default="impute")
    parser.add_argument("--input", type=Path, help="2D numeric .npy array; missing entries are NaN")
    parser.add_argument("--output", type=Path, help="Completed .npy array")
    parser.add_argument("--device", default="cpu", help="For example: cpu, cuda, cuda:0")
    parser.add_argument("--precision", choices=("amp_bfloat16", "float32"), default="amp_bfloat16")
    parser.add_argument("--max-rows", type=int, default=2000, help="Context rows; 0 uses the full axis")
    parser.add_argument("--max-cols", type=int, default=1000, help="Context columns; 0 uses the full axis")
    parser.add_argument("--local-dir", type=Path, help="Offline bundle root containing matrixformer/")
    parser.add_argument("--revision", default="main", help="Hugging Face revision; use a commit hash to pin it")
    args = parser.parse_args()
    if args.max_rows < 0 or args.max_cols < 0:
        parser.error("Context limits must be nonnegative")
    if args.output is not None and args.output.suffix != ".npy":
        parser.error("--output must have a .npy suffix")
    if args.input is not None and args.output is not None and args.input.resolve() == args.output.resolve():
        parser.error("--output must differ from --input")
    return args


def main():
    args = parse_args()
    if args.input is None:
        generator = np.random.default_rng(42)
        matrix = generator.normal(size=(12, 3)) @ generator.normal(size=(3, 9))
        missing = generator.choice(matrix.size, size=30, replace=False)
        matrix.flat[missing] = np.nan
    else:
        matrix = np.load(args.input, allow_pickle=False)
    if not np.issubdtype(matrix.dtype, np.number) or np.iscomplexobj(matrix):
        raise ValueError("Input must contain real numbers, with NaN marking missing entries")
    matrix = np.asarray(matrix, dtype=np.float64)
    if matrix.ndim != 2 or min(matrix.shape) == 0 or np.isinf(matrix).any():
        raise ValueError("Expected a nonempty 2D matrix with finite observations or NaN")
    if not np.isfinite(matrix).any():
        raise ValueError("At least one observed entry is required")
    manifest_path = artifact_path("matrixformer/artifacts.json", args.local_dir, args.revision)
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    general = load_expert("impute", manifest, args)
    if args.model == "plus":
        specialist = load_expert("recsys", manifest, args)
        imputer = TwoExpertRouter(
            general, specialist, policy="observed_blend", normalization="column"
        )
    else:
        imputer = general
    completed = imputer.impute(matrix.copy())
    if not np.isfinite(completed).all():
        raise RuntimeError("Model returned nonfinite predictions")
    observed = np.isfinite(matrix)
    if not np.array_equal(completed[observed], matrix[observed]):
        raise RuntimeError("Model changed observed entries")
    report = {
        "model": args.model,
        "device": args.device,
        "precision": args.precision,
        "shape": list(matrix.shape),
        "missing_entries": int(np.isnan(matrix).sum()),
        "all_predictions_finite": True,
        "observed_entries_preserved": True,
    }
    if args.model == "plus":
        report["blend_decision"] = imputer.last_result.decisions["observed"]
    if args.output is not None:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        np.save(args.output, completed, allow_pickle=False)
        report["output"] = str(args.output)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()