from typing import Any, List, Optional
import numpy as np


class StandardScaler:
    def __init__(self,
                 means: np.ndarray = None,
                 stds: np.ndarray = None,
                 replace_nan_token: Any = None):

        self.means = means
        self.stds = stds
        self.replace_nan_token = replace_nan_token

    def fit(self, X):
        X = np.array(X).astype(float)
        self.means = np.nanmean(X, axis=0)
        self.stds = np.nanstd(X, axis=0)
        self.means = np.where(np.isnan(self.means), np.zeros(self.means.shape), self.means)
        self.stds = np.where(np.isnan(self.stds), np.ones(self.stds.shape), self.stds)
        self.stds = np.where(self.stds == 0, np.ones(self.stds.shape), self.stds)
        return self

    def transform(self, X) -> np.ndarray:
        X = np.array(X).astype(float)
        transformed_with_nan = (X - self.means) / self.stds
        transformed_with_none = np.where(np.isnan(transformed_with_nan), self.replace_nan_token, transformed_with_nan)
        return transformed_with_none

    def inverse_transform(self, X):
        X = np.array(X).astype(float)
        transformed_with_nan = X * self.stds + self.means
        transformed_with_none = np.where(np.isnan(transformed_with_nan), self.replace_nan_token, transformed_with_nan)
        return transformed_with_none

    def inverse_transform_variance(self, X):
        X = np.array(X).astype(float)
        transformed_with_nan = X * (self.stds ** 2)
        transformed_with_none = np.where(np.isnan(transformed_with_nan), self.replace_nan_token, transformed_with_nan)
        return transformed_with_none

