// Copyright 2015 National ICT Australia Limited (NICTA)
//
// 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.

#include "MatpowerParserPlugin.h"

#include "Bus.h"
#include "Branch.h"
#include "CommonBranch.h"
#include "Gen.h"
#include "Network.h"
#include "Zip.h"
#include "YamlSupport.h"

#include <cmath>
#include <fstream>
#include <map>
#include <numeric>
#include <boost/spirit/include/qi.hpp>
#include <boost/spirit/include/phoenix_operator.hpp>
#include <boost/spirit/include/phoenix_function.hpp>
#include <boost/spirit/include/phoenix_statement.hpp>
#include <boost/spirit/include/phoenix.hpp>
#include <boost/spirit/include/support_multi_pass.hpp>

namespace qi = boost::spirit::qi;
namespace ascii = boost::spirit::ascii;
namespace phoenix = boost::phoenix;

typedef std::istreambuf_iterator<char> BaseIterator;
typedef boost::spirit::multi_pass<BaseIterator> ForwardIterator;

using namespace Sgt;

namespace
{
    typedef ForwardIterator Iterator;
    typedef decltype(ascii::blank) SpaceType;
    typedef std::vector<std::vector<double>> Matrix;
    typedef Matrix::value_type Row;

    struct MpData
    {
        double MVABase;
        Matrix bus;
        Matrix gen;
        Matrix branch;
        Matrix genCost;
    };

    struct Gram : qi::grammar<Iterator, SpaceType>
    {
        Gram(MpData& data) : Gram::base_type(start_)
        {
            statementTerm_ = (qi::eol | qi::lit(';'));

            eol_ = -(qi::lit('%') >> *(qi::char_-qi::eol)) >> qi::eol;
            ignore_ = -(qi::lit('%') >> *(qi::char_-qi::eol)) >> qi::eol;
            other_ = (*(qi::char_-statementTerm_) >> statementTerm_);

            rowSep_ = (qi::lit(';') >> -eol_) | eol_;
            row_ = +(qi::double_ >> -qi::lit(','));
            matrix_ = "[" >> -eol_ >> (row_ % rowSep_) >> -(eol_ ^ qi::lit(';')) >> "]";

            topFunction_ = qi::lit("function") >> qi::lit("mpc") >> qi::lit("=") >> *(qi::char_-statementTerm_)
                           >> statementTerm_;
            MVABase_ = (qi::lit("mpc.baseMVA") >> qi::lit('=') >> qi::double_ >> statementTerm_)
                       [phoenix::ref(data.MVABase) = qi::_1];
            busMatrix_ = (qi::lit("mpc.bus") >> qi::lit('=') >> matrix_ >> statementTerm_)
                         [phoenix::ref(data.bus) = qi::_1];
            genMatrix_ = (qi::lit("mpc.gen") >> qi::lit('=') >> matrix_ >> statementTerm_)
                         [phoenix::ref(data.gen) = qi::_1];
            branchMatrix_ = (qi::lit("mpc.branch") >> qi::lit('=') >> matrix_ >> statementTerm_)
                            [phoenix::ref(data.branch) = qi::_1];
            genCostMatrix_ = (qi::lit("mpc.gencost") >> qi::lit('=') >> matrix_ >> statementTerm_)
                             [phoenix::ref(data.genCost) = qi::_1];

            start_ = *ignore_ >> topFunction_ >> *(ignore_ | MVABase_ | busMatrix_ | genMatrix_ | branchMatrix_ |
                                                   genCostMatrix_ | other_);

            // To debug: e.g.
            // BOOST_SPIRIT_DEBUG_NODE(busMatrix_); debug(busMatrix_);
        }

        qi::rule<Iterator, SpaceType> statementTerm_;

        qi::rule<Iterator, SpaceType> eol_;
        qi::rule<Iterator, SpaceType> ignore_;
        qi::rule<Iterator, SpaceType> other_;

        // Matrix parser.
        qi::rule<Iterator, SpaceType> rowSep_;
        qi::rule<Iterator, Row(), SpaceType> row_;
        qi::rule<Iterator, Matrix(), SpaceType> matrix_;

        qi::rule<Iterator, SpaceType> topFunction_;
        qi::rule<Iterator, SpaceType> MVABase_;
        qi::rule<Iterator, SpaceType> busMatrix_;
        qi::rule<Iterator, SpaceType> genMatrix_;
        qi::rule<Iterator, SpaceType> branchMatrix_;
        qi::rule<Iterator, SpaceType> genCostMatrix_;

        qi::rule<Iterator, SpaceType> start_;
    };

    // Voltage is in p.u., with the base being specified in KV by the KVBase field.
    template<typename T> T pu2kV(T V, double kVBase)
    {
        return V * kVBase;
    }

    // Shunt admittance is in MW @ 1 p.u. voltage.
    template<typename T> T YBusShunt2Siemens(T Y, double kVBase)
    {
        return Y / (kVBase * kVBase);
    }

    // Branch admittance is in p.u., that is, KVBase^2/MVABase.
    // KVBase is referenced to the "from" bus, the second bus, bus1.
    template<typename T> T YBranch2Siemens(T Y, double kVBase, double MVABase)
    {
        return Y * MVABase / (kVBase * kVBase);
    }

    double deg2Rad(double deg)
    {
        return pi * deg / 180.0;
    }
}

namespace Sgt
{

    struct MpBusInfo
    {
        std::size_t id; // C1, Not nec. consecutive.
        int type; // C2, 1 = PQ, 2 = PV, 3 = SL, 4 = isolated.
        double Pd; // C3, MW
        double Qd; // C4, MVAr
        double Gs; // C5, MW demand @ V = 1 pu
        double Bs; // C6, MVAr demand @ V = 1 pu
        int busArea; // C7
        double VMag; // C8, pu
        double VAngDeg; // C9, pu
        double kVBase; // C10, kV
        int zone; // C11
        double VMagMax; // C12, pu
        double VMagMin; // C13, pu
        double lam_P; // C14, u/pu
        double lam_Q; // C15, u/pu
        double mu_VMax; // C16, u/pu
        double mu_VMin; // C17, u/pu
    };

    struct MpBranchInfo
    {
        std::size_t busIdF;
        std::size_t busIdT;
        double R;
        double X;
        double b;
        double rateA;
        double rateB;
        double rateC;
        double tap;
        double shiftDeg;
        int status;
        double angMinDeg;
        double angMaxDeg;
    };

    struct MpGenInfo
    {
        std::size_t busId;
        double Pg; // MW
        double Qg; // MVAr
        double QMax; // MVAr
        double QMin; // MVAr
        double Vg; // pu
        double MVABase; // MVA
        int status;
        double PMax; // MVA
        double PMin; // MVA
    };

    struct MpGenCostInfo
    {
        int model;
        double startup;
        double shutdown;
        std::vector<double> costs;
    };

    namespace
    {
        std::size_t nDigits(std::size_t n)
        {
            return n == 0 ? 1 : static_cast<std::size_t>(std::floor(std::log10(n)) + 1);
        }

        std::string zeroPadded(std::size_t n, std::size_t nDigits)
        {
            auto result = std::to_string(n);
            size_t len = result.length();
            if (len < nDigits) result.insert(0, nDigits - len, '0');
            return result;
        }

        std::string getBusName(std::size_t id, std::map<std::size_t, std::string>& map, std::size_t nDigits)
        {
            return map.insert(std::make_pair(id, "bus_" + zeroPadded(id, nDigits))).first->second;
        }

        std::string getZipName(std::size_t iZip, const std::string& busName, std::size_t nDigits)
        {
            return "zip_" + zeroPadded(iZip, nDigits) + "_" + busName;
        }

        std::string getGenName(std::size_t iGen, const std::string& busName, std::size_t nDigits)
        {
            return "gen_" + zeroPadded(iGen, nDigits) + "_" + busName;
        }

        std::string getBranchName(std::size_t iBranch, const std::string& busName0, const std::string& busName1,
                std::size_t nDigits)
        {
            return "branch_" + zeroPadded(iBranch, nDigits) + "_" + busName0 + "_" + busName1;
        }

        bool checkNzCplx(Complex x)
        {
            static constexpr double eps = std::numeric_limits<double>::epsilon();
            return std::abs(x.real()) >= eps || std::abs(x.imag()) >= eps;
        }

        bool checkNzZip(Complex SConst, Complex YConst)
        {

            return checkNzCplx(SConst) || checkNzCplx(YConst);
        }
    }

