import numpy as np
import matplotlib.pyplot as plt
from neuronova_sit import SitA
import pygame
import random
import sys
import os
import time


kontrola_chyb = pygame.init()
if kontrola_chyb[1] > 0:
    print("Nalezeno {0} chyb, konec hry.".format(kontrola_chyb[1]))
    sys.exit(-1)  # ukonceni herniho okna
else:
    print("Start PyGame OK.")  # start hry


# Zakladni nastaveni herniho prostredi
herni_prostredi = pygame.display.set_mode((200, 200)) # nastaveni rozmeru herni plochy
pygame.display.set_caption("Had")
time.sleep(0.5)

# Nastaveni barev v prostredi
cervena = pygame.Color(255, 0, 0)  # konec hry
zelena = pygame.Color(34, 139, 34)  # had
modra = pygame.Color(0,0,255)
cerna = pygame.Color(0, 0, 0)  # skore
fialova = pygame.Color(139, 0, 255)  # jidlo
stribrna = pygame.Color(196, 199, 206)
zlata = pygame.Color(207, 181, 59)  # pozadi

smery = ['UP', 'RIGHT', 'LEFT', 'DOWN']  # 4 mozne smery, kterymi se had muze pohybovat
vyber_smeru = random.SystemRandom()
smer = vyber_smeru.choice(smery)  # nahodny vyber smeru z listu polozek
zmena_smeru = smer  # zmena smeru na novy

# FPS frames per second controller, objekt pro sledovani casu
fpsController = pygame.time.Clock()

# pozice hada na pocatku hry
had_pozice = [100, 50]  # umisteni "hlavy" hada na pocatku hry
had_telo = [[100, 50], [90, 30], [80, 10]]  # umisteni tela hada na pocatku hry

jidlo_pozice = [random.randrange(0, 19) * 10, random.randrange(0, 19) * 10]  # pocatecni pozici jidla vybereme nahodne
nove_jidlo = True  # jidlo umisteno na herni plose
had_knockout = False  # dokud je had nazivu, hodnota Boolean objektu je nastavena na False
skore = 0
epizod_celkem = 0  # pocet epizod mezi jednotlivymi treninky
delka_epizody = 0
testovani = False
zbyva_iteraci = 100000
pocet_ulozenych_zazitku = 0
odmena_epizoda = 0
odmena_jidlo = 4
odmena_smrt = -4
odmena_jinak = -0.001
odmeny = [float(odmena_jidlo), float(odmena_smrt), float(odmena_jinak)]  # nastaveni odmen
batch = 500
nahoru = 0
doprava = 1
doleva = 2
dolu = 3
nahodny_smer = [nahoru, doprava, doleva, dolu]

pocatecni_stav = np.zeros(4)
pocatecni_stav[0] = 1  #definovani smeru prvniho pohybu
pohyb = np.zeros([4])  # docs.scipy: Return a new array of given shape and type, filled with zeros.

from collections import deque
zkusenosti_list = deque()  # kontejner pro skladovani hernich zazitku

# predcasne ukonceni hry
def konec_hry():
    font = pygame.font.SysFont('comicsansms', 30, bold=True)
    konecHry = font.render("Konec hry", True, cervena)
    konecHry2 = konecHry.get_rect()
    konecHry2.midtop = (100, 80)
    herni_prostredi.blit(konecHry, konecHry2)
    skore1(0)
    pygame.display.flip()  # aktualizace GUI
    time.sleep(2)
    pygame.quit()  # ukonceni PyGame
    sys.exit()


# prevod obrazku do sedotonove stupnice
def greyscale_prevod(plocha):
    for x in range(200):
        for y in range(200):
            cervena, zelena, modra, alfa = plocha.get_at((x, y))  # ziskani hodnoty barvy jednotliveho pixelu
            L = 0.3 * cervena + 0.59 * zelena + 0.11 * modra
            greyscale = (L, L, L, alfa)
            plocha.set_at((x, y), greyscale)
    return plocha


# ulozeni obrazku ze hry a predzpracovani pro neuronovou sit
from PIL import Image
from PIL import ImageOps
#from PIL import ImageGrab #Windows
#import pyscreenshot as ImageGrab #Linux
def zaznamenej_stavy():
    foto = pygame.image.tostring(herni_prostredi, 'RGB', False)  # argumenty(Surface, format, flipped)
    obrazek = Image.frombytes('RGB', (200, 200), foto)
    #greyscale_prevod(obrazek)
    obrazek = obrazek.convert('L')
    obrazek = obrazek.resize((70,70))  # prevedeni na stejne rozmery
    obrazek = ImageOps.invert(obrazek)
    obrazek = obrazek.convert('1')
    obrazek_nn = np.asarray(obrazek.getdata(), dtype=np.float64).reshape((obrazek.size[0], obrazek.size[1]))
    obrazek_nn2 = np.asarray(obrazek_nn, dtype=np.uint8)
    obrazek_uloz = Image.fromarray(obrazek_nn2, mode='L')
    obrazek_uloz.save('vstup.jpg')
    return obrazek_nn

