script.py

import argparse
import json
import logging
import math
from collections import defaultdict
from pathlib import Path
import matplotlib.pyplot as plt
import mne
import neurokit2 as nk
import numpy as np
import scipy.interpolate
from scipy.signal import resample
from tqdm import tqdm


mne.set_log_level("ERROR")

FS = 256
MIN_RPEAKS = 50

logger = logging.getLogger(__name__)


def json_friendly(x):
    """Convert numpy values to normal Python values before saving as JSON."""
    if isinstance(x, dict):
        return {k: json_friendly(v) for k, v in x.items()}

    if isinstance(x, list):
        return [json_friendly(v) for v in x]

    if isinstance(x, tuple):
        return [json_friendly(v) for v in x]

    if isinstance(x, np.ndarray):
        return json_friendly(x.tolist())

    if isinstance(x, np.integer):
        return int(x)

    if isinstance(x, np.floating):
        x = float(x)
        return None if math.isnan(x) or math.isinf(x) else x

    if isinstance(x, float):
        return None if math.isnan(x) or math.isinf(x) else x

    return x


def get_ecg_channel(raw):
    """Find the ECG channel in an MNE Raw object."""
    ecg_channels = [
        ch for ch in raw.ch_names
        if "ecg" in ch.lower() or "ekg" in ch.lower()
    ]

    if len(ecg_channels) == 0:
        raise ValueError(f"No ECG channel found. Available channels: {raw.ch_names}")

    return ecg_channels[0]


def clean_ecg(raw, fs=FS):
    """
    Extract and clean ECG.

    """
    ecg_channel = get_ecg_channel(raw)
    ecg = raw.get_data(picks=[ecg_channel]).flatten()

    original_fs = raw.info["sfreq"]

    if not np.isclose(original_fs, fs):
        n_samples = int(round(len(ecg) * fs / original_fs))
        ecg = resample(ecg, n_samples)

    ecg_clean = nk.ecg_clean(ecg, sampling_rate=fs, method="biosppy")

    info = {
        "ecg_channel": ecg_channel,
        "original_fs": float(original_fs),
        "target_fs": fs,
        "n_samples": len(ecg_clean),
    }

    return ecg_clean, ecg, info


def detect_rpeaks(ecg_clean, fs=FS):
    """Detect R-peaks using NeuroKit2."""
    _, rpeaks = nk.ecg_peaks(
        ecg_clean,
        sampling_rate=fs,
        method="nabian2018"
    )

    rpeaks["ECG_R_Peaks"] = np.asarray(rpeaks["ECG_R_Peaks"], dtype=int)
    rpeaks["sampling_rate"] = fs

    return rpeaks


def remove_bad_rpeaks(rpeaks, fs=FS, mad_threshold=3.0, min_rr_ms=300, max_rr_ms=2000):
    """
    Remove unlikely R-peaks based on RR intervals.

    """
    r = np.asarray(rpeaks["ECG_R_Peaks"], dtype=int)

    if len(r) < 3:
        return rpeaks, {
            "n_rpeaks_before": len(r),
            "n_rpeaks_after": len(r),
            "removed_percentage": 0.0,
        }

    rr_ms = np.diff(r) / fs * 1000

    median_rr = np.median(rr_ms)
    mad_rr = np.median(np.abs(rr_ms - median_rr))

    if mad_rr == 0:
        mad_outliers = np.zeros(len(rr_ms), dtype=bool)
    else:
        mad_outliers = np.abs(rr_ms - median_rr) / mad_rr > mad_threshold

    physiological_outliers = (rr_ms < min_rr_ms) | (rr_ms > max_rr_ms)
    bad_rr = mad_outliers | physiological_outliers

    # For each bad interval, remove the second R-peak.
    bad_peak_positions = np.where(bad_rr)[0] + 1
    bad_peak_positions = np.unique(bad_peak_positions)

    keep = np.ones(len(r), dtype=bool)
    keep[bad_peak_positions] = False

    r_clean = r[keep]

    report = {
        "n_rpeaks_before": int(len(r)),
        "n_rpeaks_after": int(len(r_clean)),
        "n_removed_rpeaks": int(len(bad_peak_positions)),
        "removed_percentage": float(len(bad_peak_positions) / len(r) * 100),
        "median_rr_ms": float(median_rr),
        "mad_rr_ms": float(mad_rr),
    }

    rpeaks_clean = dict(rpeaks)
    rpeaks_clean["ECG_R_Peaks"] = r_clean

    return rpeaks_clean, report


