pip install pandas numpy matplotlib seaborn loguru tenacity

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import datetime
from collections import deque
import time
import random
from typing import Dict, Any, List, Optional
from loguru import logger
from tenacity import retry, wait_exponential, stop_after_attempt, before_sleep_log
import sys

# Configure loguru logger
logger.remove()
logger.add(sys.stdout, format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>")

def create_pnl_state(start_date: str, end_date: str, exchanges: List[str]) -> Dict[str, Any]:
    """
    Initializes the state dictionary for PnL reconciliation.

    Parameters
    ----------
    start_date : str
        The start date for the reconciliation period in 'YYYY-MM-DD' format.
    end_date : str
        The end date for the reconciliation period in 'YYYY-MM-DD' format.
    exchanges : List[str]
        A list of exchange names to be reconciled.

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

    Examples
    --------
    >>> state = create_pnl_state('2023-01-01', '2023-01-07', ['ExchangeA', 'ExchangeB'])
    >>> 'reconciliation_period' in state
    True
    """
    logger.info(f"Initializing PnL state for {exchanges} from {start_date} to {end_date}")
    state = {
        'reconciliation_period': {
            'start_date': pd.to_datetime(start_date),
            'end_date': pd.to_datetime(end_date)
        },
        'exchanges': exchanges,
        'trade_data': {exchange: pd.DataFrame() for exchange in exchanges},
        'pnl_data': {exchange: pd.DataFrame() for exchange in exchanges},
        'reconciliation_report': pd.DataFrame(),
        'metrics': {},
        'config': {
            'trade_columns': ['timestamp', 'symbol', 'side', 'price', 'quantity', 'fees', 'exchange']
        }
    }
    logger.debug("PnL state initialized successfully.")
    return state

@retry(wait=wait_exponential(multiplier=1, min=4, max=10), stop=stop_after_attempt(3), before_sleep=before_sleep_log(logger, 'warning'))
def simulate_trade_data(state: Dict[str, Any], exchange_name: str, num_trades: int, symbols: List[str]) -> Dict[str, Any]:
    """
    Simulates realistic trade data for a given exchange.

    Parameters
    ----------
    state : Dict[str, Any]
        The current state dictionary.
    exchange_name : str
        The name of the exchange for which to generate data.
    num_trades : int
        The number of trades to simulate.
    symbols : List[str]
        A list of trading symbols (e.g., 'BTC/USD').

    Returns
    -------
    Dict[str, Any]
        The updated state dictionary with simulated trade data.

    Examples
    --------
    >>> state = create_pnl_state('2023-01-01', '2023-01-07', ['ExchangeA'])
    >>> updated_state = simulate_trade_data(state, 'ExchangeA', 10, ['BTC/USD'])
    >>> not updated_state['trade_data']['ExchangeA'].empty
    True
    """
    logger.info(f"Simulating {num_trades} trades for {exchange_name}.")
    try:
        start_date = state['reconciliation_period']['start_date']
        end_date = state['reconciliation_period']['end_date']
        time_range = (end_date - start_date).days * 24 * 60 * 60 # total seconds

        data = []
        for _ in range(num_trades):
            timestamp_offset = random.uniform(0, time_range)
            timestamp = start_date + datetime.timedelta(seconds=timestamp_offset)
            symbol = random.choice(symbols)
            side = random.choice(['BUY', 'SELL'])
            price = round(random.uniform(1000.0, 50000.0), 2)
            quantity = round(random.uniform(0.01, 5.0), 4)
            fees = round(price * quantity * random.uniform(0.0001, 0.001), 2) # 0.01% to 0.1% fee
            data.append({
                'timestamp': timestamp,
                'symbol': symbol,
                'side': side,
                'price': price,
                'quantity': quantity,
                'fees': fees,
                'exchange': exchange_name
            })

        df = pd.DataFrame(data)
        df['timestamp'] = pd.to_datetime(df['timestamp'])
        df = df.sort_values(by='timestamp').reset_index(drop=True)

        # Add random jitter to timestamps to simulate real-world slight differences
        if random.random() > 0.5: # Apply jitter sometimes
            df['timestamp'] = df['timestamp'] + pd.to_timedelta(np.random.normal(0, 10, len(df)), unit='ms')

        state['trade_data'][exchange_name] = df
        logger.debug(f"Successfully simulated trade data for {exchange_name}. Trades generated: {len(df)}")

        # Introduce a random delay to simulate network latency for retry mechanism
        time.sleep(random.uniform(0.1, 0.5) + random.uniform(0, 0.2)) # Base delay + jitter

    except Exception as e:
        logger.error(f"Error simulating trade data for {exchange_name}: {e}")
        raise # Re-raise to trigger retry
    return state

def calculate_pnl_single_exchange(state: Dict[str, Any], exchange_name: str) -> Dict[str, Any]:
    """
    Calculates the PnL for a single exchange based on its trade data.

    Parameters
    ----------
    state : Dict[str, Any]
        The current state dictionary.
    exchange_name : str
        The name of the exchange for which to calculate PnL.

    Returns
    -------
    Dict[str, Any]
        The updated state dictionary with calculated PnL data.

    Examples
    --------
    >>> state = create_pnl_state('2023-01-01', '2023-01-07', ['ExchangeA'])
    >>> state = simulate_trade_data(state, 'ExchangeA', 10, ['BTC/USD'])
    >>> updated_state = calculate_pnl_single_exchange(state, 'ExchangeA')
    >>> 'daily_pnl' in updated_state['pnl_data']['ExchangeA'].columns
    True
    """
    logger.info(f"Calculating PnL for {exchange_name}.")
    trade_df = state['trade_data'][exchange_name].copy()

    if trade_df.empty:
        logger.warning(f"No trade data found for {exchange_name}. PnL will be empty.")
        state['pnl_data'][exchange_name] = pd.DataFrame(columns=['date', 'symbol', 'realized_pnl', 'total_fees'])
        return state

    pnl_records = []
    # Group by symbol to calculate PnL for each asset
    for symbol, group in trade_df.groupby('symbol'):
        holdings_cost_basis = deque() # (quantity, cost_price)
        realized_pnl = 0.0
        total_fees = 0.0

        for _, trade in group.iterrows():
            trade_value = trade['price'] * trade['quantity']
            total_fees += trade['fees']

            if trade['side'] == 'BUY':
                holdings_cost_basis.append((trade['quantity'], trade['price']))
            elif trade['side'] == 'SELL':
                quantity_to_sell = trade['quantity']
                while quantity_to_sell > 0 and holdings_cost_basis:
                    held_quantity, held_cost = holdings_cost_basis[0]
                    if held_quantity <= quantity_to_sell:
                        realized_pnl += (trade['price'] - held_cost) * held_quantity
                        quantity_to_sell -= held_quantity
                        holdings_cost_basis.popleft()
                    else:
                        realized_pnl += (trade['price'] - held_cost) * quantity_to_sell
                        holdings_cost_basis[0] = (held_quantity - quantity_to_sell, held_cost)
                        quantity_to_sell = 0
            # For simplicity, we ignore unrealized PnL and focus on realized for reconciliation

        pnl_records.append({
            'date': trade_df['timestamp'].max().normalize(), # Using max trade date for simplicity
            'symbol': symbol,
            'realized_pnl': realized_pnl - total_fees,
            'total_fees': total_fees
        })

    pnl_df = pd.DataFrame(pnl_records)
    if not pnl_df.empty:
        pnl_df['date'] = pd.to_datetime(pnl_df['date'])
    state['pnl_data'][exchange_name] = pnl_df
    logger.debug(f"PnL calculated for {exchange_name}. Total PnL: {pnl_df['realized_pnl'].sum() if not pnl_df.empty else 0:.2f}")
    return state

def reconcile_pnl(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Reconciles PnL data across multiple exchanges.

    Parameters
    ----------
    state : Dict[str, Any]
        The current state dictionary.

    Returns
    -------
    Dict[str, Any]
        The updated state dictionary with the reconciliation report.

    Examples
    --------
    >>> state = create_pnl_state('2023-01-01', '2023-01-07', ['ExchangeA', 'ExchangeB'])
    >>> state = simulate_trade_data(state, 'ExchangeA', 10, ['BTC/USD'])
    >>> state = simulate_trade_data(state, 'ExchangeB', 12, ['BTC/USD'])
    >>> state = calculate_pnl_single_exchange(state, 'ExchangeA')
    >>> state = calculate_pnl_single_exchange(state, 'ExchangeB')
    >>> updated_state = reconcile_pnl(state)
    >>> 'discrepancy' in updated_state['reconciliation_report'].columns
    True
    """
    logger.info("Starting PnL reconciliation across exchanges.")
    combined_pnl = pd.DataFrame()

    # Consolidate PnL data from all exchanges
    for exchange_name in state['exchanges']:
        pnl_df = state['pnl_data'][exchange_name].copy()
        if not pnl_df.empty:
            pnl_df['exchange'] = exchange_name
            combined_pnl = pd.concat([combined_pnl, pnl_df])

    if combined_pnl.empty:
        logger.warning("No PnL data to reconcile. Reconciliation report will be empty.")
        state['reconciliation_report'] = pd.DataFrame(columns=['date', 'symbol', 'exchange_a_pnl', 'exchange_b_pnl', 'discrepancy'])
        return state

    # Pivot to compare PnL side-by-side
    pivot_columns = ['date', 'symbol']
    if len(state['exchanges']) < 2:
        logger.warning("Cannot reconcile PnL with less than two exchanges. Report will show individual PnL.")
        state['reconciliation_report'] = combined_pnl
        return state

    first_exchange = state['exchanges'][0]
    second_exchange = state['exchanges'][1]

    # Use a more robust merging approach for multiple exchanges
    reconciliation_df = pd.DataFrame()
    if not combined_pnl.empty:
        # Ensure 'date' and 'symbol' are correct types for merging
        combined_pnl['date'] = pd.to_datetime(combined_pnl['date']).dt.normalize()

        reconciliation_df = combined_pnl.pivot_table(
            index=['date', 'symbol'],
            columns='exchange',
            values='realized_pnl'
        ).reset_index()
        reconciliation_df.columns.name = None # Remove columns name 'exchange'

        # Fill NaNs with 0 where an exchange might not have trades for a symbol/date
        for ex in state['exchanges']:
            if ex not in reconciliation_df.columns:
                 reconciliation_df[ex] = 0.0 # Add missing exchange columns if needed
            reconciliation_df[ex] = reconciliation_df[ex].fillna(0.0)

        # Calculate discrepancy for the first two exchanges as a primary check
        if first_exchange in reconciliation_df.columns and second_exchange in reconciliation_df.columns:
            reconciliation_df['discrepancy'] = reconciliation_df[first_exchange] - reconciliation_df[second_exchange]
            reconciliation_df['percentage_discrepancy'] = np.where(
                reconciliation_df[first_exchange] != 0,
                (reconciliation_df['discrepancy'] / reconciliation_df[first_exchange]) * 100,
                np.where(
                    reconciliation_df[second_exchange] != 0,
                    (reconciliation_df['discrepancy'] / reconciliation_df[second_exchange]) * 100,
                    0.0
                )
            )
        else:
            logger.warning(f"Could not find both {first_exchange} and {second_exchange} in pivoted PnL for discrepancy calculation.")
            reconciliation_df['discrepancy'] = 0.0
            reconciliation_df['percentage_discrepancy'] = 0.0

    state['reconciliation_report'] = reconciliation_df
    logger.debug(f"PnL reconciliation completed. Total discrepancies found: {len(reconciliation_df[reconciliation_df['discrepancy'].abs() > 0])}")
    return state

def generate_reconciliation_report_summary(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Generates a summary of the PnL reconciliation report.

    Parameters
    ----------
    state : Dict[str, Any]
        The current state dictionary.

    Returns
    -------
    Dict[str, Any]
        The updated state dictionary with summary metrics.

    Examples
    --------
    >>> state = create_pnl_state('2023-01-01', '2023-01-07', ['ExchangeA', 'ExchangeB'])
    >>> state = simulate_trade_data(state, 'ExchangeA', 10, ['BTC/USD'])
    >>> state = simulate_trade_data(state, 'ExchangeB', 12, ['BTC/USD'])
    >>> state = calculate_pnl_single_exchange(state, 'ExchangeA')
    >>> state = calculate_pnl_single_exchange(state, 'ExchangeB')
    >>> state = reconcile_pnl(state)
    >>> updated_state = generate_reconciliation_report_summary(state)
    >>> 'total_pnl_exchangeA' in updated_state['metrics']
    True
    """
    logger.info("Generating reconciliation report summary.")
    report_df = state['reconciliation_report']
    metrics = {}

    if report_df.empty:
        logger.warning("Reconciliation report is empty. No summary metrics generated.")
        return state

    first_exchange = state['exchanges'][0]
    second_exchange = state['exchanges'][1]

    for ex in state['exchanges']:
        if ex in report_df.columns:
            metrics[f'total_pnl_{ex}'] = report_df[ex].sum()

    if 'discrepancy' in report_df.columns:
        metrics['total_absolute_discrepancy'] = report_df['discrepancy'].abs().sum()
        metrics['max_absolute_discrepancy'] = report_df['discrepancy'].abs().max()
        metrics['average_percentage_discrepancy'] = report_df['percentage_discrepancy'].abs().mean()
        metrics['count_of_discrepancies'] = len(report_df[report_df['discrepancy'].abs() > 0.001]) # Threshold for non-zero

    state['metrics'] = metrics
    logger.debug("Reconciliation report summary generated.")
    return state

def identify_discrepancy_sources(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Identifies potential sources of discrepancies by comparing trade data details.

    Parameters
    ----------
    state : Dict[str, Any]
        The current state dictionary.

    Returns
    -------
    Dict[str, Any]
        The updated state dictionary with insights into discrepancy sources.

    Examples
    --------
    >>> state = create_pnl_state('2023-01-01', '2023-01-07', ['ExchangeA', 'ExchangeB'])
    >>> state = simulate_trade_data(state, 'ExchangeA', 10, ['BTC/USD'])
    >>> state = simulate_trade_data(state, 'ExchangeB', 12, ['BTC/USD'])
    >>> state = calculate_pnl_single_exchange(state, 'ExchangeA')
    >>> state = calculate_pnl_single_exchange(state, 'ExchangeB')
    >>> state = reconcile_pnl(state)
    >>> updated_state = identify_discrepancy_sources(state)
    >>> 'trade_volume_comparison' in updated_state['metrics']
    True
    """
    logger.info("Attempting to identify discrepancy sources.")
    trade_data_comparisons = {}

    if len(state['exchanges']) < 2:
        logger.warning("Insufficient exchanges to compare trade data for discrepancy sources.")
        return state

    all_trades = pd.DataFrame()
    for ex in state['exchanges']:
        df = state['trade_data'][ex].copy()
        if not df.empty:
            df['date'] = df['timestamp'].dt.normalize()
            all_trades = pd.concat([all_trades, df])

    if all_trades.empty:
        logger.warning("No trade data available to identify discrepancy sources.")
        return state

    # Compare trade counts and volumes per symbol and date
    trade_summary = all_trades.groupby(['date', 'symbol', 'exchange']).agg(
        trade_count=('timestamp', 'count'),
        total_quantity=('quantity', 'sum'),
        total_value=('price', lambda x: (x * all_trades.loc[x.index, 'quantity']).sum()),
        total_fees=('fees', 'sum')
    ).reset_index()

    pivot_trade_summary = trade_summary.pivot_table(
        index=['date', 'symbol'],
        columns='exchange',
        values=['trade_count', 'total_quantity', 'total_value', 'total_fees']
    ).reset_index()

    pivot_trade_summary.columns = [f'{col[0]}_{col[1]}' if col[1] else col[0] for col in pivot_trade_summary.columns.values]
    pivot_trade_summary = pivot_trade_summary.fillna(0)

    first_ex = state['exchanges'][0]
    second_ex = state['exchanges'][1]

    # Calculate differences for key metrics
    metrics_to_compare = ['trade_count', 'total_quantity', 'total_value', 'total_fees']
    for metric in metrics_to_compare:
        col1 = f'{metric}_{first_ex}'
        col2 = f'{metric}_{second_ex}'
        if col1 in pivot_trade_summary.columns and col2 in pivot_trade_summary.columns:
            pivot_trade_summary[f'{metric}_diff'] = pivot_trade_summary[col1] - pivot_trade_summary[col2]
            pivot_trade_summary[f'{metric}_abs_diff'] = pivot_trade_summary[f'{metric}_diff'].abs()
            pivot_trade_summary[f'{metric}_pct_diff'] = np.where(
                pivot_trade_summary[col1] != 0,
                (pivot_trade_summary[f'{metric}_diff'] / pivot_trade_summary[col1]) * 100,
                np.where(
                    pivot_trade_summary[col2] != 0,
                    (pivot_trade_summary[f'{metric}_diff'] / pivot_trade_summary[col2]) * 100,
                    0.0
                )
            )

    state['metrics']['trade_volume_comparison'] = pivot_trade_summary
    logger.debug("Discrepancy source identification completed.")
    return state

# 5.1. Initialize State
logger.info("DEMO: Initializing state.")
start_date_str = '2023-01-01'
end_date_str = '2023-01-07'
exchanges_to_reconcile = ['Exchange_Alpha', 'Exchange_Beta']

pnl_reconciliation_state = create_pnl_state(start_date_str, end_date_str, exchanges_to_reconcile)
display(pnl_reconciliation_state['reconciliation_period'])

# 5.2. Simulate Trade Data for Each Exchange
logger.info("DEMO: Simulating trade data.")
symbols_traded = ['BTC/USD', 'ETH/USD', 'LTC/USD']

# Simulate for Exchange Alpha
pnl_reconciliation_state = simulate_trade_data(pnl_reconciliation_state, 'Exchange_Alpha', 500, symbols_traded)

# Simulate for Exchange Beta (with some intentional differences for discrepancies)
pnl_reconciliation_state = simulate_trade_data(pnl_reconciliation_state, 'Exchange_Beta', 520, symbols_traded)

# Introduce an intentional discrepancy in one of the exchanges
if not pnl_reconciliation_state['trade_data']['Exchange_Beta'].empty:
    # Change price of a few trades to create PnL difference
    num_discrepant_trades = min(5, len(pnl_reconciliation_state['trade_data']['Exchange_Beta']))
    indices_to_modify = pnl_reconciliation_state['trade_data']['Exchange_Beta'].sample(n=num_discrepant_trades).index
    pnl_reconciliation_state['trade_data']['Exchange_Beta'].loc[indices_to_modify, 'price'] *= (1 + np.random.uniform(-0.01, 0.01, num_discrepant_trades)) # +/- 1% difference
    logger.warning(f"Introduced intentional PnL discrepancy for {num_discrepant_trades} trades in Exchange_Beta.")

print("\n--- Sample Trade Data for Exchange_Alpha ---")
display(pnl_reconciliation_state['trade_data']['Exchange_Alpha'].head())
print("\n--- Sample Trade Data for Exchange_Beta ---")
display(pnl_reconciliation_state['trade_data']['Exchange_Beta'].head())

# 5.3. Calculate PnL for Each Exchange
logger.info("DEMO: Calculating PnL.")
for ex in exchanges_to_reconcile:
    pnl_reconciliation_state = calculate_pnl_single_exchange(pnl_reconciliation_state, ex)

print("\n--- PnL Data for Exchange_Alpha ---")
display(pnl_reconciliation_state['pnl_data']['Exchange_Alpha'])
print("\n--- PnL Data for Exchange_Beta ---")
display(pnl_reconciliation_state['pnl_data']['Exchange_Beta'])

# 5.4. Reconcile PnL and Generate Report
logger.info("DEMO: Reconciling PnL.")
pnl_reconciliation_state = reconcile_pnl(pnl_reconciliation_state)

print("\n--- PnL Reconciliation Report ---")
display(pnl_reconciliation_state['reconciliation_report'])

# Filter and display discrepancies
discrepancies_df = pnl_reconciliation_state['reconciliation_report']
if 'discrepancy' in discrepancies_df.columns:
    significant_discrepancies = discrepancies_df[discrepancies_df['discrepancy'].abs() > 0.01] # Define a threshold
    print("\n--- Significant PnL Discrepancies (Absolute value > 0.01) ---")
    display(significant_discrepancies)
else:
    print("No 'discrepancy' column found in the reconciliation report.")

# 5.5. Generate Reconciliation Report Summary
logger.info("DEMO: Generating summary metrics.")
pnl_reconciliation_state = generate_reconciliation_report_summary(pnl_reconciliation_state)

print("\n--- Reconciliation Summary Metrics ---")
for key, value in pnl_reconciliation_state['metrics'].items():
    if isinstance(value, (int, float)):
        print(f"{key}: {value:.2f}")
    else:
        print(f"{key}: {value}")

# 5.6. Visualize PnL and Discrepancies
logger.info("DEMO: Visualizing PnL and discrepancies.")

report_df = pnl_reconciliation_state['reconciliation_report']

if not report_df.empty and len(exchanges_to_reconcile) >= 2:
    first_ex = exchanges_to_reconcile[0]
    second_ex = exchanges_to_reconcile[1]

    # Plot 1: Daily PnL Comparison
    plt.figure(figsize=(14, 7))
    plt.plot(report_df['date'], report_df[first_ex], label=f'{first_ex} PnL', marker='o')
    plt.plot(report_df['date'], report_df[second_ex], label=f'{second_ex} PnL', marker='x')
    plt.title('Daily PnL Comparison Across Exchanges')
    plt.xlabel('Date')
    plt.ylabel('Realized PnL')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()

    # Plot 2: Daily Absolute Discrepancy
    if 'discrepancy' in report_df.columns:
        plt.figure(figsize=(14, 7))
        plt.bar(report_df['date'], report_df['discrepancy'].abs(), color='salmon')
        plt.title('Daily Absolute PnL Discrepancy')
        plt.xlabel('Date')
        plt.ylabel('Absolute Discrepancy')
        plt.grid(True)
        plt.tight_layout()
        plt.show()

    # Plot 3: Discrepancy Distribution
    if 'discrepancy' in report_df.columns and not report_df['discrepancy'].empty:
        plt.figure(figsize=(10, 6))
        sns.histplot(report_df['discrepancy'], bins=30, kde=True)
        plt.title('Distribution of PnL Discrepancies')
        plt.xlabel('Discrepancy Amount')
        plt.ylabel('Frequency')
        plt.grid(True)
        plt.tight_layout()
        plt.show()
else:
    logger.warning("Not enough data or exchanges for PnL comparison plots.")

# 5.7. Identify Discrepancy Sources Visualization
logger.info("DEMO: Visualizing potential discrepancy sources.")

pnl_reconciliation_state = identify_discrepancy_sources(pnl_reconciliation_state)
trade_comp_df = pnl_reconciliation_state['metrics'].get('trade_volume_comparison')

if trade_comp_df is not None and not trade_comp_df.empty:
    print("\n--- Trade Volume Comparison (potential discrepancy sources) ---")
    display(trade_comp_df.head())

    # Visualize total quantity difference by symbol
    if 'total_quantity_diff' in trade_comp_df.columns:
        plt.figure(figsize=(12, 6))
        sns.barplot(x='symbol', y='total_quantity_diff', data=trade_comp_df.groupby('symbol')['total_quantity_diff'].sum().reset_index())
        plt.title('Total Quantity Difference by Symbol Across Exchanges')
        plt.xlabel('Symbol')
        plt.ylabel('Quantity Difference')
        plt.xticks(rotation=45)
        plt.grid(axis='y')
        plt.tight_layout()
        plt.show()

    # Visualize total fees difference by symbol
    if 'total_fees_diff' in trade_comp_df.columns:
        plt.figure(figsize=(12, 6))
        sns.barplot(x='symbol', y='total_fees_diff', data=trade_comp_df.groupby('symbol')['total_fees_diff'].sum().reset_index())
        plt.title('Total Fees Difference by Symbol Across Exchanges')
        plt.xlabel('Symbol')
        plt.ylabel('Fees Difference')
        plt.xticks(rotation=45)
        plt.grid(axis='y')
        plt.tight_layout()
        plt.show()
else:
    logger.warning("No trade volume comparison data to visualize discrepancy sources.")