    void MatpowerParserPlugin::parse(const YAML::Node& nd, Network& netw, const ParserBase& parser) const
    {
        assertFieldPresent(nd, "default_kV_base");

        double default_kVBase = parser.expand<double>(nd["default_kV_base"]);

        double VScale = 1.0;
        auto ndVScale = nd["scale_V_by"];
        if (ndVScale)
        {
            VScale = ndVScale.as<double>();
        }
        double PScale = 1.0;
        auto ndPScale = nd["scale_P_by"];
        if (ndPScale)
        {
            PScale = ndPScale.as<double>();
        }
        double RScale = VScale * VScale / PScale;
        double GScale = 1.0 / RScale;

        YAML::Node busNamesNd = nd["bus_names"];
        std::map<size_t, std::string> busNames;
        if (busNamesNd)
        {
            for (auto keyVal : busNamesNd)
            {
                busNames[keyVal.first.as<size_t>()] = keyVal.second.as<std::string>();
            }
        }

        auto ndInputFile = nd["input_file"];
        auto ndEmbeddedInput = nd["embedded_input"];

        sgtAssert(static_cast<bool>(ndInputFile) != static_cast<bool>(ndEmbeddedInput),
                "Parsing matpower data: \"input_file\" or \"embedded_input\" must be present, but not both");

        std::unique_ptr<std::istream> inStream;
        if (ndInputFile)
        {
            std::string inputFName = parser.expand<std::string>(ndInputFile);
            auto inFStream = new std::fstream(inputFName);
            sgtAssert(inFStream->is_open(), "Could not open the matpower input file " << inputFName << ".");
            inStream.reset(inFStream);
            sgtLogMessage() << "Parsing matpower input file " << inputFName << "." << std::endl;
        }
        else
        {
            inStream.reset(new std::stringstream(ndEmbeddedInput.as<std::string>()));
            sgtLogMessage() << "Parsing matpower embedded input." << std::endl;
        }

        // Parse in the raw matpower data.
        MpData data;

        {
            Gram gram(data);

            // Iterate over stream input:
            BaseIterator inBegin(*inStream);

            ForwardIterator fwdBegin = boost::spirit::make_default_multi_pass(inBegin);
            ForwardIterator fwdEnd;

            while (fwdBegin != fwdEnd)
            {
                qi::phrase_parse(fwdBegin, fwdEnd, gram, ascii::blank);
            }
        }

        sgtLogMessage() << "Parsed " << data.bus.size() << " buses, " << data.gen.size() << " generators and "
                        << data.branch.size() << " branches." << std::endl;

        // Extract the bus data.
        std::size_t maxBusId = 0;
        std::vector<MpBusInfo> busVec;
        {
            busVec.reserve(data.bus.size());
            for (const auto& row : data.bus)
            {
                busVec.push_back(MpBusInfo{});
                MpBusInfo& busInfo = busVec.back();

                busInfo.id = static_cast<std::size_t>(row[0]);
                if (maxBusId < busInfo.id) maxBusId = busInfo.id;
                busInfo.type = static_cast<int>(row[1]);
                busInfo.Pd = row[2];
                busInfo.Qd = row[3];
                busInfo.Gs = row[4];
                busInfo.Bs = row[5];
                busInfo.busArea = static_cast<int>(row[6]);
                busInfo.VMag = row[7];
                busInfo.VAngDeg = row[8];
                busInfo.kVBase = row[9] > 1e-6 ? row[9] : default_kVBase;
                busInfo.zone = static_cast<int>(row[10]);
                busInfo.VMagMax = row[11];
                busInfo.VMagMin = row[12];
                if (data.bus[0].size() > 13)
                {
                    busInfo.lam_P = row[13];
                    busInfo.lam_Q = row[14];
                    busInfo.mu_VMax = row[15];
                    busInfo.mu_VMin = row[16];
                }
            }
        }
        std::size_t busNDigits = nDigits(maxBusId);
        std::size_t zipNDigits = busNDigits; // Upper bound.

        // Extract the gen data.
        std::vector<MpGenInfo> genVec;
        {
            genVec.reserve(data.gen.size());
            for (const auto& row : data.gen)
            {
                genVec.push_back(MpGenInfo{});
                MpGenInfo& genInfo = genVec.back();

                genInfo.busId = static_cast<std::size_t>(row[0]);
                genInfo.Pg = row[1];
                genInfo.Qg = row[2];
                genInfo.QMax = row[3];
                genInfo.QMin = row[4];
                genInfo.Vg = row[5];
                genInfo.MVABase = row[6];
                genInfo.status = static_cast<int>(row[7]);
                genInfo.PMax = row[8];
                genInfo.PMin = row[9];
                // There are other fields, but at present we are ignoring them.
            }
        }
        std::size_t genNDigits = nDigits(genVec.size() - 1);

        // Extract the branch data.
        std::vector<MpBranchInfo> branchVec;
        {
            branchVec.reserve(data.branch.size());
            for (const auto& row : data.branch)
            {
                branchVec.push_back(MpBranchInfo{});
                MpBranchInfo& branchInfo = branchVec.back();

                branchInfo.busIdF = static_cast<std::size_t>(row[0]);
                branchInfo.busIdT = static_cast<std::size_t>(row[1]);
                branchInfo.R = row[2];
                branchInfo.X = row[3];
                branchInfo.b = row[4];
                branchInfo.rateA = row[5];
                branchInfo.rateB = row[6];
                branchInfo.rateC = row[7];
                branchInfo.tap = row[8];
                branchInfo.shiftDeg = row[9];
                branchInfo.status = static_cast<int>(row[10]);
                branchInfo.angMinDeg = row[11];
                branchInfo.angMaxDeg = row[12];
            }
        }
        std::size_t branchNDigits = nDigits(branchVec.size() - 1);

        // Extract the genCost data.
        if (data.genCost.size() > 0)
        {
            if (data.genCost.size() == 2 * data.gen.size())
            {
                sgtLogWarning() << "Reactive generator costs not yet implemented. Ignoring them." << std::endl;
            }
            else
            {
                sgtAssert(data.genCost.size() == data.gen.size(),
                        "There are a different number of generators to generator costs.");
            }
        }

        std::vector<MpGenCostInfo> genCostVec;
        {
            genCostVec.reserve(data.gen.size());
            if (data.genCost.size() > 0)
            {
                for (std::size_t i = 0; i < data.gen.size(); ++i)
                {
                    auto row = data.genCost[i];
                    genCostVec.push_back(MpGenCostInfo{});
                    MpGenCostInfo& genCostInfo = genCostVec.back();

                    genCostInfo.model = static_cast<int>(row[0]);
                    genCostInfo.startup = row[1];
                    genCostInfo.shutdown = row[2];
                    auto nCost = static_cast<size_t>(row[3]);
                    genCostInfo.costs.reserve(nCost);
                    for (std::size_t j = 0; j < nCost; ++j)
                    {
                        genCostInfo.costs.push_back(row[4 + j]);
                    }
                }
            }
        }

        // Network:
        netw.setPBase(data.MVABase);

        // Buses:
        std::size_t nZip = 0;
        for (const auto& busInfo : busVec)
        {
            std::string busId = getBusName(busInfo.id, busNames, busNDigits);
            std::unique_ptr<Bus> bus(
                new Bus(busId, {Phase::BAL}, {VScale * Complex(busInfo.kVBase, 0.0)}, PScale * busInfo.kVBase));
            BusType type = BusType::BAD;
            bool isInService = true;
            switch (busInfo.type)
            {
                case 1:
                    type = BusType::PQ;
                    break;
                case 2:
                    type = BusType::PV;
                    break;
                case 3:
                    type = BusType::SL;
                    break;
                case 4:
                    type = BusType::BAD;
                    isInService = false;
                    break;
                default:
                    break;
            }
            bus->setType(type);
            bus->setIsInService(isInService);

            bus->setVMagMin(busInfo.VMagMin <= -infinity
                    ? -infinity
                    : VScale * pu2kV(busInfo.VMagMin, busInfo.kVBase));
            bus->setVMagMax(busInfo.VMagMax >= infinity
                    ? infinity
                    : VScale * pu2kV(busInfo.VMagMax, busInfo.kVBase));

            double VMag = pu2kV(busInfo.VMag, busInfo.kVBase);
            double VAng = deg2Rad(busInfo.VAngDeg);
            bus->setV({VScale * std::polar(VMag, VAng)});

            netw.addBus(std::move(bus));

            Complex SConst = Complex(busInfo.Pd, busInfo.Qd); // Already in MW.
            Complex YConst = YBusShunt2Siemens(Complex(busInfo.Gs, busInfo.Bs), busInfo.kVBase);
            if (checkNzZip(SConst, YConst))
            {
                std::string zipId = getZipName(nZip++, getBusName(busInfo.id, busNames, busNDigits), zipNDigits);
                std::unique_ptr<Zip> zip(new Zip(zipId, {Phase::BAL}));
                zip->setYConst(arma::Mat<Complex>({GScale * YConst}));
                zip->setSConst(arma::Mat<Complex>({PScale * SConst}));
                netw.addZip(std::move(zip), busId);
            }
        } // Buses

        // If there is no slack bus, Matpower assigns the first PV bus as slack. We need to do the same.
        auto it = std::find_if(netw.buses().cbegin(), netw.buses().cend(),
                [](const ConstComponentPtr<Bus>& bus)->bool{return bus->type() == BusType::SL;});
        if (it == netw.buses().cend())
        {
            sgtLogWarning()
                    << "There is no slack bus defined on the network. Setting the type of the first PV bus to SL."
                    << std::endl;
            auto itb = std::find_if(netw.buses().begin(), netw.buses().end(),
                                    [](const ComponentPtr<Bus>& bus)->bool{return bus->type() == BusType::PV;});
            assert(itb != netw.buses().end());
            (**itb).setType(BusType::SL);
        }
        else
        {
            sgtLogMessage() << "The slack bus is " << (**it).id() << std::endl;
        }

        // Gens:
        std::vector<Gen*> genCompVec;
        genCompVec.reserve(genVec.size());
        for (std::size_t i = 0; i < genVec.size(); ++i)
        {
            MpGenInfo& genInfo = genVec[i];
            std::string genId = getGenName(i, getBusName(genInfo.busId, busNames, busNDigits), genNDigits);
            std::unique_ptr<Gen> gen(new Gen(genId, {Phase::BAL}));
            genCompVec.push_back(gen.get());

            std::string busId = busNames.at(genInfo.busId);

            gen->setIsInService(genInfo.status);

            gen->setInServiceS({PScale * Complex(genInfo.Pg, genInfo.Qg)});

            gen->setPMin(genInfo.PMin <= -infinity ? -infinity : PScale * genInfo.PMin);
            gen->setPMax(genInfo.PMax >= infinity ? infinity : PScale * genInfo.PMax);
            gen->setQMin(genInfo.QMin <= -infinity ? -infinity : PScale * genInfo.QMin);
            gen->setQMax(genInfo.QMax >= infinity ? infinity : PScale * genInfo.QMax);

            gen->setC0(0.0);
            gen->setC1(0.0);
            gen->setC2(0.0);

            if (gen->isInService())
            {
                auto bus = netw.buses()[busId];
                bus->setVMagSetpoint({VScale * pu2kV(genInfo.Vg, bus->VBase())});
            }

            netw.addGen(std::move(gen), busId);
        } // Gens

        // Branches:
        for (std::size_t i = 0; i < branchVec.size(); ++i)
        {
            const MpBranchInfo& branchInfo = branchVec[i];

            std::string bus0Name = getBusName(branchInfo.busIdF, busNames, busNDigits);
            std::string bus1Name = getBusName(branchInfo.busIdT, busNames, busNDigits);
            std::string branchName = getBranchName(i, bus0Name, bus1Name, branchNDigits);

            std::unique_ptr<CommonBranch> branch(new CommonBranch(branchName));

            branch->setIsInService(branchInfo.status);

            auto bus0 = netw.buses()[bus0Name];
            auto bus1 = netw.buses()[bus1Name];

            double tap = (std::abs(branchInfo.tap) < 1e-6 ? 1.0 : branchInfo.tap) * bus0->VBase() / bus1->VBase();
            branch->setTapRatio(std::polar(tap, deg2Rad(branchInfo.shiftDeg)));
            branch->setYSeries(GScale * YBranch2Siemens(1.0 / Complex(branchInfo.R, branchInfo.X), bus1->VBase(),
                               data.MVABase));
            branch->setYShunt(GScale * YBranch2Siemens(Complex(0.0, branchInfo.b), bus1->VBase(), data.MVABase));
            branch->setRateA(branchInfo.rateA == 0.0 ? infinity : PScale * branchInfo.rateA);
            branch->setRateB(branchInfo.rateB == 0.0 ? infinity : PScale * branchInfo.rateB);
            branch->setRateC(branchInfo.rateC == 0.0 ? infinity : PScale * branchInfo.rateC);
            branch->setAngMin(branchInfo.angMinDeg * pi / 180.0);
            branch->setAngMax(branchInfo.angMaxDeg * pi / 180.0);

            netw.addBranch(std::move(branch), bus0Name, bus1Name);
        }

        // Add generator costs, if they exist.
        for (std::size_t i = 0; i < genCostVec.size(); ++i)
        {
            const MpGenCostInfo& genCostInfo = genCostVec[i];

            bool isBad = false;

            if (genCostInfo.costs.size() > 3)
            {
                sgtLogWarning() << "Can't have more than three costs for generator. Ignoring costs." << std::endl;
                isBad = true;
            }

            if (genCostInfo.model != 2)
            {
                sgtLogWarning() << "Can only use model 2 for generator costs. Ignoring costs." << std::endl;
                isBad = true;
            }

            if (!isBad)
            {
                Gen& gen = *genCompVec[i];
                gen.setCStartup(genCostInfo.startup);
                gen.setCShutdown(genCostInfo.shutdown);
                auto nCost = genCostInfo.costs.size();
                if (nCost >= 1)
                {
                    gen.setC0(genCostInfo.costs[2]);
                }
                if (nCost >= 2)
                {
                    gen.setC1(genCostInfo.costs[1]);
                }
                if (nCost >= 3)
                {
                    gen.setC2(genCostInfo.costs[0]);
                }
            }
        }
    } // parse(...)
}


