try:
        tag     = parser.expect(values.W_ContinuationPromptTag)
        handler = parser.expect(values.W_Object)
    except EndOfInput:
        pass

from pycket                    import values, values_parameter, values_string, values_struct, prims
from pycket.arity              import Arity
from pycket.cont               import continuation, loop_label, call_cont, Barrier, Cont, NilCont, Prompt
from pycket.error              import SchemeException
from pycket.argument_parser    import ArgParser, EndOfInput
from pycket.prims.expose       import default, expose, expose_val, procedure, make_procedure
from pycket.prims.plumber      import current_plumber_param
from pycket.util               import add_copy_method
from rpython.rlib              import jit
from rpython.rlib.objectmodel  import specialize

@add_copy_method(copy_method="_clone")
class DynamicWindValueCont(Cont):

    _immutable_fields_ = ['pre', 'post']

    def __init__(self, pre, post, env, cont):
        Cont.__init__(self, env, cont)
        self.pre  = pre
        self.post = post

    def has_unwind(self):
        return True

    def unwind(self, env, cont):
        return self.post.call([], env, cont)

    def has_rewind(self):
        return True

    def rewind(self, env, cont):
        return self.pre.call([], env, cont)

    def plug_reduce(self, _vals, env):
        cont = dynamic_wind_post_cont(_vals, env, self.prev)
        return self.post.call([], env, cont)

@continuation
def call_handler_cont(proc, args, env, cont, _vals):
    return proc.call(args, env, cont)

@continuation
def post_build_exception(env, cont, _vals):
    from pycket.interpreter import check_one_val
    barrier = values.w_true
    val = check_one_val(_vals)
    return raise_exception(val, barrier, env, cont)

def convert_runtime_exception(exn, env, cont):
    from pycket               import values_string
    from pycket.values        import W_ContinuationMarkSet
    from pycket.prims.control import raise_exception
    message = values_string.W_String.fromstr_utf8(exn.msg)
    marks   = W_ContinuationMarkSet(cont, values.w_default_continuation_prompt_tag)
    cont    = post_build_exception(env, cont)
    args    = exn.extra_args

    return exn.get_exn_type().constructor.call([message, marks] + args, env, cont)

def convert_os_error(exn, env, cont):
    from pycket               import values_string
    from pycket.values        import W_ContinuationMarkSet
    from pycket.prims.control import raise_exception
    message = values_string.W_String.fromstr_utf8(exn.strerror)
    marks   = W_ContinuationMarkSet(cont, values.w_default_continuation_prompt_tag)
    cont    = post_build_exception(env, cont)

    return prims.general.exn_fail.constructor.call([message, marks], env, cont)

@jit.unroll_safe
def scan_continuation(curr, prompt_tag, look_for=None, escape=False):
    """
    Segment a continuation based on a given continuation-prompt-tag.
    The head of the continuation, up to and including the desired continuation
    prompt is reversed (in place), and the tail is returned un-altered.

    The hint value |look_for| is used to determine when the continuation being
    installed is a prefix of the extant continuation.
    In this case, installing the continuation is much simpler, as the expensive
    merge operation needed to find common substructure is the two continuation is
    not needed.
    """
    handlers = False
    xs = []
    while isinstance(curr, Cont):
        if curr is look_for:
            return None, handlers
        handlers |= isinstance(curr, DynamicWindValueCont)
        xs.append(curr)
        if isinstance(curr, Prompt) and curr.tag is prompt_tag:
            break
        curr = curr.prev
        if not escape and not jit.isvirtual(curr):
            return _scan_continuation(curr, prompt_tag, look_for, xs, handlers)
    return xs, handlers

@jit.elidable
def _scan_continuation(curr, prompt_tag, look_for, xs, handlers):
    """
    Variant of scan_continuation which is elidable.
    scan_continuation switchs to using this function when all the virtual
    continuation frames are exhausted.
    """
    while isinstance(curr, Cont):
        if curr is look_for:
            return None, handlers
        handlers |= isinstance(curr, DynamicWindValueCont)
        xs.append(curr)
        if isinstance(curr, Prompt) and curr.tag is prompt_tag:
            break
        curr = curr.prev
    return xs, handlers

@jit.elidable
def find_merge_point(c1, c2):
    i = len(c1) - 1
    j = len(c2) - 1
    while i >= 0 and j >= 0 and c1[i] is c2[j]:
        i -= 1
        j -= 1
    unwind = []
    rewind = []

    r1 = c1[-1] if i == len(c1) - 1 else c1[i+1]
    r2 = c2[-1] if j == len(c2) - 1 else c2[j+1]

    while i >= 0:
        if isinstance(c1[i], DynamicWindValueCont):
            unwind.append(c1[i])
        i -= 1
    while j >= 0:
        if isinstance(c2[j], DynamicWindValueCont):
            rewind.append(c2[j])
        j -= 1
    return r1, r2, unwind, rewind

@jit.elidable
def find_handlers(cont, target):
    result = None
    while cont is not target:
        assert isinstance(cont, Cont)
        if isinstance(cont, DynamicWindValueCont):
            if not result:
                result = []
            result.append(cont)
        cont = cont.prev
    return result

@jit.unroll_safe
def install_continuation_fast_path(current_cont, args, has_handlers, env, cont):
    from pycket.interpreter import return_multi_vals, return_void

    if not has_handlers:
        args = values.Values.make(args)
        return return_multi_vals(args, env, cont)

    unwind = find_handlers(current_cont, cont)
    cont = return_args_cont(args, env, cont)
    unwind = [x for x in reversed(unwind)]
    cont = do_unwind_cont(unwind, env, cont)
    return return_void(env, cont)

@jit.unroll_safe
def install_continuation(cont, prompt_tag, args, env, current_cont, extend=False, escape=False):
    from pycket.interpreter import return_multi_vals, return_void

    # Find the common merge point for two continuations
    # The extend option controls whether or not we remove frames from the
    # existing continuation, or simply stack the new continuation on top.
    # This is what differentiates call-with-current-continuation from
    if extend:
        _   , rewind = find_continuation_prompt(prompt_tag, cont, direction='rewind')
        base, unwind = current_cont, None
        stop         = None
    else:
        head1, handlers = scan_continuation(
                current_cont, prompt_tag, look_for=cont, escape=escape)
        if head1 is None:
            return install_continuation_fast_path(current_cont, args, handlers, env, cont)
        head2, _ = scan_continuation(cont, prompt_tag)
        base, stop, unwind, rewind = find_merge_point(head1, head2)

    # Append the continuations at the appropriate prompt
    if base is not None:
        cont = cont.append(base, stop=stop, upto=prompt_tag)

    # Fast path if no unwinding is required (avoids continuation allocation)
    if not unwind and not rewind:
        args = values.Values.make(args)
        return return_multi_vals(args, env, cont)

    # NOTE: All the continuations pushed here expect no arguments
    # They simply wrap functions which we could call directly, but using
    # continuations to perform the desired sequencing is easier.
    cont = return_args_cont(args, env, cont)
    if rewind:
        cont = do_rewind_cont(rewind, env, cont)
    if unwind:
        unwind = [x for x in reversed(unwind)]
        cont = do_unwind_cont(unwind, env, cont)
    return return_void(env, cont)

@continuation
def return_args_cont(args, env, cont, _vals):
    from pycket.interpreter import return_multi_vals
    args = values.Values.make(args)
    return return_multi_vals(args, env, cont)

@continuation
def do_unwind_cont(frames, env, cont, _vals):
    return unwind_frames(frames, env, cont)

@continuation
def do_rewind_cont(frames, env, cont, _vals):
    return rewind_frames(frames, env, cont)

@jit.elidable
@specialize.arg(2)
def find_continuation_prompt(tag, cont, direction=None):
    wind_frames = [] if direction is not None else None
    while cont is not None:
        if isinstance(cont, Prompt) and cont.tag is tag:
            if direction is not None:
                return cont, wind_frames
            return cont
        if direction is not None and getattr(cont, "has_" + direction)():
            assert isinstance(cont, Cont)
            wind_frames.append(cont)
        cont = cont.get_previous_continuation()
    if direction is not None:
        return None, wind_frames
    return None

# NOTE do_unwind_frame and do_rewind_frame expect to be invoked with return_void
def do_unwind_frames(frames, index, env, final_cont):
    from pycket.interpreter import return_void
    if index >= len(frames):
        return return_void(env, final_cont)
    frame = frames[index]
    ctxt  = unwind_next_frame_cont(frames, index, final_cont, env, frame.prev)
    return frame.unwind(env, ctxt)

def do_rewind_frames(frames, index, env, final_cont):
    from pycket.interpreter import return_void
    if index < 0:
        return return_void(env, final_cont)
    frame = frames[index]
    ctxt = rewind_next_frame_cont(frames, index, final_cont, env, frame.prev)
    return frame.rewind(env, ctxt)

@continuation
def unwind_next_frame_cont(frames, index, final_cont, env, cont, _vals):
    return do_unwind_frames(frames, index + 1, env, final_cont)

@continuation
def rewind_next_frame_cont(frames, index, final_cont, env, cont, _vals):
    return do_rewind_frames(frames, index - 1, env, final_cont)

def unwind_frames(frames, env, final_cont):
    return do_unwind_frames(frames, 0, env, final_cont)

def rewind_frames(frames, env, final_cont):
    return do_rewind_frames(frames, len(frames) - 1, env, final_cont)

@expose("continuation-prompt-available?", [values.W_ContinuationPromptTag, default(values.W_Continuation, None)], simple=False)
def cont_prompt_avail(tag, continuation, env, cont):
    from pycket.interpreter import return_value
    if continuation is not None:
        cont = continuation.cont
    prompt = find_continuation_prompt(tag, cont)
    available = values.W_Bool.make(prompt is not None)
    return return_value(available, env, cont)

@continuation
def dynamic_wind_pre_cont(value, pre, post, env, cont, _vals):
    cont = DynamicWindValueCont(pre, post, env, cont)
    return value.call([], env, cont)

@continuation
def dynamic_wind_value_cont(post, env, cont, _vals):
    return post.call([], env, dynamic_wind_post_cont(_vals, env, cont))

@continuation
def dynamic_wind_post_cont(val, env, cont, _vals):
    from pycket.interpreter import return_multi_vals
    return return_multi_vals(val, env, cont)

@expose("dynamic-wind", [procedure, procedure, procedure], simple=False)
def dynamic_wind(pre, value, post, env, cont):
    return pre.call([], env, dynamic_wind_pre_cont(value, pre, post, env, cont))

@expose(["call/cc", "call-with-current-continuation"],
        [procedure, default(values.W_ContinuationPromptTag, None)],
        simple=False, extra_info=True)
