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

from qgis.core import *
from PyQt4.QtGui import *
from PyQt4.QtCore import *
import os
from PyQt4 import QtGui, uic
from freewat.freewat_utils import getVectorLayerByName, getVectorLayerNames, getModelsInfoLists, getModelInfoByName, getTransportModelsByName
# load flopy and createGrid
from flopy.modflow import *
from flopy.utils import *
import freewat.createGrid_utils
import numpy as np
import matplotlib.pyplot as plt
#
FORM_CLASS, _ = uic.loadUiType(os.path.join( os.path.dirname(__file__), 'ui/ui_plotWaterUnitBudgets.ui') )
#
class plotWaterUnitBudgetDialog(QDialog, FORM_CLASS):
    def __init__(self, iface):
        QDialog.__init__(self)
        self.iface = iface
        self.setupUi(self)
        self.buttonBox.rejected.connect(self.reject)
        self.buttonBox.button(QDialogButtonBox.Ok).clicked.connect(self.plotBudgets)
        self.cmbModelName.currentIndexChanged.connect(self.reloadFarms)
        self.manageGui()
##
##
    def manageGui(self):
        self.cmbModelName.clear()
        layerNameList = getVectorLayerNames()
        (modelNameList, pathList) =  getModelsInfoLists(layerNameList)

        self.cmbModelName.addItems(modelNameList)


    def reloadFarms(self):
        # Retrieve number of farms
        self.modelName = self.cmbModelName.currentText()
        (self.pathfile, nsp ) = getModelInfoByName(self.modelName)
        #
        layerNameList = getVectorLayerNames()
        for mName in layerNameList:
            if mName == 'WaterDemandSites_'+ self.modelName:
                waterunitsTable = getVectorLayerByName(mName)

            if mName == 'timetable_' + self.modelName:
                timelayer = getVectorLayerByName(mName)
                ## TO DO:
                # The faster method is the following but it seems
                # that it counts only geom features, no one in time_table:
                # nsp = int(timelayer.featureCount())
                ftit = timelayer.getFeatures()
                timesteps  = 0
                #tslist = []
                for f in ftit:
                    timesteps = timesteps + int(f['ts'])
                    #tslist.append(f['time_steps'])
        self.timesteps = timesteps
        wu_list = []
        for f in waterunitsTable.getFeatures():
            wu_list.append(str(f['Farm_ID']))

        wu_list.sort()
        self.nfarm = len(wu_list)
        self.cmbWaterUnit.clear()
        self.cmbWaterUnit.addItems(wu_list)

