import numpy as np
import matplotlib
from matplotlib.ticker import FuncFormatter
import matplotlib.pyplot as plt
import scipy.stats as stats

def resample_signal(time, signal, dtime=None):
	"""
	*Resample the initial signal at a constant time interval.*
  
	:param time: Time vector of your signal
	:param signal: Signal vector
	:dtime: New time step
	:returns:            
	- **rescaled_time** - Uniformally rescaled time vector
	- **rescaled_signal** - Rescaled signal

	.. note:: 	If a dtime is given, the interpolation is made to 
				have a signal with a time interval of dtime.
				Else, the dt is the smallest time interval between
				two values of the signal.
	"""

	if dtime is None :
		time_steps = np.diff(time)
		dtime = time_steps.min()
	
	rescaled_time = time - time[0]
	rescaled_time = np.arange(rescaled_time[0],rescaled_time[-1],dtime)
	

	rescaled_signal = np.interp(rescaled_time, time - time[0], signal)
	
	return rescaled_time, rescaled_signal

def calc_autocorrelation_time(time, signal, threshold=0.2):
	"""
	*Estimate the autocorrelation time at a given threshold.*
	
	:param time: Time vector of your signal
	:param signal: Signal vector
	:param threshold: Threshold under which the signal is correlated
	:type threshold: float

	:returns:

	- **autocorrelation_time** - Minimum time step to capture the signal at a
								 correltion under thethreshold
	"""

	signal -= np.mean(signal)
	correlation = np.correlate(signal, signal, 'full')
	correlation = abs(correlation[int((correlation.size - 1) / 2):])
	correlation /= correlation[0]

	autocorrelation_time = time[np.argmax(correlation < threshold)]

	return autocorrelation_time


def show_autocorrelation_time(time, signal, threshold=0.2):
	"""
	*Plot the autocorrelation function of the signal.*
	
	:param time: Time vector of your signal
	:param signal: Signal vector
	:param threshold: Autocorrelation threshold

	:returns:

	- **fig** - Figure of the result

	"""

	signal -= np.mean(signal)
	correlation = np.correlate(signal, signal, 'full')
	correlation = abs(correlation[int((correlation.size - 1) / 2):])
	correlation /= correlation[0]

	autocorrelation_time = time[np.argmax(correlation < threshold)]

	figout = plt.figure()
	plt.plot(time, correlation)
	plt.axvline(autocorrelation_time, c='r')
	plt.xlabel('Time (s)')
	plt.ylabel('Correlation')
	plt.semilogx()
	plt.grid()

	return figout

def sort_spectral_power(time, signal):
	"""
	*Determine the harmonic power contribution of the signal.*
	
	:param time: Time vector of your signal
	:param signal: Signal vector

	returns:

	- **harmonic_power** - Harmonic power of the signal
	- **total_power** - Total spectral power of the signal


	.. note:: 	It calculates the Power Spectral Density (PSD) of the complete
				signal and of a downsampled version of the signal. The 
				difference of the two PSD contains only harmonic components.
	"""

	frequency1, psd1 = power_spectral_density(time, signal)
	deltaf1 = 1. / ((time[1]-time[0]) * 2)

	downsampled_time = time[0:-1:2]
	downsampled_signal = signal[0:-1:2]
	frequency2, psd2 = power_spectral_density(downsampled_time, downsampled_signal)
	deltaf2 = 1. / ((downsampled_time[1] - downsampled_time[0]) * 2)

	psd2 = np.interp(frequency1,frequency2,psd2) * 2

	harmonic_power = np.cumsum(np.abs(psd1 - psd2) * 2)[-1] * deltaf1
	total_power = np.cumsum(psd1)[-1] * deltaf1

	return harmonic_power, total_power

