Forex_bot/charts.py

"""Chart generation — produces a PNG bytes object suitable for Telegram reply_photo."""
from __future__ import annotations

import io

import matplotlib
matplotlib.use("Agg")  # non-interactive backend — must be set before other imports

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import mplfinance as mpf
import numpy as np
import pandas as pd


_STYLE = mpf.make_mpf_style(
    base_mpf_style="nightclouds",
    rc={
        "font.size": 9,
        "axes.titlesize": 11,
        "axes.labelsize": 8,
    },
    facecolor="#1a1a2e",
    figcolor="#1a1a2e",
    gridcolor="#2a2a3e",
)


def _rsi_series(close: pd.Series, period: int = 14) -> pd.Series:
    delta = close.diff()
    gain = delta.clip(lower=0.0)
    loss = -delta.clip(upper=0.0)
    avg_gain = gain.ewm(alpha=1.0 / period, adjust=False).mean()
    avg_loss = loss.ewm(alpha=1.0 / period, adjust=False).mean()
    rs = avg_gain / avg_loss.replace(0.0, pd.NA)
    return (100.0 - (100.0 / (1.0 + rs))).fillna(100.0)


def _price_ylim(data: pd.DataFrame, pad_pct: float = 0.15) -> tuple[float, float]:
    """Return a padded Y-axis range so candle bodies/wicks are never clipped.

    Without this, mplfinance sometimes auto-zooms too tight and the candles
    appear as flat lines because the visible high-low range is a tiny fraction
    of the axis height.

    pad_pct=0.15 adds 15 % of the visible price range above and below.
    """
    lo = float(data["Low"].min())
    hi = float(data["High"].max())
    rng = hi - lo
    if rng == 0:
        rng = hi * 0.01  # fallback for perfectly flat data
    pad = rng * pad_pct
    return lo - pad, hi + pad


def generate(
    df: pd.DataFrame,
    pair: str,
    tf: str,
    verdict: str = "",
    confidence: int = 0,
    candles: int = 80,
) -> bytes:
    """Render a candlestick chart with EMA/BB overlays and RSI panel.

    Returns raw PNG bytes for direct use with Telegram reply_photo.

    Key fix: we explicitly set the price-panel Y limits with padding so
    candle bodies always take up the full vertical space and are never crushed.
    """
    # ── Compute indicators on full df for warm-up, then slice for display ──
    close = df["close"]
    ema20_full  = close.ewm(span=20, adjust=False).mean()
    ema50_full  = close.ewm(span=50, adjust=False).mean()
    mid_full    = close.rolling(20).mean()
    std_full    = close.rolling(20).std(ddof=0)
    bb_upper_full = mid_full + 2.0 * std_full
    bb_lower_full = mid_full - 2.0 * std_full
    rsi_full    = _rsi_series(close)

    n    = min(candles, len(df))
    data = df.tail(n).copy()
    # mplfinance requires capitalised column names
    data.columns = [c.capitalize() for c in data.columns]

    ema20    = ema20_full.tail(n).values
    ema50    = ema50_full.tail(n).values
    bb_upper = bb_upper_full.tail(n).values
    bb_lower = bb_lower_full.tail(n).values
    rsi_vals = rsi_full.tail(n).values

    # ── Determine padded price range BEFORE plotting ──
    ylim = _price_ylim(data)

    addplots = [
        mpf.make_addplot(ema20,    color="#00bfff", width=1.5),
        mpf.make_addplot(ema50,    color="#ffa500", width=1.5),
        mpf.make_addplot(bb_upper, color="#888888", width=0.9,
                         linestyle="--", alpha=0.7),
        mpf.make_addplot(bb_lower, color="#888888", width=0.9,
                         linestyle="--", alpha=0.7),
        mpf.make_addplot(rsi_vals, panel=1, color="#b388ff", width=1.2,
                         ylabel="RSI"),
    ]

    verdict_label = f"  {verdict} {confidence}%" if verdict else ""
    title = f"{pair}  {tf}{verdict_label}"

    # ── Plot — taller figure + tight_layout gives price panel full room ──
    fig, axes = mpf.plot(
        data,
        type="candle",
        style=_STYLE,
        addplot=addplots,
        volume=False,
        panel_ratios=(4, 1),      # give price panel even more vertical share
        title=title,
        returnfig=True,
        figsize=(12, 7),          # taller figure so candles breathe
        tight_layout=False,       # must be False to avoid mplfinance spurious vertical line bug
        ylim=ylim,                # ← the critical fix: explicit padded range
    )

    # ── RSI panel decoration ──
    # axes order: [price_ax, price_twin, rsi_ax, ...]
    rsi_ax = axes[2]
    rsi_ax.axhline(70, color="#ff5252", linewidth=0.8, linestyle="--", alpha=0.8)
    rsi_ax.axhline(50, color="#777777", linewidth=0.5, linestyle="--", alpha=0.5)
    rsi_ax.axhline(30, color="#69f0ae", linewidth=0.8, linestyle="--", alpha=0.8)
    rsi_ax.set_ylim(0, 100)
    rsi_ax.yaxis.set_major_locator(ticker.MultipleLocator(20))

    # ── EMA / BB legend on price panel ──
    price_ax = axes[0]
    price_ax.plot([], [], color="#00bfff", linewidth=1.5, label="EMA 20 (short-term trend)")
    price_ax.plot([], [], color="#ffa500", linewidth=1.5, label="EMA 50 (long-term trend)")
    price_ax.plot([], [], color="#888888", linewidth=0.9, linestyle="--",
                  label="Bollinger Bands")
    price_ax.legend(
        loc="upper left",
        fontsize=7,
        facecolor="#1a1a2e",
        edgecolor="#444466",
        labelcolor="white",
    )

    fig.subplots_adjust(left=0.08, right=0.92, top=0.92, bottom=0.10, hspace=0.05)

    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=140, facecolor=fig.get_facecolor())
    plt.close(fig)
    buf.seek(0)
    return buf.read()