def callcc(proc, prompt_tag, env, cont, extra_call_info):
    kont = [values.W_Continuation(cont, prompt_tag)]
    return proc.call_with_extra_info(kont, env, cont, extra_call_info)

@continuation
def call_with_escape_continuation_cont(env, cont, _vals):
    # Does not do anything currently. Solely to ensure call/ec does not invoke
    # its procedure in tail position
    from pycket.interpreter import return_multi_vals
    return return_multi_vals(_vals, env, cont)

@expose(["call/ec", "call-with-escape-continuation"],
        [procedure, default(values.W_ContinuationPromptTag, None)],
        simple=False, extra_info=True)
def call_with_escape_continuation(proc, prompt_tag, env, cont, extra_call_info):
    assert prompt_tag is None, "NYI"
    cont = call_with_escape_continuation_cont(env, cont)
    return proc.call_with_extra_info([values.W_EscapeContinuation(cont)],
                                     env, cont, extra_call_info)

@expose("call-with-composable-continuation",
        [procedure, default(values.W_ContinuationPromptTag, values.w_default_continuation_prompt_tag)],
        simple=False, extra_info=True)
def call_with_composable_continuation(proc, prompt_tag, env, cont, extra_call_info):
    kont = [values.W_ComposableContinuation(cont, prompt_tag)]
    return proc.call_with_extra_info(kont, env, cont, extra_call_info)

@expose("make-continuation-prompt-tag", [default(values.W_Symbol, None)])
def make_continuation_prompt_tag(sym):
    return values.W_ContinuationPromptTag(sym)

@expose("default-continuation-prompt-tag", [])
def default_continuation_prompt_tag():
    return values.w_default_continuation_prompt_tag

@expose("unsafe-root-continuation-prompt-tag", [])
def unsafe_root_configuraion_prompt_tag():
    return values.w_root_continuation_prompt_tag

def abort_current_continuation(args, env, cont):
    from pycket.interpreter import return_multi_vals
    if not args:
        raise SchemeException("abort-current-continuation: expected 1 or more args")
    tag, args = args[0], args[1:]
    if not isinstance(tag, values.W_ContinuationPromptTag):
        raise SchemeException("abort-current-continuation: expected prompt-tag for argument 0")

    prompt, frames = find_continuation_prompt(tag, cont, direction='unwind')
    if prompt is None:
        raise SchemeException("abort-current-continuation: no such prompt exists")
    handler = prompt.handler
    cont    = prompt.get_previous_continuation()
    assert cont is not None
    if handler is None:
        if not args:
            raise SchemeException("abort-current-continuation: expected thunk as argument 1")
        cont = Prompt(tag, None, env, cont)
        handler = args[0]
        args = []
    if frames:
        cont = call_handler_cont(handler, args, env, cont)
        return unwind_frames(frames, env, cont)
    return handler.call(args, env, cont)

expose("abort-current-continuation", simple=False)(abort_current_continuation)

@expose("force-exit", [values.W_Object])
def force_exit(v):
    return _force_exit(v)

def _force_exit(v):
    from pycket.error import ExitException

    if isinstance(v, values.W_Fixnum) and (0 <= v.value <= 255):
        raise ExitException(v)
    else:
        raise ExitException(values.W_Fixnum(0))

@continuation
def force_exit_cont(v, env, cont, _vals):
    return _force_exit(v)

@make_procedure("initial-exit-handler", [values.W_Object], arity=Arity.ZERO, simple=False)
def initial_exit_handler(v, env, cont):
    from pycket.prims.plumber import do_plumber_flush_all
    root_plumber = current_plumber_param.get(cont) ## ??
    return do_plumber_flush_all(root_plumber, env , force_exit_cont(v, env, cont))

exit_handler_param = values_parameter.W_Parameter(initial_exit_handler)
expose_val("exit-handler", exit_handler_param)

@make_procedure("error-value-string-handler", [values.W_Object, values.W_Fixnum])
def error_value_string_handler(v, repr_length):
    # FIXME: truncate it with the given representation length
    return values_string.W_String.make(v.tostring())

error_val_string_handler_param = values_parameter.W_Parameter(error_value_string_handler)
expose_val("error-value->string-handler", error_val_string_handler_param)

@continuation
def exit_cont(env, cont, _vals):
    from pycket.interpreter import return_value
    return return_value(values.w_void, env, cont)

@expose("exit", [default(values.W_Object, values.w_true)], simple=False)
def exit(v, env, cont):
    exit_handler = exit_handler_param.get(cont)
    return exit_handler.call([v], env, exit_cont(env, cont))

@make_procedure("default-error-escape-handler", [], simple=False)
def default_error_escape_handler(env, cont):
    from pycket.prims.general import do_void
    args = [values.w_default_continuation_prompt_tag, do_void.w_prim]
    return abort_current_continuation(args, env, cont)

error_escape_handler_param = values_parameter.W_Parameter(default_error_escape_handler)

expose_val("error-escape-handler", error_escape_handler_param)

def is_exn(v):
    from pycket.prims.general import exn
    return (isinstance(v, values_struct.W_Struct) and (exn.has_subtype(v.struct_type())))

def is_user_exn(v):
    from pycket.prims.general import exn_fail_user
    return (isinstance(v, values_struct.W_Struct) and (exn_fail_user.has_subtype(v.struct_type())))

def get_exn_message(exn, env, cont):
    offset = exn.struct_type().get_offset(exn.struct_type())
    original_field_num = 0 # this is for the message field in exceptions
    message_field_index = values.W_Fixnum(original_field_num-offset)

    return exn.struct_type().accessor.call([exn, message_field_index], env, display_escape_cont(exn, env, cont))

def display_stack_trace(port, cont):
    from pycket.prims.continuation_marks import cms_context
    context = cms_context.w_prim.call_interpret([values.W_ContinuationMarkSet(cont, values.w_default_continuation_prompt_tag)])
    if isinstance(context, values.W_Cons):
        port.write("Error Trace:\n")
        total_frames_to_show = 10
        count = 0
        while isinstance(context, values.W_Cons):
            if count >= total_frames_to_show:
                break
            port.write("-- %s\n" % context.car().tostring()[:1000])
            context = context.cdr()
            count += 1

@make_procedure("default-error-display-handler", [values_string.W_String, values.W_Object], simple=False)
def default_error_display_handler(msg, exn_object, env, cont):
    from pycket.prims.input_output import current_error_param, return_void
    port = current_error_param.get(cont)

    assert isinstance(port, values.W_OutputPort)
    if is_exn(exn_object):
        port.write("%s : %s\n" % (exn_object.struct_type().name.tostring(), msg.tostring()))
    else:
        port.write("exception : %s\n" % (msg.tostring()))

    if not is_user_exn(exn_object):
        display_stack_trace(port, cont)
    return return_void(env, cont)

error_display_handler_param = values_parameter.W_Parameter(default_error_display_handler)

expose_val("error-display-handler", error_display_handler_param)

@continuation
def display_escape_cont(exn, env, cont, _vals):
    from pycket.interpreter import check_one_val
    message = check_one_val(_vals)

    display_handler = error_display_handler_param.get(cont) # parameterize this to default first
    escape_handler = error_escape_handler_param.get(cont) # this one too

    # display, then escape
    return display_handler.call([message, exn], env, call_handler_cont(escape_handler, [], env, cont))

@make_procedure("default-uncaught-exception-handler", [values.W_Object], simple=False)
def default_uncaught_exception_handler(exn, env, cont):
    from pycket.env import w_global_config
    w_global_config.set_error_exit(exn)
    # racket/src/cs/rumble/error.ss

    #FIXME : handle Breaks
    if is_exn(exn):
        return get_exn_message(exn, env, display_escape_cont(exn, env, cont))
    else:
        from pycket.interpreter import return_value
        return return_value(values_string.W_String.make(exn.tostring()), env, display_escape_cont(exn, env, cont))

uncaught_exception_handler_param = values_parameter.W_Parameter(default_uncaught_exception_handler)

expose_val("uncaught-exception-handler", uncaught_exception_handler_param)

@make_procedure("default-continuation-prompt-handler", [procedure], simple=False)
def default_continuation_prompt_handler(proc, env, cont):
    return proc.call([], env, cont)

@expose("call-with-continuation-prompt", simple=False, arity=Arity.geq(1), extra_info=True)
def call_with_continuation_prompt(args, env, cont, extra_call_info):
    if not args:
        raise SchemeException("call-with-continuation-prompt: not given enough values")
    parser  = ArgParser("call-with-continuation-prompt", args)
    tag     = values.w_default_continuation_prompt_tag
    handler = values.w_false
    fun     = parser.expect(values.W_Object)

    args = parser.expect_many(values.W_Object)
    if not fun.iscallable():
        raise SchemeException("call-with-continuation-prompt: not given callable function")
    if handler is not values.w_false and not handler.iscallable():
        raise SchemeException("call-with-continuation-prompt: not given callable handler")
    if handler is values.w_false:
        handler = None
    cont = Prompt(tag, handler, env, cont)
    return fun.call_with_extra_info(args, env, cont, extra_call_info)

@expose("call-with-continuation-barrier", [procedure], simple=False, extra_info=True)
def call_with_continuation_barrier(proc, env, cont, calling_app):
    # TODO: Implementation
    cont = Barrier(env, cont)
    return proc.call_with_extra_info([], env, cont, calling_app)

def raise_exception(v, barrier, env, cont):
    # TODO: Handle case where barrier is not #t
    assert barrier is values.w_true

    handler = None
    k = cont
    while k is not None:
        handler, next_cont = k.find_cm(values.exn_handler_key)
        if handler is not None:
            break
        k = next_cont if next_cont else k.get_previous_continuation()
    else:
        raise SchemeException("uncaught exception:\n %s" % v.tostring())

    if not handler.iscallable():
        raise SchemeException("provided handler is not callable")

    return handler.call([v], env, cont)

expose("raise", [values.W_Object, default(values.W_Object, values.w_true)], simple=False)(raise_exception)

@expose("raise-argument-error", arity=Arity.geq(3))
def raise_arg_err(args):
    nargs = len(args)
    if nargs < 3:
        raise SchemeException("raise-argument-error: expected at least 3 arguments")
    name = args[0]
    if not isinstance(name, values.W_Symbol):
        raise SchemeException("raise-argument-error: expected symbol as the first argument")
    expected = args[1]
    if not isinstance(expected, values_string.W_String):
        raise SchemeException("raise-argument-error: expected string as the second argument")
    if nargs == 3:
        v = args[2]
        raise SchemeException("%s: expected %s but got %s" % (
            name.utf8value, expected.as_str_utf8(), v.tostring()))
    else:
        bad_v = args[2]
        if not isinstance(bad_v, values.W_Fixnum):
            raise SchemeException("raise-argument-error: expected number as the third argument")
        value = bad_v.value
        if value >= nargs - 3:
            raise SchemeException("raise-argument-error: out of bounds number as the third argument")
        v = args[value + 3]
        # FIXME: actually print the other arguments
        raise SchemeException("%s: contract violation\n  expected: %s\n  given: %s\n argument position: %s"%(
             name.utf8value, expected.as_str_utf8(), v.tostring(), value + 1))

