#!/usr/bin/env python3
"""
Q-Factor Energy Balance System (Verified Lean Specification)

This implementation follows the formal specification in:
0-Core-Formalism/lean/Semantics/Semantics/QFactor.lean

The Lean module provides:
- Q-Factor energy balance equation: Q = (E_flash + E_enthalpy + E_recovered - W_demon) / (E_work + E_loss)
- Energy balance optimization
- Bind primitive for Q-Factor transitions
- Invariant preservation theorems

This Python shim provides:
- JSON serialization for energy balance state
- Result wrapping for Lean function calls
- History deque for Q-Factor transitions
- No logic (all logic defined in Lean specification)
"""

import json
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
from collections import deque

sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "4-Infrastructure"))
from lib.q16 import Q16_ONE, Q16_SCALE, from_q16, q16_add, q16_div, q16_ge, q16_gt, q16_mul, q16_sub, to_q16

try:
    from temporal_spatial_ram import TemporalSpatialRAMSystem, NodePosition, TemporalSpatialResource
    _HAS_TS_RAM = True
except ImportError:
    _HAS_TS_RAM = False
    print("[!] Temporal-spatial RAM system not available")
@dataclass
class EnergyBalance:
    """Energy balance components (Lean: EnergyBalance)"""
    flashEnergy: int  # Q16_16 - Burst computation energy
    enthalpy: int  # Q16_16 - Steady-state energy
    recoveredEnergy: int  # Q16_16 - Energy from optimizations
    demonWork: int  # Q16_16 - Landauer limit for erasure
    workEnergy: int  # Q16_16 - Useful work energy
    energyLoss: int  # Q16_16 - Waste energy
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'flashEnergy': from_q16(self.flashEnergy),
            'enthalpy': from_q16(self.enthalpy),
            'recoveredEnergy': from_q16(self.recoveredEnergy),
            'demonWork': from_q16(self.demonWork),
            'workEnergy': from_q16(self.workEnergy),
            'energyLoss': from_q16(self.energyLoss)
        }


@dataclass
class QFactorState:
    """Q-Factor state (Lean: QFactorState)"""
    agentId: int  # UInt64
    balance: EnergyBalance
    qFactor: int  # Q16_16 - Current Q-factor
    targetQ: int  # Q16_16 - Target Q-factor (≈1.05)
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'agentId': self.agentId,
            'balance': self.balance.to_dict(),
            'qFactor': from_q16(self.qFactor),
            'targetQ': from_q16(self.targetQ)
        }


@dataclass
class QFactorAction:
    """Q-Factor optimization action (Lean: QFactorAction)"""
    agentId: int  # UInt64
    flashEnergyDelta: int  # Q16_16 - Change in flash energy
    enthalpyDelta: int  # Q16_16 - Change in enthalpy
    recoveredEnergyDelta: int  # Q16_16 - Change in recovered energy
    workEnergyDelta: int  # Q16_16 - Change in work energy
    energyLossDelta: int  # Q16_16 - Change in energy loss
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'agentId': self.agentId,
            'flashEnergyDelta': from_q16(self.flashEnergyDelta),
            'enthalpyDelta': from_q16(self.enthalpyDelta),
            'recoveredEnergyDelta': from_q16(self.recoveredEnergyDelta),
            'workEnergyDelta': from_q16(self.workEnergyDelta),
            'energyLossDelta': from_q16(self.energyLossDelta)
        }


@dataclass
class QFactorBind:
    """Q-Factor bind result (Lean: QFactorBind)"""
    lawful: bool
    qFactorBefore: int  # Q16_16
    qFactorAfter: int  # Q16_16
    energySurplus: int  # Q16_16
    invariant: str
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'lawful': self.lawful,
            'qFactorBefore': from_q16(self.qFactorBefore),
            'qFactorAfter': from_q16(self.qFactorAfter),
            'energySurplus': from_q16(self.energySurplus),
            'invariant': self.invariant
        }


# ═══════════════════════════════════════════════════════════════════════════
# Lean Function Implementations (verified by specification)
# ═══════════════════════════════════════════════════════════════════════════

def calculateQFactor(balance: EnergyBalance) -> int:
    """Calculate Q-Factor from energy balance (Lean: calculateQFactor)"""
    numerator = balance.flashEnergy + balance.enthalpy + balance.recoveredEnergy - balance.demonWork
    denominator = balance.workEnergy + balance.energyLoss
    if denominator > 0:
        return q16_div(numerator, denominator)
    return 0


