#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright (C) California Institute of Technology 2023
#
# This file is part of pygwb.

import argparse
import configparser
import distutils
import glob
import logging
import os

import numpy as np
import pandas as pd
import pycondor
from pycondor import Dagman

from pygwb.html import pygwb_html
from pygwb.workflow import Dagman, IsotropicWorkflow, Job, makedir

__process_name__ = 'pygwb_create_workflow'
__version__ = '1.2.2'


def make_workflow(args=None):

    pygwb_wf = IsotropicWorkflow('pygwb_dagman', 
                                 config_file=args.configfile,
                                 basedir=args.basedir,
                                 config_overrides=args.config_overrides)

    seglist_pruned = pygwb_wf.setup_segments()

    # check if we want segment-level results
    if pygwb_wf.config.has_option('general', 'plot_segment_results'):
        segment_results = bool(distutils.util.strtobool(pygwb_wf.config['general']['plot_segment_results']))
    else:
        segment_results = False

    all_jobs = []
    pipe_jobs = []

    cumu_dir = os.path.join(pygwb_wf.dagman.result_dir, 'combined_results/')
    makedir(cumu_dir)

    first_start = int(seglist_pruned[0][0])
    last_start = 0

    all_coh_files = []
    all_data_files = []
    all_param_files = []
    for seg in seglist_pruned:
        start = int(seg[0])
        end = int(seg[1])
        dur = int(end-start)
        seg_dir = os.path.join(pygwb_wf.dagman.result_dir, f'{start}-{dur}/')
        makedir(seg_dir)

        last_start = start

        pipe_job, pipe_output = pygwb_wf.create_pygwb_pipe_job(parents=[], 
            output_path=seg_dir, 
            t0=start, tf=end
            )
        all_jobs.append(pipe_job)
        pipe_jobs.append(pipe_job)
        param_file, coh_file, data_file = pipe_output
        all_param_files.append(param_file)
        all_coh_files.append(coh_file)
        all_data_files.append(data_file)

        if segment_results:
            # extra vars
            combine_job, combine_output = pygwb_wf.create_pygwb_combine_job(
                t0=start, tf=end,
                parents=[pipe_job], output_path=seg_dir,
                data_path=data_file, coherence_path=coh_file, param_file=param_file,
                alpha=pygwb_wf.alpha, fref=pygwb_wf.fref)
            all_jobs.append(combine_job)

            # extra vars
            csd_file, coh_file, dsc_file = combine_output
            stats_job = pygwb_wf.create_pygwb_stats_job(
                t0=start, tf=end,
                parents=[combine_job], output_path=seg_dir,
                csd_file=csd_file,
                dsc_file=dsc_file,
                coh_file=coh_file,
                param_file=param_file)
            all_jobs.append(stats_job)

    # make cumulative results
    # FIX ME - how the param file is passed
    cumu_combine_job, cumu_combine_output = pygwb_wf.create_pygwb_combine_job(
        t0=pygwb_wf.t0, tf=pygwb_wf.tf,
        parents=pipe_jobs, output_path=cumu_dir,
        data_path=all_data_files, coherence_path=all_coh_files, param_file=all_param_files[-1],
        alpha=pygwb_wf.alpha, fref=pygwb_wf.fref)
    all_jobs.append(cumu_combine_job)

    csd_file, coh_file, dsc_file = cumu_combine_output
    cumu_stats_job = pygwb_wf.create_pygwb_stats_job(
        t0=pygwb_wf.t0, tf=pygwb_wf.tf,
        parents=[cumu_combine_job], output_path=cumu_dir,
        csd_file=csd_file,
        dsc_file=dsc_file,
        coh_file=coh_file,
        param_file=all_param_files[-1]) # Also FIX ME (param_file)
    all_jobs.append(cumu_stats_job)

    pygwb_html(pygwb_wf.dagman.base_dir, config=pygwb_wf.config_path, segment_results=segment_results)

    # write out segment list and job file
    seg_dir = os.path.join(pygwb_wf.dagman.result_dir, 'segment_lists/')
    makedir(seg_dir)
    start_list = np.array(seglist_pruned).T[0]
    end_list = np.array(seglist_pruned).T[1]
    dur_list = end_list - start_list
    ones_list = np.ones(len(dur_list), dtype=int)
    df = pd.DataFrame(data={'one': ones_list,
                            'start': start_list,
                            'end': end_list,
                            'dur': dur_list})
    jobfile_loc = os.path.join(seg_dir, "jobfile.dat") 
    df.to_csv(jobfile_loc, sep = " ", header=False, index=False)

    # make html page
    html_job = pygwb_wf.create_pygwb_html_job( 
        t0=pygwb_wf.t0, tf=pygwb_wf.tf, parents=all_jobs, segment_results=segment_results,
        output_path=pygwb_wf.dagman.base_dir, plot_path=pygwb_wf.dagman.result_dir)

    # parse the cache file and filter workflow
    if pygwb_wf.cache_file_path:
        logging.info('Parsing cache file...')
        cache = pygwb_wf.dagman.load_cache(pygwb_wf.cache_file_path)
        if len(cache):
            logging.info('Filtering out pre-computed jobs from workflow...')
            pygwb_wf.dagman.filter_and_remove_from_cache(cache)
        else:
            logging.info('No files were found in cache file, skipping...')

    new_cache_path = os.path.join(pygwb_wf.dagman.base_dir, 'pygwb_cache.txt')
    logging.info(f'Writing out new cache file at {new_cache_path}')
    pygwb_wf.dagman.save_output_loc(output_loc=new_cache_path)

    pygwb_wf.dagman.build(fancyname=False)

    if args.submit:
        logging.info('Submitting to condor...')
        pygwb_wf.dagman.submit_dag()
    elif args.run_locally:
        pygwb_wf.dagman.run_jobs_serially()
    else:
        logging.info('PyGWB Isotropic Workflow can be submitted to condor by running the following:'
                      f'\n\n\t  condor_submit_dag {pygwb_wf.dagman.submit_file}\n')

    logging.info('Done')


def main(args=None):
    logging.basicConfig(format='%(asctime)s:%(levelname)s : %(message)s',
                    level=logging.WARNING)
    logger = logging.getLogger()

    parser = argparse.ArgumentParser(description=__doc__,
                                     prog=__process_name__)
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='increase verbose output')
    parser.add_argument('--debug', action='store_true',
                        help='increase verbose output to debug')
    parser.add_argument('-V', '--version', action='version',
                        version=__version__)
    parser.add_argument('--basedir', type=os.path.abspath, required=True,
            help='Build directory for condor logging')
    parser.add_argument('--configfile', type=str, required=True,
            help='config file with workflow parameters')
    parser.add_argument('--config-overrides', type=str, nargs='+', default=[],
            help='config overrides listed in the form SECTION:KEY:VALUE')
    parser.add_argument('--submit', action='store_true',
            help='Submit workflow automatically')
    parser.add_argument('--run_locally', action='store_true',
            help='Run job on local universe')
    args = parser.parse_args(args=args)

    if args.verbose:
        logger.setLevel(logging.INFO)
    if args.debug:
        logger.setLevel(logging.DEBUG)

    # call the above function
    make_workflow(args=args)

# allow be be run on the command line
if __name__ == "__main__":
    main(args=None)
