"""End-to-end correctness invariants for the reversal_5d pipeline (no network).

Each test maps 1:1 to one of the ten review checks. Naming convention: the
*execution date* ``d`` is the market session on which a target is actually
filled at the open; the *signal date* ``t`` is the session whose close formed
that target. The documented convention is ``d = next(t)`` (see
``docs/portfolio_trading_cost_model.md``), so ``close[d-1] == close[t]``.
"""

import numpy as np
import pandas as pd

from pipeline.alpha.compute import compute_alpha, investable_universe_mask
from pipeline.portfolio.construct import construct_positions
from pipeline.portfolio.discretize import repair_exposure, round_to_valid_lot
from pipeline.portfolio.market_rules import MarketRule
from pipeline.portfolio.costs import SimpleProportionalCostModel
from pipeline.portfolio.constraints import (
    SuspensionConstraint,
    VolumeCapConstraint,
)
from pipeline.portfolio.simulator import ReferenceSimulator


_SYMBOLS = ("sh600000", "sz000001", "sh688981", "sz300750")


def _panel(n_days=12, symbols=_SYMBOLS, start="2024-01-01", seed=0,
           distinct_open=True):
    """Contiguous long-format DATA frame with all columns the pipeline needs.

    Open and close differ (so overnight vs intraday PnL terms are separable),
    and the calendar is gap-free so each session is the next session's ``d-1``.
    """
    dates = pd.date_range(start, periods=n_days)
    rng = np.random.default_rng(seed)
    frames = []
    for i, sym in enumerate(symbols):
        close = np.abs(50.0 + i * 10 + np.cumsum(rng.standard_normal(n_days))) + 5.0
        open_ = close * (0.99 + 0.02 * rng.random(n_days)) if distinct_open else close.copy()
        preclose = np.concatenate([[close[0]], close[:-1]])
        frames.append(pd.DataFrame({
            "symbol_id": sym,
            "symbol_name": sym,
            "date": dates,
            "open": open_,
            "high": np.maximum(open_, close),
            "low": np.minimum(open_, close),
            "close": close,
            "preclose": preclose,
            "volume": 1_000_000.0,
            "amount": 1_000_000.0 * close,
            "tradestatus": 1,
            "isST": 0,
        }))
    return pd.concat(frames, ignore_index=True)


# --- 1. reversal signal uses only close[d-1] and earlier ---------------------

def test_reversal_signal_does_not_peek_at_future_closes():
    data = _panel(n_days=12)
    base = compute_alpha(data, "rev", "reversal_rank", lookback=5)

    # Perturb every close strictly AFTER an interior signal date t; the weight
    # dated t (executed at d = t+1) must be unchanged — it may use close[t]
    # (== close[d-1]) and earlier only.
    t = sorted(data["date"].unique())[6]
    future = data.copy()
    mask = future["date"] > t
    future.loc[mask, ["open", "high", "low", "close"]] *= 1.5
    perturbed = compute_alpha(future, "rev", "reversal_rank", lookback=5)

    b = base[base["date"] <= t].set_index(["symbol_id", "date"])["weight"]
    p = perturbed[perturbed["date"] <= t].set_index(["symbol_id", "date"])["weight"]
    pd.testing.assert_series_equal(b.sort_index(), p.sort_index())


# --- 2. the executed (fill/PnL) date is the open-execution date d = next(t) --

def test_fill_date_is_next_session_open_execution_date():
    data = _panel(n_days=8)
    weights = compute_alpha(data, "c", "reversal_rank", lookback=3)
    weights = weights.rename(columns={"alpha_name": "combo_name"})
    weights["combo_name"] = "c"
    pos = construct_positions(weights, data, booksize=1e6, portfolio_name="run1")
    fills, pnl = ReferenceSimulator(cost_bps=5, slippage_bps=5).run(pos, data)

    sessions = sorted(data["date"].unique())
    nxt = {s: sessions[i + 1] for i, s in enumerate(sessions[:-1])}
    # Every executed date equals the session AFTER some position (signal) date.
    pos_dates = set(pos["date"].unique())
    exec_dates = set(pnl["date"].unique())
    assert exec_dates == {nxt[t] for t in pos_dates if t in nxt}

    # Execution price is the open of the execution date, not the signal close.
    opn = data.pivot_table(index="date", columns="symbol_id", values="open",
                           aggfunc="first").sort_index()
    d = sorted(exec_dates)[1]
    row = fills[(fills["date"] == d) & (fills["traded_shares"] != 0)].iloc[0]
    sym = row["symbol_id"]
    expected_cost = abs(row["traded_shares"]) * opn.loc[d, sym] * (5 + 5) / 1e4
    assert np.isclose(row["trade_cost"], expected_cost)


# --- 3. PnL identity: overnight(old book) + intraday(new book) - cost --------

