# Install necessary packages
!pip install -qqq pandas matplotlib seaborn scipy ping3

# Suppress warnings for cleaner output
import warnings
warnings.filterwarnings('ignore')

import time
import random
import logging
from collections import deque
import datetime

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats
# from ping3 import ping, errors # Commented out for simulation purposes

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

def create_latency_state(
    exchanges: list[str],
    base_latency_ms: float = 50.0,
    latency_std_dev_ms: float = 5.0
) -> dict:
    """
    Initializes the state dictionary for latency benchmarking.

    Parameters
    ----------
    exchanges : list[str]
        A list of exchange names to be benchmarked.
    base_latency_ms : float, optional
        The baseline latency in milliseconds to simulate, defaults to 50.0.
    latency_std_dev_ms : float, optional
        The standard deviation for latency simulation in milliseconds,
        defaults to 5.0.

    Returns
    -------
    dict
        An initialized state dictionary containing simulation parameters,
        latency data, and retry counts.
    """
    logger.info("Creating initial latency state...")
    state = {
        "exchanges": exchanges,
        "simulation_params": {
            "base_latency_ms": base_latency_ms,
            "latency_std_dev_ms": latency_std_dev_ms,
            "packet_loss_rate": 0.05, # 5% packet loss
            "jitter_range_ms": (0.1, 2.0) # Min and Max jitter to add
        },
        "latency_data": {exchange: [] for exchange in exchanges},
        "retry_counts": {exchange: 0 for exchange in exchanges},
        "total_attempts": {exchange: 0 for exchange in exchanges},
        "simulated_rtt_history": deque(maxlen=1000) # For rolling window stats
    }
    logger.debug(f"Initial state created: {state}")
    return state

def simulate_network_request(
    state: dict,
    exchange_name: str,
    max_retries: int = 3
) -> tuple[float | None, int]:
    """
    Simulates a network request to an exchange, calculating RTT with jitter and packet loss.
    Includes exponential backoff for retries.

    Parameters
    ----------
    state : dict
        The current state dictionary.
    exchange_name : str
        The name of the exchange to simulate the request to.
    max_retries : int, optional
        The maximum number of retries for a request if packet loss occurs, defaults to 3.

    Returns
    -------
    tuple[float | None, int]
        - The simulated RTT in milliseconds, or None if the request fails after max_retries.
        - The number of attempts made for this request.
    """
    params = state["simulation_params"]
    base_latency = params["base_latency_ms"]
    std_dev = params["latency_std_dev_ms"]
    packet_loss_rate = params["packet_loss_rate"]
    jitter_min, jitter_max = params["jitter_range_ms"]

    attempts = 0
    while attempts <= max_retries:
        attempts += 1
        state["total_attempts"][exchange_name] += 1

        # Simulate packet loss
        if random.random() < packet_loss_rate:
            logger.warning(f"Packet loss simulated for {exchange_name} (attempt {attempts}/{max_retries + 1})")
            if attempts <= max_retries:
                # Exponential backoff
                wait_time = (2 ** (attempts - 1)) + random.uniform(0, 0.1) # Add random jitter to backoff
                logger.info(f"Retrying {exchange_name} in {wait_time:.2f} seconds...")
                time.sleep(wait_time)
                state["retry_counts"][exchange_name] += 1
            continue

        # Simulate base latency with normal distribution
        simulated_rtt = np.random.normal(base_latency, std_dev)
        simulated_rtt = max(0.1, simulated_rtt) # Ensure RTT is not negative

        # Add random jitter
        jitter = random.uniform(jitter_min, jitter_max)
        simulated_rtt += jitter

        logger.debug(f"Simulated RTT for {exchange_name}: {simulated_rtt:.2f}ms (attempt {attempts})")
        state["simulated_rtt_history"].append(simulated_rtt)
        return simulated_rtt, attempts

    logger.error(f"Request to {exchange_name} failed after {max_retries + 1} attempts due to persistent packet loss.")
    return None, attempts

def run_latency_simulation(
    state: dict,
    iterations: int,
    request_delay_ms: float = 100.0
) -> dict:
    """
    Runs the network latency simulation for a given number of iterations.

    Parameters
    ----------
    state : dict
        The current state dictionary.
    iterations : int
        The total number of simulation iterations to perform.
    request_delay_ms : float, optional
        The delay between successive requests in milliseconds, defaults to 100.0.

    Returns
    -------
    dict
        The updated state dictionary with collected latency data.
    """
    logger.info(f"Starting latency simulation for {iterations} iterations...")
    exchanges = state["exchanges"]

    for i in range(iterations):
        logger.debug(f"Simulation iteration {i+1}/{iterations}")
        for exchange in exchanges:
            rtt, attempts = simulate_network_request(state, exchange)
            if rtt is not None:
                state["latency_data"][exchange].append(rtt)
            else:
                logger.warning(f"Request to {exchange} failed after {attempts} attempts in iteration {i+1}.")

            # Removed time.sleep here to allow faster simulation for demonstration purposes.
            # If real-time simulation with delays is required, this can be re-added or adjusted.

    logger.info("Latency simulation completed.")
    return state

