#!/usr/bin/env python

import cogent
from numpy import mean, std
from cogent.app.stride import stride_xtra
from cogent.struct.selection import select, einput
from cogent.maths.stats.util import Numbers
from cogent.maths.stats.test import ANOVA_one_way

PDB_FILE = '1UI9_bfac_bits.pdb'

print "Loading structure: %s" % (PDB_FILE,)
structure = cogent.LoadStructure(PDB_FILE)
print "Running Strid"
print stride_xtra(structure)
print "Remove residues which have no secondary-structure annotation"
ss_residues = select(structure, 'R', None, 'ne', 'STRIDE_SS', xtra=True)
print "Transfer bfactor values from atoms to residues."
ss_residues.propagateData(mean, 'A', 'bfactor')
print "Group residues based on secondary-structure"
classes = ss_residues.splitChildren('STRIDE_SS', xtra=True)
print "Got secondary-structure groups: %s" % (" ".join(classes.keys()))
results = []
print "Get vectors of bits for the residues in each group."
for cls, ents in classes.iteritems():
    residues = einput(ents.values(), 'R')
    residue_bits = Numbers(residues.getData('bfactor'))
    results.append(residue_bits)
    print cls, residue_bits.Mean, '+/-', residue_bits.StandardDeviation

ar = ANOVA_one_way(results)
print "Result of the ANOVA one way test: F=%s , right tail=%s" % (ar[3], ar[-1])