from typing import Literal

import numpy as np
import numpy.typing as npt

from laddu.data import Dataset, Event


class VariableExpression:
    def __and__(self, rhs: VariableExpression) -> VariableExpression: ...
    def __or__(self, rhs: VariableExpression) -> VariableExpression: ...
    def __invert__(self) -> VariableExpression: ...


class Mass:
    def __init__(self, constituents: list[int]) -> None: ...
    def value(self, event: Event) -> float: ...
    def value_on(self, dataset: Dataset) -> npt.NDArray[np.float64]: ...


class CosTheta:
    def __init__(
        self,
        beam: int,
        recoil: list[int],
        daughter: list[int],
        resonance: list[int],
        frame: Literal[
            'Helicity',
            'HX',
            'HEL',
            'GottfriedJackson',
            'Gottfried Jackson',
            'GJ',
            'Gottfried-Jackson',
        ] = 'Helicity',
    ) -> None: ...
    def value(self, event: Event) -> float: ...
    def value_on(self, dataset: Dataset) -> npt.NDArray[np.float64]: ...
    def __eq__(self, value: float) -> VariableExpression: ...
    def __lt__(self, value: float) -> VariableExpression: ...
    def __gt__(self, value: float) -> VariableExpression: ...
    def __le__(self, value: float) -> VariableExpression: ...
    def __ge__(self, value: float) -> VariableExpression: ...


class Phi:
    def __init__(
        self,
        beam: int,
        recoil: list[int],
        daughter: list[int],
        resonance: list[int],
        frame: Literal[
            'Helicity',
            'HX',
            'HEL',
            'GottfriedJackson',
            'Gottfried Jackson',
            'GJ',
            'Gottfried-Jackson',
        ] = 'Helicity',
    ) -> None: ...
    def value(self, event: Event) -> float: ...
    def value_on(self, dataset: Dataset) -> npt.NDArray[np.float64]: ...
    def __eq__(self, value: float) -> VariableExpression: ...
    def __lt__(self, value: float) -> VariableExpression: ...
    def __gt__(self, value: float) -> VariableExpression: ...
    def __le__(self, value: float) -> VariableExpression: ...
    def __ge__(self, value: float) -> VariableExpression: ...


class Angles:
    costheta: CosTheta
    phi: Phi

    def __init__(
        self,
        beam: int,
        recoil: list[int],
        daughter: list[int],
        resonance: list[int],
        frame: Literal[
            'Helicity',
            'HX',
            'HEL',
            'GottfriedJackson',
            'Gottfried Jackson',
            'GJ',
            'Gottfried-Jackson',
        ] = 'Helicity',
    ) -> None: ...


class PolAngle:
    def __init__(
        self,
        beam: int,
        recoil: list[int],
        beam_polarization: int,
    ) -> None: ...
    def value(self, event: Event) -> float: ...
    def value_on(self, dataset: Dataset) -> npt.NDArray[np.float64]: ...
    def __eq__(self, value: float) -> VariableExpression: ...
    def __lt__(self, value: float) -> VariableExpression: ...
    def __gt__(self, value: float) -> VariableExpression: ...
    def __le__(self, value: float) -> VariableExpression: ...
    def __ge__(self, value: float) -> VariableExpression: ...


class PolMagnitude:
    def __init__(
        self,
        beam_polarization: int,
    ) -> None: ...
    def value(self, event: Event) -> float: ...
    def value_on(self, dataset: Dataset) -> npt.NDArray[np.float64]: ...
    def __eq__(self, value: float) -> VariableExpression: ...
    def __lt__(self, value: float) -> VariableExpression: ...
    def __gt__(self, value: float) -> VariableExpression: ...
    def __le__(self, value: float) -> VariableExpression: ...
    def __ge__(self, value: float) -> VariableExpression: ...


class Polarization:
    pol_magnitude: PolMagnitude
    pol_angle: PolAngle

    def __init__(self, beam: int, recoil: list[int], beam_polarization: int) -> None: ...


class Mandelstam:
    def __init__(
        self,
        p1: list[int],
        p2: list[int],
        p3: list[int],
        p4: list[int],
        channel: Literal['s', 't', 'u'],
    ) -> None: ...
    def value(self, event: Event) -> float: ...
    def value_on(self, dataset: Dataset) -> npt.NDArray[np.float64]: ...
    def __eq__(self, value: float) -> VariableExpression: ...
    def __lt__(self, value: float) -> VariableExpression: ...
    def __gt__(self, value: float) -> VariableExpression: ...
    def __le__(self, value: float) -> VariableExpression: ...
    def __ge__(self, value: float) -> VariableExpression: ...
