import os
import gradio as gr
import pandas as pd
import numpy as np
from typing import Dict, List

from openai import OpenAI
from dotenv import load_dotenv

# Load environment variables
load_dotenv()


# 1) Вкажіть свій OpenAI ключ
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))


##############################################################################
# 1. Вихідні дані: JSON із "хінтами"
##############################################################################
classes_json = {
    "Pain": [
        "ache", "aches", "hurts", "pain", "painful", "sore"
        # ...
    ],
    "Chest pain": [
        "aches in my chest", "chest pain", "chest hurts", "sternum pain"
    ],
    "Physical Activity": [
        "exercise", "walking", "running", "biking"
    ],
    "Office visit": [
        "appointment scheduled", "annual checkup", "office visit"
    ],
    # ...
}

##############################################################################
# 2. Глобальні змінні (спрощено)
##############################################################################
df = None
embeddings = None
class_signatures = None

##############################################################################
# 3. Функція для завантаження даних
##############################################################################
def load_data(csv_path: str = "messages.csv", emb_path: str = "embeddings.npy"):
    global df, embeddings
    df_local = pd.read_csv(csv_path)
    emb_local = np.load(emb_path)
    assert len(df_local) == len(emb_local), "CSV і embeddings різної довжини!"

    df_local["Target"] = "Unlabeled"

    # Нормалізація embeddings
    emb_local = (emb_local - emb_local.mean(axis=0)) / emb_local.std(axis=0)

    df = df_local
    embeddings = emb_local

##############################################################################
# 4. Виклик OpenAI для отримання одного embedding
##############################################################################
def get_openai_embedding(text: str, model_name: str = "text-embedding-3-small") -> list:
    response = client.embeddings.create(
        input=text,
        model=model_name
    )
    return response.data[0].embedding

##############################################################################
# 5. Отримати embeddings для списку фраз (хінтів) і усереднити
##############################################################################
def embed_hints(hint_list: List[str], model_name: str) -> np.ndarray:
    emb_list = []
    for hint in hint_list:
        emb = get_openai_embedding(hint, model_name=model_name)
        emb_list.append(emb)
    return np.array(emb_list, dtype=np.float32)

##############################################################################
# 6. Будуємо signatures для кожного класу
##############################################################################
def build_class_signatures(model_name: str):
    global class_signatures
    signatures = {}
    for cls_name, hints in classes_json.items():
        if not hints:
            continue
        arr = embed_hints(hints, model_name=model_name)
        signatures[cls_name] = arr.mean(axis=0)
    class_signatures = signatures
    return "Signatures побудовано!"

##############################################################################
# 7. Функція класифікації одного рядка (dot product)
##############################################################################
def predict_class(text_embedding: np.ndarray, signatures: Dict[str, np.ndarray]) -> str:
    best_label = "Unknown"
    best_score = float("-inf")
    for cls, sign in signatures.items():
        score = np.dot(text_embedding, sign)
        if score > best_score:
            best_score = score
            best_label = cls
    return best_label

##############################################################################
# 8. Класифікація відфільтрованих рядків
##############################################################################
def classify_rows(filter_substring: str):
    global df, embeddings, class_signatures

    if class_signatures is None:
        return "Спочатку збудуйте signatures!"

    if df is None or embeddings is None:
        return "Дані не завантажені! Спочатку викличте load_data."

    if filter_substring:
        filtered_idx = df[df["Message"].str.contains(filter_substring, case=False, na=False)].index
    else:
        filtered_idx = df.index

    for i in filtered_idx:
        emb_vec = embeddings[i]
        pred = predict_class(emb_vec, class_signatures)
        df.at[i, "Target"] = pred

    result_df = df.loc[filtered_idx, ["Message", "Target"]].copy()
    return result_df.reset_index(drop=True)

##############################################################################
# 9. Збереження CSV
##############################################################################
def save_data():
    global df
    if df is None:
        return "Дані відсутні!"
    df.to_csv("messages_with_labels.csv", index=False)
    return "Файл 'messages_with_labels.csv' збережено!"

##############################################################################
# 10. Gradio UI
##############################################################################
def ui_load_data(csv_path, emb_path):
    load_data(csv_path, emb_path)
    return f"Data loaded from {csv_path} and {emb_path}. Rows: {len(df)}"

def ui_build_signatures(model_name):
    msg = build_class_signatures(model_name)
    return msg

def ui_classify_data(filter_substring):
    result = classify_rows(filter_substring)
    if isinstance(result, str):
        return result
    return result

def ui_save_data():
    return save_data()

def main():
    import gradio as gr

    with gr.Blocks() as demo:
        gr.Markdown("# SDC Classifier з Gradio")
        gr.Markdown("## 1) Завантаження даних")

        with gr.Row():
            csv_input = gr.Textbox(value="messages.csv", label="CSV-файл")
            emb_input = gr.Textbox(value="embeddings.npy", label="Numpy Embeddings")
            load_btn = gr.Button("Load data")

        load_output = gr.Label(label="Loading result")

        load_btn.click(fn=ui_load_data, inputs=[csv_input, emb_input], outputs=load_output)

        gr.Markdown("## 2) Побудова Class Signatures")
        # openai_key_in = gr.Textbox(label="OpenAI API Key", type="password")
        model_choice = gr.Dropdown(choices=["text-embedding-3-large","text-embedding-3-small"],
                                   value="text-embedding-3-small", label="OpenAI model")
        build_btn = gr.Button("Build signatures")
        build_out = gr.Label(label="Signatures")

        build_btn.click(fn=ui_build_signatures, inputs=[model_choice], outputs=build_out)

        gr.Markdown("## 3) Класифікація")
        filter_in = gr.Textbox(label="Filter substring (optional)")
        classify_btn = gr.Button("Classify")
        classify_out = gr.Dataframe(label="Result (Message / Target)")

        classify_btn.click(fn=ui_classify_data, inputs=[filter_in], outputs=[classify_out])

        gr.Markdown("## 4) Зберегти CSV")
        save_btn = gr.Button("Save labeled data")
        save_out = gr.Label()

        save_btn.click(fn=ui_save_data, inputs=[], outputs=save_out)

        gr.Markdown("""
        ### Опис:
        1. Натисніть 'Load data', щоб завантажити ваші дані (CSV + embeddings).
        2. Укажіть OpenAI API модель, натисніть 'Build signatures'.
        3. Вкажіть фільтр (необов'язково), натисніть 'Classify'. 
           Отримаєте таблицю з полем Target. 
        4. 'Save labeled data' збереже 'messages_with_labels.csv'.
        """)
    
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
    # demo.launch()


if __name__ == "__main__":
    main()