# ============================================================
# EVALUATOR — TinyKAN + XAI
# ============================================================
import torch, torch.nn as nn, numpy as np, gc, os
from kan import KAN
from sentence_transformers import SentenceTransformer
from dataclasses import dataclass
from typing import List

torch.set_num_threads(int(os.environ.get('TORCH_THREADS', '1')))

class TinyKANDistilled(nn.Module):
    def __init__(self, input_dim=9, grid=3):
        super().__init__()
        self.kan_main = KAN(width=[[input_dim,0],[16,0],[8,0],[1,0]], grid=grid)
        self.linear_branch = nn.Sequential(nn.Linear(input_dim,8), nn.GELU(), nn.Linear(8,1))
        self.gate = nn.Sequential(nn.Linear(input_dim,8), nn.ReLU(), nn.Linear(8,2), nn.Softmax(dim=1))
        self.skip = nn.Linear(input_dim, 1)
    def forward(self, x):
        k = self.kan_main(x); l = self.linear_branch(x); g = self.gate(x)
        return (g * torch.cat([k,l],dim=1)).sum(dim=1,keepdim=True) + 0.05*self.skip(x)
    def get_gate_weights(self, x):
        with torch.no_grad(): return self.gate(x).cpu().numpy()

@dataclass
class EvalResult:
    score: float; score_pct: float; cefr: str; confidence: float
    gate_kan: float; gate_linear: float; features: np.ndarray
    importance: np.ndarray; feedback: str; xai_explanation: str