def meetsTargetQ(state: QFactorState) -> bool:
    """Check if Q-Factor meets target threshold (Lean: meetsTargetQ)"""
    return q16_ge(state.qFactor, state.targetQ)


def hasNetEnergyGain(state: QFactorState) -> bool:
    """Check if Q-Factor indicates net energy gain (Lean: hasNetEnergyGain)"""
    return q16_gt(state.qFactor, Q16_ONE)  # Q > 1.0


def energySurplus(balance: EnergyBalance) -> int:
    """Calculate energy surplus (positive if net gain) (Lean: energySurplus)"""
    totalGain = balance.flashEnergy + balance.enthalpy + balance.recoveredEnergy
    totalCost = balance.demonWork + balance.workEnergy + balance.energyLoss
    return totalGain - totalCost


def energyEfficiencyFromBalance(balance: EnergyBalance) -> int:
    """Calculate energy efficiency: η = E_work / (E_work + E_loss) (Lean: energyEfficiencyFromBalance)"""
    totalEnergyCost = balance.workEnergy + balance.energyLoss
    if totalEnergyCost > 0:
        return q16_div(balance.workEnergy, totalEnergyCost)
    return 0


def recoveryRatio(balance: EnergyBalance) -> int:
    """Calculate recovery ratio: η_rec = E_recovered / (E_flash + E_enthalpy) (Lean: recoveryRatio)"""
    totalInputEnergy = balance.flashEnergy + balance.enthalpy
    if totalInputEnergy > 0:
        return q16_div(balance.recoveredEnergy, totalInputEnergy)
    return 0


