# -*- coding: utf-8 -*-
"""
=====================================
Testing Perfomance of an Extrapolator
=====================================

This tests the speed of it all.
"""

# General imports
from astropy import units as u
from astropy.table import Table
from mayavi import mlab
import numpy as np

# Module imports
from solarbextrapolation.extrapolators import PotentialExtrapolator
from solarbextrapolation.example_data_generator import generate_example_data, dummyDataToMap
from solarbextrapolation.visualisation_functions import visualise

# The input parameters:
lis_grid_shapes = [ [ 20, 20, 20 ] ]#, [ 20, 20, 20 ]]#, [ 30, 30, 30 ]]#, [ 100, 100, 100 ]]#[ 10, 10, 10 ],[ 50, 50, 50 ], [ 100, 100, 100 ], [ 200, 200, 200 ] ]
xrange = u.Quantity([ -10.0, 10.0 ] * u.arcsec)
yrange = u.Quantity([ -10.0, 10.0 ] * u.arcsec)
zrange = u.Quantity([ 0,     20.0 ] * u.arcsec)

# Manual Pole Details
arrA0 = [ u.Quantity([ 25, 25 ] * u.percent), 10.0 * u.percent,  0.2 * u.T ]
arrA1 = [ u.Quantity([ 75, 75 ] * u.percent), 10.0 * u.percent, -0.2 * u.T ]
arrA2 = [ u.Quantity([ 25, 75 ] * u.percent), 10.0 * u.percent,  0.1 * u.T ]
arrA3 = [ u.Quantity([ 75, 25 ] * u.percent), 10.0 * u.percent, -0.1 * u.T ]

# Generate the datasets and maps
lis_maps = []
lis_extrapolators = []

# A table for storing the data
t = Table(names=('grid size', 'time (ave)', 'time (std)'), meta={'name': 'times tables'}, dtype=('S24', 'f8', 'f8'))
t['time (ave)'].unit = u.s
t['time (std)'].unit = u.s

# A list of all the necessary parameters to perform an extrapolation
lis_datasets = []
for shape in lis_grid_shapes:
    lis_datasets.append([ str(shape), shape[2], zrange,
                          dummyDataToMap(generate_example_data(shape[0:2], xrange, yrange, arrA0, arrA1, arrA2, arrA3), xrange, yrange) ])
int_trials = 2 # The times to repeat each extrapolation.

# Iterate through the extrapolations
for extrapolation in lis_datasets:
    # Setup the extrapolator and table
    aPotExt = PotentialExtrapolator(extrapolation[3], zshape=extrapolation[1], zrange=extrapolation[2])

    # List to store the trial
    lis_times = []

    # Run the extrapolation without numba for each dataset (map and ranges).
    for i in range(0, int_trials):
        aMap3D = aPotExt.extrapolate(enable_numba=False)
        lis_times.append(aMap3D.meta['extrapolator_duration'])
    t.add_row([extrapolation[0], np.round(np.average(lis_times), 2), np.round(np.std(lis_times), 2)])

    # Run the extrapolation with numba for each dataset (map and ranges).
    for i in range(0, int_trials):
        aMap3D = aPotExt.extrapolate(enable_numba=True)
        lis_times.append(aMap3D.meta['extrapolator_duration'])
    t.add_row(['(numba)'+extrapolation[0], np.round(np.average(lis_times), 2), np.round(np.std(lis_times), 2)])

# Show the data table
print t
print '\n\n'


# Visualise
visualise(aMap3D,
          boundary=lis_datasets.pop()[3],
          volume_units=[1.0*u.arcsec, 1.0*u.arcsec, 1.0*u.Mm],
          show_boundary_axes=False,
          boundary_units=[1.0*u.arcsec, 1.0*u.arcsec],
          show_volume_axes=True,
          debug=False)
mlab.show()

@@ -258,32 +258,8 @@
 ble%0A
-from mayavi import mlab%0A
 impo
@@ -455,74 +455,8 @@
 oMap
-%0Afrom solarbextrapolation.visualisation_functions import visualise
 %0A%0A# 
@@ -2752,298 +2752,4 @@
 '%0A%0A%0A
-# Visualise%0Avisualise(aMap3D,%0A          boundary=lis_datasets.pop()%5B3%5D,%0A          volume_units=%5B1.0*u.arcsec, 1.0*u.arcsec, 1.0*u.Mm%5D,%0A          show_boundary_axes=False,%0A          boundary_units=%5B1.0*u.arcsec, 1.0*u.arcsec%5D,%0A          show_volume_axes=True,%0A          debug=False)%0Amlab.show()%0A

remove all mayavi imports from the performance test example