# 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__
from quant60.portable_core import build_rebalance_plan, convert_symbol
# __PORTABLE_CORE_BUNDLE_END__

import math


class PriceHistoryUnavailable(RuntimeError):
    pass


class UnmanagedPositionError(RuntimeError):
    pass


class ClockStateError(RuntimeError):
    pass


# __BASELINE_CONFIG_START__
CONFIG = {
    "universe_mode": "pit_index",
    "index_symbol": "000905.XSHG",
    "qmt_sector_name": "中证500",
    "dedicated_account_required": True,
    "exclude_st": True,
    "universe": [
        "600000.XSHG",
        "000001.XSHE",
        "300750.XSHE",
        "000333.XSHE",
    ],
    "lookback": 20,
    # The first session close supplies the signal; the second session open
    # supplies the hosted execution boundary, matching the QMT daily wrapper.
    "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": 0.00001,
    "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