#ifndef __CLUSTER_FOREST_HPP__
#define __CLUSTER_FOREST_HPP__


#include <boost/pending/disjoint_sets.hpp>

#include "Transcript.hpp"
#include "TranscriptCluster.hpp"
#include "SalmonSpinLock.hpp"

#include <unordered_set>
#include <vector>
#include <mutex>

/** A forest of transcript clusters */
class ClusterForest {
public:
    ClusterForest(size_t numTranscripts, std::vector<Transcript>& refs) :
        rank_(std::vector<size_t>(numTranscripts, 0)),
        parent_(std::vector<size_t>(numTranscripts, 0)),
        disjointSets_(&rank_[0], &parent_[0]),
        clusters_(std::vector<TranscriptCluster>(numTranscripts))
    {
        // Initially make a unique set for each transcript
        for(size_t tnum = 0; tnum < numTranscripts; ++tnum) {
            disjointSets_.make_set(tnum);
            clusters_[tnum].members_.push_front(tnum);
            clusters_[tnum].addMass(refs[tnum].mass());
        }
    }

    template <typename FragT>
    void mergeClusters(typename std::vector<FragT>::iterator start,
                       typename std::vector<FragT>::iterator finish) {
        // Use a lock_guard to ensure this is a locked (and exception-safe) operation
#if defined __APPLE__
        spin_lock::scoped_lock sl(clusterMutex_);
#else
        std::lock_guard<std::mutex> lock(clusterMutex_);
#endif
        size_t firstCluster, otherCluster;
        auto firstTranscriptID = start->transcriptID();
        ++start;

        for (auto it = start; it != finish; ++it) {
            firstCluster = disjointSets_.find_set(firstTranscriptID);
            otherCluster = disjointSets_.find_set(it->transcriptID());
            if (otherCluster != firstCluster)  {
                disjointSets_.link(firstCluster, otherCluster);
                auto parentClust = disjointSets_.find_set(it->transcriptID());
                auto childClust = (parentClust == firstCluster)  ? otherCluster : firstCluster;
                if (parentClust == firstCluster or parentClust == otherCluster) {
                    clusters_[parentClust].merge(clusters_[childClust]);
                    clusters_[childClust].deactivate();
                } else { std::cerr << "DANGER\n"; }
            }
        }
    }


    template <typename FragT>
    void mergeClusters(typename std::vector<FragT*>::iterator start,
                       typename std::vector<FragT*>::iterator finish) {
        // Use a lock_guard to ensure this is a locked (and exception-safe) operation
#if defined __APPLE__
        spin_lock::scoped_lock sl(clusterMutex_);
#else
        std::lock_guard<std::mutex> lock(clusterMutex_);
#endif

        auto firstTranscriptID = (*start)->transcriptID();
        decltype(firstTranscriptID) firstCluster, otherCluster;
        ++start;

        for (auto it = start; it != finish; ++it) {
            firstCluster = disjointSets_.find_set(firstTranscriptID);
            otherCluster = disjointSets_.find_set((*it)->transcriptID());
            if (otherCluster != firstCluster)  {
                disjointSets_.link(firstCluster, otherCluster);
                auto parentClust = disjointSets_.find_set((*it)->transcriptID());
                auto childClust = (parentClust == firstCluster)  ? otherCluster : firstCluster;
                if (parentClust == firstCluster or parentClust == otherCluster) {
                    clusters_[parentClust].merge(clusters_[childClust]);
                    clusters_[childClust].deactivate();
                } else { std::cerr << "DANGER\n"; }
            }
        }
    }


    /*
    void mergeClusters(AlignmentBatch<ReadPair>::iterator start, AlignmentBatch<ReadPair>::iterator finish) {
        // Use a lock_guard to ensure this is a locked (and exception-safe) operation
        std::lock_guard<std::mutex> lock(clusterMutex_);
        size_t firstCluster, otherCluster;
        auto firstTranscriptID = start->read1->core.tid;
        ++start;

        for (auto it = start; it != finish; ++it) {
            firstCluster = disjointSets_.find_set(firstTranscriptID);
            otherCluster = disjointSets_.find_set(it->read1->core.tid);
            if (otherCluster != firstCluster)  {
                disjointSets_.link(firstCluster, otherCluster);
                auto parentClust = disjointSets_.find_set(it->read1->core.tid);
                auto childClust = (parentClust == firstCluster)  ? otherCluster : firstCluster;
                if (parentClust == firstCluster or parentClust == otherCluster) {
                    clusters_[parentClust].merge(clusters_[childClust]);
                    clusters_[childClust].deactivate();
                } else { std::cerr << "DANGER\n"; }
            }
        }
    }
*/
    void updateCluster(size_t memberTranscript, size_t newCount, double logNewMass, bool updateCount) {
        // Use a lock_guard to ensure this is a locked (and exception-safe) operation
#if defined __APPLE__
        spin_lock::scoped_lock sl(clusterMutex_);
#else
        std::lock_guard<std::mutex> lock(clusterMutex_);
#endif
        auto clusterID = disjointSets_.find_set(memberTranscript);
        auto& cluster = clusters_[clusterID];
        if (updateCount) {
            cluster.incrementCount(newCount);
        }
        cluster.addMass(logNewMass);
    }

    std::vector<TranscriptCluster*> getClusters() {
        std::vector<TranscriptCluster*> clusters;
        std::unordered_set<size_t> observedReps;
        for (size_t i = 0; i < clusters_.size(); ++i) {
            auto rep = disjointSets_.find_set(i);
            if (observedReps.find(rep) == observedReps.end()) {
                if (!clusters_[rep].isActive()) {
                    std::cerr << "returning a non-active cluster!\n";
                    std::exit(1);
                }
                clusters.push_back(&clusters_[rep]);
                observedReps.insert(rep);
            }
        }
        return clusters;
    }
private:
    std::vector<size_t> rank_;
    std::vector<size_t> parent_;
    boost::disjoint_sets<size_t*, size_t*> disjointSets_;
    std::vector<TranscriptCluster> clusters_;
    #if defined __APPLE__
    spin_lock clusterMutex_;
    #else
    std::mutex clusterMutex_;
    #endif
};

#endif // __CLUSTER_FOREST_HPP__
