# ── Part 3: Run evaluations ──

def route_bert(problem_text):
    """Route using BERT classifier."""
    inputs = tokenizer(problem_text, truncation=True, max_length=512, return_tensors="pt")
    with torch.no_grad():
        logits = bert_model(**inputs).logits
    pred_class = torch.argmax(logits, dim=-1).item()
    tier = pred_class + 1  # classes are 0-4 → tiers 1-5
    probs = torch.softmax(logits, dim=-1)[0]
    confidence = float(probs[pred_class])
    return tier, confidence

def route_v11(problem_text):
    """Route using v11 XGBoost + isotonic calibration."""
    feats = extract_features(problem_text)
    feat_vec = np.array([float(feats.get(k, 0.0)) for k in v11_feat_keys], dtype=np.float32).reshape(1,-1)
    tier_probs = {}
    for t in range(1, 6):
        p_raw = v11_tier_clfs[t].predict_proba(feat_vec)[0, 1]
        p_cal = float(v11_tier_calibs[t].transform([p_raw])[0])
        tier_probs[t] = p_cal
    # Find cheapest tier with P(success) >= 0.65
    for t in range(1, 6):
        if tier_probs[t] >= 0.65:
            return t, tier_probs[t], tier_probs
    return 5, tier_probs[5], tier_probs

TASK_FLOOR = {"quick_answer":1,"coding":3,"research":3,"document_drafting":2,
              "legal_regulated":4,"tool_heavy":2,"retrieval_heavy":2,"long_horizon":3,"unknown_ambiguous":3}

def classify_task(text):
    r = text.lower()
    if any(k in r for k in ["contract","legal","compliance","gdpr","privacy"]): return "legal_regulated"
    if any(k in r for k in ["debug","fix","bug","implement","refactor","code","function"]): return "coding"
    if any(k in r for k in ["research","find sources","literature","investigate"]): return "research"
    if any(k in r for k in ["search","fetch","api","query","database"]): return "tool_heavy"
    if any(k in r for k in ["plan","roadmap","orchestrate","migrate","deploy"]): return "long_horizon"
    if any(k in r for k in ["draft","write","compose","document"]): return "document_drafting"
    if any(k in r for k in ["what is","explain","define","briefly"]): return "quick_answer"
    return "unknown_ambiguous"

policies = defaultdict(lambda: {"success":0,"cost":0.0,"n":0})

print("\n[4] Evaluating all policies on SWE-Router...")
for iid, model_results in traces.items():
    problem = next(iter(model_results.values()))['problem']
    task_type = classify_task(problem)
    floor = TASK_FLOOR.get(task_type, 2)

    # Oracle
    resolved = [(m, r) for m, r in model_results.items() if r['resolved']]
    if resolved:
        cheapest = min(resolved, key=lambda x: TIER_COST.get(MODEL_TIER[x[0]], 1.0))
        policies['oracle']['success'] += 1
        policies['oracle']['cost'] += cheapest[1]['cost']
    else:
        policies['oracle']['cost'] += min(r['cost'] for r in model_results.values())
    policies['oracle']['n'] += 1

    # Always frontier
    f_model = 'claude-opus-4.7'
    if f_model in model_results:
        policies['frontier']['success'] += int(model_results[f_model]['resolved'])
        policies['frontier']['cost'] += model_results[f_model]['cost']
    policies['frontier']['n'] += 1

    # BERT router
    bert_tier, bert_conf = route_bert(problem)
    bert_tier = max(bert_tier, floor)  # enforce safety floor
    m_bert = TIER_TO_SWE.get(bert_tier, 'claude-opus-4.7')
    if m_bert in model_results:
        policies['bert']['success'] += int(model_results[m_bert]['resolved'])
        policies['bert']['cost'] += model_results[m_bert]['cost']
    else:
        policies['bert']['success'] += int(model_results.get('claude-opus-4.7',{}).get('resolved',0))
        policies['bert']['cost'] += model_results.get('claude-opus-4.7',{}).get('cost',0.3)
    policies['bert']['n'] += 1

    # v11 XGBoost
    v11_tier, v11_conf, v11_probs = route_v11(problem)
    v11_tier = max(v11_tier, floor)  # enforce safety floor
    m_v11 = TIER_TO_SWE.get(v11_tier, 'claude-opus-4.7')
    if m_v11 in model_results:
        policies['v11_xgboost']['success'] += int(model_results[m_v11]['resolved'])
        policies['v11_xgboost']['cost'] += model_results[m_v11]['cost']
    else:
        policies['v11_xgboost']['success'] += int(model_results.get('claude-opus-4.7',{}).get('resolved',0))
        policies['v11_xgboost']['cost'] += model_results.get('claude-opus-4.7',{}).get('cost',0.3)
    policies['v11_xgboost']['n'] += 1

    # BERT + feedback (escalate on failure)
    if m_bert in model_results and model_results[m_bert]['resolved']:
        policies['bert_feedback']['success'] += 1
        policies['bert_feedback']['cost'] += model_results[m_bert]['cost']
    else:
        # Escalate one tier
        up_tier = min(bert_tier + 1, 5)
        m_up = TIER_TO_SWE.get(up_tier, 'claude-opus-4.7')
        if m_up in model_results and model_results[m_up]['resolved']:
            policies['bert_feedback']['success'] += 1
            policies['bert_feedback']['cost'] += model_results.get(m_bert,{}).get('cost',0.01)
            policies['bert_feedback']['cost'] += model_results[m_up]['cost']
        else:
            # Last resort: frontier
            if f_model in model_results and model_results[f_model]['resolved']:
                policies['bert_feedback']['success'] += 1
                policies['bert_feedback']['cost'] += model_results.get(m_bert,{}).get('cost',0.01)
                policies['bert_feedback']['cost'] += model_results[f_model]['cost']
            else:
                policies['bert_feedback']['cost'] += model_results.get(m_bert,{}).get('cost',0.01)
    policies['bert_feedback']['n'] += 1

    # v11 + feedback (escalate on failure)
    if m_v11 in model_results and model_results[m_v11]['resolved']:
        policies['v11_feedback']['success'] += 1
        policies['v11_feedback']['cost'] += model_results[m_v11]['cost']
    else:
        up_tier = min(v11_tier + 1, 5)
        m_up = TIER_TO_SWE.get(up_tier, 'claude-opus-4.7')
        if m_up in model_results and model_results[m_up]['resolved']:
            policies['v11_feedback']['success'] += 1
            policies['v11_feedback']['cost'] += model_results.get(m_v11,{}).get('cost',0.01)
            policies['v11_feedback']['cost'] += model_results[m_up]['cost']
        else:
            if f_model in model_results and model_results[f_model]['resolved']:
                policies['v11_feedback']['success'] += 1
                policies['v11_feedback']['cost'] += model_results.get(m_v11,{}).get('cost',0.01)
                policies['v11_feedback']['cost'] += model_results[f_model]['cost']
            else:
                policies['v11_feedback']['cost'] += model_results.get(m_v11,{}).get('cost',0.01)
    policies['v11_feedback']['n'] += 1

    # Always cheap
    c_model = 'deepseek-v4-flash'
    if c_model in model_results:
        policies['always_cheap']['success'] += int(model_results[c_model]['resolved'])
        policies['always_cheap']['cost'] += model_results[c_model]['cost']
    policies['always_cheap']['n'] += 1