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

from adding_noise import *
from quantifying_noise import *



def load_image(filename):
    image_data = None

    with Image.open(filename) as img:
        img.load()
        image_data = np.asarray(img)

    return image_data


def main():
    points = 10

    snr_g = np.zeros(points)
    psnr_g = np.zeros(points)
    stddev = np.arange(1,points+1)

    snr_sp = np.zeros(points)
    psnr_sp = np.zeros(points)
    amount = 0.01*np.arange(1,points+1)

    subjects = 40
    images_per_subject = 10
    images = subjects * images_per_subject

    for subject in range(1, subjects+1):
        for i in range(1, images_per_subject+1):
            image = load_image(f"./att_faces/s{subject}/{i}.pgm")
            for j in range(points):
                noisy_image = add_gaussian_noise(image, stddev[j])
                snr_g[j] += snr(image, noisy_image)

                noisy_image = add_salt_pepper(image, amount[j])
                snr_sp[j] += snr(image, noisy_image)

    snr_g /= images
    snr_sp /= images

    plt.subplot(211)
    plt.stem(stddev, snr_g)
    plt.subplot(212)
    plt.stem(amount, snr_sp)
    plt.show()


if __name__ == "__main__":
    main()