# self_correction.py
from typing import Dict, Any

class SelfCorrector:
    """
    Implements the "Self-Correction Loop" from Paper 2.
    It diagnoses low scores and maps them to correction plans.
    """
    
    def __init__(self, threshold: float = 3.0):
        self.threshold = threshold

    def is_good_enough(self, v_fitness_scores: Dict[str, int]) -> bool:
        """
        Checks if the solution is good enough to be presented.
        Checks if *all* scores are at or above the threshold.
        """
        print(f"Checking scores {v_fitness_scores} against threshold {self.threshold}")
        for score in v_fitness_scores.values():
            if score < self.threshold:
                print("Score is too low. Initiating Self-Correction.")
                return False
        print("Score is good. Solution accepted.")
        return True

    def get_correction_plan(self, v_fitness_json: Dict[str, Any]) -> str:
        """
        Implements the "Diagnostic Error-to-Belbin Role Mapping" (Paper 2).
        It analyzes the full v_fitness JSON (with justifications)
        and generates a "Chain-of-Thought" correction prompt.
        """
        
        # 1. Find the lowest-scoring criterion
        lowest_score = 5
        lowest_metric = "None"
        for metric, data in v_fitness_json.items():
            if data.get('score', 5) < lowest_score:
                lowest_score = data.get('score', 5)
                lowest_metric = metric
        
        failure_justification = v_fitness_json.get(lowest_metric, {}).get('justification', "No justification provided.")
        
        # 2. Map low score to a failure diagnosis (from Paper 2)
        if lowest_metric == "Novelty":
            failure_diagnosis = f"Ideation Failure (Low {lowest_metric}). The judge's feedback was: '{failure_justification}'"
        elif lowest_metric == "Usefulness_Feasibility":
            failure_diagnosis = f"Compositional Error (Low {lowest_metric}). The judge's feedback was: '{failure_justification}'"
        elif lowest_metric == "Cultural_Appropriateness":
            failure_diagnosis = f"Sensitivity Error (Low {lowest_metric}). The judge's feedback was: '{failure_justification}'"
        else:
            failure_diagnosis = f"General Quality Failure (Low {lowest_metric}). The judge's feedback was: '{failure_justification}'"

        # 3. Generate the "Chain-of-Thought" correction prompt (Paper 2, section 4.9.2)
        correction_prompt = f"""
        YOUR PREVIOUS ATTEMPT FAILED.
        
        FAILURE ANALYSIS:
        Your solution was evaluated and received a very low score for {lowest_metric}.
        {failure_diagnosis}

        YOUR TASK:
        You MUST re-generate a new solution. This new solution must *specifically* address this failure.
        1.  **Analyze the Failure**: Briefly explain *why* the previous solution failed to be {lowest_metric.lower()}.
        2.  **Formulate a New Plan**: Outline a new plan that directly fixes this specific failure.
        3.  **Write the Corrected Solution**: Generate the full, new solution based on this plan.
        """
        return correction_prompt