#!/usr/bin/env python3
"""Read paired observed and synthetic three-component SAC waveforms.

Observed filename convention:
    NET.STA.LOC.CHA.<quantity>
Example:
    BK.BARR.00.HNE.disp

Synthetic filename convention:
    NET.STA.LOC.CHA.<quantity>.syn
Example:
    BK.BARR.00.HNE.disp.syn

The final channel-code character identifies the component:
    E = east-west, N = north-south, Z = vertical

Station azimuth and distance are read once per NET.STA.LOC group from the
observed Z-component SAC header. Stations are sorted by distance by default.

Timing note:
    Synthetic waveforms were already shifted by ttshift in the inversion code.
    Consequently, ttshift is extracted and logged but is not applied again to
    observed plotting times by default. Use --apply-ttshift only for a deliberate
    diagnostic comparison.
"""

from __future__ import annotations

import argparse
import csv
import math
import sys

import matplotlib.pyplot as plt
import numpy as np
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Tuple

try:
    from obspy import Trace, UTCDateTime, read
    from obspy.signal.cross_correlation import correlate, xcorr_max
except ImportError as exc:
    raise SystemExit(
        "ObsPy is required. Install it in the active Python environment "
        "before running this script."
    ) from exc


COMPONENTS = ("E", "N", "Z")
GLOBAL_AMPLITUDE_PLOT_PATCH_V1 = True
TRACE_NORMALIZATION_PATCH_V2 = True
AZIMUTH_TIME_ANNOTATION_PATCH_V3 = True
MULTIPAGE_PUBLICATION_PATCH_V4 = True
SELECTION_FILTER_PATCH_V5 = True
LABEL_TEXT_PATCH_V6 = True
OBSERVED_TRACE_MAX_ANNOTATION_PATCH_V7 = True
TRACE_MAX_LAYOUT_FONT_PATCH_V8 = True
TRACE_ROW_LABEL_LAYOUT_PATCH_V9 = True
NETWORK_STATION_DIST_BASELINE_PATCH_V10 = True
VARIANCE_REDUCTION_LABEL_PATCH_V11 = True
VR_LAYOUT_FONT_PATCH_V12 = True
TTSHIFT_EXTRACTION_PATCH_V13 = True
APPLY_TTSHIFT_TO_OBSERVED_PATCH_V14A = True
TTSHIFT_ALREADY_APPLIED_TO_SYNTHETIC_PATCH_V15 = True
RESIDUAL_CC_TTSHIFT_UPDATE_PATCH_V16A = True
SAC_UNDEFINED = -12345.0
SAC_UNDEFINED_TOLERANCE = 0.01


@dataclass(frozen=True, order=True)
class WaveformKey:
    """Unique identifier for one network/station/location data set."""

    network: str
    station: str
    location: str
    quantity: str

    @property
    def seed_prefix(self) -> str:
        return f"{self.network}.{self.station}.{self.location}"


@dataclass
class WaveformPair:
    """Observed and synthetic traces for one component."""

    observed: Trace
    synthetic: Trace
    observed_path: Path
    synthetic_path: Path


@dataclass
class StationWaveforms:
    """Three-component waveform pairs and station-level SAC geometry."""

    key: WaveformKey
    components: Dict[str, WaveformPair] = field(default_factory=dict)
    azimuth_deg: float | None = None
    distance_km: float | None = None
    metadata_source_path: Path | None = None


