import streamlit as st
from transformer_lens import HookedTransformer, utils
from io import StringIO
import sys
import torch
from functools import partial
import plotly.offline as pyo
import plotly.graph_objs as go
import numpy as np
import plotly.express as px
import circuitsvis as cv




# Little bit of front end for model selector

# Radio buttons
model_name = st.sidebar.radio("Model (only use patching for\nsmall (<4L) models due to memory limits)", [
    "gelu-1l",
    "gelu-2l",
    "gelu-3l",
    "gelu-4l",
    "attn-only-1l",
    "attn-only-2l",
    "attn-only-3l",
    "attn-only-4l",
    "solu-1l",
    "solu-2l",
    "solu-3l",
    "solu-4l",
    "solu-6l",
    "solu-8l",
    "solu-10l",
    "solu-12l",
    "gpt2-small",
    "gpt2-medium",
    #"gpt2-large",
    #"gpt2-xl",
    "roneneldan/TinyStories-1Layer-21M",
    "roneneldan/TinyStories-33M",
    ], index=1)


# Backend code

model = HookedTransformer.from_pretrained(model_name)

def predict_next_token(prompt):
    """
    roneneldan/TinyStories-1M
    roneneldan/TinyStories-3M
    roneneldan/TinyStories-8M
    roneneldan/TinyStories-28M
    roneneldan/TinyStories-33M
    roneneldan/TinyStories-1Layer-21M
    roneneldan/TinyStories-2Layers-33M
    roneneldan/TinyStories-Instruct-1M
    roneneldan/TinyStories-Instruct-3M
    roneneldan/TinyStories-Instruct-8M
    roneneldan/TinyStories-Instruct-28M
    roneneldan/TinyStories-Instruct-33M
    roneneldan/TinyStories-Instuct-1Layer-21M
    roneneldan/TinyStories-Instruct-2Layers-33M
    """
    logits = model(prompt)[0,-1]
    answer_index = logits.argmax()
    answer = model.tokenizer.decode(answer_index)
    answer = f"<b>|{answer}|</b> (answer by {model.cfg.model_name})"
    return answer

def test_prompt(prompt, answer):
    output = StringIO()
    sys.stdout = output
    utils.test_prompt(prompt, answer, model)
    output = output.getvalue()
    return output

def compute_residual_stream_patch(clean_prompt=None, answer=None, corrupt_prompt=None, corrupt_answer=None, layers=None):
    model.reset_hooks()
    clean_answer_index = model.tokenizer.encode(answer)[0]
    corrupt_answer_index = model.tokenizer.encode(corrupt_answer)[0]
    clean_tokens = model.to_str_tokens(clean_prompt)
    _, corrupt_cache = model.run_with_cache(corrupt_prompt)
    # Patching function
    def patch_residual_stream(activations, hook, layer="blocks.6.hook_resid_post", pos=5):
        activations[:, pos, :] = corrupt_cache[layer][:, pos, :]
        return activations
    # Compute logit diffs
    n_layers = len(layers)
    n_pos = len(clean_tokens)
    patching_effect = torch.zeros(n_layers, n_pos)
    for l, layer in enumerate(layers):
        for pos in range(n_pos):
            fwd_hooks = [(layer, partial(patch_residual_stream, layer=layer, pos=pos))]
            prediction_logits = model.run_with_hooks(clean_prompt, fwd_hooks=fwd_hooks)[0, -1]
            patching_effect[l, pos] = prediction_logits[clean_answer_index] - prediction_logits[corrupt_answer_index]
    return patching_effect

