# -*- 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



##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.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)

    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 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
        #
        sim.write_file()

        mp.write_name_file()

        mp.write_input()
        mp.run_model()

        mp.write_input()

        mp.run_model()

        # Get MODPATH output
        # load the endpoint data
        modelpth = mp.model_ws
        endfile = os.path.join(modelpth, mp.sim.endpoint_file)
        endobj = EndpointFile(endfile)
        ept = endobj.get_alldata()
        # # load the pathline data if requested
        if 'line' in params['SimulationType']:
            pthfile = os.path.join(modelpth, mp.sim.pathline_file)
            pthobj = PathlineFile(pthfile)
            plines = pthobj.get_alldata()


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

        # get grid origin to set X,Y as geographic instead of model coordinates
        xcorner, ycorner = getVectorLayerByName(lay).extent().xMinimum(), getVectorLayerByName(lay).extent().yMinimum()

        # 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'])])
                feat.setGeometry(QgsGeometry.fromPoint(QgsPoint(pt['x'] + xcorner, pt['y'] + ycorner)))
                (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(pathfile, 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=" + 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)

            dbName = os.path.join(pathfile, modelName + '.sqlite')
            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')

