import argparse
import os
import sys
import random
from datetime import datetime

import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm
from torch.utils.data import DataLoader

from .model import SeperatedModels
from .model import (
    MASK_TOKEN, MAP_W, MAP_H, EPOCHS, VQ_GAMMA
)
from .utils import (
    compute_remaining, MAP_SIZE, summarize_codebook_hits
)
from .sample import full_generate
from .dataset import GinkaSeperatedDataset, rect_mask, ensure_wall_connection
from shared.image import matrix_to_image_cv, annotate

# 图块 ID 定义：
# 0. 空地   1. 墙壁   2. 普通门   3. 资源   4. 怪物   5. 入口   6. 机关门   7. 掩码（MASK_TOKEN）

# 三阶段级联地图生成训练脚本
#
# 整体架构：
#   VQ-VAE（三组独立编码器 vq1/vq2/vq3）将三阶段地图上下文分别编码为离散潜变量，
#   再由三个独立 VectorQuantizer 分别量化为 z_q1/z_q2/z_q3；
#   三个独立 MaskGIT（mg1/mg2/mg3）分别以各自阶段 z_q 和 struct_inject 为条件，
#   逐阶段迭代解码地图图块序列。
#
# 三阶段生成目标：
#   stage1 → floor / wall（地图骨架）
#   stage2 → door / monster / entrance（功能性实体）
#   stage3 → resource（资源点）

device = torch.device(
    "cuda:0" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)

disable_tqdm = not sys.stdout.isatty()

# 训练与推理超参
VQ_BETA = 0.5 # 承诺损失权重
STAGE1_CE_WEIGHT = 1.0 # Stage1 CE 损失权重
STAGE2_CE_WEIGHT = 1.0 # Stage2 CE 损失权重
STAGE3_CE_WEIGHT = 1.0 # Stage3 CE 损失权重
MG_Z_DROPOUT = 0.1 # z 隐变量 Dropout 概率
BATCH_SIZE = 64 # 每批样本数
CHECKPOINT = 20 # 每隔多少 epoch 保存检查点并执行验证
GENERATE_STEP = 18 # MaskGIT 采样步数
SEED_SAMPLE_STEPS = 24 # 墙壁种子生长采样步数
SUBSET_WEIGHTS = (0.5, 0.3, 0.2) # 每个子集的概率

def _str2bool(v: str):
    if isinstance(v, bool): return v
    if v.lower() in ('true', '1', 'yes'): return True
    if v.lower() in ('false', '0', 'no'): return False
    raise argparse.ArgumentTypeError(f"布尔值应为 True/False，收到: {v!r}")

def parse_arguments():
    parser = argparse.ArgumentParser(description="三阶段级联训练")
    parser.add_argument("--resume", type=_str2bool, default=False)
    parser.add_argument("--state", type=str, default="", help="续训时检查点路径")
    parser.add_argument("--train", type=str, default="ginka-dataset.json")
    parser.add_argument("--validate", type=str, default="ginka-eval.json")
    parser.add_argument("--load_optim", type=_str2bool, default=True)
    return parser.parse_args()

def cross_entropy_loss(logits, target, mask):
    # logits: [B, L, C]，target: [B, L]，mask: [B, L] bool（True = 参与 loss）
    loss = F.cross_entropy(logits.permute(0, 2, 1), target, reduction='none')
    masked = loss[mask]
    if masked.numel() == 0:
        return torch.tensor(0.0, device=logits.device, requires_grad=True)
    return masked.mean()

def apply_z_dropout(
    z_q: torch.Tensor,
    mask_embedding: nn.Parameter,
    drop_prob: float
) -> torch.Tensor:
    # 以 drop_prob 概率将 z_q 中的码字替换为可学习 mask 嵌入
    # z_q: [B, L, d_z], mask_embedding: [1, 1, d_z]
    mask = torch.rand(z_q.shape[0], z_q.shape[1], 1, device=z_q.device) < drop_prob
    return z_q * (~mask).float() + mask_embedding * mask.float()

def quantize_stage_latents(
    models: SeperatedModels,
    z_e1: torch.Tensor,
    z_e2: torch.Tensor,
    z_e3: torch.Tensor
) -> tuple:
    z_q1, _, commit_loss1, _, code_hits1, entropy1 = models.quantizer1(z_e1)
    z_q2, _, commit_loss2, _, code_hits2, entropy2 = models.quantizer2(z_e2)
    z_q3, _, commit_loss3, _, code_hits3, entropy3 = models.quantizer3(z_e3)

    commit_loss = (commit_loss1 + commit_loss2 + commit_loss3) / 3
    entropy_loss = entropy1 + entropy2 + entropy3
    code_hits = (code_hits1, code_hits2, code_hits3)
    return (z_q1, z_q2, z_q3), commit_loss, code_hits, entropy_loss

