import cv2
import numpy as np
from multi_plot import multi_plot

# table of [luminance, chrominance, chrominance]
std_quantization_tab = [
  [
    [16,  11,  10,  16,  24,  40,  51,  61],
    [12,  12,  14,  19,  26,  58,  60,  55],
    [14,  13,  16,  24,  40,  57,  69,  56],
    [14,  17,  22,  29,  51,  87,  80,  62],
    [18,  22,  37,  56,  68, 109, 103,  77],
    [24,  35,  55,  64,  81, 104, 113,  92],
    [49,  64,  78,  87, 103, 121, 120, 101],
    [72,  92,  95,  98, 112, 100, 103,  99]
  ],
  [
    [17,  18,  24,  47,  99,  99,  99,  99],
    [18,  21,  26,  66,  99,  99,  99,  99],
    [24,  26,  56,  99,  99,  99,  99,  99],
    [47,  66,  99,  99,  99,  99,  99,  99],
    [99,  99,  99,  99,  99,  99,  99,  99],
    [99,  99,  99,  99,  99,  99,  99,  99],
    [99,  99,  99,  99,  99,  99,  99,  99],
    [99,  99,  99,  99,  99,  99,  99,  99]
  ],
  [
    [17,  18,  24,  47,  99,  99,  99,  99],
    [18,  21,  26,  66,  99,  99,  99,  99],
    [24,  26,  56,  99,  99,  99,  99,  99],
    [47,  66,  99,  99,  99,  99,  99,  99],
    [99,  99,  99,  99,  99,  99,  99,  99],
    [99,  99,  99,  99,  99,  99,  99,  99],
    [99,  99,  99,  99,  99,  99,  99,  99],
    [99,  99,  99,  99,  99,  99,  99,  99]
  ]
]

generated_quantization_tab = np.zeros((3, 8, 8))

#generowanie tabel dla kwantyzacji o odpowiedniej jakosci
def makeTables(quality):
  factor = quality;

  if quality < 1:
    factor = 1
  if quality > 99:
    factor = 99
  if quality < 50:
    quality = 5000 / factor
  else:
    quality = 200 - factor*2

  for i in range(0, 8):
    for j in range(0, 8):
      luma_quality = (std_quantization_tab[0][i][j] * quality + 50) / 100
      chroma_quality = (std_quantization_tab[1][i][j] * quality + 50) / 100

      # Limit 1 <= q <= 255
      if luma_quality < 1:
        luma_quality = 1
      if luma_quality > 255:
        luma_quality = 255
      generated_quantization_tab[0][i][j] = luma_quality

      if chroma_quality < 1:
        chroma_quality = 1
      if chroma_quality > 255:
        chroma_quality = 255
      generated_quantization_tab[1][i][j] = chroma_quality
      generated_quantization_tab[2][i][j] = chroma_quality



def kompresuj(obraz_oryginalny, quality):
    #generowanie tabel dla kwantyzacji o odpowiedniej jakosci
    makeTables(quality)

    wysokosc, szerokosc = obraz_oryginalny.shape[:2]

    # ustawienie wymiarow obrazu na wielokrotnosc 8
    if (szerokosc % 8) != 0:
        filler = obraz_oryginalny[:,szerokosc-1:,:]
        for i in range(8 - (szerokosc % 8)):
            obraz_oryginalny = np.append(obraz_oryginalny, filler, 1)

    if (wysokosc % 8) != 0:
        filler = obraz_oryginalny[wysokosc-1:,:,:]
        for i in range(8 - (wysokosc % 8)):
            obraz_oryginalny = np.append(obraz_oryginalny, filler, 0)

    wysokosc, szerokosc = obraz_oryginalny.shape[:2]

    obraz_po_konwersji = obraz_oryginalny.copy()
    #konwersja na przestrzen YCrCb, Y=luminancja, Cr/Cb=czerwony/niebieski-roznica chrominancji
    obraz_po_konwersji = cv2.cvtColor(obraz_po_konwersji, cv2.COLOR_BGR2YCR_CB)

    konwersja = np.empty(shape=(wysokosc, szerokosc, 3))

    # DCT na blokach 8x8
    for startY in range(0, wysokosc, 8):
        for startX in range(0, szerokosc, 8):
            for c in range(0, 3):
                blok = obraz_po_konwersji[startY:startY+8, startX:startX+8, c:c+1].reshape(8,8)

                # DCT
                blokf = np.float32(blok)     # konwersja na float
                dct = cv2.dct(blokf)         # DCT

                # kwantyzacja
                blokq = np.around(np.divide(dct, generated_quantization_tab[c]))
                blokq = np.multiply(blokq, generated_quantization_tab[c])

                # zapis do wynikowego obrazu
                for y in range(8):
                    for x in range(8):
                        konwersja[startY+y, startX+x, c] = blokq[y, x]

    # odwrotna DCT
    for startY in range(0, wysokosc, 8):
        for startX in range(0, szerokosc, 8):
            for c in range(0, 3):
                blok = konwersja[startY:startY+8, startX:startX+8, c:c+1].reshape(8,8)

                blokf = np.float32(blok)     # konwersja na float
                idct = cv2.idct(blokf)        # odwrotna DCT
                np.place(idct, idct>255.0, 255.0)
                np.place(idct, idct<0.0, 0.0)
                blok = np.uint8(np.around(idct))

                # zapis do wynikowego obrazu
                for y in range(8):
                    for x in range(8):
                        obraz_po_konwersji[startY+y, startX+x, c] = blok[y, x]

    # konwersja na przestrzen BGR
    obraz_po_konwersji = cv2.cvtColor(obraz_po_konwersji, cv2.COLOR_YCR_CB2BGR)
    return obraz_po_konwersji

if __name__ == "__main__":
    obraz_oryginalny = cv2.imread("balls.tif", cv2.IMREAD_COLOR)

    # quality - 100 - najlepsza jakosc, 1 - najgorsza
    obrazy = [
        kompresuj(obraz_oryginalny, 100),
        kompresuj(obraz_oryginalny, 75),
        kompresuj(obraz_oryginalny, 50),
        kompresuj(obraz_oryginalny, 25),
        kompresuj(obraz_oryginalny, 10),
        kompresuj(obraz_oryginalny, 0)
    ]

    cv2.imshow("Oryginalny obraz", obraz_oryginalny)

    plot = multi_plot(800, 1200, 2, 3, obrazy, ["100%","75%","50%","25%","10%","1%"], (0,0,0), True, (255,255,255))
    cv2.imshow("Konwersja DCT + kwantyzacja + odwrotna DCT", plot)

    cv2.waitKey(0)
    cv2.destroyAllWindows()