import os
import numpy as np
import matplotlib.pyplot as plt
import streamlit as st
import tempfile
import soundfile as sf
from scipy.io import wavfile
from scipy import signal

# ตั้งค่าหน้าเพจ
st.set_page_config(
    page_title="Speech Separation Demo App",
    page_icon="🎧",
    layout="wide",
    initial_sidebar_state="expanded"
)

# สไตล์ CSS
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        color: #1E88E5;
        text-align: center;
    }
    .sub-header {
        font-size: 1.5rem;
        color: #42A5F5;
    }
    .success-text {
        color: #4CAF50;
        font-weight: bold;
    }
    .info-text {
        color: #2196F3;
    }
</style>
""", unsafe_allow_html=True)

# ส่วนหัวแอป
st.markdown('<p class="main-header">🎧 Speech Separation Demo</p>', unsafe_allow_html=True)
st.markdown("""
แอปนี้แสดงการสาธิตการแยกเสียงพูดจากเสียงรบกวน โดยใช้การประมวลผลสัญญาณดิจิทัลพื้นฐาน 
ตัวอย่างนี้ใช้ตัวกรองความถี่แบบช่องผ่าน (Bandpass filter) เพื่อลดเสียงรบกวน
""")
st.markdown("---")

# สร้างโฟลเดอร์สำหรับเก็บไฟล์ชั่วคราว
os.makedirs("./temp", exist_ok=True)

# ฟังก์ชันเพิ่มเสียงรบกวน
def add_noise(signal, noise_factor=0.01):
    """เพิ่มเสียงรบกวนให้กับสัญญาณเสียง"""
    noise = np.random.normal(0, noise_factor, signal.shape)
    return signal + noise

# ฟังก์ชันแยกเสียงพูดด้วยตัวกรองความถี่แบบช่องผ่าน
def filter_speech(audio_signal, sample_rate, lowcut=300, highcut=3000, order=5):
    """
    แยกเสียงพูดด้วยตัวกรองความถี่แบบช่องผ่าน
    เสียงพูดมนุษย์ส่วนใหญ่มีความถี่อยู่ระหว่าง 300-3000 Hz
    """
    nyquist = 0.5 * sample_rate
    low = lowcut / nyquist
    high = highcut / nyquist
    
    # ออกแบบตัวกรอง Butterworth
    b, a = signal.butter(order, [low, high], btype='band')
    
    # ใช้ตัวกรองกับสัญญาณ
    filtered_signal = signal.filtfilt(b, a, audio_signal)
    
    return filtered_signal

def process_audio(audio_file):
    """ประมวลผลไฟล์เสียงและแยกเสียงพูด"""
    # บันทึกไฟล์อัปโหลดไปยังไฟล์ชั่วคราว
    with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as tmp_file:
        tmp_file.write(audio_file.getvalue())
        tmp_file_path = tmp_file.name
    
    try:
        # โหลดไฟล์เสียง
        sample_rate, waveform = wavfile.read(tmp_file_path)
        
        # เปลี่ยนเป็น float และ normalize ถ้าจำเป็น
        if waveform.dtype == np.int16:
            waveform = waveform.astype(np.float32) / 32768.0
        elif waveform.dtype == np.int32:
            waveform = waveform.astype(np.float32) / 2147483648.0
            
        # ถ้าเป็น stereo ให้แปลงเป็น mono
        if len(waveform.shape) > 1 and waveform.shape[1] > 1:
            waveform = np.mean(waveform, axis=1)
        
        # แสดงข้อมูลไฟล์เสียง
        st.markdown('<p class="sub-header">ข้อมูลไฟล์เสียง</p>', unsafe_allow_html=True)
        col1, col2 = st.columns(2)
        with col1:
            st.write(f"🔊 Sample rate: {sample_rate} Hz")
        with col2:
            st.write(f"⏱️ ความยาว: {len(waveform)/sample_rate:.2f} วินาที")
        
        # ตัดเสียงถ้ายาวเกินไป (สำหรับ demo)
        max_length = 10 * sample_rate  # สูงสุด 10 วินาที
        if len(waveform) > max_length:
            st.warning(f"ตัดเสียงให้เหลือ 10 วินาทีแรกเพื่อให้การประมวลผลเร็วขึ้น")
            waveform = waveform[:max_length]
        
        st.markdown('<p class="sub-header">เสียงต้นฉบับ</p>', unsafe_allow_html=True)
        
        # แสดงกราฟเสียงต้นฉบับ
        fig, ax = plt.subplots(figsize=(10, 2))
        ax.plot(waveform)
        ax.set_title("Original Waveform")
        ax.set_xlabel("Samples")
        ax.set_ylabel("Amplitude")
        ax.grid(True, alpha=0.3)
        st.pyplot(fig)
        
        # ให้เล่นเสียงต้นฉบับ
        st.audio(tmp_file_path, format="audio/wav")
        
        # เพิ่มเสียงรบกวน (สำหรับการทดสอบถ้าต้องการ)
        st.markdown('<p class="sub-header">ตัวเลือกการเพิ่มเสียงรบกวน</p>', unsafe_allow_html=True)
        add_noise_option = st.checkbox("เพิ่มเสียงรบกวนเพื่อการทดสอบ", value=False)
        
        if add_noise_option:
            noise_factor = st.slider("ระดับเสียงรบกวน", min_value=0.001, max_value=0.1, value=0.01, step=0.001)
            noisy_waveform = add_noise(waveform, noise_factor)
            
            st.markdown('<p class="sub-header">เสียงที่มีเสียงรบกวน</p>', unsafe_allow_html=True)
            
            # แสดงกราฟเสียงที่มีเสียงรบกวน
            fig, ax = plt.subplots(figsize=(10, 2))
            ax.plot(noisy_waveform)
            ax.set_title("Noisy Waveform")
            ax.set_xlabel("Samples")
            ax.set_ylabel("Amplitude")
            ax.grid(True, alpha=0.3)
            st.pyplot(fig)
            
            # บันทึกเสียงที่มีเสียงรบกวน
            noisy_path = "./temp/noisy_audio.wav"
            sf.write(noisy_path, noisy_waveform, sample_rate)
            st.audio(noisy_path, format="audio/wav")
        else:
            noisy_waveform = waveform  # ใช้เสียงต้นฉบับถ้าไม่เพิ่มเสียงรบกวน
        
        # ตั้งค่าพารามิเตอร์ตัวกรอง
        st.markdown('<p class="sub-header">ปรับแต่งตัวกรองความถี่</p>', unsafe_allow_html=True)
        col1, col2, col3 = st.columns(3)
        
        with col1:
            lowcut = st.slider("ความถี่ต่ำสุด (Hz)", min_value=50, max_value=500, value=300, step=10)
        with col2:
            highcut = st.slider("ความถี่สูงสุด (Hz)", min_value=1000, max_value=8000, value=3000, step=100)
        with col3:
            filter_order = st.slider("ระดับความแรงของตัวกรอง", min_value=1, max_value=10, value=5, step=1)
        
        # ทำนาย (แยกเสียง)
        st.markdown('<p class="sub-header">กำลังแยกเสียง...</p>', unsafe_allow_html=True)
        
        # แสดง spinner ขณะประมวลผล
        with st.spinner('กำลังแยกเสียง...'):
            # ใช้ตัวกรองความถี่แบบช่องผ่าน
            separated_waveform = filter_speech(noisy_waveform, sample_rate, lowcut, highcut, filter_order)
            
            # Normalize
            separated_waveform = separated_waveform / np.max(np.abs(separated_waveform))
        
        st.markdown('<p class="sub-header">ผลลัพธ์การแยกเสียง</p>', unsafe_allow_html=True)
        st.markdown('<p class="success-text">แยกเสียงสำเร็จ! 🎉</p>', unsafe_allow_html=True)
        
        # แสดงกราฟเสียงที่แยกแล้ว
        fig, ax = plt.subplots(figsize=(10, 2))
        ax.plot(separated_waveform, color='green')
        ax.set_title("Separated Waveform")
        ax.set_xlabel("Samples")
        ax.set_ylabel("Amplitude")
        ax.grid(True, alpha=0.3)
        st.pyplot(fig)
        
        # บันทึกเสียงที่แยกแล้ว
        output_path = "./temp/separated_audio.wav"
        sf.write(output_path, separated_waveform, sample_rate)
        
        # เล่นเสียงที่แยกแล้ว
        st.markdown("### เสียงที่แยกแล้ว")
        st.audio(output_path, format="audio/wav")
        
        # เปรียบเทียบเสียงต้นฉบับและเสียงที่แยกแล้ว
        st.markdown("### เปรียบเทียบเสียงต้นฉบับและเสียงที่แยกแล้ว")
        
        # แสดงกราฟเปรียบเทียบ
        fig, ax = plt.subplots(2, 1, figsize=(10, 4), sharex=True)
        ax[0].plot(noisy_waveform)
        ax[0].set_title("Original/Noisy Waveform")
        ax[0].set_ylabel("Amplitude")
        ax[0].grid(True, alpha=0.3)
        
        ax[1].plot(separated_waveform, color='green')
        ax[1].set_title("Separated Waveform")
        ax[1].set_xlabel("Samples")
        ax[1].set_ylabel("Amplitude")
        ax[1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        st.pyplot(fig)
        
        # แสดงการวิเคราะห์ความถี่ (สเปกโตรแกรม)
        st.markdown("### การวิเคราะห์ความถี่")
        
        fig, ax = plt.subplots(2, 1, figsize=(10, 6))
        
        # สเปกโตรแกรมของเสียงที่มีเสียงรบกวน
        ax[0].specgram(noisy_waveform, Fs=sample_rate, NFFT=1024, noverlap=512, cmap='viridis')
        ax[0].set_title("Spectrogram of Noisy Audio")
        ax[0].set_ylabel("Frequency (Hz)")
        
        # สเปกโตรแกรมของเสียงที่แยกแล้ว
        ax[1].specgram(separated_waveform, Fs=sample_rate, NFFT=1024, noverlap=512, cmap='viridis')
        ax[1].set_title("Spectrogram of Separated Audio")
        ax[1].set_xlabel("Time (s)")
        ax[1].set_ylabel("Frequency (Hz)")
        
        plt.tight_layout()
        st.pyplot(fig)
        
        # ให้ดาวน์โหลดไฟล์เสียงที่แยกแล้ว
        with open(output_path, "rb") as file:
            st.download_button(
                label="⬇️ ดาวน์โหลดเสียงที่แยกแล้ว",
                data=file,
                file_name="separated_audio.wav",
                mime="audio/wav"
            )
    
    except Exception as e:
        st.error(f"เกิดข้อผิดพลาด: {e}")
    
    finally:
        # ลบไฟล์ชั่วคราว
        if os.path.exists(tmp_file_path):
            os.remove(tmp_file_path)

# ส่วนหลักของแอป
def main():
    # แสดงข้อมูลแอป
    st.sidebar.image("https://img.icons8.com/fluency/96/000000/audio-wave.png", width=100)
    st.sidebar.title("เกี่ยวกับแอปนี้")
    st.sidebar.info(
        """
        แอปนี้เป็นเดโมอย่างง่ายของการแยกเสียงพูดออกจากเสียงรบกวน โดยใช้เทคนิคการประมวลผลสัญญาณดิจิทัลพื้นฐาน
        
        **วิธีใช้:**
        1. อัปโหลดไฟล์เสียงที่มีเสียงรบกวน
        2. (ตัวเลือก) เพิ่มเสียงรบกวนสำหรับการทดสอบ
        3. ปรับแต่งพารามิเตอร์ของตัวกรองความถี่
        4. ดูผลลัพธ์และดาวน์โหลดเสียงที่แยกแล้ว
        
        """
    )
    

    # อัปโหลดไฟล์เสียง
    st.markdown('<p class="sub-header">อัปโหลดไฟล์เสียง</p>', unsafe_allow_html=True)
    uploaded_file = st.file_uploader("เลือกไฟล์เสียงที่มีเสียงรบกวน", type=["wav", "mp3", "ogg", "flac"])
    
    if uploaded_file is not None:
        st.markdown('<p class="success-text">อัปโหลดไฟล์สำเร็จ! กำลังเริ่มประมวลผล...</p>', unsafe_allow_html=True)
        process_audio(uploaded_file)

# รันแอป
if __name__ == "__main__":
    main()