"""Paired evaluation under a fixed per-episode permutation of task slots.

The physical environment remains in canonical task order.  Only the actor
observation/action interface is permuted, including coalition IDs and action
masks.  A task-slot permutation-equivariant policy should therefore preserve
its physical decisions and performance.  ``permutation_scope=deliver`` keeps
the historical homogeneous-Deliver test; ``all`` also moves the heterogeneous
Explore/Idle/Resupply slots.
"""

import argparse
import copy
import csv
import json
import subprocess
import sys
import traceback
from pathlib import Path

import numpy as np
import torch
import yaml

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

from config import get_args, get_config
from runner.shared.macf_runner import MACFEnvRunner, _t2n
from train.eval import _close_runner, _load_actor, _safe_name
from train.train import add_my_args, make_eval_env


def _build_config(config_file):
    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)
    config = parser.parse_args(["--config_file", str(config_file)])
    get_config(config)
    return config


def _resolve_checkpoint(value, manifest_dir):
    path = Path(value).expanduser()
    if not path.is_absolute():
        path = manifest_dir / path
    if path.is_dir():
        path = path / "actor.pt"
    if not path.is_file():
        raise FileNotFoundError(path)
    return path.resolve()


def _task_permutation(env, scenario_id, permutation_seed, permutation_scope):
    task_num = len(env.env.tasks)
    presented_to_canonical = np.arange(task_num, dtype=np.int64)
    if permutation_scope in ("deliver", "explore"):
        selected_type = 1 if permutation_scope == "deliver" else 0
        permuted_slots = np.asarray(
            [
                index
                for index, task in enumerate(env.env.tasks)
                if int(task.type) == selected_type
            ],
            dtype=np.int64,
        )
    elif permutation_scope == "all":
        permuted_slots = np.arange(task_num, dtype=np.int64)
    else:
        raise ValueError(f"Unsupported permutation_scope={permutation_scope!r}")
    shuffled = permuted_slots.copy()
    np.random.default_rng(permutation_seed + scenario_id).shuffle(shuffled)
    presented_to_canonical[permuted_slots] = shuffled
    canonical_to_presented = np.empty_like(presented_to_canonical)
    canonical_to_presented[presented_to_canonical] = np.arange(task_num)
    return presented_to_canonical, canonical_to_presented


def _map_normalized_coalitions(values, canonical_to_presented, valid=None):
    task_num = len(canonical_to_presented)
    canonical = np.rint(values * max(1, task_num - 1)).astype(np.int64)
    canonical = np.clip(canonical, 0, task_num - 1)
    mapped = canonical_to_presented[canonical] / float(max(1, task_num - 1))
    if valid is not None:
        mapped = np.where(valid, mapped, values)
    return mapped.astype(values.dtype, copy=False)


def _present_actor_inputs(
    obs,
    available_actions,
    presented_to_canonical,
    canonical_to_presented,
    cap_dim,
):
    presented_obs = {key: np.array(value, copy=True) for key, value in obs.items()}
    presented_obs["tasks"] = np.take(
        presented_obs["tasks"], presented_to_canonical, axis=-2
    )
    ego_coalition_index = 2 + cap_dim
    presented_obs["ego"][..., ego_coalition_index] = _map_normalized_coalitions(
        presented_obs["ego"][..., ego_coalition_index],
        canonical_to_presented,
    )
    other_coalition_index = cap_dim + 4
    valid_others = np.abs(presented_obs["others"]).sum(axis=-1) > 0
    presented_obs["others"][..., other_coalition_index] = (
        _map_normalized_coalitions(
            presented_obs["others"][..., other_coalition_index],
            canonical_to_presented,
            valid=valid_others,
        )
    )
    presented_available = np.take(
        np.asarray(available_actions), presented_to_canonical, axis=-1
    )
    return presented_obs, presented_available


