import keras
from keras.models import Sequential
from keras.layers.convolutional import Convolution2D, MaxPooling2D
from keras.layers.advanced_activations import LeakyReLU
from keras.layers.core import Flatten, Dense, Lambda

from collections import deque
import numpy as np
import cv2
import glob
import matplotlib.pyplot as plt
from IPython.display import HTML
from moviepy.editor import VideoFileClip

# matplotlib inline

from utils import load_weights, process_output, draw_boxes, preprocess, remove_duplicates

keras.backend.set_image_dim_ordering('th')

model = Sequential()

#normalize data
model.add(Lambda(lambda x: x/128 - 1,input_shape=(3,448,448)))

#layer 1
model.add(Convolution2D(16, 3, 3,border_mode='same',subsample=(1,1)))
model.add(LeakyReLU(alpha=0.1))
model.add(MaxPooling2D(pool_size=(2, 2)))

#layer 2
model.add(Convolution2D(32,3,3 ,border_mode='same'))
model.add(LeakyReLU(alpha=0.1))
model.add(MaxPooling2D(pool_size=(2, 2),border_mode='valid'))

#layer 3
model.add(Convolution2D(64,3,3 ,border_mode='same'))
model.add(LeakyReLU(alpha=0.1))
model.add(MaxPooling2D(pool_size=(2, 2),border_mode='valid'))

#layer 4
model.add(Convolution2D(128,3,3 ,border_mode='same'))
model.add(LeakyReLU(alpha=0.1))
model.add(MaxPooling2D(pool_size=(2, 2),border_mode='valid'))

#layer 5
model.add(Convolution2D(256,3,3 ,border_mode='same'))
model.add(LeakyReLU(alpha=0.1))
model.add(MaxPooling2D(pool_size=(2, 2),border_mode='valid'))

#layer 6
model.add(Convolution2D(512,3,3 ,border_mode='same'))
model.add(LeakyReLU(alpha=0.1))
model.add(MaxPooling2D(pool_size=(2, 2),border_mode='valid'))

#layer 7
model.add(Convolution2D(1024,3,3 ,border_mode='same'))
model.add(LeakyReLU(alpha=0.1))

#layer 8
model.add(Convolution2D(1024,3,3 ,border_mode='same'))
model.add(LeakyReLU(alpha=0.1))

#layer 9
model.add(Convolution2D(1024,3,3 ,border_mode='same'))
model.add(LeakyReLU(alpha=0.1))
model.add(Flatten())

#layer 10
model.add(Dense(256))

#layer 11
model.add(Dense(4096))
model.add(LeakyReLU(alpha=0.1))

#output layer
model.add(Dense(1470))

model.summary()

load_weights(model,'./yolo-tiny.weights')

test_image = cv2.imread('test_images/test3.jpg')
test_image = cv2.cvtColor(test_image, cv2.COLOR_BGR2RGB)

plt.imshow(test_image[200:600, 400:1200])
test_image.shape

processed, padhw, shavedim, resized = preprocess(test_image, shave=True)
plt.imshow(resized)
padhw

prediction = model.predict(np.array([processed]))[0]
boxes = process_output(prediction, padhw=padhw, shaved=True)
len(boxes)

boxes = remove_duplicates(boxes, test_image)
img =  draw_boxes(boxes, test_image)
plt.figure(figsize=[10,10])
plt.imshow(img)

class Pipeline:
    def __init__(self):
        self.boxes = []
        self.dropped = 0
        self.history = deque(maxlen=8)
        self.first_frames = True

    def apply_threshold(self, boxes):
        if len(boxes) == 0 and len(self.history) > 0:
            self.history.popleft()
        else:
            # create blank heatmap to store the most recent boxes
            heatmap = np.zeros([720, 1280], np.int)
            for box in boxes:
                # Add += 1 for all pixels inside each box
                heatmap[box.y1:box.y2, box.x1:box.x2] += 1

            # add current heatmap to deque of heatmaps
            self.history.append(heatmap)

            # make sure to only use heatmap detection if there is enough information
            if len(self.history) > 4:
                self.first_frames = False
            else:
                self.first_frames = True

            # add heatmaps together to make full heatmap
            if not self.first_frames:
                full_heatmap = np.zeros([720, 1280], np.int)
                for preheat in self.history:
                    full_heatmap = np.add(full_heatmap, preheat)

                # if the center point of a box is not in an area the heatmap recognizes, remove that box
                new_boxes = []
                for box in boxes:
                    if full_heatmap[int(box.y), int(box.x)] > 2:
                        new_boxes.append(box)
                return new_boxes

        return boxes

    def pipeline(self, img):
        # each detection is ordered in terms of priority. A detection higher on the list will override a similar
        # detection lower down.
        # initial detection preformed on entire image.
        transposed, padhw, shavedim, resized = preprocess(img, shave=False)
        yolo_output = model.predict(np.array([transposed]))[0]
        boxes = process_output(yolo_output, threshold=0.20, padhw=padhw, shaved=False, shavedim=shavedim)

        # preform detection on each box in the stored previous boxes
        for box in self.boxes:
            offset = 150
            shavedim = [box.y1 - offset, box.y2 + offset, box.x1 - offset, box.x2 + offset]
            transposed, padhw, shavedim, resized = preprocess(img, shave=True, shavedim=shavedim)
            yolo_output = model.predict(np.array([transposed]))[0]
            boxes2 = process_output(yolo_output, threshold=0.35, padhw=padhw, shaved=True, shavedim=shavedim)
            boxes += boxes2

            # last detection preformed picture shaved to just the highway ahead
        transposed, padhw, shavedim, resized = preprocess(img, shave=True, shavedim=(350, 500, 500, 1000))
        yolo_output = model.predict(np.array([transposed]))[0]
        boxes2 = process_output(yolo_output, threshold=0.30, padhw=padhw, shaved=True, shavedim=shavedim)
        boxes += boxes2

        # remove duplacate boxes
        boxes = remove_duplicates(boxes, img)
        boxes = self.apply_threshold(boxes)
        if (len(boxes) < len(self.boxes)) and self.dropped < 3:
            self.dropped += 1
            boxes = self.boxes
        else:
            self.dropped = 0

        drawn = draw_boxes(boxes, img)
        self.boxes = boxes
        return drawn

    pipe = Pipeline()
    test_image2 = cv2.imread('test_images/test1.jpg')
    test_image2 = cv2.cvtColor(test_image2, cv2.COLOR_BGR2RGB)
    test_image2 = pipe.pipeline(test_image2)
    plt.imshow(test_image2)
    plt.figure()
    test_image3 = cv2.imread('test_images/test5.jpg')
    test_image3 = cv2.cvtColor(test_image3, cv2.COLOR_BGR2RGB)
    test_image3 = pipe.pipeline(test_image3)
    plt.imshow(test_image3)


