J'ai entraîné un modèle qui retrouve la formule d'une courbe dessinée à la main

# Curvy

**[English](#english) · [Français](#français)**

---

<a id="english"></a>

## English

Compact symbolic regression. Give it a **2D point cloud** — a curve drawn
freehand, or a set of measurements — and it returns the **mathematical formulas**
that describe it, ranked on a simplicity/accuracy Pareto front.

31.6 M parameters, five Python dependencies, one consumer GPU.

### The idea in one sentence

The model **never predicts a numeric value**. It predicts a *skeleton*, such as
`C * sin(C * x) + C`, and the constants are fitted afterwards by numerical
optimisation.

Guessing *the shape* is a language problem; finding *the numbers* is an
optimisation problem. Two jobs, two tools.

```
point cloud ──▶ affine normalisation ──▶ transformer encoder (a set of points)
                                                  │
                                             cross-attention
                                                  ▼
                                        autoregressive decoder
                                                  │
                                    beam search under an arity mask
                                                  ▼
                                          N candidate skeletons
                                                  │
                                   constant fitting (scipy least squares)
                                                  ▼
                                   Pareto front (complexity, R²)
                                                  ▼
                                          3 to 5 formulas
```

### Measured results

On **1 960 formulas explicitly withheld from training**, plus 29 hand-written
ones. The score is the share of curves recovered at R² ≥ 0.99, judged against
the exact curve on held-out points, never against the noisy samples.

| | oracle | polynomial | **Curvy** |
|---|---|---|---|
| interpolation | 0.791 | 0.669 | **0.684** |
| **extrapolation** | 0.360 | 0.069 | **0.187** |

*(beam 8; sampling noise is ±2.2 points at 95 % on n = 1960)*

The "oracle" is the same measurement given the true skeleton: the score of
someone handed the answer. It is not 100 % because the samples are noisy and
constant fitting sometimes fails.

**In interpolation, `np.polyfit` is indistinguishable from this model** — within
the error margin, in one millisecond, without a GPU. The project does not earn
its keep there. It earns it outside the observed window, where a polynomial
diverges and a formula stays true: **2.7× better**, and on the 29 hand-written
formulas the polynomial scores **zero on all 29** in extrapolation.

On the Runge function `1/(1+25x²)`, the textbook counter-example to polynomial
interpolation: **Curvy recovers it six times out of six, `np.polyfit` zero times
out of six.**

### What it cannot do

- **Curves that fold back.** Circle, heart, loop. The system predicts `y = f(x)`;
  such a curve has two `y` for one `x`. Structural, not a training issue.
- **Discontinuities.** Step, floor, sawtooth: **0.000**. No discontinuous
  operator in the vocabulary.
- **Fast oscillations.** Constant fitting loses the frequency beyond roughly two
  oscillations per unit width. The model can propose the right formula and still
  fail, because the holes cannot be filled.
- **Deep formulas.** Past eight tree levels a polynomial takes the lead, and a
  137-character expression is no more readable than eight decimal coefficients.

### Getting started

```bash
uv venv --python 3.12 .venv
uv pip install --index-url https://download.pytorch.org/whl/cu128 \
    torch==2.9.1 numpy scipy sympy matplotlib
make setup     # editable install + pytest and ruff
make env       # what compute the machine actually offers
make test      # 230 tests
```

Heavy artefacts (venv, datasets, checkpoints) live outside the repository, under
`CURVY_DATA_ROOT`. Set it before anything else:

```bash
export CURVY_DATA_ROOT=~/curvy-data
```

`make help` lists every target.

#### Reproducing the model

```bash
# 1. Skeletons. The only symbolic step, paid once. ~75 s on 12 cores.
make data N=2000000 SEED=42

# 2. Freeze the test set BEFORE training, or the final number means nothing.
#    The script refuses to overwrite an existing set and prints a SHA-256 that
#    every published measurement must cite. See benchmarks/testset-v1.json.
.venv/bin/python scripts/build_testset.py

# 3. Train. ~6 h for 40 000 steps on an RTX 3090.
./scripts/launch_run.sh exp-005 v1 40000
CURVY_REPRENDRE=1 ./scripts/launch_run.sh exp-005 v1 40000   # resume

# 4. Evaluate on the frozen set. The SHA-256 is checked on every run.
.venv/bin/python scripts/eval_testset.py --run exp-005 --preset v1 --beam 8

# 5. Serve the browser demo.
make demo      # http://127.0.0.1:8001
```

The launcher excludes the test-set skeletons by default, runs inside `tmux`, and
caps the CPU footprint at five workers at low priority. That last point is not
comfort: an eight-worker run once took the machine's network down. A checkpoint
is written every 1 000 steps together with the optimiser state, so handing the
GPU back to someone else costs at most 1 000 steps, never the run.

#### Before training on a shared machine

The GPU allow-list is set **by UUID**, not by index: a UUID survives
renumbering, an index does not. Put yours — `nvidia-smi -L` prints it — in the
`CURVY_CUDA_ALLOW` variable of the `Makefile` and of `scripts/launch_run.sh`.
An allow-list that matches nothing **raises** rather than silently falling back
to the CPU.

Worth knowing: the CUDA runtime numbers GPUs by compute capability while
`nvidia-smi` numbers them by PCI bus. `cuda:0` and `nvidia-smi`'s `0` need not
be the same card.

### Layout

| | |
|---|---|
| `curvy/data/` | grammar, sampling, canonicalisation, point clouds, dataset |
| `curvy/model/` | the transformer |
| `curvy/train/` | training loop, metrics, checkpoints |
| `curvy/infer/` | beam search, key/value cache, constant fitting, Pareto front |
| `curvy/serve/` | HTTP service, rate limiter, drawing pipeline |
| `scripts/` | generation, evaluation, figures, launchers |
| `tests/` | 230 tests |

The HTTP service deliberately uses the standard library rather than a web
framework: five dependencies is a feature, and a service that does one thing
does not need a router.

### License

MIT. See [LICENSE](LICENSE).

---

<a id="français"></a>

## Français

Régression symbolique compacte. On lui donne un **nuage de points 2D**, une
courbe tracée à main levée ou un jeu de mesures, et il rend les **formules
mathématiques** qui le décrivent, classées sur un front de Pareto
simplicité / précision.

31,6 M de paramètres, cinq dépendances Python, une carte grand public.

### L'idée en une phrase

Le modèle ne prédit **jamais** de valeur numérique. Il prédit un *squelette*,
`C * sin(C * x) + C`, et les constantes sont ajustées après coup par
optimisation numérique.

Deviner *la forme* est un problème de langage, trouver *les nombres* est un
problème d'optimisation. Deux métiers, deux outils.

```
nuage de points ──▶ normalisation affine ──▶ encodeur transformer (un ensemble)
                                                       │
                                                 attention croisée
                                                       ▼
                                             décodeur autorégressif
                                                       │
                                     beam search sous masque d'arité
                                                       ▼
                                            N squelettes candidats
                                                       │
                                   ajustement des constantes (scipy)
                                                       ▼
                                    front de Pareto (complexité, R²)
                                                       ▼
                                              3 à 5 formules
```

### Résultats mesurés

Sur **1 960 formules explicitement retirées de l'entraînement**, plus 29 écrites
à la main. Le score est la part de courbes retrouvées à R² ≥ 0,99, jugé contre
la courbe exacte sur des points tenus à l'écart, jamais contre les points
bruités.

| | oracle | polynôme | **Curvy** |
|---|---|---|---|
| interpolation | 0,791 | 0,669 | **0,684** |
| **extrapolation** | 0,360 | 0,069 | **0,187** |

*(beam 8 ; le bruit d'échantillonnage vaut ±2,2 points à 95 % sur n = 1960)*

L'« oracle » est la même mesure appliquée au vrai squelette : le score de
quelqu'un à qui on aurait donné la solution. Il ne vaut pas 100 % parce que les
points sont bruités et que l'ajustement des constantes échoue parfois.

**En interpolation, `np.polyfit` est indiscernable de ce modèle**, à l'intérieur
de la marge d'erreur, en une milliseconde et sans GPU. Le projet ne se justifie
pas là. Il se justifie hors de la fenêtre observée, où un polynôme diverge et où
une formule reste vraie : **2,7 fois mieux**, et sur les 29 formules écrites à
la main le polynôme rend **zéro sur 29** en extrapolation.

Sur la fonction de Runge `1/(1+25x²)`, le contre-exemple classique de
l'interpolation polynomiale : **Curvy la retrouve six fois sur six,
`np.polyfit` zéro fois sur six.**

### Ce qu'il ne sait pas faire

- **Les courbes qui reviennent en arrière.** Cercle, cœur, boucle. Le système
  prédit `y = f(x)` ; une telle courbe a deux `y` pour un même `x`. C'est
  structurel, pas un défaut d'entraînement.
- **Les discontinuités.** Marche, plancher, dent de scie : **0,000**. Aucun
  opérateur discontinu dans le vocabulaire.
- **Les oscillations rapides.** L'ajustement perd la fréquence au-delà d'environ
  deux oscillations par unité de largeur. Le modèle peut proposer la bonne
  formule et échouer quand même, faute de pouvoir remplir les trous.
- **Les formules profondes.** Au-delà de huit niveaux d'arbre, le polynôme
  reprend l'avantage, et une expression de 137 caractères n'est pas plus lisible
  que huit coefficients décimaux.

### Démarrage

```bash
uv venv --python 3.12 .venv
uv pip install --index-url https://download.pytorch.org/whl/cu128 \
    torch==2.9.1 numpy scipy sympy matplotlib
make setup     # installation en editable + pytest et ruff
make env       # ce que la machine offre réellement comme calcul
make test      # 230 tests
```

Les artefacts lourds (venv, jeux de données, points de reprise) vivent hors du
dépôt, sous `CURVY_DATA_ROOT`. À définir avant toute chose :

```bash
export CURVY_DATA_ROOT=~/curvy-data
```

`make help` liste toutes les cibles.

#### Reproduire le modèle

```bash
# 1. Les squelettes. Seule étape symbolique, payée une fois. ~75 s sur 12 cœurs.
make data N=2000000 SEED=42

# 2. Geler le jeu de test AVANT d'entraîner, sinon le chiffre final ne vaut
#    rien. Le script refuse d'écraser un jeu existant et publie une empreinte
#    SHA-256 que toute mesure doit citer. Voir benchmarks/testset-v1.json.
.venv/bin/python scripts/build_testset.py

# 3. Entraîner. ~6 h pour 40 000 étapes sur une RTX 3090.
./scripts/launch_run.sh exp-005 v1 40000
CURVY_REPRENDRE=1 ./scripts/launch_run.sh exp-005 v1 40000   # reprendre

# 4. Évaluer sur le jeu figé. L'empreinte est vérifiée à chaque exécution.
.venv/bin/python scripts/eval_testset.py --run exp-005 --preset v1 --beam 8

# 5. Servir la démo navigateur.
make demo      # http://127.0.0.1:8001
```

Le lanceur exclut par défaut les squelettes du jeu de test, tourne dans `tmux`,
et bride l'empreinte processeur à cinq travailleurs en priorité basse. Ce
dernier point n'est pas du confort : un entraînement à huit travailleurs a déjà
fait tomber le réseau de la machine. Un point de reprise est écrit toutes les
1 000 étapes avec l'état de l'optimiseur, si bien que rendre la carte à
quelqu'un d'autre coûte au pire 1 000 étapes, jamais l'entraînement.

#### Avant d'entraîner sur une machine partagée

La liste blanche des GPU se règle **par UUID** et non par index : un UUID survit
à une renumérotation, un index non. Renseignez le vôtre, que `nvidia-smi -L`
affiche, dans la variable `CURVY_CUDA_ALLOW` du `Makefile` et de
`scripts/launch_run.sh`. Une liste qui ne correspond à rien **lève une erreur**
plutôt que de basculer en silence sur le CPU.

À savoir : le runtime CUDA numérote les cartes par puissance quand `nvidia-smi`
les numérote par bus PCI. `cuda:0` et le `0` de `nvidia-smi` ne désignent pas
forcément la même carte.

### Organisation

| | |
|---|---|
| `curvy/data/` | grammaire, tirage, canonicalisation, nuages, jeu de données |
| `curvy/model/` | le transformer |
| `curvy/train/` | boucle d'entraînement, métriques, points de reprise |
| `curvy/infer/` | beam search, cache clé/valeur, ajustement, front de Pareto |
| `curvy/serve/` | service HTTP, limiteur de débit, traitement du tracé |
| `scripts/` | génération, évaluation, figures, lanceurs |
| `tests/` | 230 tests |

Le service HTTP s'appuie délibérément sur la bibliothèque standard plutôt que
sur un cadriciel web : cinq dépendances est une qualité, et un service qui fait
une seule chose n'a pas besoin d'un routeur.

### Licence

MIT. Voir [LICENSE](LICENSE).

MIT License

Copyright (c) 2026 Billy Girboux

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

[build-system]
requires = ["setuptools>=68"]
build-backend = "setuptools.build_meta"

[project]
name = "curvy"
version = "0.0.1"
description = "Régression symbolique compacte : d'un nuage de points 2D vers un front de Pareto de formules"
readme = "README.md"
requires-python = ">=3.10"
authors = [{ name = "Billy Girboux" }]
license = "MIT"
license-files = ["LICENSE"]
dependencies = [
    "torch>=2.5",
    "numpy>=1.26",
    "scipy>=1.11",
    "sympy>=1.12",
    "matplotlib>=3.8",
]

[project.optional-dependencies]
dev = ["pytest>=8", "ruff>=0.6"]

[tool.setuptools.packages.find]
include = ["curvy*"]

[tool.ruff]
line-length = 100
target-version = "py310"

[tool.ruff.lint]
select = ["E", "F", "I", "UP", "B", "SIM", "NPY"]
ignore = ["E501"]

[tool.pytest.ini_options]
testpaths = ["tests"]
addopts = "-q --strict-markers"
markers = [
    "slow: test lent, exclu du cycle rapide (make test)",
    "gpu: nécessite un accélérateur",
]

# Curvy — cibles principales.
# Le venv et les artefacts lourds vivent hors du dépôt, voir CURVY_DATA_ROOT.
PY      := .venv/bin/python
SEED    ?= 42
RUN     ?= exp-005
PRESET  ?= v1
# 127.0.0.1 par défaut. `make demo HOST=0.0.0.0` pour ouvrir au réseau local.
HOST    ?= 127.0.0.1
PORT    ?= 8001
BEAM    ?= 48
DEVICE  ?= auto
N       ?= 500000

# Le runtime CUDA numérote les GPU par puissance et non par bus PCI ; on force
# l'ordre PCI pour que `cuda:0` désigne la même carte que dans nvidia-smi.
export CUDA_DEVICE_ORDER := PCI_BUS_ID

# Liste blanche des GPU utilisables, par UUID plutôt que par index : un UUID
# survit à une renumérotation, un index non. À renseigner avec le vôtre, que
# `nvidia-smi -L` affiche. Une liste qui ne correspond à rien lève une erreur
# plutôt que de basculer en silence sur le CPU.
export CURVY_CUDA_ALLOW := GPU-1234abcd-0000-0000-0000-000000000000

# Vérifie la VRAM libre et signale les squatteurs avant tout entraînement.
gpu-check:       ## État de la 3090 + processus qui l'occupent
	$(PY) -m curvy.cli_gpu

.PHONY: help setup env gpu-check data train train-tmux curves eval demo test lint fmt clean

help:            ## Liste les cibles
	@grep -E '^[a-z-]+:.*?## ' $(MAKEFILE_LIST) | awk -F':.*?## ' '{printf "  \033[36m%-10s\033[0m %s\n", $$1, $$2}'

setup:           ## Installe le paquet en editable + outils de dev
	uv pip install --python $(PY) -e ".[dev]"

env:             ## Affiche l'environnement de calcul retenu (livrable Phase 0)
	$(PY) -m curvy.cli_env

data:            ## Génère l'ensemble de squelettes (Phase 1)
	$(PY) -m curvy.data.generate --n $(N) --seed $(SEED)

train: gpu-check ## Entraîne un modèle (Phase 4) — refuse de démarrer si la VRAM est prise
	$(PY) -m curvy.train.run --seed $(SEED)

train-tmux: gpu-check ## Lance les entraînements longs dans tmux (survit à la déconnexion)
	./scripts/launch_training.sh

curves:          ## Trace les courbes d'un run : make curves RUN=exp-001
	$(PY) scripts/plot_training.py --run $(RUN)

eval:            ## Évalue sur le jeu de test figé (Phase 6)
	$(PY) -m curvy.eval.run --device $(DEVICE)

demo:            ## Sert la démo navigateur sur la 3060 (Phase 8) — http://127.0.0.1:8001
	$(PY) -m curvy.serve.app --run $(RUN) --preset $(PRESET) --host $(HOST) --port $(PORT) --beam $(BEAM)

test:            ## Tests rapides
	$(PY) -m pytest -m "not slow"

lint:            ## Ruff (lint + format check)
	$(PY) -m ruff check curvy tests scripts
	$(PY) -m ruff format --check curvy tests scripts

fmt:             ## Ruff format
	$(PY) -m ruff format curvy tests scripts

clean:           ## Supprime les caches Python
	find . -name __pycache__ -type d -prune -exec rm -rf {} +
	rm -rf .pytest_cache .ruff_cache

# Les jeux de données et les points de reprise vivent sous CURVY_DATA_ROOT,
# hors de l'arbre de travail. Git ne les voit donc pas, et les deux règles
# ci-dessous ne servent que si l'on pose CURVY_DATA_ROOT sur un chemin local.
#
# Elles sont ancrées par un `/` initial, et ce détail compte : sans ancrage,
# `data/` attrape n'importe quel répertoire de ce nom à n'importe quelle
# profondeur, `curvy/data/` compris. Git ignore en silence — `git add` sur un
# fichier ignoré ne renvoie aucune erreur — et un paquet entier peut ainsi
# rester invisible pendant des jours, ses propres commits compris.
.venv
/data/
/runs/
*.ckpt
*.pt
*.onnx
*.npz
*.npy

# Python
__pycache__/
*.py[cod]
*.egg-info/
.pytest_cache/
.ruff_cache/

# Divers
.DS_Store

{
  "fichier": "testset-v1.jsonl.gz",
  "sha256": "168d2de58d16a33607149cc91a7722e842c32750772b20f01fe30c7ba9a657c3",
  "graine": 20260820,
  "sel_de_partition": "curvy-test-v1",
  "n_total": 2134,
  "n_tenu_a_lecart": 1960,
  "n_hors_distribution": 174,
  "par_formule": {
    "gaussienne": 6,
    "runge": 6,
    "sinus_haute_frequence": 6,
    "sigmoide_raide": 6,
    "oscillation_amortie": 6,
    "sinus_de_inverse": 6,
    "logarithme_decale": 6,
    "racine_de_valeur_absolue": 6,
    "hyperbole_raide": 6,
    "sinus_cardinal": 6,
    "chirp": 6,
    "x_sinus_de_inverse": 6,
    "croissance_exponentielle": 6,
    "pointe_etroite": 6,
    "coude": 6,
    "arctangente": 6,
    "onde_triangulaire": 6,
    "bessel_j0": 6,
    "polynome_degre_7": 6,
    "deux_echelles": 6,
    "marche_douce_decalee": 6,
    "marche": 6,
    "plancher": 6,
    "par_morceaux": 6,
    "dent_de_scie": 6,
    "x_puissance_x": 6,
    "weierstrass_tronquee": 6,
    "erf": 6,
    "cloche_asymetrique": 6
  },
  "partition": {
    "sel": "curvy-test-v1",
    "n_entrainement": 253082,
    "n_test": 1998,
    "par_profondeur": {
      "3": {
        "entrainement": 1,
        "test": 0
      },
      "4": {
        "entrainement": 7,
        "test": 0
      },
      "5": {
        "entrainement": 194,
        "test": 48
      },
      "6": {
        "entrainement": 9840,
        "test": 400
      },
      "7": {
        "entrainement": 61759,
        "test": 600
      },
      "8": {
        "entrainement": 181281,
        "test": 950
      }
    },
    "profondeurs_sans_reserve": [
      3,
      4
    ]
  },
  "reel_canvas": "non construit — outil de capture en Phase 8"
}

"""Curvy — régression symbolique de petite taille.

Entrée : un nuage de points 2D. Sortie : un front de Pareto de formules
mathématiques candidates, de la plus simple à la plus précise.
"""

__version__ = "0.0.1"

"""``python -m curvy.cli_compile_check`` — ``torch.compile`` est-il utilisable ici ?

La spec demande de valider ``torch.compile`` sur un module jouet avant de
l'intégrer au training loop. Un module jouet à **forme fixe** ne prouve
pourtant pas grand-chose : notre encodeur reçoit des nuages de 20 à 200 points,
donc des formes **variables**. Le mode de défaillance réaliste n'est pas
« ça plante », c'est « ça recompile à chaque nouvelle taille » et l'entraînement
devient plus lent qu'en eager, silencieusement.

Ce script teste donc trois choses :
1. compile fonctionne et donne le même résultat qu'en eager (tolérance notée) ;
2. combien de recompilations une forme variable déclenche ;
3. le gain réel en débit, forme fixe, mesuré et non supposé.
"""

from __future__ import annotations

import time

import torch
from torch import nn

from curvy.devices import pick_device
from curvy.precision import configure_precision, precision_report
from curvy.seeding import seed_everything


class Toy(nn.Module):
    """Bloc jouet représentatif : projection + attention SDPA + MLP."""

    def __init__(self, d: int = 256, heads: int = 4) -> None:
        super().__init__()
        self.heads = heads
        self.qkv = nn.Linear(d, 3 * d)
        self.proj = nn.Linear(d, d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
        self.norm = nn.LayerNorm(d)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, n, d = x.shape
        q, k, v = self.qkv(self.norm(x)).chunk(3, dim=-1)
        shape = (b, n, self.heads, d // self.heads)
        q, k, v = (t.view(shape).transpose(1, 2) for t in (q, k, v))
        # SDPA : attention fusionnée de PyTorch, pas de réimplémentation manuelle.
        a = torch.nn.functional.scaled_dot_product_attention(q, k, v)
        x = x + self.proj(a.transpose(1, 2).reshape(b, n, d))
        return x + self.mlp(x)


def _bench(fn, x: torch.Tensor, device: torch.device, iters: int = 30) -> float:
    sync = torch.cuda.synchronize if device.type == "cuda" else (lambda: None)
    for _ in range(5):
        fn(x)
    sync()
    t0 = time.perf_counter()
    for _ in range(iters):
        fn(x)
    sync()
    return (time.perf_counter() - t0) / iters


def main() -> int:
    seed_everything(42)
    configure_precision(tf32=True)
    info = pick_device("auto")
    device = info.device
    print("device    :", info)
    print("précision :", precision_report(device))
    print()

    model = Toy().to(device).eval()
    x = torch.randn(32, 128, 256, device=device)

    with torch.no_grad():
        ref = model(x)

    print("--- 1. compile fonctionne-t-il, et donne-t-il le même résultat ? ---")
    t_compile0 = time.perf_counter()
    compiled = torch.compile(model)
    with torch.no_grad():
        got = compiled(x)
    warm = time.perf_counter() - t_compile0
    diff = (ref - got).abs().max().item()
    print(f"première passe (compilation incluse) : {warm:.2f} s")
    print(f"écart max eager vs compiled          : {diff:.3e}")
    print(f"verdict                              : {'OK' if diff < 1e-3 else 'ÉCART SUSPECT'}")
    print()

    print("--- 2. formes variables : combien de recompilations ? ---")
    counter = {"n": 0}
    try:
        import torch._dynamo as dynamo

        dynamo.reset()
        compiled_dyn = torch.compile(model)
        sizes = [20, 47, 96, 128, 200, 47, 128]
        t0 = time.perf_counter()
        with torch.no_grad():
            for n in sizes:
                compiled_dyn(torch.randn(8, n, 256, device=device))
        dt = time.perf_counter() - t0
        stats = dynamo.utils.counters.get("stats", {})
        counter["n"] = stats.get("unique_graphs", -1)
        print(f"tailles testées : {sizes}")
        print(f"graphes uniques compilés : {counter['n']}")
        print(f"temps total : {dt:.2f} s")
    except Exception as exc:  # noqa: BLE001
        print(f"instrumentation dynamo indisponible : {type(exc).__name__}: {exc}")
    print()

    print("--- 3. gain réel à forme fixe ---")
    with torch.no_grad():
        t_eager = _bench(model, x, device)
        t_comp = _bench(compiled, x, device)
    print(f"eager    : {t_eager * 1e3:.3f} ms/iter")
    print(f"compiled : {t_comp * 1e3:.3f} ms/iter")
    print(f"gain     : x{t_eager / t_comp:.2f}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

"""``python -m curvy.cli_env`` — carte d'identité de l'environnement de calcul.

Sert de livrable vérifiable pour la Phase 0 et d'en-tête pour toute fiche
d'expérience : on ne consigne jamais une métrique sans savoir sur quoi elle a
été mesurée.
"""

from __future__ import annotations

import json
import platform
import sys

from curvy.config import DATA_ROOT, REPO_ROOT
from curvy.devices import describe_backends, pick_device


def main() -> int:
    info = {
        "python": sys.version.split()[0],
        "platform": platform.platform(),
        "repo_root": str(REPO_ROOT),
        "data_root": str(DATA_ROOT),
        **describe_backends(),
    }
    print(json.dumps(info, indent=2, ensure_ascii=False))
    print()
    print("device retenu :", pick_device("auto"))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

"""``python -m curvy.cli_gpu`` — contrôle avant vol de la VRAM.

À lancer avant tout entraînement (le ``Makefile`` en fait une dépendance de la
cible ``train``). Répond à trois questions :

1. Les GPU autorisés par ``CURVY_CUDA_ALLOW`` ont-ils assez de VRAM libre ?
2. Qui l'occupe, le cas échéant ?
3. Faut-il tuer ces processus ? — uniquement avec ``--kill``, et **jamais** sur
   un GPU absent de l'allowlist. Le cas le plus fréquent est un run Curvy
   zombie d'une session tmux précédente ; le second est un serveur d'inférence
   qu'on a le droit d'arrêter.

Cet outil n'importe volontairement pas ``torch`` : il doit pouvoir dire
« la carte est pleine » sans avoir besoin d'y créer un contexte CUDA — la 3090
saturée nous a déjà montré que ce n'était pas toujours possible.
"""

from __future__ import annotations

import argparse
import os
import signal
import subprocess
import sys
import time

MIN_FREE_GIB_DEFAULT = 20.0


def _smi(query: str, entity: str = "gpu") -> list[list[str]]:
    """Interroge nvidia-smi et retourne les lignes découpées."""
    flag = "--query-gpu" if entity == "gpu" else "--query-compute-apps"
    out = subprocess.run(
        ["nvidia-smi", f"{flag}={query}", "--format=csv,noheader,nounits"],
        capture_output=True,
        text=True,
        check=True,
    ).stdout
    return [[c.strip() for c in line.split(",")] for line in out.splitlines() if line.strip()]


def _allowed() -> set[str] | None:
    raw = os.environ.get("CURVY_CUDA_ALLOW", "").strip()
    return {t.strip() for t in raw.split(",") if t.strip()} or None


def _is_allowed(index: str, uuid: str, allow: set[str] | None) -> bool:
    return allow is None or index in allow or uuid in allow


def _wait_gone(pid: int, timeout_s: float) -> bool:
    """Attend la disparition d'un PID, sans busy-wait agressif."""
    deadline = time.monotonic() + timeout_s
    while time.monotonic() < deadline:
        try:
            os.kill(pid, 0)
        except ProcessLookupError:
            return True
        except PermissionError:  # existe mais ne nous appartient pas
            return False
        time.sleep(0.2)
    return False


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument(
        "--min-free-gib",
        type=float,
        default=float(os.environ.get("CURVY_MIN_FREE_GIB", MIN_FREE_GIB_DEFAULT)),
        help="VRAM libre exigée sur au moins un GPU autorisé",
    )
    ap.add_argument(
        "--kill",
        action="store_true",
        help="tue les processus occupant les GPU AUTORISÉS (SIGTERM puis SIGKILL)",
    )
    args = ap.parse_args(argv)

    allow = _allowed()
    gpus = _smi("index,name,uuid,memory.total,memory.free")
    apps = _smi("pid,gpu_uuid,used_gpu_memory,process_name", entity="apps")

    by_uuid: dict[str, list[list[str]]] = {}
    for pid, uuid, used, name in apps:
        by_uuid.setdefault(uuid, []).append([pid, used, name])

    print(f"allowlist : {sorted(allow) if allow else '(aucune — tous les GPU)'}")
    print(f"seuil     : {args.min_free_gib:.1f} Gio libres exigés\n")

    best_free = 0.0
    killed: list[str] = []

    for index, name, uuid, total_mib, free_mib in gpus:
        ok = _is_allowed(index, uuid, allow)
        free_gib = float(free_mib) / 1024
        tag = "AUTORISÉ" if ok else "réservé — ne pas toucher"
        print(
            f"[{tag}] cuda:{index} {name}  {free_gib:.1f}/{float(total_mib) / 1024:.1f} Gio libres"
        )
        for pid, used, pname in by_uuid.get(uuid, []):
            print(f"           occupé par PID {pid} — {used} MiB — {pname}")
            if args.kill and ok:
                print(f"           -> SIGTERM {pid}")
                try:
                    os.kill(int(pid), signal.SIGTERM)
                except ProcessLookupError:
                    pass
                except PermissionError:
                    print(f"           -> PID {pid} n'appartient pas à l'utilisateur, ignoré")
                    continue
                if not _wait_gone(int(pid), 20.0):
                    print(f"           -> toujours vivant, SIGKILL {pid}")
                    try:
                        os.kill(int(pid), signal.SIGKILL)
                        _wait_gone(int(pid), 10.0)
                    except ProcessLookupError:
                        pass
                killed.append(f"{pid} ({pname})")
            elif args.kill and not ok:
                print("           -> GPU hors allowlist : PROCESSUS ÉPARGNÉ")
        if ok:
            best_free = max(best_free, free_gib)

    if killed:
        print(f"\nprocessus arrêtés : {', '.join(killed)}")
        best_free = max(
            (
                float(free) / 1024
                for i, _n, u, _t, free in _smi("index,name,uuid,memory.total,memory.free")
                if _is_allowed(i, u, allow)
            ),
            default=0.0,
        )
        print(f"VRAM libre après nettoyage : {best_free:.1f} Gio")

    if best_free < args.min_free_gib:
        print(
            f"\nÉCHEC : {best_free:.1f} Gio libres < {args.min_free_gib:.1f} Gio exigés.\n"
            f"Relancer avec `--kill` pour libérer les GPU autorisés.",
            file=sys.stderr,
        )
        return 1
    print(f"\nOK : {best_free:.1f} Gio libres sur un GPU autorisé.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

"""``python -m curvy.cli_smoke`` — validation du bout en bout avant d'entraîner.

Le but n'est pas d'apprendre quoi que ce soit mais de vérifier, en quelques
secondes, que la chaîne complète tient : dataset -> collate -> encodeur ->
décodeur -> loss -> backward. Et de **mesurer** la VRAM et le débit, plutôt que
de les estimer.
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import torch
from torch import nn
from torch.utils.data import DataLoader

from curvy.config import DATASET_DIR
from curvy.data.dataset import BucketedBatches, CurvyStream, collate
from curvy.devices import pick_device
from curvy.model.config import PRESETS
from curvy.model.curvy import CurvyModel, count_parameters
from curvy.precision import bf16_supported, configure_precision
from curvy.seeding import DEFAULT_SEED, seed_everything
from curvy.tokenizer.vocab import PAD_ID


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--preset", choices=sorted(PRESETS), default="small")
    ap.add_argument("--batch-size", type=int, default=256)
    ap.add_argument("--steps", type=int, default=30)
    ap.add_argument(
        "--warmup",
        type=int,
        default=10,
        help="steps exclus de la mesure — sinon on chronomètre surtout le "
        "remplissage initial du tampon de regroupement (10 s par worker)",
    )
    ap.add_argument("--workers", type=int, default=6)
    ap.add_argument("--seed", type=int, default=DEFAULT_SEED)
    ap.add_argument("--skeletons", type=Path, default=DATASET_DIR / "skeletons-v1.jsonl.gz")
    ap.add_argument(
        "--no-bucketing", action="store_true", help="désactive le regroupement par longueur"
    )
    args = ap.parse_args(argv)

    seed_everything(args.seed)
    configure_precision(tf32=True)
    info = pick_device("auto")
    device = info.device
    cfg = PRESETS[args.preset]

    model = CurvyModel(cfg).to(device)
    params = count_parameters(model)
    opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
    loss_fn = nn.CrossEntropyLoss(ignore_index=PAD_ID)

    stream = CurvyStream(args.skeletons, seed=args.seed)
    if args.no_bucketing:
        loader = DataLoader(
            stream,
            batch_size=args.batch_size,
            num_workers=args.workers,
            collate_fn=collate,
            pin_memory=device.type == "cuda",
            persistent_workers=args.workers > 0,
            prefetch_factor=4 if args.workers > 0 else None,
        )
    else:
        loader = DataLoader(
            BucketedBatches(stream, args.batch_size),
            batch_size=None,  # le dataset produit déjà des batches
            num_workers=args.workers,
            pin_memory=device.type == "cuda",
            persistent_workers=args.workers > 0,
            prefetch_factor=4 if args.workers > 0 else None,
        )

    use_bf16 = bf16_supported(device)
    autocast = (
        torch.autocast(device.type, dtype=torch.bfloat16)
        if use_bf16 and device.type == "cuda"
        else torch.autocast(device.type, enabled=False)
    )
    if device.type == "cuda":
        torch.cuda.reset_peak_memory_stats(device)

    it = iter(loader)
    losses: list[float] = []
    n_tokens = 0
    n_points = 0
    n_slots = 0
    t_data = 0.0
    t0 = time.perf_counter()

    for step in range(args.steps):
        td = time.perf_counter()
        batch = next(it).to(device)
        t_data += time.perf_counter() - td

        tokens_in = batch.tokens[:, :-1]
        target = batch.tokens[:, 1:]
        with autocast:
            logits = model(batch.points, batch.point_mask, tokens_in, batch.token_mask[:, :-1])
            loss = loss_fn(logits.reshape(-1, logits.size(-1)).float(), target.reshape(-1))
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()

        losses.append(loss.detach().item())
        if step < args.warmup:
            continue
        n_tokens += int((target != PAD_ID).sum())
        n_points += int((~batch.point_mask).sum())
        n_slots += int(batch.point_mask.numel())

    if device.type == "cuda":
        torch.cuda.synchronize()
    elapsed = time.perf_counter() - t0

    report = {
        "preset": args.preset,
        "device": str(info),
        "parametres": params["TOTAL"],
        "bf16": use_bf16,
        "batch_size": args.batch_size,
        "bucketing": not args.no_bucketing,
        "remplissage_inutile": None,
        "steps": args.steps,
        "warmup": args.warmup,
        "duree_s": round(elapsed, 2),
        "s_par_step": round(elapsed / args.steps, 4),
        "part_attente_donnees": f"{100 * t_data / elapsed:.1f} %",
        "tokens_par_s": round(n_tokens / elapsed),
        "points_par_s": round(n_points / elapsed),
        "loss_initiale": round(losses[0], 4),
        "loss_finale": round(losses[-1], 4),
    }
    report["remplissage_inutile"] = f"{100 * (1 - n_points / n_slots):.1f} %"
    if device.type == "cuda":
        report["vram_pic_Mio"] = round(torch.cuda.max_memory_allocated(device) / 2**20)
        report["vram_reservee_Mio"] = round(torch.cuda.max_memory_reserved(device) / 2**20)
    print(json.dumps(report, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

"""Chemins et constantes globales du projet.

Tout ce qui est volumineux — venv, jeux de données, points de reprise —
vit hors du dépôt, sous ``CURVY_DATA_ROOT``. Le défaut convient à la machine
d'origine ; ailleurs, il faut le définir.
"""

from __future__ import annotations

import os
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent

#: Racine des artefacts lourds. Surchargeable par la variable d'environnement
#: ``CURVY_DATA_ROOT`` pour rendre le projet portable.
DATA_ROOT = Path(os.environ.get("CURVY_DATA_ROOT", "~/curvy-data")).expanduser()

DATASET_DIR = DATA_ROOT / "data"
RUNS_DIR = DATA_ROOT / "runs"
CACHE_DIR = DATA_ROOT / "cache"

DOCS_DIR = REPO_ROOT / "docs"
FIGURES_DIR = DOCS_DIR / "article" / "figures"
LOGS_DIR = REPO_ROOT / "logs"


def ensure_dirs() -> None:
    """Crée les répertoires d'artefacts s'ils manquent."""
    for d in (DATASET_DIR, RUNS_DIR, CACHE_DIR, FIGURES_DIR, LOGS_DIR / "raw"):
        d.mkdir(parents=True, exist_ok=True)

"""Sélection du device de calcul.

Le projet doit tourner à l'identique sur trois cibles :
  - serveur Linux multi-GPU CUDA (machine de développement réelle) ;
  - MacBook Apple Silicon (backend MPS, cible annoncée dans la spec) ;
  - CPU (CI, machines sans accélérateur).

Piège rencontré le 2026-08-19 : par défaut le runtime CUDA numérote les GPU
par ordre de puissance décroissante (``CUDA_DEVICE_ORDER=FASTEST_FIRST``), pas
par ordre de bus PCI. Sur cette machine ``cuda:0`` désignait donc la RTX 3090
— entièrement occupée par un autre processus — et non la RTX 3060 libre.
``pick_device`` choisit sur la mémoire libre réelle, ce qui rend la question

Politique d'allocation : sur cette machine, la RTX 3090 est allouée
au projet et la RTX 3060 est réservée à d'autres travaux. La contrainte est
appliquée mécaniquement par la variable d'environnement ``CURVY_CUDA_ALLOW``
(liste d'UUID ou d'index CUDA), positionnée par le ``Makefile`` — et non par
la discipline de l'opérateur.
"""

from __future__ import annotations

import os
from dataclasses import dataclass

import torch

__all__ = ["DeviceInfo", "pick_device", "describe_backends"]


@dataclass(frozen=True)
class DeviceInfo:
    """Device retenu et pourquoi."""

    device: torch.device
    name: str
    free_bytes: int | None
    reason: str

    def __str__(self) -> str:
        free = "n/a" if self.free_bytes is None else f"{self.free_bytes / 2**30:.1f} Gio libres"
        return f"{self.device} ({self.name}, {free}) — {self.reason}"


def _normalise_uuid(value: str) -> str:
    """Ramène un UUID de GPU à une forme comparable.

    Piège (2026-08-19) : ``nvidia-smi`` écrit ``GPU-1234abcd-...`` tandis que
    ``torch.cuda.get_device_properties(i).uuid`` écrit ``1234abcd-...``, sans
    le préfixe. Comparer les deux tels quels ne matche jamais — et l'allowlist
    se vidait donc silencieusement.
    """
    return value.strip().lower().removeprefix("gpu-")


def _allowlist() -> set[str] | None:
    """Ensemble des GPU autorisés, ou ``None`` si aucune restriction.

    Lu depuis ``CURVY_CUDA_ALLOW`` : liste séparée par des virgules d'index
    CUDA (``0``, ``1``) et/ou d'UUID, avec ou sans le préfixe ``GPU-``. Les
    UUID sont préférables : ils survivent aux renumérotations.
    """
    raw = os.environ.get("CURVY_CUDA_ALLOW", "").strip()
    if not raw:
        return None
    return {_normalise_uuid(tok) for tok in raw.split(",") if tok.strip()}


def _cuda_candidates() -> list[tuple[int, str, int]]:
    """(index, nom, octets libres) pour chaque GPU CUDA réellement interrogeable.

    Un GPU saturé échoue à l'initialisation de son contexte : on l'écarte
    silencieusement plutôt que de faire tomber tout le programme.

    Subtilité coûteuse (bug du 2026-08-19) : ``mem_get_info(i)`` bascule le
    device courant sur ``i`` *avant* d'échouer, et ne le restaure pas. Une
    sonde ratée laissait donc ``current_device()`` sur le GPU mort, et le
    premier ``torch.cuda.synchronize()`` sans argument explosait très loin de
    la cause. On restaure explicitement le device courant après chaque sonde.
    """
    allow = _allowlist()
    out: list[tuple[int, str, int]] = []
    for i in range(torch.cuda.device_count()):
        props = torch.cuda.get_device_properties(i)
        name = props.name
        if (
            allow is not None
            and str(i) not in allow
            and _normalise_uuid(str(props.uuid)) not in allow
        ):
            continue  # GPU réservé à d'autres travaux : on ne le sonde même pas
        try:
            free, _total = torch.cuda.mem_get_info(i)
        except Exception:  # contexte CUDA impossible à créer -> GPU inutilisable
            continue
        out.append((i, name, free))
    if out:
        # Ramène le device courant sur un GPU sain, quoi qu'aient fait les sondes.
        torch.cuda.set_device(out[0][0])
    return out


def pick_device(preference: str = "auto", min_free_gib: float = 2.0) -> DeviceInfo:
    """Retourne le device à utiliser.

    ``preference`` vaut ``auto``, ``cuda``, ``mps``, ``cpu`` ou un identifiant
    explicite du type ``cuda:1``. En mode ``auto`` on prend le GPU CUDA offrant
    le plus de mémoire libre, à défaut MPS, à défaut le CPU.
    """
    if preference not in ("auto", "cuda", "mps", "cpu"):
        dev = torch.device(preference)
        return DeviceInfo(dev, preference, None, "imposé explicitement par l'appelant")

    if preference in ("auto", "cuda") and torch.cuda.is_available():
        usable = [c for c in _cuda_candidates() if c[2] >= min_free_gib * 2**30]
        if usable:
            idx, name, free = max(usable, key=lambda c: c[2])
            torch.cuda.set_device(idx)  # sinon les appels sans device explicite visent cuda:0
            return DeviceInfo(
                torch.device(f"cuda:{idx}"),
                name,
                free,
                f"GPU CUDA le plus libre parmi {torch.cuda.device_count()} détecté(s)",
            )
        if preference == "cuda" or _allowlist() is not None:
            # Repli silencieux sur CPU interdit : une allowlist mal orthographiée
            # nous a déjà fait « choisir » le CPU sans le dire. Mieux vaut casser.
            raise RuntimeError(
                f"Aucun GPU autorisé n'a {min_free_gib} Gio libres. "
                f"CURVY_CUDA_ALLOW={os.environ.get('CURVY_CUDA_ALLOW', '(non défini)')} ; "
                f"GPU visibles et sondables : {_cuda_candidates()}. "
                f"Forcer le CPU avec preference='cpu' si c'est voulu."
            )

    if preference in ("auto", "mps") and torch.backends.mps.is_available():
        return DeviceInfo(
            torch.device("mps"), "Apple Silicon (MPS)", None, "backend MPS disponible"
        )

    if preference == "mps":
        raise RuntimeError("MPS demandé mais indisponible (machine non-Apple ou torch sans MPS).")

    return DeviceInfo(torch.device("cpu"), "CPU", None, "aucun accélérateur exploitable")


def describe_backends() -> dict[str, object]:
    """Instantané des backends, pour les fiches d'expérience."""
    info: dict[str, object] = {
        "torch": torch.__version__,
        "cuda_build": torch.version.cuda,
        "cuda_available": torch.cuda.is_available(),
        "cuda_device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0,
        "mps_built": torch.backends.mps.is_built(),
        "mps_available": torch.backends.mps.is_available(),
        "CUDA_DEVICE_ORDER": os.environ.get("CUDA_DEVICE_ORDER", "(non défini -> FASTEST_FIRST)"),
        "CUDA_VISIBLE_DEVICES": os.environ.get("CUDA_VISIBLE_DEVICES", "(non défini)"),
        "CURVY_CUDA_ALLOW": os.environ.get("CURVY_CUDA_ALLOW", "(non défini -> tous les GPU)"),
    }
    if torch.cuda.is_available():
        info["cuda_devices"] = [
            {"index": i, "name": n, "free_gib": round(f / 2**30, 2)}
            for i, n, f in _cuda_candidates()
        ]
    return info

"""Précision numérique : bf16 et TF32.

Ampere (RTX 3090, sm_86) supporte nativement le bf16 et le TF32 :

- **bf16** a la même plage d'exposant que le fp32, donc pas d'``inf`` sur les
  logits et **aucun ``GradScaler``** — contrairement au fp16. C'est le bon
  défaut pour l'entraînement.
- **TF32** est une troncature de la mantisse appliquée aux matmuls et aux
  convolutions en fp32. On y gagne beaucoup de débit pour une perte de
  précision sans effet à notre échelle. PyTorch le désactive par défaut depuis
  la 1.12 : il faut l'activer explicitement.
"""

from __future__ import annotations

import torch

__all__ = ["bf16_supported", "configure_precision", "precision_report"]


def bf16_supported(device: torch.device) -> bool:
    """Le device sait-il faire du bf16 sans émulation ?"""
    if device.type == "cuda":
        return torch.cuda.is_bf16_supported()
    if device.type == "cpu":
        return hasattr(torch, "bfloat16")
    return False  # MPS : bf16 partiel selon les versions, on ne parie pas dessus


def configure_precision(tf32: bool = True) -> None:
    """Active TF32 pour les matmuls et cuDNN. Sans effet hors CUDA."""
    if not torch.cuda.is_available():
        return
    torch.backends.cuda.matmul.allow_tf32 = tf32
    torch.backends.cudnn.allow_tf32 = tf32
    # API récente, plus explicite que le booléen historique.
    if hasattr(torch, "set_float32_matmul_precision"):
        torch.set_float32_matmul_precision("high" if tf32 else "highest")


def precision_report(device: torch.device) -> dict[str, object]:
    """Instantané, pour l'en-tête des fiches d'expérience."""
    rep: dict[str, object] = {
        "device": str(device),
        "bf16_supported": bf16_supported(device),
        "float32_matmul_precision": torch.get_float32_matmul_precision(),
    }
    if torch.cuda.is_available():
        rep["tf32_matmul"] = torch.backends.cuda.matmul.allow_tf32
        rep["tf32_cudnn"] = torch.backends.cudnn.allow_tf32
    if device.type == "cuda":  # et non `cuda.is_available()` : le device peut être CPU
        props = torch.cuda.get_device_properties(device)
        rep["compute_capability"] = f"sm_{props.major}{props.minor}"
        rep["device_name"] = props.name
    return rep

"""Reproductibilité : une seule fonction pour tout semer.

Règle projet n°5 : toute expérience doit être rejouable. La graine est
enregistrée dans chaque fiche d'expérience et dans chaque checkpoint.

Arbitrage numpy (2026-08-19) : ruff (NPY002) réclame l'abandon de
``np.random.seed`` au profit d'un ``Generator`` explicite, et il a raison pour
*notre* code. Mais l'état global hérité reste le seul levier sur les
bibliothèques tierces qui l'utilisent en interne (scipy en particulier, qu'on
emploiera pour l'ajustement des constantes en Phase 5). On fait donc les deux :
on sème l'état global *et* on expose un ``Generator`` que le code de Curvy doit
"""

from __future__ import annotations

import os
import random

import numpy as np
import torch

DEFAULT_SEED = 42

__all__ = ["DEFAULT_SEED", "make_rng", "seed_everything", "worker_seed_fn"]


def make_rng(seed: int = DEFAULT_SEED) -> np.random.Generator:
    """Générateur numpy moderne — à utiliser dans tout le code de Curvy."""
    return np.random.default_rng(seed)


def seed_everything(seed: int = DEFAULT_SEED, deterministic: bool = False) -> np.random.Generator:
    """Sème python, numpy (état global hérité) et torch (CPU + tous les GPU).

    ``deterministic=True`` force les noyaux cuDNN déterministes : plus lent,
    réservé aux runs de vérification bit-à-bit.

    Retourne le ``Generator`` numpy à utiliser par l'appelant.
    """
    os.environ["PYTHONHASHSEED"] = str(seed)
    random.seed(seed)
    # Volontaire : sème l'état global pour les bibliothèques tierces (scipy...).
    np.random.seed(seed)  # noqa: NPY002
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
    return make_rng(seed)


def worker_seed_fn(base_seed: int):
    """Fabrique un ``worker_init_fn`` pour DataLoader.

    Sans ça, chaque worker hérite du même état RNG et génère exactement les
    mêmes exemples — bug classique et silencieux des pipelines synthétiques.
    """

    def _init(worker_id: int) -> None:
        s = base_seed + worker_id
        random.seed(s)
        np.random.seed(s % (2**32))  # noqa: NPY002
        torch.manual_seed(s)

    return _init

"""Sous-paquet curvy.data — à remplir."""

"""Canonicalisation des squelettes.

Objectif : que deux écritures d'une même famille de fonctions donnent le même
arbre, pour ne pas apprendre plusieurs fois la même chose.

Pourquoi pas sympy — la raison est importante et vaut d'être répétée : dans
notre algèbre, deux ``C`` du même squelette sont **deux paramètres
indépendants**. sympy, qui les voit comme un même symbole, « simplifierait »
``C*x + C*x`` en ``2*C*x``, ce qui est faux ici. Il faut donc nos propres
règles. Elles ont aussi l'avantage d'être des dizaines de fois plus rapides,
ce qui compte : la canonicalisation tourne sur des millions d'arbres.

Règles appliquées de bas en haut jusqu'au point fixe :

1. tout sous-arbre sans ``x`` se replie sur ``C`` (couvre ``unaire(C) → C`` et
   ``C op C → C``) ;
2. ``sub(u, u) → C`` (nul, donc constant) ;
3. involutions et annulations : ``inv(inv u) → u``, ``exp(log u) → u``,
   ``log(exp u) → u``, ``abs(abs u) → abs u``, ``abs(sq u) → sq u``,
   ``sq(abs u) → sq u``, ``sqrt(sq u) → abs u`` ;
4. ``sub(u, C) → add(u, C)`` — ``C`` est libre et de signe quelconque ;
5. dans une chaîne de ``mul`` (resp. ``add``), au plus un facteur (resp. terme)
   constant : ``mul(C, mul(C, u)) → mul(C, u)`` ;
6. termes semblables d'une somme fusionnés **si l'un au moins porte un facteur
   constant libre** : ``C*u + C*u → C*u`` et ``C*u + u → C*u``, puisque la
   somme de deux constantes libres est une constante libre. Attention à la
   symétrique fausse : ``u + u`` vaut ``2u``, qui est une fonction **fixe** et
   non une famille — le fusionner en ``C*u`` ajouterait un paramètre libre et
   changerait la classe de fonctions. ``sin(x + x)`` est ``sin(2x)``, ce n'est
   pas ``sin(C*x)`` ;
7. arguments des opérateurs commutatifs triés selon un ordre total.

L'enveloppe de racine ``C * (…) + C`` est posée en dernier, après
avoir retiré du sommet du corps tout ce qu'elle absorbe déjà.
"""

from __future__ import annotations

from curvy.data.expr import Node, has_x, size, to_prefix

__all__ = ["canonicalise", "canonical_key", "strip_absorbable_root", "wrap_root"]

C: Node = ("C",)

_INVOLUTIONS = {
    ("inv", "inv"): lambda inner: inner,
    ("exp", "log"): lambda inner: inner,
    ("log", "exp"): lambda inner: inner,
    ("abs", "abs"): lambda inner: ("abs", inner),
    ("abs", "sq"): lambda inner: ("sq", inner),
    ("sq", "abs"): lambda inner: ("sq", inner),
    ("sqrt", "sq"): lambda inner: ("abs", inner),
}


def _sort_key(node: Node) -> tuple:
    """Ordre total sur les sous-arbres : d'abord la taille, puis la séquence."""
    return (size(node), tuple(to_prefix(node)))


def _flatten(node: Node, op: str) -> list[Node]:
    """Aplatit une chaîne d'opérateurs associatifs identiques."""
    out: list[Node] = []
    stack = [node]
    while stack:
        cur = stack.pop()
        if cur[0] == op:
            stack.extend(cur[1:])
        else:
            out.append(cur)
    return out


def _rebuild(op: str, parts: list[Node]) -> Node:
    node = parts[0]
    for p in parts[1:]:
        node = (op, node, p)
    return node


def _strip_free_scale(term: Node) -> tuple[bool, Node]:
    """Sépare un éventuel facteur constant libre du reste du terme.

    ``mul(C, sin(x))`` -> ``(True, sin(x))`` ; ``sin(x)`` -> ``(False, sin(x))``.
    """
    if term[0] != "mul":
        return False, term
    factors = _flatten(term, "mul")
    if C not in factors:
        return False, term
    rest = [f for f in factors if f != C]
    if not rest:
        return True, C
    return True, _rebuild("mul", sorted(rest, key=_sort_key))


def _merge_like_terms(parts: list[Node]) -> list[Node]:
    """Fusionne les termes semblables d'une somme (règle 6)."""
    groups: dict[Node, list[bool]] = {}
    for term in parts:
        scaled, base = _strip_free_scale(term)
        groups.setdefault(base, []).append(scaled)
    out: list[Node] = []
    for base, scales in groups.items():
        if len(scales) == 1:
            out.append(("mul", C, base) if scales[0] else base)
        elif any(scales):
            # Au moins une constante libre dans le groupe : elle absorbe tout.
            out.append(("mul", C, base))
        else:
            # Aucune constante libre : `u + u` vaut `2u`, non fusionnable.
            out.extend([base] * len(scales))
    return out


def canonicalise(node: Node) -> Node:
    """Forme canonique d'un sous-arbre. Idempotente."""
    tok = node[0]
    if tok in ("x", "C"):
        return node

    children = tuple(canonicalise(c) for c in node[1:])
    node = (tok, *children)

    # 1. Tout sous-arbre sans x est une constante, quelle que soit sa forme.
    if not has_x(node):
        return C

    # 3. Involutions et annulations.
    if len(children) == 1:
        inner = children[0]
        rule = _INVOLUTIONS.get((tok, inner[0]))
        if rule is not None:
            return canonicalise(rule(inner[1]))

    if tok == "sub":
        left, right = children
        # 2. u - u = 0
        if left == right:
            return C
        # 4. u - C ≡ u + C (C est libre et de signe quelconque)
        if right == C:
            return canonicalise(("add", left, C))

    if tok in ("add", "mul"):
        parts = _flatten(node, tok)
        # 5. Au plus un élément constant dans la chaîne.
        consts = [p for p in parts if p == C]
        parts = [p for p in parts if p != C]
        if consts:
            parts.append(C)
        if tok == "add":
            # 6. Termes semblables (hors le terme constant isolé, déjà traité).
            merged = _merge_like_terms([p for p in parts if p != C])
            parts = merged + ([C] if C in parts else [])
        if len(parts) == 1:
            return parts[0]
        # 7. Ordre total sur les arguments commutatifs.
        parts.sort(key=_sort_key)
        return _rebuild(tok, parts)

    return node


def strip_absorbable_root(body: Node) -> Node:
    """Retire du sommet du corps ce que l'enveloppe de racine absorbe déjà.

    L'enveloppe ``C * (…) + C`` rend redondants, **au sommet du corps
    uniquement**, un facteur constant, un terme constant, et même un
    ``sub(C, u)`` — le signe étant absorbé par le ``C`` multiplicatif.
    Plus bas dans l'arbre ces formes sont significatives et sont conservées.
    """
    while True:
        tok = body[0]
        if tok in ("add", "mul") and C in body[1:]:
            body = body[1] if body[2] == C else body[2]
            continue
        if tok == "sub" and (body[1] == C or body[2] == C):
            body = body[2] if body[1] == C else body[1]
            continue
        return body


def wrap_root(body: Node) -> Node:
    """Enveloppe canonique ``C * body + C``."""
    return ("add", ("mul", C, body), C)


def canonical_key(node: Node) -> str:
    """Clé de déduplication exacte."""
    return " ".join(to_prefix(node))

"""Flux d'entraînement : squelettes stratifiés + nuages tirés en ligne.

Rien n'est stocké sur disque à part la liste des squelettes. Chaque exemple est
fabriqué à la volée, si bien qu'un même squelette ne produit jamais deux fois
le même nuage — il n'y a donc pas de notion d'epoch, et pas de sur-apprentissage
au sens habituel.
"""

from __future__ import annotations

from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import IterableDataset, get_worker_info

from curvy.data.expr import from_prefix
from curvy.data.generate import load_skeletons
from curvy.data.pointcloud import CloudConfig, sample_cloud
from curvy.data.weighting import DEFAULT_DEPTH_TARGET, stratified_weights
from curvy.tokenizer.vocab import MAX_SEQ_LEN, PAD_ID, encode

__all__ = [
    "Batch",
    "BucketedBatches",
    "CurvyStream",
    "ValidationExample",
    "collate",
    "make_validation_set",
]


@dataclass
class Batch:
    points: torch.Tensor  # (B, N, 2) float32
    point_mask: torch.Tensor  # (B, N) bool — True = position de remplissage
    tokens: torch.Tensor  # (B, L) int64, <bos> … <eos> puis <pad>
    token_mask: torch.Tensor  # (B, L) bool — True = <pad>

    def to(self, device: torch.device) -> Batch:
        return Batch(
            self.points.to(device, non_blocking=True),
            self.point_mask.to(device, non_blocking=True),
            self.tokens.to(device, non_blocking=True),
            self.token_mask.to(device, non_blocking=True),
        )


class CurvyStream(IterableDataset):
    """Flux infini d'exemples (nuage, squelette encodé)."""

    def __init__(
        self,
        skeleton_path: Path,
        seed: int,
        cloud_cfg: CloudConfig | None = None,
        depth_target: dict[int, float] | None = None,
        max_seq_len: int = MAX_SEQ_LEN,
        max_retries: int = 12,
        exclure: frozenset[str] | None = None,
        garder: frozenset[str] | None = None,
    ) -> None:
        items = load_skeletons(skeleton_path)
        # `exclure` porte les squelettes du jeu de test (Phase 6) :
        # sans ce filtre, le jeu de test mesure de la restitution. `garder` fait
        # l'inverse et sert à construire le jeu de test lui-même.
        if exclure and garder:
            raise ValueError("`exclure` et `garder` sont exclusifs l'un de l'autre")
        self.n_exclus = 0
        if exclure:
            avant = len(items)
            items = [it for it in items if it["prefix"] not in exclure]
            self.n_exclus = avant - len(items)
        elif garder:
            avant = len(items)
            items = [it for it in items if it["prefix"] in garder]
            self.n_exclus = avant - len(items)
            if not items:
                raise ValueError("`garder` ne retient aucun squelette")
        self.prefixes = [it["prefix"] for it in items]
        self.depths = [it["depth"] for it in items]
        self.weights = stratified_weights(self.depths, depth_target or DEFAULT_DEPTH_TARGET)
        self.seed = seed
        self.cloud_cfg = cloud_cfg or CloudConfig()
        self.max_seq_len = max_seq_len
        self.max_retries = max_retries
        self._trees: list | None = None

    def _tree(self, i: int):
        if self._trees is None:
            self._trees = [None] * len(self.prefixes)
        if self._trees[i] is None:
            self._trees[i] = from_prefix(self.prefixes[i].split())
        return self._trees[i]

    def __iter__(self) -> Iterator[tuple[np.ndarray, list[int]]]:
        info = get_worker_info()
        wid = 0 if info is None else info.id
        rng = np.random.default_rng(self.seed + 7919 * wid)
        idx_pool = np.arange(len(self.prefixes))
        while True:
            i = int(rng.choice(idx_pool, p=self.weights))
            node = self._tree(i)
            for _ in range(self.max_retries):
                cloud, _ = sample_cloud(rng, node, self.cloud_cfg)
                if cloud is not None:
                    break
            else:
                continue  # squelette récalcitrant : on passe au suivant
            ids = encode(node)
            if len(ids) > self.max_seq_len:
                continue
            pts = np.stack([cloud.x, cloud.y], axis=1).astype(np.float32)
            yield pts, ids


def collate(samples: list[tuple[np.ndarray, list[int]]]) -> Batch:
    """Remplissage à la volée. Les masques valent True sur le remplissage."""
    n_max = max(len(p) for p, _ in samples)
    l_max = max(len(t) for _, t in samples)
    b = len(samples)

    points = torch.zeros(b, n_max, 2, dtype=torch.float32)
    point_mask = torch.ones(b, n_max, dtype=torch.bool)
    tokens = torch.full((b, l_max), PAD_ID, dtype=torch.long)
    token_mask = torch.ones(b, l_max, dtype=torch.bool)

    for i, (pts, ids) in enumerate(samples):
        points[i, : len(pts)] = torch.from_numpy(pts)
        point_mask[i, : len(pts)] = False
        tokens[i, : len(ids)] = torch.tensor(ids, dtype=torch.long)
        token_mask[i, : len(ids)] = False
    return Batch(points, point_mask, tokens, token_mask)


@dataclass
class ValidationExample:
    """Un exemple de validation, avec la **vérité terrain** nécessaire au R².

    ``y`` est ce que le modèle voit (bruité, normalisé) ; ``y_clean`` est la
    valeur exacte de la fonction aux mêmes abscisses, également normalisée.
    C'est contre ``y_clean`` que le R² se mesure (précision du
    2026-08-19) — mesurer contre ``y`` plafonnerait au niveau de bruit qu'on a
    soi-même injecté.
    """

    points: np.ndarray
    ids: list[int]
    node: object
    x: np.ndarray
    y: np.ndarray
    y_clean: np.ndarray
    depth: int


def make_validation_set(
    skeleton_path: Path,
    n: int,
    seed: int,
    cloud_cfg: CloudConfig | None = None,
    garder: frozenset[str] | None = None,
    un_nuage_par_squelette: bool = False,
) -> list[ValidationExample]:
    """Jeu de validation **figé** : mêmes exemples à chaque run, donc courbes
    comparables entre expériences.

    ``garder`` restreint le tirage à un sous-ensemble de squelettes — c'est
    ainsi qu'on bâtit le jeu de test de la Phase 6, à partir des seuls
    squelettes tenus à l'écart de l'entraînement.

    ``un_nuage_par_squelette`` parcourt les squelettes au lieu de les tirer :
    chaque squelette apparaît exactement une fois. Le tirage stratifié
    surreprésenterait sinon les strates profondes au sein d'une réserve déjà
    construite pour être équilibrée.

    Dans ce mode, chaque squelette a droit à ``max_retries`` tentatives de
    nuage, comme dans le flux d'entraînement. Une seule tentative perdrait
    **30 % des squelettes** (mesuré : 349/500 contre 492/500 à douze essais) —
    et pas les plus difficiles, seulement les moins chanceux du premier coup.
    Le jeu de test serait alors biaisé vers les squelettes commodes.
    """
    from curvy.data.expr import depth as tree_depth
    from curvy.data.expr import evaluate

    stream = CurvyStream(skeleton_path, seed=seed, cloud_cfg=cloud_cfg, garder=garder)
    rng = np.random.default_rng(seed)
    items = stream.prefixes
    weights = stream.weights
    out: list[ValidationExample] = []
    ordre = rng.permutation(len(items)) if un_nuage_par_squelette else None
    curseur = 0
    while len(out) < n:
        if ordre is not None:
            if curseur >= len(ordre):
                break  # réserve épuisée : on rend moins que `n`, et ça se voit
            i = int(ordre[curseur])
            curseur += 1
        else:
            i = int(rng.choice(len(items), p=weights))
        node = stream._tree(i)
        essais = stream.max_retries if un_nuage_par_squelette else 1
        cloud = None
        for _ in range(essais):
            cloud, _ = sample_cloud(rng, node, stream.cloud_cfg)
            if cloud is not None:
                break
        if cloud is None:
            continue
        ids = encode(node)
        if len(ids) > stream.max_seq_len:
            continue
        with np.errstate(over="ignore", invalid="ignore", divide="ignore"):
            clean = (evaluate(node, cloud.x, cloud.consts) - cloud.y_offset) / cloud.y_scale
        if not np.isfinite(clean).all():
            continue
        pts = np.stack([cloud.x, cloud.y], axis=1).astype(np.float32)
        out.append(ValidationExample(pts, ids, node, cloud.x, cloud.y, clean, tree_depth(node)))
    return out


class BucketedBatches(IterableDataset):
    """Regroupe les exemples de tailles voisines avant de former les batches.

    Le nombre de points est tiré uniformément dans [20, 200] : en remplissant
    chaque batch jusqu'à son maximum, **45,1 % des emplacements alloués sont du
    remplissage** (mesuré le 2026-08-19). Comme l'attention de l'encodeur coûte
    O(B·N²), on paie presque le double du nécessaire.

    **Tampon glissant, et non par blocs.** La première version tamponnait
    ``16 × batch_size`` exemples, les triait, puis livrait tous les batches d'un
    coup. Résultat mesuré : 73,7 % du temps passé à attendre les données, et un
    entraînement 2,3 fois plus lent qu'*avec* le remplissage inutile. Le
    générateur produit 404 exemples/s par cœur ; un tampon de 8192 met 20 s à se
    remplir, pendant lesquelles le GPU ne fait rien.

    Ici le tampon est rempli une seule fois, puis maintenu : chaque batch en
    retire ``batch_size`` exemples de longueurs voisines et en réinjecte autant.
    La consommation devient régulière, la rafale disparaît.

    Le décalage de départ est tiré au hasard à chaque batch : sans lui, le
    modèle verrait toujours les nuages les plus courts en premier, ce qui
    corrélerait la taille effective du batch au pas d'entraînement.
    """

    def __init__(self, stream: CurvyStream, batch_size: int, pool_factor: int = 8) -> None:
        self.stream = stream
        self.batch_size = batch_size
        self.pool_factor = pool_factor

    def __iter__(self) -> Iterator[Batch]:
        info = get_worker_info()
        wid = 0 if info is None else info.id
        rng = np.random.default_rng(self.stream.seed + 104729 * wid)
        src = iter(self.stream)
        pool_size = self.batch_size * self.pool_factor
        pool = [next(src) for _ in range(pool_size)]
        while True:
            pool.sort(key=lambda s: len(s[0]))
            start = int(rng.integers(0, len(pool) - self.batch_size + 1))
            batch = pool[start : start + self.batch_size]
            del pool[start : start + self.batch_size]
            pool.extend(next(src) for _ in range(self.batch_size))
            yield collate(batch)

"""Arbres d'expression, notation préfixe, évaluation.

Un nœud est un tuple ``(token, *enfants)`` : ``("add", ("mul", ("C",), ("x",)),
("C",))``. C'est immuable, hachable, et se compare directement — trois
propriétés dont la canonicalisation et la déduplication se servent beaucoup.

La notation préfixe rend toute séquence vérifiable par un simple
compteur d'arité, ce qui sert deux fois : valider les cibles à la génération, et
masquer les tokens impossibles pendant le beam search.
"""

from __future__ import annotations

from collections.abc import Iterator, Sequence

import numpy as np

from curvy.data.grammar import (
    ARITY,
    COMPLEXITY_COST,
    INFIX_SYMBOL,
    UNARY_RENDER,
)

Node = tuple  # ("token", *enfants)

__all__ = [
    "Node",
    "complexity",
    "const_name_iter",
    "count_constants",
    "depth",
    "evaluate",
    "from_prefix",
    "has_x",
    "iter_nodes",
    "prefix_is_complete",
    "prefix_validity",
    "size",
    "to_infix",
    "to_prefix",
]


# --- structure ---------------------------------------------------------------


def to_prefix(node: Node) -> list[str]:
    out: list[str] = []
    stack = [node]
    while stack:
        cur = stack.pop()
        out.append(cur[0])
        stack.extend(reversed(cur[1:]))
    return out


def from_prefix(tokens: Sequence[str]) -> Node:
    """Reconstruit l'arbre. Lève ``ValueError`` si la séquence est mal formée."""
    pos = 0

    def build() -> Node:
        nonlocal pos
        if pos >= len(tokens):
            raise ValueError("séquence préfixe incomplète")
        tok = tokens[pos]
        pos += 1
        if tok not in ARITY:
            raise ValueError(f"token inconnu : {tok!r}")
        children = tuple(build() for _ in range(ARITY[tok]))
        return (tok, *children)

    root = build()
    if pos != len(tokens):
        raise ValueError(f"{len(tokens) - pos} token(s) en trop après l'arbre")
    return root


def prefix_validity(tokens: Sequence[str]) -> int | None:
    """Nombre de sous-arbres encore attendus, ou ``None`` si la séquence est morte.

    ``0`` signifie « arbre complet ». C'est exactement la quantité dont le beam
    search a besoin pour savoir quels tokens sont légaux à l'étape suivante.
    """
    remaining = 1
    for tok in tokens:
        if tok not in ARITY or remaining == 0:
            return None
        remaining += ARITY[tok] - 1
    return remaining


def prefix_is_complete(tokens: Sequence[str]) -> bool:
    return prefix_validity(tokens) == 0


def iter_nodes(node: Node) -> Iterator[Node]:
    stack = [node]
    while stack:
        cur = stack.pop()
        yield cur
        stack.extend(cur[1:])


def size(node: Node) -> int:
    return sum(1 for _ in iter_nodes(node))


def depth(node: Node) -> int:
    return 1 + max((depth(c) for c in node[1:]), default=0)


def count_constants(node: Node) -> int:
    return sum(1 for n in iter_nodes(node) if n[0] == "C")


def has_x(node: Node) -> bool:
    return any(n[0] == "x" for n in iter_nodes(node))


def complexity(node: Node) -> int:
    """Coût pondéré, pour l'axe « simplicité » du front de Pareto."""
    return sum(COMPLEXITY_COST[n[0]] for n in iter_nodes(node))


# --- rendu -------------------------------------------------------------------


def to_infix(node: Node, const_names: Iterator[str] | None = None) -> str:
    """Rendu infixe parsable par sympy.

    ``const_names`` fournit un nom distinct par occurrence de ``C`` (``c0``,
    ``c1``, …) : indispensable, puisque deux ``C`` du même squelette sont deux
    paramètres indépendants et non la même valeur.
    """
    tok = node[0]
    if tok == "x":
        return "x"
    if tok == "C":
        return next(const_names) if const_names is not None else "C"
    if tok in INFIX_SYMBOL:
        left = to_infix(node[1], const_names)
        right = to_infix(node[2], const_names)
        return f"({left} {INFIX_SYMBOL[tok]} {right})"
    return UNARY_RENDER[tok].format(to_infix(node[1], const_names))


def const_name_iter(prefix: str = "c") -> Iterator[str]:
    i = 0
    while True:
        yield f"{prefix}{i}"
        i += 1


# --- évaluation --------------------------------------------------------------

#: Seuil sous lequel un dénominateur est considéré comme une singularité.
INV_EPS = 1e-3


def evaluate(node: Node, x: np.ndarray, consts: Sequence[float]) -> np.ndarray:
    """Évalue l'arbre sur ``x``, en consommant ``consts`` dans l'ordre préfixe.

    Aucune protection numérique : les domaines interdits produisent ``nan`` ou
    ``inf``, et c'est le filtre de dégénérescence qui décide. Masquer
    les singularités ici reviendrait à apprendre au modèle des fonctions qui
    n'existent pas.
    """
    idx = 0

    def go(n: Node) -> np.ndarray:
        nonlocal idx
        tok = n[0]
        if tok == "x":
            return x
        if tok == "C":
            if idx >= len(consts):
                raise ValueError("pas assez de constantes fournies")
            idx += 1
            return np.full_like(x, consts[idx - 1])
        if tok == "add":
            return go(n[1]) + go(n[2])
        if tok == "sub":
            return go(n[1]) - go(n[2])
        if tok == "mul":
            return go(n[1]) * go(n[2])
        u = go(n[1])
        if tok == "sin":
            return np.sin(u)
        if tok == "cos":
            return np.cos(u)
        if tok == "tanh":
            return np.tanh(u)
        if tok == "exp":
            return np.exp(u)
        if tok == "sq":
            return u * u
        if tok == "cube":
            return u * u * u
        if tok == "abs":
            return np.abs(u)
        if tok == "log":
            return np.where(u > 0, np.log(np.where(u > 0, u, 1.0)), np.nan)
        if tok == "sqrt":
            return np.where(u >= 0, np.sqrt(np.where(u >= 0, u, 0.0)), np.nan)
        if tok == "inv":
            return np.where(
                np.abs(u) > INV_EPS, 1.0 / np.where(np.abs(u) > INV_EPS, u, 1.0), np.nan
            )
        raise ValueError(f"token non évaluable : {tok!r}")

    with np.errstate(over="ignore", invalid="ignore", divide="ignore"):
        y = go(node)
    if idx != len(consts):
        raise ValueError(f"{len(consts) - idx} constante(s) non consommée(s)")
    return np.asarray(y, dtype=np.float64)

"""Génération de l'ensemble de squelettes (Phase 1).

Rappel du découpage : le coût symbolique — échantillonnage d'arbre,
canonicalisation, filtrage, déduplication — est payé **une fois par squelette
unique**, hors ligne. Les nuages de points sont tirés **en ligne** à
l'entraînement, si bien qu'un même squelette ne produit jamais deux fois le
même exemple. « 2M d'exemples » n'est donc pas la bonne unité de mesure ; les
chiffres qui comptent sont le nombre de squelettes uniques et le débit.

Déduplication à deux niveaux :

1. **exacte**, sur la forme canonique ;
2. **numérique**, sur une empreinte de la courbe évaluée avec un jeu de
   constantes fixé. Elle rattrape ce que la canonicalisation ne voit pas —
   ``mul(x, x)`` et ``sq(x)`` sont structurellement différents et
   numériquement identiques. C'est une heuristique : elle a des faux négatifs
   (deux membres différents d'une même famille ne collident pas), jamais de
   faux positifs à la tolérance choisie.
"""

from __future__ import annotations

import argparse
import gzip
import hashlib
import json
import multiprocessing as mp
import time
from collections import Counter
from pathlib import Path

import numpy as np

from curvy.config import DATASET_DIR, ensure_dirs
from curvy.data.expr import (
    Node,
    complexity,
    count_constants,
    depth,
    evaluate,
    to_prefix,
)
from curvy.data.pointcloud import DENSE_N, CloudConfig, normalise_y, sample_cloud
from curvy.data.sample import SamplerConfig, sample_skeleton
from curvy.seeding import DEFAULT_SEED, make_rng

#: Jeu de constantes fixe pour l'empreinte numérique. Fixe et non aléatoire :
#: deux squelettes ne sont comparables que s'ils sont évalués au même endroit.
FINGERPRINT_CONSTS = (1.7, -0.93, 2.31, 0.61, -1.42, 3.07, 0.84)
FINGERPRINT_DECIMALS = 4
#: Nombre de tirages de constantes accordés à un squelette avant de le juger
#: non viable. Sans ça, `inv(C + x)` serait jeté dès qu'un tirage malchanceux
#: place le pôle dans le domaine — alors que le squelette est bon.
VIABILITY_ATTEMPTS = 8


def fingerprint(skeleton: Node) -> str | None:
    """Empreinte de la courbe, ou ``None`` si elle n'est pas évaluable."""
    k = count_constants(skeleton)
    consts = list(FINGERPRINT_CONSTS[:k])
    if len(consts) < k:  # squelette plus riche que la table : pas d'empreinte
        return None
    x = np.linspace(-1.0, 1.0, DENSE_N)
    y = evaluate(skeleton, x, consts)
    if not np.isfinite(y).all():
        return None
    y_norm, _, _ = normalise_y(y)
    q = np.round(y_norm, FINGERPRINT_DECIMALS)
    return hashlib.blake2b(q.tobytes(), digest_size=16).hexdigest()


def _worker(task: tuple[int, int]) -> tuple[list[dict], Counter]:
    """Produit des squelettes candidats viables. Un flux RNG par worker."""
    seed, n_attempts = task
    rng = make_rng(seed)
    scfg, ccfg = SamplerConfig(), CloudConfig()
    stats: Counter = Counter()
    seen: dict[str, int] = {}
    out: list[dict] = []

    for _ in range(n_attempts):
        stats["tirages"] += 1
        sk = sample_skeleton(rng, scfg)
        if sk is None:
            stats["rejet_squelette"] += 1
            continue
        key = " ".join(to_prefix(sk))
        if key in seen:
            # La multiplicité n'est pas un déchet : c'est une mesure de la
            # probabilité a priori du squelette, dont on aura besoin pour
            # repondérer l'échantillonnage à l'entraînement.
            seen[key] += 1
            stats["doublon_local"] += 1
            continue
        seen[key] = 1

        # Viabilité : le squelette doit produire au moins un nuage valide.
        reason = None
        for _ in range(VIABILITY_ATTEMPTS):
            cloud, reason = sample_cloud(rng, sk, ccfg)
            if cloud is not None:
                break
        if reason is not None:
            stats[f"rejet_{reason}"] += 1
            continue

        stats["retenu"] += 1
        out.append(
            {
                "prefix": key,
                "depth": depth(sk),
                "n_consts": count_constants(sk),
                "complexity": complexity(sk),
                "fingerprint": fingerprint(sk),
                "count": 0,  # complété après la boucle
            }
        )
    for it in out:
        it["count"] = seen[it["prefix"]]
    return out, stats


def generate(n_attempts: int, seed: int, workers: int, out_path: Path) -> dict:
    ensure_dirs()
    per_worker = max(1, n_attempts // workers)
    tasks = [(seed + 1000 * i, per_worker) for i in range(workers)]

    t0 = time.perf_counter()
    with mp.Pool(workers) as pool:
        results = pool.map(_worker, tasks)
    elapsed = time.perf_counter() - t0

    stats: Counter = Counter()
    for _, s in results:
        stats.update(s)

    # Déduplication exacte, puis numérique (au profit du plus simple).
    by_key: dict[str, dict] = {}
    for items, _ in results:
        for it in items:
            prev = by_key.get(it["prefix"])
            if prev is None:
                by_key[it["prefix"]] = it
            else:
                prev["count"] += it["count"]
    stats["doublon_global"] = stats["retenu"] - len(by_key)

    by_fp: dict[str, dict] = {}
    kept: list[dict] = []
    for it in sorted(by_key.values(), key=lambda d: (d["complexity"], d["prefix"])):
        fp = it["fingerprint"]
        if fp is None:
            kept.append(it)
            continue
        if fp in by_fp:
            by_fp[fp]["count"] += it["count"]  # la multiplicité revient au survivant
            stats["doublon_numerique"] += 1
            continue
        by_fp[fp] = it
        kept.append(it)

    kept.sort(key=lambda d: (d["complexity"], d["prefix"]))
    with gzip.open(out_path, "wt", encoding="utf-8") as fh:
        for it in kept:
            fh.write(json.dumps(it, ensure_ascii=False) + "\n")

    summary = {
        "squelettes_uniques": len(kept),
        "tirages": stats["tirages"],
        "duree_s": round(elapsed, 2),
        "debit_tirages_par_s": round(stats["tirages"] / elapsed, 1),
        "workers": workers,
        "seed": seed,
        "fichier": str(out_path),
        "rejets": {k: v for k, v in sorted(stats.items()) if k.startswith(("rejet", "doublon"))},
        "distribution_profondeur": dict(sorted(Counter(d["depth"] for d in kept).items())),
        "distribution_constantes": dict(sorted(Counter(d["n_consts"] for d in kept).items())),
        "multiplicite_par_profondeur": {
            str(d): round(
                sum(k["count"] for k in kept if k["depth"] == d)
                / max(1, sum(1 for k in kept if k["depth"] == d)),
                1,
            )
            for d in sorted({k["depth"] for k in kept})
        },
    }
    return summary


def load_skeletons(path: Path) -> list[dict]:
    with gzip.open(path, "rt", encoding="utf-8") as fh:
        return [json.loads(line) for line in fh]


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--n", type=int, default=200_000, help="nombre de tirages d'arbres")
    ap.add_argument("--seed", type=int, default=DEFAULT_SEED)
    ap.add_argument("--workers", type=int, default=mp.cpu_count())
    ap.add_argument("--out", type=Path, default=DATASET_DIR / "skeletons-v1.jsonl.gz")
    args = ap.parse_args(argv)

    summary = generate(args.n, args.seed, args.workers, args.out)
    print(json.dumps(summary, indent=2, ensure_ascii=False))

    total_rejets = sum(summary["rejets"].values())
    print(f"\ntaux de rejet global : {100 * total_rejets / summary['tirages']:.1f} %")
    print(f"squelettes uniques   : {summary['squelettes_uniques']}")
    print(
        f"débit                : {summary['debit_tirages_par_s']:.0f} tirages/s "
        f"sur {summary['workers']} cœurs"
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

"""Grammaire v1 des expressions.

Un squelette est un arbre d'expression où les constantes numériques sont
remplacées par un placeholder unique ``C``. Le modèle ne prédit que cette
structure ; les valeurs sont trouvées après coup par optimisation.

Le vocabulaire est délibérément minuscule — 18 tokens — ce qui garde les
séquences courtes et le décodeur petit.
"""

from __future__ import annotations

PAD, BOS, EOS = "<pad>", "<bos>", "<eos>"
SPECIAL = (PAD, BOS, EOS)

BINARY = ("add", "sub", "mul")
UNARY = ("sin", "cos", "exp", "log", "sqrt", "abs", "tanh", "sq", "cube", "inv")
LEAVES = ("x", "C")

OPERATORS = BINARY + UNARY
VOCAB: tuple[str, ...] = SPECIAL + BINARY + UNARY + LEAVES

TOKEN_TO_ID = {tok: i for i, tok in enumerate(VOCAB)}
ID_TO_TOKEN = dict(enumerate(VOCAB))

ARITY: dict[str, int] = (
    {tok: 2 for tok in BINARY} | {tok: 1 for tok in UNARY} | {tok: 0 for tok in LEAVES}
)

#: Profondeur maximale du *corps* de l'expression, hors enveloppe de racine.
MAX_BODY_DEPTH = 6
#: Constantes libres autorisées dans le corps. Les 2 de l'enveloppe s'ajoutent.
MAX_BODY_CONSTANTS = 5
MAX_CONSTANTS = MAX_BODY_CONSTANTS + 2

#: Coût de complexité par token, pour le front de Pareto (Phase 5).
#: Une composition transcendante coûte plus cher qu'une addition : entre deux
#: candidats de même R², on veut proposer le plus lisible.
COMPLEXITY_COST: dict[str, int] = {
    "x": 1,
    "C": 1,
    "add": 1,
    "sub": 1,
    "mul": 2,
    "sq": 2,
    "cube": 3,
    "abs": 2,
    "sqrt": 3,
    "inv": 3,
    "tanh": 4,
    "sin": 4,
    "cos": 4,
    "exp": 4,
    "log": 4,
}

#: Rendu infixe, pour l'affichage et le passage à sympy.
INFIX_SYMBOL = {"add": "+", "sub": "-", "mul": "*"}
UNARY_RENDER = {
    "sin": "sin({0})",
    "cos": "cos({0})",
    "exp": "exp({0})",
    "log": "log({0})",
    "sqrt": "sqrt({0})",
    "abs": "Abs({0})",
    "tanh": "tanh({0})",
    "sq": "({0})**2",
    "cube": "({0})**3",
    "inv": "1/({0})",
}

assert set(COMPLEXITY_COST) == set(OPERATORS) | set(LEAVES)
assert set(UNARY_RENDER) == set(UNARY)

"""Du squelette au nuage de points, avec augmentation réaliste.

Le tirage des points suit la **courbure de la courbe** et non une loi uniforme :
un stylo ralentit dans les virages, donc les points s'y accumulent. C'est l'une
des trois composantes qui manquaient à la spec, et sans doute celle qui compte
le plus pour la robustesse au tracé à main levée.

Le bruit corrélé est appliqué **en espace d'indice** et non en espace ``x``.
Ce n'est pas une approximation par paresse : le tremblement de la main est
corrélé dans le *temps*, et le stylo avance le long de l'abscisse curviligne —
l'indice du point est donc plus proche du temps que ne l'est ``x``.
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass, field

import numpy as np

from curvy.data.expr import Node, count_constants, evaluate

__all__ = [
    "sample_cloud_fn",
    "CloudConfig",
    "PointCloud",
    "RejectReason",
    "normalise_y",
    "sample_cloud",
]

#: Grille dense servant au filtrage des dégénérescences et au calcul de densité.
DENSE_N = 512


class RejectReason:
    """Causes de rejet, comptées séparément pour être publiables."""

    NON_FINITE = "non_fini"  # nan/inf sur le domaine (log, sqrt, inv hors domaine)
    EXPLOSION = "explosion"  # |y| > seuil
    CONSTANT = "constante"  # y quasi constant : pas d'information
    DEGENERATE_NOISE = "bruit_degenere"  # étendue de y nulle après bruit
    UNIDENTIFIABLE = "bruit_excessif"  # le bruit a effacé la formule
    SPIKE = "pic_isole"  # plate partout sauf une singularité qui fixe l'échelle
    ALL = (NON_FINITE, EXPLOSION, CONSTANT, DEGENERATE_NOISE, UNIDENTIFIABLE, SPIKE)


@dataclass(frozen=True)
class CloudConfig:
    n_points_min: int = 20
    n_points_max: int = 200
    #: |C| tiré log-uniformément : sur [-1, 1], une constante doit pouvoir
    #: engendrer aussi bien une pente douce qu'une oscillation rapide.
    const_log_range: tuple[float, float] = (0.05, 20.0)
    max_abs_y: float = 1e6
    #: Étendue relative minimale de y : en dessous, la courbe est une constante.
    min_y_range: float = 1e-6
    #: Fraction maximale de la grille dense autorisée à être non finie.
    max_nonfinite_frac: float = 0.0

    # --- augmentation (chaque composante désactivable pour l'ablation) ---
    use_curvature_density: bool = True
    use_white_noise: bool = True
    use_correlated_drift: bool = True
    use_x_jitter: bool = True
    use_gaps: bool = True
    use_quantisation: bool = True

    white_sigma_range: tuple[float, float] = (1e-3, 8e-2)
    drift_sigma_range: tuple[float, float] = (1e-3, 1.2e-1)
    drift_length_range: tuple[float, float] = (0.05, 0.35)  # en fraction du nuage
    #: R² minimal entre la courbe exacte et les points bruités, tous deux
    #: normalisés. En dessous, le bruit a effacé la formule : l'exemple n'est
    #: plus de l'augmentation, c'est du bruit d'étiquetage. Constaté sur la
    #: figure d'aperçu du jalon Phase 1, où plusieurs nuages ne décrivaient
    #: manifestement plus la formule qui les avait engendrés.
    min_identifiability_r2: float = 0.95
    #: Étendue minimale du décile central de la courbe **normalisée**. Le
    #: filtre « constante » regarde l'étendue totale et se laisse berner par
    #: une quasi-singularité : la courbe est plate partout, le pic fixe
    #: l'échelle, et après normalisation il ne reste qu'un trait horizontal
    #: avec une valeur aberrante. Constaté sur la figure d'aperçu du jalon.
    min_central_spread: float = 0.20
    p_gap: float = 0.25
    gap_width_range: tuple[float, float] = (0.08, 0.30)
    quantisation_levels: tuple[int, ...] = field(default=(128, 256, 512, 1024))
    p_quantisation: float = 0.4


@dataclass
class PointCloud:
    x: np.ndarray
    y: np.ndarray
    consts: list[float]
    y_scale: float  # y_normalisé = (y_brut - y_offset) / y_scale
    y_offset: float
    n_points: int


def _sample_constants(rng: np.random.Generator, k: int, cfg: CloudConfig) -> list[float]:
    lo, hi = cfg.const_log_range
    mag = np.exp(rng.uniform(np.log(lo), np.log(hi), size=k))
    sign = rng.choice((-1.0, 1.0), size=k)
    return list(mag * sign)


def _curvature_density(y: np.ndarray, rng: np.random.Generator) -> np.ndarray:
    """Densité d'échantillonnage ∝ (1 + |dy/dx|)^α, mélangée à l'uniforme."""
    slope = np.abs(np.gradient(y))
    slope = slope / (slope.max() + 1e-12)
    alpha = rng.uniform(0.5, 2.5)
    mix = rng.uniform(0.2, 0.8)  # part d'uniforme, pour ne jamais tout concentrer
    dens = mix + (1.0 - mix) * (slope**alpha)
    return dens / dens.sum()


def _sample_x(
    rng: np.random.Generator, n: int, dense_x: np.ndarray, dense_y: np.ndarray, cfg: CloudConfig
) -> np.ndarray:
    """Positions en x, irrégulières, avec trous éventuels et bornes garanties.

    Les bornes -1 et +1 sont toujours présentes : à l'inférence, la
    normalisation envoie de toute façon le minimum et le maximum observés sur
    -1 et +1. Un nuage d'entraînement qui ne les atteindrait pas
    créerait un décalage train/test.
    """
    if cfg.use_curvature_density:
        p = _curvature_density(dense_y, rng)
    else:
        p = np.full(len(dense_x), 1.0 / len(dense_x))

    if cfg.use_gaps and rng.random() < cfg.p_gap:
        width = rng.uniform(*cfg.gap_width_range)
        start = rng.uniform(0.0, 1.0 - width)
        lo, hi = -1.0 + 2.0 * start, -1.0 + 2.0 * (start + width)
        p = np.where((dense_x > lo) & (dense_x < hi), 0.0, p)
        if p.sum() <= 0:
            p = np.full(len(dense_x), 1.0 / len(dense_x))
        p = p / p.sum()

    idx = rng.choice(len(dense_x), size=max(n - 2, 1), replace=True, p=p)
    x = np.concatenate([[-1.0, 1.0], dense_x[idx]])

    if cfg.use_x_jitter:
        step = 2.0 / len(dense_x)
        x = x + rng.normal(0.0, step * rng.uniform(0.3, 1.5), size=x.shape)
        x = np.clip(x, -1.0, 1.0)
        x[0], x[1] = -1.0, 1.0

    return np.sort(x)


def _correlated_drift(rng: np.random.Generator, n: int, cfg: CloudConfig) -> np.ndarray:
    """Bruit gaussien passé au filtre passe-bas : la dérive de la main."""
    length = max(2, int(rng.uniform(*cfg.drift_length_range) * n))
    raw = rng.normal(size=n + 4 * length)
    k = np.exp(-0.5 * (np.arange(-2 * length, 2 * length + 1) / length) ** 2)
    k /= k.sum()
    smooth = np.convolve(raw, k, mode="same")[2 * length : 2 * length + n]
    std = smooth.std()
    return smooth / std if std > 1e-12 else smooth


def normalise_y(y: np.ndarray) -> tuple[np.ndarray, float, float]:
    """Ramène y dans [-1, 1]. Retourne aussi l'affine pour l'inverser."""
    lo, hi = float(y.min()), float(y.max())
    span = hi - lo
    if span < 1e-12:
        return np.zeros_like(y), 1.0, lo
    offset = (hi + lo) / 2.0
    scale = span / 2.0
    return (y - offset) / scale, scale, offset


def sample_cloud(
    rng: np.random.Generator, skeleton: Node, cfg: CloudConfig | None = None
) -> tuple[PointCloud | None, str | None]:
    """Un nuage de points, ou ``(None, raison_du_rejet)``."""
    cfg = cfg or CloudConfig()
    consts = _sample_constants(rng, count_constants(skeleton), cfg)

    def f(xs: np.ndarray) -> np.ndarray:
        return evaluate(skeleton, xs, consts)

    return sample_cloud_fn(rng, f, cfg, consts=consts)


def sample_cloud_fn(
    rng: np.random.Generator,
    f: Callable[[np.ndarray], np.ndarray],
    cfg: CloudConfig | None = None,
    consts: list[float] | None = None,
) -> tuple[PointCloud | None, str | None]:
    """Même chaîne, pour une fonction quelconque plutôt qu'un squelette.

    Sert au sous-ensemble **hors distribution** du jeu de test (Phase 6) : les
    formules écrites à la main n'ont pas toutes de représentation dans la
    grammaire, mais doivent traverser exactement le même bruit, la même
    densification par courbure et le même filtre d'identifiabilité. Sans quoi
    on comparerait deux protocoles au lieu de deux jeux de formules.
    """
    cfg = cfg or CloudConfig()
    consts = [] if consts is None else consts

    dense_x = np.linspace(-1.0, 1.0, DENSE_N)
    dense_y = f(dense_x)

    finite = np.isfinite(dense_y)
    if (~finite).mean() > cfg.max_nonfinite_frac:
        return None, RejectReason.NON_FINITE
    if np.abs(dense_y[finite]).max() > cfg.max_abs_y:
        return None, RejectReason.EXPLOSION
    span = float(dense_y[finite].max() - dense_y[finite].min())
    if span < cfg.min_y_range * max(1.0, float(np.abs(dense_y[finite]).max())):
        return None, RejectReason.CONSTANT

    dense_norm, _, _ = normalise_y(dense_y[finite])
    spread = float(np.percentile(dense_norm, 95) - np.percentile(dense_norm, 5))
    if spread < cfg.min_central_spread:
        return None, RejectReason.SPIKE

    n = int(rng.integers(cfg.n_points_min, cfg.n_points_max + 1))
    x = _sample_x(rng, n, dense_x, np.where(finite, dense_y, 0.0), cfg)
    y = f(x)
    ok = np.isfinite(y)
    if not ok.all():
        x, y = x[ok], y[ok]
        if len(x) < cfg.n_points_min:
            return None, RejectReason.NON_FINITE

    y_clean = y.copy()
    y_span = float(y.max() - y.min())
    if cfg.use_white_noise:
        sigma = np.exp(rng.uniform(*np.log(cfg.white_sigma_range)))
        y = y + rng.normal(0.0, sigma * max(y_span, 1e-9), size=y.shape)
    if cfg.use_correlated_drift:
        sigma = np.exp(rng.uniform(*np.log(cfg.drift_sigma_range)))
        y = y + sigma * max(y_span, 1e-9) * _correlated_drift(rng, len(y), cfg)

    y_norm, scale, offset = normalise_y(y)
    if float(y_norm.max() - y_norm.min()) < 1e-9:
        return None, RejectReason.DEGENERATE_NOISE

    # La formule doit rester identifiable à partir des points bruités.
    clean_norm = (y_clean - offset) / scale
    ss_res = float(np.sum((y_norm - clean_norm) ** 2))
    ss_tot = float(np.sum((y_norm - y_norm.mean()) ** 2))
    r2 = 1.0 - ss_res / ss_tot if ss_tot > 1e-12 else 0.0
    if r2 < cfg.min_identifiability_r2:
        return None, RejectReason.UNIDENTIFIABLE

    if cfg.use_quantisation and rng.random() < cfg.p_quantisation:
        levels = int(rng.choice(cfg.quantisation_levels))
        y_norm = np.round(y_norm * levels) / levels
        x = np.round((x + 1.0) / 2.0 * levels) / levels * 2.0 - 1.0

    return PointCloud(x, y_norm, consts, scale, offset, len(x)), None

"""Échantillonnage d'arbres d'expression.

Processus de branchement volontairement **sous-critique** : l'espérance du
nombre d'enfants vaut ``0,35 × 1 + 0,30 × 2 = 0,95 < 1``, donc les arbres
restent petits en moyenne et la coupure de profondeur ne sert que pour la
queue de distribution. Un processus sur-critique produirait une majorité
d'arbres butant sur la profondeur maximale, donc une distribution dégénérée.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from curvy.data.canonical import canonicalise, strip_absorbable_root, wrap_root
from curvy.data.expr import Node, count_constants, depth, has_x
from curvy.data.grammar import BINARY, MAX_BODY_CONSTANTS, MAX_BODY_DEPTH, UNARY

__all__ = ["SamplerConfig", "sample_skeleton"]


@dataclass(frozen=True)
class SamplerConfig:
    max_body_depth: int = MAX_BODY_DEPTH
    max_body_constants: int = MAX_BODY_CONSTANTS
    p_leaf: float = 0.35
    p_unary: float = 0.35
    p_x_given_leaf: float = 0.75
    #: Poids par opérateur unaire ; ``sq``/``sin`` sont plus fréquents dans la
    #: nature que ``cube`` ou ``log``, et le dataset doit refléter cette réalité
    #: plutôt qu'une uniforme artificielle.
    unary_weights: tuple[float, ...] = (
        1.4,  # sin
        1.0,  # cos
        0.8,  # exp
        0.6,  # log
        0.7,  # sqrt
        0.6,  # abs
        1.0,  # tanh
        1.4,  # sq
        0.7,  # cube
        0.8,  # inv
    )


def _sample_tree(rng: np.random.Generator, remaining: int, cfg: SamplerConfig) -> Node:
    if remaining <= 1:
        return ("x",) if rng.random() < cfg.p_x_given_leaf else ("C",)
    u = rng.random()
    if u < cfg.p_leaf:
        return ("x",) if rng.random() < cfg.p_x_given_leaf else ("C",)
    if u < cfg.p_leaf + cfg.p_unary:
        w = np.asarray(cfg.unary_weights, dtype=np.float64)
        op = UNARY[rng.choice(len(UNARY), p=w / w.sum())]
        return (op, _sample_tree(rng, remaining - 1, cfg))
    op = BINARY[rng.integers(len(BINARY))]
    return (op, _sample_tree(rng, remaining - 1, cfg), _sample_tree(rng, remaining - 1, cfg))


def sample_skeleton(rng: np.random.Generator, cfg: SamplerConfig | None = None) -> Node | None:
    """Un squelette canonique enveloppé, ou ``None`` si le tirage est rejeté.

    Rejets possibles à ce stade : arbre sans ``x`` (donc constant), corps trop
    profond ou trop riche en constantes après canonicalisation.
    """
    cfg = cfg or SamplerConfig()
    body = _sample_tree(rng, cfg.max_body_depth, cfg)
    body = strip_absorbable_root(canonicalise(body))
    if not has_x(body):
        return None
    if depth(body) > cfg.max_body_depth:
        return None
    if count_constants(body) > cfg.max_body_constants:
        return None
    return wrap_root(body)

"""Partition entraînement / test des squelettes (Phase 6).

Un jeu de test n'a de valeur que si ses squelettes sont **exclus de
l'entraînement**. Jusqu'ici ce n'était le cas d'aucun run : le jeu de
validation était tiré du même fichier que le flux d'entraînement, et on a
mesuré qu'à 10,2 M tirages la probabilité qu'un de ses squelettes n'ait jamais
été vu vaut 5·10⁻⁸ en profondeur 8, et zéro en dessous. Tous les chiffres
publiés avant la Phase 6 mesurent donc de la **restitution**, pas de la
généralisation.

## Deux choix, et leurs raisons

**1. La réserve commence à la profondeur 5.** Il n'existe qu'*un* squelette de
profondeur 3 — `C*x+C`, la droite — et sept en profondeur 4. Les mettre au test
ne mesurerait pas une généralisation : ça retirerait de l'entraînement les
formes les plus courantes du produit. On tient donc à l'écart uniquement là où
la strate est peuplée (242 squelettes en profondeur 5, 182 231 en profondeur 8),
et **on l'écrit dans le rapport** plutôt que de laisser croire à une réserve
uniforme. Les formes peu profondes sont couvertes autrement, par le
sous-ensemble hors distribution écrit à la main.

**2. L'appartenance au test dépend du squelette, pas du fichier.** Elle est
tirée d'un hachage stable de la notation préfixe : le même squelette tombera
toujours du même côté, quel que soit l'ordre du fichier, sa taille, ou une
régénération future avec une autre graine. Une partition par indice de ligne
se casserait à la première régénération du jeu de données — silencieusement.
"""

from __future__ import annotations

import hashlib
from collections import defaultdict
from dataclasses import dataclass

__all__ = ["Partition", "RESERVE_PAR_PROFONDEUR", "SEL", "partitionner", "valeur_de_hachage"]

#: Sel du hachage. Le changer redistribue **toute** la partition : c'est
#: équivalent à changer de jeu de test, et ça invalide toute comparaison
#: antérieure. À ne faire que délibérément, en versionnant le nom.
SEL = "curvy-test-v1"

#: Nombre de squelettes tenus à l'écart, par profondeur. Rien avant la
#: profondeur 5 : voir l'en-tête du module.
#:
#: Les tailles sont dimensionnées sur la **précision par profondeur**, pas sur
#: le taux global. À 100 exemples, une strate se mesure à ±9,6 points près
#: (intervalle à 95 %) — de quoi lire une tendance, pas de quoi comparer deux
#: runs. À 400-950, on tombe à ±3 à ±5 points.
#:
#: La profondeur 5 est plafonnée par la grammaire elle-même : il n'existe que
#: **242 squelettes** de cette profondeur, en réserver 48 en retire déjà 20 %
#: de l'entraînement. Cette strate restera à ±14 points, quel que soit le
#: budget. C'est une limite de la grammaire, pas du jeu de test.
#:
#: Ajouter des nuages par squelette ne remplacerait pas des squelettes : six
#: nuages du même squelette ne valent pas six échantillons indépendants, c'est
#: la variance entre formules qui domine.
RESERVE_PAR_PROFONDEUR = {5: 48, 6: 400, 7: 600, 8: 950}


@dataclass
class Partition:
    entrainement: list[dict]
    test: list[dict]

    @property
    def prefixes_de_test(self) -> frozenset[str]:
        return frozenset(it["prefix"] for it in self.test)

    def rapport(self) -> dict:
        par_prof: dict[int, dict[str, int]] = defaultdict(lambda: {"entrainement": 0, "test": 0})
        for it in self.entrainement:
            par_prof[it["depth"]]["entrainement"] += 1
        for it in self.test:
            par_prof[it["depth"]]["test"] += 1
        return {
            "sel": SEL,
            "n_entrainement": len(self.entrainement),
            "n_test": len(self.test),
            "par_profondeur": {str(d): par_prof[d] for d in sorted(par_prof)},
            "profondeurs_sans_reserve": sorted(d for d in par_prof if par_prof[d]["test"] == 0),
        }


def valeur_de_hachage(prefix: str, sel: str = SEL) -> float:
    """Un réel de [0, 1) déterminé par le squelette seul, stable entre machines."""
    digest = hashlib.sha256(f"{sel}\x00{prefix}".encode()).digest()
    return int.from_bytes(digest[:8], "big") / 2**64


def partitionner(
    items: list[dict],
    reserve: dict[int, int] | None = None,
    sel: str = SEL,
) -> Partition:
    """Sépare les squelettes en entraînement et test.

    Pour chaque profondeur, les ``n`` squelettes de plus petite valeur de
    hachage partent au test. Déterministe, indépendant de l'ordre du fichier.
    """
    reserve = RESERVE_PAR_PROFONDEUR if reserve is None else reserve
    par_prof: dict[int, list[dict]] = defaultdict(list)
    for it in items:
        par_prof[it["depth"]].append(it)

    en_test: set[str] = set()
    for profondeur, n in reserve.items():
        candidats = par_prof.get(profondeur, [])
        if n > len(candidats):
            raise ValueError(
                f"profondeur {profondeur} : {n} squelettes demandés au test, "
                f"{len(candidats)} disponibles"
            )
        classes = sorted(candidats, key=lambda it: valeur_de_hachage(it["prefix"], sel))
        en_test.update(it["prefix"] for it in classes[:n])

    return Partition(
        entrainement=[it for it in items if it["prefix"] not in en_test],
        test=[it for it in items if it["prefix"] in en_test],
    )

"""Jeu de test figé de la Phase 6 : ce que le modèle n'a jamais vu.

Trois sous-ensembles, mesurés séparément et jamais agrégés :

1. **tenu à l'écart** — mêmes grammaire et générateur que l'entraînement, mais
   squelettes explicitement exclus du flux (``curvy.data.split``). Répond à
   « le modèle généralise-t-il à des formules inédites de sa propre famille ? » ;
2. **hors distribution** — formules écrites à la main. Certaines sont dans la
   grammaire mais qu'un tirage aléatoire ne produirait jamais (une gaussienne,
   une sinusoïde à 30 rad, la fonction de Runge) ; d'autres en sortent
   franchement (une marche, un plancher, une fonction par morceaux). Répond à
   « que se passe-t-il quand on sort du bac à sable ? » ;
3. **réel** — tracés capturés au canvas. N'existe pas encore : l'outil de
   capture est en Phase 8.

Toutes traversent **le même** pipeline de nuage que l'entraînement — bruit
blanc, dérive corrélée, densité liée à la courbure, trous, quantification
 — via ``sample_cloud_fn``. Sans quoi on comparerait deux protocoles
au lieu de deux familles de formules.
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from math import erf, factorial

import numpy as np

__all__ = ["FORMULES_A_LA_MAIN", "FormuleAlaMain"]


@dataclass(frozen=True)
class FormuleAlaMain:
    nom: str
    f: Callable[[np.ndarray], np.ndarray]
    #: Notation préfixe si la grammaire v1 peut l'écrire, sinon ``None``.
    #: Vérifié par test contre les budgets réels, jamais affirmé à la main.
    prefixe: str | None
    commentaire: str


def _morceaux(x: np.ndarray) -> np.ndarray:
    return np.where(x < 0.0, -(x**2), x**2)


def _dent_de_scie(x: np.ndarray) -> np.ndarray:
    return 3.0 * x - np.floor(3.0 * x)


def _puissance(x: np.ndarray) -> np.ndarray:
    # x^x n'est réel que sur x > 0 ; on décale pour rester défini partout.
    u = (x + 1.0) / 2.0 + 0.05
    return u**u


def _weierstrass(x: np.ndarray) -> np.ndarray:
    return sum(0.5**n * np.cos(3.0**n * np.pi * x) for n in range(6))


#: Le catalogue. `prefixe` renseigné = la grammaire sait l'écrire ; c'est le
#: modèle qu'on met en défaut, pas le vocabulaire. `prefixe = None` = la
#: grammaire ne peut pas, et l'échec attendu mesure une limite assumée.
FORMULES_A_LA_MAIN: tuple[FormuleAlaMain, ...] = (
    # --- exprimables, mais qu'un tirage ne produirait jamais ---
    FormuleAlaMain(
        "gaussienne",
        lambda x: np.exp(-8.0 * x**2),
        "exp mul C sq x",
        "la cloche : partout en science, jamais tirée par le générateur",
    ),
    FormuleAlaMain(
        "runge",
        lambda x: 1.0 / (1.0 + 25.0 * x**2),
        "inv add C mul C sq x",
        "le contre-exemple classique de l'interpolation polynomiale",
    ),
    FormuleAlaMain(
        "sinus_haute_frequence",
        lambda x: np.sin(30.0 * x),
        "sin mul C x",
        "structure triviale, constante inatteignable : mesuré le 2026-08-20, "
        "l'ajustement retrouve C jusqu'à 12 avec 6 essais et jusqu'à 20 avec 50, "
        "puis échoue. Le modèle peut proposer LE bon squelette et scorer 0.",
    ),
    FormuleAlaMain(
        "sigmoide_raide",
        lambda x: np.tanh(20.0 * x),
        "tanh mul C x",
        "quasi-marche, mais dérivable — la grammaire l'a",
    ),
    FormuleAlaMain(
        "oscillation_amortie",
        lambda x: np.exp(-3.0 * x) * np.sin(10.0 * x),
        "mul exp mul C x sin mul C x",
        "produit de deux motifs courants, rare en tirage",
    ),
    FormuleAlaMain(
        "sinus_de_inverse",
        lambda x: np.sin(1.0 / np.where(np.abs(x) < 0.05, 0.05, x)),
        "sin inv x",
        "pathologique près de 0 : oscille infiniment vite. L'appelable borne "
        "|x| à 0,05, il n'est donc pas exactement sin(1/x) — le squelette "
        "plafonne à R² 0,90 contre lui, et c'est voulu.",
    ),
    FormuleAlaMain(
        "logarithme_decale",
        lambda x: np.log(x + 1.5),
        "log add x C",
        "translation, pour sortir du domaine où log explose",
    ),
    FormuleAlaMain(
        "racine_de_valeur_absolue",
        lambda x: np.sqrt(np.abs(x)),
        "sqrt abs x",
        "pointe en 0, dérivée infinie",
    ),
    # --- pôles et asymptotes : aucun exemple jusqu'ici, et c'est un mode
    #     d'échec à part entière — la courbe part à l'infini dans le domaine ---
    FormuleAlaMain(
        "hyperbole_raide",
        lambda x: 1.0 / (x - 1.15),
        "inv add x C",
        "pôle JUSTE en dehors du domaine, donc pente très forte à droite. Un pôle "
        "*dans* le domaine a été essayé et retiré : le générateur le rejette "
        "systématiquement (filtre d'explosion à 1e6, aucun point non fini toléré). "
        "Ce n'est pas une lacune — un tracé au canvas est borné par l'écran, la "
        "démo ne recevra jamais de valeur infinie.",
    ),
    FormuleAlaMain(
        "sinus_cardinal",
        lambda x: np.sin(10.0 * x) / np.where(np.abs(x) < 1e-6, 1e-6, x),
        "mul sin mul C x inv x",
        "trou apparent en 0 alors que la fonction y est régulière",
    ),
    FormuleAlaMain(
        "chirp",
        lambda x: np.sin(25.0 * x**2),
        "sin mul C sq x",
        "fréquence variable : le motif change le long du tracé",
    ),
    FormuleAlaMain(
        "x_sinus_de_inverse",
        lambda x: x * np.sin(1.0 / np.where(np.abs(x) < 0.03, 0.03, x)),
        "mul x sin inv x",
        "oscillation dont l'amplitude s'annule — continue mais non dérivable en 0",
    ),
    FormuleAlaMain(
        "croissance_exponentielle",
        lambda x: np.exp(5.0 * x),
        "exp mul C x",
        "trois ordres de grandeur sur le domaine : la normalisation écrase la gauche",
    ),
    FormuleAlaMain(
        "pointe_etroite",
        lambda x: np.exp(-200.0 * x**2),
        "exp mul C sq x",
        "quasi nulle partout sauf sur 10 % du domaine : la densité par courbure est mise à l'épreuve",
    ),
    FormuleAlaMain(
        "coude",
        lambda x: np.abs(x - 0.2),
        "abs add x C",
        "dérivée discontinue en un point, valeur continue",
    ),
    # --- hors grammaire, modes d'échec distincts ------------------------------
    FormuleAlaMain(
        "arctangente",
        lambda x: np.arctan(8.0 * x),
        None,
        "saturation douce, très proche de tanh : le modèle devrait s'en tirer par substitution",
    ),
    FormuleAlaMain(
        "onde_triangulaire",
        lambda x: 2.0 * np.abs(2.0 * (1.5 * x - np.floor(1.5 * x + 0.5))) - 1.0,
        None,
        "périodique, continue, dérivée discontinue partout",
    ),
    FormuleAlaMain(
        "bessel_j0",
        lambda x: np.sum(
            [(-1.0) ** k / (factorial(k) ** 2) * (4.0 * x / 2.0) ** (2 * k) for k in range(12)],
            axis=0,
        ),
        None,
        "fonction spéciale oscillante amortie, absente du vocabulaire",
    ),
    FormuleAlaMain(
        "polynome_degre_7",
        lambda x: 0.5 * x**7 - 1.2 * x**5 + 0.9 * x**3 - 0.3 * x,
        None,
        "exactement ce que la baseline sait faire de mieux : le cas où elle doit gagner",
    ),
    FormuleAlaMain(
        "deux_echelles",
        lambda x: np.sin(2.0 * x) + 0.08 * np.sin(40.0 * x),
        None,
        "un motif lent et un rapide superposés : que rend le front de Pareto ?",
    ),
    FormuleAlaMain(
        "marche_douce_decalee",
        lambda x: 1.0 / (1.0 + np.exp(-25.0 * (x - 0.35))),
        None,
        "sigmoïde logistique décalée : la grammaire n'a ni exp(-u) ni décalage interne bon marché",
    ),
    # --- hors grammaire : l'échec est une limite assumée, pas un bug ---
    FormuleAlaMain(
        "marche",
        lambda x: np.sign(x),
        None,
        "discontinue : aucun opérateur du vocabulaire ne saute",
    ),
    FormuleAlaMain(
        "plancher",
        lambda x: np.floor(3.0 * x),
        None,
        "constante par morceaux, pas de `floor` au vocabulaire",
    ),
    FormuleAlaMain(
        "par_morceaux",
        _morceaux,
        None,
        "deux lois recollées en 0 : la grammaire n'a pas de conditionnelle",
    ),
    FormuleAlaMain(
        "dent_de_scie",
        _dent_de_scie,
        None,
        "périodique et discontinue, cumule les deux difficultés",
    ),
    FormuleAlaMain(
        "x_puissance_x",
        _puissance,
        None,
        "exposant variable : pas de puissance générale au vocabulaire",
    ),
    FormuleAlaMain(
        "weierstrass_tronquee",
        _weierstrass,
        None,
        "somme de six cosinus : profondeur très au-delà du budget",
    ),
    FormuleAlaMain(
        "erf",
        lambda x: np.vectorize(erf)(2.0 * x),
        None,
        "fonction spéciale, absente du vocabulaire (proche de tanh)",
    ),
    FormuleAlaMain(
        "cloche_asymetrique",
        lambda x: np.exp(-8.0 * (x - 0.3) ** 2) - 0.5 * np.exp(-20.0 * (x + 0.5) ** 2),
        None,
        "deux cloches décalées : profondeur et constantes au-delà du budget",
    ),
)

"""Pondération des squelettes à l'entraînement.

La déduplication inverse la distribution de profondeur : 71 % des squelettes
uniques sont à la profondeur maximale, alors que 35 % des *tirages* étaient de
profondeur minimale. Tirer uniformément dans l'ensemble unique reviendrait à
n'entraîner que sur des expressions complexes, et à échouer sur le cas d'usage
principal — une droite, une parabole, une sinusoïde.

La pondération est **stratifiée par profondeur**, avec une cible explicite.
L'alternative testée d'abord — un poids ``count ** tau`` — a été mesurée puis
abandonnée : entre ``tau = 0,75`` et ``tau = 1,25``, la part de la profondeur 3
passe de 5 % à 80 %. Un paramètre avec une falaise pareille n'est pas réglable.
"""

from __future__ import annotations

from collections import defaultdict
from collections.abc import Sequence

import numpy as np

__all__ = ["DEFAULT_DEPTH_TARGET", "stratified_weights", "describe_weights"]

#: Part de l'entraînement allouée à chaque profondeur de squelette.
#: Délibérément plus riche en formes simples que l'ensemble dédupliqué, moins
#: que la distribution brute de l'échantillonneur. À ablater en Phase 4.
DEFAULT_DEPTH_TARGET: dict[int, float] = {
    3: 0.05,  # C*x + C — l'affine, un seul squelette mais un cas très fréquent
    4: 0.08,
    5: 0.12,
    6: 0.20,
    7: 0.25,
    8: 0.30,
}


def stratified_weights(depths: Sequence[int], target: dict[int, float] | None = None) -> np.ndarray:
    """Poids de tirage par squelette réalisant la distribution cible.

    Dans chaque strate de profondeur, les squelettes sont équiprobables : la
    multiplicité ne sert plus qu'à décrire la strate, pas à départager ses
    membres. Les profondeurs absentes de la cible reçoivent un poids nul ;
    celles absentes des données voient leur part redistribuée.
    """
    target = target or DEFAULT_DEPTH_TARGET
    d = np.asarray(depths)
    members: dict[int, np.ndarray] = defaultdict(lambda: np.array([], dtype=int))
    for depth in np.unique(d):
        members[int(depth)] = np.flatnonzero(d == depth)

    present = {k: v for k, v in target.items() if len(members.get(k, ())) > 0}
    total = sum(present.values())
    if total <= 0:
        raise ValueError("aucune profondeur de la cible n'est présente dans les données")

    w = np.zeros(len(d), dtype=np.float64)
    for depth, share in present.items():
        idx = members[depth]
        w[idx] = (share / total) / len(idx)
    return w


def describe_weights(depths: Sequence[int], weights: np.ndarray) -> dict[int, float]:
    """Part effective de chaque profondeur, pour vérifier que la cible est tenue."""
    d = np.asarray(depths)
    return {int(k): round(100 * float(weights[d == k].sum()), 2) for k in sorted(np.unique(d))}

"""Sous-paquet curvy.tokenizer — à remplir (Phase 2)."""

"""Encodage des squelettes en séquences d'identifiants, et masque d'arité.

Le vocabulaire tient en 18 tokens. Aucun littéral numérique n'y
figure : le modèle ne prédit jamais de valeur, seulement une structure
.

La pièce importante de ce module n'est pas ``encode``/``decode`` mais
``legal_mask``. En notation préfixe, un simple compteur d'arité dit
à chaque pas quels tokens peuvent encore mener à un arbre complet. En masquant
les logits des autres pendant le beam search, on obtient une garantie et non
une espérance : **le taux de sorties syntaxiquement invalides doit être
exactement 0 %**. Si la mesure de la Phase 5 dit autre chose, le bug est ici.
"""

from __future__ import annotations

from collections.abc import Sequence

import numpy as np

from curvy.data.expr import Node, from_prefix, to_prefix
from curvy.data.grammar import (
    ARITY,
    BOS,
    EOS,
    ID_TO_TOKEN,
    MAX_CONSTANTS,
    PAD,
    TOKEN_TO_ID,
    VOCAB,
)

PAD_ID, BOS_ID, EOS_ID = TOKEN_TO_ID[PAD], TOKEN_TO_ID[BOS], TOKEN_TO_ID[EOS]
VOCAB_SIZE = len(VOCAB)

#: Arité par identifiant ; -1 pour les tokens spéciaux, qui n'en ont pas.
ARITY_BY_ID = np.array([ARITY.get(tok, -1) for tok in VOCAB], dtype=np.int64)
IS_OPERAND = ARITY_BY_ID >= 0
CONST_ID = TOKEN_TO_ID["C"]

#: Longueur maximale d'une séquence, `<bos>` et `<eos>` compris.
#: Un arbre de profondeur 8 pourrait en théorie compter 255 nœuds, mais la
#: canonicalisation et le budget de constantes ramènent le maximum très en
#: deçà. **Mesuré** sur les 255 080 squelettes du dataset v1 : maximum 43,
#: médiane 17, p99 à 29 — aucune séquence au-delà de 48. On fixe 48 plutôt
#: qu'une puissance de deux confortable : le coût de l'attention du décodeur
#: est quadratique en cette longueur.
MAX_SEQ_LEN = 48

__all__ = [
    "BOS_ID",
    "CONST_ID",
    "EOS_ID",
    "MAX_SEQ_LEN",
    "PAD_ID",
    "VOCAB_SIZE",
    "DecodeState",
    "decode",
    "encode",
    "legal_mask",
]


def encode(skeleton: Node, add_special: bool = True) -> list[int]:
    ids = [TOKEN_TO_ID[t] for t in to_prefix(skeleton)]
    return [BOS_ID, *ids, EOS_ID] if add_special else ids


def decode(ids: Sequence[int]) -> Node:
    """Reconstruit l'arbre. Ignore ``<bos>``/``<pad>``, s'arrête à ``<eos>``."""
    toks: list[str] = []
    for i in ids:
        tok = ID_TO_TOKEN[int(i)]
        if tok == EOS:
            break
        if tok in (BOS, PAD):
            continue
        toks.append(tok)
    return from_prefix(toks)


class DecodeState:
    """Suit ce qu'une séquence préfixe partielle autorise encore.

    ``remaining`` est le nombre de sous-arbres encore attendus. Il vaut 1 au
    début, 0 quand l'arbre est complet, et jamais négatif tant qu'on respecte
    le masque.
    """

    __slots__ = ("remaining", "n_consts", "n_emitted")

    def __init__(self) -> None:
        self.remaining = 1
        self.n_consts = 0
        self.n_emitted = 0

    def copy(self) -> DecodeState:
        s = DecodeState()
        s.remaining, s.n_consts, s.n_emitted = self.remaining, self.n_consts, self.n_emitted
        return s

    def advance(self, token_id: int) -> DecodeState:
        tok = ID_TO_TOKEN[int(token_id)]
        if tok in (BOS, PAD, EOS):
            return self
        self.remaining += ARITY[tok] - 1
        self.n_consts += tok == "C"
        self.n_emitted += 1
        return self

    @property
    def complete(self) -> bool:
        return self.remaining == 0


def legal_mask(
    state: DecodeState, max_len: int = MAX_SEQ_LEN, max_consts: int = MAX_CONSTANTS
) -> np.ndarray:
    """Masque booléen des tokens légaux au pas suivant.

    Trois contraintes, toutes nécessaires :

    1. **complétude** — ``<eos>`` est interdit tant que l'arbre n'est pas
       complet, et *seul* ``<eos>`` est permis une fois qu'il l'est ;
    2. **budget de longueur** — un token n'est légal que s'il reste assez de
       place pour fermer tous les sous-arbres qu'il ouvre (au minimum une
       feuille par sous-arbre restant) ;
    3. **budget de constantes** — ``C`` disparaît une fois le quota atteint
      , sinon le problème d'ajustement de la Phase 5 devient
       arbitrairement dur.
    """
    mask = np.zeros(VOCAB_SIZE, dtype=bool)

    if state.complete:
        mask[EOS_ID] = True
        return mask

    # `budget` = tokens d'arbre encore émettables avant d'atteindre max_len,
    # en réservant une place pour <eos> (et une pour le <bos> déjà émis).
    budget = max_len - 2 - state.n_emitted
    if budget <= 0:
        return mask  # séquence morte : le beam search doit écarter ce faisceau

    for tid in np.flatnonzero(IS_OPERAND):
        new_remaining = state.remaining + ARITY_BY_ID[tid] - 1
        if new_remaining > budget - 1:
            continue  # impossible de refermer l'arbre dans le budget
        if tid == CONST_ID and state.n_consts >= max_consts:
            continue
        mask[tid] = True
    return mask

"""Sous-paquet curvy.model — à remplir."""

"""Configurations d'architecture.

Deux tailles, et l'ordre dans lequel on s'en sert compte : on valide la chaîne
complète avec ``SMALL`` (~5M) avant de lancer ``V1`` (~30M). Découvrir un bug
de dataloader après huit heures d'entraînement est une erreur évitable.
"""

from __future__ import annotations

from dataclasses import dataclass

from curvy.tokenizer.vocab import MAX_SEQ_LEN, VOCAB_SIZE


@dataclass(frozen=True)
class ModelConfig:
    d_model: int = 512
    n_heads: int = 8
    n_encoder_layers: int = 6
    n_decoder_layers: int = 6
    dim_feedforward: int = 1024
    dropout: float = 0.1
    vocab_size: int = VOCAB_SIZE
    max_seq_len: int = MAX_SEQ_LEN
    #: Encodage des coordonnées d'entrée. ``linear`` suit la spec ; ``fourier``
    #: ajoute des features sinusoïdales, connues pour aider les réseaux à lire
    #: des coordonnées brutes. À départager par ablation en Phase 4.
    point_encoding: str = "linear"
    n_fourier_bands: int = 8
    #: Encodage positionnel sur les points d'entrée. Désactivé = le nuage est
    #: traité comme un **ensemble**.
    point_positional: bool = False


#: ~5M de paramètres — validation du bout en bout en quelques minutes.
SMALL = ModelConfig(
    d_model=256, n_heads=4, n_encoder_layers=4, n_decoder_layers=4, dim_feedforward=512
)

#: ~30M de paramètres — cible v1. Les dimensions de la spec (d_model 512,
#: 6 + 6 couches) donnent 44M avec un FFN standard à 4×d_model ; on descend le
#: FFN à 2×d_model, ce qui ramène à ~31M sans toucher aux dimensions annoncées.
V1 = ModelConfig()

PRESETS = {"small": SMALL, "v1": V1}

"""Architecture encodeur-décodeur.

- **Encodeur** : transformer sur l'ensemble des points. Chaque point ``(x, y)``
  est projeté en un vecteur, puis les points s'attendent mutuellement. Par
  défaut **aucun encodage positionnel** : un nuage est un ensemble, pas une
  séquence — et l'information d'ordre est déjà portée par la valeur de ``x``
  elle-même.
- **Décodeur** : autorégressif sur les tokens du squelette, cross-attention
  vers l'encodeur.

On s'appuie sur ``nn.TransformerEncoder``/``TransformerDecoder`` de PyTorch,
qui utilisent le SDPA fusionné en interne. Réimplémenter l'attention à la main
n'apporterait ici que des occasions de se tromper sur les masques.
"""

from __future__ import annotations

import math

import torch
from torch import Tensor, nn

from curvy.model.config import ModelConfig

__all__ = ["CurvyModel", "PointEmbedding", "count_parameters"]


class PointEmbedding(nn.Module):
    """``(x, y)`` -> vecteur de dimension ``d_model``.

    En mode ``fourier``, les coordonnées sont d'abord développées en
    sinus/cosinus à plusieurs fréquences. Un réseau lit très mal des
    coordonnées brutes ; ce développement est le remède standard. Il reste
    optionnel pour que l'ablation de la Phase 4 puisse le chiffrer.
    """

    def __init__(self, cfg: ModelConfig) -> None:
        super().__init__()
        self.mode = cfg.point_encoding
        if self.mode == "fourier":
            bands = 2.0 ** torch.arange(cfg.n_fourier_bands)
            self.register_buffer("bands", bands, persistent=False)
            in_dim = 2 + 4 * cfg.n_fourier_bands
        elif self.mode == "linear":
            in_dim = 2
        else:
            raise ValueError(f"point_encoding inconnu : {cfg.point_encoding!r}")
        self.proj = nn.Linear(in_dim, cfg.d_model)

    def forward(self, points: Tensor) -> Tensor:
        if self.mode == "fourier":
            angles = points.unsqueeze(-1) * self.bands * math.pi  # (B, N, 2, K)
            feats = torch.cat([points, angles.sin().flatten(-2), angles.cos().flatten(-2)], dim=-1)
            return self.proj(feats)
        return self.proj(points)


class SinusoidalPositions(nn.Module):
    """Encodage positionnel classique, pour les tokens du décodeur."""

    def __init__(self, d_model: int, max_len: int) -> None:
        super().__init__()
        pos = torch.arange(max_len).unsqueeze(1)
        div = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
        pe = torch.zeros(max_len, d_model)
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe, persistent=False)

    def forward(self, x: Tensor) -> Tensor:
        return x + self.pe[: x.size(1)].unsqueeze(0)


class CurvyModel(nn.Module):
    def __init__(self, cfg: ModelConfig) -> None:
        super().__init__()
        self.cfg = cfg
        self.point_embed = PointEmbedding(cfg)
        self.point_pos = SinusoidalPositions(cfg.d_model, 4096) if cfg.point_positional else None

        enc_layer = nn.TransformerEncoderLayer(
            d_model=cfg.d_model,
            nhead=cfg.n_heads,
            dim_feedforward=cfg.dim_feedforward,
            dropout=cfg.dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,  # pre-norm : bien plus stable sans warmup agressif
        )
        self.encoder = nn.TransformerEncoder(
            enc_layer,
            cfg.n_encoder_layers,
            norm=nn.LayerNorm(cfg.d_model),
            # Le chemin « nested tensor » est incompatible avec norm_first ;
            # le laisser à True ne produit qu'un avertissement à chaque
            # construction de modèle.
            enable_nested_tensor=False,
        )

        self.token_embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.token_pos = SinusoidalPositions(cfg.d_model, cfg.max_seq_len)
        dec_layer = nn.TransformerDecoderLayer(
            d_model=cfg.d_model,
            nhead=cfg.n_heads,
            dim_feedforward=cfg.dim_feedforward,
            dropout=cfg.dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.decoder = nn.TransformerDecoder(
            dec_layer, cfg.n_decoder_layers, norm=nn.LayerNorm(cfg.d_model)
        )
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size)
        # Le partage des poids entre embedding et sortie est gratuit et régularise.
        self.head.weight = self.token_embed.weight

        self.apply(self._init)

    @staticmethod
    def _init(module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, std=0.02)

    def encode_points(self, points: Tensor, point_mask: Tensor) -> Tensor:
        """``points`` (B, N, 2), ``point_mask`` (B, N) — True sur le remplissage."""
        h = self.point_embed(points)
        if self.point_pos is not None:
            h = self.point_pos(h)
        return self.encoder(h, src_key_padding_mask=point_mask)

    def decode(
        self, memory: Tensor, memory_mask: Tensor, tokens_in: Tensor, tokens_mask: Tensor
    ) -> Tensor:
        h = self.token_pos(self.token_embed(tokens_in))
        # Masque causal **booléen** et non flottant : mélanger un masque
        # d'attention float et un masque de padding bool déclenche un
        # avertissement de dépréciation et, à terme, un chemin non fusionné.
        causal = torch.ones(
            tokens_in.size(1), tokens_in.size(1), dtype=torch.bool, device=tokens_in.device
        ).triu(1)
        h = self.decoder(
            h,
            memory,
            tgt_mask=causal,
            tgt_key_padding_mask=tokens_mask,
            memory_key_padding_mask=memory_mask,
            tgt_is_causal=True,
        )
        return self.head(h)

    def forward(
        self, points: Tensor, point_mask: Tensor, tokens_in: Tensor, tokens_mask: Tensor
    ) -> Tensor:
        memory = self.encode_points(points, point_mask)
        return self.decode(memory, point_mask, tokens_in, tokens_mask)


def count_parameters(model: nn.Module) -> dict[str, int]:
    """Compte exact par bloc — la spec demande le chiffre, pas une estimation."""
    groups: dict[str, int] = {}
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        block = name.split(".")[0]
        groups[block] = groups.get(block, 0) + p.numel()
    groups["TOTAL"] = sum(v for k, v in groups.items() if k != "TOTAL")
    return groups

"""Sous-paquet curvy.train — à remplir."""

"""Sauvegarde et reprise.

Un entraînement doit survivre à une déconnexion SSH, à un OOM provoqué par un
processus voisin, et à une coupure. On sauvegarde donc l'état complet — modèle,
optimiseur, ordonnanceur, pas courant, graine — et pas seulement les poids.
"""

from __future__ import annotations

import json
from pathlib import Path

import torch

__all__ = ["latest_checkpoint", "load_checkpoint", "save_checkpoint"]


def save_checkpoint(
    path: Path, *, model, optimizer, scheduler, step: int, config: dict, best: float
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")
    torch.save(
        {
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict(),
            "scheduler": scheduler.state_dict(),
            "step": step,
            "config": config,
            "best": best,
        },
        tmp,
    )
    tmp.replace(path)  # atomique : jamais de checkpoint tronqué
    (path.parent / "config.json").write_text(json.dumps(config, indent=2, ensure_ascii=False))


def load_checkpoint(path: Path, *, model, optimizer=None, scheduler=None) -> dict:
    state = torch.load(path, map_location="cpu", weights_only=False)
    model.load_state_dict(state["model"])
    if optimizer is not None:
        optimizer.load_state_dict(state["optimizer"])
    if scheduler is not None:
        scheduler.load_state_dict(state["scheduler"])
    return state


def latest_checkpoint(run_dir: Path) -> Path | None:
    p = run_dir / "last.pt"
    return p if p.exists() else None

"""Hyperparamètres d'entraînement.

Tout ce qui influence un résultat figure ici et est sérialisé dans le
checkpoint **et** dans la fiche d'expérience. Une expérience dont on ne peut
pas relire les hyperparamètres n'a pas eu lieu.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from pathlib import Path

from curvy.config import DATASET_DIR, RUNS_DIR
from curvy.data.weighting import DEFAULT_DEPTH_TARGET
from curvy.seeding import DEFAULT_SEED


@dataclass
class TrainConfig:
    run_name: str = "exp-001"
    preset: str = "small"
    seed: int = DEFAULT_SEED

    steps: int = 20_000
    batch_size: int = 512
    lr: float = 3e-4
    weight_decay: float = 0.01
    betas: tuple[float, float] = (0.9, 0.95)
    warmup_steps: int = 500
    min_lr_ratio: float = 0.05
    grad_clip: float = 1.0

    workers: int = 8
    bucketing: bool = False  # mesuré perdant aujourd'hui,
    bf16: bool = True
    compile_model: bool = False  # mesuré à x0,93 en Phase 0

    log_every: int = 50
    eval_every: int = 1_000
    ckpt_every: int = 1_000
    val_size: int = 512
    val_seed: int = 777  # figé et distinct de `seed` : le val ne bouge jamais

    skeletons: Path = DATASET_DIR / "skeletons-v1.jsonl.gz"
    #: Exclut du flux les squelettes réservés au jeu de test. Faux
    #: par défaut pour ne pas réécrire silencieusement l'histoire des runs
    #: exp-001 à exp-003, qui ont tourné sans. La valeur est écrite dans
    #: `config.json` : c'est elle qui dit si les chiffres d'un run sur le jeu
    #: de test veulent dire quelque chose.
    exclure_test: bool = False
    runs_dir: Path = RUNS_DIR
    depth_target: dict[int, float] = field(default_factory=lambda: dict(DEFAULT_DEPTH_TARGET))

    @property
    def run_dir(self) -> Path:
        return self.runs_dir / self.run_name

    def to_dict(self) -> dict:
        d = asdict(self)
        for k, v in d.items():
            if isinstance(v, Path):
                d[k] = str(v)
        return d

"""Boucle d'entraînement.

Volontairement courte : la configuration est ailleurs (``config.py``), les
métriques ailleurs (``metrics.py``), la reprise ailleurs (``checkpoint.py``).
Ce fichier ne fait qu'orchestrer.
"""

from __future__ import annotations

import json
import math
import time
from pathlib import Path

import torch
from torch import nn
from torch.utils.data import DataLoader

from curvy.data.dataset import BucketedBatches, CurvyStream, collate, make_validation_set
from curvy.devices import pick_device
from curvy.model.config import PRESETS
from curvy.model.curvy import CurvyModel, count_parameters
from curvy.precision import bf16_supported, configure_precision, precision_report
from curvy.seeding import make_rng, seed_everything
from curvy.tokenizer.vocab import PAD_ID
from curvy.train.checkpoint import latest_checkpoint, load_checkpoint, save_checkpoint
from curvy.train.config import TrainConfig
from curvy.train.metrics import evaluate_model, token_accuracy

__all__ = ["Trainer"]


def cosine_with_warmup(step: int, cfg: TrainConfig) -> float:
    """Facteur multiplicatif du learning rate."""
    if step < cfg.warmup_steps:
        return (step + 1) / cfg.warmup_steps
    progress = (step - cfg.warmup_steps) / max(1, cfg.steps - cfg.warmup_steps)
    progress = min(1.0, progress)
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return cfg.min_lr_ratio + (1.0 - cfg.min_lr_ratio) * cosine


class Trainer:
    def __init__(self, cfg: TrainConfig, resume: bool = True) -> None:
        self.cfg = cfg
        seed_everything(cfg.seed)
        configure_precision(tf32=True)
        self.device = pick_device("auto").device
        self.rng = make_rng(cfg.seed + 1)

        model_cfg = PRESETS[cfg.preset]
        self.model = CurvyModel(model_cfg).to(self.device)
        self.params = count_parameters(self.model)
        if cfg.compile_model:
            self.model = torch.compile(self.model)

        self.opt = torch.optim.AdamW(
            self.model.parameters(),
            lr=cfg.lr,
            betas=cfg.betas,
            weight_decay=cfg.weight_decay,
        )
        self.sched = torch.optim.lr_scheduler.LambdaLR(
            self.opt, lambda s: cosine_with_warmup(s, cfg)
        )
        self.loss_fn = nn.CrossEntropyLoss(ignore_index=PAD_ID)
        self.use_bf16 = cfg.bf16 and bf16_supported(self.device) and self.device.type == "cuda"

        cfg.run_dir.mkdir(parents=True, exist_ok=True)
        self.log_path = cfg.run_dir / "log.jsonl"
        self.step = 0
        self.best = -1.0

        ckpt = latest_checkpoint(cfg.run_dir) if resume else None
        if ckpt is not None:
            state = load_checkpoint(
                ckpt, model=self.model, optimizer=self.opt, scheduler=self.sched
            )
            self.step = state["step"]
            self.best = state.get("best", -1.0)
            self._log({"event": "reprise", "step": self.step, "checkpoint": str(ckpt)})

        self.val = make_validation_set(cfg.skeletons, cfg.val_size, seed=cfg.val_seed)
        self.loader = self._make_loader()
        self._log(
            {
                "event": "partition",
                "exclure_test": cfg.exclure_test,
                "n_squelettes_exclus": self._n_exclus,
            }
        )

    #: Un batch écarté de temps en temps n'est pas une avarie : le modèle
    #: continue d'apprendre normalement, et sauter 3 batches sur 11 000 ne se
    #: voit dans aucune métrique (mesuré sur exp-005, 2026-08-20). Ce qui doit
    #: arrêter un run, c'est un **taux**, pas un total — sinon un run sain finit
    #: par atteindre n'importe quel compteur cumulatif.
    #:
    #: Premier seuil posé à « 5 au total » : il aurait tué exp-005 au step
    #: 28 500 alors que son taux R² dépassait celui du run de contrôle.
    FENETRE_INCIDENTS = 1_000
    MAX_INCIDENTS_PAR_FENETRE = 20  # 2 % des steps de la fenêtre

    def _incident(self, genre: str, batch, valeur: float) -> None:
        """Consigne un batch non fini, le sauvegarde, et arrête si ça se répète.

        Le batch fautif est écrit sur disque : sans lui, la cause reste
        indevinable. Il pèse moins d'un mégaoctet, et c'est la seule occasion de
        l'attraper — le flux ne repasse jamais deux fois au même endroit.
        """
        recents = getattr(self, "_incidents_recents", None)
        if recents is None:
            recents = self._incidents_recents = []
        recents.append(self.step)
        # On ne garde que la fenêtre glissante.
        seuil_bas = self.step - self.FENETRE_INCIDENTS
        self._incidents_recents = recents = [s for s in recents if s > seuil_bas]
        self._n_incidents = getattr(self, "_n_incidents", 0) + 1
        chemin = self.cfg.run_dir / f"batch-non-fini-{self.step}.pt"
        try:
            torch.save(
                {
                    "points": batch.points.detach().cpu(),
                    "point_mask": batch.point_mask.detach().cpu(),
                    "tokens": batch.tokens.detach().cpu(),
                    "token_mask": batch.token_mask.detach().cpu(),
                },
                chemin,
            )
        except Exception as exc:  # ne jamais faire tomber le run sur la sauvegarde
            chemin = f"échec de sauvegarde : {exc}"

        pts = batch.points.detach().float()
        self._log(
            {
                "event": "incident",
                "genre": genre,
                "step": self.step,
                "valeur": valeur,
                "n_incidents": self._n_incidents,
                "n_dans_la_fenetre": len(recents),
                "batch_sauve": str(chemin),
                "points_min": round(float(pts.min()), 4),
                "points_max": round(float(pts.max()), 4),
                "points_finis": bool(torch.isfinite(pts).all()),
                "tokens_max": int(batch.tokens.max()),
            }
        )
        if len(recents) >= self.MAX_INCIDENTS_PAR_FENETRE:
            self._log(
                {
                    "event": "abandon",
                    "raison": (
                        f"{len(recents)} batches non finis en {self.FENETRE_INCIDENTS} steps"
                    ),
                    "step": self.step,
                }
            )
            raise RuntimeError(
                f"{len(recents)} batches non finis en {self.FENETRE_INCIDENTS} steps "
                f"— arrêt. Batches sauvés dans {self.cfg.run_dir}"
            )

    def _prefixes_de_test(self) -> frozenset[str] | None:
        """Les squelettes réservés au test, ou ``None`` si le run n'exclut rien.

        Calculé à partir du fichier de squelettes lui-même : la partition est
        une fonction du hachage de chaque squelette, il n'y a donc pas de
        fichier d'index à tenir synchronisé.
        """
        if not self.cfg.exclure_test:
            self._n_exclus = 0
            return None
        from curvy.data.generate import load_skeletons
        from curvy.data.split import partitionner

        prefixes = partitionner(load_skeletons(self.cfg.skeletons)).prefixes_de_test
        self._n_exclus = len(prefixes)
        return prefixes

    def _make_loader(self) -> DataLoader:
        cfg = self.cfg
        stream = CurvyStream(
            cfg.skeletons,
            seed=cfg.seed + 100 * self.step,
            depth_target=cfg.depth_target,
            exclure=self._prefixes_de_test(),
        )
        common = dict(
            num_workers=cfg.workers,
            pin_memory=self.device.type == "cuda",
            persistent_workers=cfg.workers > 0,
            prefetch_factor=4 if cfg.workers > 0 else None,
        )
        if cfg.bucketing:
            return DataLoader(BucketedBatches(stream, cfg.batch_size), batch_size=None, **common)
        return DataLoader(stream, batch_size=cfg.batch_size, collate_fn=collate, **common)

    def _log(self, record: dict) -> None:
        record.setdefault("t", round(time.time(), 3))
        with self.log_path.open("a", encoding="utf-8") as fh:
            fh.write(json.dumps(record, ensure_ascii=False) + "\n")
        print(json.dumps(record, ensure_ascii=False), flush=True)

    def _autocast(self):
        if self.use_bf16:
            return torch.autocast("cuda", dtype=torch.bfloat16)
        return torch.autocast(self.device.type, enabled=False)

    def run(self) -> None:
        cfg = self.cfg
        self._log(
            {
                "event": "demarrage",
                "config": cfg.to_dict(),
                "parametres": self.params["TOTAL"],
                "device": str(self.device),
                "precision": precision_report(self.device),
                "val_size": len(self.val),
            }
        )
        it = iter(self.loader)
        self.model.train()
        t0 = time.perf_counter()
        window: list[float] = []
        n_tokens = 0

        while self.step < cfg.steps:
            batch = next(it).to(self.device)
            tokens_in, target = batch.tokens[:, :-1], batch.tokens[:, 1:]
            with self._autocast():
                logits = self.model(
                    batch.points, batch.point_mask, tokens_in, batch.token_mask[:, :-1]
                )
                loss = self.loss_fn(logits.reshape(-1, logits.size(-1)).float(), target.reshape(-1))
            self.opt.zero_grad(set_to_none=True)

            # Un seul batch à loss non finie suffit à tuer le run : le clip de
            # gradient renvoie NaN, multiplie TOUS les gradients par NaN, et
            # l'optimiseur écrit NaN dans tous les poids. Plus rien ne revient
            # ensuite. exp-005 s'est entraîné 3 500 steps sur du NaN sans que
            # rien ne l'arrête (2026-08-20).
            if not torch.isfinite(loss):
                self.step += 1
                self._incident("loss_non_finie", batch, float(loss.detach()))
                continue

            loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), cfg.grad_clip)

            # La loss peut être finie et un gradient déborder quand même.
            if not torch.isfinite(grad_norm):
                self.step += 1
                self._incident("gradient_non_fini", batch, float(loss.detach()))
                continue

            self.opt.step()
            self.sched.step()
            self.step += 1

            window.append(loss.detach().item())
            n_tokens += int((target != PAD_ID).sum())

            if self.step % cfg.log_every == 0:
                dt = time.perf_counter() - t0
                tok_acc, seq_acc = token_accuracy(logits.detach().float(), target)
                self._log(
                    {
                        "event": "train",
                        "step": self.step,
                        "loss": round(sum(window) / len(window), 5),
                        "token_acc": round(tok_acc, 4),
                        "seq_acc_tf": round(seq_acc, 4),
                        "lr": round(self.sched.get_last_lr()[0], 8),
                        "grad_norm": round(float(grad_norm), 4),
                        "tokens_par_s": round(n_tokens / dt),
                        "s_par_step": round(dt / cfg.log_every, 4),
                        "vram_Mio": (
                            round(torch.cuda.max_memory_allocated(self.device) / 2**20)
                            if self.device.type == "cuda"
                            else None
                        ),
                    }
                )
                window, n_tokens, t0 = [], 0, time.perf_counter()

            if self.step % cfg.eval_every == 0 or self.step == cfg.steps:
                t_eval = time.perf_counter()
                rep = evaluate_model(self.model, self.val, self.device, self.rng)
                self._log(
                    {
                        "event": "eval",
                        "step": self.step,
                        "duree_s": round(time.perf_counter() - t_eval, 1),
                        **rep.as_dict(),
                    }
                )
                if rep.r2_rate > self.best:
                    self.best = rep.r2_rate
                    self._save("best.pt")
                t0 = time.perf_counter()

            if self.step % cfg.ckpt_every == 0:
                self._save("last.pt")

        self._save("last.pt")
        self._log({"event": "fin", "step": self.step, "meilleur_r2_rate": round(self.best, 4)})

    def _save(self, name: str) -> None:
        save_checkpoint(
            Path(self.cfg.run_dir) / name,
            model=self.model,
            optimizer=self.opt,
            scheduler=self.sched,
            step=self.step,
            config=self.cfg.to_dict(),
            best=self.best,
        )

"""Métriques d'entraînement et d'évaluation.

Trois familles, jamais agrégées entre elles :

- **token** : accuracy en teacher forcing. Courbe lisse, utile en continu,
  ne dit presque rien sur la qualité réelle ;
- **séquence** : le squelette décodé est-il exactement celui attendu. Diagnostic
  honnête mais pessimiste — plusieurs squelettes décrivent le même nuage ;
- **R²** : la courbe proposée retrouve-t-elle la fonction ? C'est la métrique
  principale. Les constantes sont ajustées sur 80 % des points **observés**
  (bruités), le R² est mesuré sur les 20 % restants contre la valeur **exacte**
  de la fonction génératrice.

Chaque évaluation rapporte l'**oracle** : la même procédure appliquée au vrai
squelette. Sans lui, un taux de 55 % est illisible.

Le R² est mesuré **deux fois**, avec deux façons de tenir des points à l'écart :

- **interpolation** — 20 % de points tirés au hasard. On juge le remplissage
  entre les points observés ;
- **extrapolation** — les 20 % d'abscisses les plus à droite. On juge la
  prédiction *au-delà* de ce qui a été vu.

Le second n'est pas un raffinement : c'est le seul endroit où une formule bat
franchement un polynôme ajusté. Mesuré le 2026-08-19 sur le vrai squelette,
0,785 en interpolation contre 0,428 en extrapolation ; un polynôme à degré
honnête tombe lui de 0,670 à **0,088**. Ne mesurer qu'en interpolation, c'est
se comparer là où on est le plus faible (cf. `docs/benchmarks/results.md`).
"""

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np
import torch

from curvy.data.expr import Node, evaluate, to_prefix
from curvy.infer.decode import greedy_decode, ids_to_node
from curvy.infer.fit import fit_constants, r_squared
from curvy.tokenizer.vocab import PAD_ID

__all__ = ["EvalReport", "evaluate_model", "token_accuracy"]

R2_THRESHOLD = 0.99


def token_accuracy(logits: torch.Tensor, target: torch.Tensor) -> tuple[float, float]:
    """(accuracy token, accuracy séquence en teacher forcing)."""
    pred = logits.argmax(dim=-1)
    valid = target != PAD_ID
    correct = (pred == target) & valid
    tok = float(correct.sum()) / max(1, int(valid.sum()))
    seq_ok = ((pred == target) | ~valid).all(dim=1)
    return tok, float(seq_ok.float().mean())


@dataclass
class EvalReport:
    n: int = 0
    token_acc: float = 0.0
    seq_acc_teacher: float = 0.0
    seq_acc_greedy: float = 0.0
    r2_rate: float = 0.0
    r2_median: float = 0.0
    r2_rate_oracle: float = 0.0
    r2_rate_extrap: float = 0.0
    r2_rate_extrap_oracle: float = 0.0
    invalid_rate: float = 0.0
    fit_failed_rate: float = 0.0
    per_depth: dict[int, float] = field(default_factory=dict)

    def as_dict(self) -> dict:
        d = {k: v for k, v in self.__dict__.items() if k != "per_depth"}
        d["per_depth"] = {str(k): round(v, 4) for k, v in sorted(self.per_depth.items())}
        return {k: (round(v, 4) if isinstance(v, float) else v) for k, v in d.items()}


def _r2_holdout(
    node: Node,
    x: np.ndarray,
    y_obs: np.ndarray,
    truth: np.ndarray,
    rng: np.random.Generator,
    holdout: float = 0.2,
    mode: str = "interpolation",
) -> tuple[float, bool]:
    """Ajuste sur 80 % des points observés, mesure sur 20 % contre la vérité.

    ``mode="interpolation"`` tire les points tenus à l'écart au hasard.
    ``mode="extrapolation"`` tient à l'écart les abscisses les plus à droite :
    on ajuste sur la partie gauche de la courbe et on doit prédire la suite.
    """
    n = len(x)
    n_hold = max(3, int(round(holdout * n)))
    if mode == "extrapolation":
        order = np.argsort(x)
        hold, keep = order[-n_hold:], order[:-n_hold]
    else:
        idx = rng.permutation(n)
        hold, keep = idx[:n_hold], idx[n_hold:]
    res = fit_constants(node, x[keep], y_obs[keep], rng)
    if not res.ok:
        return float("-inf"), True
    with np.errstate(over="ignore", invalid="ignore", divide="ignore"):
        pred = evaluate(node, x[hold], res.consts)
    return r_squared(truth[hold], pred), False


@torch.no_grad()
def evaluate_model(model, examples, device, rng, batch_size: int = 128) -> EvalReport:
    """``examples`` : liste de ``ValidationExample`` (cf. curvy.data.dataset)."""
    from curvy.data.dataset import collate

    model.eval()
    rep = EvalReport(n=len(examples))
    tok_num = tok_den = 0
    seq_teacher = seq_greedy = 0
    invalid = fit_failed = 0
    r2s: list[float] = []
    r2s_oracle: list[float] = []
    r2s_extrap: list[float] = []
    r2s_extrap_oracle: list[float] = []
    by_depth: dict[int, list[bool]] = {}

    for start in range(0, len(examples), batch_size):
        chunk = examples[start : start + batch_size]
        batch = collate([(ex.points, ex.ids) for ex in chunk]).to(device)

        logits = model(
            batch.points, batch.point_mask, batch.tokens[:, :-1], batch.token_mask[:, :-1]
        )
        target = batch.tokens[:, 1:]
        pred = logits.argmax(dim=-1)
        valid = target != PAD_ID
        tok_num += int(((pred == target) & valid).sum())
        tok_den += int(valid.sum())
        seq_teacher += int(((pred == target) | ~valid).all(dim=1).sum())

        decoded = greedy_decode(model, batch.points, batch.point_mask)
        for ex, ids in zip(chunk, decoded, strict=True):
            node = ids_to_node(ids)
            if node is None:
                invalid += 1
                r2s.append(float("-inf"))
                r2s_extrap.append(float("-inf"))
                by_depth.setdefault(ex.depth, []).append(False)
                continue
            if to_prefix(node) == to_prefix(ex.node):
                seq_greedy += 1
            r2, failed = _r2_holdout(node, ex.x, ex.y, ex.y_clean, rng)
            fit_failed += failed
            r2s.append(r2)
            by_depth.setdefault(ex.depth, []).append(r2 >= R2_THRESHOLD)
            r2_or, _ = _r2_holdout(ex.node, ex.x, ex.y, ex.y_clean, rng)
            r2s_oracle.append(r2_or)

            # Extrapolation : mêmes candidats, mais les points tenus à l'écart
            # sont les abscisses les plus à droite. C'est là que la structure
            # rapporte et qu'un polynôme s'effondre (cf. en-tête de module).
            r2_ex, _ = _r2_holdout(node, ex.x, ex.y, ex.y_clean, rng, mode="extrapolation")
            r2s_extrap.append(r2_ex)
            r2_ex_or, _ = _r2_holdout(ex.node, ex.x, ex.y, ex.y_clean, rng, mode="extrapolation")
            r2s_extrap_oracle.append(r2_ex_or)

    arr = np.array(r2s)
    rep.token_acc = tok_num / max(1, tok_den)
    rep.seq_acc_teacher = seq_teacher / rep.n
    rep.seq_acc_greedy = seq_greedy / rep.n
    rep.r2_rate = float((arr >= R2_THRESHOLD).mean())
    rep.r2_median = float(np.median(arr[np.isfinite(arr)])) if np.isfinite(arr).any() else -1.0
    rep.r2_rate_oracle = float((np.array(r2s_oracle) >= R2_THRESHOLD).mean()) if r2s_oracle else 0.0
    rep.r2_rate_extrap = float((np.array(r2s_extrap) >= R2_THRESHOLD).mean()) if r2s_extrap else 0.0
    rep.r2_rate_extrap_oracle = (
        float((np.array(r2s_extrap_oracle) >= R2_THRESHOLD).mean()) if r2s_extrap_oracle else 0.0
    )
    rep.invalid_rate = invalid / rep.n
    rep.fit_failed_rate = fit_failed / rep.n
    rep.per_depth = {d: float(np.mean(v)) for d, v in by_depth.items()}
    model.train()
    return rep

"""``python -m curvy.train.run`` — lance ou reprend un entraînement.

La configuration est assemblée **par recoupement** entre les arguments de la
ligne de commande et les champs de ``TrainConfig``, et non par une liste
recopiée à la main. La liste recopiée a déjà coûté un run : ``--exclure-test``
existait, était accepté sans erreur, et n'arrivait jamais jusqu'à la config —
l'entraînement voyait donc les squelettes du jeu de test (2026-08-20). Un
drapeau qui ne fait rien en silence est pire qu'un drapeau absent.
"""

from __future__ import annotations

import argparse
import dataclasses
from pathlib import Path

from curvy.train.config import TrainConfig
from curvy.train.loop import Trainer


def main(argv: list[str] | None = None) -> int:
    cfg = TrainConfig()
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--run-name", default=cfg.run_name)
    ap.add_argument("--preset", default=cfg.preset)
    ap.add_argument("--seed", type=int, default=cfg.seed)
    ap.add_argument("--steps", type=int, default=cfg.steps)
    ap.add_argument("--batch-size", type=int, default=cfg.batch_size)
    ap.add_argument("--lr", type=float, default=cfg.lr)
    ap.add_argument("--warmup-steps", type=int, default=cfg.warmup_steps)
    ap.add_argument("--workers", type=int, default=cfg.workers)
    ap.add_argument("--eval-every", type=int, default=cfg.eval_every)
    ap.add_argument("--log-every", type=int, default=cfg.log_every)
    ap.add_argument("--val-size", type=int, default=cfg.val_size)
    ap.add_argument("--bucketing", action="store_true")
    ap.add_argument("--compile", dest="compile_model", action="store_true")
    ap.add_argument("--no-resume", action="store_true")
    ap.add_argument("--skeletons", type=Path, default=cfg.skeletons)
    ap.add_argument(
        "--exclure-test",
        dest="exclure_test",
        action="store_true",
        help="retire du flux les squelettes réservés au jeu de test",
    )
    args = ap.parse_args(argv)
    cfg = config_depuis_arguments(args)
    Trainer(cfg, resume=not args.no_resume).run()
    return 0


#: Arguments qui pilotent le lancement plutôt que la configuration du run.
HORS_CONFIG = frozenset({"no_resume"})


def config_depuis_arguments(args: argparse.Namespace) -> TrainConfig:
    """Assemble la ``TrainConfig`` par recoupement avec les champs du dataclass.

    Tout argument dont le nom correspond à un champ y est transporté. Un
    argument qui n'en vise aucun est une erreur franche : c'est le seul moyen de
    ne pas se retrouver avec un drapeau qui ne fait rien.
    """
    champs = {f.name for f in dataclasses.fields(TrainConfig)}
    fournis = {k: v for k, v in vars(args).items() if k not in HORS_CONFIG}
    orphelins = set(fournis) - champs
    if orphelins:
        raise ValueError(
            f"arguments sans champ correspondant dans TrainConfig : {sorted(orphelins)} — "
            "ajouter le champ, ou le déclarer dans HORS_CONFIG s'il pilote le lancement"
        )
    return TrainConfig(**fournis)


if __name__ == "__main__":
    raise SystemExit(main())

"""Sous-paquet curvy.eval — à remplir."""

"""Sous-paquet curvy.infer — à remplir."""

"""Décodage incrémental avec cache clé/valeur.

**Le problème, mesuré.** `nn.TransformerDecoder` n'a pas d'API incrémentale : au
pas ``t``, le décodage glouton ou en beam lui repasse le préfixe **entier**, et
tout est recalculé pour les ``t-1`` positions déjà traitées. Sur ``T = 40``
tokens, cela fait ``T(T+1)/2 = 820`` positions calculées au lieu de ``40``.

Profilage du 2026-08-20 : le décodage pèse **98,5 %** d'une requête à beam 8, et
dans ce décodage les passes du décodeur pèsent **99 %** (le masque d'arité en
Python, 1 %). C'est donc le seul poste qui vaille d'être optimisé.

**Ce que fait ce module.** Il rejoue à la main le corps d'une
``nn.TransformerDecoderLayer`` en ``norm_first``, en **réutilisant les
sous-modules du modèle entraîné** — `self_attn`, `multihead_attn`, `linear1/2`,
`norm1/2/3`. Aucun poids n'est copié ni remappé : c'est le même modèle, parcouru
autrement. Deux caches :

- les projections **clé/valeur de l'auto-attention**, allongées d'un token par
  pas ;
- les projections **clé/valeur de la cross-attention**, calculées **une seule
  fois** — la mémoire de l'encodeur ne change pas d'un pas à l'autre.

**Le risque est de se tromper sur les masques ou l'ordre des normalisations**,
et l'en-tête de `curvy.model.curvy` le dit : réimplémenter l'attention n'apporte
que des occasions de se tromper. La parade est un test d'équivalence stricte —
les logits du chemin incrémental doivent coïncider avec ceux du chemin complet,
sur des entrées aléatoires, à la tolérance flottante près. Sans ce test, ce
module ne devrait pas exister.
"""

from __future__ import annotations

from dataclasses import dataclass, field

import torch
import torch.nn.functional as F
from torch import Tensor

__all__ = ["CacheDecodeur", "pas_incremental"]


@dataclass
class CacheDecodeur:
    """État de décodage d'un batch. Une entrée par couche."""

    #: clés/valeurs de l'auto-attention, (B, têtes, t, dim_tête)
    self_k: list[Tensor] = field(default_factory=list)
    self_v: list[Tensor] = field(default_factory=list)
    #: clés/valeurs de la cross-attention, invariantes
    memoire_k: list[Tensor] = field(default_factory=list)
    memoire_v: list[Tensor] = field(default_factory=list)
    longueur: int = 0

    def reordonner(self, index: Tensor) -> None:
        """Réordonne le cache selon la provenance des faisceaux retenus.

        Indispensable en beam search : au pas suivant, le faisceau ``j`` peut
        descendre de n'importe quel faisceau du pas précédent. Sans ce
        réagencement, chaque faisceau hériterait du passé d'un autre — et rien
        ne le signalerait, les formes restant valides.
        """
        for liste in (self.self_k, self.self_v, self.memoire_k, self.memoire_v):
            for i, t in enumerate(liste):
                liste[i] = t.index_select(0, index)


def _projeter(module, entree: Tensor, quoi: str) -> Tensor:
    """q, k ou v d'une ``nn.MultiheadAttention`` à partir de ses poids groupés."""
    d = module.embed_dim
    dec = {"q": 0, "k": d, "v": 2 * d}[quoi]
    poids = module.in_proj_weight[dec : dec + d]
    biais = None if module.in_proj_bias is None else module.in_proj_bias[dec : dec + d]
    return F.linear(entree, poids, biais)


def _en_tetes(t: Tensor, n_tetes: int) -> Tensor:
    b, longueur, d = t.shape
    return t.view(b, longueur, n_tetes, d // n_tetes).transpose(1, 2)


def _fusionner(t: Tensor) -> Tensor:
    b, n_tetes, longueur, dim = t.transpose(1, 2).transpose(1, 2).shape
    return t.transpose(1, 2).reshape(b, longueur, n_tetes * dim)


@torch.no_grad()
def pas_incremental(
    model,
    memoire: Tensor,
    memoire_mask: Tensor,
    token: Tensor,
    cache: CacheDecodeur,
) -> Tensor:
    """Logits du prochain token, pour ``token`` (B, 1). Met le cache à jour.

    Ne renvoie **que** la dernière position : c'est tout ce dont un décodage
    autorégressif a besoin, et c'est précisément ce que le chemin complet
    recalcule inutilement pour tout le préfixe.
    """
    couches = model.decoder.layers
    premiere = cache.longueur == 0
    if premiere:
        cache.self_k = [None] * len(couches)
        cache.self_v = [None] * len(couches)
        cache.memoire_k = [None] * len(couches)
        cache.memoire_v = [None] * len(couches)

    # Position absolue du token courant. `SinusoidalPositions` ajoute
    # `pe[:len]` ; ici on ne traite qu'une position, celle d'indice
    # `cache.longueur`. Et le modèle ne met PAS l'embedding à l'échelle par
    # sqrt(d) — l'ajouter aurait produit des logits plausibles et faux.
    h = model.token_embed(token) + model.token_pos.pe[cache.longueur].view(1, 1, -1)

    pad_memoire = memoire_mask.unsqueeze(1).unsqueeze(2) if memoire_mask is not None else None

    for i, couche in enumerate(couches):
        n_tetes = couche.self_attn.num_heads

        # --- auto-attention (norm_first) ---
        x = couche.norm1(h)
        q = _en_tetes(_projeter(couche.self_attn, x, "q"), n_tetes)
        k = _en_tetes(_projeter(couche.self_attn, x, "k"), n_tetes)
        v = _en_tetes(_projeter(couche.self_attn, x, "v"), n_tetes)
        if cache.self_k[i] is not None:
            k = torch.cat([cache.self_k[i], k], dim=2)
            v = torch.cat([cache.self_v[i], v], dim=2)
        cache.self_k[i], cache.self_v[i] = k, v
        # Aucun masque causal : le cache ne contient que le passé.
        a = F.scaled_dot_product_attention(q, k, v)
        h = h + couche.self_attn.out_proj(_fusionner(a))

        # --- cross-attention, clés/valeurs calculées une seule fois ---
        x = couche.norm2(h)
        if cache.memoire_k[i] is None:
            cache.memoire_k[i] = _en_tetes(_projeter(couche.multihead_attn, memoire, "k"), n_tetes)
            cache.memoire_v[i] = _en_tetes(_projeter(couche.multihead_attn, memoire, "v"), n_tetes)
        q = _en_tetes(_projeter(couche.multihead_attn, x, "q"), n_tetes)
        a = F.scaled_dot_product_attention(
            q,
            cache.memoire_k[i],
            cache.memoire_v[i],
            attn_mask=~pad_memoire if pad_memoire is not None else None,
        )
        h = h + couche.multihead_attn.out_proj(_fusionner(a))

        # --- réseau à propagation avant ---
        x = couche.norm3(h)
        h = h + couche.linear2(couche.dropout(couche.activation(couche.linear1(x))))

    cache.longueur += 1
    if model.decoder.norm is not None:
        h = model.decoder.norm(h)
    return model.head(h)

"""Décodage contraint par l'arité.

Le masque de ``curvy.tokenizer.vocab`` est appliqué aux logits à chaque pas :
toute séquence produite est donc syntaxiquement valide **par construction**
. Deux décodeurs, le même masque :

- ``greedy_decode`` — un candidat, utilisé par les métriques d'entraînement
  parce qu'il est rapide ;
- ``beam_search`` — ``k`` candidats, ce que le produit livre réellement
 . La différence n'est pas cosmétique : la baseline
  polynomiale a le droit d'essayer huit degrés et de garder le meilleur, alors
  qu'un décodage glouton n'a droit qu'à un essai. Comparer les deux, c'est
  comparer à armes inégales.
"""

from __future__ import annotations

import numpy as np
import torch

from curvy.data.expr import Node, from_prefix
from curvy.data.grammar import ID_TO_TOKEN
from curvy.infer.cache import CacheDecodeur, pas_incremental
from curvy.tokenizer.vocab import BOS_ID, EOS_ID, MAX_SEQ_LEN, DecodeState, legal_mask

__all__ = ["beam_search", "greedy_decode", "ids_to_node"]


def ids_to_node(ids: list[int]) -> Node | None:
    toks = [ID_TO_TOKEN[i] for i in ids]
    try:
        return from_prefix(toks)
    except ValueError:
        return None


@torch.no_grad()
def greedy_decode(
    model, points: torch.Tensor, point_mask: torch.Tensor, max_len: int = MAX_SEQ_LEN
) -> list[list[int]]:
    """Décodage glouton batché. Retourne les tokens d'arbre, sans `<bos>`/`<eos>`."""
    device = points.device
    b = points.size(0)
    memory = model.encode_points(points, point_mask)

    states = [DecodeState() for _ in range(b)]
    seqs: list[list[int]] = [[] for _ in range(b)]
    done = np.zeros(b, dtype=bool)
    tokens = torch.full((b, 1), BOS_ID, dtype=torch.long, device=device)
    pad_mask = torch.zeros_like(tokens, dtype=torch.bool)

    for _ in range(max_len - 1):
        logits = model.decode(memory, point_mask, tokens, pad_mask)[:, -1]  # (B, V)
        masks = np.stack([legal_mask(s, max_len=max_len) for s in states])
        allowed = torch.from_numpy(masks).to(device)
        logits = logits.masked_fill(~allowed, float("-inf"))
        nxt = logits.argmax(dim=-1)

        for i in range(b):
            if done[i]:
                continue
            tid = int(nxt[i])
            if tid == EOS_ID or not masks[i].any():
                done[i] = True
                continue
            seqs[i].append(tid)
            states[i].advance(tid)
        if done.all():
            break
        tokens = torch.cat([tokens, nxt.unsqueeze(1)], dim=1)
        pad_mask = torch.zeros_like(tokens, dtype=torch.bool)
    return seqs


@torch.no_grad()
def beam_search(
    model,
    points: torch.Tensor,
    point_mask: torch.Tensor,
    beam: int = 8,
    max_len: int = MAX_SEQ_LEN,
    length_penalty: float = 0.0,
    cache: bool = True,
) -> list[list[tuple[list[int], float]]]:
    """Beam search batché sous masque d'arité.

    Retourne, pour chaque exemple, au plus ``beam`` candidats ``(tokens, score)``
    triés du meilleur au moins bon. Les tokens n'incluent ni ``<bos>`` ni
    ``<eos>``, comme ``greedy_decode``.

    ``length_penalty`` divise la log-vraisemblance par ``len ** alpha``. À 0 le
    score est la log-vraisemblance brute, qui favorise mécaniquement les
    squelettes courts — ce qui n'est pas neutre ici, la faiblesse mesurée du
    modèle étant précisément sur les squelettes profonds.

    ``cache`` active le décodage incrémental (``curvy.infer.cache``) : sans lui,
    chaque pas repasse le préfixe entier au décodeur et recalcule tout ce qui a
    déjà été calculé. Le mettre à ``False`` rejoue le chemin d'origine — c'est
    ce que fait le test d'équivalence.
    """
    device = points.device
    b = points.size(0)
    memory = model.encode_points(points, point_mask)
    mem_dim = memory.size(-1)

    # (B, K, ...) aplati en (B*K, ...) : chaque faisceau est une ligne du batch.
    mem_k = memory.unsqueeze(1).expand(b, beam, memory.size(1), mem_dim)
    mem_k = mem_k.reshape(b * beam, memory.size(1), mem_dim)
    pmask_k = point_mask.unsqueeze(1).expand(b, beam, point_mask.size(1))
    pmask_k = pmask_k.reshape(b * beam, point_mask.size(1))

    states = [DecodeState() for _ in range(b * beam)]
    seqs: list[list[int]] = [[] for _ in range(b * beam)]
    # -inf sur les faisceaux 1..K-1 au premier pas : sans ça, les K faisceaux
    # partent identiques et le top-K rend K fois le même candidat.
    scores = torch.full((b, beam), float("-inf"), device=device)
    scores[:, 0] = 0.0
    vivant = np.zeros((b, beam), dtype=bool)
    vivant[:, 0] = True

    fini: list[list[tuple[list[int], float]]] = [[] for _ in range(b)]
    tokens = torch.full((b * beam, 1), BOS_ID, dtype=torch.long, device=device)
    etat_cache = CacheDecodeur() if cache else None

    for pas in range(max_len - 1):
        if not vivant.any():
            break
        if etat_cache is not None:
            # Seul le dernier token entre : tout le préfixe est dans le cache.
            dernier = tokens[:, -1:] if pas else tokens
            logits = pas_incremental(model, mem_k, pmask_k, dernier, etat_cache)[:, -1]
        else:
            pad_mask = torch.zeros_like(tokens, dtype=torch.bool)
            logits = model.decode(mem_k, pmask_k, tokens, pad_mask)[:, -1]
        logp = torch.log_softmax(logits.float(), dim=-1)

        masks = np.stack([legal_mask(s, max_len=max_len) for s in states])
        allowed = torch.from_numpy(masks).to(device)
        logp = logp.masked_fill(~allowed, float("-inf"))
        # Un faisceau mort ou déjà retiré ne doit pas repeupler le top-K.
        logp = logp.masked_fill(
            ~torch.from_numpy(vivant.reshape(-1)).to(device).unsqueeze(1), float("-inf")
        )

        total = scores.reshape(-1, 1) + logp  # (B*K, V)
        total = total.reshape(b, beam * total.size(-1))
        k_eff = min(beam, total.size(-1))
        meilleurs, plats = torch.topk(total, k_eff, dim=-1)

        n_seqs, n_states, n_scores = [], [], np.zeros((b, beam), dtype=bool)
        # Pour chaque emplacement du pas suivant, l'indice de la ligne dont il
        # descend. Identité par défaut : un emplacement non repourvu garde son
        # cache, qui ne sera de toute façon plus lu.
        provenance = torch.arange(b * beam, dtype=torch.long)
        n_tokens = torch.full(
            (b * beam, tokens.size(1) + 1), BOS_ID, dtype=torch.long, device=device
        )
        nouveaux_scores = torch.full((b, beam), float("-inf"), device=device)

        for i in range(b):
            place = 0
            for rang in range(k_eff):
                sc = float(meilleurs[i, rang])
                if sc == float("-inf"):
                    continue
                plat = int(plats[i, rang])
                src, tid = divmod(plat, logp.size(-1))
                ligne = i * beam + src

                if tid == EOS_ID:
                    seq = list(seqs[ligne])
                    if seq:
                        pen = len(seq) ** length_penalty if length_penalty else 1.0
                        fini[i].append((seq, sc / pen))
                    continue
                if place >= beam:
                    continue
                cible = i * beam + place
                n_seqs.append((cible, [*seqs[ligne], tid]))
                etat = states[ligne].copy()
                etat.advance(tid)
                n_states.append((cible, etat))
                n_tokens[cible, : tokens.size(1)] = tokens[ligne]
                n_tokens[cible, tokens.size(1)] = tid
                provenance[cible] = ligne
                nouveaux_scores[i, place] = sc
                n_scores[i, place] = True
                place += 1

        if etat_cache is not None:
            # Le faisceau `cible` descend du faisceau `source` : son cache doit
            # suivre. Sans ce réagencement chaque faisceau hériterait du passé
            # d'un autre, sans que rien ne le signale — les formes restent
            # valides et les logits plausibles.
            etat_cache.reordonner(provenance.to(device))

        for cible, seq in n_seqs:
            seqs[cible] = seq
        for cible, etat in n_states:
            states[cible] = etat
        for i in range(b):
            for j in range(beam):
                if not n_scores[i, j]:
                    seqs[i * beam + j] = []
                    states[i * beam + j] = DecodeState()
        tokens, scores, vivant = n_tokens, nouveaux_scores, n_scores

    return [sorted(c, key=lambda t: -t[1])[:beam] for c in fini]

"""Ajustement des constantes d'un squelette sur un nuage de points.

Implémente l'idée d'initialisation linéaire sous une
forme un peu plus forte : une **projection variable**.

Tout squelette a la forme ``a * body(x) + b`` (enveloppe de racine).
Pour un jeu de constantes internes donné, ``a`` et ``b`` sont solution exacte
d'une régression linéaire — inutile de les chercher par descente. On n'optimise
donc que les ``k - 2`` constantes internes, et le problème non linéaire perd
deux dimensions. Pour les squelettes les plus simples (`C*x + C`,
`C*sin(x) + C`), il ne reste **rien** à optimiser : la solution est exacte et
immédiate.

``scipy.optimize.least_squares`` sert de moteur ; c'est l'implémentation de
référence, celle contre laquelle un éventuel Levenberg-Marquardt
batché sur GPU devra être validé.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.optimize import least_squares

from curvy.data.expr import Node, count_constants, evaluate

__all__ = ["frequences_candidates", "FitResult", "fit_constants", "r_squared"]

_BIG = 1e6


@dataclass
class FitResult:
    consts: list[float]
    r2_fit: float
    n_restarts_used: int
    ok: bool


def r_squared(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """R² classique. Retourne ``-inf`` si la prédiction n'est pas finie."""
    if not np.isfinite(y_pred).all():
        return float("-inf")
    ss_tot = float(np.sum((y_true - y_true.mean()) ** 2))
    if ss_tot < 1e-15:
        return 1.0 if np.allclose(y_pred, y_true) else float("-inf")
    return 1.0 - float(np.sum((y_true - y_pred) ** 2)) / ss_tot


_C: Node = ("C",)


def _unwrap_root(node: Node) -> Node | None:
    """``body`` si l'arbre a la forme ``add(mul(C, body), C)``, sinon ``None``.

    Tous les squelettes du dataset ont cette forme, mais **pas** ceux
    que produit le décodeur tant que le modèle n'a rien appris : il émet
    n'importe quel arbre syntaxiquement valide. Supposer l'enveloppe était un
    bug — trouvé au premier run de rodage, sur un `mul(x, C)` sans enveloppe.
    """
    if node[0] != "add" or node[2] != _C:
        return None
    scale = node[1]
    if scale[0] != "mul":
        return None
    if scale[1] == _C:
        return scale[2]
    if scale[2] == _C:
        return scale[1]
    return None


def _solve_affine(basis: np.ndarray, y: np.ndarray) -> tuple[float, float, np.ndarray]:
    """Meilleurs ``a, b`` tels que ``a*basis + b ≈ y``, en une seule opération."""
    design = np.stack([basis, np.ones_like(basis)], axis=1)
    coef, *_ = np.linalg.lstsq(design, y, rcond=None)
    a, b = float(coef[0]), float(coef[1])
    return a, b, a * basis + b


def frequences_candidates(x: np.ndarray, y: np.ndarray, n_max: int = 2) -> list[float]:
    """Pulsations dominantes du signal, lues par transformée de Fourier.

    **Pourquoi.** L'initialisation log-uniforme cherchait la fréquence au hasard
    dans ``[0,05 ; 20]`` — la plage de génération. Mesuré le
    2026-08-20 sur ``C·sin(C·x)+C`` avec le squelette **exact** : l'ajustement
    retrouve la fréquence jusqu'à 12 rad, puis décroche (R² 0,058 à 16 rad,
    0,002 à 30). Le paysage d'optimisation d'une fréquence est plein de minima
    locaux ; on ne l'atteint pas par tirage, on doit le viser.

    Or la fréquence est **lisible directement dans les points**. Une FFT sur
    ``y`` rééchantillonné coûte quelques microsecondes et donne le pic. Sur
    ``[-1, 1]``, un signal de ``k`` cycles sur le domaine correspond à une
    pulsation ``ω = π·k``.

    Retourne au plus ``n_max`` pulsations, la plus énergique d'abord. Liste vide
    si le signal n'a pas de pic exploitable — auquel cas on retombe sur le
    tirage aléatoire, qui reste correct en dessous de 12 rad.

    **Résolution.** Les raies d'une FFT sont espacées de ``2π/largeur``, soit
    π ≈ 3,14 rad sur ``[-1, 1]`` : une fréquence tombant entre deux raies serait
    lue à 1,6 rad près. Le pic est donc affiné par **interpolation parabolique**
    sur ses deux voisins — six lignes, et l'erreur tombe sous 0,2 rad. Ça compte
    aux fréquences élevées, où le bassin de convergence est étroit.
    """
    n = len(x)
    if n < 8:
        return []
    ordre = np.argsort(x)
    xs, ys = x[ordre], y[ordre]
    largeur = float(xs[-1] - xs[0])
    if not np.isfinite(largeur) or largeur < 1e-12:
        return []

    # La FFT exige un pas régulier ; les nuages ne le sont pas.
    grille = np.linspace(xs[0], xs[-1], max(64, n))
    yg = np.interp(grille, xs, ys)
    yg = yg - yg.mean()
    if not np.isfinite(yg).all() or float(np.abs(yg).max()) < 1e-12:
        return []

    spectre = np.abs(np.fft.rfft(yg * np.hanning(len(yg))))
    if len(spectre) < 3:
        return []
    spectre[0] = 0.0  # la composante continue est déjà absorbée par l'affine
    pics = np.argsort(spectre)[::-1][:n_max]

    out: list[float] = []
    for k in pics:
        if spectre[k] <= 0.0:
            continue
        out.append(2.0 * np.pi * _affiner_pic(spectre, int(k)) / largeur)
    return [w for w in out if 1e-3 < w < 1e4]


#: Opérateurs dont une constante interne est une pulsation.
_PERIODIQUES = frozenset({"sin", "cos"})


def _contient_periodique(node: Node) -> bool:
    if node[0] in _PERIODIQUES:
        return True
    return any(_contient_periodique(enfant) for enfant in node[1:])


def _affiner_pic(spectre: np.ndarray, k: int) -> float:
    """Position du pic entre les raies, par parabole sur ``k-1, k, k+1``.

    Sur trois points d'une parabole équidistants, le sommet est décalé de
    ``(g - d) / (2·(g - 2c + d))`` raie. Formule classique d'estimation
    spectrale ; on la borne à une demi-raie, au-delà c'est que le pic n'est pas
    parabolique et l'affinage n'a pas de sens.
    """
    # La raie 1 a pour voisine la composante continue, qu'on a mise à zéro :
    # la parabole s'appuierait sur un creux artificiel et tirerait l'estimation
    # vers le haut. Mesuré : 3,0 rad lu 3,64 avec, 3,14 sans.
    if k <= 1 or k >= len(spectre) - 1:
        return float(k)
    g, c, d = float(spectre[k - 1]), float(spectre[k]), float(spectre[k + 1])
    denom = g - 2.0 * c + d
    if abs(denom) < 1e-15:
        return float(k)
    delta = 0.5 * (g - d) / denom
    return float(k) + float(np.clip(delta, -0.5, 0.5))


def fit_constants(
    node: Node,
    x: np.ndarray,
    y: np.ndarray,
    rng: np.random.Generator,
    n_restarts: int = 6,
    max_nfev: int = 120,
    spectral: bool = True,
) -> FitResult:
    """Meilleur jeu de constantes trouvé, avec plusieurs initialisations.

    ``spectral`` ajoute des points de départ tirés d'une FFT des données plutôt
    que du hasard — voir ``frequences_candidates``. Désactivable pour mesurer ce
    qu'il apporte.
    """
    k = count_constants(node)
    body = _unwrap_root(node)
    if body is None or count_constants(body) != k - 2:
        # Arbre hors forme canonique : pas de projection variable possible,
        # on optimise directement toutes les constantes.
        return _fit_generic(node, x, y, rng, k, n_restarts, max_nfev)

    n_inner = k - 2

    with np.errstate(over="ignore", invalid="ignore", divide="ignore"):
        if n_inner == 0:
            basis = evaluate(body, x, [])
            if not np.isfinite(basis).all():
                return FitResult([], float("-inf"), 0, False)
            a, b, pred = _solve_affine(basis, y)
            return FitResult([a, b], r_squared(y, pred), 1, True)

        def residual(inner: np.ndarray) -> np.ndarray:
            basis = evaluate(body, x, list(inner))
            if not np.isfinite(basis).all():
                return np.full_like(y, _BIG)
            _, _, pred = _solve_affine(basis, y)
            return pred - y

        # Points de départ déterministes, essayés avant le tirage aléatoire :
        # le point neutre, puis la pulsation lue dans les données placée tour à
        # tour à chaque position (on ignore laquelle porte la fréquence).
        amorces: list[np.ndarray] = [np.ones(n_inner)]
        # Une pulsation n'a de sens que dans un opérateur périodique. Sur un
        # `inv(x + C)` ou un `exp(C·x)`, l'amorce spectrale est du bruit qui
        # coûte un essai — et l'essai aléatoire qu'elle remplacerait, lui,
        # trouvait. Mesuré : sans ce filtre, +49 % de temps d'ajustement pour un
        # gain concentré sur les seuls squelettes trigonométriques.
        if spectral and _contient_periodique(body):
            for omega in frequences_candidates(x, y):
                for i in range(n_inner):
                    depart = np.ones(n_inner)
                    depart[i] = omega
                    amorces.append(depart)

        # Les amorces spectrales s'AJOUTENT au budget aléatoire, elles ne le
        # remplacent pas. Première version : elles le consommaient, et
        # `hyperbole_raide` — un `inv(x + C)` sans aucune fréquence à lire —
        # est passé de 0,833 à 0,000, cinq tirages utiles ayant cédé la place à
        # trois amorces inutiles. La sortie anticipée sur R² > 0,9999 fait que
        # ce budget élargi ne coûte rien quand une amorce tombe juste.
        budget = len(amorces) + max(0, n_restarts - 1)

        best: tuple[float, list[float]] | None = None
        used = 0
        for attempt in range(budget):
            if attempt < len(amorces):
                x0 = amorces[attempt]
            else:
                # Repli log-uniforme sur la plage de génération.
                mag = np.exp(rng.uniform(np.log(0.05), np.log(20.0), size=n_inner))
                x0 = mag * rng.choice((-1.0, 1.0), size=n_inner)
            used += 1
            try:
                sol = least_squares(residual, x0, max_nfev=max_nfev, method="lm")
            except Exception:
                continue
            basis = evaluate(body, x, list(sol.x))
            if not np.isfinite(basis).all():
                continue
            a, b, pred = _solve_affine(basis, y)
            score = r_squared(y, pred)
            if best is None or score > best[0]:
                best = (score, [a, *sol.x.tolist(), b])
            if best[0] > 0.9999:  # inutile de continuer à chercher
                break

    if best is None:
        return FitResult([], float("-inf"), used, False)
    return FitResult(best[1], best[0], used, True)


def _fit_generic(
    node: Node,
    x: np.ndarray,
    y: np.ndarray,
    rng: np.random.Generator,
    k: int,
    n_restarts: int,
    max_nfev: int,
) -> FitResult:
    """Ajustement sans projection variable, pour un arbre hors forme canonique."""
    with np.errstate(over="ignore", invalid="ignore", divide="ignore"):
        if k == 0:
            pred = evaluate(node, x, [])
            return FitResult([], r_squared(y, pred), 1, np.isfinite(pred).all())

        def residual(c: np.ndarray) -> np.ndarray:
            pred = evaluate(node, x, list(c))
            if not np.isfinite(pred).all():
                return np.full_like(y, _BIG)
            return pred - y

        best: tuple[float, list[float]] | None = None
        used = 0
        for attempt in range(n_restarts):
            if attempt == 0:
                x0 = np.ones(k)
            else:
                mag = np.exp(rng.uniform(np.log(0.05), np.log(20.0), size=k))
                x0 = mag * rng.choice((-1.0, 1.0), size=k)
            used += 1
            try:
                sol = least_squares(residual, x0, max_nfev=max_nfev, method="lm")
            except Exception:
                continue
            pred = evaluate(node, x, list(sol.x))
            score = r_squared(y, pred)
            if best is None or score > best[0]:
                best = (score, sol.x.tolist())
            if best[0] > 0.9999:
                break

    if best is None:
        return FitResult([], float("-inf"), used, False)
    return FitResult(best[1], best[0], used, True)


def estimer_bruit(x: np.ndarray, y: np.ndarray) -> float:
    """Écart-type du bruit, estimé **sans connaître la fonction**.

    Pseudo-résidus de Gasser, Sroka et Jennen-Steinmetz (1986) : pour chaque
    point intérieur, on interpole ses deux voisins et on regarde de combien le
    point s'en écarte. Sur une fonction lisse, cet écart est dominé par le
    bruit ; la courbure n'y contribue qu'au second ordre.

    ``ε_i = a_i·y_{i-1} + b_i·y_{i+1} − y_i`` avec les poids de l'interpolation
    linéaire aux abscisses réelles, puis ``σ² = moyenne(c_i²·ε_i²)`` où ``c_i``
    normalise la variance. **La pondération par les espacements est
    indispensable ici** : les nuages du projet ne sont jamais à pas régulier
    (densité liée à la courbure, trous, jitter), et la différence
    seconde naïve prendrait l'irrégularité pour du bruit.

    Retourne ``0.0`` si l'estimation n'a pas de sens (trop peu de points,
    abscisses confondues).
    """
    n = len(x)
    if n < 5:
        return 0.0
    ordre = np.argsort(x)
    xs, ys = np.asarray(x, dtype=float)[ordre], np.asarray(y, dtype=float)[ordre]

    dx_g = xs[1:-1] - xs[:-2]
    dx_d = xs[2:] - xs[1:-1]
    ecart = dx_g + dx_d
    valide = ecart > 1e-12
    if not valide.any():
        return 0.0

    a = np.where(valide, dx_d / np.where(valide, ecart, 1.0), 0.0)
    b = np.where(valide, dx_g / np.where(valide, ecart, 1.0), 0.0)
    eps = a * ys[:-2] + b * ys[2:] - ys[1:-1]
    c2 = 1.0 / (a**2 + b**2 + 1.0)
    val = c2[valide] * eps[valide] ** 2
    if not val.size or not np.isfinite(val).all():
        return 0.0
    return float(np.sqrt(max(0.0, val.mean())))

"""Des candidats du beam search aux 3-5 formules livrées (Phase 5).

Trois étapes, dans cet ordre :

1. **ajuster** les constantes de chaque candidat sur les points observés ;
2. **sélectionner** celui qu'on annonce comme réponse principale ;
3. **retenir le front de Pareto** — les candidats qu'aucun autre ne domine à la
   fois en simplicité et en précision. C'est ce que voit l'utilisateur.

Une règle gouverne tout le module : **la sélection ne regarde jamais les points
tenus à l'écart.** Un candidat est jugé sur son ajustement aux points observés,
comme à l'usage réel où il n'y a pas de vérité terrain. C'est exactement le
biais reproché à la baseline polynomiale le 2026-08-19 — choisir son degré
d'après le résultat final — et il serait malhonnête de se l'autoriser ici.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from curvy.data.expr import Node, complexity
from curvy.infer.fit import estimer_bruit, fit_constants

__all__ = [
    "Candidat",
    "ajuster_candidats",
    "front_de_pareto",
    "selectionner",
    "selectionner_selon_bruit",
]


@dataclass
class Candidat:
    node: Node
    consts: list[float]
    r2_fit: float
    complexite: int
    score_modele: float = 0.0

    @property
    def n_consts(self) -> int:
        return len(self.consts)


def _ajuster_un(args):
    """Un ajustement isolé, picklable — cible du pool de processus."""
    node, x, y, graine = args
    from curvy.seeding import make_rng

    res = fit_constants(node, x, y, make_rng(graine))
    return (list(res.consts), float(res.r2_fit), bool(res.ok))


def ajuster_candidats(
    nodes: list[Node],
    x: np.ndarray,
    y: np.ndarray,
    rng: np.random.Generator,
    scores: list[float] | None = None,
    executor=None,
) -> list[Candidat]:
    """Ajuste chaque candidat sur ``(x, y)`` observés. Les échecs sont écartés.

    ``executor`` — un pool de processus optionnel. Les ajustements sont
    **indépendants** et dominent la latence : les paralléliser est le
    seul moyen d'élargir le beam sans quitter le budget d'une seconde. Mesuré à
    beam 24 : 1015 ms en série. Les threads n'aideraient pas — la fonction de
    résidu est du Python, donc tenue par le GIL.

    Sans pool, le comportement est strictement celui d'avant : c'est le chemin
    des tests et des scripts de mesure.
    """
    valides = [(i, n) for i, n in enumerate(nodes) if n is not None]
    if not valides:
        return []

    if executor is None:
        resultats = []
        for _, node in valides:
            r = fit_constants(node, x, y, rng)
            resultats.append((list(r.consts), float(r.r2_fit), bool(r.ok)))
    else:
        # Une graine par candidat, dérivée du rng appelant : le parallélisme ne
        # doit pas rendre le résultat dépendant de l'ordre d'arrivée.
        graines = [int(g) for g in rng.integers(0, 2**31 - 1, size=len(valides))]
        taches = [(n, x, y, g) for (_, n), g in zip(valides, graines, strict=True)]
        resultats = list(executor.map(_ajuster_un, taches))

    out: list[Candidat] = []
    for (i, node), (consts, r2, ok) in zip(valides, resultats, strict=True):
        if not ok or not np.isfinite(r2):
            continue
        out.append(
            Candidat(
                node=node,
                consts=consts,
                r2_fit=r2,
                complexite=complexity(node),
                score_modele=float(scores[i]) if scores is not None else 0.0,
            )
        )
    return out


def selectionner(cands: list[Candidat], tol: float = 0.0) -> Candidat | None:
    """La réponse principale : le meilleur ajustement, à parcimonie égale.

    Parmi les candidats dont le R² d'ajustement est à ``tol`` du meilleur, on
    retient **le plus simple**.

    ``tol = 0`` par défaut, et c'est une **mesure, pas une intuition**. J'avais
    prévu qu'une tolérance protégerait du sur-ajustement au bruit — choisir le
    maximum parmi 16 candidats, c'est en principe choisir celui qui épouse le
    mieux le bruit. Mesuré le 2026-08-20 sur 512 exemples, c'est faux : la
    tolérance ne fait que coûter, et sur les deux métriques à la fois.

    ===========  ==============  ==============
    tol          interpolation   extrapolation
    ===========  ==============  ==============
    0            **0,719**       **0,262**
    0,002        0,709           0,238
    0,005        0,691           0,229
    0,02         0,606           0,203
    ===========  ==============  ==============

    L'explication tient sans doute au masque d'arité : profondeur bornée,
    5 constantes au plus, vocabulaire de 18 tokens. Le beam n'a pas de quoi
    fabriquer un candidat assez tordu pour épouser le bruit — la grammaire fait
    déjà le travail de régularisation qu'on croyait devoir refaire ici.
    Le départage par complexité reste actif en cas d'égalité exacte.
    """
    if not cands:
        return None
    meilleur = max(c.r2_fit for c in cands)
    proches = [c for c in cands if c.r2_fit >= meilleur - tol]
    return min(proches, key=lambda c: (c.complexite, -c.r2_fit))


def selectionner_selon_bruit(
    cands: list[Candidat],
    x: np.ndarray,
    y: np.ndarray,
    marge: float = 1.0,
) -> Candidat | None:
    """Le plus simple des candidats dont le résidu est compatible avec le bruit.

    **L'hypothèse.** Prendre le maximum de R² d'ajustement revient, par
    construction, à retenir le candidat qui **épouse le mieux le bruit** : au-delà
    du niveau de bruit, tout R² supplémentaire est du sur-ajustement. La bonne
    règle serait donc « parmi ceux qui expliquent les données *jusqu'au bruit*,
    le plus simple ».

    Une tolérance **fixe** avait été essayée le 2026-08-20 et dégradait tout
    (0,691 contre 0,711 à 0,005). C'était attendu après coup : chaque tracé a son
    propre niveau de bruit, une tolérance uniforme est forcément trop large sur
    un nuage propre et trop étroite sur un nuage bruité. Ici le seuil est **tiré
    des données elles-mêmes**, par les pseudo-résidus de `estimer_bruit`.

    ``marge`` fixe la tolérance autour du bruit estimé : un candidat est accepté
    si son résidu quadratique moyen ne dépasse pas ``σ·(1 + marge)``.

    Retombe sur le maximum de R² si le bruit n'est pas estimable ou si aucun
    candidat n'atteint le niveau du bruit — auquel cas aucun ne « suffit », et
    prendre le meilleur reste le moins mauvais choix.
    """
    if not cands:
        return None
    sigma = estimer_bruit(x, y)
    variance = float(np.var(y))
    if sigma <= 0.0 or variance <= 1e-15:
        return selectionner(cands)

    # r2 = 1 − résidu²/variance : le seuil de bruit se convertit en seuil de R².
    seuil_r2 = 1.0 - (sigma * (1.0 + marge)) ** 2 / variance
    suffisants = [c for c in cands if c.r2_fit >= seuil_r2]
    if not suffisants:
        return selectionner(cands)
    return min(suffisants, key=lambda c: (c.complexite, -c.r2_fit))


def front_de_pareto(cands: list[Candidat]) -> list[Candidat]:
    """Candidats non dominés, du plus simple au plus précis.

    ``a`` domine ``b`` si ``a`` est au moins aussi simple **et** au moins aussi
    précis, avec un avantage strict quelque part. Les doublons de squelette sont
    écartés en amont : deux fois la même formule ne fait pas deux propositions.
    """
    vus: dict[tuple, Candidat] = {}
    for c in cands:
        cle = tuple(_aplatir(c.node))
        if cle not in vus or c.r2_fit > vus[cle].r2_fit:
            vus[cle] = c
    uniques = sorted(vus.values(), key=lambda c: (c.complexite, -c.r2_fit))

    front: list[Candidat] = []
    meilleur_r2 = -np.inf
    for c in uniques:
        if c.r2_fit > meilleur_r2:
            front.append(c)
            meilleur_r2 = c.r2_fit
    return front


def _aplatir(node: Node) -> list[str]:
    out = [node[0]]
    for enfant in node[1:]:
        out.extend(_aplatir(enfant))
    return out

"""Rendre une formule lisible, avec ses constantes et dans les bonnes unités.

Deux besoins que le squelette ne couvre pas :

1. **les valeurs.** `C * sin(C * x) + C` décrit une famille ; l'utilisateur veut
   `2.31*sin(4.07*x) - 0.5` ;
2. **les unités.** Les constantes sont ajustées en coordonnées **normalisées** —
   `x` ramené dans `[-1, 1]`, `y` centré-réduit. Rendre la formule
   telle quelle serait **faux** pour qui a fourni des mesures : elle ne
   s'appliquerait pas à ses abscisses. Il faut la composer avec les deux
   affines.

Pour un tracé au canvas la question ne se pose pas — les pixels ne sont pas une
unité qui intéresse quelqu'un. Elle se pose dès qu'on importe des données.

sympy sert ici et **seulement ici** : à mettre en forme une expression déjà
décidée. Il reste écarté de la canonicalisation, où il est sémantiquement faux
(cf. l'en-tête de `curvy.data.canonical` : il « simplifierait » `C*x + C*x` en
`2*C*x`, alors que les deux `C` sont indépendants).
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from curvy.data.expr import Node, const_name_iter, to_infix

__all__ = ["Affines", "formule_lisible"]


@dataclass(frozen=True)
class Affines:
    """Les transformations appliquées avant l'ajustement.

    ``x_norm = (x - x_centre) / x_demi`` et ``y = y_norm * y_echelle + y_decalage``.
    """

    x_centre: float = 0.0
    x_demi: float = 1.0
    y_echelle: float = 1.0
    y_decalage: float = 0.0

    @property
    def identite(self) -> bool:
        return (
            self.x_centre == 0.0
            and self.x_demi == 1.0
            and self.y_echelle == 1.0
            and self.y_decalage == 0.0
        )


def _arrondi_lisible(v: float, chiffres: int = 4) -> float:
    """Arrondit sans écraser les petites valeurs.

    ``round(1.2e-5, 4)`` vaut ``0.0``, ce qui transformerait une constante en
    zéro et changerait la formule. On arrondit donc à un nombre de **chiffres
    significatifs**, pas de décimales.
    """
    if v == 0.0 or not np.isfinite(v):
        return float(v)
    return float(f"%.{chiffres}g" % v)


def formule_lisible(
    node: Node,
    consts: list[float],
    affines: Affines | None = None,
    chiffres: int = 4,
) -> str:
    """L'expression avec ses constantes, dans les unités d'origine si fournies.

    Retombe sur un rendu textuel simple si sympy échoue — une mise en forme
    ratée ne doit jamais faire échouer une prédiction correcte.
    """
    noms = const_name_iter()
    infixe = to_infix(node, noms)
    valeurs = {f"c{i}": _arrondi_lisible(v, chiffres) for i, v in enumerate(consts)}

    try:
        import sympy

        x = sympy.Symbol("x")
        expr = sympy.sympify(infixe, locals={"x": x})
        expr = expr.subs({sympy.Symbol(k): sympy.Float(v) for k, v in valeurs.items()})

        if affines is not None and not affines.identite:
            # x du modèle = (x réel - centre) / demi, puis y réel = y*échelle + décalage
            expr = expr.subs(x, (x - sympy.Float(affines.x_centre)) / sympy.Float(affines.x_demi))
            expr = expr * sympy.Float(affines.y_echelle) + sympy.Float(affines.y_decalage)

        # Surtout PAS de `nsimplify` : il rationalise les flottants et rend
        # `231*sin(407*x/100)/100 - 1/2` là où on veut `2.31*sin(4.07*x) - 0.5`.
        if expr.count_ops() < 40:
            expr = sympy.expand(expr)

        # `evalf` n'arrondit que le résultat des opérations, pas les constantes
        # déjà présentes : la composition avec les affines laissait sortir
        # `exp(-0.0454533333333333*x)`. On arrondit chaque flottant de l'arbre.
        flottants = expr.atoms(sympy.Float)
        # Un terme négligeable devant les autres est du bruit d'arithmétique
        # flottante, pas une constante : `4.9*x^2 - 1.776e-15` doit se lire
        # `4.9*x^2`. Le critère est RELATIF — une donnée à l'échelle du
        # micromètre a des constantes légitimement minuscules.
        echelle = max((abs(float(f)) for f in flottants), default=0.0)
        seuil = 1e-10 * echelle
        expr = expr.xreplace(
            {
                f: sympy.Float(
                    0.0 if abs(float(f)) < seuil else _arrondi_lisible(float(f), chiffres)
                )
                for f in flottants
            }
        )
        rendu = str(expr.evalf(chiffres))
        return rendu.replace("**", "^")
    except Exception:
        # Repli : substitution textuelle, toujours juste même si moins jolie.
        rendu = infixe
        for nom, v in valeurs.items():
            rendu = rendu.replace(nom, repr(v))
        return rendu.replace("**", "^")

"""Service d'inférence : du tracé au front de Pareto (Phase 8)."""

"""Serveur de la démo : une page, deux endpoints, zéro dépendance ajoutée.

`http.server` de la bibliothèque standard plutôt que FastAPI : le projet tient
sur cinq dépendances, et en ajouter trois pour un endpoint JSON et une page HTML
serait disproportionné. `make demo` doit marcher sans rien installer.

**Ce service n'écrit rien sur disque.** Il est destiné à être exposé en ligne
avec l'article, sans authentification ni limite de débit : toute écriture serait
une croissance non bornée offerte au premier robot venu — un tracé pesant 3,7
Kio, dix requêtes par seconde rempliraient un gigaoctet en huit heures. Un
utilisateur qui veut garder son tracé le **télécharge** depuis son navigateur.

**Politique GPU.** Le service tourne sur la 3060, jamais sur
la 3090 d'entraînement, et il **plafonne explicitement sa VRAM** : la carte est
partagée, on ne présume pas qu'elle est vide, et un OOM provoqué par le voisin
doit rendre une erreur, pas faire tomber le serveur.

    .venv/bin/python -m curvy.serve.app --run exp-005 --preset v1
"""

from __future__ import annotations

import argparse
import json
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import numpy as np

from curvy.config import REPO_ROOT, RUNS_DIR
from curvy.seeding import make_rng
from curvy.serve.debit import Limiteur
from curvy.serve.pipeline import formules_depuis_trace

#: UUID de la 3060 : carte de service. Jamais la 3090.
#: Carte réservée au service. À renseigner avec la vôtre (`nvidia-smi -L`).
GPU_SERVICE = "GPU-xxxxxxxx-CHANGEZ-MOI"

#: Fraction de la VRAM que le service s'autorise. La carte héberge d'autres
#: travaux ; on ne prend pas ce dont on n'a pas besoin.
FRACTION_VRAM = 0.35

WEB = REPO_ROOT / "web"


def _urls(host: str, port: int) -> list[str]:
    """Les adresses réellement tapables, pas celle passée en argument.

    `0.0.0.0` n'est pas une adresse qu'on met dans un navigateur : afficher
    l'adresse de liaison telle quelle envoie l'utilisateur dans le mur.
    """
    if host != "0.0.0.0":  # noqa: S104
        return [f"http://{host}:{port}/"]
    import socket

    urls = [f"http://127.0.0.1:{port}/"]
    try:
        # Une connexion UDP n'émet aucun paquet : elle sert seulement à demander
        # au noyau par quelle interface il sortirait. Adresse hors du réseau
        # local pour ne présumer d'aucun plan d'adressage.
        prise = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        prise.connect(("203.0.113.1", 1))  # RFC 5737, réservée à la documentation
        urls.append(f"http://{prise.getsockname()[0]}:{port}/")
        prise.close()
    except OSError:
        pass
    return urls


class Etat:
    """Modèle chargé une fois, partagé par les requêtes."""

    modele = None
    device = None
    rng = None
    beam = 8
    #: Pool d'ajustement. Les ajustements de constantes dominent la latence
    #: et sont indépendants : les paralléliser est ce qui permet
    #: d'élargir le beam sans quitter le budget d'une seconde.
    pool = None
    limiteur = None
    #: Origine autorisée pour les appels croisés. Vide = aucun en-tête CORS,
    #: donc seule une page servie par ce serveur peut l'appeler.
    origine = ""
    #: Adresses depuis lesquelles `X-Forwarded-For` est cru. Vide = aucune.
    proxys: frozenset[str] = frozenset()


def charger(run: str, preset: str, checkpoint: str, beam: int, limiteur=None) -> None:
    import os

    import torch

    os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
    os.environ.setdefault("CURVY_CUDA_ALLOW", GPU_SERVICE)

    from curvy.devices import pick_device
    from curvy.model.config import PRESETS
    from curvy.model.curvy import CurvyModel

    choix = pick_device("auto")
    Etat.device = choix.device
    if Etat.device.type == "cuda":
        torch.cuda.set_per_process_memory_fraction(FRACTION_VRAM, Etat.device)
    Etat.modele = CurvyModel(PRESETS[preset]).to(Etat.device)
    etat = torch.load(RUNS_DIR / run / checkpoint, map_location="cpu", weights_only=False)
    Etat.modele.load_state_dict(etat["model"])
    Etat.modele.eval()
    Etat.rng = make_rng(0)
    Etat.beam = beam
    Etat.limiteur = limiteur

    # `spawn` et non `fork` : le processus parent porte un contexte CUDA, qu'un
    # fork dupliquerait dans un état indéfini. Les enfants ne font que du calcul
    # numpy/scipy et n'ont pas besoin du GPU.
    import multiprocessing
    from concurrent.futures import ProcessPoolExecutor

    n_procs = max(2, min(8, (os.cpu_count() or 4) - 2))
    Etat.pool = ProcessPoolExecutor(
        max_workers=n_procs, mp_context=multiprocessing.get_context("spawn")
    )
    print(f"pool    : {n_procs} processus d'ajustement", flush=True)
    print(f"modèle  : {run}/{checkpoint}, step {etat.get('step', '?')}", flush=True)
    print(f"device  : {choix}", flush=True)
    print(f"VRAM    : plafonnée à {FRACTION_VRAM:.0%} de la carte", flush=True)


class Handler(BaseHTTPRequestHandler):
    #: Coût de chaque endpoint, en jetons. Une prédiction mobilise le GPU
    #: pendant ~779 ms ; un rééchantillonnage coûte 4 ms et ne touche pas la
    #: carte. Les facturer pareil reviendrait soit à brider l'exploration de la
    #: courbe, soit à laisser passer les requêtes coûteuses.
    COUTS = {"/api/formules": 1.0, "/api/courbe": 0.1}

    #: Coût d'un fichier statique. `do_GET` était la seule route non limitée :
    #: exposé derrière un tunnel public, un robot pouvait boucler dessus sans
    #: rien consommer. Le tarif est bas (une page en fait ~25 avant d'attendre)
    #: parce que la page est autonome — un affichage = une seule requête.
    COUT_STATIQUE = 0.2

    def _adresse(self) -> str:
        """L'adresse du client, en tenant compte d'un proxy de confiance.

        **Sans ceci, la limitation par adresse s'effondre derrière un reverse
        proxy ou un tunnel** : toutes les requêtes arrivent de `127.0.0.1`, donc
        le monde entier partage un seul seau — un visiteur suffirait à bloquer
        tous les autres.

        `X-Forwarded-For` est fourni par le client et **ne doit jamais être cru
        sans condition** : n'importe qui pourrait s'inventer une adresse par
        requête et contourner le quota. On ne l'honore donc que si la connexion
        vient d'un proxy explicitement déclaré, et on prend la **dernière**
        valeur de la chaîne — celle ajoutée par ce proxy, la seule qu'il n'a pas
        recopiée du client.
        """
        directe = self.client_address[0] if self.client_address else "inconnu"
        if directe not in Etat.proxys:
            return directe
        chaine = self.headers.get("X-Forwarded-For", "")
        maillons = [m.strip() for m in chaine.split(",") if m.strip()]
        return maillons[-1] if maillons else directe

    def _debit_ok(self, cout: float | None = None) -> bool:
        """Vrai si la requête passe. Répond 429 elle-même sinon."""
        if Etat.limiteur is None:
            return True
        if cout is None:
            cout = self.COUTS.get(self.path, 1.0)
        attente = Etat.limiteur.verifier(self._adresse(), time.monotonic(), cout)
        if attente <= 0.0:
            return True
        secondes = max(1, int(attente + 0.999))
        corps = json.dumps(
            {
                "erreur": "trop de requêtes — le service est partagé",
                "reessayer_dans_s": secondes,
            },
            ensure_ascii=False,
        ).encode("utf-8")
        self.send_response(429)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Retry-After", str(secondes))
        self.send_header("Content-Length", str(len(corps)))
        self.end_headers()
        self.wfile.write(corps)
        return False

    def _cors(self) -> None:
        """Autorise l'appel depuis une page hébergée ailleurs.

        Nécessaire si la page vit sur un site personnel et que le service
        tourne sur une autre machine. L'origine autorisée est un paramètre :
        `*` ouvre à tous, ce qui est acceptable pour un service public sans
        authentification ni écriture, mais doit rester un choix explicite.
        """
        if Etat.origine:
            self.send_header("Access-Control-Allow-Origin", Etat.origine)
            self.send_header("Access-Control-Allow-Headers", "Content-Type")
            self.send_header("Access-Control-Max-Age", "86400")

    def do_OPTIONS(self) -> None:  # noqa: N802
        """Requête préalable du navigateur avant un appel croisé."""
        self.send_response(204)
        self._cors()
        self.send_header("Access-Control-Allow-Methods", "POST, GET, OPTIONS")
        self.send_header("Content-Length", "0")
        self.end_headers()

    def _json(self, code: int, charge: dict) -> None:
        corps = json.dumps(charge, ensure_ascii=False).encode("utf-8")
        self.send_response(code)
        self._cors()
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(corps)))
        self.end_headers()
        self.wfile.write(corps)

    def do_GET(self) -> None:  # noqa: N802
        if not self._debit_ok(self.COUT_STATIQUE):
            return
        chemin = "index.html" if self.path in ("/", "") else self.path.lstrip("/")
        fichier = (WEB / chemin).resolve()
        if not fichier.is_file() or WEB.resolve() not in fichier.parents:
            self._json(404, {"erreur": "introuvable"})
            return
        corps = fichier.read_bytes()
        types = {".html": "text/html; charset=utf-8", ".js": "text/javascript", ".css": "text/css"}
        self.send_response(200)
        self.send_header("Content-Type", types.get(fichier.suffix, "application/octet-stream"))
        self.send_header("Content-Length", str(len(corps)))
        self.end_headers()
        self.wfile.write(corps)

    def do_POST(self) -> None:  # noqa: N802
        # Le débit est vérifié AVANT de lire le corps : sinon un robot ferait
        # transiter des mégaoctets pour se voir refuser ensuite.
        if not self._debit_ok():
            return
        taille = int(self.headers.get("Content-Length", 0))
        if taille > 2_000_000:
            self._json(413, {"erreur": "tracé trop volumineux"})
            return
        try:
            charge = json.loads(self.rfile.read(taille) or b"{}")
            points = np.asarray(charge.get("points", []), dtype=float)
        except Exception as exc:
            self._json(400, {"erreur": f"corps illisible : {exc}"})
            return
        if self.path != "/api/courbe" and (points.ndim != 2 or points.shape[1] != 2):
            self._json(400, {"erreur": "attendu : une liste de couples [x, y]"})
            return

        if self.path == "/api/courbe":
            self._courbe(charge)
            return
        if self.path != "/api/formules":
            self._json(404, {"erreur": "endpoint inconnu"})
            return

        try:
            # `source: "donnees"` : l'ordonnée n'est pas retournée (un fichier de
            # mesures n'a pas l'axe inversé d'un canvas) et la formule est rendue
            # dans les unités fournies, pas dans les coordonnées normalisées.
            importe = charge.get("source") == "donnees"
            rep = formules_depuis_trace(
                points,
                Etat.modele,
                Etat.device,
                Etat.rng,
                beam=Etat.beam,
                retourner_y=not importe,
                executor=Etat.pool,
            )
        except Exception as exc:  # un OOM du voisin ne doit pas tuer le serveur
            self._json(503, {"erreur": f"{type(exc).__name__} : {exc}"})
            return

        self._json(
            200,
            {
                "ok": rep.ok,
                "raison": rep.raison,
                "univalue": rep.univalue,
                "n_points": rep.n_points,
                "latence_ms": rep.latence_ms,
                "y_scale": rep.y_scale,
                "y_offset": rep.y_offset,
                "x_min": rep.x_min,
                "x_max": rep.x_max,
                "domaine_apercu": rep.domaine_apercu,
                "formules": [
                    {
                        "expression": f.expression,
                        "valorisee": f.valorisee,
                        "complexite": f.complexite,
                        "r2": f.r2,
                        "constantes": f.constantes,
                        "prefixe": f.prefixe,
                        "principale": f.principale,
                        "apercu": f.apercu,
                    }
                    for f in rep.formules
                ],
            },
        )

    #: Bornes du rééchantillonnage à la demande. Le domaine est plafonné parce
    #: qu'une formule évaluée trop loin ne donne plus que des infinis, et que
    #: rien ne justifie de calculer ce que personne ne peut lire.
    DOMAINE_MAX = 1000.0
    N_COURBE_MAX = 2000

    def _courbe(self, charge: dict) -> None:
        """Rééchantillonne une formule déjà trouvée sur un domaine plus large.

        Le navigateur en a besoin dès qu'on dézoome au-delà de l'aperçu initial.
        Aucun modèle, aucun ajustement : juste une évaluation, donc quelques
        centaines de microsecondes.

        **C'est le serveur qui évalue, pas le navigateur.** Réimplémenter la
        grammaire en JavaScript créerait deux vérités qui divergeraient au
        premier opérateur ajouté — même raison que pour l'aperçu initial.
        """
        from curvy.data.expr import evaluate, from_prefix

        try:
            noeud = from_prefix(str(charge.get("prefixe", "")).split())
            consts = [float(v) for v in charge.get("constantes", [])]
            domaine = min(abs(float(charge.get("domaine", 3.0))), self.DOMAINE_MAX)
            n = min(int(charge.get("n", 600)), self.N_COURBE_MAX)
        except Exception as exc:
            self._json(400, {"erreur": f"requête de courbe illisible : {exc}"})
            return
        if domaine <= 0 or n < 2:
            self._json(400, {"erreur": "domaine ou nombre de points invalide"})
            return

        grille = np.linspace(-domaine, domaine, n)
        with np.errstate(over="ignore", invalid="ignore", divide="ignore"):
            try:
                y = evaluate(noeud, grille, consts)
            except Exception as exc:
                self._json(400, {"erreur": f"évaluation impossible : {exc}"})
                return
        self._json(
            200,
            {
                "domaine": domaine,
                "apercu": [None if not np.isfinite(v) else round(float(v), 5) for v in y],
            },
        )

    def log_message(self, fmt: str, *args) -> None:
        print(f"{self.address_string()} {fmt % args}", flush=True)


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--run", default="exp-005")
    ap.add_argument("--preset", default="v1")
    ap.add_argument("--checkpoint", default="best.pt")
    #: 48, rendu abordable par le cache clé/valeur (514 ms contre 820 sans).
    #: Progression mesurée du taux rendu : 0,693 à beam 8, 0,719 à 24, **0,742
    #: à 48**. Et surtout, le poste « aucun candidat n'était bon » tombe de 5,7
    #: à **0,6 point** : le rappel atteint 0,811 pour un oracle à 0,816. Le
    #: modèle propose donc déjà le plafond de la tâche ; tout ce qui manque est
    #: dans la sélection.
    ap.add_argument("--beam", type=int, default=48)
    #: 8001 et non 8000 : le port 8000 est le défaut de trop d'outils, et la
    #: collision se manifeste par un « Address already in use » au démarrage.
    ap.add_argument("--port", type=int, default=8001)
    #: Par défaut on n'écoute QUE en local. Exposer un serveur d'inférence sur
    #: le réseau est une décision, pas un réglage par défaut.
    ap.add_argument("--host", default="127.0.0.1")
    #: Débit par adresse : rafale puis régime permanent. Les défauts laissent
    #: passer cinq requêtes d'affilée — un humain qui dessine plusieurs courbes
    #: — puis une toutes les deux secondes.
    ap.add_argument("--rafale", type=float, default=5.0)
    ap.add_argument("--par-seconde", type=float, default=0.5)
    #: Débit global. Le GPU soutient environ 1,3 prédiction par seconde ; on
    #: plafonne un peu au-dessus pour que la file se vide, pas pour la remplir.
    ap.add_argument("--rafale-globale", type=float, default=20.0)
    ap.add_argument("--par-seconde-global", type=float, default=2.0)
    ap.add_argument("--sans-limite", action="store_true", help="désactive la limitation")
    ap.add_argument(
        "--origine",
        default="",
        help="origine autorisée pour les appels croisés, ex. https://exemple.fr ou *",
    )
    ap.add_argument(
        "--proxy-de-confiance",
        default="",
        help=(
            "adresses d'un reverse proxy ou d'un tunnel, séparées par des virgules, "
            "dont l'en-tête X-Forwarded-For sera cru — ex. 127.0.0.1. "
            "SANS CELA, la limitation par adresse ne distingue plus les visiteurs."
        ),
    )
    args = ap.parse_args(argv)

    limiteur = (
        None
        if args.sans_limite
        else Limiteur(
            par_adresse=(args.rafale, args.par_seconde),
            global_=(args.rafale_globale, args.par_seconde_global),
        )
    )
    Etat.origine = args.origine
    Etat.proxys = frozenset(a.strip() for a in args.proxy_de_confiance.split(",") if a.strip())
    charger(args.run, args.preset, args.checkpoint, args.beam, limiteur)
    serveur = ThreadingHTTPServer((args.host, args.port), Handler)
    if Etat.proxys:
        print(f"proxy   : X-Forwarded-For cru depuis {sorted(Etat.proxys)}", flush=True)
    elif args.host == "127.0.0.1":
        print(
            "proxy   : aucun. Si un reverse proxy est devant, ajouter\n"
            "          --proxy-de-confiance 127.0.0.1, sinon tous les visiteurs\n"
            "          partageront un seul quota.",
            flush=True,
        )
    if limiteur is None:
        print("débit   : AUCUNE LIMITE — à réserver à un usage local", flush=True)
    else:
        print(
            f"débit   : {args.rafale:.0f} en rafale puis {args.par_seconde:.2g}/s par adresse ; "
            f"{args.rafale_globale:.0f} puis {args.par_seconde_global:.2g}/s au total",
            flush=True,
        )
    print("", flush=True)
    for url in _urls(args.host, args.port):
        print(f"démo : {url}", flush=True)
    if args.host == "0.0.0.0":  # noqa: S104
        print(
            "\nOUVERT AU RÉSEAU. Le service n'écrit rien sur disque et le débit est\n"
            "  limité, mais il n'y a AUCUNE AUTHENTIFICATION : qui atteint ce port\n"
            "  consomme le GPU, dans la limite du quota.\n"
            "  Le pare-feu devrait n'autoriser que le sous-réseau voulu, plutôt que\n"
            "  toutes les interfaces de la machine :\n"
            f"    sudo ufw allow from <votre-sous-réseau>/24 to any port {args.port} proto tcp",
            flush=True,
        )
    # Un SIGTERM n'exécute aucun `finally` : sans ce gestionnaire, les
    # processus du pool survivent au serveur, rattachés à init. Constaté le
    # 2026-08-22 — huit orphelins à 534 Mio, soit 4,3 Gio, après un simple
    # `kill`. Sur un service qu'on redémarre, la fuite s'accumule.
    import signal

    def _arreter(signum, cadre):  # noqa: ARG001
        print(f"\narrêt (signal {signum})", flush=True)
        raise KeyboardInterrupt

    signal.signal(signal.SIGTERM, _arreter)
    signal.signal(signal.SIGINT, _arreter)

    try:
        serveur.serve_forever()
    except KeyboardInterrupt:
        print("arrêt", flush=True)
    finally:
        if Etat.pool is not None:
            Etat.pool.shutdown(wait=True, cancel_futures=True)
            print("pool arrêté", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

"""Limitation de débit pour le service public.

Le service est destiné à être mis en ligne avec l'article, sans
authentification. Il n'écrit plus rien sur disque (journal du 2026-08-22), donc
le risque n'est plus la croissance : c'est la **disponibilité**. Une requête
`/api/formules` coûte ~779 ms de GPU ; un robot à dix requêtes par seconde
monopolise la 3060 et la démo ne répond plus à personne.

**Deux niveaux, parce qu'ils protègent contre deux choses différentes.**

- **par adresse** — empêche un client seul de tout prendre ;
- **global** — protège la carte quand la charge vient de partout à la fois. Sans
  lui, cent adresses respectant chacune leur quota suffiraient à saturer le GPU.

**Le seau à jetons** plutôt qu'un compteur par fenêtre : il autorise une petite
rafale — un humain qui dessine trois courbes d'affilée ne doit pas être bloqué —
tout en bornant le débit moyen. Un compteur par fenêtre laisse au contraire
passer deux fois le quota à cheval sur deux fenêtres.

**La mémoire est bornée.** Un dictionnaire indexé par adresse grandit avec le
nombre d'adresses distinctes : c'est exactement la croissance non bornée qu'on
vient de retirer du disque, transposée en RAM. Les entrées inactives sont donc
purgées, et leur nombre est plafonné.
"""

from __future__ import annotations

import threading
from dataclasses import dataclass, field

__all__ = ["Limiteur", "SeauAJetons"]


@dataclass
class SeauAJetons:
    """Un seau qui se remplit à débit constant et se vide à chaque requête."""

    capacite: float
    par_seconde: float
    jetons: float = field(default=0.0)
    #: `None` et non `0.0` : une horloge qui démarre à zéro est indistinguable
    #: d'une sentinelle à zéro, et le temps écoulé n'était jamais crédité au
    #: premier rechargement. Bogue attrapé par les tests avant toute mise en
    #: ligne.
    dernier: float | None = None

    def __post_init__(self) -> None:
        self.jetons = float(self.capacite)

    def consommer(self, maintenant: float, cout: float = 1.0) -> float:
        """Retourne 0 si la requête passe, sinon l'attente en secondes."""
        if self.dernier is None:
            self.dernier = maintenant
        ecoule = max(0.0, maintenant - self.dernier)
        self.dernier = maintenant
        self.jetons = min(self.capacite, self.jetons + ecoule * self.par_seconde)
        if self.jetons >= cout:
            self.jetons -= cout
            return 0.0
        manque = cout - self.jetons
        return manque / self.par_seconde if self.par_seconde > 0 else float("inf")


class Limiteur:
    """Limitation par adresse et globale, avec purge des adresses inactives."""

    def __init__(
        self,
        par_adresse: tuple[float, float] = (5.0, 0.5),
        global_: tuple[float, float] = (20.0, 2.0),
        max_adresses: int = 4096,
        oubli_s: float = 900.0,
    ) -> None:
        """``(capacite, par_seconde)`` pour chacun des deux niveaux.

        Défauts : une rafale de 5 requêtes par adresse puis une toutes les deux
        secondes ; globalement une rafale de 20 puis deux par seconde — le GPU
        en soutient environ 1,3, le reste attend dans le seau plutôt que d'être
        refusé.
        """
        self.par_adresse = par_adresse
        self.global_ = global_
        self.max_adresses = max_adresses
        self.oubli_s = oubli_s
        self._seaux: dict[str, SeauAJetons] = {}
        self._seau_global = SeauAJetons(*global_)
        self._verrou = threading.Lock()

    def _purger(self, maintenant: float) -> None:
        """Oublie les adresses inactives. Sans cela, le dictionnaire grandit
        avec le nombre d'adresses distinctes — la croissance non bornée qu'on
        vient de retirer du disque, transposée en mémoire."""
        morts = [
            ip
            for ip, s in self._seaux.items()
            if s.dernier is not None and maintenant - s.dernier > self.oubli_s
        ]
        for ip in morts:
            del self._seaux[ip]
        if len(self._seaux) > self.max_adresses:
            # Purge d'urgence : on garde les plus récentes.
            recents = sorted(self._seaux.items(), key=lambda kv: -(kv[1].dernier or 0.0))
            self._seaux = dict(recents[: self.max_adresses // 2])

    def verifier(self, adresse: str, maintenant: float, cout: float = 1.0) -> float:
        """0 si la requête passe, sinon l'attente conseillée en secondes.

        Le seau global n'est débité **que si** le seau de l'adresse a accepté :
        sinon un client bloqué continuerait à épuiser le quota commun.
        """
        with self._verrou:
            self._purger(maintenant)
            seau = self._seaux.get(adresse)
            if seau is None:
                seau = self._seaux[adresse] = SeauAJetons(*self.par_adresse)
            attente = seau.consommer(maintenant, cout)
            if attente > 0.0:
                return attente
            attente_globale = self._seau_global.consommer(maintenant, cout)
            if attente_globale > 0.0:
                # On rend le jeton pris à l'adresse : la requête n'a pas eu lieu.
                seau.jetons = min(seau.capacite, seau.jetons + cout)
                return attente_globale
            return 0.0

    @property
    def adresses_suivies(self) -> int:
        return len(self._seaux)

"""Un tracé de canvas devient des formules.

Séparé du HTTP pour être testable sans serveur. Le chemin complet :

    points du canvas (pixels)
      → validité (assez de points ? univalué ?)
      → normalisation IDENTIQUE à celle de l'entraînement
      → sous-échantillonnage vers la plage vue à l'entraînement
      → beam search sous masque d'arité
      → ajustement des constantes sur les points observés
      → front de Pareto

**Le point qui compte, et qui n'est pas cosmétique :** la normalisation doit
être exactement celle du générateur (`normalise_y`), sinon le modèle
reçoit une entrée d'une autre distribution que celle sur laquelle il a été
entraîné, et tout ce qu'on a mesuré ne s'applique plus.
"""

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np

from curvy.data.expr import to_infix
from curvy.data.pointcloud import normalise_y

#: Plage de tailles de nuage vue à l'entraînement (CloudConfig).
N_MIN, N_MAX = 20, 200

#: En dessous, le tracé n'est pas exploitable : on refuse plutôt que de rendre
#: une formule fabriquée sur trois points.
N_MIN_TRACE = 12

#: Étendue minimale de y après normalisation. Même seuil que le filtre
#: `DEGENERATE_NOISE` du générateur : une courbe plate n'est pas une courbe.
#:
#: Ce n'est pas un raffinement. Un cercle trié par x, doublons moyennés, devient
#: **exactement constant** — les moitiés haute et basse s'annulent. Et
#: `r_squared(constante, constante)` vaut **1,0** : la démo affichait donc
#: « R² 1,0000 » avec une formule dénuée de sens, juste à côté de
#: l'avertissement « ce tracé revient en arrière ». Le chiffre détruisait
#: l'avertissement. Vérifié aussi avec un modèle aux poids aléatoires, qui
#: obtenait le même 1,0000 (2026-08-20).
ETENDUE_MIN = 1e-9


#: Nombre de points de la courbe rendue au navigateur pour l'affichage.
N_APERCU = 600

#: Demi-largeur du domaine d'aperçu, en coordonnées normalisées. Les données
#: occupent `[-1, 1]` ; on rend la courbe sur `[-DOMAINE_APERCU, DOMAINE_APERCU]`
#: pour que le navigateur puisse dézoomer **sans aller-retour serveur**.
#:
#: C'est là qu'est l'intérêt : hors de la fenêtre observée, un polynôme rend
#: zéro sur 29 formules hors distribution sur 29 (mesuré le 2026-08-20), tandis
#: qu'une formule tient. La démo n'affichait jusqu'ici que l'intérieur du cadre,
#: c'est-à-dire précisément la zone où la baseline nous égale.
DOMAINE_APERCU = 3.0


@dataclass
class Formule:
    expression: str
    complexite: int
    r2: float
    constantes: list[float]
    prefixe: str
    #: L'expression avec ses constantes **valorisées**, et dans les unités
    #: d'origine quand elles ont un sens (import de données). Le squelette
    #: décrit une famille ; c'est cette ligne-ci que l'utilisateur veut lire.
    valorisee: str = ""
    #: La réponse principale retenue : meilleur R² d'ajustement,
    #: départagé par la simplicité. C'est elle qu'on affiche d'abord — le front
    #: est trié par complexité croissante, donc son premier élément est le plus
    #: simple ET le moins précis. L'afficher par défaut montrait le pire
    #: candidat en premier.
    principale: bool = False
    #: La courbe ajustée, échantillonnée sur `[-DOMAINE_APERCU, DOMAINE_APERCU]`
    #: en coordonnées normalisées — les données occupent le sous-intervalle
    #: `[-1, 1]`, le reste est de l'extrapolation.
    #: Renvoyée par le serveur plutôt que recalculée côté navigateur : la
    #: grammaire n'existe qu'ici, la réimplémenter en JavaScript créerait deux
    #: vérités qui divergeraient au premier opérateur ajouté.
    apercu: list[float] = field(default_factory=list)


@dataclass
class Reponse:
    ok: bool
    raison: str = ""
    #: Faux si le tracé revient en arrière : `y = f(x)` ne peut pas le décrire.
    #: C'est la limite de la v1, et la mesure qui décidera du mode paramétrique
    #:. On la rapporte, on ne la masque pas.
    univalue: bool = True
    n_points: int = 0
    formules: list[Formule] = field(default_factory=list)
    latence_ms: float = 0.0
    #: Affine appliquée à y, pour que le navigateur puisse revenir en pixels.
    y_scale: float = 1.0
    y_offset: float = 0.0
    #: Bornes en x du tracé d'origine, en pixels.
    x_min: float = 0.0
    x_max: float = 0.0
    #: Demi-largeur du domaine couvert par `apercu`, en unités normalisées.
    domaine_apercu: float = DOMAINE_APERCU


def est_univalue(x: np.ndarray, tolerance: float = 0.02) -> bool:
    """Le tracé avance-t-il toujours dans le même sens en x ?

    ``tolerance`` est exprimée en fraction de la largeur totale : une main qui
    tremble revient de quelques pixels en arrière sans que le tracé cesse d'être
    une fonction. Au-delà, c'est un vrai retour — une boucle, un cercle, un
    caractère.
    """
    if len(x) < 2:
        return True
    largeur = float(x.max() - x.min())
    if largeur < 1e-9:
        return False
    reculs = np.diff(x)
    recul_max = float(-reculs.min()) if reculs.min() < 0 else 0.0
    avance = float(reculs.sum())
    return recul_max <= tolerance * largeur and avance != 0.0


def preparer(
    points: np.ndarray, retourner_y: bool = True
) -> tuple[np.ndarray, np.ndarray, float, float, float, float] | None:
    """Points bruts → (x, y) normalisés comme à l'entraînement, plus les affines.

    ``retourner_y`` inverse l'ordonnée : vrai pour un canvas, dont l'axe croît
    vers le bas, faux pour des mesures importées, où le signe est celui des
    données.
    """
    if len(points) < N_MIN_TRACE:
        return None
    x = points[:, 0].astype(float)
    y = points[:, 1].astype(float)
    if retourner_y:
        y = -y  # le canvas compte vers le bas

    ordre = np.argsort(x, kind="stable")
    x, y = x[ordre], y[ordre]

    # Deux points à la même abscisse ne peuvent pas coexister dans y = f(x) :
    # on garde leur moyenne plutôt que d'en jeter un au hasard.
    x_uniques, index = np.unique(np.round(x, 6), return_inverse=True)
    if len(x_uniques) < N_MIN_TRACE:
        return None
    y_moyens = np.bincount(index, weights=y) / np.bincount(index)

    largeur = x_uniques.max() - x_uniques.min()
    if largeur < 1e-9:
        return None
    x_centre = float((x_uniques.max() + x_uniques.min()) / 2.0)
    x_demi = float(largeur / 2.0)
    x_norm = (x_uniques - x_centre) / x_demi
    y_norm, echelle, decalage = normalise_y(y_moyens)
    if float(y_norm.max() - y_norm.min()) < ETENDUE_MIN:
        return None

    if len(x_norm) > N_MAX:
        # Sous-échantillonnage régulier : au-delà de 200 points le modèle n'a
        # jamais rien vu de tel, et l'attention coûte O(N²).
        idx = np.linspace(0, len(x_norm) - 1, N_MAX).round().astype(int)
        x_norm, y_norm = x_norm[idx], y_norm[idx]
    return x_norm, y_norm, echelle, decalage, x_centre, x_demi


def formules_depuis_trace(
    points: np.ndarray,
    modele,
    device,
    rng: np.random.Generator,
    beam: int = 8,
    max_formules: int = 5,
    retourner_y: bool = True,
    executor=None,
) -> Reponse:
    """Le chemin complet, du tracé aux formules classées."""
    import time

    import torch

    from curvy.data.dataset import collate
    from curvy.data.expr import evaluate, from_prefix, to_prefix
    from curvy.infer.decode import beam_search, ids_to_node
    from curvy.infer.pareto import ajuster_candidats, front_de_pareto, selectionner
    from curvy.infer.rendu import Affines, formule_lisible
    from curvy.tokenizer.vocab import encode

    t0 = time.perf_counter()
    brut_x = points[:, 0].astype(float)
    univalue = est_univalue(brut_x)

    prepare = preparer(points, retourner_y=retourner_y)
    if prepare is None:
        if not univalue:
            raison = (
                "ce tracé revient en arrière, et une fois ramené à une fonction y = f(x) "
                "il devient plat : il n'y a rien à décrire. Un cercle, une boucle ou un "
                "caractère demandent le mode paramétrique, prévu en v2."
            )
        elif len(points) < N_MIN_TRACE:
            raison = f"tracé trop court : {len(points)} points, {N_MIN_TRACE} minimum."
        else:
            raison = (
                "tracé dégénéré : soit vertical, soit horizontal. Une droite horizontale "
                "est y = constante, ce que la grammaire ne modélise pas — et une verticale "
                "n'est pas une fonction."
            )
        return Reponse(ok=False, raison=raison, univalue=univalue, n_points=len(points))
    x, y, echelle, decalage, x_centre, x_demi = prepare

    pts = np.stack([x, y], axis=1).astype(np.float32)
    faux_ids = encode(from_prefix(["add", "mul", "C", "x", "C"]))
    batch = collate([(pts, faux_ids)]).to(device)
    with torch.no_grad():
        candidats = beam_search(modele, batch.points, batch.point_mask, beam=beam)[0]

    noeuds = [ids_to_node(seq) for seq, _ in candidats]
    ajustes = ajuster_candidats(noeuds, x, y, rng, executor=executor)
    front_complet = front_de_pareto(ajustes)
    retenue = selectionner(ajustes)
    front = front_complet[:max_formules]
    # La réponse principale ne doit jamais tomber hors de la troncature : c'est
    # celle qu'on annonce.
    if retenue is not None and all(c is not retenue for c in front):
        front = front[: max_formules - 1] + [retenue]

    # Les affines ne s'expriment que pour des données importées : les pixels
    # d'un canvas ne sont pas une unité qui intéresse quelqu'un.
    affines = (
        None
        if retourner_y
        else Affines(x_centre=x_centre, x_demi=x_demi, y_echelle=echelle, y_decalage=decalage)
    )

    grille = np.linspace(-DOMAINE_APERCU, DOMAINE_APERCU, N_APERCU)
    formules = []
    for c in front:
        with np.errstate(over="ignore", invalid="ignore", divide="ignore"):
            courbe = evaluate(c.node, grille, c.consts)
        # `null` en JSON pour les points non finis : le navigateur y coupera le
        # trait au lieu de dessiner une droite vers l'infini.
        apercu = [None if not np.isfinite(v) else round(float(v), 5) for v in courbe]
        formules.append(
            Formule(
                expression=to_infix(c.node),
                valorisee=formule_lisible(c.node, c.consts, affines),
                complexite=c.complexite,
                r2=round(c.r2_fit, 6),
                constantes=[round(v, 6) for v in c.consts],
                prefixe=" ".join(to_prefix(c.node)),
                principale=c is retenue,
                apercu=apercu,
            )
        )

    return Reponse(
        ok=True,
        univalue=univalue,
        n_points=len(x),
        formules=formules,
        latence_ms=round(1000 * (time.perf_counter() - t0), 1),
        y_scale=float(echelle),
        y_offset=float(decalage),
        x_min=float(brut_x.min()),
        x_max=float(brut_x.max()),
    )

"""Phase 6 — construit et **gèle** le jeu de test.

Règle non négociable : ce script s'exécute **avant** de regarder le moindre
résultat, et son produit ne bouge plus. Il écrit un fichier accompagné de son
empreinte SHA-256 ; toute mesure publiée doit citer cette empreinte, faute de
quoi rien ne prouve qu'on n'a pas rejoué la construction jusqu'à obtenir un
jeu flatteur.

    .venv/bin/python scripts/build_testset.py

Le jeu produit ne contient **aucun** squelette vu à l'entraînement, à condition
que l'entraînement ait tourné avec ``--exclure-test`` (). Les runs
exp-001 à exp-003 sont antérieurs : leurs chiffres sur ce jeu ne vaudraient
rien, et le script le rappelle en fin de sortie.
"""

from __future__ import annotations

import argparse
import gzip
import hashlib
import json
from pathlib import Path

import numpy as np

from curvy.config import DATASET_DIR, ensure_dirs
from curvy.data.dataset import make_validation_set
from curvy.data.expr import to_prefix
from curvy.data.generate import load_skeletons
from curvy.data.pointcloud import sample_cloud_fn
from curvy.data.split import SEL, partitionner
from curvy.data.testset import FORMULES_A_LA_MAIN
from curvy.seeding import make_rng

GRAINE = 20260820  # figée une fois pour toutes : la date du gel


def _empreinte(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for bloc in iter(lambda: fh.read(1 << 20), b""):
            h.update(bloc)
    return h.hexdigest()


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--skeletons", type=Path, default=DATASET_DIR / "skeletons-v1.jsonl.gz")
    ap.add_argument("--out", type=Path, default=DATASET_DIR / "testset-v1.jsonl.gz")
    ap.add_argument("--n-hors-distribution", type=int, default=6, help="nuages par formule")
    args = ap.parse_args(argv)
    ensure_dirs()

    if args.out.exists():
        print(f"REFUS : {args.out} existe déjà — un jeu de test gelé ne se régénère pas.")
        print(f"        empreinte actuelle : {_empreinte(args.out)}")
        print("        Le supprimer à la main est une décision, pas une commande de routine.")
        return 1

    items = load_skeletons(args.skeletons)
    part = partitionner(items)
    print(json.dumps(part.rapport(), indent=2, ensure_ascii=False))

    lignes: list[dict] = []

    # --- 1. tenu à l'écart : même grammaire, squelettes jamais vus ------------
    reserve = make_validation_set(
        args.skeletons,
        n=len(part.test),
        seed=GRAINE,
        garder=part.prefixes_de_test,
        un_nuage_par_squelette=True,
    )
    for ex in reserve:
        lignes.append(
            {
                "sous_ensemble": "tenu_a_lecart",
                "nom": " ".join(to_prefix(ex.node)),
                "dans_la_grammaire": True,
                "profondeur": ex.depth,
                "x": [round(float(v), 6) for v in ex.x],
                "y": [round(float(v), 6) for v in ex.y],
                "y_exact": [round(float(v), 6) for v in ex.y_clean],
                "ids": ex.ids,
            }
        )
    print(f"\ntenu à l'écart : {len(reserve)} exemples sur {len(part.test)} squelettes réservés")

    # --- 2. hors distribution : formules écrites à la main --------------------
    rng = make_rng(GRAINE)
    par_formule: dict[str, int] = {}
    for formule in FORMULES_A_LA_MAIN:
        obtenus = 0
        for _ in range(200):
            if obtenus >= args.n_hors_distribution:
                break
            cloud, _ = sample_cloud_fn(rng, formule.f)
            if cloud is None:
                continue
            with np.errstate(all="ignore"):
                exact = (np.asarray(formule.f(cloud.x)) - cloud.y_offset) / cloud.y_scale
            if not np.isfinite(exact).all():
                continue
            lignes.append(
                {
                    "sous_ensemble": "hors_distribution",
                    "nom": formule.nom,
                    "dans_la_grammaire": formule.prefixe is not None,
                    "prefixe": formule.prefixe,
                    "commentaire": formule.commentaire,
                    "x": [round(float(v), 6) for v in cloud.x],
                    "y": [round(float(v), 6) for v in cloud.y],
                    "y_exact": [round(float(v), 6) for v in exact],
                }
            )
            obtenus += 1
        par_formule[formule.nom] = obtenus

    manquantes = {k: v for k, v in par_formule.items() if v < args.n_hors_distribution}
    print(f"hors distribution : {sum(par_formule.values())} exemples")
    if manquantes:
        # Une formule sous-représentée doit se voir, pas disparaître en silence.
        print(f"  formules incomplètes (filtre d'identifiabilité) : {manquantes}")

    # --- écriture et gel -----------------------------------------------------
    with gzip.open(args.out, "wt", encoding="utf-8") as fh:
        for ligne in lignes:
            fh.write(json.dumps(ligne, ensure_ascii=False, sort_keys=True) + "\n")

    manifeste = {
        "fichier": args.out.name,
        "sha256": _empreinte(args.out),
        "graine": GRAINE,
        "sel_de_partition": SEL,
        "n_total": len(lignes),
        "n_tenu_a_lecart": len(reserve),
        "n_hors_distribution": sum(par_formule.values()),
        "par_formule": par_formule,
        "partition": part.rapport(),
        "reel_canvas": "non construit — outil de capture en Phase 8",
    }
    chemin_manifeste = Path("docs/benchmarks/testset-v1.json")
    chemin_manifeste.write_text(
        json.dumps(manifeste, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
    )

    print(f"\nécrit   : {args.out}")
    print(f"sha256  : {manifeste['sha256']}")
    print(f"manifeste : {chemin_manifeste}")
    print(
        "\nRAPPEL : exp-001 à exp-003 ont vu ces squelettes à l'entraînement.\n"
        "         Les mesurer sur ce jeu ne dirait rien. Il faut un run lancé\n"
        "         avec la partition."
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

"""Phase 5 — que rapporte le beam search, et à quel prix ?

Compare le décodage glouton (un candidat) au beam search (``k`` candidats
ajustés puis départagés) sur les deux protocoles : interpolation et
extrapolation. Le candidat annoncé est choisi **sans jamais regarder les points
tenus à l'écart** — même règle que la validation croisée imposée à la baseline
polynomiale.

    .venv/bin/python scripts/eval_beam.py --run exp-003 --preset v1 --beams 1,4,8,16
"""

from __future__ import annotations

import argparse
import time

import numpy as np
import torch

from curvy.config import DATASET_DIR, RUNS_DIR
from curvy.data.dataset import collate, make_validation_set
from curvy.data.expr import evaluate
from curvy.devices import pick_device
from curvy.infer.decode import beam_search, ids_to_node
from curvy.infer.fit import r_squared
from curvy.infer.pareto import ajuster_candidats, front_de_pareto, selectionner
from curvy.model.config import PRESETS
from curvy.model.curvy import CurvyModel
from curvy.seeding import make_rng, seed_everything

TOLS = (0.0, 0.002, 0.005, 0.02)


def _split(x: np.ndarray, mode: str, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]:
    n = len(x)
    n_hold = max(3, int(round(0.2 * n)))
    if mode == "extrapolation":
        order = np.argsort(x)
        return order[:-n_hold], order[-n_hold:]
    idx = rng.permutation(n)
    return idx[n_hold:], idx[:n_hold]


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--run", required=True)
    ap.add_argument("--preset", required=True, choices=sorted(PRESETS))
    ap.add_argument("--checkpoint", default="best.pt")
    ap.add_argument("--beams", default="1,4,8,16")
    ap.add_argument("--val-size", type=int, default=512)
    ap.add_argument("--val-seed", type=int, default=777)
    ap.add_argument("--batch-size", type=int, default=64)
    args = ap.parse_args(argv)

    seed_everything(args.val_seed)
    choix = pick_device("auto")
    device = choix.device
    model = CurvyModel(PRESETS[args.preset]).to(device)
    state = torch.load(
        RUNS_DIR / args.run / args.checkpoint, map_location="cpu", weights_only=False
    )
    model.load_state_dict(state["model"])
    model.eval()
    print(f"{args.run}/{args.checkpoint} — step {state.get('step', '?')} — {choix}")

    val = make_validation_set(
        DATASET_DIR / "skeletons-v1.jsonl.gz", args.val_size, seed=args.val_seed
    )
    beams = [int(b) for b in args.beams.split(",")]

    print(
        f"\n{'beam':>5} {'mode':>15} "
        + " ".join(f"tol={t:<7}" for t in TOLS)
        + f" {'candidats':>10} {'s/exemple':>10}"
    )
    resultats = {}
    for k in beams:
        # Décodage : une seule passe par valeur de beam, réutilisée par les deux protocoles.
        t0 = time.perf_counter()
        tous_candidats = []
        for start in range(0, len(val), args.batch_size):
            chunk = val[start : start + args.batch_size]
            batch = collate([(ex.points, ex.ids) for ex in chunk]).to(device)
            tous_candidats.extend(beam_search(model, batch.points, batch.point_mask, beam=k))
        t_decode = time.perf_counter() - t0

        for mode in ("interpolation", "extrapolation"):
            rng = make_rng(args.val_seed)
            succes = {t: [] for t in TOLS}
            n_cands, t_fit = [], time.perf_counter()
            for ex, cands in zip(val, tous_candidats, strict=True):
                keep, hold = _split(ex.x, mode, rng)
                nodes = [ids_to_node(seq) for seq, _ in cands]
                scores = [sc for _, sc in cands]
                ajustes = ajuster_candidats(nodes, ex.x[keep], ex.y[keep], rng, scores)
                n_cands.append(len(ajustes))
                for tol in TOLS:
                    c = selectionner(ajustes, tol=tol)
                    if c is None:
                        succes[tol].append(False)
                        continue
                    with np.errstate(all="ignore"):
                        pred = evaluate(c.node, ex.x[hold], c.consts)
                    succes[tol].append(r_squared(ex.y_clean[hold], pred) >= 0.99)
            dt = (time.perf_counter() - t_fit + t_decode) / len(val)
            taux = " ".join(f"{np.mean(succes[t]):<11.4f}" for t in TOLS)
            print(f"{k:>5} {mode:>15} {taux} {np.mean(n_cands):>10.1f} {dt:>10.3f}")
            resultats[(k, mode)] = {t: float(np.mean(succes[t])) for t in TOLS}

    print("\nRappels mesurés le 2026-08-19 (docs/benchmarks/results.md) :")
    print("   polynome, degre par validation croisee : interpolation 0.670, extrapolation 0.088")
    print("   vrai squelette (oracle)                : interpolation 0.803, extrapolation 0.430")

    # Exemple de front de Pareto, pour montrer ce que le produit livrerait.
    ex = val[0]
    keep, _ = _split(ex.x, "interpolation", make_rng(args.val_seed))
    cands = ajuster_candidats(
        [ids_to_node(s) for s, _ in tous_candidats[0]], ex.x[keep], ex.y[keep], make_rng(1)
    )
    from curvy.data.expr import to_infix

    print(f"\nExemple de front de Pareto (exemple 0, beam {beams[-1]}) :")
    for c in front_de_pareto(cands):
        print(f"   complexite {c.complexite:>3}  R2_ajust {c.r2_fit:>8.4f}  {to_infix(c.node)}")
    print(f"   vraie formule                          : {to_infix(ex.node)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

"""Évalue un checkpoint déjà entraîné, sans reprendre l'entraînement.

Sert à mesurer un run terminé avec des métriques ajoutées **après** lui — par
exemple l'extrapolation, ajoutée le 2026-08-19 alors qu'exp-001 et exp-002
étaient déjà finis.

    .venv/bin/python scripts/eval_checkpoint.py --run exp-002 --preset v1
"""

from __future__ import annotations

import argparse
import json

import torch

from curvy.config import DATASET_DIR, RUNS_DIR
from curvy.data.dataset import make_validation_set
from curvy.devices import pick_device
from curvy.model.config import PRESETS
from curvy.model.curvy import CurvyModel, count_parameters
from curvy.seeding import make_rng, seed_everything
from curvy.train.metrics import evaluate_model


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--run", required=True)
    ap.add_argument("--preset", required=True, choices=sorted(PRESETS))
    ap.add_argument("--checkpoint", default="best.pt")
    ap.add_argument("--val-size", type=int, default=512)
    ap.add_argument("--val-seed", type=int, default=777)
    ap.add_argument("--batch-size", type=int, default=64)
    args = ap.parse_args(argv)

    seed_everything(args.val_seed)
    path = RUNS_DIR / args.run / args.checkpoint
    choix = pick_device("auto")
    device = choix.device
    model = CurvyModel(PRESETS[args.preset]).to(device)
    state = torch.load(path, map_location="cpu", weights_only=False)
    model.load_state_dict(state["model"])
    step = state.get("step", "?")
    total = count_parameters(model)["TOTAL"]
    print(f"{args.run}/{args.checkpoint} — step {step}, {total:,} paramètres")
    print(f"device : {choix}")

    val = make_validation_set(
        DATASET_DIR / "skeletons-v1.jsonl.gz", args.val_size, seed=args.val_seed
    )
    rep = evaluate_model(model, val, device, make_rng(args.val_seed), batch_size=args.batch_size)
    d = rep.as_dict()
    print(json.dumps(d, ensure_ascii=False, indent=2))
    print()
    print("Lecture :")
    print(f"  interpolation : modèle {d['r2_rate']:.4f}  vs oracle {d['r2_rate_oracle']:.4f}")
    print(
        f"  extrapolation : modèle {d['r2_rate_extrap']:.4f}  vs oracle {d['r2_rate_extrap_oracle']:.4f}"
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

"""De quoi l'écart au plafond est-il fait ? (question de Billy, 2026-08-20)

Le modèle rend 8 candidats et on en retient un. Quand le résultat échoue, deux
causes possibles, et elles appellent des remèdes opposés :

- **aucun des 8 candidats n'était bon** → le modèle ne sait pas proposer. C'est
  un problème de capacité ou de recherche, et plus de paramètres aideraient ;
- **un candidat était bon mais la sélection l'a écarté** → agrandir le modèle ne
  servirait à rien, il faut corriger le choix.

On mesure donc trois chiffres sur le même jeu :

- **retenu** — ce que le produit rend vraiment (sans regarder la
  réponse) ;
- **rappel@k** — au moins un des k candidats aurait réussi. C'est le plafond de
  la sélection, atteignable sans toucher au modèle ;
- **oracle** — le vrai squelette. Le plafond de la tâche.

    .venv/bin/python scripts/eval_rappel.py --run exp-005 --preset v1
"""

from __future__ import annotations

import argparse
from collections import defaultdict

import numpy as np
import torch

from curvy.config import DATASET_DIR, RUNS_DIR
from curvy.data.dataset import collate, make_validation_set
from curvy.data.expr import evaluate
from curvy.devices import pick_device
from curvy.infer.decode import beam_search, ids_to_node
from curvy.infer.fit import fit_constants, r_squared
from curvy.infer.pareto import ajuster_candidats, selectionner
from curvy.model.config import PRESETS
from curvy.model.curvy import CurvyModel
from curvy.seeding import make_rng

SEUIL = 0.99


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--run", required=True)
    ap.add_argument("--preset", required=True)
    ap.add_argument("--checkpoint", default="best.pt")
    ap.add_argument("--beam", type=int, default=8)
    ap.add_argument("--val-size", type=int, default=512)
    ap.add_argument("--val-seed", type=int, default=777)
    ap.add_argument("--batch-size", type=int, default=64)
    args = ap.parse_args(argv)

    choix = pick_device("auto")
    modele = CurvyModel(PRESETS[args.preset]).to(choix.device)
    etat = torch.load(RUNS_DIR / args.run / args.checkpoint, map_location="cpu", weights_only=False)
    modele.load_state_dict(etat["model"])
    modele.eval()
    print(f"{args.run}/{args.checkpoint} step {etat.get('step', '?')} — {choix}")

    val = make_validation_set(
        DATASET_DIR / "skeletons-v1.jsonl.gz", args.val_size, seed=args.val_seed
    )

    candidats = []
    for start in range(0, len(val), args.batch_size):
        chunk = val[start : start + args.batch_size]
        batch = collate([(ex.points, ex.ids) for ex in chunk]).to(choix.device)
        candidats.extend(beam_search(modele, batch.points, batch.point_mask, beam=args.beam))

    rng = make_rng(args.val_seed)
    par_prof: dict[int, dict[str, list[bool]]] = defaultdict(
        lambda: {"retenu": [], "rappel": [], "oracle": [], "exact": []}
    )
    for ex, cands in zip(val, candidats, strict=True):
        n = len(ex.x)
        n_hold = max(3, int(round(0.2 * n)))
        idx = rng.permutation(n)
        hold, keep = idx[:n_hold], idx[n_hold:]

        noeuds = [ids_to_node(s) for s, _ in cands]
        ajustes = ajuster_candidats(noeuds, ex.x[keep], ex.y[keep], rng)

        def score(c, ex=ex, hold=hold):  # liés tôt : B023
            with np.errstate(all="ignore"):
                return r_squared(ex.y_clean[hold], evaluate(c.node, ex.x[hold], c.consts))

        retenu = selectionner(ajustes)
        d = par_prof[ex.depth]
        d["retenu"].append(retenu is not None and score(retenu) >= SEUIL)
        d["rappel"].append(any(score(c) >= SEUIL for c in ajustes))
        res = fit_constants(ex.node, ex.x[keep], ex.y[keep], rng)
        ok_or = False
        if res.ok:
            with np.errstate(all="ignore"):
                ok_or = (
                    r_squared(ex.y_clean[hold], evaluate(ex.node, ex.x[hold], res.consts)) >= SEUIL
                )
        d["oracle"].append(ok_or)
        from curvy.data.expr import to_prefix

        vrai = to_prefix(ex.node)
        d["exact"].append(any(n_ is not None and to_prefix(n_) == vrai for n_ in noeuds))

    print(
        f"\n{'prof':>5} {'n':>5} {'retenu':>8} {'rappel@' + str(args.beam):>10} "
        f"{'oracle':>8} {'vrai squelette dans le beam':>28}"
    )
    tot = defaultdict(list)
    for d in sorted(par_prof):
        v = par_prof[d]
        for k in v:
            tot[k] += v[k]
        print(
            f"{d:>5} {len(v['retenu']):>5} {np.mean(v['retenu']):>8.3f} "
            f"{np.mean(v['rappel']):>10.3f} {np.mean(v['oracle']):>8.3f} {np.mean(v['exact']):>28.3f}"
        )
    print(
        f"{'tous':>5} {len(tot['retenu']):>5} {np.mean(tot['retenu']):>8.3f} "
        f"{np.mean(tot['rappel']):>10.3f} {np.mean(tot['oracle']):>8.3f} {np.mean(tot['exact']):>28.3f}"
    )

    retenu, rappel, oracle = np.mean(tot["retenu"]), np.mean(tot["rappel"]), np.mean(tot["oracle"])
    print(f"\nDécomposition de l'écart au plafond ({100 * (oracle - retenu):.1f} points) :")
    print(
        f"  perdu à la SÉLECTION      : {100 * (rappel - retenu):>5.1f} pts "
        f"— un bon candidat était là, on ne l'a pas pris"
    )
    print(
        f"  perdu à la PROPOSITION    : {100 * (oracle - rappel):>5.1f} pts "
        f"— aucun des {args.beam} candidats n'était bon"
    )
    print("\nLe premier poste se corrige sans toucher au modèle.")
    print("Le second est le seul que plus de paramètres pourrait réduire.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

"""Quelle règle de sélection récupère les points perdus ? (2026-08-20)

Mesuré la veille : sur les 10,5 points qui séparent le produit de l'oracle,
**4,9 sont perdus à la sélection** — un bon candidat figurait parmi les huit et
n'a pas été retenu. Ce poste se corrige sans toucher au modèle.

Le script ajuste les candidats **une seule fois** par exemple, puis applique
toutes les règles au même jeu ajusté : leur comparaison est donc exacte, à
candidats et à ajustements identiques, et non pas approchée d'un run à l'autre.

Les règles en lice, et pourquoi :

- ``max_r2`` — l'actuelle. Maximum du R² d'ajustement ;
- ``validation_croisee`` — on impose la validation croisée au polynôme pour
  choisir son degré, et on s'autorise le maximum brut pour nos candidats. La
  même règle des deux côtés supprime l'asymétrie ;
- ``penalite_complexite`` — critère d'information : le R² payé au prix du
  nombre de constantes ;
- ``modele`` — la log-vraisemblance du beam, c'est-à-dire l'avis du réseau,
  jamais utilisé jusqu'ici ;
- ``modele_puis_r2`` — le réseau départage les candidats proches en R².

    .venv/bin/python scripts/eval_selection.py --run exp-005 --preset v1
"""

from __future__ import annotations

import argparse
from collections import defaultdict

import numpy as np
import torch

from curvy.config import DATASET_DIR, RUNS_DIR
from curvy.data.dataset import collate, make_validation_set
from curvy.data.expr import evaluate
from curvy.devices import pick_device
from curvy.infer.decode import beam_search, ids_to_node
from curvy.infer.fit import fit_constants, r_squared
from curvy.infer.pareto import Candidat, ajuster_candidats, selectionner_selon_bruit
from curvy.model.config import PRESETS
from curvy.model.curvy import CurvyModel
from curvy.seeding import make_rng

SEUIL = 0.99

#: Marges essayées autour du bruit estimé. Plusieurs valeurs plutôt qu'une :
#: un résultat négatif doit distinguer « l'idée est mauvaise » de « l'idée est
#: bonne mais mal réglée ».
MARGES_BRUIT = (0.0, 0.5, 1.0, 2.0)

NOMS_REGLES = [
    "max_r2",
    "penalite_complexite",
    "modele",
    "modele_puis_r2",
    "validation_croisee",
    *[f"bruit_marge_{m}" for m in MARGES_BRUIT],
]


def _cv_score(c: Candidat, x: np.ndarray, y: np.ndarray, rng, k: int = 4) -> float:
    """R² d'un candidat sur des points d'ajustement tenus à l'écart, en k plis.

    Ne regarde jamais les points d'évaluation : c'est la règle imposée au
    polynôme, appliquée à nous.
    """
    n = len(x)
    if n < 4 * k:
        return c.r2_fit
    idx = rng.permutation(n)
    plis = np.array_split(idx, k)
    scores = []
    for pli in plis:
        tr = np.setdiff1d(idx, pli)
        res = fit_constants(c.node, x[tr], y[tr], rng)
        if not res.ok:
            continue
        with np.errstate(all="ignore"):
            scores.append(r_squared(y[pli], evaluate(c.node, x[pli], res.consts)))
    return float(np.mean(scores)) if scores else float("-inf")


def regles(cands: list[Candidat], x, y, rng) -> dict[str, Candidat | None]:
    """Un candidat retenu par règle, tous jugés sur le même jeu ajusté."""
    if not cands:
        return {}
    out: dict[str, Candidat | None] = {}
    out["max_r2"] = min(
        (c for c in cands if c.r2_fit >= max(k.r2_fit for k in cands)),
        key=lambda c: c.complexite,
    )
    out["penalite_complexite"] = max(cands, key=lambda c: c.r2_fit - 0.002 * len(c.consts))
    out["modele"] = max(cands, key=lambda c: c.score_modele)
    meilleur = max(c.r2_fit for c in cands)
    proches = [c for c in cands if c.r2_fit >= meilleur - 0.01]
    out["modele_puis_r2"] = max(proches, key=lambda c: c.score_modele)
    cv = {id(c): _cv_score(c, x, y, rng) for c in cands}
    out["validation_croisee"] = max(cands, key=lambda c: cv[id(c)])
    for marge in MARGES_BRUIT:
        out[f"bruit_marge_{marge}"] = selectionner_selon_bruit(cands, x, y, marge=marge)
    return out


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--run", required=True)
    ap.add_argument("--preset", required=True)
    ap.add_argument("--checkpoint", default="best.pt")
    ap.add_argument("--beam", type=int, default=8)
    ap.add_argument("--val-size", type=int, default=512)
    ap.add_argument("--val-seed", type=int, default=777)
    ap.add_argument("--batch-size", type=int, default=64)
    args = ap.parse_args(argv)

    choix = pick_device("auto")
    modele = CurvyModel(PRESETS[args.preset]).to(choix.device)
    etat = torch.load(RUNS_DIR / args.run / args.checkpoint, map_location="cpu", weights_only=False)
    modele.load_state_dict(etat["model"])
    modele.eval()
    print(f"{args.run}/{args.checkpoint} step {etat.get('step', '?')} — {choix}")

    val = make_validation_set(
        DATASET_DIR / "skeletons-v1.jsonl.gz", args.val_size, seed=args.val_seed
    )
    cands_par_ex = []
    for start in range(0, len(val), args.batch_size):
        chunk = val[start : start + args.batch_size]
        batch = collate([(ex.points, ex.ids) for ex in chunk]).to(choix.device)
        cands_par_ex.extend(beam_search(modele, batch.points, batch.point_mask, beam=args.beam))

    rng = make_rng(args.val_seed)
    reussites: dict[str, list[bool]] = defaultdict(list)
    for ex, cands in zip(val, cands_par_ex, strict=True):
        n = len(ex.x)
        nh = max(3, int(round(0.2 * n)))
        idx = rng.permutation(n)
        hold, keep = idx[:nh], idx[nh:]
        noeuds = [ids_to_node(s) for s, _ in cands]
        scores = [sc for _, sc in cands]
        ajustes = ajuster_candidats(noeuds, ex.x[keep], ex.y[keep], rng, scores)

        def note(c, ex=ex, hold=hold):
            if c is None:
                return False
            with np.errstate(all="ignore"):
                return r_squared(ex.y_clean[hold], evaluate(c.node, ex.x[hold], c.consts)) >= SEUIL

        choisis = regles(ajustes, ex.x[keep], ex.y[keep], rng)
        for nom, c in choisis.items():
            reussites[nom].append(note(c))
        if not choisis:
            for nom in NOMS_REGLES:
                reussites[nom].append(False)
        reussites["rappel@k"].append(any(note(c) for c in ajustes))

    base = np.mean(reussites["max_r2"])
    print(f"\n{'règle':>22} {'taux':>8} {'écart':>8}")
    for nom in [*NOMS_REGLES, "rappel@k"]:
        v = np.mean(reussites[nom])
        marque = (
            "   <- actuelle"
            if nom == "max_r2"
            else ("   <- plafond atteignable" if nom == "rappel@k" else "")
        )
        print(f"{nom:>22} {v:>8.4f} {100 * (v - base):>+7.1f}{marque}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

"""Évalue sur le jeu de test **figé** (Phase 6).

Trois concurrents, exactement le même protocole :

- **oracle** — le vrai squelette, constantes réajustées. Plafond de la tâche ;
- **polynôme** — degré 1 à 8 choisi par validation croisée sur les seuls points
  d'ajustement. La baseline à battre ;
- **Curvy** — beam search, candidat choisi sans regarder les points tenus à
  l'écart. Omis si ``--run`` n'est pas fourni : les deux premiers ne
  demandent aucun modèle et peuvent être mesurés avant qu'un run existe.

Résultats rendus **par sous-ensemble et par profondeur**, jamais agrégés en un
chiffre unique : « tenu à l'écart » et « hors distribution » ne répondent pas à
la même question.

    .venv/bin/python scripts/eval_testset.py                     # oracle + polynôme
    .venv/bin/python scripts/eval_testset.py --run exp-005 --preset v1
"""

from __future__ import annotations

import argparse
import gzip
import hashlib
import json
from collections import defaultdict
from pathlib import Path

import numpy as np

from curvy.config import DATASET_DIR, RUNS_DIR
from curvy.data.expr import evaluate, from_prefix
from curvy.infer.fit import fit_constants, r_squared
from curvy.seeding import make_rng

SEUIL = 0.99
DEGRES = range(1, 9)


def charger(path: Path) -> tuple[list[dict], str]:
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for bloc in iter(lambda: fh.read(1 << 20), b""):
            h.update(bloc)
    with gzip.open(path, "rt", encoding="utf-8") as fh:
        return [json.loads(ligne) for ligne in fh], h.hexdigest()


def _decoupe(x: np.ndarray, mode: str, rng: np.random.Generator):
    n = len(x)
    n_hold = max(3, int(round(0.2 * n)))
    if mode == "extrapolation":
        o = np.argsort(x)
        return o[:-n_hold], o[-n_hold:]
    idx = rng.permutation(n)
    return idx[n_hold:], idx[:n_hold]


def _degre_par_cv(x, y, k=5):
    n = len(x)
    idx = np.arange(n)
    plis = np.array_split(idx, k)
    best, best_mse = 1, np.inf
    for deg in DEGRES:
        err = []
        for pli in plis:
            tr = np.setdiff1d(idx, pli)
            if len(tr) <= deg + 1 or len(pli) == 0:
                continue
            with np.errstate(all="ignore"):
                try:
                    p = np.polyval(np.polyfit(x[tr], y[tr], deg), x[pli])
                except Exception:
                    continue
            if np.all(np.isfinite(p)):
                err.append(float(np.mean((p - y[pli]) ** 2)))
        if err and np.mean(err) < best_mse:
            best_mse, best = float(np.mean(err)), deg
    return best


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--testset", type=Path, default=DATASET_DIR / "testset-v1.jsonl.gz")
    ap.add_argument("--run", default=None)
    ap.add_argument("--preset", default="v1")
    ap.add_argument("--checkpoint", default="best.pt")
    ap.add_argument("--beam", type=int, default=8)
    ap.add_argument("--batch-size", type=int, default=64)
    ap.add_argument("--seed", type=int, default=20260820)
    args = ap.parse_args(argv)

    lignes, empreinte = charger(args.testset)
    print(f"jeu de test : {args.testset.name}")
    print(f"sha256      : {empreinte}")
    manifeste = Path("docs/benchmarks/testset-v1.json")
    if manifeste.exists():
        attendu = json.loads(manifeste.read_text())["sha256"]
        etat = "conforme au manifeste" if attendu == empreinte else "!! DIVERGENT DU MANIFESTE !!"
        print(f"              {etat}")
    print(f"exemples    : {len(lignes)}\n")

    modele = None
    if args.run:
        import torch

        from curvy.data.dataset import collate
        from curvy.devices import pick_device
        from curvy.infer.decode import beam_search, ids_to_node
        from curvy.infer.pareto import ajuster_candidats, selectionner
        from curvy.model.config import PRESETS
        from curvy.model.curvy import CurvyModel
        from curvy.tokenizer.vocab import encode

        choix = pick_device("auto")
        modele = CurvyModel(PRESETS[args.preset]).to(choix.device)
        etat = torch.load(
            RUNS_DIR / args.run / args.checkpoint, map_location="cpu", weights_only=False
        )
        modele.load_state_dict(etat["model"])
        modele.eval()
        cfg_run = json.loads((RUNS_DIR / args.run / "config.json").read_text())
        exclu = cfg_run.get("exclure_test", False)
        print(f"modèle      : {args.run}/{args.checkpoint}, step {etat.get('step', '?')} — {choix}")
        print(f"exclure_test: {exclu}" + ("" if exclu else "  <-- CHIFFRES SANS VALEUR"))
        print()

        # Décodage batché : un seul passage sur tout le jeu.
        candidats: list[list] = []
        for start in range(0, len(lignes), args.batch_size):
            chunk = lignes[start : start + args.batch_size]
            faux_ids = encode(from_prefix(["add", "mul", "C", "x", "C"]))
            batch = collate(
                [
                    (
                        np.stack([np.array(lg["x"]), np.array(lg["y"])], 1).astype(np.float32),
                        faux_ids,
                    )
                    for lg in chunk
                ]
            ).to(choix.device)
            candidats.extend(beam_search(modele, batch.points, batch.point_mask, beam=args.beam))

    resultats: dict = defaultdict(lambda: defaultdict(list))
    for i, ligne in enumerate(lignes):
        x = np.array(ligne["x"])
        y = np.array(ligne["y"])
        exact = np.array(ligne["y_exact"])
        sous = ligne["sous_ensemble"]
        cle = ligne["profondeur"] if sous == "tenu_a_lecart" else ligne["nom"]
        rng = make_rng(args.seed + i)

        for mode in ("interpolation", "extrapolation"):
            keep, hold = _decoupe(x, mode, rng)

            # --- polynôme ---
            deg = _degre_par_cv(x[keep], y[keep])
            with np.errstate(all="ignore"):
                try:
                    pred = np.polyval(np.polyfit(x[keep], y[keep], deg), x[hold])
                    ok_poly = r_squared(exact[hold], pred) >= SEUIL
                except Exception:
                    ok_poly = False
            resultats[(sous, mode, "polynome")][cle].append(ok_poly)

            # --- oracle : seulement si le vrai squelette existe ---
            prefixe = ligne.get("nom") if sous == "tenu_a_lecart" else ligne.get("prefixe")
            if sous == "hors_distribution" and prefixe:
                prefixe = f"add mul C {prefixe} C"
            if prefixe:
                node = from_prefix(prefixe.split())
                res = fit_constants(node, x[keep], y[keep], rng)
                ok_or = False
                if res.ok:
                    with np.errstate(all="ignore"):
                        ok_or = r_squared(exact[hold], evaluate(node, x[hold], res.consts)) >= SEUIL
                resultats[(sous, mode, "oracle")][cle].append(ok_or)

            # --- Curvy ---
            if modele is not None:
                noeuds = [ids_to_node(s) for s, _ in candidats[i]]
                ajustes = ajuster_candidats(noeuds, x[keep], y[keep], rng)
                c = selectionner(ajustes)
                ok_m = False
                if c is not None:
                    with np.errstate(all="ignore"):
                        ok_m = r_squared(exact[hold], evaluate(c.node, x[hold], c.consts)) >= SEUIL
                resultats[(sous, mode, "curvy")][cle].append(ok_m)

    concurrents = ["oracle", "polynome"] + (["curvy"] if modele is not None else [])
    for sous in ("tenu_a_lecart", "hors_distribution"):
        for mode in ("interpolation", "extrapolation"):
            dispo = [c for c in concurrents if (sous, mode, c) in resultats]
            if not dispo:
                continue
            print(f"=== {sous} — {mode} ===")
            cles = sorted({k for c in dispo for k in resultats[(sous, mode, c)]}, key=str)
            large = max(len(str(k)) for k in cles) + 2
            print(f"{'':{large}} " + " ".join(f"{c:>16}" for c in dispo))
            for cle in cles:
                vals = []
                for c in dispo:
                    v = resultats[(sous, mode, c)].get(cle, [])
                    # n par concurrent : l'oracle n'existe pas pour les formules
                    # hors grammaire, afficher un n commun ferait croire à une
                    # comparaison qui n'a pas lieu.
                    vals.append(f"{np.mean(v):>10.3f} ({len(v):>2})" if v else f"{'-':>16}")
                print(f"{str(cle):{large}} " + " ".join(vals))

            tot = []
            for c in dispo:
                v = [x for vs in resultats[(sous, mode, c)].values() for x in vs]
                tot.append(f"{np.mean(v):>10.3f} ({len(v):>2})" if v else f"{'-':>16}")
            print(f"{'TOTAL':{large}} " + " ".join(tot))

            # Comparaison à périmètre égal : restreinte aux clés où TOUS les
            # concurrents ont un chiffre. Sans ça, on compare l'oracle sur les
            # formules exprimables au polynôme sur toutes, ce qui ne veut rien dire.
            communes = {k for k in cles if all(resultats[(sous, mode, c)].get(k) for c in dispo)}
            if communes and len(communes) < len(cles):
                comm = []
                for c in dispo:
                    v = [x for k in communes for x in resultats[(sous, mode, c)][k]]
                    comm.append(f"{np.mean(v):>10.3f} ({len(v):>2})")
                print(f"{'  à périmètre égal':{large}} " + " ".join(comm))
                print(f"{'':{large}}   ({len(communes)} clés communes sur {len(cles)})")

            n_ref = sum(len(v) for v in resultats[(sous, mode, dispo[-1])].values())
            marge = 1.96 * (0.25 / max(n_ref, 1)) ** 0.5
            print(f"{'':{large}} (±{100 * marge:.1f} pt à 95 % sur n={n_ref})\n")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

#!/usr/bin/env bash
# Lance UN entraînement dans tmux, log complet vers logs/raw/.
# Usage : ./scripts/launch_run.sh <run-name> <preset> <steps> [workers]
#
# Passe --exclure-test PAR DÉFAUT : les squelettes réservés au jeu
# de test sortent du flux. Sans ça, les chiffres du run sur le jeu de test ne
# valent rien — c'est arrivé une fois, le 2026-08-20, et le run a été avorté
# après huit secondes. Pour reproduire un run d'avant la Phase 6 :
# CURVY_SANS_EXCLUSION=1 ./scripts/launch_run.sh ...
#
# CURVY_REPRENDRE=1 reprend au dernier checkpoint au lieu de repartir de zéro.
# Le checkpoint porte le modèle, l'optimiseur, le scheduler et le compteur de
# steps : la reprise retrouve la position exacte dans le cosine, ce n'est pas
# un redémarrage déguisé. Sert quand on libère le GPU pour autre chose.
#
# Empreinte CPU volontairement limitée : la machine a 12 cœurs et les fait
# tourner pour autre chose. Un run qui les prend tous a déjà fait tomber le
# réseau de la machine (2026-08-19). Défaut 5 workers, `nice 10`, threads
# torch plafonnés — : 5 workers suffisent à alimenter le GPU.
# Le run survit à une déconnexion SSH ; `--resume` (défaut) reprendrait au
# dernier checkpoint après une coupure.
set -euo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$ROOT"

RUN="${1:?usage: launch_run.sh <run-name> <preset> <steps>}"
PRESET="${2:?preset manquant}"
STEPS="${3:?steps manquant}"
WORKERS="${4:-5}"
# Le drapeau est écrit dans le config.json du run : c'est lui qui dira si les
# mesures sur le jeu de test veulent dire quelque chose.
EXCLUSION="--exclure-test"
[ -n "${CURVY_SANS_EXCLUSION:-}" ] && EXCLUSION=""
# Par défaut on repart de zéro : un run neuf ne doit jamais ramasser par
# accident le checkpoint d'un homonyme.
REPRISE="--no-resume"
[ -n "${CURVY_REPRENDRE:-}" ] && REPRISE=""

export CUDA_DEVICE_ORDER=PCI_BUS_ID
export CURVY_CUDA_ALLOW=GPU-1234abcd-0000-0000-0000-000000000000

# Plafonds CPU : le GPU est le client, pas le patron de cette machine.
export OMP_NUM_THREADS=2
export MKL_NUM_THREADS=2

SESSION="curvy-train"
STAMP="$(date '+%Y%m%d-%H%M%S')"
LOG="logs/raw/${STAMP}-${RUN}-${PRESET}.log"

tmux has-session -t "$SESSION" 2>/dev/null && {
  echo "session tmux '$SESSION' déjà active — attache-toi avec : tmux attach -t $SESSION"
  exit 1
}

# Refus net si la VRAM est prise : mieux vaut ne pas démarrer qu'OOM à 3 h de run.
.venv/bin/python -m curvy.cli_gpu

tmux new-session -d -s "$SESSION" -c "$ROOT" "
  set -o pipefail
  echo '=== ${RUN} : preset ${PRESET}, ${STEPS} steps, ${WORKERS} workers, nice 10 ==='
  echo '=== exclusion du jeu de test : ${EXCLUSION:-AUCUNE} ==='
  echo '=== reprise : ${REPRISE:-OUI, au dernier checkpoint} ==='
  nice -n 10 .venv/bin/python -m curvy.train.run \
      --run-name ${RUN} --preset ${PRESET} --steps ${STEPS} \
      --eval-every 1000 --log-every 100 --batch-size 512 \
      --workers ${WORKERS} ${EXCLUSION} ${REPRISE} \
      2>&1 | tee ${LOG}
  echo '=== terminé ==='
  sleep 3600
"
echo "session tmux '$SESSION' lancée pour ${RUN}."
echo "  suivre  : tmux attach -t $SESSION   (détacher : Ctrl-b d)"
echo "  log     : ${LOG}"
echo "  courbes : .venv/bin/python scripts/plot_training.py --run ${RUN}"
echo "  CPU     : ${WORKERS} workers + 1 process principal, nice 10, sur 12 cœurs"
echo "  test    : ${EXCLUSION:-AUCUNE EXCLUSION (les chiffres sur le jeu de test seront nuls)}"
echo "  reprise : ${REPRISE:-au dernier checkpoint}"

#!/usr/bin/env bash
# Lance les entraînements dans tmux, logs complets vers logs/raw/.
# Les runs survivent à une déconnexion SSH ; `make train` les reprendrait au
# dernier checkpoint en cas de coupure.
set -euo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$ROOT"

export CUDA_DEVICE_ORDER=PCI_BUS_ID
export CURVY_CUDA_ALLOW=GPU-1234abcd-0000-0000-0000-000000000000

SESSION="curvy-train"
STAMP="$(date '+%Y%m%d-%H%M%S')"

tmux has-session -t "$SESSION" 2>/dev/null && {
  echo "session tmux '$SESSION' déjà active — attache-toi avec : tmux attach -t $SESSION"
  exit 1
}

tmux new-session -d -s "$SESSION" -c "$ROOT" "
  set -o pipefail
  echo '=== Contrôle VRAM avant lancement ==='
  .venv/bin/python -m curvy.cli_gpu || exit 1

  echo '=== exp-001 : preset small, 20 000 steps ==='
  .venv/bin/python -m curvy.train.run \
      --run-name exp-001 --preset small --steps 20000 \
      --eval-every 1000 --log-every 100 --batch-size 512 --no-resume \
      2>&1 | tee logs/raw/${STAMP}-exp-001-small.log

  echo '=== exp-002 : preset v1, 20 000 steps ==='
  .venv/bin/python -m curvy.train.run \
      --run-name exp-002 --preset v1 --steps 20000 \
      --eval-every 1000 --log-every 100 --batch-size 512 --no-resume \
      2>&1 | tee logs/raw/${STAMP}-exp-002-v1.log

  echo '=== terminé ==='
  sleep 3600
"
echo "session tmux '$SESSION' lancée."
echo "  suivre  : tmux attach -t $SESSION   (détacher : Ctrl-b d)"
echo "  logs    : logs/raw/${STAMP}-exp-00*.log"
echo "  courbes : .venv/bin/python scripts/plot_training.py --run exp-001"

#!/usr/bin/env bash
# logcmd.sh — exécute une commande, affiche sa sortie ET l'archive dans logs/commands.md
# Usage: ./scripts/logcmd.sh "description courte" -- <commande...>
# Règle projet n°1 : aucun chiffre dans la doc qui ne vienne d'ici.
set -uo pipefail
DESC="$1"; shift
[ "${1:-}" = "--" ] && shift
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
LOG="$ROOT/logs/commands.md"
TS="$(date '+%Y-%m-%d %H:%M:%S %Z')"
OUT="$(mktemp)"
"$@" > "$OUT" 2>&1
RC=$?
LINES=$(wc -l < "$OUT")
{
  echo ""
  echo "## [$TS] $DESC"
  echo ""
  echo '```console'
  echo "\$ $*"
  if [ "$LINES" -gt 120 ]; then
    RAW="logs/raw/$(date '+%Y%m%d-%H%M%S')-$(echo "$DESC" | tr -cd '[:alnum:]' | cut -c1-30).log"
    cp "$OUT" "$ROOT/$RAW"
    head -60 "$OUT"
    echo "[... $((LINES-60)) lignes tronquées, sortie complète : $RAW ...]"
  else
    cat "$OUT"
  fi
  echo '```'
  echo ""
  echo "*code retour : $RC*"
} >> "$LOG"
cat "$OUT"; rm -f "$OUT"
exit $RC

"""Courbes d'entraînement depuis le ``log.jsonl`` d'un run.

Trace ce qui compte, et **l'oracle en trait tireté** sur le panneau du R² : sans
la référence de plafond, un taux de 45 % est illisible.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

from curvy.config import FIGURES_DIR, RUNS_DIR, ensure_dirs  # noqa: E402


def load(path: Path) -> tuple[list[dict], list[dict]]:
    train, ev = [], []
    for line in path.read_text(encoding="utf-8").splitlines():
        rec = json.loads(line)
        if rec.get("event") == "train":
            train.append(rec)
        elif rec.get("event") == "eval":
            ev.append(rec)
    return train, ev


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--run", default="exp-001")
    ap.add_argument("--runs-dir", type=Path, default=RUNS_DIR)
    ap.add_argument("--out", type=Path, default=None)
    args = ap.parse_args(argv)

    ensure_dirs()
    log = args.runs_dir / args.run / "log.jsonl"
    train, ev = load(log)
    if not train:
        print(f"aucune donnée dans {log}")
        return 1
    out = args.out or FIGURES_DIR / f"{args.run}-courbes.png"

    fig, axes = plt.subplots(2, 2, figsize=(13, 8))
    fig.suptitle(f"Curvy — {args.run} — {len(train)} points de log", fontsize=13)

    ax = axes[0, 0]
    ax.plot([r["step"] for r in train], [r["loss"] for r in train], lw=1.2)
    ax.set_title("loss (cross-entropy, teacher forcing)")
    ax.set_xlabel("step")
    ax.set_yscale("log")
    ax.grid(alpha=0.3)

    ax = axes[0, 1]
    ax.plot(
        [r["step"] for r in train], [r["token_acc"] for r in train], lw=1.2, label="token (train)"
    )
    if ev:
        ax.plot(
            [r["step"] for r in ev],
            [r["token_acc"] for r in ev],
            lw=1.4,
            marker="o",
            ms=3,
            label="token (val)",
        )
        ax.plot(
            [r["step"] for r in ev],
            [r["seq_acc_greedy"] for r in ev],
            lw=1.4,
            marker="s",
            ms=3,
            label="séquence exacte (val, glouton)",
        )
    ax.set_title("accuracy")
    ax.set_xlabel("step")
    ax.set_ylim(0, 1)
    ax.legend(fontsize=8)
    ax.grid(alpha=0.3)

    ax = axes[1, 0]
    if ev:
        ax.plot(
            [r["step"] for r in ev],
            [r["r2_rate"] for r in ev],
            lw=1.8,
            marker="o",
            ms=4,
            color="tab:green",
            label="Curvy : R² ≥ 0.99",
        )
        ax.plot(
            [r["step"] for r in ev],
            [r["r2_rate_oracle"] for r in ev],
            lw=1.4,
            ls="--",
            color="gray",
            label="oracle (vrai squelette)",
        )
        # Extrapolation : absente des runs antérieurs au 2026-08-19, on ne trace
        # la courbe que si le log la porte.
        if any("r2_rate_extrap" in r for r in ev):
            ax.plot(
                [r["step"] for r in ev if "r2_rate_extrap" in r],
                [r["r2_rate_extrap"] for r in ev if "r2_rate_extrap" in r],
                lw=1.8,
                marker="^",
                ms=4,
                color="tab:purple",
                label="Curvy : R² ≥ 0.99 en extrapolation",
            )
            ax.plot(
                [r["step"] for r in ev if "r2_rate_extrap_oracle" in r],
                [r["r2_rate_extrap_oracle"] for r in ev if "r2_rate_extrap_oracle" in r],
                lw=1.2,
                ls="--",
                color="tab:purple",
                alpha=0.5,
                label="oracle en extrapolation",
            )
        ax.axhline(0.5, color="tab:red", ls=":", lw=1.2, label="jalon Phase 4 (50 %)")
        # Baselines polynomiales mesurées le 2026-08-19 (docs/benchmarks/results.md).
        ax.axhline(0.670, color="tab:orange", ls="-.", lw=1.2, label="polynôme, interpolation")
        ax.axhline(0.088, color="tab:orange", ls=":", lw=1.2, label="polynôme, extrapolation")
    ax.set_title("MÉTRIQUE PRINCIPALE — taux de R² ≥ 0.99 (points tenus à l'écart)")
    ax.set_xlabel("step")
    ax.set_ylim(0, 1)
    ax.legend(fontsize=8)
    ax.grid(alpha=0.3)

    ax = axes[1, 1]
    if ev and ev[-1].get("per_depth"):
        depths = sorted(ev[-1]["per_depth"], key=int)
        for d in depths:
            ax.plot(
                [r["step"] for r in ev],
                [r.get("per_depth", {}).get(d, float("nan")) for r in ev],
                lw=1.2,
                marker=".",
                label=f"profondeur {d}",
            )
    ax.set_title("R² ≥ 0.99 par profondeur de squelette")
    ax.set_xlabel("step")
    ax.set_ylim(0, 1)
    ax.legend(fontsize=7, ncol=2)
    ax.grid(alpha=0.3)

    fig.tight_layout(rect=(0, 0, 1, 0.95))
    fig.savefig(out, dpi=140)
    fig.savefig(out.with_suffix(".svg"))
    print(f"figure écrite : {out}")
    if ev:
        last = ev[-1]
        print(
            f"dernier eval (step {last['step']}) : R²≥0.99 = {100 * last['r2_rate']:.1f} % "
            f"(oracle {100 * last['r2_rate_oracle']:.1f} %), "
            f"séquence exacte = {100 * last['seq_acc_greedy']:.1f} %"
        )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

"""Aperçu visuel du dataset — jalon de la Phase 1.

Deux colonnes volontairement différentes :

- **gauche** : squelettes tirés uniformément dans l'ensemble dédupliqué ;
- **droite** : squelettes tirés avec la pondération par multiplicité
  (``tau = 0,5``).

La comparaison est le sujet de la figure : à gauche ce que produit un pipeline
naïf, à droite ce que le modèle doit réellement.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import sympy  # noqa: E402

from curvy.config import DATASET_DIR, FIGURES_DIR, ensure_dirs  # noqa: E402
from curvy.data.expr import const_name_iter, evaluate, from_prefix, to_infix  # noqa: E402
from curvy.data.generate import load_skeletons  # noqa: E402
from curvy.data.pointcloud import CloudConfig, sample_cloud  # noqa: E402
from curvy.data.weighting import describe_weights, stratified_weights  # noqa: E402
from curvy.seeding import DEFAULT_SEED, make_rng  # noqa: E402


def latex_of(prefix: str, consts: list[float]) -> str:
    """Formule lisible, constantes numériques substituées."""
    node = from_prefix(prefix.split())
    names = const_name_iter()
    infix = to_infix(node, names)
    subs = {f"c{i}": round(float(c), 2) for i, c in enumerate(consts)}
    try:
        expr = sympy.sympify(infix).subs(subs)
        return f"${sympy.latex(expr)}$"
    except Exception:
        return infix


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--skeletons", type=Path, default=DATASET_DIR / "skeletons-v1.jsonl.gz")
    ap.add_argument("--seed", type=int, default=DEFAULT_SEED)

    ap.add_argument("--out", type=Path, default=FIGURES_DIR / "phase1-echantillons.png")
    args = ap.parse_args(argv)

    ensure_dirs()
    rng = make_rng(args.seed)
    items = load_skeletons(args.skeletons)
    depths = [it["depth"] for it in items]
    weights = stratified_weights(depths)
    print("distribution de profondeur visée à l'entraînement :", describe_weights(depths, weights))

    picks = [
        ("tirage uniforme sur les uniques", rng.choice(len(items), size=5, replace=False)),
        ("tirage stratifié", rng.choice(len(items), size=5, replace=False, p=weights)),
    ]

    fig, axes = plt.subplots(5, 2, figsize=(13, 15))
    fig.suptitle(
        "Curvy — Phase 1 : 10 exemples du dataset\n"
        "à gauche, tirage uniforme sur les squelettes uniques ; "
        "à droite, tirage stratifié par profondeur",
        fontsize=13,
    )
    dense_x = np.linspace(-1.0, 1.0, 400)

    for col, (label, idxs) in enumerate(picks):
        for row, i in enumerate(idxs):
            ax = axes[row, col]
            it = items[int(i)]
            node = from_prefix(it["prefix"].split())
            for _ in range(30):
                cloud, _ = sample_cloud(rng, node, CloudConfig())
                if cloud is not None:
                    break
            if cloud is None:
                ax.set_axis_off()
                continue
            clean = evaluate(node, dense_x, cloud.consts)
            clean_n = np.where(np.isfinite(clean), clean, np.nan)
            clean_n = (clean_n - cloud.y_offset) / cloud.y_scale
            ax.plot(dense_x, clean_n, lw=1.2, color="tab:orange", alpha=0.9, label="courbe exacte")
            ax.scatter(cloud.x, cloud.y, s=9, color="tab:blue", alpha=0.75, label="points bruités")
            ax.set_title(
                f"{latex_of(it['prefix'], cloud.consts)}\n"
                f"prof. {it['depth']} · {it['n_consts']} const. · "
                f"complexité {it['complexity']} · mult. {it['count']} · {cloud.n_points} pts",
                fontsize=8.5,
            )
            ax.set_ylim(-1.6, 1.6)
            ax.tick_params(labelsize=7)
            if row == 0:
                ax.legend(fontsize=7, loc="upper right")
            if row == 0:
                ax.text(
                    0.02,
                    1.28,
                    label,
                    transform=ax.transAxes,
                    fontsize=11,
                    fontweight="bold",
                    color="tab:red",
                )
    fig.tight_layout(rect=(0, 0, 1, 0.96))
    fig.savefig(args.out, dpi=140)
    fig.savefig(args.out.with_suffix(".svg"))
    print(f"figure écrite : {args.out}")
    print(f"               {args.out.with_suffix('.svg')}")

    for label, idxs in picks:
        print(f"\n--- {label} ---")
        for i in idxs:
            it = items[int(i)]
            print(f"  prof {it['depth']} · mult {it['count']:>7} · {it['prefix']}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

"""Que sait dessiner la grammaire v1 ? Sonde d'expressivité.

Le modèle ne peut trouver que ce que sa grammaire contient. Ce script prend des
formes que tout le monde a en tête — cercle, cœur, carré — et vérifie, pour
chacune, si elle est **exprimable** dans la grammaire v1 : arbre valide,
profondeur de corps ≤ 6, constantes internes ≤ 5.

C'est une question sur la grammaire, pas sur le modèle : indépendante de
l'entraînement, vérifiable tout de suite.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

from curvy.config import FIGURES_DIR, ensure_dirs  # noqa: E402
from curvy.data.canonical import canonicalise, strip_absorbable_root, wrap_root  # noqa: E402
from curvy.data.expr import (  # noqa: E402
    count_constants,
    depth,
    evaluate,
    from_prefix,
    to_prefix,
)
from curvy.data.grammar import MAX_BODY_CONSTANTS, MAX_BODY_DEPTH  # noqa: E402


def check(prefix: str) -> dict:
    """Le squelette tient-il dans les budgets de la grammaire v1 ?"""
    try:
        body = from_prefix(prefix.split())
    except ValueError as exc:
        return {"ok": False, "raison": f"non parsable : {exc}"}
    canon = strip_absorbable_root(canonicalise(body))
    d, k = depth(canon), count_constants(canon)
    ok = d <= MAX_BODY_DEPTH and k <= MAX_BODY_CONSTANTS
    raison = ""
    if d > MAX_BODY_DEPTH:
        raison += f"profondeur {d} > {MAX_BODY_DEPTH} ; "
    if k > MAX_BODY_CONSTANTS:
        raison += f"{k} constantes internes > {MAX_BODY_CONSTANTS} ; "
    return {
        "ok": ok,
        "profondeur_corps": d,
        "constantes_internes": k,
        "squelette": " ".join(to_prefix(wrap_root(canon))),
        "raison": raison.removesuffix(" ; "),
    }


#: (nom, x(t) en préfixe, constantes de x, y(t) en préfixe, constantes de y, commentaire)
PARAMETRIQUES = [
    (
        "Cercle",
        "mul C cos x",
        [1.0],
        "mul C sin x",
        [1.0],
        "x = C·cos t, y = C·sin t",
    ),
    (
        "Cœur (cardioïde)",
        "mul cos x sub C sin x",
        [1.0],
        "mul sin x sub C sin x",
        [1.0],
        "x = (C − sin t)·cos t, y = (C − sin t)·sin t",
    ),
    (
        "Cœur « classique »",
        "cube sin x",
        [],
        "sub sub sub mul C cos x mul C cos mul C x mul C cos mul C x mul C cos mul C x",
        [13 / 16, 5 / 16, 2.0, 2 / 16, 3.0, 1 / 16, 4.0],
        "x = sin³t, y = 13cos t − 5cos 2t − 2cos 3t − cos 4t",
    ),
    (
        "Lemniscate (∞)",
        "mul C cos x",
        [1.0],
        "mul C mul sin x cos x",
        [1.4],
        "x = C·cos t, y = C·sin t·cos t",
    ),
    (
        "Astroïde (carré à côtés creux)",
        "cube cos x",
        [],
        "cube sin x",
        [],
        "x = cos³t, y = sin³t",
    ),
    (
        # PIÈGE : sign(cos t), sign(sin t) ne prend que QUATRE valeurs — les
        # quatre coins. La « figure de carré » qu'on croit voir n'est que les
        # segments tracés entre ces quatre points. Conservé exprès : c'est
        # exactement le genre d'erreur qu'une figure fait commettre.
        "« Carré » par le signe (PIÈGE : 4 points)",
        "mul abs cos x inv cos x",
        [],
        "mul abs sin x inv sin x",
        [],
        "sign(cos t), sign(sin t) — 4 valeurs seulement, pas une courbe",
    ),
    (
        "Carré (vraie construction, norme max)",
        "mul cos x inv mul C add add abs cos x abs sin x abs sub abs cos x abs sin x",
        [0.5],
        "mul sin x inv mul C add add abs cos x abs sin x abs sub abs cos x abs sin x",
        [0.5],
        "x = cos t / max(|cos t|,|sin t|), avec max(a,b) = (a+b+|a−b|)/2",
    ),
]

#: Formes définies comme y = f(x), donc univaluées.
FONCTIONS = [
    ("Squircle d'ordre 4", "sqrt sqrt sub C sq sq x", [1.0], "y = (1 − x⁴)^(1/4)"),
    ("Squircle d'ordre 8", "sqrt sqrt sqrt sub C sq sq sq x", [1.0], "y = (1 − x⁸)^(1/8)"),
    ("Demi-cercle", "sqrt sub C sq x", [1.0], "y = √(1 − x²)"),
]


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--out", type=Path, default=FIGURES_DIR / "grammaire-formes.png")
    args = ap.parse_args(argv)
    ensure_dirs()

    print("=" * 78)
    print("FORMES PARAMÉTRIQUES  x(t), y(t)  —  nécessitent le mode paramétrique (Phase 7)")
    print("=" * 78)
    results = []
    for name, px, cx, py, cy, note in PARAMETRIQUES:
        rx, ry = check(px), check(py)
        ok = rx["ok"] and ry["ok"]
        results.append((name, px, cx, py, cy, note, ok, rx, ry))
        flag = "OUI" if ok else "NON"
        print(f"\n[{flag}] {name}  —  {note}")
        for label, r in (("x(t)", rx), ("y(t)", ry)):
            det = f"profondeur {r.get('profondeur_corps')}, {r.get('constantes_internes')} const."
            print(f"      {label} : {det}" + (f"  -> {r['raison']}" if r.get("raison") else ""))

    print("\n" + "=" * 78)
    print("FORMES UNIVALUÉES  y = f(x)  —  exprimables dès la v1")
    print("=" * 78)
    for name, pf, _cf, note in FONCTIONS:
        r = check(pf)
        flag = "OUI" if r["ok"] else "NON"
        print(f"\n[{flag}] {name}  —  {note}")
        print(
            f"      profondeur {r.get('profondeur_corps')}, "
            f"{r.get('constantes_internes')} const."
            + (f"  -> {r['raison']}" if r.get("raison") else "")
        )

    # --- figure ---
    n = len(PARAMETRIQUES) + len(FONCTIONS)
    cols = 3
    rows = (n + cols - 1) // cols
    fig, axes = plt.subplots(rows, cols, figsize=(13, 4.2 * rows))
    axes = np.atleast_2d(axes).ravel()
    t = np.linspace(-np.pi, np.pi, 2000)

    for ax, (name, px, cx, py, cy, note, ok, _rx, _ry) in zip(axes, results, strict=False):
        with np.errstate(all="ignore"):
            xs = evaluate(from_prefix(px.split()), t, cx)
            ys = evaluate(from_prefix(py.split()), t, cy)
        good = np.isfinite(xs) & np.isfinite(ys)
        colour = "tab:green" if ok else "tab:red"
        ax.plot(xs[good], ys[good], lw=1.6, color=colour)
        ax.set_title(
            f"{'✓' if ok else '✗'} {name}\n{note}",
            fontsize=9,
            color="black" if ok else "tab:red",
        )
        ax.set_aspect("equal", adjustable="datalim")
        ax.grid(alpha=0.3)
        ax.tick_params(labelsize=7)

    xg = np.linspace(-0.999, 0.999, 2000)
    for ax, (name, pf, cf, note) in zip(axes[len(results) :], FONCTIONS, strict=False):
        with np.errstate(all="ignore"):
            y = evaluate(from_prefix(pf.split()), xg, cf)
        r = check(pf)
        colour = "tab:green" if r["ok"] else "tab:red"
        ax.plot(xg, y, lw=1.6, color=colour)
        ax.plot(xg, -y, lw=1.6, color=colour, alpha=0.6)
        ax.set_title(
            f"{'✓' if r['ok'] else '✗'} {name}\n{note}",
            fontsize=9,
            color="black" if r["ok"] else "tab:red",
        )
        ax.set_aspect("equal", adjustable="datalim")
        ax.grid(alpha=0.3)
        ax.tick_params(labelsize=7)

    for ax in axes[n:]:
        ax.set_axis_off()

    fig.suptitle(
        "Ce que la grammaire v1 sait dessiner (vert) et ne sait pas (rouge)\n"
        "profondeur de corps ≤ 6, constantes internes ≤ 5",
        fontsize=12,
    )
    fig.tight_layout(rect=(0, 0, 1, 0.95))
    fig.savefig(args.out, dpi=140)
    fig.savefig(args.out.with_suffix(".svg"))
    print(f"\nfigure écrite : {args.out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())