from __future__ import annotations
from dataclasses import dataclass
import numpy as np
import charged_shells.units_and_constants as uc


Array = np.ndarray


@dataclass(kw_only=True)
class ModelParams:
    R: float | Array
    kappa: float | Array = None
    kappaR: float | Array = None
    c0: float | Array = None
    epsilon: float = 80
    temperature: float = 293

    def __post_init__(self):
        self.kappa, self.kappaR, self.c0 = screening_calculator(self.R, self.temperature, self.epsilon,
                                                                self.kappa, self.kappaR, self.c0)


def bjerrum(temp: float, epsilon: float) -> float:
    return uc.CONSTANTS.e0 ** 2 / (4 * np.pi * epsilon * uc.CONSTANTS.epsilon0 * uc.CONSTANTS.Boltzmann * temp)


def kappa_from_concentration(c0: float, temp: float, epsilon: float) -> float:
    return np.sqrt(8 * np.pi * bjerrum(temp, epsilon) * c0)


def concentration_from_kappa(kappa: float, temp: float, epsilon: float):
    # TODO: what is with the units of c0 calculated here?
    return kappa ** 2 / (8 * np.pi * bjerrum(temp, epsilon))


def screening_calculator(radius: float, temp: float, epsilon: float,
                         kappa: float = None, kappaR: float = None, c0: float = None) -> (float, float, float):
    if kappa is not None:
        return kappa, kappa * radius, concentration_from_kappa(kappa, temp, epsilon)
    elif kappaR is not None:
        return kappaR / radius, kappaR, concentration_from_kappa(kappaR / radius, temp, epsilon)
    elif c0 is not None:
        kappa = kappa_from_concentration(c0, temp, epsilon)
        return kappa, kappa * radius, c0
    raise ValueError('One of the arguments kappa, kappaR or c0 should be different from None.')