from google.colab import drive
drive.mount('/content/drive')

import tensorflow as tf
from tensorflow.keras.applications import EfficientNetB0
from tensorflow.keras.layers import GlobalAveragePooling2D, Dropout, Dense, BatchNormalization
from tensorflow.keras.models import Model
from tensorflow.keras.regularizers import l2
from tensorflow.keras.preprocessing import image_dataset_from_directory
import matplotlib.pyplot as plt
import numpy as np
from energyCnV import EnergyMonitor


# Dataset paths
train_dir = " " # Training data URI
val_dir = " " # Validiation or testing data URI
IMG_SIZE = (224, 224)
BATCH_SIZE = 32

# Load datasets
train_dataset = image_dataset_from_directory(
    train_dir,
    shuffle=True,
    batch_size=BATCH_SIZE,
    image_size=IMG_SIZE,
    seed=42
)

val_dataset = image_dataset_from_directory(
    val_dir,
    shuffle=True,
    batch_size=BATCH_SIZE,
    image_size=IMG_SIZE,
    seed=42
)

# Data augmentation
data_augmentation = tf.keras.Sequential([
    tf.keras.layers.RandomFlip('horizontal'),
    tf.keras.layers.RandomRotation(0.2),
    tf.keras.layers.RandomZoom(0.3),
])

# EfficientNet preprocessing
preprocess_input = tf.keras.applications.efficientnet.preprocess_input

# Model builder
def build_fall_model():
    input_shape = IMG_SIZE + (3,)
    base_model = EfficientNetB0(include_top=False, input_shape=input_shape, weights="imagenet")
    base_model.trainable = False  # Freeze base model initially

    inputs = tf.keras.Input(shape=input_shape)
    x = data_augmentation(inputs)
    x = preprocess_input(x)
    x = base_model(x, training=False)
    x = GlobalAveragePooling2D()(x)
    x = BatchNormalization()(x)
    x = Dropout(0.4)(x)
    outputs = Dense(1, activation='sigmoid', kernel_regularizer=l2(0.001))(x)

    model = Model(inputs, outputs)
    return model, base_model

# Build and compile model
model, base_model = build_fall_model()

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
              loss='binary_crossentropy',
              metrics=['accuracy'])

# Initial training
initial_epochs = 10
history = model.fit(train_dataset, validation_data=val_dataset, epochs=initial_epochs)

# Fine-tuning
base_model.trainable = True
fine_tune_at = 150

for layer in base_model.layers[:fine_tune_at]:
    layer.trainable = False

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),
              loss='binary_crossentropy',
              metrics=['accuracy'])

fine_tune_epochs = 5
total_epochs = initial_epochs + fine_tune_epochs

history_fine = model.fit(train_dataset, validation_data=val_dataset,
                         epochs=total_epochs, initial_epoch=history.epoch[-1]+1)

# Plot Accuracy and Loss
acc = history.history['accuracy'] + history_fine.history['accuracy']
val_acc = history.history['val_accuracy'] + history_fine.history['val_accuracy']

loss = history.history['loss'] + history_fine.history['loss']
val_loss = history.history['val_loss'] + history_fine.history['val_loss']

epochs_range = range(len(acc))

plt.figure(figsize=(16, 6))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, acc, label='Training Accuracy')
plt.plot(epochs_range, val_acc, label='Validation 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 Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')

plt.show()

from tensorflow.keras.preprocessing import image

img_path = " " # Test Image URI
img = image.load_img(img_path, target_size=IMG_SIZE)
img_array = image.img_to_array(img)
img_array = np.expand_dims(img_array, axis=0)
img_array = preprocess_input(img_array)

plt.imshow(img)
plt.axis("off")
plt.show()

prediction = model.predict(img_array)
print(prediction)

if prediction[0] < 0.5:
    print("Prediction: 🚨 Fall Detected! 🚨")
else:
    print("Prediction: ✅ No Fall Detected.")