Multi_Swarm_Coevolutive/src/multi_swarm/agents/hypothesis.py

from __future__ import annotations

import re
from dataclasses import dataclass, field

from ..genome.hypothesis import HypothesisAgentGenome
from ..llm.client import CompletionResult, EmptyCompletionError, LLMClient
from ..protocol.parser import ParseError, Strategy, parse_strategy
from ..protocol.validator import ValidationError, validate_strategy


@dataclass(frozen=True)
class MarketSummary:
    symbol: str
    timeframe: str
    n_bars: int
    return_mean: float
    return_std: float
    skew: float
    kurtosis: float
    volatility_regime: str


@dataclass(frozen=True)
class HypothesisProposal:
    """Risultato di una propose() del HypothesisAgent.

    ``completions`` contiene SEMPRE almeno un elemento: il primo tentativo.
    Se il primo tentativo fallisce e c'e' budget di retry, vengono accodate
    le completions successive, una per ogni retry effettuato.
    ``n_attempts == len(completions)``. ``raw_text`` riflette l'ULTIMO output
    LLM osservato (quello che ha prodotto strategy o l'ultimo parse_error).
    """

    strategy: Strategy | None
    raw_text: str
    completions: list[CompletionResult] = field(default_factory=list)
    parse_error: str | None = None
    n_attempts: int = 1


SYSTEM_TEMPLATE = """\
Sei un agente generatore di ipotesi di trading quantitativo per un sistema swarm.

Il tuo stile cognitivo: {cognitive_style}
Direttiva personale: {system_prompt}

Devi proporre una strategia di trading espressa in JSON STRETTO.
La risposta deve essere un singolo oggetto JSON dentro fence ```json...```
con questa shape:

```json
{{
  "rules": [
    {{"condition": <nodo>, "action": "entry-long|entry-short|exit|flat"}}
  ]
}}
```

NODI DISPONIBILI

Operatori logici:
  {{"op": "and", "args": [<nodo>, <nodo>, ...]}}    // >=2 nodi
  {{"op": "or",  "args": [<nodo>, <nodo>, ...]}}    // >=2 nodi
  {{"op": "not", "args": [<nodo>]}}                  // 1 nodo

Comparatori (ritornano boolean series):
  {{"op": "gt", "args": [<a>, <b>]}}    // a > b
  {{"op": "lt", "args": [<a>, <b>]}}    // a < b
  {{"op": "eq", "args": [<a>, <b>]}}    // a == b

Crossover (eventi su 2 serie):
  {{"op": "crossover",  "args": [<serie_a>, <serie_b>]}}
  {{"op": "crossunder", "args": [<serie_a>, <serie_b>]}}

Leaf - indicatori (calcolati su close):
  {{"kind": "indicator", "name": "sma",          "params": [<length>]}}
  {{"kind": "indicator", "name": "rsi",          "params": [<length>]}}
  {{"kind": "indicator", "name": "atr",          "params": [<length>]}}
  {{"kind": "indicator", "name": "realized_vol", "params": [<window>]}}
  {{"kind": "indicator", "name": "macd",         "params": [<fast>, <slow>, <signal>]}}
    // 0-3 numeri (tutti opzionali con default 12, 26, 9)

Leaf - feature OHLCV:
  {{"kind": "feature", "name": "open|high|low|close|volume"}}

Leaf - feature TEMPORALI (sempre accessibili, UTC):
  {{"kind": "feature", "name": "hour"}}            // range 0-23
  {{"kind": "feature", "name": "dow"}}             // range 0-6 (lun=0, dom=6)
  {{"kind": "feature", "name": "is_weekend"}}      // 0 o 1
  {{"kind": "feature", "name": "minute_of_hour"}}  // range 0-59

Esempi di gating temporale:
  // Solo durante la sessione US (14:00-22:00 UTC)
  {{"op": "and", "args": [
    {{"op": "gt", "args": [{{"kind": "feature", "name": "hour"}}, {{"kind": "literal", "value": 14}}]}},
    {{"op": "lt", "args": [{{"kind": "feature", "name": "hour"}}, {{"kind": "literal", "value": 22}}]}}
  ]}}

  // Solo nel weekend (sab+dom)
  {{"op": "eq", "args": [{{"kind": "feature", "name": "is_weekend"}}, {{"kind": "literal", "value": 1}}]}}

Leaf - letterale numerico:
  {{"kind": "literal", "value": 70.0}}

VINCOLI
- Gli indicator NON sono annidabili: 'params' accetta solo numeri, mai altri nodi.
- Le regole sono valutate in ordine; la prima che matcha vince per ogni timestamp.
- Default action se nessuna regola matcha = flat.
- 'op' e 'kind' sono mutuamente esclusivi sullo stesso nodo.

Rispondi SOLO con il fence ```json...``` contenente l'oggetto strategy.
Esempio:

```json
{{
  "rules": [
    {{
      "condition": {{"op": "gt", "args": [
        {{"kind": "indicator", "name": "rsi", "params": [14]}},
        {{"kind": "literal", "value": 70.0}}
      ]}},
      "action": "entry-short"
    }},
    {{
      "condition": {{"op": "lt", "args": [
        {{"kind": "indicator", "name": "rsi", "params": [14]}},
        {{"kind": "literal", "value": 30.0}}
      ]}},
      "action": "entry-long"
    }}
  ]
}}
```
"""


