# Install necessary libraries
!pip install pandas tqdm

import logging
import datetime
import time
import random
from collections import deque
from typing import Dict, Any, List, Optional

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm.notebook import tqdm

# Configure basic logging for the notebook
# This ensures that log messages are displayed in a human-readable format
# and can be captured easily.
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def create_trade_log_state() -> Dict[str, Any]:
    """
    Initializes the state dictionary for trade reason logging.

    Returns
    -------
    dict
        An initial state dictionary with an empty list for 'trade_reasons'.

    Examples
    --------
    >>> state = create_trade_log_state()
    >>> isinstance(state, dict)
    True
    >>> 'trade_reasons' in state
    True
    >>> isinstance(state['trade_reasons'], list)
    True
    """
    logger.info("Initializing trade log state.")
    return {
        "trade_reasons": []
    }

def log_trade_reason(
    state: Dict[str, Any],
    instrument: str,
    direction: str,
    price: float,
    quantity: float,
    reason_code: str,
    description: str,
    tags: Optional[List[str]] = None
) -> Dict[str, Any]:
    """
    Logs a human-readable trade reason into the state.

    Parameters
    ----------
    state : dict
        Current state dictionary, must contain 'trade_reasons' list.
    instrument : str
        The trading instrument (e.g., 'AAPL', 'EURUSD').
    direction : str
        The direction of the trade ('BUY' or 'SELL').
    price : float
        The price at which the trade occurred.
    quantity : float
        The quantity traded.
    reason_code : str
        A short, predefined code for the reason (e.g., 'MOMENTUM_BREAKOUT').
    description : str
        A detailed, human-readable description of the trade reason.
    tags : Optional[List[str]], optional
        Optional list of tags for further categorization, defaults to None.

    Returns
    -------
    dict
        The updated state dictionary with the new trade reason appended.

    Examples
    --------
    >>> state = create_trade_log_state()
    >>> state = log_trade_reason(state, 'MSFT', 'BUY', 150.25, 100.0,
    ...                          'SUPPORT_BOUNCE', 'Price bounced off daily support level.',
    ...                          tags=['technical'])
    >>> len(state['trade_reasons']) == 1
    True
    >>> state['trade_reasons'][0]['instrument'] == 'MSFT'
    True
    """
    if 'trade_reasons' not in state:
        logger.error("State dictionary missing 'trade_reasons' list. Initializing.")
        state['trade_reasons'] = []

    trade_entry = {
        "timestamp": datetime.datetime.now().isoformat(),
        "instrument": instrument,
        "direction": direction,
        "price": price,
        "quantity": quantity,
        "reason_code": reason_code,
        "description": description,
        "tags": tags if tags is not None else []
    }
    state['trade_reasons'].append(trade_entry)
    logger.debug(f"Logged trade reason for {instrument} ({direction}): {reason_code}")
    return state

def get_trade_reasons_df(state: Dict[str, Any]) -> pd.DataFrame:
    """
    Converts the logged trade reasons from the state into a pandas DataFrame.

    Parameters
    ----------
    state : dict
        Current state dictionary, must contain 'trade_reasons' list.

    Returns
    -------
    pd.DataFrame
        A DataFrame containing all logged trade reasons.
        Returns an empty DataFrame if no reasons are logged.

    Examples
    --------
    >>> state = create_trade_log_state()
    >>> df = get_trade_reasons_df(state)
    >>> isinstance(df, pd.DataFrame)
    True
    >>> df.empty
    True
    >>> state = log_trade_reason(state, 'GOOG', 'SELL', 2500.0, 50.0,
    ...                          'RESISTANCE_REJECT', 'Price rejected at daily resistance.',
    ...                          tags=['technical'])
    >>> df = get_trade_reasons_df(state)
    >>> not df.empty
    True
    """
    if not state.get('trade_reasons'):
        logger.warning("No trade reasons found in state. Returning empty DataFrame.")
        return pd.DataFrame()

    df = pd.DataFrame(state['trade_reasons'])
    df['timestamp'] = pd.to_datetime(df['timestamp'])
    logger.info(f"Successfully converted {len(df)} trade reasons to DataFrame.")
    return df

