# -- Code Cell --
# ONE CELL: load CSVs -> train/valid tf.data -> label encoding -> ready for model.fit + submission template

import os, numpy as np, pandas as pd, tensorflow as tf

# --- paths (adjust only if your files are elsewhere) ---
TRAIN_CSV = "train.csv"
TEST_CSV  = "test.csv"
IMAGES_DIR = "images"  # images/ directory

# --- load csvs ---
train_df = pd.read_csv(TRAIN_CSV)
test_df  = pd.read_csv(TEST_CSV)

# --- make absolute/normalized paths (works whether image_path already contains 'images/' or not) ---
def fix_path(p):
    p = str(p).replace("\\", "/")
    if p.startswith(IMAGES_DIR + "/") or p.startswith(IMAGES_DIR):
        return p
    return f"{IMAGES_DIR}/{p}"

train_df["image_path"] = train_df["image_path"].apply(fix_path)
test_df["image_path"]  = test_df["image_path"].apply(fix_path)

# --- label encoding (string -> int) + mappings for later decoding ---
class_names = sorted(train_df["label"].unique().tolist())  # ['bishop','knight','pawn','queen','rook']
name2id = {n:i for i,n in enumerate(class_names)}
id2name = {i:n for n,i in name2id.items()}
train_df["label_id"] = train_df["label"].map(name2id).astype(int)

print("Classes:", class_names)
print("Train size:", len(train_df), "| Test size:", len(test_df))

# --- split train -> train/valid (like Fashion-MNIST style) ---
VALID_SIZE = 5000 if len(train_df) > 5000 else max(1, int(0.1 * len(train_df)))
train_df = train_df.sample(frac=1, random_state=42).reset_index(drop=True)
valid_df = train_df.iloc[:VALID_SIZE].copy()
train_df2 = train_df.iloc[VALID_SIZE:].copy()

# --- tf.data image loader ---
IMG_SIZE = (224, 224)
BATCH = 32
AUTOTUNE = tf.data.AUTOTUNE

def decode_img(path):
    img = tf.io.read_file(path)
    img = tf.image.decode_image(img, channels=3, expand_animations=False)
    img = tf.image.resize(img, IMG_SIZE)
    img = tf.cast(img, tf.float32) / 255.0
    return img

def load_xy(path, y):
    return decode_img(path), y

def make_ds(df, training=True):
    paths = df["image_path"].values.astype(str)
    ys = df["label_id"].values.astype(np.int32)
    ds = tf.data.Dataset.from_tensor_slices((paths, ys)).map(load_xy, num_parallel_calls=AUTOTUNE)
    if training:
        ds = ds.shuffle(min(2000, len(df)), seed=42, reshuffle_each_iteration=True)
    return ds.batch(BATCH).prefetch(AUTOTUNE)

train_ds = make_ds(train_df2, training=True)
valid_ds = make_ds(valid_df, training=False)

# test dataset (no labels)
test_paths = test_df["image_path"].values.astype(str)
test_ds = tf.data.Dataset.from_tensor_slices(test_paths).map(decode_img, num_parallel_calls=AUTOTUNE).batch(BATCH).prefetch(AUTOTUNE)

# --- what you "need" now ---
# train_ds, valid_ds, test_ds, class_names, name2id, id2name, train_df2, valid_df, test_df

# --- OPTIONAL: submission helper (after you have a trained `model`) ---
# preds = model.predict(test_ds)                       # (N, num_classes)
# pred_ids = preds.argmax(axis=1)                      # (N,)
# pred_labels = [id2name[i] for i in pred_ids]
# submission = pd.DataFrame({"id": test_df["id"], "label": pred_labels})
# submission.to_csv("submission.csv", index=False)
# print("Wrote submission.csv")


# -- Code Cell --
X_train = train_df2["image_path"].values
y_train = train_df2["label_id"].values

X_valid = valid_df["image_path"].values
y_valid = valid_df["label_id"].values

X_test  = test_df["image_path"].values


# -- Code Cell --
class_names[y_train[1]]

# -- Code Cell --
print(IMG_SIZE)
print(model.input_shape)

# -- Code Cell --
tf.random.set_seed(42)
model = tf.keras.Sequential()
model.add(tf.keras.layers.Input(shape=[224, 224, 3]))
model.add(tf.keras.layers.Flatten())
model.add(tf.keras.layers.Dense(300, activation="relu"))
model.add(tf.keras.layers.Dense(100, activation="relu"))
model.add(tf.keras.layers.Dense(5, activation="softmax"))

# -- Code Cell --
model.compile(loss="sparse_categorical_crossentropy",
optimizer="sgd",
metrics=["accuracy"])

# -- Code Cell --
history = model.fit(
    train_ds,
    epochs=30,
    validation_data=valid_ds
)

# -- Code Cell --
NUM_CLASSES = len(class_names)
INPUT_SHAPE = (*IMG_SIZE, 3)

base = tf.keras.applications.EfficientNetB0(
    include_top=False,
    weights="imagenet",
    input_shape=INPUT_SHAPE
)
base.trainable = False

model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=INPUT_SHAPE),
    base,
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dropout(0.3),
    tf.keras.layers.Dense(NUM_CLASSES, activation="softmax")
])

model.compile(
    optimizer=tf.keras.optimizers.Adam(1e-3),
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"]
)

history = model.fit(train_ds, epochs=10, validation_data=valid_ds)

# fine-tuning (aproape mereu ridică accuracy mult)
base.trainable = True
model.compile(
    optimizer=tf.keras.optimizers.Adam(1e-5),
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"]
)

history2 = model.fit(train_ds, epochs=10, validation_data=valid_ds)


# -- Code Cell --
train_df["label"].value_counts()