#!pip install numpy pandas matplotlib seaborn loguru

import logging
import collections
import random
import time
from typing import Dict, Any, List

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

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

sns.set_theme(style="whitegrid")

def create_market_state(initial_price: float, volatility: float, base_liquidity: int, random_seed: int = None) -> Dict[str, Any]:
    """
    Initializes the market state with given parameters.

    Parameters
    ----------
    initial_price : float
        The starting price of the asset.
    volatility : float
        The standard deviation of price changes, representing market fluctuations.
    base_liquidity : int
        The average number of shares available to be filled at any given moment.
    random_seed : int, optional
        A seed for the random number generator to ensure reproducibility, defaults to None.

    Returns
    -------
    Dict[str, Any]
        An initialized dictionary containing the market state.

    Examples
    --------
    >>> market = create_market_state(initial_price=100.0, volatility=0.5, base_liquidity=500)
    >>> 'current_price' in market
    True
    """
    if random_seed is not None:
        random.seed(random_seed)
        np.random.seed(random_seed)

    market_state = {
        'current_price': initial_price,
        'volatility': volatility,
        'base_liquidity': base_liquidity,
        'price_history': [(0, initial_price)], # (timestamp, price)
        'fill_history': [] # (timestamp, filled_qty, fill_price)
    }
    logger.info(f"Market state initialized: {market_state}")
    return market_state

def create_order_state(total_quantity: int, displayed_quantity: int) -> Dict[str, Any]:
    """
    Initializes the iceberg order state.

    Parameters
    ----------
    total_quantity : int
        The total number of shares to be traded.
    displayed_quantity : int
        The maximum quantity of shares visible in the order book at any one time.

    Returns
    -------
    Dict[str, Any]
        An initialized dictionary containing the iceberg order's state.

    Examples
    --------
    >>> order = create_order_state(total_quantity=10000, displayed_quantity=1000)
    >>> 'remaining_quantity' in order
    True
    """
    order_state = {
        'total_quantity': total_quantity,
        'displayed_quantity': min(displayed_quantity, total_quantity), # Displayed cannot exceed total
        'remaining_quantity': total_quantity,
        'executed_quantity': 0,
        'fill_details': [], # List of {'timestamp', 'filled_qty', 'fill_price', 'slice_size'}
        'status': 'ACTIVE'
    }
    logger.info(f"Order state initialized: {order_state}")
    return order_state

def _simulate_market_fill(market_state: Dict[str, Any], quantity_to_fill: int) -> tuple[int, float]:
    """
    Simulates market fill for a given quantity, returning filled quantity and price.

    Parameters
    ----------
    market_state : Dict[str, Any]
        The current market environment dictionary.
    quantity_to_fill : int
        The number of shares attempted to be filled in this slice.

    Returns
    -------
    tuple[int, float]
        A tuple containing the actual `filled_quantity` and the `fill_price`.
    """
    # Simulate market price fluctuation
    price_change = np.random.normal(0, market_state['volatility'])
    market_state['current_price'] += price_change
    market_state['current_price'] = max(0.01, market_state['current_price']) # Ensure price is not negative

    # Simulate available liquidity for this slice
    # Liquidity can fluctuate around the base_liquidity
    available_liquidity = max(0, int(market_state['base_liquidity'] * (1 + np.random.uniform(-0.3, 0.3))))

    # Simulate fill quantity: max of what's requested, available liquidity, and a random fill rate
    fill_rate = random.uniform(0.7, 1.0) # A slice might not get 100% filled instantly

    # Calculate base filled quantity from rate and requested quantity
    base_filled_from_rate = int(quantity_to_fill * fill_rate)

    # Ensure that if quantity_to_fill is 1, and there is liquidity, at least 1 is filled.
    # This prevents an infinite loop for the last share if fill_rate makes it to 0.
    if quantity_to_fill == 1 and available_liquidity >= 1 and base_filled_from_rate == 0:
        filled_quantity_from_rate = 1
    else:
        filled_quantity_from_rate = base_filled_from_rate

    filled_quantity = min(quantity_to_fill, available_liquidity, filled_quantity_from_rate)

    fill_price = market_state['current_price'] * (1 + random.uniform(-0.001, 0.001)) # Slight price variation for fill

    logger.debug(f"Simulated fill: requested={quantity_to_fill}, available_liquidity={available_liquidity}, filled={filled_quantity}, price={fill_price:.2f}")
    return filled_quantity, fill_price

