import gradio as gr
import pandas as pd
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import networkx as nx
import matplotlib.pyplot as plt

# Sentence-BERT 모델 로드
model = SentenceTransformer('all-MiniLM-L6-v2')

# 직원 데이터를 분석하여 교육 프로그램을 추천하고 그래프를 그리는 함수
def analyze_data(employee_file, program_file):
    # 직원 데이터와 교육 프로그램 데이터 불러오기
    employee_df = pd.read_csv(employee_file.name)
    program_df = pd.read_csv(program_file.name)

    # 직원 역량과 프로그램 학습 목표를 벡터화
    employee_skills = employee_df['current_skills'].tolist()
    program_skills = program_df['skills_acquired'].tolist()
    employee_embeddings = model.encode(employee_skills)
    program_embeddings = model.encode(program_skills)

    # 유사도 계산
    similarities = cosine_similarity(employee_embeddings, program_embeddings)

    # 직원별 추천 프로그램 리스트
    recommendations = []
    for i, employee in employee_df.iterrows():
        recommended_programs = []
        for j, program in program_df.iterrows():
            if similarities[i][j] > 0.5:  # 유사도 임계값 기준
                recommended_programs.append(f"{program['program_name']} ({program['duration']})")
        
        if recommended_programs:
            recommendation = f"직원 {employee['employee_name']}의 추천 프로그램: {', '.join(recommended_programs)}"
        else:
            recommendation = f"직원 {employee['employee_name']}에게 적합한 프로그램이 없습니다."
        
        recommendations.append(recommendation)

    # 결과 텍스트
    result_text = "\n".join(recommendations)

    # 네트워크 그래프 생성
    G = nx.Graph()
    for employee in employee_df['employee_name']:
        G.add_node(employee, type='employee')

    for program in program_df['program_name']:
        G.add_node(program, type='program')

    for i, employee in employee_df.iterrows():
        for j, program in program_df.iterrows():
            if similarities[i][j] > 0.5:  # 유사도 임계값
                G.add_edge(employee['employee_name'], program['program_name'])

    # 그래프 시각화
    plt.figure(figsize=(10, 8))
    pos = nx.spring_layout(G)
    nx.draw(G, pos, with_labels=True, node_color='skyblue', node_size=2000, font_size=10, font_weight='bold')
    plt.title("직원과 프로그램 간의 관계")
    plt.tight_layout()

    return result_text, plt.gcf()

# Gradio 인터페이스 정의
def main(employee_file, program_file):
    return analyze_data(employee_file, program_file)

# Gradio 블록
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("# HybridRAG 시스템")
            gr.Markdown("두 개의 CSV 파일을 업로드하여 분석을 진행하세요.")
            
            employee_file = gr.File(label="직원 데이터 업로드")
            program_file = gr.File(label="교육 프로그램 데이터 업로드")
            analyze_button = gr.Button("분석 시작")
            output_text = gr.Textbox(label="분석 결과")

            analyze_button.click(main, inputs=[employee_file, program_file], outputs=[output_text])

        with gr.Column(scale=2):
            gr.Markdown("### 정보 패널")
            gr.Markdown("업로드된 데이터에 대한 분석 및 결과를 여기에 표시합니다.")
            
            # 시각화 차트 출력
            chart_output = gr.Plot(label="시각화 차트")
            
            # 분석 버튼 클릭 시 차트 업데이트
            analyze_button.click(main, inputs=[employee_file, program_file], outputs=[output_text, chart_output])

# Gradio 인터페이스 실행
demo.launch()