import torch
import torch.nn as nn
import gymnasium as gym

# ── 高层 Manager ──────────────────────────────
class Manager(nn.Module):
    def __init__(self, state_dim, goal_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, 128), nn.ReLU(),
            nn.Linear(128, goal_dim)   # 输出子目标向量
        )
    def forward(self, state):
        return self.net(state)

# ── 低层 Worker ───────────────────────────────
class Worker(nn.Module):
    def __init__(self, state_dim, goal_dim, action_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim + goal_dim, 128), nn.ReLU(),
            nn.Linear(128, action_dim)
        )
    def forward(self, state, goal):
        x = torch.cat([state, goal], dim=-1)
        return self.net(x)

# ── 内在奖励（子目标完成度）─────────────────────
def intrinsic_reward(next_state, goal):
    return -torch.norm(next_state - goal, dim=-1)

# ── 训练主循环 ────────────────────────────────
env = gym.make("MiniGrid-FourRooms-v0")
manager = Manager(state_dim=64, goal_dim=8)
worker  = Worker(state_dim=64, goal_dim=8, action_dim=env.action_space.n)

manager_opt = torch.optim.Adam(manager.parameters(), lr=1e-3)
worker_opt  = torch.optim.Adam(worker.parameters(),  lr=1e-3)

k = 10  # Manager 每 k 步重新设定子目标
for episode in range(5000):
    obs, _ = env.reset()
    state = torch.tensor(obs["image"].flatten(), dtype=torch.float32)
    goal  = manager(state.unsqueeze(0)).squeeze(0).detach()

    ep_reward = 0
    for step in range(200):
        # Worker 选动作
        logits = worker(state.unsqueeze(0), goal.unsqueeze(0))
        action = torch.distributions.Categorical(logits=logits).sample().item()

        next_obs, ext_reward, done, _, _ = env.step(action)
        next_state = torch.tensor(next_obs["image"].flatten(), dtype=torch.float32)

        # 内在奖励训练 Worker
        i_reward = intrinsic_reward(next_state.unsqueeze(0), goal.unsqueeze(0))
        worker_loss = -i_reward.mean()
        worker_opt.zero_grad(); worker_loss.backward(); worker_opt.step()

        # 每 k 步更新 Manager
        if (step + 1) % k == 0:
            new_goal = manager(next_state.unsqueeze(0)).squeeze(0)
            manager_loss = -torch.tensor(ext_reward)  # 用外部奖励训练 Manager
            manager_opt.zero_grad(); manager_loss.backward(); manager_opt.step()
            goal = new_goal.detach()

        state = next_state
        ep_reward += ext_reward
        if done: break

    if episode % 100 == 0:
        print(f"Episode {episode}, Reward: {ep_reward:.2f}")

# 每条样本：游戏局势描述 → 高质量战术指令
sft_sample = {
    "instruction": "当前局势如下：\n我方经济领先800金，敌方打野刚被击杀，上路二塔血量30%。\n请给出接下来30秒的战术决策。",
    "output": "立即集合上路，趁敌方打野复活CD（约40s）强推二塔。上单负责正面吸引仇恨，射手和法师输出塔体，辅助插眼防止敌方支援。推塔后立即撤退不强留。"
}

dpo_sample = {
    "prompt": "当前局势：我方全员低血量（均低于40%），但敌方高地塔血量仅剩10%。请给出决策。",
    "chosen":   "立即撤退回城补给，切勿强攻。低血量强推高地风险极高，等待下一个团战机会。",
    "rejected": "直接强推高地，塔血量很低一波可以拆掉。"
    # rejected 看似有道理但忽略了低血量被反杀的风险
}

# 方案A：GPT-4 合成（推荐，快速获取大量数据）
import openai

def generate_sft_sample(scenario: str) -> dict:
    resp = openai.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "你是王者荣耀职业教练，给出专业战术决策"},
            {"role": "user",   "content": f"局势：{scenario}\n请给出高质量战术建议"}
        ]
    )
    return {"instruction": scenario, "output": resp.choices[0].message.content}

