# coding: utf-8
"""JoinQuant hosted wrapper.

Upload the generated bundle, not this source file, to JoinQuant.  The bundle
tool replaces the marked import block with the exact portable core source.
"""

# __PORTABLE_CORE_BUNDLE_START__
# Inlined by tools/bundle_platforms.py; do not edit this region.
import math
import re


_CANONICAL_EXCHANGES = {
    "SH": "XSHG",
    "SSE": "XSHG",
    "XSHG": "XSHG",
    "SZ": "XSHE",
    "SZE": "XSHE",
    "SZSE": "XSHE",
    "XSHE": "XSHE",
    "BJ": "XBSE",
    "BSE": "XBSE",
    "XBSE": "XBSE",
}

_SHORT_EXCHANGES = {
    "XSHG": "SH",
    "XSHE": "SZ",
    "XBSE": "BJ",
}


def _sum_numeric(values):
    """Sum numerics without trusting a hosted runtime's global ``sum`` name."""

    total = 0.0
    for value in values:
        total += value
    return total


def _infer_exchange(code):
    if code.startswith(("4", "8")):
        return "XBSE"
    if code.startswith(("5", "6", "9")):
        return "XSHG"
    return "XSHE"


def _split_symbol(symbol):
    if not isinstance(symbol, str):
        raise TypeError("symbol must be a string")
    value = symbol.strip().upper().replace(" ", "")
    if not value:
        raise ValueError("symbol must not be empty")

    match = re.match(r"^(\d{6})\.(XSHG|XSHE|XBSE|SH|SZ|BJ)$", value)
    if match:
        return match.group(1), _CANONICAL_EXCHANGES[match.group(2)]

    match = re.match(r"^(XSHG|XSHE|XBSE|SH|SZ|BJ)[\.:]?(\d{6})$", value)
    if match:
        return match.group(2), _CANONICAL_EXCHANGES[match.group(1)]

    match = re.match(r"^(\d{6})[\.:](SSE|SZE|SZSE|BSE)$", value)
    if match:
        return match.group(1), _CANONICAL_EXCHANGES[match.group(2)]

    if re.match(r"^\d{6}$", value):
        return value, _infer_exchange(value)
    raise ValueError("unsupported A-share symbol: %s" % symbol)


def normalize_symbol(symbol, target="canonical"):
    """Convert common A-share codes to canonical/JQ, QMT, or Qlib format."""

    code, exchange = _split_symbol(symbol)
    target_name = str(target).strip().lower()
    if target_name in ("canonical", "joinquant", "jq"):
        return "%s.%s" % (code, exchange)
    if target_name in ("qmt", "xtquant"):
        return "%s.%s" % (code, _SHORT_EXCHANGES[exchange])
    if target_name == "qlib":
        return "%s%s" % (_SHORT_EXCHANGES[exchange], code)
    raise ValueError("unsupported target platform: %s" % target)


def convert_symbol(symbol, target="canonical"):
    """Alias kept as the stable adapter-facing name."""

    return normalize_symbol(symbol, target)


def momentum_score(prices, lookback=20, skip=0):
    """Return trailing simple momentum, or ``None`` if history is insufficient.

    ``skip=0`` uses the latest completed decision bar.  Set ``skip=1`` only
    when intentionally implementing a 20-1 convention.  A lookback of N needs
    N+skip+1 observations.
    """

    lookback = int(lookback)
    skip = int(skip)
    if lookback <= 0 or skip < 0:
        raise ValueError("lookback must be positive and skip non-negative")
    values = list(prices)
    if len(values) < lookback + skip + 1:
        return None
    end_index = len(values) - 1 - skip
    start_index = end_index - lookback
    start = float(values[start_index])
    end = float(values[end_index])
    if (
        start <= 0.0
        or end <= 0.0
        or not math.isfinite(start)
        or not math.isfinite(end)
    ):
        return None
    return end / start - 1.0


def rank_momentum(price_history, lookback=20, skip=0, top_n=None):
    """Rank ``{symbol: closes}`` by momentum with deterministic tie-breaking."""

    ranked = []
    for symbol in sorted(price_history):
        score = momentum_score(price_history[symbol], lookback, skip)
        if score is not None and math.isfinite(score):
            ranked.append((normalize_symbol(symbol), score))
    ranked.sort(key=lambda item: (-item[1], item[0]))
    if top_n is not None:
        ranked = ranked[: max(0, int(top_n))]
    return ranked


