Chan/ChanMacro/fetchers/ohlcv.py

"""
fetchers/ohlcv.py — Fetches BTC daily OHLCV from the existing data_provider service.

Also pre-computes EMA20/60/120, ATR(14), BB width, ADX(14).
"""

from datetime import date as Date, datetime, timedelta
from typing import Optional
import logging

import pandas as pd
import numpy as np
import requests

from .base import BaseFetcher
from config import config


class OHLCVFetcher(BaseFetcher):
    """Fetches BTC daily K-line data from data_provider API."""

    def __init__(self, provider_url: Optional[str] = None):
        super().__init__(timeout=30, max_retries=3)
        self.provider_url = provider_url or config.provider_url
        self.symbol = config.btc_symbol
        self.logger = logging.getLogger(__name__)

    def fetch(self, target_date: Optional[Date] = None) -> pd.DataFrame:
        """
        Fetch daily OHLCV for BTC. Returns DataFrame with computed indicators.

        Fetches enough history (200 bars) to compute EMAs/ATR/BB/ADX accurately.
        """
        url = f"{self.provider_url}/api/candles"
        params = {
            "symbol": self.symbol,
            "tf": "1d",
            "limit": 200,
        }
        resp = requests.get(url, params=params, timeout=self.timeout)
        resp.raise_for_status()
        data = resp.json()

        if not data:
            self.logger.warning("OHLCV API returned empty data")
            return pd.DataFrame()

        df = pd.DataFrame(data)
        df["timestamp"] = pd.to_datetime(df["timestamp"], unit="ms", utc=True)
        df["date"] = df["timestamp"].dt.date
        df = df.drop_duplicates(subset="date").sort_values("date").reset_index(drop=True)

        # Rename columns to match expected format
        df = df.rename(columns={
            "open": "open", "high": "high", "low": "low", "close": "close",
            "volume": "volume",
        })

        # Compute indicators
        df = self._add_indicators(df)

        return df

    def _add_indicators(self, df: pd.DataFrame) -> pd.DataFrame:
        """Add EMA, ATR, BB, ADX indicators."""
        close = df["close"].astype(float)
        high = df["high"].astype(float)
        low = df["low"].astype(float)

        # EMAs
        df["ema20"] = close.ewm(span=20, adjust=False).mean()
        df["ema60"] = close.ewm(span=60, adjust=False).mean()
        df["ema120"] = close.ewm(span=120, adjust=False).mean()

        # ATR(14)
        tr1 = high - low
        tr2 = (high - close.shift(1)).abs()
        tr3 = (low - close.shift(1)).abs()
        tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
        df["atr_14"] = tr.rolling(14).mean()

        # Bollinger Bands width
        sma20 = close.rolling(20).mean()
        std20 = close.rolling(20).std()
        df["bb_width"] = (2 * std20) / sma20 * 100  # as percentage

        # ADX(14)
        df["adx_14"] = self._compute_adx(df, period=14)

        return df

    @staticmethod
    def _compute_adx(df: pd.DataFrame, period: int = 14) -> pd.Series:
        """Compute ADX from OHLC data."""
        high = df["high"].astype(float)
        low = df["low"].astype(float)
        close = df["close"].astype(float)

        plus_dm = high.diff()
        minus_dm = low.diff().abs() * -1
        plus_dm = plus_dm.where(plus_dm > 0, 0)
        minus_dm = minus_dm.where(minus_dm < 0, 0).abs()

        tr1 = high - low
        tr2 = (high - close.shift(1)).abs()
        tr3 = (low - close.shift(1)).abs()
        tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)

        atr = tr.rolling(period).mean()
        plus_di = 100 * (plus_dm.rolling(period).mean() / atr)
        minus_di = 100 * (minus_dm.rolling(period).mean() / atr)

        dx = (abs(plus_di - minus_di) / (plus_di + minus_di)) * 100
        adx = dx.rolling(period).mean()
        return adx

    def store(self, db_path: str, records: list[dict]) -> int:
        """Store OHLCV records into SQLite. Not used directly — see store_df."""
        return 0

    def store_df(self, df: pd.DataFrame, db_path: Optional[str] = None) -> int:
        """Store the DataFrame into the ohlcv_daily table."""
        import sqlite3
        db_path = db_path or config.db_path
        conn = sqlite3.connect(db_path)

        count = 0
        for _, row in df.iterrows():
            if pd.isna(row.get("date")):
                continue
            date_str = str(row["date"])
            try:
                conn.execute("""
                    INSERT OR REPLACE INTO ohlcv_daily
                    (date, symbol, open, high, low, close, volume,
                     ema20, ema60, ema120, atr_14, bb_width, adx_14)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """, (
                    date_str, self.symbol,
                    float(row["open"]), float(row["high"]),
                    float(row["low"]), float(row["close"]),
                    float(row.get("volume", 0)),
                    float(row["ema20"]) if not pd.isna(row.get("ema20")) else None,
                    float(row["ema60"]) if not pd.isna(row.get("ema60")) else None,
                    float(row["ema120"]) if not pd.isna(row.get("ema120")) else None,
                    float(row["atr_14"]) if not pd.isna(row.get("atr_14")) else None,
                    float(row["bb_width"]) if not pd.isna(row.get("bb_width")) else None,
                    float(row["adx_14"]) if not pd.isna(row.get("adx_14")) else None,
                ))
                count += 1
            except Exception as e:
                self.logger.debug(f"Skip row {date_str}: {e}")

        conn.commit()
        conn.close()
        self.logger.info(f"Stored {count} OHLCV rows")
        return count