__author__ = "V.A. Sole - ESRF Data Analysis"
__contact__ = "sole@esrf.fr"
__license__ = "MIT"
__copyright__ = "European Synchrotron Radiation Facility, Grenoble, France"

import sys
import os
import numpy
import logging
from PyMca5 import StackPluginBase
from PyMca5.PyMcaGui import PyMcaQt as qt
from PyMca5.PyMcaGui import FFTAlignmentWindow
from PyMca5.PyMcaMath import ImageRegistration
from PyMca5.PyMcaMath.fitting import SpecfitFuns
from PyMca5.PyMcaGui.misc import CalculationThread
from PyMca5.PyMcaIO import ArraySave
from PyMca5.PyMcaGui.io import PyMcaFileDialogs
from PyMca5.PyMcaIO import specfilewrapper
from PyMca5.PyMcaGui import HDF5Widget

_logger = logging.getLogger(__name__)

try:
    from PyMca5.PyMcaGui.math import SIFTAlignmentWindow
    sift = SIFTAlignmentWindow.sift
    ocl = SIFTAlignmentWindow.silx.opencl.ocl
    SIFT = True
except Exception:
    _logger.info("SIFTAlignmentWindow not successful")
    SIFT = False

try:
    import h5py
    HDF5 = True
except Exception:
    HDF5 = False