//     Copyright (c) 2012 Vadym Kliuchnikov sqct(dot)software(at)gmail(dot)com, Dmitri Maslov, Michele Mosca
//
//     This file is part of SQCT.
//
//     SQCT is free software: you can redistribute it and/or modify
//     it under the terms of the GNU Lesser General Public License as published by
//     the Free Software Foundation, either version 3 of the License, or
//     (at your option) any later version.
//
//     SQCT is distributed in the hope that it will be useful,
//     but WITHOUT ANY WARRANTY; without even the implied warranty of
//     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//     GNU Lesser General Public License for more details.
//
//     You should have received a copy of the GNU Lesser General Public License
//     along with SQCT.  If not, see <http://www.gnu.org/licenses/>.
//

#include <fstream>
#include <algorithm>
#include <cassert>
#include <iostream>
#include <iterator>
#include "epsilonnet.h"
#include "output.h"

#include <boost/range.hpp>
using boost::make_iterator_range;

using namespace std;

/// \brief Header of the file with epsilon net
struct epsilonnetHeader
{
    size_t nodes_count;
};

std::ostream& operator<<(std::ostream& out, const enetNode& node )
{
    out << "enetNode[" << node.ipxx << "," << node.ipQxx << ","
        << node.num_offset << "," << node.compl_offset << "]";
    return out;
}

bool epsilonnet::loadFromFile(const char* filename)
{
    epsilonnetHeader eh;
    ifstream ifs( filename, ios_base::binary );
    if( !ifs )
        return false;
    ifs.read( (char*) &eh,sizeof(eh));
    nodes.clear();
    nodes.resize( eh.nodes_count );
    ifs.read( (char*) &nodes[0], nodes.size() * sizeof(enetNode) );
    numbers.clear();
    numbers.resize( nodes.back().num_offset );
    ifs.read( (char*) &numbers[0], numbers.size() * sizeof(ri) );
    ifs.close();
    denominator_exponent = denominatorExponent2();
    return true;
}

epsilonnet::epsilonnet()
{
    enetNode nd={0,0,0,0};
    nodes.push_back(nd);
}

size_t epsilonnet::nodesCount(const char* filename) const
{
    epsilonnetHeader eh;
    ifstream ifs( filename, ios_base::binary );
    ifs.read( (char*) &eh,sizeof(eh));
    ifs.close();
    return eh.nodes_count;
}

void epsilonnet::saveToFile  (const char* filename) const
{
    if( numbers.size() == 0 )
        return;

    epsilonnetHeader eh;
    eh.nodes_count = nodes.size();
    ofstream ofs( filename, ios_base::binary );
    ofs.write( (const char*) &eh,sizeof(eh));

    assert( nodes.back().ipxx == 0 );
    assert( nodes.back().ipQxx == 0 );
    assert( nodes.back().num_offset == numbers.size() );

    ofs.write( (const char*) &nodes[0], nodes.size() * sizeof(enetNode) );
    ofs.write( (const char*) &numbers[0], numbers.size() * sizeof(ri) );

    ofs.close();
}

/// \brief Comparator based on pointers data
template< class T>
struct eq_ptr
{
    bool operator() ( const T* a , const T* b )
    {
        return *a == *b;
    }
};

/// \brief Modified stub of back_insert_iterator from http://www.cplusplus.com/reference/std/iterator/back_insert_iterator/
template <class Container>
  class back_insert_iterator_ptr :
    public iterator<output_iterator_tag,void,void,void,void>
{
protected:
  Container* container;

public:
  typedef Container container_type;
  explicit back_insert_iterator_ptr (Container& x) : container(&x) {}
  back_insert_iterator_ptr<Container>& operator= (const typename Container::value_type* value)
    { container->push_back(*value); return *this; }
  back_insert_iterator_ptr<Container>& operator* ()
    { return *this; }
  back_insert_iterator_ptr<Container>& operator++ ()
    { return *this; }
  back_insert_iterator_ptr<Container> operator++ (int)
    { return *this; }
};

void epsilonnet::addNode( const pair< ip_type,ip_type>& ip, const nodeRangesPtr &ranges)
{
    nodes.back().ipxx = ip.first;
    nodes.back().ipQxx = ip.second;

    eq_ptr< ri > eq_ri;
    back_insert_iterator_ptr< vector<ri> > biit( numbers );

    int size = numbers.size();
    unique_copy( ranges.nums_begin, ranges.nums_end , biit , eq_ri );
    nodes.back().compl_offset = numbers.size() - size;
    unique_copy( ranges.nums_compl_begin, ranges.nums_compl_end , biit , eq_ri );

    enetNode en = {0,0,numbers.size(),0};
    nodes.push_back(en);
}

void epsilonnet::addNode(epsilonnet::ip_type ipxx, epsilonnet::ip_type ipQxx, const nodeRanges &ranges)
{
    nodes.back().ipxx = ipxx;
    nodes.back().ipQxx = ipQxx;

    auto bi = back_inserter( numbers );
    int size = numbers.size();
    unique_copy( ranges.nums_begin, ranges.nums_end ,bi );
    nodes.back().compl_offset = numbers.size() - size;
    unique_copy( ranges.nums_compl_begin, ranges.nums_compl_end , bi );

    enetNode en = {0,0,numbers.size(),0};
    nodes.push_back(en);
}

void epsilonnet::getNode(size_t node_id, nodeRanges &ranges) const
{
    const enetNode& node = nodes[ node_id ];
    ranges.nums_begin = numbers.begin() + node.num_offset;
    ranges.nums_end = numbers.begin() + complOffset( node_id );
    ranges.nums_compl_begin = ranges.nums_end;
    ranges.nums_compl_end = numbers.begin() + nodes[ node_id + 1 ].num_offset;
}

size_t epsilonnet::complOffset(size_t nodeId) const
{
    const enetNode& node = nodes[ nodeId ];
    return node.num_offset + node.compl_offset;
}

int epsilonnet::denominatorExponent2() const
{
    size_t offset = complOffset(0);
    auto denom = numbers[offset].ipxx() + numbers[0].ipxx();
    int val = ri::gde2( denom );
    // overflow check
    decltype( denom ) d = 1;
    d <<= val;
    assert( d == denom && "Sum of ipxx for number and complementary number is not a power of 2." );
    //denominator_exponent = val;
    return val;
}

int epsilonnet::sde() const
{
    int res = ::sde( 2 * denominatorExponent2(), numbers[0].abs2().gde() );
    if( res == std::numeric_limits<int>::max() ) res = 0;
    return res;
}

double epsilonnet::findExhaustiveApproximation(const epsilonnet::vector2double &vec, epsilonnet::vi &result) const
{
    //denominator_exponent = denominatorExponent2();
    vi current;
    double current_dist = 1.0;
    double best_dist = 1.0;
    vector2double canonical_vec = vec;
    bool conj_1 = false, conj_2 = false;
    int w_pow1 = 0, w_pow2 = 0;
    canonical_vec.first = canonical( vec.first, w_pow1, conj_1 );
    canonical_vec.second = canonical( vec.second, w_pow2, conj_2 );

    for( int i = 0; i < nodes.size(); ++i )
    {
        current_dist = findExhaustiveApproximation( canonical_vec, current, i );
        if( best_dist > current_dist )
        {
            best_dist = current_dist;
            result  = current;
        }
    }

//    cout << result.de << ":" << endl;
//    cout <<  canonical_vec.first << result.d[0].toComplex( result.de ) << endl;
//    cout <<  canonical_vec.second << result.d[1].toComplex( result.de ) << endl;

    if( conj_1 ) result.d[0].conjugate_eq();
    if( conj_2 ) result.d[1].conjugate_eq();
    result.d[0].mul_eq_w( w_pow1 );
    result.d[1].mul_eq_w( w_pow2 );
    result.de = denominator_exponent;

    return best_dist;
}

typedef  epsilonnet::vector2double vd;

/// \brief Inline Euclidean distance computation
inline double dist_squared( const vd& a, const vd& b )
{
    double d1 = a.first.real() - b.first.real();
    double d2 = a.second.real() - b.second.real();
    double d3 = a.first.imag() - b.first.imag();
    double d4 = a.second.imag() - b.second.imag();
    return d1*d1 +d2*d2 +d3*d3 + d4*d4;
}

/// \brief Inline Euclidean distance computation
inline double dist_squared( const vd& a, double fre,  double fim , double sre ,  double sim )
{
    double d1 = a.first.real() - fre;
    double d2 = a.second.real() - sre;
    double d3 = a.first.imag() - fim;
    double d4 = a.second.imag() - sim;
    return d1*d1 +d2*d2 +d3*d3 + d4*d4;
}

/// \todo Remove code duplication by changing one of epsilonnet::findExhaustiveApproximation
double epsilonnet::findExhaustiveApproximation(const epsilonnet::vector2double &vec, epsilonnet::vi &result, int node_id) const
{
    double best = 3.0;
    nodeRanges nr;
    getNode( node_id, nr );
    auto ic = nr.nums_begin;
    auto jc = nr.nums_compl_begin;
    bool tw = true;
    for( auto i = nr.nums_begin; i != nr.nums_end ;++i )
    {
        double iim, ire;
        i->toComplex( denominator_exponent, ire, iim );
        for( auto j = nr.nums_compl_begin; j != nr.nums_compl_end; ++j )
        {
            double jim, jre;
            j->toComplex( denominator_exponent, jre, jim  );
            double d1 = dist_squared( vec, ire, iim, jre, jim );
            double d2 = dist_squared( vec, jre, jim, ire, iim );

            if( d1 < best )
            {
                best = d1;
                ic = i; jc = j;
                tw = false;
            }

            if( d2 < best )
            {
                best = d2;
                ic = i; jc = j;
                tw = true;
            }
        }
    }

    if( tw )
    {
        result.d[0] = *jc;
        result.d[1] = *ic;
        //result.de = denominator_exponent; // -avoid extra operations
    }
    else
    {
        result.d[0] = *ic;
        result.d[1] = *jc;
        //result.de = denominator_exponent; // -avoid extra operations
    }

    return best;
}

