import numpy as np
import math
import sys
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
import time
import scipy.sparse


def pad_to_power_2(x_raw):
    N = x_raw.shape[0]
    N_ceil = pow(2, math.ceil(math.log(N, 2)))
    if N - N_ceil == 0:
        return x_raw, N
    else:
        x_append = np.zeros(N_ceil - N)
        return np.append(x_raw, x_append), N_ceil


def dft(x):
    x = np.asarray(x, dtype=np.complex_)
    N = x.shape[0]
    n_small = np.asarray(range(0, N), dtype=np.complex_)
    k = n_small.reshape((N, 1))
    coe = np.exp(-2j * np.pi * k * n_small / N)
    return np.dot(coe, x)


def dft_inv(X):
    X = np.asarray(X, dtype=np.complex_)
    N = X.shape[0]
    n_small = np.asarray(range(0, N), dtype=np.complex_)
    k = n_small.reshape((N, 1))
    coe = np.exp(2j * np.pi * k * n_small / N)
    return np.dot(coe, X)


def fft(x_raw):
    x_raw = np.asarray(x_raw)
    x, N = pad_to_power_2(x_raw)

    if N <= 2:
        return dft(x)
    else:
        X_even = fft(x[::2])
        X_odd = fft(x[1::2])
        coe = np.exp(-2j * np.pi * np.arange(N) / N)
        return np.concatenate((X_even + coe[:int(N / 2)] * X_odd,
                               X_even + coe[int(N / 2):] * X_odd))


def fft_inv_helper(x_raw):
    x_raw = np.asarray(x_raw)
    x, N = pad_to_power_2(x_raw)

    if N <= 2:
        return dft_inv(x)
    else:
        X_even = fft_inv_helper(x[::2])
        X_odd = fft_inv_helper(x[1::2])
        coe = np.exp(2j * np.pi * np.arange(N) / N)
        return np.concatenate([X_even + coe[:int(N / 2)] * X_odd,
                               X_even + coe[int(N / 2):] * X_odd])


def fft_inv(x_raw):
    x, N = pad_to_power_2(x_raw)
    coe = 1 / N
    result = coe * np.array(fft_inv_helper(x_raw))
    return result


def two_d_dft(signal: np.array):
    row, column = signal.shape
    columns_transformation = []
    for c in range(column):
        columns_transformation.append(dft(signal[:, c]))

    final_result = []
    columns_transformation = np.asarray(columns_transformation, dtype=np.complex_).T
    for r in range(row):
        final_result.append(dft(columns_transformation[r, :]))

    final_result = np.asarray(final_result, dtype=np.complex_)
    zero_padding = np.zeros((columns_transformation.shape[0] - row, final_result.shape[1]), dtype=np.complex_)
    return np.append(final_result, zero_padding, 0)


def two_d_fft(signal: np.array):
    row, column = signal.shape
    columns_transformation = []
    for c in range(column):
        columns_transformation.append(fft(signal[:, c]))

    final_result = []
    columns_transformation = np.asarray(columns_transformation, dtype=np.complex_).T
    for r in range(row):
        final_result.append(fft(columns_transformation[r, :]))

    final_result = np.asarray(final_result, dtype=np.complex_)
    zero_padding = np.zeros((columns_transformation.shape[0] - row, final_result.shape[1]), dtype=np.complex_)
    return np.append(final_result, zero_padding, 0)


def two_d_fft_inverse(signal: np.array):
    row, column = signal.shape
    columns_inverse = []
    for c in range(column):
        columns_inverse.append(fft_inv(signal[:, c]))

    final_result = []
    columns_inverse = np.asarray(columns_inverse, dtype=np.complex_).T
    for r in range(row):
        final_result.append(fft_inv(columns_inverse[r, :]))

    final_result = np.asarray(final_result, dtype=np.complex_)
    zero_padding = np.zeros((columns_inverse.shape[0] - row, final_result.shape[1]), dtype=np.complex_)
    return np.append(final_result, zero_padding, 0)


def mode1(img_data_raw):
    sfft = two_d_fft(img_data_raw)
    # the 2dfft generated by our function

    plt.figure(1)
    plt.subplot(121)
    plt.imshow(img_data_raw, plt.cm.gray)

    plt.subplot(122)
    plt.imshow(np.abs(sfft), norm=LogNorm(vmin=5))
    plt.show()


def filter_high(r, img_data_raw):
    sfft = two_d_fft(img_data_raw)
    # the 2dfft generated by our function

    w, h = sfft.shape
    sfft[int(r * w):int(w * (1 - r))] = 0
    sfft[:, int(h * r):int(h * (1 - r))] = 0
    new = two_d_fft_inverse(sfft)
    img_w, img_h = img_data_raw.shape
    new_im = new[0:img_w, 0:img_h]

    plt.figure(1)
    plt.subplot(121)
    plt.imshow(img_data_raw, plt.cm.gray)

    plt.subplot(122)
    plt.imshow(new_im.real, plt.cm.gray)
    plt.show()


