Small edits to make code nicer.
[htsworkflow.git] / htswanalysis / src / GetReadsInSnps / getReadsInSnps.cpp
1 #include <sys/types.h>
2 #include <dirent.h>
3 #include <errno.h>
4 #include <iostream>
5 #include <fstream>
6 #include <vector>
7 #include <map>
8 #include <queue>
9 #include <math.h>
10 #include <string>
11
12 #include <gsl/gsl_statistics.h>
13
14 #define WINDOW 25
15
16 using namespace std;
17
18 void split (const string& text, const string& separators, vector<string>& words);
19 char *strrevcomp(const string& input);
20
21 class Loci {
22   public:
23     string chr;
24     unsigned int pos;
25
26     Loci(string chr, unsigned int pos) { this->chr = chr; this->pos = pos; }
27     Loci(const Loci& l) { this->chr = l.chr; this->pos = l.pos; }
28     Loci& operator=(const Loci& l) { this->chr = l.chr; this->pos = l.pos; return *this; }
29
30     bool operator<(const Loci& a) const { if(this->chr == a.chr) { return this->pos < a.pos; } else { return this->chr < a.chr; } }
31
32 };
33
34 class Read : public Loci {
35   public:
36     string seq;
37
38     Read(string chr, unsigned int pos, string seq) : Loci(chr,pos) { this->seq = seq; }
39     Read(const Read& r) : Loci(r) { this->seq = r.seq; }
40     Read& operator=(const Read& r) { this->chr = r.chr; this->pos = r.pos; this->seq = r.seq; return *this;}
41 };
42
43 class SNP : public Loci {
44   public:
45
46     string name;
47     char reference_base;
48     unsigned int A;
49     unsigned int C;
50     unsigned int G;
51     unsigned int T;
52     unsigned int N;
53     unsigned int total;
54
55     SNP(string name, string chr, unsigned int pos, char reference_base) : Loci(chr,pos) { 
56       this->name = name; 
57       this->A = 0; this->C = 0; this->G = 0; this->T = 0; this->total = 0;
58       this->reference_base = reference_base; 
59     }
60
61     SNP(const SNP& h) : Loci(h) {
62       this->name = h.name; 
63       this->A = h.A; this->C = h.C; this->G = h.G; this->T = h.T; this->total = h.total;
64       this->reference_base = h.reference_base; 
65     }
66
67     SNP& operator=(const SNP& h) {
68       this->name = h.name; 
69       this->chr = h.chr; 
70       this->pos = h.pos; 
71       this->A = h.A; this->C = h.C; this->G = h.G; this->T = h.T; this->total = h.total;
72       this->reference_base = h.reference_base; 
73       return *this;
74     }
75
76     void add_read(char nuc) {
77       switch(nuc) {
78         case 'a':
79         case 'A':
80           A++; break;
81         case 'c':
82         case 'C':
83           C++; break;
84         case 'g':
85         case 'G':
86           G++; break;
87         case 't':
88         case 'T':
89           T++; break;
90         default:
91           N++; break;
92       }
93       total++;
94     }
95 };
96
97 ostream &operator<<( ostream &out, const SNP &h ) {
98   out << h.name.c_str() << "\t" << h.chr.c_str() << "\t" << h.pos << "\t" << h.reference_base << "\t" << h.A << "\t" << h.C << "\t" << h.G << "\t" << h.T << "\t" << h.total;
99   return out;
100 }
101
102 typedef vector<Read> Reads;
103 typedef vector<SNP> SNPs;
104
105 void read_snps(const char* filename, SNPs& snps) {
106   string delim("\t");
107
108   ifstream feat(filename);
109   size_t N = 0;
110   while(feat.peek() != EOF) {
111     char line[1024];
112     feat.getline(line,1024,'\n');
113     N++;
114     string line_str(line);
115     vector<string> fields;
116     split(line_str, delim, fields);
117     if(fields.size() != 4) { cerr << "Error (" << filename << "): wrong number of fields in feature list (line " << N << " has " << fields.size() << " fields)\n"; }
118
119     string name = fields[0];
120     string chr = fields[1];
121     unsigned int pos = atoi(fields[2].c_str());
122     char base = (fields[3])[0];
123
124     SNP snp(name,chr,pos,base);
125     snps.push_back(snp);
126   } 
127
128   //sort the features so we can run through it once
129   std::stable_sort(snps.begin(),snps.end());
130   feat.close();
131
132   cerr << "Found and sorted " << snps.size() << " snps." << endl;
133 }
134
135 void read_align_file(char* filename, Reads& features) {
136   string delim(" \n");
137   string location_delim(":");
138   char strand_str[2]; strand_str[1] = '\0';
139   ifstream seqs(filename);
140   string name("");
141   while(seqs.peek() != EOF) {
142     char line[2048];
143     seqs.getline(line,2048,'\n');
144
145     string line_str(line);
146     vector<string> fields;
147     split(line_str, delim, fields);
148     if(fields.size() != 7) { continue; }
149
150  
151     vector<string> location; split(fields[3], location_delim, location);
152     string chr = location[0];
153     if(chr == "NA") { continue; }
154     int pos = atoi(location[1].c_str());
155     bool strand = ((fields[4].c_str())[0] == 'F')?0:1;
156
157     string seq;
158     if(strand == 0) {
159       seq = fields[0]; 
160     } else {
161       seq = strrevcomp(fields[0]);
162     }
163     Read read(chr,pos,seq); 
164     features.push_back(read);
165   }
166   seqs.close(); 
167
168   //sort the data so we can run through it once
169   std::sort(features.begin(),features.end());
170   cerr << "Found and sorted " << features.size() << " reads." << endl;
171 }
172
173 void count_read_at_snps(SNPs& snps, Reads& data) {
174   SNPs::iterator snp_it = snps.begin();
175
176   //assume, for now, that all reads have the same length
177   unsigned int read_len = data[0].seq.length();
178   
179   for(Reads::iterator i = data.begin(); i != data.end(); ++i) {
180     //skip to first feature after read
181     string start_chr = snp_it->chr;
182     while(snp_it != snps.end() && *snp_it < *i) {
183       snp_it++;
184     }
185     
186     //stop if we have run out of features.
187     if(snp_it == snps.end()) { break; }
188
189     if(i->pos + read_len > snp_it->pos && i->pos <= snp_it->pos) {
190       snp_it->add_read(i->seq[snp_it->pos - i->pos]);
191     }
192   }
193 }
194
195 int main(int argc, char** argv) {
196   if(argc != 3) { cerr << "Usage: " << argv[0] << " snp_file read_file\n"; exit(1); }
197
198   char snp_filename[1024]; strcpy(snp_filename,argv[1]);
199   char read_filename[1024]; strcpy(read_filename,argv[2]);
200
201   SNPs snps; read_snps(snp_filename, snps);
202   Reads reads; read_align_file(read_filename, reads);
203
204   count_read_at_snps(snps, reads);
205
206   for(SNPs::iterator i = snps.begin(); i != snps.end(); ++i) {
207     cout << (*i) << endl;
208   }
209 }
210
211 void split (const string& text, const string& separators, vector<string>& words) {
212
213     size_t n     = text.length ();
214     size_t start = text.find_first_not_of (separators);
215
216     while (start < n) {
217         size_t stop = text.find_first_of (separators, start);
218         if (stop > n) stop = n;
219         words.push_back (text.substr (start, stop-start));
220         start = text.find_first_not_of (separators, stop+1);
221     }
222 }
223
224 char *strrevcomp(const string& input)
225 {
226   char* str = new char[input.length()];
227   strcpy(str,input.c_str());
228
229   char *p1, *p2;
230
231   if (! str || ! *str)
232     return str;
233
234   for (p1 = str, p2 = str + strlen(str) - 1; p2 > p1; ++p1, --p2) {
235     *p1 ^= *p2;
236     *p2 ^= *p1;
237     *p1 ^= *p2;
238   }
239
240   for (p1 = str; p1 < str + strlen(str); ++p1) {
241     if(*p1 == 'a' || *p1 == 'A') { *p1 = 'T'; } 
242     else if(*p1 == 'c' || *p1 == 'C') { *p1 = 'G'; } 
243     else if(*p1 == 'g' || *p1 == 'G') { *p1 = 'C'; } 
244     else if(*p1 == 't' || *p1 == 'T') { *p1 = 'A'; } 
245   }
246
247   return str;
248 }
249