import time
import h2o_methods, h2o_print as h2p, h2o_sandbox, h2o_args
from h2o_test import verboseprint, dump_json
from h2o_xl import Key

###################
# REST API ACCESSORS
# TODO: remove .json

def jobs(self, job_key=None, timeoutSecs=10, **kwargs):
    '''
    Fetch all the jobs or a single job from the /Jobs endpoint.
    '''
    params_dict = {
        # 'job_key': job_key
    }
    h2o_methods.check_params_update_kwargs(params_dict, kwargs, 'jobs', True)
    result = self.do_json_request('3/Jobs.json', timeout=timeoutSecs, params=params_dict)
    return result

# TODO: add delays, etc.
# if a key= is passed, it does a frames on that key while polling (intermediate model results?)
def poll_job(self, job_key, timeoutSecs=10, retryDelaySecs=0.5, key=None, **kwargs):
    '''
    Poll a single job from the /Jobs endpoint until it is "status": "DONE" or "CANCELLED" or "FAILED" or we time out.
    '''
    params_dict = {}
    # merge kwargs into params_dict
    h2o_methods.check_params_update_kwargs(params_dict, kwargs, 'poll_job', False)

    start_time = time.time()
    pollCount = 0
    while True:
        result = self.do_json_request('3/Jobs.json/' + job_key, timeout=timeoutSecs, params=params_dict)
        # print 'Job: ', dump_json(result)

        if key:
            frames_result = self.frames(key=key)
            print 'frames_result for key:', key, dump_json(result)

        jobs = result['jobs'][0]
        description = jobs['description']
        dest = jobs['dest']
        dest_name = dest['name']
        msec = jobs['msec']
        status = jobs['status']
        progress = jobs['progress']
        print description, \
            "dest_name:", dest_name, \
            "\tprogress:", "%-10s" % progress, \
            "\tstatus:", "%-12s" % status, \
            "\tmsec:", msec

        if status=='DONE' or status=='CANCELLED' or status=='FAILED':
            h2o_sandbox.check_sandbox_for_errors()
            return result

        # what about 'CREATED'
        # FIX! what are the other legal polling statuses that we should check for?

        if not h2o_args.no_timeout and (time.time() - start_time > timeoutSecs):
            h2o_sandbox.check_sandbox_for_errors()
            emsg = "Job:", job_key, "timed out in:", timeoutSecs

            # for debug
            a = h2o.nodes[0].get_cloud()
            print "cloud.json:", dump_json(a)
            raise Exception(emsg)
            print emsg
            return None

        # check every other poll, for now
        if (pollCount % 2) == 0:
            h2o_sandbox.check_sandbox_for_errors()

        time.sleep(retryDelaySecs)
        pollCount += 1

def import_files(self, path, timeoutSecs=180):
    '''
    Import a file or files into h2o.  The 'file' parameter accepts a directory or a single file.
    192.168.0.37:54323/ImportFiles.html?file=%2Fhome%2F0xdiag%2Fdatasets
    '''
    a = self.do_json_request('3/ImportFiles.json',
        timeout=timeoutSecs,
        params={"path": path}
    )
    verboseprint("\nimport_files result:", dump_json(a))
    h2o_sandbox.check_sandbox_for_errors()
    return a

# key is required
# FIX! for now h2o doesn't support regex here. just key or list of keys
# FIX! default turn off intermediateResults until NPE is fixed.
def parse(self, key, hex_key=None, columnTypeDict=None,
          timeoutSecs=300, retryDelaySecs=0.2, initialDelaySecs=None, pollTimeoutSecs=180,
          noise=None, benchmarkLogging=None, noPoll=False, intermediateResults=False, **kwargs):
    '''
    Parse an imported raw file or files into a Frame.
    '''
    # these should override what parse setup gets below
    params_dict = {
        'source_frames': None,
        'destination_frame': hex_key,
        'parse_type': None, # file type
        'separator': None,
        'single_quotes': None,
        'check_header': None, # forces first line to be seen as column names
        'number_columns': None,
        'column_names': None, # a list
        'column_types': None, # a list. or can use columnTypeDict param (see below)
        'na_strings' : None, # a list
        'chunk_size': None,
        # are these two no longer supported?
        'delete_on_done': None,
        'blocking': None,
    }

    # if key is a list, create a comma separated string
    # list or tuple but not string
    if not isinstance(key, basestring):
        # it's a list of some kind (tuple ok?)
        # if len(key) > 1:
        #     print "I noticed you're giving me a list of > 1 keys %s to parse:" % len(key), key

        # len 1 is ok here. 0 not. what if None or [None] here
        if not key:
            raise Exception("key seems to be bad in parse. Should be list or string. %s" % key)
        # have to put double quotes around the individual list items (single not legal)
        source_frames = "[" + ",".join(map((lambda x: '"' + x + '"'), key)) + "]"

    else:
        # what if None here
        source_frames = '["' + key + '"]' # quotes required on key

    params_dict['source_frames'] = source_frames

    # merge kwargs into params_dict
    # =None overwrites params_dict

    # columnTypeDict not used here
    h2o_methods.check_params_update_kwargs(params_dict, kwargs, 'parse before setup merge', print_params=False)
    # Call ParseSetup?source_frames=[keys] . . .

    # if benchmarkLogging:
    #     cloudPerfH2O.get_log_save(initOnly=True)

    # h2o_methods.check_params_update_kwargs(params_dict, kwargs, 'parse_setup', print_params=True)
    params_setup = {'source_frames': source_frames}
    setup_result = self.do_json_request(jsonRequest="3/ParseSetup.json", cmd='post', timeout=timeoutSecs, postData=params_setup)
    h2o_sandbox.check_sandbox_for_errors()
    verboseprint("ParseSetup result:", dump_json(setup_result))

    # this should match what we gave as input?
    if setup_result['source_frames']:
        # should these be quoted?
        source_framesStr = "[" + ",".join([('"%s"' % src['name']) for src in setup_result['source_frames'] ]) + "]"
    else:
        source_framesStr = None

    # I suppose we need a way for parameters to parse() to override these
    # should it be an array or a dict?
    if setup_result['column_names']:
        # single quotes not legal..need double quotes
        columnNamesStr = "[" + ",".join(map((lambda x: '"' + x + '"'), setup_result['column_names'])) + "]"
    else:
        columnNamesStr = None

    columnTypes = setup_result['column_types']
    assert columnTypes is not None, "%s %s" % ("column_types:", columnTypes)

    if setup_result['na_strings']:
        # single quotes not legal..need double quotes
        naStrings = "[" + ",".join(map((lambda x: '"' + x + '"' if x != None else '""'), setup_result['na_strings'])) + "]"
    else:
        naStrings = None

    # dict parameter to update columnTypeDict?
    # but we don't pass columnNames like this?
    ct = setup_result['column_types']
    columnNames = setup_result['column_names']
    if columnTypeDict:
        for k,v in columnTypeDict.iteritems():
            if isinstance(k, int):
                # if a column index
                if k>=0 and k<len(ct):
                    ct[k] = v
                else:
                    raise Exception("bad col index %s in columnTypeDict param %s" % (k, columnTypeDict))
            # if a column name
            elif isinstance(k, basestring):
                # find the index
                if k not in columnNames:
                    raise Exception("bad col name %s in columnTypeDict param %s. columnNames: %s" % (k, columnTypeDict, columnNames))
                ci = columnNames.index(k)
                ct[ci] = v
            else:
                raise Exception("%s %s should be int or string" % (k, type(k)))

    columnTypesStr = "[" + ",".join(map((lambda x: '"' + x + '"'), ct)) + "]"

    parse_params = {
        'source_frames': source_framesStr,
        'destination_frame': setup_result['destination_frame'],
        'parse_type': setup_result['parse_type'],
        'separator': setup_result['separator'],
        'single_quotes': setup_result['single_quotes'],
        'check_header': setup_result['check_header'],
        'number_columns': setup_result['number_columns'],
        'column_names': columnNamesStr,
        'column_types': columnTypesStr,
        'na_strings': naStrings,
        'chunk_size': setup_result['chunk_size'],
        # No longer supported? how come these aren't in setup_result?
        'delete_on_done': params_dict['delete_on_done'],
        'blocking': params_dict['blocking'],
    }
    # HACK: if there are too many column names..don't print! it is crazy output
    # just check the output of parse setup. Don't worry about columnNames passed as params here.
    tooManyColNamesToPrint = setup_result['column_names'] and len(setup_result['column_names']) > 2000
    if tooManyColNamesToPrint:
        h2p.yellow_print("Not printing the parameters to Parse because the columnNames are too lengthy.")
        h2p.yellow_print("See sandbox/commands.log")

    # merge params_dict into parse_params
    # don't want =None to overwrite parse_params
    h2o_methods.check_params_update_kwargs(parse_params, params_dict, 'parse after merge into parse setup',
        print_params=not tooManyColNamesToPrint, ignoreNone=True)

    print "parse source_frames is length:", len(parse_params['source_frames'])
    # This can be null now? parseSetup doesn't return default colnames?
    # print "parse column_names is length:", len(parse_params['column_names'])

    # none of the kwargs passed to here!
    parse_result = self.do_json_request( jsonRequest="3/Parse.json", cmd='post', postData=parse_params, timeout=timeoutSecs)
    verboseprint("Parse result:", dump_json(parse_result))

    job_key = parse_result['job']['key']['name']
    hex_key = parse_params['destination_frame']

    # TODO: dislike having different shapes for noPoll and poll
    if noPoll:
        # ??
        h2o_sandbox.check_sandbox_for_errors()
        # return self.jobs(job_key)
        return parse_result

    # does Frame also, while polling
    if intermediateResults:
        key = hex_key
    else:
        key = None

    job_result = self.poll_job(job_key, timeoutSecs=timeoutSecs, key=key)

    if job_result:
        jobs = job_result['jobs'][0]
        description = jobs['description']
        dest = jobs['dest']
        msec = jobs['msec']
        status = jobs['status']
        progress = jobs['progress']
        dest_key = dest['name']

        # can condition this with a parameter if some FAILED are expected by tests.
        if status=='FAILED':
            print dump_json(job_result)
            raise Exception("Taking exception on parse job status: %s %s %s %s %s" % \
                (status, progress, msec, dest_key, description))

        return self.frames(dest_key)
    else:
        # ? we should always get a job_json result
        raise Exception("parse didn't get a job_result when it expected one")
        # return None

