import numpy as np
import png

class Image:
    def __init__(self, x_pixels=0, y_pixels=0, num_channels=0, filename=''):
        # Recive como argumento ya sea el nombre del archivo o 
        # x_pixels, y_pixels y num_channels
        self.input_path = 'input/'
        self.output_path = 'output/'
        if x_pixels and y_pixels and num_channels:
            self.x_pixels = x_pixels
            self.y_pixels = y_pixels
            self.num_channels = num_channels
            self.array = np.zeros((x_pixels, y_pixels, num_channels))
        elif filename:
            self.array = self.read_image(filename)
            self.x_pixels, self.y_pixels, self.num_channels = self.array.shape
        else:
            raise ValueError("""
                    Debes ingresar un nombre de archivo,\n
                    o, especificar las dimensiones de la imagen.
                    """)

    def read_image(self, filename, gamma=2.2):
        """
            Lee imagen PNG RGN. Retorna un numpy array 3D, según Y, X, channel.
            Valores son float, la gamma esta decodificada.
        """
        im = png.Reader(self.input_path + filename).asFloat()
        resized_image = np.vstack(list(im[2]))
        resized_image.resize(im[1], im[0], 3)
        resized_image = resized_image ** gamma
        return resized_image

    def write_image(self, output_file_name, gamma=2.2):
        """
            Nunpy array 3D (Y, X, channel) valores entre 0 y 1, los escribe en el png
        """
        im = np.clip(self.array, 0, 1)
        y, x = self.array.shape[0], self.array.shape[1]
        im = im.reshape(y, x*3)
        writer = png.Writer(x, y)
        with open(self.output_path + output_file_name, 'wb') as f:
            writer.write(f, 255*(im**(1/gamma)))
        self.array.resize(y, x, 3)


if __name__ == '__main__':
    im = Image(filename='lake.png')
    im.write_image('test.png')