import sys
import numpy as np
from PIL import Image
from tifffile import imwrite, imread

print("Loading image", file=sys.stderr)

inf = sys.stdin.buffer
outf = sys.stdout.buffer

def create_circular_mask(h, w, center=None, radius=None):
  if center is None: # use the middle of the image
    center = (int(w/2), int(h/2))
  if radius is None: # use the smallest distance between the center and image walls
    radius = min(center[0], center[1], w-center[0], h-center[1])
  Y, X = np.ogrid[:h, :w]
  dist_from_center = np.sqrt((X - center[0])**2 + (Y-center[1])**2)
  mask = dist_from_center <= radius
  return mask

def create_rhombus_mask(h, w, center=None, radius=None):
  if center is None: # use the middle of the image
    center = (int(w/2), int(h/2))
  if radius is None: # use the smallest distance between the center and image walls
    radius = min(center[0], center[1], w-center[0], h-center[1])
  Y, X = np.ogrid[:h, :w]
  dist_from_center = abs(X - center[0]) + abs(Y-center[1])
  mask = dist_from_center <= radius
  return mask

n = 16
# box = np.ones([512,512]).astype(int)
box = np.zeros([512,512]).astype(int)
box[256-n:256+n, 256-n:256+n] = 1
# box = create_circular_mask(512, 512, radius=n)
#box = create_rhombus_mask(512, 512, radius=n)
print(np.sum(box), file=sys.stderr)

image = imread(inf)
print(image.shape, file=sys.stderr)
r = image[:,:,0]
g = image[:,:,1]
b = image[:,:,2]

np.set_printoptions(suppress=True, linewidth=np.inf)

def next_power(x):
  return 2**(x-1).bit_length()

def pad(array, target_shape):
  return np.pad(array, [(0, target_shape[i] - array.shape[i]) for i in range(len(array.shape))])

r = np.asarray(r, dtype='int64')
g = np.asarray(g, dtype='int64')
b = np.asarray(b, dtype='int64')

pw = image.shape[0]
#pw = next_power(image.size[0])
#ph = next_power(image.size[1])

#r = pad(r,(ph,pw))
#g = pad(g,(ph,pw))
#b = pad(b,(ph,pw))

def taxi(t):
  return abs((t%1)*2-1)*2-1

def dftmat(func,n,m=0):
  if m == 0:
    m = n
  mat = np.matrix(np.zeros((m,n), dtype=type(func(0))))
  for s in range(n):
    for f in range(m):
      mat[f,s] = func(s*f/n)
  return mat

def rt(f):
  return lambda x: f(x+0.25)

def fzip(f,g):
  return lambda x: f(x) + g(x)*1j

def cos(x):
  return np.cos(x*np.pi*2)

def hadamard(m):
  return np.matrix([
    [ 1-((i&j).bit_count()&1)*2 for j in range(0,m) ]
    for i in range(0,m)
  ])

def hadamard_sequency_ordered(m):
  mat = hadamard(m)
  result = [0] * m
  for r in np.array(mat):
    s = 0
    l = r[0]
    for c in r[1:]:
      if l != c:
        s += 1
      l = c
    result[s] = r
  return np.matrix(result)

def save(name,r,g,b):
  print("Saving image", file=sys.stderr)
  imwrite(name, np.dstack((r,g,b)), photometric='rgb')
  #Image.fromarray(np.dstack((r,g,b))).save(name, "png");

def mul2d(a,m):
  for i in range(len(a)):
    a[i] = a[i] * m
  a = np.transpose(a)
  for i in range(len(a)):
    a[i] = a[i] * m

print("Genrating hadamard matrix", file=sys.stderr)
#m = hadamard(pw)
# m = hadamard_sequency_ordered(pw)
# m = dftmat(fzip(cos,rt(cos)), pw)
m = dftmat(fzip(taxi,rt(taxi)), pw)
m = np.matrix(np.fft.fftshift(m))
# tm = m
# tm = np.transpose(np.conj(m))
tm = np.linalg.inv(m)

r = np.asarray(r, dtype=m.dtype)
g = np.asarray(g, dtype=m.dtype)
b = np.asarray(b, dtype=m.dtype)

print("Performing hadamard transform 1/3 (red plane)", file=sys.stderr)
mul2d(r,m)
r = r * (box)
mul2d(r,tm)
print("Performing hadamard transform 2/3 (green plane)", file=sys.stderr)
mul2d(g,m)
g = g * (box)
mul2d(g,tm)
print("Performing hadamard transform 3/3 (blue plane)", file=sys.stderr)
mul2d(b,m)
b = b * (box)
mul2d(b,tm)

r = np.asarray(r, dtype='int64')
g = np.asarray(g, dtype='int64')
b = np.asarray(b, dtype='int64')

save(outf, r,g,b)