#pragma once

#include <string>
#include <random>
#include <vector>

#include "nifty/graph/opt/common/proposal_generators/proposal_generator_base.hxx"


namespace nifty{
namespace graph{
namespace opt{
namespace common{


    /**
     * @brief      Watershed proposal generator for lifted_multicut::FusionMoveBased
     *
     * @tparam     OBJECTIVE  { description }
     */
    template<class OBJECTIVE>
    class RandomNodeColorProposalGenerator : 
        public ProposalGeneratorBase<OBJECTIVE>{

        
    public:
        typedef OBJECTIVE ObjectiveType;
        typedef typename ObjectiveType::GraphType GraphType;
        typedef typename GraphType:: template NodeMap<uint64_t>  ProposalType;
        typedef typename GraphType:: template NodeMap<bool>  IsUsed;    
            
        struct SettingsType{
            std::size_t numberOfColors{std::size_t(2)};
        };

        RandomNodeColorProposalGenerator(
            const ObjectiveType & objective, 
            const std::size_t numberOfThreads,
            const SettingsType & settings  = SettingsType()
        )
        :   objective_(objective),
            numberOfThreads_(numberOfThreads),
            settings_(settings),
            gens_(numberOfThreads_),
            colorDist_(0,settings.numberOfColors-1)
        {
            // use thread index as seed
            for(auto i=0; i<numberOfThreads_; ++i){
                gens_[i] = std::mt19937(i);
            }
            this->reset();
        }
        void reset(){

        }
        virtual ~RandomNodeColorProposalGenerator(){}

        virtual void generateProposal(
            const ProposalType & currentBest, 
            ProposalType & proposal, 
            const std::size_t tid
        ){  

            const auto & graph = objective_.graph();
            auto & gen = gens_[tid];
   

            // set to unused
            // and copy current best
            graph.forEachNode([&](const uint64_t node){
                proposal[node] = colorDist_(gen);
            });
        }
    private:
        const ObjectiveType & objective_;
        std::size_t numberOfThreads_;
        SettingsType settings_;
        std::vector<std::mt19937> gens_;
        std::uniform_int_distribution<>  colorDist_;
    }; 


} // namespace nifty::graph::opt::common
} // namespace nifty::graph::opt
} // namespace nifty::graph
} // namespace nifty
