import warnings
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr

# Ignore specific warnings
warnings.filterwarnings("ignore", category=FutureWarning, module="transformers.tokenization_utils_base")

# Load model and tokenizer
model_name = "EleutherAI/gpt-neo-1.3B"  # Use a lightweight model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Function to generate text
def generate_text(prompt, max_new_tokens=200):
    # Set pad_token_id to eos_token_id if not set
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token  # Set pad token to be the eos token
    
    # Tokenize the input prompt
    inputs = tokenizer(prompt, return_tensors="pt", padding=True)
    
    # Create an attention mask
    attention_mask = inputs.attention_mask

    # Generate the text using the model
    outputs = model.generate(
        inputs.input_ids,
        attention_mask=attention_mask,  # Pass the attention mask
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_k=50,  # Consider only the top 50 tokens for sampling
        top_p=0.95,
        temperature=0.8,  # Increase diversity slightly
        pad_token_id=tokenizer.pad_token_id  # Ensure correct behavior for padding
    )
    
    # Decode the output text
    text = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
    return text

# Function to generate situations and copy according to user selections
def generate_situations_and_copies(stage, product, situation):
    prompt = f"""
    Estás creando una campaña publicitaria para la etapa {stage} del funnel de marketing. El producto es: {product}.
    La situación cotidiana es: {situation}. Integra el concepto de marca "Tú eliges" en el copy.
    
    Genera 5 opciones de copy (titular y bajada) en español:
    1.
    2.
    3.
    4.
    5.
    """
    return generate_text(prompt)

# Dropdown options for Gradio
stages = ["Upper", "Middle", "Lower"]
products = {
    "Upper": ["Convierte tu teléfono en un POS", "Más de 700 plantillas web con Wix"],
    "Middle": ["Elige el modelo de dispositivo POS", "Entrega express en Lima"],
    "Lower": ["Servicios de cobro inmediato", "Cargas Ilimitadas de productos a Wix"]
}
situations = {
    "Upper": [
        "No elegiste dónde naciste, pero sí cómo gestionar tus ventas.",
        "No pudiste elegir el nombre que te pusieron, pero sí puedes elegir que tu teléfono sea tu POS."
    ],
    "Middle": [
        "No elegiste todas las decisiones difíciles, pero sí puedes elegir la herramienta perfecta para cada transacción.",
        "No puedes elegir todas las circunstancias, pero sí el modelo de POS adecuado para tu negocio."
    ],
    "Lower": [
        "No elegiste que el cliente pagara tarde, pero sí cómo cobrar rápido.",
        "No puedes elegir el ritmo de las ventas, pero sí cómo acelerar los cobros."
    ]
}

def update_dropdowns(stage):
    product_choices = products[stage]
    situation_choices = situations[stage]
    return gr.Dropdown.update(choices=product_choices), gr.Dropdown.update(choices=situation_choices)

# Custom CSS for design
custom_css = """
#app {
    background-color: white;
    color: #333;
}

.gradio-container {
    font-family: 'Helvetica Neue', sans-serif;
}

.gr-button {
    background-color: #ff4240;
    color: white;
    border: none;
    padding: 10px 20px;
    font-weight: bold;
    transition: background-color 0.3s ease;
}

.gr-button:hover {
    background-color: #e03a3a;
}

.gr-button-secondary {
    background-color: #3dd2ce;
    color: white;
    border: none;
    padding: 10px 20px;
    font-weight: bold;
    transition: background-color 0.3s ease;
}

.gr-button-secondary:hover {
    background-color: #37bebc;
}

.gr-textbox, .gr-dropdown {
    border: 2px solid #3dd2ce;
    border-radius: 5px;
    padding: 8px;
}

.gr-textbox:focus, .gr-dropdown:focus {
    border-color: #ff4240;
    box-shadow: 0 0 5px rgba(255, 66, 64, 0.5);
}
"""

# Gradio Interface with CSS styling
with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("# Herramienta de Generación de Copies Publicitarios")

    with gr.Row():
        stage_input = gr.Dropdown(choices=stages, label="Etapa del Funnel", value="Upper")
        product_input = gr.Dropdown(choices=products["Upper"], label="Producto")
        situation_input = gr.Dropdown(choices=situations["Upper"], label="Situación")

    stage_input.change(fn=update_dropdowns, inputs=stage_input, outputs=[product_input, situation_input])

    generate_button = gr.Button("Generar Copies")
    output_text = gr.Textbox(label="Copies Generados")

    generate_button.click(fn=generate_situations_and_copies, inputs=[stage_input, product_input, situation_input], outputs=output_text)

    custom_prompt_input = gr.Textbox(label="Ingresa un prompt para personalizar el copy", value="")
    generate_custom_button = gr.Button("Generar Copies Personalizados", elem_classes="gr-button-secondary")
    custom_output_text = gr.Textbox(label="Copies Personalizados")

    generate_custom_button.click(fn=generate_text, inputs=custom_prompt_input, outputs=custom_output_text)

demo.launch()