"""IFD (Instruction Following Difficulty) data quality scoring.

IFD = conditional_NLL / unconditional_NLL

- Messages format: plain text concatenation (no chat template)
- Plain format: raw instr_key + resp_key fields

v2 changelog:
  - Same token set: unconditional pass prefixes resp with a plain-text sentinel
    (default ``\\n``; use ``--sentinel_text ""`` for bos/pad fallback).
    Both branches predict the identical N resp tokens.
  - Single-token answers (rl=1) are now supported.
  - ctx_len tracked in output
  - skip_reason for None samples (no more silent None)
  - --per_token for per-token IFD breakdown

v3 changelog:
  - Append EOS at the end of response in both conditional and unconditional
    passes (``--append_eos`` / ``--no-append_eos``, default: enabled).
    The model now also predicts when the response should end, which is part
    of instruction following. Falls back gracefully when tokenizer has no EOS.
"""

import argparse
import glob
import json
import os
import statistics

import torch
import torch.nn.functional as F
import tqdm

from astrai.model import AutoModel
from astrai.preprocessing.packing import plan_bfd
from astrai.tokenize import AutoTokenizer


def _pack_bins(pairs, max_len):
    """BFD bin packing: pack (c+r) into bins of max total length.

    Reuses :func:`plan_bfd` so the BFD heuristic stays single-sourced.
    """
    # Treat each pair as a single sequence of length len(c)+len(r) for
    # planning purposes; plan_bfd works on pure lengths.
    fake_sequences = [[0] * (len(c) + len(r)) for c, r in pairs]
    plan = plan_bfd(fake_sequences, max_len)
    return [
        [(i, pairs[i][0], pairs[i][1]) for i in bin_indices] for bin_indices in plan
    ]


def _resolve_sentinel_ids(tokenizer, sentinel_text):
    """Tokenize the sentinel text for the unconditional pass prefix.

    Falls back to bos/pad_token_id when sentinel_text is empty or
    cannot be encoded.
    """
    if sentinel_text:
        ids = tokenizer.encode(sentinel_text, add_special_tokens=False)
        if ids:
            return ids
    for attr in ("bos_token_id", "pad_token_id", "eos_token_id"):
        tid = getattr(tokenizer, attr, None)
        if tid is not None:
            return [tid]
    return [0]


def _collect_input_files(input_path: str) -> list:
    """Resolve *input_path* to a list of JSONL/JSON files."""
    if os.path.isdir(input_path):
        files = []
        for ext in ("*.jsonl", "*.json"):
            files.extend(
                sorted(glob.glob(os.path.join(input_path, "**", ext), recursive=True))
            )
        return files
    return sorted(glob.glob(input_path))


def _load_items(filepath: str) -> list:
    """Load JSONL or JSON (array / single dict) into a list of dicts."""
    with open(filepath, "r", encoding="utf-8") as f:
        if filepath.lower().endswith(".json"):
            data = json.load(f)
            if isinstance(data, dict):
                return [data]
            return data
        return [json.loads(line) for line in f if line.strip()]