def test_daily_pnl_matches_overnight_plus_intraday_minus_cost():
    data = _panel(n_days=8)
    weights = compute_alpha(data, "c", "reversal_rank", lookback=3)
    weights = weights.rename(columns={"alpha_name": "combo_name"})
    weights["combo_name"] = "c"
    pos = construct_positions(weights, data, booksize=1e6, portfolio_name="run1")
    fills, pnl = ReferenceSimulator(cost_bps=5, slippage_bps=5).run(pos, data)

    opn = data.pivot_table(index="date", columns="symbol_id", values="open", aggfunc="first").sort_index()
    cls = data.pivot_table(index="date", columns="symbol_id", values="close", aggfunc="first").sort_index()
    sessions = list(cls.index)

    prev_close_of = {sessions[i]: sessions[i - 1] for i in range(1, len(sessions))}

    for d in sorted(pnl["date"].unique()):
        day = fills[fills["date"] == d]
        prev = day.set_index("symbol_id")["prev_shares"]
        realized = day.set_index("symbol_id")["realized_shares"]
        cost = day["trade_cost"].sum()

        intraday = float((realized * (cls.loc[d] - opn.loc[d]).reindex(realized.index)).sum())
        # Overnight gap on the OLD book is taken from the previous *executed*
        # date's close. With a gap-free calendar and daily execution that is the
        # immediately preceding session; the first executed date has no prior
        # book so the term is naturally zero (prev_shares == 0 there).
        pc = prev_close_of.get(d)
        if pc is not None and (prev != 0).any():
            overnight = float((prev * (opn.loc[d] - cls.loc[pc]).reindex(prev.index)).sum())
        else:
            overnight = 0.0
        expected = overnight + intraday - cost
        got = float(pnl[pnl["date"] == d]["pnl"].iloc[0])
        assert np.isclose(got, expected, rtol=1e-6, atol=1e-3), (d, got, expected)


# --- 4. realized shares (not target shares) are threaded into the next day ---

def test_realized_not_target_threaded_forward():
    data = _panel(n_days=6)
    weights = compute_alpha(data, "c", "reversal_rank", lookback=2)
    weights = weights.rename(columns={"alpha_name": "combo_name"})
    weights["combo_name"] = "c"
    pos = construct_positions(weights, data, booksize=1e8, portfolio_name="run1")

    # A tight volume cap forces partial fills, so realized != target on most
    # names — exactly the case where threading target vs realized diverges.
    fills, _ = ReferenceSimulator(
        constraints=[VolumeCapConstraint(max_frac=1e-6)]
    ).run(pos, data)

    wide_prev = fills.pivot_table(index="date", columns="symbol_id", values="prev_shares", aggfunc="first")
    wide_real = fills.pivot_table(index="date", columns="symbol_id", values="realized_shares", aggfunc="first")
    exec_dates = list(wide_prev.index)
    assert len(exec_dates) >= 2
    # Today's prev_shares == yesterday's realized_shares for every name.
    for a, b in zip(exec_dates[:-1], exec_dates[1:]):
        prev_today = wide_prev.loc[b].dropna()
        real_yest = wide_real.loc[a].reindex(prev_today.index).fillna(0.0)
        pd.testing.assert_series_equal(
            prev_today.astype(float), real_yest.astype(float), check_names=False
        )
    # And realized actually diverged from target (cap bit), so the test is real.
    assert (fills["realized_shares"] != fills["target_shares"]).any()


# --- 5. blocked trades create zero traded_shares and zero trade_cost ---------

def test_blocked_trade_has_zero_shares_and_zero_cost():
    data = _panel(n_days=6)
    # Suspend one name on every session so any attempt to trade it is blocked.
    data.loc[data["symbol_id"] == "sz000001", "tradestatus"] = 0
    weights = compute_alpha(data, "c", "reversal_rank", lookback=2)
    weights = weights.rename(columns={"alpha_name": "combo_name"})
    weights["combo_name"] = "c"
    pos = construct_positions(weights, data, booksize=1e6, portfolio_name="run1")

    fills, _ = ReferenceSimulator(
        constraints=[SuspensionConstraint()], cost_bps=5, slippage_bps=5
    ).run(pos, data)

    blocked = fills[fills["blocked"] == 1]
    assert (blocked["traded_shares"] == 0).all()
    assert (blocked["trade_cost"] == 0.0).all()
    # The suspended name never trades and never accrues cost.
    susp = fills[fills["symbol_id"] == "sz000001"]
    assert (susp["traded_shares"] == 0).all()
    assert (susp["trade_cost"] == 0.0).all()


# --- 6. liquid universe uses only information known before open[d] -----------

