from enum import Enum
from pathlib import Path

import numpy
import pywt
from matplotlib import pyplot
from sorcery import assigned_names

if __name__ == "__main__":

    class CoeffFormatEnum(Enum):
        (wavedec, wavedec2, wavedecn) = assigned_names()

    def wavelet_denoise(
        im,
        mother_wavelet: str = "db1",  # Daubechies wavelet 1
        levels: int = 4,
        keep: float = 1 / 1e2,  # percent
    ):
        """

        :param mother_wavelet:
        :type mother_wavelet:
        :param levels:
        :type levels:
        :param keep:
        :type keep:
        :param im:
        :type im:"""

        coef = pywt.wavedec2(im, wavelet=mother_wavelet, level=levels)

        coef_array, coef_slices = pywt.coeffs_to_array(coef)

        Csort = numpy.sort(numpy.abs(coef_array.reshape(-1)))

        coef_filt = pywt.array_to_coeffs(
            coef_array
            * (
                numpy.abs(coef_array) > Csort[int(numpy.floor((1 - keep) * len(Csort)))]
            ),
            coef_slices,
            output_format=CoeffFormatEnum.wavedec2.value,
        )

        recon = pywt.waverec2(coef_filt, wavelet=mother_wavelet)

        return recon

    def main(im_raw):
        """

        :param im_raw:
        :type im_raw:"""
        pyplot.figure()
        pyplot.imshow(im_raw, pyplot.cm.gray)
        pyplot.title("Original image")

        im_denoised = wavelet_denoise(im_raw)

        pyplot.figure()
        pyplot.imshow(im_denoised, pyplot.cm.gray)
        pyplot.title("Reconstructed Image")

    im22 = (
        pyplot.imread(
            str(Path.home() / "Data" / "Datasets" / "Denoise" / "moonlanding.png")
        )
        .astype(float)
        .mean(-1)
    )
    main(im22)

    pyplot.show()
