package org.rki.readtrainer;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.util.LinkedList;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.CommandLineParser;
import org.apache.commons.cli.GnuParser;
import org.apache.commons.cli.HelpFormatter;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.ParseException;
import org.rki.readclassifier.Reader;

public class Trainer {
	public static boolean showTime;
	public static double startTimeMillis;
	public static int kmersize = 4;
	public static AtomicInteger[][] matrix1;
	public static AtomicInteger[][] matrix2;
	
	public static FileInputStream a1 = null;
	public static FileInputStream a2 = null;
	public static FileInputStream b1 = null;
	public static FileInputStream b2 = null;
	
	public static String namea = "";
	public static String nameb = "";
	
	public static FileOutputStream out = null;
	public static int threads = 4;
	
	public static double absentval = 0.0001;
	
	public static void main(String[] args) {
		try {
			init();
			parseOptions(args);
			process();
			finish();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public static void finish() {
		if(showTime) {
			long currentTime = System.currentTimeMillis();
			System.out.println("Time taken: " + (currentTime - startTimeMillis) + " ms");
		}
	}
	
	public static void process() throws Exception {
		BufferedWriter outWriter = new BufferedWriter(new OutputStreamWriter(out));
		matrix1 = new AtomicInteger[(int)Math.pow(4, kmersize)][4];
		clearMatrix(matrix1);
		trainSet(a1, a2, matrix1);
		writeMatrix(outWriter, namea, matrix1);
		matrix2 = new AtomicInteger[(int)Math.pow(4, kmersize)][4];
		clearMatrix(matrix2);
		trainSet(b1, b2, matrix2);
		writeMatrix(outWriter, nameb, matrix2);
		outWriter.close();
		showScore();
	}

	public static void showScore() {
		double score = 0;
		for(int i=0; i<matrix1.length; i++) {
			int rowsum1 = 0;
			for(int j=0; j<4; j++) {
				rowsum1 += matrix1[i][j].get();
			}
			int rowsum2 = 0;
			for(int j=0; j<4; j++) {
				rowsum2 += matrix2[i][j].get();
			}
			for(int j=0; j<4; j++) {
				score += Math.abs((matrix1[i][j].get()*4.0/(double)rowsum1)-(matrix2[i][j].get()*4.0/(double)rowsum2));
			}
		}
		score = score / Math.pow(4.0d, kmersize+1.0d);
		System.out.println("Matrix difference (expected quality) score: " + score);
	}
	
	public static void writeMatrix(BufferedWriter writer, String name, AtomicInteger[][] matrix) throws IOException {
		writer.write("transition_matrix\tk=");
		writer.write(Integer.toString(kmersize));
		writer.write("\tspecies_name=");
		writer.write(name);
		writer.write("\n");
		for(int i=0; i<matrix.length; i++) {
			int rowsum = 0;
			for(int j=0; j<4; j++) {
				rowsum += matrix[i][j].get();
			}
			if(rowsum == 0) {
				writer.write("None\tNone\tNone\tNone\n");
				continue;
			}
			for(int j=0; j<4; j++) {
				double val = (double)matrix[i][j].get();
				if(val == 0)
					val = absentval;
				val = val/(double)rowsum*4.0d;
				val = Math.log(val);
				writer.write(Double.toString(val));
				if(j != 3) {
					writer.write("\t");
				}
			}
			writer.write("\n");
		}
	}
	
	public static void clearMatrix(AtomicInteger[][] matrix) {
		for(int i=0; i<(int)Math.pow(4, kmersize); i++) {
			for(int j=0; j<4; j++) {
				matrix[i][j] = new AtomicInteger(0);
			}
		}
	}
	
	public static void trainSet(FileInputStream f1, FileInputStream f2, AtomicInteger[][] matrix) throws InterruptedException {
		Reader reader = new Reader(f1, f2, 10000, 50000, threads);
		reader.start();
		LinkedList<TrainerThread> trainers = new LinkedList<TrainerThread>();
		for(int i=0; i<threads; i++) {
			trainers.add(new TrainerThread(reader, matrix));
		}
		for(TrainerThread trainer : trainers) {
			trainer.start();
		}
		for(TrainerThread trainer : trainers) {
			trainer.join();
		}
		reader.join();
	}
		
	public static void init() throws Exception {
		startTimeMillis = System.currentTimeMillis();
	}
		
	public static void parseOptions(String[] args) throws ParseException, FileNotFoundException {
		Options options = new Options();
		options.addOption("a1", "a1", true, "Input fastq/fasta file containing first reads from organism a");
		options.addOption("a2", "a2", true, "Input fastq/fasta file containing second reads from organism a");
		options.addOption("b1", "b1", true, "Input fastq/fasta file containing first reads from organism b");
		options.addOption("b2", "b2", true, "Input fastq/fasta file containing second reads from organism b");		
		options.addOption("na", "name_a", true, "Name for organism a (default: inferred from first read file name)");
		options.addOption("nb", "name_b", true, "Name for organism b (default: inferred from first read file name)");
		options.addOption("o", "datafile", true, "output file");
		options.addOption("k", "kmer", true, "K-mer size (default: " + kmersize + ")");
		options.addOption("t", "threads", true, "Number of threads (default: " + threads + ")");
		options.addOption("T", "show_time", false, "Show processing time");
		options.addOption("h", "help", false, "Print usage instructions");
		
		CommandLineParser parser = new GnuParser();
		CommandLine cmd = parser.parse(options, args);
		
		if(cmd.hasOption("h")) {
			HelpFormatter formatter = new HelpFormatter();
			formatter.printHelp("java -jar trainer.jar", options);
			System.exit(0);
		}
		if(!cmd.hasOption("a1")) {
			System.out.println("Input file for the first read for organism a must be provided.");
			System.exit(1);
		}
		if(!cmd.hasOption("b1")) {
			System.out.println("Input file for the first read for organism b must be provided.");
			System.exit(1);
		}
		if((cmd.hasOption("a2") && !cmd.hasOption("b2"))) {
			System.out.println("If the second read file is provided for one organism, it has to also be provided for the other.");
			System.exit(1);
		}
		if((cmd.hasOption("b2") && !cmd.hasOption("a2"))) {
			System.out.println("If the second read file is provided for one organism, it has to also be provided for the other.");
			System.exit(1);
		}
		
		if(cmd.hasOption("k")) {
			kmersize = Integer.parseInt(cmd.getOptionValue("k"));
		}
		
		a1 = new FileInputStream(new File(cmd.getOptionValue("a1")));
		if(cmd.hasOption("a2"))
			a2 = new FileInputStream(new File(cmd.getOptionValue("a2")));
		b1 = new FileInputStream(new File(cmd.getOptionValue("b1")));
		if(cmd.hasOption("b2"))
			b2 = new FileInputStream(new File(cmd.getOptionValue("b2")));

		out = new FileOutputStream(new File(cmd.getOptionValue("o")));
		
		if(cmd.hasOption("na")) {
			namea = cmd.getOptionValue("na");
		} else {
			namea = new File(cmd.getOptionValue("a1")).getName().split("\\.")[0];
		}
		if(cmd.hasOption("nb")) {
			nameb = cmd.getOptionValue("nb");
		} else {
			nameb = new File(cmd.getOptionValue("b1")).getName().split("\\.")[0];
		}
		
		if(cmd.hasOption("t")) {
			threads = Integer.parseInt(cmd.getOptionValue("t"));
		}
		showTime = cmd.hasOption("T");
	}
		
	public static int calculateVal(char[] kmer) {
		int res = 0;
		for(char base : kmer) {
			res = res << 2;
			switch (base) {
			case 'A':
				break;
			case 'C':
				res |= 1;
				break;
			case 'G':
				res |= 2;
				break;
			case 'T':
				res |= 3;
				break;
			default:
				break;
			}
		}
		return res;
	}
}
