A_Share_DP/src/ashare_dp/wyckoff/annotate.py

"""Walk-forward Wyckoff phase/event annotations for chart overlay."""

from __future__ import annotations

from datetime import date

from ashare_dp.domain.wyckoff import OHLCVFrame, WyckoffCycle, WyckoffEvent, WyckoffPhase
from ashare_dp.wyckoff.cycle import CycleEngine
from ashare_dp.wyckoff.event import EventEngine
from ashare_dp.wyckoff.features import FeatureEngine
from ashare_dp.wyckoff.phase import PhaseEngine

_MIN_BARS = {"1d": 40, "1w": 26, "1M": 18}

_NOTABLE_EVENTS = {
    WyckoffEvent.PS.value,
    WyckoffEvent.SC.value,
    WyckoffEvent.AR.value,
    WyckoffEvent.ST.value,
    WyckoffEvent.SPRING.value,
    WyckoffEvent.TEST.value,
    WyckoffEvent.SOS.value,
    WyckoffEvent.LPS.value,
    WyckoffEvent.JUMP.value,
    WyckoffEvent.BACKUP.value,
    WyckoffEvent.BC.value,
    WyckoffEvent.UTAD.value,
    WyckoffEvent.SOW.value,
    WyckoffEvent.LPSY.value,
}


def _slice_frame(frame: OHLCVFrame, end_idx: int) -> OHLCVFrame:
    n = end_idx + 1
    return OHLCVFrame(
        ts_code=frame.ts_code,
        timeframe=frame.timeframe,
        trade_dates=frame.trade_dates[:n],
        open=frame.open[:n],
        high=frame.high[:n],
        low=frame.low[:n],
        close=frame.close[:n],
        volume=frame.volume[:n],
        amount=frame.amount[:n] if frame.amount else [],
    )


def _compress_phases(points: list[tuple[str, str]]) -> list[dict]:
    """points: [(date_iso, phase), ...] → segments."""
    if not points:
        return []
    segs: list[dict] = []
    start, phase = points[0]
    prev = start
    for d, p in points[1:]:
        if p != phase:
            segs.append({"start": start, "end": prev, "phase": phase})
            start, phase = d, p
        prev = d
    segs.append({"start": start, "end": prev, "phase": phase})
    return segs


def annotate_frame(frame: OHLCVFrame, step: int | None = None) -> dict:
    """Pure annotation: phase bands + event markers + latest levels.

    ``step`` defaults by timeframe to keep interactive charts snappy.
    """
    tf = frame.timeframe
    min_bars = _MIN_BARS.get(tf, 30)
    if step is None:
        step = {"1d": 2, "1w": 1, "1M": 1}.get(tf, 2)

    empty = {
        "phases": [],
        "events": [],
        "levels": {},
        "bars": len(frame),
        "timeframe": tf,
    }
    if frame.empty or len(frame) < min_bars:
        return empty

    feat_eng = FeatureEngine()
    cycle_eng = CycleEngine()
    phase_eng = PhaseEngine()
    event_eng = EventEngine()

    phase_points: list[tuple[str, str]] = []
    events: list[dict] = []
    last_event: str | None = None
    levels: dict = {}

    # Ensure last bar is always evaluated
    indices = list(range(min_bars - 1, len(frame), step))
    if indices[-1] != len(frame) - 1:
        indices.append(len(frame) - 1)

    for i in indices:
        sub = _slice_frame(frame, i)
        f = feat_eng.run(sub, tf)
        c = cycle_eng.run(f, tf)
        p = phase_eng.run(c, f, tf)
        e = event_eng.run(c, p, f, tf)

        d = str(frame.trade_dates[i])[:10]
        phase = p.payload.get("phase") or WyckoffPhase.NONE.value
        phase_points.append((d, phase))

        cur = e.payload.get("current_event") or WyckoffEvent.NONE.value
        if cur in _NOTABLE_EVENTS and cur != last_event:
            events.append({
                "date": d,
                "event": cur,
                "price": float(frame.close[i]),
                "low": float(frame.low[i]),
                "high": float(frame.high[i]),
            })
            last_event = cur
        elif cur == WyckoffEvent.NONE.value:
            last_event = None

        if i == len(frame) - 1 and not f.payload.get("insufficient"):
            levels = {
                k: f.payload.get(k)
                for k in (
                    "range_high", "range_low", "ma20", "ma60",
                    "swing_high", "swing_low", "close",
                )
                if f.payload.get(k) is not None
            }
            levels["phase"] = phase
            levels["cycle"] = c.payload.get("cycle")
            levels["current_event"] = cur

    return {
        "phases": _compress_phases(phase_points),
        "events": events,
        "levels": levels,
        "bars": len(frame),
        "timeframe": tf,
    }


