import pandas as pd     # Data manipulation and analysis
import numpy as np      # Numerical operations, especially for random number generation
import plotly.graph_objects as go # Interactive charting with Plotly Graph Objects
from plotly.subplots import make_subplots # For creating subplots in Plotly figures
import datetime         # Date and time manipulation
import random           # Generation of pseudo-random numbers

def generate_ohlcv(n_bars: int = 200, start_price: float = 50_000,
                   freq: str = "1h") -> pd.DataFrame:
    """
    Generates synthetic OHLCV price data using a random walk model.

    Parameters
    ----------
    n_bars      : int
        The number of candlestick bars to generate.
    start_price : float
        The initial price level for the series.
    freq        : str
        A Pandas frequency string specifying the interval for the datetime index (e.g., '1h', '1d').

    Returns
    -------
    pd.DataFrame
        A DataFrame containing 'datetime', 'open', 'high', 'low', 'close', and 'volume' columns.
    """
    np.random.seed(42) # Ensure reproducibility of random price movements
    closes  = [start_price]
    for _ in range(n_bars - 1):
        # Generate a random return for price movement to simulate a random walk
        ret = np.random.normal(0.0002, 0.012)
        closes.append(closes[-1] * (1 + ret))

    ohlcv_data = []
    for i, close in enumerate(closes):
        # Determine the open price based on the previous bar's close, or the start price for the first bar
        open_  = closes[i-1] if i > 0 else close
        # Introduce noise to high and low prices relative to open and close to simulate intra-bar volatility
        noise  = close * 0.005
        high   = max(open_, close) + abs(np.random.normal(0, noise))
        low    = min(open_, close) - abs(np.random.normal(0, noise))
        # Generate random volume data for the bar
        volume = np.random.uniform(50, 500)
        ohlcv_data.append({"open": open_, "high": high, "low": low, "close": close, "volume": volume})

    # Create a datetime index and construct the DataFrame.
    # Utilizes timezone-aware datetime.datetime.now(datetime.UTC) for modern datetime handling.
    dates = pd.date_range(end=datetime.datetime.now(datetime.timezone.utc), periods=n_bars, freq=freq)
    df    = pd.DataFrame(ohlcv_data, index=dates).reset_index().rename(columns={"index": "datetime"})
    return df


ohlcv_df = generate_ohlcv(200)
print(f"Generated {len(ohlcv_df)} bars. Price range: ${ohlcv_df['low'].min():.0f} – ${ohlcv_df['high'].max():.0f}.")
print("Sample of generated OHLCV data:")
print(ohlcv_df.head(3).to_string())