@torch.no_grad()
def _run_episode(
    config,
    runner,
    env,
    scenario_id,
    permuted,
    permutation_seed,
    permutation_scope,
):
    scenario_seed = int(config.eval_seed) + int(scenario_id)
    numpy_state = np.random.get_state()
    np.random.seed(scenario_seed)
    try:
        env.seed(scenario_seed)
        obs, active_masks, available_actions = env.reset(scenario_id)
        task_num = len(env.env.tasks)
        if permuted:
            presented_to_canonical, canonical_to_presented = _task_permutation(
                env, scenario_id, permutation_seed, permutation_scope
            )
        else:
            presented_to_canonical = np.arange(task_num, dtype=np.int64)
            canonical_to_presented = presented_to_canonical.copy()

        rnn_states = np.zeros(
            (config.num_agents, config.recurrent_N, config.hidden_size),
            dtype=np.float32,
        )
        masks = np.ones((config.num_agents, 1), dtype=np.float32)
        last_actions = np.zeros((config.num_agents, 1), dtype=np.float32)
        episode_reward = 0.0
        final_info = None

        runner.trainer.prep_rollout()
        for step in range(config.max_steps):
            actor_obs, actor_available = _present_actor_inputs(
                obs,
                available_actions,
                presented_to_canonical,
                canonical_to_presented,
                config.cap_dim,
            )
            actions, next_rnn_states = runner.trainer.policy.act(
                actor_obs,
                rnn_states,
                masks,
                available_actions=actor_available,
                deterministic=True,
            )
            presented_actions = _t2n(actions).astype(np.int64)
            next_rnn_states = _t2n(next_rnn_states)
            canonical_actions = presented_to_canonical[
                presented_actions.reshape(-1)
            ].reshape(presented_actions.shape)
            canonical_actions = (
                active_masks * canonical_actions
                + (1.0 - active_masks) * last_actions
            )

            (
                obs,
                rewards,
                dones,
                info,
                active_masks,
                _,
                available_actions,
            ) = env.step(canonical_actions)
            final_info = info
            episode_reward += float(np.mean(rewards))

            done_mask = np.asarray(dones, dtype=bool)
            next_rnn_states[done_mask] = 0.0
            masks = np.ones((config.num_agents, 1), dtype=np.float32)
            masks[done_mask] = 0.0
            rnn_states = next_rnn_states
            last_actions = canonical_actions
            if bool(np.all(done_mask)) or info.get("end", 0):
                break
    finally:
        np.random.set_state(numpy_state)

    if final_info is None:
        raise RuntimeError("Episode ended without environment info.")
    delivery_tasks = [
        task for task in final_info["tasks"] if int(task["type"]) == 1
    ]
    completed = int(sum(bool(task["completed"]) for task in delivery_tasks))
    return {
        "scenario_id": int(scenario_id),
        "scenario_seed": scenario_seed,
        "mode": "permuted" if permuted else "original",
        "success": int(completed == len(delivery_tasks)),
        "tasks_completed": completed,
        "task_completion_ratio": completed / max(1, len(delivery_tasks)),
        "episode_length": int(step + 1),
        "episode_reward": float(episode_reward),
        "permutation": presented_to_canonical.tolist(),
        "permutation_scope": permutation_scope,
    }


def _evaluate_one(task):
    (
        model,
        base_config_file,
        manifest_dir,
        output_root,
        episodes,
        permutation_seed,
        permutation_scope,
        cuda,
    ) = task
    base_config = _build_config(base_config_file)
    config = copy.deepcopy(base_config)
    for key, value in (model.get("overrides") or {}).items():
        setattr(config, key, value)
    config.cuda = int(cuda)
    config.num_agents = config.num_uav + config.num_ugv
    config.use_eval = True
    config.use_render = False
    config.model_dir = None
    config.device = (
        torch.device(f"cuda:{cuda}")
        if int(cuda) >= 0
        else torch.device("cpu")
    )
    config.experiment_name = _safe_name(model["name"] + "_slot_permutation")
    model_output = Path(output_root) / _safe_name(model["name"])
    model_output.mkdir(parents=True, exist_ok=True)
    config.eval_output_dir = str(model_output)
    config.run_dir = model_output / "runner"

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

    envs = make_eval_env(config)
    config.envs = envs
    config.eval_envs = envs
    runner = MACFEnvRunner(config)
    checkpoint = _resolve_checkpoint(model["checkpoint"], Path(manifest_dir))
    _load_actor(runner.policy, checkpoint, config.device)
    records = []
    try:
        env = envs.envs[0]
        for scenario_id in range(episodes):
            records.append(
                _run_episode(
                    config,
                    runner,
                    env,
                    scenario_id,
                    False,
                    permutation_seed,
                    permutation_scope,
                )
            )
            records.append(
                _run_episode(
                    config,
                    runner,
                    env,
                    scenario_id,
                    True,
                    permutation_seed,
                    permutation_scope,
                )
            )
    finally:
        _close_runner(runner, envs)

    for record in records:
        record["model"] = model["name"]
    with (model_output / "episodes.json").open("w", encoding="utf-8") as file:
        json.dump(records, file, ensure_ascii=False, indent=2)
    return {"model": model["name"], "checkpoint": str(checkpoint), "records": records}


def _write_csv(path, records):
    if not records:
        return
    fields = []
    for record in records:
        for key in record:
            if key not in fields:
                fields.append(key)
    with path.open("w", newline="", encoding="utf-8") as file:
        writer = csv.DictWriter(file, fieldnames=fields)
        writer.writeheader()
        writer.writerows(records)


