from clustering_algorithm.cluster_flashes import load_dat, to_ecef
import numpy as np

center_geo = np.array([40.4463980, -104.6368130, 1000.00]) # COLMA center
lma_center = to_ecef(np.expand_dims(center_geo, axis=0))

n_grid_points, grid_spacing = 18, 0.2559
grid_start_lat = center_geo[0] - (n_grid_points / 2) * grid_spacing
grid_start_lon = center_geo[1] - (n_grid_points / 2) * grid_spacing

data, start_time = load_dat(os.path.join(data_dir, fname), min_stations=min_stations, max_chi_squared=max_chi_squared, max_altitude=max_altitude)

data_grid_lat = np.floor( (data[:,1] - grid_start_lat) / grid_spacing )
data_grid_lon = np.floor( (data[:,2] - grid_start_lon) / grid_spacing )
spatial_data = to_ecef(data[:,1:4]) - center

sources = np.zeros((data.shape[0], 10))
sources[:,0] = data[:,0]
sources[:,1:4] = spatial_data
sources[:,4] = data[:,5]
sources[:,5] = data_grid_lat
sources[:,6] = data_grid_lon
sources[:,7:] = data[:,1:4] # retain initial geodetic coordinates