"""Laboratoire d'indicateurs : invente des indicateurs inedits par evolution de formules (programmation genetique).

Ingredients : tous les etats du marche deja calcules (champs Polymarket des 4 couches, pack technique,
positionnement, primes, sentiment), chacun ramene a un z-score. Operations : addition, difference, produit,
rapport, min/max, signe, tangente hyperbolique, et "porte" (A si B > 0, sinon -A) pour les interactions.

A chaque cycle d'apprentissage :
  1. les indicateurs deja inventes sont re-evalues UNIQUEMENT sur les donnees arrivees apres leur invention
     (le seul vrai test) : valide s'ils tiennent, retire s'ils cessent de marcher -> adaptation continue ;
  2. une population de formules evolue (selection, croisement, mutation) sur le debut de la fenetre, et les
     meilleures doivent confirmer sur la fin (validation) ; les formules qui ne font que recopier un indicateur
     existant sont rejetees ;
  3. les indicateurs valides deviennent des entrees de l'indicateur composite et du meta-etiquetage.
La note d'une formule = IC partiel (correlation de rang avec le rendement futur, momentum recent retire),
penalisee par sa complexite : on prefere une formule simple qui tient a une formule baroque qui colle au passe.
"""
import json
import math
import random
import time

import numpy as np
from scipy.stats import rankdata

UNARY = ("neg", "abs", "sign", "tanh", "sq", "slog")
BINARY = ("add", "sub", "mul", "div", "max", "min", "gate")
MAX_SIZE, MAX_DEPTH = 13, 4
CONTROLS = ("ta_r1z", "ta_r15z")
SCHEMA = """
CREATE TABLE IF NOT EXISTS lab(
  id INTEGER PRIMARY KEY, formula TEXT UNIQUE, horizon TEXT, norms TEXT, created INTEGER,
  sign INTEGER, ic_train REAL, ic_val REAL, ic_fresh REAL, t_fresh REAL, fresh_days REAL,
  status TEXT, updated INTEGER          -- status : candidat | valide | retire
);
"""


# ---------------------------------------------------------------- formules

def apply(op, a, b=None):
    with np.errstate(all="ignore"):
        if op == "neg":
            r = -a
        elif op == "abs":
            r = np.abs(a)
        elif op == "sign":
            r = np.sign(a)
        elif op == "tanh":
            r = np.tanh(a)
        elif op == "sq":
            r = np.sign(a) * a * a
        elif op == "slog":
            r = np.sign(a) * np.log1p(np.abs(a))
        elif op == "add":
            r = a + b
        elif op == "sub":
            r = a - b
        elif op == "mul":
            r = a * b
        elif op == "div":
            r = a / (1 + np.abs(b))
        elif op == "max":
            r = np.maximum(a, b)
        elif op == "min":
            r = np.minimum(a, b)
        else:  # gate : interaction "A si B > 0, sinon -A"
            r = np.where(b > 0, a, -a)
    return np.clip(np.nan_to_num(r, nan=0.0, posinf=0.0, neginf=0.0), -1e6, 1e6)


def evaluate(f, terms):
    if f[0] == "x":
        return terms[f[1]]
    if f[0] in UNARY:
        return apply(f[0], evaluate(f[1], terms))
    return apply(f[0], evaluate(f[1], terms), evaluate(f[2], terms))


def size(f):
    return 1 if f[0] == "x" else 1 + sum(size(c) for c in f[1:])


def depth(f):
    return 0 if f[0] == "x" else 1 + max(depth(c) for c in f[1:])


def leaves(f):
    return {f[1]} if f[0] == "x" else set().union(*(leaves(c) for c in f[1:]))


def paths(f, p=()):
    yield p
    if f[0] != "x":
        for i, c in enumerate(f[1:], start=1):
            yield from paths(c, p + (i,))


def get(f, p):
    for i in p:
        f = f[i]
    return f


def put(f, p, sub):
    if not p:
        return sub
    f = list(f)
    f[p[0]] = put(f[p[0]], p[1:], sub)
    return f


