"""GTAV 场景识别 - 基于 ResNet 的楼层分类

使用 timm 库和预训练 ResNet 进行楼层分类 (floor 1/2/3)
"""

import random
import numpy as np
import sys
import time

from pathlib import Path
from typing import Optional, Union
from dataclasses import dataclass, field
from tclogger import TCLogger, add_fills
from PIL import Image

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

import timm
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

logger = TCLogger("FloorRecognizerV2", use_prefix=True, use_prefix_ms=True)

MODULE_DIR = Path(__file__).parent
IMAGES_DIR = MODULE_DIR / "images"
CACHE_DIR = MODULE_DIR / "cache"


def log_header(header: str):
    header_str = add_fills(header, "=", total_width=80)
    logger.hint(f"{header_str}", use_prefix=False)


@dataclass
class MatchResult:
    """匹配结果"""

    floor: int  # 匹配到的楼层 (1, 2, 3)
    confidence: float  # 置信度 (0.0 - 1.0)
    score: float  # 匹配分数
    details: dict = field(default_factory=dict)


class FloorDataset(Dataset):
    """楼层数据集"""

    def __init__(self, samples: list[tuple[Path, int]], transform=None):
        """
        Args:
            samples: [(image_path, floor_label), ...] 其中 floor_label 是 0/1/2 (对应楼层 1/2/3)
            transform: 图像变换
        """
        self.samples = samples
        self.transform = transform

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

    def __getitem__(self, idx):
        path, label = self.samples[idx]
        image = Image.open(path).convert("RGB")
        if self.transform:
            image = self.transform(image)
        return image, label


class FloorClassifier(nn.Module):
    """基于 ResNet 的楼层分类器"""

    def __init__(self, num_classes: int = 3, model_name: str = "resnet18"):
        super().__init__()
        # 使用 timm 加载预训练模型
        self.backbone = timm.create_model(model_name, pretrained=True, num_classes=0)
        # 获取特征维度
        self.feature_dim = self.backbone.num_features
        # 分类头
        self.classifier = nn.Sequential(
            nn.Dropout(0.3),
            nn.Linear(self.feature_dim, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, num_classes),
        )

    def forward(self, x):
        features = self.backbone(x)
        return self.classifier(features)

    def extract_features(self, x):
        """提取特征（用于 KNN 等）"""
        return self.backbone(x)