# 每张图 2 行 x 4 列：MASKED | PREDICTED | MASKED | PREDICTED
def make_compare_grid(samples_src, samples_pred, map_keys, tile_dict, TILE_SIZE):
    SEP = 3
    img_h = MAP_H * TILE_SIZE
    img_w = MAP_W * TILE_SIZE
    cols = 4
    rows = 2
    grid = np.ones((rows * img_h + (rows + 1) * SEP, cols * img_w + (cols + 1) * SEP, 3), dtype=np.uint8) * 255
    for r in range(rows):
        row_y = SEP + r * (img_h + SEP)
        for c in range(cols):
            is_src = (c % 2 == 0)
            idx = r * 2 + c // 2
            col_x = SEP + c * (img_w + SEP)
            data = samples_src[idx] if is_src else samples_pred[idx]
            grid[row_y:row_y + img_h, col_x:col_x + img_w] = annotate(
                matrix_to_image_cv(data, tile_dict, TILE_SIZE), map_keys[idx], y=14
            )
    # 每两张图片之间的分隔线
    for c in range(cols - 1):
        line_x = (c + 1) * (img_w + SEP) + SEP // 2
        grid[:, line_x:line_x + SEP, :] = 180
    return grid

def visualize_seed_growth(
    train_dataset: GinkaSeperatedDataset,
    models: SeperatedModels,
    device: torch.device,
    tile_dict,
    epoch: int
):
    save_dir = f"result/seperated/e{epoch}"
    os.makedirs(save_dir, exist_ok=True)
    samples_src = []
    samples_pred = []
    map_keys = []
    for _ in range(20):
        sample = train_dataset.random_sample_map()
        struct_t = sample["struct_inject"].to(device).reshape(1, -1)
        target_density_t = sample["target_density"].to(device).reshape(1, -1)
        enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE)
        enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE)
        enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE)
        enc1_np = enc1_t.cpu().numpy().reshape(MAP_H, MAP_W)
        with torch.no_grad():
            z_e1 = models.vq1(enc1_t)
            z_e2 = models.vq2(enc2_t)
            z_e3 = models.vq3(enc3_t)
            z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3)
            z1, z2, z3 = z_q
            inp = torch.full((1, MAP_SIZE), 0, dtype=torch.long, device=device)
            seed_count = random.randint(7, 14)
            seed_idx = torch.randperm(MAP_SIZE, device=device)[:seed_count]
            inp[0, seed_idx] = 1
            inp_vis = inp.clone()
            inp_vis[inp_vis == 0] = MASK_TOKEN
            _, _, merged = full_generate(
                inp, z1, z2, z3, struct_t, target_density_t, models,
                steps=SEED_SAMPLE_STEPS, seed_mode=True, stage1_method="growth"
            )
        samples_src.append(inp_vis.cpu().numpy().reshape(MAP_H, MAP_W))
        samples_pred.append(merged[0])
        map_keys.append(sample["map_name"])
    for i in range(5):
        grid = make_compare_grid(
            samples_src[i*4:i*4+4], samples_pred[i*4:i*4+4],
            map_keys[i*4:i*4+4], tile_dict, 32
        )
        cv2.imwrite(f"{save_dir}/seed_growth_{i}.png", grid)

def visualize_seed_maskgit(
    train_dataset: GinkaSeperatedDataset,
    models: SeperatedModels,
    device: torch.device,
    tile_dict,
    epoch: int
):
    save_dir = f"result/seperated/e{epoch}"
    os.makedirs(save_dir, exist_ok=True)
    samples_src = []
    samples_pred = []
    map_keys = []
    for _ in range(20):
        sample = train_dataset.random_sample_map()
        struct_t = sample["struct_inject"].to(device).reshape(1, -1)
        target_density_t = sample["target_density"].to(device).reshape(1, -1)
        enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE)
        enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE)
        enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE)
        enc1_np = enc1_t.cpu().numpy().reshape(MAP_H, MAP_W)
        with torch.no_grad():
            z_e1 = models.vq1(enc1_t)
            z_e2 = models.vq2(enc2_t)
            z_e3 = models.vq3(enc3_t)
            z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3)
            z1, z2, z3 = z_q
            inp = torch.full((1, MAP_SIZE), MASK_TOKEN, dtype=torch.long, device=device)
            _, _, merged = full_generate(
                inp, z1, z2, z3, struct_t, target_density_t, models,
                steps=GENERATE_STEP, seed_mode=True, stage1_method="maskgit"
            )
        samples_src.append(inp.cpu().numpy().reshape(MAP_H, MAP_W))
        samples_pred.append(merged[0])
        map_keys.append(sample["map_name"])
    for i in range(5):
        grid = make_compare_grid(
            samples_src[i*4:i*4+4], samples_pred[i*4:i*4+4],
            map_keys[i*4:i*4+4], tile_dict, 32
        )
        cv2.imwrite(f"{save_dir}/seed_maskgit_{i}.png", grid)

