data_us.py

"""
data_us.py — US market data layer (yfinance), replacing baostock/pytdx.

Levels & history limits (Yahoo Finance API constraints):
    daily   : 10 years          (weekly / monthly are resampled from daily
                                 by chan_multilevel.resample_weekly/_monthly)
    60m     : last 730 days  (fetched as "1h" interval with explicit start/end dates)
    30m/15m : last 60 days
    5m      : last 60 days
    1m      : last 7 days only → too short for Chan decomposition, NOT used.
              MultiLevelChan handles a missing 1m level gracefully (skips it).

Output schema (identical to the original A-share loaders):
    date, open, close, high, low, volume, amount
`amount` (turnover) is approximated as close × volume (Yahoo has no turnover field).

All downloads are cached to parquet under ./_cache_us/<TICKER>/<level>.parquet
and refreshed when stale (daily: >12h old, intraday: >2h old) or on force=True.
"""
from __future__ import annotations

import os
import threading
import time
import traceback

import pandas as pd

import paths

# yfinance uses a shared SQLite cache (peewee) for timezone lookups.
# When multiple threads call yf.Ticker().history() simultaneously the DB
# gets locked and raises peewee.OperationalError, stalling prefetch and
# freezing the "Run analysis" button. Serialise the yfinance connect/lookup
# phase with a process-wide lock — Yahoo's own rate-limit is the real
# bottleneck anyway, so the extra serialisation costs almost nothing.
_YF_LOCK = threading.Lock()

CACHE_DIR = os.environ.get("CHAN_CACHE_DIR", paths.CACHE_DIR)

LEVELS = {
    # level: (yfinance interval, period_or_days)
    # For "60m" Yahoo requires explicit start/end dates (not a period string)
    # when fetching more than ~60 days back; we pass days as an int sentinel.
    "d":   ("1d",  "10y"),
    "60m": ("1h",  "730d"),    # use explicit start/end — "period='730d'" is rejected by Yahoo for 1h
    "30m": ("30m", "60d"),
    "15m": ("15m", "60d"),
    "5m":  ("5m",  "60d"),
    "1m":  ("1m",  "7d"),    # only 7 days available; short but usable for the
                             # finest nested-interval confirmation when present
}

_STALE_SECONDS = {"d": 12 * 3600, "60m": 2 * 3600, "30m": 2 * 3600,
                  "15m": 2 * 3600, "5m": 2 * 3600, "1m": 1800}


def _cache_path(ticker: str, level: str) -> str:
    d = os.path.join(CACHE_DIR, ticker.upper().replace("/", "_"))
    os.makedirs(d, exist_ok=True)
    return os.path.join(d, f"{level}.parquet")


def _normalize(df: pd.DataFrame) -> pd.DataFrame:
    """yfinance frame → engine schema (date/open/close/high/low/volume/amount)."""
    if df is None or len(df) == 0:
        return pd.DataFrame(columns=["date", "open", "close", "high", "low", "volume", "amount"])
    d = df.copy()
    if isinstance(d.columns, pd.MultiIndex):           # yf>=0.2 returns MultiIndex sometimes
        d.columns = [c[0] if isinstance(c, tuple) else c for c in d.columns]
    d = d.reset_index()
    # index column may be 'Date' or 'Datetime'
    for cand in ("Datetime", "Date", "index"):
        if cand in d.columns:
            d = d.rename(columns={cand: "date"})
            break
    d.columns = [str(c).lower() for c in d.columns]
    keep = {"date", "open", "high", "low", "close", "volume"}
    d = d[[c for c in d.columns if c in keep]]
    d["date"] = pd.to_datetime(d["date"])
    # strip timezone so comparisons with naive Timestamps in the engine work
    try:
        d["date"] = d["date"].dt.tz_localize(None)
    except (TypeError, AttributeError):
        pass
    d = d.dropna(subset=["open", "high", "low", "close"])
    d = d.sort_values("date").reset_index(drop=True)
    d["amount"] = d["close"] * d.get("volume", 0)
    return d[["date", "open", "close", "high", "low", "volume", "amount"]]


def load_level(ticker: str, level: str, force: bool = False) -> pd.DataFrame:
    """Load one level for a ticker, using parquet cache when fresh."""
    assert level in LEVELS, f"unknown level {level}"
    path = _cache_path(ticker, level)
    if not force and os.path.exists(path):
        age = time.time() - os.path.getmtime(path)
        if age < _STALE_SECONDS[level]:
            try:
                return pd.read_parquet(path)
            except Exception:
                pass
    try:
        import yfinance as yf
        from datetime import datetime, timedelta
        interval, period = LEVELS[level]
        # Acquire lock before any yfinance call — the shared peewee/SQLite
        # timezone cache raises "database is locked" under concurrent access.
        with _YF_LOCK:
            if isinstance(period, int):
                # Yahoo rejects period strings for hourly data older than ~60 days.
                # Use explicit start/end timestamps instead.
                end_dt = datetime.utcnow()
                start_dt = end_dt - timedelta(days=period)
                raw = yf.Ticker(ticker).history(start=start_dt, end=end_dt,
                                                interval=interval,
                                                auto_adjust=True, actions=False)
            else:
                raw = yf.Ticker(ticker).history(period=period, interval=interval,
                                                auto_adjust=True, actions=False)
        df = _normalize(raw)
        if len(df):
            df.to_parquet(path, index=False)
            return df
    except Exception:
        traceback.print_exc()
    # network failed → fall back to stale cache if any
    if os.path.exists(path):
        try:
            return pd.read_parquet(path)
        except Exception:
            pass
    return pd.DataFrame(columns=["date", "open", "close", "high", "low", "volume", "amount"])


# Full nested-interval set (区间套): the more sub-levels confirm, the more
# precise the buy/sell point. We fetch the deepest Yahoo allows. 1m has only
# 7 days of history — included when present, skipped gracefully otherwise.
# Downloads are parallel + cached, so the extra levels cost little wall-time.
FULL_LEVELS = ("d", "60m", "30m", "15m", "5m", "1m")
FAST_LEVELS = FULL_LEVELS  # default everywhere; alias kept for older callers


def load_levels(ticker: str, levels=FAST_LEVELS, force: bool = False) -> dict:
    return {lvl: load_level(ticker, lvl, force=force) for lvl in levels}


def load_all_levels(ticker: str, force: bool = False) -> dict:
    """Return {'d':…, '60m':…, '30m':…, '15m':…, '5m':…} (1m intentionally absent)."""
    return {lvl: load_level(ticker, lvl, force=force) for lvl in LEVELS}


def prefetch(tickers, levels=FAST_LEVELS, force: bool = False, workers: int = 5,
             budget_s: int = 45):
    """Download all (ticker, level) pairs in parallel with a hard time budget.
    Yahoo rate-limits datacenter IPs; without a budget one throttled request
    could hang the whole Run-analysis click. Whatever isn't fetched in time is
    skipped — the engine analyzes from daily/cached data and the next run
    picks up the rest."""
    from concurrent.futures import ThreadPoolExecutor, wait
    jobs = [(t, lvl) for t in tickers for lvl in levels]
    ex = ThreadPoolExecutor(max_workers=workers)
    futs = [ex.submit(load_level, t, lvl, force) for t, lvl in jobs]
    done, not_done = wait(futs, timeout=budget_s)
    ex.shutdown(wait=False, cancel_futures=True)
    return len(done), len(not_done)


def last_daily_date(ticker: str):
    df = load_level(ticker, "d")
    return None if df.empty else pd.Timestamp(df["date"].iloc[-1])