# working_main_system.py
import sys
import os
import pickle
import json
import re
import numpy as np
from datetime import datetime
from typing import Dict, List, Set, Tuple
from collections import Counter, defaultdict

sys.path.append('src')

from core_system import AgenticAISystem
from hybrid_neural_system import HybridNeuralSymbolicSystem
from pure_medical_qlearning import PureMedicalQLearning
from rl_feedback import RLRewardSystem, QFeedbackLoop
from dynamic_agents import AgentCreator
from rare_handler import RareHandler
from pure_data_interactive_session import PureDataInteractiveSession

from agentic_layer import PatientDocumentProcessor, AgentOutputProcessor

from enhanced_mimic_learning import enhance_interactive_session_learning, enhance_interactive_session_with_rag_llm
from rag_llm_integration import integrate_rag_llm_with_system

from model_persistence import ModelPersistenceManager

class CompleteAgenticWorkflowSystem:
    
    def __init__(self):
        self.hybrid_system = None
        self.qlearning_system = None
        self.agent_creator = None
        self.document_processor = PatientDocumentProcessor()
        self.agentic_effectiveness_cache = {}
        
        print("Complete Agentic Workflow System initialized")
    
    def process_patient_documents_to_case(self, document_paths: List[str]) -> Dict:
        
        print(f"PROCESSING {len(document_paths)} PATIENT DOCUMENTS...")
        
        unified_case = self.document_processor.process_multiple_documents(document_paths)
        
        print(f"UNIFIED PATIENT CASE CREATED:")
        print(f"   Patient ID: {unified_case['patient_id']}")
        print(f"   Primary symptoms: {unified_case['symptoms']}")
        print(f"   Diagnoses found: {len(unified_case['diagnoses'])}")
        print(f"   Procedures documented: {len(unified_case['procedures'])}")
        print(f"   Medications listed: {len(unified_case['medications'])}")
        print(f"   Clinical notes: {len(unified_case['clinical_notes'])}")
        print(f"   Document sources: {unified_case['document_sources']}")
        
        return unified_case
    
    def create_enhanced_agents_with_agentic_layer(self, patient_case: Dict, 
                                                neural_system: HybridNeuralSymbolicSystem) -> List[Dict]:
        
        print(f"CREATING ENHANCED AGENTS WITH AGENTIC LAYER...")
        
        agentic_analysis = neural_system.get_agentic_analysis(patient_case)
        
        if not agentic_analysis.get('available'):
            print("Agentic layer not available, using standard agents")
            return []
        
        enhanced_agents = []
        agent_representations = agentic_analysis.get('agent_representations', [])
        
        for batch_agents in agent_representations:
            for agent_rep in batch_agents:
                
                enhanced_agent = {
                    'agent_id': f"agentic_agent_{agent_rep.agent_id}",
                    'Fi': agent_rep.focus_concepts,
                    'Pi': agent_rep.discovery_patterns,
                    'effi': agent_rep.effectiveness,
                    'Li': agent_rep.concept_connections,
                    'Hi': agent_rep.success_history,
                    'Ni': agent_rep.neural_features,
                    'original_representation': agent_rep,
                    'agentic_layer_created': True
                }
                
                enhanced_agents.append(enhanced_agent)
        
        print(f"Created {len(enhanced_agents)} enhanced agents from agentic layer")
        return enhanced_agents
    
    def fill_qlearning_patterns_in_agents(self, enhanced_agents: List[Dict], 
                                        qlearning_system: PureMedicalQLearning) -> List[Dict]:
        
        print("Filling Q-learning patterns in enhanced agents...")
        
        for agent in enhanced_agents:
            focus_concepts = list(agent['Fi'])
            
            qlearning_patterns = {}
            for concept in focus_concepts:
                if hasattr(qlearning_system, 'get_concept_relationships'):
                    relationships = qlearning_system.get_concept_relationships(concept, top_k=3)
                    qlearning_patterns[concept] = {
                        'relationships': relationships,
                        'discovery_strategy': f"explore_{concept}_connections"
                    }
                else:
                    qlearning_patterns[concept] = {
                        'relationships': [],
                        'discovery_strategy': f"investigate_{concept}"
                    }
            
            agent['Pi'] = qlearning_patterns
        
        print(f"Filled Q-learning patterns for {len(enhanced_agents)} agents")
        return enhanced_agents
    
    def calculate_enhanced_effectiveness(self, agent: Dict, 
                                       qlearning_system: PureMedicalQLearning,
                                       mimic_frequencies: Dict) -> float:
        
        focus_concepts = list(agent['Fi'])
        if not focus_concepts:
            return 0.5
        
        total_effectiveness = 0
        for concept in focus_concepts:
            q_value = 0
            if hasattr(qlearning_system, 'get_concept_relationships'):
                relationships = qlearning_system.get_concept_relationships(concept)
                if relationships:
                    q_value = max([rel[1] for rel in relationships], default=0)
            
            q_component = min(max(q_value, 0), 100) / 100
            
            freq_value = mimic_frequencies.get(concept.lower(), 0)
            freq_component = min(freq_value, 50) / 50
            
            neural_conf = agent.get('effi', 0.5)
            
            concept_effectiveness = (q_component + freq_component + neural_conf) / 3
            total_effectiveness += concept_effectiveness
        
        enhanced_effectiveness = total_effectiveness / len(focus_concepts)
        return min(max(enhanced_effectiveness, 0), 1)
    
    def generate_agentic_questions(self, enhanced_agents: List[Dict], 
                                 patient_case: Dict) -> List[Dict]:
        
        print("GENERATING QUESTIONS WITH ENHANCED AGENTIC AGENTS...")
        
        all_questions = []
        
        for agent in enhanced_agents:
            focus_concepts = list(agent['Fi'])[:3]
            effectiveness = agent['effi']
            agent_id = agent['agent_id']
            
            diagnosis = patient_case.get('admission_info', {}).get('diagnosis', 'condition')
            
            for i, concept in enumerate(focus_concepts):
                if effectiveness > 0.7:
                    priority = "HIGH PRIORITY"
                elif effectiveness > 0.4:
                    priority = "MEDIUM PRIORITY" 
                else:
                    priority = "EXPLORATORY"
                
                if 'family' in concept.lower() or 'history' in concept.lower():
                    question = f"[{priority}] Family history or genetic risk factors relevant to {diagnosis}?"
                elif 'pain' in concept.lower():
                    question = f"[{priority}] Detailed pain assessment: onset, character, severity, radiation for {diagnosis}?"
                elif 'medication' in concept.lower():
                    question = f"[{priority}] Current medications, dosages, and compliance for {diagnosis} management?"
                elif 'test' in concept.lower() or 'lab' in concept.lower():
                    question = f"[{priority}] Recent laboratory results and diagnostic studies for {diagnosis}?"
                else:
                    question = f"[{priority}] Clinical significance of {concept} in context of {diagnosis}?"
                
                question_info = {
                    'agent_id': agent_id,
                    'question': question,
                    'focus_concept': concept,
                    'effectiveness': effectiveness,
                    'priority_rank': i + 1,
                    'agentic_generated': True
                }
                all_questions.append(question_info)
        
        print(f"Generated {len(all_questions)} agentic questions from {len(enhanced_agents)} agents")
        return all_questions

