import OpenImageIO as oiio
import pandas as pd
import time
import os
from matplotlib import pyplot as plt


# Setup the cache
#cache = oiio.ImageCache.create (shared=True)
#cache.attribute ("max_open_files", 500)


def get_attr_info_buf(filename):
    """Gets compression from the filename"""
    starttime = time.time()
    buf = oiio.ImageBuf(filename)
    spec = buf.spec()
    res = spec.get_attribute('XResolution')
    #cache.invalidate(filename)
    endtime = time.time()
    return endtime - starttime

    
def get_attr_info_hndl(filename):
    """Gets compression from the image filehandle"""
    starttime = time.time()
    hndl = oiio.ImageInput.open(filename)
    spec = hndl.spec()
    compression = spec.get_attribute('XResolution')
    hndl.close()
    endtime = time.time()
    return endtime - starttime
    
    
def generate_graph(output, datadict):
    """Given the specified datadict, generate the graph"""
    df = pd.DataFrame(datadict)
    plot = df.plot(kind='line')
    fig = plot.get_figure()
    fig.savefig(output, dpi=72)
    plt.close(fig)
    
    
def main(rootdir, graphoutput):
    """Main function. Finds all the files in the rootdir and processes them. 
    Make Sure all the files in there are images!"""
    buftimes = []
    hndltimes = []
    files = os.listdir(rootdir)
    for f in files:
        image = os.path.join(rootdir, f)
        #hndltimes.append(get_attr_info_hndl(image))
        buftimes.append(get_attr_info_buf(image))   
    
    # Graph the output
    datadict = {'ImageBuf': buftimes}
    generate_graph(graphoutput, datadict)
    
    
if __name__ == '__main__':
    import sys
 
    # argv[1] = Path to imagedir
    # argv[2] = Output filename for graph
    main(sys.argv[1], sys.argv[2])