def compute_attn_patch(clean_prompt=None, answer=None, corrupt_prompt=None, corrupt_answer=None):
    use_attn_result_prev = model.cfg.use_attn_result
    model.cfg.use_attn_result = True
    clean_answer_index = model.tokenizer.encode(answer)[0]
    corrupt_answer_index = model.tokenizer.encode(corrupt_answer)[0]
    clean_tokens = model.to_str_tokens(clean_prompt)
    _, corrupt_cache = model.run_with_cache(corrupt_prompt)
    # Patching function
    def patch_head_result(activations, hook, head=None, pos=None):
        activations[:, pos, head, :] = corrupt_cache[hook.name][:, pos, head, :]
        return activations

    n_layers = model.cfg.n_layers
    n_heads = model.cfg.n_heads
    n_pos = len(clean_tokens)
    patching_effect = torch.zeros(n_layers*n_heads, n_pos)
    for layer in range(n_layers):
        for head in range(n_heads):
          for pos in range(n_pos):
              fwd_hooks = [(f"blocks.{layer}.attn.hook_result", partial(patch_head_result, head=head, pos=pos))]
              prediction_logits = model.run_with_hooks(clean_prompt, fwd_hooks=fwd_hooks)[0, -1]
              patching_effect[n_heads*layer+head, pos] = prediction_logits[clean_answer_index] - prediction_logits[corrupt_answer_index]
    model.cfg.use_attn_result = use_attn_result_prev
    return patching_effect

def imshow(tensor, xlabel="X", ylabel="Y", zlabel=None, xticks=None, yticks=None, c_midpoint=0.0, c_scale="RdBu", **kwargs):
    tensor = utils.to_numpy(tensor)
    xticks = [str(x) for x in xticks]
    yticks = [str(y) for y in yticks]
    labels = {"x": xlabel, "y": ylabel}
    if zlabel is not None:
        labels["color"] = zlabel
    fig = px.imshow(tensor, x=xticks, y=yticks, labels=labels, color_continuous_midpoint=c_midpoint,
                    color_continuous_scale=c_scale, **kwargs)
    return fig

def plot_residual_stream_patch(clean_prompt=None, answer=None, corrupt_prompt=None, corrupt_answer=None):
    layers = ["blocks.0.hook_resid_pre", *[f"blocks.{i}.hook_resid_post" for i in range(model.cfg.n_layers)]]
    clean_tokens = model.to_str_tokens(clean_prompt)
    token_labels = [f"(pos {i:2}) {t}" for i, t in enumerate(clean_tokens)]
    patching_effect = compute_residual_stream_patch(clean_prompt=clean_prompt, answer=answer, corrupt_prompt=corrupt_prompt, corrupt_answer=corrupt_answer, layers=layers)
    fig = imshow(patching_effect, xticks=token_labels, yticks=layers, xlabel="Position", ylabel="Layer",
       zlabel="Logit Difference", title="Patching residual stream at specific layer and position")
    return fig

def plot_attn_patch(clean_prompt=None, answer=None, corrupt_prompt=None, corrupt_answer=None):
    clean_tokens = model.to_str_tokens(clean_prompt)
    n_layers = model.cfg.n_layers
    n_heads = model.cfg.n_heads
    layerhead_labels = [f"{l}.{h}" for l in range(n_layers) for h in range(n_heads)]
    token_labels = [f"(pos {i:2}) {t}" for i, t in enumerate(clean_tokens)]
    patching_effect = compute_attn_patch(clean_prompt=clean_prompt, answer=answer, corrupt_prompt=corrupt_prompt, corrupt_answer=corrupt_answer)
    return imshow(patching_effect, xticks=token_labels, yticks=layerhead_labels, xlabel="Position", ylabel="Layer.Head",
           zlabel="Logit Difference", title=f"Patching attention outputs for specific layer, head, and position", width=600, height=300+200*n_layers)


# Frontend code
st.title("Simple Trafo Mech Int")
st.subheader("Transformer Mechanistic Interpretability")
st.markdown("Powered by [TransformerLens](https://github.com/neelnanda-io/TransformerLens/)")
st.markdown("For _what_ these plots are, and _why_, see this [tutorial](https://docs.google.com/document/d/1e6cs8d9QNretWvOLsv_KaMp6kSPWpJEW0GWc0nwjqxo/).")

# Predict next token
st.header("Predict the next token")
st.markdown("Just a simple test UI, enter a prompt and the model will predict the next token")
prompt_simple = st.text_input("Prompt:", "Today, the weather is", key="prompt_simple")

if "prompt_simple_output" not in st.session_state:
    st.session_state.prompt_simple_output = None

if st.button("Run model", key="key_button_prompt_simple"):
    res = predict_next_token(prompt_simple)
    st.session_state.prompt_simple_output = res

if st.session_state.prompt_simple_output:
    st.markdown(st.session_state.prompt_simple_output, unsafe_allow_html=True)


