import streamlit as st
import librosa
import numpy as np
import tempfile
from collections import defaultdict
from audiorecorder import audiorecorder

st.set_page_config(page_title="Bass Guitar Tablature Generator", layout="centered")

# Title
st.title("🎸 Real-Time Bass Guitar Tab Generator")
st.markdown("Upload or record your bass guitar audio and generate tab notation in standard tuning (E1, A1, D2, G2).")

# Note mapping (EADG - standard 4-string bass)
note_names = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
open_strings_bass = {
    4: "E1",
    3: "A1",
    2: "D2",
    1: "G2"
}

def generate_note_to_fretboard_map_bass():
    note_to_positions = defaultdict(list)
    for string, open_note in open_strings_bass.items():
        base_note = open_note[:-1]
        base_octave = int(open_note[-1])
        start_index = note_names.index(base_note)
        for fret in range(0, 21):
            note_index = (start_index + fret) % 12
            octave_shift = (start_index + fret) // 12
            note = note_names[note_index] + str(base_octave + octave_shift)
            note_to_positions[note].append((string, fret))
    return note_to_positions

note_to_frets_bass = generate_note_to_fretboard_map_bass()

# Audio input
uploaded_file = st.file_uploader("Upload a bass audio clip (WAV or MP3)", type=["wav", "mp3"])
audio = audiorecorder("🎧 Click to Record", "⏸ Stop Recording")

if audio.tobytes():
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
        tmp.write(audio.tobytes())
        uploaded_file = tmp.name

if uploaded_file:
    st.audio(uploaded_file)

    # Load and analyze
    y, sr = librosa.load(uploaded_file)
    pitches, magnitudes = librosa.piptrack(y=y, sr=sr)

    notes_detected = []
    for i in range(pitches.shape[1]):
        index = magnitudes[:, i].argmax()
        pitch = pitches[index, i]
        if pitch > 0:
            note_index = int(np.round(12 * np.log2(pitch / 440.0) + 69))
            note_name = note_names[note_index % 12] + str((note_index // 12) - 1)
            if note_name not in notes_detected:
                notes_detected.append(note_name)

    # Tablature rendering
    tab_lines = {1: "G|", 2: "D|", 3: "A|", 4: "E|"}

    for note in notes_detected:
        positions = note_to_frets_bass.get(note, [])
        if not positions:
            continue
        string, fret = sorted(positions, key=lambda x: x[1])[0]  # pick lowest fret
        for s in tab_lines:
            if s == string:
                tab_lines[s] += f"--{fret:2}"  # spacing for 2-digit
            else:
                tab_lines[s] += "----"

    st.subheader("🎼 Generated Tablature:")
    st.code("\n".join(tab_lines[s] for s in sorted(tab_lines)), language="text")