def capped_target_weights(
    scores, max_weight=0.20, gross_target=1.0, min_score=0.0
):
    """Long-only score-proportional weights with iterative cap redistribution."""

    max_weight = float(max_weight)
    gross_target = float(gross_target)
    min_score = float(min_score)
    if not math.isfinite(max_weight) or not 0.0 < max_weight <= 1.0:
        raise ValueError("max_weight must be finite and lie in (0, 1]")
    if not math.isfinite(gross_target) or not 0.0 <= gross_target <= 1.0:
        raise ValueError("gross_target must be finite and lie in [0, 1]")
    if not math.isfinite(min_score):
        raise ValueError("min_score must be finite")

    if hasattr(scores, "items"):
        items = list(scores.items())
    else:
        items = list(scores)
    normalized = {}
    for symbol, raw_score in items:
        name = normalize_symbol(symbol)
        score = float(raw_score)
        if not math.isfinite(score):
            raise ValueError("score must be finite for %s" % name)
        normalized[name] = score

    weights = dict((symbol, 0.0) for symbol in normalized)
    active = dict(
        (symbol, score)
        for symbol, score in normalized.items()
        if score > min_score and score > 0.0
    )
    remaining = min(gross_target, len(active) * max_weight)

    while active and remaining > 1e-15:
        score_sum = _sum_numeric(active.values())
        if score_sum <= 0.0:
            break
        capped = []
        for symbol in sorted(active):
            proposed = remaining * active[symbol] / score_sum
            if proposed >= max_weight - 1e-15:
                weights[symbol] = max_weight
                capped.append(symbol)
        if not capped:
            for symbol in sorted(active):
                weights[symbol] = remaining * active[symbol] / score_sum
            remaining = 0.0
            break
        for symbol in capped:
            active.pop(symbol)
            remaining -= max_weight
        remaining = max(0.0, remaining)
    return weights


def round_board_lot(quantity, lot_size=100):
    """Round an absolute long position down to a board-lot quantity."""

    lot_size = int(lot_size)
    if lot_size <= 0:
        raise ValueError("lot_size must be positive")
    if not math.isfinite(float(quantity)):
        raise ValueError("quantity must be finite")
    quantity = max(0, int(quantity))
    return quantity // lot_size * lot_size


def _is_star_market(symbol):
    code, exchange = _split_symbol(symbol)
    return exchange == "XSHG" and code.startswith(("688", "689"))


def target_quantities(
    weights, prices, equity, lot_size=100, cash_buffer=0.02
):
    """Convert target weights to affordable board-lot long quantities."""

    equity = float(equity)
    cash_buffer = float(cash_buffer)
    if (
        not math.isfinite(equity)
        or equity < 0.0
        or not math.isfinite(cash_buffer)
        or not 0.0 <= cash_buffer < 1.0
    ):
        raise ValueError("invalid equity or cash_buffer")
    positive_weights = []
    for raw_symbol, raw_weight in weights.items():
        weight = float(raw_weight)
        if not math.isfinite(weight):
            raise ValueError(
                "weight must be finite for %s" % normalize_symbol(raw_symbol)
            )
        positive_weights.append(max(0.0, weight))
    requested_gross = _sum_numeric(positive_weights)
    gross_ceiling = 1.0 - cash_buffer
    scale = (
        min(1.0, gross_ceiling / requested_gross)
        if requested_gross > 0.0
        else 1.0
    )
    result = {}
    for raw_symbol in sorted(weights):
        symbol = normalize_symbol(raw_symbol)
        weight = float(weights[raw_symbol])
        if not math.isfinite(weight):
            raise ValueError("weight must be finite for %s" % symbol)
        weight = max(0.0, weight)
        price = prices.get(raw_symbol)
        if price is None:
            price = prices.get(symbol)
        if price is None:
            raise KeyError("missing price for %s" % symbol)
        price = float(price)
        if price <= 0.0 or not math.isfinite(price):
            raise ValueError("invalid price for %s" % symbol)
        # Weights are absolute equity weights.  cash_buffer is a ceiling, not
        # a second multiplicative haircut: gross=.95 and buffer=.02 stays .95.
        budget = equity * weight * scale
        quantity = round_board_lot(budget / price, lot_size)
        # STAR Market auction orders must contain at least 200 shares.  Keep
        # the portable planner conservative and on the configured 100-share
        # grid even though quantities above 200 may increment by one share.
        if _is_star_market(symbol) and 0 < quantity < 200:
            quantity = 0
        result[symbol] = quantity
    return result


