import spaces
import gradio as gr
import torch
import numpy as np
import cv2
from PIL import Image
from pathlib import Path
from torchvision.transforms import transforms

from models import DNSteerableLeNet, VanillaLeNet
from gradcam import GradCAM_ECNN


CLASS_NAMES  = ["FRI", "FRII", "Compact", "Bent"]
IMSIZE       = 150
DEVICE       = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
DATAMEAN     = 0.0005
DATASTD      = 0.0142

MODEL_PATHS = {
    "DNSteerableLeNet (G-CNN)": "models/FIRST_dnlenet.pt",
    "VanillaLeNet (baseline)": "models/FIRST_lenet.pt",
}

CLASS_DESC = {
    "FRI": (
        "**FR-I** — Jets are brightest near the core and fade outward. "
        "Typically associated with lower radio luminosity."
    ),
    "FRII": (
        "**FR-II** — Bright hotspots at the ends of extended lobes with an edge-brightened morphology."
    ),
    "Compact": (
        "**Compact** — Radio emission is concentrated into a compact unresolved or barely resolved source with no prominent jet or lobe structure."
    ),
    "Bent": (
        "**Bent** — Radio jets or lobes are curved due to interactions with the surrounding environment, commonly seen in cluster galaxies."
    ),
}

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# ── preprocessing (matches training pipeline) ──────────────────────────────
def preprocess(image):
    if image.ndim == 3:
        image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    image = Image.fromarray(image.astype(np.uint8))
    
    image = transforms.CenterCrop(IMSIZE)(image)
    image = transforms.Pad((0, 0, 1, 1))(image)
    image = transforms.ToTensor()(image)
    image = transforms.Normalize((DATAMEAN,), (DATASTD,))(image)

    return image.unsqueeze(0)


# ── model loader (cached) ──────────────────────────────────────────────────
_model_cache = {}

def load_model(model_name):
    if model_name in _model_cache:
        return _model_cache[model_name]
    path = MODEL_PATHS[model_name]
    if "DNSteerable" in model_name:
        model = DNSteerableLeNet(1, 4, IMSIZE + 1, kernel_size=5, N=16)
    else:
        model = VanillaLeNet(1, 4, IMSIZE + 1, kernel_size=5)

    state = torch.load(path, map_location=DEVICE)
    model.load_state_dict(state)
    model = model.to(DEVICE)
    model.eval()
    _model_cache[model_name] = model
    return model

_cam_cache = {}

def get_gradcam(model_name):
    if model_name not in _cam_cache:
        _cam_cache[model_name] = GradCAM_ECNN(load_model(model_name))
    return _cam_cache[model_name]

# ── inference ──────────────────────────────────────────────────────────────
@spaces.GPU
def predict(image, model_name):
    if image is None:
        return None, None, "Upload an image to classify."

    model = load_model(model_name)
    gcam = get_gradcam(model_name)

    img_tensor = preprocess(image).to(DEVICE)
    cam, pred_idx, probs = gcam.generate(img_tensor)

    # ── confidence output ──
    pred_name = CLASS_NAMES[pred_idx]
    conf = probs[pred_idx] * 100

    prediction_html = f"""
    <div class="prediction-card">
        <h3 class="section-title">Prediction</h3>
        <div class="prediction-name">{pred_name}</div>
        <hr>
        <h3 class="section-title">Confidence</h3>
        <div class="confidence-value">{conf:.1f}%</div>
        <div class="confidence-bar">
            <div class="confidence-fill" style="width:{conf:.1f}%"></div>
        </div>
        <div class="confidence-scale">
            <span>0%</span>
            <span>50%</span>
            <span>100%</span>
        </div>
    </div>
    """

    # ── overlay ──
    orig = img_tensor[0, 0].detach().cpu().numpy()
    orig = (orig - orig.min()) / (orig.max() - orig.min() + 1e-8)
    orig_uint8  = (orig * 255).astype(np.uint8)
    heatmap     = cv2.applyColorMap((cam * 255).astype(np.uint8), cv2.COLORMAP_INFERNO)
    heatmap_rgb = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
    orig_rgb    = cv2.cvtColor(orig_uint8, cv2.COLOR_GRAY2RGB)
    overlay     = cv2.addWeighted(orig_rgb, 0.55, heatmap_rgb, 0.45, 0)

    # ── explanation ──
    EXPLANATIONS = {
        "FRI": ["/gradio_api/file=assets/fri1.png", "/gradio_api/file=assets/fri2.png"],
        "FRII": ["/gradio_api/file=assets/frii1.png", "/gradio_api/file=assets/frii2.png"],
        "Compact": ["/gradio_api/file=assets/compact1.png", "/gradio_api/file=assets/compact2.png"],
        "Bent": ["/gradio_api/file=assets/bent1.png", "/gradio_api/file=assets/bent2.png"],
    }
    images_html = f"""
<div style="
    display:flex;
    gap:12px;
    justify-content:center;
    align-items:center;
    flex-wrap:wrap;
">
    {
        "".join(
            f'<img src="{img}" style="width:120px; border-radius:8px;">'
            for img in EXPLANATIONS[pred_name]
        )
    }
</div>
"""
    explanation = (
        f"### Predicted: {pred_name} ({conf:.1f}% confidence)\n\n"
        f"{CLASS_DESC[pred_name]}\n\n"
        f"The heatmap shows which pixels most influenced this decision — "
        f"brighter regions had greater weight.\n\n"
        f"{images_html}"
    )
    return (
        gr.update(value=prediction_html, visible=True),
        Image.fromarray(overlay),
        gr.update(value=explanation, visible=True),
    )


