#!/usr/bin/env python3
"""Train v10 router on REAL SWE-Router execution data.

This is the big one: 500 tasks x 8 models = 4000 real outcomes.
We learn which model succeeds on which task, at what cost.
"""
import sys, json, random, pickle, math
from collections import defaultdict
from datasets import load_dataset
import numpy as np

print("="*80)
print("TRAINING v10 ROUTER ON REAL SWE-ROUTER DATA")
print("="*80)

# Load all SWE-Router traces
MODELS = ['claude-opus-4.7','gpt-5-mini','gpt-5-nano','gpt-5.2',
          'gemini-2.5-pro','gemini-3-pro','deepseek-v3.2','deepseek-v4-flash']

MODEL_TIER = {
    'deepseek-v4-flash': 1, 'gpt-5-nano': 1,
    'gpt-5-mini': 2, 'deepseek-v3.2': 2,
    'gemini-2.5-pro': 3,
    'claude-opus-4.7': 4, 'gpt-5.2': 4,
    'gemini-3-pro': 5,
}

TIER_COST = {1:0.01, 2:0.05, 3:0.15, 4:0.30, 5:0.50}

print("\n[1] Loading SWE-Router traces...")
traces = defaultdict(dict)
for model in MODELS:
    ds = load_dataset(f'SWE-Router/swebench-verified-{model}', split='test')
    for row in ds:
        iid = row['instance_id']
        traces[iid][model] = {
            'resolved': row['resolved'],
            'cost': float(row['instance_cost']),
            'api_calls': int(row['api_calls']),
            'problem': row['problem_statement'],
        }
    print(f"  {model}: loaded")

print(f"\n  Total tasks: {len(traces)}")
print(f"  Total traces: {sum(len(v) for v in traces.values())}")

# ─── Feature Engineering ────────────────────────────────────────────────
print("\n[2] Engineering features from problem statements...")

# Keyword sets for feature extraction
CODE_KW = ["python","javascript","code","function","bug","debug","refactor","implement","test",
           "compile","runtime","segfault","thread","async","class","module","import","error","traceback"]
LEGAL_KW = ["contract","legal","compliance","gdpr","privacy","policy","regulatory","liability"]
RESEARCH_KW = ["research","investigate","compare","analyze","survey","paper"]
TOOL_KW = ["search","fetch","retrieve","query","api","database","scrape","aggregate"]
CRITICAL_KW = ["critical","production","urgent","emergency","live","deployed","safety","security"]
SIMPLE_KW = ["typo","simple","quick","brief","minor","small","easy","trivial","just"]
LONG_KW = ["plan","project","roadmap","orchestrate","migrate","pipeline","deploy","architecture"]
MATH_KW = ["calculate","compute","solve","equation","formula","optimize","probability"]

def extract_features(problem_text):
    r = problem_text.lower()
    feats = {
        'req_len': len(problem_text),
        'num_words': len(problem_text.split()),
        'has_code': int(any(k in r for k in CODE_KW)),
        'n_code': sum(1 for k in CODE_KW if k in r),
        'has_legal': int(any(k in r for k in LEGAL_KW)),
        'has_research': int(any(k in r for k in RESEARCH_KW)),
        'has_tool': int(any(k in r for k in TOOL_KW)),
        'has_critical': int(any(k in r for k in CRITICAL_KW)),
        'has_simple': int(any(k in r for k in SIMPLE_KW)),
        'has_long': int(any(k in r for k in LONG_KW)),
        'has_math': int(any(k in r for k in MATH_KW)),
        'has_error_msg': int('error' in r or 'traceback' in r or 'exception' in r),
        'has_file_path': int('/' in r and ('.' in r.split('/')[0] if '/' in r else False)),
        'n_lines': problem_text.count('\n') + 1,
        'has_version': int('version' in r or 'update' in r or 'upgrade' in r),
        'has_add': int('add' in r or 'new' in r or 'create' in r),
        'has_fix': int('fix' in r or 'bug' in r or 'issue' in r or 'broken' in r),
        'has_change': int('change' in r or 'modify' in r or 'update' in r),
        'has_remove': int('remove' in r or 'delete' in r or 'drop' in r),
        'has_test': int('test' in r or 'spec' in r or 'assert' in r),
        'has_doc': int('doc' in r or 'readme' in r or 'comment' in r),
        # SWE-specific features
        'has_see_also': int('see also' in r or 'related' in r),
        'has_steps_to_reproduce': int('steps to reproduce' in r or 'reproduce' in r),
    }
    return feats

# ─── Build Training Data ────────────────────────────────────────────────
print("\n[3] Building training data...")