def valid_index(x, n):
    try:
        if x is None:
            return False

        x = float(x)

        if np.isnan(x) or np.isinf(x):
            return False

        x = int(round(x))
        return 0 <= x < n

    except Exception:
        return False


def to_index(x):
    return int(round(float(x)))


def get_wave(waves, name, i):
    try:
        return waves[name][i]
    except Exception:
        return np.nan


def angle_between_points(left, centre, right, signal, fs=FS):
    """
    Angle at the centre point.

    The x-axis is time in ms and the y-axis is ECG amplitude in microvolts.
    """
    signal_uv = signal * 1e6

    p1 = np.array([left / fs * 1000, signal_uv[left]])
    p2 = np.array([centre / fs * 1000, signal_uv[centre]])
    p3 = np.array([right / fs * 1000, signal_uv[right]])

    v1 = p1 - p2
    v2 = p3 - p2

    denom = np.linalg.norm(v1) * np.linalg.norm(v2)

    if denom == 0:
        return np.nan

    cos_angle = np.dot(v1, v2) / denom
    cos_angle = np.clip(cos_angle, -1, 1)

    return float(np.degrees(np.arccos(cos_angle)))


def extract_morphology(ecg_clean, rpeaks, fs=FS):
    """Extract ECG morphology features from delineated ECG waves."""
    features = {
        "qrs_duration_ms": [],
        "st_segment_ms": [],
        "pr_interval_ms": [],
        "pr_segment_ms": [],
        "qt_interval_ms": [],
        "p_amp_uV": [],
        "q_amp_uV": [],
        "r_amp_uV": [],
        "s_amp_uV": [],
        "t_amp_uV": [],
        "q_angle_deg": [],
        "r_angle_deg": [],
        "s_angle_deg": [],
    }

    r = np.asarray(rpeaks["ECG_R_Peaks"], dtype=int)

    if len(r) == 0:
        return features

    try:
        _, waves = nk.ecg_delineate(
            ecg_clean,
            rpeaks,
            sampling_rate=fs,
            method="dwt"
        )
    except Exception as e:
        logger.warning("ECG delineation failed: %s", e)
        return features

    n = len(ecg_clean)

    for i, r_peak in enumerate(r):
        p_on = get_wave(waves, "ECG_P_Onsets", i)
        p_peak = get_wave(waves, "ECG_P_Peaks", i)
        p_off = get_wave(waves, "ECG_P_Offsets", i)

        q_peak = get_wave(waves, "ECG_Q_Peaks", i)

        r_on = get_wave(waves, "ECG_R_Onsets", i)
        r_off = get_wave(waves, "ECG_R_Offsets", i)

        s_peak = get_wave(waves, "ECG_S_Peaks", i)

        t_on = get_wave(waves, "ECG_T_Onsets", i)
        t_peak = get_wave(waves, "ECG_T_Peaks", i)
        t_off = get_wave(waves, "ECG_T_Offsets", i)

        points = [
            p_on, p_peak, p_off,
            q_peak, r_on, r_peak, r_off,
            s_peak, t_on, t_peak, t_off,
        ]

        if not all(valid_index(x, n) for x in points):
            continue

        p_on = to_index(p_on)
        p_peak = to_index(p_peak)
        p_off = to_index(p_off)
        q_peak = to_index(q_peak)
        r_on = to_index(r_on)
        r_peak = to_index(r_peak)
        r_off = to_index(r_off)
        s_peak = to_index(s_peak)
        t_on = to_index(t_on)
        t_peak = to_index(t_peak)
        t_off = to_index(t_off)

        qrs_ms = (r_off - r_on) / fs * 1000
        st_ms = (t_on - r_off) / fs * 1000
        pr_interval_ms = (r_on - p_on) / fs * 1000
        pr_segment_ms = (r_on - p_off) / fs * 1000
        qt_ms = (t_off - r_on) / fs * 1000

        # Skip clearly impossible delineations.
        if any(x <= 0 for x in [qrs_ms, pr_interval_ms, qt_ms]):
            continue

        features["qrs_duration_ms"].append(float(qrs_ms))
        features["st_segment_ms"].append(float(st_ms))
        features["pr_interval_ms"].append(float(pr_interval_ms))
        features["pr_segment_ms"].append(float(pr_segment_ms))
        features["qt_interval_ms"].append(float(qt_ms))

        features["p_amp_uV"].append(float(ecg_clean[p_peak] * 1e6))
        features["q_amp_uV"].append(float(ecg_clean[q_peak] * 1e6))
        features["r_amp_uV"].append(float(ecg_clean[r_peak] * 1e6))
        features["s_amp_uV"].append(float(ecg_clean[s_peak] * 1e6))
        features["t_amp_uV"].append(float(ecg_clean[t_peak] * 1e6))

        features["q_angle_deg"].append(
            angle_between_points(p_peak, q_peak, r_peak, ecg_clean, fs)
        )

        features["r_angle_deg"].append(
            angle_between_points(q_peak, r_peak, s_peak, ecg_clean, fs)
        )

        features["s_angle_deg"].append(
            angle_between_points(r_peak, s_peak, t_peak, ecg_clean, fs)
        )

    return features


