package federated

// accountant_test.go pins the accountant math, because the accountant correctness IS the privacy claim. The
// tests check the per release RDP closed form, composition monotonicity, the RDP to (epsilon, delta)
// conversion, and a comparison against a published Google differential privacy RDP accountant number quoted
// as a FIXED CONSTANT (never by calling Python or any external tool). They also prove the noise sampler
// cannot run.

import (
	"context"
	"errors"
	"math"
	"testing"

	"github.com/missionworkforce/missionmatch/internal/platform/config"
)

// almostEqual reports whether a and b agree to the given absolute tolerance.
func almostEqual(a, b, tol float64) bool { return math.Abs(a-b) <= tol }

// TestPerReleaseRDPClosedForm pins the per release Gaussian RDP at hand computed (sigma, alpha) points. The
// closed form is alpha / (2 * sigma^2) for the L2 sensitivity 1 mechanism (noise multiplier sigma).
func TestPerReleaseRDPClosedForm(t *testing.T) {
	cases := []struct {
		sigma, alpha, want float64
	}{
		// sigma=2, alpha=2: 2 / (2 * 4) = 0.25.
		{2, 2, 0.25},
		// sigma=1, alpha=10: 10 / (2 * 1) = 5.
		{1, 10, 5},
		// sigma=0.5, alpha=3: 3 / (2 * 0.25) = 6.
		{0.5, 3, 6},
		// sigma=4, alpha=8: 8 / (2 * 16) = 0.25.
		{4, 8, 0.25},
	}
	for _, c := range cases {
		got := gaussianRDP(c.sigma, c.alpha)
		if !almostEqual(got, c.want, 1e-12) {
			t.Fatalf("gaussianRDP(sigma=%v, alpha=%v) = %v, want %v", c.sigma, c.alpha, got, c.want)
		}
	}
}

// TestNoiseMultiplierConvention pins the documented convention that the only quantity the RDP curve depends
// on is sigma = NoiseSigma / ClipNorm, the noise in units of the sensitivity. A release with NoiseSigma 4 and
// ClipNorm 2 has the same privacy as NoiseSigma 2 and ClipNorm 1, both multiplier 2.
func TestNoiseMultiplierConvention(t *testing.T) {
	a := GaussianRelease{NoiseSigma: 4, ClipNorm: 2}
	b := GaussianRelease{NoiseSigma: 2, ClipNorm: 1}
	if !almostEqual(a.noiseMultiplier(), b.noiseMultiplier(), 1e-12) {
		t.Fatalf("noise multiplier convention broken: %v vs %v", a.noiseMultiplier(), b.noiseMultiplier())
	}
	if !almostEqual(a.noiseMultiplier(), 2, 1e-12) {
		t.Fatalf("noise multiplier = %v, want 2", a.noiseMultiplier())
	}
}

// TestRDPToEpsilonConversionClosedForm pins the RDP to (epsilon, delta) conversion at one order against the
// hand computed closed form epsilon(alpha) = rdp(alpha) + log(1/delta) / (alpha - 1).
func TestRDPToEpsilonConversionClosedForm(t *testing.T) {
	// rdp=0.25 at alpha=2, delta=1e-5: 0.25 + ln(1e5) / 1 = 0.25 + 11.512925465 = 11.762925465.
	got := epsilonAtOrder(0.25, 2, 1e-5)
	want := 0.25 + math.Log(1/1e-5)/(2-1)
	if !almostEqual(got, want, 1e-12) {
		t.Fatalf("epsilonAtOrder(0.25, 2, 1e-5) = %v, want %v", got, want)
	}
	// A second independent point: rdp=5 at alpha=10, delta=1e-6: 5 + ln(1e6)/9.
	got2 := epsilonAtOrder(5, 10, 1e-6)
	want2 := 5 + math.Log(1/1e-6)/(10-1)
	if !almostEqual(got2, want2, 1e-12) {
		t.Fatalf("epsilonAtOrder(5, 10, 1e-6) = %v, want %v", got2, want2)
	}
}