def order_deltas(
    targets, current, sellable=None, lot_size=100
):
    """Return signed board-lot orders; negative sells respect sellable quantity."""

    canonical_targets = {}
    canonical_current = {}
    canonical_sellable = {}
    for symbol, quantity in targets.items():
        canonical_targets[normalize_symbol(symbol)] = max(0, int(quantity))
    for symbol, quantity in current.items():
        canonical_current[normalize_symbol(symbol)] = max(0, int(quantity))
    if sellable is not None:
        for symbol, quantity in sellable.items():
            canonical_sellable[normalize_symbol(symbol)] = max(0, int(quantity))

    result = {}
    symbols = sorted(set(canonical_targets) | set(canonical_current))
    for symbol in symbols:
        target = round_board_lot(canonical_targets.get(symbol, 0), lot_size)
        if _is_star_market(symbol) and 0 < target < 200:
            target = 0
        held = canonical_current.get(symbol, 0)
        delta = target - held
        if delta > 0:
            delta = round_board_lot(delta, lot_size)
            if _is_star_market(symbol) and 0 < delta < 200:
                delta = 0
        elif delta < 0:
            available = held
            if sellable is not None:
                available = min(held, canonical_sellable.get(symbol, 0))
            requested = min(-delta, available)
            # A complete liquidation may submit the entire odd-lot balance.
            # This also covers the STAR exception for a remaining balance
            # below 200 shares.  Partial orders stay on the configured grid.
            full_liquidation = (
                target == 0
                and requested == held
                and available >= held
            )
            if full_liquidation:
                delta = -held
            else:
                delta = -round_board_lot(requested, lot_size)
                if _is_star_market(symbol) and 0 < -delta < 200:
                    delta = 0
        if delta:
            result[symbol] = delta
    return result


def build_rebalance_plan(
    price_history,
    current,
    sellable,
    equity,
    lookback=20,
    skip=0,
    top_n=10,
    max_weight=0.20,
    gross_target=1.0,
    cash_buffer=0.02,
    lot_size=100,
):
    """Build the complete portable signal -> weight -> quantity -> order plan."""

    ranked = rank_momentum(price_history, lookback, skip, top_n)
    scores = dict(ranked)
    weights = capped_target_weights(scores, max_weight, gross_target, 0.0)
    for raw_symbol in price_history:
        symbol = normalize_symbol(raw_symbol)
        if symbol not in scores:
            scores[symbol] = momentum_score(
                price_history[raw_symbol], lookback, skip
            )
        if symbol not in weights:
            weights[symbol] = 0.0
    prices = {}
    for raw_symbol, history in price_history.items():
        if not history:
            continue
        prices[normalize_symbol(raw_symbol)] = history[-1]
    targets = target_quantities(
        weights, prices, equity, lot_size, cash_buffer
    )
    orders = order_deltas(targets, current, sellable, lot_size)
    return {
        "scores": scores,
        "weights": weights,
        "targets": targets,
        "orders": orders,
    }


__all__ = [
    "build_rebalance_plan",
    "capped_target_weights",
    "convert_symbol",
    "momentum_score",
    "normalize_symbol",
    "order_deltas",
    "rank_momentum",
    "round_board_lot",
    "target_quantities",
]
# __PORTABLE_CORE_BUNDLE_END__

import math


class PriceHistoryUnavailable(RuntimeError):
    pass


class UnmanagedPositionError(RuntimeError):
    pass


class ClockStateError(RuntimeError):
    pass


