# -*- 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, 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_viewOutput.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.reloadFields)
        # self.cmbStressPeriod.currentIndexChanged.connect(self.reloadSteps)
        # 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 reloadFields(self):
        # Insert list of transport model(s) according with selected flow model:
        self.cmbTransportName.clear()

        layerNameList = getVectorLayerNames()
        try:
            list_of_models = getTransportModelsByName(self.cmbModelName.currentText())
        except:
            list_of_models = []

        self.cmbTransportName.addItems(list_of_models)



    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')

        # Model for Selecting more than one SP
        # Insert in ComboBox
        model = QtGui.QStandardItemModel(nsp, 1)# nsp rows, 1 col
        it0 = QtGui.QStandardItem('-- Select -- ')
        for i, sp  in enumerate(SPitems):
            item = QtGui.QStandardItem(str(sp))
            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.modelSP = model
        self.cmbStressPeriodList.setModel(model)
        self.cmbStressPeriodList.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

        # --- Post processing
        # Note you may have to set compiler type.
        # 'l' suits for OS X
        # mread = ModflowHdsRead(ml,compiler='l')
        #
        #
        if self.chkFlowModel.isChecked():
            # USINg HeadFile object
            model = modelName
            hdobj = HeadFile(pathfile+'/'+ model +'.hds', text = 'head', precision='single')


            kstpkper = hdobj.get_kstpkper()

##            htot = hdobj.get_data(kstpkper = (kstp, kper))
        if self.chkTransportModel.isChecked():
            model = self.cmbTransportName.currentText()
            nspec = int( self.txtSpecies.text())
            ucnobj = UcnFile(pathfile+'/'+ 'MT3D00%i.UCN'%nspec, precision='single')
            times = ucnobj.get_times()
            kstpkper = ucnobj.get_kstpkper()


        if self.groupSingleSP.isChecked():
            kper = int(self.cmbStressPeriod.currentText()) - 1
            kstp = int(self.txtTimeStep.text()) - 1

            if self.chkFlowModel.isChecked():
                harray = hdobj.get_data(kstpkper = (kstp, kper))
            if self.chkTransportModel.isChecked():
                transpTime = 0
                for i, k in enumerate(kstpkper):
                    if k[1] == kper and k[0] == kstp :
                        transpTime = times[i]
                harray = ucnobj.get_data(totim=transpTime)

            if vector_output:
                fileName = model + '_sp_' + str(kper +1 )+ '_ts_' + str(kstp + 1 )
                self.writeGridOutput(fileName, layerGrid, harray)

            else:
                for ilay in range(0,nlay):
                    # write raster
                    fileName = pathfile+'/'+ model +'_lay_' + str(ilay +1) + '_sp_' + str(kper +1 )+ '_ts_' + str(kstp + 1 ) + '.asc'
                    if self.chkTransportModel.isChecked():
                        comp = self.txtSpecies.text()
                        fileName = pathfile+'/'+ model + 'species_'+ comp +'_lay_' + str(ilay +1) + '_sp_' + str(kper +1 )+ '_ts_' + str(kstp + 1 ) + '.asc'
                    self.writeRasterOutput(fileName, layerGrid, ilay, ncol, nrow, delc, modelName, harray )

        # Option of Multi-Stress Period
        if self.groupMultiSP.isChecked():

            # Retrieve list of selected StressPeriods from SP combo
            splist = []
            i = 1
            while self.modelSP.item(i):
                if self.modelSP.item(i).checkState() == 2:
                    valore = self.modelSP.item(i).data(0)
                    splist.append( int(valore) - 1 )
                i += 1

            for j, kk in enumerate(kstpkper):

                if kk[1] in splist:
                    kper = kk[1]
                    kstp = kk[0]
                    if self.chkFlowModel.isChecked():
                        harray = hdobj.get_data(kstpkper = (kstp, kper))
                    if self.chkTransportModel.isChecked():
                        harray = ucnobj.get_data(kstpkper = (kstp, kper))

                    if vector_output:
                        fileName = model + '_sp_' + str(kper +1 )+ '_ts_' + str(kstp + 1 )
                        self.writeGridOutput(fileName, layerGrid, harray)

                    else:
                        for ilay in range(0,nlay):
                            # write raster
                            fileName = pathfile+'/'+ model +'_lay_' + str(ilay +1) + '_sp_' + str(kper + 1)+ '_ts_' + str(kstp + 1) + '.asc'
                            if self.chkTransportModel.isChecked():
                                comp = self.txtSpecies.text()
                                fileName = pathfile+'/'+ model + 'species_'+ comp +'_lay_' + str(ilay +1) + '_sp_' + str(kper + 1)+ '_ts_' + str(kstp + 1) + '.asc'

                            self.writeRasterOutput(fileName, layerGrid, ilay, ncol, nrow, delc, modelName, harray )

        # Close the dialog
        QDialog.reject(self)
