Blame for core/triage.py

import torch
import numpy as np
from dataclasses import dataclass
from typing import List, Dict, Tuple

# ==========================================
# 1. CONSTANTS (Aligned with V7/V5 Specs)
# ==========================================
PHONEMES = [
    'p', 'pʰ', 'b', 'bʰ', 't', 'tʰ', 'd', 'dʰ', 't̪', 't̪ʰ', 'd̪', 'd̪ʱ',
    'ʈ', 'ɖ', 'k', 'kʰ', 'g', 'ɡʰ', 'q', 'ʔ', 'f', 'v', 's', 'z',
    'ʃ', 'ʒ', 'ʂ', 'ʐ', 'x', 'ɣ', 'h', 'θ', 'ð', 'tʃ', 'dʒ', 'm',
    'mʰ', 'n', 'nʰ', 'ŋ', 'ɳ', 'l', 'lʰ', 'r', 'rʰ', 'ɽ', 'ɽʰ', 'w',
    'j', 'a', 'aː', 'i', 'iː', 'u', 'uː', 'e', 'eː', 'o', 'oː',
    'ə', 'æ', 'ɔ', 'ɪ', 'ʊ'
]
P2I = {p: i for i, p in enumerate(PHONEMES)}

ERRORS = ["Sub", "Front", "Stop", "Glide", "Clust", "Epen", "Block", "Prolong", "Repeat"]

# Mapping Error Index (0-8) to Major Type Category
# 0-4 = Articulation, 5 = Motor, 6-8 = Fluency
MAJOR_MAP = {
    0: "Artic", 1: "Artic", 2: "Artic", 3: "Artic", 4: "Artic",
    5: "Motor", 6: "Fluency", 7: "Fluency", 8: "Fluency"
}

# Clinical Priority: Treat Motor issues first, then Fluency, then Artic
PRIORITY = {"Motor": 3, "Fluency": 2, "Artic": 1}

@dataclass
class ClinicalDiagnosis:
    id: int
    error_idx: int
    error_name: str
    major_type: str
    severity: float
    active_indices: List[int]
    active_phonemes: List[str]
    position_idx: int

