#Import the necessary libraries

import string
import math
import operator

####################################
#                                  #
#       define functions           #
#                                  #
####################################


#********************************************************************
# Function   : rdcsv                                                *
# Purpose    : Read in data from a csv file                         *
# Parameters : name (string)                                        *
# Returns    : list containing a list of data in each line of file  *
#********************************************************************
def rdcsv(name):
   f = open(name,"r")
   data = []
   tmp = f.readline()
   while tmp !="":
       onerow = tmp.split(",")
       last = len(onerow)-1
       k = len(onerow[last])
       if k!= 1:
          onerow[last] = onerow[last][0:k-1]
       if len(onerow) !=  0:
          data.append(onerow)
       tmp = f.readline()
   f.close()
   return(data)

#********************************************************************
# Function   : outcsv                                               *
# Purpose    : Write data out to a csv file                         *
# Parameters : name (string)                                        *
#              data (list of data points)                           *
# Returns    : n/a                                                  *
#********************************************************************
def outcsv(name,data):
   f = open(name,"w")
   for i in range(0,len(data)):
       tmp = ""
       for j in range(0,len(data[i])):
           tmp = tmp + str(data[i][j]) + ","
       tmp = tmp[0:len(tmp)-1]+"\n"
       f.writelines(tmp)
   f.close()

#********************************************************************
# Function   : norm                                                 *
# Purpose    : Integrate the data using Simpson's Rule              *
# Parameters : data (list of data points)                           *
# Returns    : integral of the data                                 *
#********************************************************************
def norm(data):
   sum=0
   dx = abs(data[1][0]-data[0][0])
   for i in range(0,len(data)):
       sum+=data[i][1]*2**((i % 2)+1)
   return(sum*dx/3.0)

#********************************************************************
# Function   : normalize                                            *
# Purpose    : Normalize data so that the integral is 1             *
# Parameters : data (list of data points)                           *
# Returns    : list of data points                                  *
#********************************************************************
def normalize(data):
   area = norm(data)
   for i in range(0,len(data)):
       data[i][1] = data[i][1]/area
   return(data)

#********************************************************************
# Function   : agm                                                  *
# Purpose    : Compute the arithmetic-geometric mean M(a,b)         *
# Parameters : a,b (float)                                          *
# Returns    : float                                                *
#********************************************************************
def agm(a,b):
   a1 = a
   b1 = b
   for i in range(0,5):
       a2 = (a1+b1)/2.0
       b2 = math.sqrt(a1*b1)
       a1=a2
       b1=b2
   return(a1)

#********************************************************************
# Function   : intensity                                            *
# Purpose    : Compute the intensity function for the ideal powder  *
#              spectrum (nonaxially symmetric case)                 *
# Parameters : delta,d1,d2,d3 (float)                               *
# Returns    : float                                                *
#********************************************************************
def intensity(delta,d3,d2,d1):
   if (delta<=d3):
       return(0)
   elif (delta<d2):
       m = math.sqrt((d1-d2)*(delta-d3))
       kappa = math.sqrt(((d1-delta)*(d2-d3))/((delta-d3)*(d1-d2)))
       return(1/(2*m*kappa*agm(1,math.sqrt(1-1/(kappa*kappa)))))
   elif (delta==d2):
       return(intensity(d2+.0001,d3,d2,d1))
   elif (delta<d1):
       m = math.sqrt((d1-d2)*(delta-d3))
       kappa = math.sqrt(((d1-delta)*(d2-d3))/((delta-d3)*(d1-d2)))
       return(1/(2*m*agm(1,math.sqrt(1-kappa*kappa))))
   else:
       return(0)