def get_df_value(df, col):
    try:
        return float(df[col].iloc[0])
    except Exception:
        return np.nan


def extract_hrv(rpeaks, fs=FS):
    """Extract time, frequency and nonlinear HRV features."""
    r = np.asarray(rpeaks["ECG_R_Peaks"], dtype=int)

    features = {
        "MeanNN": np.nan,
        "SDNN": np.nan,
        "RMSSD": np.nan,
        "SDSD": np.nan,
        "MinNN": np.nan,
        "MaxNN": np.nan,
        "LF": np.nan,
        "HF": np.nan,
        "LFHF": np.nan,
        "SD1": np.nan,
        "SD2": np.nan,
        "SD1SD2": np.nan,
        "CSI": np.nan,
        "CVI": np.nan,
        "SpecEn": np.nan,
        "RenEn_alpha2": np.nan,
        "LZC": np.nan,
    }

    for m in [1, 2]:
        for r_frac in [0.1, 0.2]:
            features[f"SampEn_m{m}_r{r_frac}"] = np.nan
            features[f"FuzzyEn_m{m}_r{r_frac}"] = np.nan

    for m in [3, 4, 5]:
        features[f"PermEn_m{m}"] = np.nan

    for c in [5, 6, 7]:
        features[f"DispEn_c{c}"] = np.nan

    for kmax in [5, 10, 20]:
        features[f"HFD_k{kmax}"] = np.nan

    if len(r) < 4:
        return features

    peaks = {"ECG_R_Peaks": r}
    rr = np.diff(r) / fs
    rr_std = np.std(rr)

    try:
        hrv_time = nk.hrv_time(peaks, sampling_rate=fs, show=False)

        features["MeanNN"] = get_df_value(hrv_time, "HRV_MeanNN")
        features["SDNN"] = get_df_value(hrv_time, "HRV_SDNN")
        features["RMSSD"] = get_df_value(hrv_time, "HRV_RMSSD")
        features["SDSD"] = get_df_value(hrv_time, "HRV_SDSD")
        features["MinNN"] = get_df_value(hrv_time, "HRV_MinNN")
        features["MaxNN"] = get_df_value(hrv_time, "HRV_MaxNN")

    except Exception as e:
        logger.warning("Time-domain HRV failed: %s", e)

    try:
        hrv_freq = nk.hrv_frequency(
            peaks,
            sampling_rate=fs,
            interpolation_rate=4,
            psd_method="welch",
            normalize=True,
            show=False
        )

        features["LF"] = get_df_value(hrv_freq, "HRV_LF")
        features["HF"] = get_df_value(hrv_freq, "HRV_HF")
        features["LFHF"] = get_df_value(hrv_freq, "HRV_LFHF")

    except Exception as e:
        logger.warning("Frequency-domain HRV failed: %s", e)

    try:
        hrv_non = nk.hrv_nonlinear(peaks, sampling_rate=fs, show=False)

        features["SD1"] = get_df_value(hrv_non, "HRV_SD1")
        features["SD2"] = get_df_value(hrv_non, "HRV_SD2")
        features["SD1SD2"] = get_df_value(hrv_non, "HRV_SD1SD2")
        features["CSI"] = get_df_value(hrv_non, "HRV_CSI")
        features["CVI"] = get_df_value(hrv_non, "HRV_CVI")

    except Exception as e:
        logger.warning("Nonlinear HRV failed: %s", e)

    for m in [1, 2]:
        for r_frac in [0.1, 0.2]:
            tolerance = r_frac * rr_std

            try:
                value, _ = nk.entropy_sample(
                    rr,
                    delay=1,
                    dimension=m,
                    tolerance=tolerance
                )
                features[f"SampEn_m{m}_r{r_frac}"] = float(value)
            except Exception:
                pass

            try:
                value, _ = nk.entropy_fuzzy(
                    rr,
                    delay=1,
                    dimension=m,
                    tolerance=tolerance,
                    n=2
                )
                features[f"FuzzyEn_m{m}_r{r_frac}"] = float(value)
            except Exception:
                pass

    try:
        value, _ = nk.entropy_renyi(rr, alpha=2)
        features["RenEn_alpha2"] = float(value)
    except Exception:
        pass

    for m in [3, 4, 5]:
        try:
            value, _ = nk.entropy_permutation(
                rr,
                delay=1,
                dimension=m
            )
            features[f"PermEn_m{m}"] = float(value)
        except Exception:
            pass

    for c in [5, 6, 7]:
        try:
            value, _ = nk.entropy_dispersion(
                rr,
                delay=1,
                dimension=2,
                c=c
            )
            features[f"DispEn_c{c}"] = float(value)
        except Exception:
            pass

    rr_interp = interpolate_rr(rr)

    if rr_interp is not None:
        try:
            psd = nk.signal_psd(rr_interp, sampling_rate=4)
            power = psd["Power"].values.astype(float)

            if power.sum() > 0:
                power = power / power.sum()
                spen, _ = nk.entropy_shannon(freq=power)
                features["SpecEn"] = float(spen / np.log2(len(power)))

        except Exception:
            pass

        for kmax in [5, 10, 20]:
            try:
                value, _ = nk.fractal_higuchi(rr_interp, k_max=kmax)
                features[f"HFD_k{kmax}"] = float(value)
            except Exception:
                pass

        try:
            value, _ = nk.complexity_lempelziv(
                rr_interp,
                symbolize="mean"
            )
            features["LZC"] = float(value)
        except Exception:
            pass

    return features


