package jalview.ws.jws2;

import jalview.api.AlignCalcWorkerI;
import jalview.datamodel.AlignmentAnnotation;
import jalview.datamodel.Annotation;
import jalview.gui.AlignFrame;
import jalview.ws.jws2.jabaws2.Jws2Instance;
import jalview.ws.params.WsParamSetI;

import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.TreeSet;

import compbio.data.sequence.RNAStructReader.AlifoldResult;
import compbio.data.sequence.RNAStructScoreManager;
import compbio.data.sequence.Range;
import compbio.data.sequence.Score;
import compbio.metadata.Argument;

public class RNAalifoldClient extends JabawsAlignCalcWorker implements
		AlignCalcWorkerI 
{
	
	// test
	
	String methodName;
	
	AlignFrame af;
	
	// keeps track of whether the RNAalifold result includes base contact probabilities
	boolean bpScores;
	
	public RNAalifoldClient(Jws2Instance sh, AlignFrame alignFrame,
			WsParamSetI preset, List<Argument> paramset)
	{
		super(sh, alignFrame, preset, paramset);
		
		if (arguments == null) arguments = new ArrayList<Argument>();
		arguments.add(sh.getRunnerConfig().getArgumentByOptionName("-p"));

		
		af = alignFrame;
		methodName = sh.serviceType;
		
		// defult false. Which one here?
		// submitGaps = true;
		nucleotidesAllowed = true;
		proteinAllowed = false;
		
		arguments.add(sh.getRunnerConfig().getArgumentByOptionName("-p"));
	}
	
	@Override
	public String getServiceActionText()
	{
		return "Submitting RNA alignment for Secondary Structure prediction using "
				+ "RNAalifold Service";
	}


	@Override
	public void updateResultAnnotation(boolean immediate)
	{

		if (immediate || !calcMan.isWorking(this) && scoremanager != null) 
		{
			
			List<AlignmentAnnotation> ourAnnot = new ArrayList<AlignmentAnnotation>();

			// Unpack the ScoreManager
			List<String> structs = ((RNAStructScoreManager) scoremanager).getStructs();
			List<TreeSet<Score>> data = ((RNAStructScoreManager) scoremanager).getData();
			
			// test to see if this data object contains base pair contacts 
			Score fscore = data.get(0).first();
		  this.bpScores = (fscore.getMethod().equals(
					AlifoldResult.contactProbabilities.toString()));

			// Add annotations for the mfe Structure 
		  if (bpScores) 
		  	createAnnotationRowforScoreHolder(ourAnnot, getCalcId(), structs.get(1),
		  			data.get(0), data.get(1));
		  else 
		  	createAnnotationRowforScoreHolder(ourAnnot, getCalcId(), structs.get(1),
		  			data.get(1));
		  
		  // add annotation for the consensus sequence alignment
			createAnnotationRowforScoreHolder(ourAnnot, getCalcId(), structs.get(0), null);
			
			// Not loop for the rest of the Annotations
			if (structs.size() > 2) {
				for (int i = 2; i < structs.size(); i++) {
					// I can't think of a nice way of presenting the ensembleValues data
					//  so I wont for now.
					if (!data.get(i).first().getMethod().equals(
							AlifoldResult.ensembleValues.toString())) {
						createAnnotationRowforScoreHolder(ourAnnot, getCalcId(), structs.get(i),
								data.get(i));
					}
				}
			}
			
			
			if (ourAnnot.size() > 0) {
				
				// Modify the visible annotation on the alignment viewport with the
        // new alignment annotation rows created.
				updateOurAnnots(ourAnnot);
				// ap.adjustAnnotationHeight();
			}
		}
	}
	
	// just for the base pair contact annotation. It uses a second score object.
	protected void createAnnotationRowforScoreHolder(
			List<AlignmentAnnotation> ourAnnot, String calcId, 
			String struct, TreeSet<Score> data, TreeSet<Score> descriptionData) {
		
		String typename = data.first().getMethod().toString();
		
		AlignmentAnnotation annotation = alignViewport.getAlignment()
				.findOrCreateAnnotation(typename, calcId, false, null, null);
		
		constructAnnotationFromContactProbabilities(annotation, struct, data);
		
		String description = constructAlignmentAnnotationDescription(descriptionData.first());
		if (description.length() == 0) description = typename;
		annotation.description = description;
		
		// dan test
		annotation.belowAlignment = false;
		
		annotation.validateRangeAndDisplay();
		
		ourAnnot.add(annotation);
	}
	
	protected void createAnnotationRowforScoreHolder(
			List<AlignmentAnnotation> ourAnnot, String calcId, 
			String struct, TreeSet<Score> data)
	{
		/* If contactProbability information is returned from RNAalifold it is stored
		 * in the first TreeSet<Score> object corresponding to the String Id which
		 * holds the consensus alignment. The method enumeration is then updated to
		 * AlifoldResult.contactProbabilties. This line (hack) recreates the same 
		 * data object as was overwritten with the contact probabilites data.
		 */
		if (data == null) data = compbio.data.sequence.RNAStructReader
				.newEmptyScore(AlifoldResult.consensusAlignment);
		
		String typename = data.first().getMethod().toString();
		
		AlignmentAnnotation annotation = alignViewport.getAlignment()
					.findOrCreateAnnotation(typename, calcId, false, null, null);
				
		// construct annotation from ScoreHolder (unpacked into struct and data)
		if (bpScores && data.first().getMethod().equals(
				AlifoldResult.contactProbabilities.toString())) 
			constructAnnotationFromContactProbabilities(annotation, struct, data);
		
		else
			// if bpScores is false the TreeSet<Score> data should always contain
			//  a single Score object
			constructAnnotationFromStructureString(annotation, struct, data.first());
		
		/*  update annotation description with the free Energy, frequency in ensemble 
		 *  or other data where appropriate.
		 *  
		 *  Doesnt deal with AlifoldResult.ensembleValues, the free energy of ensemble
		 *  and frequency of mfe structure in ensemble. How to deal with these? 
		 */
		String description = constructAlignmentAnnotationDescription(data.first());
		if (description.length() == 0) description = typename;
		annotation.description = description;
		
		// dan test
		annotation.belowAlignment = false;
		
		annotation.validateRangeAndDisplay();
		
		ourAnnot.add(annotation);
	}
	
	
	
	private AlignmentAnnotation constructAnnotationFromStructureString(
			AlignmentAnnotation annotation, String struct, Score score) 
	{
		
		Annotation[] anns = new Annotation[struct.length()];
		
		for (int i = 0; i < struct.length(); i++) {
			anns[i] = new Annotation(struct.substring(i, i+1), "",
					struct.charAt(i), Float.NaN);
		}
		
		annotation.graph = 0; // No graph
		annotation.annotations = anns;

		
		return annotation;
		
	}

	private String constructAlignmentAnnotationDescription(Score score) {
		String description = "";
		String datatype = score.getMethod();
		
		if (datatype.equals(AlifoldResult.mfeStructure.toString()) || 
				datatype.equals(AlifoldResult.centroidStructure.toString())) {
			description = MessageFormat.format("Energy: {0} = {1} + {2}", 
					score.getScores().get(0), score.getScores().get(1), score.getScores().get(2));
		}
		else if (datatype.equals(AlifoldResult.contactProbabilityStructure.toString())) {
			description = MessageFormat.format("Energy: {0}  Frequency: {1}", 
					score.getScores().get(0), score.getScores().get(1));
		}
		else if (datatype.equals(AlifoldResult.stochBTStructure.toString())) {
			if (score.getScores().size() > 0) {
				description = MessageFormat.format("Probability: {0}  Energy: {1}", 
						score.getScores().get(0), score.getScores().get(1));
			}
			else description = "Stochastic Backtrack Structure";
		}
		else if (datatype.equals(AlifoldResult.MEAStucture.toString())) {
			description = MessageFormat.format("Maximum Expected Accuracy Values: '{' {0} MEA={1} '}",
						score.getScores().get(0), score.getScores().get(1));
		}
		
		return description;
	}
	
	
	private AlignmentAnnotation constructAnnotationFromContactProbabilities(
			AlignmentAnnotation annotation, String struct, TreeSet<Score> data) 
	{
		Annotation[] anns = new Annotation[struct.length()];

		TreeMap<Range, Float> basePairs = null;
		// The base pair probabilities are stored in a set in scoreholder. we want a map
		basePairs = new TreeMap<Range, Float>();
		for (Score score : data) {
			// The Score objects contain a set of size one containing the range and
			//  an ArrayList<float> of size one containing the probabilty
			basePairs.put(score.getRanges().first(), new Float(score.getScores().get(0)));
		}
		for (int i = 0; i < struct.length(); i++) {

			// Return all the contacts associated with position i
			List<Range> contacts = isContact(basePairs, i+1);

			if (contacts.size() == 0) {
				anns[i] = new Annotation(struct.substring(i, i+1), "", struct.charAt(i), 0f);
			}
			else if (contacts.size() == 1) {
				// There is only one contact associated with this base
				float prob = basePairs.get(contacts.get(0));
				anns[i] = new Annotation(struct.substring(i, i+1), "", struct.charAt(i), prob);
			}
			else if (contacts.size() > 1) {
				// For now we will simply deal with alternate contact information by mentioning its
				//  existance in the description
				float prob = basePairs.get(contacts.get(0));
				anns[i] = new Annotation(struct.substring(i, i+1), "This base has alternate contacts",
						struct.charAt(i), prob);
			}
		}
		
		annotation.annotations = anns;
		
		return annotation;
	}
	

	// Check whether, at position i there is a base contact and return all the
	//  contacts at this position. Should be in order of descending probability.
	private List<Range> isContact(TreeMap<Range, Float> basePairs, int i) {
		
		List<Range> contacts = new ArrayList<Range>();
		
		for (Range contact : basePairs.keySet()) {
			// finds the contacts associtated with position i ordered by the natural
			//  ordering of the Scores TreeSet in ScoreManager which is, descending probability
			if (contact.from == i || contact.to == i) contacts.add(contact);
		}
		
		return contacts;
	}
	

	public String getCalcId()
  {
    return SequenceAnnotationWSClient.AAConCalcId;
  }
	
}