/*
 * Segmentation.cpp
 *
 * Copyright (C) 2012 Björn Voß

 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or any later version.

 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.

 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

 * Authors: Thorsten Bischler and Björn Voß
 */

#include "Segmentation.h"
#include <iostream>
#include <sstream>
#include <math.h>
#include <limits>
#include <algorithm>

#define DEBUG
extern int threads;

using namespace std;

Segmentation::Segmentation(DataStorage data, int kmax, int cpmax, int prim_window, double read_cutoff, double read_cutoff2,  double pvalue_cutoff, double mean_cov_cutoff, std::string gff_file, std::string score_file, int kmin, bool verbose, double ratio, double non_transcript_cov_cutoff, int zmax, int zfrac)
: data(data), n(data.fragment_size), genome_size(data.genome_size), kmax(kmax), cpmax(cpmax), prim_window(prim_window), read_cutoff(read_cutoff), read_cutoff2(read_cutoff2), pvalue_cutoff(pvalue_cutoff), mean_cov_cutoff(mean_cov_cutoff), gff_file(gff_file), score_file(score_file), kmin(kmin), verbose(verbose), ratio(ratio), non_transcript_cov_cutoff(non_transcript_cov_cutoff), zmax(zmax), zfrac(zfrac)
{}

void Segmentation::compute_G(vector<vector<double> >& G) {
    
    vector<double> cr (n,0);
    vector<double> cq (n,0);
    
    cr[0] = log2(data.get_src(0)+1.0);
    cq[0] = pow(log2(data.get_src(0)+1.0), 2);
    
    for (int i = 1; i < n; i++) {
        cr[i] = cr[i-1] + log2(data.get_src(i)+1.0);
        cq[i] = cq[i-1] + pow(log2(data.get_src(i)+1.0), 2);
    }
        
#pragma omp parallel num_threads(threads)
    {
#pragma omp for
        for (int k = 0; k < kmax; k++) {
            G.at(0).at(k) = cq[k] - pow(cr[k], 2)/double(k+1);
        }
#pragma omp for
        for (int i = 1; i < n; i++) {
            for (int k = 0; k < kmax; k++) {
                G.at(i).at(k) = numeric_limits<double>::max();
            }
        }
    }
    cout << "Progress: " << 100/kmax << "%\r" << flush;
    
    for (int k = 0; k < kmax-1; k++) {
        
#pragma omp parallel for num_threads(threads)
        for (int i = 0; i < n-(k+1); i++) {
            double crk = cr[i+k+1]-cr[i];
            double cqk = cq[i+k+1]-cq[i];
            G.at(i+1).at(k) = cqk - pow(crk,2)/double(k+1);
        }
        cout << "Progress: " << (k+2)*100/kmax << "%\r" << flush;
    }
}


void Segmentation::compute_N(vector<int >& N) {
    int putative_starts = 0;
    int max_gap = -1;
    {
    	int prev_prs = -(prim_window+1);
        
      	for (int i = 0; i < n; i++) {
            bool pval_met = false;
            for (int s = i - 13; s <= i - 11; s++) {
                if (s >= 0) {
                    if (data.get_pval(s) <= pvalue_cutoff) { pval_met = true; }
                }
            }
            
            if ( ((pval_met == true && data.get_prs(i) >= read_cutoff) || data.get_prs(i) >= read_cutoff2) && data.get_prs(i) >= ratio * data.get_prc(i) )  {
                prev_prs = i;
                putative_starts++;
            }
            
            if ( i-prev_prs > max_gap ) { max_gap = i-prev_prs; }
            
            N.at(i) = prev_prs;
            
            for (int j = N.at(i); j>0 && i-N.at(j) <= prim_window; j--) {
                
                if (data.get_prs(N.at(j)) > data.get_prs(N.at(i))) {
                    N.at(i) = N.at(j);
                }
                
            }
      	}
    }
    N.at(n) = n;
    
    cout << "Found " << putative_starts << " putative TSS! (max. gap = " << max_gap << ")" << endl;

    if (putative_starts * 2 +1 < cpmax) {
	cpmax = putative_starts * 2 +1;
	cout << "NOTE: Max. number of change points lowered to " << cpmax << endl;
    }

    if (max_gap > kmax * 2) {

	kmax = max_gap / 1.5;
	cout << "NOTE: Max. transcript size raised to " << kmax << endl;

    }

}