# skore hry - v prubehu hry se objevuje v levem hornim rohu, po ukonceni hry celkove skore uprostred herniho okna
def skore1(volba=1):
    skore_font = pygame.font.SysFont('comicsansms', 20, bold=True)
    skore_oznameni = skore_font.render('Skóre : {0}'.format(skore), True, cerna)
    skore_rectangle = skore_oznameni.get_rect()
    if volba == 1:
        skore_rectangle.midtop = (50, 20)
    else:
        skore_rectangle.midtop = (100, 100)
    herni_prostredi.blit(skore_oznameni, skore_rectangle)


# restartovani hry z puvodniho stavu
def restartuj_hru():
    global had_x, had_y, skore, delka_epizody, odmena_epizoda, jidlo_pozice, akce, stav, nasledujici_stav, had_knockout
    had_y = [had_telo[0][0], had_telo[1][0], had_telo[2][0]]
    had_x = [had_telo[0][1], had_telo[1][1], had_telo[2][1]]
    skore = 0
    delka_epizody = 0
    odmena_epizoda = 0
    had_knockout = False
    jidlo_pozice = [random.randrange(0, 19) * 10, random.randrange(0, 19) * 10]
    akce = random.randint(0, 3) # nahodny vyber akce - smeru
    #  pocatecni nastaveni stavu
    stav = [zaznamenej_stavy(), zaznamenej_stavy()]
    nasledujici_stav = [zaznamenej_stavy(), zaznamenej_stavy()]

    #  aktualizace herniho prostredi
    herni_prostredi.fill(zlata)
    for pozice in range(0, len(had_x)):
        herni_prostredi.blit(had, (had_x[pozice], had_y[pozice]))
    pygame.draw.rect(herni_prostredi, fialova, pygame.Rect(jidlo_pozice[0], jidlo_pozice[1], 20, 20))
    #skore1()
    pygame.display.update()

def tisk():
    import csv
    with open('./dataset/epochy_presnost.csv', 'r') as csvfile:
        x = []
        y = []
        i = 0
        soubor = csv.reader(csvfile, delimiter=',')
        for sloupec in soubor:
            x.append(i)
            i+=1
            y.append(sloupec[1])
    plt.plot(x, y)
    plt.yticks([])
    plt.xlabel('Počet epoch')
    plt.ylabel('Přesnost')
    plt.title('Graf závislosti počtu epoch na přesnosti')
    plt.show()

def ukoncit_hru():
    global zaznam, zbyva_iteraci, skore, delka_epizody, odmena_epizoda, \
        testovani, zazitky_pamet, pocet_ulozenych_zazitku, epizod_celkem
    epizod_celkem += 1
    testovani = True
    if testovani:
        # ulozeni informaci o aktualni epizode
        zaznam.do_csv('data_testovani.csv', [skore, epizod_celkem, delka_epizody, "%.3f" % odmena_epizoda])
        print('Skore epizody: %s Epizod celkem: %s Delka epizody: %s Odmena za epizodu: %s' % (skore, epizod_celkem,
                                                                            delka_epizody, "%.3f" % odmena_epizoda))
    testovani = False

    if skore >= 1:
        pocet_ulozenych_zazitku += len(zazitky_pamet)
        print('Pocet ulozenych vzorku: %s Skore epizody: %s Delka epizody: %s' % (pocet_ulozenych_zazitku, skore, delka_epizody))
        # ulozeni uspesnych epizod do datasetu
        zaznam.do_csv('data_trenovani.csv', [skore, epizod_celkem, delka_epizody, odmena_epizoda])
        for stav, akce, odmena, nasledujici_stav, ukonceni in zazitky_pamet:
            agent.uloz_zazitek(stav, akce, odmena, nasledujici_stav, ukonceni)

    # trenovani neuronove_site
    if len(zkusenosti_list) >= batch:
        pocet_ulozenych_zazitku = 0
        print('Celkem epizod: %d' % epizod_celkem)
        epizod_celkem = 0
        if zbyva_iteraci == 0:
            nn.save()
            sys.exit(0)
        agent.trenovani()

        zbyva_iteraci -= 1 if zbyva_iteraci != -1 else 0
        #po treninku zacne testovani
        testovani = True

    zazitky_pamet = [] #list pro ukladani prozitych stavu do pameti
    #skore1()
    pygame.display.update()
    restartuj_hru() #restartovani epizody


