# 
#   This file is part of Mantis, a Multivariate ANalysis Tool for Spectromicroscopy.
# 
#   Copyright (C) 2011 Mirna Lerotic, 2nd Look
#   http://2ndlookconsulting.com
#   License: GNU GPL v3
#
#   Mantis is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   any later version.
#
#   Mantis is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details <http://www.gnu.org/licenses/>.


from __future__ import division
from __future__ import print_function

import numpy as np
import os
import re

verbose = 0

title = 'Ncb'
extension = ['*.ncb']
read_types = ['image', 'stack']
write_types = ['image', 'stack', 'results']


def identify(filename):
    try:
        basename, extension = os.path.splitext(filename)
        dat_fn = basename + '.dat'
        f = open(str(dat_fn), 'r')
        f.close()
        return True
    except:
        return False


def GetFileStructure(FileName):
    return None


# ----------------------------------------------------------------------
def read(filename, self, selection=None, *args, **kwargs):
    basename, extension = os.path.splitext(filename)
    dat_fn = basename + '.dat'

    f = open(str(dat_fn), 'r')
    lines = f.readlines()
    temp = lines[0].split()

    self.n_cols = int(temp[0])
    self.n_rows = int(temp[1])
    scale = float(temp[2])

    # print self.n_cols, self.n_rows, scale

    temp = lines[1].split()
    x_start = float(temp[0])
    x_stop = float(temp[1])

    temp = lines[2].split()
    y_start = float(temp[0])
    y_stop = float(temp[1])

    self.n_ev = int(lines[3])

    # print self.n_ev

    self.ev = np.zeros((self.n_ev))
    self.data_dwell = np.zeros((self.n_ev))
    filename_list = []
    for i in range(self.n_ev):
        self.ev[i] = float(lines[4 + i])

    # print self.ev

    for i in range(self.n_ev):
        self.ev[i] = float(lines[4 + i])

    try:
        for i in range(self.n_ev):
            temp = lines[4 + self.n_ev + i].split()

            filename_list.append(temp[0])
            self.data_dwell[i] = float(temp[2])
    except:
        self.data_dwell = np.ones((self.n_ev))

    f.close()

    if scale < 0:
        dataformat = np.float32
    else:
        dataformat = np.int16

    # print 'data format', dataformat

    f = open(str(filename), 'rb')
    big_array = np.fromfile(f, dataformat, self.n_cols * self.n_rows * self.n_ev)
    f.close()

    if scale > 0 and scale != 1:
        image_stack = big_array.astype(float) / scale
        print("data rescaled by ", 1. / scale)
    else:
        image_stack = big_array.astype(float)

    if (x_start > x_stop):
        image_stack = image_stack[::-1, :, :]
        t = x_start
        x_start = x_stop
        x_stop = t
        print('x data reversed')

    if (y_start > y_stop):
        image_stack = image_stack[:, ::-1, :]
        t = y_start
        y_start = y_stop
        y_stop = t
        print('y data reversed')

    xstep = (x_stop - x_start) / (self.n_cols - 1)
    self.x_dist = np.arange(x_start, x_stop + xstep, xstep)

    ystep = (y_stop - y_start) / (self.n_rows - 1)
    self.y_dist = np.arange(y_start, y_stop + ystep, ystep)

    self.absdata = np.empty((self.n_cols, self.n_rows, self.n_ev))

    self.absdata = np.reshape(image_stack, (self.n_cols, self.n_rows, self.n_ev), order='F')

    self.fill_h5_struct_from_stk()

    return


# ----------------------------------------------------------------------
def read_ncb4D(self, filenames):
    if verbose: print('First energy stack:', filenames[0])

    data_files = natural_sort(filenames)

    neng = len(data_files)
    if verbose: print('Number of energies ', neng)

    self.stack4D = []
    theta = []

    for i in range(neng):
        # Read the data
        thisstack, thistheta = read_ncb_data(self, data_files[i])

        # Check if we have negative angles, if yes convert to 0-360deg
        for ith in range(len(thistheta)):
            if thistheta[ith] < 0:
                thistheta[ith] = thistheta[ith] + 360

        self.stack4D.append(thisstack)
        theta.append(thistheta)

        dims = thisstack.shape

    self.stack4D = np.array(self.stack4D)
    self.stack4D = np.transpose(self.stack4D, axes=(1, 2, 0, 3))

    self.theta = theta[0]
    self.n_theta = len(theta)

    self.absdata = self.stack4D[:, :, :, 0]

    self.data_dwell = np.ones((neng))

    self.fill_h5_struct_from_stk()

    return


