import json, time
import numpy as np
from implicit_fft import implicit_dft, implicit_idft, implicit_convolve

rng = np.random.default_rng(2297)
rows = []
# Prediction 1: identity holds for every admissible tile size.
for L, M in [(7, 16), (31, 64), (100, 256), (257, 1024)]:
    x = rng.normal(size=L) + 1j * rng.normal(size=L)
    for m in [1, 2, 4, 8, 16, 32]:
        if M % m:
            continue
        got = implicit_dft(x, M, m)
        padded = np.pad(x, (0, M-L))
        ref = np.conj(np.fft.fft(np.conj(padded)))
        rows.append({'kind':'dft', 'L':L, 'M':M, 'm':m,
                     'relerr':float(np.linalg.norm(got-ref)/np.linalg.norm(ref))})
        back = implicit_idft(got, L, m)
        rows.append({'kind':'inverse', 'L':L, 'M':M, 'm':m,
                     'relerr':float(np.linalg.norm(back-x)/np.linalg.norm(x))})

# Prediction 2: convolution is exact independently of the split.
conv_rows = []
for L, K in [(13, 5), (31, 17), (64, 33), (127, 65)]:
    x, g = rng.normal(size=L), rng.normal(size=K)
    ref = np.convolve(x, g)
    M = 1
    while M < L + K - 1:
        M *= 2
    for m in [1, 2, 4, 8, 16, 32, 64]:
        if M % m:
            continue
        got = implicit_convolve(x, g, m)
        conv_rows.append({'L':L, 'K':K, 'M':M, 'm':m,
                          'relerr':float(np.linalg.norm(got-ref)/np.linalg.norm(ref))})

# Prediction 3: explicit padded storage / tile storage is M/(ceil(L/m)m).
storage = []
for L, M in [(4096, 8192), (4096, 16384), (16384, 32768), (16384, 65536)]:
    for m in [16, 32, 64, 128, 256, 512]:
        if M % m:
            continue
        tile_len = ((L + m - 1) // m) * m
        storage.append({'L':L, 'M':M, 'm':m, 'explicit_elems':M,
                        'implicit_tile_elems':tile_len,
                        'ratio':M/tile_len})

# Small timing comparison; this NumPy implementation is a correctness MVP,
# not a fused production kernel, so timing is reported honestly.
def explicit_dft(x, M):
    return np.conj(np.fft.fft(np.conj(np.pad(x, (0, M-len(x))))))

def median_time(fn, reps=3):
    vals = []
    for _ in range(reps):
        t = time.perf_counter(); fn(); vals.append(time.perf_counter()-t)
    return float(np.median(vals))

timing = []
for L, K, m in [(257, 65, 16), (1025, 257, 32), (2049, 513, 64)]:
    M = 1
    while M < L + K - 1:
        M *= 2
    x = rng.normal(size=L); g = rng.normal(size=K)
    b = median_time(lambda: explicit_dft(x, M) * explicit_dft(g, M))
    i = median_time(lambda: implicit_dft(x, M, m) * implicit_dft(g, M, m))
    timing.append({'L':L, 'K':K, 'M':M, 'm':m,
                   'explicit_sec':b, 'implicit_sec':i, 'speedup':b/i})

out = {'dft_inverse':rows, 'convolution':conv_rows,
       'storage':storage, 'timing':timing}
with open('results.json', 'w') as f:
    json.dump(out, f, indent=2)
print('max dft/inverse error:', max(r['relerr'] for r in rows))
print('max convolution error:', max(r['relerr'] for r in conv_rows))
print('storage ratios:', min(x['ratio'] for x in storage), max(x['ratio'] for x in storage))
print('timing:', timing)