# -- Code Cell --
import os
import numpy as np
import pandas as pd
import tensorflow as tf

from sklearn.model_selection import train_test_split

SEED = 42
IMG_SIZE = 224
BATCH_SIZE = 32
EPOCHS = 10
LR = 3e-4

IMAGES_DIR = "images"

train_df = pd.read_csv("train.csv")
test_df  = pd.read_csv("test.csv")

classes = sorted(train_df["label"].unique())
class_to_idx = {c:i for i,c in enumerate(classes)}
idx_to_class = {i:c for c,i in class_to_idx.items()}

train_df["label_idx"] = train_df["label"].map(class_to_idx).astype(int)

tr_df, va_df = train_test_split(
    train_df, test_size=0.2, random_state=SEED, stratify=train_df["label"]
)

print("Classes:", classes)
print("GPU:", tf.config.list_physical_devices("GPU"))


# -- Code Cell --
from tensorflow.keras.preprocessing.image import ImageDataGenerator

train_gen = ImageDataGenerator(preprocessing_function=preprocess, rotation_range=8, zoom_range=0.05, horizontal_flip=True)
val_gen   = ImageDataGenerator(preprocessing_function=preprocess)

train_flow = train_gen.flow_from_dataframe(
    tr_df,
    directory=IMAGES_DIR,
    x_col="image_path",
    y_col="label_idx",
    target_size=(IMG_SIZE, IMG_SIZE),
    class_mode="raw",          # returnează label ca int
    batch_size=BATCH_SIZE,
    shuffle=True,
    seed=SEED
)

val_flow = val_gen.flow_from_dataframe(
    va_df,
    directory=IMAGES_DIR,
    x_col="image_path",
    y_col="label_idx",
    target_size=(IMG_SIZE, IMG_SIZE),
    class_mode="raw",
    batch_size=BATCH_SIZE,
    shuffle=False
)

test_flow = val_gen.flow_from_dataframe(
    test_df,
    directory=IMAGES_DIR,
    x_col="image_path",
    y_col=None,
    target_size=(IMG_SIZE, IMG_SIZE),
    class_mode=None,           # fără label
    batch_size=BATCH_SIZE,
    shuffle=False
)


# -- Code Cell --
NUM_CLASSES = len(classes)

base = tf.keras.applications.ResNet50(
    include_top=False,
    weights="imagenet",
    input_shape=(IMG_SIZE, IMG_SIZE, 3)
)
preprocess = tf.keras.applications.resnet.preprocess_input
base.trainable = False

inputs = tf.keras.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
x = base(inputs, training=False)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.Dropout(0.2)(x)
outputs = tf.keras.layers.Dense(NUM_CLASSES, activation="softmax")(x)

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

model.compile(
    optimizer=tf.keras.optimizers.Adam(LR),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(),
    metrics=["accuracy"]
)

ckpt = tf.keras.callbacks.ModelCheckpoint(
    "best.keras", monitor="val_accuracy", mode="max", save_best_only=True
)

history = model.fit(
    train_flow,
    validation_data=val_flow,
    epochs=EPOCHS,
    callbacks=[ckpt]
)


# -- Code Cell --
base.trainable = True

# îngheață primele ~70% straturi
fine_tune_at = int(len(base.layers) * 0.7)
for layer in base.layers[:fine_tune_at]:
    layer.trainable = False

model.compile(
    optimizer=tf.keras.optimizers.Adam(LR * 0.1),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(),
    metrics=["accuracy"]
)

history_ft = model.fit(
    train_flow,
    validation_data=val_flow,
    epochs=max(5, EPOCHS//2),
    callbacks=[ckpt]
)


# -- Code Cell --
best = tf.keras.models.load_model("best.keras")

probs = best.predict(test_flow, verbose=1)
pred_idx = probs.argmax(axis=1)
pred_labels = [idx_to_class[i] for i in pred_idx]

submission = pd.DataFrame({
    "id": test_df["id"].values,
    "label": pred_labels
})
submission.to_csv("submission2.csv", index=False)
submission.head()