# TODO: remove .json
def frames(self, key=None, timeoutSecs=60, **kwargs):
    if not (key is None or isinstance(key, (basestring, Key))):
        raise Exception("frames: key should be string or Key type %s %s" % (type(key), key))

    params_dict = {
        'find_compatible_models': 0,
        'row_offset': 0, # is offset working yet?
        'row_count': 5,
    }
    '''
    Return a single Frame or all of the Frames in the h2o cluster.  The
    frames are contained in a list called "frames" at the top level of the
    result.  Currently the list is unordered.
    TODO:
    When find_compatible_models is implemented then the top level
    dict will also contain a "models" list.
    '''
    h2o_methods.check_params_update_kwargs(params_dict, kwargs, 'frames', False)

    # key can be type Key? (from h2o_xl) str(key) should return
    if key:
        if isinstance(key, Key):
            keyStr = key.frame
        else:
            keyStr = key
        result = self.do_json_request('3/Frames.json/' + keyStr, timeout=timeoutSecs, params=params_dict)
    else:
        result = self.do_json_request('3/Frames.json', timeout=timeoutSecs, params=params_dict)
    return result

# TODO: remove .json
def columns(self, key, timeoutSecs=10, **kwargs):
    '''
    Return the columns for a single Frame in the h2o cluster.
    '''
    params_dict = {
        'offset': 0,
        'len': 100
    }
    h2o_methods.check_params_update_kwargs(params_dict, kwargs, 'columns', True)

    result = self.do_json_request('3/Frames.json/' + key + '/columns', timeout=timeoutSecs, params=params_dict)
    return result

# TODO: remove .json
'''
Return a single column for a single Frame in the h2o cluster.
'''
def column(self, key, column, timeoutSecs=10, **kwargs):
    params_dict = {
        'offset': 0,
        'len': 100
    }
    h2o_methods.check_params_update_kwargs(params_dict, kwargs, 'column', True)

    result = self.do_json_request('3/Frames.json/' + key + '/columns/' + column, timeout=timeoutSecs, params=params_dict)
    return result

# TODO: remove .json
def summary(self, key, column="C1", timeoutSecs=10, **kwargs):
    '''
    Return the summary for a single column for a single Frame in the h2o cluster.
    '''
    params_dict = {
        # 'offset': 0,
        # 'len': 100
    }
    h2o_methods.check_params_update_kwargs(params_dict, kwargs, 'summary', True)

    result = self.do_json_request('3/Frames.json/%s/columns/%s/summary' % (key, column), timeout=timeoutSecs, params=params_dict)
    h2o_sandbox.check_sandbox_for_errors()
    return result

def delete_frame(self, key, ignoreMissingKey=True, timeoutSecs=60, **kwargs):
    '''
    Delete a frame on the h2o cluster, given its key.
    '''
    assert key is not None, '"key" parameter is null'

    result = self.do_json_request('/3/Frames.json/' + key, cmd='delete', timeout=timeoutSecs)

    # TODO: look for what?
    if not ignoreMissingKey and 'f00b4r' in result:
        raise ValueError('Frame key not found: ' + key)
    return result

def delete_frames(self, timeoutSecs=60, **kwargs):
    '''
    Delete all frames on the h2o cluster.
    '''
    parameters = { }
    result = self.do_json_request('/3/Frames.json', cmd='delete', timeout=timeoutSecs)
    return result

# TODO: remove .json
def model_builders(self, algo=None, timeoutSecs=10, **kwargs):
    '''
    Return a model builder or all of the model builders known to the
    h2o cluster.  The model builders are contained in a dictionary
    called "model_builders" at the top level of the result.  The
    dictionary maps algorithm names to parameters lists.  Each of the
    parameters contains all the metdata required by a client to
    present a model building interface to the user.

    if parameters = True, return the parameters?
    '''
    params_dict = {}
    h2o_methods.check_params_update_kwargs(params_dict, kwargs, 'model_builders', False)

    request = '3/ModelBuilders.json'
    if algo:
        request += "/" + algo

    result = self.do_json_request(request, timeout=timeoutSecs, params=params_dict)
    # verboseprint(request, "result:", dump_json(result))
    h2o_sandbox.check_sandbox_for_errors()
    return result

def validate_model_parameters(self, algo, training_frame, parameters, timeoutSecs=60, **kwargs):
    '''
    Check a dictionary of model builder parameters on the h2o cluster
    using the given algorithm and model parameters.
    '''
    assert algo is not None, '"algo" parameter is null'
    # Allow this now: assert training_frame is not None, '"training_frame" parameter is null'
    assert parameters is not None, '"parameters" parameter is null'

    model_builders = self.model_builders(timeoutSecs=timeoutSecs)
    assert model_builders is not None, "/ModelBuilders REST call failed"
    assert algo in model_builders['model_builders']
    builder = model_builders['model_builders'][algo]

    # TODO: test this assert, I don't think this is working. . .
    if training_frame is not None:
        frames = self.frames(key=training_frame)
        assert frames is not None, "/Frames/{0} REST call failed".format(training_frame)

        key_name = frames['frames'][0]['key']['name']
        assert key_name==training_frame, \
            "/Frames/{0} returned Frame {1} rather than Frame {2}".format(training_frame, key_name, training_frame)

        parameters['training_frame'] = training_frame

    # TODO: add parameter existence checks
    # TODO: add parameter value validation

    # FIX! why ignoreH2oError here?
    result = self.do_json_request('/3/ModelBuilders.json/' + algo + "/parameters", cmd='post',
        timeout=timeoutSecs, postData=parameters, ignoreH2oError=True, noExtraErrorCheck=True)

    verboseprint("model parameters validation: " + repr(result))
    return result

# should training_frame be required? or in parameters. same with destination_frame
# because validation_frame is in parameters
# destination_frame is old, model_id is new?
def build_model(self, algo, training_frame, parameters, destination_frame=None, model_id=None,
    timeoutSecs=60, noPoll=False, **kwargs):

    if 'destination_key' in kwargs:
        raise Exception('Change destination_key in build_model() to model_id')

    '''
    Build a model on the h2o cluster using the given algorithm, training
    Frame and model parameters.
    '''
    assert algo is not None, '"algo" parameter is null'
    assert training_frame is not None, '"training_frame" parameter is null'
    assert parameters is not None, '"parameters" parameter is null'

    # why always check that the algo is in here?
    model_builders = self.model_builders(timeoutSecs=timeoutSecs)
    assert model_builders is not None, "/ModelBuilders REST call failed"
    assert algo in model_builders['model_builders'], "%s %s" % (algo, [k for k in model_builders['model_builders']])
    builder = model_builders['model_builders'][algo]

    # TODO: test this assert, I don't think this is working. . .
    frames = self.frames(key=training_frame)
    assert frames is not None, "/Frames/{0} REST call failed".format(training_frame)

    key_name = frames['frames'][0]['frame_id']['name']
    assert key_name==training_frame, \
        "/Frames/{0} returned Frame {1} rather than Frame {2}".format(training_frame, key_name, training_frame)
    parameters['training_frame'] = training_frame

    if destination_frame is not None:
        print "destination_frame should be replaced by model_id now"
        parameters['model_id'] = destination_frame

    if model_id is not None:
        parameters['model_id'] = model_id

    print "build_model parameters", parameters
    start = time.time()
    result1 = self.do_json_request('/3/ModelBuilders.json/' + algo, cmd='post',
        timeout=timeoutSecs, postData=parameters)
    # make get overwritten after polling
    elapsed = time.time() - start
    verboseprint("build_model result", dump_json(result1))

    if noPoll:
        result = result1
    elif ('validation_error_count' in result1) and (result1['validation_error_count']>0):
        h2p.yellow_print("parameter error in model_builders: %s" % result1)
        # parameters validation failure
        # TODO: add schema_type and schema_version into all the schemas to make this clean to check
        result = result1
        # don't bother printing a time message
    elif 'exception_msg' in result1:
        h2p.yellow_print("exception msg in model_builders: %s" % result1['exception_msg'])
        result = result1
    else:
        job_result = result1['job']
        job_key = job_result['key']['name']
        verboseprint("build_model job_key: " + repr(job_key))

        job_result = self.poll_job(job_key, timeoutSecs=timeoutSecs)
        verboseprint(job_result)

        elapsed = time.time() - start
        print "ModelBuilders", algo,  "end on", training_frame, 'took', time.time() - start, 'seconds'
        print "%d pct. of timeout" % ((elapsed/timeoutSecs) * 100)

        if job_result:
            jobs = job_result['jobs'][0]
            description = jobs['description']
            dest = jobs['dest']
            msec = jobs['msec']
            status = jobs['status']
            progress = jobs['progress']

            # can condition this with a parameter if some FAILED are expected by tests.
            if status=='FAILED':
                print dump_json(job_result)
                raise Exception("Taking exception on build_model job status: %s %s %s %s" % \
                    (status, progress, msec, description))

            result = job_result
        else:
            # ? we should always get a job_json result
            raise Exception("build_model didn't get a job_result when it expected one")
            # return None

    verboseprint("result:", result)
    h2o_sandbox.check_sandbox_for_errors()
    result['python_elapsed'] = elapsed
    return result

def compute_model_metrics(self, model, frame, timeoutSecs=60, **kwargs):
    '''
    Score a model on the h2o cluster on the given Frame and return only the model metrics.
    '''
    assert model is not None, '"model" parameter is null'
    assert frame is not None, '"frame" parameter is null'

    models = self.models(key=model, timeoutSecs=timeoutSecs)
    assert models is not None, "/Models REST call failed"
    assert models['models'][0]['model_id']['name'] == model, "/Models/{0} returned Model {1} rather than Model {2}".format(model, models['models'][0]['key']['name'], model)

    # TODO: test this assert, I don't think this is working. . .
    frames = self.frames(key=frame)
    assert frames is not None, "/Frames/{0} REST call failed".format(frame)

    print "frames:", dump_json(frames)
    # is the name not there?
    # assert frames['frames'][0]['model_id']['name'] == frame, "/Frames/{0} returned Frame {1} rather than Frame {2}".format(frame, models['models'][0]['key']['name'], frame)

    result = self.do_json_request('/3/ModelMetrics.json/models/' + model + '/frames/' + frame, cmd='post', timeout=timeoutSecs)

    mm = result['model_metrics'][0]
    verboseprint("model metrics: " + repr(mm))
    h2o_sandbox.check_sandbox_for_errors()
    return mm