class KANEvaluator:
    FEATURE_NAMES = ['cos_MiniLM','cos_MPNet','cos_RoBERTa',
                     '|MiniLM-MPNet|','|MiniLM-RoBERTa|','|MPNet-RoBERTa|',
                     'mean(cos)','max(cos)','min(cos)']
    CEFR_LEVELS = [(0.90,'C2'),(0.80,'C1'),(0.70,'B2'),(0.55,'B1'),(0.40,'A2'),(0.0,'A1')]
    ENCODERS = {'MiniLM':'sentence-transformers/all-MiniLM-L6-v2',
                'MPNet':'sentence-transformers/all-mpnet-base-v2',
                'RoBERTa':'sentence-transformers/stsb-roberta-base'}

    def __init__(self, model_path, device=None):
        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = TinyKANDistilled(9,3).to(self.device)
        self.model.load_state_dict(torch.load(model_path, map_location=self.device, weights_only=False))
        self.model.eval()
        print(f"[Evaluator] TinyKAN loaded ✅ ({sum(p.numel() for p in self.model.parameters()):,} params)")
        self.encoders = {}
        for name, path in self.ENCODERS.items():
            enc = SentenceTransformer(path, device=self.device)
            enc.eval()
            self._quantize(enc, name)
            self.encoders[name] = enc
            gc.collect()
        print("[Evaluator] Encoders loaded ✅ (dynamically quantized, int8)")

    def _quantize(self, sentence_transformer, name):
        """Applies dynamic int8 quantization to the Linear layers of the
        underlying HF transformer to reduce resident memory on CPU.
        Falls back silently (keeps full precision) if unsupported."""
        try:
            inner = sentence_transformer[0].auto_model
            quantized = torch.quantization.quantize_dynamic(
                inner, {torch.nn.Linear}, dtype=torch.qint8)
            sentence_transformer[0].auto_model = quantized
        except Exception as e:
            print(f"[Evaluator] Quantization skipped for {name}: {e}")

    def _features(self, s1, s2):
        cos = []
        for enc in self.encoders.values():
            e1 = enc.encode([s1], convert_to_numpy=True)
            e2 = enc.encode([s2], convert_to_numpy=True)
            e1 /= np.maximum(np.linalg.norm(e1,axis=1,keepdims=True),1e-8)
            e2 /= np.maximum(np.linalg.norm(e2,axis=1,keepdims=True),1e-8)
            cos.append(float((e1*e2).sum()))
        c1,c2,c3 = cos
        return np.array([c1,c2,c3,abs(c1-c2),abs(c1-c3),abs(c2-c3),(c1+c2+c3)/3,max(cos),min(cos)],dtype=np.float32)

    def _predict(self, feat):
        x = torch.tensor(feat,dtype=torch.float32).unsqueeze(0).to(self.device)
        with torch.no_grad(): return float(torch.sigmoid(self.model(x)).cpu().numpy()[0][0])

    def _importance(self, feat):
        base = self._predict(feat)
        imp = np.array([abs(base - self._predict(np.where(np.arange(9)==i, 0, feat))) for i in range(9)])
        return imp / imp.sum() if imp.sum() > 0 else imp

    def _cefr(self, score):
        for t,l in self.CEFR_LEVELS:
            if score >= t: return l
        return 'A1'

    STRICT_CATEGORIES = {'grammar', 'conjugation', 'spelling'}

    @staticmethod
    def _normalize(s):
        import re
        return re.sub(r"[^\w\s']", '', s.lower()).split()

    def _surface_penalty(self, learner, reference, source_sentence=None):
        """For exercise categories that test a specific grammatical form
        (not a paraphrase), pure semantic similarity is largely blind to
        small but critical errors, since a missing function word barely
        shifts the sentence embedding. Two mechanisms address this:
        (1) if the answer is essentially identical to the flawed source
        sentence, no correction was attempted at all — the score is
        collapsed near zero regardless of semantic closeness; (2) otherwise,
        a continuous word-overlap ratio against the reference is squared,
        so partial corrections are penalised more than a linear measure
        would, while an exact match keeps full credit."""
        import difflib
        lw, rw = self._normalize(learner), self._normalize(reference)
        if lw == rw:
            return 1.0
        if source_sentence:
            sw = self._normalize(source_sentence)
            if lw == sw:
                return 0.10
        ratio = difflib.SequenceMatcher(None, lw, rw).ratio()
        penalty = ratio ** 2
        if len(lw) != len(rw):
            # a differing word count usually means a grammatical element was
            # dropped or added — exactly what these exercises test — so the
            # credit is capped regardless of how similar the rest reads.
            penalty = min(penalty, 0.5)
        return max(0.15, penalty)

    def evaluate(self, reference, learner, category=None, source_sentence=None):
        feat = self._features(reference, learner)
        raw_score = self._predict(feat)
        penalty = self._surface_penalty(learner, reference, source_sentence) if category in self.STRICT_CATEGORIES else 1.0
        score = raw_score * penalty
        cefr = self._cefr(score)
        imp = self._importance(feat)
        x = torch.tensor(feat,dtype=torch.float32).unsqueeze(0).to(self.device)
        gates = self.model.get_gate_weights(x)[0]
        top_idx = np.argmax(imp)
        top_feat = self.FEATURE_NAMES[top_idx]

        if score >= 0.90: fb = f"Excellent — Score {score*100:.1f}% | Level {cefr}"
        elif score >= 0.75: fb = f"Good — Score {score*100:.1f}% | Level {cefr}"
        elif score >= 0.55: fb = f"Satisfactory — Score {score*100:.1f}% | Level {cefr}"
        elif score >= 0.40: fb = f"Needs improvement — Score {score*100:.1f}% | Level {cefr}"
        else: fb = f"Insufficient — Score {score*100:.1f}% | Level {cefr}"

        xai = f"Most influential feature: '{top_feat}' ({imp[top_idx]*100:.1f}% importance)"
        if top_idx < 3: xai += f" — Model relies on {['MiniLM','MPNet','RoBERTa'][top_idx]} cosine similarity."
        elif top_idx < 6: xai += " — Encoder divergence detected (semantic difficulty signal)."
        else: xai += f" — {['Mean','Maximum','Minimum'][top_idx-6]} cosine is the decisive statistic."
        if penalty <= 0.10:
            xai += f" Note: semantic similarity alone gave {raw_score*100:.1f}%; the answer matches the uncorrected source sentence, so no correction was made — score collapsed accordingly."
        elif penalty < 1.0:
            xai += f" Note: semantic similarity alone gave {raw_score*100:.1f}%; partial mismatch against the expected correction reduced the final score."

        return EvalResult(score=score, score_pct=score*100, cefr=cefr,
                         confidence=min(min([abs(score-t) for t,_ in self.CEFR_LEVELS])*2,1.0),
                         gate_kan=float(gates[0]), gate_linear=float(gates[1]),
                         features=feat, importance=imp, feedback=fb, xai_explanation=xai)