#include "htio2/Cast.h"
#include "htio2/FastqIO.h"
#include "htio2/FastqSeq.h"
#include "htio2/PairedEndIO.h"
#include "htio2/PlainFileHandle.h"
#include "htio2/HeaderUtil.h"

#define USE_header_format
#define USE_se_pe
#include "htqc/Options.h"
#include "htqc/App.h"
#include "htqc/MultiSeqFile.h"

#include "htio2/JUCE-3.0.8/JuceHeader.h"

#include <set>

// We don't care strand in this program, so simply let them identical.
#define PE_STRAND_FILE htio2::STRAND_FWD, htio2::STRAND_FWD
#define PE_STRAND_OBJ htio2::STRAND_FWD, htio2::STRAND_FWD

using namespace std;
using namespace htqc;
using namespace htio2::juce;

string OPT_dir_out;
htio2::Option OPT_dir_out_ENTRY("out", 'o', OPTION_GROUP_IO,
                                &OPT_dir_out, htio2::Option::FLAG_NONE,
                                "Output directory.", "DIR");
string OPT_file_barcode;
htio2::Option OPT_file_barcode_ENTRY("barcode", '\0', OPTION_GROUP_IO,
                                     &OPT_file_barcode, htio2::Option::FLAG_NONE,
                                     "Tab-delimited file containing three columns which are project, sample and barcode. Output will be written to output_dir/project/sample_[12].fastq",
                                     "FILE");

int OPT_mismatch(0);
htio2::Option OPT_mismatch_ENTRY("mismatch", '\0', "Algorithm Parameters",
                                 &OPT_mismatch, htio2::Option::FLAG_NONE,
                                 "Allowed number of barcode mismatch.", "INT");

void parse_options(int argc, char** argv)
{
    htio2::OptionParser option_parser;
    option_parser.add_option(OPT_files_in_ENTRY);
    option_parser.add_option(OPT_dir_out_ENTRY);
    option_parser.add_option(OPT_file_barcode_ENTRY);

    option_parser.add_option(OPT_mismatch_ENTRY);

    option_parser.add_option(OPT_se_ENTRY);
    option_parser.add_option(OPT_pe_ENTRY);
    option_parser.add_option(OPT_pe_itlv_ENTRY);
    option_parser.add_option(OPT_pe_strand_ENTRY);
    option_parser.add_option(OPT_header_format_ENTRY);
    option_parser.add_option(OPT_header_sra_ENTRY);

    option_parser.add_option(OPT_quiet_ENTRY);
    option_parser.add_option(OPT_help_ENTRY);
    option_parser.add_option(OPT_version_ENTRY);

    option_parser.parse_options(argc, argv);

    // show help
    if (OPT_help)
    {
        cout << endl
            << argv[0] << " - demultiplex reads by barcode"
            << endl << endl
            << option_parser.format_document();
        exit(EXIT_SUCCESS);
    }

    if (OPT_version) show_version_and_exit();
}

void validate_options()
{
    ensure_files_in();
    ensure_se_pe();

    if (!OPT_dir_out.length())
    {
        cerr << "ERROR: output directory not specified" << endl;
        exit(EXIT_FAILURE);
    }

    ensure_header_format();

    if (OPT_header_format != htio2::HEADER_1_8)
    {
        cerr << "ERROR: only CASAVA 1.8 headers can be used for demultiplex" << endl
            << "your header format is: " << htio2::to_string(OPT_header_format) << endl;
        exit(EXIT_FAILURE);
    }
}

void show_options()
{
    show_se_pe();
    show_files_in();
    cout << "# output dir: " << OPT_dir_out << endl;
    cout << "# barcode file: " << OPT_file_barcode << endl;
    cout << "# barcode mismatch: " << OPT_mismatch << endl;
    show_header_format();
}

void read_barcode(map<string, pair<string, string> > & barcodes)
{
    if (!OPT_quiet)
        cout << "read barcode file" << endl;
    barcodes.clear();

    htio2::PlainFileHandle IN(OPT_file_barcode.c_str(), htio2::READ);
    string buffer;

    while (IN.read_line(buffer))
    {
        size_t tab1 = buffer.find_first_of('\t');
        size_t tab2 = buffer.find_first_of('\t', tab1 + 1);
        string proj(buffer.substr(0, tab1));
        string sample(buffer.substr(tab1 + 1, tab2 - tab1 - 1));
        string barcode(buffer.substr(tab2 + 1));

        if (!barcodes.insert(make_pair(barcode, make_pair(proj, sample))).second)
        {
            cerr << "ERROR: duplicate barcode: " << barcode << endl;
            exit(EXIT_FAILURE);
        }
    }
}

