#!python
""" CMAC animation generator. """

import argparse
import datetime
import glob
import os
import shutil
import subprocess


def datespan(start_date, end_date, delta=datetime.timedelta(days=1)):
    """ Retrieves all dates between the start and end date. """
    current_date = start_date
    while current_date < end_date:
        yield current_date
        current_date += delta


def main():
    """ Script that takes a start date and an end date and creates
    mp4 for each day, between and including the start and end date, for
    the cmac_four_panel plot. Also an option to change frame rate. """
    parser = argparse.ArgumentParser(
        description='Create a mp4 of cmac_four_panel plot for a each day.')
    parser.add_argument(
        'site_save_name', type=str,
        help='Radar site save name for the mp4s, example: sgpxsaprcmacsurI5')
    parser.add_argument(
        'field', type=str,
        help='Field name of pngs to be converted to mp4.')
    parser.add_argument(
        'start_date', type=str,
        help="Start date string to do mp4 creation. Example:'20170909'")
    parser.add_argument(
        'end_date', type=str,
        help="End date string to do mp4 creation. Example:'20170928'")
    parser.add_argument(
        '-fr', '--frame_rate', type=float, default=1.0,
        help='Frame rate to create mp4 with. Example: 2.0')
    parser.add_argument(
        '-p', '--path', type=str,
        default=None, help='Target path')
    parser.add_argument(
        '-s', '--source_path', default=None,
        help='Source path')

    args = parser.parse_args()

    site = args.site_save_name

    if args.path is None:
        path = ('/lustre/or-hydra/cades-arm/proj-shared/'
                + site + '.mp4/')
    else:
        path = args.path

    # Creates an mp4s directory if it does not exist.
    if not os.path.exists(path):
        os.makedirs(path)
        subprocess.call('chmod -R g+rw ' + path, shell=True)

    # Loops through each day between the start and end date.
    start = datetime.datetime.strptime(args.start_date, "%Y%m%d")
    stop = datetime.datetime.strptime(args.end_date, "%Y%m%d")
    for date_time in datespan(start, stop, delta=datetime.timedelta(days=1)):
        date = datetime.datetime.strftime(date_time, '%Y%m%d')
        if args.source_path is None:
            files = glob.glob(
                '/lustre/or-hydra/cades-arm/proj-shared/'
                + 'sgpxsaprcmacsur' + site + '.c1.png/' + date
                + '**/cmac*', recursive=True)
        else:
            files = glob.glob(args.source_path + '/' + date + '**/'
                              + args.field + '*', recursive=True)

        # If a date has no image files, loop goes to the next date.
        if not files:
            print(date + ' has no image files. Moving to the next date.')
            continue

        # Creates a directory with corresponding date for each mp4.
        date_path = path + date + '/'
        if not os.path.exists(date_path):
            os.makedirs(date_path)
            subprocess.call('chmod -R g+rw ' + date_path, shell=True)

        # Creates a temporary image folder to store a day at a time.
        os.makedirs(path + 'temp_images/')
        subprocess.call('chmod -R g+rw ' + path + 'temp_images/',
                        shell=True)
        # Copies the files found with glob and places in the temp directory.
        for file in files:
            shutil.copy(file, path + 'temp_images/')

        # Runs ffmpeg and creates an mp4 from all the image files in the
        # temporary image directory.
        os.system("cat " + path + "temp_images/*.png |"
                  + " ffmpeg -f image2pipe -r " + str(args.frame_rate)
                  + " -i - -movflags faststart -pix_fmt yuv420p -vf"
                  + " 'scale=trunc(iw/2)*2:trunc(ih/2)*2' -y "
                  + date_path + args.field + args.site_save_name
                  + date + ".mp4")

        # Removes temp_images directory to start fresh with each day.
        shutil.rmtree(path + 'temp_images/')

if __name__ == '__main__':
    main()
