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 - ALWAYS RANDOM"""
    
    # Get base working pattern - ALWAYS RANDOM
    base_pattern = get_random_working_pattern(genre, complexity)
    
    # Add random variations to make it even more unique
    base_pattern = create_variations(base_pattern, genre)
    
    # 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 random timestamp to ensure uniqueness
    import time
    timestamp = int(time.time() * 1000) % 10000  # Last 4 digits of timestamp
    
    # Add comment based on prompt
    comment_line = f"// {prompt[:50]}{'...' if len(prompt) > 50 else ''} #{timestamp}"
    
    return f"{comment_line}\n{base_pattern}"

def create_visual_code(style="reactive", genre="techno"):
    """Create working visual code with genre-appropriate randomization"""
    
    # Genre-specific visual styles
    genre_visuals = {
        "techno": [
            """osc(12, 0.1, 0.8)
.color(2, 0.5, 1.5)
.kaleid(4)
.rotate(0, 0.1)
.out()""",
            """shape(4, 0.5, 0.01)
.repeat(4, 3)
.color(1.8, 0.8, 2)
.kaleid(6)
.out()""",
            """osc(16, 0.05, 1)
.pixelate(24, 18)
.color(2, 1, 0.8)
.contrast(1.8)
.out()"""
        ],
        "house": [
            """osc(6, 0.1, 1.2)
.color(1.5, 1.8, 0.8)
.modulate(noise(2), 0.2)
.kaleid(8)
.out()""",
            """shape(6, 0.4, 0.02)
.repeat(3, 2)
.rotate(0, 0.05)
.color(1.2, 1.5, 1.8)
.out()""",
            """noise(2, 0.1)
.color(1.8, 1.2, 1.5)
.modulate(osc(3, 0.08), 0.3)
.kaleid(4)
.out()"""
        ],
        "ambient": [
            """noise(3, 0.05)
.color(0.8, 1.5, 1.8)
.modulate(osc(1, 0.1), 0.4)
.contrast(1.2)
.out()""",
            """osc(4, 0.02, 1.5)
.color(1.2, 1.8, 1.5)
.modulate(noise(1), 0.5)
.kaleid(3)
.out()""",
            """shape(8, 0.2, 0.005)
.repeat(2, 2)
.color(0.8, 1.8, 1.2)
.rotate(0, 0.02)
.out()"""
        ],
        "breakbeat": [
            """osc(20, 0.2, 0.5)
.pixelate(16, 12)
.color(2, 1.5, 0.5)
.contrast(2)
.out()""",
            """shape(3, 0.8, 0.05)
.repeat(6, 4)
.color(1.8, 2, 1)
.kaleid(5)
.out()""",
            """noise(4, 0.3)
.color(2, 1.8, 0.8)
.modulate(osc(8, 0.1), 0.6)
.pixelate(20, 15)
.out()"""
        ],
        "experimental": [
            """noise(5, 0.4)
.color(2, 0.8, 1.8)
.modulate(osc(15, 0.3), 0.8)
.pixelate(8, 6)
.contrast(2.5)
.out()""",
            """osc(25, 0.4, 0.2)
.color(1.8, 2, 0.5)
.kaleid(12)
.rotate(0, 0.2)
.out()""",
            """shape(7, 0.9, 0.1)
.repeat(8, 6)
.color(2, 1, 2)
.modulate(noise(6), 0.7)
.out()"""
        ],
        "hiphop": [
            """osc(8, 0.1, 1)
.color(1.8, 1.5, 0.8)
.pixelate(32, 24)
.contrast(1.6)
.out()""",
            """shape(4, 0.6, 0.02)
.repeat(4, 3)
.color(2, 1.8, 1)
.kaleid(4)
.out()""",
            """noise(3, 0.2)
.color(1.5, 1.8, 0.8)
.modulate(osc(4, 0.1), 0.3)
.pixelate(28, 20)
.out()"""
        ]
    }
    
    # Get random visual from genre-specific options
    visuals = genre_visuals.get(genre, genre_visuals["techno"])
    return random.choice(visuals)

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

{create_visual_code("random", genre)}

"""
    
    # 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, 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 random temperature and seed for variety
        import time
        random_seed = int(time.time() * 1000) % 1000  # Random seed based on timestamp
        
        system_prompt = f"""Generate working Strudel live coding pattern for: {prompt}
Genre: {genre}, Complexity: {complexity}, Variation: {random_seed}

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 unique {complexity} {genre} pattern:"""
        
        try:
            outputs = code_generator(
                system_prompt,
                max_length=len(system_prompt.split()) + 80,
                temperature=random.uniform(0.7, 1.2),  # Random temperature for variety
                do_sample=True,
                top_p=random.uniform(0.8, 0.95),  # Random top_p for variety
                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} #{random_seed}\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 - ALWAYS RANDOM
        strudel_code = generate_working_strudel_code(prompt, genre, complexity)
    
    # Create complete working template with random visuals based on genre
    complete_code = create_complete_strudel_code(strudel_code, include_visuals, genre)
    
    return complete_code

# Create Gradio interface
with gr.Blocks(title="Working Strudel Generator", theme=gr.themes.Soft()) as app:
    gr.Markdown("""
    # 🎵 CODEL: AI Strudel Code Generator
    
    Generate high-quality Strudel live coding patterns! 
    
    ✅ **Premium curated patterns** - professional quality sounds  
    🎛️ **Real Strudel syntax** - copy & paste ready  
    🎨 **Random genre-specific visuals** - different every time!  
    🎵 **6 genres** with authentic patterns  
 
    
    **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
                )
            
            generate_btn = gr.Button("Generate 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 = [
            ["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,
            use_ai
        ],
        outputs=output_code
    )

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