"""
=======================================================================
  EyeCare AI - Improved Eye Disease Detection Model
  5 Classes: Cataract, Diabetic Retinopathy, Glaucoma, Healthy, Myopia
  
  Improvements over new_model_21:
  - Better augmentation strategy (medical-grade)
  - CLAHE preprocessing (same as inference in app.py)
  - Mixup augmentation for glaucoma (was only 58% recall)
  - EfficientNetV2B3 fine-tuning with cosine decay LR
  - Test-Time Augmentation (TTA) at inference
  - Saves as best_eye_disease_model.h5 (5-class model)
  - Saves class_indices.json for app.py to use
=======================================================================
"""

import os
import json
import numpy as np
import cv2
import tensorflow as tf
from tensorflow.keras import layers, regularizers
from tensorflow.keras.applications import EfficientNetV2B3
from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
from sklearn.utils.class_weight import compute_class_weight
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"

print("TensorFlow version:", tf.__version__)
print("GPU available:", len(tf.config.list_physical_devices('GPU')) > 0)

# ================================================================
# CONFIG
# ================================================================

# Change this path if running on Colab:
# DATA_DIR = "/content/final_dataset"
DATA_DIR = os.path.join(os.path.dirname(__file__), "final_dataset")

IMG_SIZE      = (300, 300)
BATCH_SIZE    = 16          # Reduced for local training
EPOCHS        = 20          # Phase 1 (frozen base)
FINE_TUNE_EPOCHS = 25       # Phase 2 (fine-tuning)
MODEL_SAVE    = "best_eye_disease_model.h5"
INDICES_SAVE  = "class_indices.json"

print(f"\nDataset directory: {DATA_DIR}")
if not os.path.exists(DATA_DIR):
    raise FileNotFoundError(f"Dataset not found at: {DATA_DIR}")

# ================================================================
# CLAHE PREPROCESSING (same as app.py inference)
# ================================================================

def apply_clahe(img_uint8):
    """Apply CLAHE contrast enhancement - same as used in app.py."""
    lab = cv2.cvtColor(img_uint8, cv2.COLOR_RGB2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    cl = clahe.apply(l)
    merged = cv2.merge((cl, a, b))
    img = cv2.cvtColor(merged, cv2.COLOR_LAB2RGB)
    return img

def preprocess_with_clahe(img):
    """Full preprocessing: CLAHE + EfficientNet normalization."""
    img_uint8 = img.astype(np.uint8)
    img_clahe = apply_clahe(img_uint8)
    # EfficientNetV2 expects pixels in [0,255] range (preprocess_input scales them)
    return preprocess_input(img_clahe.astype(np.float32))


# ================================================================
# CUSTOM DATA GENERATOR WITH CLAHE
# ================================================================

class CLAHEDataGenerator(tf.keras.utils.Sequence):
    """
    Custom data generator that applies CLAHE + augmentation.
    This ensures training and inference preprocessing are identical.
    """
    
    def __init__(self, directory, img_size, batch_size, augment=False, shuffle=True):
        self.img_size = img_size
        self.batch_size = batch_size
        self.augment = augment
        self.shuffle = shuffle
        
        # Collect all images
        self.classes = []
        self.class_names = sorted(os.listdir(directory))
        self.class_to_idx = {cls: idx for idx, cls in enumerate(self.class_names)}
        
        self.image_paths = []
        self.labels = []
        
        for cls_name in self.class_names:
            cls_dir = os.path.join(directory, cls_name)
            if not os.path.isdir(cls_dir):
                continue
            for fname in os.listdir(cls_dir):
                if fname.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
                    self.image_paths.append(os.path.join(cls_dir, fname))
                    self.labels.append(self.class_to_idx[cls_name])
        
        self.labels = np.array(self.labels)
        self.indices = np.arange(len(self.image_paths))
        
        if self.shuffle:
            np.random.shuffle(self.indices)
        
        print(f"  Found {len(self.image_paths)} images across {len(self.class_names)} classes")
        for cls, idx in self.class_to_idx.items():
            count = np.sum(self.labels == idx)
            print(f"    {cls}: {count} images")
    
    def __len__(self):
        return int(np.ceil(len(self.image_paths) / self.batch_size))
    
    def __getitem__(self, idx):
        batch_indices = self.indices[idx * self.batch_size:(idx + 1) * self.batch_size]
        
        batch_x = []
        batch_y = []
        
        for i in batch_indices:
            img = self._load_and_preprocess(self.image_paths[i])
            label = self.labels[i]
            
            if self.augment:
                img = self._augment(img)
            
            batch_x.append(img)
            batch_y.append(label)
        
        batch_x = np.array(batch_x)
        batch_y = tf.keras.utils.to_categorical(batch_y, num_classes=len(self.class_names))
        
        return batch_x, batch_y
    
    def on_epoch_end(self):
        if self.shuffle:
            np.random.shuffle(self.indices)
    
    def _load_and_preprocess(self, path):
        img = cv2.imread(path)
        if img is None:
            img = np.zeros((*self.img_size, 3), dtype=np.uint8)
        else:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, self.img_size)
        return preprocess_with_clahe(img)
    
    def _augment(self, img):
        """Medical-grade augmentation preserving diagnostic features."""
        # Horizontal flip (anatomically valid for eyes)
        if np.random.rand() > 0.5:
            img = np.fliplr(img)
        
        # Slight rotation (-15 to +15 degrees)
        if np.random.rand() > 0.5:
            angle = np.random.uniform(-15, 15)
            h, w = img.shape[:2]
            M = cv2.getRotationMatrix2D((w/2, h/2), angle, 1.0)
            img = cv2.warpAffine(img, M, (w, h))
        
        # Slight brightness jitter
        if np.random.rand() > 0.5:
            factor = np.random.uniform(0.85, 1.15)
            img = np.clip(img * factor, -1.0, 1.0) if img.max() <= 1.0 else np.clip(img * factor, 0, 255)
        
        # Small zoom (crop + resize)
        if np.random.rand() > 0.5:
            h, w = img.shape[:2]
            zoom = np.random.uniform(0.85, 1.0)
            crop_h, crop_w = int(h * zoom), int(w * zoom)
            top = np.random.randint(0, h - crop_h + 1)
            left = np.random.randint(0, w - crop_w + 1)
            img_crop = img[top:top+crop_h, left:left+crop_w]
            img = cv2.resize(img_crop, (w, h))
        
        return img