def _summarize(result):
    by_mode = {
        mode: {
            record["scenario_id"]: record
            for record in result["records"]
            if record["mode"] == mode
        }
        for mode in ("original", "permuted")
    }
    scenario_ids = sorted(set(by_mode["original"]) & set(by_mode["permuted"]))
    original = [by_mode["original"][index] for index in scenario_ids]
    permuted = [by_mode["permuted"][index] for index in scenario_ids]
    success_to_failure = sum(
        a["success"] == 1 and b["success"] == 0
        for a, b in zip(original, permuted)
    )
    failure_to_success = sum(
        a["success"] == 0 and b["success"] == 1
        for a, b in zip(original, permuted)
    )
    return {
        "model": result["model"],
        "checkpoint": result["checkpoint"],
        "num_episodes": len(scenario_ids),
        "original_success": float(np.mean([row["success"] for row in original])),
        "permuted_success": float(np.mean([row["success"] for row in permuted])),
        "success_delta": float(
            np.mean([row["success"] for row in permuted])
            - np.mean([row["success"] for row in original])
        ),
        "success_to_failure": int(success_to_failure),
        "failure_to_success": int(failure_to_success),
        "original_completion": float(
            np.mean([row["task_completion_ratio"] for row in original])
        ),
        "permuted_completion": float(
            np.mean([row["task_completion_ratio"] for row in permuted])
        ),
        "original_length": float(
            np.mean([row["episode_length"] for row in original])
        ),
        "permuted_length": float(
            np.mean([row["episode_length"] for row in permuted])
        ),
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config_file", required=True)
    parser.add_argument("--models_file", required=True)
    parser.add_argument("--output_dir", required=True)
    parser.add_argument("--episodes", type=int, default=128)
    parser.add_argument("--permutation_seed", type=int, default=760000)
    parser.add_argument(
        "--permutation_scope",
        choices=("deliver", "explore", "all"),
        default="deliver",
    )
    parser.add_argument("--cuda", type=int, default=1)
    parser.add_argument("--max_workers", type=int, default=2)
    parser.add_argument("--model_name_pattern")
    parser.add_argument("--detach", action="store_true")
    args = parser.parse_args()

    config_file = Path(args.config_file).resolve()
    models_file = Path(args.models_file).resolve()
    output_dir = Path(args.output_dir).resolve()
    if args.detach:
        output_dir.parent.mkdir(parents=True, exist_ok=True)
        command = [
            sys.executable,
            str(Path(__file__).resolve()),
            *[argument for argument in sys.argv[1:] if argument != "--detach"],
        ]
        log_path = output_dir.parent / f"{output_dir.name}.log"
        with log_path.open("w", encoding="utf-8") as log_file:
            process = subprocess.Popen(
                command,
                cwd=REPO_ROOT,
                stdin=subprocess.DEVNULL,
                stdout=log_file,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
        pid_path = output_dir.parent / f"{output_dir.name}.pid"
        pid_path.write_text(f"{process.pid}\n", encoding="utf-8")
        print(
            f"Detached permutation eval pid={process.pid} "
            f"output_dir={output_dir}"
        )
        return
    output_dir.mkdir(parents=True, exist_ok=False)
    with models_file.open("r", encoding="utf-8") as file:
        manifest = yaml.safe_load(file) or {}
    models = manifest["models"]
    if args.model_name_pattern:
        models = [
            model
            for model in models
            if args.model_name_pattern in model["name"]
        ]
        if not models:
            raise ValueError(
                f"No model matched {args.model_name_pattern!r}."
            )

    # Initialize the shared CUDA context before worker threads construct
    # independent runners.  Concurrent lazy initialization can otherwise make
    # one thread observe an incomplete CUDA driver state.
    if args.cuda >= 0:
        if not torch.cuda.is_available():
            raise RuntimeError(f"CUDA device {args.cuda} requested but unavailable.")
        torch.cuda.set_device(args.cuda)
        torch.empty(1, device=f"cuda:{args.cuda}")

    tasks = [
        (
            model,
            str(config_file),
            str(models_file.parent),
            str(output_dir),
            args.episodes,
            args.permutation_seed,
            args.permutation_scope,
            args.cuda,
        )
        for model in models
    ]
    # Evaluate models sequentially inside each process.  Each runner mutates
    # global NumPy/Torch RNG state and owns a CUDA model, so threads would make
    # paired scenarios nondeterministic and race CUDA initialization.  Campaign
    # supervisors can still run independent permutation scopes in parallel.
    results = []
    for task in tasks:
        try:
            result = _evaluate_one(task)
        except Exception:
            (output_dir / "failure.log").write_text(
                traceback.format_exc(), encoding="utf-8"
            )
            raise
        results.append(result)
        print(f"Completed {task[0]['name']}", flush=True)

    summaries = [_summarize(result) for result in results]
    summaries.sort(key=lambda row: row["model"])
    _write_csv(output_dir / "summary.csv", summaries)
    with (output_dir / "summary.json").open("w", encoding="utf-8") as file:
        json.dump(summaries, file, ensure_ascii=False, indent=2)
    print(f"Wrote {output_dir}", flush=True)


if __name__ == "__main__":
    main()