def create_simulated_market_data(
    start_date: str,
    end_date: str,
    instrument: str,
    initial_price: float,
    volatility: float
) -> pd.DataFrame:
    """
    Generates simulated market price data.

    Parameters
    ----------
    start_date : str
        Start date for the simulation (YYYY-MM-DD).
    end_date : str
        End date for the simulation (YYYY-MM-DD).
    instrument : str
        The name of the instrument.
    initial_price : float
        The starting price.
    volatility : float
        A factor for price fluctuations.

    Returns
    -------
    pd.DataFrame
        DataFrame with 'timestamp' and 'price' columns.

    Examples
    --------
    >>> df_sim = create_simulated_market_data('2023-01-01', '2023-01-05', 'TEST', 100.0, 0.01)
    >>> isinstance(df_sim, pd.DataFrame)
    True
    >>> 'price' in df_sim.columns
    True
    >>> len(df_sim) > 0
    True
    """
    dates = pd.date_range(start=start_date, end=end_date, freq='H')
    prices = [initial_price]
    for _ in range(1, len(dates)):
        change = np.random.normal(0, volatility) * prices[-1]
        prices.append(prices[-1] + change)

    df = pd.DataFrame({
        'timestamp': dates,
        'instrument': instrument,
        'price': prices
    })
    logger.info(f"Generated {len(df)} simulated data points for {instrument}.")
    return df

def simulate_trading_strategy(
    state: Dict[str, Any],
    market_data: pd.DataFrame,
    window_size: int = 10
) -> Dict[str, Any]:
    """
    Simulates a simple trading strategy and logs trade reasons.

    The strategy generates a 'BUY' signal when price crosses above a simple moving average (SMA),
    and a 'SELL' signal when price crosses below the SMA.

    Parameters
    ----------
    state : dict
        Current state dictionary for trade logging.
    market_data : pd.DataFrame
        DataFrame with 'timestamp', 'instrument', and 'price' columns.
    window_size : int, optional
        The window size for the simple moving average, defaults to 10.

    Returns
    -------
    dict
        The updated state dictionary with newly logged trade reasons.

    Examples
    --------
    >>> initial_state = create_trade_log_state()
    >>> sim_data = create_simulated_market_data('2023-01-01', '2023-01-02', 'TEST', 100.0, 0.005)
    >>> updated_state = simulate_trading_strategy(initial_state, sim_data, window_size=5)
    >>> isinstance(updated_state, dict)
    True
    >>> len(updated_state['trade_reasons']) >= 0
    True
    """
    df = market_data.copy()
    instrument = df['instrument'].iloc[0]
    df['SMA'] = df['price'].rolling(window=window_size).mean()

    # Initialize previous price/SMA relationship
    prev_price_above_sma = False
    if not df.iloc[:window_size].empty:
      if df['price'].iloc[window_size - 1] > df['SMA'].iloc[window_size - 1]:
          prev_price_above_sma = True

    # Simulate with random jitter for realism
    for i in tqdm(range(window_size, len(df)), desc=f"Simulating {instrument} trades"):
        current_row = df.iloc[i]
        current_price = current_row['price']
        current_sma = current_row['SMA']

        # Add random jitter to simulate processing delay or market dynamics
        time.sleep(0.001 * random.uniform(0.8, 1.2))

        if current_price > current_sma and not prev_price_above_sma:
            # Price crossed above SMA: BUY signal
            state = log_trade_reason(
                state,
                instrument,
                'BUY',
                current_price,
                random.choice([10, 20, 50]), # Random quantity
                'SMA_CROSS_UP',
                f"Price {current_price:.2f} crossed above {window_size}-period SMA {current_sma:.2f}.",
                tags=['technical', 'trend_following']
            )
            prev_price_above_sma = True
        elif current_price < current_sma and prev_price_above_sma:
            # Price crossed below SMA: SELL signal
            state = log_trade_reason(
                state,
                instrument,
                'SELL',
                current_price,
                random.choice([10, 20, 50]), # Random quantity
                'SMA_CROSS_DOWN',
                f"Price {current_price:.2f} crossed below {window_size}-period SMA {current_sma:.2f}.",
                tags=['technical', 'trend_following']
            )
            prev_price_above_sma = False

    logger.info(f"Finished simulating trades for {instrument}.")
    return state

