brain.py

"""
ForensicAI 2nd Brain — adaptive state manager.
Tracks case types, evidence confidence, and forensic insights across sessions.
"""

from __future__ import annotations

import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

BRAIN_PATH = Path("/app/forensic_brain_state.json")

CASE_TYPES = [
    "homicide", "sexual_assault", "burglary", "fraud",
    "drug_offense", "digital_crime", "arson", "hit_and_run",
]

DEFAULT: dict = {
    "session_count":         0,
    "total_analyses":        0,
    "total_evidence_items":  0,
    "case_frequency":        {c: 0 for c in CASE_TYPES},
    "confidence_history":    [],
    "complexity_history":    [],
    "avg_confidence":        0.70,
    "avg_complexity":        0.50,
    "dominant_case":         "default",
    "pattern_generation":    1,
    "feed_rate":             0.037,
    "kill_rate":             0.060,
    "last_updated":          None,
    "insights":              [],
    "evolution_log":         [],
}

_CASE_PARAMS = {
    "homicide":       (0.025, 0.055),
    "sexual_assault": (0.037, 0.061),
    "digital_crime":  (0.029, 0.057),
    "drug_offense":   (0.040, 0.060),
    "arson":          (0.033, 0.058),
    "burglary":       (0.038, 0.062),
    "fraud":          (0.026, 0.053),
    "hit_and_run":    (0.035, 0.059),
    "default":        (0.037, 0.060),
}


def load() -> dict:
    if BRAIN_PATH.exists():
        try:
            return json.loads(BRAIN_PATH.read_text())
        except Exception:
            pass
    return DEFAULT.copy()


def _save(state: dict) -> None:
    BRAIN_PATH.parent.mkdir(parents=True, exist_ok=True)
    state["last_updated"] = datetime.now(timezone.utc).isoformat()
    BRAIN_PATH.write_text(json.dumps(state, indent=2))


def _recompute(state: dict) -> None:
    dom = state["dominant_case"]
    base_f, base_k = _CASE_PARAMS.get(dom, _CASE_PARAMS["default"])
    conf  = max(0.0, min(1.0, state["avg_confidence"]))
    cmplx = max(0.0, min(1.0, state["avg_complexity"]))
    state["feed_rate"] = round(base_f + cmplx * 0.008, 5)
    state["kill_rate"] = round(base_k + conf  * 0.006, 5)


def update(
    case_type:    Optional[str]   = None,
    confidence:   Optional[float] = None,
    complexity:   Optional[float] = None,
    n_evidence:   int             = 0,
    insight:      Optional[str]   = None,
) -> dict:
    state = load()
    state["session_count"]  += 1
    state["total_analyses"] += 1
    state["total_evidence_items"] += n_evidence

    if confidence is not None:
        h = state["confidence_history"]
        h.append(round(confidence, 3))
        state["confidence_history"] = h[-200:]
        state["avg_confidence"] = round(sum(h) / len(h), 3)

    if complexity is not None:
        h = state["complexity_history"]
        h.append(round(complexity, 3))
        state["complexity_history"] = h[-200:]
        state["avg_complexity"] = round(sum(h) / len(h), 3)

    if case_type and case_type in state["case_frequency"]:
        state["case_frequency"][case_type] += 1
        state["dominant_case"] = max(
            state["case_frequency"], key=lambda k: state["case_frequency"][k]
        )

    if insight:
        state["insights"].append({
            "ts":   datetime.now(timezone.utc).isoformat(),
            "text": insight,
        })
        state["insights"] = state["insights"][-30:]

    _recompute(state)

    if state["session_count"] % 15 == 0:
        old = state["pattern_generation"]
        state["pattern_generation"] += 1
        state["evolution_log"].append({
            "ts":         datetime.now(timezone.utc).isoformat(),
            "from_gen":   old,
            "to_gen":     state["pattern_generation"],
            "feed":       state["feed_rate"],
            "kill":       state["kill_rate"],
            "dominant":   state["dominant_case"],
        })
        state["evolution_log"] = state["evolution_log"][-20:]

    _save(state)
    return state


def reset() -> dict:
    state = DEFAULT.copy()
    _save(state)
    return state