double epsilonnet::findExhaustiveApproximation(const epsilonnet::vector2double &vec, epsilonnet::vi &result, int node_id, double &bdist) const
{
    nodeRanges nr;
    getNode( node_id, nr );
    bool found = false;
    auto ic = nr.nums_begin;
    auto jc = nr.nums_compl_begin;
    bool tw = true;
    for( auto i = nr.nums_begin; i != nr.nums_end ;++i )
    {
        double iim, ire;
        i->toComplex( denominator_exponent, ire, iim );
        for( auto j = nr.nums_compl_begin; j != nr.nums_compl_end; ++j )
        {
            double jim, jre;
            j->toComplex( denominator_exponent, jre, jim  );
            double d1 = dist_squared( vec, ire, iim, jre, jim );
            double d2 = dist_squared( vec, jre, jim, ire, iim );

            if( d1 < bdist )
            {
                found = true;
                bdist = d1;
                ic = i; jc = j;
                tw = false;
            }

            if( d2 < bdist )
            {
                found = true;
                bdist = d2;
                ic = i; jc = j;
                tw = true;
            }
        }
    }

    if( found )
    {
        if( tw )
        {
            result.d[0] = *jc;
            result.d[1] = *ic;
            result.de = denominator_exponent;
        }
        else
        {
            result.d[0] = *ic;
            result.d[1] = *jc;
            result.de = denominator_exponent;
        }
    }

    return bdist;
}

//// Code that was used to iterate through epsilon net. Does not required now

//private:
//    /// \brief Executes functor for all complex values in epsilon net
//    template< class T>
//    void for_all_complex( T& functor ) const;
//    /// \brief Executes functor for all complex values in epsilon net within node
//    template< class T>
//    void for_all_complex( T& functor, int node_id ) const;
//    /// \brief Executes functor for all complex values in epsilon net within node with fixed first element
//    template< class T>
//    void for_all_complex( T& functor, nodeRanges& nr, ri& a ) const;
//    /// \brief Executes functor for a, b and elements that are of the form \f$ a\omega^k, a\omega^k \f$
//    template< class T>
//    void for_all_complex( T& functor, std::complex<double>& a, ri& b ) const;

//template< class T>
//void epsilonnet::for_all_complex( T& functor ) const
//{
//    denominator_exponent = denominatorExponent2();
//    for( int i = 0; i < nodes.size() - 1 ; ++i )
//        for_all_complex( functor, i );
//}

//template< class T>
//void epsilonnet::for_all_complex( T& functor, int node_id ) const
//{
//    nodeRanges nr;
//    getNode( node_id, nr );
//    for( auto i : make_iterator_range( nr.nums_begin, nr.nums_end) )
//    {
//        ri a(i);
//        if( i.is_im_eq0() )
//        {
//            for_all_complex( functor, nr, a );
//        }
//        else
//        {
//            for_all_complex( functor, nr, a );
//            a.conjugate_eq();
//            for_all_complex( functor, nr, a );
//        }
//    }
//}

//template< class T>
//void epsilonnet::for_all_complex( T& functor, nodeRanges& nr, ri& a ) const
//{
//    for( int i = 0; i < 8 ; ++i )
//    {
//        a.mul_eq_w();
//        auto first = a.toComplex( denominator_exponent );
//        for( auto j : make_iterator_range(nr.nums_compl_begin, nr.nums_compl_end) )
//        {
//            ri b(j);
//            if( j.is_im_eq0() )
//            {
//                for_all_complex( functor, first, b );
//            }
//            else
//            {
//                for_all_complex( functor, first, b );
//                b.conjugate_eq();
//                for_all_complex( functor, first, b );
//            }
//        }
//    }
//}

//template< class T>
//void epsilonnet::for_all_complex( T& functor, std::complex<double>& first, ri& b ) const
//{
//    int max_ph = 1;
//    if( all_phases ) max_ph = 8;

//    for( int i = 0; i < max_ph ; ++i )
//    {
//        b.mul_eq_w();
//        auto second = b.toComplex( denominator_exponent );
//        functor( first, second );
//        functor( second, first );
//    }
//}


//=======================================================================
// Copyright 2014-2015 David Simmons-Duffin.
// Distributed under the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#include <algorithm>
#include <iostream>
#include <iomanip>
#include <fstream>
#include <ostream>
#include <string>
#include "omp.h"
//Tweak to allow Ubuntu-14.04/gcc-4.8.4 and similar environments to compile
#define BOOST_NO_CXX11_SCOPED_ENUMS
#include <boost/filesystem.hpp>
#undef BOOST_NO_CXX11_SCOPED_ENUMS
#include "boost/filesystem.hpp"
#include "boost/date_time/posix_time/posix_time.hpp"
#include "boost/program_options.hpp"
#include "types.h"
#include "Timers.h"
#include "SDP.h"
#include "parse.h"
#include "SDPSolver.h"

using std::cout;
using std::cerr;
using std::endl;
using std::setfill;
using std::setw;

using boost::filesystem::path;
using boost::posix_time::second_clock;

namespace po = boost::program_options;

Timers timers;

int solveSDP(const vector<path> &sdpFiles,
             const path &outFile,
             const path &checkpointFileIn,
             const path &checkpointFileOut,
             SDPSolverParameters parameters) {
  // Set the default precision of all Real numbers to that specified
  // by the 'precision' parameter.
  mpf_set_default_prec(parameters.precision);

  // Set cout to print the appropriate number of digits
  cout.precision(min(static_cast<int>(parameters.precision * 0.31 + 5), 30));

  // Ensure all the Real parameters have the appropriate precision
  parameters.resetPrecision();

  // Set the maximum number of threads for parallel loops
  omp_set_num_threads(parameters.maxThreads);

  cout << "SDPB started at " << second_clock::local_time() << endl;
  for (auto const& sdpFile: sdpFiles) {
    cout << "SDP file        : " << sdpFile        << endl;
  }
  cout << "out file        : " << outFile        << endl;
  cout << "checkpoint in   : " << checkpointFileIn << endl;
  cout << "checkpoint out  : " << checkpointFileOut << endl;

  cout << "\nParameters:\n";
  cout << parameters << endl;

  // Read an SDP from sdpFile and create a solver for it
  SDPSolver solver(readBootstrapSDP(sdpFiles), parameters);

  if (exists(checkpointFileIn))
    solver.loadCheckpoint(checkpointFileIn);

  timers["Solver runtime"].start();
  timers["Last checkpoint"].start();
  SDPSolverTerminateReason reason = solver.run(checkpointFileOut);
  timers["Solver runtime"].stop();

  cout << "-----" << setfill('-') << setw(116) << std::left << reason << endl;
  cout << endl;
  cout << "primalObjective = " << solver.primalObjective << endl;
  cout << "dualObjective   = " << solver.dualObjective   << endl;
  cout << "dualityGap      = " << solver.dualityGap      << endl;
  cout << "primalError     = " << solver.primalError     << endl;
  cout << "dualError       = " << solver.dualError       << endl;
  cout << endl;

  if (!parameters.noFinalCheckpoint)
    solver.saveCheckpoint(checkpointFileOut);
  timers["Last checkpoint"].stop();
  solver.saveSolution(reason, outFile);

  cout << endl << timers;

  return 0;
}

