# -*- 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 math import exp
import matplotlib.pyplot as pl
import csv

class CropGrowthModule():
    '''
    Module for crop modelling of farm process

    INPUT
    -----
    t_min
    t_max
    t_base
    lai_max
    gdd_em
    alpha_1
    alpha_2
    gdd_lai_max
    rad
    rue
    hi_ref
    tr_act
    tr_max
    ky


    OUTPUT
    ------
    gdd
    agb_pot
    yh
    agb_act
    y_act
    '''

    def __init__(self, t_min  = None, t_max = None, t_base = None, lai_max = None, gdd_em = None, alpha_1 = None, alpha_2 = None, gdd_lai_max = None, rad = None, rue = None, hi_ref = None, tr_act = None, tr_max = None, ky = None, time_data = None, time_interval = None):

        #Input
        self.t_min = t_min
        self.t_max = t_max
        self.t_base = t_base
        self.lai_max = lai_max
        self.gdd_em = gdd_em
        self.alpha_1 = alpha_1
        self.alpha_2 = alpha_2
        self.gdd_lai_max = gdd_lai_max
        self.rad = rad
        self.rue = rue
        self.hi_ref = hi_ref
        self.tr_act = tr_act
        self.tr_max = tr_max
        self.ky = ky
        self.time_data = time_data
        self.time_interval = time_interval

        #Output
        self.gdd = None
        self.agb_pot = None
        self.yh = None
        self.agb_act = None
        self.y_act = None


    def run_model(self):
        '''
        takes all the inputs and calculate the equations
        '''
        t_avg = []

        # calculate t_avg from t_max and t_min
        for i in range(len(self.t_min)):
            t_avg.append((self.t_min[i] + self.t_max[i])/2.0)

        # define hu (daily heat units)
        hu = []
        for i in range(len(t_avg)):
            if t_avg[i] < self.t_base:
                hu.append(0)
            else:
                hu.append(t_avg[i] - self.t_base)

        # calculate gdd (growing degree days)
        gdd = [hu[0]]
        for i in range(1,len(hu)):
            gdd.append(gdd[i-1] + hu[i])

        self.gdd = gdd

        # calculate lai_pot (potential leaf area index at day t)
        lai_pot = []
        for i in range(len(self.gdd)):
			if self.gdd[i] <= self.gdd_em:
				lai_pot.append(0)
			else:			   
				if self.gdd[i] <= self.gdd[i] * self.lai_max:
					lai_pot.append(self.lai_max * (((self.gdd[i] - self.gdd_em) / self.gdd_lai_max)**2 * exp(2/self.alpha_1 * (1- ((self.gdd[i] - self.gdd_em) / self.gdd_lai_max)**self.alpha_1))))
				else:
					lai_pot.append(self.lai_max * (((self.gdd[i] - self.gdd_em) / self.gdd_lai_max)**2 * exp(2/self.alpha_2 * (1- ((self.gdd[i] - self.gdd_em) / self.gdd_lai_max)**self.alpha_2))))

        # calculate ipar () (daily amount of crop-intercepted photosinthetically active radiation)
        ipar = []
        # calculate agb-pot (daily increase in optimum above-ground biomass)
        d_agb_pot = []
        for i in range(len(self.rad)):
            ipar.append(0.5*self.rad[i]*(1-exp(-(0.65*lai_pot[i]))))
            d_agb_pot.append(self.rue * ipar[i])

        # calculate agb_pot (potential above-ground biomass at day t)
        # FINAL RESULT
        agb_pot = [d_agb_pot[0]]
        for i in range(1,len(d_agb_pot)):
            agb_pot.append(agb_pot[i-1] + d_agb_pot[i])

        # calculate agb_poth (above-ground biomass at harvest)
        agb_poth = agb_pot[-1]

        # calculate crop Yield at harvest
        ## FINAL RESULT
        yh = agb_poth * self.hi_ref

        # return retults
        self.agb_pot = agb_pot
        self.yh = yh


        # calculate ws (daily water stress factor)
        ws = []
        for i in range(len(self.tr_act)):
            ws.append(1-(self.tr_act[i] / self.tr_max[i]))

        # calculate wss (cumulated water stress factor)
        wss = sum(self.tr_act) / sum(self.tr_max)

        # calculate abg_act (daily actual above-ground biomass)
        ## FINAL RESULT
        agb_act = []
        for i in range(len(ws)):
            agb_act.append(self.agb_pot[i]*ws[i])

        # calculate y_act (water limited yield at harvest)
        ## FINAL RESULT
        y_act = self.yh * (1-self.ky * (1-wss))

        # return retults
        self.agb_act = agb_act
        self.y_act = y_act


    def write_output(self, filename):
        '''
        write the output csv files
        file is made by 4 columns:
            1. time series of the data as pure day indexes (like 1,2,3)
            2. time series of real dates (like 2016-01-17)
            3. agb_act values calculated in the run_model method
            4. agb_pot values calculated in the run_model method
        '''

        with open(filename, 'wb') as csvfile:
            fieldnames = ['sp_index', 'date', 'agb_act', 'agb_pot']
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

            writer.writeheader()

            for i in range(len(self.time_data)):
                writer.writerow({'sp_index': self.time_data[i], 'date': self.time_interval[i], 'agb_act': self.agb_act[i], 'agb_pot': self.agb_pot[i]})


    def create_plot(self, crop_name, plot1, plot2):
        '''
        create png plot for each crop

        crop_name: String. The name of the single crop
        plot1: String. The path of the first plot (time vs yvalue1)
        plot2: String. The path of the second plot (time vs yvalue2)
        '''

        # tvalue is the Date object for xaxis
        tvalue = self.time_interval
        # this is agb_act
        yvalue1 = self.agb_act
        # this is agb_pot
        yvalue2 =  self.agb_pot

        self.fig, ax = pl.subplots()
        # plot Date and not just numbers and autoformat axis
        ax.plot_date(tvalue, yvalue1, '--o')
        self.fig.autofmt_xdate()
        ax.set_xlabel('Time (day)')
        ax.set_ylabel('Daily actual above-gound biomass ')
        ax.set_title(crop_name)
        ax.grid()

        self.fig2, ax2 = pl.subplots()
        # plot Date and not just numbers and autoformat axis
        ax2.plot_date(tvalue, yvalue2, '--o')
        self.fig2.autofmt_xdate()
        ax2.set_xlabel('Time (day)')
        ax2.set_ylabel('Potential above-gound biomass ')
        ax2.set_title(crop_name)
        ax2.grid()

        pl.savefig(plot1)
        pl.savefig(plot2)

        return self.fig, self.fig2