void Segmentation::writeTransToFile(int start, int end, ofstream &outfile) {
    if (outfile.is_open()){
        outfile << "SeqID\tRNAseg\ttranscript\t" << start << "\t" << end << "\t.\t" << data.strand << "\t.\tcolour=255 99 71" << endl;
    }
    else {
        cerr << "Output file not open" << endl;
        exit(1);
    }
}

void Segmentation::writeNonTransToFile(int start, int end, ofstream &outfile) {
    if (outfile.is_open()){
        outfile << "SeqID\tRNAseg\tnon_transcript\t" << start << "\t" << end << "\t.\t" << data.strand << "\t.\tcolour=0 71 255" << endl;
        
    }
    else {
        cerr << "Output file not open" << endl;
        exit(1);
    }
}

void Segmentation::writeOutputFile(vector<int >& N, vector<vector <int> >& th, string gff_file, string score_file) {
    
    ofstream gff_file_stream(gff_file.c_str());
    ofstream score_file_stream(score_file.c_str());
    
    score_file_stream <<"cp_number\tprim_starts_trans\tprim_starts_nonTrans\tsec_cov_trans\tsec_cov_nonTrans\toptScore" << endl;
    
    cout << "Writing result file..." << endl;

    for(int cp_number = 0; cp_number < cpmax; cp_number++) {
        
        if (th.at(cp_number).at(cp_number) == -1 || th.at(cp_number).at(0) == 2*n) { continue;} // Skip changepoint numbers without valid segementations
        
        gff_file_stream << cp_number+1 << " change points" << endl;
        double trans_score_sec = 0.0;
        double non_trans_score_sec = 0.0;
        double trans_score_prim = 0.0;
        double non_trans_score_prim = 0.0;
        
        //start and end of the current segment depending on the position with max prim_starts
        SegBound seg_bound;
        
        //output for the first segment where cp = 0
        seg_bound.start = 0;
        seg_bound.end = abs(th.at(cp_number).at(0)) - 1;
        
        if ( seg_bound.start - N.at(seg_bound.start) <= prim_window ) {   // data.get_prs(seg_bound.start) >= read_cutoff) {
            trans_score_sec += data.get_sum_src(seg_bound.start, seg_bound.end);
            trans_score_prim += data.get_prs(seg_bound.start);
            if (data.strand == '+') {
                writeTransToFile(data.start_pos, seg_bound.end + data.start_pos, gff_file_stream);
            }
            else{
                writeTransToFile(data.end_pos - seg_bound.end, data.end_pos, gff_file_stream);
            }
        }
        else {
            non_trans_score_sec += data.get_sum_src(seg_bound.start, seg_bound.end);
            non_trans_score_prim += data.get_prs(seg_bound.start);
            if (data.strand == '+') {
                writeNonTransToFile(data.start_pos, seg_bound.end + data.start_pos, gff_file_stream);
            }
            else{
                writeNonTransToFile(data.end_pos - seg_bound.end, data.end_pos, gff_file_stream);
            }
        }
        
        
        //output for the second to the last segment if available
        for (int cp = 1; cp <= cp_number; cp++) {
            
            seg_bound.start = abs(th.at(cp_number).at(cp-1));
            seg_bound.end = abs(th.at(cp_number).at(cp)) - 1;
            
            if(th.at(cp_number).at(cp-1) > 0) {
                trans_score_sec += data.get_sum_src(seg_bound.start, seg_bound.end);
                trans_score_prim += data.get_prs(seg_bound.start);
                if (data.strand == '+') {
                    writeTransToFile(seg_bound.start + data.start_pos, seg_bound.end + data.start_pos, gff_file_stream);
                }
                else{
                    writeTransToFile(data.end_pos - seg_bound.end, data.end_pos - seg_bound.start, gff_file_stream);
                }
            }
            else {
                non_trans_score_sec += data.get_sum_src(seg_bound.start, seg_bound.end);
                non_trans_score_prim += data.get_prs(seg_bound.start);
                if (data.strand == '+') {
                    writeNonTransToFile(seg_bound.start + data.start_pos, seg_bound.end + data.start_pos, gff_file_stream);
                }
                else{
                    writeNonTransToFile(data.end_pos - seg_bound.end, data.end_pos - seg_bound.start, gff_file_stream);
                }
            }
        }
        score_file_stream << cp_number+1 << "\t" << trans_score_prim << "\t" << non_trans_score_prim << "\t" << trans_score_sec << "\t" << non_trans_score_sec << "\t" << (double(trans_score_prim) * (double(trans_score_sec)/double(non_trans_score_sec))) << endl;
    }
    cout << "Result file complete!" << endl;
    
    gff_file_stream.close();
    score_file_stream.close();
}