int main(int argc, char** argv) {
  vector<path> sdpFiles;
  path outFile;
  path checkpointFileIn;
  path checkpointFileOut;
  path paramFile;

  SDPSolverParameters parameters;

  po::options_description basicOptions("Basic options");
  basicOptions.add_options()
    ("help,h", "Show this helpful message.")
    ("sdpFile,s",
     po::value< vector<path> >(&sdpFiles)->required(),
     "SDP data file(s) in XML format. Use this option repeatedly to specify multiple data files.")
    ("paramFile,p",
     po::value<path>(&paramFile),
     "Any parameter can optionally be set via this file in key=value "
     "format. Command line arguments override values in the parameter "
     "file.")
    ("outFile,o",
     po::value<path>(&outFile),
     "The optimal solution is saved to this file in Mathematica "
     "format. Defaults to sdpFile with '.out' extension.")
    ("checkpointFile,c",
     po::value<path>(&checkpointFileOut),
     "Checkpoints are saved to this file every checkpointInterval. Defaults "
     "to sdpFile with '.ck' extension.")
    ("initialCheckpointFile,i",
     po::value<path>(&checkpointFileIn),
     "The initial checkpoint to load. Defaults to checkpointFile.")
    ;

  po::options_description solverParamsOptions("Solver parameters");
  solverParamsOptions.add_options()
    ("precision",
     po::value<int>(&parameters.precision)->default_value(400),
     "Precision in binary digits.  GMP will round up to the nearest "
     "multiple of 64 (or 32 on older systems).")
    ("maxThreads",
     po::value<int>(&parameters.maxThreads)->default_value(4),
     "Maximum number of threads to use for parallel calculation.")
    ("checkpointInterval",
     po::value<int>(&parameters.checkpointInterval)->default_value(3600),
     "Save checkpoints to checkpointFile every checkpointInterval seconds.")
    ("noFinalCheckpoint",
     po::bool_switch(&parameters.noFinalCheckpoint)->default_value(false),
     "Don't save a final checkpoint after terminating (useful when debugging).")
    ("findPrimalFeasible",
     po::bool_switch(&parameters.findPrimalFeasible)->default_value(false),
     "Terminate once a primal feasible solution is found.")
    ("findDualFeasible",
     po::bool_switch(&parameters.findDualFeasible)->default_value(false),
     "Terminate once a dual feasible solution is found.")
    ("detectPrimalFeasibleJump",
     po::bool_switch(&parameters.detectPrimalFeasibleJump)->default_value(false),
     "Terminate if a primal-step of 1 is taken. This often indicates that a "
     "primal feasible solution would be found if the precision were high "
     "enough. Try increasing either primalErrorThreshold or precision "
     "and run from the latest checkpoint.")
    ("detectDualFeasibleJump",
     po::bool_switch(&parameters.detectDualFeasibleJump)->default_value(false),
     "Terminate if a dual-step of 1 is taken. This often indicates that a "
     "dual feasible solution would be found if the precision were high "
     "enough. Try increasing either dualErrorThreshold or precision "
     "and run from the latest checkpoint.")
    ("maxIterations",
     po::value<int>(&parameters.maxIterations)->default_value(500),
     "Maximum number of iterations to run the solver.")
    ("maxRuntime",
     po::value<int>(&parameters.maxRuntime)->default_value(86400),
     "Maximum amount of time to run the solver in seconds.")
    ("dualityGapThreshold",
     po::value<Real>(&parameters.dualityGapThreshold)->default_value(Real("1e-30")),
     "Threshold for duality gap (roughly the difference in primal and dual "
     "objective) at which the solution is considered "
     "optimal. Corresponds to SDPA's epsilonStar.")
    ("primalErrorThreshold",
     po::value<Real>(&parameters.primalErrorThreshold)->default_value(Real("1e-30")),
     "Threshold for feasibility of the primal problem. Corresponds to SDPA's "
     "epsilonBar.")
    ("dualErrorThreshold",
     po::value<Real>(&parameters.dualErrorThreshold)->default_value(Real("1e-30")),
     "Threshold for feasibility of the dual problem. Corresponds to SDPA's epsilonBar.")
    ("initialMatrixScalePrimal",
     po::value<Real>(&parameters.initialMatrixScalePrimal)->default_value(Real("1e20")),
     "The primal matrix X begins at initialMatrixScalePrimal times the "
     "identity matrix. Corresponds to SDPA's lambdaStar.")
    ("initialMatrixScaleDual",
     po::value<Real>(&parameters.initialMatrixScaleDual)->default_value(Real("1e20")),
     "The dual matrix Y begins at initialMatrixScaleDual times the "
     "identity matrix. Corresponds to SDPA's lambdaStar.")
    ("feasibleCenteringParameter",
     po::value<Real>(&parameters.feasibleCenteringParameter)->default_value(Real("0.1")),
     "Shrink the complementarity X Y by this factor when the primal and dual "
     "problems are feasible. Corresponds to SDPA's betaStar.")
    ("infeasibleCenteringParameter",
     po::value<Real>(&parameters.infeasibleCenteringParameter)->default_value(Real("0.3")),
     "Shrink the complementarity X Y by this factor when either the primal "
     "or dual problems are infeasible. Corresponds to SDPA's betaBar.")
    ("stepLengthReduction",
     po::value<Real>(&parameters.stepLengthReduction)->default_value(Real("0.7")),
     "Shrink each newton step by this factor (smaller means slower, more "
     "stable convergence). Corresponds to SDPA's gammaStar.")
    ("choleskyStabilizeThreshold",
     po::value<Real>(&parameters.choleskyStabilizeThreshold)->default_value(Real("1e-40")),
     "Adds stabilizing terms to the cholesky decomposition of the schur complement "
     "matrix for diagonal entries which are smaller than this threshold times the "
     "geometric mean of other diagonal entries. Somewhat higher choleskyStabilizeThreshold "
     "can improve numerical stability but if the threshold is large enough that a high "
     "proportion of eigenvalues are being stabilized, the computation will slow substantially.")
    ("maxComplementarity",
     po::value<Real>(&parameters.maxComplementarity)->default_value(Real("1e100")),
     "Terminate if the complementarity mu = Tr(X Y)/dim(X) exceeds this value.")
    ;

  po::options_description cmdLineOptions;
  cmdLineOptions.add(basicOptions).add(solverParamsOptions);

  po::variables_map variablesMap;

  try {
    po::store(po::parse_command_line(argc, argv, cmdLineOptions), variablesMap);

    if (variablesMap.count("help")) {
      cout << cmdLineOptions << endl;
      return 0;
    }

    if (variablesMap.count("paramFile")) {
      paramFile = variablesMap["paramFile"].as<path>();
      std::ifstream ifs(paramFile.string().c_str());
      po::store(po::parse_config_file(ifs, solverParamsOptions), variablesMap);
    }

    po::notify(variablesMap);

    if (!variablesMap.count("outFile")) {
      outFile = sdpFiles[0];
      outFile.replace_extension("out");
    }

    if (!variablesMap.count("checkpointFile")) {
      checkpointFileOut = sdpFiles[0];
      checkpointFileOut.replace_extension("ck");
    }

    if (!variablesMap.count("initialCheckpointFile")) {
      checkpointFileIn = checkpointFileOut;
    }

    std::ofstream ofs(outFile.string().c_str());
    ofs.close();
    if (!ofs) {
      cerr << "Cannot write to outFile." << endl;
      return 1;
    }
  } catch(po::error& e) {
    cerr << "ERROR: " << e.what() << endl;
    cerr << cmdLineOptions << endl;
    return 1;
  }

  return solveSDP(sdpFiles, outFile, checkpointFileIn, checkpointFileOut, parameters);
}


#include "lir.hpp"

#include <iostream>
#include <unordered_map>
#include <unordered_set>

#include <boost/range/adaptor/reversed.hpp>

