import tensorflow as tf

X = dane_x[0].shape[0]
Y = dane_x[0].shape[1]
N = X*Y

def splot(obraz, waga, bias):
    return tf.nn.relu( tf.nn.bias_add( tf.nn.conv2d(obraz, waga, strides=[1,1,1,1], padding="SAME"), bias ) )

def max_pool(obraz, k):
    return tf.nn.max_pool( obraz, ksize=[1,k,k,1], strides=[1,k,k,1], padding="SAME")

def zmienna_wagowa(rozmiar, sd):
    z = tf.truncated_normal(rozmiar, stddev=sd)
    return tf.Variable(z)

def zmienna_przesuniecia(rozmiar):
    z = tf.constant(0, shape=rozmiar, dtype=tf.float32)
    return tf.Variable(z)

def wagi(filtr_x, filtr_y, k, warstwa):

    wagi = {
        'splot_1': zmienna_wagowa([filtr_x, filtr_y, 1, 32], 0.1),
        'splot_2': zmienna_wagowa([filtr_x, filtr_y, 32, 64], 0.1),
        'polaczenie': zmienna_wagowa([ (X // k )* (Y // k )*64, warstwa ], 0.1),
        'wyjscie': zmienna_wagowa([warstwa, 2], 0.1)
    }
    return wagi

def przesuniecia(warstwa):
    przes = {
        'splot_1': zmienna_przesuniecia([32]),
        'splot_2': zmienna_przesuniecia([64]),
        'polaczenie': zmienna_przesuniecia([warstwa]),
        'wyjscie': zmienna_przesuniecia([2])
    }
    return przes

def siec(obraz, filtr_x, filtr_y, k, wielkosc_warstwy, dp):


    obraz = tf.reshape(obraz, shape=[-1, X, Y, 1])
    S1 = splot(obraz, wagi(3,3,2,512)['splot_1'], przesuniecia(512)['splot_1'])
    S1 = max_pool(S1, k)
    S1 = tf.nn.dropout(S1, dp )

    S2 = splot(S1, wagi(3,3,2,512)['splot_2'], przesuniecia(512)['splot_2'])
    S2 = max_pool(S2, k)
    S2 = tf.nn.dropout(S2, dp)

    pol = tf.reshape(S2, [-1, wagi(3,3,2,512)['polaczenie'].get_shape().as_list()[0] ])
    pol = tf.nn.relu( tf.nn.bias_add( tf.matmul( pol, wagi(3,3,2,512)['polaczenie'] ),
                                     przesuniecia(512)['polaczenie']) )
    pol = tf.nn.dropout(pol, dp)

    wyjscie = tf.add( tf.matmul( pol, wagi(3,3,2,512)['wyjscie']), przesuniecia(512)['wyjscie'])
    return wyjscie

def podaj_dane(zbiornik_X, zbiornik_Y, rozmiar_probki, krok):

    start = krok*rozmiar_probki % ( len(zbiornik_X) - rozmiar_probki )
    batchX = zbiornik_X[start:start+rozmiar_probki]
    batchY = zbiornik_Y[start:start+rozmiar_probki]

    return batchX, batchY

learning_rate = 1e-3
rozmiar_probki = 20
liczba_iteracji = int(10e4)
x = tf.placeholder(tf.float32, [None, X, Y])
y = tf.placeholder(tf.float32, [None, 2])
keep_prob = tf.placeholder(tf.float32)
dropout = 0.75

co_ile_wyswietlamy = 1200

przewidywania = siec(x, 3, 3, 2, 512, keep_prob)

koszt = tf.reduce_mean( tf.nn.softmax_cross_entropy_with_logits(przewidywania, y))
optymalizacja = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(koszt)

init = tf.initialize_all_variables()

czy_dobrze = tf.equal(tf.argmax(przewidywania,1), tf.argmax(y,1))
dokladnosc = tf.reduce_mean( tf.cast( czy_dobrze, tf.float32) )

with tf.Session() as sess:
    sess.run(init)
    krok = 0

    while krok*rozmiar_probki < liczba_iteracji:
        batch_X, batch_Y = podaj_dane(dane_x, dane_y, rozmiar_probki, krok)
        sess.run(optymalizacja, feed_dict={x: batch_X, y: batch_Y, keep_prob: dropout})

        if krok % co_ile_wyswietlamy:
            dokl = sess.run(dokladnosc, feed_dict={x:batch_X, y:batch_Y, keep_prob: 1.})
            strata = sess.run(kosz, feed_dict={x:batch_X, y:batch_Y, keep_prob: 1.})

            print("Iteracja nr " + str(krok*rozmiar_probki) + ", dokladnosc to {:.3f}".format(dokl), \
                 ", strata to {:.3f}".format(strata))

        krok += 1