def generate_trades(ohlcv_df: pd.DataFrame, n_trades: int = 15) -> pd.DataFrame:
    """
    Generates synthetic trade records aligned with the OHLCV data.

    Each trade record specifies an entry bar, an exit bar, corresponding prices,
    and a reason for trade closure.

    Parameters
    ----------
    ohlcv_df : pd.DataFrame
        The OHLCV DataFrame generated by `generate_ohlcv`.
    n_trades : int
        The number of synthetic trades to generate.

    Returns
    -------
    pd.DataFrame
        A DataFrame containing trade entry/exit records, including 'entry_time',
        'exit_time', 'direction', 'entry_price', 'exit_price', 'pnl_pct', and 'reason'.
    """
    random.seed(7) # Ensure reproducibility of random trade generation
    trades = []
    used_bars = set() # Tracks bars occupied by trades to prevent overlaps
    n_bars = len(ohlcv_df)

    for _ in range(n_trades):
        # Select a random, non-overlapping entry bar index.
        # Ensures sufficient subsequent bars for trade duration.
        entry_idx = random.randint(5, n_bars - 20)

        # Add a safeguard to prevent infinite loops if no non-overlapping slot is found
        max_attempts = 1000
        attempts = 0
        while any(idx in used_bars for idx in range(entry_idx, entry_idx + 16)):
            entry_idx = random.randint(5, n_bars - 20)
            attempts += 1
            if attempts > max_attempts:
                # If too many attempts, skip this trade and try for the next one
                # print(f"Warning: Could not find a non-overlapping entry for trade {_ + 1} after {max_attempts} attempts. Skipping trade.")
                break

        # If we broke out of the while loop due to max_attempts, skip this iteration
        if attempts > max_attempts:
            continue

        # Determine a random holding period and calculate the corresponding exit index.
        hold_bars = random.randint(3, 15)
        exit_idx  = min(entry_idx + hold_bars, n_bars - 1) # Ensure exit_idx does not exceed DataFrame bounds

        # Extract entry and exit prices from the OHLCV data.
        entry_price = float(ohlcv_df.iloc[entry_idx]["close"])
        exit_price  = float(ohlcv_df.iloc[exit_idx]["close"])
        direction   = random.choice(["long", "short"])

        # Calculate PnL percentage based on trade direction.
        if direction == "long":
            pnl_pct = ((exit_price - entry_price) / entry_price) * 100
        else: # "short" direction
            pnl_pct = ((entry_price - exit_price) / entry_price) * 100

        # Assign a reason for trade closure based on the calculated PnL.
        reason = "take_profit" if pnl_pct > 0 else "stop_loss"

        # Append the trade details to the list.
        trades.append({
            "entry_time":   ohlcv_df.iloc[entry_idx]["datetime"],
            "exit_time":    ohlcv_df.iloc[exit_idx]["datetime"],
            "direction":    direction,
            "entry_price":  round(entry_price, 2),
            "exit_price":   round(exit_price, 2),
            "pnl_pct":      round(pnl_pct, 2),
            "reason":       reason,
        })
        # Mark bars within the trade's duration as used to prevent subsequent overlaps.
        used_bars.update(range(entry_idx, exit_idx + 1))

    return pd.DataFrame(trades)


trades_df = generate_trades(ohlcv_df)
print(f"Generated {len(trades_df)} synthetic trades.")
print("Sample of generated trade data:")
print(trades_df[["direction","entry_price","exit_price","pnl_pct","reason"]].head(5).to_string(index=False))

def build_candlestick_trace(df: pd.DataFrame) -> go.Candlestick:
    """
    Constructs a Plotly Candlestick trace from OHLCV data.

    Parameters
    ----------
    df : pd.DataFrame
        DataFrame containing 'datetime', 'open', 'high', 'low', 'close' columns.

    Returns
    -------
    go.Candlestick
        A Plotly Candlestick trace object for price visualization.
    """
    return go.Candlestick(
        x=df["datetime"],
        open=df["open"], high=df["high"],
        low=df["low"],   close=df["close"],
        increasing_line_color="#2ecc71", # Green for increasing candles
        decreasing_line_color="#e74c3c", # Red for decreasing candles
        name="Price",
        showlegend=False, # Candlestick trace does not require a legend entry
    )


def build_entry_markers(trades: pd.DataFrame) -> list[go.Scatter]:
    """
    Constructs entry marker traces for long and short trades.

    Long entries are represented by an upward-pointing triangle (▲) positioned below the entry price.
    Short entries are represented by a downward-pointing triangle (▼) positioned above the entry price.

    Parameters
    ----------
    trades : pd.DataFrame
        DataFrame containing trade records with 'entry_time', 'entry_price', and 'direction' columns.

    Returns
    -------
    list[go.Scatter]
        A list of `go.Scatter` traces, specifically one for each trade direction (long and short).
    """
    longs  = trades[trades["direction"] == "long"]
    shorts = trades[trades["direction"] == "short"]
    traces = []

    if not longs.empty:
        traces.append(go.Scatter(
            x=longs["entry_time"], y=longs["entry_price"] * 0.997, # Position slightly below entry price
            mode="markers",
            marker=dict(symbol="triangle-up", size=14, color="#3498db", line=dict(width=1, color="#fff")),
            name="Long Entry", legendgroup="entries", # Group legend entries for better organization
            hovertemplate="<b>LONG ENTRY</b><br>Price: $%{y:,.2f}<extra></extra>", # Custom hover text
        ))

    if not shorts.empty:
        traces.append(go.Scatter(
            x=shorts["entry_time"], y=shorts["entry_price"] * 1.003, # Position slightly above entry price
            mode="markers",
            marker=dict(symbol="triangle-down", size=14, color="#e67e22", line=dict(width=1, color="#fff")),
            name="Short Entry", legendgroup="entries",
            hovertemplate="<b>SHORT ENTRY</b><br>Price: $%{y:,.2f}<extra></extra>",
        ))

    return traces


