import functools
import io
import json
import logging
import math
import pathlib
import typing

import beartype
import einops
import einops.layers.torch
import gradio as gr
import numpy as np
import saev.activations
import saev.config
import saev.nn
import saev.visuals
import torch
from jaxtyping import Bool, Float, Int, UInt8, jaxtyped
from PIL import Image, ImageDraw
from torch import Tensor

import constants
import data
import modeling

logger = logging.getLogger("app.py")

####################
# Global Constants #
####################


MAX_FREQ = 1e-2
"""Maximum frequency. Any feature that fires more than this is ignored."""

RESIZE_SIZE = 512
"""Resize shorter size to this size in pixels."""

CROP_SIZE = (448, 448)
"""Crop size in pixels."""

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
"""Hardware accelerator, if any."""

CWD = pathlib.Path(".")
"""Current working directory."""

N_SAE_LATENTS = 3
"""Number of SAE latents to show."""

N_LATENT_EXAMPLES = 4
"""Number of examples per SAE latent to show."""

##########
# Models #
##########


@functools.cache
def load_sae(device: str) -> saev.nn.SparseAutoencoder:
    """
    Loads a sparse autoencoder from disk.
    """
    sae_ckpt_fpath = CWD / "assets" / "sae.pt"
    sae = saev.nn.load(str(sae_ckpt_fpath))
    sae.to(device).eval()
    return sae


@functools.cache
def load_clf() -> torch.nn.Module:
    # /home/stevens.994/projects/saev/checkpoints/contrib/semseg/lr_0_001__wd_0_001/model_step8000.pt
    head_ckpt_fpath = CWD / "assets" / "clf.pt"
    with open(head_ckpt_fpath, "rb") as fd:
        kwargs = json.loads(fd.readline().decode())
        buffer = io.BytesIO(fd.read())

    model = torch.nn.Linear(**kwargs)
    state_dict = torch.load(buffer, weights_only=True, map_location=DEVICE)
    model.load_state_dict(state_dict)
    model = model.to(DEVICE).eval()
    return model


####################
# Global Variables #
####################


@beartype.beartype
def load_tensor(path: str | pathlib.Path) -> Tensor:
    return torch.load(path, weights_only=True, map_location="cpu")


@functools.cache
def load_tensors() -> tuple[
    Int[Tensor, "d_sae k"],
    UInt8[Tensor, "d_sae k n_patches"],
    Bool[Tensor, " d_sae"],
]:
    """
    Loads the tensors for the SAE for ADE20K.
    """
    top_img_i = load_tensor(CWD / "assets" / "top_img_i.pt")
    top_values = load_tensor(CWD / "assets" / "top_values_uint8.pt")
    sparsity = load_tensor(CWD / "assets" / "sparsity.pt")

    mask = torch.ones(sparsity.shape, dtype=bool)
    mask = mask & (sparsity < MAX_FREQ)

    return top_img_i, top_values, mask


############
# Datasets #
############


@jaxtyped(typechecker=beartype.beartype)
def add_highlights(
    img: Image.Image,
    patches: Float[np.ndarray, " n_patches"],
    *,
    upper: int | None = None,
    opacity: float = 0.9,
) -> Image.Image:
    breakpoint()
    if not len(patches):
        return img

    iw_np, ih_np = int(math.sqrt(len(patches))), int(math.sqrt(len(patches)))
    iw_px, ih_px = img.size
    pw_px, ph_px = iw_px // iw_np, ih_px // ih_np
    assert iw_np * ih_np == len(patches)

    # Create a transparent overlay
    overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(overlay)

    # Using semi-transparent red (255, 0, 0, alpha)
    for p, val in enumerate(patches):
        assert upper is not None
        val /= upper + 1e-9
        x_np, y_np = p % iw_np, p // ih_np
        draw.rectangle(
            [
                (x_np * pw_px, y_np * ph_px),
                (x_np * pw_px + pw_px, y_np * ph_px + ph_px),
            ],
            fill=(int(val * 256), 0, 0, int(opacity * val * 256)),
        )

    # Composite the original image and the overlay
    return Image.alpha_composite(img.convert("RGBA"), overlay)


#######################
# Inference Functions #
#######################


@beartype.beartype
class Example(typing.TypedDict):
    """Represents an example image and its associated label.

    Used to store examples of SAE latent activations for visualization.
    """

    index: int
    """Dataset index."""
    orig_url: str
    """The URL or path to access the original example image."""
    highlighted_url: str
    """The URL or path to access the SAE-highlighted image."""
    seg_url: str
    """Base64-encoded version of the colored segmentation map."""


