# Build the Model
model = tf.keras.Sequential()
model.add(Conv1D(filters=64, kernel_size=3, activation='relu', input_shape=(signal_length, 1), name = 'first_conv'))
model.add(MaxPooling1D(pool_size=2))
model.add(Dropout(0.2))

model.add(Conv1D(filters = 128, kernel_size = 5, activation = 'relu', name = 'second_conv'))
model.add(BatchNormalization())
model.add(MaxPooling1D(pool_size = 2))
model.add(Dropout(0.2))

model.add(Conv1D(filters = 256, kernel_size = 10, activation = 'relu', name='last_conv'))
model.add(BatchNormalization())
model.add(MaxPooling1D(pool_size = 2))
model.add(Dropout(0.2))

model.add(Flatten())

model.add(Dense(units=128, activation='relu'))
model.add(Dropout(0.2))

model.add(Dense(units=5, activation='softmax')) # 5 classes for N, S, V, F, Q

# Compile the Model
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

# Print Model.output
model.summary()

# Train the Model
history = model.fit(X_train, y_train, epochs=5, batch_size=32, validation_data=(X_test, y_test))
loss, accuracy = model.evaluate(X_test, y_test)
print("Loss: ", loss)
print("Accuracy: ", accuracy)

plt.plot(history.history['accuracy'], label='train accuracy')
plt.plot(history.history['val_accuracy'], label='val accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import random

def grad_cam(model, layer_name, data, pred_index):
    # Define a model that outputs the layer's activations and the final prediction
    grad_model = tf.keras.models.Model(
        inputs=[model.inputs], 
        outputs=[model.get_layer(layer_name).output, model.output]
    )

    with tf.GradientTape() as tape:
        conv_outputs, predictions = grad_model(data)
        loss = predictions[:, pred_index]

    grads = tape.gradient(loss, conv_outputs)
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))  # Average over the height, width, and channels

    conv_outputs = conv_outputs[0]
    heatmap = tf.reduce_mean(tf.multiply(pooled_grads, conv_outputs), axis=-1)
    heatmap = np.maximum(heatmap, 0) / (np.max(heatmap) + tf.keras.backend.epsilon())  # Avoid division by zero
    return heatmap

def plot_gradcam(ecg_data, cam, prediction, class_names, title='GradCAM'):
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(ecg_data.flatten(), alpha=0.8, label='ECG Signal')

    cax = ax.imshow(np.expand_dims(cam, axis=0), cmap='jet', aspect='auto', alpha=0.5, extent=[0, len(ecg_data.flatten()), ax.get_ylim()[0], ax.get_ylim()[1]])
    fig.colorbar(cax, ax=ax, orientation='horizontal', label='Activation Intensity')
    ax.set_title(title)
    ax.set_xlabel('Time (ms)')
    ax.set_ylabel('Activation')
    ax.legend()

    text = f'Classification: {class_names[np.argmax(prediction)]}nPrediction: {class_names[np.argmax(prediction)]}nConfidence: {np.max(prediction):.2f}'
    print(text)
    plt.show()

# Example usage
sample_index = random.randint(0, X_test.shape[0] - 1)
sample_data = X_test[sample_index:sample_index + 1]
class_index = np.argmax(y_test[sample_index])

cam = grad_cam(model, 'last_conv', sample_data, class_index)
prediction = model.predict(sample_data)
class_names = list(id_to_label.values())

for i in range(10):
    sample_index = random.randint(0, X_test.shape[0] - 1)
    sample_data = X_test[sample_index:sample_index + 1]
    class_index = np.argmax(y_test[sample_index])
    cam = grad_cam(model, 'last_conv', sample_data, class_index)
    prediction = model.predict(sample_data)
    plot_gradcam(sample_data[0], cam, prediction, class_names, title=f'GradCAM for Sample {i + 1}')