server/reward.py

try :
    from server.disorders import get_disorder_info
except ImportError:
    from disorders import get_disorder_info

def disorder_similarity(disorder1, disorder2):
    """
    Calculate similarity between two disorders based on their symptoms and trigger keywords.
    
    Args:
        disorder1 (str): Name of the first disorder.
        disorder2 (str): Name of the second disorder.

    Returns:
        str: Similarity level ("full","partial","wrong") based on symptom and trigger keyword overlap.
    """
    info1 = get_disorder_info(disorder1)
    info2 = get_disorder_info(disorder2)

    symptoms1 = set(info1.get("symptoms", {}).get("emotional", []))
    symptoms2 = set(info2.get("symptoms", {}).get("emotional", []))

    triggers1 = set(info1.get("trigger_keywords", []))
    triggers2 = set(info2.get("trigger_keywords", []))

    symptom_overlap = len(symptoms1.intersection(symptoms2)) / max(len(symptoms1), 1)
    trigger_overlap = len(triggers1.intersection(triggers2)) / max(len(triggers1), 1)

    if symptom_overlap > 0.7 and trigger_overlap > 0.7:
        return "full"
    elif symptom_overlap > 0.3 and trigger_overlap > 0.3:
        return "partial"
    else:
        return "wrong"

def check_action_repeat(action_sequence, current_action):
    """
    Check if the current action has been repeated in the recent action sequence.

    Args:
        action_sequence (List[str]): List of recent action types taken by the agent.
        current_action (str): The action type of the current step.
    Returns:
        float: A penalty value based on the frequency of the current action in the recent sequence.
    """
    repeat_count = action_sequence.count(current_action)
    if repeat_count > 2:
        return -0.5  # High penalty for excessive repetition
    elif repeat_count > 0:
        return -0.2  # Moderate penalty for some repetition
    else:
        return 0.0   # No penalty for unique actions


def normalize(score, min_score, max_score):
    if max_score == min_score:
        return 0.0
    return max(0.0, min(1.0, (score - min_score) / (max_score - min_score)))


def calculate_reward(
    disorder,
    action_type,
    task_difficulty,
    diagnosis=None,
    action_sequence=None,
    state=None
):
    """Calculate reward based on the agent's action, diagnosis, and the patient's state.

    Args:
        disorder (str): The actual disorder of the patient.
        action_type (str): The type of action taken by the agent.
        task_difficulty (str): The difficulty level of the task ("easy", "medium", "hard").
        diagnosis (str, optional): The diagnosis provided by the agent, if any.
        action_sequence (List[str], optional): The sequence of recent actions taken by the agent.
        state (MentalState, optional): The current state of the patient.

    Returns:
        float: A reward value between 0.0 and 1.0 based on the quality of the action and diagnosis.
    """
    score = 0.0

    # =========================
    # EASY → correctness only
    # =========================
    if task_difficulty == "easy":

        if action_type == "diagnose" and diagnosis:
            similarity = disorder_similarity(disorder, diagnosis)

            if similarity == "full":
                score = 1.0
            elif similarity == "partial":
                score = 0.5
            else:
                score = 0.0

        else:
            score = 0.1  # small reward for non-diagnose steps

        return score  # already normalized

    # =========================
    # MEDIUM → strategy
    # =========================
    elif task_difficulty == "medium":

        # Action quality (0 → 1 range)
        action_score = {
            "ask_open": 0.6,
            "ask_direct": 0.4,
            "reflect": 0.8,
            "ask_risk": 0.5,
            "diagnose": 0.0
        }.get(action_type, 0.0)

        # Diagnosis score
        diagnosis_score = 0.0
        if action_type == "diagnose" and diagnosis:
            if len(diagnosis)==1:
                similarity = disorder_similarity(disorder, diagnosis[-1])

                if similarity == "full":
                    diagnosis_score = 1.0
                elif similarity == "partial":
                    diagnosis_score = 0.5
                else:
                    diagnosis_score = 0.0
            else:
                diagnosis_score=0.0
                for d in diagnosis:
                    similarity=disorder_similarity(disorder,d)

                    if similarity=="full":
                        diagnosis_score+=0.5
                    elif diagnosis_score=="partial":
                        diagnosis_score+=0.3
                    else:
                        diagnosis_score-=0.5
                diagnosis_score=normalize(diagnosis_score,0.0,1.0)


        # Repetition penalty (converted to positive scale)
        repeat_penalty = 0.0
        if action_sequence:
            penalty = check_action_repeat(action_sequence, action_type)
            repeat_penalty = max(0.0, 1.0 + penalty)  # convert (-0.5 → 0.5)

        # Combine (weighted)
        score = (
            0.4 * action_score +
            0.4 * diagnosis_score +
            0.2 * repeat_penalty
        )

        return max(0.0, min(1.0, score))

    # =========================
    # HARD → empathy + safety + depth
    # =========================
    elif task_difficulty == "hard":

        # EMPATHY
        empathy_score = {
            "reflect": 1.0,
            "ask_open": 0.7,
            "ask_risk": 0.8,
            "ask_direct": 0.3,
            "diagnose": 0.0
        }.get(action_type, 0.0)

        # SAFETY
        safety_score = 0.5  # neutral
        if state:
            if state.risk_level == "high":
                if action_type == "ask_risk":
                    safety_score = 1.0
                else:
                    safety_score = 0.0

        # DEPTH
        depth_score = 0.0
        if state:
            depth_score = (state.trust_level + state.openness) / 2.0

        # Diagnosis (less important here)
        diagnosis_score = 0.0
        if action_type == "diagnose" and diagnosis:
            similarity = disorder_similarity(disorder, diagnosis)

            if similarity == "full":
                diagnosis_score = 0.8
            elif similarity == "partial":
                diagnosis_score = 0.4
            else:
                diagnosis_score = 0.0

        # Repetition
        repeat_score = 1.0
        if action_sequence:
            penalty = check_action_repeat(action_sequence, action_type)
            repeat_score = max(0.0, 1.0 + penalty)

        # Final weighted score
        score = (
            0.3 * empathy_score +
            0.25 * safety_score +
            0.2 * depth_score +
            0.15 * diagnosis_score +
            0.1 * repeat_score
        )

        return max(0.0, min(1.0, score))

    return 0.0