# __BASELINE_CONFIG_START__
# Generated from configs/baseline.json; do not edit this region.
CONFIG = {'universe_mode': 'pit_index',
 'index_symbol': '000905.XSHG',
 'qmt_sector_name': '\u4e2d\u8bc1500',
 'dedicated_account_required': True,
 'exclude_st': True,
 'universe': ['600000.XSHG', '000001.XSHE', '300750.XSHE', '000333.XSHE'],
 'lookback': 20,
 'skip': 0,
 'top_n': 20,
 'max_weight': 0.05,
 'gross_target': 0.95,
 'cash_buffer': 0.02,
 'lot_size': 100,
 'participation_rate': 0.1,
 'slippage_bps': 2.0,
 'commission_rate': 0.0002,
 'minimum_commission': 5.0,
 'stamp_duty_rate': 0.0005,
 'transfer_fee_rate': 1e-05,
 'rebalance_schedule': 'weekly_first_close'}
# __BASELINE_CONFIG_END__


def initialize(context):
    """Configure a daily-open clock that executes a first-week-close signal."""
    set_option("avoid_future_data", True)
    set_option("use_real_price", True)
    g.quant60_config = dict(CONFIG)
    if g.quant60_config.get("rebalance_schedule") != "weekly_first_close":
        raise ValueError("unsupported rebalance_schedule")
    if g.quant60_config.get("universe_mode") not in ("pit_index", "fixed"):
        raise ValueError("unsupported universe_mode")
    if (
        g.quant60_config.get("universe_mode") == "pit_index"
        and not g.quant60_config.get("dedicated_account_required")
    ):
        raise ValueError("pit_index mode requires a dedicated account")
    set_benchmark(g.quant60_config["index_symbol"])
    commission = (
        float(g.quant60_config["commission_rate"])
        + float(g.quant60_config["transfer_fee_rate"])
    )
    set_order_cost(
        OrderCost(
            open_tax=0.0,
            close_tax=float(g.quant60_config["stamp_duty_rate"]),
            open_commission=commission,
            close_commission=commission,
            close_today_commission=0.0,
            min_commission=float(
                g.quant60_config["minimum_commission"]
            ),
        ),
        type="stock",
    )
    # JoinQuant applies half of PriceRelatedSlippage on each side. The local
    # slippage_bps setting is one-sided, hence the explicit factor of two.
    set_slippage(
        PriceRelatedSlippage(
            2.0 * float(g.quant60_config["slippage_bps"]) / 10000.0
        ),
        type="stock",
    )
    set_option(
        "order_volume_ratio",
        float(g.quant60_config["participation_rate"]),
    )
    g.quant60_last_plan = None
    g.quant60_last_universe = None
    g.quant60_last_universe_as_of = None
    g.quant60_skipped_orders = []
    g.quant60_pending_first_session = None
    g.quant60_last_callback_date = None
    run_daily(rebalance, time="open")


def _close_history(symbol, count):
    raw = attribute_history(
        symbol,
        count,
        unit="1d",
        fields=["close"],
        skip_paused=False,
        df=False,
        fq="pre",
    )
    values = list(raw.get("close", [])) if hasattr(raw, "get") else []
    if len(values) != int(count):
        raise PriceHistoryUnavailable(
            "%s needs %d completed daily closes, got %d"
            % (symbol, int(count), len(values))
        )
    output = []
    for value in values:
        try:
            number = float(value)
        except (TypeError, ValueError):
            raise PriceHistoryUnavailable(
                "%s contains a missing/non-numeric close" % symbol
            )
        if not math.isfinite(number) or number <= 0.0:
            raise PriceHistoryUnavailable(
                "%s contains a missing/non-positive close" % symbol
            )
        output.append(number)
    return output


def _close_histories(symbols, count):
    raw = history(
        int(count),
        unit="1d",
        field="close",
        security_list=list(symbols),
        df=False,
        skip_paused=False,
        fq="pre",
    )
    if not hasattr(raw, "get"):
        raise PriceHistoryUnavailable("history returned an invalid payload")
    output = {}
    for symbol in symbols:
        values = list(raw.get(symbol, []))
        if len(values) != int(count):
            raise PriceHistoryUnavailable(
                "%s needs %d completed daily closes, got %d"
                % (symbol, int(count), len(values))
            )
        clean = []
        for value in values:
            try:
                number = float(value)
            except (TypeError, ValueError):
                raise PriceHistoryUnavailable(
                    "%s contains a missing/non-numeric close" % symbol
                )
            if not math.isfinite(number) or number <= 0.0:
                raise PriceHistoryUnavailable(
                    "%s contains a missing/non-positive close" % symbol
                )
            clean.append(number)
        output[convert_symbol(symbol, "canonical")] = clean
    return output