def power_representative_frequency(time, signal, threshold=0.8):
	"""
	*Calculate the frequency that captures a level of spectral power.*

	:param time: Time vector of your signal
	:param signal: Signal vector
	:param threshold: Level of representativity of the spectral power

	:returns:

	- **representative_frequency** - Frequency above which the power threshold is reached

	.. note:: 	It calculates the cumulative power spectral density and returns
				the frequency that reaches the threshold of spectral power.
	"""
	
	frequency, power_density = power_spectral_density(time, signal)

	cumulative_psd = np.cumsum(power_density)
	cumulative_psd /= cumulative_psd[-1]
	representative_frequency = frequency[np.argmax(cumulative_psd>threshold)]

	return representative_frequency

def show_power_representative_frequency(time, signal, threshold=0.80):
	"""
	*Plot the power spectral density of the signal.*
	
	:param time: Time vector of your signal
	:param signal: Signal vector
	:param threshold: Power representative frequency threshold

	:returns:

	- **fig** - Figure of the result

	"""

	frequency, power_density = power_spectral_density(time, signal)
	cumulative_psd = np.cumsum(power_density)
	cumulative_psd /= cumulative_psd[-1]
	representative_frequency = frequency[np.argmax(cumulative_psd>threshold)]

	figout = plt.figure()
	plt.plot(frequency, cumulative_psd)
	plt.axvline(representative_frequency, c='r')
	plt.xlabel('frequency (Hz)')
	plt.ylabel('PSD (%)')
	plt.grid()
	
	return figout



def ks_test_distrib(data, distribution='normal'):
	"""
	*Calculate the correlation score of the signal with the distribution*

	:param data: array of values
	:param distribution: kind of distribution the values follow to test

	:returns:

	- **score** -Minimum score over the height of the ks test
	- **position** -Index of the height at which the min. of the test is found
	- **height** -Corresponding heiht where the min. is found
	- **scale** -Scale parameter of the lognormal fitting

	"""
	data_shape = np.shape(data)
	if len(data_shape) > 1:
		p_value = np.zeros(data_shape[1])
		for j in range(data_shape[1]):
			data_1d = data[:,j]

			mu = np.average(data_1d)
			sigma = np.std(data_1d)
			
			# if data_1d.min() <= 0:
			# 	p_value[j]=1
			# else:
			if distribution == "lognormal":
				sigma, loc, scale = stats.lognorm.fit(data_1d, floc=0)
				stat, p_value[j] = stats.kstest(data_1d, 'lognorm', args=(sigma, loc, scale))
				# if p_value[j]<5.0e-2:
				# 	print(p_value[j])
				# 	print(temperature.min())
				# 	print(temperature.max())

			elif distribution == "normal":
				loc, sigma = stats.norm.fit(data_1d, loc=300)
				stat, p_value[j] = stats.kstest(data_1d, 'norm', args=(loc, sigma))
				scale = None
		position = np.argmin(p_value)
		height = position/data_shape[1]
	else:
		p_value = 0

		data_1d = data
		mu = np.average(data_1d)
		sigma = np.std(data_1d)
		if distribution == "lognormal":
			sigma, loc, scale = stats.lognorm.fit(data_1d, floc=0)
			stat, p_value = stats.kstest(data_1d, 'lognorm', args=(sigma, loc, scale))
			# if p_value[j]<5.0e-2:
			# 	print(p_value[j])
			# 	print(temperature.min())
			# 	print(temperature.max())

		elif distribution == "normal":
			loc, sigma = stats.norm.fit(data_1d, loc=300)
			stat, p_value = stats.kstest(data_1d, 'norm', args=(loc, sigma))
			scale=None
		position = np.argmin(p_value)
		height = None
	
	#print(p_value)	
	score = np.min(p_value)
	
	


	

	return score, position, height, scale

