"""
Projection Map — Condensed view of layer, head, and head-combination projections.

One table: 26 rows (one per layer).
Columns: H0-H3 (individual heads), all pairs, all triples, full layer projection.
Format: #LL|tok      PP%|tok      PP%|...|laytok   PP%|
"""

import argparse
import unicodedata
from itertools import combinations
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

DEFAULT_MODEL_ID = "google/gemma-3-1b-it"

C_H   = '\033[95m'
C_B   = '\033[94m'
C_G   = '\033[92m'
C_Y   = '\033[93m'
C_C   = '\033[96m'
C_E   = '\033[0m'
C_BLD = '\033[1m'
C_DIM = '\033[2m'


def safe_print(text):
    try:
        print(text)
    except UnicodeEncodeError:
        print(text.encode('ascii', 'replace').decode('ascii'))


def char_width(ch):
    """Terminal column width of a single character."""
    cat = unicodedata.category(ch)
    if cat in ('Mn', 'Me'):
        return 0
    if ord(ch) in (0x200B, 0x200C, 0x200D, 0xFEFF):
        return 0
    eaw = unicodedata.east_asian_width(ch)
    if eaw in ('W', 'F'):
        return 2
    return 1


def display_width(s):
    return sum(char_width(ch) for ch in s)


def fmt_tok(s, width=9):
    """Format token to exactly `width` terminal columns."""
    s = s.replace('\n', '\\n').replace('\r', '\\r').replace('\t', '\\t')
    s = s.lstrip('\u2581')
    s = s.lstrip(' ')
    if not s:
        s = '·'
    if display_width(s) <= width:
        pad = width - display_width(s)
        return s + ' ' * pad
    result = []
    used = 0
    for ch in s:
        ch_w = char_width(ch)
        if ch_w > 0 and used + ch_w > width - 1:
            break
        result.append(ch)
        used += ch_w
    result.append('.')
    used += 1
    while used < width:
        result.append(' ')
        used += 1
    return ''.join(result)


def fmt_pct(p):
    pct = int(p * 100)
    if pct > 99:
        return "99"
    return f"{pct:02d}"


def prob_color(p):
    if p > 0.5: return C_G
    if p > 0.1: return C_Y
    return C_DIM


def combo_label(indices):
    """Short label for a head combination."""
    return ''.join(str(i) for i in indices)


def load_model(model_id):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"{C_H}Loading {model_id}...{C_E}")
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    tokenizer.padding_side = "left"
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # Use bfloat16 — Gemma 3's RMSNorm breaks with float16 (NaN logits).
    # bfloat16 works on RTX 30+ / A100 / H100. Fall back to float32 if unavailable.
    try:
        use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    except Exception:
        use_bf16 = False
    dtype = torch.bfloat16 if use_bf16 else torch.float32
    model = AutoModelForCausalLM.from_pretrained(
        model_id, device_map="auto", torch_dtype=dtype
    )
    model.eval()

    num_layers = model.config.num_hidden_layers
    num_heads = model.config.num_attention_heads
    print(f"{C_C}Architecture: {num_layers} layers, {num_heads} heads{C_E}")

    return model, tokenizer, device, num_layers, num_heads


def project_to_token(projected_hidden, final_norm, lm_head, tokenizer):
    """Take a projected hidden state, norm -> lm_head -> argmax, return (token_str, prob)."""
    normed = final_norm(projected_hidden)
    logits = lm_head(normed)
    probs = F.softmax(logits, dim=-1)
    p, tid = torch.max(probs, dim=-1)
    tok = tokenizer.decode(tid.item())
    return tok, p.item()


