import library
import json

@library.export
def init(args):
    model = [[9.2, 0.21, 0.21],
             [8.2, 0.22, 0.21],
             [7.2, 1.21, 2.41],
             [1.2, 2.21, 0.29]]
    library.put("model", model)
    ROUND = 0
    library.put("ROUND", ROUND)
    alpha = 0.2
    library.put("alpha", alpha)

@library.export
def clientUpload(args):
    # get client model
    client = json.loads(args["data"])

    # client round
    k = "round" + str(client["round"])

    # save model to buckets
    library.put_bucket(k, client["model"])

    # if enough models
    if library.count_bucket(k) > 20:
        ROUND = library.get("ROUND")

        # check client rounds == current rounds
        if ROUND != client["round"]:
            return False

        # set round to -1 to prevent clients uploading to this bucket
        library.put("ROUND", -1)

        model = library.get("model")

        list_weights = library.get_bucket(k)
        model = updateModel(model, list_weights)

        # save calculated model and restore round
        library.put("model", model)
        library.put("ROUND", ROUND+1)
    return True

def updateModel(model, list_weights):
    """
        list_weights : 3D list of shape : (clientNumber,modelOuter, modelInner)
        It contains all the models for each client
    """

    # this part will change developer to developer
    # one can just take avg
    # or one can discard smallest and largest than take average
    # this example just takes avg without use of external library

    alpha = library.get("alpha")

    # getting shape of 3D array
    number_clients = len(list_weights)
    size_outer = len(list_weights[0])
    size_inner = len(list_weights[0][0])

    # constructing a new 2D array of zeros of same size
    newModel = [ [0 for j in range(size_inner)] for i in range(size_outer)]

    # validate new created shape
    assert(len(newModel) == size_outer)
    assert(len(newModel[0]) == size_inner)

    # sum for all the clients
    for weights in list_weights:
        for outerIndex, outerList in enumerate(weights):
            for innerIndex, innerVal in enumerate(outerList):
                newModel[outerIndex][innerIndex] += innerVal

    # average it by number of clients
    for outerIndex, outerList in enumerate(newModel):
        for innerIndex, innerVal in enumerate(outerList):
            newModel[outerIndex][innerIndex] /= number_clients

    # now update the model using the learning rate using below formula
    # model = (1-a) * model  + a * new_model
    # Prev. part and next part could be merged for efficiency but readability they implemented with two loops

    # Iterate over model

    for outerIndex, outerList in enumerate(newModel):
        for innerIndex, innerVal in enumerate(outerList):
            model[outerIndex][innerIndex] *= 1-alpha
            model[outerIndex][innerIndex] += alpha * newModel[outerIndex][innerIndex]
    # Finally update round number

    return model

@library.export
def getModel(args):
    return library.get("model")

@library.export
def getRound(args):
    return library.get("ROUND")

[((10, 4, 10, 31), 'library.put', 'library.put', ({(10, 16, 10, 23): '"""model"""', (10, 25, 10, 30): 'model'}, {}), "('model', model)", False, 'import library\n'), ((12, 4, 12, 31), 'library.put', 'library.put', ({(12, 16, 12, 23): '"""ROUND"""', (12, 25, 12, 30): 'ROUND'}, {}), "('ROUND', ROUND)", False, 'import library\n'), ((14, 4, 14, 31), 'library.put', 'library.put', ({(14, 16, 14, 23): '"""alpha"""', (14, 25, 14, 30): 'alpha'}, {}), "('alpha', alpha)", False, 'import library\n'), ((19, 13, 19, 37), 'json.loads', 'json.loads', ({(19, 24, 19, 36): "args['data']"}, {}), "(args['data'])", False, 'import json\n'), ((25, 4, 25, 42), 'library.put_bucket', 'library.put_bucket', ({(25, 23, 25, 24): 'k', (25, 26, 25, 41): "client['model']"}, {}), "(k, client['model'])", False, 'import library\n'), ((59, 12, 59, 32), 'library.get', 'library.get', ({(59, 24, 59, 31): '"""alpha"""'}, {}), "('alpha')", False, 'import library\n'), ((100, 11, 100, 31), 'library.get', 'library.get', ({(100, 23, 100, 30): '"""model"""'}, {}), "('model')", False, 'import library\n'), ((104, 11, 104, 31), 'library.get', 'library.get', ({(104, 23, 104, 30): '"""ROUND"""'}, {}), "('ROUND')", False, 'import library\n'), ((28, 7, 28, 30), 'library.count_bucket', 'library.count_bucket', ({(28, 28, 28, 29): 'k'}, {}), '(k)', False, 'import library\n'), ((29, 16, 29, 36), 'library.get', 'library.get', ({(29, 28, 29, 35): '"""ROUND"""'}, {}), "('ROUND')", False, 'import library\n'), ((36, 8, 36, 32), 'library.put', 'library.put', ({(36, 20, 36, 27): '"""ROUND"""', (36, 29, 36, 31): '(-1)'}, {}), "('ROUND', -1)", False, 'import library\n'), ((38, 16, 38, 36), 'library.get', 'library.get', ({(38, 28, 38, 35): '"""model"""'}, {}), "('model')", False, 'import library\n'), ((40, 23, 40, 44), 'library.get_bucket', 'library.get_bucket', ({(40, 42, 40, 43): 'k'}, {}), '(k)', False, 'import library\n'), ((44, 8, 44, 35), 'library.put', 'library.put', ({(44, 20, 44, 27): '"""model"""', (44, 29, 44, 34): 'model'}, {}), "('model', model)", False, 'import library\n'), ((45, 8, 45, 37), 'library.put', 'library.put', ({(45, 20, 45, 27): '"""ROUND"""', (45, 29, 45, 36): '(ROUND + 1)'}, {}), "('ROUND', ROUND + 1)", False, 'import library\n')]