import json
import numpy as np
from boundary_radial import bars_from_radii, radial_loss, match_loss

np.set_printoptions(precision=6, suppress=True)

def fit_slope(x, y):
    return float(np.polyfit(np.asarray(x), np.asarray(y), 1)[0])

def main():
    target = [3.0, 5.0]
    # Prediction 1: identical boundary barcodes have exactly zero matching loss.
    zero = radial_loss(target, target, unmatched=1.0)

    # Prediction 2: a uniform radial displacement eps changes both endpoints,
    # hence each matched circle costs 2|eps|, and two circles cost 4|eps|.
    eps = np.linspace(0.0, 1.0, 11)
    perturbed = [radial_loss([3.0 + e, 5.0 + e], target, unmatched=10.0)
                 for e in eps]
    slope = fit_slope(eps[1:], perturbed[1:])
    max_abs_fit_error = float(np.max(np.abs(np.asarray(perturbed) - 4*eps)))

    # Prediction 3: with an empty target and unmatched penalty lambda, each
    # predicted interval contributes lambda, so the loss is q*lambda.
    lambdas = [0.25, 0.5, 1.0, 2.0]
    qs = [0, 1, 2, 3, 4]
    unmatched_table = []
    for lam in lambdas:
        vals = [radial_loss([3.0 + 2*j for j in range(q)], [], unmatched=lam)
                for q in qs]
        unmatched_table.append(vals)
    # Fit all nonzero q/lambda values to q*lambda.
    observed = np.array(unmatched_table)
    expected = np.array([[q*lam for q in qs] for lam in lambdas])
    unmatched_max_error = float(np.max(np.abs(observed-expected)))

    # Matching is order-invariant: a permutation of target intervals has zero cost.
    permutation_loss = radial_loss([3.0, 5.0, 7.0], [7.0, 3.0, 5.0], unmatched=10.0)

    # Small baseline-vs-idea proxy: noisy candidate boundaries. The standard
    # endpoint MSE and persistence L1 both recover the true radial vector here;
    # this intentionally reports a sanity comparison, not a CNN claim.
    rng = np.random.default_rng(7)
    true = np.array([3.0, 5.0, 7.0])
    noisy = true + rng.normal(0, 0.30, size=true.shape)
    baseline_endpoint_l1 = float(np.sum(np.abs(noisy-true)))
    idea_bar_loss = radial_loss(noisy.tolist(), true.tolist(), unmatched=10.0)

    result = {
        "zero_loss": zero,
        "perturbation_eps": eps.tolist(),
        "perturbation_loss": perturbed,
        "predicted_perturbation_slope": 4.0,
        "observed_perturbation_slope": slope,
        "perturbation_max_abs_fit_error": max_abs_fit_error,
        "unmatched_lambdas": lambdas,
        "unmatched_qs": qs,
        "unmatched_observed": observed.tolist(),
        "unmatched_expected": expected.tolist(),
        "unmatched_max_abs_error": unmatched_max_error,
        "permutation_loss": permutation_loss,
        "noisy_radii": noisy.tolist(),
        "baseline_endpoint_l1": baseline_endpoint_l1,
        "idea_persistence_loss": idea_bar_loss,
    }
    print(json.dumps(result, indent=2))
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)

    assert zero < 1e-12
    assert abs(slope - 4.0) < 1e-8 and max_abs_fit_error < 1e-8
    assert unmatched_max_error < 1e-12
    assert permutation_loss < 1e-12

if __name__ == "__main__":
    main()