fkrelloc.py

#!/usr/bin/env python3
"""
Steven J. Gibbons 2026/01/11
copilot was used to generate parts of this code.

Calculates the relative locations of two seismic events in Cartesian geometry
(with units in km). Requires CC-traces in SAC format, e.g. generated by the ccdtest
program. An input file needs to specify the slowness vectors describing the outward
going wavefield for each station used.

fkrelloc.py
Reads a station/phase/SAC-path list (5 columns) and can:
 1) Plot each SAC file individually via ObsPy.
 2) Build a common time axis (union or intersection), interpolate/pad all traces,
    and plot an overlay of all traces on the same figure.
 3) Attach array coordinates (x=sx, y=sy) to each trace and perform a broadband
    f–k (frequency–slowness) scan using ObsPy's array_processing, then plot and/or
    export the accumulated f–k grid to NetCDF for external tools.

Columns (whitespace-separated):
1) sac_path (string: absolute or relative if --allow-relative-paths)
2) station (string)
3) phase (string)
4) sx (float)  # array x or y depending on --coord-order
5) sy (float)  # array y or x depending on --coord-order

Prerequisite:
  pip install obspy
Optional for NetCDF:
  pip install netCDF4
  (fallback to scipy.io.netcdf, last-resort NPZ)
"""
import argparse
import os
from typing import List, Dict, Tuple, Optional
import numpy as np
from obspy.core.util import AttribDict  # coordinates attachment


# ------------------------------- Parsing --------------------------------------
def parse_line(line: str, line_number: int,
               allow_relative_paths: bool = False,
               base_dir: Optional[str] = None,
               require_file_exists: bool = True) -> Tuple[Dict, str]:
    parts = line.strip().split()
    if len(parts) != 5:
        return {}, f"Line {line_number}: expected 5 columns, got {len(parts)} -> {parts}"
    sac_path, station, phase, sx_str, sy_str = parts
    try:
        sx = -float(sx_str)
        sy = -float(sy_str)
    except ValueError:
        return {}, f"Line {line_number}: sx/sy not valid floats: '{sx_str}', '{sy_str}'"

    if not os.path.isabs(sac_path):
        if not allow_relative_paths:
            return {}, f"Line {line_number}: SAC path is not absolute (use --allow-relative-paths): '{sac_path}'"
        sac_path = os.path.normpath(os.path.join(base_dir or "", sac_path))

    if require_file_exists and not os.path.exists(sac_path):
        return {}, f"Line {line_number}: SAC file does not exist: '{sac_path}'"

    return {"sac_path": sac_path, "station": station, "phase": phase, "sx": sx, "sy": sy}, ""


def read_station_phase_list(filename: str,
                            allow_missing_files: bool = False,
                            allow_relative_paths: bool = False) -> List[Dict]:
    if not os.path.exists(filename):
        raise FileNotFoundError(f"List file not found: {filename}")
    base_dir = os.path.dirname(os.path.abspath(filename))
    records, errors = [], []
    with open(filename, "r", encoding="utf-8") as f:
        for i, raw in enumerate(f, start=1):
            line = raw.strip()
            if not line or line.startswith("#"):
                continue
            rec, err = parse_line(line, i, allow_relative_paths, base_dir, require_file_exists=(not allow_missing_files))
            if err:
                if allow_missing_files and ("does not exist" in err):
                    parts = line.split()
                    sac_path, station, phase, sx_str, sy_str = parts
                    if not os.path.isabs(sac_path) and allow_relative_paths:
                        sac_path = os.path.normpath(os.path.join(base_dir, sac_path))
                    rec = {"sac_path": sac_path, "station": station, "phase": phase,
                           "sx": float(sx_str), "sy": float(sy_str)}
                    errors.append(f"WARNING (line {i}): {err}")
                    records.append(rec)
                else:
                    errors.append(err)
            else:
                records.append(rec)
    if errors and not allow_missing_files:
        raise ValueError("Errors encountered while parsing the list:\n" + "\n".join(errors))
    if errors and allow_missing_files:
        print("Warnings during parsing:")
        for e in errors:
            print(" " + e)
    return records


