from __future__ import annotations

import numpy as np
from scipy import special as special

from ..utils import kB


class Weight:
    type = 'Line shape'
    name = 'Weighting factor'
    equation = r'A * [0.5 \pm 0.5 erf[(x-T_{0})/\DeltaT]] + A_{0}'
    params = ['T_{0}', r'\DeltaT', 'A', 'A_{0}']
    choices = [('Direction', 'sign', {'increase': '+', 'decrease': '-'})]
    bounds = [(0, None), (0, None), (None, None), (None, None)]

    @staticmethod
    def func(x: np.ndarray | float, t0: float, dt: float, amp: float, off: float, sign: str = '+') -> np.ndarray | float:
        if sign not in '+-':
            raise ValueError(f"`value` is `+` or `-`, not {sign}")

        error_func = 1
        if sign == '+':
            error_func += special.erf((x-t0)/dt)
        else:
            error_func -= special.erf((x - t0) / dt)

        return amp * error_func / 2. + off


class HendricksonBray:
    type = 'Line shape'
    name = 'Hendrickson-Bray'
    equation = 'AB/(B+(A-B)exp[-E/(kT)]) + w_0'
    params = ['A', 'B', 'E', 'w_{0}']
    bounds = [(0, None)] * 4

    @staticmethod
    def func(x: np.ndarray | float, a: float, b: float, e: float, w0: float) -> np.ndarray | float:
        return a*b / (b + (a-b)*np.exp(-e/kB/x)) + w0
