#!/usr/bin/env python
from __future__ import annotations

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

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from dovla_cil.data.datasets import CILDataset  # noqa: E402
from dovla_cil.eval.lattice_eval import _validation_group_ids  # noqa: E402
from dovla_cil.eval.maniskill_policy_rollout import _numeric_action_values  # noqa: E402
from dovla_cil.models.dovla import (  # noqa: E402
    DoVLAConfig,
    DoVLAModel,
    load_model_state,
    vectorize_toy_observation,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Export policy BC targets chosen by a trained field on CIL action lattices."
    )
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--dataset", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--split", choices=("train", "val", "all"), default="train")
    parser.add_argument("--batch-groups", type=int, default=32)
    parser.add_argument(
        "--exclude-types",
        default="expert",
        help="Comma-separated candidate_type values to exclude before field selection.",
    )
    parser.add_argument("--max-groups", type=int, default=None)
    args = parser.parse_args(argv)

    if args.batch_groups <= 0:
        parser.error("--batch-groups must be positive")
    if args.max_groups is not None and args.max_groups <= 0:
        parser.error("--max-groups must be positive when provided")

    try:
        import torch
    except ImportError as exc:  # pragma: no cover
        raise ImportError("export_field_selected_policy_targets.py requires torch") from exc

    checkpoint = torch.load(
        args.checkpoint,
        map_location=_resolve_device(args.device),
        weights_only=False,
    )
    model_config = DoVLAConfig(**checkpoint["model_config"])
    if model_config.observation_mode != "state":
        raise ValueError("field-selected target export currently supports state observations only")
    device = _resolve_device(args.device)
    model = DoVLAModel(model_config).to(device)
    load_model_state(model, checkpoint)
    model.eval()

    dataset = CILDataset(args.dataset)
    trainer_config = checkpoint.get("trainer_config", {})
    val_ids = set(
        _validation_group_ids(
            dataset.group_ids,
            val_fraction=float(trainer_config.get("val_fraction", 0.2)),
            seed=int(trainer_config.get("seed", 0)),
        )
    )
    if args.split == "train":
        group_ids = [group_id for group_id in dataset.group_ids if group_id not in val_ids]
    elif args.split == "val":
        group_ids = [group_id for group_id in dataset.group_ids if group_id in val_ids]
    else:
        group_ids = list(dataset.group_ids)
    if args.max_groups is not None:
        group_ids = group_ids[: args.max_groups]

    excluded = {item.strip() for item in args.exclude_types.split(",") if item.strip()}
    targets: dict[str, dict[str, Any]] = {}
    counts: dict[str, int] = {}
    with torch.no_grad():
        for start in range(0, len(group_ids), args.batch_groups):
            for group_id in group_ids[start : start + args.batch_groups]:
                records = [
                    record
                    for record in dataset.get_group(group_id)
                    if record.candidate_type not in excluded
                ]
                if not records:
                    records = dataset.get_group(group_id)
                if not records:
                    continue
                obs = torch.tensor(
                    [
                        vectorize_toy_observation(
                            records[0].observation_inline or {},
                            obs_dim=model_config.obs_dim,
                        )
                    ]
                    * len(records),
                    dtype=torch.float32,
                    device=device,
                )
                actions = torch.tensor(
                    [_numeric_action_values(record) for record in records],
                    dtype=torch.float32,
                    device=device,
                )
                field = model.forward_field(
                    obs,
                    [record.instruction for record in records],
                    actions,
                )
                best_idx = int(torch.argmax(field["potential"].reshape(len(records))).item())
                best = records[best_idx]
                counts[best.candidate_type] = counts.get(best.candidate_type, 0) + 1
                targets[group_id] = {
                    "record_id": best.record_id,
                    "candidate_type": best.candidate_type,
                    "task_id": best.task_id,
                    "score": float(best.reward.score),
                    "rank_within_group": best.rank_within_group,
                }

    payload = {
        "checkpoint": str(args.checkpoint),
        "dataset": str(args.dataset),
        "split": args.split,
        "excluded_candidate_types": sorted(excluded),
        "num_groups": len(group_ids),
        "num_targets": len(targets),
        "selected_candidate_type_counts": counts,
        "targets": targets,
    }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(payload, indent=2) + "\n")
    print(json.dumps({k: v for k, v in payload.items() if k != "targets"}, indent=2))
    print(f"Wrote {args.out}")
    return 0


def _resolve_device(device: str) -> str:
    if device != "auto":
        return device
    try:
        import torch
    except ImportError:  # pragma: no cover
        return "cpu"
    return "cuda" if torch.cuda.is_available() else "cpu"


if __name__ == "__main__":
    raise SystemExit(main())