import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Model configuration
MODEL_NAME = "Text2MotionPrompter/Text2MotionPrompter"

# Pre-load tokenizer at startup
print("📥 Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
print("✅ Tokenizer loaded successfully!")

# Pre-load model to CPU at startup
print(f"📥 Loading model to CPU: {MODEL_NAME}...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype="auto",
    device_map="cpu"
)
print("✅ Model loaded to CPU successfully!")

# Prompt template
TEMPLATE = """
# Role
You are an expert in 3D motion analysis, animation timing, and choreography. Your task is to analyze textual action descriptions to estimate execution time and standardize the language for motion generation systems.

# Task
Analyze the user-provided [Input Action] and generate a structured JSON response containing a duration estimate and a refined caption.

# Instructions

### 1. Duration Estimation (frame_count)
- Analyze the complexity, speed, and physical constraints of the described action.
- Estimate the time required to perform the action in a **smooth, natural, and realistic manner**.
- Calculate the total duration in frames based on a **30 fps** (frames per second) standard.
- Output strictly as an Integer.

### 2. Caption Refinement (short_caption)
- Generate a refined, grammatically correct version of the input description in **English**.
- **Strict Constraints**:
    - You must **PRESERVE** the original sequence of events (chronological order).
    - You must **RETAIN** all original spatial modifiers (e.g., "left," "upward," "quickly").
    - **DO NOT** add new sub-actions or hallucinate details not present in the input.
    - **DO NOT** delete any specific movements.
- The goal is to improve clarity and flow while maintaining 100% semantic fidelity to the original request.

### 3. Output Format
- Return **ONLY** a raw JSON object.
- Do not use Markdown formatting (i.e., do not use ```json ... ```).
- Ensure the JSON is valid and parsable.

# JSON Structure
{{
    "duration": <Integer, frames at 30fps>,
    "short_caption": "<String, the refined English description>"
}}

# Input
{}
"""

@spaces.GPU(duration=120)
def generate_motion_prompt(action_input: str, max_new_tokens: int = 512) -> str:
    """
    Generate motion prompt from action description.
    
    Args:
        action_input: The action description to analyze
        max_new_tokens: Maximum number of tokens to generate
        
    Returns:
        Generated JSON response with duration and refined caption
    """
    if not action_input.strip():
        return "Please enter an action description."
    
    # Move model to GPU
    model.to("cuda")
    
    # Prepare the prompt
    messages = [
        {"role": "user", "content": TEMPLATE.format(action_input)}
    ]
    
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    
    model_inputs = tokenizer([text], return_tensors="pt").to("cuda")
    
    # Generate response
    with torch.no_grad():
        generated_ids = model.generate(
            **model_inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
        )
    
    # Decode only the new tokens
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
    content = tokenizer.decode(output_ids, skip_special_tokens=True)
    
    return content


# Custom CSS for a distinctive look
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;600&family=Space+Grotesk:wght@400;500;600&display=swap');

.gradio-container {
    font-family: 'Space Grotesk', sans-serif !important;
    background: linear-gradient(135deg, #0f0f1a 0%, #1a1a2e 50%, #16213e 100%) !important;
    min-height: 100vh;
}

.main-title {
    background: linear-gradient(90deg, #00d4ff, #7c3aed, #f472b6);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    font-size: 2.5rem !important;
    font-weight: 600 !important;
    text-align: center;
    margin-bottom: 0.5rem !important;
    font-family: 'Space Grotesk', sans-serif !important;
}

.subtitle {
    color: #94a3b8 !important;
    text-align: center;
    font-size: 1.1rem !important;
    margin-bottom: 2rem !important;
}

.input-container, .output-container {
    background: rgba(30, 41, 59, 0.7) !important;
    border: 1px solid rgba(100, 116, 139, 0.3) !important;
    border-radius: 16px !important;
    backdrop-filter: blur(10px);
}

textarea, input[type="text"] {
    background: rgba(15, 23, 42, 0.8) !important;
    border: 1px solid rgba(100, 116, 139, 0.4) !important;
    color: #e2e8f0 !important;
    font-family: 'JetBrains Mono', monospace !important;
    border-radius: 12px !important;
}

textarea:focus, input[type="text"]:focus {
    border-color: #7c3aed !important;
    box-shadow: 0 0 0 3px rgba(124, 58, 237, 0.2) !important;
}

.primary-btn {
    background: linear-gradient(135deg, #7c3aed 0%, #a855f7 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
    padding: 12px 32px !important;
    border-radius: 12px !important;
    font-size: 1rem !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(124, 58, 237, 0.4) !important;
}

.primary-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(124, 58, 237, 0.6) !important;
}

.example-btn {
    background: rgba(51, 65, 85, 0.6) !important;
    border: 1px solid rgba(100, 116, 139, 0.4) !important;
    color: #cbd5e1 !important;
    border-radius: 8px !important;
    transition: all 0.2s ease !important;
}

.example-btn:hover {
    background: rgba(71, 85, 105, 0.8) !important;
    border-color: #7c3aed !important;
}

label {
    color: #94a3b8 !important;
    font-weight: 500 !important;
}

.output-json {
    font-family: 'JetBrains Mono', monospace !important;
    background: rgba(15, 23, 42, 0.9) !important;
    color: #22d3ee !important;
    padding: 1.5rem !important;
    border-radius: 12px !important;
    border: 1px solid rgba(34, 211, 238, 0.2) !important;
}

.footer {
    text-align: center;
    color: #64748b;
    margin-top: 2rem;
    padding: 1rem;
    font-size: 0.9rem;
}

.slider-container input[type="range"] {
    accent-color: #7c3aed !important;
}
"""

# Build the Gradio interface
with gr.Blocks(css=custom_css, theme=gr.themes.Base()) as demo:
    gr.HTML("""
        <h1 class="main-title">🎬 Text2Motion Prompter</h1>
        <p class="subtitle">Transform action descriptions into structured motion prompts for 3D animation</p>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            action_input = gr.Textbox(
                label="🎯 Action Description",
                placeholder="Enter an action description (e.g., 'walking forward', '向前走', 'jumping up and down')",
                lines=4,
                elem_classes=["input-container"]
            )
            
            max_tokens = gr.Slider(
                label="🔧 Max Tokens",
                minimum=128,
                maximum=2048,
                value=512,
                step=64,
                elem_classes=["slider-container"]
            )
            
            submit_btn = gr.Button(
                "✨ Generate Motion Prompt",
                variant="primary",
                elem_classes=["primary-btn"]
            )
            
            gr.Markdown("### 📝 Examples")
            examples = gr.Examples(
                examples=[
                    ["走路"],
                    ["A person walks forward slowly, then turns left"],
                    ["跳跃并挥手"],
                    ["Running quickly and then stopping suddenly"],
                    ["坐下来休息一会儿，然后站起来"],
                    ["Dancing with arms raised above the head"],
                ],
                inputs=action_input,
                elem_id="examples"
            )
        
        with gr.Column(scale=1):
            output = gr.Textbox(
                label="📤 Generated Motion Prompt (JSON)",
                lines=12,
                elem_classes=["output-container", "output-json"]
            )
            
            gr.Markdown("""
            ### 📖 Output Format
            The model generates a JSON response containing:
            - **duration**: Estimated frames at 30fps
            - **short_caption**: Refined English description
            """)
    
    gr.HTML("""
        <div class="footer">
            <p>Powered by <strong>Text2MotionPrompter</strong> | Using Hugging Face ZeroGPU</p>
        </div>
    """)
    
    # Set up the action
    submit_btn.click(
        fn=generate_motion_prompt,
        inputs=[action_input, max_tokens],
        outputs=output
    )
    
    # Allow Enter key to submit
    action_input.submit(
        fn=generate_motion_prompt,
        inputs=[action_input, max_tokens],
        outputs=output
    )

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