"""
Define the resolution functions for the data.

This defines classes for 1D and 2D resolution calculations.
"""
from scipy.special import erf
from numpy import sqrt
import numpy as np

SLIT_SMEAR_POINTS = 500
MINIMUM_RESOLUTION = 1e-8
# TODO: Q_EXTEND_STEPS is much bigger than necessary
Q_EXTEND_STEPS = 30   # number of extra q points above and below


class Resolution1D(object):
    """
    Abstract base class defining a 1D resolution function.

    *q* is the set of q values at which the data is measured.

    *q_calc* is the set of q values at which the theory needs to be evaluated.
    This may extend and interpolate the q values.

    *apply* is the method to call with I(q_calc) to compute the resolution
    smeared theory I(q).
    """
    q = None
    q_calc = None
    def apply(self, Iq_calc):
        """
        Smear *Iq_calc* by the resolution function, returning *Iq*.
        """
        raise NotImplementedError("Subclass does not define the apply function")

class Perfect1D(Resolution1D):
    """
    Resolution function to use when there is no actual resolution smearing
    to be applied.  It has the same interface as the other resolution
    functions, but returns the identity function.
    """
    def __init__(self, q):
        self.q_calc = self.q = q

    def apply(self, Iq_calc):
        return Iq_calc

class Pinhole1D(Resolution1D):
    r"""
    Pinhole aperture with q-dependent gaussian resolution.

    *q* points at which the data is measured.

    *dq* gaussian 1-sigma resolution at each data point.

    *q_calc* is the list of points to calculate, or None if this should
    be autogenerated from the *q,dq*.
    """
    def __init__(self, q, q_width, q_calc=None):
        #TODO: maybe add min_step=np.inf
        #*min_step* is the minimum point spacing to use when computing the
        #underlying model.  It should be on the order of
        #$\tfrac{1}{10}\tfrac{2\pi}{d_\text{max}}$ to make sure that fringes
        #are computed with sufficient density to avoid aliasing effects.

        # Protect against calls with q_width=0.  The extend_q function will
        # not extend the q if q_width is 0, but q_width must be non-zero when
        # constructing the weight matrix to avoid division by zero errors.
        # In practice this should never be needed, since resolution should
        # default to Perfect1D if the pinhole geometry is not defined.
        self.q, self.q_width = q, q_width
        self.q_calc = pinhole_extend_q(q, q_width) if q_calc is None else q_calc
        self.weight_matrix = pinhole_resolution(self.q_calc,
                self.q, np.maximum(q_width, MINIMUM_RESOLUTION))

    def apply(self, Iq_calc):
        return apply_resolution_matrix(self.weight_matrix, Iq_calc)

class Slit1D(Resolution1D):
    """
    Slit aperture with a complicated resolution function.

    *q* points at which the data is measured.

    *qx_width* slit width

    *qy_height* slit height
    """
    def __init__(self, q, qx_width, qy_width):
        if np.isscalar(qx_width):
            qx_width = qx_width*np.ones(len(q))
        if np.isscalar(qy_width):
            qy_width = qy_width*np.ones(len(q))
        self.q, self.qx_width, self.qy_width = [
            v.flatten() for v in (q, qx_width, qy_width)]
        self.q_calc = slit_extend_q(q, qx_width, qy_width)
        self.weight_matrix = \
            slit_resolution(self.q_calc, self.q, self.qx_width, self.qy_width)

    def apply(self, Iq_calc):
        return apply_resolution_matrix(self.weight_matrix, Iq_calc)


def apply_resolution_matrix(weight_matrix, Iq_calc):
    """
    Apply the resolution weight matrix to the computed theory function.
    """
    #print "apply shapes",Iq_calc.shape, self.weight_matrix.shape
    Iq = np.dot(Iq_calc[None,:], weight_matrix)
    #print "result shape",Iq.shape
    return Iq.flatten()

