"""
Streamlit App - Metode Penilaian Kemiripan Bacaan Al-Qur'an
pada Pembelajaran DIROSA Menggunakan Representasi Audio WavLM dan DTW
"""

import os
import tempfile
import re
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import streamlit as st
import torch

from scoring import SimilarityScorer
from correlation_analysis import (
    run_correlation_analysis,
    plot_correlation_bar,
    plot_scatter_best_layer,
    plot_heatmap,
    plot_pairing_diagram
)

# ---------------------------------------------------------------------------
# Model caching
# ---------------------------------------------------------------------------

@st.cache_resource(show_spinner="Memuat model WavLM dan Pipeline ...")
def load_pipeline():
    """Load pipeline components once and cache them.
    
    Model WavLM otomatis berjalan di GPU (CUDA) jika tersedia,
    atau fallback ke CPU jika GPU tidak terdeteksi.
    Deteksi device dilakukan oleh WavLMEncoder secara otomatis.
    """
    scorer = SimilarityScorer(
        model_name="./wavlm-base-plus",
        distance_metric="cosine",
        sakoe_chiba_ratio=0.1,
        normalize_dtw=True
    )
    return scorer


def get_device_info() -> str:
    """Kembalikan string info device yang sedang digunakan (GPU/CPU)."""
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        return f"⚡ GPU — {gpu_name}"
    return "🖥️ CPU (GPU tidak terdeteksi / PyTorch tanpa CUDA)"


def run_pipeline(ref_path: str, test_path: str, use_vad: bool = True,
                 layer_indices: list = None):
    """Run full pipeline using SimilarityScorer."""
    scorer = load_pipeline()
    
    # Delegate the processing to SimilarityScorer's multi-layer handler
    detailed_data = scorer.compute_detailed_similarity(
        audio_path1=ref_path, 
        audio_path2=test_path, 
        use_vad=use_vad, 
        layer_indices=layer_indices
    )
    
    results = detailed_data["results"]
    waveforms = detailed_data["waveforms"]
    
    return (
        results, 
        waveforms["ref_raw"], 
        waveforms["ref_vad"], 
        waveforms.get("ref_normalized", waveforms["ref_vad"]),
        waveforms["test_raw"], 
        waveforms["test_vad"],
        waveforms.get("test_normalized", waveforms["test_vad"]),
    )



def plot_alignment(warping_path):
    """Create a simple DTW alignment line plot."""
    path = np.array(warping_path)
    fig, ax = plt.subplots(figsize=(6, 3.5))
    ax.plot(path[:, 0], path[:, 1], linewidth=0.8, color="black")
    ax.set_xlabel("Frame Referensi")
    ax.set_ylabel("Frame Peserta")
    ax.set_title("Alignment DTW")
    fig.tight_layout()
    return fig


def plot_dtw_heatmap(dtw_matrix: np.ndarray, warping_path):
    """Heatmap of the accumulated DTW cost matrix with the warping path overlay."""
    # Remove the padding row/col used during DP (index 0)
    matrix = dtw_matrix[1:, 1:]
    # Replace inf with max finite value for colour mapping
    finite_vals = matrix[np.isfinite(matrix)]
    if finite_vals.size > 0:
        matrix = np.where(np.isfinite(matrix), matrix, finite_vals.max())

    path = np.array(warping_path)
    fig, ax = plt.subplots(figsize=(6, 5))
    im = ax.imshow(matrix.T, origin="lower", aspect="auto", cmap="magma_r",
                   interpolation="nearest")
    ax.plot(path[:, 0], path[:, 1], color="cyan", linewidth=1.0, alpha=0.85)
    ax.set_xlabel("Frame Referensi")
    ax.set_ylabel("Frame Peserta")
    ax.set_title("DTW Cost Matrix & Warping Path")
    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="Accumulated Cost")
    fig.tight_layout()
    return fig


