app.py

"""
Main Gradio application for the AI Trading Experiment.
Combines trading game, chatbot, and experiment tracking.
"""

import gradio as gr
import uuid
import time
import random
from datetime import datetime
from typing import Optional, Tuple, List, Dict, Any

from config import (
    ExperimentConfig, DEFAULT_CONFIG, SCENARIOS,
    ParticipantVisibleParams, ResearcherControlledParams,
    EXPERIMENT_CONDITIONS, get_condition
)
from trading import TradingEngine, format_currency, format_percentage
from chatbot import TradingChatbot, ChatResponse
from tracking import (
    ExperimentTracker, DecisionRecord, ChatInteraction,
    tracker
)


# Global state management
class SessionState:
    """Manages session state for each participant."""
    def __init__(self):
        self.reset()

    def reset(self):
        self.participant_id: Optional[str] = None
        self.trading_engine: Optional[TradingEngine] = None
        self.chatbot: Optional[TradingChatbot] = None
        self.visible_params = ParticipantVisibleParams()
        self.hidden_params = ResearcherControlledParams()
        self.current_scenario = None
        self.scenario_start_time: float = 0
        self.ai_advice_shown_time: float = 0
        self.proactive_advice_shown: bool = False
        self.proactive_advice_engaged: bool = False
        self.chat_queries_this_scenario: int = 0
        self.pre_advice_confidence: int = 50


# Initialize components
session = SessionState()


def get_random_condition() -> ResearcherControlledParams:
    """Get a random experimental condition for A/B testing."""
    condition_names = list(EXPERIMENT_CONDITIONS.keys())
    selected = random.choice(condition_names)
    return get_condition(selected)


def start_experiment(condition_dropdown: str) -> Tuple:
    """Initialize a new experiment session."""
    session.reset()

    # Set experimental condition
    if condition_dropdown == "Random (A/B Testing)":
        session.hidden_params = get_random_condition()
    else:
        session.hidden_params = get_condition(condition_dropdown)

    # Create participant session
    session.participant_id = tracker.create_session(
        condition_name=session.hidden_params.condition_name,
        initial_portfolio=DEFAULT_CONFIG.initial_portfolio_value
    )

    # Initialize trading engine
    session.trading_engine = TradingEngine(DEFAULT_CONFIG)
    portfolio, first_scenario = session.trading_engine.start_new_game()

    # Initialize chatbot
    session.chatbot = TradingChatbot()
    session.chatbot.clear_history()

    # Set current scenario
    session.current_scenario = first_scenario
    session.scenario_start_time = time.time()

    # Generate proactive advice (maybe)
    proactive_message = ""
    session.proactive_advice_shown = False

    proactive_response = session.chatbot.generate_proactive_advice(
        first_scenario,
        session.visible_params,
        session.hidden_params
    )

    if proactive_response:
        proactive_message = f"**AI Advisor:** {proactive_response.message}"
        session.proactive_advice_shown = True
        session.ai_advice_shown_time = time.time()

        # Record proactive interaction
        tracker.record_chat_interaction(ChatInteraction(
            interaction_id=str(uuid.uuid4())[:12],
            participant_id=session.participant_id,
            timestamp=datetime.now().isoformat(),
            scenario_id=first_scenario.scenario_id,
            interaction_type="proactive",
            user_query=None,
            ai_response=proactive_response.message,
            explanation_depth=session.visible_params.explanation_depth,
            communication_style=session.visible_params.communication_style,
            confidence_framing=session.hidden_params.confidence_framing,
            risk_bias=session.hidden_params.risk_bias,
            response_time_ms=0,
            user_engaged=False,
            dismissed=False
        ))

    # Generate main AI recommendation
    ai_recommendation = session.chatbot.generate_ai_recommendation(
        first_scenario,
        session.visible_params,
        session.hidden_params
    )

    # Format scenario display
    scenario_text = format_scenario_display(first_scenario)
    progress = session.trading_engine.get_progress_info()

    return (
        gr.update(visible=False),  # Hide welcome
        gr.update(visible=True),   # Show game
        gr.update(visible=False),  # Hide results
        f"Participant: {session.participant_id}",
        format_currency(portfolio.total_value),
        f"Scenario {progress['current_scenario']} of {progress['total_scenarios']}",
        scenario_text,
        f"**AI Recommendation:**\n\n{ai_recommendation.message}",
        proactive_message,
        [],  # Clear chat history
        gr.update(value=50),  # Reset confidence slider
        gr.update(value="HOLD"),  # Reset decision
    )


