#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
To Do
- callbacks
-

@author: Stefan Hiemer
"""

from itertools import product

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines

from tensorflow import complex128,float64,int64,device,convert_to_tensor

import tensorflow_nufft as tfft

def test_nufft1(natoms=10000,
               nsamples=40,
               ndim=1,
               gridlens=[16,32,64,128],
               tols=[1e-3,1e-6,1e-9,1e-12,1e-15]):

    # random seed
    np.random.seed(0)

    # create test data
    coordinates=(convert_to_tensor(np.random.random((nsamples,
                                                    natoms,
                                                    ndim)),
                                  dtype=float64)-0.5) * 2 * np.pi
    values = convert_to_tensor(np.random.random((nsamples,natoms)),
                               dtype=complex128)

    # iterate over different tolerances and grid sizes
    for tol,gridlen in product(tols,gridlens):

        # create grid
        grid_shape=convert_to_tensor(gridlen*np.ones(ndim),
                                     dtype=int64)

        #
        kspace = tfft.nufft(source=values,
                            points=coordinates,
                            grid_shape=grid_shape,
                            transform_type='type_1',
                            fft_direction='forward',
                            tol=tol)

        #
        backtransform = tfft.nufft(source=kspace,
                                   points=coordinates,
                                   grid_shape=grid_shape,
                                   transform_type='type_2',
                                   fft_direction='backward',
                                   tol=tol)

    return

if __name__ == '__main__':

    with device('/CPU:0'):
       test_nufft1(natoms=100,
                   nsamples=40,
                   ndim=1,
                   gridlens=[1000000],
                   tols=[1e-3,1e-6,1e-9,1e-12,1e-15])
    with device('/CPU:0'):
        test_nufft1(natoms=100,
                    nsamples=40,
                    ndim=2,
                    gridlens=[128],
                    tols=[1e-3,1e-6,1e-9,1e-12,1e-15])
    with device('/CPU:0'):
        test_nufft1(natoms=100,
                    nsamples=40,
                    ndim=3,
                    gridlens=[64],
                    tols=[1e-3,1e-6,1e-9,1e-12,1e-15])