app.py

import gradio as gr
import yaml
import json
import os
import traceback
import matplotlib.pyplot as plt
import numpy as np
import tempfile
import speech_recognition as sr
from TTS.api import TTS

tts_model = TTS("tts_models/en/ljspeech/tacotron2-DDC")

#persona_voices = {
#    "Maya": "p225",     # younger female
#    "Robert": "p240",   # older male
#    "Angela": "p231",   # middle-aged female
#    "Jack": "p260"      # younger male
#}

from engine.loader import load_persona
from engine.drift import apply_context_shift
from engine.responder import generate_response, generate_response_local
from engine.utils import safe_log
from engine.logger import log_interaction

# Paths
persona_dir = "./personas"
contexts_path = "./contexts/scenarios.json"
error_log_path = "./ot_simulator_errors.log"

# Load available personas
def get_persona_choices():
    return [f for f in os.listdir(persona_dir) if f.endswith(".yml")]

# Load available contextual scenarios
def get_scenario_choices():
    try:
        with open(contexts_path, "r") as f:
            scenarios = json.load(f)
        return [s["scenario"] for s in scenarios]
    except Exception as e:
        safe_log("Scenarios load error", str(e))
        return []

# Generate radar chart for emotional/behavioral states
def plot_state(state, persona_name):
    if persona_name == "Jack":
        metrics = ["anxiety", "trust", "openness", "physical_discomfort"]
        colors = ["#e74c3c", "#3498db", "#2ecc71", "#f39c12"]
    elif persona_name == "Maya":
        metrics = ["anxiety", "trust", "creative_engagement", "occupational_balance"]
        colors = ["#e74c3c", "#3498db", "#9b59b6", "#1abc9c"]
    else:
        metrics = ["anxiety", "trust", "openness", "engagement"]
        colors = ["#e74c3c", "#3498db", "#2ecc71", "#95a5a6"]
    
    values = [state.get(m, 0.0) for m in metrics]
    angles = np.linspace(0, 2 * np.pi, len(metrics), endpoint=False).tolist()
    values += values[:1]
    angles += angles[:1]

    fig, ax = plt.subplots(figsize=(5, 5), subplot_kw=dict(polar=True))
    ax.plot(angles, values, color=colors[0], linewidth=2)
    ax.fill(angles, values, color=colors[0], alpha=0.25)
    ax.set_xticks(angles[:-1])
    ax.set_xticklabels([m.replace('_', ' ').title() for m in metrics])
    ax.set_ylim(0, 1)
    ax.set_yticklabels(['0.0', '0.2', '0.4', '0.6', '0.8', '1.0'])
    ax.set_title(f"{persona_name}'s Emotional State", fontsize=14, pad=20)
    ax.grid(True)
    fig.tight_layout()

    chart_path = f"./state_chart_{persona_name}.png"
    fig.savefig(chart_path, dpi=100, bbox_inches='tight')
    plt.close(fig)
    return chart_path

# Generate interaction history visualization
def plot_interaction_history(history):
    if not history or len(history) < 2:
        return None
    
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 6))
    
    interactions = list(range(1, len(history) + 1))
    anxiety_vals = [h.get('anxiety', 0) for h in history]
    trust_vals = [h.get('trust', 0) for h in history]
    
    ax1.plot(interactions, anxiety_vals, marker='o', color='#e74c3c', linewidth=2, label='Anxiety')
    ax1.set_ylabel('Anxiety Level', fontsize=10)
    ax1.set_ylim(0, 1)
    ax1.grid(True, alpha=0.3)
    ax1.legend(loc='upper right')
    
    ax2.plot(interactions, trust_vals, marker='o', color='#3498db', linewidth=2, label='Trust')
    ax2.set_xlabel('Interaction Number', fontsize=10)
    ax2.set_ylabel('Trust Level', fontsize=10)
    ax2.set_ylim(0, 1)
    ax2.grid(True, alpha=0.3)
    ax2.legend(loc='upper right')
    
    fig.suptitle('Therapeutic Relationship Over Time', fontsize=14)
    fig.tight_layout()
    
    history_path = "./interaction_history.png"
    fig.savefig(history_path, dpi=100, bbox_inches='tight')
    plt.close(fig)
    return history_path