@beartype.beartype
class SaeActivation(typing.TypedDict):
    """Represents the activation pattern of a single SAE latent across patches.

    This captures how strongly a particular SAE latent fires on different patches of an input image.
    """

    latent: int
    """The index of the SAE latent being measured."""

    highlighted_url: str
    """The image with the colormaps applied."""

    activations: list[float]
    """The activation values of this latent across different patches. Each value represents how strongly this latent fired on a particular patch."""

    examples: list[Example]
    """Top examples for this latent."""


@beartype.beartype
def get_image(i: int) -> tuple[str, str, int]:
    img_sized = data.to_sized(data.get_image(i))
    seg_sized = data.to_sized(data.get_seg(i))
    seg_u8_sized = data.to_u8(seg_sized)
    seg_img_sized = data.u8_to_img(seg_u8_sized)

    return data.img_to_base64(img_sized), data.img_to_base64(seg_img_sized), i


@beartype.beartype
@torch.inference_mode
def get_sae_activations(image_i: int, patches: list[int]) -> list[SaeActivation]:
    """
    Given a particular cell, returns some highlighted images showing what feature fires most on this cell.
    """
    if not patches:
        return []

    split_vit, vit_transform = modeling.load_vit(DEVICE)
    sae = load_sae(DEVICE)

    img = data.get_image(image_i)

    x_BCWH = vit_transform(img)[None, ...].to(DEVICE)

    x_BPD = split_vit.forward_start(x_BCWH)
    x_BPD = (
        x_BPD.clamp(-1e-5, 1e5) - (constants.DINOV2_IMAGENET1K_MEAN).to(DEVICE)
    ) / constants.DINOV2_IMAGENET1K_SCALAR

    # Need to pick out the right patches
    # + 1 + 4 for 1 [CLS] token and 4 register tokens
    x_PD = x_BPD[0, [p + 1 + 4 for p in patches]]
    _, f_x_PS, _ = sae(x_PD)

    f_x_S = einops.reduce(f_x_PS, "patches n_latents -> n_latents", "sum")
    logger.info("Got SAE activations.")

    top_img_i, top_values, mask = load_tensors()

    latents = torch.argsort(f_x_S, descending=True).cpu()
    latents = latents[mask[latents]][:N_SAE_LATENTS].tolist()

    sae_activations = []
    for latent in latents:
        pairs, seen_i_im = [], set()
        for i_im, values_p in zip(top_img_i[latent].tolist(), top_values[latent]):
            if i_im in seen_i_im:
                continue

            pairs.append((i_im, values_p))
            seen_i_im.add(i_im)
            if len(pairs) >= N_LATENT_EXAMPLES:
                break

        # How to scale values.
        upper = None
        if top_values[latent].numel() > 0:
            upper = top_values[latent].max().item()

        examples = []
        for i_im, values_p in pairs:
            seg_sized = data.to_sized(data.get_seg(i_im))
            img_sized = data.to_sized(data.get_image(i_im))

            seg_u8_sized = data.to_u8(seg_sized)
            seg_img_sized = data.u8_to_img(seg_u8_sized)

            highlighted_sized = add_highlights(
                img_sized, values_p.float().numpy(), upper=upper
            )

            examples.append({
                "index": i_im,
                "orig_url": data.img_to_base64(img_sized),
                "highlighted_url": data.img_to_base64(highlighted_sized),
                "seg_url": data.img_to_base64(seg_img_sized),
            })

        sae_activations.append({
            "latent": latent,
            "examples": examples,
        })

    return sae_activations


@torch.inference_mode
def get_true_labels(image_i: int) -> Image.Image:
    seg = human_dataset[image_i]["segmentation"]
    image = seg_to_img(seg)
    return image


@torch.inference_mode
def get_pred_labels(i: int) -> list[Image.Image | list[int]]:
    sample = vit_dataset[i]
    x = sample["image"][None, ...].to(device)
    x_BPD = rest_of_vit.forward_start(x)
    x_BPD = rest_of_vit.forward_end(x_BPD)

    x_WHD = einops.rearrange(x_BPD, "() (w h) dim -> w h dim", w=16, h=16)

    logits_WHC = head(x_WHD)

    pred_WH = logits_WHC.argmax(axis=-1)
    preds = einops.rearrange(pred_WH, "w h -> (w h)").tolist()
    return [seg_to_img(upsample(pred_WH)), preds]


@beartype.beartype
def unscaled(x: float, max_obs: float) -> float:
    """Scale from [-10, 10] to [10 * -max_obs, 10 * max_obs]."""
    return map_range(x, (-10.0, 10.0), (-10.0 * max_obs, 10.0 * max_obs))


@beartype.beartype
def map_range(
    x: float,
    domain: tuple[float | int, float | int],
    range: tuple[float | int, float | int],
):
    a, b = domain
    c, d = range
    if not (a <= x <= b):
        raise ValueError(f"x={x:.3f} must be in {[a, b]}.")
    return c + (x - a) * (d - c) / (b - a)


