"""Evaluate one or more MACF actor checkpoints on a fixed scenario suite."""

import copy
import csv
import json
import re
import sys
from datetime import datetime
from pathlib import Path

import numpy as np
import torch
import yaml

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

from config import get_args, get_config
from runner.shared.macf_runner import MACFEnvRunner
from runner.shared.gymrt2a_runner import GymRT2ARunner
from runner.qmix.macf_runner import MACFRunner as QMixRunner
from train.train import add_my_args, make_eval_env


SUITE_CONFIG_KEYS = {
    "map_size",
    "num_uav",
    "num_ugv",
    "uav_perception_rg",
    "ugv_perception_rg",
    "target_num",
    "max_steps",
    "episode_time_limit",
    "cap_dim",
    "eval_seed",
    "eval_reset_mode",
    "eval_episodes",
    "n_eval_rollout_threads",
    "rewards_config_file",
    "gymrt2a_map_size",
    "gymrt2a_num_agents",
    "gymrt2a_objects",
    "gymrt2a_view_range",
    "gymrt2a_comm_range",
    "gymrt2a_share_intention",
}


def parse_args(args):
    parser = get_args()
    parser.add_argument("--scenario_name", type=str, default="MyEnv")
    parser.add_argument("--num_landmarks", type=int, default=3)
    add_my_args(parser)
    parser.add_argument(
        "--models_file",
        type=str,
        required=True,
        help="YAML file containing the checkpoints and their config overrides",
    )
    parser.add_argument(
        "--output_root",
        type=str,
        default=None,
        help="override the output root declared in the model manifest",
    )
    parser.add_argument(
        "--model_name_pattern",
        default=None,
        help="optional regular expression selecting a subset of manifest models",
    )
    parser.add_argument(
        "--progress_file",
        default=None,
        help="write atomic per-episode progress for the currently evaluated model",
    )
    return parser.parse_args(args)


def _load_manifest(path):
    manifest_path = Path(path).resolve()
    with manifest_path.open("r", encoding="utf-8") as file:
        manifest = yaml.safe_load(file) or {}
    models = manifest.get("models")
    if not isinstance(models, list) or not models:
        raise ValueError("models_file must contain a non-empty 'models' list.")
    return manifest_path, manifest


def _resolve_checkpoint(value, manifest_dir):
    path = Path(value).expanduser()
    candidates = [path, manifest_dir / path] if not path.is_absolute() else [path]
    for candidate in candidates:
        if candidate.is_dir():
            candidate = candidate / "actor.pt"
        if candidate.is_file():
            return candidate.resolve()
    raise FileNotFoundError(f"Actor checkpoint not found: {value}")


def _safe_name(name):
    safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", str(name)).strip("._")
    if not safe:
        raise ValueError(f"Invalid model name: {name!r}")
    return safe


def _select_device(config):
    if config.cuda >= 0 and torch.cuda.is_available():
        device = torch.device(f"cuda:{config.cuda}")
        if config.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        device = torch.device("cpu")
    torch.set_num_threads(config.n_training_threads)
    return device


def _load_actor(policy, checkpoint_path, device):
    try:
        checkpoint = torch.load(
            checkpoint_path, map_location=device, weights_only=True
        )
    except TypeError:
        checkpoint = torch.load(checkpoint_path, map_location=device)
    if hasattr(policy, "load_actor_checkpoint"):
        policy.load_actor_checkpoint(checkpoint)
    else:
        policy.actor.load_state_dict(checkpoint)


def _close_runner(runner, envs):
    envs.close()
    runner.writter.close()
    if hasattr(runner, "writer") and runner.writer is not runner.writter:
        runner.writer.flush()
        runner.writer.close()