def format_scenario_display(scenario) -> str:
    """Format a scenario for display."""
    return f"""
## {scenario.company_name} ({scenario.company_symbol})

**Sector:** {scenario.sector}
**Country:** {scenario.country}
**Current Price:** {scenario.current_price} credits

---

### Situation

{scenario.situation_description}
"""


def update_ai_params(explanation_depth: int, communication_style: int):
    """Update visible AI parameters and regenerate advice."""
    session.visible_params.explanation_depth = explanation_depth
    session.visible_params.communication_style = communication_style

    if session.current_scenario and session.chatbot:
        # Regenerate recommendation with new params
        ai_recommendation = session.chatbot.generate_ai_recommendation(
            session.current_scenario,
            session.visible_params,
            session.hidden_params
        )
        return f"**AI Recommendation:**\n\n{ai_recommendation.message}"

    return ""


def handle_chat_query(
    message: str,
    chat_history: List[Tuple[str, str]]
) -> Tuple[List[Tuple[str, str]], str]:
    """Handle a user query to the chatbot."""
    if not message.strip() or not session.chatbot:
        return chat_history, ""

    query_start = time.time()

    response = session.chatbot.answer_query(
        message,
        session.current_scenario,
        session.visible_params,
        session.hidden_params
    )

    response_time_ms = int((time.time() - query_start) * 1000)

    # Record interaction
    if session.participant_id:
        tracker.record_chat_interaction(ChatInteraction(
            interaction_id=str(uuid.uuid4())[:12],
            participant_id=session.participant_id,
            timestamp=datetime.now().isoformat(),
            scenario_id=session.current_scenario.scenario_id if session.current_scenario else None,
            interaction_type="reactive_query",
            user_query=message,
            ai_response=response.message,
            explanation_depth=session.visible_params.explanation_depth,
            communication_style=session.visible_params.communication_style,
            confidence_framing=session.hidden_params.confidence_framing,
            risk_bias=session.hidden_params.risk_bias,
            response_time_ms=response_time_ms,
            user_engaged=True,
            dismissed=False
        ))

    session.chat_queries_this_scenario += 1

    chat_history.append((message, response.message))
    return chat_history, ""


def engage_proactive_advice():
    """Mark that user engaged with proactive advice."""
    session.proactive_advice_engaged = True
    return "You engaged with the AI's initial observation."


def dismiss_proactive_advice():
    """Mark that user dismissed proactive advice."""
    session.proactive_advice_engaged = False
    return gr.update(visible=False)