##
    def plotBudgets(self):

        # ------------ Load the model name and model data data  ------------
        ## TO be concluded
        # Get nfarms and nrows
        fm_selected =  int(self.cmbWaterUnit.currentText())
        nfarm = self.nfarm
        nrows = (self.timesteps)*nfarm

        # DELIVERY SPECIFIC BUDGET (FDS)
        if self.chkFds.isChecked():
            FDSfile = os.path.join(self.pathfile, 'FDS.OUT')
            # FDSfile = self.pathfile + '\\' + 'FDS.OUT'
            dtype=[('PER', '<i8'), ('STP', '<i8'), ('DAYS', '<f8'), ('FID', '<i8'),  ('OFE', '<f8'),('TFDR-INI', '<f8'),('NR-SWD-INI', '<f8'),('R-SWD-INI', '<f8'),('QREQ-INI', '<f8'),
            ('TFDR-FIN', '<f8'), ('NR-SWD-FIN', '<f8'), ('R-SWD-FIN', '<f8'), ('QREQ-FIN', '<f8'), ('Q-FIN', '<f8'), ('DEF-FLAG', '<i8') ]
            aa = np.genfromtxt(FDSfile , skip_header = 1, dtype = dtype)
            # Dictionary of budget terms for each farm
            fm_dict = {}
            for fid in range(1,nfarm+1):
                f1 = np.recarray((nrows/nfarm,),dtype = dtype)
                f1['DAYS'][0] = aa['DAYS'][fid-1]
                f1['R-SWD-INI'][0] = aa['RSWDINI'][fid-1]
                f1['NR-SWD-INI'][0] = aa['NRSWDINI'][fid-1]
                f1['TFDR-INI'][0] = aa['TFDRINI'][fid-1]
                f1['Q-FIN'][0] = aa['QFIN'][fid-1]
                j = 1
                fmcounter = nfarm + 1 - fid
                #print fmcounter
                for i in range(fid-1,len(aa)-fmcounter,nfarm):
                    f1['DAYS'][j] = aa['DAYS'][i+nfarm]
                    f1['NR-SWD-INI'][j] = aa['NRSWDINI'][i+nfarm]
                    f1['R-SWD-INI'][j] = aa['RSWDINI'][i+nfarm]
                    f1['TFDR-INI'][j]= aa['TFDRINI'][i+nfarm]
                    f1['Q-FIN'][j] = aa['QFIN'][i+nfarm]
                    j += 1

                fm_dict[fid] = f1
                del f1
            # Plotting
            fig, ax = plt.subplots()
            ax.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['TFDR-INI'] , 'b-', label = 'TFDR')
            if max(abs(fm_dict[fm_selected]['R-SWD-INI'])) > 0:
                ax.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['R-SWD-INI'] , 'g*-', label = 'Surface Water Delivery')
            if max(abs(fm_dict[fm_selected]['Q-FIN'])) > 0:
                ax.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-FIN'] , 'rs', label = 'Well Pumping')
            if max(abs(fm_dict[fm_selected]['NR-SWD-INI'])) > 0:
                ax.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['NR-SWD-INI'] , 'go', label = 'External Delivery')
            ax.set_xlabel('Time (days)')
            ax.set_ylabel('Water Flow  (L^3/T)')
            ax.set_title('Water Unit n.%i'%fm_selected)
            ax.legend().draggable()
            ax.grid()

        # BUDGET DETAILS  ------------
        if self.chkFbdIn.isChecked() or self.chkFbdOut.isChecked():
            FBDfile = os.path.join(self.pathfile, 'FB_DETAILS.OUT')
            # FBDfile = self.pathfile + '\\' + 'FB_DETAILS.OUT'
            dtype=[('PER', '<i8'), ('STP', '<i8'), ('DAYS', '<f8'), ('FID', '<i8'),  ('Q-p-in','<f8'),('Q-nrd-in','<f8'),('Q-srd-in','<f8'),
            ('Q-rd-in','<f8'),('Q-wells-in', '<f8'),('Q-egw-in','<f8'), ('Q-tgw-in', '<f8'),('Q-ext-in', '<f8'), ('Q-tot-in', '<f8'),('Q-ei-out', '<f8'), ('Q-ep-out', '<f8'),('Q-egw-out','<f8'),('Q-ti-out','<f8'),('Q-tp-out', '<f8'), ('Q-tgw-out', '<f8'), ('Q-run-out', '<f8'),('Q-dp-out','<f8'),('Q-nrd-out', '<f8'), ('Q-srd-out', '<f8'), ('Q-rd-out', '<f8'), ('Q-wells-out', '<f8'), ('Q-tot-out', '<f8')]
            aa = np.genfromtxt(FBDfile , skip_header = 1, dtype = dtype)
            # FB (INFLOW)
            if self.chkFbdIn.isChecked():
                # Dictionary of budget terms for each farm
                fm_dict = {}
                for fid in range(1,nfarm+1):
                    f1 = np.recarray((nrows/nfarm,),dtype = dtype)
                    f1['DAYS'][0] = aa['DAYS'][fid-1]
                    f1['Q-p-in'][0] = aa['Qpin'][fid-1]
                    f1['Q-nrd-in'][0] = aa['Qnrdin'][fid-1]
                    f1['Q-srd-in'][0] = aa['Qsrdin'][fid-1]
                    f1['Q-wells-in'][0] = aa['Qwellsin'][fid-1]
                    f1['Q-egw-in'][0] = aa['Qegwin'][fid-1]
                    f1['Q-tgw-in'][0] = aa['Qtgwin'][fid-1]
                    f1['Q-ext-in'][0] = aa['Qextin'][fid-1]
                    f1['Q-tot-in'][0] = aa['Qtotin'][fid-1]
                    j = 1
                    fmcounter = nfarm + 1 - fid
                    for i in range(fid-1,len(aa)-fmcounter,nfarm):
                        f1['DAYS'][j] = aa['DAYS'][i+nfarm]
                        f1['Q-p-in'][j] = aa['Qpin'][i+nfarm]
                        f1['Q-nrd-in'][j] = aa['Qnrdin'][i+nfarm]
                        f1['Q-srd-in'][j] = aa['Qsrdin'][i+nfarm]
                        f1['Q-wells-in'][j] = aa['Qwellsin'][i+nfarm]
                        f1['Q-egw-in'][j] = aa['Qegwin'][i+nfarm]
                        f1['Q-tgw-in'][j] = aa['Qtgwin'][i+nfarm]
                        f1['Q-ext-in'][j] = aa['Qextin'][i+nfarm]
                        f1['Q-tot-in'][j] = aa['Qtotin'][i+nfarm]
                        j += 1

                    fm_dict[fid] = f1
                    del f1
                # Plotting
                fig2, ax2 = plt.subplots()
                ax2.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-p-in'] , 'b-', label = 'Precipitation')
                ax2.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-tot-in'] , 'k-', label = 'Total inflow into Water Unit')
                if max(abs(fm_dict[fm_selected]['Q-nrd-in'])) > 0:
                    ax2.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-nrd-in'] , 'g*', label = 'Non-routed Water Delivery')
                if max(abs(fm_dict[fm_selected]['Q-srd-in'])) > 0:
                    ax2.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-srd-in'] , 'r*-', label = 'Surface Water Delivery')
                if max(abs(fm_dict[fm_selected]['Q-wells-in'])) > 0:
                    ax2.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-wells-in'] , 'go-', label = 'Well Pumping')
                if max(abs(fm_dict[fm_selected]['Q-egw-in'])) > 0:
                    ax2.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-egw-in'] , 'r-', label = 'Evaporation from GW')
                if max(abs(fm_dict[fm_selected]['Q-tgw-in'])) > 0:
                    ax2.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-tgw-in'] , 'g-', label = 'Transpiration from GW')
                if max(abs(fm_dict[fm_selected]['Q-ext-in'])) > 0:
                    ax2.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-ext-in'] , 'rs-', label = 'External Water Delivery')
                ax2.set_xlabel('Time (days)')
                ax2.set_ylabel('Water Flow  (L^3/T)')
                ax2.set_title('Inflows to Water Unit n.%i'%fm_selected)
                ax2.legend().draggable()
                ax2.grid()

            # BUDGET DETAILS  FB (OUTFLOW)
            if self.chkFbdOut.isChecked():
                # Dictionary of budget terms for each farm
                fm_dict = {}
                for fid in range(1,nfarm+1):
                    f1 = np.recarray((nrows/nfarm,),dtype = dtype)
                    f1['DAYS'][0] = aa['DAYS'][fid-1]
                    f1['Q-ei-out'][0] = aa['Qeiout'][fid-1]
                    f1['Q-ep-out'][0] = aa['Qepout'][fid-1]
                    f1['Q-egw-out'][0] = aa['Qegwout'][fid-1]
                    f1['Q-ti-out'][0] = aa['Qtiout'][fid-1]
                    f1['Q-tp-out'][0] = aa['Qtpout'][fid-1]
                    f1['Q-tgw-out'][0] = aa['Qtgwout'][fid-1]
                    f1['Q-run-out'][0] = aa['Qrunout'][fid-1]
                    f1['Q-dp-out'][0] = aa['Qdpout'][fid-1]
                    f1['Q-tot-out'][0] = aa['Qtotout'][fid-1]
                    j = 1
                    fmcounter = nfarm + 1 - fid
                    for i in range(fid-1,len(aa)-fmcounter,nfarm):
                        f1['DAYS'][j] = aa['DAYS'][i+nfarm]
                        f1['Q-ei-out'][j] = aa['Qeiout'][i+nfarm]
                        f1['Q-ep-out'][j] = aa['Qepout'][i+nfarm]
                        f1['Q-egw-out'][j] = aa['Qegwout'][i+nfarm]
                        f1['Q-ti-out'][j] = aa['Qtiout'][i+nfarm]
                        f1['Q-tp-out'][j] = aa['Qtpout'][i+nfarm]
                        f1['Q-tgw-out'][j] = aa['Qtgwout'][i+nfarm]
                        f1['Q-run-out'][j] = aa['Qrunout'][i+nfarm]
                        f1['Q-dp-out'][j] = aa['Qdpout'][i+nfarm]
                        f1['Q-tot-out'][j] = aa['Qtotout'][i+nfarm]
                        j += 1

                    fm_dict[fid] = f1
                    del f1
                # Plotting
                fig3, ax3 = plt.subplots()
                ax3.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-tot-out'] , 'k-', label = 'Total outflow from Water Unit')
                if max(abs(fm_dict[fm_selected]['Q-ei-out'])) > 0:
                    ax3.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-ei-out'] , 'g*-', label = 'Evap. from irrigation')
                if max(abs(fm_dict[fm_selected]['Q-ep-out'])) > 0:
                    ax3.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-ep-out'] , 'b*-', label = 'Evap. from precipitation')
                if max(abs(fm_dict[fm_selected]['Q-egw-out'])) > 0:
                    ax3.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-egw-out'] , 'r*-', label = 'Evaporation from GW')

                if max(abs(fm_dict[fm_selected]['Q-ti-out'])) > 0:
                    ax3.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-ti-out'] , 'go-', label = 'Transp. from Irrigation')
                if max(abs(fm_dict[fm_selected]['Q-tp-out'])) > 0:
                    ax3.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-tp-out'] , 'go-', label = 'Transp. from Precipitation')

                if max(abs(fm_dict[fm_selected]['Q-tgw-out'])) > 0:
                    ax3.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-tgw-out'] , 'bo-', label = 'Transp. from GW')


                if max(abs(fm_dict[fm_selected]['Q-run-out'])) > 0:
                    ax3.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-run-out'] , 'ms-', label = 'Runoff')

                if max(abs(fm_dict[fm_selected]['Q-dp-out'])) > 0:
                    ax3.plot(fm_dict[fm_selected]['DAYS'],fm_dict[fm_selected]['Q-dp-out'] , 'mo-', label = 'Deep percolation')

                ax3.set_xlabel('Time (days)')
                ax3.set_ylabel('Water Flow  (L^3/T)')
                ax3.set_title('Outflows out Water Unit n.%i'%fm_selected)
                ax3.legend().draggable()
                ax3.grid()

        plt.show()
