from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from PIL import Image


def fig4_10(image_filename="butter.jpeg", pause_between=True, max_image_dimension=None):
    # low rank approximations using SVD

    # k=np.array([5, 10, 20, 40, 80, 160, 320])
    k = np.array([10, 25, 50, 100], dtype=int)

    # read in image and display it
    # x=np.asarray(Image.open('flower.jpg'))
    image_path = Path(__file__).resolve().parent / image_filename
    image = Image.open(image_path).convert("RGB")
    if max_image_dimension is not None and max(image.size) > max_image_dimension:
        image.thumbnail((max_image_dimension, max_image_dimension), Image.Resampling.LANCZOS)
    x = np.asarray(image)
    I = x.astype(float)
    # I = np.asarray(x, dtype=float) / 255.0  # equivalent to im2double(x); MATLAB requires the image processing toolbox
    # I = np.asarray(x, dtype=np.uint8)  # equivalent to im2uint8(x); MATLAB requires the image processing toolbox
    size_of_image = I.shape
    print(f"size_of_image = {size_of_image}")

    # get(gcf)
    fig1 = plt.figure(1, figsize=(5.60, 4.93))
    fig1.clear()
    ax1 = fig1.add_subplot(111)
    # x=np.asarray(Image.open('ff.jpeg'))
    ax1.imshow(x)
    ax1.set_aspect("equal", adjustable="box")
    ax1.axis("off")
    say = "Original Image"
    ax1.text(5, -50, say, color="black", fontsize=18, fontweight="bold")
    ax1.set_ylim(x.shape[0] - 0.5, -80)
    fig1.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.98)
    # pause

    # separate color planes
    m = len(I[:, 0, 0])
    n = len(I[0, :, 0])
    IRed = I[:, :, 0]
    IGreen = I[:, :, 1]
    IBlue = I[:, :, 2]

    # check out some of the values in red plane
    # print(IRed[0:5,0:5])

    # compose A and then find its SVD
    A = np.vstack((IRed, IGreen, IBlue))
    size_A = A.shape
    print(f"size_A = {size_A}")
    U, S, Vh = np.linalg.svd(A, full_matrices=False)
    sigma = S.copy()
    V = Vh.T
    ss = A.shape
    rs = min(ss[0], ss[1])

    # compressed versions
    fig2 = plt.figure(2, figsize=(5.60, 4.93))
    for ik in range(len(k)):

        # zero out singular values bigger than k
        Sa = np.zeros_like(sigma)
        Sa[0 : k[ik]] = sigma[0 : k[ik]]
        Aa = (U * Sa) @ Vh

        fig2.clear()
        # get(gcf)
        fig2.set_size_inches(5.60, 4.93, forward=True)
        ax2 = fig2.add_subplot(111)

        # reassemble color planes
        PP = np.zeros((m, n, 3))
        PP[:, :, 0] = Aa[0:m, :]
        PP[:, :, 1] = Aa[m : 2 * m, :]
        PP[:, :, 2] = Aa[2 * m : 3 * m, :]

        # convert floats to 8 bit integers
        AUint8 = np.clip(np.rint(PP), 0, 255).astype(np.uint8)
        ax2.imshow(AUint8)

        ax2.set_aspect("equal", adjustable="box")
        ax2.axis("off")
        err = sigma[k[ik]] / sigma[0]
        errz = np.round(err, 3)
        say = f"k = {k[ik]}          Rel Error = {errz:g}"
        ax2.text(5, -70, say, color="black", fontsize=18, fontweight="bold")
        ax2.set_ylim(AUint8.shape[0] - 0.5, -105)
        fig2.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.98)
        fig2.canvas.draw_idle()
        if pause_between:
            plt.show(block=False)
            input("Press Enter to continue...")
        # if ik==0:
        #     fig2.savefig('/Users/mark/Desktop/flower1.eps')
        # elif ik==1:
        #     fig2.savefig('/Users/mark/Desktop/flower2.eps')
        # elif ik==2:
        #     fig2.savefig('/Users/mark/Desktop/flower3.eps')
        # elif ik==3:
        #     fig2.savefig('/Users/mark/Desktop/flower4.eps')

    # compute and then plot relative error
    fig1 = plt.figure(1)
    fig1.set_size_inches(6.58, 2.30, forward=True)
    id_values = np.arange(1, ss[1])
    err = sigma[1:ss[1]] / sigma[0]
    fig1.clear()
    # get(gcf)
    fig1.set_size_inches(6.58, 2.30, forward=True)
    ax1 = fig1.add_subplot(111)
    ax1.semilogy(id_values, err, "r", linewidth=2)
    ax1.set_yticks([1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1])
    ax1.set_xlabel("k-axis", fontsize=14)
    ax1.set_ylabel(" Relative Error", fontsize=14)
    ax1.grid(True, which="major")
    ax1.grid(False, which="minor", axis="y")
    for spine in ax1.spines.values():
        spine.set_visible(True)
    ax1.tick_params(labelsize=14)
    for label in ax1.get_xticklabels() + ax1.get_yticklabels():
        label.set_fontweight("bold")
    ax1.xaxis.label.set_fontweight("bold")
    ax1.yaxis.label.set_fontweight("bold")
    fig1.subplots_adjust(top=0.96, bottom=0.23, right=0.98, left=0.11)

    plt.show()


if __name__ == "__main__":
    fig4_10()