def show_temperature_distribution(temperature_recording, height,distribution='normal'):
	"""
	*Plot the temperature distribution and the fitting curve*

	:param temperature_recording: Temperature as a function of height and time
	:param height: Height in the plan40
	:param distribution: Type of distribution for the fitting method

	:returns:

	- **fig** - Figure of the result
	"""
	hH = np.linspace(0,1,len(temperature_recording[0,:]))
	index = np.abs(hH-height).argmin()
	data = temperature_recording[:,index]
	
	
	
	abscissa = np.linspace(0,np.amax(data),num=50)

	if distribution == 'normal':
		loc, sigma = stats.norm.fit(data, loc=300)
		pdf_fitted = stats.norm.pdf(abscissa, loc, sigma)

	if distribution == 'lognormal':
		sigma, loc, scale = stats.lognorm.fit(data, floc=300)
		pdf_fitted = stats.lognorm.pdf(abscissa, sigma, loc, scale)

	
	figout = plt.figure()
	plt.hist(data, density=True, bins=30, color='lightgray')
	plt.plot(abscissa, pdf_fitted)
	plt.xlabel('T (K)')
	
	return figout


def duration_for_uncertainty(time, signal, target=10, confidence=0.95, distribution='normal'):
	"""
	*Give suggestion of simulation duration of a plan40 calculation.*
	
	:param time: Time vector of your signal
	:param signal: Signal vector
	:param target: Desired amplitude of the confidence interval
	:param confidence: Level of confidence of the interval
	:param distribution: Type of distribution of the signal to make the interval

	:returns:

	- **duration** -Duration of the signal
	"""

	dtime = calc_autocorrelation_time(time, signal)
	sigma = np.std(signal)
	
	if distribution == "normal":
		z_coef = stats.norm.ppf(1 - (1 - confidence)/2)
		n_timesteps = int((2 * z_coef * sigma / target) ** 2)

	

	duration = dtime * n_timesteps

	return duration

def uncertainty_from_duration(dtime, sigma, duration, confidence=0.95, distribution='normal'):
	"""
	*Give confidence interval length of a plan40 calculation.*

	:param dtime: Time step of your solutions
	:param sigma: Standard deviation of your signal
	:param duration: Desired duration of the signal
	:param confidence: Level of confidence of the interval
	:param distribution: Type of distribution of the signal to make the interval

	:returns:

	- **length** - Length of the confidence interval in K
	"""
	
	if distribution == "normal":
		z_coef = stats.norm.ppf(1 - (1 - confidence)/2)
		#print(z_coef)
		n_timesteps = int(duration / dtime)
		length = 2 * z_coef * sigma / np.sqrt(n_timesteps)

	return length

	
def convergence_cartography(time, signal, **kwargs):
	"""
	*Create a cartography of the convergence of the confidence interval in a simulation.*

	:param time: Time vector of your signal
	:param signal: Signal vector


	==**kwargs==
	
		:param max_time: Maximal simulation duration
		:param interlen: Maximal interval length

	:returns:
		- **fig** - Figure of the cartography

	"""
	max_time = kwargs.get('max_time', None)
	interlen = kwargs.get('interlen', None)

	if max_time is None: max_time = 0.05
	if interlen is None: interlen = 50

	dtime = calc_autocorrelation_time(time, signal)
	sigma = np.std(signal)
	durations = np.arange(0.001, max_time, 0.001)
	interval_lengths = np.arange(0,interlen)
	level = np.zeros((len(durations), len(interval_lengths)))
	for idx, dur in enumerate(durations):
		for idy, length in enumerate(interval_lengths):
			sqrtn = int(np.sqrt(dur/dtime))
			z_star = length * sqrtn / (2 * sigma)
			dec_per = stats.norm.cdf(z_star)
			level[idx, idy] = 1 - (1 - dec_per) * 2

	fig = plt.figure()
	fig = plt.contour(durations,interval_lengths,np.transpose(level),[0.6,
																	  0.8,
																	  0.9, 
																	  0.95,
																	  0.975])
	#plt.legend()
	plt.xlabel('Simulation duration (s)')
	plt.ylabel('Interval amplitude (K)')
	plt.clabel(fig, inline=1, fontsize=10, colors='k', fmt={0.6:'60%',
															0.8:'80%',
															0.9:'90%',
															0.95:'95%',
															0.975:'97.5%'})
	plt.title('Contour of confidence interval(simulation duration, interval amplitude)')
	plt.grid(True)
	return fig




