"""
训练数据结构模块

实现2048游戏的训练数据结构，包括：
1. 棋盘状态的对数变换
2. 二面体群D4的8种变换（棋盘压缩）
3. 训练数据的内存缓存和硬盘持久化
4. 与PyTorch生态的集成
"""

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import pickle
import hashlib
import os
from typing import Tuple, List, Dict, Optional, Union
from dataclasses import dataclass
from pathlib import Path
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class TrainingExample:
    """单个训练样本的数据结构"""
    board_state: np.ndarray  # 棋盘状态 (H, W)
    action: int  # 动作 (0:上, 1:下, 2:左, 3:右)
    value: float  # 该状态-动作对的价值
    canonical_hash: str  # 规范化后的哈希值


class BoardTransform:
    """棋盘变换工具类，实现二面体群D4的8种变换"""
    
    @staticmethod
    def log_transform(board: np.ndarray) -> np.ndarray:
        """
        对棋盘进行对数变换
        
        Args:
            board: 原始棋盘状态，包含2的幂次数字
            
        Returns:
            对数变换后的棋盘，空位为0，其他位置为log2(value)
        """
        result = np.zeros_like(board, dtype=np.int32)
        mask = board > 0
        result[mask] = np.log2(board[mask]).astype(np.int32)
        return result
    
    @staticmethod
    def inverse_log_transform(log_board: np.ndarray) -> np.ndarray:
        """
        对数变换的逆变换
        
        Args:
            log_board: 对数变换后的棋盘
            
        Returns:
            原始棋盘状态
        """
        result = np.zeros_like(log_board, dtype=np.int32)
        mask = log_board > 0
        result[mask] = (2 ** log_board[mask]).astype(np.int32)
        return result
    
    @staticmethod
    def rotate_90(matrix: np.ndarray) -> np.ndarray:
        """顺时针旋转90度"""
        return np.rot90(matrix, k=-1)
    
    @staticmethod
    def flip_horizontal(matrix: np.ndarray) -> np.ndarray:
        """水平翻转"""
        return np.fliplr(matrix)
    
    @classmethod
    def get_all_transforms(cls, matrix: np.ndarray) -> List[np.ndarray]:
        """
        获取二面体群D4的所有8种变换
        
        Args:
            matrix: 输入矩阵
            
        Returns:
            包含8种变换结果的列表
        """
        transforms = []
        
        # 原始图像 (R0)
        r0 = matrix.copy()
        transforms.append(r0)
        
        # 旋转90° (R90)
        r90 = cls.rotate_90(r0)
        transforms.append(r90)
        
        # 旋转180° (R180)
        r180 = cls.rotate_90(r90)
        transforms.append(r180)
        
        # 旋转270° (R270)
        r270 = cls.rotate_90(r180)
        transforms.append(r270)
        
        # 水平翻转 (F)
        f = cls.flip_horizontal(r0)
        transforms.append(f)
        
        # 翻转后旋转90° (F+R90)
        fr90 = cls.rotate_90(f)
        transforms.append(fr90)
        
        # 翻转后旋转180° (F+R180)
        fr180 = cls.rotate_90(fr90)
        transforms.append(fr180)
        
        # 翻转后旋转270° (F+R270)
        fr270 = cls.rotate_90(fr180)
        transforms.append(fr270)
        
        return transforms
    
    @classmethod
    def get_canonical_form(cls, matrix: np.ndarray) -> Tuple[np.ndarray, int]:
        """
        获取矩阵的规范形式（字典序最小的变换）

        Args:
            matrix: 输入矩阵

        Returns:
            (规范形式矩阵, 变换索引)
        """
        transforms = cls.get_all_transforms(matrix)

        # 将每个变换拉平为1D向量并比较字典序
        flattened = [t.flatten() for t in transforms]

        # 找到字典序最小的索引
        min_idx = 0
        min_flat = flattened[0]

        for i, flat in enumerate(flattened[1:], 1):
            # 逐元素比较字典序
            if cls._is_lexicographically_smaller(flat, min_flat):
                min_idx = i
                min_flat = flat

        return transforms[min_idx], min_idx

    @staticmethod
    def _is_lexicographically_smaller(a: np.ndarray, b: np.ndarray) -> bool:
        """
        检查数组a是否在字典序上小于数组b

        Args:
            a, b: 要比较的数组

        Returns:
            如果a < b则返回True
        """
        for i in range(min(len(a), len(b))):
            if a[i] < b[i]:
                return True
            elif a[i] > b[i]:
                return False
        # 如果前面都相等，较短的数组更小
        return len(a) < len(b)
    
    @classmethod
    def compute_hash(cls, matrix: np.ndarray) -> str:
        """
        计算矩阵规范形式的哈希值
        
        Args:
            matrix: 输入矩阵
            
        Returns:
            哈希字符串
        """
        canonical, _ = cls.get_canonical_form(matrix)
        # 使用规范形式的字节表示计算哈希
        return hashlib.md5(canonical.tobytes()).hexdigest()