# ----------------------------------------------------------------------
def read_ncb_data(self, filename):
    basename, extension = os.path.splitext(filename)
    dat_fn = basename + '.dat'

    f = open(str(dat_fn), 'r')
    lines = f.readlines()
    temp = lines[0].split()

    self.n_cols = int(temp[0])
    self.n_rows = int(temp[1])
    scale = float(temp[2])

    if verbose:
        print('self.n_cols, self.n_rows, scale', self.n_cols, self.n_rows, scale)

    temp = lines[1].split()
    x_start = float(temp[0])
    x_stop = float(temp[1])

    temp = lines[2].split()
    y_start = float(temp[0])
    y_stop = float(temp[1])

    self.n_theta = int(lines[3])

    angles = np.zeros((self.n_theta))
    self.data_dwell = np.zeros((self.n_theta))
    filename_list = []
    for i in range(self.n_theta):
        angles[i] = float(lines[4 + i])

    for i in range(self.n_theta):
        angles[i] = float(lines[4 + i])

    try:
        for i in range(self.n_theta):
            temp = lines[4 + self.n_ev + i].split()

            filename_list.append(temp[0])
            self.data_dwell[i] = float(temp[2])
    except:
        self.data_dwell = np.ones((self.n_theta))

    f.close()

    if scale < 0:
        dataformat = np.float32
    else:
        dataformat = np.int16

    f = open(str(filename), 'rb')
    big_array = np.fromfile(f, dataformat, self.n_cols * self.n_rows * self.n_theta)
    f.close()

    if scale > 0 and scale != 1:
        image_stack = big_array.astype(float) / scale
        if verbose: print("data rescaled by ", 1. / scale)
    else:
        image_stack = big_array.astype(float)

    if (x_start > x_stop):
        image_stack = image_stack[::-1, :, :]
        t = x_start
        x_start = x_stop
        x_stop = t
        if verbose: print('x data reversed')

    if (y_start > y_stop):
        image_stack = image_stack[:, ::-1, :]
        t = y_start
        y_start = y_stop
        y_stop = t
        if verbose: print('y data reversed')

    xstep = (x_stop - x_start) / (self.n_cols - 1)
    self.x_dist = np.arange(x_start, x_stop + xstep, xstep)

    ystep = (y_stop - y_start) / (self.n_rows - 1)
    self.y_dist = np.arange(y_start, y_stop + ystep, ystep)

    image_stack = np.reshape(image_stack, (self.n_cols, self.n_rows, self.n_theta), order='F')

    return image_stack, angles


