#!/usr/bin/env python
# encoding: utf-8

from __future__ import (absolute_import, division, print_function, unicode_literals)

import argparse
import os
import tempfile
from subprocess import Popen, PIPE
from textwrap import dedent
from astra import log
from time import time

APOGEE_VERSION = os.getenv("APOGEE_VER")


_safe_encoding = lambda s: s.decode("utf-8").encode("utf-8")


def aspcap_parameter_file_contents(args, input_dir):
    """ Return the contents of an ASPCAP parameter file, given the parsed arguments. """

    #assert input_dir.startswith(os.getcwd())
    #relative_input_dir = input_dir[len(os.getcwd()):]
    cwd = "{}/".format(os.getcwd())
    assert args.output_dir.startswith(cwd)
    relative_output_dir = args.output_dir[len(cwd):]

    # TODO: Which arguments should be put into a parameter file, and which
    #       should be put to the command line arguments for aspcaprun?
    kwds = dict(input_dir=input_dir, relative_output_dir=relative_output_dir)
    kwds.update(args.__dict__)
    return dedent("""
                  apogee_version  {apogee_version}
                  apvisit    0
                  apred_vers '{apred_version}'
                  telescope '{telescope}'
                  instrument '{instrument}'
                  apstar_vers '{apstar_version}'
                  aspcap_vers '{aspcap_version}'
                  aspcap_config '{aspcap_config}'
                  ncpus             {ncpus: >2.0f}
                  queue              {queue}
                  qname       '{qname}'
                  qgroup      '{qgroup}'
                  caldir       '{caldir}'
                  typedef struct {{
                    char field[24];
                    char outfield[24];
                  }} ASPCAP;
                  ASPCAP {input_dir} {relative_output_dir}
                  """.format(**kwds)).strip()


def aspcap_argument_string(planfile_path, args):

    aspcap_arguments = [planfile_path, args.flag, int(args.clobber)]
    aspcap_arguments = ",".join(["'{0}'".format(ea) for ea in aspcap_arguments])

    if args.noelem: aspcap_arguments += ",/noelem"
    if args.noplot: aspcap_arguments += ",/noplot"
    if args.elemplot: aspcap_arguments += ",/elemplot"

    return aspcap_arguments



