#!/usr/bin/env python3
"""
Batch Scheduler QUBO Builder

Encodes job-to-backend assignment as a QUBO problem.
This is where QPU/SA becomes meaningful (vs 3-var router).

Variables:
    x_{i,b} ∈ {0,1} for job i, backend b ∈ {cpu, gpu, qpu}

Constraints (as quadratic penalties):
    - One-hot: each job assigned to exactly one backend
    - VRAM capacity: GPU jobs don't exceed VRAM limit
    - Thermal: penalize GPU when hot
    - Deadline: penalize late jobs

Objective:
    min Σ cost(job_i, backend_b) × x_{i,b} + penalties

Usage:
    from flowtorch.batch_scheduler_qubo import BatchSchedulerQUBO, Job

    jobs = [Job(id=0, vram_mb=100, deadline_ms=500, se=0.3), ...]
    scheduler = BatchSchedulerQUBO(jobs, gpu_vram_mb=4000, gpu_temp_c=55)
    qubo_dict = scheduler.build_qubo()

    # Solve with existing infrastructure
    from scripts.qubo_solve import solve_qubo_neal
    result = solve_qubo_neal(qubo_dict)
"""

import json
import hashlib
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Optional
from pathlib import Path


@dataclass
class Job:
    """A job to be scheduled."""
    id: int
    vram_mb: float              # GPU VRAM required (0 for CPU-only jobs)
    deadline_ms: float          # Deadline from now
    se: float                   # Symbolic entropy (0-1)
    expected_runtime_cpu_ms: float = 100.0
    expected_runtime_gpu_ms: float = 50.0
    expected_runtime_qpu_ms: float = 200.0  # Includes network latency
    priority: float = 1.0       # Higher = more important to meet deadline


@dataclass
class SchedulerConfig:
    """Penalty weights for QUBO formulation."""
    # Hard constraint penalties (must dominate objective to prevent violations)
    # Rule: λ >= 10 × N × max_objective_swing_per_job
    # With max objective ~10 per job, for 64 jobs: λ >= 10 × 64 × 10 = 6400
    lambda_onehot: float = 10000.0    # One-hot constraint (HARD)
    lambda_vram: float = 5000.0       # VRAM capacity (HARD with slack bits)

    # Soft penalties (part of objective, can trade off)
    lambda_thermal: float = 10.0      # Thermal cost
    lambda_deadline: float = 20.0     # Deadline miss
    lambda_qpu_latency: float = 5.0   # QPU network overhead

    # Thresholds
    gpu_temp_threshold: float = 70.0  # Start penalizing above this
    se_gpu_threshold: float = 0.4     # Below this, GPU is overkill
    se_qpu_threshold: float = 0.7     # Below this, QPU is overkill


BACKENDS = ['cpu', 'gpu', 'qpu']
N_BACKENDS = len(BACKENDS)

# VRAM quantization chunk size (MB)
VRAM_CHUNK_MB = 64