def load_preprocessed_mimic_data(pkl_path="preprocessed_mimic_data.pkl"):
    
    print(f"Loading preprocessed MIMIC data from {pkl_path}...")
    
    if not os.path.exists(pkl_path):
        print(f"Preprocessed data file not found: {pkl_path}")
        print("Please run mimic_processor.py first to create the preprocessed data!")
        return None
    
    try:
        with open(pkl_path, 'rb') as f:
            data_package = pickle.load(f)
        
        patient_cases = data_package['patient_cases']
        experts = data_package['experts']
        preprocessing_stats = data_package['preprocessing_stats']
        
        print(f"Successfully loaded preprocessed MIMIC data!")
        print(f"   Processed on: {data_package['timestamp']}")
        print(f"   Patients: {len(patient_cases):,}")
        print(f"   Experts: {len(experts)}")
        
        return patient_cases, experts, data_package
        
    except Exception as e:
        print(f"Failed to load preprocessed data: {e}")
        return None

def run_complete_agentic_training_system(pkl_path="preprocessed_mimic_data.pkl", 
                                       patient_documents: List[str] = None):
    
    print("COMPLETE AGENTIC TRAINING SYSTEM - AAAI PAPER IMPLEMENTATION")
    print("="*80)
    
    agentic_system = CompleteAgenticWorkflowSystem()
    
    if patient_documents:
        print("1. PROCESSING PATIENT DOCUMENTS...")
        patient_case_from_docs = agentic_system.process_patient_documents_to_case(patient_documents)
        print(f"Patient case created from {len(patient_documents)} documents")
    else:
        patient_case_from_docs = None
    
    print("\n2. LOADING PREPROCESSED MIMIC DATA...")
    load_result = load_preprocessed_mimic_data(pkl_path)
    if not load_result:
        print("Failed to load MIMIC data")
        return None
    
    patient_cases, mimic_experts, data_package = load_result
    print(f"Loaded {len(patient_cases):,} MIMIC cases")
    
    if patient_case_from_docs:
        patient_cases.append(patient_case_from_docs)
        print(f"Added document-based case to training data")
    
    print("\n3. INITIALIZING COMPLETE SYSTEM WITH AGENTIC LAYER...")
    
    hybrid_system = HybridNeuralSymbolicSystem()
    qlearning_system = PureMedicalQLearning()
    rl_system = RLRewardSystem()
    feedback_loop = QFeedbackLoop(qlearning_system, hybrid_system)
    agent_creator = AgentCreator(qlearning_system, patient_cases)
    rare_handler = RareHandler(patient_cases)
    
    agentic_system.hybrid_system = hybrid_system
    agentic_system.qlearning_system = qlearning_system
    agentic_system.agent_creator = agent_creator
    
    print("All systems initialized with agentic layer integration")
    
    print("\n4. TRAINING PHASE WITH AGENTIC LAYER...")
    
    print("   Training Q-learning from MIMIC data...")
    learned_concepts = qlearning_system.learn_from_mimic_data_only(patient_cases)
    q_scores = qlearning_system.train_qlearning(1200)
    print(f"   Q-learning: {len(learned_concepts)} concepts learned")
    
    mimic_frequencies = {}
    all_concepts = []
    for case in patient_cases:
        all_concepts.extend(case.get('diagnoses', []))
        all_concepts.extend(case.get('procedures', []))
        all_concepts.extend(case.get('medications', []))
    
    concept_counts = Counter(all_concepts)
    for concept, count in concept_counts.items():
        mimic_frequencies[str(concept).lower()] = count
    
    print(f"   Built MIMIC frequencies for {len(mimic_frequencies)} concepts")
    
    print("   Learning enhanced patterns...")
    missing_patterns = agent_creator.learn_missing_patterns()
    concept_clusters = agent_creator.discover_clusters()
    rare_relationships = rare_handler.find_rare_critical_relationships()
    
    print("\n5. NEURAL TRAINING WITH AGENTIC LAYER...")
    
    print("   Building neural vocabulary (Target: 4,005 concepts)...")
    vocab_size = hybrid_system.build_concept_vocabulary(patient_cases)
    print(f"   Neural vocabulary: {vocab_size:,} concepts")
    
    print("   Creating neural training examples...")
    training_examples = hybrid_system.create_training_examples(patient_cases)
    print(f"   Created {len(training_examples)} training examples")
    
    print("   Training neural component with agentic layer...")
    print("       (25 epochs with agentic layer L4 - AAAI paper specification)")
    neural_training_history = hybrid_system.train_neural_component(epochs=25)
    print(f"   Neural + agentic training complete! {len(neural_training_history)} epochs")
    
    print("\n6. TESTING COMPLETE AGENTIC REASONING...")
    
    test_patient = patient_case_from_docs if patient_case_from_docs else patient_cases[0]
    print(f"   Testing with patient: {test_patient['patient_id']}")
    
    test_experts = mimic_experts[:2]
    complete_result = hybrid_system.neural_enhanced_reasoning(test_patient, test_experts)
    
    print(f"   Complete reasoning: Neural + Symbolic + Agentic fusion")
    print(f"   Agentic layer active: {complete_result.get('agentic_layer_active', False)}")
    
    print("\n7. CREATING ENHANCED AGENTS WITH AGENTIC LAYER...")
    
    enhanced_agents = agentic_system.create_enhanced_agents_with_agentic_layer(
        test_patient, hybrid_system
    )
    
    if enhanced_agents:
        enhanced_agents = agentic_system.fill_qlearning_patterns_in_agents(
            enhanced_agents, qlearning_system
        )
        
        for agent in enhanced_agents:
            enhanced_eff = agentic_system.calculate_enhanced_effectiveness(
                agent, qlearning_system, mimic_frequencies
            )
            agent['effi_calculated'] = enhanced_eff
        
        print(f"Enhanced agents with Q-learning patterns and effectiveness:")
        for i, agent in enumerate(enhanced_agents[:3]):
            print(f"   Agent {i+1}: {agent['agent_id']}")
            print(f"     Focus concepts (Fi): {list(agent['Fi'])[:3]}")
            print(f"     Effectiveness (effi): {agent['effi']:.3f}")
            print(f"     Calculated effectiveness: {agent.get('effi_calculated', 0):.3f}")
            print(f"     Q-learning patterns (Pi): {len(agent['Pi'])} concepts")
    
    print("\n8. GENERATING AGENTIC QUESTIONS...")
    
    if enhanced_agents:
        agentic_questions = agentic_system.generate_agentic_questions(enhanced_agents, test_patient)
        
        print(f"Generated {len(agentic_questions)} agentic questions:")
        for i, q in enumerate(agentic_questions[:5], 1):
            print(f"   {i}. {q['question']}")
            print(f"      Agent: {q['agent_id']}, Effectiveness: {q['effectiveness']:.3f}")
    
    print("\n9. INTEGRATING WITH EXISTING SYSTEMS...")
    
    agent_creator = integrate_rag_llm_with_system(agent_creator, patient_cases)
    
    interactive_session = PureDataInteractiveSession(agent_creator, patient_cases)
    interactive_session = enhance_interactive_session_learning(interactive_session, patient_cases)
    
    print("Complete system integration finished")
    
    print("\nCOMPLETE AGENTIC SYSTEM RESULTS:")
    print("="*80)
    
    print(f"AAAI Paper Implementation Status:")
    print(f"   Neural Architecture: 6 layers (L1-L5 + Agentic L4)")
    print(f"   Vocabulary Size: {vocab_size:,} concepts")
    print(f"   Total Parameters: {hybrid_system.neural_model._count_parameters():,}")
    print(f"   Agentic Layer Parameters: {hybrid_system.neural_model.agentic_layer._count_parameters():,}")
    print(f"   Agent Representation: a_i = (Fi, Pi, effi, Li, Hi, Ni)")
    print(f"   Enhanced Effectiveness Formula: (Q-learning + MIMIC + Neural) / 3")
    print(f"   Dynamic Agent Creation: Variable |M| agents")
    print(f"   Document Processing: Multi-document patient input")
    
    print(f"\nEnhanced Agents Analysis:")
    print(f"   • Agents created: {len(enhanced_agents) if enhanced_agents else 0}")
    print(f"   • Agentic questions generated: {len(agentic_questions) if enhanced_agents else 0}")
    print(f"   • Q-learning patterns integrated: Complete")
    print(f"   • Neural attention weights: Complete")
    print(f"   • MIMIC frequency analysis: Complete")
    
    print(f"\nDocument Processing Capability:")
    if patient_case_from_docs:
        print(f"   Processed {len(patient_documents)} patient documents")
        print(f"   Created unified patient case: {patient_case_from_docs['patient_id']}")
        print(f"   Integrated with MIMIC training data")
    else:
        print(f"   Ready to process patient documents (provide document_paths)")
    
    return {
        'agentic_system': agentic_system,
        'patient_cases': patient_cases,
        'qlearning_system': qlearning_system,
        'hybrid_system': hybrid_system,
        'enhanced_agents': enhanced_agents if enhanced_agents else [],
        'agentic_questions': agentic_questions if enhanced_agents else [],
        'complete_result': complete_result,
        'neural_training_history': neural_training_history,
        'mimic_frequencies': mimic_frequencies,
        'rare_relationships': rare_relationships,
        'vocab_size': vocab_size,
        'patient_case_from_docs': patient_case_from_docs,
        
        'agentic_layer_integrated': True,
        'neural_symbolic_fusion': True,
        'document_processing_enabled': True,
        'aaai_paper_implementation': True,
        'training_complete': True,
        'real_mimic_data': True,
        'fast_training': True
    }