class TriageEngineV2:
    def __init__(self, device='cpu'):
        self.device = device

    def scan_tensor(self, tensor_3d: np.ndarray) -> List[ClinicalDiagnosis]:
        """
        Scans the 64x9x3 AMLM output map for clinical syndromes.
        Returns a sorted list of diagnoses.
        """
        diagnoses = []
        uid_counter = 0

        # tensor_3d shape: (3, 64, 9) -> (Channels, Phonemes, Errors)
        
        for err_idx in range(9):
            # Check maximum probability for this error type across all phonemes/positions
            slice_max = np.max(tensor_3d[:, :, err_idx])

            # STRICT FILTERING: Threshold > 0.6 means the AMLM is confident
            if slice_max > 0.6:
                # Find which phonemes are affected
                # We collapse the channel dim (positions) to see if the phoneme is affected anywhere
                hits = np.where(tensor_3d[:, :, err_idx].max(axis=0) > 0.6)[0]
                
                # Filter specific weird edge cases (e.g. Gliding on vowels) if necessary
                # Logic from Integration.ipynb: valid_hits = [int(h) for h in hits if not (err_idx == 3 and h >= 48)]
                valid_hits = [int(h) for h in hits]

                # NOISE REDUCTION: Cap at Top 10 worst errors per category
                if len(valid_hits) > 10:
                    severity_scores = []
                    for h in valid_hits:
                        score = tensor_3d[:, h, err_idx].max()
                        severity_scores.append((score, h))
                    severity_scores.sort(key=lambda x: x[0], reverse=True)
                    valid_hits = [x[1] for x in severity_scores[:10]]

                valid_hits.sort()

                if not valid_hits: continue

                # Determine dominant position (Initial, Medial, or Final)
                pos_sums = tensor_3d[:, valid_hits, err_idx].sum(axis=1)
                dom_pos = np.argmax(pos_sums)

                diag = ClinicalDiagnosis(
                    id=uid_counter,
                    error_idx=err_idx,
                    error_name=ERRORS[err_idx],
                    major_type=MAJOR_MAP[err_idx],
                    severity=float(slice_max),
                    active_indices=valid_hits,
                    active_phonemes=[PHONEMES[i] for i in valid_hits],
                    position_idx=int(dom_pos)
                )
                diagnoses.append(diag)
                uid_counter += 1

        # Sort by Priority (Motor > Fluency > Artic) then Severity
        return sorted(diagnoses, key=lambda x: (PRIORITY[x.major_type], x.severity), reverse=True)

    def prepare_aeg_input(self, focus_diag: ClinicalDiagnosis, all_diagnoses: List[ClinicalDiagnosis], user_stats: List[float]) -> Dict[str, torch.Tensor]:
        """
        Generates the 162-Dim Input Vector for the AEG Neural Network.
        Adheres strictly to 'Data Generator.py' logic.
        """
        # 1. Cluster Mask (Targets)
        # Which sounds should we target in therapy? (The ones in the focus diagnosis)
        cluster_mask = [0.0] * 64
        target_set = set(focus_diag.active_indices)
        for idx in focus_diag.active_indices:
            cluster_mask[idx] = 1.0

        # 2. Trap Mask (Forbidden Sounds)
        # Sounds we must NOT use to avoid negative reinforcement.
        trap_mask = [0.0] * 64

        def attempt_trap(idx):
            # Logic: A sound cannot be a Trap if it is currently a Target.
            if idx not in target_set:
                trap_mask[idx] = 1.0

        # A. Standard Traps (Rule-based from Data Generator)
        if focus_diag.error_name == "Front": # If Fronting K->T
            # Trap T/D so we don't ask user to say them while trying to learn K
            if any(p in ['k', 'g', 'kʰ', 'ɡʰ'] for p in focus_diag.active_phonemes):
                for x in ['t', 'd', 'tʰ', 'dʰ']:
                    if x in P2I: attempt_trap(P2I[x])

        # B. Comorbidity Traps (Dynamic)
        # If user also has a Lisp (Sub S), Trap 'S' even if we are treating Fronting.
        for other in all_diagnoses:
            if other.id != focus_diag.id:
                for idx in other.active_indices:
                    attempt_trap(idx)

        # 3. Phonetic Features (8-Dim Vector)
        # Describes the target sound (e.g., +Stop, +Velar)
        feat = [0.0] * 8
        t_str = "".join(focus_diag.active_phonemes)

        if any(x in t_str for x in "pbtdkgʈɖʔ"): feat[0] = 1.0 # Stop
        if any(x in t_str for x in "fvszʃʒxhθð"): feat[1] = 1.0 # Fricative
        if any(x in t_str for x in "lrw"): feat[2] = 1.0 # Liquid
        if any(x in t_str for x in "mnŋ"): feat[3] = 1.0 # Nasal
        if any(x in t_str for x in "bdgvzʒðmnlrw"): feat[4] = 1.0 # Voiced
        if any(x in t_str for x in "pbfvmθð"): feat[5] = 1.0 # Front
        if any(x in t_str for x in "kgŋxɣhʔ"): feat[6] = 1.0 # Back
        if len(focus_diag.active_indices) > 1 or focus_diag.error_name == "Clust": feat[7] = 1.0 # Cluster

        # 4. Flags & Stats
        trap_count = sum(trap_mask)
        # user_stats expects: [Severity, Fatigue, Age, Delta]
        sev, fat, age, delta = user_stats

        # Binary Flags for the Network
        is_systemic = 1.0 if len(focus_diag.active_indices) > 1 else 0.0
        is_trap = 1.0 if trap_count > 0 else 0.0
        is_comorb = 1.0 if len(all_diagnoses) > 1 else 0.0
        is_child = 1.0 if age < 0.3 else 0.0 # Normalized age (0.3 ~ 7 years old)
        is_maint = 1.0 if sev <= 0.2 else 0.0
        is_fatigued = 1.0 if fat > 0.8 else 0.0

        flags = [is_systemic, is_trap, is_comorb, is_child, is_maint, is_fatigued]

        maj_int = {"Artic": 0, "Fluency": 1, "Motor": 2}[focus_diag.major_type]

        # Convert to PyTorch Tensors on correct Device
        return {
            "error_idx": torch.tensor([focus_diag.error_idx]).to(self.device),
            "major_idx": torch.tensor([maj_int]).to(self.device),
            "pos_idx": torch.tensor([focus_diag.position_idx]).to(self.device),
            "cluster": torch.tensor([cluster_mask], dtype=torch.float32).to(self.device),
            "trap": torch.tensor([trap_mask], dtype=torch.float32).to(self.device),
            "feat": torch.tensor([feat], dtype=torch.float32).to(self.device),
            "flags": torch.tensor([flags], dtype=torch.float32).to(self.device),
            "stats": torch.tensor([user_stats], dtype=torch.float32).to(self.device)
        }