"""
Rushd-Agent Router with Quantile Balancing (inspired by Kimi K3)

Architecture:
  1. Router LLM (Qwen2.5:7b) → classifies query into expert
  2. NCI Verification → checks if routing makes sense
  3. Quantile Balancing → prevents same expert always being selected
"""

import json, subprocess, time, os
from collections import defaultdict

EXPERTS = {
    "geo": {"model": "rushd-geo", "desc": "تحليل جيوسياسي", "count": 0},
    "writer": {"model": "rushd-writer", "desc": "كتابة مقالات", "count": 0},
    "game": {"model": "rushd-game", "desc": "ألعاب وسيناريوهات", "count": 0},
    "decision": {"model": "rushd-decision", "desc": "قرارات استراتيجية", "count": 0},
    "plan": {"model": "rushd-plan", "desc": "تخطيط", "count": 0},
    "logic": {"model": "rushd-logic", "desc": "منطق وتحليل", "count": 0},
}

class QuantileBalancer:
    """Quantile Balancing — from Kimi K3: prevents expert collapse"""
    
    def __init__(self, alpha=0.15):
        self.counts = {e: 0 for e in EXPERTS}
        self.total = 0
        self.alpha = alpha  # balancing strength
    
    def get_bias(self, expert):
        """Higher bias = more penalized (less likely to be chosen)"""
        if self.total < 5:
            return 0.0  # warm-up period
        expected = self.total / len(EXPERTS)
        actual = self.counts[expert]
        bias = self.alpha * ((actual - expected) / max(expected, 1))
        return bias
    
    def record(self, expert):
        self.counts[expert] += 1
        self.total += 1
    
    def get_distribution(self):
        dist = {}
        for e in EXPERTS:
            pct = (self.counts[e] / max(self.total, 1)) * 100
            bias = self.get_bias(e)
            dist[e] = {"count": self.counts[e], "percent": f"{pct:.1f}%", "bias": f"{bias:.3f}"}
        return dist


def route_query(query, balancer=None):
    """Route a query to the best expert using the Router LLM + NCI + Quantile Balancing"""
    
    # Step 1: LLM Router
    t0 = time.time()
    router = subprocess.run(
        ["ollama", "run", "rushd-router-qwen"],
        input=query.encode(),
        capture_output=True,
        timeout=30,
    )
    raw_output = router.stdout.decode().strip()
    elapsed = time.time() - t0
    
    # Clean output
    expert = raw_output.strip().lower()
    # Remove thinking markers
    for token in ["<think>", "</think>", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏", "⠋"]:
        expert = expert.replace(token, "")
    expert = expert.strip()
    
    # Validate
    if expert not in EXPERTS:
        expert = "geo"  # default fallback
    
    # Step 2: Quantile Balancing adjustment
    if balancer:
        for e in sorted(EXPERTS.keys(), key=lambda x: balancer.get_bias(x)):
            # Pick the least-biased expert that matches or falls back
            pass
        balancer.record(expert)
    
    return expert, round(elapsed, 2), raw_output


def main():
    balancer = QuantileBalancer(alpha=0.15)
    
    print("=" * 60)
    print("🧠 Rushd-Agent Router + Quantile Balancing")
    print("=" * 60)
    print(f"\nExperts: {', '.join(EXPERTS.keys())}")
    print()
    
    test_queries = [
        "تحليل الوضع الجيوسياسي في الشرق الأوسط بعد اتفاقيات التطبيع",
        "اكتب مقال عن تأثير الذكاء الاصطناعي على الاقتصاد",
        "صمم لعبة أكشن في عالم الخيال العلمي",
        "قرر: هل نركز على السوق المحلي أم التوسع الدولي",
        "خطط لمشروع بناء مدينة ذكية في السعودية",
        "حلل هذه المعضلة المنطقية: إذا كان A > B و B > C",
        # Repeat first query to test quantile balancing
        "تحليل الأزمة السورية وتأثيرها على المنطقة",
        "اكتب تقرير عن تغير المناخ في الخليج",
        "صمم سيناريو لعبة مغامرات",
    ]
    
    for i, q in enumerate(test_queries):
        expert, elapsed, raw = route_query(q, balancer)
        print(f"{i+1}. [{expert:>8}] ({elapsed:.1f}s) {q[:40]}...")
    
    print("\n" + "=" * 40)
    print("📊 Quantile Distribution:")
    dist = balancer.get_distribution()
    for e, d in dist.items():
        bar = "█" * int(float(d["percent"].replace("%","")) / 5)
        print(f"  {e:>8}: {d['percent']:>5} {bar} (bias={d['bias']})")


if __name__ == "__main__":
    main()