import numpy as np

def normalize_s2(img: np.ndarray, factor: int = 10_000) -> np.ndarray:
    """Normalizes the given S2 L2A image dividing the raw reflectance values by
    the given factor, then clipping the values to the range [0, 1].

    Args:
        img (np.ndarray): The image to normalize.
        factor (int, optional): The factor to divide the raw reflectance values
            by. Defaults to 10_000.

    Returns:
        np.ndarray: The normalized image.
    """
    return np.clip(img / float(factor), 0, 1)


def impute_nan(arr):
    nan_times = np.where(np.isnan(arr).any(-1).any(-1).any(-1))[0]

    for month in nan_times:
      i = 1
      if month < 4:
        while np.isnan(arr[month]).any():
          if month+i>=6:
            arr[month][np.isnan(arr[month])] = 0.0
          else:
            arr[month][np.isnan(arr[month])] = arr[month+i][np.isnan(arr[month])]
            i+=1
      elif month>=4:
        while np.isnan(arr[month]).any():
          if month-i<0:
            arr[month][np.isnan(arr[month])] = 0.0
          else:
            arr[month][np.isnan(arr[month])] = arr[month-i][np.isnan(arr[month])]
            i+=1
    return arr