import gradio as gr
from transformers import pipeline
from gtts import gTTS
from moviepy.editor import ImageSequenceClip, CompositeVideoClip, ImageClip, AudioFileClip, VideoFileClip, TextClip, concatenate_videoclips
from PIL import Image, ImageDraw, ImageFont
import numpy as np
from scipy.io.wavfile import write

# Load and Initialize Models
script_generator = pipeline("text-generation", model="gpt2", truncation=True, max_length=100)

def generate_comedy_script(prompt):
    script = script_generator(prompt)[0]['generated_text']
    return script

def text_to_speech(script):
    tts = gTTS(text=script, lang='en')
    audio_file = 'output.mp3'
    tts.save(audio_file)
    return audio_file

def create_images_from_script(script):
    lines = script.split('. ')
    image_paths = []
    for i, line in enumerate(lines):
        img = Image.new('RGB', (800, 400), color=(0, 0, 0))
        d = ImageDraw.Draw(img)
        fnt = ImageFont.load_default()
        d.text((10, 180), line, font=fnt, fill=(255, 255, 255))
        image_path = f'/tmp/image_{i}.png'
        img.save(image_path)
        image_paths.append(image_path)
    return image_paths

def generate_animation_with_images(script):
    image_paths = create_images_from_script(script)
    clips = []
    for i, img_path in enumerate(image_paths):
        image_clip = ImageClip(img_path).set_duration(3).set_position(('center', 'center'))
        clips.append(image_clip)
    final_video = concatenate_videoclips(clips, method="compose")
    final_video.write_videofile("/tmp/final_video.mp4", fps=24)
    return "/tmp/final_video.mp4"

def combine_audio_video(video_path, audio_path):
    video = VideoFileClip(video_path)
    audio = AudioFileClip(audio_path)
    final_video = video.set_audio(audio)
    final_video.write_videofile("/tmp/final_comedy_video.mp4", fps=24)
    return "/tmp/final_comedy_video.mp4"

def generate_comedy_and_animation(prompt):
    script = generate_comedy_script(prompt)
    audio_file = text_to_speech(script)
    video_file = generate_animation_with_images(script)
    final_video = combine_audio_video(video_file, audio_file)
    return script, audio_file, final_video

def generate_sine_wave(frequency, duration, sample_rate=44100, amplitude=0.5):
    t = np.linspace(0, duration, int(sample_rate * duration), False)
    wave = amplitude * np.sin(2 * np.pi * frequency * t)
    return wave

def generate_kids_music(theme, output_music_file="kids_music.wav"):
    sample_rate = 44100
    duration = 2
    c4_wave = generate_sine_wave(261.63, duration)
    d4_wave = generate_sine_wave(293.66, duration)
    e4_wave = generate_sine_wave(329.63, duration)
    wave = np.concatenate([c4_wave, d4_wave, e4_wave])
    audio_wave = np.int16(wave * 32767)
    write(output_music_file, sample_rate, audio_wave)
    return output_music_file

def generate_kids_animation_with_music(theme, output_video_file="kids_animation.mp4"):
    music_file = generate_kids_music(theme)
    clips = []
    for i in range(5):