/**
 * @file: Mlp.cpp
 *
 * Implementation of the bare-bones MLP computations. See header file
 * for documentation.
 */

#include "Mlp.hpp"

#include <iostream>
#include <sstream>
#include <stdexcept>
#include <cmath>

using std::ostream;
using std::ostringstream;
using std::runtime_error;
using std::endl;

using namespace jymlp;

/** Local helper function: Compute total number of weights. 
 * TODO: Move to some utils.cpp if need arises.
 */
static size_t total_nweights(const vector<size_t>& nneur){
    if (nneur.size() == 0) return 0;
    size_t tot = 0;
    for (size_t i=1; i<nneur.size(); ++i){
        tot += nneur[i] * (nneur[i-1] + 1);
    }
    return tot;
}

static
string toString(ActF a){
    switch(a){
    case actf::Unset:
        return "Unset";
    case actf::linear:
        return "lin";
    case actf::hyptan:
        return "tanh";
    case actf::logsig:
        return "logsig";
    default:
        throw runtime_error("Cannot convert ActF to string.");
    }
}

using namespace jymlp;

// -------------------------- Constructors:

Mlp::Mlp(){
    // Uninitialized! No layers, no weights, nothing!
}

Mlp::Mlp(const vector<size_t>& inneur){
    nneur = inneur;
    weights.resize(total_nweights(nneur));
    actf.resize(nneur.size());
    for (double & w : weights){
        w = 0.0;
    }
    actf[0] = actf::Unset;
    for (size_t i=1; i<actf.size(); ++i){
        actf[i]=(i<actf.size()-1)?actf::hyptan:actf::linear;
    }
}

Mlp::Mlp(const vector<size_t>& inneur,
         const vector<ActF>& iactf,
         const vector<double>& iweights){
    nneur = inneur;
    actf = iactf;
    weights = iweights;
}

Mlp::Mlp(const Mlp &other){
    nneur = other.nneur;
    actf = other.actf;
    weights = other.weights;
}

Mlp::Mlp(const vector<size_t>& inneur, const vector<ActF>& iactf, SynapticRandomizer& sr){
    nneur = inneur;
    actf = iactf;
    weights.resize(total_nweights(nneur));
    initRnd(1.0,sr);
}

// -------------------------- Protected:

void Mlp::initRnd(double a, SynapticRandomizer & sr){
    for(double & w : weights){
        w = sr.rndU(a);
    }
}

// FIXME: Rethink the computations - need less space after all?
size_t Mlp::getWorkspaceSize(){
    size_t maxn = 0;
    for (size_t i=0;i<getNLayers();++i){
        if (nneur[i] > maxn) maxn = nneur[i];
    }
    return maxn*getNLayers();
}

// -------------------------- Public:

size_t Mlp::getNHiddenNeurons() const{
    size_t res = 0;
    for (size_t i=1;i<getNLayers()-1;++i){
        res += nneur[i];
    }
    return res;
}

size_t Mlp::getNumNonzeroWeights() const{
    size_t res = 0;
    for(const double & w : weights){
        if (w != 0.0) res++;
    }
    return res;
}

size_t Mlp::getNumConnectedInputs() const{
    size_t res = 0;
    for(size_t iin = 0; iin<nneur[0]; ++iin){
        bool isConnected = false;
        for(size_t ineur = 0 ; ineur < nneur[1] ; ++ineur){
            // Indexing the weights in the packed representation is a
            // bit tedious and prone to errors. TODO: easier access
            // functions; possibly functions like
            // abscolsum(layer,col). Then just
            // "res+=abscolsum(1,iin)>0.0?1:0" in here. ... instead of
            // this obscure passage (which can "short circuit", though):
            if (weights[ineur*(nneur[0]+1)+1+iin] != 0.0){
                isConnected = true;
                break; // No need to look further.
            }
            // .. and, of course, this is one of the representation
            // issues in evolutionary MLPs. We could have a separate
            // bit string for active inputs, but we decided to follow
            // the "straightforward" storage format here, for
            // resemblance to pre-existing MLP codes.
        }
        if (isConnected) res++;
    }
    return res;
}


string
Mlp::prettyPrintLayerSizes(){
    ostringstream oss("");
    size_t nlay = getNLayers();
    for (size_t i=0;i<nlay;++i){
        oss << nneur[i];
        if (i<nlay-1) {oss << "-";}
    }
    return oss.str();
}

string
Mlp::prettyPrintLayerActivations(){
   ostringstream oss("");
    size_t nlay = getNLayers();
    oss << "in-";
    for (size_t i=1;i<nlay;++i){
        oss << toString(actf[i]);
        if (i<nlay-1) {oss << "-";}
    }
    return oss.str();
}