def execute_iceberg_slice(order_state: Dict[str, Any], market_state: Dict[str, Any], current_timestamp: float) -> Dict[str, Any]:
    """
    Attempts to execute a single slice of the iceberg order.

    Parameters
    ----------
    order_state : Dict[str, Any]
        The current state of the iceberg order.
    market_state : Dict[str, Any]
        The current state of the market.
    current_timestamp : float
        The timestamp at which this slice is being executed.

    Returns
    -------
    Dict[str, Any]
        The updated `order_state` dictionary.

    Examples
    --------
    >>> market = create_market_state(100, 0.1, 100)
    >>> order = create_order_state(1000, 100)
    >>> updated_order = execute_iceberg_slice(order, market, 1.0)
    >>> updated_order['executed_quantity'] >= 0
    True
    """
    if order_state['remaining_quantity'] <= 0:
        order_state['status'] = 'COMPLETED'
        logger.info(f"Order {order_state['total_quantity']} completed. No more slices to execute.")
        return order_state

    # The quantity to display for this slice is the minimum of displayed_quantity and remaining_quantity
    slice_quantity = min(order_state['displayed_quantity'], order_state['remaining_quantity'])
    logger.info(f"Attempting to execute a slice of {slice_quantity} shares at {current_timestamp:.2f}s.")

    try:
        filled_qty, fill_price = _simulate_market_fill(market_state, slice_quantity)

        order_state['executed_quantity'] += filled_qty
        order_state['remaining_quantity'] -= filled_qty

        market_state['price_history'].append((current_timestamp, market_state['current_price']))
        if filled_qty > 0:
            market_state['fill_history'].append((current_timestamp, filled_qty, fill_price))
            order_state['fill_details'].append({
                'timestamp': current_timestamp,
                'filled_qty': filled_qty,
                'fill_price': fill_price,
                'slice_size': slice_quantity
            })
            logger.info(f"Slice filled: {filled_qty} shares at {fill_price:.2f} (Market Price: {market_state['current_price']:.2f}). Remaining: {order_state['remaining_quantity']}")
        else:
            logger.warning(f"Slice of {slice_quantity} shares received 0 fill. Remaining: {order_state['remaining_quantity']}")

    except Exception as e:
        logger.error(f"Error during slice execution at {current_timestamp:.2f}s: {e}")
        # Implement a simple retry mechanism with jitter if desired, for now just log and continue
        time.sleep(random.uniform(0.1, 0.5)) # Small random delay before potential next attempt

    if order_state['remaining_quantity'] <= 0:
        order_state['status'] = 'COMPLETED'
        logger.info(f"Iceberg order of {order_state['total_quantity']} shares fully completed!")

    return order_state

def split_and_execute_iceberg_order(
    order_state: Dict[str, Any],
    market_state: Dict[str, Any],
    max_steps: int = 1000,
    max_time_seconds: int = 600,
    min_delay_between_slices: float = 0.1,
    max_delay_between_slices: float = 1.0
) -> tuple[Dict[str, Any], Dict[str, Any]]:
    """
    Orchestrates the splitting and execution of an iceberg order.

    Parameters
    ----------
    order_state : Dict[str, Any]
        The initial state of the iceberg order.
    market_state : Dict[str, Any]
        The initial state of the market.
    max_steps : int, optional
        The maximum number of execution attempts for slices, defaults to 1000.
    max_time_seconds : int, optional
        The maximum total time allowed for execution in simulated seconds, defaults to 600.
    min_delay_between_slices : float, optional
        Minimum random delay between slice executions, defaults to 0.1 seconds.
    max_delay_between_slices : float, optional
        Maximum random delay between slice executions, defaults to 1.0 seconds.

    Returns
    -------
    tuple[Dict[str, Any], Dict[str, Any]]
        A tuple containing the final `order_state` and `market_state`.

    Examples
    --------
    >>> market = create_market_state(100, 0.1, 100)
    >>> order = create_order_state(1000, 100)
    >>> final_order, final_market = split_and_execute_iceberg_order(order, market, max_steps=50)
    >>> final_order['executed_quantity'] >= 0
    True
    """
    current_time = 0.0
    step_count = 0
    start_real_time = time.time()

    logger.info(f"Starting iceberg order execution for {order_state['total_quantity']} shares...")

    while order_state['remaining_quantity'] > 0 and step_count < max_steps and current_time < max_time_seconds:
        step_count += 1
        logger.debug(f"Execution Step {step_count} at simulated time {current_time:.2f}s")

        order_state = execute_iceberg_slice(order_state, market_state, current_time)

        if order_state['status'] == 'COMPLETED':
            break

        # Simulate time passing between slices with random jitter
        delay = random.uniform(min_delay_between_slices, max_delay_between_slices)
        current_time += delay
        # time.sleep(0.01) # Small real-time sleep for demonstration purposes if needed, removed for faster simulation

    end_real_time = time.time()
    elapsed_real_time = end_real_time - start_real_time

    if order_state['remaining_quantity'] > 0:
        order_state['status'] = 'PARTIALLY_FILLED_TIMEOUT' if current_time >= max_time_seconds else 'PARTIALLY_FILLED_STEPS'
        logger.warning(f"Iceberg order partially filled. Remaining: {order_state['remaining_quantity']} shares. Status: {order_state['status']}")
    else:
        logger.info(f"Iceberg order fully executed in {step_count} steps and {current_time:.2f} simulated seconds. (Real time elapsed: {elapsed_real_time:.2f}s)")

    return order_state, market_state

