#!/usr/bin/env python3
"""
Grok Logic Layer - ISL-Safe Implementation
Structured reasoning with artifacts, not text-based chain-of-thought
"""

from typing import List, Dict, Any
from lego_moe import LegoMoE
import re
import json

class GrokLogicLayer:
    """
    ISL-compatible reasoning layer
    Uses structured artifacts, not prose reasoning
    """
    
    def __init__(self, lego_moe: LegoMoE):
        self.moe = lego_moe
        self.reasoning_artifacts = {}
        print("✓ Grok Logic Layer initialized (ISL-safe mode)")
    
    def parse_intent(self, query: str) -> Dict[str, Any]:
        """
        Parse query into structured intent vector
        
        Returns structured intent, not text
        """
        query_lower = query.lower()
        
        # Classify type
        if any(word in query_lower for word in ['calculate', 'compute', 'multiply', 'add', 'divide', 'subtract', 'times']):
            qtype = 'math_arithmetic'
        elif any(word in query_lower for word in ['if', 'then', 'all', 'some', 'implies']):
            qtype = 'logic_deduction'
        elif any(word in query_lower for word in ['code', 'function', 'debug', 'fix']):
            qtype = 'code_analysis'
        else:
            qtype = 'general_qa'
        
        # Extract entities
        numbers = [float(n) for n in re.findall(r'\d+\.?\d*', query)]
        
        # Extract operators (including word forms)
        operators = re.findall(r'[\+\-\*/×÷]', query)
        if 'times' in query_lower or 'multiply' in query_lower:
            operators.append('*')
        if 'divided' in query_lower or 'divide' in query_lower:
            operators.append('/')
        if 'plus' in query_lower:
            operators.append('+')
        if 'minus' in query_lower:
            operators.append('-')
        
        # Determine complexity
        complexity_score = len(query.split()) + len(numbers) * 2
        
        # Force math type if we have numbers and operators
        if len(numbers) >= 2 and len(operators) >= 1:
            qtype = 'math_arithmetic'
        
        return {
            "type": qtype,
            "complexity": complexity_score,
            "entities": {
                "numbers": numbers,
                "operators": operators
            },
            "requires_decomposition": complexity_score > 15
        }
    
    def decompose_to_graph(self, query: str, intent: Dict) -> Dict[str, Any]:
        """
        Decompose into structured task graph
        NOT text-based sub-problems
        """
        if not intent['requires_decomposition']:
            return {
                "nodes": [{
                    "id": 0,
                    "task": query,
                    "type": intent['type'],
                    "entities": intent['entities']  # Pass entities through
                }],
                "edges": [],
                "execution_order": [0]
            }
        
        # For math: decompose into operations
        if intent['type'] == 'math_arithmetic':
            return self._decompose_math(query, intent)
        
        # For logic: decompose into premises and conclusion
        elif intent['type'] == 'logic_deduction':
            return self._decompose_logic(query)
        
        # Default: single node
        return {
            "nodes": [{
                "id": 0,
                "task": query,
                "type": intent['type'],
                "entities": intent['entities']
            }],
            "edges": [],
            "execution_order": [0]
        }
    
    def _decompose_math(self, query: str, intent: Dict) -> Dict:
        """Decompose math problem into operation graph"""
        numbers = intent['entities']['numbers']
        operators = intent['entities']['operators']
        
        if not numbers or not operators:
            return {
                "nodes": [{"id": 0, "task": query, "type": "math_arithmetic"}],
                "edges": [],
                "execution_order": [0]
            }
        
        # Build operation sequence
        nodes = []
        for i, (num, op) in enumerate(zip(numbers[:-1], operators)):
            nodes.append({
                "id": i,
                "operation": op,
                "operands": [num, numbers[i+1]],
                "type": "math_operation"
            })
        
        return {
            "nodes": nodes,
            "edges": [(i, i+1) for i in range(len(nodes)-1)],
            "execution_order": list(range(len(nodes)))
        }
    
    def _decompose_logic(self, query: str) -> Dict:
        """Decompose logic problem into premises and conclusion"""
        # Simple heuristic: split on periods and "therefore"
        sentences = [s.strip() for s in re.split(r'[.?!]', query) if s.strip()]
        
        nodes = []
        for i, sentence in enumerate(sentences):
            node_type = "conclusion" if any(word in sentence.lower() for word in ['therefore', 'thus', 'so']) else "premise"
            nodes.append({
                "id": i,
                "statement": sentence,
                "type": node_type
            })
        
        return {
            "nodes": nodes,
            "edges": [(i, len(nodes)-1) for i in range(len(nodes)-1)],
            "execution_order": list(range(len(nodes)))
        }
    
    def plan_execution(self, task_graph: Dict) -> Dict[str, Any]:
        """
        Create execution DAG with expert routing
        """
        execution_plan = {
            "steps": [],
            "expert_routing": {},
            "verification_rules": []
        }
        
        for node in task_graph['nodes']:
            step = {
                "node_id": node['id'],
                "expert": self._select_expert(node),
                "input": node,
                "verification": self._get_verification_rules(node)
            }
            execution_plan['steps'].append(step)
        
        return execution_plan
    
    def _select_expert(self, node: Dict) -> str:
        """Select expert for node"""
        node_type = node.get('type', 'general')
        
        if 'math' in node_type:
            return 'mathematical_reasoning'
        elif 'logic' in node_type:
            return 'logical_inference'
        elif 'code' in node_type:
            return 'code_generation'
        else:
            return 'general'
    
    def _get_verification_rules(self, node: Dict) -> List[str]:
        """Get verification rules for node"""
        rules = ["bounds_check"]
        
        if node.get('type') == 'math_operation':
            rules.extend(["reverse_check", "unit_check"])
        elif node.get('type') == 'premise':
            rules.append("consistency_check")
        
        return rules
    
    def execute(self, execution_plan: Dict) -> Dict[str, Any]:
        """Execute plan using LEGO-MoE"""
        results = {
            "step_results": [],
            "final_state": None
        }
        
        for step in execution_plan['steps']:
            # Execute via LEGO-MoE
            node_input = step['input']
            
            # For simple math with 2 numbers and 1 operator, calculate directly
            if node_input.get('type') == 'math_arithmetic':
                entities = node_input.get('entities', {})
                numbers = entities.get('numbers', [])
                operators = entities.get('operators', [])
                
                if len(numbers) == 2 and len(operators) == 1:
                    # Direct calculation
                    result = self._execute_math_op(operators[0], numbers)
                else:
                    # Fallback to task execution
                    task = node_input.get('task', '')
                    result = self.moe.generate(task, verbose=False)
            elif node_input.get('type') == 'math_operation':
                # Direct calculation
                op = node_input['operation']
                operands = node_input['operands']
                result = self._execute_math_op(op, operands)
            else:
                # Route to LEGO-MoE
                task = node_input.get('task') or node_input.get('statement')
                result = self.moe.generate(task, verbose=False)
            
            results['step_results'].append({
                "step_id": step['node_id'],
                "result": result,
                "verified": False  # Will verify next
            })
        
        return results
    
    def _execute_math_op(self, operator: str, operands: List[float]) -> float:
        """Execute math operation"""
        if operator == '+':
            return operands[0] + operands[1]
        elif operator == '-':
            return operands[0] - operands[1]
        elif operator == '*' or operator == '×':
            return operands[0] * operands[1]
        elif operator == '/':
            return operands[0] / operands[1] if operands[1] != 0 else None
        return None
    
    def verify(self, results: Dict, task_graph: Dict) -> Dict[str, Any]:
        """
        Verify results with constraint checks
        NOT multiple re-solves
        """
        verification = {
            "passed": True,
            "checks": [],
            "confidence": 1.0
        }
        
        for step_result in results['step_results']:
            result_value = step_result['result']
            
            # Bounds check
            if isinstance(result_value, (int, float)):
                bounds_ok = -1e10 < result_value < 1e10
                verification['checks'].append({
                    "type": "bounds_check",
                    "passed": bounds_ok
                })
                if not bounds_ok:
                    verification['passed'] = False
                    verification['confidence'] *= 0.5
        
        return verification
    
    def synthesize(self, verified_results: Dict, verification: Dict) -> Dict[str, Any]:
        """
        Synthesize final answer
        NO reasoning dump
        """
        if not verified_results['step_results']:
            return {
                "answer": "Unable to process query",
                "confidence": 0.0
            }
        
        # Get final result
        final_result = verified_results['step_results'][-1]['result']
        
        return {
            "answer": str(final_result),
            "confidence": verification['confidence'],
            "verified": verification['passed']
        }
    
    def self_correct(self, query: str, result: Dict, max_iterations: int = 2) -> Dict:
        """
        Self-correction with hard cap
        Downgrades confidence if uncertain
        """
        for i in range(max_iterations):
            if result['verified'] and result['confidence'] > 0.8:
                return result
            
            # Re-execute with alternative approach
            # (simplified for now)
            result['confidence'] *= 0.9
        
        # If still uncertain, downgrade
        if result['confidence'] < 0.5:
            result['note'] = "Uncertain, partial verification"
        
        return result
    
    def reason(self, query: str, verbose: bool = True) -> Dict[str, Any]:
        """
        Main reasoning pipeline
        Returns structured result, NOT text trace
        """
        if verbose:
            print(f"\n{'='*70}")
            print(f"🧠 GROK LOGIC LAYER (ISL-Safe)")
            print(f"{'='*70}")
            print(f"Query: {query}\n")
        
        # 1. Understand (structured intent)
        intent = self.parse_intent(query)
        if verbose:
            print(f"Intent: {intent['type']} (complexity: {intent['complexity']})")
        
        # 2. Decompose (task graph)
        task_graph = self.decompose_to_graph(query, intent)
        if verbose:
            print(f"Task graph: {len(task_graph['nodes'])} nodes")
        
        # 3. Plan (execution DAG)
        execution_plan = self.plan_execution(task_graph)
        if verbose:
            print(f"Execution plan: {len(execution_plan['steps'])} steps")
        
        # 4. Execute
        results = self.execute(execution_plan)
        if verbose:
            print(f"Executed: {len(results['step_results'])} steps")
        
        # 5. Verify
        verification = self.verify(results, task_graph)
        if verbose:
            print(f"Verification: {'PASSED' if verification['passed'] else 'FAILED'}")
        
        # 6. Synthesize
        final_answer = self.synthesize(results, verification)
        
        # 7. Self-correct if needed
        final_answer = self.self_correct(query, final_answer)
        
        if verbose:
            print(f"\n{'='*70}")
            print(f"Answer: {final_answer['answer']}")
            print(f"Confidence: {final_answer['confidence']:.2f}")
            print(f"{'='*70}\n")
        
        # Store artifacts (machine-readable)
        self.reasoning_artifacts = {
            "intent": intent,
            "task_graph": task_graph,
            "execution_plan": execution_plan,
            "results": results,
            "verification": verification
        }
        
        return final_answer


def main():
    """Test ISL-safe Grok Logic Layer"""
    print("\n🧠 Grok Logic Layer - ISL-Safe Implementation\n")
    
    # Initialize
    print("Loading LEGO-MoE...")
    moe = LegoMoE()
    
    print("\nInitializing Grok Logic Layer...")
    grok = GrokLogicLayer(moe)
    
    # Test queries
    test_queries = [
        "What is 15 times 23?",
        "What is 100 divided by 4?",
    ]
    
    for query in test_queries:
        result = grok.reason(query)
        
        print(f"\nReasoning Artifacts (structured):")
        print(json.dumps(grok.reasoning_artifacts['intent'], indent=2))


if __name__ == "__main__":
    main()