#! /usr/bin/pythonw
# (Change the above to your pythin distribution)
#
#
# Routine to fit Chebychev Polynomials to the A* linear attenuation factor in a cylinder for a given Mu*R
# Craig M. Brown 09/19, NIST Center for Neutron Research
#

# Must have have numpy  and matplotlib installed

import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button, RadioButtons
import numpy as np
from numpy.polynomial import Chebyshev as T

#Tabulation of A* from https://archive.org/details/InternationalTablesForX-rayCrystallographyVol2/page/n309
# and Table 6.3.3.2 in Volume C of the International Tables for Crystallography (p. 523 in the 1995 ed.)

Astar = {0.0 : [1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000],
        0.1 : [1.1843,1.1843,1.1842,1.1840,1.1838,1.1835,1.1832,1.1828,1.1823,1.1818,1.1813,1.1808,1.1802,1.1798,1.1793,1.1790,1.1787,1.1785,1.1785],
        0.2 : [1.4009,1.4007,1.4002,1.3995,1.3984,1.3970,1.3953,1.3934,1.3912,1.3889,1.3865,1.3841,1.3818,1.3796,1.3777,1.3761,1.3749,1.3741,1.3739],
        0.3 : [1.6548,1.6544,1.6531,1.6510,1.6481,1.6443,1.6398,1.6347,1.6290,1.6230,1.6169,1.6108,1.6049,1.5994,1.5946,1.5906,1.5876,1.5857,1.5851],
        0.4 : [1.9522,1.9513,1.9485,1.9439,1.9376,1.9296,1.9201,1.9094,1.8979,1.8857,1.8733,1.8611,1.8495,1.8388,1.8293,1.8215,1.8157,1.8121,1.8108],
        0.5 : [2.2996,2.2979,2.2926,2.2840,2.2721,2.2572,2.2398,2.2204,2.1996,2.1781,2.1564,2.1352,2.1152,2.0969,2.0809,2.0677,2.0579,2.0518,2.0497],
        0.6 : [2.7047,2.7017,2.6926,2.6775,2.6570,2.6317,2.6023,2.5701,2.5359,2.5010,2.4662,2.4327,2.4012,2.3728,2.3480,2.3277,2.3126,2.3033,2.3001],
        0.7 : [3.1762,3.1712,3.1561,3.1315,3.0982,3.0575,3.0111,2.9607,2.9081,2.8549,2.8028,2.7530,2.7068,2.6653,2.6295,2.6003,2.5786,2.5651,2.5606],
        0.8 : [3.7236,3.7157,3.6919,3.6532,3.6015,3.5392,3.4691,3.3941,3.3169,3.2400,3.1656,3.0953,3.0307,2.9732,2.9239,2.8839,2.8542,2.8359,2.8297],
        0.9 : [4.3578,4.3456,4.3093,4.2507,4.1733,4.0812,3.9792,3.8718,3.7629,3.6560,3.5538,3.4584,3.3717,3.2951,3.2299,3.1772,3.1383,3.1142,3.1061],
        1.0 : [5.0907,5.0724,5.0185,4.9323,4.8196,4.6877,4.5439,4.3948,4.2461,4.1022,3.9664,3.8413,3.7286,3.6298,3.5462,3.4790,3.4295,3.3990,3.3886],
        1.1 : [5.9356,5.9089,5.8305,5.7065,5.5466,5.3624,5.1649,4.9636,4.7660,4.5776,4.4022,4.2424,4.0998,3.9759,3.8717,3.7882,3.7269,3.6891,3.6763],
        1.2 : [6.9070,6.8690,6.7570,6.5820,6.3600,6.1090,5.8436,5.5782,5.3219,5.0811,4.8598,4.6604,4.4842,4.3322,4.2051,4.1038,4.0295,3.9838,3.9682],
        1.3 : [8.0210,7.9670,7.8100,7.5680,7.2660,6.9290,6.5810,6.2380,5.9125,5.6110,5.3376,5.0938,4.8805,4.6976,4.5456,4.4248,4.3365,4.2821,4.2636],
        1.4 : [9.2940,9.2190,9.0030,8.6740,8.2680,7.8260,7.3760,6.9420,6.5360,6.1660,5.8341,5.5413,5.2873,5.0711,4.8922,4.7506,4.6471,4.5835,4.5619],
        1.5 : [10.7460,10.6430,10.3490,9.9070,9.3720,8.8000,8.2300,7.6890,7.1920,6.7440,6.3480,6.0020,5.7036,5.4516,5.2441,5.0804,4.9609,4.8875,4.8625],
        1.6 : [12.3970,12.2570,11.8620,11.2760,10.5810,9.8520,9.1410,8.4770,7.8770,7.3440,6.8770,6.4730,6.1280,5.8385,5.6007,5.4136,5.2773,5.1935,5.1650],
        1.7 : [14.2670,14.0800,13.5550,12.7880,11.8970,10.9820,10.1060,9.3040,8.5890,7.9630,7.4200,6.9550,6.5610,6.2310,5.9610,5.7499,5.5960,5.5014,5.4691],
        1.8 : [16.3790,16.1310,15.4410,14.4500,13.3230,12.1890,11.1250,10.1680,9.3270,8.6000,7.9760,7.4460,7.0000,6.6280,6.3260,6.0890,5.9166,5.8107,5.7746],
        1.9 : [18.7600,18.4300,17.5300,16.2670,14.8580,13.4700,12.1940,11.0660,10.0890,9.2530,8.5440,7.9460,7.4440,7.0300,6.6930,6.4300,6.2390,6.1210,6.0810],
        2.0 : [21.4300,21.0000,19.8400,18.2400,16.5000,14.8240,13.3110,11.9950,10.8710,9.9210,9.1220,8.4520,7.8950,7.4350,7.0630,6.7730,6.5620,6.4330,6.3890],
        2.1 : [24.4100,23.8700,22.3900,20.3800,18.2500,16.2470,14.4720,12.9530,11.6730,10.6020,9.7090,8.9650,8.3490,7.8430,7.4360,7.1180,6.8870,6.7450,6.6970],
        2.2 : [27.7400,27.0400,25.1700,22.6900,20.1100,17.7400,15.6750,13.9380,12.4930,11.2950,10.3040,9.4840,8.8080,8.2550,7.8100,7.4640,7.2130,7.0590,7.0060],
        2.3 : [31.44, 30.55, 28.20, 25.16, 22.07, 19.29,16.92, 14.947, 13.328, 11.999, 10.906, 10.008, 9.271, 8.669, 8.187, 7.812, 7.540, 7.372, 7.315],
        2.4 : [35.3, 34.2, 31.4, 27.7, 24.1, 20.9, 18.2, 15.987, 14.177, 12.711, 11.515, 10.537, 9.736, 9.086, 8.565, 8.161, 7.868, 7.687, 7.625],
        2.5 : [40.06, 38.65,35.05, 30.59, 26.28, 22.56,19.50, 17.03, 15.040,13.433, 12.130, 11.069, 10.205, 9.505,8.945, 8.511, 8.196, 8.002 ,7.935],
        2.6 : [44.7,43.1,38.7,33.4,28.5,24.2,20.8,18.1,15.9,14.2,12.8,11.6,10.7,9.93,9.33,8.88,8.53,8.33,8.27],
        2.7 : [50.1,48.1,42.8,36.5,30.8,26,22.2,19.2,16.8,14.9,13.4,12.2,11.2,10.4,9.72,9.23,8.87,8.64,8.58],
        2.8 : [56.1,53.5,47.2,39.7,33.2,27.8,23.6,20.3,17.7,15.6,14,12.7,11.7,10.8,10.1,9.59,9.2,8.96,8.9],
        2.9 : [62.5,59.4,51.8,43.1,35.7,29.6,25,21.4,18.6,16.4,14.6,13.2,12.1,11.2,10.5,9.95,9.53,9.28,9.21],
        3.0 : [69.5,65.8,56.7,46.6,38.2,31.5,26.4,22.5,19.5,17.1,15.3,13.8,12.6,11.6,10.9,10.3,9.86,9.6,9.53],
        3.1 : [77.2,72.6,61.8,50.3,40.8,33.4,27.9,23.7,20.4,17.9,15.9,14.3,13.1,12.1,11.3,10.7,10.2,9.92,9.84],
        3.2 : [85.4,79.9,67.3,54,43.5,35.4,29.3,24.8,21.4,18.7,16.6,14.9,13.6,12.5,11.7,11,10.5,10.2,10.2],
        3.3 : [94.2,87.6,72.9,57.9,46.2,37.3,30.8,26,22.3,19.4,17.2,15.5,14.1,12.9,12,11.4,10.9,10.6,10.5],
        3.4 : [104,95.9,78.9,61.9,48.9,39.3,32.3,27.1,23.2,20.2,17.9,16,14.6,13.4,12.4,11.7,11.2,10.9,10.8],
        3.5 : [114,105,85,65.9,51.7,41.3,33.8,28.3,24.2,21,18.5,16.6,15.1,13.8,12.8,12.1,11.5,11.2,11.1],
        3.6 : [125,114,91.4,70.1,54.6,43.3,35.3,29.5,25.1,21.7,19.2,17.1,15.6,14.2,13.2,12.5,11.9,11.5,11.4],
        3.7 : [136,124,98,74.4,57.4,45.4,36.9,30.7,26.1,22.5,19.8,17.7,16.1,14.7,13.6,12.8,12.2,11.8,11.7],
        3.8 : [149,134,105,78.7,60.4,47.5,38.4,31.9,27,23.3,20.5,18.3,16.6,15.1,14,13.2,12.6,12.2,12.1],
        3.9 : [162,145,112,83.1,63.3,49.5,39.9,33.1,28,24.1,21.2,18.8,17.1,15.6,14.4,13.6,12.9,12.5,12.4],
        4.0 : [175,156,119,87.6,66.3,51.6,41.5,34.3,28.9,24.9,21.8,19.4,17.6,16,14.8,13.9,13.2,12.8,12.7],
        4.1 : [190,168,127,92.7,69.3,53.8,43.1,35.5,29.9,25.7,22.5,20,18.1,16.4,15.2,14.3,13.6,13.1,13],
        4.2 : [206,180,134,96.7,72.3,55.9,44.6,36.7,30.9,26.5,23.2,20.6,18.6,16.9,15.6,14.6,13.9,13.5,13.3],
        4.3 : [222,193,142,101,75.4,58.0,46.2,37.9,31.8,27.3,23.8,21.1,19.1,17.3,16,15,14.3,13.8,13.7],
        4.4 : [239,206,150,106,78.5,60.2,47.8,39.1,32.8,28.1,24.5,21.7,19.6,17.8,16.4,15.4,14.6,14.1,14]}

