import struct
import numpy as np
import collections

def _write_ply_point(fp, x,y,z, color=None, normal=None, binary=False):
  args = [x,y,z]
  if color is not None:
    args += [int(color[0]), int(color[1]), int(color[2])]
  if normal is not None:
    args += [normal[0],normal[1],normal[2]]
  if binary:
    fmt = '<fff'
    if color is not None:
      fmt = fmt + 'BBB'
    if normal is not None:
      fmt = fmt + 'fff'
    fp.write(struct.pack(fmt, *args))
  else:
    fmt = '%f %f %f'
    if color is not None:
      fmt = fmt + ' %d %d %d'
    if normal is not None:
      fmt = fmt + ' %f %f %f'
    fmt += '\n'
    fp.write(fmt % tuple(args))

def _write_ply_triangle(fp, i0,i1,i2, binary):
  if binary:
    fp.write(struct.pack('<Biii', 3,i0,i1,i2))
  else:
    fp.write('3 %d %d %d\n' % (i0,i1,i2))

def _write_ply_header_line(fp, str, binary):
  if binary:
    fp.write(str.encode())
  else:
    fp.write(str)

def write_ply(path, verts, trias=None, color=None, normals=None, binary=False):
  if verts.shape[1] != 3:
    raise Exception('verts has to be of shape Nx3')
  if trias is not None and trias.shape[1] != 3:
    raise Exception('trias has to be of shape Nx3')
  if color is not None and not callable(color) and not isinstance(color, np.ndarray) and color.shape[1] != 3:
    raise Exception('color has to be of shape Nx3 or a callable')

  mode = 'wb' if binary else 'w'
  with open(path, mode) as fp:
    _write_ply_header_line(fp, "ply\n", binary)
    if binary:
      _write_ply_header_line(fp, "format binary_little_endian 1.0\n", binary)
    else:
      _write_ply_header_line(fp, "format ascii 1.0\n", binary)
    _write_ply_header_line(fp, "element vertex %d\n" % (verts.shape[0]), binary)
    _write_ply_header_line(fp, "property float32 x\n", binary)
    _write_ply_header_line(fp, "property float32 y\n", binary)
    _write_ply_header_line(fp, "property float32 z\n", binary)
    if color is not None:
      _write_ply_header_line(fp, "property uchar red\n", binary)
      _write_ply_header_line(fp, "property uchar green\n", binary)
      _write_ply_header_line(fp, "property uchar blue\n", binary)
    if normals is not None:
      _write_ply_header_line(fp, "property float32 nx\n", binary)
      _write_ply_header_line(fp, "property float32 ny\n", binary)
      _write_ply_header_line(fp, "property float32 nz\n", binary)
    if trias is not None:
      _write_ply_header_line(fp, "element face %d\n" % (trias.shape[0]), binary)
      _write_ply_header_line(fp, "property list uchar int32 vertex_indices\n", binary)
    _write_ply_header_line(fp, "end_header\n", binary)

    for vidx, v in enumerate(verts):
      if color is not None:
        if callable(color):
          c = color(vidx)
        elif color.shape[0] > 1:
          c = color[vidx]
        else:
          c = color[0]
      else:
        c = None
      if normals is None:
        n = None
      else:
        n = normals[vidx]
      _write_ply_point(fp, v[0],v[1],v[2], c, n, binary)

    if trias is not None:
      for t in trias:
        _write_ply_triangle(fp, t[0],t[1],t[2], binary)

def faces_to_triangles(faces):
  new_faces = []
  for f in faces:
    if f[0] == 3:
      new_faces.append([f[1], f[2], f[3]])
    elif f[0] == 4:
      new_faces.append([f[1], f[2], f[3]])
      new_faces.append([f[3], f[4], f[1]])
    else:
      raise Exception('unknown face count %d', f[0])
  return new_faces