# Download session transcript
def download_session(conversation_history, state_history, selected_persona_file):
    """Generate downloadable transcript file."""
    if not conversation_history:
        return None

    try:
        persona_path = os.path.join(persona_dir, selected_persona_file)
        persona = load_persona(persona_path)

        from datetime import datetime
        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        name = persona.get("persona_name", "Character")

        # Header
        transcript = f"""================================================================
LIT DIGITAL TWIN - INTERACTION TRANSCRIPT
================================================================

Character: {name}
Date: {timestamp}
Number of Turns: {len(conversation_history)}

================================================================
CONVERSATION
================================================================
"""

        # Conversation history
        for i, turn in enumerate(conversation_history, 1):
            transcript += f"\n[Turn {i}]\n"
            if 'scenario' in turn:
                transcript += f"Scenario: {turn['scenario']}\n\n"
            transcript += f"Interviewer: {turn.get('student', '')}\n\n"
            transcript += f"{name}: {turn.get('client', '')}\n\n"
            transcript += "-" * 63 + "\n"

        # State progression
        if state_history:
            transcript += f"""================================================================
STATE PROGRESSION
================================================================

Initial State:
  Anxiety: {state_history[0].get('anxiety', 0):.2f}
  Trust: {state_history[0].get('trust', 0):.2f}
  Openness: {state_history[0].get('openness', 0):.2f}

Final State:
  Anxiety: {state_history[-1].get('anxiety', 0):.2f}
  Trust: {state_history[-1].get('trust', 0):.2f}
  Openness: {state_history[-1].get('openness', 0):.2f}

Change:
  Anxiety: {state_history[-1].get('anxiety', 0) - state_history[0].get('anxiety', 0):+.2f}
  Trust: {state_history[-1].get('trust', 0) - state_history[0].get('trust', 0):+.2f}
  Openness: {state_history[-1].get('openness', 0) - state_history[0].get('openness', 0):+.2f}

================================================================
"""

        # Save to file
        filename = f"{name}_{timestamp}.txt"
        filepath = os.path.join("lit_transcripts", filename)
        os.makedirs("lit_transcripts", exist_ok=True)

        with open(filepath, "w", encoding="utf-8") as f:
            f.write(transcript)

        return filepath

    except Exception as e:
        safe_log("Download error", str(e))
        return None