def grow(rng, names, d):
    if d == 0 or rng.random() < 0.3:
        return ["x", rng.choice(names)]
    if rng.random() < 0.3:
        return [rng.choice(UNARY), grow(rng, names, d - 1)]
    return [rng.choice(BINARY), grow(rng, names, d - 1), grow(rng, names, d - 1)]


def mutate(rng, f, names):
    p = rng.choice(list(paths(f)))
    node = get(f, p)
    r = rng.random()
    if node[0] == "x" and r < 0.5:
        new = ["x", rng.choice(names)]
    elif node[0] in UNARY and r < 0.5:
        new = [rng.choice(UNARY), node[1]]
    elif node[0] in BINARY and r < 0.5:
        new = [rng.choice(BINARY), node[1], node[2]]
    else:
        new = grow(rng, names, 2)
    return put(f, p, new)


def crossover(rng, a, b):
    return put(a, rng.choice(list(paths(a))), get(b, rng.choice(list(paths(b)))))


def ok_shape(f):
    return size(f) <= MAX_SIZE and depth(f) <= MAX_DEPTH and f[0] != "x"


def render(f, label):
    if f[0] == "x":
        return label(f[1])
    sym = {"neg": "-{}", "abs": "|{}|", "sign": "signe({})", "tanh": "tanh({})", "sq": "carre({})", "slog": "log({})",
           "add": "({} + {})", "sub": "({} - {})", "mul": "({} x {})", "div": "({} / (1+|{}|))",
           "max": "max({}, {})", "min": "min({}, {})", "gate": "({} si {} > 0, sinon l'inverse)"}[f[0]]
    return sym.format(*(render(c, label) for c in f[1:]))


# ---------------------------------------------------------------- donnees et notation

def column(rows, f):
    return np.array([np.nan if r.get(f) is None else r[f] for r in rows], dtype=float)


def zscore(a, mean, sd):
    z = np.clip((a - mean) / sd, -5, 5)
    z[np.isnan(z)] = 0.0
    return z


class Scorer:
    """IC partiel rapide (rangs, momentum recent retire) d'une serie de valeurs contre le rendement futur."""

    def __init__(self, rows, ycol):
        y = column(rows, ycol)
        C = np.column_stack([np.ones(len(rows))] + [rankdata(np.nan_to_num(column(rows, c))) for c in CONTROLS])
        self.q, _ = np.linalg.qr(C)
        self.ry = self.res(rankdata(y))
        self.ny = np.linalg.norm(self.ry)

    def res(self, v):
        return v - self.q @ (self.q.T @ v)

    def ic(self, x):
        rx = self.res(rankdata(x))
        n = np.linalg.norm(rx)
        return float(rx @ self.ry / (n * self.ny)) if n > 0 and self.ny > 0 else 0.0


def t_stat(ic, n_eff):
    return ic * math.sqrt(max(n_eff - 2, 1)) / math.sqrt(max(1 - ic * ic, 1e-9))


def n_eff_of(rows, x, step, X):
    """Episodes independants : limites par le chevauchement des rendements et par la lenteur de l'indicateur."""
    times = len({r["t"] for r in rows})
    first_asset = rows[0]["asset"]
    seq = np.array([v for r, v in zip(rows, x) if r["asset"] == first_asset])
    tau = 1
    if len(seq) > 50 and seq.std() > 0:
        s = (seq - seq.mean()) / seq.std()
        lag = 1
        while lag < len(s) // 2:
            if float(np.mean(s[:-lag] * s[lag:])) < 1 / math.e:
                break
            lag *= 2
        tau = lag
    return min(times * min(1.0, step / X), times / tau)


# ---------------------------------------------------------------- cycle du laboratoire

def load(con, status=("candidat", "valide")):
    con.executescript(SCHEMA)
    q = ",".join("?" * len(status))
    return [{"id": i, "formula": json.loads(f), "horizon": h, "norms": json.loads(n), "created": c, "sign": s,
             "ic_train": a, "ic_val": b, "ic_fresh": fr, "t_fresh": tf, "fresh_days": fd, "status": st}
            for i, f, h, n, c, s, a, b, fr, tf, fd, st in con.execute(
                f"SELECT id, formula, horizon, norms, created, sign, ic_train, ic_val, ic_fresh, t_fresh, fresh_days, "
                f"status FROM lab WHERE status IN ({q}) ORDER BY id", status)]


