/*
 * nvbio
 * Copyright (c) 2011-2014, NVIDIA CORPORATION. All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *    * Redistributions of source code must retain the above copyright
 *      notice, this list of conditions and the following disclaimer.
 *    * Redistributions in binary form must reproduce the above copyright
 *      notice, this list of conditions and the following disclaimer in the
 *      documentation and/or other materials provided with the distribution.
 *    * Neither the name of the NVIDIA CORPORATION nor the
 *      names of its contributors may be used to endorse or promote products
 *      derived from this software without specific prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <cstdlib>
#include <nvbio/basic/thrust_view.h>
#include <nvbio/basic/primitives.h>
#include <nvbio/basic/atomics.h>
#include <nvbio/basic/cuda/sort.h>
#include <nvbio/basic/cuda/primitives.h>
#include <nvbio/strings/seeds.h>

#include "assembly_types.h"

/* K-mer Extraction/Counting/Uniqueness Functionality */

// extract the string set sequence id from the kmer coordinates
struct get_kmer_seq_id : public thrust::unary_function<SequenceSetKmerCoord,uint32>
{
    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    uint32 operator() (const SequenceSetKmerCoord kmer_coord) const
    {
        return kmer_coord.x;
    }
};

struct get_kmer_reg_id
{
    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    uint32 operator() (const SequenceSetKmerCoord kmer_coord) const
    {
        return kmer_coord.w;
    }
};

// extract the kmer size from the kmer coordinates
struct get_kmer_size : public thrust::unary_function<SequenceSetKmerCoord,uint32>
{
    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    uint32 operator() (const SequenceSetKmerCoord kmer_coord) const
    {
        return kmer_coord.z - kmer_coord.y;
    }
};

struct get_global_id : public thrust::unary_function<SequenceSetKmerCoord,uint32>
{
	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    uint32 operator() (const SequenceSetKmerCoord coord) const
    {
        return coord.z;
    }
};

// store the global kmer id in its coordinate
struct set_global_id
{
	SequenceSetKmerCoord* coords;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	set_global_id(SequenceSetKmerCoord* coords) : coords(coords) { }

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    void operator() (const uint32 idx) const
    {
        coords[idx].z = idx;
    }
};

// store the global kmer id and the active region id in its coordinate
struct set_global_id_region_id
{
	SequenceSetKmerCoord* coords;
	const uint32* active_region_ids;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	set_global_id_region_id(SequenceSetKmerCoord* coords, const uint32* _active_region_ids) :
		coords(coords), active_region_ids(_active_region_ids) { }

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    void operator() (const uint32 idx) const
    {
		SequenceSetKmerCoord& coord = coords[idx];
		coord.z = idx; // since the kmer size is fixed, can reuse this field
        coord.w = active_region_ids[coord.x];
    }
};

// maps from the coordinate of a kmer
// to the global id of its prefix
struct get_prefix_global_id : public thrust::unary_function<uint32,uint32>
{
    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    uint32 operator() (const SequenceSetKmerCoord coord) const
    {
    	return coord.z + coord.x;
    }
};

// maps from the id of a kmer
// to the global id of its prefix
struct get_prefix_global_id_by_idx : public thrust::unary_function<uint32,uint32>
{
	const SequenceSetKmerCoord* coords;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	get_prefix_global_id_by_idx(const SequenceSetKmerCoord* _coords):
	coords(_coords) { }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    uint32 operator() (const uint32 idx) const
    {
    	const SequenceSetKmerCoord coord = coords[idx];
    	return coord.z + coord.x;
    }
};

// maps from the id of a kmer
// to the global id of its suffix
struct get_suffix_global_id_by_idx : public thrust::unary_function<uint32,uint32>
{
	const SequenceSetKmerCoord* coords;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	get_suffix_global_id_by_idx(const SequenceSetKmerCoord* _coords):
	coords(_coords) { }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    uint32 operator() (const uint32 idx) const
    {
    	const SequenceSetKmerCoord coord = coords[idx];
    	return coord.z + coord.x + 1;
    }
};

// creates a map from the id of a kmer
// to the global id of its prefix
struct compute_prefix_global_id
{
	const SequenceSetKmerCoord* coords;
	uint32* prefix_ids;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	compute_prefix_global_id(const SequenceSetKmerCoord* _coords, uint32* _prefix_ids):
	coords(_coords), prefix_ids(_prefix_ids) { }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    void operator() (const uint32 id) const
    {
    	const SequenceSetKmerCoord coord = coords[id];
    	prefix_ids[id] = coord.z + coord.x;
    }
};

// creates a map from the id of a kmer
// to the global id of its suffix
struct compute_suffix_global_id
{
	const SequenceSetKmerCoord* coords;
	uint32* suffix_ids;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	compute_suffix_global_id(const SequenceSetKmerCoord* _coords, uint32* _suffix_ids):
	coords(_coords), suffix_ids(_suffix_ids) { }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    void operator() (const uint32 id) const
    {
    	const SequenceSetKmerCoord coord = coords[id];
    	suffix_ids[id] = coord.z + coord.x + 1;
    }
};

// maps kmer global id to its UID
struct global_to_uid : public thrust::unary_function<uint64,uint64>
{
	const uint32* uid_map;
	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	global_to_uid(const uint32* _uid_map) : uid_map(_uid_map) { }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    uint32 operator() (const uint32 global_id) const
    {
    	return uid_map[global_id];
    }
};

// creates a map from the global id to the id in the sorted array
struct global_to_sorted_id
{
	uint32* global_to_sorted_id_map;
	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	global_to_sorted_id(uint32* _map) : global_to_sorted_id_map(_map) { }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    void operator() (const thrust::tuple<uint32, SequenceSetKmerCoord> idx_coord) const
    {
    	const uint32 idx = thrust::get<0>(idx_coord);
    	const SequenceSetKmerCoord coord = thrust::get<1>(idx_coord);
    	global_to_sorted_id_map[coord.z] = idx; // TODO: not coalesced
    }
};

// compute the kmer keys: 64-bit compacted kmer sequence
template <typename string_set_type>
struct kmers_to_64b_functor
{
    typedef typename string_set_type::string_type sequence;
    const string_set_type   string_set;
    const uint32            kmer_size;
	const uint32            dna_symbol_size;
	const uint32            symbol_mask;

	NVBIO_HOST_DEVICE
    kmers_to_64b_functor(const uint32 _kmer_size, const uint32 _dna_symbol_size, const string_set_type _string_set) :
        kmer_size(_kmer_size), dna_symbol_size(_dna_symbol_size), symbol_mask((1u << dna_symbol_size) - 1u),
        string_set(_string_set) {}

    NVBIO_HOST_DEVICE
    uint64 operator() (const SequenceSetKmerCoord kmer_coord) const
    {
    	const sequence seq = string_set[kmer_coord.x];
        const uint32 seq_pos = kmer_coord.y;
        const uint32 seq_len = seq.length();

        uint64 kmer_key = 0u;
        for (uint32 i = 0; i < kmer_size; i++) {
        	kmer_key |= uint64(seq_pos + i < seq_len ? (seq[seq_pos + i] & symbol_mask) : 0u) << ((kmer_size-1-i)*dna_symbol_size); //(i*dna_symbol_size);
        }
        return kmer_key;
    }
};

// equality based on kmer key and sequence id coordinate
struct kmer_key_sid_eq
{
	const uint64* keys;
	const SequenceSetKmerCoord* coords;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	kmer_key_sid_eq(const uint64* _keys, const SequenceSetKmerCoord* _coords):
		keys(_keys), coords(_coords) { }

	NVBIO_HOST_DEVICE bool operator()(const uint32 i, const uint32 j)
	{
		return keys[i] == keys[j] && coords[i].x == coords[j].x;
	}
};

// equality based on kmer key and region id coordinate
struct kmer_key_rid_eq
{
	const uint64* keys;
	const SequenceSetKmerCoord* coords;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	kmer_key_rid_eq(const uint64* _keys, const SequenceSetKmerCoord* _coords):
		keys(_keys), coords(_coords) { }

