import sys
sys.path.append('/path/file/to/mantid/python/API') # /opt/Mantid/bin in linux systems
from mantid.simpleapi import LoadNexus, LoadSassena, Transpose, Rebin, Scale, SassenaFFT, ConvolveWorkspaces, Fit, mtd, DSFinterp

rootd='.'

LoadNexus(Filename='{0}/elastic.nxs'.format(rootd), OutputWorkspace='elastic')

 LoadNexus(Filename='{0}/exp200K.nxs'.format(rootd), OutputWorkspace='exp200K')

parametervalues='0.02 0.03 0.04 0.05 0.06 0.07 0.08 0.09 0.10 0.11 0.12 0.13 0.14 0.15'
for K in parametervalues.split():
    # Load spectra into Mantid workspace. Each file contains the intermediate structure factor I(Q,t) spectra for
    # 9 different values of Q (0.3, 0.5, 0.7,..,1.9). The names of the resulting workspaces are incSM0.03, .., incSM0.15
    LoadSassena(Filename='{0}/K{1}/T_200/fqt_inc_run1_rms2first.h5'.format(rootd,K), TimeUnit=1.0, OutputWorkspace='incSMK{0}'.format(K))

Transpose(InputWorkspace='incSMK0.11_fqt.Re', OutputWorkspace='incSMK0.11_fqt.Re')   # from I(Q,t) to I(t,Q)
Rebin(InputWorkspace='incSMK0.11_fqt.Re', Params=[0.2,0.2,2.0], OutputWorkspace='incSMK0.11_fqt.Re') # Rebin in Q to (0.3, 0.5,..,1.9)
Transpose(InputWorkspace='incSMK0.11_fqt.Re', OutputWorkspace='incSMK0.11_fqt.Re')   # from I(t,Q) back to I(Q,t)

Scale(InputWorkspace='incSMK0.11_fqt.Re', factor=0.5, Operation='Multiply', OutputWorkspace='incSMK0.11_fqt.Re')

SassenaFFT(InputWorkspace='incSMK{0}'.format(K), FFTonlyRealpart=1, DetailedBalance=1, Temp=200)

Rebin(InputWorkspace='incSMK{0}_sqw'.format(K), Params=[-0.2,0.0004,0.2], OutputWorkspace='incSMK{0}_sqw'.format(K))

ws=mtd['incSMK{0}_sqw'.format(K)] # simulated S(Q,E)
# Remove the elastic line for the spectra at each Q.
for iQ in range(9):
    Q=0.03+0.02*iQ # not neccessary but just to remind the Q-value
    SQE=ws.dataY(ih)
    SQE[499]=SQE[498]  
    SQE[500]=SQE[501]

parametervalues='0.02 0.03 0.04 0.05 0.06 0.07 0.08 0.09 0.10 0.11 0.12 0.13 0.14'.split()
inputworkspaces=['simSMK{0}'.format(K) for K in parametervalues]  # S(Q,E) from the simulations
targetvalues=[0.03+0.001*i for i in range(101)] #target K varies from 0.03 to 0.13 every 0.001
outputworkspaces=['out{0}'.format(K) for K in targetvalues] # interpolated S(Q,E)

DSFinterp(ParameterValues=[float(K) for K in parametervalues], Workspaces=inputworkspaces, LoadErrors=0,\
    TargetParameters=targetvalues, OutputWorkspaces=outputworkspaces,\
    LocalRegression=1, RegressionWindow=4, RegressionType='quadratic')

chi2values=[]
for K in targetvalues:
    # Our fitting model is a elastic line at Q=1.9 plus our simulated S(Q=1.9,E) plus a linear background
    fit_string    ='name=TabulatedFunction,Workspace=elastic,WorkspaceIndex=8,Scaling=1.0;'
    fit_string +='name=TabulatedFunction,Workspace=out{0},WorkspaceIndex=8,Scaling=1.0;'.format(K)
    fit_string +='name=LinearBackground,A0=0.0,A1=0.0'
    # Algorithm Fit carries out the fitting of the previous model to the experimental S(Q=1.9,E), which is index 8 of workspace exp200K.
    # Notice the energy range is [-0.13, 0.1] meV. Outside this energy range, the signal to noise ration of the experimental S(Q=1.9, E)
    # is not very reliable.
    Fit(fit_string, InputWorkspace='exp200K', WorkspaceIndex=8, StartX=-0.13, EndX=0.1, CreateOutput = 1 )
    ws=mtd['exp200K_Parameters'] #Workspace resulting from the fitting containing the optimal value of the fitting parameters and the Chi2
    chi2=ws.row(4)['Value']
    print 'K=', K, ' Chi2=', chi2
    chi2values.append(chi2)

for outws in outputworkspaces:  DeleteWorkspace(outws)

buf='#K Chi2\n'
for i in range(len(targetvalues)):  buf += '{0} {1}\n'.format(targetvalues[i], chi2values[i])
open('/tmp/chi2vsK.dat', 'w').write(buf)

guess={'Scaling':1.0, 'Intensity':1.0, 'K':0.065}

parametervalues='0.02 0.03 0.04 0.05 0.06 0.07 0.08 0.09 0.10 0.11 0.12 0.13 0.14'
inputworkspaces=' '.join( ['simSMK{0}'.format(K) for K in parametervalues.split()])

for iw in [8,7,6,5,4,3,2,1,0]:
    Q=0.3+iw*0.2
    # Our model is a elastic line plus a  structure factor interpolated from the simulated S(Q,E), plus a linear background
    fit_string = 'name=TabulatedFunction,Workspace=elastic,WorkspaceIndex={0},Scaling={1},constraints=(0.0001<Scaling);'.format(iw,guess['Scaling'])+\
        'name=DSFinterp1DFit,InputWorkspaces="{0}",ParameterValues="{1}",'.format(inputworkspaces,parametervalues) +\
        'LoadErrors=0,LocalRegression=1,RegressionType=quadratic,RegressionWindow=4,' +\
        'WorkspaceIndex={0},Intensity={1},TargetParameter={2},'.format(iw,guess['Intensity'],guess['K']) +\
        'constraints=(0.0001<Intensity);' +\
        'name=LinearBackground,A0=0.0,A1=0.0'
    Fit(fit_string, InputWorkspace='exp200K', WorkspaceIndex=iw, StartX=-0.13, EndX=0.1, CreateOutput = 1 )
    ws=mtd['exp200K_Parameters'] # Workspace resulting from the fitting containing the optimal value of the fitting parameters and the Chi2
    Koptimal=ws.row(2)['Value']
    print 'Q=', Q, 'L={0:5.1f}'.format( 6.2832/Q), 'K={0:6.4f}'.format(Koptimal), 'Chi2={0:3.1f}'.format(ws.row(5)['Value'])
    # Use the optimal intensities of the elastic line and simulated quasi-elastic line as guesses for the fitting of the S(Q,E) for the next Q
    guess['Scaling']=ws.row(0)['Value']
    guess['Intensity']=ws.row(1)['Value']