#!/usr/bin/env python
"""
Main execution script to demonstrate the full, end-to-end DANCE-ST pipeline.

This script sets up a complete, simulated environment and runs the entire three-phase
workflow using the updated agent architecture and the Constraint-Potential Diffusion
methodology.

It demonstrates:
1.  Setup of a mock knowledge graph and system state with a constraint-critical "hot spot".
2.  Initialization of all agents and estimators.
3.  Execution of the full agent-based workflow via the AgentCoordinator.
4.  Real-time logging output showing each phase in action.
"""
import os
import logging
import json
import networkx as nx
import numpy as np
import tensorflow as tf
import joblib
from sklearn.preprocessing import StandardScaler

# --- Import all necessary DANCE-ST components from the Core module ---
# This assumes the script is run from the project root where 'Core' is a subdirectory.
try:
    from Core.agents import AgentCoordinator, KnowledgeGraphManagementAgent, DataAgent, FusionAndProjectionAgent
    from Core.real_models import NeuralEstimator, SymbolicEstimator
except ImportError:
    print("FATAL: Could not import DANCE-ST components from the 'Core' module.")
    print("Please ensure this script is in the project's root directory and the 'Core' folder is present.")
    exit(1)

# --- Configure Logging ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("DANCEST.MainRunner")

# --- Environment Setup Functions ---

def setup_mock_estimators():
    """
    Creates temporary mock model and scaler files for the NeuralEstimator.
    In a real scenario, these paths would point to pre-trained artifacts.
    """
    logger.info("Setting up mock Keras model and scaler for demonstration...")
    # 1. Create a simple Keras model
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=(21,)), # Matches C-MAPSS feature count for realism
        tf.keras.layers.Dense(32, activation='relu'),
        tf.keras.layers.Dropout(0.5),
        tf.keras.layers.Dense(1)
    ])
    model_path = "mock_model.keras"
    model.save(model_path)

    # 2. Create a dummy scaler
    scaler = StandardScaler()
    scaler.fit(np.random.rand(100, 21)) # Fit with random data
    scaler_path = "mock_scaler.joblib"
    joblib.dump(scaler, scaler_path)
    
    # 3. Define symbolic model parameters
    symbolic_params = {
        "base_degradation_rate": 0.015,
        "base_rate_uncertainty": 0.003,
        "activation_energy_eV": 0.6
    }

    logger.info("Mock estimators created successfully.")
    return model_path, scaler_path, symbolic_params

def create_knowledge_graph(num_nodes=400):
    """Creates a mock knowledge graph, e.g., a grid for a turbine surface."""
    logger.info(f"Creating a mock knowledge graph with {num_nodes} nodes...")
    G = nx.grid_2d_graph(int(np.sqrt(num_nodes)), int(np.sqrt(num_nodes)))
    # Relabel nodes to be simple integers for easier state mapping
    node_mapping = {node: i for i, node in enumerate(G.nodes())}
    G = nx.relabel_nodes(G, node_mapping)
    return G

def simulate_system_state(graph, constraints):
    """
    Creates a dictionary representing the current state of all components in the graph.
    Includes a "hot spot" to trigger Phase 1.
    """
    logger.info("Simulating system state with a constraint-critical 'hot spot'...")
    num_nodes = len(graph.nodes)
    temp_limit = constraints.get("temperature_limit", {}).get("value", 1200.0)
    
    # Most nodes are at a stable operating temperature
    stable_temp = temp_limit * 0.8
    system_state = {node: stable_temp for node in graph.nodes()}

    # Create a "hot spot" where one node is very close to the limit
    hot_spot_node = num_nodes // 2
    hot_spot_temp = temp_limit * 0.99
    system_state[hot_spot_node] = hot_spot_temp
    
    logger.info(f"Node {hot_spot_node} is set to a critical temperature of {hot_spot_temp:.2f}°C (Limit: {temp_limit}°C).")
    return system_state

def load_constraints(filepath="constraints.json"):
    """Loads the physical constraints from a JSON file."""
    logger.info(f"Loading physical constraints from {filepath}...")
    try:
        with open(filepath, 'r') as f:
            return json.load(f)
    except FileNotFoundError:
        logger.error(f"FATAL: Constraints file not found at '{filepath}'.")
        exit(1)

def cleanup_mock_files(model_path, scaler_path):
    """Removes the temporary files created for the demonstration."""
    logger.info("Cleaning up temporary mock files...")
    if os.path.exists(model_path):
        os.remove(model_path)
    if os.path.exists(scaler_path):
        os.remove(scaler_path)
    logger.info("Cleanup complete.")

# --- Main Execution ---

def main():
    """
    Main function to initialize and run the full DANCE-ST agent-based system.
    """
    # 1. Setup Environment
    model_path, scaler_path, symbolic_params = setup_mock_estimators()
    constraints = load_constraints()
    knowledge_graph = create_knowledge_graph()

    # 2. Instantiate Estimators
    try:
        neural_estimator = NeuralEstimator(model_path, scaler_path)
        symbolic_estimator = SymbolicEstimator(symbolic_params)
    except Exception as e:
        logger.error(f"Failed to initialize estimators: {e}")
        cleanup_mock_files(model_path, scaler_path)
        return

    # 3. Instantiate Agent System
    coordinator = AgentCoordinator()
    
    # Create and register all agents
    kgma = KnowledgeGraphManagementAgent("KGMA", coordinator, knowledge_graph)
    data_agent = DataAgent("DataAgent", coordinator, neural_estimator, symbolic_estimator)
    fusion_projection_agent = FusionAndProjectionAgent("FusionAndProjectionAgent", coordinator, constraints)
    
    coordinator.register_agent(kgma)
    coordinator.register_agent(data_agent)
    coordinator.register_agent(fusion_projection_agent)

    # 4. Simulate a System State and Start the Workflow
    system_state = simulate_system_state(knowledge_graph, constraints)
    
    # This single call kicks off the entire three-phase pipeline via the agent coordinator.
    # The 'k' parameter tells Phase 1 how many critical nodes to select.
    coordinator.start_workflow(system_state, constraints, k=50)

    # 5. Cleanup
    cleanup_mock_files(model_path, scaler_path)

if __name__ == "__main__":
    main()

