"""Util functions for codebook features."""
import re
import typing
from dataclasses import dataclass
from functools import partial
from typing import Optional

import numpy as np
import torch
import torch.nn.functional as F
from termcolor import colored
from tqdm import tqdm


@dataclass
class CodeInfo:
    """Dataclass for codebook info."""

    code: int
    layer: int
    head: Optional[int]
    cb_at: Optional[str] = None

    # for patching interventions
    pos: Optional[int] = None
    code_pos: Optional[int] = -1

    # for description & regex-based interpretation
    description: Optional[str] = None
    regex: Optional[str] = None
    prec: Optional[float] = None
    recall: Optional[float] = None
    num_acts: Optional[int] = None

    def __post_init__(self):
        """Convert to appropriate types."""
        self.code = int(self.code)
        self.layer = int(self.layer)
        if self.head:
            self.head = int(self.head)
        if self.pos:
            self.pos = int(self.pos)
        if self.code_pos:
            self.code_pos = int(self.code_pos)
        if self.prec:
            self.prec = float(self.prec)
            assert 0 <= self.prec <= 1
        if self.recall:
            self.recall = float(self.recall)
            assert 0 <= self.recall <= 1
        if self.num_acts:
            self.num_acts = int(self.num_acts)

    def check_description_info(self):
        """Check if the regex info is present."""
        assert self.num_acts is not None and self.description is not None
        if self.regex is not None:
            assert self.prec is not None and self.recall is not None

    def check_patch_info(self):
        """Check if the patch info is present."""
        # TODO: pos can be none for patching
        assert self.pos is not None and self.code_pos is not None

    def __repr__(self):
        """Return the string representation."""
        repr = f"CodeInfo(code={self.code}, layer={self.layer}, head={self.head}, cb_at={self.cb_at}"
        if self.pos is not None or self.code_pos is not None:
            repr += f", pos={self.pos}, code_pos={self.code_pos}"
        if self.description is not None:
            repr += f", description={self.description}"
        if self.regex is not None:
            repr += f", regex={self.regex}, prec={self.prec}, recall={self.recall}"
        if self.num_acts is not None:
            repr += f", num_acts={self.num_acts}"
        repr += ")"
        return repr


def logits_to_pred(logits, tokenizer, k=5):
    """Convert logits to top-k predictions."""
    sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
    probs = sorted_logits.softmax(dim=-1)
    topk_preds = [tokenizer.convert_ids_to_tokens(e) for e in sorted_indices[:, -1, :k]]
    topk_preds = [
        tokenizer.convert_tokens_to_string([e]) for batch in topk_preds for e in batch
    ]
    return [(topk_preds[i], probs[:, -1, i].item()) for i in range(len(topk_preds))]


def patch_codebook_ids(
    corrupted_codebook_ids, hook, pos, cache, cache_pos=None, code_idx=None
):
    """Patch codebook ids with cached ids."""
    if cache_pos is None:
        cache_pos = pos
    if code_idx is None:
        corrupted_codebook_ids[:, pos] = cache[hook.name][:, cache_pos]
    else:
        for code_id in range(32):
            if code_id in code_idx:
                corrupted_codebook_ids[:, pos, code_id] = cache[hook.name][
                    :, cache_pos, code_id
                ]
            else:
                corrupted_codebook_ids[:, pos, code_id] = -1

    return corrupted_codebook_ids


def logits_to_ave_logit_diff(logits, answer_tokens, per_prompt=False):
    """Calculate the average logit difference between the answer and the other token."""
    # Only the final logits are relevant for the answer
    final_logits = logits[:, -1, :]
    answer_logits = final_logits.gather(dim=-1, index=answer_tokens)
    answer_logit_diff = answer_logits[:, 0] - answer_logits[:, 1]
    if per_prompt:
        return answer_logit_diff
    else:
        return answer_logit_diff.mean()