def build_exit_markers(trades: pd.DataFrame) -> list[go.Scatter]:
    """
    Constructs exit marker traces, with color-coding based on the trade close reason.

    - 'take_profit' exits are denoted by a green circle.
    - 'stop_loss' exits are denoted by a red circle.
    - Additional reasons (e.g., 'flip') would be orange, though not present in the current synthetic data.

    Parameters
    ----------
    trades : pd.DataFrame
        DataFrame containing trade records with 'exit_time', 'exit_price', 'pnl_pct', and 'reason' columns.

    Returns
    -------
    list[go.Scatter]
        A list of `go.Scatter` traces, each categorized by exit reason and annotated with PnL.
    """
    colour_map = {"take_profit": "#2ecc71", "stop_loss": "#e74c3c"}
    traces     = []

    for reason, colour in colour_map.items():
        subset = trades[trades["reason"] == reason]
        if subset.empty:
            continue
        traces.append(go.Scatter(
            x=subset["exit_time"], y=subset["exit_price"],
            mode="markers+text", # Display both markers and PnL text
            marker=dict(symbol="circle", size=10, color=colour,
                        line=dict(width=1.5, color="#fff")),
            text=[f"{p:+.2f}%" for p in subset["pnl_pct"]], # PnL percentage formatted as text label
            textposition="top center",
            textfont=dict(size=9, color=colour),
            name=reason.replace("_", " ").title(), # Formatted name for legend
            hovertemplate=(
                f"<b>{reason.replace('_',' ').upper()}</b><br>" # Hover text for the reason
                "Price: $%{y:,.2f}<br>PnL: %{text}<extra></extra>" # Hover text for price and PnL
            ),
        ))

    return traces


def build_trade_connector_lines(trades: pd.DataFrame) -> list[go.Scatter]:
    """
    Generates lines connecting each trade's entry point to its corresponding exit point.

    This visualization aids in comprehending the duration and trajectory of each trade.
    Lines are color-coded based on the trade's PnL: green for profitable trades and red for losing trades.

    Parameters
    ----------
    trades : pd.DataFrame
        DataFrame containing trade records with 'entry_time', 'exit_time', 'entry_price', 'exit_price', and 'pnl_pct' columns.

    Returns
    -------
    list[go.Scatter]
        A list of `go.Scatter` traces, each representing a trade connector line.
    """
    traces = []
    for _, t in trades.iterrows():
        # Determine the line color based on PnL (green for profit, red for loss)
        colour = "#2ecc71" if t["pnl_pct"] > 0 else "#e74c3c"
        traces.append(go.Scatter(
            x=[t["entry_time"], t["exit_time"]],
            y=[t["entry_price"], t["exit_price"]],
            mode="lines",
            line=dict(color=colour, width=1, dash="dot"),
            showlegend=False, # Connector lines do not require a legend entry
            hoverinfo="skip", # No hover information is displayed for connector lines
        ))
    return traces


# Confirmation of chart function definition
# print("Chart construction functions successfully defined.") # Removed informal print statement