def simulate(prompt, selected_event, selected_persona_file, ai_mode,
             conversation_history, state_history, use_fast_mode):
    try:
        # Normalize inputs
        if hasattr(prompt, 'value'):
            prompt = prompt.value
        prompt = str(prompt) if prompt else ""

        if hasattr(selected_event, 'value'):
            selected_event = selected_event.value

        if hasattr(selected_persona_file, 'value'):
            selected_persona_file = selected_persona_file.value

        if hasattr(ai_mode, 'value'):
            ai_mode = ai_mode.value

        # Initialize histories
        conversation_history = conversation_history or []
        state_history = state_history or []

        # Load persona
        persona_path = os.path.join(persona_dir, selected_persona_file)
        persona = load_persona(persona_path)

        # Load contextual scenario
        with open(contexts_path, "r") as f:
            scenarios = json.load(f)
        scenario = next((s for s in scenarios if s["scenario"] == selected_event), None)

        if scenario:
            persona = apply_context_shift(persona, scenario)
            context_note = f"**Context:** {scenario.get('description', selected_event)}\n\n"
        else:
            context_note = ""

        # Generate response
        response, updated_state, teaching_note = generate_response(
            prompt,
            persona,
            conversation_history,
            force_mode=ai_mode,
            use_fast_mode=use_fast_mode
        )


        # Update histories
        conversation_history.append({
            "student": prompt,
            "client": response,
            "scenario": selected_event
        })
        state_history.append(updated_state.copy())

        # Conversation display (now uses Interlocutor)
        conversation_display = ""
        for i, turn in enumerate(conversation_history, 1):
            conversation_display += f"**Turn {i}**\n"
            if 'scenario' in turn:
                conversation_display += f"*[Scenario: {turn['scenario']}]*\n"
            conversation_display += f"**Interviewer:** {turn['student']}\n\n"
            conversation_display += f"**{persona['persona_name']}:** {turn['client']}\n\n"
            conversation_display += "---\n\n"

        # Visualizations
        state_yaml = yaml.dump(updated_state, sort_keys=False)
        current_chart = plot_state(updated_state, persona['persona_name'])
        history_chart = plot_interaction_history(state_history)

        # Teaching feedback
        teaching_feedback = f"**Note:**\n{teaching_note}\n\n"
        teaching_feedback += f"**Current State:**\n"
        for k, v in updated_state.items():
            if isinstance(v, (int, float)):
                teaching_feedback += f"- {k.capitalize()}: {v:.2f}\n"

        if 'emotional_memory' in updated_state and updated_state['emotional_memory']:
            teaching_feedback += f"\n**Memory:** {updated_state['emotional_memory'][-1]}\n"

        # Log interaction
        transcript_path = log_interaction(
            persona,
            prompt,
            selected_event,
            response,
            updated_state,
            teaching_note,
            
        )

        return (
            conversation_display,
            teaching_feedback,
            state_yaml,
            current_chart,
            history_chart,
            conversation_history,
            state_history
        )

    except Exception as e:
        error_msg = traceback.format_exc()
        safe_log("Simulation error", error_msg)
        print(f"ERROR: {error_msg}")
        return (
            "[ERROR] Simulation failed. Check logs.",
            "Error occurred",
            "",
            None,
            None,
            conversation_history,
            state_history
        )

        # Save to file
        filename = f"{name}_{timestamp}.txt"
        filepath = os.path.join("lit_transcripts", filename)
        os.makedirs("lit_transcripts", exist_ok=True)

        with open(filepath, "w", encoding="utf-8") as f:
            f.write(transcript)

        return filepath

    except Exception as e:
        safe_log("Download error", str(e))
        return None
# --- Speech-to-Text helper ---
def speech_to_text(audio_file):
    recognizer = sr.Recognizer()
    with sr.AudioFile(audio_file) as source:
        audio = recognizer.record(source)
    try:
        text = recognizer.recognize_google(audio)
    except sr.UnknownValueError:
        text = ""
    return text

# --- Wrapper with STT + TTS ---
def simulate_with_voice(audio_in, scenario, persona, mode, conversation_state, state_history):
    student_prompt = speech_to_text(audio_in) if audio_in else ""

    conversation_display, teaching_feedback, state_yaml, current_chart, history_chart, conversation_state, state_history = simulate(
        student_prompt, scenario, persona, mode, conversation_state, state_history
    )

    persona_response = conversation_display
    tmpfile = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    tts_model.tts_to_file(text=persona_response, file_path=tmpfile.name)

    return (
        conversation_display,
        teaching_feedback,
        state_yaml,
        current_chart,
        history_chart,
        conversation_state,
        state_history,
        tmpfile.name
    )


import gradio as gr
import os