WaveformData = Dict[WaveformKey, StationWaveforms]


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Read paired observed and synthetic E/N/Z SAC waveforms, read "
            "station geometry from the observed Z component, and print a "
            "distance-sorted station summary."
        )
    )
    parser.add_argument(
        "--data-dir",
        type=Path,
        required=True,
        help="Directory containing the observed and synthetic SAC files.",
    )
    parser.add_argument(
        "--reg-inv-out",
        type=Path,
        required=True,
        help=(
            "Inversion output containing per-trace Variance reduction values. "
            "A relative path is resolved inside --data-dir."
        ),
    )
    parser.add_argument(
        "--vr-decimals",
        type=int,
        default=2,
        help="Decimal places for VR in trace labels (default: 2).",
    )
    parser.add_argument(
        "--ttshift-decimals",
        type=int,
        default=2,
        help="Decimal places for ttshift in the terminal log (default: 2).",
    )
    parser.add_argument(
        "--apply-ttshift",
        dest="apply_ttshift",
        action="store_true",
        help=(
            "Apply ttshift to observed plotting times for a diagnostic "
            "comparison. Synthetic waveforms were already shifted by "
            "the inversion code."
        ),
    )
    parser.add_argument(
        "--no-apply-ttshift",
        dest="apply_ttshift",
        action="store_false",
        help=(
            "Do not apply ttshift during plotting (default); synthetic "
            "waveforms were already shifted by the inversion code."
        ),
    )
    parser.set_defaults(apply_ttshift=False)
    parser.add_argument(
        "--ttshift-sign",
        type=float,
        choices=(-1.0, 1.0),
        default=1.0,
        help=(
            "Multiplier applied to stored ttshift before adding it to the "
            "observed time axis; choose -1 or 1 (default: 1)."
        ),
    )
    parser.add_argument(
        "--make-correlation-analysis",
        action=argparse.BooleanOptionalAction,
        default=True,
        help=(
            "Compute residual observed/synthetic cross correlation and proposed "
            "inversion ttshift updates (default: on)."
        ),
    )
    parser.add_argument(
        "--cc-max-lag-s", type=float, default=10.0,
        help="Maximum positive-correlation lag search in seconds (default: 10).",
    )
    parser.add_argument(
        "--cc-min-correlation", type=float, default=0.0,
        help="Minimum maximum-positive CC required to accept an update (default: 0).",
    )
    parser.add_argument(
        "--cc-min-vr-improvement", type=float, default=0.0,
        help=(
            "Minimum VR improvement in percentage points required to accept an "
            "update (default: 0)."
        ),
    )
    parser.add_argument(
        "--cc-output-csv", type=Path,
        default=Path("ttshift_correlation.csv"),
        help="Detailed per-trace correlation CSV output.",
    )
    parser.add_argument(
        "--ttshift-update-file", type=Path,
        default=Path("ttshift_update.txt"),
        help="Station-level inversion ttshift update file.",
    )
    parser.add_argument(
        "--ttshift-update-gps-line",
        default="29 gps_weightIn ngl_disp_select_up0.txt",
        help="Final GPS line appended to ttshift_update.txt; empty omits it.",
    )
    parser.add_argument(
        "--quantity",
        default="disp",
        help=(
            "Observed-data filename extension, for example disp, vel, or acc "
            "(default: disp)."
        ),
    )
    parser.add_argument(
        "--station-order",
        choices=("azimuth", "distance", "name"),
        default="azimuth",
        help=(
            "Station order from top to bottom: increasing SAC azimuth, "
            "increasing distance, or alphabetical NET.STA.LOC "
            "(default: azimuth)."
        ),
    )
    parser.add_argument(
        "--eq-origin-time",
        default="2024-12-05T18:44:16.460000",
        help=(
            "Earthquake origin time used as zero for waveform plotting "
            "(default: 2024-12-05T18:44:16.460000)."
        ),
    )
    parser.add_argument("--make-plot", action="store_true")
    parser.add_argument("--plot-file", type=Path,
                        default=Path("ff_3c_waveform_plot.pdf"))
    parser.add_argument(
        "--normalization",
        choices=("station", "trace", "global", "none"),
        default="station",
        help=(
            "Amplitude scaling mode (default: station). station uses one "
            "combined maximum per network.station.location; trace uses "
            "each observed and synthetic component trace maximum; global "
            "uses one maximum for the complete data set; none applies no "
            "normalization."
        ),
    )
    parser.add_argument("--input-unit-factor", type=float, default=1.0,
                        help="Multiply SAC samples before plotting; m to cm = 100")
    parser.add_argument(
        "--display-unit", default="cm",
        help=(
            "Amplitude unit used in diagnostics and Max. amp. annotations "
            "(default: cm). Apply numerical conversion with "
            "--input-unit-factor."
        ),
    )
    parser.add_argument("--station-spacing", type=float, default=1.0)
    parser.add_argument("--waveform-half-height", type=float, default=0.42)
    parser.add_argument("--line-width", type=float, default=0.8)
    parser.add_argument("--figure-width", type=float, default=12.0)
    parser.add_argument("--row-height", type=float, default=0.34)
    parser.add_argument(
        "--time-min", type=float, default=None,
        help="Minimum plotted time relative to origin in seconds.",
    )
    parser.add_argument(
        "--time-max", type=float, default=None,
        help="Maximum plotted time relative to origin in seconds.",
    )
    parser.add_argument(
        "--station-label-size", type=float, default=7.0,
        help="Station-code label size (default: 7).",
    )
    parser.add_argument(
        "--vr-font-size", type=float, default=7.0,
        help="Per-trace VR label font size (default: 7).",
    )
    parser.add_argument(
        "--az-font-size", type=float, default=6.0,
        help="AZ label font size (default: 6).",
    )
    parser.add_argument(
        "--dist-font-size", type=float, default=6.0,
        help="DIST label font size (default: 6).",
    )
    parser.add_argument(
        "--metadata-label-size", type=float, default=6.0,
        help="Azimuth/distance label size (default: 6).",
    )
    parser.add_argument(
        "--component-title-size", type=float, default=14.0,
        help="Component title size, for example East (default: 14).",
    )
    parser.add_argument(
        "--axis-label-size", type=float, default=10.0,
        help="Axis-label font size (default: 10).",
    )
    parser.add_argument(
        "--tick-label-size", type=float, default=9.0,
        help="Numeric time-axis tick-label size (default: 9).",
    )
    parser.add_argument(
        "--legend-font-size", type=float, default=9.0,
        help="Observed/synthetic legend font size (default: 9).",
    )
    parser.add_argument(
        "--footer-font-size", type=float, default=7.0,
        help="Footer font size (default: 7).",
    )
    parser.add_argument(
        "--max-amp-font-size", type=float, default=6.0,
        help="Per-trace Max. amp. font size (default: 6).",
    )
    parser.add_argument(
        "--max-amp-decimals", type=int, default=2,
        help="Decimal places in Max. amp. labels (default: 2).",
    )
    parser.add_argument(
        "--max-amp-x", type=float, default=0.01,
        help="Max. amp. x position in axes fraction (default: 0.01).",
    )
    parser.add_argument(
        "--max-amp-y-offset", type=float, default=0.18,
        help=(
            "Downward Max. amp. offset as a fraction of station spacing "
            "(default: 0.18)."
        ),
    )
    parser.add_argument(
        "--station-name-x", type=float, default=0.99,
        help="Station-name x position in East axes fraction (default: 0.99).",
    )
    parser.add_argument(
        "--station-name-y-offset", type=float, default=0.10,
        help="Upward station-name offset in row spacing (default: 0.10).",
    )
    parser.add_argument(
        "--vr-x", type=float, default=0.01,
        help="VR x position in axes fraction (default: 0.01).",
    )
    parser.add_argument(
        "--vr-y-offset", type=float, default=0.10,
        help="Upward VR offset in row spacing (default: 0.10).",
    )
    parser.add_argument(
        "--az-x", type=float, default=0.99,
        help="AZ x position in Vertical axes fraction (default: 0.99).",
    )
    parser.add_argument(
        "--az-y-offset", type=float, default=0.10,
        help="Upward AZ offset in row spacing (default: 0.10).",
    )
    parser.add_argument(
        "--dist-x", type=float, default=0.99,
        help="DIST x position in Vertical axes fraction (default: 0.99).",
    )
    parser.add_argument(
        "--dist-y-offset", type=float, default=0.06,
        help="Downward DIST offset in row spacing (default: 0.06).",
    )
    parser.add_argument(
        "--legend-loc", default="upper right",
        help="Matplotlib legend location (default: upper right).",
    )
    parser.add_argument(
        "--stations-per-page", type=int, default=18,
        help=(
            "Maximum stations on each figure page (default: 18). "
            "Use 0 to place all stations on one page."
        ),
    )
    parser.add_argument(
        "--y-axis-label", default="",
        help=(
            "Left-axis label; use an empty string to suppress it "
            "(default: no y-axis label)."
        ),
    )
    parser.add_argument(
        "--annotate-station-absmax", action="store_true",
        help="Add the station combined absolute maximum below AZ and DIS.",
    )
    parser.add_argument(
        "--footer", action=argparse.BooleanOptionalAction, default=True,
        help="Show reproducibility metadata in the page footer (default: on).",
    )
    parser.add_argument(
        "--include-stations", default="",
        help=(
            "Comma-separated station codes or NET.STA.LOC identifiers to include. "
            "Empty means include all stations before other filters."
        ),
    )
    parser.add_argument(
        "--exclude-stations", default="",
        help=(
            "Comma-separated station codes or NET.STA.LOC identifiers to exclude."
        ),
    )
    parser.add_argument(
        "--station-file", type=Path, default=None,
        help=(
            "Optional text file containing one station code or NET.STA.LOC per "
            "line. Blank lines and lines beginning with # are ignored."
        ),
    )
    parser.add_argument(
        "--min-distance-km", type=float, default=None,
        help="Minimum station distance in km (inclusive).",
    )
    parser.add_argument(
        "--max-distance-km", type=float, default=None,
        help="Maximum station distance in km (inclusive).",
    )
    parser.add_argument(
        "--min-azimuth-deg", type=float, default=None,
        help="Minimum station azimuth in degrees (inclusive).",
    )
    parser.add_argument(
        "--max-azimuth-deg", type=float, default=None,
        help="Maximum station azimuth in degrees (inclusive).",
    )
    parser.add_argument(
        "--components", default="E,N,Z",
        help=(
            "Comma-separated components to plot, chosen from E,N,Z "
            "(default: E,N,Z)."
        ),
    )
    return parser.parse_args()


def parse_observed_filename(path: Path, quantity: str) -> Tuple[WaveformKey, str, str]:
    """Parse an observed filename and return key, component, and channel."""
    parts = path.name.split(".")
    if len(parts) != 5:
        raise ValueError(
            f"Expected NET.STA.LOC.CHA.{quantity}, but found: {path.name}"
        )

    network, station, location, channel, file_quantity = parts
    if file_quantity != quantity:
        raise ValueError(
            f"Expected quantity '{quantity}', but found '{file_quantity}' "
            f"in {path.name}"
        )
    if not channel:
        raise ValueError(f"Empty channel code in {path.name}")

    component = channel[-1].upper()
    if component not in COMPONENTS:
        raise ValueError(
            f"Channel {channel!r} does not end in E, N, or Z: {path.name}"
        )

    key = WaveformKey(network, station, location, quantity)
    return key, component, channel


def read_one_sac(path: Path) -> Trace:
    """Read exactly one SAC trace from a file."""
    stream = read(str(path), format="SAC")
    if len(stream) != 1:
        raise ValueError(
            f"Expected one SAC trace in {path}, but ObsPy read {len(stream)} traces."
        )
    return stream[0]


def validate_trace_identity(
    trace: Trace,
    key: WaveformKey,
    expected_channel: str,
    path: Path,
) -> None:
    """Check populated SAC/SEED identifiers against the filename."""
    expected = {
        "network": key.network,
        "station": key.station,
        "location": key.location,
        "channel": expected_channel,
    }
    for field_name, expected_value in expected.items():
        actual_value = str(getattr(trace.stats, field_name, "") or "")
        # Some SAC files do not carry complete SEED identifiers. Only compare
        # fields that ObsPy populated.
        if actual_value and actual_value != expected_value:
            raise ValueError(
                f"{path}: filename {field_name}={expected_value!r}, but SAC "
                f"header has {field_name}={actual_value!r}."
            )


def read_required_sac_float(trace: Trace, header_name: str, path: Path) -> float:
    """Return one finite, defined numeric SAC header value."""
    sac = getattr(trace.stats, "sac", None)
    if sac is None:
        raise ValueError(f"{path}: SAC header block is missing.")

    value = sac.get(header_name, None)
    if value is None:
        raise ValueError(f"{path}: required SAC header '{header_name}' is missing.")

    try:
        number = float(value)
    except (TypeError, ValueError) as exc:
        raise ValueError(
            f"{path}: SAC header '{header_name}' is not numeric: {value!r}"
        ) from exc

    if not math.isfinite(number):
        raise ValueError(
            f"{path}: SAC header '{header_name}' is not finite: {number!r}"
        )
    if math.isclose(
        number,
        SAC_UNDEFINED,
        rel_tol=0.0,
        abs_tol=SAC_UNDEFINED_TOLERANCE,
    ):
        raise ValueError(
            f"{path}: SAC header '{header_name}' has the undefined SAC value "
            f"{number:g}."
        )
    return number


