import numpy as np


def wavelet(image, width=None):
  shape = image.shape
  if len(shape) != 2 or shape[0] != shape[1]:
    raise(f"Expected quadratic image but got data with shape: {shape}")

  if width is None:
    width = shape[0]

  if shape[0] < width:
    raise(f"Image with shape {shape} is smaller then the requested size {width}.")

  result = image.copy()

  half = width // 2
  for y in range(half):
    for x in range(half):
      px = 2 * x
      py = 2 * y
      result[y, x] = image[py, px]
      result[y, x + half] = image[py, px + 1] - image[py, px]
      result[y + half, x] = image[py + 1, px] - image[py, px]
      result[y + half, x + half] = image[py + 1, px + 1] - image[py, px]

  return result


def refine(image):
  shape = image.shape
  if len(shape) != 2 or shape[0] != shape[1]:
    raise(f"Expected quadratic image but got data with shape: {shape}")

  width = shape[0]
  while width > 1:
    image = wavelet(image, width)
    width = width // 2

  return image



def rotate(image, n=1):
  shape = image.shape
  if len(shape) != 2 or shape[0] != shape[1]:
    raise(f"Expected quadratic image but got data with shape: {shape}")

  result = image.copy()
  w = shape[0] - 1

  if n == 0:   # id
    p = lambda y, x: (y, x)
  elif n == 1:
    p = lambda y, x: (y, w - x)
  elif n == 2:
    p = lambda y, x: (w - y, x)
  elif n == 3:
    p = lambda y, x: (w - y, w - x)
  elif n == 4:
    p = lambda y, x: (x, y)
  elif n == 5:
    p = lambda y, x: (x, w - y)
  elif n == 6:
    p = lambda y, x: (w - x, y)
  elif n == 7:
    p = lambda y, x: (w - x, w - y)
  else:
    raise f"Unexpected n={n} for n in {0,1,2,3,4,5,6}"
    

  for y in range(w):
    for x in range(w):
      result[x,y] = image[y,x]

  return result