	NVBIO_HOST_DEVICE bool operator()(const uint32 i, const uint32 j)
	{
		return keys[i] == keys[j] && coords[i].w == coords[j].w;
	}
};

// equality based on sequence id and position
struct kmer_pos_sid_eq
{
	const SequenceSetKmerCoord* coords;
	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	kmer_pos_sid_eq(const SequenceSetKmerCoord* _coords): coords(_coords) { }

	NVBIO_HOST_DEVICE bool operator()(const uint32 i, const uint32 j)
	{
		const SequenceSetKmerCoord ci = coords[i];
		const SequenceSetKmerCoord cj = coords[j];
		return ci.y == cj.y && ci.x == cj.x;
	}
};

// equality based on kmer key only
struct kmer_uid_eq
{
	const uint32* keys;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	kmer_uid_eq(const uint32* _keys) : keys(_keys) { }

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	bool operator()(const uint32 i, const uint32 j)
	{
		return keys[i] == keys[j];
	}
};

// extract i-th word from the kmer in a string set
template <typename string_set_type>
struct kmer_word_extractor_functor
{
	typedef typename string_set_type::string_type sequence;
	const string_set_type   string_set;
	const uint32            dna_symbol_size;
	const uint32            symbol_offset;
	const uint32            symbols_per_word;
	const SequenceSetKmerCoord* coords;
	const uint32 			i;

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    kmer_word_extractor_functor(const string_set_type _string_set, const uint32 _dna_symbol_size,
    		const SequenceSetKmerCoord* _coords, const uint32 _i) :
    	dna_symbol_size(_dna_symbol_size),
    	symbol_offset(uint32(8u*sizeof(uint32)) - dna_symbol_size),
    	symbols_per_word(uint32(8u * sizeof(uint32))/dna_symbol_size),
    	coords(_coords), string_set(_string_set), i(_i) {}

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    uint32 operator() (const uint32 kmer_idx) const
    {
    	SequenceSetKmerCoord kmer_coord = coords[kmer_idx];
    	const uint32 kmer_pos = kmer_coord.y;
    	const uint32 kmer_size = kmer_coord.z - kmer_coord.y; // TODO: z-field is currently used for the global id, use kmer_size if fixed
    	const sequence seq = string_set[kmer_coord.x];
    	const uint32 seq_len = seq.length();

		uint32 word = 0u;
		for (uint32 j = 0; j < kmer_size; j++) {
			const uint32 jj = kmer_pos + i*symbols_per_word + j;
			const uint32 c = jj < seq_len ? char_to_iupac16(dna_to_char(seq[jj])) : 0u;
			word |= (c << (symbol_offset - j*dna_symbol_size));
		}
    	return word;
    }
};

// add tuples of kmer and sequence counts
struct kmer_count_tuple_sum
{
	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	thrust::tuple<uint32, uint32> operator() (const thrust::tuple<uint32, uint32>& x, const thrust::tuple<uint32, uint32>& y) const
    {
        return thrust::tuple<uint32, uint32>(thrust::get<0>(x) + thrust::get<0>(y), thrust::get<1>(x) + thrust::get<1>(y));
    }
};

// check if a kmer occurs more than once in at least one sequence
// given a tuple of the number of sequences that contain it and
// the number of times it occurs in all the sequences
struct is_unique_kmer
{
	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    bool operator() (const thrust::tuple<uint32, uint32>& nSeq_nOcc) const
    {
        return thrust::get<0>(nSeq_nOcc) >= thrust::get<1>(nSeq_nOcc);
    }
};

// given the last id of a kmer in the group of consecutive equal kmers
// mark all identical consecutive kmers as unique in the map
struct mark_kmer_uniqueness
{
	uint8* is_unique_map;
	const uint64* sorted_kmer_keys;
	const SequenceSetKmerCoord* sorted_coords;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	mark_kmer_uniqueness(uint8* _is_unique_map, const uint64* _sorted_kmer_keys, const SequenceSetKmerCoord* _sorted_coords) :
	is_unique_map(_is_unique_map), sorted_kmer_keys(_sorted_kmer_keys), sorted_coords(_sorted_coords) { }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    void operator() (const uint32 kmer_id) const
    {
    	const uint64 kmer_key = sorted_kmer_keys[kmer_id];
    	const uint32 rid = sorted_coords[kmer_id].w;

    	int32 idx = kmer_id;
    	while(idx >= 0) {
    		if(sorted_coords[idx].w != rid || sorted_kmer_keys[idx] != kmer_key) break;
    		is_unique_map[idx] = 1; //TODO: thrust documentation states first, but last in experiments
    		idx--;
    	}
    }
};

struct store_kmer_unique_ids
{
	uint32* unique_id_map;
	const uint64* sorted_kmer_keys;
	const SequenceSetKmerCoord* sorted_coords;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	store_kmer_unique_ids(uint32* _unique_id_map, const uint64* _sorted_kmer_keys, const SequenceSetKmerCoord* _sorted_coords) :
	unique_id_map(_unique_id_map), sorted_kmer_keys(_sorted_kmer_keys), sorted_coords(_sorted_coords) { }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    void operator() (const thrust::tuple<uint32, uint32>& uid_kid) const
    {
    	const uint32 uid = thrust::get<0>(uid_kid);
    	const uint32 kid = thrust::get<1>(uid_kid);
    	const uint64 kmer_key = sorted_kmer_keys[kid];
    	const uint32 rid = sorted_coords[kid].w;

    	int32 idx = kid;
    	while(idx >= 0) {
    		if(sorted_coords[idx].w != rid || sorted_kmer_keys[idx] != kmer_key) break;
    		unique_id_map[idx] = uid;
    		idx--; //TODO: thrust documentation states first, but last in experiments
    	}
    }
};

//struct populate_unique_kmer_data
//{
//	const uint64* sorted_kmer_keys;
//	const SequenceSetKmerCoord* sorted_coords;
//	const uint32* unique_kmer_idxs;
//
//	// fill in unique kmer annotations
//	uint8* unique_flag_map;
//	uint32* unique_UID_map;
//
//	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
//	populate_unique_kmer_data(const uint64* _sorted_kmer_keys, const SequenceSetKmerCoord* _coords,
//			const uint32* _unique_kmer_idxs, uint8* _unique_flag_map, uint32* _unique_UID_map) :
//	sorted_kmer_keys(_sorted_kmer_keys), sorted_coords(_coords), unique_kmer_idxs(_unique_kmer_idxs),
//	unique_flag_map(_unique_flag_map), unique_UID_map (_unique_UID_map) { }
//
//    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
//    void operator() (const uint32 uid) const
//    {
//    	const uint32 unique_kmer_idx = unique_kmer_idxs[uid];
//    	const uint64 key = sorted_kmer_keys[unique_kmer_idx]; // find the group of coordinates with this key and region
//    	const uint32 region = sorted_coords[unique_kmer_idx].w;
//    	uint32 global_kmer_id = sorted_coords[unique_kmer_idx].z;
//    	unique_flag_map[global_kmer_id] = 1;
//    	unique_UID_map[global_kmer_id] = uid;
//
//    	int32 idx = unique_kmer_idx - 1;
//    	while(idx >= 0) {
//    		if(sorted_kmer_keys[idx] != key || sorted_coords[idx].w != region) break;
//    		global_kmer_id = sorted_coords[idx].z;
//    		unique_flag_map[global_kmer_id] = 1; //TODO: thrust documentation states first, but last in experiments
//    		unique_UID_map[global_kmer_id] = uid;
//    		idx--;
//    	}
//    }
//};

struct populate_unique_kmer_data
{
	const uint32* kmer_counts;
	const SequenceSetKmerCoord* sorted_coords;
	const uint32* unique_kmer_idxs;