namespace algrad {
namespace compiler {

using Live_set = std::unordered_set<unsigned>;

Live_set
get_live_out(std::vector<Live_set>& liveIn, lir::Program const& program, lir::Block& bb)
{
    Live_set ret;
    for (int logical = 0; logical < 2; ++logical) {
        for (auto succ : (logical ? bb.logicalSuccessors() : bb.linearizedSuccessors())) {
            for (auto e : liveIn[succ->id()])
                if ((logical && program.temp_info(e).reg_class == lir::RegClass::vgpr) ||
                    (!logical && program.temp_info(e).reg_class != lir::RegClass::vgpr))
                    ret.insert(e);
            int index = -1;
            unsigned i = 0;
            for (auto pred : (logical ? succ->logicalPredecessors() : succ->linearizedPredecessors())) {
                if (pred == &bb)
                    index = i;
                ++i;
            }

            for (auto& insn : succ->instructions()) {
                if (insn->opCode() != lir::OpCode::phi)
                    break;
                auto reg_class = program.temp_info(insn->getDefinition(0).temp()).reg_class;
                if ((logical && reg_class == lir::RegClass::vgpr) || (!logical && reg_class != lir::RegClass::vgpr))
                    ret.insert(insn->getOperand(index).temp());
            }
        }
    }
    return ret;
}
std::vector<Live_set>
compute_live_in(lir::Program& program)
{
    std::vector<Live_set> live_in(program.blocks().size());

    for (;;) {
        bool changed = false;
        for (auto& bb : boost::adaptors::reverse(program.blocks())) {
            auto live = get_live_out(live_in, program, *bb);
            for (auto& insn : boost::adaptors::reverse(bb->instructions())) {
                auto defCount = insn->definitionCount();
                for (unsigned i = 0; i < defCount; ++i) {
                    auto it = live.find(insn->getDefinition(i).temp());
                    if (it != live.end())
                        live.erase(it);
                }

                if (insn->opCode() != lir::OpCode::phi) {
                    auto op_count = insn->operandCount();
                    for (unsigned i = 0; i < op_count; ++i) {
                        if (insn->getOperand(i).is_temp()) {
                            insn->getOperand(i).setKill(live.find(insn->getOperand(i).temp()) == live.end());
                        }
                    }
                    for (unsigned i = 0; i < op_count; ++i) {
                        if (insn->getOperand(i).is_temp()) {
                            if (live.find(insn->getOperand(i).temp()) == live.end())
                                live.insert(insn->getOperand(i).temp());
                        }
                    }
                }
            }
            if (live != live_in[bb->id()]) {
                live_in[bb->id()] = live;
                changed = true;
            }
        }
        if (!changed)
            break;
    }
    return live_in;
}

void
insert_copies(lir::Program& program)
{
    auto live_in = compute_live_in(program);
    for (auto& bb : program.blocks()) {
        std::vector<std::unique_ptr<lir::Inst>> instructions;
        instructions.reserve(bb->instructions().size());
        auto live = get_live_out(live_in, program, *bb);

        for (auto it = bb->instructions().end(); it != bb->instructions().begin();) {
            --it;
            auto& insn = *it;
            bool need_move = false;
            auto def_count = insn->definitionCount();
            for (std::size_t i = 0; i < def_count; ++i) {
                auto const& def = insn->getDefinition(i);
                if (def.is_temp() && def.isFixed())
                    need_move = true;
                auto it = live.find(def.temp());
                if (it != live.end())
                    live.erase(it);
            }

            auto op_count = insn->operandCount();
            for (std::size_t i = 0; i < op_count; ++i) {
                auto arg = insn->getOperand(i);
                if (arg.is_temp()) {
                    if (arg.isFixed())
                        need_move = true;
                    live.insert(arg.temp());
                }
            }

            instructions.push_back(std::move(insn));
            if (need_move && !live.empty()) {
                auto copy = std::make_unique<lir::Inst>(lir::OpCode::parallel_copy, live.size(), live.size());
                unsigned idx = 0;
                for (auto e : live) {
                    copy->getOperand(idx) = lir::Arg{e};
                    copy->getDefinition(idx) = lir::Arg{e};
                    ++idx;
                }
                instructions.push_back(std::move(copy));
            }
        }

        std::reverse(instructions.begin(), instructions.end());
        bb->instructions() = std::move(instructions);
    }
}

void
fix_ssa_rename_visit(lir::Program& program, lir::Block& block, std::vector<bool>& visited,
                     std::vector<unsigned>& renames, std::vector<std::pair<unsigned, unsigned>>& undo, bool logical)
{
    if (visited[block.id()])
        return;
    visited[block.id()] = true;

    auto undo_size = undo.size();
    for (auto& insn : block.instructions()) {
        if (insn->opCode() != lir::OpCode::phi) {
            auto op_count = insn->operandCount();
            for (std::size_t i = 0; i < op_count; ++i) {
                auto id = insn->getOperand(i).temp();
                if ((logical && program.temp_info(id).reg_class != lir::RegClass::vgpr) ||
                    (!logical && program.temp_info(id).reg_class == lir::RegClass::vgpr))
                    continue;
                if (renames[id] == ~0U)
                    std::terminate();

                insn->getOperand(i).set_temp(renames[id]);
            }
        }

        auto def_count = insn->definitionCount();
        for (std::size_t i = 0; i < def_count; ++i) {
            auto id = insn->getDefinition(i).temp();
            if ((logical && program.temp_info(id).reg_class != lir::RegClass::vgpr) ||
                (!logical && program.temp_info(id).reg_class == lir::RegClass::vgpr))
                continue;
            undo.push_back({id, renames[id]});
            if (renames[id] == ~0U) {
                renames[id] = id;
            } else {
                auto new_temp = program.allocate_temp(program.temp_info(id).reg_class, program.temp_info(id).size);
                renames[id] = new_temp;
                insn->getDefinition(i).set_temp(new_temp);
            }
        }
    }

    for (auto succ : (logical ? block.logicalSuccessors() : block.linearizedSuccessors())) {
        auto index = lir::findBlock(logical ? succ->logicalPredecessors() : succ->linearizedPredecessors(), &block);
        for (auto& inst : succ->instructions()) {
            if (inst->opCode() != lir::OpCode::phi)
                break;
            if ((logical && program.temp_info(inst->getDefinition(0).temp()).reg_class != lir::RegClass::vgpr) ||
                (!logical && program.temp_info(inst->getDefinition(0).temp()).reg_class == lir::RegClass::vgpr))
                continue;

            auto id = inst->getOperand(index).temp();
            if (renames.at(id) == ~0U)
                std::terminate();

            inst->getOperand(index).set_temp(renames[id]);
        }
    }

    if (logical) {
        for (auto succ : block.logicalSuccessors())
            fix_ssa_rename_visit(program, *succ, visited, renames, undo, logical);
    } else {
        for (auto succ : block.linearizedSuccessors())
            fix_ssa_rename_visit(program, *succ, visited, renames, undo, logical);
    }

    for (std::size_t i = undo.size(); i > undo_size; --i) {
        renames[undo[i - 1].first] = undo[i - 1].second;
    }

    undo.resize(undo_size);
}

void
fix_ssa_rename(lir::Program& program)
{
    std::vector<bool> visited(program.blocks().size());
    std::vector<unsigned> renames(program.allocated_temp_count(), ~0U);
    std::vector<std::pair<unsigned, unsigned>> undo;
    fix_ssa_rename_visit(program, *program.blocks()[0], visited, renames, undo, false);
    std::fill(visited.begin(), visited.end(), false);
    std::fill(renames.begin(), renames.end(), ~0U);
    fix_ssa_rename_visit(program, *program.blocks()[0], visited, renames, undo, true);
}
void
fix_ssa(lir::Program& program)
{
    fix_ssa_rename(program);
}
bool
allowed(std::vector<bool> const& forbidden, unsigned index, unsigned size)
{
    for (unsigned i = 0; i < size; ++i)
        if (forbidden[i + index])
            return false;
    return true;
}

void
set_forbidden(std::vector<bool>& forbidden, unsigned index, unsigned size, bool value)
{
    for (unsigned i = 0; i < size; ++i)
        forbidden[i + index] = value;
}

std::vector<int>
color_registers(lir::Program& program)
{
    std::vector<std::unordered_set<unsigned>> ig(program.allocated_temp_count());
    std::vector<int> colors(program.allocated_temp_count(), -1);
    auto live_in = compute_live_in(program);

    for (auto& bb : program.blocks()) {
        std::vector<bool> colors_used(2048);
        for (auto e : live_in[bb->id()]) {
            auto size = program.temp_info(e).size;
            for (std::size_t i = 0; i < size; ++i)
                colors_used[colors[e] + i] = true;
        }
        for (auto it = bb->instructions().begin(); it != bb->instructions().end(); ++it) {
            if ((*it)->opCode() != lir::OpCode::phi) {
                auto op_count = (*it)->operandCount();
                for (std::size_t i = 0; i < op_count; ++i) {
                    auto& arg = (*it)->getOperand(i);
                    if (arg.is_temp()) {
                        if (arg.kill()) {
                            set_forbidden(colors_used, colors[arg.temp()], program.temp_info(arg.temp()).size, false);
                        }
                        arg.setFixed(lir::PhysReg{static_cast<unsigned>(colors[arg.temp()])});
                    }
                }
            }

            auto def_count = (*it)->definitionCount();
            for (std::size_t i = 0; i < def_count; ++i) {
                auto& def = (*it)->getDefinition(i);
                if (colors[def.temp()] < 0) {
                    std::vector<bool> forbidden = colors_used;
                    int c = -1;
                    if (def.isFixed()) {
                        c = def.physReg().reg;
                    }
                    if (it + 1 != bb->instructions().end() && (*it)->opCode() == lir::OpCode::parallel_copy) {
                        auto op_count_2 = it[1]->operandCount();
                        for (unsigned j = 0; j < op_count_2; ++j) {
                            auto const& arg2 = it[1]->getOperand(j);
                            if (arg2.isFixed()) {
                                if (def.temp() != arg2.temp())
                                    set_forbidden(forbidden, arg2.physReg().reg, program.temp_info(arg2.temp()).size,
                                                  true);
                                else
                                    c = arg2.physReg().reg;
                            }
                        }
                    }
                    if (c == -1 && (*it)->opCode() == lir::OpCode::parallel_copy) {
                        auto& prev_arg = (*it)->getOperand(i);
                        if (prev_arg.temp() && !prev_arg.kill()) {
                            std::cerr << prev_arg.temp() << " has no kill\n";
                        }
                        if (allowed(forbidden, prev_arg.physReg().reg, program.temp_info(prev_arg.temp()).size)) {
                            c = prev_arg.physReg().reg;
                        } else
                            std::cerr << "preferred move failed for " << def.temp() << " " << prev_arg.temp() << "\n";
                    }
                    if (c == -1 && (*it)->opCode() == lir::OpCode::phi) {
                        int candidate = -1;
                        auto op_count = (*it)->operandCount();
                        for (std::size_t j = 0; j < op_count; ++j) {
                            if ((*it)->getOperand(j).is_temp() && colors[(*it)->getOperand(j).temp()] >= 0)
                                candidate = colors[(*it)->getOperand(j).temp()];
                        }
                        if (candidate >= 0 && allowed(forbidden, candidate, program.temp_info(def.temp()).size)) {
                            c = candidate;
                        }
                    }

                    if (c == -1) {
                        c = program.temp_info(def.temp()).reg_class == lir::RegClass::vgpr ? 1024 : 0;
                        while (!allowed(forbidden, c, program.temp_info(def.temp()).size))
                            c += program.temp_info(def.temp()).size;
                    }

                    set_forbidden(colors_used, c, program.temp_info(def.temp()).size, true);
                    colors[def.temp()] = c;
                }
                def.setFixed(lir::PhysReg{static_cast<unsigned>(colors[def.temp()])});
            }
        }
    }
    for (auto& bb : program.blocks()) {
        for (auto it = bb->instructions().begin();
             it != bb->instructions().end() && (*it)->opCode() == lir::OpCode::phi; ++it) {
            auto op_count = (*it)->operandCount();
            for (std::size_t i = 0; i < op_count; ++i) {
                auto& arg = (*it)->getOperand(i);
                if (arg.is_temp()) {
                    arg.setFixed(lir::PhysReg{static_cast<unsigned>(colors[arg.temp()])});
                }
            }
        }
    }
    return colors;
}

bool
has_logical_phis(lir::Program& program, lir::Block& block)
{
    for (auto it = block.instructions().begin();
         it != block.instructions().end() && (*it)->opCode() == lir::OpCode::phi; ++it)
        if (program.temp_info((*it)->getDefinition(0).temp()).reg_class == lir::RegClass::vgpr)
            return true;
    return false;
}

bool
has_linearized_phis(lir::Program& program, lir::Block& block)
{
    for (auto it = block.instructions().begin();
         it != block.instructions().end() && (*it)->opCode() == lir::OpCode::phi; ++it)
        if (program.temp_info((*it)->getDefinition(0).temp()).reg_class != lir::RegClass::vgpr)
            return true;
    return false;
}

void
destroy_phis(lir::Program& program)
{
    for (auto& bb : program.blocks()) {
        for (auto succ : bb->linearizedSuccessors())
            if (has_linearized_phis(program, *succ))
                std::terminate();

        std::vector<std::pair<lir::Arg, lir::Arg>> args;
        for (auto succ : bb->logicalSuccessors()) {
            auto index = lir::findOrInsertBlock(succ->logicalPredecessors(), bb.get());
            for (auto it = succ->instructions().begin();
                 it != succ->instructions().end() && (*it)->opCode() == lir::OpCode::phi; ++it) {
                if (program.temp_info((*it)->getDefinition(0).temp()).reg_class != lir::RegClass::vgpr)
                    continue;
                // if ((*it)->getOperand(index).physReg().reg == (*it)->getDefinition(0).physReg().reg)
                //    continue;
                args.emplace_back((*it)->getOperand(index), (*it)->getDefinition(0));
            }
        }
        if (args.empty())
            continue;

        auto inst = std::make_unique<lir::Inst>(lir::OpCode::parallel_copy, args.size(), args.size());
        for (std::size_t i = 0; i < args.size(); ++i) {
            inst->getOperand(i) = args[i].first;
            inst->getDefinition(i) = args[i].second;
        }

        auto it = --bb->instructions().end();
        bb->instructions().insert(it, std::move(inst));
    }

    for (auto& bb : program.blocks()) {
        auto it = bb->instructions().begin();
        while (it != bb->instructions().end() && (*it)->opCode() == lir::OpCode::phi)
            ++it;
        bb->instructions().erase(bb->instructions().begin(), it);
    }
}

void
allocateRegisters(lir::Program& program)
{
    insert_copies(program);
    fix_ssa(program);
    color_registers(program);
    destroy_phis(program);
}
}
}


// graph-tool -- a general graph modification and manipulation thingy
//
// Copyright (C) 2006-2018 Tiago de Paula Peixoto <tiago@skewed.de>
//
// This program is free software; you can redistribute it and/or
// modify it under the terms of the GNU General Public License
// as published by the Free Software Foundation; either version 3
// of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

#include "graph_filtering.hh"
#include "graph_python_interface.hh"

#include <boost/python.hpp>
#include <boost/graph/astar_search.hpp>

#include "graph.hh"
#include "graph_selectors.hh"
#include "graph_util.hh"

#include "graph_astar.hh"

#include "coroutine.hh"
#include "graph_python_interface.hh"

using namespace std;
using namespace boost;
using namespace graph_tool;

template <class T>
python::object operator |(const python::object& a, const T& b)
{
    return a / b;
}

struct do_astar_search
{
    template <class Graph, class DistanceMap, class PredMap, class Visitor>
    void operator()(Graph& g, size_t s, DistanceMap dist, PredMap pred,
                    boost::any aweight, Visitor vis, pair<AStarCmp, AStarCmb> cmp,
                    pair<python::object, python::object> range,
                    python::object h, GraphInterface& gi) const
    {
        typedef typename graph_traits<Graph>::edge_descriptor edge_t;
        typedef typename property_traits<DistanceMap>::value_type dtype_t;
        dtype_t z = python::extract<dtype_t>(range.first);
        dtype_t i = python::extract<dtype_t>(range.second);
        checked_vector_property_map<default_color_type,
                                    decltype(get(vertex_index, g))>
            color(get(vertex_index, g));
        checked_vector_property_map<dtype_t,
                                    decltype(get(vertex_index, g))>
            cost(get(vertex_index, g));
        DynamicPropertyMapWrap<dtype_t, edge_t> weight(aweight,
                                                       edge_properties());
        astar_search(g, vertex(s, g), AStarH<Graph, dtype_t>(gi, g, h),
                     vis, pred, cost, dist, weight, get(vertex_index, g), color,
                     cmp.first, cmp.second, i, z);
   }
};

struct do_astar_search_fast
{
    template <class Graph, class DistanceMap, class WeightMap,
              class Visitor>
    void operator()(Graph& g, size_t s, DistanceMap dist,
                    WeightMap weight, Visitor vis,
                    pair<python::object, python::object> range,
                    python::object h, GraphInterface& gi) const
    {
        typedef typename property_traits<DistanceMap>::value_type dtype_t;
        dtype_t z = python::extract<dtype_t>(range.first);
        dtype_t i = python::extract<dtype_t>(range.second);
        astar_search(g, vertex(s, g), AStarH<Graph, dtype_t>(gi, g, h),
                     weight_map(weight).distance_map(dist).distance_zero(z).
                     distance_inf(i).visitor(vis));
   }
};

