import sys
import os
import scoop
import time
import datetime
import pidly
from common_pidly import setup_idl_environment
from scoop import utils


def eboss_fit(args):
    flog = makelog(args)
    scoop.logging.basicConfig(filename=flog, format="[%(asctime)-15s] %(module)-9s %(levelname)-7s %(message)s")
    scoop.logger.info("EBOSS fitting spectra")
    scoop.logger.info(args)
    scoop.logger.info( "Log file: " + flog )
    idl = pidly.IDL()
    setup_idl_environment(idl)
    idl_command = "eboss, '{}', '{}', cpu=[{},{}], /logging, logfile='{}'".format(args[2],args[3],args[0]+1,args[1],flog)
    scoop.logger.info("Computing command:")
    scoop.logger.info(idl_command)
    idl(idl_command)  
    idl_feedback = idl( idl_command, ret=True)  
    idl.close()   
    scoop.logger.info("================== IDL FEEDBACK ==================")
    for line in idl_feedback.split("\n"):
        if not re.search('% Compiled module:', line):
            scoop.logger.info(line)
    scoop.logger.info("====================== END =======================")

    scoop.logger.info( "Finished: ID {} from {} nIDs".format(args[0]+1, args[1]) )
    return

def survey_fit(args):
    flog = makelog(args)
    scoop.logging.basicConfig(filename=flog, format="[%(asctime)-15s] %(module)-9s %(levelname)-7s %(message)s")
    scoop.logger.info("{} fitting spectra".format(args[6]))
    scoop.logger.info(args)
    scoop.logger.info( "Log file: " + flog )
    idl = pidly.IDL()
    setup_idl_environment(idl)
    idl_command = "survey, '{}', '{}', cpu=[{},{}], /logging, logfile='{}', survey='{}'".format(args[2],args[3],args[0]+1,args[1],flog,args[6])
    scoop.logger.info("Computing command:")
    scoop.logger.info(idl_command)
    idl(idl_command)  
    idl_feedback = idl( idl_command, ret=True)  
    idl.close()   
    scoop.logger.info("================== IDL FEEDBACK ==================")
    for line in idl_feedback.split("\n"):
        if not re.search('% Compiled module:', line):
            scoop.logger.info(line)
    scoop.logger.info("====================== END =======================")

    scoop.logger.info( "Finished: ID {} from {} nIDs".format(args[0]+1, args[1]) )
    return

def makelog(args):

    datetime_stamp = datetime.datetime.now().isoformat()[:-7]
    log_name = 'log_{}_{:02d}{:02d}_{}.log'.format(args[4], args[0] + 1, args[1], datetime_stamp).replace('-','').replace(':','')
    log_name = os.path.join(os.path.dirname(args[5]), log_name)

    with open(log_name, 'w') as log:
        log.write( 'Current dir: {} \n'.format(os.getcwd()) )
        log.write( 'File list: {} \n'.format(args[2]) )
        log.write( 'Hostname: {} \n'.format(utils.socket.gethostname()) )
        log.write( 'IP: {} \n'.format(utils.ip) )
        log.write( 'ID: {} \n'.format(args[0]+1) )
        log.write( 'nIDs: {} \n'.format(args[1]) )
        
    return log_name

"""
    Usage examples:
    (for server)
        python -m scoop -n 7 multiprocess.py eboss_fit 7 /sdss15/dr15/eboss/spectro/redux/v5_10_0/spectra/lite/ /data1/rcsed_v2/SDSS/EBOSS/MILES/ /u/vgorad/work/Data/logs/ > eboss_fit.log 2>&1
        python -m scoop -n 7 multiprocess.py eboss_fit 7 /sdss15/dr15/eboss/spectro/redux/v5_10_0/spectra/lite/ /data1/rcsed_v2/SDSS/EBOSS/MILES/ /u/vgorad/work/Data/logs/ > eboss_fit.log 2>&1
        python -m scoop -n 7 multiprocess.py eboss_fit 7 /sdss15/dr15/eboss/spectro/redux/v5_10_0/spectra/lite/ /data1/rcsed_v2/SDSS/EBOSS/MILES/ /u/vgorad/work/Data/logs/ hecto_270> eboss_fit.log 2>&1
"""
if __name__ == '__main__':

    try:
        task = sys.argv[1]
        proc_n = int(sys.argv[2])
        file_list = sys.argv[3]
        outpath = sys.argv[4]
        log_path = sys.argv[5]
        survey = sys.argv[6]
    except IndexError:
        print utils.getEnv()
        scoop.logger.error('Launch requires 5/6 arguments, got {:d}'.format(len(sys.argv) - 1))
        sys.exit(1)

    proc_ids = range(proc_n)
    # args is like [(0, 4, 'filename', 'outpath', 'function_name', 'logpath'), (1, 4, 'filename', 'outpath', 'function_name', 'logpath'), (2, 4, ...), ...]
    args = zip( proc_ids, [proc_n] * proc_n, [file_list] * proc_n, [outpath] * proc_n, [task] * proc_n, [log_path] * proc_n, [survey] * proc_n )
    
    if task == 'eboss_fit':
        returnValues = list(
            scoop.futures.map_as_completed(eboss_fit, args))
    if task == 'survey_fit':
        returnValues = list(
            scoop.futures.map_as_completed(survey_fit, args))

    else:
        scoop.logger.error("Wrong task {}".format(task))
        sys.exit(1)