import base64
import functools
import io
import json
import logging
import math
import os
import pathlib
import random
import typing

import beartype
import einops.layers.torch
import gradio as gr
import numpy as np
import open_clip
import requests
import saev.nn
import torch
from jaxtyping import Float, jaxtyped
from PIL import Image, ImageDraw
from torch import Tensor
from torchvision import transforms

log_format = "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s"
logging.basicConfig(level=logging.INFO, format=log_format)
logger = logging.getLogger("app.py")


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

DEBUG = False
"""Whether we are debugging."""

n_sae_latents = 5
"""Number of SAE latents to show."""

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

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

vit_ckpt = "ViT-B-16/openai"
"""CLIP checkpoint."""

n_patches_per_img: int = 196
"""Number of patches per image in vit_ckpt."""

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

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


r2_url = "https://pub-289086e849214430853bc87bd8964988.r2.dev/"


logger.info("Set global constants.")


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

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

    orig_url: str
    """The URL or path to access the original example image."""
    highlighted_url: typing.NotRequired[str]
    """The URL or path to access the SAE-highlighted image."""
    target: int
    """Class ID."""


@beartype.beartype
class SaeLatent(typing.TypedDict):
    """Represents a single SAE latent."""

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

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

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


###########
# Helpers #
###########


@beartype.beartype
def get_cache_dir() -> str:
    """
    Get cache directory from environment variables, defaulting to the current working directory (.)

    Returns:
        A path to a cache directory (might not exist yet).
    """
    cache_dir = ""
    for var in ("HF_HOME", "HF_HUB_CACHE"):
        cache_dir = cache_dir or os.environ.get(var, "")
    return cache_dir or "."


@beartype.beartype
def load_model(fpath: str | pathlib.Path, *, device: str = "cpu") -> torch.nn.Module:
    """
    Loads a linear layer from disk.
    """
    with open(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)
    return model


@beartype.beartype
@functools.lru_cache(maxsize=512)
def get_dataset_img(i: int) -> Image.Image:
    return Image.open(requests.get(r2_url + image_fpaths[i], stream=True).raw)


@beartype.beartype
def to_sized(img: Image.Image) -> Image.Image:
    # Copied from contrib/classification/transforms.py:for_webapp()
    w, h = img.size
    if w > h:
        resize_w = int(w * 512 / h)
        resize_px = (resize_w, 512)

        margin_x = (resize_w - 448) // 2
        crop_px = (margin_x, 32, 448 + margin_x, 480)
    else:
        resize_h = int(h * 512 / w)
        resize_px = (512, resize_h)
        margin_y = (resize_h - 448) // 2
        crop_px = (32, margin_y, 480, 448 + margin_y)

    return img.resize(resize_px, resample=Image.Resampling.BICUBIC).crop(crop_px)


@beartype.beartype
def img_to_base64(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.save(buf, format="webp", lossless=True)
    b64 = base64.b64encode(buf.getvalue())
    s64 = b64.decode("utf8")
    return "data:image/webp;base64," + s64


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


@jaxtyped(typechecker=beartype.beartype)
class SplitClip(torch.nn.Module):
    def __init__(self, *, n_end_layers: int):
        super().__init__()

        if vit_ckpt.startswith("hf-hub:"):
            clip, _ = open_clip.create_model_from_pretrained(
                vit_ckpt, cache_dir=get_cache_dir()
            )
        else:
            arch, ckpt = vit_ckpt.split("/")
            clip, _ = open_clip.create_model_from_pretrained(
                arch, pretrained=ckpt, cache_dir=get_cache_dir()
            )
        model = clip.visual
        model.proj = None
        model.output_tokens = True  # type: ignore
        self.vit = model.eval()
        assert not isinstance(self.vit, open_clip.timm_model.TimmModel)

        self.n_end_layers = n_end_layers

    @staticmethod
    def _expand_token(token, batch_size: int):
        return token.view(1, 1, -1).expand(batch_size, -1, -1)

    def forward_start(self, x: Float[Tensor, "batch channels width height"]):
        x = self.vit.conv1(x)  # shape = [*, width, grid, grid]
        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]

        # class embeddings and positional embeddings
        x = torch.cat(
            [self._expand_token(self.vit.class_embedding, x.shape[0]).to(x.dtype), x],
            dim=1,
        )
        # shape = [*, grid ** 2 + 1, width]
        x = x + self.vit.positional_embedding.to(x.dtype)

        x = self.vit.patch_dropout(x)
        x = self.vit.ln_pre(x)
        for r in self.vit.transformer.resblocks[: -self.n_end_layers]:
            x = r(x)
        return x

    def forward_end(self, x: Float[Tensor, "batch n_patches dim"]):
        for r in self.vit.transformer.resblocks[-self.n_end_layers :]:
            x = r(x)

        x = self.vit.ln_post(x)
        pooled, _ = self.vit._global_pool(x)
        if self.vit.proj is not None:
            pooled = pooled @ self.vit.proj

        return pooled


