#!/usr/bin/env python3
"""
dual_quat_model_selector.py — Mass semantic model selector using dual quaternions.

Each model's semantic mass is a dual quaternion:
  Real part (q_real): compressive components (H, I, C, quality)
  Dual part (q_dual): anti-compressive components (R, L, noise, cost)

The chiral ratio χ = |q_real|² / (|q_real|² + |q_dual|²) determines
whether to keep (χ > 0.5) or drop (χ < 0.5) each model.

The panel is a dual quaternion matrix. Selection = project onto the
real subspace, keeping only models with χ > threshold.

Usage:
  python3 dual_quat_model_selector.py [--threshold 0.5] [--problem-type lean_proof]
"""

import argparse
import json
import math
from dataclasses import dataclass, field
from typing import Optional


# ── Dual Quaternion ──────────────────────────────────────────────────────

@dataclass
class DualQuat:
    """Dual quaternion: q = q_real + ε·q_dual, where ε² = 0."""
    # Real part (4 components)
    rw: float = 0.0  # scalar
    rx: float = 0.0  # i
    ry: float = 0.0  # j
    rz: float = 0.0  # k
    # Dual part (4 components)
    dw: float = 0.0  # ε·scalar
    dx: float = 0.0  # ε·i
    dy: float = 0.0  # ε·j
    dz: float = 0.0  # ε·k

    @property
    def real_norm_sq(self) -> float:
        """|q_real|² = rw² + rx² + ry² + rz²"""
        return self.rw**2 + self.rx**2 + self.ry**2 + self.rz**2

    @property
    def dual_norm_sq(self) -> float:
        """|q_dual|² = dw² + dx² + dy² + dz²"""
        return self.dw**2 + self.dx**2 + self.dy**2 + self.dz**2

    @property
    def chi(self) -> float:
        """Chiral ratio χ = |q_real|² / (|q_real|² + |q_dual|²)"""
        r = self.real_norm_sq
        d = self.dual_norm_sq
        total = r + d
        if total == 0:
            return 0.5  # critical balance
        return r / total

    @property
    def is_compressible(self) -> bool:
        """χ > 0.5: real (compressive) dominates"""
        return self.chi > 0.5

    @property
    def mass(self) -> float:
        """Total semantic mass = |q_real| + |q_dual|"""
        return math.sqrt(self.real_norm_sq) + math.sqrt(self.dual_norm_sq)

    def __mul__(self, other: 'DualQuat') -> 'DualQuat':
        """Dual quaternion multiplication (Hamilton product + dual extension)."""
        # Real part: q1_real × q2_real
        rw = self.rw*other.rw - self.rx*other.rx - self.ry*other.ry - self.rz*other.rz
        rx = self.rw*other.rx + self.rx*other.rw + self.ry*other.rz - self.rz*other.ry
        ry = self.rw*other.ry - self.rx*other.rz + self.ry*other.rw + self.rz*other.rx
        rz = self.rw*other.rz + self.rx*other.ry - self.ry*other.rx + self.rz*other.rw
        # Dual part: q1_real × q2_dual + q1_dual × q2_real
        dw = self.rw*other.dw + self.rx*other.dx + self.ry*other.dy + self.rz*other.dz + \
             self.dw*other.rw - self.dx*other.rx - self.dy*other.ry - self.dz*other.rz
        dx = self.rw*other.dx - self.rx*other.dw + self.ry*other.dz - self.rz*other.dy + \
             self.dw*other.rx + self.dx*other.rw + self.dy*other.rz - self.dz*other.ry
        dy = self.rw*other.dy - self.rx*other.dz - self.ry*other.dw + self.rz*other.dx + \
             self.dw*other.ry - self.dx*other.rz + self.dy*other.rw + self.dz*other.rx
        dz = self.rw*other.dz + self.rx*other.dy - self.ry*other.dx - self.rz*other.dw + \
             self.dw*other.rz + self.dx*other.ry - self.dy*other.rx + self.dz*other.rw
        return DualQuat(rw, rx, ry, rz, dw, dx, dy, dz)

    def conjugate(self) -> 'DualQuat':
        """Quaternion conjugate: q* = (rw, -rx, -ry, -rz, dw, -dx, -dy, -dz)"""
        return DualQuat(self.rw, -self.rx, -self.ry, -self.rz,
                       self.dw, -self.dx, -self.dy, -self.dz)

    def to_dict(self) -> dict:
        return {
            "real": [self.rw, self.rx, self.ry, self.rz],
            "dual": [self.dw, self.dx, self.dy, self.dz],
            "chi": round(self.chi, 4),
            "mass": round(self.mass, 4),
            "compressible": self.is_compressible,
        }