	// fill in unique kmer annotations
	uint8* unique_flag_map;
	uint32* unique_UID_map;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	populate_unique_kmer_data(const uint32* _kmer_counts, const SequenceSetKmerCoord* _coords,
			const uint32* _unique_kmer_idxs, uint8* _unique_flag_map, uint32* _unique_UID_map) :
	kmer_counts(_kmer_counts), sorted_coords(_coords), unique_kmer_idxs(_unique_kmer_idxs),
	unique_flag_map(_unique_flag_map), unique_UID_map (_unique_UID_map) { }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    void operator() (const uint32 uid) const
    {
    	const uint32 unique_kmer_idx = unique_kmer_idxs[uid];
    	const uint32 count = kmer_counts[uid];
    	uint32 global_kmer_id = sorted_coords[unique_kmer_idx].z;
    	unique_flag_map[global_kmer_id] = 1;
    	unique_UID_map[global_kmer_id] = uid;

    	for(uint32 i = 1; i < count; i++) {
    		global_kmer_id = sorted_coords[unique_kmer_idx - i].z; //TODO: thrust bug: pointer to last instead of first
    		unique_flag_map[global_kmer_id] = 1; // TODO: not coalesced
    		unique_UID_map[global_kmer_id] = uid;
    	}
    }
};

// marks kmers that have a prefix satisfying the following:
// the prefix is unique and has a matching kmer
// (i.e. it is not the last (k-1)mer in the sequence)
struct mark_unique_prefix_kmer
{
	uint8* unique_prefix_map;
	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	mark_unique_prefix_kmer(uint8* _unique_prefix_map) : unique_prefix_map(_unique_prefix_map) { }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    void operator() (const thrust::tuple<uint8, uint64>& uflag_prefix) const
    {
    	const uint8 is_unique_prefix = thrust::get<0>(uflag_prefix);
    	const uint32 kmer_id = thrust::get<1>(uflag_prefix);

    	if((is_unique_prefix == 1) && (kmer_id != (uint32) -1)) {
    		unique_prefix_map[kmer_id] = 1; // TODO: writes not coalesced
    	}
    }
};

// ---------- Super-kmers: chains of consecutive kmers

// given a non-unique kmer N, finds the closest unique predecessor and successor kmers in the sequence, P and S
// returns the new coordinate for the chain starting with the P and ending with S [P...N...S]
// returns a dummy coordinate (N.x, -1, -1, 0) if no such previous unique kmer exists
// TODO: could use shared memory here
template <typename string_set_type>
struct extract_super_kmers_functor
{
	typedef typename string_set_type::string_type sequence;
	const string_set_type   string_set;
	const uint32* repeat_global_ids; // global ids of the repeat kmers
	const uint8* global_unique_flag_map; // unsorted
	const uint32* global_UID_map;
	const uint32* global_sorted_idx_map;
	const SequenceSetKmerCoord* coords; // sorted
	const uint32 kmer_size;

	// structures to populate
	SequenceSetKmerCoord* super_coords;
	uint32* prefix_uids;
	uint32* suffix_uids;

	NVBIO_HOST_DEVICE
	extract_super_kmers_functor(const uint32* _repeat_global_ids,
			const uint8* _is_unique_map, const uint32* _global_UID_map, const uint32* _sorted_idx_map,
			const SequenceSetKmerCoord* _coords, const string_set_type _string_set, const uint32 _kmer_size,
			SequenceSetKmerCoord* _super_coords, uint32* _prefix_uids, uint32* _suffix_uids):
				repeat_global_ids(_repeat_global_ids), global_unique_flag_map(_is_unique_map), global_UID_map(_global_UID_map),
				global_sorted_idx_map(_sorted_idx_map),
				coords(_coords), string_set(_string_set), kmer_size(_kmer_size),
				super_coords(_super_coords), prefix_uids(_prefix_uids), suffix_uids(_suffix_uids){ }

	NVBIO_HOST_DEVICE
	void operator()(const uint32 idx) const {
		const uint32 global_id = repeat_global_ids[idx];
		const SequenceSetKmerCoord coord = coords[global_sorted_idx_map[global_id]];
		const sequence seq = string_set[coord.x];
		const uint32 seq_len = seq.length();
		SequenceSetKmerCoord super_coord = make_uint4(coord.x, (uint32)-1, (uint32)-1, coord.w);
		uint32 prefix_id = (uint32) -1;
		uint32 suffix_id = (uint32) -1;

		// find the unique predecessor
		uint32 i = 1u;
		while(i <= coord.y) { // search only this sequence
			if(global_unique_flag_map[global_id-i] == 1) {
				super_coord.y = coord.y - i;
				super_coord.z = coord.y + kmer_size;
				prefix_id = global_UID_map[global_id-i];
				break;
			}
			i++;
		}

		if(super_coord.y != (uint32) -1) { // otherwise this chain will be ignored
			// find the unique successor
			i = 1u;
			while(i < seq_len - coord.y - kmer_size + 1) { // search only until the end of the sequence
				if(global_unique_flag_map[global_id+i] == 1) {
					super_coord.z = coord.y + i + kmer_size;
					suffix_id = global_UID_map[global_id+i];
					break;
				}
				i++;
			}
		}

		super_coords[idx] = super_coord;
		prefix_uids[idx] = prefix_id;
		suffix_uids[idx] = suffix_id;
	}
};

template <typename string_set_type>
struct print_super_kmer_functor
{
	typedef typename string_set_type::string_type sequence;
	const string_set_type   string_set;
	const SequenceSetKmerCoord* coords;

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    print_super_kmer_functor(const string_set_type _string_set, const SequenceSetKmerCoord* _coords) :
    	coords(_coords), string_set(_string_set) {}

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    void operator() (const uint32 kmer_idx) const
    {
    	SequenceSetKmerCoord kmer_coord = coords[kmer_idx];
    	const uint32 kmer_pos = kmer_coord.y;
    	const uint32 kmer_size = kmer_coord.z - kmer_coord.y;
    	const sequence seq = string_set[kmer_coord.x];
    	const uint32 seq_len = seq.length();

    	uint8 kmer_seq[14];
    	for (uint32 k = 0; k < kmer_size; k++) {
			uint8 c = iupac16_to_char(char_to_iupac16(dna_to_char(seq[kmer_pos + k])));
			kmer_seq[k] = c;
		}
    	kmer_seq[kmer_size] = '\0';

    	printf("id: %llu, reg %u, seq %s \n", kmer_idx, kmer_coord.w, kmer_seq);
    }
};

// returns true if the super kmer does not represent a valid chain of kmers
// based on the value of its unique predecessor id
struct is_invalid_super_kmer
{
	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	bool operator()(const uint32 chain_prefix_id)
	{
		return chain_prefix_id == (uint32) -1;
	}
};

struct is_unique_suffix_id : public thrust::unary_function<uint32,bool>
{
	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	bool operator()(const uint32 suffix_id)
	{
		return suffix_id != (uint32) -1;
	}
};

struct is_unique_uid : public thrust::unary_function<uint32,uint8>
{
	const uint32 max_unique_uid;
	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	is_unique_uid(const uint32 _max_unique_uid) : max_unique_uid(_max_unique_uid) { }

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	uint8 operator()(const uint32 uid)
	{
		return uid < max_unique_uid;
	}
};

struct is_repeat_kmer
{
	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	bool operator()(const uint8 flag)
	{
		return (flag != 1);
	}
};

// collapse kmers that start at the same position and belong to the same sequence
// by just returning the one with max length
struct collapse_same_start_kmers
{
	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	thrust::tuple<SequenceSetKmerCoord, uint32, uint32> operator() (
			const thrust::tuple<SequenceSetKmerCoord, uint32, uint32>& x,
			const thrust::tuple<SequenceSetKmerCoord, uint32, uint32>& y) const
    {
		const SequenceSetKmerCoord i = thrust::get<0>(x);
		const SequenceSetKmerCoord j = thrust::get<0>(y);

		if(i.z > j.z) {
			return x;
		}
		return y;
    }
};

