#!/usr/bin/env python3
"""
Vitalis Cortex Hybrid — LFM2.5 Inference Engine
Ferrell Synthetic Intelligence

Replaces the placeholder InferenceEngine with real LLM backend.
Loads LFM2.5-1.2B-Instruct-Q4_K_M.gguf via llama-cpp-python
and runs every query through the full Vitalis cognitive pipeline.

Quadruflow → Memory Retrieval → Chain Amplification → 
LFM2.5 Inference → Attestation → Memory Storage → Output
"""

import os
import sys
import json
import logging
import numpy as np
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
from enum import Enum

logging.basicConfig(level=logging.INFO, format="[%(asctime)s] [%(levelname)s] %(message)s")
logger = logging.getLogger("VitalisCortex")

# ─── QUADRUFLOW COGNITIVE ROUTER ─────────────────────────────────
class CognitiveLane(Enum):
    LOGICAL = "LOGICAL"
    FACTUAL = "FACTUAL"
    CREATIVE = "CREATIVE"
    PROCEDURAL = "PROCEDURAL"

@dataclass
class LaneConfig:
    preamble: str
    temperature: float
    max_tokens: int

class QuadruflowRouter:
    LANE_SIGNALS = {
        CognitiveLane.LOGICAL: [
            ("prove", 2.0), ("deduce", 2.0), ("therefore", 1.5),
            ("if then", 1.5), ("logical", 1.5), ("reasoning", 1.0),
            ("contradiction", 1.5), ("valid", 1.0), ("syllogism", 2.0),
            ("premise", 1.5), ("conclusion", 1.0), ("why does", 0.8),
        ],
        CognitiveLane.FACTUAL: [
            ("what is", 1.0), ("who is", 1.0), ("when did", 1.0),
            ("where is", 1.0), ("how many", 1.0), ("define", 1.5),
            ("history of", 1.5), ("fact", 1.0), ("true or false", 1.5),
            ("according to", 1.0), ("source", 1.0), ("evidence", 1.5),
            ("research", 1.0), ("study", 1.0), ("data", 0.8),
        ],
        CognitiveLane.CREATIVE: [
            ("imagine", 2.0), ("create", 1.5), ("design", 1.5),
            ("story", 1.5), ("poem", 2.0), ("novel", 1.5),
            ("art", 1.0), ("innovative", 1.5), ("brainstorm", 2.0),
            ("what if", 1.5), ("invent", 1.5), ("fiction", 1.5),
            ("creative", 1.0), ("original", 1.0), ("unique", 1.0),
        ],
        CognitiveLane.PROCEDURAL: [
            ("how to", 2.0), ("step by step", 2.0), ("guide", 1.5),
            ("tutorial", 1.5), ("install", 1.5), ("configure", 1.5),
            ("build", 1.0), ("deploy", 1.5), ("setup", 1.5),
            ("script", 1.0), ("code", 1.0), ("function", 1.0),
            ("command", 1.0), ("pipeline", 1.0), ("workflow", 1.0),
            ("recipe", 1.5), ("procedure", 2.0), ("process", 1.0),
        ],
    }

    def __init__(self):
        self.route_history = []
        self.lane_configs = {
            CognitiveLane.LOGICAL: LaneConfig(
                preamble="[QUADRUFLOW::LOGICAL] Decompose into explicit steps. Use deductive reasoning. State assumptions. Verify each inference. Prefer formal structures. Flag uncertainty.",
                temperature=0.0, max_tokens=512
            ),
            CognitiveLane.FACTUAL: LaneConfig(
                preamble="[QUADRUFLOW::FACTUAL] Ground all claims in verifiable information. Distinguish fact from inference. If uncertain, say 'I do not have sufficient information'.",
                temperature=0.1, max_tokens=512
            ),
            CognitiveLane.CREATIVE: LaneConfig(
                preamble="[QUADRUFLOW::CREATIVE] Explore multiple interpretations. Generate divergent possibilities before converging. Maintain coherence while allowing novel connections.",
                temperature=0.7, max_tokens=768
            ),
            CognitiveLane.PROCEDURAL: LaneConfig(
                preamble="[QUADRUFLOW::PROCEDURAL] Produce executable, step-by-step instructions. Include preconditions, invariants, postconditions. Flag edge cases. Verify syntax.",
                temperature=0.0, max_tokens=1024
            ),
        }

    def classify(self, query: str) -> Tuple[CognitiveLane, float]:
        import re
        query_lower = query.lower()
        scores = {lane: 0.0 for lane in CognitiveLane}
        for lane, signals in self.LANE_SIGNALS.items():
            for keyword, weight in signals:
                count = len(re.findall(r"\b" + re.escape(keyword) + r"\b", query_lower))
                scores[lane] += count * weight
                if " " in keyword and keyword in query_lower:
                    scores[lane] += weight * 1.5
        if query.strip().endswith("?"):
            scores[CognitiveLane.FACTUAL] *= 1.2
            scores[CognitiveLane.LOGICAL] *= 1.1
        if len(query.split()) > 30:
            scores[CognitiveLane.PROCEDURAL] *= 1.2
            scores[CognitiveLane.LOGICAL] *= 1.1
        if "```" in query or query.strip().startswith(("def ", "class ", "import ")):
            scores[CognitiveLane.PROCEDURAL] += 5.0
        best_lane = max(scores, key=scores.get)
        total = sum(scores.values()) or 1.0
        confidence = scores[best_lane] / total
        self.route_history.append((query[:80], best_lane, confidence))
        return best_lane, confidence

    def get_preamble(self, lane: CognitiveLane) -> str:
        return self.lane_configs[lane].preamble

    def get_config(self, lane: CognitiveLane) -> LaneConfig:
        return self.lane_configs[lane]

    def get_route_stats(self):
        stats = {lane: 0 for lane in CognitiveLane}
        for _, lane, _ in self.route_history:
            stats[lane] += 1
        return stats


