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) -> Expression: ...
    def __mul__(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) -> Expression: ...
    def __mul__(self, other: AmplitudeID | Expression) -> Expression: ...

class Amplitude: ...

class Manager:
    def __init__(self) -> None: ...
    def register(self, amplitude: Amplitude) -> AmplitudeID: ...
    def load(self, 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], expression: Expression) -> npt.NDArray[np.complex128]: ...

class NLL:
    parameters: list[str]
    def __init__(self, manager: Manager, ds_data: Dataset, ds_mc: Dataset) -> None: ...
    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], expression: Expression) -> float: ...
    def project(self, parameters: list[float], expression: Expression) -> npt.NDArray[np.float64]: ...

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