def interpolate_rr(rr, target_fs=4):
    """Interpolate RR intervals to a regular time axis."""
    if len(rr) < 3:
        return None

    rr_time = np.cumsum(rr)

    if rr_time[-1] <= rr_time[0]:
        return None

    n_samples = int((rr_time[-1] - rr_time[0]) * target_fs)

    if n_samples < 3:
        return None

    new_time = np.linspace(rr_time[0], rr_time[-1], n_samples)

    try:
        f = scipy.interpolate.interp1d(
            rr_time,
            rr,
            kind="linear",
            fill_value="extrapolate"
        )

        return f(new_time)

    except Exception:
        return None


def plot_ecg(raw_ecg, clean_ecg, fs, file_name, output_dir):
    """Save a quick plot of raw and cleaned ECG."""
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    n = min(len(raw_ecg), len(clean_ecg))
    t = np.arange(n) / fs

    plt.figure(figsize=(12, 5))
    plt.plot(t, raw_ecg[:n] * 1e6, label="Raw ECG", alpha=0.6)
    plt.plot(t, clean_ecg[:n] * 1e6, label="Cleaned ECG", alpha=0.8)
    plt.xlabel("Time (s)")
    plt.ylabel("Amplitude (µV)")
    plt.title("Raw and cleaned ECG")
    plt.legend()
    plt.tight_layout()

    out_file = output_dir / f"{Path(file_name).stem}_ecg.png"
    plt.savefig(out_file, dpi=300)
    plt.close()


def parse_condition(file_name):
    """
    Detect condition from file name.

    examples:
    SubID_Pre_Ictal_1.edf
    SubID_Inter_Ictal_1.edf
    """
    name = file_name.lower()

    if "pre_ictal" in name or "preictal" in name or "pre-ictal" in name:
        return "preictal"

    if "post_ictal" in name or "postictal" in name or "post-ictal" in name:
        return "postictal"

    if "inter_ictal" in name or "interictal" in name or "inter-ictal" in name:
        return "interictal"

    if "ictal" in name:
        return "ictal"

    return None


def get_patient_id(file_name):
    """
    Extract patient ID from file name.

    """
    return file_name.split("_")[0]


def get_data_files(input_dir):
    input_dir = Path(input_dir)
    return sorted(
        list(input_dir.glob("*.edf")) +
        list(input_dir.glob("*.EDF")) +
        list(input_dir.glob("*.fif")) +
        list(input_dir.glob("*.fif.gz"))
    )

