from abc import abstractmethod
from typing import Any

import equinox as eqx
from jaxtyping import Array, Float


class AbstractPriorProjector(eqx.Module, strict=True):
    """
    Abstract class for prior projectors.
    """

    @abstractmethod
    def initialize(
        self,
        init_state: Any | None = None,
    ) -> Any:
        raise NotImplementedError(
            "Abstract method not implemented. "
            "Please implement the initialize method in the subclass."
        )

    @abstractmethod
    def __call__(
        self,
        ref_positions: Float[Array, "n_atoms 3"] | Float[Array, "n_walkers n_atoms 3"],
        state: Any,
        bias_constant: Float,
    ) -> tuple[Float[Array, "n_atoms 3"] | Float[Array, "n_walkers n_atoms 3"], Any]:
        raise NotImplementedError("Abstract method not implemented.")


class AbstractEnsemblePriorProjector(eqx.Module, strict=True):
    """
    Abstract class for ensemble prior projectors.
    """

    projectors: eqx.AbstractVar[list[AbstractPriorProjector]]

    def initialize(
        self,
        init_states: list[Any] | None = None,
    ) -> list[Any]:
        if init_states is None:
            init_states = [None] * len(self.projectors)
        states = []
        for i, projector in enumerate(self.projectors):
            states.append(projector.initialize(init_states[i]))
        return states

    @abstractmethod
    def __call__(
        self,
        ref_positions: Float[Array, "n_walkers n_atoms 3"],
        states: list[Any],
        bias_constant: Float,
    ) -> tuple[Float[Array, "n_walkers n_atoms 3"], list[Any]]:
        raise NotImplementedError("Abstract method not implemented.")
