from fastai.learner import load_learner
from fastai.vision.core import PILImage
from fastai.vision.all import DisplayedTransform
from pathlib import Path
import gradio as gr
import PIL
import sys
import os
import albumentations as A
import numpy as np

# Print debug information
print(f"Python version: {sys.version}")
print(f"PIL version: {PIL.__version__}")
print(f"fastai path: {Path(__file__).parent}")
print(f"Current working directory: {os.getcwd()}")

# Create a custom transform class that integrates Albumentations with fastai
class AlbumentationsTransform(DisplayedTransform):
    split_idx, order = None, 2  # Apply to both train and validation sets
    
    def __init__(self, train_aug, valid_aug=None):
        self.train_aug = train_aug
        self.valid_aug = valid_aug or A.Compose([
            A.SmallestMaxSize(max_size=224),
            A.CenterCrop(height=224, width=224)
        ])
        
    def before_call(self, b, split_idx):
        self.split_idx = split_idx
        
    def encodes(self, img: PILImage):
        # Use validation augmentation for inference
        aug = self.valid_aug
        # Convert to numpy array and ensure it's uint8 before augmentation
        img_array = np.array(img).astype(np.uint8)
        aug_img = aug(image=img_array)['image']
        # Convert back to PIL Image
        return PILImage.create(aug_img)

# Define the albumentations transforms
def get_train_transform():
    return A.Compose([
        A.SmallestMaxSize(max_size=256),
        A.RandomCrop(height=224, width=224),
        A.HorizontalFlip(p=0.5),
        A.RandomBrightnessContrast(p=0.3),
    ])

def get_valid_transform():
    return A.Compose([
        A.SmallestMaxSize(max_size=256),
        A.CenterCrop(height=224, width=224),
    ])

try:
    learn = load_learner('export.pkl')
    print(f"Model loaded successfully: {type(learn)}")
    
    # Define friendly label names
    label_mapping = {
        'positive': "Slate",
        'negative': "No Slate"
    }
    
    # Get original labels
    original_labels = learn.dls.vocab
    print(f"Original labels: {original_labels}")
    
    def predict(img):
        # Convert to PILImage and apply transforms
        img_pil = PILImage.create(img)
        # Create transform with both train and valid transforms (will use valid for inference)
        transform = AlbumentationsTransform(get_train_transform(), get_valid_transform())
        transformed_img = transform.encodes(img_pil)
        
        # Get prediction
        pred, pred_idx, probs = learn.predict(transformed_img)
        
        # Map the original labels to friendly names
        return {label_mapping[original_labels[i]]: float(probs[i]) for i in range(len(original_labels))}

    # Create a more attractive interface with custom styling
    with gr.Blocks(css="footer {visibility: hidden}") as demo:
        gr.Markdown("# 🎬 Image Classifier: Slate or No Slate?")
        gr.Markdown("Upload an image and let AI tell you whether or not it contains a film slate!")
        
        with gr.Row():
            with gr.Column(scale=1):
                # Remove webcam by setting sources to only 'upload' and 'clipboard'
                input_image = gr.Image(
                    type="pil",
                    label="Upload Image",
                    sources=["upload", "clipboard"]
                )
                submit_btn = gr.Button("Classify", variant="primary")
            
            with gr.Column(scale=1):
                output = gr.Label(num_top_classes=3, label="Predictions")
                
        with gr.Accordion("About", open=False):
            gr.Markdown("""
            ## How it works
            This app uses a deep learning model trained with fastai to classify images.
            
            ## Tips for best results
            - Use clear, well-lit images
            - Center the subject in the frame
            - Supported categories: """ + ", ".join([label_mapping[label] for label in original_labels]))
        
        # Set up the prediction flow
        submit_btn.click(
            fn=predict,
            inputs=input_image,
            outputs=output
        )

        # Allow image input to trigger prediction as well
        input_image.change(
            fn=predict,
            inputs=input_image,
            outputs=output
        )

        # Use all images in the examples folder as examples
        examples_dir = Path("examples").resolve()
        example_paths = [str(p) for p in examples_dir.glob("*") if p.is_file()]

        gr.Examples(
           examples=example_paths,
           inputs=input_image
        )

        allowed_paths = [str(examples_dir)]

        # Launch the app, allow access to the entire examples directory for @Gradio
        demo.launch(allowed_paths=allowed_paths) # ssr_mode=False
except Exception as e:
    print(f"Error occurred: {e}")
    import traceback
    traceback.print_exc()