import jalview.workers.FeatureCounterI;
import jalview.workers.AlignmentAnnotationFactory;

/*
 * Example script that registers two alignment annotation calculators
 * - one that counts residues in a column with Pfam annotation
 * - one that counts only charged residues with Pfam annotation
 * To try this, first load uniref50.fa from the examples folder, then load features
 * from examples/exampleFeatures.txt, before running this script from the Groovy console.
 
 * Modify this example as required to count by column any desired value that can be 
 * derived from the residue and sequence features at each position of an alignment.
 */

/*
 * A closure that returns true for any Charged residue
 */
def isCharged = { residue ->
    switch(residue) {
        case ['D', 'd', 'E', 'e', 'H', 'h', 'K', 'k', 'R', 'r']:
            return true
    }
    false
} 

/*
 * A closure that returns 1 if sequence features include type 'Pfam', else 0
 * Argument should be a list of SequenceFeature 
 */
def hasPfam = { features -> 
    for (sf in features)
    {
        /*
         * Here we inspect the type of the sequence feature.
         * You can also test sf.description, sf.score, sf.featureGroup,
         * sf.strand, sf.phase, sf.begin, sf.end
         * or sf.getValue(attributeName) for GFF 'column 9' properties
         */
        if ("Pfam".equals(sf.type))
        {
            return true
        }
    }
    false
}

/*
 * Closure that counts residues with a Pfam feature annotation
 * Parameters are
 * - the name (label) for the alignment annotation
 * - the description (tooltip) for the annotation
 * - a closure (groovy function) that tests whether to include a residue
 * - a closure that tests whether to increment count based on sequence features  
 */
def getColumnCounter = { name, desc, acceptResidue, acceptFeatures ->
    [
     getName: { name }, 
     getDescription: { desc },
     getMinColour: { [0, 255, 255] }, // cyan
     getMaxColour: { [0, 0, 255] }, // blue
     count: 
         { res, feats -> 
            def c = 0
            if (acceptResidue.call(res))
            {
                if (acceptFeatures.call(feats))
                {
                    c++
                }
            }
            c
         }
     ] as FeatureCounterI
}

/*
 * Define an annotation that counts any residue with Pfam domain annotation
 */
def pfamAnnotation = getColumnCounter("Pfam", "Count of residues with Pfam domain annotation", {true}, hasPfam)

/*
 * Define an annotation that counts charged residues with Pfam domain annotation
 */
def chargedPfamAnnotation = getColumnCounter("Pfam charged", "Count of charged residues with Pfam domain annotation", isCharged, hasPfam)

/*
 * Register the annotations
 */
AlignmentAnnotationFactory.newCalculator(pfamAnnotation) 
AlignmentAnnotationFactory.newCalculator(chargedPfamAnnotation)