import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split, WeightedRandomSampler
from torchvision import datasets, transforms
from PIL import Image, ImageDraw
import os
import argparse

from model import (CropLowerRightTriangle, GarageDoorCNN, TRIANGLE_CROP_WIDTH,
                   TRIANGLE_CROP_HEIGHT, RESIZE_DIM)


def check_crop():
    """Saves a sample cropped image for debugging purposes and exits."""
    # Find a sample image from your dataset
    SAMPLE_IMAGE_DIR = 'data/labelled/open'
    if not os.path.isdir(SAMPLE_IMAGE_DIR) or not os.listdir(SAMPLE_IMAGE_DIR):
        print(f"Error: Cannot find sample image in '{SAMPLE_IMAGE_DIR}'.")
        print("Please ensure the directory exists and contains images.")
        return

    sample_image_name = os.listdir(SAMPLE_IMAGE_DIR)[0]
    sample_image_path = os.path.join(SAMPLE_IMAGE_DIR, sample_image_name)

    print(f"Creating debug crop from image: {sample_image_path}")

    # Load the image
    image = Image.open(sample_image_path)

    # Create the transform
    cropper = CropLowerRightTriangle(triangle_width=TRIANGLE_CROP_WIDTH, triangle_height=TRIANGLE_CROP_HEIGHT)

    # Apply the transform
    cropped_image = cropper(image)

    # Save the result
    output_path = "cropped_debug_output.png"
    cropped_image.save(output_path)
    print(f"Debug image saved to '{output_path}'.")


def train_model():
    # --- Hyperparameters and Configuration ---
    DATA_DIR = 'data/labelled'
    MODEL_SAVE_PATH = 'garage_door_cnn.pth'
    NUM_EPOCHS = 10
    BATCH_SIZE = 32
    LEARNING_RATE = 0.001

    # --- Data Preparation ---
    # Define transforms
    data_transforms = transforms.Compose([
        CropLowerRightTriangle(triangle_width=TRIANGLE_CROP_WIDTH, triangle_height=TRIANGLE_CROP_HEIGHT),
        transforms.Resize((RESIZE_DIM, RESIZE_DIM)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    # Load dataset with ImageFolder
    full_dataset = datasets.ImageFolder(DATA_DIR, transform=data_transforms)
    
    # Split into training and validation sets
    train_size = int(0.8 * len(full_dataset))
    val_size = len(full_dataset) - train_size
    train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size])

    # --- Handle Class Imbalance ---
    # Get labels for training set
    train_labels = [full_dataset.targets[i] for i in train_dataset.indices]
    
    # Get class counts
    class_counts = torch.bincount(torch.tensor(train_labels))
    
    # Compute weight for each class (inverse of count)
    class_weights = 1. / class_counts.float()
    
    # Assign a weight to each sample in the training set
    sample_weights = torch.tensor([class_weights[label] for label in train_labels])
    
    # Create a WeightedRandomSampler to balance the classes during training
    sampler = WeightedRandomSampler(weights=sample_weights, num_samples=len(sample_weights), replacement=True)

    # The sampler will handle shuffling, so shuffle must be False for the DataLoader
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, sampler=sampler)
    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)

    # --- Model, Loss, Optimizer ---
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    model = GarageDoorCNN(resize_dim=RESIZE_DIM).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

    # --- Training Loop ---
    print("Starting training...")
    for epoch in range(NUM_EPOCHS):
        model.train()
        running_loss = 0.0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * inputs.size(0)

        epoch_loss = running_loss / len(train_dataset)
        print(f"Epoch {epoch+1}/{NUM_EPOCHS}, Training Loss: {epoch_loss:.4f}")

        # --- Validation Loop ---
        model.eval()
        val_loss = 0.0
        corrects = 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                val_loss += loss.item() * inputs.size(0)
                _, preds = torch.max(outputs, 1)
                corrects += torch.sum(preds == labels.data)

        val_epoch_loss = val_loss / len(val_dataset)
        val_epoch_acc = corrects.double() / len(val_dataset)
        print(f"Validation Loss: {val_epoch_loss:.4f}, Accuracy: {val_epoch_acc:.4f}")

    # --- Save the trained model ---
    torch.save(model.state_dict(), MODEL_SAVE_PATH)
    print(f"Model saved to {MODEL_SAVE_PATH}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Train a CNN for garage door detection or check the image crop.")
    parser.add_argument('--check-crop', action='store_true', help='Save a sample cropped image and exit.')
    args = parser.parse_args()

    if args.check_crop:
        check_crop()
    else:
        # Check if data directory exists
        if not os.path.isdir('data/labelled/open') or not os.path.isdir('data/labelled/closed'):
            print("Error: Data directories 'data/open' and 'data/closed' not found.")
            print("Please create them and place your image snapshots inside.")
        else:
            train_model()