def isQFactorActionLawful(state: QFactorState, action: QFactorAction) -> bool:
    """Check if Q-Factor action is lawful (Lean: isQFactorActionLawful)"""
    workPositive = q16_gt(state.balance.workEnergy + action.workEnergyDelta, 0)
    lossReasonable = q16_ge(action.energyLossDelta, -state.balance.energyLoss // 2)
    recoveredReasonable = q16_ge(action.recoveredEnergyDelta, 0) or q16_ge(action.recoveredEnergyDelta, -state.balance.recoveredEnergy // 2)
    return workPositive and lossReasonable and recoveredReasonable


def updateEnergyBalance(balance: EnergyBalance, action: QFactorAction) -> EnergyBalance:
    """Update energy balance from action (Lean: updateEnergyBalance)"""
    return EnergyBalance(
        flashEnergy=q16_add(balance.flashEnergy, action.flashEnergyDelta),
        enthalpy=q16_add(balance.enthalpy, action.enthalpyDelta),
        recoveredEnergy=q16_add(balance.recoveredEnergy, action.recoveredEnergyDelta),
        demonWork=balance.demonWork,  # Constant for now
        workEnergy=q16_add(balance.workEnergy, action.workEnergyDelta),
        energyLoss=q16_add(balance.energyLoss, action.energyLossDelta)
    )


def qFactorBind(state: QFactorState, action: QFactorAction) -> QFactorBind:
    """Bind primitive for Q-Factor optimization (Lean: qFactorBind)"""
    lawful = isQFactorActionLawful(state, action)
    newBalance = updateEnergyBalance(state.balance, action) if lawful else state.balance
    qFactorBefore = state.qFactor
    qFactorAfter = calculateQFactor(newBalance) if lawful else state.qFactor
    surplus = energySurplus(newBalance) if lawful else 0
    
    return QFactorBind(
        lawful=lawful,
        qFactorBefore=qFactorBefore,
        qFactorAfter=qFactorAfter,
        energySurplus=surplus,
        invariant="energy_balance_satisfied" if lawful else "energy_constraint_violated"
    )


class QFactorSystem:
    """
    Q-Factor energy balance system (Python shim wrapping Lean specification).
    
    All core logic is defined in 0-Core-Formalism/lean/Semantics/Semantics/QFactor.lean
    """
    
    def __init__(self):
        self.agentStates: Dict[int, QFactorState] = {}
        self.qFactorHistory: List[Dict[str, Any]] = []
        self.temporalSpatialSystem: Optional[TemporalSpatialRAMSystem] = None
        
        if _HAS_TS_RAM:
            self.temporalSpatialSystem = TemporalSpatialRAMSystem()
        
        print("[QFactor] Initialized (Lean specification)")
    
    def initializeAgent(self, agentId: int, flashEnergy: float, enthalpy: float, workEnergy: float, energyLoss: float, targetQ: float = 1.05) -> Dict[str, Any]:
        """Initialize agent Q-Factor state"""
        balance = EnergyBalance(
            flashEnergy=to_q16(flashEnergy),
            enthalpy=to_q16(enthalpy),
            recoveredEnergy=to_q16(0.0),
            demonWork=to_q16(20.0),  # Landauer limit
            workEnergy=to_q16(workEnergy),
            energyLoss=to_q16(energyLoss)
        )
        
        qFactor = calculateQFactor(balance)
        state = QFactorState(
            agentId=agentId,
            balance=balance,
            qFactor=qFactor,
            targetQ=to_q16(targetQ)
        )
        
        self.agentStates[agentId] = state
        
        # Register node in temporal-spatial system if available
        if self.temporalSpatialSystem:
            self.temporalSpatialSystem.registerNode(
                nodeId=agentId,
                x=0.0, y=0.0, z=0.0,  # Origin for now
                physicalRAM=workEnergy,
                currentTime=0.0
            )
        
        return {
            'agentId': agentId,
            'state': state.to_dict()
        }
    
    def submitQFactorAction(self, action: QFactorAction) -> Dict[str, Any]:
        """Submit Q-Factor action for processing (Lean specification)"""
        if action.agentId not in self.agentStates:
            return {'error': 'Agent not initialized'}
        
        currentState = self.agentStates[action.agentId]
        bindResult = qFactorBind(currentState, action)
        
        if bindResult.lawful:
            newBalance = updateEnergyBalance(currentState.balance, action)
            
            # Adjust work energy based on temporal-spatial resources if available
            if self.temporalSpatialSystem:
                tsResources = self.temporalSpatialSystem.getNodeResources(action.agentId)
                if tsResources:
                    # Add temporal-spatial RAM to work energy (proximity bonus)
                    spatialBonus = tsResources['resources']['spatialRAM'] * 0.1
                    temporalBonus = tsResources['resources']['temporalRAM'] * 0.01
                    newBalance = EnergyBalance(
                        flashEnergy=newBalance.flashEnergy,
                        enthalpy=newBalance.enthalpy,
                        recoveredEnergy=newBalance.recoveredEnergy,
                        demonWork=newBalance.demonWork,
                        workEnergy=newBalance.workEnergy + to_q16(spatialBonus + temporalBonus),
                        energyLoss=newBalance.energyLoss
                    )
            
            newState = QFactorState(
                agentId=currentState.agentId,
                balance=newBalance,
                qFactor=calculateQFactor(newBalance),
                targetQ=currentState.targetQ
            )
            self.agentStates[action.agentId] = newState
            
            # Record Q-Factor history
            self.qFactorHistory.append({
                'agentId': action.agentId,
                'action': action.to_dict(),
                'bindResult': bindResult.to_dict(),
                'stateBefore': currentState.to_dict(),
                'stateAfter': newState.to_dict(),
                'timestamp': time.time()
            })
        
        return {
            'success': bindResult.lawful,
            'bindResult': bindResult.to_dict(),
            'state': self.agentStates[action.agentId].to_dict() if bindResult.lawful else currentState.to_dict()
        }
    
    def getAgentState(self, agentId: int) -> Optional[Dict[str, Any]]:
        """Get current agent Q-Factor state"""
        if agentId in self.agentStates:
            return self.agentStates[agentId].to_dict()
        return None
    
    def getQFactorHistory(self, agentId: Optional[int] = None, limit: int = 10) -> List[Dict[str, Any]]:
        """Get Q-Factor history"""
        if agentId is not None:
            filtered = [h for h in self.qFactorHistory if h['agentId'] == agentId]
            return filtered[-limit:]
        return self.qFactorHistory[-limit:]
    
    def calculateSystemEfficiency(self, agentId: int) -> Optional[float]:
        """Calculate system efficiency for agent"""
        if agentId not in self.agentStates:
            return None
        
        state = self.agentStates[agentId]
        efficiency = energyEfficiencyFromBalance(state.balance)
        return from_q16(efficiency)
    
    def printSystemState(self):
        """Print system state"""
        print("\n" + "="*60)
        print("Q-FACTOR ENERGY BALANCE SYSTEM STATE")
        print("="*60)
        
        print(f"\n📊 Active Agents: {len(self.agentStates)}")
        for agentId, state in self.agentStates.items():
            print(f"\n  Agent {agentId}:")
            print(f"    Q-Factor: {from_q16(state.qFactor):.3f} (target: {from_q16(state.targetQ):.3f})")
            print(f"    Net energy gain: {hasNetEnergyGain(state)}")
            print(f"    Energy surplus: {from_q16(energySurplus(state.balance)):.3f}")
            print(f"    Efficiency: {from_q16(energyEfficiencyFromBalance(state.balance)):.3f}")
            print(f"    Recovery ratio: {from_q16(recoveryRatio(state.balance)):.3f}")
            
            # Add temporal-spatial resource display if available
            if self.temporalSpatialSystem:
                tsResources = self.temporalSpatialSystem.getNodeResources(agentId)
                if tsResources:
                    print(f"    Temporal-Spatial RAM: {tsResources['resources']['totalRAM']:.3f}")
                    print(f"      - Temporal: {tsResources['resources']['temporalRAM']:.3f}")
                    print(f"      - Spatial: {tsResources['resources']['spatialRAM']:.3f}")
        
        print(f"\n📜 Q-Factor History: {len(self.qFactorHistory)} entries")
        
        print("\n" + "="*60)


def main():
    """Test Q-Factor system"""
    system = QFactorSystem()
    
    print("[Test 1] Initialize agent with energy balance...")
    initResult = system.initializeAgent(
        agentId=1,
        flashEnergy=100.0,
        enthalpy=50.0,
        workEnergy=80.0,
        energyLoss=10.0,
        targetQ=1.05
    )
    print(f"  Agent 1 initialized: Q-Factor={initResult['state']['qFactor']:.3f}")
    
    print("\n[Test 2] Submit Q-Factor action (increase recovered energy)...")
    action1 = QFactorAction(
        agentId=1,
        flashEnergyDelta=to_q16(10.0),
        enthalpyDelta=to_q16(5.0),
        recoveredEnergyDelta=to_q16(15.0),
        workEnergyDelta=to_q16(10.0),
        energyLossDelta=to_q16(2.0)
    )
    result1 = system.submitQFactorAction(action1)
    print(f"  Result: Lawful={result1['success']}")
    if result1['success']:
        print(f"  Q-Factor before: {result1['bindResult']['qFactorBefore']:.3f}")
        print(f"  Q-Factor after: {result1['bindResult']['qFactorAfter']:.3f}")
        print(f"  Energy surplus: {result1['bindResult']['energySurplus']:.3f}")
    
    print("\n[Test 3] Submit Q-Factor action (optimize energy loss)...")
    action2 = QFactorAction(
        agentId=1,
        flashEnergyDelta=to_q16(0.0),
        enthalpyDelta=to_q16(0.0),
        recoveredEnergyDelta=to_q16(5.0),
        workEnergyDelta=to_q16(5.0),
        energyLossDelta=to_q16(-5.0)  # Reduce energy loss
    )
    result2 = system.submitQFactorAction(action2)
    print(f"  Result: Lawful={result2['success']}")
    if result2['success']:
        print(f"  Q-Factor after: {result2['state']['qFactor']:.3f}")
        print(f"  Energy surplus: {from_q16(energySurplus(system.agentStates[1].balance)):.3f}")
    
    print("\n[Test 4] Submit invalid Q-Factor action (negative work energy)...")
    action3 = QFactorAction(
        agentId=1,
        flashEnergyDelta=to_q16(0.0),
        enthalpyDelta=to_q16(0.0),
        recoveredEnergyDelta=to_q16(0.0),
        workEnergyDelta=to_q16(-100.0),  # Invalid: negative work energy
        energyLossDelta=to_q16(0.0)
    )
    result3 = system.submitQFactorAction(action3)
    print(f"  Result: Lawful={result3['success']}")
    print(f"  Invariant: {result3['bindResult']['invariant']}")
    
    print("\n[Test 5] Check if target Q is met...")
    state = system.agentStates[1]
    targetMet = meetsTargetQ(state)
    netGain = hasNetEnergyGain(state)
    print(f"  Target Q met: {targetMet}")
    print(f"  Net energy gain: {netGain}")
    
    print("\n[Test 6] Calculate system efficiency...")
    efficiency = system.calculateSystemEfficiency(agentId=1)
    if efficiency is not None:
        print(f"  System efficiency: {efficiency:.3f}")
    
    print("\n[System State]")
    system.printSystemState()


if __name__ == '__main__':
    main()