def normalize_patched_logit_diff(
    patched_logit_diff,
    base_average_logit_diff,
    corrupted_average_logit_diff,
):
    """Normalize the patched logit difference."""
    # Subtract corrupted logit diff to measure the improvement,
    # divide by the total improvement from clean to corrupted to normalise
    # 0 means zero change, negative means actively made worse,
    # 1 means totally recovered clean performance, >1 means actively *improved* on clean performance
    return (patched_logit_diff - corrupted_average_logit_diff) / (
        base_average_logit_diff - corrupted_average_logit_diff
    )


def features_to_tokens(cb_key, cb_acts, num_codes, code=None):
    """Return the set of token ids each codebook feature activates on."""
    codebook_ids = cb_acts[cb_key]

    if code is None:
        features_tokens = [[] for _ in range(num_codes)]
        for i in tqdm(range(codebook_ids.shape[0])):
            for j in range(codebook_ids.shape[1]):
                for k in range(codebook_ids.shape[2]):
                    features_tokens[codebook_ids[i, j, k]].append((i, j))
    else:
        idx0, idx1, _ = np.where(codebook_ids == code)
        features_tokens = list(zip(idx0, idx1))

    return features_tokens


def color_str(s: str, color: str, html: bool):
    """Color the string for html or terminal."""
    if html:
        return f"<span style='color:{color}'>{s}</span>"
    else:
        return colored(s, color)


def color_tokens_red_automata(tokens, red_idx, html=False):
    """Separate states with a dash and color red the tokens in red_idx."""
    ret_string = ""
    itr_over_red_idx = 0
    tokens_enumerate = enumerate(tokens)
    if tokens[0] == "<|endoftext|>":
        next(tokens_enumerate)
        if red_idx[0] == 0:
            itr_over_red_idx += 1
    for i, c in tokens_enumerate:
        if i % 2 == 1:
            ret_string += "-"
        if itr_over_red_idx < len(red_idx) and i == red_idx[itr_over_red_idx]:
            ret_string += color_str(c, "red", html)
            itr_over_red_idx += 1
        else:
            ret_string += c
    return ret_string


def color_tokens_red(tokens, red_idx, n=3, html=False):
    """Color red the tokens in red_idx."""
    ret_string = ""
    last_colored_token_idx = -1
    for i in red_idx:
        c_str = tokens[i]
        if i <= last_colored_token_idx + 2 * n + 1:
            ret_string += "".join(tokens[last_colored_token_idx + 1 : i])
        else:
            ret_string += "".join(
                tokens[last_colored_token_idx + 1 : last_colored_token_idx + n + 1]
            )
            ret_string += " ... "
            ret_string += "".join(tokens[i - n : i])
        ret_string += color_str(c_str, "red", html)
        last_colored_token_idx = i
    ret_string += "".join(
        tokens[
            last_colored_token_idx + 1 : min(last_colored_token_idx + n, len(tokens))
        ]
    )
    return ret_string


def prepare_example_print(
    example_id,
    example_tokens,
    tokens_to_color_red,
    html,
    color_red_fn=color_tokens_red,
):
    """Format example to print."""
    example_output = color_str(example_id, "green", html)
    example_output += (
        ": "
        + color_red_fn(example_tokens, tokens_to_color_red, html=html)
        + ("<br>" if html else "\n")
    )
    return example_output


