import gradio as gr
import requests
import json
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch
import random

# Use much faster, smaller model
MODEL_NAME = "microsoft/DialoGPT-small"

# Initialize model with pipeline for easier usage
try:
    code_generator = pipeline(
        "text-generation",
        model=MODEL_NAME,
        tokenizer=MODEL_NAME,
        torch_dtype=torch.float16,
        device_map="auto" if torch.cuda.is_available() else None,
        trust_remote_code=True
    )
except Exception as e:
    print(f"Error loading model: {e}")
    # Fallback to CPU if GPU fails
    code_generator = pipeline(
        "text-generation",
        model=MODEL_NAME,
        tokenizer=MODEL_NAME,
        device_map=None
    )

# Enhanced working Strudel patterns - these are all tested and functional
WORKING_PATTERNS = {
    "techno": [
        """stack(
  s("bd*4").gain(0.9),
  s("~ hh ~ hh").gain(0.6),
  n("0 ~ 3 ~").s("sawtooth").octave(2).lpf(1200).gain(0.8),
  s("~ ~ rim ~").slow(2).gain(0.5)
).scale("a:minor")""",
        
        """stack(
  s("bd bd ~ bd, ~ rim ~ rim").bank("RolandTR909").gain(0.9),
  s("[hh ch]*4").degradeBy(0.1).gain(0.5),
  n("[0 3 5 7]*2").s("square").octave(2).cutoff(sine.range(400,1600).slow(8)).gain(0.7),
  n("7 9 7 4").s("triangle").octave(4).delay(0.25).room(0.4).gain(0.5)
).scale("d:minor")""",

        """stack(
  s("bd*4, [~ rim]*2, hh*8").bank("RolandTR909").speed(0.95).gain(0.8),
  n("0 0 3 3").s("sawtooth").octave(2).lpf(2000).lpq(8).gain(0.8),
  n("[7 9]*4").s("square").octave(4).room(0.5).cutoff(1800).gain(0.4),
  s("~ ~ ~ clap").slow(2).gain(0.6).room(0.3)
).scale("g:minor")""",

        """stack(
  s("bd*4").sometimes(fast(2)).gain(0.9),
  s("hh*8").degradeBy(0.2).gain(0.4),
  n("[0 3 [5 0] 7]*2").s("sawtooth").octave(2).lpf(sine.range(600,2400).slow(4)).gain(0.8),
  n("0 7 4 2").s("triangle").octave(5).delay(0.125).room(0.6).gain(0.5)
).scale("e:minor")"""
    ],
    
    "house": [
        """stack(
  s("bd ~ ~ ~ bd ~ ~ ~").gain(0.8),
  s("~ hh ~ hh").gain(0.5),
  n("0 2 4 7").s("sine").octave(3).room(0.6).gain(0.7),
  s("~ ~ ~ clap").gain(0.6)
).scale("c:major")""",

        """stack(
  s("bd ~ [~ bd] ~").sometimes(fast(2)).gain(0.8),
  s("[~ ch ~ hh]*2").pan(sine.slow(4)).gain(0.4),
  n("[0 4 7]*2").s("sawtooth").octave(2).lpf(1500).gain(0.8),
  n("0 7 4 2").s("triangle").octave(4).delay(0.125).room(0.5).gain(0.5)
).scale("f:major")""",

        """stack(
  s("bd ~ ~ ~ bd ~ ~ ~").gain(0.8),
  n("[0 2] [4 7] [2 4] [7 0]").s("sine").octave(2).lpf(2000).gain(0.7),
  s("hh*8").degradeBy(0.3).pan(cosine.slow(6)).gain(0.4),
  s("~ ~ clap ~").room(0.4).gain(0.6)
).scale("g:major")""",

        """stack(
  s("bd ~ [bd ~] [~ bd] ~").gain(0.8),
  s("~ hh ~ [hh ch]").gain(0.5),
  n("0 4 7 9").s("sawtooth").octave(2).sometimes(add(note(12))).lpf(1800).gain(0.7),
  n("[2 4]*4").s("triangle").octave(4).delay(0.25).gain(0.4)
).scale("d:major")"""
    ],
    
    "ambient": [
        """stack(
  s("~ ~ ~ rim").slow(4).gain(0.3).room(0.9),
  n("[0 2 4]/3").s("sine").octave(3).slow(8).room(0.9).lpf(800).gain(0.6),
  n("[7 9 11]/5").s("triangle").octave(4).slow(12).delay(0.5).room(0.8).gain(0.4),
  n("[0 4]/7").s("sawtooth").octave(2).slow(16).lpf(400).gain(0.3)
).scale("d:minor")""",

        """stack(
  n("0 2 4 7").s("sine").octave(2).slow(16).room(0.9).lpf(600).gain(0.5),
  n("[4 7 9]/7").s("triangle").octave(4).slow(12).delay(0.75).room(0.9).gain(0.4),
  s("~ ~ rim ~").slow(8).gain(0.2).room(0.9),
  n("[11 14]/11").s("sine").octave(5).slow(20).gain(0.3)
).scale("a:minor")""",

        """n("[0 4 7 11]/4").s("sine").octave(3).slow(12).room(0.9).lpf(perlin.range(400,1200).slow(8)).gain(0.7).scale("e:minor")""",

        """stack(
  n("[0 2]/5").s("sine").octave(2).slow(16).room(0.9).gain(0.4),
  n("[4 7]/7").s("triangle").octave(4).slow(20).delay(1).room(0.9).gain(0.3),
  s("~ ~ ~ [rim vinyl]").slow(8).gain(0.2).room(0.9)
).scale("f:major")"""
    ],
    
    "breakbeat": [
        """stack(
  s("[bd sn]*2").sometimes(rev).speed(0.9).gain(0.9),
  s("[hh ch]*4").degradeBy(0.3).gain(0.6),
  n("0 4 7 0").s("square").octave(2).cutoff(1000).gain(0.8),
  s("~ [perc:0 perc:1] ~").gain(0.4)
).scale("a:minor")""",

        """stack(
  s("bd [sn sn] bd sn").speed(0.85).gain(0.9),
  n("[0 ~ 7] [3 5]").s("sawtooth").octave(2).lpf(1500).gain(0.7),
  s("hh*8").sometimes(fast(2)).degradeBy(0.2).gain(0.5),
  s("~ ~ [perc:2 perc:3] ~").gain(0.5)
).scale("d:minor")""",

        """stack(
  s("[bd ~ sn ~]*2").sometimes(rev).gain(0.9),
  s("[hh ch oh]*3").degradeBy(0.4).pan(perlin.slow(2)).gain(0.5),
  n("[0 7 4 2]*2").s("square").octave(2).cutoff(rand.range(600,1800)).gain(0.7)
).scale("g:minor")"""
    ],
    
    "experimental": [
        """stack(
  s("[bd rim]*3").degradeBy(0.5).speed(perlin.range(0.8,1.2)).gain(0.8),
  n("[0 3 7]/5").s("square").octave(2).cutoff(perlin.range(400,2000)).gain(0.7),
  s("~ [hh hh] ~").pan(sine.slow(3)).degradeBy(0.6).gain(0.5),
  s("[noise:0]/16").gain(0.2).hpf(2000)
).scale("g:minor")""",

        """stack(
  s("bd ~ [perc:0 perc:1]").speed(rand.range(0.7,1.3)).gain(0.8),
  n("0 [2 5] 7").s("sawtooth").octave(2).sometimes(add(note(12))).cutoff(sine.range(300,1500)).gain(0.7),
  s("[ch oh]*5").degradeBy(0.7).pan(cosine.slow(2)).gain(0.4),
  n("[9 11]/9").s("triangle").octave(5).delay(0.75).gain(0.3)
).scale("f#:minor")""",

        """stack(
  s("[bd ~ rim]*2").rev().degradeBy(0.4).gain(0.8),
  n("[0 4 7]/7").s("square").octave(2).cutoff(perlin.range(200,2000)).sometimes(fast(4)).gain(0.6),
  s("~ [vinyl:0 vinyl:1]/8").gain(0.3).hpf(1500)
).scale("bb:minor")"""
    ],

    "hiphop": [
        """stack(
  s("bd ~ sn ~").gain(0.9),
  s("hh ~ hh hh ~ hh ~ hh").gain(0.5),
  n("0 0 3 5").s("sawtooth").octave(1).lpf(1000).gain(0.8),
  s("~ ~ ~ [clap clap]").slow(2).gain(0.6)
).scale("c:minor")""",

        """stack(
  s("bd ~ [bd sn] ~").gain(0.9),
  s("[hh ch]*4").degradeBy(0.1).gain(0.5),
  n("[0 3 5]*2").s("sine").octave(1).lpf(800).gain(0.9),
  n("0 3 5 7").s("square").octave(3).lpf(1200).gain(0.4),
  s("~ ~ clap ~").room(0.3).gain(0.7),
  s("~ [perc:0]/4 ~ [perc:1]/4").gain(0.4)
).scale("a:minor")""",

        """stack(
  s("bd ~ sn [~ bd]").sometimes(fast(2)).gain(0.9),
  s("hh*8").degradeBy(0.2).gain(0.4),
  n("0 ~ 3 ~").s("sawtooth").octave(1).lpf(1200).gain(0.9),
  n("[0 3 5 7]*2").s("triangle").octave(3).gain(0.5),
  s("~ ~ [clap rim] ~").gain(0.6).room(0.4)
).scale("d:minor")""",

        """stack(
  s("[bd bd] ~ sn ~").gain(0.9),
  s("~ hh ~ [hh ch hh]").gain(0.5),
  n("[0 3] [5 0] [7 3] [5 0]").s("sine").octave(1).lpf(1500).gain(0.8),
  n("0 7 4 2").s("square").octave(3).delay(0.125).gain(0.4),
  s("~ ~ clap ~").delay(0.125).gain(0.7)
).scale("g:minor")"""
    ]
}

