#!/usr/bin/env luajit
--
-- Plot trajectory of dimers and solvent.
-- Copyright © 2013 Peter Colberg.
-- For conditions of distribution and use, see copyright notice in LICENSE.
--

local hdf5 = require("hdf5")
local vtk  = require("vtk")
local ffi  = require("ffi")

local file = hdf5.open_file(arg[1])
local dataset = file:open_dataset("trajectory/dimer/position/value")
local filespace = dataset:get_space()
local filedims = filespace:get_simple_extent_dims()
local dataset_s = file:open_dataset("trajectory/solvent/position/value")
local filespace_s = dataset_s:get_space()
local filedims_s = filespace_s:get_simple_extent_dims()

local Nd = filedims[2]
local Ns = filedims_s[2]
local X, Y, Z
do
  local group = file:open_group("trajectory/dimer/box")
  local dset = group:open_dataset("value")
  local buf = ffi.new("double[2][3]")
  local dims = ffi.new("hsize_t[2]", 2, 3)
  local memspace = hdf5.create_simple_space(2, dims)
  dset:read("native_double", memspace, nil, nil, buf)
  X = buf[1][0] - buf[0][0]
  Y = buf[1][1] - buf[0][1]
  Z = buf[1][2] - buf[0][2]
end
local sigma
do
  local group = file:open_group("trajectory/dimer/species")
  local attr = group:open_attribute("diameter")
  local buf = ffi.new("double[2]")
  attr:read("native_double", buf)
  sigma = {buf[0], buf[1]}
end

local array = vtk.float_array()
array:set_number_of_components(3)
array:set_number_of_tuples(Nd)
local pos = ffi.cast("float (*)[3]", array:get_pointer())
local polydata = vtk.poly_data()
do
  local points = vtk.points()
  for i = 0, Nd-1 do pos[i][0], pos[i][1], pos[i][2] = 0, 0, 0 end
  points:set_data(array)
  polydata:set_points(points)
end
do
  local array = vtk.float_array()
  array:set_number_of_tuples(Nd)
  polydata:get_point_data():set_scalars(array)
  local color = array:get_pointer()
  for i = 0, Nd-1, 2 do color[i], color[i + 1] = 0, 1 end
end
do
  local array = vtk.float_array()
  array:set_number_of_components(3)
  array:set_number_of_tuples(Nd)
  polydata:get_point_data():set_vectors(array)
  local scale = ffi.cast("float (*)[3]", array:get_pointer())
  for i = 0, Nd-1, 2 do
    scale[i][0] = sigma[1]
    scale[i][1] = 0
    scale[i][2] = 0
    scale[i + 1][0] = sigma[2]
    scale[i + 1][1] = 0
    scale[i + 1][2] = 0
  end
end
local spheresource = vtk.sphere_source()
spheresource:set_center(0, 0, 0)
spheresource:set_radius(0.5)
spheresource:set_phi_resolution(40)
spheresource:set_theta_resolution(40)
local sphereglyph = vtk.glyph3d()
sphereglyph:set_source(spheresource:get_output())
sphereglyph:set_input(polydata)
sphereglyph:set_color_mode("scalar")
sphereglyph:set_scale_mode("vector")
sphereglyph:update()
local spheremapper = vtk.poly_data_mapper()
spheremapper:set_input_connection(sphereglyph:get_output_port())
local sphereactor = vtk.actor()
sphereactor:set_mapper(spheremapper)

local array_s = vtk.float_array()
array_s:set_number_of_components(3)
array_s:set_number_of_tuples(Ns)
local pos_s = ffi.cast("float (*)[3]", array_s:get_pointer())
local polydata_s = vtk.poly_data()
do
  local points = vtk.points()
  for i = 0, Ns-1 do pos_s[i][0], pos_s[i][1], pos_s[i][2] = 0, 0, 0 end
  points:set_data(array_s)
  polydata_s:set_points(points)
end
local vertexglyph = vtk.vertex_glyph_filter()
vertexglyph:set_input(polydata_s)
local vertexmapper = vtk.poly_data_mapper()
vertexmapper:set_input_connection(vertexglyph:get_output_port())
local vertexactor = vtk.actor()
vertexactor:set_mapper(vertexmapper)

local renderer = vtk.renderer()
renderer:add_actor(sphereactor)
renderer:add_actor(vertexactor)
renderer:set_background(0.1, 0.2, 0.3)
local render_window = vtk.render_window()
render_window:add_renderer(renderer)
local window_to_image_filter
local ogg_theora_writer
local render_window_interactor
if arg[2] then
  render_window:set_off_screen_rendering(true)
  render_window:set_size(720, 720)
  window_to_image_filter = vtk.window_to_image_filter()
  window_to_image_filter:set_input(render_window)
  ogg_theora_writer = vtk.ogg_theora_writer()
  ogg_theora_writer:set_file_name(arg[2])
  ogg_theora_writer:set_input_connection(window_to_image_filter:get_output_port())
else
  render_window_interactor = vtk.render_window_interactor()
  render_window_interactor:set_render_window(render_window)
end

local text_actor = vtk.text_actor()
text_actor:get_text_property()
text_actor:get_text_property():set_font_size(24)
text_actor:set_position2(10,40)
text_actor:get_text_property():set_color(1, 1, 1)
renderer:add_actor2d(text_actor)

