from dataclasses import dataclass
from random import randrange
from scipy.optimize import basinhopping
import math
import numpy as np

#from ROOT import TFile, TNtupleD, gROOT

import matplotlib.pyplot as plt

M = np.diag([1, 1, -1]) 

@dataclass
class Vertexer:

	nodes: np.ndarray
	c: float

	def __post_init__(self):
		# Calculate valid input range
		self.min_dt = np.amin(np.linalg.norm(self.nodes - 
			np.average(self.nodes, axis=0), axis=1)) / self.c

		self.max_dt = 0
		for n in self.nodes:
			for p in self.nodes:
				dist = np.linalg.norm(n - p)

				if dist > self.max_dt:
					self.max_dt = dist 
					
		self.max_dt /= self.c

	def reconstruct(self, times):
		# Normalize times
		times -= times.min()

		#print(times.min())

		#print(times[0] - times[1], " " , times[0] - times[2], " " , times[1] - times[2])

		if(not (times < 0.00001491746).all()):			
			print("Yay")
		#	raise Exception()

		# max = 0.00000745873 # min = 

		A = np.append(self.nodes, np.reshape(times, (-1, 1)) * self.c, axis=1)

		def ssr_error(point):
			# Return SSR error
		#	print('Error: ', np.sum(((np.linalg.norm(self.nodes - point, axis=1) / self.c)
		#                                   - times)**2))
			return np.sum(((np.linalg.norm(self.nodes - point, axis=1) / self.c)  
				   - times)**2)

		def lorentz(a, b):
			# Return Lorentzian Inner-Product
			return np.sum(a * (b @ M), axis=-1)

		b = lorentz(A, A) * 0.5
		C = np.linalg.solve(A, np.ones(3))
		D = np.linalg.solve(A, b)

		roots = np.roots([lorentz(C, C),
						 (lorentz(C, D) - 1) * 2, 
						  lorentz(D, D)])

		solutions = []
		for root in roots:
			X, Y, T = M @ np.linalg.solve(A, root + b)
			solutions.append(np.array([X,Y]))
			#print(X,Y)
		
		solution = max(solutions, key = ssr_error)

		print("ORIGINAL ERR: ", ssr_error(solution))

		result = basinhopping(ssr_error, solution, niter=10000, minimizer_kwargs={'method':'L-BFGS-B'})
		print("NEW RESULT: ", result)
		print("NEW ERR: ", ssr_error(result.x))

		#print(ssr_error(solution))
		return(solution, ssr_error(solution))	
	
# Simulate sources to test code
x_1 = randrange(10); y_1 = randrange(10); #//z_1 = randrange(10)
x_2 = randrange(10); y_2 = randrange(10); #//z_2 = randrange(10)
x_3 = randrange(10); y_3 = randrange(10); #//z_3 = randrange(10)
#x_4 = randrange(10); y_4 = randrange(10); //z_4 = randrange(10)

# Pick source to be at random location
x = randrange(10); y = randrange(10); #//z = randrange(1000)

# Set velocity
c = 299792 # km/ns

# Generate simulated source
t_1 = math.sqrt( (x - x_1)**2 + (y - y_1)**2 ) / c
t_2 = math.sqrt( (x - x_2)**2 + (y - y_2)**2 ) / c
t_3 = math.sqrt( (x - x_3)**2 + (y - y_3)**2 ) / c
#t_4 = math.sqrt( (x - x_4)**2 + (y - y_4)**2 ) / c

print('Actual:', x, y)
myVertexer = Vertexer(np.array([[x_1, y_1],[x_2, y_2],[x_3, y_3]]), c)
print(myVertexer.reconstruct(np.array([t_1, t_2, t_3])))