from __future__ import annotations

from dataclasses import dataclass
import math

import numpy as np

from .models import LocalProblem, PlanAction, PlanDelta, SolverResult


@dataclass(frozen=True)
class _State:
    x: np.ndarray
    z: np.ndarray
    u: np.ndarray


class ADMMLiteSolver:
    def __init__(self, *, rho: float = 1.0, tolerance: float = 1e-6, max_iterations: int = 100) -> None:
        if rho <= 0:
            raise ValueError("rho must be positive")
        self.rho = rho
        self.tolerance = tolerance
        self.max_iterations = max_iterations

    def solve(self, problems: list[LocalProblem]) -> SolverResult:
        if not problems:
            raise ValueError("at least one local problem is required")

        horizons = {len(problem.preferred_dispatch) for problem in problems}
        if len(horizons) != 1:
            raise ValueError("all problems must share the same dispatch horizon")

        horizon = horizons.pop()
        lower = np.array([problem.lower_bounds or [-math.inf] * horizon for problem in problems], dtype=float)
        upper = np.array([problem.upper_bounds or [math.inf] * horizon for problem in problems], dtype=float)
        preferred = np.array([problem.preferred_dispatch for problem in problems], dtype=float)
        weights = np.array([problem.quadratic_weight for problem in problems], dtype=float).reshape(-1, 1)

        x = preferred.copy()
        z = preferred.mean(axis=0)
        u = np.zeros_like(x)
        primal_residual = float("inf")
        dual_residual = float("inf")

        for iteration in range(1, self.max_iterations + 1):
            previous_z = z.copy()
            x = np.clip((weights * preferred + self.rho * (z - u)) / (weights + self.rho), lower, upper)
            z = np.mean(x + u, axis=0)
            u = u + x - z

            primal_residual = float(np.max(np.linalg.norm(x - z, axis=1)))
            dual_residual = float(np.linalg.norm(z - previous_z) * math.sqrt(len(problems)))
            if primal_residual <= self.tolerance and dual_residual <= self.tolerance:
                break

        deltas: list[PlanDelta] = []
        for index, problem in enumerate(problems):
            actions = [
                PlanAction(resource_id=f"{problem.site_id}:{slot}", action="dispatch", value=float(value), metadata={"slot": slot})
                for slot, value in enumerate(x[index])
            ]
            deltas.append(
                PlanDelta(
                    site_id=problem.site_id,
                    adapter_id=f"adapter:{problem.site_id}",
                    revision=iteration,
                    round_id=problem.round_id,
                    parent_revision=max(0, iteration - 1),
                    actions=actions,
                    metadata={"solver": "admm-lite", "round_id": problem.round_id},
                )
            )

        return SolverResult(
            consensus=[float(value) for value in z],
            iterations=iteration,
            primal_residual=primal_residual,
            dual_residual=dual_residual,
            deltas=deltas,
            metadata={"rho": self.rho, "tolerance": self.tolerance, "max_iterations": self.max_iterations},
        )