class ImageAlignmentStackPlugin(StackPluginBase.StackPluginBase):
    def __init__(self, stackWindow, **kw):
        if _logger.getEffectiveLevel() == logging.DEBUG:
            StackPluginBase.pluginBaseLogger.setLevel(logging.DEBUG)
        StackPluginBase.StackPluginBase.__init__(self, stackWindow, **kw)
        self.methodDict = {'FFT Alignment':[self._fftAlignment,
                                            "Align using FFT",
                                            None]}
        self.__methodKeys = ['FFT Alignment']
        if SIFT:
            key = 'SIFT Alignment'
            self.methodDict[key] = [self._siftAlignment,
                                    "Align using SIFT Algorithm",
                                    None]
            self.__methodKeys.append(key)
        key = 'From File Alignment'
        self.methodDict[key] = [self._shiftFromFile,
                                "Align using shifts from file",
                                None]
        self.__methodKeys.append(key)
        self.widget = None

    def stackUpdated(self):
        self.widget = None

    #Methods implemented by the plugin
    def getMethods(self):
        return self.__methodKeys

    def getMethodToolTip(self, name):
        return self.methodDict[name][1]

    def getMethodPixmap(self, name):
        return self.methodDict[name][2]

    def applyMethod(self, name):
        return self.methodDict[name][0]()

    def _fftAlignment(self):
        stack = self.getStackDataObject()
        if stack is None:
            return
        mcaIndex = stack.info.get('McaIndex')
        if not (mcaIndex in [0, -1, 2]):
            raise IndexError("1D index must be 0, 2, or -1")

        if self.widget is None:
            self.widget = FFTAlignmentWindow.FFTAlignmentDialog()
        self.widget.setStack(stack)
        ret = self.widget.exec()
        if ret:
            ddict = self.widget.getParameters()
            self.widget.setDummyStack()
            offsets = [ddict['Dim 0']['offset'], ddict['Dim 1']['offset']]
            widths = [ddict['Dim 0']['width'], ddict['Dim 1']['width']]
            if mcaIndex == 0:
                reference = stack.data[ddict['reference_index']]
            else:
                reference = ddict['reference_image']
            crop = False
            if ddict['file_use']:
                filename = ddict['file_name']
            else:
                filename = None
            if filename is not None:
                self.__hdf5 = self.initializeHDF5File(filename)

            if _logger.getEffectiveLevel() == logging.DEBUG:
                shifts = self.calculateShiftsFFT(stack,
                                                 reference,
                                                 offsets=offsets,
                                                 widths=widths,
                                                 crop=crop)
                result = self.shiftStack(stack,
                                         shifts,
                                         crop=crop,
                                         filename=filename)
            else:
                result = self.__calculateShiftsFFT(stack,
                                                   reference,
                                                   offsets=offsets,
                                                   widths=widths,
                                                   crop=crop)
                # result[0] contains the string "Exception" in case
                # of error. However, direct comparison will raise an
                # error
                if isinstance(result[0], str):
                    if result[0] == 'Exception':
                        # exception occurred
                        raise Exception(result[1], result[2], result[3])
                else:
                    shifts = result
                result = self.__shiftStack(stack,
                                           shifts,
                                           crop=crop,
                                           filename=filename)
                if result is not None:
                    # exception occurred
                    raise Exception(result[1], result[2], result[3])
            if filename is not None:
                hdf = self.__hdf5
                alignmentGroup = hdf['/entry_000/Alignment']
                outputShifts = self.getHDF5BufferIntoGroup(alignmentGroup,
                                                     shape=(stack.data.shape[mcaIndex], 2),
                                                     name="shifts",
                                                     dtype=numpy.float32)
                outputShifts[:,:] = shifts
                attributes={'interpretation':'image'}
                referenceFrame = self.getHDF5BufferIntoGroup(alignmentGroup,
                                                     shape=reference.shape,
                                                     name="reference_frame",
                                                     dtype=numpy.float32,
                                                     attributes=attributes)
                referenceFrame[:,:] = reference[:,:]
                maskFrame = self.getHDF5BufferIntoGroup(alignmentGroup,
                                                   shape=reference.shape,
                                                   name="reference_mask",
                                                   dtype=numpy.uint8,
                                                   attributes=attributes)

                maskData = numpy.zeros(reference.shape, dtype=numpy.uint8)
                maskData[offsets[0]:offsets[0] + widths[0], offsets[1] : offsets[1] + widths[1]] = 1
                maskFrame[:,:] = maskData[:,:]
                # fill the axes information
                dataGroup = hdf['/entry_000/Data']
                try:
                    activeCurve = self.getActiveCurve()
                    if activeCurve is None:
                        activeCurve = self.getAllCurves()[0]
                    x, y, legend, info = activeCurve
                    dataGroup[info['xlabel']] = numpy.array(x, dtype=numpy.float32)
                    dataGroup[info['xlabel']].attrs['axis'] = numpy.int32(1)
                    axesAttribute = '%s:dim_1:dim_2' % info['xlabel']
                except Exception:
                    if _logger.getEffectiveLevel() == logging.DEBUG:
                        raise
                    dataGroup['dim_0'] = numpy.arange(stack.data.shape[mcaIndex]).astype(numpy.float32)
                    dataGroup['dim_0'].attrs['axis'] = numpy.int32(1)
                    axesAttribute = 'dim_0:dim_1:dim_2'
                dataGroup['dim_1'] = numpy.arange(reference.shape[0]).astype(numpy.float32)
                dataGroup['dim_1'].attrs['axis'] = numpy.int32(2)
                dataGroup['dim_2'] = numpy.arange(reference.shape[1]).astype(numpy.float32)
                dataGroup['dim_2'].attrs['axis'] = numpy.int32(3)
                dim2 = numpy.arange(reference.shape[1]).astype(numpy.float32)
                dataGroup['data'].attrs['axes'] = axesAttribute
                self.finishHDF5File(hdf)
            else:
                self.setStack(stack)

    def __calculateShiftsFFT(self, *var, **kw):
        self._progress = 0.0
        thread = CalculationThread.CalculationThread(\
                calculation_method=self.calculateShiftsFFT,
                calculation_vars=var,
                calculation_kw=kw)
        thread.start()
        CalculationThread.waitingMessageDialog(thread,
                                               message="Please wait. Calculation going on.",
                                               parent=self.widget,
                                               modal=True,
                                               update_callback=self._waitingCallback)
        return thread.result

    def __shiftStack(self, *var, **kw):
        self._progress = 0.0
        thread = CalculationThread.CalculationThread(\
                calculation_method=self.shiftStack,
                calculation_vars=var,
                calculation_kw=kw)
        thread.start()
        CalculationThread.waitingMessageDialog(thread,
                                               message="Please wait. Calculation going on.",
                                               parent=self.widget,
                                               modal=True,
                                               update_callback=self._waitingCallback)
        return thread.result

    def __calculateShiftsSIFT(self, *var, **kw):
        self._progress = 0.0
        thread = CalculationThread.CalculationThread(\
                calculation_method=self.calculateShiftsSIFT,
                calculation_vars=var,
                calculation_kw=kw)
        thread.start()
        CalculationThread.waitingMessageDialog(thread,
                                               message="Please wait. Calculation going on.",
                                               parent=self.widget,
                                               modal=True,
                                               update_callback=self._waitingCallback)
        return thread.result

    def _waitingCallback(self):
        ddict = {}
        ddict['message'] = "Calculation Progress = %d %%" % self._progress
        return ddict

    def _siftAlignment(self):
        if not SIFT:
            try:
                import pyopencl
            except Exception:
                raise ImportError("PyOpenCL does not seem to be installed on your system")
        if ocl is None:
            raise ImportError("PyOpenCL does not seem to be installed on your system")
        stack = self.getStackDataObject()
        if stack is None:
            return
        mcaIndex = stack.info.get('McaIndex')
        if not (mcaIndex in [0, 2, -1]):
            raise IndexError("Unsupported 1D index %d" % mcaIndex)
        widget = SIFTAlignmentWindow.SIFTAlignmentDialog()
        widget.setStack(stack)
        mask = self.getStackSelectionMask()
        widget.setSelectionMask(mask)
        ret = widget.exec()
        if ret:
            ddict = widget.getParameters()
            widget.setDummyStack()
            reference = ddict['reference_image']
            mask = ddict['mask']
            if ddict['file_use']:
                filename = ddict['file_name']
            else:
                filename = None
            if filename is not None:
                self.__hdf5 = self.initializeHDF5File(filename)
            crop = False
            device = ddict['opencl_device']
            if _logger.getEffectiveLevel() == logging.DEBUG:
                result = self.calculateShiftsSIFT(stack, reference, mask=mask, device=device,
                                                  crop=crop, filename=filename)
            else:
                result = self.__calculateShiftsSIFT(stack, reference, mask=mask, device=device,
                                                    crop=crop, filename=filename)
                if result is not None:
                    if len(result):
                        if isinstance(result[0], str):
                            if result[0] == 'Exception':
                                # exception occurred
                                raise Exception(result[1], result[2], result[3])
            if filename is None:
                self.setStack(stack)

    def calculateShiftsSIFT(self, stack, reference, mask=None, device=None, crop=None,
                            sigma=None,
                            filename=None):
        mask = self.getStackSelectionMask()
        if mask is not None:
            if mask.sum() == 0:
                mask = None
        if device is None:
            siftInstance = sift.LinearAlign(reference.astype(numpy.float32),
                                            devicetype="cpu",
                                            init_sigma=sigma)
        else:
            siftInstance = sift.LinearAlign(reference.astype(numpy.float32),
                                            deviceid=device,
                                            init_sigma=sigma)
        data = stack.data
        mcaIndex = stack.info['McaIndex']
        if not (mcaIndex in [0, 2, -1]):
             raise IndexError("Unsupported 1D index %d" % mcaIndex)
        total = float(data.shape[mcaIndex])
        if filename is not None:
            hdf = self.__hdf5
            dataGroup = hdf['/entry_000/Data']
            attributes = {}
            attributes['interpretation'] = "image"
            attributes['signal'] = numpy.int32(1)
            outputStack = self.getHDF5BufferIntoGroup(dataGroup,
                                                      shape=(data.shape[mcaIndex],
                                                            reference.shape[0],
                                                            reference.shape[1]),
                                                      name="data",
                                                      dtype=numpy.float32,
                                                      attributes=attributes)
        shifts = numpy.zeros((data.shape[mcaIndex], 2), dtype=numpy.float32)
        if mcaIndex == 0:
            for i in range(data.shape[mcaIndex]):
                _logger.debug("SIFT Shifting image %d", i)
                result = siftInstance.align(data[i].astype(numpy.float32), shift_only=True, return_all=True)
                _logger.debug("Index = %d shift = %.4f, %.4f",
                              i, result['offset'][0], result['offset'][1])
                if filename is None:
                    stack.data[i] = result['result']
                else:
                    outputStack[i] = result['result']
                shifts[i, 0] = result['offset'][0]
                shifts[i, 1] = result['offset'][1]
                self._progress = (100 * i) / total
        else:
            image2 = numpy.zeros(reference.shape, dtype=numpy.float32)
            for i in range(data.shape[mcaIndex]):
                _logger.debug("SIFT Shifting image %d", i)
                image2[:, :] = data[:, :, i]
                result = siftInstance.align(image2, shift_only=True, return_all=True)
                _logger.debug("Index = %d shift = %.4f, %.4f",
                              i, result['offset'][0], result['offset'][1])
                if filename is None:
                    stack.data[:, :, i] = result['result']
                else:
                    outputStack[i] = result['result']
                shifts[i, 0] = result['offset'][0]
                shifts[i, 1] = result['offset'][1]
                self._progress = (100 * i) / total
        if filename is not None:
            hdf = self.__hdf5
            alignmentGroup = hdf['/entry_000/Alignment']
            outputShifts = self.getHDF5BufferIntoGroup(alignmentGroup,
                                                 shape=(stack.data.shape[mcaIndex], 2),
                                                 name="shifts",
                                                 dtype=numpy.float32)
            outputShifts[:,:] = shifts
            attributes={'interpretation':'image'}
            referenceFrame = self.getHDF5BufferIntoGroup(alignmentGroup,
                                                 shape=reference.shape,
                                                 name="reference_frame",
                                                 dtype=numpy.float32,
                                                 attributes=attributes)
            referenceFrame[:,:] = reference[:,:]
            maskFrame = self.getHDF5BufferIntoGroup(alignmentGroup,
                                               shape=reference.shape,
                                               name="reference_mask",
                                               dtype=numpy.uint8,
                                               attributes=attributes)

            if mask is None:
                maskData = numpy.ones(reference.shape, dtype=numpy.uint8)
            else:
                maskData = mask
            maskFrame[:,:] = maskData[:,:]
            # fill the axes information
            dataGroup = hdf['/entry_000/Data']
            try:
                activeCurve = self.getActiveCurve()
                if activeCurve is None:
                    activeCurve = self.getAllCurves()[0]
                x, y, legend, info = activeCurve
                dataGroup[info['xlabel']] = numpy.array(x, dtype=numpy.float32)
                dataGroup[info['xlabel']].attrs['axis'] = numpy.int32(1)
                axesAttribute = '%s:dim_1:dim_2' % info['xlabel']
            except Exception:
                if _logger.getEffectiveLevel() == logging.DEBUG:
                    raise
                dataGroup['dim_0'] = numpy.arange(stack.data.shape[mcaIndex]).astype(numpy.float32)
                dataGroup['dim_0'].attrs['axis'] = numpy.int32(1)
                axesAttribute = 'dim_0:dim_1:dim_2'
            dataGroup['dim_1'] = numpy.arange(reference.shape[0]).astype(numpy.float32)
            dataGroup['dim_1'].attrs['axis'] = numpy.int32(2)
            dataGroup['dim_2'] = numpy.arange(reference.shape[1]).astype(numpy.float32)
            dataGroup['dim_2'].attrs['axis'] = numpy.int32(3)
            dim2 = numpy.arange(reference.shape[1]).astype(numpy.float32)
            dataGroup['data'].attrs['axes'] = axesAttribute
            self.finishHDF5File(hdf)

    def calculateShiftsFFT(self, stack, reference, offsets=None, widths=None, crop=False):
        _logger.debug("Offsets = %s", offsets)
        _logger.debug("Widths = %s", widths)
        data = stack.data
        if offsets is None:
            offsets = [0.0, 0.0]
        if widths is None:
            widths = [reference.shape[0], reference.shape[1]]
        fft2Function = numpy.fft.fft2
        if 1:
            DTYPE = numpy.float32
        else:
            DTYPE = numpy.float64
        image2 = numpy.zeros((widths[0], widths[1]), dtype=DTYPE)
        shape = image2.shape

        USE_APODIZATION_WINDOW = False
        apo = [10, 10]
        if USE_APODIZATION_WINDOW:
            # use apodization window
            window = numpy.outer(SpecfitFuns.slit([0.5, shape[0] * 0.5, shape[0] - 4 * apo[0], apo[0]],
                                          numpy.arange(float(shape[0]))),
                                 SpecfitFuns.slit([0.5, shape[1] * 0.5, shape[1] - 4 * apo[1], apo[1]],
                                          numpy.arange(float(shape[1])))).astype(DTYPE)
        else:
            window = numpy.zeros((shape[0], shape[1]), dtype=DTYPE)
            window[apo[0]:shape[0] - apo[0], apo[1]:shape[1] - apo[1]] = 1
        image2[:,:] = window * reference[offsets[0]:offsets[0]+widths[0],
                                         offsets[1]:offsets[1]+widths[1]]
        image2fft2 = fft2Function(image2)
        mcaIndex = stack.info.get('McaIndex')
        shifts = numpy.zeros((data.shape[mcaIndex], 2), numpy.float64)
        image1 = numpy.zeros(image2.shape, dtype=DTYPE)
        total = float(data.shape[mcaIndex])
        if mcaIndex == 0:
            for i in range(data.shape[mcaIndex]):
                image1[:,:] = window * data[i][offsets[0]:offsets[0]+widths[0],
                                               offsets[1]:offsets[1]+widths[1]]

                image1fft2 = fft2Function(image1)
                shifts[i] = ImageRegistration.measure_offset_from_ffts(image2fft2,
                                                                       image1fft2)
                _logger.debug("Index = %d shift = %.4f, %.4f",
                              i, shifts[i][0], shifts[i][1])
                self._progress = (100 * i) / total
        elif mcaIndex in [2, -1]:
            for i in range(data.shape[mcaIndex]):
                image1[:,:] = window * data[:,:,i][offsets[0]:offsets[0]+widths[0],
                                               offsets[1]:offsets[1]+widths[1]]

                image1fft2 = fft2Function(image1)
                shifts[i] = ImageRegistration.measure_offset_from_ffts(image2fft2,
                                                                       image1fft2)
                _logger.debug("Index = %d shift = %.4f, %.4f",
                              i, shifts[i][0], shifts[i][1])
                self._progress = (100 * i) / total
        else:
            raise IndexError("Only stacks of images or spectra supported. 1D index should be 0 or 2")
        return shifts

    def _shiftFromFile(self):
        stack = self.getStackDataObject()
        if stack is None:
            return
        data = stack.data
        mcaIndex = stack.info.get('McaIndex')
        if not (mcaIndex in [0, -1, 2]):
            raise IndexError("1D index must be 0, 2, or -1")
        filefilter = ['HDF5 Files (*.h5 *.nxs *.hdf *.hdf5)', 'CSV 2-column (*.csv)', 'ASCII 2-column (*)']
        filename, ffilter = PyMcaFileDialogs.\
                    getFileList(parent=None,
                        filetypelist=filefilter,
                        message='Load',
                        mode='OPEN',
                        single=True,
                        getfilter=True,
                        currentfilter=filefilter[0])
        if len(filename):
            _logger.debug("file name = %s file filter = %s", filename, ffilter)
        else:
            _logger.debug("nothing selected")
            return
        filename = filename[0]
        if ffilter.startswith('HDF5'):
            # browse
            self.__hdf5Dialog = qt.QDialog()
            self.__hdf5Dialog.setWindowTitle('Select your data set by a double click')
            self.__hdf5Dialog.mainLayout = qt.QVBoxLayout(self.__hdf5Dialog)
            self.__hdf5Dialog.mainLayout.setContentsMargins(0, 0, 0, 0)
            self.__hdf5Dialog.mainLayout.setSpacing(0)
            fileModel = HDF5Widget.FileModel()
            fileView = HDF5Widget.HDF5Widget(fileModel)
            with h5py.File(filename, "r") as hdfFile:
                fileModel.appendPhynxFile(hdfFile, weakreference=True)
                self.__shiftsDataset = None
                fileView.sigHDF5WidgetSignal.connect(self._hdf5WidgetSlot)
                self.__hdf5Dialog.mainLayout.addWidget(fileView)
                self.__hdf5Dialog.resize(400, 200)
                ret = self.__hdf5Dialog.exec()
                if not ret:
                    return
                shifts = hdfFile[self.__shitfsDataset].value
        else:
            sf = specfilewrapper.Specfile(filename)
            nScans = len(sf)
            targetScan = None
            for scan in sf:
                if scan.lines() ==  data.shape[stack.info['McaIndex']]:
                    targetScan = scan
                    break
            if targetScan is None:
                scan = None
                sf = None
                raise IOError("Number of read lines does not match stack shape")
            shifts = targetScan.data()
            targetScan = None
            scan = None
            sf = None
            if shifts.shape[0] == 3 and\
               shifts.shape[1] == data.shape[stack.info['McaIndex']]:
                # one column was added (point number)
                shifts = shifts[1:].T

        filename = None
        if not isinstance(data, numpy.ndarray):
            filefilter = ['HDF5 Files (*.h5)']
            filename = PyMcaFileDialogs.\
                        getFileList(parent=None,
                        filetypelist=filefilter,
                        message='Select output file',
                        mode='SAVE',
                        single=True,
                        getfilter=False,
                        currentfilter=filefilter[0])
            if len(filename):
                filename = filename[0]
                _logger.debug("file name = %s", filename)
            else:
                raise IOError("No output file selected")
        if filename is not None:
            self.__hdf5 = self.initializeHDF5File(filename)
        crop = False
        if _logger.getEffectiveLevel() == logging.DEBUG:
            result = self.shiftStack(stack,
                                     shifts,
                                     crop=crop,
                                     filename=filename)
        else:
            result = self.__shiftStack(stack,
                                       shifts,
                                       crop=crop,
                                       filename=filename)
            if result is not None:
                # exception occurred
                raise Exception(result[1], result[2], result[3])

        if filename is not None:
            hdf = self.__hdf5
            alignmentGroup = hdf['/entry_000/Alignment']
            outputShifts = self.getHDF5BufferIntoGroup(alignmentGroup,
                                                 shape=(stack.data.shape[mcaIndex], 2),
                                                 name="shifts",
                                                 dtype=numpy.float32)
            outputShifts[:,:] = shifts
            attributes={'interpretation':'image'}
            # fill the axes information
            dataGroup = hdf['/entry_000/Data']
            if mcaIndex == 0:
                reference_shape = data[0].shape
            else:
                reference_shape = data.shape[0], data.shape[1]
            try:
                activeCurve = self.getActiveCurve()
                if activeCurve is None:
                    activeCurve = self.getAllCurves()[0]
                x, y, legend, info = activeCurve
                dataGroup[info['xlabel']] = numpy.array(x, dtype=numpy.float32)
                dataGroup[info['xlabel']].attrs['axis'] = numpy.int32(1)
                axesAttribute = '%s:dim_1:dim_2' % info['xlabel']
            except Exception:
                if _logger.getEffectiveLevel() == logging.DEBUG:
                    raise
                dataGroup['dim_0'] = numpy.arange(stack.data.shape[mcaIndex]).astype(numpy.float32)
                dataGroup['dim_0'].attrs['axis'] = numpy.int32(1)
                axesAttribute = 'dim_0:dim_1:dim_2'
            dataGroup['dim_1'] = numpy.arange(reference_shape[0]).astype(numpy.float32)
            dataGroup['dim_1'].attrs['axis'] = numpy.int32(2)
            dataGroup['dim_2'] = numpy.arange(reference_shape[1]).astype(numpy.float32)
            dataGroup['dim_2'].attrs['axis'] = numpy.int32(3)
            dim2 = numpy.arange(reference_shape[1]).astype(numpy.float32)
            dataGroup['data'].attrs['axes'] = axesAttribute
            self.finishHDF5File(hdf)
        else:
            self.setStack(stack)

    def _hdf5WidgetSlot(self, ddict):
        if ddict['event'] == "itemDoubleClicked":
            if ddict['type'].lower() in ['dataset']:
                self.__shitfsDataset = ddict['name']
                self.__hdf5Dialog.accept()

    def shiftStack(self, stack, shifts, crop=False, filename=None):
        """
        """
        data = stack.data
        mcaIndex = stack.info['McaIndex']
        if mcaIndex not in [0, 2, -1]:
             raise IndexError("Only stacks of images or spectra supported. 1D index should be 0 or 2")
        if mcaIndex == 0:
            shape = data[mcaIndex].shape
        else:
            shape = data.shape[0], data.shape[1]
        d0_start, d0_end, d1_start, d1_end = \
                  ImageRegistration.get_crop_indices(shape,
                                                     shifts[:, 0],
                                                     shifts[:, 1])
        window = numpy.zeros(shape, numpy.float32)
        window[d0_start:d0_end, d1_start:d1_end] = 1.0
        self._progress = 0.0
        total = float(data.shape[mcaIndex])
        if filename is not None:
            hdf = self.__hdf5
            dataGroup = hdf['/entry_000/Data']
            attributes = {}
            attributes['interpretation'] = "image"
            attributes['signal'] = numpy.int32(1)
            outputStack = self.getHDF5BufferIntoGroup(dataGroup,
                                                      shape=(data.shape[mcaIndex],
                                                            shape[0],
                                                            shape[1]),
                                                      name="data",
                                                      dtype=numpy.float32,
                                                      attributes=attributes)
        for i in range(data.shape[mcaIndex]):
            #tmpImage = ImageRegistration.shiftFFT(data[i], shifts[i])
            if mcaIndex == 0:
                tmpImage = ImageRegistration.shiftBilinear(data[i], shifts[i])
                #tmpImage = ImageRegistration.shiftImage(data[i], -shifts[i], method="fft")
                if filename is None:
                    stack.data[i] = tmpImage * window
                else:
                    outputStack[i] = tmpImage * window
            else:
                tmpImage = ImageRegistration.shiftBilinear(data[:,:,i], shifts[i])
                if filename is None:
                    stack.data[:, :, i] = tmpImage * window
                else:
                    outputStack[i] = tmpImage * window
            _logger.debug("Index %d bilinear shifted", i)
            self._progress = (100 * i) / total

    def initializeHDF5File(self, fname):
        #for the time being overwriting
        if os.path.exists(fname):
            os.remove(fname)
        hdf = h5py.File(fname, 'w')
        entryName = "entry_000"
        nxEntry = hdf.require_group(entryName)
        if 'NX_class' not in nxEntry.attrs:
            nxEntry.attrs['NX_class'] = 'NXentry'.encode('utf-8')
        nxEntry['title'] = numpy.string_("PyMca saved 3D Array".encode('utf-8'))
        nxEntry['start_time'] = numpy.string_(ArraySave.getDate().encode('utf-8'))

        alignmentGroup = nxEntry.require_group('Alignment')
        dataGroup = nxEntry.require_group('Data')
        dataGroup.attrs['NX_class'] = 'NXdata'.encode('utf-8')
        return hdf

    def finishHDF5File(self, hdf):
        #add final date
        toplevelEntry = hdf["entry_000"]
        toplevelEntry['end_time'] = numpy.string_(ArraySave.getDate().encode('utf-8'))
        hdf.flush()
        hdf.close()

    def getHDF5BufferIntoGroup(self, h5Group, shape,
                               name="data", dtype=numpy.float32,
                               attributes=None,
                               compression=None,
                               shuffle=False,
                               chunks=None,
                               chunk_cache=None):
        dataset = h5Group.require_dataset(name,
                                          shape=shape,
                                          dtype=dtype,
                                          chunks=chunks,
                                          shuffle=shuffle,
                                          compression=compression)
        if attributes is None:
            attributes = {}
        for attribute in attributes:
            dataset.attrs[attribute] = attributes[attribute]
        return dataset