// computes the max prefix overlap with the closest lexicographically smaller super-kmer
// assumes that the function is not called for idx == 0
// returns the number of kmers by which the chains overlap
// no overlap if the previous kmer is from a different region
template <typename string_set_type>
struct find_max_kmer_overlaps
{
	typedef typename string_set_type::string_type sequence;
	const string_set_type   string_set;
	const SequenceSetKmerCoord* coords; // sorted in lexicographic order
	const uint32 kmer_size;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	find_max_kmer_overlaps(const string_set_type _string_set, const SequenceSetKmerCoord* _coords,
			const uint32 _kmer_size): string_set(_string_set), coords(_coords), kmer_size(_kmer_size) { }

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	uint32 operator() (const uint32 idx) const
	{
		const SequenceSetKmerCoord coord_curr = coords[idx];
		const SequenceSetKmerCoord coord_prev = coords[idx-1];

		if(coord_curr.w != coord_prev.w) return 0;

		const sequence seq_curr = string_set[coord_curr.x];
		const sequence seq_prev = string_set[coord_prev.x];
		const uint32 max_overlap =
				(coord_curr.z - coord_curr.y) > (coord_prev.z - coord_prev.y) ?
				(coord_prev.z - coord_prev.y) : (coord_curr.z - coord_curr.y);

		uint32 overlap = 0u;
		for (uint32 j = 0; j < max_overlap; j++)
		{
			if(seq_curr[coord_curr.y + j] == seq_prev[coord_prev.y +j]) {
				overlap++;
			} else {
				break;
			}
		}
		uint32 n_ovp = (overlap >= kmer_size) ? overlap - kmer_size + 1 : 0;
		return n_ovp;
	}
};

// returns the number of kmers that did not overlap with the lexicographic
// super-kmer predecessor and are not unique
// i.e. total number of kmers in the chain - number of overlapping kmers
// does not count the unique first and last kmers
struct num_non_overlapped_repeats : public thrust::unary_function<uint32,uint32>
{
	const SequenceSetKmerCoord* coords;
	const uint32* suffix_kmer_ids;
	const uint32* overlaps;
	const uint32 kmer_size;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	num_non_overlapped_repeats(const SequenceSetKmerCoord* _coords,
			const uint32* _suffix_kmer_ids,
			const uint32* _overlaps,
			const uint32 _kmer_size):
		coords(_coords), suffix_kmer_ids(_suffix_kmer_ids), overlaps(_overlaps), kmer_size(_kmer_size) { }

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	uint32 operator() (const uint32 idx) const
	{
		const SequenceSetKmerCoord coord = coords[idx];
		uint32 len = coord.z - coord.y;
		const uint32 n_kmers = (len >= kmer_size) ? len - kmer_size + 1 : 0;

		uint32 n_unique = 0;
		if(overlaps[idx] == 0) {
			n_unique++; // skip the header unique kmer if this is the first chain starting with this kmer
						// otherwise, it will be included in the overlap
		}
		if(overlaps[idx] != n_kmers && suffix_kmer_ids[idx] != (uint32) -1) { // if not all the kmers overlapped
			n_unique++; // skip the last unique kmer
		}
		return n_kmers - overlaps[idx] - n_unique;
	}
};

struct num_adj_repeats : public thrust::unary_function<uint32,uint32>
{
	const SequenceSetKmerCoord* coords;
	const uint32* suffix_kmer_ids;
	const uint32* overlaps;
	const uint32 kmer_size;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	num_adj_repeats(const SequenceSetKmerCoord* _coords,
			const uint32* _suffix_kmer_ids,
			const uint32* _overlaps,
			const uint32 _kmer_size):
		coords(_coords), suffix_kmer_ids(_suffix_kmer_ids), overlaps(_overlaps), kmer_size(_kmer_size) { }

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	uint32 operator() (const uint32 idx) const
	{
		const SequenceSetKmerCoord coord = coords[idx];
		uint32 len = coord.z - coord.y;
		const uint32 n_kmers = (len >= kmer_size) ? len - kmer_size + 1 : 0; // TODO: remove len check

		uint32 n_repeat_kmers_to_extract = n_kmers - overlaps[idx];
		if(overlaps[idx] == 0) {
			n_repeat_kmers_to_extract--; // skip the first unique kmer
		}
		if(overlaps[idx] != n_kmers && suffix_kmer_ids[idx] != (uint32) -1) { // if not all the kmers overlapped
			n_repeat_kmers_to_extract--; // skip the last unique kmer
		}

		// number of kmers to extract will be the same as number of edges, plus RU edge if exists
		uint32 n_adj = n_repeat_kmers_to_extract;
		if(overlaps[idx] != n_kmers && suffix_kmer_ids[idx] != (uint32) -1) { // if not all the kmers overlapped
			n_adj++; // add the last R-U edge
		}

		return n_adj;
	}
};

template <typename string_set_type>
struct extract_non_overlapped_repeats
{
	typedef typename string_set_type::string_type sequence;
	const string_set_type   string_set;
	const SequenceSetKmerCoord* super_coords;
	const uint32* super_prefix_global_ids;
	const uint32* super_suffix_global_ids;
	const uint32* chain_offsets;
	const uint32* adj_extraction_offsets;
	const uint32* chain_overlaps;
	const uint32 kmer_size;
	const uint32 uid_offset;
	const uniform_seeds_functor<> seeder;

	SequenceSetKmerCoord* kmers_out;
	uint32* prefix_uids;
	uint32* suffix_uids;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	extract_non_overlapped_repeats(const string_set_type _string_set,
			const SequenceSetKmerCoord* _chain_coords,
			const uint32* _chain_pref_gids, const uint32* _chain_suffix_gids,
			const uint32* _chain_offsets,
			const uint32* _adj_extraction_offsets,
			const uint32* _chain_overlaps,
			const uint32 _kmer_size, const uint32 _uid_offset,
			SequenceSetKmerCoord* _kmers_out, uint32* _prefix_uids, uint32* _suffix_uids):
		string_set(_string_set),
		super_coords(_chain_coords),
		super_prefix_global_ids(_chain_pref_gids),
		super_suffix_global_ids(_chain_suffix_gids),
		chain_offsets(_chain_offsets),
		adj_extraction_offsets(_adj_extraction_offsets),
		chain_overlaps(_chain_overlaps),
		kmer_size(_kmer_size), uid_offset(_uid_offset),
		kmers_out(_kmers_out),
		prefix_uids(_prefix_uids), suffix_uids(_suffix_uids),
		seeder(_kmer_size, 1u) { }

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	void operator() (const uint32 idx) const
	{
		const SequenceSetKmerCoord chain_coord = super_coords[idx];
		const sequence seq = string_set[chain_coord.x];
		const uint32 seq_len = seq.length();
		const uint32 offset = chain_offsets[idx] + uid_offset;
		const uint32 adj_offset = adj_extraction_offsets[idx];

		const uint8 unique_suffix = (super_suffix_global_ids[idx] != (uint32) -1);
		const uint32 n_kmers_to_extract = chain_offsets[idx+1] - chain_offsets[idx];

		uint32 len = chain_coord.z - chain_coord.y;
		const uint32 n_kmers = (len >= kmer_size) ? len - kmer_size + 1 : 0; // TODO: remove len check

		// if everything overlapped or nothing to extract and no possible R-U edge
		if(chain_overlaps[idx] == n_kmers || (n_kmers_to_extract + unique_suffix == 0)) return;

		const uint32 kmer_id = (chain_coord.z - unique_suffix) - n_kmers_to_extract - kmer_size + 1;
		for (uint32 j = 0; j < n_kmers_to_extract; j++) {
			const uint2 kmer = seeder.seed(seq_len, kmer_id + j);
			kmers_out[offset + j] = make_uint4(chain_coord.x, kmer.x, kmer.y, chain_coord.w);

			suffix_uids[adj_offset + j] = offset + j;
			if(j < n_kmers_to_extract - 1) {
				prefix_uids[adj_offset + j + 1] = offset + j;
			} else {
				if(unique_suffix == 1) {
					prefix_uids[adj_offset + j + 1] = offset + j; // the last kmer is a prefix only to the RU (k+1)mer
				}
			}
		}
		if(unique_suffix == 1) {
			suffix_uids[adj_offset + n_kmers_to_extract] = super_suffix_global_ids[idx];
		}

		// find the first prefix
		if(chain_overlaps[idx] == 0 || chain_overlaps[idx] == 1) {
			// overlap is just the unique header kmer or none
			// use the unique id of the kmer
			prefix_uids[adj_offset] = super_prefix_global_ids[idx];
		} else {
			uint32 pred_chain_idx = idx - 1;
			while(1) {
				if(chain_overlaps[pred_chain_idx] >= chain_overlaps[idx]) { // not stored by the previous chain // >=2
					pred_chain_idx--;
				} else {
					// ignores the unique predecessor that overlapped and was not extracted
					prefix_uids[adj_offset] = uid_offset + chain_offsets[pred_chain_idx] + (chain_overlaps[idx] - chain_overlaps[pred_chain_idx]) - 2; // last overlapped kmer
					break;
				}
			}
		}
	}
};

