# tanh-like
from .tanh import Tanh
from torch.nn.functional import tanh
from .softsign import SoftSign
from torch.nn.functional import softsign

# relu-like
from .relu import ReLU
from torch.nn.functional import relu
from .silu import SiLU
from torch.nn.functional import silu
from .gelu import GELU
from torch.nn.functional import gelu
from .elu import ELU
from torch.nn.functional import elu
from .softplus import SoftPlus
from torch.nn.functional import softplus
from .mish import Mish
from torch.nn.functional import mish

from typing import Callable

def parse_activ_f(activ_str: str) -> Callable:
    """
    Parses string into proper activation function.
    """
    # Parse activation
    match activ_str:
        # tanh-like
        case "tanh":
            return Tanh().f
        case "softsign":
            return SoftSign().f
        # relu-like
        case "relu":
            return ReLU().f
        case "silu":
            return SiLU().f
        case "gelu":
            return GELU().f
        case "elu":
            return ELU().f
        case "softplus":
            return SoftPlus().f
        case "mish":
            return Mish().f
        case _:
            raise NotImplementedError(f"Activation {activ_str} is not implemented yet")

def parse_activ_f_TC(activ_str: str) -> Callable:
    """
    Parses string into proper activation function.
    """
    # Parse activation
    match activ_str:
        # tanh-like
        case "tanh":
            return tanh
        case "softsign":
            return softsign
        # relu-like
        case "relu":
            return relu
        case "silu":
            return silu
        case "gelu":
            return gelu
        case "elu":
            return elu
        case "softplus":
            return softplus
        case "mish":
            return mish
        case _:
            raise NotImplementedError(f"Activation {activ_str} is not implemented yet")

def parse_activ_df(activ_str: str, order=1):
    """
    Parses string into proper activation function derivative of the given order.
    """
    # Parse activation gradient
    match activ_str:
        # tanh-like
        case "tanh":
            return lambda x: Tanh().df(x, order=order)
        case "softsign":
            return lambda x: SoftSign().df(x, order=order)
        # relu-like
        case "relu":
            return lambda x: ReLU().df(x, order=order)
        case "silu":
            return lambda x: SiLU().df(x, order=order)
        case "gelu":
            return lambda x: GELU().df(x, order=order)
        case "elu":
            return lambda x: ELU().df(x, order=order)
        case "softplus":
            return lambda x: SoftPlus().df(x, order=order)
        case "mish":
            return lambda x: Mish().df(x, order=order)
        case _:
            raise NotImplementedError(f"Gradient of the activation {activ_str} is not implemented yet")