def read_optional_sac_float(trace: Trace, header_name: str) -> float | None:
    """Return an optional finite, defined SAC value, or None."""
    sac = getattr(trace.stats, "sac", None)
    if sac is None:
        return None
    value = sac.get(header_name, None)
    if value is None:
        return None
    try:
        number = float(value)
    except (TypeError, ValueError):
        return None
    if not math.isfinite(number):
        return None
    if math.isclose(number, SAC_UNDEFINED, rel_tol=0.0,
                    abs_tol=SAC_UNDEFINED_TOLERANCE):
        return None
    return number


def get_trace_timing(trace: Trace, path: Path) -> Dict[str, object]:
    """Collect sample timing and selected SAC reference headers."""
    npts = int(trace.stats.npts)
    delta = float(trace.stats.delta)
    if npts <= 0:
        raise ValueError(f"{path}: trace has no samples (npts={npts}).")
    if not math.isfinite(delta) or delta <= 0.0:
        raise ValueError(f"{path}: invalid sampling interval delta={delta!r}.")

    b = read_required_sac_float(trace, "b", path)
    e = read_required_sac_float(trace, "e", path)
    expected_e = b + (npts - 1) * delta
    e_tolerance = max(1.0e-5, abs(delta) * 1.0e-3)
    if not math.isclose(e, expected_e, rel_tol=0.0, abs_tol=e_tolerance):
        raise ValueError(
            f"{path}: inconsistent SAC timing: e={e:.9g}, but "
            f"b + (npts - 1) * delta={expected_e:.9g}; "
            f"npts={npts}, delta={delta:.9g}, b={b:.9g}."
        )

    return {
        "npts": npts,
        "delta": delta,
        "b": b,
        "e": e,
        "starttime": trace.stats.starttime,
        "o": read_optional_sac_float(trace, "o"),
        "a": read_optional_sac_float(trace, "a"),
        "t0": read_optional_sac_float(trace, "t0"),
        "t1": read_optional_sac_float(trace, "t1"),
    }


def validate_pair_timing(key: WaveformKey, component: str,
                         pair: WaveformPair) -> None:
    """Require matching sample count/rate; allow different absolute starts."""
    observed = get_trace_timing(pair.observed, pair.observed_path)
    synthetic = get_trace_timing(pair.synthetic, pair.synthetic_path)
    mismatches = []
    if observed["npts"] != synthetic["npts"]:
        mismatches.append("npts")

    obs_delta = float(observed["delta"])
    syn_delta = float(synthetic["delta"])
    delta_tolerance = max(1.0e-9, abs(obs_delta) * 1.0e-6)
    if not math.isclose(obs_delta, syn_delta, rel_tol=0.0,
                        abs_tol=delta_tolerance):
        mismatches.append("delta")

    if mismatches:
        raise ValueError(
            f"Sampling mismatch for {key.seed_prefix} component {component}: "
            f"{', '.join(mismatches)}\n"
            f"  observed  {pair.observed_path.name}: "
            f"npts={observed['npts']} delta={obs_delta:.9g}\n"
            f"  synthetic {pair.synthetic_path.name}: "
            f"npts={synthetic['npts']} delta={syn_delta:.9g}"
        )


def trace_time_relative_to_origin(trace: Trace,
                                  eq_origin_time: UTCDateTime) -> List[float]:
    """Return each sample time in seconds relative to earthquake origin."""
    first_sample = float(trace.stats.starttime - eq_origin_time)
    delta = float(trace.stats.delta)
    return [first_sample + index * delta for index in range(int(trace.stats.npts))]


def format_optional_sac_value(value: object) -> str:
    """Format an optional SAC reference-time header for diagnostics."""
    return "undefined" if value is None else f"{float(value):.3f}"

def set_station_geometry(station_data: StationWaveforms) -> None:
    """Read azimuth and distance once from the observed Z-component trace."""
    z_pair = station_data.components["Z"]
    source_path = z_pair.observed_path
    source_trace = z_pair.observed

    azimuth = read_required_sac_float(source_trace, "az", source_path)
    distance = read_required_sac_float(source_trace, "dist", source_path)

    # Accept the equivalent endpoint and store a conventional [0, 360) value.
    if math.isclose(azimuth, 360.0, rel_tol=0.0, abs_tol=1.0e-6):
        azimuth = 0.0
    if not 0.0 <= azimuth < 360.0:
        raise ValueError(
            f"{source_path}: SAC az must satisfy 0 <= az < 360 degrees; "
            f"found {azimuth:g}."
        )
    if distance < 0.0:
        raise ValueError(
            f"{source_path}: SAC dist must be nonnegative; found {distance:g} km."
        )

    station_data.azimuth_deg = azimuth
    station_data.distance_km = distance
    station_data.metadata_source_path = source_path


def read_waveform_pairs(data_dir: Path, quantity: str) -> WaveformData:
    """Discover and read complete observed/synthetic three-component sets."""
    data_dir = data_dir.expanduser().resolve()
    if not data_dir.is_dir():
        raise NotADirectoryError(f"Waveform directory does not exist: {data_dir}")

    observed_paths = sorted(data_dir.glob(f"*.{quantity}"))
    if not observed_paths:
        raise FileNotFoundError(
            f"No observed '*.{quantity}' files were found in {data_dir}"
        )

    data: WaveformData = {}
    for observed_path in observed_paths:
        key, component, channel = parse_observed_filename(observed_path, quantity)
        synthetic_path = Path(f"{observed_path}.syn")
        if not synthetic_path.is_file():
            raise FileNotFoundError(
                f"Synthetic file is missing for {observed_path.name}: "
                f"expected {synthetic_path.name}"
            )

        observed = read_one_sac(observed_path)
        synthetic = read_one_sac(synthetic_path)
        validate_trace_identity(observed, key, channel, observed_path)
        validate_trace_identity(synthetic, key, channel, synthetic_path)

        station_data = data.setdefault(key, StationWaveforms(key=key))
        if component in station_data.components:
            raise ValueError(
                f"Duplicate {component}-component data for {key.seed_prefix}"
            )
        station_data.components[component] = WaveformPair(
            observed=observed,
            synthetic=synthetic,
            observed_path=observed_path,
            synthetic_path=synthetic_path,
        )

    missing_messages = []
    for key, station_data in sorted(data.items()):
        missing = [c for c in COMPONENTS if c not in station_data.components]
        if missing:
            missing_messages.append(
                f"{key.seed_prefix}: missing component(s) {', '.join(missing)}"
            )
    if missing_messages:
        raise ValueError(
            "Each network.station.location must have E, N, and Z components:\n  "
            + "\n  ".join(missing_messages)
        )

    # Geometry is stored once per NET.STA.LOC. The observed Z trace is the
    # authoritative metadata source for this stage of the workflow.
    for station_data in data.values():
        set_station_geometry(station_data)
        for component in COMPONENTS:
            validate_pair_timing(
                station_data.key,
                component,
                station_data.components[component],
            )

    return data


def ordered_stations(
    data: WaveformData,
    station_order: str,
) -> List[StationWaveforms]:
    """Return stations in top-to-bottom plotting and reporting order."""
    stations = list(data.values())
    identity = lambda item: (
        item.key.network, item.key.station, item.key.location)
    if station_order == "azimuth":
        return sorted(
            stations,
            key=lambda item: (float(item.azimuth_deg),) + identity(item),
        )
    if station_order == "distance":
        return sorted(
            stations,
            key=lambda item: (float(item.distance_km),) + identity(item),
        )
    return sorted(stations, key=identity)


