import time
import datetime
import random
import math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.gridspec import GridSpec
from collections import deque

# External libraries for data manipulation and visualization.
# `time` and `datetime` are standard libraries for time-related operations.
# `random` is used for simulating price movements in the demo.
# `math` provides mathematical functions.
# `numpy` is essential for numerical operations, especially with arrays.
# `pandas` is used for data structuring and analysis.
# `matplotlib.pyplot` is the primary plotting interface.
# `matplotlib.animation` enables dynamic chart updates.
# `matplotlib.gridspec` assists in arranging subplots within a figure.
# `collections.deque` provides an efficient double-ended queue for rolling data windows.

class LivePnLTracker:
    """
    Manages and tracks real-time Profit & Loss (PnL) for active financial positions.

    This class maintains a historical record of unrealized and realized PnL, calculates
    maximum drawdown, and provides structured data for visualization.

    Parameters
    ----------
    max_points : int
        The maximum number of data points to retain in the rolling PnL windows.
    """

    def __init__(self, max_points: int = 300):
        self.max_points = max_points

        # Rolling time series data structures for PnL components.
        # `deque` is utilized for efficient appending and popping from either end,
        # ensuring a fixed-size window of historical data.
        self.timestamps:      deque = deque(maxlen=max_points)
        self.unrealised_pnl:  deque = deque(maxlen=max_points)
        self.realised_pnl:    deque = deque(maxlen=max_points)
        self.total_pnl:       deque = deque(maxlen=max_points)
        self.prices:          deque = deque(maxlen=max_points) # Note: This deque is not currently used to store prices.

        # Dictionary to store details of currently open positions.
        # Keys are instrument symbols, values are dictionaries containing entry price,
        # position direction ('long' or 'short'), and quantity.
        self.positions:  dict  = {}

        # Cumulative realized PnL from all closed positions.
        self.closed_pnl: float = 0.0

        # Variables for drawdown tracking.
        # `peak_pnl` records the highest total PnL achieved.
        # `max_dd` stores the maximum percentage drawdown encountered.
        self.peak_pnl:   float = 0.0
        self.max_dd:     float = 0.0

    def add_position(self, symbol: str, entry: float,
                     direction: str, qty: float) -> None:
        """
        Registers a new open position for PnL tracking.

        Parameters
        ----------
        symbol : str
            The identifier for the financial instrument (e.g., "BTCUSDT").
        entry : float
            The average entry price of the position.
        direction : str
            The direction of the position ('long' or 'short').
        qty : float
            The quantity or size of the position.
        """
        self.positions[symbol] = {
            "entry":     entry,
            "direction": direction,
            "qty":       qty,
        }

    def close_position(self, symbol: str, exit_price: float) -> float:
        """
        Closes a tracked position and updates the cumulative realized PnL.

        Parameters
        ----------
        symbol : str
            The identifier of the financial instrument to close.
        exit_price : float
            The price at which the position is closed.

        Returns
        -------
        float
            The realized PnL in the base currency (e.g., USD) for the closed position.
        """
        if symbol not in self.positions:
            return 0.0

        pos = self.positions.pop(symbol) # Remove position from active tracking

        if pos["direction"] == "long":
            pnl = (exit_price - pos["entry"]) * pos["qty"]
        else:
            pnl = (pos["entry"] - exit_price) * pos["qty"]

        self.closed_pnl += pnl # Accumulate to total realized PnL
        return pnl

    def compute_unrealised_pnl(self, current_prices: dict) -> float:
        """
        Calculates the aggregate unrealized PnL for all currently open positions.

        Parameters
        ----------
        current_prices : dict
            A dictionary mapping instrument symbols to their current market prices.
            Example: `{"BTCUSDT": 68500.0, "ETHUSDT": 3600.0}`.

        Returns
        -------
        float
            The total unrealized PnL in the base currency.
        """
        total_unrealized = 0.0
        for symbol, pos in self.positions.items():
            # Use current price if available, otherwise use entry price (implies no change yet).
            price = current_prices.get(symbol, pos["entry"])
            if pos["direction"] == "long":
                total_unrealized += (price - pos["entry"]) * pos["qty"]
            else:
                total_unrealized += (pos["entry"] - price) * pos["qty"]
        return total_unrealized

    def update(self, current_prices: dict) -> None:
        """
        Records a new PnL state snapshot based on the latest market prices.

        This method computes current unrealized PnL, updates total PnL,
        and tracks maximum drawdown, appending these values to their respective
        historical data deques.

        Parameters
        ----------
        current_prices : dict
            A dictionary mapping instrument symbols to their latest market prices.
        """
        current_time = datetime.datetime.utcnow()
        unrealized_pnl = self.compute_unrealised_pnl(current_prices)
        total_pnl = self.closed_pnl + unrealized_pnl

        # Drawdown Tracking Logic
        # Update the highest PnL achieved to date (peak).
        self.peak_pnl = max(self.peak_pnl, total_pnl)

        # Calculate and update maximum drawdown if a positive peak PnL exists.
        if self.peak_pnl > 0:
            # Drawdown is expressed as a percentage of the peak PnL.
            current_drawdown_percentage = ((total_pnl - self.peak_pnl) / abs(self.peak_pnl)) * 100
            self.max_dd = min(self.max_dd, current_drawdown_percentage)
        elif total_pnl < 0: # If peak is 0 or negative, and current PnL is negative, drawdown is -100% relative to 0 or even worse.
            # For simplicity, if PnL is consistently negative from the start, we can cap initial drawdown at -100% or use the PnL value itself.
            # Here, we track it as the minimum negative percentage relative to the initial zero or negative peak.
            # A more robust solution might initialize peak_pnl with the first total_pnl if it's positive.
            pass # Max_dd remains at initial 0.0 or lowest recorded negative.

        # Append current PnL metrics to their historical deques.
        self.timestamps.append(current_time)
        self.unrealised_pnl.append(unrealized_pnl)
        self.realised_pnl.append(self.closed_pnl)
        self.total_pnl.append(total_pnl)

    def get_series(self) -> tuple:
        """
        Retrieves the historical PnL time series data.

        Returns
        -------
        tuple
            A tuple containing four lists: (timestamps, unrealized_pnl, realized_pnl, total_pnl).
        """
        return (list(self.timestamps), list(self.unrealised_pnl),
                list(self.realised_pnl), list(self.total_pnl))

