#!/usr/bin/env python
# -*- coding:iso-8859-1

"""
This module contains the utility functions for the main use cases:
  * run simulations (locally or on some remote hosts via ssh access),
  * get data from hosts
  * generate plots
  * run T-tests
  * list runs made so far

The functions all expect the name of the simulation folder which should
have a file called simulation.conf in it.
A simulation folder will in the end contain the following dirs:

  * jobs (configurations files for all needed jobs)
  * data (all log files)
  * plots (generated PDFs)

"""

'''
Copyright (c) 2014 Nicolas Höning

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
'''

import os
import os.path as osp
import sys
from shutil import copy, rmtree
from ConfigParser import ConfigParser
from subprocess import Popen
import signal
import fjd

from stosim.sim import utils, job_creator
from stosim.analysis import plotter, tester


we_exited = [False]  # this helps us to stop using FJD if the user quit
def signal_handler(signal, frame):
    ''' gently exiting, e.g. when CTRL-C was pressed in FJD.  '''
    we_exited[0] = True


def run(simfolder):
    ''' The main function to start running simulations

        :param string simfolder: relative path to simfolder
        :returns: True if successful, False otherwise
    '''
    print('*' * 80)
    sim_name = utils.get_simulation_name(simfolder, "{}/stosim.conf".format(simfolder))
    print("Running simulation {}".format(sim_name))
    print('*' * 80)
    print('')

    if not osp.exists("%s/stosim.conf" % simfolder):
        print "[StoSim] %s/stosim.conf does not exist!" % simfolder
        utils.usage()
        return False

    # prepare all jobs to be run by FJD
    fjd_dir = fjd.utils.ensure_wdir(sim_name)
    fjd.utils.empty_queues(sim_name)
    for job in os.listdir("{}/jobs".format(simfolder)):
        copy("{}/jobs/{}".format(simfolder, job),
             "{}/jobqueue".format(fjd_dir))

    # now decide if recruiting is done in a local network or on a PBS cluster
    scheduler = utils.get_scheduler(simfolder)
    if scheduler == 'fjd':
        # let FJD handle it in local network (default: only local PC)
        if os.path.exists('{}/remote.conf'.format(simfolder)):
            copy('{}/remote.conf'.format(simfolder), fjd_dir)
        Popen('fjd-recruiter --project {} hire'.format(sim_name), shell=True).wait()
        if not we_exited[0]:
            Popen('fjd-dispatcher --project {} --end_on_empty_queue --interval {}'\
                .format(sim_name, utils.get_interval(simfolder)), shell=True).wait()
        # when recruiter got remote.conf, clean up in fjd dir
        if os.path.exists('{}/remote.conf'.format(simfolder)):
            os.remove('{}/remote.conf'.format(fjd_dir))

    elif scheduler == 'pbs':
        # queue the PBS jobs we created on a PBS job scheduler (e.g. clusters 
        # running Torque or PBS Pro). These simply start FJD workers.
        for job in [j for j in os.listdir('{}/jobs'.format(simfolder)) if j.endswith('.pbs')]:
            Popen('qsub {}'.format('{}/jobs/{}'.format(simfolder, job)), shell=True).wait()
        # Now we start dispatching
        Popen('fjd-dispatcher --project {} --end_on_empty_queue --interval {}'\
            .format(sim_name, utils.get_interval(simfolder)), shell=True).wait()

    return True


def resume(simfolder):
    sim_name = utils.get_simulation_name(simfolder, "{}/stosim.conf".format(simfolder))
    Popen('fjd-dispatcher --project {} --interval {}'\
          .format(sim_name, utils.get_interval(simfolder)), shell=True).wait()
    return True


def check(simfolder):
    """
    Check status of simulation
    """
    scheduler = utils.get_scheduler(simfolder)
    sim_name = utils.get_simulation_name(simfolder, "{}/stosim.conf".format(simfolder))
    if scheduler == 'fjd':
        Popen('fjd-dispatcher --project {} --interval {}'\
            .format(sim_name, utils.get_interval(simfolder)), shell=True).wait()
    elif scheduler == 'pbs':
        Popen('showq -u {}'.format(), shell=True).wait()

    return True


