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

import wavelet


pos = 0
first = True
firstImage = True


def imagePreview(fileName):
  img = pydicom.dcmread(fileName).pixel_array
  plt.imshow(img)
  plt.savefig("graphics/plain.png")
  w = wavelet.wavelet(img)
  plt.imshow(w)
  plt.savefig("graphics/wavelet.png")
  w = wavelet.refine(img)
  plt.imshow(w)
  plt.savefig("graphics/wavelet_full.png")


labels = []
prediction = []

titles = []

film = {}

with open("prediction.csv") as f:
  for row in csv.reader(f, delimiter=","):
    if first:
      first = False
      titles = row[8:12]
      continue

    if len(row) < 9:
      continue

    #if firstImage:
    #  firstImage = False
    #  fileName = row[0]
    #  imagePreview(fileName)


    y = [float(row[8]), float(row[9]), float(row[10]), float(row[11])]
    z = [float(row[12]), float(row[13]), float(row[14]), float(row[15])]
    labels.append(y)
    prediction.append(z)

    k = f"{row[1]}_{row[2]}"
    if k in film:
      x = film[k]
      film[k] = (x[0] + 1, x[1], x[2] + [z])
    else:
      film[k] = (1, y, [z])


    pos += 1

plt.show()


filmx = []
for k in film.keys():
  x = film[k]
  s = len(x[2])
  w = [0 for _ in range(s)]
  a = 0
  b = s - 1
  f = 1.0
  while a <= b:
    w[a] = f
    w[b] = f
    #f = f * 1.5
    a = a + 1
    b = b - 1

  s = sum(w)
  w = [a / s for a in w]

  z = [sum([f * a[n] for f, a in zip(w, x[2])]) for n in range(4)]


  #filmx.append( (x[0], x[1], [a / x[0] for a in x[2]]) )
  filmx.append( (x[0], x[1], z) )
  #filmx.append( x )

#filmx = [(1, y, z) for y, z in zip(labels, prediction)]


for n in range(4):
  keys = set([x[1][n] for x in filmx])
  boxes = [[] for _ in range(5)]
  for x in filmx:
    k = int(x[1][n]) - 1
    boxes[k].append(x[2][n])


  #plt.scatter([x[1][n] for x in filmx], [x[2][n] for x in filmx], c=[int(x[2][n] + 0.5) for x in filmx], label=titles[n])

  for y in [1.5,2.5,3.5,4.5]:
    plt.plot([0,5], [y, y], color="#cccccc")

  plt.boxplot(boxes)

  plt.savefig(f"graphics/prediction_{titles[n]}.pdf")
  plt.close()

  diffs = [abs(z[1][n] - z[2][n]) < 0.5 for z in filmx]
  good = 0
  bad = 0
  for x in diffs:
    if x:
      good += 1
    else:
      bad += 1

  print(f"{titles[n]}: {good} <-> {bad}")