def visualize_mask_growth(
    train_dataset: GinkaSeperatedDataset,
    models: SeperatedModels,
    device: torch.device,
    tile_dict,
    epoch: int
):
    save_dir = f"result/seperated/e{epoch}"
    os.makedirs(save_dir, exist_ok=True)
    samples_src = []
    samples_pred = []
    map_keys = []
    for _ in range(20):
        sample = train_dataset.random_sample_map()
        raw_map = sample["raw_map"].cpu().numpy().reshape(MAP_H, MAP_W)
        enc1_np = sample["encoder_stage1"].cpu().numpy().reshape(MAP_H, MAP_W)
        enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE)
        enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE)
        enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE)
        struct_t = sample["struct_inject"].to(device).reshape(1, -1)
        target_density_t = sample["target_density"].to(device).reshape(1, -1)

        # 矩形分块掩码 + 确保墙壁连通性，与训练时一致
        ratio = random.uniform(0.3, 0.9)
        rmask = rect_mask(ratio).reshape(MAP_H, MAP_W)
        rmask = ensure_wall_connection(rmask, enc1_np)
        inp = torch.tensor(raw_map.flatten(), dtype=torch.long, device=device).reshape(1, MAP_SIZE)
        inp[0, rmask.flatten()] = MASK_TOKEN
        masked_np = inp.cpu().numpy().reshape(MAP_H, MAP_W)

        with torch.no_grad():
            z_e1 = models.vq1(enc1_t)
            z_e2 = models.vq2(enc2_t)
            z_e3 = models.vq3(enc3_t)
            z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3)
            z1, z2, z3 = z_q
            _, _, merged = full_generate(
                inp, z1, z2, z3, struct_t, target_density_t, models,
                stage1_method="growth"
            )
        samples_src.append(masked_np)
        samples_pred.append(merged[0])
        map_keys.append(sample["map_name"])
    for i in range(5):
        grid = make_compare_grid(
            samples_src[i*4:i*4+4], samples_pred[i*4:i*4+4],
            map_keys[i*4:i*4+4], tile_dict, 32
        )
        cv2.imwrite(f"{save_dir}/mask_growth_{i}.png", grid)

def visualize_mask_maskgit(
    train_dataset: GinkaSeperatedDataset,
    models: SeperatedModels,
    device: torch.device,
    tile_dict,
    epoch: int
):
    save_dir = f"result/seperated/e{epoch}"
    os.makedirs(save_dir, exist_ok=True)
    samples_src = []
    samples_pred = []
    map_keys = []
    for _ in range(20):
        sample = train_dataset.random_sample_map()
        raw_map = sample["raw_map"].cpu().numpy().reshape(MAP_H, MAP_W)
        enc1_np = sample["encoder_stage1"].cpu().numpy().reshape(MAP_H, MAP_W)
        enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE)
        enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE)
        enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE)
        struct_t = sample["struct_inject"].to(device).reshape(1, -1)
        target_density_t = sample["target_density"].to(device).reshape(1, -1)

        # 矩形分块掩码 + 确保墙壁连通性，与训练时一致
        ratio = random.uniform(0.3, 0.9)
        rmask = rect_mask(ratio).reshape(MAP_H, MAP_W)
        rmask = ensure_wall_connection(rmask, enc1_np)
        inp = torch.tensor(raw_map.flatten(), dtype=torch.long, device=device).reshape(1, MAP_SIZE)
        inp[0, rmask.flatten()] = MASK_TOKEN
        masked_np = inp.cpu().numpy().reshape(MAP_H, MAP_W)

        with torch.no_grad():
            z_e1 = models.vq1(enc1_t)
            z_e2 = models.vq2(enc2_t)
            z_e3 = models.vq3(enc3_t)
            z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3)
            z1, z2, z3 = z_q
            _, _, merged = full_generate(
                inp, z1, z2, z3, struct_t, target_density_t, models,
                stage1_method="maskgit"
            )
        samples_src.append(masked_np)
        samples_pred.append(merged[0])
        map_keys.append(sample["map_name"])
    for i in range(5):
        grid = make_compare_grid(
            samples_src[i*4:i*4+4], samples_pred[i*4:i*4+4],
            map_keys[i*4:i*4+4], tile_dict, 32
        )
        cv2.imwrite(f"{save_dir}/mask_maskgit_{i}.png", grid)