def run_with_agentic_persistence(pkl_path="preprocessed_mimic_data.pkl",
                               patient_documents: List[str] = None):
    
    print("COMPLETE AGENTIC SYSTEM WITH MODEL PERSISTENCE")
    print("="*70)
    
    persistence = ModelPersistenceManager()
    
    model_status = persistence.check_trained_models_exist()
    
    if all(model_status.values()):
        print("TRAINED MODELS FOUND - LOADING FROM DISK")
        
        loaded_system = persistence.load_complete_trained_system()
        
        if loaded_system:
            print("System loaded successfully!")
            
            if patient_documents:
                print(f"Processing new patient documents...")
                doc_processor = PatientDocumentProcessor()
                new_case = doc_processor.process_multiple_documents(patient_documents)
                loaded_system['new_patient_case'] = new_case
                print(f"New patient case processed: {new_case['patient_id']}")
            
            return loaded_system
        else:
            print("Loading failed - will retrain")
    
    print("Training complete agentic system...")
    
    training_results = run_complete_agentic_training_system(pkl_path, patient_documents)
    
    if training_results and training_results.get('training_complete'):
        print("Saving complete agentic system...")
        saved_files = persistence.save_complete_trained_system(training_results)
        
        if saved_files:
            print("Complete agentic system saved!")
            training_results['model_files_saved'] = saved_files
        else:
            print("Model saving failed")
    
    return training_results