class Zapis:
    def __init__(self):
        self.path = 'dataset/'
        if not os.path.exists(self.path):
            os.makedirs(self.path)  # vytvoreni adresare pro ukladani datasetu
        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()

zaznam = Zapis()
zazitky_pamet = list()

had = pygame.Surface((20,20))
had.fill(zelena)
had_y = [had_telo[0][0], had_telo[1][0], had_telo[2][0]]
had_x = [had_telo[0][1], had_telo[1][1], had_telo[2][1]]
# nahodny vyber akce - smeru
akce = random.randint(0, 3)
stav = [zaznamenej_stavy(), zaznamenej_stavy()]
nasledujici_stav = [zaznamenej_stavy(), zaznamenej_stavy()]
nn = SitA()


# zahajeni trenovani z prikazoveho radku
import argparse
parser = argparse.ArgumentParser(description='Zahajeni trenovani z prikazoveho radku.')
parser.add_argument('-trenink', action='store_true', help='trenovani agenta')
parser.add_argument('-m', '--mode', help='Trenink / Hra', required=True)
parser.add_argument('-l', '--model', type=str, help='Spusteni predtrenovaneho modelu.', required=False)
args = vars(parser.parse_args())
if args['model'] == 'sitA_model.h5':
    from keras.models import load_model
    print("Model nacten")
    load_model('sitA_model.h5')


class Agent:
    def __init__(self):
        self.pocet_treninku = 0
        if args['mode'] == 'Hra':  # mod Hra
            print("Zvolen herni mod")
            pozorovani = 2000000
            self.epsilon = 0.0001
            print("Epsilon: ", self.epsilon)
            print("Nacteni vah ulozeneho modelu")
            nn.load("sitA_vahy.h5")  # nacteni vah predtrenovaneho modelu
        else:  # mod Trenink
            print("Zvolen treninkovy mod")
            pozorovani = 3000
            self.epsilon = 0.9999
            print("Epsilon: ", self.epsilon)
        # strategie epsilon-greedy, vyssi hodnota znamena vyssi pravdepodobnost volby nahodne akce a podporu explorace
        # na ukor exploatace

    def trenovani(self):
        if self.epsilon >= 0:   # snizeni hodnoty epsilon povede ke snizeni explorace a zvyseni exploatace
            self.epsilon -= 0.001
        else:
            epsilon = 0.0001
        self.pocet_treninku += 1
        print('Pocet treninku: ', self.pocet_treninku, ' Epsilon:', self.epsilon)
        batch1 = self.vzorky_trenink()
        nn.trenovani(batch1) # trenink site ze ziskanych vzorku


    def zvol_akci(self, stav):
        pohyb = np.zeros([4])
        if (random.random() <= self.epsilon):
            #print("Zvolena nahodna akce.")
            akce = random.choice(nahodny_smer)  # nahodny vyber z listu moznych akci [0=nahoru, 1=doprava, 2=doleva, 3=dolu]
            pohyb[akce] = 1
            return akce
        else:
            #print("Zvolena akce s maximalni Q-hodnotou")
            Q_hodnoty = nn.predict(stav)
            maxQ = np.argmax(Q_hodnoty)  # vybere maximalni Q-hodnotu z moznych a priradi akci
            akce = maxQ
            pohyb[akce] = 1
            return akce


    def uloz_zazitek(self, stav, akce, odmena, nasledujici_stav, ukonceni):
        zkusenosti_list.append({'stav': stav, 'akce': akce, 'odmena': odmena, 'nasledujici_stav': nasledujici_stav,
                                'ukonceni': ukonceni})


    def vzorky_trenink(self):
        vzorky = []
        for i in range(batch):
            polozka = zkusenosti_list.popleft()  # vyradi se zazitek z leve casti deque
            vzorky.append(polozka)
        return np.asarray(vzorky)

agent = Agent()
pocet_kroku = 0
pruzkum = 2000000
pozorovani = 3000
# informace o stavu, v jakem se agent nachazi
stavA = ""
if pocet_kroku <= pozorovani:
    stavA = "pozorovani"
elif pocet_kroku > pozorovani and pocet_kroku <= pozorovani + pruzkum:
    stavA = "pruzkum"
else:
    stavA = "trenovani"
