"""Shared long-only order sizing and cash-funded execution for replay and trading."""

from __future__ import annotations

import math

import numpy as np
import pandas as pd

ENGINE_VERSION = "cash-ledger-pit-v3"

# A held US stock that stops printing (cash acquisition, delisting) is sold at
# its last print once it has gone this many sessions without one and is no
# longer an index member; a halted member stays frozen. Without this the
# position stays at its last mark forever and the capital is never
# reinvested. A-shares are excluded: a suspended CSI 300 name can leave the
# index while it still cannot be sold.
SETTLE_AFTER_SESSIONS = 5


def portfolio_value(holdings: dict, prices: dict, cash: float) -> float:
    return cash + sum(shares * prices.get(t, 0.0) for t, shares in holdings.items())


def delisting_settlements(market, holdings: dict, sessions_without_print: dict,
                          members: pd.Series | None) -> set:
    """Held names to sell at their last print today. ``members`` is the day's
    index membership; tickers outside it (ETFs) are never settled."""
    if market != "us" or members is None:
        return set()
    return {t for t in holdings
            if sessions_without_print.get(t, 0) >= SETTLE_AFTER_SESSIONS and not members.get(t, True)}


def lot_rule(ticker: str, lot_size: int) -> tuple[int, int]:
    # STAR orders start at 200 shares, then increase by one (SSE rule 6.7).
    return (1, 200) if lot_size == 100 and ticker.startswith("688") and ticker.endswith(".SS") else (lot_size, lot_size)


def market_permissions(market, day, raw: pd.Series, previous: pd.Series, volume: pd.Series | None = None):
    fresh = raw.notna() & np.isfinite(raw) & (raw > 0)
    if volume is not None:
        fresh &= volume.reindex(raw.index).gt(0)
    buy, sell = fresh.copy(), fresh.copy()
    tax = 0.0
    if market == "cn":
        limits = pd.Series(.10, index=raw.index)
        for t in raw.index:
            if str(t).startswith("688") or (str(t).startswith("30") and day >= pd.Timestamp("2020-08-24")):
                limits[t] = .20
        change = raw / previous.reindex(raw.index) - 1
        # A closing print at the limit does not establish an available fill.
        buy &= ~(change >= limits - .002)
        sell &= ~(change <= -limits + .002)
        # MOF announcement 2023-39: half-rate from 2023-08-28, sells only.
        tax = .0005 if day >= pd.Timestamp("2023-08-28") else .001
    return set(raw.index[buy]), set(raw.index[sell]), tax


def _allocate_lots(desired: dict, holdings: dict, prices: dict, cash: float,
                   lot_size: int, reserve: float) -> dict:
    """Whole-lot share targets: round each move toward zero, then spend the
    remaining cash one lot at a time on the names furthest below target,
    as long as the lot brings them nearer (nearest-lot rounding)."""
    final, spend = {}, 0.0
    for t, want in desired.items():
        cur = holdings.get(t, 0.0)
        step, minimum = lot_rule(t, lot_size)
        if want <= 0:
            final[t] = 0.0                          # full exit, odd lot included
        else:
            delta = int((want - cur) / step) * step
            final[t] = cur + (delta if abs(delta) >= minimum else 0)
        spend += (final[t] - cur) * prices[t]
    budget = cash - spend - reserve
    short = sorted(((desired[t] - final[t]) * prices[t], t) for t in desired if desired[t] > 0)
    for gap, t in reversed(short):
        step, minimum = lot_rule(t, lot_size)
        extra = minimum if final[t] == holdings.get(t, 0) else step
        lot_value = extra * prices[t]
        if gap > 0.5 * lot_value and lot_value <= budget:
            final[t] += extra
            budget -= lot_value
    return final


