#! /bin/env python
# coding: utf8

import numpy as np
import mms_tools_lib as mms
import os
import datetime
import matplotlib.pyplot as plt
import scipy.interpolate as scpinterp
import scipy.optimize as scpopt

class calc_harvey(object) :
	"""
	harvey algorithm routine setup with MMS initialization
	"""
	def __init__(self, orbit_files = None, data_files = None, orbite_vars = None, difftime_vars = None) :
		
		self.info = {}
		#self._init_default_info()
		
		if orbit_files.any : self.info['orbit_files'] = orbit_files
		if data_files.any : self.info['data_files'] = data_files
		if orbite_vars.any : self.info['orbite_vars'] = orbite_vars
		if difftime_vars.any : self.info['data_cross'] = difftime_vars
	
		self.satpostot = np.array([mms.donnee(fichier).get_data_as_field(varz = self.info['orbite_vars'], withtime = True) for fichier in self.info['orbit_files']])
		
	def evaluate_harvey(self, ij = [0,0], td = None, tf = None) :
	  
		if not td : td = mms.donnee(self.info['data_files'][0]).donnee['epoch'][0]
		if not tf : tf = mms.donnee(self.info['data_files'][0]).donnee['epoch'][-1]
		#ij = [0,0]
		
		def g2(ij) :
			t1 = td + datetime.timedelta(seconds = ij[0])
			t2 = tf - datetime.timedelta(seconds = ij[1])
		
			Btot = np.array([mms.donnee(Bsat).get_data_as_field(varz = self.info['data_cross'], t1 = t1, t2 = t2, withtime = True) for Bsat in self.info['data_files']])
			res = harvey(pos_vel = self.satpostot, data_timing = Btot)
			return res
		
		self.harvey = g2(ij)
		self.results = self.harvey.results

	def optimize_crossing(self) :
		"""
		This function helps finding the best differential timings in order to estimate Harvey method
		NOTE : THIS IS TOTALLY NOT OPTIMIZED!! Many useless calls are made. and probably useless.... :(
		"""
		td = mms.donnee(self.info['data_files'][0]).donnee['epoch'][0]
		tf = mms.donnee(self.info['data_files'][0]).donnee['epoch'][-1]
		maxvar = (tf - td).total_seconds()
	
		def g(ij) : #input time in seconds to change the correlation interval, return max correlation coefficient
			t1 = td + datetime.timedelta(seconds = ij[0])
			t2 = tf - datetime.timedelta(seconds = ij[1])
		
			Btot = np.array([mms.donnee(Bsat).get_data_as_field(varz = self.info['data_cross'], t1 = t1, t2 = t2, withtime = True) for Bsat in self.info['data_files']])
			res = harvey(data_timing = Btot, calc_harvey = False)
			return 1. - np.mean([max(corr) for corr in res.results['correlation_functions']])
	
	
		ij = np.array([0,0])
		#methodes = ['Powell' , 'L-BFGS-B', 'TNC', 'COBYLA', 'SLSQP', 'trust-ncg']
		#for meth in methodes :
			#print 'calculating with ', meth
			#res = scpopt.minimize(g, ij, bounds = [(0., maxvar/2.), (0, maxvar/2.)] , method = meth)
			#print res
		self.corr = np.zeros([20,20])
		for i1,i in enumerate(np.linspace(0,10,20)) :
			for j1,j in enumerate(np.linspace(0,10,20)) :
				print i1,j1
				ij = [i,j]
				self.corr[i1,j1] = g(ij)
		
		#res = scpopt.basinhopping(g, ij, niter = 20, stepsize = 0.5)
		print self.corr
	
	def _init_default_info(self):
	
		self.info['orbite_vars'] = ['mec_rx_gse', 'mec_ry_gse', 'mec_rz_gse', 'mec_vx_gse', 'mec_vy_gse', 'mec_vz_gse']
		self.info['data_cross'] = ['fgm_by_gse_brst_l2']
		
		rep = 'time_table_20150908_100807_20150908_100844'
		baserep = os.curdir
		rep = os.path.join(baserep, rep)
		mecsat1 = os.path.join(rep, 'obj7.cef')
		mecsat2 = os.path.join(rep, 'obj8.cef')
		mecsat3 = os.path.join(rep, 'obj9.cef')
		mecsat4 = os.path.join(rep, 'obj10.cef')
		self.info['orbit_files'] = [mecsat1, mecsat2, mecsat3, mecsat4]

		Bsat1 = os.path.join(rep, 'obj3.cef')
		Bsat2 = os.path.join(rep, 'obj4.cef')
		Bsat3 = os.path.join(rep, 'obj5.cef')
		Bsat4 = os.path.join(rep, 'obj6.cef')
		self.info['data_files'] = [Bsat1, Bsat2, Bsat3, Bsat4]


	