def predict(self, model, frame, timeoutSecs=60, **kwargs):
    assert model is not None, '"model" parameter is null'
    assert frame is not None, '"frame" parameter is null'

    models = self.models(key=model, timeoutSecs=timeoutSecs)
    assert models is not None, "/Models REST call failed"

    # FIX! what is right here now?
    # assert models['models'][0]['key']['name'] == model, "/Models/{0} returned Model {1} rather than Model {2}".format(model, models['models'][0]['key']['name'], model)

    # TODO: test this assert, I don't think this is working. . .
    frames = self.frames(key=frame)
    assert frames is not None, "/Frames/{0} REST call failed".format(frame)

    # FIX! what is right here now?
    # assert frames['frames'][0]['key']['name'] == frame, "/Frames/{0} returned Frame {1} rather than Frame {2}".format(frame, frames['frames'][0]['key']['name'], frame)

    result = self.do_json_request('/3/Predictions.json/models/' + model + '/frames/' + frame, cmd='post', timeout=timeoutSecs)

    h2o_sandbox.check_sandbox_for_errors()
    return result

def model_metrics(self, timeoutSecs=60, **kwargs):
    '''
    ModelMetrics list.
    '''
    result = self.do_json_request('/3/ModelMetrics.json', cmd='get', timeout=timeoutSecs)
    h2o_sandbox.check_sandbox_for_errors()
    return result

# TODO: remove .json
def models(self, key=None, timeoutSecs=10, **kwargs):
    '''
    Return all of the models in the h2o cluster, or a single model given its key.
    The models are contained in a list called "models" at the top level of the
    result.  Currently the list is unordered.
    TODO:
    When find_compatible_frames is implemented then the top level
    dict will also contain a "frames" list.
    '''
    params_dict = {
        'find_compatible_frames': False
    }
    h2o_methods.check_params_update_kwargs(params_dict, kwargs, 'models', True)

    if key:
        # result = self.do_json_request('3/Models.json', timeout=timeoutSecs, params=params_dict)
        # print "for ray:", dump_json(result)
        result = self.do_json_request('3/Models.json/' + key, timeout=timeoutSecs, params=params_dict)
    else:
        result = self.do_json_request('3/Models.json', timeout=timeoutSecs, params=params_dict)

    verboseprint("models result:", dump_json(result))
    h2o_sandbox.check_sandbox_for_errors()
    return result

def delete_model(self, key, ignoreMissingKey=True, timeoutSecs=60, **kwargs):
    '''
    Delete a model on the h2o cluster, given its key.
    '''
    assert key is not None, '"key" parameter is null'

    result = self.do_json_request('/3/Models.json/' + key, cmd='delete', timeout=timeoutSecs)

    # TODO: look for what?
    if not ignoreMissingKey and 'f00b4r' in result:
        raise ValueError('Model key not found: ' + key)

    verboseprint("delete_model result:", dump_json(result))
    return result

def delete_models(self, timeoutSecs=60, **kwargs):
    '''
    Delete all models on the h2o cluster.
    '''
    parameters = { }
    result = self.do_json_request('/3/Models.json', cmd='delete', timeout=timeoutSecs)
    return result

def endpoints(self, timeoutSecs=60, **kwargs):
    '''
    Fetch the list of REST API endpoints.
    '''
    parameters = { }
    result = self.do_json_request('/3/Metadata/endpoints.json', cmd='get', timeout=timeoutSecs)
    return result

def endpoint_by_number(self, num, timeoutSecs=60, **kwargs):
    '''
    Fetch the metadata for the given numbered REST API endpoint.
    '''
    parameters = { }
    result = self.do_json_request('/3/Metadata/endpoints.json/' + str(num), cmd='get', timeout=timeoutSecs)
    return result

def schemas(self, timeoutSecs=60, **kwargs):
    '''
    Fetch the list of REST API schemas.
    '''
    parameters = { }
    result = self.__do_json_request('/3/Metadata/schemas.json', cmd='get', timeout=timeoutSecs)

    return result

def schema(self, schemaname, timeoutSecs=60, **kwargs):
    '''
    Fetch the metadata for the given named REST API schema (e.g., FrameV2).
    '''
    parameters = { }
    result = self.__do_json_request('/3/Metadata/schemas.json/' + schemaname, cmd='get', timeout=timeoutSecs)

    return result


# -*- encoding: utf-8 -*-
"""
Rapids expressions. These are helper classes for H2OFrame.

:copyright: (c) 2016 H2O.ai
:license:   Apache License Version 2.0 (see LICENSE for details)
"""
from __future__ import division, print_function, absolute_import, unicode_literals

import collections
import copy
import gc
import math
import sys
import time
import numbers

import tabulate

import h2o
from h2o.backend.connection import H2OConnectionError
from h2o.utils.compatibility import *  # NOQA
from h2o.utils.compatibility import repr2, viewitems, viewvalues
from h2o.utils.shared_utils import _is_fr, _py_tmp_key
from h2o.model.model_base import ModelBase
from h2o.expr_optimizer import optimize

