import tensorflow as tf
from tensorflow.keras.applications import EfficientNetB0
from tensorflow.keras.applications.efficientnet import preprocess_input
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, Dropout, BatchNormalization
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import ReduceLROnPlateau, EarlyStopping
import numpy as np

TRAIN_DIR = r'C:\Python_CV\mritest\Training'
TEST_DIR = r'C:\Python_CV\mritest\Testing'
IMG_SIZE = (224, 224)
BATCH_SIZE = 16  # ลด Batch Size ลงเพื่อเพิ่ม Generalization

# 1. Data Augmentation แบบเข้มข้น
train_datagen = ImageDataGenerator(
    preprocessing_function=preprocess_input,
    rotation_range=15,
    width_shift_range=0.1,
    height_shift_range=0.1,
    shear_range=0.1,
    zoom_range=0.15,
    horizontal_flip=True,
    fill_mode='nearest'
)

test_datagen = ImageDataGenerator(preprocessing_function=preprocess_input)

train_gen = train_datagen.flow_from_directory(
    TRAIN_DIR, target_size=IMG_SIZE, batch_size=BATCH_SIZE, class_mode='categorical', shuffle=True
)

test_gen = test_datagen.flow_from_directory(
    TEST_DIR, target_size=IMG_SIZE, batch_size=BATCH_SIZE, class_mode='categorical', shuffle=False
)

# 2. สร้างโครงสร้างโมเดล
base_model = EfficientNetB0(weights='imagenet', include_top=False, input_shape=(224, 224, 3))

# Unfreeze 40 ชั้นสุดท้ายของ EfficientNetB0 เพื่อ Fine-Tune
base_model.trainable = True
for layer in base_model.layers[:-40]:
    layer.trainable = False

x = base_model.output
x = GlobalAveragePooling2D()(x)
x = BatchNormalization()(x)
x = Dense(256, activation='relu')(x)
x = Dropout(0.4)(x)  # ลด Overfitting
outputs = Dense(4, activation='softmax')(x)

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

# 3. คอมไพล์โมเดลด้วย Learning Rate ต่ำสำหรับ Fine-tuning
model.compile(
    optimizer=Adam(learning_rate=1e-4),
    loss='categorical_crossentropy',
    metrics=['accuracy']
)

# 4. Callbacks ปรับ Learning Rate อัตโนมัติเมื่อ Accuracy เริ่มนิ่ง
callbacks = [
    ReduceLROnPlateau(monitor='val_accuracy', factor=0.3, patience=3, verbose=1, min_lr=1e-6),
    EarlyStopping(monitor='val_accuracy', patience=7, restore_best_weights=True)
]

# 5. เทรนโมเดล
print("Starting Fine-Tuning Training...")
history = model.fit(
    train_gen,
    epochs=25,
    validation_data=test_gen,
    callbacks=callbacks
)

# 6. ประเมินผลความแม่นยำ
test_loss, test_acc = model.evaluate(test_gen)
print(f"\n>>> Final Test Accuracy: {test_acc * 100:.2f}% <<<")

# บันทึกโมเดลไว้ในโฟลเดอร์โครงการ
model.save(r'C:\Python_CV\mritest\efficientnetb0_finetuned_brain_mri.keras')
print("เซฟโมเดลแบบ .keras เรียบร้อยแล้ว!")

from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# 1. พยากรณ์ผลบนชุด Testing Data
test_gen.reset()
y_pred_prob = model.predict(test_gen, verbose=1)
y_pred = np.argmax(y_pred_prob, axis=1)
y_true = test_gen.classes
class_labels = list(test_gen.class_indices.keys())

# 2. พิมพ์รายงาน Classification Report (Precision, Recall, F1-score)
print("\n================ Classification Report ================")
print(classification_report(y_true, y_pred, target_names=class_labels))

# 3. วาดกราฟ Confusion Matrix
cm = confusion_matrix(y_true, y_pred)
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
            xticklabels=class_labels, yticklabels=class_labels)
plt.title('Fine-Tuned EfficientNetB0 - Confusion Matrix')
plt.xlabel('Predicted Label')
plt.ylabel('True Label')
plt.show()

# พล็อต กราฟ Loss & Accuracy
acc = history.history['accuracy']
val_acc = history.history['val_accuracy']
loss = history.history['loss']
val_loss = history.history['val_loss']
epochs_range = range(len(acc))

plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, acc, label='Training Accuracy')
plt.plot(epochs_range, val_acc, label='Validation/Test Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')

plt.subplot(1, 2, 2)
plt.plot(epochs_range, loss, label='Training Loss')
plt.plot(epochs_range, val_loss, label='Validation/Test Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.show()