package beastfx.app.seqgen;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import beast.base.core.BEASTInterface;
import beast.base.core.BEASTObject;
import beast.base.core.Description;
import beast.base.core.Input;
import beast.base.core.Input.Validate;
import beast.base.evolution.alignment.Alignment;
import beast.base.evolution.alignment.Sequence;
import beast.base.parser.XMLParser;
import beast.base.parser.XMLParserException;
import beast.base.parser.XMLProducer;
import beast.base.util.FileUtils;

@Description("Helper for Sequence Simulator, allows specifying template input file and destination output file")
public class MergeDataWith extends BEASTObject {
	final public Input<String> templateFileInput = new Input<>("template","name of template file. " +
			"The template file should be a valid beast2 XML file with a single (possibly empty) alignment. " +
			"This alignment will be replaced by the alignment generated by the SequenceSimulator. " +
			"Any occurrance of $(n) will be replaced by the iteration number of the SequenceSimulator. " +
			"For example, calling a log filename \"output$(n).log\" will be replaced by " +
			"\"output1.log\" at first iteration, \"output2.log\" at seccond iteration, etc..", Validate.REQUIRED);
	final public Input<String> outputFileInput = new Input<>("output","name of the output file. " +
			"Instances of $(n) are replaced by iteration number.", Validate.REQUIRED);

	File templateFile;
	
	@Override
	public void initAndValidate() {
		templateFile = new File(templateFileInput.get());
		if (!(templateFile.exists())) {
			throw new RuntimeException("Template file " + templateFileInput.get() + " does not exist"); 
		}
	} // initAndValidate
	
	
	void process(Alignment data, int iteration) throws IOException, XMLParserException, IllegalArgumentException, IllegalAccessException {
		// read template
		String templateXML = FileUtils.load(templateFile);
		// RRB: the following line prevents use of templates in input XML that use var="n", 
		// which is widely used, so commented out
		// templateXML = templateXML.replaceAll("\\$\\(n\\)", iteration+"");
		XMLParser parser = new XMLParser();
		BEASTInterface b = parser.parseBareFragment(templateXML, false);

		// replace alignment
		Alignment a = getAlignment(b);
		List<Sequence> sequences = a.sequenceInput.get();
		sequences.clear();
		sequences.addAll(data.sequenceInput.get());

		// write file
		String outputFile = outputFileInput.get();
		outputFile = outputFile.replaceAll("\\$\\(n\\)", iteration+"");
		FileWriter outfile = new FileWriter(outputFile);
		
		Set<BEASTInterface> beastObjects = new HashSet<>();
		String xml = new XMLProducer().toXML(b, beastObjects);
        outfile.write(xml);
        outfile.close();
	} // process


	private Alignment getAlignment(BEASTInterface b) throws IllegalArgumentException, IllegalAccessException {
		Alignment a = null;
		for (BEASTInterface i : b.listActiveBEASTObjects()) {
			if (i.getClass().equals(Alignment.class)){
				return (Alignment) i;
			} else {
				a = getAlignment(i);
				if (a != null) {
					return a;
				}
			}
		}
		return null;
	} // getAlignment
	
}