void a_star_search(GraphInterface& g, size_t source, boost::any dist_map,
                   boost::any pred_map, boost::any weight, python::object vis,
                   python::object cmp, python::object cmb, python::object zero,
                   python::object inf, python::object h)
{
    typedef typename property_map_type::
        apply<int64_t, GraphInterface::vertex_index_map_t>::type pred_t;
    pred_t pred = any_cast<pred_t>(pred_map);
    run_action<graph_tool::all_graph_views,mpl::true_>()
        (g, std::bind(do_astar_search(),  std::placeholders::_1, source,
                      std::placeholders::_2, pred, weight,
                      AStarVisitorWrapper(g, vis), make_pair(AStarCmp(cmp),
                                                             AStarCmb(cmb)),
                      make_pair(zero, inf), h, std::ref(g)),
         writable_vertex_properties())(dist_map);
}

#ifdef HAVE_BOOST_COROUTINE

class AStarGeneratorVisitor : public astar_visitor<>
{
public:
    AStarGeneratorVisitor(GraphInterface& gi,
                          coro_t::push_type& yield)
        : _gi(gi), _yield(yield) {}

    template <class Edge, class Graph>
    void edge_relaxed(const Edge& e, Graph& g)
    {
        auto gp = retrieve_graph_view<Graph>(_gi, g);
        _yield(boost::python::object(PythonEdge<Graph>(gp, e)));
    }

private:
    GraphInterface& _gi;
    coro_t::push_type& _yield;
};

#endif // HAVE_BOOST_COROUTINE

boost::python::object astar_search_generator(GraphInterface& g,
                                             size_t source,
                                             boost::any dist_map,
                                             boost::any weight,
                                             python::object cmp,
                                             python::object cmb,
                                             python::object zero,
                                             python::object inf,
                                             python::object h)
{
#ifdef HAVE_BOOST_COROUTINE
    auto dispatch = [&](auto& yield)
        {
            AStarGeneratorVisitor vis(g, yield);
            run_action<graph_tool::all_graph_views,mpl::true_>()
               (g, std::bind(do_astar_search(),  std::placeholders::_1, source,
                             std::placeholders::_2, dummy_property_map(), weight,
                             vis, make_pair(AStarCmp(cmp), AStarCmb(cmb)),
                             make_pair(zero, inf), h, std::ref(g)),
                writable_vertex_properties())(dist_map);
        };
    return boost::python::object(CoroGenerator(dispatch));
#else
    throw GraphException("This functionality is not available because boost::coroutine was not found at compile-time");
#endif
}

boost::python::object astar_search_generator_fast(GraphInterface& g,
                                                  size_t source,
                                                  boost::any dist_map,
                                                  boost::any weight,
                                                  python::object zero,
                                                  python::object inf,
                                                  python::object h)
{
#ifdef HAVE_BOOST_COROUTINE
    auto dispatch = [&](auto& yield)
        {
            AStarGeneratorVisitor vis(g, yield);
            run_action<graph_tool::all_graph_views,mpl::true_>()
               (g, std::bind(do_astar_search_fast(),  std::placeholders::_1, source,
                             std::placeholders::_2, std::placeholders::_3,
                             vis, make_pair(zero, inf), h, std::ref(g)),
                writable_vertex_scalar_properties(),
                edge_scalar_properties())(dist_map, weight);
        };
    return boost::python::object(CoroGenerator(dispatch));
#else
    throw GraphException("This functionality is not available because boost::coroutine was not found at compile-time");
#endif
}

class AStarArrayVisitor: public astar_visitor<>
{
public:
    AStarArrayVisitor(std::vector<std::array<size_t, 2>>& edges)
        : _edges(edges) {}

    template <class Edge, class Graph>
    void edge_relaxed(const Edge& e, Graph& g)
    {
        _edges.push_back({{source(e, g), target(e,g)}});
    }

private:
    std::vector<std::array<size_t, 2>>& _edges;
};

boost::python::object astar_search_array(GraphInterface& g,
                                         size_t source,
                                         boost::any dist_map,
                                         boost::any weight,
                                         python::object cmp,
                                         python::object cmb,
                                         python::object zero,
                                         python::object inf,
                                         python::object h)
{
    std::vector<std::array<size_t, 2>> edges;
    AStarArrayVisitor vis(edges);
    run_action<graph_tool::all_graph_views,mpl::true_>()
        (g, std::bind(do_astar_search(),  std::placeholders::_1, source,
                      std::placeholders::_2, dummy_property_map(), weight,
                      vis, make_pair(AStarCmp(cmp), AStarCmb(cmb)),
                      make_pair(zero, inf), h, std::ref(g)),
         writable_vertex_properties())(dist_map);
    return wrap_vector_owned<size_t,2>(edges);
}

boost::python::object astar_search_array_fast(GraphInterface& g,
                                              size_t source,
                                              boost::any dist_map,
                                              boost::any weight,
                                              python::object zero,
                                              python::object inf,
                                              python::object h)
{
    std::vector<std::array<size_t, 2>> edges;
    AStarArrayVisitor vis(edges);
    run_action<graph_tool::all_graph_views,mpl::true_>()
        (g, std::bind(do_astar_search_fast(),  std::placeholders::_1, source,
                      std::placeholders::_2, std::placeholders::_3,
                      vis, make_pair(zero, inf), h, std::ref(g)),
         writable_vertex_scalar_properties(),
         edge_scalar_properties())(dist_map, weight);
    return wrap_vector_owned<size_t,2>(edges);
}

void export_astar()
{
    using namespace boost::python;
    def("astar_search", &a_star_search);
    def("astar_generator", &astar_search_generator);
    def("astar_generator_fast", &astar_search_generator_fast);
    def("astar_array", &astar_search_array);
    def("astar_array_fast", &astar_search_array_fast);
}


#include <mpi.h>
#include <stdio.h>
#include <Eigen/Dense>
#include <Eigen/Sparse>

#include <iostream>
#include <fstream>

#include <fstream>
#include <string>
#include <algorithm>
#include <random>
#include <chrono>
#include <memory>
#include <cmath>
#include <stdlib.h>

#include <getopt.h>

#include <unsupported/Eigen/SparseExtra>

#include "linop.h"
#include "macau_mpi.h"
#include "macau.h"

extern "C" {
  #include "dsparse.h"
}

using namespace Eigen;
using namespace std;

void usage() {
   printf("Usage:\n");
   printf("  macau_mpi --train <train_file> --row-features <feature-file> [options]\n");
   printf("Optional:\n");
   printf("  --test    test_file  test data (for computing RMSE)\n");
   printf("  --burnin        200  number of samples to discard\n");
   printf("  --nsamples      800  number of samples to collect\n");
   printf("  --num-latent     96  number of latent dimensions\n");
   printf("  --precision     5.0  precision of observations\n");
   printf("  --lambda-beta  10.0  initial value of lambda beta\n");
   printf("  --tol          1e-6  tolerance for CG\n");
   printf("  --output    results  prefix for result files\n");
}

bool file_exists(const char *fileName)
{
   std::ifstream infile(fileName);
   return infile.good();
}

int get_num_omp_threads() {
   int nt = 0;
#pragma omp parallel
   {
#pragma omp single
      {
         nt = omp_get_num_threads();
      }
   }
   return nt;
}

void die(std::string message, int world_rank) {
   if (world_rank == 0) {
      std::cout << message;
   }
   MPI_Finalize();
   exit(1);
}

std::unique_ptr<SparseFeat> load_bcsr(const char* filename) {
   SparseBinaryMatrix* A = read_sbm(filename);
   SparseFeat* sf = new SparseFeat(A->nrow, A->ncol, A->nnz, A->rows, A->cols);
   free_sbm(A);
   std::unique_ptr<SparseFeat> sf_ptr(sf);
   return sf_ptr;
}

// var for MPI
int* rhs_for_rank = NULL;
double* rec     = NULL;
int* sendcounts = NULL;
int* displs     = NULL;