# ----------------------------------------------------------------------
#  This procedure writes  a whole stack (3d (E,x,y) array) to a binary file
#  with associated *.dat file to track paramaters
def write(filename, stack, data_type):  # ,norm):
    #
    # def operate_on_Narray(A, B, function):
    #     try:
    #         return [operate_on_Narray(a, b, function) for a, b in zip(A, B)]
    #     except TypeError as e:
    #         # Not iterable
    #         return function(A, B)
    # import matplotlib.pyplot as plt
    # def export_figure_matplotlib(name,arr, dpi=100, resize_fact=1, plt_show=False):
    #     fig = plt.figure(frameon=False)
    #     fig.set_size_inches(arr.shape[1] / dpi, arr.shape[0] / dpi)
    #     ax = plt.Axes(fig, [0., 0., 1., 1.])
    #     ax.set_axis_off()
    #     fig.add_axes(ax)
    #     ax.imshow(arr, cmap='gray', interpolation='none')
    #     if plt_show:
    #         plt.show()
    #     else:
    #         plt.savefig(name+'.tif', dpi=(dpi * resize_fact))
    #         plt.close()
    print('Writing .ncb stack:', filename)

    basename, extension = os.path.splitext(filename)
    dat_fn = basename + '.dat'
    filename = basename + '.ncb'

    image_stack = np.transpose(stack.absdata, axes=(1, 0, 2))
    '''Experimental OD filtering and edge substraction'''
    # if norm:
    #
    #     image_stack = np.transpose(stack.od3d, axes=(1,0,2))
    #     #image_stack = stack.od3d.copy()
    #     OD_filter = np.amax(image_stack, 2)
    #     #returning the ODmask as tif
    #     mask = np.where((OD_filter > 2.5), 0, np.ones(np.shape(image_stack[:,:,0])))
    #     export_figure_matplotlib(basename+"od_mask",mask)
    #
    #     for i in range(np.size(image_stack, 2)):
    #         image_stack[:,:,i] = np.where((OD_filter > 2.5),0, image_stack[:,:,i])
    #     #image_stack = np.where(image_stack > 2,0, image_stack) #OD > 2 filter
    #     e0 = 0
    #     num = 5
    #     ##pre-edge
    #     ave = image_stack[:,:,e0]
    #     for i in range(num-1):
    #         ave = ave + image_stack[:,:,e0+i+1]
    #     ave = ave / num
    #
    #     #export pre-edge_average for contour plot.
    #     pre = np.transpose(stack.absdata, axes=(1,0,2))[:,:,e0]
    #     for i in range(num-1):
    #         pre = pre + np.transpose(stack.absdata, axes=(1,0,2))[:,:,e0+i+1]
    #     pre = pre / num
    #     export_figure_matplotlib(basename,pre)
    #
    #     for i in range(np.size(image_stack, 2)):
    #         image_stack[:,:,i] = image_stack[:,:,i]-ave
    #
    #     #post-edge
    #     ave = image_stack[:,:,-1]
    #     for i in range(num-1):
    #         ave = ave + image_stack[:,:,e0-i-2]
    #     ave = (ave / num) + 0.0001
    #     #print(ave)
    #     for i in range(np.size(image_stack, 2)):
    #         image_stack[:,:,i] = image_stack[:,:,i] - operate_on_Narray(ave, image_stack[:,:,i], lambda a, b: a/((1+ np.exp(-b+0.5*a)**(25*1/a))))

    image_stack = np.reshape(image_stack, (stack.n_cols * stack.n_rows * stack.n_ev), order='F')

    tmax = np.amax(image_stack)
    tmin = np.amin(image_stack)
    test = np.amax([abs(tmin), abs(tmax)])
    scale = 1.
    if test > 3e4 or test < 1e3:
        scale = 10. ** (3 - np.fix(np.log10(test)))

    # Save image stack to a binary .dat file
    f = open(str(filename), 'wb')
    if scale != 1.0:
        print('Scaling the data by ', scale)
        saveddata = image_stack * scale
        saveddata.astype(np.int16).tofile(f)
    else:
        scale = -1.0
        image_stack.astype(np.float32).tofile(f)
    f.close()

    # print 'imagedims', image_stack.shape

    f = open(str(dat_fn), 'w')

    print('\t%d\t%d\t%.6f' % (stack.n_rows, stack.n_cols, scale), file=f)

    x_start = stack.y_dist[0]
    x_stop = stack.y_dist[-1]
    if x_start != 0.:
        x_stop = x_stop - x_start
        x_start = 0.

    print('\t%.6f\t%.6f' % (x_start, x_stop), file=f)

    y_start = stack.x_dist[0]
    y_stop = stack.x_dist[-1]
    if y_start != 0.:
        y_stop = y_stop - y_start
        y_start = 0.
    print('\t%.6f\t%.6f' % (y_start, y_stop), file=f)

    print('\t%d' % (stack.n_ev), file=f)

    for i in range(stack.n_ev):
        print('\t%.6f' % (stack.ev[i]), file=f)

    for i in range(stack.n_ev):
        thisstr = 'image' + str(i + 1)
        print('%s\t%.6f\t%.6f' % (thisstr, stack.ev[i], stack.data_dwell[i]), file=f)

    f.close()

    return


# ----------------------------------------------------------------------
def natural_sort(l):
    convert = lambda text: int(text) if text.isdigit() else text.lower()
    alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
    return sorted(l, key=alphanum_key)


