import gradio as gr
from sae_auto_interp.sae import Sae
from sae_auto_interp.utils import maybe_load_llava_model, load_single_sae
from sae_auto_interp.features.features import upsample_mask
import torch
from transformers import AutoTokenizer
from PIL import Image
import spaces

CITATION_BUTTON_TEXT = """
@misc{zhang2024largemultimodalmodelsinterpret,
      title={Large Multi-modal Models Can Interpret Features in Large Multi-modal Models},
      author={Kaichen Zhang and Yifei Shen and Bo Li and Ziwei Liu},
      year={2024},
      eprint={2411.14982},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2411.14982},
}
"""

INSTRUCTIONS = """
## Instructions to use the demo
You can use this demo to :
    1. Visualize the activations of the model for a given image.
    2. Generate text with a specific feature clamped to a certain value.

### Visualization of Activations
1. Upload an image. (or use an example)
2. Click on the "Submit" button to visualize the activations. The top-100 features will be displayed. (It might contains lots of low level features that activates on many patterns so explainable features might not rank very high)
3. Use the slider to select a feature number.
4. Click on the "Visualize" button to see the activation of that feature.

### Steering Model
1. Use the slider to select a feature number.
2. Use the number input to select the feature strength.
3. Type the text input.
4. Upload an image. (optional)
5. Click on the "Submit" button to generate text with the selected feature clamped to the selected strength.
"""

cached_tensor = None
topk_indices = None

sunglasses_file_path = "assets/sunglasses.jpg"
greedy_file_path = "assets/greedy.jpg"
railway_file_path = "assets/railway.jpg"
happy_file_path = "assets/happy.jpg"

@spaces.GPU
def generate_activations(image):
    prompt = "<image>"
    inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
    global cached_tensor, topk_indices

    def hook(module: torch.nn.Module, _, outputs):
        global cached_tensor, topk_indices
        # Maybe unpack tuple outputs
        if isinstance(outputs, tuple):
            unpack_outputs = list(outputs)
        else:
            unpack_outputs = list(outputs)
        latents = sae.pre_acts(unpack_outputs[0])
        # When the tokenizer is llama and text is None (image only)
        # I skip the first bos tokens
        if "llama" in tokenizer.name_or_path:
            latents = latents[:, 1:, :]

        topk = torch.topk(
            latents, k=sae.cfg.k, dim=-1
        )
        # make all other values 0
        result = torch.zeros_like(latents)
        # results (bs, seq, num_latents)
        result.scatter_(-1, topk.indices, topk.values)
        cached_tensor = result.detach().cpu()
        topk_indices = (
            latents.squeeze(0).mean(dim=0).topk(k=100).indices.detach().cpu()
        )

    handles = [hooked_module.register_forward_hook(hook)]
    try:
        with torch.no_grad():
            outputs = model(
                input_ids=inputs["input_ids"].to("cuda"),
                pixel_values=inputs["pixel_values"].to("cuda"),
                image_sizes=inputs["image_sizes"].to("cuda"),
                attention_mask=inputs["attention_mask"].to("cuda"),
            )
    finally:
        for handle in handles:
            handle.remove()
    
    torch.cuda.empty_cache()
    return topk_indices
    

def visualize_activations(image, feature_num):
    base_img_tokens = 576
    patch_size = 24
    # Using Cached tensor
    # select the feature_num-th feature
    # Then keeping the first 576 tokens
    base_image_activations = cached_tensor[0, :base_img_tokens, feature_num].view(patch_size, patch_size)

    upsampled_image_mask = upsample_mask(base_image_activations, (336, 336))
    

    background = Image.new("L", (336, 336), 0).convert("RGB")

    # Somehow as I looked closer into the llava-hf preprocessing code,
    # I found out that they don't use the padded image as the base image feat
    # but use the simple resized image. This is different from original llava but
    # we align to llava-hf for now as we use llava-hf
    resized_image = image.resize((336, 336))
    activation_images = Image.composite(background, resized_image, upsampled_image_mask).convert("RGB")

    return activation_images

@spaces.GPU
def clamp_features_max(
    sae: Sae, feature: int, hooked_module: torch.nn.Module, k: float = 10
):
    def hook(module: torch.nn.Module, _, outputs):
        # Maybe unpack tuple outputs
        if isinstance(outputs, tuple):
            unpack_outputs = list(outputs)
        else:
            unpack_outputs = list(outputs)
        latents = sae.pre_acts(unpack_outputs[0])
        # Only clamp the feature for the first forward
        if latents.shape[1] != 1:
            latents[:, :, feature] = k
        top_acts, top_indices = sae.select_topk(latents)
        sae_out = sae.decode(top_acts[0], top_indices[0]).unsqueeze(0).to(torch.float16)
        unpack_outputs[0] = sae_out
        if isinstance(outputs, tuple):
            outputs = tuple(unpack_outputs)
        else:
            outputs = unpack_outputs[0]
        return outputs

    handles = [hooked_module.register_forward_hook(hook)]

    return handles

