import tensorflow as tf
import numpy as np
import os
import tarfile
import pickle
from typing import Dict, Any

def load_cifar10_from_local(local_path):
    """从本地cifar-10-python.tar.gz文件加载CIFAR-10数据集"""
    """Load the CIFAR-10 dataset from the local cifar-10-python.tar.gz file."""
    
    if not os.path.exists(local_path):
        raise FileNotFoundError(f"未找到CIFAR-10数据集文件: {local_path}")
    
    # 解压目录
    # Extraction directory
    extract_dir = 'cifar-10-batches-py'
    
    # 如果解压目录不存在，则解压文件
    # If the extraction directory does not exist, extract the file.
    if not os.path.exists(extract_dir):
        print("解压CIFAR-10数据集...Extract CIFAR-10")
        with tarfile.open(local_path, 'r:gz') as tar:
            tar.extractall()
        print("解压完成!Extraction completed")
    else:
        print("CIFAR-10数据集已解压，跳过解压步骤The CIFAR-10 dataset has been extracted, skipping the extraction step.")
    
    return extract_dir

def unpickle(file):
    """读取CIFAR-10的pickle文件"""
    """Read the CIFAR-10 pickle files"""
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict

def create_cifar10_tfrecord(local_path='cifar-10-python.tar.gz'):
    """从本地文件创建CIFAR-10数据集的TFRecord文件"""
    """Create TFRecord files for the CIFAR-10 dataset from local files."""
    
    # 从本地文件加载数据集
    # Load the dataset from local files.
    cifar_dir = load_cifar10_from_local(local_path)
    
    # 加载训练数据
    # Load training data
    x_train = []
    y_train = []
    
    # 加载5个训练批次
    # Load 5 train batch
    for i in range(1, 6):
        batch_file = os.path.join(cifar_dir, f'data_batch_{i}')
        dict = unpickle(batch_file)
        
        # 获取图像数据和标签
        # Get data and label
        data = dict[b'data']
        labels = dict[b'labels']
        
        x_train.append(data)
        y_train.extend(labels)
    
    x_train = np.concatenate(x_train, axis=0)
    y_train = np.array(y_train)
    
    # 加载测试数据
    # Load test data
    test_file = os.path.join(cifar_dir, 'test_batch')
    dict_test = unpickle(test_file)
    x_test = dict_test[b'data']
    y_test = np.array(dict_test[b'labels'])
    
    # 重塑图像形状为32x32x3
    # Reshape picture to 32x32x3
    x_train = x_train.reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
    x_test = x_test.reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
    
    # 创建输出目录
    # Create output directory
    os.makedirs('./tfrecords', exist_ok=True)
    
    def _bytes_feature(value):
        """转换为字节特征"""
        """Convert to byte feature"""
        if isinstance(value, type(tf.constant(0))):
            value = value.numpy()
        return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

    def _int64_feature(value):
        """转换为整数特征"""
        """Convert to Int feature"""
        return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
    
    # 写入训练集TFRecord
    # Write trainset TFRecord
    with tf.io.TFRecordWriter('./tfrecords/cifar10_train.tfrecord') as writer:
        for i in range(len(x_train)):
            image = x_train[i]
            label = y_train[i]
            
            # 确保图像数据是uint8类型
            # Make sure the image data type is uint8.
            image = image.astype(np.uint8)
            
            feature = {
                'image': _bytes_feature(tf.compat.as_bytes(image.tobytes())),
                'label': _int64_feature(label),
                'height': _int64_feature(32),
                'width': _int64_feature(32),
                'depth': _int64_feature(3)
            }
            example = tf.train.Example(features=tf.train.Features(feature=feature))
            writer.write(example.SerializeToString())
    
    # 写入验证集TFRecord
    # Write valset TFRecord
    with tf.io.TFRecordWriter('./tfrecords/cifar10_val.tfrecord') as writer:
        for i in range(len(x_test)):
            image = x_test[i]
            label = y_test[i]
            
            image = image.astype(np.uint8)
            
            feature = {
                'image': _bytes_feature(tf.compat.as_bytes(image.tobytes())),
                'label': _int64_feature(label),
                'height': _int64_feature(32),
                'width': _int64_feature(32),
                'depth': _int64_feature(3)
            }
            example = tf.train.Example(features=tf.train.Features(feature=feature))
            writer.write(example.SerializeToString())
    
    print(f"TFRecord文件创建完成: 训练样本 {len(x_train)}, 验证样本 {len(x_test)}")
    return len(x_train), len(x_test)