class ExprNode(object):
    """
    Composable Expressions: This module contains code for the lazy expression DAG.

    Execution Overview
    ------------------
      The job of ExprNode is to provide a layer of indirection to H2OFrame instances that
      are built of arbitrarily large, lazy expression DAGs. In order to do this job well,
      ExprNode must also track top-level entry points to the such DAGs, maintain a sane
      amount of state to know which H2OFrame instances are temporary (or not), and maintain
      a cache of H2OFrame properties (nrows, ncols, types, names, few rows of data).

      Top-Level Entry Points
      ----------------------
        An expression is declared top-level if it
          A) Computes and returns an H2OFrame to some on-demand call from somewhere
          B) An H2OFrame instance has more referrers than the 4 needed for the usual flow
             of python execution (see MAGIC_REF_COUNT below for more details).

      Sane Amount of State
      --------------------
        Instances of H2OFrame live and die by the state contained in the _ex field. The three
        pieces of state -- _op, _children, _cache -- are the fewest pieces of state (and no
        fewer) needed to unambiguously track temporary H2OFrame instances and prune
        them according to the usual scoping laws of python.

        If _cache._id is None, then this DAG has never been sent over to H2O, and there's
        nothing more to do when the object goes out of scope.

        If _cache._id is not None, then there has been some work done by H2O to compute the
        big data object sitting in H2O to which _id points. At the time that __del__ is
        called on this object, a determination to throw out the corresponding data in H2O or
        to keep that data is made by the None'ness of _children.

        tl;dr:
          If _cache._id is not None and _children is None, then do not delete in H2O cluster
          If _cache._id is not None and _children is not None, then do delete in H2O cluster

      H2OCache
      --------
        To prevent several unnecessary REST calls and unnecessary execution, a few of the
        oft-needed attributes of the H2OFrame are cached for convenience. The primary
        consumers of these cached values are __getitem__, __setitem__, and a few other
        H2OFrame ops that do argument validation or exchange (e.g. colnames for indices).
        There are more details available under the H2OCache class declaration.
    """

    # Magical count-of-5:   (get 2 more when looking at it in debug mode)
    #  2 for _get_ast_str frame, 2 for _get_ast_str local dictionary list, 1 for parent
    MAGIC_REF_COUNT = 5 if sys.gettrace() is None else 7  # M = debug ? 7 : 5

    # Flag to control application of local expression tree optimizations
    __ENABLE_EXPR_OPTIMIZATIONS__ = True

    def __init__(self, op="", *args):
        # assert isinstance(op, str), op
        self._op = op  # Base opcode string
        self._children = tuple(
            a._ex if _is_fr(a) else a for a in args)  # ast children; if not None and _cache._id is not None then tmp
        self._cache = H2OCache()  # ncols, nrows, names, types
        # try to fuse/simplify expression
        if self.__ENABLE_EXPR_OPTIMIZATIONS__:
            self._optimize()

    def _eager_frame(self):
        if not self._cache.is_empty(): return
        if self._cache._id is not None: return  # Data already computed under ID, but not cached locally
        self._eval_driver(True)

    def _eager_scalar(self):  # returns a scalar (or a list of scalars)
        if not self._cache.is_empty():
            assert self._cache.is_scalar()
            return self
        assert self._cache._id is None
        self._eval_driver(False)
        assert self._cache._id is None
        assert self._cache.is_scalar()
        return self._cache._data

    def _eager_map_frame(self):  # returns a scalar (or a list of scalars)
        self._eval_driver(False)
        return self._cache

    def _eval_driver(self, top):
        exec_str = self._get_ast_str(top)
        res = ExprNode.rapids(exec_str)
        if 'scalar' in res:
            if isinstance(res['scalar'], list):
                self._cache._data = [float(x) for x in res['scalar']]
            else:
                self._cache._data = None if res['scalar'] is None else float(res['scalar'])
        if 'string' in res: self._cache._data = res['string']
        if 'funstr' in res: raise NotImplementedError
        if 'key' in res:
            self._cache.nrows = res['num_rows']
            self._cache.ncols = res['num_cols']
        if 'map_keys' in res:
            self._cache.map_keys = res['map_keys']
            self._cache.frames = res['frames']
        return self

    def _optimize(self):
        while True:
            opt = optimize(self)
            if opt is not None:
                opt(ctx=None)
            else:
                break

    # Recursively build a rapids execution string.  Any object with more than
    # MAGIC_REF_COUNT referrers will be cached as a temp until the next client GC
    # cycle - consuming memory.  Do Not Call This except when you need to do some
    # other cluster operation on the evaluated object.  Examples might be: lazy
    # dataset time parse vs changing the global timezone.  Global timezone change
    # is eager, so the time parse as to occur in the correct order relative to
    # the timezone change, so cannot be lazy.
    #
    def _get_ast_str(self, top):
        if not self._cache.is_empty():  # Data already computed and cached; could a "false-like" cached value
            return str(self._cache._data) if self._cache.is_scalar() else self._cache._id
        if self._cache._id is not None:
            return self._cache._id  # Data already computed under ID, but not cached
        assert isinstance(self._children,tuple)
        exec_str = "({} {})".format(self._op, " ".join([ExprNode._arg_to_expr(ast) for ast in self._children]))
        gc_ref_cnt = len(gc.get_referrers(self))
        if top or gc_ref_cnt >= ExprNode.MAGIC_REF_COUNT:
            self._cache._id = _py_tmp_key(append=h2o.connection().session_id)
            exec_str = "(tmp= {} {})".format(self._cache._id, exec_str)
        return exec_str

    @staticmethod
    def _arg_to_expr(arg):
        if arg is None:
            return "[]"  # empty list
        if isinstance(arg, ExprNode):
            return arg._get_ast_str(False)
        if isinstance(arg, ASTId):
            return str(arg)
        if isinstance(arg, (list, tuple, range)):
            return "[%s]" % " ".join(repr2(x) for x in arg)
        if isinstance(arg, slice):
            start = 0 if arg.start is None else arg.start
            stop = float("nan") if arg.stop is None else arg.stop
            step = 1 if arg.step is None else arg.step
            assert start >= 0 and step >= 1 and (math.isnan(stop) or stop >= start + step)
            if step == 1:
                return "[%d:%s]" % (start, str(stop - start))
            else:
                return "[%d:%s:%d]" % (start, str((stop - start + step - 1) // step), step)
        if isinstance(arg, ModelBase):
            return arg.model_id
        # Number representation without Py2 L suffix enforced
        if isinstance(arg, numbers.Integral):
            return repr2(arg).strip('L')
        return repr2(arg)

    def __del__(self):
        try:
            if self._cache._id is not None and self._children is not None:
                ExprNode.rapids("(rm {})".format(self._cache._id))
        except (AttributeError, H2OConnectionError):
            pass

    def arg(self, idx):
        return self._children[idx]

    def args(self):
        return self._children

    def narg(self):
        return len(self._children)

    @staticmethod
    def _collapse_sb(sb):
        return ' '.join("".join(sb).replace("\n", "").split()).replace(" )", ")")

    def _debug_print(self, pprint=True):
        return "".join(self._2_string(sb=[])) if pprint else ExprNode._collapse_sb(self._2_string(sb=[]))

    def _to_string(self):
        return ' '.join(["(" + self._op] + [ExprNode._arg_to_expr(a) for a in self._children] + [")"])

    def _2_string(self, depth=0, sb=None):
        sb += ['\n', " " * depth, "(" + self._op, " "]
        if self._children is not None:
            for child in self._children:
                if _is_fr(child) and child._ex._cache._id is None:
                    child._ex._2_string(depth + 2, sb)
                elif _is_fr(child):
                    sb += ['\n', ' ' * (depth + 2), child._ex._cache._id]
                elif isinstance(child, ExprNode):
                    child._2_string(depth + 2, sb)
                else:
                    sb += ['\n', ' ' * (depth + 2), str(child)]
        sb += ['\n', ' ' * depth + ") "] + ['\n'] * (depth == 0)  # add a \n if depth == 0
        return sb

    def __repr__(self):
        return "<Expr(%s)%s%s>" % (
            " ".join([self._op] + [repr(x) for x in (self._children or [])]),
            "#%s" % self._cache._id if self._cache._id else "",
            "; scalar" if self._cache.is_scalar() else "",
        )

    @staticmethod
    def rapids(expr):
        """
        Execute a Rapids expression.

        :param expr: The rapids expression (ascii string).

        :returns: The JSON response (as a python dictionary) of the Rapids execution
        """
        return h2o.api("POST /99/Rapids", data={"ast": expr, "session_id": h2o.connection().session_id})

class ASTId:
    def __init__(self, name=None):
        if name is None:
            raise ValueError("Attempted to make ASTId with no name.")
        self.name = name

    def __repr__(self):
        return self.name

class H2OCache(object):
    def __init__(self):
        self._id = None
        self._nrows = -1
        self._ncols = -1
        self._types = None  # col types
        self._names = None  # col names
        self._data = None  # ordered dict of cached rows, or a scalar
        self._l = 0  # nrows cached

    @property
    def nrows(self):
        return self._nrows

    @nrows.setter
    def nrows(self, value):
        self._nrows = value

    def nrows_valid(self):
        return self._nrows >= 0

    @property
    def ncols(self):
        return self._ncols

    @ncols.setter
    def ncols(self, value):
        self._ncols = value

    def ncols_valid(self):
        return self._ncols >= 0

    @property
    def names(self):
        return self._names

    @names.setter
    def names(self, value):
        self._names = value

    def names_valid(self):
        return self._names is not None

    @property
    def types(self):
        return self._types

    @types.setter
    def types(self, value):
        self._types = value

    def types_valid(self):
        return self._types is not None

    @property
    def scalar(self):
        return self._data if self.is_scalar() else None

    @scalar.setter
    def scalar(self, value):
        self._data = value

    def __len__(self):
        return self._l

    def is_empty(self):
        return self._data is None

    def is_scalar(self):
        return not isinstance(self._data, dict)

    def is_valid(self):
        return (  # self._id is not None and
                not self.is_empty() and
                self.nrows_valid() and
                self.ncols_valid() and
                self.names_valid() and
                self.types_valid())

    def fill(self, rows=10, rows_offset=0, cols=-1, full_cols=-1, cols_offset=0, light=False):
        assert self._id is not None
        if self._data is not None:
            if rows <= len(self):
                return
        req_params = {
            "row_count": rows,
            "row_offset": rows_offset,
            "column_count" : cols,
            "full_column_count" : full_cols,
            "column_offset" : cols_offset
        }
        if light:
            endpoint = "/3/Frames/%s/light"
        else:
            endpoint = "/3/Frames/%s"
        res = h2o.api("GET " + endpoint % self._id, data=req_params)["frames"][0]
        self._l = rows
        self._nrows = res["rows"]
        self._ncols = res["total_column_count"]
        self._names = [c["label"] for c in res["columns"]]
        self._types = dict(zip(self._names, [c["type"] for c in res["columns"]]))
        self._fill_data(res)

    def _fill_data(self, json):
        self._data = collections.OrderedDict()
        for c in json["columns"]:
            c.pop('__meta')  # Redundant description ColV3
            c.pop('domain_cardinality')  # Same as len(c['domain'])
            sdata = c.pop('string_data')
            if sdata:
                c['data'] = sdata  # Only use data field; may contain either [str] or [real]
            # Data (not string) columns should not have a string in them.  However,
            # our NaNs are encoded as string literals "NaN" as opposed to the bare
            # token NaN, so the default python json decoder does not convert them
            # to math.nan.  Do that now.
            else:
                if c['data'] and (len(c['data']) > 0):  # orc file parse can return frame with zero rows
                    c['data'] = [float('nan') if x == "NaN" else x for x in c['data']]
            if c['data']:
                self._data[c.pop('label')] = c  # Label used as the Key
        return self

    #---- pretty printing ----

    def _tabulate(self, tablefmt="simple", rollups=False, rows=10):
        """Pretty tabulated string of all the cached data, and column names"""
        if not self.is_valid(): self.fill(rows=rows)
        # Pretty print cached data
        d = collections.OrderedDict()
        # If also printing the rollup stats, build a full row-header
        if rollups:
            col = next(iter(viewvalues(self._data)))  # Get a sample column
            lrows = len(col['data'])  # Cached rows being displayed
            d[""] = ["type", "mins", "mean", "maxs", "sigma", "zeros", "missing"] + list(map(str, range(lrows)))
        # For all columns...
        for k, v in viewitems(self._data):
            x = v['data']  # Data to display
            t = v["type"]  # Column type
            if t == "enum":
                domain = v['domain']  # Map to cat strings as needed
                x = ["" if math.isnan(idx) else domain[int(idx)] for idx in x]
            elif t == "time":
                x = ["" if math.isnan(z) else time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime(z / 1000)) for z in x]
            if rollups:  # Rollups, if requested
                mins = v['mins'][0] if v['mins'] and v["type"] != "enum" else None
                maxs = v['maxs'][0] if v['maxs'] and v["type"] != "enum" else None
                #Cross check type with mean and sigma. Set to None if of type enum.
                if v['type'] == "enum":
                    v['mean'] = v['sigma'] = v['zero_count'] = None
                x = [v['type'], mins, v['mean'], maxs, v['sigma'], v['zero_count'], v['missing_count']] + x
            d[k] = x  # Insert into ordered-dict
        return tabulate.tabulate(d, headers="keys", tablefmt=tablefmt)

    def flush(self):  # flush everything but the frame_id
        fr_id = self._id
        self.__dict__ = H2OCache().__dict__.copy()
        self._id = fr_id
        return self

    def fill_from(self, cache):
        assert isinstance(cache, H2OCache)
        cur_id = self._id
        self.__dict__ = copy.copy(cache.__dict__)  # copy.deepcopy is buggy :( https://bugs.python.org/issue16251
        self._data = None
        self._id = cur_id

    def dummy_fill(self):
        self._id = "dummy"
        self._nrows = 0
        self._ncols = 0
        self._names = []
        self._types = {}
        self._data = {}


#!/usr/bin/python

import argparse
import boto
import os, time, sys, socket
import h2o_cmd
import h2o2 as h2o
import h2o_hosts
import json
import commands
import traceback

'''
    Simple EC2 utility:

       * to setup clooud of 5 nodes: ./ec2_cmd.py create --instances 5
       * to terminated the cloud   : ./ec2_cmd.py terminate --hosts <host file returned by previous command>
'''

DEFAULT_NUMBER_OF_INSTANCES = 4
DEFAULT_HOSTS_FILENAME = 'ec2-config-{0}.json'
DEFAULT_REGION = 'us-east-1'
DEFAULT_INSTANCE_NAME='node_{0}'.format(os.getenv('USER'))
ADVANCED_SSH_OPTIONS='-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o GlobalKnownHostsFile=/dev/null'

'''
Default EC2 instance setup
'''
DEFAULT_EC2_INSTANCE_CONFIGS = {
  'us-east-1':{
              'image_id'        : 'ami-cf5132a6', #'ami-30c6a059', #'ami-b85cc4d1', # 'ami-cd9a11a4',
              'security_groups' : [ 'MrJenkinsTest' ],
              'key_name'        : 'mrjenkins_test',
              'instance_type'   : 'm1.xlarge',
              'region'          : 'us-east-1',
              'pem'             : '~/.ec2/keys/mrjenkins_test.pem',
              'username'        : '0xdiag',
              'aws_credentials' : '~/.ec2/AwsCredentials.properties',
              'hdfs_config'     : '~/.ec2/core-site.xml',
             },
  'us-west-1':{
              'image_id'        : 'ami-a6cbe6e3', # 'ami-cd9a11a4',
              'security_groups' : [ 'MrJenkinsTest' ],
              'key_name'        : 'mrjenkins_test',
              'instance_type'   : 'm1.xlarge',
              'pem'             : '~/.ec2/keys/mrjenkins_test.pem',
              'username'        : '0xdiag',
              'aws_credentials' : '~/.ec2/AwsCredentials.properties',
              'hdfs_config'     : '~/.ec2/core-site.xml',
              'region'          : 'us-west-1',
             },
}

''' Memory mappings for instance kinds '''
MEMORY_MAPPING = {
    'm1.small'     : { 'xmx' : 1  },
    'm1.medium'    : { 'xmx' : 3  },
    'm1.large'     : { 'xmx' : 5  },  # $0.24/hr
    'm1.xlarge'    : { 'xmx' : 12 },  # $0.48/hr
    'm2.xlarge'    : { 'xmx' : 13 },  # $0.41/hr
    'm2.2xlarge'   : { 'xmx' : 30 },  # $0.82/hr
    'm2.4xlarge'   : { 'xmx' : 64 },  # $1.64/hr
    'm3.medium'    : { 'xmx' : 3  },
    'm3.large'     : { 'xmx' : 5  },
    'm3.xlarge'    : { 'xmx' : 12 },  # $0.50/hr
    'm3.2xlarge'   : { 'xmx' : 26 },  # $1.00/hr
    'c1.medium'    : { 'xmx' : 1  },
    'c1.xlarge'    : { 'xmx' : 6  },  # $0.58/hr
    'c3.large'     : { 'xmx' : 3  },
    'c3.xlarge'    : { 'xmx' : 5  },  # $0.30/hr
    'c3.2xlarge'   : { 'xmx' : 12 },
    'c3.4xlarge'   : { 'xmx' : 27 },
    'c3.8xlarge'   : { 'xmx' : 56 },
    'cc2.8xlarge'  : { 'xmx' : 56 },
    'hi1.4xlarge'  : { 'xmx' : 56 },  # $3.10/hr 60.5GB dram. 8 cores/8 threads. 2TB SSD. 10GE
    'cr1.8xlarge'  : { 'xmx' : 230},  # $4.60/hr 244GB dram. 2 E5-2670. 240GB SSD. 10GE
    'g2.2xlarge'   : { 'xmx' : 12 },
    'i2.xlarge'    : { 'xmx' : 27 },
    'i2.2xlarge'   : { 'xmx' : 57 },
    'cg1.4xlarge'  : { 'xmx' : 19 },
    'i2.4xlarge'   : { 'xmx' : 116},
    'hs1.8xlarge'  : { 'xmx' : 112},
    'i2.8xlarge'   : { 'xmx' : 236},
}

''' EC2 API default configuration. The corresponding values are replaces by EC2 user config. '''
EC2_API_RUN_INSTANCE = {
'image_id'        :None,
'min_count'       :1,
'max_count'       :1,
'key_name'        :None,
'security_groups' :None,
'user_data'       :None,
'addressing_type' :None,
'instance_type'   :None,
'placement'       :None,
'monitoring_enabled':False,
'subnet_id'       :None,
'block_device_map':None,
'disable_api_termination':False,
'instance_initiated_shutdown_behavior':None
}

def inheritparams(parent, kid):
    newkid = {}
    for k,v in kid.items():
        if parent.has_key(k):
            newkid[k] = parent[k]

    return newkid

def find_file(base):
    f = base
    if not os.path.exists(f): f = os.path.expanduser(base)
    if not os.path.exists(f): f = os.path.expanduser("~")+ '/' + base
    if not os.path.exists(f):
        return None
    return f

''' Returns a boto connection to given region '''
def ec2_connect(region):
    check_required_env_variables()
    import boto.ec2
    conn = boto.ec2.connect_to_region(region)
    if not conn:
        raise Exception("\033[91m[ec2] Cannot create EC2 connection into {0} region!\033[0m".format(region))

    return conn

def check_required_env_variables():
    ok = True
    if not os.environ['AWS_ACCESS_KEY_ID']:
        warn("AWS_ACCESS_KEY_ID need to be defined!")
        ok = False
    if not os.environ['AWS_SECRET_ACCESS_KEY']:
        warn("AWS_SECRET_ACCESS_KEY need to be defined!")
        ok = False

    if not ok: raise Exception("\033[91m[ec2] Missing AWS environment variables!\033[0m")

''' Run number of EC2 instance.
Waits forthem and optionaly waits for ssh service.
'''
def run_instances(count, ec2_config, region, waitForSSH=True, tags=None):
    '''Create a new reservation for count instances'''

    ec2params = inheritparams(ec2_config, EC2_API_RUN_INSTANCE)
    ec2params.setdefault('min_count', count)
    ec2params.setdefault('max_count', count)

    reservation = None
    conn = ec2_connect(region)
    try:
        reservation = conn.run_instances(**ec2params)
        log('Reservation: {0}'.format(reservation.id))
        log('Waiting for {0} EC2 instances {1} to come up, this can take 1-2 minutes.'.format(len(reservation.instances), reservation.instances))
        start = time.time()
        time.sleep(1)
        for instance in reservation.instances:
            while instance.update() == 'pending':
               time.sleep(1)
               h2o_cmd.dot()

            if not instance.state == 'running':
                raise Exception('\033[91m[ec2] Error waiting for running state. Instance is in state {0}.\033[0m'.format(instance.state))

        log('Instances started in {0} seconds'.format(time.time() - start))
        log('Instances: ')
        for inst in reservation.instances: log("   {0} ({1}) : public ip: {2}, private ip: {3}".format(inst.public_dns_name, inst.id, inst.ip_address, inst.private_ip_address))

        if waitForSSH:
            # kbn: changing to private address, so it should fail if not in right domain
            # used to have the public ip address
            wait_for_ssh([ i.private_ip_address for i in reservation.instances ])

        # Tag instances
        try:
            if tags:
                conn.create_tags([i.id for i in reservation.instances], tags)
        except:
            warn('Something wrong during tagging instances. Exceptions IGNORED!')
            print sys.exc_info()
            pass

        return reservation
    except:
        print "\033[91mUnexpected error\033[0m :", sys.exc_info()
        if reservation:
            terminate_reservation(reservation, region)
        raise

''' Wait for ssh port
'''
def ssh_live(ip, port=22):
    return h2o_cmd.port_live(ip,port)

def terminate_reservation(reservation, region):
    terminate_instances([ i.id for i in reservation.instances ], region)

def terminate_instances(instances, region):
    '''terminate all the instances given by its ids'''
    if not instances: return
    conn = ec2_connect(region)
    log("Terminating instances {0}.".format(instances))
    conn.terminate_instances(instances)
    log("Done")

def stop_instances(instances, region):
    '''stop all the instances given by its ids'''
    if not instances: return
    conn = ec2_connect(region)
    log("Stopping instances {0}.".format(instances))
    conn.stop_instances(instances)
    log("Done")

def start_instances(instances, region):
    '''Start all the instances given by its ids'''
    if not instances: return
    conn = ec2_connect(region)
    log("Starting instances {0}.".format(instances))
    conn.start_instances(instances)
    log("Done")

def reboot_instances(instances, region):
    '''Reboot all the instances given by its ids'''
    if not instances: return
    conn = ec2_connect(region)
    log("Rebooting instances {0}.".format(instances))
    conn.reboot_instances(instances)
    log("Done")

def wait_for_ssh(ips, port=22, skipAlive=True, requiredsuccess=3):
    ''' Wait for ssh service to appear on given hosts'''
    log('Waiting for SSH on following hosts: {0}'.format(ips))
    for ip in ips:
        if not skipAlive or not ssh_live(ip, port):
            log('Waiting for SSH on instance {0}...'.format(ip))
            count = 0
            while count < requiredsuccess:
                if ssh_live(ip, port):
                    count += 1
                else:
                    count = 0
                time.sleep(1)
                h2o_cmd.dot()

def dump_hosts_config(ec2_config, reservation, filename=DEFAULT_HOSTS_FILENAME, save=True, h2o_per_host=1):
    if not filename: filename=DEFAULT_HOSTS_FILENAME

    cfg = {}
    f = find_file(ec2_config['aws_credentials'])
    if f: cfg['aws_credentials'] = f
    else: warn_file_miss(ec2_config['aws_credentials'])
    f = find_file(ec2_config['pem'])
    if f: cfg['key_filename'] = f
    else: warn_file_miss(ec2_config['key_filename'])
    f = find_file(ec2_config['hdfs_config'])
    if f: cfg['hdfs_config']  = f
    else: warn_file_miss(ec2_config['hdfs_config'])
    cfg['username']        = ec2_config['username']
    cfg['use_flatfile']    = True
    cfg['h2o_per_host']    = h2o_per_host
    cfg['java_heap_GB']    = MEMORY_MAPPING[ec2_config['instance_type']]['xmx']
    cfg['java_extra_args'] = '' # No default Java arguments '-XX:MaxDirectMemorySize=1g'
    cfg['base_port']       = 54321
    cfg['ip'] = [ i.private_ip_address for i in reservation.instances ]
    cfg['ec2_instances']   = [ { 'id': i.id, 'private_ip_address': i.private_ip_address, 'public_ip_address': i.ip_address, 'public_dns_name': i.public_dns_name } for i in reservation.instances ]
    cfg['ec2_reservation_id']  = reservation.id
    cfg['ec2_region']      = ec2_config['region']
    # cfg['redirect_import_folder_to_s3_path'] = True
    # New! we can redirect import folder to s3n thru hdfs, now (ec2)
    # cfg['redirect_import_folder_to_s3n_path'] = True
    # kbn 9/28/14..change it back to s3 to see what breaks
    cfg['redirect_import_folder_to_s3_path'] = True
    # put ssh commands into comments
    cmds = get_ssh_commands(ec2_config, reservation)
    idx  = 1
    for cmd in cmds:
        cfg['ec2_comment_ssh_{0}'.format(idx)] = cmd
        idx += 1

    if save:
        # save config
        filename = filename.format(reservation.id)
        with open(filename, 'w+') as f:
            f.write(json.dumps(cfg, indent=4))

        log("Host config dumped into {0}".format(filename))
        log("To terminate instances call:")
        log("\033[93mpython ec2_cmd.py terminate --hosts {0}\033[0m".format(filename))
        log("To start H2O cloud call:")
        log("\033[93mpython ec2_cmd.py start_h2o --hosts {0}\033[0m".format(filename))
        log("To watch cloud in browser follow address:")
        log("   http://{0}:{1}".format(reservation.instances[0].public_dns_name, cfg['base_port']))

    return (cfg, filename)

def get_ssh_commands(ec2_config, reservation, ssh_options=""):
    cmds = []
    if not ssh_options: ssh_options = ""
    for i in reservation.instances:
        cmds.append( "ssh -i ~/.ec2/keys/mrjenkins_test.pem {1} ubuntu@{0}".format(i.private_ip_address,ssh_options) )
    return cmds

def dump_ssh_commands(ec2_config, reservation):
    cmds = get_ssh_commands(ec2_config, reservation)
    for cmd in cmds:
        print cmd

# for cleaning /tmp after it becomes full, or any one string command (can separate with semicolon)
def execute_using_ssh_commands(ec2_config, reservation, command_string='df'):
    if not command_string: log("Nothing to execute. Exiting...")
    cmds = get_ssh_commands(ec2_config, reservation, ADVANCED_SSH_OPTIONS)
    for cmd in cmds:
        c = cmd + " '" + command_string + "'"
        print "\n"+c
        ret,out = commands.getstatusoutput(c)
        print out

def load_ec2_region(region):
    for r in DEFAULT_EC2_INSTANCE_CONFIGS:
        if r == region:
            return region

    raise Exception('\033[91m[ec2] Unsupported EC2 region: {0}. The available regions are: {1}\033[0m'.format(region, [r for r in DEFAULT_EC2_INSTANCE_CONFIGS ]))

def load_ec2_config(config_file, region, instance_type=None, image_id=None):
    if config_file:
        f = find_file(config_file)
        with open(f, 'rb') as fp:
             ec2_cfg = json.load(fp)
    else:
        ec2_cfg = {}

    for k,v in DEFAULT_EC2_INSTANCE_CONFIGS[region].items():
        ec2_cfg.setdefault(k, v)

    if instance_type: ec2_cfg['instance_type'] = instance_type
    if image_id     : ec2_cfg['image_id'     ] = image_id

    return ec2_cfg

def load_ec2_reservation(reservation, region):
    conn = ec2_connect(region)
    lr   = [ r for r in conn.get_all_instances() if r.id == reservation ]
    if not lr: raise Exception('Reservation id {0} not found !'.format(reservation))

    return lr[0]

def load_hosts_config(config_file):
    f = find_file(config_file)
    with open(f, 'rb') as fp:
         host_cfg = json.load(fp)
    return host_cfg

def log(msg):
    print "\033[92m[ec2] \033[0m", msg

def warn(msg):
    print "\033[92m[ec2] \033[0m \033[91m{0}\033[0m".format(msg)

def warn_file_miss(f):
    warn("File {0} is missing! Please update the generated config manually.".format(f))

def invoke_hosts_action(action, hosts_config, args, ec2_reservation=None):
    ids = [ inst['id'] for inst in hosts_config['ec2_instances'] ]
    ips = [ inst['private_ip_address'] for inst in hosts_config['ec2_instances'] ]
    region = hosts_config['ec2_region']

    if (action == 'terminate'):
        terminate_instances(ids, region)
    elif (action == 'stop'):
        stop_instances(ids, region)
    elif (action == 'reboot'):
        reboot_instances(ids, region)
        wait_for_ssh(ips, skipAlive=False, requiredsuccess=10)
    elif (action == 'start'):
        start_instances(ids, region)
        # FIXME after start instances receive new IPs: wait_for_ssh(ips)
    elif (action == 'distribute_h2o'):
        pass
    elif (action == 'start_h2o'):
        try:
            h2o.config_json = args.hosts
            log("Starting H2O cloud...")
            h2o_hosts.build_cloud_with_hosts(timeoutSecs=120, retryDelaySecs=5)
            h2o.touch_cloud()
            log("Cloud started. Let's roll!")
            log("You can start for example here \033[93mhttp://{0}:{1}\033[0m".format(hosts_config['ec2_instances'][0]['public_dns_name'],hosts_config['base_port']))
            if args.timeout:
                log("Cloud will shutdown after {0} seconds or use Ctrl+C to shutdown it.".format(args.timeout))
                time.sleep(args.timeout)
            else:
                log("To kill the cloud please use Ctrl+C as usual.")
                while (True): time.sleep(3600)
        except:
            print traceback.format_exc()
        finally:
            log("Goodbye H2O cloud...")
            h2o.tear_down_cloud()
            log("Cloud is gone.")

    elif (action == 'stop_h2o'):
        pass
    elif (action == 'clean_tmp'):
        execute_using_ssh_commands(hosts_config, ec2_reservation, command_string='sudo rm -rf /tmp/*; df')
    elif (action == 'nexec'):
        execute_using_ssh_commands(hosts_config, ec2_reservation, command_string=args.cmd)

def report_reservations(region, reservation_id=None):
    conn = ec2_connect(region)
    reservations = conn.get_all_instances()
    if reservation_id: reservations = [i for i in reservations if i.id == reservation_id ]
    log('Reservations:')
    for r in reservations: report_reservation(r); log('')

def report_reservation(r):
    log('  Reservation : {0}'.format(r.id))
    log('  Instances   : {0}'.format(len(r.instances)))
    for i in r.instances:
        log('    [{0} : {4}] {5} {1}/{2}/{3} {6}'.format(i.id, i.public_dns_name, i.ip_address, i.private_ip_address,i.instance_type, format_state(i.state), format_name(i.tags)))

def format_state(state):
    if state == 'stopped': return '\033[093mSTOPPED\033[0m'
    if state == 'running': return '\033[092mRUNNING\033[0m'
    if state == 'terminated': return '\033[090mTERMINATED\033[0m'
    return state.upper()

def format_name(tags):
    if 'Name' in tags: return '\033[91m<{0}>\033[0m'.format(tags['Name'])
    else: return '\033[94m<NONAME>\033[0m'

def merge_reservations(reservations):
    pass

def create_tags(**kwargs):
    tags = { }
    for key,value in kwargs.iteritems():
        tags[key] = value

    return tags

def main():
    parser = argparse.ArgumentParser(description='H2O EC2 instances launcher')
    parser.add_argument('action', choices=['help', 'demo', 'create', 'terminate', 'stop', 'reboot', 'start', 'distribute_h2o', 'start_h2o', 'show_defaults', 'dump_reservation', 'show_reservations', 'clean_tmp','nexec'], help='EC2 instances action!')
    parser.add_argument('-c', '--config',    help='Configuration file to configure NEW EC2 instances (if not specified default is used - see "show_defaults")', type=str, default=None)
    parser.add_argument('-i', '--instances', help='Number of instances to launch', type=int, default=DEFAULT_NUMBER_OF_INSTANCES)
    parser.add_argument('-H', '--hosts',     help='Hosts file describing existing "EXISTING" EC2 instances ', type=str, default=None)
    parser.add_argument('-r', '--region',    help='Specifies target create region', type=str, default=DEFAULT_REGION)
    parser.add_argument('--reservation',     help='Reservation ID, for example "r-1824ec65"', type=str, default=None)
    parser.add_argument('--name',            help='Name for launched instances', type=str, default=DEFAULT_INSTANCE_NAME)
    parser.add_argument('--timeout',         help='Timeout in seconds.', type=int, default=None)
    parser.add_argument('--instance_type',   help='Enfore a type of EC2 to launch (e.g., m2.2xlarge).', type=str, default=None)
    parser.add_argument('--cmd',             help='Shell command to be executed by nexec.', type=str, default=None)
    parser.add_argument('--image_id',        help='Override defautl image_id', type=str, default=None)
    parser.add_argument('--h2o_per_host',    help='Number of JVM launched per node', type=int, default=1)
    args = parser.parse_args()

    ec2_region = load_ec2_region(args.region)
    if (args.action == 'help'):
        parser.print_help()
    elif (args.action == 'create' or args.action == 'demo'):
        ec2_config = load_ec2_config(args.config, ec2_region, args.instance_type, args.image_id)
        tags       = create_tags(Name=args.name)
        log("EC2 region : {0}".format(ec2_region))
        log("EC2 itype  : {0}".format(ec2_config['instance_type']))
        log("EC2 ami    : {0}".format(ec2_config['image_id']))
        log("EC2 config : {0}".format(ec2_config))
        log("Instances  : {0}".format(args.instances))
        log("Tags       : {0}".format(tags))
        reservation = run_instances(args.instances, ec2_config, ec2_region, tags=tags)

        hosts_cfg, filename   = dump_hosts_config(ec2_config, reservation, filename=args.hosts, h2o_per_host=args.h2o_per_host)
        dump_ssh_commands(ec2_config, reservation)
        if (args.action == 'demo'):
            args.hosts = filename
            try:
                invoke_hosts_action('start_h2o', hosts_cfg, args)
            finally:
                invoke_hosts_action('terminate', hosts_cfg, args)

    elif (args.action == 'show_defaults'):
        print
        print "\033[92mConfig\033[0m : {0}".format(json.dumps(DEFAULT_EC2_INSTANCE_CONFIGS,indent=2))
        print "\033[92mInstances\033[0m         : {0}".format(DEFAULT_NUMBER_OF_INSTANCES)
        print "\033[92mSupported regions\033[0m : {0}".format( [ i for i in DEFAULT_EC2_INSTANCE_CONFIGS ] )
        print
    elif (args.action == 'merge_reservations'):
        merge_reservations(args.reservations, args.region)
    elif (args.action == 'dump_reservation'):
        ec2_config = load_ec2_config(args.config, ec2_region)
        ec2_reservation = load_ec2_reservation(args.reservation, ec2_region)
        dump_hosts_config(ec2_config, ec2_reservation, filename=args.hosts)
    elif (args.action == 'show_reservations'):
        report_reservations(args.region, args.reservation)
    else:
        if args.hosts:
            hosts_config = load_hosts_config(args.hosts)
            if hosts_config['ec2_reservation_id']: ec2_reservation = load_ec2_reservation(hosts_config['ec2_reservation_id'], ec2_region)
            else: ec2_reservation = None
        elif args.reservation: # TODO allows for specifying multiple reservations and merge them
            ec2_config      = load_ec2_config(args.config, ec2_region)
            ec2_reservation = load_ec2_reservation(args.reservation, ec2_region)
            hosts_config,_  = dump_hosts_config(ec2_config, ec2_reservation, save=False)
        invoke_hosts_action(args.action, hosts_config, args, ec2_reservation)
        if (args.action == 'terminate' and args.hosts):
            log("Deleting {0} host file.".format(args.hosts))
            os.remove(args.hosts)

if __name__ == '__main__':
    main()


# -*- encoding: utf-8 -*-
"""
Collection of utilities for debugging.

:copyright: (c) 2016 H2O.ai
:license:   Apache License Version 2.0 (see LICENSE for details)
"""
from __future__ import division, print_function, absolute_import, unicode_literals

import colorama
import inspect
import re
import sys
from colorama import Style, Fore
from types import ModuleType

from h2o.exceptions import H2OJobCancelled, H2OSoftError
from h2o.utils.compatibility import *  # NOQA
from h2o.utils.compatibility import viewkeys

# Nothing to import; this module's only job is to install an exception hook for debugging.
__all__ = ()

def get_tb():
    return get_tb.tb

def err(msg=""):
    """Helper function for printing to stderr."""
    print(msg, file=sys.stderr)

# In case any other module installs its own exception hook, try to play nicely and use that when appropriate
_prev_except_hook = sys.excepthook

def _except_hook(exc_type, exc_value, exc_tb):
    """
    This is an advanced exception-handling hook function, that is designed to supercede the standard Python's
    exception handler. It offers several enhancements:
        * Clearer and more readable format for the exception message and the traceback.
        * Decorators are filtered out from the traceback (if they declare their implementation function to
          have name "decorator_invisible").
        * Local variables in all execution frames are also printed out.
        * Print out server-side stacktrace for exceptions that carry this information (e.g. H2OServerError).

    Some documentation about the objects used herein:

    "traceback" type (types.TracebackType):  -- stack traceback of an exception
        tb_frame:  execution frame object at the current level
        tb_lasti:  index of the last attempted instruction in the bytecode
        tb_lineno: line number in the source code
        tb_next:   next level in the traceback stack (towards the frame where exception occurred)

    "frame" type (types.FrameType):  -- the execution frame
        f_back: previous stack frame (toward the caller)
        f_code: code object being executed
        f_locals: dictionary of all local variables
        f_globals: dictionary of all global variables
        f_builtins: dictionary of built-in names
        f_lineno: line number
        f_lasti: current instruction being executed (index into the f_code.co_code)
        f_restricted: ?
        f_trace: (function) function called at the start of each code line (settable!)
        f_exc_type, f_exc_value, f_exc_traceback: [Py2 only] most recent exception triple

    "code" type (types.CodeType):  -- byte-compiled code
        co_name: function name
        co_argcount: number of positional arguments
        co_nlocals: number of local variables inside the function
        co_varnames: names of all local variables
        co_cellvars: names of variables referenced by nested functions
        co_freevars: names of free variables used by nested functions
        co_code: raw bytecode
        co_consts: literals used in the bytecode
        co_names: names used in the bytecode
        co_filename: name of the file that contains the function
        co_firstlineno: line number where the function starts
        co_lnotab: encoded offsets to line numbers within the bytecode
        co_stacksize: required stack size
        co_flags: function flags. Bit 2 = function uses *args; bit 3 = function uses **kwargs

    Note that when the Python code re-raises an exception (by calling `raise` without arguments, which semantically
    means "raise the exception again as if it wasn't caught"), then tb.tb_lineno will contain the line number of the
    original exception, whereas tb.tb_frame.f_lineno will be the line number where the exception was re-raised.

    :param exc_type: exception class
    :param exc_value: the exception instance object
    :param exc_tb: stacktrace at the point of the exception. The "traceback" type is actually a linked list,
                   and this object represents the beginning of this list (i.e. corresponds to the execution
                   frame of the outermost expression being evaluated). We need to walk down the list (by repeatedly
                   moving to exc_tb.tb_next) in order to find the execution frame where the actual exception occurred.
    """
    if isinstance(exc_value, H2OJobCancelled):
        return
    if isinstance(exc_value, H2OSoftError):
        _handle_soft_error(exc_type, exc_value, exc_tb)
    else:
        _prev_except_hook(exc_type, exc_value, exc_tb)

    # Everythin else is disabled for now, because it generates too much output due to bugs in H2OFrame implementation
    return

    import linecache
    if not exc_tb:  # Happens on SyntaxError exceptions
        sys.__excepthook__(exc_type, exc_value, exc_tb)
        return
    get_tb.tb = exc_tb

    err("\n================================ EXCEPTION INFO ================================\n")
    if exc_type != type(exc_value):
        err("Exception type(s): %s / %s" % (exc_type, type(exc_value)))

    # If the exception contains .stacktrace attribute (representing server-side stack trace) then print it too.
    for arg in exc_value.args:
        if hasattr(arg, "stacktrace"):
            err("[SERVER STACKTRACE]")
            for line in arg.stacktrace:
                err("    %s" % line.strip())
            err()

    # Print local frames
    err("[LOCAL FRAMES]")
    err("Omitted: imported modules, class declarations, __future__ features, None-valued")
    tb = exc_tb
    while tb:
        tb_line = tb.tb_lineno
        frame = tb.tb_frame
        frame_file = frame.f_code.co_filename
        frame_func = frame.f_code.co_name
        frame_locl = frame.f_locals
        tb = tb.tb_next
        if frame_func == "decorator_invisible": continue
        if frame_func == "__getattribute__": continue
        if not frame_locl: continue
        err("\n  Within %s() line %s in file %s:" % (frame_func, tb_line, frame_file))
        for key in sorted(viewkeys(frame_locl), reverse=True):
            if key.startswith("__") and key.endswith("__"): continue
            value = frame_locl[key]
            if value is None: continue  # do not print uninitialized variables
            if hasattr(value, "__class__"):
                if value.__class__ is ModuleType: continue  # omit imported modules
                if value.__class__ is type: continue  # omit class declarations (new-style classes only)
                if value.__class__ is print_function.__class__: continue  # omit __future__ declarations
            try:
                strval = str(value)
                n_lines = strval.count("\n")
                if n_lines > 1:
                    strval = "%s... (+ %d line%s)" % (strval[:strval.index("\n")], n_lines - 1,
                                                      "s" if n_lines > 2 else "")
                err("%25s: %s" % (key, strval))
            except:
                err("%25s: <UNABLE TO PRINT VALUE>" % key)
    err()

    # Print the traceback
    err("[STACKTRACE]")
    last_file = None
    tb = exc_tb
    prev_info = None
    skip_frames = 0
    while tb:
        tb_lineno = tb.tb_lineno
        frame = tb.tb_frame
        frame_file = frame.f_code.co_filename
        frame_func = frame.f_code.co_name
        frame_glob = frame.f_globals
        tb = tb.tb_next
        if frame_func == "decorator_invisible": continue
        if frame_func == "__getattribute__": continue
        if (tb_lineno, frame_file) == prev_info:
            skip_frames += 1
            continue
        else:
            if skip_frames:
                err("    %20s   ... +%d nested calls" % ("", skip_frames))
            skip_frames = 0
            prev_info = (tb_lineno, frame_file)

        if frame_file != last_file:
            last_file = frame_file
            err("\n  File %s:" % frame_file)
        line_txt = linecache.getline(frame_file, tb_lineno, frame_glob)
        err("    %20s() #%04d  %s" % (frame_func, tb_lineno, line_txt.strip()))
    if skip_frames:
        err("    %20s   ... +%d nested calls" % ("", skip_frames))
    err()

    # Print the exception message
    tb = exc_tb
    while tb.tb_next:
        tb = tb.tb_next
    print("[EXCEPTION]", file=sys.stderr)
    print("  %s: %s" % (exc_value.__class__.__name__, str(exc_value)), file=sys.stderr)
    print("  at line %d in %s\n" % (tb.tb_lineno, tb.tb_frame.f_code.co_filename), file=sys.stderr)

    # There was a warning in {https://docs.python.org/2/library/sys.html} that storing traceback object in a local
    # variable may cause a circular reference; so we explicitly delete these vars just in case.
    del tb
    del exc_tb

# Install the enhanced exception hook into the system.
# The original exception hook is stored at sys.__excepthook__, and it will get called if our custom exception
# handling function itself raises an exception.
sys.excepthook = _except_hook

def _handle_soft_error(exc_type, exc_value, exc_tb):
    colorama.init()
    err((Fore.LIGHTRED_EX + "%s: %s" + Style.RESET_ALL) % (exc_type.__name__, exc_value))

    # Convert to the list of frames
    tb = exc_tb
    frames = []
    while tb:
        frames.append(tb.tb_frame)
        tb = tb.tb_next

    i0 = len(frames) - 1 - getattr(exc_value, "skip_frames", 0)
    indent = " " * (len(exc_type.__name__) + 2)
    for i in range(i0, 0, -1):
        co = frames[i].f_code
        func = _find_function_from_code(frames[i - 1], co)
        fullname = _get_method_full_name(func) if func else "???." + co.co_name
        highlight = getattr(exc_value, "var_name", None) if i == i0 else None
        args_str = _get_args_str(func, highlight=highlight)
        indent_len = len(exc_type.__name__) + len(fullname) + 6
        line = Fore.LIGHTBLACK_EX + indent + ("in " if i == i0 else "   ")
        line += (Fore.CYAN + fullname + Fore.LIGHTBLACK_EX if i == i0 else fullname) + "("
        line += _wrap(args_str + ") line %d" % frames[i].f_lineno, indent=indent_len)
        line += Style.RESET_ALL
        err(line)
    err()

    colorama.deinit()

def _get_method_full_name(func):
    """
    Return fully qualified function name.

    This method will attempt to find "full name" of the given function object. This full name is either of
    the form "<class name>.<method name>" if the function is a class method, or "<module name>.<func name>"
    if it's a regular function. Thus, this is an attempt to back-port func.__qualname__ to Python 2.

    :param func: a function object.

    :returns: string with the function's full name as explained above.
    """
    # Python 3.3 already has this information available...
    if hasattr(func, "__qualname__"): return func.__qualname__

    module = inspect.getmodule(func)
    if module is None:
        return "?.%s" % getattr(func, "__name__", "?")
    for cls_name in dir(module):
        cls = getattr(module, cls_name)
        if not inspect.isclass(cls): continue
        for method_name in dir(cls):
            cls_method = getattr(cls, method_name)
            if cls_method == func:
                return "%s.%s" % (cls_name, method_name)
    if hasattr(func, "__name__"):
        return "%s.%s" % (module.__name__, func.__name__)
    return "<unknown>"

def _find_function_from_code(frame, code):
    """
    Given a frame and a compiled function code, find the corresponding function object within the frame.

    This function addresses the following problem: when handling a stacktrace, we receive information about
    which piece of code was being executed in the form of a CodeType object. That objects contains function name,
    file name, line number, and the compiled bytecode. What it *doesn't* contain is the function object itself.

    So this utility function aims at locating this function object, and it does so by searching through objects
    in the preceding local frame (i.e. the frame where the function was called from). We expect that the function
    should usually exist there -- either by itself, or as a method on one of the objects.

    :param types.FrameType frame: local frame where the function ought to be found somewhere.
    :param types.CodeType code: the compiled code of the function to look for.

    :returns: the function object, or None if not found.
    """
    def find_code(iterable, depth=0):
        if depth > 3: return  # Avoid potential infinite loops, or generally objects that are too deep.
        for item in iterable:
            if item is None: continue
            found = None
            if hasattr(item, "__code__") and item.__code__ == code:
                found = item
            elif isinstance(item, type) or isinstance(item, ModuleType):  # class / module
                try:
                    found = find_code((getattr(item, n, None) for n in dir(item)), depth + 1)
                except Exception:
                    # Sometimes merely getting module's attributes may cause an exception. For example :mod:`six.moves`
                    # is such an offender...
                    continue
            elif isinstance(item, (list, tuple, set)):
                found = find_code(item, depth + 1)
            elif isinstance(item, dict):
                found = find_code(item.values(), depth + 1)
            if found: return found
    return find_code(frame.f_locals.values()) or find_code(frame.f_globals.values())

def _get_args_str(func, highlight=None):
    """
    Return function's declared arguments as a string.

    For example for this function it returns "func, highlight=None"; for the ``_wrap`` function it returns
    "text, wrap_at=120, indent=4". This should usually coincide with the function's declaration (the part
    which is inside the parentheses).
    """
    if not func: return ""
    s = str(inspect.signature(func))[1:-1]
    if highlight:
        s = re.sub(r"\b%s\b" % highlight, Style.BRIGHT + Fore.WHITE + highlight + Fore.LIGHTBLACK_EX + Style.NORMAL, s)
    return s

def _wrap(text, wrap_at=120, indent=4):
    """
    Return piece of text, wrapped around if needed.

    :param text: text that may be too long and then needs to be wrapped.
    :param wrap_at: the maximum line length.
    :param indent: number of spaces to prepend to all subsequent lines after the first.
    """
    out = ""
    curr_line_length = indent
    space_needed = False
    for word in text.split():
        if curr_line_length + len(word) > wrap_at:
            out += "\n" + " " * indent
            curr_line_length = indent
            space_needed = False
        if space_needed:
            out += " "
            curr_line_length += 1
        out += word
        curr_line_length += len(word)
        space_needed = True
    return out


#!/usr/bin/python

# "Avoid locker or centralized resource by hard-wiring the port mapping within range"
# "implied by max # of ports used per job, max # of executors per machine, and # of machines."
# "Map of source determines port. in/out using env variables"
print "\njenkins_h2o_port_allocate...."

import socket, os, subprocess
USED_HOSTNAMES = [
    'mr-0xb1',
    'mr-0xb4',

    'mr-0x2',
    'mr-0x3',
    'mr-0x4',
    'mr-0x5',
    'mr-0x6',
    'mr-0x7',
    'mr-0x8',
    'mr-0x9',
    'mr-0x10',

    'mr-0xd4',
    'mr-0xd5',
    'mr-0xd6',
    'mr-0xd7',
    'mr-0xd8',
    'mr-0xd9',
    'mr-0xd10',
    'Kevin-Ubuntu3',
]

# maximum number of ports a job uses 10 = 5 jvms * 2 ports per h2o jvm (current max known)
PORTS_PER_SLOT = 10
DEFAULT_BASE_PORT = 54340
EXECUTOR_NUM = 8

def jenkins_h2o_port_allocate():
    """
    input: jenkins environment variable EXECUTOR_NUMBER
    output: creates ./BASE_PORT.sh, that you should 'source ./PORT.sh'
        (can't see the env. variables directly from python?)

    which will create os environment variables H2O_PORT and H2O_PORT_OFFSET (legacy)

        internal state for this script that can be updated:
            USED_HOSTNAMES (list of machine names),
            PORTS_PER_SLOT (max per any job),
            DEFAULT_BASE_PORT
        If you modify any of the internal state, you may introduce contention between
        new jenkins jobs and running jenkins jobs. (might not!)
        You should stop/start all jobs (or ignore failures) if you modify internal state here.
        Hence, no parameters to avoid living dangerously!
    """
    if os.environ.has_key("EXECUTOR_NUMBER"):
        # this will fail if it's not an integer
        executor = int(os.environ["EXECUTOR_NUMBER"])
    else:
        executor = 1 # jenkins starts with 1

    print "jenkins EXECUTOR_NUMBER:", executor
    if executor<0 or executor>=EXECUTOR_NUM:
        raise Exception("executor: %s wrong? Expecting 1-8 jenkins executors on a machine (0-7 exp.)" % executor)

    h2oPort = DEFAULT_BASE_PORT
    h2oPortOffset = 0
    hostname = socket.gethostname()
    if hostname not in USED_HOSTNAMES:
        print "WARNING: this hostname: %s isn't in my list. You should add it?" % hostname
        print "Will use default base port"
    else:
        hostnameIndex = USED_HOSTNAMES.index(hostname)
        h2oPortOffset = PORTS_PER_SLOT * (executor + hostnameIndex)
        h2oPort += h2oPortOffset

    print "Possible h2o base_port range is %s to %s" % \
        (DEFAULT_BASE_PORT, DEFAULT_BASE_PORT + (PORTS_PER_SLOT * EXECUTOR_NUM * len(USED_HOSTNAMES)) - 2)
    print "Possible h2o ports used ranged is %s to %s" % \
        (DEFAULT_BASE_PORT, DEFAULT_BASE_PORT + (PORTS_PER_SLOT * EXECUTOR_NUM * len(USED_HOSTNAMES)) - 1)
    print "want to 'export H2O_PORT=%s'" % h2oPort
    print "want to 'export H2O_PORT_OFFSET=%s # legacy'" % h2oPortOffset

    f = open('H2O_BASE_PORT.sh','w')
    f.write('export H2O_PORT=%s\n' % h2oPort)

    f.write('export H2O_PORT_OFFSET=%s # legacy\n' % h2oPortOffset)
    f.close()
    print "\nNow please:\nsource ./H2O_BASE_PORT.sh"

if __name__ == "__main__":
    jenkins_h2o_port_allocate()

"""
This auto-magics the manual allocation I did when parallelized the current 8-way jenkins jobs,
2 per machine, on the jenkins mr-0xd4 that dispatches to mr-0xd5 thru mr-0xd9

The rationale for a global allocate requires understanding what machines a jenkins master/slave can be on,
and what machines they send h2o jars to.

at 0xdata:

A jenkins master is a member of a group of machines. Jenkins can send the python or other test to another slave machine, and then the test can dispatch h2o either locally, or to other machines in the group.

it can target h2o.jar's anywhere in that group, or dispatch a job to a slave in that group that might do the same.

We currently have two such groups, with one jenkins master in each group (mr-0xb4 and mr-0xd4)
(update: let's just say it's all one big group. Not worth optimizing for subgroup knowlege)
So using
    (hostname offset in the list of total hostnames)  * (EXECUTOR_NUMBER-1 * PORTS_PER_SLOT)

Will give a unique offset from the default 54340 base, for the job, regardless of which jenkins (master or slave) starts it in the group and where the h2o targest are (which is controlled by the config.json used in the job)

all cloud builds done in a job (one or more) use the same offset.

Dispatching tests from your laptop..will they collide with jenkins?
If the host machine is not in the list, like a laptop, then the offset is 0. (54340 will be used). I suppose jenkins could shift it's base_port to be at least 10 above 54340, so existing scripts that users have, that use 54340, won't be stepped on by jenkins. 54340 could be the jenkins base port.

EC2:
I suppose if the tests are used in ec2, we only do one h2o jar per machine, (or multijvm) so no conflict if 54340 is used. (or 54340). We typically want fast EC2 results, so don't overload target machines?. I suppose an EC2 machine list could be created in this script if we started overloading EC2 machines also

PORTS_PER_SLOT is 10 right now, since the most a job will do is 5 h2o jvms.

I guess to ease the transition, I could leave the H2O_PORT_OFFSET as the api to build_cloud(), and have another python script look at the current ho2 IP and EXECUTOR_NUMBER env variable from jenkins

Notes:
Right now, assuming the subnet octet range from a group is 160-180 or 181-190 works. 164 is an oddball case (out of the ten range for it's group)

I guess I could just put a list of IPs for the jenkins groups that exist, and find the group your in, and then get a "group index" from that list. That's robust and easily maintainable.

This algorithm keeps the total port range in use = (max # of executors per jenkins master or slave) * PORTS_PER_SLOT * (# of machines in a group)
Using 2 executors per machine is nice. 4 is about the max that works well with h2o. so 4 * 10 * 10 = 400 ports
that would be 54340 thru 54721

NICE POSSIBILITES: If we know that ubuntu or other services need to reserve ports that are in our range, we can put in mappings to other ports for those values, or shift the port range or whatever...i.e. we can adjust the algorithm in one place. If the 54340 base is not good, that's set in h2o.py..currently tests don't modify base_port (except for some cloud tests we don't run in jenkins, that do more than 5 jvms on a single machine)

I suppose the tool could output the exact port to use, rather than an offset to h2o.py's default.  Maybe initially will output both, so h2o.py can migrate
i.e. environment variables H2O_PORT_OFFSET and H2O_PORT (= 5321 + H2O_PORT_OFFSET)

UPDATE: To allow for dispatching h2o to any machine in any jenkins group, we can have just one group list that has all possible machines. Makes the used port range twice as big (800) but that's okay. It's like going to a 255.255.0.0 network!

Detail:
Jenkins has global environment variables

This one is useful
EXECUTOR_NUMBER    The unique number that identifies the current executor (among executors of the same machine) that's carrying out this build. This is the number you see in the "build executor status", except that the number starts from 0, not 1.

Now each slave machine can have multiple executors, in addition to the master.

So since in a grand scheme, we don't know who's creating h2o.jars on target machines, from which machine, (jenkins master or slave)...
it means we want a global h2o port allocation (assuming that scraping an h2o port from OS allocation is ugly)

I have cases on 164 jenkins that send the python job to jenkins slave 174, which dispatches h2o jars to 175-180, Or dispatch to YARN on hadoop clusters, but we don't care about ports there, we get told the ip/port by the h2odriver.

Since the pool of machines in a group is fixed, we have the EXECUTOR_NUMBER which is the parallelism per machine (jenkins master or slave), and we

Will give a unique offset to 54340
We can call it a "PORT_SLOT" and pass it as a environment variable like the current "export H2O_PORT_OFFSET=40"
that the build_cloud() uses to offset the default base_port. I suppose PORTS_PER_SLOT can be fixed in build_cloud() so it's the same for all jobs (so jobs don't step over each other.
"""