import json
import gradio as gr
import spaces
import wbgtopic
import plotly.graph_objects as go
import plotly.express as px
import plotly.figure_factory as ff
import nltk
import numpy as np
import pandas as pd
from collections import Counter
from scipy import stats
import torch
from wordcloud import WordCloud
from topic_translator import translate_topics
from nltk.tokenize import sent_tokenize, word_tokenize
from nltk.sentiment import SentimentIntensityAnalyzer
from sklearn.cluster import KMeans

# GPU 설정
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# NLTK 필요 데이터 다운로드
try:
    nltk.download('punkt', quiet=True)
    nltk.download('vader_lexicon', quiet=True)
except Exception as e:
    print(f"NLTK 데이터 다운로드 중 오류 발생: {e}")

SAMPLE_TEXT = """
The three reportedly discussed the Stargate Project, a large-scale AI initiative led by OpenAI, SoftBank, and U.S. software giant Oracle. The project aims to invest $500 billion over the next four years in building new AI infrastructure in the U.S. The U.S. government has shown a strong commitment to the initiative, with President Donald Trump personally announcing it at the White House the day after his inauguration last month. If Samsung participates, the project will lead to a Korea-U.S.-Japan AI alliance.
The AI sector requires massive investments and extensive resources, including advanced models, high-performance AI chips to power the models, and large-scale data centers to operate them. Nvidia and TSMC currently dominate the AI sector, but a partnership between Samsung, SoftBank, and OpenAI could pave the way for a competitive alternative.
"""

# WBGDocTopic 초기화 시 device 지정
clf = wbgtopic.WBGDocTopic(device=device)

def safe_process(func):
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            print(f"Error in {func.__name__}: {str(e)}")
            return None
    return wrapper

@safe_process
def analyze_text_sections(text):
    sentences = sent_tokenize(text)
    sections = [' '.join(sentences[i:i+3]) for i in range(0, len(sentences), 3)]
    section_topics = []
    
    for section in sections:
        topics = clf.suggest_topics(section)[0]
        section_topics.append(topics)
    
    return section_topics

@safe_process
def calculate_topic_correlations(topics):
    topic_scores = {}
    for topic in topics:
        topic_scores[topic['label']] = topic['score_mean']
    
    if len(topic_scores) < 2:
        return np.array([[1]]), list(topic_scores.keys())
    
    correlation_matrix = np.corrcoef(list(topic_scores.values()))
    return correlation_matrix, list(topic_scores.keys())

@safe_process
def perform_sentiment_analysis(text):
    sia = SentimentIntensityAnalyzer()
    sentences = sent_tokenize(text)
    sentiments = [sia.polarity_scores(sent) for sent in sentences]
    return pd.DataFrame(sentiments)

@safe_process
def create_topic_clusters(topics):
    if len(topics) < 3:
        return np.zeros(len(topics))
    
    X = np.array([[t['score_mean'], t['score_std']] for t in topics])
    kmeans = KMeans(n_clusters=min(3, len(topics)), random_state=42)
    clusters = kmeans.fit_predict(X)
    return clusters

@safe_process
def create_main_charts(topics):
    bar_fig = go.Figure()
    bar_fig.add_trace(go.Bar(
        x=[t['label'] for t in topics],
        y=[t['score'] for t in topics],
        name='관련도',
        marker_color='rgb(55, 83, 109)'
    ))
    bar_fig.update_layout(
        title='주제 분석 결과',
        height=500,
        xaxis_title='주제',
        yaxis_title='관련도 (%)',
        template='plotly_white'
    )
    
    radar_fig = go.Figure()
    radar_fig.add_trace(go.Scatterpolar(
        r=[t['score'] for t in topics],
        theta=[t['label'] for t in topics],
        fill='toself',
        name='주제 분포'
    ))
    radar_fig.update_layout(
        title='주제 레이더 차트',
        height=500,
        template='plotly_white'
    )
    
    return bar_fig, radar_fig

@safe_process
def create_correlation_heatmap(corr_matrix, labels):
    fig = go.Figure(data=go.Heatmap(
        z=corr_matrix,
        x=labels,
        y=labels,
        colorscale='Viridis'
    ))
    fig.update_layout(
        title='주제 간 상관관계',
        height=500,
        template='plotly_white'
    )
    return fig