_RANGE_CYCLES = {
    WyckoffCycle.ACCUMULATION.value,
    WyckoffCycle.RE_ACCUMULATION.value,
    WyckoffCycle.DISTRIBUTION.value,
    WyckoffCycle.RE_DISTRIBUTION.value,
}


def _build_range_zones(
    price_frame: OHLCVFrame,
    cycle_segs: list[dict],
    levels: dict | None = None,
) -> list[dict]:
    """Build price boxes (high/low × date span) for accum/distrib ranges."""
    if price_frame.empty:
        return []
    dates = [str(d)[:10] for d in price_frame.trade_dates]
    highs = price_frame.high
    lows = price_frame.low
    zones: list[dict] = []

    for seg in cycle_segs or []:
        cy = seg.get("cycle")
        if cy not in _RANGE_CYCLES:
            continue
        start, end = seg["start"], seg["end"]
        idxs = [i for i, d in enumerate(dates) if start <= d <= end]
        if not idxs:
            # weekly bar date may sit between daily bars — take nearest window
            i0 = next((i for i, d in enumerate(dates) if d >= start), None)
            if i0 is None:
                continue
            i1 = next((i for i, d in enumerate(dates) if d > end), len(dates)) - 1
            idxs = list(range(i0, max(i0, i1) + 1))
        if not idxs:
            continue
        # pad short weekly hits to at least ~1 week of dailies for visibility
        if len(idxs) < 5 and idxs[-1] + 1 < len(dates):
            extra = min(5 - len(idxs), len(dates) - 1 - idxs[-1])
            idxs = list(range(idxs[0], idxs[-1] + 1 + max(0, extra)))
        hi = max(highs[i] for i in idxs)
        lo = min(lows[i] for i in idxs)
        if hi <= lo:
            continue
        zones.append({
            "kind": cy,
            "start": dates[idxs[0]],
            "end": dates[idxs[-1]],
            "high": float(hi),
            "low": float(lo),
            "current": False,
        })

    # Always expose the latest trading-range box from feature snapshot
    levels = levels or {}
    rh, rl = levels.get("range_high"), levels.get("range_low")
    if rh is not None and rl is not None and float(rh) > float(rl):
        look = min(60, len(dates))
        cy = levels.get("cycle") or "Unknown"
        if cy not in _RANGE_CYCLES:
            # Phase B/C in a range → treat as accumulation-style TR for display
            ph = levels.get("phase") or ""
            if ph in ("A", "B", "C"):
                cy = WyckoffCycle.ACCUMULATION.value
            elif ph in ("D", "E") and float(levels.get("close") or 0) < float(rh):
                cy = WyckoffCycle.ACCUMULATION.value
            else:
                cy = "Range"
        zones.append({
            "kind": cy,
            "start": dates[-look],
            "end": dates[-1],
            "high": float(rh),
            "low": float(rl),
            "current": True,
        })

    return zones


