# coding: utf-8
class TopHat :
    def __init__ (self, ks ,k0) :
        self.ks = ks
        self.k0 = k0

    def __call__ (self, k) :
        import numpy as np
        f = np.ones (k.shape)
        f[k < self.ks - self.k0/2] = 0 
        f[k > self.ks + self.k0/2] = 0 
        return f

class CutOff :
    def __init__ (self, kc) :
        self.kc = kc

    def __call__ (self, k) :
        import numpy as np
        f = np.ones (k.shape)
        f[k > self.kc] = self.kc**2 / k[k > self.kc]**2
        return f

class RandomScalarField2D:
    def __init__ (self, nx, lx, ny, ly, scalar_max, scalar_min, ksk0, kcks) :
        import numpy as np
	dx = lx / nx
	dy = ly / ny
        self.X, self.Y = np.meshgrid(np.arange(0,nx*dx,dx), np.arange(0,ny*dy,dy))
        self.KX, self.KY = np.meshgrid(2*np.pi*np.fft.fftfreq(nx, dx), 2*np.pi*np.fft.fftfreq(ny, dy))
        print (np.arange(0,nx*dx,dx), np.arange(0,ny*dy,dy))
        print (2*np.pi*np.fft.fftfreq(nx, dx), 2*np.pi*np.fft.fftfreq(ny, dy))
        self.k = np.sqrt(self.KX**2 + self.KY**2) + 1e-15
        self.smax = scalar_max
        self.smin = scalar_min

        self.Nmax = max(nx, ny)
        # self.k0 = self.k.max() / self.Nmax
        self.k0 = 2 * np.pi / lx
	print self.k0
        # self.k0 = min (np.fft.fftfreq(nx, dx)[1], np.fft.fftfreq(ny, dy)[1])
	# print self.k0
        self.ks = ksk0 * self.k0
        self.kc = kcks * self.ks

	print self.ks, self.kc
	print self.k

        self.f = TopHat(self.ks, self.k0)
        self.F = CutOff(self.kc)
        
        self.Phi = np.sqrt(self.f(self.k)/(4*np.pi*self.k**2)) * np.exp(2*np.pi*(1j)*np.random.uniform(0, 1, self.k.shape))

	print np.count_nonzero(self.Phi)
	print np.nonzero(self.Phi)
	print self.Phi[np.nonzero(self.Phi)]

        self.iPhi = np.fft.ifft2 (self.Phi) 
	print self.iPhi
	print np.angle(self.iPhi)
	print np.angle(self.iPhi).min()
	print np.angle(self.iPhi).max()
        self.phi = np.zeros(self.iPhi.shape)
        self.phi[np.angle(self.iPhi) >= 0 ] = self.smax
        self.phi[np.angle(self.iPhi) < 0 ] = self.smin

        self.Phi_ = self.F(self.k) * np.fft.fft2(self.phi)
        self.phi_ = np.fft.ifft2(self.Phi_)

    def energy_spectrum (self) :
        import numpy as np
        E = np.zeros (self.Nmax + 1)
        for i in np.arange(self.Nmax + 1):
            k = self.k[np.logical_and(
                    self.k >= (i-0.5)*self.k0 , self.k < (i+0.5)*self.k0 )]
            Phi_n = self.f(k)/(4*np.pi*k**2)
            E[i] = Phi_n.sum()
        E /= E.sum() 
        E += 10e-8

        E_ = np.zeros (self.Nmax + 1)
        Phi_ = np.fft.fft2(self.phi)
        for i in np.arange(self.Nmax + 1):
            Phi_n =  Phi_[np.logical_and(
                    self.k >= (i-0.5)*self.k0 , self.k < (i+0.5)*self.k0 )]
            E_[i] = (np.abs(Phi_n)**2).sum()
        E_ /= E_.sum() 
        E_ += 10e-8

        E__ = np.zeros (self.Nmax + 1)
        for i in np.arange(self.Nmax + 1):
            Phi_n =  self.Phi_[np.logical_and(
                    self.k >= (i-0.5)*self.k0 , self.k < (i+0.5)*self.k0 )]
            E__[i] = (np.abs(Phi_n)**2).sum()
        E__ /= E__.sum() 
        E__ += 10e-8

        return E, E_, E__


class RandomScalarField3D:
    def __init__ (self, nx, dx, scalar_max, scalar_min, ksk0, kcks) :
        import numpy as np
        self.X, self.Y, self.Z = np.meshgrid(np.arange(0,nx*dx,dx), np.arange(0,nx*dx,dx), np.arange(0,nx*dx,dx))
        self.KX, self.KY, self.KZ = np.meshgrid(np.fft.fftfreq(nx, dx), np.fft.fftfreq(nx, dx), np.fft.fftfreq(nx, dx))

        print (np.arange(0,nx*dx,dx))

        print (np.fft.fftfreq(nx, dx))

        self.k = np.sqrt(self.KX**2 + self.KY**2 + self.KZ**2) + 1e-15

        self.smax = scalar_max

        self.smin = scalar_min

        self.Nmax = nx

        self.k0 = self.k.max() / self.Nmax
        self.ks = ksk0 * self.k0
        self.kc = kcks * self.ks

        self.f = TopHat(self.ks, self.k0)
        self.F = CutOff(self.kc)
        
        self.Phi = np.sqrt(self.f(self.k)/(4*np.pi*self.k**2)) * np.exp(2*np.pi*(1j)*np.random.uniform(0, 1, self.k.shape))

        self.iPhi = np.fft.ifftn (self.Phi) 
        self.phi = np.zeros(self.iPhi.shape)
        self.phi[self.iPhi >= 0 ] = self.smax
        self.phi[self.iPhi < 0 ] = self.smin

        self.Phi_ = self.F(self.k) * np.fft.fftn(self.phi)
        self.phi_ = np.fft.ifftn(self.Phi_)
	'''
	'''