ui.py

import gradio as gr
import src.config as configs
from constants import TAB_NAMES, MODEL_TYPE_MAP, OUTPUT_FORM_MAP
from src.display.formatting import render_leaderboard_html
from src.display.css_html_js import get_leaderboard_table_html, custom_css
import pandas as pd
from constants import LEADERBOARD_REQUIRED_COLUMNS

def render_pretty_leaderboard_html(df):
    """
    Renders a pretty leaderboard table using badge and gauge.
    Supports both ['Model', 'Score'] and ['Model Name', 'Overall'] columns.
    Sorts by score descending and rounds for display.
    """
    # Flexible column mapping
    col_map = {}
    if "Model" in df.columns:
        col_map["Model"] = "Model"
    elif "Model Name" in df.columns:
        col_map["Model"] = "Model Name"
    else:
        return "<div style='color:red'>DataFrame must have a 'Model' or 'Model Name' column.</div>"
    if "Score" in df.columns:
        col_map["Score"] = "Score"
    elif "Overall" in df.columns:
        col_map["Score"] = "Overall"
    else:
        return "<div style='color:red'>DataFrame must have a 'Score' or 'Overall' column.</div>"

    # Example mappings for demonstration (expand as needed)
    model_type_map = MODEL_TYPE_MAP
    output_form_map = OUTPUT_FORM_MAP

    # Copy and rename for uniformity
    df2 = df.copy()
    df2 = df2.rename(columns={col_map["Model"]: "Model", col_map["Score"]: "Score"})

    # 매핑 전후로 누락된 모델명을 출력 (디버깅용)
    missing_type = set(df2["Model"]) - set(model_type_map.keys())
    missing_output = set(df2["Model"]) - set(output_form_map.keys())
    if missing_type:
        print("Model Type 매핑 누락:", missing_type)
    if missing_output:
        print("Output Form 매핑 누락:", missing_output)

    # Add badge columns
    df2["Model Type"] = df2["Model"].map(model_type_map).fillna("open")
    df2["Output Form"] = df2["Model"].map(output_form_map).fillna("normal")
    # Drop NA, sort, round
    df2 = df2[["Model", "Score", "Model Type", "Output Form"]].dropna()
    df2["Score"] = pd.to_numeric(df2["Score"], errors="coerce").round(2)
    df2 = df2.sort_values("Score", ascending=False).reset_index(drop=True)
    
    return get_leaderboard_table_html(df2)