static string doPrettyPrint(const vector<size_t> & nneur, const vector<double> weights){
    ostringstream oss("");
    // TODO: Can only do plaintext as of now. LaTeX could be nice?
    size_t wcount = 0;
    size_t nlay = nneur.size();
    for (size_t lay = 1; lay < nlay ; ++lay){
        for (size_t ineur = 0; ineur < (size_t) nneur[lay]; ++ineur){
            for (size_t iw = 0; iw < (size_t) nneur[lay-1] + 1; ++iw){
                oss << weights[wcount++];
                if (iw<(size_t)nneur[lay-1]) { oss << " "; } else {oss << endl;}
            }
        }
        oss << endl;
    }
    return oss.str();
}

string
Mlp::prettyPrintWeights() const{
    return doPrettyPrint(nneur,weights);
}

string
Mlp::prettyPrintGradient(const vector <double> & grad) const{
    return doPrettyPrint(nneur,grad);
}


string
Mlp::prettyPrint(){
    ostringstream oss;
    oss << prettyPrintLayerSizes();
    oss << " (" << prettyPrintLayerActivations() << ")" << endl;
    oss << prettyPrintWeights() << endl;
    return oss.str();
}

void
Mlp::toStream(ostream & o){
    size_t nlay = getNLayers();
    o << "1 "; // 'version ID'
    o << nlay;
    for (size_t i=0;i<nlay;++i){
        o  << " " << nneur[i];
    }
    // FIXME: See if init is correct for actf[0]!!
    for (size_t i=0;i<nlay;++i){
        o  << " " << actf[i];
    }
    for (size_t i=0;i<weights.size();++i){
        o << " " << weights[i];
    }
}

void
Mlp::fromStream(istream & ins){
    int version_id;
    size_t s;
    ins >> version_id; // Expect 0 so far..
    switch (version_id){
    case 1:
        ins >> s;
        nneur.resize(s);
        for (size_t i = 0;i<nneur.size();++i){
            if (ins.eof()) throw runtime_error("Too few values for MLP");
            ins >> nneur[i];
        }
        actf.resize(nneur.size());
        for (size_t i = 0;i<actf.size();++i){
            if (ins.eof()) throw runtime_error("Too few values for MLP");
            int en;
            ins >> en;
            actf[i] = static_cast<ActF>(en);
        }
        weights.resize(total_nweights(nneur));
        for (size_t i = 0;i<weights.size();++i){
            if (ins.eof()) throw runtime_error("Too few values for MLP");
            ins >> weights[i];
        }
        break;
    default:
        throw runtime_error("Unknown version_id");
    }
}

// MLP evaluations

/* Activation function on forward pass */
static double act (ActF af, double s){
    switch(af){
    case actf::linear:
        return s;
    case actf::logsig:
        return 1.0/(1.0 + exp(-s));
    case actf::hyptan:
        return 2.0/(1.0 + exp(-2.0 * s)) -1.0;
    default:
        throw runtime_error("Unknown activation function");
    }
}

/* Activation function derivative on backward pass; fp must be the
 * previously computed value
 */
static double actD (ActF af, double fp){
    switch(af){
    case actf::linear:
        return 1.0;
    case actf::logsig:
        return fp * (1.0 - fp);
    case actf::hyptan:
        return (1.0 - fp * fp);
    default:
        throw runtime_error("Unknown activation function");
    }
}


void
Mlp::forward(const vector<double> &input, double * workspace) const{
    /* Make a copy of the input data. Costly, but simplifies backward loop: */
    for (size_t i=0;i<input.size();++i){
        workspace[i] = input[i];
    }

    const double *oprev = workspace;    // The "o^0" of our iteration equation.
    double *othis = workspace+nneur[0]; // The "o^1" of our iteration equation.
    const double *wthis = weights.data();

    //cerr << prettyPrintWeights();

    // Our layer index lay==0 for input, lay==1 for first hidden layer etc.
    // NLayers includes input, so for 2-3-2 ANN NLayers=3, and C++ indices for layers
    // are 0, 1, and 2.
    for (size_t lay = 1; lay < getNLayers(); ++lay){
        size_t nprev = nneur[lay-1];
        size_t nthis = nneur[lay];
        for (size_t neur = 0; neur < nthis; ++neur){
            othis[neur] = wthis[0]; // bias;
            //cerr << "actf(" << othis[neur];
            // careful with the indices here:
            for(size_t conn = 0; conn < nprev; ++conn){
                othis[neur] += wthis[1+conn] * oprev[conn];
                //cerr << " + (" << wthis[1+conn] << "*" << oprev[conn] << ")";
            }
            othis[neur] = act(actf[lay],othis[neur]);
            //cerr << ") =" << othis[neur] << endl;
            wthis += (nprev + 1); // next weight matrix row.
        }
        oprev = othis;       // update previous layer ptr
        othis += nthis;      // next workspace vector.
    }
}

