#!/bin/env python

import argparse
import os
import shutil
import sys
from pathlib import Path, PosixPath

import numpy as np


def sortingFunction(item):
    if type(item) == PosixPath:
        return int(item.stem.split("-")[2])
    else:
        return int(item.split("-")[2])
    
def find_corresponding_file(time, list_files):
    list_times = [sortingFunction(file) for file in list_files]
    index = -1
    correct_file = list_times[index]
    while time < correct_file:
        index = index - 1
        correct_file = list_times[index]
    correct_file = list_files[index ]
    return correct_file

def main():
    dag_parser = argparse.ArgumentParser()
    dag_parser.add_argument(
        "--subfile", help="Submission file.", action="store", type=str
    ) #"condor/Simulated_Data_New_Pipeline.sub"
    dag_parser.add_argument(
        "--data_path", help="Path to data files folder.", action="store", type=Path
    ) #"PROJECTS/SMDC_2021/100_day_test_pygwb/MDC_Generation_2/output/"
    dag_parser.add_argument(
        "--job_duration", help="Each job duration in seconds.", action="store", type=int, required=False
    ) #default job duration is entire file.
    dag_parser.add_argument(
        "--parentdir", help="Starting folder.", action="store", type=Path, required=False
    )
    dag_parser.add_argument(
        "--param_file", help="Path to parameters.ini file.", action="store", type=str, required=False
    )
    dag_parser.add_argument(
        "--dag_name", help="Dag file name.", action="store", type=str, required=False
    )
    dag_args = dag_parser.parse_args() 
    
    if not dag_args.parentdir:
        dag_args.parentdir = Path(os.path.abspath('./'))
    if not dag_args.param_file:
        dag_args.param_file = os.path.abspath('../parameters_mock.ini')
    if not dag_args.dag_name:
        dag_args.dag_name = "condor_simulated_100_day_MDC_2.dag"
    
    # Files
    files_H_sorted = [f for f in dag_args.data_path.glob("H*.*") if f.is_file()]
    files_L_sorted = [f for f in dag_args.data_path.glob("L*.*") if f.is_file()]
    
    files_H_sorted.sort(key=sortingFunction)
    files_L_sorted.sort(key=sortingFunction)
    
    n_Files = np.size(files_H_sorted)
    
    filelength = int(files_H_sorted[int(n_Files/2) + 1].stem.split("-")[2]) - int(files_H_sorted[int(n_Files/2)].stem.split("-")[2])
    
    if not dag_args.job_duration:
        try:
            dag_args.job_duration = int(filelength)
        except:
            raise ValueError("Unexpected format for file name -- please set job_duration explicitly.")
            
    # Initiating times of jobs    
    times_in_DAGfile = np.arange(0, n_Files*filelength, dag_args.job_duration)
    
    #lil_times = times_in_DAGfile[0:7] #Use this for testing and bug fixing
    
    outputdir = dag_args.parentdir / "output"
    logdir = outputdir / "condorLogs"
    
    # Make directories
    logdir.mkdir(parents=True, exist_ok=True)
    outputdir.mkdir(parents=True, exist_ok=True)
    
    dag = outputdir / dag_args.dag_name

    # Get the local executable
    executable = shutil.which('pygwb_pipe')
    
    #executable = '/home/max.lalleman/public_html/Code/New_Clone_pyGWB/pygwb/pygwb_pipe/pygwb_pipe' #Used if no pyGWB env installed. 
    
    at_the_end = False
    
    # Writing the DAG file
    
    with open(dag, "w") as dagfile:
    
        for idx, time in enumerate(times_in_DAGfile):
            
            if time > (n_Files-1)*filelength:
                path_H1 = os.path.abspath(files_H_sorted[-1])
                path_L1 = os.path.abspath(files_L_sorted[-1])
            
            elif time % filelength == 0:
                index_file = int(time/filelength)
                path_H1 = os.path.abspath(files_H_sorted[index_file])
                path_L1 = os.path.abspath(files_L_sorted[index_file])               
            
            else:
                H1_file = find_corresponding_file(time, files_H_sorted)
                H1_file_stem = H1_file.stem
                L1_stem = H1_file_stem.replace('H', 'L')
                L1_file = H1_file.with_stem(L1_stem)
                path_H1 = os.path.abspath(H1_file) #dag_args.data_path / path_H1
                path_L1 = os.path.abspath(L1_file) #dag_args.data_path / path_L1
                
            next_file_time = sortingFunction(path_H1) + filelength
            if time + dag_args.job_duration > n_Files*filelength:
                at_the_end = True

            elif time + dag_args.job_duration > next_file_time:
                path_H1 = Path(path_H1)
                index_file = files_H_sorted.index(path_H1)
                path_H1 = str(path_H1)
                path_H1 = [path_H1, os.path.abspath(files_H_sorted[index_file + 1])]
                path_L1 = [path_L1, os.path.abspath(files_L_sorted[index_file + 1])]
                
                path_H1 = '[\\\\\\"{0}\\\\\\",\\\\\\"{1}\\\\\\"]'.format(path_H1[0], path_H1[1])
                path_L1 = '[\\\\\\"{0}\\\\\\",\\\\\\"{1}\\\\\\"]'.format(path_L1[0], path_L1[1])
            
            else:
                pass
                
            if at_the_end:
                tf_value = sortingFunction(files_H_sorted[-1]) + filelength
                dagfile.write(f"JOB {idx} {os.path.abspath(dag_args.subfile)}\n")
                dagfile.write(
                    f'VARS {idx} job="{idx}" executable="{executable}" '
                    f'ARGS="--t0 {time} --tf {tf_value} '
                    f'--H1 {path_H1} --L1 {path_L1} --output_path {outputdir} '
                    f'--param_file {os.path.abspath(dag_args.param_file)}" logdir="{logdir}"\n'
                )
                dagfile.write("\n")
                
            else:
                dagfile.write(f"JOB {idx} {os.path.abspath(dag_args.subfile)}\n")
                dagfile.write(
                    f'VARS {idx} job="{idx}" executable="{executable}" '
                    f'ARGS="--t0 {time} --tf {time+dag_args.job_duration} '
                    f'--H1 {path_H1} --L1 {path_L1} --output_path {outputdir} '
                    f'--param_file {os.path.abspath(dag_args.param_file)}" logdir="{logdir}"\n'
                )
                dagfile.write("\n")
            

if __name__ == "__main__":
    main()