@torch.inference_mode
def get_modified_labels(
    i: int,
    latent1: int,
    latent2: int,
    latent3: int,
    value1: float,
    value2: float,
    value3: float,
) -> list[Image.Image | list[int]]:
    sample = vit_dataset[i]
    x = sample["image"][None, ...].to(device)
    x_BPD = rest_of_vit.forward_start(x)

    x_hat_BPD, f_x_BPS, _ = sae(x_BPD)

    err_BPD = x_BPD - x_hat_BPD

    values = torch.tensor(
        [
            unscaled(float(value), top_values[latent].max().item())
            for value, latent in [
                (value1, latent1),
                (value2, latent2),
                (value3, latent3),
            ]
        ],
        device=device,
    )
    f_x_BPS[..., torch.tensor([latent1, latent2, latent3], device=device)] = values

    # Reproduce the SAE forward pass after f_x
    modified_x_hat_BPD = (
        einops.einsum(
            f_x_BPS,
            sae.W_dec,
            "batch patches d_sae, d_sae d_vit -> batch patches d_vit",
        )
        + sae.b_dec
    )
    modified_BPD = err_BPD + modified_x_hat_BPD

    modified_BPD = rest_of_vit.forward_end(modified_BPD)

    logits_BPC = head(modified_BPD)
    pred_P = logits_BPC[0].argmax(axis=-1)
    pred_WH = einops.rearrange(pred_P, "(w h) -> w h", w=16, h=16)
    return seg_to_img(upsample(pred_WH)), pred_P.tolist()


@jaxtyped(typechecker=beartype.beartype)
@torch.inference_mode
def upsample(
    x_WH: Int[Tensor, "width_ps height_ps"],
) -> UInt8[Tensor, "width_px height_px"]:
    return (
        torch.nn.functional.interpolate(
            x_WH.view((1, 1, 16, 16)).float(),
            scale_factor=28,
        )
        .view((448, 448))
        .type(torch.uint8)
    )


with gr.Blocks() as demo:
    image_number = gr.Number(label="Validation Example")

    input_image_base64 = gr.Text(label="Image in Base64")
    true_labels_base64 = gr.Text(label="Labels in Base64")

    get_input_image_btn = gr.Button(value="Get Input Image")
    get_input_image_btn.click(
        get_image,
        inputs=[image_number],
        outputs=[input_image_base64, true_labels_base64, image_number],
        api_name="get-image",
    )

    # input_image = gr.Image(
    #     label="Input Image",
    #     sources=["upload", "clipboard"],
    #     type="pil",
    #     interactive=True,
    # )
    # patch_numbers = gr.CheckboxGroup(label="Image Patch", choices=list(range(256)))
    # top_latent_numbers = gr.CheckboxGroup(label="Top Latents")
    # top_latent_numbers = [
    #     gr.Number(label="Top Latents #{j+1}") for j in range(n_sae_latents)
    # ]
    # sae_example_images = [
    #     gr.Image(label=f"Latent #{j}, Example #{i + 1}", format="png")
    #     for i in range(n_sae_examples)
    #     for j in range(n_sae_latents)
    # ]

    patches_json = gr.JSON(label="Patches", value=[])
    activations_json = gr.JSON(label="Activations", value=[])

    get_sae_activations_btn = gr.Button(value="Get SAE Activations")
    get_sae_activations_btn.click(
        get_sae_activations,
        inputs=[image_number, patches_json],
        outputs=[activations_json],
        api_name="get-sae-examples",
    )
    # semseg_image = gr.Image(label="Semantic Segmentaions", format="png")
    # semseg_colors = gr.CheckboxGroup(
    #     label="Sem Seg Colors", choices=list(range(1, 151))
    # )

    # get_pred_labels_btn = gr.Button(value="Get Pred. Labels")
    # get_pred_labels_btn.click(
    #     get_pred_labels,
    #     inputs=[image_number],
    #     outputs=[semseg_image, semseg_colors],
    #     api_name="get-pred-labels",
    # )

    # get_true_labels_btn = gr.Button(value="Get True Label")
    # get_true_labels_btn.click(
    #     get_true_labels,
    #     inputs=[image_number],
    #     outputs=semseg_image,
    #     api_name="get-true-labels",
    # )

    # latent_numbers = [gr.Number(label=f"Latent {i + 1}") for i in range(3)]
    # value_sliders = [
    #     gr.Slider(label=f"Value {i + 1}", minimum=-10, maximum=10) for i in range(3)
    # ]

    # get_modified_labels_btn = gr.Button(value="Get Modified Label")
    # get_modified_labels_btn.click(
    #     get_modified_labels,
    #     inputs=[image_number] + latent_numbers + value_sliders,
    #     outputs=[semseg_image, semseg_colors],
    #     api_name="get-modified-labels",
    # )

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