# -*- 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

#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_obh


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.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.manageGui()


        self.cmb_SP.currentIndexChanged.connect(self.Refresh)
        self.cmb_SP.currentIndexChanged.connect(self.CreatePlot)



    def manageGui(self):


        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 FILE
        self.hob_path = os.path.join(pathfile, self.modelName + '.hob')

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

        # 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)


        # First run ---
        self.Refresh()
        self.CreatePlot()


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

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



    def saveStatistics(self):

        self.mp = ModelPath()
        self.path_statistics = os.path.join(self.mp, "{}_{}_{}.{}".format(self.modelName, 'residuals_statistics_sp', self.spNumber, '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, "{}_{}_{}.{}".format(self.modelName, 'residuals_sp', self.spNumber, '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 k, v in self.d.items():
                if k == self.spNumber:
                    # print self.d[k]['residuals']
                    for itm, itm2 in zip(self.d[k]['residuals'], self.d[k]['labels']):
                        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 CreatePlot(self):

        # clear the axes
        self.axes.clear()


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

        try:
            # get the obh file path
            obh_path = os.path.join(pathfile, self.modelName + '.obh')
        except:
            QMessageBox.warning(self.iface.mainWindow(), self.tr("Warning"), self.tr("No OBH file found! Be sure that you run the Modflow simulation with Observation package active"))
            return

        # convert the obh file into a np array
        obh = np.genfromtxt(obh_path , skip_header = 1, dtype=None)


        # dictionary filled with the obh file values, residual calculation and stress period corresponding number taken from the related MDO layer
        self.row_data = {}

        self.row_data['simulated'] = [i[0] for i in obh]
        self.row_data['observed'] = [i[1] for i in obh]
        self.row_data['labels'] = [i[2] for i in obh]
        self.row_data['residuals'] = [x - y for x, y in zip(self.row_data['simulated'], self.row_data['observed'])]
        self.row_data['sp'] = []
        # update the sp key with values of the mdo obseration FILE
        for i in self.total_sp:
            self.row_data['sp'].append(int(i))


        # dictionary with key = SP and values are, other dictionary with values lists
        self.d = {}
        for i, v in enumerate(self.row_data['sp']):
            self.d.setdefault(v, {'simulated':[],'observed':[], 'labels':[], 'residuals':[]})
            self.d[v]['simulated'].append(self.row_data['simulated'][i])
            self.d[v]['observed'].append(self.row_data['observed'][i])
            self.d[v]['labels'].append(self.row_data['labels'][i])
            self.d[v]['residuals'].append(self.row_data['residuals'][i])


        # connect the corresponding stress period with the plot

        # list for the min and max values (axes limits)
        x_y = []

        for k,v in self.d.items():
            if k == self.spNumber:
                self.axes.scatter(self.d[k]['observed'], self.d[k]['simulated'], c='r')
                # get the min and max for the axes limits
                if min(self.d[k]['observed']) < min(self.d[k]['simulated']):
                    x_y.append(min(self.d[k]['observed']))
                else:
                    x_y.append(min(self.d[k]['simulated']))
                if max(self.d[k]['observed']) > max(self.d[k]['simulated']):
                    x_y.append(max(self.d[k]['observed']))
                else:
                    x_y.append(max(self.d[k]['simulated']))



                # calculate the single statistics
                res_count = len(self.d[k]['residuals'])
                res_mean = np.mean(self.d[k]['residuals'])
                res_sd = np.std(self.d[k]['residuals'])
                res_min = min(self.d[k]['residuals'])
                res_max = max(self.d[k]['residuals'])

                res_abs = np.mean(np.abs(self.d[k]['residuals']))
                res_stderr = np.sqrt(np.mean([((i - res_mean)**2) / (res_count - 1) for i in self.d[k]['residuals']]))
                res_rms = np.sqrt(np.mean(self.d[k]['residuals'])**2)
                res_nrms = res_rms / (max(self.d[k]['observed']) - min(self.d[k]['observed']))
                (res_corr, p) = stats.pearsonr(self.d[k]['observed'], self.d[k]['simulated'])

                # 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_nrms), str(res_corr)]
        }




        # labels on each point of the plot if the checkbox is checked
        if self.chk_labels.isChecked():
            for k,v in self.d.items():
                if k == self.spNumber:
                    for i, txt in enumerate(self.d[k]['labels']):
                        self.axes.annotate(txt, xy=(self.d[k]['observed'][i], self.d[k]['simulated'][i]))



        # plot title
        self.axes.set_title('Residuals of the {} Stress Period'.format(self.spNumber))


        # set the axes limits
        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')

        # legend
        self.axes.legend(loc='upper left')


        # 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))
