#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Jie Ren
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# External input data retain the separate rights recorded in the input JSON.

"""An educational single-template GW150914 matched-filter calculation.

Keep this file beside gw150914-matched-filter-inputs.json. Install the accompanying
requirements. Run --self-test without downloading data; run normally to fetch
hash-verified public inputs, or use --offline --cache-dir /path/to/existing/data.
The default cache is under the operating system temporary directory. Results go
to standard output unless --output is supplied. No random numbers are used.

This is not a search or parameter-estimation pipeline. Read the accompanying
article and manifest before interpreting the peak. In particular, the supplied
NR waveform was selected and adjusted using knowledge of this event.
"""

import argparse
import hashlib
import json
import platform
from pathlib import Path
import sys
import tempfile
import urllib.request

import h5py
import numpy as np
import scipy
from scipy import signal


STEM = "gw150914-matched-filter"
HERE = Path(__file__).resolve().parent


def require(condition, message):
    """Unlike assert, this check is not disabled by python -O."""
    if not condition:
        raise ValueError(message)


def digest(path, algorithm="sha256"):
    h = hashlib.new(algorithm)
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def verify_input(path, record):
    require(digest(path) == record["sha256"], f"SHA-256 mismatch: {path}")
    if "release_md5" in record:
        require(digest(path, "md5") == record["release_md5"],
                f"Release MD5 mismatch: {path}")


def obtain_input(record, cache, offline):
    """Never silently replace a corrupt cached input."""
    path = cache / record["filename"]
    if path.exists():
        verify_input(path, record)
        return path
    require(not offline, f"Missing input in offline mode: {path}")
    cache.mkdir(parents=True, exist_ok=True)
    temporary = None
    try:
        with tempfile.NamedTemporaryFile(dir=cache, delete=False) as target:
            temporary = Path(target.name)
            with urllib.request.urlopen(record["url"], timeout=60) as response:
                while True:
                    chunk = response.read(1024 * 1024)
                    if not chunk:
                        break
                    target.write(chunk)
        verify_input(temporary, record)
        temporary.replace(path)
        return path
    finally:
        if temporary is not None and temporary.exists():
            temporary.unlink()


def read_strain(path, record, settings):
    fs = settings["sample_rate_hz"]
    n = fs * settings["duration_seconds"]
    with h5py.File(path, "r") as handle:
        dataset = handle["strain/Strain"]
        data = dataset[:]
        require(handle["meta/Detector"][()].decode() == record["role"],
                "Detector metadata mismatch")
        require(int(handle["meta/GPSstart"][()]) == settings["gps_start"],
                "GPS start mismatch")
        require(int(handle["meta/Duration"][()]) == settings["duration_seconds"],
                "Duration mismatch")
        require(int(dataset.attrs["Npoints"]) == n and data.shape == (n,),
                "Sample count mismatch")
        require(float(dataset.attrs["Xspacing"]) == 1 / fs and
                int(dataset.attrs["Xstart"]) == settings["gps_start"],
                "Strain time grid mismatch")
        quality = handle["quality/simple/DQmask"][:]
        injections = handle["quality/injections/Injmask"][:]
    require(np.isfinite(data).all(), "Nonfinite strain")
    require(quality.shape == injections.shape == (settings["duration_seconds"],),
            "Quality flag cadence mismatch")
    require(np.all(quality == record["expected_dqmask"]), "Unexpected quality flags")
    require(np.all(injections == record["expected_injmask"]), "Unexpected injection flags")
    # Bits 0, 1, 2 mean NO CBC, burst, and detector-characterization injections.
    # L1 mask 23 is valid here: bit 3 is clear because a CW injection was present.
    require(np.all((injections & 7) == 7), "Transient hardware injection present")
    return data, {"samples": n, "dqmask": int(quality[0]),
                  "injmask": int(injections[0]), "no_transient_hardware_injections": True,
                  "continuous_wave_hardware_injection_present": bool((injections[0] & 8) == 0)}