def load_raw_file(file_path):
    file_path = Path(file_path)
    suffix = file_path.suffix.lower()

    if suffix == ".edf":
        return mne.io.read_raw_edf(file_path, preload=True, verbose="ERROR")

    if suffix == ".fif" or file_path.name.lower().endswith(".fif.gz"):
        return mne.io.read_raw_fif(file_path, preload=True, verbose="ERROR")

    raise ValueError(f"Unsupported file type: {file_path}")

def group_by_patient(files):
    patient_files = defaultdict(list)

    for file_path in files:
        patient_id = get_patient_id(file_path.name)
        patient_files[patient_id].append(file_path)

    return patient_files


def patient_has_required_files(files):
    conditions = {parse_condition(f.name) for f in files}
    conditions.discard(None)

    has_interictal = "interictal" in conditions
    has_event = any(c in conditions for c in ["preictal", "ictal", "postictal"])

    return has_interictal and has_event


def process_file(file_path, fs=FS, save_plot=False, plot_dir=None):
    """Process one FIF/EDF file and return ECG features."""
    raw = load_raw_file(file_path)

    ecg_clean, ecg_raw, clean_info = clean_ecg(raw, fs=fs)

    if save_plot and plot_dir is not None:
        plot_ecg(ecg_raw, ecg_clean, fs, file_path.name, plot_dir)

    rpeaks = detect_rpeaks(ecg_clean, fs=fs)
    rpeaks, rpeak_report = remove_bad_rpeaks(rpeaks, fs=fs)

    n_rpeaks = len(rpeaks["ECG_R_Peaks"])

    if n_rpeaks < MIN_RPEAKS:
        raise ValueError(f"Too few R-peaks detected: {n_rpeaks}")

    morphology = extract_morphology(ecg_clean, rpeaks, fs=fs)
    hrv = extract_hrv(rpeaks, fs=fs)

    features = {
        "file_name": file_path.name,
        "duration_sec": len(ecg_clean) / fs,
        "sampling_rate": fs,
        "num_rpeaks": n_rpeaks,
        "cleaning": clean_info,
        "rpeak_filtering": rpeak_report,
        "morphology": morphology,
        "hrv": hrv,
    }

    return features


def process_data(input_dir, output_dir, fs=FS, save_plots=False):
    input_dir = Path(input_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    plot_dir = output_dir / "plots" if save_plots else None

    data_files = get_data_files(input_dir)
    logger.info("Found %d EDF/FIF files", len(data_files))
    
    patient_files = group_by_patient(data_files)
    
    valid_patients = {
        patient: files
        for patient, files in patient_files.items()
        if patient_has_required_files(files)
    }

    logger.info("Found %d patients with event and interictal files", len(valid_patients))

    for patient, files in tqdm(valid_patients.items(), desc="Processing patients"):
        patient_features = defaultdict(list)

        for file_path in files:
            condition = parse_condition(file_path.name)

            if condition is None:
                logger.warning("Could not detect condition for %s", file_path.name)
                continue

            try:
                features = process_file(
                    file_path,
                    fs=fs,
                    save_plot=save_plots,
                    plot_dir=plot_dir
                )

                features["patient_id"] = patient
                features["condition"] = condition

                patient_features[condition].append(features)

            except Exception as e:
                logger.warning("Skipping %s: %s", file_path.name, e)

        for condition, features in patient_features.items():
            if len(features) == 0:
                continue

            out_file = output_dir / f"{patient}_{condition}_ECG_features.json"

            with open(out_file, "w", encoding="utf-8") as f:
                json.dump(
                    json_friendly(features),
                    f,
                    indent=4,
                    ensure_ascii=False,
                    allow_nan=False
                )

            logger.info("Saved %s", out_file)


def main():
    parser = argparse.ArgumentParser(
        description="Extract ECG morphology and HRV features from FIF files."
    )

    parser.add_argument(
        "--input_dir",
        required=True,
        help="Folder containing FIF files."
    )

    parser.add_argument(
        "--output_dir",
        required=True,
        help="Folder where feature files will be saved."
    )

    parser.add_argument(
        "--fs",
        type=int,
        default=FS,
        help="Target sampling frequency. Default is 256 Hz."
    )

    parser.add_argument(
        "--save_plots",
        action="store_true",
        help="Save raw/cleaned ECG plots."
    )

    args = parser.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(levelname)s: %(message)s"
    )

    process_data(
        input_dir=args.input_dir,
        output_dir=args.output_dir,
        fs=args.fs,
        save_plots=args.save_plots
    )


if __name__ == "__main__":
    main()