def parse_tfrecord_fn(example_proto, config: Dict[str, Any], is_training: bool = True):
    """解析TFRecord示例"""
    """Parse TFRecord examples"""
    feature_description = {
        'image': tf.io.FixedLenFeature([], tf.string),
        'label': tf.io.FixedLenFeature([], tf.int64),
        'height': tf.io.FixedLenFeature([], tf.int64),
        'width': tf.io.FixedLenFeature([], tf.int64),
        'depth': tf.io.FixedLenFeature([], tf.int64)
    }
    
    example = tf.io.parse_single_example(example_proto, feature_description)
    
    # 解析图像数据
    # Parse image data
    image = tf.io.decode_raw(example['image'], tf.uint8)
    image = tf.reshape(image, [32, 32, 3])
    image = tf.cast(image, tf.float32) / 255.0
    
    # 数据增强
    # Data augmentation
    if is_training and config.get('data_augmentation', True):
        if config.get('horizontal_flip', True):
            image = tf.image.random_flip_left_right(image)
        
        if config.get('brightness_delta', 0) > 0:
            image = tf.image.random_brightness(image, config['brightness_delta'])
        
        if config.get('contrast_range'):
            image = tf.image.random_contrast(image, 
                                           config['contrast_range'][0], 
                                           config['contrast_range'][1])
    
    # 标准化处理
    # Normalization
    image = (image - tf.constant([0.4914, 0.4822, 0.4465])) / tf.constant([0.2023, 0.1994, 0.2010])
    
    label = tf.cast(example['label'], tf.int32)
    
    return image, label

def create_dataset(config: Dict[str, Any], is_training: bool = True):
    """创建TensorFlow数据集管道"""
    """Create TensorFlow dataset pipe"""
    tfrecord_path = config['train_tfrecord_path'] if is_training else config['val_tfrecord_path']
    
    # 检查TFRecord文件是否存在
    # Check TFRecord file
    if not os.path.exists(tfrecord_path):
        raise FileNotFoundError(f"TFRecord文件不存在: {tfrecord_path}")
    
    dataset = tf.data.TFRecordDataset(tfrecord_path)
    dataset = dataset.map(
        lambda x: parse_tfrecord_fn(x, config, is_training),
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    )
    
    if is_training:
        dataset = dataset.shuffle(buffer_size=10000)
    
    dataset = dataset.batch(config.get('batch_size', 128))
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    
    return dataset

import tensorflow as tf
from tensorflow.keras import layers, Model

class BasicBlock(layers.Layer):
    """ResNet基础残差块"""
    """ResNet Basic Residual Block"""
    def __init__(self, filter_num, stride=1):
        super(BasicBlock, self).__init__()
        
        self.conv1 = layers.Conv2D(filter_num, (3, 3), strides=stride, 
                                   padding='same', use_bias=False)
        self.bn1 = layers.BatchNormalization()
        self.relu = layers.ReLU()
        
        self.conv2 = layers.Conv2D(filter_num, (3, 3), strides=1, 
                                   padding='same', use_bias=False)
        self.bn2 = layers.BatchNormalization()
        
        # 捷径连接处理
        # Connection Handling
        if stride != 1:
            self.downsample = tf.keras.Sequential([
                layers.Conv2D(filter_num, (1, 1), strides=stride, use_bias=False),
                layers.BatchNormalization()
            ])
        else:
            self.downsample = lambda x: x
    
    def call(self, inputs, training=None):
        residual = inputs
        
        x = self.conv1(inputs)
        x = self.bn1(x, training=training)
        x = self.relu(x)
        
        x = self.conv2(x)
        x = self.bn2(x, training=training)
        
        identity = self.downsample(residual)
        
        output = tf.nn.relu(x + identity)
        return output