def compute_trades(holdings: dict, cash: float, target_weights: dict,
                   prices: dict, min_trade_value: float = 50.0,
                   integer_shares: bool = False, lot_size: int = 1,
                   frozen: set | None = None, resize_band: float = 0.0) -> list[dict]:
    """
    Compute trades needed to move from current holdings to target weights.

    If integer_shares=True, positions move in whole multiples of ``lot_size``
    (e.g. 100-share A-share board lots): each move is rounded toward zero and
    leftover cash then buys one more lot for the names furthest below target
    when that lot brings them nearer; a position whose target is zero is
    always sold in full, odd lot included.
    A held position with a nonzero target is not resized while it is within
    ``resize_band`` (fraction of its target value) of the target.
    Tickers in ``frozen`` (no fresh price today) are neither bought nor sold.

    Returns list of {ticker, shares_delta, direction, est_value}.
    """
    weights = np.array(list(target_weights.values()), dtype=float)
    if not np.isfinite(weights).all() or (weights < 0).any() or weights.sum() > 1 + 1e-8:
        raise ValueError("Targets must be finite, long-only and cash-funded (sum <= 1)")
    total = portfolio_value(holdings, prices, cash)
    frozen = frozen or set()
    desired = {}
    for ticker in sorted(set(target_weights) | set(holdings)):
        if ticker in frozen:
            continue
        price = prices.get(ticker, 0.0)
        if not math.isfinite(price) or price <= 0:
            continue
        target_weight = target_weights.get(ticker, 0.0)
        target_value = total * target_weight
        current_shares = holdings.get(ticker, 0.0)
        if (current_shares > 0 and target_weight > 0
                and abs(target_value - current_shares * price) < resize_band * target_value):
            continue
        desired[ticker] = target_value / price

    targets = (_allocate_lots(desired, holdings, prices, cash, lot_size,
                              reserve=0.002 * total + 10.0 * len(desired))
               if integer_shares else desired)

    raw = []
    for ticker, target_shares in targets.items():
        price = prices[ticker]
        current_shares = holdings.get(ticker, 0.0)
        delta = target_shares - current_shares
        trade_value = abs(delta * price)
        if trade_value < min_trade_value and target_shares > 0:
            continue
        if abs(delta) < 1e-10:
            continue
        if integer_shares and abs(delta) < 1:
            continue

        raw.append({
            "ticker": ticker,
            "shares_delta": int(round(delta)) if integer_shares else delta,
            "direction": "BUY" if delta > 0 else "SELL",
            "est_value": round(trade_value, 2),
            "price": round(price, 6),
            "target_shares": int(round(target_shares)) if integer_shares else round(target_shares, 4),
            "current_shares": int(current_shares) if integer_shares else round(current_shares, 4),
        })

    return raw


def _per_trade_commission(
    shares: float,
    price: float,
    tx_cost: float,
    fixed_fee: float,
    fee_base: float = 0.0,
    fee_per_share: float = 0.0,
) -> float:
    """Commission for one trade.

    Matches the IBKR-style tiered formula used by the backtest engine:
        commission = bps_cost + max(fixed_fee, fee_base + fee_per_share * shares)
    With fee_base=0 and fee_per_share=0 this degenerates to the flat
    fixed-fee model (legacy behavior).
    """
    bps_cost = abs(shares) * price * tx_cost
    per_trade = fee_base + fee_per_share * abs(shares)
    floor = max(fixed_fee, per_trade)
    return bps_cost + floor


