# -*- coding: utf-8 -*-
"""
CUT L通道 模型定义
==================
轻量级 Generator + PatchGAN Discriminator + PatchNCE 对比损失。
仅处理单通道 L (亮度) 图像，输入尺寸 140×440。
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


# ============================================================
# Generator: UNet 风格 encoder-decoder
# ============================================================

class ResBlock(nn.Module):
    """残差块"""
    def __init__(self, ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.ReflectionPad2d(1),
            nn.Conv2d(ch, ch, 3),
            nn.InstanceNorm2d(ch),
            nn.ReLU(True),
            nn.ReflectionPad2d(1),
            nn.Conv2d(ch, ch, 3),
            nn.InstanceNorm2d(ch),
        )

    def forward(self, x):
        return x + self.conv(x)


class LChannelGenerator(nn.Module):
    """
    L通道生成器: 1ch → 1ch
    输入: (B, 1, 140, 440) 合成车牌L通道
    输出: (B, 1, 140, 440) 真实风格L通道
    """
    def __init__(self, n_res=6):
        super().__init__()
        # Encoder
        self.enc1 = self._conv_block(1, 64, 7, 1, 3)       # 140×440
        self.enc2 = self._conv_block(64, 128, 3, 2, 1)       # 70×220
        self.enc3 = self._conv_block(128, 256, 3, 2, 1)      # 35×110

        # Bottleneck
        self.res_blocks = nn.Sequential(*[ResBlock(256) for _ in range(n_res)])

        # Decoder
        self.dec2 = self._upconv_block(256, 128)              # 70×220
        self.dec1 = self._upconv_block(128, 64)               # 140×440
        self.out = nn.Sequential(
            nn.ReflectionPad2d(3),
            nn.Conv2d(64, 1, 7),
            nn.Tanh(),
        )

    def _conv_block(self, ci, co, k, s, p):
        return nn.Sequential(
            nn.ReflectionPad2d(p),
            nn.Conv2d(ci, co, k, s),
            nn.InstanceNorm2d(co),
            nn.ReLU(True),
        )

    def _upconv_block(self, ci, co):
        return nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            nn.ReflectionPad2d(1),
            nn.Conv2d(ci, co, 3),
            nn.InstanceNorm2d(co),
            nn.ReLU(True),
        )

    def forward(self, x, return_features=False):
        """
        Args:
            x: (B, 1, H, W) 归一化到 [-1, 1]
            return_features: 返回中间特征用于 PatchNCE loss
        """
        features = []

        e1 = self.enc1(x)        # (B, 64, 140, 440)
        features.append(e1)

        e2 = self.enc2(e1)       # (B, 128, 70, 220)
        features.append(e2)

        e3 = self.enc3(e2)       # (B, 256, 35, 110)
        features.append(e3)

        b = self.res_blocks(e3)  # (B, 256, 35, 110)

        d2 = self.dec2(b)        # (B, 128, 70, 220)
        d1 = self.dec1(d2)       # (B, 64, 140, 440)

        out = self.out(d1)       # (B, 1, 140, 440)

        if return_features:
            return out, features
        return out


# ============================================================
# Discriminator: PatchGAN 70×70
# ============================================================

class PatchGANDiscriminator(nn.Module):
    """
    PatchGAN 判别器，70×70 感受野。
    输入: (B, 1, 140, 440) L通道
    输出: (B, 1, H/16, W/16) 真假概率图
    """
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(1, 64, 4, 2, 1),           # 70×220
            nn.LeakyReLU(0.2, True),
            nn.Conv2d(64, 128, 4, 2, 1),          # 35×110
            nn.InstanceNorm2d(128),
            nn.LeakyReLU(0.2, True),
            nn.Conv2d(128, 256, 4, 2, 1),          # 17×55
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2, True),
            nn.Conv2d(256, 512, 4, 1, 1),          # 16×54
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, True),
            nn.Conv2d(512, 1, 4, 1, 1),            # 15×53
        )

    def forward(self, x):
        return self.model(x)


# ============================================================
# PatchNCE 对比学习
# ============================================================

class PatchNCELoss(nn.Module):
    """
    PatchNCE 对比损失。
    对编码器的每一层特征图做 patch-wise 对比学习。
    """
    def __init__(self, nce_layers=(0, 1, 2), nce_t=0.07, n_patches=256):
        super().__init__()
        self.nce_layers = nce_layers
        self.nce_t = nce_t
        self.n_patches = n_patches
        self.cross_entropy = nn.CrossEntropyLoss()

    def forward(self, feat_q, feat_k):
        """
        Args:
            feat_q: 生成图编码特征列表 [(B,C,H,W), ...]
            feat_k: 原图编码特征列表 [(B,C,H,W), ...]
        Returns:
            loss: scalar
        """
        total_loss = 0.0
        n_layers = 0

        for layer_idx in self.nce_layers:
            if layer_idx >= len(feat_q) or layer_idx >= len(feat_k):
                continue

            fq = feat_q[layer_idx]  # (B, C, H, W)
            fk = feat_k[layer_idx]  # (B, C, H, W)

            B, C, H, W = fq.shape
            n_patches = min(self.n_patches, H * W)

            # 随机采样 patch 位置
            idx_h = torch.randint(0, H, (n_patches,), device=fq.device)
            idx_w = torch.randint(0, W, (n_patches,), device=fq.device)

            # query patches
            q = fq[:, :, idx_h, idx_w]  # (B, C, N)
            q = q.permute(0, 2, 1).reshape(B * n_patches, C)  # (B*N, C)

            # key patches (positive pairs)
            k = fk[:, :, idx_h, idx_w]  # (B, C, N)
            k = k.permute(0, 2, 1).reshape(B * n_patches, C)  # (B*N, C)

            # 负样本：同一 batch 内其他图片的同一位置
            q_norm = F.normalize(q, dim=1)  # (B*N, C)
            k_norm = F.normalize(k, dim=1)  # (B*N, C)

            # 计算相似度矩阵
            logits = torch.mm(q_norm, k_norm.T) / self.nce_t  # (B*N, B*N)
            labels = torch.arange(B * n_patches, device=q.device)

            loss = self.cross_entropy(logits, labels)
            total_loss += loss
            n_layers += 1

        return total_loss / max(n_layers, 1)


# ============================================================
# 模型工具
# ============================================================

def init_weights(model):
    """Xavier 初始化"""
    for m in model.modules():
        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
            nn.init.xavier_normal_(m.weight, 0.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)


def count_params(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)