/* 
     Description: 处理intervals

     Copyright : All right reserved by ICT

     Author : Zhang Zhonghai
     Date : 2019/11/24
*/

#include "interval.h"

#include <algorithm>
#include <sstream>
#include <fstream>
#include <string>
#include <iostream>

#include <htslib/sam.h>

#include "utils.h"
#include "bam_wrap.h"

using std::min;
using std::max;
using std::string;
using std::ifstream;
using std::stringstream;

using namespace std;

// 构造函数
Interval::Interval() : Interval(0, 0) {}
Interval::Interval(int64_t l, int64_t r) : left(l), right(r) {}

// 比较函数
bool Interval::operator<(const Interval& other) {
    if (left == other.left) {
        return right < other.right;
    }
    return left < other.left;
}

// 是否有重叠
bool Interval::overlaps(const Interval &other) {
    return left <= other.right && right >= other.left;
}

// 两个interval的合并
Interval& Interval::spanWith(const Interval &other) {
    left = min(left, other.left);
    right = max(right, other.right);
    return *this;
}

// 返回两个interval的交集，不改变当前interval
Interval Interval::intersect(const Interval &that) const {
    Interval val;
    val.left = max(left, that.left);
    val.right = min(right, that.right);
    return val;
}


/*
 * 合并两个interval arr，取相交区域的交集, interval arr都是排序后的
 */
void Interval::IntersectIntervals(const IntervalArr &a_arr, 
                                  const IntervalArr &b_arr, 
                                  IntervalArr *r_arr) {
    if (a_arr.size() < 1 || b_arr.size() < 1) return;
    int ai=0, bi=0;
    const Interval *last, *cur;
    if (a_arr[ai].left < b_arr[bi].left) last = &a_arr[ai++];
    else last = &b_arr[bi++];
    while (ai < a_arr.size() && bi < b_arr.size()) {
        if (a_arr[ai].left < b_arr[bi].left) cur = &a_arr[ai++];
        else cur = &b_arr[bi++];
        if (last->right < cur->left) {
            last = cur; continue;
        } else if (last->right > cur->right) {
            r_arr->push_back(*cur); 
        } else {
            r_arr->push_back(Interval(cur->left, last->right));
            last = cur;
        }
    }
    const IntervalArr *arrp;
    int ii;
    if (ai < a_arr.size()) { arrp = &a_arr; ii = ai;} 
    else { arrp = &b_arr; ii = bi; }
    const IntervalArr &arr = *arrp;
    while(ii < arr.size()) {
        cur = &arr[ii++];
        if (last->right < cur->left) {
            break;
        } else if (last->right > cur->right) {
            r_arr->push_back(*cur); 
        } else {
            r_arr->push_back(Interval(cur->left, last->right));
            break;
        }
    }
}

/*
 * 合并两个interval arr，取并集
 */
void Interval::UnionIntervals(const IntervalArr &a_arr,
                              const IntervalArr &b_arr,
                              IntervalArr *r_arr) {
    Interval tmp;
    const Interval *cur;
    Interval *last;
    int ai=0, bi=0;
    if (a_arr.size() < 1) { *r_arr = b_arr; return; }
    if (b_arr.size() < 1) { *r_arr = a_arr; return; }
    r_arr->clear();

    if (a_arr[ai].left < b_arr[bi].left) tmp = a_arr[ai++];
    else tmp = b_arr[bi++];
    last = &tmp;
    while(ai < a_arr.size() && bi < b_arr.size()) {
        if (a_arr[ai].left < b_arr[bi].left) cur = &a_arr[ai++];
        else cur = &b_arr[bi++];
        if (last->right < cur->left) {
            r_arr->push_back(*last);
            *last = *cur;
        } else {
            last->right = max(last->right, cur->right);
        }
    }
    const IntervalArr *arrp;
    int ii;
    if (ai < a_arr.size()) { arrp = &a_arr; ii = ai; } 
    else { arrp = &b_arr; ii = bi; }
    const IntervalArr &arr = *arrp;

    while(ii < arr.size()) {
        cur = &arr[ii++];
        if (last->right < cur->left) {
            r_arr->push_back(*last);
            *last = *cur;
        } else {
            last->right = max(last->right, cur->right);
        }
    }
    r_arr->push_back(*last);
}

/*
 * 将有read覆盖的区域和参数提供的interval文件中的区域做一个交集
 */