@expose("raise-arguments-error", arity=Arity.geq(2))
def raise_args_err(args):
    name = args[0]
    assert isinstance(name, values.W_Symbol)
    message = args[1]
    assert isinstance(message, values_string.W_String)
    from rpython.rlib.rstring import StringBuilder
    error_msg = StringBuilder()
    error_msg.append(name.utf8value)
    error_msg.append(": ")
    error_msg.append(message.as_str_utf8())
    error_msg.append("\n")
    i = 2
    while i + 1 < len(args):
        field = args[i]
        assert isinstance(field, values_string.W_String)
        v = args[i+1]
        assert isinstance(v, values.W_Object)
        error_msg.append("%s: %s\n" % (field.as_str_utf8(), v.tostring()))
        i += 2
    raise SchemeException(error_msg.build())

@expose("raise-mismatch-error", arity=Arity.geq(3))
def raise_mismatch_err(args):
    name = args[0]
    assert isinstance(name, values.W_Symbol)
    message = args[1]
    assert isinstance(message, values_string.W_String)
    v = args[2]
    assert isinstance(v, values.W_Object)
    from rpython.rlib.rstring import StringBuilder
    error_msg = StringBuilder()
    error_msg.append(name.utf8value)
    error_msg.append(": ")
    error_msg.append(message.as_str_utf8())
    error_msg.append(v.tostring())
    i = 3
    while i + 1 < len(args):
        message = args[i]
        assert isinstance(message, values_string.W_String)
        error_msg.append(message.as_str_utf8())
        v = args[i+1]
        assert isinstance(v, values.W_Object)
        error_msg.append(v.tostring())
        i += 2
    raise SchemeException(error_msg.build())

@expose("raise-type-error", [values.W_Symbol, values_string.W_String, values.W_Object])
def raise_type_error(name, expected, v):
    raise SchemeException("%s: expected %s in %s" % (name.tostring(), expected.tostring(), v.tostring()))


from gwssh import *
from twisted.python import components
import cPickle as Pickle
from twisted.conch.ssh import session

class PickleProtocol(SBProtocol):
    def write_prompt(self):

class PickleTunnelSession(SshTunnelSession):
	WRAPPROTOCOL = PickleProtocol

class PickleSession(ShellSession):
	ISESSION = PickleTunnelSession

class PickleUser(ShellUser):
    SESSIONCLASS = PickleSession
    def sendReply(self, reply, inreply, **kw):
        if 'msg' in kw:
            s = Pickle.dumps(kw['msg'])
            self.con.write(s)

class PickleAuthServer(SBSSHUserAuthServer):
	USERCLASS = PickleUser

class PickleGW(ShellGW):
    PORTALISE = True
    USERCLASS = PickleUser
    DEFAULT_PORT = 9023
    PROTOCOL = 'pickle'
    USESSL = False
    CONFIG_EXTRA = []

    class FactoryClass(ShellGW.FactoryClass):
        services = {
            'ssh-userauth': PickleAuthServer,
            'ssh-connection': connection.SSHConnection
        }


self.sess.run(self.optimizer, feed_dict={self.critic_state_in: state_t_batch, self.critic_action_in:action_batch, self.q_value_in: y_i_batch})

import numpy as np
import tensorflow as tf
import math

TAU = 0.001
LEARNING_RATE= 0.001
BATCH_SIZE = 64
class CriticNet:
    """ Critic Q value model of the DDPG algorithm """
    def __init__(self,num_states,num_actions):

        self.g=tf.Graph()
        with self.g.as_default():
            self.sess = tf.InteractiveSession()

            #critic_q_model parameters:
            self.W1_c, self.B1_c, self.W2_c, self.W2_action_c, self.B2_c, self.W3_c, self.B3_c,\
            self.critic_q_model, self.critic_state_in, self.critic_action_in = self.create_critic_net(num_states, num_actions)

            #create target_q_model:
            self.t_W1_c, self.t_B1_c, self.t_W2_c, self.t_W2_action_c, self.t_B2_c, self.t_W3_c, self.t_B3_c,\
            self.t_critic_q_model, self.t_critic_state_in, self.t_critic_action_in = self.create_critic_net(num_states, num_actions)

            self.q_value_in=tf.placeholder("float",[None,1]) #supervisor
            self.l2_regularizer_loss = 0.0001*tf.reduce_sum(tf.pow(self.W2_c,2))+ 0.0001*tf.reduce_sum(tf.pow(self.B2_c,2))
            self.cost=tf.pow(self.critic_q_model-self.q_value_in,2)/BATCH_SIZE + self.l2_regularizer_loss#/tf.to_float(tf.shape(self.q_value_in)[0])
            self.optimizer = tf.train.AdamOptimizer(learning_rate=LEARNING_RATE).minimize(self.cost)

            #action gradient to be used in actor network:
            #from simple actor net:
            self.act_grad_v = tf.gradients(self.critic_q_model, self.critic_action_in)
            self.action_gradients = [self.act_grad_v[0]/tf.to_float(tf.shape(self.act_grad_v[0])[0])] #this is just divided by batch size
            #from simple actor net:
            self.check_fl = self.action_gradients

            #initialize all tensor variable parameters:
            self.sess.run(tf.initialize_all_variables())

            #To make sure critic and target have same parmameters copy the parameters:
            # copy target parameters
            self.sess.run([
				self.t_W1_c.assign(self.W1_c),
				self.t_B1_c.assign(self.B1_c),
				self.t_W2_c.assign(self.W2_c),
				self.t_W2_action_c.assign(self.W2_action_c),
				self.t_B2_c.assign(self.B2_c),
				self.t_W3_c.assign(self.W3_c),
				self.t_B3_c.assign(self.B3_c)
			])

            self.update_target_critic_op = [
                self.t_W1_c.assign(TAU*self.W1_c+(1-TAU)*self.t_W1_c),
                self.t_B1_c.assign(TAU*self.B1_c+(1-TAU)*self.t_B1_c),
                self.t_W2_c.assign(TAU*self.W2_c+(1-TAU)*self.t_W2_c),
                self.t_W2_action_c.assign(TAU*self.W2_action_c+(1-TAU)*self.t_W2_action_c),
                self.t_B2_c.assign(TAU*self.B2_c+(1-TAU)*self.t_B2_c),
                self.t_W3_c.assign(TAU*self.W3_c+(1-TAU)*self.t_W3_c),
                self.t_B3_c.assign(TAU*self.B3_c+(1-TAU)*self.t_B3_c)
            ]

    def create_critic_net(self, num_states=4, num_actions=1):
        N_HIDDEN_1 = 400
        N_HIDDEN_2 = 300
        critic_state_in = tf.placeholder("float",[None,num_states])
        critic_action_in = tf.placeholder("float",[None,num_actions])

        W1_c = tf.Variable(tf.random_uniform([num_states,N_HIDDEN_1],-1/math.sqrt(num_states),1/math.sqrt(num_states)))
        B1_c = tf.Variable(tf.random_uniform([N_HIDDEN_1],-1/math.sqrt(num_states),1/math.sqrt(num_states)))
        W2_c = tf.Variable(tf.random_uniform([N_HIDDEN_1,N_HIDDEN_2],-1/math.sqrt(N_HIDDEN_1+num_actions),1/math.sqrt(N_HIDDEN_1+num_actions)))
        W2_action_c = tf.Variable(tf.random_uniform([num_actions,N_HIDDEN_2],-1/math.sqrt(N_HIDDEN_1+num_actions),1/math.sqrt(N_HIDDEN_1+num_actions)))
        B2_c= tf.Variable(tf.random_uniform([N_HIDDEN_2],-1/math.sqrt(N_HIDDEN_1+num_actions),1/math.sqrt(N_HIDDEN_1+num_actions)))
        W3_c= tf.Variable(tf.random_uniform([N_HIDDEN_2,1],-0.003,0.003))
        B3_c= tf.Variable(tf.random_uniform([1],-0.003,0.003))

        H1_c=tf.nn.softplus(tf.matmul(critic_state_in,W1_c)+B1_c)
        H2_c=tf.nn.tanh(tf.matmul(H1_c,W2_c)+tf.matmul(critic_action_in,W2_action_c)+B2_c)

        critic_q_model=tf.matmul(H2_c,W3_c)+B3_c

        return W1_c, B1_c, W2_c, W2_action_c, B2_c, W3_c, B3_c, critic_q_model, critic_state_in, critic_action_in

    def train_critic(self, state_t_batch, action_batch, y_i_batch ):

    def evaluate_target_critic(self,state_t_1,action_t_1):
        return self.sess.run(self.t_critic_q_model, feed_dict={self.t_critic_state_in: state_t_1, self.t_critic_action_in: action_t_1})

    def compute_delQ_a(self,state_t,action_t):
#        print 'check grad number'
#        print len(ch)
#        raw_input("Press Enter to continue...")
        return self.sess.run(self.action_gradients, feed_dict={self.critic_state_in: state_t,self.critic_action_in: action_t})

    def update_target_critic(self):
        self.sess.run(self.update_target_critic_op)


user_list = []
        item_list = []
        test_ranking_filename = "%s/%s.test.negative.%s" % (self.data_dir, self.arg.dataset_name, self.neg_num)
        with open(test_ranking_filename) as f:
            for line in f:
                token = line.split('\t')
                user_list.append([int(token[0])])
                item_list.append(list(map(lambda _: int(_), token[1:])))
        self.ranking_user = np.array(user_list)
        self.ranking_item = np.array(item_list)

# author: yf lu

import os
from collections import defaultdict
import numpy as np
from torch.utils.data import Dataset