# ── Semantic Mass → Dual Quaternion ─────────────────────────────────────

def mass_to_dualquat(mass: dict) -> DualQuat:
    """Convert semantic mass vector to dual quaternion.

    Mapping:
      Real (compressive): H (reasoning), I (invariant), C (closure), quality
      Dual (anti-compressive): R (residual/scar), L (latency), noise, cost
    """
    H = mass.get("H", 0)  # reasoning depth
    I = mass.get("I", 0)  # invariant pressure
    C = mass.get("C", 0)  # closure cost
    R = mass.get("R", 0)  # residual scar
    L = mass.get("L", 0)  # latency cost
    Q = mass.get("Q", 0)  # quality (derived)

    # Real part: compressive components
    rw = H * 0.4 + I * 0.3 + C * 0.2 + Q * 0.1
    rx = H * 0.3 - I * 0.1  # reasoning - invariant coupling
    ry = C * 0.2 + Q * 0.1  # closure + quality coupling
    rz = H * 0.1 + C * 0.1  # reasoning + closure coupling

    # Dual part: anti-compressive components
    dw = R * 0.5 + L * 0.3  # scar + latency
    dx = R * 0.3 - L * 0.1  # scar - latency coupling
    dy = L * 0.2 + R * 0.1  # latency + scar coupling
    dz = R * 0.1 + L * 0.1  # scar + latency coupling

    return DualQuat(rw, rx, ry, rz, dw, dx, dy, dz)


# ── Model Registry ──────────────────────────────────────────────────────

MODELS = {
    "deepseek-v4-pro": {
        "provider": "deepseek",
        "cost_per_1m": (0.27, 1.10),
        "mass": {"H": 0.9, "L": 0.1, "I": 0.8, "R": 0.1, "C": 0.9, "Q": 0.85},
        "strengths": ["math", "code", "reasoning"],
    },
    "deepseek-v4-flash": {
        "provider": "deepseek",
        "cost_per_1m": (0.07, 0.28),
        "mass": {"H": 0.3, "L": 0.05, "I": 0.4, "R": 0.3, "C": 0.4, "Q": 0.35},
        "strengths": ["speed", "iteration"],
    },
    "claude-opus": {
        "provider": "anthropic",
        "cost_per_1m": (15.0, 75.0),
        "mass": {"H": 0.95, "L": 0.3, "I": 0.9, "R": 0.05, "C": 0.85, "Q": 0.9},
        "strengths": ["formal_proof", "synthesis", "lean"],
    },
    "gpt-5.5": {
        "provider": "openai",
        "cost_per_1m": (10.0, 30.0),
        "mass": {"H": 0.7, "L": 0.2, "I": 0.7, "R": 0.15, "C": 0.8, "Q": 0.7},
        "strengths": ["breadth", "ops", "implementation"],
    },
    "cohere-command": {
        "provider": "cohere",
        "cost_per_1m": (2.50, 10.0),
        "mass": {"H": 0.6, "L": 0.15, "I": 0.8, "R": 0.1, "C": 0.7, "Q": 0.65},
        "strengths": ["structured_output", "tool_use", "rag"],
    },
    "gemini-pro": {
        "provider": "google",
        "cost_per_1m": (1.25, 5.0),
        "mass": {"H": 0.5, "L": 0.1, "I": 0.6, "R": 0.2, "C": 0.6, "Q": 0.55},
        "strengths": ["multimodal", "different_perspective"],
    },
    "glm-5.2": {
        "provider": "zhipu",
        "cost_per_1m": (1.0, 4.0),
        "mass": {"H": 0.4, "L": 0.15, "I": 0.5, "R": 0.25, "C": 0.5, "Q": 0.45},
        "strengths": ["architectural_diversity", "753B"],
    },
    "kimi-k2.6": {
        "provider": "moonshot",
        "cost_per_1m": (0.50, 2.0),
        "mass": {"H": 0.45, "L": 0.1, "I": 0.55, "R": 0.2, "C": 0.55, "Q": 0.5},
        "strengths": ["code", "different_training"],
    },
}