def tkn_print(
    ll,
    tokens,
    separate_states,
    n=3,
    max_examples=100,
    randomize=False,
    html=False,
    return_example_list=False,
):
    """Format and prints the tokens in ll."""
    if randomize:
        raise NotImplementedError("Randomize not yet implemented.")
    indices = range(len(ll))
    print_output = [] if return_example_list else ""
    curr_ex = ll[0][0]
    total_examples = 0
    tokens_to_color_red = []
    color_red_fn = (
        color_tokens_red_automata if separate_states else partial(color_tokens_red, n=n)
    )
    for idx in indices:
        if total_examples > max_examples:
            break
        i, j = ll[idx]

        if i != curr_ex and curr_ex >= 0:
            curr_ex_output = prepare_example_print(
                curr_ex,
                tokens[curr_ex],
                tokens_to_color_red,
                html,
                color_red_fn,
            )
            total_examples += 1
            if return_example_list:
                print_output.append((curr_ex_output, len(tokens_to_color_red)))
            else:
                print_output += curr_ex_output
            curr_ex = i
            tokens_to_color_red = []
        tokens_to_color_red.append(j)
    curr_ex_output = prepare_example_print(
        curr_ex,
        tokens[curr_ex],
        tokens_to_color_red,
        html,
        color_red_fn,
    )
    if return_example_list:
        print_output.append((curr_ex_output, len(tokens_to_color_red)))
    else:
        print_output += curr_ex_output
        asterisk_str = "********************************************"
        print_output += color_str(asterisk_str, "green", html)
    total_examples += 1

    return print_output


def print_ft_tkns(
    ft_tkns,
    tokens,
    separate_states=False,
    n=3,
    start=0,
    stop=1000,
    indices=None,
    max_examples=100,
    freq_filter=None,
    randomize=False,
    html=False,
    return_example_list=False,
):
    """Print the tokens for the codebook features."""
    indices = list(range(start, stop)) if indices is None else indices
    num_tokens = len(tokens) * len(tokens[0])
    codes, token_act_freqs, token_acts = [], [], []
    for i in indices:
        tkns = ft_tkns[i]
        freq = (len(tkns), 100 * len(tkns) / num_tokens)
        if freq_filter is not None and freq[1] > freq_filter:
            continue
        codes.append(i)
        token_act_freqs.append(freq)
        if len(tkns) > 0:
            tkn_acts = tkn_print(
                tkns,
                tokens,
                separate_states,
                n=n,
                max_examples=max_examples,
                randomize=randomize,
                html=html,
                return_example_list=return_example_list,
            )
            token_acts.append(tkn_acts)
        else:
            token_acts.append("")
    return codes, token_act_freqs, token_acts


def patch_in_codes(run_cb_ids, hook, pos, code, code_pos=None):
    """Patch in the `code` at `run_cb_ids`."""
    pos = slice(None) if pos is None else pos
    code_pos = slice(None) if code_pos is None else code_pos

    if code_pos == "append":
        assert pos == slice(None)
        run_cb_ids = F.pad(run_cb_ids, (0, 1), mode="constant", value=code)
    if isinstance(pos, typing.Iterable) or isinstance(pos, typing.Iterable):
        for p in pos:
            run_cb_ids[:, p, code_pos] = code
    else:
        run_cb_ids[:, pos, code_pos] = code
    return run_cb_ids


def get_cb_layer_name(cb_at, layer_idx, head_idx=None):
    """Get the layer name used to store hooks/cache."""
    if head_idx is None:
        return f"blocks.{layer_idx}.{cb_at}.codebook_layer.hook_codebook_ids"
    else:
        return f"blocks.{layer_idx}.{cb_at}.codebook_layer.codebook.{head_idx}.hook_codebook_ids"


def get_cb_layer_names(layer, patch_types, n_heads):
    """Get the layer names used to store hooks/cache."""
    layer_names = []
    attn_added, mlp_added = False, False
    if "attn_out" in patch_types:
        attn_added = True
        for head in range(n_heads):
            layer_names.append(
                f"blocks.{layer}.attn.codebook_layer.codebook.{head}.hook_codebook_ids"
            )
    if "mlp_out" in patch_types:
        mlp_added = True
        layer_names.append(f"blocks.{layer}.mlp.codebook_layer.hook_codebook_ids")

    for patch_type in patch_types:
        # match patch_type of the pattern attn_\d_head_\d
        attn_head = re.match(r"attn_(\d)_head_(\d)", patch_type)
        if (not attn_added) and attn_head and attn_head[1] == str(layer):
            layer_names.append(
                f"blocks.{layer}.attn.codebook_layer.codebook.{attn_head[2]}.hook_codebook_ids"
            )
        mlp = re.match(r"mlp_(\d)", patch_type)
        if (not mlp_added) and mlp and mlp[1] == str(layer):
            layer_names.append(f"blocks.{layer}.mlp.codebook_layer.hook_codebook_ids")

    return layer_names