# For each task, we know which models succeeded.
# Ground truth: optimal_tier = cheapest tier where at least one model succeeded
# Features: extracted from problem statement

all_feat_keys = None
training_data = []
tier_labels = {1:[],2:[],3:[],4:[],5:[]}
cost_labels = []

for iid, model_results in traces.items():
    problem = next(iter(model_results.values()))['problem']
    feats = extract_features(problem)
    
    if all_feat_keys is None:
        all_feat_keys = sorted(feats.keys())
    
    feat_vec = [float(feats.get(k, 0.0)) for k in all_feat_keys]
    
    # Determine ground truth: which tiers succeeded?
    tier_success = {}
    for model, result in model_results.items():
        tier = MODEL_TIER[model]
        if tier not in tier_success:
            tier_success[tier] = False
        if result['resolved']:
            tier_success[tier] = True
    
    # Optimal tier = cheapest that succeeded
    optimal_tier = 5
    for t in range(1, 6):
        if tier_success.get(t, False):
            optimal_tier = t
            break
    
    # Per-tier success labels
    for t in range(1, 6):
        tier_labels[t].append(int(tier_success.get(t, False)))
    
    training_data.append({
        'features': feat_vec,
        'optimal_tier': optimal_tier,
        'tier_success': tier_success,
        'cost': min(r['cost'] for r in model_results.values()),
    })

print(f"  Training samples: {len(training_data)}")
print(f"  Features: {len(all_feat_keys)}")
print(f"  Optimal tier distribution:")
opt_dist = defaultdict(int)
for t in training_data:
    opt_dist[t['optimal_tier']] += 1
for tier in sorted(opt_dist.keys()):
    print(f"    Tier {tier}: {opt_dist[tier]} ({opt_dist[tier]/len(training_data)*100:.1f}%)")
print(f"  Per-tier success rates:")
for t in range(1,6):
    s = sum(tier_labels[t])
    print(f"    Tier {t}: {s}/{len(training_data)} = {s/len(training_data)*100:.1f}%")

# ─── Train XGBoost Models ────────────────────────────────────────────────
print("\n[4] Training XGBoost per-tier success predictors...")

from xgboost import XGBClassifier
from sklearn.calibration import IsotonicRegression
from sklearn.model_selection import cross_val_score
import warnings
warnings.filterwarnings('ignore')

X = np.array([t['features'] for t in training_data], dtype=np.float32)
y_tier = {t: np.array(tier_labels[t]) for t in range(1,6)}
y_optimal = np.array([t['optimal_tier'] for t in training_data])

tier_clfs = {}
tier_calibs = {}
tier_cv_scores = {}

for t in range(1, 6):
    y = y_tier[t]
    n_pos = y.sum()
    n_neg = len(y) - n_pos
    
    # Scale pos weight for imbalanced data
    spw = max(1, n_neg / max(n_pos, 1))
    
    clf = XGBClassifier(
        n_estimators=200, max_depth=5, learning_rate=0.05,
        subsample=0.8, colsample_bytree=0.8,
        scale_pos_weight=spw,
        eval_metric='logloss', use_label_encoder=False,
        random_state=42,
    )
    
    # CV score
    try:
        scores = cross_val_score(clf, X, y, cv=5, scoring='f1')
        tier_cv_scores[t] = scores.mean()
    except:
        tier_cv_scores[t] = 0.0
    
    clf.fit(X, y)
    
    # Calibrate
    p_raw = clf.predict_proba(X)[:, 1]
    cal = IsotonicRegression(out_of_bounds='clip')
    cal.fit(p_raw, y)
    
    tier_clfs[t] = clf
    tier_calibs[t] = cal
    
    p_cal = cal.transform(p_raw)
    brier = np.mean((p_cal - y) ** 2)
    
    print(f"  Tier {t}: n_pos={n_pos}, CV_f1={tier_cv_scores[t]:.3f}, Brier={brier:.4f}")

# ─── Train Direct Optimal-Tier Predictor ────────────────────────────────
print("\n[5] Training direct optimal-tier predictor...")

from xgboost import XGBRegressor

opt_clf = XGBClassifier(
    n_estimators=300, max_depth=6, learning_rate=0.05,
    subsample=0.8, colsample_bytree=0.8,
    eval_metric='mlogloss', use_label_encoder=False,
    random_state=42, num_class=5,
)
opt_clf.fit(X, y_optimal - 1)  # 0-indexed
opt_pred = opt_clf.predict(X) + 1
opt_acc = np.mean(opt_pred == y_optimal)
print(f"  Direct optimal-tier accuracy: {opt_acc:.3f}")
print(f"  Confusion (predicted vs actual):")
from collections import Counter
for actual_tier in range(1, 6):
    mask = y_optimal == actual_tier
    if mask.sum() > 0:
        pred_dist = Counter(opt_pred[mask].tolist())
        print(f"    Actual tier {actual_tier}: {dict(pred_dist)}")

