#!/usr/bin/env python3
"""
Robust PSD Estimation Tool.

This script estimates the Power Spectral Density (PSD) of gravitational wave data.
It generalizes legacy functionality by leveraging the SGN framework for:
1. Robust data fetching (GWOSC).
2. Modern PSD estimators (MGM, Recursive).
3. Flexible output formatting via PSDSink.
"""

from __future__ import annotations

import argparse
import sys
from typing import List, Optional

# SGN & Family
from sgn.apps import Pipeline
from sgn.base import SourceElement
from sgndrift.transforms.psd import MGMPSD, RecursivePSD
from sgnligo.sinks.psd_sink import PSDSink
from sgnligo.sources.gwosc import GWOSCSource


def build_pipeline(
    start: float,
    end: float,
    detectors: List[str],
    sample_rate: int,
    fft_length: float,
    overlap: float,
    estimator_type: str,
    output_fmt: str,
    write_interval: Optional[float] = None,
    source_element: Optional[SourceElement] = None,
    verbose: bool = False,
) -> Pipeline:
    """Constructs the processing pipeline for the requested interval."""
    pipe = Pipeline()

    # 1. Source
    if source_element:
        source = source_element
        source_name = source.name
    else:
        source_name = f"gwosc_{int(start)}"
        source = GWOSCSource(
            name=source_name,
            start=start,
            end=end,
            detectors=detectors,
            sample_rate=sample_rate,
            cache_data=True,
            verbose=verbose,
        )
    pipe.insert(source)

    # 2. Per-Detector Processing
    for ifo in detectors:
        # A. Estimator
        est_name = f"psd_{ifo}"
        if estimator_type.lower() == "recursive":
            est = RecursivePSD(
                name=est_name,
                sink_pad_names=("in",),
                source_pad_names=("out",),
                fft_length=fft_length,
                overlap=overlap,
                sample_rate=int(sample_rate),
            )
        else:
            # Default to MGM
            est = MGMPSD(
                name=est_name,
                sink_pad_names=("in",),
                source_pad_names=("out",),
                fft_length=fft_length,
                overlap=overlap,
                sample_rate=int(sample_rate),
            )

        # B. Sink
        # Format filename: defaults to "psd_{detector}_{start}.xml.gz"
        fname = output_fmt.format(
            detector=ifo,
            start="{gps}",
            duration=int(end - start),
            rate=int(sample_rate),
        )

        # PSDSink usually expects the IFO name as the sink pad
        # to correctly label the XML element
        sink = PSDSink(
            name=f"sink_{ifo}",
            sink_pad_names=(ifo,),
            fname=fname,
            write_interval=write_interval,
            verbose=verbose,
        )

        pipe.insert(est, sink)

        # C. Link
        # Source[IFO] -> Est[in]
        # Est[out] -> Sink[IFO]

        # Check source pads
        if ifo in source.srcs:
            src_pad = source.srcs[ifo]
        else:
            if verbose:
                print(
                    f"Warning: Source pad '{ifo}' not found in source '{source.name}'. Skipping {ifo}."
                )
            continue

        pipe.link({est.snks["in"]: src_pad, sink.snks[ifo]: est.srcs["out"]})

    return pipe


def estimate_psd(
    start: float,
    end: float,
    detectors: List[str],
    output_fmt: str = "psd_{detector}_{start}.xml.gz",
    estimator: str = "mgm",
    fft_length: float = 4.0,
    overlap: float = 0.5,
    sample_rate: float = 4096.0,
    write_interval: Optional[float] = None,
    verbose: bool = False,
):
    """
    Main driver for PSD estimation.
    Runs the pipeline over the user-specified [start, end) interval.
    """
    duration = end - start
    if duration < (fft_length * 4):
        print(f"Error: Duration ({duration}s) is too short for stable PSD estimation.")
        return

    if verbose:
        print(f"Processing interval: {start} - {end} ({duration}s)")
        print(f"Detectors: {detectors}")
        print(f"Estimator: {estimator}")
        if write_interval:
            print(f"Write Interval: {write_interval}s")

    try:
        pipe = build_pipeline(
            start=start,
            end=end,
            detectors=detectors,
            sample_rate=int(sample_rate),
            fft_length=fft_length,
            overlap=overlap,
            estimator_type=estimator,
            output_fmt=output_fmt,
            write_interval=write_interval,
            verbose=verbose,
        )
        pipe.run()
        if verbose:
            print("Pipeline finished successfully.")

    except Exception as e:
        print(f"Pipeline failed: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)


def parse_args():
    parser = argparse.ArgumentParser(description="Estimate GW PSDs.")
    parser.add_argument(
        "--detectors", type=str, required=True, help="Comma-separated list (e.g. H1,L1)"
    )
    parser.add_argument("--start", type=float, required=True, help="GPS Start Time")
    parser.add_argument("--end", type=float, required=True, help="GPS End Time")

    parser.add_argument(
        "--output-fmt",
        type=str,
        default="psd_{detector}_{start}.xml.gz",
        help="Output filename format. Available keys: {detector}, {start}, {duration}, {rate}",
    )

    parser.add_argument(
        "--estimator",
        type=str,
        default="mgm",
        choices=["mgm", "recursive"],
        help="PSD Estimation algorithm",
    )
    parser.add_argument(
        "--fft-length", type=float, default=4.0, help="FFT length in seconds"
    )
    parser.add_argument(
        "--overlap", type=float, default=0.5, help="Overlap fraction (0-1)"
    )
    parser.add_argument(
        "--sample-rate", type=float, default=4096.0, help="Sampling rate in Hz"
    )
    parser.add_argument(
        "--write-interval",
        type=float,
        default=None,
        help="Interval in seconds to write intermediate PSDs to disk.",
    )

    parser.add_argument(
        "-v", "--verbose", action="store_true", help="Enable verbose logging"
    )

    return parser.parse_args()


def main():
    args = parse_args()
    dets = [d.strip() for d in args.detectors.split(",")]

    estimate_psd(
        start=args.start,
        end=args.end,
        detectors=dets,
        output_fmt=args.output_fmt,
        estimator=args.estimator,
        fft_length=args.fft_length,
        overlap=args.overlap,
        sample_rate=args.sample_rate,
        write_interval=args.write_interval,
        verbose=args.verbose,
    )


if __name__ == "__main__":
    main()