def summarize_execution_results(order_state: Dict[str, Any], market_state: Dict[str, Any], initial_market_price: float) -> Dict[str, Any]:
    """
    Calculates and returns summary statistics of the iceberg order execution.

    Parameters
    ----------
    order_state : Dict[str, Any]
        The final state of the iceberg order after execution.
    market_state : Dict[str, Any]
        The final state of the market after execution.
    initial_market_price : float
        The market price at the beginning of the execution.

    Returns
    -------
    Dict[str, Any]
        A dictionary containing summary statistics of the execution.

    Examples
    --------
    >>> market = create_market_state(100, 0.1, 100)
    >>> order = create_order_state(1000, 100)
    >>> final_order, final_market = split_and_execute_iceberg_order(order, market, max_steps=10)
    >>> summary = summarize_execution_results(final_order, final_market, 100.0)
    >>> 'total_executed_value' in summary
    True
    """
    fill_details = order_state['fill_details']
    if not fill_details:
        return {
            'total_quantity_requested': order_state['total_quantity'],
            'total_executed_quantity': 0,
            'fill_percentage': 0.0,
            'average_fill_price': np.nan,
            'total_executed_value': 0.0,
            'total_simulated_time_taken': 0.0,
            'slippage_per_share': np.nan,
            'final_status': order_state['status']
        }

    total_executed_quantity = sum(d['filled_qty'] for d in fill_details)
    total_executed_value = sum(d['filled_qty'] * d['fill_price'] for d in fill_details)
    average_fill_price = total_executed_value / total_executed_quantity if total_executed_quantity > 0 else np.nan
    max_timestamp = max(d['timestamp'] for d in fill_details)

    # Slippage: difference between average fill price and initial market price
    slippage_per_share = average_fill_price - initial_market_price if average_fill_price else np.nan

    summary = {
        'total_quantity_requested': order_state['total_quantity'],
        'total_executed_quantity': total_executed_quantity,
        'fill_percentage': (total_executed_quantity / order_state['total_quantity']) * 100,
        'average_fill_price': average_fill_price,
        'total_executed_value': total_executed_value,
        'total_simulated_time_taken': max_timestamp,
        'slippage_per_share': slippage_per_share,
        'final_status': order_state['status']
    }
    logger.info(f"Execution summary generated: {summary}")
    return summary

# --- Scenario 1: Typical Iceberg Order Execution ---

logger.info("\n--- Running Demonstration Scenario 1: Typical Execution ---")

# 1. Initialize Market and Order States
initial_market_price = 100.0
market_state_1 = create_market_state(initial_price=initial_market_price, volatility=0.05, base_liquidity=200, random_seed=42)
order_state_1 = create_order_state(total_quantity=5000, displayed_quantity=500)

# 2. Execute the Iceberg Order
final_order_state_1, final_market_state_1 = split_and_execute_iceberg_order(
    order_state_1, market_state_1, max_steps=200, max_time_seconds=120, min_delay_between_slices=0.05, max_delay_between_slices=0.5
)

# 3. Summarize Results
summary_results_1 = summarize_execution_results(final_order_state_1, final_market_state_1, initial_market_price)
print("\n--- Execution Summary (Scenario 1) ---")
display(pd.DataFrame([summary_results_1]))

# 4. Visualization of Market Price and Fills (Scenario 1)

price_df_1 = pd.DataFrame(final_market_state_1['price_history'], columns=['timestamp', 'price'])
fill_df_1 = pd.DataFrame(final_market_state_1['fill_history'], columns=['timestamp', 'filled_qty', 'fill_price'])

fig, ax1 = plt.subplots(figsize=(14, 7))

color = 'tab:blue'
ax1.set_xlabel('Simulated Time (seconds)')
ax1.set_ylabel('Market Price', color=color)
ax1.plot(price_df_1['timestamp'], price_df_1['price'], color=color, label='Market Price', alpha=0.7)
ax1.tick_params(axis='y', labelcolor=color)

