#!/usr/bin/env python3
"""
PSD Visualization Tool.

Generates comparison plots of Power Spectral Densities (PSDs).

Features:
- Supports multiple input files (LIGO LW XML, HDF5, TXT).
- Auto-detects Time Evolution (Drift) using internal metadata.
- Smart Labeling (Relative Time).
- Sequential Colormaps for Drift Plots.
"""

import argparse
import sys
import os
import re
from typing import List, Optional, Any

import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
from gwpy.frequencyseries import FrequencySeries


def load_psd(filepath: str, ifo: Optional[str] = None) -> FrequencySeries:
    """Loads a PSD FrequencySeries from a file using GWpy."""
    if not os.path.exists(filepath):
        print(f"Error: File not found at {filepath}")
        sys.exit(1)

    try:
        # For XML files, explicit IFO helps GWpy disambiguate
        if ifo:
            return FrequencySeries.read(filepath, instrument=ifo)
        else:
            return FrequencySeries.read(filepath)
    except Exception as e:
        print(f"Error reading {filepath}: {e}")
        print(f"Tip: If {filepath} is an XML file, try providing --ifos")
        sys.exit(1)


def make_plot(
    file_list: List[str],
    output_path: str,
    ifos: Optional[List[str]] = None,
    labels: Optional[List[str]] = None,
    title: Optional[str] = None,
    xlim: Optional[List[float]] = None,
    ylim: Optional[List[float]] = None,
    asd_mode: bool = False,
    sort_time: bool = True,
):
    """Generates and saves the PSD comparison plot."""

    # --- 1. Load Data & Extract Metadata ---
    loaded_items = []

    for i, fpath in enumerate(file_list):
        # IFO hint from CLI args
        ifo_hint = ifos[i] if ifos and i < len(ifos) else None

        # Load Series
        series = load_psd(fpath, ifo_hint)

        # Extract GPS Time from Object Attributes
        gps_time = None
        if series.epoch is not None:
            # gwpy epoch is usually an astropy Time object
            try:
                gps_time = series.epoch.gps
            except AttributeError:
                # Fallback if it's just a number
                gps_time = float(series.epoch)

        # Extract IFO Name (Prefer Channel Name -> Filename -> Hint)
        detected_ifo = None
        if series.channel:
            # Typical format: "H1:DMT-PSD..."
            cname = str(series.channel)
            if ":" in cname:
                detected_ifo = cname.split(":")[0]

        if not detected_ifo:
            # Fallback regex on filename
            match = re.search(r"([HLV]1)", os.path.basename(fpath))
            if match:
                detected_ifo = match.group(1)

        # Final IFO resolution
        final_ifo = ifo_hint if ifo_hint else detected_ifo

        loaded_items.append(
            {
                "series": series,
                "gps": gps_time,
                "ifo": final_ifo,
                "path": fpath,
                "explicit_label": labels[i] if labels and i < len(labels) else None,
            }
        )

    # --- 2. Sort by Time ---
    if sort_time:
        # Sort by GPS, handling None (put at start or end as preferred)
        loaded_items.sort(key=lambda x: x["gps"] if x["gps"] is not None else -1.0)

    # --- 3. Determine Plot Mode (Drift vs Comparison) ---
    unique_ifos = set(item["ifo"] for item in loaded_items if item["ifo"])
    # If strictly one IFO is found across multiple files, assume Drift Mode
    is_drift_mode = (len(unique_ifos) == 1) and (len(loaded_items) > 1)

    # --- 4. Setup Plotting ---
    plt.figure(figsize=(12, 8))

    # Colors
    num_plots = len(loaded_items)
    if is_drift_mode:
        # Sequential (Viridis) for time evolution
        colors = cm.viridis(np.linspace(0, 1, num_plots))
    else:
        # Qualitative (Tab10) for distinct items
        colors = cm.tab10(np.linspace(0, 1, min(num_plots, 10)))

    # Base Time (Start of sequence)
    t0 = loaded_items[0]["gps"] if loaded_items and loaded_items[0]["gps"] else None

    for i, item in enumerate(loaded_items):
        series = item["series"]

        # ASD Transform
        if asd_mode:
            # GWpy propagates units automatically
            series = series**0.5
            series.name = "Amplitude Spectral Density"

        # Construct Label
        if item["explicit_label"]:
            label = item["explicit_label"]
        elif is_drift_mode and item["gps"] is not None and t0 is not None:
            # Smart Relative Label: "H1 (+16s)"
            dt = item["gps"] - t0
            ifo_str = item["ifo"] if item["ifo"] else "PSD"
            label = f"{ifo_str} (T+{dt:.1f}s)"
        else:
            # Standard Label: "H1 (Channel)" or Filename
            base = item["ifo"] if item["ifo"] else os.path.basename(item["path"])
            # Only append GPS to label if we aren't in smart drift mode
            if item["gps"]:
                base += f" @ {int(item['gps'])}"
            label = base

        # Plot
        plt.loglog(
            series.frequencies.value,
            series.value,
            label=label,
            alpha=0.8,
            linewidth=1.5,
            color=colors[i % len(colors)],
        )

    # --- 5. Styling ---
    plt.xlabel("Frequency [Hz]", fontsize=12)

    if asd_mode:
        plt.ylabel(r"ASD [$\mathrm{Strain} / \sqrt{\mathrm{Hz}}$]", fontsize=12)
    else:
        plt.ylabel(r"PSD [$\mathrm{Strain}^2 / \mathrm{Hz}$]", fontsize=12)

    plt.grid(True, which="both", alpha=0.2, ls="-")

    if xlim:
        plt.xlim(xlim)
    else:
        plt.xlim(10, 2048)

    if ylim:
        plt.ylim(ylim)

    if title:
        plt.title(title, fontsize=14)
    else:
        mode = "Drift Evolution" if is_drift_mode else "Noise Comparison"
        if t0 and is_drift_mode:
            mode += f" (Start GPS: {int(t0)})"
        plt.title(mode, fontsize=14)

    # Legend
    legend = plt.legend(loc="upper right", framealpha=0.9, fontsize=10)
    if is_drift_mode:
        legend.set_title("Time Evolution")

    try:
        plt.savefig(output_path, dpi=150, bbox_inches="tight")
        print(f"Plot saved to: {output_path}")
    except Exception as e:
        print(f"Error saving plot: {e}")