def plot_waveforms(
    raw: np.ndarray,
    vad: np.ndarray | None,
    normalized: np.ndarray | None,
    sr: int,
    title: str,
):
    """Plot two preprocessing stages vertically for thesis report.

    Stages shown:
      1. Sebelum Pre-processing (raw waveform)
      2. Setelah Pre-processing Lengkap (VAD + normalisasi amplitudo)

    All subplots share a fixed Y-axis [-1, 1] so amplitude differences
    before/after normalisation are visually obvious.
    """

    def _fmt_dur(n_samples: int) -> str:
        """Format duration string: 'Durasi: xx,xx detik'."""
        dur = n_samples / sr
        return f"Durasi: {dur:,.2f} detik".replace(",", "X").replace(".", ",").replace("X", ".")

    # Build stage list ---------------------------------------------------
    stages: list[tuple[np.ndarray, str, str]] = []

    # Stage 1 – Raw
    stages.append((
        raw,
        f"{title} — Sebelum Pre-processing  ({_fmt_dur(len(raw))})",
        "#4A90D9",
    ))

    # Stage 2 – After full pre-processing (normalised)
    if normalized is not None:
        stages.append((
            normalized,
            f"{title} — Setelah Pre-processing Lengkap  ({_fmt_dur(len(normalized))})",
            "#2ECC71",
        ))
    elif vad is not None:
        # Fallback: show VAD result as final stage if normalised is absent
        stages.append((
            vad,
            f"{title} — Setelah Pre-processing Lengkap  ({_fmt_dur(len(vad))})",
            "#2ECC71",
        ))

    n_plots = len(stages)
    fig, axes = plt.subplots(
        n_plots, 1,
        figsize=(8, 1.7 * n_plots + 0.6),
        sharex=False,
        sharey=True,
        constrained_layout=True,
    )
    if n_plots == 1:
        axes = [axes]

    for i, (data, label, color) in enumerate(stages):
        ax = axes[i]
        t = np.arange(len(data)) / sr
        ax.plot(t, data, linewidth=0.35, color=color)
        ax.set_title(label, fontsize=11, fontweight="bold", pad=6)
        ax.set_ylabel("Amplitudo", fontsize=10)
        ax.set_xlim(t[0], t[-1])
        ax.set_ylim(-1, 1)
        ax.tick_params(labelsize=9)
        ax.grid(True, linewidth=0.3, alpha=0.5)

    axes[-1].set_xlabel("Waktu (detik)", fontsize=10)

    return fig


def interpret_score(score: float) -> str:
    """Interpret normalized similarity score."""
    if score >= 80:
        return "Sangat Mirip (>80)"
    elif score >= 65:
        return "Mirip (65 - 80)"
    elif score >= 50:
        return "Cukup Mirip (50 - 65)"
    else:
        return "Kurang Mirip (<50)"


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

st.set_page_config(
    page_title="Penilaian Kemiripan Bacaan Al-Qur'an - DIROSA WavLM-DTW",
    layout="centered",
)

# Tampilkan info device di sidebar
with st.sidebar:
    st.markdown("### ⚙️ Info Sistem")
    st.info(f"**Device:** {get_device_info()}")
    st.caption("Model WavLM berjalan di GPU jika PyTorch CUDA tersedia, " \
               "atau fallback ke CPU secara otomatis.")

st.markdown(
    """
    <style>
    /* Mengubah max-width dari block container bawaan Streamlit */
    .block-container {
        max-width: 1000px !important;
    }
    </style>
    """,
    unsafe_allow_html=True
)

st.markdown(
    "<h3 style='text-align:center;'>"
    "Metode Penilaian Kemiripan Bacaan Al-Qur'an<br>"
    "pada Pembelajaran DIROSA &mdash; WavLM + DTW"
    "</h3>",
    unsafe_allow_html=True,
)

st.divider()

tab1, tab2, tab3 = st.tabs(["Single Processing", "Batch Processing (Folder)", "Analisis Korelasi (Overview)"])

