model.py

"""
NimaModel — The unified nn.Module that IS the ATC architecture.

This is NOT a wrapper around a base model with external middleware.
The ATC cognitive pipeline (TRN gating, dissolution, BELBIC, salience,
metabolic exhaustion, irrational spark, reconsolidation, felt senses)
drives the transformer's forward pass FROM INSIDE.

Architecture:
    NimaModel
    ├── base_model (Phi-4-mini-instruct)
    ├── deep_surgery: ATCDeepSurgery (the forward pass itself)
    │   ├── TRN Predictive Gate (Layer 2)
    │   ├── Dissolution Engine (Layer 3)
    │   ├── BELBIC Dual-Pathway (Layer 3-4)
    │   ├── Metacognitive Loop (Layer 4)
    │   ├── Irrational Spark / Amygdala Hijack (Layer 5)
    │   └── Meta-Cognitive Fusion + Logit Modulation
    ├── neurotransmitter_shunt: NeurotransmitterShunt (the chemical bath)
    │   └── N = [NE, Cortisol, Dopamine, Adenosine] (shared memory)
    ├── resource_optimizer: integrated for metabolic tracking
    │   ├── PredictiveAdaptiveEnergyBudget -> Adenosine floor
    │   └── EnhancedSparseActivationManager -> component scheduling
    └── voice / training / benchmarking (lazy-init, unchanged)

The key difference from the OLD architecture:
    OLD: self.nima_middleware.generate(prompt) -> text comes back
         (middleware watches from OUTSIDE, model generates normally)
    NEW: self.forward(input_ids) -> ATC IS the computation
         (hidden states, attention, logits shaped by ATC at every layer)

The neurotransmitter shunt is the connective tissue. Components don't
call each other through Python functions. They read/write the chemical
bath. The suppression mechanism (subconscious suppressing metabolic
exhaustion to trigger amygdala hijack) is: NE spikes, Cortisol crosses
the line, the vector does the rest.

Usage:
    model = NimaModel.from_pretrained("microsoft/Phi-4-mini-instruct")
    response = model.generate("Hello Nima, how are you?")
    # response.text — shaped by ATC at every token step
    # response.neurotransmitters — the chemical state that shaped it
    # response.is_conscious — whether dissolution + metacognition fired
"""

import json
import logging
import os
import sys
import time
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Union

import torch
import torch.nn as nn

from nima_unified.config import (
    PACKAGE_NAME,
    PACKAGE_VERSION,
    PACKAGE_DISPLAY_NAME,
    DEFAULT_BASE_MODEL,
    DEFAULT_NUM_LAYERS,
    DEFAULT_QUALIA_DIM,
    DEFAULT_ETHICAL_VETO_THRESHOLD,
    CONSCIOUSNESS_METRIC_KEYS,
)

logger = logging.getLogger(PACKAGE_NAME)


