from scitbx.math import basic_statistics

from giant.structure.select import non_h, protein, backbone, sidechains

########################################################################################

class BfactorStatistics(object):
    """B-factor statstics for a structure"""
    def __init__(self, all=None, protein=None, backbone=None, sidechain=None):
        self.all       = all
        self.protein   = protein
        self.backbone  = backbone
        self.sidechain = sidechain

    def to_z_score(self, b_vals, method='backbone'):
        """Convert a set of B-factors to Z-scores"""
        assert method in ['all','protein','backbone','sidechain']
        if method == 'all':       stats = self.all
        if method == 'protein':   stats = self.protein
        if method == 'backbone':  stats = self.backbone
        if method == 'sidechain': stats = self.sidechain
        if stats.biased_standard_deviation == 0.0:
            # No variation - return list of zeroes
            return b_vals-b_vals

        return (b_vals - stats.mean)/stats.biased_standard_deviation

    @classmethod
    def from_pdb(cls, pdb_input=None, pdb_hierarchy=None):
        """Calculate the b-factor statistics of a model"""

        assert [pdb_input, pdb_hierarchy].count(None)==1,'Provide pdb_input OR pdb_hierarchy'
        if pdb_input: pdb_hierarchy = pdb_input.construct_hierarchy()

        cache = pdb_hierarchy.atom_selection_cache()

        all_b       = non_h(hierarchy=pdb_hierarchy, cache=cache, copy=True).atoms().extract_b()
        protein_b   = protein(hierarchy=pdb_hierarchy, cache=cache, copy=True).atoms().extract_b()
        backbone_b  = backbone(hierarchy=pdb_hierarchy, cache=cache, copy=True).atoms().extract_b()
        sidechain_b = sidechains(hierarchy=pdb_hierarchy, cache=cache, copy=True).atoms().extract_b()

        return cls(all       = basic_statistics(all_b),
                   protein   = basic_statistics(protein_b),
                   backbone  = basic_statistics(backbone_b),
                   sidechain = basic_statistics(sidechain_b) )

    def show(self):
        if self.all:
            print '================>'
            print 'All Atoms:'
            print self.all.show()
        if self.protein:
            print '================>'
            print 'Protein Atoms:'
            print self.protein.show()
        if self.backbone:
            print '================>'
            print 'Backbone Atoms:'
            print self.backbone.show()
        if self.sidechain:
            print '================>'
            print 'Sidechain Atoms:'
            print self.sidechain.show()

########################################################################################

def normalise_b_factors_to_z_scores(pdb_input=None, pdb_hierarchy=None, b_factor_statistics=None, method='backbone'):
    """Calculate the b-factor statistics of the model"""

    assert [pdb_input, pdb_hierarchy].count(None)==1,'Provide pdb_input OR pdb_hierarchy'
    if pdb_input: pdb_hierarchy = pdb_input.construct_hierarchy()

    if not b_factor_statistics: b_factor_statistics = BfactorStatistics.from_pdb(pdb_hierarchy=pdb_hierarchy)
    new_b = b_factor_statistics.to_z_score(b_vals=pdb_hierarchy.atoms().extract_b(), method=method)

    output_h = pdb_hierarchy.deep_copy()
    output_h.atoms().set_b(new_b)

    return output_h

########################################################################################