with tab1:
    st.markdown("#### Uji Audio Individu")
    
    # Upload section
    col_ref, col_test = st.columns(2)
    
    with col_ref:
        st.subheader("Audio Referensi")
        ref_file = st.file_uploader(
            "Upload audio referensi",
            type=["wav"],
            key="ref",
            label_visibility="collapsed",
        )
    
    with col_test:
        st.subheader("Audio Peserta")
        test_file = st.file_uploader(
            "Upload audio peserta",
            type=["wav"],
            key="test",
            label_visibility="collapsed",
        )
    
    st.write("")  # spacer
    
    use_vad = st.checkbox("Aktifkan VAD (Voice Activity Detection)", value=True, key="vad_single",
                          help="Menghapus bagian hening di awal dan akhir audio sebelum diproses.")
    
    show_diagnostics = st.checkbox("Tampilkan diagnostik DTW", value=False,
                                   help="Menampilkan metrik internal DTW untuk analisis lanjutan.")

    st.write("")
    st.markdown("#### Parameter Model")
    select_all_layers = st.checkbox("Pilih Semua Layer (1-12)")
    
    if select_all_layers:
        sel_layers_single = list(range(1, 13))
        st.multiselect(
            "Pilih Layer WavLM",
            options=list(range(1, 13)),
            default=list(range(1, 13)),
            disabled=True,
            help="Semua layer telah dipilih."
        )
    else:
        sel_layers_single = st.multiselect(
            "Pilih Layer WavLM",
            options=list(range(1, 13)),
            default=[9, 10, 11, 12],
            help="Pilih satu atau lebih layer transformer WavLM (1-12) untuk diekstrak menjadi representasi khusus masing-masing layer."
        )
    
    st.write("")
    btn = st.button("Proses Penilaian Single", use_container_width=True)
    
    # ---------------------------------------------------------------------------
    # Processing & results
    # ---------------------------------------------------------------------------
    
    if btn:
        if ref_file is None or test_file is None:
            st.warning("Upload kedua file audio terlebih dahulu.")
        elif not sel_layers_single:
            st.warning("Pilih minimal satu layer WavLM.")
        else:
            # Save uploaded files to temp directory
            tmp_dir = tempfile.mkdtemp()
            ref_path = os.path.join(tmp_dir, "ref.wav")
            test_path = os.path.join(tmp_dir, "test.wav")
    
            with open(ref_path, "wb") as f:
                f.write(ref_file.getbuffer())
            with open(test_path, "wb") as f:
                f.write(test_file.getbuffer())
    
            with st.spinner("Memproses audio ..."):
                ((results, 
                 wf_ref_raw, wf_ref_vad, wf_ref_norm,
                 wf_test_raw, wf_test_vad, wf_test_norm)) = run_pipeline(
                    ref_path, test_path, use_vad,
                    layer_indices=sel_layers_single,
                )
    
            sr = 16_000  # pipeline target sample rate
    
            # --- Input Validation (Panjang Audio & Keheningan) ----------------
            ref_raw_np = wf_ref_raw.squeeze().numpy()
            test_raw_np = wf_test_raw.squeeze().numpy()
            ref_vad_np = wf_ref_vad.squeeze().numpy() if use_vad else None
            test_vad_np = wf_test_vad.squeeze().numpy() if use_vad else None
            ref_norm_np = wf_ref_norm.squeeze().numpy() if use_vad else None
            test_norm_np = wf_test_norm.squeeze().numpy() if use_vad else None

            ref_dur = len(ref_raw_np) / sr
            test_dur = len(test_raw_np) / sr
            if ref_dur < 0.5 or test_dur < 0.5:
                st.warning("**Peringatan Validation:** Salah satu atau kedua audio sangat pendek (< 0.5 detik). Hasil DTW mungkin menjadi kurang representatif.")
            
            silence_threshold = 0.005
            if np.max(np.abs(ref_raw_np)) < silence_threshold or np.max(np.abs(test_raw_np)) < silence_threshold:
                st.warning("**Peringatan Validation:** Terdeteksi audio yang hampir tidak bersuara (near-silent). VAD dan DTW kemungkinan kesulitan mencocokkan pola.")
            
            # --- Audio Preview ------------------------------------------------
            st.divider()
            st.subheader("Preview Audio & Tahapan Pre-Processing")
            st.caption(
                "Menampilkan dua tahap sinyal audio: "
                "**Sebelum Pre-processing** (audio asli) dan "
                "**Setelah Pre-processing Lengkap** (VAD + normalisasi amplitudo ke [-1, 1]).  \n"
                "Seluruh grafik menggunakan skala sumbu Y tetap **[-1, 1]** agar perbedaan "
                "amplitudo sebelum dan sesudah normalisasi terlihat jelas."
            )

            # --- Referensi ---
            st.markdown("##### Audio Referensi")
            fig_ref = plot_waveforms(ref_raw_np, ref_vad_np, ref_norm_np, sr, "Referensi")
            st.pyplot(fig_ref, use_container_width=True)
            plt.close(fig_ref)
            if use_vad and ref_vad_np is not None:
                dur_raw_ref = len(ref_raw_np) / sr
                dur_trim_ref = len(ref_vad_np) / sr
                delta_ref = dur_raw_ref - dur_trim_ref
                st.caption(
                    f"📐 **Efek Pre-processing (VAD):** "
                    f"Durasi asli = {dur_raw_ref:.2f} detik → "
                    f"Setelah VAD = {dur_trim_ref:.2f} detik "
                    f"(terpotong {delta_ref:.2f} detik)"
                )
            st.audio(ref_norm_np if use_vad else ref_raw_np, sample_rate=sr)

            # --- Peserta ---
            st.markdown("##### Audio Peserta")
            fig_test = plot_waveforms(test_raw_np, test_vad_np, test_norm_np, sr, "Peserta")
            st.pyplot(fig_test, use_container_width=True)
            plt.close(fig_test)
            if use_vad and test_vad_np is not None:
                dur_raw_test = len(test_raw_np) / sr
                dur_trim_test = len(test_vad_np) / sr
                delta_test = dur_raw_test - dur_trim_test
                st.caption(
                    f"📐 **Efek Pre-processing (VAD):** "
                    f"Durasi asli = {dur_raw_test:.2f} detik → "
                    f"Setelah VAD = {dur_trim_test:.2f} detik "
                    f"(terpotong {delta_test:.2f} detik)"
                )
            st.audio(test_norm_np if use_vad else test_raw_np, sample_rate=sr)
    
            # --- Aggregated Results -------------------------------------------
            st.divider()
            st.subheader("Ringkasan Hasil Penilaian (Agregasi)")
            
            scores = {layer: results[layer]["score"] for layer in sel_layers_single}
            mean_score = sum(scores.values()) / len(scores)
            best_layer = max(scores, key=scores.get)
            best_score = scores[best_layer]
            
            # --- Durasi Audio dalam ms ---
            ref_dur_ms = round(len(ref_raw_np) / sr * 1000)
            test_dur_ms = round(len(test_raw_np) / sr * 1000)
            ref_dur_vad_ms = round(len(ref_vad_np) / sr * 1000) if ref_vad_np is not None else ref_dur_ms
            test_dur_vad_ms = round(len(test_vad_np) / sr * 1000) if test_vad_np is not None else test_dur_ms

            col_dur1, col_dur2 = st.columns(2)
            with col_dur1:
                st.metric("Durasi Audio Referensi", f"{ref_dur_ms} ms",
                          delta=f"{ref_dur_vad_ms} ms setelah VAD" if use_vad else None,
                          delta_color="off")
            with col_dur2:
                st.metric("Durasi Audio Peserta", f"{test_dur_ms} ms",
                          delta=f"{test_dur_vad_ms} ms setelah VAD" if use_vad else None,
                          delta_color="off")

            col_agg1, col_agg2 = st.columns(2)
            with col_agg1:
                st.metric("Skor Rata-rata (Agregasi)", f"{mean_score:.2f} / 100")
                st.info(f"**Interpretasi:** {interpret_score(mean_score)}")
            with col_agg2:
                st.metric(f"Skor Tertinggi (Layer {best_layer})", f"{best_score:.2f} / 100")
               
                
            st.caption("Skor di atas berbasis kalibrasi sigmoid pada jarak _cosine_ DTW.")

            # --- Results Breakdown --------------------------------------------
            st.divider()
            st.subheader("Detail per Layer")
        
            layer_tabs = st.tabs([f"Layer {l}" for l in sel_layers_single])
            
            for idx, layer in enumerate(sel_layers_single):
                with layer_tabs[idx]:
                    layer_data = results[layer]
                    score = layer_data["score"]
                    warping_path = layer_data["warping_path"]
                    dtw_matrix = layer_data["dtw_matrix"]
                    diagnostics = layer_data["diagnostics"]

                    with st.container():
                        st.metric(f"Skor Kemiripan (Layer {layer})", f"{score:.2f} / 100")
            
                    st.write("")  # spacer
            
                    with st.container():
                        fig = plot_alignment(warping_path)
                        st.pyplot(fig, use_container_width=True)
                        plt.close(fig)
            
                    # --- DTW Diagnostics (optional) -----------------------------------
                    if show_diagnostics:
                        st.divider()
                        st.subheader(f"Diagnostik DTW - Layer {layer}")
            
                        d = diagnostics  # shorthand
                        c1, c2, c3 = st.columns(3)
                        c1.metric("Raw DTW Distance", f"{d['raw_dtw_distance']:.6f}")
                        c2.metric("Normalized Distance", f"{d['normalized_distance']:.6f}")
                        c3.metric("Path Length", d["path_length"])
            
                        c4, c5, c6 = st.columns(3)
                        c4.metric("Frames Referensi", d["num_frames_ref"])
                        c5.metric("Frames Peserta", d["num_frames_test"])
                        c6.metric("Sakoe-Chiba Ratio", d["sakoe_chiba_ratio"])
            
                        c7, c8, c9 = st.columns(3)
                        c7.metric("Durasi Ref (detik)", f"{d['ref_duration_sec']:.3f}")
                        c8.metric("Durasi Peserta (detik)", f"{d['test_duration_sec']:.3f}")
                        c9.metric("Rasio Durasi", f"{d['duration_ratio']:.4f}")
            
                        # Heatmap DTW
                        st.write("")
                        fig_hm = plot_dtw_heatmap(dtw_matrix, warping_path)
                        st.pyplot(fig_hm, use_container_width=True)
                        plt.close(fig_hm)


