import math, json, time
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score


def mobius(n):
    x, distinct = n, 0
    p = 2
    while p * p <= x:
        if x % p == 0:
            x //= p; distinct += 1
            if x % p == 0: return 0
            while x % p == 0: x //= p
        p += 1
    if x > 1: distinct += 1
    return -1 if distinct % 2 else 1


def lie_dim(k, d):
    return sum(mobius(q) * d ** (k // q) for q in range(1, k + 1) if k % q == 0) // k


def tail(z, N):
    term, partial = 1.0, 1.0
    for k in range(1, N + 1):
        term *= z / k
        partial += term
    return math.exp(z) - partial


def proxy(ells, r, N):
    return float(sum(tail(r * e, N) for e in ells))


def cost(m, N, d):
    return m * sum(lie_dim(k, d) for k in range(1, N + 1))


def select(ells, d, r, eps, Nmax=5):
    m, = [len(ells)]
    feasible = [(cost(m, N, d), N, proxy(ells, r, N))
                for N in range(1, Nmax + 1) if proxy(ells, r, N) <= eps]
    return min(feasible) if feasible else None


def variation_quantile_segments(x, m):
    ds = np.linalg.norm(np.diff(x, axis=0), axis=1)
    v = np.r_[0.0, np.cumsum(ds)]
    cuts = [0] + [int(np.searchsorted(v, q * v[-1] / m, side="left")) for q in range(1, m)] + [len(x)-1]
    cuts = np.maximum.accumulate(np.asarray(cuts)); cuts[-1] = len(x)-1
    ells = np.array([v[cuts[j+1]] - v[cuts[j]] for j in range(m)])
    return cuts, ells


def local_features(x, cuts, degree=2):
    # First-level increments and antisymmetric degree-2 log-signature area.
    feats = []
    for a, b in zip(cuts[:-1], cuts[1:]):
        p = x[a:b+1]; dx = np.diff(p, axis=0)
        inc = p[-1] - p[0]
        f = list(inc) + [float(b-a) / max(1, len(x)-1)]
        if degree >= 2:
            area = np.zeros((x.shape[1], x.shape[1]))
            prefix = np.zeros(x.shape[1])
            for u in dx:
                area += 0.5 * (np.outer(prefix, u) - np.outer(u, prefix))
                prefix += u
            f += [area[0,1]] if x.shape[1] >= 2 else []
        feats.extend(f)
    return np.asarray(feats, dtype=np.float32)


def make_data(n=700, length=64, seed=7):
    rng = np.random.default_rng(seed); X=[]; y=[]
    t=np.linspace(0,1,length)
    for i in range(n):
        cls=i%2; freq=3 if cls==0 else 7
        phase=rng.uniform(0,2*np.pi); amp=1+rng.normal(0,.04)
        s=amp*np.sin(2*np.pi*freq*t+phase)+rng.normal(0,.10,length)
        # second channel makes signed area informative for local signatures
        q=amp*np.cos(2*np.pi*freq*t+phase)+rng.normal(0,.10,length)
        X.append(np.stack([s,q],1)); y.append(cls)
    return np.asarray(X),np.asarray(y)


def features_fixed(X, m, degree):
    out=[]
    for x in X:
        cuts=np.linspace(0,len(x)-1,m+1).astype(int)
        out.append(local_features(x,cuts,degree))
    return np.asarray(out)


def run_classifier(X,y,F):
    a,b=train_test_split(np.arange(len(y)),test_size=.3,random_state=19,stratify=y)
    clf=LogisticRegression(max_iter=500,C=1.0).fit(F[a],y[a])
    return accuracy_score(y[b],clf.predict(F[b]))


def main():
    np.random.seed(7)
    # Core math verification: exact equal-variation predictions.
    print('LIE_DIMS', [lie_dim(k,2) for k in range(1,6)])
    print('TAIL_SWEEP')
    rows=[]
    for z in [0.25,0.5,1.0,2.0,3.0]:
        vals=[tail(z,N) for N in range(1,6)]
        rows.append((z,vals))
        print(z, ['%.8g'%v for v in vals])
    # Prediction 1: for fixed total rL, equal m gives m R_N(rL/m), decreasing with m.
    print('PREDICTION_EQUAL_M')
    for ztot in [1.,3.,6.]:
        vals=[proxy([ztot/m]*m,1.,2) for m in [1,2,4,8,16]]
        print(ztot, ['%.8g'%v for v in vals])
    # Prediction 2: raising N strictly decreases tail, with factorial-like ratios.
    print('PREDICTION_DEGREE')
    for z in [.5,1.,2.]:
        vals=[tail(z,N) for N in range(1,6)]
        print(z, ['%.8g'%v for v in vals], 'ratios', ['%.5g'%(vals[i+1]/vals[i]) for i in range(4)])
    # Prediction 3: adaptive choice transitions toward larger m for oscillatory/high variation paths.
    print('SELECTION_SWEEP')
    for L in [1.,2.,4.,8.]:
        eps=.05; picks=[]
        for m in [1,2,4,8,16,32]:
            # equal-variation proxy, retaining lowest cost across degrees
            cand=select([L/m]*m,2,1.,eps)
            if cand: picks.append((m,cand[1],cand[0],cand[2]))
        print('L',L,'feasible_lowest_per_m',picks)
    X,y=make_data()
    results={}
    t=time.time(); results['raw_8patch']=run_classifier(X,y,features_fixed(X,8,1))
    results['fixed_sig_4_degree2']=run_classifier(X,y,features_fixed(X,4,2))
    # adaptive: choose degree using observed proxy, and use 8 local tokens; representation is compact.
    adaptive=[]; choices=[]
    for x in X:
        cuts,ells=variation_quantile_segments(x,8)
        c=select(ells,2,1.,eps=.08)
        N=2 if c is None else c[1]
        choices.append(N); adaptive.append(local_features(x,cuts,N))
    results['adaptive_8_local']=run_classifier(X,y,np.asarray(adaptive))
    results['adaptive_degree_counts']={str(k):int(choices.count(k)) for k in sorted(set(choices))}
    results['elapsed_sec']=time.time()-t
    print('RESULTS',json.dumps(results,sort_keys=True))

if __name__=='__main__': main()