def execute_trades(state: dict, trades: list[dict], prices: dict,
                   tx_cost: float = 0.001, fixed_fee: float = 0.0,
                   fee_base: float = 0.0, fee_per_share: float = 0.0,
                   trade_date: str = "", integer_shares: bool = False,
                   lot_size: int = 1, sell_tax: float = 0.0) -> None:
    """Execute trades: update holdings and cash in state, append to trade_log.

    Sells release cash before buys. Every buy must fit cash after all fees,
    including fractional-share orders.

    Per-trade commission supports both the legacy flat ``fixed_fee`` and
    the IBKR-style tiered ``max(fixed_fee, fee_base + fee_per_share*shares)``
    schedule used by the backtest engine.
    """
    holdings = state["holdings"]
    cash = state["cash"]

    # Sell first so buys cannot implicitly borrow against unsettled target changes.
    sells = [t for t in trades if t["shares_delta"] < 0]
    buys = [t for t in trades if t["shares_delta"] > 0]
    ordered = sells + buys

    us_fee_per_share = fee_per_share
    for trade in ordered:
        ticker = trade["ticker"]
        delta = trade["shares_delta"]
        price = prices.get(ticker)
        if price is None or not math.isfinite(price) or price <= 0:
            continue
        if delta < 0:
            delta = -min(-delta, holdings.get(ticker, 0.0))
        if abs(delta) < 1e-10:
            continue
        cost = abs(delta * price)
        # The per-share schedule applies to US-listed shares only; HK lines
        # (often < $1/share) pay the proportional cost and the ticket floor.
        fee_per_share = 0.0 if ticker.endswith(".HK") else us_fee_per_share
        commission = _per_trade_commission(
            abs(delta), price, tx_cost, fixed_fee, fee_base, fee_per_share,
        )

        if delta > 0:
            # BUY — resize to the cash available after fees.
            if (cost + commission) > cash:
                # Try buying fewer shares that we can afford, accounting for
                # the per-share variable component of the commission.
                affordable_price_unit = price * (1 + tx_cost) + fee_per_share
                if affordable_price_unit <= 0:
                    continue
                # Both branches of max(floor, base + per-share fee) must fit.
                affordable = min((cash - fixed_fee) / (price * (1 + tx_cost)),
                                 (cash - fee_base) / affordable_price_unit)
                if integer_shares:
                    step, minimum = lot_rule(ticker, lot_size)
                    affordable = max(0, math.floor(affordable / step)) * step
                    if affordable < minimum:
                        continue
                else:
                    affordable = max(0.0, affordable * (1 - 1e-12))
                if affordable <= 1e-10:
                    continue
                delta = affordable
                cost = abs(delta * price)
                commission = _per_trade_commission(
                    delta, price, tx_cost, fixed_fee, fee_base, fee_per_share,
                )
            cash -= (cost + commission)
            holdings[ticker] = holdings.get(ticker, 0.0) + delta
        else:
            # SELL
            commission += cost * sell_tax
            if commission > cash + cost:
                continue
            cash += (cost - commission)
            holdings[ticker] = holdings.get(ticker, 0.0) + delta  # delta is negative

        # Remove zero holdings
        if ticker in holdings and abs(holdings[ticker]) < (0.5 if integer_shares else 1e-10):
            del holdings[ticker]

        state["trade_log"].append({
            "date": trade_date,
            "action": "BUY" if delta > 0 else "SELL",
            "ticker": ticker,
            "shares": abs(delta),
            "price": round(price, 6),
            "value": cost,
            "commission": commission,
        })

    if cash < -1e-7:
        raise ValueError("Execution would create an unfunded cash balance")
    state["cash"] = max(cash, 0.0)
    if integer_shares:
        state["holdings"] = {k: int(round(v)) for k, v in holdings.items() if abs(v) >= 0.5}
    else:
        state["holdings"] = {k: v for k, v in holdings.items() if abs(v) >= 1e-10}


def rebalance(state: dict, targets: dict, prices: dict, fresh: set, *,
              transaction_cost: float = 0.001, fixed_fee: float = 0.0,
              fee_base: float = 0.0, fee_per_share: float = 0.0,
              min_trade_value: float = 50.0, resize_band: float = 0.10,
              integer_shares: bool = False, lot_size: int = 1,
              trade_date: str = "", can_buy: set | None = None,
              can_sell: set | None = None, sell_tax: float = 0.0) -> list[dict]:
    value = portfolio_value(state["holdings"], prices, state["cash"])
    floor = max(min_trade_value, .001 * value) if min_trade_value else 0.0
    frozen = (set(targets) | set(state["holdings"])) - fresh
    orders = compute_trades(state["holdings"], state["cash"], targets, prices,
                            floor, integer_shares, lot_size, frozen, resize_band)
    orders = [o for o in orders
              if ((o["shares_delta"] > 0 and (can_buy is None or o["ticker"] in can_buy))
                  or (o["shares_delta"] < 0 and (can_sell is None or o["ticker"] in can_sell)))]
    first = len(state["trade_log"])
    execute_trades(state, orders, prices, transaction_cost, fixed_fee, fee_base,
                   fee_per_share, trade_date, integer_shares, lot_size, sell_tax)
    return state["trade_log"][first:]