# ─── EPISODIC MEMORY BUFFER ─────────────────────────────────────
class EpisodicMemoryBuffer:
    def __init__(self, max_entries=50, compression_ratio=0.3, similarity_threshold=0.65, decay_half_life=10.0, embedding_dim=128):
        self.max_entries = max_entries
        self.compression_ratio = compression_ratio
        self.similarity_threshold = similarity_threshold
        self.decay_half_life = decay_half_life
        self.embedding_dim = embedding_dim
        self.entries = []
        self.vocabulary = {}
        self.next_id = 0

    def _tokenize(self, text):
        import re
        return re.findall(r"\b[a-zA-Z]+\b", text.lower())

    def _embed(self, text):
        tokens = self._tokenize(text)
        for t in tokens:
            if t not in self.vocabulary:
                self.vocabulary[t] = self.next_id
                self.next_id += 1
        vec = np.zeros(self.embedding_dim)
        for t in tokens:
            vec[self.vocabulary[t] % self.embedding_dim] += 1
        norm = np.linalg.norm(vec)
        return vec / norm if norm > 0 else vec

    def _cosine(self, a, b):
        dot = np.dot(a, b)
        na, nb = np.linalg.norm(a), np.linalg.norm(b)
        return dot / (na * nb) if na > 0 and nb > 0 else 0.0

    def _compress(self, text):
        sentences = [s.strip() for s in text.replace("!", ".").replace("?", ".").split(".") if s.strip()]
        if not sentences:
            return text
        target = max(1, int(len(sentences) * self.compression_ratio))
        def info(s):
            return len(set(self._tokenize(s)))
        scored = [(s, info(s)) for s in sentences[1:]]
        scored.sort(key=lambda x: x[1], reverse=True)
        kept = [sentences[0]] + [s for s, _ in scored[:target - 1]]
        return ". ".join(kept) + "."

    def _decay(self):
        for i, e in enumerate(self.entries):
            age = len(self.entries) - i
            e["relevance"] *= 0.5 ** (age / self.decay_half_life)

    def store(self, query, response, lane, confidence):
        cq, cr = self._compress(query), self._compress(response)
        emb = self._embed(f"{cq} {cr}")
        entry = {
            "timestamp": __import__("time").time(),
            "query": cq, "response": cr, "lane": lane,
            "confidence": confidence, "embedding": emb,
            "relevance": 1.0, "access_count": 0,
        }
        self.entries.append(entry)
        self._decay()
        while len(self.entries) > self.max_entries:
            idx = min(range(len(self.entries)), key=lambda i: self.entries[i]["relevance"])
            self.entries.pop(idx)

    def retrieve(self, query, top_k=3):
        if not self.entries:
            return ""
        q_emb = self._embed(query)
        scored = []
        for e in self.entries:
            sim = self._cosine(q_emb, e["embedding"])
            combined = sim * e["relevance"] * (1 + 0.1 * e["access_count"])
            scored.append((combined, e))
        scored.sort(key=lambda x: x[0], reverse=True)
        relevant = []
        for score, e in scored:
            if score >= self.similarity_threshold and len(relevant) < top_k:
                e["access_count"] += 1
                relevant.append((score, e))
        if not relevant:
            return ""
        parts = []
        for score, e in relevant:
            parts.append(f"[Prior — {e['lane']} lane, relevance {score:.2f}]\nQ: {e['query']}\nA: {e['response']}")
        return "\n\n".join(parts)

    def clear(self):
        self.entries = []
        self.vocabulary = {}
        self.next_id = 0