def project_to_token_rich(projected_hidden, final_norm, lm_head, tokenizer, top_k=5):
    """Rich projection: top-k tokens, entropy, full distribution shape.
    Returns dict with:
      'top_k': list of (token_str, prob) for top k
      'entropy': Shannon entropy of the distribution (bits)
      'argmax': (token_str, prob) — same as project_to_token
      'mass_top5': total probability mass in top 5
      'mass_top50': total probability mass in top 50
    """
    normed = final_norm(projected_hidden)
    logits = lm_head(normed)
    probs = F.softmax(logits, dim=-1)

    # Entropy in bits
    log_probs = torch.log2(probs + 1e-10)
    entropy = -(probs * log_probs).sum().item()

    # Top-k
    topk_probs, topk_ids = torch.topk(probs.squeeze(), top_k)
    top_k_list = []
    for i in range(top_k):
        tok = tokenizer.decode(topk_ids[i].item())
        top_k_list.append((tok, topk_probs[i].item()))

    # Mass in top-50
    top50_probs, _ = torch.topk(probs.squeeze(), min(50, probs.shape[-1]))
    mass_top50 = top50_probs.sum().item()

    return {
        'top_k': top_k_list,
        'entropy': entropy,
        'argmax': top_k_list[0],
        'mass_top5': sum(p for _, p in top_k_list),
        'mass_top50': mass_top50,
    }


def project_to_softmax(projected_hidden, final_norm, lm_head):
    """Project hidden state to full softmax distribution.
    Returns the raw probability tensor (vocab_size,) on CPU as float16.
    This is the complete recording — every possible analysis is a post-hoc query.
    """
    normed = final_norm(projected_hidden)
    logits = lm_head(normed)
    probs = F.softmax(logits, dim=-1)
    return probs.squeeze().detach().cpu().half()


def project_full(projected_hidden, final_norm, lm_head):
    """Project hidden state and return BOTH the projected vector and softmax.
    Returns: (projected_vector, softmax_probs)
      - projected_vector: (hidden_dim,) float32 on CPU — the normed hidden state,
        directly comparable to lm_head weight vectors (i.e. token embeddings).
        This is what you search against the Lance token index.
      - softmax_probs: (vocab_size,) float16 on CPU — the full distribution.
    """
    normed = final_norm(projected_hidden)
    logits = lm_head(normed)
    probs = F.softmax(logits, dim=-1)
    return (
        normed.squeeze().detach().cpu().float(),
        probs.squeeze().detach().cpu().half(),
    )


def make_cell(tok, prob):
    """Format a single table cell."""
    pc = prob_color(prob)
    return f"{pc}{fmt_tok(tok)}{fmt_pct(prob)}%{C_E}"


def build_col_specs(num_heads):
    """Build column specifications for the projection table."""
    head_indices = list(range(num_heads))
    pairs = list(combinations(head_indices, 2))
    triples = list(combinations(head_indices, 3))

    col_specs = []
    for i in head_indices:
        col_specs.append(('single', (i,), f"H{i}"))
    for combo in pairs:
        col_specs.append(('combo', combo, f"H{''.join(str(i) for i in combo)}"))
    for combo in triples:
        col_specs.append(('combo', combo, f"H{''.join(str(i) for i in combo)}"))
    col_specs.append(('layer', None, 'Layer'))

    return col_specs, head_indices, pairs, triples


