#!/usr/bin/env python
""" this script provides basic utilities for testing the effect of translation of simulated 2d vortex data
on the wavelet coeffs for a 2D dwt"""
from pylab import *
from pywt import *
from dwt3 import *
from numpy.linalg import norm
from time import sleep
img = complex(0,1)


#makes freq array with 5/3 power law and random phase<- need proper 2d pawer law
#def mkarray2(n,p=-1.0):
#  A = zeros((n,n),dtype=complex64)
#  for i in range(n/2):
#    for j in range(n/2):
#       A[i,j] = complex((1+i**2+j**2)**(p/2.0))*e**(2*pi*img*uniform())
#       A[n-i-1,n-j-1] = conj(A[i,j])
#       A[n-i-1,j] = complex((1+i**2+j**2)**(p/2.0))*e**(2*pi*img*uniform())
#       A[i,n-j-1] = conj(A[i+n/2,j])     
#  return A


def mkvort(n,p=.9,size=5):
  A = zeros((n,n))
  for i in range(n):
    for j in range(n):
      if p < uniform():
        A[i,j] = 5.0*randn()
  return A
        
def prntstats(coeffs):
  for lev in coeffs:
    data = 0
    for j in lev:
      data = data + mean(abs(j))
      print max
    m.append(data)
  return m

def prnt(coeffs):
  i = 0
  for lev in coeffs:
    print 'lev = ' + str(i)
    j = 0
    for arry in lev:  
      print 'dim  = ' + str(j), arry.round(),'\n'
      j = j+1
    print '\n'
    i = i+1
      
def plt(coeffs,lev,dim):
  plot((coeffs[lev][dim]).flatten())
     
def pltall(coeffs1,coeffs2):
  i = 0
  clf()
  for lev in range(len(coeffs1)):
    print 'lev = ' + str(i)
    j = 0
    for dim in range(len(coeffs1[lev])):  
      print 'dim  = ' + str(j),'\n'
      j = j+1
      plot((coeffs1[lev][dim]).flatten())
      plot((coeffs2[lev][dim]).flatten())
      raw_input()
      clf()
    print '\n'
    i = i+1    
      
def prntstats(coeffs):
  i = 0
  for lev in coeffs:
    print 'lev = ' + str(i)
    j = 0
    for arry in lev:  
      print 'dim  = ' + str(j)
      j = j+1
      print numpy.max(abs(arry)), mean(abs(arry))
    print '\n'
    i = i+1
    

def sim(n,wvname='db2'):
  a = mkvort(n)
  coeffs1 = wavedec2(a, wvname,mode= 'per')
  a = fftshift(a) #shift data
  coeffs2 = wavedec2(a, wvname,mode='per') 
  return coeffs1,coeffs2
  
  
def test(n,wvname='db2'):
  c,d = sim(n,wvname='db2')
  print 'unshifted stats:', prntstats(c)
  print 'shifted stats:', prntstats(d)
  print 'plotting level 0 coeffs'
  plt(c,0,0)
  plt(d,0,0)
  print 'type return when done'
  raw_input()
  
  
if __name__ == '__main__':
  test(30)