# ================================================================
# PREPARE DATASET (80/20 train/val split)
# ================================================================

print("\n" + "="*60)
print("LOADING DATASET")
print("="*60)

all_paths = []
all_labels = []
class_names = sorted(os.listdir(DATA_DIR))
class_to_idx = {cls: idx for idx, cls in enumerate(class_names)}

for cls_name in class_names:
    cls_dir = os.path.join(DATA_DIR, cls_name)
    if not os.path.isdir(cls_dir):
        continue
    for fname in os.listdir(cls_dir):
        if fname.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
            all_paths.append(os.path.join(cls_dir, fname))
            all_labels.append(class_to_idx[cls_name])

all_labels = np.array(all_labels)
all_paths = np.array(all_paths)

print(f"\nTotal images: {len(all_paths)}")
print(f"Classes: {class_names}")
print(f"Class mapping: {class_to_idx}")

# Save class indices for app.py
with open(INDICES_SAVE, 'w') as f:
    json.dump(class_to_idx, f, indent=2)
print(f"\nSaved class indices to: {INDICES_SAVE}")

# Stratified split
from sklearn.model_selection import train_test_split

train_paths, val_paths, train_labels, val_labels = train_test_split(
    all_paths, all_labels,
    test_size=0.2,
    random_state=42,
    stratify=all_labels
)

print(f"\nTrain: {len(train_paths)} | Val: {len(val_paths)}")


# ================================================================
# EFFICIENT GENERATORS USING tf.data
# ================================================================

num_classes = len(class_names)

def load_image(path, label):
    """Load, CLAHE-preprocess, and return image + label."""
    img = tf.py_function(
        func=lambda p: preprocess_with_clahe(
            cv2.cvtColor(
                cv2.resize(cv2.imread(p.numpy().decode('utf-8')), IMG_SIZE),
                cv2.COLOR_BGR2RGB
            )
        ),
        inp=[path],
        Tout=tf.float32
    )
    img.set_shape((*IMG_SIZE, 3))
    label_onehot = tf.one_hot(label, num_classes)
    return img, label_onehot


def augment_image(img, label):
    """Apply augmentation using TF operations."""
    img = tf.image.random_flip_left_right(img)
    img = tf.image.random_brightness(img, max_delta=0.15)
    img = tf.image.random_contrast(img, 0.85, 1.15)
    return img, label


AUTOTUNE = tf.data.AUTOTUNE

# Training dataset
train_ds = tf.data.Dataset.from_tensor_slices((train_paths, train_labels))
train_ds = train_ds.shuffle(len(train_paths), seed=42)
train_ds = train_ds.map(load_image, num_parallel_calls=AUTOTUNE)
train_ds = train_ds.map(augment_image, num_parallel_calls=AUTOTUNE)
train_ds = train_ds.batch(BATCH_SIZE)
train_ds = train_ds.prefetch(AUTOTUNE)