// TestSingleReleaseEpsilonHandComputed pins the full min over orders conversion for one Gaussian release at a
// few (sigma, ClipNorm, delta) points against hand computed closed form values. The accountant grid is the
// shipped coarse grid; the hand value is the min over THAT grid (computed independently in the test loop), so
// this checks the accountant reproduces its own conversion exactly.
func TestSingleReleaseEpsilonHandComputed(t *testing.T) {
	cases := []struct {
		sigma, clip, delta float64
	}{
		{2, 1, 1e-5},
		{1, 1, 1e-5},
		{4, 2, 1e-6}, // multiplier 2
		{3, 1, 1e-7},
	}
	for _, c := range cases {
		acct, err := newRDPAccountant(100, c.delta, nil)
		if err != nil {
			t.Fatalf("newRDPAccountant: %v", err)
		}
		mult := c.sigma / c.clip
		// Hand compute the min over the same grid the accountant uses.
		want := math.Inf(1)
		for _, alpha := range acct.orders {
			rdp := alpha / (2 * mult * mult)
			eps := rdp + math.Log(1/c.delta)/(alpha-1)
			if eps < want {
				want = eps
			}
		}
		got, err := acct.Charge(GaussianRelease{NoiseSigma: c.sigma, ClipNorm: c.clip})
		if err != nil {
			t.Fatalf("Charge: %v", err)
		}
		if !almostEqual(got, want, 1e-9) {
			t.Fatalf("single release eps (sigma=%v clip=%v delta=%v) = %v, want %v", c.sigma, c.clip, c.delta, got, want)
		}
	}
}

// TestCompositionMonotonicity asserts the reported epsilon is strictly nondecreasing as releases accumulate,
// over a sequence of identical releases. Each release adds positive RDP at every order, so the accumulation
// grows and the converted epsilon can only rise.
func TestCompositionMonotonicity(t *testing.T) {
	acct, err := newRDPAccountant(1e9, 1e-6, nil) // huge budget so no refusal interferes with the monotonicity check
	if err != nil {
		t.Fatalf("newRDPAccountant: %v", err)
	}
	prev := 0.0
	for i := 0; i < 200; i++ {
		eps, err := acct.Charge(GaussianRelease{NoiseSigma: 2, ClipNorm: 1})
		if err != nil {
			t.Fatalf("Charge at release %d: %v", i, err)
		}
		if eps < prev {
			t.Fatalf("epsilon decreased at release %d: %v < %v", i, eps, prev)
		}
		if i > 0 && eps <= prev {
			t.Fatalf("epsilon not strictly increasing at release %d: %v <= %v", i, eps, prev)
		}
		prev = eps
	}
}

