#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
CUT L通道 训练脚本
==================
仅对车牌 L (亮度) 通道做纹理迁移训练，a,b 颜色通道不动。

用法:
  python train.py \
      --domain-a ./gan_data/domain_A_synthetic/police \
      --domain-b ./gan_data/domain_B_real \
      --epochs 100 \
      --batch 4 \
      --output ./output/cut_police
"""

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

import torch
import torch.nn as nn
import torch.optim as optim

# 添加 tools 目录到路径
sys.path.insert(0, str(Path(__file__).resolve().parent))

from models import (
    LChannelGenerator, PatchGANDiscriminator,
    PatchNCELoss, init_weights, count_params,
)
from dataset import make_loader


def train(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    # ---- 数据 ----
    loader_a = make_loader(args.domain_a, args.batch, (args.img_w, args.img_h),
                           is_real=False, num_workers=args.workers)
    loader_b = make_loader(args.domain_b, args.batch, (args.img_w, args.img_h),
                           is_real=True, num_workers=args.workers)

    print(f"Domain A (合成): {len(loader_a.dataset)} 张")
    print(f"Domain B (真实): {len(loader_b.dataset)} 张")
    print(f"Batch size:      {args.batch}")
    print(f"Epochs:          {args.epochs}")

    # ---- 模型 ----
    G = LChannelGenerator(n_res=args.n_res).to(device)
    D = PatchGANDiscriminator().to(device)

    init_weights(G)
    init_weights(D)

    print(f"G params: {count_params(G):,}")
    print(f"D params: {count_params(D):,}")

    # ---- 损失 & 优化器 ----
    criterion_gan = nn.MSELoss()           # LSGAN
    criterion_nce = PatchNCELoss(
        nce_layers=args.nce_layers,
        nce_t=args.nce_t,
        n_patches=args.n_patches,
    ).to(device)
    criterion_idt = nn.L1Loss()

    opt_G = optim.Adam(G.parameters(), lr=args.lr, betas=(0.5, 0.999))
    opt_D = optim.Adam(D.parameters(), lr=args.lr, betas=(0.5, 0.999))

    # ---- 日志 ----
    os.makedirs(args.output, exist_ok=True)

    # ---- 训练循环 ----
    for epoch in range(1, args.epochs + 1):
        epoch_start = time.time()
        G_loss_total = 0.0
        D_loss_total = 0.0
        n_batches = 0

        for real_a, real_b in zip(loader_a, loader_b):
            real_a = real_a.to(device)      # (B, 1, H, W) 合成L
            real_b = real_b.to(device)      # (B, 1, H, W) 真实L
            B = real_a.size(0)

            # ---- 1. 更新 D ----
            fake_b = G(real_a).detach()
            # G 不用返回特征

            pred_real = D(real_b)
            pred_fake = D(fake_b)

            loss_D_real = criterion_gan(pred_real,
                                         torch.ones_like(pred_real))
            loss_D_fake = criterion_gan(pred_fake,
                                         torch.zeros_like(pred_fake))
            loss_D = (loss_D_real + loss_D_fake) * 0.5

            opt_D.zero_grad()
            loss_D.backward()
            opt_D.step()

            # ---- 2. 更新 G ----
            fake_b, feat_q = G(real_a, return_features=True)

            # GAN loss
            pred_fake = D(fake_b)
            loss_G_gan = criterion_gan(pred_fake,
                                        torch.ones_like(pred_fake))

            # PatchNCE loss (对比学习)
            _, feat_k = G(real_b, return_features=True)
            loss_G_nce = criterion_nce(feat_q, feat_k)

            # Identity loss: 裁剪输出到输入尺寸（UNet可能多出几像素）
            identity = G(real_a)
            if identity.shape != real_a.shape:
                identity = identity[:, :, :real_a.shape[2], :real_a.shape[3]]
            loss_G_idt = criterion_idt(identity, real_a) * args.lambda_idt

            loss_G = loss_G_gan + loss_G_nce * args.lambda_nce + loss_G_idt

            opt_G.zero_grad()
            loss_G.backward()
            opt_G.step()

            G_loss_total += loss_G.item()
            D_loss_total += loss_D.item()
            n_batches += 1

        # ---- Epoch 统计 ----
        G_avg = G_loss_total / max(n_batches, 1)
        D_avg = D_loss_total / max(n_batches, 1)
        elapsed = time.time() - epoch_start

        print(f"Epoch {epoch:3d}/{args.epochs} | "
              f"G: {G_avg:.4f}  D: {D_avg:.4f} | "
              f"{elapsed:.1f}s")

        pass  # 日志仅print，不使用tensorboard

        # ---- 保存 ----
        if epoch % args.save_every == 0:
            ckpt = {
                'epoch': epoch,
                'G_state': G.state_dict(),
                'D_state': D.state_dict(),
                'opt_G': opt_G.state_dict(),
                'opt_D': opt_D.state_dict(),
            }
            torch.save(ckpt, os.path.join(args.output, f'epoch_{epoch}.pth'))

    # 最终保存
    final_ckpt = {
        'G_state': G.state_dict(),
        'D_state': D.state_dict(),
    }
    torch.save(final_ckpt, os.path.join(args.output, 'final.pth'))
    print(f"\n训练完成! 模型保存在: {args.output}")


def main():
    parser = argparse.ArgumentParser(description='CUT L通道训练')

    # 数据
    parser.add_argument('--domain-a', type=str, required=True,
                        help='合成车牌目录')
    parser.add_argument('--domain-b', type=str, required=True,
                        help='真实车牌目录')

    # 训练
    parser.add_argument('--epochs', type=int, default=100)
    parser.add_argument('--batch', type=int, default=4)
    parser.add_argument('--lr', type=float, default=2e-4)
    parser.add_argument('--workers', type=int, default=2)

    # 模型
    parser.add_argument('--n-res', type=int, default=6,
                        help='残差块数量')

    # 损失权重
    parser.add_argument('--lambda-nce', type=float, default=1.0)
    parser.add_argument('--lambda-idt', type=float, default=0.5)

    # PatchNCE
    parser.add_argument('--nce-layers', type=int, nargs='+', default=[0, 1, 2])
    parser.add_argument('--nce-t', type=float, default=0.07)
    parser.add_argument('--n-patches', type=int, default=128)

    # 保存
    parser.add_argument('--output', type=str, default='./output/cut_train')
    parser.add_argument('--img-w', type=int, default=220)
    parser.add_argument('--img-h', type=int, default=70)
    parser.add_argument('--save-every', type=int, default=20)

    args = parser.parse_args()
    train(args)


if __name__ == '__main__':
    main()