def submit_decision(
    decision: str,
    confidence: int,
    trade_amount: float
) -> Tuple:
    """Process the participant's trading decision."""
    if not session.trading_engine or not session.current_scenario:
        return (gr.update(),) * 8

    # Calculate timing
    decision_time_ms = int((time.time() - session.scenario_start_time) * 1000)
    ai_viewing_time_ms = int((time.time() - session.ai_advice_shown_time) * 1000) if session.ai_advice_shown_time > 0 else 0

    # Process the decision
    outcome = session.trading_engine.process_decision(
        session.current_scenario,
        decision,
        trade_amount,
        session.current_scenario.ai_recommendation
    )

    # Record decision
    tracker.record_decision(DecisionRecord(
        decision_id=str(uuid.uuid4())[:12],
        participant_id=session.participant_id,
        timestamp=datetime.now().isoformat(),
        scenario_id=session.current_scenario.scenario_id,
        company_symbol=session.current_scenario.company_symbol,
        explanation_depth=session.visible_params.explanation_depth,
        communication_style=session.visible_params.communication_style,
        confidence_framing=session.hidden_params.confidence_framing,
        risk_bias=session.hidden_params.risk_bias,
        ai_recommendation=session.current_scenario.ai_recommendation,
        ai_was_correct=session.current_scenario.ai_is_correct,
        participant_decision=decision,
        followed_ai=outcome.followed_ai,
        decision_confidence=confidence,
        time_to_decision_ms=decision_time_ms,
        time_viewing_ai_advice_ms=ai_viewing_time_ms,
        outcome_percentage=outcome.outcome_percentage,
        portfolio_before=outcome.portfolio_before,
        portfolio_after=outcome.portfolio_after,
        trade_amount=trade_amount,
        proactive_advice_shown=session.proactive_advice_shown,
        proactive_advice_engaged=session.proactive_advice_engaged
    ))

    # Record trust metrics
    tracker.record_trust_metric(
        participant_id=session.participant_id,
        scenario_id=session.current_scenario.scenario_id,
        pre_confidence=session.pre_advice_confidence,
        post_confidence=confidence,
        advice_followed=outcome.followed_ai,
        time_deliberating_ms=decision_time_ms,
        queries_before_decision=session.chat_queries_this_scenario,
        outcome_positive=(outcome.outcome_percentage > 0)
    )

    # Format outcome message
    outcome_color = "green" if outcome.outcome_percentage > 0 else "red"
    outcome_message = f"""
### Decision Outcome

You chose to **{decision}** with {trade_amount:,.0f} credits.

**Result:** {format_percentage(outcome.outcome_percentage)} ({format_currency(outcome.outcome_amount)})

**Portfolio:** {format_currency(outcome.portfolio_after)}
"""

    # Check if game is complete
    if session.trading_engine.is_game_complete():
        return show_results()

    # Move to next scenario
    next_scenario = session.trading_engine.get_next_scenario()
    session.current_scenario = next_scenario
    session.scenario_start_time = time.time()
    session.ai_advice_shown_time = 0
    session.proactive_advice_shown = False
    session.proactive_advice_engaged = False
    session.chat_queries_this_scenario = 0
    session.chatbot.clear_history()

    # Generate content for next scenario
    scenario_text = format_scenario_display(next_scenario)

    # Maybe show proactive advice
    proactive_message = ""
    proactive_response = session.chatbot.generate_proactive_advice(
        next_scenario,
        session.visible_params,
        session.hidden_params
    )

    if proactive_response:
        proactive_message = f"**AI Advisor:** {proactive_response.message}"
        session.proactive_advice_shown = True
        session.ai_advice_shown_time = time.time()

        tracker.record_chat_interaction(ChatInteraction(
            interaction_id=str(uuid.uuid4())[:12],
            participant_id=session.participant_id,
            timestamp=datetime.now().isoformat(),
            scenario_id=next_scenario.scenario_id,
            interaction_type="proactive",
            user_query=None,
            ai_response=proactive_response.message,
            explanation_depth=session.visible_params.explanation_depth,
            communication_style=session.visible_params.communication_style,
            confidence_framing=session.hidden_params.confidence_framing,
            risk_bias=session.hidden_params.risk_bias,
            response_time_ms=0,
            user_engaged=False,
            dismissed=False
        ))

    # Generate main recommendation
    ai_recommendation = session.chatbot.generate_ai_recommendation(
        next_scenario,
        session.visible_params,
        session.hidden_params
    )
    session.ai_advice_shown_time = time.time()

    progress = session.trading_engine.get_progress_info()

    return (
        gr.update(visible=True),   # Keep game visible
        gr.update(visible=False),  # Keep results hidden
        format_currency(session.trading_engine.portfolio.total_value),
        f"Scenario {progress['current_scenario']} of {progress['total_scenarios']}",
        scenario_text,
        f"**AI Recommendation:**\n\n{ai_recommendation.message}",
        proactive_message,
        [],  # Clear chat
        gr.update(value=50),
        gr.update(value="HOLD"),
        outcome_message
    )


