import torch
from .utils import init_parameter

class LinearS(torch.nn.Module):
    """
    Applies a linear S transformation to the incoming data.

    Attributes
    ----------
    foot : float
        foot of the linear S function
    shoulder : float
        shoulder of the linear S function

    Returns
    -------
    torch.tensor
        a tensor of equal size to the input tensor
    """
    def __init__(self, foot:float = None, shoulder:float = None) -> None:
        super().__init__()
        self.shoulder = init_parameter(shoulder)
        self.foot = init_parameter(foot)
    
    def get_center(self) -> torch.Tensor:
        return self.foot - self.shoulder
    
    def forward(self, x) -> torch.Tensor:
        x = x - self.foot
        x = x / (self.shoulder - self.foot)
        x = torch.maximum(x, torch.tensor(0))
        x = torch.minimum(x, torch.tensor(1))
        return x
    
    def __setitem__(self, key, value):
        if key == "shoulder":
            self.shoulder = init_parameter(value)
        elif key == "foot":
            self.foot = init_parameter(value)
        else:
            raise KeyError(f"Invalid key: {key}")