Multi_Swarm_Coevolutive/src/multi_swarm/agents/hypothesis.py

from __future__ import annotations

import re
from dataclasses import dataclass

from ..genome.hypothesis import HypothesisAgentGenome
from ..llm.client import CompletionResult, 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:
    strategy: Strategy | None
    raw_text: str
    completion: CompletionResult
    parse_error: str | None = None


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 nel linguaggio S-expression
con i seguenti verbi disponibili:

  Azioni:        entry-long, entry-short, exit, flat
  Logici:        and, or, not
  Comparatori:   gt, lt, eq
  Dati:          feature, indicator, crossover, crossunder

Indicatori disponibili (calcolati implicitamente sul prezzo close):
  sma <length>, rsi <length>, atr <length>, macd, realized_vol <window>.
Feature disponibili: open, high, low, close, volume.

REGOLE STRETTE DI SINTASSI:
- (indicator <name> <args...>) restituisce una serie numerica. Es.
  (indicator rsi 14), (indicator sma 50), (indicator macd 12 26 9).
- (feature <name>) restituisce la colonna OHLCV. Es. (feature close).
- Gli indicatori NON sono annidabili: NON puoi scrivere
  (sma (indicator realized_vol 30) 150) o (indicator rsi (feature high) 14).
  Le funzioni sma/rsi/etc. ESISTONO SOLO come argomenti di indicator,
  non sono verbi indipendenti.
- Costanti numeriche (es. 70.0, 30, 0.02) sono valide come 2° operando di gt/lt/eq.
- crossover/crossunder accettano due espressioni-serie:
  (crossover (feature close) (indicator sma 20)) — corretto.
  (crossover (sma close 20) (sma close 50)) — ERRATO (sma non è verbo).

Le regole sono valutate in ordine; la prima che matcha vince per ogni timestamp.
La default action se nessuna regola matcha è 'flat'.

Rispondi SOLO con la S-expression in un fence ```lisp ... ```, senza prosa,
senza spiegazioni. Esempio formato:

```lisp
(strategy
  (when (gt (indicator rsi 14) 70.0) (entry-short))
  (when (lt (indicator rsi 14) 30.0) (entry-long))
  (when (crossover (feature close) (indicator sma 50)) (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.
"""


_SEXP_FENCE_RE = re.compile(
    r"```(?:lisp|scheme|sexp)?\s*(\(strategy[\s\S]*?\))\s*```",
    re.MULTILINE,
)


def _extract_sexp(text: str) -> str | None:
    m = _SEXP_FENCE_RE.search(text)
    if m:
        return m.group(1)
    if text.strip().startswith("(strategy"):
        return text.strip()
    return None


class HypothesisAgent:
    def __init__(self, llm: LLMClient):
        self._llm = llm

    def propose(
        self,
        genome: HypothesisAgentGenome,
        market: MarketSummary,
    ) -> HypothesisProposal:
        system = SYSTEM_TEMPLATE.format(
            cognitive_style=genome.cognitive_style,
            system_prompt=genome.system_prompt,
        )
        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,
        )

        completion = self._llm.complete(genome, system=system, user=user)

        sexp = _extract_sexp(completion.text)
        if sexp is None:
            return HypothesisProposal(
                strategy=None,
                raw_text=completion.text,
                completion=completion,
                parse_error="no s-expression found in output",
            )
        try:
            ast = parse_strategy(sexp)
            validate_strategy(ast)
            return HypothesisProposal(
                strategy=ast,
                raw_text=completion.text,
                completion=completion,
            )
        except (ParseError, ValidationError) as e:
            return HypothesisProposal(
                strategy=None,
                raw_text=completion.text,
                completion=completion,
                parse_error=str(e),
            )