MENU_TEXT = "Image Alignment Tool"
def getStackPluginInstance(stackWindow, **kw):
    ob = ImageAlignmentStackPlugin(stackWindow)
    return ob


import os

import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm

from src.data_processing import vec2img, img2vec
from src.utils import DATA_DIR
# ----- Dictionary learning -----
# Inspired from https://github.com/louismartin/dictionary-learning

def random_patches(Xtr, n_patches, patch_width):
    """Sample patches from images in Xtr"""
    # We will work on images and patches of shape (None, width, height, 3)
    Xtr = vec2img(Xtr)
    n_samples, img_width, _, _ = Xtr.shape

    # Choose a random image for each atom
    img_indexes = np.random.randint(n_samples, size=n_patches)
    img_indexes = img_indexes.repeat(patch_width * patch_width)\
                             .reshape(n_patches, patch_width, patch_width)

    # Coordinates of top left corner of each atom
    max_top_left = img_width - patch_width + 1
    top_left_X = np.random.randint(max_top_left, size=n_patches)
    top_left_Y = np.random.randint(max_top_left, size=n_patches)
    top_left_X = top_left_X.repeat(patch_width * patch_width)\
                           .reshape(n_patches, patch_width, patch_width)
    top_left_Y = top_left_Y.repeat(patch_width * patch_width)\
                           .reshape(n_patches, patch_width, patch_width)

    # patches coordinates strictly speaking
    X, Y = np.meshgrid(range(patch_width), range(patch_width))
    X = np.expand_dims(X, axis=0).repeat(n_patches, axis=0)
    Y = np.expand_dims(Y, axis=0).repeat(n_patches, axis=0)
    X += top_left_X
    Y += top_left_Y

    patches = Xtr[img_indexes, Y, X, :]
    # Cast back from (n_patches, width, height, 3) to (n_patches, n_pixels)
    patches = img2vec(patches)
    return patches