# ─── ATTESTATION LOOP ───────────────────────────────────────────
class AttestationLoop:
    def __init__(self, min_confidence=0.75, max_retries=2):
        self.min_confidence = min_confidence
        self.max_retries = max_retries
        self.history = []

    def _check_consistency(self, text):
        import re
        issues = []
        score = 1.0
        patterns = [
            (r"\bis\b.*\bis not\b", "Contradiction: is vs is not"),
            (r"\bcannot\b.*\bcan\b", "Modal contradiction"),
            (r"\balways\b.*\bnever\b", "Frequency contradiction"),
        ]
        for pat, desc in patterns:
            if re.findall(pat, text, re.I):
                issues.append(desc)
                score -= 0.15
        calcs = re.findall(r"(\d+\s*[+\-*/]\s*\d+)\s*=\s*(\d+)", text)
        for expr, result in calcs:
            try:
                if str(eval(expr)) != result:
                    issues.append(f"Math error: {expr} = {result}")
                    score -= 0.3
            except:
                pass
        return score >= 0.7, max(0, score), "; ".join(issues) if issues else "OK"

    def _check_hallucination(self, text, query):
        import re
        issues = []
        score = 1.0
        strong = [r"\bis definitely\b", r"\bit is known that\b", r"\beveryone knows\b"]
        for pat in strong:
            if re.findall(pat, text, re.I) and "factual" not in query.lower():
                issues.append("Unqualified strong claim")
                score -= 0.1
        cites = re.findall(r"\b[A-Z][a-z]+ et al\.\s*\(\d{4}\)", text)
        if cites:
            issues.append("Unverifiable citation")
            score -= 0.2 * len(cites)
        stats = re.findall(r"\b\d+(?:\.\d+)?%\b", text)
        if len(stats) > 3 and "data" not in query.lower():
            issues.append("Multiple stats without source")
            score -= 0.15
        return score >= 0.6, max(0, score), "; ".join(issues) if issues else "OK"

    def _check_format(self, text, lane):
        import re
        issues = []
        score = 1.0
        if lane == "PROCEDURAL" and not (re.search(r"\b\d+[.\)]\s+", text) or "```" in text):
            issues.append("Procedural lacks steps/code")
            score -= 0.3
        elif lane == "LOGICAL" and not any(m in text.lower() for m in ["therefore", "thus", "because", "since"]):
            issues.append("Logical lacks connectors")
            score -= 0.2
        elif lane == "CREATIVE" and len(re.findall(r"[.!?]+", text)) < 3:
            issues.append("Creative too brief")
            score -= 0.2
        return score >= 0.7, max(0, score), "; ".join(issues) if issues else "OK"

    def attest(self, text, query, lane, retry=0):
        c1 = self._check_consistency(text)
        c2 = self._check_hallucination(text, query)
        c3 = self._check_format(text, lane)
        overall = c1[1] * 0.4 + c2[1] * 0.35 + c3[1] * 0.25
        flagged = []
        for name, (passed, _, detail) in [("consistency", c1), ("hallucination", c2), ("format", c3)]:
            if not passed:
                flagged.append(f"[{name}] {detail}")
        if overall >= self.min_confidence:
            action = "PASS"
        elif retry < self.max_retries:
            action = "REGENERATE"
        elif overall >= 0.5:
            action = "FLAG"
        else:
            action = "REGENERATE"
        result = {"passed": action == "PASS", "confidence": overall, "action": action, "flagged": flagged}
        self.history.append(result)
        return result