class DataUtil(Dataset):
    def __init__(self, arg, root_dir, data_type):
        self.arg = arg
        self.root_dir = root_dir
        self.data_type = data_type
        self.u_items = defaultdict(set)
        self.i_users = defaultdict(set)
        self.data_dir = os.path.join(self.root_dir + '/data/', self.arg.dataset_name)  # ../data/yelp/
        self.data_size = 0
        self.filename = "%s/%s.%s.rating.712" % (self.data_dir, self.arg.dataset_name, self.data_type)
        self.data = defaultdict(int)  # more efficient
        self.neg_num = self.arg.eval_num
        if data_type == 'eval':
            self.get_data_4ranking()
        elif data_type == 'att_analysis':
            self.filename = "%s/%s.att.analysis" % (self.data_dir, self.arg.dataset_name)
            self.load_data()
        else:
            self.load_data()

    def load_data(self):
        with open(self.filename) as f:
            for line in f:
                token = line.split('\t')  # user_id \t item_id \t label \t act_list
                user = int(token[0])
                item = int(token[1])
                label = int(token[2])
                act = token[3].strip()
                self.data[(user, item, label, act)] = 1
                if label == 1:  # positive
                    self.u_items[user].add(item)
                    self.i_users[item].add(user)
                self.data_size += 1

        self.idx2user = np.zeros((self.data_size,), dtype=np.int32)
        self.idx2item = np.zeros((self.data_size,), dtype=np.int32)
        self.idx2label = np.zeros((self.data_size,), dtype=np.int32)
        self.idx2act = np.zeros((self.data_size,), dtype=object)
        for idx, (u, i, l, a) in enumerate(self.data):
            self.idx2user[idx] = u
            self.idx2item[idx] = i
            self.idx2label[idx] = l
            self.idx2act[idx] = list(map(lambda x: int(x), a.split(' ')))

    def __len__(self):
        return self.data_size

    def __getitem__(self, idx):
        """
        for wechat dataset
        :param idx:
        :return:
        """
        user = self.idx2user[idx]
        item = self.idx2item[idx]
        label = self.idx2label[idx]
        act = self.idx2act[idx]

        sample = {
            'user': user,
            'item': item,
            'label': label,
            'act': act,
        }

        return sample

    def get_data_4ranking(self):

    def load_social_data(self):
        social_graph_filename = "%s/%s.social.graph" % (self.data_dir, self.arg.dataset_name)
        social_relation = defaultdict(list)
        with open(social_graph_filename) as f:
            for line in f:
                token = line.split('\t')  # user_id \t user_id
                social_relation[int(token[0])].append(int(token[1]))
                social_relation[int(token[1])].append(int(token[0]))
        return social_relation

    def load_biz_data(self):
        user_biz_file = "%s/%s.user.biz" % (self.data_dir, self.arg.dataset_name)
        item_biz_file = "%s/%s.item.biz" % (self.data_dir, self.arg.dataset_name)
        user_bizs = defaultdict(list)
        item_bizs = defaultdict(list)
        with open(user_biz_file) as f:
            for line in f:
                token = line.split('\t')  # user_id \t biz_id
                user_bizs[int(token[0])].append(int(token[1]))
        with open(item_biz_file) as f:
            for line in f:
                token = line.split('\t')  # item_id \t biz_id
                item_bizs[int(token[0])].append(int(token[1]))
        return user_bizs, item_bizs


        user = self.idx2user[idx]
        item = self.idx2item[idx]
        label = self.idx2label[idx]
        act = self.idx2act[idx]

        sample = {
            'user': user,
            'item': item,
            'label': label,
            'act': act,
        }

        return sample

# author: yf lu

import os
from collections import defaultdict
import numpy as np
from torch.utils.data import Dataset

class DataUtil(Dataset):
    def __init__(self, arg, root_dir, data_type):
        self.arg = arg
        self.root_dir = root_dir
        self.data_type = data_type
        self.u_items = defaultdict(set)
        self.i_users = defaultdict(set)
        self.data_dir = os.path.join(self.root_dir + '/data/', self.arg.dataset_name)  # ../data/yelp/
        self.data_size = 0
        self.filename = "%s/%s.%s.rating.712" % (self.data_dir, self.arg.dataset_name, self.data_type)
        self.data = defaultdict(int)  # more efficient
        self.neg_num = self.arg.eval_num
        if data_type == 'eval':
            self.get_data_4ranking()
        elif data_type == 'att_analysis':
            self.filename = "%s/%s.att.analysis" % (self.data_dir, self.arg.dataset_name)
            self.load_data()
        else:
            self.load_data()

    def load_data(self):
        with open(self.filename) as f:
            for line in f:
                token = line.split('\t')  # user_id \t item_id \t label \t act_list
                user = int(token[0])
                item = int(token[1])
                label = int(token[2])
                act = token[3].strip()
                self.data[(user, item, label, act)] = 1
                if label == 1:  # positive
                    self.u_items[user].add(item)
                    self.i_users[item].add(user)
                self.data_size += 1

        self.idx2user = np.zeros((self.data_size,), dtype=np.int32)
        self.idx2item = np.zeros((self.data_size,), dtype=np.int32)
        self.idx2label = np.zeros((self.data_size,), dtype=np.int32)
        self.idx2act = np.zeros((self.data_size,), dtype=object)
        for idx, (u, i, l, a) in enumerate(self.data):
            self.idx2user[idx] = u
            self.idx2item[idx] = i
            self.idx2label[idx] = l
            self.idx2act[idx] = list(map(lambda x: int(x), a.split(' ')))

    def __len__(self):
        return self.data_size

    def __getitem__(self, idx):
        """
        for wechat dataset
        :param idx:
        :return:
        """

    def get_data_4ranking(self):
        user_list = []
        item_list = []
        test_ranking_filename = "%s/%s.test.negative.%s" % (self.data_dir, self.arg.dataset_name, self.neg_num)
        with open(test_ranking_filename) as f:
            for line in f:
                token = line.split('\t')
                user_list.append([int(token[0])])
                item_list.append(list(map(lambda _: int(_), token[1:])))
        self.ranking_user = np.array(user_list)
        self.ranking_item = np.array(item_list)

    def load_social_data(self):
        social_graph_filename = "%s/%s.social.graph" % (self.data_dir, self.arg.dataset_name)
        social_relation = defaultdict(list)
        with open(social_graph_filename) as f:
            for line in f:
                token = line.split('\t')  # user_id \t user_id
                social_relation[int(token[0])].append(int(token[1]))
                social_relation[int(token[1])].append(int(token[0]))
        return social_relation

    def load_biz_data(self):
        user_biz_file = "%s/%s.user.biz" % (self.data_dir, self.arg.dataset_name)
        item_biz_file = "%s/%s.item.biz" % (self.data_dir, self.arg.dataset_name)
        user_bizs = defaultdict(list)
        item_bizs = defaultdict(list)
        with open(user_biz_file) as f:
            for line in f:
                token = line.split('\t')  # user_id \t biz_id
                user_bizs[int(token[0])].append(int(token[1]))
        with open(item_biz_file) as f:
            for line in f:
                token = line.split('\t')  # item_id \t biz_id
                item_bizs[int(token[0])].append(int(token[1]))
        return user_bizs, item_bizs


from .common import *
from .zeromq import *

def listen(host, port):
  socket = server_socket((host, port))
  server = Server(socket)
  return server

def listen_on_random_port(host):
  socket = server_socket((host, -1))
  server = Server(socket)
  return server

def connect(host, port):


array = []
        for x in data:
                if x != '0' or 0:
                        array.append(x)
        return array

def average(amount):
        answer = []
        for x in range(amount):
                array = raw_input().split()
                array = cleanData(array)
                total, average, y = 0,0,0
                while y < len(array):
                        total += int(array[y])
                        y+=1
                        average = "%.02f" % ((float(total)) / float(len(array)))
                        average = int(round(float(average)))
                answer.append(str(average))
        print(' '.join(answer))
average(input())


while y < len(array):
                        total += int(array[y])
                        y+=1
                        average = "%.02f" % ((float(total)) / float(len(array)))
                        average = int(round(float(average)))

def cleanData(data):
        array = []
        for x in data:
                if x != '0' or 0:
                        array.append(x)
        return array

def average(amount):
        answer = []
        for x in range(amount):
                array = raw_input().split()
                array = cleanData(array)
                total, average, y = 0,0,0

                answer.append(str(average))
        print(' '.join(answer))
average(input())


from ..state.prekeystore import PreKeyStore
from ..state.prekeyrecord import PreKeyRecord
from ..invalidkeyidexception import InvalidKeyIdException

class InMemoryPreKeyStore(PreKeyStore):
    def __init__(self):

    def loadPreKey(self, preKeyId):
        if preKeyId not in self.store:
            raise InvalidKeyIdException("No such prekeyRecord!")

        return PreKeyRecord(serialized=self.store[preKeyId])

    def storePreKey(self, preKeyId, preKeyRecord):
        self.store[preKeyId] = preKeyRecord.serialize()

    def containsPreKey(self, preKeyId):
        return preKeyId in self.store

    def removePreKey(self, preKeyId):
        if preKeyId in self.store:
            del self.store[preKeyId]


self._max_services = max_services
        self._max_results = max_results
        self._token = token

import logging
import suds, suds.sudsobject
from suds.sax.element import Element

from django.utils.translation import ugettext_lazy
from django.utils.translation import ugettext as _
from django.conf import settings

from molly.apps.places.models import Entity
from molly.apps.places.providers import BaseMapsProvider

logger = logging.getLogger(__name__)