def summarize_latency_metrics(state: dict) -> pd.DataFrame:
    """
    Summarizes the collected latency metrics for each exchange.

    Parameters
    ----------
    state : dict
        The current state dictionary with collected latency data and retry counts.

    Returns
    -------
    pd.DataFrame
        A Pandas DataFrame summarizing the key latency metrics for each exchange.
    """
    logger.info("Summarizing latency metrics...")
    summary_data = []

    for exchange, latencies in state["latency_data"].items():
        total_attempts = state["total_attempts"].get(exchange, 0)
        successful_pings = len(latencies)
        packet_loss_count = state["retry_counts"].get(exchange, 0)
        # If packet_loss_count is the only thing we're counting for failed attempts,
        # then total_attempts should be successful_pings + actual_failed_attempts.
        # For now, let's derive packet_loss_rate from successful pings vs total attempts.
        # Note: If simulate_network_request can fail completely, total_attempts might be higher.
        # Let's use a more robust way to calculate packet loss based on what was *not* successful.

        if total_attempts > 0:
            actual_packet_loss_rate = 1 - (successful_pings / total_attempts)
        else:
            actual_packet_loss_rate = 0.0

        if latencies:
            summary = {
                "Exchange": exchange,
                "Mean_RTT_ms": np.mean(latencies),
                "Median_RTT_ms": np.median(latencies),
                "StdDev_RTT_ms": np.std(latencies),
                "Min_RTT_ms": np.min(latencies),
                "Max_RTT_ms": np.max(latencies),
                "Successful_Pings": successful_pings,
                "Total_Attempts": total_attempts,
                "Packet_Loss_Rate": actual_packet_loss_rate * 100, # as percentage
                "Total_Retries": state["retry_counts"].get(exchange, 0)
            }
        else:
            # Handle cases where no successful pings occurred for an exchange
            summary = {
                "Exchange": exchange,
                "Mean_RTT_ms": np.nan,
                "Median_RTT_ms": np.nan,
                "StdDev_RTT_ms": np.nan,
                "Min_RTT_ms": np.nan,
                "Max_RTT_ms": np.nan,
                "Successful_Pings": successful_pings,
                "Total_Attempts": total_attempts,
                "Packet_Loss_Rate": actual_packet_loss_rate * 100,
                "Total_Retries": state["retry_counts"].get(exchange, 0)
            }
        summary_data.append(summary)

    df_summary = pd.DataFrame(summary_data)
    logger.info("Latency metrics summarized.")
    return df_summary

def calculate_rolling_metrics(state: dict, window_size: int = 100) -> dict:
    """
    Calculates rolling window statistics from the simulated RTT history.

    Parameters
    ----------
    state : dict
        The current state dictionary.
    window_size : int, optional
        The number of recent RTT samples to include in the rolling window, defaults to 100.

    Returns
    -------
    dict
        A dictionary containing the rolling mean, std dev, min, and max RTT
        for the current window. Returns NaNs if window_size is greater
        than available data.
    """
    logger.debug(f"Calculating rolling metrics for window size: {window_size}")
    history = list(state["simulated_rtt_history"])

    if len(history) < window_size:
        logger.warning(f"Not enough data for rolling window of size {window_size}. Available: {len(history)}")
        return {
            "rolling_mean": np.nan,
            "rolling_std": np.nan,
            "rolling_min": np.nan,
            "rolling_max": np.nan
        }

    recent_rtts = history[-window_size:]

    metrics = {
        "rolling_mean": np.mean(recent_rtts),
        "rolling_std": np.std(recent_rtts),
        "rolling_min": np.min(recent_rtts),
        "rolling_max": np.max(recent_rtts)
    }
    logger.debug(f"Rolling metrics: {metrics}")
    return metrics

# 1. Initialize State
EXCHANGES = ["Exchange A", "Exchange B", "Exchange C"]
initial_state = create_latency_state(exchanges=EXCHANGES, base_latency_ms=60.0, latency_std_dev_ms=8.0)

# 2. Run Simulation
SIMULATION_ITERATIONS = 50
SIMULATION_DELAY_MS = 50.0
logger.info(f"Running a basic simulation for {SIMULATION_ITERATIONS} iterations with {SIMULATION_DELAY_MS}ms delay...")
updated_state = run_latency_simulation(initial_state, SIMULATION_ITERATIONS, SIMULATION_DELAY_MS)

# 3. Summarize Metrics
latency_summary_df = summarize_latency_metrics(updated_state)

print("\n--- Latency Simulation Summary ---")
display(latency_summary_df)

