from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class IntrinsicFunction:
    name: str
    arguments: list[str]
    min_args: int
    max_args: int
    description: str


INTRINSIC_FUNCTIONS = {
    "sqrt": IntrinsicFunction(
        name="sqrt",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Square Root",
    ),
    "log": IntrinsicFunction(
        name="log",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Logarithm",
    ),
    "exp": IntrinsicFunction(
        name="exp",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Exponential",
    ),
    "sin": IntrinsicFunction(
        name="sin",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Sine",
    ),
    "cos": IntrinsicFunction(
        name="cos",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Cosine",
    ),
    "tan": IntrinsicFunction(
        name="tan",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Tangent",
    ),
    "cot": IntrinsicFunction(
        name="cot",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Cotangent",
    ),
    "sinc": IntrinsicFunction(
        name="sinc",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Sin(x)/x Function",
    ),
    "asin": IntrinsicFunction(
        name="asin",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Arc sine",
    ),
    "acos": IntrinsicFunction(
        name="acos",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Arc cosine",
    ),
    "atan": IntrinsicFunction(
        name="atan",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Arc tangent",
    ),
    "atan2": IntrinsicFunction(
        name="atan2",
        arguments=["y", "x"],
        min_args=2,
        max_args=2,
        description="Arc tangent of y/x",
    ),
    "sinh": IntrinsicFunction(
        name="sinh",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Hyperbolic sine",
    ),
    "cosh": IntrinsicFunction(
        name="cosh",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Hyperbolic cosine",
    ),
    "tanh": IntrinsicFunction(
        name="tanh",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Hyperbolic tangent",
    ),
    "coth": IntrinsicFunction(
        name="coth",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Hyperbolic cotangent",
    ),
    "asinh": IntrinsicFunction(
        name="asinh",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Hyperbolic arc sine",
    ),
    "acosh": IntrinsicFunction(
        name="acosh",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Hyperbolic arc cosine",
    ),
    "atanh": IntrinsicFunction(
        name="atanh",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Hyperbolic arc tangent",
    ),
    "acoth": IntrinsicFunction(
        name="acoth",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Hyperbolic arc cotangent",
    ),
    "abs": IntrinsicFunction(
        name="abs",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Absolute Value",
    ),
    "factorial": IntrinsicFunction(
        name="factorial",
        arguments=["n"],
        min_args=1,
        max_args=1,
        description="Factorial",
    ),
    "ran": IntrinsicFunction(
        name="ran",
        arguments=[],
        min_args=0,
        max_args=0,
        description="Random number between 0 and 1",
    ),
    "ran_gauss": IntrinsicFunction(
        name="ran_gauss",
        arguments=["sig_cut"],
        min_args=0,
        max_args=1,
        description="Gaussian distributed random number",
    ),
    "int": IntrinsicFunction(
        name="int",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Nearest integer with magnitude less than x",
    ),
    "nint": IntrinsicFunction(
        name="nint",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Nearest integer to x",
    ),
    "sign": IntrinsicFunction(
        name="sign",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="1 if x positive, -1 if negative, 0 if zero",
    ),
    "floor": IntrinsicFunction(
        name="floor",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Nearest integer less than x",
    ),
    "ceiling": IntrinsicFunction(
        name="ceiling",
        arguments=["x"],
        min_args=1,
        max_args=1,
        description="Nearest integer greater than x",
    ),
    "modulo": IntrinsicFunction(
        name="modulo",
        arguments=["a", "p"],
        min_args=2,
        max_args=2,
        description="a - floor(a/p) * p. Will be in range [0, p]",
    ),
    "mass_of": IntrinsicFunction(
        name="mass_of",
        arguments=["A"],
        min_args=1,
        max_args=1,
        description="Mass of particle A",
    ),
    "charge_of": IntrinsicFunction(
        name="charge_of",
        arguments=["A"],
        min_args=1,
        max_args=1,
        description="Charge, in units of the elementary charge, of particle A",
    ),
    "anomalous_moment_of": IntrinsicFunction(
        name="anomalous_moment_of",
        arguments=["A"],
        min_args=1,
        max_args=1,
        description="Anomalous magnetic moment of particle A",
    ),
    "species": IntrinsicFunction(
        name="species",
        arguments=["A"],
        min_args=1,
        max_args=1,
        description="Species ID of A",
    ),
}

RESERVED_NAMES = {
    "beam",
    "beginning",
    "calc_reference_orbit",
    "call",
    "combine_consecutive_elements",
    "debug_marker",
    "end_file",
    "expand_lattice",
    "merge_elements",
    "no_digested",
    "no_superimpose",
    "parameter",
    "parser_debug",
    "particle_start",
    "print",
    "redef",
    "remove_elements",
    "return",
    "root",
    "slice_lattice",
    "start_branch_at",
    "superimpose",
    "title",
    "use",
    "use_local_lat_file",
    "write_digested",
}

BUILTIN_CONSTANTS = {
    "twopi",
    "fourpi",
    "pi",
    "e",
    "e_log",
    "sqrt_2",
    "degrad",
    "degrees",
    "raddeg",
    "m_electron",
    "m_muon",
    "m_pion_0",
    "m_pion_charged",
    "m_proton",
    "m_deuteron",
    "m_neutron",
    "c_light",
    "r_e",
    "r_p",
    "e_charge",
    "h_planck",
    "h_bar_planck",
    "fine_struct_const",
    "anom_moment_electron",
    "anom_moment_proton",
    "anom_moment_neutron",
    "anom_moment_muon",
    "anom_moment_deuteron",
    "anom_moment_he3",
}

ALL_BUILTIN = frozenset(BUILTIN_CONSTANTS | RESERVED_NAMES | set(INTRINSIC_FUNCTIONS))
