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)
        
        # 画像データとTagを取得します。
        # 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()