from datasets import load_dataset, concatenate_datasets, Dataset
from sentence_transformers import SentenceTransformer
import numpy as np
import pandas as pd
import gradio as gr
import kagglehub
from kagglehub import KaggleDatasetAdapter
import faiss
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import unicodedata
import random


# ==================================================
# Cargar datasets
# ==================================================

# Dataset 1: MyDramaList (5000 dramas)
mydramalist = kagglehub.load_dataset(
    KaggleDatasetAdapter.PANDAS,
    "anittasaju/complete-5000-dramas-from-mydramalist-review-info",
    "Top_5000_popular_drama_details_from_mydramalist.csv"
)
mydramalist = Dataset.from_pandas(mydramalist)

# Dataset 2: Netflix Movies and Shows
netflix_movies_shows = load_dataset("harshi321/netflix-movies_shows")["train"]

# ==================================================
# Preprocesamiento de datos
# ==================================================

# Renombrar columnas para unificar los datasets
mydramalist = mydramalist.rename_column("name", "title").rename_column("content", "description")

# Filtrar por país (Corea del Sur)
def filter_kdramas(dataset):
    if 'country' in dataset.features:
        return dataset.filter(lambda x: x['country'] == "South Korea" if x['country'] else False)
    return dataset

# Aplicar filtro a los datasets
kdramas1 = filter_kdramas(netflix_movies_shows)
kdramas2 = filter_kdramas(mydramalist)

# Eliminar columnas innecesarias (incluyendo 'rating')
columns_to_remove = ["Unnamed: 0", "no_of_reviews", "aka_names", "screenwriter", "director",
                     "no_of_viewers", "end_date", "start_date", "year", "duration", "no_of_rating",
                     "rank", "popularity", "content_rating", "where_to_watch", "main_role",
                     "support_role", "no_of_extracted_reviews", "Total_sentences",
                     "POSITIVE_people_sentiment", "POSITIVE_sentences", "NEGATIVE_people_sentiment",
                     "NEGATIVE_sentences", "rating"]  # Eliminar 'rating'

for dataset in [kdramas1, kdramas2]:
    dataset = dataset.remove_columns([col for col in columns_to_remove if col in dataset.features])

# Asegurar que todos los datasets tengan la columna 'type'
def add_type_column(dataset, default_type="TV Show"):
    if 'type' not in dataset.features:
        dataset = dataset.map(lambda x: {"type": default_type})
    return dataset

kdramas1 = add_type_column(kdramas1, "TV Show")
kdramas2 = add_type_column(kdramas2, "TV Show")

# Asegurar que todos los datasets tengan la columna 'genres'
def add_genres_column(dataset, default_genres="Unknown"):
    if 'genres' not in dataset.features:
        dataset = dataset.map(lambda x: {"genres": default_genres})
    return dataset

kdramas1 = add_genres_column(kdramas1, "Unknown")
kdramas2 = add_genres_column(kdramas2, "Unknown")

# Asegurar que todos los datasets tengan las mismas columnas y tipos
def align_datasets(dataset1, dataset2):
    # Obtener las columnas comunes
    common_columns = set(dataset1.features.keys()).intersection(set(dataset2.features.keys()))

    # Mantener solo las columnas comunes
    dataset1 = dataset1.select_columns(list(common_columns))
    dataset2 = dataset2.select_columns(list(common_columns))

    return dataset1, dataset2

# Alinear los datasets
kdramas1, kdramas2 = align_datasets(kdramas1, kdramas2)

# Eliminar columnas adicionales que no se usan
kdramas1 = kdramas1.remove_columns(['rating'])
kdramas2 = kdramas2.remove_columns(['rating'])

# Combinar todos los datasets
kdramas = concatenate_datasets([kdramas1, kdramas2])

# Eliminar duplicados basados en el título
kdramas_df = kdramas.to_pandas().drop_duplicates(subset=['title'])
kdramas = Dataset.from_pandas(kdramas_df)