# Test prompt
st.header("Verbose test prompt")
st.markdown("Enter a prompt and the correct answer, the model will run the prompt and print the results")

prompt = st.text_input("Prompt:", "The most popular programming language is", key="prompt")
answer = st.text_input("Answer:", " Java", key="answer")

if "test_prompt_output" not in st.session_state:
    st.session_state.test_prompt_output = None

if st.button("Run model", key="key_button_test_prompt"):
    res = test_prompt(prompt, answer)
    st.session_state.test_prompt_output = res
    
if st.session_state.test_prompt_output:
    st.code(st.session_state.test_prompt_output)


# Residual stream patching

st.header("Residual stream patching")
st.markdown("Enter a clean prompt, correct answer, corrupt prompt and corrupt answer, the model will compute the patching effect")

default_clean_prompt = "Her name was Alex Hart. Tomorrow at lunch time Alex"
default_clean_answer = "Hart"
default_corrupt_prompt = "Her name was Alex Carroll. Tomorrow at lunch time Alex"
default_corrupt_answer = "Carroll"

clean_prompt   = st.text_input("Clean Prompt:",   default_clean_prompt)
clean_answer   = st.text_input("Correct Answer:", default_clean_answer)
corrupt_prompt = st.text_input("Corrupt Prompt:", default_corrupt_prompt)
corrupt_answer = st.text_input("Corrupt Answer:", default_corrupt_answer)

if "residual_stream_patch_out" not in st.session_state:
    st.session_state.residual_stream_patch_out = None

if st.button("Run model", key="key_button_residual_stream_patch"):
    fig = plot_residual_stream_patch(clean_prompt=clean_prompt, answer=clean_answer, corrupt_prompt=corrupt_prompt, corrupt_answer=corrupt_answer)
    st.session_state.residual_stream_patch_out = fig

if st.session_state.residual_stream_patch_out:
    st.plotly_chart(st.session_state.residual_stream_patch_out)


# Attention head output

st.header("Attention head output patching")
st.markdown("Enter a clean prompt, correct answer, corrupt prompt and corrupt answer, the model will compute the patching effect")

clean_prompt_attn   = st.text_input("Clean Prompt:",   default_clean_prompt, key="key2_clean_prompt_attn")
clean_answer_attn   = st.text_input("Correct Answer:", default_clean_answer, key="key2_clean_answer_attn")
corrupt_prompt_attn = st.text_input("Corrupt Prompt:", default_corrupt_prompt, key="key2_corrupt_prompt_attn")
corrupt_answer_attn = st.text_input("Corrupt Answer:", default_corrupt_answer, key="key2_corrupt_answer_attn")

if "attn_head_patch_out" not in st.session_state:
    st.session_state.attn_head_patch_out = None

if st.button("Run model", key="key_button_attn_head_patch"):
    fig = plot_attn_patch(clean_prompt=clean_prompt_attn, answer=clean_answer_attn, corrupt_prompt=corrupt_prompt_attn, corrupt_answer=corrupt_answer_attn)
    st.session_state.attn_head_patch_out = fig

if st.session_state.attn_head_patch_out:
    st.plotly_chart(st.session_state.attn_head_patch_out)


# Attention Head Visualization

st.header("Attention Pattern Visualization")
st.markdown("Powered by [CircuitsVis](https://github.com/alan-cooney/CircuitsVis)")
st.markdown("Enter a prompt, show attention patterns")

default_prompt_attn = "Her name was Alex Hart. Tomorrow at lunch time Alex"
prompt_attn   = st.text_input("Prompt:",   default_prompt_attn)

if "attn_html" not in st.session_state:
    st.session_state.attn_html = None

if st.button("Run model", key="key_button_attention_head"):
    _, cache = model.run_with_cache(prompt_attn)
    st.session_state.attn_html = []
    for layer in range(model.cfg.n_layers):
        html = cv.attention.attention_patterns(tokens=model.to_str_tokens(prompt_attn),
                                attention=cache[f'blocks.{layer}.attn.hook_pattern'][0])
        st.session_state.attn_html.append(html.show_code())

if st.session_state.attn_html:
    for layer in range(len(st.session_state.attn_html)):
        st.write(f"Attention patterns Layer {layer}:")
        st.components.v1.html(st.session_state.attn_html[layer], height=500)