#!/usr/bin/env python3

import os
import sys
import time
import pprint
import argparse

import multiprocessing       as mp
import radical.utils         as ru
import radical.synapse       as rs
import radical.synapse.utils as rsu


# FIXME: support tags as flags


# ------------------------------------------------------------------------------
#
def run_samples (url, time, cpu_flops, cpu_time,
                 sto_src, sto_in, sto_tgt, sto_out,
                 sto_buf, mem_rss, n_samples):

    samples  = list()

    # assume 1 sample per second -- but that is not interpreted anyway...
    for n in range(n_samples):

        # append TIME sample
        if time:
            samples.append(['time', float(n+1), {'real'       : time}])

        # append CPU sample (set efficiency to 1)
        if cpu_time or cpu_flops:
            samples.append(['cpu', float(n+1),  {'time'       : cpu_time,
                                                 'flops'      : cpu_flops,
                                                 'efficiency' : 1}])

        # append STO sample
        if sto_in or sto_out:
            samples.append(['sto', float(n+1),  {'src'        : sto_src,
                                                 'rsize'      : sto_in,
                                                 'tgt'        : sto_tgt,
                                                 'wsize'      : sto_out,
                                                 'buf'        : sto_buf}])

        # append MEM sample
        if mem_rss:
            samples.append(['mem', float(n+1),  {'size'       : mem_rss}])

    if not samples:
        usage('no samples defined')

    info, ret, out = rs.emulate (samples=samples)
    pprint.pprint (info)

    rsu.store_profile (info, mode='emu', url=url)


# ------------------------------------------------------------------------------
#
def usage (msg=None, noexit=False):

    if  msg:
        print "\n      Error: %s" % msg

    print """
      usage     : %s -m <mode> [-n <processes>] [-f <flops>]
                               [-t <seconds>]   [-c <seconds>]
                               [-I <source>]    [-i <input>]
                               [-O <target>]    [-o <output>]
                               [-b <bufsize]    [-r <memory>]
                               [-s <samples>]   [-u <url>]

      examples  : %s -m sample -f 10000000

      mode(s)   :

        help    : show this message
        sample  : run the specified load as emulation samples


      arguments :

        -t      : number of seconds to spend sleeping
        -c      : number of seconds to spend computing
        -f      : number of flops to emulate
        -I      : name   of file  to read from
        -i      : number of bytes to read from disk
        -O      : name   of file  to write to
        -o      : number of bytes to write to disk
        -b      : number of bytes to use for dist I/O buffering
        -r      : number of bytes to allocate (RSS)
        -s      : number of samples to run with the above configuration
        -n      : number of processes which each run the given samples
        -u      : database URL (empty string to disable db access)


      Notes     :

        The default mode is 'sample'.

""" % (sys.argv[0], sys.argv[0])

    if  msg:
        sys.exit (1)

    sys.exit (0)


# ------------------------------------------------------------------------------
#
def main():
    parser = argparse.ArgumentParser (add_help=False)

    parser.add_argument('-m', '--mode',      dest='mode')
    parser.add_argument('-n', '--procs',     dest='procs')
    parser.add_argument('-t', '--time',      dest='time')
    parser.add_argument('-c', '--cputime',   dest='cpu_time')
    parser.add_argument('-f', '--flops',     dest='cpu_flops')
    parser.add_argument('-I', '--source',    dest='sto_src')
    parser.add_argument('-i', '--input',     dest='sto_in')
    parser.add_argument('-O', '--target',    dest='sto_tgt')
    parser.add_argument('-o', '--output',    dest='sto_out')
    parser.add_argument('-b', '--buffer',    dest='sto_buf')
    parser.add_argument('-r', '--memory',    dest='mem_rss')
    parser.add_argument('-s', '--samples',   dest='samples')
    parser.add_argument('-u', '--url',       dest='url')
    parser.add_argument('-h', '--help',      dest='help', action="store_true")

    arguments, args = parser.parse_known_args ()

    mode      = arguments.mode
    url       = arguments.url
    procs     = arguments.procs
    time      = arguments.time
    cpu_time  = arguments.cpu_time
    cpu_flops = arguments.cpu_flops
    sto_src   = arguments.sto_src
    sto_in    = arguments.sto_in
    sto_tgt   = arguments.sto_tgt
    sto_out   = arguments.sto_out
    sto_buf   = arguments.sto_buf
    mem_rss   = arguments.mem_rss
    samples   = arguments.samples

    if not procs     : procs     = 1
    if not time      : time      = 0
    if not cpu_time  : cpu_time  = 0
    if not cpu_flops : cpu_flops = 0
    if not sto_src   : sto_src   = None
    if not sto_in    : sto_in    = 0
    if not sto_tgt   : sto_tgt   = None
    if not sto_out   : sto_out   = 0
    if not sto_buf   : sto_buf   = 0
    if not mem_rss   : mem_rss   = 0
    if not samples   : samples   = 1

    procs     = int  (procs    )
    time      = float(time     )
    cpu_time  = int  (cpu_time )
    cpu_flops = int  (cpu_flops)
    sto_in    = int  (sto_in   )
    sto_out   = int  (sto_out  )
    sto_buf   = int  (sto_buf  )
    mem_rss   = int  (mem_rss  )
    samples   = int  (samples  )

    if  arguments.help:
        usage ()

    if  not mode:
        mode = 'sample'

    if  mode in ['help']:
        usage ()

    if not url:
        url = os.environ.get('RADICAL_SYNAPSE_DBURL')

    elif mode == 'help': usage(noexit=True)
    else               : usage("unknown mode '%s'" % mode)

    handles = list()
    for _ in range(procs):
        args = [url, time, cpu_flops, cpu_time,
                sto_src, sto_in, sto_tgt, sto_out,
                sto_buf, mem_rss, samples]
        p = mp.Process(target=run_samples, args=args)
        p.start()
        handles.append(p)

    for p in handles:
        p.join()


# ------------------------------------------------------------------------------
#
if __name__ == '__main__':

    base  = os.path.basename(os.getcwd())
    prof  = os.environ.get('RP_PROF', './%s.prof' % base)
    pfd   = None
    comp  = 'radical.synapse'
    tid   = ru.get_thread_name()
    uid   = os.environ.get('RP_UNIT_ID', base)
    state = 'AGENT_EXECUTING'


    try:
        if prof:
            pfd = open(prof, 'a')
            pfd.write("%.4f,%s,%s,%s,%s,%s,\n"
                     % (time.time(), 'app_start', comp, tid, uid, state))
            pfd.flush()
            pfd.close()

        main()

    finally:
        if prof:
            pfd = open(prof, 'a')
            pfd.write("%.4f,%s,%s,%s,%s,%s,\n"
                     % (time.time(), 'app_stop', comp, tid, uid, state))
            pfd.flush()
            pfd.close()


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

