import streamlit as st
import librosa
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from datetime import datetime

# Load the pre-trained ResNet model
@st.cache_resource
def load_model():
    model = tf.keras.models.load_model('Heart_ResNet.h5')
    return model

model = load_model()

# Initialize session state
if 'page' not in st.session_state:
    st.session_state.page = '🏠 Home'
if 'theme' not in st.session_state:
    st.session_state.theme = 'Light Green'
if 'history' not in st.session_state:
    st.session_state.history = []

# Custom CSS for theme
def apply_theme():
    if st.session_state.theme == "Light Green":
        st.markdown("""
        <style>
        body, .stApp { background-color: #e8f5e9; }
        .stApp { color: #004d40; }
        .stButton > button, .stFileUpload > div { 
            background-color: #004d40; 
            color: white; 
        }
        .stButton > button:hover, .stFileUpload > div:hover {
            background-color: #00332c;
        }
        </style>
        """, unsafe_allow_html=True)
    else:
        st.markdown("""
        <style>
        body, .stApp { background-color: #e0f7fa; }
        .stApp { color: #006064; }
        .stButton > button, .stFileUpload > div { 
            background-color: #006064; 
            color: white; 
        }
        .stButton > button:hover, .stFileUpload > div:hover {
            background-color: #004d40;
        }
        </style>
        """, unsafe_allow_html=True)

# Sidebar navigation
with st.sidebar:
    st.title("Heartbeat Analysis 🩺")
    st.session_state.page = st.radio(
        "Navigation",
        ["🏠 Home", "⚙️ Settings", "👤 Profile"],
        index=["🏠 Home", "⚙️ Settings", "👤 Profile"].index(st.session_state.page)
    )

# Audio processing function
def process_audio(file_path):
    SAMPLE_RATE = 22050
    DURATION = 10
    input_length = int(SAMPLE_RATE * DURATION)
    
    X, sr = librosa.load(file_path, sr=SAMPLE_RATE, duration=DURATION)
    
    if len(X) < input_length:
        pad_width = input_length - len(X)
        X = np.pad(X, (0, pad_width), mode='constant')
    
    mfccs = np.mean(librosa.feature.mfcc(y=X, sr=sr, n_mfcc=52,
                                        n_fft=512, hop_length=256).T, axis=0)
    return mfccs, X, sr

def classify_audio(filepath):
    mfccs, waveform, sr = process_audio(filepath)
    features = mfccs.reshape(1, 52, 1)
    preds = model.predict(features)
    class_names = ["artifact", "murmur", "normal"]
    result = {name: float(preds[0][i]) for i, name in enumerate(class_names)}
    
    # Store in history
    st.session_state.history.append({
        'date': datetime.now().strftime("%Y-%m-%d %H:%M"),
        'file': filepath,
        'result': result
    })
    
    return result, waveform, sr

# Page rendering functions
def home_page():
    st.title("Heartbeat Analysis")
    uploaded_file = st.file_uploader("Upload your heartbeat audio", type=["wav", "mp3"])
    
    if uploaded_file is not None:
        st.audio(uploaded_file.read(), format='audio/wav')
        uploaded_file.seek(0)
        
        if st.button("Analyze Now"):
            with st.spinner('Analyzing...'):
                with open("temp.wav", "wb") as f:
                    f.write(uploaded_file.getbuffer())
                
                results, waveform, sr = classify_audio("temp.wav")
                
                st.subheader("Analysis Results")
                cols = st.columns(3)
                labels = {
                    'artifact': "🚨 Artifact",
                    'murmur': "💔 Murmur",
                    'normal': "❤️ Normal"
                }
                
                for (label, value), col in zip(results.items(), cols):
                    with col:
                        st.metric(labels[label], f"{value*100:.2f}%")
                
                st.subheader("Heartbeat Waveform")
                fig, ax = plt.subplots(figsize=(10, 3))
                librosa.display.waveshow(waveform, sr=sr, ax=ax)
                ax.set_title("Audio Waveform Analysis")
                st.pyplot(fig)

def settings_page():
    st.title("Settings")
    new_theme = st.selectbox(
        "Select Theme",
        ["Light Green", "Light Blue"],
        index=0 if st.session_state.theme == "Light Green" else 1
    )
    
    if new_theme != st.session_state.theme:
        st.session_state.theme = new_theme
        st.experimental_rerun()

def profile_page():
    st.title("Medical Profile")
    with st.expander("Personal Information", expanded=True):
        col1, col2 = st.columns(2)
        with col1:
            st.write("**Name:** Kpetaa Patrick")
            st.write("**Age:** 35")
        with col2:
            st.write("**Blood Type:** O+")
            st.write("**Last Checkup:** 2025-06-17")
    
    st.subheader("Analysis History")
    if not st.session_state.history:
        st.write("No previous analyses found")
    else:
        for analysis in reversed(st.session_state.history):
            with st.expander(f"Analysis from {analysis['date']}"):
                st.write(f"File: {analysis['file']}")
                st.write("Results:")
                for label, value in analysis['result'].items():
                    st.progress(value, text=f"{label.capitalize()}: {value*100:.2f}%")

# Main app logic
apply_theme()

if st.session_state.page == "🏠 Home":
    home_page()
elif st.session_state.page == "⚙️ Settings":
    settings_page()
elif st.session_state.page == "👤 Profile":
    profile_page()