class LiveDepartureBoardPlacesProvider(BaseMapsProvider):
    _WSDL_URL = "http://example.com"
    _ATTRIBUTION = { 'title': _("Powered by National Rail Enquiries"),
                  'url': "http://www.nationalrail.co.uk",
                  'picture': settings.STATIC_URL + "places/images/powered-by-nre.png" }

    def __init__(self, token, max_services=15, max_results=1):

    def delayed(self, eta, sta, etd, std):
        """
        Try and figure out if a service is delayed based on free text values
        for estimated/scheduled times of arrival/depature
        """
        if eta in ('Delayed', 'Cancelled') or etd in ('Delayed', 'Cancelled'):
            # Easy case
            return True
        else:
            # More complex case, have to parse time stamps
            try:
                # Compare scheduled and expected arrival times
                schedh, schedm = sta.split(':')
                exph, expm = eta.rstrip('*').split(':')
            except ValueError:
                pass
            else:
                # Minutes since midnight, as we can't compare time objects
                sched_msm = int(schedh) * 60 + int(schedm)
                exp_msm = int(exph) * 60 + int(expm)

                if exp_msm < sched_msm:
                    # Deal with wraparound at midnight
                    sched_msm += 1440

                if exp_msm - sched_msm >= 5:
                    return True

            try:
                # Compare scheduled and expected departure times
                schedh, schedm = std.split(':')
                exph, expm = etd.rstrip('*').split(':')
            except ValueError:
                pass
            else:
                # Minutes since midnight, as we can't compare time objects
                sched_msm = int(schedh) * 60 + int(schedm)
                exp_msm = int(exph) * 60 + int(expm)

                if exp_msm < sched_msm:
                    # Deal with wraparound at midnight
                    sched_msm += 1440

                if exp_msm - sched_msm >= 5:
                    return True

            return False

    def augment_metadata(self, entities, board='departures', **kwargs):
        station_entities = []
        for entity in entities:
            if not entity.identifiers.get('crs'):
                continue
            station_entities.append(entity)

        station_entities = station_entities[:self._max_results]
        if not station_entities:
            return

        try:
            ldb = suds.client.Client(self._WSDL_URL, soapheaders=Element('AccessToken').insert(Element('TokenValue').setText(self._token)))
        except Exception, e:
            logger.warning("Could not instantiate suds client for live departure board.", exc_info=True, extra={'wsdl_url': self._WSDL_URL})
            self._add_error(station_entities)
            return

        for entity in station_entities:
            try:
                if board == 'arrivals':
                    db = ldb.service.GetArrivalBoard(self._max_services, entity.identifiers['crs'])
                else:
                    db = ldb.service.GetDepartureBoard(self._max_services, entity.identifiers['crs'])
                db = self.transform_suds(db)
                entity.metadata['ldb'] = db
                entity.metadata['service_details'] = lambda s: self.service_details(s, entity)
                entity.metadata['ldb_service'] = lambda s: self.transform_suds(ldb.service.GetServiceDetails(s))
                entity.metadata['service_type'] = 'ldb'

                # Show bus services too
                if board == 'arrivals':
                    db = self.transform_suds(
                        ldb.service.GetDepartureBoard(self._max_services,
                                                     entity.identifiers['crs']))

                if 'trainServices' in db:
                    for service in db['trainServices']['service']:
                        service['problems'] = self.delayed(
                            service.get('eta', ''), service.get('sta', ''),
                            service.get('etd', ''), service.get('std', ''))

                if 'busServices' in db:
                    for service in db['busServices']['service']:
                        entity.metadata['real_time_information'] = {
                            'services':
                                [{
                                    'service': 'BUS',
                                    'destination': service['destination']['location'][0]['locationName'],
                                    'next': service['std'],
                                    'following': [],
                                }]
                        }

            except Exception, e:
                logger.warning("Could not retrieve departure board for station: %r", entity.identifiers.get('crs'),
                               exc_info=True)
                self._add_error((entity,))
            entity.metadata['meta_refresh'] = 60

    def service_details(self, service, entity):
        try:
            service = entity.metadata['ldb_service'](service)
        except suds.WebFault as f:
            if f.fault['faultstring'] == 'Unexpected server error: Invalid length for a Base-64 char array.':
                raise Http404
            else:
                return({'error': f.fault['faultstring']})
        if service is None:
            return None

        # Trains can split and join, which makes figuring out the list of
        # calling points a bit difficult. The LiveDepartureBoards documentation
        # details how these should be handled. First, we build a list of all
        # the calling points on the "through" train.
        calling_points = service['previousCallingPoints']['callingPointList'][0]['callingPoint'] if len(service['previousCallingPoints']) else []

        # Then attach joining services to our thorough route in the correct
        # point, but only if there is a list of previous calling points
        if len(service['previousCallingPoints']):
            for i, points in enumerate(service['previousCallingPoints']['callingPointList'][1:]):
                # If the other half of the split ends up here, it's not a split, but
                # a change between rail replacement bus of train
                for j, point in enumerate(calling_points):
                    if points['callingPoint'][-1]['crs'] == point['crs']:
                        if j > 0:
                            point['joining'] = points['callingPoint']
                        else:
                            point['service_change'] = service['previousCallingPoints']['callingPointList'][i]['_serviceType']
                            calling_points = points['callingPoint'] + calling_points
                            break

        # Add our current station
        calling_points += [{
            'locationName': service['locationName'],
            'crs': service['crs'],
            'st': service['std'] if 'std' in service else service['sta'],
            'et': service['etd'] if 'etd' in service else service['eta'] if 'eta' in service else '',
            'at': service['atd'] if 'atd' in service else '',
        }]

        if len(service['subsequentCallingPoints']):
            if service['serviceType'] != service['subsequentCallingPoints']['callingPointList'][0]['_serviceType']:
                calling_points[-1]['service_change'] = ugettext_lazy(service['subsequentCallingPoints']['callingPointList'][0]['_serviceType'])

        if len(service['previousCallingPoints']):
            if service['serviceType'] != service['previousCallingPoints']['callingPointList'][0]['_serviceType']:
                calling_points[-1]['service_change'] = ugettext_lazy(service['serviceType'])

        # Now add services going forward
        if len(service['subsequentCallingPoints']):
            calling_points += service['subsequentCallingPoints']['callingPointList'][0]['callingPoint']

        # And do the same with splitting services
        if len(service['subsequentCallingPoints']):
            for i, points in enumerate(service['subsequentCallingPoints']['callingPointList'][1:]):
                # Now we have to handle changes between trains and rail replacement services
                if service['subsequentCallingPoints']['callingPointList'][i]['callingPoint'][-1]['crs'] == points['callingPoint'][0]['crs']:
                    calling_points += points['callingPoint']
                    points['callingPoint'][0]['service_change'] = ugettext_lazy(points['_serviceType'])
                else:
                    for point in calling_points:
                        if points['callingPoint'][0]['crs'] == point['crs']:
                            point['splitting'] = {'destination': points['callingPoint'][-1]['locationName'], 'list': points['callingPoint']}

        sources = [calling_points[0]['locationName']]
        for point in calling_points:
            if 'joining' in point:
                sources.append(point['joining'][0]['locationName'])

        destinations = [calling_points[-1]['locationName']]
        for point in calling_points:
            if 'splitting' in point:
                destinations.append(point['splitting']['destination'])

        stop_entities = []

        # Now get a list of the entities for the stations (if they exist)
        # to plot on a map, and figure out if this stop is delayed or not
        for point in calling_points:

            point['problems'] = self.delayed(
                point.get('et', ''), point.get('st', ''), '', '')

            if 'joining' in point:
                for jpoint in point['joining']:

                    jpoint['problems'] = self.delayed(
                        jpoint.get('et', ''), jpoint.get('st', ''), '', '')

                    point_entity = Entity.objects.filter(_identifiers__scheme='crs', _identifiers__value=str(jpoint['crs']))
                    if len(point_entity):
                        point_entity = point_entity[0]
                        jpoint['entity'] = point_entity
                        stop_entities.append(point_entity)
                        jpoint['stop_num'] = len(stop_entities)

            point_entity = Entity.objects.filter(_identifiers__scheme='crs', _identifiers__value=str(point['crs']))
            if len(point_entity):
                point_entity = point_entity[0]
                point['entity'] = point_entity
                stop_entities.append(point_entity)
                point['stop_num'] = len(stop_entities)

            if 'splitting' in point:
                for spoint in point['splitting']['list']:

                    spoint['problems'] = self.delayed(
                        spoint.get('et', ''), spoint.get('st', ''), '', '')

                    point_entity = Entity.objects.filter(_identifiers__scheme='crs', _identifiers__value=str(spoint['crs']))
                    if len(point_entity):
                        point_entity = point_entity[0]
                        spoint['entity'] = point_entity
                        stop_entities.append(point_entity)
                        spoint['stop_num'] = len(stop_entities)

        if 'std' in service:
            title = service['std'] + ' ' + service['locationName'] + ' to ' + ' and '.join(destinations)
        else:
            # This service arrives here
            title = service['sta'] + ' from ' + ' and '.join(sources)

        messages = service.get('adhocAlerts', [])

        if 'disruptionReason' in service:
            messages.append(service['disruptionReason'])

        if 'overdueMessage' in service:
            messages.append(service['overdueMessage'])

        return {
            'title': title,
            'entities': stop_entities,
            'ldb': service,
            'calling_points': calling_points,
            'has_timetable': True,
            'has_realtime': True,
            'operator': service.get('operator'),
            'platform': service.get('platform'),
            'messages': messages
        }

    def transform_suds(self, o):
        if isinstance(o, suds.sudsobject.Object):
            return dict((k, self.transform_suds(v)) for k,v in o)
        elif isinstance(o, list):
            return map(self.transform_suds, o)
        else:
            return o

    def _add_error(self, entities):
        for entity in entities:
            entity.metadata['ldb'] = {'error': True}


try:
            service = entity.metadata['ldb_service'](service)
        except suds.WebFault as f:
            if f.fault['faultstring'] == 'Unexpected server error: Invalid length for a Base-64 char array.':
                raise Http404
            else:
                return({'error': f.fault['faultstring']})

import logging
import suds, suds.sudsobject
from suds.sax.element import Element

from django.utils.translation import ugettext_lazy
from django.utils.translation import ugettext as _
from django.conf import settings

from molly.apps.places.models import Entity
from molly.apps.places.providers import BaseMapsProvider

logger = logging.getLogger(__name__)

