from math import log2
from random import sample


class Doubles:
    EPS, INF = 1e-6, 1e9

    @classmethod
    def smaller(cls, value1, value2):
        return value1 + cls.EPS < value2

    @classmethod
    def equals(cls, value1, value2):
        return value1 - cls.EPS < value2 and value1 + cls.EPS > value2

    @classmethod
    def isZero(cls, value):
        return cls.equals(value, 0.0)

    @classmethod
    def isNegative(cls, value):
        return cls.smaller(value, 0.0)

    @classmethod
    def minusInfinity(cls):
        return -cls.INF

    @classmethod
    def plusInfinity(cls):
        return +cls.INF


def separateData(data, ratio=0.8, debug=False):
    first, second = [], []
    selected = sample(range(len(data)), int(ratio * (len(data) - 1)))

    if debug:
        print('dataSize: {:d}'.format(len(data)))
        print('proportion: {:f}; firstSize: {:d}'.format(ratio, len(selected)))

    for index in range(len(data)):
        if index in selected:
            first.append(data[index])
        else:
            second.append(data[index])

    return first, second


def waitForEnter(debug=False):
    while debug and input():
        pass


def separateClasses(data):
    return [value[:-1] for value in data], [value[-1] for value in data]


def attachClasses(data, dataClasses):
    assert len(data) == len(dataClasses)
    return [data[index] + [dataClasses[index]] for index in range(len(data))]


def informationGain(node, attribute, trainingSet, trainingSetClasses,
                    featureCount, classCount, debug=False):
    if not node.subset:
        return 0.0

    result = entropy(node.subset, trainingSetClasses, classCount)

    if debug:
        print('Labels:', end='')
        for index in node.subset:
            print(' {:d}'.format(trainingSetClasses[index]), end='')
        print('\nattribute: {:d}'.format(attribute), end='; ')

    a, b, median = divideSet(node, trainingSet, attribute)

    if a:
        proportion = len(a) / len(node.subset)
        result -= entropy(a, trainingSetClasses, classCount) * proportion

    if b:
        proportion = len(b) / len(node.subset)
        result -= entropy(b, trainingSetClasses, classCount) * proportion

    """
    if debug:
        print('size(a) = {:d}; size(b) = {:d}; total = {:d};\tI = {:f}'.
              format(len(a), len(b), len(node.subset), result))
        waitForEnter(debug)
    """

    if Doubles.isNegative(result):
        print('NEGATIVE information gain!!!')
        exit()

    return result


def entropy(subset, trainingSetClasses, classCount, debug=False):
    classesPresent = [trainingSetClasses[element] for element in subset]
    classSizes = [classesPresent.count(c) for c in range(classCount)]
    result = 0.0

    """
    if debug:
        print('Labels:', end='')
        for index in subset:
            print(' {:d}'.format(trainingSetClasses[index]), end='')
        print('\n\n# subset indices = {:d}'.format(len(subset)))
    """

    for c in range(classCount):
        if classSizes[c]:
            p_c = classSizes[c] / len(subset)
            result -= p_c * log2(p_c)

            if debug:
                print('{:d} {:d}, {:f}'.format(c, classSizes[c], p_c))

    if Doubles.isNegative(result):
        print('NEGATIVE entropy!!!')
        exit()

    """
    if debug:
        print('H = {:f}'.format(result))
        waitForEnter(debug)
    """

    return result


def divideSet(node, trainingSet, attribute, debug=False):
    a, b = [], []
    values = [trainingSet[element][attribute] for element in node.subset]
    median = sorted(values)[(len(values) - 1) // 2]

    for element in node.subset:
        if Doubles.smaller(trainingSet[element][attribute], median):
            a.append(element)
        else:
            b.append(element)

    if debug:
        print('attribute = {:d}; values: '.format(attribute), end='')
        for value in values:
            print(str(value), end=' ')
        print()

        print('a: ', end='')
        for v in a:
            print('{:d}({:d})'.format(v, trainingSet[v][attribute]), end=', ')
        print('\na.size() = {:d}'.format(len(a)))

        print('b: ', end='')
        for v in b:
            print('{:d}({:d})'.format(v, trainingSet[v][attribute]), end=', ')
        print('\nb.size() = {:d}'.format(len(b)))

    return a, b, median
