#===============================================================================
# Copyright 2014-2017 Intel Corporation All Rights Reserved.
#
# The source code,  information  and material  ("Material") contained  herein is
# owned by Intel Corporation or its  suppliers or licensors,  and  title to such
# Material remains with Intel  Corporation or its  suppliers or  licensors.  The
# Material  contains  proprietary  information  of  Intel or  its suppliers  and
# licensors.  The Material is protected by  worldwide copyright  laws and treaty
# provisions.  No part  of  the  Material   may  be  used,  copied,  reproduced,
# modified, published,  uploaded, posted, transmitted,  distributed or disclosed
# in any way without Intel's prior express written permission.  No license under
# any patent,  copyright or other  intellectual property rights  in the Material
# is granted to  or  conferred  upon  you,  either   expressly,  by implication,
# inducement,  estoppel  or  otherwise.  Any  license   under such  intellectual
# property rights must be express and approved by Intel in writing.
#
# Unless otherwise agreed by Intel in writing,  you may not remove or alter this
# notice or  any  other  notice   embedded  in  Materials  by  Intel  or Intel's
# suppliers or licensors in any way.
#===============================================================================

## <a name="DAAL-EXAMPLE-PY-IMPLICIT_ALS_CSR_DISTRIBUTED"></a>
## \example impl_als_csr_distr.py

from os import environ
from os.path import join as jp

import numpy

from daal import step1Local, step2Master, step3Local, step4Local

import daal.algorithms.implicit_als.prediction.ratings as ratings
import daal.algorithms.implicit_als.training as training
import daal.algorithms.implicit_als.training.init as init

from daal.data_management import KeyValueDataCollection, HomogenNumericTable

import inspect, sys, os.path
utils_folder = os.path.realpath(os.path.abspath(jp(os.path.split(inspect.getfile(inspect.currentframe()))[0], "..")))
if utils_folder not in sys.path:
    sys.path.insert(0, utils_folder)
from utils import createSparseTable, computePartialModelBlocksToNode, printALSRatings

DAAL_PREFIX = jp('..', 'data')

# Input data set parameters
nBlocks = 4

# Number of observations in training data set blocks
trainDatasetFileNames = [
    jp(DAAL_PREFIX, 'distributed', 'implicit_als_csr_1.csv'),
    jp(DAAL_PREFIX, 'distributed', 'implicit_als_csr_2.csv'),
    jp(DAAL_PREFIX, 'distributed', 'implicit_als_csr_3.csv'),
    jp(DAAL_PREFIX, 'distributed', 'implicit_als_csr_4.csv')
]

# Number of observations in transposed training data set blocks
transposedTrainDatasetFileNames = [
    jp(DAAL_PREFIX, 'distributed', 'implicit_als_trans_csr_1.csv'),
    jp(DAAL_PREFIX, 'distributed', 'implicit_als_trans_csr_2.csv'),
    jp(DAAL_PREFIX, 'distributed', 'implicit_als_trans_csr_3.csv'),
    jp(DAAL_PREFIX, 'distributed', 'implicit_als_trans_csr_4.csv')
]

usersPartition = [0] * (nBlocks + 1)
itemsPartition = [0] * (nBlocks + 1)

# Algorithm parameters
nUsers = 46        # Full number of users
nFactors = 2       # Number of factors
maxIterations = 5  # Number of iterations in the implicit ALS training algorithm

dataTable = [0] * nBlocks
transposedDataTable = [0] * nBlocks

predictedRatings = [[0] * nBlocks for x in range(nBlocks)]

userFactorsToNodes = [None] * nBlocks
itemFactorsToNodes = [None] * nBlocks

itemsPartialResultLocal = [0] * nBlocks
usersPartialResultLocal = [0] * nBlocks

step1LocalResult = [0] * nBlocks
step2MasterResult = None
step3LocalResult = [0] * nBlocks
step4LocalInput = [0] * nBlocks


def readData(block):
    global dataTable, transposedDataTable

    # Read trainDatasetFileName from a file and create a numeric table to store the input data
    dataTable[block] = createSparseTable(trainDatasetFileNames[block])

    # Read trainDatasetFileName from a file and create a numeric table to store the input data
    transposedDataTable[block] = createSparseTable(transposedTrainDatasetFileNames[block])


def initializeModel(block):
    global itemsPartialResultLocal

    # Create an algorithm object to initialize the implicit ALS model with the fastCSR method
    initAlgorithm = init.Distributed(step=step1Local, method=init.fastCSR)
    initAlgorithm.parameter.fullNUsers = nUsers
    initAlgorithm.parameter.nFactors = nFactors
    initAlgorithm.parameter.seed += block

    # Pass a training data set to the algorithm
    initAlgorithm.input.set(init.data, transposedDataTable[block])

    # Initialize the implicit ALS model
    res = initAlgorithm.compute()
    partialModelLocal = res.get(init.partialModel)

    itemsPartialResultLocal[block] = training.DistributedPartialResultStep4()
    itemsPartialResultLocal[block].set(training.outputOfStep4ForStep1, partialModelLocal)


