import streamlit as st
from st_audiorec import st_audiorec
import matplotlib.pyplot as plt
import sounddevice as sd
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torchaudio
import wave
import io
from scipy.io import wavfile
import pydub
import time
import os
import atexit
import librosa
import torchaudio.functional as F
import torchaudio.transforms as T

# MODEL LOADING and INITIALISATION
n_fft = 1024
win_length = None
hop_length = 32
# Input tensor shape was ([1,16000])
class SnoreNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.transform = torchaudio.transforms.Spectrogram(n_fft = n_fft,win_length = win_length,hop_length = hop_length,center = True ,pad_mode = "reflect",power = 2.0,)
        self.fc1 = nn.Linear(257013, 512)
        self.act1 = nn.Tanh()
        self.fc2 = nn.Linear(512, 2)
        self.logs1 = nn.LogSoftmax(dim=1)
    
    def forward(self, raw_audio_tensor):
        # print(raw_audio_tensor.shape)
        spectrogram = self.transform(raw_audio_tensor)
        # print(spectrogram.shape)
        spectrogram = spectrogram.reshape(spectrogram.size(0), -1)
        # print(spectrogram.shape)
        output = self.act1(self.fc1(spectrogram))
        output = self.fc2(output)
        output = torch.abs(self.logs1(output))
        return output
model = SnoreNet()
model.load_state_dict(torch.load('snoreNetv1.pt'))
model.eval()
# Audio parameters
def process_data(waveform_chunks):
    snore = 0
    other = 0
    for chunk in waveform_chunks:
        input_tensor = torch.tensor(chunk).unsqueeze(0).to(torch.float32)
        # st.write(input_tensor[0][98])
        with torch.no_grad():
            result = model(input_tensor)
        # st.write(result)
        if np.abs(result[0][0]) > np.abs(result[0][1]):
            other += 1
        else:
            snore += 1
    return snore, other
st.sidebar.markdown(
    """
    <div align="justify">
        <h4>ABOUT</h4>
        <p>Transform your sleep experience with the cutting-edge Snore Detector by Hypermind Labs!
        Discover the power to monitor and understand your nighttime sounds like never before.
        Take control of your sleep quality and uncover the secrets of your peaceful slumber with our innovative app.</p>
    </div>
    """,
    unsafe_allow_html=True,
)
st.title('Real-Time Snore Detection App 😴')
uploaded_file = st.file_uploader("Upload Sample", type=["wav"])
if uploaded_file is not None:
    st.write("Analsysing...")
    audio_bytes = uploaded_file.getvalue()
    audio_array = np.frombuffer(audio_bytes, dtype=np.int16)
    chunk_size = 80000
    num_chunks = len(audio_array) // chunk_size
    waveform_chunks = np.array_split(audio_array[:num_chunks * chunk_size], num_chunks)
    snore, other = process_data(waveform_chunks)
    total = snore + other
    snore_percentage = (snore / total) * 100
    other_percentage = (other / total) * 100
    categories = ["Snore"]
    percentages = [snore_percentage]
    plt.figure(figsize=(8, 4))
    plt.barh(categories, percentages, color=['#ff0033'])
    plt.xlabel('Percentage')
    plt.title('Percentage of Snoring')
    plt.xlim(0, 100)
    for i, percentage in enumerate(percentages):
        plt.text(percentage, i, f' {percentage:.2f}%', va='center')
    st.write("DONE")
    st.pyplot(plt)