def values(spec, rows):
    terms = {name: zscore(column(rows, name), m, sd) for name, (m, sd) in spec["norms"].items()}
    return evaluate(spec["formula"], terms)


def add_features(rows, specs):
    """Ajoute aux lignes les indicateurs valides (lab_<id>) ; retourne leurs noms."""
    names = []
    for spec in specs:
        v = values(spec, rows)
        name = f"lab_{spec['id']}"
        for r, x in zip(rows, v):
            r[name] = float(x)
        names.append(name)
    return names


def value_row(spec, r):
    """Valeur d'un indicateur du laboratoire pour une seule ligne (moteur en direct)."""
    return float(values(spec, [r])[0])


def reassess(con, rows, step, horizons, now):
    """Le seul vrai test : chaque indicateur n'est juge que sur les donnees posterieures a son invention."""
    out = []
    for spec in load(con):
        X = horizons.get(spec["horizon"])
        ycol = f"y_{spec['horizon']}"
        fresh = [r for r in rows if r["t"] > spec["created"] and r.get(ycol) is not None
                 and all(r.get(c) is not None for c in CONTROLS)]
        days = len({r["t"] for r in fresh}) * step / 86400
        status, ic, t = spec["status"], None, None
        if X and len(fresh) >= 300:
            x = values(spec, fresh)
            ic = Scorer(fresh, ycol).ic(x) * spec["sign"]  # positif = va dans le sens decouvert
            t = t_stat(ic, n_eff_of(fresh, x, step, X))
            if days >= 3 and t >= 2:
                status = "valide"
            elif days >= 7 and t < 1:
                status = "retire"  # a cesse de marcher (ou n'a jamais marche hors de son echantillon)
        con.execute("UPDATE lab SET ic_fresh = ?, t_fresh = ?, fresh_days = ?, status = ?, updated = ? WHERE id = ?",
                    (ic, t, days, status, now, spec["id"]))
        spec.update(ic_fresh=ic, t_fresh=t, fresh_days=days, status=status)
        out.append(spec)
    con.commit()
    return out