def build_trade_chart(ohlcv_df: pd.DataFrame, trades_df: pd.DataFrame,
                       title: str = "Trade Annotations",
                       show_volume: bool = True) -> go.Figure:
    """
    Assembles the complete annotated trading chart.

    Parameters
    ----------
    ohlcv_df    : pd.DataFrame
        DataFrame containing OHLCV price data.
    trades_df   : pd.DataFrame
        DataFrame containing trade records.
    title       : str
        The title string for the chart.
    show_volume : bool
        A boolean flag indicating whether to include a volume sub-panel below the main chart.

    Returns
    -------
    go.Figure
        A Plotly Figure object displaying the annotated candlestick chart.
    """
    # Determine subplot layout based on the inclusion of the volume sub-panel
    row_heights = [0.75, 0.25] if show_volume else [1.0]
    n_rows      = 2 if show_volume else 1

    fig = make_subplots(
        rows=n_rows, cols=1,
        shared_xaxes=True, # Enables synchronized zooming across subplots
        vertical_spacing=0.03,
        row_heights=row_heights,
    )

    # ── Candlestick Trace Integration ──────────────────────────────────────────
    fig.add_trace(build_candlestick_trace(ohlcv_df), row=1, col=1)

    # ── Trade Annotations Integration (Entry/Exit Markers and Connectors) ─────
    # Add entry markers for both long and short trades
    for trace in build_entry_markers(trades_df):
        fig.add_trace(trace, row=1, col=1)
    # Add exit markers, including PnL annotations
    for trace in build_exit_markers(trades_df):
        fig.add_trace(trace, row=1, col=1)
    # Add connector lines between trade entry and exit points
    for trace in build_trade_connector_lines(trades_df):
        fig.add_trace(trace, row=1, col=1)

    # ── Volume Bars Integration (Optional) ─────────────────────────────────────
    if show_volume:
        # Determine volume bar colors based on price movement (close vs. open)
        vol_colours = ["#2ecc71" if c >= o else "#e74c3c"
                       for c, o in zip(ohlcv_df["close"], ohlcv_df["open"])]
        fig.add_trace(go.Bar(
            x=ohlcv_df["datetime"], y=ohlcv_df["volume"],
            marker_color=vol_colours, name="Volume", showlegend=False,
        ), row=2, col=1)

    # ── Chart Layout Configuration ─────────────────────────────────────────────
    fig.update_layout(
        title=dict(text=title, font=dict(size=16)), # Chart title and font size
        height=650,
        paper_bgcolor="#1a1a2e", # Background color of the entire figure
        plot_bgcolor="#1a1a2e",   # Background color of the plotting area
        font=dict(color="#ccc"), # Default font color for text elements
        xaxis_rangeslider_visible=False, # Hides the range slider for a cleaner aesthetic
        legend=dict(orientation="h", y=1.02, x=0), # Positions the legend horizontally at the top
        margin=dict(l=40, r=20, t=60, b=40), # Adjusts plot margins for optimal display
    )
    # Configure x and y axes grid lines for enhanced readability
    fig.update_xaxes(showgrid=False, gridcolor="#2d2d44")
    fig.update_yaxes(showgrid=True,  gridcolor="#2d2d44")

    return fig


# ── Chart Generation and Display Execution ─────────────────────────────────
# Instantiate the chart with generated OHLCV data and trade records.
fig = build_trade_chart(
    ohlcv_df, trades_df,
    title="BTC/USDT — 1H — Trade Annotations",
    show_volume=True,
)

# ── Performance Statistics Overlay ─────────────────────────────────────────
# Calculate trade performance metrics.
n_wins   = (trades_df["pnl_pct"] > 0).sum()
n_losses = (trades_df["pnl_pct"] <= 0).sum()
total    = len(trades_df)

# Add an annotation to display calculated trade statistics directly on the chart.
fig.add_annotation(
    text=f"Total Trades: {total} | Wins: {n_wins} | Losses: {n_losses} | "
         f"Win Rate: {n_wins/total*100:.0f}%",
    xref="paper", yref="paper", x=0.01, y=0.98, # Positions the annotation in the top-left corner
    showarrow=False, bgcolor="#2d2d44", font=dict(color="#fff", size=11),
    bordercolor="#444", borderwidth=1,
)

# Display the generated Plotly figure.
fig.show()

# ── Chart Export to Local File ────────────────────────────────────────────
# Exports the interactive chart to an HTML file in the current working directory.
fig.write_html("trade_chart.html")
print(f"Chart exported to trade_chart.html.")
print(f"Number of trades plotted: {len(trades_df)}.")