import qrcode
import numpy as np
from PIL import Image
import base64

correction_map = {
  'L': qrcode.constants.ERROR_CORRECT_L,
  'M': qrcode.constants.ERROR_CORRECT_M,
  'Q': qrcode.constants.ERROR_CORRECT_Q,
  'H': qrcode.constants.ERROR_CORRECT_H,
}

backgrounds = ['white', 'black', 'gray', 'noise']

bg_map = {
    'white': 255,
    'black': 0,
    'gray': 128,
}

def create_code(text, module_size, margin, background, error_correction, centered, submodule_prop, split=False):
  qr = qrcode.QRCode(
    version=1,
    error_correction=correction_map[error_correction],
    box_size=module_size,
    border=margin if margin > 0 else 1,
  )
  qr.add_data(text)
  qr.make(fit=True)

  img = qr.make_image(fill_color="black", back_color="white")

  # find smallest image size multiple of 256 that can fit qr
  offset_min = 8 * module_size
  w, h = img.size

  mask = None

  if margin == 0:
    corner_size = 9 * module_size
    # make a mask that hides the margin when not in the corners
    mask = np.ones((h, w)).astype(bool)
    mask[corner_size:-corner_size, :] = False
    mask[:, corner_size:-corner_size] = False
    mask[-corner_size:, -corner_size:] = False
    mask[module_size:-module_size, module_size:-module_size] = True

  if submodule_prop != 1:
    submodule_size = round(module_size * submodule_prop)
    k = (module_size - submodule_size) // 2
    # qr_array = np.array(img).reshape((h // module_size, w // module_size, module_size, module_size))
    new_mask = np.zeros((h // module_size, module_size, w // module_size, module_size)).astype(bool)
    new_mask[:, k:k+submodule_size, :, k:k+submodule_size] = True
    mask = np.logical_and(mask, new_mask.reshape((h, w)))

  w = (w + 255 + offset_min) // 256 * 256
  h = (h + 255 + offset_min) // 256 * 256

  # create new image with background chosen by user
  if background == 'noise':
    # noise = np.random.randint(0, 256, (h // module_size, w // module_size), dtype=np.uint8)
    noise = np.random.normal(128, 64, (h // module_size, w // module_size)).astype(np.uint8)
    # noise = (noise // 128) * 255
    noise = np.round(noise / 128) * 128
    # clamp values
    noise = np.clip(noise, 0, 255)
    
    bg = Image.fromarray(noise.astype(np.uint8)).resize((w, h), Image.NEAREST)
  else:
    bg = Image.new('L', (w, h), bg_map[background])

  # paste qr code in the center of the image
  if centered:
    coords = ((w - img.size[0]) // 2, (h - img.size[1]) // 2)
  else:
    # paste it aligned on module size, closest to center
    coords = ((w - img.size[0]) // 2 // module_size * module_size, (h - img.size[1]) // 2 // module_size * module_size)

  if mask is not None:
    new_img = Image.new('L', img.size, bg_map['gray'] if background == 'noise' else bg_map[background])
    new_img.paste(img, mask=Image.fromarray(mask))
    bg.paste(new_img, coords)
  else:
    bg.paste(img, coords)

  if split:
    # isolate the 3 qr markers, paste them onto a new gray image
    m_size = module_size * 9

    new_bg = Image.new('L', (w, h), bg_map['gray'])

    # create mask for the 3 markers from a numpy array
    mask = np.zeros((h, w), dtype=bool)
    mask[coords[1]:coords[1] + m_size, coords[0]:coords[0] + m_size] = True
    mask[coords[1]:coords[1] + m_size, coords[0] + img.size[0] - m_size:coords[0] + img.size[0]] = True
    mask[coords[1] + img.size[1] - m_size:coords[1] + img.size[1], coords[0]:coords[0] + m_size] = True

    # paste the 3 markers on the new image
    new_bg.paste(bg, mask=Image.fromarray(mask))

    return bg, new_bg
  else:
    return bg,