from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from PIL import Image

# ============================================================
# ΡΥΘΜΙΣΕΙΣ
# ============================================================

# Βάλε εδώ τη διαδρομή του φακέλου όπου αποσυμπίεσες το dataset.
DATASET_DIR = Path(r"C:\Users\anpen\Downloads\archive")

# Επιλογές: "auto", "images", "csv"
DATA_FORMAT = "auto"

# Μέγεθος εικόνας MNIST
IMAGE_SIZE = (20, 20)

# Κανονικοποίηση pixels από [0, 255] σε [0, 1]
NORMALIZE = True

# Εμφάνιση παραδειγμάτων και στατιστικών
SHOW_PLOTS = True

# ============================================================
# ΒΟΗΘΗΤΙΚΕΣ ΣΥΝΑΡΤΗΣΕΙΣ
# ============================================================

IMAGE_EXTENSIONS = {
    ".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"
}


def find_label(folder: Path, root: Path):
    """
    Εντοπίζει την ετικέτα 0-9 από το όνομα ενός φακέλου.

    Υποστηρίζει δομές όπως:
        MNIST/0/image.png
        MNIST/train/0/image.png
        MNIST/test/7/image.png
    """
    for parent in folder.parents:
        if parent == root.parent:
            break

        if parent.name.isdigit():
            label = int(parent.name)
            if 0 <= label <= 9:
                return label

        if parent == root:
            break

    if folder.name.isdigit():
        label = int(folder.name)
        if 0 <= label <= 9:
            return label

    return None


def load_image_dataset(root: Path):
    """Διαβάζει εικόνες από υποφακέλους με ονόματα 0-9."""
    images = []
    labels = []
    skipped = 0

    files = sorted(
        p for p in root.rglob("*")
        if p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS
    )

    if not files:
        raise FileNotFoundError(
            f"Δεν βρέθηκαν εικόνες στον φάκελο: {root}"
        )

    for file_path in files:
        # Η ετικέτα πρέπει να προκύπτει από φάκελο 0-9.
        label = None

        for parent in file_path.parents:
            if parent == root.parent:
                break

            if parent.name.isdigit():
                value = int(parent.name)
                if 0 <= value <= 9:
                    label = value
                    break

            if parent == root:
                break

        if label is None:
            skipped += 1
            continue

        try:
            with Image.open(file_path) as img:
                img = img.convert("L")
                img = img.resize(IMAGE_SIZE)
                array = np.asarray(img, dtype=np.uint8)

            images.append(array)
            labels.append(label)

        except Exception as exc:
            print(f"Παράλειψη {file_path}: {exc}")
            skipped += 1

    if not images:
        raise ValueError(
            "Δεν βρέθηκαν έγκυρες εικόνες με labels 0-9. "
            "Έλεγξε τη δομή των φακέλων."
        )

    X = np.stack(images)
    y = np.asarray(labels, dtype=np.int64)

    print(f"Εικόνες που διαβάστηκαν: {len(X)}")
    print(f"Αρχεία που παραλείφθηκαν: {skipped}")

    return X, y


def load_csv_dataset(root: Path):
    """
    Διαβάζει CSV αρχεία με labels και pixels.

    Υποστηρίζει:
      - Στήλη label, digit, class ή target.
      - Μορφή με 784 στήλες pixels για εικόνες 28x28.
      - Ένα ή περισσότερα CSV αρχεία σε υποφακέλους.
    """
    csv_files = sorted(root.rglob("*.csv"))

    if not csv_files:
        raise FileNotFoundError(
            f"Δεν βρέθηκαν CSV αρχεία στον φάκελο: {root}"
        )

    all_X = []
    all_y = []

    label_names = {"label", "digit", "class", "target"}

    for csv_path in csv_files:
        print(f"Ανάγνωση CSV: {csv_path}")

        df = pd.read_csv(csv_path)

        if df.empty:
            print("  Το αρχείο είναι κενό.")
            continue

        # Εντοπισμός της στήλης ετικέτας.
        label_col = next(
            (
                col for col in df.columns
                if str(col).strip().lower() in label_names
            ),
            None
        )

        # Υποστήριξη ονομάτων στηλών όπως pixel0, pixel1, ...
        if label_col is None:
            label_col = next(
                (
                    col for col in df.columns
                    if str(col).strip().lower() in {
                        "label", "digit", "class", "target"
                    }
                ),
                None
            )

        # Εναλλακτικά, αναγνώριση της πρώτης στήλης ως label,
        # εφόσον οι υπόλοιπες στήλες περιέχουν pixels.
        if label_col is None:
            first_col = df.columns[0]
            first_values = pd.to_numeric(
                df[first_col], errors="coerce"
            )

            if (
                first_values.notna().all()
                and first_values.between(0, 9).all()
                and len(df.columns) > 1
            ):
                label_col = first_col

        if label_col is None:
            raise ValueError(
                f"Δεν εντοπίστηκε στήλη label στο {csv_path}. "
                "Χρειάζεται στήλη label/digit/class/target "
                "ή πρώτη στήλη με τιμές 0-9."
            )

        y = pd.to_numeric(
            df[label_col], errors="coerce"
        ).to_numpy()

        pixel_df = df.drop(columns=[label_col]).apply(
            pd.to_numeric, errors="coerce"
        )

        # Αφαιρούμε στήλες που δεν περιέχουν αριθμητικά pixels.
        pixel_df = pixel_df.dropna(axis=1, how="all")

        if pixel_df.shape[1] == 0:
            raise ValueError(
                f"Δεν βρέθηκαν στήλες pixels στο {csv_path}."
            )

        X = pixel_df.to_numpy(dtype=np.float32)

        # Διατηρούμε μόνο έγκυρα labels 0-9 και γραμμές
        # που περιέχουν αριθμητικές τιμές σε όλα τα pixels.
        valid = (
            np.isfinite(y)
            & (y >= 0)
            & (y <= 9)
            & (y == np.floor(y))
            & np.isfinite(X).all(axis=1)
        )

        X = X[valid]
        y = y[valid].astype(np.int64)

        # Μετατροπή 400 pixels σε εικόνες 20x20.
        if X.shape[1] == 400:
            X = X.reshape(-1, 20, 20)

        # Εναλλακτικά, υποστήριξη 785 στηλών συνολικά
        # έχει ήδη καλυφθεί από την αφαίρεση της στήλης label.
        elif X.shape[1] != 20 * 20:
            raise ValueError(
                f"Το {csv_path} έχει {X.shape[1]} στήλες pixels. "
                "Το συγκεκριμένο script περιμένει 400 pixels "
                "ανά εικόνα (20x20)."
            )

        all_X.append(X)
        all_y.append(y)

    if not all_X:
        raise ValueError("Δεν διαβάστηκαν έγκυρα δεδομένα από CSV.")

    X = np.concatenate(all_X, axis=0)
    y = np.concatenate(all_y, axis=0)

    print(f"Δείγματα που διαβάστηκαν: {len(X)}")

    return X, y