def run_live_pnl_plot(tracker: LivePnLTracker,
                       price_source_fn,
                       interval_ms: int = 500,
                       max_frames: int = 100) -> None:
    """
    Executes a Matplotlib animation to visualize live PnL data.

    Parameters
    ----------
    tracker : LivePnLTracker
        An instance of `LivePnLTracker` with registered positions.
    price_source_fn : Callable[[], dict]
        A callable function that returns a dictionary of current market prices
        (`{symbol: price}`). This function simulates or fetches live price data.
    interval_ms : int, optional
        The delay between frames in milliseconds. Default is 500ms.
    max_frames : int, optional
        The total number of frames to generate before the animation stops.
        If `None`, the animation runs indefinitely. Default is 100 frames.
    """
    fig = plt.figure(figsize=(14, 7), facecolor="#1a1a2e") # Initialize figure with dark background
    gs  = GridSpec(2, 2, figure=fig, hspace=0.4, wspace=0.3) # Define a 2x2 grid for subplots

    # Assign subplots to specific grid locations
    ax_total    = fig.add_subplot(gs[0, :]) # Total PnL (spans both columns)
    ax_unreal   = fig.add_subplot(gs[1, 0]) # Unrealized PnL
    ax_dd       = fig.add_subplot(gs[1, 1]) # Drawdown

    # Apply common styling to all axes for a consistent dark theme
    for ax in [ax_total, ax_unreal, ax_dd]:
        ax.set_facecolor("#1a1a2e") # Dark background for the plot area
        ax.tick_params(colors="#ccc") # Light grey tick labels
        ax.spines["bottom"].set_color("#444") # Dark grey bottom spine
        ax.spines["left"].set_color("#444")   # Dark grey left spine
        ax.spines["top"].set_visible(False)   # Hide top spine
        ax.spines["right"].set_visible(False) # Hide right spine
        ax.title.set_color("#fff")            # White title text
        ax.xaxis.label.set_color("#ccc")      # Light grey x-axis label
        ax.yaxis.label.set_color("#ccc")      # Light grey y-axis label

    frame_count = [0] # Counter for animation frames

    def animate(i):
        """
        Update function for the Matplotlib animation. Called every `interval_ms`.
        """
        # Stop animation if max_frames limit is reached
        if max_frames and frame_count[0] >= max_frames:
            # Close the figure to properly terminate the animation and free resources.
            plt.close(fig)
            return

        # Update tracker with the latest simulated/live prices
        prices = price_source_fn()
        tracker.update(prices)
        frame_count[0] += 1

        # Retrieve updated PnL series from the tracker
        ts, unreal, realised, total = tracker.get_series()

        # Ensure there is sufficient data for plotting
        if len(ts) < 2:
            return

        # ── Total PnL Chart Update ─────────────────────────────────────────
        ax_total.clear()
        ax_total.set_facecolor("#1a1a2e") # Maintain dark background
        # Determine line color based on final total PnL
        pnl_color = "#2ecc71" if total[-1] >= 0 else "#e74c3c" # Green for profit, red for loss
        ax_total.plot(range(len(total)), total, color=pnl_color, linewidth=2) # Plot total PnL line
        ax_total.fill_between(range(len(total)), total, 0,
                               alpha=0.15, color=pnl_color) # Fill area under the curve
        ax_total.axhline(0, color="#555", linewidth=1) # Horizontal line at zero PnL
        ax_total.set_title(f"Total PnL (USD)   Current: ${total[-1]:+.2f}   "
                           f"Max Drawdown: {tracker.max_dd:.2f}%", color="#fff") # Dynamic title
        ax_total.set_ylabel("PnL (USD)", color="#ccc") # Y-axis label
        ax_total.tick_params(colors="#ccc") # Light grey tick parameters
        ax_total.set_xticks([]) # Hide x-axis ticks for cleanliness

        # ── Unrealized PnL Chart Update ────────────────────────────────────
        ax_unreal.clear()
        ax_unreal.set_facecolor("#1a1a2e") # Maintain dark background
        ax_unreal.plot(range(len(unreal)), unreal, color="#3498db", linewidth=1.5) # Plot unrealized PnL
        ax_unreal.axhline(0, color="#555", linewidth=1) # Horizontal line at zero PnL
        ax_unreal.set_title(f"Unrealized PnL: ${unreal[-1]:+.2f}", color="#fff") # Dynamic title
        ax_unreal.tick_params(colors="#ccc") # Light grey tick parameters
        ax_unreal.set_xticks([]) # Hide x-axis ticks
        ax_unreal.set_ylabel("PnL (USD)", color="#ccc")

        # ── Drawdown Plot Update ───────────────────────────────────────────
        ax_dd.clear()
        ax_dd.set_facecolor("#1a1a2e") # Maintain dark background

        if len(total) > 0:
            # Calculate running peaks from the total PnL series
            current_peaks = np.maximum.accumulate(np.array(total))

            # Calculate drawdown percentages relative to the running peak.
            # Handle cases where current_peaks might be zero to avoid division by zero.
            # If peak is 0 or negative, drawdown is considered 0 or undefined for positive percentage calculations.
            drawdown_percentages = np.where(current_peaks > 0, (total - current_peaks) / current_peaks * 100, 0)

            # Fill the area below 0% for negative drawdowns
            ax_dd.fill_between(range(len(total)), drawdown_percentages, 0,
                               where=drawdown_percentages < 0, color="#e74c3c", alpha=0.6) # Red fill for drawdown area
            ax_dd.axhline(0, color="#555", linewidth=1) # Reference line for 0% drawdown
            ax_dd.set_title(f"Drawdown (%): Max {tracker.max_dd:.2f}%", color="#fff")
        else:
            ax_dd.set_title("Drawdown Data Not Available", color="#fff")

        ax_dd.tick_params(colors="#ccc")
        ax_dd.set_xticks([]) # Hide x-axis ticks
        ax_dd.set_ylabel("Drawdown (%)", color="#ccc")

        # Set overall figure title with current UTC time
        fig.suptitle(f"Live PnL Monitor — {datetime.datetime.utcnow().strftime('%H:%M:%S')} UTC",
                     color="#fff", fontsize=13)

    # Initialize and run the animation
    # `cache_frame_data=False` prevents caching of frames, which is good for truly live data.
    ani = animation.FuncAnimation(fig, animate, interval=interval_ms, cache_frame_data=False)
    plt.show() # Display the plot
    return ani # Return the animation object if further control is needed

