biomind_comprehensive_diagnostic.py

"""
BIOMIND MMLU Comprehensive Diagnostic Script
===========================================

Comprehensive diagnostic tool for analyzing failed MMLU questions.
Provides detailed analysis of model responses, confidence calibration,
and cognitive control loop behavior.
"""

import sys
import os
import json
import re
from typing import Dict, List, Any, Tuple
from datetime import datetime

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from biomind_recursive_dmn import BIOMINDRecursiveDMN
from real_dataset_loader import RealDatasetLoader
from integrated_biomind_evaluation import IntegratedBIOMINDEvaluator

class BIOMINDDiagnosticAnalyzer:
    """Comprehensive diagnostic analyzer for BIOMIND MMLU failures"""

    def __init__(self):
        self.specialists = ['qwen_math_expert', 'qwen_general_reasoner', 'tiny_llama_planner', 'tiny_llama_critic']
        self.biomind_system = BIOMINDRecursiveDMN(
            self.specialists,
            max_reflection_cycles=15,
            use_real_models=True,
            use_trained_router=True
        )
        self.evaluator = IntegratedBIOMINDEvaluator(self.specialists)

    def diagnose_failed_questions(self, num_questions: int = 20, output_file: str = None) -> Dict[str, Any]:
        """Run diagnostic analysis on MMLU questions and identify failure patterns"""

        print("[SEARCH] BIOMIND MMLU Diagnostic Analysis")
        print("=" * 60)

        # Load questions
        loader = RealDatasetLoader()
        questions = loader.load_mmlu_subject('abstract_algebra', split='test', max_samples=num_questions)

        if not questions:
            print("[FAIL] Failed to load MMLU questions")
            return {}

        print(f"[OK] Loaded {len(questions)} questions for diagnostic analysis")

        # Run diagnostic evaluation
        diagnostic_results = []
        failure_patterns = {
            'parsing_errors': [],
            'model_hallucinations': [],
            'confidence_miscalibration': [],
            'routing_issues': [],
            'cognitive_loop_failures': []
        }

        for i, question_data in enumerate(questions):
            print(f"\n[LAB] Analyzing Question {i+1}/{len(questions)}")
            print("-" * 50)

            result = self._analyze_single_question(question_data, i+1)
            diagnostic_results.append(result)

            # Categorize failure patterns
            if not result['is_correct']:
                self._categorize_failure(result, failure_patterns)

        # Generate diagnostic report
        report = self._generate_diagnostic_report(diagnostic_results, failure_patterns)

        # Save to file if requested
        if output_file:
            with open(output_file, 'w') as f:
                json.dump(report, f, indent=2, default=str)
            print(f"? Diagnostic report saved to: {output_file}")

        return report

    def _analyze_single_question(self, question_data: Dict[str, Any], question_num: int) -> Dict[str, Any]:
        """Perform detailed analysis of a single question"""

        question = question_data['question']
        choices = question_data['choices']
        answer_letters = ['A', 'B', 'C', 'D', 'E', 'F']
        expected_answer_letter = answer_letters[question_data['answer']]

        print(f"? Question: {question}")
        print(f"? Choices: {choices}")
        print(f"[OK] Expected Answer: {expected_answer_letter}")

        # Run BIOMIND evaluation with detailed logging
        result = self.evaluator.evaluate_with_reflection(
            question=question,
            choices=choices,
            expected_answer=expected_answer_letter,
            subject='abstract_algebra'
        )

        predicted = result.get('predicted_answer', 'No answer')
        confidence = result.get('confidence', 0)
        cycles = result.get('reflection_cycles', 0)
        is_correct = predicted.upper() == expected_answer_letter.upper()

        print(f"? Predicted: {predicted} (confidence: {confidence:.3f})")
        print(f"[CHART] Cycles: {cycles}")
        print(f"[TARGET] Correct: {'[OK]' if is_correct else '[FAIL]'}")

        # Extract detailed cognitive loop information
        cognitive_analysis = self._analyze_cognitive_loop(result, question_data)

        # Analyze model responses
        response_analysis = self._analyze_model_responses(result, question_data)

        # Confidence calibration check
        confidence_analysis = self._analyze_confidence_calibration(confidence, is_correct, cognitive_analysis)

        return {
            'question_num': question_num,
            'question': question,
            'choices': choices,
            'expected': expected_answer_letter,
            'predicted': predicted,
            'is_correct': is_correct,
            'confidence': confidence,
            'cycles': cycles,
            'cognitive_analysis': cognitive_analysis,
            'response_analysis': response_analysis,
            'confidence_analysis': confidence_analysis,
            'raw_result': result
        }

    def _analyze_cognitive_loop(self, result: Dict[str, Any], question_data: Dict[str, Any]) -> Dict[str, Any]:
        """Analyze the cognitive control loop behavior"""

        analysis = {
            'specialist_selection': result.get('selected_specialist', 'Unknown'),
            'routing_method': result.get('routing_method', 'Unknown'),
            'cycles_used': result.get('reflection_cycles', 0),
            'termination_reason': result.get('termination_reason', 'Unknown'),
            'convergence_achieved': False,
            'policy_confidence_history': [],
            'answer_confidence_history': []
        }

        # Check for convergence
        if 'cognitive_state' in result:
            cognitive_state = result['cognitive_state']
            analysis['convergence_achieved'] = cognitive_state.get('converged', False)

            # Extract confidence histories if available
            if 'confidence_history' in cognitive_state:
                history = cognitive_state['confidence_history']
                analysis['policy_confidence_history'] = [h.get('policy', 0) for h in history]
                analysis['answer_confidence_history'] = [h.get('answer', 0) for h in history]

        return analysis

    def _analyze_model_responses(self, result: Dict[str, Any], question_data: Dict[str, Any]) -> Dict[str, Any]:
        """Analyze the model responses for correctness and patterns"""

        analysis = {
            'response_quality': 'unknown',
            'parsing_success': False,
            'answer_extraction_method': 'unknown',
            'response_patterns': [],
            'mathematical_correctness': False,
            'logical_consistency': False
        }

        # Extract execution history if available
        if 'execution_history' in result:
            history = result['execution_history']
            responses = []

            for cycle_data in history:
                if 'specialist_response' in cycle_data:
                    response = cycle_data['specialist_response']
                    responses.append(response)

                    # Analyze response content
                    response_analysis = self._analyze_response_content(response, question_data)
                    analysis['response_patterns'].extend(response_analysis['patterns'])

            # Overall response quality assessment
            if responses:
                analysis['response_quality'] = self._assess_response_quality(responses, question_data)

        # Check parsing success
        predicted = result.get('predicted_answer', '')
        expected = ['A', 'B', 'C', 'D', 'E', 'F'][question_data['answer']]

        if predicted.upper() == expected.upper():
            analysis['parsing_success'] = True
            analysis['answer_extraction_method'] = 'correct'
        else:
            analysis['parsing_success'] = False
            analysis['answer_extraction_method'] = self._identify_parsing_method(result)

        return analysis

    def _analyze_response_content(self, response: str, question_data: Dict[str, Any]) -> Dict[str, Any]:
        """Analyze the content of a single model response"""

        analysis = {'patterns': []}

        response_lower = response.lower()

        # Check for mathematical reasoning patterns
        if any(word in response_lower for word in ['calculate', 'compute', 'solve', 'therefore', 'thus']):
            analysis['patterns'].append('mathematical_reasoning')

        # Check for boxed answers
        if '\\boxed{' in response or '$\\boxed{' in response:
            analysis['patterns'].append('latex_boxed_answer')

        # Check for explicit answer statements
        if 'final answer' in response_lower:
            analysis['patterns'].append('explicit_answer_statement')

        # Check for choice letters
        choice_letters = ['A', 'B', 'C', 'D', 'E', 'F']
        found_letters = [letter for letter in choice_letters if letter in response]
        if found_letters:
            analysis['patterns'].append(f'choice_letters_found: {found_letters}')

        # Check for numerical answers
        import re
        numbers = re.findall(r'\d+', response)
        if numbers:
            analysis['patterns'].append(f'numerical_answers: {numbers}')

        return analysis

    def _assess_response_quality(self, responses: List[str], question_data: Dict[str, Any]) -> str:
        """Assess the overall quality of model responses"""

        if not responses:
            return 'no_responses'

        # Check if any response contains the correct mathematical reasoning
        expected_answer = ['A', 'B', 'C', 'D', 'E', 'F'][question_data['answer']]
        correct_choice = question_data['choices'][question_data['answer']]

        quality_score = 0

        for response in responses:
            response_lower = response.lower()

            # Check if response mentions key mathematical concepts
            if any(concept in response_lower for concept in ['order', 'element', 'group', 'cycle', 'lcm']):
                quality_score += 1

            # Check if response shows calculation steps
            if '=' in response or 'lcm' in response_lower:
                quality_score += 1

            # Check if response contains the correct numerical answer
            if correct_choice in response:
                quality_score += 2

        if quality_score >= 3:
            return 'high_quality'
        elif quality_score >= 1:
            return 'medium_quality'
        else:
            return 'low_quality'

    def _identify_parsing_method(self, result: Dict[str, Any]) -> str:
        """Identify how the answer was extracted"""

        # This would need to be enhanced based on the actual parsing logic
        # For now, return a placeholder
        return 'regex_pattern_matching'

    def _analyze_confidence_calibration(self, confidence: float, is_correct: bool,
                                      cognitive_analysis: Dict[str, Any]) -> Dict[str, Any]:
        """Analyze confidence calibration"""

        analysis = {
            'confidence_level': self._categorize_confidence(confidence),
            'calibration_quality': 'unknown',
            'should_be_lower': False,
            'should_be_higher': False
        }

        # Check calibration
        if is_correct and confidence < 0.5:
            analysis['calibration_quality'] = 'underconfident'
            analysis['should_be_higher'] = True
        elif not is_correct and confidence > 0.8:
            analysis['calibration_quality'] = 'overconfident'
            analysis['should_be_lower'] = True
        elif (is_correct and confidence >= 0.7) or (not is_correct and confidence <= 0.4):
            analysis['calibration_quality'] = 'well_calibrated'
        else:
            analysis['calibration_quality'] = 'moderately_calibrated'

        return analysis

    def _categorize_confidence(self, confidence: float) -> str:
        """Categorize confidence level"""

        if confidence >= 0.9:
            return 'very_high'
        elif confidence >= 0.7:
            return 'high'
        elif confidence >= 0.5:
            return 'medium'
        elif confidence >= 0.3:
            return 'low'
        else:
            return 'very_low'

    def _categorize_failure(self, result: Dict[str, Any], failure_patterns: Dict[str, List]) -> None:
        """Categorize the type of failure"""

        confidence = result['confidence']
        cognitive = result['cognitive_analysis']
        response = result['response_analysis']

        # Confidence miscalibration
        if result['confidence_analysis']['should_be_lower']:
            failure_patterns['confidence_miscalibration'].append({
                'question_num': result['question_num'],
                'confidence': confidence,
                'reason': 'overconfident_wrong_answer'
            })

        # Parsing errors
        if not response['parsing_success'] and response['response_quality'] == 'high_quality':
            failure_patterns['parsing_errors'].append({
                'question_num': result['question_num'],
                'reason': 'correct_response_wrong_parsing'
            })

        # Model hallucinations
        if response['response_quality'] == 'low_quality':
            failure_patterns['model_hallucinations'].append({
                'question_num': result['question_num'],
                'reason': 'poor_mathematical_reasoning'
            })

        # Cognitive loop failures
        if cognitive['cycles_used'] >= 10 and not cognitive['convergence_achieved']:
            failure_patterns['cognitive_loop_failures'].append({
                'question_num': result['question_num'],
                'cycles': cognitive['cycles_used'],
                'reason': 'no_convergence'
            })

    def _generate_diagnostic_report(self, results: List[Dict], failure_patterns: Dict) -> Dict[str, Any]:
        """Generate comprehensive diagnostic report"""

        total_questions = len(results)
        correct_answers = sum(1 for r in results if r['is_correct'])
        accuracy = correct_answers / total_questions * 100

        # Confidence analysis
        confidences = [r['confidence'] for r in results]
        avg_confidence = sum(confidences) / len(confidences)

        correct_confidences = [r['confidence'] for r in results if r['is_correct']]
        incorrect_confidences = [r['confidence'] for r in results if not r['is_correct']]

        avg_correct_confidence = sum(correct_confidences) / len(correct_confidences) if correct_confidences else 0
        avg_incorrect_confidence = sum(incorrect_confidences) / len(incorrect_confidences) if incorrect_confidences else 0

        # Cognitive loop analysis
        cycles_used = [r['cycles'] for r in results]
        avg_cycles = sum(cycles_used) / len(cycles_used)

        # Response quality analysis
        quality_counts = {}
        for result in results:
            quality = result['response_analysis']['response_quality']
            quality_counts[quality] = quality_counts.get(quality, 0) + 1

        report = {
            'timestamp': datetime.now().isoformat(),
            'summary': {
                'total_questions': total_questions,
                'accuracy': accuracy,
                'average_confidence': avg_confidence,
                'average_cycles': avg_cycles,
                'correct_confidence_avg': avg_correct_confidence,
                'incorrect_confidence_avg': avg_incorrect_confidence
            },
            'failure_analysis': failure_patterns,
            'response_quality_distribution': quality_counts,
            'detailed_results': results,
            'recommendations': self._generate_recommendations(results, failure_patterns)
        }

        return report

    def _generate_recommendations(self, results: List[Dict], failure_patterns: Dict) -> List[str]:
        """Generate improvement recommendations based on diagnostic findings"""

        recommendations = []

        # Confidence calibration issues
        miscalibration_count = len(failure_patterns['confidence_miscalibration'])
        if miscalibration_count > 0:
            recommendations.append(f"Address confidence miscalibration ({miscalibration_count} cases): Implement better confidence estimation in the cognitive control loop")

        # Parsing issues
        parsing_errors = len(failure_patterns['parsing_errors'])
        if parsing_errors > 0:
            recommendations.append(f"Fix answer parsing ({parsing_errors} cases): Improve regex patterns and DistilBERT answer extraction for mathematical responses")

        # Model quality issues
        hallucinations = len(failure_patterns['model_hallucinations'])
        if hallucinations > 0:
            recommendations.append(f"Improve model responses ({hallucinations} cases): Fine-tune math models on abstract algebra or use more capable models")

        # Cognitive loop issues
        loop_failures = len(failure_patterns['cognitive_loop_failures'])
        if loop_failures > 0:
            recommendations.append(f"Optimize cognitive control ({loop_failures} cases): Adjust convergence criteria and specialist selection logic")

        # General recommendations
        accuracy = sum(1 for r in results if r['is_correct']) / len(results) * 100
        if accuracy < 50:
            recommendations.append("Overall accuracy is low: Consider using more capable base models or implementing ensemble methods")

        avg_confidence = sum(r['confidence'] for r in results) / len(results)
        if avg_confidence < 0.6:
            recommendations.append("Average confidence is low: Investigate if models are being too conservative or if confidence estimation needs calibration")

        return recommendations

