# Install necessary libraries
!pip install pandas matplotlib seaborn

import logging
import time
import random
from collections import deque, defaultdict
from typing import Dict, Any, List, Tuple

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__)

def create_position_state(initial_positions: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Initializes a dictionary to hold the local position state.

    Parameters
    ----------
    initial_positions : List[Dict[str, Any]]
        A list of dictionaries, each containing 'asset', 'quantity', and 'avg_price'.

    Returns
    -------
    Dict[str, Any]
        A dictionary representing the initial local position state.
        Example: {'BTC': {'quantity': 0.5, 'avg_price': 30000.0}}
    """
    state = {}
    for pos in initial_positions:
        asset = pos.get('asset')
        quantity = pos.get('quantity', 0.0)
        avg_price = pos.get('avg_price', 0.0)
        if asset:
            state[asset] = {'quantity': quantity, 'avg_price': avg_price}
            logger.info(f"Initialized local position for {asset}: {quantity} @ {avg_price}")
        else:
            logger.warning("Skipping position initialization due to missing 'asset' key.")
    return state

def create_exchange_state(initial_balances: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Initializes a dictionary to hold the exchange balance state.

    Parameters
    ----------
    initial_balances : List[Dict[str, Any]]
        A list of dictionaries, each containing 'asset' and 'balance'.

    Returns
    -------
    Dict[str, Any]
        A dictionary representing the initial exchange balance state.
        Example: {'BTC': {'balance': 0.5}, 'USD': {'balance': 10000.0}}
    """
    state = {}
    for bal in initial_balances:
        asset = bal.get('asset')
        balance = bal.get('balance', 0.0)
        if asset:
            state[asset] = {'balance': balance}
            logger.info(f"Initialized exchange balance for {asset}: {balance}")
        else:
            logger.warning("Skipping balance initialization due to missing 'asset' key.")
    return state

def get_exchange_positions(exchange_state: Dict[str, Any], max_retries: int = 3, base_delay: float = 0.1) -> Dict[str, float]:
    """
    Simulates fetching current positions from an external exchange with retry logic.

    Parameters
    ----------
    exchange_state : Dict[str, Any]
        The current simulated state of the exchange.
    max_retries : int, optional
        The maximum number of retry attempts, defaults to 3.
    base_delay : float, optional
        The initial delay in seconds before retrying, defaults to 0.1.

    Returns
    -------
    Dict[str, float]
        A dictionary where keys are asset symbols and values are their quantities on the exchange.
        Returns an empty dictionary if fetching fails after retries.
    """
    for i in range(max_retries):
        try:
            # Simulate network latency or API call failure
            if random.random() < 0.2 and i < max_retries - 1: # 20% chance of failure, but not on last attempt
                raise ConnectionError("Simulated API connection error.")

            positions = {asset: data['balance'] for asset, data in exchange_state.items() if data['balance'] > 0}
            logger.info(f"Successfully fetched exchange positions on attempt {i+1}.")
            return positions
        except ConnectionError as e:
            delay = base_delay * (2 ** i) + random.uniform(0, 0.1) # Exponential backoff with jitter
            logger.warning(f"Attempt {i+1} failed: {e}. Retrying in {delay:.2f} seconds...")
            time.sleep(delay)
        except Exception as e:
            logger.error(f"An unexpected error occurred while fetching exchange positions: {e}")
            return {}

    logger.error("Failed to fetch exchange positions after multiple retries.")
    return {}

def get_local_positions(local_state: Dict[str, Any]) -> Dict[str, float]:
    """
    Simulates retrieving current local positions from an internal system.

    Parameters
    ----------
    local_state : Dict[str, Any]
        The current simulated state of the local system.

    Returns
    -------
    Dict[str, float]
        A dictionary where keys are asset symbols and values are their quantities in the local system.
    """
    positions = {asset: data['quantity'] for asset, data in local_state.items()}
    logger.info("Successfully fetched local positions.")
    return positions

def reconcile_positions(local_positions: Dict[str, float],
                        exchange_positions: Dict[str, float],
                        tolerance: float = 1e-6) -> List[Dict[str, Any]]:
    """
    Compares local and exchange positions to identify discrepancies.

    Parameters
    ----------
    local_positions : Dict[str, float]
        A dictionary of local asset positions.
    exchange_positions : Dict[str, float]
        A dictionary of exchange asset positions.
    tolerance : float, optional
        The allowable difference between quantities to consider them reconciled, defaults to 1e-6.

    Returns
    -------
    List[Dict[str, Any]]
        A list of dictionaries, each describing a discrepancy (asset, type, local_qty, exchange_qty, diff).
    """
    discrepancies = []
    all_assets = set(local_positions.keys()).union(set(exchange_positions.keys()))

    for asset in all_assets:
        local_qty = local_positions.get(asset, 0.0)
        exchange_qty = exchange_positions.get(asset, 0.0)

        if abs(local_qty - exchange_qty) > tolerance:
            diff = exchange_qty - local_qty
            discrepancy_type = ''
            if local_qty == 0.0 and exchange_qty > 0.0:
                discrepancy_type = 'missing_local' # Exchange has it, local doesn't
                logger.warning(f"Discrepancy for {asset}: Missing locally. Local: {local_qty}, Exchange: {exchange_qty}")
            elif exchange_qty == 0.0 and local_qty > 0.0:
                discrepancy_type = 'missing_exchange' # Local has it, exchange doesn't
                logger.warning(f"Discrepancy for {asset}: Missing on exchange. Local: {local_qty}, Exchange: {exchange_qty}")
            else:
                discrepancy_type = 'quantity_mismatch' # Quantities don't match
                logger.warning(f"Discrepancy for {asset}: Quantity mismatch. Local: {local_qty}, Exchange: {exchange_qty}")

            discrepancies.append({
                'asset': asset,
                'type': discrepancy_type,
                'local_quantity': local_qty,
                'exchange_quantity': exchange_qty,
                'difference': diff
            })
        else:
            logger.debug(f"Positions for {asset} are reconciled. Local: {local_qty}, Exchange: {exchange_qty}")

    if not discrepancies:
        logger.info("No discrepancies found. All positions reconciled.")
    return discrepancies

def adjust_position(exchange_state: Dict[str, Any],
                    asset: str,
                    action: str,
                    quantity: float,
                    price: float) -> Dict[str, Any]:
    """
    Simulates adjusting a position on an exchange.

    Parameters
    ----------
    exchange_state : Dict[str, Any]
        The current simulated exchange state to be updated.
    asset : str
        The asset symbol to adjust.
    action : str
        The type of adjustment ('buy', 'sell').
    quantity : float
        The quantity to adjust by.
    price : float
        The price at which the adjustment is made (for logging).

    Returns
    -------
    Dict[str, Any]
        The updated exchange state.
    """
    if asset not in exchange_state:
        exchange_state[asset] = {'balance': 0.0}
        logger.info(f"Added {asset} to exchange state for adjustment.")

    if action.lower() == 'buy':
        exchange_state[asset]['balance'] += quantity
        logger.info(f"Simulated BUY order for {quantity} {asset} at {price}. Exchange balance for {asset} is now {exchange_state[asset]['balance']}")
    elif action.lower() == 'sell':
        exchange_state[asset]['balance'] -= quantity
        logger.info(f"Simulated SELL order for {quantity} {asset} at {price}. Exchange balance for {asset} is now {exchange_state[asset]['balance']}")
    else:
        logger.error(f"Invalid adjustment action: {action}. Must be 'buy' or 'sell'.")

    return exchange_state

def track_reconciliation_metrics(metrics_state: Dict[str, Any],
                                 discrepancies: List[Dict[str, Any]],
                                 reconciliation_time: float,
                                 adjustments_made: int) -> Dict[str, Any]:
    """
    Tracks key metrics related to the reconciliation process.

    Parameters
    ----------
    metrics_state : Dict[str, Any]
        The dictionary holding the current reconciliation metrics.
    discrepancies : List[Dict[str, Any]]
        The list of discrepancies found in the current run.
    reconciliation_time : float
        The time taken for the reconciliation process.
    adjustments_made : int
        The number of adjustments performed.

    Returns
    -------
    Dict[str, Any]
        The updated metrics state.
    """
    if 'total_runs' not in metrics_state:
        metrics_state = defaultdict(int)

    metrics_state['total_runs'] += 1
    metrics_state['total_discrepancies_found'] += len(discrepancies)
    metrics_state['total_adjustments_made'] += adjustments_made
    metrics_state['total_reconciliation_time'] += reconciliation_time

    logger.info(f"Metrics updated: Run {metrics_state['total_runs']}, Discrepancies: {len(discrepancies)}, Adjustments: {adjustments_made}, Time: {reconciliation_time:.4f}s")
    return dict(metrics_state)

# 1. Initialize Simulated Data

# Local system state
initial_local_positions = [
    {'asset': 'BTC', 'quantity': 0.5, 'avg_price': 30000.0},
    {'asset': 'ETH', 'quantity': 2.0, 'avg_price': 1800.0},
    {'asset': 'XRP', 'quantity': 100.0, 'avg_price': 0.5}
]
local_state = create_position_state(initial_local_positions)

# Exchange state (with some initial discrepancies)
initial_exchange_balances = [
    {'asset': 'BTC', 'balance': 0.49},
    {'asset': 'ETH', 'balance': 2.1},
    {'asset': 'USD', 'balance': 10000.0}
]
exchange_state = create_exchange_state(initial_exchange_balances)

# Initialize metrics state
metrics_state = defaultdict(int)

# 2. Perform Initial Reconciliation

start_time = time.time()
local_positions = get_local_positions(local_state)
exchange_positions = get_exchange_positions(exchange_state)

discrepancies_round1 = reconcile_positions(local_positions, exchange_positions, tolerance=0.01)
reconciliation_time_round1 = time.time() - start_time

print("\n--- Initial Reconciliation Results ---")
if discrepancies_round1:
    df_discrepancies_round1 = pd.DataFrame(discrepancies_round1)
    display(df_discrepancies_round1)
else:
    print("No discrepancies found in the initial reconciliation.")

metrics_state = track_reconciliation_metrics(metrics_state, discrepancies_round1, reconciliation_time_round1, 0)

# 3. Visualize Initial Discrepancies

if 'df_discrepancies_round1' in locals():
    plt.figure(figsize=(10, 6))
    sns.barplot(x='asset', y='difference', hue='type', data=df_discrepancies_round1)
    plt.axhline(0, color='grey', linestyle='--')
    plt.title('Initial Position Discrepancies (Exchange Qty - Local Qty)')
    plt.xlabel('Asset')
    plt.ylabel('Quantity Difference')
    plt.legend(title='Discrepancy Type')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.show()
else:
    print("No discrepancies to visualize.")

# 4. Demonstrate Position Adjustment

adjustments_count = 0
for disc in discrepancies_round1:
    asset = disc['asset']
    diff = disc['difference'] # exchange_qty - local_qty

    if abs(diff) > 0.01: # Only adjust if difference is significant
        action = 'buy' if diff > 0 else 'sell'
        qty_to_adjust = abs(diff)
        current_price = local_state.get(asset, {}).get('avg_price', 1.0) # Use local avg price for simulation
        exchange_state = adjust_position(exchange_state, asset, action, qty_to_adjust, current_price)
        adjustments_count += 1
        # In a real system, you'd also update the local_state after a successful adjustment
        # For this demo, we'll see the effect on the *next* reconciliation run

metrics_state = track_reconciliation_metrics(metrics_state, discrepancies_round1, 0.0, adjustments_count)
print(f"\n--- Adjustments Made: {adjustments_count} ---")

# 5. Perform Second Reconciliation After Adjustments

start_time = time.time()
local_positions_after_adj = get_local_positions(local_state)
exchange_positions_after_adj = get_exchange_positions(exchange_state)

discrepancies_round2 = reconcile_positions(local_positions_after_adj, exchange_positions_after_adj, tolerance=0.01)
reconciliation_time_round2 = time.time() - start_time

print("\n--- Second Reconciliation Results (After Adjustments) ---")
if discrepancies_round2:
    df_discrepancies_round2 = pd.DataFrame(discrepancies_round2)
    display(df_discrepancies_round2)
else:
    print("No discrepancies found in the second reconciliation. Adjustments were successful.")

metrics_state = track_reconciliation_metrics(metrics_state, discrepancies_round2, reconciliation_time_round2, 0)

# 6. Edge Case Testing: Missing Position on one side

print("\n--- Edge Case: Missing Position ---")
# Scenario: Local has LTC, Exchange does not.
local_state_edge = create_position_state([
    {'asset': 'LTC', 'quantity': 1.0, 'avg_price': 100.0}
])
exchange_state_edge = create_exchange_state([
    {'asset': 'USD', 'balance': 500.0}
])

local_positions_edge = get_local_positions(local_state_edge)
exchange_positions_edge = get_exchange_positions(exchange_state_edge)

discrepancies_edge = reconcile_positions(local_positions_edge, exchange_positions_edge)

if discrepancies_edge:
    df_discrepancies_edge = pd.DataFrame(discrepancies_edge)
    display(df_discrepancies_edge)
else:
    print("No discrepancies found for the edge case.")

# Scenario: Exchange has DOGE, Local does not.
local_state_edge2 = create_position_state([
    {'asset': 'USD', 'quantity': 100.0, 'avg_price': 1.0}
])
exchange_state_edge2 = create_exchange_state([
    {'asset': 'DOGE', 'balance': 5000.0}
])

local_positions_edge2 = get_local_positions(local_state_edge2)
exchange_positions_edge2 = get_exchange_positions(exchange_state_edge2)

discrepancies_edge2 = reconcile_positions(local_positions_edge2, exchange_positions_edge2)

if discrepancies_edge2:
    df_discrepancies_edge2 = pd.DataFrame(discrepancies_edge2)
    display(df_discrepancies_edge2)

# 7. Summarize Metrics

print("\n--- Overall Reconciliation Metrics ---")
metrics_df = pd.DataFrame([metrics_state])
display(metrics_df)

# Plotting reconciliation time over runs (simplified as we have few runs here)
# For a real system, you'd collect data over many runs.
run_times = [reconciliation_time_round1, reconciliation_time_round2] # Add more if you run it multiple times
rp_metrics = {
    'Run': [1, 2],
    'Time (s)': run_times,
    'Discrepancies Found': [len(discrepancies_round1), len(discrepancies_round2)]
}
rp_metrics_df = pd.DataFrame(rp_metrics)

plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
sns.barplot(x='Run', y='Time (s)', data=rp_metrics_df, palette='viridis')
plt.title('Reconciliation Time per Run')
plt.xlabel('Reconciliation Run')
plt.ylabel('Time Taken (s)')
plt.grid(axis='y', linestyle='--', alpha=0.7)

plt.subplot(1, 2, 2)
sns.barplot(x='Run', y='Discrepancies Found', data=rp_metrics_df, palette='magma')
plt.title('Discrepancies Found per Run')
plt.xlabel('Reconciliation Run')
plt.ylabel('Number of Discrepancies')
plt.grid(axis='y', linestyle='--', alpha=0.7)

plt.tight_layout()
plt.show()