void gen_output_dir(const map<string, pair<string, string> > & barcodes)
{
    if (!OPT_quiet) cout << "create output directories" << endl;

    File p_out(OPT_dir_out.c_str());

    if (!p_out.isDirectory())
    {
        Result re = p_out.createDirectory();
        if (re.failed())
        {
            fprintf(stderr, "ERROR: failed to create output directory \"%s\": %s\n",
                    OPT_dir_out.c_str(), re.getErrorMessage().toRawUTF8());
            exit(EXIT_FAILURE);
        }
    }

    set<string> projs;
    for (map<string, pair<string, string> >::const_iterator it = barcodes.begin(); it != barcodes.end(); it++)
    {
        projs.insert(it->second.first);
    }

    for (set<string>::const_iterator it = projs.begin(); it != projs.end(); it++)
    {
        File p_out_proj = p_out.getChildFile(it->c_str());

        // remove existing directory
        if (p_out_proj.isDirectory())
        {
            if (!OPT_quiet)
                printf("  remove existing dir \"%s\"\n", p_out_proj.getFullPathName().toRawUTF8());

            if (!p_out_proj.deleteRecursively())
            {
                fprintf(stderr, "ERROR: failed to remove dir \"%s\"\n", p_out_proj.getFullPathName().toRawUTF8());
                exit(EXIT_FAILURE);
            }
        }

        // create new directory
        if (!OPT_quiet)
            cout << "  create output dir: " << p_out_proj.getFullPathName().toStdString() << endl;

        Result re = p_out_proj.createDirectory();
        if (re.failed())
        {
            fprintf(stderr, "ERROR: failed to create output directory \"%s\": %s\n",
                    p_out_proj.getFullPathName().toRawUTF8(), re.getErrorMessage().toRawUTF8());
            exit(EXIT_FAILURE);
        }
    }
}

bool check_barcode(const map<string, pair<string, string> > & barcodes,
                   const string& barcode,
                   string& proj,
                   string& sample)
{
    map<string, pair<string, string> >::const_iterator it = barcodes.find(barcode);

    if (it != barcodes.end())
    {
        proj = it->second.first;
        sample = it->second.second;
        return true;
    }
    else if (OPT_mismatch)
    {
        for (it = barcodes.begin(); it != barcodes.end(); it++)
        {
            const string& ref_barcode = it->first;
            if (ref_barcode.length() != barcode.length()) continue;
            size_t diff = 0;
            for (size_t i = 0; i < ref_barcode.length(); i++)
            {
                if (ref_barcode[i] != barcode[i]) diff++;
            }

            if (diff <= OPT_mismatch)
            {
                proj = it->second.first;
                sample = it->second.second;
                return true;
            }
        }
        return false;
    }
    else
    {
        return false;
    }
}

int64_t n_accept = 0;
int64_t n_seq = 0;

string gen_out_file_one(const string& proj, const string& sample)
{
    File prefix = File(OPT_dir_out).getChildFile(String(proj)).getChildFile(String(sample));
    string file_out;
    add_fastq_suffix_se(prefix.getFullPathName().toStdString(), OPT_compress, file_out);
    check_output_overwrite(OPT_files_in, file_out);
    return file_out;
}

void gen_out_file_pair(const string& proj, const string& sample, string& f1, string& f2)
{
    File prefix = File(OPT_dir_out).getChildFile(String(proj)).getChildFile(String(sample));
    add_fastq_suffix_pe(prefix.getFullPathName().toStdString(), OPT_compress, f1, f2);
    check_output_overwrite(OPT_files_in, f1);
    check_output_overwrite(OPT_files_in, f2);
}

string gen_fail_file_one()
{
    File prefix = File(OPT_dir_out).getChildFile("demul_failed");
    string file_out;
    add_fastq_suffix_se(prefix.getFullPathName().toStdString(), OPT_compress, file_out);
    check_output_overwrite(OPT_files_in, file_out);
    return file_out;
}

void gen_fail_file_pair(string& f1, string& f2)
{
    File prefix = File(OPT_dir_out).getChildFile("demul_failed");
    add_fastq_suffix_pe(prefix.getFullPathName().toStdString(), OPT_compress, f1, f2);
    check_output_overwrite(OPT_files_in, f1);
    check_output_overwrite(OPT_files_in, f2);
}