class FloorRecognizerV2:
    """基于 ResNet 的楼层识别器"""

    def __init__(
        self,
        floor_count: int = 3,
        images_dir: Optional[Path] = None,
        cache_dir: Optional[Path] = None,
        model_name: str = "resnet18",
        device: Optional[str] = None,
    ):
        self.floor_count = floor_count
        self.images_dir = images_dir or IMAGES_DIR
        self.cache_dir = cache_dir or CACHE_DIR
        self.model_path = self.cache_dir / "floor_classifier_resnet.pth"
        self.model_name = model_name

        # 设备
        if device is None:
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        else:
            self.device = torch.device(device)
        logger.note(f"使用设备: {self.device}")

        # 模型
        self.model = FloorClassifier(num_classes=floor_count, model_name=model_name)
        self.model.to(self.device)

        # 数据变换（使用 timm 的默认配置）
        self.data_config = resolve_data_config({}, model=self.model.backbone)
        self.transform = create_transform(**self.data_config)

        # 推理用的简单变换（不做数据增强）
        self.eval_transform = create_transform(**self.data_config, is_training=False)

    def _collect_samples(self) -> list[tuple[Path, int]]:
        """收集所有训练样本"""
        samples = []
        locations_dir = self.cache_dir / "locations"

        for floor in range(1, self.floor_count + 1):
            label = floor - 1  # 转为 0-indexed
            for d in locations_dir.glob(f"*floor_{floor}_minimap"):
                if not d.is_dir():
                    continue
                for f in d.glob("*.jpg"):
                    samples.append((f, label))

        return samples

    def train(
        self,
        epochs: int = 20,
        batch_size: int = 16,
        lr: float = 1e-4,
        save: bool = True,
    ):
        """训练分类器"""
        samples = self._collect_samples()
        if not samples:
            logger.warn("没有找到训练样本")
            return

        # 统计各类别样本数
        class_counts = [0] * self.floor_count
        for _, label in samples:
            class_counts[label] += 1
        for floor in range(1, self.floor_count + 1):
            logger.note(f"楼层 {floor} 样本数: {class_counts[floor-1]}")

        # 计算类别权重（处理样本不平衡）
        total = sum(class_counts)
        class_weights = [total / (c + 1) for c in class_counts]
        class_weights = torch.tensor(class_weights, dtype=torch.float32).to(self.device)
        class_weights = class_weights / class_weights.sum() * self.floor_count

        # 创建数据增强变换
        train_transform = create_transform(
            **self.data_config,
            is_training=True,
            auto_augment="rand-m9-mstd0.5",
        )

        dataset = FloorDataset(samples, transform=train_transform)
        dataloader = DataLoader(
            dataset, batch_size=batch_size, shuffle=True, num_workers=0
        )

        # 冻结 backbone 的前几层
        for name, param in self.model.backbone.named_parameters():
            if "layer4" not in name and "fc" not in name:
                param.requires_grad = False

        # 优化器
        optimizer = torch.optim.AdamW(
            filter(lambda p: p.requires_grad, self.model.parameters()),
            lr=lr,
            weight_decay=0.01,
        )

        # 学习率调度器
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)

        # 损失函数
        criterion = nn.CrossEntropyLoss(weight=class_weights)

        # 训练循环
        self.model.train()
        for epoch in range(epochs):
            total_loss = 0.0
            correct = 0
            total = 0

            for images, labels in dataloader:
                images = images.to(self.device)
                labels = labels.to(self.device)

                optimizer.zero_grad()
                outputs = self.model(images)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()

                total_loss += loss.item()
                _, predicted = outputs.max(1)
                correct += predicted.eq(labels).sum().item()
                total += labels.size(0)

            scheduler.step()
            acc = correct / total if total > 0 else 0
            logger.mesg(
                f"Epoch {epoch+1}/{epochs}: Loss={total_loss:.4f}, Acc={acc:.1%}"
            )

        if save:
            self._save_model()

    def _save_model(self):
        """保存模型"""
        self.model_path.parent.mkdir(parents=True, exist_ok=True)
        torch.save(
            {
                "model_state_dict": self.model.state_dict(),
                "model_name": self.model_name,
                "floor_count": self.floor_count,
            },
            self.model_path,
        )
        logger.okay(f"模型已保存: {self.model_path}")

    def load_model(self) -> bool:
        """加载模型"""
        if not self.model_path.exists():
            return False
        try:
            checkpoint = torch.load(self.model_path, map_location=self.device)
            self.model.load_state_dict(checkpoint["model_state_dict"])
            self.model.eval()
            logger.okay("模型已加载")
            return True
        except Exception as e:
            logger.warn(f"加载失败: {e}")
            return False

    def recognize(self, image: Union[np.ndarray, Image.Image]) -> MatchResult:
        """识别楼层"""
        self.model.eval()

        # 转换图像
        if isinstance(image, np.ndarray):
            # BGR to RGB
            if len(image.shape) == 3 and image.shape[2] == 3:
                image = image[:, :, ::-1]
            image = Image.fromarray(image)

        # 应用变换
        image_tensor = self.eval_transform(image).unsqueeze(0).to(self.device)

        # 推理
        with torch.no_grad():
            outputs = self.model(image_tensor)
            probs = F.softmax(outputs, dim=1).squeeze()

        # 解析结果
        probs_np = probs.cpu().numpy()
        predicted_label = probs_np.argmax()
        predicted_floor = predicted_label + 1  # 转为 1-indexed
        confidence = probs_np[predicted_label]

        # 构建详情
        details = {f"floor_{i+1}": float(probs_np[i]) for i in range(self.floor_count)}

        return MatchResult(
            floor=int(predicted_floor),
            confidence=float(confidence),
            score=float(confidence),
            details=details,
        )

    def recognize_file(self, path: Union[str, Path]) -> Optional[MatchResult]:
        """从文件识别"""
        try:
            image = Image.open(path).convert("RGB")
            return self.recognize(image)
        except Exception as e:
            logger.warn(f"读取图像失败: {e}")
            return None

    def evaluate(self, test_dir: Path, expected_floor: int) -> dict:
        """评估准确率"""
        files = list(Path(test_dir).glob("*.jpg"))
        if not files:
            return {"accuracy": 0, "total": 0, "correct": 0}

        correct = sum(
            1
            for f in files
            if (r := self.recognize_file(f)) and r.floor == expected_floor
        )
        return {
            "expected_floor": expected_floor,
            "total": len(files),
            "correct": correct,
            "accuracy": correct / len(files),
        }

    def __repr__(self) -> str:
        return f"FloorRecognizerV2(model={self.model_name}, device={self.device})"