class ResNet(Model):
    """ResNet模型"""
    """ResNet model"""
    def __init__(self, layer_dims, num_classes=10):
        super(ResNet, self).__init__()
        
        # 预处理层
        # Preprocessing Layer
        self.stem = tf.keras.Sequential([
            layers.Conv2D(64, (3, 3), strides=1, padding='same', use_bias=False),
            layers.BatchNormalization(),
            layers.ReLU(),
        ])
        
        # 构建残差层
        # Construct the Residual Layer
        self.layer1 = self._build_resblock(64, layer_dims[0], stride=1)
        self.layer2 = self._build_resblock(128, layer_dims[1], stride=2)
        self.layer3 = self._build_resblock(256, layer_dims[2], stride=2)
        self.layer4 = self._build_resblock(512, layer_dims[3], stride=2)
        
        # 分类器
        # Classifier
        self.avgpool = layers.GlobalAveragePooling2D()
        self.fc = layers.Dense(num_classes)
    
    def _build_resblock(self, filter_num, blocks, stride=1):
        """构建残差块序列"""
        """Construct the sequence of residual blocks"""
        res_blocks = tf.keras.Sequential()
        res_blocks.add(BasicBlock(filter_num, stride))
        
        for _ in range(1, blocks):
            res_blocks.add(BasicBlock(filter_num, stride=1))
        
        return res_blocks
    
    def call(self, inputs, training=None):
        x = self.stem(inputs, training=training)
        x = self.layer1(x, training=training)
        x = self.layer2(x, training=training)
        x = self.layer3(x, training=training)
        x = self.layer4(x, training=training)
        x = self.avgpool(x)
        x = self.fc(x)
        return x

def resnet18(num_classes=10):
    """创建ResNet-18模型"""
    """Construct ResNet-18 model"""
    return ResNet([2, 2, 2, 2], num_classes)

import mlflow
import mlflow.tensorflow
import json
import matplotlib.pyplot as plt
from datetime import datetime