# ── PnL Tracker Initialization ─────────────────────────────────────────────
# Instantiate the LivePnLTracker to store and compute PnL metrics.
tracker = LivePnLTracker(max_points=200) # Retain up to 200 data points in historical deques.

# Register initial open positions with the tracker.
# Example: A long position in BTCUSDT and a short position in ETHUSDT.
tracker.add_position("BTCUSDT", entry=68_000.0, direction="long",  qty=0.15)
tracker.add_position("ETHUSDT", entry=3_500.0,  direction="short", qty=2.0)

# ── Simulated Price Source Definition ──────────────────────────────────────
# Initialize current prices for the simulated instruments.
_prices = {"BTCUSDT": 68_000.0, "ETHUSDT": 3_500.0}

def simulated_price_source() -> dict:
    """
    Generates simulated market prices with random walk characteristics.

    This function mimics live exchange data feeds by applying small, random
    Gaussian perturbations to the last known price of each instrument.

    Returns
    -------
    dict
        A dictionary mapping instrument symbols to their randomly walked prices.
    """
    # Apply a Gaussian random walk to BTCUSDT price.
    _prices["BTCUSDT"] = max(1, _prices["BTCUSDT"] * (1 + random.gauss(0.0002, 0.008)))
    # Apply a Gaussian random walk to ETHUSDT price.
    _prices["ETHUSDT"] = max(1, _prices["ETHUSDT"] * (1 + random.gauss(0.0001, 0.010)))
    return dict(_prices)