def run_with_patient_documents(document_paths: List[str]):
    
    print(f"RUNNING COMPLETE AGENTIC SYSTEM WITH PATIENT DOCUMENTS")
    print(f"   Documents to process: {len(document_paths)}")
    
    results = run_with_agentic_persistence(
        pkl_path="preprocessed_mimic_data.pkl",
        patient_documents=document_paths
    )
    
    if results:
        print(f"\nRESULTS WITH PATIENT DOCUMENTS:")
        print(f"   System training: {'Complete' if results.get('training_complete') else 'Failed'}")
        print(f"   Agentic layer: {'Active' if results.get('agentic_layer_integrated') else 'Inactive'}")
        print(f"   Document processing: {'Enabled' if results.get('document_processing_enabled') else 'Disabled'}")
        print(f"   Enhanced agents: {len(results.get('enhanced_agents', []))}")
        print(f"   Agentic questions: {len(results.get('agentic_questions', []))}")
        
        if results.get('patient_case_from_docs'):
            case = results['patient_case_from_docs']
            print(f"\nProcessed Patient Case:")
            print(f"   Patient ID: {case['patient_id']}")
            print(f"   Symptoms: {case['symptoms']}")
            print(f"   Source documents: {len(case.get('document_sources', []))}")
    
    return results