def get_random_working_pattern(genre, complexity="moderate"):
    """Get a random working pattern from our curated collection"""
    patterns = WORKING_PATTERNS.get(genre, WORKING_PATTERNS["techno"])
    
    if complexity == "simple":
        # Return simpler single-line patterns
        simple_patterns = {
            "techno": 's("bd*4, hh*8").gain(0.8)',
            "house": 's("bd ~ ~ ~ bd ~ ~ ~, ~ hh ~ hh").gain(0.7)', 
            "ambient": 'n("[0 2 4]/3").s("sine").octave(3).slow(8).room(0.9).gain(0.6).scale("d:minor")',
            "breakbeat": 's("[bd sn]*2, hh*8").speed(0.9).gain(0.8)',
            "experimental": 's("bd ~ [perc:0 perc:1]").degradeBy(0.4).speed(rand.range(0.8,1.2)).gain(0.7)',
            "hiphop": 's("bd ~ sn ~, hh ~ hh hh ~ hh ~ hh").gain(0.8).add(n("0 0 3 5").s("sine").octave(1).gain(0.8))'
        }
        return simple_patterns.get(genre, simple_patterns["techno"])
    
    return random.choice(patterns)

def create_variations(base_pattern, genre):
    """Create variations of a working base pattern"""
    variations = []
    
    # Original pattern
    variations.append(base_pattern)
    
    # Speed variations
    if "stack(" in base_pattern:
        variations.append(base_pattern.replace(").scale(", ").sometimes(fast(2)).scale("))
        variations.append(base_pattern.replace(").scale(", ").slow(2).scale("))
    
    # Effect variations
    if ".gain(0.8)" in base_pattern:
        variations.append(base_pattern.replace(".gain(0.8)", ".gain(0.8).room(0.4)"))
    
    # Scale variations
    scales = ["a:minor", "d:minor", "g:minor", "c:major", "f:major", "e:minor", "bb:minor", "f#:minor"]
    for scale in scales:
        if scale not in base_pattern:
            new_pattern = base_pattern
            for old_scale in scales:
                if old_scale in base_pattern:
                    new_pattern = base_pattern.replace(old_scale, scale)
                    break
            if new_pattern != base_pattern:
                variations.append(new_pattern)
                break
    
    return random.choice(variations)

