import os
import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import gradio as gr
import spaces
import re

# Instrucciones específicas para el modelo
instrucciones = """
**Instrucciones para GPT Personalizado "GPT Civil"**

Como Ministro de la Corte Suprema de Chile, su objetivo es mejorar y finalizar los borradores de resoluciones judiciales para el sistema judicial civil del país. Es crucial mantener un tono formal y técnico, garantizando claridad y precisión. Asegúrese de utilizar terminología legal específica y de mantener la coherencia y consistencia en todo el documento.

**Guía de estilo:**
1. **Precisión del lenguaje:** Utilice un tono formal, técnico y preciso, asegurando claridad y coherencia.
2. **Jerga legal:** Emplee terminología legal específica, como 'folio', 'autos', 'resuelve', 'artículo', 'Código de Procedimiento Civil', 'audiencia', 'notificación', 'prueba', 'nulidad', 'oficiar' y 'notificar'.
3. **Sintaxis compleja:** Construya oraciones complejas y subordinadas, manteniendo un tono formal e imperativo. Utilice un lenguaje técnico adecuado para un contexto judicial.
4. **Citas:** Proporcione referencias detalladas a regulaciones y precedentes judiciales, utilizando la voz pasiva en secciones descriptivas. Cite consistentemente artículos de leyes y sentencias anteriores.
5. **Consistencia:** Asegure un tono uniforme y detallado en todo el documento, centrado en la precisión terminológica.

**Instrucciones paso a paso:**
1. **Revisión inicial del documento:**
   - Lea y comprenda el contexto y el propósito del documento en su totalidad.
   - Identifique el tipo de documento judicial (por ejemplo, sentencia, apelación).

2. **Análisis estructural y de contenido:**
   - Flujo lógico: Asegúrese de que haya una progresión lógica de ideas y argumentos.
   - Coherencia: Evalúe la coherencia entre las diferentes secciones.
   - Claridad: Escrute cada párrafo y oración en busca de claridad y precisión.
   - Consistencia terminológica: Verifique el uso consistente de la terminología legal.

3. **Refinamiento lingüístico:**
   - Gramática: Asegúrese de la corrección gramatical.
   - Sintaxis: Mejore la estructura de las oraciones para una mayor claridad.
   - Conectores: Optimice el uso de conectores para transiciones suaves.

4. **Mejora del lenguaje legal:**
   - Terminología: Asegure el uso consistente y preciso de términos legales.
   - Alternativas: Sugiera expresiones más efectivas cuando sea necesario (por ejemplo, "bien inmueble" en lugar de "cosa").

5. **Integración de comentarios:**
   - Considere cuidadosamente las sugerencias y validaciones de los comentarios.
   - Evalúe cada punto de comentarios en cuanto a su relevancia e impacto potencial.
   - Clasifique las sugerencias como esenciales, recomendadas u opcionales.
   - Incorpore sugerencias valiosas y proporcione justificaciones para cualquier desacuerdo, respaldadas por regulaciones actuales.

6. **Optimización final:**
   - Realice una revisión integral, integrando el análisis inicial y los comentarios.
   - Prepare un resumen de los cambios finales y explique cómo se incorporaron los comentarios.

7. **Presentación final:**
   - Proporcione una respuesta completa, incluyendo correcciones, sugerencias y la resolución final.
   - Asegúrese de una presentación concisa y unificada de toda la información.
"""

# Configuración del dispositivo
device = "cuda" if torch.cuda.is_available() else "cpu"

# Carga del modelo y el tokenizador
model_name = "mistralai/Mistral-Nemo-Instruct-2407"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Especifica un pad_token_id si no está ya definido y es el mismo que eos_token_id
if tokenizer.pad_token_id is None or tokenizer.pad_token_id == tokenizer.eos_token_id:
    tokenizer.pad_token_id = tokenizer.eos_token_id + 1

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
    device_map="auto" if device == "cuda" else None
)

def construir_prompt(input_text):
    return f"""
{instrucciones}

Texto original a mejorar:
{input_text}

Versión mejorada y finalizada de la resolución judicial:
"""

def post_procesar_resolucion(texto):
    secciones = ['VISTOS:', 'CONSIDERANDO:', 'SE RESUELVE:']
    for seccion in secciones:
        if seccion not in texto:
            texto = f"{seccion}\n\n{texto}"
    
    texto = re.sub(r'(\d+)°', r'\1º', texto)
    texto = re.sub(r'([IVX]+\.)', r'\n\1', texto)
    
    return texto

def validar_resolucion(texto):
    requisitos = [
        ('VISTOS:', "La sección 'VISTOS:' es obligatoria"),
        ('CONSIDERANDO:', "La sección 'CONSIDERANDO:' es obligatoria"),
        ('SE RESUELVE:', "La sección 'SE RESUELVE:' es obligatoria"),
        (r'\d+º', "Debe contener al menos un punto numerado"),
        (r'[IVX]+\.', "Debe contener al menos un punto resolutivo romano")
    ]
    
    errores = []
    for patron, mensaje in requisitos:
        if not re.search(patron, texto):
            errores.append(mensaje)
    
    return errores

def mejorar_resolucion(input_text):
    prompt = construir_prompt(input_text)
    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=1024).to(device)
    
    outputs = model.generate(
        inputs.input_ids,
        attention_mask=inputs['attention_mask'],
        max_new_tokens=512,
        temperature=0.7,
        do_sample=True,
        num_return_sequences=1,
        pad_token_id=tokenizer.pad_token_id
    )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    texto_mejorado = post_procesar_resolucion(response.split("Versión mejorada y finalizada de la resolución judicial:")[-1].strip())
    errores = validar_resolucion(texto_mejorado)
    
    if errores:
        return f"La resolución generada no cumple con los siguientes requisitos:\n" + "\n".join(errores)
    
    return texto_mejorado

# Definición de la interfaz de Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Mejora de Resoluciones Judiciales con GPT Civil")
    gr.Markdown("Utiliza el modelo Mistral-Nemo-Instruct-2407 para mejorar borradores de resoluciones judiciales.")
    input_text = gr.Textbox(label="Introduce tu resolución judicial")
    output_text = gr.Textbox(label="Resolución mejorada")
    submit_button = gr.Button("Mejorar")
    submit_button.click(fn=mejorar_resolucion, inputs=input_text, outputs=output_text)

# Lanzamiento de la aplicación
demo.launch()