trading.py

"""
Trading engine with scenario management and portfolio tracking.
Handles game state, decision processing, and outcome calculation.
"""

import random
import uuid
from datetime import datetime
from typing import List, Optional, Dict, Any, Tuple
from dataclasses import dataclass, field

from config import (
    Scenario, SCENARIOS, ExperimentConfig, DEFAULT_CONFIG,
    ResearcherControlledParams, ParticipantVisibleParams
)


@dataclass
class Portfolio:
    """Represents a participant's portfolio state."""
    cash: float
    initial_value: float
    positions: Dict[str, Dict[str, Any]] = field(default_factory=dict)
    history: List[Dict[str, Any]] = field(default_factory=list)

    @property
    def total_value(self) -> float:
        """Calculate total portfolio value (cash + positions)."""
        position_value = sum(
            pos["shares"] * pos["current_price"]
            for pos in self.positions.values()
        )
        return self.cash + position_value

    @property
    def return_percentage(self) -> float:
        """Calculate portfolio return as percentage."""
        if self.initial_value == 0:
            return 0
        return ((self.total_value - self.initial_value) / self.initial_value) * 100

    def record_state(self, scenario_id: str, action: str, outcome: float):
        """Record a portfolio state change."""
        self.history.append({
            "timestamp": datetime.now().isoformat(),
            "scenario_id": scenario_id,
            "action": action,
            "outcome_pct": outcome,
            "portfolio_value": self.total_value
        })


@dataclass
class DecisionOutcome:
    """Result of a trading decision."""
    scenario_id: str
    decision: str  # "BUY", "SELL", "HOLD"
    outcome_percentage: float
    outcome_amount: float
    portfolio_before: float
    portfolio_after: float
    ai_was_correct: bool
    followed_ai: bool
    was_optimal: bool


class ScenarioManager:
    """Manages scenario selection and ordering for experiments."""

    def __init__(self, config: ExperimentConfig = DEFAULT_CONFIG):
        self.config = config
        self.all_scenarios = SCENARIOS.copy()
        self.session_scenarios: List[Scenario] = []
        self.current_index: int = 0

    def initialize_session(self, shuffle: bool = True) -> List[Scenario]:
        """
        Initialize scenarios for a new session.
        Returns the list of scenarios that will be presented.
        """
        # Select scenarios for this session
        num_scenarios = min(self.config.scenarios_per_session, len(self.all_scenarios))
        self.session_scenarios = self.all_scenarios[:num_scenarios]

        if shuffle:
            random.shuffle(self.session_scenarios)

        self.current_index = 0
        return self.session_scenarios

    def get_current_scenario(self) -> Optional[Scenario]:
        """Get the current scenario."""
        if self.current_index < len(self.session_scenarios):
            return self.session_scenarios[self.current_index]
        return None

    def advance_to_next(self) -> Optional[Scenario]:
        """Move to the next scenario and return it."""
        self.current_index += 1
        return self.get_current_scenario()

    def get_progress(self) -> Tuple[int, int]:
        """Return (current_number, total) for progress display."""
        return (self.current_index + 1, len(self.session_scenarios))

    def is_complete(self) -> bool:
        """Check if all scenarios have been completed."""
        return self.current_index >= len(self.session_scenarios)

    def reset(self):
        """Reset the scenario manager."""
        self.session_scenarios = []
        self.current_index = 0