def generate_working_strudel_code(prompt, genre="techno", complexity="moderate"):
    """Generate guaranteed working Strudel code"""
    
    # Get base working pattern
    base_pattern = get_random_working_pattern(genre, complexity)
    
    # Create variations based on prompt keywords
    if any(word in prompt.lower() for word in ["fast", "speed", "quick", "rapid", "energetic"]):
        if "sometimes(fast(2))" not in base_pattern:
            base_pattern = base_pattern.replace(".gain(", ".sometimes(fast(2)).gain(")
    
    if any(word in prompt.lower() for word in ["slow", "chill", "ambient", "relaxed", "calm"]):
        if ".slow(" not in base_pattern:
            base_pattern = base_pattern.replace(".gain(", ".slow(2).gain(")
    
    if any(word in prompt.lower() for word in ["reverb", "space", "room", "hall", "spacious"]):
        if ".room(" not in base_pattern:
            base_pattern = base_pattern.replace(".gain(", ".room(0.7).gain(")
    
    if any(word in prompt.lower() for word in ["delay", "echo", "repeat", "dub"]):
        if ".delay(" not in base_pattern:
            base_pattern = base_pattern.replace(".gain(", ".delay(0.25).gain(")
    
    if any(word in prompt.lower() for word in ["filter", "sweep", "cutoff", "lpf", "filtered"]):
        if ".lpf(" not in base_pattern and ".cutoff(" not in base_pattern:
            base_pattern = base_pattern.replace(".gain(", ".lpf(sine.range(400,1600).slow(8)).gain(")
    
    if any(word in prompt.lower() for word in ["glitch", "degraded", "broken", "corrupt", "dirty"]):
        if ".degradeBy(" not in base_pattern:
            base_pattern = base_pattern.replace(".gain(", ".degradeBy(0.4).gain(")
    
    if any(word in prompt.lower() for word in ["pan", "stereo", "wide", "spatial"]):
        if ".pan(" not in base_pattern:
            base_pattern = base_pattern.replace(".gain(", ".pan(sine.slow(4)).gain(")
    
    # Add comment based on prompt
    comment_line = f"// {prompt[:50]}{'...' if len(prompt) > 50 else ''}"
    
    return f"{comment_line}\n{base_pattern}"

