//! [mapsembler_extensions2]
//
// code inspired by assembly_forensics from gatb-tools

#if UNORDERED_MAP_TR1_FOUND == 1
//#if __clang__
#include <tr1/unordered_map>
#else
#include <unordered_map>
#endif


#include <mapsembler2_extremities.hpp>
#include <set>
#include <string>


#if UNORDERED_MAP_TR1_FOUND == 1
//#if ! __clang__
using namespace tr1;
#endif
using namespace std;

// global variables because yeah I code like that
size_t kmerSize;


#if UNORDERED_MAP_TR1_FOUND == 1
std::tr1::unordered_map<string, int> extremities_abundances; // init a hash table for recording abundance of extension
#else
std::unordered_map<string, int> extremities_abundances; // init a hash table for recording abundance of extension
#endif

int debug = 0;
int min_solid_subkmer; // by default, require that a k-mer be seen >= 3 times to be a sub-kmer;
FILE *output_file;

/********************************************************************************/

// some explicit kmer normalization code, I know GATB has those, but I'm too lazy to look up the docs (Erwan please please make the doxygen searchable)

string reverse_complement(string s)
{
    string t;
    for (int i = 0; i < s.size(); i ++)
    {
        switch (s[i])
        {
            case 'A':
                t += "T"; break;
            case 'T':
                t += "A"; break;
            case 'G':
                t += "C"; break;
            case 'C':
                t += "G"; break;
        }
    }
    return string ( t.rbegin(), t.rend() );
}

string normalize(string s)
{
    string t = reverse_complement(s);
    return (t < s) ? t : s;
}


/********************************************************************************/


// We take a k+1-mer s and simulate all possibilities from the three events:
// 
// - there is a substitution in one of the last k characters of s
// - there is a deletion at one of the characters of s
//   *except at the very end if it's a right kmer (don't want to get the previous kmer)*
//   *except at the very beginning if it's a left kmer (don't want to get the next kmer)*
//   *except if the deleted character makes the resulting k-mer exactly identical to next/previous kmer to the extremal kmer*
// - there is an insertion in of the last k-1 characters of s, 
//   *except at the very end if it's a right kmer (don't want to extend that kmer)*
//   *except at the very beginning if it's a left kmer (don't want to extend that kmer)*
//
// For each possibility extract the prefix (if it's a left kmer) or suffix (if it's a right kmer) 
// of length k of the string resulting from the simulated event.
//
// So the resulting strings are not exactly of edit distance 1 of s, 
// (but one can see that there does exists an intermediate string that is of edit distance 1).
// Furthermore all returned strings are distinct.
set<string> all_kmers_at_distance_1(string s, bool is_first_kmer = true)
{
    set<string> res;
    int size = s.size();
    string alphabet = "ACTG";
    int a = alphabet.size();
    int offset = is_first_kmer ? 0 : 1;

    int start_extremal_kmer = offset;
    int end_extremal_kmer = size - (1-offset); // immediately after the last character

    // substitutions
    for (int i = start_extremal_kmer; i < end_extremal_kmer; i++)
    {
        for (int j = 0; j < a; j++)
        {
            char c = alphabet[j];
            if (s[i] == c)
                continue;
            string t(s);
            t[i] = c;
            res.insert(t.substr(offset, kmerSize));
            //printf("substitution of %s: %s\n",s.c_str(),t.substr(offset, kmerSize).c_str());
        }
    }
    
    // deletions
    for (int i = 1; i < size - 1; i++)
    {
        string t = s.substr(0, i) + s.substr(i+1, size - i - 1);

        // make sure the resulting kmer is not identical to the next/previous kmer to the extremal kmer
        if (is_first_kmer && s.substr(1, size-1).compare(t) == 0)
            continue;
        if ((!is_first_kmer) && s.substr(0, size-1).compare(t) == 0)
            continue;
        
        res.insert(t);
        //printf("deletion of %s: %s\n",s.c_str(),t.c_str());
    }

    // insertions
    for (int i = 1; i < size - 1; i++)
    {
        for (int j = 0; j < a; j++)
        {
            char c = alphabet[j];
            string u(1, c);
            string t = s.substr(0, i) + u + s.substr(i+1, size - i - 1);
            res.insert(t.substr(offset, kmerSize));
            //printf("insertion of %s: %s\n",s.c_str(),t.substr(offset, kmerSize).c_str());
        }
    }

    return res;
}

/********************************************************************************/

void write_starter(string starter_seq, string starter_name){
    fprintf(output_file,">%s\n%s\n", starter_name.c_str(), starter_seq.c_str());
}

void write_subkmer(string subkmer, string direction, int subkmer_number, int abundance)
{
    fprintf(output_file,">extrem_%s|id_%d|abundance_%d\n%s\n",\
            direction.c_str(), subkmer_number, abundance, subkmer.c_str());
}

