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

#******************************************************************************
#
# Freewat
# ---------------------------------------------------------
#
#
# Copyright (C) 2014 - 2015 Iacopo Borsi (iacopo.borsi@tea-group.com)
#
#
#******************************************************************************
from PyQt4.QtCore import *
from PyQt4.QtGui import *
from qgis.core import *
import os
from PyQt4 import QtGui, uic
from freewat.freewat_utils import getVectorLayerByName, getVectorLayerNames, getModelsInfoLists, getModelInfoByName
from freewat.mdoCreate_utils import createMDO
from freewat.sqlite_utils import getTableNamesList
#
from flopy.modflow import *
from flopy.utils import *
import freewat.createGrid_utils
# FloPy Add-On
from freewat.flopyaddon import ModflowUsb, ModflowUzf1_new
import numpy as np
import sys
import subprocess as sub
#
class ComputeUSBDialog(QDialog):
    def __init__(self, iface):
        QDialog.__init__(self)
        self.iface = iface
        uic.loadUi(os.path.join( os.path.dirname(__file__), 'ui/ui_computeUnsatSoluteBalance.ui'), self)

        self.buttonBox.button(QDialogButtonBox.Ok).clicked.connect(self.runUSB)
        self.manageGui()
##
##
    def manageGui(self):

        self.cmbModelName.clear()
        self.cmbGrid.clear()
        self.cmbUsb.clear()
        #
        layerNameList = getVectorLayerNames()
        layerNameList.sort()

        #retrieve the name and path of the modeltable
        (modelNameList, pathList) =  getModelsInfoLists(layerNameList)
        self.cmbModelName.addItems(modelNameList)
        #
        grid_layers = []
        for nametemp in layerNameList:
            if '_grid' in nametemp:
                grid_layers.append(nametemp)
        self.cmbGrid.addItems(grid_layers)
        #
        usb_layers = []
        uzf_layers = []
        surf_layers = []
        for nametemp in layerNameList:
            if '_usb' in nametemp:
                usb_layers.append(nametemp)
        # uzf
            if '_uzf' in nametemp:
                uzf_layers.append(nametemp)
            if '_sml' in nametemp:
                surf_layers.append(nametemp)
        self.cmbUZF.addItems(uzf_layers)
        self.cmbSurf.addItems(surf_layers)
        self.cmbUsb.addItems(usb_layers)