# 方案B：从职业选手录像中提取（更真实，但需要人工标注）
# 截帧 → 状态识别 → 对应操作序列 → 人工筛选高质量片段

# 使用 LLaMA-Factory（推荐，支持 Qwen2.5 + LoRA，配置简单）
# pip install llamafactory

# dataset_info.json 中注册数据集
{
  "game_sft": {
    "file_name": "game_sft.json",
    "formatting": "alpaca",
    "columns": {"prompt": "instruction", "response": "output"}
  }
}

from trl import DPOTrainer, DPOConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# 加载 SFT 后的模型作为起点
base_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
model = PeftModel.from_pretrained(base_model, "./sft_output")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")

# DPO 配置
training_args = DPOConfig(
    beta=0.1,                    # KL 惩罚系数，控制偏离 SFT 模型的程度
    learning_rate=5e-5,
    num_train_epochs=1,
    per_device_train_batch_size=2,
    output_dir="./dpo_output",
)

# 加载偏好数据（chosen / rejected 对）
from datasets import load_dataset
dpo_dataset = load_dataset("json", data_files="game_dpo.json")["train"]

trainer = DPOTrainer(
    model=model,
    args=training_args,
    train_dataset=dpo_dataset,
    tokenizer=tokenizer,
)
trainer.train()

import json
from transformers import pipeline

model_sft = pipeline("text-generation", model="./sft_output")
model_dpo = pipeline("text-generation", model="./dpo_output")

def evaluate_model(model_pipeline, test_cases):
    scores = {"executable": 0, "reasonable": 0, "total": len(test_cases)}
    for case in test_cases:
        output = model_pipeline(case["prompt"], max_new_tokens=200)[0]["generated_text"]
        # 自动评估：可执行性（规则校验）
        if not any(bad in output for bad in ["低血量强推", "1v5", "送人头"]):
            scores["executable"] += 1
    return scores

# 对比实验
print("SFT 模型：", evaluate_model(model_sft, test_cases))
print("DPO 模型：", evaluate_model(model_dpo, test_cases))

from dataclasses import dataclass
from typing import List

@dataclass
class HeroState:
    name: str
    hp_ratio: float      # 当前血量百分比
    position: str        # "上路/中路/下路/jungle/高地"
    skill_ready: bool    # 大招是否就绪
    economy: int         # 金币数

@dataclass
class GameState:
    time_sec: int
    our_team: List[HeroState]
    enemy_team: List[HeroState]
    map_events: List[str]   # ["龙刷新30s后", "上路一塔已破", "敌方打野现身下路"]
    score_diff: int         # 我方经济领先值（负数为落后）

def serialize_state(state: GameState) -> str:
    lines = [f"当前时间：{state.time_sec // 60}分{state.time_sec % 60}秒"]
    lines.append(f"经济差：{'领先' if state.score_diff > 0 else '落后'} {abs(state.score_diff)} 金币")
    lines.append("我方状态：")
    for h in state.our_team:
        skill = "大招就绪" if h.skill_ready else "大招CD中"
        lines.append(f"  - {h.name}：血量{h.hp_ratio:.0%}，位于{h.position}，{skill}")
    lines.append("敌方动态：")
    for h in state.enemy_team:
        lines.append(f"  - {h.name}：血量{h.hp_ratio:.0%}，位于{h.position}")
    lines.append("地图事件：" + "；".join(state.map_events))
    return "\n".join(lines)

SYSTEM_PROMPT = """你是王者荣耀顶级战术指挥官，擅长宏观决策。
请严格按以下 JSON 格式输出，不要输出其他内容：
{
  "reasoning": "分析过程（50字内）",
  "objective": "当前主要目标",
  "priority": "high/medium/low",
  "actions": [
    {"hero": "英雄名", "action": "具体行动", "position": "目标位置"}
  ],
  "warning": "需要注意的威胁（可为null）"
}"""