# -------------------------- Individual Trace Plot -----------------------------
def plot_sac_records_with_obspy(records: List[Dict], outdir: Optional[str] = None, show: bool = False,
                                detrend: bool = True, taper: bool = True,
                                bandpass: Optional[Tuple[float, float]] = None,
                                time_window: Optional[Tuple[float, float]] = None,
                                dpi: int = 150) -> None:
    """
    Plot each SAC stream. If --outdir is given, save BOTH PNG and PDF per stream.
    """
    try:
        from obspy import read
    except ImportError:
        print("ERROR: ObsPy is not installed. Install with `pip install obspy` and try again.")
        return
    if outdir:
        os.makedirs(outdir, exist_ok=True)

    def safe(s: str) -> str:
        return "".join(c if c.isalnum() or c in "._-" else "_" for c in s)

    n_ok, n_err = 0, 0
    for idx, rec in enumerate(records, start=1):
        sac_path = rec.get("sac_path")
        station = rec.get("station", f"STA{idx}")
        phase = rec.get("phase", f"PHASE{idx}")
        if not sac_path or (not os.path.exists(sac_path)):
            print(f"Skip [{idx}]: Missing SAC file: {sac_path}")
            n_err += 1
            continue
        try:
            st = read(sac_path)
        except Exception as e:
            print(f"ERROR reading [{idx}] {sac_path}: {e}"); n_err += 1; continue
        try:
            for tr in st:
                if detrend: tr.detrend("linear")
                if taper: tr.taper(max_percentage=0.05, type="cosine")
                if bandpass:
                    fmin, fmax = bandpass
                    if fmin <= 0 or fmax <= fmin:
                        print(f"WARNING [{idx}]: invalid bandpass ({fmin}, {fmax}); skipping filter.")
                    else:
                        tr.filter("bandpass", freqmin=fmin, freqmax=fmax, corners=4, zerophase=True)
                if time_window:
                    tmin, tmax = time_window
                    start = tr.stats.starttime
                    tr.trim(start + tmin, start + tmax, pad=False)
        except Exception as e:
            print(f"ERROR preprocessing [{idx}] {sac_path}: {e}"); n_err += 1; continue
        try:
            outfile_png = outfile_pdf = None
            if outdir:
                base = f"{safe(station)}_{safe(phase)}"
                # ensure unique filename if collision
                candidate_png = f"{base}.png"
                candidate_pdf = f"{base}.pdf"
                if os.path.exists(os.path.join(outdir, candidate_png)) or \
                   os.path.exists(os.path.join(outdir, candidate_pdf)):
                    candidate_png = f"{base}_{idx}.png"
                    candidate_pdf = f"{base}_{idx}.pdf"
                outfile_png = os.path.join(outdir, candidate_png)
                outfile_pdf = os.path.join(outdir, candidate_pdf)

                # Save PNG
                st.plot(size=(800, 600), dpi=dpi, outfile=outfile_png, equal_scale=False)
                print(f"Saved waveform PNG: {outfile_png}")
                # Save PDF
                st.plot(size=(800, 600), dpi=dpi, outfile=outfile_pdf, equal_scale=False)
                print(f"Saved waveform PDF: {outfile_pdf}")

            if show and not outdir:
                st.plot(equal_scale=False)
        except Exception as e:
            print(f"ERROR plotting [{idx}] {sac_path}: {e}"); n_err += 1; continue
        n_ok += 1
    print(f"Plotting complete. Success: {n_ok}, Errors: {n_err}")