// TestPublishedGoogleRDPAccountant compares the accountant against a published Google differential privacy
// RDP accountant number, QUOTED AS A FIXED CONSTANT (never by calling any external tool).
//
// SOURCE AND DERIVATION. For the Gaussian mechanism with noise multiplier sigma and no subsampling (q = 1),
// the Renyi differential privacy at order alpha is alpha / (2 * sigma^2), and the Google dp_accounting
// RdpAccountant converts accumulated RDP to (epsilon, delta) by epsilon = min over alpha of
// rdp(alpha) + log(1/delta) / (alpha - 1) (Mironov 2017, Proposition 3 and Corollary 3; the same formula the
// Google library uses in python/dp_accounting/rdp/rdp_privacy_accountant.py). Cited in docs/AI_RESEARCH.md,
// the federated learning with differential privacy section:
//
//	https://github.com/google/differential-privacy/tree/main/python/dp_accounting
//	Mironov, Renyi Differential Privacy, CSF 2017, arXiv:1702.07476.
//
// FIXED PUBLISHED CONSTANT. For sigma = 1.0, a single Gaussian query (no subsampling), delta = 1e-5, the
// RDP accountant over a dense order grid yields the analytic optimum epsilon = 5.2985261 (optimal order
// near alpha = 5.80). This is the value the Google RdpAccountant reports for the same inputs; it is reproduced
// here ONLY as a quoted constant, not by calling the library. The accountant's shipped coarse grid lands at
// 5.3025851 (optimal order 6 on the grid), which is a SOUND UPPER BOUND: the min over a coarser grid is never
// below the dense optimum, so the accountant can only overstate epsilon, never understate it (the safe
// direction for a privacy claim).
func TestPublishedGoogleRDPAccountant(t *testing.T) {
	const publishedDenseOptimum = 5.2985261 // sigma=1, single query, delta=1e-5, dense grid analytic optimum

	// Reproduce the dense grid optimum here (a denser order set than the shipped grid) and check it matches
	// the published constant. This proves the conversion formula itself is the published one.
	dense := make([]float64, 0, 200000)
	for alpha := 1.01; alpha <= 2000; alpha += 0.01 {
		dense = append(dense, alpha)
	}
	denseBest := math.Inf(1)
	for _, alpha := range dense {
		rdp := alpha / (2 * 1.0 * 1.0)
		eps := rdp + math.Log(1/1e-5)/(alpha-1)
		if eps < denseBest {
			denseBest = eps
		}
	}
	if !almostEqual(denseBest, publishedDenseOptimum, 1e-4) {
		t.Fatalf("dense grid conversion = %v, published Google RDP accountant optimum = %v", denseBest, publishedDenseOptimum)
	}

	// The shipped coarse grid must be a SOUND UPPER BOUND of the published optimum (never below it), and close.
	acct, err := newRDPAccountant(100, 1e-5, nil)
	if err != nil {
		t.Fatalf("newRDPAccountant: %v", err)
	}
	coarse, err := acct.Charge(GaussianRelease{NoiseSigma: 1, ClipNorm: 1})
	if err != nil {
		t.Fatalf("Charge: %v", err)
	}
	if coarse < publishedDenseOptimum-1e-6 {
		t.Fatalf("coarse grid epsilon %v is BELOW the published optimum %v, unsound (understates privacy loss)", coarse, publishedDenseOptimum)
	}
	if coarse > publishedDenseOptimum+0.05 {
		t.Fatalf("coarse grid epsilon %v is too far above the published optimum %v, grid too sparse", coarse, publishedDenseOptimum)
	}
}

// TestRemainingHeadroom checks Remaining returns the budget headroom and never goes negative.
func TestRemainingHeadroom(t *testing.T) {
	acct, err := newRDPAccountant(10, 1e-6, nil)
	if err != nil {
		t.Fatalf("newRDPAccountant: %v", err)
	}
	if rem := acct.Remaining(); !almostEqual(rem, 10, 1e-9) {
		t.Fatalf("fresh Remaining = %v, want 10", rem)
	}
	if _, err := acct.Charge(GaussianRelease{NoiseSigma: 2, ClipNorm: 1}); err != nil {
		t.Fatalf("Charge: %v", err)
	}
	if rem := acct.Remaining(); rem < 0 || rem >= 10 {
		t.Fatalf("Remaining after one charge = %v, want in [0, 10)", rem)
	}
}

// TestBudgetRefusal asserts the pure accountant refuses a release that would exceed the epsilon budget,
// returns ErrBudgetExceeded, and leaves the accumulation untouched (a refused release writes nothing).
func TestBudgetRefusal(t *testing.T) {
	// A tiny budget so a single noisy release blows it: sigma small (large RDP), tight delta, small epsilon.
	acct, err := newRDPAccountant(0.5, 1e-7, nil)
	if err != nil {
		t.Fatalf("newRDPAccountant: %v", err)
	}
	epsBefore := acct.currentEpsilon()
	_, err = acct.Charge(GaussianRelease{NoiseSigma: 0.1, ClipNorm: 1}) // multiplier 0.1, huge RDP
	if !errors.Is(err, ErrBudgetExceeded) {
		t.Fatalf("expected ErrBudgetExceeded, got %v", err)
	}
	if epsAfter := acct.currentEpsilon(); !almostEqual(epsAfter, epsBefore, 1e-12) {
		t.Fatalf("a refused release mutated the accumulation: before %v after %v", epsBefore, epsAfter)
	}
}

