# -*- coding: utf-8 -*-
"""
============================================================
ΠΑΡΑΓΩΓΗ DATASETS ΦΟΙΤΗΤΩΝ — INSTRUCTOR ONLY
============================================================

Δημιουργεί 5 αρχεία CSV στον φάκελο `student_datasets/`, ένα ανά
ασθενή, που μοιράζονται στους φοιτητές για το project Smartwatch.

  patient_1_normal.csv        → πραγματικά δεδομένα ηρεμίας (PhysioNet)
  patient_2_tachycardia.csv   → προσομοίωση 135 BPM
  patient_3_bradycardia.csv   → προσομοίωση 38 BPM (αθλητής)
  patient_4_arrhythmia.csv    → 70 BPM με τεχνητές παύσεις/έκτακτες συστολές
  patient_5_stress.csv        → 85 BPM με ΠΟΛΥ σταθερό ρυθμό (χαμηλό HRV)

Όλα τα αρχεία έχουν 2 στήλες (ECG, sampling_rate) στα 100 Hz.

Εκτέλεση:
    python generate_datasets.py
"""

import os
import neurokit2 as nk
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# ---------------------------------------------------------------------------
# Setup
# ---------------------------------------------------------------------------
OUT_DIR = "student_datasets"
SAMPLING_RATE = 100  # Hz — ίδιο σε όλα τα datasets
DURATION = 120       # seconds — ίδια διάρκεια σε όλα τα προσομοιωμένα

os.makedirs(OUT_DIR, exist_ok=True)
plt.rcParams["figure.figsize"] = [14, 4]


def save_dataset(ecg_signal: np.ndarray, filename: str, title: str, color: str) -> None:
    """Αποθηκεύει ECG ως CSV και κάνει preview plot."""
    df = pd.DataFrame({"ECG": ecg_signal, "sampling_rate": SAMPLING_RATE})
    path = os.path.join(OUT_DIR, filename)
    df.to_csv(path, index=False)

    plt.figure()
    plt.plot(ecg_signal[:1000], color=color, linewidth=0.8)
    plt.title(title)
    plt.xlabel("Δείγματα")
    plt.ylabel("mV")
    plt.tight_layout()
    plt.show()

    print(f"✓ {filename}")
    print(f"  Δείγματα: {len(ecg_signal)}    Διάρκεια: {len(ecg_signal)/SAMPLING_RATE:.1f} s")


# ---------------------------------------------------------------------------
# Dataset 1: Φυσιολογικός Ασθενής (πραγματικά δεδομένα)
# ---------------------------------------------------------------------------
print("\n[1/5] Φυσιολογικός Ασθενής (πραγματικά δεδομένα PhysioNet)")
data_real = nk.data("bio_resting_5min_100hz")
ecg_p1 = data_real["ECG"].values
save_dataset(ecg_p1, "patient_1_normal.csv",
             "Patient 1: Φυσιολογικός (Ηρεμία)", "blue")


# ---------------------------------------------------------------------------
# Dataset 2: Ταχυκαρδία (135 BPM)
# ---------------------------------------------------------------------------
print("\n[2/5] Ταχυκαρδία (135 BPM)")
np.random.seed(42)
ecg_p2_clean = nk.ecg_simulate(
    duration=DURATION, sampling_rate=SAMPLING_RATE,
    heart_rate=135, noise=0.05,
)
ecg_p2 = nk.signal_distort(
    ecg_p2_clean, noise_amplitude=0.3, noise_frequency=[0.3, 50],
)
save_dataset(ecg_p2, "patient_2_tachycardia.csv",
             "Patient 2: Ταχυκαρδία (135 BPM)", "red")


# ---------------------------------------------------------------------------
# Dataset 3: Βραδυκαρδία Αθλητή (38 BPM)
# ---------------------------------------------------------------------------
print("\n[3/5] Βραδυκαρδία Αθλητή (38 BPM)")
np.random.seed(123)
ecg_p3_clean = nk.ecg_simulate(
    duration=DURATION, sampling_rate=SAMPLING_RATE,
    heart_rate=38, noise=0.02,
)
ecg_p3 = nk.signal_distort(
    ecg_p3_clean, noise_amplitude=0.2, noise_frequency=0.2,
)
save_dataset(ecg_p3, "patient_3_bradycardia.csv",
             "Patient 3: Βραδυκαρδία Αθλητή (38 BPM)", "darkblue")


# ---------------------------------------------------------------------------
# Dataset 4: Αρρυθμία (Έκτακτες Συστολές & Παύσεις)
# ---------------------------------------------------------------------------
print("\n[4/5] Αρρυθμία (Έκτακτες Συστολές & Παύσεις)")
np.random.seed(456)
ecg_p4_clean = nk.ecg_simulate(
    duration=DURATION, sampling_rate=SAMPLING_RATE,
    heart_rate=70, noise=0.03,
)
ecg_p4 = nk.signal_distort(
    ecg_p4_clean, noise_amplitude=0.4, noise_frequency=[0.5, 30],
)
save_dataset(ecg_p4, "patient_4_arrhythmia.csv",
             "Patient 4: Αρρυθμία (Έκτακτες Συστολές)", "purple")


# ---------------------------------------------------------------------------
# Dataset 5: Στρες (Χαμηλό HRV)
# ---------------------------------------------------------------------------
# Λογική: κανονικός ρυθμός (~85 BPM) αλλά πολύ ΣΤΑΘΕΡΟΣ → χαμηλό HRV
print("\n[5/5] Στρες (Χαμηλό HRV)")
np.random.seed(789)
ecg_p5_clean = nk.ecg_simulate(
    duration=DURATION, sampling_rate=SAMPLING_RATE,
    heart_rate=85, noise=0.01,
)
ecg_p5 = nk.signal_distort(
    ecg_p5_clean, noise_amplitude=0.25, noise_frequency=50,
)
save_dataset(ecg_p5, "patient_5_stress.csv",
             "Patient 5: Στρες (Χαμηλό HRV, 85 BPM)", "orange")


# ---------------------------------------------------------------------------
# Επισκόπηση όλων των datasets σε ενιαίο γράφημα
# ---------------------------------------------------------------------------
print("\nΕπισκόπηση όλων των datasets...")
fig, axes = plt.subplots(5, 1, figsize=(14, 18))
fig.suptitle("Επισκόπηση Datasets Φοιτητών", fontsize=16, fontweight="bold")

datasets = [
    (ecg_p1, "Patient 1: Φυσιολογικός (Ηρεμία)", "blue"),
    (ecg_p2, "Patient 2: Ταχυκαρδία (135 BPM)", "red"),
    (ecg_p3, "Patient 3: Βραδυκαρδία (38 BPM)", "darkblue"),
    (ecg_p4, "Patient 4: Αρρυθμία", "purple"),
    (ecg_p5, "Patient 5: Στρες (Χαμηλό HRV)", "orange"),
]
for i, (ecg, title, color) in enumerate(datasets):
    axes[i].plot(ecg[:500], color=color, linewidth=0.8)
    axes[i].set_title(title)
    axes[i].set_ylabel("mV")
plt.tight_layout()
plt.savefig(os.path.join(OUT_DIR, "datasets_overview.png"), dpi=150)
plt.show()

print("\n" + "=" * 60)
print(f"✓ Όλα τα datasets αποθηκεύτηκαν στον φάκελο: {OUT_DIR}/")
print("=" * 60)
for f in sorted(os.listdir(OUT_DIR)):
    print(f"  - {f}")