# ----------------- Common Axis Alignment (NumPy interpolation) ----------------
def build_common_time_stream(records: List[Dict], target_fs: Optional[float] = None,
                             align_mode: str = "union", pad_value: float = 0.0,
                             pick_trace: str = "first", verbose: bool = True,
                             anti_alias: bool = False, xy_units: str = "km",
                             coord_order: str = "xy", invert_x: bool = False, invert_y: bool = False):
    """
    Build an ObsPy Stream where each Trace is interpolated onto a shared time axis.
    Attach coordinates per ObsPy array_processing (coordsys='xy'): x=east (km), y=north (km).
    """
    try:
        from obspy import read, Stream
    except ImportError:
        print("ERROR: ObsPy is not installed. Install with `pip install obspy` and try again.")
        return None, None, None, None

    coord_scale = 1.0 if xy_units.lower() == "km" else 1e-3  # meters -> km
    traces, fs_list, starts, ends, labels = [], [], [], [], []

    # Read and collect metadata
    for idx, rec in enumerate(records, start=1):
        sac_path = rec.get("sac_path")
        station = rec.get("station", f"STA{idx}")
        phase = rec.get("phase", f"PHASE{idx}")
        sx_raw = rec.get("sx", 0.0) * coord_scale
        sy_raw = rec.get("sy", 0.0) * coord_scale

        # Apply order and sign controls here
        if coord_order == "xy":
            x = sx_raw; y = sy_raw
        else:  # 'yx' means list gives (north,east)
            x = sy_raw; y = sx_raw
        if invert_x: x = -x
        if invert_y: y = -y

        label = f"{station}:{phase}"
        if not sac_path or (not os.path.exists(sac_path)):
            if verbose: print(f"Skip [{idx}]: Missing SAC file: {sac_path}")
            continue
        try:
            st = read(sac_path)
        except Exception as e:
            if verbose: print(f"ERROR reading [{idx}] {sac_path}: {e}")
            continue
        if len(st) == 0:
            if verbose: print(f"WARNING [{idx}]: Empty stream from {sac_path}")
            continue

        tr = max(st, key=lambda t: t.stats.npts) if pick_trace == "maxlen" else st[0]
        fs = 1.0 / tr.stats.delta

        tr.stats.station = station or tr.stats.station
        tr.stats._phase = phase

        coord = AttribDict()
        coord.x = x  # km (east)
        coord.y = y  # km (north)
        coord.elevation = getattr(tr.stats, "elevation", 0.0)  # km
        tr.stats.coordinates = coord

        print ( label, tr.stats.starttime, tr.stats.endtime )
        traces.append(tr); fs_list.append(fs); starts.append(tr.stats.starttime); ends.append(tr.stats.endtime)
        labels.append(label)

    if not traces:
        if verbose: print("No readable SAC traces found in the provided records.")
        return Stream(), None, None, None

    # Global window
    if align_mode.lower() == "intersection":
        global_start = max(starts); global_end = min(ends)
        if global_end <= global_start:
            raise ValueError("Intersection window is empty—no common overlap among traces.")
    else:
        global_start = min(starts); global_end = max(ends)

    if target_fs is None:
        target_fs = max(fs_list)

    duration = float(global_end - global_start)
    npts = int(np.floor(duration * target_fs)) + 1
    print ('global_start = ', global_start )
    print ('global_end   = ', global_end   )
    print ('duration     = ', duration     )
    # exit()
    npts = max(npts, 2)

    # Common grid (seconds since epoch)
    t0_new = global_start.timestamp
    t_new = t0_new + np.arange(npts, dtype=float) / float(target_fs)

    aligned = Stream()
    for tr, label in zip(traces, labels):
        try:
            tr_work = tr.copy()
            native_fs = 1.0 / tr_work.stats.delta
            if anti_alias and target_fs < native_fs:
                cutoff = 0.45 * (target_fs / 2.0); cutoff = max(cutoff, 0.1)
                tr_work.filter("lowpass", freq=cutoff, corners=4, zerophase=True)
            t0_old = tr_work.stats.starttime.timestamp
            dt_old = tr_work.stats.delta
            t_old = t0_old + np.arange(tr_work.stats.npts, dtype=float) * dt_old
            y_old = tr_work.data.astype(float)
            y_new = np.interp(t_new, t_old, y_old, left=pad_value, right=pad_value)

            tr_work.data = y_new.astype(tr_work.data.dtype, copy=False)
            tr_work.stats.delta = 1.0 / target_fs
            tr_work.stats.sampling_rate = target_fs
            tr_work.stats.starttime = global_start
            tr_work.stats.npts = npts

            tr_work.stats._plot_label = label
            aligned.append(tr_work)
        except Exception as e:
            if verbose: print(f"ERROR aligning trace {label}: {e}")
            continue

    return aligned, global_start, target_fs, npts