def filter_low(r, img_data_raw):
    sfft = two_d_fft(img_data_raw)
    # the 2dfft generated by our function

    w, h = sfft.shape
    sfft[0:int(w * r), 0:int(h * r)] = 0
    sfft[0:int(w * r), int(h * 1 - r):h] = 0
    sfft[int(w * 1 - r):w, int(h * 1 - r):h] = 0
    sfft[int(w * 1 - r):w, 0:int(h * r)] = 0
    new = two_d_fft_inverse(sfft)
    img_w, img_h = img_data_raw.shape
    new_im = new[0:img_w, 0:img_h]

    plt.figure(1)
    plt.subplot(121)
    plt.imshow(img_data_raw, plt.cm.gray)

    plt.subplot(122)
    plt.imshow(new_im.real, plt.cm.gray)
    plt.show()


def compress(img_data_raw):
    w, h = img_data_raw.shape
    img_fft = two_d_fft(img_data_raw)
    compression_percentage = np.linspace(0, 0.95, 6)
    fig, axis = plt.subplots(2, 3)
    for i in range(0, 6):
        compressed_fft = high_frequency_compression(img_fft, compression_percentage[i])
        img = two_d_fft_inverse(compressed_fft).real[:w, :h]
        axis[int(i / 3)][int(i % 3)].imshow(abs(img), cmap='gray')
        axis[int(i / 3)][int(i % 3)].set_title(str(int(compression_percentage[i] * 100)) + "% compression")

    plt.show()


def high_frequency_compression(img_fft, removal_percentage):
    w, h = img_fft.shape
    p = math.sqrt((1 - removal_percentage) / 4)

    img_fft[int(w * p):int(- w * p)] = 0
    img_fft[:, int(h * p):int(- h * p)] = 0

    print("----- compression percentage: ", removal_percentage, " -----")
    print("number of zeros: ", np.count_nonzero(img_fft))
    print("zero percentage: ", 1 - (np.count_nonzero(img_fft) / (w * h)))

    csr = scipy.sparse.csr_matrix(img_fft)
    scipy.sparse.save_npz("frequency_compression_compressed_" + str(round(removal_percentage, 2)), csr, compressed=False)

    return img_fft


def smallest_coefficient_compression(img_fft, removal_percentage):
    width, height = img_fft.shape
    num_remove = int(width * height * removal_percentage)
    sorted_index = np.argsort(abs(img_fft), axis=None)[:num_remove]
    np.put(img_fft, sorted_index, np.zeros((num_remove, 1)))

    print("----- compression percentage: ", removal_percentage, " -----")
    print("number of zeros: ", np.count_nonzero(img_fft))
    print("zero percentage: ", 1 - (np.count_nonzero(img_fft) / (width * height)))

    csr = scipy.sparse.csr_matrix(img_fft)
    scipy.sparse.save_npz("coefficient_compression_" + str(round(removal_percentage, 2)), csr, compressed=False)

    return img_fft


def run_time_evaluation():
    size = [2 ** 5, 2 ** 6, 2 ** 7, 2 ** 8, 2 ** 9, 2 ** 10]
    fft_avg = []
    fft_std = []
    dft_avg = []
    dft_std = []
    for s in size:
        data = np.random.rand(s, s)
        fft_time_result = []
        dft_time_result = []
        for i in range(10):
            start_time = time.time()
            two_d_fft(data)
            end_time = time.time()
            fft_time_result.append(end_time - start_time)

            start_time = time.time()
            two_d_dft(data)
            end_time = time.time()
            dft_time_result.append(end_time - start_time)

        fft_avg.append(np.mean(fft_time_result))
        fft_std.append(np.std(fft_time_result) * 2)

        dft_avg.append(np.mean(dft_time_result))
        dft_std.append(np.std(dft_time_result) * 2)

    print("----- runtime result -----")
    for i in range(6):
        print("data size: ", size[i], "*", size[i])
        print("fft mean time: ", fft_avg[i], "; fft std deviation: ", fft_std[i])
        print("dft mean time: ", dft_avg[i], "; dft std deviation: ", dft_std[i])

    plt.clf()
    plt.errorbar(size, fft_avg, yerr=fft_std, fmt="b.", label="fft")
    plt.errorbar(size, dft_avg, yerr=dft_std, fmt="r.", label="dft")
    plt.title("runtime comparison")
    plt.xlabel("width of data")
    plt.ylabel("time/s")
    plt.legend()
    plt.show()


if __name__ == '__main__':
    mode = 1
    img = "moonlanding.png"
    index = 1
    while index < len(sys.argv):
        if sys.argv[index] == "-m":
            index += 1
            mode = int(sys.argv[index])
        elif sys.argv[index] == "-i":
            index += 1
            img = sys.argv[index]
        else:
            print("error in argument")
            exit()

        index += 1

    img_data_raw = plt.imread(img).astype(float)

    if mode == 1:
        mode1(img_data_raw)
        exit()

    if mode == 2:
        filter_high(0.05, img_data_raw)
        exit()

    if mode == 3:
        compress(img_data_raw)
        exit()

    if mode == 4:
        run_time_evaluation()
        exit()
