from mantid.api import *
import numpy as np
import scipy.special as sps


def ICGauss(a,b,R,s,t) :
    # Based on Jack Carpenter notes from "Elements of Slow Neutron Scattering"
    # Modified to take Gaussian sigma directly using recursive formula (DLA)

    prefactor = a/2.0/np.sqrt(2.0*np.pi)/s
    exp_a = np.exp(-a*t)
    c1 = (1.0 - R)*a**2
    d = a - b    
    c2 = 2.0*R*(a**2*b)/d**3

    prefactor_gn = np.sqrt(np.pi/2.0)*s
    ga0 = prefactor_gn*np.exp(0.5*a**2*s**2)*sps.erfc((a*s - t/s)/np.sqrt(2.0))
    gb0 = prefactor_gn*np.exp(0.5*b**2*s**2)*sps.erfc((b*s - t/s)/np.sqrt(2.0))
    ga1 = s**2*(-a*ga0 + np.exp(a*t - t**2/2.0/s**2))
    ga2 = s**2*(ga0 - a*ga1 - t*np.exp(a*t - t**2/2.0/s**2))    

    term1 = (c1*t**2 - c2*(1.0 + d*t + 0.5*d**2*t**2))*exp_a*ga0
    term2 = c2*np.exp(-b*t)*gb0
    term3 = (2*c1*t - c2*d*(1.0+d*t))*exp_a*ga1
    term4 = (c1 - c2*d**2/2.0)*exp_a*ga2

    return prefactor*(term1 + term2 + term3 + term4)


# Class definition
class ICxGmonitor2(IFunction1D):
 	# Use Ikeda-Carpenter parameters at source, scale to monitor 2 position using measured monitor 1 Gaussian width (sigma)
	# Correct for monitor thickness and time bins - assume both monitors the same

	def init(self):
		self.declareParameter("Intensity",1.0)
		self.declareParameter("Tpeak",1000.0)
		self.declareParameter("Alpha",0.2)
		self.declareParameter("Beta",0.04)
		self.declareParameter("R",0.6)
		self.declareParameter("Sigma_m1",3.0)
		self.declareParameter("L_fermi",11.61)
		self.declareParameter("L_m1",11.835)
		self.declareParameter("L_m2",18.5)
		self.declareParameter("Thickness_mon",0.01)
		self.declareParameter("Tbin",1.0)
		self.declareParameter("E_i",100.0)
		
	def function1D(self,xvals):
		# Load parameters
		I0=self.getParameterValue("Intensity")	
		tp = self.getParameterValue("Tpeak")
		a=self.getParameterValue("Alpha")	
		b = self.getParameterValue("Beta")
		R=self.getParameterValue("R")	
		s_m1 = self.getParameterValue("Sigma_m1")
		L_fermi = self.getParameterValue("L_fermi")
		L_m1 = self.getParameterValue("L_m1")
		L_m2 = self.getParameterValue("L_m2")
		Thick_mon = self.getParameterValue("Thickness_mon")
		Tbin = self.getParameterValue("Tbin")
		E_i = self.getParameterValue("E_i")

		f_source = L_m2/L_fermi - 1.0
		f_mon = L_m2/L_m1

		a_m2 = a/f_source
		b_m2 = b/f_source

		tmean = 3.0/a_m2 + R/b_m2

		s_tbin = Tbin/3.4641   # sqrt(12) for top hat
		s_tof = 1e6*Thick_mon/np.sqrt(E_i/5.2276e-6)/3.4641  # thickness/v_i in microseconds
		sigma_m1 = np.sqrt(np.abs(s_m1**2 - s_tof**2 - s_tbin**2))
		sigma_m2 = sigma_m1*f_mon
		s_m2 = np.sqrt(sigma_m2**2 + s_tof**2 + s_tbin**2)		

		return I0*ICGauss(a_m2,b_m2,R,s_m2,tp+tmean-xvals) 

# Register algorithm with Mantid
FunctionFactory.subscribe(ICxGmonitor2)

