import os
import argparse 

from rawls.rawls import Rawls
from rawls.stats import RawlsStats

def main():

    parser = argparse.ArgumentParser(description="concat and mean data and compare with 1 image only")

    parser.add_argument('--folder', type=str, help="folder which contains rawls files")
    parser.add_argument('--n', type=int, help="number of images to concat and mean")
    parser.add_argument('--output', type=str, help="output data folder")

    args = parser.parse_args()

    p_folder = args.folder
    p_n      = args.n
    p_output = args.output


    files = os.listdir(p_folder)

    paths = [ os.path.join(p_folder, f)  for f in files ]

    if not os.path.exists(p_output):
        os.makedirs(p_output)

    print('Loading 1 rawls file')
    rawls_one = Rawls.load(paths[0])
    print('Saving 1 rawls file into png')
    rawls_one.save(os.path.join(p_output, 'rawls_one.png' ))


    print('Loading {0} rawls files'.format(p_n))
    rawls_n = RawlsStats.load(paths[0:p_n])
    print('Saving {0} rawls files into png'.format(p_n))
    rawls_n.mean().save(os.path.join(p_output, 'rawls_{0}.png'.format(p_n)))


if __name__ == "__main__":
    main()