def initialize_atoms(Xtr, n_atoms, atom_width):
    """Initialize random atoms with high energy for dictionary learning
    returned atoms have shape (n_atoms, atom_width * atom_width * 3)
    """
    # We will initialize more atoms than needed in order to take those with
    # the highest energies
    atoms = random_patches(Xtr, 100 * n_atoms, atom_width)
    # Select atoms with highest energy
    energies = np.sum(atoms**2, axis=1)
    indexes = np.argsort(energies)[::-1][:n_atoms]
    atoms = atoms[indexes]
    return atoms

def sparsity_projection(coefs, sparsity):
    """Keep only k largest coefs on each row"""
    # Find the kth largest elements (pivots)
    pivots = np.partition(abs(coefs), -sparsity, axis=1)[:, -sparsity]
    to_keep = abs(coefs) >= pivots[np.newaxis].T
    # Only keep these kth largest coefficients
    # (careful sometimes more than k elements kept due to equal values)
    coefs = coefs * to_keep
    return coefs

def sparse_coding(data, atoms, coefs, sparsity, iterations=100):
    """Get the sparse representation of data on atoms represented by coefs
    Shapes:
        data: (n_samples, n_pixels)
        atoms: (n_atoms, n_pixels)
        coefs: (n_samples, n_atoms)
    """
    lambd = 1/np.linalg.norm(np.dot(atoms.T, atoms))
    for i in tqdm(range(iterations)):
        error = np.dot(coefs, atoms) - data
        coefs = sparsity_projection(coefs - lambd * np.dot(error, atoms.T),
                                    sparsity)
    return coefs

