# -*- 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 PyQt4.QtCore import *
from PyQt4.QtGui import *
from qgis.core import *
import os
from PyQt4 import QtGui, uic
import matplotlib.pyplot as plt
import numpy as np
from freewat.freewat_utils import getVectorLayerByName, getVectorLayerNames, getModelsInfoLists, getModelInfoByName, getTransportModelsByName, getModelUnitsByName
# load flopy and createGrid
from flopy.modflow import *
from flopy.utils import *
import freewat.createGrid_utils
#

class viewBudgetDialog(QDialog):
    def __init__(self, iface):
        QDialog.__init__(self)
        self.iface = iface
        uic.loadUi(os.path.join( os.path.dirname(__file__), 'ui/ui_viewBudget.ui'), self)
        #self.setupUi(self)
        self.buttonBox.rejected.connect(self.reject)
        self.buttonBox.button(QDialogButtonBox.Ok).clicked.connect(self.viewBudget)

        self.cmbModelName.currentIndexChanged.connect(self.reloadTime)
        self.manageGui()

##
##
    def manageGui(self):
        self.cmbModelName.clear()
        layerNameList = getVectorLayerNames()
        (modelNameList, pathList) =  getModelsInfoLists(layerNameList)

        self.cmbModelName.addItems(modelNameList)


    def reloadTime(self):
        layerNameList = getVectorLayerNames()
        # Retrieve the model
        self.modelName = self.cmbModelName.currentText()
        (self.pathfile, self.nsp) = getModelInfoByName(self.modelName)
        SPitems = ['%s' % (i + 1) for i in range(self.nsp)]

        self.cmbStressPeriod.addItems(SPitems)

##
    def viewBudget(self):
        # LIST file of current model
        listfile = os.path.join(self.pathfile, self.modelName + '.list')
        # listfile = self.pathfile + '\\' + self.modelName + '.list'

        # Load the FloPy method
        mf_list = MfListBudget(listfile)

        # Units
        (lenuni, itmuni) = getModelUnitsByName(self.modelName)

        # Select stress period and time step
        listKstpKper = mf_list.get_kstpkper()

        # Select desired kper
        iper = int(self.cmbStressPeriod.currentText())
        # First get the all stress period and write the total budget

        budgetDictVol = {}
        budgetDictRate = {}
        # Budget always refers to last time step of a stress period
        for kperiod in range(self.nsp):
            for k,t in enumerate(listKstpKper):
                (kstp,kper) = listKstpKper[k]
                if kper == kperiod:
                    kstpkper = (kstp, kper)

            # Retrieve data
            data = mf_list.get_data(kstpkper= kstpkper)
            budgetDictVol[kperiod] = data
            data = mf_list.get_data(kstpkper= kstpkper, incremental=True)
            budgetDictRate[kperiod] = data

        # Create and Open the output file
        outputfile = os.path.join(self.pathfile, 'budget_%s.csv'%self.modelName)
        fid = open(outputfile, 'w')
        fid.write('VOLUMETRIC BUDGET of MODEL "%s" \n'%self.modelName)
        fid.write('Budget_Term, Volume StressPeriod 1, Rate StressPeriod 1')
        for k in range(1,self.nsp):
            fid.write(', Volume Stress Period %i, Rate StressPeriod %i'%(k+1,k+1))
        fid.write('\n')


        for i in range(len(data['name'])):
            fid.write('%s '%data['name'][i])
            for k in range(self.nsp):
                myarrayVol = budgetDictVol[k]
                myarrayRate = budgetDictRate[k]
                fid.write(', %10f , %10f'%(myarrayVol['value'][i], myarrayRate['value'][i]) )
            fid.write('\n')
        fid.close()

        QtGui.QMessageBox.information(None, 'Saving Volumetric Budget', 'Budget is saved as CSV file  in: \n' + outputfile )

        # Plot car chart of volumetric budget
        # Plot only SP selected in GUI
        data = budgetDictVol[iper - 1]
        plt.bar(data['index'], data['value'])
        plt.xticks(data['index'], data['name'], rotation=25, size=9.0)
        plt.ylabel('Water Volume (in %s^3)'%lenuni)
        plt.grid()
        plt.show()

        # Close the dialog
        QDialog.reject(self)