def train(device: torch.device):
    args = parse_arguments()

    result = SeperatedModels(device)

    tqdm.write(f"Device: {device}")
    model_list = [
        ("vq1", result.vq1), ("vq2", result.vq2), ("vq3", result.vq3),
        ("mg1", result.mg1), ("mg2", result.mg2), ("mg3", result.mg3),
        ("quantizer1", result.quantizer1), ("quantizer2", result.quantizer2), ("quantizer3", result.quantizer3)
    ]
    total_params = 0
    for name, m in model_list:
        n = sum(p.numel() for p in m.parameters())
        total_params += n
        tqdm.write(f"{name}: {n:,} params")
    tqdm.write(f"Total: {total_params:,} params")

    start_epoch = 0

    if args.resume:
        # 从指定检查点恢复：加载所有模型权重及训练状态
        start_epoch = result.load(args.state, load_optim=args.load_optim, map_location=device)
        tqdm.write(f"Resumed from epoch {start_epoch}: {args.state}")

    os.makedirs("result/seperated", exist_ok=True)

    dataset = GinkaSeperatedDataset(
        args.train, subset_weights=SUBSET_WEIGHTS
    )
    dataloader = DataLoader(
        dataset, batch_size=BATCH_SIZE, shuffle=True
    )

    # 预加载图块图像，键为文件名（不含扩展名），用于可视化时将 ID 映射为像素图
    tile_dict = {}
    for f in os.listdir("tiles"):
        name = os.path.splitext(f)[0]
        img = cv2.imread(f"tiles/{f}", cv2.IMREAD_UNCHANGED)
        if img is not None:
            tile_dict[name] = img

    for epoch in tqdm(range(start_epoch, EPOCHS), desc="Seperated Training", disable=disable_tqdm):
        loss_total = torch.Tensor([0]).to(device)
        loss1_total = torch.Tensor([0]).to(device)
        loss2_total = torch.Tensor([0]).to(device)
        loss3_total = torch.Tensor([0]).to(device)
        commit_total = torch.Tensor([0]).to(device)
        entropy_total = torch.Tensor([0]).to(device)
        code_hits_total = (torch.zeros(result.quantizer1.K, device=device), torch.zeros(result.quantizer2.K, device=device), torch.zeros(result.quantizer3.K, device=device))  # validate

        for batch in tqdm(dataloader, leave=False, desc="Epoch Progress", disable=disable_tqdm):
            # 三阶段各自的掩码输入序列、预测目标和编码器上下文
            inp1 = batch["input_stage1"].to(device).reshape(-1, MAP_SIZE)
            target1 = batch["target_stage1"].to(device).reshape(-1, MAP_SIZE)
            enc1 = batch["encoder_stage1"].to(device).reshape(-1, MAP_SIZE)

            inp2 = batch["input_stage2"].to(device).reshape(-1, MAP_SIZE)
            target2 = batch["target_stage2"].to(device).reshape(-1, MAP_SIZE)
            enc2 = batch["encoder_stage2"].to(device).reshape(-1, MAP_SIZE)

            inp3 = batch["input_stage3"].to(device).reshape(-1, MAP_SIZE)
            target3 = batch["target_stage3"].to(device).reshape(-1, MAP_SIZE)
            enc3 = batch["encoder_stage3"].to(device).reshape(-1, MAP_SIZE)

            # 结构条件向量：[cond_sym, cond_outer]
            struct = batch["struct_inject"].to(device)
            target_density = batch["target_density"].to(device)

            result.optimizer.zero_grad() # 训练循环

            # VQ 编码：各阶段编码器分别处理各自上下文切片
            z_e1 = result.vq1(enc1) # [B, L, d_z]
            z_e2 = result.vq2(enc2)
            z_e3 = result.vq3(enc3)

            # 三阶段分别量化，各自使用独立 codebook
            z_q, commit_loss, code_hits, entropy_loss = quantize_stage_latents(
                result, z_e1, z_e2, z_e3
            )
            z_q1, z_q2, z_q3 = z_q

            # latent dropout：训练时随机丢弃部分码字，替换为可学习 mask 嵌入
            z_q1 = apply_z_dropout(z_q1, result.latent_mask_embedding, MG_Z_DROPOUT)
            z_q2 = apply_z_dropout(z_q2, result.latent_mask_embedding, MG_Z_DROPOUT)
            z_q3 = apply_z_dropout(z_q3, result.latent_mask_embedding, MG_Z_DROPOUT)

            remain1 = compute_remaining(inp1, target_density, 1)
            remain2 = compute_remaining(inp2, target_density, 2)
            remain3 = compute_remaining(inp3, target_density, 3)

            # 三阶段 MaskGIT 前向：各阶段接收自己的 z_q、struct 和动态 remain
            logits1 = result.mg1(inp1, z_q1, struct, remain1)
            logits2 = result.mg2(inp2, z_q2, struct, remain2)
            logits3 = result.mg3(inp3, z_q3, struct, remain3)

            # 三阶段 Cross Entropy
            # Stage1: 仅对掩码位置计算 loss
            # Stage2/3: 掩码 + 空地位置均计算 loss，模型需学习哪里不能放东西
            mask1 = (inp1 == MASK_TOKEN)
            mask2 = (inp2 == MASK_TOKEN) | (inp2 == 0)
            mask3 = (inp3 == MASK_TOKEN) | (inp3 == 0)
            loss1 = cross_entropy_loss(logits1, target1, mask1)
            loss2 = cross_entropy_loss(logits2, target2, mask2)
            loss3 = cross_entropy_loss(logits3, target3, mask3)

            loss1_weighted = STAGE1_CE_WEIGHT * loss1
            loss2_weighted = STAGE2_CE_WEIGHT * loss2
            loss3_weighted = STAGE3_CE_WEIGHT * loss3
            commit_weighted = VQ_BETA * commit_loss + VQ_GAMMA * entropy_loss
            loss = loss1_weighted + loss2_weighted + loss3_weighted + commit_weighted

            loss.backward()
            result.optimizer.step()

            # detach 后累加，避免保留计算图占用显存
            loss_total += loss.detach()
            loss1_total += loss1.detach()
            loss2_total += loss2.detach()
            loss3_total += loss3.detach()
            commit_total += commit_loss.detach()
            entropy_total += entropy_loss.detach()
            code_hits_total = (code_hits_total[0] + code_hits[0].detach(), code_hits_total[1] + code_hits[1].detach(), code_hits_total[2] + code_hits[2].detach())  # accumulate train

        # 每个 epoch 结束后更新学习率
        result.scheduler.step()

        data_length = len(dataloader)
        # 总 loss 去掉 entropy_loss 以便与 CE + VQ 对齐
        loss_display = (loss_total.item() - VQ_GAMMA * entropy_total.item()) / data_length
        entropy_display = entropy_total.item() / data_length
        stats = summarize_codebook_hits(code_hits_total)
        parts = []
        for name in ["q1(stage1)", "q2(stage2)", "q3(stage3)"]:
            s = stats[name]
            parts.append(f"{s['active']}/{s['K']} ppl={s['ppl']:.1f}")
        c = stats["combined"]
        tqdm.write(
            f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] "
            f"E: {epoch + 1} | "
            f"Loss: {loss_display:.4f} | "
            f"CE: {loss1_total.item() / data_length:.4f}, {loss2_total.item() / data_length:.4f}, {loss3_total.item() / data_length:.4f} | "
            f"cmt: {commit_total.item() / data_length:.4f} ent: {entropy_display:+.4f} | "
            f"VQ: {' | '.join(parts)} | "
            f"Total: {c['active']}/{c['K']} ppl={c['ppl']:.1f} | "
            f"LR: {result.scheduler.get_last_lr()[0]:.6f}"
        )

        # 每 CHECKPOINT 个 epoch 执行可视化并保存检查点
        if (epoch + 1) % CHECKPOINT == 0:
            visualize_seed_growth(dataset, result, device, tile_dict, epoch + 1)
            visualize_seed_maskgit(dataset, result, device, tile_dict, epoch + 1)
            visualize_mask_growth(dataset, result, device, tile_dict, epoch + 1)
            visualize_mask_maskgit(dataset, result, device, tile_dict, epoch + 1)
            ckpt_path = f"result/seperated/sep-{epoch + 1}.pth"
            result.save(ckpt_path, epoch + 1, save_optim=True)
            tqdm.write(f"Saved checkpoint: {ckpt_path}")

    # 训练结束后保存最终完整权重
    final_path = "result/seperated.pth"
    result.save(final_path, EPOCHS, save_optim=False)
    tqdm.write(f"Training complete. Final model saved: {final_path}")

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