#********************************************************************
# Function   : convolution                                          *
# Purpose    : Compute the convolution of two data sets using       *
#              Simpson's Rule                                       *
# Parameters : a,b (list)                                           *
#              dx (float)                                           *
# Returns    : list of floats                                       *
#********************************************************************
def convolution(a,b,dx):
   c = []
   if (len(a)!=len(b)):
       print("ERROR")
   else:
       for i in range(0,2*len(a)):
           c.append(0)
           for j in range(0,len(a)):
               k = i-j
               if not((k<0) or (k>len(a)-1)):
                   c[i]+=a[k]*b[j]*2**((j % 2)+1)
               else:
                   c[i]+=0
           c[i] = c[i]*dx/3.0

   return(c)

#********************************************************************
# Function   : Gaussian                                             *
# Purpose    : Compute the Gaussian at x with mean mu               *
#              and standard deviation sigma                         *
# Parameters : x,mu,sigma (float)                                   *
# Returns    : float                                                *
#********************************************************************
def Gaussian(x,mu,sigma):
   pi = 4*math.atan(1)
   return(1/(sigma*math.sqrt(2*pi))*math.exp(-((x-mu)*(x-mu))/(2*sigma*sigma)))

#********************************************************************
# Function   : GaussianFirstDeriv                                   *
# Purpose    : Compute the first derivative of the Gaussian at      *
#              x with mean mu and standard deviation sigma          *
# Parameters : x,mu,sigma (float)                                   *
# Returns    : float                                                *
#********************************************************************
def GaussianFirstDeriv(x,mu,sigma):
   pi = 4*math.atan(1)
   coeff = -(x-mu)/(sigma*sigma*sigma*math.sqrt(2*pi))
   exponential = math.exp(-((x-mu)*(x-mu))/(2*sigma*sigma))
   return(coeff*exponential)

#********************************************************************
# Function   : GaussianSecondDeriv                                  *
# Purpose    : Compute the second derivative of the Gaussian at     *
#              x with mean mu and standard deviation sigma          *
# Parameters : x,mu,sigma (float)                                   *
# Returns    : float                                                *
#********************************************************************
def GaussianSecondDeriv(x,mu,sigma):
   pi = 4*math.atan(1)
   coeff = 1/(sigma*sigma*sigma*math.sqrt(2*pi))*((x-mu)*(x-mu)/(sigma*sigma)-1)
   exponential = math.exp(-((x-mu)*(x-mu))/(2*sigma*sigma))
   return(coeff*exponential)
#********************************************************************
# Function   : GaussianThirdDeriv                                  *
# Purpose    : Compute the third derivative of the Gaussian at     *
#              x with mean mu and standard deviation sigma          *
# Parameters : x,mu,sigma (float)                                   *
# Returns    : float                                                *
#********************************************************************
def GaussianThirdDeriv(x,mu,sigma):
   pi = 4*math.atan(1)
   coeff = 1/(sigma**7*math.sqrt(2*pi))*((x-mu)*(3*sigma*sigma-x*x+2*x*mu-mu*mu))
   exponential = math.exp(-((x-mu)*(x-mu))/(2*sigma*sigma))
   return(coeff*exponential)

#********************************************************************
# Function   : findzeros                                            *
# Purpose    : Compute the list of independent variable values at   *
#              which the dependent variable values cross the        *
#              independent axis                                     *
# Parameters : data (list of data points)                           *
# Returns    : list of zero crossings                               *
#********************************************************************
def findzeros(data):
   zeros = []
   for i in range(0,len(data)-1):
       if data[i][1]*data[i+1][1]<0:
           zeros.append(data[i])
   return(zeros)

#********************************************************************
# Function   : phi                                                  *
# Purpose    : Compute l^2 norm of the difference between the data  *
#              lists                                                *
# Parameters : data1 (list of data points)                          *
#              data2 (list of data points)                          *
# Returns    : float                                                *
#********************************************************************
def phi(data1,data2):
    dx = data1[1][0]-data1[0][0]
    k=abs(int(round((data1[0][0]-data2[0][0])/dx)))
    sum = 0
    for i in range(0,len(data1)):
        sum += (data1[i][1] - data2[k+i][1])**2
    norm = math.sqrt(sum*dx)
    return(norm)

