# utils/dataset_helpers.py
from __future__ import annotations
import logging, os
from pathlib import Path
from collections import Counter, defaultdict
from typing import Collection, Optional
from datetime import datetime, timezone
import json

import numpy as np
from datasets import load_dataset

logger = logging.getLogger(__name__)

# tokens we want to watch closely
_WATCH = [" nodded", " leaned"]


def _chosen_target_quotas(
    chosen_counts: Counter[str],
    strength: float,
) -> dict[str, int]:
    """Return per-token occurrence caps for chosen-token regularisation."""
    if not chosen_counts or strength <= 0:
        return dict(chosen_counts)

    # Prevent the largest outliers from dominating the distribution before
    # applying the smoother median-threshold regularisation.
    if len(chosen_counts) >= 10:
        cap_value = sorted(chosen_counts.values(), reverse=True)[9]
        capped = {
            token: min(count, cap_value)
            for token, count in chosen_counts.items()
        }
    else:
        capped = dict(chosen_counts)

    median = float(np.median(list(capped.values())))
    return {
        token: int(round(
            count
            if count <= median
            else count * (median / count) ** strength
        ))
        for token, count in capped.items()
    }


def _trim_chosen_to_quotas(
    rows: list[dict],
    quotas: dict[str, int],
    rng: np.random.Generator,
) -> list[dict]:
    """Uniformly retain chosen-token occurrences up to their global quotas."""
    counts = Counter(
        token
        for row in rows
        for token in (row.get("multi_chosen_decoded") or [])
    )
    if all(quotas.get(token, count) >= count for token, count in counts.items()):
        return [dict(row) for row in rows]

    # Randomise both row and slot traversal so quota allocation does not
    # systematically favour early examples or a token's probability rank.
    seen: Counter[str] = Counter()
    trimmed: list[dict | None] = [None] * len(rows)
    for row_idx_raw in rng.permutation(len(rows)):
        row_idx = int(row_idx_raw)
        row = rows[row_idx]
        decoded = row.get("multi_chosen_decoded") or []
        keep: set[int] = set()
        for slot_idx_raw in rng.permutation(len(decoded)):
            slot_idx = int(slot_idx_raw)
            token = decoded[slot_idx]
            if seen[token] < max(0, quotas.get(token, counts[token])):
                keep.add(slot_idx)
                seen[token] += 1

        new_row = dict(row)
        new_row["multi_chosen_decoded"] = [
            token for slot_idx, token in enumerate(decoded) if slot_idx in keep
        ]

        raw = row.get("multi_chosen_raw")
        if isinstance(raw, list) and len(raw) == len(decoded):
            new_row["multi_chosen_raw"] = [
                token for slot_idx, token in enumerate(raw) if slot_idx in keep
            ]

        trimmed[row_idx] = new_row

    return [row for row in trimmed if row is not None]