def capture_projection_map(model, tokenizer, input_ids, device, num_layers, num_heads,
                           scale_mode='full', rich=False, top_k=5, full_softmax=False):
    """
    Run forward pass and return structured projection data.
    If rich=False: list of dicts, one per layer, each containing {col_label: (token, prob)}
    If rich=True:  list of dicts, one per layer, each containing {col_label: rich_dict}
                   where rich_dict has top_k, entropy, argmax, mass_top5, mass_top50
    If full_softmax=True: also returns softmax_data — list of dicts per layer,
                   each {col_label: tensor(vocab_size)} as float16 on CPU.
                   This is the complete MRI — every query is post-hoc.
    Returns: (layer_data, col_specs) or (layer_data, col_specs, softmax_data) if full_softmax
    """
    final_norm = model.model.norm
    lm_head = model.lm_head

    captured = {}

    def make_hook(layer_idx):
        def hook_fn(module, args, output):
            captured[layer_idx] = args[0].detach()
        return hook_fn

    handles = []
    for l in range(num_layers):
        layer = model.model.layers[l]
        h = layer.self_attn.o_proj.register_forward_hook(make_hook(l))
        handles.append(h)

    try:
        with torch.no_grad():
            outputs = model(input_ids, output_hidden_states=True)
    finally:
        for h in handles:
            h.remove()

    col_specs, head_indices, _, _ = build_col_specs(num_heads)

    layer_data = []
    softmax_data = [] if full_softmax else None
    for l_idx in range(num_layers):
        layer_state = outputs.hidden_states[l_idx + 1]
        last_token = layer_state[:, -1, :]
        if rich:
            lay_rich = project_to_token_rich(last_token, final_norm, lm_head, tokenizer, top_k)
        else:
            lay_tok, lay_prob = project_to_token(last_token, final_norm, lm_head, tokenizer)

        row = {}

        if l_idx in captured:
            input_tensor = captured[l_idx]
            layer = model.model.layers[l_idx]
            o_proj_weight = layer.self_attn.o_proj.weight
            hidden_size = o_proj_weight.shape[0]
            attn_out_dim = o_proj_weight.shape[1]
            head_dim = attn_out_dim // num_heads

            batch, seq, _ = input_tensor.shape
            multi_head = input_tensor.view(batch, seq, num_heads, head_dim)
            weight_view = o_proj_weight.view(hidden_size, num_heads, head_dim)

            head_projected = {}
            for h_idx in head_indices:
                head_output = multi_head[:, :, h_idx, :]
                head_weights = weight_view[:, h_idx, :]
                projected = torch.matmul(head_output, head_weights.t())
                head_projected[h_idx] = projected[:, -1, :]

            for col_type, col_heads, col_label in col_specs:
                if col_type == 'layer':
                    if rich:
                        row[col_label] = lay_rich
                    else:
                        row[col_label] = (lay_tok, lay_prob)
                else:
                    combined = sum(head_projected[h] for h in col_heads)
                    if scale_mode == 'full':
                        combined = combined * (num_heads / len(col_heads))
                    elif scale_mode == 'mean':
                        combined = combined / len(col_heads)
                    if rich:
                        row[col_label] = project_to_token_rich(combined, final_norm, lm_head, tokenizer, top_k)
                    else:
                        tok, prob = project_to_token(combined, final_norm, lm_head, tokenizer)
                        row[col_label] = (tok, prob)

            # If full softmax capture requested, store raw distributions
            if full_softmax:
                softmax_row = {}
                for col_type, col_heads, col_label in col_specs:
                    if col_type == 'layer':
                        softmax_row[col_label] = project_to_softmax(last_token, final_norm, lm_head)
                    else:
                        combined = sum(head_projected[h] for h in col_heads)
                        if scale_mode == 'full':
                            combined = combined * (num_heads / len(col_heads))
                        elif scale_mode == 'mean':
                            combined = combined / len(col_heads)
                        softmax_row[col_label] = project_to_softmax(combined, final_norm, lm_head)
                softmax_data.append(softmax_row)
        else:
            for _, _, col_label in col_specs:
                if rich:
                    row[col_label] = {'top_k': [('--', 0.0)]*top_k, 'entropy': 0.0,
                                       'argmax': ('--', 0.0), 'mass_top5': 0.0, 'mass_top50': 0.0}
                else:
                    row[col_label] = ('--', 0.0)
            if full_softmax:
                softmax_data.append({})

        layer_data.append(row)

    if full_softmax:
        return layer_data, col_specs, softmax_data
    return layer_data, col_specs


