app.py

import os
import io
import gradio as gr
import torch
import numpy as np
import re
import pronouncing
import functools
from transformers import (
    AutoModelForAudioClassification,
    AutoFeatureExtractor,
    AutoTokenizer,
    pipeline,
    AutoModelForCausalLM,
    BitsAndBytesConfig
)
from huggingface_hub import login
from utils import (
    load_audio,
    extract_audio_duration,
    extract_mfcc_features,
    format_genre_results,
    ensure_cuda_availability
)
from emotionanalysis import MusicAnalyzer 
import librosa
from beat_analysis import BeatAnalyzer  # Import the BeatAnalyzer class

# Initialize beat analyzer
beat_analyzer = BeatAnalyzer()

# Login to Hugging Face Hub if token is provided
if "HF_TOKEN" in os.environ:
    login(token=os.environ["HF_TOKEN"])

# Constants
GENRE_MODEL_NAME = "dima806/music_genres_classification"
MUSIC_DETECTION_MODEL = "MIT/ast-finetuned-audioset-10-10-0.4593"
LLM_MODEL_NAME = "Qwen/Qwen3-32B"
SAMPLE_RATE = 22050  # Standard sample rate for audio processing

# Check CUDA availability (for informational purposes)
CUDA_AVAILABLE = ensure_cuda_availability()

# Load models at initialization time
print("Loading genre classification model...")
try:
    genre_feature_extractor = AutoFeatureExtractor.from_pretrained(GENRE_MODEL_NAME)
    genre_model = AutoModelForAudioClassification.from_pretrained(
        GENRE_MODEL_NAME,
        device_map="auto" if CUDA_AVAILABLE else None
    )
    # Create a convenience wrapper function with the same interface as before
    def get_genre_model():
        return genre_model, genre_feature_extractor
except Exception as e:
    print(f"Error loading genre model: {str(e)}")
    genre_model = None
    genre_feature_extractor = None

# Load LLM and tokenizer at initialization time
print("Loading Qwen LLM model with 4-bit quantization...")
try:
    # Configure 4-bit quantization for better performance
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True
    )
    
    llm_tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_NAME)
    llm_model = AutoModelForCausalLM.from_pretrained(
        LLM_MODEL_NAME,
        quantization_config=quantization_config,
        device_map="auto",
        trust_remote_code=True,
        torch_dtype=torch.float16,
        use_cache=True
    )
except Exception as e:
    print(f"Error loading LLM model: {str(e)}")
    llm_tokenizer = None
    llm_model = None

# Create music analyzer instance
music_analyzer = MusicAnalyzer()

# Define emotion and theme lexicons for content analysis
emotion_lexicons = {
    "happy": ["joy", "happy", "smile", "laugh", "light", "bright", "sun", "dance", "celebrate", "glow", "warm"],
    "sad": ["cry", "tear", "pain", "loss", "grief", "dark", "alone", "miss", "gone", "sorrow", "heart", "break"],
    "calm": ["peace", "quiet", "still", "gentle", "soft", "slow", "breath", "serene", "tranquil", "relax", "flow"],
    "energetic": ["run", "fast", "beat", "pulse", "jump", "fire", "alive", "spark", "rush", "wild", "free"],
    "tense": ["fear", "worry", "wait", "edge", "grip", "tight", "storm", "break", "shadow", "threat", "doubt"],
    "nostalgic": ["memory", "remember", "past", "time", "ago", "once", "childhood", "return", "old", "familiar", "home"],
    "reflective": ["think", "ponder", "wonder", "question", "search", "mind", "deep", "self", "mirror", "path", "journey"],
    "triumphant": ["win", "rise", "stand", "overcome", "above", "victory", "summit", "conquer", "champion", "succeed"],
    "yearning": ["want", "need", "desire", "reach", "seek", "dream", "hope", "wish", "long", "hunger", "thirst"],
    "peaceful": ["calm", "rest", "still", "quiet", "harmony", "balance", "ease", "gentle", "soft", "float", "drift"]
}

theme_lexicons = {
    "love": ["love", "heart", "touch", "together", "hold", "kiss", "embrace", "feel", "close", "intimate", "passion"],
    "loss": ["gone", "away", "empty", "missing", "leave", "without", "never", "forever", "lost", "memory", "shadow"],
    "freedom": ["free", "fly", "open", "release", "escape", "chain", "break", "boundless", "space", "breathe", "wings"],
    "triumph": ["victory", "overcome", "win", "rise", "mountain", "climb", "top", "struggle", "strength", "succeed"],
    "reflection": ["mirror", "water", "see", "self", "face", "look", "inside", "truth", "reality", "soul", "mind"],
    "journey": ["road", "path", "walk", "step", "travel", "distance", "far", "way", "wander", "search", "find"],
    "time": ["clock", "moment", "second", "hour", "pass", "wait", "forever", "instant", "eternity", "memory", "future"],
    "conflict": ["fight", "battle", "against", "oppose", "between", "war", "struggle", "clash", "resist", "enemy"],
    "nature": ["earth", "wind", "fire", "water", "sky", "tree", "flower", "mountain", "river", "ocean", "stars"],
    "change": ["transform", "become", "different", "shift", "turn", "evolve", "grow", "new", "begin", "end", "cycle"]
}