# ==================================================
# Modelo de embeddings y recomendación
# ==================================================

# Cargar el modelo de embeddings
model = SentenceTransformer('sentence-transformers/paraphrase-MiniLM-L6-v2')

# Calcular o cargar embeddings
# Verificar si el archivo de embeddings existe
if os.path.exists("kdrama_embeddings.npy"):
    # Cargar embeddings precalculados
    embeddings_np = np.load("kdrama_embeddings.npy")
else:
    # Calcular embeddings y guardarlos
    descriptions = kdramas["description"]
    embeddings = model.encode(descriptions, convert_to_tensor=True)
    embeddings_np = embeddings.cpu().numpy()
    np.save("kdrama_embeddings.npy", embeddings_np)
    print("¡Embeddings listos! Cada descripción ahora es un vector numérico.")

# Crear un índice FAISS para búsqueda eficiente
dimension = embeddings_np.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(embeddings_np)

# Función para recomendar K-Dramas similares
def recommend_kdramas(title, k=5):
    title_indices = [i for i, t in enumerate(kdramas['title']) if title.lower() in t.lower()]
    if not title_indices:
        return f"No se encontraron títulos similares a '{title}'."
    
    query_embedding = embeddings_np[title_indices[0]].reshape(1, -1)
    distances, similar_indices = index.search(query_embedding, k + 10)  # Ampliar el rango de búsqueda
    
    # Seleccionar aleatoriamente k índices de los 10 más similares
    selected_indices = random.sample(list(similar_indices[0][1:]), k)
    
    recommendations = []
    for i in selected_indices:
        recommended_title = kdramas["title"][i]
        recommended_type = kdramas["type"][i] if "type" in kdramas.features else "Unknown"
        recommended_genres = kdramas["genres"][i] if "genres" in kdramas.features else "Unknown"
        recommendations.append(
            f"### {recommended_title}\n"
            f"- **Tipo**: {recommended_type}\n"
            f"- **Géneros**: {recommended_genres}\n"
        )
    return "\n".join(recommendations)

# ==================================================
# Chatbot
# ==================================================

# Cargar el modelo de chat
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
model_chat = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")

# Mapeo de géneros (español sin tildes → inglés)
mapeo_generos = {
    # Géneros (sin tildes)
    "romantico": "romantic",
    "accion": "action",
    "misterio": "mystery",
    "comedia": "comedy",
    "drama": "drama",
    "crimen": "crime",
    "fantasia": "fantasy",
    "thriller": "thriller",
    "romance": "romance",
    # Palabras clave adicionales (sin tildes)
    "aventura": "adventure",
    "historico": "historical",
    # Agrega más según sea necesario
}

# Función para normalizar texto
def normalizar_texto(texto):
    # Eliminar tildes y convertir a minúsculas
    texto_normalizado = ''.join(
        letra for letra in unicodedata.normalize('NFD', texto)
        if unicodedata.category(letra) != 'Mn'
    )
    return texto_normalizado.lower()

# Función para traducir preferencias
def traducir_preferencia(entrada_usuario):
    entrada_normalizada = normalizar_texto(entrada_usuario)
    for espanol, ingles in mapeo_generos.items():
        if espanol in entrada_normalizada:
            return ingles
    return None  # Si no se encuentra una traducción

# Función para buscar K-Dramas por género
def buscar_por_genero(genero, k=5):
    # Filtrar K-Dramas que contengan el género especificado
    genre_embedding = model.encode(genero, convert_to_tensor=True).cpu().numpy()
    genre_embedding = genre_embedding.reshape(1, -1)

    # Buscar en el índice FAISS
    D, I = index.search(genre_embedding, k * 2)  # Ampliar el rango de búsqueda

    # Seleccionar aleatoriamente k índices de los resultados
    selected_indices = random.sample(list(I[0]), k)

    # Formatear las recomendaciones
    recommendations = []
    for i in selected_indices:
        recommended_title = kdramas["title"][i]
        recommended_type = kdramas["type"][i] if "type" in kdramas.features else "Unknown"
        recommended_genres = kdramas["genres"][i] if "genres" in kdramas.features else "Unknown"
        recommendations.append(
            f"{recommended_title}\n"
            f"- **Tipo**: {recommended_type}\n"
            f"- **Géneros**: {recommended_genres}\n"
        )
    return "\n".join(recommendations) if recommendations else f"No se encontraron K-Dramas del género '{genero}'."

