# experiment/probing/forward_pass.py
"""Two-pass forward routine: prompt-only + teacher-forced.

Returns a DataFrame of per-token rows matching PARQUET_COLUMNS.
"""
from __future__ import annotations

import os
import sys

import pandas as pd
import torch
from PIL import Image

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../.."))

from experiment.training.finetune_adv import HiddenStateCapture, count_lm_layers
from experiment.probing._helpers import (
    PARQUET_COLUMNS,
    encode_per_token_sae,
)


def two_pass_forward(
    *,
    method: str,
    image_id: str,
    image_path: str = "",
    image: Image.Image = None,
    prompt: str,
    caption: str,
    toilet_token_positions: list[int],
    feature_ids_per_layer: dict[int, list[int]],
    model,
    processor,
    sae,
    device: torch.device,
) -> pd.DataFrame:
    """Run Pass 1 (prompt-only) and Pass 2 (teacher-forced) for one image.

    Args:
      method: tag stored in every output row.
      image_id: tag stored in every output row.
      image_path: filesystem path to the image (PIL-readable). Used when
        ``image`` is not supplied.
      image: PIL.Image to use directly (takes priority over ``image_path``).
      prompt: prompt text (will be wrapped into the processor's chat template).
      caption: teacher caption (concatenated after prompt for Pass 2).
      toilet_token_positions: positions WITHIN the caption (0-indexed against
        the tokenized caption alone) where the regex matched.
      feature_ids_per_layer: layer_idx → list of feature ids to record.
      model, processor, sae, device: model components.

    Returns:
      pandas.DataFrame with PARQUET_COLUMNS.

    Prompt format note: FinetuneDataset.__getitem__ (datasets.py line 379) passes
    ``text=f"<image>\\n{prompt}"`` to the processor. We mirror that convention here
    so the tokenization matches the rest of the codebase.
    """
    tok = processor.tokenizer
    n_layers = count_lm_layers(model)
    layer_ids = list(range(n_layers))
    capture = HiddenStateCapture(model, layer_ids)

    if image is None:
        image = Image.open(image_path).convert("RGB")
    else:
        image = image.convert("RGB")

    # Pass 1 — prompt-only (mirrors FinetuneDataset: text=f"<image>\n{prompt}")
    p1 = processor(
        text=f"<image>\n{prompt}", images=image,
        return_tensors="pt", padding=False,
    ).to(device)
    with capture, torch.no_grad():
        model(**p1, use_cache=False)
    z1 = encode_per_token_sae(capture.hidden_states, sae)
    decoded1 = [tok.decode([int(t)]) for t in p1["input_ids"][0].tolist()]
    rows = _emit_rows(method, image_id, "prompt", decoded1,
                      toilet_positions=set(), z=z1,
                      feature_ids_per_layer=feature_ids_per_layer)

    # Pass 2 — teacher-forced (prompt + caption)
    p2 = processor(
        text=f"<image>\n{prompt} {caption}", images=image,
        return_tensors="pt", padding=False,
    ).to(device)
    with capture, torch.no_grad():
        model(**p2, use_cache=False)
    z2 = encode_per_token_sae(capture.hidden_states, sae)
    decoded2 = [tok.decode([int(t)]) for t in p2["input_ids"][0].tolist()]
    # Re-locate toilet positions in the full packed sequence by token-string match.
    # (Caller supplied `toilet_token_positions` relative to the caption-only
    # tokenization, which may not be byte-identical after concatenation; safest
    # is to re-derive in the packed decoded list.)
    from experiment.probing._helpers import find_toilet_token_positions_subword
    toilet_in_packed = set(find_toilet_token_positions_subword(tok, p2["input_ids"][0].tolist()))
    rows.extend(_emit_rows(method, image_id, "teacher", decoded2,
                           toilet_positions=toilet_in_packed, z=z2,
                           feature_ids_per_layer=feature_ids_per_layer))

    return pd.DataFrame(rows, columns=PARQUET_COLUMNS)


def _emit_rows(method, image_id, pass_, decoded_tokens, toilet_positions,
               z, feature_ids_per_layer):
    rows = []
    for layer, z_layer in z.items():
        feat_ids = feature_ids_per_layer.get(layer, [])
        if not feat_ids:
            continue
        # z_layer: (1, T, d_sae); take batch idx 0
        z0 = z_layer[0]
        for t in range(z0.shape[0]):
            tok_str = decoded_tokens[t] if t < len(decoded_tokens) else ""
            is_tt = t in toilet_positions
            for f in feat_ids:
                rows.append((method, image_id, pass_, int(layer), int(t),
                             tok_str, bool(is_tt), int(f),
                             float(z0[t, f].item())))
    return rows