"""M1 Competition 111 forecasting eval (macro MAPE) for byte-level spark-gpt OEIS checkpoints.

A PyTorch port of eval_m1_competition_mape_mlx.py from the NextTerm-440M repository (the protocol
behind the NextTerm model cards: Naive2 17.7987, NextTerm-440M 17.6239, NextTerm-47M 18.7621),
unchanged in substance:
  * each series' context is scaled to integers by the largest number of decimal places visible
    in the context, joined by commas, with a trailing comma
  * greedy decoding with EOS and PAD suppressed, stopping once `horizon` commas are generated
    (or after --max-new-tokens); series are decoded together in left-padded batches with
    per-row stopping (same greedy result as one at a time, up to batching numerics)
  * the first `horizon` generated integers, divided back by the scale, are the forecast;
    missing terms are scored as 0 (missing_policy "zero")
  * series MAPE = mean APE over the horizon (zero actuals skipped); macro MAPE = mean over series
    (the number on the model cards); median MAPE = median over series is also reported, since
    one runaway forecast (e.g. a parabola on a flat series, APE ~700) can swing the macro mean
    by several points
  * prompts of --max-context-tokens or more tokens (BOS included) are skipped

Usage: python eval_m1.py <hf_dir> <m1_competition_111.jsonl> <out_prefix>
Writes <out_prefix>.series.jsonl and <out_prefix>.summary.json.
"""
import argparse
import json
import math
from decimal import Decimal, InvalidOperation, localcontext
from statistics import mean, median

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


def parse_decimal(raw):
    try:
        value = Decimal(str(raw).strip())
    except InvalidOperation as exc:
        raise ValueError(f"Could not parse decimal value {raw!r}") from exc
    if not value.is_finite():
        raise ValueError(f"Non-finite decimal value {raw!r}")
    return value


def context_scale(raw_context):
    places = 0
    for raw in raw_context:
        exponent = parse_decimal(raw).as_tuple().exponent
        if exponent < 0:
            places = max(places, -exponent)
    return 10 ** places


def scale_decimal_to_int(value, scale):
    scaled = value * Decimal(scale)
    with localcontext() as ctx:
        ctx.prec = max(50, len(scaled.as_tuple().digits) + 10)
        rounded = scaled.to_integral_value()
    if rounded != scaled:
        raise ValueError(f"Context scale {scale} does not make {value} integral")
    return int(rounded)


def parse_generated_terms(text, limit):
    terms, current, malformed = [], [], False

    def flush():
        nonlocal malformed
        if not current:
            return
        token = "".join(current)
        current.clear()
        if token == "-":
            malformed = True
            return
        try:
            terms.append(int(token))
        except ValueError:
            malformed = True

    for ch in text:
        if len(terms) >= limit:
            break
        if ch.isdigit() or (ch == "-" and not current):
            current.append(ch)
        elif ch == ",":
            flush()
        elif ch.isspace():
            continue
        else:
            malformed = True
            flush()
            break
    if len(terms) < limit:
        flush()
    return terms[:limit], malformed


def ape(actual, prediction):
    if actual == 0:
        return None
    return float(abs(actual - prediction) / abs(actual) * Decimal(100))


def mape_for_predictions(actuals, predictions, missing_policy, fallback):
    apes, missing = [], 0
    for actual, pred in zip(actuals, predictions):
        if pred is None:
            missing += 1
            if missing_policy == "skip":
                apes.append(None)
                continue
            pred = Decimal(0) if missing_policy == "zero" else fallback
        apes.append(ape(actual, pred))
    valid = [x for x in apes if x is not None and math.isfinite(x)]
    return (mean(valid) if valid else None), apes, missing