# ViT
split_vit = SplitClip(n_end_layers=1)
split_vit = split_vit.to(device)
logger.info("Initialized CLIP ViT.")

# Linear classifier
clf_ckpt_fpath = CWD / "ckpts" / "clf.pt"
clf = load_model(clf_ckpt_fpath)
clf = clf.to(device).eval()
logger.info("Loaded linear classifier.")

# SAE
sae_ckpt_fpath = CWD / "ckpts" / "sae.pt"
sae = saev.nn.load(sae_ckpt_fpath.as_posix())
sae.to(device).eval()
logger.info("Loaded SAE.")


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

human_transform = transforms.Compose([
    transforms.Resize(512, interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.CenterCrop((448, 448)),
    transforms.ToTensor(),
    einops.layers.torch.Rearrange("channels width height -> width height channels"),
])

arch, ckpt = vit_ckpt.split("/")
_, vit_transform = open_clip.create_model_from_pretrained(
    arch, pretrained=ckpt, cache_dir=get_cache_dir()
)


with open(CWD / "data" / "image_fpaths.json") as fd:
    image_fpaths = json.load(fd)


with open(CWD / "data" / "image_labels.json") as fd:
    img_labels = json.load(fd)


logger.info("Loaded all datasets.")

#############
# Variables #
#############


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


top_img_i = load_tensor(CWD / "data" / "top_img_i.pt")
top_values = load_tensor(CWD / "data" / "top_values_uint8.pt")
sparsity = load_tensor(CWD / "data" / "sparsity.pt")


mask = torch.ones((sae.cfg.d_sae), dtype=bool)
mask = mask & (sparsity < max_frequency)


#############
# Inference #
#############


@beartype.beartype
def get_img(img_i: int) -> Example:
    img = get_dataset_img(img_i)
    img = human_transform(img)
    return {
        "orig_url": img_to_base64(Image.fromarray((img * 255).to(torch.uint8).numpy())),
        "target": img_labels[img_i],
    }


@beartype.beartype
def get_random_class_img(cls: int) -> Example:
    indices = [i for i, tgt in enumerate(img_labels) if tgt == cls]
    i = random.choice(indices)

    img = get_dataset_img(i)
    img = human_transform(img)
    return {
        "orig_url": img_to_base64(Image.fromarray((img * 255).to(torch.uint8).numpy())),
        "target": cls,
    }


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

    logger.info("Getting SAE examples for patches %s.", patches)

    img = get_dataset_img(img_i)
    x_BCWH = vit_transform(img)[None, ...].to(device)
    x_BPD = split_vit.forward_start(x_BCWH)
    # Need to add 1 to account for [CLS] token.
    vit_acts_MD = x_BPD[0, [p + 1 for p in patches]].to(device)

    _, f_x_MS, _ = sae(vit_acts_MD)
    f_x_S = f_x_MS.sum(axis=0)

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

    sae_latents = []
    for latent in latents:
        intermediates, 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

            example_img = get_dataset_img(i_im)
            intermediates.append({
                "img": example_img,
                "patches": values_p,
                "target": img_labels[i_im],
            })
            seen_i_im.add(i_im)

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

        examples = []
        for intermediate in intermediates[:n_latent_examples]:
            img_sized = to_sized(intermediate["img"])
            examples.append({
                "orig_url": img_to_base64(img_sized),
                "highlighted_url": img_to_base64(
                    add_highlights(
                        img_sized,
                        intermediate["patches"].to(float).numpy(),
                        upper=upper,
                    )
                ),
                "target": intermediate["target"],
            })

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

    return sae_latents


@torch.inference_mode
def get_pred_dist(i: int) -> dict[int, float]:
    img = get_dataset_img(i)
    x = vit_transform(img)[None, ...].to(device)
    x_BPD = split_vit.forward_start(x)
    x_BD = split_vit.forward_end(x_BPD)

    logits_BC = clf(x_BD)

    probs = torch.nn.functional.softmax(logits_BC[0], dim=0).cpu().tolist()
    return {i: prob for i, prob in enumerate(probs)}


@torch.inference_mode
def get_modified_dist(
    image_i: int,
    patches: list[int],
    latent1: int,
    latent2: int,
    latent3: int,
    value1: float,
    value2: float,
    value3: float,
) -> dict[int, float]:
    img = get_dataset_img(image_i)
    x = vit_transform(img)[None, ...].to(device)
    x_BPD = split_vit.forward_start(x)

    cls_B1D, x_BPD = x_BPD[:, :1, :], x_BPD[:, 1:, :]

    x_hat_BPD, f_x_BPS, _ = sae(x_BPD)
    err_BPD = x_BPD - x_hat_BPD

    values = torch.tensor(
        [
            unscaled(value, top_values[latent].max().item())
            for value, latent in [
                (value1, latent1),
                (value2, latent2),
                (value3, latent3),
            ]
        ],
        device=device,
    )
    patches = torch.tensor(patches, device=device)
    latents = torch.tensor([latent1, latent2, latent3], device=device)
    f_x_BPS[:, patches[:, None], latents[None, :]] = 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 = torch.cat([cls_B1D, err_BPD + modified_x_hat_BPD], axis=1)

    modified_BD = split_vit.forward_end(modified_BPD)
    logits_BC = clf(modified_BD)

    probs = torch.nn.functional.softmax(logits_BC[0], dim=0).cpu().tolist()
    return {i: prob for i, prob in enumerate(probs)}


@beartype.beartype
def unscaled(x: float | int, max_obs: float | int) -> float:
    """Scale from [-20, 20] to [20 * -max_obs, 20 * max_obs]."""
    return map_range(x, (-20.0, 20.0), (-20.0 * max_obs, 20.0 * max_obs))


@beartype.beartype
def map_range(
    x: float | int,
    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)


@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:
    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)

    colors = np.zeros((len(patches), 3), dtype=np.uint8)
    colors[:, 0] = ((patches / (upper + 1e-9)) * 255).astype(np.uint8)

    for p, (val, color) in enumerate(zip(patches, colors)):
        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=(*color, int(opacity * val * 255)),
        )

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