def calculate_std(time, signal, frequency):
	"""
	*Give the standard deviation of a signal at a given frequency.*
	
	:param time: Time vector of your signal
	:param signal: Signal vector
	:param frequency: Frequency at which values of the recording are taken

	:returns:

		- **std** - Standard deviation of the values taken from the recording
	"""

	dtime = 1 / frequency
	time -= time[0]
	rescaled_time = np.arange(0,time[-1],dtime)
	rescaled_signal = np.interp(rescaled_time, time, signal)

	std = np.std(rescaled_signal)
	
	return std

def power_spectral_density(time, signal):
	"""
	*Automate the computation of the Power Spectral Density of a signal.*

	:param time: Time vector of your signal
	:param signal: Signal vector

	:returns:

	- **frequency** -Frequency vector of the signal's power spectral density
	- **power_spectral_density** -Power spectral density of the signal
	"""

	total_time = time[-1] - time[0]
	dtime = time[1] - time[0]

	freq_resolution = 1. / total_time
	freq_nyquist = 1. / dtime

	N = len(time)
	
	frequency = np.arange(0, freq_nyquist, freq_resolution ,dtype=np.float)
	frequency = frequency[:int(N/2)-1]

	
	mean_amplitude = np.mean(signal)

	raw_fft = np.fft.fft(signal - mean_amplitude)
	power_spectral_density = np.square(np.absolute(raw_fft)[:int(N/2)-1]/ np.sqrt(N))
	power_spectral_density[1:] *= 2

	return frequency, power_spectral_density

def to_percent(y, position):
	"""
	*Rescale the y-axis to per*
    
    """
	s = str(round(100. * y))

	# The percent symbol needs escaping in latex
	if matplotlib.rcParams['text.usetex'] is True:
		return s + r'$\%$'
	else:
		return s + '%'

def plot_distributions(path="./data.txt"):
	formatter = FuncFormatter(to_percent)
	data = np.genfromtxt(path)
	S = np.shape(data)

	t_fusion = 1930.
	bin_edges = list(np.linspace(t_fusion - 200., t_fusion + 200., 21))

	for j in range(S[1]):

		temperatures = data[:,j]

		clipped_temperatures=np.clip(temperatures, bin_edges[0], bin_edges[-1])
		clipped_temperatures.sort()
		weights = np.ones_like(clipped_temperatures)/float(len(clipped_temperatures))
		lower=0
		upper=-1
		lower=np.argmax(clipped_temperatures>bin_edges[0])-1
		upper=np.argmax(clipped_temperatures==bin_edges[-1])

		if upper==0:
			upper=-1

		fig, (ax1,ax2,ax3) = plt.subplots(1,3,sharey=True, gridspec_kw={"width_ratios":[1, 5, 1]})

		rect = plt.Rectangle((0.4,0),0.2,np.sum(weights[:lower+1]))
		ax1.add_patch(rect)
		plt.sca(ax1)
		plt.xticks([0.5], ['<'+str(t_fusion-200)])
		plt.gca().yaxis.set_major_formatter(formatter)
		ax1.set_ylabel('Probability');

		ax2.hist(clipped_temperatures[lower:upper], bins=bin_edges, weights=weights[lower:upper])
		xticks = list(np.linspace(t_fusion - 200., t_fusion + 200., 11))
		xlabels = list(map(str,xticks))
		plt.sca(ax2)
		plt.xticks(xticks, xlabels, rotation=45,size=8)
		plt.gca().yaxis.set_major_formatter(formatter)


		rect = matplotlib.patches.Rectangle((0.4,0),0.2,np.sum(weights[upper:]))
		ax3.add_patch(rect)
		plt.sca(ax3)
		plt.xticks([0.5], [str(t_fusion+200)+'<'])
		plt.gca().yaxis.set_major_formatter(formatter)


		plt.savefig('pdf_h_%s.png' % j)
		plt.close()