# ------------------------------ Overlay Plot ----------------------------------
def plot_aligned_stream_overlay(aligned_stream, global_start, target_fs, npts,
                                outpath: Optional[str] = None, show: bool = False,
                                normalize: bool = False, dpi: int = 150,
                                title: Optional[str] = None):
    """
    Plot the aligned overlay. If --overlay-out is given, save BOTH PNG and PDF.
    """
    import matplotlib.pyplot as plt
    if aligned_stream is None or len(aligned_stream) == 0:
        print("Nothing to plot: aligned stream is empty."); return
    t = np.arange(npts, dtype=float) / float(target_fs)
    plt.figure(figsize=(10, 6))
    for tr in aligned_stream:
        y = tr.data.astype(float)
        label = getattr(tr.stats, "_plot_label", f"{tr.id}")
        if normalize:
            m = np.max(np.abs(y)) if y.size else 1.0
            if m > 0: y = y / m
        plt.plot(t, y, label=label, linewidth=1.0)
    plt.xlabel(f"Time (s) since {global_start.isoformat()} UTC")
    plt.ylabel("Amplitude" + (" (normalized)" if normalize else ""))
    plt.grid(True, alpha=0.3)
    if title: plt.title(title)
    plt.legend(loc="upper right", fontsize="small", ncol=2)
    plt.tight_layout()

    if outpath:
        outdir = os.path.dirname(os.path.abspath(outpath))
        os.makedirs(outdir, exist_ok=True)
        root, ext = os.path.splitext(outpath)
        png_path = outpath if ext.lower() == ".png" else root + ".png"
        pdf_path = root + ".pdf"
        plt.savefig(png_path, dpi=dpi)
        plt.savefig(pdf_path)  # vector export
        print(f"Saved overlay PNG: {png_path}")
        print(f"Saved overlay PDF: {pdf_path}")
    if show and not outpath:
        plt.show()
    plt.close()


# ------------------------- FK scan & strict store_map -------------------------
def run_fk_scan_cartesian(aligned_stream, global_start, global_end,
                          frqlow: float, frqhigh: float, win_len: float, win_frac: float,
                          sll_x: float, slm_x: float, sll_y: float, slm_y: float, sl_s: float,
                          prewhiten: int = 0, method: int = 0, verbose: bool = False):
    """
    Run ObsPy array_processing over XY slowness grid, accumulate relative power maps.
    """
    try:
        from obspy.signal.array_analysis import array_processing
        from obspy.core.utcdatetime import UTCDateTime
    except ImportError:
        print("ERROR: ObsPy is not installed. Install with `pip install obspy` and try again.")
        return None, None, None

    sx_axis = np.arange(sll_x, slm_x + sl_s * 0.5, sl_s)
    sy_axis = np.arange(sll_y, slm_y + sl_s * 0.5, sl_s)
    nx, ny = len(sx_axis), len(sy_axis)
    accum = np.zeros((ny, nx), dtype=float)

    # Print a sample of coordinates to verify geometry (x=east, y=north, km)
    print("Array coordinates sample (x=east km, y=north km):")
    for i, tr in enumerate(aligned_stream[:min(5, len(aligned_stream))]):
        c = tr.stats.coordinates
        print(f"  {i+1}: x={getattr(c, 'x', np.nan):.3f} km, y={getattr(c, 'y', np.nan):.3f} km")

    def store_map(relpow_map, t_ofs, iterno):
        nonlocal accum
        m = np.asarray(relpow_map, dtype=float)
        expected = (ny, nx)
        if m.shape != expected:
            if m.shape == (nx, ny):
                print("WARNING: relpow_map arrived as (nx, ny); transposing to (ny, nx).")
                m = m.T
            else:
                raise ValueError(f"Unexpected relpow_map shape {m.shape}; expected {expected}.")
        accum += m

    _ = array_processing(
        stream=aligned_stream,
        win_len=win_len, win_frac=win_frac,
        sll_x=sll_x, slm_x=slm_x, sll_y=sll_y, slm_y=slm_y,
        sl_s=sl_s,
        semb_thres=-1e9, vel_thres=-1e9,
        frqlow=frqlow, frqhigh=frqhigh,
        stime=UTCDateTime(global_start), etime=UTCDateTime(global_end),
        prewhiten=prewhiten, verbose=verbose,
        coordsys='xy', timestamp='julsec', method=method,
        store=store_map,
    )
    return accum, sx_axis, sy_axis