def load_ftpo_multi_dataset(
    path: Path,
    tokenizer,
    *,
    experiment_run_dir: Path | None = None,
    max_seq_len: int = 4096,
    rejected_reg_strength: float = 0.0,
    chosen_reg_strength: float = 0.0,
    min_chosen_tokens: int = 1,
    max_train_examples: int | None = None,
    stop_words: Optional[Collection[str]] = None,
    num_proc: int | None = None,
    batch_size: int = 512,
):
    """
    Parallel loader for “multi-chosen” FTPO JSONL with dual regularisation.
    Logs the counts of `_WATCH` tokens at every major stage.
    """

    if min_chosen_tokens < 1:
        min_chosen_tokens = 1

    # ------------------------------------------------------------------
    # helpers
    # ------------------------------------------------------------------
    rng = np.random.default_rng(3407)

    def _median_threshold(cts: Counter[str], strength: float) -> dict[str, float]:
        if not cts or strength <= 0:
            return {}
        med = float(np.median(list(cts.values())))
        return {t: 1.0 if c <= med else (med / c) ** strength for t, c in cts.items()}

    def _log_top(cts: Counter[str], what: str) -> None:
        head = ", ".join(f"{tok!r}:{cnt}" for tok, cnt in cts.most_common(20))
        logger.info(f"[ftpo-loader] {what} top-20 → {head}")
        logger.info(
            "          ↳ watch «%s»: %s   «%s»: %s",
            _WATCH[0], cts[_WATCH[0]],
            _WATCH[1], cts[_WATCH[1]],
        )

    # ------------------------------------------------------------------
    # stop-word list (unchanged)
    # ------------------------------------------------------------------
    if stop_words is None:
        stop_words = {
            "the","a","an","in","on","at","by","for","to","of","and","or","but",
            "if","then","else","when","where","how","why","what","who","whom",
            "this","that","these","those","is","are","was","were","be","being",
            "been","have","has","had","do","does","did","will","would","shall",
            "should","can","could","may","might","must"
        }
    stop_words = {w.lower() for w in stop_words}

    # ------------------------------------------------------------------
    # 0️⃣  raw load + shuffle
    # ------------------------------------------------------------------
    raw = load_dataset("json", data_files=str(path), split="train").shuffle(seed=3407)
    rows = list(raw)
    if not rows:
        raise ValueError(f"{path} contained no rows")

    rej_counts = Counter(r["rejected_decoded"] for r in rows)
    _log_top(rej_counts, "BEFORE")

    # ────────────────────────────────────────────────────────────────
    # 1️⃣  Capture ORIGINAL rejected-token distribution & ratios
    #     (no rows removed, no chosen trimming yet)
    # ────────────────────────────────────────────────────────────────
    rej_cts_orig = Counter(r["rejected_decoded"] for r in rows)
    _log_top(rej_cts_orig, "PRE-NORMALISATION")

    # convert to fractional “weights” via median-threshold regularisation
    med = float(np.median(list(rej_cts_orig.values())))
    w_rej = {tok: 1.0 if c <= med else (med / c) ** rejected_reg_strength
            for tok, c in rej_cts_orig.items()}
    
    # normalised ratios we *want* to keep in the final dataset
    total_weighted = sum(w_rej[t] * c for t, c in rej_cts_orig.items())
    ratio_rej = {tok: (w_rej[tok] * cnt) / total_weighted
                for tok, cnt in rej_cts_orig.items()}

    # ────────────────────────────────────────────────────────────────
    # 2️⃣  Chosen-token trimming  (build quotas *before* we cut)
    # ────────────────────────────────────────────────────────────────
    chosen_cts_orig = Counter(tok
                        for r in rows
                        for tok in (r["multi_chosen_decoded"] or []))

    _log_top(chosen_cts_orig, "ORIGINAL CHOSEN TOKENS")

    tgt_chosen = _chosen_target_quotas(
        chosen_cts_orig,
        chosen_reg_strength,
    )

    # Log the target quotas
    quota_items = sorted(tgt_chosen.items(), key=lambda x: x[1], reverse=True)[:20]
    quota_str = ", ".join(f"{tok!r}:{quota}" for tok, quota in quota_items)
    logger.info(f"[ftpo-loader] CHOSEN TARGET QUOTAS top-20 → {quota_str}")
    logger.info(
        "          ↳ watch quotas «%s»: %s (was %s)   «%s»: %s (was %s)",
        _WATCH[0], tgt_chosen.get(_WATCH[0], 0), chosen_cts_orig.get(_WATCH[0], 0),
        _WATCH[1], tgt_chosen.get(_WATCH[1], 0), chosen_cts_orig.get(_WATCH[1], 0),
    )

    rows = _trim_chosen_to_quotas(rows, tgt_chosen, rng)
    chosen_cts_trimmed = Counter(
        token
        for row in rows
        for token in (row["multi_chosen_decoded"] or [])
    )
    _log_top(chosen_cts_trimmed, "POST-CHOSEN")

    # ────────────────────────────────────────────────────────────────
    # 3️⃣  Apply min_chosen_tokens row filter
    # ────────────────────────────────────────────────────────────────
    rows = [r for r in rows if len(r["multi_chosen_decoded"]) >= min_chosen_tokens]

    _log_top(Counter(r["rejected_decoded"] for r in rows), "POST-MIN-FILTER")

    # ────────────────────────────────────────────────────────────────
    # 4️⃣  Row-level quota sampling **now** that trimming & filtering
    #     are done.  Scale the original ratios to the remaining size.
    # ────────────────────────────────────────────────────────────────
    N_final = max_train_examples or len(rows)
    target_rows = {tok: int(round(ratio_rej[tok] * N_final))
                for tok in ratio_rej}

    rng.shuffle(rows)
    selected, seen = [], defaultdict(int)
    selected_indices = set()  # Track indices instead of row objects
    
    # First pass: try to fill quotas
    for i, r in enumerate(rows):
        tok = r["rejected_decoded"]
        if seen[tok] < target_rows.get(tok, 0):
            selected.append(r)
            selected_indices.add(i)
            seen[tok] += 1
        if len(selected) >= N_final:
            break
    
    # Second pass: if still short, keep adding while maintaining proportions
    if len(selected) < N_final:
        # Build index of remaining rows by token
        remaining_by_token = defaultdict(list)
        for i, r in enumerate(rows):
            if i not in selected_indices:
                remaining_by_token[r["rejected_decoded"]].append((i, r))
        
        # Keep adding until we reach N_final
        while len(selected) < N_final:
            # Find token that's furthest below its target ratio AND has rows available
            best_tok = None
            best_ratio_diff = -1
            
            for tok, available_rows in remaining_by_token.items():
                if not available_rows:  # Skip tokens with no remaining rows
                    continue
                    
                current_ratio = seen[tok] / len(selected) if len(selected) > 0 else 0
                target_ratio = ratio_rej.get(tok, 0)
                ratio_diff = target_ratio - current_ratio
                
                if ratio_diff > best_ratio_diff:
                    best_ratio_diff = ratio_diff
                    best_tok = tok
            
            # If no tokens have available rows, we're done
            if best_tok is None:
                break
            
            # Add one row for the most underrepresented token
            idx, r = remaining_by_token[best_tok].pop()
            selected.append(r)
            selected_indices.add(idx)
            seen[best_tok] += 1
    
    rows = selected

    final_chosen_counts = Counter(
        token
        for row in rows
        for token in (row["multi_chosen_decoded"] or [])
    )
    _log_top(final_chosen_counts, "FINAL CHOSEN TOKENS")

    # ── Dump the final row subset exactly as it was read (no tokenisation) ──
    if experiment_run_dir is not None:
        ts = datetime.now(timezone.utc).astimezone()\
                 .strftime("%Y-%m-%d_%H-%M-%S")
        dump_file = experiment_run_dir / f"ftpo_training_set_used_{ts}.jsonl"
        try:
            with open(dump_file, "w", encoding="utf-8") as fh:
                for r in rows:
                    fh.write(json.dumps(r, ensure_ascii=False) + "\n")
            logger.info("[ftpo-loader] dumped %d rows → %s", len(rows), dump_file)
        except Exception as e:
            logger.warning("[ftpo-loader] failed to dump training rows: %s", e)



    _log_top(Counter(r["rejected_decoded"] for r in rows), "AFTER-SAMPLING")
    logger.info("[ftpo-loader] kept %d rows after quota sampling", len(rows))


    # ------------------------------------------------------------------
    # 5️⃣  tokenisation (unchanged section)
    # ------------------------------------------------------------------
    from datasets import Dataset
    ds = Dataset.from_list(rows)

    tokenizer.truncation_side = "left"
    num_proc = num_proc or max(1, int(os.cpu_count() / 4))

    def _tok(batch):
        out_prompt, out_chosen, out_rej, out_valid = [], [], [], []

        prompt_tok = tokenizer(
            batch["context_with_chat_template"],
            add_special_tokens=False,
            truncation=False,
            return_attention_mask=False,
        ).input_ids

        for p_ids, chosen_surf, rej_surf in zip(
            prompt_tok, batch["multi_chosen_decoded"], batch["rejected_decoded"]
        ):
            chosen_surf = chosen_surf or []
            chosen_tok_ids = [tokenizer(t, add_special_tokens=False).input_ids
                              for t in chosen_surf]
            rej_tok_ids = tokenizer(rej_surf, add_special_tokens=False).input_ids

            valid = (
                chosen_tok_ids
                and all(len(t) == 1 for t in chosen_tok_ids)
                and len(rej_tok_ids) == 1
                and rej_surf.strip().lower() not in stop_words
                and len(p_ids) + 1 <= max_seq_len
            )
            if valid and rej_tok_ids[0] in [t[0] for t in chosen_tok_ids]:
                valid = False

            out_valid.append(valid)
            if valid:
                out_prompt.append(p_ids)
                out_chosen.append([t[0] for t in chosen_tok_ids])
                out_rej.append(rej_tok_ids[0])
            else:
                out_prompt.append([0]); out_chosen.append([0]); out_rej.append(0)

        return {
            "prompt_ids":        out_prompt,
            "chosen_ids":        out_chosen,
            "rejected_token_id": out_rej,
            "__valid":           out_valid,
        }

    ds = ds.map(
        _tok, batched=True, batch_size=batch_size,
        remove_columns=ds.column_names,
        num_proc=num_proc, desc="tokenising",
    )
    ds = ds.filter(lambda ex: ex["__valid"], num_proc=num_proc, desc="filter")
    ds = ds.remove_columns("__valid")
    if len(ds) == 0:
        raise ValueError("no ftpo samples survived length / sanity checks")

    return ds.shuffle(seed=3407)
