using PyCall
using ImageTransformations: imresize
using MIRTjim: jim
using Plots: plot
# you must install these packages for the python version in your environment path, i.e. PYTHON
torch = pyimport("torch") # import torch
np = pyimport("numpy") # import numpy as np
include("utils.jl") # import useful functions
include("model.jl") # import neural network
include("train.jl")

#------------- create training data and ground truth -------------#
N = 1
Cin = 1
Cout = 1
H = 64
W = 64
x = torch.randn(N, Cin, H, W) # initialize CNN input

tmp = [
   zeros(3,20);
   0 1 0 0 0 0 1 0 0 0 1 1 1 1 0 1 1 1 1 0;
   0 1 0 0 0 0 1 0 0 0 0 1 0 0 1 0 0 1 0 0;
   0 1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 0 1 0 0;
   0 0 1 1 1 1 0 0 0 0 1 1 0 0 0 0 0 1 1 0;
   zeros(3,20)
   0 0 1 0 0 1 1 1 1 0 0 1 0 0 1 1 1 1 0 0;
   0 0 1 0 0 1 0 0 1 0 0 1 0 0 1 0 0 1 0 0;
   0 0 1 0 0 1 0 0 1 0 0 1 0 0 1 0 0 1 0 0;
   0 0 1 0 0 1 1 1 1 0 0 1 0 0 0 0 0 1 0 0;
   0 0 1 0 0 1 0 0 1 0 0 1 0 0 0 0 0 1 0 0;
   0 0 1 0 0 1 0 0 1 0 0 1 0 0 0 0 0 1 0 0;
   0 0 1 0 0 1 1 1 1 0 0 1 0 0 0 0 0 1 0 0;
   zeros(3,20)
]';

# ground truth
target = kron(tmp, ones(3,3))
target = imresize(target, 64, 64) / maximum(target)
target[target .> 0] .= 1
target = Float32.(target)
# jim(target)
target_t = reshape(target, N, Cout, H, W)
target_t = py"numpy2tensor($target_t)" # convert to tensor

#------------------ training the neural network ------------------#
model = py"My_model($Cin, $Cout)" # initialize CNN, input channel is 3, output channel is 1
print("CNN training param =", py"cal_param($model)")

trainer = py"Trainer($model, num_epochs = 500)"
trainer.train(x, target_t)

#------------------ visualize training loss ------------------#
loss = torch.load("./checkpoint_res/train_loss.pt")
plot(loss, label = "training loss", xlabel = "# of epochs", ylabel = "MSE") # converged!


#------------------ test the trained neural network ------------------#
test_model = py"My_model($Cin, $Cout)" # initialize test model
ckpt_path = pwd() * "/checkpoint_res/last_checkpoint.pt"
ckpt = py"load_ckpt($ckpt_path)"
test_model.load_state_dict(ckpt)
cnn_output = py"test_model($test_model, $x)"
cnn_output = reshape(py"tensor2numpy($cnn_output)", H, W) # convert to julia arrays
jim(jim(cnn_output, "CNN"), jim(target, "GT")) # should be really good!