# Función para recomendar K-Dramas basado en preferencias
def recomendar_kdramas_chat(entrada_usuario):
    # Traducir preferencia del usuario
    preferencia_traducida = traducir_preferencia(entrada_usuario)

    if preferencia_traducida:
        return buscar_por_genero(preferencia_traducida, k=5)
    else:
        return None  # No se encontró una preferencia válida

# Función para generar respuestas del chatbot
def generar_respuesta(entrada_usuario, historial_chat=""):
    inputs = tokenizer.encode(entrada_usuario + historial_chat, return_tensors="pt")
    respuesta_ids = model_chat.generate(inputs, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    respuesta = tokenizer.decode(respuesta_ids[:, inputs.shape[-1]:][0], skip_special_tokens=True)
    return respuesta

# Función para manejar la interacción del chatbot
def chat(entrada_usuario, historial_chat=""):
    # Agregar el mensaje del usuario al historial
    historial_chat += f"Usuario: {entrada_usuario}\n"

    # Generar una respuesta natural del chatbot
    if not historial_chat.strip():  # Si es la primera interacción
        respuesta = "¡Hola! Soy tu asistente para recomendar K-Dramas. ¿Qué tipo de K-Drama te gustaría ver hoy?"
    else:
        respuesta = generar_respuesta(entrada_usuario, historial_chat)

        # Verificar si el usuario mencionó un género
        recomendacion = recomendar_kdramas_chat(entrada_usuario)
        if recomendacion:
            respuesta = f"¡Genial! Aquí tienes algunas recomendaciones de K-Dramas que podrían gustarte:\n{recomendacion}"
        else:
            respuesta = "¿Qué tipo de K-Drama te gustaría ver? Puedo recomendarte dramas de acción, romance, misterio y más."

    # Agregar la respuesta del chatbot al historial
    historial_chat += f"Chatbot: {respuesta}\n"

    # Devolver la respuesta y el historial actualizado
    return respuesta, historial_chat, ""  # Limpiar el cuadro de texto

# ==================================================
# Interfaz de Gradio
# ==================================================

# Interfaz para el chatbot
interfaz_chatbot = gr.Interface(
    fn=chat,
    inputs=[gr.Textbox(label="Escribe tu mensaje"), gr.Textbox(label="Historial", visible=False)],
    outputs=[gr.Textbox(label="Respuesta del chatbot"), gr.Textbox(label="Mensaje", visible=False)],
    title="Chatbot Recomendador de K-Dramas",
    description="Habla con el chatbot para obtener recomendaciones personalizadas de K-Dramas.",
    allow_flagging="never"
)
# Interfaz para el recomendador tradicional
interfaz_recomendador = gr.Interface(
    theme=gr.themes.Citrus(),
    fn=recommend_kdramas,
    inputs=[
        gr.Textbox(label="Ingresa el título de un K-Drama"),
    ],
    outputs=gr.Markdown(label="K-Dramas recomendados"),
    title="Recomendador de K-Dramas",
    description="Ingresa el título de un K-Drama y recibe recomendaciones similares, indicando si son series o películas y sus géneros.",
    examples=[["Vincenzo", 5], ["Crash Landing on You", 3], ["Itaewon Class", 4]],
    allow_flagging="never",
)

# Lanzar ambas interfaces
gr.TabbedInterface(
    [interfaz_recomendador, interfaz_chatbot],  # Chatbot en la segunda pestaña
    ["Recomendador", "Chatbot"]
).launch()