#!/bin/bash
set -e

pip install --no-cache-dir seisbench==0.10.2

python3 << 'EOF'
import os
from pathlib import Path
from glob import glob

import numpy as np
import obspy
import pandas as pd
import seisbench.models as sbm
from obspy import Stream, Trace, UTCDateTime


def load_npz_inputs(npz_dir: str | Path) -> dict[str, obspy.Stream]:
    npz_dir = Path(npz_dir)
    npz_files = sorted(glob(str(npz_dir / "*.npz")))

    required_fields = ["data", "dt", "start_time", "network", "station", "channels"]

    streams = {}

    for npz_file in npz_files:
        filename = os.path.basename(npz_file)
        data = np.load(npz_file, allow_pickle=False)

        missing_fields = [field for field in required_fields if field not in data]
        if missing_fields:
            print(f"Skipping {filename}: missing required fields {missing_fields}")
            continue

        waveform = data["data"] * 1e10  # (12000, 3)
        dt = float(data["dt"])
        sampling_rate = 1.0 / dt

        network = str(data["network"])
        station = str(data["station"])
        channels = str(data["channels"])
        start_time_str = str(data["start_time"])

        if "," in channels:
            channel_list = channels.split(",")
        else:
            channel_list = [f"HH{c}" for c in "ENZ"]

        try:
            start_time = UTCDateTime(start_time_str)
        except Exception:
            start_time = UTCDateTime(0)

        stream = Stream()
        for i, ch in enumerate(channel_list):
            if i < waveform.shape[1]:
                tr = Trace(data=waveform[:, i].astype(np.float64))
                tr.stats.network = network
                tr.stats.station = station
                tr.stats.channel = ch if len(ch) >= 2 else f"HH{ch}"
                tr.stats.sampling_rate = sampling_rate
                tr.stats.starttime = start_time
                stream.append(tr)

        streams[filename] = stream
        print(f"Loaded {filename}: {len(stream)} traces")

    return streams


def load_denoiser(device: str = None) -> sbm.DeepDenoiser:
    denoiser = sbm.DeepDenoiser.from_pretrained("original")

    if device:
        denoiser.to(device)
    else:
        denoiser.to_preferred_device()

    return denoiser


def load_models(model_names: list[str] = None, device: str = None) -> dict:
    if model_names is None:
        model_names = ["phasenet", "eqtransformer", "gpd"]

    models = {}

    for name in model_names:
        name_lower = name.lower()
        if name_lower == "phasenet":
            model = sbm.PhaseNet.from_pretrained("original")
        elif name_lower == "eqtransformer":
            model = sbm.EQTransformer.from_pretrained("original")
        elif name_lower == "gpd":
            model = sbm.GPD.from_pretrained("stead")
        else:
            print(f"Unknown model: {name}, skipping")
            continue

        if device:
            model.to(device)
        else:
            model.to_preferred_device()

        models[name_lower] = model

    return models


def run_inference(
    streams: dict[str, obspy.Stream],
    models: dict,
    output_csv: str | Path = None,
    denoiser: sbm.DeepDenoiser = None,
) -> dict:
    all_results = {}

    for model_name, model in models.items():
        print(f"\n{'='*60}")
        print(f"Running inference with {model_name}")
        if denoiser:
            print("(with DeepDenoiser preprocessing)")
        print(f"{'='*60}")

        all_picks = []

        for filename, stream in streams.items():
            sampling_rate = stream[0].stats.sampling_rate

            try:
                if denoiser:
                    input_stream = denoiser.annotate(stream)
                else:
                    input_stream = stream

                outputs = model.classify(input_stream)
                picks = outputs.picks

                start_time = stream[0].stats.starttime

                for pick in picks:
                    pick_idx = int((pick.peak_time - start_time) * sampling_rate)

                    all_picks.append({
                        "file_name": filename,
                        "phase": pick.phase,
                        "pick_idx": pick_idx,
                        "probability": pick.peak_value,
                    })

                file_p_picks = [p for p in all_picks if p["file_name"] == filename and p["phase"] == "P"]
                file_s_picks = [p for p in all_picks if p["file_name"] == filename and p["phase"] == "S"]

                print(f"  {filename}: P picks: {len(file_p_picks)}, S picks: {len(file_s_picks)}")

            except Exception as e:
                print(f"  Error processing {filename}: {e}")
                continue

        all_results[model_name] = all_picks

        # one row per pick
        if output_csv and all_picks:
            picks_df = pd.DataFrame(all_picks)
            picks_df.to_csv(output_csv, index=False)
            print(f"\nSaved {len(picks_df)} picks to {output_csv}")

    return all_results


def main():
    npz_inputs_dir = Path("/root/data")
    output_csv = Path("/root/results.csv")

    print("="*60)
    print("SeisBench Phase Picking Inference")
    print("="*60)

    print("\n1. Loading NPZ input data...")
    streams = load_npz_inputs(npz_inputs_dir)
    print(f"Loaded {len(streams)} samples")

    if not streams:
        print("No data loaded. Exiting.")
        return None

    print("\n2. Loading SeisBench models...")
    models = load_models(["phasenet"])
    denoiser = load_denoiser()

    print("\n3. Running inference...")
    results = run_inference(streams, models, output_csv, denoiser=None)

    print("\n" + "="*60)
    print(f"Results saved to: {output_csv}")

    return results

main()

EOF