import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# A rooted tree is a tuple of (edge_length, child); a leaf is ().
def poly_eval(tree, ys, x=1.0):
    ys = np.asarray(ys, dtype=float)
    if len(tree) == 0:
        return np.full(len(ys), x, dtype=float)
    out = np.ones(len(ys), dtype=float)
    for w, child in tree:
        out *= w * ys + poly_eval(child, ys, x)
    return out

def ordinary_sum(tree):
    # Standard scalar sum message: edge lengths plus terminal messages.
    if len(tree) == 0:
        return 1.0
    return sum(float(w) + ordinary_sum(c) for w, c in tree)

def count_edges(tree):
    return sum(1 + count_edges(c) for _, c in tree)

def fingerprint(tree, ys):
    z = poly_eval(tree, ys)
    return np.log(np.abs(z) + 1e-12)

def make_a(vals):
    # root -> two leaves and one internal node -> two leaves
    a,b,c,d,e = vals
    return ((a, ()), (b, ()), (c, ((d, ()), (e, ()))))

def make_b(vals):
    # root -> one leaf and one internal node -> three leaves
    a,b,c,d,e = vals
    return ((a, ()), (b, ((c, ()), (d, ()), (e, ()))))

def main():
    rng = np.random.default_rng(485)
    ys = np.array([-3., -2., -1., -.5, .5, 1., 2., 3.])
    heldout_y = np.array([-.8, -.3, .3, .8, 1.7, 2.4])

    # Core claim sanity check: same ordinary sum, distinct polynomial values.
    vals = np.array([.5, .8, 1.1, 1.4, 1.7])
    ta, tb = make_a(vals), make_b(vals)
    za, zb = poly_eval(ta, ys), poly_eval(tb, ys)
    collision = np.isclose(ordinary_sum(ta), ordinary_sum(tb))
    separation = float(np.max(np.abs(za - zb)))
    assert count_edges(ta) == count_edges(tb) == 5
    assert collision and separation > 1e-6

    Xbase, Xfp, Xfp_hold, y = [], [], [], []
    for label in (0, 1):
        for _ in range(600):
            # Both classes have the same edge count and leaf count, and use
            # the same edge-length multiset: scalar sum is an exact collision.
            v = rng.permutation(vals)
            t = make_a(v) if label == 0 else make_b(v)
            Xbase.append([ordinary_sum(t)])
            Xfp.append(fingerprint(t, ys))
            Xfp_hold.append(fingerprint(t, heldout_y))
            y.append(label)
    Xbase, Xfp, Xfp_hold, y = map(np.asarray, (Xbase, Xfp, Xfp_hold, y))
    tr, te = train_test_split(np.arange(len(y)), test_size=.3, random_state=17, stratify=y)
    base = make_pipeline(StandardScaler(), LogisticRegression(random_state=0, max_iter=1000))
    fp = make_pipeline(StandardScaler(), LogisticRegression(random_state=0, max_iter=1000))
    fp_hold = make_pipeline(StandardScaler(), LogisticRegression(random_state=0, max_iter=1000))
    base.fit(Xbase[tr], y[tr]); fp.fit(Xfp[tr], y[tr]); fp_hold.fit(Xfp_hold[tr], y[tr])
    acc_base = accuracy_score(y[te], base.predict(Xbase[te]))
    acc_fp = accuracy_score(y[te], fp.predict(Xfp[te]))
    acc_hold = accuracy_score(y[te], fp_hold.predict(Xfp_hold[te]))
    print({
        'ordinary_sum_collision': bool(collision),
        'example_max_polynomial_difference': separation,
        'baseline_accuracy': round(float(acc_base), 4),
        'fingerprint_accuracy': round(float(acc_fp), 4),
        'heldout_evaluation_grid_accuracy': round(float(acc_hold), 4),
        'n_test': len(te), 'M': len(ys),
    })

if __name__ == '__main__':
    main()