def dictionary_projection(atoms):
    """Scale all atoms to unit norm"""
    norms = np.linalg.norm(atoms, axis=1)
    norms = np.tile(norms, (atoms.shape[1], 1)).T
    atoms = atoms / norms
    return atoms

def dictionary_update(data, atoms, coefs, iterations=100):
    """Update atoms given fixed coefs, to represent the data
    Shapes:
        data: (n_samples, n_pixels)
        atoms: (n_atoms, n_pixels)
        coefs: (n_samples, n_atoms)
    """
    lambd = 1/np.linalg.norm(np.dot(coefs.T, coefs))
    for i in tqdm(range(iterations)):
        error = np.dot(coefs, atoms) - data
        atoms = dictionary_projection(atoms - lambd * np.dot(coefs.T, error))
    return atoms

def learn_dictionary(Xtr, n_atoms, atom_width, plot=False):
    n_samples = Xtr.shape[0]
    atoms = initialize_atoms(Xtr, n_atoms=n_atoms, atom_width=atom_width)
    n_patches = 10 * n_samples
    data = random_patches(Xtr, n_patches, atom_width)
    coefs = np.random.rand(n_patches, n_atoms)

    # Learn the dictionary
    iterations = 100
    errors = np.zeros(2*iterations)
    for i in tqdm(range(iterations)):
        # Sparse coding
        coefs = sparse_coding(data, atoms, coefs, sparsity=5, iterations=50)
        errors[2*i] = np.linalg.norm(np.dot(coefs, atoms) - data)**2

        # Dictionary update
        atoms = dictionary_update(data, atoms, coefs, iterations=50)
        errors[2*i+1] = np.linalg.norm(np.dot(coefs, atoms) - data)**2
    path = os.path.join(DATA_DIR, "dictionary_learning_errors.npy")
    np.save(path, errors)
    if plot:
        plt.plot(errors)
    return atoms