# ── example images ─────────────────────────────────────────────────────────
EXAMPLES = {
    "FRI": ["assets/eg_fri1.png", "assets/eg_fri2.png"],
    "FRII": ["assets/eg_frii1.png", "assets/eg_frii2.png"],
    "COMPACT": ["assets/eg_compact1.png", "assets/eg_compact2.png"],
    "BENT": ["assets/eg_bent1.png", "assets/eg_bent2.png"],
}

# ── UI ─────────────────────────────────────────────────────────────────────
css = Path("style.css").read_text()

with gr.Blocks(fill_width=True,) as demo:

    # ==========================================================
    # NAVBAR
    # ==========================================================
    gr.HTML("""
    <div class="navbar">
        <div class="title">
            <h2 style="margin: 0px;">Radio Galaxy Morphology Classifier</h2>
        </div>

        <div class="nav-links">
            <a href="https://github.com/huma-m/radio-galaxy-classifier" target="_blank" style="display: inline-flex; align-items: center; gap: 8px; color: white; padding: 10px 20px; font-size: 16px; font-weight: 600; font-family: sans-serif; text-decoration: none;">
                <svg height="20" width="20" viewBox="0 0 16 16" fill="white"><path d="M8 0c4.42 0 8 3.58 8 8a8.013 8.013 0 0 1-5.45 7.59c-.4.08-.55-.17-.55-.38 0-.27.01-1.13.01-2.2 0-.75-.25-1.23-.54-1.48 1.78-.2 3.65-.88 3.65-3.95 0-.88-.31-1.59-.82-2.15.08-.2.36-1.02-.08-2.12 0 0-.67-.22-2.2.82A7.48 7.48 0 0 0 8 3c-.68 0-1.36.09-2 .27-1.53-1.03-2.2-.82-2.2-.82-.44 1.1-.16 1.92-.08 2.12-.51.56-.82 1.28-.82 2.15 0 3.06 1.86 3.75 3.64 3.95-.23.2-.44.55-.51 1.07-.46.21-1.61.55-2.33-.66-.15-.24-.6-.83-1.23-.82-.67.01-.27.38.01.53.34.19.73.9.82 1.13.16.45.68 1.31 2.69.94 0 .67.01 1.3.01 1.49 0 .21-.15.45-.55.38A7.995 7.995 0 0 1 0 8c0-4.42 3.58-8 8-8Z"/></svg>
                GitHub
            </a>
        </div>
    </div>
    """, elem_id="navbar-block")

    # ==========================================================
    # MAIN CONTENT
    # ==========================================================
    with gr.Row(equal_height=True):

        # ======================================================
        # LEFT PANEL
        # ======================================================
        with gr.Column(scale=1, elem_classes="card"):

            gr.Markdown("### Upload Radio Galaxy Image (PNG, JPG, JPEG)", elem_classes="title-text")
            image_input = gr.Image(
                type="numpy",
                image_mode="L",
                show_label=False,
                height=320,
            )

            gr.Markdown("### Select a Model")
            with gr.Row(equal_height=True, elem_id="model_row"):
                model_dropdown = gr.Dropdown(
                    choices=list(MODEL_PATHS.keys()),
                    value="DNSteerableLeNet (G-CNN)",
                    show_label=False,
                    allow_custom_value=False,
                    filterable=False,
                    elem_id="model_dropdown",
                    scale=4
                )

                classify_btn = gr.Button(
                    "Classify",
                    elem_id="classify-btn",
                    min_width=100,
                    scale=1
                )

            
            if EXAMPLES:
                gr.Markdown("### Example Images")
                with gr.Row():
                    gr.Examples(
                        examples=EXAMPLES["FRI"],
                        inputs=[image_input,],
                        examples_per_page=2,
                        label="FRI",
                        elem_id="examples"
                    )
                    gr.Examples(
                        examples=EXAMPLES["FRII"],
                        inputs=[image_input,],
                        examples_per_page=2,
                        label="FRII",
                        elem_id="examples"
                    )
                with gr.Row():
                    gr.Examples(
                        examples=EXAMPLES["COMPACT"],
                        inputs=[image_input,],
                        examples_per_page=2,
                        label="Compact",
                        elem_id="examples"
                    )
                    gr.Examples(
                        examples=EXAMPLES["BENT"],
                        inputs=[image_input,],
                        examples_per_page=2,
                        label="Bent",
                        elem_id="examples"
                    )

        # ======================================================
        # RIGHT PANEL
        # ======================================================
        with gr.Column(scale=2, elem_classes="card"):
            
            gr.Markdown("### Classification Result", elem_classes="title-text")
                
            with gr.Row(equal_height=True, elem_classes="pred-row"):
                overlay_output = gr.Image(
                    label="Grad-CAM",
                    height=320,
                    scale=2
                )
                gr.Markdown("""### Grad-CAM Explanation
                Brighter (yellow/red) regions had the greatest influence on the model's prediction, while darker (blue) regions contributed less. The model focuses on the brightest regions of the radio galaxy, which correspond to the core and lobes in FR-I and FR-II morphologies, respectively.
                """, elem_classes="gradcam-desc", scale=1)
                
            with gr.Row(equal_height=True, elem_classes="pred-row"):
                prediction_output = gr.HTML(elem_id="prediction-card", visible=False)
                explain_output = gr.Markdown(elem_id="result-md", visible=False)
                
    # ==========================================================
    # FOOTER
    # ==========================================================
    with gr.Column(elem_id="footer"):

        with gr.Accordion("Project Notes", open=False, elem_id="project-notes"):
            gr.Markdown("""
## About the Project

This application classifies radio galaxies into four morphological classes using deep learning:

- **FRI** – Core-brightened galaxies with jets that fade away from the center.
- **FRII** – Edge-brightened galaxies with prominent hotspots at the ends of their radio lobes.
- **Compact** – Small unresolved radio sources without extended jet structures.
- **Bent** – Galaxies whose jets are curved due to interactions with their surrounding environment.

---

## Dataset

The model was trained using the RadioGalaxyDataset, a curated dataset containing 2,158 grayscale radio galaxy images from the FIRST (Faint Images of the Radio Sky at Twenty-Centimeters) survey. 
The dataset combines expert-labelled sources from six published catalogues: MiraBest, Gendre, FRICAT (Capetti et al., 2017b), FR0CAT (Capetti et al., 2017a), Baldi et al. (2018), and Proctor. 
Each image is labelled as one of four radio galaxy morphologies: FRI, FRII, Compact, or Bent.

---

## Model

Multiple convolutional neural network architectures were evaluated for radio galaxy morphology classification, including **ResNet-18, DenseNet-121, VanillaLeNet,** and **DNSteerableLeNet,** a Group Equivariant CNN based on the work of **Scaife & Porter (2021)**. The final application uses **DNSteerableLeNet**, which achieved the best performance with **81% test accuracy** and a **Macro F1-score of 0.81**.

The effect of input image resolution was also investigated. Increasing the image size from **150×150** to **225×225** improved the performance of pretrained CNNs, with **ResNet-18** increasing from **69% → 78%** accuracy and **DenseNet-121** from **73% → 79%**. However, the equivariant models showed performance degradation at the higher resolution, indicating that rotationally equivariant feature extraction contributed more to performance than simply increasing input resolution. Therefore, the final model uses the **150×150** input resolution.

---

## Error Analysis

While the model performs well overall, some morphologies remain challenging.

Typical confusion occurs between:

- **Bent ↔ FRI**, as bent jets can closely resemble the edge-darkened jet structures of FR-I galaxies.
- **Bent ↔ FRII**, when curved lobes or asymmetric hotspots resemble distorted FR-II morphologies.

The Bent class is the most challenging to classify due to its high morphological variability and similarity to both FRI and FRII galaxies. These errors reflect the inherent complexity of radio galaxy morphology rather than simple model failures.
  
---
  
## References

**Dataset**
- RadioGalaxyDataset (Zenodo): https://zenodo.org/records/7351724

**Model**
- Scaife & Porter (2021), *Fanaroff–Riley Classification of Radio Galaxies Using Group Equivariant Convolutional Neural Networks:*
  https://arxiv.org/abs/2102.08252
""")

    # ==========================================================
    # EVENTS
    # ==========================================================
    classify_btn.click(
        fn=predict,
        inputs=[
            image_input,
            model_dropdown
        ],
        outputs=[
            prediction_output,
            overlay_output,
            explain_output
        ]
    )
demo.launch(server_name="0.0.0.0",server_port=7860,css=css, footer_links=[], allowed_paths=["assets"])