#!/usr/bin/env python
# encoding: utf-8
"""
Title: SIMO aggregation model implementations

simoaggr includes the aggregation functions for summing up and
calculating means, averages etc. for simulation data. These functions are
used in the simulation model chains in the same way as the prediction
functions.


*aggr.py*

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


import math
import numpy


def _nanmean(a):
    numer = numpy.nansum(a)
    denom = numpy.sum(numpy.isnan(a) == False)
    if denom == 0.0:
        return 0.0
    return numer / denom

def _nanargmax(values):
    mx = numpy.nanargmax(values)
    if numpy.isnan(mx):
        return None
    return mx

def _nanargmin(values):
    mx = numpy.nanargmin(values)
    if numpy.isnan(mx):
        return None
    return mx


def simo_sum(arg):
    """Sum operands of double type

    NB! Name changed to 'simo_sum' because python has a built-in function
    'sum'
    """
    su = numpy.nansum if arg.use_nan_funcs else numpy.nansum
    if arg.operands == 1:
        vals = arg.values[0]
        tind = arg.oper_target_index[0][:,3:]
        ssum = map(lambda x: su(vals[x[0]:x[1]]), tind)
        arg.results = numpy.array(ssum, dtype=float)
    else:
        if arg.operands > 2:
            a = numpy.array(arg.values)
            arg.results = a.sum(axis=0)
        else:
            arg.results = arg.values[0] + arg.values[1]

def substract(arg):
    """Substract two operands
    """
    if arg.operands != 2:
        arg.success = False
        arg.errors.append("Two operands required in 'substract'")
    else:
        arg.results = arg.values[0] - arg.values[1]

def move_to_level(arg):
    """Move a variable from initial level to given level (e.g. from stratum
    to stand level)
    """
    if arg.operands != 1:
        arg.success = False
        arg.errors.append("One operand required in 'move_to_level'")
    else:
        arg.results = arg.values[0]

def average(arg):
    """Calculate arithmetic average of a list of operands
    """
    mean = _nanmean if arg.use_nan_funcs else numpy.mean
    if arg.operands == 1:
        vals = arg.values[0]
        tind = arg.oper_target_index[0][:,3:]
        avg = map(lambda x: mean(vals[x[0]:x[1]]), tind)
        arg.results = numpy.array(avg, dtype=float)
    else:
        fst = arg.values[0]
        snd = arg.values[1]
        arg.results = (fst + snd) / 2

def w_average(arg):
    """Calculate weighted average of a list of operands
    """
    su = numpy.nansum if arg.use_nan_funcs else numpy.sum
    if arg.weights[0] is None:
        arg.success = False
        arg.errors.append("Values and weights required in 'w_average'")
    else:
        V = arg.values[0]
        W = arg.weights[0]
        tind = arg.oper_target_index[0][:,3:]
        wavg = map(lambda x: su(V[x[0]:x[1]]*W[x[0]:x[1]]) / su(W[x[0]:x[1]]),
                   tind)
        arg.results = numpy.array(wavg, dtype=float)

def mode(arg):
    """Calculate the mode of a list of operands"""
    arg.success = False
    arg.errors.append("Function 'mode' not implemented yet!")

def grow_percentage(arg):
    """Calculate the absolute growth of a variable with given growth
    percentage"""
    if arg.operands != 2:
        arg.success = False
        arg.errors.append("Two operands required in 'grow_percentage'")
    else:
        fst = arg.values[0]
        snd = arg.values[1]
        arg.results = fst + fst * (0.01 * snd)

def percentage(arg):
    """Calculate the percentage of a variable from another variable
    """
    if arg.operands != 2:
        arg.success = False
        arg.errors.append("Two operands required in 'percentage'")
    else:
        fst = arg.values[0]
        snd = arg.values[1]
        arg.results = fst / snd * 100.0

def divide(arg):
    """Divide the variable value with another variable value
    """
    if arg.operands != 2:
        arg.success = False
        arg.errors.append("Two operands required in 'divide'")
    else:
        fst = arg.values[0]
        snd = arg.values[1]
        nulls = snd == 0
        arg.results = fst / snd
        arg.results[nulls] = 0.0

def multiply(arg):
    """Multiply the variable value with another variable value
    """
    if arg.operands != 2:
        arg.success = False
        arg.errors.append("Two operands required in 'multiply'")
    else:
        fst = arg.values[0]
        snd = arg.values[1]
        arg.results = fst * snd

def sum_products(arg):
    """Sum the products of two operands
    """
    if arg.values[1] is None:
        arg.success = False
        arg.errors.append("two operands required in model 'sum_products'")
    else:
        su = numpy.nansum if arg.use_nan_funcs else numpy.sum
        V1 = arg.values[0]
        V2 = arg.values[1]
        tind = arg.oper_target_index[0][:,3:]
        sumprod = map(lambda x: su(V1[x[0]:x[1]]*V2[x[0]:x[1]]), tind)
        arg.results = numpy.array(sumprod, dtype=float)

def grow_d_iba(arg):
    """Grow diameter with given basal area increment
    """
    arg.success = False
    arg.errors.append("Function 'grow_d_iba' not implemented yet!")

def sum_BA(arg):
    """Sums basal areas of trees and weights them with number of tree
    """
    su = numpy.nansum if arg.use_nan_funcs else numpy.sum
    if arg.weights[0] is None:
        arg.success = False
        arg.errors.append("Values and weights required in 'sum_BA'")
    else:
        D = arg.values[0]
        N = arg.weights[0]
        tind = arg.oper_target_index[0][:,3:]
        pi = math.pi
        BA = map(lambda x: su(N[x[0]:x[1]]*0.25*pi*(D[x[0]:x[1]]*0.01)**2), tind)
        arg.results = numpy.array(BA, dtype=float)

def ba_from_d(arg):
    """Calculate basal area from diameter (at tree level cm->cm2!)
    """
    if arg.operands != 1:
        arg.success = False
        arg.errors.append("Only one operand required in 'ba_from_d'")
    else:
        d = arg.values[0]
        arg.results = 0.25 * math.pi * d ** 2

def d_from_ba(arg):
    """Calculate diameter area from basal area (at tree level cm2->cm!)
    """
    if arg.operands != 1:
        arg.success = False
        arg.errors.append("Only one operand required in 'ba_from_d'")
    else:
        ba = arg.values[0]
        arg.results = (4 * ba / math.pi) ** 0.5

def get_value_from_max_object_with_filter(arg):
    """
    Choose single value for given variable of a given level (2nd operand)
    by finding the maximum value of the 1st operand provided the 3rd operand
    fulfills the limit passed in the 4th operand. Otherwise get values from
    2nd operand with the next highest 1st operand value and so on until the filter
    limit is fulfilled. If several operand values are used, weigh by the
    filter values (3rd operand)
    """
    if arg.values[0] is None :
        arg.success = False
        arg.errors.append("Soring operand (2nd) missing for "\
                          "'get_value_from_max_object_with_filter'")
        return
    V = arg.values[1]
    W = arg.values[0]
    filter_by = arg.values[2]
    filter_limit = arg.values[3][0]
    tind = arg.oper_target_index[0][:,3:]
    arg.results = numpy.empty(tind.shape[0], dtype=float)
    arg.results.fill(numpy.NaN)
    for i in range(tind.shape[0]):
        fr, to = tind[i, 0], tind[i, 1]
        if fr != to:
            so = [[W[j], V[j], j] for j in range(fr, to)]
            so.sort(reverse=True)
            order = [v[2] for v in so]
            limit = 0.
            targets = []
            for f_ind in range(len(order)):
                limit += filter_by[order[f_ind]]
                targets.append(order[f_ind])
                if limit >= filter_limit:
                    break
            if len(targets) > 1:
                weights = filter_by[targets].copy()
                if weights.sum() > 100.:
                    last_weight = filter_limit - filter_by[targets[:-1]].sum()
                    weights[-1] = last_weight
                res = sum(V[targets] * weights) / weights.sum()
                arg.results[i] = res
            else:
                arg.results[i] = V[targets[0]]

def get_value_from_sorted_max_object_with_filter(arg):
    """
    First sorts the data given the 1st sorting variable (1st operand), 2nd
    sorting variable (2rd operand) and the value variable (3rd operand).
    Then chooses single value for the value variable (3rd operand)
    provided the 4th operand fulfills the limit passed in the 5th operand.
    Otherwise get values from 3rd operand with the next highest 2nd operand value
    and so on until the filter limit is fulfilled. If several operand values
    are used, weigh by the filter values (4th operand)
    """
    if arg.values[1] is None :
        arg.success = False
        arg.errors.append("Sorting operand (3rd) missing for "\
                          "'get_value_from_sorted_max_object_with_filter'")
        return
    # 1, 2, 3 for sorting var to sort in 1,2,3 order in reversed sorting
    S = [1/x for x in arg.values[0]]
    V = arg.values[2]
    W = arg.values[1]
    filter_by = arg.values[3]
    filter_limit = arg.values[4][0]
    tind = arg.oper_target_index[0][:,3:]
    arg.results = numpy.empty(tind.shape[0], dtype=float)
    arg.results.fill(numpy.NaN)
    for i in range(tind.shape[0]):
        fr, to = tind[i, 0], tind[i, 1]
        if fr != to:
            so = [[S[j], W[j], V[j], j] for j in range(fr, to)]
            so.sort(reverse=True)
            order = [v[3] for v in so]
            limit = 0.
            targets = []
            for f_ind in range(len(order)):
                limit += filter_by[order[f_ind]]
                targets.append(order[f_ind])
                if limit >= filter_limit:
                    break
            if len(targets) > 1:
                weights = filter_by[targets].copy()
                if weights.sum() > 100.:
                    last_weight = filter_limit - filter_by[targets[:-1]].sum()
                    weights[-1] = last_weight
                res = sum(V[targets] * weights) / weights.sum()
                arg.results[i] = res
            else:
                arg.results[i] = V[targets[0]]

def get_value_from_max_object(arg):
    """Choose single value for given variable of a given level by finding the
    maximum value of the 2nd operand
    """
    amx = _nanargmax if arg.use_nan_funcs else numpy.argmax
    if arg.weights[0] is None:
        arg.success = False
        arg.errors.append("Values and weights required in "\
                          "'get_value_from_max_object'")
    else:
        V = arg.values[0]
        W = arg.weights[0]
        tind = arg.oper_target_index[0][:,3:]
        arg.results = numpy.empty(tind.shape[0], dtype=float)
        arg.results.fill(numpy.NaN)
        for i in range(tind.shape[0]):
            if tind[i,0] != tind[i,1]:
                ind = amx(W[tind[i,0]:tind[i,1]])
                if ind is not None:
                    arg.results[i] = V[tind[i,0]:tind[i,1]][ind]

def get_value_from_min_object(arg):
    """Choose single value for given variable of a given level by finding the
    minimum value of the 2nd operand
    """
    amn = _nanargmin if arg.use_nan_funcs else numpy.argmin
    if arg.weights[0] is None:
        arg.success = False
        arg.errors.append("Values and weights required in "\
                          "'get_value_from_min_object'")
    else:
        V = arg.values[0]
        W = arg.weights[0]
        tind = arg.oper_target_index[0][:,3:]
        arg.results = numpy.empty(tind.shape[0], dtype=float)
        arg.results.fill(numpy.NaN)
        for i in range(tind.shape[0]):
            if tind[i,0] != tind[i,1]:
                ind = amn(W[tind[i,0]:tind[i,1]])
                if ind is not None:
                    arg.results[i] = V[tind[i,0]:tind[i,1]][ind]

def get_value_from_median_object(arg):
    """Choose single value for given variable of a given level by finding the
    median value of the 2nd operand
    """
    if arg.weights[0] is None:
        arg.success = False
        arg.errors.append("Values and weights required in "\
                          "'get_value_from_median_object'")
    else:
        V = arg.values[0]
        W = arg.weights[0]
        tind = arg.oper_target_index[0][:,3:]
        arg.results = numpy.empty(tind.shape[0], dtype=float)
        arg.results.fill(numpy.NaN)
        for i in range(tind.shape[0]):
            if tind[i,0] != tind[i,1]:
                v = V[tind[i,0]:tind[i,1]]
                w = W[tind[i,0]:tind[i,1]]
                n = w.size
                order = numpy.argsort(w)
                v_sorted = v[order]
                if n & 1:
                    # There is an odd number of elements
                    arg.results[i] = v_sorted[n//2]
                else:
                    arg.results[i] = (v_sorted[n//2-1]+v_sorted[n//2])/2

def __compare_second_item(a, b):
    return cmp(a[1], b[1])

def get_value_from_ba_weighted_median_tree(n, data):
    """
    Function: get_value_from_ba_weighted_median_tree

    Get a variable value from the basal area weighted median tree.
    """
    """
    if len(data[0]) != 2:
        return None, None, 'Data values and weights needed!'
    if len(data[0][0]) != len(data[0][1]):
        return None, None, 'Data vectors of unequal length!'
    values = data[0][0]
    weights = data[0][1]
    if values[0].size > 1:
        result = numpy.zeros(len(values))
        numpy.cumsum
        for i in xrange(len(values)):
            gtind = numpy.where(numpy.cumsum(weights[i]) > numpy.sum(weights[i])/2)[0]
            if gtind.size > 0:
                result[i] = values[i][gtind[0]]
    else:
        gtind = numpy.where(numpy.cumsum(weights) > numpy.sum(weights)/2)[0]
        result = values[gtind[0]]
    return result, None
    """
    arg.success = False
    arg.errors.append("Function 'get_value_from_ba_weighted_median_tree' "\
                      "not implemented yet!")

def count(arg):
    """Count the number of child objects
    """
    if arg.operands != 1:
        arg.success = False
        arg.errors.append("Only one operand allowed in 'count'")
    else:
        if arg.use_nan_funcs:
            su = numpy.nansum
            isnan = numpy.isnan
            vals = arg.values[0]
            tind = arg.oper_target_index[0][:,3:]
            cnt = map(lambda x: su(isnan(vals[x[0]:x[1]]) == False), tind)
            arg.results = numpy.array(cnt, dtype=float)
        else:
            start = arg.oper_target_index[0][:,-2]
            end = arg.oper_target_index[0][:,-1]
            arg.results = end - start

def zero_value(arg):
    """Set the values in the data vector to 0
    """
    tind = arg.oper_target_index[0]
    arg.results = numpy.zeros(numpy.size(tind, axis=0), dtype=float)

def assign_value(arg):
    """Assign given value for variable
    """
    if arg.operands != 1:
        arg.success = False
        arg.errors.append("Only one operand required in 'assign_value'")
    else:
        arg.results = arg.values[0]

def rate_of_change_abs(arg):
    """Calculate absolute rate of change
    """
    if arg.operands != 2:
        arg.success = False
        arg.errors.append("Two operands required in 'rate_of_change_abs'")
    else:
        fst = arg.values[0]
        snd = arg.values[1]
        arg.results = snd - fst

def rate_of_change_rel(arg):
    """Calculate relative rate of change
    """
    if arg.operands != 2:
        arg.success = False
        arg.errors.append("Two operands required in 'rate_of_change_rel'")
    else:
        fst = arg.values[0]
        snd = arg.values[1]
        arg.results = (snd - fst) / fst * 100.0

def random_number_0_1(arg):
    """Generate evenly distributed (0..1) random numbers and multiple
    single operand with the random numbers
    """
    if arg.operands != 1:
        arg.success = False
        arg.errors.append("Only one operand required in 'random_number_0_1'")
    elif arg.deterministic:
        N = arg.target_index.shape[0]
        arg.results = arg.values[0] * numpy.array(N*[0.5])
    else:
        N = arg.target_index.shape[0]
        arg.results = arg.values[0] * numpy.random.rand(N)

def cumulative_avg_updown(arg):
    """Calculate the average of operand values, cumulated using weight values
    until a limit defined in second operand

    Cumulation order is from from last to first (updown)
    """
    if arg.operands != 2:
        arg.success = False
        arg.errors.append("Two operands required in 'cumulative_avg_updown'")
    if arg.weights[0] is None:
        arg.success = False
        arg.errors.append("Values and weights required in "\
                          "'cumulative_avg_updown'")
    if arg.success:
        V = arg.values[0]
        W = arg.weights[0]
        limits = arg.values[1]
        tind = arg.oper_target_index[0][:,3:]
        arg.results = numpy.zeros(tind.shape[0], dtype=float)
        for i in range(tind.shape[0]):
            v = V[tind[i,0]:tind[i,1]]
            w = W[tind[i,0]:tind[i,1]]
            n = v.size
            if n == 0:
                continue
            wsum = 0.0
            msum = 0.0
            for j in range(n-1, -1, -1):
                if (w[j] + wsum) >= limits[i]:
                    delta = limits[i] - wsum
                    msum += v[j] * delta
                    wsum += delta
                    break
                else:
                    msum += v[j] * w[j]
                    wsum += w[j]
            arg.results[i] = msum / wsum

def cumulative_avg_downup(arg):
    """Calculate the average of operand values, cumulated using weight values
    until a limit defined in second operand

    Cumulation order is from from first to last (downup)
    """
    if arg.operands != 2:
        arg.success = False
        arg.errors.append("Two operands required in 'cumulative_avg_downup'")
    if arg.weights[0] is None:
        arg.success = False
        arg.errors.append("Values and weights required in "\
                          "'cumulative_avg_downup'")
    if arg.success:
        V = arg.values[0]
        W = arg.weights[0]
        limits = arg.values[1]
        tind = arg.oper_target_index[0][:,3:]
        arg.results = numpy.zeros(tind.shape[0], dtype=float)
        for i in range(tind.shape[0]):
            v = V[tind[i,0]:tind[i,1]]
            w = W[tind[i,0]:tind[i,1]]
            n = v.size
            if n == 0:
                continue
            wsum = 0.0
            msum = 0.0
            for j in range(n):
                if (w[j] + wsum) >= limits[i]:
                    delta = limits[i] - wsum
                    msum += v[j] * delta
                    wsum += delta
                    break
                else:
                    msum += v[j] * w[j]
                    wsum += w[j]
            arg.results[i] = msum / wsum

