import jalview.workers.AlignmentAnnotationFactory;
import jalview.workers.AnnotationProviderI;
import jalview.datamodel.AlignmentAnnotation;
import jalview.datamodel.Annotation;
import jalview.util.ColorUtils;
import jalview.util.Comparison;
import java.awt.Color;

/*
 * Example script to compute two alignment annotations
 * - count of Phosphorylation features
 * - count of Turn features
 * To try this, first load example file uniref50.fa and load on features file
 * exampleFeatures.txt, before running this script
 *
 * The script only needs to be run once - it will be registered by Jalview
 * and recalculated automatically when the alignment changes.
 */

/*
 * A closure that returns true if value includes "PHOSPHORYLATION"
 */
def phosCounter = { type ->    type.contains("PHOSPHORYLATION") }

/*
 * A closure that returns true if value includes "TURN"
 */
def turnCounter = { type ->    type.contains("TURN") }

/*
 * A closure that computes and returns an array of Annotation values,
 * one for each column of the alignment
 */
def getAnnotations(al, fr, counter) 
{
    def width = al.width
    def counts = new int[width] 
    def max = 0
    
    /*
     * count features in each column, record the maximum value
     */
    for (col = 0 ; col < width ; col++)
    {
        def count = 0
        for (row = 0 ; row < al.height ; row++)
        {
            seq = al.getSequenceAt(row)
            if (seq != null && col < seq.getLength())
            {
                def res = seq.getCharAt(col)
                if (!Comparison.isGap(res))
                {
                    pos = seq.findPosition(col)
                    features = fr.findFeaturesAtRes(seq, pos)
                    for (feature in features)
                    {
                        if (counter.call(feature.type))
                        {
                            count++
                        }
                    }
                }
            }
        }
        counts[col] = count
        if (count > max)
        {
            max = count
        }
    }
    
    /*
     * make the Annotation objects, with a graduated colour scale 
     * (from min value to max value) for the histogram bars 
     */
    def zero = '0' as char
    def anns = new Annotation[width] 
    for (col = 0 ; col < width ; col++)
    {
        def c = counts[col]
        if (c > 0)
        {
            Color color = ColorUtils.getGraduatedColour(c, 0, Color.cyan,
                max, Color.blue)
            anns[col] = AlignmentAnnotationFactory.newAnnotation(String.valueOf(c),
                String.valueOf(c), zero, c, color)
        }
    }
    anns
}

/*
 * Define the method that performs the calculations, and builds two
 * AlignmentAnnotation objects
 */
def annotator = 
    [ calculateAnnotation: { al, fr ->
        def phosAnns = getAnnotations(al, fr, phosCounter)
        def ann1 = AlignmentAnnotationFactory.newAlignmentAnnotation("Phosphorylation", "Count of Phosphorylation features", phosAnns)
        def turnAnns = getAnnotations(al, fr, turnCounter)
        def ann2 = AlignmentAnnotationFactory.newAlignmentAnnotation("Turn", "Count of Turn features", turnAnns)
        return [ann1, ann2]
      } 
    ] as AnnotationProviderI
    
/*
 * Register the annotation calculator with Jalview
 */
AlignmentAnnotationFactory.newCalculator(annotator)