src/core/multiversal_compute_system.py

"""
Multiversal Compute System for JARVIS-2v
Main interface for parallel universes as compute nodes with cross-universe knowledge transfer
"""

import json
import logging
import math
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from .multiversal_adapters import MultiversalAdapter, MultiversalComputeEngine, MultiversalRoutingEngine
from ..quantum.multiversal_quantum import MultiversalQuantumEngine, MultiversalExperimentConfig


logger = logging.getLogger(__name__)


@dataclass
class MultiversalQuery:
    """Query for multiversal computation"""
    query_id: str
    problem_description: str
    problem_domain: str
    complexity: float
    urgency: str  # "low", "medium", "high"
    target_outcome: Optional[str] = None
    constraints: Optional[Dict[str, Any]] = None
    max_universes: int = 5
    allow_cross_universe_transfer: bool = True
    simulation_steps: int = 10
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "query_id": self.query_id,
            "problem_description": self.problem_description,
            "problem_domain": self.problem_domain,
            "complexity": self.complexity,
            "urgency": self.urgency,
            "target_outcome": self.target_outcome,
            "constraints": self.constraints or {},
            "max_universes": self.max_universes,
            "allow_cross_universe_transfer": self.allow_cross_universe_transfer,
            "simulation_steps": self.simulation_steps
        }


@dataclass
class MultiversalSolution:
    """Solution from multiversal computation"""
    solution_id: str
    query_id: str
    primary_universe: str
    contributing_universes: List[str]
    solution_quality: float
    confidence: float
    solution_data: Dict[str, Any]
    cross_universe_insights: List[Dict[str, Any]]
    processing_time: float
    artifacts_generated: List[str]
    timestamp: float
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "solution_id": self.solution_id,
            "query_id": self.query_id,
            "primary_universe": self.primary_universe,
            "contributing_universes": self.contributing_universes,
            "solution_quality": self.solution_quality,
            "confidence": self.confidence,
            "solution_data": self.solution_data,
            "cross_universe_insights": self.cross_universe_insights,
            "processing_time": self.processing_time,
            "artifacts_generated": self.artifacts_generated,
            "timestamp": self.timestamp
        }


