#!/usr/bin/env python
## -*- coding: utf-8 -*-

"""Script to test DQR scripts
"""

import gwdetchar.io.html as gwhtml
import gwsumm
import gwsumm.tabs
import gwsumm.plot


import numpy as np
import os
import numpy as np
import argparse, logging

from pycondor import Dagman
from dqrtasks.workflow import Job

parser = argparse.ArgumentParser()
parser.add_argument('--exc-args', type=str,
                    help='')
parser.add_argument('--injection-file', type=os.path.abspath,
                    help='')
parser.add_argument('--injection-num', type=int,
                    help='')
parser.add_argument('--ifar-cut', type=float,
                    help='')
parser.add_argument('--psd-file', type=os.path.abspath,
                    help='')
parser.add_argument('--ifo', type=str,
                    help='')
parser.add_argument('--output-dir', type=os.path.abspath,
                    help='')
parser.add_argument('--no-injections', action='store_true',
                    help='')
parser.add_argument('--accounting-group', type=str,
                    help='', default='ligo.dev.o4.detchar.transient.dqr')
args = parser.parse_args()

logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO, force=True)

logging.info('Setting up directories...')

basedir = args.output_dir
if not os.path.exists(basedir):
    os.makedirs(basedir)


#Condor file directories
condor_dir =  os.path.join(basedir,'condor')
if not os.path.exists(condor_dir):
    os.makedirs(condor_dir)

submit_dir =  os.path.join(condor_dir,'submit')
if not os.path.exists(submit_dir):
    os.makedirs(submit_dir)

error_dir =  os.path.join(condor_dir,'error')
if not os.path.exists(error_dir):
    os.makedirs(error_dir)

output_dir =  os.path.join(condor_dir,'output')
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

log_dir =  os.path.join(condor_dir,'log')
if not os.path.exists(log_dir):
    os.makedirs(log_dir)    

#Result file directories
results_dir =  os.path.join(basedir,'results')
if not os.path.exists(results_dir):
    os.makedirs(results_dir)

task_dir =  os.path.join(results_dir,'task_results')
if not os.path.exists(task_dir):
    os.makedirs(task_dir)

dagman = Dagman(name='dqr_test_dagman',
                submit=submit_dir)


logging.info('Setting up jobs...')
inj_list = np.arange(args.injection_num)

json_files = []
inj_jobs = []

for inj in inj_list:

    arg_file = os.path.join(submit_dir,'args_%d.txt'%inj)
    with open(arg_file, 'w') as f:
        f.write(args.exc_args)

    inj_args = ['--exc-args', str(arg_file),
                '--injection-file', str(args.injection_file),
                '--injection-num', str(inj),
                '--ifar-cut', str(args.ifar_cut),
                '--psd-file', str(args.psd_file),
                '--ifo', str(args.ifo),
                '--output-dir', os.path.join(task_dir, 'injection_%d'%inj)] 
    if args.no_injections:
        inj_args.append('--no-injections')

    job_snr = Job(name='inj_job_%d'%(inj),
               executable='dqr-test-script',
               accounting_group = args.accounting_group,
               submit=submit_dir,
               error=error_dir,
               output=output_dir,
               log=log_dir,
               getenv=True,
               arguments=' '.join(inj_args),
               request_disk='2048MB',
               dag=dagman)

    json_files.append(os.path.join(task_dir, 'injection_%d/data.json'%inj))
    inj_jobs.append(job_snr)

pp_plot_loc = os.path.join(task_dir, 'pp_plot.png')
pp_args = ['--o', pp_plot_loc,
            '--i']
pp_args = pp_args + json_files

job_pp = Job(name='pp_job',
               executable='dqr-pp-plot',
               accounting_group = args.accounting_group,
               submit=submit_dir,
               error=error_dir,
               output=output_dir,
               log=log_dir,
               getenv=True,
               arguments=' '.join(pp_args),
               request_disk='2048MB',
               dag=dagman)
for job in inj_jobs:
    job_pp.add_parent(job)


logging.info('Making html page...')

os.chdir('/'.join(results_dir.split('/')[:-1]))
html_dir = results_dir.split('/')[-1]

tabs = []
base_tabs = []
all_plots = []

home_index = os.path.join(html_dir,'index.html')
home_plots = [os.path.join(html_dir,'pp_plot.png')]
os.symlink(pp_plot_loc, os.path.join(html_dir,'pp_plot.png'))
plot_tab = gwsumm.tabs.PlotTab('Summary',path=html_dir,index=home_index)
plot_tab.set_layout((2,2))
for plot in home_plots:
    loc = os.path.join(html_dir,plot.split('/')[-1])
    plot_tab.add_plot(loc)
    all_plots.append(plot)
base_tabs.append(plot_tab)
tabs.append(plot_tab)

len_nums = int(np.ceil(np.log10(args.injection_num)))


inj_tabs = []
if not os.path.exists(os.path.join(html_dir,'injections')):
    os.makedirs(os.path.join(html_dir,'injections'))
for inj in inj_list:
    if inj==0:
        tab0 = gwsumm.tabs.ExternalTab('Injection %s'%str(inj).zfill(len_nums),
                                   url=os.path.join(html_dir,'task_results/injection_%d'%inj),
                                   parent='Injections',path=html_dir)
        tab = tab0
    else:
        tab = gwsumm.tabs.ExternalTab('Injection %s'%str(inj).zfill(len_nums),
                                  url=os.path.join(html_dir,'task_results/injection_%d'%inj),
                                  parent=tab0.parent, path=html_dir)

    inj_tabs.append(tab)
    tabs.append(tab)
    os.symlink(os.path.join(task_dir,'injection_%d'%inj),
               os.path.join(html_dir,'injections/injection_%s'%str(inj).zfill(len_nums)))
about_tab = gwsumm.tabs.AboutTab('About',path=html_dir)
about_href = about_tab.href

# Write out tabs
logging.info('Writing html of tabs...')
for tab in base_tabs:
    tab.write_html('',tabs=tabs,about=about_href)

#for tab in inj_tabs:
#    tab.write_html(tabs=tabs,about=about_href)

#injection_tab.write_html(inj_table_full,tabs=tabs,about=about_href)
about_tab.write_html(config=None,tabs=tabs)

dagman.build(fancyname=False)

print('Submit this to condor via:')
print('    condor_submit_dag %s'%os.path.join(submit_dir,'dqr_test_dagman.submit'))











logging.info('Done')
              