def pinhole_resolution(q_calc, q, q_width):
    """
    Compute the convolution matrix *W* for pinhole resolution 1-D data.

    Each row *W[i]* determines the normalized weight that the corresponding
    points *q_calc* contribute to the resolution smeared point *q[i]*.  Given
    *W*, the resolution smearing can be computed using *dot(W,q)*.

    *q_calc* must be increasing.
    """
    edges = bin_edges(q_calc)
    edges[edges<0.0] = 0.0 # clip edges below zero
    index = (q_width>0.0) # treat perfect resolution differently
    G = erf( (edges[:,None] - q[None,:]) / (sqrt(2.0)*q_width)[None,:] )
    weights = G[1:] - G[:-1]
    # w = np.sum(weights, axis=0); print "w",w.shape, w
    weights /= np.sum(weights, axis=0)[None,:]
    return weights


def pinhole_extend_q(q, q_width):
    """
    Given *q* and *q_width*, find a set of sampling points *q_calc* so
    that each point I(q) has sufficient support from the underlying
    function.
    """
    # If using min_step, you could do something like:
    #     q_calc = np.arange(min, max, min_step)
    #     index = np.searchsorted(q_calc, q)
    #     q_calc[index] = q
    # A refinement would be to assign q to the nearest neighbour.  A further
    # refinement would be to guard multiple q points between q_calc points,
    # either by removing duplicates or by only moving the values nearest the
    # edges.  For really sparse samples, you will want to remove all remaining
    # points that are not within three standard deviations of a measured q.
    q_min, q_max = np.min(q), np.max(q)
    extended_min, extended_max = np.min(q - 3*q_width), np.max(q + 3*q_width)
    nbins_low, nbins_high = Q_EXTEND_STEPS, Q_EXTEND_STEPS
    if extended_min < q_min:
        q_low = np.linspace(extended_min, q_min, nbins_low+1)[:-1]
        q_low = q_low[q_low>0.0]
    else:
        q_low = []
    if extended_max > q_max:
        q_high = np.linspace(q_max, extended_max, nbins_high+1)[:-1]
    else:
        q_high = []
    return np.concatenate([q_low, q, q_high])


def slit_resolution(q_calc, q, qx_width, qy_width):
    edges = bin_edges(q_calc) # Note: requires q > 0
    edges[edges<0.0] = 0.0 # clip edges below zero
    qy_min, qy_max = 0.0, edges[-1]

    # Make q_calc into a row vector, and q, qx_width, qy_width into columns
    # Make weights match [ q_calc X q ]
    weights = np.zeros((len(q),len(q_calc)),'d')
    q_calc = q_calc[None,:]
    q, qx_width, qy_width, edges = [
        v[:,None] for v in (q, qx_width, qy_width, edges)]

    # Loop for width (height is analytical).
    # Condition: height >>> width, otherwise, below is not accurate enough.
    # Smear weight numerical iteration for width>0 when height>0.
    # When width = 0, the numerical iteration will be skipped.
    # The resolution calculation for the height is done by direct integration,
    # assuming the I(q'=sqrt(q_j^2-(q+shift_w)^2)) is constant within
    # a q' bin, [q_high, q_low].
    # In general, this weight numerical iteration for width>0 might be a rough
    # approximation, but it must be good enough when height >>> width.
    E_sq = edges**2
    y_points = SLIT_SMEAR_POINTS if np.any(qy_width>0) else 1
    qy_step = 0 if y_points == 1 else qy_width/(y_points-1)
    for k in range(-y_points+1,y_points):
        qy = np.clip(q + qy_step*k, qy_min, qy_max)
        qx_low = qy
        qx_high = sqrt(qx_low**2 + qx_width**2)
        in_x = (q_calc>=qx_low)*(q_calc<=qx_high)
        qy_sq = qy**2
        weights += (sqrt(E_sq[1:]-qy_sq) - sqrt(qy_sq - E_sq[:-1]))*in_x
    w = np.sum(weights, axis=1); print "w",w.shape, w
    weights /= np.sum(weights, axis=1)[:,None]
    return weights