# ------------------------------- FK Plot --------------------------------------
def plot_fk_map(fk_map: np.ndarray, sx_axis: np.ndarray, sy_axis: np.ndarray,
                outpath: Optional[str] = None, dpi: int = 150, title: Optional[str] = None,
                normalize: bool = False):
    """
    Plot FK map (relative power) vs slowness components (sx east, sy north). Also show AZ & BAZ.
    """
    import matplotlib.pyplot as plt
    if fk_map is None or fk_map.size == 0:
        print("Nothing to plot: FK map is empty."); return

    iy, ix = np.unravel_index(np.argmax(fk_map), fk_map.shape)
    sx_peak = sx_axis[ix]; sy_peak = sy_axis[iy]
    s_mag = (sx_peak**2 + sy_peak**2)**0.5
    velocity = (1.0 / s_mag) if s_mag > 0 else np.inf

    # AZ measured clockwise from North to slowness vector; BAZ = AZ + 180°.
    az_deg  = (np.degrees(np.arctan2(sx_peak, sy_peak)) + 360.0) % 360.0
    baz_deg = (az_deg + 180.0) % 360.0
    print(f"FK peak at (sx={sx_peak:.3f}, sy={sy_peak:.3f}) -> AZ={az_deg:.1f}°, BAZ={baz_deg:.1f}°")

    data = fk_map; cbar_label = "Relative power (unitless)"
    if normalize:
        m = np.nanmax(fk_map); 
        if m > 0: data = fk_map / m; cbar_label = "Normalized power (0–1)"

    extent = (sx_axis.min(), sx_axis.max(), sy_axis.min(), sy_axis.max())
    plt.figure(figsize=(8, 7))
    im = plt.imshow(data, origin="lower", extent=extent, cmap="viridis", aspect="equal")
    plt.colorbar(im, label=cbar_label)
    plt.xlabel("Distance W-E $x$ (km)"); plt.ylabel("Distance S-N $y$ (km)")
    if title: plt.title(title)

    if outpath:
        root, ext = os.path.splitext(outpath)
    else:
        root      = " "

    plt.plot(sx_peak, sy_peak, marker="x", color="red", ms=8, mew=2)
    #plt.text(sx_peak, sy_peak,
    #         f"Peak:\nAZ={az_deg:.1f}°\nBAZ={baz_deg:.1f}°\ns={s_mag:.3f} s/km\nv≈{velocity:.2f} km/s",
    #         color="red", fontsize=9, ha="left", va="bottom")
    plt.text(sx_peak, sy_peak,
             f"Peak:\nAZ={az_deg:.1f}°\nDist={s_mag:.3f} km\n{root}\n",
             color="red", fontsize=9, ha="left", va="bottom")
    plt.grid(True, alpha=0.3); plt.tight_layout()

    if outpath:
        outdir = os.path.dirname(os.path.abspath(outpath)); os.makedirs(outdir, exist_ok=True)
        root, ext = os.path.splitext(outpath)
        png_path = outpath if ext.lower() == ".png" else root + ".png"
        pdf_path = root + ".pdf"
        plt.savefig(png_path, dpi=dpi); plt.savefig(pdf_path)
        print(f"Saved f–k PNG: {png_path}"); print(f"Saved f–k PDF: {pdf_path}")
        txt_path = root + ".txt"
        f = open( txt_path, "w" )
        line    = root + " "
        line   += "{:.4f}".format( sx_peak ).rjust(8) + " "
        line   += "{:.4f}".format( sy_peak ).rjust(8) + " "
        line   += "{:.1f}".format( az_deg ).rjust(6) + " "
        line   += "{:.4f}".format( s_mag ).rjust(8) + "\n"
        f.write( line )
        f.close()
    else:
        plt.show()
    plt.close()


