# -*- 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
#from ui_viewOutput import Ui_viewOutput
from freewat.freewat_utils import getVectorLayerByName, getVectorLayerNames, getModelsInfoLists, getModelInfoByName, getTransportModelsByName
#
# load flopy and createGrid
from flopy.modflow import *
from flopy.utils import *
# Import the code for the dialog
from freewat.flopyaddon import ModflowZbd
import numpy as np
import freewat.createGrid_utils
import sys
import subprocess as sub

#
FORM_CLASS, _ = uic.loadUiType(os.path.join( os.path.dirname(__file__), 'ui/ui_runZoneBudget.ui') )
#
class runZoneBudgetDialog(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.runZoneBudget)

        # self.cmbStressPeriod.currentIndexChanged.connect(self.reloadSteps)
        # self.cancelButton.clicked.connect(self.stopProcessing)
        self.manageGui()

##
##
    def manageGui(self):
        # Retrieve list of Models and Zone MDOs, and insert them in Comboboxes
        layerNameList = getVectorLayerNames()
        #
        self.cmbModelName.clear()
        (modelNameList, pathList) =  getModelsInfoLists(layerNameList)
        self.cmbModelName.addItems(modelNameList)
        #
        zoneList = []
        for mName in layerNameList:
            if '_zone' in mName:
                zoneList.append(mName)
        self.cmbZone.addItems(zoneList)
##
    def getProgramLocation(self):
        try:
            # Retrieve programs location
            layerNameList = getVectorLayerNames()
            for layname in layerNameList:
                if layname == 'prg_locations_'+ self.cmbModelName.currentText():
                    locslayer = getVectorLayerByName(layname)
                    dirdict = {}
                    for ft in locslayer.getFeatures():
                        dirdict[ft['code']] = ft['executable']

            return dirdict
        except:
            message  = ''' You didn't enter any location for Executable files !!
                        Open Program Locations and fill in the right path
                        to your codes  '''
            QtGui.QMessageBox.information(None, 'Warning !!', message  )


##
    def runZoneBudget(self):

        # ------------ Load input data  ------------
        modelName = self.cmbModelName.currentText()
        (pathfile, nsp) = getModelInfoByName(modelName)
        zonelayerName = self.cmbZone.currentText()
        #
        # -------- Load MODFLOW Directory
        dirdict = self.getProgramLocation()
        zoneExe = dirdict['ZONE']

        # --- MODFLOW files needed
        namefile = os.path.join(pathfile, modelName + '.nam')


        # --- Load MODEL and write Zone File
        # Load Model
        ml = Modflow.load(namefile, version='mf2005', model_ws= pathfile )
        nrow, ncol, nlay, nper = ml.nrow_ncol_nlay_nper

        # Retrieve zone array
        zonelayer = getVectorLayerByName(zonelayerName)
        nzn = 1
        zone_dict = {}
        for i in range(nlay):
            zone_dict[i] = np.ones(shape = (nrow,ncol))

        for ft in zonelayer.getFeatures():
            nr = ft['ROW']-1
            nc = ft['COL']-1
            for i in range(nlay):
                ar = zone_dict[i]
                ar[nr,nc] = ft['zone_' + 'lay_' + str(i+1)]
                # Update numebr of zones
                nzn = ar.max()

        # Build Zone Budget FloPy Object, and write ZON file
        zb = ModflowZbd(ml, nzn = nzn, zone_dict = zone_dict, exe_name = zoneExe)
        zb.write_file()

        # Run Zone Budget
        try:
            from flopy.utils import ZoneBudget
            from flopy.utils import read_zbarray
        except:

            # messaggio
            message  = ''' You are attempting to get the Zone Budget but you have installed an
                        old version of FloPy.
                        Please, upgrade FloPy to version 3.2.6 (minimum)
                        and come back here later ... '''

            QtGui.QMessageBox.information(None, 'Zone Budget Available !!', message  )
            return

        # Get .zon and .cbc file
        zone_file = os.path.join(pathfile, modelName + '.zon')
        cbc_f = str(os.path.join(pathfile, modelName + '.cbc'))

        zon = read_zbarray(zone_file)
        zb = ZoneBudget(cbc_f, zon )
        zonebudget_output = modelName + '_zonebudget.csv'
        zb.to_csv(os.path.join(pathfile, zonebudget_output))

        # messaggio
        message  = 'Zone Budget has been saved as CSV file in:\n %s '%zonebudget_output

        QtGui.QMessageBox.information(None, 'Zone Budget done!', message  )



        #zb.run_zonebudget()


        # Close the dialog
        QDialog.reject(self)
