import math

import numpy
import pandas
import scipy
import sklearn.decomposition
from django.core import serializers
from matplotlib import pyplot as plt


class ExtendedDataFrame(pandas.DataFrame):
    def heat_plot(self, filename=None):
        plt.clf()
        column_labels = self.index.tolist()
        row_labels = column_labels

        _, ax = plt.subplots()
        _ = plt.gcf()

        # heat map
        _ = ax.pcolor(self, cmap=plt.cm.YlOrRd)

        # put the major ticks at the middle of each cell
        ax.set_xticks(numpy.arange(self.shape[0]) + 0.5, minor=False)
        ax.set_yticks(numpy.arange(self.shape[1]) + 0.5, minor=False)
        # want a more natural, table-like display
        ax.invert_yaxis()
        ax.xaxis.tick_top()
        ax.set_xticklabels(row_labels, minor=False)
        ax.set_yticklabels(column_labels, minor=False)
        plt.xticks(rotation=90)

        if filename:
            plt.savefig(filename)
        else:
            return plt.show()

    def scatter(self, with_labels=False):
        if with_labels:
            fig, ax = plt.subplots()
            self.plot(kind='scatter', x=self.columns[0], y=self.columns[1], ax=ax)
            for k, v in self.iterrows():
                ax.annotate(k, v)
        else:
            self.plot(kind='scatter', x=self.columns[0], y=self.columns[1])

    def get_diagonal(self, n=1):
        mat = numpy.array(self)
        return pandas.Series(numpy.diagonal(mat, n))

    def variation_coefficient(self):
        seq = []
        for _, row in self.T.iteritems():
            r = row[~row.isnull()]
            seq.append(r.std() / r.mean())
        return pandas.Series(seq, index=self.index)

        # return self.T.std() / self.T.mean()

    def ks_test(self, columns=None):
        if not columns:
            columns = self.columns[:2]
        return scipy.stats.ks_2samp(self[columns[0]], self[columns[1]])


def print_cseg(c):
    return '< {} >'.format(' '.join(map(str, c)))


def reduce_dimensions(matrix):
    pca = sklearn.decomposition.PCA(n_components=2)
    pca.fit(matrix)
    return pca.transform(matrix)


def numpy_to_float(number):
    if abs(number) == numpy.inf:
        return None
    else:
        return float(number)


def get_histogram_data(series, bins=6):
    _hist_amount, _hist_range = numpy.histogram(series.fillna(0), bins=bins)
    return list(zip(list(map(float, _hist_range)), list(map(int, _hist_amount))))


def serialize_model(query_set, t='json'):
    data = serializers.serialize(t, query_set.objects.all())
    print(len(data))
    with open(query_set.__name__ + '-serialized.' + t, 'w') as out:
        out.write(data)


def deserialize(filename, save=False):
    t = filename.split('.')[-1]
    with open(filename, 'r') as out:
        data = out.readline()

    x = []
    for obj in serializers.deserialize(t, data):
        if not type(obj.object).objects.filter(id=obj.object.id) and save:
            obj.save()
        else:
            x.append(obj)
    if not save:
        return x


def sample_size(population_size, confidence_level=95, confidence_interval=5):
    # SUPPORTED CONFIDENCE LEVELS: 50%, 68%, 90%, 95%, and 99%
    confidence_level_constant = [50, .67], [68, .99], [90, 1.64], [95, 1.96], [99, 2.57]
    Z = 0.0
    p = 0.5
    e = confidence_interval / 100.0
    N = population_size
    n_0 = 0.0
    n = 0.0

    # LOOP THROUGH SUPPORTED CONFIDENCE LEVELS AND FIND THE NUM STD
    # DEVIATIONS FOR THAT CONFIDENCE LEVEL
    for i in confidence_level_constant:
        if i[0] == confidence_level:
            Z = i[1]

    if Z == 0.0:
        return -1

    # CALC SAMPLE SIZE
    n_0 = ((Z ** 2) * p * (1 - p)) / (e ** 2)

    # ADJUST SAMPLE SIZE FOR FINITE POPULATION
    n = n_0 / (1 + ((n_0 - 1) / float(N)))

    return int(math.ceil(n))  # THE SAMPLE SIZE


def get_first(sequence):
    if not isinstance(sequence, list):
        raise AttributeError('The argument must be of list type')
    return sequence[0]


def get_last(sequence):
    if not isinstance(sequence, list):
        raise AttributeError('The argument must be of list type')
    return sequence[-1]


def get_distinct(queryset, k):
    return [el[k] for el in queryset.values(k).distinct()]


def apply_fn_to_matrix(matrix, fn):
    size = len(matrix.T)
    return numpy.matrix([fn(matrix[:, i]) for i in range(size)])


def round_tolist(matrix, n=2):
    return matrix.round(n).tolist()[0]


def flat_list(seq):
    return [item for sublist in seq for item in sublist]