/* Return location of output layer in the workspace. Observe: We have
   made a copy of the input, which is taken into account here. */
static double * outputLayerPtr(const vector <size_t> & nneur, double * ws){
    for(size_t lay=0;lay<nneur.size()-1;++lay){
        ws += nneur[lay];
    }
    return ws;
}

static const double * outputLayerPtr(const vector <size_t> & nneur,
                                     const double * ws){
    for(size_t lay=0;lay<nneur.size()-1;++lay){
        ws += nneur[lay];
    }
    return ws;
}


/* Return location of layer gradient. Input is not in the gradient! */
static double * outputLayerGradPtr(const vector <size_t> & nneur, double * g){
    for(size_t lay=1;lay<nneur.size()-1;++lay){
        g += nneur[lay]*(1+nneur[lay-1]);
    }
    return g;
}

/* Return location of layer gradient. Input is not in the gradient! */
static const double * outputLayerGradPtr(const vector <size_t> & nneur, const double * g){
    for(size_t lay=1;lay<nneur.size()-1;++lay){
        g += nneur[lay]*(1+nneur[lay-1]);
    }
    return g;
}


void
Mlp::errorVec(const vector<double> &target, double * workspace) const{
    double * o = outputLayerPtr(nneur, workspace);
    for(size_t i=0; i<getNumOutputs(); ++i){
        o[i] -= target[i];
    }
}

vector <double>
Mlp::copyOutputVec(double * workspace) const {
    size_t nout = getNumOutputs();
    const double * o = outputLayerPtr(nneur, workspace);
    vector <double> res(nout);
    for(size_t i=0;i<nout;++i){
        res[i] = o[i];
    }
    return res;
}

size_t
Mlp::outputVecAsClassIndex(const double * workspace) const{
    size_t nout = getNumOutputs();
    const double * o = outputLayerPtr(nneur, workspace);
    size_t ind0 = 0;
    for(size_t i=0;i<nout;++i){
        if (o[i] > o[ind0]) ind0 = i;
    }
    return ind0 + 1;
}


void
Mlp::backwardEucSq(double coeff,
                   double *destE,
                   double *destG,
                   double * workspace,
                   ErrT errortype) const{

    double * o = outputLayerPtr(nneur, workspace);
    double normc=0.;
    // This part is different for different norms:
    if (errortype == errt::q2a2) {
        /* Error term is: coeff * (1/2) * ||e||^2 */
        for(size_t i=0;i<getNumOutputs();++i){normc += o[i] * o[i];}
        *destE += coeff * .5 * normc;
    } else if (errortype == errt::q2a1) {
        /* Error term is: coeff * ||e|| */
        for(size_t i=0;i<getNumOutputs();++i){normc += o[i] * o[i];}
        *destE += coeff * sqrt(normc);
    } else if (errortype == errt::q1a1) {
        /* Error term is: coeff * ||e||_1 */
        for(size_t i=0;i<getNumOutputs();++i){normc += abs(o[i]);}
        *destE += coeff * normc;
    } else {
        throw runtime_error("Unknown errortype");
    }

    // If gradient need not be computed, we're done:
    if (destG==nullptr) return;

    // After this, we start modifying the workspace.

    // This is the last part different between norms:
    if (errortype == errt::q2a2) {
        /* Apply gradient of (1/2) * ||e||^2   */
        /* .. which is ||e||, so nothing is to be done */
    } else if (errortype == errt::q2a1) {
        /* Apply gradient of ||e||. (Not defined at zero.) */
        if (normc>.000001) // FIXME: Tolerance? Smoothing?
            for(size_t i=0;i<getNumOutputs();++i){o[i] /= normc;}
    } else if (errortype == errt::q1a1) {
        /* Apply gradient of ||e||_1. (Not defined at zero.) */
        for(size_t i=0;i<getNumOutputs();++i){o[i] = abs(o[i]);}
    } else {
        throw runtime_error("Unknown errortype");
    }

    /* Start from the last layer. \xi^L in the equations of
     * Kärkkäinen&Heikkola.
     */
    size_t lay=getNLayers()-1;
    if (actf[lay] != actf::linear)
        for(size_t i=0; i<getNumOutputs(); ++i)
            o[i] = actD(actf[lay], o[i]) * o[i];

    // using just ptrs in following. This ok? So far, so good...
    double *g = outputLayerGradPtr(nneur, destG);
    const double *wnext = outputLayerGradPtr(nneur, weights.data());

    // here  G{L} = coeff * [sum(out{L+1},2), out{L+1} * out{L}'];

    // Pointers to indices: this "l", next "l+1", and prev "l-1"
    double * onext = NULL;
    double * othis = o; // Start with l=L
    double * oprev = othis - nneur[nneur.size()-2];

    /* In Matlab, the following goes as (easily as):
     * for L = lastL-1:-1:1
     *   out{L+1} = activD{L}(out{L+1}) .* (W{L+1}(:,2:end)' * out{L+2});
     *   G{L} = (1.0/ndata) * [sum(out{L+1},2), out{L+1} * out{L}'];
     * end
     */

    // Other layers, L-1, L-2, ... 0:
    for(size_t backloop=0; backloop<getNLayers()-1; ++backloop){
        size_t nthis = nneur[lay];
        size_t nprev = nneur[lay-1];

        /* Outer product \xi^{l} (1 o^{l-1}^T) yields derivatives: We
         * add it to the result using the linear coefficient given
         * (could be, for example 1/N):
         */
        double *pg = g;
        for (size_t neur = 0; neur < nthis; ++neur){
            double mul = coeff * othis[neur];
            // For bias, we multiply the extended vector (1 o^(l-1)^T)
            *(pg++) += mul;
            for (size_t conn = 0; conn < nprev; ++conn){
                *(pg++) += mul * oprev[conn];
            }
        }

        // Update indices here. Done, when the input layer is reached.
        lay--;
        if (lay==0) break;

        size_t nnext = nneur[lay+1];
        nthis = nneur[lay];
        nprev = nneur[lay-1];

        oprev -= nneur[lay-1];
        onext = othis;
        othis -= nneur[lay];
        //        onext -= nneur[lay+1];

        /* Apply activation derivative to produce next "\xi^{l}" in
         * workspace:
         *
         * Compute "o^{l} := Diag{(F^l)'(W^l (1 o^{l-1})^T)" =
         * "Diag{(F^l)'(o^{l})" in place.
         */
        for (size_t k = 0; k < nthis; ++k){
            othis[k] = actD(actf[lay],othis[k]);
        }
        // These multiply the elements of W_1^{l+1}^T \xi^{l+1}
        for (size_t i = 0; i < nthis; ++i){
            double r = 0.;
            for (size_t j = 0; j < nnext; ++j) {
                size_t ind = j*nnext + (1+i);
                r += wnext[ind] * onext[j]; // indexing the transpose
            }
            othis[i] = othis[i] * r; // Accounts for the diag multiply.
        }
        wnext -= nneur[lay+1] * (nneur[lay] + 1);
        g     -= nneur[lay] * (nneur[lay-1] + 1);
    }
}

