/* NiuTrans.Tensor - an open-source tensor library
 * Copyright (C) 2018, Natural Language Processing Lab, Northestern University.
 * All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*
 * $Created by: XIAO Tong (xiaotong@mail.neu.edu.cn) 2018-07-12
 */

#include "XNet.h"
#include "XNoder.h"
#include "XBackwardLoss.h"
#include "XBackwardMath.h"
#include "XBackwardFunc.h"
#include "XBackwardShape.h"
#include "../tensor/XName.h"

namespace nts{

unsigned int netIDGlobal = 0;
MUTEX_HANDLE netMutex;

/* generate a network id */
unsigned int MakeNetID()
{
    if(netIDGlobal == 0)
        MUTEX_INIT(netMutex);

    MUTEX_LOCK(netMutex);
    netIDGlobal += 3;
    unsigned int id = netIDGlobal;
    MUTEX_UNLOCK(netMutex);

    return id;
}

void XNetClearAll()
{
    MUTEX_DELE(netMutex);
}

/* constructor */
XNet::XNet()
{
    nodes.Clear();
}

/* de-constructor */
XNet::~XNet()
{
}

/* clear the network */
void XNet::Clear()
{
    nodes.Clear();
    gradNodes.Clear();
    outputs.Clear();
    inputs.Clear();
}

/* 
backward propagation to obtain gradient wrt. the loss/error function 
>> root - root node (output) of the network
>> gold - gold standard for the output
>> loss - name of loss function
*/
void XNet::Backward(XTensor &root, XTensor &gold, LOSS_FUNCTION_NAME loss)
{
    XList roots(1);
    roots.Add(&root);

    XList golds(1);
    golds.Add(&gold);

    Backward(roots, golds, loss);
}

/* 
backward propagation to obtain gradient
>> root - root node (output) of the network
>> loss - name of loss function
*/
void XNet::Backward(XTensor &root, LOSS_FUNCTION_NAME loss)
{
    XList roots(1);
    roots.Add(&root);

    XList golds(1);
    golds.Add(NULL);

    Backward(roots, golds, loss);
}

/* 
backward propagation to obtain gradient wrt. the loss/error function
with a number of root nodes 
>> root - a list of root nodes (output) of the network
>> gold - a list of gold standard for the output
>> loss - name of loss function
*/
void XNet::Backward(XList &roots, XList &golds, LOSS_FUNCTION_NAME loss)
{
    Traverse(roots);

    for(int i = 0; i < nodes.count; i++){
        XTensor * node = (XTensor*)nodes.Get(i);
        node->visitMark = NODE_UNFINISHED;
    }

    XLossGrad lossGrad;

    /* we start with the gradient with respect to the loss for output layers */
    for(int i = 0; i < roots.count; i++){
        XTensor * root = (XTensor*)roots.Get(i);
        XTensor * gold = (XTensor*)golds.Get(i);
        XLink &income = root->income;
        int funcID = income.typeID;
        void * params = income.params;

        /* we compute dE/dx if the output is generated by an activation function y = f(x).
           Note that we do not need to obtain dE/dy here because it is no use in the 
           folloing process of back-propagation */
        if(gold != NULL && income.tailNum == 1 && (funcID & FUNCTION_BASE)){
            XTensor * x = income.tails[0];
            XNoder::MakeGrad(x);
            lossGrad.Compute(gold, root, x, NULL, x->grad, funcID, params, loss);
            root->visitMark = NODE_FINISHED;
        }
        /* we compuate dE/dy (y is the output) if no predefined activation function is used */
        else{
            XNoder::MakeGrad(root);
            lossGrad.Compute(gold, root, root->grad, loss);
        }
    }
    
    /* back-propagation from output to input */
    for(int i = nodes.count - 1; i >= 0; i--){
        XTensor * node = (XTensor*)nodes.Get(i);

        if(node->mem != NULL){
            CheckNTErrors(node->mem->bufUsed < BUF_PITCH, "Illegal access of buffer!");
        }

        if(node->visitMark == NODE_FINISHED)
            continue;

        BackwardNode(node);  
    }
}

/* 
backward propagation to obtain gradient
with a number of root nodes 
>> root - a list of root nodes (output) of the network
>> loss - name of loss function
*/
void XNet::Backward(XList &roots, LOSS_FUNCTION_NAME loss)
{
    XList golds(roots.count);
    for(int i = 0; i < roots.count; i++)
        golds.Add(NULL);

    Backward(roots, golds, loss);
}

/* 
backward computation for a given node 
>> node - the node keeps the result of an operation (e.g., activation function)
*/
void XNet::BackwardNode(XTensor * node)
{
    if(node == NULL || node->visitMark == NODE_FINISHED)
        return;

    if(!XNoder::IsLeaf(node)){
        /* post processing for parent nodes */
        BackwardNodePost(node);

        /* process the current node */
        if(XMathGrad::IsMathOP(node))
            XMathGrad::MakeGrad(node);
        else if(XFuncGrad::IsFunc(node))
            XFuncGrad::MakeGrad(node);
        else if(XShapeGrad::IsShapeOP(node))
            XShapeGrad::MakeGrad(node);
        else{
            ShowNTErrors("Wrong node type!");
        }
    }
}

/* 
backward computation (in post processing) for a given node 
>> node - the node whose parent nodes are not processed yet. So
          we do the job at the child node.
*/
void XNet::BackwardNodePost(XTensor * node)
{
    bool isSplitList = false;
    XLink &outgo = node->outgo;
    for(int i = 0; i < outgo.tailNum; i++){
        if(outgo.tails[i]->income.typeID == SHAPE_SPLIT_LIST)
            isSplitList = true;
    }

    if(isSplitList)
        XShapeGrad::PostProcessing(node, SHAPE_SPLIT_LIST);
}

/* 
traverse the net and find the topological order by 
depth-first search (Tarjan's algorithm) 
>> root - root node (or output of the net)
*/
void XNet::Traverse(XTensor &root)
{
    XList roots(1);
    roots.Add(&root);

    Traverse(roots);
}

/* 
traverse the net and find the topological order by 
depth-first search (Tarjan's algorithm) 
>> roots - a list of roots (or output nodes)
*/
void XNet::Traverse(XList &roots)
{
    id = MakeNetID();
    nodes.Clear();
 
    for (int i = 0; i < roots.count; i++)
        TarjanVisit((XTensor*)roots.Get(i), nodes, id);

    for(int i = 0; i < nodes.count; i++){
        XTensor * node = (XTensor*)nodes.Get(i);
        if(XNoder::IsRoot(node))
            outputs.Add(node);
        if(XNoder::IsLeaf(node))
            inputs.Add(node);
        if(XNoder::IsGrad(node))
            gradNodes.Add(node);
    }
}

/* 
depth-first search given a node (Tarjan's algorithm for topological ordering)
>> node - the node to visit (mark 0:unvisited, 1:visiting, 2:done)
>> orders - topological order of the nodes
>> code - code of the network
*/
void XNet::TarjanVisit(XTensor * node, XList &orders, const unsigned int code)
{
    if(node == NULL)
        return;

    //fprintf(stderr, "%d\n", node->id);
    if(node->visitMark == code + 1){
        ShowNTErrors("There is a circle in the network\n");
    }
    else if(node->visitMark <= code){
        node->visitMark = code + 1;
        XLink &income = node->income;
        for(int i = 0; i < income.tailNum; i++){
            XTensor * child = income.tails[i];
            if(child == NULL)
                continue;
            TarjanVisit(child, orders, code);
        }
        node->visitMark = code + 2;
        orders.Add(node);
    }
}

/* 
dump network information 
>> file - the file for dumping
*/
void XNet::Dump(FILE * file)
{
    for(int i = 0; i < nodes.count; i++){
        XTensor * node =  (XTensor*)nodes.Get(i);
        fprintf(file, "node %d: %d\n", i, node->id);
        node->Dump(file, "tensor: ");
        if(node->grad != NULL)
            node->grad->Dump(file, "grad: ");
        else
            fprintf(file, "no gradient!\n");
        fprintf(file, "\n");
    }
}

}