'''
ECE 8527
Anna K. Chau
Intro to Machine Learning

Final Project -- Convolutional Neural Network (CNN)
--> script to *test* CNN + get predictions

usage: 
python dnn_predict.py <model path> <test .npy> <output predictions path>

'''
import numpy as np
import argparse
import pandas as pd
from tensorflow.keras.models import load_model

def load_and_predict(model_path, input_file):
    # load
    model = load_model(model_path)
    input_data = np.load(input_file)
    predictions = np.round(model.predict(input_data)).astype(int)

    return predictions

def save_predictions_to_csv(predictions, output_file):
    headers = ["1dAVb", "RBBB", "LBBB", "SB", "AF", "ST"]
    df = pd.DataFrame(predictions, columns=headers)
    df.to_csv(output_file, index=False)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Load saved model and make predictions on input data')
    parser.add_argument('model_path', type=str, help='Path to the saved model file (.h5)')
    parser.add_argument('input_file', type=str, help='Path to the input file containing data in .npy format')
    parser.add_argument('output_file', type=str, help='Path to the output CSV file to save predictions')
    args = parser.parse_args()

    predictions = load_and_predict(args.model_path, args.input_file)

    save_predictions_to_csv(predictions, args.output_file)

    print(f"Predictions saved to {args.output_file}")