# 1. Initialize State
EXCHANGES = ["Exchange A", "Exchange B", "Exchange C"]
initial_state = create_latency_state(exchanges=EXCHANGES, base_latency_ms=60.0, latency_std_dev_ms=8.0)

# 2. Run Simulation
SIMULATION_ITERATIONS = 50
SIMULATION_DELAY_MS = 50.0
logger.info(f"Running a basic simulation for {SIMULATION_ITERATIONS} iterations with {SIMULATION_DELAY_MS}ms delay...")
updated_state = run_latency_simulation(initial_state, SIMULATION_ITERATIONS, SIMULATION_DELAY_MS)

# 3. Summarize Metrics (Optional, as the plot directly uses updated_state)
# latency_summary_df = summarize_latency_metrics(updated_state)

# print("\n--- Latency Simulation Summary ---")
# display(latency_summary_df)

plt.figure(figsize=(12, 6))
sns.histplot(pd.DataFrame({"RTT_ms": updated_state["simulated_rtt_history"]}), x="RTT_ms", kde=True, stat="density", linewidth=0)
plt.title('Overall Simulated RTT Distribution')
plt.xlabel('Round Trip Time (ms)')
plt.ylabel('Density')
plt.grid(True, linestyle='--', alpha=0.6)
plt.show()

plt.figure(figsize=(12, 6))
for exchange in EXCHANGES:
    if updated_state["latency_data"][exchange]:
        sns.kdeplot(updated_state["latency_data"][exchange], label=f'{exchange}', fill=True, alpha=0.5)
plt.title('Latency Distribution per Exchange')
plt.xlabel('Round Trip Time (ms)')
plt.ylabel('Density')
plt.legend(title='Exchange')
plt.grid(True, linestyle='--', alpha=0.6)
plt.show()

# Box plot for better comparison
all_latencies = []
for exchange, latencies in updated_state["latency_data"].items():
    for rtt in latencies:
        all_latencies.append({"Exchange": exchange, "RTT_ms": rtt})

if all_latencies:
    df_all_latencies = pd.DataFrame(all_latencies)
    plt.figure(figsize=(12, 6))
    sns.boxplot(x='Exchange', y='RTT_ms', data=df_all_latencies)
    plt.title('Latency Distribution (Box Plot) per Exchange')
    plt.xlabel('Exchange')
    plt.ylabel('Round Trip Time (ms)')
    plt.grid(axis='y', linestyle='--', alpha=0.6)
    plt.show()
else:
    logger.warning("No latency data available to create box plots.")

rolling_window_size = 50 # Example window size
rolling_metrics_history = []

# Re-run simulation but collect rolling metrics at each step
# Reset state to ensure clean run for rolling metrics collection
EXCHANGES = ["Exchange A", "Exchange B", "Exchange C"]
initial_state_rolling = create_latency_state(exchanges=EXCHANGES, base_latency_ms=60.0, latency_std_dev_ms=8.0)

SIMULATION_ITERATIONS_ROLLING = 50
for i in range(SIMULATION_ITERATIONS_ROLLING):
    for exchange in EXCHANGES:
        rtt, attempts = simulate_network_request(initial_state_rolling, exchange)
        if rtt is not None:
            initial_state_rolling["latency_data"][exchange].append(rtt)

    # Calculate rolling metrics after each full iteration across exchanges
    if len(initial_state_rolling["simulated_rtt_history"]) >= rolling_window_size:
        metrics = calculate_rolling_metrics(initial_state_rolling, rolling_window_size)
        rolling_metrics_history.append(metrics)
    else:
        # Append NaNs if not enough data for rolling window yet
        rolling_metrics_history.append({
            "rolling_mean": np.nan,
            "rolling_std": np.nan,
            "rolling_min": np.nan,
            "rolling_max": np.nan
        })

df_rolling_metrics = pd.DataFrame(rolling_metrics_history).dropna()

if not df_rolling_metrics.empty:
    plt.figure(figsize=(15, 7))

    plt.plot(df_rolling_metrics.index, df_rolling_metrics['rolling_mean'], label='Rolling Mean RTT')
    plt.fill_between(df_rolling_metrics.index,
                     df_rolling_metrics['rolling_mean'] - df_rolling_metrics['rolling_std'],
                     df_rolling_metrics['rolling_mean'] + df_rolling_metrics['rolling_std'],
                     color='blue', alpha=0.1, label='Rolling Std Dev (Mean ± Std)')
    plt.plot(df_rolling_metrics.index, df_rolling_metrics['rolling_min'], label='Rolling Min RTT', linestyle='--')
    plt.plot(df_rolling_metrics.index, df_rolling_metrics['rolling_max'], label='Rolling Max RTT', linestyle='--')

    plt.title(f'Rolling Latency Metrics (Window Size: {rolling_window_size})')
    plt.xlabel('Simulation Step')
    plt.ylabel('Round Trip Time (ms)')
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.show()
else:
    logger.warning("Not enough data to plot rolling metrics or all values are NaN.")