struct count_overlapped_adjacencies
{
	const uint32* adj_extraction_offsets;
	const uint32* overlaps;
	uint32* counts;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	count_overlapped_adjacencies(const uint32* _adj_extraction_offsets, const uint32* _overlaps, uint32* _counts):
		adj_extraction_offsets(_adj_extraction_offsets),
		overlaps(_overlaps),
		counts(_counts) { }

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	void operator() (const uint32 idx) const
	{
		int32 n_to_update = overlaps[idx] - 1;
		if(n_to_update <= 0) return;

		uint32 pred_chain_idx = idx - 1;
		while(n_to_update > 0) {
			const int32 pred_overlap = overlaps[pred_chain_idx] - 1;
			if(pred_overlap >= n_to_update) { // everything is shared with the pred
				pred_chain_idx--;
			} else {
				// update the counts in this pred chain
				const uint32 n_update_pred = n_to_update - (pred_overlap > 0 ? pred_overlap : 0);
				const uint32 pred_offset = adj_extraction_offsets[pred_chain_idx];
				for(uint32 i = 0; i < n_update_pred; i++) {
					nvbio::atomic_add(&counts[pred_offset + i], 1);
				}
				n_to_update -= n_update_pred;
			}
		}
	}
};

struct extract_repeat_adjacencies
{
	const uint32* node_offsets;
	const uint32* prefix_uids;
	const uint32* suffix_uids;
	const uint32* edge_counts;
	uint32* node_adj_map;
	uint32* edge_counts_shuffled;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	extract_repeat_adjacencies(const uint32* _node_offsets, const uint32* _prefix_uids,
			const uint32* _suffix_uids, const uint32* _edge_counts,
			uint32* _node_adj_map, uint32* _edge_counts_shuffled):
		node_offsets(_node_offsets),
		prefix_uids(_prefix_uids),
		suffix_uids(_suffix_uids),
		edge_counts(_edge_counts),
		node_adj_map(_node_adj_map),
		edge_counts_shuffled(_edge_counts_shuffled){ }

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	void operator() (const uint32 prefix_idx) const
	{
		const uint32 from_uid = prefix_uids[prefix_idx];
		const uint32 offset = node_offsets[from_uid];
		const uint32 num_edges = node_offsets[from_uid+1] - node_offsets[from_uid];

		uint32 i = 0;
		while(i < num_edges) { // at most this many prefixes to handle
			if(prefix_idx < i) break;
			const uint32 uid = prefix_uids[prefix_idx - i];
			if(uid != from_uid) break; // handled all the prefixes with the given UID
			node_adj_map[offset + i] = suffix_uids[prefix_idx - i];
			edge_counts_shuffled[offset + i] = edge_counts[prefix_idx - i];
			i++;
		}
	}
};

struct extract_unique_adjacencies
{
	const uint32* node_offsets;
	const uint32* prefix_uids;
	const uint32* suffix_uids;
	const uint32* edge_counts;
	uint32* node_adj_map;
	uint32* edge_counts_shuffled;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	extract_unique_adjacencies(const uint32* _node_offsets, const uint32* _prefix_uids,
			const uint32* _suffix_uids, const uint32* _edge_counts,
			uint32* _node_adj_map, uint32* _edge_counts_shuffled):
		node_offsets(_node_offsets),
		prefix_uids(_prefix_uids),
		suffix_uids(_suffix_uids),
		edge_counts(_edge_counts),
		node_adj_map(_node_adj_map),
		edge_counts_shuffled(_edge_counts_shuffled){ }

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	void operator() (const uint32 prefix_idx) const
	{
		const uint32 from_uid = prefix_uids[prefix_idx];
		const uint32 offset_last = node_offsets[from_uid+1] - 1;
		const uint32 num_edges = node_offsets[from_uid+1] - node_offsets[from_uid];

		uint32 i = 0;
		uint32 n_unique = 0;
		while(i < num_edges) { // at most this many prefixes to handle
			if(prefix_idx < i) break;
			const uint32 uid = prefix_uids[prefix_idx - i];
			if(uid != from_uid) break; // handled all the prefixes with the given UID
			if(suffix_uids[prefix_idx - i] == (uint32) -1) { // repeat (TODO: fill with -1)
				i++;
				continue;
			}
			node_adj_map[offset_last - n_unique] = suffix_uids[prefix_idx - i]; // fill out from the end
			edge_counts_shuffled[offset_last - n_unique] = edge_counts[prefix_idx - i];
			i++;
			n_unique++;
		}
	}
};

// converts a coordinate to the corresponding sequence -- used for debugging
template <typename string_set_type>
struct super_coord_to_seq
{
	typedef typename string_set_type::string_type sequence;
	const string_set_type   string_set;
	const SequenceSetKmerCoord* super_coords;
	const uint32* offsets;
	uint8* sequences;

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	super_coord_to_seq(const string_set_type _string_set,
			const SequenceSetKmerCoord* _super_coords, const uint32* _offsets,
			uint8* _sequences) :
		string_set(_string_set), super_coords(_super_coords), offsets(_offsets),
		sequences(_sequences) {}

	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	void operator() (const uint64 idx) const
	{
		const SequenceSetKmerCoord coord = super_coords[idx];
		const sequence seq = string_set[coord.x];
		const uint32 seq_pos = coord.y;
		const uint32 offset = offsets[idx];

		for (uint32 i = 0; i < coord.z - coord.y; i++) {
			uint8 c = dna_to_char(seq[seq_pos + i]);
			sequences[offset + i] = c;
		}
	}
};

// generate the coordinates of each kmer in the given set
template <typename string_set_type>
void D_KmerSet<string_set_type>::gen_kmer_coords()
{
	n_kmers = enumerate_string_set_seeds(
			string_set,
			uniform_seeds_functor<>(kmer_size, 1u),
			coords);

	// record the index into the kmer array in the coordinates vector
	thrust::for_each(
			thrust::make_counting_iterator<uint32>(0u),
			thrust::make_counting_iterator<uint32>(0u) + n_kmers,
			set_global_id_region_id(plain_view(coords), plain_view(active_region_ids)));
}

// generate the 64-bit compacted kmer sequence keys
// assumes the coordinates have already been generated
template <typename string_set_type>
void D_KmerSet<string_set_type>::gen_kmer_64b_keys()
{
	kmers_64b.resize(n_kmers);
	thrust::transform(
			coords.begin(),
			coords.begin() + n_kmers,
			kmers_64b.begin(),
			kmers_to_64b_functor<string_set_type>(kmer_size, 2, string_set));
}

// sort the kmers based on the 64-bit kmer sequence key only
// assumes the keys have already been generated
template <typename string_set_type>
void D_KmerSet<string_set_type>::sort_kmers_by_64b_keys()
{
	// identical keys from the same sequence will be consecutive
	// identical keys from the same region will also be consecutive
	thrust::stable_sort_by_key(
				kmers_64b.begin(),
				kmers_64b.begin() + n_kmers,
				coords.begin());
}

// stable sort the kmers based on the 64-bit kmer sequence key
// assumes the kmers are sorted by seq id only already
// assumes the keys have already been generated
template <typename string_set_type>
template <typename meta_iterator_type>
void D_KmerSet<string_set_type>::sort_kmers_by_64b_keys_meta(const meta_iterator_type meta_data)
{
	thrust::stable_sort_by_key(
			kmers_64b.begin(),
			kmers_64b.begin() + n_kmers,
			thrust::make_zip_iterator(thrust::make_tuple(coords.begin(), meta_data)));
}