def read_ply(path):
  with open(path, 'rb') as f:
    # parse header
    line = f.readline().decode().strip()
    if line != 'ply':
      raise Exception('Header error')
    n_verts = 0
    n_faces = 0
    vert_types = {}
    vert_bin_format = []
    vert_bin_len = 0
    vert_bin_cols = 0
    line = f.readline().decode()
    parse_vertex_prop = False
    while line.strip() != 'end_header':
      if 'format' in line:
        if 'ascii' in line:
          binary = False
        elif 'binary_little_endian' in line:
          binary = True
        else:
          raise Exception('invalid ply format')
      if 'element face' in line:
        splits = line.strip().split(' ')
        n_faces = int(splits[-1])
        parse_vertex_prop = False
      if 'element camera' in line:
        parse_vertex_prop = False
      if 'element vertex' in line:
        splits = line.strip().split(' ')
        n_verts = int(splits[-1])
        parse_vertex_prop = True
      if parse_vertex_prop and 'property' in line:
        prop = line.strip().split()
        if prop[1] == 'float':
          vert_bin_format.append('f4')
          vert_bin_len += 4
          vert_bin_cols += 1
        elif prop[1] == 'uchar':
          vert_bin_format.append('B')
          vert_bin_len += 1
          vert_bin_cols += 1
        else:
          raise Exception('invalid property')
        vert_types[prop[2]] = len(vert_types)
      line = f.readline().decode()

    # parse content
    if binary:
      sz = n_verts * vert_bin_len
      fmt = ','.join(vert_bin_format)
      verts = np.ndarray(shape=(1, n_verts), dtype=np.dtype(fmt), buffer=f.read(sz))
      verts = verts[0].astype(vert_bin_cols*'f4,').view(dtype='f4').reshape((n_verts,-1))
      faces = []
      for idx in range(n_faces):
        fmt = '<Biii'
        length = struct.calcsize(fmt)
        dat = f.read(length)
        vals = struct.unpack(fmt, dat)
        faces.append(vals)
      faces = faces_to_triangles(faces)
      faces = np.array(faces, dtype=np.int32)
    else:
      verts = []
      for idx in range(n_verts):
        vals = [float(v) for v in f.readline().decode().strip().split(' ')]
        verts.append(vals)
      verts = np.array(verts, dtype=np.float32)
      faces = []
      for idx in range(n_faces):
        splits = f.readline().decode().strip().split(' ')
        n_face_verts = int(splits[0]) 
        vals = [int(v) for v in splits[0:n_face_verts+1]]
        faces.append(vals)
      faces = faces_to_triangles(faces)
      faces = np.array(faces, dtype=np.int32)

  xyz = None
  if 'x' in vert_types and 'y' in vert_types and 'z' in vert_types:
    xyz = verts[:,[vert_types['x'], vert_types['y'], vert_types['z']]]
  colors = None
  if 'red' in vert_types and 'green' in vert_types and 'blue' in vert_types:
    colors = verts[:,[vert_types['red'], vert_types['green'], vert_types['blue']]]
    colors /= 255
  normals = None
  if 'nx' in vert_types and 'ny' in vert_types and 'nz' in vert_types:
    normals = verts[:,[vert_types['nx'], vert_types['ny'], vert_types['nz']]]

  return xyz, faces, colors, normals


def _read_obj_split_f(s):
  parts = s.split('/')
  vidx = int(parts[0]) - 1
  if len(parts) >= 2 and len(parts[1]) > 0:
    tidx = int(parts[1]) - 1
  else:
    tidx = -1
  if len(parts) >= 3 and len(parts[2]) > 0:
    nidx = int(parts[2]) - 1
  else:
    nidx = -1
  return vidx, tidx, nidx

def read_obj(path):
  with open(path, 'r') as fp:
    lines = fp.readlines()

  verts = []
  colors = []
  fnorms = []
  fnorm_map = collections.defaultdict(list)
  faces = []
  for line in lines:
    line = line.strip()
    if line.startswith('#') or len(line) == 0:
      continue

    parts = line.split()
    if line.startswith('v '):
      parts = parts[1:]
      x,y,z = float(parts[0]), float(parts[1]), float(parts[2])
      if len(parts) == 4 or len(parts) == 7:
        w = float(parts[3])
        x,y,z = x/w, y/w, z/w
      verts.append((x,y,z))
      if len(parts) >= 6:
        r,g,b = float(parts[-3]), float(parts[-2]), float(parts[-1])
        rgb.append((r,g,b))

    elif line.startswith('vn '):
      parts = parts[1:]
      x,y,z = float(parts[0]), float(parts[1]), float(parts[2])
      fnorms.append((x,y,z))

    elif line.startswith('f '):
      parts = parts[1:]
      if len(parts) != 3:
        raise Exception('only triangle meshes supported atm')
      vidx0, tidx0, nidx0 = _read_obj_split_f(parts[0])
      vidx1, tidx1, nidx1 = _read_obj_split_f(parts[1])
      vidx2, tidx2, nidx2 = _read_obj_split_f(parts[2])

      faces.append((vidx0, vidx1, vidx2))
      if nidx0 >= 0:
        fnorm_map[vidx0].append( nidx0 )
      if nidx1 >= 0:
        fnorm_map[vidx1].append( nidx1 )
      if nidx2 >= 0:
        fnorm_map[vidx2].append( nidx2 )

  verts = np.array(verts)
  colors = np.array(colors)
  fnorms = np.array(fnorms)
  faces = np.array(faces)
  
  # face normals to vertex normals
  norms = np.zeros_like(verts)
  for vidx in fnorm_map.keys():
    ind = fnorm_map[vidx]
    norms[vidx] = fnorms[ind].sum(axis=0)
  N = np.linalg.norm(norms, axis=1, keepdims=True)
  np.divide(norms, N, out=norms, where=N != 0)

  return verts, faces, colors, norms