"""Continuous → tradable position discretization and exposure repair.

Pure-numpy, no I/O. Two steps:

1. :func:`round_to_valid_lot` — snap continuous target shares to the nearest
   *valid resting position* given the per-name lot rule AND the current holding
   (``prev_shares``). Rounding is state-dependent: a target below the board
   minimum cannot *open* a fresh lot, but an existing holding may always be
   reduced to 0, and a 科创 odd-lot residual may be sold whole.

2. :func:`repair_exposure` — a two-stage greedy that drives net exposure to ~0
   (Stage A) and gross exposure to the booksize (Stage B) while minimizing the
   dollar-space tracking error ``sum((v_i - v_target_i)**2)``. Splitting the two
   stages avoids the oscillation a single mixed loop suffers (gross repair
   breaking net neutrality and vice versa). O(N log N) via lazy heaps.

A "move" adjusts one name by ±``increment`` shares (or closes it to 0 at the
lattice boundary); names are never opened during repair and never flip sign.
"""

from __future__ import annotations

import heapq
import itertools

import numpy as np


def round_to_valid_lot(
    target: np.ndarray,
    prev_shares: np.ndarray,
    min_open: np.ndarray,
    increment: np.ndarray,
    sell_is_odd_full: np.ndarray | None = None,
) -> np.ndarray:
    """Snap continuous target shares to valid integer resting positions.

    The valid resting lattice for a name is ``{0} ∪ {min_open + k·increment :
    k ≥ 0}`` on each side. Rounding depends on the current holding:

    * **Opening** (no holding on the target side) a magnitude below ``min_open``
      is not allowed → snaps to 0.
    * **Holding** the same side, a sub-minimum target snaps to the nearer of
      ``{0, min_open}`` (so a reduction can rest at the minimum lot or close out).
    * A full liquidation to 0 is always valid (covers the 科创 odd-lot sell:
      a residual ``< min_open`` can only be sold whole, i.e. to 0).
    * Sign is never flipped unless ``target`` itself flipped sign.

    Args:
        target: Continuous signed target shares, length N.
        prev_shares: Current signed integer holding, length N.
        min_open: Per-name minimum open size, length N.
        increment: Per-name share increment (> 0), length N.
        sell_is_odd_full: Unused for resting validity (odd-lot sells already
            resolve to 0); accepted for API symmetry and documentation.

    Returns:
        ``int64`` array of valid resting positions, length N.
    """
    target = np.asarray(target, dtype=np.float64)
    prev = np.asarray(prev_shares, dtype=np.int64)
    min_open = np.asarray(min_open, dtype=np.float64)
    increment = np.asarray(increment, dtype=np.float64)

    sign = np.sign(target)
    mag = np.abs(target)

    # Lattice magnitude for mag >= min_open: min_open + round((mag-min_open)/inc)*inc
    k = np.maximum(np.round((mag - min_open) / increment), 0.0)
    lattice_mag = min_open + k * increment

    holding_same_side = (prev != 0) & (np.sign(prev) == sign) & (sign != 0)

    # Sub-minimum handling: opening -> 0; holding same side -> nearer of {0, min_open}.
    sub_min = mag < min_open
    sub_min_mag = np.where(
        holding_same_side & (mag >= 0.5 * min_open), min_open, 0.0
    )
    final_mag = np.where(sub_min, sub_min_mag, lattice_mag)
    rounded = sign * final_mag
    return rounded.astype(np.int64)


def _exposures(q: np.ndarray, price: np.ndarray) -> tuple[np.ndarray, float, float]:
    v = q.astype(np.float64) * price
    return v, float(v.sum()), float(np.abs(v).sum())