class BatchSchedulerQUBO:
    """
    Builds QUBO for batch job scheduling.

    Variable encoding:
        x[i * N_BACKENDS + b] = 1 iff job i assigned to backend b
        where b: 0=cpu, 1=gpu, 2=qpu
    """

    def __init__(
        self,
        jobs: List[Job],
        gpu_vram_mb: float = 4000.0,
        gpu_temp_c: float = 50.0,
        cpu_load: float = 0.5,
        qpu_available: bool = True,
        config: Optional[SchedulerConfig] = None
    ):
        self.jobs = jobs
        self.n_jobs = len(jobs)

        self.gpu_vram_mb = gpu_vram_mb
        self.gpu_temp_c = gpu_temp_c
        self.cpu_load = cpu_load
        self.qpu_available = qpu_available

        self.config = config or SchedulerConfig()

        # VRAM slack bits for proper capacity constraint
        # Quantize VRAM into chunks, add slack bits to represent unused capacity
        # NOTE: floor division means effective capacity is slightly less than gpu_vram_mb
        self.vram_cap_chunks = int(gpu_vram_mb // VRAM_CHUNK_MB)
        self.effective_vram_mb = self.vram_cap_chunks * VRAM_CHUNK_MB
        self.n_vram_slack = max(1, (self.vram_cap_chunks + 1).bit_length())  # ceil(log2(cap+1))

        # Variable layout:
        # [0, n_jobs * 3): job-backend assignments x_{i,b}
        # [n_jobs * 3, n_jobs * 3 + n_vram_slack): VRAM slack bits s_k
        self.n_assignment_vars = self.n_jobs * N_BACKENDS
        self.n_vars = self.n_assignment_vars + self.n_vram_slack

        self._vram_slack_start = self.n_assignment_vars

        # QUBO coefficients
        self.linear: Dict[int, float] = {}
        self.quadratic: Dict[Tuple[int, int], float] = {}

    def _var_idx(self, job_idx: int, backend_idx: int) -> int:
        """Get variable index for job i, backend b."""
        return job_idx * N_BACKENDS + backend_idx

    def _vram_slack_idx(self, k: int) -> int:
        """Get variable index for VRAM slack bit k."""
        return self._vram_slack_start + k

    def _add_linear(self, idx: int, value: float):
        """Add to linear coefficient."""
        self.linear[idx] = self.linear.get(idx, 0.0) + value

    def _add_quadratic(self, i: int, j: int, value: float):
        """Add to quadratic coefficient (upper triangle only)."""
        if i > j:
            i, j = j, i
        if i == j:
            self._add_linear(i, value)
        else:
            key = (i, j)
            self.quadratic[key] = self.quadratic.get(key, 0.0) + value

    def _add_onehot_constraints(self):
        """
        Each job must be assigned to exactly one backend.

        Penalty: λ(1 - Σ_b x_{i,b})² = λ(1 - 2Σx + ΣΣx_i x_j)

        For each job i:
            Linear: -2λ for each x_{i,b}
            Quadratic: +2λ for each pair (x_{i,b1}, x_{i,b2})
            Constant: +λ (ignored in QUBO)
        """
        λ = self.config.lambda_onehot

        for job_idx in range(self.n_jobs):
            # Linear terms: -2λ for each backend option
            for b in range(N_BACKENDS):
                idx = self._var_idx(job_idx, b)
                self._add_linear(idx, -2 * λ)

            # Quadratic terms: +2λ for each pair within same job
            for b1 in range(N_BACKENDS):
                for b2 in range(b1 + 1, N_BACKENDS):
                    idx1 = self._var_idx(job_idx, b1)
                    idx2 = self._var_idx(job_idx, b2)
                    self._add_quadratic(idx1, idx2, 2 * λ)

            # Self-interaction (x² = x for binary): +λ per variable
            for b in range(N_BACKENDS):
                idx = self._var_idx(job_idx, b)
                self._add_linear(idx, λ)

    def _add_vram_constraint(self):
        """
        GPU VRAM capacity constraint using slack bits.

        Proper QUBO formulation:
            Σᵢ wᵢ·x_{i,gpu} + Σₖ 2^k·sₖ = C

        where:
            wᵢ = ceil(vram_i / VRAM_CHUNK_MB) in chunks
            C = floor(VRAM_max / VRAM_CHUNK_MB) capacity in chunks
            sₖ = slack bits representing unused capacity

        Penalty: λ(Σwᵢx_{i,gpu} + Σ2^k·sₖ - C)²

        This expands to pure quadratic form that the solver can optimize.
        Slack bits absorb unused capacity, enforcing sum ≤ C.
        """
        λ = self.config.lambda_vram
        gpu_idx = BACKENDS.index('gpu')
        C = self.vram_cap_chunks

        # Build list of (var_index, weight) for GPU assignments
        gpu_terms: List[Tuple[int, float]] = []
        for i, job in enumerate(self.jobs):
            w_i = max(1, int((job.vram_mb + VRAM_CHUNK_MB - 1) // VRAM_CHUNK_MB))  # ceil
            gpu_terms.append((self._var_idx(i, gpu_idx), float(w_i)))

        # Slack bits: s_k has weight 2^k
        slack_terms: List[Tuple[int, float]] = []
        for k in range(self.n_vram_slack):
            slack_terms.append((self._vram_slack_idx(k), float(2 ** k)))

        # Combine all terms
        all_terms = gpu_terms + slack_terms
        const = -float(C)

        # Expand (Σ aᵢxᵢ + const)² into QUBO form
        # = const² + 2·const·Σaᵢxᵢ + Σaᵢ²xᵢ² + 2·Σ_{i<j} aᵢaⱼxᵢxⱼ
        # Since xᵢ² = xᵢ for binary vars:
        # Linear: aᵢ² + 2·const·aᵢ
        # Quadratic: 2·aᵢ·aⱼ for i<j

        for idx, a_i in all_terms:
            linear_coef = λ * (a_i * a_i + 2.0 * const * a_i)
            self._add_linear(idx, linear_coef)

        for idx_a in range(len(all_terms)):
            i, a_i = all_terms[idx_a]
            for idx_b in range(idx_a + 1, len(all_terms)):
                j, a_j = all_terms[idx_b]
                self._add_quadratic(i, j, λ * 2.0 * a_i * a_j)

    def _add_thermal_penalty(self):
        """
        Penalize GPU usage when GPU is hot.

        Linear penalty on each x_{i,gpu} proportional to temperature.
        """
        λ = self.config.lambda_thermal
        threshold = self.config.gpu_temp_threshold
        gpu_idx = BACKENDS.index('gpu')

        if self.gpu_temp_c > threshold:
            # Thermal stress factor (0 at threshold, 1 at threshold+30)
            stress = min(1.0, (self.gpu_temp_c - threshold) / 30.0)

            for job_idx in range(self.n_jobs):
                idx = self._var_idx(job_idx, gpu_idx)
                self._add_linear(idx, λ * stress)

    def _add_deadline_penalties(self):
        """
        Penalize backend choices that would miss deadline.

        QUBO-compatible: precompute lateness as constant per (job, backend).
        Since deadline_i is fixed at solve time, max(0, runtime - deadline) is a
        constant that we multiply by x_{i,b}.

        Penalty: λ × Σᵢ Σ_b late_cost[i][b] × x_{i,b}

        This is linear in decision variables (QUBO-friendly).
        """
        λ = self.config.lambda_deadline

        for job_idx, job in enumerate(self.jobs):
            runtimes = [
                job.expected_runtime_cpu_ms,
                job.expected_runtime_gpu_ms,
                job.expected_runtime_qpu_ms
            ]

            for b, runtime in enumerate(runtimes):
                # Precompute lateness as constant (0 if on time, positive if late)
                lateness = max(0.0, runtime - job.deadline_ms) / max(job.deadline_ms, 1.0)
                if lateness > 0:
                    idx = self._var_idx(job_idx, b)
                    penalty = λ * lateness * job.priority
                    self._add_linear(idx, penalty)

    def _add_se_mismatch_penalty(self):
        """
        Penalize using heavy backends for light tasks (and vice versa).

        - Low Se + GPU = overkill
        - Low Se + QPU = very overkill
        - High Se + CPU = underpowered
        """
        gpu_idx = BACKENDS.index('gpu')
        qpu_idx = BACKENDS.index('qpu')
        cpu_idx = BACKENDS.index('cpu')

        for job_idx, job in enumerate(self.jobs):
            se = job.se

            # Low Se tasks shouldn't use GPU/QPU
            if se < self.config.se_gpu_threshold:
                overkill = self.config.se_gpu_threshold - se
                self._add_linear(self._var_idx(job_idx, gpu_idx), 5.0 * overkill)
                self._add_linear(self._var_idx(job_idx, qpu_idx), 10.0 * overkill)

            # High Se tasks should avoid CPU
            if se > self.config.se_qpu_threshold:
                underpowered = se - self.config.se_qpu_threshold
                self._add_linear(self._var_idx(job_idx, cpu_idx), 5.0 * underpowered)

    def _add_runtime_objective(self):
        """
        Base objective: minimize total expected runtime.

        Normalized so that CPU baseline = 1.0.
        """
        for job_idx, job in enumerate(self.jobs):
            # Normalize by CPU runtime
            cpu_baseline = job.expected_runtime_cpu_ms

            costs = [
                1.0,  # CPU (baseline)
                job.expected_runtime_gpu_ms / cpu_baseline,
                job.expected_runtime_qpu_ms / cpu_baseline
            ]

            for b, cost in enumerate(costs):
                idx = self._var_idx(job_idx, b)
                self._add_linear(idx, cost)

    def _add_qpu_unavailable_penalty(self):
        """If QPU is unavailable, make it infinitely expensive."""
        if not self.qpu_available:
            qpu_idx = BACKENDS.index('qpu')
            for job_idx in range(self.n_jobs):
                idx = self._var_idx(job_idx, qpu_idx)
                self._add_linear(idx, 1000.0)  # Effectively forbidden

    def build_qubo(self) -> Dict:
        """
        Build the complete QUBO dictionary.

        Returns:
            Dict with 'n', 'linear', 'quadratic' compatible with qubo_solve.py
        """
        # Reset coefficients
        self.linear = {}
        self.quadratic = {}

        # Build QUBO terms
        self._add_onehot_constraints()
        self._add_vram_constraint()
        self._add_thermal_penalty()
        self._add_deadline_penalties()
        self._add_se_mismatch_penalty()
        self._add_runtime_objective()
        self._add_qpu_unavailable_penalty()

        # Convert to qubo_solve.py format
        linear_dict = {str(k): v for k, v in self.linear.items()}
        quadratic_dict = {f"{i},{j}": v for (i, j), v in self.quadratic.items()}

        return {
            "n": self.n_vars,
            "linear": linear_dict,
            "quadratic": quadratic_dict,
            "metadata": {
                "n_jobs": self.n_jobs,
                "n_backends": N_BACKENDS,
                "backends": BACKENDS,
                "n_assignment_vars": self.n_assignment_vars,
                "n_vram_slack": self.n_vram_slack,
                "vram_cap_chunks": self.vram_cap_chunks,
                "vram_chunk_mb": VRAM_CHUNK_MB,
                "gpu_vram_mb_requested": self.gpu_vram_mb,
                "gpu_vram_mb_effective": self.effective_vram_mb,
                "gpu_temp_c": self.gpu_temp_c,
                "qpu_available": self.qpu_available,
                "penalty_weights": {
                    "lambda_onehot": self.config.lambda_onehot,
                    "lambda_vram": self.config.lambda_vram,
                    "lambda_thermal": self.config.lambda_thermal,
                    "lambda_deadline": self.config.lambda_deadline,
                },
            }
        }

    def decode_solution(self, bits: List[int]) -> Dict:
        """
        Decode QUBO solution to job assignments with constraint verification.

        Args:
            bits: Binary solution vector of length n_vars

        Returns:
            Dict with 'assignments' list, 'summary' stats, and 'verification'
        """
        if len(bits) != self.n_vars:
            raise ValueError(f"Expected {self.n_vars} bits, got {len(bits)}")

        assignments = []
        backend_counts = {b: 0 for b in BACKENDS}
        total_vram = 0.0
        total_vram_chunks = 0
        onehot_violations = 0

        for job_idx, job in enumerate(self.jobs):
            assigned_backend = None
            backends_selected = 0

            for b in range(N_BACKENDS):
                idx = self._var_idx(job_idx, b)
                if bits[idx] == 1:
                    backends_selected += 1
                    if assigned_backend is not None:
                        # Constraint violation: multiple backends
                        assigned_backend = f"INVALID({assigned_backend}+{BACKENDS[b]})"
                    else:
                        assigned_backend = BACKENDS[b]
                        backend_counts[BACKENDS[b]] += 1
                        if BACKENDS[b] == 'gpu':
                            total_vram += job.vram_mb
                            total_vram_chunks += max(1, int((job.vram_mb + VRAM_CHUNK_MB - 1) // VRAM_CHUNK_MB))

            if assigned_backend is None:
                assigned_backend = "UNASSIGNED"
                onehot_violations += 1
            elif backends_selected > 1:
                onehot_violations += 1

            assignments.append({
                "job_id": job.id,
                "backend": assigned_backend,
                "se": job.se,
                "vram_mb": job.vram_mb,
                "deadline_ms": job.deadline_ms
            })

        # Decode slack bits to get unused capacity
        slack_value = 0
        for k in range(self.n_vram_slack):
            if bits[self._vram_slack_idx(k)] == 1:
                slack_value += 2 ** k

        # Verify VRAM constraint: used_chunks + slack should equal capacity
        vram_constraint_lhs = total_vram_chunks + slack_value
        vram_constraint_satisfied = (vram_constraint_lhs == self.vram_cap_chunks)
        vram_over = total_vram_chunks > self.vram_cap_chunks

        # Check for invalid slack value (solver got weird)
        slack_max_representable = (1 << self.n_vram_slack) - 1
        slack_value_invalid = slack_value > self.vram_cap_chunks
        lhs_exceeds_cap = vram_constraint_lhs > self.vram_cap_chunks

        # Count violations: any failure of the equality constraint is a violation
        # (not just "over capacity" - slack mismatch is also a violation)
        vram_violation = 0 if vram_constraint_satisfied else 1

        return {
            "assignments": assignments,
            "summary": {
                "backend_counts": backend_counts,
                "total_gpu_vram_mb": total_vram,
                "vram_utilization": total_vram / self.effective_vram_mb if self.effective_vram_mb > 0 else 0,
                "constraint_violations": onehot_violations + vram_violation
            },
            "verification": {
                "onehot_violations": onehot_violations,
                "vram_used_chunks": total_vram_chunks,
                "vram_slack_value": slack_value,
                "vram_slack_max_representable": slack_max_representable,
                "vram_slack_value_invalid": slack_value_invalid,
                "vram_cap_chunks": self.vram_cap_chunks,
                "vram_constraint_lhs": vram_constraint_lhs,
                "vram_constraint_lhs_exceeds_cap": lhs_exceeds_cap,
                "vram_constraint_satisfied": vram_constraint_satisfied,
                "vram_over_capacity": vram_over,
                "vram_used_mb": total_vram,
                "vram_cap_mb_requested": self.gpu_vram_mb,
                "vram_cap_mb_effective": self.effective_vram_mb,
            }
        }


def generate_synthetic_jobs(n_jobs: int, seed: int = 42) -> List[Job]:
    """Generate synthetic jobs for testing."""
    import random
    random.seed(seed)

    jobs = []
    for i in range(n_jobs):
        # Mix of job types
        job_type = random.choice(['light', 'medium', 'heavy', 'urgent'])

        if job_type == 'light':
            se = random.uniform(0.05, 0.3)
            vram = random.uniform(0, 100)
            deadline = random.uniform(500, 2000)
            cpu_ms, gpu_ms, qpu_ms = 50, 30, 250
        elif job_type == 'medium':
            se = random.uniform(0.3, 0.6)
            vram = random.uniform(100, 500)
            deadline = random.uniform(200, 1000)
            cpu_ms, gpu_ms, qpu_ms = 200, 80, 300
        elif job_type == 'heavy':
            se = random.uniform(0.6, 0.95)
            vram = random.uniform(500, 2000)
            deadline = random.uniform(100, 500)
            cpu_ms, gpu_ms, qpu_ms = 500, 100, 350
        else:  # urgent
            se = random.uniform(0.4, 0.8)
            vram = random.uniform(200, 800)
            deadline = random.uniform(50, 200)
            cpu_ms, gpu_ms, qpu_ms = 150, 50, 400

        jobs.append(Job(
            id=i,
            vram_mb=vram,
            deadline_ms=deadline,
            se=se,
            expected_runtime_cpu_ms=cpu_ms * random.uniform(0.8, 1.2),
            expected_runtime_gpu_ms=gpu_ms * random.uniform(0.8, 1.2),
            expected_runtime_qpu_ms=qpu_ms * random.uniform(0.8, 1.2),
            priority=random.uniform(0.5, 2.0) if job_type == 'urgent' else 1.0
        ))

    return jobs


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Batch Scheduler QUBO Builder")
    parser.add_argument('--jobs', '-n', type=int, default=16, help='Number of jobs')
    parser.add_argument('--seed', type=int, default=42, help='Random seed')
    parser.add_argument('--gpu-vram', type=float, default=4000, help='GPU VRAM in MB')
    parser.add_argument('--gpu-temp', type=float, default=55, help='GPU temperature in C')
    parser.add_argument('--output', '-o', type=str, default=None, help='Output QUBO JSON file')
    parser.add_argument('--solve', action='store_true', help='Solve immediately with SA')

    args = parser.parse_args()

    print(f"Generating {args.jobs} synthetic jobs...")
    jobs = generate_synthetic_jobs(args.jobs, args.seed)

    print(f"Building QUBO ({args.jobs} jobs × 3 backends = {args.jobs * 3} variables)...")
    scheduler = BatchSchedulerQUBO(
        jobs=jobs,
        gpu_vram_mb=args.gpu_vram,
        gpu_temp_c=args.gpu_temp,
        qpu_available=False  # Simulated only for now
    )

    qubo = scheduler.build_qubo()

    print(f"  Variables: {qubo['n']}")
    print(f"  Linear terms: {len(qubo['linear'])}")
    print(f"  Quadratic terms: {len(qubo['quadratic'])}")

    # Save QUBO
    if args.output:
        output_path = Path(args.output)
    else:
        output_path = Path(__file__).parent.parent / "artifacts" / "qubo" / f"batch_scheduler_{args.jobs}jobs.json"

    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, 'w') as f:
        json.dump(qubo, f, indent=2)
    print(f"  QUBO saved to: {output_path}")

    # Solve if requested
    if args.solve:
        print("\nSolving with SimulatedAnnealing...")
        import sys
        sys.path.insert(0, str(Path(__file__).parent.parent / "scripts"))

        try:
            from qubo_solve import solve_qubo_neal
            result = solve_qubo_neal(qubo, seed=args.seed, sweeps=1000, num_reads=10)

            bits = result['bits']  # Already a list
            decoded = scheduler.decode_solution(bits)

            print(f"\nSolution (energy={result['energy']:.2f}):")
            print(f"  Backend distribution: {decoded['summary']['backend_counts']}")
            print(f"  GPU VRAM usage: {decoded['summary']['total_gpu_vram_mb']:.0f}/{args.gpu_vram} MB")
            print(f"  Constraint violations: {decoded['summary']['constraint_violations']}")

            # Verification results
            v = decoded['verification']
            print(f"\nVerification:")
            print(f"  One-hot violations: {v['onehot_violations']}")
            print(f"  VRAM used: {v['vram_used_chunks']} chunks ({v['vram_used_mb']:.0f} MB)")
            print(f"  VRAM capacity: {v['vram_cap_chunks']} chunks ({v['vram_cap_mb']:.0f} MB)")
            print(f"  VRAM slack: {v['vram_slack_value']} chunks")
            print(f"  VRAM constraint LHS: {v['vram_constraint_lhs']} (should = {v['vram_cap_chunks']})")
            print(f"  VRAM constraint satisfied: {v['vram_constraint_satisfied']}")
            print(f"  VRAM over capacity: {v['vram_over_capacity']}")

            print("\nAssignments:")
            for a in decoded['assignments'][:10]:  # Show first 10
                print(f"  Job {a['job_id']:2d}: {a['backend']:4s} (Se={a['se']:.2f}, VRAM={a['vram_mb']:.0f}MB)")
            if len(decoded['assignments']) > 10:
                print(f"  ... and {len(decoded['assignments']) - 10} more")

        except Exception as e:
            print(f"  Solve failed: {e}")
            print("  (Install dwave-samplers: pip install dwave-samplers dimod)")