void
Mlp::weightDecaySq(double coeff,
                   double *destE,
                   double *destG,
                   bool excludeOutputBias) const{
    size_t iw=0; // Weights organized in packed rows [bias+connections]
    size_t lastL=getNLayers();
    for (size_t ilay = 1;ilay<getNLayers();++ilay){
        for(size_t ineur = 0;ineur < nneur[ilay];++ineur){
            size_t iconn = 0;
            if ((ilay==lastL) && (excludeOutputBias)) {
                iw++;
                iconn++;
            }
            for(;iconn < (nneur[ilay-1]+1);++iconn){
                (*destE) += .5 * coeff * weights[iw] * weights[iw];
                if (destG != nullptr){
                    destG[iw] += coeff * weights[iw]; // optional deriv
                }
                iw++;
            }
        }
    }
}

void
Mlp::weightDecayAbs(double coeff,
                   double *destE,
                   double *destG,
                   bool excludeOutputBias) const{
    size_t iw=0; // Weights organized in packed rows [bias+connections]
    size_t lastL=getNLayers();
    for (size_t ilay = 1;ilay<getNLayers();++ilay){
        for(size_t ineur = 0;ineur < nneur[ilay];++ineur){
            size_t iconn = 0;
            if ((ilay==lastL) && (excludeOutputBias)) {
                iw++;
                iconn++;
            }
            for(;iconn < (nneur[ilay-1]+1);++iconn){
                double absw = abs(weights[iw]);
                (*destE) += coeff * absw;
                if (destG != nullptr){
                    double signw = 0.0;
                    if (weights[iw]<0.0){
                        signw = -1.0;
                    } else if (weights[iw]>0.0) {
                        signw = 1.0;
                    }
                    destG[iw] += coeff * signw; // optional deriv
                }
                iw++;
            }
        }
    }
}


void
Mlp::update(double coeff, const double *wupd){
    for (size_t iw = 0; iw < getNWeights(); ++iw){
        weights[iw] += coeff * wupd[iw];
    }
}




#if 0
// FIXME: Unnecessary!?
void
Mlp::addErrorAndGradient(const vector<double> &input,
                         const vector<double> &target,
                         double coefficient,
                         double * sqerror,
                         double * error,
                         vector<double> * mseGrad,
                         vector<double> * meeGrad) const{
    cerr << "WARNING: Mlp::addErrorAndGradient() not implemented!" << endl;
}
#endif