/*********************                                                        */
/*! \file ParallelSolver.h
 ** \verbatim
 ** Top contributors (to current version):
 **   Haoze Wu
 ** This file is part of the Marabou project.
 ** Copyright (c) 2017-2024 by the authors listed in the file AUTHORS
 ** in the top-level source directory) and their institutional affiliations.
 ** All rights reserved. See the file COPYING in the top-level source
 ** directory for licensing information.\endverbatim
 **
 ** [[ Add lengthier description here ]]

**/

#ifndef __ParallelSolver_h__
#define __ParallelSolver_h__

#include "GurobiWrapper.h"

#include <atomic>
#include <boost/lockfree/queue.hpp>
#include <boost/thread.hpp>
#include <mutex>

namespace NLR {

class Layer;
class LayerOwner;
struct NeuronIndex;

class ParallelSolver
{
public:
    typedef boost::lockfree::queue<GurobiWrapper *, boost::lockfree::fixed_sized<true>> SolverQueue;

    /*
      Arguments for the spawned thread. This is needed because Boost::thread does
      not seem to support functions with more than 7 arguments.
    */
    struct ThreadArgument
    {
        ThreadArgument( GurobiWrapper *gurobi,
                        Layer *layer,
                        const Map<unsigned, Layer *> *layers,
                        unsigned index,
                        double currentLb,
                        double currentUb,
                        bool cutoffInUse,
                        double cutoffValue,
                        LayerOwner *layerOwner,
                        SolverQueue &freeSolvers,
                        std::mutex &mtx,
                        std::atomic_bool &infeasible,
                        std::atomic_uint &tighterBoundCounter,
                        std::atomic_uint &signChanges,
                        std::atomic_uint &cutoffs,
                        bool skipTightenLb,
                        bool skipTightenUb )
            : _gurobi( gurobi )
            , _layer( layer )
            , _layers( layers )
            , _index( index )
            , _currentLb( currentLb )
            , _currentUb( currentUb )
            , _cutoffInUse( cutoffInUse )
            , _cutoffValue( cutoffValue )
            , _layerOwner( layerOwner )
            , _freeSolvers( freeSolvers )
            , _mtx( mtx )
            , _infeasible( infeasible )
            , _tighterBoundCounter( tighterBoundCounter )
            , _signChanges( signChanges )
            , _cutoffs( cutoffs )
            , _skipTightenLb( skipTightenLb )
            , _skipTightenUb( skipTightenUb )
            , _lastFixedNeuron( NULL )
        {
        }

        ThreadArgument( GurobiWrapper *gurobi,
                        Layer *layer,
                        unsigned index,
                        double currentLb,
                        double currentUb,
                        bool cutoffInUse,
                        double cutoffValue,
                        LayerOwner *layerOwner,
                        SolverQueue &freeSolvers,
                        std::mutex &mtx,
                        std::atomic_bool &infeasible,
                        std::atomic_uint &tighterBoundCounter,
                        std::atomic_uint &signChanges,
                        std::atomic_uint &cutoffs,
                        bool skipTightenLb,
                        bool skipTightenUb )
            : _gurobi( gurobi )
            , _layer( layer )
            , _layers( NULL )
            , _index( index )
            , _currentLb( currentLb )
            , _currentUb( currentUb )
            , _cutoffInUse( cutoffInUse )
            , _cutoffValue( cutoffValue )
            , _layerOwner( layerOwner )
            , _freeSolvers( freeSolvers )
            , _mtx( mtx )
            , _infeasible( infeasible )
            , _tighterBoundCounter( tighterBoundCounter )
            , _signChanges( signChanges )
            , _cutoffs( cutoffs )
            , _skipTightenLb( skipTightenLb )
            , _skipTightenUb( skipTightenUb )
            , _lastFixedNeuron( NULL )
        {
        }

        ThreadArgument( GurobiWrapper *gurobi,
                        Layer *layer,
                        unsigned index,
                        double currentLb,
                        double currentUb,
                        bool cutoffInUse,
                        double cutoffValue,
                        LayerOwner *layerOwner,
                        SolverQueue &freeSolvers,
                        std::mutex &mtx,
                        std::atomic_bool &infeasible,
                        std::atomic_uint &tighterBoundCounter,
                        std::atomic_uint &signChanges,
                        std::atomic_uint &cutoffs,
                        NeuronIndex *lastFixedNeuron )
            : _gurobi( gurobi )
            , _layer( layer )
            , _layers( NULL )
            , _index( index )
            , _currentLb( currentLb )
            , _currentUb( currentUb )
            , _cutoffInUse( cutoffInUse )
            , _cutoffValue( cutoffValue )
            , _layerOwner( layerOwner )
            , _freeSolvers( freeSolvers )
            , _mtx( mtx )
            , _infeasible( infeasible )
            , _tighterBoundCounter( tighterBoundCounter )
            , _signChanges( signChanges )
            , _cutoffs( cutoffs )
            , _lastFixedNeuron( lastFixedNeuron )
        {
        }

        ThreadArgument( Layer *layer,
                        const Map<unsigned, Layer *> *layers,
                        SolverQueue &freeSolvers,
                        std::mutex &mtx,
                        std::atomic_bool &infeasible,
                        std::atomic_uint &tighterBoundCounter,
                        std::atomic_uint &signChanges,
                        std::atomic_uint &cutoffs,
                        unsigned lastIndexOfRelaxation,
                        unsigned targetIndex,
                        boost::thread *threads,
                        const Map<GurobiWrapper *, unsigned> *solverToIndex )
            : _layer( layer )
            , _layers( layers )
            , _freeSolvers( freeSolvers )
            , _mtx( mtx )
            , _infeasible( infeasible )
            , _tighterBoundCounter( tighterBoundCounter )
            , _signChanges( signChanges )
            , _cutoffs( cutoffs )
            , _lastIndexOfRelaxation( lastIndexOfRelaxation )
            , _targetIndex( targetIndex )
            , _threads( threads )
            , _solverToIndex( solverToIndex )
        {
        }

        GurobiWrapper *_gurobi;
        Layer *_layer;
        const Map<unsigned, Layer *> *_layers;
        unsigned _index;
        double _currentLb;
        double _currentUb;
        bool _cutoffInUse;
        double _cutoffValue;
        LayerOwner *_layerOwner;
        SolverQueue &_freeSolvers;
        std::mutex &_mtx;
        std::atomic_bool &_infeasible;
        std::atomic_uint &_tighterBoundCounter;
        std::atomic_uint &_signChanges;
        std::atomic_uint &_cutoffs;
        bool _skipTightenLb;
        bool _skipTightenUb;
        NeuronIndex *_lastFixedNeuron;
        unsigned _lastIndexOfRelaxation;
        unsigned _targetIndex;
        boost::thread *_threads;
        const Map<GurobiWrapper *, unsigned> *_solverToIndex;
    };

    /*
      Free all the solvers in the SolverQueue
    */
    static void clearSolverQueue( SolverQueue &freeSolvers );

    static void enqueueSolver( SolverQueue &solvers, GurobiWrapper *solver );
};

} // namespace NLR

#endif // __ParallelSolver_h__