with tab2:
    st.markdown("#### Batch Processing (Dari Folder Lokal)")
    st.info("Fitur ini akan memproses semua audio di folder `audio peserta` dan membandingkannya dengan folder `audio referensi`.")
    
    use_vad_batch = st.checkbox("Aktifkan VAD", value=True, key="vad_batch")
    
    # Langsung pakai semua layer (1-12)
    sel_layers = list(range(1, 13))
    
    btn_batch = st.button("Jalankan Batch Processing", use_container_width=True)
    
    if btn_batch:
        if not sel_layers:
            st.warning("Pilih minimal satu layer untuk diproses.")
        else:
            peserta_dir = Path("audio peserta")
            referensi_dir = Path("audio referensi")
            
            if not peserta_dir.exists() or not referensi_dir.exists():
                st.error("Folder `audio peserta` atau `audio referensi` tidak ditemukan di direktori saat ini.")
            else:
                with st.spinner("Memproses seluruh audio dalam batch..."):
                    scorer = load_pipeline()
                    
                    pesertas = sorted(
                        peserta_dir.glob("peserta *"), 
                        key=lambda x: int(re.search(r"\d+", x.name).group()) if re.search(r"\d+", x.name) else 0
                    )
                    
                    rows = []
                    prog_bar = st.progress(0)
                    total_p = len(pesertas)
                    
                    for idx_p, p in enumerate(pesertas):
                        audios = [f for f in p.glob("*.wav") if re.search(r"\d+", f.name)]
                        audios = sorted(audios, key=lambda x: int(re.search(r"\d+", x.name).group()))
                        
                        if not audios:
                            continue
                            
                        for audio in audios:
                            ref_audio = referensi_dir / audio.name
                            if not ref_audio.exists(): continue
                            
                            detailed_data = scorer.compute_detailed_similarity(
                                audio_path1=str(ref_audio), 
                                audio_path2=str(audio), 
                                use_vad=use_vad_batch, 
                                layer_indices=sel_layers
                            )
                            
                            row = {
                                "Peserta": p.name,
                                "File": audio.name,
                            }
                            
                            for layer in sel_layers:
                                l_res = detailed_data["results"][layer]
                                row[f"Score L{layer}"] = round(l_res["score"], 2)
                                row[f"Dist L{layer}"] = round(l_res["dtw_distance"], 4)
                            
                            rows.append(row)
                            
                        prog_bar.progress((idx_p + 1) / total_p)
                
                if rows:
                    df = pd.DataFrame(rows)
                    st.success("Batch processing selesai!")
                    st.dataframe(df, use_container_width=True)
                    
                    csv = df.to_csv(index=False).encode("utf-8")
                    st.download_button(
                        label="Download Hasil CSV",
                        data=csv,
                        file_name="hasil_batch_multi_layer.csv",
                        mime="text/csv",
                        use_container_width=True
                    )
                else:
                    st.warning("Tidak ada data valid yang diproses.")