class harvey(object) :
	"""
	Differential timing calculation
	Input : Time + x,y,z satelite positions +  vx,vy,vz satelite velocities (total = 7 columns) * 4 Satellites * N data
	This must be provided as a numpy array such as : pos_vel[0,:,:] provides data for satellite 1
	      pos_vel[2,:,3] provides all z positions of satellite 3.
	      Structure: pos_vel = np.zeros([n_satellites ( = 4), data size, number of expected variables ( = 7)])
	data_timing: set of data used to correlate the timing, for automatic differential timings deduction
	      Structure: data_timing[1,:,1] provide all data for satellite 2
			  np.zeros([n_satellite (= 4), size of data, 2 (for time and data value)])
	NOTE: RESULTS ARE PROVIDED IN SELF.RESULTS
	"""
	def __init__(self, pos_vel = np.zeros([4,10,7]), data_timing = np.zeros([4,10,2]), calc_harvey = True) :
		
		self.results = {}
		
		self._time_reference = data_timing[0][0,0]
		self.data_timing = data_timing
		self._find_crossing()
		if calc_harvey : 
			self.pos_vel = pos_vel
			self._calc()
	
		
	def _calc(self) :
		"""
		Methodology :
		1. for all spacecraft : interpolate the position of each spacecraft at t = timing
		2. find mesocentre position
		3. derive distance each SC from mesocentre
		4. Build covariance matrix of delta_to_mesocentre
		5. 
		"""
		j = 0 #### USE ONLY FIRST LINE OF DATA : COULD BE IMPROVED!
		new_pos = np.zeros([4,3])
		for i, sc_data in enumerate(self.pos_vel) : 
			delta_t = (self.timing[i] - sc_data[:,0])[j].total_seconds()
			new_pos[i,:] = sc_data[j,1:4] + sc_data[j,4:7] * delta_t
		
		pos_vel_meso = new_pos.mean(axis = 0)
		#print 'pos mesocentre : ', '\n', pos_vel_meso
		 
		#### à vérifier : calcul de delta_sc_pos : utiliser new pos ou old pos?
		delta_sc_pos = new_pos[:,:] - pos_vel_meso[:]
		#print 'delta sc to meso : ', '\n', delta_sc_pos
		
		matrice = np.mat(np.zeros([3,3]))
		for i in range(3) :
			for j in range(3) :
				matrice[i,j] = (delta_sc_pos[:,i] * delta_sc_pos[:,j]).mean()
		invr = np.linalg.inv(matrice)
		
		vect_base = np.zeros(3)
		for i in range(3) :
			for j in range(3) :
				vect_base[i] = vect_base[i] + (self.timing_seconds[:] * delta_sc_pos[:,j]).mean() * invr[i,j]
		
		self.results['normal_velocity'] = 1./np.linalg.norm(vect_base)
		self.results['normal_vector'] = vect_base * self.results['normal_velocity']

		
	def _find_crossing(self, bestfit = False) :
		"""
		1. interpolation des données de B2,3 et 4 sur abscisse de 1
		2. Input data to self.getDiffTiming()
		"""
		data = np.zeros([4, len(self.data_timing[0][:,0])])
		data[0,:] = self.data_timing[0][:,1]
		for i, sc_data in enumerate(self.data_timing[1:]):
			data[i+1,:] = mms.time_interpolation(self.data_timing[0][:,0], sc_data[:,0], sc_data[:,1])
		
		if bestfit : 
			ind = self.search_best_fit()
			self.getDiffTiming(time_axis = self.data_timing[0][ind:-ind,0], data = data[:,ind:-ind])
		else : self.getDiffTiming(time_axis = self.data_timing[0][:,0], data = data)
		
	
	def getDiffTiming(self, time_axis = np.zeros(100), data = np.zeros([4,100])):
		"""
		Input time_axis = refence time abscisse
		Input data = all interpolated data for all spacecraft
		"""
		self.timing = []
		self.results['differential_timings'] = []
		self.results['correlation_functions'] = []
		ord1 = self._normalize(data[0,:])/len(data[0,:])
		for sc in data:
			res = np.correlate(ord1, sc, mode = 'same')
			self.timing.append(time_axis[res.argmax()])
			self.results['correlation_functions'].append(res)
			
		self.results['differential_timings'] = np.array([(time - self.timing[0]).total_seconds() for time in self.timing])
		
		self.timing_seconds = np.array([(time - self._time_reference).total_seconds() for time in self.timing]) #FOR SC position interpolation!!

	def _normalize(self, data) :
		return (data - data.mean())/data.std()
	
	def search_best_fit(self, n = 10):
		"""
		Search for best correlation coefficient
		Method: cut data into n time blocks
		1st step: calculate correlation coefficient
		2nd step: remove extreme left and right blocks -> recorrelate
		3rd: repeat
		"""
		block  = int(len(self.data_timing[0][:,0]) / (2.*n))
		maxcorr = np.zeros([n,4])
		# Set all data to same time axis (the one of Sat1):	
		data = np.zeros([4, len(self.data_timing[0][:,0])])
		data[0,:] = self.data_timing[0][:,1]
		for i, sc_data in enumerate(self.data_timing[1:]):
			data[i+1,:] = mms.time_interpolation(self.data_timing[0][:,0], sc_data[:,0], sc_data[:,1])
		
		maxcorr[0,:] = [np.correlate(self._normalize(data[0,:])/len(data[0,:]), self._normalize(data[i,:])).max() for i in [0,1,2,3]]
		for j in range(1,n):
			d = block*j
			maxcorr[j,:] = [np.correlate(self._normalize(data[0, d:-d])/len(data[0,d:-d]), self._normalize(data[i, d:-d]), 'same').max() for i in [0,1,2,3]]
		
		best_indice = np.array([np.linalg.norm(i) for i in maxcorr]).argmax()
		return block*best_indice
	
	def plot(self) :
		
		plt.figure()
		barre = [self.data_timing[0][:,1].min(), self.data_timing[0][:,1].max()]
		
		for i, champ in enumerate(self.data_timing) : 
			plt.plot(champ[:,0], champ[:,1], label = str(i))
			plt.plot([self.timing[i], self.timing[i]], barre)
		plt.legend()
		plt.show()
	
	def plot_hodo(self) :
		
		champ2 = mms.time_interpolation(self.data_timing[0][:,0], self.data_timing[1][:,0], self.data_timing[1][:,1])
		plt.figure()
		plt.plot(self.data_timing[0][:,1], champ2)
		plt.show()
	
	#def _find_crossing_old(self, val = 0) :

		#self.timing = np.array([sc_data[(sc_data[:,1] - val).argmin(), 0] for sc_data in self.data_timing])
		#self.timing_seconds = np.array([(time - self.time_reference).total_seconds() for time in self.timing])