# TFZ-based scoring 
TFZ_BINS = [
    0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5, 7.0,
    7.5, 8.0, 8.5, 9.0, 9.5, 10.0, 10.5, 11.0, 11.5, 12.0, 12.5, 13.0, 13.5,
    14.0, 14.5, 15.0,
    ]

TFZ_POLAR_SOLVED_FRACS = [
    0, 0, 0, 0, 0, 0.0163934426, 0.0059171598, 0.018766756, 0.0751252087,
    0.1077170418, 0.2623762376, 0.4738955823, 0.679245283, 0.9090909091,
    0.9385964912, 0.9903846154, 0.981981982, 1, 1, 1, 1, 1, 1, 0.9836065574,
    1, 1, 1, 1, 1, 1, 1,
    ]

TFZ_NONPOLAR_SOLVED_FRACS = [
    0, 0, 0, 0, 0, 0, 0, 0, 0.0042735043, 0.0077319588, 0.0354609929,
    0.0585215606, 0.1129251701, 0.2610497238, 0.6229946524, 0.8595317726,
    0.9399293286, 0.9819494585, 0.9928057554, 0.9928315412, 0.9925650558,
    0.996031746, 0.9927272727, 1, 0.9957983193, 1, 1, 1, 1, 1, 1 
    ]

# Fit: TFZ= -3 -> 15
# Use: TFZ=  2 -> 7.5
TFZ_POLAR_SOLVED_FRAC_CHEBYSHEV_COEFS = [
    1.0406546848089813, 0.6422230701725635, -0.026212793078773954,
    -0.18917346102528998, 0.03663524734377194, 0.11441748389589021,
    -0.023259469645663312, -0.06270170527196411, 0.023205093752784692,
    0.04313152013342015, -0.018415540276329505, -0.030795941717043043,
    0.011369088794049376, 0.01266040544297032, -0.016988349865962007,
    -0.020467498706326505, 0.004135891367235603, 0.0024823963204214966,
    -0.010840019119182868, -0.011797082201127044,
    ]


# Fit: TFZ= -3 -> 18
# Use: TFZ=  4 -> 10
TFZ_NONPOLAR_SOLVED_FRAC_CHEBYSHEV_COEFS = [
 -0.19552326319346666, 1.8649851395036177, -1.2150936747492587,
 0.9067488456954561, -0.9570223688009899, 0.9891869618763561,
 -0.7576200376815242, 0.49364441176419094, -0.34986647887756894,
 0.257829571743423, -0.055414541217036574, -0.17279616706198767,
 0.328353624714475, -0.4129911067675695, 0.5335599782204036,
 -0.6552756209909941, 0.7191178086864504, -0.7136431968840204,
 0.7066944667633923, -0.7018771386606003, 0.6525623963512635,
 -0.5684631934530187, 0.47036727433418174, -0.4075157167744924,
 0.31671053444831, -0.23712711445882487, 0.1443721102351674,
 -0.10822995600373658, 0.061040826040314036, -0.0344753971178862,
 ]

import operator

class TFZvsProbabilityCurve(object):
    """
    Expansion of empirical probability on Chebyshev polynomials
    """
    
    def __init__(self, polynomial, tfz_min, tfz_max):
        
        self.polynomial = polynomial
        self.tfz_min = tfz_min
        self.tfz_max = tfz_max
        
        
    def probability_for(self, tfz):
        
        if tfz < self.tfz_min:
            return 0.0
        
        elif self.tfz_min <= tfz and tfz < self.tfz_max:
            return min( max( self.polynomial.f( tfz ), 0.0 ), 1.0 )
        
        else:
            return 1.0
        
        
class TFZvsProbabilityCurveCollection(object):
    """
    Select appropriate scoring curve from a selection 
    """
    
    def __init__(self, condition_curve_pairs):
        
        self.condition_curve_pairs = condition_curve_pairs
        
        
    def probability_curve_for(self, description):
        
        for ( condition, curve ) in self.condition_curve_pairs:
            if condition( description = description ):
                return curve
            
        else:
            raise RuntimeError, "No appropriate curve found"
        
        
    @classmethod
    def PolarNonpolar(cls, polar, nonpolar):
        
        return cls(
            condition_curve_pairs = [
                ( cls.select_nonpolar, nonpolar ),
                ( cls.select_polar, polar ),
                ]
            )
        
    @staticmethod
    def select_nonpolar(description):
        
        return ( 1 < len( description.peaks )
            or description.space_group_info.number_of_continuous_allowed_origin_shifts() == 0 )
        
        
    @staticmethod
    def select_polar(description):
        
        return ( len( description.peaks ) <= 1
            and 1 <= description.space_group_info.number_of_continuous_allowed_origin_shifts() )
        
        
class TFZBasedProbabilityPeakDescription(object):
    """
    Stores data about a peak in structure that is used in TFZ-based scoring
    """
    
    def __init__(self, tfz, annotation = "component"):
        
        self.tfz = tfz
        self.annotation = annotation
        
        
    def __str__(self):
        
        return "%s TFZ=%.2f" % ( self.annotation, self.tfz )
        
        
class TFZBasedProbabilityStructureDescription(object):
    """
    Stores data about a structure that is used in TFZ-based scoring
    """
    
    def __init__(self, space_group_info, peaks):
        
        self.space_group_info = space_group_info
        self.peaks = peaks
        
        
class TFZBasedProbability(object):
    """
    Calculates probability scores for a structure
    """
    
    def __init__(self, collection, description):
        
        self.curve = collection.probability_curve_for( description = description )
        self.peaks = description.peaks
        self.individuals = [
            self.curve.probability_for( tfz = p.tfz ) for p in self.peaks
            ]
        
    
    @property
    def total(self):
         
        return reduce( operator.mul, self.individuals, 1.0 )
        
        
    def __str__(self):
        
        lines = [
            "%d. %s => P(i)=%.2f" % ( index, peak, prob )
            for ( index, ( peak, prob ) )
            in enumerate( zip( self.peaks, self.individuals ), start = 1 )
            ]
        lines.append( "P(total)=%.2f" % self.total )
        
        return "\n".join( lines )
        
        
from scitbx.math import chebyshev_polynome
from scitbx.array_family import flex

tfz_polar_single = TFZvsProbabilityCurve(
    polynomial = chebyshev_polynome( 20, -3, 15, flex.double( TFZ_POLAR_SOLVED_FRAC_CHEBYSHEV_COEFS ) ),
    tfz_min = 2.0,
    tfz_max = 7.5,
    )

tfz_nonpolar = TFZvsProbabilityCurve(
    polynomial = chebyshev_polynome( 30, -3, 18, flex.double( TFZ_NONPOLAR_SOLVED_FRAC_CHEBYSHEV_COEFS ) ),
    tfz_min = 4.0,
    tfz_max = 10.0,
    )

TFZ_POLAR_NONPOLAR = TFZvsProbabilityCurveCollection.PolarNonpolar(
    polar = tfz_polar_single,
    nonpolar = tfz_nonpolar,
    )