momentum-system/tests/unit/test_recap_extract.py

"""Pure pieces of the Ross recap extractor: verbatim-quote verification, chunking, and the
merge of per-chunk results. The Ollama call itself is exercised by the gold-set evaluation."""

from collections import Counter

from youtube_claims.recap.extract import (chunk_bundle, clean_ticker, fix_spoken_prices, merge_records,
                                          scanner_prices, verify_record)
from youtube_claims.recap.quotes import Transcript, transcript_text_from_bundle

TX = Transcript(
    "So I got in on that dip right there, and I got a nice little trade as it rallied back up. "
    "Then we come back up right here looking for a curl through VWAP. It fails. Doesn't happen. "
    "I took a small loss right there. GLND. This is not an easy one to day trade because it's very thickly traded. "
    "Generally, the right approach that I have is to avoid stocks that are giving us those big topping tails.")


# ---------------------------------------------------------------- quotes
def test_exact_quote_is_found_and_snapped_to_transcript_words():
    got = TX.find("looking for a curl through vwap")
    assert got == "looking for a curl through VWAP."


def test_quote_matching_ignores_punctuation_and_case():
    assert TX.find("It fails, doesn't happen, I took a small loss") == "It fails. Doesn't happen. I took a small loss"


def test_one_miscopied_word_in_a_long_quote_is_tolerated():
    q = "This is not an easy one to trade because it's very thickly traded"   # dropped "day"
    got = TX.find(q)
    assert got is not None and got.startswith("This is not an easy one") and "thickly traded" in got


def test_invented_quote_is_rejected():
    assert TX.find("I made ten thousand dollars on the open with a huge short squeeze") is None
    # a real 40-char prefix with an invented tail is NOT enough (the old prefix check accepted it)
    assert TX.find("So I got in on that dip right there, and then I sold the whole thing at twelve dollars") is None


def test_too_short_quote_is_not_evidence():
    assert TX.find("it fails") is None
    assert TX.find("") is None


def test_transcript_text_from_bundle_strips_timestamps():
    b = "video_id: x\nTRANSCRIPT (mm:ss text):\n[00:01] Hello there.\n[00:05] Second line."
    assert transcript_text_from_bundle(b) == "Hello there. Second line."


# ---------------------------------------------------------------- verification
def _raw():
    return {
        "video_kind": "daily_recap", "session_date": "2026-09-24",
        "day": {"net_pnl": "$2,992.62", "green": True, "market_temperature": "Cold",
                "market_quote": "totally made up sentence about the market being slow today"},
        "trades": [
            {"ticker_heard": "a pus", "ticker_guess": "$apus", "setup": "dip_and_curl",
             "setup_quote": "I got in on that dip right there", "entry_price": "5.50", "outcome": "WIN",
             "why_entered": [{"reason": "dip", "quote": "So I got in on that dip right there"},
                             {"reason": "fake", "quote": "because the level two was stacked with buyers everywhere"}],
             "stock_context": {"catalyst": "headline", "catalyst_quote": "a quote that does not exist anywhere here",
                               "other": ["topping tails", "rocket ship"]}},
            {"ticker_heard": "ZZZ", "ticker_guess": "ZZZ", "setup": "bull_flag",
             "setup_quote": "I bought the bull flag at the high of day for a big win",
             "why_entered": [{"reason": "x", "quote": "entirely invented reason for entering this trade"}]},
        ],
        "passed_on": [
            {"ticker_heard": "glnd", "ticker_guess": "GLND", "reason": "grinder",
             "quote": "not an easy one to day trade because it's very thickly traded"},
            {"ticker_heard": "pmax", "ticker_guess": "PMAX", "reason": "halts",
             "quote": "you cannot possibly manage your risk on these halting stocks"},
        ],
        "rules_stated": [{"rule": "avoid topping tails",
                          "quote": "avoid stocks that are giving us those big topping tails"},
                         {"rule": "fake", "quote": "always buy the first red candle after a halt resume"}],
    }