void process_starters(string starters, bool init = true)
{
    printf("processing starters file %s %s\n", starters.c_str(), init ? "(init)" : "");


    // We declare a Bank instance defined by a list of filenames
    BankFasta b (starters);

    // We create a sequence iterator for the bank
    BankFasta::Iterator* itSeq = new BankFasta::Iterator (b);

    //  We create a sequence iterator that notifies listeners every N sequences
    SubjectIterator<Sequence> iter (itSeq, 1000);

    // for each contig, hash its extremities
    for (iter.first(); !iter.isDone(); iter.next())
    {
        string seq = iter->toString();
        string starter_name = iter->getComment();

        if (seq.size() < kmerSize)
        {
            printf("starter %s is shorter than kmer size, ignored\n", starter_name.c_str());
            continue;
        }

        // get first and last kmers
        string first_kplusonemer_raw = seq.substr(0, kmerSize + 1);
        string last_kplusonemer_raw = seq.substr(seq.size() - kmerSize - 1, kmerSize + 1);

        string both[2] = {first_kplusonemer_raw, last_kplusonemer_raw};
        
        if (!init) write_starter(seq, starter_name);
        for (int i = 0; i < 2; i++) // i=0: first kmer, i=1: last kmer
        {
            // get mutated extremal kmers
            set<string> subkmers = all_kmers_at_distance_1(both[i], i == 0);

            // also add non-mutated versions
            if (i == 0)
                subkmers.insert(seq.substr(0, kmerSize));
            else
                subkmers.insert(seq.substr(seq.size() - kmerSize, kmerSize));

            for (set<string>::iterator it = subkmers.begin(); it != subkmers.end(); it++)
            {
                string subkmer = *it;
                //printf("processing subkmer: %s\n",subkmer.c_str());
                string normalized_subkmer = normalize(subkmer);

                int solid_subkmer_number = 0;

                if (init)
                {
                    extremities_abundances[normalized_subkmer] = 0;
                    //printf("adding kmer: %s to extremities abundances\n",normalized_subkmer.c_str());
                }
                else
                {
                    if (extremities_abundances[normalized_subkmer] >= min_solid_subkmer)
                    {
                        write_subkmer(subkmer, i == 0 ? "left" : "right", solid_subkmer_number++, extremities_abundances[normalized_subkmer]);
                    }
                }

            }
        }          
   }
}



void process_reads(vector<string> read_files)
{
    printf("reading reads file(s): ");
    for (unsigned n=0; n<read_files.size(); ++n)
        printf("%s ",read_files.at( n ).c_str());
    printf("\n");
    
    // We declare a Bank instance defined by a list of filenames
    BankFasta b (read_files);
    BankFasta::Iterator* itSeq = new BankFasta::Iterator (b);
    SubjectIterator<Sequence> iter (itSeq, 100000);
    iter.addObserver (new ProgressTimer (b.estimateNbItems(), "Iterating reads"));

    for (iter.first(); !iter.isDone(); iter.next())
    {
        string seq = iter->toString();
        int size = seq.size();

        if (size < kmerSize)
            continue; // normally this check is useless due to the next for.. but that's normally.. (when a read is smaller than k, I get wrong behavior without that test, so that's a TODO; e.g. in the following test file /scratch/reads/ion-torrent/blue-corrected/minia/sga_test/all_tags.fasta on rayan-1)

        for (int i = 0; i < (size - kmerSize + 1); i++)
        {
            string kmer = normalize(seq.substr(i, kmerSize));
            if (extremities_abundances.find(kmer) == extremities_abundances.end())
                continue;
            //printf("found kmer: %s in extremities abundances\n",kmer.c_str());
            extremities_abundances[kmer] ++;
        } 
    }
}






// We define some constant strings for names of command line parameters
static const char* STR_FOO = "-foo";

/*********************************************************************
** METHOD  :
** PURPOSE :
** INPUT   :
** OUTPUT  :
** RETURN  :
** REMARKS :
*********************************************************************/
mapsembler2_extremities::mapsembler2_extremities ()  : Tool ("mapsembler2_extremities")
{
    OptionsParser parser ("mapsembler2_extremities");
    parser.push_front (new OptionOneParam ("--k", "kmer size that will be used for mapsembler2", true));
    parser.push_front (new OptionOneParam ("--starters", "starters fasta file", true));
    parser.push_front (new OptionOneParam ("--reads", "reads dataset file name. Several reads sets can be provided, surrounded by \"'s and separated by a space (e.g. --reads \"reads1.fa reads2.fa\")", true));
    parser.push_front (new OptionOneParam ("--output", "output substarters file name", true));
    parser.push_front (new OptionOneParam ("--min-solid-subkmer", "minimim abundance to keep a subkmer",  false, "3"));
    parser.push_front (new OptionNoParam ("-debug", "debugging", false));
    getParser()->push_front (parser);
}


/*********************************************************************
 ** METHOD  :
 ** PURPOSE : Split a string into a vector of strings wrt to a delimiter
 ** INPUT   :
 ** OUTPUT  :
 ** RETURN  :
 ** REMARKS :
 *********************************************************************/
vector<string> split(const string &s, char delim) {
    vector<std::string> elems;
    stringstream ss(s);
    string item;
    while (std::getline(ss, item, delim)) {
        elems.push_back(item);
    }
    return elems;
}

/*********************************************************************
** METHOD  :
** PURPOSE :
** INPUT   :
** OUTPUT  :
** RETURN  :
** REMARKS :
*********************************************************************/
void mapsembler2_extremities::execute ()
{
    // We can do here anything we want.
    // For further information about the Tool class, please have a look
    // on the ToyTool snippet  (http://gatb-core.gforge.inria.fr/snippets_tools.html)

    // We gather some statistics.
    getInfo()->add (1, "input");
    getInfo()->add (1, &LibraryInfo::getInfo());

    /** we get the kmer size chosen by the end user. */
    kmerSize = getInput()->getInt ("--k");
    
    string starters_file = getInput()->getStr("--starters");
    string reads_file = getInput()->getStr("--reads");
    string output_filename = getInput()->getStr("--output");
    min_solid_subkmer = getInput()->getInt("--min-solid-subkmer");
    debug = getParser()->saw("--debug");

    output_file = fopen(output_filename.c_str(),"w");

    process_starters(starters_file, true);
    
    process_reads(split(reads_file, ' '));
    process_starters(starters_file, false);

    fclose(output_file);
}
