# -*- 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 pyspatialite import dbapi2 as sqlite3
from freewat.freewat_utils import getVectorLayerNames, getModelsInfoLists, getModelInfoByName, getModelNspecByName, getTransportModelsByName
from freewat.sqlite_utils import getTableNamesList

#
#Load the interface
FORM_CLASS, _ = uic.loadUiType(os.path.join( os.path.dirname(__file__), 'ui/ui_createLKTLayer.ui') )
#
class CreateLKTLayerDialog(QDialog, FORM_CLASS):
    def __init__(self, iface):
        QDialog.__init__(self)
        self.iface = iface
        self.setupUi(self)
        # link the flow model with related transport model(s):
        #If the content of the cmbFlowName combo box containing the name of the flow model changes,
        #then the reloadFields function is recalled
        self.cmbFlowName.currentIndexChanged.connect(self.reloadFields)
        #The label of the 'OK' button is changed to 'Run'
        self.buttonBox.button(QDialogButtonBox.Ok).setText("Run")
        #If the 'Run' button is clicked, the createLKTTables function below is recalled
        self.buttonBox.button(QDialogButtonBox.Ok).clicked.connect(self.createLKTTables)
        #If the 'Cancel' button in the buttonBox is clicked, the Dialog closes
        self.buttonBox.button(QDialogButtonBox.Cancel).clicked.connect(self.reject)
        #Recall the manageGui function
        self.manageGui()