@safe_process
def create_topic_evolution(section_topics):
    if not section_topics or len(section_topics) == 0:
        return go.Figure()

    fig = go.Figure()
    for topic in section_topics[0]:
        try:
            topic_scores = [topics[topic['label']]['score_mean'] 
                          for topics in section_topics]
            fig.add_trace(go.Scatter(
                x=list(range(len(section_topics))),
                y=topic_scores,
                name=topic['label'],
                mode='lines+markers'
            ))
        except Exception as e:
            print(f"Error processing topic {topic['label']}: {e}")
            continue
    
    fig.update_layout(
        title='주제 변화 추이',
        xaxis_title='섹션',
        yaxis_title='관련도',
        height=500,
        template='plotly_white'
    )
    return fig

@safe_process
def create_confidence_gauge(topics):
    fig = go.Figure()
    for i, topic in enumerate(topics):
        fig.add_trace(go.Indicator(
            mode="gauge+number",
            value=topic['confidence'],
            title={'text': topic['label']},
            domain={'row': 0, 'column': i, 'x': [i/len(topics), (i+1)/len(topics)]}
        ))
    fig.update_layout(
        grid={'rows': 1, 'columns': len(topics)},
        height=400,
        template='plotly_white'
    )
    return fig

@safe_process
def process_results(results):
    if not results or not results[0]:
        return []
    
    topics = results[0]
    top_topics = sorted(topics, key=lambda x: x['score_mean'], reverse=True)[:5]
    
    formatted_topics = []
    for topic in top_topics:
        formatted_topic = {
            'label': translate_topics.get(topic['label'], topic['label']),
            'score': round(topic['score_mean'] * 100, 1),
            'confidence': round((1 - topic['score_std']) * 100, 1)
        }
        formatted_topics.append(formatted_topic)
    
    return formatted_topics

@spaces.GPU(enable_queue=True, duration=50)
def process_all_analysis(text):
    try:
        # 기본 주제 분석
        raw_results = clf.suggest_topics(text)
        topics = process_results(raw_results)
        
        # 추가 분석
        section_topics = analyze_text_sections(text)
        corr_matrix, labels = calculate_topic_correlations(topics)
        sentiments = perform_sentiment_analysis(text)
        clusters = create_topic_clusters(topics)
        
        # 차트 생성
        bar_chart, radar_chart = create_main_charts(topics)
        heatmap = create_correlation_heatmap(corr_matrix, labels)
        evolution_chart = create_topic_evolution(section_topics)
        gauge_chart = create_confidence_gauge(topics)
        
        return {
            'topics': topics,
            'bar_chart': bar_chart,
            'radar_chart': radar_chart,
            'heatmap': heatmap,
            'evolution': evolution_chart,
            'gauge': gauge_chart,
            'sentiments': sentiments.to_dict() if sentiments is not None else {},
            'clusters': clusters.tolist() if clusters is not None else []
        }
    except Exception as e:
        print(f"Analysis error: {str(e)}")
        return {
            'error': str(e),
            'topics': [],
            'bar_chart': go.Figure(),
            'radar_chart': go.Figure(),
            'heatmap': go.Figure(),
            'evolution': go.Figure(),
            'gauge': go.Figure(),
            'sentiments': {},
            'clusters': []
        }

# Gradio 인터페이스
with gr.Blocks(title="고급 문서 주제 분석기") as demo:
    gr.Markdown("## 📊 고급 문서 주제 분석기")
    gr.Markdown("문서를 입력하면 다양한 분석 결과를 시각화하여 보여줍니다.")
    
    with gr.Row():
        text = gr.Textbox(
            value=SAMPLE_TEXT,
            label="분석할 텍스트",
            placeholder="여기에 분석할 텍스트를 입력하세요",
            lines=8
        )
    
    with gr.Row():
        submit_btn = gr.Button("분석 시작", variant="primary")
    
    with gr.Tabs():
        with gr.TabItem("주요 분석"):
            with gr.Row():
                plot1 = gr.Plot(label="주제 분포")
                plot2 = gr.Plot(label="레이더 차트")
        
        with gr.TabItem("상세 분석"):
            with gr.Row():
                plot3 = gr.Plot(label="상관관계 히트맵")
                plot4 = gr.Plot(label="주제 변화 추이")
            
        with gr.TabItem("신뢰도 분석"):
            plot5 = gr.Plot(label="신뢰도 게이지")
            
        with gr.TabItem("감성 분석"):
            plot6 = gr.Plot(label="감성 분석 결과")
            
    with gr.Row():
        output = gr.JSON(label="상세 분석 결과")
    
    submit_btn.click(
        fn=process_all_analysis,
        inputs=[text],
        outputs=[output, plot1, plot2, plot3, plot4, plot5, plot6]
    )

if __name__ == "__main__":
    demo.queue(concurrency_count=1)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True
    )