# ============================================================
# 测试
# ============================================================


def test_train():
    """训练模型"""
    log_header("训练模型")

    recognizer = FloorRecognizerV2()
    recognizer.train(epochs=20, batch_size=16, lr=1e-4)


def test_accuracy():
    """测试准确率"""
    log_header("准确率测试")

    recognizer = FloorRecognizerV2()
    if not recognizer.load_model():
        logger.warn("未找到模型，先训练...")
        recognizer.train(epochs=20)

    logger.note(f"{recognizer}")

    locations_dir = CACHE_DIR / "locations"
    results = []

    for d in sorted(locations_dir.glob("*_minimap")):
        for floor in range(1, 4):
            if f"floor_{floor}" in d.name:
                r = recognizer.evaluate(d, floor)
                results.append(r)
                logger.mesg(
                    f"{d.name}: {r['accuracy']:.1%} ({r['correct']}/{r['total']})"
                )
                break

    if results:
        total = sum(r["total"] for r in results)
        correct = sum(r["correct"] for r in results)
        logger.okay(f"总体: {correct/total:.1%} ({correct}/{total})")


def stratified_k_fold(
    samples: list[tuple[Path, int]], n_splits: int = 5, seed: int = 42
):
    """手动实现分层 K-Fold 交叉验证

    Args:
        samples: [(path, label), ...]
        n_splits: 折数
        seed: 随机种子

    Yields:
        (train_indices, val_indices) for each fold
    """
    random.seed(seed)

    # 按类别分组
    class_indices: dict[int, list[int]] = {}
    for i, (_, label) in enumerate(samples):
        if label not in class_indices:
            class_indices[label] = []
        class_indices[label].append(i)

    # 打乱每个类别的索引
    for indices in class_indices.values():
        random.shuffle(indices)

    # 为每个 fold 分配样本
    folds: list[list[int]] = [[] for _ in range(n_splits)]
    for label, indices in class_indices.items():
        for i, idx in enumerate(indices):
            folds[i % n_splits].append(idx)

    # 生成 train/val 划分
    for fold_idx in range(n_splits):
        val_indices = folds[fold_idx]
        train_indices = []
        for i, fold in enumerate(folds):
            if i != fold_idx:
                train_indices.extend(fold)
        yield train_indices, val_indices


