import matplotlib as mpl
import matplotlib.pyplot as plt
import glob
import numpy as np
from scipy.optimize import curve_fit
import traceback
from lmfit import minimize, Parameters
from uncertainties import ufloat
import pandas as pd
import Settings
import sys

# todo: check program with several different settings
# todo: solve the problems with manual fitting

# todo: adjust the plotting function to write the parameters better in the tight_layout
# todo: implement R2' calculations, sorting, and adding them to the plots
#       R'2 = 1 - MQres/MQTcorr; MQTcor = SQTcor / I-1; MQres = SQres / I-p
#       I: degrees of freedom. p: number of parameters in the model
#       SQTcor = sum(yi - mean(y)) ^ 2
#       SQres = chisqr
# todo: check why the R2 on the automatic linear fit is buggy

# todo: Check the results file, if the correct parameters are being recorded
# todo: check if converting GP and Eta to ndarrays at the beginning breaks anything.

# todo: Remove the prev_extracted setting and try to guess this parameter

# todo: remove the debugging setting. Just use the debugging tools.

class Fitter:
    def __init__(self, filename, settings, do_fit=True):
        self.VISC_LIMIT = 10000000
        self.l_first_point = 0
        self.l_last_point = -1
        self.nl_first_point = 0
        self.nl_last_point = -1
        self.manip = FileManip()
        self.filename = filename
        self.settings = settings
        self.model = self.settings.NL_FITTING_METHOD
        self.l_R2 = 0
        self.nl_R2 = 0
        self.wait = float(settings.WAIT)
        self.fixed_fp = settings.FIXED_FP_NL
        if not self.fixed_fp:
            self.max_fp = int(settings.MAX_FP_NL)
        else:
            self.max_fp = 0

        self.lin_done = False
        self.nl_done = False

        if self.model == 'Cross':
            #self.param_names = ['eta_0', 'eta_inf', 'GP_b', 'n']
            self.param_names = 'eta_0 eta_inf GP_b n'
        elif self.model == 'Carreau':
            #self.param_names = ['eta_0', 'eta_inf', 'GP_b', 'n']
            self.param_names = 'eta_0 eta_inf GP_b n'
        elif self.model == 'Carreau-Yasuda':
            #self.param_names = ['eta_0', 'eta_inf', 'lambda', 'a', 'n']
            self.param_names = 'eta_0 eta_inf lambda a n'
        else:
            raise NameError(f'Did not understand model {self.model}')

        self.param_names_lin = ['Int', 'Slp']  # todo: check if this is the correct order.

        if self.settings.DO_LIN:
            self.int = 50
            self.int_err = 0.1
            self.slp = 0
            self.slp_err = 0
        elif self.settings.DO_NL:
            if self.model == 'Carreau' or self.model == 'Cross' or self.model == 'Carreau-Yasuda':
                self.params = [0, 0, 0, 0]
                self.param_errs = [0, 0, 0, 0]
            else:
                raise ValueError(f'Unknown model: {model}')

        try:
            if self.settings.PREV_EXTRACTED:
                self.GP, self.Eta = self.manip.ExtractData_pd(filename)
                self.GP = np.array(self.GP)
                self.Eta = np.array(self.Eta)
            else:
                self.GP, self.Eta = self.manip.ExtractData(filename)
                self.GP = np.array(self.GP)
                self.Eta = np.array(self.Eta)
        except ValueError:
            self.manip.logger(filename, 'Failed to open')
            raise ValueError(f'!!!! No Flow Curve data was found! Re-export the data on file{filename}.')
        except KeyError:
            self.manip.logger(filename, 'Failed to open')
            raise ValueError(f'!!!! No Flow Curve data was found! Re-export the data on file{filename}.')

        if len(self.GP) != len(self.Eta):
            self.manip.logger(self.filename, 'Failed to open')
            raise ValueError(f'!!!! GP and Eta have different lengths. '
                             f'Re-export {filename} or fix the problem manually.')

        if do_fit:
            self._fit()

    def _fit(self):  # Uses fit_curve. Does not provide an R2 value.
        if self.settings.DO_LIN:
            if self.settings.AUTO_LIN:
                self.automatic_lin_fitting(True)
            else:  # todo: plot, save and ask for the required points
                self.manual_fit(0, -1, 'Linear')
        if self.settings.DO_NL:
            if self.settings.AUTO_NL:
                self.automatic_nl_fitting(True)
            else:
                self.manual_fit(0, -1, self.settings.NL_FITTING_METHOD, True)

    def fit(self):
        if self.settings.DO_LIN:
            if self.settings.AUTO_LIN:
                self.automatic_lin_fitting_lm(True)
            else:  # todo: plot, save and ask for the required points
                self.manual_fit(0, -1, 'Linear')
        if self.settings.DO_NL:
            if self.settings.AUTO_NL:
                self.automatic_nl_fitting_lm(True)
            else:
                self.manual_fit(0, -1, self.settings.NL_FITTING_METHOD, True)

    @staticmethod
    def fit_Carreau(GP, eta_0, eta_inf, GP_b, n):
        """Eta = eta_inf + (eta_0 - eta_inf) / (1+(GP/GP_b)**2)**(n/2)
        GP_b is a constant with the dimension of time and n is a dimensionless constant"""
        return eta_inf + (eta_0 - eta_inf) / (1 + (GP / GP_b) ** 2) ** (n / 2)

    @staticmethod
    def fit_Cross(GP, eta_0, eta_inf, GP_b, n):
        return eta_inf + (eta_0 - eta_inf) / (1 + (GP / GP_b) ** n)

    @staticmethod
    def fit_PowerLaw(GP, k, n):
        """Power Law: eta = k * GP ** (n-1)"""
        return k * GP ** (n - 1)

    @staticmethod
    def fit_CarreauYasuda(GP, eta_0, eta_inf, lbda, a, n):
        """Carreau-Yasuda: eta(GP) = eta_inf + (eta_0 - eta_inf)(1+(lambda * GP)**a)**((n-1)/a)"""
        return eta_inf + (eta_0 - eta_inf) / (1 + (lbda * GP) ** a) ** ((n - 1) / a)

    @staticmethod
    def fit_lin(x, a, b):
        """Simple function for a linear fit, with a as the linear coefficient and b the angular coefficient."""
        return a + b * x

    @staticmethod
    def carr_uncertainty(GP, eta0, etainf, GPb, n, eta0_err, etainf_err, GPb_err, n_err):
        """Uses the uncertainty package to calculate the Carreau model values. GP
        can be a numpy array, which returns two lists of values and errors, a float64,
        float or int and returns a tuple (val, err)"""
        f_eta0 = ufloat(eta0, eta0_err)
        f_etainf = ufloat(etainf, etainf_err)
        f_GPb = ufloat(GPb, GPb_err)
        f_n = ufloat(n, n_err)
        Carr = f_etainf + (f_eta0 - f_etainf) / (1 + (GP / f_GPb) ** 2) ** (f_n / 2)

        # Extracts all val +- err pairs if GP is an ndarray
        if type(GP) is np.ndarray:
            Carr_val = [a.nominal_value for a in Carr]
            Carr_err = [a.std_dev for a in Carr]

        # If GP is numeric, separates the two values.
        if (type(GP) is np.float64) or (type(GP) is float) or (type(GP) is int):
            Carr_val = Carr.nominal_value
            Carr_err = Carr.std_dev

        return Carr_val, Carr_err

    @staticmethod
    def cross_uncertainty(GP, eta0, etainf, GPb, n, eta0_err, etainf_err, GPb_err, n_err):
        f_eta0 = ufloat(eta0, eta0_err)
        f_etainf = ufloat(etainf, etainf_err)
        f_GPb = ufloat(GPb, GPb_err)
        f_n = ufloat(n, n_err)
        Cross = f_etainf + (f_eta0 - f_etainf) / (1 + (GP / f_GPb) ** f_n)

        # Extracts all val +- err pairs if GP is an ndarray
        if type(GP) is np.ndarray:
            Cross_val = [a.nominal_value for a in Cross]
            Cross_err = [a.std_dev for a in Cross]

        # If GP is numeric, separates the two values.
        if (type(GP) is np.float64) or (type(GP) is float) or (type(GP) is int):
            Cross_val = Cross.nominal_value
            Cross_err = Cross.std_dev

        return Cross_val, Cross_err

    @staticmethod
    def carryas_uncertainty(GP, eta0, etainf, lbda, a, n, eta0_err, etainf_err, lbda_err, a_err, n_err):
        f_eta0 = ufloat(eta0, eta0_err)
        f_etainf = ufloat(etainf, etainf_err)
        f_n = ufloat(n, n_err)
        f_lbda = ufloat(lbda, lbda_err)
        f_a = ufloat(a, a_err)
        CY = f_etainf + (f_eta0 - f_etainf) * (1 + (f_lbda * GP) ** f_a) ** ((f_n - 1) / f_a)

        # Extracts all val +- err pairs if GP is an ndarray
        if type(GP) is np.ndarray:
            CY_val = [a.nominal_value for a in CY]
            CY_err = [a.std_dev for a in CY]

        # If GP is numeric, separates the two values.
        if (type(GP) is np.float64) or (type(GP) is float) or (type(GP) is int):
            CY_val = CY.nominal_value
            CY_err = CY.std_dev

        return CY_val, CY_err

    # todo: alterar o termo int para outro valor para impedir que haja um clash.
    def lm_curvefit(self, GP, Eta, do_lin=False):
        params = Parameters()
        SStot = sum((Eta - np.mean(Eta)) ** 2)
        if do_lin:  # todo: Check why R2 is very weird here.
            params.add('Int', 50, vary=True, min=0)
            params.add('Slp', 0, vary=False)
            fit = minimize(self.residual_lin, params, args=(GP, Eta))
            slp = fit.params['Slp'].value
            int = fit.params['Int'].value
            slp_err = fit.params['Slp'].stderr
            int_err = fit.params['Int'].stderr
            chisqr = fit.chisqr
            R2 = 1 - fit.chisqr / SStot
            return [slp, int], [slp_err, int_err], R2
        elif self.model == 'Carreau':
            params.add('eta_0', 100, vary=True, min=0)
            params.add('eta_inf', 1, vary=True, min=0)
            params.add('GP_b', 5, vary=True, min=0)
            params.add('n', 1, vary=True, min=0)
            fit = minimize(self.residual, params, args=(GP, Eta))
            params = [fit.params[par].value for par in fit.params]
            param_errs = [fit.params[par].stderr for par in fit.params]
            R2 = 1 - fit.chisqr / SStot
            return params, param_errs, R2
        elif self.model == 'Cross':
            params.add('eta_0', 100, vary=True, min=0)
            params.add('eta_inf', 1, vary=True, min=0)
            params.add('GP_b', 5, vary=True, min=0)
            params.add('n', 1, vary=True, min=0)
            fit = minimize(self.residual, params, args=(GP, Eta))
            params = [fit.params[par].value for par in fit.params]
            param_errs = [fit.params[par].stderr for par in fit.params]
            R2 = 1 - fit.chisqr / SStot
            return params, param_errs, R2
        elif self.model == 'Carreau-Yasuda':
            params.add('eta_0', 100, vary=True, min=0)
            params.add('eta_inf', 1, vary=True, min=0)
            params.add('lbda', 5, vary=True, min=0)
            params.add('a', 1, vary=True, min=0)
            params.add('n', 1, vary=True, min=0)
            fit = minimize(self.residual, params, args=(GP, Eta))
            params = [fit.params[par].value for par in fit.params]
            param_errs = [fit.params[par].stderr for par in fit.params]
            SSres = fit.chisqr
            R2 = 1 - SSres / SStot
            return params, param_errs, R2

    def residual(self, params, x, dataset):
        if self.model == 'Carreau':
            mod = self.fit_Carreau(x, params['eta_0'], params['eta_inf'], params['GP_b'], params['n'])
        elif self.model == 'Cross':
            mod = self.fit_Cross(x, params['eta_0'], params['eta_inf'], params['GP_b'], params['n'])
        elif self.model == 'Carreau-Yasuda':
            mod = self.fit_CarreauYasuda(x, params['eta_0'], params['eta_inf'], params['lbda'], params['a'],
                                         params['n'])
        resid = dataset - mod
        return resid

    def residual_lin(self, params, x, dataset):
        if type(x) == list:
            x = np.array(x)
        mod = params['Int'] + params['Slp'] * x
        resid = dataset - mod
        return resid

    def automatic_lin_fitting_lm(self, save=True):
        length = len(self.GP)
        fittings = []

        # Go through several possible ranges to fit, and fit them, then get the best fit
        for first_point in range(0, length//3, 1):
            for last_point in range(first_point + 3, length // 2, 1):
                GP_arr = np.array(self.GP[first_point:last_point + 1])  # todo: check if this conversion is necessary
                Eta_arr = np.array(self.Eta[first_point:last_point + 1])
                try:
                    #popt, pcov = curve_fit(self.fit_lin, GP_arr, Eta_arr, p0=(30, 0),
                    #                   bounds=(0, [self.VISC_LIMIT, 0.0001]))
                    params, param_errs, R2 = self.lm_curvefit(GP_arr, Eta_arr, do_lin=True)
                except:  # todo: test here and find what types of errors can occur
                    print(f'Error while using linear fit for file {self.filename}')
                    print(traceback.format_exc())
                    self.manip.logger(self.filename, 'Generic')

                #perr = np.sqrt(np.diag(pcov))
                fittings.append((first_point, last_point, params, param_errs, R2))

        if self.settings.LIN_SORTING_METHOD == 'by_error':
            fittings.sort(key=lambda x: np.log(x[3][0]))
        elif self.settings.LIN_SORTING_METHOD == 'by_error_length':
            fittings.sort(key=lambda x: np.log(x[2][1]) / (x[1] - x[0]))
        elif self.settings.LIN_SORTING_METHOD == 'by_R2':
            fittings.sort(key=lambda x: x[4])

        self.l_first_point = fittings[0][0]
        # todo: add variable names to first and last points of linear and nl
        self.l_last_point = fittings[0][1]
        self.int = fittings[0][2][1]
        self.int_err = fittings[0][3][1]
        self.l_R2 = fittings[0][4]
        self.lin_done = True

        if self.settings.DEBUG:
            print('Debug: fittings_sorted: ', fittings)
            print('Debug: a: ', self.int)
            print('Debug: aerr: ', self.int_err)

        if save:
            self.manip.record_fit('linear', self.int, self.int_err, silent=False,
                                  extra=f"{fittings[0][0]};{fittings[0][1]};")
            # todo: check how this was done before, for consistency

        return self.int, self.int_err, self.l_R2

    def automatic_lin_fitting(self, save=True):
        """Goes through all the files, fits them and selects the best fit according to two algorithms.
        First, it selects two points, a beginning and an end point, the first starting at point 0
        and going to a third of the curve. The second, starting at points to the right,
        going until the middle of the curve.
        Then, it fits the data by fixing the slope at 0 and goes through every possible combination
        of the first and second points.
        It selects the data based on two criteria:
        1. sorting = 'by_error': finds the minimal error. Tends to select less points overall and
            gives a fitting with a less than ideal representation overall.
        2. sorting = 'by_error_length': divides the error by how many points were used in the fit.
            May result in a higher overall error, but gives a better representation of the curve.
        """

        length = len(self.GP)
        fittings = []

        # Go through several possible ranges to fit, and fit them, then get the best fit
        for first_point in range(0, length//3, 1):
            for last_point in range(first_point + 3, length // 2, 1):
                GP_arr = np.array(self.GP[first_point:last_point + 1])
                Eta_arr = np.array(self.Eta[first_point:last_point + 1])
                try:
                    popt, pcov = curve_fit(self.fit_lin, GP_arr, Eta_arr, p0=(30, 0),
                                       bounds=(0, [self.VISC_LIMIT, 0.0001]))
                except:  # todo: test here and find what types of errors can occur
                    print(f'Error while using linear fit for file {self.filename}')
                    print(traceback.format_exc())
                    self.manip.logger(self.filename, 'Generic')

                perr = np.sqrt(np.diag(pcov))
                fittings.append((first_point, last_point, popt, perr))

        if self.settings.LIN_SORTING_METHOD == 'by_error':
            fittings.sort(key=lambda x: np.log(x[3][0]))  # gets perr of eta_0
        elif self.settings.LIN_SORTING_METHOD == 'by_error_length':
            fittings.sort( key=lambda x: np.log(x[3][0]) / (x[1] - x[0]) ) # divides perr by last-first

        self.int = fittings[0][2][0]
        self.int_err = fittings[0][3][0]
        self.l_first_point = fittings[0][0]  # todo: add variable names to first and last points of linear and nl
        self.l_last_point = fittings[0][1]
        self.lin_done = True

        if self.settings.DEBUG:
            print('Debug: fittings_sorted: ', fittings)
            print('Debug: a: ', self.int)
            print('Debug: aerr: ', self.int_err)

        if save:
            self.manip.record_fit(self.filename, self.int, self.int_err, silent=False,
                                  extra=f"{fittings[0][0]};{fittings[0][1]};")

        return self.int, self.int_err

    # todo: change from curve_fit to lm_fit.
    # todo: calculate R2 for all fittings and add it in the end to the class
    # todo: add options to sort by R2.

    def automatic_nl_fitting_lm(self, save=True):
        fittings = []
        try:
            max_range = len(self.GP) // self.max_fp
        except ZeroDivisionError:
            max_range = 1

        for first_point in range(0, max_range, 1):
            GP_arr = np.array(self.GP[first_point:])
            Eta_arr = np.array(self.Eta[first_point:])
            nonlinear_has_error = ''
            try:
                params, param_errs, R2 = self.lm_curvefit(GP_arr, Eta_arr, do_lin=False)
            except FloatingPointError:  # todo: check if these exceptions work
                print('!!!! Overflow detected on one of the parameters. Could not determine all parameters')
                nonlinear_has_error = ';param_overflow_during_fitting'
                self.manip.logger(self.filename, 'Overflow')
            except RuntimeError:
                print('!!!! Overflow detected on one of the parameters. Could not determine all parameters')
                nonlinear_has_error = ';param_overflow_during_fitting'
                self.manip.logger(self.filename, 'Overflow')
            except OverflowError:
                print('!!!! Overflow detected on one of the parameters.')
                self.manip.logger(self.filename, 'Overflow')

            fittings.append((first_point, params, param_errs, R2))

        if self.settings.NL_SORTING_METHOD == 'eta_0':
            fittings.sort(key=lambda x: x[2][0])
        elif self.settings.NL_SORTING_METHOD == 'overall':
            # fittings.sort(key=lambda x: x[2][0] + x[2][1] + x[2][2] + x[2][3])
            fittings.sort(key=lambda x: sum(x[2]))  # sums the errors
        elif self.settings.NL_SORTING_METHOD == 'R2':
            fittings.sort(key=lambda x: x[3])
        else:
            raise ValueError(f'Could not understand the sorting method {self.settings.NL_SORTING_METHOD}')

        self.nl_first_point = fittings[0][0]
        self.params = fittings[0][1]
        self.param_errs = fittings[0][2]
        self.nl_R2 = fittings[0][3]

        if save:  # todo: check here to return a good destination file
            try:
                self.manip.record_fit(
                    self.filename, self.params[0],
                    self.param_errs[0], silent=False,
                    extra=f"{fittings[0][0]};{fittings[0][1]};nonlinear_auto_{self.settings.NL_FITTING_METHOD};"
                          f"{nonlinear_has_error}", fdest_name=self.settings.NL_FITTING_METHOD + '.csv'
                )
            except UnboundLocalError:
                print('Unable to write to file because the subroutine did not return the fitting parameters')
                print(traceback.format_exc())
                self.manip.record_fit(self.filename, 0, 0, extra=f'nonlinear_auto_{self.settings.NL_FITTING_METHOD};'
                                                                 f'unable_to_find_viscosity',
                                      fdest_name=self.settings.NL_FITTING_METHOD + '.csv')
                self.manip.logger(self.filename, 'No Viscosity')

        self.nl_done = True
        return self.nl_first_point, self.params, self.param_errs, self.nl_R2

    def automatic_nl_fitting(self, save=True):
        fittings = []
        try:
            max_range = len(self.GP) // self.max_fp
        except ZeroDivisionError:
            max_range = len(self.GP)

        for first_point in range(0, max_range, 1):
            GP_arr = np.array(self.GP[first_point:])
            Eta_arr = np.array(self.Eta[first_point:])
            nonlinear_has_error = ''
            try:
                if self.settings.NL_FITTING_METHOD == 'Carreau':
                    popt, pcov = curve_fit(self.fit_Carreau, GP_arr, Eta_arr, bounds=(0, np.inf))
                elif self.settings.NL_FITTING_METHOD == 'Cross':
                    popt, pcov = curve_fit(self.fit_Cross, GP_arr, Eta_arr, bounds=(0, np.inf))
                elif self.settings.NL_FITTING_METHOD == 'Carreau-Yasuda':
                    popt, pcov = curve_fit(self.fit_CarreauYasuda, GP_arr, Eta_arr)
                else:
                    raise ValueError(f'Model not present: {self.settings.NL_FITTING_METHOD}')

            except FloatingPointError:
                print('!!!! Overflow detected on one of the parameters. Could not determine all parameters')
                nonlinear_has_error = ';param_overflow_during_fitting'
                self.manip.logger(self.filename, 'Overflow')
                continue
            except RuntimeError:
                print('!!!! Overflow detected on one of the parameters. Could not determine all parameters')
                nonlinear_has_error = ';param_overflow_during_fitting'
                self.manip.logger(self.filename, 'Overflow')
                continue
            except OverflowError:
                print('!!!! Overflow detected on one of the parameters.')
                self.manip.logger(self.filename, 'Overflow')
                continue

            perr = np.sqrt(np.diag(pcov))
            fittings.append((first_point, popt, perr))

            if self.settings.DEBUG:
                fitting_params_str = ' '.join([str(round(i, 2)) + '+/-' +
                                               str(round(j,2)) for i, j in zip(popt, perr) ])
                                                # 'a+/-aerr b+/-berr ...'
                print(f"{self.settings.NL_FITTING_METHOD} fitting: {fitting_params_str}")

        if self.settings.NL_SORTING_METHOD == 'eta_0':
            fittings.sort(key=lambda x: x[2][0])
        elif self.settings.NL_SORTING_METHOD == 'overall':
            #fittings.sort(key=lambda x: x[2][0] + x[2][1] + x[2][2] + x[2][3])
            fittings.sort(key=lambda x: sum(x[2])) # sums the errors
        else:
            raise ValueError(f'Could not understand the sorting method {self.settings.NL_SORTING_METHOD}')

        self.nl_first_point = fittings[0][0]
        self.params = fittings[0][1]
        self.param_errs = fittings[0][2]

        if save:  # todo: check here to return a good destination file
            try:
                self.manip.record_fit(
                    self.filename, self.params[0],
                    self.param_errs[0], silent=False,
                    extra=f"{fittings[0][0]};{fittings[0][1]};nonlinear_auto_{self.settings.NL_FITTING_METHOD};"
                          f"{nonlinear_has_error}", fdest_name=self.settings.NL_FITTING_METHOD + '.csv'
                )
            except UnboundLocalError:
                print('Unable to write to file because the subroutine did not return the fitting parameters')
                print(traceback.format_exc())
                self.manip.record_fit(self.filename, 0, 0, extra=f'nonlinear_auto_{self.settings.NL_FITTING_METHOD};'
                                                                 f'unable_to_find_viscosity',
                                      fdest_name=self.settings.NL_FITTING_METHOD+'.csv')
                self.manip.logger(self.filename, 'No Viscosity')

        self.nl_done = True
        return self.nl_first_point, self.params, self.param_errs

    # TODO: check if the bounds are correct
    # TODO: increment this function to be able to accept multiple fittings
    def manual_fit(self, first, last, fit_types, save=True):
        GP_arr = np.array(self.GP[first:last + 1])
        Eta_arr = np.array(self.Eta[first:last + 1])
        fittings = []

        for type in fit_types:
            if 'Linear' in type:
                popt, pcov = curve_fit(self.fit_lin, GP_arr, Eta_arr, p0=(30, 0),
                                       bounds=(0, [self.VISC_LIMIT, 0.0001]))
            elif 'Carreau' in type:
                popt, pcov = curve_fit(self.fit_Carreau, GP_arr, Eta_arr, p0=(30, 0),
                                       bounds=(0, np.inf))
            elif 'Cross' in type:
                popt, pcov = curve_fit(self.fit_Carreau, GP_arr, Eta_arr, p0=(30, 0),
                                       bounds=(0, np.inf))
            elif 'Carreau-Yasuda' in type:
                popt, pcov = curve_fit(self.fit_CarreauYasuda, GP_arr, Eta_arr, p0=(30, 0),
                                       bounds=(0, np.inf))
            else:
                raise NameError(f'Could not understand the list fit_types {fit_types}')

            perr = np.sqrt(np.diag(pcov))

            self.params = popt  # Will be continuously overwritten. todo: will this be a problem?
            self.param_errs = perr

            if self.settings.DEBUG:
                # 'a+/-aerr b+/-berr ...'
                fitting_params_str = ' '.join([str(round(i, 2)) + '+/-' +
                                               str(round(j, 2)) for i, j in zip(popt, perr)])
                print(f"{fit_types} fitting: {fitting_params_str}")

            if save:  # todo: check here to return a good destination file
                self.manip.record_fit(self.settings.NL_FITTING_METHOD, self.params,
                                      self.param_errs, silent=False)

            fittings.append((type, popt, perr))

        return fittings

    # todo: use AnchoredText instead of text to write the fitting results;
    # from mpl_toolkits.axes_grid.anchored_artists import AnchoredText
    # at = AnchoredText("Figure 1a",
    #                   prop=dict(size=8), frameon=True,
    #                   loc=2,
    #                   )
    # at.patch.set_boxstyle("round,pad=0.,rounding_size=0.2")
    # ax.add_artist(at)
    def plot_error_graphs(self): # todo: If it has both plots, make them side by side
        from mpl_toolkits.axes_grid.anchored_artists import AnchoredText
        TEXT_FILENAME_X = 0.1
        TEXT_PARAMS_X = 0.3
        TEXT_Y = 0.98
        x = np.logspace(np.log10(self.GP[0]), np.log10(self.GP[-1]))

        if self.settings.DEBUG:
            print('Debug: x', x)
            print('Debug: params', self.params)
            print('Debug: GP', self.GP, 'Eta', self.Eta)

        if self.nl_done:
            if self.model == 'Carreau':
                y, yerr = self.carr_uncertainty(x, *self.params, *self.param_errs)
            elif self.model == 'Cross':
                y, yerr = self.cross_uncertainty(x, *self.params, *self.param_errs)
            elif self.model == 'Carreau-Yasuda':
                y, yerr = self.carryas_uncertainty(x, *self.params, *self.param_errs)
        if self.lin_done:
            y_l, yerr_l = np.ones(len(x)) * self.int, np.ones(len(x)) * self.int_err
            # Creates a horizontal line with n points

        if self.nl_done and self.lin_done:
            fig, [axn, axl] = plt.subplots(ncols=2, nrows=1, figsize=(12, 4))
        elif self.nl_done and not self.lin_done:
            fig, axn = plt.subplots(ncols=1, nrows=1, figsize=(6, 4))
        elif not self.nl_done and not self.lin_done:
            fig, axl = plt.subplot(ncols=1, nrows=1, figsize=(6, 4))

        if self.nl_done:
            axn.set_xscale('log')
            axn.set_yscale('log')
            axn.plot(self.GP, self.Eta, linewidth=0, marker='o', markersize=5)
            axn.errorbar(x, y, yerr=yerr)
            axn.annotate(str(self.nl_first_point + 1), (self.GP[self.nl_first_point], self.Eta[self.nl_first_point]), color='red')
            if self.nl_last_point == -1:
                axn.annotate(str(len(self.GP)), (self.GP[self.nl_last_point], self.Eta[self.nl_last_point]),
                             color='red')  # todo: check this function
            else:
                axn.annotate(str(self.nl_last_point), (self.GP[self.nl_last_point], self.Eta[self.nl_last_point]), color='red')
            model_param_names = 'Model: ' + self.model + ' Params: ' + self.param_names
            param_text = " ".join([str(round(par, 2)) + '+/-' + str(round(err, 2))
                                   for par, err in zip(self.params, self.param_errs)])

            total_text = f'{self.filename}\n{model_param_names}\n{param_text}\n$R^2$={round(self.nl_R2, 2)}'
            anchored_text = AnchoredText(total_text, loc=3, frameon=True, prop={'fontsize':'small'})
            axn.add_artist(anchored_text)

        if self.lin_done:
            axl.set_xscale('log')
            axl.set_yscale('log')
            axl.plot(self.GP, self.Eta, linewidth=0, marker='o', markersize=5)
            axl.errorbar(x, y_l, yerr=yerr_l)
            axl.annotate(str(self.l_first_point + 1), (self.GP[self.l_first_point], self.Eta[self.l_first_point]), color='red')
            if self.l_last_point == -1:
                axl.annotate(str(len(self.GP)), (self.GP[self.l_last_point], self.Eta[self.l_last_point]),
                             color='red')  # todo: check this function
            else:
                axl.annotate(str(self.l_last_point), (self.GP[self.l_last_point], self.Eta[self.l_last_point]), color='red')

            model_param_names = 'Model: Linear. Params: Intercept'
            param_text = f"int = {self.int}+/-{self.int_err}"
            total_text = f'{self.filename}\n{model_param_names}\n{param_text}\n$R^2$={round(self.l_R2, 2)}'
            anchored_text = AnchoredText(total_text, loc=3, frameon=True, prop={'fontsize':'small'})
            axl.add_artist(anchored_text)

        plt.tight_layout()

        if self.settings.SAVE_GRAPHS:
            fig.savefig(self.filename[:-4] + '.png')
            print('Figure saved.')
        if not self.settings.INLINE_GRAPHS and self.settings.PLOT_GRAPHS:
            plt.draw()
            plt.pause(self.wait)
            #plt.clf()
            plt.close(fig)
        elif self.settings.PLOT_GRAPHS:
            plt.show()
        return

class FileManip:
    #def __init__(self, sett):
    #    self.settings = sett

    @staticmethod
    def ExtractData(fname, FC_segment=0):
        """Opens the file fname and extracts the data based on where it finds the word 'Eta' and 'GP', these being
        the Viscosity and the Shear Rate (gamma point). If the file has multiple segments, for example, when multiple
        experiments were done in succession, FC_segment indicates which of those experiments was a Flow Curve."""
        fhand = open(fname, 'r')
        GP = []
        Eta = []
        column_eta = 0
        column_gp = 0
        # FC_segment = '3'

        # while FC_segment == 0:
        #    FC_segment = input("What is the segment that has the flow curves? (eg. [1], 2, 3) If you do not know, don't write anything. ")
        #    if FC_segment == '':
        #        print(fhand.read())
        #    elif FC_segment.isnumeric():
        #        break
        #    else:
        #        print('Not a valid number')

        for line in fhand:
            if line.startswith(';'):
                column_names = line.rstrip().split(';')
                # if settings['DEBUG']:
                #     print('Debug: column names', column_names)
                for i, column in enumerate(column_names):
                    if 'Eta' in column and 'Eta*' not in column:
                        column_eta = i
                        #if settings['DEBUG']:
                        #    print('Debug: Found Eta at', column_eta)
                    if 'GP' in column:
                        column_gp = i
                        #if settings['DEBUG']:
                        #    print('Debug: Found GP at', column_gp)
            try:
                GP.append(float(line.replace(',', '.').split(';')[column_gp]))
                Eta.append(float(line.replace(',', '.').split(';')[column_eta]))
            except:
                pass

            # if line.startswith(FC_segment + '|'):
            #    line = line.rstrip()
            #    num, gp, tau, eta, *rest = line.replace(',','.').split(';')
            #    GP.append(float(gp))
            #    Eta.append(float(eta))
            #    #print(line)

        fhand.close()
        if len(GP) == 0:
            # print('!!!!No Flow Curve data was found! Re-export the data on file', fname)
            raise ValueError
        # return pd.Series(GP), pd.Series(Eta)
        # if settings['DEBUG']:
        #     print('Debug: Extracted Data: GP:', GP, 'Eta:', Eta)
        return GP, Eta

    @staticmethod
    def ExtractData_pd(fname):
        """Uses pandas do extract the data if it was exported using the data extraction tool"""
        pd_temp = pd.read_csv(fname, delimiter=';', encoding='latin1', decimal=',')
        pd_temp = pd_temp[pd_temp > 0].dropna()

        col_GP = ''
        col_Eta = ''

        for col in pd_temp.columns:
            if 'GP' in col:
                col_GP = col
                # print('achou GP em', col)
            if 'Eta' in col:
                col_Eta = col
                # print('achou Eta em', col)

        GP = pd_temp[col_GP].tolist()
        Eta = pd_temp[col_Eta].tolist()
        return GP, Eta

    @staticmethod
    def record_fit(name, eta0, eta0_err, silent=False, extra='', fdest_name='results.csv'):
        if not silent:
            print(f"{name}: Intercept={eta0} +- {eta0_err}. Extra={extra}")
            #print(name + ':', 'Intercept', eta0, '+-', eta0_err, extra)

        with open(fdest_name, 'a', encoding='utf-8') as fdest:
            #fdest.write(name + ';' + str(eta0) + ';' + str(eta0_err) + ';' + extra + '\n')
            fdest.write(f"{name};{eta0};{eta0_err};{extra}\n")

    @staticmethod
    def select_files():
        files = []
        extension = input('What is the file extension? txt, dat, etc:\n')
        allfiles = glob.glob('*.' + extension)
        print(*[str(num) + ')' + file + '\n' for num, file in enumerate(allfiles)], sep='')
        while True:
            file_to_add = input('Which file to add? Number, nothing to continue or "quit" to exit: ')
            if file_to_add == 'quit':
                return []
            elif file_to_add == '':
                break
            else:
                try:
                    files.append(allfiles[int(file_to_add)])
                except IndexError:
                    print('Invalid value')
                except ValueError:
                    print('Enter a number, not text')
        if len(files) == 0:
            print('No file was selected! The program will now quit.')
            return files
        print('====Selected files:====')
        print(*[file + '\n' for file in files], sep='', end='')
        print('=======================')
        return files

    @staticmethod
    def logger(file, type, extra=''):
        with open('log', 'a') as log:
            if type == 'Overflow':
                log.write(f'Parameter overflow while trying to fit file {file}: {extra}\n')
            if type == 'No Viscosity':
                log.write(f'Unable to find viscosity for file {file}\n')
            if type == 'Failed to open':
                log.write(f'Failed to open file {file}. Re-export the data.')
            else:  # type == 'Generic'
                log.write(f'Error while processing {file}: {extra}\n')

def test():
    settings = Settings.Settings()
    settings.NL_FITTING_METHOD = 'Carreau-Yasuda'
    filename = 'CF_Sac50-3--0.csv'
    fit = Fitter(filename, settings, do_fit=False)
    fit.automatic_nl_fitting_lm(save=True)
    print(fit.model, fit.nl_R2, *fit.params)

    return fit

def main():
    settings = Settings.Settings()
    manip = FileManip()
    settings.print_settings()
    do_change = input('Do you want to change the settings? y/[n]')
    if do_change == 'y':
        settings.edit_settings()

    if settings.TREAT_ALL:
        files = glob.glob(f'*.{settings.EXT}')
        if len(files) == 0:
            print(f'No files with the extension {settings.EXT} found.'
                  f' Please select them manually or change EXT accordingly.')
            files = manip.select_files()

    else:
        files = manip.select_files()

    if len(files) == 0:
        print('No files selected. Quitting.')
        sys.exit()

    for file in files:
        try:
            fit = Fitter(file, settings, do_fit=True)
        except ValueError:  # todo: debug and check what would be needed here.
            print(f'Skipping {file}: Value Error')
            continue
        except KeyError:
            print(f'Skipping {file} Key Error')
            continue
            #print(traceback.format_exc())

        if settings.PLOT_GRAPHS or settings.SAVE_GRAPHS:
            try:
                fit.plot_error_graphs()
            except OverflowError:  # todo: write which parameter has overflown
                print('!!!! Overflow detected on one of the parameters. Could not plot the data')
                nonlinear_has_error = ';param_overflow_during_fitting'
                # todo: log this
            except UnboundLocalError:
                print('Not able to write to file because the subroutine did not return the fitting parameters')
                # todo: log this

            #fit.plot_error_graphs(file[:-4] + '_lin_' + file[-4:], fit.params, fit.first_point, fit.last_point,
            #                      model=fit.settings.NL_FITTING_METHOD, param_names=[''])

                # # except:
                # #     print('Error found while plotting the linear fit')
                # #     print(traceback.format_exc())
                # #     lin_has_error = 'error_during_fitting'
                # record(file, a, aerr,
                #        extra='linear automatic;FP=' + str(lin_points[0]) + 'LP=' + str(lin_points[1]) +
                #              lin_has_error, fdest_name='linear.csv')

            #         if settings['PLOT_GRAPHS']:
            #             plot_error_graphs(file[:-4] + '_carr_' + file[-4:], GP, Eta,
            #                               params=np.concatenate((popt, perr)),
            #                               first_point=nl_first, model=settings['NL_FITTING_METHOD'],
            #                               param_names=Param_names_errs[settings['NL_FITTING_METHOD']])
            #     except OverflowError:  # todo: write which parameter has overflown
            #         print('!!!! Overflow detected on one of the parameters. Could not plot the data')
            #         nonlinear_has_error = ';param_overflow_during_fitting'
            #     try:
            #         record(file, popt[0], perr[0], extra='nonlinear_auto_' + settings['NL_FITTING_METHOD'] +
            #                                              nonlinear_has_error,
            #                fdest_name=settings['NL_FITTING_METHOD'] + '.csv')
            #     except UnboundLocalError:
            #         print('Not able to write to file because the subroutine did not return the fitting parameters')
            #         record(file, 0, 0, extra='nonlinear_auto_' + settings['NL_FITTING_METHOD'] + ';' +
            #                                  'unable_to_find_viscosity',
            #                fdest_name=settings['NL_FITTING_METHOD'] + '.csv')
            #         with open('log', 'a') as log:
        #     #             log.write('Unable to find viscosity for file ' + file + '\n')
        # except:
        #     print('!!!!We have encountered an error while processing file', file)
        #     print(traceback.format_exc())
        #     with open('log', 'a') as log:
        #         log.write('Error while processing ' + file + '\n')
        #         log.write(traceback.format_exc())

if __name__ == '__main__':
    fit = main()
    #main()


def EDMFanalysis(numberofsensor,numberofmodelinstance,meanuncertaintydata,stduncertaintydata,ID_to_be_excluded,finalpathprediction,finalpathmeas,sidakvalue,uncmultiplier,indextext):
    """
        Perform EDMF analysis with all uploaded files and intermediate results

        Developed by :   <NAME> (ETH Singapore)
        Contact :        <EMAIL>
        Date:            August 03, 2021

        INPUTS:
            numberofsensor : number of sensors involved in the studied problem.
            numberofmodelinstance : number of initial model instances involved in the studied problem.
            ID_to_be_excluded: index of measurements to be excluded for EDMF analysis.
            finalpathprediction : directory of the Prediction(Geo).xlsx.
            finalpathmeas: directory of the Measurements(Geo).xlsx.
            sidakvalue: target reliability of identification
            uncmultiplier : multiplier for uncertainty to assess sensitivity of results to uncertainty definitions
            indextext: index of measurements to be excluded for EDMF analysis in the original format.
        OUTPUTS:
            prediction : predictions made with candidate models.
            measurement : measurement data.
            finalresults: logic results of initial model instances: 0 represents falsified model; 1 represents candidate model
            IMS: initial model instances
            CMS: candidate models
            CMSID: location reference of candidate models
        NOTE:
            These inputs will be automatically read from other functions. No maunal actions are needed.
        """

    import numpy
    import os
    import xlrd
    from openpyxl import load_workbook
    import pandas
    from scipy.stats import norm

    #########for validation
    fullset = numpy.arange(numberofsensor)+1
    ID_to_be_included_ = set(ID_to_be_excluded) ^ set(fullset)
    ID_to_be_included = numpy.zeros((len(ID_to_be_included_), 1))
    for i in range(0, len(ID_to_be_included_)):
        ID_to_be_included[i][0] = (list(ID_to_be_included_)[i])
    ID_to_be_included = ID_to_be_included.astype(numpy.int64)-1

    #####load uncertainty excel
    excelfileuncertainty1_ = meanuncertaintydata
    excelfileuncertainty2_ = stduncertaintydata

    excelfileprediction = pandas.read_excel(finalpathprediction,sheet_name = 'Prediction', engine = 'openpyxl')

    excelfilemeas = pandas.read_excel(finalpathmeas,header = None,sheet_name = 'Measurement', engine = 'openpyxl')

    ##########read uncertainty
    excelsheetuncertainty1 = excelfileuncertainty1_[0:numberofmodelinstance, 0:numberofsensor]
    ucombinemean = excelsheetuncertainty1

    excelsheetuncertainty2 = excelfileuncertainty2_[0:numberofmodelinstance, 0:numberofsensor]
    ucombinesd = excelsheetuncertainty2

    ###########calculate sidak
    lowerbound = numpy.zeros((numberofmodelinstance, numberofsensor))
    upperbound = numpy.zeros((numberofmodelinstance, numberofsensor))

    sidak = sidakvalue ** (1 / len(ID_to_be_included))

    for i in range(0,numberofmodelinstance):
        for s in range(0,numberofsensor):
            bound1_ = norm.ppf((1-sidak)/2,loc = ucombinemean[i][s]*uncmultiplier,scale = ucombinesd[i][s])
            bound2_ = norm.ppf(1-(1-sidak)/2,loc = ucombinemean[i][s]*uncmultiplier,scale = ucombinesd[i][s])
            lowerbound[i][s] = bound1_
            upperbound[i][s] = bound2_

    #######falsification
    #######load predictions and measurements from other functions or excel
    excelsheetprediction = excelfileprediction.to_numpy()
    prediction = numpy.zeros((numberofmodelinstance, numberofsensor))
    for i in range(0,numberofmodelinstance):
        for s in range(0,numberofsensor):
            excelsheetpredictionvalue_= excelsheetprediction[i,s+2]
            prediction[i][s] = excelsheetpredictionvalue_

    excelsheetmeas = excelfilemeas.to_numpy()
    measurement = numpy.zeros((numberofsensor,1))
    for s in range(0,numberofsensor):
        excelsheetmeasvalue_= excelsheetmeas[0,s]
        measurement[s][0] = excelsheetmeasvalue_

    falsification = numpy.zeros((numberofmodelinstance, numberofsensor))

    for i in range(0,numberofmodelinstance):
        for s in range(0,numberofsensor):
            residual = prediction[i][s] - measurement[s][0]
            falsification_ = (residual<upperbound[i][s] and residual>lowerbound[i][s])
            if falsification_==True:
                falsification[i][s] = 1
            else:
                falsification[i][s] = 0

    ##########for validation
    falsificationfinal = numpy.zeros((numberofmodelinstance,len(ID_to_be_included)))
    for i in range(0,len(ID_to_be_included)):
        for s in range(0,numberofmodelinstance):
            falsificationfinal[s][i] = falsification[s,int(ID_to_be_included[i][0])]

    final = sum(numpy.transpose(falsificationfinal))
    is_candidate = numpy.zeros((numberofmodelinstance,1))

    for i in range(0,numberofmodelinstance):
        if final[i] == float(len(ID_to_be_included)):
            is_candidate[i][0] = 1
        else:
            is_candidate[i][0] = 0

    ###write to excel
    finalpathCMS = finalpathprediction
    savefile = load_workbook(finalpathCMS)
    savesheets = savefile.sheetnames
    sheetCMS = savefile[savesheets[1]]

    for s in range(0,numberofmodelinstance):
        sheetCMS.cell(row = s+3, column = 5).value = is_candidate[s][0]

    savefile.save(finalpathCMS)

    ###########load IMS
    excelsheetIMS = excelfileprediction.to_numpy()[0:numberofmodelinstance, 0:2]
    ims = numpy.zeros((numberofmodelinstance, 2))
    for i in range(0,numberofmodelinstance):
        for s in range(0,2):
            excelsheetIMSvalue_= excelsheetIMS[i,s]
            ims[i][s] = excelsheetIMSvalue_

    ########CMS
    cmsID = numpy.where(is_candidate == 1)
    cms = ims[cmsID[0], :]

    ########write to excel
    from f_resultsexcel_Geo import f_resultsexcel_Geo
    f_resultsexcel_Geo(ims, prediction, is_candidate,uncmultiplier,sidakvalue, indextext)

    return prediction,measurement,is_candidate,ims,cms,cmsID,sidakvalue,indextext,uncmultiplier


import scipy.io
import scipy.misc
from glob import glob
import os
import numpy as np
from ops import *
import tensorflow as tf
from tensorflow import contrib
from menpo_functions import *
from logging_functions import *
from data_loading_functions import *

class DeepHeatmapsModel(object):

    """facial landmark localization Network"""

    def __init__(self, mode='TRAIN', train_iter=100000, batch_size=10, learning_rate=1e-3, adam_optimizer=True,
                 momentum=0.95, step=100000, gamma=0.1, reg=0, weight_initializer='xavier', weight_initializer_std=0.01,
                 bias_initializer=0.0, image_size=256, c_dim=3, num_landmarks=68, sigma=1.5, scale=1, margin=0.25,
                 bb_type='gt', approx_maps=True, win_mult=3.33335, augment_basic=True, basic_start=0,
                 augment_texture=False, p_texture=0., augment_geom=False, p_geom=0., artistic_step=-1, artistic_start=0,
                 output_dir='output', save_model_path='model', save_sample_path='sample', save_log_path='logs',
                 test_model_path='model/deep_heatmaps-50000', pre_train_path='model/deep_heatmaps-50000',load_pretrain=False,
                 img_path='data', test_data='full', valid_data='full', valid_size=0, log_valid_every=5,
                 train_crop_dir='crop_gt_margin_0.25', img_dir_ns='crop_gt_margin_0.25_ns',
                 print_every=100, save_every=5000, sample_every=5000, sample_grid=9, sample_to_log=True,
                 debug_data_size=20, debug=False, epoch_data_dir='epoch_data', use_epoch_data=False, menpo_verbose=True):

        # define some extra parameters

        self.log_histograms = False  # save weight + gradient histogram to log
        self.save_valid_images = True  # sample heat maps of validation images
        self.log_artistic_augmentation_probs = False  # save p_texture & p_geom to log
        self.sample_per_channel = False  # sample heatmaps separately for each landmark
        self.approx_maps_gpu = False  # create heat-maps on gpu. NOT RECOMMENDED. TODO: REMOVE

        # for fine-tuning, choose reset_training_op==True. when resuming training, reset_training_op==False
        self.reset_training_op = False

        self.allocate_once = True  # create batch images/landmarks/maps zero arrays only once

        self.fast_img_gen = True

        self.compute_nme = True  # compute normalized mean error

        self.config = tf.ConfigProto()
        self.config.gpu_options.allow_growth = True

        # sampling and logging parameters
        self.print_every = print_every  # print losses to screen + log
        self.save_every = save_every  # save model
        self.sample_every = sample_every  # save images of gen heat maps compared to GT
        self.sample_grid = sample_grid  # number of training images in sample
        self.sample_to_log = sample_to_log  # sample images to log instead of disk
        self.log_valid_every = log_valid_every  # log validation loss (in epochs)

        self.debug = debug
        self.debug_data_size = debug_data_size
        self.use_epoch_data = use_epoch_data
        self.epoch_data_dir = epoch_data_dir

        self.load_pretrain = load_pretrain
        self.pre_train_path = pre_train_path

        self.mode = mode
        self.train_iter = train_iter
        self.learning_rate = learning_rate

        self.image_size = image_size
        self.c_dim = c_dim
        self.batch_size = batch_size

        self.num_landmarks = num_landmarks

        self.save_log_path = save_log_path
        self.save_sample_path = save_sample_path
        self.save_model_path = save_model_path
        self.test_model_path = test_model_path
        self.img_path=img_path

        self.momentum = momentum
        self.step = step  # for lr decay
        self.gamma = gamma  # for lr decay
        self.reg = reg  # weight decay scale

        self.weight_initializer = weight_initializer  # random_normal or xavier
        self.weight_initializer_std = weight_initializer_std
        self.bias_initializer = bias_initializer
        self.adam_optimizer = adam_optimizer

        self.sigma = sigma  # sigma for heatmap generation
        self.scale = scale  # scale for image normalization 255 / 1 / 0
        self.win_mult = win_mult  # gaussian filter size for cpu/gpu approximation: 2 * sigma * win_mult + 1
        self.approx_maps_cpu = approx_maps  # create heat-maps by inserting gaussian filter around landmark locations

        self.test_data = test_data  # if mode is TEST, this choose the set to use full/common/challenging/test/art
        self.train_crop_dir = train_crop_dir
        self.img_dir_ns = os.path.join(img_path, img_dir_ns)
        self.augment_basic = augment_basic  # perform basic augmentation (rotation,flip,crop)
        self.augment_texture = augment_texture  # perform artistic texture augmentation (NS)
        self.p_texture = p_texture  # initial probability of artistic texture augmentation
        self.augment_geom = augment_geom  # perform artistic geometric augmentation
        self.p_geom = p_geom  # initial probability of artistic geometric augmentation
        self.artistic_step = artistic_step  # increase probability of artistic augmentation every X epochs
        self.artistic_start = artistic_start  # min epoch to start artistic augmentation
        self.basic_start = basic_start  # min epoch to start basic augmentation

        self.valid_size = valid_size
        self.valid_data = valid_data

        # load image, bb and landmark data using menpo
        self.bb_dir = os.path.join(img_path, 'Bounding_Boxes')
        self.bb_dictionary = load_bb_dictionary(self.bb_dir, mode, test_data=self.test_data)

        if self.use_epoch_data:
            epoch_0 = os.path.join(self.epoch_data_dir, '0')
            self.img_menpo_list = load_menpo_image_list(
                img_path, train_crop_dir=epoch_0, img_dir_ns=None, mode=mode, bb_dictionary=self.bb_dictionary,
                image_size=self.image_size,test_data=self.test_data, augment_basic=False, augment_texture=False,
                augment_geom=False, verbose=menpo_verbose)
        else:
            self.img_menpo_list = load_menpo_image_list(
                img_path, train_crop_dir, self.img_dir_ns, mode, bb_dictionary=self.bb_dictionary,
                image_size=self.image_size, margin=margin, bb_type=bb_type, test_data=self.test_data,
                augment_basic=(augment_basic and basic_start == 0),
                augment_texture=(augment_texture and artistic_start == 0 and p_texture > 0.), p_texture=p_texture,
                augment_geom=(augment_geom and artistic_start == 0 and p_geom > 0.), p_geom=p_geom,
                verbose=menpo_verbose)

        if mode == 'TRAIN':

            train_params = locals()
            print_training_params_to_file(train_params)  # save init parameters

            self.train_inds = np.arange(len(self.img_menpo_list))

            if self.debug:
                self.train_inds = self.train_inds[:self.debug_data_size]
                self.img_menpo_list = self.img_menpo_list[self.train_inds]

            if valid_size > 0:

                self.valid_bb_dictionary = load_bb_dictionary(self.bb_dir, 'TEST', test_data=self.valid_data)
                self.valid_img_menpo_list = load_menpo_image_list(
                    img_path, train_crop_dir, self.img_dir_ns, 'TEST', bb_dictionary=self.valid_bb_dictionary,
                    image_size=self.image_size, margin=margin, bb_type=bb_type, test_data=self.valid_data,
                    verbose=menpo_verbose)

                np.random.seed(0)
                self.val_inds = np.arange(len(self.valid_img_menpo_list))
                np.random.shuffle(self.val_inds)
                self.val_inds = self.val_inds[:self.valid_size]

                self.valid_img_menpo_list = self.valid_img_menpo_list[self.val_inds]

                if self.approx_maps_cpu:
                    self.valid_images_loaded, self.valid_gt_maps_loaded, self.valid_landmarks_loaded =\
                        load_images_landmarks_approx_maps(
                            self.valid_img_menpo_list, np.arange(self.valid_size), primary=True, image_size=self.image_size,
                            num_landmarks=self.num_landmarks, c_dim=self.c_dim, scale=self.scale, win_mult=self.win_mult,
                            sigma=self.sigma, save_landmarks=True)
                else:
                    self.valid_images_loaded, self.valid_gt_maps_loaded, self.valid_landmarks_loaded =\
                        load_images_landmarks_maps(
                            self.valid_img_menpo_list, np.arange(self.valid_size), primary=True, image_size=self.image_size,
                            c_dim=self.c_dim, num_landmarks=self.num_landmarks, scale=self.scale, sigma=self.sigma,
                            save_landmarks=True)

                if self.allocate_once:
                    self.valid_landmarks_pred = np.zeros([self.valid_size, self.num_landmarks, 2]).astype('float32')

                if self.valid_size > self.sample_grid:
                    self.valid_gt_maps_loaded = self.valid_gt_maps_loaded[:self.sample_grid]
            else:
                self.val_inds = None

            self.epoch_inds_shuffle = train_val_shuffle_inds_per_epoch(
                self.val_inds, self.train_inds, train_iter, batch_size, save_log_path)

    def add_placeholders(self):

        if self.mode == 'TEST':
            self.images = tf.placeholder(
                tf.float32, [None, self.image_size, self.image_size, self.c_dim], 'images')

            self.heatmaps_small = tf.placeholder(
                tf.float32, [None, int(self.image_size/4), int(self.image_size/4), self.num_landmarks], 'heatmaps_small')
            self.lms_small = tf.placeholder(tf.float32, [None, self.num_landmarks, 2], 'lms_small')
            self.pred_lms_small = tf.placeholder(tf.float32, [None, self.num_landmarks, 2], 'pred_lms_small')

        elif self.mode == 'TRAIN':
            self.images = tf.placeholder(
                tf.float32, [None, self.image_size, self.image_size, self.c_dim], 'train_images')

            self.heatmaps_small = tf.placeholder(
                tf.float32, [None, int(self.image_size/4), int(self.image_size/4), self.num_landmarks], 'train_heatmaps_small')

            self.train_lms_small = tf.placeholder(tf.float32, [None, self.num_landmarks, 2], 'train_lms_small')
            self.train_pred_lms_small = tf.placeholder(tf.float32, [None, self.num_landmarks, 2], 'train_pred_lms_small')

            self.valid_lms_small = tf.placeholder(tf.float32, [None, self.num_landmarks, 2], 'valid_lms_small')
            self.valid_pred_lms_small = tf.placeholder(tf.float32, [None, self.num_landmarks, 2], 'valid_pred_lms_small')

            self.p_texture_log = tf.placeholder(tf.float32, [])
            self.p_geom_log = tf.placeholder(tf.float32, [])

            self.sparse_hm_small = tf.placeholder(tf.float32, [None, int(self.image_size/4), int(self.image_size/4), 1])

            if self.sample_to_log:
                row = int(np.sqrt(self.sample_grid))
                self.log_image_map = tf.placeholder(
                    tf.uint8, [None,row * int(self.image_size/4), 3 * row *int(self.image_size/4), self.c_dim], 'sample_img_map')
                if self.sample_per_channel:
                    row = np.ceil(np.sqrt(self.num_landmarks)).astype(np.int64)
                    self.log_map_channels = tf.placeholder(
                        tf.uint8, [None, row * int(self.image_size/4), 2 * row * int(self.image_size/4), self.c_dim],
                        'sample_map_channels')

    def heatmaps_network(self, input_images, reuse=None, name='pred_heatmaps'):

        with tf.name_scope(name):

            if self.weight_initializer == 'xavier':
                weight_initializer = contrib.layers.xavier_initializer()
            else:
                weight_initializer = tf.random_normal_initializer(stddev=self.weight_initializer_std)

            bias_init = tf.constant_initializer(self.bias_initializer)

            with tf.variable_scope('heatmaps_network'):
                with tf.name_scope('primary_net'):

                    l1 = conv_relu_pool(input_images, 5, 128, conv_ker_init=weight_initializer, conv_bias_init=bias_init,
                                        reuse=reuse, var_scope='conv_1')
                    l2 = conv_relu_pool(l1, 5, 128, conv_ker_init=weight_initializer, conv_bias_init=bias_init,
                                        reuse=reuse, var_scope='conv_2')
                    l3 = conv_relu(l2, 5, 128, conv_ker_init=weight_initializer, conv_bias_init=bias_init,
                                   reuse=reuse, var_scope='conv_3')

                    l4_1 = conv_relu(l3, 3, 128, conv_dilation=1, conv_ker_init=weight_initializer,
                                     conv_bias_init=bias_init, reuse=reuse, var_scope='conv_4_1')
                    l4_2 = conv_relu(l3, 3, 128, conv_dilation=2, conv_ker_init=weight_initializer,
                                     conv_bias_init=bias_init, reuse=reuse, var_scope='conv_4_2')
                    l4_3 = conv_relu(l3, 3, 128, conv_dilation=3, conv_ker_init=weight_initializer,
                                     conv_bias_init=bias_init, reuse=reuse, var_scope='conv_4_3')
                    l4_4 = conv_relu(l3, 3, 128, conv_dilation=4, conv_ker_init=weight_initializer,
                                     conv_bias_init=bias_init, reuse=reuse, var_scope='conv_4_4')

                    l4 = tf.concat([l4_1, l4_2, l4_3, l4_4], 3, name='conv_4')

                    l5_1 = conv_relu(l4, 3, 256, conv_dilation=1, conv_ker_init=weight_initializer,
                                     conv_bias_init=bias_init, reuse=reuse, var_scope='conv_5_1')
                    l5_2 = conv_relu(l4, 3, 256, conv_dilation=2, conv_ker_init=weight_initializer,
                                     conv_bias_init=bias_init, reuse=reuse, var_scope='conv_5_2')
                    l5_3 = conv_relu(l4, 3, 256, conv_dilation=3, conv_ker_init=weight_initializer,
                                     conv_bias_init=bias_init, reuse=reuse, var_scope='conv_5_3')
                    l5_4 = conv_relu(l4, 3, 256, conv_dilation=4, conv_ker_init=weight_initializer,
                                     conv_bias_init=bias_init, reuse=reuse, var_scope='conv_5_4')

                    l5 = tf.concat([l5_1, l5_2, l5_3, l5_4], 3, name='conv_5')

                    l6 = conv_relu(l5, 1, 512, conv_ker_init=weight_initializer,
                                   conv_bias_init=bias_init, reuse=reuse, var_scope='conv_6')
                    l7 = conv_relu(l6, 1, 256, conv_ker_init=weight_initializer,
                                   conv_bias_init=bias_init, reuse=reuse, var_scope='conv_7')
                    primary_out = conv(l7, 1, self.num_landmarks, conv_ker_init=weight_initializer,
                                            conv_bias_init=bias_init, reuse=reuse, var_scope='conv_8')

                self.all_layers = [l1, l2, l3, l4, l5, l6, l7, primary_out]

                return primary_out

    def build_model(self):
        self.pred_hm_p = self.heatmaps_network(self.images,name='heatmaps_prediction')

    def build_hm_generator(self):  # TODO: remove
        # generate heat-maps using:
        # a sparse base (matrix of zeros with 1's in landmark locations) and convolving with a gaussian filter
        print ("*** using convolution to create heat-maps. use this option only with GPU support ***")

        # create gaussian filter
        win_small = int(self.win_mult * self.sigma)
        x_small, y_small = np.mgrid[0:2*win_small+1, 0:2*win_small+1]

        gauss_small = (8. / 3) * self.sigma * gaussian(x_small, y_small, win_small, win_small, sigma=self.sigma)
        gauss_small = tf.constant(gauss_small, tf.float32)
        gauss_small = tf.reshape(gauss_small, [2 * win_small + 1, 2 * win_small + 1, 1, 1])

        # convolve sparse map with gaussian
        self.filt_hm_small = tf.nn.conv2d(self.sparse_hm_small, gauss_small, strides=[1, 1, 1, 1], padding='SAME')
        self.filt_hm_small = tf.transpose(
            tf.concat(tf.split(self.filt_hm_small, self.batch_size, axis=0), 3), [3, 1, 2, 0])

    def create_loss_ops(self):  # TODO: calculate NME on resized maps to 256

        def l2_loss_norm_eyes(pred_landmarks, real_landmarks, normalize=True, name='NME'):

            with tf.name_scope(name):
                with tf.name_scope('real_pred_landmarks_rmse'):
                    landmarks_rms_err = tf.reduce_mean(
                        tf.sqrt(tf.reduce_sum(tf.square(pred_landmarks - real_landmarks), axis=2)), axis=1)
                if normalize:
                    with tf.name_scope('inter_pupil_dist'):
                        with tf.name_scope('left_eye_center'):
                            p1 = tf.reduce_mean(tf.slice(real_landmarks, [0, 42, 0], [-1, 6, 2]), axis=1)
                        with tf.name_scope('right_eye_center'):
                            p2 = tf.reduce_mean(tf.slice(real_landmarks, [0, 36, 0], [-1, 6, 2]), axis=1)

                        eye_dist = tf.sqrt(tf.reduce_sum(tf.square(p1 - p2), axis=1))

                    return landmarks_rms_err / eye_dist
                else:
                    return landmarks_rms_err

        if self.mode is 'TRAIN':
            primary_maps_diff = self.pred_hm_p-self.heatmaps_small
            self.total_loss = 1000.*tf.reduce_mean(tf.square(primary_maps_diff))

            # add weight decay
            self.total_loss += self.reg * tf.add_n(
                [tf.nn.l2_loss(v) for v in tf.trainable_variables() if 'bias' not in v.name])

            if self.compute_nme:
                self.nme_loss = tf.reduce_mean(l2_loss_norm_eyes(self.train_pred_lms_small,self.train_lms_small))

            if self.valid_size > 0 and self.compute_nme:
                self.valid_nme_loss = tf.reduce_mean(l2_loss_norm_eyes(self.valid_pred_lms_small,self.valid_lms_small))

        elif self.mode == 'TEST' and self.compute_nme:
            self.nme_per_image = l2_loss_norm_eyes(self.pred_lms_small, self.lms_small)
            self.nme_loss = tf.reduce_mean(self.nme_per_image)

    def predict_landmarks_in_batches(self, image_paths, session):

        num_batches = int(1.*len(image_paths)/self.batch_size)
        if num_batches == 0:
            batch_size = len(image_paths)
            num_batches = 1
        else:
            batch_size = self.batch_size

        img_inds = np.arange(len(image_paths))
        for j in range(num_batches):
            batch_inds = img_inds[j * batch_size:(j + 1) * batch_size]

            batch_images, _, batch_lms_small = \
                load_images_landmarks_maps(
                    self.img_menpo_list, batch_inds, primary=True, image_size=self.image_size,
                    c_dim=self.c_dim, num_landmarks=self.num_landmarks, scale=self.scale, sigma=self.sigma,
                    save_landmarks=self.compute_nme)

            batch_maps_small_pred = session.run(self.pred_hm_p, {self.images: batch_images})
            batch_pred_landmarks = batch_heat_maps_to_landmarks(
                batch_maps_small_pred, batch_size=batch_size, image_size=int(self.image_size/4),
                num_landmarks=self.num_landmarks)

            if j == 0:
                all_pred_landmarks = batch_pred_landmarks.copy()
                all_gt_landmarks = batch_lms_small.copy()
            else:
                all_pred_landmarks = np.concatenate((all_pred_landmarks,batch_pred_landmarks),0)
                all_gt_landmarks = np.concatenate((all_gt_landmarks, batch_lms_small), 0)

        reminder = len(image_paths)-num_batches*batch_size

        if reminder > 0:
            reminder_inds = img_inds[-reminder:]

            batch_images, _, batch_lms_small = \
                load_images_landmarks_maps(
                    self.img_menpo_list, reminder_inds, primary=True, image_size=self.image_size,
                    c_dim=self.c_dim, num_landmarks=self.num_landmarks, scale=self.scale, sigma=self.sigma,
                    save_landmarks=self.compute_nme)

            batch_maps_small_pred = session.run(self.pred_hm_p, {self.images: batch_images})
            batch_pred_landmarks = batch_heat_maps_to_landmarks(
                batch_maps_small_pred, batch_size=reminder, image_size=int(self.image_size/4),
                num_landmarks=self.num_landmarks)

            all_pred_landmarks = np.concatenate((all_pred_landmarks, batch_pred_landmarks), 0)
            all_gt_landmarks = np.concatenate((all_gt_landmarks, batch_lms_small), 0)

        return all_pred_landmarks, all_gt_landmarks

    def predict_landmarks_in_batches_loaded(self, images, session):

        num_images = int(images.shape[0])
        num_batches = int(1.*num_images/self.batch_size)
        if num_batches == 0:
            batch_size = num_images
            num_batches = 1
        else:
            batch_size = self.batch_size

        for j in range(num_batches):

            batch_images = images[j * batch_size:(j + 1) * batch_size,:,:,:]
            batch_maps_small_pred = session.run(self.pred_hm_p, {self.images: batch_images})
            if self.allocate_once:
                batch_heat_maps_to_landmarks_alloc_once(
                    batch_maps=batch_maps_small_pred,
                    batch_landmarks=self.valid_landmarks_pred[j * batch_size:(j + 1) * batch_size, :, :],
                    batch_size=batch_size, image_size=int(self.image_size/4), num_landmarks=self.num_landmarks)
            else:
                batch_pred_landmarks = batch_heat_maps_to_landmarks(
                    batch_maps_small_pred, batch_size=batch_size, image_size=int(self.image_size/4),
                    num_landmarks=self.num_landmarks)

                if j == 0:
                    all_pred_landmarks = batch_pred_landmarks.copy()
                else:
                    all_pred_landmarks = np.concatenate((all_pred_landmarks, batch_pred_landmarks), 0)

        reminder = num_images-num_batches*batch_size
        if reminder > 0:

            batch_images = images[-reminder:, :, :, :]
            batch_maps_small_pred = session.run(self.pred_hm_p, {self.images: batch_images})
            if self.allocate_once:
                batch_heat_maps_to_landmarks_alloc_once(
                    batch_maps=batch_maps_small_pred,
                    batch_landmarks=self.valid_landmarks_pred[-reminder:, :, :],
                    batch_size=reminder, image_size=int(self.image_size/4), num_landmarks=self.num_landmarks)
            else:
                batch_pred_landmarks = batch_heat_maps_to_landmarks(
                    batch_maps_small_pred, batch_size=reminder, image_size=int(self.image_size/4),
                    num_landmarks=self.num_landmarks)

                all_pred_landmarks = np.concatenate((all_pred_landmarks, batch_pred_landmarks), 0)

        if not self.allocate_once:
            return all_pred_landmarks

    def create_summary_ops(self):

        self.batch_summary_op = tf.summary.scalar('l_total', self.total_loss)

        if self.compute_nme:
            l_nme = tf.summary.scalar('l_nme', self.nme_loss)
            self.batch_summary_op = tf.summary.merge([self.batch_summary_op, l_nme])

        if self.log_histograms:
            var_summary = [tf.summary.histogram(var.name, var) for var in tf.trainable_variables()]
            grads = tf.gradients(self.total_loss, tf.trainable_variables())
            grads = list(zip(grads, tf.trainable_variables()))
            grad_summary = [tf.summary.histogram(var.name + '/grads', grad) for grad, var in grads]
            activ_summary = [tf.summary.histogram(layer.name, layer) for layer in self.all_layers]
            self.batch_summary_op = tf.summary.merge([self.batch_summary_op, var_summary, grad_summary, activ_summary])

        if self.augment_texture and self.log_artistic_augmentation_probs:
            p_texture_summary = tf.summary.scalar('p_texture', self.p_texture_log)
            self.batch_summary_op = tf.summary.merge([self.batch_summary_op, p_texture_summary])

        if self.augment_geom and self.log_artistic_augmentation_probs:
            p_geom_summary = tf.summary.scalar('p_geom', self.p_geom_log)
            self.batch_summary_op = tf.summary.merge([self.batch_summary_op, p_geom_summary])

        if self.valid_size > 0 and self.compute_nme:
            self.valid_summary = tf.summary.scalar('valid_l_nme', self.valid_nme_loss)

        if self.sample_to_log:
            img_map_summary =tf.summary.image('compare_map_to_gt',self.log_image_map)
            if self.sample_per_channel:
                map_channels_summary = tf.summary.image('compare_map_channels_to_gt', self.log_map_channels)
                self.img_summary = tf.summary.merge([img_map_summary, map_channels_summary])
            else:
                self.img_summary = img_map_summary
            if self.valid_size >= self.sample_grid:
                img_map_summary_valid = tf.summary.image('compare_map_to_gt_valid', self.log_image_map)
                if self.sample_per_channel:
                    map_channels_summary_valid = tf.summary.image('compare_map_channels_to_gt_valid', self.log_map_channels)
                    self.img_summary_valid = tf.summary.merge([img_map_summary_valid, map_channels_summary_valid])
                else:
                    self.img_summary_valid = img_map_summary_valid

    def eval(self):

        self.add_placeholders()
        # build model
        self.build_model()
        self.create_loss_ops()

        if self.debug:
            self.img_menpo_list = self.img_menpo_list[:np.min([self.debug_data_size, len(self.img_menpo_list)])]

        num_images = len(self.img_menpo_list)
        img_inds = np.arange(num_images)

        sample_iter = np.ceil(1. * num_images / self.sample_grid).astype('int')

        with tf.Session(config=self.config) as sess:

            # load trained parameters
            print ('loading test model...')
            saver = tf.train.Saver()
            saver.restore(sess, self.test_model_path)

            _, model_name = os.path.split(self.test_model_path)

            gt_provided = self.img_menpo_list[0].has_landmarks  # check if GT landmarks provided

            for i in range(sample_iter):

                batch_inds = img_inds[i * self.sample_grid:(i + 1) * self.sample_grid]

                if not gt_provided:
                    batch_images = load_images(self.img_menpo_list, batch_inds, image_size=self.image_size,
                                               c_dim=self.c_dim, scale=self.scale)

                    batch_maps_small_pred = sess.run(self.pred_hm_p, {self.images: batch_images})

                    batch_maps_gt = None
                else:
                    # TODO: add option for approx maps + allocate once
                    batch_images, batch_maps_gt, _ = \
                        load_images_landmarks_maps(
                            self.img_menpo_list, batch_inds, primary=True, image_size=self.image_size,
                            c_dim=self.c_dim, num_landmarks=self.num_landmarks, scale=self.scale, sigma=self.sigma,
                            save_landmarks=False)

                    batch_maps_small_pred = sess.run(self.pred_hm_p, {self.images: batch_images})

                sample_path_imgs = os.path.join(
                    self.save_sample_path, model_name +'-'+ self.test_data+'-sample-%d-to-%d-1.png' % (
                        i * self.sample_grid, (i + 1) * self.sample_grid))

                merged_img = merge_images_landmarks_maps_gt(
                    batch_images.copy(), batch_maps_small_pred, batch_maps_gt, image_size=self.image_size,
                    num_landmarks=self.num_landmarks, num_samples=self.sample_grid, scale=self.scale, circle_size=0,
                    fast=self.fast_img_gen)

                scipy.misc.imsave(sample_path_imgs, merged_img)

                if self.sample_per_channel:
                    map_per_channel = map_comapre_channels(
                        batch_images.copy(), batch_maps_small_pred,batch_maps_gt, image_size=int(self.image_size/4),
                        num_landmarks=self.num_landmarks, scale=self.scale)

                    sample_path_channels = os.path.join(
                        self.save_sample_path, model_name + '-' + self.test_data + '-sample-%d-to-%d-3.png' % (
                            i * self.sample_grid, (i + 1) * self.sample_grid))

                    scipy.misc.imsave(sample_path_channels, map_per_channel)

                print ('saved %s' % sample_path_imgs)

            if self.compute_nme and self.test_data in ['full', 'challenging', 'common', 'training', 'test']:
                print ('\n Calculating NME on: ' + self.test_data + '...')
                pred_lms, lms_gt = self.predict_landmarks_in_batches(self.img_menpo_list, sess)
                nme = sess.run(self.nme_loss, {self.pred_lms_small: pred_lms, self.lms_small: lms_gt})
                print ('NME on ' + self.test_data + ': ' + str(nme))

    def train(self):
        # set random seed
        tf.set_random_seed(1234)
        np.random.seed(1234)
        # build a graph
        # add placeholders
        self.add_placeholders()
        # build model
        self.build_model()
        # create loss ops
        self.create_loss_ops()
        # create summary ops
        self.create_summary_ops()

        # create optimizer and training op
        global_step = tf.Variable(0, trainable=False)
        lr = tf.train.exponential_decay(self.learning_rate,global_step, self.step, self.gamma, staircase=True)
        if self.adam_optimizer:
            optimizer = tf.train.AdamOptimizer(lr)
        else:
            optimizer = tf.train.MomentumOptimizer(lr, self.momentum)

        train_op = optimizer.minimize(self.total_loss,global_step=global_step)

        # TODO: remove
        if self.approx_maps_gpu:  # create heat-maps using tf convolution. use only with GPU support!
            self.build_hm_generator()

        with tf.Session(config=self.config) as sess:

            tf.global_variables_initializer().run()

            # load pre trained weights if load_pretrain==True
            if self.load_pretrain:
                print
                print('*** loading pre-trained weights from: '+self.pre_train_path+' ***')
                loader = tf.train.Saver()
                loader.restore(sess, self.pre_train_path)
                print("*** Model restore finished, current global step: %d" % global_step.eval())

            # for fine-tuning, choose reset_training_op==True. when resuming training, reset_training_op==False
            if self.reset_training_op:
                print ("resetting optimizer and global step")
                opt_var_list = [optimizer.get_slot(var, name) for name in optimizer.get_slot_names()
                                 for var in tf.global_variables() if optimizer.get_slot(var, name) is not None]
                opt_var_list_init = tf.variables_initializer(opt_var_list)
                opt_var_list_init.run()
                sess.run(global_step.initializer)

            # create model saver and file writer
            summary_writer = tf.summary.FileWriter(logdir=self.save_log_path, graph=tf.get_default_graph())
            saver = tf.train.Saver()

            print
            print('*** Start Training ***')

            # initialize some variables before training loop
            resume_step = global_step.eval()
            num_train_images = len(self.img_menpo_list)
            batches_in_epoch = int(float(num_train_images) / float(self.batch_size))
            epoch = int(resume_step / batches_in_epoch)
            img_inds = self.epoch_inds_shuffle[epoch, :]
            p_texture = self.p_texture
            p_geom = self.p_geom
            artistic_reload = False
            basic_reload = True
            log_valid = True
            log_valid_images = True

            if self.allocate_once:
                batch_images = np.zeros([self.batch_size, self.image_size, self.image_size, self.c_dim]).astype('float32')
                batch_lms_small = np.zeros([self.batch_size, self.num_landmarks, 2]).astype('float32')
                batch_lms_small_pred = np.zeros([self.batch_size, self.num_landmarks, 2]).astype('float32')
                if self.approx_maps_gpu:
                    batch_hm_base_small = np.zeros((self.batch_size * self.num_landmarks,
                                                    int(self.image_size/4), int(self.image_size/4), 1)).astype('float32')
                else:
                    batch_maps_small = np.zeros((self.batch_size, int(self.image_size/4),
                                                 int(self.image_size/4), self.num_landmarks)).astype('float32')

            if self.approx_maps_cpu:
                gaussian_filt = create_gaussian_filter(sigma=self.sigma, win_mult=self.win_mult)

            for step in range(resume_step, self.train_iter):

                j = step % batches_in_epoch  # j==0 if we finished an epoch

                if step > resume_step and j == 0:  # if we finished an epoch and this isn't the first step
                    epoch += 1
                    img_inds = self.epoch_inds_shuffle[epoch, :]  # get next shuffled image inds
                    artistic_reload = True
                    log_valid = True
                    log_valid_images = True
                    if self.use_epoch_data:
                        epoch_dir = os.path.join(self.epoch_data_dir, str(epoch))
                        self.img_menpo_list = load_menpo_image_list(
                            self.img_path, train_crop_dir=epoch_dir, img_dir_ns=None, mode=self.mode,
                            bb_dictionary=self.bb_dictionary, image_size=self.image_size, test_data=self.test_data,
                            augment_basic=False, augment_texture=False, augment_geom=False)

                # add basic augmentation (if basic_start > 0 and augment_basic is True)
                if basic_reload and (epoch >= self.basic_start) and self.basic_start > 0 and self.augment_basic:
                    basic_reload = False
                    self.img_menpo_list = reload_menpo_image_list(
                        self.img_path, self.train_crop_dir, self.img_dir_ns, self.mode, self.train_inds,
                        image_size=self.image_size, augment_basic=self.augment_basic,
                        augment_texture=(self.augment_texture and epoch >= self.artistic_start), p_texture=p_texture,
                        augment_geom=(self.augment_geom and epoch >= self.artistic_start), p_geom=p_geom)
                    print ("****** adding basic augmentation ******")

                # increase artistic augmentation probability
                if ((epoch % self.artistic_step == 0 and epoch >= self.artistic_start and self.artistic_step != -1)
                    or (epoch == self.artistic_start)) and (self.augment_geom or self.augment_texture)\
                        and artistic_reload:

                    artistic_reload = False

                    if epoch == self.artistic_start:
                        print ("****** adding artistic augmentation ******")
                        print ("****** augment_geom: " + str(self.augment_geom) + ", p_geom: " + str(p_geom) + " ******")
                        print ("****** augment_texture: " + str(self.augment_texture) + ", p_texture: " +
                               str(p_texture) + " ******")

                    if epoch % self.artistic_step == 0 and self.artistic_step != -1:
                        print ("****** increasing artistic augmentation probability ******")

                        p_geom = 1.- 0.95 ** (epoch/self.artistic_step)
                        p_texture = 1. - 0.95 ** (epoch/self.artistic_step)

                        print ("****** augment_geom: " + str(self.augment_geom) + ", p_geom: " + str(p_geom) + " ******")
                        print ("****** augment_texture: " + str(self.augment_texture) + ", p_texture: " +
                               str(p_texture) + " ******")

                    self.img_menpo_list = reload_menpo_image_list(
                        self.img_path, self.train_crop_dir, self.img_dir_ns, self.mode, self.train_inds,
                        image_size=self.image_size, augment_basic=(self.augment_basic and epoch >= self.basic_start),
                        augment_texture=self.augment_texture, p_texture=p_texture,
                        augment_geom=self.augment_geom, p_geom=p_geom)

                # get batch images
                batch_inds = img_inds[j * self.batch_size:(j + 1) * self.batch_size]

                if self.approx_maps_gpu:  # TODO: remove
                    if self.allocate_once:
                        load_images_landmarks_alloc_once(
                            self.img_menpo_list, batch_inds, images=batch_images, landmarks_small=batch_lms_small,
                            landmarks=None, primary=True, image_size=self.image_size, scale=self.scale)

                        create_heat_maps_base_alloc_once(
                            landmarks_small=batch_lms_small.astype(int), landmarks=None,
                            hm_small=batch_hm_base_small, hm_large=None, primary=True, num_images=self.batch_size,
                            num_landmarks=self.num_landmarks)
                    else:
                        batch_images, batch_lms_small = load_images_landmarks(
                            self.img_menpo_list, batch_inds, primary=True, image_size=self.image_size, c_dim=self.c_dim,
                            num_landmarks=self.num_landmarks, scale=self.scale)

                        batch_hm_base_small = create_heat_maps_base(
                            landmarks_small=batch_lms_small.astype(int), landmarks=None, primary=True,
                            num_images=self.batch_size, image_size=self.image_size, num_landmarks=self.num_landmarks)

                    batch_maps_small = sess.run(self.filt_hm_small, {self.sparse_hm_small: batch_hm_base_small})
                elif self.approx_maps_cpu:
                    if self.allocate_once:
                        load_images_landmarks_approx_maps_alloc_once(
                            self.img_menpo_list, batch_inds, images=batch_images, maps_small=batch_maps_small,
                            maps=None, landmarks=batch_lms_small, primary=True, image_size=self.image_size,
                            num_landmarks=self.num_landmarks, scale=self.scale, gauss_filt_small=gaussian_filt,
                            win_mult=self.win_mult, sigma=self.sigma, save_landmarks=self.compute_nme)
                    else:
                        batch_images, batch_maps_small, batch_lms_small = load_images_landmarks_approx_maps(
                            self.img_menpo_list, batch_inds, primary=True, image_size=self.image_size,
                            num_landmarks=self.num_landmarks, c_dim=self.c_dim, scale=self.scale,
                            gauss_filt_small=gaussian_filt, win_mult=self.win_mult, sigma=self.sigma,
                            save_landmarks=self.compute_nme)
                else:
                    if self.allocate_once:
                        load_images_landmarks_maps_alloc_once(
                            self.img_menpo_list, batch_inds, images=batch_images, maps_small=batch_maps_small,
                            landmarks=batch_lms_small, maps=None, primary=True, image_size=self.image_size,
                            num_landmarks=self.num_landmarks, scale=self.scale, sigma=self.sigma,
                            save_landmarks=self.compute_nme)
                    else:
                        batch_images, batch_maps_small, batch_lms_small = load_images_landmarks_maps(
                            self.img_menpo_list, batch_inds, primary=True, image_size=self.image_size, c_dim=self.c_dim,
                            num_landmarks=self.num_landmarks, scale=self.scale, sigma=self.sigma,
                            save_landmarks=self.compute_nme)

                feed_dict_train = {self.images: batch_images, self.heatmaps_small: batch_maps_small}

                sess.run(train_op, feed_dict_train)

                # save to log and print status
                if step == resume_step or (step + 1) % self.print_every == 0:

                    # log probability of artistic augmentation
                    if self.log_artistic_augmentation_probs and (self.augment_geom or self.augment_texture):
                        if self.augment_geom and not self.augment_texture:
                            art_augment_prob_dict = {self.p_geom_log: p_geom}
                        elif self.augment_texture and not self.augment_geom:
                            art_augment_prob_dict = {self.p_texture_log: p_texture}
                        else:
                            art_augment_prob_dict = {self.p_texture_log: p_texture, self.p_geom_log: p_geom}

                    # train data log
                    if self.compute_nme:
                        batch_maps_small_pred = sess.run(self.pred_hm_p, {self.images: batch_images})
                        if self.allocate_once:
                            batch_heat_maps_to_landmarks_alloc_once(
                                batch_maps=batch_maps_small_pred, batch_landmarks=batch_lms_small_pred,
                                batch_size=self.batch_size, image_size=int(self.image_size/4),
                                num_landmarks=self.num_landmarks)
                        else:
                            batch_lms_small_pred = batch_heat_maps_to_landmarks(
                                batch_maps_small_pred, self.batch_size, image_size=int(self.image_size/4),
                                num_landmarks=self.num_landmarks)

                        train_feed_dict_log = {
                            self.images: batch_images, self.heatmaps_small: batch_maps_small,
                            self.train_lms_small: batch_lms_small, self.train_pred_lms_small: batch_lms_small_pred}
                        if self.log_artistic_augmentation_probs and (self.augment_geom or self.augment_texture):
                            train_feed_dict_log.update(art_augment_prob_dict)

                        summary, l_t, l_nme = sess.run(
                            [self.batch_summary_op, self.total_loss, self.nme_loss], train_feed_dict_log)

                        print (
                            'epoch: [%d] step: [%d/%d] primary loss: [%.6f] NME: [%.6f]' % (
                                epoch, step + 1, self.train_iter, l_t, l_nme))
                    else:
                        train_feed_dict_log = {self.images: batch_images, self.heatmaps_small: batch_maps_small}
                        if self.log_artistic_augmentation_probs and (self.augment_geom or self.augment_texture):
                            train_feed_dict_log.update(art_augment_prob_dict)

                        summary, l_t = sess.run(
                            [self.batch_summary_op, self.total_loss], train_feed_dict_log)

                        print (
                            'epoch: [%d] step: [%d/%d] primary loss: [%.6f]' % (
                                epoch, step + 1, self.train_iter, l_t))

                    summary_writer.add_summary(summary, step)

                    # valid data log
                    if self.valid_size > 0 and (log_valid and epoch % self.log_valid_every == 0)\
                            and self.compute_nme:
                        log_valid = False

                        if self.allocate_once:
                            self.predict_landmarks_in_batches_loaded(self.valid_images_loaded, sess)
                            valid_feed_dict_log = {
                                self.valid_lms_small: self.valid_landmarks_loaded,
                                self.valid_pred_lms_small: self.valid_landmarks_pred}
                        else:
                            valid_pred_lms = self.predict_landmarks_in_batches_loaded(self.valid_images_loaded, sess)
                            valid_feed_dict_log = {
                                self.valid_lms_small: self.valid_landmarks_loaded,
                                self.valid_pred_lms_small: valid_pred_lms}

                        v_summary,l_v_nme = sess.run([self.valid_summary, self.valid_nme_loss], valid_feed_dict_log)
                        summary_writer.add_summary(v_summary, step)

                        print (
                            'epoch: [%d] step: [%d/%d] valid NME: [%.6f]' % (
                                epoch, step + 1, self.train_iter, l_v_nme))

                # save model
                if (step + 1) % self.save_every == 0:
                    saver.save(sess, os.path.join(self.save_model_path, 'deep_heatmaps'), global_step=step + 1)
                    print ('model/deep-heatmaps-%d saved' % (step + 1))

                # save images. TODO: add option to allocate once
                if step == resume_step or (step + 1) % self.sample_every == 0:

                    if not self.compute_nme:
                        batch_maps_small_pred = sess.run(self.pred_hm_p,  {self.images: batch_images})
                        batch_lms_small_pred=None

                    merged_img = merge_images_landmarks_maps_gt(
                        batch_images.copy(), batch_maps_small_pred, batch_maps_small,
                        landmarks=batch_lms_small_pred, image_size=self.image_size,
                        num_landmarks=self.num_landmarks, num_samples=self.sample_grid, scale=self.scale,
                        circle_size=0, fast=self.fast_img_gen)

                    if self.sample_per_channel:
                        map_per_channel = map_comapre_channels(
                            batch_images.copy(), batch_maps_small_pred,batch_maps_small,
                            image_size=int(self.image_size/4), num_landmarks=self.num_landmarks, scale=self.scale)

                    if self.sample_to_log:
                        if self.sample_per_channel:
                            summary_img = sess.run(
                                self.img_summary, {self.log_image_map: np.expand_dims(merged_img, 0),
                                                   self.log_map_channels: np.expand_dims(map_per_channel, 0)})
                        else:
                            summary_img = sess.run(
                                self.img_summary, {self.log_image_map: np.expand_dims(merged_img, 0)})

                        summary_writer.add_summary(summary_img, step)

                        if (self.valid_size >= self.sample_grid) and self.save_valid_images and\
                                (log_valid_images and epoch % self.log_valid_every == 0):
                            log_valid_images=False

                            batch_maps_small_pred_val = sess.run(
                                self.pred_hm_p, {self.images: self.valid_images_loaded[:self.sample_grid]})

                            merged_img = merge_images_landmarks_maps_gt(
                                self.valid_images_loaded[:self.sample_grid].copy(), batch_maps_small_pred_val,
                                self.valid_gt_maps_loaded, image_size=self.image_size,
                                num_landmarks=self.num_landmarks, num_samples=self.sample_grid,
                                scale=self.scale, circle_size=0, fast=self.fast_img_gen)

                            if self.sample_per_channel:
                                map_per_channel = map_comapre_channels(
                                    self.valid_images_loaded[:self.sample_grid].copy(), batch_maps_small_pred_val,
                                    self.valid_gt_maps_loaded, image_size=int(self.image_size/4),
                                    num_landmarks=self.num_landmarks, scale=self.scale)

                                summary_img = sess.run(
                                    self.img_summary_valid, {self.log_image_map: np.expand_dims(merged_img, 0),
                                                       self.log_map_channels: np.expand_dims(map_per_channel, 0)})
                            else:
                                summary_img = sess.run(
                                    self.img_summary_valid, {self.log_image_map: np.expand_dims(merged_img, 0)})
                            summary_writer.add_summary(summary_img, step)

                    else:
                        sample_path_imgs = os.path.join(self.save_sample_path,'epoch-%d-train-iter-%d-1.png'
                                                        % (epoch, step + 1))
                        scipy.misc.imsave(sample_path_imgs, merged_img)
                        if self.sample_per_channel:
                            sample_path_ch_maps = os.path.join(self.save_sample_path, 'epoch-%d-train-iter-%d-3.png'
                                                               % (epoch, step + 1))
                            scipy.misc.imsave(sample_path_ch_maps, map_per_channel)

            print('*** Finished Training ***')

    def get_maps_image(self, test_image, reuse=None):
        self.add_placeholders()
        # build model
        pred_hm_p = self.heatmaps_network(self.images,reuse=reuse)

        with tf.Session(config=self.config) as sess:
            # load trained parameters
            saver = tf.train.Saver()
            saver.restore(sess, self.test_model_path)
            _, model_name = os.path.split(self.test_model_path)

            test_image = test_image.pixels_with_channels_at_back().astype('float32')
            if self.scale is '255':
                test_image *= 255
            elif self.scale is '0':
                test_image = 2 * test_image - 1

            test_image_map = sess.run(pred_hm_p, {self.images: np.expand_dims(test_image,0)})

        return test_image_map


import random
import numpy as np
from statistics import mean

# Using "random" to create a dataset
def create_dataset(hm, variance, step=2, correlation=False):
    val = 1
    sy = []
    for i in range(hm):
        y = val + random.randrange(-variance, variance)
        sy.append(y)
        if correlation and correlation == 'pos':
            val += step
        elif correlation and correlation == 'neg':
            val -= step

    sx = [i for i in range(len(sy))]

    return np.array(sx, dtype=np.float64), np.array(sy, dtype=np.float64)

# finding best fit slope and intercept of dataset
def best_fit_slope_and_intercept(sx, sy):
    m = (((mean(sx) * mean(sy)) - mean(sx * sy)) /
         ((mean(sx) * mean(sx)) - mean(sx * sx)))

    b = mean(sy) - m * mean(sx)

    return m, b

# dataset - coefficient of determination
def coefficient_of_determination(sy_orig, sy_line):
    y_mean_line = [mean(sy_orig) for y in sy_orig]

    squared_error_regr = sum((sy_line - sy_orig) * (sy_line - sy_orig))
    squared_error_y_mean = sum((y_mean_line - sy_orig) * (y_mean_line - sy_orig))

    print(squared_error_regr)
    print(squared_error_y_mean)

    r_squared = 1 - (squared_error_regr / squared_error_y_mean)

    return r_squared

def get_result():
    sx, sy = create_dataset(20, 10, 3)
    m, b = best_fit_slope_and_intercept(sx, sy)
    regression_line = [(m * x) + b for x in sx]
    r_squared = coefficient_of_determination(sy, regression_line)
    print('Rsquared result: ')
    return r_squared