def read_template(path, record, fs, n):
    table = np.loadtxt(path)
    require(table.shape == (record["expected_samples"], 2), "Template shape mismatch")
    time, waveform = table.T
    require(np.isfinite(table).all(), "Nonfinite template")
    require(np.allclose(np.diff(time), 1 / fs, rtol=0, atol=1e-14),
            "Template sampling does not match strain sampling")
    waveform = waveform.copy()
    first, last = round(0.1 * fs), round(0.005 * fs)
    waveform[:first] *= np.sin(np.linspace(0, np.pi / 2, first)) ** 2
    waveform[-last:] *= np.cos(np.linspace(0, np.pi / 2, last)) ** 2
    padded = np.zeros(n)
    padded[:len(waveform)] = waveform
    return padded, float(time[0])


def offsource_psd(data, fs, seconds, intervals):
    """Mean one-sided periodograms; no segment crosses the excluded event."""
    length = seconds * fs
    window = signal.windows.hann(length, sym=False)
    periodograms = []
    for low, high in intervals:
        for start in range(low * fs, high * fs - length + 1, length // 2):
            chunk = data[start:start + length]
            transform = np.fft.rfft((chunk - chunk.mean()) * window)
            power = 2 / fs * abs(transform) ** 2 / np.sum(window ** 2)
            power[[0, -1]] /= 2  # DC and Nyquist have no distinct negative partners.
            periodograms.append(power)
    require(len(periodograms) > 1, "Too few offsource PSD segments")
    psd = np.mean(periodograms, axis=0)
    require(np.all(np.isfinite(psd) & (psd > 0)), "Nonpositive/nonfinite PSD")
    return np.fft.rfftfreq(length, 1 / fs), psd, len(periodograms)


def matched_filter(data, template, dt, psd, band):
    """Return complex SNR and template norm, keeping only interior positive bins.

    NumPy uses exp(-2*pi*i*f*t). Its dt*rfft is the complex conjugate of the
    site's positive-exponent transform for real data. Therefore D*conj(H) below
    equals conj(D_site)*H_site. The positive-exponent ifft gives a positive lag
    for a delayed template. Do NOT mirror these bins into negative frequencies:
    their missing mirror is what retains both orthogonal quadratures.
    """
    n = len(data)
    require(len(template) == n and n % 2 == 0, "Need equal, even input lengths")
    frequencies = np.fft.rfftfreq(n, dt)
    keep = (frequencies >= band[0]) & (frequencies <= band[1])
    keep[[0, -1]] = False
    require(np.any(keep), "Empty frequency band")
    require(np.all(np.isfinite(psd[keep]) & (psd[keep] > 0)), "Invalid in-band PSD")
    dft = dt * np.fft.rfft(data)
    hft = dt * np.fft.rfft(template)
    delta_f = 1 / (n * dt)
    norm = np.sqrt(4 * delta_f * np.sum(abs(hft[keep]) ** 2 / psd[keep]))
    require(norm > 0, "Zero template norm")
    positive = np.zeros(len(frequencies), complex)
    positive[keep] = dft[keep] * hft[keep].conj() / psd[keep]
    # ifft divides by n; restore it to obtain the Riemann sum 4*delta_f*sum.
    snr = 4 * delta_f * n * np.fft.ifft(positive, n=n) / norm
    return snr, float(norm)


def self_checks(tolerance=1e-10):
    """Analytic line norm, explicit DFT, time-domain sums, and signed delays.

    The direct DFT has the site's POSITIVE sign and no FFT-library transform.
    This deliberately challenges the sign/normalization in matched_filter.
    """
    n, dt = 256, 1 / 128
    sample = np.arange(n)
    # Mixed even/odd bins avoid the half-period sign degeneracy of all-odd lines.
    lines, amplitudes = np.array([3, 10, 23]), np.array([1.0, 0.4, 0.2])
    q = sum(a * np.cos(2 * np.pi * k * sample / n + 0.17 * k)
            for k, a in zip(lines, amplitudes))
    frequency = np.arange(n // 2 + 1) / (n * dt)
    psd = 2 * (1 + (frequency / 20) ** 2)
    band = [frequency[1], frequency[-2]]
    analytic_norm = np.sqrt(n * dt * np.sum(amplitudes ** 2 / psd[lines]))
    _, fft_norm = matched_filter(q, q, dt, psd, band)
    errors = {"analytic_line_norm_relative": abs(fft_norm / analytic_norm - 1)}
    # Full direct DFT and Parseval, including both frequency signs.
    direct_matrix = np.exp(2j * np.pi * np.outer(sample, sample) / n)
    qplus = dt * (direct_matrix @ q)
    time_energy = dt * np.sum(q ** 2)
    fourier_energy = np.sum(abs(qplus) ** 2) / (n * dt)
    errors["direct_dft_parseval_relative"] = abs(fourier_energy / time_energy - 1)
    direct_max = 0.0
    norm_max = 0.0
    recovered = []
    keep = np.arange(1, n // 2)
    for shift, amplitude in [(37, 0.7), (-29, -1.3)]:
        # Pure array indexing builds the injection, independent of Fourier phases.
        data = amplitude * q[(sample - shift) % n]
        z, norm = matched_filter(data, q, dt, psd, band)
        dplus = dt * (direct_matrix @ data)
        direct = 4 / (n * dt) / analytic_norm * np.array([
            np.sum(dplus[keep].conj() * qplus[keep] / psd[keep]
                   * np.exp(2j * np.pi * keep * lag / n)) for lag in sample])
        direct_max = max(direct_max, float(np.max(abs(z - direct)) / (abs(amplitude) * norm)))
        norm_max = max(norm_max, abs(z[shift % n].real / (amplitude * analytic_norm) - 1))
        peak = int(np.argmax(abs(z)))
        signed_peak = peak if peak < n // 2 else peak - n
        require(signed_peak == shift, "Signed delay recovery failed")
        require(abs(z[peak].imag) / norm < tolerance, "Injection quadrature is nonzero")
        recovered.append({"injected_samples": shift, "recovered_samples": signed_peak,
                          "injected_amplitude": amplitude})
    errors["all_lags_direct_dft_relative"] = direct_max
    errors["signed_injection_amplitude_relative"] = norm_max
    # With flat one-sided PSD=2, (d|q)=dt*sum(d*q) after excluding zero endpoints.
    z, norm = matched_filter(q, q, dt, np.full(n // 2 + 1, 2.0), band)
    time_correlation = dt * np.array([np.sum(q * q[(sample - lag) % n]) for lag in sample])
    errors["all_lags_time_sum_relative"] = float(np.max(abs(z.real * norm - time_correlation)) / time_energy)
    # PSD density units and window normalization checked against scipy.signal.welch.
    noise_like = np.sin(0.31 * sample ** 2) + np.cos(0.71 * sample)
    f, manual, count = offsource_psd(noise_like, 32, 2, [[0, 8]])
    f2, library = signal.welch(noise_like, fs=32, window="hann", nperseg=64,
                               noverlap=32, detrend="constant", scaling="density", average="mean")
    require(np.array_equal(f, f2) and count == 7, "Welch comparison grid mismatch")
    errors["welch_library_relative"] = float(np.max(abs(manual - library)) / np.max(library))
    require(all(value < tolerance for value in errors.values()), f"Independent check failed: {errors}")
    return {"status": "passed", "tolerance": tolerance, "errors": errors,
            "signed_delay_checks": recovered, "random_seed": None}


def analyze(manifest, paths):
    settings = manifest["analysis"]
    fs = settings["sample_rate_hz"]
    n = fs * settings["duration_seconds"]
    records = {item["role"]: item for item in manifest["inputs"]}
    template, first_time = read_template(paths["template"], records["template"], fs, n)
    strain, metadata = {}, {}
    for detector in ["H1", "L1"]:
        strain[detector], metadata[detector] = read_strain(paths[detector], records[detector], settings)
    # The full data window is exactly one over the template's event-time support.
    window = signal.windows.tukey(n, alpha=0.125)
    conditioned = {key: (value - value.mean()) * window for key, value in strain.items()}
    frequency = np.fft.rfftfreq(n, 1 / fs)
    # A shift of the padded array moves its FIRST template sample. Subtract the
    # file's first-column origin to express the release waveform's reference time.
    peak_times = np.arange(n) / fs - first_time
    search_low, search_high = settings["peak_search_seconds_from_start"]
    selected = np.flatnonzero((peak_times >= search_low) & (peak_times <= search_high))
    reference_from_start = settings["gps_reference"] - settings["gps_start"]
    runs = []
    for configuration in settings["configurations"]:
        result = {"configuration": configuration, "detectors": {}}
        for detector in ["H1", "L1"]:
            coarse_f, coarse_psd, count = offsource_psd(
                strain[detector], fs, configuration["psd_segment_seconds"],
                settings["offsource_intervals_seconds_from_start"])
            psd = np.interp(frequency, coarse_f, coarse_psd)
            z, norm = matched_filter(conditioned[detector], template, 1 / fs,
                                     psd, configuration["band_hz"])
            peak = int(selected[np.argmax(abs(z[selected]))])
            require(peak not in (selected[0], selected[-1]), "Peak touches search boundary")
            result["detectors"][detector] = {
                "peak_snr": float(abs(z[peak])), "template_norm": norm,
                "reference_time_offset_seconds": float(peak_times[peak] - reference_from_start),
                "template_start_lag_samples": peak, "psd_periodograms": count}
        result["hanford_minus_livingston_seconds"] = (
            result["detectors"]["H1"]["reference_time_offset_seconds"]
            - result["detectors"]["L1"]["reference_time_offset_seconds"])
        runs.append(result)
    return {"data_checks": metadata, "sample_spacing_seconds": 1 / fs,
            "gps_reference": settings["gps_reference"], "template_first_time_seconds": first_time,
            "runs": runs}


def compare_results(actual, reference, tolerances):
    """Check numerical reproduction; these are NOT astrophysical uncertainties."""
    require(actual["provenance"] == reference["provenance"], "Provenance differs from reference")
    require(actual["data_checks"] == reference["data_checks"], "Decoded data checks differ")
    require(len(actual["runs"]) == len(reference["runs"]), "Configuration count differs")
    for run, expected in zip(actual["runs"], reference["runs"]):
        require(run["configuration"] == expected["configuration"], "Configurations differ")
        for detector in ["H1", "L1"]:
            got, want = run["detectors"][detector], expected["detectors"][detector]
            for key in ["peak_snr", "template_norm"]:
                require(abs(got[key] - want[key]) <= tolerances["rerun_snr_absolute"],
                        f"Numerical rerun differs: {detector} {key}")
            require(abs(got["reference_time_offset_seconds"] - want["reference_time_offset_seconds"])
                    <= tolerances["rerun_time_seconds_absolute"], "Peak time differs")
            require(got["psd_periodograms"] == want["psd_periodograms"] and
                    got["template_start_lag_samples"] == want["template_start_lag_samples"],
                    "Discrete analysis output differs")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--cache-dir", type=Path,
                        default=Path(tempfile.gettempdir()) / "blackhole-gw150914-data")
    parser.add_argument("--offline", action="store_true", help="Use existing verified files only")
    parser.add_argument("--self-test", action="store_true", help="Independent checks only; no downloads")
    parser.add_argument("--output", type=Path, help="Write JSON here instead of standard output")
    parser.add_argument("--compare", type=Path, help="Compare against the published results JSON")
    args = parser.parse_args()
    manifest_path = HERE / f"{STEM}-inputs.json"
    manifest = json.loads(manifest_path.read_text())
    checks = self_checks(manifest["numerical_tolerances"]["independent_check_relative"])
    result = {"schema_version": 1, "purpose": manifest["purpose"], "self_checks": checks,
              "environment": {"platform": platform.platform(), "python": platform.python_version(), "numpy": np.__version__,
                              "scipy": scipy.__version__, "h5py": h5py.__version__}}
    if not args.self_test:
        paths = {record["role"]: obtain_input(record, args.cache_dir, args.offline)
                 for record in manifest["inputs"]}
        result["provenance"] = {"script_sha256": digest(Path(__file__)),
                                "manifest_sha256": digest(manifest_path),
                                "requirements_sha256": digest(HERE / f"{STEM}-requirements.txt"),
                                "input_sha256": {record["role"]: record["sha256"] for record in manifest["inputs"]}}
        result.update(analyze(manifest, paths))
        if args.compare:
            compare_results(result, json.loads(args.compare.read_text()), manifest["numerical_tolerances"])
            result["reference_comparison"] = "passed numerical reproduction tolerances"
    else:
        require(args.compare is None, "--compare needs a data run, not --self-test")
    output = json.dumps(result, indent=2, sort_keys=True, allow_nan=False) + "\n"
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(output)
    else:
        sys.stdout.write(output)


if __name__ == "__main__":
    main()
