# advanced.py
# by James Fulford
# put advanced methods here


from _basic import *
import copy

# categorical data


def bucket(dataset, acs=str):
    """
    Partitions dataset into buckets
    defined by unique results of acs function.
    """
    final = {}
    for entry in dataset:
        buck = acs(entry)
        if(buck not in final.keys()):
            final[buck] = []
        final[buck].append(copy.deepcopy(entry))
    return final


def chi(dataset, partition_fn, filter_fn, significance=0.05):
    """Returns chi squared statistic of the subset generated by filter_fn.
    Categories are sorted out by partition_fn.
    (Generates expected values from full dataset.)
    """
    # Calculate proportions of each partition to full dataset
    freq = frequency(dataset, acs=partition_fn)
    cnt = count(dataset)
    expect = {}  # will hold expected proportion for each partition.
    for key in freq.keys():
        expect[key] = freq[key] / float(cnt)

    buckets = bucket(dataset, acs=partition_fn)
    subset = {}
    for key in buckets.keys():
        subset[key] = filter(filter_fn, buckets[key])

    # gotta make sure everything is still okay
    assert set(expect.keys()) is set(buckets.keys())

    # start calculations
    sums = []
    for key in subset.keys():  # or expect.keys(); we know they are the same.
        value = count(subset[key])  # frequency of key in subset
        if(value < 5):
            # if too few entries in a partition,
            # chi square is not effective anymore.
            message = "Chi-square Error: " + key
            message += " has too few entries in the subset. ( < 5 )"
            raise Exception(message)
        sums.append((value - expect[key]) ** 2)
    chi2_stat = summation(sums)

    # Hypothesis testing
    critical_value = scipy.stats.chi2.isf(significance, df)
    if critical_value < chi2_stat:
        print("")
    return summation(sums)


def df(buckets):
    """Returns the degrees of freedom of the dataset.
    number of observations (entries) - number of parameters (fields)
    the number of values that are free to vary
    equation: sum(deviations(dataset)) = 0
    remember linear algebra: one equation, n variables: n-1 are free variables!
    https://www.youtube.com/watch?v=VIlVWeUQ0vs
    """
    return (count(buckets) - 1)


def correlation(datasets, labels=None):
    """Finds linear correlation coefficient between
    each of the provided datasets (must supply at least 2)
    If labels are provided, will associate by index.
    """
    pass