def build_prompt(state: GameState) -> str:
    state_text = serialize_state(state)
    return f"{state_text}\n\n请给出当前最优战术决策："

import json
from openai import OpenAI  # Qwen2.5 兼容 OpenAI 接口

# 本地部署 Qwen2.5（ollama）
client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")

def get_decision(state: GameState) -> dict:
    prompt = build_prompt(state)
    response = client.chat.completions.create(
        model="qwen2.5:7b",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": prompt}
        ],
        temperature=0.3,   # 低温保证决策稳定性
    )
    text = response.choices[0].message.content
    return json.loads(text)

# 示例输出：
# {
#   "reasoning": "经济领先且敌方打野现身下路，上中路压制时机成熟",
#   "objective": "推倒上路二塔",
#   "priority": "high",
#   "actions": [
#     {"hero": "李白", "action": "压制上路", "position": "上路二塔"},
#     {"hero": "诸葛亮", "action": "跟进支援", "position": "上路"}
#   ],
#   "warning": "注意敌方打野可能绕后"
# }

# 自动校验：动作是否可执行（英雄血量/位置合理性）
def validate_decision(decision: dict, state: GameState) -> dict:
    issues = []
    our_heroes = {h.name: h for h in state.our_team}
    for action in decision["actions"]:
        hero = our_heroes.get(action["hero"])
        if hero and hero.hp_ratio < 0.3:
            issues.append(f"{action['hero']} 血量过低（{hero.hp_ratio:.0%}），不宜进攻")
    decision["issues"] = issues
    decision["executable"] = len(issues) == 0
    return decision

# 评测指标收集
def evaluate_batch(states, ground_truth_actions):
    results = []
    for state, gt in zip(states, ground_truth_actions):
        decision = get_decision(state)
        validated = validate_decision(decision, state)
        results.append({
            "executable_rate": validated["executable"],
            "objective_match": decision["objective"] == gt["objective"],
            "latency_ms": ...  # 记录推理耗时
        })
    return results

import torch
from sentence_transformers import SentenceTransformer

encoder = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")

def alignment_reward(g_text: str, trajectory_summary: str, weight: float = 0.3) -> float:
    """
    计算 LLM 目标语言 与 RL 实际轨迹摘要 的语义相似度，作为额外奖励信号
    g_text:             LLM 输出的目标，如 "集合上路推二塔"
    trajectory_summary: RL 轨迹自动生成摘要，如 "英雄移动到上路并攻击防御塔"
    """
    g_embed  = encoder.encode(g_text,             convert_to_tensor=True)
    t_embed  = encoder.encode(trajectory_summary, convert_to_tensor=True)
    similarity = torch.nn.functional.cosine_similarity(g_embed, t_embed, dim=0)
    return weight * similarity.item()   # 相似度高 → 奖励高，促使 RL 行为与 LLM 目标对齐

def summarize_trajectory(states: list, actions: list) -> str:
    """将 RL 轨迹转为自然语言，供 LLM 理解"""
    action_map = {0: "向上移动", 1: "向下移动", 2: "向左移动",
                  3: "向右移动", 4: "攻击", 5: "使用技能", 6: "撤退"}
    action_counts = {}
    for a in actions:
        desc = action_map.get(a, "未知")
        action_counts[desc] = action_counts.get(desc, 0) + 1
    summary_parts = [f"{desc}{count}次" for desc, count in action_counts.items()]
    final_pos = states[-1].get("position", "未知位置") if states else "未知"
    return f"英雄执行了：{'、'.join(summary_parts)}，最终到达{final_pos}"

import torch.nn as nn