# Problem type → dimension weights (which H/I/R/L/C components matter most)
PROBLEM_WEIGHTS = {
    "lean_proof": {"H": 1.0, "I": 0.9, "C": 0.8, "R": 0.3, "L": 0.2, "Q": 0.9},
    "code_bug": {"H": 0.6, "I": 0.5, "C": 0.7, "R": 0.4, "L": 0.3, "Q": 0.7},
    "architecture": {"H": 0.8, "I": 0.7, "C": 0.9, "R": 0.2, "L": 0.2, "Q": 0.8},
    "math_novel": {"H": 1.0, "I": 0.8, "C": 0.7, "R": 0.3, "L": 0.2, "Q": 0.85},
    "infra_debug": {"H": 0.5, "I": 0.4, "C": 0.6, "R": 0.5, "L": 0.3, "Q": 0.6},
    "hard_wall": {"H": 1.0, "I": 1.0, "C": 1.0, "R": 0.5, "L": 0.5, "Q": 1.0},
}


# ── Model Selection ──────────────────────────────────────────────────────

def select_models(problem_type: str, threshold: float = 0.5,
                  max_models: int = 8) -> dict:
    """Select models using dual quaternion chiral ratio.

    Args:
        problem_type: One of lean_proof, code_bug, architecture, math_novel, infra_debug, hard_wall
        threshold: Minimum χ to keep a model (default 0.5)
        max_models: Maximum models to select

    Returns:
        dict with selected models, their χ values, and cost estimate
    """
    weights = PROBLEM_WEIGHTS.get(problem_type, PROBLEM_WEIGHTS["code_bug"])

    # Compute dual quaternion for each model, weighted by problem type
    model_dqs = []
    for name, config in MODELS.items():
        mass = config["mass"]
        # Apply problem-type weights
        weighted_mass = {k: mass.get(k, 0) * weights.get(k, 0) for k in
                        set(mass.keys()) | set(weights.keys())}
        dq = mass_to_dualquat(weighted_mass)
        model_dqs.append((name, dq, config))

    # Sort by χ (highest first = most compressive)
    model_dqs.sort(key=lambda x: x[1].chi, reverse=True)

    # Select models with χ > threshold
    selected = []
    for name, dq, config in model_dqs:
        if len(selected) >= max_models:
            break
        if dq.chi >= threshold:
            selected.append({
                "model": name,
                "provider": config["provider"],
                "chi": round(dq.chi, 4),
                "mass": round(dq.mass, 4),
                "cost_per_1m": config["cost_per_1m"],
                "strengths": config["strengths"],
                "dq": dq.to_dict(),
            })

    # Compute panel dual quaternion (sum of selected models)
    panel_dq = DualQuat()
    for _, dq, _ in [(n, d, c) for n, d, c in model_dqs if any(s["model"] == n for s in selected)]:
        panel_dq = DualQuat(
            panel_dq.rw + dq.rw, panel_dq.rx + dq.rx,
            panel_dq.ry + dq.ry, panel_dq.rz + dq.rz,
            panel_dq.dw + dq.dw, panel_dq.dx + dq.dx,
            panel_dq.dy + dq.dy, panel_dq.dz + dq.dz,
        )

    # Estimate cost (assume 1000 tokens in, 2000 tokens out per model)
    total_cost = 0
    for s in selected:
        cost_in, cost_out = s["cost_per_1m"]
        total_cost += (1000 / 1_000_000) * cost_in + (2000 / 1_000_000) * cost_out

    return {
        "problem_type": problem_type,
        "threshold": threshold,
        "selected": selected,
        "dropped": [n for n, _, _ in model_dqs if not any(s["model"] == n for s in selected)],
        "panel_chi": round(panel_dq.chi, 4),
        "panel_mass": round(panel_dq.mass, 4),
        "estimated_cost": round(total_cost, 4),
        "model_count": len(selected),
    }


