"""Reading an eval log.

Two layouts are accepted (web/router.js reads the same two):

  wide  one row per prompt:
        prompt_id, prompt, tiny_correct, tiny_cost, tiny_latency, tiny_confidence, small_correct, ...
        Required per model: <model>_correct and <model>_cost. Optional: <model>_latency, <model>_confidence.
  long  one row per (prompt, model):
        prompt_id, model, correct, cost [, latency] [, confidence] [, prompt]

`correct` is 0/1 or a score in [0, 1]. `cost` is in any unit (dollars, cents, tokens); the report
uses the unit as given. `prompt` (the text) is optional and only needed to learn a router when
there is no confidence column. Prompts that do not have a row for every model are dropped.
Models are ordered from cheapest to most expensive by mean cost.
"""
from __future__ import annotations

import csv
import io
import math
from dataclasses import dataclass, field
from typing import Optional

import numpy as np

ALIASES = {
    "id": ("prompt_id", "id", "question_id", "example_id", "row_id", "prompt id"),
    "prompt": ("prompt", "question", "input", "text", "query"),
    "model": ("model", "model_name", "system"),
    "correct": ("correct", "is_correct", "score", "accuracy", "acc", "quality", "right"),
    "cost": ("cost", "cost_usd", "price", "usd", "dollars", "tokens", "total_tokens"),
    "latency": ("latency", "latency_ms", "latency_s", "time", "seconds", "ms"),
    "confidence": ("confidence", "conf", "p_correct", "probability", "prob", "self_confidence"),
}


@dataclass
class EvalLog:
    models: list                      # cheap -> expensive
    ids: list
    correct: np.ndarray               # (n, m) in [0, 1]
    cost: np.ndarray                  # (n, m)
    latency: Optional[np.ndarray]     # (n, m) or None
    confidence: Optional[np.ndarray]  # (n, m) or None
    prompts: Optional[list]           # n strings or None
    dropped: int = 0                  # prompts without a row for every model
    warnings: list = field(default_factory=list)

    @property
    def n(self) -> int:
        return self.correct.shape[0]

    @property
    def m(self) -> int:
        return self.correct.shape[1]

    @property
    def binary(self) -> bool:
        return bool(np.all((self.correct == 0) | (self.correct == 1)))


def _norm(name: str) -> str:
    return name.strip().lower().replace("-", "_").replace(" ", "_").lstrip("")


def _num(s):
    s = (s or "").strip().lower()
    if s in ("", "na", "nan", "none", "null", "-"):
        return math.nan
    if s in ("true", "yes", "y", "t", "correct", "pass"):
        return 1.0
    if s in ("false", "no", "n", "f", "incorrect", "fail", "wrong"):
        return 0.0
    try:
        return float(s)
    except ValueError:
        return math.nan


def _find(cols, key):
    for a in ALIASES[key]:
        if a in cols:
            return a
    return None


def parse_csv_text(text: str):
    if text.startswith(""):
        text = text[1:]
    sample = text[:20000]
    delim = ","
    best = -1
    for d in (",", ";", "\t", "|"):
        lines = [l for l in sample.split("\n")[:20] if l.strip()]
        if not lines:
            continue
        counts = [l.count(d) for l in lines]
        if counts[0] == 0:
            continue
        score = sum(1 for c in counts if c == counts[0]) * 1000 + counts[0]
        if score > best:
            best, delim = score, d
    rows = list(csv.reader(io.StringIO(text), delimiter=delim))
    rows = [r for r in rows if any(c.strip() for c in r)]
    if not rows:
        raise ValueError("The file is empty.")
    header = [_norm(h) for h in rows[0]]
    return header, rows[1:]


def load_text(text: str) -> EvalLog:
    header, rows = parse_csv_text(text)
    cols = {h: i for i, h in enumerate(header)}
    model_col = _find(cols, "model")
    if model_col is not None and _find(cols, "correct") is not None:
        return _load_long(header, cols, rows)
    return _load_wide(header, cols, rows)


def load_csv(path: str) -> EvalLog:
    with open(path, encoding="utf-8-sig", newline="") as f:
        return load_text(f.read())


def _finish(models, ids, correct, cost, latency, confidence, prompts, dropped, warnings):
    correct = np.asarray(correct, dtype=float)
    cost = np.asarray(cost, dtype=float)
    if correct.ndim != 2 or correct.shape[0] == 0:
        raise ValueError("No complete prompts found: every prompt needs a correct and a cost value for every model.")
    if np.nanmin(correct) < 0 or np.nanmax(correct) > 1:
        raise ValueError("The correct/score column must be 0/1 or a number between 0 and 1.")
    if latency is not None:
        latency = np.asarray(latency, dtype=float)
        if np.isnan(latency).any():
            warnings.append("Some latency values are missing; they were treated as 0.")
            latency = np.nan_to_num(latency)
    if confidence is not None:
        confidence = np.asarray(confidence, dtype=float)
        if np.isnan(confidence).any():
            warnings.append("Some confidence values are missing; they were treated as 0 (never accepted).")
            confidence = np.nan_to_num(confidence)
        if confidence.min() < 0 or confidence.max() > 1:
            lo, hi = confidence.min(), confidence.max()
            confidence = (confidence - lo) / (hi - lo) if hi > lo else np.zeros_like(confidence)
            warnings.append("Confidence values were outside [0, 1] and have been rescaled to it.")
    # order models cheap -> expensive by mean cost (ties by name), identical to the JS
    mean_cost = cost.mean(axis=0)
    order = sorted(range(len(models)), key=lambda i: (float(mean_cost[i]), models[i]))
    models = [models[i] for i in order]
    correct = correct[:, order]
    cost = cost[:, order]
    latency = latency[:, order] if latency is not None else None
    confidence = confidence[:, order] if confidence is not None else None
    if prompts is not None and all(not p for p in prompts):
        prompts = None
    return EvalLog(models, ids, correct, cost, latency, confidence, prompts, dropped, warnings)


