#!/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 constants import EPSILON, TREELABEL, STRATUMLABEL, STANDLABEL
from collections import defaultdict

import constants

class SilvicultureData(object):
    def __init__(self, silviculturetype, data, variables, parameters,
                 ptable_values, cash_flow_table, keepemptystrata=False):
        super(SilvicultureData, self).__init__()
        self.valid = True
        # original data
        self.data = data
        # data indices and
        # controlling data constructed from parameter, variable
        # and data values
        if silviculturetype == 'tending':
            inddict = constants.tendingind
            pardict = {}
            datadict = constants.tendingdata
        elif silviculturetype == 'planting':
            inddict ={}
            pardict = {}
            datadict = constants.plantingdata
        elif silviculturetype == 'supplemental_planting':
            inddict = {}
            pardict = {}
            datadict = constants.supplemental_plantingdata

        self.ind = inddict
        self.cdata = {}
        self.__set_cdata('par', pardict, parameters)
        self.__set_cdata('data', datadict, data)
        try:
            self.cdata['targetdensity'] = ptable_values[0][0][0]
            if silviculturetype == 'tending':
                self.cdata['preferenceorder'] = list(ptable_values[0][1])
        except IndexError:
            self.valid = False

    def separate_strata_and_remove_treeless(self, treecount):
        """
        Separate the stratums into artificially and naturally created and remove
        stratums without trees.
        """
        pind = self.ind['tree_parentid']
        spind = self.ind['stratum_sp']
        idind = self.ind['stratum_id']
        Nind = self.ind['stratum_N']
        artind = self.ind['stratum_artif']
        data = self.data
        natstrat = []
        artstrat = []
        for istratum, stratumdata in enumerate(data[STRATUMLABEL]):
            processstratum = False
            curspecies = stratumdata[spind]
            sid = stratumdata[idind]
            # Check that stratum has trees in it
            if sid in data[TREELABEL][:, pind]:
                processstratum=True
            # If there's trees in the stratum, process stratum
            if processstratum:
                if stratumdata[artind] == 0:
                    natstrat.append([istratum, stratumdata[3], stratumdata[spind]])
                else:
                    artstrat.append([istratum, stratumdata[3], stratumdata[spind]])
            else:
                targetdensity = self.cdata['targetdensity']
                if stratumdata[Nind] < 1:
                    # If stratum's N is 0, remove the stratum
                    stratumdata[-1] = 2
                elif stratumdata[Nind] < (treecount-targetdensity):
                    # Remove the whole stratum, set tree number as 0 and tag as 2
                    treecount = treecount - stratumdata[Nind]
                    stratumdata[Nind] = 0
                    stratumdata[-1] = 2
                # Remove only a portion of the trees in the stratum
                else:
                    stratumdata[Nind] = stratumdata[Nind] - (treecount - targetdensity)
                    # update tree number for the stratum, indicate modification
                    treecount = targetdensity
                    if stratumdata[Nind] < 1:
                        # If stratum's N is less than 1, remove the whole stratum
                        stratumdata[Nind] = 0
                        stratumdata[-1] = 2
                    else: stratumdata[-1] = 1
        return natstrat, artstrat, treecount

    def get_real_stem_numbers_for_stratums(self):
        """
        Construct a dictionary, where number key is stratum id and number of
        stems is the value
        """
        pind = self.ind['tree_parentid']
        nind = self.ind['tree_n']
        treelist = self.data[TREELABEL]
        stem_number = {}
        for tree in treelist:
            tid = tree[pind]
            if not stem_number.has_key(tid):
                stem_number[tid] = 0
            stem_number[tid] += tree[nind]
        return stem_number

    def get_diamclasses_for_stratums(self):
        """
        Construct a dictionary, where number key is stratum id and number of
        diameter classes is the value
        """
        pind = self.ind['tree_parentid']
        treelist = self.data[TREELABEL]
        treenums = {}
        for tree in treelist:
            tid = tree[pind]
            if not treenums.has_key(tid):
                treenums[tid] = 0
            treenums[tid] += 1
        return treenums

    def create_sorted_objectlist(self):
        """
        Construct a sorted list of objects by given sort attribute.
        """
        dind = self.ind['tree_d']
        idind = self.ind['tree_id']
        objectarray = self.data[TREELABEL]
        sortedarray = objectarray[numpy.lexsort((objectarray[:,idind], objectarray[:,dind]))]
        return sortedarray

    def __set_cdata(self, dtype, vars, data):
        for varname, vardef in vars.iteritems():
            if dtype == 'data':
                self.cdata[varname] = data[vardef[0]][vardef[1], vardef[2]]
            elif dtype in ('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_stem_number(self, treecount, origtreecount):
        """
        Update stem number after supplemental planting and indicate
        modifications
        """
        if abs(treecount - origtreecount) >= EPSILON:
            # update the tree number for the stand and indicate modification
            self.data[STANDLABEL][0, 3] = treecount
            self.data[STANDLABEL][0, -1] = 1
            return 1

    def targetdensity_reached(self, treecount):
        """
        Is target density already reached
        """
        targetdensity = self.cdata['targetdensity']
        if treecount < targetdensity or abs(treecount - targetdensity) < EPSILON:
            # reached target density, exit
            self.data[STANDLABEL][0, 3] = treecount
            self.data[STANDLABEL][0, -1] = 1
            return 1