# --------------------------- NetCDF writer for FK -----------------------------
def write_fk_to_netcdf(
    outpath: str,
    fk_map: np.ndarray,
    dx_axis: np.ndarray,
    dy_axis: np.ndarray,
    meta: dict,
):
    """
    Write FK grid to NetCDF.

    Dimensions:
      dy: distance south-to-north (km)
      dx: distance west-to-east (km)
    Variables:
      dy(dy), dx(dx), power(dy, dx)

    'meta' can include keys like:
      method, frqlow, frqhigh, win_len, win_frac, stime, etime, az_peak, baz_peak,
      s_mag_peak, v_peak, notes, coordsys
    """
    root, ext = os.path.splitext(outpath)
    # Try netCDF4 first
    try:
        from netCDF4 import Dataset  # type: ignore
        ds = Dataset(outpath, "w")
        try:
            ds.createDimension("dy", dy_axis.size)
            ds.createDimension("dx", dx_axis.size)
            v_dy = ds.createVariable("dy", "f8", ("dy",))
            v_dx = ds.createVariable("dx", "f8", ("dx",))
            v_dy.units = "km"; v_dy.standard_name = "dist_south_to_north_km"
            v_dx.units = "km"; v_dx.standard_name = "dist_west_to_east_km"
            v_dy[:] = dy_axis
            v_dx[:] = dx_axis

            v_pow = ds.createVariable("power", "f8", ("dy", "dx"))
            v_pow.units = "1"  # unitless relative power
            v_pow.long_name = "Accumulated relative power (ObsPy array_processing store)"
            v_pow[:, :] = fk_map  # (ny, nx): rows=dy, cols=dx

            ds.title = "Frequency–slowness (f–k) grid"
            ds.history = "Created by analyze_cc_traces.py"
            ds.coordinates = "dy dx"
            ds.coordsys = meta.get("coordsys", "xy (x=east, y=north)")
            for k, v in meta.items():
                try:
                    setattr(ds, k, v if isinstance(v, (str, int, float)) else str(v))
                except Exception:
                    pass
        finally:
            ds.close()
        print(f"Saved FK NetCDF: {outpath}")
        return
    except Exception as e:
        print(f"netCDF4 not available or failed ({e}). Trying scipy.io.netcdf...")

    # Fallback to scipy.io.netcdf (NetCDF-3 classic)
    try:
        from scipy.io import netcdf  # type: ignore
        with netcdf.netcdf_file(outpath, "w") as ds:
            ds.createDimension("dy", dy_axis.size)
            ds.createDimension("dx", dx_axis.size)
            v_dy = ds.createVariable("dy", "d", ("dy",))
            v_dx = ds.createVariable("dx", "d", ("dx",))
            v_dy[:] = dy_axis; v_dx[:] = dx_axis
            v_pow = ds.createVariable("power", "d", ("dy", "dx"))
            v_pow[:] = fk_map
            ds.title = "Frequency–slowness (f–k) grid"
            ds.coordsys = meta.get("coordsys", "xy (x=east, y=north)")
        print(f"Saved FK NetCDF (classic): {outpath}")
        return
    except Exception as e:
        print(f"Classic NetCDF write failed ({e}). Saving as NPZ instead...")

    # Last resort: NPZ
    npz_path = root + ".npz"
    np.savez_compressed(npz_path, dy=dy_axis, dx=dx_axis, power=fk_map, **meta)
    print(f"Saved fallback NPZ: {npz_path}")


