from WorkingPyDemo import *
def produce_error_array(encoded_string, list_dic, bins, use_diff):
    """
    This function decodes the encoded_matrix.
    Input:
    list_dic        (num_dic + 1,): a list of huffman coding table 
    encoded_matrix  (512, 640): encoded matrix
    bins            (num_bins - 1,): a list of threshold to cut the bins
    use_diff (bool): whether or not to use the alternate differencing system
    
    Return:
    decode_matrix   (512, 640): decoded matrix
    """
    change_matrix = np.zeros((512,640))
    A = np.array([[3,0,-1],[0,3,3],[1,-3,-4]]) # the matrix for system of equation
    # change the dictionary back to list
    # !!!!!WARNING!!!! has to change this part, everytime you change the number of bins
    the_keys0 = list(list_dic[0].keys())
    the_values0 = list(list_dic[0].values())
    
    the_keys1 = list(list_dic[1].keys())
    the_values1 = list(list_dic[1].values())
    
    the_keys2 = list(list_dic[2].keys())
    the_values2 = list(list_dic[2].values())
    
    the_keys3 = list(list_dic[3].keys())
    the_values3 = list(list_dic[3].values())
    
    the_keys4 = list(list_dic[4].keys())
    the_values4 = list(list_dic[4].values())
    
    #Matrix system of points that will be used to solve the least squares fitting hyperplane
    points = np.array([[-1,-1,1], [-1,0,1], [-1,1,1], [0,-1,1]])
    
    decode_matrix = np.zeros((512,640))
    # loop through all the element in the matrix
    for i in range(decode_matrix.shape[0]):
        for j in range(decode_matrix.shape[1]):
            # if it's the very first pixel on the image
            if i == 0 and j == 0:
                colorvalue, encoded_string = decode_string(encoded_string,the_keys=the_keys0, the_values=the_values0)
                decode_matrix[i][j] = colorvalue
                
            # if it's on the boundary (any of the 4 edges)
            elif i == 0 or i == decode_matrix.shape[0]-1 or j == 0 or j == decode_matrix.shape[1]-1:
                colorvalue, encoded_string = decode_string(encoded_string,the_keys=the_keys0, the_values=the_values0)
                decode_matrix[i][j] = colorvalue + decode_matrix[0][0]
            # if not the boundary
            else:
                # predict the image with the known pixel value
                z0 = decode_matrix[i-1][j-1]
                z1 = decode_matrix[i-1][j]
                z2 = decode_matrix[i-1][j+1]
                z3 = decode_matrix[i][j-1]
                y0 = int(-z0+z2-z3)
                y1 = int(z0+z1+z2)
                y2 = int(-z0-z1-z2-z3)
                y = np.vstack((y0,y1,y2))

                    
                f, difference, rank, s = la.lstsq(points, [z0,z1,z2,z3], rcond=None) 
                difference = difference.astype(int)
                    
                predict = np.round(np.round(np.linalg.solve(A,y)[-1][0],1))
                # add on the difference by searching the dictionary
                # !!!!!WARNING!!!! has to change this part, eveytime you change the number of bins
                if difference <= bins[0]:
                    colorvalue, encoded_string = decode_string(encoded_string,the_keys=the_keys1, the_values=the_values1)
                    decode_matrix[i][j] = colorvalue + int(predict)
                    change_matrix[i][j] = colorvalue
                elif difference <= bins[1] and difference > bins[0]:
                    colorvalue, encoded_string = decode_string(encoded_string,the_keys=the_keys2, the_values=the_values2)
                    decode_matrix[i][j] = colorvalue + int(predict)
                    change_matrix[i][j] = colorvalue
                elif difference <= bins[2] and difference > bins[1]:
                    colorvalue, encoded_string = decode_string(encoded_string,the_keys=the_keys3, the_values=the_values3)
                    decode_matrix[i][j] = colorvalue + int(predict)
                    change_matrix[i][j] = colorvalue
                else:
                    colorvalue, encoded_string = decode_string(encoded_string,the_keys=the_keys4, the_values=the_values4)
                    decode_matrix[i][j] = colorvalue + int(predict)
                    change_matrix[i][j] = colorvalue
    return decode_matrix.astype(int), change_matrix.astype(int)
def color_adjust(visual_array):
    min_of_errors = np.min(visual_array)
    adjusted_array = visual_array - min_of_errors
    adjusted_array = np.round(adjusted_array*255/np.max(adjusted_array))
    return adjusted_array
if __name__ == "__main__":
    scenes = file_extractor("averaged_images(11)")
    images = image_extractor(scenes)
    list_dic = np.load("second_dict.npy", allow_pickle="TRUE")
    bins = [21,32,48]
    print(images)
    encoded_string2 = bytes_to_bitstring(read_from_file(images[0][:-5] + "_Compressed.txt"))
    original_array, error_array = produce_error_array(encoded_string2, list_dic, bins, False)
    adjusted_errors = color_adjust(abs(error_array))
    original_array_adjusted = color_adjust(original_array)
    # print(adjusted_errors)
    print(error_array)
    plt.subplot(131)
    plt.imshow(original_array_adjusted,cmap='gray',vmin = 0, vmax=255)
    plt.subplot(132)
    plt.imshow(adjusted_errors,cmap = 'gray', vmin = 0, vmax=255)




    # encoded_string2 = bytes_to_bitstring(read_from_file(images[0][:-5] + "_Compressed.txt"))
    # original_array, error_array = produce_error_array(encoded_string2, list_dic, bins, False)
    # adjusted_errors = color_adjust(abs(error_array))
    # original_array_adjusted = color_adjust(original_array)
    # plt.subplot(133)
    # plt.imshow(adjusted_errors,cmap = 'gray', vmin = 0, vmax=255)



    plt.show()