def computeStep1Local(block, partialResultLocal):
    global step1LocalResult

    # Create an algorithm object to perform first step of the implicit ALS training algorithm on local-node data
    algorithm = training.Distributed(step=step1Local)
    algorithm.parameter.nFactors = nFactors

    # Set input objects for the algorithm
    algorithm.input.set(training.partialModel, partialResultLocal[block].get(training.outputOfStep4ForStep1))

    # Compute partial results of the first step on local nodes
    # DistributedPartialResultStep1 class from training
    res = algorithm.compute()
    step1LocalResult[block] = res


def computeStep2Master():
    global step2MasterResult

    # Create an algorithm object to perform second step of the implicit ALS training algorithm
    algorithm = training.Distributed(step=step2Master)
    algorithm.parameter.nFactors = nFactors

    # Set the partial results of the first local step of distributed computations
    # as input for the master-node algorithm
    for i in range(nBlocks):
        algorithm.input.add(training.inputOfStep2FromStep1, step1LocalResult[i])

    # Compute a partial result on the master node from the partial results on local nodes
    # DistributedPartialResultStep2 class from training
    res = algorithm.compute()
    step2MasterResult = res.get(training.outputOfStep2ForStep4)


def computeStep3Local(block, offsets, partialResultLocal, factorsToNodes):
    global step3LocalResult, step4LocalInput

    # Create an algorithm object to perform third step of the implicit ALS training algorithm on local-node data
    algorithm = training.Distributed(step=step3Local)
    algorithm.parameter.nFactors = nFactors

    offset_array = numpy.array([[offsets[block]]], dtype=numpy.float64)
    offsetTable = HomogenNumericTable(offset_array)

    # Set input objects for the algorithm
    algorithm.input.setModel(training.partialModel, partialResultLocal[block].get(training.outputOfStep4ForStep3))
    algorithm.input.setCollection(training.partialModelBlocksToNode, factorsToNodes[block])
    algorithm.input.setTable(training.offset, offsetTable)

    # Compute partial results of the third step on local nodes
    # DistributedPartialResultStep3 class from training
    res = algorithm.compute()
    step3LocalResult[block] = res.get(training.outputOfStep3ForStep4)

    # Prepare input objects for the fourth step of the distributed algorithm
    for i in range(nBlocks):
        step4LocalInput[i][block] = step3LocalResult[block][i]


def computeStep4Local(block, dataTable, partialResultLocal):

    # Create an algorithm object to perform fourth step of the implicit ALS training algorithm on local-node data
    algorithm = training.Distributed(step=step4Local)
    algorithm.parameter.nFactors = nFactors

    # Set input objects for the algorithm
    algorithm.input.setModels(training.partialModels, step4LocalInput[block])
    algorithm.input.setTable(training.partialData, dataTable[block])
    algorithm.input.setTable(training.inputOfStep4FromStep2, step2MasterResult)

    # Build the implicit ALS partial model on the local node
    # DistributedPartialResultStep4 class from training
    res = algorithm.compute()
    partialResultLocal[block] = res


def trainModel():

    for iteration in range(maxIterations):

        # Update partial users factors
        for i in range(nBlocks):
            computeStep1Local(i, itemsPartialResultLocal)

        computeStep2Master()

        for i in range(nBlocks):
            computeStep3Local(i, itemsPartition, itemsPartialResultLocal, itemFactorsToNodes)

        for i in range(nBlocks):
            computeStep4Local(i, dataTable, usersPartialResultLocal)

        # Update partial items factors
        for i in range(nBlocks):
            computeStep1Local(i, usersPartialResultLocal)

        computeStep2Master()

        for i in range(nBlocks):
            computeStep3Local(i, usersPartition, usersPartialResultLocal, userFactorsToNodes)

        for i in range(nBlocks):
            computeStep4Local(i, transposedDataTable, itemsPartialResultLocal)


def testModel(i, j):
    # Create an algorithm object to predict ratings based in the implicit ALS partial models
    algorithm = ratings.Distributed(step=step1Local, method=ratings.defaultDense)
    algorithm.parameter.nFactors = nFactors

    # Set input objects for the algorithm
    algorithm.input.set(ratings.usersPartialModel, usersPartialResultLocal[i].get(training.outputOfStep4))
    algorithm.input.set(ratings.itemsPartialModel, itemsPartialResultLocal[j].get(training.outputOfStep4))

    # Predict ratings and retrieve the algorithm results
    algorithm.compute()

    # Result class from ratings
    res = algorithm.finalizeCompute()
    return res.get(ratings.prediction)


def printResults():

    for i in range(nBlocks):
        for j in range(nBlocks):
            print("Ratings for users block {}, items block {} :".format(i, j))
            printALSRatings(usersPartition[i], itemsPartition[j], predictedRatings[i][j])

if __name__ == "__main__":
    for i in range(nBlocks):
        readData(i)
        initializeModel(i)
        step4LocalInput[i] = KeyValueDataCollection()

    computePartialModelBlocksToNode(
        nBlocks, dataTable, transposedDataTable, usersPartition, itemsPartition,
        userFactorsToNodes, itemFactorsToNodes
    )

    trainModel()

    for i in range(nBlocks):
        for j in range(nBlocks):
            predictedRatings[i][j] = testModel(i, j)
    printResults()