int main(int argc, char** argv) {
   // Initialize the MPI environment
   MPI_Init(NULL, NULL);
   // Get the number of processes
   int world_size, world_rank;
   MPI_Comm_size(MPI_COMM_WORLD, &world_size);
   MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);

   // Get the name of the processor
   char processor_name[MPI_MAX_PROCESSOR_NAME];
   int name_len;
   MPI_Get_processor_name(processor_name, &name_len);

   char* fname_train         = NULL;
   char* fname_test          = NULL;
   char* fname_row_features  = NULL;
   std::string output_prefix = std::string("result");
   double precision   = 5.0;
   double lambda_beta = 10.0;
   double tol         = 1e-6;
   int burnin     = 200;
   int nsamples   = 800;
   int num_latent = 96;

   // reading command line arguments
   while (1) {
      static struct option long_options[] =
      {
         {"train",      required_argument, 0, 't'},
         {"test",       required_argument, 0, 'e'},
         {"row-features", required_argument, 0, 'r'},
         {"precision",  required_argument, 0, 'p'},
         {"burnin",     required_argument, 0, 'b'},
         {"nsamples",   required_argument, 0, 'n'},
         {"output",     required_argument, 0, 'o'},
         {"num-latent", required_argument, 0, 'l'},
         {"lambda-beta",required_argument, 0, 'a'},
         {"tol",        required_argument, 0, 'c'},
         {0, 0, 0, 0}
      };
      int option_index = 0;
      int c = getopt_long(argc, argv, "t:e:r:p:b:n:o:a:c:", long_options, &option_index);
      if (c == -1)
         break;

      switch (c) {
         case 'a': lambda_beta   = strtod(optarg, NULL); break;
         case 'b': burnin        = strtol(optarg, NULL, 10); break;
         case 'c': tol           = atof(optarg); break;
         case 'e': fname_test    = optarg; break;
         case 'l': num_latent    = strtol(optarg, NULL, 10); break;
         case 'n': nsamples      = strtol(optarg, NULL, 10); break;
         case 'o': output_prefix = std::string(optarg); break;
         case 'p': precision     = strtod(optarg, NULL); break;
         case 'r': fname_row_features = optarg; break;
         case 't': fname_train = optarg; break;
         case '?':
         default:
           if (world_rank == 0)
              usage();
           MPI_Finalize();
           exit(1);
      }
   }
   if (fname_train == NULL || fname_row_features == NULL) {
      if (world_rank == 0) {
         printf("[ERROR]\nMissing parameters '--matrix' or '--row-features'.\n");
         usage();
      }
      MPI_Finalize();
      exit(1);
   }
   if (world_rank == 0) {
      printf("Train data:    '%s'\n", fname_train);
      printf("Test data:     '%s'\n", fname_test==NULL ?"" :fname_test);
      printf("Row features:  '%s'\n", fname_row_features);
      printf("Output prefix: '%s'\n", output_prefix.c_str());
      printf("Burn-in:       %d\n", burnin);
      printf("Samples:       %d\n", nsamples);
      printf("Num-latents:   %d\n", num_latent);
      printf("Precision:     %.1f\n", precision);
      printf("Lambda-beta:   %.1f\n", lambda_beta);
      printf("tol:           %.1e\n", tol);
   }
   if ( ! file_exists(fname_train) ) {
      die(std::string("[ERROR]\nTrain data file '") + fname_train + "' not found.\n", world_rank);
   }
   if ( ! file_exists(fname_row_features) ) {
      die(std::string("[ERROR]\nRow feature file '") + fname_row_features + "' not found.\n", world_rank);
   }
   if ( (fname_test != NULL) && ! file_exists(fname_test) ) {
      die(std::string("[ERROR]\nTest data file '") + fname_test + "' not found.\n", world_rank);
   }

   int n_omp_threads = get_num_omp_threads();
   rhs_for_rank = new int[world_size];
   split_work_mpi(num_latent, world_size, rhs_for_rank);

   // Print off a hello world message
   printf("Processor %s, rank %d"
          " out of %d processors using %d OpenMP threads for %d RHS.\n",
          processor_name, world_rank, world_size, n_omp_threads, rhs_for_rank[world_rank]);

   // Step 1. Loading data
   //std::unique_ptr<SparseFeat> row_features = load_bcsr(fname_row_features);
   auto row_features = load_bcsr(fname_row_features);
   if (world_rank == 0) {
      printf("Row features:   [%d x %d].\n", row_features->rows(), row_features->cols());
   }
   sendcounts = new int[world_size];
   displs     = new int[world_size];
   int sum = 0;
   for (int n = 0; n < world_size; n++) {
      sendcounts[n] = rhs_for_rank[n] * row_features->cols();
      displs[n]     = sum;
      sum          += sendcounts[n];
   }
   rec = new double[sendcounts[world_rank]];

   SparseDoubleMatrix* Y     = NULL;
   SparseDoubleMatrix* Ytest = NULL;

   Macau* macau = new Macau(num_latent);
   macau->setPrecision(precision);
   macau->setSamples(burnin, nsamples);
   macau->setVerbose(true);

   // 1) row prior with side information
   int nfeat    = row_features->rows();
   auto prior_u = new MacauPrior<SparseFeat>(num_latent, row_features, false);
   prior_u->setLambdaBeta(lambda_beta);
   prior_u->setTol(tol);
   auto prior_v = new BPMFPrior(num_latent);

   // 2) activity data (read_sdm)
   Y = read_sdm(fname_train);

   if (nfeat != Y->nrow) {
      die(std::string("[ERROR]\nNumber of rows (" +
                      std::to_string(Y->nrow) +
                      ") in train must be equal to number of rows in row-features (" +
                      std::to_string(row_features->rows()) +
                      ")."),
          world_rank);
   }

   // 3) create Macau object
   macau->setRelationData(Y->rows, Y->cols, Y->vals, Y->nnz, Y->nrow, Y->ncol);

   // test data
   if (fname_test != NULL) {
      Ytest = read_sdm(fname_test);
      if (Ytest->nrow != Y->nrow || Ytest->ncol != Y->ncol) {
         die(std::string("[ERROR]\nSize of train (") +
                         std::to_string(Y->nrow) + " x " +
                         std::to_string(Y->ncol) + ") must be equal to size of test (" +
                         std::to_string(Ytest->nrow) + " x " +
                         std::to_string(Ytest->ncol) + ").",
             world_rank);
      }
      macau->setRelationDataTest(Ytest->rows, Ytest->cols, Ytest->vals, Ytest->nnz, Ytest->nrow, Ytest->ncol);
   }
   std::unique_ptr<ILatentPrior> u_ptr(prior_u);
   std::unique_ptr<ILatentPrior> v_ptr(prior_v);
   macau->addPrior( u_ptr );
   macau->addPrior( v_ptr );

   if (world_rank == 0) {
      printf("Training data:  %ld [%d x %d]\n", Y->nnz, Y->nrow, Y->ncol);
      if (Ytest != NULL) {
         printf("Test data:      %ld [%d x %d]\n", Ytest->nnz, Ytest->nrow, Ytest->ncol);
      } else {
         printf("Test data:      --\n");
      }
   }

   run_macau_mpi(macau, world_rank);

   // save results
   if (world_rank == 0) {
      VectorXd yhat_raw     = macau->getPredictions();
      VectorXd yhat_sd_raw  = macau->getStds();
      MatrixXd testdata_raw = macau->getTestData();

      std::string fname_pred = output_prefix + "-predictions.csv";
      std::ofstream predfile;
      predfile.open(fname_pred);
      predfile << "row,col,y,y_pred,y_pred_std\n";
      for (int i = 0; i < yhat_raw.size(); i++) {
         predfile << to_string( (int)testdata_raw(i,0) );
         predfile << "," << to_string( (int)testdata_raw(i,1) );
         predfile << "," << to_string( testdata_raw(i,2) );
         predfile << "," << to_string( yhat_raw(i) );
         predfile << "," << to_string( yhat_sd_raw(i) );
         predfile << "\n";
      }
      predfile.close();
      printf("Saved predictions into '%s'.\n", fname_pred.c_str());
   }

   // Finalize the MPI environment.
   delete macau;
   MPI_Finalize();
   return 0;
}

void run_macau_mpi(
      Macau* macau,
      int world_rank)
{
   /* adapted from Macau.run() */
   macau->init();
   if (world_rank == 0 && macau->verbose) {
      std::cout << "Sampling" << std::endl;
   }

   const int num_rows = macau->Y.rows();
   const int num_cols = macau->Y.cols();
   macau->predictions     = VectorXd::Zero( macau->Ytest.nonZeros() );
   macau->predictions_var = VectorXd::Zero( macau->Ytest.nonZeros() );

   auto start = tick();
   for (int i = 0; i < macau->burnin + macau->nsamples; i++) {
      if (world_rank == 0 && macau->verbose && i == macau->burnin) {
         printf(" ====== Burn-in complete, averaging samples ====== \n");
      }
      auto starti = tick();

      if (world_rank == 0) {
         // sample latent vectors
         macau->noise->sample_latents(macau->priors[0], *macau->samples[0], macau->Yt, macau->mean_rating, *macau->samples[1], macau->num_latent);
         macau->noise->sample_latents(macau->priors[1], *macau->samples[1], macau->Y,  macau->mean_rating, *macau->samples[0], macau->num_latent);
         //macau->priors[0]->sample_latents(*macau->samples[0], macau->Yt, macau->mean_rating, *macau->samples[1], macau->alpha, macau->num_latent);
         //macau->priors[1]->sample_latents(*macau->samples[1], macau->Y,  macau->mean_rating, *macau->samples[0], macau->alpha, macau->num_latent);
      }

      // Sample hyperparams
      update_prior_mpi( *(MacauPrior<SparseFeat>*) macau->priors[0].get(), *macau->samples[0], world_rank);
      if (world_rank == 0) {
         macau->priors[1]->update_prior(*macau->samples[1]);

         auto eval = eval_rmse(macau->Ytest, (i < macau->burnin) ? 0 : (i - macau->burnin), macau->predictions, macau->predictions_var,
                               *macau->samples[1], *macau->samples[0], macau->mean_rating);

         auto endi = tick();
         auto elapsed = endi - start;
         double samples_per_sec = (i + 1) * (num_rows + num_cols) / elapsed;
         double elapsedi = endi - starti;

         if (macau->verbose) {
           macau->printStatus(i, eval.first, eval.second, elapsedi, samples_per_sec);
         }
         macau->rmse_test = eval.second;
      }
   }
}

void update_prior_mpi(MacauPrior<SparseFeat> &prior, const Eigen::MatrixXd &U, int world_rank) {
   if (world_rank == 0) {
      // residual (Uhat is later overwritten):
      prior.Uhat.noalias() = U - prior.Uhat;
      MatrixXd BBt = A_mul_At_combo(prior.beta);
      // sampling Gaussian
      std::tie(prior.mu, prior.Lambda) = CondNormalWishart(prior.Uhat, prior.mu0, prior.b0, prior.WI + prior.lambda_beta * BBt, prior.df + prior.beta.cols());
   }
   sample_beta_mpi(prior, U, world_rank);
   if (world_rank == 0) {
      compute_uhat(prior.Uhat, *prior.F, prior.beta);
      prior.lambda_beta = sample_lambda_beta(prior.beta, prior.Lambda, prior.lambda_beta_nu0, prior.lambda_beta_mu0);
   }
}

void sample_beta_mpi(MacauPrior<SparseFeat> &prior, const Eigen::MatrixXd &U, int world_rank) {
   const int num_latent = prior.beta.rows();
   const int num_feat = prior.beta.cols();
   MatrixXd Ft_y(0,0);

   if (world_rank == 0) {
      // Ft_y = (U .- mu + Normal(0, Lambda^-1)) * F + sqrt(lambda_beta) * Normal(0, Lambda^-1)
      // Ft_y is [ D x F ] matrix
      MatrixXd tmp = (U + MvNormal_prec_omp(prior.Lambda, U.cols())).colwise() - prior.mu;
      Ft_y = A_mul_B(tmp, *prior.F);
      MatrixXd tmp2 = MvNormal_prec_omp(prior.Lambda, num_feat);

      #pragma omp parallel for schedule(static)
      for (int f = 0; f < num_feat; f++) {
         for (int d = 0; d < num_latent; d++) {
            Ft_y(d, f) += sqrt(prior.lambda_beta) * tmp2(d, f);
         }
      }
      Ft_y.transposeInPlace();
   }

   MPI_Bcast(& prior.lambda_beta, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
   MPI_Bcast(& prior.tol,         1, MPI_DOUBLE, 0, MPI_COMM_WORLD);

   // sending Ft_y
   MPI_Scatterv(Ft_y.data(), sendcounts, displs, MPI_DOUBLE, rec, sendcounts[world_rank], MPI_DOUBLE, 0, MPI_COMM_WORLD);
   int nrhs = rhs_for_rank[world_rank];
   MatrixXd RHS(nrhs, num_feat), result(nrhs, num_feat);

#pragma omp parallel for schedule(static)
   for (int f = 0; f < num_feat; f++) {
      for (int d = 0; d < nrhs; d++) {
         RHS(d, f) = rec[f + d * num_feat];
      }
   }
   // solving
   solve_blockcg(result, *prior.F, prior.lambda_beta, RHS, prior.tol, 32, 8);
   result.transposeInPlace();
   MPI_Gatherv(result.data(), nrhs*num_feat, MPI_DOUBLE, Ft_y.data(), sendcounts, displs, MPI_DOUBLE, 0, MPI_COMM_WORLD);
   if (world_rank == 0) {
      //prior.beta = Ft_y.transpose();
#pragma omp parallel for schedule(static)
      for (int f = 0; f < num_feat; f++) {
         for (int d = 0; d < num_latent; d++) {
            prior.beta(d, f) = Ft_y(f, d);
         }
      }
   }
}