from typing import Literal

import numpy as np
import numpy.typing as npt

from laddu.data import Dataset, Event
from laddu.utils.vectors import Vec3, Vec4

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

class Topology:
    def __init__(
        self,
        k1: str | list[str],
        k2: str | list[str],
        k3: str | list[str],
        k4: str | list[str],
    ) -> None: ...
    @staticmethod
    def missing_k1(
        k2: str | list[str], k3: str | list[str], k4: str | list[str]
    ) -> Topology: ...
    @staticmethod
    def missing_k2(
        k1: str | list[str], k3: str | list[str], k4: str | list[str]
    ) -> Topology: ...
    @staticmethod
    def missing_k3(
        k1: str | list[str], k2: str | list[str], k4: str | list[str]
    ) -> Topology: ...
    @staticmethod
    def missing_k4(
        k1: str | list[str], k2: str | list[str], k3: str | list[str]
    ) -> Topology: ...
    def k1_names(self) -> list[str] | None: ...
    def k2_names(self) -> list[str] | None: ...
    def k3_names(self) -> list[str] | None: ...
    def k4_names(self) -> list[str] | None: ...
    def com_boost_vector(self, event: Event) -> Vec3: ...
    def k1(self, event: Event) -> Vec4: ...
    def k2(self, event: Event) -> Vec4: ...
    def k3(self, event: Event) -> Vec4: ...
    def k4(self, event: Event) -> Vec4: ...
    def k1_com(self, event: Event) -> Vec4: ...
    def k2_com(self, event: Event) -> Vec4: ...
    def k3_com(self, event: Event) -> Vec4: ...
    def k4_com(self, event: Event) -> Vec4: ...

class Mass:
    def __init__(self, constituents: str | list[str]) -> None: ...
    def value(self, event: Event) -> float: ...
    def value_on(self, dataset: Dataset) -> npt.NDArray[np.float64]: ...
    def __eq__(self, value: float) -> VariableExpression: ...  # type: ignore[override]
    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 CosTheta:
    def __init__(
        self,
        topology: Topology,
        daughter: str | list[str],
        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: ...  # type: ignore[override]
    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,
        topology: Topology,
        daughter: str | list[str],
        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: ...  # type: ignore[override]
    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,
        topology: Topology,
        daughter: str | list[str],
        frame: Literal[
            'Helicity',
            'HX',
            'HEL',
            'GottfriedJackson',
            'Gottfried Jackson',
            'GJ',
            'Gottfried-Jackson',
        ] = 'Helicity',
    ) -> None: ...

class PolAngle:
    def __init__(self, topology: Topology, pol_angle: str) -> None: ...
    def value(self, event: Event) -> float: ...
    def value_on(self, dataset: Dataset) -> npt.NDArray[np.float64]: ...
    def __eq__(self, value: float) -> VariableExpression: ...  # type: ignore[override]
    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, pol_magnitude: str) -> None: ...
    def value(self, event: Event) -> float: ...
    def value_on(self, dataset: Dataset) -> npt.NDArray[np.float64]: ...
    def __eq__(self, value: float) -> VariableExpression: ...  # type: ignore[override]
    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,
        topology: Topology,
        *,
        pol_magnitude: str,
        pol_angle: str,
    ) -> None: ...

class Mandelstam:
    def __init__(
        self, topology: Topology, channel: Literal['s', 't', 'u', '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: ...  # type: ignore[override]
    def __lt__(self, value: float) -> VariableExpression: ...
    def __gt__(self, value: float) -> VariableExpression: ...
    def __le__(self, value: float) -> VariableExpression: ...
    def __ge__(self, value: float) -> VariableExpression: ...