def print_summary(data: WaveformData, station_order: str, eq_origin_time: UTCDateTime) -> None:
    """Print station geometry and observed/synthetic trace information."""
    stations = ordered_stations(data, station_order)
    print(
        f"Read {len(stations)} complete three-component station set(s); "
        f"order={station_order}."
    )
    print(f"Earthquake origin time (plot t=0): {eq_origin_time}")

    for index, station_data in enumerate(stations, start=1):
        key = station_data.key
        print(
            f"\n{index:3d}  {key.seed_prefix}  quantity={key.quantity}  "
            f"az={station_data.azimuth_deg:.2f} deg  "
            f"dist={station_data.distance_km:.3f} km"
        )
        print(
            f"     metadata source: "
            f"{station_data.metadata_source_path.name}"
        )
        for component in COMPONENTS:
            pair = station_data.components[component]
            obs_timing = get_trace_timing(pair.observed, pair.observed_path)
            syn_timing = get_trace_timing(pair.synthetic, pair.synthetic_path)
            obs_start_rel = float(pair.observed.stats.starttime - eq_origin_time)
            syn_start_rel = float(pair.synthetic.stats.starttime - eq_origin_time)
            syn_minus_obs = syn_start_rel - obs_start_rel
            print(
                f"     {component}: "
                f"obs npts={obs_timing['npts']} "
                f"dt={float(obs_timing['delta']):g} s "
                f"b={float(obs_timing['b']):.3f} e={float(obs_timing['e']):.3f}; "
                f"syn npts={syn_timing['npts']} "
                f"dt={float(syn_timing['delta']):g} s "
                f"b={float(syn_timing['b']):.3f} e={float(syn_timing['e']):.3f}; "
                f"sampling=match"
            )
            print(
                f"        relative to origin: obs_start={obs_start_rel:.6f} s  "
                f"syn_start={syn_start_rel:.6f} s  "
                f"syn_minus_obs={syn_minus_obs:+.6f} s"
            )
            print(
                f"        obs starttime={obs_timing['starttime']}  "
                f"o={format_optional_sac_value(obs_timing['o'])} "
                f"a={format_optional_sac_value(obs_timing['a'])} "
                f"t0={format_optional_sac_value(obs_timing['t0'])} "
                f"t1={format_optional_sac_value(obs_timing['t1'])}"
            )
            print(
                f"        syn starttime={syn_timing['starttime']}  "
                f"o={format_optional_sac_value(syn_timing['o'])} "
                f"a={format_optional_sac_value(syn_timing['a'])} "
                f"t0={format_optional_sac_value(syn_timing['t0'])} "
                f"t1={format_optional_sac_value(syn_timing['t1'])}"
            )




def parse_station_tokens(value: str) -> set[str]:
    """Parse a comma-separated station selection into uppercase tokens."""
    return {token.strip().upper() for token in value.split(",") if token.strip()}


def read_station_selection_file(path: Path | None) -> set[str]:
    """Read station codes or NET.STA.LOC identifiers from a text file."""
    if path is None:
        return set()
    path = path.expanduser()
    if not path.is_file():
        raise FileNotFoundError(f"Station selection file does not exist: {path}")
    tokens: set[str] = set()
    for line_number, raw_line in enumerate(
            path.read_text(encoding="utf-8").splitlines(), start=1):
        line = raw_line.strip()
        if not line or line.startswith("#"):
            continue
        token = line.split("#", 1)[0].strip().upper()
        if not token:
            continue
        if "," in token:
            raise ValueError(
                f"{path}:{line_number}: use one station identifier per line")
        tokens.add(token)
    return tokens


def parse_components(value: str) -> Tuple[str, ...]:
    """Parse and validate a unique component list while preserving order."""
    components: List[str] = []
    for token in value.split(","):
        component = token.strip().upper()
        if not component:
            continue
        if component not in COMPONENTS:
            raise ValueError(
                f"Invalid component {component!r}; choose from E,N,Z")
        if component not in components:
            components.append(component)
    if not components:
        raise ValueError("--components must select at least one of E,N,Z")
    return tuple(components)


def station_matches(item: StationWaveforms, tokens: set[str]) -> bool:
    """Match either station code or full NET.STA.LOC identifier."""
    if not tokens:
        return False
    station = item.key.station.upper()
    seed_prefix = item.key.seed_prefix.upper()
    return station in tokens or seed_prefix in tokens


def filter_waveform_data(data: WaveformData,
                         args: argparse.Namespace) -> WaveformData:
    """Apply explicit station and geometric filters before scaling/pagination."""
    include = parse_station_tokens(args.include_stations)
    include.update(read_station_selection_file(args.station_file))
    exclude = parse_station_tokens(args.exclude_stations)

    numeric_values = (
        ("--min-distance-km", args.min_distance_km),
        ("--max-distance-km", args.max_distance_km),
        ("--min-azimuth-deg", args.min_azimuth_deg),
        ("--max-azimuth-deg", args.max_azimuth_deg),
    )
    for option, value in numeric_values:
        if value is not None and not math.isfinite(value):
            raise ValueError(f"{option} must be finite")
    if (args.min_distance_km is not None and args.max_distance_km is not None
            and args.min_distance_km > args.max_distance_km):
        raise ValueError("--min-distance-km cannot exceed --max-distance-km")
    if (args.min_azimuth_deg is not None and args.max_azimuth_deg is not None
            and args.min_azimuth_deg > args.max_azimuth_deg):
        raise ValueError("--min-azimuth-deg cannot exceed --max-azimuth-deg")
    for option, value in (("--min-azimuth-deg", args.min_azimuth_deg),
                          ("--max-azimuth-deg", args.max_azimuth_deg)):
        if value is not None and not 0.0 <= value <= 360.0:
            raise ValueError(f"{option} must be between 0 and 360 degrees")

    selected: WaveformData = {}
    rejected: List[Tuple[StationWaveforms, str]] = []
    for key, item in data.items():
        reason = None
        if include and not station_matches(item, include):
            reason = "not in include selection"
        elif station_matches(item, exclude):
            reason = "explicitly excluded"
        elif (args.min_distance_km is not None
              and float(item.distance_km) < args.min_distance_km):
            reason = f"distance below {args.min_distance_km:g} km"
        elif (args.max_distance_km is not None
              and float(item.distance_km) > args.max_distance_km):
            reason = f"distance above {args.max_distance_km:g} km"
        elif (args.min_azimuth_deg is not None
              and float(item.azimuth_deg) < args.min_azimuth_deg):
            reason = f"azimuth below {args.min_azimuth_deg:g} deg"
        elif (args.max_azimuth_deg is not None
              and float(item.azimuth_deg) > args.max_azimuth_deg):
            reason = f"azimuth above {args.max_azimuth_deg:g} deg"

        if reason is None:
            selected[key] = item
        else:
            rejected.append((item, reason))

    available_codes = {item.key.station.upper() for item in data.values()}
    available_ids = {item.key.seed_prefix.upper() for item in data.values()}
    unmatched = sorted(include - available_codes - available_ids)
    if unmatched:
        print(
            "WARNING: requested station identifiers were not found: "
            + ", ".join(unmatched),
            file=sys.stderr,
        )

    print("\nStation selection:")
    print(f"  available: {len(data)}")
    print(f"  selected:  {len(selected)}")
    print(f"  rejected:  {len(rejected)}")
    for item in ordered_stations(selected, args.station_order):
        print(
            f"    SELECT {item.key.seed_prefix}  "
            f"az={item.azimuth_deg:.2f} deg  dist={item.distance_km:.3f} km")
    for item, reason in sorted(
            rejected,
            key=lambda pair: (
                float(pair[0].azimuth_deg), pair[0].key.seed_prefix)):
        print(f"    REJECT {item.key.seed_prefix}: {reason}")

    if not selected:
        raise ValueError("No stations remain after applying selection filters")
    return selected


def waveform_array(trace: Trace, path: Path, factor: float) -> np.ndarray:
    """Return finite signed samples in requested display units."""
    values = np.asarray(trace.data, dtype=np.float64) * factor
    if values.size == 0:
        raise ValueError(f"{path}: empty waveform array")
    bad = np.flatnonzero(~np.isfinite(values))
    if bad.size:
        raise ValueError(f"{path}: {bad.size} nonfinite samples; first={bad[0]}")
    return values


def amplitude_scales(data: WaveformData, factor: float,
                     components: Tuple[str, ...]) -> Dict[str, object]:
    """Find per-trace, station, and global absolute maxima.

    Absolute values are used only to calculate scale factors. The original
    signed waveform samples are retained for plotting.
    """
    if not math.isfinite(factor) or factor == 0:
        raise ValueError("--input-unit-factor must be finite and nonzero")

    stations: Dict[WaveformKey, float] = {}
    traces: Dict[Tuple[WaveformKey, str, str], float] = {}
    global_obs = 0.0
    global_syn = 0.0

    for key, item in data.items():
        station_obs = 0.0
        station_syn = 0.0
        for component in components:
            pair = item.components[component]
            observed = waveform_array(
                pair.observed, pair.observed_path, factor)
            synthetic = waveform_array(
                pair.synthetic, pair.synthetic_path, factor)
            observed_absmax = float(np.max(np.abs(observed)))
            synthetic_absmax = float(np.max(np.abs(synthetic)))

            traces[(key, component, "observed")] = observed_absmax
            traces[(key, component, "synthetic")] = synthetic_absmax
            station_obs = max(station_obs, observed_absmax)
            station_syn = max(station_syn, synthetic_absmax)

        stations[key] = max(station_obs, station_syn)
        global_obs = max(global_obs, station_obs)
        global_syn = max(global_syn, station_syn)

    combined = max(global_obs, global_syn)
    if combined <= 0:
        raise ValueError("All observed and synthetic samples are zero")

    return {
        "station": stations,
        "trace": traces,
        "observed": global_obs,
        "synthetic": global_syn,
        "combined": combined,
    }