void do_se(const map<string, pair<string, string> > & barcodes)
{
    htio2::FastqSeq curr_seq;
    htio2::HeaderParser parser(OPT_header_format, OPT_header_sra);
    string proj;
    string sample;

    int64_t n_accept = 0;
    int64_t n_seq = 0;

    string file_fail = gen_fail_file_one();

    htio2::FastqIO FAIL(file_fail, htio2::WRITE, OPT_compress);

    for (size_t file_i = 0; file_i < OPT_files_in.size(); file_i++)
    {
        htio2::FastqIO IN(OPT_files_in[file_i], htio2::READ, OPT_compress);

        while (IN.next_seq(curr_seq))
        {
            parser.parse(curr_seq.id, curr_seq.desc);

            if (check_barcode(barcodes, parser.barcode, proj, sample))
            {
                string file_out = gen_out_file_one(proj, sample);
                htio2::FastqIO OUT(file_out, htio2::APPEND, OPT_compress);
                OUT.write_seq(curr_seq);
                n_accept++;
            }
            else
            {
                FAIL.write_seq(curr_seq);
            }

            n_seq++;
            if (!OPT_quiet && n_accept % LOG_BLOCK == 0)
                cout << "  " << n_seq << " sequences, " << n_accept << " has barcode" << endl;
        }
    }

    if (!OPT_quiet)
        cout << "  " << n_seq << " sequences, " << n_accept << " has barcode" << endl;
}

void process_one_pe(const htio2::FastqSeq& seq1,
                    const htio2::FastqSeq& seq2,
                    htio2::PairedEndIO& fh_fail,
                    const map<string, pair<string, string> > & barcodes,
                    htio2::HeaderParser& parser1,
                    htio2::HeaderParser& parser2)
{
    n_seq++;

    string proj;
    string sample;

    parser1.parse(seq1.id, seq1.desc);
    parser2.parse(seq2.id, seq2.desc);

    if (parser1.barcode != parser2.barcode)
    {
        cerr << "ERROR: paired-end has conflict barcodes: " << endl
            << parser1.barcode << endl
            << parser2.barcode << endl;
        exit(EXIT_FAILURE);
    }

    if (check_barcode(barcodes, parser1.barcode, proj, sample))
    {
        n_accept++;

        if (OPT_pe_itlv)
        {
            string file_out = gen_out_file_one(proj, sample);
            htio2::PairedEndIO fh_out(file_out,
                                     htio2::APPEND,
                                     PE_STRAND_FILE,
                                     OPT_compress);
            fh_out.write_pair(seq1, seq2, PE_STRAND_OBJ);
        }
        else
        {
            string f1;
            string f2;
            gen_out_file_pair(proj, sample, f1, f2);
            htio2::PairedEndIO fh_out(f1,
                                     f2,
                                     htio2::APPEND,
                                     PE_STRAND_FILE,
                                     OPT_compress);
            fh_out.write_pair(seq1, seq2, PE_STRAND_OBJ);
        }
    }
    else
    {
        fh_fail.write_pair(seq1, seq2);
    }

    if (!OPT_quiet && n_accept % LOG_BLOCK == 0)
        cout << "  " << n_seq << " paired-ends, " << n_accept << " has barcode" << endl;
}

void do_pe(const map<string, pair<string, string> > & barcodes)
{
    htio2::FastqSeq curr1;
    htio2::FastqSeq curr2;
    htio2::HeaderParser parser1(OPT_header_format, OPT_header_sra);
    htio2::HeaderParser parser2(OPT_header_format, OPT_header_sra);

    if (!OPT_quiet)
        cout << "traverse input files" << endl;

    if (OPT_pe_itlv)
    {
        // create handle for failed output
        string file_fail = gen_fail_file_one();

        htio2::PairedEndIO fh_fail(file_fail,
                                   htio2::WRITE,
                                   PE_STRAND_FILE,
                                   OPT_compress);

        // traverse input files
        htqc::MultiSeqFilePE fh_in(OPT_files_in, PE_STRAND_FILE, OPT_compress);

        while (fh_in.next_pair(curr1, curr2, PE_STRAND_OBJ))
        {
            process_one_pe(curr1, curr2, fh_fail, barcodes, parser1, parser2);
        }
    }
    else
    {
        string file_fail1;
        string file_fail2;
        gen_fail_file_pair(file_fail1, file_fail2);

        htio2::PairedEndIO fh_fail(file_fail1,
                                   file_fail2,
                                   htio2::WRITE,
                                   PE_STRAND_FILE,
                                   OPT_compress);

        vector<string> files_in_a;
        vector<string> files_in_b;
        separate_paired_files(OPT_files_in, files_in_a, files_in_b);

        // traverse input file pairs
        htqc::MultiSeqFilePE fh_in(files_in_a,
                                   files_in_b,
                                   PE_STRAND_FILE,
                                   OPT_compress);

        while (fh_in.next_pair(curr1, curr2, PE_STRAND_OBJ))
        {
            process_one_pe(curr1, curr2, fh_fail, barcodes, parser1, parser2);
        }
    }

    if (!OPT_quiet)
        cout << "done: " << n_seq << " paired-ends, " << n_accept << " has barcode" << endl;
}

void execute()
{
    map<string, pair<string, string> > barcodes;

    read_barcode(barcodes);
    gen_output_dir(barcodes);

    if (OPT_se)
        do_se(barcodes);
    else if (OPT_pe)
        do_pe(barcodes);
    else
        throw runtime_error("neither single-end nor paired-end");
}
