import pickle
import theano
import theano.tensor as T
from pylearn2.models import mlp
from pylearn2.training_algorithms import sgd
from pylearn2.training_algorithms.learning_rule import AdaGrad
from pylearn2.training_algorithms.learning_rule import Momentum
from pylearn2.training_algorithms.sgd import ExponentialDecay
from pylearn2.training_algorithms.sgd import LinearDecay
from pylearn2.training_algorithms.learning_rule import AdaDelta

from pylearn2.termination_criteria import EpochCounter
from pylearn2.datasets.dense_design_matrix import DenseDesignMatrix
from pylearn2.costs.cost import Cost
from pylearn2.costs.cost import DefaultDataSpecsMixin
import numpy as np
from random import randint
import os
import matplotlib.pyplot as plt
from features import logfbank


class XOR(DenseDesignMatrix):
    def __init__(self):
        self.class_names = ['0', '1']
        X = [[randint(0, 1), randint(0, 1)] for _ in range(1000)]
        y = []
        for a, b in X:
            if a + b == 1:
                y.append([0, 1])
            else:
                y.append([1, 0])
        X = np.array(X)
        y = np.array(y)
        super(XOR, self).__init__(X=X, y=y)


class NegativeLogLikelihoodCost(DefaultDataSpecsMixin, Cost):
    supervised = True

    def expr(self, model, data, **kwargs):
        space, source = self.get_data_specs(model)
        space.validate(data)

        inputs, targets = data

        outputs = model.fprop(inputs)

        loss = -(targets * T.log(outputs)).sum(axis=1)

        return loss.mean()

left_context = 10
right_context = 5
keywords = ["she", "had"]

def getDataForFrames(data, X, y):
    fbanks = data[0]
    segments = data[1]
    feat_cnt = len(fbanks[0])
    for start, end, class_id in segments:

        if randint(0, 1) == 0:
            for i in xrange(start, end):
                features = np.empty(0)
                for j in xrange(-left_context, right_context):
                    if i + j < 0 or i + j >= len(fbanks):
                        features = np.append(features, np.zeros(feat_cnt))
                    else:
                        features = np.append(features, fbanks[i + j])
                result = np.zeros(len(keywords) + 1)
                result[class_id] = 1.0

                X.append(features)
                y.append(result)

        if randint(0, 1) == 0:
            for i in xrange(start, end):
                features = np.empty(0)
                for j in xrange(-left_context, right_context):
                    if i + j < start or i + j >= end:
                        features = np.append(features, np.zeros(feat_cnt))
                    else:
                        features = np.append(features, fbanks[i + j])
                result = np.zeros(len(keywords) + 1)
                result[class_id] = 1.0

                X.append(features)
                y.append(result)


def getWindowedFeats(fbanks):
    X = []
    feat_cnt = len(fbanks[0])
    for i in xrange(len(fbanks)):
        features = np.empty(0)
        for j in xrange(-left_context, right_context):
            if i + j < 0 or i + j >= len(fbanks):
                features = np.append(features, np.zeros(feat_cnt))
            else:
                features = np.append(features, fbanks[i + j])
        X.append(features)
    return X

def getDataFromPath(path, maxFiles = 1e9):
    X = []
    y = []
    for file in os.listdir(path):
        maxFiles -= 1
        if maxFiles < 0:
            break
        data = np.load(path + "/" + file)
        getDataForFrames(data, X, y)

    X = np.array(X)
    y = np.array(y)

    return X, y

class SHEHAD(DenseDesignMatrix):
    feat_cnt = 0

    def __init__(self, path):
        self.class_names = ["she", "had", "filler"]
        X, y = getDataFromPath(path)
        # print X.shape
        self.feat_cnt = len(X[0])
        super(SHEHAD, self).__init__(X=X, y=y)


def test(model):
    X, y = getDataFromPath("/Users/evgeny/data/TEST")

    confusion = np.zeros([3, 3])

    ypred = np.log(model.fprop(theano.shared(np.array(X), name='inputs')).eval())

    cnt = np.zeros(3)


    for a, b in zip(ypred, y):
        pos = np.argmax(b)
        i = np.argmax(a)
        #for i in xrange(3):
        confusion[pos][i] += 1
        # cnt[pos] += 1

    for i in xrange(3):
        for j in xrange(3):
            print "%.0f" % confusion[i][j],