#!/usr/bin/env python3
"""
hypoxic_burden
==============

Python port of ``HypoxicBurden.m``.

Calculates the hypoxic burden of an SpO2 signal. 

INPUT
-----
spo2       : 1-D numeric array / sequence containing the SpO2 signal sampled at
             1 Hz. May contain NaN to mark artefacts.
subject_id : (str) currently analysed subject ID.
wave       : (str) assessment wave.
threshold  : (float) threshold defining a desaturation event. Default: 2 (%).

OUTPUT
------
A dict (1 row) with the following fields, mirroring the MATLAB output table:
    ID              : subject ID
    Wave            : wave ID
    Threshold       : used threshold
    AREA            : inverted event burden area, i.e. area between the
                      event baseline (100%) and the desaturation curve
                      (percentage * time in seconds)
    TST             : assessed total sleep time minus artefacts (seconds)
    TotalAREA       : total raw SpO2 area under the full non-NaN signal
                      (diagnostic; not used for HB)
    HB              : hypoxic burden (AREA / TST)
    numEvents       : number of desaturation events
    events_per_hour : number of desaturation events per hour

REFERENCES
----------
Code adapted from Esmaeili N et al., Hypoxic Burden Based on Automatically
Identified Desaturations Is Associated with Adverse Health Outcomes.
Ann Am Thorac Soc. 2023 Nov;20(11):1633-1641.
doi: 10.1513/AnnalsATS.202303-248OC. PMID: 37531573; PMCID: PMC10632930.

Original MATLAB author: Antoine Weihs <antoine.weihs@uni-greifswald.de>
Python porting : June Christoph Kang <cnsla@korea.ac.kr>
License: MIT.
"""

from __future__ import annotations

import argparse
import warnings

import numpy as np


# --------------------------------------------------------------------------- #
# Helpers reproducing MATLAB built-ins
# --------------------------------------------------------------------------- #
def _trapz(y):
    """MATLAB trapz with unit spacing. trapz of <2 samples is 0."""
    y = np.asarray(y, dtype=float)
    if y.size < 2:
        return 0.0
    return float(np.trapz(y))


def _round_half_away_from_zero(x):
    """MATLAB round(..., TieBreaker='fromzero'): round half away from zero."""
    return float(np.sign(x) * np.floor(np.abs(x) + 0.5))


def _findpeaks(x):
    """
    Reproduce MATLAB ``findpeaks`` (locations only), returning the 1-based
    positions of local maxima.

    A peak is a sample strictly greater than its two neighbours. Endpoints are
    never peaks. Flat-topped peaks (plateaus) return the centre of the plateau
    (the lower of the two central indices when the plateau has even length),
    matching MATLAB's behaviour. In practice ENS_AVG is an average of many real
    SpO2 curves, so exact plateaus do not occur; this branch is for parity only.
    """
    x = np.asarray(x, dtype=float)
    n = x.size
    locs = []
    i = 1
    while i < n - 1:
        if x[i] > x[i - 1]:
            # rising into i: extend across any flat plateau
            j = i
            while j < n - 1 and x[j + 1] == x[j]:
                j += 1
            if j < n - 1 and x[j + 1] < x[j]:
                # falling after the plateau -> it is a peak
                pk = (i + j) // 2  # plateau centre (lower-middle if even)
                locs.append(pk + 1)  # 1-based, like MATLAB
            i = j + 1
        else:
            i += 1
    return np.asarray(locs, dtype=int)