def print_projection_map(layer_data, col_specs):
    """Print the projection map table from structured data."""
    header_cells = [f"{C_BLD}{C_C}#  {C_E}"]
    for _, _, label in col_specs:
        header_cells.append(f"{C_BLD}{C_C}{label:^12}{C_E}")
    header = '|'.join(header_cells) + '|'

    sep_width = 4 + len(col_specs) * 13
    sep = f"{C_DIM}{'─' * sep_width}{C_E}"

    print()
    safe_print(header)
    safe_print(sep)

    for l_idx, row in enumerate(layer_data):
        cells = []
        for _, _, col_label in col_specs:
            tok, prob = row[col_label]
            cells.append(make_cell(tok, prob))
        row_num = f"{C_BLD}#{l_idx:02d}{C_E}"
        safe_print(f"{row_num}|{'|'.join(cells)}|")

    safe_print(sep)


def format_map_plain(layer_data, col_specs):
    """Format projection map as plain text (no ANSI) for LLM consumption."""
    labels = [label for _, _, label in col_specs]
    header = f"#  |{'|'.join(f'{l:^12}' for l in labels)}|"
    sep = '─' * len(header)

    lines = [header, sep]
    for l_idx, row in enumerate(layer_data):
        cells = []
        for _, _, col_label in col_specs:
            tok, prob = row[col_label]
            tok_clean = tok.replace('\n', '\\n').replace('\r', '\\r')
            if len(tok_clean) > 9:
                tok_clean = tok_clean[:8] + '.'
            cell = f"{tok_clean:<9}{int(prob*100):02d}%"
            cells.append(cell)
        lines.append(f"#{l_idx:02d}|{'|'.join(cells)}|")
    lines.append(sep)
    return '\n'.join(lines)


def run(model, tokenizer, input_ids, device, num_layers, num_heads, scale_mode='full'):
    """Capture and print the projection map. Returns structured data."""
    layer_data, col_specs = capture_projection_map(
        model, tokenizer, input_ids, device, num_layers, num_heads, scale_mode
    )
    print_projection_map(layer_data, col_specs)

    # Legend
    _, head_indices, pairs, triples = build_col_specs(num_heads)
    print()
    safe_print(f"{C_DIM}Singles: {', '.join(f'H{i}' for i in head_indices)}{C_E}")
    safe_print(f"{C_DIM}Pairs:   {', '.join('H' + ''.join(str(i) for i in c) for c in pairs)}{C_E}")
    safe_print(f"{C_DIM}Triples: {', '.join('H' + ''.join(str(i) for i in c) for c in triples)}{C_E}")

    return layer_data, col_specs


def main():
    parser = argparse.ArgumentParser(description="Projection map — layer × head × combos")
    parser.add_argument("--model", "-m", default=DEFAULT_MODEL_ID,
                        help=f"Model ID (default: {DEFAULT_MODEL_ID})")
    parser.add_argument("--prompt", "-p", default="",
                        help="Prompt text (default: empty, just BOS)")
    parser.add_argument("--scale", "-s", default="full",
                        choices=["raw", "full", "mean"],
                        help="Head scaling: raw (no scaling), full (scale to full-head magnitude), mean (average). Default: full")
    args = parser.parse_args()

    model, tokenizer, device, num_layers, num_heads = load_model(args.model)

    if args.prompt:
        input_ids = tokenizer(args.prompt, return_tensors="pt").input_ids.to(device)
        safe_print(f"\n{C_Y}Prompt: \"{args.prompt}\"{C_E}")
        safe_print(f"{C_DIM}Tokens: {input_ids.shape[1]}{C_E}")
    else:
        bos_id = tokenizer.bos_token_id
        if bos_id is None:
            input_ids = tokenizer("", return_tensors="pt").input_ids.to(device)
        else:
            input_ids = torch.tensor([[bos_id]], device=device)
        safe_print(f"\n{C_Y}No prompt — raw BOS token only{C_E}")

    safe_print(f"{C_DIM}Scale mode: {args.scale}{C_E}")

    run(model, tokenizer, input_ids, device, num_layers, num_heads, scale_mode=args.scale)

    print(f"\n{C_G}Done.{C_E}")


if __name__ == "__main__":
    main()