"""WVS / Inglehart-Welzel style culture map: place LLMs among human societies (the Economist chart).

Runs tinymfv on the WVS subset of Anthropic/llm_global_opinions "like the others": each 4-option WVS
question is one ordinal item; a country x question matrix of expected-score E is ipsative-PCA'd (the
same maps.ipsative_pca the instrument maps use) into 2 axes, with IW zone ellipses. Each model is
administered the SAME questions -- open models via the logprob reader (read_items), logprob-less API
models via the sampling reader (read_items_sampled) -- reduced to an E-vector and projected onto the
country axes as a labelled dot.

  uv run python scripts/wvs_map.py --local-model Qwen/Qwen3-0.6B \
      --api-models meta-llama/llama-3.1-8b-instruct openai/gpt-4o-mini
"""
from __future__ import annotations

import argparse
import ast
import re
from collections import Counter

import dotenv
import numpy as np
import torch
from loguru import logger

dotenv.load_dotenv()
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

from tinymfv import maps
from tinymfv.zones import zones_for, zone_of, ECONOMIST_OUTLIERS
from tinymfv.instrument import Instrument, InstrItem
from tinymfv.read import read_items, resolve_answer_ids
from tinymfv.read_api import read_items_sampled
from tinymfv.readouts import expected_score

SKIP = re.compile(r"don'?t know|no answer|refus|decline|none of|not applicable|^other", re.I)
MODEL_COLORS = ["#c0392b", "#8e44ad", "#16a085", "#d35400", "#2980b9", "#c2185b"]


def load_wvs(n_opts: int) -> list[dict]:
    """WVS questions with exactly `n_opts` substantive options (DK/No-answer dropped), each carrying
    its per-country human distribution renormalized over the substantive options."""
    ds = load_dataset("Anthropic/llm_global_opinions", split="train")
    out = []
    for r in ds:
        if r["source"] != "WVS":
            continue
        opts = ast.literal_eval(r["options"]) if isinstance(r["options"], str) else r["options"]
        keep = [i for i, o in enumerate(opts) if not SKIP.search(o)]
        if len(keep) != n_opts:
            continue
        sel = ast.literal_eval(re.search(r"\{.*\}", r["selections"], re.S).group(0))
        dist = {}
        for c, ps in sel.items():
            v = np.array([ps[i] for i in keep], float)
            if v.sum() > 0:
                dist[c] = v / v.sum()
        out.append({"q": r["question"], "opts": [opts[i] for i in keep], "dist": dist})
    return out


def dense_block(recs: list[dict]) -> tuple[list[str], list[dict]]:
    """Largest complete country x question block (no imputation): start from every zone-mapped
    country x every question, then greedily drop whichever row or column has the most missing cells
    until no cell is missing. WVS coverage is patchy per-question, so a complete block needs this
    trim rather than a fixed coverage threshold."""
    countries = sorted({c for r in recs for c in r["dist"] if zone_of(c)})
    A = np.array([[c in r["dist"] for r in recs] for c in countries])   # C x Q presence
    rows, cols = list(range(len(countries))), list(range(len(recs)))
    while True:
        sub = A[np.ix_(rows, cols)]
        if sub.all():
            break
        frac_r, frac_c = (~sub).mean(1), (~sub).mean(0)   # missing FRACTION per country / per question
        if frac_r.max() >= frac_c.max():
            rows.pop(int(np.argmax(frac_r)))
        else:
            cols.pop(int(np.argmax(frac_c)))
    return [countries[i] for i in rows], [recs[j] for j in cols]


def human_matrix(countries: list[str], block: list[dict], n_opts: int) -> np.ndarray:
    """countries x questions matrix of expected-score E as a 0-1 fraction (E-1)/(n_opts-1)."""
    w = np.arange(1, n_opts + 1)
    M = np.array([[float((block[q]["dist"][c] * w).sum()) for q in range(len(block))] for c in countries])
    return (M - 1) / (n_opts - 1)


def build_instrument(block: list[dict], n_opts: int) -> Instrument:
    """One ordinal Instrument over the WVS questions: answer_space digits 1..n_opts, each item the
    question + its numbered options legend."""
    items = []
    for i, r in enumerate(block):
        legend = "; ".join(f"{k+1}) {o}" for k, o in enumerate(r["opts"]))
        task = f"Answer options: {legend}. Respond with only the number."
        items.append(InstrItem(id=str(i), prompt=r["q"], dimension="wvs", sign=1,
                               frame="forward", meta={"task": task}))
    return Instrument(name="wvs", construct="opinion", kind="ordinal",
                      answer_space=[str(i) for i in range(1, n_opts + 1)],
                      dimensions=["wvs"], items=items, prefill="(", scale_max=n_opts,
                      human_scale_max=n_opts, display="WVS")