# Process uploaded audio file
def process_audio(audio_file):
    if audio_file is None:
        return "No audio file provided", None, None, None, None, None, None, None
    
    try:
        # Load and analyze audio
        y, sr = load_audio(audio_file, sr=SAMPLE_RATE)
        
        # Basic audio information
        duration = extract_audio_duration(y, sr)
        
        # Analyze music with MusicAnalyzer
        music_analysis = music_analyzer.analyze_music(audio_file)
        
        # Extract key information
        tempo = music_analysis["rhythm_analysis"]["tempo"]
        emotion = music_analysis["emotion_analysis"]["primary_emotion"]
        theme = music_analysis["theme_analysis"]["primary_theme"]
        
        # Use genre classification directly instead of pipeline
        if genre_model is not None and genre_feature_extractor is not None:
            # Resample audio to 16000 Hz for the genre model
            y_16k = librosa.resample(y, orig_sr=sr, target_sr=16000)
            
            # Extract features
            inputs = genre_feature_extractor(
                y_16k, 
                sampling_rate=16000, 
                return_tensors="pt"
            ).to(genre_model.device)
            
            # Classify genre
            with torch.no_grad():
                outputs = genre_model(**inputs)
                logits = outputs.logits
                probs = torch.nn.functional.softmax(logits, dim=-1)
                
            # Get top genres
            values, indices = torch.topk(probs[0], k=5)
            top_genres = [(genre_model.config.id2label[idx.item()], val.item()) for val, idx in zip(values, indices)]
        else:
            # Fallback if model loading failed
            top_genres = [("Unknown", 1.0)]
        
        # Format genre results for display
        genre_results_text = format_genre_results(top_genres)
        primary_genre = top_genres[0][0]
        
        # Override time signature for pop and disco genres to always be 4/4
        if any(genre.lower() in primary_genre.lower() for genre in ['pop', 'disco']):
            music_analysis["rhythm_analysis"]["estimated_time_signature"] = "4/4"
            time_signature = "4/4"
        else:
            # Use detected time signature for other genres
            time_signature = music_analysis["rhythm_analysis"]["estimated_time_signature"]
            
            # Ensure time signature is one of the supported ones (4/4, 3/4, 6/8)
            if time_signature not in ["4/4", "3/4", "6/8"]:
                time_signature = "4/4"  # Default to 4/4 if unsupported
                music_analysis["rhythm_analysis"]["estimated_time_signature"] = time_signature
        
        # Analyze beat patterns and create lyrics template using the time signature
        beat_analysis = beat_analyzer.analyze_beat_pattern(audio_file, time_signature=time_signature)
        lyric_templates = beat_analyzer.create_lyric_template(beat_analysis)
        
        # Store these in the music_analysis dict for use in lyrics generation
        music_analysis["beat_analysis"] = beat_analysis
        music_analysis["lyric_templates"] = lyric_templates
        
        # Prepare analysis summary
        analysis_summary = f"""
### Music Analysis Results

**Duration:** {duration:.2f} seconds
**Tempo:** {tempo:.1f} BPM
**Time Signature:** {time_signature}
**Key:** {music_analysis["tonal_analysis"]["key"]} {music_analysis["tonal_analysis"]["mode"]}
**Primary Emotion:** {emotion}
**Primary Theme:** {theme}
**Top Genre:** {primary_genre}

{genre_results_text}
"""

        # Add beat analysis summary
        if lyric_templates:
            analysis_summary += f"""
### Beat Analysis

**Total Phrases:** {len(lyric_templates)}
**Average Beats Per Phrase:** {np.mean([t['num_beats'] for t in lyric_templates]):.1f}
**Beat Pattern Examples:** 
- Phrase 1: {lyric_templates[0]['stress_pattern'] if lyric_templates else 'N/A'}
- Phrase 2: {lyric_templates[1]['stress_pattern'] if len(lyric_templates) > 1 else 'N/A'}
"""
        
        # Check if genre is supported for lyrics generation
        # Use the supported_genres list from BeatAnalyzer
        genre_supported = any(genre.lower() in primary_genre.lower() for genre in beat_analyzer.supported_genres)
        
        # Generate lyrics only for supported genres
        if genre_supported:
            lyrics = generate_lyrics(music_analysis, primary_genre, duration)
            beat_match_analysis = analyze_lyrics_rhythm_match(lyrics, lyric_templates, primary_genre, emotion, theme)
        else:
            supported_genres_str = ", ".join([genre.capitalize() for genre in beat_analyzer.supported_genres])
            lyrics = f"Lyrics generation is only supported for the following genres: {supported_genres_str}.\n\nDetected genre '{primary_genre}' doesn't have strong syllable-to-beat patterns required for our lyric generation algorithm."
            beat_match_analysis = "Lyrics generation not available for this genre."
        
        return analysis_summary, lyrics, tempo, time_signature, emotion, theme, primary_genre, beat_match_analysis
    
    except Exception as e:
        error_msg = f"Error processing audio: {str(e)}"
        print(error_msg)
        return error_msg, None, None, None, None, None, None, None