class ScoreCalculator:
    """分数计算工具类"""
    
    @staticmethod
    def calculate_tile_value(tile_log: int) -> int:
        """
        计算单个瓦片的累积分数价值
        
        Args:
            tile_log: 瓦片的对数值 (log2(tile_value))
            
        Returns:
            累积分数价值
        """
        if tile_log <= 1:  # 对应原始值2或更小
            return 0
        
        # V(N) = (log2(N) - 1) * N，其中N = 2^tile_log
        n = 2 ** tile_log
        return (tile_log - 1) * n
    
    @classmethod
    def calculate_board_score(cls, log_board: np.ndarray) -> int:
        """
        计算整个棋盘的累积分数
        
        Args:
            log_board: 对数变换后的棋盘
            
        Returns:
            总累积分数
        """
        total_score = 0
        for tile_log in log_board.flatten():
            if tile_log > 0:
                total_score += cls.calculate_tile_value(tile_log)
        return total_score


class TrainingDataCache:
    """训练数据的内存缓存系统"""
    
    def __init__(self, max_size: int = 1000000):
        """
        初始化缓存
        
        Args:
            max_size: 最大缓存条目数
        """
        self.max_size = max_size
        self.cache: Dict[str, TrainingExample] = {}
        self.access_order: List[str] = []  # 用于LRU淘汰
        
    def get(self, key: str) -> Optional[TrainingExample]:
        """获取缓存项"""
        if key in self.cache:
            # 更新访问顺序
            self.access_order.remove(key)
            self.access_order.append(key)
            return self.cache[key]
        return None
    
    def put(self, key: str, example: TrainingExample) -> None:
        """添加或更新缓存项"""
        if key in self.cache:
            # 更新现有项
            self.cache[key] = example
            self.access_order.remove(key)
            self.access_order.append(key)
        else:
            # 添加新项
            if len(self.cache) >= self.max_size:
                # LRU淘汰
                oldest_key = self.access_order.pop(0)
                del self.cache[oldest_key]
            
            self.cache[key] = example
            self.access_order.append(key)
    
    def update_if_better(self, key: str, example: TrainingExample) -> bool:
        """如果新样本的价值更高，则更新缓存"""
        existing = self.get(key)
        if existing is None or example.value > existing.value:
            self.put(key, example)
            return True
        return False
    
    def size(self) -> int:
        """返回缓存大小"""
        return len(self.cache)
    
    def clear(self) -> None:
        """清空缓存"""
        self.cache.clear()
        self.access_order.clear()
    
    def get_all_examples(self) -> List[TrainingExample]:
        """获取所有缓存的训练样本"""
        return list(self.cache.values())


class TrainingDataPersistence:
    """训练数据的硬盘持久化系统"""
    
    def __init__(self, data_dir: str = "data/training"):
        """
        初始化持久化系统

        Args:
            data_dir: 数据存储目录
        """
        self.data_dir = Path(data_dir)
        self.data_dir.mkdir(parents=True, exist_ok=True)
        
    def save_cache(self, cache: TrainingDataCache, filename: str) -> None:
        """
        保存缓存到硬盘
        
        Args:
            cache: 要保存的缓存
            filename: 文件名
        """
        filepath = self.data_dir / f"{filename}.pkl"
        examples = cache.get_all_examples()
        
        with open(filepath, 'wb') as f:
            pickle.dump(examples, f)
        
        logger.info(f"Saved {len(examples)} examples to {filepath}")
    
    def load_cache(self, filename: str) -> List[TrainingExample]:
        """
        从硬盘加载训练数据
        
        Args:
            filename: 文件名
            
        Returns:
            训练样本列表
        """
        filepath = self.data_dir / f"{filename}.pkl"
        
        if not filepath.exists():
            logger.warning(f"File {filepath} does not exist")
            return []
        
        with open(filepath, 'rb') as f:
            examples = pickle.load(f)
        
        logger.info(f"Loaded {len(examples)} examples from {filepath}")
        return examples
    
    def save_examples_batch(self, examples: List[TrainingExample], 
                           batch_name: str) -> None:
        """
        批量保存训练样本
        
        Args:
            examples: 训练样本列表
            batch_name: 批次名称
        """
        self.save_cache(TrainingDataCache(), batch_name)
        # 直接保存examples列表
        filepath = self.data_dir / f"{batch_name}.pkl"
        with open(filepath, 'wb') as f:
            pickle.dump(examples, f)
        
        logger.info(f"Saved batch {batch_name} with {len(examples)} examples")
    
    def list_saved_files(self) -> List[str]:
        """列出所有保存的数据文件"""
        return [f.stem for f in self.data_dir.glob("*.pkl")]