def print_trace_amplitude_summary(data: WaveformData, station_order: str,
                                  scales: Dict[str, object], unit: str,
                                  components: Tuple[str, ...]) -> None:
    """Print observed and synthetic absmax for every component trace."""
    print("\nAbsolute maximum amplitude per trace "
          "(absmax only; signed data retained for plotting):")
    trace_scales = scales["trace"]
    station_scales = scales["station"]
    for item in ordered_stations(data, station_order):
        key = item.key
        print(
            f"  {key.seed_prefix}  "
            f"station_combined_absmax={station_scales[key]:.8g} {unit}"
        )
        for component in components:
            observed_absmax = trace_scales[(key, component, "observed")]
            synthetic_absmax = trace_scales[(key, component, "synthetic")]
            print(
                f"    {component}: "
                f"observed_absmax={observed_absmax:.8g} {unit}  "
                f"synthetic_absmax={synthetic_absmax:.8g} {unit}"
            )


def resolve_reg_inv_out_path(data_dir: Path, reg_inv_out: Path) -> Path:
    """Resolve REG_INV_OUT, interpreting a relative name inside DATA_DIR."""
    path = reg_inv_out.expanduser()
    if not path.is_absolute():
        path = data_dir.expanduser().resolve() / path
    return path.resolve()


def read_variance_reductions(path: Path) -> Dict[str, float]:
    """Read observed-filename to variance-reduction-percent mappings."""
    import re

    if not path.is_file():
        raise FileNotFoundError(f"REG_INV_OUT file does not exist: {path}")

    pattern = re.compile(
        r"^\s*Station\s+(?P<filename>\S+?):?\s+.*?"
        r"Variance\s+reduction\s+"
        r"(?P<vr>[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[Ee][-+]?\d+)?)\s*$",
        re.IGNORECASE,
    )
    result: Dict[str, float] = {}
    for line_number, line in enumerate(
            path.read_text(encoding="utf-8", errors="replace").splitlines(),
            start=1):
        match = pattern.match(line)
        if match is None:
            continue
        filename = match.group("filename").rstrip(":")
        vr = float(match.group("vr"))
        if not math.isfinite(vr):
            raise ValueError(
                f"{path}:{line_number}: nonfinite VR for {filename}")
        if filename in result and not math.isclose(
                result[filename], vr, rel_tol=0.0, abs_tol=1.0e-8):
            raise ValueError(
                f"{path}:{line_number}: conflicting VR values for {filename}: "
                f"{result[filename]:g} and {vr:g}")
        result[filename] = vr

    if not result:
        raise ValueError(
            f"No Station ... Variance reduction ... records found in {path}")
    return result


def read_time_shifts(path: Path) -> Dict[str, float]:
    """Read observed-filename to inversion ttshift mappings in seconds.

    Expected example:
      Station BK.BARR.00.HNN.disp comp N lat 39.304771 lon -123.197479
      ttshift -0.900000, sacshift 0, numGFset 1
    """
    import re

    if not path.is_file():
        raise FileNotFoundError(f"REG_INV_OUT file does not exist: {path}")

    pattern = re.compile(
        r"^\s*Station\s+(?P<filename>\S+?):?\s+.*?"
        r"\bttshift\s+"
        r"(?P<ttshift>[-+]?(?:\d+(?:\.\d*)?|\.\d+)"
        r"(?:[Ee][-+]?\d+)?)\s*,?",
        re.IGNORECASE,
    )

    result: Dict[str, float] = {}
    for line_number, line in enumerate(
            path.read_text(encoding="utf-8", errors="replace").splitlines(),
            start=1):
        match = pattern.match(line)
        if match is None:
            continue
        filename = match.group("filename").rstrip(":")
        ttshift = float(match.group("ttshift"))
        if not math.isfinite(ttshift):
            raise ValueError(
                f"{path}:{line_number}: nonfinite ttshift for {filename}")
        if filename in result and not math.isclose(
                result[filename], ttshift, rel_tol=0.0, abs_tol=1.0e-8):
            raise ValueError(
                f"{path}:{line_number}: conflicting ttshift values for "
                f"{filename}: {result[filename]:g} and {ttshift:g}")
        result[filename] = ttshift

    if not result:
        raise ValueError(
            f"No 'Station ... ttshift ...' records found in {path}")
    return result


def validate_time_shifts(data: WaveformData,
                         components: Tuple[str, ...],
                         ttshift_by_filename: Dict[str, float],
                         source_path: Path) -> None:
    """Require one ttshift record for each selected observed trace."""
    missing: List[str] = []
    for station_data in data.values():
        for component in components:
            filename = station_data.components[component].observed_path.name
            if filename not in ttshift_by_filename:
                missing.append(filename)
    if missing:
        shown = ", ".join(sorted(missing)[:10])
        if len(missing) > 10:
            shown += f", ... ({len(missing)} total)"
        raise ValueError(
            f"Missing ttshift record(s) in {source_path}: {shown}")


def print_time_shift_summary(data: WaveformData,
                             station_order: str,
                             components: Tuple[str, ...],
                             ttshift_by_filename: Dict[str, float],
                             decimals: int,
                             origin: UTCDateTime,
                             apply_ttshift: bool,
                             ttshift_sign: float) -> None:
    """Print stored and applied inversion shifts for every selected trace."""
    print("\nInversion time shift per selected observed trace:")
    print(
        "  NOTE: synthetic waveforms were already shifted by ttshift "
        "in the inversion code.")
    print(
        f"  apply_ttshift={apply_ttshift}  ttshift_sign={ttshift_sign:+.0f}  "
        "rule: observed_plot_time = SAC_time + sign * ttshift")
    for item in ordered_stations(data, station_order):
        for component in components:
            pair = item.components[component]
            filename = pair.observed_path.name
            ttshift = ttshift_by_filename[filename]
            applied = ttshift_sign * ttshift if apply_ttshift else 0.0
            original_start = float(pair.observed.stats.starttime - origin)
            shifted_start = original_start + applied
            print(
                f"  {filename}: "
                f"ttshift={ttshift:.{decimals}f} s  "
                f"applied={applied:+.{decimals}f} s  "
                f"start={original_start:.{decimals}f} -> "
                f"{shifted_start:.{decimals}f} s")

def validate_variance_reductions(data: WaveformData,
                                 components: Tuple[str, ...],
                                 vr_by_filename: Dict[str, float],
                                 source_path: Path) -> None:
    """Require a VR record for every selected observed component trace."""
    missing: List[str] = []
    for station_data in data.values():
        for component in components:
            filename = station_data.components[component].observed_path.name
            if filename not in vr_by_filename:
                missing.append(filename)
    if missing:
        shown = ", ".join(sorted(missing)[:10])
        if len(missing) > 10:
            shown += f", ... ({len(missing)} total)"
        raise ValueError(
            f"Missing VR record(s) in {source_path}: {shown}")

def print_variance_reduction_summary(data: WaveformData,
                                     station_order: str,
                                     components: Tuple[str, ...],
                                     vr_by_filename: Dict[str, float],
                                     decimals: int) -> None:
    """Print VR for every selected observed component trace."""
    print("\nVariance reduction per selected observed trace:")
    for item in ordered_stations(data, station_order):
        for component in components:
            filename = item.components[component].observed_path.name
            vr = vr_by_filename[filename]
            print(
                f"  {filename}: VR={vr:.{decimals}f} %")

def observed_plot_time(pair: WaveformPair,
                       origin: UTCDateTime,
                       args: argparse.Namespace) -> np.ndarray:
    """Return observed plotting times.

    By default no shift is added because synthetic waveforms were already
    shifted by ttshift in the inversion code. The optional observed shift is
    retained only for deliberate diagnostic comparisons.
    """
    time = np.asarray(
        trace_time_relative_to_origin(pair.observed, origin),
        dtype=float,
    )
    if args.apply_ttshift:
        ttshift = args.ttshift_by_filename[pair.observed_path.name]
        time = time + args.ttshift_sign * ttshift
    return time