// TestNonFiniteReleaseRejected asserts a release with a non positive clip norm is rejected rather than
// silently producing a non finite RDP.
func TestNonFiniteReleaseRejected(t *testing.T) {
	acct, err := newRDPAccountant(10, 1e-6, nil)
	if err != nil {
		t.Fatalf("newRDPAccountant: %v", err)
	}
	if _, err := acct.Charge(GaussianRelease{NoiseSigma: 2, ClipNorm: 0}); err == nil {
		t.Fatalf("expected an error for a non positive clip norm, got nil")
	}
}

// TestNoiseSamplerCannotRunFlagOff is the activation blocker test. With MM_FEDERATED_DP_ENABLED unset (the
// default false config), the active aggregator cannot even be constructed, so there is no object on which to
// call Noise: the only reachable path returns ErrDisabled. This proves the naive float sampler cannot run on
// a flag off path.
func TestNoiseSamplerCannotRunFlagOff(t *testing.T) {
	cfg := config.Federated{DPEnabled: false, DPEpsilon: 10, DPDelta: 1e-7}
	_, err := NewActiveAggregator(cfg, nil)
	if !errors.Is(err, ErrDisabled) {
		t.Fatalf("flag off: NewActiveAggregator must refuse with ErrDisabled, got %v", err)
	}
}

// TestNoiseSamplerIsBlockedEvenWhenFlagOn asserts that even when the flag is ON, the noise sampler is the
// activation blocker: it returns ErrNotImplemented and samples nothing. The naive float Gaussian is absent
// from the whole package, so an activated path still cannot add noise until the Mironov safe sampler is
// written.
func TestNoiseSamplerIsBlockedEvenWhenFlagOn(t *testing.T) {
	cfg := config.Federated{DPEnabled: true, DPEpsilon: 10, DPDelta: 1e-7}
	agg, err := NewActiveAggregator(cfg, nil)
	if err != nil {
		t.Fatalf("flag on: NewActiveAggregator should build, got %v", err)
	}
	_, err = agg.Noise([]float64{1, 2, 3}, 1.0)
	if !errors.Is(err, ErrNotImplemented) {
		t.Fatalf("Noise must be the activation blocker (ErrNotImplemented), got %v", err)
	}
	// Clip is the deterministic safe step and does run: a vector outside the ball is scaled to the clip norm.
	clipped := agg.Clip([]float64{3, 4}, 1.0) // norm 5, clip to 1 => (0.6, 0.8)
	cn := math.Hypot(clipped[0], clipped[1])
	if !almostEqual(cn, 1.0, 1e-12) {
		t.Fatalf("Clip did not bound the vector to the clip norm: got norm %v", cn)
	}
}

// staticReader satisfies Accountant so the interface stays exercised; it confirms rdpAccountant is a valid
// Accountant implementation.
var _ Accountant = (*rdpAccountant)(nil)

// TestSpendGateFlagOffRefuses asserts the ledger backed gate refuses every Charge with ErrDisabled while the
// flag is off, writing nothing. The store is nil because the gate must short circuit before touching it.
func TestSpendGateFlagOffRefuses(t *testing.T) {
	cfg := config.Federated{DPEnabled: false, DPEpsilon: 10, DPDelta: 1e-7}
	gate := NewSpendGate(cfg, nil, nil)
	_, err := gate.Charge(context.Background(), Release{ModelVersion: "v1.4.0", Gaussian: GaussianRelease{NoiseSigma: 2, ClipNorm: 1}})
	if !errors.Is(err, ErrDisabled) {
		t.Fatalf("flag off gate: Charge must refuse with ErrDisabled before touching the store, got %v", err)
	}
}