# -*- coding: utf-8 -*-

#******************************************************************************
#
# Freewat
# ---------------------------------------------------------
#
#
# Copyright (C) 2014 - 2015 Iacopo Borsi (iacopo.borsi@tea-group.com)
#
# This source is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 2 of the License, or (at your option)
# any later version.
#
# This code is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
# details.
#
# A copy of the GNU General Public License is available on the World Wide Web
# at <http://www.gnu.org/licenses/>. You can also obtain it by writing
# to the Free Software Foundation, 51 Franklin Street, Suite 500 Boston,
# MA 02110-1335 USA.
#
#******************************************************************************


import os
from PyQt4.QtCore import *
from PyQt4.QtGui import *
from PyQt4 import QtGui, uic
from qgis.core import *
from PyQt4 import QtCore
from qgis.gui import QgsMessageBar

#matplotlib imports
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
if matplotlib.__version__ >= '1.5.0':
    from matplotlib.backends.backend_qt4agg import NavigationToolbar2QT as NavigationToolbar
else:
    from matplotlib.backends.backend_qt4agg import NavigationToolbar2QTAgg as NavigationToolbar
from matplotlib.figure import Figure

# numpy ans scipy imports
import numpy as np
import scipy.stats as stats


from freewat.freewat_utils import getVectorLayerNames, getModelsInfoLists, getModelInfoByName, ModelPath, load_hob, load_obh, getTimeTableInfo, getTSFromOffset


FORM_CLASS, _ = uic.loadUiType(os.path.join( os.path.dirname(__file__), 'ui/ui_SimplePlotCalibration.ui') )

