#!/usr/bin/env python3

description = """This script is a wrapper around Snakemake calls.  It is provided to
ease the basic usage of the pipeline. Snakemake provides a lot of
functionnalities which can be parsed and used using 'snakemake --help'
directly from the command line.
"""
from shutil import copyfile, copymode
import os
import yaml
import os.path as op
import datetime
import glob
import h5py
import numpy as np

def parse_config(config, ignore_unit=True):
    """ Return dict from configuration file
    """
    def str2num(v):
        try:
            v_ = float(v) if "." in v else int(v)
        except:
            v_ = v
        return v_
            
    cfg = yaml.load(open(config, "r"), Loader=yaml.BaseLoader)
    for k,v in cfg.items():
        #if isinstance(v, (list, tuple)) and len(v)==2 and ignore_unit:
        #    cfg[k] = str2num(v[0]) # ignore units for now
        #else:
        if isinstance(v, (list, tuple)):
            cfg[k] = [str2num(vi) for vi in v]
        else:
            cfg[k] = str2num(v)
        
    return cfg


if __name__ == "__main__":

    import argparse
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('--init', type=str, default="", 
                        help= "Initialize a new run directory INIT")
    parser.add_argument('--info', action='store_true',
                        help= "Print run configuration summary")
    parser.add_argument('--data', type=str, default="", 
                        help= "Link existing ancillary data DATA or download them if DATA=download")
    parser.add_argument('--apccluster', action='store_true',
                        help= "Print command to run on apccluster")
    parser.add_argument('--rewrite-gb-history', action='store_true',
                        help= "Rewrite GB files timestamps")
    
    args = parser.parse_args()

    if args.init:
        files = ["Snakefile", "config.yml", 
                 "lisanode_config.py", "LDC_graph.py",
                 "param_vgb.yml", "param_mbhb.yml", 
                 "igb-pipeline/Snakefile", "igb-pipeline/config.yml",
                 "dgb-pipeline/Snakefile", "dgb-pipeline/config.yml",
                 "dgb-pipeline/param_dgb.yml", "igb-pipeline/param_igb.yml",
                 "runsim"]
        src = os.path.join(os.getcwd(), "test")
        dest = args.init
        os.makedirs(dest, exist_ok=True)
        os.makedirs(os.path.join(dest, "igb-pipeline"), exist_ok=True)
        os.makedirs(os.path.join(dest, "dgb-pipeline"), exist_ok=True)

        for f in files:
            if f=="runsim":
                copyfile(os.path.join(os.getcwd(), f),
                         os.path.join(dest, f))
                copymode(os.path.join(os.getcwd(), f),os.path.join(dest, f))
            else:
                copyfile(os.path.join(src, f),
                         os.path.join(dest, f))

    else:
        cfg_instru = parse_config("config.yml")
        cfg_instru_d = parse_config("dgb-pipeline/config.yml")
        cfg_instru_i = parse_config("igb-pipeline/config.yml")
        k_instru = ["t_min", "t_max",
                    "nominal_arm_length", "initial_position", "initial_rotation",
                    "travel_time_order"]

    if args.data:
        param_file = ["param_vgb.yml", "param_mbhb.yml", "dgb-pipeline/param_dgb.yml", "igb-pipeline/param_igb.yml"]
        if args.data == "download":
            print("not yet available") # to go wget in dest
        else:
            for cfg in param_file:
                cfgfile = cfg 
                cfg_src = parse_config(cfgfile)
                cfg_src["catalogs"] = os.path.join(args.data, os.path.basename(cfg_src["catalogs"]))
                yaml.dump(cfg_src, open(cfgfile, "w"), default_flow_style=False)
    
    if args.info:
        print("Instrumental configuration is: ")
        for k in k_instru:
            print("%s: %s"%(k, cfg_instru[k]))
            if cfg_instru[k]!= cfg_instru_d[k] or cfg_instru[k]!= cfg_instru_i[k]:
                print("WARNING: general config doesn't match GB ones")
                print("%s: %s !=%s or !=%s"%(k, cfg_instru[k],
                                             cfg_instru_d[k], cfg_instru_i[k]))

        print("Source configuration is: ")
        if "mbhb" in cfg_instru["sources"]:
            cfg_mbhb = parse_config("param_mbhb.yml")
            print("Will include %s MBHB from %s at dt=%s s"%(cfg_mbhb["nsource"],
                                                             cfg_mbhb["catalogs"],
                                                             cfg_instru["dt"] ))
            print("-----> using %d batch jobs"%cfg_instru["nbatch"])
        if "vgb" in cfg_instru["sources"]:
            cfg_vgb = parse_config("param_vgb.yml")
            print("Will include %d VGB from %s at dt=%s s"%(cfg_vgb["nsource"],
                                                            cfg_vgb["catalogs"],
                                                            cfg_instru["dt"]))
            print("-----> using %d batch jobs"%cfg_instru["nbatch"])

        if "dgb" in cfg_instru_d["sources"]:
            cfg_dgb = parse_config("dgb-pipeline/param_dgb.yml")
            print("Will include %d GB from %s at dt=%s s "%(cfg_dgb["nsource"],
                                                            cfg_dgb["catalogs"],
                                                            cfg_instru_d["dt"] ))
            print("-----> using %d batch jobs"%cfg_instru_d["nbatch"])
        if "igb" in cfg_instru_i["sources"]:
            cfg_igb = parse_config("igb-pipeline/param_igb.yml")
            print("Will include %d GB from %s at dt=%s s "%(cfg_igb["nsource"],
                                                            cfg_igb["catalogs"],
                                                            cfg_instru_i["dt"] ))
            print("-----> using %s batch jobs"%cfg_instru_i["nbatch"])
        print("Output will be saved in %s "%(cfg_instru["dirname"]))
        print("(and in dgb-pipeline/%s for GB strains)"%(cfg_instru_d["dirname"]))
        print("(and in igb-pipeline/%s for GB strains)"%(cfg_instru_i["dirname"]))
            
    if args.apccluster:
        max_batch = max(int(cfg_instru_d["nbatch"]),
                        int(cfg_instru_i["nbatch"]),
                        int(cfg_instru["nbatch"]))
        cmd = 'snakemake --cores %d --use-singularity --default-res mem_mb=8000 --cluster "qsub -l nodes=1:ppn=1,mem={resources.mem_mb}mb,walltime=96:00:00"'%max_batch
        print(cmd)

    if args.rewrite_gb_history:
        for pref in ["dgb", "igb"]:
            gbdir = "%s-pipeline"%pref
            if pref=="dgb":
                gbd = op.join(gbdir, cfg_instru_d["dirname"])
            else:
                gbd = op.join(gbdir, cfg_instru_i["dirname"])
            cdate = os.path.getmtime(op.join(gbd, "%s-y.h5"%pref))
            cdate = datetime.datetime.fromtimestamp(cdate)

            tmp_batch = glob.glob(op.join(gbd , "%s--*.h5"%pref))
            delta = datetime.timedelta(0.1)
            for f in tmp_batch:
                os.system("touch --date='%s' %s"%(cdate-delta, f))

            catalog = glob.glob(op.join(gbd , "%s.npy"%pref))
            for f in catalog:
                os.system("touch --date='%s' %s"%(cdate-2*delta, f))

            configs = ["%s/config.yml"%gbdir,
                       "%s/param_%s.yml"%(gbdir, pref), 
                       "%s/Snakefile"%gbdir]
            for f in configs:
                os.system("touch --date='%s' %s"%(cdate-3*delta, f))



