# The following line helps with future compatibility with Python 3
# print must now be used as a function, e.g print('Hello','World')
from __future__ import (absolute_import, division, print_function, unicode_literals)

# import mantid algorithms, numpy and matplotlib
from mantid.simpleapi import *
import matplotlib.pyplot as plt
import numpy as np


# Set path for data
#data_nexus = "/Users/torbennielsen/Desktop/tmp3/McStasScript_Instr/LoKI/Union_abs_logger_event/3_16/nD_be/1E7_nr15/mccode.h5"

#data_nexus = "/Users/torbennielsen/Projects/23_OFF_detector/Mantid/mccode1E8_4cm.h5"
data_nexus = "/Users/torbennielsen/Desktop/mccode1E9.h5"
# Load McStas event data 
ws = Load(data_nexus)


# Rebin TOF
sample = Rebin(ws[0],'1000', False)

# Convert to lambda
sample = ConvertUnits(sample, 'Wavelength')
#sample = ConvertUnits(ws[0], 'Wavelength')
#sample = ConvertUnits(InputWorkspace=ws[0] , Target='Wavelength', ConvertFromPointData=False)


# Rebin lamba
binning = '1.75,0.001,16.5'
sample = Rebin(sample, binning)

# Load McStas normalization spectrum
mcstasMonitor = CloneWorkspace(mtd['Edet0.dat_ws'])

# Convert McStas data to Mantid histogram
mcstasMonitor = ConvertToHistogram(mcstasMonitor)
mcstasMonitor = rebin(mcstasMonitor, binning)
monitor     = CreateWorkspace(mcstasMonitor.dataX(0), mcstasMonitor.dataY(0), mcstasMonitor.dataE(0), UnitX='Wavelength')

# Setup q grid

binningMASK	= '0.005,0.005,0.55'


#Reduce to I(q) and normailzed
reduced_norm = Q1D(sample, binningMASK, WavelengthAdj='monitor')

# Save file
SaveAscii(InputWorkspace='reduced_norm', Filename='Mantid_Reduced_Iq_ver2.dat', Separator='Tab', ScientificFormat=True)


# Remove zero and inf
data_q = np.column_stack((reduced_norm.readX(0)[:-1], reduced_norm.readY(0)))
data_q = data_q[np.logical_not(np.isnan(data_q[:, 1]))]
data_q = data_q[np.logical_not(np.isinf(data_q[:, 1]))]
data_q = data_q[data_q[:, 1] != 0]

# Plot
plt.plot(data_q[:,0], data_q[:,1], 'k-', lw=2)
plt.yscale('log')
plt.xlabel('q [ 1/Å ]')
plt.ylabel('I(q) [a.u.]')
plt.show()