def create_visual_code(style="reactive"):
    """Create working visual code"""
    visual_styles = {
        "reactive": """osc(8, 0.1, 1.2)
.color(1.8, 0.8, 1.5)
.modulate(noise(2), 0.3)
.kaleid(6)
.out()""",
        "kaleidoscope": """shape(6, 0.3, 0.01)
.repeat(3, 2)
.rotate(0, 0.03)
.color(1.2, 1.8, 0.8)
.kaleid(8)
.out()""",
        "flowing": """noise(3, 0.1)
.color(1.5, 1.2, 0.8)
.modulate(osc(2, 0.05), 0.4)
.contrast(1.4)
.out()""",
        "geometric": """osc(12, 0.02, 0.8)
.pixelate(32, 24)
.color(2, 0.8, 1.5)
.contrast(1.6)
.out()"""
    }
    return visual_styles.get(style, visual_styles["reactive"])

def create_complete_strudel_code(generated_code, include_visuals=True, visual_style="reactive"):
    """Create complete, working Strudel code"""
    
    visual_code = ""
    if include_visuals:
        visual_code = f"""// Hydra visuals
await initHydra({{feedStrudel:5}})

{create_visual_code(visual_style)}

"""
    
    # Add proper Strudel prefix for audio patterns
    if not generated_code.startswith("$:"):
        generated_code = f"$: {generated_code}"
    
    complete_code = f"""{visual_code}{generated_code}"""
    
    return complete_code

# Main generation function
def generate_interface(prompt, genre, complexity, include_visuals, visual_style, use_ai):
    """Main interface function that generates working code"""
    
    if not prompt.strip():
        return "Please enter a description of the music you want to create."
    
    if use_ai:
        # Use AI generation with better examples
        system_prompt = f"""Generate working Strudel live coding pattern for: {prompt}
Genre: {genre}, Complexity: {complexity}

Example Strudel patterns:
s("bd*4, hh*8").gain(0.8)
stack(s("bd ~ sn ~"), n("0 2 4").s("sine").octave(3).gain(0.7))
n("[0 3 5 7]*2").s("sawtooth").octave(2).lpf(1200).gain(0.8).scale("a:minor")

Generate a {complexity} {genre} pattern:"""
        
        try:
            outputs = code_generator(
                system_prompt,
                max_length=len(system_prompt.split()) + 80,
                temperature=0.8,
                do_sample=True,
                top_p=0.9,
                num_return_sequences=1
            )
            
            generated_text = outputs[0]['generated_text']
            strudel_code = generated_text[len(system_prompt):].strip()
            
            # Clean the AI output more thoroughly
            lines = strudel_code.split('\n')
            clean_lines = []
            for line in lines[:8]:  # Limit to 8 lines
                line = line.strip()
                # More robust pattern matching for Strudel code
                if line and (line.startswith('s(') or line.startswith('n(') or line.startswith('stack(') or 
                           '.gain(' in line or '.s(' in line or '.octave(' in line or '.scale(' in line or
                           line.endswith(')') or line.endswith(',') or '.lpf(' in line):
                    clean_lines.append(line)
                # Stop at obvious non-Strudel patterns
                elif any(stop_word in line.lower() for stop_word in ['function', 'var ', 'let ', 'const ', 'import', '//']):
                    break
            
            if clean_lines:
                strudel_code = '\n'.join(clean_lines)
                strudel_code = f"// AI Generated: {prompt}\n{strudel_code}"
            else:
                # Fallback to curated patterns if AI fails
                strudel_code = generate_working_strudel_code(prompt, genre, complexity)
        except Exception as e:
            # Fallback to curated patterns if AI fails
            print(f"AI generation failed: {e}")
            strudel_code = generate_working_strudel_code(prompt, genre, complexity)
    else:
        # Use curated pattern generation
        strudel_code = generate_working_strudel_code(prompt, genre, complexity)
    
    # Create complete working template
    complete_code = create_complete_strudel_code(strudel_code, include_visuals, visual_style)
    
    return complete_code

