"""Sinkhorn optimal transport for batched GPU assignment.

Ported directly from agents/scheduler-agent/app/sinkhorn.py — no changes needed.
Feature-flagged via config.enable_sinkhorn.
"""

from __future__ import annotations

import numpy as np


class SinkhornSolver:
    """Entropy-regularized optimal transport solver for GPU assignment."""

    def __init__(self, epsilon: float = 0.1, max_iterations: int = 100):
        self.epsilon = epsilon
        self.max_iterations = max_iterations

    def solve(self, costs: np.ndarray, supply: np.ndarray, demand: np.ndarray) -> np.ndarray:
        """Solve optimal transport problem using Sinkhorn algorithm.

        Args:
            costs: Cost matrix [n_jobs, n_gpus]
            supply: Job demands [n_jobs] (typically all 1.0)
            demand: GPU capacities [n_gpus]

        Returns:
            Assignment matrix [n_jobs, n_gpus] with soft assignments
        """
        n_jobs, n_gpus = costs.shape

        K = np.exp(-costs / self.epsilon)

        u = np.ones(n_jobs) / n_jobs
        v = np.ones(n_gpus) / n_gpus

        for _ in range(self.max_iterations):
            u_prev = u.copy()
            v = demand / (K.T @ u + 1e-10)
            u = supply / (K @ v + 1e-10)
            if np.allclose(u, u_prev, rtol=1e-6):
                break

        return np.diag(u) @ K @ np.diag(v)

    def assign_jobs_to_gpus(self, job_costs: list[list[float]],
                            gpu_capacities: list[int]) -> list[dict[int, float]]:
        """Assign jobs to GPUs using Sinkhorn algorithm.

        Args:
            job_costs: Cost vectors per job [job][gpu]
            gpu_capacities: Capacity of each GPU

        Returns:
            List of {gpu_idx: probability} per job
        """
        if not job_costs or not gpu_capacities:
            return []

        n_jobs = len(job_costs)
        costs = np.array(job_costs)
        supply = np.ones(n_jobs)
        demand = np.array(gpu_capacities, dtype=float)
        demand = demand / demand.sum() * n_jobs

        assignment = self.solve(costs, supply, demand)

        result = []
        for job_idx in range(n_jobs):
            job_assignment = {}
            for gpu_idx in range(len(gpu_capacities)):
                prob = assignment[job_idx, gpu_idx]
                if prob > 1e-6:
                    job_assignment[gpu_idx] = float(prob)
            result.append(job_assignment)

        return result