def slit_extend_q(q, qx_width, qy_width):
    return q


def bin_edges(x):
    """
    Determine bin edges from bin centers, assuming that edges are centered
    between the bins.

    Note: this uses the arithmetic mean, which may not be appropriate for
    log-scaled data.
    """
    if len(x) < 2 or (np.diff(x)<0).any():
        raise ValueError("Expected bins to be an increasing set")
    edges = np.hstack([
        x[0]  - 0.5*(x[1]  - x[0]),  # first point minus half first interval
        0.5*(x[1:] + x[:-1]),        # mid points of all central intervals
        x[-1] + 0.5*(x[-1] - x[-2]), # last point plus half last interval
        ])
    return edges

############################################################################
# unit tests
############################################################################
import unittest

class ResolutionTest(unittest.TestCase):

    def setUp(self):
        self.x = 0.001*np.arange(1,11)
        self.y = self.Iq(self.x)

    def Iq(self, q):
        "Linear function for resolution unit test"
        return 12.0 - 1000.0*q

    def test_perfect(self):
        """
        Perfect resolution and no smearing.
        """
        resolution = Perfect1D(self.x)
        Iq_calc = self.Iq(resolution.q_calc)
        output = resolution.apply(Iq_calc)
        np.testing.assert_equal(output, self.y)

    def test_slit_zero(self):
        """
        Slit smearing with perfect resolution.
        """
        resolution = Slit1D(self.x, 0., 0.)
        Iq_calc = self.Iq(resolution.q_calc)
        output = resolution.apply(Iq_calc)
        np.testing.assert_equal(output, self.y)

    @unittest.skip("slit smearing is still broken")
    def test_slit(self):
        """
        Slit smearing with height 0.005
        """
        resolution = Slit1D(self.x, 0., 0.005)
        Iq_calc = self.Iq(resolution.q_calc)
        output = resolution.apply(Iq_calc)
        # The following commented line was the correct output for even bins [see smearer.cpp for details]
        #answer = [ 9.666,  9.056,  8.329,  7.494,  6.642,  5.721,  4.774,  3.824,  2.871, 2.   ]
        answer = [ 9.0618,  8.6401,  8.1186,  7.1391,  6.1528,  5.5555,     4.5584,  3.5606,  2.5623, 2.    ]
        np.testing.assert_allclose(output, answer, atol=1e-4)

    def test_pinhole_zero(self):
        """
        Pinhole smearing with perfect resolution
        """
        resolution = Pinhole1D(self.x, 0.0*self.x)
        Iq_calc = self.Iq(resolution.q_calc)
        output = resolution.apply(Iq_calc)
        np.testing.assert_equal(output, self.y)

    def test_pinhole(self):
        """
        Pinhole smearing with dQ = 0.001 [Note: not dQ/Q = 0.001]
        """
        resolution = Pinhole1D(self.x, 0.001*np.ones_like(self.x),
                               q_calc=self.x)
        Iq_calc = 12.0-1000.0*resolution.q_calc
        output = resolution.apply(Iq_calc)
        answer = [ 10.44785079,   9.84991299,   8.98101708,
                  7.99906585,   6.99998311,   6.00001689,
                  5.00093415,   4.01898292,   3.15008701,   2.55214921]
        np.testing.assert_allclose(output, answer, atol=1e-8)

