import numpy as np
import tensorflow as tf

# Load trained model
MODEL_PATH = "ptbxl_ecg_cnn_model.keras"
model = tf.keras.models.load_model(MODEL_PATH, compile=False)

# Diagnostic labels
TARGET_NAMES = ["NORM", "MI", "STTC", "CD", "HYP"]

# Default thresholds (can be replaced with tuned ones)
DEFAULT_THRESHOLDS = {
    "NORM": 0.5,
    "MI": 0.5,
    "STTC": 0.5,
    "CD": 0.5,
    "HYP": 0.5
}


def predict_ecg(ecg_signal, thresholds=DEFAULT_THRESHOLDS):
    """
    Predict cardiac abnormalities from a 12-lead ECG.

    Parameters
    ----------
    ecg_signal : np.ndarray
        Shape (1000, 12), preprocessed ECG signal
    thresholds : dict
        Thresholds for each class

    Returns
    -------
    predicted_labels : list
        List of predicted diagnostic labels
    probabilities : dict
        Probability per diagnostic class
    """

    if ecg_signal.shape != (1000, 12):
        raise ValueError("ECG signal must have shape (1000, 12)")

    # Add batch dimension
    ecg_signal = np.expand_dims(ecg_signal, axis=0)

    # Model prediction
    probs = model.predict(ecg_signal, verbose=0)[0]

    predicted_labels = []
    probabilities = {}

    for i, label in enumerate(TARGET_NAMES):
        probabilities[label] = float(probs[i])
        if probs[i] >= thresholds[label]:
            predicted_labels.append(label)

    return predicted_labels, probabilities


# Example usage (for testing only)
if __name__ == "__main__":
    dummy_ecg = np.random.randn(1000, 12)
    labels, probs = predict_ecg(dummy_ecg)

    print("Predicted labels:", labels)
    print("Probabilities:")
    for k, v in probs.items():
        print(f"{k}: {v:.3f}")