def extract_all_patches(Xtr, patch_width):
    """Extract all non overlapping patches of the images in Xtr"""
    # We will work on images and patches of shape (None, width, height, 3)
    Xtr = vec2img(Xtr)
    img_width = Xtr.shape[1]
    patch_per_side = img_width // patch_width
    n_patches = patch_per_side**2

    coords = np.arange(0, img_width - patch_width + 1, patch_width)
    top_left_X, top_left_Y = np.meshgrid(coords, coords)
    top_left_X = top_left_X.flatten()
    top_left_Y = top_left_Y.flatten()
    top_left_X = top_left_X.repeat(patch_width * patch_width)\
                           .reshape(n_patches, patch_width, patch_width)
    top_left_Y = top_left_Y.repeat(patch_width * patch_width)\
                           .reshape(n_patches, patch_width, patch_width)
    X, Y = np.meshgrid(range(patch_width), range(patch_width))
    X = np.expand_dims(X, axis=0).repeat(n_patches, axis=0)
    Y = np.expand_dims(Y, axis=0).repeat(n_patches, axis=0)
    X += top_left_X
    Y += top_left_Y

    patches = Xtr[:, Y, X, :]
    return patches

class Dictionary:
    def __init__(self, n_atoms, atom_width):
        self.n_atoms = n_atoms
        self.atom_width = atom_width
        filename = "atoms_{}_{}.npy".format(self.atom_width, self.n_atoms)
        self.save_path = os.path.join(DATA_DIR, filename)

    @property
    def weights_available(self):
        return os.path.exists(self.save_path)

    def fit(self, Xtr):
        self.atoms = learn_dictionary(Xtr,
                                      n_atoms=self.n_atoms,
                                      atom_width=self.atom_width,
                                      plot=False)

    def get_representation(self, Xtr, iterations=200):
        n_samples = Xtr.shape[0]

        # Extract all non overlapping patches of each image
        patches = extract_all_patches(Xtr, patch_width=self.atom_width)
        n_patches = patches.shape[1]
        patches = patches.reshape(n_samples * n_patches,
                                  self.atom_width, self.atom_width, 3)
        patches = img2vec(patches)

        # Get sparse representation of all patches
        coefs = np.random.rand(patches.shape[0], self.n_atoms)
        coefs = sparse_coding(patches, self.atoms, coefs,
                              sparsity=self.n_atoms//2, iterations=iterations)
        coefs = coefs.reshape(n_samples, self.n_atoms * n_patches)
        return coefs

    def save(self):
        np.save(self.save_path, self.atoms)

    def load(self):
        self.atoms = np.load(self.save_path)


import base64
import streamlit as st
from pathlib import Path
from . import resources
import os

def render_svg(path):
    """Renders the given svg string."""
    with open(path, 'r+') as file:
        svg = file.read()
    b64 = base64.b64encode(svg.encode('utf-8')).decode("utf-8")
    html = r'' % b64
    return html

def header():
    st.set_page_config(layout="wide")
    col1, col2 = st.columns(2)
    title = """

Experiment Tracker
"""

    col1.markdown(title, unsafe_allow_html=True)
    p = Path(resources.__path__._path[0])
    col2.write(render_svg(os.path.join(p,"logo_csem.svg")), unsafe_allow_html=True)


import hydra
import torch
import re
import matplotlib.pyplot as plt
import re
from hydra.utils import instantiate, get_class, call
from hydra.core.config_store import ConfigStore
from hydra_main import FourDVarNetHydraRunner
import pytorch_lightning as pl
import pandas as pd
from pathlib import Path
import traceback
import hydra_config
from IPython.display import display, Markdown, Latex, HTML

s = """
## TODO
Dev
- [ ]  list xps
- [ ] from xp get dir and ckpt
"""
display(Markdown(s))

def main():
    try:
        fn = test_models_on_diff_noises
        # fn = data_stats
        # fn = plot_results

        locals().update(fn())
    except Exception as e:
        print('I am here')
        print(traceback.format_exc())
    finally:
        return locals()

cs = ConfigStore.instance()

def bst_ckpt(dirpath, glob='*'):
    return min(Path(dirpath).glob('version_*/checkpoints/*'), key=lambda p: float(re.match('.+val_loss=(.+)\.ckpt', str(p)).group(1)))

def get_cfg(xp_cfg, overrides=None):
    overrides = overrides if overrides is not None else []
    with hydra.initialize_config_dir(str(Path('hydra_config').absolute())):
        cfg = hydra.compose(config_name='main', overrides=
            [
                f'xp={xp_cfg}',
                'file_paths=dgx_ifremer',
                'entrypoint=train',
            ] + overrides
        )

    return cfg

def get_model(xp_cfg, ckpt, dm=None, add_overrides=None):
    overrides = []
    if add_overrides is not None:
        overrides =  overrides + add_overrides
    cfg = get_cfg(xp_cfg, overrides)
    lit_mod_cls = get_class(cfg.lit_mod_cls)
    if dm is None:
        dm = instantiate(cfg.datamodule)
    runner = FourDVarNetHydraRunner(cfg.params, dm, lit_mod_cls)
    mod = runner._get_model(ckpt)
    return mod

def get_dm(xp_cfg, setup=True, add_overrides=None):
    overrides = []
    if add_overrides is not None:
        overrides = overrides + add_overrides
    cfg = get_cfg(xp_cfg, overrides)
    dm = instantiate(cfg.datamodule)
    if setup:
        dm.setup()
    return dm

# for mn in models:
#     print(mn)
#     get_model(mn)

# for dmn in dms:
#     print(dmn)
#     get_dm(dmn)

import importlib
import calibration.dataset
import hydra_main
importlib.reload(calibration.dataset)
importlib.reload(hydra_main)
import json

def foo():

        prepro_cfg={
                '_target_': 'calibration.dataset.weight_error',
                'w': 0.1,
        }
        or_str = str(prepro_cfg).replace(" ", "").replace("'","")
        dm = get_dm('errs', add_overrides=[f'+datamodule.item_prepro={or_str}'])
        batch = next(iter(dm.val_dataloader()))
        obs = batch[2]
        im = plt.imshow(obs[1,3,...])
        plt.colorbar(im)

def plot_results():
    noise_level_df = pd.read_csv('noise_levels.csv')
    print(noise_level_df.to_markdown())
    df = pd.read_csv('sst_mods_metrics3.csv')
    # df = pd.read_csv('no_sst_metrics3.csv').loc[lambda _df: ~_df.model.str.contains('cal') | _df.model.str.contains('5nad')]

    metrics = [
        'test_spatial_res_imp',
        'test_nrmse_ratio_glob',
        'test_mse_ratio_glob',
        'test_nrmse_swath',
        'test_mse_swath'
    ]
    for m in metrics:
        (
                df[['dm', 'model', m]]
                .pivot(columns='model', values=m, index='dm')
                .rename({'errs_0.0_sst': 'errs_0.000_sst', 'errs_0.0': 'errs_0.000'})
                .sort_index()
                .plot(title=m, figsize=(10, 5), logy=True)
        )

    return locals()

def data_stats():
    noise_level = []
    for dm_name in [
        'errs_0.0',
        'errs_0.001',
        'errs_0.01',
        'errs_0.1',
        'errs_0.33',
        'errs_0.66',
        'errs_1.0',
        '5nad',
    ]:
        dm = get_dm(dm_name)
        dm.setup()
        it = iter(dm.val_dataloader())
        next(it)
        oi, mask, obs, gt, obs_gt  = next(it)
        nanobs = obs.where(mask, torch.full_like(obs, float('nan')))
        print(dm_name, torch.nanmean((nanobs - obs_gt)**2).sqrt())
        noise_level.append({'dm': dm_name, 'rmse': torch.nanmean((nanobs - obs_gt)**2).sqrt().item()})
        # img = plt.imshow((nanobs-obs_gt)[1,3,...])
        # plt.colorbar(img)
        # plt.show()

        # img = plt.imshow(obs.where(mask, torch.full_like(obs,float('nan')))[1,3,...])
        # plt.colorbar(img)
        # plt.show()

    noise_level_df = pd.DataFrame(noise_level)
    noise_level_df.to_csv('noise_levels.csv')
    return locals()

def test_models_on_diff_noises():
    try:
        metrics = []
        xps = cs.list('xp')
        for cfg in xps:
            if '5nad' in cfg: continue
            print(cfg)
            m = re.match(r'qxp(?P<xp_number>\d+)_(?P<xp_name>.+)\.yaml', cfg)
            xp_n = m.group('xp_number')
            xp_name = m.group('xp_name')

            # inference_configs:
            xp_dir = Path('xp_roll')
            data = {
                '5nad': '+xp/qfebvre/ds=five_nadirs',
                'roll': '+xp/qfebvre/ds=new_noisy_swot_roll',
                'syst': '+xp/qfebvre/ds=new_noisy_swot_syst_errs',
            }

            ws = {
                'w0.000': 'datamodule.item_prepro.w=0.000',
                'w0.001': 'datamodule.item_prepro.w=0.001',
                'w0.01': 'datamodule.item_prepro.w=0.01',
                'w0.1': 'datamodule.item_prepro.w=0.1',
                'w1.': 'datamodule.item_prepro.w=1',
            }

            ckpt = bst_ckpt(xp_dir / xp_name)

            for obs, obs_or in data.items():
                for w, w_or in ws.items():
                    print(obs, w)
                    if (obs == '5nad') and not (w=='w1.'): continue

                    xp_cfg = get_cfg(cfg)
                    dm = get_dm(cfg, setup=False, add_overrides=[obs_or, w_or])
                    mod = get_model(cfg, ckpt, dm, add_overrides=[obs_or, w_or])
                    trainer = pl.Trainer(gpus=[3])
                    trainer.test(mod, dataloaders=dm.test_dataloader())
                    metrics.append(
                        {
                            'model': xp_name,
                            'xp_n': xp_n,
                            'obs_train': xp_cfg.datamodule.obs_mask_var,
                            'obs_test': obs + w,
                            'sst': xp_cfg.params.sst,
                            'cal': xp_cfg.params.loss_loc,
                            'aug': xp_cfg.params.aug_state,
                            'metrics': mod.latest_metrics,
                        }
                    )

                df = pd.DataFrame(
                    metrics
                ).pipe(
                        lambda _df:_df.assign(**pd.DataFrame(list(_df.metrics)).to_dict(orient='list'))
                ).drop('metrics', axis=1)
                df.to_csv('metrics_roll2.csv')
    except Exception as e:

        print(traceback.format_exc())

    finally:

        return locals()