#********************************************************************
# Function   : edgedetect                                           *
# Purpose    : Compute possible independent variable values at which*
#              the data has an edge and at which the data will have *
#              a maximum                                            *
# Parameters : data (list of data points)                           *
# Returns    : list of first deriv. zero crossings                  *
#              list of second deriv. zero crossings                 *
#              list of minimum values of the second deriv.          *
#********************************************************************
def edgedetect(data):
    #compute stepsize dx from the data set
    dx = abs(data[1][0]-data[0][0])

    #extract the dependent variable values from the data set for analysis
    n = len(data)
    dep_var_values = []
    for i in range(0,n):
       dep_var_values.append(data[i][1])

    #generate first and second derivatives of the broadening function
    gx0 = -n*dx/2.0 
    gaussian_first_deriv = []
    gaussian_second_deriv = []
    indep_var_values = []
    for i in range(0,n):
       indep_var_values.append(gx0+i*dx)
       #make sigma large to smooth noise
       sigma  = 5.0
       gaussian_first_deriv.append(GaussianFirstDeriv(gx0+i*dx,0,sigma))
       gaussian_second_deriv.append(GaussianSecondDeriv(gx0+i*dx,0,sigma))

    #compute the convolutions
    convolved_first_deriv = convolution(dep_var_values,gaussian_first_deriv,dx)
    convolved_second_deriv = convolution(dep_var_values,gaussian_second_deriv,dx)

    #Reconstruct data points
    gaussian_first_deriv_pts=[]
    gaussian_second_deriv_pts=[]
    for i in range(0,n):
       gaussian_first_deriv_pts.append([indep_var_values[i],gaussian_first_deriv[i]])
       gaussian_second_deriv_pts.append([indep_var_values[i],gaussian_second_deriv[i]])

    convol_first_deriv_pts=[]
    convol_second_deriv_pts=[]
    for i in range(0,len(convolved_first_deriv)):
       convol_first_deriv_pts.append([indep_var_values[0]+data[0][0]+i*dx,convolved_first_deriv[i]])
       convol_second_deriv_pts.append([indep_var_values[0]+data[0][0]+i*dx,convolved_second_deriv[i]])

    #find the zeros of the first and second derivatives
    first_deriv_zeros= findzeros(convol_first_deriv_pts)
    second_deriv_zeros= findzeros(convol_second_deriv_pts)

    #find the region of the minimum of the second derivative
    convol_second_deriv_pts_sorted = sorted(convol_second_deriv_pts,key=operator.itemgetter(1))
    min_second_deriv = convol_second_deriv_pts_sorted[0:5]

    return(first_deriv_zeros,second_deriv_zeros,min_second_deriv)

#********************************************************************
# Function   : gridsearch                                           *
# Purpose    : Use a simple grid search to optimize the fitting.    *
# Parameters : data (list of data points)                           *
#              delta33,delta22,delta11,sigma (float)                *
# Returns    : list of data for fitted, simulated spectrum          *
#********************************************************************
def gridsearch(data,delta33,delta22,delta11,sigma):
    n = len(data) - (len(data) % 2)
    dx = abs(data[1][0]-data[0][0])
    small = [1.0,0]
    numpts = 20
    step = 2.0/float(numpts)
    for i in range(0,numpts+1):
        stdev = sigma-1+i*step

        #produce the ideal spectrum
        x0 = data[0][0]
        simulated_dep = []
        simulated_indep = []
        for i in range(0,n):
            simulated_indep.append(x0+dx*i)
            simulated_dep.append(intensity(x0+dx*i,delta33,delta22,delta11))

        #produce the gaussian
        gx0 = -n*dx/2.0 
        gaussian_indep = []
        gaussian_dep = []
        for i in range(0,n):
           gaussian_indep.append(gx0+i*dx)
           gaussian_dep.append(Gaussian(gx0+i*dx,0,stdev))

        #do the convolution
        broadened = convolution(simulated_dep,gaussian_dep,dx)

        #add dependent coordinates
        broadened_pts=[]
        gaussian_pts=[]
        for i in range(0,n):
            gaussian_pts.append([gaussian_indep[i],gaussian_dep[i]])
        for i in range(0,len(broadened)):
            broadened_pts.append([gx0+simulated_indep[0]+i*dx,broadened[i]])

        broadened_pts=normalize(broadened_pts)
        r = phi(data,broadened_pts)
        print "||data - broadened|| = "+str(r)
        print [r,stdev]
        
        #compare and update
        if r < small[0]:
            small = [r,stdev]
            print small
    print small
    return(broadened_pts)