# ── Simulation Execution ───────────────────────────────────────────────────
print("Executing 60-frame PnL simulation...")
for frame_idx in range(60):
    # Update the tracker with prices from the simulated source in each frame.
    tracker.update(simulated_price_source())

    # Simulate closing a position at a specific frame (e.g., frame 40).
    if frame_idx == 39: # Corresponds to the 40th frame (0-indexed)
        closed_pnl_eth = tracker.close_position("ETHUSDT", exit_price=3_440.0)
        print(f"  INFO: ETHUSDT position closed at frame {frame_idx+1} with realized PnL: ${closed_pnl_eth:+.2f}")

# Retrieve the complete historical PnL series after the simulation.
ts, unreal, realised, total = tracker.get_series()

# ── Static Summary Chart Generation ────────────────────────────────────────
# Create a figure and a 1x2 grid of subplots for static analysis.
fig, axes = plt.subplots(1, 2, figsize=(14, 5), facecolor="#1a1a2e")
fig.suptitle("PnL Tracker — 60-Frame Simulation Summary", color="#fff", fontsize=13)

# Apply common styling to subplots.
for ax in axes:
    ax.set_facecolor("#1a1a2e") # Dark background
    ax.tick_params(colors="#ccc") # Light grey ticks
    for spine in ax.spines.values():
        spine.set_color("#444") # Dark grey spines

# Plot 1: Total PnL over the simulation duration.
# Determine plot color based on the final total PnL.
final_pnl_color = "#2ecc71" if total[-1] >= 0 else "#e74c3c"
axes[0].plot(total, color=final_pnl_color, linewidth=2)
axes[0].fill_between(range(len(total)), total, 0, alpha=0.2, color=final_pnl_color)
axes[0].axhline(0, color="#555") # Horizontal zero line
axes[0].set_title(f"Total PnL  Final: ${total[-1]:+.2f}", color="#fff")
axes[0].set_ylabel("PnL (USD)", color="#ccc")
axes[0].set_xlabel("Simulation Frame", color="#ccc")

# Plot 2: Comparison of Unrealized vs. Realized PnL.
axes[1].plot(unreal,   label="Unrealized PnL", color="#3498db", linewidth=1.5)
axes[1].plot(realised, label="Realized PnL",   color="#2ecc71", linewidth=1.5, linestyle="--")
axes[1].axhline(0, color="#555") # Horizontal zero line
axes[1].set_title("Unrealized vs. Realized PnL", color="#fff")
axes[1].legend(facecolor="#2d2d44", labelcolor="#fff", frameon=True, edgecolor="#555") # Legend styling
axes[1].set_ylabel("PnL (USD)", color="#ccc")
axes[1].set_xlabel("Simulation Frame", color="#ccc")

plt.tight_layout() # Adjust layout to prevent overlaps

# Save the generated figure to a file.
# The `facecolor` argument ensures the dark background is preserved in the saved image.
# plt.savefig("/home/claude/live_pnl_simulation.png", dpi=150, bbox_inches="tight",
#             facecolor="#1a1a2e") # Commented out for Colab environment where local file paths might not be directly accessible by default.

plt.show() # Display the static plots.

# Output final PnL metrics to console.
print(f"\nSimulation Summary:")
print(f"  Final Total PnL      : ${total[-1]:+.2f}")
print(f"  Final Unrealized PnL : ${unreal[-1]:+.2f}")
print(f"  Final Realized PnL   : ${realised[-1]:+.2f}")
print(f"  Maximum Drawdown     : {tracker.max_dd:.2f}%")

# Example of how to run the live animation (this code is not executed in this block).
# ani = run_live_pnl_plot(tracker, simulated_price_source, interval_ms=100, max_frames=200)
# plt.show()

# Execute the live PnL animation.
# The 'tracker' and 'simulated_price_source' are inherited from the previous simulation section.
# `interval_ms` controls the update frequency (e.g., 100ms = 10 updates per second).
# `max_frames` limits the animation to 200 updates for demonstration purposes.
print("Starting live PnL animation... (This will run for 200 frames)")
# ani = run_live_pnl_plot(tracker, simulated_price_source, interval_ms=100, max_frames=200)

# To display the animation in a browser environment (like Jupyter/Colab),
# the animation object needs to be implicitly rendered or saved.
# For inline display, the 'ani' object itself often triggers the display.
# If not, use IPython.display.HTML(ani.to_jshtml()) or save it.

# Note: In some environments, the animation might not display inline automatically.
# If the plot window doesn't appear, ensure the backend is set correctly or
# explicitly save the animation to a file (e.g., ani.save('live_pnl.gif')).
# For Colab, `plt.show()` within `run_live_pnl_plot` should render it.