import streamlit as st
import numpy as np
import onnxruntime as rt
import soundfile as sf

# Load the ONNX model
model_path = 'model.onnx'  # Replace with the actual path to your model
model_inference = rt.InferenceSession(model_path)

# Function to preprocess audio data
def preprocess_audio(audio_file):
    raw_values, rate = sf.read(audio_file)
    mu = np.nanmean(raw_values)
    std = np.nanstd(raw_values)
    audio_data = (raw_values - mu) / std
    audio_data = np.pad(audio_data, (0, 22050 - len(audio_data)), 'constant').reshape(1, -1, 1).astype(np.float32)
    return audio_data

# Function to preprocess clinical data
def preprocess_clinical_data(age, height, weight, reported_cough_dur, heart_rate, temperature, sex, tb_prior, tb_prior_Pul, tb_prior_Extrapul, tb_prior_Unknown, hemoptysis, weight_loss, smoke_lweek, fever, night_sweats):
    sex_Female = 1 if sex == 'Female' else 0
    sex_Male = 1 if sex == 'Male' else 0
    tb_prior_No = 1 if tb_prior == 'No' else 0
    tb_prior_Not_sure = 1 if tb_prior == 'Not sure' else 0
    tb_prior_Yes = 1 if tb_prior == 'Yes' else 0
    tb_prior_Pul_No = 1 if tb_prior_Pul == 'No' else 0
    tb_prior_Pul_Yes = 1 if tb_prior_Pul == 'Yes' else 0
    tb_prior_Extrapul_No = 1 if tb_prior_Extrapul == 'No' else 0
    tb_prior_Extrapul_Yes = 1 if tb_prior_Extrapul == 'Yes' else 0
    tb_prior_Unknown_No = 1 if tb_prior_Unknown == 'No' else 0
    tb_prior_Unknown_Yes = 1 if tb_prior_Unknown == 'Yes' else 0
    hemoptysis_No = 1 if hemoptysis == 'No' else 0
    hemoptysis_Yes = 1 if hemoptysis == 'Yes' else 0
    weight_loss_No = 1 if weight_loss == 'No' else 0
    weight_loss_Yes = 1 if weight_loss == 'Yes' else 0
    smoke_lweek_No = 1 if smoke_lweek == 'No' else 0
    smoke_lweek_Yes = 1 if smoke_lweek == 'Yes' else 0
    fever_No = 1 if fever == 'No' else 0
    fever_Yes = 1 if fever == 'Yes' else 0
    night_sweats_No = 1 if night_sweats == 'No' else 0
    night_sweats_Yes = 1 if night_sweats == 'Yes' else 0

    clinical_data = [age, height, weight, reported_cough_dur, heart_rate, temperature,
                     sex_Female, sex_Male, tb_prior_No, tb_prior_Not_sure, tb_prior_Yes,
                     tb_prior_Pul_No, tb_prior_Pul_Yes, tb_prior_Extrapul_No, tb_prior_Extrapul_Yes,
                     tb_prior_Unknown_No, tb_prior_Unknown_Yes, hemoptysis_No, hemoptysis_Yes,
                     weight_loss_No, weight_loss_Yes, smoke_lweek_No, smoke_lweek_Yes,
                     fever_No, fever_Yes, night_sweats_No, night_sweats_Yes]
    clinical_data = np.array(clinical_data).reshape(1, -1).astype(np.float32)
    return clinical_data

# Main function to run the app
def main():
    st.title('TB Cough Sound Analysis')

    st.write('**Step 1: Upload Cough Sound Recording**')
    audio_file = st.file_uploader("Upload a cough sound recording", type=['wav', 'mp3'])

    if audio_file is not None:
        st.audio(audio_file)

        st.write('**Step 2: Enter Clinical Data**')
        age = st.slider('Age', min_value=0, max_value=100, value=30, step=1)
        height = st.slider('Height (in cm)', min_value=0.0, max_value=300.0, value=170.0, step=0.1)
        weight = st.slider('Weight (in kg)', min_value=0.0, max_value=200.0, value=70.0, step=0.1)
        reported_cough_dur = st.slider('Reported Cough Duration (in days)', min_value=0, max_value=100, value=10, step=1)
        heart_rate = st.slider('Heart Rate (beats per minute)', min_value=0, max_value=200, value=80, step=1)
        temperature = st.slider('Body Temperature (in °C)', min_value=0.0, max_value=40.0, value=37.0, step=0.1)
        sex = st.radio('Sex', ['Male', 'Female'])
        tb_prior = st.selectbox('TB Prior', ['No', 'Not sure', 'Yes'], index=0)
        tb_prior_Pul = st.selectbox('TB Prior - Pulmonary', ['No', 'Yes'], index=0)
        tb_prior_Extrapul = st.selectbox('TB Prior - Extrapulmonary', ['No', 'Yes'], index=0)
        tb_prior_Unknown = st.selectbox('TB Prior - Unknown', ['No', 'Yes'], index=0)
        hemoptysis = st.selectbox('Hemoptysis', ['No', 'Yes'], index=0)
        weight_loss = st.selectbox('Weight Loss', ['No', 'Yes'], index=0)
        smoke_lweek = st.selectbox('Smoking Last Week', ['No', 'Yes'], index=0)
        fever = st.selectbox('Fever', ['No', 'Yes'], index=0)
        night_sweats = st.selectbox('Night Sweats', ['No', 'Yes'], index=0)

        if st.button('Predict'):
            audio_data = preprocess_audio(audio_file)
            clinical_data = preprocess_clinical_data(age, height, weight, reported_cough_dur, heart_rate, temperature,
                                                     sex, tb_prior, tb_prior_Pul, tb_prior_Extrapul, tb_prior_Unknown, hemoptysis, weight_loss, smoke_lweek, fever,
                                                     night_sweats)
            input_name = model_inference.get_inputs()[0].name
            input_name2 = model_inference.get_inputs()[1].name
            label_name = model_inference.get_outputs()[0].name
            onnx_pred = model_inference.run([label_name], {input_name: audio_data, input_name2: clinical_data})
            result = onnx_pred[0]

            # Interpret prediction
            if result >= 0.5:
                prediction_text = "Tuberculosis is found based on the audio data and clinical information."
                probability_label = "Positive Probability"
            else:
                prediction_text = "No Tuberculosis is found based on the audio data and clinical information."
                probability_label = "Negative Probability"

            # Display prediction in a popup
            st.markdown(f"<h2 style='text-align:center;'>Prediction:</h2>", unsafe_allow_html=True)
            st.markdown(f"<h3 style='text-align:center; font-weight:bold;'>{prediction_text}</h3>", unsafe_allow_html=True)
            st.markdown(f"<h4 style='text-align:center;'>{probability_label}: {result}</h4>", unsafe_allow_html=True)

if __name__ == '__main__':
    main()