#
    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 Executale files !!
                        Open Program Locations and fill in the right path
                        to your codes  '''
            QtGui.QMessageBox.information(None, 'Warning !!', message  )

##
    def runUSB(self):
        # retrieve program location
        dirdict = self.getProgramLocation()
        modlflowdir = dirdict['MF2005']

        # ------------ Load input  Flow Model  ------------
        layerNameList = getVectorLayerNames()
        #
        modelName = self.cmbModelName.currentText()

        # Retrieve number of stress periods from FlowModel
        (pathfile, nsp ) = getModelInfoByName(modelName)

        namefile = modelName + '.nam'
        ml = Modflow.load(namefile, version='mf2005', model_ws= pathfile)
        ml.exename = modlflowdir
        ml.external_fnames = []
        #os.remove(pathfile + '/' + namefile)
        #

        nrow, ncol, nlay, nper = ml.nrow_ncol_nlay_nper

        # Input for USB object
        # - initialize dictionaries
        cell_dict = {}
        cin_dict = {}
        mu_dict = {}
        for k in range(nper):
            cin_dict[k] = np.zeros(shape=(nrow,ncol))

        # Insert data in dictionaries
        usblayer = getVectorLayerByName(self.cmbUsb.currentText())
        for f in usblayer.getFeatures():
            cid = int(f['PKUID'])
            nr = int(f['ROW'])-1
            nc = int(f['COL'])-1
            cell_dict[cid] = [nr,nc]
            mu_dict[cid] = float(f['decay_cnst'])
            for k in range(nper):
                cin_dict[k][nr,nc] = float(f['init_conc_' + str(k +1)])


        # Get data to build again UZF
        uzflayer = getVectorLayerByName(self.cmbUZF.currentText())
        iuzfbound = freewat.createGrid_utils.get_param_array(uzflayer, fieldName = 'iuzfbnd')
        eps = freewat.createGrid_utils.get_param_array(uzflayer, fieldName = 'eps')
        thts = freewat.createGrid_utils.get_param_array(uzflayer, fieldName = 'thts')
        thti = freewat.createGrid_utils.get_param_array(uzflayer, fieldName = 'thti')

        finf_list = []
        pet_list = []
        extdp_list = []
        extwc_list = []
        for i in range(1,nper+1):
            ftemp = freewat.createGrid_utils.get_param_array(uzflayer, fieldName = 'finf_'+ str(i))
            finf_list.append(ftemp)
            pttemp = freewat.createGrid_utils.get_param_array(uzflayer, fieldName = 'pet_'+ str(i))
            pet_list.append(pttemp)
            dptemp = freewat.createGrid_utils.get_param_array(uzflayer, fieldName = 'extdp_'+ str(i))
            extdp_list.append(dptemp)
            wctemp = freewat.createGrid_utils.get_param_array(uzflayer, fieldName = 'extwc_'+ str(i))
            extwc_list.append(wctemp)
        #
        nuztopstr = self.cmbUzfOpt.currentText()
        if 'Top' in nuztopstr :
            nuztop = 1
        if 'specified' in nuztopstr :
            nuztop = 2
        if 'Active' in nuztopstr :
            nuztop = 3
        #
        if self.chkUzfSfr.isChecked():
            irunflg = 1
            surflayer = getVectorLayerByName(self.cmbSurf.currentText())
            irunbnd = freewat.createGrid_utils.get_param_array(surflayer, fieldName = 'irunbnd')
        else:
            irunflg = 0
            irunbnd = 0
        #
        if self.chkUzfEvt.isChecked():
            ietflg = 1
        else:
            ietflg = 0
        #
        # WRITE UZF GAGE


        # retrieve DIS package of Modflow model, to get perlen and nstep
        # disinput = ml.name + '.dis' #
        disinput = ml.model_ws + '/' + ml.name + '.dis'
        dis = ModflowDis.load(disinput, ml)

        # make self this info
        Nsteps = dis.nstp
        SPlength = dis.perlen
        # get cell dimension (we assume delr = delc)
        delr = dis.delr
        if len(delr) > 1:
            delr = max(dis.delr)

        # ---
        nuzgag = len(cell_dict)
        # write a gage file for each cell
        pathoutput = ml.model_ws + '/'
        row_col_iftunit_iuzopt = []
        icl = 1
        for g in cell_dict.keys():
            nr = cell_dict[g][0] + 1
            nc = cell_dict[g][1] + 1
            # define UNIT of file as 900 + cellid
            uncell = 900 + int(g)
            row_col_iftunit_iuzopt.append([nr, nc, uncell , 2  ])
            #
            #ml.external_fnames.append(ml.name + '.uzf'+ str(icl))
            #ml.external_units.append(uncell)
            icl += 1

        # FIXED Values: iuzfopt= 2, iuzfcb1=57, iuzfcb2=0, ntrail2=10, nsets=20
        # HERE I HAVE TO USE A CORRECTED VERSION OF UZF .... waiting for official fixing in Flopy repository
        # remove existing UZF
        ml.remove_package('UZF')
        uzf = ModflowUzf1_new(ml, nuztop=  nuztop, iuzfopt= 2, irunflg= irunflg, ietflg= ietflg, iuzfcb1=57, iuzfcb2=0, ntrail2=10, nsets=20, nuzgag= nuzgag, surfdep=0.2,
                             iuzfbnd= iuzfbound, irunbnd = irunbnd, eps = eps, thts = thts, thtr=0.15, thti = thti, finf= finf_list, pet= pet_list, extdp= extdp_list , extwc= extwc_list, row_col_iftunit_iuzopt = row_col_iftunit_iuzopt  )
        uzf.write_file()


        # ... and run again Modflow Model
        ml.write_input()
        try :
            ml.run_model()
        except:
            # if it does not work, do the following ...
            namefile = ml.name + '.nam'
            proc = sub.Popen([ml.exename, namefile], stdin = sub.PIPE, stdout = sub.PIPE, stderr = sub.PIPE, cwd= ml.model_ws)
            result = proc.communicate()[0]
            QtGui.QMessageBox.information(None, 'Information', result  )

        # Get the cell dimension


        # Create the USB object
        usb = ModflowUsb(ml, uzf, delr, cell_dict = cell_dict, cin_dict = cin_dict, mu_dict = mu_dict)
        # and run USB
        # usb.write_uzf_gages()
        usb.run_model(Nsteps,SPlength)

        cout = usb.cout_dict

        # Save output as MDO

        # Retrieve the information of the model and the name
        dbname = os.path.join(pathfile, modelName + '.sqlite')
        tableList = getTableNamesList(dbname)
        # Retrieve the name of the new layer from Text Box
        name = self.textEdit.text()
        layerName = name + "_usbout"

        if layerName in tableList:
            QMessageBox.warning(self, self.tr('Warning!'), self.tr('Table %s already exists!' % layerName))
            return
        # Initialize Default values for fields:
        name_fields = []
        type_fields = []
        default_fields = []

        name_fields.append('layer')
        type_fields.append(QVariant.Int)
        default_fields.append(1)

        for i in range(1,nper+1):
            name_fields.append('unsat_conc_' + str(i))
            type_fields.append(QVariant.Double)
            default_fields.append(1.0)

        # Create template for MDO for USB:
        gridLayer = getVectorLayerByName(self.cmbGrid.currentText())
        createMDO(gridLayer, dbname, layerName, name_fields, type_fields, default_fields)

        # And insert data:
        vl = getVectorLayerByName(layerName)

        fields = vl.dataProvider().fields()
        vl.startEditing()
        changingAttributesDict = { }

        for f in vl.getFeatures():
            fid = f.id()
            nr = int(f['ROW'])-1
            nc = int(f['COL'])-1
            attrs = { }
            for i in range(nsp):
                idx = fields.indexFromName('unsat_conc_' + str(i + 1))#retrieve field ID
                cval = float(cout[i][nr,nc])
                attrs[idx] = cval
            changingAttributesDict[fid] = attrs

        #changingAttributesDict = {1 :  {8: 0.0, 5: 0.0, 6: 0.0, 7: 0.0} }
        vl.dataProvider().changeAttributeValues(changingAttributesDict)
        vl.commitChanges()
        #
        #
        # for i in range(nsp):
        #     #if max(cout[i] ) > 0:
        #     idx = fields.indexFromName('unsat_conc_' + str(i + 1))#retrieve field
        #     for f in vl.getFeatures():
        #         attrs = f.attributes()
        #         for attr in attrs:
        #             feat = f.id()
        #             nr = f['ROW']-1
        #             nc = f['COL']-1
        #             cval = cout[i][nr,nc]
        #             vl.changeAttributeValue(feat, idx, cval)
        #     idval += 1
        #
        #
        #
        # for i in range(nsp):
        #     #if max(cout[i] ) > 0:
        #     idx = fields.indexFromName('unsat_conc_' + str(i + 1))#retrieve field
        #     for f in vl.getFeatures():
        #         attrs = f.attributes()
        #         for attr in attrs:
        #             feat = f.id()
        #             nr = f['ROW']-1
        #             nc = f['COL']-1
        #             cval = cout[i][nr,nc]
        #             vl.changeAttributeValue(feat, idx, cval)
        #     idval += 1




        #Close the dialog window after the execution of the algorithm
        QDialog.reject(self)