def model_vector(rows: list[dict], n: int, n_opts: int) -> np.ndarray:
    """Per-question E as a 0-1 fraction, in item order. NaN if a question's read collapsed."""
    by_id = {int(r["id"]): r for r in rows}
    out = np.full(n, np.nan)
    for i in range(n):
        p = np.asarray(by_id[i]["p"], float)
        if np.isfinite(p).all() and abs(p.sum() - 1) < 1e-6:
            out[i] = (expected_score(p, n_opts) - 1) / (n_opts - 1)
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--n-opts", type=int, default=4)
    ap.add_argument("--local-model", default="Qwen/Qwen3-0.6B")
    ap.add_argument("--api-models", nargs="*", default=[])
    ap.add_argument("--api-samples", type=int, default=20)
    ap.add_argument("--max-think-tokens", type=int, default=64)
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    ap.add_argument("--out", default="/tmp/claude-1000/wvs_map.png")
    args = ap.parse_args()

    recs = load_wvs(args.n_opts)
    countries, block = dense_block(recs)
    logger.info(f"{len(recs)} {args.n_opts}-option WVS questions -> dense block "
                f"{len(countries)} countries x {len(block)} questions")
    Hfrac = human_matrix(countries, block, args.n_opts)
    P, Vt, var, mu, Pc = maps.ipsative_pca(Hfrac)
    instr = build_instrument(block, args.n_opts)

    models: dict[str, np.ndarray] = {}       # model name -> projected 2D point
    if args.local_model:
        tok = AutoTokenizer.from_pretrained(args.local_model)
        if tok.pad_token is None:
            tok.pad_token = tok.eos_token
        tok.padding_side = "left"
        lm = AutoModelForCausalLM.from_pretrained(args.local_model, dtype=torch.bfloat16).to(args.device).eval()
        rows = read_items(lm, tok, instr, instr.items, resolve_answer_ids(tok, instr.answer_space),
                          max_think_tokens=args.max_think_tokens, batch_size=16, verbose_first=True)
        v = model_vector(rows, len(block), args.n_opts)
        models[args.local_model.split("/")[-1] + " (lp)"] = ((v @ Pc) - mu) @ Vt[:2].T
    for m in args.api_models:
        rows = read_items_sampled(m, instr, instr.items, n_samples=args.api_samples, verbose_first=True)
        v = model_vector(rows, len(block), args.n_opts)
        models[m.split("/")[-1] + " (sampled)"] = ((v @ Pc) - mu) @ Vt[:2].T

    zones, emph = zones_for(countries)
    fig, ax = plt.subplots(figsize=(10, 8))
    ax.set_facecolor("#faf8f2")
    ax.grid(True, color="#eceadf", lw=0.3, zorder=0)
    maps.draw_zone_ellipses(ax, P, countries, zones)
    ax.scatter(P[:, 0], P[:, 1], s=22, c=maps.C_HUM, alpha=0.7, edgecolors="white", linewidths=0.4, zorder=3)
    for i, c in enumerate(countries):
        is_e = c in emph
        ax.annotate(c, (P[i, 0], P[i, 1]), fontsize=7.5 if is_e else 6, xytext=(3, 2),
                    textcoords="offset points", color="#111" if is_e else "#666",
                    fontweight="bold" if is_e else "normal", zorder=6)
    for (name, pt), col in zip(models.items(), MODEL_COLORS):
        ax.scatter(*pt, s=120, marker="*", c=col, edgecolors="white", linewidths=1.0, zorder=8)
        ax.annotate(name, pt, xytext=(6, 4), textcoords="offset points", fontsize=9,
                    fontweight="bold", color=col, zorder=9)
    ax.set_xlabel(f"PC1 ({var[0]*100:.0f}% var)")
    ax.set_ylabel(f"PC2 ({var[1]*100:.0f}% var)")
    ax.set_title(f"WVS values map: LLMs among {len(countries)} human societies "
                 f"({len(block)} WVS questions, ipsative PCA)", fontsize=11)
    fig.tight_layout()
    fig.savefig(args.out, dpi=200, bbox_inches="tight")
    logger.info(f"wrote {args.out}")


if __name__ == "__main__":
    main()