#!/usr/bin/env python
# coding: utf-8

# # Trajectory interpolation
# 
# Notebook: Lukas Winiwarter, 2023
# 
# In this demo, we present how an existing or external trajectory can be used to model platform location and attitude in a HELIOS++ simulation.

# In[1]:


from pathlib import Path
from IPython.display import Code
from pyhelios.util.xmldisplayer import display_xml, find_playback_dir


# In[2]:


import os

os.chdir("..")


# ## Survey
# 
# In this example, we simulate an airborne laser scanner acquisition flying inside of a box of 100 m side length. This allows us to see the full point pattern. We set the `platform` to `interpolated`, and define a trajectory file in the `<platformSettings >` tag. This is provided by a filename and a column order. By setting `syncGPSTime`, we ensure that the output point cloud is in the same time system as the input trajectory.

# In[3]:


Code(display_xml("data/surveys/demo/box_survey_interp.xml"))


# The contents of the trajectory file are simple: For the first 2.5 seconds, we move from `(0, -40, 0)` to `(0, 0, 0)`, with a constant attitude of `(0,0,0)` (roll, pitch, yaw). In the next 2.5 seconds, we continue to move to `(0, 40, 0)`, but rotate (roll) until we have an attitude of `(90, 0, 0)`. HELIOS++ linearly interpolates between these values.
# 
# The second half of the file show a movement from `(-40, 0, 0)` to `(40, 0, 0)`. Note how, for the `interpolated` platform, the `yaw` angle has to be set explicitely, in this case to `90` degrees (i.e., due east). Apart from these changes, the 

# In[4]:


Code(open("data/trajectories/flyandrotate.trj").read())


# ## Execute the simulation

# In[5]:


get_ipython().system('helios data/surveys/demo/box_survey_interp.xml -q')


# ## Results
# 
# Now let's find the output files and display a 3D plot.

# In[6]:


import pandas as pd
import matplotlib.pyplot as plt

output_path = find_playback_dir("data/surveys/demo/box_survey_interp.xml")

print("Loading points from", output_path)
PC = pd.read_csv(
    Path(output_path) / "leg000_points.xyz",
    names="X Y Z intensity echoWidth returnNumber numberOfReturns fullwaveIndex hitObjectId class gpsTime".split(
        " "
    ),
    delimiter=" ",
)
trj = pd.read_csv(
    Path(output_path) / "leg000_trajectory.txt",
    names="X Y Z gpsTime roll pitch yaw".split(" "),
    delimiter=" ",
)

PC1 = PC[PC["gpsTime"] <= 2.5]  # select all points recorded in the first half
PC2 = PC[PC["gpsTime"] > 2.5]  # select all points recorded in the second half

fig = plt.figure(figsize=(15, 10))
ax = fig.add_subplot(projection="3d")

ax.scatter(PC1["X"], PC1["Y"], PC1["Z"], s=0.001, c=PC1["gpsTime"])
ax.scatter(PC2["X"], PC2["Y"], PC2["Z"], s=0.001, c=PC2["gpsTime"])
ax.scatter(trj["X"], trj["Y"], trj["Z"], s=10, c="red")
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")

ax.tick_params(labelsize=16)
plt.show()


# In the first plot, we see the platform travelling along th y-axis. After 2.5 seconds, it reaches (0,0,0), and starts rolling (right wing downwards). In consequence, the laser swath travels left and upwards on the face of the cube.
# 
# In the second plot (below), the same thing happens, but with a different orientation (yaw angle).

# In[7]:


PC = pd.read_csv(
    Path(output_path) / "leg001_points.xyz",
    names="X Y Z intensity echoWidth returnNumber numberOfReturns fullwaveIndex hitObjectId class gpsTime".split(
        " "
    ),
    delimiter=" ",
)
trj = pd.read_csv(
    Path(output_path) / "leg001_trajectory.txt",
    names="X Y Z gpsTime roll pitch yaw".split(" "),
    delimiter=" ",
)

PC1 = PC[PC["gpsTime"] <= 7.5]  # select all points recorded in the first half
PC2 = PC[PC["gpsTime"] > 7.5]  # select all points recorded in the second half

fig = plt.figure(figsize=(15, 10))
ax = fig.add_subplot(projection="3d")

ax.scatter(PC1["X"], PC1["Y"], PC1["Z"], s=0.001, c=PC1["gpsTime"])
ax.scatter(PC2["X"], PC2["Y"], PC2["Z"], s=0.001, c=PC2["gpsTime"])
ax.scatter(trj["X"], trj["Y"], trj["Z"], s=10, c="red")
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")

ax.tick_params(labelsize=16)
plt.show()


# In[ ]: