from __future__ import print_function
import numpy as np
import matplotlib.pyplot as plt
import pygame
import random
import sys
import os
import time
import keras
from keras.layers import *
from keras.utils import np_utils
from keras.optimizers import *
from keras import optimizers
from keras.models import Sequential
from keras.models import model_from_json
from keras.layers import Dense, Flatten, Dropout, Activation, Conv2D
from keras.layers.normalization import BatchNormalization
from keras.initializers import normal, identity
from keras.utils import plot_model
from keras.models import load_model
import pydot
from keras.callbacks import TensorBoard
import csv
import json

seed = 5
np.random.seed(seed) # zaruka stejnych vysledku

diskontni_faktor = 0.8
epochy = 1
batch_size = 32  #velikost minibatche
obrazek_x, obrazek_y = 70, 70
obrazek = 2  # kanal
input_shape = (obrazek, obrazek_x, obrazek_y)


class LossHistory(keras.callbacks.Callback):
    def on_train_begin(self, logs={}):
        self.losses = []

    def on_batch_end(self, batch, logs={}):
        self.losses.append(logs.get('loss'))

history = LossHistory()


class SitA:
    def __init__(self):
        # architektura neuronove site
        self.model = Sequential()
        self.model.add(Conv2D(32, kernel_size=(8, 8), strides=(4, 4), padding='valid', data_format='channels_first',
                              activation='relu', input_shape=input_shape))  # 2*70*70
        self.model.add(BatchNormalization(axis=1, input_shape=input_shape))
        self.model.add(Conv2D(64, kernel_size=(4, 4), strides=(2, 2), padding='valid', data_format='channels_first',
                              activation='relu', input_shape=input_shape))
        self.model.add(BatchNormalization(axis=1))
        self.model.add(Conv2D(64, kernel_size=(2, 2), strides=(2, 2), padding='valid', data_format='channels_first',
                              activation='relu', input_shape=input_shape))
        self.model.add(Flatten())
        self.model.add(Dropout(0.15))
        self.model.add(Dense(512, bias_initializer='zeros', activation='relu'))
        self.model.add(Dense(4))

        from keras.utils import plot_model
        plot_model(self.model, to_file='model_SitA.png', show_shapes=True)  # vizualizace modelu SitA

        sgd = optimizers.SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)
        self.model.compile(loss='mse', optimizer=sgd, metrics=['accuracy'])

    def trenovani(self, batch):
        x_data = list()
        y_data = list()
        for udaj in batch:
            x_data.append(udaj['stav'].astype(np.float64))
            odhad_dalsiho_stavu = self.predict(udaj['nasledujici_stav'].astype(np.float64)).ravel()
            nova_q_hodnota = np.max(odhad_dalsiho_stavu)

            list1 = list(self.predict(udaj['stav'])[0])
            if udaj['ukonceni']:
                list1[udaj['akce']] = udaj['odmena']
            else:
                list1[udaj['akce']] = udaj['odmena'] + diskontni_faktor * nova_q_hodnota
            y_data.append(list1)

        x_data = np.asarray(x_data).squeeze()  # odstraneni 1D polozek => zmenseni pole
        y_data = np.asarray(y_data).squeeze()
        zaznam_vyvoje = self.model.fit(x_data, y_data, batch_size=batch_size, epochs=epochy, verbose=1,
                                       validation_split=0.1, callbacks=[history])

        tisk.do_csv('ztraty.csv', history.losses)  # zaznam vyvoje ztrat
        tisk.do_csv('epochy_presnost.csv', (epochy, zaznam_vyvoje.history['acc']))


        #print(odhad_dalsiho_stavu)

        # graf vyvoje presnosti
        plt.plot(zaznam_vyvoje.history['acc'])
        plt.plot(zaznam_vyvoje.history['val_acc'])
        plt.title('Presnost modelu')
        plt.ylabel('Presnost')
        plt.xlabel('Epochy')
        plt.legend(['trenink', 'testovani'], loc='upper right')
        #plt.show()

        # graf vyvoje ztrat
        plt.plot(zaznam_vyvoje.history['loss'])
        plt.plot(zaznam_vyvoje.history['val_loss'])
        plt.title('Ztraty modelu')
        plt.ylabel('Ztrata')
        plt.xlabel('Epochy')
        plt.legend(['trenink', 'testovani'], loc='upper right')
        #plt.show()


    def predict(self, stav):
        return self.model.predict(stav.astype(np.float64), batch_size=1)

    def save(self):
        from keras.models import model_from_yaml
        yaml = self.model.to_yaml()  # ulozeni pouze architektury modelu
        with open("model.yaml", "w") as yaml_file:
            yaml_file.write(yaml)
        self.model.save('sitA_model.h5')  # ulozeni modelu vcetne vah a konfigurace
        self.model.save_weights('sitA_vahy.h5', overwrite=True)
        with open("model.json", "w") as outfile:
            json.dump(self.model.to_json(), outfile)

    def load(self, path):
        self.model.load_weights(path)


class Soubor:
    def __init__(self):
        self.path = 'dataset/'
        if not os.path.exists(self.path):
            os.makedirs(self.path)
        self.save_file = self.path

    def do_csv(self, nazev_souboru, radek):
        try:
            soubor = open(self.path + nazev_souboru, 'a')
        except IOError:
            return None
        string = ','.join([str(hodnota) for hodnota in radek])
        string = string + '\n' if not string.endswith('\n') else ''
        soubor.write(string)
        soubor.close()
tisk = Soubor()