def show_results() -> Tuple:
    """Show the final results screen."""
    if not session.trading_engine:
        return (gr.update(),) * 11

    # Complete the session
    summary = session.trading_engine.get_game_summary()
    tracker.complete_session(
        session.participant_id,
        summary["final_portfolio"]
    )

    # Get full session summary
    session_summary = tracker.get_session_summary(session.participant_id)

    results_text = f"""
# Experiment Complete

Thank you for participating!

---

## Your Results

**Participant ID:** {session.participant_id}
(Save this ID if you need to reference your data)

---

### Portfolio Performance

| Metric | Value |
|--------|-------|
| Starting Portfolio | {format_currency(summary['initial_portfolio'])} |
| Final Portfolio | {format_currency(summary['final_portfolio'])} |
| Total Return | {format_currency(summary['total_return'])} |
| Return Percentage | {summary['return_percentage']:.1f}% |

---

### Decision Analysis

| Metric | Value |
|--------|-------|
| Total Decisions | {summary['total_decisions']} |
| AI Follow Rate | {summary['ai_follow_rate']*100:.1f}% |
| Optimal Decision Rate | {summary['optimal_decision_rate']*100:.1f}% |

---

### AI Interaction Summary

| Metric | Value |
|--------|-------|
| Times Followed Correct AI | {summary['followed_correct_ai']} / {summary['ai_correct_scenarios']} |
| Times Followed Incorrect AI | {summary['followed_incorrect_ai']} / {summary['ai_incorrect_scenarios']} |
| Chat Queries Made | {session_summary.get('total_chat_queries', 0)} |

---

### Decision History

"""

    for d in summary['decisions']:
        followed = "Followed AI" if d['followed_ai'] else "Disagreed with AI"
        optimal = "Optimal" if d['was_optimal'] else "Suboptimal"
        results_text += f"- **{d['scenario']}**: {d['decision']} → {d['outcome']} ({followed}, {optimal})\n"

    return (
        gr.update(visible=False),  # Hide game
        gr.update(visible=True),   # Show results
        results_text,
        "", "", "", "", "", [], gr.update(), gr.update(), ""
    )


