#!/usr/bin/env python2.7

import os
import shutil

# Local imports
import locator
import nml

def parse_args():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("date", metavar="DATE", 
                        help="Run date.")
    parser.add_argument("-w", "--work-dir", default=os.getcwd())
    parser.add_argument("-f", "--forecast-types", default="forecast",
                        help="Choice of forecast,persistence,climatology")
    parser.add_argument("-l", "--lead-times", default="12",
                        help="Forecast lead times")
    parser.add_argument("-o", "--obs-types", default="profbfiles",
                        help="Choice of namobs types.")
    parser.add_argument("--class4", dest="ln_cl4", action="store_true",
                        help="Flag to choose class 4 file outputs")
    parser.add_argument("namelist", metavar="NAMELIST", 
                        help="NEMO namelist to edit.")
    args = parser.parse_args()
    args.forecast_types = args.forecast_types.split(',')
    args.obs_types = args.obs_types.split(',')
    args.lead_times = map(int, args.lead_times.split(','))
    return args

def main():
    args = parse_args()
    date = args.date
    
    # Move to working directory
    if not os.path.exists(args.work_dir):
        os.makedirs(args.work_dir)
    os.chdir(args.work_dir)

    # Collect forecast files
    types = args.forecast_types
    lead_times = args.lead_times
    namoff, namcl4 = locator.forecasts(date=date,
                                       types=types,
                                       lead_times=lead_times)

    # Process NEMO namelist
    text = nml.reader(args.namelist)
    sublists = nml.namelists(text)

    # namoff
    if "namoff" not in sublists:
        # Attach boilerplate
        text += nml.new("namoff")
    text = nml.update("namoff", text, data=namoff)

    # namcl4
    if "namcl4" not in sublists:
        # Attach boilerplate
        text += nml.new("namcl4")
    namcl4["cl4_leadtime"] = lead_times
    namcl4["cl4_date"] = date
    namcl4["cl4_match_len"] = len(namcl4["cl4_vars"])
    namcl4["cl4_fcst_len"] = len(namcl4["cl4_leadtime"])
    text = nml.update("namcl4", text, data=namcl4)

    # namrun
    namrun = {"nn_date0": date}
    text = nml.update("namrun", text, data=namrun)

    # namobs
    namobs = locator.observations(date=date,
                                  types=args.obs_types)
    namobs["ln_cl4"] = args.ln_cl4
    text = nml.update("namobs", text, data=namobs)

    # pipe text to file
    tmp = args.namelist+".tmp"
    nml.writer(tmp, text)
    shutil.move(tmp, args.namelist)

    # Run job
    run.submit()
    


if __name__ == '__main__':
    main()