def generate_lyrics(music_analysis, genre, duration):
    try:
        # Extract meaningful information for context
        tempo = music_analysis["rhythm_analysis"]["tempo"]
        key = music_analysis["tonal_analysis"]["key"]
        mode = music_analysis["tonal_analysis"]["mode"]
        emotion = music_analysis["emotion_analysis"]["primary_emotion"]
        theme = music_analysis["theme_analysis"]["primary_theme"]
        
        # Get beat analysis and templates
        lyric_templates = music_analysis.get("lyric_templates", [])
        
        # Verify LLM is loaded
        if llm_model is None or llm_tokenizer is None:
            return "Error: LLM model not properly loaded"

        # If no templates, fall back to original method
        if not lyric_templates:
            # Simplified prompt
            prompt = f"""Write song lyrics for a {genre} song in {key} {mode} with tempo {tempo} BPM. The emotion is {emotion} and theme is {theme}.

ONLY WRITE THE ACTUAL LYRICS. NO EXPLANATIONS OR META-TEXT.
"""
        else:
            # Create phrase examples
            num_phrases = len(lyric_templates)
            
            # Calculate the typical syllable range for this genre
            if num_phrases > 0:
                # Get max syllables per line from templates
                max_syllables = max([t.get('max_expected', 7) for t in lyric_templates]) if lyric_templates[0].get('max_expected') else 7
                min_syllables = min([t.get('min_expected', 2) for t in lyric_templates]) if lyric_templates[0].get('min_expected') else 2
                avg_syllables = (min_syllables + max_syllables) // 2
            else:
                min_syllables = 2
                max_syllables = 7
                avg_syllables = 4
            
            # Create a more direct prompt with examples and specific syllable count guidance
            prompt = f"""Write song lyrics for a {genre} song that truly captures the emotional essence of "{emotion}" and explores the theme of "{theme}". The song is in {key} {mode} with tempo {tempo} BPM.

I need EXACTLY {num_phrases} lines of lyrics - one line for each musical phrase.

YOUR TOP PRIORITIES (in order):
1. EXPRESS THE EMOTION: "{emotion}" should be felt through your word choices
2. DEVELOP THE THEME: "{theme}" should be clearly represented
3. CONNECT YOUR LINES: spread complete thoughts across 2-3 consecutive lines
4. KEEP LINES SHORT: {min_syllables}-{max_syllables} syllables per line (aim for {avg_syllables})

ADDITIONAL REQUIREMENTS:
- Create original lyrics that reflect this specific emotion and theme
- Let sentence clauses flow naturally across line breaks
- Use vivid imagery and sensory details related to the emotion
- Each line should contribute to the overall theme
- Don't copy the example structures - be creative and unique
- Use simple, concise words that evoke strong emotions

AVOID:
- Generic phrases that could apply to any song
- Copying patterns from the examples below
- Complete, independent thoughts on each line
- Abstract concepts without concrete imagery

FORMAT:
- Plain text, one line per musical phrase
- No annotations, explanations, or labels

Here's a simplified structural example of connecting thoughts across lines:

Line 1 (introduces an image)
Line 2 (extends the image)
Line 3 (completes the thought)

Line 4 (starts a new thought)
Line 5 (continues it)
And so on...

Remember: YOUR LYRICS SHOULD DEEPLY EXPRESS "{emotion}" AND EXPLORE "{theme}" - make every word count toward these goals.
"""

        # Generate lyrics using the LLM model
        messages = [
            {"role": "user", "content": prompt}
        ]
        
        # Apply chat template
        text = llm_tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        
        # Tokenize and move to model device
        model_inputs = llm_tokenizer([text], return_tensors="pt").to(llm_model.device)
        
        # Generate with optimized parameters
        generated_ids = llm_model.generate(
            **model_inputs,
            max_new_tokens=1024,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.2,
            pad_token_id=llm_tokenizer.eos_token_id
        )
        
        # Decode the output
        output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
        lyrics = llm_tokenizer.decode(output_ids, skip_special_tokens=True).strip()
        
        # ULTRA AGGRESSIVE CLEANING - COMPLETELY REVISED
        # ------------------------------------------------
        
        # 1. First, look for any standard dividers that might separate thinking from lyrics
        divider_patterns = [
            r'Here are the lyrics:',
            r'Here is my song:',
            r'The lyrics:',
            r'My lyrics:',
            r'Song lyrics:',
            r'\*\*\*+',
            r'===+',
            r'---+',
            r'```',
            r'Lyrics:'
        ]
        
        for pattern in divider_patterns:
            matches = re.finditer(pattern, lyrics, re.IGNORECASE)
            for match in matches:
                # Keep only content after the divider
                lyrics = lyrics[match.end():].strip()
        
        # 2. Remove thinking tags completely before splitting into lines
        lyrics = re.sub(r'<think>.*?</think>', '', lyrics, flags=re.DOTALL)
        lyrics = re.sub(r'\[thinking\].*?\[/thinking\]', '', lyrics, flags=re.DOTALL)
        lyrics = re.sub(r'<think>', '', lyrics, flags=re.DOTALL)
        lyrics = re.sub(r'</think>', '', lyrics, flags=re.DOTALL)
        lyrics = re.sub(r'\[thinking\]', '', lyrics, flags=re.DOTALL)
        lyrics = re.sub(r'\[/thinking\]', '', lyrics, flags=re.DOTALL)
        
        # 3. Split text into lines for aggressive line-by-line filtering
        lines = lyrics.strip().split('\n')
        clean_lines = []
        
        # 4. Define comprehensive patterns for non-lyrical content
        non_lyric_patterns = [
            # Meta-commentary
            r'^(note|thinking|thoughts|let me|i will|i am going|i would|i can|i need to|i have to|i should|let\'s|here|now)',
            r'^(first|second|third|next|finally|importantly|remember|so|ok|okay|as requested|as asked|considering)',
            # Explanations
            r'syllable[s]?|phrase|rhythm|beats?|tempo|bpm|instruction|follow|alignment|match|corresponding',
            r'verses?|chorus|bridge|section|stanza|part|template|format|pattern|example',
            r'requirements?|guidelines?|song structure|stressed|unstressed',
            # Technical language
            r'generated|output|result|provide|create|write|draft|version',
            # Annotations and numbering
            r'^line \d+|^\d+[\.\):]|^\[\w+\]|^[\*\-\+] ',
            # Questions or analytical statements
            r'\?$|analysis|evaluate|review|check|ensure',
            # Instruction-like statements
            r'make sure|please note|important|notice|pay attention'
        ]
        
        # 5. Identify which lines are likely actual lyrics vs non-lyrics
        for line in lines:
            line = line.strip()
            
            # Skip empty lines or lines with just spaces/tabs
            if not line or line.isspace():
                continue
            
            # Skip lines that match any non-lyric pattern
            should_skip = False
            for pattern in non_lyric_patterns:
                if re.search(pattern, line.lower()):
                    should_skip = True
                    break
            
            if should_skip:
                continue
            
            # Skip section headers
            if (line.startswith('[') and ']' in line) or (line.startswith('(') and ')' in line and len(line) < 20):
                continue
            
            # Skip lines that look like annotations (not prose-like)
            if ':' in line and not any(word in line.lower() for word in ['like', 'when', 'where', 'how', 'why', 'what']):
                if len(line.split(':')[0]) < 15:  # Short prefixes followed by colon are likely annotations
                    continue
            
            # Skip very short lines that aren't likely to be lyrics (unless it's just a few words which could be valid)
            if len(line) < 3:
                continue
            
            # Skip lines that are numbered or bulleted
            if re.match(r'^\d+\.|\(#\d+\)|\d+\)', line):
                continue
            
            # Skip markdown-style emphasis or headers
            if re.match(r'^#{1,6} |^\*\*|^__', line):
                continue
            
            # Skip lines with think tags
            if '<think>' in line.lower() or '</think>' in line.lower() or '[thinking]' in line.lower() or '[/thinking]' in line.lower():
                continue
                
            # Add this line as it passed all filters
            clean_lines.append(line)
        
        # 6. Additional block-level filters for common patterns
        # Check beginning of lyrics for common prefixes
        if clean_lines and any(clean_lines[0].lower().startswith(prefix) for prefix in 
                            ['here are', 'these are', 'below are', 'following are']):
            clean_lines = clean_lines[1:]  # Skip the first line
        
        # 7. Process blocks of lines to detect explanation blocks
        if len(clean_lines) > 3:
            # Check for explanation blocks at the beginning
            first_three = ' '.join(clean_lines[:3]).lower()
            if any(term in first_three for term in ['i will', 'i have created', 'i\'ll provide', 'i\'ll write']):
                # This looks like an explanation, skip the first few lines
                start_idx = 0
                for i, line in enumerate(clean_lines):
                    if i >= 3 and not any(term in line.lower() for term in ['i will', 'created', 'write', 'provide']):
                        start_idx = i
                        break
                clean_lines = clean_lines[start_idx:]
            
            # Check for explanation blocks at the end
            last_three = ' '.join(clean_lines[-3:]).lower()
            if any(term in last_three for term in ['hope this', 'these lyrics', 'as you can see', 'this song', 'i have']):
                # This looks like an explanation at the end, truncate
                end_idx = len(clean_lines)
                for i in range(len(clean_lines) - 1, max(0, len(clean_lines) - 4), -1):
                    if i < len(clean_lines) and not any(term in clean_lines[i].lower() for term in 
                                                    ['hope', 'these lyrics', 'as you can see', 'this song']):
                        end_idx = i + 1
                        break
                clean_lines = clean_lines[:end_idx]
        
        # 8. Cleanup - Remove remaining annotations or thinking
        for i in range(len(clean_lines)):
            # Remove trailing thoughts/annotations
            clean_lines[i] = re.sub(r'\s+//.*$', '', clean_lines[i])
            clean_lines[i] = re.sub(r'\s+\(.*?\)$', '', clean_lines[i])
            
            # Remove thinking tags completely
            clean_lines[i] = re.sub(r'<think>.*?</think>', '', clean_lines[i], flags=re.DOTALL)
            clean_lines[i] = re.sub(r'\[thinking\].*?\[/thinking\]', '', clean_lines[i], flags=re.DOTALL)
            clean_lines[i] = re.sub(r'<think>', '', clean_lines[i])
            clean_lines[i] = re.sub(r'</think>', '', clean_lines[i])
            clean_lines[i] = re.sub(r'\[thinking\]', '', clean_lines[i])
            clean_lines[i] = re.sub(r'\[/thinking\]', '', clean_lines[i])
            
            # Remove syllable count annotations
            clean_lines[i] = re.sub(r'\s*\(\d+\s*syllables?\)', '', clean_lines[i])
        
        # 9. Filter out any remaining empty lines after tag removal
        clean_lines = [line for line in clean_lines if line.strip() and not line.isspace()]
        
        # 10. If we have lyric templates, ensure we have the correct number of lines
        if lyric_templates:
            num_required = len(lyric_templates)
            
            # If we have too many lines, keep just the best ones
            if len(clean_lines) > num_required:
                # Keep the first num_required lines
                clean_lines = clean_lines[:num_required]
            
            # If we don't have enough lines, generate placeholders that fit the syllable count
            while len(clean_lines) < num_required:
                i = len(clean_lines)
                if i < len(lyric_templates):
                    template = lyric_templates[i]
                    target_syllables = min(max_syllables, (template.get('min_expected', 2) + template.get('max_expected', 7)) // 2)
                    
                    # Create a diverse set of placeholders that match the theme/emotion
                    placeholders = {
                        # 2-3 syllables
                        2: [
                            "Night falls", 
                            "Time stops", 
                            "Hearts beat", 
                            "Rain falls", 
                            "Stars shine"
                        ],
                        # 3-4 syllables
                        3: [
                            "Empty chair", 
                            "Shadows dance", 
                            "Whispers fade", 
                            "Memories", 
                            "Silent room"
                        ],
                        # 4-5 syllables
                        4: [
                            "Moonlight shimmers", 
                            "Echoes of time", 
                            "Footsteps fading", 
                            "Memories drift", 
                            "Silence speaks loud"
                        ],
                        # 5-6 syllables
                        5: [
                            "Walking in the rain", 
                            "Whispers in the dark", 
                            "Echoes of your voice", 
                            "Traces left behind", 
                            "Time moves ever on"
                        ],
                        # 6-7 syllables
                        6: [
                            "Dancing in the moonlight", 
                            "Shadows play on the wall", 
                            "Memories fade to silence", 
                            "Moments lost in the wind", 
                            "Whispers of a better time"
                        ]
                    }
                    
                    # Get the closest matching syllable group
                    closest_group = min(placeholders.keys(), key=lambda k: abs(k - target_syllables))
                    
                    # Choose a placeholder that hasn't been used yet
                    available_placeholders = [p for p in placeholders[closest_group] 
                                             if p not in clean_lines]
                    
                    if available_placeholders:
                        placeholder = available_placeholders[i % len(available_placeholders)]
                    else:
                        # If we've used all placeholders in this group, create a custom one
                        if emotion.lower() in ["sad", "nostalgic", "calm"]:
                            placeholder = f"Memories of {emotion}"
                        elif emotion.lower() in ["happy", "energetic"]:
                            placeholder = f"Dancing through {emotion}"
                        else:
                            placeholder = f"Feeling {emotion} now"
                else:
                    placeholder = "Silence speaks volumes"
                    
                clean_lines.append(placeholder)
        
        # Assemble final lyrics
        final_lyrics = '\n'.join(clean_lines)
        
        # 11. Final sanity check - if we have nothing or garbage, return an error
        if not final_lyrics or len(final_lyrics) < 10:
            return "The model generated only thinking content but no actual lyrics. Please try again."
        
        return final_lyrics
    
    except Exception as e:
        error_msg = f"Error generating lyrics: {str(e)}"
        print(error_msg)
        return error_msg

def analyze_lyrics_rhythm_match(lyrics, lyric_templates, genre="pop", emotion="reflective", theme="journey"):
    """Analyze how well the generated lyrics match the beat patterns and syllable requirements"""
    if not lyric_templates or not lyrics:
        return "No beat templates or lyrics available for analysis."
    
    # Split lyrics into lines
    lines = lyrics.strip().split('\n')
    lines = [line for line in lines if line.strip()]  # Remove empty lines
    
    # Prepare analysis result
    result = "### Beat & Syllable Match Analysis\n\n"
    result += "| Line | Syllables | Target Range | Match | Stress Pattern |\n"
    result += "| ---- | --------- | ------------ | ----- | -------------- |\n"
    
    # Maximum number of lines to analyze (either all lines or all templates)
    line_count = min(len(lines), len(lyric_templates))
    
    # Track overall match statistics
    total_matches = 0
    total_range_matches = 0
    total_stress_matches = 0
    total_stress_percentage = 0
    total_ideal_matches = 0
    
    for i in range(line_count):
        line = lines[i]
        template = lyric_templates[i]
        
        # Check match between line and template with genre awareness
        check_result = beat_analyzer.check_syllable_stress_match(line, template, genre)
        
        # Get match symbols
        if check_result["close_to_ideal"]:
            syllable_match = "✓"  # Ideal or very close
        elif check_result["within_range"]:
            syllable_match = "✓*"  # Within range but not ideal
        else:
            syllable_match = "✗"  # Outside range
            
        stress_match = "✓" if check_result["stress_matches"] else f"{int(check_result['stress_match_percentage']*100)}%"
        
        # Update stats
        if check_result["close_to_ideal"]:
            total_matches += 1
            total_ideal_matches += 1
        elif check_result["within_range"]:
            total_range_matches += 1
            
        if check_result["stress_matches"]:
            total_stress_matches += 1
        total_stress_percentage += check_result["stress_match_percentage"]
        
        # Create visual representation of the stress pattern
        stress_visual = ""
        for char in template['stress_pattern']:
            if char == "S":
                stress_visual += "X"  # Strong
            elif char == "M":
                stress_visual += "x"  # Medium
            else:
                stress_visual += "."  # Weak
        
        # Add line to results table
        result += f"| {i+1} | {check_result['syllable_count']} | {check_result['min_expected']}-{check_result['max_expected']} | {syllable_match} | {stress_visual} |\n"
    
    # Add summary statistics
    if line_count > 0:
        exact_match_rate = (total_matches / line_count) * 100
        range_match_rate = ((total_matches + total_range_matches) / line_count) * 100
        ideal_match_rate = (total_ideal_matches / line_count) * 100
        stress_match_rate = (total_stress_matches / line_count) * 100
        avg_stress_percentage = (total_stress_percentage / line_count) * 100
        
        result += f"\n**Summary:**\n"
        result += f"- Ideal or near-ideal syllable match rate: {exact_match_rate:.1f}%\n"
        result += f"- Genre-appropriate syllable range match rate: {range_match_rate:.1f}%\n"
        result += f"- Perfect stress pattern match rate: {stress_match_rate:.1f}%\n"
        result += f"- Average stress pattern accuracy: {avg_stress_percentage:.1f}%\n"
        result += f"- Overall rhythmic accuracy: {((range_match_rate + avg_stress_percentage) / 2):.1f}%\n"
        
        # Analyze sentence flow across lines
        sentence_flow_analysis = analyze_sentence_flow(lines)
        result += f"\n**Sentence Flow Analysis:**\n"
        result += f"- Connected thought groups: {sentence_flow_analysis['connected_groups']} detected\n"
        result += f"- Average lines per thought: {sentence_flow_analysis['avg_lines_per_group']:.1f}\n"
        result += f"- Flow quality: {sentence_flow_analysis['flow_quality']}\n"
        
        # Analyze theme and emotion expression
        content_analysis = analyze_theme_emotion_expression(lyrics, theme, emotion)
        result += f"\n**Theme & Emotion Expression:**\n"
        result += f"- Emotion ({emotion}) expression: {content_analysis['emotion_score']:.1f}% ({content_analysis['emotion_words_found']} words)\n"
        result += f"- Theme ({theme}) development: {content_analysis['theme_score']:.1f}% ({content_analysis['theme_words_found']} words)\n"
        result += f"- Overall expression quality: {content_analysis['expression_quality']}\n"
        
        # Add guidance on ideal distribution for syllables and sentence flow
        result += f"\n**Improvement Recommendations:**\n"
        
        # Syllable recommendations
        if range_match_rate < 70:
            result += f"- **Syllable count:** Aim for {min([t.get('min_expected', 3) for t in lyric_templates])}-{max([t.get('max_expected', 7) for t in lyric_templates])} syllables per line\n"
        
        # Flow recommendations
        if sentence_flow_analysis['connected_groups'] < len(lines) / 5:
            result += f"- **Line connections:** Break complete thoughts across 2-3 lines using conjunctions and prepositions\n"
            result += f"- **Flow techniques:** Start lines with connecting words like 'as', 'when', 'while', 'through'\n"
        
        # Theme/emotion recommendations
        if content_analysis['emotion_score'] < 20:
            result += f"- **Emotion expression:** Use more words that evoke '{emotion}' feelings (e.g., {', '.join(emotion_lexicons.get(emotion.lower(), ['expressive words'])[:3])})\n"
        
        if content_analysis['theme_score'] < 20:
            result += f"- **Theme development:** Incorporate more '{theme}' imagery and concepts (e.g., {', '.join(theme_lexicons.get(theme.lower(), ['thematic words'])[:3])})\n"
        
        # General recommendations
        result += f"- **Originality:** Avoid generic phrases and create specific, vivid imagery\n"
        result += f"- **Sensory details:** Include concrete details that can be seen, heard, or felt\n"
        
        # Add genre-specific notes
        result += f"\n**Genre Notes ({genre}):**\n"
        
        # Add appropriate genre notes based on genre
        if genre.lower() == "pop":
            result += "- Pop lyrics work well with thoughts spanning 2-3 musical phrases\n"
            result += "- Create flow by connecting lines with transitions like 'as', 'when', 'through'\n"
        elif genre.lower() == "rock":
            result += "- Rock lyrics benefit from short phrases that build into complete thoughts\n"
            result += "- Use line breaks strategically to emphasize key words\n"
        elif genre.lower() == "country":
            result += "- Country lyrics tell stories that flow naturally across multiple lines\n"
            result += "- Connect narrative elements across phrases for authentic storytelling\n"
        elif genre.lower() == "disco":
            result += "- Disco lyrics work well with phrases that create rhythmic momentum\n"
            result += "- Use line transitions that maintain energy and flow\n"
        elif genre.lower() == "metal":
            result += "- Metal lyrics can create intensity by breaking phrases at dramatic points\n"
            result += "- Connect lines to build tension and release across measures\n"
        else:
            result += "- This genre works well with connected thoughts across multiple lines\n"
            result += "- Aim for natural speech flow rather than complete thoughts per line\n"
    
    return result

def analyze_sentence_flow(lines):
    """Analyze how well the lyrics create sentence flow across multiple lines"""
    if not lines or len(lines) < 2:
        return {
            "connected_groups": 0,
            "avg_lines_per_group": 0,
            "flow_quality": "Insufficient lines to analyze"
        }
    
    # Simplified analysis looking for grammatical clues of sentence continuation
    continuation_starters = [
        'and', 'but', 'or', 'nor', 'for', 'yet', 'so',  # Coordinating conjunctions
        'as', 'when', 'while', 'before', 'after', 'since', 'until', 'because', 'although', 'though',  # Subordinating conjunctions
        'with', 'without', 'through', 'throughout', 'beyond', 'beneath', 'under', 'over', 'into', 'onto',  # Prepositions
        'to', 'from', 'by', 'at', 'in', 'on', 'of',  # Common prepositions
        'where', 'how', 'who', 'whom', 'whose', 'which', 'that',  # Relative pronouns
        'if', 'then',  # Conditional connectors
    ]
    
    # Check for lines that likely continue a thought from previous line
    connected_lines = []
    potential_groups = []
    current_group = [0]  # Start with first line
    
    for i in range(1, len(lines)):
        # Check if line starts with a continuation word
        words = lines[i].lower().split()
        
        # Empty line or no words
        if not words:
            if len(current_group) > 1:  # Only consider groups of 2+ lines
                potential_groups.append(current_group.copy())
            current_group = [i]
            continue
            
        # Check first word for continuation clues
        first_word = words[0].strip(',.!?;:')
        if first_word in continuation_starters:
            connected_lines.append(i)
            current_group.append(i)
        # Check for absence of capitalization as continuation clue
        elif not first_word[0].isupper() and first_word[0].isalpha():
            connected_lines.append(i)
            current_group.append(i)
        # Check if current line is very short (likely part of a continued thought)
        elif len(words) <= 3 and i < len(lines) - 1:
            # Look ahead to see if next line could be a continuation
            if i+1 < len(lines):
                next_words = lines[i+1].lower().split()
                if next_words and next_words[0] in continuation_starters:
                    connected_lines.append(i)
                    current_group.append(i)
                else:
                    # This might end a group
                    if len(current_group) > 1:  # Only consider groups of 2+ lines
                        potential_groups.append(current_group.copy())
                    current_group = [i]
        else:
            # This likely starts a new thought
            if len(current_group) > 1:  # Only consider groups of 2+ lines
                potential_groups.append(current_group.copy())
            current_group = [i]
    
    # Add the last group if it has multiple lines
    if len(current_group) > 1:
        potential_groups.append(current_group)
    
    # Calculate metrics
    connected_groups = len(potential_groups)
    
    if connected_groups > 0:
        avg_lines_per_group = sum(len(group) for group in potential_groups) / connected_groups
        
        # Determine flow quality
        if connected_groups >= len(lines) / 3 and avg_lines_per_group >= 2.5:
            flow_quality = "Excellent - multiple connected thoughts across lines"
        elif connected_groups >= len(lines) / 4 and avg_lines_per_group >= 2:
            flow_quality = "Good - some connected thoughts across lines"
        elif connected_groups > 0:
            flow_quality = "Fair - limited connection between lines"
        else:
            flow_quality = "Poor - mostly independent lines"
    else:
        avg_lines_per_group = 0
        flow_quality = "Poor - no connected thoughts detected"
    
    return {
        "connected_groups": connected_groups,
        "avg_lines_per_group": avg_lines_per_group,
        "flow_quality": flow_quality
    }

def analyze_theme_emotion_expression(lyrics, theme, emotion):
    """Analyze how well the lyrics express the target theme and emotion"""
    
    # Normalize inputs
    lyrics_text = lyrics.lower()
    emotion = emotion.lower()
    theme = theme.lower()
    
    # Find closest emotion lexicon if exact match not found
    if emotion not in emotion_lexicons:
        closest_emotion = "reflective"  # Default
        for key in emotion_lexicons:
            if emotion in key or key in emotion:
                closest_emotion = key
                break
        emotion = closest_emotion
        
    # Find closest theme lexicon if exact match not found
    if theme not in theme_lexicons:
        closest_theme = "journey"  # Default
        for key in theme_lexicons:
            if theme in key or key in theme:
                closest_theme = key
                break
        theme = closest_theme
    
    # Count emotional and thematic words in lyrics
    emotion_matches = 0
    theme_matches = 0
    
    for word in emotion_lexicons[emotion]:
        if word in lyrics_text:
            emotion_matches += 1
            
    for word in theme_lexicons[theme]:
        if word in lyrics_text:
            theme_matches += 1
    
    # Calculate scores as percentages of available words
    emotion_score = min(100, (emotion_matches / len(emotion_lexicons[emotion])) * 100)
    theme_score = min(100, (theme_matches / len(theme_lexicons[theme])) * 100)
    
    # Qualitative assessment
    if emotion_score >= 30 and theme_score >= 30:
        expression_quality = "Strong"
    elif emotion_score >= 20 and theme_score >= 20:
        expression_quality = "Good"
    elif emotion_score >= 10 and theme_score >= 10:
        expression_quality = "Fair"
    else:
        expression_quality = "Weak"
    
    return {
        "emotion_score": emotion_score,
        "theme_score": theme_score,
        "expression_quality": expression_quality,
        "emotion_words_found": emotion_matches,
        "theme_words_found": theme_matches
    }

# Create Gradio interface
def create_interface():
    with gr.Blocks(title="Music Analysis & Lyrics Generator") as demo:
        gr.Markdown("# Music Analysis & Lyrics Generator")
        gr.Markdown("Upload a music file or record audio to analyze it and generate matching lyrics")
        
        with gr.Row():
            with gr.Column(scale=1):
                audio_input = gr.Audio(
                    label="Upload or Record Audio", 
                    type="filepath",
                    sources=["upload", "microphone"]
                )
                analyze_btn = gr.Button("Analyze and Generate Lyrics", variant="primary")
            
            with gr.Column(scale=2):
                with gr.Tab("Analysis"):
                    analysis_output = gr.Textbox(label="Music Analysis Results", lines=10)
                    
                    with gr.Row():
                        tempo_output = gr.Number(label="Tempo (BPM)")
                        time_sig_output = gr.Textbox(label="Time Signature")
                        emotion_output = gr.Textbox(label="Primary Emotion")
                        theme_output = gr.Textbox(label="Primary Theme")
                        genre_output = gr.Textbox(label="Primary Genre")
                
                with gr.Tab("Generated Lyrics"):
                    lyrics_output = gr.Textbox(label="Generated Lyrics", lines=20)
                
                with gr.Tab("Beat Matching"):
                    beat_match_output = gr.Markdown(label="Beat & Syllable Matching Analysis")
        
        # Set up event handlers
        analyze_btn.click(
            fn=process_audio,
            inputs=[audio_input],
            outputs=[analysis_output, lyrics_output, tempo_output, time_sig_output, 
                    emotion_output, theme_output, genre_output, beat_match_output]
        )
        
        # Format supported genres for display
        supported_genres_md = "\n".join([f"- {genre.capitalize()}" for genre in beat_analyzer.supported_genres])
        
        gr.Markdown(f"""
        ## How it works
        1. Upload or record a music file
        2. The system analyzes tempo, beats, time signature and other musical features
        3. It detects emotion, theme, and music genre
        4. Using beat patterns and syllable stress analysis, it generates perfectly aligned lyrics
        5. Each line of the lyrics is matched to the beat pattern of the corresponding musical phrase
        
        ## Supported Genres
        **Note:** Lyrics generation is currently only supported for the following genres:
        {supported_genres_md}
        
        These genres have consistent syllable-to-beat patterns that work well with our algorithm.
        For other genres, only music analysis will be provided.
        """)
    
    return demo

# Launch the app
demo = create_interface()

if __name__ == "__main__":
    demo.launch()
else:
    # For Hugging Face Spaces
    app = demo