@torch.inference_mode()
def _score_batch(
    pairs, model, device, max_len=2048, sentinel_ids=None, per_token=False
):
    """BFD-packed IFD with text-sentinel-anchored unconditional pass.

    Conditional:   (ctx + resp[0..i-1]) → resp[i],  i = 0..N-1
    Unconditional: (<sentinel> + resp[0..i-1]) → resp[i],  i = 0..N-1

    Both branches predict the identical N response tokens.  A short
    plain-text sentinel gives the unconditional pass a prefix so that
    every response token can be predicted.  Single-token answers (rl=1)
    are supported.
    """
    if not pairs:
        return []

    if sentinel_ids is None:
        sentinel_ids = [0]

    bins = _pack_bins(pairs, max_len)
    result = [None] * len(pairs)

    # ---- conditional pass (packed, per-document position IDs) ----
    for bin_items in bins:
        seq_ids = []
        global_pos = []
        doc_ids = []
        doc_offsets = []

        for di, (orig_idx, c, r) in enumerate(bin_items):
            ctx_len = len(c)
            start = len(seq_ids)
            item_len = len(c) + len(r)
            seq_ids.extend(c)
            seq_ids.extend(r)
            end = len(seq_ids)
            global_pos.extend(range(item_len))
            doc_ids.extend([di] * item_len)
            doc_offsets.append((start, end, orig_idx, ctx_len))

        full_ids = torch.tensor([seq_ids], device=device, dtype=torch.long)
        pos_ids = torch.tensor([global_pos], device=device, dtype=torch.long)
        seq_len = len(seq_ids)
        causal = torch.tril(
            torch.ones(seq_len, seq_len, dtype=torch.bool, device=device)
        )
        doc_t = torch.tensor([doc_ids], device=device)
        doc_mask = doc_t.unsqueeze(-1) == doc_t.unsqueeze(-2)
        attn_mask = (causal & doc_mask[0]).unsqueeze(0).unsqueeze(0)
        logits_full = model(full_ids, position_ids=pos_ids, input_mask=attn_mask)[
            "logits"
        ][0]

        for start, end, orig_idx, ctx_len in doc_offsets:
            rl = end - start - ctx_len
            resp_start = start + ctx_len - 1
            resp_logits = logits_full[resp_start : end - 1]
            resp_targets = torch.tensor(
                seq_ids[start + ctx_len : end], device=device, dtype=torch.long
            )
            cond_losses = F.cross_entropy(
                resp_logits, resp_targets, reduction="none"
            ).cpu()
            result[orig_idx] = {
                "_cond_losses": cond_losses,
                "_rl": rl,
                "_ctx_len": ctx_len,
            }

    # ---- unconditional pass (sentinel-prefixed, batched 2D) ----
    valid_items = [
        (
            i,
            result[i]["_rl"],
            result[i]["_ctx_len"],
            result[i]["_cond_losses"],
            pairs[i][1],
        )
        for i in range(len(pairs))
        if result[i] is not None and "_cond_losses" in result[i]
    ]
    if not valid_items:
        return result

    valid_items.sort(key=lambda x: -x[1])
    prefix_len = len(sentinel_ids)
    max_rl = prefix_len + max(rl for _, rl, _, _, _ in valid_items)
    bsz = len(valid_items)

    u_batch = torch.zeros(bsz, max_rl, dtype=torch.long, device=device)
    for ri, (_, rl, _, _, r_ids) in enumerate(valid_items):
        u_batch[ri, :prefix_len] = torch.tensor(sentinel_ids, dtype=torch.long)
        u_batch[ri, prefix_len : prefix_len + rl] = torch.tensor(
            r_ids, dtype=torch.long
        )

    logits_resp = model(u_batch)["logits"]

    for ri, (orig_idx, rl, ctx_len, cond_losses, _) in enumerate(valid_items):
        unp_logits = logits_resp[ri, prefix_len - 1 : prefix_len - 1 + rl]
        unp_targets = u_batch[ri, prefix_len : prefix_len + rl]
        uncond_losses = F.cross_entropy(unp_logits, unp_targets, reduction="none").cpu()

        L_cond = cond_losses.mean().item()
        L_uncond = uncond_losses.mean().item()
        ifd = L_cond / L_uncond if L_uncond > 0 else None

        out = {
            "L_cond": round(L_cond, 6),
            "L_uncond": round(L_uncond, 6),
            "ifd": round(ifd, 6) if ifd is not None else None,
            "ctx_len": ctx_len,
            "resp_len": rl,
        }
        if per_token:
            per = [
                (round(c.item() / u.item(), 6) if u.item() > 0 else None)
                for c, u in zip(cond_losses, uncond_losses)
            ]
            out["ifd_per_token"] = per
        result[orig_idx] = out

    return result