def cb_layer_name_to_info(layer_name):
    """Get the layer info from the layer name."""
    layer_name_split = layer_name.split(".")
    layer_idx = int(layer_name_split[1])
    cb_at = layer_name_split[2]
    if cb_at == "mlp":
        head_idx = None
    else:
        head_idx = int(layer_name_split[5])
    return cb_at, layer_idx, head_idx


def get_hooks(code, cb_at, layer_idx, head_idx=None, pos=None):
    """Get the hooks for the codebook features."""
    hook_fns = [
        partial(patch_in_codes, pos=pos, code=code[i]) for i in range(len(code))
    ]
    return [
        (get_cb_layer_name(cb_at[i], layer_idx[i], head_idx[i]), hook_fns[i])
        for i in range(len(code))
    ]


def run_with_codes(
    input, cb_model, code, cb_at, layer_idx, head_idx=None, pos=None, prepend_bos=True
):
    """Run the model with the codebook features patched in."""
    hook_fns = [
        partial(patch_in_codes, pos=pos, code=code[i]) for i in range(len(code))
    ]
    cb_model.reset_codebook_metrics()
    cb_model.reset_hook_kwargs()
    fwd_hooks = [
        (get_cb_layer_name(cb_at[i], layer_idx[i], head_idx[i]), hook_fns[i])
        for i in range(len(cb_at))
    ]
    with cb_model.hooks(fwd_hooks, [], True, False) as hooked_model:
        patched_logits, patched_cache = hooked_model.run_with_cache(
            input, prepend_bos=prepend_bos
        )
    return patched_logits, patched_cache


def in_hook_list(list_of_arg_tuples, layer, head=None):
    """Check if the component specified by `layer` and `head` is in the `list_of_arg_tuples`."""
    # if head is not provided, then checks in MLP
    for arg_tuple in list_of_arg_tuples:
        if head is None:
            if arg_tuple.cb_at == "mlp" and arg_tuple.layer == layer:
                return True
        else:
            if (
                arg_tuple.cb_at == "attn"
                and arg_tuple.layer == layer
                and arg_tuple.head == head
            ):
                return True
    return False


# def generate_with_codes(input, code, cb_at, layer_idx, head_idx=None, pos=None, disable_other_comps=False):
def generate_with_codes(
    input,
    cb_model,
    list_of_code_infos=(),
    disable_other_comps=False,
    automata=None,
    generate_kwargs=None,
):
    """Model's generation with the codebook features patched in."""
    if generate_kwargs is None:
        generate_kwargs = {}
    hook_fns = [
        partial(patch_in_codes, pos=tupl.pos, code=tupl.code)
        for tupl in list_of_code_infos
    ]
    fwd_hooks = [
        (get_cb_layer_name(tupl.cb_at, tupl.layer, tupl.head), hook_fns[i])
        for i, tupl in enumerate(list_of_code_infos)
    ]
    cb_model.reset_hook_kwargs()
    if disable_other_comps:
        for layer, cb in cb_model.all_codebooks.items():
            for head_idx, head in enumerate(cb[0].codebook):
                if not in_hook_list(list_of_code_infos, layer, head_idx):
                    head.set_hook_kwargs(
                        disable_topk=1, disable_for_tkns=[-1], keep_k_codes=False
                    )
            if not in_hook_list(list_of_code_infos, layer):
                cb[1].set_hook_kwargs(
                    disable_topk=1, disable_for_tkns=[-1], keep_k_codes=False
                )
    with cb_model.hooks(fwd_hooks, [], True, False) as hooked_model:
        gen = hooked_model.generate(input, **generate_kwargs)
    return automata.seq_to_traj(gen)[0] if automata is not None else gen