class GoalConditionedPPO(nn.Module):
    """目标条件化 PPO：将 LLM 目标嵌入拼接到状态中"""
    def __init__(self, state_dim: int, goal_dim: int, action_dim: int):
        super().__init__()
        self.actor = nn.Sequential(
            nn.Linear(state_dim + goal_dim, 256), nn.ReLU(),
            nn.Linear(256, 128),                  nn.ReLU(),
            nn.Linear(128, action_dim)
        )
        self.critic = nn.Sequential(
            nn.Linear(state_dim + goal_dim, 256), nn.ReLU(),
            nn.Linear(256, 128),                  nn.ReLU(),
            nn.Linear(128, 1)
        )

    def forward(self, state: torch.Tensor, goal_embed: torch.Tensor):
        x = torch.cat([state, goal_embed], dim=-1)
        return self.actor(x), self.critic(x)

    def get_total_reward(self, env_reward: float, g_text: str,
                         traj_summary: str) -> float:
        r_align = alignment_reward(g_text, traj_summary)
        return env_reward + r_align

REFLECT_PROMPT = """你是游戏 AI 的反思模块。
目标：{goal}
实际执行轨迹：{trajectory}
环境奖励：{reward}

请分析：
1. 行为是否与目标一致？
2. 失败原因是什么？
3. 下次应该如何调整目标或策略？

以 JSON 格式输出：aligned"""

class ReflexionMemory:
    """存储反思历史，注入到下一轮 LLM 规划的 Prompt 中"""
    def __init__(self, max_size: int = 5):
        self.memory = []
        self.max_size = max_size

    def add(self, reflection: dict):
        self.memory.append(reflection)
        if len(self.memory) > self.max_size:
            self.memory.pop(0)

    def to_prompt(self) -> str:
        if not self.memory:
            return ""
        lines = ["历史反思经验："]
        for i, r in enumerate(self.memory[-3:], 1):   # 只取最近3条
            lines.append(f"{i}. {r.get('next_adjustment', '')}")
        return "\n".join(lines)

def reflect(llm_client, goal: str, trajectory: str, reward: float) -> dict:
    import json
    prompt = REFLECT_PROMPT.format(goal=goal, trajectory=trajectory, reward=reward)
    resp = llm_client.chat.completions.create(
        model="qwen2.5:7b",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.2
    )
    return json.loads(resp.choices[0].message.content)

from openai import OpenAI

def train(env, llm_client, agent, memory: ReflexionMemory, episodes: int = 2000):
    for ep in range(episodes):
        obs, _ = env.reset()
        state_text = serialize_state(obs)           # 项目④的序列化器复用

        # Layer 1：LLM 规划（注入 Reflexion Memory）
        history_prompt = memory.to_prompt()
        plan_prompt = f"{history_prompt}\n{state_text}\n请给出当前战术目标（一句话）："
        g_text = llm_client.chat.completions.create(
            model="qwen2.5:7b",
            messages=[{"role": "user", "content": plan_prompt}],
            temperature=0.3
        ).choices[0].message.content.strip()

        g_embed = torch.tensor(encoder.encode(g_text), dtype=torch.float32)

        # Layer 2：RL 执行
        states, actions, total_reward = [], [], 0
        state_tensor = torch.tensor(obs["image"].flatten(), dtype=torch.float32)

        for step in range(200):
            logits, _ = agent(state_tensor.unsqueeze(0), g_embed.unsqueeze(0))
            action = torch.distributions.Categorical(logits=logits).sample().item()
            next_obs, r_env, done, _, _ = env.step(action)

            traj_summary = summarize_trajectory(states, actions)
            r_total = agent.get_total_reward(r_env, g_text, traj_summary)

            states.append(obs); actions.append(action)
            total_reward += r_total
            state_tensor = torch.tensor(next_obs["image"].flatten(), dtype=torch.float32)
            if done: break

        # Layer 3：Reflexion 反思 → 更新 Memory
        traj_summary = summarize_trajectory(states, actions)
        reflection = reflect(llm_client, g_text, traj_summary, total_reward)
        memory.add(reflection)

        if ep % 100 == 0:
            aligned = reflection.get("aligned", False)
            print(f"Ep {ep:4d} | Reward: {total_reward:.2f} | "
                  f"Goal: {g_text[:20]}... | Aligned: {aligned}")