# 1. Initialize the trade log state
trade_log_state = create_trade_log_state()

# 2. Log some initial manual trade reasons (before simulation)
trade_log_state = log_trade_reason(
    trade_log_state, 'ETHUSD', 'BUY', 1800.50, 0.5,
    'MANUAL_ENTRY', 'Initial discretionary buy order based on market sentiment.',
    tags=['manual', 'sentiment']
)
trade_log_state = log_trade_reason(
    trade_log_state, 'BTCUSD', 'SELL', 30500.00, 0.1,
    'RISK_REDUCTION', 'Reducing exposure due to upcoming economic news.',
    tags=['risk_management', 'macro']
)

# 3. Create simulated market data for multiple instruments
sim_data_a = create_simulated_market_data('2023-01-01', '2023-01-10', 'STOCK_A', 100.0, 0.005)
sim_data_b = create_simulated_market_data('2023-01-01', '2023-01-10', 'STOCK_B', 50.0, 0.01)
sim_data_c = create_simulated_market_data('2023-01-01', '2023-01-10', 'STOCK_C', 200.0, 0.003)

# Combine data for simulation
combined_sim_data = pd.concat([sim_data_a, sim_data_b, sim_data_c])

# 4. Simulate trading strategy and log reasons
trade_log_state = simulate_trading_strategy(trade_log_state, sim_data_a, window_size=5)
trade_log_state = simulate_trading_strategy(trade_log_state, sim_data_b, window_size=15)
trade_log_state = simulate_trading_strategy(trade_log_state, sim_data_c, window_size=8)

# 5. Retrieve all logged trade reasons into a DataFrame
all_trade_reasons_df = get_trade_reasons_df(trade_log_state)

print("\n--- All Logged Trade Reasons ---")
display(all_trade_reasons_df.head(10))
print(f"Total unique reason codes: {all_trade_reasons_df['reason_code'].nunique()}")
print(f"Total unique instruments: {all_trade_reasons_df['instrument'].nunique()}")

if not all_trade_reasons_df.empty:
    # Plot 1: Distribution of Trade Reason Codes
    plt.figure(figsize=(12, 6))
    sns.countplot(y='reason_code', data=all_trade_reasons_df, order=all_trade_reasons_df['reason_code'].value_counts().index, palette='viridis')
    plt.title('Distribution of Trade Reason Codes')
    plt.xlabel('Number of Trades')
    plt.ylabel('Reason Code')
    plt.tight_layout()
    plt.show()

    # Plot 2: Trades Over Time by Instrument
    plt.figure(figsize=(14, 7))
    sns.scatterplot(
        x='timestamp',
        y='price',
        hue='instrument',
        size='quantity',
        style='direction',
        data=all_trade_reasons_df.sort_values('timestamp'),
        sizes=(50, 500), # Size range for quantity
        alpha=0.7
    )
    plt.title('Trades Over Time by Instrument, Price, and Quantity')
    plt.xlabel('Timestamp')
    plt.ylabel('Price')
    plt.xticks(rotation=45)
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.tight_layout()
    plt.show()

    # Plot 3: Trade Direction Breakdown by Instrument
    plt.figure(figsize=(12, 6))
    sns.countplot(x='instrument', hue='direction', data=all_trade_reasons_df, palette='coolwarm')
    plt.title('Trade Direction Breakdown by Instrument')
    plt.xlabel('Instrument')
    plt.ylabel('Number of Trades')
    plt.legend(title='Direction')
    plt.tight_layout()
    plt.show()

    # Plot 4: Average Quantity by Reason Code
    avg_qty_by_reason = all_trade_reasons_df.groupby('reason_code')['quantity'].mean().sort_values(ascending=False).reset_index()
    plt.figure(figsize=(12, 6))
    sns.barplot(x='quantity', y='reason_code', data=avg_qty_by_reason, palette='plasma')
    plt.title('Average Trade Quantity by Reason Code')
    plt.xlabel('Average Quantity')
    plt.ylabel('Reason Code')
    plt.tight_layout()
    plt.show()
else:
    logger.info("No trade reasons to visualize.")