def evolve(rows, feats, h, X, step, rng, pop_size=120, generations=12, label=str):
    """Une campagne d'evolution pour l'horizon h. Retourne les formules qui confirment en validation."""
    ycol = f"y_{h}"
    data = [r for r in rows if r.get(ycol) is not None and all(r.get(c) is not None for c in CONTROLS)]
    if len(data) < 2000:
        return []
    times = sorted({r["t"] for r in data})
    cut = times[int(len(times) * 0.7)]
    train = [r for r in data if r["t"] + X <= cut]
    val = [r for r in data if r["t"] >= cut]
    if len(train) < 1000 or len(val) < 300:
        return []
    sub = train[::max(1, len(train) // 40000)]  # sous-echantillon pour la vitesse
    names, norms, T_sub, T_val = [], {}, {}, {}
    for f in feats:
        a = column(train, f)
        if np.isnan(a).mean() > 0.5:
            continue
        m, sd = float(np.nanmean(a)), float(np.nanstd(a))
        if not sd or not np.isfinite(sd):
            continue
        names.append(f)
        norms[f] = (m, sd)
        T_sub[f] = zscore(column(sub, f), m, sd)
        T_val[f] = zscore(column(val, f), m, sd)
    if len(names) < 3:
        return []
    sc = Scorer(sub, ycol)
    cache = {}

    def fitness(f):
        k = json.dumps(f)
        if k not in cache:
            ic = sc.ic(evaluate(f, T_sub))
            cache[k] = (abs(ic) - 0.003 * size(f), ic)
        return cache[k]

    pop = [grow(rng, names, 3) for _ in range(pop_size)]
    pop = [f for f in pop if ok_shape(f)] or [["neg", ["x", names[0]]]]
    for _ in range(generations):
        scored = sorted(pop, key=lambda f: -fitness(f)[0])
        elite = scored[:pop_size // 10]
        nxt = list(elite)
        while len(nxt) < pop_size:
            a = max(rng.sample(scored, 4), key=lambda f: fitness(f)[0])  # tournoi
            if rng.random() < 0.6:
                b = max(rng.sample(scored, 4), key=lambda f: fitness(f)[0])
                child = crossover(rng, a, b)
            else:
                child = mutate(rng, a, names)
            if ok_shape(child):
                nxt.append(child)
        pop = nxt
    best, seen = [], set()
    for f in sorted(pop, key=lambda f: -fitness(f)[0]):
        k = json.dumps(f)
        if k in seen:
            continue
        seen.add(k)
        best.append(f)
        if len(best) >= 8:
            break
    sv = Scorer(val, ycol)
    found = []
    for f in best:
        fit, ic_tr = fitness(f)
        xv = evaluate(f, T_val)
        ic_v = sv.ic(xv)
        t_v = t_stat(ic_v * np.sign(ic_tr), n_eff_of(val, xv, step, X))
        if t_v < 2.5:
            continue  # ne confirme pas sur la periode de validation
        # rejeter une simple copie d'un indicateur existant : ce ne serait pas une invention
        xs = evaluate(f, T_sub)
        if xs.std() == 0:
            continue
        closest = max((abs(float(np.corrcoef(xs, T_sub[n])[0, 1])) if T_sub[n].std() > 0 else 0, n) for n in names)
        if closest[0] > 0.95:
            continue
        used = {n: norms[n] for n in leaves(f)}
        found.append({"formula": f, "horizon": h, "norms": used, "sign": int(np.sign(ic_tr)) or 1,
                      "ic_train": ic_tr, "ic_val": ic_v, "t_val": t_v, "closest": closest,
                      "text": render(f, label)})
    return found


def run(con, rows, feats, step, horizons, save, label=str, max_active=40):
    """Cycle complet. Retourne (indicateurs valides a utiliser comme entrees, lignes de rapport)."""
    now = int(time.time())
    con.executescript(SCHEMA)
    report = []
    if save:
        reassess(con, rows, step, horizons, now)
        rng = random.Random(now)
        base = [f for f in feats if not f.startswith(("lab_", "tg_"))]
        for h, X in horizons.items():
            for c in evolve(rows, base, h, X, step, rng, label=label):
                cur = con.execute("INSERT OR IGNORE INTO lab(formula, horizon, norms, created, sign, ic_train, "
                                  "ic_val, status, updated) VALUES(?,?,?,?,?,?,?,?,?)",
                                  (json.dumps(c["formula"]), h, json.dumps(c["norms"]), now, c["sign"],
                                   c["ic_train"], c["ic_val"], "candidat", now))
                if cur.rowcount:
                    report.append(f"  + nouveau candidat ({h}) : {c['text']}  [IC entrainement {c['ic_train']:+.3f}, "
                                  f"validation {c['ic_val']:+.3f}, t={c['t_val']:.1f}]")
        # ne garder que les candidats les plus prometteurs
        cand = con.execute("SELECT id FROM lab WHERE status = 'candidat' ORDER BY ABS(ic_val) DESC").fetchall()
        for (i,) in cand[max_active:]:
            con.execute("UPDATE lab SET status = 'retire', updated = ? WHERE id = ?", (now, i))
        con.commit()
    valid = load(con, ("valide",))
    for s in sorted(load(con), key=lambda s: (s["status"] != "valide", -(s["t_fresh"] or 0))):
        fresh = (f"apres invention : IC {s['ic_fresh']:+.3f} t={s['t_fresh']:+.1f} sur {s['fresh_days']:.1f} j"
                 if s["ic_fresh"] is not None else "apres invention : pas encore assez de donnees")
        report.append(f"  [{s['status']:<8}] lab_{s['id']:<4} {s['horizon']:<4} {render(s['formula'], label)}  | "
                      f"validation {s['ic_val']:+.3f} | {fresh}")
    return valid, report