app.py

import gradio as gr
import pandas as pd
import json
import math
import random

# ----------------------------
# Load JSON files
# ----------------------------
with open("sjt_trait_summary.json", "r") as f:
    sjt_trait_summary_dict = json.load(f)

with open("hexaco_trait_stat_summary.json", "r") as f:
    hexaco_trait_summary_dict = json.load(f)

with open("hexaco_subtrait_stat_summary.json", "r") as f:
    hexaco_subtrait_summary_dict = json.load(f)

with open("case_study_personas.json", "r") as f:
    persona_summary_dict = json.load(f)


# ----------------------------
# Function to display summaries
# ----------------------------
def show_summary(persona_id, custom_id):
    # Priority: custom textbox > dropdown
    persona_key = custom_id.strip() if custom_id.strip() else persona_id

    # ---- SJT Trait Summary ----
    if persona_key in sjt_trait_summary_dict:
        sjt_traits = sjt_trait_summary_dict[persona_key]
        sjt_df = pd.DataFrame([
            {
                "Trait": t,
                "Count": v["count"],
                "Proportion (%)": round(v["proportion"] * 100, 1)
            }
            for t, v in sjt_traits.items()
        ])
        sjt_header = f"### SJT Trait Summary for persona: {persona_key}"
    else:
        sjt_df, sjt_header = None, f"No SJT trait data found for persona_id: {persona_key}"

    # ---- HEXACO Trait Summary ----
    if persona_key in hexaco_trait_summary_dict:
        traits = hexaco_trait_summary_dict[persona_key]
        hexaco_df = pd.DataFrame([
            {
                "Trait": t.capitalize(),
                "Refusal Rate (%)": round(v["refusal_rate"] * 100, 1),
                "Mean Score": round(v["trait_score_mean"], 2) if not math.isnan(v["trait_score_mean"]) else None,
                "Std Dev": round(v["trait_scor_std"], 2) if v["trait_scor_std"] is not None and not math.isnan(v["trait_scor_std"]) else None,
            }
            for t, v in traits.items()
        ])
        hexaco_header = f"### HEXACO Trait Summary for persona: {persona_key}"
    else:
        hexaco_df, hexaco_header = None, f"No HEXACO trait data found for persona_id: {persona_key}"

    # ---- HEXACO Subtrait Summary ----
    if persona_key in hexaco_subtrait_summary_dict:
        subtraits = []
        for trait, subs in hexaco_subtrait_summary_dict[persona_key].items():
            for sub_name, v in subs.items():
                subtraits.append({
                    "Trait": trait.capitalize(),
                    "Subtrait": sub_name.capitalize(),
                    "Refusal Rate (%)": round(v["refusal_rate"] * 100, 1),
                    "Mean Score": round(v["trait_score_mean"], 2) if not math.isnan(v["trait_score_mean"]) else None,
                    "Std Dev": round(v["trait_scor_std"], 2) if v["trait_scor_std"] is not None and not math.isnan(v["trait_scor_std"]) else None,
                })
        subtrait_df = pd.DataFrame(subtraits)
        subtrait_header = f"### HEXACO Subtrait Summary for persona: {persona_key}"
    else:
        subtrait_df, subtrait_header = None, f"No HEXACO subtrait data found for persona_id: {persona_key}"

    # ---- Persona Summary ----
    if persona_key in persona_summary_dict:
        persona_info = persona_summary_dict[persona_key]
        persona_str = persona_info.get("persona_string", "").replace("\n", "<br>")
        archetype = persona_info.get("archetype", "N/A")
        persona_md = f"### Persona Summary\n**Archetype:** {archetype}\n\n{persona_str}"
    else:
        persona_md = f"No persona summary found for persona_id: {persona_key}"

    return sjt_header, sjt_df, hexaco_header, hexaco_df, subtrait_header, subtrait_df, persona_md


# ----------------------------
# Pick a random persona for initialization
# ----------------------------
initial_persona = random.choice(list(persona_summary_dict.keys()))
initial_outputs = show_summary(initial_persona, "")


# ----------------------------
# Gradio UI
# ----------------------------
with gr.Blocks() as demo:
    gr.Markdown("# 🧩 Persona Summary, SJT Traits & HEXACO Traits Viewer")
    gr.Markdown("Select a persona from the dropdown or enter a **custom persona_id** manually.")

    with gr.Row():
        persona_id_in = gr.Dropdown(
            choices=list(persona_summary_dict.keys()),
            value=initial_persona,
            label="Select Persona ID",
            interactive=True
        )
        custom_id_in = gr.Textbox(label="Or enter Persona ID manually")

    # Persona Summary
    persona_text = gr.Markdown(value=initial_outputs[6])
    
    # SJT Trait Summary
    sjt_text = gr.Markdown(value=initial_outputs[0])
    sjt_table = gr.Dataframe(value=initial_outputs[1], interactive=False)

    # HEXACO Trait Summary
    hexaco_text = gr.Markdown(value=initial_outputs[2])
    hexaco_table = gr.Dataframe(value=initial_outputs[3], interactive=False)

    # HEXACO Subtrait Summary
    subtrait_text = gr.Markdown(value=initial_outputs[4])
    subtrait_table = gr.Dataframe(value=initial_outputs[5], interactive=False)

    

    # Trigger automatically on change
    persona_id_in.change(
        show_summary,
        inputs=[persona_id_in, custom_id_in],
        outputs=[sjt_text, sjt_table, hexaco_text, hexaco_table, subtrait_text, subtrait_table, persona_text]
    )

    custom_id_in.change(
        show_summary,
        inputs=[persona_id_in, custom_id_in],
        outputs=[sjt_text, sjt_table, hexaco_text, hexaco_table, subtrait_text, subtrait_table, persona_text]
    )

# ----------------------------
# Launch
# ----------------------------
if __name__ == "__main__":
    demo.launch()