# -*- 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 freewat.freewat_utils import getVectorLayerByName, getVectorLayerNames, getModelsInfoLists, getModelInfoByName, getModelNlayByName, getTransportModelsByName, ComboStyledItemDelegate
# load flopy and createGrid
from flopy.modflow import *
from flopy.utils import *
import freewat.createGrid_utils
from freewat.mdoCreate_utils import createMDOLarge, createMDO
#
FORM_CLASS, _ = uic.loadUiType(os.path.join( os.path.dirname(__file__), 'ui/ui_viewSWIOutput.ui') )
#
class viewOutputDialog(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.viewOutput)

        self.cmbModelName.currentIndexChanged.connect(self.reloadTime)
        self.cmbModelName.currentIndexChanged.connect(self.reloadLayers)
        # self.cancelButton.clicked.connect(self.stopProcessing)
        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 table
        isok = 0
        for mName in layerNameList:
            if mName == 'timetable_' + self.cmbModelName.currentText():
                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()
                nsp = 0
                #tslist = []
                for f in ftit:
                    nsp = nsp + 1
                    #tslist.append(f['time_steps'])

        SPitems = ['%s' % (i + 1) for i in range(nsp)]

        #self.listStressPeriod.addItems(SPitems)
        self.cmbStressPeriod.addItems(SPitems)
        #self.cmbStressPeriod.currentIndexChanged.connect(self.reloadSteps)
        self.txtTimeStep.setText('1')

    def reloadLayers(self):
        # Model for Selecting more than one Layer
        self.nlay = getModelNlayByName(str(self.cmbModelName.currentText()))
        LAYitems = ['%s' % (i + 1) for i in range(self.nlay)]
        # Insert in ComboBox
        model = QtGui.QStandardItemModel(self.nlay, 1)# nsp rows, 1 col
        it0 = QtGui.QStandardItem('-- Select -- ')
        for i, ll  in enumerate(LAYitems):
            item = QtGui.QStandardItem(str(ll))
            item.setFlags(Qt.ItemIsUserCheckable | Qt.ItemIsEnabled)
            if i == 0:
                item.setData(Qt.Checked, Qt.CheckStateRole)
            else:
                item.setData(Qt.Unchecked, Qt.CheckStateRole)
            model.setItem(i+1, 0, item)
        model.setItem(0,it0)
        self.modelLAY = model
        self.cmbLayersList.setModel(model)
        self.cmbLayersList.setItemDelegate(ComboStyledItemDelegate())

##
    def writeRasterOutput(self, fileName, layerGrid, ilay, ncol, nrow, delc, modelName, harray ):
        rstfile = open(fileName, 'w')
        # Data for this layer
        h11 = harray[ilay]
        # Header
        extent = layerGrid.extent()
        rstfile.write('NCOLS  %i \n'%ncol)
        rstfile.write('NROWS  %i \n'%nrow)
        rstfile.write('XLLCORNER  %f \n'%extent.xMinimum())
        rstfile.write('YLLCORNER  %f \n'%extent.yMinimum())
        rstfile.write('CELLSIZE   %f  \n'%delc)
        rstfile.write('NODATA_VALUE   -9999.0 \n')

        # Array of selected Head (kper, kstp)
        # rstfile.write('NODATA_VALUE   -9999.0 \n')

        for j in range(0, len(h11)):
            for k in range(0, len(h11[j])):
                rstfile.write(str(h11[j][k])+'  ')
            rstfile.write('\n')

        rstfile.close()

        fileInfo = QFileInfo(fileName)
        baseName = fileInfo.baseName()
        rlayer = QgsRasterLayer(fileName, baseName)
        QgsMapLayerRegistry.instance().addMapLayer(rlayer)
##
    ## -- Create new grid with results inside:
    def writeGridOutput(self, layerName, layerGrid,  h3darray ):
        dbName = self.pathfile + '/' + self.modelName + '.sqlite'

        name_fields = [] #['BORDER' ,'ACTIVE','TOP', 'BOTTOM', 'THICKNESS', 'STRT', 'KX', 'KY', 'KZ', 'SS', 'SY', 'NT', 'NE', 'WETDRY']
        type_fields = [] #[QVariant.Int , QVariant.Int, QVariant.Double,  QVariant.Double , QVariant.Double, QVariant.Double, QVariant.Double, QVariant.Double, QVariant.Double, QVariant.Double, QVariant.Double, QVariant.Double, QVariant.Double, QVariant.Double ]
        default_fields = [] #[0, 1, top, bottom, (top - bottom ), 1.0, 0.001, 0.001, 0.0001, 0.001, 0.1, 1, 1, -0.01]

        for i in range(1,self.nlay+1):
            name_fields.append('lay_' + str(i) )
            type_fields.append(QVariant.Double)
            default_fields.append(10.0)

