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: ...

class Expression:
    parameters: list[str]
    free_parameters: list[str]
    fixed_parameters: list[str]
    n_free: int
    n_fixed: int
    n_parameters: int

    def load(self, dataset: Dataset) -> Evaluator: ...
    def real(self) -> Expression: ...
    def imag(self) -> Expression: ...
    def conj(self) -> Expression: ...
    def norm_sqr(self) -> Expression: ...
    def fix(self, name: str, value: float) -> Expression: ...
    def free(self, name: str) -> Expression: ...
    def rename_parameter(self, old: str, new: str) -> Expression: ...
    def rename_parameters(self, mapping: dict[str, str]) -> Expression: ...
    def __add__(self, other: Expression | int) -> Expression: ...
    def __radd__(self, other: Expression | int) -> Expression: ...
    def __sub__(self, other: Expression) -> Expression: ...
    def __rsub__(self, other: Expression) -> Expression: ...
    def __mul__(self, other: Expression) -> Expression: ...
    def __rmul__(self, other: Expression) -> Expression: ...
    def __truediv__(self, other: Expression) -> Expression: ...
    def __rtruediv__(self, other: Expression) -> Expression: ...
    def __neg__(self) -> Expression: ...

def One() -> Expression: ...
def Zero() -> Expression: ...
def expr_sum(amplitudes: Sequence[Expression]) -> Expression: ...
def expr_product(amplitudes: Sequence[Expression]) -> Expression: ...
def constant(name: str, value: float) -> ParameterLike: ...
def parameter(name: str) -> ParameterLike: ...

class Evaluator:
    parameters: list[str]
    free_parameters: list[str]
    fixed_parameters: list[str]
    n_free: int
    n_fixed: int
    n_parameters: int

    def fix(self, name: str, value: float) -> Evaluator: ...
    def free(self, name: str) -> Evaluator: ...
    def rename_parameter(self, old: str, new: str) -> Evaluator: ...
    def rename_parameters(self, mapping: dict[str, str]) -> Evaluator: ...
    def activate(self, name: str | list[str], *, strict: bool = True) -> None: ...
    def activate_all(self) -> None: ...
    def deactivate(self, name: str | list[str], *, strict: bool = True) -> None: ...
    def deactivate_all(self) -> None: ...
    def isolate(self, name: str | list[str], *, strict: bool = True) -> 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) -> Expression: ...

__all__ = [
    'Evaluator',
    'Expression',
    'One',
    'ParameterLike',
    'TestAmplitude',
    'Zero',
    'breit_wigner',
    'common',
    'constant',
    'expr_product',
    'expr_sum',
    'kmatrix',
    'parameter',
    'phase_space',
    'ylm',
    'zlm',
]
