package jalview.ws.jws2;

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

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.TreeSet;

import compbio.data.msa.jaxws.Align;
import compbio.data.sequence.Range;
import compbio.data.sequence.Score;
import compbio.data.sequence.RNAStructScoreManager;
import compbio.data.sequence.ScoreManager;
import compbio.metadata.Argument;

public class RNAalifoldClient extends JabawsAlignCalcWorker implements
		AlignCalcWorkerI 
{
	
	// test
	
	String methodName;
	
	AlignFrame af;
	
	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>();
			//		ourAnnots = new ArrayList<AlignmentAnnotation>();

			// So I don't have to do any more casting
			List<String> structs = ((RNAStructScoreManager) scoremanager).getStructs();
			List<TreeSet<Score>> data = ((RNAStructScoreManager) scoremanager).getData();
			
			
			// I think this will never find an annotation at the moment. It will always create...
			AlignmentAnnotation annotation = alignViewport.getAlignment()
					.findOrCreateAnnotation("Consensus Structure", getCalcId(), false, null , null);
			
			// construct Annotation from scoremanager
			
			// Deal with the consensus structure and (?)BasePair Probabilities
			Annotation[] anns = new Annotation[structs.get(1).length()];

			// check if the first Score object is populated with base pair probabilities
			Score fscore = data.get(0).first();
			boolean BPScores = (fscore.getScores().size() > 0 
					&& fscore.getRanges() != null);

			TreeMap<Range, Float> basePairs = null;
			if (BPScores) {
				// The base pair probabilities are stored in a set in scoremanager... we want a map
				basePairs = new TreeMap<Range, Float>();
				for (Score score : data.get(0)) {
					// 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)));
				}
			}

			// ignoring the Consensus alignemnt for now, get the Structure and make an AlignmentAnnotation
			String struct = structs.get(1); // get(1)
			for (int i = 0; i < struct.length(); i++) {

				if (BPScores) {
					// 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);
					}
				}
				else {
					// Same as the first if from the previous block
					anns[i] = new Annotation(struct.substring(i, i+1), "", struct.charAt(i), 0f);
				}
			}

			System.out.println("size of anns: " + anns.length);
			
			// Set the annotation to the AlignmentAnnotation object
			annotation.annotations = anns;
			
			// Set the probability
			annotation.setScore(data.get(1).first().getScores().get(0));
			
			

			// old
//			AlignmentAnnotation annot = new AlignmentAnnotation("Consensus Structure", "Free Energy", anns);
			
			
			
			
			System.out.println("RNAalifoldClient - annotation:\n");
			for (Annotation ann : annotation.annotations) {
				System.out.print(ann.toString()+"|");
			}
			System.out.println();

			// Instead of this look at existing methods for creating annotations
			ourAnnot.add(annotation);
			
			
			if (ourAnnot.size() > 0) {
				updateOurAnnots(ourAnnot);
			}
			
		}
	}
	
	// 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;
  }
	
}