# cifar10_resnet.py
import os
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.model_selection import train_test_split
import seaborn as sns

CLASS_NAMES = [
    "airplane", "automobile", "bird", "cat", "deer",
    "dog", "frog", "horse", "ship", "truck"
]

def create_results_folder():
    os.makedirs("results", exist_ok=True)
    
def save_sample_images(x_train, y_train):
    plt.figure(figsize=(12, 8))
    
    for i in range(15):
        plt.subplot(3, 5, i + 1)
        plt.imshow(x_train[i])
        label = np.argmax(y_train[i]) if len(y_train[i].shape) > 0 else y_train[i]
        plt.title(CLASS_NAMES[label])
        plt.axis("off")
    
    plt.tight_layout()
    plt.savefig("results/cifar10_samples.png", dpi=300, bbox_inches="tight")
    plt.close()
    
def save_training_history(history):
    # Accuracy plot
    plt.figure(figsize=(8, 6))
    plt.plot(history.history["accuracy"], label="Training Accuracy")
    plt.plot(history.history["val_accuracy"], label="Validation Accuracy")
    plt.xlabel("Epoch")
    plt.ylabel("Accuracy")
    plt.title("Training and Validation Accuracy")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("results/training_accuracy.png", dpi=300, bbox_inches="tight")
    plt.close()

    # Loss plot
    plt.figure(figsize=(8, 6))
    plt.plot(history.history["loss"], label="Training Loss")
    plt.plot(history.history["val_loss"], label="Validation Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.title("Training and Validation Loss")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("results/training_loss.png", dpi=300, bbox_inches="tight")
    plt.close()
    
def evaluate_and_save_results(model, x_test, y_test):
    y_pred_probs = model.predict(x_test)
    y_pred = np.argmax(y_pred_probs, axis=1)
    y_true = np.argmax(y_test, axis=1)

    # Classification report
    report = classification_report(y_true, y_pred, target_names=CLASS_NAMES)
    with open("results/classification_report.txt", "w", encoding="utf-8") as f:
        f.write(report)

    print("\nClassification Report:\n")
    print(report)

    # Confusion matrix
    cm = confusion_matrix(y_true, y_pred)

    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
                xticklabels=CLASS_NAMES, yticklabels=CLASS_NAMES)
    plt.xlabel("Predicted Label")
    plt.ylabel("True Label")
    plt.title("Confusion Matrix - CIFAR-10")
    plt.tight_layout()
    plt.savefig("results/confusion_matrix.png", dpi=300, bbox_inches="tight")
    plt.close()

    return y_true, y_pred, y_pred_probs

def save_example_predictions(model, x_test, y_test):
    y_pred_probs = model.predict(x_test[:15])
    y_pred = np.argmax(y_pred_probs, axis=1)
    y_true = np.argmax(y_test[:15], axis=1)

    plt.figure(figsize=(12, 8))
    for i in range(15):
        plt.subplot(3, 5, i + 1)
        plt.imshow(x_test[i])
        plt.title(f"T: {CLASS_NAMES[y_true[i]]}\nP: {CLASS_NAMES[y_pred[i]]}")
        plt.axis("off")

    plt.tight_layout()
    plt.savefig("results/example_predictions.png", dpi=300, bbox_inches="tight")
    plt.close()

def load_and_preprocess_data():
    """Load and preprocess CIFAR-10 dataset"""
    (x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
    
    # Normalize pixel values
    x_train = x_train.astype('float32') / 255.0
    x_test = x_test.astype('float32') / 255.0
    
    # One-hot encode labels
    y_train = keras.utils.to_categorical(y_train, 10)
    y_test = keras.utils.to_categorical(y_test, 10)
    
    return (x_train, y_train), (x_test, y_test)

def create_residual_block(x, filters, downsample=False):
    """Create a residual block with optional downsampling"""
    strides = 2 if downsample else 1
    y = layers.BatchNormalization()(x)
    y = layers.Activation('relu')(y)
    
    # First convolution
    y = layers.Conv2D(filters, 3, strides=strides, padding='same', 
                      kernel_initializer='he_normal')(y)
    y = layers.BatchNormalization()(y)
    y = layers.Activation('relu')(y)
    
    # Second convolution
    y = layers.Conv2D(filters, 3, strides=1, padding='same', 
                      kernel_initializer='he_normal')(y)
    
    # Shortcut connection
    if downsample or x.shape[-1] != filters:
        x = layers.Conv2D(filters, 1, strides=strides, padding='same', 
                         kernel_initializer='he_normal')(x)
    
    return layers.Add()([x, y])

def build_resnet_model(input_shape=(32, 32, 3), num_classes=10):
    """Build ResNet-20 model for CIFAR-10"""
    inputs = keras.Input(shape=input_shape)
    
    # Data augmentation
    x = layers.RandomFlip("horizontal")(inputs)
    x = layers.RandomRotation(0.1)(x)
    x = layers.RandomZoom(0.1)(x)
    x = layers.RandomContrast(0.2)(x)
    
    # Initial convolution
    x = layers.Conv2D(16, 3, padding='same', kernel_initializer='he_normal')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    
    # Stage 1 (16 filters)
    for _ in range(3):
        x = create_residual_block(x, 16)
    
    # Stage 2 (32 filters)
    x = create_residual_block(x, 32, downsample=True)
    for _ in range(2):
        x = create_residual_block(x, 32)
    
    # Stage 3 (64 filters)
    x = create_residual_block(x, 64, downsample=True)
    for _ in range(2):
        x = create_residual_block(x, 64)
    
    # Final layers
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    x = layers.GlobalAveragePooling2D()(x)
    x = layers.Dropout(0.4)(x)
    outputs = layers.Dense(num_classes, activation='softmax')(x)
    
    model = keras.Model(inputs, outputs)
    return model

def train_model():
    """Train the ResNet model"""
    # Load data
    (x_train, y_train), (x_test, y_test) = load_and_preprocess_data()
    
    # Create validation split
    
    x_train, x_val, y_train, y_val = train_test_split(
        x_train, y_train, test_size=0.1, random_state=42
    )
    
    # Build model
    model = build_resnet_model()
    
    # Compile model
    model.compile(
        optimizer=keras.optimizers.AdamW(learning_rate=0.001, weight_decay=0.0001),
        loss='categorical_crossentropy',
        metrics=['accuracy']
    )
    
    # Callbacks
    callbacks = [
        keras.callbacks.ReduceLROnPlateau(
            monitor='val_loss', factor=0.5, patience=5, min_lr=1e-7
        ),
        keras.callbacks.EarlyStopping(
            monitor='val_accuracy', patience=10, restore_best_weights=True
        ),
        keras.callbacks.ModelCheckpoint(
            'best_resnet_model.h5', save_best_only=True, monitor='val_accuracy'
        )
    ]
    
    # Train model
    history = model.fit(
        x_train, y_train,
        batch_size=128,
        epochs=100,
        validation_data=(x_val, y_val),
        callbacks=callbacks,
        verbose=1
    )
    
    return model, history, (x_test, y_test)

if __name__ == "__main__":
    create_results_folder()

    # Load dataset
    (x_train, y_train), (x_test, y_test) = load_and_preprocess_data()

    # Save sample dataset images
    save_sample_images(x_train, y_train)

    # Load the already trained best model
    model = keras.models.load_model("best_resnet_model.h5")

    # Evaluate and save outputs
    y_true, y_pred, y_pred_probs = evaluate_and_save_results(model, x_test, y_test)

    # Save example predictions
    save_example_predictions(model, x_test, y_test)

    print("\nAll results saved in the 'results' folder.")