import gradio as gr
import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from TTS.api import TTS
import numpy as np
import librosa
import soundfile as sf
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D
import io
import base64
import os

# Initialize text generation model (GPT-2)
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

# Initialize TTS model
tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC")

# Function to generate text using GPT-2
def generate_text(prompt, max_length=100):
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    output = model.generate(input_ids, max_length=max_length, num_return_sequences=1, no_repeat_ngram_size=2)
    return tokenizer.decode(output[0], skip_special_tokens=True)

# Function to generate speech using TTS
def generate_speech(text):
    output_path = "generated_speech.wav"
    tts.tts_to_file(text=text, file_path=output_path)
    return output_path

# Function to create a more advanced character animation
def create_character_animation(text):
    def update_point(num, data, line):
        line.set_data(data[:2, :num])
        line.set_3d_properties(data[2, :num])
        return line

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # Generate a more complex 3D path for the character
    t = np.linspace(0, 20, 100)
    x = np.sin(t)
    y = np.cos(t)
    z = t/10

    data = np.array([x, y, z])
    line, = ax.plot(data[0, 0:1], data[1, 0:1], data[2, 0:1])

    # Setting the axes properties
    ax.set_xlim3d([-1.0, 1.0])
    ax.set_xlabel('X')
    ax.set_ylim3d([-1.0, 1.0])
    ax.set_ylabel('Y')
    ax.set_zlim3d([0.0, 2.0])
    ax.set_zlabel('Z')

    # Adding text
    ax.text2D(0.05, 0.95, text[:50], transform=ax.transAxes)

    ani = animation.FuncAnimation(fig, update_point, frames=100, fargs=(data, line), interval=100, blit=False)
    
    # Save animation as gif
    ani.save('character_animation.gif', writer='pillow')
    
    return 'character_animation.gif'

# Main function to generate comedy animation
def generate_comedy_animation(prompt):
    script = generate_text(f"Write a short comedy script about {prompt}: ")
    animation_path = create_character_animation(script)
    speech_path = generate_speech(script)
    return script, animation_path, speech_path

# Main function to generate kids music animation
def generate_kids_music_animation(theme):
    lyrics = generate_text(f"Write a short children's song about {theme}: ")
    animation_path = create_character_animation(lyrics)
    speech_path = generate_speech(lyrics)
    return lyrics, animation_path, speech_path

# Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("## Advanced Character Animation Generator with Open Source Models")
    
    with gr.Tab("Comedy Animation"):
        comedy_prompt = gr.Textbox(label="Enter comedy prompt")
        comedy_generate_btn = gr.Button("Generate Comedy Animation")
        comedy_script = gr.Textbox(label="Generated Comedy Script")
        comedy_animation = gr.Image(label="Comedy Animation")
        comedy_audio = gr.Audio(label="Comedy Speech")

        comedy_generate_btn.click(
            generate_comedy_animation,
            inputs=comedy_prompt,
            outputs=[comedy_script, comedy_animation, comedy_audio]
        )
    
    with gr.Tab("Kids Music Animation"):
        kids_theme = gr.Textbox(label="Enter kids music theme")
        kids_generate_btn = gr.Button("Generate Kids Music Animation")
        kids_lyrics = gr.Textbox(label="Generated Lyrics")
        kids_animation = gr.Image(label="Kids Music Animation")
        kids_audio = gr.Audio(label="Kids Speech")

        kids_generate_btn.click(
            generate_kids_music_animation,
            inputs=kids_theme,
            outputs=[kids_lyrics, kids_animation, kids_audio]
        )

app.launch()