import concurrent.futures
import math
import os
from functools import lru_cache
from io import BytesIO

os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

import gradio as gr
import numpy as np
import pandas as pd
import requests
import torch
import torch.nn.functional as F
import torchvision.models as models
from matplotlib import colormaps
from PIL import Image, ImageDraw, ImageFilter


torch.set_num_threads(max(1, min(4, os.cpu_count() or 1)))

APP_TITLE = "Model Layers Viewer"
MOMA_URL = "https://media.githubusercontent.com/media/MuseumofModernArt/collection/main/Artworks.csv"
HEADERS = {
    "User-Agent": (
        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
        "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0 Safari/537.36"
    )
}
MOMA_SAMPLE_SIZE = 36
MOMA_CANDIDATES = 140
NOTEBOOK_LAYERS = [0, 2, 4, 7, 14, 18]
CHANNELS_PER_LAYER = 8
OVERLAY_LAYER = 14
PREVIEW_SIZE = 360
TILE_SIZE = 120


def stage_name(layer_index):
    if layer_index <= 2:
        return "Early vision"
    if layer_index <= 7:
        return "Textures and small shapes"
    if layer_index <= 14:
        return "Object parts and composition"
    return "Compact semantic evidence"


def stage_explanation(layer_index):
    if layer_index <= 2:
        return "These first filters respond to edges, contrast, color changes, and simple marks."
    if layer_index <= 7:
        return "The model is combining simple marks into repeated textures, corners, curves, and local motifs."
    if layer_index <= 14:
        return "At this depth, each activation cell responds to a wider region of the original image."
    return "The final feature blocks are very small spatial maps that summarize evidence for the image as a whole."


def make_fallback_images():
    samples = []
    size = 384

    def canvas(bg):
        return Image.new("RGB", (size, size), bg)

    img = canvas((232, 225, 210))
    draw = ImageDraw.Draw(img)
    for i, color in enumerate([(199, 44, 65), (25, 93, 138), (240, 174, 54), (42, 120, 83)]):
        draw.rectangle((30 + i * 38, 40 + i * 48, 250 + i * 20, 118 + i * 55), fill=color)
    samples.append({"img": img, "title": "Fallback: stacked color blocks", "artist": "generated", "url": ""})

    img = canvas((22, 25, 30))
    draw = ImageDraw.Draw(img)
    for radius, color in zip(
        range(170, 10, -24),
        [(231, 76, 60), (241, 196, 15), (52, 152, 219), (46, 204, 113), (155, 89, 182)],
    ):
        draw.ellipse((192 - radius, 192 - radius, 192 + radius, 192 + radius), outline=color, width=14)
    samples.append({"img": img, "title": "Fallback: concentric rings", "artist": "generated", "url": ""})

    img = canvas((246, 244, 238))
    draw = ImageDraw.Draw(img)
    for x in range(-120, size + 120, 34):
        draw.line((x, 0, x + 190, size), fill=(35, 63, 99), width=9)
        draw.line((x + 16, 0, x + 206, size), fill=(218, 83, 44), width=4)
    samples.append({"img": img, "title": "Fallback: diagonal line field", "artist": "generated", "url": ""})

    img = canvas((245, 245, 242))
    draw = ImageDraw.Draw(img)
    points = [(40, 280), (85, 110), (150, 250), (205, 70), (260, 220), (340, 115)]
    draw.line(points, fill=(26, 26, 26), width=12, joint="curve")
    for x, y in points:
        draw.ellipse((x - 18, y - 18, x + 18, y + 18), fill=(35, 125, 161))
    samples.append({"img": img.filter(ImageFilter.SMOOTH), "title": "Fallback: bold path drawing", "artist": "generated", "url": ""})

    return samples


@lru_cache(maxsize=1)
def load_model():
    weights = models.MobileNet_V2_Weights.DEFAULT
    model = models.mobilenet_v2(weights=weights)
    model.eval()
    return model, weights.transforms()


@lru_cache(maxsize=1)
def layer_choices():
    model, _ = load_model()
    return [f"{index}: {layer.__class__.__name__}" for index, layer in enumerate(model.features)]


def rgb_image(image):
    if image is None:
        return None
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    return image.convert("RGB")


def fit_image(image, max_side=PREVIEW_SIZE):
    image = image.copy().convert("RGB")
    image.thumbnail((max_side, max_side), Image.Resampling.LANCZOS)
    return image


def normalize_map(values):
    values = values.astype(np.float32)
    low = float(values.min())
    high = float(values.max())
    return (values - low) / (high - low + 1e-8)


def colorize(values, cmap_name="magma"):
    norm = normalize_map(values)
    rgba = colormaps[cmap_name](norm)
    rgb = (rgba[:, :, :3] * 255).astype(np.uint8)
    return Image.fromarray(rgb)


