import json
import gradio as gr
import spaces
import wbgtopic
import plotly.graph_objects as go
import plotly.express as px
import numpy as np
import pandas as pd
import nltk
from nltk.tokenize import sent_tokenize, word_tokenize
from nltk.sentiment import SentimentIntensityAnalyzer
from sklearn.cluster import KMeans
import torch

# GPU 사용 가능 시 설정
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

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

# 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.
"""

def safe_process(func):
    """
    예외 발생 시 로그를 남기고 None을 반환하는 데코레이터.
    Gradio 인터페이스가 예외로 인해 중단되지 않도록 도와줍니다.
    """
    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 parse_wbg_results(raw_output):
    """
    wbgtopic.WBGDocTopic의 suggest_topics() 결과를
    'label', 'score_mean', 'score_std' 구조의 리스트로 통일해서 반환한다.

    반환 구조 예시:
    [
      {
        "label": "Agriculture",
        "score_mean": 0.32,
        "score_std": 0.05
      },
      ...
    ]
    """
    # 디버그: 실제 결과 구조를 확인해보세요
    print(">>> DEBUG: raw_output =", raw_output)

    # 결과가 비었으면 빈 리스트 반환
    if not raw_output:
        return []

    first_item = raw_output[0]

    # (1) 이미 'label' 키가 있는 딕셔너리 형태라면
    # 예: [{"label": "...", "score": ...}, ...] 혹은 {"label": "...", "score_mean": ...}
    if isinstance(first_item, dict) and ("label" in first_item):
        parsed_list = []
        for item in raw_output:
            label = item.get("label", "")
            # score_mean / score_std가 이미 있으면 사용
            # 없으면 score 등에서 추론
            score_mean = item.get("score_mean", None)
            score_std = item.get("score_std", None)

            # 예: score만 있는 경우
            if score_mean is None and "score" in item:
                # 점수가 0~1 범위인지 0~100 범위인지 확인 필요
                # 우선 그대로 float 처리
                score_mean = float(item["score"])
            if score_mean is None:
                score_mean = 0.0

            if score_std is None:
                score_std = 0.0

            parsed_list.append({
                "label": label,
                "score_mean": float(score_mean),
                "score_std": float(score_std)
            })
        return parsed_list

    # (2) 토픽 이름: 점수 형태의 딕셔너리가 있는 경우
    # 예: [{"Agriculture": 0.22, "Climate Change": 0.55}, ...]
    if isinstance(first_item, dict):
        # raw_output가 여러 dict를 담고 있을 수 있으므로, 하나로 합치거나
        # 혹은 첫 번째 dict만 파싱할지 결정해야 함.
        # 일단 여기서는 합치는 방식으로 시연:
        merged = {}
        for d in raw_output:
            for k, v in d.items():
                # 키 중복 시 마지막 값으로 overwrite
                merged[k] = v

        parsed_list = []
        for label, val in merged.items():
            parsed_list.append({
                "label": label,
                "score_mean": float(val),
                "score_std": 0.0
            })
        return parsed_list

    # 예상치 못한 구조인 경우
    return []

@safe_process
def analyze_text_sections(text):
    """
    텍스트를 여러 섹션(예: 3문장씩)으로 나누고,
    각 섹션별로 suggest_topics() 결과를 parse_wbg_results()로 파싱해 리스트로 모은다.
    """
    sentences = sent_tokenize(text)
    # 3문장씩 묶어서 하나의 섹션을 구성
    sections = [' '.join(sentences[i:i+3]) for i in range(0, len(sentences), 3)]

    section_topics = []
    for section in sections:
        raw_sec = clf.suggest_topics(section)
        parsed_sec = parse_wbg_results(raw_sec)
        section_topics.append(parsed_sec)

    return section_topics

@safe_process
def calculate_topic_correlations(topic_dicts):
    """
    topic_dicts: [{'label': ..., 'score_mean': ..., 'score_std': ...}, ...]

    주제별 score_mean만 뽑아서 상관관계를 구한다.
    실제로는 '서로 다른 문서'들에 대한 상관을 구하는 것이 타당하나,
    여기서는 예시로 단일 텍스트의 서로 다른 토픽들 간 점수 상관도를 계산한다.
    """
    if len(topic_dicts) < 2:
        return np.array([[1.0]]), ["Insufficient topics"]

    labels = [d['label'] for d in topic_dicts]
    scores = [d['score_mean'] for d in topic_dicts]

    if len(scores) < 2:
        return np.array([[1.0]]), ["Insufficient topics"]

    corr_matrix = np.corrcoef(scores)
    return corr_matrix, labels

@safe_process
def perform_sentiment_analysis(text):
    """
    NLTK VADER를 사용해 문장별 감성 점수를 계산한다.
    반환값은 pandas DataFrame 형식.
    """
    sia = SentimentIntensityAnalyzer()
    sents = sent_tokenize(text)
    results = [sia.polarity_scores(s) for s in sents]
    return pd.DataFrame(results)

@safe_process
def create_topic_clusters(topic_dicts):
    """
    score_mean, score_std 2차원으로 KMeans 클러스터링.
    토픽 수가 3개 미만이면 trivially 0번 클러스터로 처리.
    """
    if len(topic_dicts) < 3:
        return [0] * len(topic_dicts)

    X = []
    for t in topic_dicts:
        X.append([t['score_mean'], t.get('score_std', 0.0)])

    X = np.array(X)
    if X.shape[0] < 3:
        return [0] * X.shape[0]

    kmeans = KMeans(n_clusters=min(3, X.shape[0]), random_state=42)
    clusters = kmeans.fit_predict(X)
    return clusters.tolist()

@safe_process
def create_main_charts(topic_dicts):
    """
    바 차트와 레이더 차트를 생성.
    'score_mean'을 0~1로 보고, 100배 하여 퍼센트로 시각화.
    """
    if not topic_dicts:
        return go.Figure(), go.Figure()

    labels = [t['label'] for t in topic_dicts]
    scores = [t['score_mean'] * 100 for t in topic_dicts]

    # 바 차트
    bar_fig = go.Figure(
        data=[go.Bar(x=labels, y=scores, marker_color='rgb(55, 83, 109)')]
    )
    bar_fig.update_layout(
        title='주제 분석 결과',
        xaxis_title='주제',
        yaxis_title='관련도(%)',
        template='plotly_white',
        height=500,
    )

    # 레이더 차트
    radar_fig = go.Figure()
    radar_fig.add_trace(go.Scatterpolar(
        r=scores,
        theta=labels,
        fill='toself',
        name='주제 분포'
    ))
    radar_fig.update_layout(
        title='주제 레이더 차트',
        template='plotly_white',
        height=500,
        polar=dict(radialaxis=dict(visible=True)),
        showlegend=False
    )
    return bar_fig, radar_fig

@safe_process
def create_correlation_heatmap(corr_matrix, labels):
    """
    상관관계 행렬을 히트맵으로 시각화.
    만약 데이터가 부족하면 안내 문구만 표시.
    """
    if corr_matrix.ndim == 0:
        # 스칼라(0차원)이면 2차원 배열로 바꿔줌
        corr_matrix = np.array([[corr_matrix]])

    if corr_matrix.shape == (1, 1):
        # 데이터가 부족한 경우
        fig = go.Figure()
        fig.add_annotation(text="Not enough topics for correlation", showarrow=False)
        return fig

    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):
    """
    섹션별 토픽 점수 변화를 라인 차트로 나타낸다.
    section_topics: [[{'label':..., 'score_mean':...}, ...], [...], ...]
    """
    fig = go.Figure()
    if not section_topics or len(section_topics) == 0:
        return fig

    if not section_topics[0]:
        return fig

    # 첫 섹션의 토픽들을 기준으로, 각 섹션마다 해당 토픽이 존재하면 점수를 추출
    for topic_dict in section_topics[0]:
        label = topic_dict['label']
        score_list = []
        for sec_list in section_topics:
            match = next((d for d in sec_list if d['label'] == label), None)
            if match:
                score_list.append(match['score_mean'])
            else:
                score_list.append(0.0)

        fig.add_trace(go.Scatter(
            x=list(range(len(section_topics))),
            y=score_list,
            name=label,
            mode='lines+markers'
        ))

    fig.update_layout(
        title='섹션별 주제 변화 추이',
        xaxis_title='섹션',
        yaxis_title='score_mean',
        height=500,
        template='plotly_white'
    )
    return fig

@safe_process
def create_confidence_gauge(topic_dicts):
    """
    각 토픽의 신뢰도를 게이지 형태로 표시.
    여기서는 (1 - score_std) * 100 단순 공식 사용.
    """
    if not topic_dicts:
        return go.Figure()

    fig = go.Figure()
    num_topics = len(topic_dicts)

    for i, t in enumerate(topic_dicts):
        conf_val = 100.0 * (1.0 - t.get("score_std", 0.0))
        fig.add_trace(go.Indicator(
            mode="gauge+number",
            value=conf_val,
            title={'text': t['label']},
            domain={'row': 0, 'column': i}
        ))

    fig.update_layout(
        grid={'rows': 1, 'columns': num_topics},
        height=400,
        template='plotly_white'
    )
    return fig

@spaces.GPU()
def process_all_analysis(text):
    """
    전체 텍스트에 대한 토픽 분석, 섹션 분석, 상관관계, 감성분석, 클러스터링 등을 수행한 뒤
    JSON 결과와 Plotly 차트들을 반환한다.
    """
    try:
        # 1) 전체 텍스트 대상 토픽 분석
        raw_results = clf.suggest_topics(text)
        all_topics = parse_wbg_results(raw_results)

        # 2) score_mean 기준 내림차순 정렬 후 상위 5개
        sorted_topics = sorted(all_topics, key=lambda x: x['score_mean'], reverse=True)
        top_topics = sorted_topics[:5]

        # 3) 섹션 별 분석
        section_topics = analyze_text_sections(text)

        # 4) 추가 분석(상관관계, 감성분석, 클러스터)
        corr_matrix, corr_labels = calculate_topic_correlations(all_topics)
        sentiments_df = perform_sentiment_analysis(text)
        clusters = create_topic_clusters(all_topics)

        # 5) 차트 생성
        bar_chart, radar_chart = create_main_charts(top_topics)
        heatmap = create_correlation_heatmap(corr_matrix, corr_labels)
        evolution_chart = create_topic_evolution(section_topics)
        gauge_chart = create_confidence_gauge(top_topics)

        # 6) JSON 형태로 묶어서 반환(문자열 키만 사용)
        results = {
            "top_topics": top_topics,                         # 상위 5개 토픽
            "clusters": clusters,                             # 클러스터 결과
            "sentiments": sentiments_df.to_dict(orient="records")  # 감성 분석
        }

        return (
            results,          # JSON output
            bar_chart,        # plot1
            radar_chart,      # plot2
            heatmap,          # plot3
            evolution_chart,  # plot4
            gauge_chart,      # plot5
            go.Figure()       # plot6 (필요 시 감성분석 그래프 사용)
        )

    except Exception as e:
        print(f"Analysis error: {str(e)}")
        empty_fig = go.Figure()
        return (
            {"error": str(e), "topics": []},
            empty_fig,
            empty_fig,
            empty_fig,
            empty_fig,
            empty_fig,
            empty_fig
        )

######################################################
#                Gradio UI Definition               #
######################################################

with gr.Blocks(title="고급 문서 주제 분석기") as demo:
    gr.Markdown("## 고급 문서 주제 분석기")
    gr.Markdown(
        "텍스트를 입력한 뒤, **분석 시작** 버튼을 눌러주세요. "
        "주요 토픽 분석, 상관관계, 신뢰도 게이지, 감성분석 결과 등을 확인할 수 있습니다."
    )

    with gr.Row():
        text_input = gr.Textbox(
            value=SAMPLE_TEXT,
            label="분석할 텍스트 입력",
            lines=8
        )
    with gr.Row():
        submit_btn = gr.Button("분석 시작", variant="primary")

    with gr.Tabs():
        with gr.TabItem("주요 분석"):
            with gr.Row():
                plot1 = gr.Plot(label="주제 분포(Bar Chart)")
                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_json = gr.JSON(label="상세 분석 결과(JSON)")

    submit_btn.click(
        fn=process_all_analysis,
        inputs=[text_input],
        outputs=[output_json, plot1, plot2, plot3, plot4, plot5, plot6]
    )

if __name__ == "__main__":
    demo.queue(max_size=1)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,  # 공개 링크 필요 시 True
        debug=True
    )