def test_verify_drops_unquoted_items_and_snaps_the_rest():
    st = Counter()
    rec = verify_record(_raw(), TX, st)
    assert [t["ticker_guess"] for t in rec["trades"]] == ["APUS"]       # ZZZ had no real quote -> dropped
    t = rec["trades"][0]
    assert t["entry_price"] is None and t["outcome"] == "win"          # 5.50 is never said -> dropped
    assert [w["reason"] for w in t["why_entered"]] == ["dip"]
    assert t["stock_context"]["catalyst"] is None                       # its quote was invented
    assert t["stock_context"]["other"] == ["topping tails"]
    assert [p["ticker_guess"] for p in rec["passed_on"]] == ["GLND"]
    assert rec["passed_on"][0]["quote"].startswith("not an easy one to day trade")
    assert len(rec["rules_stated"]) == 1
    assert rec["day"]["net_pnl"] == 2992.62 and rec["day"]["market_temperature"] == "cold"
    assert rec["day"]["market_quote"] is None
    assert st["dropped_trades"] == 1 and st["dropped_passed_on"] == 1 and st["dropped_rules"] == 1


def test_clean_ticker():
    assert clean_ticker("$apus") == "APUS"
    assert clean_ticker("K I T T") == "KITT"
    assert clean_ticker("null") is None and clean_ticker("Greenland stocks") is None


# ---------------------------------------------------------------- chunking
def test_short_bundle_is_one_chunk():
    b = "video_id: x\nTRANSCRIPT (mm:ss text):\n[00:01] hi"
    assert chunk_bundle(b, 10_000) == [b]


def test_long_bundle_chunks_keep_header_and_overlap():
    head = "video_id: x\nupload_date: 20260924\nTRANSCRIPT (mm:ss text):"
    lines = [f"[{i // 60:02d}:{i % 60:02d}] line number {i} " + "x" * 80 for i in range(200)]
    b = head + "\n" + "\n".join(lines)
    chunks = chunk_bundle(b, 4000, overlap_lines=3)
    assert len(chunks) > 3
    assert all(c.startswith(head) for c in chunks)
    body = [c.split("TRANSCRIPT (mm:ss text):\n", 1)[1].splitlines() for c in chunks]
    assert body[0][-3:] == body[1][:3]                                   # overlap carried over
    seen = {ln for part in body for ln in part}
    assert seen == set(lines)                                            # nothing lost


# ---------------------------------------------------------------- merge
def _trade(tk, setup_q, reasons=(), **kw):
    return {"ticker_heard": tk, "ticker_guess": tk, "setup": kw.pop("setup", "other"), "setup_quote": setup_q,
            "entry_price": kw.pop("entry_price", None), "exit_price": None, "outcome": kw.pop("outcome", None),
            "stock_context": {"catalyst": None, "other": kw.pop("other", [])},
            "why_entered": [{"reason": r, "quote": q} for r, q in reasons], "tape_reads": [], "why_exited": [],
            "mistake_or_lesson": []}


def test_merge_dedupes_the_overlap_and_unions_fields():
    a = {"video_kind": "daily_recap", "session_date": "2026-09-24",
         "day": {"net_pnl": 100.0, "market_temperature": "cold", "market_quote": None},
         "trades": [_trade("APUS", "i bought the dip right there", [("dip", "i bought the dip right there")],
                           other=["grinder"])],
         "passed_on": [{"ticker_guess": "GLND", "ticker_heard": "glnd", "reason": "r", "quote": "q1"},
                       {"ticker_guess": "PFSA", "ticker_heard": "pfsa", "reason": "r", "quote": "q2"}],
         "rules_stated": [{"rule": "a", "quote": "cut losses quickly and let winners run"}]}
    b = {"video_kind": "daily_recap", "session_date": None,
         "day": {"net_pnl": None, "market_temperature": "cold", "market_quote": "slow market today"},
         "trades": [_trade("APUS", "i bought the dip right there and it rallied", entry_price=5.5, outcome="win",
                           other=["sympathy"]),
                    _trade("PFSA", "curl through vwap", setup="vwap_break")],
         "passed_on": [{"ticker_guess": "GLND", "ticker_heard": "glnd", "reason": "r", "quote": "q1 again"}],
         "rules_stated": [{"rule": "a", "quote": "cut losses quickly and let winners run"},
                          {"rule": "b", "quote": "size down when the market is cold"}]}
    m = merge_records([a, b])
    assert [t["ticker_guess"] for t in m["trades"]] == ["APUS", "PFSA"]
    apus = m["trades"][0]
    assert apus["entry_price"] == 5.5 and apus["outcome"] == "win"          # filled from the later chunk
    assert apus["stock_context"]["other"] == ["grinder", "sympathy"]
    assert [p["ticker_guess"] for p in m["passed_on"]] == ["GLND", "PFSA"]   # GLND deduped across chunks
    assert len(m["rules_stated"]) == 2
    assert m["day"] == {"net_pnl": 100.0, "market_temperature": "cold", "market_quote": "slow market today"}
    assert m["session_date"] == "2026-09-24"


