from collections.abc import Sequence

import numpy as np
import numpy.typing as npt

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

class ParameterLike: ...

def AmplitudeOne() -> Expression: ...
def AmplitudeZero() -> Expression: ...
def amplitude_sum(
    amplitudes: Sequence[AmplitudeID | Expression]
    | Sequence[AmplitudeID]
    | Sequence[Expression],
) -> Expression: ...
def amplitude_product(
    amplitudes: Sequence[AmplitudeID | Expression]
    | Sequence[AmplitudeID]
    | Sequence[Expression],
) -> Expression: ...
def constant(value: float) -> ParameterLike: ...
def parameter(name: str) -> ParameterLike: ...

class AmplitudeID:
    def real(self) -> Expression: ...
    def imag(self) -> Expression: ...
    def conj(self) -> Expression: ...
    def norm_sqr(self) -> Expression: ...
    def __add__(self, other: AmplitudeID | Expression | int) -> Expression: ...
    def __radd__(self, other: AmplitudeID | Expression | int) -> Expression: ...
    def __sub__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __rsub__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __mul__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __rmul__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __truediv__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __rtruediv__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __neg__(self) -> Expression: ...

class Expression:
    def real(self) -> Expression: ...
    def imag(self) -> Expression: ...
    def conj(self) -> Expression: ...
    def norm_sqr(self) -> Expression: ...
    def __add__(self, other: AmplitudeID | Expression | int) -> Expression: ...
    def __radd__(self, other: AmplitudeID | Expression | int) -> Expression: ...
    def __sub__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __rsub__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __mul__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __rmul__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __truediv__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __rtruediv__(self, other: AmplitudeID | Expression) -> Expression: ...
    def __neg__(self) -> Expression: ...

class Amplitude: ...

class Manager:
    parameters: list[str]

    def __init__(self) -> None: ...
    def register(self, amplitude: Amplitude) -> AmplitudeID: ...
    def model(self, expression: Expression | AmplitudeID) -> Model: ...

class Model:
    parameters: list[str]

    def __init__(self) -> None: ...
    def load(self, dataset: Dataset) -> Evaluator: ...
    def __getstate__(self) -> object: ...
    def __setstate__(self, state: object) -> None: ...

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],
        *,
        threads: int | None = None,
    ) -> npt.NDArray[np.complex128]: ...
    def evaluate_batch(
        self,
        parameters: list[float] | npt.NDArray[np.float64],
        indices: list[int] | npt.NDArray[np.integer],
        *,
        threads: int | None = None,
    ) -> npt.NDArray[np.complex128]: ...
    def evaluate_gradient(
        self,
        parameters: list[float] | npt.NDArray[np.float64],
        *,
        threads: int | None = None,
    ) -> npt.NDArray[np.complex128]: ...
    def evaluate_gradient_batch(
        self,
        parameters: list[float] | npt.NDArray[np.float64],
        indices: list[int] | npt.NDArray[np.integer],
        *,
        threads: int | None = None,
    ) -> npt.NDArray[np.complex128]: ...

def TestAmplitude(name: str, re: ParameterLike, im: ParameterLike) -> Amplitude: ...

__all__ = [
    'Amplitude',
    'AmplitudeID',
    'AmplitudeOne',
    'AmplitudeZero',
    'Evaluator',
    'Expression',
    'Manager',
    'Model',
    'ParameterLike',
    'TestAmplitude',
    'amplitude_sum',
    'breit_wigner',
    'common',
    'constant',
    'kmatrix',
    'parameter',
    'phase_space',
    'ylm',
    'zlm',
]
