import streamlit as st
import numpy as np
import librosa
import soundfile as sf
import crepe
from fpdf import FPDF
import os

st.set_page_config(page_title="Guitar Tab Generator", layout="centered")
st.title("🎸 Guitar Audio to Tablature Generator")

# Define mapping from string to base note
standard_tuning = {
    6: 'E2',
    5: 'A2',
    4: 'D3',
    3: 'G3',
    2: 'B3',
    1: 'E4',
}

notes = ['C', 'C#', 'D', 'D#', 'E', 'F',
         'F#', 'G', 'G#', 'A', 'A#', 'B']

# Convert frequency to note
def freq_to_note(freq):
    A4 = 440
    if freq == 0:
        return "-"
    half_steps = round(12 * np.log2(freq / A4))
    note_index = (half_steps + 9) % 12
    octave = 4 + ((half_steps + 9) // 12)
    return notes[note_index] + str(octave)

# Generate fretboard dictionary
def generate_fretboard():
    fretboard = {}
    for string, base_note in standard_tuning.items():
        base_idx = notes.index(base_note[:-1])
        octave = int(base_note[-1])
        fretboard[string] = []
        for fret in range(0, 13):  # frets 0-12
            n = (base_idx + fret) % 12
            o = octave + (base_idx + fret) // 12
            fretboard[string].append(notes[n] + str(o))
    return fretboard

# Match note to fret
def note_to_fret(note, fretboard):
    for string in fretboard:
        for fret, n in enumerate(fretboard[string]):
            if n == note:
                return string, fret
    return None, None

# Display tab as strings
def generate_tab(notes_played):
    fretboard = generate_fretboard()
    tab_lines = {1: "e|", 2: "B|", 3: "G|", 4: "D|", 5: "A|", 6: "E|"}
    for note in notes_played:
        string, fret = note_to_fret(note, fretboard)
        for i in range(6, 0, -1):
            if i == string:
                tab_lines[i] += f"-{fret}-"
            else:
                tab_lines[i] += "---"
    return "\n".join([tab_lines[i] for i in range(1, 7)][::-1])

# Upload audio
uploaded_file = st.file_uploader("🎙️ Upload a mono guitar recording (WAV, 16kHz recommended)", type=["wav"])
if uploaded_file:
    with open("temp_audio.wav", "wb") as f:
        f.write(uploaded_file.read())

    y, sr = librosa.load("temp_audio.wav", sr=16000, mono=True)

    st.info("⏳ Detecting pitch...")
    _, frequency, confidence, _ = crepe.predict(y, sr, viterbi=True)
    thresholded = [(f, c) for f, c in zip(frequency, confidence) if c > 0.8]

    if not thresholded:
        st.error("❌ No strong pitch detected. Try a cleaner mono guitar track.")
    else:
        filtered_freqs = [f for f, c in thresholded[::10]]  # reduce to fewer notes
        notes_detected = [freq_to_note(f) for f in filtered_freqs]
        tab_result = generate_tab(notes_detected)

        st.subheader("🎼 Generated Guitar Tablature:")
        st.code(tab_result)

        # Export as PDF
        if st.button("📄 Export as PDF"):
            pdf = FPDF()
            pdf.add_page()
            pdf.set_font("Courier", size=12)
            for line in tab_result.split("\n"):
                pdf.cell(200, 10, txt=line, ln=True)
            pdf.output("Guitar_Tab.pdf")
            with open("Guitar_Tab.pdf", "rb") as f:
                st.download_button("Download PDF", f, "Guitar_Tab.pdf")