def main():
    p = argparse.ArgumentParser()
    p.add_argument("hf_dir")
    p.add_argument("data")
    p.add_argument("out_prefix")
    p.add_argument("--max-new-tokens", type=int, default=1024)
    p.add_argument("--max-context-tokens", type=int, default=2048)
    p.add_argument("--missing-policy", choices=["zero", "last_context", "skip"], default="zero")
    p.add_argument("--batch-size", type=int, default=128,
                   help="series decoded together (greedy, left-padded, per-row stopping)")
    args = p.parse_args()

    tok = AutoTokenizer.from_pretrained(args.hf_dir)
    model = AutoModelForCausalLM.from_pretrained(args.hf_dir, dtype=torch.float32).cuda().eval()
    bos, eos, pad = tok.bos_token_id, tok.eos_token_id, tok.pad_token_id
    comma = ord(",")
    suppress = torch.tensor([t for t in (eos, pad) if t is not None], device="cuda")

    rows = [json.loads(line) for line in open(args.data) if line.strip()]
    series = []
    skipped_long = 0
    for row in rows:
        raw_context = [str(x) for x in row["context"]]
        context = [parse_decimal(x) for x in raw_context]
        scale = context_scale(raw_context)
        prompt_text = ",".join(str(scale_decimal_to_int(v, scale)) for v in context) + ","
        prompt = [bos] + list(prompt_text.encode("utf-8"))
        assert tok(prompt_text).input_ids == prompt, "tokenizer disagrees with BOS + bytes"
        if args.max_context_tokens > 0 and len(prompt) >= args.max_context_tokens:
            skipped_long += 1
            continue
        series.append((row, raw_context, context, scale, prompt, int(row["horizon"])))

    @torch.no_grad()
    def greedy_batch(batch):
        """Greedy decode a left-padded batch with a KV cache, EOS/PAD suppressed; each row
        stops once it has generated `horizon` commas (or after --max-new-tokens)."""
        width = max(len(s[4]) for s in batch)
        ids = torch.full((len(batch), width), pad, dtype=torch.long, device="cuda")
        mask = torch.zeros((len(batch), width), dtype=torch.long, device="cuda")
        for r, s in enumerate(batch):
            ids[r, width - len(s[4]):] = torch.tensor(s[4], device="cuda")
            mask[r, width - len(s[4]):] = 1
        pos = (mask.cumsum(-1) - 1).clamp(min=0)
        out = model(ids, attention_mask=mask, position_ids=pos, use_cache=True)
        past, logits = out.past_key_values, out.logits[:, -1]
        horizon = torch.tensor([s[5] for s in batch], device="cuda")
        commas = torch.zeros(len(batch), dtype=torch.long, device="cuda")
        done = torch.zeros(len(batch), dtype=torch.bool, device="cuda")
        last_pos = pos[:, -1]
        generated = []
        for _ in range(args.max_new_tokens):
            logits = logits.float()
            logits[:, suppress] = float("-inf")
            nxt = logits.argmax(-1)
            generated.append(torch.where(done, torch.full_like(nxt, pad), nxt))
            commas += (nxt == comma) & ~done
            done |= commas >= horizon
            if bool(done.all()):
                break
            mask = torch.cat([mask, torch.ones_like(mask[:, :1])], dim=1)
            last_pos = last_pos + 1
            out = model(nxt[:, None], attention_mask=mask, position_ids=last_pos[:, None],
                        past_key_values=past, use_cache=True)
            past, logits = out.past_key_values, out.logits[:, -1]
        gen = torch.stack(generated, dim=1).tolist()
        return [bytes(t for t in g if t < 256).decode("utf-8", "replace") for g in gen]

    texts = []
    for b in range(0, len(series), args.batch_size):
        texts.extend(greedy_batch(series[b:b + args.batch_size]))

    records = []
    for (row, raw_context, context, scale, prompt, horizon), text in zip(series, texts):
        target = [parse_decimal(str(x)) for x in row["target"]]
        naive2 = [parse_decimal(str(x)) for x in row["naive2_forecast"]]
        terms, malformed = parse_generated_terms(text, horizon)
        predictions = [Decimal(t) / Decimal(scale) for t in terms] + [None] * (horizon - len(terms))
        fallback = context[-1] if context else Decimal(0)
        mape, apes, missing = mape_for_predictions(target, predictions, args.missing_policy, fallback)
        naive2_mape, naive2_apes, _ = mape_for_predictions(target, naive2, "skip", fallback)
        records.append({
            "row_index": row["row_index"], "series_name": row["series_name"], "period": row["period"],
            "horizon": horizon, "scale": scale, "prompt_tokens": len(prompt),
            "target": [str(x) for x in target],
            "prediction": [str(x) if x is not None else None for x in predictions],
            "parsed_terms": len(terms), "missing_terms": missing, "malformed": malformed,
            "mape": mape, "apes": apes, "naive2_mape": naive2_mape, "naive2_apes": naive2_apes,
            "generated_text": text,
        })

    def finite(rs, key):
        return [float(r[key]) for r in rs if r[key] is not None and math.isfinite(float(r[key]))]

    def macro(rs, key):
        vals = finite(rs, key)
        return mean(vals) if vals else None

    def med(rs, key):
        vals = finite(rs, key)
        return median(vals) if vals else None

    by_period = {}
    for period in sorted({r["period"] for r in records}):
        rs = [r for r in records if r["period"] == period]
        by_period[period] = {"series": len(rs),
                             "model_median_series_mape": med(rs, "mape"),
                             "model_macro_mape": macro(rs, "mape"),
                             "naive2_median_series_mape": med(rs, "naive2_mape"),
                             "naive2_macro_mape": macro(rs, "naive2_mape")}
    summary = {
        "checkpoint": args.hf_dir,
        "series_evaluated": len(records),
        "skipped_long": skipped_long,
        "model_median_series_mape": med(records, "mape"),
        "model_macro_mape": macro(records, "mape"),
        "naive2_median_series_mape": med(records, "naive2_mape"),
        "naive2_macro_mape": macro(records, "naive2_mape"),
        "parsed_full_series": sum(1 for r in records if r["parsed_terms"] >= r["horizon"]),
        "total_missing_terms": sum(r["missing_terms"] for r in records),
        "malformed_series": sum(1 for r in records if r["malformed"]),
        "by_period": by_period,
    }
    with open(args.out_prefix + ".series.jsonl", "w") as f:
        for r in records:
            f.write(json.dumps(r) + "\n")
    with open(args.out_prefix + ".summary.json", "w") as f:
        json.dump(summary, f, indent=2)
    print(json.dumps({k: summary[k] for k in ["series_evaluated", "skipped_long",
                                              "model_median_series_mape", "naive2_median_series_mape",
                                              "model_macro_mape", "naive2_macro_mape", "parsed_full_series",
                                              "total_missing_terms", "malformed_series"]}))


if __name__ == "__main__":
    main()