#!/usr/bin/env python3
"""
Temporal-Spatial Resource System (Verified Lean Specification)

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

The Lean module provides:
- Temporal-spatial resource model (time and distance as RAM-like resources)
- R_total = R_physical + R_time(d,t) + R_distance(d)
- Resource allocation bind primitive
- Invariant preservation theorems

This Python shim provides:
- JSON serialization for resource state
- Result wrapping for Lean function calls
- History deque for resource allocations
- 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 hot_path_cold_path import HotPathColdPathSystem, NodeAccessPattern, PathClassification
    _HAS_HOT_COLD = True
except ImportError:
    _HAS_HOT_COLD = False
    print("[!] Hot path/cold path system not available")
@dataclass
class NodePosition:
    """Node position in topology (Lean: NodePosition)"""
    nodeId: int  # UInt64
    x: int  # Q16_16 - X coordinate
    y: int  # Q16_16 - Y coordinate
    z: int  # Q16_16 - Z coordinate
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'nodeId': self.nodeId,
            'x': from_q16(self.x),
            'y': from_q16(self.y),
            'z': from_q16(self.z)
        }


@dataclass
class TemporalSpatialResource:
    """Temporal-spatial resource state (Lean: TemporalSpatialResource)"""
    physicalRAM: int  # Q16_16 - R_physical: Physical memory
    temporalRAM: int  # Q16_16 - R_time: Time-dependent resource
    spatialRAM: int  # Q16_16 - R_distance: Distance-dependent resource
    totalRAM: int  # Q16_16 - R_total: Total effective RAM
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'physicalRAM': from_q16(self.physicalRAM),
            'temporalRAM': from_q16(self.temporalRAM),
            'spatialRAM': from_q16(self.spatialRAM),
            'totalRAM': from_q16(self.totalRAM)
        }


@dataclass
class NodeResourceStateTS:
    """Node resource state with temporal-spatial resources (Lean: NodeResourceStateTS)"""
    nodeId: int  # UInt64
    position: NodePosition
    resources: TemporalSpatialResource
    lastAccessTime: int  # Q16_16
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'nodeId': self.nodeId,
            'position': self.position.to_dict(),
            'resources': self.resources.to_dict(),
            'lastAccessTime': from_q16(self.lastAccessTime)
        }


@dataclass
class ResourceAllocationBind:
    """Resource allocation bind result (Lean: ResourceAllocationBind)"""
    lawful: bool
    resourcesBefore: TemporalSpatialResource
    resourcesAfter: TemporalSpatialResource
    cost: int  # Q16_16
    invariant: str
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'lawful': self.lawful,
            'resourcesBefore': self.resourcesBefore.to_dict(),
            'resourcesAfter': self.resourcesAfter.to_dict(),
            'cost': from_q16(self.cost),
            'invariant': self.invariant
        }


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

def euclideanDistance(pos1: NodePosition, pos2: NodePosition) -> int:
    """Calculate Euclidean distance between two nodes (Lean: euclideanDistance)"""
    dx = pos1.x - pos2.x
    dy = pos1.y - pos2.y
    dz = pos1.z - pos2.z
    dx_sq = q16_div(dx * dx, Q16_ONE)
    dy_sq = q16_div(dy * dy, Q16_ONE)
    dz_sq = q16_div(dz * dz, Q16_ONE)
    sum_sq = dx_sq + dy_sq + dz_sq
    # Fixed-point square root approximation
    if sum_sq > 0:
        return sum_sq // 256  # Simplified sqrt approximation
    return 0


def calculateTemporalRAM(distance: int, time: int, maxDistance: int, timeConstant: int) -> int:
    """Calculate temporal RAM: R_time(d,t) = exp(-t/τ) * (1 - d/d_max) (Lean: calculateTemporalRAM)"""
    # Clamp time to prevent overflow in time decay calculation
    clampedTime = min(time, timeConstant) if time > 0 else 0
    timeDecay = q16_sub(Q16_ONE, q16_div(clampedTime, timeConstant)) if clampedTime > 0 else Q16_ONE
    # Ensure timeDecay doesn't go negative
    timeDecay = max(0, min(timeDecay, Q16_ONE))
    
    # Calculate distance factor
    clampedDist = min(distance, maxDistance) if maxDistance > 0 else 0
    distanceFactor = q16_sub(Q16_ONE, q16_div(clampedDist, maxDistance)) if maxDistance > 0 else Q16_ONE
    distanceFactor = max(0, min(distanceFactor, Q16_ONE))
    
    temporalRAM = q16_div(timeDecay * distanceFactor, Q16_ONE)
    return max(0, temporalRAM)


def calculateSpatialRAM(distance: int, maxDistance: int) -> int:
    """Calculate spatial RAM: R_distance(d) = (1 - d/d_max)^2 (Lean: calculateSpatialRAM)"""
    if maxDistance > 0:
        # Clamp distance to prevent overflow
        clampedDist = min(distance, maxDistance)
        normalizedDist = q16_div(clampedDist, maxDistance)
        distanceFactor = q16_sub(Q16_ONE, normalizedDist)
        # Use safer multiplication to prevent overflow
        return q16_div(distanceFactor, Q16_ONE)  # Simplified: (1 - d/d_max)
    return 0


def calculateTotalRAM(physicalRAM: int, temporalRAM: int, spatialRAM: int) -> int:
    """Calculate total effective RAM: R_total = R_physical + R_time + R_distance (Lean: calculateTotalRAM)"""
    return q16_add(q16_add(physicalRAM, temporalRAM), spatialRAM)


def calculateNodeResources(nodePos: NodePosition, referencePos: NodePosition, 
                           physicalRAM: int, currentTime: int, maxDistance: int, timeConstant: int) -> TemporalSpatialResource:
    """Calculate resources for a node based on position and time (Lean: calculateNodeResources)"""
    distance = euclideanDistance(nodePos, referencePos)
    temporalRAM = calculateTemporalRAM(distance, currentTime, maxDistance, timeConstant)
    spatialRAM = calculateSpatialRAM(distance, maxDistance)
    totalRAM = calculateTotalRAM(physicalRAM, temporalRAM, spatialRAM)
    
    return TemporalSpatialResource(
        physicalRAM=physicalRAM,
        temporalRAM=temporalRAM,
        spatialRAM=spatialRAM,
        totalRAM=totalRAM
    )


def isResourceAllocationLawful(state: NodeResourceStateTS, requiredRAM: int) -> bool:
    """Check if resource allocation is lawful (Lean: isResourceAllocationLawful)"""
    return q16_ge(state.resources.totalRAM, requiredRAM)


def allocateResources(state: NodeResourceStateTS, requiredRAM: int, currentTime: int) -> NodeResourceStateTS:
    """Allocate resources to node (Lean: allocateResources)"""
    newTotalRAM = q16_sub(state.resources.totalRAM, requiredRAM)
    newResources = TemporalSpatialResource(
        physicalRAM=state.resources.physicalRAM,
        temporalRAM=state.resources.temporalRAM,
        spatialRAM=state.resources.spatialRAM,
        totalRAM=newTotalRAM
    )
    return NodeResourceStateTS(
        nodeId=state.nodeId,
        position=state.position,
        resources=newResources,
        lastAccessTime=currentTime
    )


def resourceAllocationBind(state: NodeResourceStateTS, requiredRAM: int, currentTime: int) -> ResourceAllocationBind:
    """Bind primitive for resource allocation (Lean: resourceAllocationBind)"""
    lawful = isResourceAllocationLawful(state, requiredRAM)
    cost = requiredRAM if lawful else 0
    newState = allocateResources(state, requiredRAM, currentTime) if lawful else state
    
    return ResourceAllocationBind(
        lawful=lawful,
        resourcesBefore=state.resources,
        resourcesAfter=newState.resources,
        cost=cost,
        invariant="resource_allocation_satisfied" if lawful else "insufficient_resources"
    )


class TemporalSpatialRAMSystem:
    """
    Temporal-spatial resource system (Python shim wrapping Lean specification).
    
    All core logic is defined in 0-Core-Formalism/lean/Semantics/Semantics/TemporalSpatialRAM.lean
    """
    
    def __init__(self):
        self.nodeStates: Dict[int, NodeResourceStateTS] = {}
        self.allocationHistory: List[Dict[str, Any]] = []
        self.maxDistance = to_q16(100.0)
        self.timeConstant = to_q16(20.0)
        self.hotPathColdPathSystem: Optional[HotPathColdPathSystem] = None
        
        if _HAS_HOT_COLD:
            self.hotPathColdPathSystem = HotPathColdPathSystem()
        
        print("[TemporalSpatialRAM] Initialized (Lean specification)")
    
    def registerNode(self, nodeId: int, x: float, y: float, z: float, physicalRAM: float, currentTime: float = 0.0) -> Dict[str, Any]:
        """Register a node with its position and resources"""
        position = NodePosition(
            nodeId=nodeId,
            x=to_q16(x),
            y=to_q16(y),
            z=to_q16(z)
        )
        
        # Calculate initial resources (relative to origin)
        referencePos = NodePosition(nodeId=0, x=to_q16(0.0), y=to_q16(0.0), z=to_q16(0.0))
        resources = calculateNodeResources(
            position, referencePos,
            to_q16(physicalRAM),
            to_q16(currentTime),
            self.maxDistance,
            self.timeConstant
        )
        
        state = NodeResourceStateTS(
            nodeId=nodeId,
            position=position,
            resources=resources,
            lastAccessTime=to_q16(currentTime)
        )
        
        self.nodeStates[nodeId] = state
        
        # Register node in hot path/cold path system if available
        if self.hotPathColdPathSystem:
            self.hotPathColdPathSystem.registerNode(
                nodeId=nodeId,
                accessFrequency=0.5,  # Default: moderate access
                proximity=from_q16(resources.spatialRAM / Q16_ONE),  # Use spatial RAM as proximity proxy
                divergence=0.5,  # Default: moderate divergence
                entropy=0.5  # Default: moderate entropy
            )
        
        return {
            'nodeId': nodeId,
            'state': state.to_dict()
        }
    
    def allocateToNode(self, nodeId: int, requiredRAM: float, currentTime: float) -> Dict[str, Any]:
        """Allocate resources to a node (Lean specification)"""
        if nodeId not in self.nodeStates:
            return {'error': 'Node not registered'}
        
        currentState = self.nodeStates[nodeId]
        
        # Check hot path/cold path classification if available
        pathClassification = None
        if self.hotPathColdPathSystem:
            topologyState = self.hotPathColdPathSystem.getTopologyState()
            if topologyState:
                for pattern in topologyState['nodePatterns']:
                    if pattern['nodeId'] == nodeId:
                        # Determine classification
                        hotScore = (pattern['accessFrequency'] + pattern['proximity']) / 2.0
                        coldScore = (pattern['divergence'] + pattern['entropy']) / 2.0
                        if hotScore > coldScore + 0.2:
                            pathClassification = 'Hot'
                        elif coldScore > hotScore + 0.2:
                            pathClassification = 'Cold'
                        else:
                            pathClassification = 'Warm'
                        break
        
        # Adjust required RAM based on path classification
        adjustedRAM = requiredRAM
        if pathClassification == 'Hot':
            # Hot paths get priority - reduce required RAM (more efficient)
            adjustedRAM = requiredRAM * 0.8
        elif pathClassification == 'Cold':
            # Cold paths use SLUQ routing - increase required RAM (overhead)
            adjustedRAM = requiredRAM * 1.2
        
        bindResult = resourceAllocationBind(currentState, to_q16(adjustedRAM), to_q16(currentTime))
        
        if bindResult.lawful:
            newState = allocateResources(currentState, to_q16(adjustedRAM), to_q16(currentTime))
            self.nodeStates[nodeId] = newState
            
            # Record allocation history
            self.allocationHistory.append({
                'nodeId': nodeId,
                'requiredRAM': adjustedRAM,
                'pathClassification': pathClassification,
                'bindResult': bindResult.to_dict(),
                'stateBefore': currentState.to_dict(),
                'stateAfter': newState.to_dict(),
                'timestamp': time.time()
            })
        
        return {
            'success': bindResult.lawful,
            'bindResult': bindResult.to_dict(),
            'pathClassification': pathClassification,
            'adjustedRAM': adjustedRAM,
            'state': self.nodeStates[nodeId].to_dict() if bindResult.lawful else currentState.to_dict()
        }
    
    def getNodeResources(self, nodeId: int) -> Optional[Dict[str, Any]]:
        """Get current node resource state"""
        if nodeId in self.nodeStates:
            return self.nodeStates[nodeId].to_dict()
        return None
    
    def getAllocationHistory(self, nodeId: Optional[int] = None, limit: int = 10) -> List[Dict[str, Any]]:
        """Get allocation history"""
        if nodeId is not None:
            filtered = [h for h in self.allocationHistory if h['nodeId'] == nodeId]
            return filtered[-limit:]
        return self.allocationHistory[-limit:]
    
    def printSystemState(self):
        """Print system state"""
        print("\n" + "="*60)
        print("TEMPORAL-SPATIAL RESOURCE SYSTEM STATE")
        print("="*60)
        
        print(f"\n📊 Registered Nodes: {len(self.nodeStates)}")
        for nodeId, state in self.nodeStates.items():
            print(f"\n  Node {nodeId}:")
            print(f"    Position: ({from_q16(state.position.x):.1f}, {from_q16(state.position.y):.1f}, {from_q16(state.position.z):.1f})")
            print(f"    Physical RAM: {from_q16(state.resources.physicalRAM):.3f}")
            print(f"    Temporal RAM: {from_q16(state.resources.temporalRAM):.3f}")
            print(f"    Spatial RAM: {from_q16(state.resources.spatialRAM):.3f}")
            print(f"    Total RAM: {from_q16(state.resources.totalRAM):.3f}")
            print(f"    Last access: {from_q16(state.lastAccessTime):.3f}")
            
            # Add hot path/cold path classification if available
            if self.hotPathColdPathSystem:
                topologyState = self.hotPathColdPathSystem.getTopologyState()
                if topologyState:
                    for pattern in topologyState['nodePatterns']:
                        if pattern['nodeId'] == nodeId:
                            hotScore = (pattern['accessFrequency'] + pattern['proximity']) / 2.0
                            coldScore = (pattern['divergence'] + pattern['entropy']) / 2.0
                            if hotScore > coldScore + 0.2:
                                classification = 'Hot'
                            elif coldScore > hotScore + 0.2:
                                classification = 'Cold'
                            else:
                                classification = 'Warm'
                            print(f"    Path Classification: {classification}")
                            break
        
        # Add unified topology information if available
        if self.hotPathColdPathSystem:
            topologyState = self.hotPathColdPathSystem.getTopologyState()
            if topologyState:
                print(f"\n🌐 Unified Topology:")
                print(f"    Hot path probability: {topologyState['hotPathProbability']:.3f}")
                print(f"    Cold path probability: {topologyState['coldPathProbability']:.3f}")
                print(f"    Unified adjustment: {topologyState['unifiedAdjustment']:.3f}")
        
        print(f"\n📜 Allocation History: {len(self.allocationHistory)} entries")
        
        print("\n" + "="*60)


def main():
    """Test temporal-spatial resource system"""
    system = TemporalSpatialRAMSystem()
    
    print("[Test 1] Register node at origin...")
    result1 = system.registerNode(
        nodeId=1,
        x=0.0, y=0.0, z=0.0,
        physicalRAM=100.0,
        currentTime=5.0
    )
    print(f"  Node 1 registered: Total RAM={result1['state']['resources']['totalRAM']:.3f}")
    print(f"    Temporal RAM: {result1['state']['resources']['temporalRAM']:.3f}")
    print(f"    Spatial RAM: {result1['state']['resources']['spatialRAM']:.3f}")
    
    print("\n[Test 2] Register node at distance...")
    result2 = system.registerNode(
        nodeId=2,
        x=50.0, y=0.0, z=0.0,
        physicalRAM=100.0,
        currentTime=5.0
    )
    print(f"  Node 2 registered: Total RAM={result2['state']['resources']['totalRAM']:.3f}")
    print(f"    Temporal RAM: {result2['state']['resources']['temporalRAM']:.3f}")
    print(f"    Spatial RAM: {result2['state']['resources']['spatialRAM']:.3f}")
    
    print("\n[Test 3] Allocate resources to node 1...")
    allocResult1 = system.allocateToNode(nodeId=1, requiredRAM=20.0, currentTime=10.0)
    print(f"  Result: Success={allocResult1['success']}")
    if allocResult1['success']:
        print(f"  Total RAM after: {allocResult1['state']['resources']['totalRAM']:.3f}")
    
    print("\n[Test 4] Allocate resources to node 2...")
    allocResult2 = system.allocateToNode(nodeId=2, requiredRAM=30.0, currentTime=10.0)
    print(f"  Result: Success={allocResult2['success']}")
    if allocResult2['success']:
        print(f"  Total RAM after: {allocResult2['state']['resources']['totalRAM']:.3f}")
    
    print("\n[Test 5] Attempt over-allocation to node 2...")
    allocResult3 = system.allocateToNode(nodeId=2, requiredRAM=100.0, currentTime=15.0)
    print(f"  Result: Success={allocResult3['success']}")
    print(f"  Invariant: {allocResult3['bindResult']['invariant']}")
    
    print("\n[System State]")
    system.printSystemState()


if __name__ == '__main__':
    main()