#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Christoph Federrath

import argparse
import numpy as np
import cfpack as cfp
from cfpack.defaults import *


def myfunc(x):
    f = np.sin(x) + np.sin(10*x) + 2
    return f

def myfunc2d(x,y):
    f = np.sin(x) + np.sin(10*y) + 2
    return f


# ===== the following applies in case we are running this in script mode =====
if __name__ == "__main__":

    # parse command line arguments
    parser = argparse.ArgumentParser(description='Fourier analysis demo.')
    parser.add_argument("-t", "--type", default='1d', help="Type of analysis (can be '1d', '2d', or <HDF5 filename>)")
    args = parser.parse_args()

    if args.type == '1d':

        # define function grid
        num = 200
        x = np.linspace(0, 2*np.pi, num=num)
        f = myfunc(x)

        # plot y(x)
        cfp.plot(x=x, y=f, xlabel=r"$x$", ylabel=r"$f(x)$", save="1d_data.pdf")

        # compute FFT, define wavenumbers
        # f_ft = np.fft.fft(f) / num # note that without "norm", we have to normalise ourselves by the number of datapoints
        f_ft = np.fft.fft(f, norm='forward')
        k = np.fft.fftfreq(num) * num

        # shift 0-frequency to the centre
        k = np.fft.fftshift(k)
        f_ft = np.fft.fftshift(f_ft)

        # compute spectrum
        k_spect = k[num//2:] # these are only the positive wavennumbers (for plotting the spectrum below)
        f_ft_pos = f_ft[k >= 0] # Fourier-transformed y where k >= 0
        f_ft_neg = f_ft[k <= 0][1:] # Fourier-transformed y where k <= 0 (excluding the Nyquist frequency; here k = -100)
        f_ft_neg = np.flip(f_ft_neg) # fold the negative modes onto the positive-mode indices
        f_spect = abs(f_ft_pos)**2 + abs(f_ft_neg)**2 # this is now the spectrum
        f_spect[0] /= 2 # correct the k=0 mode, which should not be counted twice, but we did in the previous line

        # plot the power spectrum
        cfp.plot(x=k_spect, y=f_spect, linestyle="", marker="+", xlabel=r"$k$", ylabel=r"$P(k)$",
                 save="1d_data_power_spectrum.pdf")

        # filter out high-frequency modes and plot again
        index = abs(k) > 5 # get all indices where |k| > 5
        f_ft[index] = 0.0 # set them to zero to filter away high-frequency modes
        # first shift k=0 back to the start, and feed that into the inverse FT

        f_filtered = np.fft.ifft(np.fft.ifftshift(f_ft), norm='forward')
        # plot filtered data
        cfp.plot(x=x, y=f_filtered.real, xlabel=r"$x$", ylabel=r"Filtered $y$",
                 save="1d_data_Fourier_filtered.pdf")

        stop()

    if args.type == '2d':

        # define function grid
        num = 50
        x, y = cfp.get_2d_coords(cmin=[0,0], cmax=[2*np.pi,2*np.pi], ndim=[num,num], cell_centred=False)
        f = myfunc2d(x,y)

        # plot function
        cfp.plot_map(f, xlim=[0,2*np.pi], ylim=[0,2*np.pi],
                     xlabel=r"$x$", ylabel=r"$y$", cmap_label=r"$f(x,y)$", save="2d_data.pdf")

        # 2D FFT
        f_ft = np.fft.fft2(f, norm='forward')
        f_ft = np.fft.fftshift(f_ft) # shift 0-frequency to the centre

        # create k grid, such that k=0 is in the centre, i.e., kx, ky are both running in [-num, num-1]
        kx, ky = cfp.get_2d_coords(cmin=[-num//2,-num//2], cmax=[num//2-1,num//2-1], ndim=[num,num], cell_centred=False)

        # plot spectrum
        cfp.plot_map(abs(f_ft), xlim=[np.min(kx),np.max(kx)], ylim=[np.min(ky),np.max(ky)],
                     xlabel=r"$k_x$", ylabel=r"$k_y$", cmap_label=r"$P(k_x,k_y)$",
                     save="2d_data_power_spectrum.pdf")

        # filter and plot inverse FT
        ind = abs(ky) > 5 # filter only on ky
        f_ft[ind] = 0.0 # filter away high-frequency modes
        # first shift k=0 back to the start, and feed that into the inverse FT
        f_new = np.fft.ifft2(np.fft.ifftshift(f_ft), norm='forward')
        # plot filtered data
        cfp.plot_map(f_new.real, xlim=[0,2*np.pi], ylim=[0,2*np.pi],
                     xlabel=r"$x$", ylabel=r"$y$", cmap_label=r"Filtered $f(x,y)$",
                     save="2d_data_Fourier_filtered.pdf")

        stop()

    # data file; and args.type contains the filename
    if args.type != '1d' and args.type != '2d':

        # read column density map from file (filename is input via argparse -t)
        from cfpack import hdfio
        data = hdfio.read(args.type, 'dens_proj_xy')

        # grab only the first two dims
        d = data[:,:,0]

        # define number of cells
        num = d.shape[0]

        # plot function
        cfp.plot_map(d, log=True, xlabel=r"$x$", ylabel=r"$y$", cmap_label=r"Data", show=True)

        # FFT
        d_ft = np.fft.fft2(d, norm='forward')
        d_ft = np.fft.fftshift(d_ft) # shift 0-freq to centre

        # create k grid, such that 0-freq is in the centre
        kx, ky = cfp.get_2d_coords(cmin=[-num//2,-num//2], cmax=[num//2-1,num//2-1], ndim=[num,num], cell_centred=False)

        # plot spectrum
        cfp.plot_map(abs(d_ft), xlim=[np.min(kx),np.max(kx)], ylim=[np.min(ky),np.max(ky)],
                     xlabel=r"$k_x$", ylabel=r"$k_y$", cmap_label=r"$P(k_x,k_y)$", log=True, show=True)

        # filter and inverse FT
        k = np.sqrt(kx**2 + ky**2) # create circle selection in k-space
        ind = k > 100 # filter away everything that is outside the k=100 circle in k-space
        d_ft[ind] = 0.0
        d_new = np.fft.ifft2(np.fft.ifftshift(d_ft), norm='forward')

        # plot filtered spectrum
        cfp.plot_map(abs(d_ft), xlim=[np.min(kx),np.max(kx)], ylim=[np.min(ky),np.max(ky)],
                     xlabel=r"$k_x$", ylabel=r"$k_y$", cmap_label=r"Filtered $P(k_x,k_y)$", log=True, show=True)

        # plot filtered data
        cfp.plot_map(d_new.real, log=True, xlabel=r"$x$", ylabel=r"$y$",
                     cmap_label=r"Filtered data", show=True)

        stop()