######## Imports ########
#### Standard Library ####
import time
#### Third Party ####
import numpy as np
from PIL import Image
import matplotlib
from matplotlib import pyplot as plt
#### CSRK ####
from csrk import HybridGP

######## Setup ########
# Image location
FNAME_ORIGINAL = "originals/Nighthawks_2993_5180.jpg"
FNAME_OUT = "csrk_Nighthawks.jpg"
# GP stuff
train_width = 800
#train_width = 480
sample_width = 3996
#sample_width = 1920

## Hyperparameters ##
# For reconstructions of hand-drawn images that have been scanned,
#   resolution trumps all. Keep scale at 2.0 and order at 1.
scale = np.full(2,2.0)
whitenoise = 0.
order = 1

######## Load Image ########
# Load image pixel data
tic = time.perf_counter()
raw = Image.open(FNAME_ORIGINAL)
raw_YCbCr = raw.convert('YCbCr')
pixels_YCbCr = raw_YCbCr.load()
width, height = raw_YCbCr.size
toc = time.perf_counter()
print(f"Pillow initialization time: {toc-tic:.6f} s")

######## Generate Training Data for GP ########
tic = time.perf_counter()
# Define what we want to do wth the GP
train_height = int(train_width * (height / width))
train_height = min(train_height, height)
train_width = min(train_width, width)
sample_height = int(sample_width * (height / width))

# Define training x data
training_x_loc = np.arange(train_width)
training_y_loc = np.arange(train_height)
training_x_loc, training_y_loc = np.meshgrid(training_x_loc, training_y_loc)
training_x_loc = training_x_loc.flatten()
training_y_loc = training_y_loc.flatten()
X = np.column_stack([training_x_loc, training_y_loc])
# Define Y_err
Y_err = np.zeros(X.shape[0])

# Define sample data
eval_x_loc = np.arange(sample_width) * (train_width / sample_width)
eval_y_loc = np.arange(sample_height) * (train_height / sample_height)
eval_x_loc, eval_y_loc = np.meshgrid(eval_x_loc, eval_y_loc)
eval_x_loc = eval_x_loc.flatten()
eval_y_loc = eval_y_loc.flatten()
X_eval = np.column_stack((eval_x_loc, eval_y_loc)).astype(float)
    
# Define block width / height
block_width = int(np.floor(width / train_width))
block_height = int(np.floor(height / train_height))
block_size = block_height * block_height

# Define block "stamp"
block_x = np.arange(block_width).astype(int)
block_y = np.arange(block_height).astype(int)
block_x, block_y = np.meshgrid(block_x, block_y)
block_x = block_x.flatten()
block_y = block_y.flatten()

# Define Y data
Y_uv        = []
Y_Cb        = []
Y_Cr        = []
for (i, j) in X:
    x_start = i * block_width
    y_start = j * block_height
    x = block_x + x_start
    y = block_y + y_start
    lums, Cbs, Crs = 0, 0, 0
    for (xi, yi) in zip(x,y):
        lum, cb, cr = pixels_YCbCr[xi, yi]
        lums += lum
        Cbs += cb
        Crs += cr
    Y_uv.append(    lums    / block_size)
    Y_Cb.append(    Cbs     / block_size)
    Y_Cr.append(    Crs     / block_size)

Y_uv = np.asarray(Y_uv)
Y_Cb = np.asarray(Y_Cb)
Y_Cr = np.asarray(Y_Cr)
# Change X to float
X = X.astype(float)

toc = time.perf_counter()
print(f"Setting up training data time: {toc-tic:.6f} s!")

#### Train Gaussian Processes ####

# uv
print(f"Training lum GP on {train_width} by {train_height} grid!")
tic = time.perf_counter()
gp_uv = HybridGP(
    X,
    Y_uv,
    Y_err,
    scale,
    whitenoise,
    order,
)
toc = time.perf_counter()
print(f"time: {toc-tic:.6f} s")
# Cb
print(f"Training Cb GP on {train_width} by {train_height} grid!")
tic = time.perf_counter()
gp_Cb = HybridGP(
    X,
    Y_Cb,
    Y_err,
    scale,
    whitenoise,
    order,
)
toc = time.perf_counter()
print(f"time: {toc-tic:.6f} s")
# Cr
print(f"Training Cr GP on {train_width} by {train_height} grid!")
tic = time.perf_counter()
gp_Cr = HybridGP(
    X,
    Y_Cr,
    Y_err,
    scale,
    whitenoise,
    order,
)
toc = time.perf_counter()
print(f"time: {toc-tic:.6f} s")

#### Evaluate sample data ####

# Eval uv
print(f"Evaluating uv GP on {sample_width} by {sample_height} grid!")
tic = time.perf_counter()
y_eval_uv = gp_uv.predict_mean(X_eval)
toc = time.perf_counter()
print(f"time: {toc-tic:.6f} s")

# Eval Cr
print(f"Evaluating Cr GP on {sample_width} by {sample_height} grid!")
tic = time.perf_counter()
y_eval_Cb = gp_Cb.predict_mean(X_eval)
toc = time.perf_counter()
print(f"time: {toc-tic:.6f} s")

# Eval Cb
print(f"Evaluating Cb GP on {sample_width} by {sample_height} grid!")
tic = time.perf_counter()
y_eval_Cr = gp_Cr.predict_mean(X_eval)
toc = time.perf_counter()
print(f"time: {toc-tic:.6f} s")

y_eval_uv = y_eval_uv.reshape((sample_height,sample_width)).astype('u1')
y_eval_Cb = y_eval_Cb.reshape((sample_height,sample_width)).astype('u1')
y_eval_Cr = y_eval_Cr.reshape((sample_height,sample_width)).astype('u1')
y_eval = np.stack([y_eval_uv, y_eval_Cb, y_eval_Cr], axis=2)

#### Manage output image ####
# Create new image in Pillow
new_image = Image.fromarray(
    y_eval,
    "YCbCr"
)
# Save image
new_image.save(FNAME_OUT)
