import sys
from primitives import LocalProblem, SharedVariables, SafetyPolicy

class ADMMSolver:
    def __init__(self, rho=1.0, max_iter=50, epsilon=1e-3):
        self.rho = rho
        self.max_iter = max_iter
        self.epsilon = epsilon
        
    def run(self, robots, safety_policy):
        z = 0.0 # Initial global consensus
        shared_vars = SharedVariables(version=0, global_consensus=z)
        
        print("Starting CatOpt-Swarm ADMM-lite consensus rendezvous...")
        
        for iteration in range(self.max_iter):
            old_z = shared_vars.global_consensus
            
            # 1. Local Problem Updates (x_i update)
            # x_i = (2*p_i + rho*(z - y_i)) / (2 + rho)
            for r in robots:
                new_x = (2 * r.target_pos + self.rho * (shared_vars.global_consensus - r.dual_var)) / (2 + self.rho)
                if not safety_policy.verify(r.target_pos, new_x):
                    print(f"Safety violation for {r.robot_id}: travel bound exceeded.")
                    sys.exit(1)
                r.current_pos = new_x
            
            # 2. Shared Variables Update (z update)
            # z = 1/N * sum(x_i + y_i)
            sum_x_y = sum([r.current_pos + r.dual_var for r in robots])
            new_z = sum_x_y / len(robots)
            shared_vars.global_consensus = new_z
            shared_vars.version += 1
            
            # 3. Dual Variables Update (y_i update)
            primal_residual = 0.0
            for r in robots:
                r.dual_var += (r.current_pos - shared_vars.global_consensus)
                primal_residual += abs(r.current_pos - shared_vars.global_consensus)
                
            dual_residual = abs(shared_vars.global_consensus - old_z) * self.rho
            
            print(f"Iter {iteration:02d}: Z={shared_vars.global_consensus:.3f} | PrimalRes={primal_residual:.4f} | DualRes={dual_residual:.4f}")
            
            if primal_residual < self.epsilon and dual_residual < self.epsilon:
                print("Convergence achieved!")
                return shared_vars.global_consensus
                
        print("Failed to converge within max iterations.")
        sys.exit(1)

if __name__ == "__main__":
    robots = [
        LocalProblem(robot_id="drone-1", target_pos=0.0),
        LocalProblem(robot_id="drone-2", target_pos=10.0)
    ]
    policy = SafetyPolicy(max_travel_distance=15.0)
    
    solver = ADMMSolver(rho=1.5)
    final_z = solver.run(robots, policy)
    print(f"Final agreed rendezvous point: {final_z:.3f}")
    assert abs(final_z - 5.0) < 0.01, "Rendezvous should be exactly in the middle!"