# Q, dQ, I(Q) for Igor default sphere model.
# combines CMSphere5.txt and CMSphere5smearsphere.txt from sasview/tests
# TODO: move test data into its own file?
SPHERE_RESOLUTION_TEST_DATA = """\
0.001278 0.0002847 2538.41176383
0.001562 0.0002905 2536.91820405
0.001846 0.0002956 2535.13182479
0.002130 0.0003017 2533.06217813
0.002414 0.0003087 2530.70378586
0.002698 0.0003165 2528.05024192
0.002982 0.0003249 2525.10408349
0.003266 0.0003340 2521.86667499
0.003550 0.0003437 2518.33907750
0.003834 0.0003539 2514.52246995
0.004118 0.0003646 2510.41798319
0.004402 0.0003757 2506.02690988
0.004686 0.0003872 2501.35067884
0.004970 0.0003990 2496.38678318
0.005253 0.0004112 2491.16237596
0.005537 0.0004237 2485.63911673
0.005821 0.0004365 2479.83657083
0.006105 0.0004495 2473.75676948
0.006389 0.0004628 2467.40145990
0.006673 0.0004762 2460.77293372
0.006957 0.0004899 2453.86724627
0.007241 0.0005037 2446.69623838
0.007525 0.0005177 2439.25775219
0.007809 0.0005318 2431.55421398
0.008093 0.0005461 2423.58785521
0.008377 0.0005605 2415.36158137
0.008661 0.0005750 2406.87009473
0.008945 0.0005896 2398.12841186
0.009229 0.0006044 2389.13360806
0.009513 0.0006192 2379.88958042
0.009797 0.0006341 2370.39776774
0.010080 0.0006491 2360.69528793
0.010360 0.0006641 2350.85169027
0.010650 0.0006793 2340.42023633
0.010930 0.0006945 2330.11206013
0.011220 0.0007097 2319.20109972
0.011500 0.0007251 2308.43503981
0.011780 0.0007404 2297.44820179
0.012070 0.0007558 2285.83853677
0.012350 0.0007713 2274.41290746
0.012640 0.0007868 2262.36219581
0.012920 0.0008024 2250.51169731
0.013200 0.0008180 2238.45596231
0.013490 0.0008336 2225.76495666
0.013770 0.0008493 2213.29618391
0.014060 0.0008650 2200.19110751
0.014340 0.0008807 2187.34050325
0.014620 0.0008965 2174.30529864
0.014910 0.0009123 2160.61632548
0.015190 0.0009281 2147.21038112
0.015470 0.0009440 2133.62023580
0.015760 0.0009598 2119.37907426
0.016040 0.0009757 2105.45234903
0.016330 0.0009916 2090.86319102
0.016610 0.0010080 2076.60576032
0.016890 0.0010240 2062.19214565
0.017180 0.0010390 2047.10550219
0.017460 0.0010550 2032.38715621
0.017740 0.0010710 2017.52560123
0.018030 0.0010880 2001.99124318
0.018310 0.0011040 1986.84662060
0.018600 0.0011200 1971.03389745
0.018880 0.0011360 1955.61395119
0.019160 0.0011520 1940.08291563
0.019450 0.0011680 1923.87672225
0.019730 0.0011840 1908.10656374
0.020020 0.0012000 1891.66297192
0.020300 0.0012160 1875.66789021
0.020580 0.0012320 1859.56357196
0.020870 0.0012490 1842.79468290
0.021150 0.0012650 1826.50064489
0.021430 0.0012810 1810.11533702
0.021720 0.0012970 1793.06840882
0.022000 0.0013130 1776.51153580
0.022280 0.0013290 1759.87201249
0.022570 0.0013460 1742.57354412
0.022850 0.0013620 1725.79397319
0.023140 0.0013780 1708.35831550
0.023420 0.0013940 1691.45256069
0.023700 0.0014110 1674.48561783
0.023990 0.0014270 1656.86525366
0.024270 0.0014430 1639.79847285
0.024550 0.0014590 1622.68887088
0.024840 0.0014760 1604.96421100
0.025120 0.0014920 1587.85768129
0.025410 0.0015080 1569.99297335
0.025690 0.0015240 1552.84580279
0.025970 0.0015410 1535.54074115
0.026260 0.0015570 1517.75249337
0.026540 0.0015730 1500.40115023
0.026820 0.0015900 1483.03632237
0.027110 0.0016060 1465.05942429
0.027390 0.0016220 1447.67682181
0.027670 0.0016390 1430.46495191
0.027960 0.0016550 1412.49232282
0.028240 0.0016710 1395.13182318
0.028520 0.0016880 1377.93439837
0.028810 0.0017040 1359.99528971
0.029090 0.0017200 1342.67274512
0.029370 0.0017370 1325.55375609
"""