ax2 = ax1.twinx() # instantiate a second axes that shares the same x-axis
color = 'tab:red'
ax2.set_ylabel('Fill Price', color=color)  # we already handled the x-label with ax1
ax2.scatter(fill_df_1['timestamp'], fill_df_1['fill_price'], color=color, s=50, label='Fill Price', marker='x')
ax2.tick_params(axis='y', labelcolor=color)

# Add vertical lines for each fill event
for _, row in fill_df_1.iterrows():
    ax1.axvline(row['timestamp'], color='gray', linestyle='--', alpha=0.3)

fig.tight_layout() # otherwise the right y-label is slightly clipped
plt.title('Market Price and Fill Events Over Time (Scenario 1)')
ax1.legend(loc='upper left')
ax2.legend(loc='upper right')
plt.grid(True, linestyle='--', alpha=0.6)
plt.show()

# 5. Visualization of Order Execution Progression (Scenario 1)

# Prepare data for progression plot
progression_data = []
remaining_qty = order_state_1['total_quantity']
executed_qty = 0
current_time = 0.0

# Initial state
progression_data.append({'timestamp': 0.0, 'Remaining Quantity': remaining_qty, 'Executed Quantity': executed_qty})

for detail in final_order_state_1['fill_details']:
    current_time = detail['timestamp']
    executed_qty += detail['filled_qty']
    remaining_qty -= detail['filled_qty']
    progression_data.append({'timestamp': current_time, 'Remaining Quantity': remaining_qty, 'Executed Quantity': executed_qty})

progression_df_1 = pd.DataFrame(progression_data)

plt.figure(figsize=(14, 7))
sns.lineplot(x='timestamp', y='value', hue='variable', data=pd.melt(progression_df_1, ['timestamp']), linewidth=2.5)
plt.title('Iceberg Order Quantity Progression (Scenario 1)')
plt.xlabel('Simulated Time (seconds)')
plt.ylabel('Quantity (shares)')
plt.grid(True, linestyle='--', alpha=0.6)
plt.legend(title='Quantity Type')
plt.show()

# --- Scenario 2: High Volatility, Lower Liquidity ---

logger.info("\n--- Running Demonstration Scenario 2: High Volatility ---")

# 1. Initialize Market and Order States with different parameters
initial_market_price_2 = 100.0
market_state_2 = create_market_state(initial_price=initial_market_price_2, volatility=0.2, base_liquidity=100, random_seed=123)
order_state_2 = create_order_state(total_quantity=5000, displayed_quantity=500)

# 2. Execute the Iceberg Order
final_order_state_2, final_market_state_2 = split_and_execute_iceberg_order(
    order_state_2, market_state_2, max_steps=200, max_time_seconds=120, min_delay_between_slices=0.05, max_delay_between_slices=0.5
)

# 3. Summarize Results
summary_results_2 = summarize_execution_results(final_order_state_2, final_market_state_2, initial_market_price_2)
print("\n--- Execution Summary (Scenario 2: High Volatility) ---")
display(pd.DataFrame([summary_results_2]))

# 4. Visualization of Market Price and Fills (Scenario 2)

price_df_2 = pd.DataFrame(final_market_state_2['price_history'], columns=['timestamp', 'price'])
fill_df_2 = pd.DataFrame(final_market_state_2['fill_history'], columns=['timestamp', 'filled_qty', 'fill_price'])

fig, ax1 = plt.subplots(figsize=(14, 7))

color = 'tab:blue'
ax1.set_xlabel('Simulated Time (seconds)')
ax1.set_ylabel('Market Price', color=color)
ax1.plot(price_df_2['timestamp'], price_df_2['price'], color=color, label='Market Price', alpha=0.7)
ax1.tick_params(axis='y', labelcolor=color)

ax2 = ax1.twinx() # instantiate a second axes that shares the same x-axis
color = 'tab:red'
ax2.set_ylabel('Fill Price', color=color)  # we already handled the x-label with ax1
ax2.scatter(fill_df_2['timestamp'], fill_df_2['fill_price'], color=color, s=50, label='Fill Price', marker='x')
ax2.tick_params(axis='y', labelcolor=color)

# Add vertical lines for each fill event
for _, row in fill_df_2.iterrows():
    ax1.axvline(row['timestamp'], color='gray', linestyle='--', alpha=0.3)

fig.tight_layout() # otherwise the right y-label is slightly clipped
plt.title('Market Price and Fill Events Over Time (Scenario 2: High Volatility)')
ax1.legend(loc='upper left')
ax2.legend(loc='upper right')
plt.grid(True, linestyle='--', alpha=0.6)
plt.show()