'''
    DS2000
    Spring 2022
    Sample code from class -- decompressing compressed image files
    
    The csv file starts with the compressed format:
        num, color, num, color
        num, color, num, color
    
    We wrote three functions to generate the image:
        1. read from the csv file into a 2d list
        2. decompress the 2d list (3, red ===> red, red, red)
        3. plot the decompressed 2d list
'''
import csv
import matplotlib.pyplot as plt

IMG1 = "compressed.csv"
IMG2 = "compressed2.csv"

def read_csv(filename):
    ''' Function: read_csv
        parameter: filename, a string
        Returns: 2d list of strings, the contents of the file
    '''
    data = []
    with open(filename, "r") as infile:
        csvfile = csv.reader(infile, delimiter = ",")
        for row in csvfile:
            data.append(row)
    return data   


def decompress_img(lst):
    ''' Function: decompress_img
        Parameters: 2d list of strings, a compressed image
        Returns: an uncompressed version of the iamge, 2d list of strings
    '''
    decompressed = []
    for row in lst:
        uncompressed_row = []
        for i in range(0, len(row), 2):
            num = int(row[i])
            color = row[i + 1]
            for j in range(num):
                uncompressed_row.append(color)
        decompressed.append(uncompressed_row)
    return decompressed
    
def plot_img(lst):
    ''' function: plot_img
        Parameter: 2d list of strings, uncompressed image
        returns: nothing, just plots the image
    '''
    height = len(lst)
    for i in range(len(lst)):
        for j in range(len(lst[i])):
             plt.plot(j, height - i, "s", color = lst[i][j])
             

def main():
    # Step one: gather data -- ask the user which image to decompress and plot
    img = input("Which image to decompress, 1 or 2?\n")
    if img == "1":
        imgfile = IMG1
    else:
        imgfile = IMG2
        
    # Open the file and get 2 d list of its contents
    compressed = read_csv(imgfile)
  
    # Call our two functions to generate the image
    decompressed = decompress_img(compressed)
    plot_img(decompressed)
    
   
main()