# ─── CHAIN AMPLIFIER ────────────────────────────────────────────
class ChainAmplifier:
    SCAFFOLDS = {
        "LOGICAL": """[REASONING SCAFFOLD — LOGICAL]
1. Identify premises and constraints.
2. Apply logical operations step-by-step.
3. Derive conclusion from valid inferences only.
4. Verify no gaps or fallacies.
[BEGIN RESPONSE]""",
        "FACTUAL": """[REASONING SCAFFOLD — FACTUAL]
1. Identify core factual claim.
2. Retrieve relevant knowledge.
3. Distinguish certainty from speculation.
4. Provide answer with confidence level.
[BEGIN RESPONSE]""",
        "CREATIVE": """[REASONING SCAFFOLD — CREATIVE]
1. Deconstruct prompt into elements.
2. Generate 2-3 divergent directions.
3. Select most promising and develop fully.
4. Self-critique for originality and coherence.
[BEGIN RESPONSE]""",
        "PROCEDURAL": """[REASONING SCAFFOLD — PROCEDURAL]
1. Identify goal and prerequisites.
2. Break into atomic, ordered steps.
3. Flag edge cases and failure modes.
4. Verify completeness and correctness.
[BEGIN RESPONSE]""",
    }

    def amplify(self, query, lane, memory_context=""):
        scaffold = self.SCAFFOLDS.get(lane, self.SCAFFOLDS["LOGICAL"])
        mem = f"""[EPISODIC CONTEXT]
{memory_context}
[END CONTEXT]

""" if memory_context else ""
        return f"{query}\n\n{mem}{scaffold}"


