import numpy as np
import tensorflow as tf
from tensorflow.keras.applications import InceptionV3
from tensorflow.keras import layers, models
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.utils import to_categorical
from sklearn.metrics import classification_report, confusion_matrix
import seaborn as sns
import matplotlib.pyplot as plt

# ------------------ Load Data ------------------
(X_train, y_train), (X_test, y_test) = cifar10.load_data()

# Normalize
X_train = X_train / 255.0
X_test = X_test / 255.0

# Resize to 299x299
X_train = tf.image.resize(X_train, (299,299))
X_test = tf.image.resize(X_test, (299,299))

# One-hot encoding
y_train_cat = to_categorical(y_train, 10)
y_test_cat = to_categorical(y_test, 10)

# ------------------ Load InceptionV3 ------------------
base_model = InceptionV3(weights='imagenet',
                         include_top=False,
                         input_shape=(299,299,3))

# Freeze all layers
for layer in base_model.layers:
    layer.trainable = False

# ------------------ Custom Model ------------------
x = base_model.output
x = layers.GlobalAveragePooling2D()(x)
x = layers.BatchNormalization()(x)
x = layers.Dense(256, activation='relu')(x)
x = layers.Dropout(0.5)(x)
output = layers.Dense(10, activation='softmax')(x)

model = models.Model(inputs=base_model.input, outputs=output)

# Compile
model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

# Train (Feature Extraction)
model.fit(X_train, y_train_cat, epochs=5, batch_size=64, validation_split=0.1)

# ------------------ Fine Tuning using startswith ------------------

for layer in base_model.layers:
    if (layer.name.startswith('mixed7') or
        layer.name.startswith('mixed8') or
        layer.name.startswith('mixed9') or
        layer.name.startswith('mixed10')):
        layer.trainable = True

# Recompile with low LR
model.compile(optimizer=tf.keras.optimizers.Adam(1e-5),
              loss='categorical_crossentropy',
              metrics=['accuracy'])

# Train again
model.fit(X_train, y_train_cat, epochs=3, batch_size=64)

# ------------------ Evaluation ------------------

# Accuracy
loss, acc = model.evaluate(X_test, y_test_cat)
print("Test Accuracy:", acc)

# Predictions
y_pred_probs = model.predict(X_test)
y_pred = np.argmax(y_pred_probs, axis=1)

# Flatten labels
y_test_flat = y_test.flatten()

# Precision, Recall, F1-score
print("\nClassification Report:\n")
print(classification_report(y_test_flat, y_pred))

# Confusion Matrix
cm = confusion_matrix(y_test_flat, y_pred)

plt.figure(figsize=(8,6))
sns.heatmap(cm, annot=True, fmt='d')
plt.xlabel("Predicted")
plt.ylabel("Actual")
plt.title("Confusion Matrix")
plt.show()