def plot_waveforms(data: WaveformData, station_order: str,
                   origin: UTCDateTime, scales: Dict[str, object],
                   args: argparse.Namespace) -> None:
    """Create consistently scaled E/N/Z pages and a multipage PDF."""
    from matplotlib.backends.backend_pdf import PdfPages

    if args.station_spacing <= 0 or args.waveform_half_height <= 0:
        raise ValueError("station spacing and waveform half height must be positive")
    if args.stations_per_page < 0:
        raise ValueError("--stations-per-page must be zero or a positive integer")
    if args.time_min is not None and not math.isfinite(args.time_min):
        raise ValueError("--time-min must be finite")
    if args.time_max is not None and not math.isfinite(args.time_max):
        raise ValueError("--time-max must be finite")
    if (args.time_min is not None and args.time_max is not None
            and args.time_min >= args.time_max):
        raise ValueError("--time-min must be less than --time-max")
    if args.station_label_size <= 0 or args.metadata_label_size <= 0:
        raise ValueError("station and metadata label sizes must be positive")
    text_sizes = (
        args.station_label_size,
        args.vr_font_size,
        args.az_font_size,
        args.dist_font_size,
        args.metadata_label_size,
        args.component_title_size,
        args.axis_label_size,
        args.tick_label_size,
        args.legend_font_size,
        args.footer_font_size,
        args.max_amp_font_size,
    )
    if any(not math.isfinite(size) or size <= 0 for size in text_sizes):
        raise ValueError("all text-size options must be finite and positive")
    if args.max_amp_decimals < 0:
        raise ValueError("--max-amp-decimals must be zero or positive")
    if not math.isfinite(args.max_amp_x):
        raise ValueError("--max-amp-x must be finite")
    if not math.isfinite(args.max_amp_y_offset):
        raise ValueError("--max-amp-y-offset must be finite")
    row_label_positions = (
        args.station_name_x, args.station_name_y_offset,
        args.vr_x, args.vr_y_offset,
        args.az_x, args.az_y_offset,
        args.dist_x, args.dist_y_offset,
    )
    if any(not math.isfinite(value) for value in row_label_positions):
        raise ValueError("all row-label position options must be finite")

    stations = ordered_stations(data, station_order)
    plot_components = args.plot_components
    global_max = float(scales["combined"])

    # Determine one common data interval before creating pages. Every page uses
    # identical x limits, enabling direct page-to-page comparison.
    data_starts: List[float] = []
    data_ends: List[float] = []
    for item in stations:
        for component in plot_components:
            pair = item.components[component]
            observed_time = observed_plot_time(pair, origin, args)
            synthetic_time = np.asarray(
                trace_time_relative_to_origin(pair.synthetic, origin),
                dtype=float,
            )
            data_starts.extend((
                float(observed_time[0]), float(synthetic_time[0])))
            data_ends.extend((
                float(observed_time[-1]), float(synthetic_time[-1])))

    data_time_min = min(data_starts)
    data_time_max = max(data_ends)
    plot_time_min = data_time_min if args.time_min is None else args.time_min
    plot_time_max = data_time_max if args.time_max is None else args.time_max
    if plot_time_max < data_time_min or plot_time_min > data_time_max:
        raise ValueError(
            f"Requested time window [{plot_time_min:g}, {plot_time_max:g}] s "
            f"does not overlap data [{data_time_min:g}, {data_time_max:g}] s")

    page_size = len(stations) if args.stations_per_page == 0 else args.stations_per_page
    pages = [stations[index:index + page_size]
             for index in range(0, len(stations), page_size)]
    output = args.plot_file.expanduser()
    output.parent.mkdir(parents=True, exist_ok=True)
    is_pdf = output.suffix.lower() == ".pdf"
    pdf = PdfPages(output) if is_pdf else None
    written_files: List[Path] = []

    try:
        for page_number, page_stations in enumerate(pages, start=1):
            nsta = len(page_stations)
            height = max(6.0, args.row_height * nsta + 2.1)
            fig, axes = plt.subplots(
                1, len(plot_components),
                figsize=(args.figure_width, height),
                sharex=True,
                sharey=True,
            )
            axes = np.atleast_1d(axes)
            baselines = (
                np.arange(nsta - 1, -1, -1, dtype=float)
                * args.station_spacing
            )

            for row, item in enumerate(page_stations):
                base = baselines[row]
                for column, component in enumerate(plot_components):
                    pair = item.components[component]
                    observed = waveform_array(
                        pair.observed, pair.observed_path,
                        args.input_unit_factor)
                    synthetic = waveform_array(
                        pair.synthetic, pair.synthetic_path,
                        args.input_unit_factor)
                    observed_time = observed_plot_time(
                        pair, origin, args)
                    synthetic_time = np.asarray(
                        trace_time_relative_to_origin(pair.synthetic, origin),
                        dtype=float,
                    )

                    plot_height = (
                        args.waveform_half_height * args.station_spacing)
                    if args.normalization == "global":
                        observed_gain = synthetic_gain = plot_height / global_max
                    elif args.normalization == "station":
                        denominator = (
                            float(scales["station"][item.key]) or global_max)
                        observed_gain = synthetic_gain = plot_height / denominator
                    elif args.normalization == "trace":
                        observed_max = float(
                            scales["trace"][(item.key, component, "observed")])
                        synthetic_max = float(
                            scales["trace"][(item.key, component, "synthetic")])
                        observed_gain = (
                            plot_height / observed_max
                            if observed_max > 0.0 else 0.0)
                        synthetic_gain = (
                            plot_height / synthetic_max
                            if synthetic_max > 0.0 else 0.0)
                    else:
                        observed_gain = synthetic_gain = 1.0

                    # Preserve waveform polarity; abs() is used only for scales.
                    axes[column].plot(
                        observed_time,
                        base + observed * observed_gain,
                        color="black", lw=args.line_width, zorder=3,
                    )
                    axes[column].plot(
                        synthetic_time,
                        base + synthetic * synthetic_gain,
                        color="#e63223", lw=args.line_width, zorder=4,
                    )
                    vr_percent = args.vr_by_filename[pair.observed_path.name]
                    axes[column].text(
                        args.vr_x,
                        base + args.vr_y_offset * args.station_spacing,
                        f"VR: {vr_percent:.{args.vr_decimals}f} %",
                        transform=axes[column].get_yaxis_transform(),
                        ha="left", va="bottom",
                        fontsize=args.vr_font_size,
                        fontweight="normal", color="black", clip_on=True,
                    )
                    if component == "E":
                        axes[column].text(
                            args.station_name_x,
                            base + args.station_name_y_offset * args.station_spacing,
                            f"{item.key.network}.{item.key.station}",
                            transform=axes[column].get_yaxis_transform(),
                            ha="right", va="bottom",
                            fontsize=args.station_label_size,
                            fontweight="normal", color="black", clip_on=True,
                        )
                    observed_absmax = float(
                        scales["trace"][(item.key, component, "observed")])
                    max_amp_text = (
                        f"Max. amp.: {observed_absmax:.{args.max_amp_decimals}f} "
                        f"{args.display_unit}")
                    axes[column].text(
                        args.max_amp_x,
                        base - args.max_amp_y_offset * args.station_spacing,
                        max_amp_text,
                        transform=axes[column].get_yaxis_transform(),
                        ha="left", va="top",
                        fontsize=args.max_amp_font_size,
                        fontweight="normal",
                        color="0.25",
                        clip_on=True,
                    )

                if "Z" in plot_components:
                    vertical_axis = axes[plot_components.index("Z")]
                    vertical_axis.text(
                        args.az_x,
                        base + args.az_y_offset * args.station_spacing,
                        f"AZ: {item.azimuth_deg:.0f} deg.",
                        transform=vertical_axis.get_yaxis_transform(),
                        ha="right", va="bottom",
                        fontsize=args.az_font_size,
                        fontweight="normal", clip_on=True,
                    )
                    vertical_axis.text(
                        args.dist_x,
                        base - args.dist_y_offset * args.station_spacing,
                        f"DIST: {item.distance_km:.1f} km",
                        transform=vertical_axis.get_yaxis_transform(),
                        ha="right", va="top",
                        fontsize=args.dist_font_size,
                        fontweight="normal", clip_on=True,
                    )
                    if args.annotate_station_absmax:
                        station_absmax = float(scales["station"][item.key])
                        vertical_axis.text(
                            args.dist_x,
                            base - (args.dist_y_offset + 0.10) * args.station_spacing,
                            f"Station max. {station_absmax:.4g} "
                            f"{args.display_unit}",
                            transform=vertical_axis.get_yaxis_transform(),
                            ha="left", va="top",
                            fontsize=args.metadata_label_size,
                            fontweight="normal", clip_on=True,
                        )

            component_titles = {"E": "East", "N": "North", "Z": "Vertical"}
            for axis, component in zip(axes, plot_components):
                title = component_titles[component]
                axis.set_title(
                    title, fontsize=args.component_title_size,
                    fontweight="normal")
                axis.spines["top"].set_visible(False)
                axis.spines["right"].set_visible(False)
                axis.set_xlim(plot_time_min, plot_time_max)
                axis.set_xlabel(
                    "Elapsed time from EQ origin time (s)",
                    fontsize=args.axis_label_size,
                    fontweight="normal")
                axis.tick_params(
                    axis="x", labelsize=args.tick_label_size)

            margin = 0.55 * args.station_spacing
            axes[0].set_ylim(-margin, baselines[0] + margin)
            axes[0].set_yticks(baselines)
            axes[0].set_yticklabels(["" for _ in page_stations])
            axes[0].set_ylabel(
                args.y_axis_label, fontsize=args.axis_label_size,
                fontweight="normal")

            observed_handle = plt.Line2D(
                [], [], color="black", lw=args.line_width, label="Observed")
            synthetic_handle = plt.Line2D(
                [], [], color="#e63223", lw=args.line_width,
                label="Synthetic")
            fig.legend(
                handles=[observed_handle, synthetic_handle],
                loc="upper center", bbox_to_anchor=(0.47, 0.995),
                frameon=False, ncol=2,
                fontsize=args.legend_font_size,
            )

            if args.normalization == "station":
                scale_description = (
                    "station: combined observed+synthetic E/N/Z absmax")
            elif args.normalization == "trace":
                scale_description = "trace: each waveform uses its own absmax"
            elif args.normalization == "global":
                scale_description = (
                    f"global: absmax={global_max:.6g} {args.display_unit}")
            else:
                scale_description = "none: physical signed amplitudes"

            if args.footer:
                footer = (
                    f"Origin: {origin}   |   Order: increasing {station_order}   |   "
                    f"Normalization: {scale_description}   |   "
                f"Components: {','.join(plot_components)}   |   "
                    f"ttshift: {'on' if args.apply_ttshift else 'off'} "
                    f"(sign={args.ttshift_sign:+.0f})   |   "
                    f"Observed: black; Synthetic: red   |   "
                    f"Page {page_number}/{len(pages)}")
                fig.text(
                    0.5, 0.012, footer, ha="center", va="bottom",
                    fontsize=args.footer_font_size,
                    fontweight="normal")

            # Leave room for figure legend, footer, and right-side metadata.
            fig.tight_layout(rect=(0.03, 0.045, 0.99, 0.965), w_pad=0.7)

            if pdf is not None:
                pdf.savefig(fig, bbox_inches="tight")
            else:
                page_file = output.with_name(
                    f"{output.stem}_page{page_number:02d}{output.suffix}")
                fig.savefig(page_file, bbox_inches="tight")
                written_files.append(page_file)
            plt.close(fig)
    finally:
        if pdf is not None:
            pdf.close()

    if pdf is not None:
        print(f"Multipage waveform plot: {output} ({len(pages)} page(s))")
    else:
        print(f"Waveform plot pages: {len(written_files)}")
        for page_file in written_files:
            print(f"  {page_file}")


