#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Dec  2 14:04:52 2020

@author: madugadzama
"""

'''WEEK 9'''

''' Demonstration of how correlations between parameters were investigated for a given experiment. This involved creation of a Hesisan matrix.'''



#Other imports
import numpy as np
import time

import matplotlib.pyplot as plt
from timeit import default_timer as timer 


#starting a timer
t1_start = time.process_time()


#define the Lorentzian Linewidth
def linewidth(x, s_bgd, s_0, x_0, d_x):
    return s_bgd + s_0/(1+4*((x-x_0)/d_x)**2)

# setting true paramter values
background = 2
s_0 = 8 # height of peak
x_0 = 50 # centre of peak
d_x = 10 # angular width of peak

#setting the number of repeats for the obtaining of our 'experimental data' 
number_of_runs = 10

#setting the variance of noise
noise_variance= 4.2

# setting the x values and corresponding linewidth values
x = np.linspace(0,100,200)
y = linewidth(x,background,s_0,x_0,d_x) # this is the linewidth before noisy is added


'''I am just going to use a variance in the noise of 2.2'''


y_1 = np.loadtxt('exp_1.txt')
y_2 =np.loadtxt('exp_2.txt')
y_3 =np.loadtxt('exp_3.txt')
y_4 =np.loadtxt('exp_4.txt') 
y_5 = np.loadtxt('exp_5.txt')
y_6 =np.loadtxt('exp_6.txt')
y_7 =np.loadtxt('exp_7.txt')
y_8 =np.loadtxt('exp_8.txt') 
y_9 = np.loadtxt('exp_9.txt')
y_10 =np.loadtxt('exp_10.txt')


# Finding the mean and standard error of 

mean = [] # this is the data we are plotting originally as our 'noisy data' 
standard_error = []

# using the zip function method to find mean and std
a = list(zip(y_1,y_2,y_3,y_4,y_5,y_6,y_7,y_8,y_9,y_10))

for i in range(len(a)):
    mean.append((np.sum(a[i])/number_of_runs))
    standard_error.append(np.std(a[i])/np.sqrt(number_of_runs)) 

#plottinf the noisy lorentzian with error bars 
plt.errorbar(x, mean, yerr=standard_error, marker = '.', linestyle = 'None')
plt.xlabel('Scattering Angle')
plt.ylabel('Diffracted Intensity')
plt.show()     

#making an array of standard error values squared
standard_error_squared = [] 
for i in range(len(standard_error)):
    standard_error_squared.append(standard_error[i]**2)


#turn lists into numpy arrays
standard_error = np.array(standard_error)

#defining our cost function to minimise
def chi_squared(x, s_bgd, s_0, x_0, d_x):
    return np.sum((mean -linewidth(x, s_bgd, s_0, x_0, d_x))**2/ standard_error**2)

def least_squares(x, s_bgd, s_0, x_0, d_x):
    return np.sum((mean -linewidth(x, s_bgd, s_0, x_0, d_x)**2))




'''Calculations part'''

height_array = np.arange(3, 12.00000001,1) # peak height 
background_array = np.arange(-2,7.00000001,1)#background


# need to import the parameter values and chi-squared values to find the values of the other
#... parameters (which have been optimised) when at min chi-squared. 



# note in these text files, the final number shows us which variance was used. If it ends with 1, var=2.2. If it ends with 2, var=6.2. If it ends with 3, var=6.2
chi_array = np.loadtxt('week_9_collection_chi_squared_background_peak_height3.txt')
centre_array = np.loadtxt('week_9_collection_chi_squared_background_peak_height_1_centre_value_3.txt')
width_array = np.loadtxt('week_9_collection_chi_squared_background_peak_height_1_peak_width_value_3.txt')


#now making code to find the values of the other parameters at the min chi_squared
# we want to find the indices of the minimum value
chi_array = np.reshape(chi_array,(len(height_array),len(background_array))) #here we reshape the array so it matches the one used in the contour plotter
centre_array =  np.reshape(centre_array ,(len(height_array),len(background_array)))
width_array =  np.reshape(width_array,(len(height_array),len(background_array)))

#^^ we know that the rows correspond to the height array and the columns to the background to the 
ind = np.unravel_index(np.argmin(chi_array, axis=None), chi_array.shape) # this gives us the index of the minimum value 

centre_array_at_min_chi_squared = centre_array[ind] # this gives us the value of x0 when at min chi-squared
width_array_at_min_chi_squared = width_array[ind] # this gives us the value of dx when at min chi-squared




#print(centre_array_at_min_chi_squared,width_array_at_min_chi_squared)


'''Hessian Matrix creation'''

H_1 = 2* (np.sum(1/standard_error**2))
H_2 = 2*np.sum((1/standard_error**2)*(1/(1+4*((x-centre_array_at_min_chi_squared)/width_array_at_min_chi_squared)**2)))
H_3 =  2*np.sum((1/standard_error**2)*(1/(1+4*((x-centre_array_at_min_chi_squared)/width_array_at_min_chi_squared)**2)))
H_4 =  2*np.sum((1/standard_error**2)*(1/((1+4*((x-centre_array_at_min_chi_squared)/width_array_at_min_chi_squared)**2)**2)))

hessian_matrix = np.empty((2,2))
hessian_matrix[0,0] = H_1
hessian_matrix[0,1] = H_2
hessian_matrix[1,0]= H_3
hessian_matrix[1,1] = H_4

print('Hessain Matrix is:\n',hessian_matrix)

'''curvature matrix and covariance matrix'''

curvature_matrix = 0.5*hessian_matrix
covariance_matrix = np.linalg.inv(curvature_matrix)

print('Covariance Matrix is:\n',covariance_matrix)

'''Now to make a correlation matrix for variance 2.2'''
correlation_matrix = np.empty((2,2))
correlation_matrix[0,0] = correlation_matrix[1,1] = 1 # by definition this uses diagonals of 1
correlation_matrix[0,1] = covariance_matrix[0,1]/np.sqrt(covariance_matrix[0,0]*covariance_matrix[1,1])
correlation_matrix[1,0] = covariance_matrix[1,0]/ np.sqrt(covariance_matrix[0,0]*covariance_matrix[1,1])
print('Correlation matrix is:\n',correlation_matrix)
