"""memory_engine.py - Apprentissage contextuel par trade fermé.

Chaque trade est associé à un contexte (régime, bucket RSI, sentiment).
Une distribution Beta à décroissance exponentielle est entretenue par contexte.
Le StrategyEngine consulte cette mémoire pour veto les contextes perdants.
"""
import threading
import time
from dataclasses import dataclass
from typing import Optional

import database as db


@dataclass
class TradeContext:
    regime: str
    rsi_bucket: int
    sentiment_bucket: int
    entry_price: float
    exit_price: float
    pnl_pct: float
    duration_min: int
    exit_reason: str


def rsi_to_bucket(rsi: float) -> int:
    if rsi < 20:
        return 0
    if rsi < 40:
        return 1
    if rsi < 60:
        return 2
    if rsi < 80:
        return 3
    return 4


def sentiment_to_bucket(score: float) -> int:
    if score < -0.4:
        return -2
    if score < -0.1:
        return -1
    if score < 0.1:
        return 0
    if score < 0.4:
        return 1
    return 2


class ContextualMemory:
    DECAY = 0.995
    MIN_TRADES = 8
    VETO_THRESHOLD = 0.42

    def __init__(self):
        self._lock = threading.Lock()
        self.cells: dict = {}
        self.trade_log: list = []
        self._load_from_db()

    def _load_from_db(self):
        try:
            rows = db.load_memory_cells()
            for r in rows:
                key = (r["regime"], int(r["rsi_bucket"]))
                self.cells[key] = {
                    "wins": float(r["wins"]),
                    "total": float(r["total"]),
                    "last_update": float(r.get("last_update", time.time())),
                }
            if self.cells:
                print(f"[MEMORY] {len(self.cells)} contextes chargés.")
        except Exception as e:
            print(f"[MEMORY] Aucune mémoire existante : {e}")

    def _save_cell(self, key):
        try:
            cell = self.cells[key]
            db.save_memory_cell(
                key[0], key[1],
                cell["wins"], cell["total"], cell["last_update"],
            )
        except Exception as e:
            print(f"[MEMORY] Sauvegarde échouée : {e}")

    def record_trade(self, ctx: TradeContext):
        key = (ctx.regime, ctx.rsi_bucket)
        with self._lock:
            cell = self.cells.setdefault(
                key, {"wins": 0.0, "total": 0.0, "last_update": time.time()}
            )
            cell["wins"] = cell["wins"] * self.DECAY + (1.0 if ctx.pnl_pct > 0 else 0.0)
            cell["total"] = cell["total"] * self.DECAY + 1.0
            cell["last_update"] = time.time()
            self._save_cell(key)

            self.trade_log.append(ctx)
            if len(self.trade_log) > 500:
                self.trade_log = self.trade_log[-500:]

            wr = (cell["wins"] / cell["total"]) * 100 if cell["total"] > 0 else 0
            print(
                f"[MEMORY] {ctx.regime} RSI[{ctx.rsi_bucket}] "
                f"PnL={ctx.pnl_pct:+.2f}% → WR={wr:.0f}% (n={cell['total']:.1f})"
            )

    def expected_win_rate(self, regime: str, rsi_bucket: int) -> Optional[float]:
        with self._lock:
            cell = self.cells.get((regime, rsi_bucket))
            if not cell or cell["total"] < self.MIN_TRADES:
                return None
            return cell["wins"] / cell["total"]

    def should_trade(self, regime: str, rsi_bucket: int):
        wr = self.expected_win_rate(regime, rsi_bucket)
        if wr is None:
            return True, "apprentissage en cours"
        if wr < self.VETO_THRESHOLD:
            return False, f"win rate {wr*100:.0f}% < {self.VETO_THRESHOLD*100:.0f}%"
        return True, f"win rate {wr*100:.0f}%"

    def stats(self) -> dict:
        with self._lock:
            cells_list = []
            for (regime, rsi_bucket), cell in self.cells.items():
                if cell["total"] < 1:
                    continue
                wr = cell["wins"] / cell["total"] if cell["total"] > 0 else 0
                cells_list.append({
                    "regime": regime,
                    "rsi_bucket": rsi_bucket,
                    "win_rate": wr,
                    "samples": cell["total"],
                    "last_update": cell["last_update"],
                })
            cells_list.sort(key=lambda c: -c["win_rate"])
            return {
                "total_trades_logged": len(self.trade_log),
                "contexts_learned": len(cells_list),
                "best_context": cells_list[0] if cells_list else None,
                "worst_context": cells_list[-1] if cells_list else None,
                "all_cells": cells_list,
                "veto_threshold": self.VETO_THRESHOLD,
                "min_samples": self.MIN_TRADES,
            }


MEMORY = ContextualMemory()