def annotate_symbol(
    ts_code: str,
    freq: str,
    end_date: date | None = None,
    lookback: int = 180,
) -> dict:
    """IO + annotate for one symbol (used by API).

    For daily charts, phase bands come from **weekly** structure (Wyckoff
    primary timeframe), while event markers / levels come from daily.
    """
    from ashare_dp.wyckoff.io import latest_daily_trade_date, load_frames_batch

    if freq not in ("1d", "1w", "1M"):
        raise ValueError(f"unsupported freq: {freq}")
    ed = end_date or latest_daily_trade_date()
    empty = {
        "ts_code": ts_code,
        "freq": freq,
        "phases": [],
        "events": [],
        "levels": {},
        "zones": [],
        "bars": 0,
        "phase_source": freq,
    }
    if ed is None:
        return empty

    if freq == "1d":
        daily_frames = load_frames_batch("1d", ed, lookback, ts_codes=[ts_code])
        weekly_frames = load_frames_batch("1w", ed, max(60, lookback // 3), ts_codes=[ts_code])
        daily = daily_frames.get(ts_code)
        weekly = weekly_frames.get(ts_code)
        if daily is None:
            return empty
        d_ann = annotate_frame(daily)
        w_ann = annotate_frame(weekly) if weekly is not None else {"phases": []}
        cycles = _cycle_segments(weekly) if weekly is not None else []
        levels = d_ann.get("levels") or {}
        # Prefer weekly cycle on the latest levels for zone labeling
        if cycles:
            levels = {**levels, "cycle": cycles[-1].get("cycle") or levels.get("cycle")}
            # latest non-None weekly phase
            for p in reversed(w_ann.get("phases") or []):
                if p.get("phase") not in (None, "None"):
                    levels = {**levels, "phase": p["phase"]}
                    break
        return {
            "ts_code": ts_code,
            "freq": freq,
            "end_date": ed.isoformat(),
            "phases": w_ann.get("phases") or [],
            "events": d_ann.get("events") or [],
            "levels": d_ann.get("levels") or {},
            "zones": _build_range_zones(daily, cycles, levels),
            "bars": d_ann.get("bars", 0),
            "phase_source": "1w",
            "cycles": cycles,
        }

    frames = load_frames_batch(freq, ed, lookback, ts_codes=[ts_code])
    frame = frames.get(ts_code)
    if frame is None:
        return empty
    out = annotate_frame(frame)
    out["ts_code"] = ts_code
    out["freq"] = freq
    out["end_date"] = ed.isoformat()
    out["phase_source"] = freq
    out["cycles"] = _cycle_segments(frame)
    out["zones"] = _build_range_zones(frame, out["cycles"], out.get("levels") or {})
    if freq == "1M":
        # Monthly chart: cycle bands are more meaningful than phase
        if not any(p.get("phase") not in (None, "None") for p in out["phases"]):
            out["phases"] = [
                {"start": c["start"], "end": c["end"], "phase": c["cycle"]}
                for c in out["cycles"]
                if c.get("cycle") and c["cycle"] != "Unknown"
            ]
    return out


def _cycle_segments(frame: OHLCVFrame, step: int | None = None) -> list[dict]:
    """Walk-forward cycle labels compressed to segments."""
    tf = frame.timeframe
    min_bars = _MIN_BARS.get(tf, 30)
    if step is None:
        step = {"1d": 3, "1w": 1, "1M": 1}.get(tf, 2)
    if frame.empty or len(frame) < min_bars:
        return []

    feat_eng = FeatureEngine()
    cycle_eng = CycleEngine()
    points: list[tuple[str, str]] = []
    indices = list(range(min_bars - 1, len(frame), step))
    if indices[-1] != len(frame) - 1:
        indices.append(len(frame) - 1)
    for i in indices:
        sub = _slice_frame(frame, i)
        f = feat_eng.run(sub, tf)
        c = cycle_eng.run(f, tf)
        points.append((str(frame.trade_dates[i])[:10], c.payload.get("cycle") or "Unknown"))
    segs = _compress_phases(points)
    return [{"start": s["start"], "end": s["end"], "cycle": s["phase"]} for s in segs]