class Pinhole1DTest(unittest.TestCase):

    def setUp(self):
        # sample q, dq, I(q) calculated by NIST Igor SANS package
        self.data = np.loadtxt(SPHERE_RESOLUTION_TEST_DATA.split('\n')).T
        self.pars = dict(scale=1.0, background=0.01, radius=60.0,
                         solvent_sld=6.3, sld=1)

    def Iq(self, q, dq):
        from sasmodels import core
        from sasmodels.models import sphere

        model = core.load_model(sphere)
        resolution = Pinhole1D(q, dq)
        kernel = core.make_kernel(model, [resolution.q_calc])
        Iq_calc = core.call_kernel(kernel, self.pars)
        result = resolution.apply(Iq_calc)
        return result

    def test_sphere(self):
        """
        Compare pinhole resolution smearing with NIST Igor SANS
        """
        q, dq, answer = self.data
        output = self.Iq(q, dq)
        np.testing.assert_allclose(output, answer, rtol=0.006)

    #TODO: is all sas data sampled densely enough to support resolution calcs?
    @unittest.skip("suppress sparse data test; not supported by current code")
    def test_sphere_sparse(self):
        """
        Compare pinhole resolution smearing with NIST Igor SANS on sparse data
        """
        q, dq, answer = self.data[:, ::20]  # Take every nth point
        output = self.Iq(q, dq)
        np.testing.assert_allclose(output, answer, rtol=0.006)


def main():
    """
    Run tests given is sys.argv.

    Returns 0 if success or 1 if any tests fail.
    """
    import sys
    import xmlrunner

    suite = unittest.TestSuite()
    suite.addTest(unittest.defaultTestLoader.loadTestsFromModule(sys.modules[__name__]))

    runner = xmlrunner.XMLTestRunner(output='logs')
    result = runner.run(suite)
    return 1 if result.failures or result.errors else 0


############################################################################
# usage demo
############################################################################

def _eval_demo_1d(resolution, title):
    from sasmodels import core
    from sasmodels.models import cylinder
    ## or alternatively:
    # cylinder = core.load_model_definition('cylinder')
    model = core.load_model(cylinder)

    kernel = core.make_kernel(model, [resolution.q_calc])
    Iq_calc = core.call_kernel(kernel, {'length':210, 'radius':500})
    Iq = resolution.apply(Iq_calc)

    import matplotlib.pyplot as plt
    plt.loglog(resolution.q_calc, Iq_calc, label='unsmeared')
    plt.loglog(resolution.q, Iq, label='smeared', hold=True)
    plt.legend()
    plt.title(title)
    plt.xlabel("Q (1/Ang)")
    plt.ylabel("I(Q) (1/cm)")

def demo_pinhole_1d():
    q = np.logspace(-3,-1,400)
    dq = 0.1*q
    resolution = Pinhole1D(q, dq)
    _eval_demo_1d(resolution, title="10% dQ/Q Pinhole Resolution")

def demo_slit_1d():
    q = np.logspace(-3,-1,400)
    qx_width = 0.005
    qy_width = 0.0
    resolution = Slit1D(q, qx_width, qy_width)
    _eval_demo_1d(resolution, title="0.005 Qx Slit Resolution")

def demo():
    import matplotlib.pyplot as plt
    plt.subplot(121)
    demo_pinhole_1d()
    plt.subplot(122)
    demo_slit_1d()
    plt.show()


if __name__ == "__main__":
    #demo()
    main()