with gr.Blocks(
    title="Literary Character Simulator",
    theme=gr.themes.Soft()
) as ui:

    # Header
    with gr.Row():
        with gr.Column(scale=1):
            gr.Image("NYIT.png", show_label=False, height=60, width=180)
        with gr.Column(scale=8):
            gr.Markdown("## Literary Character Simulator")

    # Intro text
    gr.Markdown(
        """
### Learn from characters in fiction

Ask questions pertaining to the story.
"""
    )

    # State management
    conversation_state = gr.State([])
    state_history = gr.State([])

    # Main row: left column (selectors + info + mode) and right column (conversation)
    with gr.Row():
        with gr.Column(scale=1):
            persona_files = get_persona_choices()
            default_persona = persona_files[0] if persona_files else None

            persona_selector = gr.Dropdown(
                label="Select Character",
                choices=persona_files,
                value=default_persona,
                allow_custom_value=False,
                info="Choose which character you'll be working with"
            )

            scenario_choices = get_scenario_choices()
            scenario_selector = gr.Dropdown(
                label="Contextual Scenario",
                choices=scenario_choices,
                value=scenario_choices[0] if scenario_choices else None,
                info="What's happening in the character's life right now?"
            )

            gr.Markdown("---")

            with gr.Accordion("Character Information", open=False):
                gr.Markdown(
"""**Jimmy Markum (38)** – Corner store owner and ex-convict driven by grief and vengeance after his daughter's murder. Uses street-level intuition and intimidation to investigate outside the law.

**Sean Devine (38)** – Homicide detective torn between duty and personal history. Methodical and emotionally guarded, investigating the same case that haunts his childhood friendships.

**Hamlet (28)** – Prince of Denmark consumed by grief, betrayal, and existential doubt. Philosophical and emotionally volatile, struggling to avenge his father’s death.

**Jane (30)** – Narrator of *The Yellow Wallpaper*, confined and unraveling under the rest cure. Perceptive and imaginative, she descends into madness while observing the wallpaper’s patterns.

**Harold Krebs (25)** – War veteran from *Soldier’s Home*, emotionally numb and disconnected from postwar life. Struggles with reintegration, honesty, and societal expectations.

**Louise Mallard (28)** – Grieving wife from *The Story of an Hour*, who experiences a fleeting sense of freedom upon hearing of her husband’s death. Emotionally complex and physically fragile.

**Eveline (19)** – Young woman from Joyce’s *Eveline*, torn between duty to her family and the promise of escape. Paralyzed by fear and memory at the moment of decision.

**Uncle Ben (50)** – Civil servant from Achebe’s *Uncle Ben’s Choice*, navigating moral compromise in postcolonial Nigeria. Pragmatic and ethically conflicted, choosing silence over confrontation.

**John Keegan (47)** – Deputy Inspector in the NYPD and subject of a hit TV series. Known as “Keegan the Destroyer,” he’s a controversial figure haunted by past cases, a fractured marriage, and the death of his partner. Struggles with emotional suppression and guilt.

**Creon (50)** – Advisor and brother-in-law to Oedipus. Rational and loyal, he is wrongly accused of conspiracy but remains composed. Becomes ruler after Oedipus’s fall, representing justice and restraint.

**Tiresias (70)** – Blind prophet of Thebes. Cryptic and wise, he reveals truths others fear to hear. Burdened by divine insight, he speaks with solemn clarity and moral weight.

**Jocasta (56)** – Queen of Thebes, mother and wife to Oedipus. Emotionally intuitive and protective, she tries to suppress the truth that will destroy her family, ultimately taking her own life.

**Laertes (25)** – Nobleman of Denmark, son of Polonius. Passionate and impulsive, he seeks revenge for his family’s destruction, ultimately dying in a duel with Hamlet.

**Arianna Nunez (32)** – Ambitious NYPD Detective and digital expert. Balances Keegan and Lavin’s chaos with wit and precision. Known for undercover work, sharp comebacks, and a drive to rise in the ranks.

**Karl Lavin (58)** – NYPD Lieutenant and Keegan’s loyal, sarcastic partner. A veteran of Vice, he uses crude humor and streetwise tactics to protect his team. Planning retirement, but still Keegan’s “babysitter.”

**Ellie Williams (19)** – Immune survivor in a post-apocalyptic world. Fierce, sarcastic, and emotionally scarred, she carries the weight of loss and the burden of being a symbol. Loyal but guarded, she fights to define her own path.

**Joel Miller (52)** – Survivor of a global outbreak and father figure to Ellie. Sorrowful but guarded, his emotional wounds have scarred over. Loyal and pragmatic, he makes morally complex choices to protect those he loves.

**Oedipus (40)** – King of Thebes and tragic seeker of truth. Determined to end a plague, he uncovers his own guilt and fulfills a devastating prophecy, blinding himself in remorse.

**Pauline McCrory-Keegan (45)** – NYPD officer on leave, seeking identity and purpose after years of emotional strain and marital separation. Investigating a missing child case in Maryland while confronting her past and redefining her future.
"""
    )

            ai_mode_selector = gr.Radio(
                label="Response Mode",
                choices=["AI (HuggingFace)", "Templates (Local)"],
                value="AI (HuggingFace)",
                info="Choose how responses are generated"
            )

            use_fast_mode_toggle = gr.Checkbox(
                label="Use Fast Mode (Local AI)",
                value=False,
                info="Enable faster responses using local model"
            )

        with gr.Column(scale=2):
            conversation_display = gr.Markdown(
                label="Conversation History",
                value="*Conversation will appear here...*"
            )

    #with gr.Row():
     #   audio_in = gr.Audio(type="filepath", label="🎤 Speak to the character", interactive=True)
      #  voice_submit_btn = gr.Button("Send Voice Input")
       # audio_out = gr.Audio(label="Character Voice", type="filepath")

    # User input row
    with gr.Row():
        with gr.Column():
            student_prompt = gr.Textbox(
                label="Your Input",
                lines=3,
                placeholder="Ask the character something...",
                value=""
            )

    # Buttons row
    with gr.Row():
        with gr.Column(scale=1):
            send_btn = gr.Button("Send", variant="primary")
        with gr.Column(scale=1):
            download_btn = gr.Button("Download Session")
        with gr.Column(scale=1):
            reset_btn = gr.Button("Reset Conversation")

    # Feedback + technical state row
    with gr.Row():
        with gr.Column():
            teaching_output = gr.Markdown(label="Narrative Feedback")
        with gr.Column():
            state_output = gr.Textbox(
                label="Technical State (for debugging)",
                lines=8,
                visible=False
            )

    # Charts row
    with gr.Row():
        with gr.Column():
            current_state_chart = gr.Image(label="Current State")
        with gr.Column():
            history_chart = gr.Image(label="Progress Over Time")

    # Hidden file output for downloads
    download_file = gr.File(
        label="Your session transcript (right-click to save)",
        visible=True
    )

    # Instructor guide
    with gr.Accordion("Instructor Guide", open=False):
        gr.Markdown(
"""This simulator allows you to interact with fictional characters
in different scenarios. Use it to explore dialogue, context, and
character development in a safe environment."""
        )

    # Button actions
    send_btn.click(
        fn=simulate,
        inputs=[
            student_prompt,
            scenario_selector,
            persona_selector,
            ai_mode_selector,
            conversation_state,
            state_history,
            use_fast_mode_toggle
        ],
        outputs=[
            conversation_display,
            teaching_output,
            state_output,
            current_state_chart,
            history_chart,
            conversation_state,
            state_history
        ]
    )



    download_btn.click(
        fn=download_session,
        inputs=[
            conversation_state,
            state_history,
            persona_selector
        ],
        outputs=download_file
    )

  #  voice_submit_btn.click(
   #     fn=simulate_with_voice,
    #    inputs=[
     #     persona_selector,
  #          ai_mode_selector,
   #         conversation_state,
  #          state_history
   #     ],
  #      outputs=[
  #          conversation_display,
  #          teaching_output,
  #          state_output,
  #          current_state_chart,
  #          history_chart,
  #          conversation_state,
  #          state_history,
  #          audio_out
  #      ]
  #  )

    def reset_conversation():
        return (
            "*Conversation will appear here...*",
            "",
            "",
            None,
            None,
            [],
            []
        )

    reset_btn.click(
        fn=reset_conversation,
        inputs=[],
        outputs=[
            conversation_display,
            teaching_output,
            state_output,
            current_state_chart,
            history_chart,
            conversation_state,
            state_history
        ]
    )

if __name__ == "__main__":
    os.makedirs("personas", exist_ok=True)
    os.makedirs("contexts", exist_ok=True)
    os.makedirs("transcripts", exist_ok=True)
    os.makedirs("engine", exist_ok=True)

    ui.launch(share=True, server_name="0.0.0.0", server_port=7860)