do
  local points = vtk.points()
  local Nd = 8
  points:set_number_of_points(Nd)
  points:set_point(0, 0, 0, 0)
  points:set_point(1, X, 0, 0)
  points:set_point(2, 0, Y, 0)
  points:set_point(3, 0, 0, Z)
  points:set_point(4, X, Y, 0)
  points:set_point(5, 0, Y, Z)
  points:set_point(6, X, 0, Z)
  points:set_point(7, X, Y, Z)
  local cells = vtk.cell_array()
  do
    local polyline = vtk.poly_line()
    local point_ids = polyline:get_point_ids()
    point_ids:set_number_of_ids(4)
    point_ids:set_id(0, 0)
    point_ids:set_id(1, 1)
    point_ids:set_id(2, 4)
    point_ids:set_id(3, 7)
    cells:insert_next_cell(polyline)
    local polyline = vtk.poly_line()
    local point_ids = polyline:get_point_ids()
    point_ids:set_number_of_ids(4)
    point_ids:set_id(0, 1)
    point_ids:set_id(1, 6)
    point_ids:set_id(2, 7)
    point_ids:set_id(3, 5)
    cells:insert_next_cell(polyline)
    local polyline = vtk.poly_line()
    local point_ids = polyline:get_point_ids()
    point_ids:set_number_of_ids(4)
    point_ids:set_id(0, 6)
    point_ids:set_id(1, 3)
    point_ids:set_id(2, 5)
    point_ids:set_id(3, 2)
    cells:insert_next_cell(polyline)
    local polyline = vtk.poly_line()
    local point_ids = polyline:get_point_ids()
    point_ids:set_number_of_ids(4)
    point_ids:set_id(0, 3)
    point_ids:set_id(1, 0)
    point_ids:set_id(2, 2)
    point_ids:set_id(3, 4)
    cells:insert_next_cell(polyline)
  end
  local polydata = vtk.poly_data()
  polydata:set_points(points)
  polydata:set_lines(cells)
  local mapper = vtk.poly_data_mapper()
  mapper:set_input(polydata)
  local actor = vtk.actor()
  actor:set_mapper(mapper)
  renderer:add_actor(actor)
end

local count = ffi.new("hsize_t[3]", 1, Nd, 3)
local start = ffi.new("hsize_t[3]", 0, 0, 0)
local dims = ffi.new("hsize_t[2]", Nd, 3)
local memspace = hdf5.create_simple_space(2, dims)

local count_s = ffi.new("hsize_t[3]", 1, Ns, 3)
local start_s = ffi.new("hsize_t[3]", 0, 0, 0)
local dims_s = ffi.new("hsize_t[2]", Ns, 3)
local memspace_s = hdf5.create_simple_space(2, dims_s)

if render_window_interactor then render_window_interactor:initialize() end
if ogg_theora_writer then ogg_theora_writer:start() end
local mass = {math.pow(sigma[1], 3), math.pow(sigma[2], 3)}
local mass_inv = 1 / (mass[1] + mass[2])
repeat
  for n = 0, filedims[1] - 1 do
    if render_window_interactor and not render_window_interactor:process_events() then return end
    start[0] = n
    filespace:select_hyperslab("set", start, nil, count)
    dataset:read("native_float", memspace, filespace, nil, pos)
    for i = 0, Nd-1, 2 do
      local dx = pos[i + 1][0] - pos[i][0]
      local dy = pos[i + 1][1] - pos[i][1]
      local dz = pos[i + 1][2] - pos[i][2]
      dx = dx - math.floor(dx / X + 0.5) * X
      dy = dy - math.floor(dy / Y + 0.5) * Y
      dz = dz - math.floor(dz / Z + 0.5) * Z
      local x_cm = pos[i][0] + mass[2] * mass_inv * dx
      local y_cm = pos[i][1] + mass[2] * mass_inv * dy
      local z_cm = pos[i][2] + mass[2] * mass_inv * dz
      x_cm = x_cm - math.floor(x_cm / X) * X
      y_cm = y_cm - math.floor(y_cm / Y) * Y
      z_cm = z_cm - math.floor(z_cm / Z) * Z
      pos[i][0] = x_cm - mass[2] * mass_inv * dx
      pos[i][1] = y_cm - mass[2] * mass_inv * dy
      pos[i][2] = z_cm - mass[2] * mass_inv * dz
      pos[i + 1][0] = x_cm + mass[1] * mass_inv * dx
      pos[i + 1][1] = y_cm + mass[1] * mass_inv * dy
      pos[i + 1][2] = z_cm + mass[1] * mass_inv * dz
    end
    polydata:modified()
    start_s[0] = n
    filespace_s:select_hyperslab("set", start_s, nil, count_s)
    dataset_s:read("native_float", memspace_s, filespace_s, nil, pos_s)
    for i = 0, Ns-1 do
      pos_s[i][0] = pos_s[i][0] - math.floor(pos_s[i][0] / X) * X
      pos_s[i][1] = pos_s[i][1] - math.floor(pos_s[i][1] / Y) * Y
      pos_s[i][2] = pos_s[i][2] - math.floor(pos_s[i][2] / Z) * Z
    end
    polydata_s:modified()
    text_actor:set_input(tostring(n))
    render_window:render()
    if window_to_image_filter then
      window_to_image_filter:modified()
      ogg_theora_writer:write()
    end
  end
until not render_window_interactor
if ogg_theora_writer then ogg_theora_writer:finish() end