# -*- 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, read_data, pop_message, ModelPath
from datetime import datetime, timedelta, date
import numpy as np
from freewat.crop_utils import CropGrowthModule
import matplotlib.pyplot as pl
import csv

FORM_CLASS, _ = uic.loadUiType(os.path.join(os.path.dirname(__file__), 'ui/ui_runCrop.ui'))

class runCropDialog(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.runCrop)
        self.manageGui()

    def manageGui(self):

        # get the modeltable as QgsVectorLayer
        layerMap = QgsMapLayerRegistry.instance().mapLayers()
        for k, v in layerMap.items():
            if 'modeltable_' in k:
                self.modeltable = v

        # get the layer list from the TOC
        layerNameList = getVectorLayerNames()

        # get the crop model table name and fill the combobox with the crop model models
        self.cmbCropModel.clear()
        for lay in layerNameList:
            if 'crop_models_table' in lay:
                n = 'crop_models_table'
                self.cropmodeltable = getVectorLayerByName(n)

        # fill the combobox with the crop_model_name values
        try:
            for model in self.cropmodeltable.getFeatures():
                self.cmbCropModel.addItems([model['crop_model_name']])
        except Exception as e:
            self.close()
            pop_message('No Crop Model Table found in TOC.\n Please create it and then run the model', 'warning')

        # climate data table
        clim_layers= []
        for clm in layerNameList:
            if 'ClimateData' in clm:
                clim_layers.append(clm)
        # filter the climate combobox with only the ClimateTable layer
        self.cmbClimate.addItems(clim_layers)


        # soil_crop table
        soil_crop = []
        for i in layerNameList:
            if 'soil_crop' in i:
                soil_crop.append(i)
        # filter the Crop layer combobox with only the Crop Layer (vector)
        self.cmbSoilCrop.addItems(soil_crop)


    def runCrop(self):

        # Load all the Date/Time information

        # initial date of the model from the modeltable and convert it to a datetime object
        for i in self.modeltable.getFeatures():
            init_date = i['initial_date']
        date0 = datetime.strptime(init_date, "%Y-%m-%d")
        date0 = date0.date()


        # initial seeding and harvesting DATE from the Crop from the combobox
        for i in self.cropmodeltable.getFeatures():
            if self.cmbCropModel.currentText() in i['crop_model_name']:
                init_data_ts = i['ts']
                init_data_th = i['th']

        # convert seeding and harvesting DATE to date objects
        self.date_ts0 = datetime.strptime(init_data_ts, "%Y-%m-%d")
        self.date_ts0 = self.date_ts0.date()

        self.date_th0 = datetime.strptime(init_data_th, "%Y-%m-%d")
        self.date_th0 = self.date_th0.date()

        # date calculations with the seeding, harvesting and initial datetime objects
        # difference from seeding date respect to initial date
        diff_ts = self.date_ts0 - date0
        diff_ts = diff_ts.days

        # difference from seeding date respect to initial date
        diff_th = self.date_th0 - date0
        diff_th = diff_th.days

        # total TS number (in days)
        nts = (diff_th - diff_ts) + 1

        # list of interval of Dates between ts0 and th0 as datetime object (useful for the plot creation)
        dates_interval = [self.date_ts0 + timedelta(n) for n in range(int(((self.date_th0 + timedelta(1)) - self.date_ts0).days))]


        # check if the initial date (from the model table) is smaller than the other ones
        if date0 > self.date_ts0 or date0 > self.date_th0:
            QMessageBox.warning(self, self.tr('Warning!'), self.tr('Check the input date and the intial date of the model in the modeltable'))
            return


        # throw error if the harvest time is smaller than the seeding time
        if self.date_th0 < self.date_ts0:
            QMessageBox.warning(self, self.tr('Warning!'), self.tr('Harvest time has to be AFTER seed time'))
            return

        # LOAD CLIMATE DATA

        # get the climate data from the climate table
        climateTable = getVectorLayerByName(self.cmbClimate.currentText())

        # initialize the list of the parameters
        Tmin = []
        Tmax = []
        Rad = []

        # populate the lists with float values of the climate table
        for i in climateTable.getFeatures():
            Tmin.append(float(i['MinT']))
            Tmax.append(float(i['MaxT']))
            Rad.append(float(i['Rad']))

        # update the lists with only values of the correct time steps interval (from seeding to harvesting)
        Tmin = Tmin[diff_ts: diff_th + 1]
        Tmax = Tmax[diff_ts: diff_th + 1]
        Rad = Rad[diff_ts: diff_th + 1]


        # LOAD CROP DATA created in the other dialog (so from the table name_cgm)

        # get the crop table as QgsVectorLayer
        self.crop_type = self.cmbCropModel.currentText()
        self.crop_type += '_cgm'

        layerMap = QgsMapLayerRegistry.instance().mapLayers()
        for k, v in layerMap.items():
            if self.crop_type in k:
                self.crop_table = v

        # list of crop names
        self.crp_name = [i['crop_name'] for i in self.crop_table.getFeatures()]

        # temporary list
        tmp_ll = []
        for i in self.crop_table.getFeatures():
            tmp_ll.append(i.attributes()[4:])



        # create dictionary with KEYS = CROP ID and values are the list for each key
        self.table_dic = {k[0] : k[1:] for k in tmp_ll}


        # nrow and ncol from the SoilCrop layer
        # get the crop table from the combobox
        soil_crop_table =  getVectorLayerByName(self.cmbSoilCrop.currentText())

        # build nrow and ncol
        nrow = []
        ncol = []

        for i in soil_crop_table.getFeatures():
            nrow.append(int(i['ROW']))
            ncol.append(int(i['COL']))

        nrow = max(nrow)
        ncol = max(ncol)


        # LOAD ET_ARRAY DATA

        # load the ET_ARRAY.OUT file from the working directory
        self.pathfile = ModelPath()
        et = os.path.join(self.pathfile, 'ET_ARRAY.OUT')

        # from the loaded file return a dictionary with KEYS = TS and VALUES = the blocks (array) of the file diff(s) are zero based, so -1 is needed
        et_dic = read_data(et, nrow, diff_ts - 1 , diff_th - 1)


        # get the information from the CROP LAYER
        crop_id = []
        # fill list looping over CropId
        for i in soil_crop_table.getFeatures():
            crop_id.append(int(i['crop_id']))

        # transform the crop id list with only its sorted unique values
        crop_id = sorted(list(set(crop_id)))

        # initialize a zero array of the crop id (array dimension correspond to the crop layer)
        crop_id_array = np.zeros(shape = (nrow, ncol))

        # replace the zeros of the array with the corresponding crop_id values of the crop layer
        for f in soil_crop_table.getFeatures():
            nr = int(f['ROW'])
            nc = int(f['COL'])
            # -1 is needed for the zero based indexing
            crop_id_array[nr - 1, nc - 1] = int(f['crop_id'])

        # initialize empty lists of the final Tr and ncel values
        Tr = []
        ncel = []

        # loop over the time step number
        for t in range(diff_ts - 1, diff_th):
            # atr is array of the ET_ARRAY.OUT file with blocks for each time step (et_dic[t])
            atr = et_dic[t]
            # initialize Tr_temp and ncel_temp lists that are lists of only the crop id != 0
            Tr_temp = [0.0 for j in crop_id]
            ncel_temp = [0 for j in crop_id]
            # loop over row and col (crop layer dimension)
            for nr in range(nrow):
                for nc in range(ncol):
                    id_crop = crop_id_array[nr, nc]
                    for k, c in enumerate(crop_id):
                        if c == id_crop:
                            ncel_temp[k] = ncel_temp[k] + 1
                            Tr_temp[k] = Tr_temp[k] + atr[nr, nc]
            for k in range(len(Tr_temp)):
                if ncel_temp[k] != 0:
                    Tr_temp[k] = Tr_temp[k] / ncel_temp[k]
            # Tr and ncel are lists of list: for each time step and for each crop id
            Tr.append(Tr_temp)
            ncel.append(ncel_temp)


        # create parent folder for the crops output
        crops_datapath = os.path.join(self.pathfile, 'crops_output')
        if os.path.exists(crops_datapath):
            pass
        else:
            if not os.access(crops_datapath, os.F_OK):
                os.mkdir(crops_datapath)

        # create single child folder for all the crops models
        datapath = os.path.join(crops_datapath, self.cmbCropModel.currentText())
        if os.path.exists(datapath):
            pass
        else:
            if not os.access(datapath, os.F_OK):
                os.mkdir(datapath)

        # initialize the crop growth module object
        cgm_dict = {}

        for j, cpid in enumerate(crop_id) :

            t_base = self.table_dic[cpid][0]
            lai_max = self.table_dic[cpid][1]
            gdd_em = self.table_dic[cpid][2]
            alpha_1 = self.table_dic[cpid][3]
            alpha_2 = self.table_dic[cpid][4]
            gdd_lai_max = self.table_dic[cpid][5]
            rue = self.table_dic[cpid][6]
            hi_ref = self.table_dic[cpid][7]
            ky  =  self.table_dic[cpid][8]

            time_data = [t for t in range (diff_ts, diff_th + 1)]

            Tr_crop = [0 for jj in range(nts)]

            for i in range(nts):
                Tr_crop[i] = Tr[i][j]

            # INITIALIZE THE CGM OBJECT (ALL THE INPUT DATA ARE READY)
            cgm = CropGrowthModule(t_min  = Tmin, t_max = Tmax, t_base = t_base, lai_max = lai_max, gdd_em = gdd_em, alpha_1 = alpha_1, alpha_2 = alpha_2, gdd_lai_max = gdd_lai_max, rad = Rad, rue = rue, hi_ref = hi_ref, tr_act = Tr_crop, tr_max = Tr_crop, ky = ky, time_data = time_data, time_interval = dates_interval)

            # run the CGM model for this cgm object
            cgm.run_model()

            # path plot1
            plt1 = os.path.join(datapath, 'agb_act_{}.png').format(self.crp_name[j])
            plt2 = os.path.join(datapath, 'agb_pot_{}.png').format(self.crp_name[j])

            # save plot for each crop_id and for each cgm output (2 plot for each crop_id)
            cgm.create_plot(self.crp_name[j], plt1, plt2)

            # Write input for this cgm object
            filename = os.path.join(datapath, 'cropmodel_{}.csv'.format(self.crp_name[j]))

            cgm.write_output(filename)

            cgm_dict[cpid] = cgm

        # output yield csv file within the correct crop model folder
        fileout = os.path.join(datapath, 'yield.csv')

        with open(fileout, 'wb') as csvfile:
            fieldnames = ['crop_id', 'crop_name', 'yh', 'y_act']
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

            writer.writeheader()

##            for i, j in enumerate(cgm_dict.keys()):
##                print i, j

            for jj, kk in enumerate(cgm_dict.keys()):
                # writer.writerow({'crop_id': kk, 'crop_name': self.crp_name[kk - 1], 'yh': cgm_dict[kk].yh, 'y_act': cgm_dict[kk].y_act})
                writer.writerow({'crop_id': kk, 'crop_name': self.crp_name[jj], 'yh': cgm_dict[kk].yh, 'y_act': cgm_dict[kk].y_act})

        pop_message(self.tr('Output correctly saved in folder {}'.format(datapath)), 'information')