@spaces.GPU
def generate_with_clamp(feature_idx, feature_strength, text, image, chat_history):
    if not isinstance(feature_idx, int):
        feature_idx = int(feature_idx)
    if not isinstance(feature_strength, float):
        feature_strength = float(feature_strength)
    
    conversation = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": text},
            ],
        },
    ]
    if image is not None:
        conversation[0]["content"].append(
            {"type": "image"},
        )

        chat_history.append({"role": "user", "content": gr.Image(value=image)})
    chat_history.append({"role": "user", "content": text})
    

    prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)

    inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
    handles = clamp_features_max(sae, feature_idx, hooked_module, k=feature_strength)
    try:
        with torch.no_grad():
            output = model.generate(**inputs, max_new_tokens=512)
        cont = output[:, inputs["input_ids"].shape[-1] :]
    finally:
        for handle in handles:
            handle.remove()

    text = processor.batch_decode(cont, skip_special_tokens=True)[0]
    chat_history.append(
        {
            "role": "assistant",
            "content": text,
        }
    )

    return chat_history


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Large Multi-modal Models Can Interpret Features in Large Multi-modal Models

        🔍 [ArXiv Paper](https://arxiv.org/abs/2411.14982) | 🏠 [LMMs-Lab Homepage](https://lmms-lab.framer.ai) | 🤗 [Huggingface Collections](https://huggingface.co./collections/lmms-lab/llava-sae-674026e4e7bc8c29c70bc3a3) | [GitHub Repo](https://github.com/EvolvingLMMs-Lab/multimodal-sae)
        """
    )
    with gr.Accordion("ℹ️ Instructions", open=False):
        gr.Markdown(INSTRUCTIONS)

    with gr.Tabs(elem_classes="tab-buttons") as tabs:
        with gr.TabItem("Visualization of Activations", elem_id="visualization", id=0):
            with gr.Row():
                with gr.Column():
                    image = gr.Image(type="pil", interactive=True, label="Sample Image")
                    topk_features = gr.Textbox(value=topk_indices, placeholder="Top 100 Features", label="Top 100 Features")
                    with gr.Row():
                        clear_btn = gr.ClearButton([image, topk_features], value="Clear")
                        submit_btn = gr.Button("Submit", variant="primary")
                        submit_btn.click(generate_activations, inputs=[image], outputs=[topk_features])
                with gr.Column():
                    output = gr.Image(label="Activation Visualization")
                    feature_num = gr.Slider(1, 131072, 1, 1, label="Feature Number", interactive=True)
                    visualize_btn = gr.Button("Visualize", variant="primary")
                    visualize_btn.click(visualize_activations, inputs=[image, feature_num], outputs=[output])
                
            dummy_text = gr.Textbox(visible=False, label="Explanation")
            gr.Examples(
                [
                    ["assets/sunglasses.jpg", 10, "Sunglasses"],
                    ["assets/greedy.jpg", 14, "Greedy eating"],
                    ["assets/railway.jpg", 28, "Railway tracks"],
                ],
                inputs=[image, feature_num, dummy_text],
                label="Examples",
            )

        with gr.TabItem("Steering Model", elem_id="steering", id=2):
            chatbot = gr.Chatbot(type="messages")
            with gr.Row(variant="compact", equal_height=True):
                feature_num = gr.Slider(1, 131072, 1, 1, label="Feature Number", interactive=True)
                feature_strength = gr.Number(value=50, label="Feature Strength", interactive=True)
            with gr.Row(variant="compact", equal_height=True):
                text_input = gr.Textbox(label="Text Input", placeholder="Type here", interactive=True)
                image_input = gr.Image(type="pil", label="Image Input", interactive=True, height=250)
            with gr.Row():
                chatbot_clear = gr.ClearButton([text_input, image_input, chatbot], value="Clear")
                chatbot_submit = gr.Button("Submit", variant="primary")
                chatbot_submit.click(
                    generate_with_clamp,
                    inputs=[feature_num, feature_strength, text_input, image_input, chatbot],
                    outputs=[chatbot],
                )
            gr.Examples(
                [
                    [19379, 50, "Look at this image, what is your feeling right now?", happy_file_path],
                    [14, 50, "Tell me a story about Alice and Bob", None],
                    [108692, 50, "What is your feeling right now?", None],
                ],
                inputs=[feature_num, feature_strength, text_input, image_input],
                label="Examples",
            )
            

    with gr.Row():
        with gr.Accordion("📙 Citation", open=False):
            gr.Markdown("```bib\n" + CITATION_BUTTON_TEXT + "\n```")


if __name__ == "__main__":
    tokenizer = AutoTokenizer.from_pretrained("llava-hf/llama3-llava-next-8b-hf")
    sae = load_single_sae("lmms-lab/llama3-llava-next-8b-hf-sae-131k", "model.layers.24")
    model, processor = maybe_load_llava_model(
        "llava-hf/llama3-llava-next-8b-hf",
        rank=0,
        dtype=torch.float16,
        hf_token=None
    )
    hooked_module = model.language_model.get_submodule("model.layers.24")

    demo.launch()