"""在 DelivCoal 环境中训练 HeteroMRTA 项目提供的 CFM 分配网络。

这里直接使用原项目的 AttentionNet，并按智能体顺序选择任务。场景、可见任务、
可选动作、环境推进、奖励和统计指标仍由 DelivCoal 提供。
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from datetime import datetime
from pathlib import Path

import numpy as np
import torch
from torch.distributions import Categorical

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

from experiments.reproduction.delivcoal_adapter_common import (
    DEFAULT_CONFIG,
    aggregate_rows,
    atomic_json,
    config_dict,
    load_env_config,
    seed_everything,
    summarize_episode,
)
from envs.delivcoal.env_discrete import DiscreteActionEnv


CFM_SOURCE = REPO_ROOT / "third_party/heteromrta"
if str(CFM_SOURCE) not in sys.path:
    sys.path.insert(0, str(CFM_SOURCE))

from attention import AttentionNet  # noqa: E402


def now():
    return datetime.now().isoformat(timespec="seconds")


def model_input_dims(config):
    """根据 DelivCoal 当前的观测格式计算 CFM 的输入维度。"""
    agent_dim = 2 + int(config.cap_dim) + 5
    if (
        bool(config.use_coalition_delivery)
        and getattr(config, "coalition_task_feature_mode", "compact") == "compact"
    ):
        task_dim = 6
    else:
        task_dim = 2 + int(config.cap_dim) + 1 + 1 + 3 + 1
    return agent_dim, task_dim


def atomic_torch_save(path, payload):
    """先写临时文件，再替换正式 checkpoint，避免读到未写完的文件。"""
    path = Path(path)
    temporary = path.with_suffix(path.suffix + ".tmp")
    torch.save(payload, temporary)
    os.replace(temporary, path)


def team_agent_features(env):
    """读取所有智能体的完整特征，供 CFM 依次分配任务。"""
    return env.get_global_state()["agents"].astype(np.float32, copy=True)


def choose_actions(network, env, obs, active_mask, available_actions, device,
                   greedy=False, record=False):
    """按 active_ids 顺序为智能体选择任务，并返回训练所需的输入快照。"""
    n_agents, n_actions = available_actions.shape
    actions = np.asarray(
        [agent.coalition for agent in env.env.agents], dtype=np.int64
    )
    agent_features = team_agent_features(env)
    transitions = []
    active_ids = np.flatnonzero(np.asarray(active_mask).reshape(-1) > 0.5)

    # CFM 在同一个决策时刻逐个处理智能体。
    for agent_id in active_ids:
        available = np.asarray(available_actions[agent_id], dtype=bool)
        if not available.any():
            continue
        # AttentionNet 会使用任务序列的第 0 行汇总信息。先把可选任务移到序列前面，
        # 网络选出序列位置后，再通过 task_order 转回 DelivCoal 的任务编号。
        task_order = np.concatenate(
            [np.flatnonzero(available), np.flatnonzero(~available)]
        ).astype(np.int64)
        ordered_available = available[task_order]
        tasks = torch.as_tensor(
            obs["tasks"][agent_id, task_order], dtype=torch.float32, device=device
        ).unsqueeze(0)
        agents = torch.as_tensor(
            agent_features, dtype=torch.float32, device=device
        ).unsqueeze(0)
        mask = torch.as_tensor(
            ~ordered_available, dtype=torch.bool, device=device
        ).reshape(1, 1, n_actions)
        index = torch.tensor([[[agent_id]]], dtype=torch.long, device=device)
        with torch.no_grad():
            # greedy 用于确定性评估；训练时从概率分布中采样。
            probabilities, _ = network(tasks, agents, mask, index)
            distribution = Categorical(probs=probabilities)
            action = (
                torch.argmax(probabilities, dim=-1)
                if greedy
                else distribution.sample()
            )
        sequence_action = int(action.item())
        action_id = int(task_order[sequence_action])
        actions[agent_id] = action_id
        if record:
            # 保存动作发生时的输入。后续智能体会看到更新后的 agent_features。
            transitions.append(
                (
                    obs["tasks"][agent_id, task_order].copy(),
                    agent_features.copy(),
                    (~ordered_available).copy(),
                    int(agent_id),
                    sequence_action,
                )
            )
        # 第 8 项保存联盟编号。立刻写回本次选择，让后续智能体看到前面的分配结果。
        agent_features[agent_id, 8] = action_id / max(1, n_actions - 1)
    return actions, transitions


def rollout(network, config, seed, device, greedy, record=False,
            policy_seed=None):
    """运行一个完整 episode，并按需保存其中每次任务选择。"""
    seed_everything(seed)
    env = DiscreteActionEnv(config)
    env.seed(seed)
    obs, active, available = env.reset(0)
    if policy_seed is not None:
        # 环境已经用 seed 固定场景。这里单独设置采样种子，让同一场景的多次
        # rollout 可以抽到不同动作。
        seed_everything(policy_seed)
    total_reward = 0.0
    sent_total = 0.0
    transitions = []
    info = {}
    length = 0
    for step in range(int(config.max_steps)):
        actions, current = choose_actions(
            network, env, obs, active, available, device, greedy, record
        )
        obs, rewards, dones, info, active, _, available = env.step(actions)
        total_reward += float(np.mean(rewards))
        sent_total += float(info.get("sent", 0.0))
        transitions.extend(current)
        length = step + 1
        if bool(np.all(dones)):
            break
    row = summarize_episode(info, total_reward, length, sent_total)
    row["seed"] = int(seed)
    return row, transitions


def update_policy(network, optimizer, samples, device, batch_size):
    """用每次 rollout 相对组内平均回报的差值更新 CFM。"""
    optimizer.zero_grad(set_to_none=True)
    rollout_terms = []
    entropy_terms = []
    decision_count = 0
    for transitions, advantage in samples:
        # advantage 是该 rollout 回报减去同组 rollout 的平均回报。
        if not transitions:
            continue
        log_prob_parts = []
        entropy_parts = []
        for start in range(0, len(transitions), batch_size):
            # transition 中保存的是当时实际输入网络的数据，重新前向计算所选动作概率。
            batch = transitions[start : start + batch_size]
            tasks = torch.as_tensor(
                np.stack([item[0] for item in batch]), device=device
            )
            agents = torch.as_tensor(
                np.stack([item[1] for item in batch]), device=device
            )
            masks = torch.as_tensor(
                np.stack([item[2] for item in batch]), device=device
            ).unsqueeze(1)
            indices = torch.as_tensor(
                [item[3] for item in batch], dtype=torch.long, device=device
            ).reshape(-1, 1, 1)
            actions = torch.as_tensor(
                [item[4] for item in batch], dtype=torch.long, device=device
            )
            probabilities, _ = network(tasks, agents, masks, indices)
            distribution = Categorical(probs=probabilities)
            log_prob_parts.append(distribution.log_prob(actions))
            entropy_parts.append(distribution.entropy())
        mean_log_prob = torch.cat(log_prob_parts).mean()
        # 提高高于平均回报动作的概率，降低低于平均回报动作的概率。
        rollout_terms.append(-float(advantage) * mean_log_prob)
        entropy_terms.append(torch.cat(entropy_parts).mean())
        decision_count += len(transitions)
    if not rollout_terms:
        return {"loss": 0.0, "entropy": 0.0, "grad_norm": 0.0,
                "decisions": 0}
    policy_loss = torch.stack(rollout_terms).mean()
    entropy = torch.stack(entropy_terms).mean()
    loss = policy_loss
    loss.backward()
    # 限制梯度大小，避免单组回报差值过大导致更新突变。
    grad_norm = torch.nn.utils.clip_grad_norm_(network.parameters(), 10.0)
    optimizer.step()
    return {
        "loss": float(loss.detach().cpu()),
        "entropy": float(entropy.detach().cpu()),
        "grad_norm": float(torch.as_tensor(grad_norm).detach().cpu()),
        "decisions": int(decision_count),
    }


@torch.no_grad()
def evaluate(network, config, device, seeds):
    training = network.training
    network.eval()
    rows = [rollout(network, config, seed, device, True)[0] for seed in seeds]
    if training:
        network.train()
    return {"episodes": rows, "summary": aggregate_rows(rows)}


def checkpoint_payload(network, optimizer, scheduler, episode, args, config):
    agent_input_dim, task_input_dim = model_input_dims(config)
    return {
        "method": "cfm-delivcoal",
        "episode": int(episode),
        "model": network.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
        "agent_input_dim": agent_input_dim,
        "task_input_dim": task_input_dim,
        "embedding_dim": int(args.embedding_dim),
        "args": vars(args),
        "environment": config_dict(config),
    }


def train(args):
    output = Path(args.output).resolve()
    models = output / "models"
    trainers = output / "trainers"
    output.mkdir(parents=True, exist_ok=True)
    models.mkdir(exist_ok=True)
    trainers.mkdir(exist_ok=True)
    config = load_env_config(args.env_config)
    seed_everything(args.seed)
    device = torch.device(
        f"cuda:{args.cuda}" if args.cuda >= 0 and torch.cuda.is_available() else "cpu"
    )
    agent_input_dim, task_input_dim = model_input_dims(config)
    network = AttentionNet(
        agent_input_dim, task_input_dim, args.embedding_dim
    ).to(device)
    optimizer = torch.optim.Adam(network.parameters(), lr=args.lr)
    scheduler = torch.optim.lr_scheduler.StepLR(
        optimizer, step_size=args.lr_decay_interval, gamma=args.lr_decay
    )
    status = {
        "name": f"cfm-delivcoal-s{args.seed}",
        "state": "training",
        "created_at": now(),
        "method": "CFM",
        "environment": "DelivCoal",
        "environment_config": config_dict(config),
        "seed": args.seed,
        "completed_episodes": 0,
        "total_episodes": args.episodes,
    }
    atomic_json(output / "status.json", status)
    atomic_json(output / "reproduction_config.json", {
        "method": "official HeteroMRTA AttentionNet / CFM",
        "adapter": "DelivCoal sequential allocation adapter",
        "args": vars(args),
        "environment": config_dict(config),
    })
    metrics_path = output / "metrics.jsonl"
    started = time.monotonic()
    for episode in range(1, args.episodes + 1):
        network.train()
        rows = []
        samples = []
        scenario_seed = args.seed * 1_000_000 + episode
        for pomo_id in range(args.pomo_size):
            row, transitions = rollout(
                network, config, scenario_seed, device, False, record=True,
                policy_seed=scenario_seed + 10_000 * (pomo_id + 1),
            )
            rows.append(row)
            samples.append((transitions, 0.0))
        rewards = np.asarray([row["episode_reward"] for row in rows])
        baseline = float(rewards.mean())
        samples = [
            (transitions, float(reward - baseline))
            for (transitions, _), reward in zip(samples, rewards)
        ]
        update = update_policy(network, optimizer, samples, device, args.batch_size)
        scheduler.step()
        payload = checkpoint_payload(
            network, optimizer, scheduler, episode, args, config
        )
        if episode % args.actor_save_interval == 0 or episode == args.episodes:
            atomic_torch_save(models / f"actor_ep{episode:06d}.pt", {
                key: payload[key] for key in (
                    "method", "episode", "model", "agent_input_dim",
                    "task_input_dim", "embedding_dim", "environment"
                )
            })
        if episode % args.trainer_save_interval == 0 or episode == args.episodes:
            atomic_torch_save(trainers / f"trainer_ep{episode:06d}.pt", payload)
        evaluation = None
        if episode % args.eval_interval == 0 or episode == args.episodes:
            seeds = range(args.validation_seed,
                          args.validation_seed + args.validation_episodes)
            evaluation = evaluate(network, config, device, seeds)
        record = {
            "episode": episode,
            "environment_steps": int(episode * args.pomo_size * config.max_steps),
            "train": aggregate_rows(rows),
            "update": update,
            "validation": evaluation,
            "recorded_at": now(),
        }
        with metrics_path.open("a", encoding="utf-8") as file:
            file.write(json.dumps(record, ensure_ascii=False) + "\n")
        status["completed_episodes"] = episode
        status["updated_at"] = now()
        atomic_json(output / "status.json", status)
        elapsed = time.monotonic() - started
        fps = episode * args.pomo_size * config.max_steps / max(elapsed, 1e-6)
        eta = elapsed / episode * (args.episodes - episode)
        print(
            f"rollout {episode}/{args.episodes}, timesteps "
            f"{episode * args.pomo_size * config.max_steps}, FPS {fps:.1f}, "
            f"time cost {elapsed:.1f}s, expected time {eta:.1f}s, "
            f"reward {rewards.mean():.4f}, success "
            f"{np.mean([row['success'] for row in rows]):.3f}",
            flush=True,
        )
    status.update(state="completed", finished_at=now(), actor_checkpoint=str(
        models / f"actor_ep{args.episodes:06d}.pt"
    ))
    atomic_json(output / "status.json", status)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", required=True)
    parser.add_argument("--env-config", default=str(DEFAULT_CONFIG))
    parser.add_argument("--seed", type=int, default=10)
    parser.add_argument("--episodes", type=int, default=1500)
    parser.add_argument("--pomo-size", type=int, default=4)
    parser.add_argument("--embedding-dim", type=int, default=128)
    parser.add_argument("--lr", type=float, default=1e-5)
    parser.add_argument("--lr-decay-interval", type=int, default=32)
    parser.add_argument("--lr-decay", type=float, default=0.98)
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--actor-save-interval", type=int, default=1)
    parser.add_argument("--trainer-save-interval", type=int, default=100)
    parser.add_argument("--eval-interval", type=int, default=50)
    parser.add_argument("--validation-episodes", type=int, default=32)
    parser.add_argument("--validation-seed", type=int, default=10_000)
    parser.add_argument("--cuda", type=int, default=0)
    return parser.parse_args()


if __name__ == "__main__":
    train(parse_args())