# SPDX-FileCopyrightText: 2024 The eminus developers
# SPDX-License-Identifier: Apache-2.0
from collections.abc import Callable
from dataclasses import dataclass
from functools import cached_property
from typing import Any, TypeAlias

from numpy import floating
from numpy.typing import NDArray

_Float: TypeAlias = floating
_ArrayReal: TypeAlias = NDArray[_Float]

@dataclass
class Coefficients:
    theta: _ArrayReal
    omega: float
    e1: float
    e2: float
    e3: float
    e4: float
    e5: float
    d1: float
    d2: float
    d3: float
    d4: float
    d5: float
    c1: float
    c2: float
    c3: float
    b1: float
    b2: float
    b3: float
    b4: float
    a1: float = ...
    a2: float = ...
    a3: float = ...
    a4: float = ...
    a5: float = ...
    a6: float = ...
    @cached_property
    def a0(self) -> float: ...
    @cached_property
    def b5(self) -> float: ...
    @cached_property
    def a(self) -> _ArrayReal: ...
    @property
    def dadtheta(self) -> _ArrayReal: ...
    @cached_property
    def b(self) -> _ArrayReal: ...
    @property
    def dbdtheta(self) -> _ArrayReal: ...
    @cached_property
    def c(self) -> _ArrayReal: ...
    @property
    def dcdtheta(self) -> _ArrayReal: ...
    @cached_property
    def d(self) -> _ArrayReal: ...
    @property
    def dddtheta(self) -> _ArrayReal: ...
    @cached_property
    def e(self) -> _ArrayReal: ...
    @cached_property
    def dedtheta(self) -> _ArrayReal: ...

@dataclass
class Zeta0Coeffs(Coefficients):
    omega: float = ...
    b1: float = ...
    b2: float = ...
    b3: float = ...
    b4: float = ...
    c1: float = ...
    c2: float = ...
    c3: float = ...
    d1: float = ...
    d2: float = ...
    d3: float = ...
    d4: float = ...
    d5: float = ...
    e1: float = ...
    e2: float = ...
    e3: float = ...
    e4: float = ...
    e5: float = ...

@dataclass
class Zeta1Coeffs(Coefficients):
    omega: float = ...
    b1: float = ...
    b2: float = ...
    b3: float = ...
    b4: float = ...
    c1: float = ...
    c2: float = ...
    c3: float = ...
    d1: float = ...
    d2: float = ...
    d3: float = ...
    d4: float = ...
    d5: float = ...
    e1: float = ...
    e2: float = ...
    e3: float = ...
    e4: float = ...
    e5: float = ...

@dataclass
class PhiParams:
    g1: float = ...
    g2: float = ...
    g3: float = ...
    lambda1: float = ...
    lambda2: float = ...

def lda_xc_ksdt(
    n: _ArrayReal,
    T: float = ...,
    zeta0_coeffs: Callable[..., Coefficients] = ...,
    zeta1_coeffs: Callable[..., Coefficients] = ...,
    phi_params: Callable[..., Any] = ...,
    **kwargs: Any,
) -> tuple[_ArrayReal, _ArrayReal, None]: ...
def lda_xc_ksdt_spin(
    n: _ArrayReal,
    zeta: _ArrayReal,
    T: float = ...,
    zeta0_coeffs: Callable[..., Coefficients] = ...,
    zeta1_coeffs: Callable[..., Coefficients] = ...,
    phi_params: Callable[..., Any] = ...,
    **kwargs: Any,
) -> tuple[_ArrayReal, _ArrayReal, None]: ...
