#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""
STDP (Spike-Timing-Dependent Plasticity) for knowledge retrieval sequences.

When items A then B are retrieved in sequence and the task succeeds,
the edge A->B is strengthened ("A predicts the need for B").
When the task fails, the edge is weakened.

Temporal window: items retrieved within 120s of each other form pairs.
Asymmetry: A->B (causal) gets stronger boost than B->A (anti-causal).

Phase 3 AGI refinement — Studies #69-71 (Dan & Poo, Markram, Song).
"""
import json
import logging
import math
from datetime import datetime, timezone
from pulse.core.paths import get_state_dir

log = logging.getLogger("pulse.stdp")

STDP_DIR = get_state_dir() / "stdp"
SESSIONS_FILE = STDP_DIR / "session_retrievals.json"
TEMPORAL_WINDOW_S = 120  # Items within 120s form pairs
_HEBB_WINDOW = 300  # Item 58: 5-minute biological co-occurrence window for Hebbian learning
LTP_DELTA = 0.08         # Causal strengthening (A before B, task success)
LTD_DELTA = -0.04        # Anti-causal / failure weakening
ANTI_CAUSAL_RATIO = 0.3  # B->A gets 30% of A->B signal
_SILENT_SYNAPSE_THRESHOLD = 3  # F10: new edges need 3+ co-activations before becoming active

# Item 24: Eligibility traces — reward-modulated STDP
_ELIGIBILITY_TRACE: list = []  # [(edge_id, delta, timestamp)]
_MAX_TRACE = 20
_TRACE_WINDOW = 60  # seconds

# P10-33: Burst detection — 6+ co-activations in 5s = burst LTP (double delta)
_BURST_TIMESTAMPS: list = []
_BURST_THRESHOLD = 6
_BURST_WINDOW = 5.0  # seconds

_BCM_MAX_EXPECTED = 20  # Expected max events per minute for BCM normalisation


def _get_bcm_theta() -> float:
    """BCM sliding modification threshold theta_M (Bienenstock, Cooper & Munro 1982).

    theta_M = (recent_count / MAX_EXPECTED)^2
    High activity -> high theta -> LTP harder to induce (homeostatic).
    Low  activity -> low  theta -> LTP easier (lower bar).
    """
    import time
    now = time.time()
    try:
        recent_count = len([t for t in _BURST_TIMESTAMPS if now - t < 60])
        theta = (recent_count / _BCM_MAX_EXPECTED) ** 2
        return max(0.01, min(1.0, theta))  # Clamp to (0.01, 1.0]
    except Exception:
        return 0.25  # Neutral fallback

def _check_burst() -> bool:
    """Return True when 6+ STDP events occurred within the last 5 seconds."""
    import time
    now = time.time()
    _BURST_TIMESTAMPS.append(now)
    while _BURST_TIMESTAMPS and _BURST_TIMESTAMPS[0] < now - _BURST_WINDOW:
        _BURST_TIMESTAMPS.pop(0)
    recent = [t for t in _BURST_TIMESTAMPS if now - t < _BURST_WINDOW]
    return len(recent) >= _BURST_THRESHOLD

def _apply_reward_modulation(reward: float):
    """Retroactively strengthen recent STDP updates by reward signal."""
    import time
    now = time.time()
    try:
        from pulse.core.brain_db import get_conn
        conn = get_conn()
        for edge_id, delta, ts in list(_ELIGIBILITY_TRACE):
            if now - ts < _TRACE_WINDOW and reward > 0:
                try:
                    conn.execute(
                        "UPDATE graph_edges SET confidence = MIN(1.0, confidence + ?) WHERE id = ?",
                        (round(delta * reward * 0.5, 4), edge_id)
                    )
                except Exception:
                    pass
        conn.commit()
    except Exception:
        pass  # fail-open
    _ELIGIBILITY_TRACE.clear()

def record_retrieval(session_id: str, item_id: str):
    """Record that an item was retrieved in this session (called from brain_search)."""
    STDP_DIR.mkdir(parents=True, exist_ok=True)
    now = datetime.now(timezone.utc).isoformat()
    try:
        sessions = {}
        if SESSIONS_FILE.exists():
            sessions = json.loads(SESSIONS_FILE.read_text(encoding="utf-8"))
        seq = sessions.setdefault(session_id, [])
        # Avoid duplicate consecutive entries
        if seq and seq[-1].get("item_id") == item_id:
            return
        seq.append({"item_id": item_id, "ts": now})
        # Cap session length
        if len(seq) > 50:
            sessions[session_id] = seq[-50:]
        # Cap total sessions (keep last 24)
        if len(sessions) > 24:
            sorted_keys = sorted(sessions.keys())
            for k in sorted_keys[:-24]:
                del sessions[k]
        SESSIONS_FILE.write_text(json.dumps(sessions), encoding="utf-8")
    except Exception as e:
        log.debug("STDP record failed: %s", e)

def apply_stdp(session_id: str, success: bool) -> int:
    """Apply STDP learning on task completion. Returns count of edges updated."""
    try:
        from pulse.core.brain_db import get_conn
    except ImportError:
        return 0
    if not SESSIONS_FILE.exists():
        return 0
    try:
        sessions = json.loads(SESSIONS_FILE.read_text(encoding="utf-8"))
    except Exception:
        return 0
    seq = sessions.get(session_id, [])
    if len(seq) < 2:
        return 0
    conn = get_conn()
    updated = 0
    # Form temporal pairs within the window
    _burst_active = False
    try: _burst_active = _check_burst()
    except Exception: pass
    for i in range(len(seq)):
        for j in range(i + 1, len(seq)):
            try:
                ts_i = datetime.fromisoformat(seq[i]["ts"].replace("Z", "+00:00"))
                ts_j = datetime.fromisoformat(seq[j]["ts"].replace("Z", "+00:00"))
                delta_s = abs((ts_j - ts_i).total_seconds())
            except (ValueError, TypeError):
                continue
            if delta_s > TEMPORAL_WINDOW_S:
                break  # Beyond window
            pre_id = seq[i]["item_id"]
            post_id = seq[j]["item_id"]
            if pre_id == post_id:
                continue
            # Causal direction: pre -> post (A retrieved before B)
            # P1-4 FIX: Exponential temporal decay — closer retrievals get stronger signal.
            # Item 139: Error correction -- failure LTD proportional to surprise (P33)
            if success:
                raw_delta = LTP_DELTA
            else:
                try:
                    from pulse.brain.stdp_annealing import compute_failure_ltd as _cfl139
                    raw_delta = _cfl139(LTD_DELTA)
                except Exception:
                    raw_delta = LTD_DELTA
            causal_delta = raw_delta * math.exp(-delta_s / 30.0)
            try:
                if _burst_active:
                    causal_delta *= 2.0  # P10-33: burst LTP
            except Exception:
                pass
            # N49: BCM sliding threshold — single gate replaces redundant metaplasticity+frequency multipliers
            try:
                theta_M = _get_bcm_theta()
                if raw_delta > 0:
                    if causal_delta > theta_M * LTP_DELTA:
                        pass   # LTP — above threshold, keep delta as-is
                    else:
                        causal_delta = -abs(causal_delta) * 0.5  # Below threshold -> LTD
            except Exception:
                pass
            # Item 27: Monopoly penalty — if post_id dominates graph, halve confidence delta
            if _check_monopoly(post_id, conn):
                causal_delta *= 0.5
            updated += _upsert_temporal_edge(conn, pre_id, post_id, causal_delta)
            # Anti-causal: post -> pre — real STDP (Bi & Poo 1998): anti-causal → LTD on success
            if causal_delta > 0:
                anti_delta = -abs(causal_delta) * ANTI_CAUSAL_RATIO  # anti-causal → LTD
            else:
                anti_delta = abs(causal_delta) * ANTI_CAUSAL_RATIO * 0.5  # failure → mild B→A LTP
            updated += _upsert_temporal_edge(conn, post_id, pre_id, anti_delta)
    # Item 41: Triplet STDP — A->C shortcuts for A->B->C motifs
    try:
        _seen = set()
        _pairs = [(seq[i]["item_id"], seq[j]["item_id"]) for i in range(len(seq))
                  for j in range(i + 1, len(seq)) if seq[i]["item_id"] != seq[j]["item_id"]]
        for _pair_pre, _pq in _pairs:
            if (_pair_pre, _pq) in _seen: continue
            _seen.add((_pair_pre, _pq))
            for _row in conn.execute("SELECT to_node FROM graph_edges WHERE from_node=? AND edge_type='CAUSAL' AND confidence>0.3 LIMIT 3", (_pq,)).fetchall():
                if _row[0] != _pair_pre:
                    conn.execute("INSERT OR IGNORE INTO graph_edges (from_node,to_node,edge_type,confidence) VALUES (?,?,'CAUSAL', 0.3)", (_pair_pre, _row[0]))
    except Exception:
        pass  # fail-open
    conn.commit()
    # Clean up session after processing
    if session_id in sessions:
        del sessions[session_id]
        try:
            SESSIONS_FILE.write_text(json.dumps(sessions), encoding="utf-8")
        except Exception:
            pass
    log.info("STDP: %d edges updated for session %s (success=%s)", updated, session_id, success)
    return updated

def _check_monopoly(node_id: str, conn) -> bool:
    """Item 27: Return True if node has >30% of all graph edges (monopoly hub)."""
    try:
        total = conn.execute("SELECT COUNT(*) FROM graph_edges").fetchone()[0]
        node_edges = conn.execute(
            "SELECT COUNT(*) FROM graph_edges WHERE to_node = ?", (node_id,)
        ).fetchone()[0]
        return total > 0 and (node_edges / total) > 0.3
    except Exception:
        return False

def _upsert_temporal_edge(conn, from_id: str, to_id: str, delta: float) -> int:
    """Create or update a TEMPORAL edge. F10: new edges start silent (conf=0)."""
    existing = conn.execute(
        "SELECT id, confidence, metadata FROM graph_edges WHERE from_node=? AND to_node=? AND edge_type='TEMPORAL'",
        (from_id, to_id)
    ).fetchone()
    if existing:
        try:
            meta = json.loads(existing["metadata"] or "{}")
        except Exception:
            meta = {}
        coact = meta.get("coactivation_count", 1) + 1
        meta["coactivation_count"] = coact
        # Promote from silent only after threshold reached
        current_conf = existing["confidence"]
        if coact >= _SILENT_SYNAPSE_THRESHOLD:
            if delta < 0:
                # Item 45: Depotentiation vs LTD distinction
                try:
                    import time as _t45
                    last_pot = meta.get("last_potentiated", 0)
                    if _t45.time() - last_pot < 300:
                        # Depotentiation: revert recently-potentiated edge to zero
                        new_conf = 0.0
                    else:
                        # LTD: gradual long-term depression
                        new_conf = max(0.0, current_conf * 0.7)
                except Exception:
                    new_conf = max(0.0, current_conf * 0.7)
            else:
                new_conf = max(0.0, min(1.0, max(current_conf, 0.01) + delta))
                # Track potentiation timestamp for depotentiation logic
                try:
                    import time as _t45p
                    meta["last_potentiated"] = _t45p.time()
                except Exception:
                    pass
        else:
            new_conf = current_conf  # stays silent
        if new_conf < 0.01 and coact >= _SILENT_SYNAPSE_THRESHOLD:
            conn.execute("DELETE FROM graph_edges WHERE id=?", (existing["id"],))
        else:
            conn.execute(
                "UPDATE graph_edges SET confidence=?, metadata=? WHERE id=?",
                (round(new_conf, 4), json.dumps(meta), existing["id"])
            )
        # Item 24: Record to eligibility trace
        try:
            import time as _t
            _ELIGIBILITY_TRACE.append((existing["id"], delta, _t.time()))
            if len(_ELIGIBILITY_TRACE) > _MAX_TRACE:
                _ELIGIBILITY_TRACE.pop(0)
        except Exception:
            pass
    else:
        if delta > 0:
            now = datetime.now(timezone.utc).isoformat()
            # F10: new edges start silent — confidence=0.0 until 3+ co-activations
            conn.execute(
                "INSERT INTO graph_edges (from_node, to_node, edge_type, confidence, metadata) "
                "VALUES (?, ?, 'TEMPORAL', ?, ?)",
                (from_id, to_id, 0.0,
                 json.dumps({"source": "stdp", "created": now, "coactivation_count": 1,
                             "natural_confidence": round(max(0.1, delta), 4)}))
            )
        # Item 24: Record new edge to eligibility trace
        try:
            import time as _tnew
            lastrowid = conn.execute(
                "SELECT id FROM graph_edges WHERE from_node=? AND to_node=? AND edge_type='TEMPORAL'",
                (from_id, to_id)
            ).fetchone()
            if lastrowid:
                _ELIGIBILITY_TRACE.append((lastrowid[0], delta, _tnew.time()))
                if len(_ELIGIBILITY_TRACE) > _MAX_TRACE:
                    _ELIGIBILITY_TRACE.pop(0)
        except Exception:
            pass
    return 1