def calculate_vr(observed: np.ndarray, synthetic: np.ndarray) -> float:
    """Return variance reduction in percent for equal-length arrays."""
    denominator = float(np.sum(observed ** 2))
    if denominator <= 0.0:
        return float("nan")
    residual = observed - synthetic
    return (1.0 - float(np.sum(residual ** 2)) / denominator) * 100.0


def common_comparison_arrays(pair: WaveformPair,
                             origin: UTCDateTime,
                             time_min: float | None,
                             time_max: float | None) -> Tuple[np.ndarray, np.ndarray, float, float, float]:
    """Interpolate synthetic data onto observed sample times in the common window."""
    observed_time = np.asarray(
        trace_time_relative_to_origin(pair.observed, origin), dtype=float)
    synthetic_time = np.asarray(
        trace_time_relative_to_origin(pair.synthetic, origin), dtype=float)
    observed = np.asarray(pair.observed.data, dtype=np.float64)
    synthetic = np.asarray(pair.synthetic.data, dtype=np.float64)

    start = max(float(observed_time[0]), float(synthetic_time[0]))
    end = min(float(observed_time[-1]), float(synthetic_time[-1]))
    if time_min is not None:
        start = max(start, float(time_min))
    if time_max is not None:
        end = min(end, float(time_max))
    if start >= end:
        raise ValueError(
            f"No common correlation window for {pair.observed_path.name}: "
            f"{start:g} to {end:g} s")

    mask = (observed_time >= start) & (observed_time <= end)
    grid = observed_time[mask]
    observed_common = observed[mask]
    if grid.size < 3:
        raise ValueError(
            f"Too few common samples for {pair.observed_path.name}: {grid.size}")
    synthetic_common = np.interp(grid, synthetic_time, synthetic)
    delta = float(pair.observed.stats.delta)
    return observed_common, synthetic_common, delta, float(grid[0]), float(grid[-1])


def shift_observed_for_vr(observed: np.ndarray,
                          correlation_lag_samples: int) -> Tuple[np.ndarray, np.ndarray]:
    """Shift observed by -CC lag and return only valid, nonwrapped overlap."""
    sample_shift = -int(correlation_lag_samples)
    shifted = np.roll(observed, sample_shift)
    valid = np.ones(observed.size, dtype=bool)
    if sample_shift > 0:
        valid[:sample_shift] = False
    elif sample_shift < 0:
        valid[sample_shift:] = False
    return shifted, valid


def analyze_trace_correlation(pair: WaveformPair,
                              origin: UTCDateTime,
                              current_ttshift: float,
                              args: argparse.Namespace) -> Dict[str, object]:
    """Compute residual lag, CC, and VR change for one component trace."""
    observed, synthetic, delta, window_start, window_end = common_comparison_arrays(
        pair, origin, args.time_min, args.time_max)
    max_lag_samples = int(round(args.cc_max_lag_s / delta))
    max_lag_samples = min(max_lag_samples, observed.size - 2)
    if max_lag_samples < 0:
        raise ValueError("--cc-max-lag-s produces an invalid lag window")

    correlation = correlate(
        observed, synthetic, max_lag_samples, method="direct")
    middle = (len(correlation) - 1) // 2
    cc_zero_lag = float(correlation[middle])
    cc_lag_samples, cc_max_positive = xcorr_max(
        correlation, abs_max=False)
    _, cc_max_absolute = xcorr_max(correlation, abs_max=True)
    cc_lag_samples = int(cc_lag_samples)
    residual_lag_s = cc_lag_samples * delta
    proposed_ttshift = current_ttshift + residual_lag_s

    vr_before = calculate_vr(observed, synthetic)
    shifted_observed, valid = shift_observed_for_vr(
        observed, cc_lag_samples)
    vr_after = calculate_vr(shifted_observed[valid], synthetic[valid])
    vr_improvement = vr_after - vr_before
    accepted = bool(
        float(cc_max_positive) >= args.cc_min_correlation
        and vr_improvement >= args.cc_min_vr_improvement
        and cc_lag_samples != 0
    )

    return {
        "observed_filename": pair.observed_path.name,
        "cc_zero_lag": cc_zero_lag,
        "cc_max_positive": float(cc_max_positive),
        "cc_max_absolute": float(cc_max_absolute),
        "cc_lag_samples": cc_lag_samples,
        "residual_lag_s": residual_lag_s,
        "current_ttshift_s": current_ttshift,
        "proposed_ttshift_s": proposed_ttshift,
        "vr_before_percent": vr_before,
        "vr_after_percent": vr_after,
        "vr_improvement_points": vr_improvement,
        "accepted_trace": accepted,
        "comparison_start_s": window_start,
        "comparison_end_s": window_end,
        "comparison_npts": int(observed.size),
    }


