#!/usr/bin/python3

from sklearn.cluster import KMeans
import numpy as np
import math,sys


generator_func_name = "%s_generator" % sys.argv[2]
norm_func_name = "%s_generator" % sys.argv[2]

textdata = open(sys.argv[1]).readlines()
normal_groups = 4
normals_from_data = False
#textdata = open("lc.bai").readlines()
#normals_from_data = True
#normal_groups = 16
points = []
point_normals = []
lines = []
line_normals = []

# Read in the basic format data that our tooling generated
for j in textdata:
  st = j.strip()
  st = st.replace('_','')
  st = st.replace(' ','')
  st = st.split('},')
  for entry in st:
    if len(entry) > 0 and entry[0]=='{':
      entry = entry.replace('{','')
      entry = entry.replace('}','')
      nv = entry.split(',')
      for q in range(len(nv)-1):
        nv[q+1] = float(nv[q+1])
      points.append(nv)

# Generate a normal vector based on the centrepoint for each point
# This is actually a pretty dumb way to do this.... but for the purposes of
# our trees, will actually suffice - our tree normals are "always" away from the horizontal
# center - so we're just going to generate a point normal of [0, 0, z]
if normals_from_data:
  pass
else:
  for j in points:
    point_normals.append([j[1], 0, j[3]])

# Okay, now we generate lines
for j in range(1, len(points)):
  if points[j][0] != 'MoveTo':
    l = [points[j-1], points[j]]
    ln = []
    if normals_from_data:
      ln = points[j][4:]
    else:
      for p in range(3):
        ln.append((point_normals[j-1][p]+point_normals[j][p])/2)
    lnlen = math.sqrt(ln[0]*ln[0]+ln[1]*ln[1]+ln[2]*ln[2])
    if lnlen > 0:
      for p in range(3):
        ln[p] /= lnlen
    lines.append(l)
    line_normals.append(ln)

# Now, we use kmeans clustering to get a couple of "normal sets" - this is the performance
# tradeoff we take.... the idea being that we don't want to "generate" anything in
# gsbasic, but we DO want to cull for performance!
print("'Starting kmenas....",len(line_normals))
kmeans = KMeans(n_clusters = normal_groups).fit(line_normals)
print("'Kmeans done")

print("function %s()" % generator_func_name)
print("dim litobj[%d]" % normal_groups)
for n in range(normal_groups):
  ll=[-1,-1,-1]
  print("litobj[%d] = { _" % (n+1))
  fc=""
  for seg in range(len(lines)):
    if kmeans.labels_[seg] == n:
      if ll != [lines[seg][0][1], lines[seg][0][2], lines[seg][0][3]]:
        print("%s{MoveTo, %f,%f,%f} _" % (fc, lines[seg][0][1], lines[seg][0][2], lines[seg][0][3]))
      print(",{DrawTo, %f,%f,%f} _" % (lines[seg][1][1], lines[seg][1][2], lines[seg][1][3]))
      ll = [lines[seg][1][1], lines[seg][1][2], lines[seg][1][3]]
      fc=","
  print("}")
print("return litobj")
print("endfunction")
print("")

print("function %s()" % norm_func_name)
print("dim norms[%d]" % normal_groups)
for n in range(normal_groups):
  print("norms[%d] = { %f, %f, %f }" % (n+1, kmeans.cluster_centers_[n][0],  kmeans.cluster_centers_[n][1], kmeans.cluster_centers_[n][2]))
print("return norms")
print("endfunction")
