#!/usr/bin/env python3
"""Evaluate held-out loss with the slow grouped 64-wide MoE runtime patch."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import numpy as np

from heapr.eval.loss import evaluate_token_cache
from heapr.grouped_model import (
    ExpandedGroupedLagunaContext,
    GroupedLagunaContext,
    RepackedExpandedGroupedLagunaContext,
)
from heapr.model_utils import build_max_memory, load_causal_lm, validate_model_device_placement
from heapr.utils import require_torch
from heapr.utils import write_json


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-id", required=True)
    parser.add_argument("--cache-path", required=True)
    parser.add_argument("--group-mask", required=True)
    parser.add_argument("--group-indices", required=True)
    parser.add_argument("--output-path", required=True)
    parser.add_argument("--revision")
    parser.add_argument("--dtype", default="bfloat16")
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--max-chunks", type=int)
    parser.add_argument("--routing-mode", choices=["parent", "expanded", "repacked"], default="parent")
    parser.add_argument(
        "--child-budget-mode",
        choices=["backfill", "layer-scaled", "parent-weighted"],
        default="backfill",
    )
    parser.add_argument("--gpu-memory-per-device")
    parser.add_argument("--max-gpu-memory")
    parser.add_argument("--max-cpu-memory")
    parser.add_argument("--offload-folder")
    parser.add_argument("--allow-cpu-offload", action="store_true")
    parser.add_argument("--cache-implementation", default="static")
    parser.add_argument("--no-cache", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.offload_folder and not args.allow_cpu_offload:
        raise ValueError("--offload-folder requires --allow-cpu-offload")
    max_memory = build_max_memory(
        gpu_memory_per_device=args.gpu_memory_per_device,
        max_gpu_memory=args.max_gpu_memory,
        max_cpu_memory=args.max_cpu_memory,
        allow_cpu_offload=args.allow_cpu_offload,
    )
    torch = require_torch()
    requested_gpu_count = torch.cuda.device_count() if args.gpu_memory_per_device else None
    model = load_causal_lm(
        args.model_id,
        revision=args.revision,
        dtype=args.dtype,
        max_memory=max_memory,
        offload_folder=args.offload_folder if args.allow_cpu_offload else None,
        use_cache=not args.no_cache,
        cache_implementation=args.cache_implementation,
    )
    validate_model_device_placement(
        model,
        allow_cpu_offload=args.allow_cpu_offload,
        requested_gpu_count=requested_gpu_count,
    )
    keep_mask = np.load(args.group_mask)
    group_indices = np.load(args.group_indices)
    context_cls = {
        "parent": GroupedLagunaContext,
        "expanded": ExpandedGroupedLagunaContext,
        "repacked": RepackedExpandedGroupedLagunaContext,
    }[args.routing_mode]
    context_kwargs = {}
    if context_cls is RepackedExpandedGroupedLagunaContext:
        context_kwargs["child_budget_mode"] = args.child_budget_mode
    with context_cls(model, keep_group_masks=keep_mask, group_indices=group_indices, **context_kwargs):
        metrics = evaluate_token_cache(
            model,
            args.cache_path,
            batch_size=args.batch_size,
            max_chunks=args.max_chunks,
            use_cache=not args.no_cache,
            cache_implementation=args.cache_implementation,
        )
    metrics["grouped_routing_mode"] = args.routing_mode
    metrics["child_budget_mode"] = args.child_budget_mode
    write_json(args.output_path, metrics)
    print(metrics)


if __name__ == "__main__":
    main()