def run_correlation_analysis(data: WaveformData,
                             station_order: str,
                             components: Tuple[str, ...],
                             origin: UTCDateTime,
                             args: argparse.Namespace) -> List[Dict[str, object]]:
    """Analyze every selected trace and print detailed diagnostics."""
    if not math.isfinite(args.cc_max_lag_s) or args.cc_max_lag_s < 0.0:
        raise ValueError("--cc-max-lag-s must be finite and nonnegative")
    if not math.isfinite(args.cc_min_correlation):
        raise ValueError("--cc-min-correlation must be finite")
    if not math.isfinite(args.cc_min_vr_improvement):
        raise ValueError("--cc-min-vr-improvement must be finite")

    rows: List[Dict[str, object]] = []
    print("\nResidual observed/synthetic correlation analysis:")
    print(
        "  Synthetic waveforms already include inversion ttshift; residual CC "
        "is diagnostic and does not alter the figure.")
    for item in ordered_stations(data, station_order):
        for component in components:
            pair = item.components[component]
            current_ttshift = args.ttshift_by_filename[pair.observed_path.name]
            result = analyze_trace_correlation(
                pair, origin, current_ttshift, args)
            result.update({
                "network": item.key.network,
                "station": item.key.station,
                "location": item.key.location,
                "component": component,
                "azimuth_deg": float(item.azimuth_deg),
                "distance_km": float(item.distance_km),
                "station_selected_component": False,
                "station_update_accepted": False,
                "station_output_ttshift_s": current_ttshift,
            })
            rows.append(result)
            print(
                f"  {pair.observed_path.name}: "
                f"CC0={result['cc_zero_lag']:.4f}  "
                f"CCmax+={result['cc_max_positive']:.4f}  "
                f"lag={result['cc_lag_samples']:+d} samples "
                f"({result['residual_lag_s']:+.3f} s)  "
                f"ttshift={current_ttshift:+.3f} -> "
                f"{result['proposed_ttshift_s']:+.3f} s  "
                f"VR={result['vr_before_percent']:.2f} -> "
                f"{result['vr_after_percent']:.2f} %  "
                f"dVR={result['vr_improvement_points']:+.2f}")
    return rows


def select_station_updates(data: WaveformData,
                           rows: List[Dict[str, object]],
                           args: argparse.Namespace) -> Dict[WaveformKey, Dict[str, object]]:
    """Choose the component with maximum post-shift VR for each station."""
    selected: Dict[WaveformKey, Dict[str, object]] = {}
    for item in data.values():
        station_rows = [
            row for row in rows
            if row["network"] == item.key.network
            and row["station"] == item.key.station
            and row["location"] == item.key.location
        ]
        finite_rows = [row for row in station_rows
                       if math.isfinite(float(row["vr_after_percent"]))]
        if not finite_rows:
            raise ValueError(
                f"No finite post-shift VR for {item.key.seed_prefix}")
        best = max(finite_rows, key=lambda row: float(row["vr_after_percent"]))
        accepted = bool(best["accepted_trace"])
        output_shift = (
            float(best["proposed_ttshift_s"])
            if accepted else float(best["current_ttshift_s"])
        )
        best["station_selected_component"] = True
        for row in station_rows:
            row["station_update_accepted"] = accepted
            row["station_output_ttshift_s"] = output_shift
        selected[item.key] = {
            "best_component": best["component"],
            "accepted": accepted,
            "output_ttshift_s": output_shift,
            "vr_after_percent": best["vr_after_percent"],
            "cc_max_positive": best["cc_max_positive"],
        }
        print(
            f"  Station {item.key.seed_prefix}: best={best['component']}  "
            f"accepted={accepted}  output_ttshift={output_shift:+.3f} s  "
            f"VR_after={float(best['vr_after_percent']):.2f} %")
    return selected


def write_correlation_csv(path: Path,
                          rows: List[Dict[str, object]]) -> None:
    """Write detailed per-trace analysis results."""
    path = path.expanduser()
    path.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        "network", "station", "location", "component",
        "observed_filename", "azimuth_deg", "distance_km",
        "current_ttshift_s", "cc_zero_lag", "cc_max_positive",
        "cc_max_absolute", "cc_lag_samples", "residual_lag_s",
        "proposed_ttshift_s", "vr_before_percent", "vr_after_percent",
        "vr_improvement_points", "accepted_trace",
        "station_selected_component", "station_update_accepted",
        "station_output_ttshift_s", "comparison_start_s",
        "comparison_end_s", "comparison_npts",
    ]
    with path.open("w", newline="", encoding="utf-8") as output:
        writer = csv.DictWriter(output, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)
    print(f"Correlation CSV: {path}")


def write_ttshift_update(path: Path,
                         data: WaveformData,
                         station_order: str,
                         station_updates: Dict[WaveformKey, Dict[str, object]],
                         gps_line: str) -> None:
    """Write a station-level update for E/N/Z in the legacy inversion format."""
    stations = ordered_stations(data, station_order)
    output_components = ("Z", "N", "E")
    reference_pairs = [item.components["Z"] for item in stations]
    npts_values = {int(pair.observed.stats.npts) for pair in reference_pairs}
    delta_values = {float(pair.observed.stats.delta) for pair in reference_pairs}
    if len(npts_values) != 1 or len(delta_values) != 1:
        raise ValueError(
            "ttshift_update.txt requires common npts and delta across stations")
    npts = npts_values.pop()
    delta = delta_values.pop()

    path = path.expanduser()
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as output:
        output.write(f"{len(stations) * len(output_components)} S\n")
        output.write(f"{npts} {delta:.6f}\n")
        for item in stations:
            source = item.components["Z"].observed
            sac = getattr(source.stats, "sac", None)
            if sac is None or "stla" not in sac or "stlo" not in sac:
                raise ValueError(
                    f"Missing SAC stla/stlo for {item.key.seed_prefix}")
            latitude = float(sac.stla)
            longitude = float(sac.stlo)
            shift = float(station_updates[item.key]["output_ttshift_s"])
            channel_prefix = item.components["Z"].observed.stats.channel[:-1]
            for component in output_components:
                filename = (
                    f"{item.key.seed_prefix}.{channel_prefix}{component}."
                    f"{item.key.quantity}")
                output.write(
                    f"{latitude:.6f} {longitude:.6f} {filename} "
                    f"{component} {shift:.3f} 1\n")
        if gps_line:
            output.write(gps_line.rstrip("\n") + "\n")
    print(f"ttshift update file: {path}")

def main() -> int:
    args = parse_arguments()
    try:
        eq_origin_time = UTCDateTime(args.eq_origin_time)
        args.plot_components = parse_components(args.components)
        if args.vr_decimals < 0:
            raise ValueError("--vr-decimals must be zero or positive")
        if args.ttshift_decimals < 0:
            raise ValueError("--ttshift-decimals must be zero or positive")
        if args.ttshift_sign not in (-1.0, 1.0):
            raise ValueError("--ttshift-sign must be either -1 or 1")
        reg_inv_out_path = resolve_reg_inv_out_path(
            args.data_dir, args.reg_inv_out)
        args.vr_by_filename = read_variance_reductions(reg_inv_out_path)
        args.ttshift_by_filename = read_time_shifts(reg_inv_out_path)
        all_data = read_waveform_pairs(args.data_dir, args.quantity)
        data = filter_waveform_data(all_data, args)
        validate_variance_reductions(
            data, args.plot_components, args.vr_by_filename,
            reg_inv_out_path)
        validate_time_shifts(
            data, args.plot_components, args.ttshift_by_filename,
            reg_inv_out_path)
        print(f"REG_INV_OUT: {reg_inv_out_path}")
        print(f"Variance reduction records: {len(args.vr_by_filename)}")
        print(f"Time-shift records: {len(args.ttshift_by_filename)}")
        print(f"Selected components: {','.join(args.plot_components)}")
        print_variance_reduction_summary(
            data, args.station_order, args.plot_components,
            args.vr_by_filename, args.vr_decimals)
        print_time_shift_summary(
            data, args.station_order, args.plot_components,
            args.ttshift_by_filename, args.ttshift_decimals,
            eq_origin_time, args.apply_ttshift, args.ttshift_sign)
        print_summary(data, args.station_order, eq_origin_time)
        scales = amplitude_scales(
            data, args.input_unit_factor, args.plot_components)
        print("\nAmplitude diagnostics (absmax; signed data retained):")
        print(f"  global observed:  {scales['observed']:.8g} {args.display_unit}")
        print(f"  global synthetic: {scales['synthetic']:.8g} {args.display_unit}")
        print(f"  global combined:  {scales['combined']:.8g} {args.display_unit}")
        print_trace_amplitude_summary(
            data, args.station_order, scales, args.display_unit,
            args.plot_components)
        if args.make_correlation_analysis:
            correlation_rows = run_correlation_analysis(
                data, args.station_order, args.plot_components,
                eq_origin_time, args)
            print("\nStation-level ttshift update selection:")
            station_updates = select_station_updates(
                data, correlation_rows, args)
            write_correlation_csv(args.cc_output_csv, correlation_rows)
            write_ttshift_update(
                args.ttshift_update_file, data, args.station_order,
                station_updates, args.ttshift_update_gps_line)
        if args.make_plot:
            plot_waveforms(data, args.station_order, eq_origin_time, scales, args)
    except (FileNotFoundError, NotADirectoryError, OSError, ValueError) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())