def _decision_universe(context):
    config = g.quant60_config
    if config.get("universe_mode") == "fixed":
        values = list(config["universe"])
    else:
        as_of = _as_date(
            getattr(context, "previous_date", None),
            "previous_date",
        )
        values = get_index_stocks(
            config["index_symbol"],
            date=as_of,
        )
        if not values:
            raise PriceHistoryUnavailable(
                "no PIT index members for %s on %s"
                % (config["index_symbol"], as_of.isoformat())
            )
        g.quant60_last_universe_as_of = as_of.isoformat()
    normalized = sorted(
        set(convert_symbol(symbol, "joinquant") for symbol in values)
    )
    if not normalized:
        raise PriceHistoryUnavailable("decision universe is empty")
    g.quant60_last_universe = list(normalized)
    return normalized


def _filter_st_universe(context, symbols):
    if not g.quant60_config.get("exclude_st", True):
        return list(symbols)
    as_of = _as_date(
        getattr(context, "previous_date", None),
        "previous_date",
    )
    raw = get_extras(
        "is_st",
        list(symbols),
        end_date=as_of,
        count=1,
        df=False,
    )
    if not hasattr(raw, "get"):
        raise PriceHistoryUnavailable(
            "get_extras is_st returned an invalid payload"
        )
    output = []
    for symbol in symbols:
        values = list(raw.get(symbol, []))
        if len(values) != 1:
            raise PriceHistoryUnavailable(
                "%s needs one PIT is_st value on %s"
                % (symbol, as_of.isoformat())
            )
        try:
            numeric = float(values[0])
        except (TypeError, ValueError):
            raise PriceHistoryUnavailable(
                "%s contains a non-boolean PIT is_st value" % symbol
            )
        if not math.isfinite(numeric) or numeric not in (0.0, 1.0):
            raise PriceHistoryUnavailable(
                "%s contains an invalid PIT is_st value" % symbol
            )
        if not bool(numeric):
            output.append(symbol)
    if not output:
        raise PriceHistoryUnavailable(
            "PIT ST exclusion removed the entire decision universe"
        )
    g.quant60_last_universe = list(output)
    return output


def _position_amount(position, name, default=0):
    value = getattr(position, name, default)
    return int(value or 0)


def _snapshot_portfolio(context, managed):
    current = {}
    sellable = {}
    managed = set(convert_symbol(symbol, "canonical") for symbol in managed)
    dynamic = g.quant60_config.get("universe_mode") == "pit_index"
    positions = getattr(context.portfolio, "positions", {})
    for symbol, position in positions.items():
        canonical = convert_symbol(str(symbol), "canonical")
        quantity = _position_amount(position, "total_amount")
        if quantity and canonical not in managed and not dynamic:
            raise UnmanagedPositionError(
                "account contains unmanaged position %s; use a dedicated "
                "strategy account or reconcile ownership before trading" % canonical
            )
        if canonical not in managed and not dynamic:
            continue
        current[canonical] = quantity
        sellable[canonical] = _position_amount(
            position, "closeable_amount", current[canonical]
        )
    equity = float(getattr(context.portfolio, "total_value", 0.0) or 0.0)
    return current, sellable, equity


def _is_star_board(symbol):
    code = str(symbol).split(".", 1)[0]
    return code.startswith("688") or code.startswith("689")


def _order_request(symbol, delta):
    """Build a fail-closed hosted order request from current market data.

    JoinQuant rejects an unprotected market order for STAR Market securities.
    A limit at the exchange daily bound supplies equivalent worst-price
    protection while keeping the rebalance target expressed in exact shares.
    """
    try:
        item = get_current_data()[symbol]
    except Exception as exc:
        g.quant60_skipped_orders.append(
            {"symbol": symbol, "reason": "current_data_unavailable", "detail": str(exc)}
        )
        return None
    if bool(getattr(item, "paused", False)):
        g.quant60_skipped_orders.append(
            {"symbol": symbol, "reason": "paused"}
        )
        return None
    if not _is_star_board(symbol):
        return {"style": None}

    field = "high_limit" if int(delta) > 0 else "low_limit"
    raw_price = getattr(item, field, None)
    try:
        protection_price = float(raw_price)
    except (TypeError, ValueError):
        protection_price = float("nan")
    if (
        not math.isfinite(protection_price)
        or protection_price <= 0.0
        or protection_price >= 10000.0
    ):
        g.quant60_skipped_orders.append(
            {
                "symbol": symbol,
                "reason": "star_protection_price_invalid",
                "detail": field,
            }
        )
        return None
    try:
        style = LimitOrderStyle(protection_price)
    except Exception as exc:
        g.quant60_skipped_orders.append(
            {
                "symbol": symbol,
                "reason": "star_protection_style_unavailable",
                "detail": str(exc),
            }
        )
        return None
    return {"style": style}