class LiveDepartureBoardPlacesProvider(BaseMapsProvider):
    _WSDL_URL = "http://example.com"
    _ATTRIBUTION = { 'title': _("Powered by National Rail Enquiries"),
                  'url': "http://www.nationalrail.co.uk",
                  'picture': settings.STATIC_URL + "places/images/powered-by-nre.png" }

    def __init__(self, token, max_services=15, max_results=1):
        self._max_services = max_services
        self._max_results = max_results
        self._token = token

    def delayed(self, eta, sta, etd, std):
        """
        Try and figure out if a service is delayed based on free text values
        for estimated/scheduled times of arrival/depature
        """
        if eta in ('Delayed', 'Cancelled') or etd in ('Delayed', 'Cancelled'):
            # Easy case
            return True
        else:
            # More complex case, have to parse time stamps
            try:
                # Compare scheduled and expected arrival times
                schedh, schedm = sta.split(':')
                exph, expm = eta.rstrip('*').split(':')
            except ValueError:
                pass
            else:
                # Minutes since midnight, as we can't compare time objects
                sched_msm = int(schedh) * 60 + int(schedm)
                exp_msm = int(exph) * 60 + int(expm)

                if exp_msm < sched_msm:
                    # Deal with wraparound at midnight
                    sched_msm += 1440

                if exp_msm - sched_msm >= 5:
                    return True

            try:
                # Compare scheduled and expected departure times
                schedh, schedm = std.split(':')
                exph, expm = etd.rstrip('*').split(':')
            except ValueError:
                pass
            else:
                # Minutes since midnight, as we can't compare time objects
                sched_msm = int(schedh) * 60 + int(schedm)
                exp_msm = int(exph) * 60 + int(expm)

                if exp_msm < sched_msm:
                    # Deal with wraparound at midnight
                    sched_msm += 1440

                if exp_msm - sched_msm >= 5:
                    return True

            return False

    def augment_metadata(self, entities, board='departures', **kwargs):
        station_entities = []
        for entity in entities:
            if not entity.identifiers.get('crs'):
                continue
            station_entities.append(entity)

        station_entities = station_entities[:self._max_results]
        if not station_entities:
            return

        try:
            ldb = suds.client.Client(self._WSDL_URL, soapheaders=Element('AccessToken').insert(Element('TokenValue').setText(self._token)))
        except Exception, e:
            logger.warning("Could not instantiate suds client for live departure board.", exc_info=True, extra={'wsdl_url': self._WSDL_URL})
            self._add_error(station_entities)
            return

        for entity in station_entities:
            try:
                if board == 'arrivals':
                    db = ldb.service.GetArrivalBoard(self._max_services, entity.identifiers['crs'])
                else:
                    db = ldb.service.GetDepartureBoard(self._max_services, entity.identifiers['crs'])
                db = self.transform_suds(db)
                entity.metadata['ldb'] = db
                entity.metadata['service_details'] = lambda s: self.service_details(s, entity)
                entity.metadata['ldb_service'] = lambda s: self.transform_suds(ldb.service.GetServiceDetails(s))
                entity.metadata['service_type'] = 'ldb'

                # Show bus services too
                if board == 'arrivals':
                    db = self.transform_suds(
                        ldb.service.GetDepartureBoard(self._max_services,
                                                     entity.identifiers['crs']))

                if 'trainServices' in db:
                    for service in db['trainServices']['service']:
                        service['problems'] = self.delayed(
                            service.get('eta', ''), service.get('sta', ''),
                            service.get('etd', ''), service.get('std', ''))

                if 'busServices' in db:
                    for service in db['busServices']['service']:
                        entity.metadata['real_time_information'] = {
                            'services':
                                [{
                                    'service': 'BUS',
                                    'destination': service['destination']['location'][0]['locationName'],
                                    'next': service['std'],
                                    'following': [],
                                }]
                        }

            except Exception, e:
                logger.warning("Could not retrieve departure board for station: %r", entity.identifiers.get('crs'),
                               exc_info=True)
                self._add_error((entity,))
            entity.metadata['meta_refresh'] = 60

    def service_details(self, service, entity):

        if service is None:
            return None

        # Trains can split and join, which makes figuring out the list of
        # calling points a bit difficult. The LiveDepartureBoards documentation
        # details how these should be handled. First, we build a list of all
        # the calling points on the "through" train.
        calling_points = service['previousCallingPoints']['callingPointList'][0]['callingPoint'] if len(service['previousCallingPoints']) else []

        # Then attach joining services to our thorough route in the correct
        # point, but only if there is a list of previous calling points
        if len(service['previousCallingPoints']):
            for i, points in enumerate(service['previousCallingPoints']['callingPointList'][1:]):
                # If the other half of the split ends up here, it's not a split, but
                # a change between rail replacement bus of train
                for j, point in enumerate(calling_points):
                    if points['callingPoint'][-1]['crs'] == point['crs']:
                        if j > 0:
                            point['joining'] = points['callingPoint']
                        else:
                            point['service_change'] = service['previousCallingPoints']['callingPointList'][i]['_serviceType']
                            calling_points = points['callingPoint'] + calling_points
                            break

        # Add our current station
        calling_points += [{
            'locationName': service['locationName'],
            'crs': service['crs'],
            'st': service['std'] if 'std' in service else service['sta'],
            'et': service['etd'] if 'etd' in service else service['eta'] if 'eta' in service else '',
            'at': service['atd'] if 'atd' in service else '',
        }]

        if len(service['subsequentCallingPoints']):
            if service['serviceType'] != service['subsequentCallingPoints']['callingPointList'][0]['_serviceType']:
                calling_points[-1]['service_change'] = ugettext_lazy(service['subsequentCallingPoints']['callingPointList'][0]['_serviceType'])

        if len(service['previousCallingPoints']):
            if service['serviceType'] != service['previousCallingPoints']['callingPointList'][0]['_serviceType']:
                calling_points[-1]['service_change'] = ugettext_lazy(service['serviceType'])

        # Now add services going forward
        if len(service['subsequentCallingPoints']):
            calling_points += service['subsequentCallingPoints']['callingPointList'][0]['callingPoint']

        # And do the same with splitting services
        if len(service['subsequentCallingPoints']):
            for i, points in enumerate(service['subsequentCallingPoints']['callingPointList'][1:]):
                # Now we have to handle changes between trains and rail replacement services
                if service['subsequentCallingPoints']['callingPointList'][i]['callingPoint'][-1]['crs'] == points['callingPoint'][0]['crs']:
                    calling_points += points['callingPoint']
                    points['callingPoint'][0]['service_change'] = ugettext_lazy(points['_serviceType'])
                else:
                    for point in calling_points:
                        if points['callingPoint'][0]['crs'] == point['crs']:
                            point['splitting'] = {'destination': points['callingPoint'][-1]['locationName'], 'list': points['callingPoint']}

        sources = [calling_points[0]['locationName']]
        for point in calling_points:
            if 'joining' in point:
                sources.append(point['joining'][0]['locationName'])

        destinations = [calling_points[-1]['locationName']]
        for point in calling_points:
            if 'splitting' in point:
                destinations.append(point['splitting']['destination'])

        stop_entities = []

        # Now get a list of the entities for the stations (if they exist)
        # to plot on a map, and figure out if this stop is delayed or not
        for point in calling_points:

            point['problems'] = self.delayed(
                point.get('et', ''), point.get('st', ''), '', '')

            if 'joining' in point:
                for jpoint in point['joining']:

                    jpoint['problems'] = self.delayed(
                        jpoint.get('et', ''), jpoint.get('st', ''), '', '')

                    point_entity = Entity.objects.filter(_identifiers__scheme='crs', _identifiers__value=str(jpoint['crs']))
                    if len(point_entity):
                        point_entity = point_entity[0]
                        jpoint['entity'] = point_entity
                        stop_entities.append(point_entity)
                        jpoint['stop_num'] = len(stop_entities)

            point_entity = Entity.objects.filter(_identifiers__scheme='crs', _identifiers__value=str(point['crs']))
            if len(point_entity):
                point_entity = point_entity[0]
                point['entity'] = point_entity
                stop_entities.append(point_entity)
                point['stop_num'] = len(stop_entities)

            if 'splitting' in point:
                for spoint in point['splitting']['list']:

                    spoint['problems'] = self.delayed(
                        spoint.get('et', ''), spoint.get('st', ''), '', '')

                    point_entity = Entity.objects.filter(_identifiers__scheme='crs', _identifiers__value=str(spoint['crs']))
                    if len(point_entity):
                        point_entity = point_entity[0]
                        spoint['entity'] = point_entity
                        stop_entities.append(point_entity)
                        spoint['stop_num'] = len(stop_entities)

        if 'std' in service:
            title = service['std'] + ' ' + service['locationName'] + ' to ' + ' and '.join(destinations)
        else:
            # This service arrives here
            title = service['sta'] + ' from ' + ' and '.join(sources)

        messages = service.get('adhocAlerts', [])

        if 'disruptionReason' in service:
            messages.append(service['disruptionReason'])

        if 'overdueMessage' in service:
            messages.append(service['overdueMessage'])

        return {
            'title': title,
            'entities': stop_entities,
            'ldb': service,
            'calling_points': calling_points,
            'has_timetable': True,
            'has_realtime': True,
            'operator': service.get('operator'),
            'platform': service.get('platform'),
            'messages': messages
        }

    def transform_suds(self, o):
        if isinstance(o, suds.sudsobject.Object):
            return dict((k, self.transform_suds(v)) for k,v in o)
        elif isinstance(o, list):
            return map(self.transform_suds, o)
        else:
            return o

    def _add_error(self, entities):
        for entity in entities:
            entity.metadata['ldb'] = {'error': True}


"""Get *all* tasks of the sub dags"""
    qry_sub_dag = (
        session.query(TaskInstance)
        .filter(TaskInstance.dag_id.in_(sub_dag_run_ids), TaskInstance.execution_date.in_(confirmed_dates))
        .filter(or_(TaskInstance.state.is_(None), TaskInstance.state != state))
    )
    return qry_sub_dag

"""Marks tasks APIs."""

import datetime
from typing import Iterable

from sqlalchemy import or_
from sqlalchemy.orm import contains_eager

from airflow.models.baseoperator import BaseOperator
from airflow.models.dagrun import DagRun
from airflow.models.taskinstance import TaskInstance
from airflow.operators.subdag import SubDagOperator
from airflow.utils import timezone
from airflow.utils.session import provide_session
from airflow.utils.state import State
from airflow.utils.types import DagRunType

def _create_dagruns(dag, execution_dates, state, run_type):
    """
    Infers from the dates which dag runs need to be created and does so.

    :param dag: the dag to create dag runs for
    :param execution_dates: list of execution dates to evaluate
    :param state: the state to set the dag run to
    :param run_type: The prefix will be used to construct dag run id: {run_id_prefix}__{execution_date}
    :return: newly created and existing dag runs for the execution dates supplied
    """
    # find out if we need to create any dag runs
    dag_runs = DagRun.find(dag_id=dag.dag_id, execution_date=execution_dates)
    dates_to_create = list(set(execution_dates) - {dag_run.execution_date for dag_run in dag_runs})

    for date in dates_to_create:
        dag_run = dag.create_dagrun(
            execution_date=date,
            start_date=timezone.utcnow(),
            external_trigger=False,
            state=state,
            run_type=run_type,
        )
        dag_runs.append(dag_run)

    return dag_runs