def _evaluate_model(base_config, model, checkpoint, output_dir, device):
    config = copy.deepcopy(base_config)
    overrides = model.get("overrides", {}) or {}
    if not isinstance(overrides, dict):
        raise ValueError(f"overrides for {model['name']!r} must be a mapping.")
    forbidden = SUITE_CONFIG_KEYS.intersection(overrides)
    if forbidden:
        raise ValueError(
            f"Model {model['name']!r} overrides suite settings: {sorted(forbidden)}"
        )
    for key, value in overrides.items():
        setattr(config, key, value)

    trainer_name = str(
        getattr(config, "trainer_algorithm_name", None) or config.algorithm_name
    ).lower()
    if trainer_name in {"qmix", "vdn", "hatt_qmix"}:
        if str(config.env_name) in ("gymrt2a", "mrt2a"):
            raise NotImplementedError("QMIX-family evaluation currently targets capcoal/delivcoal.")
        config.num_agents = config.num_uav + config.num_ugv
        runner_class = QMixRunner
    elif str(config.env_name) in ("gymrt2a", "mrt2a"):
        config.num_agents = int(config.gymrt2a_num_agents)
        runner_class = GymRT2ARunner
    else:
        config.num_agents = config.num_uav + config.num_ugv
        runner_class = MACFEnvRunner
    config.use_eval = True
    config.use_render = False
    config.model_dir = None
    config.device = device
    config.experiment_name = _safe_name(model["name"])
    config.eval_output_dir = str(output_dir)
    config.eval_progress_file = getattr(base_config, "progress_file", None)
    config.run_dir = output_dir / "runner"

    torch.manual_seed(config.seed)
    torch.cuda.manual_seed_all(config.seed)
    np.random.seed(config.seed)

    envs = make_eval_env(config)
    config.envs = envs
    config.eval_envs = envs
    runner = None
    try:
        runner = runner_class(config)
        _load_actor(runner.policy, checkpoint, device)
        return runner.eval(total_num_steps=int(model.get("step", 0)))
    finally:
        if runner is not None:
            _close_runner(runner, envs)
        else:
            envs.close()


def _write_suite_results(output_dir, records, metadata):
    fields = []
    for record in records:
        for key in record:
            if key not in fields:
                fields.append(key)
    with (output_dir / "summary.csv").open(
        "w", newline="", encoding="utf-8"
    ) as file:
        writer = csv.DictWriter(file, fieldnames=fields)
        writer.writeheader()
        writer.writerows(records)
    with (output_dir / "summary.json").open("w", encoding="utf-8") as file:
        json.dump(
            {"metadata": metadata, "models": records},
            file,
            ensure_ascii=False,
            indent=2,
        )


def main(args=None):
    config = parse_args(sys.argv[1:] if args is None else args)
    get_config(config)
    manifest_path, manifest = _load_manifest(config.models_file)

    suite_name = _safe_name(manifest.get("name", "macf-eval"))
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    output_root = Path(
        config.output_root or manifest.get("output_root", "eval_results")
    ).expanduser()
    output_dir = (output_root / f"{suite_name}_{timestamp}").resolve()
    output_dir.mkdir(parents=True, exist_ok=False)

    device = _select_device(config)
    models = manifest["models"]
    if config.model_name_pattern:
        pattern = re.compile(config.model_name_pattern)
        models = [model for model in models if pattern.search(str(model.get("name", "")))]
        if not models:
            raise ValueError(
                f"model_name_pattern matched no models: {config.model_name_pattern!r}"
            )

    records = []
    seen_names = set()
    for model in models:
        if not isinstance(model, dict) or "name" not in model or "checkpoint" not in model:
            raise ValueError("Each model requires 'name' and 'checkpoint'.")
        safe_name = _safe_name(model["name"])
        if safe_name in seen_names:
            raise ValueError(f"Duplicate model name: {model['name']!r}")
        seen_names.add(safe_name)

        checkpoint = _resolve_checkpoint(
            model["checkpoint"], manifest_path.parent
        )
        model_output = output_dir / "models" / safe_name
        model_output.mkdir(parents=True)
        print(f"Evaluating {model['name']}: {checkpoint}")
        summary = _evaluate_model(
            config, model, checkpoint, model_output, device
        )
        records.append(
            {
                "name": str(model["name"]),
                "checkpoint": str(checkpoint),
                **summary,
            }
        )

    metadata = {
        "name": suite_name,
        "created_at": datetime.now().isoformat(timespec="seconds"),
        "config_file": str(Path(config.config_file).resolve())
        if config.config_file
        else None,
        "models_file": str(manifest_path),
        "eval_seed": int(config.eval_seed),
        "eval_reset_mode": getattr(config, "eval_reset_mode", None),
        "eval_episodes": int(config.eval_episodes),
        "eval_deterministic": bool(
            getattr(config, "eval_deterministic", True)
        ),
        "delivery_tasks_initially_discovered": bool(
            getattr(config, "delivery_tasks_initially_discovered", False)
        ),
        "mask_explore_actions": bool(
            getattr(config, "mask_explore_actions", False)
        ),
        "max_steps": int(config.max_steps),
        "episode_time_limit": int(
            getattr(config, "episode_time_limit", None) or config.max_steps
        ),
        "device": str(device),
        "environment": {
            key: getattr(config, key) for key in sorted(SUITE_CONFIG_KEYS)
        },
    }
    _write_suite_results(output_dir, records, metadata)
    print(f"Evaluation suite written to {output_dir}")
    return output_dir


if __name__ == "__main__":
    main()