def kl_div(logits1, logits2, pos=-1, reduction="batchmean"):
    """Calculate the KL divergence between the logits at `pos`."""
    logits1_last, logits2_last = logits1[:, pos, :], logits2[:, pos, :]
    # calculate kl divergence between clean and mod logits last
    return F.kl_div(
        F.log_softmax(logits1_last, dim=-1),
        F.log_softmax(logits2_last, dim=-1),
        log_target=True,
        reduction=reduction,
    )


def JSD(logits1, logits2, pos=-1, reduction="batchmean"):
    """Compute the Jensen-Shannon divergence between two distributions."""
    if len(logits1.shape) == 3:
        logits1, logits2 = logits1[:, pos, :], logits2[:, pos, :]

    probs1 = F.softmax(logits1, dim=-1)
    probs2 = F.softmax(logits2, dim=-1)

    total_m = (0.5 * (probs1 + probs2)).log()

    loss = 0.0
    loss += F.kl_div(
        total_m,
        F.log_softmax(logits1, dim=-1),
        log_target=True,
        reduction=reduction,
    )
    loss += F.kl_div(
        total_m,
        F.log_softmax(logits2, dim=-1),
        log_target=True,
        reduction=reduction,
    )
    return 0.5 * loss


def residual_stream_patching_hook(resid_pre, hook, cache, position: int):
    """Patch in the codebook features at `position` from `cache`."""
    clean_resid_pre = cache[hook.name]
    resid_pre[:, position, :] = clean_resid_pre[:, position, :]
    return resid_pre


def find_code_changes(cache1, cache2, pos=None):
    """Find the codebook codes that are different between the two caches."""
    for k in cache1.keys():
        if "codebook" in k:
            c1 = cache1[k][0, pos]
            c2 = cache2[k][0, pos]
            if not torch.all(c1 == c2):
                print(cb_layer_name_to_info(k), c1.tolist(), c2.tolist())
                print(cb_layer_name_to_info(k), c1.tolist(), c2.tolist())


def common_codes_in_cache(cache_codes, threshold=0.0):
    """Get the common code in the cache."""
    codes, counts = torch.unique(cache_codes, return_counts=True, sorted=True)
    counts = counts.float() * 100
    counts /= cache_codes.shape[1]
    counts, indices = torch.sort(counts, descending=True)
    codes = codes[indices]
    indices = counts > threshold
    codes, counts = codes[indices], counts[indices]
    return codes, counts


def parse_code_info_string(
    info_str: str, cb_at="attn", pos=None, code_pos=-1
) -> CodeInfo:
    """Parse the code info string.

    The format of the `info_str` is:
    `code: 0, layer: 0, head: 0, occ_freq: 0.0, train_act_freq: 0.0`.
    """
    code, layer, head, occ_freq, train_act_freq = info_str.split(", ")
    code = int(code.split(": ")[1])
    layer = int(layer.split(": ")[1])
    head = int(head.split(": ")[1]) if head else None
    occ_freq = float(occ_freq.split(": ")[1])
    train_act_freq = float(train_act_freq.split(": ")[1])
    return CodeInfo(code, layer, head, pos=pos, code_pos=code_pos, cb_at=cb_at)


def parse_concept_codes_string(info_str: str, pos=None, code_append=False):
    """Parse the concept codes string."""
    code_info_strs = info_str.strip().split("\n")
    concept_codes = []
    layer, head = None, None
    code_pos = "append" if code_append else -1
    for code_info_str in code_info_strs:
        concept_codes.append(
            parse_code_info_string(code_info_str, pos=pos, code_pos=code_pos)
        )
        if code_append:
            continue
        if layer == concept_codes[-1].layer and head == concept_codes[-1].head:
            code_pos -= 1
        else:
            code_pos = -1
        concept_codes[-1].code_pos = code_pos
        layer, head = concept_codes[-1].layer, concept_codes[-1].head
    return concept_codes