class MultiversalComputeSystem:
    """Main system for multiversal computing with parallel universe simulation"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.storage_path = Path(config.get("multiverse", {}).get("storage_path", "./multiverse"))
        self.storage_path.mkdir(parents=True, exist_ok=True)
        
        # Initialize core engines
        self.multiverse_engine = MultiversalComputeEngine(config)
        self.routing_engine = MultiversalRoutingEngine(
            config.get("bits", {}).get("y_bits", 16),
            config.get("bits", {}).get("z_bits", 8),
            config.get("bits", {}).get("x_bits", 8),
            config.get("bits", {}).get("u_bits", 16)
        )
        
        # Initialize quantum engine for multiversal experiments
        artifacts_path = config.get("artifacts", {}).get("storage_path", "./artifacts")
        self.quantum_engine = MultiversalQuantumEngine(artifacts_path, None, self.multiverse_engine)
        
        # Session management
        self.active_queries: Dict[str, MultiversalQuery] = {}
        self.solutions: Dict[str, MultiversalSolution] = {}
        
        # Performance metrics
        self.performance_metrics = {
            "total_queries": 0,
            "successful_solutions": 0,
            "cross_universe_transfers": 0,
            "average_processing_time": 0.0,
            "multiverse_health": 0.0
        }
        
        # Load existing state
        self._load_system_state()
    
    def process_multiversal_query(self, query: MultiversalQuery) -> MultiversalSolution:
        """Process a query using multiversal computation"""
        
        start_time = time.time()
        
        print(f"🌌 Processing multiversal query: {query.query_id}")
        print(f"   Domain: {query.problem_domain}")
        print(f"   Complexity: {query.complexity}")
        print(f"   Max universes: {query.max_universes}")
        
        # Store active query
        self.active_queries[query.query_id] = query
        
        # Step 1: Find or create relevant universes
        target_universes = self._find_or_create_relevant_universes(query)
        
        # Step 2: Route to best universes using interference patterns
        routed_universes = self.routing_engine.route_to_parallel_universes(
            query.to_dict(),
            target_universes,
            target_universe=f"target_{query.query_id}"
        )
        
        # If no universes were routed (shouldn't happen), use all target universes
        if not routed_universes:
            print(f"⚠️  No universes routed, using all target universes")
            for adapter in target_universes:
                routed_universes.append((adapter, 0.5))  # Default interference weight
        
        # Step 3: Execute computation across universes
        universe_results = []
        for adapter, interference_weight in routed_universes:
            result = self._execute_universe_computation(adapter, query, interference_weight)
            universe_results.append(result)
        
        # Step 4: Cross-universe knowledge transfer (if enabled)
        cross_universe_insights = []
        if query.allow_cross_universe_transfer:
            cross_universe_insights = self._perform_cross_universe_transfer(
                universe_results, query
            )
        
        # Step 5: Synthesize final solution
        solution = self._synthesize_multiversal_solution(
            query, universe_results, cross_universe_insights, time.time() - start_time
        )
        
        # Update metrics
        self._update_performance_metrics(solution, time.time() - start_time)
        
        # Store solution
        self.solutions[solution.solution_id] = solution
        
        # Clean up active query
        self.active_queries.pop(query.query_id, None)
        
        # Save system state
        self._save_system_state()
        
        print(f"✅ Solution generated in {solution.processing_time:.2f}s")
        print(f"   Primary universe: {solution.primary_universe}")
        print(f"   Quality: {solution.solution_quality:.2f}")
        print(f"   Contributing universes: {len(solution.contributing_universes)}")
        
        return solution
    
    def _find_or_create_relevant_universes(self, query: MultiversalQuery) -> List[MultiversalAdapter]:
        """Find or create universes relevant to the query"""
        
        relevant_universes = []
        
        # Try to find existing relevant universes
        successful_universes = self.multiverse_engine.find_successful_universes(
            query.problem_domain, similarity_threshold=0.5
        )
        
        # Create adapters for successful universes
        for universe_id, reach_factor in successful_universes:
            adapter = MultiversalAdapter(
                id=f"adapter_{universe_id}",
                task_tags=[query.problem_domain],
                y_bits=[0] * 16,
                z_bits=[0] * 8,
                x_bits=[0] * 8,
                universe_id=universe_id,
                universe_bits=self.routing_engine.generate_universe_signature({
                    "type": query.problem_domain,
                    "complexity": query.complexity,
                    "domain": query.problem_domain
                }),
                interference_weight=reach_factor,
                coherence_level=self.multiverse_engine.universes.get(universe_id, {}).coherence_level or 0.8
            )
            relevant_universes.append(adapter)
        
        # Create new universes if needed
        while len(relevant_universes) < query.max_universes:
            new_universe_id = self.multiverse_engine.create_parallel_universe(
                parent_universe_id="base_multiverse",
                decision_point=f"query_{query.query_id}_branch_{len(relevant_universes)}",
                problem_context={
                    "domain": query.problem_domain,
                    "complexity": query.complexity,
                    "query_id": query.query_id
                }
            )
            
            new_adapter = MultiversalAdapter(
                id=f"adapter_{new_universe_id}",
                task_tags=[query.problem_domain],
                y_bits=[0] * 16,
                z_bits=[0] * 8,
                x_bits=[0] * 8,
                universe_id=new_universe_id,
                universe_bits=self.routing_engine.generate_universe_signature({
                    "type": query.problem_domain,
                    "complexity": query.complexity,
                    "domain": query.problem_domain
                }),
                interference_weight=0.5,  # Default weight for new universes
                coherence_level=0.8  # Default coherence for new universes
            )
            relevant_universes.append(new_adapter)
        
        return relevant_universes[:query.max_universes]
    
    def _execute_universe_computation(self, adapter: MultiversalAdapter, 
                                    query: MultiversalQuery, interference_weight: float) -> Dict[str, Any]:
        """Execute computation in a specific universe.

        For most domains, this uses the existing multiversal simulation.

        For protein_folding, this runs REAL physics-based folding (not mock).
        """

        if query.problem_domain == "protein_folding":
            return self._execute_protein_folding_universe(adapter, query, interference_weight)

        # Simulate universe evolution
        evolution_result = self.multiverse_engine.simulate_universe_evolution(
            adapter.universe_id, steps=query.simulation_steps
        )

        # Generate universe-specific solution
        solution_quality = min(1.0, (adapter.coherence_level * interference_weight *
                                   (0.5 + query.complexity * 0.5)))

        # Add some randomness
        solution_quality += (hash(adapter.universe_id + query.query_id) % 100) / 1000
        solution_quality = max(0.1, min(1.0, solution_quality))

        universe_solution = {
            "universe_id": adapter.universe_id,
            "adapter_id": adapter.id,
            "solution_quality": solution_quality,
            "interference_weight": interference_weight,
            "coherence_level": adapter.coherence_level,
            "evolution_result": evolution_result,
            "universe_insights": f"Universe {adapter.universe_id} approach to {query.problem_description}",
            "processing_metadata": {
                "query_domain": query.problem_domain,
                "complexity": query.complexity,
                "urgency": query.urgency,
                "timestamp": time.time()
            }
        }

        return universe_solution

    def _execute_protein_folding_universe(
        self,
        adapter: MultiversalAdapter,
        query: MultiversalQuery,
        interference_weight: float,
    ) -> Dict[str, Any]:
        """REAL protein folding in a single universe.

        Expected query.constraints:
          - sequence: str (required)
          - steps_per_universe: int (optional)
          - t_start: float (optional)
          - t_end: float (optional)
          - seed: int (optional)

        Returns a universe result with best_energy, final_energy, best_structure, and artifact_path.
        """
        from ..multiversal.protein_folding_engine import ProteinFoldingEngine

        constraints = query.constraints or {}
        sequence = (constraints.get("sequence") or "").strip().upper()
        if not sequence:
            raise ValueError("protein_folding requires constraints.sequence")

        valid = set("ACDEFGHIKLMNPQRSTVWY")
        if not all(aa in valid for aa in sequence):
            raise ValueError("Invalid amino acid sequence for protein_folding")

        steps = int(constraints.get("steps_per_universe", 5000))
        t_start = float(constraints.get("t_start", 2.0))
        t_end = float(constraints.get("t_end", 0.2))

        # Deterministic per-universe seed unless provided
        seed = int(constraints.get("seed") or (abs(hash(adapter.universe_id + query.query_id)) % 2_147_483_647))

        # Keep universe evolution for continuity with the multiverse engine
        evolution_result = self.multiverse_engine.simulate_universe_evolution(
            adapter.universe_id, steps=query.simulation_steps
        )

        engine = ProteinFoldingEngine(artifacts_dir="./protein_folding_artifacts")
        initial = engine.initialize_extended_chain(sequence, seed=seed)

        log_every = max(1, min(250, steps // 10))
        folding_result = engine.metropolis_anneal(
            initial,
            steps=steps,
            t_start=t_start,
            t_end=t_end,
            seed=seed,
            log_every=log_every,
        )

        best_energy = float(folding_result["best_energy"])
        final_energy = float(folding_result["final_energy"])
        acceptance_rate = float(folding_result["acceptance_rate"])

        # Map energy to quality in [0,1]; lower energy => higher quality
        # The exp transform keeps it stable for large magnitudes.
        solution_quality = 1.0 / (1.0 + math.exp(best_energy / 10.0))
        solution_quality *= (0.5 + 0.5 * interference_weight) * adapter.coherence_level
        solution_quality = max(0.0, min(1.0, solution_quality))

        artifact_path = engine.save_artifact(
            run_id=f"{query.query_id}_{adapter.universe_id}",
            payload={
                "query_id": query.query_id,
                "universe_id": adapter.universe_id,
                "sequence": sequence,
                "initial_structure": initial,
                "best_structure": folding_result["best_structure"],
                "best_energy": best_energy,
                "final_energy": final_energy,
                "acceptance_rate": acceptance_rate,
                "trajectory": folding_result["trajectory"],
                "parameters": {
                    "steps": steps,
                    "t_start": t_start,
                    "t_end": t_end,
                    "seed": seed,
                },
            },
            filename_prefix="protein_fold_universe",
        )

        logger.info(
            "protein_folding universe=%s best_energy=%.6f final_energy=%.6f acc=%.3f artifact=%s",
            adapter.universe_id,
            best_energy,
            final_energy,
            acceptance_rate,
            artifact_path,
        )

        return {
            "universe_id": adapter.universe_id,
            "adapter_id": adapter.id,
            "solution_quality": solution_quality,
            "interference_weight": interference_weight,
            "coherence_level": adapter.coherence_level,
            "evolution_result": evolution_result,
            "universe_insights": f"Protein folding trajectory with best_energy={best_energy:.6f}",
            "protein_folding": {
                "sequence": sequence,
                "seed": seed,
                "steps": steps,
                "t_start": t_start,
                "t_end": t_end,
                "best_energy": best_energy,
                "final_energy": final_energy,
                "acceptance_rate": acceptance_rate,
                "best_structure": folding_result["best_structure"].to_dict(),
                "artifact_path": artifact_path,
            },
            "processing_metadata": {
                "query_domain": query.problem_domain,
                "complexity": query.complexity,
                "urgency": query.urgency,
                "timestamp": time.time(),
            },
        }
    
    def _perform_cross_universe_transfer(self, universe_results: List[Dict[str, Any]], 
                                      query: MultiversalQuery) -> List[Dict[str, Any]]:
        """Perform cross-universe knowledge transfer"""
        
        cross_universe_insights = []
        
        # Get knowledge from successful universes
        target_problem = {
            "domain": query.problem_domain,
            "complexity": query.complexity,
            "description": query.problem_description
        }
        
        for result in universe_results:
            universe_id = result["universe_id"]
            
            # Borrow knowledge from this universe
            borrowed_knowledge = self.quantum_engine.get_cross_universe_knowledge(
                query.problem_domain, target_problem
            )
            
            if borrowed_knowledge.get("knowledge_borrowed"):
                for knowledge in borrowed_knowledge["knowledge_borrowed"]:
                    if knowledge["source_universe"] != universe_id:
                        insight = {
                            "type": "cross_universe_transfer",
                            "source_universe": knowledge["source_universe"],
                            "target_universe": universe_id,
                            "insight": f"Cross-universe insight from {knowledge['source_universe']}",
                            "relevance_score": knowledge["relevance_score"],
                            "echo_strength": knowledge["echo_artifact"]["echo_strength"]
                        }
                        cross_universe_insights.append(insight)
        
        return cross_universe_insights
    
    def _synthesize_multiversal_solution(self, query: MultiversalQuery, 
                                       universe_results: List[Dict[str, Any]],
                                       cross_universe_insights: List[Dict[str, Any]],
                                       processing_time: float) -> MultiversalSolution:
        """Synthesize final solution from multiverse computation"""
        
        # Find primary universe (highest quality)
        best_result = max(universe_results, key=lambda x: x["solution_quality"])
        
        # Calculate overall solution quality
        quality_scores = [r["solution_quality"] for r in universe_results]
        avg_quality = sum(quality_scores) / len(quality_scores)
        quality_variance = sum((q - avg_quality) ** 2 for q in quality_scores) / len(quality_scores)
        
        # Confidence based on consensus between universes
        confidence = max(0.1, min(1.0, 1.0 - quality_variance))
        
        # Boost confidence if cross-universe insights available
        if cross_universe_insights:
            confidence = min(1.0, confidence + 0.2)
        
        # Collect contributing universes
        contributing_universes = [r["universe_id"] for r in universe_results]
        
        # Generate solution data
        solution_data = {
            "primary_approach": best_result["universe_insights"],
            "multiverse_consensus": avg_quality,
            "universe_results": universe_results,
            "recommendations": self._generate_recommendations(universe_results, cross_universe_insights),
            "next_steps": self._suggest_next_steps(query, best_result)
        }
        
        # Create artifacts
        artifacts_generated = []
        if query.problem_domain == "cancer_treatment":
            # Special handling for cancer treatment queries
            cancer_artifact = self._generate_cancer_treatment_artifacts(query, universe_results)
            artifacts_generated.append(cancer_artifact)
        
        return MultiversalSolution(
            solution_id=f"solution_{uuid.uuid4().hex[:8]}",
            query_id=query.query_id,
            primary_universe=best_result["universe_id"],
            contributing_universes=contributing_universes,
            solution_quality=avg_quality,
            confidence=confidence,
            solution_data=solution_data,
            cross_universe_insights=cross_universe_insights,
            processing_time=processing_time,
            artifacts_generated=artifacts_generated,
            timestamp=time.time()
        )
    
    def _generate_recommendations(self, universe_results: List[Dict[str, Any]], 
                                cross_universe_insights: List[Dict[str, Any]]) -> List[str]:
        """Generate recommendations based on multiverse results"""
        
        recommendations = []
        
        # Analyze universe consensus
        qualities = [r["solution_quality"] for r in universe_results]
        avg_quality = sum(qualities) / len(qualities)
        
        if avg_quality > 0.8:
            recommendations.append("High confidence solution - proceed with primary approach")
        elif avg_quality > 0.6:
            recommendations.append("Moderate confidence - consider secondary approaches")
        else:
            recommendations.append("Low confidence - explore alternative strategies")
        
        # Check for cross-universe consensus
        if cross_universe_insights:
            recommendations.append("Cross-universe insights available - review for additional strategies")
        
        # Quality-based recommendations
        best_result = max(universe_results, key=lambda x: x["solution_quality"])
        worst_result = min(universe_results, key=lambda x: x["solution_quality"])
        
        if best_result["solution_quality"] - worst_result["solution_quality"] > 0.3:
            recommendations.append("High variance in universe results - focus on top-performing approaches")
        
        return recommendations
    
    def _suggest_next_steps(self, query: MultiversalQuery, best_result: Dict[str, Any]) -> List[str]:
        """Suggest next steps based on query and results"""
        
        next_steps = []
        
        if query.urgency == "high":
            next_steps.append("Immediate action recommended based on multiverse analysis")
        
        if query.complexity > 0.8:
            next_steps.append("High complexity detected - consider breaking into sub-problems")
        
        if best_result["solution_quality"] > 0.9:
            next_steps.append("Excellent solution quality - proceed with confidence")
        else:
            next_steps.append("Monitor results and prepare alternative approaches")
        
        if query.problem_domain == "cancer_treatment":
            next_steps.append("Consult with medical professionals before implementing treatment changes")
        
        return next_steps
    
    def _generate_cancer_treatment_artifacts(self, query: MultiversalQuery, 
                                           universe_results: List[Dict[str, Any]]) -> str:
        """Generate special artifacts for cancer treatment queries"""
        
        # Create multiversal cancer simulation
        config = MultiversalExperimentConfig(
            experiment_type="multiversal_cancer_simulation",  # Required parameter
            universe_count=min(len(universe_results), 10),
            parallel_simulations=5,
            cross_universe_transfer=True,
            interference_amplification=True
        )
        
        simulation_result = self.quantum_engine.run_multiversal_cancer_simulation(config)
        
        return simulation_result["experiment_id"]
    
    def get_system_status(self) -> Dict[str, Any]:
        """Get current status of the multiversal compute system"""
        
        multiverse_overview = self.multiverse_engine.get_multiverse_overview()
        
        status = {
            "system_health": self.performance_metrics["multiverse_health"],
            "multiverse_overview": multiverse_overview,
            "performance_metrics": self.performance_metrics,
            "active_queries": len(self.active_queries),
            "total_solutions": len(self.solutions),
            "recent_solutions": [
                {
                    "solution_id": sol.solution_id,
                    "query_id": sol.query_id,
                    "quality": sol.solution_quality,
                    "confidence": sol.confidence,
                    "timestamp": sol.timestamp
                }
                for sol in list(self.solutions.values())[-5:]  # Last 5 solutions
            ],
            "timestamp": time.time()
        }
        
        return status
    
    def _update_performance_metrics(self, solution: MultiversalSolution, processing_time: float):
        """Update system performance metrics"""
        
        self.performance_metrics["total_queries"] += 1
        
        if solution.solution_quality > 0.7:
            self.performance_metrics["successful_solutions"] += 1
        
        if solution.cross_universe_insights:
            self.performance_metrics["cross_universe_transfers"] += 1
        
        # Update average processing time
        total_queries = self.performance_metrics["total_queries"]
        current_avg = self.performance_metrics["average_processing_time"]
        self.performance_metrics["average_processing_time"] = (
            (current_avg * (total_queries - 1) + processing_time) / total_queries
        )
        
        # Update multiverse health
        success_rate = self.performance_metrics["successful_solutions"] / total_queries
        multiverse_health = multiverse_overview = self.multiverse_engine.get_multiverse_overview()["multiverse_health"]
        self.performance_metrics["multiverse_health"] = (success_rate + multiverse_health) / 2
    
    def _load_system_state(self):
        """Load system state from disk"""
        state_file = self.storage_path / "system_state.json"
        if state_file.exists():
            try:
                with open(state_file, 'r') as f:
                    data = json.load(f)
                    
                self.performance_metrics.update(data.get("performance_metrics", {}))
                self.solutions = {}
                for sol_data in data.get("solutions", []):
                    solution = MultiversalSolution(**sol_data)
                    self.solutions[solution.solution_id] = solution
                    
            except (json.JSONDecodeError, TypeError):
                pass
    
    def _save_system_state(self):
        """Save system state to disk"""
        state_file = self.storage_path / "system_state.json"
        
        state_data = {
            "performance_metrics": self.performance_metrics,
            "solutions": [sol.to_dict() for sol in self.solutions.values()],
            "last_updated": time.time()
        }
        
        with open(state_file, 'w') as f:
            json.dump(state_data, f, indent=2)
    
    def get_solution(self, solution_id: str) -> Optional[MultiversalSolution]:
        """Retrieve a specific solution by ID"""
        return self.solutions.get(solution_id)
    
    def list_solutions(self, limit: int = 10) -> List[MultiversalSolution]:
        """List recent solutions"""
        return list(self.solutions.values())[-limit:]


__all__ = ["MultiversalComputeSystem", "MultiversalQuery", "MultiversalSolution"]