# --------------------------------------------------------------------------- #
# Main function
# --------------------------------------------------------------------------- #
def hypoxic_burden(spo2, subject_id="", wave="", threshold=2.0):
    """
    Compute the hypoxic burden of an SpO2 signal (1 Hz).

    Parameters
    ----------
    spo2 : array_like
        1-D SpO2 signal sampled at 1 Hz (may contain NaN for artefacts).
    subject_id : str
        Subject ID.
    wave : str
        Assessment wave.
    threshold : float
        Desaturation threshold in %. Must be positive. Default 2.

    Returns
    -------
    dict
        See module docstring for the field description.
    """
    spo2 = np.asarray(spo2, dtype=float).ravel()
    if not np.isreal(spo2).all():
        raise ValueError("SpO2 must be real.")
    if not (np.isscalar(threshold) or np.ndim(threshold) == 0) or threshold <= 0:
        raise ValueError("threshold must be a positive scalar.")
    threshold = float(threshold)

    n = spo2.size

    # ------------------------------------------------------------------ #
    # 1) Detect desaturation events
    #    Slide a 10-sample window; if its minimum drops more than
    #    `threshold` below the window's first sample, record the location
    #    of that minimum and jump 10 samples ahead.
    #    (Events are stored as MATLAB-style 1-based indices.)
    # ------------------------------------------------------------------ #
    desat_events = []  # 1-based indices, exactly as the MATLAB code stores them
    l_index = 1  # MATLAB 1-based loop counter
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)  # all-NaN slices
        while l_index < n - 10:
            temp = spo2[l_index - 1: l_index - 1 + 10]  # SpO2(l_index : l_index+9)
            mn = np.nanmin(temp) if np.any(~np.isnan(temp)) else np.nan
            # NaN comparisons are False, matching MATLAB (min ignores NaN).
            if mn < spo2[l_index - 1] - threshold:
                # first position (1-based) of the minimum inside the window
                p1 = int(np.argmax(temp == mn)) + 1
                desat_events.append(l_index + p1)
                l_index += 10
            else:
                l_index += 1

    # ------------------------------------------------------------------ #
    # 2) Total area under the curve (non-NaN samples)
    # ------------------------------------------------------------------ #
    total_area = _trapz(spo2[~np.isnan(spo2)])

    tst = int(np.count_nonzero(~np.isnan(spo2)))  # seconds

    if len(desat_events) == 0:
        return {
            "ID": subject_id,
            "Wave": wave,
            "Threshold": threshold,
            "AREA": 0.0,
            "TST": tst,
            "TotalAREA": total_area,
            "HB": 0.0,
            "numEvents": 0,
            "events_per_hour": 0.0,
        }

    # ------------------------------------------------------------------ #
    # 3) Build 201-sample windows centred on every event
    #    (event-100 .. event+100, inclusive => 201 samples).
    #    Windows that fall outside the signal stay all-NaN.
    # ------------------------------------------------------------------ #
    n_events = len(desat_events)
    desat_windows = np.full((n_events, 201), np.nan, dtype=float)
    for k, e in enumerate(desat_events):  # e is 1-based
        # The MATLAB test reads "e-100 < 0"; a window must lie fully inside the
        # signal, i.e. 1 <= e-100 and e+100 <= n. (e == 100 would index 0 and
        # crash in MATLAB, so such events never occur in valid runs; using
        # e-100 < 1 here is equivalent for any data MATLAB does not crash on.)
        if e - 100 < 1 or e + 100 > n:
            continue
        desat_windows[k, :] = spo2[(e - 100) - 1: (e + 100)]  # SpO2(e-100 : e+100)

    dirty = desat_windows.copy()  # keep every row, NaNs included

    # Keep only windows that are completely free of NaN for the ensemble mean.
    clean_mask = ~np.isnan(desat_windows).any(axis=1)
    clean_windows = desat_windows[clean_mask, :]
    if clean_windows.shape[0] == 0:
        raise RuntimeError("No clean window to calculate desaturation window")

    # ------------------------------------------------------------------ #
    # 4) Ensemble-average desaturation curve and shoulder detection
    # ------------------------------------------------------------------ #
    ens_avg = clean_windows.mean(axis=0)  # length 201, no NaN by construction

    # peak before the nadir: search the first 100 samples (MATLAB 1:100)
    locs = _findpeaks(ens_avg[0:100])
    if locs.size == 0:
        raise RuntimeError("No pre-desaturation peak found in ensemble average.")
    if locs.size > 1:
        before = locs[locs < 90]
        peak1 = int(before[-1]) if before.size else int(locs[-1])
    else:
        peak1 = int(locs[0])

    # peak after the nadir: search samples 101:201, offset back to full index
    locs = _findpeaks(ens_avg[100:201])
    if locs.size == 0:
        raise RuntimeError("No post-desaturation peak found in ensemble average.")
    peak2 = int(locs[0]) + 100

    # ------------------------------------------------------------------ #
    # 5) Define the search window (10% padding on either shoulder) and
    #    integrate every (dirty) window over that span.
    #
    #    IMPORTANT CORRECTION:
    #    The previous Python port integrated the raw SpO2 values in `row`,
    #    which is the lower-side/absolute area under the SpO2 curve. The actual
    #    hypoxic burden is the inverted area above the desaturation curve and
    #    below the event-related baseline. With the baseline taken as the local
    #    event-window peak (per-row max), this is the quantity referred to as
    #    invNormAREA in the side/original calculation, so it is now the main
    #    AREA/HB.
    # ------------------------------------------------------------------ #
    pad = _round_half_away_from_zero((peak2 - peak1) * 0.1)
    padded1 = max(int(peak1 - pad), 1)
    padded2 = min(int(peak2 + pad), 201)  # width(desatWindows) == 201

    cropped = dirty[:, (padded1 - 1): padded2]  # MATLAB dirty(:, padded1:padded2)

    burden_area = 0.0
    lower_side_area = 0.0  # diagnostic only; intentionally not returned as HB
    for row in cropped:
        valid = row[~np.isnan(row)]
        if valid.size == 0:
            continue

        # Old/misplaced calculation: raw lower-side area under SpO2.
        lower_side_area += _trapz(valid)

        # Correct calculation: inverted burden area. The event-related baseline
        # is the highest valid SpO2 value in this subject-specific event window.
        # Clip tiny numerical negatives so a flat/near-flat segment contributes
        # zero burden rather than a negative value.
        baseline = float(np.nanmax(valid))
        inverted_desaturation = np.maximum(baseline - valid, 0.0)
        burden_area += _trapz(inverted_desaturation)

    hb = burden_area / tst if tst else 0.0
    events_per_hour = n_events / (tst / 60.0 / 60.0) if tst else 0.0

    return {
        "ID": subject_id,
        "Wave": wave,
        "Threshold": threshold,
        "AREA": burden_area,        # corrected/inverted area; invNormAREA numerator
        "TST": tst,                 # seconds
        "TotalAREA": total_area,    # raw full-signal area; diagnostic only
        "HB": hb,                 # corrected HB = AREA / TST
        "numEvents": n_events,      # count
        "events_per_hour": events_per_hour,
    }


