"""Independent NumPy/SciPy implementation of the documented AFTERLIGHT DSP recipe.

Usage: python scripts/verify_numerics.py --report /path/to/report.json
This checks numerical reproducibility, not a wallet signature, detection significance,
or the scientific interpretation of the chosen processing recipe.
"""
from pathlib import Path
import argparse
import hashlib
import json
import platform
import sys
import numpy as np
import scipy
from scipy import signal, fft

FS = 4096
START = 1126259447
EVENT = 1126259462.4
REPORT_TOLERANCE = 1e-5
SUITE_TOLERANCE = 1e-8


def sha(data):
    return hashlib.sha256(data).hexdigest()


def load_inputs(data_dir):
    manifest = json.loads((data_dir / 'manifest.json').read_text())
    arrays = {}
    for f in manifest['rawBinaryFiles']:
        contents = (data_dir / f['file']).read_bytes()
        if sha(contents) != f['sha256'] or len(contents) != f['bytes']:
            raise ValueError(f"Input integrity check failed: {f['file']}")
        x = np.frombuffer(contents, dtype='<f8')
        if len(x) != FS * 32 or not np.isfinite(x).all():
            raise ValueError('Unexpected or nonfinite strain data')
        arrays[f['detector']] = x
    return manifest, arrays


def index(seconds):
    return int(np.floor((EVENT + seconds - START) * FS + .5))


def valid_config(c):
    return (isinstance(c.get('low'), (int, float)) and isinstance(c.get('high'), (int, float))
            and np.isfinite(c['low']) and np.isfinite(c['high']) and c['low'] >= 20
            and c['high'] <= 800 and c['high'] - c['low'] >= 20
            and isinstance(c.get('whiten'), bool) and isinstance(c.get('notches'), bool))