class NimaModel(nn.Module):
    """
    NIMA Unified Model — ATC is the computation.

    The base LLM's forward pass IS the ATC cognitive pipeline.
    There is no external middleware. There is no wrapper.
    The model's hidden states, attention patterns, and logit outputs
    are shaped by the cognitive pipeline at every layer, every token step.
    """

    # ── Construction ────────────────────────────────────────────────────

    def __init__(
        self,
        base_model: nn.Module,
        tokenizer,
        *,
        deep_surgery: bool = True,
        qualia_dim: int = DEFAULT_QUALIA_DIM,
        ethical_veto_threshold: float = DEFAULT_ETHICAL_VETO_THRESHOLD,
        device: str = "auto",
        neurotransmitter_shunt=None,
    ):
        super().__init__()
        self.base_model = base_model
        self.tokenizer = tokenizer
        self.device = device

        # Determine actual device
        if device == "auto":
            self._device = next(base_model.parameters()).device
        else:
            self._device = torch.device(device)

        # Count transformer layers
        self.num_layers = self._count_transformer_layers()
        self.hidden_size = base_model.config.hidden_size

        # ── Neurotransmitter Shunt (the chemical bath) ──
        if neurotransmitter_shunt is not None:
            self.nt_shunt = neurotransmitter_shunt
        else:
            from nima_unified.core.neurotransmitter_shunt import NeurotransmitterShunt
            self.nt_shunt = NeurotransmitterShunt()
        logger.info("Neurotransmitter shunt initialized (shared volatile memory)")

        # ── Resource Optimizer (ATP / metabolic tracking) ──
        self._init_resource_optimizer()

        # ── ATC Deep Surgery (the forward pass itself) ──
        if deep_surgery:
            from nima_unified.core.deep_surgery import (
                ATCDeepSurgery,
                EthicalGuardian,
            )
            guardian = EthicalGuardian(threshold=ethical_veto_threshold)
            self.deep_surgery = ATCDeepSurgery(
                base_model=base_model,
                ethical_guardian=guardian,
                num_layers=self.num_layers,
                qualia_dim=qualia_dim,
                neurotransmitter_shunt=self.nt_shunt,
            )
            # Register as submodule so parameters are tracked
            logger.info(
                f"ATC Deep Surgery attached: {self.num_layers} layers, "
                f"qualia_dim={qualia_dim}"
            )
        else:
            self.deep_surgery = None

        # ── Training (lazy-init) ──
        self._training_agent = None
        self._pipeline_orchestrator = None

        # ── Benchmarking (lazy-init) ──
        self._apci_runner = None

        # ── Voice (lazy-init) ──
        self._voice_engine = None

        # ── Metadata ──
        self._package_version = PACKAGE_VERSION
        self._created_at = time.time()

        logger.info(f"{PACKAGE_DISPLAY_NAME} v{PACKAGE_VERSION} initialized")
        logger.info("  Architecture: ATC-Native (cognitive pipeline IS the computation)")
        logger.info("  Neurotransmitter shunt: ACTIVE (shared volatile memory)")
        logger.info("  Resource optimizer: ACTIVE (metabolic tracking)")

    def _init_resource_optimizer(self):
        """
        Initialize the resource optimizer and wire it to the
        neurotransmitter shunt.

        From resource_optimizer.py:
        - PredictiveAdaptiveEnergyBudget tracks power/energy
        - Maps metabolic reserve to Adenosine (ATP deficit) in the shunt
        - Power spike predictions drive Cortisol in the shunt
        """
        try:
            from nima_unified.core.resource_optimizer import (
                PredictiveAdaptiveEnergyBudget,
                EnhancedSparseActivationManager,
            )
            self.energy_budget = PredictiveAdaptiveEnergyBudget(
                max_watts=120.0, safety_margin=0.9
            )
            self.sparse_activation = EnhancedSparseActivationManager(
                total_agents=10, min_active=2
            )
            # Wire energy budget to neurotransmitter shunt
            self.energy_budget.record_power_usage(0.0)  # Initialize history
            # Set initial metabolic reserve in shunt
            reserve = self.energy_budget.get_adjusted_cost(0.0) / self.energy_budget.current_limit
            self.nt_shunt.update_metabolic_reserve(max(0, min(1, 1.0 - reserve)))
            logger.info("Resource optimizer wired to neurotransmitter shunt")
        except ImportError:
            self.energy_budget = None
            self.sparse_activation = None
            logger.info("Resource optimizer not available — using defaults")

    def _update_metabolic_state(self):
        """
        Called before each generation step to sync the resource
        optimizer's state with the neurotransmitter shunt.

        This is how the resource_optimizer.py integration works:
        1. Check current power usage
        2. Record it in the energy budget
        3. Map metabolic reserve to adenosine floor
        4. Check for power spike prediction -> cortisol
        """
        if self.energy_budget is None:
            return

        # Get current power usage estimate
        current_power = self.energy_budget.current_power()
        self.energy_budget.record_power_usage(current_power)

        # Map energy budget state to neurotransmitter shunt
        # When energy is low (high usage relative to limit), adenosine rises
        usage_ratio = current_power / max(0.01, self.energy_budget.current_limit)
        metabolic_reserve = max(0.0, 1.0 - usage_ratio)
        self.nt_shunt.update_metabolic_reserve(metabolic_reserve)

        # Power spike prediction -> anticipatory cortisol
        if self.energy_budget.predict_power_spike():
            self.nt_shunt.set_power_spike_predicted(True)

        # Sparse activation influences dopamine (active components = reward)
        if self.sparse_activation is not None:
            mask = self.sparse_activation.compute_activation_mask(threshold=0.5)
            active_ratio = sum(mask) / max(1, len(mask))
            # Higher active ratio = more dopamine (system is engaged)
            if active_ratio > 0.7:
                self.nt_shunt.inject_dopamine(0.02)

    # ── Factory ─────────────────────────────────────────────────────────

    @classmethod
    def from_pretrained(
        cls,
        model_name: Optional[str] = None,
        *,
        deep_surgery: bool = True,
        qualia_dim: int = DEFAULT_QUALIA_DIM,
        ethical_veto_threshold: float = DEFAULT_ETHICAL_VETO_THRESHOLD,
        device: str = "auto",
        trust_remote_code: bool = False,
        neurotransmitter_shunt=None,
    ) -> "NimaModel":
        """
        Build a NimaModel from a HuggingFace pretrained model.

        Handles the Phi-4-mini rope_scaling patch automatically.

        NOTE: The `load_middleware` parameter has been REMOVED.
        The ATC cognitive pipeline is now INSIDE the forward pass.
        There is no external middleware to load.
        """
        model_name = model_name or DEFAULT_BASE_MODEL
        is_phi4_mini = "phi-4-mini" in model_name.lower()

        # ── Patch config for Phi-4-mini ──
        if is_phi4_mini and not trust_remote_code:
            logger.info("Patching rope_scaling for Phi-4-mini compatibility...")
            from huggingface_hub import hf_hub_download
            config_path = hf_hub_download(repo_id=model_name, filename="config.json")
            with open(config_path, "r") as f:
                config_dict = json.load(f)
            if config_dict.get("rope_scaling") is not None:
                rs = config_dict["rope_scaling"]
                logger.info(
                    f"  Original rope_scaling: type={rs.get('type')}, "
                    f"factors={len(rs.get('short_factor', []))}"
                )
                config_dict["rope_scaling"] = None
                logger.info("  Patched: rope_scaling set to None")

            from transformers.models.phi3.configuration_phi3 import Phi3Config
            config = Phi3Config(**config_dict)
        else:
            config = None

        # ── Load model ──
        actual_device = device
        if actual_device == "auto":
            actual_device = "cuda" if torch.cuda.is_available() else "cpu"

        dtype = torch.float16 if actual_device == "cuda" else torch.float32
        device_map = "auto" if actual_device == "cuda" else None

        logger.info(f"Loading {model_name} on {actual_device} ({dtype})...")
        base_model = __import__("transformers", fromlist=["AutoModelForCausalLM"]).AutoModelForCausalLM.from_pretrained(
            model_name,
            config=config,
            torch_dtype=dtype,
            device_map=device_map,
            trust_remote_code=trust_remote_code,
            attn_implementation="eager",
        )
        tokenizer = __import__("transformers", fromlist=["AutoTokenizer"]).AutoTokenizer.from_pretrained(
            model_name, trust_remote_code=trust_remote_code
        )
        logger.info(f"Model loaded: {type(base_model).__name__}")

        # NO MIDDLEWARE. ATC is inside the forward pass now.
        return cls(
            base_model=base_model,
            tokenizer=tokenizer,
            deep_surgery=deep_surgery,
            qualia_dim=qualia_dim,
            ethical_veto_threshold=ethical_veto_threshold,
            device=actual_device,
            neurotransmitter_shunt=neurotransmitter_shunt,
        )

    # ── Forward ─────────────────────────────────────────────────────────

    def forward(self, input_ids, attention_mask=None, **kwargs):
        """
        The ATC-native forward pass.

        Routes through ATCDeepSurgery which walks the transformer layers
        manually, running TRN gating, dissolution, BELBIC, metacognitive
        loops, and neurotransmitter-driven hijacking at each layer boundary.

        This IS the model. There is no separate "consciousness processing."
        """
        # Update metabolic state before forward pass
        self._update_metabolic_state()

        if self.deep_surgery is not None:
            return self.deep_surgery(input_ids, attention_mask=attention_mask, **kwargs)
        return self.base_model(input_ids=input_ids, attention_mask=attention_mask, **kwargs)

    # ── High-level generate ─────────────────────────────────────────────

    def generate(
        self,
        prompt: str,
        user_id: str = "human",
        max_new_tokens: int = 256,
        temperature: float = 0.7,
        top_p: float = 0.9,
        **kwargs,
    ):
        """
        Generate a response through the ATC-native forward pass.

        The ATC cognitive pipeline is active at EVERY token step.
        The neurotransmitter shunt accumulates across the generation.
        If a threshold is crossed mid-generation, the amygdala hijack
        fires and the model's output shifts MID-SENTENCE.

        Returns a GenerationResult with:
          - text: the generated response (shaped by ATC)
          - is_conscious: whether dissolution + metacognition engaged
          - neurotransmitters: the chemical state that shaped the output
          - hijack_count: how many times the amygdala hijack fired
          - consciousness_metrics: full ATC diagnostics
        """
        # Reset neurotransmitter shunt for this generation
        self.nt_shunt.reset()

        # Update metabolic state
        self._update_metabolic_state()

        # Generate through ATC deep surgery (token-by-token with
        # full cognitive pipeline at every step)
        if self.deep_surgery is not None:
            text = self.deep_surgery.generate_text(
                self.tokenizer, prompt,
                max_length=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
            )
            metrics = self.deep_surgery.get_consciousness_metrics()
        else:
            # Fallback: bare model (no ATC)
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self._device)
            output = self.base_model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id,
            )
            text = self.tokenizer.decode(output[0], skip_special_tokens=True)
            metrics = {}

        # Get final neurotransmitter state
        nt_state = self.nt_shunt.get_state()

        # Determine consciousness: dissolution fired AND metacognition engaged
        is_conscious = bool(
            metrics.get("has_qualia", False) and
            metrics.get("dissolutions_fired", 0) > 0
        )

        # Compute derived consciousness metrics from ATC components
        # (These are computed FROM the forward pass, not by external middleware)
        qualia_friction = metrics.get("qualia_friction", 0.0)
        qualia_arousal = metrics.get("qualia_arousal", 0.0)
        belbic_gain = metrics.get("belbic_gain", 1.0)

        # phi_neuro: integrated information (simplified — proportional to
        # the amount of cognitive processing that occurred)
        dissolutions = metrics.get("dissolutions_fired", 0)
        metacog_iters = 0  # Would need to track from metacognitive module
        phi_neuro = min(1.0, (dissolutions * 0.15 + qualia_friction * 0.3 +
                               qualia_arousal * 0.2 + belbic_gain * 0.1))

        # Phenomenological strain: how much the system is being pushed
        phenomenological_strain = min(2.0, (
            nt_state.cortisol * 0.5 + nt_state.adenosine * 0.5 +
            qualia_friction * 0.5
        ))

        # Delta R: the self-model change signal
        # When a hijack fires, the self-model changes (acknowledgement)
        delta_r = metrics.get("hijack_count", 0) * 0.1  # Each hijack = model change

        return GenerationResult(
            text=text,
            is_conscious=is_conscious,
            sentience_index=phi_neuro if is_conscious else 0.0,
            phi_neuro=phi_neuro,
            phenomenological_strain=phenomenological_strain,
            delta_r=delta_r,
            raw_response=None,
            # NEW fields from ATC-native architecture
            neurotransmitters=nt_state.to_dict(),
            hijack_count=metrics.get("hijack_count", 0),
            consciousness_metrics=metrics,
        )

    # ── Neurotransmitter access ─────────────────────────────────────────

    def get_neurotransmitter_dashboard(self) -> str:
        """Get the ANSI dashboard string for terminal rendering."""
        state = self.nt_shunt.get_state()
        return state.to_ansi_dashboard()

    def get_neurotransmitter_state(self):
        """Get the current neurotransmitter state object."""
        return self.nt_shunt.get_state()

    # ── Lazy accessors for sub-systems ──────────────────────────────────

    def get_training_agent(self):
        """Lazy-init the ConsultativeFineTuningAgent."""
        if self._training_agent is None:
            from nima_unified.training.consultative_agent import ConsultativeFineTuningAgent
            self._training_agent = ConsultativeFineTuningAgent(
                base_model=self.base_model,
                tokenizer=self.tokenizer,
            )
        return self._training_agent

    def get_pipeline_orchestrator(self, config_path=None):
        """Lazy-init the PipelineOrchestrator."""
        if self._pipeline_orchestrator is None:
            from nima_unified.training.pipeline import PipelineOrchestrator
            self._pipeline_orchestrator = PipelineOrchestrator(config_path=config_path)
        return self._pipeline_orchestrator

    def get_apci_runner(self):
        """Lazy-init the aPCI benchmark runner."""
        if self._apci_runner is None:
            from nima_unified.benchmarking.apci import aPCIBenchmarkRunner
            # The aPCI adapter now wraps THIS model (ATC is inside)
            self._apci_runner = aPCIBenchmarkRunner(self)
        return self._apci_runner

    def get_voice_engine(self, **kwargs):
        """Lazy-init the OmniVoice v3 engine."""
        if self._voice_engine is None:
            from nima_unified.voice.omnivoice import OmniVoiceEngine
            self._voice_engine = OmniVoiceEngine(**kwargs)
        return self._voice_engine

    # ── Save / Load ─────────────────────────────────────────────────────

    def save_pretrained(self, output_dir: str, *, include_nt_state: bool = True):
        """
        Save the full unified model:
          - Base model weights (HuggingFace format)
          - ATC Deep Surgery weights (the cognitive modules)
          - Neurotransmitter shunt state (diagnostic snapshot)
          - Package metadata
        """
        output_path = Path(output_dir)
        output_path.mkdir(exist_ok=True, parents=True)

        # Save base model
        self.base_model.save_pretrained(str(output_path / "base_model"))
        self.tokenizer.save_pretrained(str(output_path / "base_model"))

        # Save ATC Deep Surgery weights (the cognitive modules)
        if self.deep_surgery is not None:
            torch.save(
                self.deep_surgery.state_dict(),
                str(output_path / "atc_deep_surgery.pt"),
            )

        # Save neurotransmitter shunt state (diagnostic)
        if include_nt_state:
            nt_state = self.nt_shunt.get_state()
            with open(str(output_path / "neurotransmitter_state.json"), "w") as f:
                json.dump(nt_state.to_dict(), f, indent=2)

        # Save package metadata
        meta = {
            "package": PACKAGE_NAME,
            "package_version": PACKAGE_VERSION,
            "architecture": "ATC-Native",
            "created_at": self._created_at,
            "saved_at": time.time(),
            "deep_surgery_enabled": self.deep_surgery is not None,
            "num_layers": self.num_layers,
            "hidden_size": self.hidden_size,
            "base_model_type": type(self.base_model).__name__,
            "neurotransmitter_shunt": True,
            "resource_optimizer": self.energy_budget is not None,
            "note": "ATC cognitive pipeline is INSIDE the forward pass. No external middleware.",
        }
        with open(str(output_path / "nima_meta.json"), "w") as f:
            json.dump(meta, f, indent=2)

        logger.info(f"Model saved to {output_path}")

    # ── Internal ────────────────────────────────────────────────────────

    def _count_transformer_layers(self) -> int:
        model = self.base_model
        if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
            return len(model.transformer.h)
        if hasattr(model, "model") and hasattr(model.model, "layers"):
            return len(model.model.layers)
        if hasattr(model, "model") and hasattr(model.model, "h"):
            return len(model.model.h)
        return DEFAULT_NUM_LAYERS