def prepare_data(X, normalize=True):
    """Μετατρέπει τις εικόνες σε float32 και κανονικοποιεί."""
    X = X.astype(np.float32)

    if normalize:
        # Για συνηθισμένα pixels 0-255.
        if X.max() > 1.0:
            X /= 255.0

    return X


def print_statistics(X, y):
    """Εμφανίζει βασικά στατιστικά του dataset."""
    print("\n" + "=" * 45)
    print("ΣΤΑΤΙΣΤΙΚΑ DATASET")
    print("=" * 45)

    print(f"Πλήθος εικόνων: {len(X)}")
    print(f"Διαστάσεις X: {X.shape}")
    print(f"Διαστάσεις y: {y.shape}")
    print(f"Τύπος δεδομένων X: {X.dtype}")
    print(f"Ελάχιστη τιμή pixel: {X.min():.4f}")
    print(f"Μέγιστη τιμή pixel: {X.max():.4f}")

    print("\nΚατανομή ψηφίων:")

    counts = np.bincount(y, minlength=10)

    for digit, count in enumerate(counts):
        print(f"Ψηφίο {digit}: {count} εικόνες")


def show_samples(X, y, n=10):
    """Εμφανίζει έως n εικόνες με τις ετικέτες τους."""
    if len(X) == 0:
        return

    n = min(n, len(X))

    fig, axes = plt.subplots(
        1, n, figsize=(2 * n, 2.5), squeeze=False
    )

    for i in range(n):
        axes[0, i].imshow(X[i], cmap="gray")
        axes[0, i].set_title(f"Label: {y[i]}")
        axes[0, i].axis("off")

    plt.tight_layout()
    plt.show()


def show_distribution(y):
    """Γράφημα κατανομής των ψηφίων 0-9."""
    counts = np.bincount(y, minlength=10)

    plt.figure(figsize=(8, 4))
    plt.bar(range(10), counts, color="steelblue")
    plt.xticks(range(10))
    plt.xlabel("Ψηφίο")
    plt.ylabel("Πλήθος εικόνων")
    plt.title("Κατανομή ψηφίων στο MNIST")
    plt.grid(axis="y", alpha=0.3)
    plt.tight_layout()
    plt.show()


# ============================================================
# ΚΥΡΙΟ ΠΡΟΓΡΑΜΜΑ
# ============================================================

def main():
    if not DATASET_DIR.exists():
        raise FileNotFoundError(
            f"Ο φάκελος dataset δεν υπάρχει: {DATASET_DIR}\n"
            "Άλλαξε την τιμή της μεταβλητής DATASET_DIR."
        )

    print(f"Φάκελος dataset: {DATASET_DIR}")
    print(f"Μορφή δεδομένων: {DATA_FORMAT}")

    if DATA_FORMAT == "images":
        X, y = load_image_dataset(DATASET_DIR)

    elif DATA_FORMAT == "csv":
        X, y = load_csv_dataset(DATASET_DIR)

    elif DATA_FORMAT == "auto":
        # Προτεραιότητα σε εικόνες, εφόσον υπάρχουν.
        has_images = any(
            p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS
            for p in DATASET_DIR.rglob("*")
        )

        if has_images:
            X, y = load_image_dataset(DATASET_DIR)
        else:
            X, y = load_csv_dataset(DATASET_DIR)

    else:
        raise ValueError(
            "Το DATA_FORMAT πρέπει να είναι "
            "'auto', 'images' ή 'csv'."
        )

    X = prepare_data(X, normalize=NORMALIZE)

    print_statistics(X, y)

    # Οι παρακάτω μεταβλητές είναι πλέον έτοιμες για ML.
    # X: εικόνες, y: labels 0-9.

    if SHOW_PLOTS:
        show_samples(X, y, n=10)
        show_distribution(y)

    print("\nΤο dataset φορτώθηκε επιτυχώς!")


if __name__ == "__main__":
    main()