int64_t Interval::MergeIntervals(const IntervalArr &n_arr,
                                  const IntervalArr &t_arr,
                                  IntervalArr &in_arr,
                                  int64_t start_loc, // 闭区间
                                  int64_t *end_loc, // 开区间
                                  IntervalArr *r_arr) {
    IntervalArr tmp_arr;
    const int64_t end_loc_val = *end_loc;
    if (in_arr.size() < 1) { // 如果输入的interval为空，则使用tumor normal覆盖的interval
        UnionIntervals(n_arr, t_arr, &tmp_arr);
    } else {
        IntervalArr mid_arr;
        UnionIntervals(n_arr, t_arr, &mid_arr);
        IntersectIntervals(mid_arr, in_arr, &tmp_arr);
    }
    for(int i=tmp_arr.size()-1; i>=0; --i) {
        if (tmp_arr[i].left >= end_loc_val) {
            tmp_arr.pop_back(); // 删除该元素
            continue;
        }
        tmp_arr[i].right = min(tmp_arr[i].right, end_loc_val - 1); // end_loc是开区间
        break;
    }
    for (int i=0; i<tmp_arr.size(); ++i) {
        if (tmp_arr[i].right < start_loc) {
            continue;
        }
        if (tmp_arr[i].left < start_loc) {
            r_arr->push_back(Interval(start_loc, tmp_arr[i].right));
        } else {
            r_arr->push_back(tmp_arr[i]);
        }
    }

    int next_i = 0;
    while(next_i < in_arr.size() && in_arr[next_i].right < end_loc_val) ++next_i;
    if (next_i < in_arr.size()) {
        if (end_loc_val < in_arr[next_i].left) {
            *end_loc = in_arr[next_i].left; // 更新本次处理的终点
        } else {
            in_arr[next_i].left = end_loc_val; // 更新panel
        }
        int i=0, j=next_i;
        for (; j<in_arr.size(); ++i, ++j) {
            in_arr[i] = in_arr[j];
        }
        in_arr.resize(i);
    } else {
        in_arr.clear();
    }

    int64_t locus_num = 0;
    for (int i=0; i<r_arr->size(); ++i) {
        locus_num += (*r_arr)[i].right - (*r_arr)[i].left + 1; 
    }
    return locus_num;
}

/*
 * 读取interval文件
 */
void Interval::ReadInterval(const string &interval_fn, 
                            bam_hdr_t* header, 
                            int interval_padding,
                            IntervalArr *r_arr) {
    ifstream interval_fs(interval_fn);
    string one_line;
    IntervalArr intervals;
    getline(interval_fs, one_line);
    while (!interval_fs.eof()) {
        if (one_line[0] == '@') {
            getline(interval_fs, one_line);
            continue;
        }
        stringstream ss_line(one_line);
        string contig_name;
        ss_line >> contig_name;
        int itid = sam_hdr_name2tid(header, contig_name.c_str());
        if (itid < 0) error("[%s] interval file has unknown contig name [%s]\n", __func__, contig_name.c_str());
        int64_t tid = (int64_t)itid;
        tid <<= CONTIG_SHIFT;
        int64_t start, stop;
        ss_line >> start >> stop;
        // interval文件是1-based，所以这里要减去1
        intervals.push_back(Interval(tid + start - 1, tid + stop -1));
        getline(interval_fs, one_line);
    }
    sort(intervals.begin(), intervals.end());
    if (intervals.size() > 0) {
        Interval new_span(intervals[0].left-interval_padding, intervals[0].right+interval_padding);
        for (int i=1; i<intervals.size(); ++i) {
            if (intervals[i].left - interval_padding > new_span.right) {
                r_arr->push_back(new_span);
                new_span.left = intervals[i].left - interval_padding;
                new_span.right = intervals[i].right + interval_padding;
            } else {
                new_span.right = max(new_span.right, intervals[i].right + interval_padding);
            }
        }
        r_arr->push_back(new_span);
    }
    interval_fs.close();
}

/*
 * 将interval相连的区域合并
 */
void Interval::ShrinkInterval(IntervalArr *ivap) {
    if (ivap->size() < 1) return;
    IntervalArr &iva = *ivap;
    IntervalArr tiva = iva;
    iva.clear();
    Interval iv;
    iv.left = tiva[0].left;
    iv.right = tiva[0].right;
    for (int i=1; i<tiva.size(); ++i) {
        if (iv.right+1 < tiva[i].left) {
            iva.push_back(iv);
            iv.left = tiva[i].left;
        }
        iv.right = tiva[i].right;
    }
    iva.push_back(iv);
}

/*
 * 根据header信息，扩展interval
 */
Interval Interval::ExpandInterval(int64_t start, int64_t end, int expandVal, bam_hdr_t* header) {
    Interval result;
    result.left = start;
    result.right = end;

    int64_t ext_left = start - expandVal;
    int64_t ext_right = end + expandVal;
    int tid = BamWrap::bam_tid(start);
    uint32_t contig_len = header->target_len[tid];
    result.left = max(BamWrap::bam_global_pos(tid, 0), ext_left);
    result.right = min(ext_right, contig_len - 1 + BamWrap::bam_global_pos(tid, 0));

    return result;
}