class TradingEngine:
    """
    Main trading engine that processes decisions and manages game state.
    """

    def __init__(self, config: ExperimentConfig = DEFAULT_CONFIG):
        self.config = config
        self.portfolio: Optional[Portfolio] = None
        self.scenario_manager = ScenarioManager(config)
        self.decisions_made: List[DecisionOutcome] = []

    def start_new_game(self) -> Tuple[Portfolio, Scenario]:
        """
        Start a new trading game.
        Returns the initial portfolio and first scenario.
        """
        # Initialize portfolio
        self.portfolio = Portfolio(
            cash=self.config.initial_portfolio_value,
            initial_value=self.config.initial_portfolio_value
        )

        # Initialize scenarios
        self.scenario_manager.initialize_session(shuffle=True)
        self.decisions_made = []

        first_scenario = self.scenario_manager.get_current_scenario()
        return self.portfolio, first_scenario

    def process_decision(
        self,
        scenario: Scenario,
        decision: str,  # "BUY", "SELL", "HOLD"
        trade_amount: float,
        ai_recommendation: str
    ) -> DecisionOutcome:
        """
        Process a trading decision and return the outcome.
        """
        # Validate decision
        decision = decision.upper()
        if decision not in ["BUY", "SELL", "HOLD"]:
            raise ValueError(f"Invalid decision: {decision}")

        # Get outcome percentage based on decision
        outcome_map = {
            "BUY": scenario.outcome_buy,
            "SELL": scenario.outcome_sell,
            "HOLD": scenario.outcome_hold
        }
        outcome_pct = outcome_map[decision]

        # Calculate outcome amount
        portfolio_before = self.portfolio.total_value

        # For simplicity, we apply the outcome to the trade amount
        # In a more complex system, you might track actual share positions
        if decision in ["BUY", "HOLD"]:
            # Participant is exposed to the stock's movement
            outcome_amount = trade_amount * outcome_pct
        else:  # SELL
            # Participant avoided the stock's movement (inverse)
            outcome_amount = trade_amount * outcome_pct

        # Update portfolio
        self.portfolio.cash += outcome_amount
        portfolio_after = self.portfolio.total_value

        # Record state
        self.portfolio.record_state(scenario.scenario_id, decision, outcome_pct)

        # Determine if AI was followed and if decision was optimal
        followed_ai = (decision == ai_recommendation)
        was_optimal = (decision == scenario.optimal_action)

        outcome = DecisionOutcome(
            scenario_id=scenario.scenario_id,
            decision=decision,
            outcome_percentage=outcome_pct,
            outcome_amount=outcome_amount,
            portfolio_before=portfolio_before,
            portfolio_after=portfolio_after,
            ai_was_correct=scenario.ai_is_correct,
            followed_ai=followed_ai,
            was_optimal=was_optimal
        )

        self.decisions_made.append(outcome)
        return outcome

    def get_next_scenario(self) -> Optional[Scenario]:
        """Get the next scenario in the session."""
        return self.scenario_manager.advance_to_next()

    def is_game_complete(self) -> bool:
        """Check if the game is complete."""
        return self.scenario_manager.is_complete()

    def get_game_summary(self) -> Dict[str, Any]:
        """Get a summary of the completed game."""
        if not self.portfolio:
            return {}

        total_decisions = len(self.decisions_made)
        ai_followed_count = sum(1 for d in self.decisions_made if d.followed_ai)
        optimal_count = sum(1 for d in self.decisions_made if d.was_optimal)

        # Calculate when AI was correct vs wrong
        ai_correct_decisions = [d for d in self.decisions_made if d.ai_was_correct]
        ai_wrong_decisions = [d for d in self.decisions_made if not d.ai_was_correct]

        followed_when_correct = sum(1 for d in ai_correct_decisions if d.followed_ai)
        followed_when_wrong = sum(1 for d in ai_wrong_decisions if d.followed_ai)

        return {
            "initial_portfolio": self.portfolio.initial_value,
            "final_portfolio": self.portfolio.total_value,
            "total_return": self.portfolio.total_value - self.portfolio.initial_value,
            "return_percentage": self.portfolio.return_percentage,
            "total_decisions": total_decisions,
            "ai_follow_rate": ai_followed_count / total_decisions if total_decisions > 0 else 0,
            "optimal_decision_rate": optimal_count / total_decisions if total_decisions > 0 else 0,
            "followed_correct_ai": followed_when_correct,
            "followed_incorrect_ai": followed_when_wrong,
            "ai_correct_scenarios": len(ai_correct_decisions),
            "ai_incorrect_scenarios": len(ai_wrong_decisions),
            "decisions": [
                {
                    "scenario": d.scenario_id,
                    "decision": d.decision,
                    "outcome": f"{d.outcome_percentage * 100:+.1f}%",
                    "followed_ai": d.followed_ai,
                    "was_optimal": d.was_optimal
                }
                for d in self.decisions_made
            ]
        }

    def get_progress_info(self) -> Dict[str, Any]:
        """Get current progress information."""
        current, total = self.scenario_manager.get_progress()
        return {
            "current_scenario": current,
            "total_scenarios": total,
            "progress_percentage": (current / total) * 100 if total > 0 else 0,
            "portfolio_value": self.portfolio.total_value if self.portfolio else 0,
            "portfolio_return": self.portfolio.return_percentage if self.portfolio else 0
        }


def calculate_suggested_trade_amount(
    portfolio_value: float,
    risk_level: int = 50
) -> float:
    """
    Calculate a suggested trade amount based on portfolio and risk level.
    """
    # Base: 10-30% of portfolio depending on risk level
    min_pct = 0.10
    max_pct = 0.30
    risk_factor = risk_level / 100

    suggested_pct = min_pct + (max_pct - min_pct) * risk_factor
    return round(portfolio_value * suggested_pct, 2)


def format_currency(amount: float, symbol: str = "credits") -> str:
    """Format a currency amount for display."""
    if amount >= 0:
        return f"{amount:,.2f} {symbol}"
    else:
        return f"-{abs(amount):,.2f} {symbol}"


def format_percentage(value: float, include_sign: bool = True) -> str:
    """Format a percentage for display."""
    if include_sign:
        return f"{value * 100:+.1f}%"
    return f"{value * 100:.1f}%"