def compute_plan(context):
    """Return the portable plan without submitting orders."""
    config = g.quant60_config
    required = int(config["lookback"]) + int(config["skip"]) + 1
    platform_universe = _decision_universe(context)
    platform_universe = _filter_st_universe(
        context,
        platform_universe,
    )
    price_history = _close_histories(platform_universe, required)

    current, sellable, equity = _snapshot_portfolio(
        context,
        platform_universe,
    )
    if equity <= 0:
        raise ValueError("portfolio.total_value must be positive")
    plan = build_rebalance_plan(
        price_history=price_history,
        current=current,
        sellable=sellable,
        equity=equity,
        lookback=config["lookback"],
        skip=config["skip"],
        top_n=config["top_n"],
        max_weight=config["max_weight"],
        gross_target=config["gross_target"],
        cash_buffer=config["cash_buffer"],
        lot_size=config["lot_size"],
    )
    return plan


def _execute_rebalance(context):
    """Submit the portable plan's exact executable share deltas."""
    plan = compute_plan(context)
    order_deltas = plan["orders"]
    current, _, unused_equity = _snapshot_portfolio(
        context,
        g.quant60_last_universe or g.quant60_config["universe"],
    )
    del unused_equity
    all_symbols = set(current)
    all_symbols.update(order_deltas)

    # Exits first.  The requested target is current + executable delta, not the
    # unconstrained portfolio target; this preserves T+1 sellable caps and odd
    # lots from the portable plan.
    for canonical in sorted(
        all_symbols,
        key=lambda item: (
            int(order_deltas.get(item, 0)) > 0,
            item,
        ),
    ):
        delta = int(order_deltas.get(canonical, 0))
        if delta == 0:
            continue
        platform_symbol = convert_symbol(canonical, "joinquant")
        request = _order_request(platform_symbol, delta)
        if request is None:
            continue
        target = int(current.get(canonical, 0)) + delta
        if request["style"] is None:
            order_target(platform_symbol, target)
        else:
            order_target(
                platform_symbol,
                target,
                style=request["style"],
            )

    g.quant60_last_plan = plan
    return plan


def _as_date(value, name):
    if hasattr(value, "date"):
        value = value.date()
    if not hasattr(value, "isocalendar"):
        raise ClockStateError("%s must be a date/datetime" % name)
    return value


def rebalance(context):
    """Execute on the session immediately after a week's first close.

    A daily state machine is used instead of ``run_weekly(..., 2)`` so a
    one-session holiday week still executes on the next available session.
    """
    current_date = _as_date(getattr(context, "current_dt", None), "current_dt")
    previous_date = _as_date(
        getattr(context, "previous_date", None),
        "previous_date",
    )
    if g.quant60_last_callback_date == current_date:
        return None
    pending = g.quant60_pending_first_session
    is_first_session = (
        previous_date.isocalendar()[:2]
        != current_date.isocalendar()[:2]
    )
    if pending is not None:
        if previous_date != pending:
            raise ClockStateError(
                "pending first-session close is not the immediately "
                "previous trading session"
            )
        # Consume the clock token before crossing any hosted order boundary.
        # A partial API failure must not make a callback retry duplicate the
        # already accepted prefix of the order list.
        g.quant60_pending_first_session = (
            current_date if is_first_session else None
        )
        g.quant60_last_callback_date = current_date
        plan = _execute_rebalance(context)
        return plan
    if is_first_session:
        g.quant60_pending_first_session = current_date
    g.quant60_last_callback_date = current_date
    return None