mur=4.0
ncheb = 12
x=[0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70,75, 80, 85, 90]
y = Astar[mur]

#Fit to Chebychev
p,stats = T.fit(x, y, ncheb-1, full=True)
xx, yy = p.linspace()

fig, ax1 = plt.subplots()
plt.subplots_adjust(left=0.25, bottom=0.3)
s, = plt.plot(x, y, 'bo')
l, = plt.plot(xx,yy, '-r')
axes = plt.gca()

ax1.set_title('Goodness-of-fit: '+ "%.6f" % (stats[0])[0])
axes.set_xlabel('Angle (degrees)')
axes.set_ylabel('A*')

axcolor = 'lightgoldenrodyellow'
btcolor = 'lightgreen'

axnscheb = plt.axes([0.25, 0.15, 0.65, 0.03], facecolor=axcolor)
sncheb = Slider(axnscheb, '#Coeffs', 1, 12.0, valinit=ncheb, valstep=1)
axval = plt.axes([0.25, 0.1, 0.65, 0.03], facecolor=axcolor)
sval = Slider(axval, 'Mu*r', 0.0, 4.4, valinit=mur, valstep=0.1)

bx = plt.axes([0.25, 0.02, .65, 0.05])
button  = Button (bx, 'save coefficients to file',color=btcolor)