# ─── Evaluate on SWE-Router data ────────────────────────────────────────
print("\n[6] Evaluating routing policies on SWE-Router...")

from aco.classifier import TaskCostClassifier
classifier = TaskCostClassifier()

def route_v10(problem_text):
    """v10: Real-data trained router."""
    feats = extract_features(problem_text)
    feat_vec = np.array([float(feats.get(k, 0.0)) for k in all_feat_keys], dtype=np.float32).reshape(1,-1)
    
    # Method 1: Direct optimal tier prediction
    predicted_tier = int(opt_clf.predict(feat_vec)[0]) + 1
    
    # Method 2: Per-tier P(success) cascade
    tier_probs = {}
    for t in range(1, 6):
        p_raw = tier_clfs[t].predict_proba(feat_vec)[0, 1]
        p_cal = float(tier_calibs[t].transform([p_raw])[0])
        tier_probs[t] = p_cal
    
    # Find cheapest tier with P(success) > threshold
    for t in range(1, 6):
        if tier_probs[t] >= 0.5:  # 50% success threshold
            cascade_tier = t
            break
    else:
        cascade_tier = 5
    
    return predicted_tier, cascade_tier, tier_probs

# Evaluate
TIER_TO_SWE = {
    1: 'deepseek-v4-flash', 2: 'gpt-5-mini',
    3: 'gemini-2.5-pro', 4: 'claude-opus-4.7', 5: 'gemini-3-pro',
}

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

for iid, model_results in traces.items():
    problem = next(iter(model_results.values()))['problem']
    
    # 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 (tier 4)
    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
    
    # v8 (old synthetic-trained router)
    pred = classifier.classify(problem)
    from aco.router import ModelCascadeRouter
    old_router = ModelCascadeRouter(model_path="/app/router_models/router_bundle_v8.pkl")
    r8 = old_router.route(problem, "coding", pred["difficulty"], pred)
    m8 = TIER_TO_SWE.get(r8.tier, 'claude-opus-4.7')
    if m8 in model_results:
        policies['v8_synthetic']['success'] += int(model_results[m8]['resolved'])
        policies['v8_synthetic']['cost'] += model_results[m8]['cost']
    policies['v8_synthetic']['n'] += 1
    
    # v10 direct optimal-tier
    predicted_tier, cascade_tier, tier_probs = route_v10(problem)
    m10 = TIER_TO_SWE.get(predicted_tier, 'claude-opus-4.7')
    if m10 in model_results:
        policies['v10_direct']['success'] += int(model_results[m10]['resolved'])
        policies['v10_direct']['cost'] += model_results[m10]['cost']
    else:
        # Fallback to frontier
        policies['v10_direct']['success'] += int(model_results.get('claude-opus-4.7',{}).get('resolved',0))
        policies['v10_direct']['cost'] += model_results.get('claude-opus-4.7',{}).get('cost',0.3)
    policies['v10_direct']['n'] += 1
    
    # v10 cascade (per-tier P(success) > 0.5)
    m10c = TIER_TO_SWE.get(cascade_tier, 'claude-opus-4.7')
    if m10c in model_results:
        policies['v10_cascade']['success'] += int(model_results[m10c]['resolved'])
        policies['v10_cascade']['cost'] += model_results[m10c]['cost']
    else:
        policies['v10_cascade']['success'] += int(model_results.get('claude-opus-4.7',{}).get('resolved',0))
        policies['v10_cascade']['cost'] += model_results.get('claude-opus-4.7',{}).get('cost',0.3)
    policies['v10_cascade']['n'] += 1
    
    # Always cheap (tier 1)
    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

# Print results
print(f"\n\n{'='*80}")
print("REAL SWE-BENCH RESULTS WITH v10 REAL-DATA ROUTER")
print(f"{'='*80}")

fr_cost = policies['frontier']['cost'] / policies['frontier']['n']
print(f"\n{'Policy':<20} {'Success':>10} {'AvgCost':>10} {'CostRed':>10}")
print("-"*50)
for name in ['oracle','v10_direct','v10_cascade','v8_synthetic','frontier','always_cheap']:
    r = policies[name]
    sr = r['success']/r['n']
    ac = r['cost']/r['n']
    cr = (1 - ac/fr_cost)*100
    print(f"{name:<20} {sr:>10.3f} {ac:>10.4f} {cr:>9.1f}%")