def simulate_book(targets: pd.DataFrame, prices: pd.DataFrame, capital: float,
                  *, fresh_prices: pd.DataFrame | None = None,
                  contributions: pd.Series | None = None,
                  cash_returns: pd.Series | None = None,
                  can_buy: pd.DataFrame | None = None, can_sell: pd.DataFrame | None = None,
                  sell_tax: pd.Series | None = None, market: str | None = None,
                  previous_prices: pd.DataFrame | None = None,
                  volume: pd.DataFrame | None = None,
                  member: pd.DataFrame | None = None, **options):
    """Targets are executed at their row's close; stale marks never authorize
    fills, except a delisting settlement (``delisting_settlements``, which
    needs the point-in-time ``member`` panel)."""
    raw = prices if fresh_prices is None else fresh_prices.reindex_like(prices)
    marks = prices.ffill()
    previous = marks.shift(1) if previous_prices is None else previous_prices.reindex_like(prices)
    targets = targets.reindex_like(prices).fillna(0.0)
    flows = (contributions.reindex(prices.index).fillna(0.0) if contributions is not None
             else pd.Series(0.0, index=prices.index))
    rates = (cash_returns.reindex(prices.index).fillna(0.0) if cash_returns is not None
             else pd.Series(0.0, index=prices.index))
    state = {"holdings": {}, "cash": float(capital), "trade_log": []}
    equity, units, values, cash = [], [], [], []
    unit_count = float(capital)
    without_print = {}
    for i, day in enumerate(prices.index):
        quote = marks.iloc[i].dropna().to_dict()
        state["cash"] *= 1 + float(rates.iloc[i])
        before = portfolio_value(state["holdings"], quote, state["cash"])
        flow = float(flows.iloc[i])
        if flow < 0:
            raise ValueError("Only nonnegative contributions are supported")
        unit_count += flow / (before / unit_count) if before > 0 else 0.0
        state["cash"] += flow
        fresh = set(raw.columns[(raw.iloc[i].notna() & (raw.iloc[i] > 0)).to_numpy()])
        allowed_buy, allowed_sell, tax = market_permissions(market, day, raw.iloc[i], previous.iloc[i],
                                                              volume.loc[day] if volume is not None else None)
        without_print = {t: 0 if t in fresh else without_print.get(t, 0) + 1 for t in state["holdings"]}
        settle = delisting_settlements(market, state["holdings"], without_print,
                                       member.loc[day] if member is not None else None)
        row = targets.iloc[i]
        wanted = {t: w for t, w in row[row != 0].items() if t not in settle}
        log = rebalance(state, wanted, quote, fresh | settle,
                        can_buy=set(can_buy.columns[can_buy.loc[day]]) if can_buy is not None else allowed_buy,
                        can_sell=(set(can_sell.columns[can_sell.loc[day]]) if can_sell is not None
                                  else allowed_sell) | settle,
                        sell_tax=float(sell_tax.loc[day]) if sell_tax is not None else tax,
                        trade_date=str(day.date()), **options)
        for trade in log:
            if trade["ticker"] in settle:
                trade["settled"] = True
        value = portfolio_value(state["holdings"], quote, state["cash"])
        equity.append(value)
        units.append(value / unit_count * capital)
        cash.append(state["cash"])
        values.append(before)
    return {"equity": pd.Series(equity, index=prices.index),
            "unit_equity": pd.Series(units, index=prices.index),
            "cash": pd.Series(cash, index=prices.index),
            "pretrade_equity": pd.Series(values, index=prices.index),
            "trade_log": state["trade_log"], "holdings": state["holdings"]}