# %% [code]
import os
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from PIL import Image
import numpy as np
from sklearn.model_selection import train_test_split
import random


# %% [code]
main_folder = "/kaggle/input/garbage-guru-challenge/dataset/train"  

image_data = []
for subfolder in os.listdir(main_folder):
    subfolder_path = os.path.join(main_folder, subfolder)
    if os.path.isdir(subfolder_path):
        for filename in os.listdir(subfolder_path):
            if filename.endswith(".jpg"):
                file_path = os.path.join(subfolder_path, filename)
                label = subfolder
                image_data.append((file_path, label))

# Create a Pandas DataFrame
df = pd.DataFrame(image_data, columns=['file_path', 'label'])

print("DataFrame with image paths and labels:")
print(df.head())
print(f"Number of images: {len(df)}")

# %% [code]
train_df = df  # All data in the 'train' folder will be used for training initially

val_data = []
val_folder = "/kaggle/input/garbage-guru-challenge/dataset/val"  
for subfolder in os.listdir(val_folder):
    subfolder_path = os.path.join(val_folder, subfolder)
    if os.path.isdir(subfolder_path):
        for filename in os.listdir(subfolder_path):
            if filename.endswith(".jpg"):
                file_path = os.path.join(subfolder_path, filename)
                label = subfolder
                val_data.append((file_path, label))
val_df = pd.DataFrame(val_data, columns=['file_path', 'label'])

test_data = []
test_folder = "/kaggle/input/garbage-guru-challenge/dataset/test"  
for subfolder in os.listdir(test_folder):
    subfolder_path = os.path.join(test_folder, subfolder)
    if os.path.isdir(subfolder_path):
        for filename in os.listdir(subfolder_path):
            if filename.endswith(".jpg"):
                file_path = os.path.join(subfolder_path, filename)
                label = subfolder
                test_data.append((file_path, label))
test_df = pd.DataFrame(test_data, columns=['file_path', 'label'])

print("\nTraining DataFrame:")
print(train_df.head())
print(f"Number of training images: {len(train_df)}")

print("\nValidation DataFrame:")
print(val_df.head())
print(f"Number of validation images: {len(val_df)}")

print("\nTest DataFrame:")
print(test_df.head())
print(f"Number of test images: {len(test_df)}")

# %% [code]
## Visualize the Class Distribution

# %% [code]
plt.figure(figsize=(10, 6))
sns.countplot(y=train_df['label'], order=train_df['label'].value_counts().index)
plt.title('Training Data Class Distribution')
plt.xlabel('Number of Images')
plt.ylabel('Waste Category')
plt.show()

plt.figure(figsize=(10, 6))
sns.countplot(y=val_df['label'], order=val_df['label'].value_counts().index)
plt.title('Validation Data Class Distribution')
plt.xlabel('Number of Images')
plt.ylabel('Waste Category')
plt.show()

plt.figure(figsize=(10, 6))
sns.countplot(y=test_df['label'], order=test_df['label'].value_counts().index)
plt.title('Test Data Class Distribution')
plt.xlabel('Number of Images')
plt.ylabel('Waste Category')
plt.show()

# %% [code]
##  Visualize Example Images

# %% [code]
def visualize_samples(dataframe, num_samples=5):
    """Visualizes a few random images from each category in the DataFrame."""
    categories = dataframe['label'].unique()
    for category in categories:
        category_df = dataframe[dataframe['label'] == category]
        sample_filenames = random.sample(category_df['file_path'].tolist(), min(num_samples, len(category_df)))

        print(f"\nCategory: {category}")
        plt.figure(figsize=(15, 3 * num_samples))
        for i, filename in enumerate(sample_filenames):
            try:
                img = Image.open(filename)
                plt.subplot(num_samples, len(categories), categories.tolist().index(category) + 1 + i * len(categories))
                plt.imshow(img)
                plt.title(filename.split('/')[-1])
                plt.axis('off')
            except FileNotFoundError:
                print(f"Error: File not found - {filename}")
        plt.show()

# Visualize samples from the training set
visualize_samples(train_df)

# %% [code]
## Basic Data Exploration:

# %% [code]
# Get information about the training DataFrame
print("Training DataFrame Info:")
train_df.info()

# Get value counts for the 'label' column in the training DataFrame
print("\nTraining Data Label Counts:")
print(train_df['label'].value_counts())

# Check for missing values in the training DataFrame
print("\nMissing values in Training Data:")
print(train_df.isnull().sum())

# Repeat for validation and test sets
print("\nValidation DataFrame Info:")
val_df.info()
print("\nValidation Data Label Counts:")
print(val_df['label'].value_counts())
print("\nMissing values in Validation Data:")
print(val_df.isnull().sum())

print("\nTest DataFrame Info:")
test_df.info()
print("\nTest Data Label Counts:")
print(test_df['label'].value_counts())
print("\nMissing values in Test Data:")
print(test_df.isnull().sum())

# %% [code]
## Analyzing Image Dimensions:

## Let's get the image dimensions and analyze them more systematically.

# %% [code]
def get_image_dimensions(file_path):
    try:
        img = Image.open(file_path)
        width, height = img.size
        return width, height
    except FileNotFoundError:
        return None, None

train_df[['width', 'height']] = train_df['file_path'].apply(get_image_dimensions).apply(pd.Series)
val_df[['width', 'height']] = val_df['file_path'].apply(get_image_dimensions).apply(pd.Series)
test_df[['width', 'height']] = test_df['file_path'].apply(get_image_dimensions).apply(pd.Series)

print("\nTraining Data Dimensions:")
print(train_df[['width', 'height']].describe())

print("\nValidation Data Dimensions:")
print(val_df[['width', 'height']].describe())

print("\nTest Data Dimensions:")
print(test_df[['width', 'height']].describe())

# Visualize dimension distributions
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
sns.histplot(train_df['width'].dropna(), bins=50, kde=True)
plt.title('Training Image Width Distribution')
plt.xlabel('Width')
plt.ylabel('Frequency')

plt.subplot(1, 2, 2)
sns.histplot(train_df['height'].dropna(), bins=50, kde=True)
plt.title('Training Image Height Distribution')
plt.xlabel('Height')
plt.ylabel('Frequency')

plt.show()

# %% [code]
## Analyzing Class Balance:

# %% [code]
# Calculate class imbalance ratio (ratio of the largest class count to the smallest)
train_counts = train_df['label'].value_counts()
imbalance_ratio_train = train_counts.max() / train_counts.min()
print(f"\nTraining Data Imbalance Ratio: {imbalance_ratio_train:.2f}")

val_counts = val_df['label'].value_counts()
imbalance_ratio_val = val_counts.max() / val_counts.min()
print(f"Validation Data Imbalance Ratio: {imbalance_ratio_val:.2f}")

test_counts = test_df['label'].value_counts()
imbalance_ratio_test = test_counts.max() / test_counts.min()
print(f"Test Data Imbalance Ratio: {imbalance_ratio_test:.2f}")

# %% [code]