import numpy as np
import numpy.typing as npt

from laddu.amplitudes import breit_wigner, common, kmatrix, ylm, zlm
from laddu.data import Dataset

class ParameterLike: ...

def parameter(name: str) -> ParameterLike: ...
def constant(value: float) -> ParameterLike: ...

class AmplitudeID:
    def real(self) -> Expression: ...
    def imag(self) -> Expression: ...
    def norm_sqr(self) -> Expression: ...
    def __add__(self, other: AmplitudeID | Expression | int) -> Expression: ...
    def __radd__(self, other: AmplitudeID | Expression | int) -> Expression: ...
    def __mul__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __rmul__(self, other: AmplitudeID | Expression) -> Expression: ...

class Expression:
    def real(self) -> Expression: ...
    def imag(self) -> Expression: ...
    def norm_sqr(self) -> Expression: ...
    def __add__(self, other: AmplitudeID | Expression | int) -> Expression: ...
    def __radd__(self, other: AmplitudeID | Expression | int) -> Expression: ...
    def __mul__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __rmul__(self, other: AmplitudeID | Expression) -> Expression: ...

class Amplitude: ...

class Manager:
    def __init__(self) -> None: ...
    def register(self, amplitude: Amplitude) -> AmplitudeID: ...
    def load(self, expression: Expression, dataset: Dataset) -> Evaluator: ...

class Evaluator:
    parameters: list[str]
    def activate(self, name: str | list[str]) -> None: ...
    def activate_all(self) -> None: ...
    def deactivate(self, name: str | list[str]) -> None: ...
    def deactivate_all(self) -> None: ...
    def isolate(self, name: str | list[str]) -> None: ...
    def evaluate(self, parameters: list[float] | npt.NDArray[np.float64]) -> npt.NDArray[np.complex128]: ...

__all__ = [
    "AmplitudeID",
    "Expression",
    "Amplitude",
    "Manager",
    "Evaluator",
    "ParameterLike",
    "parameter",
    "constant",
    "common",
    "ylm",
    "zlm",
    "breit_wigner",
    "kmatrix",
]