with tab3:
    st.markdown("#### Analisis Korelasi (Overview)")
    st.info("Visualisasi hubungan antara skor sistem (DTW) dan penilaian Ustadz (rating).")
    
    st.markdown("### Struktur Dataset Pasangan Frasa")
    csv_path = "dataset_final.csv"
    if os.path.exists(csv_path):
        df_pairing = pd.read_csv(csv_path)
        fig_pairing = plot_pairing_diagram(df_pairing)
        st.pyplot(fig_pairing, use_container_width=True)
        plt.close(fig_pairing)
        
    st.markdown("""
    **Penjelasan Singkat:**
    - Setiap peserta membaca frasa yang sama
    - Setiap frasa memiliki satu audio referensi
    - Sistem membandingkan pasangan audio pada frasa yang sama
    """)
    st.divider()
    
    
    if not os.path.exists(csv_path):
        st.warning(f"File {csv_path} tidak ditemukan. Jalankan grid_search atau buat dataset terlebih dahulu.")
    else:
        with st.spinner("Menjalankan analisis korelasi..."):
            df, df_results = run_correlation_analysis(csv_path)
            
            # Cari best layer (Spearman tertinggi)
            best_idx = df_results['spearman_rho'].idxmax()
            best_layer_name = df_results.loc[best_idx, 'layer']
            best_spearman = df_results.loc[best_idx, 'spearman_rho']
            
            st.subheader(f"Data Korelasi Per Layer")
            st.dataframe(df_results, use_container_width=True)
            
            st.markdown(f"**Layer Terbaik:** `{best_layer_name}` dengan korelasi Spearman **{best_spearman:.4f}**")
            
            st.subheader("Bar Chart: Spearman Rho")
            fig_bar = plot_correlation_bar(df_results)
            st.pyplot(fig_bar, use_container_width=True)
            plt.close(fig_bar)
            
            st.subheader(f"Scatter Plot: {best_layer_name} vs Rating")
            fig_scatter = plot_scatter_best_layer(df, best_layer_name)
            st.pyplot(fig_scatter, use_container_width=True)
            plt.close(fig_scatter)
            
            st.subheader("Heatmap Korelasi")
            fig_hm = plot_heatmap(df_results)
            st.pyplot(fig_hm, use_container_width=True)
            plt.close(fig_hm)