void Segmentation::computeSegments() {
    
    vector<int> N (n+1,0);
    
    cout << "Computing N matrix..." << endl;
    compute_N(N);
    cout << "N matrix computed!" << endl;
 
    vector<vector<double> > G (n, vector <double> (kmax,numeric_limits<double>::max()));
    
    cout << "Computing G matrix..." << endl;
    compute_G(G);
    cout << "G matrix computed!" << endl;
    
    /* initialize for cp=0: mI.at(k).at(0) is simply G.at(0).at(k) */

    cout << "Initializing mI matrix..." << endl;

    vector<vector<double> > mI(n, vector<double> (cpmax,-1.0));
    
#pragma omp parallel num_threads(threads)
    {
#pragma omp for
        for(int k = 0; k < kmax; k++) {
            mI.at(k).at(0) = G.at(0).at(k);
        }
        
    }
    cout << "mI matrix initialized" << endl;
    
    cout << "Initializing mt matrix..." << endl;
    
    vector<vector<int> > mt (n, vector <int> (cpmax, 2*n));
    
    cout << "mt matrix initialized" << endl;
    
    cout << "Computing segmentation for up to " << cpmax << " change points... (min. mean coverage = " << mean_cov_cutoff << ")." << endl;
    
    for (int cp = 1; cp < cpmax; cp++) {
        /*  Best segmentation with cp change points from 0 to j
         is found from considering best segmentations from 0 to j-k-1
         with cp-1 segments, plus cost of segment from j-k to j. */
        
#pragma omp parallel for num_threads(threads)
        for (int j = 0; j < n; j++) {
            double zmin = -1;
            int imin = 2*n; 
            double mean = 0.0;
            int zero_count = 0;
            int zero_stretch = 0;
            int longest_zs = 0;
            
            /* find the best change point between 0 and j-1 */
            int k0 = (j < kmax) ? j : kmax;
            for (int k = 0; k < k0; k++) {
                                
                mean = ((mean * k) + data.get_src(j-k)) / (k+1);
                if (data.get_src(j-k) == 0.0) {
                    zero_count++;
                    zero_stretch++;
                } else {
                    if (zero_stretch > longest_zs) {
                        longest_zs = zero_stretch;
                    }
                    zero_stretch = 0;
                }
                
                bool is_prim = false;
                
                int pre_cp = (cp>1) ? abs(mt.at(j-k-1).at((cp-1)-1)) : 0;
                
                if (N.at(j-k) > 0) {                                            // might be -(prim_window+1) at the beginning of the region
                
                    if ( (j-k)-N.at(j-k) <= prim_window &&                       // TSS within range
                        mean >= mean_cov_cutoff &&                               // mean coverage above cutoff
                        zero_count <= (zfrac*k)/100 &&
                        longest_zs <= zmax &&
                        j - N.at(j-k) >= kmin && //N.at(j-k) - pre_cp >= kmin && // Min. size of transcript segments
                        mI.at(N.at(j-k)-1).at(cp-1) >= 0) {                      // Front segment valid
                            is_prim = true;
                    }
                    
                }
 
                /*check if the previous segment is a transcript segment,
                 if not the current segment must be one*/
                if(cp != 1) {                    
                   
                    if (is_prim == false && ( mt.at(j-k-1).at((cp-1)-1) < 0 ||       // prev. segment non-transcript
                                              mI.at(j-k-1).at(cp-1) < 0     ||       // Front segment invalid
                                              (mean >= non_transcript_cov_cutoff && longest_zs <= zmax) ) )  // mean coverage above cutoff for non-transcripts & no long empty stretch
                    {
                        continue; // go to next k
                    }
                }
                else {
                    if (is_prim == false) { // In case of a single cp it must be a primary! No constraints on front segment!
                        continue;
                    }
                }
                
                double z;
                
                /* Best segmentation from 0 to j-k-1 or N[j-k]-1, respectively*/
                if ( is_prim == true ) {
                    z = mI.at(N.at(j-k)-1).at(cp-1);
                } else {
                    z = mI.at(j-k-1).at(cp-1);
                }
                
                /* Cost of segment from j-k to j */
                z += G.at(j-k).at(k);
                
                /* Assign negative positions to non_transcripts */
                if (z < zmin || zmin < 0) {
                    if( is_prim == true ) {
                        imin = N.at(j-k);
                        zmin = z;
                    }
                    else {
                        //if (mean < non_transcript_cov_cutoff) {
                            imin = -(j-k);
                            zmin = z;
                        //}
                    }
                } /* if z */
                
            } /* for k */
            
            mI.at(j).at(cp) = zmin;
            mt.at(j).at(cp-1) = imin;
            //}
        } /* for j */
        cout << "Progress: " << cp*100/(cpmax-1) << "%\r" << flush;
    } /* for cp */
    cout << "Segment scoring complete!" << endl;

    cout << "Starting backtracing..." << endl;
    /* th: elements 0...cp-1 of the cp-th row of matrix th contain
     the cp change points; element cp has value n, which corresponds
     to a changepoint at the rightmost point */
    
    vector<vector<int> > th (cpmax, vector <int> (cpmax, -1));
    
    for(int cp = 0;  cp < cpmax; cp++) {

        
        if (mt.at(n-1).at(cp) == 2*n) {  // no valid segmentation for this number of changepoints. Go to next number of cps.
            if (verbose) {cerr << "No valid segmentation found for " << cp << " changepoints. Skipping." << endl;}
            continue;
        }
        
        /* Backtrack to get th */
        /* In the following loop i is always the change point to the right */
        int i = n;
        th.at(cp).at(cp) = i;
        
        for(int j = cp-1; j >= 0; j--) {
            if((i==0)||(abs(i)>n)) {
                th.at(cp).at(cp) = -1;
                if (verbose) {cerr << "No valid segmentation found for " << cp << " changepoints. Skipping." << endl;}
                break;
            } else {
                
                th.at(cp).at(j) = i = mt.at(abs(i)-1).at(j);
                
            }
            
        }
        
        cout << "Progress: " << (cp+1)*100/cpmax << "%\r" << flush;
    }
    cout << "Traceback complete!" << endl;
    
    bool found_no_valid_seg = true;
    
    for(int cp = 0;  cp < cpmax; cp++) {
    
        if (th.at(cp).at(cp) > -1) { found_no_valid_seg = false; }
    
    }
    
    if (found_no_valid_seg) {
        cerr << "RNAseg could not find a valid segmentation for any number of change points!\n Possible fixes: Increase kmax or decrease one or several cutoffs (-r, -R, -p, -a, -u, -z, -Z, ...)\n";
        exit;
    }
    
    stringstream gff_file_stream;
    if(gff_file == "") {
        gff_file_stream << "output.gff";
    }
    else{
        gff_file_stream << gff_file;
    }
    
    stringstream score_file_stream;
    if(score_file == "") {
        score_file_stream << "output.csv";
    }
    else{
        score_file_stream << score_file;
    }
    writeOutputFile(N,th,gff_file_stream.str(), score_file_stream.str());
    
}
