import gradio as gr
import pandas as pd
from PIL import Image

# Define your sections and labels
SECTION_LABELS = {
    "Oil Pore Related Issues": [
        "Very Large Pores (Not Red)",
        "Whiteheads (Clogged Pores)",
        "Blackheads (Clogged Pores)",
        "Shinny Skin",
        "Sebaceous Filaments (Sebum)"
    ],
    "Acne and Blemishes": [
        "Pustules",
        "Papules",
        "Nodules",
        "Cysts"
    ],
    "Redness and Irritation": [
        "Redness",
        "Irritation",
        "Itching",
        "Burning Sensation",
        "Allergic Reactions"
    ],
    "Dryness and Texture Issues": [
        "Dryness",
        "Fine Lines / Wrinkles",
        "Skin Flakes"
    ],
    "Aging and Elasticity Issues": [
        "Loose Skin",
        "Wrinkles"
    ],
    "Pigmentation Issues": [
        "Dark Spots",
        "Uneven Skin Tone",
        "Melasma",
        "Freckles"
    ]
}

# Flattened labels list
ALL_LABELS = [label for labels in SECTION_LABELS.values() for label in labels]

# Global state
images = []
current_index = 0
results = []
annotations = {}

# Core functions

def display_image(idx):
    if images:
        img = Image.open(images[idx])
        fname = images[idx].split('/')[-1]
        tick = ' ✅' if idx in annotations else ''
        caption = f"{fname} ({idx+1}/{len(images)}){tick}"
        states = annotations.get(idx, [False] * len(ALL_LABELS))
        return [img, caption] + states
    return [None, "No images uploaded"] + [False] * len(ALL_LABELS)


def navigate(delta):
    global current_index
    current_index = (current_index + delta) % len(images)
    return display_image(current_index)


def submit(*selections):
    if not images:
        # Return status and no file
        return "No image to label", None

    # Save selections
    annotations[current_index] = list(selections)
    fname = images[current_index].split('/')[-1]
    chosen = [lbl for lbl, sel in zip(ALL_LABELS, selections) if sel]
    global results
    results = [r for r in results if r['image'] != fname]
    results.append({'image': fname, 'labels': ', '.join(chosen)})

    # Write CSV
    df = pd.DataFrame(results)
    df.to_csv('image_labels.csv', index=False)

    # Return status message and CSV path
    return "Labels saved!", 'image_labels.csv'


def upload(files):
    global images, current_index, results, annotations
    images = [f.name if hasattr(f, 'name') else f for f in files]
    current_index = 0
    results = []
    annotations = {}
    outputs = display_image(0)
    # hide uploader after upload
    return outputs + [gr.update(visible=False)]


with gr.Blocks() as demo:
    gr.Markdown("## Dermo Annotator")
    file_upload = gr.File(label="Upload Images", file_count="multiple", file_types=["image"])

    # Create checkbox components in layout order
    checkbox_components = []

    # First row: first 3 sections
    with gr.Row():
        for section in list(SECTION_LABELS.keys())[:3]:
            with gr.Column():
                gr.Markdown(f"### {section}")
                for lbl in SECTION_LABELS[section]:
                    cb = gr.Checkbox(label=lbl)
                    checkbox_components.append(cb)

    # Second row: next 3 sections
    with gr.Row():
        for section in list(SECTION_LABELS.keys())[3:]:
            with gr.Column():
                gr.Markdown(f"### {section}")
                for lbl in SECTION_LABELS[section]:
                    cb = gr.Checkbox(label=lbl)
                    checkbox_components.append(cb)

    # Image display and controls
    with gr.Row():
        with gr.Column(scale=2):
            img = gr.Image(label="Image")
            caption = gr.Label(value="No images uploaded")
            with gr.Row():
                prev_btn = gr.Button("⬅️ Previous")
                next_btn = gr.Button("Next ➡️")

        with gr.Column(scale=1):
            submit_btn = gr.Button("Submit Labels")
            status = gr.Label()
            csv_downloader = gr.File(label="Download labels CSV")

    # Wire events
    file_upload.upload(
        fn=upload,
        inputs=file_upload,
        outputs=[img, caption] + checkbox_components + [file_upload]
    )

    prev_btn.click(
        fn=lambda: navigate(-1),
        outputs=[img, caption] + checkbox_components
    )

    next_btn.click(
        fn=lambda: navigate(1),
        outputs=[img, caption] + checkbox_components
    )

    submit_btn.click(
        fn=submit,
        inputs=checkbox_components,
        outputs=[status, csv_downloader]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)