"""Prepare trimmed test data for offline inspiral integration test.

This script trims the original large GWF frame files down to a 512s window
and creates the necessary auxiliary files (segments) for the test.

The chosen window [1241740448, 1241740960] includes:
- All 3 IFOs (H1, L1, V1) with valid frame data (no real gaps)
- An artificial L1 segment gap at [1241740700, 1241740708] (8s, centered)
- 34 reference events with network_snr >= 7

The output goes into sgnl-test-data/src/sgnl_test_data/data/offline/.
SVD banks are reused from sgnl-test-data/online/svd_bank/ (identical files).
The frame cache is generated dynamically at test time since it contains
absolute paths.

Usage:
    python prepare_test_data_offline.py \
        --frame-dir /path/to/original/frames \
        --output-dir /path/to/sgnl-test-data/src/sgnl_test_data/data/offline \
        --src-dir /path/to/offline-test  # for PSD and config

After running this script, generate the reference triggers by running
sgnl-inspiral on the trimmed data and save as reference_triggers.sqlite
in the output directory.
"""

import argparse
import os
import shutil

from gwpy.timeseries import TimeSeries
from ligo.lw import ligolw, lsctables
from ligo.lw import utils as ligolw_utils
from ligo.lw.utils import process as ligolw_process

GPS_START = 1241740448
GPS_END = 1241740960
DUR = GPS_END - GPS_START
MID = GPS_START + DUR // 2
GAP_HALF = 4

FRAME_SPECS = [
    {
        "site": "H",
        "type": "H1_HOFT_C01",
        "channel": "H1:DCS-CALIB_STRAIN_CLEAN_C01",
        "src_pattern": "H-H1_HOFT_C01-1241739264-4096.gwf",
    },
    {
        "site": "L",
        "type": "L1_HOFT_C01",
        "channel": "L1:DCS-CALIB_STRAIN_CLEAN_C01",
        "src_pattern": "L-L1_HOFT_C01-1241739264-4096.gwf",
    },
    {
        "site": "V",
        "type": "V1Online",
        "channel": "V1:Hrec_hoft_16384Hz",
        "src_pattern": "V-V1Online-1241740000-2000.gwf",
    },
]

IFO_SEGMENTS = {
    "H1": [(GPS_START, GPS_END)],
    "L1": [(GPS_START, MID - GAP_HALF), (MID + GAP_HALF, GPS_END)],
    "V1": [(GPS_START, GPS_END)],
}

COPY_FILES = [
    "H1L1V1-SGNL_REFERENCE_PSD-1241738590-3835.xml.gz",
    "fake_cbc.yaml",
]


def trim_frames(frame_dir, output_dir):
    """Trim GWF frame files to the 128s window."""
    for spec in FRAME_SPECS:
        src = os.path.join(frame_dir, spec["src_pattern"])
        out_name = f"{spec['site']}-{spec['type']}-{GPS_START}-{DUR}.gwf"
        out_path = os.path.join(output_dir, out_name)

        print(f"Trimming {spec['channel']} -> {out_name}")
        ts = TimeSeries.read(src, spec["channel"], start=GPS_START, end=GPS_END)
        ts.write(out_path)
        size_mb = os.path.getsize(out_path) / 1024 / 1024
        print(f"  {ts.shape[0]} samples, {size_mb:.1f} MB")


def create_segments(output_dir):
    """Create segments.xml.gz with the L1 gap."""
    xmldoc = ligolw.Document()
    xmldoc.appendChild(ligolw.LIGO_LW())

    proc_table = lsctables.New(lsctables.ProcessTable)
    xmldoc.childNodes[0].appendChild(proc_table)
    proc = ligolw_process.register_to_xmldoc(xmldoc, "prepare_test_data_offline", {})

    seg_def_table = lsctables.New(lsctables.SegmentDefTable)
    xmldoc.childNodes[0].appendChild(seg_def_table)

    seg_table = lsctables.New(lsctables.SegmentTable)
    xmldoc.childNodes[0].appendChild(seg_table)

    seg_id = 0
    for def_id, (ifo, segs) in enumerate(IFO_SEGMENTS.items()):
        seg_def = seg_def_table.RowType()
        seg_def.segment_def_id = def_id
        seg_def.process_id = proc.process_id
        seg_def.ifos = ifo
        seg_def.name = "datasegments"
        seg_def.version = 1
        seg_def.comment = ""
        seg_def.creator_db = None
        seg_def.insertion_time = None
        seg_def_table.append(seg_def)

        for start, end in segs:
            seg_row = seg_table.RowType()
            seg_row.segment_id = seg_id
            seg_row.segment_def_id = def_id
            seg_row.segment_def_cdb = None
            seg_row.process_id = proc.process_id
            seg_row.creator_db = None
            seg_row.start_time = start
            seg_row.start_time_ns = 0
            seg_row.end_time = end
            seg_row.end_time_ns = 0
            seg_table.append(seg_row)
            seg_id += 1

    path = os.path.join(output_dir, "segments.xml.gz")
    ligolw_utils.write_filename(xmldoc, path, compress="gz")
    print(f"Created {path}")
    for ifo, segs in IFO_SEGMENTS.items():
        for s, e in segs:
            print(f"  {ifo}: [{s}, {e}] ({e - s}s)")


def copy_files(src_dir, output_dir):
    """Copy reference PSD and event config."""
    for fname in COPY_FILES:
        src = os.path.join(src_dir, fname)
        dst = os.path.join(output_dir, fname)
        if not os.path.exists(src):
            print(f"WARNING: {src} not found, skipping")
            continue
        shutil.copy2(src, dst)
        size_mb = os.path.getsize(dst) / 1024 / 1024
        print(f"Copied {fname} ({size_mb:.1f} MB)")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--frame-dir",
        required=True,
        help="Directory containing original large GWF frame files",
    )
    parser.add_argument(
        "--output-dir",
        required=True,
        help="Output directory (sgnl-test-data/src/sgnl_test_data/data/offline)",
    )
    parser.add_argument(
        "--src-dir",
        required=True,
        help="Source directory with PSD and config",
    )
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    print("=== Trimming frames ===")
    trim_frames(args.frame_dir, args.output_dir)

    print("\n=== Creating segments ===")
    create_segments(args.output_dir)

    print("\n=== Copying auxiliary files ===")
    copy_files(args.src_dir, args.output_dir)

    print("\n=== Done ===")
    print(
        f"\nTo generate reference triggers, run sgnl-inspiral on the trimmed"
        f" data and save the output as reference_triggers.sqlite in"
        f" {args.output_dir}/"
    )
    print(
        "\nNote: SVD banks are reused from sgnl-test-data/online/svd_bank/"
        " (GSTLAL naming, identical content to SGNL banks)."
    )


if __name__ == "__main__":
    main()