def create_leaderboard_tab(df, key, search_leaderboard, update_modelselector_group, update_leaderboard, column_selector_value):
    """
    df: DataFrame to display
    key: "Category" or "Language"
    search_leaderboard, update_modelselector_group, update_leaderboard: handler functions
    column_selector_value: default columns to select
    """
    with gr.TabItem(
        TAB_NAMES[key],
        visible=True
    ):
        df_state = gr.State(df)

        with gr.Row():
            with gr.Column():
                search_box = gr.Textbox(label="Search Model by Name")
                group_list = df["Group"].unique().tolist()
                group_selector = gr.CheckboxGroup(
                    choices=df["Group"].unique().tolist(),
                    value=group_list,
                    label="Select Model Group"
                )
                # 필수 컬럼 항상 포함, 체크 해제 불가(disabled)
                # 선택지에서 "Model Name", "Group", "Overall" 제외
                exclude_cols = {"Model Name", "Group", "Overall"}
                selectable_columns = [col for col in df.columns.tolist()[3:] if col not in exclude_cols]
                all_columns = list(dict.fromkeys(LEADERBOARD_REQUIRED_COLUMNS + selectable_columns))
                column_selector = gr.CheckboxGroup(
                    choices=selectable_columns,
                    value=[col for col in column_selector_value if col in selectable_columns],
                    label="Select Columns"
                )

            with gr.Column():
                with gr.Accordion("Model List", open=False):
                    model_group = df["Model Name"].tolist()
                    model_selector = gr.CheckboxGroup(
                        choices=df["Model Name"].tolist(),
                        value=model_group,
                        label="Select Models"
                    )

        # badge 정보 포함 DataFrame 생성 (위쪽 테이블용)
        df_badge = df.copy()
        # Model 컬럼명 통일
        if "Model Name" in df_badge.columns:
            df_badge["Model"] = df_badge["Model Name"]
        # 예시 매핑 (아래쪽과 동일하게 확장)
        model_type_map = MODEL_TYPE_MAP
        output_form_map = OUTPUT_FORM_MAP
        df_badge["Model Type"] = df_badge["Model"].map(model_type_map).fillna("open")
        df_badge["Output Form"] = df_badge["Model"].map(output_form_map).fillna("normal")
        df_badge = df_badge.sort_values("Overall" if "Overall" in df_badge.columns else "Score", ascending=False).reset_index(drop=True)
        df_badge["Rank"] = df_badge.index + 1

        # 정렬 상태 관리용 State (한 번만 생성, 이후 재사용)
        default_sort_col = "Overall" if "Overall" in df_badge.columns else "Score"
        sort_col_state = gr.State(default_sort_col)
        sort_asc_state = gr.State(False)  # 내림차순이 기본값

        # 정렬 함수 (JS에서 넘긴 asc 값을 그대로 사용)
        def sort_and_render(col, asc, models, columns, df_):
            print(f"[sort_and_render] called: col={col}, asc={asc}, models={models}, columns={columns}")
            filtered_df = update_leaderboard(models, columns, df_, col, asc)
            # 정렬 상태를 DataFrame에 임시로 저장해 헤더에 반영
            filtered_df._sort_col = col
            filtered_df._sort_asc = asc
            return render_leaderboard_html(filtered_df.round(3)), col, asc

        leaderboard_html = render_leaderboard_html(df_badge.round(3))
        leaderboard_html_comp = gr.HTML(value=leaderboard_html, elem_id="leaderboard-table")

        # 정렬 트리거용 hidden textbox 추가
        sort_trigger = gr.Textbox(visible=False, elem_id="sort-leaderboard-trigger")

        # sort-arrow 클릭 시 항상 새로운 값으로 value를 변경하는 JS 삽입 (정렬 방향 포함)
        sort_js = """
        <script>
        (function() {
            document.addEventListener('DOMContentLoaded', function() {
                const table = document.getElementById('leaderboard-table');
                if (!table) return;
                table.addEventListener('click', function(e) {
                    const arrow = e.target.closest('.sort-arrow');
                    if (arrow) {
                        const col = arrow.getAttribute('data-col');
                        const asc = arrow.getAttribute('data-asc');
                        // 항상 새로운 값으로 value를 변경하여 change 이벤트 강제 발생
                        const trigger = document.querySelector('#sort-leaderboard-trigger input');
                        if (trigger) {
                            trigger.value = col + '|' + asc + '|' + Date.now();
                            trigger.dispatchEvent(new Event('input', { bubbles: true }));
                            trigger.dispatchEvent(new Event('change', { bubbles: true }));
                        }
                    }
                });
            });
        })();
        </script>
        """
        # 정렬 버튼 클릭 시에도 update_leaderboard를 호출하도록 wiring
        def sort_trigger_change(col_val, models, columns, df_, prev_col, prev_asc):
            print(f"[sort_trigger.change] col_val={col_val}, prev_col={prev_col}, prev_asc={prev_asc}")
            col, asc = col_val.split('|')[0], col_val.split('|')[1].lower() == "true"
            return sort_and_render(col, asc, models, columns, df_)

        sort_trigger.change(
            fn=sort_trigger_change,
            inputs=[sort_trigger, model_selector, column_selector, df_state, sort_col_state, sort_asc_state],
            outputs=[leaderboard_html_comp, sort_col_state, sort_asc_state]
        )

        # 커스텀 JS를 상단 테이블에 삽입
        leaderboard_html_comp.style = None  # gr.HTML에는 style 파라미터가 없으므로, 아래에서 삽입
        leaderboard_html_comp.value += sort_js

        # Pretty leaderboard preview (uses only 'Model' and 'Score' columns)
        pretty_html = gr.HTML(value=render_pretty_leaderboard_html(df.round(3)))

        # Define change functions for user interaction
        # 모든 UI 이벤트에서 update_leaderboard → sort_and_render → render_leaderboard_html 순으로 갱신
        def filter_and_sort_search(query, df, sort_col, sort_asc):
            print(f"[filter_and_sort_search] sort_col={sort_col}, sort_asc={sort_asc}")
            filtered_df = search_leaderboard(query, df, sort_col, sort_asc)
            # 정렬 상태를 DataFrame에 임시로 저장해 헤더에 반영
            filtered_df._sort_col = sort_col
            filtered_df._sort_asc = sort_asc
            return render_leaderboard_html(filtered_df), sort_col, sort_asc

        def filter_and_sort_model(models, columns, df, sort_col, sort_asc):
            print(f"[filter_and_sort_model] sort_col={sort_col}, sort_asc={sort_asc}")
            filtered_df = update_leaderboard(models, columns, df, sort_col, sort_asc)
            filtered_df._sort_col = sort_col
            filtered_df._sort_asc = sort_asc
            return render_leaderboard_html(filtered_df), sort_col, sort_asc

        def filter_and_sort_column(models, columns, df, sort_col, sort_asc):
            print(f"[filter_and_sort_column] sort_col={sort_col}, sort_asc={sort_asc}")
            filtered_df = update_leaderboard(models, columns, df, sort_col, sort_asc)
            filtered_df._sort_col = sort_col
            filtered_df._sort_asc = sort_asc
            return render_leaderboard_html(filtered_df), sort_col, sort_asc

        search_box.change(
            fn=filter_and_sort_search,
            inputs=[search_box, df_state, sort_col_state, sort_asc_state],
            outputs=[leaderboard_html_comp, sort_col_state, sort_asc_state]
        )

        group_selector.change(fn=update_modelselector_group, inputs=[group_selector, df_state], outputs=model_selector)
        model_selector.change(
            fn=filter_and_sort_model,
            inputs=[model_selector, column_selector, df_state, sort_col_state, sort_asc_state],
            outputs=[leaderboard_html_comp, sort_col_state, sort_asc_state]
        )

        # column_selector 변경 시에도 항상 최신 sort_col, sort_asc를 유지
        column_selector.change(
            fn=filter_and_sort_column,
            inputs=[model_selector, column_selector, df_state, sort_col_state, sort_asc_state],
            outputs=[leaderboard_html_comp, sort_col_state, sort_asc_state]
        )

        return {
            "search_box": search_box,
            "group_selector": group_selector,
            "column_selector": column_selector,
            "model_selector": model_selector,
            "leaderboard_html_comp": leaderboard_html_comp,
            "sort_trigger": sort_trigger,
            "df_state": df_state,
            "pretty_html": pretty_html
        }