def overlay_heatmap(image, values, alpha=0.46, cmap_name="magma"):
    heat = colorize(values, cmap_name).resize(image.size, Image.Resampling.BILINEAR).convert("RGBA")
    base = image.convert("RGBA")
    return Image.blend(base, heat, alpha=alpha).convert("RGB")


def fetch_moma_image(row):
    try:
        response = requests.get(row["ImageURL"], headers=HEADERS, timeout=10)
        response.raise_for_status()
        image = Image.open(BytesIO(response.content)).convert("RGB")
        image.thumbnail((640, 640), Image.Resampling.LANCZOS)
        return {
            "img": image.copy(),
            "title": str(row.get("Title") or "Untitled")[:80],
            "artist": str(row.get("Artist") or "Unknown artist")[:80],
            "url": str(row.get("ImageURL") or ""),
        }
    except Exception:
        return None


@lru_cache(maxsize=1)
def load_moma_items():
    try:
        df = pd.read_csv(
            MOMA_URL,
            usecols=["Title", "Artist", "ImageURL"],
            low_memory=False,
        )
        has_image = df["ImageURL"].notna() & df["ImageURL"].str.startswith("http", na=False)
        candidates = df[has_image].sample(
            min(MOMA_CANDIDATES, int(has_image.sum())),
            random_state=42,
        )
        rows = [row for _, row in candidates.iterrows()]

        with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
            results = list(executor.map(fetch_moma_image, rows))

        items = [item for item in results if item is not None][:MOMA_SAMPLE_SIZE]
        if len(items) >= 8:
            return items, "Loaded images from the MoMA collection data used in the notebook."
    except Exception as exc:
        return make_fallback_images(), f"MoMA images could not be loaded here ({type(exc).__name__}). Using generated fallback images."

    return make_fallback_images(), "MoMA image downloads did not return enough usable images. Using generated fallback images."


def gallery_items():
    items, _ = load_moma_items()
    return [(item["img"], f"{item['title']}\n{item['artist']}") for item in items]


def collect_activations(image):
    model, preprocess = load_model()
    x = preprocess(image).unsqueeze(0)
    activations = []

    with torch.inference_mode():
        out = x
        for layer in model.features:
            out = layer(out)
            activations.append(out.detach().cpu())

    return activations


def feature_vector(image):
    model, preprocess = load_model()
    x = preprocess(image).unsqueeze(0)
    with torch.inference_mode():
        out = model.features(x)
        pooled = F.adaptive_avg_pool2d(out, (1, 1)).flatten(1)
    return pooled.squeeze(0).cpu().numpy()


@lru_cache(maxsize=1)
def moma_feature_space():
    items, message = load_moma_items()
    features = np.vstack([feature_vector(item["img"]) for item in items])

    method = "UMAP"
    try:
        import umap

        reducer = umap.UMAP(
            n_components=2,
            n_neighbors=max(2, min(10, len(items) - 1)),
            min_dist=0.12,
            metric="cosine",
            random_state=42,
        )
        embedding = reducer.fit_transform(features)
    except Exception:
        method = "PCA fallback"
        centered = features - features.mean(axis=0, keepdims=True)
        _, _, vt = np.linalg.svd(centered, full_matrices=False)
        embedding = centered @ vt[:2].T

    return items, features, embedding, method, message


def render_embedding_image(embedding, selected_index, method):
    width, height = 720, 540
    margin = 58
    image = Image.new("RGB", (width, height), (250, 250, 248))
    draw = ImageDraw.Draw(image)

    x = embedding[:, 0]
    y = embedding[:, 1]
    x_span = float(x.max() - x.min()) or 1.0
    y_span = float(y.max() - y.min()) or 1.0
    xs = margin + ((x - x.min()) / x_span) * (width - 2 * margin)
    ys = height - margin - ((y - y.min()) / y_span) * (height - 2 * margin)

    for t in np.linspace(0, 1, 5):
        gx = margin + t * (width - 2 * margin)
        gy = margin + t * (height - 2 * margin)
        draw.line((gx, margin, gx, height - margin), fill=(226, 228, 232), width=1)
        draw.line((margin, gy, width - margin, gy), fill=(226, 228, 232), width=1)

    draw.rectangle((margin, margin, width - margin, height - margin), outline=(190, 196, 205), width=1)
    draw.text((margin, 22), f"{method} of MobileNetV2 activations for MoMA images", fill=(35, 39, 47))

    for index, (px, py) in enumerate(zip(xs, ys)):
        if index == selected_index:
            continue
        draw.ellipse((px - 6, py - 6, px + 6, py + 6), fill=(59, 130, 246), outline=(255, 255, 255), width=2)

    px = xs[selected_index]
    py = ys[selected_index]
    draw.ellipse((px - 12, py - 12, px + 12, py + 12), fill=(225, 29, 72), outline=(20, 20, 20), width=2)
    draw.text((margin, height - 36), "red point = selected image", fill=(75, 85, 99))
    return image