def run_more(simfolder):
    """ let the user make more runs on current config,
        in addition to the given data
        TODO: problematic when param values are only relevant for one subconf.
              Then, setup.create gets confused (it helps to set the --sim
              option in this case.
                Will get fixed anyway when we revamp the job distribution?)

        :param string simfolder: relative path to simfolder
        :returns: True if successful, False otherwise
    """
    simfolder = simfolder.strip('/')
    conf = utils.get_combined_conf(simfolder)

    print('''
[StoSim] Let's make {} more run(s)! Please tell me on which configurations.\n
Enter any parameter values you want to narrow down to, nothing otherwise."
'''.format(conf.getint('control', 'runs')))
    sel_params = {}
    for o in conf.options('params'):
        selected = False
        params = [p.strip() for p in conf.get('params', o).split(',')]
        if len(params) == 1:
            print("<{}> has only one value:".format(o))
            print(params[0])
            continue
        while not selected:
            choice = []
            print("<{}> ? (out of [{}])".format(o, conf.get('params', o)))
            for selection in raw_input().split(','):
                selected = True
                if selection == "":
                    pass
                elif selection in params:
                    choice.append(selection)
                else:
                    print("Sorry, {} is not a valid value.".format(selection))
                    selected = False
        if len(choice) > 0:
            sel_params[o] = choice
        else:
            print("No restriction chosen.")
    print("You selected: {}. Do this? [Y|n]\n"\
          "(Remember that configuration and code should still be the same!)"\
                .format(str(sel_params)))
    if raw_input().lower() in ["", "y"]:
        _prepare(simfolder, limit_to=sel_params, more=True)
        return run(simfolder)
    return False


def make_plots(simfolder, plot_nrs=[]):
    """ generate plots as specified in the simulation conf

        :param string simfolder: relative path to simfolder
        :param list plot_nrs: a list with plot indices. If empty, plot all
    """

    simfolder = simfolder.strip('/')

    #if osp.exists("%s/plots" % simfolder):
    #   rmtree('%s/plots' % simfolder)
    if not osp.exists("%s/plots" % simfolder):
        os.mkdir('%s/plots' % simfolder)

    # tell about what we'll do if we have at least one plot
    relevant_confs = utils.get_relevant_confs(simfolder)
    for c in relevant_confs:
        if c.has_section("figure1"):
            print('')
            print('*' * 80)
            print("[StoSim] creating plots ...")
            print('*' * 80)
            print('')
            break
    else:
        print("[StoSim] No plots specified")

    # Describe all options first.
    # These might be set in plot-settings (in each simulation config)
    # and also per-figure
    general_options = {'use-colors': bool, 'use-tex': bool, 'line-width': int,
                       'font-size': int, 'infobox-pos': str,
                       'use-y-errorbars': bool, 'errorbar-every': int
                      }
    figure_specific_options = {
                       'name': str, 'xcol': int, 'x-range': str,
                       'y-range': str, 'x-label': str, 'y-label': str,
                       'custom-script': str
                      }
    figure_specific_options.update(general_options)

    def get_opt_val(conf, d, section, option, t):
        if conf.has_option(section, option):
            val = c.get(section, option).strip()
            if t is int:
                val = c.getint(section, option)
            if t is bool:
                val = c.getboolean(section, option)
            if t is float:
                val = c.getfloat(section, option)
            # config-options with '-' are nice, but not good parameter names
            d[option.replace('-', '_')] = val

    general_settings = {}
    c = ConfigParser()
    c.read('{}/stosim.conf'.format((simfolder)))
    delim = utils.get_delimiter(c)
    for o, t in general_options.iteritems():
        get_opt_val(c, general_settings, 'plot-settings', o, t)
    general_params = []
    if c.has_option('plot-settings', 'params'):
        general_params = c.get('plot-settings', 'params').split(',')

    for c in relevant_confs:
        i = 1
        settings = general_settings.copy()
        # overwrite with plot-settings for this subsimulation
        for o, t in general_options.iteritems():
            get_opt_val(c, settings, 'plot-settings', o, t)
        if c.has_option('plot-settings', 'params'):
            general_params.extend(c.get('plot-settings', 'params').split(','))

        while c.has_section("figure%i" % i):
            if i in plot_nrs or len(plot_nrs) == 0:
                fig_settings = settings.copy()
                for o, t in figure_specific_options.iteritems():
                    get_opt_val(c, fig_settings, 'figure%i' % i, o, t)

                plot_confs = []
                j = 1
                while c.has_option("figure%i" % i, "plot%i" % j):
                    # make plot settings from conf string
                    d = utils.decode_search_from_confstr(
                            c.get('figure%i' % i, 'plot%i' % j),
                            sim=c.get('meta', 'name')
                        )
                    # then add general param settings to each plot, if
                    # not explicitly in there
                    for param in general_params:
                        if ":" in param:
                            param = param.split(':')
                            key = param[0].strip()
                            if not key in d.keys():
                                d[key] = param[1].strip()
                    # add simulation file name, then we can select accordingly
                    if c.has_option('meta', '_subconf-filename_'):
                        scfn = c.get('meta', '_subconf-filename_')
                        if scfn != '' and not 'sim' in d:
                            d['sim'] = scfn
                    # making sure all necessary plot attributes are there
                    if ('_name' in d.keys()and '_ycol' in d.keys() and
                        '_type' in d.keys()):
                        plot_confs.append(d)
                    else:
                        print('''
[StoSim] Warning: Incomplete graph specification in Experiment %s
- for plot {} in figure {}. \n
Specify at least _name and _ycol.'''.format((c.get('meta', 'name'), j, i)))
                    j += 1
                plotter.plot(filepath='%s/data' % simfolder,
                             delim=delim,
                             outfile_name='%s/plots/%s.pdf' \
                                % (simfolder, fig_settings['name']),\
                             plots=plot_confs,\
                             **fig_settings)
            i += 1


