#!/usr/bin/env python
# encoding: utf-8
"""
Title: harvestdata

Data object for harvests

*harvestdata.py*

Copyright (c) 2007 Simosol Oy. Distributed under the GNU General Public License 2.0.
"""

import math
import numpy
from collections import defaultdict

import constants
from constants import EPSILON

class HarvestData(object):
    def __init__(self, harvesttype, data, variables, parameters,
                 par_table_val, cash_flow_table, cache_obj,
                 keepemptystrata=False):
        super(HarvestData, self).__init__()
        # original data
        self.data = data
        # initialize a global cache
        self.price_cache = cache_obj
        # data indices and
        # controlling data constructed from parameter, variable
        # and data values
        if harvesttype == 'thinning':
            inddict = constants.thinningind
            pardict = constants.thinningpar
            datadict = constants.thinningdata
        elif harvesttype == 'clearcut':
            inddict = constants.clearcutind
            pardict = {}
            datadict = constants.clearcutdata
        elif harvesttype == 'seedtree':
            inddict = constants.seedtreeind
            pardict = {}
            datadict = constants.seedtreedata
        elif harvesttype == 'removeseedtree':
            inddict = constants.removeseedtreeind
            pardict = {}
            datadict = constants.removeseedtreedata
        elif harvesttype == 'standvalue':
            inddict = constants.standvalueind
            pardict = {}
            datadict = constants.standvaluedata
        elif harvesttype == 'first_thinning_stem_number':
            inddict = constants.thinningind
            pardict = constants.first_thinning_N_par
            datadict = constants.first_thinning_N_data
        elif harvesttype == 'wood_assortments':
            inddict = constants.woodassortind
            pardict = {}
            datadict = constants.woodassortdata
        elif harvesttype == 'thinning_with_removed_volume':
            inddict = constants.thinningind
            pardict = constants.thinningpar
            datadict = constants.thinning_with_removed_volumedata
        elif harvesttype == 'selection_cut':
            inddict = constants.thinningind
            pardict = constants.selectionpar
            datadict = constants.thinningdata
        self.ind = inddict
        self.cdata = {}
        self.cdata['target_var'] = None # thinning target var will be set later
        self._set_cdata('par', pardict, parameters)
        self._set_cdata('data', datadict, data)
        # stratum data
        self.stratumdata = data[constants.STRATUMLABEL]
        if not keepemptystrata:
            self._remove_empty_strata()
        self.cdata['n_strata'] = len(self.stratumdata)
        self.cdata['sp_proportion'] = {}
        if harvesttype == 'thinning' or \
           harvesttype == 'first_thinning_stem_number':
            sppref = list(par_table_val[0][0])
            fpsp = self.cdata.get('forced_preferred_sp')
            sppref = self._modify_sp_pref(sppref, fpsp)
            pp = self.cdata.get('preference_percentage')
            if pp is None:
                pp = []
            for i in xrange(len(sppref)):
                if i < len(pp):
                    self.cdata['sp_proportion'][sppref[i]] = pp[i]
                else:
                    self.cdata['sp_proportion'][sppref[i]] = 0.0
            self.cdata['regional_log_reduction'] = par_table_val[1][0][0]
        if harvesttype == 'first_thinning_stem_number':
            N_target = par_table_val[2][0][0]
            self.cdata['N_target'] = N_target
        if harvesttype == 'seedtree':
            #
            allsp = self.stratumdata[:, self.ind['stratum_sp']]
            self.cdata['seedtreespecies']=allsp

            vol = 0
            seedtreesp2 = self.cdata['regen_sp']
            for stratum in self.stratumdata:
                if stratum[2] != self.cdata['regen_sp'] and stratum[3] > vol:
                    seedtreesp2 = stratum[2]
                    vol = stratum[3]
            self.cdata['seedtreesp2']=seedtreesp2
            self.cdata['N_seedtrees'] = par_table_val[0][0][0]
            self.cdata['regional_log_reduction'] = par_table_val[1][0][0]
        ops = ('standvalue', 'clearcut', 'removeseedtree', 'wood_assortments',
               'selection_cut')
        if harvesttype in ops:
            self.cdata['regional_log_reduction'] = par_table_val[0][0][0]

        # tree data
        #     sort to ascending order by diameter
        self.treedata = data[constants.TREELABEL]
        treeids = self.treedata[:, self.ind['tree_id']]
        diam = self.treedata[:, self.ind['tree_d']]
        self.cdata['removal_order'] = numpy.lexsort((treeids, diam))
        self.cdata['n_trees'] = numpy.size(self.treedata, 0)
        self.cdata['n_children'] = self._count_dclass_by_parent()
        # cash flow
        #cftable = cash_flow_table.active_array.dict
        cftable = cash_flow_table[0].active_array.array
        keycount = cash_flow_table[0].keycount
        self._constr_pricematrix(cftable, keycount)
        # volumes
        self.result = {}
        self.result['dclass_vol'] = {}  # total volume of each diameter class
        self.result['stratum_vol'] = {}  # total volume of each stratum
        #These are the usable amounts of biomass in the stratums
        self.result['bm_branch'] = {}
        self.result['bm_stump_root'] = {}
        self.result['bm_stump'] = {}
        self.result['bm_root'] = {}
        self.result['bm_scrap'] = {}
        #These are the total amounts of biomass in the stratums
        self.result['bm_scrap_all'] = {}
        self.result['bm_root_all'] = {}
        self.result['bm_stump_all'] = {}
        self.result['bm_branch_all'] = {}
        self.result['bm_foliage'] = {}
        #     construct dictionaries for storing timber assortment volumes
        #     and values
        labels = cash_flow_table[0].classifiers['assortment'].labels
        ao = self.cdata['assortment_order'] = []
        for label in labels:
            ao.append(float(label.value))
        self.cdata['ao']=ao
        self._constr_timber_dicts(labels)
        ops = ['standvalue','wood_assortments','removeseedtree']
        if harvesttype not in ops:
            bm_cftable = cash_flow_table[1].active_array.array
            bm_keycount = cash_flow_table[1].keycount
            self._constr_bm_pricematrix(bm_cftable, bm_keycount)
            bm_labels = cash_flow_table[1].classifiers['bm_assortment'].labels
            bm_ao = self.cdata['assortment_order'] = []
            for bm_label in bm_labels:
                bm_ao.append(bm_label)
            self.cdata['bm_ao']=bm_ao
            self._constr_bm_dicts(bm_labels)

    def get_ba_from_treelist(self, start=0, stratid=None):
        """
        Function: __get_ba_from_treelist
        """
        if stratid is None:
            treeind = xrange(start, len(self.treedata))
        else:
            parentid = self.treedata[:, self.ind['tree_parentid']]
            treeind = numpy.where(parentid == stratid)[0]
        diam = self.treedata[treeind, self.ind['tree_d']]
        n = self.treedata[treeind, self.ind['tree_n']]
        return numpy.sum(math.pi * (diam * 0.005) ** 2 * n)

    def get_dgm_from_treelist(self, start=0, stratid=None):
        """
        Function: get_dgm_from_treelist
        """
        if stratid is None:
            treeind = xrange(start, len(self.treedata))
        else:
            parentid = self.treedata[:, self.ind['tree_parentid']]
            treeind = numpy.where(parentid == stratid)[0]
        diam = self.treedata[treeind, self.ind['tree_d']]
        ba = (math.pi * diam ** 2) / 4
        n = self.treedata[treeind, self.ind['tree_n']]
        return numpy.sum(diam * ba * n) / numpy.sum(ba * n)

    def get_n_from_treelist(self, start=0, stratid=None):
        if stratid is None:
            treeind = xrange(start, len(self.treedata))
        else:
            parentid = self.treedata[:, self.ind['tree_parentid']]
            treeind = numpy.where(parentid == stratid)[0]
        return numpy.sum(self.treedata[treeind, self.ind['tree_n']])

    def set_d_cut_all_below(self):
        cd = self.cdata
        cd['d_cut_all_below'] = cd['cut_all_below'] * cd['d_gm']

    def set_d_cut_all_above(self):
        cd = self.cdata
        cd['d_cut_all_above'] = cd['cut_all_above'] * cd['d_gm']

    def set_track_removal(self):
        """
        proportion of the basal area / stem number to be removed because of
        logging tracks
        """
        dist = self.cdata['track_distance']
        width = self.cdata['track_width']
        N = numpy.sum(self.stratumdata[:, self.ind['stratum_N']])
        if N > 1500:
            removal=(((100/dist)*width*100)/10000*100)/100
        else:
            removal=(3+(((100/dist)*width*100)/10000*100-3)/1000*(N-500))/100
        if removal < 0:
            removal = 0
        self.cdata['track_removal'] = removal

    def set_thinning_control_param(self, targetvar):
        self.cdata['target_var'] = targetvar.lower()
        if targetvar=='BA':
            ul = self.cdata['upper_limit']
            ll = self.cdata['lower_limit']
            td = self.cdata['target_density']
            BA_current = self.get_ba_from_treelist()
            BA_min = BA_current * (1 - (self.cdata['max_ba_removal'] / 100))
            if ll <= 0:
                ll = 0
            BA_target = (ul - ll) * td + ll
            if BA_target < BA_min:
                BA_target = BA_min
            self.cdata['target'] = BA_target
            self.cdata['to_remove'] = BA_current - BA_target
        elif targetvar=='N':
            N_current = self.get_n_from_treelist()
            N_target = self.cdata['N_target']
            min_N = self.cdata['min_N_after_thinning']
            if N_target < min_N:
                N_target = min_N
            self.cdata['target'] = N_target
            self.cdata['to_remove'] = N_current - N_target
        elif targetvar=='V_com':
            self.cdata['to_remove'] = self.cdata['v_harvested']


    def set_stump_inclusion(self):
        """
        Calculate the proportion of stumps that should be left in the ground
        Sum up the number of stems where stump diameter is greater than the
        minimum stump diameter limit
        """
        stumpdiam = self.treedata[:, self.ind['tree_d_s']]
        ind = numpy.where((stumpdiam - self.cdata['min_stump_diam']) \
                          >= -EPSILON)[0]
        n = self.treedata[ind, self.ind['tree_n']]
        stump_lift_prop =  self.cdata.get('stump_lift_prop', 1.0)
        stump_excl_num = self.cdata.get('stump_excl_num', 0.0)
        stump_incl_prop = stump_lift_prop * 1 - (stump_excl_num / n.sum())
        if stump_incl_prop < 0.0:
            stump_incl_prop = 0
        self.cdata['stump_incl_prop'] = stump_incl_prop

    def _constr_prop(self, prop):
        prop = prop / 100.0
        if prop > 1.0:
            prop = 1.0
        elif prop < 0.0:
            prop = 0.0
        return prop

    def _constr_pricematrix(self, cashflowtable, keycount):
        """
        price matrix of
        0: array of [assortment id, diameter, length, value]
        1: minimum diameter
        2: minimum length
        3: pulp label
        """
        process = True
        use_cache = bool(int(self.cdata['use_cache'])) and \
                self.price_cache.use_cache
        if use_cache:
            process, pricematrix = self.price_cache.check(cashflowtable)
            if not process and pricematrix is not None:
                self.cdata['price_matrix'] = pricematrix
        if process:
            self.cdata['price_matrix'] = {}  # pricematrix per each species
            pm = self.cdata['price_matrix']
            pulplabel = {}  # lowest assortment label (biggest number!)
            minlen, mindiam = {}, {}  # NB! start with big enough number!
            for row in cashflowtable:
                datarow, value = row[:-1], row[-1]
                key = tuple(datarow[:keycount])
                assrt, diam, length = datarow[keycount:]
                if not key in pm:
                    pm[key] = [[], None, None, None]
                    pulplabel[key] = 0
                    minlen[key] = 10000
                    mindiam[key] = 10000
                pm[key][0].append([assrt, diam, length, value])
                # update minimum diameter
                if diam < mindiam[key]: mindiam[key] = diam
                # update minimum length
                if length < minlen[key]: minlen[key] = length
                # update probable pulp label
                if assrt > pulplabel[key]: pulplabel[key] = assrt

            for key, container in pm.iteritems():
                container[0] = numpy.array(container[0])
                container[1:4] = mindiam[key], minlen[key], pulplabel[key]

            if use_cache:
                # update cashflowtable cache
                self.price_cache.update(cashflowtable, pm)

    def _constr_bm_pricematrix(self, cashflowtable, keycount):
        """
        price matrix of
        0: array of [assortment id, value]
        """
        self.cdata['bm_price_matrix'] = {}  # pricematrix per each species
        bm_pm = self.cdata['bm_price_matrix']
        for row in cashflowtable:
            datarow, value = row[:-1], row[-1]
            key = datarow[keycount-1]
            assrt = datarow[keycount]
            if not key in bm_pm:
                bm_pm[key] = [[]]
            bm_pm[key][0].append([assrt, value])

    def _constr_timber_dicts(self, assortments):
        """
        Construct dictionaries for storing the volumes and values by
        stratum / diameter class and assortment (grade). Maintains one to one
        order mapping of data objects and result objects
        """
        for level in (constants.STRATUMLABEL, constants.TREELABEL):
            gradequantities = {}
            scrapwood = {}
            bmbranch = {}
            bmstump = {}
            bmstumproot = {}
            bmscrap = {}
            bmroot = {}
            bmfoliage = {}
            bmstumpall = {}
            bmbranchall = {}
            bmscrapall = {}
            bmrootall = {}
            for ind, obj in enumerate(self.data[level]):
                gradequantities[ind] = {}
                for assortment in assortments:
                    aval = float(assortment.value)
                    gradequantities[ind][aval] = [0.0, 0.0]
                scrapwood[ind] = 0.0
                bmbranch[ind] = 0.0
                bmstump[ind] = 0.0
                bmstumproot[ind] = 0.0
                bmscrap[ind] = 0.0
                bmroot[ind] = 0.0
                bmfoliage[ind] = 0.0
                bmbranchall[ind] = 0.0
                bmstumpall[ind] = 0.0
                bmscrapall[ind] = 0.0
                bmrootall[ind] = 0.0
                if level == constants.STRATUMLABEL:
                    self.result['stratum_timber'] = gradequantities
                    self.result['stratum_scrap'] = scrapwood
                    self.result['bm_branch'] = bmbranch
                    self.result['bm_stump_root'] = bmstumproot
                    self.result['bm_stump'] = bmstump
                    self.result['bm_scrap'] = bmscrap
                    self.result['bm_root'] = bmroot
                    self.result['bm_foliage'] = bmfoliage
                    self.result['bm_scrap_all'] = bmscrapall
                    self.result['bm_root_all'] = bmrootall
                    self.result['bm_branch_all'] = bmbranchall
                    self.result['bm_stump_all'] = bmstumpall
                elif level == constants.TREELABEL:
                    self.result['dclass_timber'] = gradequantities

    def _constr_bm_dicts(self, assortments):
        """
        Construct dictionaries for storing the volumes and values by
        stratum / diameter class and assortment (grade). Maintains one to one
        order mapping of data objects and result objects
        """
        level = constants.STRATUMLABEL
        gradequantities = {}
        for ind, obj in enumerate(self.data[level]):
            gradequantities[ind] = {}
            for assortment in assortments:
                aval = float(assortment.value)
                gradequantities[ind][aval] = [0.0, 0.0]
            if level == constants.STRATUMLABEL:
                self.result['bm_timber'] = gradequantities

    def _count_dclass_by_parent(self):
        """
        Go through all diameter classes in the data and count the number
        of classes of each parent object.

        NB! The parent indice refers to the parent object's position in the
        current operation data, not the real simulation data. a.k.a the
        indice is not the actual object indice!
        """
        nclass = defaultdict(int)
        for dclass in self.treedata:
            stratumid = self.stratumdata[:, self.ind['stratum_id']]
            parentid = dclass[self.ind['tree_parentid']]
            parent_id = numpy.where(stratumid == parentid)[0][0]
            nclass[parent_id] += 1
        return nclass

    def _modify_sp_pref(self, preforder, forcedprefsp):
        """
        remove non existing species from the preference list and move the
        forced one to first in the list if set
        """
        neworder = []
        allsp = self.stratumdata[:, self.ind['stratum_sp']]
        for sp in preforder:
            if sp in allsp:
                neworder.append(sp)
        if forcedprefsp in neworder:
            popind = neworder.index(forcedprefsp)
            fpsp = neworder.pop(popind)
            neworder.insert(0, fpsp)
        return neworder

    def _remove_empty_strata(self):
        """
        flag empty strata for removal (NB! not needed in matrixsimo)
        """
        volind = self.ind['stratum_vol']
        stag = self.ind['stratum_tag']
        for stratum in self.data[constants.STRATUMLABEL]:
            if stratum[volind] == 0:
                stratum[stag] = 2

    def _set_cdata(self, dtype, vars, data):
        for varname, vardef in vars.iteritems():
            if dtype == 'data':
                if vardef[0] == 'proportion':
                    val = data[vardef[1]][vardef[2], vardef[3]]
                    self.cdata[varname] = self._constr_prop(val)
                else:
                    self.cdata[varname] = data[vardef[0]][vardef[1], vardef[2]]
            elif dtype == 'par':
                vtype, ind = vardef
                if vtype is None:
                    self.cdata[varname] = data[ind]
                elif vtype == 'bool':
                    vals = data[ind]
                    if isinstance(vals, list):
                        self.cdata[varname] = [bool(x) for x in vals]
                    else:
                        self.cdata[varname] = bool(vals)
                elif vardef[0] == 'proportion':
                    vals = data[ind]
                    if isinstance(vals, list):
                        self.cdata[varname] = [self._constr_prop(x) for x in vals]
                    else:
                        self.cdata[varname] = self._constr_prop(vals)
                else:
                    self.cdata[varname] = None

    def update_to_remove(self, cut_timber=0):
        """
        update basal area or stem number still to remove
        """
        targetvar = self.cdata['target_var']
        if targetvar == 'ba':
            ba_current = self.get_ba_from_treelist()
            self.cdata['to_remove'] = ba_current - self.cdata['target']
        elif targetvar == 'n':
            n_current = self.get_n_from_treelist()
            self.cdata['to_remove'] = n_current - self.cdata['target']
        elif targetvar == 'v_com':
            self.cdata['to_remove'] = self.cdata['v_harvested'] - cut_timber
