# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Julia 1.9.2
#     language: julia
#     name: julia-1.9
# ---

# ## OptShrink matrix denoising demo
# - 2017-11-16 Jeff Fessler, University of Michigan
# - 2023-05-21 Julia 1.9.0

using LaTeXStrings
using Random: seed!
using Plots: default, gui, plot, savefig, scatter!
using LinearAlgebra: svd, svdvals, norm, Diagonal, rank
using MIRTjim: jim, prompt
default(markersize=7, markerstrokecolor=:auto,
 labelfontsize = 18, legendfontsize = 18, tickfontsize = 14)

# you must modify the following directory path:
dircode = ENV["hw551test"]
include(joinpath(dircode, "optshrink2.jl"));

# make a matrix that has low rank
tmp = [
    zeros(1,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(1,20)
]';
X = 80 * kron(tmp, ones(5,5));

# Q: What is the maximum rank of X?
# - A. 1
# - B. 2
# - C. 3
# - D. 4
# - E. 5

clim = (0,80)
cticks = [0,80]
siz = (600, 260)
args = ( ; clim, cticks, size = siz,
 xaxis = false, yaxis = false, colorbar = :none, # book
)
p1 = jim(X,
 L"\mathrm{Original\ image\ } X \mathrm{\ with\ rank\ } %$(rank(X))";
 xlabel = " ", args...,
)

# +
# savefig(p1, "07_optshrink1x.pdf")
# -


# Data Y is X + noise
seed!(0)
Z = 7 * randn(size(X))
Y = X + Z
nrmse = (Xh) -> norm(Xh - X) / norm(X) * 100
nrmse_y = nrmse(Y)
p2 = jim(Y,
 L"\mathrm{Noisy\ image\ } Y=X+Z \mathrm{\ with\ rank\ } %$(rank(Y))";
 xlabel="NRMSE = $(round(nrmse_y, digits=1)) %",
 args...,
)

# +
# savefig(p2, "07_optshrink1y.pdf")
# -

rmax = minimum(size(X))
ps = plot(title="Singular values", widen=true,
    xaxis = (L"k", (0,rmax), [1,rank(X),rmax]),
    yaxis = (L"\sigma_k", (0,1700), [0, 267, 746, 947, 1672]))
scatter!(1:rmax, svdvals(Y), color=:red, label=L"\sigma_k(Y)", marker=:hexagon)
scatter!(1:rmax, svdvals(X), color=:blue, label=L"\sigma_k(X)")

# +
# savefig(ps, "07_optshrink1s.pdf")
# -

# try several ranks
ranks = 1:15
Xos = zeros(size(Y)..., length(ranks))
Xlr = copy(Xos)
nrmse_os = zeros(length(ranks))
nrmse_lr = copy(nrmse_os)
for (ir, r) in enumerate(ranks)
    Xos[:,:,ir] = optshrink2(Y, r)
    U,s,V = svd(Y)
    Xlr[:,:,ir] = U[:,1:r] * Diagonal(s[1:r]) * V[:,1:r]'
    nrmse_os[ir] = nrmse(Xos[:,:,ir])
    nrmse_lr[ir] = nrmse(Xlr[:,:,ir])
end

ir = 4
title = L"\mathrm{OptShrink\ image\ } \hat{X} \mathrm{\ with\ rank\ guess\ %$(ranks[ir])}"
p3 = jim(Xos[:,:,ir], title;
 xlabel="NRMSE = $(round(nrmse_os[ir], digits=1)) %", args...
)

# +
# savefig(p3, "07_optshrink1g4.pdf")
# -

pn = plot(
    xaxis = ("Rank (estimate)", extrema(ranks), [1,rank(X),maximum(ranks)]),
    yaxis = ("NRMSE %", (0, 60), [0, 7, 60]), widen = true)
scatter!(ranks, nrmse_lr, color=:magenta, label = "LR NRMSE")
scatter!(ranks, nrmse_os, color=:darkgreen, label = "OptShrink NRMSE", marker=:star)

# +
# savefig(pn, "07_optshrink1n.pdf")
# -

ir = 9
title = L"\mathrm{OptShrink\ image\ } \hat{X} \mathrm{\ with\ rank\ guess\ %$(ranks[ir])}"
p4 = jim(Xos[:,:,ir], title; args...,
 xlabel = "NRMSE = $(round(nrmse_os[ir], digits=1)) %",
)

# +
# savefig(p4, "07_optshrink1g9.pdf")
# -

rmax = minimum(size(X))
ir = 9
pss = deepcopy(ps)
label = latexstring("\\mathrm{OptShrink:\\ } \\hat{r}=$ir")
scatter!(pss, 1:rmax, svdvals(Xos[:,:,ir]); color=:green, label, marker=:star)

# +
# savefig(pss, "07_optshrink1o.pdf")
# -

# ## Below here for HW

# +
if false # delete for HW
include(dircode * "lr_schatten.jl") # for HW, put path to your lr_schatten here

# try different regularization parameters
for l2reg in 8:0.5:12
    Xr = lr_schatten(Y, 2^l2reg)
    @show l2reg, nrmse(Xr)
end

# Schatten LR with β=1000
β = 1000
Xs = lr_schatten(Y, β)
@show nrmse(Xs)

jim(Xs, "Schatten LR image with p=1/2 and β = $β"; size=(600,250),
    xlabel="NRMSE = $(round(nrmse(Xs), digits=1)) %")

#savefig("hsj68im.pdf")

psp = deepcopy(ps)
scatter!(psp, 1:rmax, svdvals(Xs), color=:magenta, label=L"\mathrm{Schatten }\ p=1/2")

#savefig(psp, "hsj68svd.pdf")
end # delete for HW
# -