def _load_long(header, cols, rows):
    c_id, c_prompt, c_model = _find(cols, "id"), _find(cols, "prompt"), _find(cols, "model")
    c_cor, c_cost, c_lat, c_conf = _find(cols, "correct"), _find(cols, "cost"), _find(cols, "latency"), _find(cols, "confidence")
    if c_cost is None:
        raise ValueError("Long format needs a cost column (cost, cost_usd, price or tokens).")
    models, seen = [], {}
    by_id = {}
    order = []
    for k, r in enumerate(rows):
        g = lambda c: r[cols[c]] if c is not None and cols[c] < len(r) else ""
        pid = g(c_id) if c_id is not None else str(k)
        model = g(c_model).strip()
        if not model:
            continue
        if model not in seen:
            seen[model] = len(models)
            models.append(model)
        if pid not in by_id:
            by_id[pid] = {"prompt": g(c_prompt), "rows": {}}
            order.append(pid)
        by_id[pid]["rows"][model] = (_num(g(c_cor)), _num(g(c_cost)),
                                     _num(g(c_lat)) if c_lat is not None else None,
                                     _num(g(c_conf)) if c_conf is not None else None)
    ids, prompts, correct, cost, lat, conf, dropped = [], [], [], [], [], [], 0
    for pid in order:
        e = by_id[pid]
        if any(mo not in e["rows"] for mo in models):
            dropped += 1
            continue
        vals = [e["rows"][mo] for mo in models]
        if any(math.isnan(v[0]) or math.isnan(v[1]) for v in vals):
            dropped += 1
            continue
        ids.append(pid)
        prompts.append(e["prompt"])
        correct.append([v[0] for v in vals])
        cost.append([v[1] for v in vals])
        if c_lat is not None:
            lat.append([v[2] for v in vals])
        if c_conf is not None:
            conf.append([v[3] for v in vals])
    warnings = []
    if dropped:
        warnings.append(f"{dropped} prompt(s) were dropped because not every model had a correct and cost value for them.")
    return _finish(models, ids, correct, cost, lat if c_lat is not None else None,
                   conf if c_conf is not None else None, prompts, dropped, warnings)


def _load_wide(header, cols, rows):
    c_id, c_prompt = _find(cols, "id"), _find(cols, "prompt")
    suffix = {}
    for key in ("correct", "cost", "latency", "confidence"):
        for a in ALIASES[key]:
            for h in header:
                if h.endswith("_" + a) and len(h) > len(a) + 1:
                    name = h[: -len(a) - 1]
                    suffix.setdefault(name, {})
                    suffix[name].setdefault(key, h)
    models = [mo for mo in suffix if "correct" in suffix[mo] and "cost" in suffix[mo]]
    # keep the order in which the models' correct columns appear
    models.sort(key=lambda mo: cols[suffix[mo]["correct"]])
    if len(models) < 2:
        raise ValueError("Could not find at least two models. Wide format needs <model>_correct and <model>_cost "
                         "columns; long format needs model, correct and cost columns.")
    has_lat = all("latency" in suffix[mo] for mo in models)
    has_conf = all("confidence" in suffix[mo] for mo in models)
    ids, prompts, correct, cost, lat, conf, dropped = [], [], [], [], [], [], 0
    for k, r in enumerate(rows):
        g = lambda c: r[cols[c]] if cols[c] < len(r) else ""
        cor = [_num(g(suffix[mo]["correct"])) for mo in models]
        co = [_num(g(suffix[mo]["cost"])) for mo in models]
        if any(math.isnan(v) for v in cor) or any(math.isnan(v) for v in co):
            dropped += 1
            continue
        ids.append(g(c_id) if c_id is not None else str(k))
        prompts.append(g(c_prompt) if c_prompt is not None else "")
        correct.append(cor)
        cost.append(co)
        if has_lat:
            lat.append([_num(g(suffix[mo]["latency"])) for mo in models])
        if has_conf:
            conf.append([_num(g(suffix[mo]["confidence"])) for mo in models])
    warnings = []
    if dropped:
        warnings.append(f"{dropped} row(s) were dropped because a correct or cost value was missing.")
    if not has_conf and any("confidence" in suffix[mo] for mo in models):
        warnings.append("Only some models have a confidence column, so confidences were ignored.")
    return _finish(models, ids, correct, cost, lat if has_lat else None, conf if has_conf else None,
                   prompts, dropped, warnings)


SAMPLE_WIDE = """prompt_id,prompt,tiny_correct,tiny_cost,tiny_latency,tiny_confidence,frontier_correct,frontier_cost,frontier_latency,frontier_confidence
q1,What is 17 times 23?,1,0.00004,310,0.91,1,0.0121,2400,0.98
q2,Summarize this contract clause: ...,0,0.00009,620,0.77,1,0.0263,4100,0.93
q3,Who painted The Night Watch?,1,0.00003,280,0.95,1,0.0097,1900,0.99
"""

SAMPLE_LONG = """prompt_id,model,correct,cost,latency,confidence,prompt
q1,tiny,1,0.00004,310,0.91,What is 17 times 23?
q1,frontier,1,0.0121,2400,0.98,What is 17 times 23?
q2,tiny,0,0.00009,620,0.77,Summarize this contract clause: ...
q2,frontier,1,0.0263,4100,0.93,Summarize this contract clause: ...
"""