import uuid
from typing import List, Any, Optional, Dict
import time
import math

class Atom:
    """
    Base class for all Periodic Table Elements.
    Represents the fundamental unit with mass (heat/complexity) and charge (valence).
    """
    def __init__(self, symbol: str, name: str, domain: str = "General"):
        self.id = str(uuid.uuid4())
        self.symbol = symbol  # e.g., 'Pr', 'To'
        self.name = name
        self.timestamp = time.time()
        self.domain = domain # e.g. "Physics", "Code", "Logic" (Video 26)
        self.heat = 0.0  # Dissonance/Energy
        self.valence = [] # Connections to other atoms (Hyper-Edges)

    def resonate(self) -> float:
        """Returns the current heat/resonance level."""
        return self.heat

    def __repr__(self):
        return f"[{self.symbol}:{self.name}]"

class Vector(Atom):
    """
    Symbol: Ve
    Represents Semantic Gravity (Embeddings).
    """
    def __init__(self, embedding: List[float]):
        super().__init__('Ve', 'Vector')
        self.embedding = embedding
        self.heat = sum(abs(x) for x in embedding) / len(embedding) if embedding else 0.0

    def distance(self, other: 'Vector') -> float:
        """Calculates cosine distance (Heat) between two vectors."""
        # Simplified for now, assuming normalized
        dot = sum(a*b for a, b in zip(self.embedding, other.embedding))
        return 1.0 - dot

class Token(Atom):
    """
    Symbol: To
    The atomic unit of semantic meaning (Text, Image Patch).
    """
    def __init__(self, content: Any, source: str = "user"):
        super().__init__('To', 'Token')
        self.content = content
        self.source = source
        # Heuristic heat: length of content
        self.heat = len(str(content)) * 0.01

class Actuator(Atom):
    """
    Symbol: Ac
    Represents a Physical Output Device (Motor, Switch, API).
    """
    def __init__(self, device_id: str, state: dict = None):
        super().__init__('Ac', device_id)
        self.device_id = device_id
        self.state = state or {"power": "off"}
        self.heat = 0.5 # Physical actions generate heat/friction

class Audio(Atom):
    """
    Symbol: Au
    Represents Sound/Voice data.
    Generated by TTS or captured from Microphone.
    """
    def __init__(self, content: bytes, source: str = "TTS"):
        super().__init__('Au', 'Audio')
        self.content = content # Raw bytes or path
        self.source = source
        self.duration = 0.0 # Placeholder
        self.heat = len(content) * 0.001 # Density of data

class Model(Atom):
    """
    Symbol: Mo
    The dense compute node (LLM, Diffusion, etc.).
    """
    def __init__(self, model_id: str, architecture: str = "Transformer", local: bool = True):
        super().__init__('Mo', model_id)
        self.local = local
        self.architecture = architecture # e.g. "Transformer", "Diffusion", "SSM"
        self.status = "idle"

    def process(self, input_atoms: List[Atom]) -> List[Atom]:
        """
        Transforms input atoms (Prompt) into output atoms (Tokens).
        This is where the actual LLM call happens.
        """
        # Placeholder for actual inference logic
        self.status = "processing"
        output_content = f"Processed {len(input_atoms)} atoms via {self.name}"
        time.sleep(0.1) # Simulate compute
        self.status = "idle"
        return [Token(output_content, source=self.name)]

class Handoff(Atom):
    """
    Symbol: Ha
    Represents a transfer of control to a specific Tool or Agent.
    Triggered when Heat > Threshold.
    """
    def __init__(self, target_tool_name: str, reason: str):
        super().__init__('Ha', target_tool_name)
        self.target = target_tool_name
        self.reason = reason
        self.status = "pending"

    def resolve(self) -> Token:
        """
        Executing the handoff means returning a Token that summarizes the action.
        The actual execution happens in the Router via the Tool map.
        """
        self.status = "active"
        return Token(f"[HANDOFF] Transferring to {self.target} due to: {self.reason}", source="Router")

class Tool(Atom):
    """
    Symbol: Fu (Function) / Co (Code)
    Executable logic that binds to specific Tokens.
    """
    def __init__(self, name: str, func: callable, description: str):
        super().__init__('Fu', name)
        self.func = func
        self.description = description

    def execute(self, *args, **kwargs) -> Atom:
        result = self.func(*args, **kwargs)
        return Token(result, source=f"Tool:{self.name}")

class Tensor(Atom):
    """
    Symbol: Te
    Represents a linear transformation (Matrix) or higher-order tensor.
    Agents apply Tensors to Vectors to change the State.
    """
    def __init__(self, matrix: List[List[float]], name: str = "Transform"):
        super().__init__('Te', name)
        self.matrix = matrix 
        # Heat = Frobenius norm (magnitude of change)
        self.heat = math.sqrt(sum(x*x for row in matrix for x in row))

    def apply(self, vector: Vector) -> Vector:
        """
        Applies the tensor (matrix multiplication) to a vector.
        v_new = T * v
        """
        if not vector.embedding:
            return vector
            
        # Simple Matrix-Vector multiplication
        new_vec = []
        for row in self.matrix:
            # Handle dimension mismatch by padding or truncating conceptually
            # For this prototype, we assume row len fits vector len or we dot product available dims
            dot = sum(a*b for a, b in zip(row, vector.embedding))
            new_vec.append(dot)
            
        return Vector(new_vec)

# Prime Resonance Mapping (Gödel Numbering)
PRIME_DOMAINS = {
    "General": 2,
    "Physics": 3,
    "Code": 5,
    "Logic": 7,
    "Prompt": 11,
    "Inference": 13,
    "External_Knowledge": 17,
    "Audio": 19,
    "Vision": 23
}

class ManifoldState:
    """
    Symbol: St (State)
    The container for the recursive loop.
    Now tracks Prime Resonance (Gödel Numbering).
    """
    def __init__(self):
        self.atoms: List[Atom] = []
        self.stress: float = 1.0 
        self.deformation_gradient: List[float] = [] 
        
        # Prime Resonance (The scalar product of context history)
        self.resonance_product: int = 1 
        
        # Context Graph (Video 15: Nodes & Edges)
        self.graph = {
            "nodes": {}, 
            "edges": [] 
        }

    def inject(self, atom: Atom, parent_id: Optional[str] = None, relation: str = "follows"):
        self.atoms.append(atom)
        
        # 1. Update Stress (Physics)
        force = atom.heat
        area = len(self.atoms)
        self.stress = force / (area + 1e-9) 
        self.deformation_gradient.append(self.stress)

        # 2. Update Resonance (Number Theory)
        prime_val = PRIME_DOMAINS.get(atom.domain, 2)
        self.resonance_product *= prime_val

        # 3. Update Graph (Topology)
        self.graph["nodes"][atom.id] = {
            "type": atom.symbol,
            "name": atom.name,
            "heat": atom.heat,
            "domain": atom.domain,
            "prime": prime_val
        }
        
        if parent_id and parent_id in self.graph["nodes"]:
            self.graph["edges"].append({
                "source": parent_id,
                "target": atom.id,
                "relation": relation,
                "weight": self.stress
            })

    def stabilize(self) -> bool:
        """Returns True if Stress is below threshold (Equilibrium)."""
        return self.stress < 0.1