"""
Fortuna

Python project to visualize uncertatinty in probabilistic exploration models.

Created on 09/06/2018

@authors: Natalia Shchukina, Graham Brew, Marco van Veen, Behrooz Bashokooh, Tobias Stål, Robert Leckenby
"""

# Import libraries

import numpy as np
import glob
from matplotlib import pyplot as plt
import pandas as pd
import xarray as xr
import pyproj as proj
from scipy.stats import norm



class Fortuna(object):
    """
    Class to load the fortuna dataset and call different methods for visualization in a web frontend.

    Args:
        There are no required arguments at the moment. Input files could be defined.
    """



    def __init__(self, **kwargs):
        """
        Method that is called when a object of the class Fortuna is initiated, it imports the data and directly creates some important variables.
        """

        # hardcode geometry
        self.size_raster = (250,162)
        self.X_corner = 390885
        self.Y_corner = 7156947
        self.dx, self.dy, self.dz = 25, 25, 100
        self.top_model = 950
        self.bottom_model = 1050


        self.base_cube = None
        self.top_cube = None
        self.base_n = None
        self.top_n = None
        self.vol = None

        # Create empty xarray dataset
        self.ds = xr.Dataset()

        self.xx = None
        self.yy = None
        self.zz = None
        self.model = None
        self.base_mean = None
        self.base_std = None
        self.top_mean = None
        self.top_std = None


        ## Initial methods to load
        self.import_data()
        self.calc_xarray()
        self.calc_stat()




    ### Methods for initiating the object

    def folder2cube(self, files):
        """
        Method to read a file.
        """
        base_set = glob.glob(files)
        cube = np.zeros(self.size_raster + (len(base_set),))
        for i, model in enumerate(base_set):
            cube[:, :, i] = np.loadtxt(model, skiprows=1).reshape(self.size_raster)
        return cube, len(base_set)


    def import_data(self):
        """
        Method to load different data objects from files.
        """
        self.base_cube, self.base_n = self.folder2cube('data/Hackaton/BaseSet/MapSimu__*.data')
        self.top_cube, self.top_n = self.folder2cube('data/Hackaton/TopSet/MapSimu__*.data')

        self.vol = pd.read_csv('data/Hackaton/VolumeDistribution/Volumes', delim_whitespace=True)


    def calc_xarray (self):
        self.xx = np.linspace(self.X_corner, self.X_corner + self.size_raster[0] * self.dx, self.size_raster[0])
        self.yy = np.linspace(self.Y_corner, self.Y_corner + self.size_raster[1] * self.dy, self.size_raster[1])
        self.zz = np.linspace(self.top_model, self.bottom_model, self.dz)

        self.model = np.linspace(0, self.top_model, self.base_n)

        self.ds.coords['X'] = self.xx
        self.ds.coords['Y'] = self.yy
        self.ds.coords['Z'] = self.zz
        self.ds.coords['MODEL'] = self.model

        self.ds['BASE'] = (('X', 'Y', 'MODEL'), self.base_cube)
        self.ds['TOP'] = (('X', 'Y', 'MODEL'), self.top_cube)


    def calc_stat (self):
        self.base_mean = self.ds['BASE'].mean(dim='MODEL')
        self.base_std = self.ds['BASE'].std(dim='MODEL')

        self.top_mean = self.ds['TOP'].mean(dim='MODEL')
        self.top_std = self.ds['TOP'].std(dim='MODEL')



    ## Data Management methods

    def load_pickle(self, path):
        return np.load(path)


    ## Methods to compute different uncertatinty cubes --> cubes to be displayed in the frontend

    def calc_lithology(self, iterations = 2):
        """
        Sample from both distributions and fill each z-stack accordingly
        """
        # create empty array
        block = np.zeros((iterations, self.size_raster[0], self.size_raster[1], self.zz.size), dtype='int8')

        for i in range(iterations):
            for j in range(self.size_raster[0]):  # size_raster[0]
                for k in range(self.size_raster[1]):

                    # sample from top and base distributions for specific x,y position
                    top = np.random.normal(self.top_mean[j, k], self.top_std[j, k])
                    base = np.random.normal(self.base_mean[j, k], self.base_std[j, k])

                    # iterate over vertical z-stack
                    for l in range(self.zz.size):

                        if self.zz[l] <= top:
                            block[i, j, k, l] = 1
                        elif self.zz[l] > base:
                            block[i, j, k, l] = 3
                        elif ((self.zz[l] > top) and (l <= base)):
                            block[i, j, k, l] = 2

        return block


    def calc_lithology_vect(self, iterations=2):
        """
        Resample from z value statistics and fill each z-stack in a lithology block accordingly.
        This is the new method with vectorized operations to speed up calculations.
        """

        # create empty array
        block = np.zeros((iterations, self.xx.size, self.yy.size, self.zz.size), dtype='int8')

        for i in range(iterations):

            # create meshgrids grid for coordinate-wise iterations
            mesh_x, mesh_y, mesh_z = np.meshgrid(np.arange(self.xx.size),
                                                 np.arange(self.yy.size),
                                                 np.arange(self.zz.size))

            # sample from top and base distributions for specific x,y position
            top = np.zeros([self.xx.size, self.yy.size])
            base = np.zeros([self.xx.size, self.yy.size])

            top[mesh_x, mesh_y] = np.random.normal(self.top_mean.values[mesh_x, mesh_y],
                                                   self.top_std.values[mesh_x, mesh_y])
            base[mesh_x, mesh_y] = np.random.normal(self.top_mean.values[mesh_x, mesh_y],
                                                    self.top_std.values[mesh_x, mesh_y])

            # compare each cell to resampled reference values
            # TODO generalize for any number of lithologies
            block[i, mesh_x, mesh_y, mesh_z] = np.where(self.zz < top[mesh_x, mesh_y], 1,
                                                        np.where(self.zz < base[mesh_x, mesh_y], 2, 3))

        return block


    ### Modifyed from GemPy!
    def calc_probability_lithology(self, cube):
        """Blocks must be just the lith blocks!"""
        lith_blocks = cube.reshape([cube.shape[0], (self.xx.size * self.yy.size * self.zz.size)])
        lith_id = np.unique(lith_blocks)
        # lith_count = np.zeros_like(lith_blocks[0:len(lith_id)])
        lith_count = np.zeros((len(np.unique(lith_blocks)), lith_blocks.shape[1]))
        for i, l_id in enumerate(lith_id):
            lith_count[i] = np.sum(lith_blocks == l_id, axis=0)
        lith_prob = lith_count / len(lith_blocks)
        return lith_prob


    ### Modyfied from GemPy!
    def calc_information_entropy(self, lith_prob):
        """Calculates information entropy for the given probability array."""
        cube = np.zeros_like(lith_prob[0])
        for l in lith_prob:
            pm = np.ma.masked_equal(l, 0)  # mask where layer prob is 0
            cube -= (pm * np.ma.log2(pm)).filled(0)
        return cube.reshape([self.xx.size, self.yy.size, self.zz.size])
        # Try numpy.flatten and numpy.ravel

    ## Simple plotting methods

    def plot_entropy(self, cube, slice=10):
        plt.imshow(cube[slice, :, :].T, origin='upperleft', cmap='viridis')
        plt.show()