// sort the kmers based on the 64-bit kmer sequence key by region
// assumes the keys have already been generated
template <typename string_set_type>
void D_KmerSet<string_set_type>::segmented_sort_kmers_by_64b_keys()
{
	// identical keys from the same sequence will be consecutive
	// identical keys from the same region will also be consecutive
	thrust::stable_sort_by_key(
				kmers_64b.begin(),
				kmers_64b.begin() + n_kmers,
				coords.begin());

	D_VectorU32::iterator region_ids = get_scratch_space(n_kmers);
	thrust::transform(
			coords.begin(),
			coords.begin() + n_kmers,
			region_ids,
			get_kmer_reg_id());

	thrust::stable_sort_by_key(
			region_ids,
			region_ids + n_kmers,
			thrust::make_zip_iterator(thrust::make_tuple(coords.begin(), kmers_64b.begin())));

	reset_scratch_space();
}

// sort kmers by key and set sequence id
// doing 2-pass sort to use the much faster thrust radix-sort
// instead of merge-sort on custom keys
template <typename string_set_type>
void D_KmerSet<string_set_type>::sort_kmers_by_64b_keys_seqid()
{
	thrust::stable_sort_by_key(
			kmers_64b.begin(),
			kmers_64b.begin() + n_kmers,
			coords.begin());
	D_VectorU32 seq_ids(n_kmers);
	thrust::transform(
			coords.begin(),
			coords.begin() + n_kmers,
			seq_ids.begin(),
			get_kmer_seq_id());
	thrust::stable_sort_by_key(
			seq_ids.begin(),
			seq_ids.begin() + n_kmers,
			thrust::make_zip_iterator(thrust::make_tuple(kmers_64b.begin(), coords.begin())));
}

// sort kmers by key and set sequence id
// sort the additional meta-data along with the kmers
// doing 2-pass sort to use the much faster thrust radix-sort
// instead of merge-sort on custom keys
template <typename string_set_type>
template <typename meta_iterator_type>
void D_KmerSet<string_set_type>::sort_kmers_by_64b_keys_seqid_meta(const meta_iterator_type meta_data)
{
	thrust::stable_sort_by_key(
			kmers_64b.begin(),
			kmers_64b.begin() + n_kmers,
			thrust::make_zip_iterator(thrust::make_tuple(coords.begin(), meta_data)));
	D_VectorU32 seq_ids(n_kmers);
	thrust::transform(
			coords.begin(),
			coords.begin() + n_kmers,
			seq_ids.begin(),
			get_kmer_seq_id());
	thrust::stable_sort_by_key(
			seq_ids.begin(),
			seq_ids.begin() + n_kmers,
			thrust::make_zip_iterator(thrust::make_tuple(kmers_64b.begin(), coords.begin(), meta_data)));
}

// sort kmers in lexicographic order - sorts the indices and returns the iterator
// kmer length can vary
// if kmer_size is set, all kmers are assumed to have this size
template <typename string_set_type>
D_VectorU32::iterator sort_kmers_lexicographic(const string_set_type string_set,
		D_VectorSetKmerCoord& coords,
		const uint32 n_kmers, const uint32 fixed_kmer_size,
		D_VectorU32& indices)
{
	const uint32 WORD_BITS = uint32(8u * sizeof(uint32));
	const uint32 SYMBOL_SIZE = 4u;
	const uint32 SYMBOLS_PER_WORD = WORD_BITS/SYMBOL_SIZE;

	// maximum size of a kmer in the set
	uint32 max_length = fixed_kmer_size;
	if(fixed_kmer_size == 0u) {
		max_length = thrust::reduce(
	            thrust::make_transform_iterator(coords.begin(), get_kmer_size()),
	            thrust::make_transform_iterator(coords.begin() + n_kmers, get_kmer_size()),
	            0u,
	            thrust::maximum<uint32>());
	}

	// maximum number of words needed to represent a kmer
	const uint32 max_words = (max_length + SYMBOLS_PER_WORD-1)/SYMBOLS_PER_WORD;

	// ----- LSD radix-sort (word by word) -----
	//uint32 n_kmers = coords.size();
	D_VectorU32 radices(2*n_kmers); // kmer words (sort keys)
	indices.resize(2*n_kmers); // kmer ids (values to be sorted)
	thrust::copy(
			thrust::make_counting_iterator<uint32>(0u),
			thrust::make_counting_iterator<uint32>(n_kmers),
	        indices.begin());

	// ping-pong buffers
	cuda::SortBuffers<uint32*,uint32*> sort_buffers;
	cuda::SortEnactor sort_enactor;
	sort_buffers.selector = 0;
	sort_buffers.keys[0] = nvbio::device_view(radices);
	sort_buffers.keys[1] = nvbio::device_view(radices) + n_kmers;
	sort_buffers.values[0] = nvbio::device_view(indices);
	sort_buffers.values[1] = nvbio::device_view(indices) + n_kmers;

	// sort in LSD order
	for (int32 word_idx = max_words-1; word_idx >= 0; --word_idx) {
		// extract the given radix word from each kmer
		thrust::transform(
				indices.begin() + sort_buffers.selector * n_kmers,
				indices.begin() + sort_buffers.selector * n_kmers + n_kmers,
				radices.begin() + sort_buffers.selector * n_kmers,
				kmer_word_extractor_functor<string_set_type>(
		                string_set,
		                SYMBOL_SIZE,
		                plain_view(coords),
		                word_idx));
		// sort the words
		sort_enactor.sort(n_kmers, sort_buffers);
	}

	return indices.begin() + sort_buffers.selector * n_kmers;
}

template <typename string_set_type>
void segmented_sort_super_kmers_lexicographic(const string_set_type string_set, uint32 n_coords,
		D_VectorSetKmerCoord& super_coords, D_VectorU32& super_prefix_ids, D_VectorU32& super_suffix_ids)
{
	D_VectorU32 ids(n_coords);
	D_VectorSetKmerCoord temp(super_coords);
	D_VectorU32 temp_pref(super_prefix_ids);
	D_VectorU32 temp_suf(super_suffix_ids);

	// first pass: sort by coord key
	D_VectorU32::iterator sorted_ids = sort_kmers_lexicographic(string_set, super_coords, n_coords, 0u, ids);
	thrust::gather(
			sorted_ids,
			sorted_ids + n_coords,
			temp.begin(),
			super_coords.begin());
	thrust::gather(
			sorted_ids,
			sorted_ids + n_coords,
			thrust::make_zip_iterator(thrust::make_tuple(temp_pref.begin(), temp_suf.begin())),
			thrust::make_zip_iterator(thrust::make_tuple(super_prefix_ids.begin(), super_suffix_ids.begin())));

	// second pass: stable sort by region id => as a result coord keys will be sorted by region
	//D_VectorU32 reg_ids(n_coords);
	thrust::transform(
			super_coords.begin(),
			super_coords.begin() + n_coords,
			ids.begin(),
			get_kmer_reg_id());

	thrust::stable_sort_by_key(
			ids.begin(),
			ids.begin() + n_coords,
			thrust::make_zip_iterator(thrust::make_tuple(super_coords.begin(), super_prefix_ids.begin(), super_suffix_ids.begin())));
}

// count the kmer occurrences across all the sequences in the set
// assumes that the keys have been sorted
template <typename string_set_type>
void D_KmerSet<string_set_type>::count_kmers_rle()
{
	// count the number of distinct keys and extract them
	kmers_64b_distinct.resize(n_kmers);
	kmers_counts.resize(n_kmers);
	D_VectorU8 d_temp_storage;
	n_distinct = cuda::runlength_encode(
			n_kmers,
			kmers_64b.begin(),
			kmers_64b_distinct.begin(),
			kmers_counts.begin(),
			d_temp_storage);
}

