# -*- 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.QtGui import *
from qgis.core import *
from qgis.gui import *
from qgis.utils import *
from PyQt4.QtCore import *

import os
from PyQt4 import QtGui, uic
import qgis.utils
import processing
import math # to calculate de angle
from freewat.createGrid_utils import transf # to transform model coordinates into geographic coordinates
import freewat.ftools_utils as ftools_utils # to extract point coordinates


##from modpath import modpathAdapter as adpater
#
import sys
import numpy as np
##import subprocess as sub

# load flopy and grid utils
from flopy.modflow import *

from flopy.modpath import *
from flopy.utils import *

import freewat.createGrid_utils
from freewat.freewat_utils import getVectorLayerByName, getVectorLayerNames, getModelsInfoLists, getModelInfoByName, deselectAll
from freewat.sqlite_utils import uploadQgisVectorLayer, checkIfTableExists
#
FORM_CLASS, _ = uic.loadUiType(os.path.join( os.path.dirname(__file__), 'ui/ui_createParticleTracking.ui'))
#


class createPartikelTracking(QDialog, FORM_CLASS):
    def __init__(self, iface):
        QDialog.__init__(self)

        # Since Freewat v.0.5 MODPATH is available ONLY with flopy 3.2.6
        from flopy import version as vs
        if vs.__version__ < '3.2.6':
            # messaggio
            message  = ''' You are running MODPATH 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.warning(None, 'Warning !!', message  )
            return

        self.iface = iface
        self.setupUi(self)
        self.modelpath = ''
        self.modelname = ''
        self.exePath = ''

        self.btnRun.clicked.connect(self.runModpath)
        self.btnPostProc.clicked.connect(self.postProcessing)
        self.btnLoad.clicked.connect(self.reloadLayers)
        self.btnApply.clicked.connect(self.updatePathStyle)

        #
        self.manageGui()
##
##
    def manageGui(self):
        self.cmbModelName.clear()
        self.layerNameList = getVectorLayerNames()

        self.layerNameList.sort()

        # Remark: here we assume the name of models table starts with "modeltable"
        (modelNameList, pathList) =  getModelsInfoLists(self.layerNameList)
        self.cmbModelName.addItems(modelNameList)

        # load MPSIM output file, if any
        self.reloadSimulation()

    def reloadLayers(self):
        path_list = []
        for name in self.layerNameList:
            if "_pathlines_lay_" in name:
                path_list.append(name)

        # insert layers name in combobox for plotting:
        self.cmbLayer.addItems(path_list)
    #
    def reloadSimulation(self):
        # get all the info necessary about working directory
        self.params = self.pass_params_to_modpath()
        self.modelName = self.params['ModelName']
        (self.modelPath, nsp ) = getModelInfoByName(self.modelName)
        modeltable = getVectorLayerByName('modeltable_' + self.modelName)
        for ft in modeltable.getFeatures():
            self.crs = ft['crs']
        # get layer name from LPF table
        lpftable = getVectorLayerByName("lpf_"+ self.modelName)
        dpLPF = lpftable.dataProvider()
        layNameList = []
        for ft in lpftable.getFeatures():
            layNameList.append(ft['name'])
        self.lay = layNameList[0]

    # --
    def pass_params_to_modpath(self):
        params = {
            'ModelName': self.cmbModelName.currentText(),
            'Package': self.cmbPackage.currentText(),
            'SimulationType': self.cmbSimType.currentText(),
            'ReferenceTime': self.txtRefTime.text(),
            'TrackingDirection': self.cmbTrackingDirection.currentText(),
            'OutputInterval': self.txtOutputTime.text(),
            'Retardation': self.chkRetardation.isChecked(),
            #'ParticleSource': self.cbParticleSource.currentText(),
            'ReleaseTime': self.txtReleaseTime.text(),
            #'SelectedFeatures': self.getSelectedFeatures()
        }
        return params
    # --
    def getProgramLocation(self):
        try:
            # Retrieve programs location
            layerNameList = self.layerNameList
            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 applyStyle(self, layer, fieldname):
        # Use the currently selected layer
        registry = QgsSymbolLayerV2Registry.instance()
        lineMeta = registry.symbolLayerMetadata("SimpleLine")
        markerMeta = registry.symbolLayerMetadata("MarkerLine")

        symbol = QgsSymbolV2.defaultSymbol(layer.geometryType())

        # Line layer
        lineLayer = lineMeta.createSymbolLayer({'width': '0.26', 'color': 'blue', 'offset': '0.0', 'penstyle': 'solid', 'use_custom_dash': '0', 'joinstyle': 'bevel', 'capstyle': 'square'})

        # Marker layer
        markerLayer = markerMeta.createSymbolLayer({'width': '0.26', 'color': 'blue', 'rotate': '1', 'placement': 'CentralPoint', 'offset': '0'})
        markerLayer.setPlacement(4)

        subSymbol = markerLayer.subSymbol()
        # Replace the default layer with our own SimpleMarker
        subSymbol.deleteSymbolLayer(0)
        triangle = registry.symbolLayerMetadata("FilledMarker").createSymbolLayer({'name': 'filled_arrowhead', 'color': 'blue', 'color_border': 'blue', 'offset': '0,0', 'size': '2.0', 'angle': '0'})
        subSymbol.appendSymbolLayer(triangle)

        # Replace the default layer with our two custom layers
        symbol.deleteSymbolLayer(0)
        symbol.appendSymbolLayer(lineLayer)
        symbol.appendSymbolLayer(markerLayer)

        label = QgsPalLayerSettings()
        label.readFromLayer(layer)
        label.enabled = True
        label.displayAll = True
        label.fieldName = fieldname
        label.placement= QgsPalLayerSettings.Line
        #label.setDataDefinedProperty(QgsPalLayerSettings.Size,True,True,'red','time')
        label.writeToLayer(layer)
        #labelingEngine = QgsPalLabeling()
        #iface.mapCanvas().mapRenderer().setLabelingEngine(labelingEngine)

        # Replace the renderer of the current layer
        renderer = QgsSingleSymbolRendererV2(symbol)
        layer.setRendererV2(renderer)
        iface.legendInterface().refreshLayerSymbology(layer)
        iface.mapCanvas().refresh()

    # --
    def runModpath(self):

        # ------------ Get input from GUI
        params = self.pass_params_to_modpath()
        modelName = params['ModelName']
        trackdir= params['TrackingDirection']
        if 'line' in params['SimulationType']:
            simtype = 'pathline'
        else:
            simtype = 'endpoint'

        packages= None
        if 'WEL' in params['Package']:
            packages = 'WEL'
        if 'RCH' in params['Package']:
            packages = 'RCH'

        # ------------Get Model's info
        # -- Load MODFLOW and MODPATH exe path from program Locations
        dirdict = self.getProgramLocation()
        mf_exe = dirdict['MF2005']
        mp_exe = dirdict['MODPATH']
        #  -- Retrieve path of Modflow Model
        (pathfile, nsp ) = getModelInfoByName(modelName)
        modeltable = getVectorLayerByName('modeltable_' + modelName)
        for ft in modeltable.getFeatures():
            crs = ft['crs']

        # -- Load Modflow model as FloPy object
        namefile = modelName + '.nam'
        ml = Modflow.load(namefile, version='mf2005', model_ws=pathfile)
        # -- get basic info
        nrow, ncol, nlay, nsp = ml.nrow_ncol_nlay_nper

        # -- get input of effective porosity from model layers
        # get layers name from LPF table
        lpftable = getVectorLayerByName("lpf_"+ modelName)
        dpLPF = lpftable.dataProvider()
        layNameList = []
        for ft in lpftable.getFeatures():
            layNameList.append(ft['name'])

        porosity = np.zeros(shape = (nlay, nrow, ncol))
        for i,lay in enumerate(layNameList):
            prstemp = freewat.createGrid_utils.get_param_array(getVectorLayerByName(lay), fieldName = 'NE')
            # update 3D array
            porosity[i, :, : ] = prstemp


        # -- Create MODPATH model object and run
        #

        mpname = modelName + 'mp'

        # Pass ipakcb to LPF
        #ml.lpf.ipakcb = 53

        mp = Modpath(mpname, exe_name=mp_exe, modflowmodel=ml, model_ws = ml.model_ws )
        mp.dis_file = modelName + '.dis'
        mp.write_name_file()
        mpbas = ModpathBas(mp, hnoflo=ml.bas6.hnoflo, hdry=ml.lpf.hdry,
                                 ibound=ml.bas6.ibound.array, prsity= porosity)


        # Reduce ROW and COL Count to 2, to avoid TOO SLOW execution
        sim = mp.create_mpsim(trackdir=trackdir, simtype=simtype,
                 ParticleColumnCount=2, ParticleRowCount=2,
                 packages= packages, start_time=float(params['ReferenceTime']))
        # If MODPATH is applied to RCH, select only cells with recharge flux not zero
        if packages == 'RCH':
            # rch flux array
            rch = ml.rch.rech
            updateFlag = True
            sim.group_region = []
            for kper in range(nsp):
                if updateFlag:
                    #inrch, flux_rech = rch.get_kper_entry(kper)
                    imin, jmin = nrow, ncol
                    imax, jmax = 0, 0
                    for i in range(nrow):
                        for j in range(ncol):
                            if rch[kper][i,j] != 0.0 :
                                if i < imin: imin = i
                                if j < jmin: jmin = j
                                if i > imax : imax = i
                                if j > jmax : jmax = j
                    #
                    sim.group_region.append([0,imin, jmin, 0, imax, jmax])
                    updateFlag = False
        #
        sim.option_flags[5] = 1
        # Writ input
        sim.write_file()
        mp.write_name_file()
        mp.write_input()
        # Run MODPATH
        mp.run_model()
        #Message
        msg = ''' MODPATH run successfully! Please, check results in your working directory.
              Open the Postprocessing ab if you want
              to run the post-processor.
              '''
        QMessageBox.information(self, self.tr('MODPATH Simulation executed!'), self.tr(msg))

        # update list of simulation files
        self.reloadSimulation()


    def postProcessing(self):
        # Get MODPATH output
        # load the PathLine data
        modelpth = self.modelPath
        mpName = self.modelName + 'mp'
        endpoint_file = mpName + '.mpend'
        endfile = os.path.join(modelpth, endpoint_file)
        endobj = EndpointFile(endfile)
        ept = endobj.get_alldata()
        # # load the pathline data if requested
        if 'line' in self.params['SimulationType']:
            pathline_file = mpName +  '.mppth'
            pthfile = os.path.join(modelpth, pathline_file)
            pthobj = PathlineFile(pthfile)
            plines = pthobj.get_alldata()


        vl = QgsVectorLayer("Point?crs=" + self.crs, "particles_points", "memory")
##        QgsMapLayerRegistry.instance().addMapLayer(vl)

        # get grid origin to set X,Y as geographic instead of model coordinates
        lay = self.lay
        cell = [feature for feature in getVectorLayerByName(lay).getFeatures()] # select cells (features) within the lay
        cellgeom = cell[0].geometry().exportToWkt() # write WKT of the first cell (feature) of layer
        cellcoord = QgsGeometry.fromWkt(cellgeom) # extract WKT coodinates
        p0, p1, p2, p3, p4 = ftools_utils.extractPoints(cellcoord) # write cell (feature) vertices coordinates

        xcorner = p3.x() # associate the model corner (x origin) to the first vertex of the first cell
        ycorner = p3.y() # associate the model corner (y origin) to the first vertex of the first cell

        angle = math.atan((p0.x()-p3.x())/(p0.y()-p3.y())) # calculate the rotation of the grid

        # Create attributes
        pr = vl.dataProvider()
        res = pr.addAttributes( [ QgsField("particleid", QVariant.Int), \
                                  QgsField("layer", QVariant.Int), QgsField("time", QVariant.Double)] )

        vl.updateFields()

        # List of modflow layers:
        layers = []
        for p in plines:
            for i,pt in enumerate(p):
                feat = QgsFeature(vl.fields())
                feat.setAttributes([int(pt['id']),int(pt['k']),float(pt['time'])])
                xl, yl = transf(xcorner, ycorner, 0, 0, angle, pt['x'], pt['y']) # define/convert the coordinates for each point into geographic with rotation based on "transf" function from createGrid_utils.py
                feat.setGeometry(QgsGeometry.fromPoint(QgsPoint(xl, yl))) # create the points from geographic coordinates
                (res, outFeats) = pr.addFeatures( [ feat ] )
                if str(pt['k']) not in layers:
                    layers.append(str(pt['k']))

        vl.commitChanges()

        path_list = []
        # Selection per layer and save in DB, and load to canvas
        for l in layers:
            exp = QgsExpression( "\"layer\" =  %s"%l )
            ids = [i.id() for i in vl.getFeatures(QgsFeatureRequest(exp))]
            vl.setSelectedFeatures(ids)
            # Upload the vlayer into DB SQlite
            mpSimName = self.txtName.text()
            newName = mpSimName + "_particles_lay_" + str(int(l) +1)
            dbName = os.path.join(modelpth, self.modelName + '.sqlite')
            uploadQgisVectorLayer(dbName, vl, newName, selected= True)
            # Retrieve the Spatialite layer and add it to mapp
            uri = QgsDataSourceURI()
            uri.setDatabase(dbName)
            schema = ''
            table = newName
            geom_column = "Geometry"
            uri.setDataSource(schema, table, geom_column)
            display_name = table

            wlayer = QgsVectorLayer(uri.uri(), display_name, 'spatialite')

            QgsMapLayerRegistry.instance().addMapLayer(wlayer)

            features = wlayer.getFeatures()
            Points= {}
            for feature in features:
                geom = feature.geometry().centroid().asPoint()
                timePoint = feature['time']
                try:
                    Points[feature['particleid']].append([timePoint,geom])
                except:
                    Points[feature['particleid']] = [[timePoint,geom]]
            #Points.sort()


            # Create the lines layer
            # # New memory line layer
            ll = QgsVectorLayer("LineString?crs=" + self.crs, "pathlines", "memory")
            pr = ll.dataProvider()
            res = pr.addAttributes( [ QgsField("particleid", QVariant.Int), QgsField("time", QVariant.Double)] )
            ll.updateFields()
##            QgsMapLayerRegistry.instance().addMapLayer(ll)

            for j,p in enumerate(Points.keys()):
                for i in range(1, len(Points[p])):
                    line = QgsFeature(ll.fields())
                    line.setGeometry(QgsGeometry.fromPolyline([Points[p][i-1][1], Points[p][i][1]]))
                    line.setAttributes([int(p),float(Points[p][i][0])])
                    #line.setAttributes([122,1.0])
                    #print i,float(Points[p][i][0]), Points[p][i-1][1], Points[p][i-1][1], Points[p][i][1]
                    (res, outFeats) = ll.dataProvider().addFeatures( [ line ] )

            ll.commitChanges()

            # last line:
##            line = QgsFeature(ll.fields())
##            line.setGeometry(QgsGeometry.fromPolyline([Points[len(Points)-1][1], Points[len(Points)][1]]))
##            line.setAttributes([len(Points),float(Points[i][0]) ])
##            (res, outFeats) = ll.dataProvider().addFeatures( [ line ] )

            #QgsMapLayerRegistry.instance().addMapLayer(ll)
            # Upload the line into DB SQlite
            newName = mpSimName + "_pathlines_lay_" + str(int(l) +1)
            #
            path_list.append(newName)
            uploadQgisVectorLayer(dbName, ll, newName)
            # Retrieve the Spatialite layer and add it to mapp
            uri = QgsDataSourceURI()
            uri.setDatabase(dbName)
            schema = ''
            table = newName
            geom_column = "Geometry"
            uri.setDataSource(schema, table, geom_column)
            display_name = table

            llayer = QgsVectorLayer(uri.uri(), display_name, 'spatialite')

            QgsMapLayerRegistry.instance().addMapLayer(llayer)



        #
        del vl, ll
        deselectAll()


    def updatePathStyle(self):
        vl = getVectorLayerByName(self.cmbLayer.currentText())
        self.applyStyle(vl, 'time')
