import tensorflow as tf
from tensorflow.keras import layers, models, regularizers
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint, EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.applications import ConvNeXtTiny
from tensorflow.keras.applications.convnext import preprocess_input
from time import time
import matplotlib.pyplot as plt

# This is the new 2026 version
# This model was trained with 13,185 images
# See CNNImageProcessor solution for creating the test images for training this model
# **************************************** I M P O R T A N T ********************************************************
# to start the environment "source tf_gpu/bin/activate"
# then type "code ."
# Train the model on EUPORIE laptop using the GPU card with WSL2.  (Windows Subsystem For Linux).  I am running Ubuntu1 22.04.2
# To launch WSL open up a command prompt, run powershell and type "wsl".
# The folder structure will be /home/pi/CNN.
# You can access the folder structure through windows explorer.  type "\\wsl$" in explorer and navigate to the folder.
# drop in the Data and Model and run the model
# There is a shell script in the Scripts folder. setup_tf_gpu.sh  Copy teh script to the CNN folder and run it.  It will
# create the venv enviroment and install python 3.10 and tensorflow (gpu)
# ********************************************************************************************************************

# Figure out if we are training in CPU or GPU
print("GPUs:", tf.config.list_physical_devices('GPU'))
# -----------------------
# ConvNeXt-Tiny Base Model
# -----------------------
modelname='convnext_20260228_90.h5.keras'

# convneXt was pretrained with 224 but our image data is 128 so we upscale our images to match the 224 requirements of the model
actualImageDimension=224
convneXtImageDimension=224

# Tensorboard
log_dir = f'logs/convnext_{int(time())}'
tensorboard = TensorBoard(log_dir=log_dir)

# -----------------------
# Configuration
# -----------------------
shuffle_count=3000
dataset_path = '/home/pi/DeepLearning/Data'
image_size = (actualImageDimension, actualImageDimension)
batch_size = 16    # try 16  was 32
image_size=(actualImageDimension, actualImageDimension)


# -----------------------
# Dataset Loading
# -----------------------

train_ds = tf.keras.preprocessing.image_dataset_from_directory(
    dataset_path,
    label_mode="binary",
    subset="training",
    validation_split=0.2,
    image_size=image_size,
    color_mode='rgb',   # IMPORTANT for grayscale datasets
    batch_size=batch_size,
    seed=50
)

val_ds = tf.keras.preprocessing.image_dataset_from_directory(
    dataset_path,
    label_mode="binary",
    subset="validation",
    validation_split=0.2,
    image_size=image_size,
    color_mode='rgb',
    batch_size=batch_size,
    seed=50
)

# -----------------------
# Data Augmentation
# -----------------------
# I do this in c#-land

def preprocess_val(x, y):
    return x, y

val_ds = (
    val_ds
    .prefetch(tf.data.AUTOTUNE)
)

train_ds = (
    train_ds
    .shuffle(3000)
    .prefetch(tf.data.AUTOTUNE)
)


# -----------------------
# ConvNeXt-Tiny Base Model
# -----------------------

base_model = ConvNeXtTiny(
    weights='imagenet',
    include_top=False,
    input_shape=(convneXtImageDimension, convneXtImageDimension, 3)
)

base_model.trainable = False  # Freeze for initial training

# -----------------------
# Build Full Model (Preprocessing Inside Model)
# -----------------------
inputs = tf.keras.Input(shape=(actualImageDimension, actualImageDimension, 3))



x = preprocess_input(inputs)
x = base_model(x)
x = layers.GlobalAveragePooling2D()(x)
x = layers.BatchNormalization()(x)
x = layers.Dense(256, activation="relu")(x)
x = layers.Dropout(0.4)(x)

# End Dense Head
outputs = layers.Dense(1, activation="sigmoid")(x)
model = tf.keras.Model(inputs, outputs)


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

model.summary()

# -----------------------
# Callbacks
# -----------------------


early_stopping = EarlyStopping(
    monitor='val_loss',
    patience=15,
    restore_best_weights=True,
    verbose=1
)

checkpointer = ModelCheckpoint(
    filepath=modelname,
    monitor='val_accuracy',
    save_best_only=True,
    verbose=1
)

lr_scheduler = ReduceLROnPlateau(
    monitor='val_loss',
    factor=0.5,
    patience=5,
    min_lr=1e-6,
    verbose=1
)

# -----------------------
# Initial Training
# -----------------------

history = model.fit(
    train_ds,
    epochs=50,
    validation_data=val_ds,
    callbacks=[tensorboard, lr_scheduler, early_stopping, checkpointer]
)

# -----------------------
# Fine-Tuning
# -----------------------

base_model.trainable = True

# Freeze early layers (recommended)
for layer in base_model.layers[:-40]:
    layer.trainable = False

model.compile(
    optimizer=Adam(1e-5),
    loss='binary_crossentropy',
    metrics=['accuracy']
)

history_fine = model.fit(
    train_ds,
    epochs=50,
    validation_data=val_ds,
    callbacks=[tensorboard, lr_scheduler, early_stopping, checkpointer]
)