def test_merge_keeps_distinct_trades_on_the_same_ticker():
    a = {"video_kind": "daily_recap", "day": {}, "passed_on": [], "rules_stated": [],
         "trades": [_trade("APUS", "first entry on the dip at the open")]}
    b = {"video_kind": "daily_recap", "day": {}, "passed_on": [], "rules_stated": [],
         "trades": [_trade("APUS", "then i got back in for the break through seven")]}
    assert len(merge_records([a, b])["trades"]) == 2


def test_merge_temperature_conflict_becomes_mixed():
    parts = [{"video_kind": "daily_recap", "day": {"market_temperature": t}, "trades": [], "passed_on": [],
              "rules_stated": []} for t in ("hot", "cold")]
    assert merge_records(parts)["day"]["market_temperature"] == "mixed"


def test_spoken_prices_are_rescaled_for_small_caps_only():
    b = ("OUR SCANNER LIST (symbol, gap %, price) by session date:\n"
         "  2026-06-23: ICCM 80.1% $7.2, BIG 5.0% $150.0\nTRANSCRIPT (mm:ss text):\n")
    px = scanner_prices(b)
    assert px == {"ICCM": 7.2, "BIG": 150.0}
    rec = {"trades": [{"ticker_guess": "ICCM", "entry_price": 699.0, "exit_price": 7.5, "stated_high": None,
                       "stated_low": None},
                      {"ticker_guess": "BIG", "entry_price": 151.0}]}
    assert fix_spoken_prices(rec, px) == 1
    assert rec["trades"][0]["entry_price"] == 6.99 and rec["trades"][1]["entry_price"] == 151.0


def test_paragraph_quotes_are_verified_then_cut_to_a_pointer():
    from youtube_claims.recap.extract import MAX_QUOTE_WORDS, _snap
    words = [f"w{i}" for i in range(120)]
    tx = Transcript(" ".join(words))
    got = _snap(tx, " ".join(words[5:100]))
    assert got == " ".join(words[5:5 + MAX_QUOTE_WORDS])


def test_passed_on_names_must_be_said():
    from youtube_claims.recap.extract import name_mentioned
    tx = Transcript("Then K I T T popped up and Greenland stocks ran but I passed on PMAX")
    assert name_mentioned(tx, "Kit", "KITT")               # spelled out
    assert name_mentioned(tx, "Greenland stocks", None)    # heard phrase
    assert name_mentioned(tx, "p max", "PMAX")
    assert not name_mentioned(tx, "benf", "BENF")          # only on the scanner list


def test_prices_must_be_spoken_in_the_transcript():
    from youtube_claims.recap.extract import price_mentioned
    tx = Transcript("I got in at 550 and sold at $6.25, the high was 7 and it traded 1,200 shares")
    assert price_mentioned(tx, 5.5) and price_mentioned(tx, 6.25) and price_mentioned(tx, 7.0)
    assert not price_mentioned(tx, 8.4)
    rec = verify_record({"trades": [{"ticker_guess": "ABC", "setup_quote": "I got in at 550 and sold at",
                                     "entry_price": 5.5, "exit_price": 9.1}]}, tx)
    assert rec["trades"][0]["entry_price"] == 5.5 and rec["trades"][0]["exit_price"] is None