import torch
from transformers import pipeline
import sys
import sys

# Cargar el pipeline de generación de texto
pipe = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-beta", torch_dtype=torch.bfloat16, device_map="auto")

# Definir el prompt para generar un JSON con eventos anidados
prompt = (
    "Genera un JSON que describa una serie de eventos consecutivos en un formato similar al siguiente:\n\n"
    "{\n"
    "  \"events\": {\n"
    "    \"event\": {\n"
    "      \"event_number\": 1,\n"
    "      \"name\": \"conflict_start\",\n"
    "      \"description\": \"Tensions escalate between Iran and Israel\",\n"
    "      \"probability\": 70,\n"
    "      \"duration_days\": 30,\n"
    "      \"subevents\": {\n"
    "        \"event\": {\n"
    "          \"event_number\": 2,\n"
    "          \"name\": \"diplomatic_failure\",\n"
    "          \"description\": \"Diplomatic negotiations fail\",\n"
    "          \"probability\": 60,\n"
    "          \"duration_days\": 15,\n"
    "          \"subevents\": {\n"
    "            \"event\": {\n"
    "              \"event_number\": 3,\n"
    "              \"name\": \"military_clash\",\n"
    "              \"description\": \"Initial military clash at the border\",\n"
    "              \"probability\": 50,\n"
    "              \"duration_days\": 10,\n"
    "              \"subevents\": {\n"
    "                \"event\": [\n"
    "                  {\n"
    "                    \"event_number\": 4,\n"
    "                    \"name\": \"escalation\",\n"
    "                    \"description\": \"Conflict escalates into full-scale war\",\n"
    "                    \"probability\": 40,\n"
    "                    \"duration_days\": 180,\n"
    "                    \"subevents\": {\n"
    "                      \"event\": [\n"
    "                        {\n"
    "                          \"event_number\": 5,\n"
    "                          \"name\": \"regional_involvement\",\n"
    "                          \"description\": \"Other Middle Eastern countries get involved\",\n"
    "                          \"probability\": 30,\n"
    "                          \"duration_days\": 365,\n"
    "                          \"subevents\": {\n"
    "                            \"event\": [\n"
    "                              {\n"
    "                                \"event_number\": 6,\n"
    "                                \"name\": \"ceasefire\",\n"
    "                                \"description\": \"International powers broker a ceasefire\",\n"
    "                                \"probability\": 20,\n"
    "                                \"duration_days\": 30\n"
    "                              },\n"
    "                              {\n"
    "                                \"event_number\": 7,\n"
    "                                \"name\": \"prolonged_conflict\",\n"
    "                                \"description\": \"Conflict continues for over a year\",\n"
    "                                \"probability\": 50,\n"
    "                                \"duration_days\": 365\n"
    "                              }\n"
    "                            ]\n"
    "                          }\n"
    "                        },\n"
    "                        {\n"
    "                          \"event_number\": 8,\n"
    "                          \"name\": \"international_intervention\",\n"
    "                          \"description\": \"UN or other international organizations intervene\",\n"
    "                          \"probability\": 25,\n"
    "                          \"duration_days\": 60\n"
    "                        }\n"
    "                      ]\n"
    "                    }\n"
    "                  },\n"
    "                  {\n"
    "                    \"event_number\": 9,\n"
    "                    \"name\": \"containment\",\n"
    "                    \"description\": \"Conflict is contained and doesn't escalate\",\n"
    "                    \"probability\": 30,\n"
    "                    \"duration_days\": 90\n"
    "                  }\n"
    "                ]\n"
    "              }\n"
    "            },\n"
    "            \"event\": {\n"
    "              \"event_number\": 10,\n"
    "              \"name\": \"sanctions\",\n"
    "              \"description\": \"Increased sanctions on Iran\",\n"
    "              \"probability\": 70,\n"
    "              \"duration_days\": 180,\n"
    "              \"subevents\": {\n"
    "                \"event\": [\n"
    "                  {\n"
    "                    \"event_number\": 11,\n"
    "                    \"name\": \"iran_retaliates\",\n"
    "                    \"description\": \"Iran retaliates with cyberattacks\",\n"
    "                    \"probability\": 40,\n"
    "                    \"duration_days\": 60\n"
    "                  },\n"
    "                  {\n"
    "                    \"event_number\": 12,\n"
    "                    \"name\": \"israel_response\",\n"
    "                    \"description\": \"Israel responds with targeted airstrikes\",\n"
    "                    \"probability\": 50,\n"
    "                    \"duration_days\": 60\n"
    "                  }\n"
    "                ]\n"
    "              }\n"
    "            }\n"
    "          }\n"
    "        },\n"
    "        \"event\": {\n"
    "          \"event_number\": 13,\n"
    "          \"name\": \"diplomatic_success\",\n"
    "          \"description\": \"Successful diplomatic negotiations\",\n"
    "          \"probability\": 40,\n"
    "          \"duration_days\": 30,\n"
    "          \"subevents\": {\n"
    "            \"event\": [\n"
    "              {\n"
    "                \"event_number\": 14,\n"
    "                \"name\": \"peace_agreement\",\n"
    "                \"description\": \"Iran and Israel sign a peace agreement\",\n"
    "                \"probability\": 20,\n"
    "                \"duration_days\": 60\n"
    "              },\n"
    "              {\n"
    "                \"event_number\": 15,\n"
    "                \"name\": \"temporary_truce\",\n"
    "                \"description\": \"A temporary truce is established\",\n"
    "                \"probability\": 30,\n"
    "                \"duration_days\": 30\n"
    "              }\n"
    "            ]\n"
    "          }\n"
    "        }\n"
    "      }\n"
    "    }\n"
    "  }\n"
    "}\n\n"
    "Ahora, genera un JSON similar con eventos anidados, pero cambia los detalles y números para hacer que sea con el input que viene a continuacion, respondiendo solo el JSON: "
)


def generate():
    # Generar el texto usando el modelo
    outputs = pipe(prompt+" "+sys.argv[1], max_new_tokens=1024, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)

    # Imprimir la salida generada
    print(outputs[0]["generated_text"])
    return outputs[0]["generated_text"]