if __name__ == '__main__' :

    parser = argparse.ArgumentParser(prog=os.path.basename(__file__),
                                     description="Runs the ASPCAP pipeline on reduced data products")
    
    # These arguments are top-level requirements for an Astra component.
    parser.add_argument("-v", "--verbose", action="store_true", default=False,
                        help="verbose output")
    parser.add_argument("-i", "--from-file", action="store_true", default=False,
                        help="specifies that the INPUT_PATH is a text file that contains a list of "
                             "input paths that are separated by new lines")
    parser.add_argument("input_path",
                        help="local path to a reduced data product, or a file that contains a list "
                             "of paths to reduced data products if the -i flag is used")
    parser.add_argument("output_dir",
                        help="directory for analysis outputs")

    # TODO: help strings.
    parser.add_argument("--flag", default="11111",
                        help="a bit flag passed to ASPCAP to instruct it what to do [TODO docs] (default: 11111)")

    parser.add_argument("--apogee-version", default="unknown",
                        help="the APOGEE version used to analyze these data products (default: unknown)")
    parser.add_argument("--apred-version", default="unknown",
                        help="the reduction version of the data products (default: unknown)")
    parser.add_argument("--apstar-version",  default="unknown",
                        help="the apStar version of the data products (default: unknown)")
    parser.add_argument("--aspcap-version", default="unknown",
                        help="the version of ASPCAP used for these data products (default: unknown)")
    parser.add_argument("--aspcap-config",  default="unknown",
                        help="the ASPCAP configuration to use for these data products (default: unknown)")
    parser.add_argument("--telescope", default="apo25m",
                        help="the telescope these data were acquired from (default: apo25m)")
    parser.add_argument("--instrument", default="apogee-n",
                        help="the instrument used to acquire these data (default: apogee-n)")
    parser.add_argument("--apvisit", default=0,
                        help="the apVisit number for these data (default: 0)")

    # These arguments are optional arguments for the doaspcap.pro file.
    parser.add_argument("--ncpus", default=4,
                        help="the nubmer of CPUs available for use (default: 4)")
    parser.add_argument("--queue", default=0)
    parser.add_argument("--qname", default="apogee",
                        help="TODO docs qname (default: apogee)")
    parser.add_argument("--qgroup", default="apogee",
                        help="TODO docs qgroup (default: apogee)")
    parser.add_argument("--caldir", default=0,
                        help="the path of the calibration directory")

    parser.add_argument("--clobber", action="store_true",
                        help="overwrite existing output files")
    parser.add_argument("--noelem", action="store_true",
                        help="do not calculate elemental abundances.")
    parser.add_argument("--noplot", action="store_true",
                        help="do not produce plots.")
    parser.add_argument("--elemplot", action="store_true",
                        help="do not produce (lots of) elemental abundance plots.")
    # TODO: Do we need the --done argument anymore?
    parser.add_argument("--done", default=None,
                        help="provide an executable for a callback when this process finishes.")

    # TODO: This is for Holt'z job management, I think, because they don't go to doaspcap.
    # TODO: Check with Holtz.
    #parser.add_argument("--host")

    # These are additional options that go to doaspcap.pro
    # TODO: A lot of work (and expert domain knowledge) is needed in order to 
    #       properly document all of these options.
    """
    parser.add_argument("--errbar", default=1)
    parser.add_argument("--renorm", default=0)
    parser.add_argument("--obscont", default=0)
    parser.add_argument("--nruns", default=1)
    # TODO: provide explicit APOGEE version from this code???
    parser.add_argument("--aspcap-vers", "--aspcap_vers", default=f"aspcap_{APOGEE_VERSION}",
                        help="the ASPCAP version used to analyze these data products")
    parser.add_argument("--results-vers", "--results_vers", default=f"aspcap_{APOGEE_VERSION}",
                        help="the (internal) results version to assign to these outputs")
    parser.add_argument("--hmask", nargs="?", const=None, type=str)
    parser.add_argument("--maskfile", nargs="?", const=None, type=str) 
    parser.add_argument("--obspixmask", nargs="?", const=None, type=str) 
    parser.add_argument("--conthighbad", default=1.1)
    parser.add_argument("--contlowbad",default=0.001)
    parser.add_argument("--conthigherr", default=1e10)
    parser.add_argument("--highbad", default=1.1)
    parser.add_argument("--lowbad", default=0.001)
    parser.add_argument("--higherr", default=1.e10)
    parser.add_argument("--commiss", action="store_true", default=False)
    parser.add_argument("--nored", action="store_true", default=False) 
    parser.add_argument("--minerr", default=0.005)
    parser.add_argument("--vacuum", action="store_true", default=False)
    parser.add_argument("--symlink", action="store_true", default=False)
    parser.add_argument("--nstars", default=0) 
    parser.add_argument("--starlist", default=0)
    parser.add_argument("--qaspcap", default=False, action="store_true")
    parser.add_argument("--doelemplot", default=0)
    parser.add_argument("--altmaskdir", default=None)
    parser.add_argument("--maxwind", default=0)
    parser.add_argument("--aspcap-root", "--aspcap_root", default=None)
    parser.add_argument("--redux-root", "--redux_root", default=None)
    parser.add_argument("--liblist-path", "--liblist_path", default=None)
    parser.add_argument("--libr-path", "--libr_path", default=None)
    """

    args = parser.parse_args()
    log.info("Parsed arguments: {0}".format(args))

    # Read input paths.
    if args.from_file:
        with open(args.input_path, "r") as fp:
            input_paths = [line.strip() for line in fp.readlines()]
    else:
        input_paths = [args.input_path]

    # Create a temporary input directory and create symbolic links to the
    # absolute input paths.
    temp_input_dir = tempfile.mkdtemp()
    log.info("Created temporary input directory {0}".format(temp_input_dir))
    for source in map(os.path.abspath, input_paths):
        destination = os.path.join(temp_input_dir, os.path.basename(source))

        os.symlink(source, destination)
        log.info("Created symbolic link {0} -> {1}".format(source, destination))

    # Create output directory if needed.
    args.output_dir = os.path.abspath(os.path.realpath(args.output_dir))
    try:
        os.makedirs(args.output_dir)
    
    except OSError:
        log.info("Output directory {0} already exists".format(args.output_dir))

    else:
        log.info("Created output directory {0}".format(args.output_dir))

    # Generate a planfile.
    planfile_path = os.path.join(temp_input_dir, "plan")
    contents = aspcap_parameter_file_contents(args, temp_input_dir)
    with open(planfile_path, "w") as fp:
        fp.write(contents)
    log.info("Created plan file at {0} with contents:\n{1}".format(planfile_path, contents))

    # Construct command line arguments.
    """
    # Trying to directly run "aspcaprun" fails spectacularly.
    # This happened even if I set the CWD for the subprocess.
    # For now let's use Holtz's "runplans".
    # TODO: Figure out wtf is going on when running aspcaprun as is
    idl_args = ",".join(["aspcaprun", aspcap_argument_string(planfile_path, args)])
    commands = ["idl", "-e", idl_args]
    log.info("Executing: {0}".format(" ".join(commands)))
    """

    commands = ["runplans", planfile_path]
    log.info("Executing: {0}".format(" ".join(commands)))

    # Execute!
    t = time()
    process = Popen(commands, stdout=PIPE, stderr=PIPE)
    stdout, stderr = process.communicate()

    log.info("Process complete after {0:.0f} seconds".format(time() - t))
    log.info("Standard output:")
    log.info(_safe_encoding(stdout))
    log.info("Standard output:")
    log.info(_safe_encoding(stderr))

    # Callback?
    if args.done is not None:
        try: 
            callback = Popen(["setdone", args.done], stdout=PIPE, stderr=PIPE)

        except: 
            log.info("Exception raised when executing callback:")

        else:
            log.info("Callback completed successfully.")

        finally:
            cb_stdout, cb_stderr = callback.communicate()

            log.info("Callback standard output:")
            log.info(_safe_encoding(cb_stdout))
            log.info("Callback standard error:")
            log.info(_safe_encoding(cb_stderr))
            

        # TODO: figure out what this is from Holtz, since we are removing something,...
        #if args.host is not None :
        #    try: os.remove(args.done+'.'+args.host)
        #    except: pass

    # Clean up temporary input directory?
    raise a


