import pydicom
import numpy as np
import matplotlib.pyplot as plt
import os
import csv
import math

import wavelet
import network
import keras
import tensorflow as tf
import random



actions = [ lambda img: wavelet.rotate(img, 1)
          , lambda img: wavelet.rotate(img, 2)
          , lambda img: wavelet.rotate(img, 3)
          , lambda img: wavelet.rotate(img, 4)
          , lambda img: wavelet.rotate(img, 5)
          , lambda img: wavelet.rotate(img, 6)
          , lambda img: wavelet.rotate(img, 7)
          ]

#while len(actions) > 4:
#  p = random.randint(0, len(actions) - 1)
#  del actions[p]

genData = []
labels = []
data = []

hSize = 1000
model = network.createModelHistogram(1000)
pos = 0
first = True
with open("mimx.csv") as f:
  for row in csv.reader(f, delimiter=","):
    if first or len(row) < 9:
      first = False
      continue

    n = f"{row[2]}"
    while len(n) < 4:
      n = f"0{n}"

    fileName = f"../Proband {row[0]}/SE00000{row[1]}/{row[0]}_{n}.dcm_blured_histo_diff.npy"


    print(pos, end="\r")
    pos += 1

    x = np.load(fileName, allow_pickle=False)[:hSize]
    y = network.toOneHot([1.0,2.0,3.0,4.0,5.0], float(row[8]))
    if int(row[7]) == 4:
      n = 30
    elif int(row[7]) < 4:
      n = 60
    else:
      n = 10
    for _ in range(n):
      noise = np.array([float(random.randint(0,4)) for _ in x])
      labels.append(y)
      data.append(x + noise)


print()
print(f"Start Train with {len(data)} items.")
data = np.array(data)
labels = np.array(labels)


with open("histogram.log", "wt") as f: 
  bestLoss = None
  for epoch in range(128):
    print(f"Epoch {epoch + 1}")
    h = model.fit(data, labels, epochs=1)
    h = h.history
    loss = h['loss']
    if bestLoss is None or loss < bestLoss:
      bestLoss = loss
      print(f"Update best loss to {bestLoss} and save model")
      f.write(f"Update best loss to {bestLoss} and save model\n")
      network.save(model, "model.keras")
print("done")