# Importing
import theano.tensor as T
import sys, os
sys.path.append("../GeMpy")

# Importing GeMpy modules
import GeMpy_core
import Visualization

# Reloading (only for development purposes)
import importlib
importlib.reload(GeMpy_core)
importlib.reload(Visualization)

# Usuful packages
import numpy as np
import pandas as pn

import matplotlib.pyplot as plt

# This was to choose the gpu
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

# Default options of printin
np.set_printoptions(precision = 6, linewidth= 130, suppress =  True)

%matplotlib inline
#%matplotlib notebook

"""
Explanation: Example 1: Sandstone Model
End of explanation
"""


# Setting extend, grid and compile
# Setting the extent
sandstone = GeMpy_core.GeMpy()

# Create Data class with raw data
sandstone.import_data( [696000,747000,6863000,6950000,-20000, 2000],[ 40, 40, 80],
                         path_f = os.pardir+"/input_data/a_Foliations.csv",
                         path_i = os.pardir+"/input_data/a_Points.csv")

"""
Explanation: First we make a GeMpy instance with most of the parameters default (except range that is given by the project). Then we also fix the extension and the resolution of the domain we want to interpolate. Finally we compile the function, only needed once every time we open the project (the guys of theano they are working on letting loading compiled files, even though in our case it is not a big deal).
General note. So far the reescaling factor is calculated for all series at the same time. GeoModeller does it individually for every potential field. I have to look better what this parameter exactly means
End of explanation
"""


sandstone.Data.Foliations.head()

"""
Explanation: All input data is stored in pandas dataframes under, self.Data.Interances and self.Data.Foliations:
End of explanation
"""


sandstone.Data.set_series({"EarlyGranite_Series":sandstone.Data.formations[-1], 
                      "BIF_Series":(sandstone.Data.formations[0], sandstone.Data.formations[1]),
                      "SimpleMafic_Series":sandstone.Data.formations[2]}, 
                       order = ["EarlyGranite_Series",
                              "BIF_Series",
                              "SimpleMafic_Series"]) 

"""
Explanation: In case of disconformities, we can define which formation belong to which series using a dictionary. Until python 3.6 is important to specify the order of the series otherwise is random
End of explanation
"""


sandstone.Data.Foliations.head()

"""
Explanation: Now in the data frame we should have the series column too
End of explanation
"""


# Create a class Grid so far just regular grid
sandstone.create_grid()
sandstone.Grid.grid

"""
Explanation: Next step is the creating of a grid. So far only regular. By default it takes the extent and the resolution given in the import_data method.
End of explanation
"""


sandstone.Plot.plot_data(series = sandstone.Data.series.columns.values[1])

"""
Explanation: Plotting raw data
The object Plot is created automatically as we call the methods above. This object contains some methods to plot the data and the results.
It is possible to plot a 2D projection of the data in a specific direction using the following method. Also is possible to choose the series you want to plot. Additionally all the key arguments of seaborn lmplot can be used.
End of explanation
"""


sandstone.set_interpolator()

"""
Explanation: Class Interpolator
This class will take the data from the class Data and calculate potential fields and block. We can pass as key arguments all the variables of the interpolation. I recommend not to touch them if you do not know what are you doing. The default values should be good enough. Also the first time we execute the method, we will compile the theano function so it can take a bit of time.
End of explanation
"""


sandstone.Plot.plot_potential_field(10, n_pf=0)

"""
Explanation: Now we could visualize the individual potential fields as follow:
Early granite
End of explanation
"""


sandstone.Plot.plot_potential_field(13, n_pf=1, cmap = "magma",  plot_data = True,
                                        verbose = 5 )

"""
Explanation: BIF Series
End of explanation
"""


sandstone.Plot.plot_potential_field(10, n_pf=2)

"""
Explanation: SImple mafic
End of explanation
"""


# Reset the block
sandstone.Interpolator.block.set_value(np.zeros_like(sandstone.Grid.grid[:,0]))

# Compute the block
sandstone.Interpolator.compute_block_model([0,1,2], verbose = 0)

sandstone.Interpolator.block.get_value(), np.unique(sandstone.Interpolator.block.get_value())

"""
Explanation: Optimizing the export of lithologies
But usually the final result we want to get is the final block. The method compute_block_model will compute the block model, updating the attribute block. This attribute is a theano shared function that can return a 3D array (raveled) using the method get_value().
End of explanation
"""


sandstone.Plot.plot_block_section(13, interpolation = 'nearest',  direction='y')
plt.savefig("sandstone_example.png")

"""
Explanation: And again after computing the model in the Plot object we can use the method plot_block_section to see a 2D section of the model
End of explanation
"""


"""Export model to VTK

Export the geology blocks to VTK for visualisation of the entire 3-D model in an
external VTK viewer, e.g. Paraview.

..Note:: Requires pyevtk, available for free on: https://github.com/firedrakeproject/firedrake/tree/master/python/evtk

**Optional keywords**:
    - *vtk_filename* = string : filename of VTK file (default: output_name)
    - *data* = np.array : data array to export to VKT (default: entire block model)
"""
vtk_filename = "noddyFunct2"

extent_x = 10
extent_y = 10
extent_z = 10

delx = 0.2
dely = 0.2
delz = 0.2
from pyevtk.hl import gridToVTK
# Coordinates
x = np.arange(0, extent_x + 0.1*delx, delx, dtype='float64')
y = np.arange(0, extent_y + 0.1*dely, dely, dtype='float64')
z = np.arange(0, extent_z + 0.1*delz, delz, dtype='float64')

# self.block = np.swapaxes(self.block, 0, 2)
gridToVTK(vtk_filename, x, y, z, cellData = {"geology" : sol})

"""
Explanation: Export  to vtk. (Under development)
End of explanation
"""


%%timeit
# Reset the block
sandstone.Interpolator.block.set_value(np.zeros_like(sandstone.Grid.grid[:,0]))

# Compute the block
sandstone.Interpolator.compute_block_model([0,1,2], verbose = 0)

"""
Explanation: Performance Analysis
One of the advantages of theano is the posibility to create a full profile of the function. This has to be included in at the time of the creation of the function. At the moment it should be active (the downside is larger compilation time and I think also a bit in the computation so be careful if you need a fast call)
CPU
The following profile is with a 2 core laptop. Nothing spectacular.
End of explanation
"""


esandstone.Interpolator._interpolate.profile.summary()

"""
Explanation: Looking at the profile we can see that most of time is in pow operation (exponential). This probably is that the extent is huge and we are doing it with too much precision. I am working on it
End of explanation
"""


%%timeit
# Reset the block
sandstone.block.set_value(np.zeros_like(sandstone.grid[:,0]))

# Compute the block
sandstone.compute_block_model([0,1,2], verbose = 0)

sandstone.block_export.profile.summary()

"""
Explanation: GPU
End of explanation
"""