print("Agent se nachazi ve stavu: " + stavA)
# nekonecna smycka hry
while True:
    delka_epizody += 1
    odmena = odmeny[2]  # odmena za pohyb prostredim
    nasledujici_stav[0] = stav[1]
    for udalost in pygame.event.get():  # prochazeni udalosti
        if udalost.type == pygame.QUIT:
            nn.save()
            konec_hry()
            sys.exit()

    # (0...nahoru, 1...doprava, 2...doleva, 3...dolu)
    # potvrzeni spravnosti smeru: pokud nahoru => nesmi dolu, pokud doprava => nesmi doleva
    smery = ['UP', 'RIGHT', 'LEFT', 'DOWN']
    vyber_smeru = random.SystemRandom()
    smer = vyber_smeru.choice(smery)   # nahodny vyber polozky z listu
    if vyber_smeru == 'RIGHT' and not smer == 'LEFT':  # pokud jde had doprava, nesmi doleva
        smer = 'RIGHT'
    if vyber_smeru == 'LEFT' and not smer == 'RIGHT':  # pokud jde had doleva, nesmi doprava
        smer = 'LEFT'
    if vyber_smeru == 'UP' and not smer == 'DOWN':  # pokud jde had nahoru, nesmi dolu
        smer = 'UP'
    if vyber_smeru == 'DOWN' and not smer == 'UP':  # pokud jde had dolu, nesmi nahoru
        smer = 'DOWN'

    if smer == 'UP':
        zmena_smeru = 'UP'
    elif smer == 'RIGHT':
        zmena_smeru = 'RIGHT'
    elif smer == 'LEFT':
        zmena_smeru = 'LEFT'
    else:
        zmena_smeru = 'DOWN'

    i = len(had_x) - 1
    while i >= 1:
        had_x[i] = had_x[i - 1]
        had_y[i] = had_y[i - 1]
        i -= 1

    if smer == 'RIGHT':
        had_x[0] += 20
    if smer == 'LEFT':
        had_x[0] -= 20
    if smer == 'UP':
        had_y[0] -= 20
    if smer == 'DOWN':
        had_y[0] += 20

    had_telo.insert(0, list(had_pozice))
    if had_x[0] == jidlo_pozice[0] and had_y[0] == jidlo_pozice[1]:
        # had nasel jidlo a zvetsi se jeho delka
        had_x.append(20)
        had_y.append(20)
        skore += 1
        nove_jidlo = False
        odmena = odmeny[0]
    else:
        had_telo.pop()


    if nove_jidlo is False:
        jidlo_pozice = [random.randrange(0, 19) * 10, random.randrange(0, 19) * 10]  # nahodne generovano nove jidlo
    nove_jidlo = True

    # moznosti ukonceni hry
    if had_x[0] > 190 or had_x[0] < 0:  # had narazi do zdi vpravo nebo vlevo
        time.sleep(0.1)
        had_pozice = [100, 50]
        jidlo_pozice = [random.randrange(0, 19) * 10, random.randrange(0, 19) * 10]
        odmena = odmeny[1]
        had_knockout = True

    if had_y[0] > 190 or had_y[0] < 0:  # had narazi do zdi nahore nebo dole
        time.sleep(0.1)
        had_pozice = [100, 50]
        jidlo_pozice = [random.randrange(0, 19) * 10, random.randrange(0, 19) * 10]
        odmena = odmeny[1]
        had_knockout = True

    for cast in had_telo[1:]:
        if had_x[0] == cast[0] and had_y[0] == cast[1]:  # pokud had narazi sam do sebe
            time.sleep(0.1)
            had_pozice = [100, 50]
            jidlo_pozice = [random.randrange(0, 19) * 10, random.randrange(0, 19) * 10]
            odmena = odmeny[1]
            had_knockout = True

    odmena_epizoda += odmena  # pricteni jednotlivych odmen k celkove odmene za celou epizodu
    herni_prostredi.fill(zlata)
    #skore1()
    pygame.display.flip()

    for pozice in range(0, len(had_x)):
        herni_prostredi.blit(had, (had_x[pozice], had_y[pozice]))
    pygame.draw.rect(herni_prostredi, fialova, pygame.Rect(jidlo_pozice[0], jidlo_pozice[1], 20, 20))
    pygame.display.update()

    # aktualizace stavu
    nasledujici_stav[1] = zaznamenej_stavy()

    if had_knockout is True:
        ukonceni = True
    else:
        ukonceni = False

    zazitky_pamet.append((np.asarray([stav]), akce, odmena, np.asarray([nasledujici_stav]), ukonceni))
    # ulozeni zkusenosti do pameti ve forme (aktualni stav, akce, odmena, nasledujici stav, ukonceni hry)
    pocet_kroku += 1


    # agent zvoli nasledujici akci z listu moznych
    akce = agent.zvol_akci(np.asarray([list(nasledujici_stav)]))

    # za jakych podminek bude hra ukoncena
    if had_knockout is True:
        ukoncit_hru()
    elif delka_epizody > 1500:
        ukoncit_hru()