import os
import numpy as np
from pathlib import Path
import torch
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms
import tensorflow as tf



IMAGE_SIZE = (150, 150)   
BATCH_SIZE = 32
VAL_SPLIT = 0.2        
SEED = 42

CLASS_NAMES = ["buildings", "forest", "glacier", "mountain", "sea", "street"]



#  PyTorch  Data Pipeline
def get_pytorch_loaders(train_dir: str, test_dir: str):
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]

    
    train_transform = transforms.Compose([
        transforms.Resize(IMAGE_SIZE),
        transforms.RandomHorizontalFlip(p=0.5),          
        transforms.RandomRotation(degrees=15),          
        transforms.RandomResizedCrop(IMAGE_SIZE,scale=(0.8, 1.0)),
        transforms.ColorJitter(brightness=0.2,contrast=0.2),
        transforms.ToTensor(),                          
        transforms.Normalize(mean, std),   
    ])

    eval_transform = transforms.Compose([
        transforms.Resize(IMAGE_SIZE),
        transforms.ToTensor(),
        transforms.Normalize(mean, std),
    ])

    full_train = datasets.ImageFolder(root=train_dir,transform=train_transform)

    # Split into train / validation
    n_val = int(len(full_train) * VAL_SPLIT)
    n_train = len(full_train) - n_val
    train_ds, val_ds = random_split(
        full_train, [n_train, n_val],
        generator=torch.Generator().manual_seed(SEED)
    )

    
    val_ds.dataset = datasets.ImageFolder(root=train_dir,transform=eval_transform)
    test_ds = datasets.ImageFolder(root=test_dir, transform=eval_transform)

    
    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True)
    val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True)
    test_loader  = DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True)

    print(f"[PyTorch] Train: {n_train} | Val: {n_val} | Test: {len(test_ds)}")
    return train_loader, val_loader, test_loader, full_train.classes





#  TensorFlow  Data Pipeline
def get_tensorflow_datasets(train_dir: str, test_dir: str):
    img_h, img_w = IMAGE_SIZE

    raw_train = tf.keras.utils.image_dataset_from_directory(
        train_dir,
        validation_split=VAL_SPLIT,
        subset="training",
        seed=SEED,
        image_size=IMAGE_SIZE,
        batch_size=BATCH_SIZE,
        label_mode="categorical", 
    )

    raw_val = tf.keras.utils.image_dataset_from_directory(
        train_dir,
        validation_split=VAL_SPLIT,
        subset="validation",
        seed=SEED,
        image_size=IMAGE_SIZE,
        batch_size=BATCH_SIZE,
        label_mode="categorical",
    )

    raw_test = tf.keras.utils.image_dataset_from_directory(
        test_dir,
        image_size=IMAGE_SIZE,
        batch_size=BATCH_SIZE,
        label_mode="categorical",
        shuffle=False,
    )
    

    class_names = raw_train.class_names

    
    normalization = tf.keras.layers.Rescaling(1.0 / 255)
    augmentation = tf.keras.Sequential([
        tf.keras.layers.RandomFlip("horizontal"),
        tf.keras.layers.RandomRotation(0.1),       
        tf.keras.layers.RandomZoom(0.2),           
        tf.keras.layers.RandomContrast(0.1),
    ])

    def preprocess_train(images, labels):
        images = normalization(images)
        images = augmentation(images, training=True)
        return images, labels

    def preprocess_eval(images, labels):
        images = normalization(images)
        return images, labels

    AUTOTUNE = tf.data.AUTOTUNE

    train_ds = (raw_train
                .map(preprocess_train, num_parallel_calls=AUTOTUNE)
                .cache()
                .shuffle(1000)
                .prefetch(AUTOTUNE))

    val_ds = (raw_val
              .map(preprocess_eval, num_parallel_calls=AUTOTUNE)
              .cache()
              .prefetch(AUTOTUNE))

    test_ds = (raw_test
               .map(preprocess_eval, num_parallel_calls=AUTOTUNE)
               .prefetch(AUTOTUNE))

    print(f"[TensorFlow] Classes: {class_names}")
    return train_ds, val_ds, test_ds, class_names