#############
# Interface #
#############


with gr.Blocks() as demo:
    ###########
    # get-img #
    ###########

    # Inputs
    number = gr.Number(label="Number", precision=0)

    # Outputs
    json_out = gr.JSON(label="get_img_out", value={})

    get_img_btn = gr.Button(value="Get Input Image")
    get_img_btn.click(
        get_img,
        inputs=[number],
        outputs=[json_out],
        api_name="get-img",
    )

    ########################
    # get-random-class-img #
    ########################

    get_random_class_image_btn = gr.Button(value="Get Random Class Image")
    get_img_btn.click(
        get_random_class_img,
        inputs=[number],
        outputs=[json_out],
        api_name="get-random-class-img",
    )

    patch_numbers = gr.CheckboxGroup(
        label="Image Patch", choices=list(range(n_patches_per_img))
    )
    get_sae_latents_btn = gr.Button(value="Get SAE Examples")
    get_sae_latents_btn.click(
        get_sae_latents,
        inputs=[number, patch_numbers],
        outputs=json_out,
        api_name="get-sae-latents",
        concurrency_limit=16,
    )

    pred_dist = gr.Label(label="Pred. Dist.")
    get_pred_dist_btn = gr.Button(value="Get Pred. Distribution")
    get_pred_dist_btn.click(
        get_pred_dist,
        inputs=[number],
        outputs=[pred_dist],
        api_name="get-preds",
    )

    latent_numbers = [gr.Number(label=f"Latent {i + 1}", precision=0) 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_dist_btn = gr.Button(value="Get Modified Label")
    get_modified_dist_btn.click(
        get_modified_dist,
        inputs=[number, patch_numbers] + latent_numbers + value_sliders,
        outputs=[pred_dist],
        api_name="get-modified",
        concurrency_limit=16,
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2, max_size=32)
    demo.launch()