package org.rki.readclassifier;

import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;

public class MMClassifier implements ReadClassifier {
	private double tmorg1[][]; 
	private double tmorg2[][];
	private int kmersize;
	private int kmermask = 0;
	
	public String org1name = "";
	public String org2name = "";
	
	@Override
	public String getName() {
		return "MMClassifier";
	}

	@Override
	public String getDescription() {
		return("Classifies reads based on Markov Model of Kmer transitions.");
	}

	@Override
	public void readData(InputStream in) throws Exception {
		BufferedReader reader = new BufferedReader(new InputStreamReader(in));
		String line = reader.readLine();
		String[] dat = line.split("\\s+");
		if(!dat[0].equals("transition_matrix")) {
			throw new Exception("Error: Wrong data file type. Found " + dat[0] + ", expected transition_matrix");
		}
		kmersize = Integer.parseInt(dat[1].split("=")[1]);
		kmermask = (int)Math.pow(4, kmersize) - 1;
		tmorg1 = new double[(int)Math.pow(4, kmersize)][4];
		tmorg2 = new double[(int)Math.pow(4, kmersize)][4];
//		Classifier.logger.log(Level.INFO, "Found kmer size: " + kmersize);
		org1name = dat[2].split("=")[1].trim();
		int currDist = 0;
		int numline = 0;
		while(reader.ready()) {
			dat = reader.readLine().split("\\s+");
			if(dat[0].equals("transition_matrix")) {
				currDist += 1;
				numline = 0;
				org2name = dat[2].split("=")[1].trim();
				continue;
			}
			int pos=0;
			for(String val : dat) {
				double dval;
				val = val.trim();
				if(val.equals("None")) {
					dval = Double.NaN;
				}
				else {
					dval = Double.parseDouble(val);
				}
				if(currDist == 0)
					tmorg1[numline][pos] = dval;
				else
					tmorg2[numline][pos] = dval;
				pos += 1;
			}
			numline += 1;
		}
		reader.close();	
	}

	@Override
	public void getScore(ReadPair reads) {
		if(reads.read1.bases.length() < kmersize) {
			System.out.println("Error: Read with length < kmer size found: ");
			System.out.println(">" + reads.read1.name);
			System.out.println(">" + reads.read1.bases);
			System.exit(1);
		}
		getReadScore(reads.read1);
		if(reads.read2 != null) {
			if(reads.read2.bases.length() < kmersize) {
				System.out.println("Error: Read with length < kmer size found: ");
				System.out.println(">" + reads.read2.name);
				System.out.println(">" + reads.read2.bases);
				System.exit(1);
			}
			getReadScore(reads.read2);
			reads.read2.score1 += reads.read1.score1;
			reads.read2.score2 += reads.read1.score2;
			reads.read1.score1 = reads.read2.score1;
			reads.read1.score2 = reads.read2.score2;
			reads.read1.score += reads.read2.score;
			reads.read2.score = reads.read1.score;
		}
	}
	
	private int[] getKmerval(char[] bases, int start) {
		int[] res = {0, 0};
		int i=0;
		int kmerval = 0;
		while(start+i<bases.length && i<kmersize) {
			int nextbase = getBaseVal(bases[start+i]);
			if(nextbase == -1) {
				start += i + 1;
				i = 0;
				kmerval = 0;
				continue;
			}
			kmerval = kmerval << 2;
			kmerval |= nextbase;
			i++;
		}
		if(i < kmersize) {
			return null;
		}
		res[0] = kmerval;
		res[1] = start + i;
		return res;
	}
	
	private void getReadScore(Read read) {
		char[] bases = read.bases.toCharArray();
		read.score = 0;
		read.score1 = 0;
		read.score2 = 0;
		double len = bases.length;
		int[] initialKmer = getKmerval(bases, 0);
		if(initialKmer == null) {
			read.invalid = true;
			return;
		}
		int kmerval = initialKmer[0];
		int i = initialKmer[1];
		for(; i<len; i++) {
			int nextbase = getBaseVal(bases[i]);
			if(nextbase == -1) {
				initialKmer = getKmerval(bases, i+1);
				if(initialKmer == null)
					break;
				kmerval = initialKmer[0];
				i = initialKmer[1];
				continue;
			}
			double ds1 = tmorg1[kmerval][nextbase];
			double ds2 = tmorg2[kmerval][nextbase];
			if(!Double.isNaN(ds1) && !Double.isNaN(ds2)) {
				read.score1 += ds1;
				read.score2 += ds2;
//				read.score1 *= ds1;
//				read.score2 *= ds2;	
			}
			kmerval = kmerval << 2;
			kmerval |= nextbase;
			kmerval &= kmermask;
		}
		read.score = read.score1 - read.score2;
		read.normalizeTo(Classifier.normalizeTo);
	}

	private int getBaseVal(char base) {
		switch (base) {
		case 'A':
			return 0;
		case 'C':
			return 1;
		case 'G':
			return 2;
		case 'T':
			return 3;
		default:
			return -1;
		}		
	}	
}