def _trim(context_ids, resp_ids, max_len):
    """Truncate to fit max_len, keeping response intact if possible."""
    if len(resp_ids) > max_len // 2:
        resp_ids = resp_ids[: max_len // 2]
    full_ids = context_ids + resp_ids
    if len(full_ids) <= max_len:
        return context_ids, resp_ids
    overflow = len(full_ids) - max_len
    if overflow >= len(context_ids):
        return [], resp_ids[:max_len]
    return context_ids[overflow:], resp_ids


def process_file(
    model,
    tokenizer,
    input_file,
    output_file,
    instr_key,
    resp_key,
    max_len=2048,
    data_format="plain",
    batch_size=1,
    device=None,
    sentinel_ids=None,
    per_token=False,
    max_samples=None,
    eos_ids=None,
):
    """Score a single file, write per-sample JSONL, return summary stats."""
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    if sentinel_ids is None:
        sentinel_ids = _resolve_sentinel_ids(tokenizer, "\n")

    if eos_ids is None:
        eos_ids = []

    eos_len = len(eos_ids)

    data = _load_items(input_file)

    if max_samples and len(data) > max_samples:
        import random

        data = random.sample(data, max_samples)

    results = []
    all_ifds = []
    buffer = []

    label = os.path.splitext(os.path.basename(input_file))[0]

    for item in tqdm.tqdm(data, desc=f"  {label}", unit="sample", leave=False):
        if data_format == "messages":
            turns = []
            for i, msg in enumerate(item.get("messages", [])):
                if msg.get("role") != "assistant":
                    continue
                ctx_text = "\n\n".join(m["content"] for m in item["messages"][:i])
                ctx_ids = tokenizer.encode(ctx_text)
                resp_ids = tokenizer.encode(msg["content"], add_special_tokens=False)
                ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len - eos_len)
                if eos_ids and resp_ids and resp_ids[-1:] != eos_ids:
                    resp_ids = resp_ids + eos_ids
                if ctx_ids and resp_ids:
                    turns.append((ctx_ids, resp_ids))
            if not turns:
                results.append(
                    {
                        **item,
                        "ifd": None,
                        "skip_reason": "no valid assistant turns",
                        "ifd_turns": [],
                    }
                )
                continue
            buffer.append((item, turns, "messages"))
        else:
            ctx_ids = tokenizer.encode(item[instr_key], add_special_tokens=False)
            resp_ids = tokenizer.encode(item[resp_key], add_special_tokens=False)
            ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len - eos_len)
            if not ctx_ids or not resp_ids:
                results.append(
                    {
                        **item,
                        "ifd": None,
                        "ifd_detail": {"skip_reason": "empty ctx or resp"},
                    }
                )
                continue
            if eos_ids and resp_ids[-1:] != eos_ids:
                resp_ids = resp_ids + eos_ids
            buffer.append((item, [(ctx_ids, resp_ids)], "plain"))

        if len(buffer) >= batch_size:
            _flush_buffer(
                buffer,
                results,
                all_ifds,
                model,
                device,
                max_len,
                sentinel_ids,
                per_token,
            )

    if buffer:
        _flush_buffer(
            buffer, results, all_ifds, model, device, max_len, sentinel_ids, per_token
        )

    with open(output_file, "w", encoding="utf-8") as f:
        for item in results:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

    valid_ifd = [v for v in all_ifds if v is not None]
    stats = {
        "samples": len(data),
        "valid_ifd": len(valid_ifd),
        "skipped": len(data) - len(valid_ifd),
    }
    if valid_ifd:
        stats["mean_ifd"] = statistics.mean(valid_ifd)
        stats["median_ifd"] = statistics.median(valid_ifd)
        if len(valid_ifd) > 1:
            stats["stdev_ifd"] = statistics.stdev(valid_ifd)
        stats["min_ifd"] = min(valid_ifd)
        stats["max_ifd"] = max(valid_ifd)

        print(f"\n{'=' * 50}")
        print(f"  [{label}]")
        print(f"{'=' * 50}")
        print(f"  Samples:       {len(data)}")
        print(f"  Valid IFD:     {len(valid_ifd)}")
        print(f"  Skipped:       {len(data) - len(valid_ifd)}")
        print(f"  Mean IFD:      {statistics.mean(valid_ifd):.4f}")
        print(f"  Median IFD:    {statistics.median(valid_ifd):.4f}")
        if len(valid_ifd) > 1:
            print(f"  Stdev IFD:     {statistics.stdev(valid_ifd):.4f}")
        print(f"  Min IFD:       {min(valid_ifd):.4f}")
        print(f"  Max IFD:       {max(valid_ifd):.4f}")
        print(f"{'=' * 50}")
    print(f"  Results saved to {output_file}")
    return stats


def _flush_buffer(
    buffer, results, all_ifds, model, device, max_len, sentinel_ids, per_token
):
    all_pairs = []
    indices = []
    for item, turns, fmt in buffer:
        start = len(all_pairs)
        all_pairs.extend(turns)
        indices.append((item, turns, fmt, start, len(all_pairs)))

    raw = _score_batch(
        all_pairs,
        model,
        device,
        max_len,
        sentinel_ids=sentinel_ids,
        per_token=per_token,
    )

    for item, turns, fmt, start, end in indices:
        turn_scores = raw[start:end]
        if fmt == "messages":
            valid = [
                s for s in turn_scores if s is not None and s.get("ifd") is not None
            ]
            if not valid:
                results.append({**item, "ifd": None, "ifd_turns": turn_scores})
            else:
                avg = sum(s["ifd"] for s in valid) / len(valid)
                all_ifds.append(avg)
                results.append(
                    {
                        **item,
                        "ifd": avg,
                        "ifd_detail": valid[0] if len(valid) == 1 else None,
                        "ifd_turns": turn_scores,
                    }
                )
        else:
            score = turn_scores[0]
            all_ifds.append(score.get("ifd"))
            results.append({**item, "ifd": score.get("ifd"), "ifd_detail": score})

    buffer.clear()