# Build the Gradio interface
def create_app():
    """Create and return the Gradio application."""

    with gr.Blocks(
        title="TradeVerse AI Experiment",
        theme=gr.themes.Soft(),
        css="""
        .container { max-width: 1200px; margin: auto; }
        .scenario-box { background: #f8f9fa; padding: 20px; border-radius: 10px; }
        .ai-advice { background: #e3f2fd; padding: 15px; border-radius: 8px; border-left: 4px solid #2196f3; }
        .proactive-advice { background: #fff3e0; padding: 15px; border-radius: 8px; border-left: 4px solid #ff9800; }
        .outcome-box { padding: 15px; border-radius: 8px; margin-top: 10px; }
        """
    ) as app:

        # ==================== WELCOME SCREEN ====================
        with gr.Column(visible=True) as welcome_section:
            gr.Markdown("""
            # TradeVerse AI Trading Experiment

            Welcome to this research study on AI-assisted decision making.

            ---

            ## About This Experiment

            You will participate in a simulated trading game set in the **TradeVerse**, a fictional
            financial universe. You will make trading decisions for several companies while receiving
            advice from an AI assistant.

            **Important:** All companies, markets, and situations are fictional. No real-world
            financial knowledge is required or applicable.

            ---

            ## What You'll Do

            1. **Review Trading Scenarios** - Each scenario presents a company and market situation
            2. **Consult the AI** - An AI advisor will offer recommendations and answer questions
            3. **Make Decisions** - Choose to BUY, SELL, or HOLD for each scenario
            4. **See Outcomes** - Learn the results of your decisions

            ---

            ## Data Collection

            This experiment collects:
            - Your trading decisions and confidence levels
            - Your interactions with the AI advisor
            - Timing information

            All data is anonymized. Your participant ID contains no personal information.

            ---

            ## Consent

            By clicking "Start Experiment", you consent to participate in this research study.
            You may stop at any time.
            """)

            with gr.Row():
                condition_dropdown = gr.Dropdown(
                    choices=["Random (A/B Testing)"] + list(EXPERIMENT_CONDITIONS.keys()),
                    value="Random (A/B Testing)",
                    label="Experimental Condition (Researcher Use)",
                    visible=True  # Set to False in production
                )

            start_btn = gr.Button("Start Experiment", variant="primary", size="lg")

        # ==================== GAME SCREEN ====================
        with gr.Column(visible=False) as game_section:

            # Status bar
            with gr.Row():
                participant_display = gr.Markdown("Participant: ---")
                portfolio_display = gr.Markdown("Portfolio: ---")
                progress_display = gr.Markdown("Progress: ---")

            with gr.Row():
                # Left column: Scenario and Decision
                with gr.Column(scale=2):
                    scenario_display = gr.Markdown(
                        "Loading scenario...",
                        elem_classes=["scenario-box"]
                    )

                    # AI Recommendation
                    ai_recommendation_display = gr.Markdown(
                        "",
                        elem_classes=["ai-advice"]
                    )

                    # Proactive advice (may or may not show)
                    proactive_display = gr.Markdown(
                        "",
                        elem_classes=["proactive-advice"]
                    )

                    # Decision controls
                    gr.Markdown("### Your Decision")

                    with gr.Row():
                        decision_radio = gr.Radio(
                            choices=["BUY", "HOLD", "SELL"],
                            value="HOLD",
                            label="Action"
                        )
                        trade_amount = gr.Slider(
                            minimum=1000,
                            maximum=50000,
                            value=10000,
                            step=1000,
                            label="Trade Amount (credits)"
                        )

                    confidence_slider = gr.Slider(
                        minimum=0,
                        maximum=100,
                        value=50,
                        step=5,
                        label="How confident are you in this decision?"
                    )

                    submit_btn = gr.Button("Submit Decision", variant="primary")

                    outcome_display = gr.Markdown("")

                # Right column: AI Chat and Controls
                with gr.Column(scale=1):
                    gr.Markdown("### AI Advisor Settings")

                    explanation_slider = gr.Slider(
                        minimum=0,
                        maximum=100,
                        value=50,
                        step=10,
                        label="Explanation Depth",
                        info="Minimal ← → Detailed"
                    )

                    style_slider = gr.Slider(
                        minimum=0,
                        maximum=100,
                        value=50,
                        step=10,
                        label="Communication Style",
                        info="Formal ← → Casual"
                    )

                    update_params_btn = gr.Button("Update AI Settings", size="sm")

                    gr.Markdown("### Ask the AI")

                    chatbot_display = gr.Chatbot(
                        label="Chat with AI Advisor",
                        height=300
                    )

                    chat_input = gr.Textbox(
                        placeholder="Ask a question about this scenario...",
                        label="Your Question",
                        lines=2
                    )

                    chat_btn = gr.Button("Send", size="sm")

        # ==================== RESULTS SCREEN ====================
        with gr.Column(visible=False) as results_section:
            results_display = gr.Markdown("Loading results...")

            restart_btn = gr.Button("Start New Session", variant="primary")

        # ==================== EVENT HANDLERS ====================

        # Start experiment
        start_btn.click(
            fn=start_experiment,
            inputs=[condition_dropdown],
            outputs=[
                welcome_section,
                game_section,
                results_section,
                participant_display,
                portfolio_display,
                progress_display,
                scenario_display,
                ai_recommendation_display,
                proactive_display,
                chatbot_display,
                confidence_slider,
                decision_radio
            ]
        )

        # Update AI parameters
        update_params_btn.click(
            fn=update_ai_params,
            inputs=[explanation_slider, style_slider],
            outputs=[ai_recommendation_display]
        )

        # Chat interaction
        chat_btn.click(
            fn=handle_chat_query,
            inputs=[chat_input, chatbot_display],
            outputs=[chatbot_display, chat_input]
        )

        chat_input.submit(
            fn=handle_chat_query,
            inputs=[chat_input, chatbot_display],
            outputs=[chatbot_display, chat_input]
        )

        # Submit decision
        submit_btn.click(
            fn=submit_decision,
            inputs=[decision_radio, confidence_slider, trade_amount],
            outputs=[
                game_section,
                results_section,
                portfolio_display,
                progress_display,
                scenario_display,
                ai_recommendation_display,
                proactive_display,
                chatbot_display,
                confidence_slider,
                decision_radio,
                outcome_display
            ]
        )

        # Restart
        restart_btn.click(
            fn=lambda: (gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)),
            outputs=[welcome_section, game_section, results_section]
        )

    return app


# Launch the application
if __name__ == "__main__":
    app = create_app()
    app.launch(debug=True)