##        try:
##            createMDOLarge(layerGrid,dbName,layerName, name_fields, type_fields, default_fields)
##        except:
        createMDO(layerGrid,dbName,layerName, name_fields, type_fields, default_fields)

        # Ora ha creato una copia del model layer, con aggiunti i nuovi attributi e valore di default.
        # devo cancellare quelli che stanno fra 'col' e 'lay_1'

        # Retrieve the Spatialite layer and add it to mapp
        uri = QgsDataSourceURI()
        uri.setDatabase(dbName)
        schema = ''
        table = layerName
        geom_column = "Geometry"
        uri.setDataSource(schema, table, geom_column)
        display_name = table

        # Get the DB layer as a Qgs layer object
        wlayer = QgsVectorLayer(uri.uri(), display_name, 'spatialite')

        # --- delete fields beetween COL and lay_1 (the ones got from model_layer MDO)
        pr = wlayer.dataProvider()
        fld = pr.fields()
        wlayer.startEditing()
        idcol = fld.indexFromName('COL')
        idlay = fld.indexFromName('lay_1')
        # If package = RIV or DRN, take ROW, COl, layer
        idx = range(idcol+1, idlay-1)

        # delete other fields
        wlayer.deleteAttributes(idx)
        wlayer.updateFields()
        wlayer.commitChanges()

        # Insert values of solution:
        request = QgsFeatureRequest().setFlags(QgsFeatureRequest.NoGeometry)
        wlayer.startEditing()
        changedAttributesDict = {}
        for feature in wlayer.getFeatures(request):
            changedAttributes = {}
            nr = int(feature['ROW'])
            nc = int(feature['COL'])

            for j in range(1,self.nlay+1):
                fieldIdx = pr.fieldNameIndex('lay_%i'%j)

                h11 = h3darray[j-1]
                changedAttributes[fieldIdx] = float(h11[nr-1][nc-1])

            changedAttributesDict[feature.id()] = changedAttributes

        pr.changeAttributeValues(changedAttributesDict)
        wlayer.commitChanges()
##
    def viewOutput(self, vector_output = False):
        # Choose Vector option, if requested:
        if self.radioVector.isChecked():
            vector_output = True
        # ------------ Load input data  ------------
        modelName = self.cmbModelName.currentText()
        self.modelName = modelName
        #

        # Remark: pathfile from modelltable
        layerNameList = getVectorLayerNames()

        for mName in layerNameList:
            if mName == 'modeltable_'+ modelName:
                modelNameTable = getVectorLayerByName(mName)
                for f in modelNameTable.getFeatures():
                    pathfile = f['working_dir']

        # Retrieve LPF table, and from there model layers and data needed for LPF
            if mName ==  "lpf_"+ modelName:
                lpftable = getVectorLayerByName(mName)
                # Number of layers
                nlay = 0
                # Get layers name from LPF table
                dpLPF = lpftable.dataProvider()

                # Create lists of layers properties
                layNameList = []

                for ft in lpftable.getFeatures():
                    attrs = ft.attributes()
                    layNameList.append(attrs[0])
                    nlay = nlay + 1

        # render pathfile as self
        self.pathfile = pathfile

        # Retrieve the basis of the grid for Extent etc.
        layerGrid = getVectorLayerByName(layNameList[0])
        # Spatial discretization:
        # Number of rows (along width)
        nrow, ncol = freewat.createGrid_utils.get_rgrid_nrow_ncol(getVectorLayerByName(layNameList[0]))
        # delc, delrow
        delr, delc = freewat.createGrid_utils.get_rgrid_delr_delc(getVectorLayerByName(layNameList[0]))

        # render nlay, nrow, ncol as self
        self.nlay, self.nrow , self.ncol = nlay, nrow, ncol

        # Get desired layers from combo (transform to 0-based):
        laySelected = []
        i = 1
        while self.modelLAY.item(i):
            if self.modelLAY.item(i).checkState() == 2:
                valore = self.modelLAY.item(i).data(0)
                laySelected.append( int(valore) - 1 )
            i += 1

        # --- Post processing
        # Note you may have to set compiler type.
        # 'l' suits for OS X
        # mread = ModflowHdsRead(ml,compiler='l')
        #
        # Get zeta surface array from CellBudget
        ## # Get Zeta Surface from ZTA file
        #
        model = modelName
        myfile = os.path.join(pathfile, modelName +'.zta')
        # check if SWI output really exists:
        if not os.path.isfile(myfile):
            # error message and return
            msg = 'There is any SWI2 output file in \n your working directory!!\nCheck and try again!'
            QtGui.QMessageBox.warning(None, 'Error!', msg)
            return
        # use CellBudget object from flopy.utils
        zfile = CellBudgetFile(myfile)
        kper = int(self.cmbStressPeriod.currentText()) - 1
        kstp = int(self.txtTimeStep.text()) - 1
        # At a given (time step, stress period) - zero-based!!
        kstpkper = (kstp, kper)
        # this is the list of zeta (len is NSRF):
        zall = zfile.get_data(kstpkper=kstpkper, text='ZETASRF  1')
        # but here we assume NSFR = 1, so:
        zeta = zall[0]
        # zeta is now collecting 1 surface for each layer
        # zeta is the aray shaped (nlay,nrow,ncol) at tstep = kstpkper[0], stress period = kstpkper[1]
        #
        # now, manage layer(s) we want to visualize
        if vector_output:
            fileName = model + '_interface_sp_' + str(kper +1 )+ '_ts_' + str(kstp + 1 )
            self.writeGridOutput(fileName, layerGrid, zeta)

        else:
            # recall: laySelected is 0-based
            for ilay in laySelected:
                # write raster
                fileName = os.path.join(pathfile, model +'_interface_lay_' + str(ilay+1) + '_sp_' + str(kper+1)+ '_ts_' + str(kstp+1) + '.asc')
                self.writeRasterOutput(fileName, layerGrid, ilay , ncol, nrow, delc, modelName, zeta )

        # Close the dialog
        QDialog.reject(self)