if __name__ == "__main__":
    example_documents = [
        "patient_history.txt",
        "lab_results.csv", 
        "clinical_notes.txt"
    ]
    
    print("Starting COMPLETE AGENTIC SYSTEM...")
    
    results = run_with_agentic_persistence(
        pkl_path="preprocessed_mimic_data.pkl",
        patient_documents=None
    )
    
    if results and results.get('training_complete'):
        if results.get('loaded_from_disk'):
            print("COMPLETE AGENTIC SYSTEM LOADED FROM DISK!")
        else:
            print("COMPLETE AGENTIC SYSTEM TRAINED SUCCESSFULLY!")
        
        print(f"\nAAAI Paper Implementation Summary:")
        print(f"   Agentic Layer (L4): {'Integrated' if results.get('agentic_layer_integrated') else 'Missing'}")
        print(f"   Neural Architecture: 6 layers total (L1-L5 + Agentic L4)")
        print(f"   Vocabulary: {results.get('vocab_size', 0):,} concepts")
        print(f"   Target Parameters: 7.03M (Paper specification)")
        print(f"   Document Processing: {'Ready' if results.get('document_processing_enabled') else 'Not ready'}")
        print(f"   Agent Formula: a_i = (Fi, Pi, effi, Li, Hi, Ni)")
        print(f"   Effectiveness Formula: (Q-learning + MIMIC + Neural) / 3")
        
        persistence = ModelPersistenceManager()
        
        model_info = persistence.get_model_info()
        if model_info.get('available'):
            print(f"\nModel Info:")
            print(f"   Training timestamp: {model_info['training_timestamp']}")
            print(f"   Patient cases used: {model_info['patient_cases_used']:,}")
            print(f"   Components available: {model_info['components_available']}")
    else:
        print(" System startup failed")