def run_diagnostic_analysis():
    """Run the complete diagnostic analysis"""

    analyzer = BIOMINDDiagnosticAnalyzer()

    # Generate timestamped output file
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = f"biomind_mmlu_comprehensive_diagnostic_{timestamp}.json"

    report = analyzer.diagnose_failed_questions(num_questions=20, output_file=output_file)

    # Print summary to console
    print("\n" + "=" * 80)
    print("[TARGET] DIAGNOSTIC SUMMARY")
    print("=" * 80)

    summary = report['summary']
    print(f"[CHART] Accuracy: {summary['accuracy']:.1f}% ({summary['correct_confidence_avg']:.3f} avg conf when correct)")
    print(f"[TARGET] Confidence: {summary['average_confidence']:.3f} overall ({summary['incorrect_confidence_avg']:.3f} when wrong)")
    print(f"[CYCLE] Cycles: {summary['average_cycles']:.1f} average")

    print(f"\n[STATS] Response Quality Distribution:")
    for quality, count in report['response_quality_distribution'].items():
        print(f"   {quality}: {count}")

    print(f"\n? Failure Patterns:")
    for pattern_type, cases in report['failure_analysis'].items():
        if cases:
            print(f"   {pattern_type}: {len(cases)} cases")

    print(f"\n[IDEA] Recommendations:")
    for rec in report['recommendations']:
        print(f"   ? {rec}")

    print(f"\n? Full report saved to: {output_file}")

if __name__ == "__main__":
    run_diagnostic_analysis()