# ---------------------------------- CLI ---------------------------------------
def main():
    parser = argparse.ArgumentParser(
        description="Analyze & plot SAC traces; perform FK using ObsPy array_processing.")
    parser.add_argument("listfile", help="ASCII file: sac_path station phase sx sy")
    parser.add_argument("--allow-missing-files", action="store_true")
    parser.add_argument("--allow-relative-paths", action="store_true")

    # Individual plots
    parser.add_argument("--plot", action="store_true")
    parser.add_argument("--outdir", default=None)
    parser.add_argument("--show", action="store_true")
    parser.add_argument("--detrend-off", action="store_true")
    parser.add_argument("--taper-off", action="store_true")
    parser.add_argument("--bandpass", nargs=2, type=float, metavar=("FMIN", "FMAX"))
    parser.add_argument("--time-window", nargs=2, type=float, metavar=("TMIN", "TMAX"))
    parser.add_argument("--dpi", type=int, default=150)

    # Alignment & overlay
    parser.add_argument("--align-common-axis", action="store_true")
    parser.add_argument("--target-fs", type=float, default=None)
    parser.add_argument("--intersection", action="store_true")
    parser.add_argument("--pad-value", type=float, default=0.0)
    parser.add_argument("--normalize", action="store_true")
    parser.add_argument("--overlay-out", default=None)
    parser.add_argument("--anti-alias", action="store_true")
    parser.add_argument("--xy-units", default="km", choices=["km", "m"])

    # Coordinate order & sign controls (critical for AZ orientation)
    parser.add_argument("--coord-order", default="xy", choices=["xy", "yx"],
                        help="Order of columns in list file: 'xy' = (sx=east, sy=north). 'yx' = (sx=north, sy=east).")
    parser.add_argument("--invert-x", action="store_true", help="Flip sign of x (east).")
    parser.add_argument("--invert-y", action="store_true", help="Flip sign of y (north).")

    # FK
    parser.add_argument("--fk", action="store_true")
    parser.add_argument("--fk-frqlow", type=float, default=1.0)
    parser.add_argument("--fk-frqhigh", type=float, default=3.0)
    parser.add_argument("--fk-win-len", type=float, default=1.0)
    parser.add_argument("--fk-win-frac", type=float, default=0.05)
    parser.add_argument("--fk-sll-x", type=float, default=-4.0)
    parser.add_argument("--fk-slm-x", type=float, default=+4.0)
    parser.add_argument("--fk-sll-y", type=float, default=-4.0)
    parser.add_argument("--fk-slm-y", type=float, default=+4.0)
    parser.add_argument("--fk-sl-s", type=float, default=0.04)
    parser.add_argument("--fk-prewhiten", type=int, default=0, choices=[0, 1])
    parser.add_argument("--fk-method", type=str, default="bf", choices=["bf", "capon"])
    parser.add_argument("--fk-out", default=None)
    parser.add_argument("--fk-normalize", action="store_true")
    parser.add_argument("--fk-nc", default=None,
                        help="Path to save the FK grid as NetCDF (coords: sx=east, sy=north).")

    args = parser.parse_args()

    # Read list
    try:
        records = read_station_phase_list(args.listfile, args.allow_missing_files, args.allow_relative_paths)
    except (FileNotFoundError, ValueError) as e:
        print(f"ERROR: {e}"); return

    print(f"Parsed {len(records)} records.")
    if records:
        for i, r in enumerate(records[:min(5, len(records))], start=1):
            print(f"{i}: station={r['station']} phase={r['phase']} sx={r['sx']} sy={r['sy']} sac_path={r['sac_path']}")

    # Individual plots (save PNG+PDF if outdir provided)
    if args.plot:
        plot_sac_records_with_obspy(records, args.outdir, args.show,
                                    detrend=(not args.detrend_off), taper=(not args.taper_off),
                                    bandpass=tuple(args.bandpass) if args.bandpass else None,
                                    time_window=tuple(args.time_window) if args.time_window else None,
                                    dpi=args.dpi)

    # Alignment (and needed for FK)
    if args.align_common_axis or args.fk:
        aligned_stream, global_start, target_fs, npts = build_common_time_stream(
            records, target_fs=args.target_fs,
            align_mode=("intersection" if args.intersection else "union"),
            pad_value=args.pad_value, pick_trace="first", verbose=True,
            anti_alias=args.anti_alias, xy_units=args.xy_units,
            coord_order=args.coord_order, invert_x=args.invert_x, invert_y=args.invert_y)
        if aligned_stream is None or len(aligned_stream) == 0:
            print("Aligned stream could not be built."); return

    # Overlay plot (save PNG+PDF if overlay_out provided)
    if args.align_common_axis:
        plot_aligned_stream_overlay(aligned_stream, global_start, target_fs, npts,
                                    outpath=args.overlay_out, show=bool(args.show and not args.plot),
                                    normalize=args.normalize, dpi=args.dpi,
                                    title="Aligned traces on common time axis")

    # FK
    if args.fk:
        global_end = aligned_stream[0].stats.endtime
        method_flag = 0 if args.fk_method == "bf" else 1
        fk_map, sx_axis, sy_axis = run_fk_scan_cartesian(
            aligned_stream, global_start, global_end,
            args.fk_frqlow, args.fk_frqhigh, args.fk_win_len, args.fk_win_frac,
            args.fk_sll_x, args.fk_slm_x, args.fk_sll_y, args.fk_slm_y, args.fk_sl_s,
            prewhiten=args.fk_prewhiten, method=method_flag, verbose=False)
        if fk_map is None:
            print("FK analysis failed."); return

        fk_map = fk_map.T

        # NetCDF export (optional)
        if args.fk_nc:
            iy, ix = np.unravel_index(np.argmax(fk_map), fk_map.shape)
            sx_peak = sx_axis[ix]
            sy_peak = sy_axis[iy]
            s_mag = (sx_peak**2 + sy_peak**2)**0.5
            v_peak = (1.0 / s_mag) if s_mag > 0 else np.inf
            az_peak  = (np.degrees(np.arctan2(sx_peak, sy_peak)) + 360.0) % 360.0  # AZ
            baz_peak = (az_peak + 180.0) % 360.0
            meta = {
                "method": args.fk_method,
                "frqlow": float(args.fk_frqlow),
                "frqhigh": float(args.fk_frqhigh),
                "win_len": float(args.fk_win_len),
                "win_frac": float(args.fk_win_frac),
                "stime": str(global_start),
                "etime": str(global_end),
                "az_peak_deg": float(az_peak),
                "baz_peak_deg": float(baz_peak),
                "s_mag_peak_s_per_km": float(s_mag),
                "v_peak_km_per_s": float(v_peak),
                "coordsys": "xy (x=east, y=north) for slowness grid",
                "note": "power(sy, sx) with sy=northing and sx=easting; unitless relative power accumulated across windows.",
            }
            write_fk_to_netcdf(args.fk_nc, fk_map, sx_axis, sy_axis, meta)

        # f–k plot (save PNG+PDF if fk_out provided)
        plot_fk_map(fk_map, sx_axis, sy_axis, outpath=args.fk_out, dpi=args.dpi,
                    title=f"Broadband f–k scan ({args.fk_frqlow}-{args.fk_frqhigh} Hz)",
                    normalize=args.fk_normalize)


if __name__ == "__main__":
    main()