import os
import shutil
import random
import numpy as np
import json
from sklearn.model_selection import train_test_split
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator

SOURCE_DIR = r"E:\Dog_Breed_Classification_model\dataset"
TARGET_DIR = r"E:\Dog_Breed_Classification_model\dataset_split"
IMG_SIZE = (224, 224)
BATCH_SIZE = 16
EPOCHS = 20
SEED = 42

def split_dataset():
    random.seed(SEED)
    for split in ['train', 'val']:
        os.makedirs(os.path.join(TARGET_DIR, split), exist_ok=True)

    for breed in os.listdir(SOURCE_DIR):
        class_path = os.path.join(SOURCE_DIR, breed)
        if not os.path.isdir(class_path):
            continue
        images = [img for img in os.listdir(class_path) if img.lower().endswith(('.jpg', '.jpeg', '.png'))]
        train_imgs, val_imgs = train_test_split(images, test_size=0.2, random_state=SEED)

        for split, img_list in zip(['train', 'val'], [train_imgs, val_imgs]):
            split_dir = os.path.join(TARGET_DIR, split, breed)
            os.makedirs(split_dir, exist_ok=True)
            for img in img_list:
                shutil.copy2(os.path.join(class_path, img), os.path.join(split_dir, img))

def create_generators():
    train_gen = ImageDataGenerator(rescale=1./255, rotation_range=15, width_shift_range=0.1,
                                   height_shift_range=0.1, horizontal_flip=True)
    val_gen = ImageDataGenerator(rescale=1./255)

    train_data = train_gen.flow_from_directory(os.path.join(TARGET_DIR, 'train'),
                                               target_size=IMG_SIZE, batch_size=BATCH_SIZE,
                                               class_mode='categorical')

    val_data = val_gen.flow_from_directory(os.path.join(TARGET_DIR, 'val'),
                                           target_size=IMG_SIZE, batch_size=BATCH_SIZE,
                                           class_mode='categorical')
    return train_data, val_data

def build_model(num_classes):
    base_model = tf.keras.applications.MobileNetV2(input_shape=(*IMG_SIZE, 3), include_top=False, weights='imagenet')
    base_model.trainable = False
    model = tf.keras.Sequential([
        base_model,
        tf.keras.layers.GlobalAveragePooling2D(),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dropout(0.3),
        tf.keras.layers.Dense(num_classes, activation='softmax')
    ])
    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
    return model

def save_model(model, class_indices):
    model.save("dog_breed_classifier.h5")
    with open("class_indices.json", "w") as f:
        json.dump(class_indices, f)

if __name__ == "__main__":
    split_dataset()
    train_data, val_data = create_generators()
    model = build_model(num_classes=len(train_data.class_indices))
    model.fit(train_data, validation_data=val_data, epochs=EPOCHS)
    save_model(model, train_data.class_indices)