# -*- coding: utf-8 -*-
"""
L通道车牌数据集加载器
=====================
Domain A: 合成车牌 (140×440) → 提取 L 通道
Domain B: 真实车牌 (140×440) → 提取 L 通道
"""

import os
import random
from pathlib import Path
from typing import Tuple

import cv2
import numpy as np
import torch
from torch.utils.data import Dataset


class LChannelPlateDataset(Dataset):
    """
    车牌 L 通道数据集。

    Args:
        root:     图片根目录
        img_size: 统一输出尺寸 (W, H)
        is_real:  True=真实车牌域, False=合成车牌域
    """

    def __init__(self, root: str, img_size=(440, 140), is_real=False):
        self.root = Path(root)
        self.img_size = img_size
        self.is_real = is_real

        self.paths = []
        if self.root.is_dir():
            for ext in ('*.jpg', '*.jpeg', '*.png'):
                self.paths.extend(self.root.glob(ext))

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, idx) -> torch.Tensor:
        path = self.paths[idx]
        img = cv2.imread(str(path))
        if img is None:
            return self.__getitem__((idx + 1) % len(self))

        # 统一尺寸
        img = cv2.resize(img, self.img_size, interpolation=cv2.INTER_LINEAR)

        # 提取 L 通道: BGR → LAB → L
        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        l_channel = lab[:, :, 0].astype(np.float32)  # (H, W)  [0, 255]

        # 归一化到 [-1, 1]
        l_channel = (l_channel / 127.5) - 1.0

        return torch.from_numpy(l_channel).unsqueeze(0)  # (1, H, W)


def make_loader(root: str, batch_size=4, img_size=(440, 140),
                is_real=False, shuffle=True, num_workers=2):
    """创建 DataLoader"""
    from torch.utils.data import DataLoader

    ds = LChannelPlateDataset(root, img_size, is_real)
    return DataLoader(ds, batch_size=batch_size, shuffle=shuffle,
                      num_workers=num_workers, pin_memory=True,
                      drop_last=True)


class LChannelInference:
    """
    L通道 CUT 推理封装：合成车牌 → 真实风格车牌。

    用法:
        infer = LChannelInference('checkpoint.pth')
        result_bgr = infer(synthetic_plate_bgr)
    """

    def __init__(self, checkpoint_path: str, device=None, n_res=6):
        from models import LChannelGenerator, init_weights

        self.device = device or torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )

        self.G = LChannelGenerator(n_res=n_res).to(self.device)
        ckpt = torch.load(checkpoint_path, map_location=self.device)
        self.G.load_state_dict(ckpt.get('G_state', ckpt))
        self.G.eval()

        # 期望输入尺寸
        self.img_size = (440, 140)

    def __call__(self, plate_bgr: np.ndarray) -> np.ndarray:
        """
        Args:
            plate_bgr: (H, W, 3) 合成车牌 BGR

        Returns:
            真实风格车牌 BGR (H, W, 3)
        """
        h, w = plate_bgr.shape[:2]
        img = cv2.resize(plate_bgr, self.img_size, interpolation=cv2.INTER_LINEAR)

        # BGR → LAB
        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB).astype(np.float32)

        # L → tensor → G → tensor → L
        l_tensor = torch.from_numpy(lab[:, :, 0] / 127.5 - 1.0)
        l_tensor = l_tensor.unsqueeze(0).unsqueeze(0).to(self.device)  # (1,1,H,W)

        with torch.no_grad():
            l_out = self.G(l_tensor).squeeze().cpu().numpy()  # (H, W)

        # 还原 L 到 [0, 255]
        l_out = (l_out + 1.0) * 127.5
        l_out = np.clip(l_out, 0, 255)

        lab[:, :, 0] = l_out
        # a, b 通道不变 → 保护颜色

        result = cv2.cvtColor(lab.astype(np.uint8), cv2.COLOR_LAB2BGR)

        # 还原原始尺寸
        if (h, w) != (self.img_size[1], self.img_size[0]):
            result = cv2.resize(result, (w, h), interpolation=cv2.INTER_LINEAR)

        return result