from typing import Literal

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, expression: Expression, parameters: list[float] | npt.NDArray[np.float64]
    ) -> 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, expression: Expression, parameters: list[float] | npt.NDArray[np.float64]) -> float: ...
    def project(
        self, expression: Expression, parameters: list[float] | npt.NDArray[np.float64]
    ) -> npt.NDArray[np.float64]: ...
    def minimize(
        self,
        expression: Expression,
        p0: list[float],
        bounds: list[tuple[float | None, float | None]] | None = None,
        method: Literal["lbfgsb", "nelder_mead"] = "lbfgsb",
        max_steps: int = 4000,
        debug: bool = False,  # noqa: FBT001, FBT002
        verbose: bool = False,  # noqa: FBT001, FBT002
        **kwargs,
    ) -> Status: ...

class Status:
    x: npt.NDArray[np.float64]
    err: npt.NDArray[np.float64] | None
    x0: npt.NDArray[np.float64]
    fx: float
    cov: npt.NDArray[np.float64] | None
    hess: npt.NDArray[np.float64] | None
    message: str
    converged: bool
    bounds: list[Bound] | None
    n_f_evals: int
    n_g_evals: int

class Bound:
    lower: float
    upper: float

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