import numpy as np
from gensim.models import Word2Vec
import gradio as gr
from sklearn.decomposition import PCA
import plotly.graph_objects as go
import nltk
from nltk.tokenize import word_tokenize
from nltk.tag import pos_tag

nltk.download('punkt')
nltk.download('averaged_perceptron_tagger')

# Word2Vec 모델 학습 함수
def train_word2vec(sentences):
    model = Word2Vec(sentences, vector_size=100, window=5, min_count=1)
    return model

def preprocess_text(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        text = file.read()
    
    # 토큰화 및 품사 태깅
    tokens = word_tokenize(text)
    tagged = pos_tag(tokens)
    
    # 명사만 추출 (NN, NNS, NNP, NNPS)
    nouns = [word.lower() for word, pos in tagged if pos.startswith('NN')]
    
    # 중복 제거 및 정렬
    unique_nouns = sorted(set(nouns))
    
    # 간단한 문장 생성 (각 명사를 개별 문장으로 취급)
    sentences = [[noun] for noun in unique_nouns]
    
    return sentences, unique_nouns

def apply_pca(word_vectors):
    pca = PCA(n_components=3)
    return pca.fit_transform(word_vectors)

def process_text(file_path, target_word):
    # 전처리
    sentences, unique_words = preprocess_text(file_path)
    
    # Word2Vec 모델 학습
    model = train_word2vec(sentences)
    
    # 각 단어의 임베딩 벡터 추출
    word_vectors = np.array([model.wv[word] for word in unique_words])
    
    # PCA로 차원 축소
    word_vectors_3d = apply_pca(word_vectors)
    
    # 색상 설정 (투명도 추가)
    colors = ['rgba(128, 128, 128, 0.3)' if word != target_word else 'rgba(255, 0, 0, 1)' for word in unique_words]
    
    # 가장 가까운 단어 10개 찾기
    if target_word in model.wv:
        similar_words = model.wv.most_similar(target_word, topn=10)
        similar_word_indices = [unique_words.index(word) for word, _ in similar_words]
        for idx in similar_word_indices:
            colors[idx] = 'rgba(0, 255, 0, 1)'  # 가까운 단어들을 초록색으로 표시
    
    # Plotly를 사용한 3D 산점도 생성
    fig = go.Figure(data=[go.Scatter3d(
        x=word_vectors_3d[:, 0],
        y=word_vectors_3d[:, 1],
        z=word_vectors_3d[:, 2],
        mode='markers+text',
        text=unique_words,
        textposition="top center",
        marker=dict(
            size=8,
            color=colors,
        )
    )])
    
    fig.update_layout(
        title="Word Embeddings 3D Visualization",
        scene=dict(
            xaxis_title="PCA 1",
            yaxis_title="PCA 2",
            zaxis_title="PCA 3"
        ),
        width=800,
        height=800
    )
    
    # 가장 가까운 단어 10개 목록 생성
    similar_words_text = ""
    if target_word in model.wv:
        similar_words_text = "가장 가까운 단어 10개:\n" + "\n".join([f"{word}: {score:.4f}" for word, score in similar_words])
    
    return fig, similar_words_text

# Gradio 인터페이스 수정
with gr.Blocks() as iface:
    gr.Markdown("# Word Embedding 3D 시각화")
    gr.Markdown("텍스트 파일(.txt)을 업로드하고 강조할 단어를 입력하세요. Word2Vec과 PCA를 사용하여 단어 임베딩을 3D로 시각화합니다. 입력한 단어는 빨간색으로, 가장 유사한 10개 단어는 초록색으로 강조됩니다. 유사한 단어 목록은 그래프 아래에 표시됩니다.")
    
    with gr.Row():
        file_input = gr.File(label="텍스트 파일 업로드 (.txt)", file_types=[".txt"])
        word_input = gr.Textbox(label="강조할 단어 입력")
    
    submit_btn = gr.Button("제출")
    
    plot_output = gr.Plot(label="Word Embedding 3D 시각화")
    similar_words_output = gr.Textbox(label="유사한 단어")
    
    submit_btn.click(
        fn=process_text,
        inputs=[file_input, word_input],
        outputs=[plot_output, similar_words_output]
    )

if __name__ == "__main__":
    iface.launch()