# Install required libraries
!pip install torch torchvision timm matplotlib

# Import necessary libraries
import os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, models
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt

# Initialize lists to store image paths and labels
image_paths = []
labels = []

# Iterate through each gender, age group, and foot type
for gender in ['male', 'female']:
    for age_group in ['21-30', '31-40', '41-50', '51-60']:
        for foot_type in ['left', 'right', 'both']:
            folder_path = os.path.join(data_dir, 'gender', gender)
            if os.path.exists(folder_path):
                for img_file in os.listdir(folder_path):
                    if img_file.endswith(('.jpg', '.png')):
                        image_paths.append(os.path.join(folder_path, img_file))
                        gender_label = 1 if gender == 'male' else 0
                        age_group_label = ['21-30', '31-40', '41-50', '51-60'].index(age_group)
                        foot_type_label = {'left': 0, 'right': 1, 'both': 2}[foot_type]
                        labels.append((gender_label, age_group_label, foot_type_label))

# Convert lists to numpy arrays
image_paths = np.array(image_paths)
labels = np.array(labels)

# Custom dataset class
class FootDataset(Dataset):
    def __init__(self, image_paths, labels, transform=None):
        self.image_paths = image_paths
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img = Image.open(self.image_paths[idx]).convert('L')  # Convert to grayscale
        if self.transform:
            img = self.transform(img)
        label = torch.tensor(self.labels[idx], dtype=torch.long)
        return img, label

# Define transformations with channel replication for Inception v3
transform = transforms.Compose([
    transforms.Resize((299, 299)),
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.repeat(3, 1, 1)),  # Replicate grayscale channel 3 times
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# Create dataset and DataLoader
dataset = FootDataset(image_paths, labels, transform=transform)
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# Load the Inception v3 model
model = models.inception_v3(weights=models.Inception_V3_Weights.DEFAULT, aux_logits=True)

# Modify the classifier for multi-label output
num_features = model.fc.in_features
model.fc = nn.Identity()

# Define new output heads for each task
model.fc_gender = nn.Linear(num_features, 2)
model.fc_age_group = nn.Linear(num_features, 4)
model.fc_foot_type = nn.Linear(num_features, 3)

# Move the model to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Define loss functions and optimizer
criterion_gender = nn.CrossEntropyLoss()
criterion_age_group = nn.CrossEntropyLoss()
criterion_foot_type = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Initialize lists to track loss for plotting
losses = []

# Forward pass function
def forward_pass(images):
    outputs = model(images)
    features = outputs.logits if isinstance(outputs, tuple) else outputs
    outputs_gender = model.fc_gender(features)
    outputs_age_group = model.fc_age_group(features)
    outputs_foot_type = model.fc_foot_type(features)
    return outputs_gender, outputs_age_group, outputs_foot_type

# Training loop
num_epochs = 50
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0

    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()

        outputs_gender, outputs_age_group, outputs_foot_type = forward_pass(images)

        loss_gender = criterion_gender(outputs_gender, labels[:, 0])
        loss_age_group = criterion_age_group(outputs_age_group, labels[:, 1])
        loss_foot_type = criterion_foot_type(outputs_foot_type, labels[:, 2])

        loss = loss_gender + loss_age_group + loss_foot_type

        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    epoch_loss = running_loss / len(train_loader)
    losses.append(epoch_loss)  # Track loss
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}')

# Plot the loss graph
plt.figure(figsize=(10, 5))
plt.plot(range(1, num_epochs + 1), losses, marker='o', linestyle='-', color='b', label='Training Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Training Loss over Epochs')
plt.legend()
plt.grid(True)
plt.show()

# Testing loop
model.eval()
correct_gender = correct_age_group = correct_foot_type = total = 0

with torch.no_grad():
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)
        outputs_gender, outputs_age_group, outputs_foot_type = forward_pass(images)

        _, predicted_gender = torch.max(outputs_gender, 1)
        _, predicted_age_group = torch.max(outputs_age_group, 1)
        _, predicted_foot_type = torch.max(outputs_foot_type, 1)

        correct_gender += (predicted_gender == labels[:, 0]).sum().item()
        correct_age_group += (predicted_age_group == labels[:, 1]).sum().item()
        correct_foot_type += (predicted_foot_type == labels[:, 2]).sum().item()
        total += labels.size(0)

print(f'Gender Accuracy: {100 * correct_gender / total:.2f}%')
print(f'Age Group Accuracy: {100 * correct_age_group / total:.2f}%')
print(f'Foot Type Accuracy: {100 * correct_foot_type / total:.2f}%')