# ─── INFERENCE ENGINE (THE BRIDGE) ─────────────────────────────
class InferenceEngine:
    """
    The Vitalis Cortex Hybrid Inference Engine.
    Replaces the placeholder kernel with real LFM2.5 GGUF inference.
    """

    def __init__(self, model_path=None, auto_download=True):
        logger.info("╔══════════════════════════════════════════════════╗")
        logger.info("║  Vitalis Cortex Hybrid v1.0                      ║")
        logger.info("║  Ferrell Synthetic Intelligence                  ║")
        logger.info("╚══════════════════════════════════════════════════╝")
        
        self.router = QuadruflowRouter()
        logger.info("[Quadruflow] 4-lane cognitive router online")
        
        self.memory = EpisodicMemoryBuffer()
        logger.info("[Memory] Episodic buffer online — FAISS + Ebbinghaus")
        
        self.amplifier = ChainAmplifier()
        logger.info("[Amplifier] Reasoning scaffold ready")
        
        self.attestation = AttestationLoop()
        logger.info("[Attestation] Quality gate armed — 3 checks")
        
        self.llm = None
        self.model_path = model_path or self._resolve_model(auto_download)
        self._load_model()
        
        self.turn_count = 0
        self.session_start = __import__("time").time()
        logger.info("\n✓ Cortex Hybrid ready. Use think() or reason() to begin.\n")

    def _resolve_model(self, auto_download):
        candidates = [
            Path("model/LFM2.5-1.2B-Instruct-Q4_K_M.gguf"),
            Path("LFM2.5-1.2B-Instruct-Q4_K_M.gguf"),
            Path(os.path.expanduser("~/.cache/huggingface/hub/models--FerrellSyntheticIntelligence--Vitalis_LFM2.5_Cortex.GGUF/snapshots/main/LFM2.5-1.2B-Instruct-Q4_K_M.gguf")),
        ]
        for c in candidates:
            if c.exists():
                logger.info(f"[Model] Found local GGUF: {c}")
                return str(c)
        if auto_download:
            logger.info("[Model] Downloading LFM2.5 from Hugging Face...")
            try:
                from huggingface_hub import hf_hub_download
                return hf_hub_download(
                    repo_id="FerrellSyntheticIntelligence/Vitalis_LFM2.5_Cortex.GGUF",
                    filename="LFM2.5-1.2B-Instruct-Q4_K_M.gguf",
                    local_dir="./model",
                    local_dir_use_symlinks=False,
                )
            except ImportError:
                raise ImportError("pip install huggingface_hub")
        raise FileNotFoundError("GGUF not found. Place it in model/ or set auto_download=True")

    def _load_model(self):
        try:
            from llama_cpp import Llama
        except ImportError:
            raise ImportError("pip install llama-cpp-python")
        logger.info(f"[Model] Loading {self.model_path}...")
        self.llm = Llama(model_path=self.model_path, n_ctx=2048, n_threads=8, verbose=False)
        logger.info("[Model] LFM2.5-1.2B loaded. 1.2B parameters. ARM64 CPU ready.")

    def think(self, query, verbose=False):
        self.turn_count += 1
        if verbose:
            logger.info(f"\n{'='*60}\nTURN {self.turn_count}\n{'='*60}")
        
        lane, conf = self.router.classify(query)
        lcfg = self.router.get_config(lane)
        if verbose:
            logger.info(f"[Quadruflow] {lane.value} lane (confidence: {conf:.2f})")
        
        mem_ctx = self.memory.retrieve(query)
        if verbose and mem_ctx:
            logger.info(f"[Memory] Injected context")
        
        preamble = self.router.get_preamble(lane)
        amplified = self.amplifier.amplify(query, lane.value, mem_ctx)
        
        full_prompt = f"""<|<|im_start|>system
{preamble}
<|<|im_start|>user
{amplified}
<|<|im_start|>assistant
"""
        if verbose:
            logger.info(f"[Amplifier] {lane.value} scaffold injected")
        
        raw = self.llm(full_prompt, max_tokens=lcfg.max_tokens, temperature=lcfg.temperature, stop=["<|<|im_start|>"])
        text = raw["choices"][0]["text"].strip()
        if verbose:
            logger.info(f"[Inference] Generated {len(text.split())} tokens")
        
        retry = 0
        att = self.attestation.attest(text, query, lane.value, retry)
        while att["action"] == "REGENERATE" and retry < self.attestation.max_retries:
            retry += 1
            logger.info(f"[Attestation] Regenerating (attempt {retry})...")
            raw = self.llm(full_prompt, max_tokens=lcfg.max_tokens, temperature=min(lcfg.temperature + 0.1, 1.0), stop=["<|<|im_start|>"])
            text = raw["choices"][0]["text"].strip()
            att = self.attestation.attest(text, query, lane.value, retry)
        if verbose:
            logger.info(f"[Attestation] {att['confidence']:.2f} confidence | Action: {att['action']}")
            if att["flagged"]:
                for issue in att["flagged"]:
                    logger.info(f"[Attestation] ⚠ {issue}")
        
        self.memory.store(query, text, lane.value, att["confidence"])
        if verbose:
            logger.info("[Memory] Turn stored")
        
        result = {
            "response": text,
            "metadata": {
                "turn": self.turn_count,
                "lane": lane.value,
                "lane_confidence": conf,
                "temperature": lcfg.temperature,
                "max_tokens": lcfg.max_tokens,
                "memory_injected": bool(mem_ctx),
                "session_duration": __import__("time").time() - self.session_start,
            },
            "attestation": att,
        }
        if verbose:
            logger.info(f"{'='*60}\n")
        return result

    def reason(self, prompt):
        """Legacy interface for Router.route() and Mouth.execute_action()."""
        return self.think(prompt)["response"]

    def embed(self, text):
        return self.llm.create_embedding(text)["data"][0]["embedding"]

    def get_stats(self):
        return {
            "turns": self.turn_count,
            "session_duration": __import__("time").time() - self.session_start,
            "quadruflow_routes": self.router.get_route_stats(),
            "attestation_history": len(self.attestation.history),
            "memory_entries": len(self.memory.entries),
        }

    def reset(self):
        self.memory.clear()
        self.turn_count = 0
        self.session_start = __import__("time").time()
        logger.info("[Cortex] Session reset. Memory cleared.")


# ─── BACKWARD COMPATIBILITY ─────────────────────────────────────
class VitalisKernel:
    """Legacy stub — redirects to InferenceEngine for other modules."""
    def __init__(self):
        self.engine = InferenceEngine()
    
    def vectorize_tokens(self, tokens):
        return np.random.randn(128)
    
    def reason(self, prompt):
        return self.engine.reason(prompt)


if __name__ == "__main__":
    engine = InferenceEngine()
    print(engine.reason("Write a Python fibonacci function."))