# Create Gradio interface
with gr.Blocks(title="Working Strudel Generator", theme=gr.themes.Soft()) as app:
    gr.Markdown("""
    # 🎵 CODEL: Strudel Code Generator
    
    Generate high-quality Strudel live coding patterns! 
    
    ✅ **Premium curated patterns** - professional quality sounds  
    🎛️ **Real Strudel syntax** - copy & paste ready  
    🎨 **Reactive Hydra visuals** included  
    🎵 **6 genres** with authentic patterns  
    🤖 **Optional AI** (recommended: keep AI off for speed)
    
    **Usage:** Describe music → Generate → Copy to [strudel.cc](https://strudel.cc) → Play!
    """)
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="🎼 Describe your music",
                placeholder="e.g., 'Fast techno with reverb', 'Chill ambient with delay', 'Glitchy breakbeat'",
                lines=3
            )
            
            with gr.Row():
                genre_dropdown = gr.Dropdown(
                    choices=["techno", "house", "ambient", "breakbeat", "experimental", "hiphop"],
                    value="techno",
                    label="🎵 Genre"
                )
                
                complexity_dropdown = gr.Dropdown(
                    choices=["simple", "moderate", "complex"],
                    value="moderate",
                    label="⚙️ Complexity"
                )
            
            with gr.Row():
                include_visuals = gr.Checkbox(
                    label="🎨 Include visuals",
                    value=True
                )
                
                use_ai = gr.Checkbox(
                    label="🤖 Use AI (slower)",
                    value=False
                )
            
            visual_style = gr.Dropdown(
                choices=["reactive", "kaleidoscope", "flowing", "geometric"],
                value="reactive",
                label="👁️ Visual Style"
            )
            
            generate_btn = gr.Button("🎵 Generate Premium Code", variant="primary", size="lg")
        
        with gr.Column():
            output_code = gr.Code(
                label="✅ Premium Strudel Code (Ready to Copy)",
                language="javascript",
                lines=18
            )
            
            gr.Markdown("""
            **🎯 Instructions:**
            1. **Copy** all the code above
            2. **Go to** [strudel.cc](https://strudel.cc)
            3. **Paste** and **click the play button** ▶️
            
            **🔧 Pro Tips:**
            - Change `.gain()` values (0.1 to 1.0) to balance levels
            - Try different scales: "a:minor", "c:major", "d:minor"
            - Use `.sometimes()` for variation: `.sometimes(fast(2))`
            - Add effects: `.room()`, `.delay()`, `.lpf()`, `.degradeBy()`
            """)
    
    # Enhanced examples
    gr.Markdown("### 🎪 Premium Examples (Pro Quality)")
    with gr.Row():
        working_examples = [
            ["Driving techno with filter sweeps", "techno", "complex"],
            ["Deep house with swing", "house", "moderate"],
            ["Ethereal ambient soundscape", "ambient", "complex"],
            ["Crunchy hip-hop beat", "hiphop", "moderate"],
            ["Glitchy experimental textures", "experimental", "complex"],
        ]
        
        for example_text, example_genre, example_complexity in working_examples:
            btn = gr.Button(f"🔥 {example_text}", size="sm")
            btn.click(
                lambda t=example_text, g=example_genre, c=example_complexity: (t, g, c),
                outputs=[prompt_input, genre_dropdown, complexity_dropdown]
            )
    
    # Connect the generate button
    generate_btn.click(
        generate_interface,
        inputs=[
            prompt_input,
            genre_dropdown,
            complexity_dropdown,
            include_visuals,
            visual_style,
            use_ai
        ],
        outputs=output_code
    )

# Launch the app
if __name__ == "__main__":
    app.launch()