from __future__ import annotations

from math import fsum
from typing import Iterable, Sequence

from .models import (
    AuditLogEntry,
    ConvergenceCertificate,
    LocalProblem,
    PlanDelta,
    SafetyPolicy,
    SharedVariables,
    SwarmSolution,
)
from .verification import verify_swarm_solution
class ADMMSolver:
    def __init__(self, rho: float = 1.0, max_iter: int = 50, epsilon: float = 1e-3, staleness_bound: int = 2):
        self.rho = rho
        self.max_iter = max_iter
        self.epsilon = epsilon
        self.staleness_bound = staleness_bound

    def _ordered_deltas(self, deltas: Iterable[PlanDelta]) -> list[PlanDelta]:
        return sorted(
            deltas,
            key=lambda delta: (delta.timestamp, delta.author, delta.contract_id, delta.sequence),
        )

    def _apply_deltas(
        self,
        robots: Sequence[LocalProblem],
        shared: SharedVariables,
        deltas: Iterable[PlanDelta],
    ) -> list[str]:
        applied: list[str] = []
        indexed = {robot.robot_id: robot for robot in robots}
        for delta in self._ordered_deltas(deltas):
            if shared.version - delta.sequence > self.staleness_bound:
                continue
            if delta.target_robot_id is None or delta.target_robot_id not in indexed:
                continue
            indexed[delta.target_robot_id].apply_updates(delta.updates)
            indexed[delta.target_robot_id].version = max(indexed[delta.target_robot_id].version, delta.sequence)
            applied.append(f"{delta.contract_id}:{delta.author}:{delta.sequence}")
        if applied:
            shared.delta_clock += len(applied)
        return applied

    def _project_for_separation(self, candidates: list[float], minimum: float) -> list[float]:
        if minimum <= 0 or len(candidates) <= 1:
            return candidates
        center = fsum(candidates) / len(candidates)
        half_span = minimum * (len(candidates) - 1) / 2.0
        return [center - half_span + index * minimum for index in range(len(candidates))]

    def solve(
        self,
        robots: Sequence[LocalProblem],
        safety_policy: SafetyPolicy,
        incoming_deltas: Iterable[PlanDelta] = (),
    ) -> SwarmSolution:
        working = [robot.model_copy(deep=True) for robot in robots]
        shared = SharedVariables(version=0, global_consensus=fsum(robot.target_pos for robot in working) / len(working))
        audit_log: list[AuditLogEntry] = []
        applied_deltas = self._apply_deltas(working, shared, incoming_deltas)

        final_primal = 0.0
        final_dual = 0.0
        converged = False

        for iteration in range(self.max_iter):
            old_consensus = shared.global_consensus
            candidate_positions: list[float] = []
            previous_positions = [robot.current_pos for robot in working]

            for robot in working:
                raw_position = (2.0 * robot.target_pos + self.rho * (shared.global_consensus - robot.dual_var)) / (2.0 + self.rho)
                step = max(-robot.max_step, min(robot.max_step, raw_position - robot.current_pos))
                candidate_positions.append(robot.current_pos + step)

            projected_positions = self._project_for_separation(candidate_positions, safety_policy.min_separation)

            for robot, new_position in zip(working, projected_positions):
                robot.current_pos = new_position
                if not safety_policy.verify(robot.initial_pos, robot.current_pos):
                    raise ValueError(f"Safety violation for {robot.robot_id}: travel bound exceeded")

            shared.global_consensus = fsum(projected_positions) / len(projected_positions)
            shared.version += 1

            final_primal = 0.0
            for robot, previous_position in zip(working, previous_positions):
                robot.dual_var += 0.5 * (robot.current_pos - shared.global_consensus)
                final_primal = max(final_primal, abs(robot.current_pos - previous_position))

            final_dual = abs(shared.global_consensus - old_consensus) * self.rho
            audit_log.append(
                AuditLogEntry(
                    iteration=iteration,
                    shared_version=shared.version,
                    consensus=shared.global_consensus,
                    primal_residual=final_primal,
                    dual_residual=final_dual,
                    applied_deltas=applied_deltas if iteration == 0 else [],
                )
            )

            if final_primal < self.epsilon and final_dual < self.epsilon:
                converged = True
                break

        certificate = ConvergenceCertificate(
            converged=converged,
            iterations=len(audit_log),
            epsilon=self.epsilon,
            final_primal_residual=final_primal,
            final_dual_residual=final_dual,
        )
        solution = SwarmSolution(
            consensus=shared.global_consensus,
            robots=working,
            shared=shared,
            audit_log=audit_log,
            certificate=certificate,
        )
        report = verify_swarm_solution(solution, safety_policy)
        if not report.passed:
            raise ValueError("; ".join(report.violations))
        return solution

    def run(
        self,
        robots: Sequence[LocalProblem],
        safety_policy: SafetyPolicy,
        incoming_deltas: Iterable[PlanDelta] = (),
    ) -> float:
        return self.solve(robots, safety_policy, incoming_deltas=incoming_deltas).consensus


def reference_scenario() -> SwarmSolution:
    robots = [
        LocalProblem(robot_id="drone-1", target_pos=0.0, current_pos=0.0),
        LocalProblem(robot_id="drone-2", target_pos=10.0, current_pos=10.0),
    ]
    policy = SafetyPolicy(max_travel_distance=15.0, min_separation=1.0, energy_budget=20.0)
    return ADMMSolver(rho=1.5, max_iter=40, epsilon=1e-4).solve(robots, policy)