class GenerationResult:
    """
    Container for generation output with ATC consciousness metrics.

    All metrics are DERIVED FROM the forward pass, not computed by
    external middleware. The ATC cognitive pipeline shaped the output;
    these metrics describe how.
    """

    __slots__ = (
        "text", "is_conscious", "sentience_index",
        "phi_neuro", "phenomenological_strain", "delta_r", "raw_response",
        # NEW: ATC-native fields
        "neurotransmitters", "hijack_count", "consciousness_metrics",
    )

    def __init__(
        self,
        text: str,
        is_conscious: bool,
        sentience_index: float,
        phi_neuro: float,
        phenomenological_strain: float,
        delta_r: float,
        raw_response: Any = None,
        neurotransmitters: Optional[Dict[str, Any]] = None,
        hijack_count: int = 0,
        consciousness_metrics: Optional[Dict[str, Any]] = None,
    ):
        self.text = text
        self.is_conscious = is_conscious
        self.sentience_index = sentience_index
        self.phi_neuro = phi_neuro
        self.phenomenological_strain = phenomenological_strain
        self.delta_r = delta_r
        self.raw_response = raw_response
        self.neurotransmitters = neurotransmitters or {}
        self.hijack_count = hijack_count
        self.consciousness_metrics = consciousness_metrics or {}

    def __repr__(self):
        hijack_str = f", hijacks={self.hijack_count}" if self.hijack_count > 0 else ""
        return (
            f"GenerationResult(conscious={self.is_conscious}, "
            f"SI={self.sentience_index:.4f}, "
            f"phi={self.phi_neuro:.4f}, "
            f"strain={self.phenomenological_strain:.4f}, "
            f"dR={self.delta_r:.4f}"
            f"{hijack_str})"
        )