def activation_energy(activation):
    fmap = activation.clamp(min=0).mean(dim=1).squeeze(0).cpu().numpy()
    return normalize_map(fmap)


def strongest_channels(activation, count):
    fmap = activation.squeeze(0).cpu()
    strengths = fmap.mean(dim=(1, 2))
    channel_ids = torch.argsort(strengths, descending=True)[: min(int(count), fmap.shape[0])].tolist()
    return [(channel, fmap[channel].numpy()) for channel in channel_ids]


def channel_grid_image(activation, count=CHANNELS_PER_LAYER):
    selected = strongest_channels(activation, count)
    cols = 4
    rows = math.ceil(len(selected) / cols)
    label_h = 20
    grid = Image.new("RGB", (cols * TILE_SIZE, rows * (TILE_SIZE + label_h)), (248, 248, 246))
    draw = ImageDraw.Draw(grid)

    for position, (channel, values) in enumerate(selected):
        x = (position % cols) * TILE_SIZE
        y = (position // cols) * (TILE_SIZE + label_h)
        tile = colorize(values, "viridis").resize((TILE_SIZE, TILE_SIZE), Image.Resampling.NEAREST)
        grid.paste(tile, (x, y))
        draw.text((x + 6, y + TILE_SIZE + 3), f"channel {channel}", fill=(35, 39, 47))

    return grid


def fixed_layer_gallery(activations):
    choices = layer_choices()
    outputs = []

    for layer_index in NOTEBOOK_LAYERS:
        activation = activations[layer_index]
        _, channels, height, width = activation.shape
        caption = (
            f"Layer {layer_index}: {stage_name(layer_index)}\n"
            f"torch.Size([1, {channels}, {height}, {width}])\n"
            f"{choices[layer_index].split(': ', 1)[1]} | {CHANNELS_PER_LAYER} channel maps"
        )
        outputs.append((channel_grid_image(activation), caption))

    return outputs


def layer_summary():
    return (
        "### Pixel activations through the model\n"
        "These are the same intermediate layers checked in the notebook: 0, 2, 4, 7, 14, and 18. Each panel is one layer. "
        "Inside each panel are several channel maps from that layer: bright cells are places where that map activated strongly. "
        "Earlier layers keep larger grids; later layers shrink to smaller, chunkier grids."
    )


def analyze_image(image):
    image = rgb_image(image)
    if image is None:
        items, _ = load_moma_items()
        image = items[0]["img"]

    activations = collect_activations(image)
    overlay_activation = activations[OVERLAY_LAYER]

    return (
        image,
        fixed_layer_gallery(activations),
        overlay_heatmap(fit_image(image), activation_energy(overlay_activation), alpha=0.48, cmap_name="magma"),
        layer_summary(),
    )


def make_umap_plot(selected_index=0):
    items, _, embedding, method, message = moma_feature_space()
    selected_index = 0 if selected_index is None else int(selected_index)
    selected_index = max(0, min(selected_index, len(items) - 1))
    plot_image = render_embedding_image(embedding, selected_index, method)

    item = items[selected_index]
    distances = np.linalg.norm(embedding - embedding[selected_index], axis=1)
    closest_ids = [idx for idx in np.argsort(distances) if idx != selected_index][:3]
    farthest_ids = [idx for idx in np.argsort(distances)[::-1] if idx != selected_index][:3]
    closest_gallery = [(items[idx]["img"], f"{items[idx]['title']}\n{items[idx]['artist']}") for idx in closest_ids]
    farthest_gallery = [(items[idx]["img"], f"{items[idx]['title']}\n{items[idx]['artist']}") for idx in farthest_ids]
    relation = (
        f"### Selected MoMA image\n"
        f"**{item['title']}** by {item['artist']}\n\n"
        f"Nearby points have similar MobileNetV2 feature vectors; faraway points have very different vectors. "
        f"This does not mean the artworks share labels. It means the neural net produced similar or different "
        f"patterns of activation after processing the pixels.\n\n"
        f"{message}"
    )
    return plot_image, relation, closest_gallery, farthest_gallery


def initialize():
    items, _ = load_moma_items()
    selected_index = 0
    selected_image, layer_gallery, overlay, summary = analyze_image(items[selected_index]["img"])
    umap_plot, relation, closest_gallery, farthest_gallery = make_umap_plot(selected_index)
    return (
        gallery_items(),
        selected_image,
        layer_gallery,
        overlay,
        summary,
        umap_plot,
        relation,
        closest_gallery,
        farthest_gallery,
        selected_index,
    )


def select_moma(evt: gr.SelectData):
    index = evt.index if isinstance(evt.index, int) else 0
    items, _ = load_moma_items()
    index = max(0, min(index, len(items) - 1))
    selected_image, layer_gallery, overlay, summary = analyze_image(items[index]["img"])
    umap_plot, relation, closest_gallery, farthest_gallery = make_umap_plot(index)
    return selected_image, layer_gallery, overlay, summary, umap_plot, relation, closest_gallery, farthest_gallery, index


def analyze_upload(image):
    selected_image, layer_gallery, overlay, summary = analyze_image(image)
    umap_plot, relation, closest_gallery, farthest_gallery = make_umap_plot(0)
    relation = (
        "### Uploaded image\n"
        "The UMAP tab is built from the MoMA image set. Uploaded images are shown in the layer viewer, "
        "but are not inserted into the precomputed MoMA map."
    )
    return selected_image, layer_gallery, overlay, summary, umap_plot, relation, closest_gallery, farthest_gallery, -1


def build_app():
    theme = gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="pink",
        neutral_hue="slate",
        radius_size="sm",
    )

    css = """
    .activation-gallery img { object-fit: contain !important; }
    .moma-gallery img { object-fit: cover !important; }
    .compact-note textarea { font-size: 0.95rem !important; }
    """

    with gr.Blocks(title=APP_TITLE, theme=theme, css=css) as demo:
        selected_index = gr.State(0)

        gr.Markdown(
            "# Model Layers Viewer\n"
            "Choose a MoMA image or upload your own. The app runs MobileNetV2 on CPU and shows pixel activation maps "
            "at the same intermediate layers used in the workshop notebook.\n\n"
            "Inspired by [Small ML: The Art of Data Discovery]"
            "(https://colab.research.google.com/drive/1_8GasNHKJpO8x-AAt93D3LfrvCLGzsbj?usp=sharing), "
            "a workshop by Sam Keene at [ITP Camp 2026](https://itp.nyu.edu/camp/2026/session/205)."
        )

        with gr.Row(equal_height=False):
            with gr.Column(scale=1, min_width=310):
                moma_gallery = gr.Gallery(
                    label="MoMA image set",
                    columns=3,
                    rows=4,
                    height=520,
                    object_fit="cover",
                    elem_classes=["moma-gallery"],
                )
                upload = gr.Image(label="Upload your own image", type="pil", sources=["upload", "clipboard"])

            with gr.Column(scale=2, min_width=520):
                with gr.Tabs():
                    with gr.Tab("Layer activations"):
                        with gr.Row(equal_height=False):
                            selected_image = gr.Image(label="Selected image", type="pil", interactive=False)
                            selected_overlay = gr.Image(label="One overlay: layer 14 activation over image", type="pil", interactive=False)
                        layer_summary_md = gr.Markdown()
                        layer_progression = gr.Gallery(
                            label="Notebook layers: activation maps across channels",
                            columns=1,
                            height=1400,
                            object_fit="contain",
                            elem_classes=["activation-gallery"],
                        )

                    with gr.Tab("UMAP relationships"):
                        umap_plot = gr.Image(label="MoMA images after MobileNetV2 activation extraction", type="pil", interactive=False)
                        relation_md = gr.Markdown()
                        with gr.Row(equal_height=False):
                            closest_gallery = gr.Gallery(
                                label="3 closest images in feature space",
                                columns=3,
                                rows=1,
                                height=300,
                                object_fit="cover",
                                elem_classes=["moma-gallery"],
                            )
                            farthest_gallery = gr.Gallery(
                                label="3 farthest images in feature space",
                                columns=3,
                                rows=1,
                                height=300,
                                object_fit="cover",
                                elem_classes=["moma-gallery"],
                            )

        demo.load(
            initialize,
            inputs=None,
            outputs=[
                moma_gallery,
                selected_image,
                layer_progression,
                selected_overlay,
                layer_summary_md,
                umap_plot,
                relation_md,
                closest_gallery,
                farthest_gallery,
                selected_index,
            ],
            show_progress="full",
        )
        moma_gallery.select(
            select_moma,
            inputs=None,
            outputs=[
                selected_image,
                layer_progression,
                selected_overlay,
                layer_summary_md,
                umap_plot,
                relation_md,
                closest_gallery,
                farthest_gallery,
                selected_index,
            ],
            show_progress="minimal",
        )
        upload.change(
            analyze_upload,
            inputs=[upload],
            outputs=[
                selected_image,
                layer_progression,
                selected_overlay,
                layer_summary_md,
                umap_plot,
                relation_md,
                closest_gallery,
                farthest_gallery,
                selected_index,
            ],
            show_progress="minimal",
        )
    return demo


if __name__ == "__main__":
    build_app().launch()