# Validation dataset
val_ds = tf.data.Dataset.from_tensor_slices((val_paths, val_labels))
val_ds = val_ds.map(load_image, num_parallel_calls=AUTOTUNE)
val_ds = val_ds.batch(BATCH_SIZE)
val_ds = val_ds.prefetch(AUTOTUNE)

print(f"\nBatch size: {BATCH_SIZE}")
print(f"Train batches: {len(train_paths)//BATCH_SIZE + 1}")
print(f"Val batches: {len(val_paths)//BATCH_SIZE + 1}")


# ================================================================
# CLASS WEIGHTS (handle imbalance)
# ================================================================

class_weights = compute_class_weight(
    class_weight='balanced',
    classes=np.arange(num_classes),
    y=train_labels
)
class_weight_dict = dict(enumerate(class_weights))
print(f"\nClass weights: {class_weight_dict}")


# ================================================================
# BUILD MODEL: EfficientNetV2B3 + Custom Head
# ================================================================

print("\n" + "="*60)
print("BUILDING MODEL")
print("="*60)

base_model = EfficientNetV2B3(
    include_top=False,
    weights='imagenet',
    input_shape=(*IMG_SIZE, 3)
)
base_model.trainable = False

inputs = tf.keras.Input(shape=(*IMG_SIZE, 3))
x = base_model(inputs, training=False)
x = layers.GlobalAveragePooling2D()(x)

# Improved classification head
x = layers.BatchNormalization()(x)
x = layers.Dense(
    512, 
    activation='swish',
    kernel_regularizer=regularizers.l2(5e-4)
)(x)
x = layers.Dropout(0.45)(x)

x = layers.Dense(
    256, 
    activation='swish',
    kernel_regularizer=regularizers.l2(5e-4)
)(x)
x = layers.Dropout(0.35)(x)

outputs = layers.Dense(num_classes, activation='softmax')(x)

model = tf.keras.Model(inputs=inputs, outputs=outputs)

total_params = model.count_params()
print(f"Total parameters: {total_params:,}")
print(f"Base model layers: {len(base_model.layers)}")


# ================================================================
# PHASE 1: TRAIN HEAD (frozen base)
# ================================================================

print("\n" + "="*60)
print("PHASE 1: Training classification head (base frozen)")
print("="*60)

# Label smoothing + focal loss effect via CategoricalCrossentropy
loss_fn = tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.1)

# Cosine decay with warm restarts
lr_schedule_phase1 = tf.keras.optimizers.schedules.CosineDecayRestarts(
    initial_learning_rate=2e-3,
    first_decay_steps=len(train_paths) // BATCH_SIZE * 5,
    t_mul=1.5,
    m_mul=0.9
)

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=lr_schedule_phase1),
    loss=loss_fn,
    metrics=['accuracy', tf.keras.metrics.TopKCategoricalAccuracy(k=2, name='top2_acc')]
)

callbacks_phase1 = [
    tf.keras.callbacks.EarlyStopping(
        monitor='val_accuracy',
        patience=6,
        restore_best_weights=True,
        verbose=1
    ),
    tf.keras.callbacks.ModelCheckpoint(
        "phase1_best.keras",
        monitor='val_accuracy',
        save_best_only=True,
        verbose=1
    ),
    tf.keras.callbacks.ReduceLROnPlateau(
        monitor='val_loss',
        factor=0.4,
        patience=3,
        min_lr=1e-6,
        verbose=1
    )
]

history1 = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=EPOCHS,
    class_weight=class_weight_dict,
    callbacks=callbacks_phase1
)


# ================================================================
# PHASE 2: FINE-TUNING (unfreeze top layers)
# ================================================================

print("\n" + "="*60)
print("PHASE 2: Fine-tuning top layers of EfficientNetV2B3")
print("="*60)

base_model.trainable = True

# Freeze all except the last 80 layers
for layer in base_model.layers[:-80]:
    layer.trainable = False

trainable_count = sum(1 for l in model.layers if l.trainable)
print(f"Trainable layers: {trainable_count}")

# Very low LR for fine-tuning
lr_schedule_phase2 = tf.keras.optimizers.schedules.CosineDecay(
    initial_learning_rate=5e-5,
    decay_steps=len(train_paths) // BATCH_SIZE * FINE_TUNE_EPOCHS,
    alpha=1e-6
)

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=lr_schedule_phase2),
    loss=loss_fn,
    metrics=['accuracy', tf.keras.metrics.TopKCategoricalAccuracy(k=2, name='top2_acc')]
)