[((54, 18, 54, 30), 'xarray.Dataset', 'xr.Dataset', ({}, {}), '()', True, 'import xarray as xr\n'), ((80, 19, 80, 35), 'glob.glob', 'glob.glob', ({(80, 29, 80, 34): 'files'}, {}), '(files)', False, 'import glob\n'), ((94, 19, 94, 97), 'pandas.read_csv', 'pd.read_csv', (), '', True, 'import pandas as pd\n'), ((98, 18, 98, 112), 'numpy.linspace', 'np.linspace', ({(98, 30, 98, 43): 'self.X_corner', (98, 45, 98, 90): 'self.X_corner + self.size_raster[0] * self.dx', (98, 92, 98, 111): 'self.size_raster[0]'}, {}), '(self.X_corner, self.X_corner + self.size_raster[0] * self.dx,\n    self.size_raster[0])', True, 'import numpy as np\n'), ((99, 18, 99, 112), 'numpy.linspace', 'np.linspace', ({(99, 30, 99, 43): 'self.Y_corner', (99, 45, 99, 90): 'self.Y_corner + self.size_raster[1] * self.dy', (99, 92, 99, 111): 'self.size_raster[1]'}, {}), '(self.Y_corner, self.Y_corner + self.size_raster[1] * self.dy,\n    self.size_raster[1])', True, 'import numpy as np\n'), ((100, 18, 100, 73), 'numpy.linspace', 'np.linspace', ({(100, 30, 100, 44): 'self.top_model', (100, 46, 100, 63): 'self.bottom_model', (100, 65, 100, 72): 'self.dz'}, {}), '(self.top_model, self.bottom_model, self.dz)', True, 'import numpy as np\n'), ((102, 21, 102, 64), 'numpy.linspace', 'np.linspace', ({(102, 33, 102, 34): '0', (102, 36, 102, 50): 'self.top_model', (102, 52, 102, 63): 'self.base_n'}, {}), '(0, self.top_model, self.base_n)', True, 'import numpy as np\n'), ((125, 15, 125, 28), 'numpy.load', 'np.load', ({(125, 23, 125, 27): 'path'}, {}), '(path)', True, 'import numpy as np\n'), ((135, 16, 135, 108), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((165, 16, 165, 94), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((195, 18, 195, 40), 'numpy.unique', 'np.unique', ({(195, 28, 195, 39): 'lith_blocks'}, {}), '(lith_blocks)', True, 'import numpy as np\n'), ((207, 15, 207, 42), 'numpy.zeros_like', 'np.zeros_like', ({(207, 29, 207, 41): 'lith_prob[0]'}, {}), '(lith_prob[0])', True, 'import numpy as np\n'), ((217, 8, 217, 75), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'from matplotlib import pyplot as plt\n'), ((218, 8, 218, 18), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'from matplotlib import pyplot as plt\n'), ((175, 18, 175, 56), 'numpy.zeros', 'np.zeros', ({(175, 27, 175, 55): '[self.xx.size, self.yy.size]'}, {}), '([self.xx.size, self.yy.size])', True, 'import numpy as np\n'), ((176, 19, 176, 57), 'numpy.zeros', 'np.zeros', ({(176, 28, 176, 56): '[self.xx.size, self.yy.size]'}, {}), '([self.xx.size, self.yy.size])', True, 'import numpy as np\n'), ((178, 34, 179, 87), 'numpy.random.normal', 'np.random.normal', ({(178, 51, 178, 87): 'self.top_mean.values[mesh_x, mesh_y]', (179, 51, 179, 86): 'self.top_std.values[mesh_x, mesh_y]'}, {}), '(self.top_mean.values[mesh_x, mesh_y], self.top_std.values[\n    mesh_x, mesh_y])', True, 'import numpy as np\n'), ((180, 35, 181, 88), 'numpy.random.normal', 'np.random.normal', ({(180, 52, 180, 88): 'self.top_mean.values[mesh_x, mesh_y]', (181, 52, 181, 87): 'self.top_std.values[mesh_x, mesh_y]'}, {}), '(self.top_mean.values[mesh_x, mesh_y], self.top_std.values[\n    mesh_x, mesh_y])', True, 'import numpy as np\n'), ((199, 28, 199, 63), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((209, 17, 209, 41), 'numpy.ma.masked_equal', 'np.ma.masked_equal', ({(209, 36, 209, 37): 'l', (209, 39, 209, 40): '0'}, {}), '(l, 0)', True, 'import numpy as np\n'), ((170, 49, 170, 72), 'numpy.arange', 'np.arange', ({(170, 59, 170, 71): 'self.xx.size'}, {}), '(self.xx.size)', True, 'import numpy as np\n'), ((171, 49, 171, 72), 'numpy.arange', 'np.arange', ({(171, 59, 171, 71): 'self.yy.size'}, {}), '(self.yy.size)', True, 'import numpy as np\n'), ((172, 49, 172, 72), 'numpy.arange', 'np.arange', ({(172, 59, 172, 71): 'self.zz.size'}, {}), '(self.zz.size)', True, 'import numpy as np\n'), ((186, 56, 186, 102), 'numpy.where', 'np.where', ({(186, 65, 186, 95): 'self.zz < base[mesh_x, mesh_y]', (186, 97, 186, 98): '2', (186, 100, 186, 101): '3'}, {}), '(self.zz < base[mesh_x, mesh_y], 2, 3)', True, 'import numpy as np\n'), ((83, 28, 83, 57), 'numpy.loadtxt', 'np.loadtxt', (), '', True, 'import numpy as np\n'), ((142, 26, 142, 83), 'numpy.random.normal', 'np.random.normal', ({(142, 43, 142, 62): 'self.top_mean[j, k]', (142, 64, 142, 82): 'self.top_std[j, k]'}, {}), '(self.top_mean[j, k], self.top_std[j, k])', True, 'import numpy as np\n'), ((143, 27, 143, 86), 'numpy.random.normal', 'np.random.normal', ({(143, 44, 143, 64): 'self.base_mean[j, k]', (143, 66, 143, 85): 'self.base_std[j, k]'}, {}), '(self.base_mean[j, k], self.base_std[j, k])', True, 'import numpy as np\n'), ((197, 35, 197, 57), 'numpy.unique', 'np.unique', ({(197, 45, 197, 56): 'lith_blocks'}, {}), '(lith_blocks)', True, 'import numpy as np\n'), ((210, 26, 210, 40), 'numpy.ma.log2', 'np.ma.log2', ({(210, 37, 210, 39): 'pm'}, {}), '(pm)', True, 'import numpy as np\n')]