@provide_session
def set_state(
    tasks: Iterable[BaseOperator],
    execution_date: datetime.datetime,
    upstream: bool = False,
    downstream: bool = False,
    future: bool = False,
    past: bool = False,
    state: str = State.SUCCESS,
    commit: bool = False,
    session=None,
):
    """
    Set the state of a task instance and if needed its relatives. Can set state
    for future tasks (calculated from execution_date) and retroactively
    for past tasks. Will verify integrity of past dag runs in order to create
    tasks that did not exist. It will not create dag runs that are missing
    on the schedule (but it will as for subdag dag runs if needed).

    :param tasks: the iterable of tasks from which to work. task.task.dag needs to be set
    :param execution_date: the execution date from which to start looking
    :param upstream: Mark all parents (upstream tasks)
    :param downstream: Mark all siblings (downstream tasks) of task_id, including SubDags
    :param future: Mark all future tasks on the interval of the dag up until
        last execution date.
    :param past: Retroactively mark all tasks starting from start_date of the DAG
    :param state: State to which the tasks need to be set
    :param commit: Commit tasks to be altered to the database
    :param session: database session
    :return: list of tasks that have been created and updated
    """
    if not tasks:
        return []

    if not timezone.is_localized(execution_date):
        raise ValueError(f"Received non-localized date {execution_date}")

    task_dags = {task.dag for task in tasks}
    if len(task_dags) > 1:
        raise ValueError(f"Received tasks from multiple DAGs: {task_dags}")
    dag = next(iter(task_dags))
    if dag is None:
        raise ValueError("Received tasks with no DAG")

    dates = get_execution_dates(dag, execution_date, future, past)

    task_ids = list(find_task_relatives(tasks, downstream, upstream))

    confirmed_dates = verify_dag_run_integrity(dag, dates)

    sub_dag_run_ids = get_subdag_runs(dag, session, state, task_ids, commit, confirmed_dates)

    # now look for the task instances that are affected

    qry_dag = get_all_dag_task_query(dag, session, state, task_ids, confirmed_dates)

    if commit:
        tis_altered = qry_dag.with_for_update().all()
        if sub_dag_run_ids:
            qry_sub_dag = all_subdag_tasks_query(sub_dag_run_ids, session, state, confirmed_dates)
            tis_altered += qry_sub_dag.with_for_update().all()
        for task_instance in tis_altered:
            task_instance.state = state
            if state in State.finished:
                task_instance.end_date = timezone.utcnow()
                task_instance.set_duration()
    else:
        tis_altered = qry_dag.all()
        if sub_dag_run_ids:
            qry_sub_dag = all_subdag_tasks_query(sub_dag_run_ids, session, state, confirmed_dates)
            tis_altered += qry_sub_dag.all()
    return tis_altered

def all_subdag_tasks_query(sub_dag_run_ids, session, state, confirmed_dates):

def get_all_dag_task_query(dag, session, state, task_ids, confirmed_dates):
    """Get all tasks of the main dag that will be affected by a state change"""
    qry_dag = (
        session.query(TaskInstance)
        .join(TaskInstance.dag_run)
        .filter(
            TaskInstance.dag_id == dag.dag_id,
            DagRun.execution_date.in_(confirmed_dates),
            TaskInstance.task_id.in_(task_ids),
        )
        .filter(or_(TaskInstance.state.is_(None), TaskInstance.state != state))
        .options(contains_eager(TaskInstance.dag_run))
    )
    return qry_dag

def get_subdag_runs(dag, session, state, task_ids, commit, confirmed_dates):
    """Go through subdag operators and create dag runs. We will only work
    within the scope of the subdag. We won't propagate to the parent dag,
    but we will propagate from parent to subdag.
    """
    dags = [dag]
    sub_dag_ids = []
    while dags:
        current_dag = dags.pop()
        for task_id in task_ids:
            if not current_dag.has_task(task_id):
                continue

            current_task = current_dag.get_task(task_id)
            if isinstance(current_task, SubDagOperator) or current_task.task_type == "SubDagOperator":
                # this works as a kind of integrity check
                # it creates missing dag runs for subdag operators,
                # maybe this should be moved to dagrun.verify_integrity
                dag_runs = _create_dagruns(
                    current_task.subdag,
                    execution_dates=confirmed_dates,
                    state=State.RUNNING,
                    run_type=DagRunType.BACKFILL_JOB,
                )

                verify_dagruns(dag_runs, commit, state, session, current_task)

                dags.append(current_task.subdag)
                sub_dag_ids.append(current_task.subdag.dag_id)
    return sub_dag_ids

def verify_dagruns(dag_runs, commit, state, session, current_task):
    """Verifies integrity of dag_runs.

    :param dag_runs: dag runs to verify
    :param commit: whether dag runs state should be updated
    :param state: state of the dag_run to set if commit is True
    :param session: session to use
    :param current_task: current task
    :return:
    """
    for dag_run in dag_runs:
        dag_run.dag = current_task.subdag
        dag_run.verify_integrity()
        if commit:
            dag_run.state = state
            session.merge(dag_run)

def verify_dag_run_integrity(dag, dates):
    """
    Verify the integrity of the dag runs in case a task was added or removed
    set the confirmed execution dates as they might be different
    from what was provided
    """
    confirmed_dates = []
    dag_runs = DagRun.find(dag_id=dag.dag_id, execution_date=dates)
    for dag_run in dag_runs:
        dag_run.dag = dag
        dag_run.verify_integrity()
        confirmed_dates.append(dag_run.execution_date)
    return confirmed_dates

def find_task_relatives(tasks, downstream, upstream):
    """Yield task ids and optionally ancestor and descendant ids."""
    for task in tasks:
        yield task.task_id
        if downstream:
            for relative in task.get_flat_relatives(upstream=False):
                yield relative.task_id
        if upstream:
            for relative in task.get_flat_relatives(upstream=True):
                yield relative.task_id

def get_execution_dates(dag, execution_date, future, past):
    """Returns dates of DAG execution"""
    latest_execution_date = dag.get_latest_execution_date()
    if latest_execution_date is None:
        raise ValueError(f"Received non-localized date {execution_date}")
    # determine date range of dag runs and tasks to consider
    end_date = latest_execution_date if future else execution_date
    if 'start_date' in dag.default_args:
        start_date = dag.default_args['start_date']
    elif dag.start_date:
        start_date = dag.start_date
    else:
        start_date = execution_date
    start_date = execution_date if not past else start_date
    if not dag.timetable.can_run:
        # If the DAG never schedules, need to look at existing DagRun if the user wants future or
        # past runs.
        dag_runs = dag.get_dagruns_between(start_date=start_date, end_date=end_date)
        dates = sorted({d.execution_date for d in dag_runs})
    elif not dag.timetable.periodic:
        dates = [start_date]
    else:
        dates = [
            info.logical_date for info in dag.iter_dagrun_infos_between(start_date, end_date, align=False)
        ]
    return dates

@provide_session
def _set_dag_run_state(dag_id, execution_date, state, session=None):
    """
    Helper method that set dag run state in the DB.

    :param dag_id: dag_id of target dag run
    :param execution_date: the execution date from which to start looking
    :param state: target state
    :param session: database session
    """
    dag_run = (
        session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.execution_date == execution_date).one()
    )
    dag_run.state = state
    if state == State.RUNNING:
        dag_run.start_date = timezone.utcnow()
        dag_run.end_date = None
    else:
        dag_run.end_date = timezone.utcnow()
    session.merge(dag_run)

@provide_session
def set_dag_run_state_to_success(dag, execution_date, commit=False, session=None):
    """
    Set the dag run for a specific execution date and its task instances
    to success.

    :param dag: the DAG of which to alter state
    :param execution_date: the execution date from which to start looking
    :param commit: commit DAG and tasks to be altered to the database
    :param session: database session
    :return: If commit is true, list of tasks that have been updated,
             otherwise list of tasks that will be updated
    :raises: ValueError if dag or execution_date is invalid
    """
    if not dag or not execution_date:
        return []

    # Mark the dag run to success.
    if commit:
        _set_dag_run_state(dag.dag_id, execution_date, State.SUCCESS, session)

    # Mark all task instances of the dag run to success.
    for task in dag.tasks:
        task.dag = dag
    return set_state(
        tasks=dag.tasks, execution_date=execution_date, state=State.SUCCESS, commit=commit, session=session
    )

@provide_session
def set_dag_run_state_to_failed(dag, execution_date, commit=False, session=None):
    """
    Set the dag run for a specific execution date and its running task instances
    to failed.

    :param dag: the DAG of which to alter state
    :param execution_date: the execution date from which to start looking
    :param commit: commit DAG and tasks to be altered to the database
    :param session: database session
    :return: If commit is true, list of tasks that have been updated,
             otherwise list of tasks that will be updated
    :raises: AssertionError if dag or execution_date is invalid
    """
    if not dag or not execution_date:
        return []

    # Mark the dag run to failed.
    if commit:
        _set_dag_run_state(dag.dag_id, execution_date, State.FAILED, session)

    # Mark only RUNNING task instances.
    task_ids = [task.task_id for task in dag.tasks]
    tis = (
        session.query(TaskInstance)
        .filter(
            TaskInstance.dag_id == dag.dag_id,
            TaskInstance.execution_date == execution_date,
            TaskInstance.task_id.in_(task_ids),
        )
        .filter(TaskInstance.state == State.RUNNING)
    )
    task_ids_of_running_tis = [task_instance.task_id for task_instance in tis]

    tasks = []
    for task in dag.tasks:
        if task.task_id not in task_ids_of_running_tis:
            continue
        task.dag = dag
        tasks.append(task)

    return set_state(
        tasks=tasks, execution_date=execution_date, state=State.FAILED, commit=commit, session=session
    )

@provide_session
def set_dag_run_state_to_running(dag, execution_date, commit=False, session=None):
    """
    Set the dag run for a specific execution date to running.

    :param dag: the DAG of which to alter state
    :param execution_date: the execution date from which to start looking
    :param commit: commit DAG and tasks to be altered to the database
    :param session: database session
    :return: If commit is true, list of tasks that have been updated,
             otherwise list of tasks that will be updated
    """
    res = []
    if not dag or not execution_date:
        return res

    # Mark the dag run to running.
    if commit:
        _set_dag_run_state(dag.dag_id, execution_date, State.RUNNING, session)

    # To keep the return type consistent with the other similar functions.
    return res


    dag_run = (
        session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.execution_date == execution_date).one()
    )
    dag_run.state = state
    if state == State.RUNNING:
        dag_run.start_date = timezone.utcnow()
        dag_run.end_date = None
    else:
        dag_run.end_date = timezone.utcnow()
    session.merge(dag_run)

"""Marks tasks APIs."""

import datetime
from typing import Iterable

from sqlalchemy import or_
from sqlalchemy.orm import contains_eager

from airflow.models.baseoperator import BaseOperator
from airflow.models.dagrun import DagRun
from airflow.models.taskinstance import TaskInstance
from airflow.operators.subdag import SubDagOperator
from airflow.utils import timezone
from airflow.utils.session import provide_session
from airflow.utils.state import State
from airflow.utils.types import DagRunType

def _create_dagruns(dag, execution_dates, state, run_type):
    """
    Infers from the dates which dag runs need to be created and does so.

    :param dag: the dag to create dag runs for
    :param execution_dates: list of execution dates to evaluate
    :param state: the state to set the dag run to
    :param run_type: The prefix will be used to construct dag run id: {run_id_prefix}__{execution_date}
    :return: newly created and existing dag runs for the execution dates supplied
    """
    # find out if we need to create any dag runs
    dag_runs = DagRun.find(dag_id=dag.dag_id, execution_date=execution_dates)
    dates_to_create = list(set(execution_dates) - {dag_run.execution_date for dag_run in dag_runs})

    for date in dates_to_create:
        dag_run = dag.create_dagrun(
            execution_date=date,
            start_date=timezone.utcnow(),
            external_trigger=False,
            state=state,
            run_type=run_type,
        )
        dag_runs.append(dag_run)

    return dag_runs