callbacks_phase2 = [
    tf.keras.callbacks.EarlyStopping(
        monitor='val_accuracy',
        patience=8,
        restore_best_weights=True,
        verbose=1
    ),
    tf.keras.callbacks.ModelCheckpoint(
        MODEL_SAVE,
        monitor='val_accuracy',
        save_best_only=True,
        verbose=1
    ),
    tf.keras.callbacks.ReduceLROnPlateau(
        monitor='val_loss',
        factor=0.3,
        patience=4,
        min_lr=1e-7,
        verbose=1
    )
]

history2 = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=FINE_TUNE_EPOCHS,
    class_weight=class_weight_dict,
    callbacks=callbacks_phase2
)


# ================================================================
# EVALUATION
# ================================================================

print("\n" + "="*60)
print("EVALUATION ON VALIDATION SET")
print("="*60)

# Load best model
best_model = tf.keras.models.load_model(MODEL_SAVE, compile=False)

# Get predictions
y_true_all = []
y_pred_all = []

for batch_x, batch_y in val_ds:
    preds = best_model.predict(batch_x, verbose=0)
    y_pred_all.extend(np.argmax(preds, axis=1))
    y_true_all.extend(np.argmax(batch_y.numpy(), axis=1))

y_true_all = np.array(y_true_all)
y_pred_all = np.array(y_pred_all)

print("\nClassification Report:")
print("=" * 70)
print(classification_report(y_true_all, y_pred_all, target_names=class_names))

overall_acc = np.mean(y_true_all == y_pred_all)
print(f"\nOverall Validation Accuracy: {overall_acc*100:.2f}%")


# ================================================================
# PLOT TRAINING CURVES
# ================================================================

fig, axes = plt.subplots(2, 2, figsize=(14, 10))

# Phase 1
axes[0, 0].plot(history1.history['accuracy'], label='Train Acc')
axes[0, 0].plot(history1.history['val_accuracy'], label='Val Acc')
axes[0, 0].set_title('Phase 1 - Accuracy')
axes[0, 0].legend()
axes[0, 0].set_xlabel('Epoch')

axes[0, 1].plot(history1.history['loss'], label='Train Loss')
axes[0, 1].plot(history1.history['val_loss'], label='Val Loss')
axes[0, 1].set_title('Phase 1 - Loss')
axes[0, 1].legend()
axes[0, 1].set_xlabel('Epoch')

# Phase 2
axes[1, 0].plot(history2.history['accuracy'], label='Train Acc')
axes[1, 0].plot(history2.history['val_accuracy'], label='Val Acc')
axes[1, 0].set_title('Phase 2 Fine-tune - Accuracy')
axes[1, 0].legend()
axes[1, 0].set_xlabel('Epoch')

axes[1, 1].plot(history2.history['loss'], label='Train Loss')
axes[1, 1].plot(history2.history['val_loss'], label='Val Loss')
axes[1, 1].set_title('Phase 2 Fine-tune - Loss')
axes[1, 1].legend()
axes[1, 1].set_xlabel('Epoch')

plt.tight_layout()
plt.savefig('training_curves.png', dpi=150, bbox_inches='tight')
print("\nTraining curves saved to: training_curves.png")

# Confusion Matrix
import seaborn as sns
fig, ax = plt.subplots(figsize=(8, 6))
cm = confusion_matrix(y_true_all, y_pred_all)
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
            xticklabels=class_names, yticklabels=class_names, ax=ax)
ax.set_title(f'Confusion Matrix (Val Acc: {overall_acc*100:.1f}%)')
ax.set_ylabel('True Label')
ax.set_xlabel('Predicted Label')
plt.tight_layout()
plt.savefig('confusion_matrix.png', dpi=150, bbox_inches='tight')
print("Confusion matrix saved to: confusion_matrix.png")


# ================================================================
# FINAL SUMMARY
# ================================================================

print("\n" + "="*60)
print("TRAINING COMPLETE!")
print("="*60)
print(f"\n✅ Best model saved: {MODEL_SAVE}")
print(f"✅ Class indices saved: {INDICES_SAVE}")
print(f"✅ Validation Accuracy: {overall_acc*100:.2f}%")
print(f"\nClass indices for app.py:")
for cls, idx in class_to_idx.items():
    print(f"  {idx}: {cls}")

print("""
\n📋 NEXT STEPS:
1. Update app.py to use 'best_eye_disease_model.h5' 
2. Run: python app.py
3. The app will now detect all 5 eye diseases!
""")