# ── Z-Transform Compression ─────────────────────────────────────────────

def compress_history(mass_stream: list[float], order: int = 2) -> dict:
    """Compress semantic mass history into Z-domain recurrence.

    μ[k] = a₁·μ[k-1] + a₂·μ[k-2] + b₀·u[k]

    Returns coefficients + state vector.
    """
    if len(mass_stream) < order + 1:
        return {"error": "need more data points", "min_required": order + 1}

    # Simple least-squares fit for AR(2) model
    n = len(mass_stream)
    # Build matrix for [a1, a2, b0] estimation
    # μ[k] = a1*μ[k-1] + a2*μ[k-2] + b0*u[k]
    # where u[k] = 1 (constant input)
    sum_mu_k1 = sum(mass_stream[k] * mass_stream[k-1] for k in range(2, n))
    sum_mu_k2 = sum(mass_stream[k] * mass_stream[k-2] for k in range(2, n))
    sum_mu_k = sum(mass_stream[k] for k in range(2, n))
    sum_mu_k1_sq = sum(mass_stream[k-1]**2 for k in range(2, n))
    sum_mu_k2_sq = sum(mass_stream[k-2]**2 for k in range(2, n))
    sum_mu_k1_k2 = sum(mass_stream[k-1] * mass_stream[k-2] for k in range(2, n))
    sum_mu_k1_u = sum(mass_stream[k-1] for k in range(2, n))
    sum_mu_k2_u = sum(mass_stream[k-2] for k in range(2, n))

    # Simplified: use Yule-Walker equations for AR(2)
    n_obs = n - 2
    if n_obs < 3:
        return {"error": "insufficient data"}

    # Autocorrelation
    mean_mu = sum(mass_stream) / n
    r0 = sum((m - mean_mu)**2 for m in mass_stream) / n
    r1 = sum((mass_stream[k] - mean_mu) * (mass_stream[k-1] - mean_mu) for k in range(1, n)) / n
    r2 = sum((mass_stream[k] - mean_mu) * (mass_stream[k-2] - mean_mu) for k in range(2, n)) / n

    if r0 == 0:
        return {"a1": 0, "a2": 0, "b0": mean_mu, "r0": 0, "r1": 0, "r2": 0}

    # Yule-Walker: [r0 r1; r1 r0] [a1; a2] = [r1; r2]
    det = r0 * r0 - r1 * r1
    if abs(det) < 1e-10:
        a1 = r1 / r0 if r0 > 0 else 0
        a2 = 0
    else:
        a1 = (r0 * r1 - r1 * r2) / det
        a2 = (r0 * r2 - r1 * r1) / det

    b0 = mean_mu * (1 - a1 - a2)

    return {
        "a1": round(a1, 6),
        "a2": round(a2, 6),
        "b0": round(b0, 6),
        "r0": round(r0, 6),
        "r1": round(r1, 6),
        "r2": round(r2, 6),
        "mean": round(mean_mu, 6),
        "stream_length": n,
    }


# ── CLI ──────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Dual quaternion model selector")
    parser.add_argument("--threshold", type=float, default=0.5, help="χ threshold")
    parser.add_argument("--problem-type", default="lean_proof",
                       choices=list(PROBLEM_WEIGHTS.keys()))
    parser.add_argument("--max-models", type=int, default=8)
    parser.add_argument("--all-types", action="store_true", help="Show all problem types")
    args = parser.parse_args()

    if args.all_types:
        for ptype in PROBLEM_WEIGHTS:
            result = select_models(ptype, args.threshold, args.max_models)
            print(f"\n{'='*60}")
            print(f"  {ptype} (threshold={args.threshold})")
            print(f"{'='*60}")
            print(f"  Selected: {result['model_count']} models, χ_panel={result['panel_chi']}")
            print(f"  Cost: ${result['estimated_cost']:.4f}")
            for s in result["selected"]:
                print(f"    {s['model']:20s} χ={s['chi']:.3f}  ${s['cost_per_1m'][0]:.2f}/${s['cost_per_1m'][1]:.2f}")
            if result["dropped"]:
                print(f"  Dropped: {', '.join(result['dropped'])}")
    else:
        result = select_models(args.problem_type, args.threshold, args.max_models)
        print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()