class SimplePlotCalibration(QDialog, FORM_CLASS):
    def __init__(self, iface):
        QDialog.__init__(self)
        self.iface = iface
        self.setupUi(self)

        # self.bar = QgsMessageBar()
        # self.bar.setSizePolicy( QSizePolicy.Minimum, QSizePolicy.Fixed )
        # # self.layout().setContentsMargins(0, 0, 0, 0)
        # self.layout().addWidget(self.bar, 0,0,1,0)

        # self.figure, self.axes = plt.subplots()
        self.figure = Figure()
        self.axes = self.figure.add_subplot(111)
        self.canvas = FigureCanvas(self.figure)
        self.mpltoolbar = NavigationToolbar(self.canvas, self.widgetPlot)
        lstActions = self.mpltoolbar.actions()
        self.mpltoolbar.removeAction(lstActions[7])
        self.layoutPlot.addWidget(self.canvas)
        self.layoutPlot.addWidget(self.mpltoolbar)

        self.btn_save.clicked.connect(self.saveStatistics)
        self.btn_save_residuals.clicked.connect(self.saveResiduals)
        self.chk_labels.stateChanged.connect(self.CreatePlot, self.cmb_TS.currentIndex())
        self.btn_bubble.clicked.connect(self.createBubbleLayer)

        # create the dictionary with SP as keys and nested dictionary with TS as keys and ranges as values
        # it has to be created here in order to have SP and corresponding TS to populate the comboboxes

        self.spts = getTimeTableInfo()

        self.cmbModelName.clear()
        layerNameList = getVectorLayerNames()
        layerNameList.sort()

        # Retrieve modelname and pathfile List
        (modelNameList, pathList) =  getModelsInfoLists(layerNameList)
        self.cmbModelName.addItems(modelNameList)
        self.modelName = self.cmbModelName.currentText()

        # get the stress period number
        (pathfile, nsp) = getModelInfoByName(self.modelName)

        # get the path of the hob and obh files
        self.hob_path = os.path.join(pathfile, self.modelName + '.hob')
        self.obh_path = os.path.join(pathfile, self.modelName + '.obh')

        # get the lists of unique SP, total SP, total_TS and Labels from the loaded hob file in the working directory
        try:
            self.unique_sp, self.total_sp, self.total_ts, self.labels, self.mod_lay = load_hob(self.hob_path)
        except:
            QMessageBox.warning(self.iface.mainWindow(), self.tr("Warning"), self.tr("No HOB file in your working directory found!"))
            return

        # get the lists of simulated, observed and labels from the OBH file
        try:
            self.sim, self.obs, self.lab = load_obh(self.obh_path)
        except:
            QMessageBox.warning(self.iface.mainWindow(), self.tr("Warning"), self.tr("No OBH file in your working directory found!"))
            return

        # create the list of corresponding TS using the getTSFromOffset function
        self.ts_list = []
        for i, j in enumerate(self.total_sp):
            self.ts_list.append(getTSFromOffset(self.total_ts[i], self.spts[j]))


        # trying to create the final dictionary with all the values
        # the final dict will be d = {SP:{TS{'sim':[], 'res':[], 'lab':[]}}}
        # sub-divided for each SP and for each TS
        self.data_dic = {}
        for i in self.total_sp:
            self.data_dic[i] = {}

        for i, j in zip(self.total_sp, self.ts_list):
            if j not in self.data_dic[i]:
                self.data_dic[i][j] = {'sim':[], 'obs':[], 'res':[], 'lab':[], 'mod_lay':[]}

        for i, j in enumerate(self.total_sp):
            if self.ts_list[i] not in self.data_dic[j][self.ts_list[i]]:
                self.data_dic[j][self.ts_list[i]]['sim'].append(self.sim[i])
                self.data_dic[j][self.ts_list[i]]['obs'].append(self.obs[i])
                self.data_dic[j][self.ts_list[i]]['res'].append(self.sim[i] - self.obs[i])
                self.data_dic[j][self.ts_list[i]]['lab'].append(self.lab[i])
                self.data_dic[j][self.ts_list[i]]['mod_lay'].append(self.mod_lay[i])


        # transform the numeric list into string to load in the QCombobox
        self.unique_sp = [str(i) for i in self.unique_sp]

        # load the unique SP in the Combobox
        self.cmb_SP.addItems(self.unique_sp)
        self.cmb_SP.addItem('--all SP--')

        self.cmb_SP.currentIndexChanged.connect(self.Refresh)
        self.cmb_TS.currentIndexChanged.connect(self.CreatePlot, self.cmb_TS.currentIndex())
        self.Refresh()


    def Refresh(self):
        # clear the axes for the next plots
        self.axes.clear()

        # get the current stress period number from the combobox
        try:
            self.spNumber = int(self.cmb_SP.currentText())
        except:
            self.spNumber = self.cmb_SP.currentText()

        if self.spNumber == '--all SP--':
            self.cmb_TS.setEnabled(False)
            self.label_TS.setEnabled(False)
        else:
            self.cmb_TS.setEnabled(True)
            self.label_TS.setEnabled(True)

        self.cmb_TS.clear()
        if not self.spNumber == '--all SP--':
            for k in self.data_dic[self.spNumber].keys():
                self.cmb_TS.addItem(str(k))

        self.cmb_TS.addItem('--all TS--')


    def CreatePlot(self, idx):
        '''
        method that create the final plot

        the input variable idx is necessary to avoid the calling of the function
        as a result of double checking the currentIndexChanged of both SP and TS
        comboboxes

        if no elements appear in the combobox the currentIndex is -1, else is
        the index of the element chosen
        '''

        # break the code if idx = -1, that is, if the index of the TS combobox
        # is not filled with anything
        if idx == -1:
            pass

        # run the function according to the UI
        else:

            if not self.spNumber == '--all SP--':
                self.tsNumber = self.data_dic[self.spNumber].keys()[0]
            # get the current TS number or string from the combobox
            try:
                self.tsNumber = int(self.cmb_TS.currentText())
            except:
                self.tsNumber = self.cmb_TS.currentText()

            # clear the axes
            self.axes.clear()

            # get information of pathfile and modelname
            self.modelName = self.cmbModelName.currentText()
            (pathfile, nsp) = getModelInfoByName(self.modelName)


            # create the residual list according to SP and TS (used to calculate the statistics)
            self.res = []
            self.filt_sim = []
            self.filt_obs = []
            self.filt_lab = []
            self.filt_mod_lay = []


            # if all the SP are chosen, without filtering by TS
            if self.spNumber == '--all SP--':
                # reset the progressbar value
                self.progressBar.setValue(0)

                for sp, values in self.data_dic.items():
                    for ts, in_values in values.items():
                        for i, j in enumerate(in_values['sim']):

                            # update the progressBar value
                            value = self.progressBar.value() + 1
                            self.progressBar.setValue(value)
                            self.axes.scatter(in_values['obs'][i], in_values['sim'][i], c='r')
                            self.res.append(in_values['res'][i])
                            self.filt_sim.append(in_values['sim'][i])
                            self.filt_obs.append(in_values['obs'][i])
                            self.filt_lab.append(in_values['lab'][i])
                            self.filt_mod_lay.append(in_values['mod_lay'][i])

                self.progressBar.setValue(100)


            # if all TS for a certain SP have been chosen
            elif self.tsNumber == '--all TS--':
                for k, v in self.data_dic[self.spNumber].items():
                    for i, j in enumerate(v['obs']):
                        self.axes.scatter(v['obs'][i], v['sim'][i], c='r')
                        self.res.append(v['res'][i])
                        self.filt_sim.append(v['sim'][i])
                        self.filt_obs.append(v['obs'][i])
                        self.filt_lab.append(v['lab'][i])
                        self.filt_mod_lay.append(v['mod_lay'][i])

                    count = len(self.res)
                    percent = i / float(count) * 100
                    self.progressBar.setValue(percent)

                self.progressBar.setValue(100)

            # it one TS for one SP is chosen
            else:
                if not self.spNumber == '--all SP--':
                    self.axes.scatter(self.data_dic[self.spNumber][self.tsNumber]['obs'], self.data_dic[self.spNumber][self.tsNumber]['sim'], c='r')
                    self.res = self.data_dic[self.spNumber][self.tsNumber]['res']
                    self.filt_sim = self.data_dic[self.spNumber][self.tsNumber]['sim']
                    self.filt_obs = self.data_dic[self.spNumber][self.tsNumber]['obs']
                    self.filt_lab = self.data_dic[self.spNumber][self.tsNumber]['lab']
                    self.filt_mod_lay = self.data_dic[self.spNumber][self.tsNumber]['mod_lay']

                    for i, j in enumerate(self.data_dic[self.spNumber][self.tsNumber]['obs']):

                        count = len(self.data_dic[self.spNumber][self.tsNumber]['obs'])
                        percent = i / float(count) * 100
                        self.progressBar.setValue(percent)

                    self.progressBar.setValue(100)

            # create a list for the axis limits (necessary for the confidence intervals)
            x_y = []

            if self.res:
                if min(self.filt_obs) < min(self.filt_sim):
                    x_y.append(min(self.filt_obs))
                else:
                    x_y.append(min(self.filt_sim))
                if max(self.filt_obs) > max(self.filt_sim):
                    x_y.append(max(self.filt_obs))
                else:
                    x_y.append(max(self.filt_sim))


            # check if self.res is not empty. This happen when for that SP no Observation are found at that TS
            if self.res:
                # calculate the single statistics
                res_count = len(self.res)
                res_mean = np.mean(self.res)
                res_sd = np.std(self.res)
                res_min = min(self.res)
                res_max = max(self.res)
                res_abs = np.mean(np.abs(self.res))
                res_stderr = np.sqrt(np.mean([((i - res_mean)**2) / (res_count - 1) for i in self.res]))
                res_rms = np.sqrt(np.mean(self.res)**2)
                res_nrms = res_rms / (max(self.filt_obs) - min(self.filt_obs))
                (res_corr, p) = stats.pearsonr(self.filt_obs, self.filt_sim)

                # confidence intervals
                t95 = stats.t.ppf(0.95, res_count)
                t90 = stats.t.ppf(0.90, res_count)


                # fill the dictionary with the calculated statistics
                self.data = {
                    'Residual Statistics':[
                        'Residual count','Mean', 'Sd', 'Min', 'Max', 'Absolute Residual Mean',\
                        'Standard Error of the Estimate', 'Residual RMS', 'Normalized RMS', \
                        'Pearson Correlation Coefficient'
                    ],
                    'Value':[
                        str(res_count), str(res_mean), str(res_sd), str(res_min), str(res_max), \
                        str(res_abs), str(res_stderr), str(res_rms), str(res_rms), str(res_nrms)
                    ]
                }


            # labels on each point of the plot if the checkbox is checked
            if self.chk_labels.isChecked():
                for i, txt in enumerate(self.filt_lab):
                    self.axes.annotate(txt, xy=(self.filt_obs[i], self.filt_sim[i]))


            self.axes.set_title('Residuals of the {} Stress Period and {} Time Step'.format(self.spNumber, self.tsNumber))

            # set the axes limits
            if x_y:
                self.axes.set_xlim(x_y[0] - 1, x_y[1] + 1)
                self.axes.set_ylim(x_y[0] - 1, x_y[1] + 1)

                # plot the 45 degree line
                self.axes.plot((x_y[0] - 1, x_y[1] + 1), (x_y[0] - 1, x_y[1] + 1), 'k', label = 'Bisector Line')
                # 95 confidence interval line
                self.axes.plot((x_y[0] - 1, x_y[1] + 1), (x_y[0] - 1, x_y[1] + 1) + t95, 'b--', label = '95% Confidence Interval')
                self.axes.plot((x_y[0] - 1, x_y[1] + 1), (x_y[0] - 1, x_y[1] + 1) - t95, 'b--')
                # 90 confidence interval line
                self.axes.plot((x_y[0] - 1, x_y[1] + 1), (x_y[0] - 1, x_y[1] + 1) + t90, 'g--', label = '90% Confidence Interval')
                self.axes.plot((x_y[0] - 1, x_y[1] + 1), (x_y[0] - 1, x_y[1] + 1) - t90, 'g--')

            # axes labels
            self.axes.set_xlabel('Observed')
            self.axes.set_ylabel('Simulated')

            # draw the final plot
            self.canvas.draw()

            # fill the QTableWidget with all the statistics
            self.statTable.setRowCount(len(self.data['Residual Statistics']))
            self.statTable.setColumnCount(2)

            # hide row numbers
            self.statTable.verticalHeader().setVisible(False)


            horHeaders = []

            for n, key in enumerate(sorted(self.data.keys())):
                horHeaders.append(self.tr(key))
                for m, item in enumerate(self.data[key]):
                    newitem = QtGui.QTableWidgetItem(item)
                    self.statTable.setItem(m, n, newitem)

            self.statTable.setHorizontalHeaderLabels(horHeaders)
            header = self.statTable.horizontalHeader()

            # resize the columns accordingly to their content
            self.statTable.resizeColumnToContents(0)
            self.statTable.resizeColumnToContents(1)
            # adjust the table dimension as
            self.statTable.setFixedWidth(self.statTable.columnWidth(0) + self.statTable.columnWidth(1))


    def saveStatistics(self):

        self.mp = ModelPath()
        self.path_statistics = os.path.join(self.mp, "{}_residuals_statistics_sp_{}_ts_{}.{}".format(self.modelName, self.spNumber, self.tsNumber, 'csv'))


        if os.path.exists(self.path_statistics):
            QMessageBox.information(self.iface.mainWindow(), self.tr("Warning"), self.tr("File already exists!"))
            return

        with open(self.path_statistics, 'wb') as f:
            for i in xrange(len(self.data['Residual Statistics'])):
                f.write("{}, {}\n".format(self.data['Residual Statistics'][i], self.data['Value'][i]))

        QMessageBox.information(self.iface.mainWindow(), self.tr("Information"), self.tr("{}\n{}\n{}".format("Statistics file", self.path_statistics, "saved in your working directory")))

    def saveResiduals(self):

        self.mp = ModelPath()
        self.path_residuals = os.path.join(self.mp, "{}_residuals_sp_{}_ts_{}.{}".format(self.modelName, self.spNumber, self.tsNumber, 'csv'))


        if os.path.exists(self.path_residuals):
            QMessageBox.information(self.iface.mainWindow(), "Warning", "File already exists!")
            return

        with open(self.path_residuals, 'wb') as ff:
            ff.write("{}, {}\n".format("Residuals", "Observation name"))
            for itm, itm2 in zip(self.res, self.filt_lab):
                ff.write("{}, {}\n".format(itm, itm2))

        QMessageBox.information(self.iface.mainWindow(), self.tr("Information"), self.tr("{}\n{}\n{}".format("Residual file", self.path_residuals, "saved in your working directory")))

    def createBubbleLayer(self):

        layerList = QgsMapLayerRegistry.instance().mapLayers()

        for k, v in layerList.items():
            if v.name().endswith('_hob'):
                self.head_obs = v
            if v.name().startswith('modeltable'):
                model = v

        for ft in model.getFeatures():
            crs = ft['crs']

        # reset the progressBar value
        self.progressBar.setValue(0)

        for num, mod in enumerate(set(self.filt_mod_lay)):

            if self.spNumber == '--all SP--':
                name = "residual_lay_all_SP"
            else:
                name = "residual_lay_{}_sp_{}_ts_{}".format(mod, self.spNumber, self.tsNumber)

            # create the point vector layer with CRS taken from the modellayer table
            vl = QgsVectorLayer("Point?crs=" + crs, name, "memory")
            pr = vl.dataProvider()

            # add the fields to the temporary vector layer
            pr.addAttributes([QgsField("OBSNAM", QVariant.String),
                              QgsField("residuals", QVariant.Double),
                              QgsField("observed", QVariant.Double),
                              QgsField("simulated", QVariant.Double),
                              QgsField("layer", QVariant.Int),
            ])

            vl.updateFields()

            # loop and add the features to the new vector layer
            for lab in self.head_obs.getFeatures():
                for i, flab in enumerate(self.filt_lab):
                    if lab["OBSNAM"] == flab:
                        if self.filt_mod_lay[i] == mod:
                            fet = QgsFeature()
                            fet.setGeometry(QgsGeometry.fromPoint(QgsPoint(lab.geometry().centroid().asPoint())))
                            fet.setAttributes([flab, self.res[i], self.filt_obs[i], self.filt_sim[i], self.filt_mod_lay[i]])
                            pr.addFeatures([fet])

                # update the progressBar values
                value = self.progressBar.value() + 1
                self.progressBar.setValue(value)

            # add the vector layer to the TOC
            QgsMapLayerRegistry.instance().addMapLayer(vl)

        self.progressBar.setValue(100)