def run_ttests(simfolder):
    '''
    Make statistical t tests

    :param string simfolder: relative path to simfolder
    '''
    c = ConfigParser()
    c.read('{}/stosim.conf'.format((simfolder)))
    delim = utils.get_delimiter(c)

    relevant_confs = utils.get_relevant_confs(simfolder)

    # tell about what we'll do if we have at least one test
    for c in relevant_confs:
        if c.has_section("ttest1"):
            print('')
            print('*' * 80)
            print("[StoSim] Running T-tests ...")
            print('*' * 80)
            print('')
            break
    else:
        print "[StoSim] No T-tests specified"

    for c in relevant_confs:
        i = 1

        while c.has_section("ttest%i" % i):
            print("Test {}:".format(c.get('ttest%i' % i, 'name').strip('"')))
            if not (c.has_option("ttest%i" % i, "set1") and
                    c.has_option("ttest%i" % i, "set2")):
                print("[StoSim] T-test {} is missing one or both"\
                      " data set descriptions.".format(i))
                break

            tester.ttest(simfolder, c, i, delim)
            i += 1


def list_data(simfolder):
    """ List the number of runs that have been made per configuration.

        :param string simfolder: relative path to simfolder
        :returns: True if successful, False otherwise
    """
    print("[StoSim] The configurations and number of runs made so far:\n")
    for sim in utils.get_subsimulation_names(utils.get_main_conf(simfolder)):
        print("Simulation: {}".format(sim))
        # get a list w/ relevant params from first-found config file
        # they should be the same
        data_dirs = os.listdir("%s/data" % simfolder)
        sim_dirs = [d for d in data_dirs if d.startswith("sim{}".format(sim))]
        if len(sim_dirs) == 0:
            print("No runs found for simulation {}\n".format(sim))
            continue
        conf = utils.get_combined_conf(simfolder)
        params = conf.options('params')
        charlen = 1 + sum([len(p) + 6 for p in params]) + 9
        print('-' * charlen)
        print("|"),
        for p in params:
            print("  {}  |".format(p)),
        print("| runs |")
        print('-' * charlen)
        # now show how much we have in each relevant dir
        for d in sim_dirs:
            print("|"),
            for p in params:
                v = d.split("_{}".format(p))[1].split("_")[0]
                print("  {}|".format(v.ljust(len(p) + 2))),
            print("| {} |".format(str(utils.runs_in_folder(simfolder, d)).rjust(4)))
            print('-' * charlen)
    return True

# -----------------------------------------------------------------------------


def _assure_writable(simfolder, more=False):
    """ check if old data is lying around, ask if it can go

        :param boolean more: when True, new data will simply be added
                             to existing data
    """
    if osp.exists("%s/data" % simfolder):
        data_content = os.listdir('%s/data' % simfolder)
        if len([f for f in data_content if not f.startswith('.')]) > 0:
            if not more:
                if '-d' in sys.argv:
                    rmtree('%s/data' % simfolder)
                else:
                    print('[StoSim] I found older log data (in {}/data).'\
                        ' Remove? [y/N]'.format(simfolder))
                    if raw_input().lower() == 'y':
                        rmtree('%s/data' % simfolder)


def _prepare(simfolder, limit_to={}, more=False):
    """ ensure that data and job directories exist, create jobs.
        limit_to can contain parameter settings we want
        to limit ourselves to (this is in case we add more data)

        :param string simfolder: relative path to simfolder
        :param dict limit_to: key-value pairs that narrow down the dataset,
                              when empty (default) all possible configs are run
        :param boolean more: when True, new data will simply be added
    """
    if not osp.exists("%s/data" % simfolder):
        os.mkdir('%s/data' % simfolder)
    if osp.exists("%s/jobs" % simfolder):
        rmtree('%s/jobs' % simfolder)
    os.mkdir('%s/jobs' % simfolder)

    conf = utils.get_main_conf(simfolder)
    job_creator.create(conf, simfolder, limit_to=limit_to, more=more)


if __name__ == "__main__":

    signal.signal(signal.SIGINT, signal_handler)
    args = utils.read_args()
    utils.check_conf(args.folder)

    # define standard program (if no options are set)
    if (args.run == args.ttests == args.more\
       == args.list == args.resume == False and args.plots is None):
        args.run = args.ttests = True
        #args.plots = []

    # do these if they are given:
    fine = True
    if args.check:
        fine = check(args.folder)
    elif args.resume:
        fine = resume(args.folder)
    elif args.more:
        fine = run_more(args.folder)
    elif args.run:
        _assure_writable(args.folder, more=args.more)
        _prepare(args.folder)
        fine = run(args.folder)
    elif args.list:
        fine = list_data(args.folder)

    if fine:
        if args.plots is not None:
            make_plots(args.folder, plot_nrs=args.plots)
        if args.ttests:
            run_ttests(args.folder)
