import base64
import html
from io import BytesIO

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import torch
from einops import rearrange
from matplotlib.colors import LinearSegmentedColormap

DECODER_HEATMAP_COLORMAP = (
    LinearSegmentedColormap.from_list(
        "decoder_heatmap",
        (
            "#ffffff",
            "#1aff66",
            "#245705",
        ),
    )
)


def _decode_product_with_mean_cross_attention(
    model,
    encoded_lattice_features: torch.Tensor,
    lattice_feature_padding_mask: torch.Tensor,
    product_decoder_inputs: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    final_decoder_cross_attention = (
        model.decoder.layers[-1].multihead_attn
    )

    captured_cross_attention_inputs = {}

    def capture_cross_attention_inputs(
        module,
        args,
        kwargs,
    ):
        captured_cross_attention_inputs["query"] = args[0]
        captured_cross_attention_inputs["key"] = args[1]
        captured_cross_attention_inputs["value"] = args[2]
        captured_cross_attention_inputs["attn_mask"] = kwargs.get(
            "attn_mask"
        )
        captured_cross_attention_inputs[
            "key_padding_mask"
        ] = kwargs.get("key_padding_mask")
        captured_cross_attention_inputs["is_causal"] = kwargs.get(
            "is_causal",
            False,
        )

    hook_handle = (
        final_decoder_cross_attention.register_forward_pre_hook(
            capture_cross_attention_inputs,
            with_kwargs=True,
        )
    )

    try:
        product_logits = model._decode_product(
            encoded_lattice_features=encoded_lattice_features,
            lattice_feature_padding_mask=(
                lattice_feature_padding_mask
            ),
            product_decoder_inputs=product_decoder_inputs,
        )
    finally:
        hook_handle.remove()

    if not captured_cross_attention_inputs:
        raise RuntimeError(
            "The final decoder cross-attention inputs were not captured."
        )

    _, cross_attention_weights_by_head = (
        final_decoder_cross_attention(
            query=captured_cross_attention_inputs["query"],
            key=captured_cross_attention_inputs["key"],
            value=captured_cross_attention_inputs["value"],
            key_padding_mask=captured_cross_attention_inputs[
                "key_padding_mask"
            ],
            need_weights=True,
            attn_mask=captured_cross_attention_inputs[
                "attn_mask"
            ],
            average_attn_weights=False,
            is_causal=captured_cross_attention_inputs[
                "is_causal"
            ],
        )
    )

    mean_cross_attention_for_current_position = (
        cross_attention_weights_by_head[
            :,
            :,
            -1,
            :,
        ].mean(dim=1)
    )

    return (
        product_logits,
        mean_cross_attention_for_current_position,
    )


def _render_cross_attention_heatmap(
    mean_cross_attention: torch.Tensor,
    a_digits_least_significant_first: list[str],
    b_digits_least_significant_first: list[str],
    maximum_attention_value: float,
) -> str:
    figure, axis = plt.subplots(
        figsize=(5.0, 4.0),
    )

    heatmap = axis.imshow(
        mean_cross_attention.numpy(),
        cmap=DECODER_HEATMAP_COLORMAP,
        vmin=0.0,
        vmax=maximum_attention_value,
        aspect="auto",
    )

    axis.set_xticks(
        range(len(b_digits_least_significant_first))
    )
    axis.set_xticklabels(
        b_digits_least_significant_first,
        fontsize=8,
    )
    axis.set_yticks(
        range(len(a_digits_least_significant_first))
    )
    axis.set_yticklabels(
        a_digits_least_significant_first,
        fontsize=8,
    )

    axis.set_xlabel(
        "b digits (least-significant first)",
        fontsize=9,
    )
    axis.set_ylabel(
        "a digits (least-significant first)",
        fontsize=9,
    )

    colorbar = figure.colorbar(
        heatmap,
        ax=axis,
        fraction=0.05,
        pad=0.03,
    )
    colorbar.set_label(
        "Mean attention weight",
        fontsize=8,
    )
    colorbar.ax.tick_params(
        labelsize=8,
    )

    figure.tight_layout(
        pad=0.6,
    )

    image_buffer = BytesIO()
    figure.savefig(
        image_buffer,
        format="png",
        dpi=130,
        bbox_inches="tight",
    )
    plt.close(figure)

    return base64.b64encode(
        image_buffer.getvalue()
    ).decode("ascii")


def generate_decoder_analysis_html(
    model,
    model_config: dict,
    tokenizer,
    a_input: str,
    b_input: str,
    sequence_length: int,
    expected_predicted_product: str,
    device: str,
) -> str:
    a_token_ids = tokenizer.encode(a_input)
    b_token_ids = tokenizer.encode(b_input)

    a_token_ids.reverse()
    b_token_ids.reverse()

    a_token_ids.extend(
        [model_config["pad_id"]]
        * (sequence_length - len(a_token_ids))
    )
    b_token_ids.extend(
        [model_config["pad_id"]]
        * (sequence_length - len(b_token_ids))
    )

    a = torch.tensor(
        [a_token_ids],
        dtype=torch.long,
        device=device,
    )
    b = torch.tensor(
        [b_token_ids],
        dtype=torch.long,
        device=device,
    )

    (
        encoded_lattice_features,
        lattice_feature_padding_mask,
    ) = model._encode_lattice(
        a,
        b,
    )

    product_decoder_inputs = torch.full(
        (1, 1),
        model.bos_id,
        dtype=torch.long,
        device=device,
    )

    generated_product_token_ids = []
    mean_cross_attention_by_product_position = []

    maximum_product_sequence_length = 2 * sequence_length + 1

    for product_position in range(
        maximum_product_sequence_length
    ):
        (
            product_logits,
            mean_cross_attention_for_current_position,
        ) = _decode_product_with_mean_cross_attention(
            model=model,
            encoded_lattice_features=encoded_lattice_features,
            lattice_feature_padding_mask=(
                lattice_feature_padding_mask
            ),
            product_decoder_inputs=product_decoder_inputs,
        )

        next_product_token = product_logits[
            :,
            -1,
        ].argmax(
            dim=-1
        )

        next_product_token_id = int(
            next_product_token.item()
        )

        if next_product_token_id == model.eos_id:
            break

        generated_product_token_ids.append(
            next_product_token_id
        )

        mean_cross_attention_lattice = rearrange(
            mean_cross_attention_for_current_position[
                0
            ].float().cpu(),
            "(row column) -> row column",
            row=sequence_length,
            column=sequence_length,
        )

        mean_cross_attention_lattice = (
            mean_cross_attention_lattice[
                : len(a_input),
                : len(b_input),
            ]
        )

        mean_cross_attention_by_product_position.append(
            mean_cross_attention_lattice
        )

        product_decoder_inputs = torch.cat(
            [
                product_decoder_inputs,
                next_product_token.unsqueeze(1),
            ],
            dim=1,
        )

    generated_product = tokenizer.decode(
        generated_product_token_ids
    )[::-1]

    if generated_product != expected_predicted_product:
        raise RuntimeError(
            "Decoder analysis did not reproduce the displayed prediction."
        )

    if not mean_cross_attention_by_product_position:
        return (
            "<p>No product digits were generated before the EOS token.</p>"
        )

    maximum_attention_value = max(
        float(mean_cross_attention.max().item())
        for mean_cross_attention
        in mean_cross_attention_by_product_position
    )

    if maximum_attention_value == 0.0:
        maximum_attention_value = 1.0

    a_digits_least_significant_first = list(
        reversed(a_input)
    )
    b_digits_least_significant_first = list(
        reversed(b_input)
    )

    table_rows = []

    for (
        product_position,
        (
            product_token_id,
            mean_cross_attention,
        ),
    ) in enumerate(
        zip(
            generated_product_token_ids,
            mean_cross_attention_by_product_position,
        )
    ):
        product_token_value = tokenizer.int_to_char.get(
            product_token_id,
            str(product_token_id),
        )

        heatmap_image = _render_cross_attention_heatmap(
            mean_cross_attention=mean_cross_attention,
            a_digits_least_significant_first=(
                a_digits_least_significant_first
            ),
            b_digits_least_significant_first=(
                b_digits_least_significant_first
            ),
            maximum_attention_value=maximum_attention_value,
        )

        table_rows.append(
            f"""
            <tr>
                <td class="decoder-position">{product_position}</td>
                <td class="decoder-value">{html.escape(product_token_value)}</td>
                <td class="decoder-heatmap">
                    <img
                        src="data:image/png;base64,{heatmap_image}"
                        alt="Mean decoder cross-attention at position {product_position}"
                        loading="lazy"
                    >
                </td>
            </tr>
            """
        )

    return f"""
    <style>
        .decoder-analysis-summary {{
            margin: 0 0 1rem 0;
        }}

        .decoder-analysis-table-container {{
            max-height: 900px;
            overflow: auto;
            border: 1px solid var(--border-color-primary);
            border-radius: 8px;
        }}

        .decoder-analysis-table {{
            width: 100%;
            border-collapse: collapse;
            background: var(--background-fill-primary);
        }}

        .decoder-analysis-table th,
        .decoder-analysis-table td {{
            padding: 0.75rem;
            text-align: center;
            vertical-align: middle;
            border-bottom: 1px solid var(--border-color-primary);
        }}

        .decoder-analysis-table th {{
            position: sticky;
            top: 0;
            z-index: 1;
            background: var(--background-fill-secondary);
        }}

        .decoder-position {{
            width: 90px;
            font-size: 1rem;
            font-weight: 600;
        }}

        .decoder-value {{
            width: 90px;
            font-size: 1.35rem;
            font-weight: 700;
        }}

        .decoder-heatmap img {{
            display: block;
            width: min(100%, 520px);
            height: auto;
            margin: 0 auto;
        }}
    </style>

    <div class="decoder-analysis-summary">
        <h3>Decoder cross-attention analysis</h3>
        <p>
            <strong>{html.escape(a_input)} × {html.escape(b_input)} =
            {html.escape(expected_predicted_product)}</strong>
        </p>
        <p>
            Final decoder layer, averaged across
            {model.decoder.layers[-1].multihead_attn.num_heads} attention
            heads. Position 0 is the least-significant generated product
            digit. All heatmaps use the same color scale.
        </p>
    </div>

    <div class="decoder-analysis-table-container">
        <table class="decoder-analysis-table">
            <thead>
                <tr>
                    <th>Position</th>
                    <th>Value</th>
                    <th>Mean cross-attention</th>
                </tr>
            </thead>
            <tbody>
                {"".join(table_rows)}
            </tbody>
        </table>
    </div>
    """