// count the kmer occurrences across all the sequences in the set
// assumes that the keys have been sorted
template <typename string_set_type>
void D_KmerSet<string_set_type>::count_kmers()
{
	// count the number of distinct keys and extract them
	kmers_64b_unique_idxs.resize(n_kmers);
	kmers_counts.resize(n_kmers);
	n_unique = thrust::reduce_by_key(
			thrust::counting_iterator<uint32>(0),
			thrust::counting_iterator<uint32>(0) + n_kmers,
			thrust::constant_iterator<uint32>(1),
			kmers_64b_unique_idxs.begin(),
			kmers_counts.begin(),
			kmer_key_rid_eq(plain_view(kmers_64b), plain_view(coords))).first - kmers_64b_unique_idxs.begin();
}

// partition kmers by index into the sorted kmer key set
// based on whether they are unique or not
// a kmer is not unique if it occurs more than once in at least one sequence
template <typename string_set_type>
void D_KmerSet<string_set_type>::partition_kmers_by_uniqueness()
{
	// count the number of distinct kmers per sequence
	D_VectorU32::iterator distinct_idxs_per_seq = get_scratch_space(n_kmers);
	D_VectorU32::iterator count_per_seq = get_scratch_space(n_kmers);
	thrust::counting_iterator<uint64> ids(0);
	uint64 n_distinct_per_seq = thrust::reduce_by_key(
			ids,
			ids + n_kmers,
			thrust::constant_iterator<uint32>(1),
			distinct_idxs_per_seq,
			count_per_seq,
			kmer_key_sid_eq(plain_view(kmers_64b), plain_view(coords))).first - distinct_idxs_per_seq;

	D_VectorU32::iterator distinct_idxs = get_scratch_space(n_distinct_per_seq);
	D_VectorU32::iterator seq_count = get_scratch_space(n_distinct_per_seq); // number of sequences containing the same kmer
	D_VectorU32::iterator kmer_counts = get_scratch_space(n_distinct_per_seq); // number of kmer occurrences across all sequences
	n_distinct = thrust::reduce_by_key(
			distinct_idxs_per_seq,
			distinct_idxs_per_seq + n_distinct_per_seq,
			thrust::make_zip_iterator(thrust::make_tuple(thrust::constant_iterator<uint32>(1), count_per_seq)),
			distinct_idxs,
			thrust::make_zip_iterator(thrust::make_tuple(seq_count, kmer_counts)),
			kmer_key_rid_eq(plain_view(kmers_64b), plain_view(coords)),
			kmer_count_tuple_sum()).first - distinct_idxs;

	// partition the distinct indices into unique and non-unique kmers
	kmers_64b_unique_idxs.resize(n_distinct);
	kmers_counts.resize(n_distinct);
	n_unique =	thrust::copy_if(
			thrust::make_zip_iterator(thrust::make_tuple(distinct_idxs, kmer_counts)),
			thrust::make_zip_iterator(thrust::make_tuple(distinct_idxs + n_distinct, kmer_counts + n_distinct)),
			thrust::make_zip_iterator(thrust::make_tuple(seq_count, kmer_counts)),
			thrust::make_zip_iterator(thrust::make_tuple(kmers_64b_unique_idxs.begin(), kmers_counts.begin())),
			is_unique_kmer()) - thrust::make_zip_iterator(thrust::make_tuple(kmers_64b_unique_idxs.begin(), kmers_counts.begin()));

	reset_scratch_space();
}

template <typename string_set_type>
void D_KmerSet<string_set_type>::gen_global_unique_map()
{
	D_VectorU8 unique_map_sorted(n_kmers);
	thrust::for_each(
				kmers_64b_unique_idxs.begin(),
				kmers_64b_unique_idxs.begin() + n_unique,
				mark_kmer_uniqueness(plain_view(unique_map_sorted), plain_view(kmers_64b), plain_view(coords)));

	global_unique_flag_map.resize(n_kmers);
	thrust::gather(
				global_to_sorted_id_map.begin(),
				global_to_sorted_id_map.begin() + n_kmers,
				unique_map_sorted.begin(),
				global_unique_flag_map.begin());
}

template <typename string_set_type>
void D_KmerSet<string_set_type>::gen_global_UID_map()
{
	D_VectorU32 id_map_sorted(n_kmers);
	thrust::fill(id_map_sorted.begin(), id_map_sorted.begin() + n_kmers, (uint32) -1);
	thrust::for_each(
				thrust::make_zip_iterator(thrust::make_tuple(thrust::make_counting_iterator<uint32>(0u), kmers_64b_unique_idxs.begin())),
				thrust::make_zip_iterator(thrust::make_tuple(thrust::make_counting_iterator<uint32>(0u) + n_unique, kmers_64b_unique_idxs.begin() + n_unique)),
				store_kmer_unique_ids(plain_view(id_map_sorted), plain_view(kmers_64b), plain_view(coords)));

	global_to_UID_map.resize(n_kmers);
	thrust::gather(
				global_to_sorted_id_map.begin(),
				global_to_sorted_id_map.begin() + n_kmers,
				id_map_sorted.begin(),
				global_to_UID_map.begin());
}

template <typename string_set_type>
void D_KmerSet<string_set_type>::gen_global_to_sorted_id_map()
{
	global_to_sorted_id_map.resize(n_kmers);
	thrust::for_each(
			thrust::make_zip_iterator(thrust::make_tuple(thrust::make_counting_iterator<uint32>(0u), coords.begin())),
			thrust::make_zip_iterator(thrust::make_tuple(thrust::make_counting_iterator<uint32>(0u) + n_kmers, coords.begin() + n_kmers)),
			global_to_sorted_id(plain_view(global_to_sorted_id_map)));
}

template <typename string_set_type>
void D_KmerSet<string_set_type>::mark_unique_kmers()
{
	global_unique_flag_map.resize(n_kmers);
	global_to_UID_map.resize(n_kmers);
	thrust::fill(global_to_UID_map.begin(), global_to_UID_map.begin() + n_kmers, (uint32) -1);
	thrust::for_each(
			thrust::make_counting_iterator<uint32>(0u),
			thrust::make_counting_iterator<uint32>(0u) + n_unique,
			populate_unique_kmer_data(plain_view(kmers_counts), plain_view(coords), plain_view(kmers_64b_unique_idxs),
					plain_view(global_unique_flag_map), plain_view(global_to_UID_map)));
}


struct test_funct : public thrust::unary_function<uint32,uint32>
{
	NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
	uint32 operator()(const uint32 uid)
	{
		return uid;
	}
};

// only keep the coordinates that have a unique prefix
template <typename string_set_type>
void D_KmerSet<string_set_type>::filter_coords_by_prefix_uniqueness(const D_VectorU8& unique_map)
{
	D_VectorSetKmerCoord unique_pref_coords(n_kmers);
	uint32 n_unique_coords = thrust::copy_if( // remove_if
			coords.begin(),
			coords.begin() + n_kmers,
			thrust::make_permutation_iterator(
					unique_map.begin(),
					thrust::make_transform_iterator(coords.begin(), get_prefix_global_id())),
			unique_pref_coords.begin(),
			thrust::identity<uint8>()) - unique_pref_coords.begin();

	thrust::copy(unique_pref_coords.begin(), unique_pref_coords.begin() + n_unique_coords, coords.begin());
	n_kmers = n_unique_coords;
}