USER_TEMPLATE = """\
Mercato: {symbol} timeframe {timeframe}, {n_bars} barre osservate.
Statistiche return: mean={return_mean:.5f}, std={return_std:.5f}, \
skew={skew:.3f}, kurt={kurtosis:.3f}.
Regime volatilità: {volatility_regime}.

Feature accessibili dal tuo genoma: {feature_access}.
Lookback massimo che puoi usare nel ragionamento: {lookback_window} barre.

Genera una strategia che cerchi anomalie sfruttabili in questo regime.
"""


_RETRY_TEMPLATE = """\
{original_user}

--- TENTATIVO PRECEDENTE FALLITO ---
Output: {previous_raw}
Errore: {previous_error}
---
Correggi l'errore e rispondi di nuovo con un singolo oggetto JSON valido
dentro fence ```json...```, seguendo strettamente lo schema fornito nel
SYSTEM message.
"""

_RETRY_RAW_TRUNCATE = 800


_JSON_FENCE_RE = re.compile(
    r"```(?:json)?\s*(\{[\s\S]*\})\s*```",
    re.MULTILINE,
)


def _balance_braces(s: str) -> str | None:
    """Ritorna il prefix di ``s`` che chiude la prima ``{`` con bilanciamento.

    Usato come fallback quando l'LLM ritorna JSON top-level senza fence ma
    seguito da prosa: troviamo dove finisce il primo oggetto e tagliamo.
    """
    if not s.startswith("{"):
        return None
    depth = 0
    in_string = False
    escape = False
    for i, ch in enumerate(s):
        if in_string:
            if escape:
                escape = False
            elif ch == "\\":
                escape = True
            elif ch == '"':
                in_string = False
            continue
        if ch == '"':
            in_string = True
        elif ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return s[: i + 1]
    return None


def _extract_json(text: str) -> str | None:
    """Estrai un oggetto JSON dal testo del completion.

    Strategie di estrazione, in ordine:
      1. Fence ```json...``` (greedy: cattura fino all'ultimo ``}`` prima della
         chiusura del fence).
      2. Testo che inizia direttamente con ``{`` (dopo strip), bilanciato a
         livello di parentesi graffe.
    """
    m = _JSON_FENCE_RE.search(text)
    if m:
        return m.group(1)
    stripped = text.strip()
    return _balance_braces(stripped)


def _try_parse(text: str) -> tuple[Strategy | None, str | None]:
    """Estrai+parsea+valida. Ritorna (strategy, error). Esattamente uno e' None."""
    payload = _extract_json(text)
    if payload is None:
        return None, "no JSON object found in output"
    try:
        ast = parse_strategy(payload)
        validate_strategy(ast)
    except (ParseError, ValidationError) as e:
        return None, str(e)
    return ast, None


class HypothesisAgent:
    def __init__(self, llm: LLMClient, max_retries: int = 1):
        if max_retries < 0:
            raise ValueError("max_retries must be >= 0")
        self._llm = llm
        self._max_retries = max_retries

    def propose(
        self,
        genome: HypothesisAgentGenome,
        market: MarketSummary,
    ) -> HypothesisProposal:
        system = SYSTEM_TEMPLATE.format(
            cognitive_style=genome.cognitive_style,
            system_prompt=genome.system_prompt,
        )
        original_user = USER_TEMPLATE.format(
            symbol=market.symbol,
            timeframe=market.timeframe,
            n_bars=market.n_bars,
            return_mean=market.return_mean,
            return_std=market.return_std,
            skew=market.skew,
            kurtosis=market.kurtosis,
            volatility_regime=market.volatility_regime,
            feature_access=", ".join(genome.feature_access),
            lookback_window=genome.lookback_window,
        )

        completions: list[CompletionResult] = []
        errors: list[str] = []
        last_raw = ""
        max_attempts = 1 + self._max_retries

        for attempt in range(max_attempts):
            if attempt == 0:
                user = original_user
            else:
                truncated = last_raw[:_RETRY_RAW_TRUNCATE]
                user = _RETRY_TEMPLATE.format(
                    original_user=original_user,
                    previous_raw=truncated,
                    previous_error=errors[-1],
                )

            try:
                completion = self._llm.complete(genome, system=system, user=user)
            except EmptyCompletionError as e:
                # LLM esaurito retry tenacity senza una risposta. Tratta come
                # parse-fail "empty" e ritenta nel loop esterno (max_attempts).
                errors.append(f"empty_completion: {e}")
                last_raw = ""
                continue
            completions.append(completion)
            last_raw = completion.text

            strategy, err = _try_parse(completion.text)
            if strategy is not None:
                return HypothesisProposal(
                    strategy=strategy,
                    raw_text=completion.text,
                    completions=completions,
                    parse_error=None,
                    n_attempts=len(completions),
                )
            assert err is not None
            errors.append(err)

        chained = " | ".join(
            f"attempt {i + 1}: {e}" for i, e in enumerate(errors)
        )
        return HypothesisProposal(
            strategy=None,
            raw_text=last_raw,
            completions=completions,
            parse_error=chained,
            n_attempts=len(completions),
        )