# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import tensorflow as tf
import numpy as np
from cvnn.layers import Convolutional
from pdb import set_trace
import sys
from scipy import signal
from scipy import linalg

# Both TF and NP results calculate fft the same way

# +
aaa = np.linspace(1.0, 10000.0, 10000)
x = aaa + 1j * aaa
x_tensor = tf.convert_to_tensor(x)

tf_fft = tf.signal.fft(x_tensor)
np_fft = np.fft.fft(x)

print(tf_fft.dtype)
print(np.all(tf_fft.numpy() == np_fft))  # Results are not exactly the same (but fair enough)
print(tf_fft.numpy()[:10])
print(np_fft[:10])
print(tf_fft.numpy() == np_fft)
print((tf_fft.numpy() - np_fft)[1])
# -

# ## Testing on 1D

# +
b = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
c = [1, 0, 1]

b_pad = tf.cast(tf.pad(b, tf.constant([[0, 2]])), tf.complex64)  # Full padding
I = tf.signal.fft(tf.cast(b_pad, tf.complex64))
paddings = tf.constant([[0, 9]])
c_pad = tf.cast(tf.pad(c, paddings), tf.complex64)
C = tf.signal.fft(c_pad)
F = tf.math.multiply(I, C)
f = tf.signal.ifft(F)
f_real = tf.cast(f, tf.int32)

# print("std_out: " + str(std_out))
print("f_real:  \t" + str(f_real.numpy()))
print("convolve:\t" + str(np.convolve(b, c)))

manual_conv = []
for i in range(len(b)-len(c)+1):
    manual_conv.append(np.sum(tf.math.multiply(c, b[i:i+3]).numpy()))
print("Manual nn conv: " + str(manual_conv))

c.reverse()
manual_conv = []
for i in range(len(b)-len(c)+1):
    manual_conv.append(np.sum(tf.math.multiply(c, b[i:i+3]).numpy()))
print("Manual fft conv:" + str(manual_conv))
# -

# ## Testing on 2D

# +
np.set_printoptions(suppress=True)
img2 = np.array([
    [10, 10, 10, 0, 0, 0],
    [10, 10, 10, 0, 0, 0],
    [10, 10, 10, 0, 0, 0],
    [10, 10, 10, 0, 0, 0],
    [10, 10, 10, 0, 0, 0],
    [10, 10, 10, 0, 0, 0]
]).astype(np.float64)
k = np.array([
        [1., 0., -1.],
        [1., 0., -1.],
        [1., 0., -1.]
    ]).astype(np.float64)
mode = 'full'

conv = Convolutional(1, (3, 3), (6, 6, 1), padding=2, input_dtype=np.float32)
conv.kernels = []
conv.kernels.append(tf.reshape(tf.cast(tf.Variable(k, name="kernel" + str(0) + "_f" + str(0)), dtype=np.float32),
                               (3, 3, 1)))
std_out = conv([img2])[..., 0]
print("manual_conv: \n" + str(std_out.numpy()))

img_tf = tf.constant(tf.reshape(img2, (1, 6, 6, 1)), dtype=tf.float64)
k_tf = tf.constant(tf.reshape(k, (3, 3, 1, 1)), dtype=tf.float64)
conv_tf = tf.nn.conv2d(img_tf, k_tf, strides=[1, 1], padding="SAME")[0, ..., 0]
print("tf_nn_conv2d: \n" + str(np.around(conv_tf.numpy())))
# set_trace()

img2_pad = tf.pad(img2.astype(np.float64), tf.constant([[0, 2], [0, 2]]))
k_pad = tf.pad(k, tf.constant([[0, 5], [0, 5]]))
I = tf.signal.fft2d(tf.cast(img2_pad, tf.complex128))
print(I)
K = tf.signal.fft2d(tf.cast(k_pad, tf.complex128))
F = tf.math.multiply(I, K)
f = tf.signal.ifft2d(F)
f_real = tf.cast(f, tf.int32)
print("manual_fft_conv: " + str(f_real))

np_fft_conv = np.array(signal.fftconvolve(img2, k, mode=mode) , np.int32)
print("sp_fft_conv_" + mode + ":\n" + str(np_fft_conv))

np_conv = np.array(signal.convolve2d(img2 , k, mode), np.int32)
print("sp_conv2d" + mode + ":\n" + str(np_conv))