def update1(val):
    global p
    global mur
    global x
    global y
    global xx
    global yy
    global ncheb
    global Astar
    mur = sval.val
    #deal with the decimal rounding errors of the slider:
    mur = int(mur*10.0)/10
    y = Astar[mur]
    p,stats = T.fit(x, y, ncheb-1, full=True)
    xx, yy = p.linspace()
    s.set_ydata(y)
    l.set_ydata(yy)
    axes.set_ylim(0, max(y)+0.1*max(y))
    ax1.set_title('Goodness-of-fit: '+ "%.6f" % (stats[0])[0])
    fig.canvas.draw_idle()

def update2(val):
    global p
    global mur
    global x
    global y
    global xx
    global yy
    global Astar
    global ncheb
    ncheb = sncheb.val
    #deal with the decimal rounding errors of the slider:
    ncheb = int(ncheb)
    y = Astar[mur]
    p,stats = T.fit(x, y, ncheb-1, full=True)
    xx, yy = p.linspace()
    s.set_ydata(y)
    l.set_ydata(yy)
    axes.set_ylim(0, max(y)+0.1*max(y))
    ax1.set_title('Goodness-of-fit: '+ "%.6f" % (stats[0])[0])
    fig.canvas.draw_idle()

def saveme(event):
    global ncheb
    global mur
    global p
    stringer = str(ncheb) + ', '
    for x in p.coef:
        stringer = stringer + str(x) + ', '
    #save the following to a file 'abs.corr'
    print('A* absorption coefficients for Mu.R = ' + str(mur))
    print(stringer)
    fileout =  "abs.corr."
    file = open(fileout, "w")
    file.write('A* absorption coefficients for Mu.R = ' + str(mur) + "\n")
    file.write(stringer)
    file.close()

button.on_clicked(saveme)
sval.on_changed(update1)
sncheb.on_changed(update2)

plt.show()