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

Harves operation implementations

*harvestop.py*

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

import math
import numpy
import copy
from collections import defaultdict
from constants import EPSILON, TREELABEL, STRATUMLABEL, STANDLABEL
from operations_ext import compute_treestem, bucking_algorithm

class SilvicultureOperation(object):
    def __init__(self, silviculturedata):
        super(SilvicultureOperation, self).__init__()
        self.sdata = silviculturedata
        self.treecount = self.sdata.cdata.get('N')
        self.origtreecount = self.treecount
        self.regen_sp = self.sdata.cdata.get('REGEN_SP')
        self.N_Sp = self.sdata.cdata.get('N_Sp')
        self.N_Ns = self.sdata.cdata.get('N_Ns')
        self.N_sb = self.sdata.cdata.get('N_sb')
    
    def remove_trees(self, stratlist, trees, treenums, sptoremove, treecount, sd):
        """
        Remove trees of given species from a stand. Trees will be removed until
        target density is reached.
    
            stratlist -
            trees -
            treenums -
            sptoremove -
            treecount -
            targetdensity -
            data -
        """
        targetdensity = sd.cdata['targetdensity']
        data = sd.data
        # If the stratums are empty, return without proceeding further
        if stratlist == []:
            return treecount
        # Create a list of species in the current stratums
        # stratlist contains lists of [stratum indice, parent id and species]
        splist = [s[2] for s in stratlist]
        # If the current species to be removed is not in the list, return
        if not sptoremove in splist:
            return treecount
        # Construct a list of the indices of the current stratums
        stratums = [s[0] for s in stratlist]
        origtreenums = copy.deepcopy(treenums)
        pind = sd.ind['tree_parentid']
        nind = sd.ind['tree_n']
        tidind = sd.ind['tree_id']
        sidind = sd.ind['stratum_id']
        spind = sd.ind['stratum_sp']
        Nind = sd.ind['stratum_N']
        # tree parentid is the actual simulation indice of parent stratum
        # (used in simulator data matrix)
        # construct a mapping from stratum id to stratum indice to find the
        # correct stratums based on parent info of the trees and mapping from
        # tree id to tree indice
        stratids = list(data[STRATUMLABEL][:,sidind])
        treeids = list(data[TREELABEL][:,tidind])
        stratum_id2ind = dict([(i, stratids.index(i)) for i in stratids])
        tree_id2ind = dict([(i, treeids.index(i)) for i in treeids])
        for tree in trees:
            stratum_id = tree[pind]
            stratum_ind = stratum_id2ind[stratum_id]
            stratum_sp = data[STRATUMLABEL][stratum_ind, spind]
            if stratum_sp == sptoremove:
                # If targetdensity is reached, stop iteration
                if abs(treecount - targetdensity) < EPSILON:  break
                tree_id = tree[tidind]
                # If the current stratum is not on the stratum list, move to next tree
                if stratum_ind not in stratums:
                    continue
                tree_ind = tree_id2ind[tree_id]
                # Remove either whole diameter class or part of it
                if (treecount - targetdensity) - tree[nind] < -EPSILON:
                    # Remove a part of the current diameter class
                    nremoved = (treecount - targetdensity)
                    treecount -= nremoved
                    data[TREELABEL][tree_ind, -1] = 1
                    data[TREELABEL][tree_ind, nind] -= nremoved
                    data[STRATUMLABEL][stratum_ind, Nind] -= nremoved
                    if data[STRATUMLABEL][stratum_ind, Nind] < EPSILON:
                        treenums[stratum_id] -= 1
                        data[TREELABEL][tree_ind, nind] = 0
                else:
                    # Remove the current diameter class and update indicator tags and
                    # tree numbers
                    treenums[stratum_id] -= 1
                    treecount -= tree[nind]
                    data[TREELABEL][tree_ind, -1] = 2
                    data[STRATUMLABEL][stratum_ind, Nind] -= tree[nind]
                    data[TREELABEL][tree_ind, nind] = 0
        # Go through the stratums and check if the stratums were totally removed
        # or modified
        for stratum_id in treenums.keys():
            stratum_ind = stratum_id2ind[stratum_id]
            if treenums[stratum_id] < EPSILON:
                data[STRATUMLABEL][stratum_ind][-1] = 2
            elif treenums[stratum_id] - origtreenums[stratum_id] < -EPSILON:
                data[STRATUMLABEL][stratum_ind][-1] = 1
        return treecount
    
    def define_stem_number_for_supplemental_planting(self, targetdensity):
        """
        Calculate treecount and define planted stem number
        """
        # Tree species (SP) used for regeneration
        SP = self.regen_sp
        # Treecount of regenerated tree species before supplemental planting
        if SP == 1:
            treecount_regensp = self.N_Sp
        elif SP == 2:
            treecount_regensp = self.N_Ns + self.N_Sp
        elif SP == 3:
            treecount_regensp = self.N_sb + self.N_Ns
        else:
            treecount_regensp = 0
        
        # Number of seedlings that will be planted
        suppseedlingsnum = targetdensity - treecount_regensp
        return suppseedlingsnum, SP