# Check numpy implementation
I = np.fft.fft2(img2_pad)
K = np.fft.fft2(tf.pad(k, tf.constant([[0, 5], [0, 5]])))
F = np.multiply(I, K)
f = np.fft.ifft2(F)
print("np_fft_conv: \n" + str(np.round(f.astype(np.float32))))


# -

# There are 2 results here and they are:
#
# - $(x*v)(n) = \sum x(m) \, v(m)$
# - $(x*v)(n) = \sum x(m) \, v(n-m)$

# ## StackOverflow Example
#
# https://stackoverflow.com/questions/40703751/using-fourier-transforms-to-do-convolution

# +
x = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 3, 0], [0, 0, 0, 1]])
y = np.array([[4, 5], [3, 4]])

print("conv:\n",  signal.convolve2d(x, y, 'full'))

s1 = np.array(x.shape)
s2 = np.array(y.shape)
size = s1 + s2 - 1      # Full padding size = (5, 5)
fsize = 2 ** np.ceil(np.log2(size)).astype(int)    # I do this to have a 2^n size to make fft faster
# fsize = (8, 8)
fslice = tuple([slice(0, int(sz)) for sz in size])  
# slice to get the values later ([0:5], [0:5])

new_x = np.fft.fft2(x, fsize)
new_y = np.fft.fft2(y, fsize)
result = np.fft.ifft2(new_x*new_y)[fslice].copy()

print("manual fft method:\n", np.array(result.real, np.int32))

print("fft:\n" , np.array(signal.fftconvolve(x, y), np.int32))
# -

# ## Complex Conv

# +
img2 = np.array([
    [10, 10, 10, 0, 0, 0],
    [10, 10, 10, 0, 0, 0],
    [10, 10, 10, 0, 0, 0],
    [10, 10, 10, 0, 0, 0],
    [10, 10, 10, 0, 0, 0],
    [10, 10, 10, 0, 0, 0]
]).astype(np.float64)
k = np.array([
        [1., 0., -1.],
        [1., 0., -1.],
        [1., 0., -1.]
    ]).astype(np.float64)
mode = 'full'

c_img = tf.Variable(tf.complex(img2, img2))
c_k = tf.Variable(tf.complex(k, np.zeros(k.shape)))
conv_tf = tf.nn.conv2d(c_img, c_k, strides=[1, 1], padding="SAME")[0, ..., 0]
print(conv_tf)
# -

# https://stackoverflow.com/questions/47577458/complex-convolution-in-tensorflow

# ## Tensorflow fftconv2d

# +
import tensorflow as tf

def _centered(arr, newshape):
    # Return the center newshape portion of the array.
    currshape = tf.shape(arr)[-2:]
    startind = (currshape - newshape) // 2
    endind = startind + newshape
    return arr[..., startind[0]:endind[0], startind[1]:endind[1]]

def fftconv(in1, in2, mode="full"):
    mode = mode.lower()
    # Reorder channels to come second (needed for fft)
    in1 = tf.transpose(in1, perm=[0, 3, 1, 2])
    in2 = tf.transpose(in2, perm=[0, 3, 1, 2])

    # Extract shapes
    s1 = tf.convert_to_tensor(tf.shape(in1)[-2:])
    s2 = tf.convert_to_tensor(tf.shape(in2)[-2:])
    shape = s1 + s2 - 1

    # Compute convolution in fourier space
    sp1 = tf.spectral.rfft2d(in1, shape)
    sp2 = tf.spectral.rfft2d(in2, shape)
    ret = tf.spectral.irfft2d(sp1 * sp2, shape)

    # Crop according to mode
    if mode == "full":
        cropped = ret
    elif mode == "same":
        cropped = _centered(ret, s1)
    elif mode == "valid":
        cropped = _centered(ret, s1 - s2 + 1)
    else:
        raise ValueError("Acceptable mode flags are 'valid',"
                         " 'same', or 'full'.")

    # Reorder channels to last
    result = tf.transpose(cropped, perm=[0, 2, 3, 1])
    return result


# -

onv_tf = fftconv(img_tf, k_tf, mode="SAME")[0, ..., 0]



StarcoderdataJupyterScriptsDedupFiltered