# Also try hybrid v10 + feedback
# v10 routes to cheap model; if it fails, escalate
print("\n\n[7] v10 + feedback cascade...")
policies_hybrid = defaultdict(lambda: {"success":0,"cost":0.0,"n":0})

for iid, model_results in traces.items():
    problem = next(iter(model_results.values()))['problem']
    predicted_tier, cascade_tier, tier_probs = route_v10(problem)
    
    # Start with cascade_tier (more conservative than direct)
    m_cascade = TIER_TO_SWE.get(cascade_tier, 'claude-opus-4.7')
    
    if m_cascade in model_results and model_results[m_cascade]['resolved']:
        # Initial model succeeded
        policies_hybrid['v10_feedback']['success'] += 1
        policies_hybrid['v10_feedback']['cost'] += model_results[m_cascade]['cost']
    elif cascade_tier < 5:
        # Failed: escalate
        up_tier = min(cascade_tier + 1, 5)
        up_model = TIER_TO_SWE.get(up_tier, 'claude-opus-4.7')
        if up_model in model_results and model_results[up_model]['resolved']:
            policies_hybrid['v10_feedback']['success'] += 1
            policies_hybrid['v10_feedback']['cost'] += model_results[m_cascade]['cost']
            policies_hybrid['v10_feedback']['cost'] += model_results[up_model]['cost']
        else:
            # Try tier 4 (frontier) as last resort
            f_model = 'claude-opus-4.7'
            if f_model in model_results and model_results[f_model]['resolved']:
                policies_hybrid['v10_feedback']['success'] += 1
                policies_hybrid['v10_feedback']['cost'] += model_results[m_cascade]['cost']
                policies_hybrid['v10_feedback']['cost'] += model_results[f_model]['cost']
            else:
                policies_hybrid['v10_feedback']['cost'] += model_results[m_cascade]['cost']
    else:
        policies_hybrid['v10_feedback']['cost'] += model_results.get(m_cascade, {}).get('cost', 0.3)
    policies_hybrid['v10_feedback']['n'] += 1
    
    # Also track 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_hybrid['oracle']['success'] += 1
        policies_hybrid['oracle']['cost'] += cheapest[1]['cost']
    policies_hybrid['oracle']['n'] += 1

    # Frontier
    f_model = 'claude-opus-4.7'
    policies_hybrid['frontier']['success'] += int(model_results[f_model]['resolved'])
    policies_hybrid['frontier']['cost'] += model_results[f_model]['cost']
    policies_hybrid['frontier']['n'] += 1

fr_cost_h = policies_hybrid['frontier']['cost'] / policies_hybrid['frontier']['n']
print(f"\n{'Policy':<20} {'Success':>10} {'AvgCost':>10} {'CostRed':>10}")
print("-"*50)
for name in ['oracle','v10_feedback','frontier']:
    r = policies_hybrid[name]
    sr = r['success']/r['n']
    ac = r['cost']/r['n']
    cr = (1-ac/fr_cost_h)*100
    print(f"{name:<20} {sr:>10.3f} {ac:>10.4f} {cr:>9.1f}%")

# Save v10 bundle
v10_bundle = {
    'tier_clfs': {str(k):v for k,v in tier_clfs.items()},
    'tier_calibrators': {str(k):v for k,v in tier_calibs.items()},
    'opt_clf': opt_clf,
    'feat_keys': all_feat_keys,
    'tier_config': {str(k):v for k,v in TIER_COST.items()},
    'version': '10.0',
    'description': 'ACO v10: Trained on REAL SWE-Router execution data (500 tasks x 8 models)',
    'training_data': 'SWE-Router/swebench-verified-*',
    'n_training': len(training_data),
    'n_features': len(all_feat_keys),
}
with open('/app/router_models/router_bundle_v10.pkl', 'wb') as f:
    pickle.dump(v10_bundle, f)
print(f"\nSaved router_bundle_v10.pkl ({os.path.getsize('/app/router_models/router_bundle_v10.pkl')/1024:.0f} KB)")

# Save results
all_results = {}
for name, r in policies.items():
    all_results[name] = {"success":r['success']/r['n'],"avg_cost":r['cost']/r['n']}
for name, r in policies_hybrid.items():
    all_results[f"hybrid_{name}"] = {"success":r['success']/r['n'],"avg_cost":r['cost']/r['n']}
all_results['v10_cv_scores'] = tier_cv_scores
all_results['v10_opt_acc'] = opt_acc
all_results['feat_keys'] = all_feat_keys
with open('/app/swe_v10_results.json', 'w') as f:
    json.dump(all_results, f, indent=2, default=str)

print(f"\nSaved swe_v10_results.json")
print("DONE!")