from pathlib import Path

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


def fig4_8(image_filename="p17.jpg", pause_between=True):
    # export: expand axes

    # read in image and calculate svd
    image_path = Path(__file__).resolve().parent / image_filename
    x = np.asarray(Image.open(image_path).convert("L"))
    A = x.astype(float)

    print(A[0:5, 0:5])

    U, S, Vh = np.linalg.svd(A, full_matrices=False)
    sigma = S.copy()
    V = Vh.T
    ss = A.shape
    print(f"ss = {ss}")

    k = np.array([10, 25, 50, 100], dtype=int)

    rs = min(ss[0], ss[1])
    print(f"rs = {rs}")

    # compressed versions
    fig = plt.figure(1, figsize=(6.26, 7.39))
    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

        fig.clear()
        # get(gcf)
        fig.set_size_inches(6.26, 7.39, forward=True)
        ax = fig.add_subplot(111)

        ax.imshow(Aa, cmap="gray", aspect="equal")

        ax.axis("off")
        err = sigma[k[ik]] / sigma[0]
        errz = np.round(err, 3)
        say = f"k = {k[ik]}          Rel Error = {errz:g}"
        ax.text(5, -40, say, color="black", fontsize=23, fontweight="bold")
        ax.set_ylim(Aa.shape[0] - 0.5, -75)
        fig.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.98)
        fig.canvas.draw_idle()
        if pause_between:
            plt.show(block=False)
            input("Press Enter to continue...")
        # if ik==0:
        #     fig.savefig('/Users/mark/Desktop/10.eps')
        # elif ik==1:
        #     fig.savefig('/Users/mark/Desktop/25.eps')
        # elif ik==2:
        #     fig.savefig('/Users/mark/Desktop/50.eps')
        # elif ik==3:
        #     fig.savefig('/Users/mark/Desktop/100.eps')

    plt.show()


if __name__ == "__main__":
    fig4_8()