@provide_session
def set_state(
    tasks: Iterable[BaseOperator],
    execution_date: datetime.datetime,
    upstream: bool = False,
    downstream: bool = False,
    future: bool = False,
    past: bool = False,
    state: str = State.SUCCESS,
    commit: bool = False,
    session=None,
):
    """
    Set the state of a task instance and if needed its relatives. Can set state
    for future tasks (calculated from execution_date) and retroactively
    for past tasks. Will verify integrity of past dag runs in order to create
    tasks that did not exist. It will not create dag runs that are missing
    on the schedule (but it will as for subdag dag runs if needed).

    :param tasks: the iterable of tasks from which to work. task.task.dag needs to be set
    :param execution_date: the execution date from which to start looking
    :param upstream: Mark all parents (upstream tasks)
    :param downstream: Mark all siblings (downstream tasks) of task_id, including SubDags
    :param future: Mark all future tasks on the interval of the dag up until
        last execution date.
    :param past: Retroactively mark all tasks starting from start_date of the DAG
    :param state: State to which the tasks need to be set
    :param commit: Commit tasks to be altered to the database
    :param session: database session
    :return: list of tasks that have been created and updated
    """
    if not tasks:
        return []

    if not timezone.is_localized(execution_date):
        raise ValueError(f"Received non-localized date {execution_date}")

    task_dags = {task.dag for task in tasks}
    if len(task_dags) > 1:
        raise ValueError(f"Received tasks from multiple DAGs: {task_dags}")
    dag = next(iter(task_dags))
    if dag is None:
        raise ValueError("Received tasks with no DAG")

    dates = get_execution_dates(dag, execution_date, future, past)

    task_ids = list(find_task_relatives(tasks, downstream, upstream))

    confirmed_dates = verify_dag_run_integrity(dag, dates)

    sub_dag_run_ids = get_subdag_runs(dag, session, state, task_ids, commit, confirmed_dates)

    # now look for the task instances that are affected

    qry_dag = get_all_dag_task_query(dag, session, state, task_ids, confirmed_dates)

    if commit:
        tis_altered = qry_dag.with_for_update().all()
        if sub_dag_run_ids:
            qry_sub_dag = all_subdag_tasks_query(sub_dag_run_ids, session, state, confirmed_dates)
            tis_altered += qry_sub_dag.with_for_update().all()
        for task_instance in tis_altered:
            task_instance.state = state
            if state in State.finished:
                task_instance.end_date = timezone.utcnow()
                task_instance.set_duration()
    else:
        tis_altered = qry_dag.all()
        if sub_dag_run_ids:
            qry_sub_dag = all_subdag_tasks_query(sub_dag_run_ids, session, state, confirmed_dates)
            tis_altered += qry_sub_dag.all()
    return tis_altered

def all_subdag_tasks_query(sub_dag_run_ids, session, state, confirmed_dates):
    """Get *all* tasks of the sub dags"""
    qry_sub_dag = (
        session.query(TaskInstance)
        .filter(TaskInstance.dag_id.in_(sub_dag_run_ids), TaskInstance.execution_date.in_(confirmed_dates))
        .filter(or_(TaskInstance.state.is_(None), TaskInstance.state != state))
    )
    return qry_sub_dag

def get_all_dag_task_query(dag, session, state, task_ids, confirmed_dates):
    """Get all tasks of the main dag that will be affected by a state change"""
    qry_dag = (
        session.query(TaskInstance)
        .join(TaskInstance.dag_run)
        .filter(
            TaskInstance.dag_id == dag.dag_id,
            DagRun.execution_date.in_(confirmed_dates),
            TaskInstance.task_id.in_(task_ids),
        )
        .filter(or_(TaskInstance.state.is_(None), TaskInstance.state != state))
        .options(contains_eager(TaskInstance.dag_run))
    )
    return qry_dag

def get_subdag_runs(dag, session, state, task_ids, commit, confirmed_dates):
    """Go through subdag operators and create dag runs. We will only work
    within the scope of the subdag. We won't propagate to the parent dag,
    but we will propagate from parent to subdag.
    """
    dags = [dag]
    sub_dag_ids = []
    while dags:
        current_dag = dags.pop()
        for task_id in task_ids:
            if not current_dag.has_task(task_id):
                continue

            current_task = current_dag.get_task(task_id)
            if isinstance(current_task, SubDagOperator) or current_task.task_type == "SubDagOperator":
                # this works as a kind of integrity check
                # it creates missing dag runs for subdag operators,
                # maybe this should be moved to dagrun.verify_integrity
                dag_runs = _create_dagruns(
                    current_task.subdag,
                    execution_dates=confirmed_dates,
                    state=State.RUNNING,
                    run_type=DagRunType.BACKFILL_JOB,
                )

                verify_dagruns(dag_runs, commit, state, session, current_task)

                dags.append(current_task.subdag)
                sub_dag_ids.append(current_task.subdag.dag_id)
    return sub_dag_ids

def verify_dagruns(dag_runs, commit, state, session, current_task):
    """Verifies integrity of dag_runs.

    :param dag_runs: dag runs to verify
    :param commit: whether dag runs state should be updated
    :param state: state of the dag_run to set if commit is True
    :param session: session to use
    :param current_task: current task
    :return:
    """
    for dag_run in dag_runs:
        dag_run.dag = current_task.subdag
        dag_run.verify_integrity()
        if commit:
            dag_run.state = state
            session.merge(dag_run)

def verify_dag_run_integrity(dag, dates):
    """
    Verify the integrity of the dag runs in case a task was added or removed
    set the confirmed execution dates as they might be different
    from what was provided
    """
    confirmed_dates = []
    dag_runs = DagRun.find(dag_id=dag.dag_id, execution_date=dates)
    for dag_run in dag_runs:
        dag_run.dag = dag
        dag_run.verify_integrity()
        confirmed_dates.append(dag_run.execution_date)
    return confirmed_dates

def find_task_relatives(tasks, downstream, upstream):
    """Yield task ids and optionally ancestor and descendant ids."""
    for task in tasks:
        yield task.task_id
        if downstream:
            for relative in task.get_flat_relatives(upstream=False):
                yield relative.task_id
        if upstream:
            for relative in task.get_flat_relatives(upstream=True):
                yield relative.task_id

def get_execution_dates(dag, execution_date, future, past):
    """Returns dates of DAG execution"""
    latest_execution_date = dag.get_latest_execution_date()
    if latest_execution_date is None:
        raise ValueError(f"Received non-localized date {execution_date}")
    # determine date range of dag runs and tasks to consider
    end_date = latest_execution_date if future else execution_date
    if 'start_date' in dag.default_args:
        start_date = dag.default_args['start_date']
    elif dag.start_date:
        start_date = dag.start_date
    else:
        start_date = execution_date
    start_date = execution_date if not past else start_date
    if not dag.timetable.can_run:
        # If the DAG never schedules, need to look at existing DagRun if the user wants future or
        # past runs.
        dag_runs = dag.get_dagruns_between(start_date=start_date, end_date=end_date)
        dates = sorted({d.execution_date for d in dag_runs})
    elif not dag.timetable.periodic:
        dates = [start_date]
    else:
        dates = [
            info.logical_date for info in dag.iter_dagrun_infos_between(start_date, end_date, align=False)
        ]
    return dates

@provide_session
def _set_dag_run_state(dag_id, execution_date, state, session=None):
    """
    Helper method that set dag run state in the DB.

    :param dag_id: dag_id of target dag run
    :param execution_date: the execution date from which to start looking
    :param state: target state
    :param session: database session
    """

@provide_session
def set_dag_run_state_to_success(dag, execution_date, commit=False, session=None):
    """
    Set the dag run for a specific execution date and its task instances
    to success.

    :param dag: the DAG of which to alter state
    :param execution_date: the execution date from which to start looking
    :param commit: commit DAG and tasks to be altered to the database
    :param session: database session
    :return: If commit is true, list of tasks that have been updated,
             otherwise list of tasks that will be updated
    :raises: ValueError if dag or execution_date is invalid
    """
    if not dag or not execution_date:
        return []

    # Mark the dag run to success.
    if commit:
        _set_dag_run_state(dag.dag_id, execution_date, State.SUCCESS, session)

    # Mark all task instances of the dag run to success.
    for task in dag.tasks:
        task.dag = dag
    return set_state(
        tasks=dag.tasks, execution_date=execution_date, state=State.SUCCESS, commit=commit, session=session
    )

@provide_session
def set_dag_run_state_to_failed(dag, execution_date, commit=False, session=None):
    """
    Set the dag run for a specific execution date and its running task instances
    to failed.

    :param dag: the DAG of which to alter state
    :param execution_date: the execution date from which to start looking
    :param commit: commit DAG and tasks to be altered to the database
    :param session: database session
    :return: If commit is true, list of tasks that have been updated,
             otherwise list of tasks that will be updated
    :raises: AssertionError if dag or execution_date is invalid
    """
    if not dag or not execution_date:
        return []

    # Mark the dag run to failed.
    if commit:
        _set_dag_run_state(dag.dag_id, execution_date, State.FAILED, session)

    # Mark only RUNNING task instances.
    task_ids = [task.task_id for task in dag.tasks]
    tis = (
        session.query(TaskInstance)
        .filter(
            TaskInstance.dag_id == dag.dag_id,
            TaskInstance.execution_date == execution_date,
            TaskInstance.task_id.in_(task_ids),
        )
        .filter(TaskInstance.state == State.RUNNING)
    )
    task_ids_of_running_tis = [task_instance.task_id for task_instance in tis]

    tasks = []
    for task in dag.tasks:
        if task.task_id not in task_ids_of_running_tis:
            continue
        task.dag = dag
        tasks.append(task)

    return set_state(
        tasks=tasks, execution_date=execution_date, state=State.FAILED, commit=commit, session=session
    )

@provide_session
def set_dag_run_state_to_running(dag, execution_date, commit=False, session=None):
    """
    Set the dag run for a specific execution date to running.

    :param dag: the DAG of which to alter state
    :param execution_date: the execution date from which to start looking
    :param commit: commit DAG and tasks to be altered to the database
    :param session: database session
    :return: If commit is true, list of tasks that have been updated,
             otherwise list of tasks that will be updated
    """
    res = []
    if not dag or not execution_date:
        return res

    # Mark the dag run to running.
    if commit:
        _set_dag_run_state(dag.dag_id, execution_date, State.RUNNING, session)

    # To keep the return type consistent with the other similar functions.
    return res