class CIFAR10Trainer:
    def __init__(self, config_path='config/training_config.json'):
        self.config = self.load_config(config_path)
        self.setup_mlflow()
        self.model = None
        self.optimizer = None
        
    def load_config(self, config_path):
        """从JSON文件加载超参数配置"""
        """Load hyperparam config from JSON file"""
        with open(config_path, 'r') as f:
            return json.load(f)
    
    def setup_mlflow(self):
        """设置MLflow实验跟踪"""
        """Set MLflow experiment tracking"""
        mlflow.set_tracking_uri("30.22.40.186:5000")
        mlflow.set_experiment("cifar10-resnet-tf2.2-local")
        
    def create_optimizer(self):
        """创建优化器"""
        """Construct optimizer"""
        if self.config['optimizer'] == 'adam':
            return tf.keras.optimizers.Adam(
                learning_rate=self.config['learning_rate']
            )
        else:
            return tf.keras.optimizers.SGD(
                learning_rate=self.config['learning_rate'],
                momentum=self.config.get('momentum', 0.9)
            )
    
    def train(self):
        """训练模型"""
        """Training model"""
        # 创建TFRecord文件（如果不存在）
        # Construct TFRecord file(if not exist)
        if not os.path.exists(self.config['train_tfrecord_path']):
            print("创建TFRecord文件...Construct TFRecord file...")
            try:
                train_samples, test_samples = create_cifar10_tfrecord(
                    self.config['cifar10_local_path']
                )
                print(f"成功加载: 训练样本 {train_samples}, 测试样本 {test_samples}.Load success: {train_samples},{test_samples}")
            except FileNotFoundError as e:
                print(f"错误: {e} error: {e}")
                print("请确保cifar-10-python.tar.gz文件存在于当前目录.Please make sure cifar-10-python.tar.gz exist in current directory")
                return
        
        with mlflow.start_run():
            # 记录超参数
            # Log hyperparam
            mlflow.log_params(self.config)
            mlflow.log_param("framework", "TensorFlow 2.2")
            mlflow.log_param("dataset", "CIFAR-10")
            mlflow.log_param("data_source", "local_file")
            mlflow.log_param("start_time", datetime.now().isoformat())
            
            # 创建数据集
            # Construct dataset
            train_dataset = create_dataset(self.config, is_training=True)
            val_dataset = create_dataset(self.config, is_training=False)
            
            # 创建模型和优化器
            # Construct model and optimizer
            self.model = resnet18(num_classes=self.config['num_classes'])
            self.optimizer = self.create_optimizer()
            
            # 编译模型
            # Compile model
            self.model.compile(
                optimizer=self.optimizer,
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['accuracy']
            )
            
            # 构建模型并显示摘要
            # Construct model and show abstract
            self.model.build((None, self.config['image_size'], 
                            self.config['image_size'], 3))
            self.model.summary()
            
            # 训练参数
            # train param
            train_loss = tf.keras.metrics.Mean(name='train_loss')
            train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
            val_loss = tf.keras.metrics.Mean(name='val_loss')
            val_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='val_accuracy')
            
            @tf.function
            def train_step(images, labels):
                with tf.GradientTape() as tape:
                    predictions = self.model(images, training=True)
                    loss = tf.keras.losses.sparse_categorical_crossentropy(
                        labels, predictions, from_logits=True)
                
                gradients = tape.gradient(loss, self.model.trainable_variables)
                self.optimizer.apply_gradients(zip(gradients, self.model.trainable_variables))
                
                train_loss(loss)
                train_accuracy(labels, predictions)
                return loss
            
            best_val_acc = 0.0
            train_losses = []
            train_accs = []
            val_losses = []
            val_accs = []
            
            print("开始训练...Start training")
            for epoch in range(self.config['epochs']):
                # 重置指标
                # Reset metrics
                train_loss.reset_state()
                train_accuracy.reset_state()
                val_loss.reset_state()
                val_accuracy.reset_state()
                
                # 训练阶段
                # Training step
                for batch, (images, labels) in enumerate(train_dataset):
                    loss = train_step(images, labels)
                    if batch % 50 == 0:
                        #print(f'Epoch {epoch+1}, Batch {batch}, Loss: {loss.numpy():.4f}')
                        print(f'Epoch {epoch+1}, Batch {batch}, Loss: {tf.reduce_mean(loss).numpy():.4f}')
                
                # 验证阶段
                # Validate step
                for images, labels in val_dataset:
                    predictions = self.model(images, training=False)
                    v_loss = tf.keras.losses.sparse_categorical_crossentropy(
                        labels, predictions, from_logits=True)
                    val_loss(v_loss)
                    val_accuracy(labels, predictions)
                
                # 计算epoch指标
                # Calculate epoch metric
                epoch_train_loss = train_loss.result().numpy()
                epoch_train_acc = train_accuracy.result().numpy()
                epoch_val_loss = val_loss.result().numpy()
                epoch_val_acc = val_accuracy.result().numpy()
                
                # 记录到MLflow
                # Log to MLflow
                mlflow.log_metric("train_loss", epoch_train_loss, step=epoch)
                mlflow.log_metric("train_accuracy", epoch_train_acc, step=epoch)
                mlflow.log_metric("val_loss", epoch_val_loss, step=epoch)
                mlflow.log_metric("val_accuracy", epoch_val_acc, step=epoch)
                mlflow.log_metric("learning_rate", 
                                self.optimizer.learning_rate.numpy(), step=epoch)
                
                # 保存训练历史
                # Save training history
                train_losses.append(epoch_train_loss)
                train_accs.append(epoch_train_acc)
                val_losses.append(epoch_val_loss)
                val_accs.append(epoch_val_acc)
                
                print(f'Epoch {epoch+1}/{self.config["epochs"]}: '
                      f'Train Loss: {epoch_train_loss:.4f}, Train Acc: {epoch_train_acc*100:.2f}%, '
                      f'Val Loss: {epoch_val_loss:.4f}, Val Acc: {epoch_val_acc*100:.2f}%')
                
                # 保存最佳模型
                # Save best model
                if epoch_val_acc > best_val_acc:
                    best_val_acc = epoch_val_acc
                    mlflow.tensorflow.log_model(self.model, "best_model")
                    print(f"新的最佳模型已保存，验证准确率: {epoch_val_acc*100:.2f}%The new best model is saved, validate accuracy: {epoch_val_acc*100:.2f}%")
            
            # 记录最终结果
            # Log final result
            mlflow.log_metric("best_val_accuracy", best_val_acc)
            mlflow.tensorflow.log_model(self.model, "final_model")
            
            # 绘制训练曲线
            # Plot the training curves。
            self.plot_training_history(train_losses, train_accs, val_losses, val_accs)
            
            print(f"训练完成！最佳验证准确率: {best_val_acc*100:.2f}%Training complete!The best validate accuracy: {best_val_acc*100:.2f}%")
    
    def plot_training_history(self, train_losses, train_accs, val_losses, val_accs):
        """绘制训练历史曲线"""
        """Plot the training history curves"""
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
        
        # 损失曲线
        # Loss curves
        ax1.plot(train_losses, label='Train Loss')
        ax1.plot(val_losses, label='Val Loss')
        ax1.set_title('Training and Validation Loss')
        ax1.set_xlabel('Epoch')
        ax1.set_ylabel('Loss')
        ax1.legend()
        ax1.grid(True)
        
        # 准确率曲线
        # Accuracy curves
        ax2.plot(train_accs, label='Train Accuracy')
        ax2.plot(val_accs, label='Val Accuracy')
        ax2.set_title('Training and Validation Accuracy')
        ax2.set_xlabel('Epoch')
        ax2.set_ylabel('Accuracy')
        ax2.legend()
        ax2.grid(True)
        
        plt.tight_layout()
        plt.savefig('training_history.png')
        mlflow.log_artifact('training_history.png')
        plt.show()