#!/usr//bin/env python
# vim: ts=8 sts=4 sw=4 expandtab
# Author: Douglas Clowes (dcl@ansto.gov.au) 2012-07-18

import math

class CubicPowderSample:
    def __init__(self):
        self.aspacing = 3.5238 # Angstroms
        self.wavelength = 2.34548 # Angstroms
        self.wavelength_error = 0.00 # Angstroms
        self.peakpos_error = 0.0 # degrees
        self.mEv = 0.0

    def set_wavelength_from_theta_pg002(self, theta):
        dd = 3.35416e-10 # For Monochromator PG 0 0 2
        h_for_joules = 6.62606957e-34
        neutron_mass = 1.674927351e-27
        ev_to_joules = 1.602176565e-19
        angle = theta
        myWavelength = 2.0 * dd * math.sin(angle * 3.141592654 / 180.0)
        if myWavelength > 0:
            temp = h_for_joules / (neutron_mass * myWavelength)
            joules = 0.5 * neutron_mass * (temp * temp)
            self.mEv = 1000 * joules / ev_to_joules
            self.wavelength = 1e10 * myWavelength
        else:
            self.mEv = 0.0
            self.wavelength = 0.0
        print "theta = %g, wl = %g, mEv = %g" % (angle, self.wavelength, self.mEv)

    
    def atod(self, the_aspacing, h, k, l):
        result = math.sqrt(the_aspacing**2 /(h**2 + k**2 + l**2))
        return result

    def peak_position(self, dspacing, wavelength):
        try:
            theta = math.asin(wavelength / (2.0 * dspacing)) * 180 / 3.141592654
            return 2.0 * theta
        except:
            print "error in peak_position(", dspacing, wavelength, ")"
            return 0.0

    def funx(self, x, mu, sigma):
        try:
            factor = 1.0 / (sigma * math.sqrt(2.0 * math.pi))
            y = factor * math.exp(-0.5 * ((x - mu) / sigma)**2)
            return y
        except:
            print "error in func(", x, mu, sigma, ")"
            return 0.0

    def calc(self, x):
        peaks = []
        peaks.append([0.5, 0.5, 0.5, 1.0])
        peaks.append([1, 0, 0, 0.5])
        peaks.append([1, 1, 0, 0.8])
        peaks.append([1, 1, 1, 0.9])
        peaks.append([2, 0, 0, 0.4])
        peaks.append([1.5, 0.5, 0.5, 0.1])
        peaks.append([2, 1, 0, 0.5])
        peaks.append([2, 1, 1, 0.4])
        peaks.append([2, 2, 0, 0.3])
        sigma = 0.1 + 0.4 * math.sin(math.pi * x / 180.0)
        y = 0.0
        for peak in peaks:
            peakpos = self.peak_position(self.atod(self.aspacing, peak[0], peak[1], peak[2]), self.wavelength + self.wavelength_error)
            y += self.funx(x, peakpos + self.peakpos_error, sigma) * peak[3]
        return y

if __name__ == '__main__':
    nickel = CubicPowderSample()
    px = []
    for i in range(200, 891):
        x = 0.1 * i
        px.append(x)
    from sys import argv
    from mpl_toolkits.axes_grid.axislines import SubplotZero
    import matplotlib.pyplot as plt
    fig = plt.figure(1)
    ax = SubplotZero(fig, 111)
    fig.add_subplot(ax)
    for direction in ["xzero", "yzero"]:
        #ax.axis[direction].set_axisline_style("-|>")
        ax.axis[direction].set_visible(True)

    for direction in ["left", "right", "bottom", "top"]:
        ax.axis[direction].set_visible(False)

    if len(argv) > 1:
        for i in range(1, len(argv)):
            print "argv[%d]:" % i, argv[i]
            nickel.set_wavelength_from_theta_pg002(float(argv[i]))
            py = []
            for x in px:
                y = nickel.calc(x)
                py.append(y)
            ax.plot(px, py)
    else:
        py = []
        for x in px:
            y = nickel.calc(x)
            py.append(y)
        ax.plot(px, py)
        print "len:", len(px), len(py)

    plt.show()