import numpy as np
import torch


class NScaler(object):
    def transform(self, data):
        return data

    def inverse_transform(self, data):
        return data


class StandardScaler:
    """
    Standard the input
    """

    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def transform(self, data):
        return (data - self.mean) / self.std

    def inverse_transform(self, data):
        if type(data) is torch.Tensor and type(self.mean) is np.ndarray:
            self.std = torch.from_numpy(self.std).to(data.device).type(data.dtype)
            self.mean = torch.from_numpy(self.mean).to(data.device).type(data.dtype)
        return (data * self.std) + self.mean


class MinMax01Scaler:
    """
    Standard the input
    """

    def __init__(self, min_value, max_value):
        self.min_value = min_value
        self.max_value = max_value

    def transform(self, data):
        return (data - self.min_value) / (self.max_value - self.min_value)

    def inverse_transform(self, data):
        if type(data) is torch.Tensor and type(self.min_value) is np.ndarray:
            self.min_value = torch.from_numpy(self.min_value).to(data.device).type(data.dtype)
            self.max_value = torch.from_numpy(self.max_value).to(data.device).type(data.dtype)
        return data * (self.max_value - self.min_value) + self.min_value


class MinMax11Scaler:
    """
    Standard the input
    """

    def __init__(self, min_value, max_value):
        self.min_value = min_value
        self.max_value = max_value

    def transform(self, data):
        return ((data - self.min_value) / (self.max_value - self.min_value)) * 2. - 1.

    def inverse_transform(self, data):
        if type(data) is torch.Tensor and type(self.min_value) is np.ndarray:
            self.min_value = torch.from_numpy(self.min_value).to(data.device).type(data.dtype)
            self.max_value = torch.from_numpy(self.max_value).to(data.device).type(data.dtype)
        return ((data + 1.) / 2.) * (self.max_value - self.min_value) + self.min_value


class ColumnMinMaxScaler:
    # Note: to use this scale, must init the min and max with column min and column max
    def __init__(self, min_value, max_value):
        self.min_value = min_value
        self.min_max = max_value - self.min_value
        self.min_max[self.min_max == 0] = 1

    def transform(self, data):
        print(data.shape, self.min_max.shape)
        return (data - self.min_value) / self.min_max

    def inverse_transform(self, data):
        if type(data) is torch.Tensor and type(self.min_value) is np.ndarray:
            self.min_max = torch.from_numpy(self.min_max).to(data.device).type(torch.float32)
            self.min_value = torch.from_numpy(self.min_value).to(data.device).type(torch.float32)
        # print(data.dtype, self.min_max.dtype, self.min.dtype)
        return data * self.min_max + self.min_value


if __name__ == '__main__':
    test_data = np.array([[0, 0, 0, 1], [0, 1, 3, 2], [0, 2, 1, 3]])
    print(test_data)
    minimum = test_data.min(axis=1)
    print(minimum, minimum.shape, test_data.shape)
    maximum = test_data.max(axis=1)
    print(maximum)
    print(test_data - minimum)
    test_data = (test_data - minimum) / (maximum - minimum)
    print(test_data)
    print(0 == 0)
    print(0.00 == 0)
    print(0 == 0.00)