def repair_exposure(
    q_round: np.ndarray,
    q_target: np.ndarray,
    price: np.ndarray,
    increment: np.ndarray,
    min_open: np.ndarray,
    prev_shares: np.ndarray,
    sell_is_odd_full: np.ndarray | None = None,
    booksize: float = 1.0,
    net_tol: float = 0.02,
    gross_tol: float = 0.02,
    max_iters: int | None = None,
) -> np.ndarray:
    """Two-stage greedy exposure repair in dollar space.

    Stage A drives ``net = sum(v_i)`` toward 0; Stage B drives ``gross =
    sum(|v_i|)`` toward ``booksize`` using only moves that keep ``|net|`` within
    its tolerance band, so Stage B cannot undo Stage A. Both stages pick, at each
    step, the admissible ±``increment`` move with the lowest tracking-error cost
    per dollar moved (``ΔTE/|Δv|`` where ``ΔTE = 2·Δv·(v_i - v_target_i) +
    Δv²``). Names that round to 0 are never re-opened here.

    Tolerances are fractions of ``booksize`` but floored to the lot granularity:
    with coarse lots (e.g. pre-2023 100-share main-board lots) exact neutrality
    is unreachable, so the floor prevents a deadlock / infinite loop.

    Args:
        q_round: Integer positions from :func:`round_to_valid_lot`, length N.
        q_target: Continuous target shares (the tracking anchor), length N.
        price: Per-name price (yuan), length N.
        increment: Per-name share increment (> 0), length N.
        min_open: Per-name minimum open size, length N.
        prev_shares: Current holding (unused directly; reserved for borrow caps).
        sell_is_odd_full: Reserved; accepted for API symmetry.
        booksize: Target gross exposure ``B``.
        net_tol: Net tolerance as a fraction of ``B``.
        gross_tol: Gross tolerance as a fraction of ``B``.
        max_iters: Hard cap on repair moves (default ``8·N``).

    Returns:
        ``int64`` repaired positions, length N.
    """
    q = np.asarray(q_round, dtype=np.int64).copy()
    price = np.asarray(price, dtype=np.float64)
    increment = np.asarray(increment, dtype=np.int64).astype(np.float64)
    min_open = np.asarray(min_open, dtype=np.int64).astype(np.float64)
    qt = np.asarray(q_target, dtype=np.float64)
    n = len(q)
    if n == 0:
        return q

    vt = np.where(np.isfinite(qt), qt, 0.0) * price  # v_target, NaN-safe
    tradable = np.isfinite(price) & (price > 0)
    step = np.where(tradable, increment * price, np.inf)  # dollar per increment

    if max_iters is None:
        max_iters = 8 * n

    # Adaptive absolute tolerances: never finer than the lot granularity.
    active_step = step[(q != 0) & tradable]
    max_step = float(active_step.max()) if active_step.size else 0.0
    min_step = float(active_step.min()) if active_step.size else 0.0
    net_tol_abs = max(net_tol * booksize, max_step)
    gross_tol_abs = max(gross_tol * booksize, min_step)
    net_band = net_tol_abs  # Stage B keeps |net| within this band

    v, net, gross = _exposures(q, price)

    def _move(i: int, grow: bool):
        """Return (dshares, dv, dte) for a grow/shrink move on name i, or None."""
        if q[i] == 0 or not tradable[i]:
            return None
        s = 1 if q[i] > 0 else -1
        if grow:
            dshares = s * int(increment[i])
        else:
            mag = abs(int(q[i]))
            if mag - increment[i] >= min_open[i]:
                dshares = -s * int(increment[i])
            else:
                dshares = -int(q[i])  # close to 0 (lattice boundary / odd lot)
        if dshares == 0:
            return None
        dv = dshares * price[i]
        dte = 2.0 * dv * (v[i] - vt[i]) + dv * dv
        return dshares, dv, dte

    def _apply(i: int, dshares: int, dv: float):
        nonlocal net, gross
        old_abs = abs(v[i])
        q[i] += dshares
        v[i] += dv
        net += dv
        gross += abs(v[i]) - old_abs

    counter = itertools.count()
    active_idx = np.nonzero((q != 0) & tradable)[0]

    # ---- Stage A: net repair -------------------------------------------------
    def _stageA_dir() -> int:
        return -1 if net > 0 else 1  # desired sign of dv

    iters = 0
    while abs(net) > net_tol_abs and iters < max_iters:
        want = _stageA_dir()  # dv sign we need
        heap: list = []
        best_key: dict[int, float] = {}
        for i in active_idx:
            i = int(i)
            # For net>0 (want dv<0): shrink longs, grow shorts. Mirror otherwise.
            grow = (q[i] < 0) if want < 0 else (q[i] > 0)
            mv = _move(i, grow)
            if mv is None:
                continue
            _, dv, dte = mv
            if np.sign(dv) != want:
                continue
            key = dte / abs(dv)
            best_key[i] = key
            heapq.heappush(heap, (key, next(counter), i, grow))
        if not heap:
            break
        progressed = False
        while heap and abs(net) > net_tol_abs and iters < max_iters:
            key, _, i, grow = heapq.heappop(heap)
            if best_key.get(i) != key:
                continue  # stale
            mv = _move(i, grow)
            if mv is None:
                best_key.pop(i, None)
                continue
            dshares, dv, dte = mv
            if np.sign(dv) != want:
                best_key.pop(i, None)
                continue
            # Don't overshoot net through 0 by more than the tolerance band.
            if abs(net + dv) > abs(net) and abs(net + dv) > net_tol_abs:
                best_key.pop(i, None)
                continue
            _apply(i, dshares, dv)
            iters += 1
            progressed = True
            if q[i] != 0:
                nk = _move(i, grow)
                if nk is not None:
                    _, ndv, ndte = nk
                    if np.sign(ndv) == want:
                        k2 = ndte / abs(ndv)
                        best_key[i] = k2
                        heapq.heappush(heap, (k2, next(counter), i, grow))
                        continue
            best_key.pop(i, None)
        if not progressed:
            break

    # ---- Stage B: gross repair (net-preserving) -----------------------------
    iters = 0
    active_idx = np.nonzero((q != 0) & tradable)[0]
    while abs(gross - booksize) > gross_tol_abs and iters < max_iters:
        grow = gross < booksize  # need more gross → grow magnitudes; else shrink
        heap = []
        best_key = {}
        for i in active_idx:
            i = int(i)
            mv = _move(i, grow)
            if mv is None:
                continue
            _, dv, dte = mv
            # Net-band filter: never push |net| past the band.
            if abs(net + dv) > net_band and abs(net + dv) >= abs(net):
                continue
            key = dte / abs(dv)
            best_key[i] = key
            heapq.heappush(heap, (key, next(counter), i, grow))
        if not heap:
            break
        progressed = False
        while heap and abs(gross - booksize) > gross_tol_abs and iters < max_iters:
            key, _, i, g = heapq.heappop(heap)
            if best_key.get(i) != key:
                continue
            mv = _move(i, g)
            if mv is None:
                best_key.pop(i, None)
                continue
            dshares, dv, dte = mv
            if abs(net + dv) > net_band and abs(net + dv) >= abs(net):
                best_key.pop(i, None)
                continue
            _apply(i, dshares, dv)
            iters += 1
            progressed = True
            if q[i] != 0:
                nk = _move(i, g)
                if nk is not None:
                    _, ndv, ndte = nk
                    if not (abs(net + ndv) > net_band and abs(net + ndv) >= abs(net)):
                        k2 = ndte / abs(ndv)
                        best_key[i] = k2
                        heapq.heappush(heap, (k2, next(counter), i, g))
                        continue
            best_key.pop(i, None)
        if not progressed:
            break

    return q