"""A module for basic statistics

Available functions:
- average
- std
- rms
- median
- print_statistics

"""

import math
from decimal import Decimal

def average(sample):
    return float(sum(sample)) / len(sample)


def squared(sample):
    return [x**2 for x in sample]


def std(sample):
    s = Decimal(str(average(squared(sample)))) - Decimal(str(average(sample)**2))
    s = float(s.sqrt())
    return s


# The root mean square (quadratic mean)
def rms(sample):
    return math.sqrt(average(squared(sample)))


def median(sample):
    sample = sorted(sample)
    N = len(sample)
    half = (N-1)/2
    # Even
    if N % 2 == 0:
        return 0.5*(sample[half] + sample[half+1])
    # Odd
    else:
        return float(sample[half])


def differential(y):
	return map(lambda b, a: b - a, y[1:], y[0:-1])


def filter_columns(columns, n_sigma):
    """Filter iteratively columns by keeping only the elements of columns of
    index i inferior to n_sigma while the convergence criterion is not
    satisfied.

    Usage 
    -----
    filtered_columns = filter_columns(columns, i, n_sigma)

    Parameters
    ----------
    columns : a list of tuples [(col1_1, col2_1),
                                (col1_2, col2_2),
                                ...]
            the first column is the column of reference for the comparison

    n_sigma : a number
            the dispersion in which the new sample must be determined

    Returns
    -------
    filtered_columns : a list of tuples
            the median (that converged if iteration < max_iteration,
            not converged otherwise)
    
    How it works
    ------------
    1. The median and the standard deviation are computed
    2. A new sample is determined by taking only values inferior to n_sigma
    3. If this new sample is the same as the original sample, the median is returned
        (the median converged)
    4. If the original sample and the new sample didn't converge at the max_iteration
        return the median (no convergence)

    """
    max_iteration = 100

    for i in range(1, max_iteration+1):
        sample = [x[0] for x in columns]
        median_value = median(sample)
        std_value = std(sample)

        new_columns = [col for col in columns if col[0] <= median_value + n_sigma * std_value]

        if columns != new_columns:
            columns = new_columns[:]
        elif i == max_iteration:
            # print("No convergence before iteration %d" % max_iteration)
            break 
        else:
            # print("Convergence at iteration %d" % i)
            break
    return new_columns


def filter_sample(pairs_of_columns, n_sigma):
    filtered_columns = filter_columns(pairs_of_columns, n_sigma)
    filtered_sample = [x[1] for x in filtered_columns]
    return filtered_sample


def compute_statistics(column, sample):
    statistics = {"column": column,
                 "average": average(sample),
                  "median": median(sample),
                     "std": std(sample)}
    return statistics


def print_statistics(column, sample):
    stats = compute_statistics(column, sample)
    # output = (stats["column"], stats["average"], stats["median"], stats["std"])
    output = stats["column"].ljust(22)
    output += "%.5g\t%.5g\t%.5g" % (stats["average"], stats["median"], stats["std"])
    print(output)
    return output


def print_statistics_after_filter(column, pairs_of_columns, n_sigma):
    filtered_sample = filter_sample(pairs_of_columns, n_sigma)
    output = print_statistics(column, filtered_sample)
    return output