##
    def manageGui(self):
        #The cmbFlowName combo box containing the name of the flow model is cleared
        self.cmbFlowName.clear()
        #The list layerNameList is filled with the names of all the vector layers available
        #in the Legend (getVectorLayerNames is a function of the freewat_utils.py module;
        #such function returns a list)
        layerNameList = getVectorLayerNames()
        #The list layerNameList is sorted
        layerNameList.sort()

        # Retrieve modelname and pathfile List
        #The tuple (modelNameList, pathList) is made of two lists:
        # -the list modelNameList is filled with the names of all the flow models available
        # -the list pathList is filled with paths of working directories where the flow models DBs are saved
        #(getModelsInfoLists is a function of the freewat_utils.py module; such function
        #takes the list layerNameList as input, looks for layers whose names start with
        #modeltable and reads the content of the name and working_dir fields in the
        #modeltables found; it returns a tuple)
        (modelNameList, pathList) =  getModelsInfoLists(layerNameList)

        #The cmbFlowName combo box is filled with the elements of the list modelNameList
        self.cmbFlowName.addItems(modelNameList)

    def reloadFields(self):
        # Insert list of transport model(s) according with selected flow model:
        #The cmbTransportName combo box containing the name of the transport model is cleared
        self.cmbTransportName.clear()
        #The list layerNameList is filled with the names of all the vector layers available
        #in the Legend (getVectorLayerNames is a function of the freewat_utils.py module;
        #such function returns a list)
        layerNameList = getVectorLayerNames()

        #The list list_of_models is filled with the names of all the transport models available
        #(getTransportModelsByName is a function of the freewat_utils.py module; such function
        #reads the content of the name field in the transport_flow_model table (where flow_model
        #is the name of the flow model read from the cmbFlowName combo box) and returns a list)
        list_of_models = getTransportModelsByName(self.cmbFlowName.currentText())
        #The cmbTransportName combo box is filled with the elements of the list list_of_models
        self.cmbTransportName.addItems(list_of_models)

    def createLKTTables(self):

        self.progressBar_2.setMinimum(0)
        self.progressBar_2.setMaximum(0)
        self.progressBar_2.setValue(0)
        QApplication.processEvents()

        #modelName is the name of the flow model read from the cmbFlowName combo box
        modelName = self.cmbFlowName.currentText()
        #transportName is the name of the transport model read from the cmbTransportName combo box
        transportName = self.cmbTransportName.currentText()

        # Retrieve the path of the working directory and the number of stress periods.
        #The tuple (pathfile, nsp ) is made of two objects:
        # -pathfile is the path of the working directiry for a specific flow model
        # -nsp is the number of stress period for a specific flow model
        #(getModelInfoByName is a function of the freewat_utils.py module; such function
        #recalls the getModelsInfoLists function to retrieve the names of all the flow
        #models available in the legend and the corresponding paths, then it retrieves
        #the path of the specific flow model selected in the cmbFlowName combo box and
        #iterates over the rows of the timetable_modelName table to get the number of
        #stress periods; it returns a tuple)
        (pathfile, nsp ) = getModelInfoByName(modelName)
        # Retrieve the number of species, the number of mobile species and the mass units.
        #The tuple (nspec, mob, mass) is made of three objects:
        # -nspec is the number of species simulated for a specific transport model
        # -mob is the number of mobile species simulated for a specific transport model
        # -mass is the mass unit used for a specific transport model
        #(getModelNspecByName is a function of the freewat_utils.py module; such function
        #takes the name of the flow model (modelName, read from the cmbFlowName combo box)
        #and the name of a transport model (transportName, read from the cmbTransportName
        #combo box) as inputs, iterates over the rows of the species_transportName table
        #to get the number of species simulated, reads the content of the is_mobile field
        #and counts only the cells whose content is 'yes', and reads the content of the
        #unit field in the transport_modelName table to get the mass unit; it returns a tuple)
        (nspec, mob, mass) = getModelNspecByName(modelName, transportName)

        #Retrieve the DB name and path
        dbname = pathfile + '/' + modelName + ".sqlite"

        #The list tableList is filled with the names of all the tables available
        #in the DB (getTableNamesList is a function of the sqlite_utils.py module;
        #such function returns a list)
        tableList = getTableNamesList(dbname)

        #Connect SpatiaLite DB
        con = sqlite3.connect(dbname)
        cur = con.cursor()


        ### Retrieve the number of lakes in the flow model (this comes from the number of rows in the ###
        ### table lak_flow_model_name, which is in the model DB, but it is not loaded in the Legend)  ###

        #lak_table is the name of the table lak_flow_model_name which is going to be queried
        lak_table = 'lak_' + modelName

        #If the lak_table is not in the DB...
        if (lak_table not in tableList):
            #...a warning message appears and nothing more is done
            QMessageBox.warning(self, self.tr('Warning!'), self.tr('Table %s does not exists!\nPlease create some lakes in your flow model!' % lak_table))
            return

        #The number of rows in the table lak_table is counted (this number is stored in the variable nlakes)
        cur.execute('SELECT COUNT(*) FROM {}'.format(lak_table))
        count = cur.fetchall()
        nlakes = format(count[0][0])


        ### Create the table lake_starting_conc_ + transportName ###

        #table1 is the name to be assigned to the table lake_starting_conc_ + transportName
        table1 = 'lake_starting_conc_' + transportName

        #column_names is the list of columns headings. It contains at the beginning
        #the heading of the first column only ('lake_ID')
        column_names=['lake_ID']

        #If the table1 table is not in the DB...
        if (table1 not in tableList):
            #...then a new table renamed table1 is created.
            #Such table contains initially the field lake_ID (integer)
            SQLstring = 'CREATE TABLE "%s" ("lake_ID" integer );'%table1
            cur.execute(SQLstring)
            #According to the number of species in the transport model, fields
            #renamed coldlak_spec_n (float) are added (where n in the n-th species)
            #and the list column_names is updated
            for n in range(nspec):
                fieldname='coldlak_spec_' + str(n+1)
                column_names.append(fieldname)
                addColumn = 'ALTER TABLE "%s" ADD COLUMN "%s" float;'%(table1, fieldname)
                cur.execute(addColumn)
        #If the table1 table is in the DB...
        elif table1 in tableList:
            #...a warning message appears and nothing more is done
            QMessageBox.warning(self, self.tr('Warning!'), self.tr('Table %s already exists!' % table1))
            return

        #cell_values is the list of values at each cell of each row of the table lake_starting_conc_ + transportName
        cell_values=[]

        #row is a new dictionary. Its content will be like:
        #row={'lake_ID':1,'coldlak_spec_1':0.0,'coldlak_spec_2':0.0,...,'coldlak_spec_n':0.0}
        #The keys of this dictionary will be the elements of the list column_names,
        #the values of this dictionary will be the elements of the list cell_values
        row={}

        #i iterates over the number of lakes (which also corresponds to the
        #number of rows which will be added in the table lake_starting_conc_ + transportName)
        for i in range(int(nlakes)):
            #The list cell_values is updated with the lake ID value
            cell_values.append(int(i+1))
            #j iterates over the number of species
            for j in range(nspec):
                #The list cell_values is updated with the value 0.0
                #(0.0 is the starting concentration for each species)
                cell_values.append(0.0)
            #The built-in function enumerate allows to iterate over the list column_names.
            #The outputs of this function are:
            # - index is the index ID of each element of the list column_names
            # - key is the value of each element of the list column_names
            for index, key in enumerate(column_names):
                #The dictionary row is filled
                row[key]=cell_values[index]
            #A new row is added into the table lake_starting_conc_ + transportName
            #(its content is nothing but the content of the dictionary row)
            columns = ', '.join(map(str, row.keys()))
            values = ', '.join(map(repr, row.values()))
            cur.execute('INSERT INTO "{0}" ({1}) VALUES ({2})'.format(table1, columns, values))
            #The list cell_values and the dictionary row are emptied
            cell_values=[]
            row={}


        ### Create the table lake_table_bc_ + transportName ###

        #table2 is the name to be assigned to the table lake_table_bc_ + transportName
        table2 = 'lake_table_bc_' + transportName

        #column_names is the list of columns headings. It contains at the beginning
        #the heading of the first column ('lake_ID') and of the second column ('bc_type') only
        column_names=['lake_ID','bc_type']

        #If the table2 table is not in the DB...
        if (table2 not in tableList):
            #...then a new table renamed table2 is created.
            #Such table contains initially the fields lake_ID (integer) and bc_type (text)
            SQLstring = 'CREATE TABLE "%s" ("lake_ID" integer,"bc_type" varchar(30) );'%table2
            cur.execute(SQLstring)
            #According to the number of species in the transport model and to the number of
            #stress periods, fields renamed conc_sp_m_spec_n (float) are added (where m is the
            #m-th stress period and n in the n-th species) and the list column_names is updated
            for n in range(nspec):
                for m in range(nsp):
                    fieldname='conc_sp_' + str(m+1) + '_spec_' + str(n+1)
                    column_names.append(fieldname)
                    addColumn = 'ALTER TABLE "%s" ADD COLUMN "%s" float;'%(table2, fieldname)
                    cur.execute(addColumn)
        #If the table2 table is in the DB...
        elif table2 in tableList:
            #...a warning message appears and nothing more is done
            QMessageBox.warning(self, self.tr('Warning!'), self.tr('Table %s already exists!' % table2))
            return

        #cell_values is the list of values at each cell of each row of the table lake_table_bc_ + transportName
        cell_values=[]

        #row is a new dictionary. Its content will be like:
        #row={'lake_ID':1,'bc_type':'','conc_sp_1_spec_1':0.0,...,'conc_sp_m_spec_1':0.0,...,'conc_sp_1_spec_n':0.0,...,'conc_sp_m_spec_n':0.0}
        #The keys of this dictionary will be the elements of the list column_names,
        #the values of this dictionary will be the elements of the list cell_values
        row={}

        #bc_options is a list of possible boundary conditions types
        bc_options=['precipitation','runoff','pumping','evaporation']

        #i iterates over the number of lakes (which also corresponds to the
        #number of rows which will be added in the table lake_table_bc_ + transportName)
        for i in range(int(nlakes)):
            #For each lake, four rows will be created. This is because each
            #lake can potentially be assigned 4 types of boundary conditions
            for j in range(4):
                #The list cell_values is updated with the lake ID value
                cell_values.append(int(i+1))
                #The list cell_values is updated with the string 'input bc type...' (for the 'bc_type' column)
                cell_values.append(bc_options[j])
                #j iterates over the number of species
                for j in range(nspec):
                    for k in range(nsp):
                        #The list cell_values is updated with the value 0.0 (0.0 is the
                        #starting concentration for each species and for each stress period)
                        cell_values.append(0.0)
                #The built-in function enumerate allows to iterate over the list column_names.
                #The outputs of this function are:
                # - index is the index ID of each element of the list column_names
                # - key is the value of each element of the list column_names
                for index, key in enumerate(column_names):
                    #The dictionary row is filled
                    row[key]=cell_values[index]
                #A new row is added into the table lake_table_bc_ + transportName
                #(its content is nothing but the content of the dictionary row)
                columns = ', '.join(map(str, row.keys()))
                values = ', '.join(map(repr, row.values()))
                cur.execute('INSERT INTO "{0}" ({1}) VALUES ({2})'.format(table2, columns, values))
                #The list cell_values and the dictionary row are emptied
                cell_values=[]
                row={}

        # Close SpatiaLiteDB
        cur.close()
        # Save the changes
        con.commit()
        # Close connection
        con.close()

        #Add the lake_starting_conc_ + transportName table into QGis map
        uri = QgsDataSourceURI()
        uri.setDatabase(dbname)
        schema = ''
        table = table1
        geom_column = None
        uri.setDataSource(schema, table, geom_column)
        display_name = table1
        tableLayer = QgsVectorLayer(uri.uri(), display_name, 'spatialite')

        QgsMapLayerRegistry.instance().addMapLayer(tableLayer)

        #Add the lake_table_bc_ + transportName table into QGis map
        uri = QgsDataSourceURI()
        uri.setDatabase(dbname)
        schema = ''
        table = table2
        geom_column = None
        uri.setDataSource(schema, table, geom_column)
        display_name = table2
        tableLayer = QgsVectorLayer(uri.uri(), display_name, 'spatialite')

        QgsMapLayerRegistry.instance().addMapLayer(tableLayer)

        self.progressBar_2.setMaximum(100)

        #Close the dialog window after the execution of the algorithm
        self.reject()
