using MLDatasets
using NumNN
using Plots
gr()

# ] add https://github.com/timholy/ProgressMeter.jl.git ;

using ProgressMeter
ProgressMeter.ijulia_behavior(:clear);

X_train, Y_train = FashionMNIST.traindata(Float64);
X_test, Y_test = FashionMNIST.testdata(Float64);

X_train = reshape(X_train, (size(X_train)[1:2]..., 1, size(X_train)[end]))
X_test = reshape(X_test, (size(X_test)[1:2]...,1,size(X_test)[end]))

Y_train = oneHot(Y_train)
Y_test = oneHot(Y_test);

X_Input = Input(X_train)
X = Conv2D(10, (3,3))(X_Input)
X = BatchNorm(dim=3)(X) #to normalize across the channels
X = Activation(:relu)(X)
X = MaxPool2D((2,2))(X)
X = Conv2D(20, (5,5))(X)
X = BatchNorm(dim=3)(X)
X = Activation(:relu)(X)
X = AveragePool2D((3,3))(X)
X = Flatten()(X)
X_Output = FCLayer(10, :softmax)(X);

model = Model(X_train,Y_train,X_Input,X_Output, 0.005; optimizer=:adam);

TestP = predict(model, X_test, Y_test);

println()
println("The accuracy of Test Data before the training process $(round(TestP[:accuracy], digits=4))")
println("The cost of Test Data before the training process $(round(TestP[:cost], digits=4))")

TrainP = predict(model, X_train, Y_train);

println()
println("The accuracy of Train Data before the training process $(round(TrainP[:accuracy], digits=4))")
println("The cost of Train Data before the training process $(round(TrainP[:cost], digits=4))")

TrainD = train(X_train, Y_train, model, 10);# testData = X_test, testLabels = Y_test);

plot(1:10, TrainD[:trainAccuracies], label="Training Accuracies")
plot!(1:10, TrainD[:trainCosts], label="Training Costs")
# plot!(1:10, TrainD[:testAccuracies], label="Test Accuracies")
# plot!(1:10, TrainD[:testCosts], label="Test Costs")
ylabel!("Epochs")

TrainP = predict(model, X_train, Y_train);

println()
println("The accuracy of Train Data before the training process $(round(TrainP[:accuracy], digits=4))")
println("The cost of Train Data before the training process $(round(TrainP[:cost], digits=4))")

TestP = predict(model, X_test, Y_test);

println()
println("The accuracy of Test Data before the training process $(round(TestP[:accuracy], digits=4))")
println("The cost of Test Data before the training process $(round(TestP[:cost], digits=4))")