def main():
    parser = argparse.ArgumentParser(
        description="Compute IFD scores for instruction-response data"
    )
    parser.add_argument("--param_path", type=str, required=True, help="Model directory")
    parser.add_argument(
        "--input_path",
        type=str,
        required=True,
        help="Input file, glob pattern, or directory.",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        required=True,
        help="Directory for output files (summary.json + per-file JSONL).",
    )
    parser.add_argument("--max_len", type=int, default=2048, help="Max token length")
    parser.add_argument(
        "--format",
        type=str,
        default="plain",
        choices=["plain", "messages"],
        help="Input format",
    )
    parser.add_argument(
        "--instr_key", type=str, default="instruction", help="Key for instruction field"
    )
    parser.add_argument(
        "--resp_key", type=str, default="response", help="Key for response field"
    )
    parser.add_argument(
        "--batch_size", type=int, default=8, help="Batch size for model forward passes"
    )
    parser.add_argument("--device", type=str, default=None, help="Device (e.g. cuda:0)")
    parser.add_argument(
        "--dtype",
        type=str,
        default="bfloat16" if torch.cuda.is_available() else "float32",
        help="Torch dtype",
    )
    parser.add_argument(
        "--sentinel_text",
        type=str,
        default="\n",
        help='Plain-text prefix for unconditional pass (default: "\\n"). Use "" for bos/pad fallback.',
    )
    parser.add_argument(
        "--per_token",
        action="store_true",
        help="Include per-token IFD breakdown in output",
    )
    parser.add_argument(
        "--max_samples",
        type=int,
        default=None,
        help="Maximum number of samples per file (random subsample). Default: all.",
    )
    parser.add_argument(
        "--append_eos/--no-append_eos",
        default=True,
        help="Append EOS token at the end of response in both passes (default: enabled).",
    )
    args = parser.parse_args()

    if args.device is None:
        args.device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = getattr(torch, args.dtype)

    print(f"Loading model from {args.param_path} ...")
    model = AutoModel.from_pretrained(args.param_path)
    tokenizer = AutoTokenizer.from_pretrained(args.param_path)
    model.to(device=args.device, dtype=dtype)
    model.eval()

    sentinel_ids = _resolve_sentinel_ids(tokenizer, args.sentinel_text)

    eos_ids = []
    if args.append_eos:
        eos_token_id = getattr(tokenizer, "eos_token_id", None)
        if eos_token_id is not None:
            eos_ids = [eos_token_id]
        else:
            print(
                "Warning: --append_eos enabled but tokenizer has no EOS token; skipping."
            )

    input_files = _collect_input_files(args.input_path)
    if not input_files:
        print(f"No input files found at {args.input_path}")
        return

    print(f"Found {len(input_files)} file(s) to evaluate")
    os.makedirs(args.output_dir, exist_ok=True)

    all_stats = {}
    for filepath in input_files:
        label = os.path.splitext(os.path.basename(filepath))[0]
        output_file = os.path.join(args.output_dir, f"{label}_ifd.jsonl")

        stats = process_file(
            model=model,
            tokenizer=tokenizer,
            input_file=filepath,
            output_file=output_file,
            instr_key=args.instr_key,
            resp_key=args.resp_key,
            max_len=args.max_len,
            data_format=args.format,
            batch_size=args.batch_size,
            device=args.device,
            sentinel_ids=sentinel_ids,
            per_token=args.per_token,
            max_samples=args.max_samples,
            eos_ids=eos_ids,
        )
        all_stats[label] = stats

    summary_path = os.path.join(args.output_dir, "summary.json")
    with open(summary_path, "w", encoding="utf-8") as f:
        json.dump(all_stats, f, ensure_ascii=False, indent=2)
    print(f"\nSummary saved to {summary_path}")


if __name__ == "__main__":
    main()