from dataclasses import dataclass
from random import randrange
import math
import numpy as np

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

@dataclass
class Vertexer:

    nodes: np.ndarray
    c: float

    def reconstruct(self, times):

        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(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(1000); y = randrange(1000); #//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])))