def parse_args():
    parser = argparse.ArgumentParser(description="Visualize and Compare GW PSDs.")
    parser.add_argument(
        "input_files", nargs="+", type=str, help="Path(s) to input PSD file(s)"
    )
    parser.add_argument(
        "-o", "--output", type=str, default="psd_plot.png", help="Output filename"
    )

    # Plotting Options
    parser.add_argument(
        "--labels",
        type=str,
        default=None,
        help="Comma-separated labels (overrides auto-detection)",
    )
    parser.add_argument(
        "--ifos",
        type=str,
        default=None,
        help="Comma-separated IFO names (hint for XML parsing)",
    )
    parser.add_argument(
        "--asd", action="store_true", help="Plot Amplitude Spectral Density (sqrt(PSD))"
    )
    parser.add_argument(
        "--no-sort", action="store_true", help="Disable automatic sorting by GPS time"
    )

    # Ranges
    parser.add_argument("--fmin", type=float, default=10.0, help="Min Frequency")
    parser.add_argument("--fmax", type=float, default=2048.0, help="Max Frequency")
    parser.add_argument("--ymin", type=float, default=None, help="Min Y-axis")
    parser.add_argument("--ymax", type=float, default=None, help="Max Y-axis")

    parser.add_argument("--title", type=str, default=None, help="Custom plot title")

    return parser.parse_args()


def main():
    args = parse_args()

    # Process comma-separated lists
    labels = [l.strip() for l in args.labels.split(",")] if args.labels else None
    ifos = [i.strip() for i in args.ifos.split(",")] if args.ifos else None

    # Handle shell globbing (Python glob if shell didn't expand)
    expanded_files = []
    import glob

    for f in args.input_files:
        if "*" in f:
            expanded_files.extend(glob.glob(f))
        else:
            expanded_files.append(f)

    # Remove duplicates and sort for consistency
    expanded_files = sorted(list(set(expanded_files)))

    if not expanded_files:
        print("No input files found.")
        sys.exit(1)

    make_plot(
        expanded_files,
        output_path=args.output,
        ifos=ifos,
        labels=labels,
        title=args.title,
        xlim=[args.fmin, args.fmax],
        ylim=[args.ymin, args.ymax] if args.ymin and args.ymax else None,
        asd_mode=args.asd,
        sort_time=not args.no_sort,
    )


if __name__ == "__main__":
    main()