// for every kmer that is not unique N
// find its closest unique predecessor P and successor S kmer in the same sequence
// and output a kmer chain (P ->... N ... -> S)
// as a new 'super' kmer coordinate
template <typename string_set_type>
void D_KmerSet<string_set_type>::extract_super_kmers()
{
	// collect the ids of the repeat kmers
	kmers_64b_repeat_idxs.resize(n_kmers);
	uint32 n_chains = thrust::copy_if(
			thrust::make_counting_iterator<uint32>(0u),
			thrust::make_counting_iterator<uint32>(0u) + n_kmers,
			global_unique_flag_map.begin(),
			kmers_64b_repeat_idxs.begin(),
			thrust::logical_not<uint8>()) - kmers_64b_repeat_idxs.begin();

	D_VectorSetKmerCoord super_kmers(n_chains);
	D_VectorU32 prefix_uids(n_chains);
	D_VectorU32 suffix_uids(n_chains);
	thrust::for_each(
			thrust::make_counting_iterator(0),
			thrust::make_counting_iterator(0) + n_chains,
			extract_super_kmers_functor<string_set_type>(
					plain_view(kmers_64b_repeat_idxs),
					plain_view(global_unique_flag_map),
					plain_view(global_to_UID_map),
					plain_view(global_to_sorted_id_map),
					plain_view(coords),
					string_set,
					kmer_size,
					plain_view(super_kmers),
					plain_view(prefix_uids),
					plain_view(suffix_uids)));

	// collapse chains that belong to the same sequence and same unique kmer predecessor
	// (these will be consecutive by construction)
	// keep the max chain length coord per same kmer and seq id
	// -- will keep one copy of the chains that is terminated with a unique kmer
	// and the longest chain amongst the ones that are not
	super_coords.resize(n_chains);
	super_prefix_uids.resize(n_chains);
	super_suffix_uids.resize(n_chains);
	uint32 n_collapsed_chains = thrust::reduce_by_key(
				thrust::make_counting_iterator(0u),
				thrust::make_counting_iterator(0u) + n_chains,
				thrust::make_zip_iterator(thrust::make_tuple(super_kmers.begin(), prefix_uids.begin(), suffix_uids.begin())),
				thrust::make_discard_iterator(),
				thrust::make_zip_iterator(thrust::make_tuple(super_coords.begin(), super_prefix_uids.begin(), super_suffix_uids.begin())),
				kmer_pos_sid_eq(plain_view(super_kmers)),
				collapse_same_start_kmers()).second - thrust::make_zip_iterator(thrust::make_tuple(super_coords.begin(), super_prefix_uids.begin(), super_suffix_uids.begin()));

	// eliminate non-unique kmers that are not preceded by a unique kmer
	D_VectorSetKmerCoord t1(n_chains);
	D_VectorU32 t2(n_chains);
	D_VectorU32 t3(n_chains);
	n_super_coords = thrust::remove_copy_if(
			thrust::make_zip_iterator(thrust::make_tuple(
					super_coords.begin(),
					super_prefix_uids.begin(),
					super_suffix_uids.begin())),
			thrust::make_zip_iterator(thrust::make_tuple(
					super_coords.begin() + n_collapsed_chains,
					super_prefix_uids.begin() + n_collapsed_chains,
					super_suffix_uids.begin() + n_collapsed_chains)),
			super_prefix_uids.begin(),
			thrust::make_zip_iterator(thrust::make_tuple(
					t1.begin(),
					t2.begin(),
					t3.begin())),
			is_invalid_super_kmer()) - thrust::make_zip_iterator(thrust::make_tuple(
					t1.begin(),
					t2.begin(),
					t3.begin()));


	thrust::copy(t1.begin(), t1.begin() + n_chains, super_coords.begin());
	thrust::copy(t2.begin(), t2.begin() + n_chains, super_prefix_uids.begin());
	thrust::copy(t3.begin(), t3.begin() + n_chains, super_suffix_uids.begin());

	//super_coords.erase(super_coords.begin() + n_super_coords, super_coords.end());
	//super_prefix_uids.erase(super_prefix_uids.begin() + n_super_coords, super_prefix_uids.end());
	//super_suffix_uids.erase(super_suffix_uids.begin() + n_super_coords, super_suffix_uids.end());

//	printf("Collapsed Filtered Super Kmer Coords: \n");
//	for(uint64 i = 0; i <  n_collapsed_chains_filtered; i++) {
//		SequenceSetKmerCoord c = (SequenceSetKmerCoord) super_coords[i];
//		printf("coord [%u %u %u %u] pref %llu suf %llu\n", c.x, c.y, c.z, c.w, (uint64) super_prefix_global_ids[i],
//				(uint64) super_suffix_global_ids[i]);
//	}

	reset_scratch_space();
}

// collapses lexicographically sorted chains of kmers by computing their prefix overlaps
// and extracting kmers not included in the overlapping region
// along with kmer adjacencies and adjacency count information
template <typename string_set_type>
void D_KmerSet<string_set_type>::collapse_and_extract_non_overlaps(D_VectorSetKmerCoord& kmers_out,
		D_VectorU32& prefix_ids_out, D_VectorU32& suffix_ids_out, D_VectorU32& adj_counts)
{
	uint32 n_super_kmers = n_super_coords;

	// find overlap with the previous chain (in sorted order)
	// overlap is measured in # of shared prefix kmers
	D_VectorU32 max_predecessor_overlaps(n_super_kmers);
	thrust::transform(
			thrust::make_counting_iterator<uint32>(1u), // skip the first chain
			thrust::make_counting_iterator<uint32>(0u) + n_super_kmers,
			max_predecessor_overlaps.begin() + 1,
			find_max_kmer_overlaps<string_set_type>(string_set, plain_view(super_coords), kmer_size));

	// find the number of kmer nodes that need to be created
	// the first kmer in a chain is a unique kmer and should be ignored
	// the last kmer in the chain might be unique, in which case it is also ignored
	thrust::device_vector<uint8> temp_storage;
	D_VectorU32 chain_extraction_offsets(n_super_kmers + 1);
	cuda::inclusive_scan(
			n_super_kmers,
			thrust::make_transform_iterator(
					thrust::make_counting_iterator<uint32>(0u),
					num_non_overlapped_repeats(
							plain_view(super_coords),
							plain_view(super_suffix_uids),
							plain_view(max_predecessor_overlaps), kmer_size)),
			chain_extraction_offsets.begin() + 1,
			thrust::plus<uint32>(),
			temp_storage);

	D_VectorU32 adj_extraction_offsets(n_super_kmers + 1);
	cuda::inclusive_scan(
			n_super_kmers,
			thrust::make_transform_iterator(
					thrust::make_counting_iterator<uint32>(0u),
					num_adj_repeats(
							plain_view(super_coords),
							plain_view(super_suffix_uids),
							plain_view(max_predecessor_overlaps), kmer_size)),
			adj_extraction_offsets.begin() + 1,
			thrust::plus<uint32>(),
			temp_storage);

	// extract non-unique kmers and the (k+1)mer prefix and suffix kmer ids from the collapsed chains
	n_repeat = chain_extraction_offsets[n_super_kmers];
	uint32 n_prefsuf_to_extract = adj_extraction_offsets[n_super_kmers];

	//printf("n_super_coords %u \n", n_super_coords);
	//printf("n_repeat %u \n", n_repeat);
	//printf("Num X-R-X edges %u \n", n_prefsuf_to_extract);

	// allocate the necessary storage and populate
	kmers_out.resize(n_unique + n_repeat);
	prefix_ids_out.resize(n_prefsuf_to_extract);
	suffix_ids_out.resize(n_prefsuf_to_extract);
	thrust::for_each(
			thrust::make_counting_iterator<uint32>(0u),
			thrust::make_counting_iterator<uint32>(0u) + n_super_kmers,
			extract_non_overlapped_repeats<string_set_type>(
					string_set,
					plain_view(super_coords),
					plain_view(super_prefix_uids),
					plain_view(super_suffix_uids),
					plain_view(chain_extraction_offsets),
					plain_view(adj_extraction_offsets),
					plain_view(max_predecessor_overlaps),
					kmer_size, n_unique,
					plain_view(kmers_out),
					plain_view(prefix_ids_out),
					plain_view(suffix_ids_out)));


	// count the number of repeated kmer adjacencies
    adj_counts.resize(n_prefsuf_to_extract);
	thrust::fill(adj_counts.begin(), adj_counts.begin() + n_prefsuf_to_extract, 1);
	thrust::for_each(
			thrust::make_counting_iterator<uint32>(1u),
			thrust::make_counting_iterator<uint32>(0u) + n_super_kmers,
			count_overlapped_adjacencies(
					plain_view(adj_extraction_offsets),
					plain_view(max_predecessor_overlaps),
					plain_view(adj_counts)));

}