integrations/quantaalpha/llm_strategy.py

"""LLM-driven strategy selection and factor weighting (QuantaAlpha API integration)."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import pandas as pd

from config.settings import PROJECT_ROOT
from integrations.quantaalpha.client import QuantaAlphaLLMClient, load_llm_config
from integrations.quantaalpha.factor_library import list_factors


STRATEGY_CATALOG = [
    "topk_dropout",
    "long_short_quantile",
    "score_weighted_topk",
    "rank_weighted",
    "soft_topk",
    "enhanced_indexing",
    "dynamic_risk_topk",
    "factor_equal_topk",
    "factor_ic_weighted_topk",
]


def _build_factor_context(catalog: pd.DataFrame, max_factors: int = 30) -> str:
    rows = catalog.head(max_factors).to_dict(orient="records")
    return json.dumps(rows, ensure_ascii=False, indent=2)


def propose_strategy_with_llm(
    factor_catalog: pd.DataFrame,
    market_context: str | None = None,
    client: QuantaAlphaLLMClient | None = None,
) -> dict[str, Any]:
    """
    Ask LLM to propose strategy type, parameters, and factor selection.
    Returns a dict compatible with strategies/registry.yaml entries.
    """
    client = client or QuantaAlphaLLMClient(load_llm_config())
    factor_json = _build_factor_context(factor_catalog)

    system = (
        "You are a quantitative portfolio strategist. "
        "Given factor metadata, choose the best strategy from the catalog and parameters. "
        f"Available strategies: {', '.join(STRATEGY_CATALOG)}. "
        "Respond in JSON with keys: strategy_name, strategy_kwargs, selected_factor_ids, "
        "signal_combine, rationale."
    )
    user = (
        f"Market context: {market_context or 'CSI300 daily alpha strategy, out-of-sample backtest'}\n\n"
        f"Factor catalog:\n{factor_json}\n\n"
        "Pick 3-10 factors if using multi-factor combine strategies."
    )

    result = client.chat_json(
        [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ]
    )
    return result


def propose_from_quantaalpha_library(
    library_path: str | Path,
    market_context: str | None = None,
    quality_filter: str | None = "high",
) -> dict[str, Any]:
    catalog = list_factors(library_path, quality_filter=quality_filter)
    if catalog.empty:
        catalog = list_factors(library_path)
    return propose_strategy_with_llm(catalog, market_context=market_context)


def build_llm_strategy_plan(
    library_path: str | Path | None = None,
    factor_panel_path: str | Path | None = None,
    market_context: str | None = None,
) -> dict[str, Any]:
    """High-level entry: LLM plan -> signal source + strategy config."""
    if library_path:
        llm_plan = propose_from_quantaalpha_library(library_path, market_context=market_context)
        signal_source = {
            "type": "quantaalpha_library",
            "path": str(library_path),
            "factor_ids": llm_plan.get("selected_factor_ids"),
            "combine": llm_plan.get("signal_combine", "ic_weighted"),
        }
    elif factor_panel_path:
        from data_pipeline.factor_loader import load_factor_panel

        panel = load_factor_panel(factor_panel_path)
        factor_cols = [c for c in panel.columns if c.startswith("factor_")]
        pseudo = pd.DataFrame({"factor_id": factor_cols, "factor_name": factor_cols, "icir": 1.0})
        llm_plan = propose_strategy_with_llm(pseudo, market_context=market_context)
        signal_source = {
            "type": "factor_panel",
            "path": str(factor_panel_path),
            "factor_cols": llm_plan.get("selected_factor_ids") or factor_cols,
            "combine": llm_plan.get("signal_combine", "equal"),
        }
    else:
        raise ValueError("Provide library_path or factor_panel_path")

    strategy_name = llm_plan.get("strategy_name", "topk_dropout")
    if strategy_name not in STRATEGY_CATALOG:
        strategy_name = "topk_dropout"

    return {
        "llm_plan": llm_plan,
        "signal_source": signal_source,
        "strategy": {
            "name": strategy_name,
            "kwargs": llm_plan.get("strategy_kwargs", {}),
        },
    }


def save_strategy_plan(plan: dict[str, Any], output_path: str | Path) -> Path:
    path = Path(output_path)
    if not path.is_absolute():
        path = PROJECT_ROOT / path
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(plan, f, ensure_ascii=False, indent=2)
    return path