"""Candidate rules. Each returns rule(i) -> weights, using only data up to i."""
import math, rotmap as R, lab

def ratio(M, a, b="BTC"):
    return [(M.idx[a][i] / M.idx[b][i]) if M.idx[a][i] and M.idx[b][i] else None for i in range(M.N)]

def ma(x, i, n):
    v = [y for y in x[max(0, i - n + 1):i + 1] if y is not None]
    return sum(v) / len(v) if len(v) >= n * 0.8 else None

def trend_state(x, n, band):
    out, s = [], False
    for i in range(len(x)):
        m = ma(x, i, n)
        if m and x[i]:
            if x[i] > m * (1 + band): s = True
            elif x[i] < m * (1 - band): s = False
        out.append(s)
    return out

def alts_w(M, i, keep=None):
    bs = [b for b in R.ALT_BASKETS if M.idx[b][i] and (keep is None or b in keep)]
    return {b: 1 / len(bs) for b in bs} if bs else {"BTC": 1.0}

def rs_gate(M, n=50, band=0.02, src="ALTS"):
    """Risk on + alts/BTC ratio in an uptrend -> alts. Risk on only -> BTC."""
    if "ALTS" not in M.idx: M.idx["ALTS"] = lab.alts_idx(M)
    st = trend_state(ratio(M, src), n, band)
    return lambda i: ({} if not M.risk_on(i) else (alts_w(M, i) if st[i] else {"BTC": 1.0}))

def breadth_gate(M, on=0.55, off=0.40):
    s = {"on": False}
    def rule(i):
        c = M.c30[i]
        if c is not None:
            if c >= on: s["on"] = True
            elif c < off: s["on"] = False
        if not M.risk_on(i): return {}
        return alts_w(M, i) if s["on"] else {"BTC": 1.0}
    return rule

def book_plus_rs(M, n=50, band=0.02):
    """Book G, plus: risk on and alts/BTC in uptrend also opens alts; the
    alts/BTC downtrend closes them."""
    if "ALTS" not in M.idx: M.idx["ALTS"] = lab.alts_idx(M)
    st = trend_state(ratio(M, "ALTS"), n, band)
    g = {"on": False, "dropped": set()}
    def rule(i):
        up = M.risk_on(i)
        if not up: g["on"] = False
        elif not g["on"] and (M.alt_bottom(i) or st[i]): g["on"], g["dropped"] = True, set()
        elif g["on"] and not st[i] and not M.alt_bottom(i): g["on"] = False
        if not g["on"]: return {"BTC": 1.0} if up else {}
        for b in R.ALT_BASKETS:
            if M.basket_top(b, i): g["dropped"].add(b)
        have = [b for b in R.ALT_BASKETS if M.idx[b][i]]
        live = [b for b in have if b not in g["dropped"]]
        w = {b: 1 / len(have) for b in live}
        if len(live) < len(have): w["BTC"] = 1 - sum(w.values())
        return w
    return rule

def basket_rs(M, n=50, band=0.02, top=None):
    """Risk on: hold every alt basket whose own ratio to BTC is in uptrend; the
    rest in BTC. Optionally only the top-k by 30d relative strength."""
    st = {b: trend_state(ratio(M, b), n, band) for b in R.ALT_BASKETS}
    def rule(i):
        if not M.risk_on(i): return {}
        on = [b for b in R.ALT_BASKETS if M.idx[b][i] and st[b][i]]
        if top and len(on) > top:
            on = sorted(on, key=lambda b: M.ret(b, i, 30) or -9)[-top:]
        if not on: return {"BTC": 1.0}
        k = max(len(on), 3)            # never more than a third per basket
        w = {b: 1 / k for b in on}
        if sum(w.values()) < 1: w["BTC"] = 1 - sum(w.values())
        return w
    return rule