# --------------------------------------------------------------------------- #
# Command-line interface
# --------------------------------------------------------------------------- #
def _load_signal(path):
    """Load the SpO2 signal from a .npy / .npz / .csv / .txt file."""
    lower = path.lower()
    if lower.endswith(".npy"):
        return np.load(path)
    if lower.endswith(".npz"):
        data = np.load(path)
        key = list(data.keys())[0]
        return data[key]
    # plain text / csv with one column of numbers
    return np.loadtxt(path, delimiter=",")


def main(argv=None):
    parser = argparse.ArgumentParser(
        description="Compute the hypoxic burden of a 1 Hz SpO2 signal "
                    "(Python port of HypoxicBurden.m)."
    )
    parser.add_argument("signal",
                        help="Path to the SpO2 signal (.npy, .npz, .csv or .txt). "
                             "Must be a 1-D array sampled at 1 Hz.")
    parser.add_argument("--id", default="", help="Subject ID.")
    parser.add_argument("--wave", default="", help="Assessment wave.")
    parser.add_argument("--threshold", type=float, default=2.0,
                        help="Desaturation threshold in %% (default: 2).")
    parser.add_argument("--output", default=None,
                        help="Optional path to write the result as a one-row CSV.")
    args = parser.parse_args(argv)

    spo2 = _load_signal(args.signal)
    result = hypoxic_burden(spo2, subject_id=args.id, wave=args.wave,
                            threshold=args.threshold)

    # Pretty print to stdout
    width = max(len(k) for k in result)
    for key, value in result.items():
        print(f"{key:<{width}} : {value}")

    if args.output:
        try:
            import pandas as pd
            pd.DataFrame([result]).to_csv(args.output, index=False)
        except ImportError:
            import csv
            with open(args.output, "w", newline="") as fh:
                writer = csv.DictWriter(fh, fieldnames=list(result))
                writer.writeheader()
                writer.writerow(result)
        print(f"\nSaved result to {args.output}")

    return result


if __name__ == "__main__":
    main()