def test_investable_universe_mask_is_causal():
    data = _panel(n_days=14)
    close = data.pivot_table(index="date", columns="symbol_id", values="close", aggfunc="first").sort_index()
    full = investable_universe_mask(data, close, top_n=10, min_history=3)

    t = sorted(data["date"].unique())[8]
    # Recompute the mask from data truncated at the signal date t: the mask row
    # for t must be identical, proving it never reads dates > t (i.e. nothing
    # from open[d=t+1] onward).
    trunc = data[data["date"] <= t]
    close_t = close.loc[:t]
    mask_t = investable_universe_mask(trunc, close_t, top_n=10, min_history=3)
    pd.testing.assert_series_equal(
        full.loc[t].sort_index(), mask_t.loc[t].sort_index(), check_names=False
    )


# --- 7. cost bps is one-way per-trade (a round trip is charged twice) --------

def test_cost_bps_is_one_way_per_trade():
    model = SimpleProportionalCostModel(cost_bps=5, slippage_bps=5)
    price = np.array([20.0])

    buy = model.compute(np.array([1000]), price, np.array([1]), date=None)
    sell = model.compute(np.array([-1000]), price, np.array([-1]), date=None)

    one_way = 1000 * 20 * (5 + 5) / 1e4
    assert np.isclose(buy[0], one_way)      # charged once on the buy leg
    assert np.isclose(sell[0], one_way)     # charged again on the sell leg
    # A full round trip (enter then exit) therefore costs ~2x the one-way rate.
    assert np.isclose(buy[0] + sell[0], 2 * one_way)


# --- 8. execution & PnL use raw tradable prices on the same scale as shares --

def test_position_value_is_shares_times_raw_price():
    data = _panel(n_days=10)
    weights = compute_alpha(data, "c", "reversal_rank", lookback=3)
    weights = weights.rename(columns={"alpha_name": "combo_name"})
    weights["combo_name"] = "c"
    pos = construct_positions(weights, data, booksize=1e6, portfolio_name="run1")

    finite = pos["price"] > 0
    # The stored value is exactly integer shares × the raw construction price —
    # no adjusted-price factor is mixed into the share→value accounting.
    expected = pos.loc[finite, "position_shares"] * pos.loc[finite, "price"]
    pd.testing.assert_series_equal(
        pos.loc[finite, "position_value"].astype(float),
        expected.astype(float),
        check_names=False,
    )


# --- 9. alpha is scale-free (adjusted prices ok); accounting uses raw units --

def test_alpha_weights_invariant_to_per_symbol_price_scaling():
    data = _panel(n_days=12)
    base = compute_alpha(data, "rev", "reversal_rank", lookback=5)

    # A qfq/hfq adjustment is (per symbol) a multiplicative rescaling of the
    # price series; pct_change is scale-free, so the alpha weights must not move.
    scaled = data.copy()
    factor = {"sh600000": 2.0, "sz000001": 0.5, "sh688981": 3.0, "sz300750": 1.25}
    for sym, f in factor.items():
        m = scaled["symbol_id"] == sym
        scaled.loc[m, ["open", "high", "low", "close"]] *= f
    scaled_alpha = compute_alpha(scaled, "rev", "reversal_rank", lookback=5)

    b = base.set_index(["symbol_id", "date"])["weight"].sort_index()
    s = scaled_alpha.set_index(["symbol_id", "date"])["weight"].sort_index()
    pd.testing.assert_series_equal(b, s)


# --- 10. repaired book stays on valid A-share lot lattices -------------------

def _on_lattice(q, min_open, increment):
    q = np.abs(np.asarray(q, dtype=np.int64))
    on = (q == 0) | ((q >= min_open) & ((q - min_open) % increment == 0))
    return bool(on.all())


def test_repair_output_stays_on_lot_lattice():
    # Pre-2023 main board has a 100-share increment (the strongest lattice
    # constraint); STAR uses min 200 / increment 1.
    symbols = np.array(["sh600000", "sz000001", "sh688981", "sz300750"], dtype=object)
    rule = MarketRule()
    on = "2022-06-01"  # pre 2023-08-10 → main-board increment is 100
    min_open, increment, odd_full, _ = rule.get_rules_vectorized(
        symbols, on, np.zeros(len(symbols), dtype=bool)
    )
    assert min_open[0] == 100 and increment[0] == 100  # main board pre-2023

    price = np.array([12.3, 8.7, 45.0, 230.0])
    prev = np.zeros(len(symbols), dtype=np.int64)
    q_target = np.array([3251.0, -7777.0, 640.0, -415.0])

    q_round = round_to_valid_lot(q_target, prev, min_open, increment, odd_full)
    assert _on_lattice(q_round, min_open, increment)

    repaired = repair_exposure(
        q_round, q_target, price, increment, min_open, prev, odd_full,
        booksize=float(np.abs(q_target * price).sum()),
    )
    assert _on_lattice(repaired, min_open, increment)
    # Repair never flips a name's sign relative to the rounded book.
    nz = q_round != 0
    assert np.all(np.sign(repaired[nz]) * np.sign(q_round[nz]) >= 0)