import abc
import time
from collections.abc import Iterable
from typing import NamedTuple

from jaxtyping import Array, Float
from liblaf.grapes.logging import autolog

from liblaf.peach import tree
from liblaf.peach.constraints import Constraint
from liblaf.peach.optim.objective import Objective

from ._types import Callback, OptimizeSolution, Params, Result, State, Stats

type Vector = Float[Array, " N"]


class SetupResult[StateT: State, StatsT: Stats](NamedTuple):
    objective: Objective
    constraints: list[Constraint]
    state: StateT
    stats: StatsT


@tree.define
class Optimizer[StateT: State, StatsT: Stats](abc.ABC):
    from ._types import OptimizeSolution as Solution
    from ._types import State, Stats

    max_steps: int = tree.field(default=256, kw_only=True)
    jit: bool = tree.field(default=False, kw_only=True)
    timer: bool = tree.field(default=False, kw_only=True)

    def init(
        self,
        objective: Objective,
        params: Params,
        *,
        constraints: Iterable[Constraint] = (),
    ) -> SetupResult[StateT, StatsT]:
        params_flat: Vector
        objective, params_flat, constraints = objective.flatten(
            params, constraints=constraints
        )
        if self.jit:
            objective = objective.jit()
        if self.timer:
            objective = objective.timer()
        assert objective.structure is not None
        state = self.State(params_flat=params_flat, structure=objective.structure)
        return SetupResult(objective, constraints, state, self.Stats())  # pyright: ignore[reportReturnType]

    @abc.abstractmethod
    def step(
        self,
        objective: Objective,
        state: StateT,
        *,
        constraints: Iterable[Constraint] = (),
    ) -> StateT:
        raise NotImplementedError

    def update_stats(
        self,
        objective: Objective,  # noqa: ARG002
        state: StateT,  # noqa: ARG002
        stats: StatsT,
        *,
        constraints: Iterable[Constraint] = (),  # noqa: ARG002
    ) -> StatsT:
        return stats

    @abc.abstractmethod
    def terminate(
        self,
        objective: Objective,
        state: StateT,
        stats: StatsT,
        *,
        constraints: Iterable[Constraint] = (),
    ) -> tuple[bool, Result]:
        raise NotImplementedError

    def postprocess(
        self,
        objective: Objective,
        state: StateT,
        stats: StatsT,
        result: Result,
        *,
        constraints: Iterable[Constraint] = (),  # noqa: ARG002
    ) -> OptimizeSolution[StateT, StatsT]:
        stats.end_time = time.perf_counter()
        solution: OptimizeSolution[StateT, StatsT] = OptimizeSolution(
            result=result, state=state, stats=stats
        )
        objective.timer_finish()
        return solution

    def minimize(
        self,
        objective: Objective,
        params: Params,
        *,
        constraints: Iterable[Constraint] = (),
        callback: Callback[StateT, StatsT] | None = None,
    ) -> OptimizeSolution[StateT, StatsT]:
        state: StateT
        stats: StatsT
        objective, constraints, state, stats = self.init(
            objective, params, constraints=constraints
        )
        done: bool = False
        n_steps: int = 0
        result: Result = Result.UNKNOWN_ERROR
        while n_steps < self.max_steps and not done:
            state = self.step(objective, state, constraints=constraints)
            n_steps += 1
            stats.n_steps = n_steps
            stats = self.update_stats(objective, state, stats, constraints=constraints)
            if callback is not None:
                callback(state, stats)
            done, result = self.terminate(
                objective, state, stats, constraints=constraints
            )
        if not done:
            result = Result.MAX_STEPS_REACHED
        solution: OptimizeSolution[StateT, StatsT] = self.postprocess(
            objective, state, stats, result, constraints=constraints
        )
        return solution

    def _warn_ignore_constraints(
        self,
        constraints: Iterable[Constraint],
    ) -> None:
        _logging_hide = True
        if constraints:
            autolog.warning(
                "Constraints are not supported by %s. Ignoring them.",
                type(self).__name__,
            )