def test_cross_validation():
    """K-Fold 交叉验证"""
    log_header("K-Fold 交叉验证")

    recognizer = FloorRecognizerV2()
    samples = recognizer._collect_samples()

    if not samples:
        logger.warn("没有找到样本")
        return

    # 5-Fold 交叉验证
    n_splits = 5
    fold_results = []

    for fold, (train_idx, val_idx) in enumerate(
        stratified_k_fold(samples, n_splits=n_splits)
    ):
        logger.note(f"Fold {fold + 1}/5")

        # 重新初始化模型
        recognizer = FloorRecognizerV2()

        # 准备数据
        train_samples = [samples[i] for i in train_idx]
        val_samples = [samples[i] for i in val_idx]

        # 训练（只用训练集）
        train_transform = create_transform(
            **recognizer.data_config,
            is_training=True,
            auto_augment="rand-m9-mstd0.5",
        )
        dataset = FloorDataset(train_samples, transform=train_transform)
        dataloader = DataLoader(dataset, batch_size=16, shuffle=True, num_workers=0)

        # 类别权重
        class_counts = [0] * recognizer.floor_count
        for _, label in train_samples:
            class_counts[label] += 1
        class_weights = [len(train_samples) / (c + 1) for c in class_counts]
        class_weights = torch.tensor(class_weights, dtype=torch.float32).to(
            recognizer.device
        )
        class_weights = class_weights / class_weights.sum() * recognizer.floor_count

        # 冻结部分层
        for name, param in recognizer.model.backbone.named_parameters():
            if "layer4" not in name and "fc" not in name:
                param.requires_grad = False

        optimizer = torch.optim.AdamW(
            filter(lambda p: p.requires_grad, recognizer.model.parameters()),
            lr=1e-4,
            weight_decay=0.01,
        )
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=15)
        criterion = nn.CrossEntropyLoss(weight=class_weights)

        # 训练 15 epochs
        recognizer.model.train()
        for epoch in range(15):
            for images, labels_batch in dataloader:
                images = images.to(recognizer.device)
                labels_batch = labels_batch.to(recognizer.device)
                optimizer.zero_grad()
                outputs = recognizer.model(images)
                loss = criterion(outputs, labels_batch)
                loss.backward()
                optimizer.step()
            scheduler.step()

        # 验证
        recognizer.model.eval()
        correct = 0
        for path, label in val_samples:
            result = recognizer.recognize_file(path)
            if result and result.floor == label + 1:
                correct += 1

        acc = correct / len(val_samples)
        fold_results.append(acc)
        logger.mesg(f"Fold {fold + 1} Accuracy: {acc:.1%}")

    avg_acc = sum(fold_results) / len(fold_results)
    std_acc = (sum((x - avg_acc) ** 2 for x in fold_results) / len(fold_results)) ** 0.5
    logger.okay(f"平均准确率: {avg_acc:.1%} ± {std_acc:.1%}")


def test_speed():
    """测试速度"""
    log_header("速度测试")

    recognizer = FloorRecognizerV2()
    if not recognizer.load_model():
        logger.warn("未找到模型，先训练...")
        recognizer.train(epochs=10)

    # 收集测试图像
    test_files = []
    for d in (CACHE_DIR / "locations").glob("*_minimap"):
        test_files.extend(list(d.glob("*.jpg"))[:5])

    if not test_files:
        logger.warn("没有测试图像")
        return

    # 预热
    recognizer.recognize_file(test_files[0])

    start = time.time()
    for f in test_files:
        recognizer.recognize_file(f)
    elapsed = time.time() - start

    logger.okay(f"平均: {elapsed/len(test_files)*1000:.1f} ms/张 ({len(test_files)} 张)")


if __name__ == "__main__":

    if len(sys.argv) > 1:
        cmd = sys.argv[1]
        if cmd == "train":
            test_train()
        elif cmd == "test":
            test_accuracy()
        elif cmd == "cv":
            test_cross_validation()
        elif cmd == "speed":
            test_speed()
    else:
        test_train()
        test_accuracy()
        test_speed()

    # python -m gtaz.recognizes_v2           # 训练+测试+速度
    # python -m gtaz.recognizes_v2 train     # 仅训练
    # python -m gtaz.recognizes_v2 test      # 仅测试
    # python -m gtaz.recognizes_v2 cv        # 交叉验证
    # python -m gtaz.recognizes_v2 speed     # 速度测试