# ----------------------------------------------------------------------
class Cncb:
    def __init__(self):
        pass

    # ----------------------------------------------------------------------
    def read_ncb(self, filename):

        print('Reading stack:', filename)

        basename, extension = os.path.splitext(filename)
        dat_fn = basename + '.dat'

        f = open(str(dat_fn), 'r')
        lines = f.readlines()
        temp = lines[0].split()

        self.n_cols = int(temp[0])
        self.n_rows = int(temp[1])
        scale = float(temp[2])

        # print self.n_cols, self.n_rows, scale

        temp = lines[1].split()
        x_start = float(temp[0])
        x_stop = float(temp[1])

        temp = lines[2].split()
        y_start = float(temp[0])
        y_stop = float(temp[1])

        self.n_ev = int(lines[3])

        # print self.n_ev

        self.ev = np.zeros((self.n_ev))
        self.data_dwell = np.zeros((self.n_ev))
        filename_list = []
        for i in range(self.n_ev):
            self.ev[i] = float(lines[4 + i])

        # print self.ev

        for i in range(self.n_ev):
            self.ev[i] = float(lines[4 + i])

        try:
            for i in range(self.n_ev):
                temp = lines[4 + self.n_ev + i].split()

                filename_list.append(temp[0])
                self.data_dwell[i] = float(temp[2])
        except:
            self.data_dwell = np.ones((self.n_ev))

        f.close()

        if scale < 0:
            dataformat = np.float32
        else:
            dataformat = np.int16

        # print 'data format', dataformat

        f = open(str(filename), 'rb')
        big_array = np.fromfile(f, dataformat, self.n_cols * self.n_rows * self.n_ev)
        f.close()

        if scale > 0 and scale != 1:
            image_stack = big_array.astype(float) / scale
            print("data rescaled by ", 1. / scale)
        else:
            image_stack = big_array.astype(float)

        if (x_start > x_stop):
            image_stack = image_stack[::-1, :, :]
            t = x_start
            x_start = x_stop
            x_stop = t
            print('x data reversed')

        if (y_start > y_stop):
            image_stack = image_stack[:, ::-1, :]
            t = y_start
            y_start = y_stop
            y_stop = t
            print('y data reversed')

        xstep = (x_stop - x_start) / (self.n_cols - 1)
        self.x_dist = np.arange(x_start, x_stop + xstep, xstep)

        ystep = (y_stop - y_start) / (self.n_rows - 1)
        self.y_dist = np.arange(y_start, y_stop + ystep, ystep)

        self.absdata = np.empty((self.n_cols, self.n_rows, self.n_ev))

        self.absdata = np.reshape(image_stack, (self.n_cols, self.n_rows, self.n_ev), order='F')

        return

    # ----------------------------------------------------------------------
    #  This procedure writes  a whole stack (3d (E,x,y) array) to a binary file
    #  with associated *.dat file to track paramaters
    def write_ncb(self, filename, data_struct):

        print('Writing .ncb stack:', filename)

        basename, extension = os.path.splitext(filename)
        dat_fn = basename + '.dat'

        image_stack = np.transpose(self.absdata, axes=(1, 0, 2))

        # image_stack = self.absdata.copy()

        image_stack = np.reshape(image_stack, (self.n_cols * self.n_rows * self.n_ev), order='F')

        tmax = np.amax(image_stack)
        tmin = np.amin(image_stack)
        test = np.amax([abs(tmin), abs(tmax)])
        scale = 1.
        if test > 3e4 or test < 1e3:
            scale = 10. ** (3 - np.fix(np.log10(test)))

        # Save image stack to a binary .dat file
        f = open(str(filename), 'wb')
        if scale != 1.0:
            print('Scaling the data by ', scale)
            saveddata = image_stack * scale
            saveddata.astype(np.int16).tofile(f)
        else:
            scale = -1.0
            image_stack.astype(np.float32).tofile(f)
        f.close()

        # print 'imagedims', image_stack.shape

        f = open(str(dat_fn), 'w')

        print('\t%d\t%d\t%.6f' % (self.n_rows, self.n_cols, scale), file=f)
        # print('\t%d\t%d\t%.6f' %(self.n_cols, self.n_rows, scale)

        x_start = self.y_dist[0]
        x_stop = self.y_dist[-1]
        if x_start != 0.:
            x_stop = x_stop - x_start
            x_start = 0.

        print('\t%.6f\t%.6f' % (x_start, x_stop), file=f)

        y_start = self.x_dist[0]
        y_stop = self.x_dist[-1]
        if y_start != 0.:
            y_stop = y_stop - y_start
            y_start = 0.
        print('\t%.6f\t%.6f' % (y_start, y_stop), file=f)

        print('\t%d' % (self.n_ev), file=f)

        for i in range(self.n_ev):
            print('\t%.6f' % (self.ev[i]), file=f)

        for i in range(self.n_ev):
            thisstr = 'image' + str(i + 1)
            print('%s\t%.6f\t%.6f' % (thisstr, self.ev[i], self.data_dwell[i]), file=f)

        f.close()

        return

    # ----------------------------------------------------------------------
    def read_ncb_data(self, filename):

        # print 'Reading stack:', filename

        basename, extension = os.path.splitext(filename)
        dat_fn = basename + '.dat'

        f = open(str(dat_fn), 'r')
        lines = f.readlines()
        temp = lines[0].split()

        self.n_cols = int(temp[0])
        self.n_rows = int(temp[1])
        scale = float(temp[2])

        if verbose:
            print('self.n_cols, self.n_rows, scale', self.n_cols, self.n_rows, scale)

        temp = lines[1].split()
        x_start = float(temp[0])
        x_stop = float(temp[1])

        temp = lines[2].split()
        y_start = float(temp[0])
        y_stop = float(temp[1])

        self.n_theta = int(lines[3])

        # print self.n_ev

        angles = np.zeros((self.n_theta))
        self.data_dwell = np.zeros((self.n_theta))
        filename_list = []
        for i in range(self.n_theta):
            angles[i] = float(lines[4 + i])

        # print self.ev

        for i in range(self.n_theta):
            angles[i] = float(lines[4 + i])

        try:
            for i in range(self.n_theta):
                temp = lines[4 + self.n_ev + i].split()

                filename_list.append(temp[0])
                self.data_dwell[i] = float(temp[2])
        except:
            self.data_dwell = np.ones((self.n_theta))

        f.close()

        if scale < 0:
            dataformat = np.float32
        else:
            dataformat = np.int16

        # print 'data format', dataformat

        f = open(str(filename), 'rb')
        big_array = np.fromfile(f, dataformat, self.n_cols * self.n_rows * self.n_theta)
        f.close()

        if scale > 0 and scale != 1:
            image_stack = big_array.astype(float) / scale
            if verbose: print("data rescaled by ", 1. / scale)
        else:
            image_stack = big_array.astype(float)

        if (x_start > x_stop):
            image_stack = image_stack[::-1, :, :]
            t = x_start
            x_start = x_stop
            x_stop = t
            if verbose: print('x data reversed')

        if (y_start > y_stop):
            image_stack = image_stack[:, ::-1, :]
            t = y_start
            y_start = y_stop
            y_stop = t
            if verbose: print('y data reversed')

        xstep = (x_stop - x_start) / (self.n_cols - 1)
        self.x_dist = np.arange(x_start, x_stop + xstep, xstep)

        ystep = (y_stop - y_start) / (self.n_rows - 1)
        self.y_dist = np.arange(y_start, y_stop + ystep, ystep)

        image_stack = np.reshape(image_stack, (self.n_cols, self.n_rows, self.n_theta), order='F')

        return image_stack, angles

    # ----------------------------------------------------------------------
    def read_ncb4D(self, filenames):

        if verbose: print('First energy stack:', filenames[0])

        data_files = natural_sort(filenames)

        neng = len(data_files)
        if verbose: print('Number of energies ', neng)

        self.stack4D = []
        theta = []

        for i in range(neng):
            # Read the data
            thisstack, thistheta = self.read_ncb_data(data_files[i])

            # Check if we have negative angles, if yes convert to 0-360deg
            for ith in range(len(thistheta)):
                if thistheta[ith] < 0:
                    thistheta[ith] = thistheta[ith] + 360

            self.stack4D.append(thisstack)
            theta.append(thistheta)

            dims = thisstack.shape

        self.stack4D = np.array(self.stack4D)
        self.stack4D = np.transpose(self.stack4D, axes=(1, 2, 0, 3))

        self.theta = theta[0]
        self.n_theta = len(theta)

        self.absdata = self.stack4D[:, :, :, 0]

        self.data_dwell = np.ones((neng))

        return
