!pip install pandas matplotlib seaborn numpy

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

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

def create_detector_state(lookback_window: int = 1000) -> dict:
    """
    Initializes the state for wash trade detection.

    Parameters
    ----------
    lookback_window : int, optional
        Size of the rolling window for trade history, defaults to 1000.

    Returns
    -------
    dict
        A dictionary containing buffers and tracking metrics.
    """
    logger.info(f"Initializing detector state with lookback window of {lookback_window}")
    return {
        'trade_history': deque(maxlen=lookback_window),
        'flagged_trades': [],
        'metrics': {
            'total_processed': 0,
            'total_flagged': 0,
            'suspicious_volume': 0.0
        }
    }

def detect_self_trades(state: dict, trade: dict) -> dict:
    """
    Identifies trades where the buyer and seller are the same.

    Parameters
    ----------
    state : dict
        Current state dictionary.
    trade : dict
        Trade data containing 'buyer_id', 'seller_id', 'price', and 'quantity'.

    Returns
    -------
    dict
        Updated state with flagged trades if self-trade detected.
    """
    state['metrics']['total_processed'] += 1

    if trade['buyer_id'] == trade['seller_id']:
        logger.warning(f"Self-trade detected! Asset: {trade['asset']}, Entity: {trade['buyer_id']}")
        flag = trade.copy()
        flag['reason'] = 'Self-Trade'
        state['flagged_trades'].append(flag)
        state['metrics']['total_flagged'] += 1
        state['metrics']['suspicious_volume'] += (trade['price'] * trade['quantity'])

    state['trade_history'].append(trade)
    return state

def detect_round_trips(state: dict, time_threshold_seconds: int = 60) -> dict:
    """
    Detects A-B-A trading patterns in the current trade window.

    Parameters
    ----------
    state : dict
        Current state.
    time_threshold_seconds : int, optional
        Seconds within which the reversal must occur, defaults to 60.

    Returns
    -------
    dict
        State with round-trip violations added to flagged_trades.
    """
    history = list(state['trade_history'])
    if len(history) < 2:
        return state

    latest = history[-1]
    for prev in reversed(history[:-1]):
        time_diff = (latest['timestamp'] - prev['timestamp']).total_seconds()
        if time_diff > time_threshold_seconds:
            break

        # Check if A sold to B and then B sold back to A for similar price/qty
        if (latest['buyer_id'] == prev['seller_id'] and
            latest['seller_id'] == prev['buyer_id'] and
            abs(latest['price'] - prev['price']) / prev['price'] < 0.01):

            logger.info("Round-trip trade pattern detected.")
            flag = latest.copy()
            flag['reason'] = 'Round-Trip'
            state['flagged_trades'].append(flag)
            state['metrics']['total_flagged'] += 1
            break

    return state

def simulate_market_data(n_trades: int = 100):
    data = []
    start_time = datetime.now()
    for i in range(n_trades):
        # Random jitter for timing
        start_time += timedelta(seconds=random.randint(1, 5)) + timedelta(milliseconds=random.random() * 100)

        # Normal trade
        trade = {
            'timestamp': start_time,
            'asset': 'BTC/USD',
            'buyer_id': f"user_{random.randint(1, 20)}",
            'seller_id': f"user_{random.randint(1, 20)}",
            'price': 50000 + random.uniform(-100, 100),
            'quantity': random.uniform(0.1, 2.0)
        }

        # Inject Wash Trade at index 50 (Self trade)
        if i == 50:
            trade['buyer_id'] = "malicious_actor"
            trade['seller_id'] = "malicious_actor"

        # Inject Round Trip at index 80-81
        if i == 81:
            prev = data[-1]
            trade['buyer_id'] = prev['seller_id']
            trade['seller_id'] = prev['buyer_id']
            trade['price'] = prev['price']

        data.append(trade)
    return data

# Run simulation
trades = simulate_market_data(200)
detector = create_detector_state()

for t in trades:
    detector = detect_self_trades(detector, t)
    detector = detect_round_trips(detector, time_threshold_seconds=30)

# Visualization
df_all = pd.DataFrame(trades)
df_flagged = pd.DataFrame(detector['flagged_trades'])

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))

sns.lineplot(data=df_all, x='timestamp', y='price', ax=ax1, label='Market Price')
if not df_flagged.empty:
    ax1.scatter(df_flagged['timestamp'], df_flagged['price'], color='red', zorder=5, label='Flagged Wash Trade')
ax1.set_title('Market Activity with Detected Wash Trades')
ax1.legend()

sns.barplot(x=['Total Trades', 'Flagged Trades'], y=[len(df_all), len(df_flagged)], ax=ax2, palette='muted')
ax2.set_title('Detection Statistics Summary')
plt.tight_layout()
plt.show()

display(pd.DataFrame([detector['metrics']]))