# -*- 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
#
from freewat.freewat_utils import getVectorLayerByName, getVectorLayerNames, getModelsInfoLists, getModelInfoByName, getTransportModelsByName
# load flopy and createGrid
from flopy.modflow import *
from flopy.utils import *
import freewat.createGrid_utils
#
FORM_CLASS, _ = uic.loadUiType(os.path.join( os.path.dirname(__file__), 'ui/ui_viewCrossSection.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.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):
        # Begin reload Time
        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.cmbStressPeriod.addItems(SPitems)
        self.txtTimeStep.setText('1')
##
    def viewOutput(self, vector_output = False):
        # ------------ 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']
                    l_units = str(f['length_unit'])
        # 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.
        topLayer = getVectorLayerByName(layNameList[0])
        bottomLayer = getVectorLayerByName(layNameList[-1])


        # 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
        #
        if self.chkFlowModel.isChecked():
            # USINg HeadFile object
            model = modelName
            hdobj = HeadFile(pathfile+'/'+ model +'.hds', text = 'head', precision='single')

            kstpkper = hdobj.get_kstpkper()

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

        # get stress period and time step from GUI
        kper = int(self.cmbStressPeriod.currentText()) - 1
        kstp = int(self.txtTimeStep.text()) - 1

        # geta data: harray is an 3D np.array shpaed (nlay,nrow,ncol)
        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(kstpkper = (kstp, kper))

        # Vertical exageration from GUI
        vex = float(self.txtVex.text())
        fixed = None
        txtFixed = ''
        notFixed = ''
        if 'col' in self.cmbCol.currentText().lower():
            txtFixed = 'COL'
            notFixed = 'ROW'
            fixed = int(self.spinBoxCol.text()) - 1
            z = harray[:,:,fixed]
            L = nrow  #Ly
            # Selection for top and bottom
            exp = QgsExpression( "\"col\" =  %s"%(fixed + 1) )
        else:
            txtFixed = 'ROW'
            fixed = int(self.spinBoxCol.text()) - 1
            notFixed = 'COL'
            z = harray[:,fixed,:]
            L = ncol#Lx
            # Selection for top and bottom
            exp = QgsExpression( "\"row\" =  %s"%(fixed + 1) )

        # Get TOP and Bottom
        toplist = []
        botlist = []
        for f in topLayer.getFeatures(QgsFeatureRequest(exp)):
            toplist.append(f['TOP'])
        for f in bottomLayer.getFeatures(QgsFeatureRequest(exp)):
            botlist.append(f['BOTTOM'])
        top = max(toplist)
        bottom = min(botlist)

        thickness = vex*(top-bottom)


        ##x = np.linspace(0, L, size(harray[0,0,:]) )
        ##y = np.linspace(0, vex,size(harray[:,0,0]) )
        ##X, Y = np.meshgrid(x, y)

        fig, ax = plt.subplots()
        cax = ax.imshow(z,interpolation='spline36', cmap = 'jet',  extent = (0,L,0,thickness))

        mezzo = (z.min()+z.max())*0.5
        cbar = fig.colorbar(cax, ticks=[z.min(), mezzo/2.0, mezzo, 3*mezzo/2.0, z.max()])
        cbar.ax.set_yticklabels(['%f'%z.min(), '%f'%(mezzo/2.0), '%f'%mezzo, '%f'%(3*(mezzo/2.0)), '%f'%z.max()])
        ax.set_title('Data at SP =%i, TS =%i, at %s = %i '%(kper + 1, kstp + 1, txtFixed, fixed + 1))
        ax.set_ylabel('Model thickness (%s)'%l_units)
        ax.set_xlabel(notFixed)

        plt.show()

        # Close the dialog
        QDialog.reject(self)