def process(x, c):
    # SciPy periodogram performs windowing, detrending, one-sided doubling and
    # spectral-density normalization, independently of the JavaScript FFT.
    starts = np.arange(0, len(x) - FS + 1, FS // 2)
    absolute = START + starts / FS
    starts = starts[~((absolute < EVENT + 1) & (absolute + 1 > EVENT - 1))]
    segments = np.stack([x[i:i + FS] for i in starts])
    frequencies, periodograms = signal.periodogram(
        segments, fs=FS, window=signal.windows.hann(FS, sym=True),
        detrend='constant', return_onesided=True, scaling='density', axis=-1)
    psd = periodograms.mean(axis=0)
    bins = fft.rfftfreq(len(x), 1 / FS)
    low, high = c['low'], c['high']
    gain = np.where((bins >= low) & (bins <= high), 1.0, 0.0)
    rising = (bins >= low - 5) & (bins < low)
    falling = (bins > high) & (bins <= high + 5)
    gain[rising] = (1 - np.cos(np.pi * (bins[rising] - low + 5) / 5)) / 2
    gain[falling] = (1 + np.cos(np.pi * (bins[falling] - high) / 5)) / 2
    if c['notches']:
        for line in (60, 120, 180):
            distance = np.abs(bins - line)
            gain[distance < 1] = 0
            taper = (distance >= 1) & (distance < 2)
            gain[taper] *= (1 - np.cos(np.pi * (distance[taper] - 1))) / 2
    if c['whiten']:
        gain /= np.sqrt(np.maximum(np.interp(bins, frequencies, psd), 1e-60) * FS / 2)
    tapered = signal.detrend(x, type='constant') * signal.windows.tukey(len(x), .25, sym=True)
    output = fft.irfft(fft.rfft(tapered) * gain, n=len(x))
    output /= np.sqrt(np.mean(output[index(-12):index(-2)] ** 2))
    return output, len(starts)


def calculate(arrays, config):
    if not valid_config(config):
        raise ValueError('Unsupported frequency configuration')
    h, count = process(arrays['H1'], config)
    l, _ = process(arrays['L1'], config)
    a, b = index(-.15), index(.05)
    offsets = np.arange(-int(.01 * FS), int(.01 * FS) + 1)
    hv = h[a:b]
    lv = np.stack([l[a + shift:b + shift] for shift in offsets])
    correlations = lv @ hv / np.sqrt(np.sum(lv * lv, axis=1) * np.sum(hv * hv))
    best = int(np.argmax(np.abs(correlations)))
    start, end = index(-.6), index(.3)
    return {
        'h1': h[start:end], 'l1': l[start:end], 'correlations': correlations,
        'lagSamples': int(offsets[best]), 'lagMilliseconds': float(offsets[best] / FS * 1000),
        'signedCorrelation': float(correlations[best]),
        'h1WindowRms': float(np.sqrt(np.mean(hv * hv))),
        'l1WindowRms': float(np.sqrt(np.mean(l[a:b] ** 2))), 'noiseSegments': count,
        'timeStart': START + start / FS - EVENT, 'dt': 1 / FS,
    }


def array_error(actual, expected):
    expected = np.asarray(expected, dtype=float)
    if actual.shape != expected.shape or not np.isfinite(expected).all():
        raise ValueError('Invalid waveform dimensions or values')
    return float(np.max(np.abs(actual - expected)))


def compare(reference, metrics, h1, l1, tolerance, metric_tolerance=None):
    metric_tolerance = tolerance if metric_tolerance is None else metric_tolerance
    errors = {name: abs(reference[name] - metrics[name]) for name in
              ('lagMilliseconds', 'signedCorrelation', 'h1WindowRms', 'l1WindowRms')}
    errors.update(h1=array_error(reference['h1'], h1), l1=array_error(reference['l1'], l1))
    lag_agrees = reference['lagSamples'] == metrics['lagSamples']
    noise_agrees = reference['noiseSegments'] == metrics['noiseSegments']
    return {'passed': bool(lag_agrees and noise_agrees and max(errors['h1'], errors['l1']) <= tolerance and max(errors[k] for k in ('lagMilliseconds', 'signedCorrelation', 'h1WindowRms', 'l1WindowRms')) <= metric_tolerance),
            'lagSamplesAgree': bool(lag_agrees), 'noiseSegmentsAgree': bool(noise_agrees),
            'maxAbsoluteErrors': errors, 'waveformTolerance': tolerance, 'metricTolerance': metric_tolerance}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    choice = parser.add_mutually_exclusive_group(required=True)
    choice.add_argument('--report', type=Path)
    choice.add_argument('--suite', type=Path)
    parser.add_argument('--data', type=Path)
    parser.add_argument('--output', type=Path)
    args = parser.parse_args()
    base = Path(__file__).resolve().parent.parent
    data_dir = args.data or (base / 'dist/data' if (base / 'dist/data').is_dir() else base / 'data')
    manifest, arrays = load_inputs(data_dir)
    if args.report:
        report = json.loads(args.report.read_text())
        p = report['payload']
        engine = json.loads((data_dir / 'engine.json').read_text())
        expected_event = {'name': 'GW150914', 'eventGPS': EVENT, 'startGPS': START,
                          'sampleRateHz': FS, 'durationSeconds': 32}
        if (p.get('schema') != 'afterlight.report.v1' or p.get('event') != expected_event
                or p.get('sources') != manifest['rawBinaryFiles'] or p.get('engine') != engine):
            raise ValueError('Unsupported source, event, or engine identity')
        w = p['results']['waveform']
        reference = calculate(arrays, p['configuration'])
        if w['dt'] != reference['dt'] or w['timeStart'] != reference['timeStart']:
            raise ValueError('The waveform time axis differs from this implementation')
        result = compare(reference, p['results'], w['h1'], w['l1'], REPORT_TOLERANCE)
        result.update(scope='Numerical reproduction only. Use the website or npm run replay for content and signature verification.',
                      configuration=p['configuration'])
    else:
        cases = json.loads(args.suite.read_text())
        results = []
        for case in cases:
            ref = calculate(arrays, case['configuration'])
            check = compare(ref, case['metrics'], case['output']['h1'], case['output']['l1'], SUITE_TOLERANCE, 1e-6)
            check['configuration'] = case['configuration']
            check['correlationCurveMaxAbsoluteError'] = array_error(ref['correlations'], case['output']['correlations'])
            check['passed'] = bool(check['passed'] and check['correlationCurveMaxAbsoluteError'] <= SUITE_TOLERANCE)
            check['waveformSamplesPerDetector'] = len(ref['h1'])
            check['referenceLagMilliseconds'] = ref['lagMilliseconds']
            results.append(check)
        result = {'passed': all(r['passed'] for r in results), 'cases': results,
                  'runtime': {'python': platform.python_version(), 'numpy': np.__version__, 'scipy': scipy.__version__}}
    if args.output:
        args.output.write_text(json.dumps(result, indent=2, allow_nan=False) + '\n')
    else:
        print(json.dumps(result, indent=2, allow_nan=False))
    return 0 if result['passed'] else 1


if __name__ == '__main__':
    try:
        sys.exit(main())
    except (ValueError, KeyError, TypeError, OSError) as exc:
        print(f'Verification failed: {exc}', file=sys.stderr)
        sys.exit(1)