#********************************************************************
# Function   : gatherinformation                                    *
# Purpose    : Collect strength and third deriv. values for use     *
#              in determining key zero-crossings of second deriv.   *
# Parameters : info (list of the results from edgedetect function   *
#              data (list of data points)                           *
# Returns    : list of data for fitted, simulated spectrum          *
#********************************************************************
def gatherinformation(info,data):
    second = info[1]

    #generate the third derivative
    dx = abs(data[1][0]-data[0][0])
    n = len(data)
    dep_var_values = []
    for i in range(0,n):
       dep_var_values.append(data[i][1])
    gx0 = -n*dx/2.0 
    gaussian_third_deriv = []
    indep_var_values = []
    for i in range(0,n):
       indep_var_values.append(gx0+i*dx)
       #make sigma large to smooth noise
       sigma  = 5.0
       gaussian_third_deriv.append(GaussianThirdDeriv(gx0+i*dx,0,sigma))
    convolved_third_deriv = convolution(dep_var_values,gaussian_third_deriv,dx)
    convol_third_deriv_pts=[]
    for i in range(0,len(convolved_third_deriv)):
       convol_third_deriv_pts.append([indep_var_values[0]+data[0][0]+i*dx,convolved_third_deriv[i]])

    third = findzeros(convol_third_deriv_pts)

    result = []
    
    for i in range(0,len(second)):
        for j in range(i,len(third)):
            if (second[i][0] > third[j][0]) and (second[i][0]<third[j+1][0]):
                 result.append([second[i][0],third[j+1][0]-third[j][0]])
                 
    #find the value of the third deriv. at each zero-crossing
    for i in range(0,len(result)):
        for j in range(0,len(convol_third_deriv_pts)):
            if (result[i][0] == convol_third_deriv_pts[j][0]):
                result[i].append(convol_third_deriv_pts[j][1])
                
    return(result)

####################################
#                                  #
#       main section of code       #
#                                  #
####################################

#
#Edge Detection routine
#

#read in data
filename = "spectrum.csv"
data = rdcsv(filename)
data.reverse()
#convert data from strings to floats
for i in range(0,len(data)):
   for j in range(0,2):
       data[i][j] = float(data[i][j])
#normalize the data set
data = normalize(data)

results = edgedetect(data)

print "The approximate zeros of the first derivative are:"
for i in range(0,len(results[0])):
    print results[0][i]
print
print "The region of the minimum of the second derivative is:"
for i in range(0,len(results[2])):
    print results[2][i]
print
print "Use the above values to approximate delta22."
print
print "Approximate zeros of the second derivative:"
print
information = gatherinformation(results,data)
print "Edge candidate\t edge strength\t sign of third derivative"
for i in range(0,len(information)):
    print str(information[i][0])+"\t "+str(information[i][1])+"\t "+str(information[i][2])
print

print "Now, use the above informatio to determine values for delta33 < delta22 < delta11."
delta33 = float(raw_input("Enter delta33: "))
delta22 = float(raw_input("Enter delta22: "))
delta11 = float(raw_input("Enter delta11: "))
sigma = float(raw_input("Enter a guess for the standard deviation of the Gaussian for broadening:"))
print

fit_simulation = gridsearch(data,delta33,delta22,delta11,sigma)
fit_simulation.reverse()
data.reverse()
outcsv("fit.csv",fit_simulation)
outcsv("newdata.csv",data)