#!/usr/bin/env python

"""Demonstrate how to solve Poisson's equation using FFT"""


from numpy import *
from numpy.fft import fft, ifft, fftfreq
import Gnuplot, Gnuplot.funcutils
import time

#some global definitions
k = 2.
N = 512
L = 10.


#Set up a charge distribution
charge = zeros(N)
x = linspace(0.,L,N)
#charge = sin(k*2.*pi/L*x)# + 0.05*random.standard_normal(N)
charge = zeros(N)
charge[N/4:N/3] = 1.
charge[2*N/3:3*N/4] = -1.
#charge[3*N/4:N] = -1.

#plot charge distribution
g = Gnuplot.Gnuplot()
#g('set data style line')
charge_data = Gnuplot.Data(x,charge,with_='line')
g.plot(charge_data,title='charge distribution')
raw_input('Press return to continue!')


#FFT charge distribution
fourier = fft(charge)
step = L/N
freq = L*fftfreq(N, d=step)
freq[0] = 0.01

#g.plot(freq,title='frequencies')
#raw_input('Press return to continue')

#plot FFT of charge distribution
dat = Gnuplot.Data(freq,abs(fourier)**2,with_='line')
g('set logscale y')
g.plot(dat,title='fft of charge distribution')
raw_input('Press return to continue!')

#plot IFFT of charge distribution
fil_q = ifft(fourier)
fil_dat = Gnuplot.Data(x,fil_q,with_='line')
g('unset logscale y')
g.plot(fil_dat,charge_data,title='filtered and original charge distribibution')
raw_input('Press return to continue!')


#Divide FFT of charge distribution by k**2
#that is equal to the fourier transform of the potential, fourier_pot
fourier_pot = zeros(N)
#ofreq = freq.copy()
#kappa = sin(freq/step)/(freq/step)
fourier_pot = fourier/freq/freq
#fourier_pot[0:N] = fourier[0:N]/freq[0:N]/freq[0:N]
#print freq,ofreq
dat_fp = Gnuplot.Data(freq,abs(fourier_pot)**2,with_='line')
g.plot(dat_fp,title='fft of potential')
raw_input('Press return to continue!')



#inverse transform yields potential in x space
pot = real_if_close(ifft(fourier_pot))
dat_pot = Gnuplot.Data(x,pot,with_='line')
g('unset logscale y')
g.plot(dat_pot,charge_data,title='Potential and charge distribution')
raw_input('Press return to continue!')



#differentiate once to obtain electric field
field = zeros(N)
field[1:N-2] = pot[0:N-3] - pot[2:N-1]
field[0] = pot[N-1] - pot[1] #for periodic boundary conditions
field[N-1] = pot[N-2] - pot[0] #for periodic boundary conditions
field /= 2.*step
dat_field = Gnuplot.Data(x,field,with_='line')
g.plot(dat_field,charge_data,title='field')
raw_input('Press return to continue!')