class Game2048Dataset(Dataset):
    """PyTorch Dataset for 2048 training data"""

    def __init__(self, examples: List[TrainingExample],
                 board_size: Tuple[int, int] = (4, 4),
                 max_tile_value: int = 17):
        """
        初始化数据集

        Args:
            examples: 训练样本列表
            board_size: 棋盘大小 (height, width)
            max_tile_value: 最大瓦片值的对数 (log2)
        """
        self.examples = examples
        self.board_size = board_size
        self.max_tile_value = max_tile_value

    def __len__(self) -> int:
        return len(self.examples)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        获取单个训练样本

        Args:
            idx: 样本索引

        Returns:
            (board_tensor, action_tensor, value_tensor)
        """
        example = self.examples[idx]

        # 将棋盘状态转换为one-hot编码
        board_tensor = self._encode_board(example.board_state)

        # 动作标签
        action_tensor = torch.tensor(example.action, dtype=torch.long)

        # 价值标签
        value_tensor = torch.tensor(example.value, dtype=torch.float32)

        return board_tensor, action_tensor, value_tensor

    def _encode_board(self, board: np.ndarray) -> torch.Tensor:
        """
        将棋盘状态编码为one-hot张量

        Args:
            board: 对数变换后的棋盘状态

        Returns:
            形状为 (max_tile_value + 1, height, width) 的张量
        """
        # 创建one-hot编码
        # 通道0: 空位 (值为0)
        # 通道1: 值为1 (对应原始值2)
        # 通道2: 值为2 (对应原始值4)
        # ...
        channels = self.max_tile_value + 1
        height, width = self.board_size

        encoded = torch.zeros(channels, height, width, dtype=torch.float32)

        for i in range(height):
            for j in range(width):
                tile_value = int(board[i, j])
                if 0 <= tile_value <= self.max_tile_value:
                    encoded[tile_value, i, j] = 1.0

        return encoded


class TrainingDataManager:
    """训练数据管理器，整合缓存、持久化和PyTorch集成"""

    def __init__(self, data_dir: str = "data/training",
                 cache_size: int = 1000000,
                 board_size: Tuple[int, int] = (4, 4)):
        """
        初始化数据管理器

        Args:
            data_dir: 数据存储目录
            cache_size: 内存缓存大小
            board_size: 棋盘大小
        """
        self.cache = TrainingDataCache(cache_size)
        self.persistence = TrainingDataPersistence(data_dir)
        self.board_size = board_size
        self.transform = BoardTransform()
        self.score_calc = ScoreCalculator()

    def add_training_example(self, board_state: np.ndarray,
                           action: int, value: float) -> str:
        """
        添加训练样本

        Args:
            board_state: 原始棋盘状态
            action: 动作
            value: 价值

        Returns:
            样本的哈希键
        """
        # 对数变换
        log_board = self.transform.log_transform(board_state)

        # 计算规范哈希
        canonical_hash = self.transform.compute_hash(log_board)

        # 创建训练样本
        example = TrainingExample(
            board_state=log_board,
            action=action,
            value=value,
            canonical_hash=canonical_hash
        )

        # 构造缓存键 (状态哈希 + 动作)
        cache_key = f"{canonical_hash}_{action}"

        # 更新缓存（如果新价值更高）
        self.cache.update_if_better(cache_key, example)

        return cache_key

    def get_pytorch_dataset(self, filter_min_value: float = 0.0) -> Game2048Dataset:
        """
        获取PyTorch数据集

        Args:
            filter_min_value: 最小价值过滤阈值

        Returns:
            PyTorch数据集
        """
        examples = [ex for ex in self.cache.get_all_examples()
                   if ex.value >= filter_min_value]

        return Game2048Dataset(examples, self.board_size)

    def get_dataloader(self, batch_size: int = 32, shuffle: bool = True,
                      filter_min_value: float = 0.0, **kwargs) -> DataLoader:
        """
        获取PyTorch DataLoader

        Args:
            batch_size: 批次大小
            shuffle: 是否打乱数据
            filter_min_value: 最小价值过滤阈值
            **kwargs: 其他DataLoader参数

        Returns:
            PyTorch DataLoader
        """
        dataset = self.get_pytorch_dataset(filter_min_value)
        return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, **kwargs)

    def save_current_cache(self, filename: str) -> None:
        """保存当前缓存到硬盘"""
        self.persistence.save_cache(self.cache, filename)

    def load_from_file(self, filename: str) -> int:
        """
        从文件加载训练数据到缓存

        Args:
            filename: 文件名

        Returns:
            加载的样本数量
        """
        examples = self.persistence.load_cache(filename)

        loaded_count = 0
        for example in examples:
            cache_key = f"{example.canonical_hash}_{example.action}"
            if self.cache.update_if_better(cache_key, example):
                loaded_count += 1

        logger.info(f"Loaded {loaded_count} examples into cache")
        return loaded_count

    def get_cache_stats(self) -> Dict[str, int]:
        """获取缓存统计信息"""
        return {
            "cache_size": self.cache.size(),
            "max_cache_size": self.cache.max_size,
            "saved_files": len(self.persistence.list_saved_files())
        }

    def merge_caches(self, other_manager: 'TrainingDataManager') -> int:
        """
        合并另一个数据管理器的缓存

        Args:
            other_manager: 另一个数据管理器

        Returns:
            合并的样本数量
        """
        merged_count = 0
        for example in other_manager.cache.get_all_examples():
            cache_key = f"{example.canonical_hash}_{example.action}"
            if self.cache.update_if_better(cache_key, example):
                merged_count += 1

        logger.info(f"Merged {merged_count} examples from other cache")
        return merged_count