'''
commands
=============

This module provides all commands you can use in StoSim:
run, resume, status, kill, run_more, make_plots, run_ttests, list_data
'''



import os
import os.path as osp
import subprocess
from ConfigParser import ConfigParser
from shutil import copy

import fjd

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


we_exited = [False]  # this helps us to stop using FJD if the user quit

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 [j for j in os.listdir("{}/jobs".format(simfolder))\
                if j.endswith('.conf')]:
        copy("{}/jobs/{}".format(simfolder, job),
             "{}/jobqueue".format(fjd_dir))
    dispatch_cmd = 'fjd-dispatcher --project {} --end_when_jobs_are_done '\
                   ' --end_callback "stosim --kill" --interval {}'\
                   .format(sim_name, utils.get_interval(simfolder))

    # 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)
        subprocess.call('fjd-recruiter --project {} hire'.format(sim_name), shell=True)
        if not we_exited[0]:
            subprocess.call(dispatch_cmd, shell=True)
        # 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')]:
            subprocess.call('qsub {}'.format('{}/jobs/{}'.format(simfolder, job)), shell=True)
        # Now we start dispatching
        subprocess.call(dispatch_cmd, shell=True)

    return True


def resume(simfolder):
    """
    (Re)start dispatching jobs 

    :param string simfolder: relative path to simfolder
    :returns: True if successful, False otherwise
    """
    sim_name = utils.get_simulation_name(simfolder, "{}/stosim.conf".format(simfolder))
    subprocess.call('fjd-dispatcher --project {} --interval {}'\
          .format(sim_name, utils.get_interval(simfolder)), shell=True)
    return True


def status(simfolder):
    """
    Check status of simulation. If on PBS scheduling, show status of nodes.
    Then, ask the fjd-dispatcher about status of jobs.

    :param string simfolder: relative path to simfolder
    :returns: True if successful, False otherwise
    """
    scheduler = utils.get_scheduler(simfolder)
    sim_name = utils.get_simulation_name(simfolder, "{}/stosim.conf".format(simfolder))
    if scheduler == 'pbs':
        num_nodes = len([n for n in os.listdir('{}/jobs'.format(simfolder))\
                        if n.endswith('.pbs')])
        if num_nodes > 0:
            print("[StoSim] State of our {} PBS computing nodes:".format(num_nodes))
            subprocess.call('echo "Waiting: $(qselect -u $USER -s W | wc -l)"', shell=True)
            subprocess.call('echo "Queued: $(qselect -u $USER -s Q | wc -l)"', shell=True)
            subprocess.call('echo "Running: $(qselect -u $USER -s R | wc -l)"', shell=True)
        else:
            print("[StoSim] No PBS computing nodes seem to be configured...")

    print("[StoSim] State of workers and jobs:")
    subprocess.call('fjd-dispatcher --project {} --status_only --interval {}'\
        .format(sim_name, utils.get_interval(simfolder)), shell=True)
    return True


def kill(simfolder):
    """
    Kill simulation
    """
    scheduler = utils.get_scheduler(simfolder)
    sim_name = utils.get_simulation_name(simfolder, "{}/stosim.conf".format(simfolder))
    if scheduler == 'fjd':
        subprocess.call('fjd-recruiter --project {} fire'.format(sim_name),
              shell=True)
    elif scheduler == 'pbs':
        subprocess.call('qselect -u $USER | xargs qdel', shell=True)

    return True



def run_more(simfolder):
    """ let the user make more runs on current config,
        in addition to the given data

        :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"]:
        utils.prepare_folders_and_jobs(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


