import numpy as np
from matplotlib import pyplot as plt
from math import pi, sin, cos, atan2, hypot
import plotly.graph_objects as go

# parameters
tolerance = 0.02        # distance from the plane to accept point
iterations = 1000       # number of iterations

"""
Explanation: <a href="https://colab.research.google.com/github/OSGeoLabBp/tutorials/blob/master/english/data_processing/lessons/ransac_plane.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
Plane fitting with RANSAC
What is RANSAC?
RANdom SAmple Consensus (RANSAC) is an iterative method to make robust non-deterministic estimation of parameters of a mathematical model from a set of data containing outliers.
Parameters
To use RANSAC algorithms first we have to define the following parameters:
data – Set of observations.
model – Model (line, circle, plane, etc.) to explain observed data points.
n – Minimum number of data points required to estimate model parameters.
k – Maximum number of iterations allowed in the algorithm.
t – Threshold value to determine data points that are fit well by model.
d – Number of close data points required to assert that a model fits well to data.

The steps of the algorithm

Set up tolerance and the number of iterations.
Select three random points from the point cloud.
Find the plane of the randomly selected points.
Find the inliers, which are closer to the plane than the tolerance.
If the number of inliers is larger than in any previous solution save the solution as the best so far.
Repeat from step 2 if the number of iterations allows.

RANSAC plane in practice
Finding points that fits into a plane is very popular a segmentation tasks in point cloud segmentation and it usually bases on RANSAC.
In the following the operation of the RANSAC plane algorithm will be presented.
Let's start with the preparation!
During the process the following modules and parameters will be used:
End of explanation
"""


n = 1000 # number of inliers
m = 600  # number of outliers
range = 100.0   # range of x, y, z coordinates from zero to range

# points on a horizontal plane
x = np.random.rand(n + m) * range
y = np.random.rand(n + m) * range
z = np.zeros(n + m)
z[0:n] = np.random.rand(n) * tolerance + range / 2.0    # on plane
z[n:n+m] = np.random.rand(m) * range                    # noise
# random rotate points to have general plane
alpha, beta = np.random.randint(10, 80, 2) / 180.0 * pi
r1 = np.array([[1.0, 0.0, 0.0], [0.0, cos(alpha), -sin(alpha)], [0.0, sin(alpha), cos(alpha)]])
r2 = np.array([[cos(beta), 0.0, sin(beta)], [0.0, 1.0, 0.0], [-sin(beta), 0.0, cos(beta)]])
r = r1.dot(r2)
# rotate points
points = np.c_[x, y, z].dot(r)

"""
Explanation: To test RANSAC method some data will be required. 
Let's generate random points close to a plane and some outliers! For the test data set first points on a horizontal plane and random outliers are genarated. In a second step points are rotated around X and Y axices by random angles between 10 and 80 degree.
End of explanation
"""


fig = go.Figure(data=[go.Scatter3d(x=points[0:n, 0], y=points[0:n, 1], 
                                   z=points[0:n, 2], mode='markers',     
                                   marker=dict(size=3, color='green')),
                      go.Scatter3d(x=points[n:n+m, 0], y=points[n:n+m, 1], 
                                   z=points[n:n+m, 2], mode='markers',     
                                   marker=dict(size=2, color='red'))])
# Add title to the axes
fig.update_layout(scene=dict(
                    aspectmode='manual', aspectratio=dict(x=1, y=1, z=1),
                    xaxis_title='X', yaxis_title='Y', zaxis_title='Z'),
                    width=700, margin=dict(r=20, b=10, l=10, t=10)
)
fig.update_layout(showlegend=False)
fig.show()

"""
Explanation: Let's display the points!
We use two visualization method a dynamic one using Plotly package (this is not visible on the GitHub page) and a static one using Matplotlib.
In Plotly you can rotate, pan and zoom the fugure.
End of explanation
"""


fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(points[n:n+m, 0], points[n:n+m, 1], points[n:n+m, 2], label='outliers')
ax.scatter(points[0:n, 0], points[0:n, 1], points[0:n, 2], label='inliers')
ax.legend()
ax.view_init(30, 70) 

"""
Explanation: Change the two angles (degrees) in view_init to change the point of view. The first angle is the azimut, the second is the elevation angle. For example a top view is (0, 90), a front view is (0, 0).
End of explanation
"""


if points.shape[1] < 4:
    points = np.c_[points, np.full((n+m, 1), 1)]    # change to homogenouos coord
best_n = 0          # number of points on the best plane so far
best_i = 0          # iteration index of best plane so far
best_inliers = np.array([]) # indices of inliers of the best plane so far
for i in np.arange(iterations):
    # select three random points
    l = []  # list of random indices for points
    while len(l) != 3:
        l = list(set(np.random.randint(n+m, size=3))) # remove repeated random integers
    pp = points[l]  # randomly selected points
    # plane equation from the three points using homogenouos coordinates
    a = np.linalg.det(pp[:,1:])
    b = -np.linalg.det(pp[:, [0, 2, 3]])
    c = np.linalg.det(pp[:, [0, 1, 3]])
    d = -np.linalg.det(pp[:, 0:3])
    plane = np.array([a, b, c, d])
    plane = plane / np.sqrt(plane[0:3].dot(plane[0:3]))   # normal to unit vector
    # select close points
    inliers = points[np.abs(np.dot(points, plane)) < tolerance]
    if inliers.shape[0] > best_n:
        # better solution found
        best_n = inliers.shape[0]
        best_i = i
        best_inliers = inliers[:,:3].copy()
        best_plane = plane.copy()
print(f'Best solution after {best_i} iterations, {best_n} points on plane: {best_plane}')

"""
Explanation: Iteration to find RANSAC plane. The plane parameters (a, b, c, d) are calculated using homogenouos coordinates.
End of explanation
"""


fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(best_inliers[:, 0], best_inliers[:, 1], best_inliers[:, 2], label='inliers')
#ax.scatter(points[:, 0], points[:, 1], points[:, 2], label='outliers')
ax.legend()
ax.view_init(30, 70)

"""
Explanation: Try to run the block above several times and check best_i and best_n to fine tune the number interations. Do you get the same result is you repeat the calculation?
End of explanation
"""


!pip install open3d -q

"""
Explanation: Tasks

Run the iteration block above and check best_i and best_n to fine tune interation number.

Find multiple planes in a point cloud
Mostly we are looking for not a single plane but the simplified model of the object in the point cloud. In this sample we try to find the faces (linear) of a geometrical object.

To manage point clouds there are several Python packages like  Open3D, PDAL, lidR, pyntcloud, etc.
In the following we will use Open3D, which can handle RANSAC in addition to point clouds.
At first install Open3D module! It may take few minutes.
End of explanation
"""


!wget -q https://github.com/OSGeoLabBp/tutorials/raw/master/english/data_processing/lessons/code/house.pcd -O sample_data/house.pcd

"""
Explanation: A small point cloud of a building will be used, let's download it from GitHub!
End of explanation
"""


import open3d as o3d
import numpy as np

def label(ang):
    """ return roof/wall/ground label from angle of normal to horizontal """
    label = '???'
    if abs(ang < 15):
        label = 'wall'
    elif 80 < abs(ang) < 100:
        label = 'ground'
    elif 35 < abs(ang) < 55:
        label = 'roof'
    return label

n_planes = 8
tol = 0.1           # tolerance distance from plane
n_points = 3        # number of random points
iterations = 300
# load point cloud
pc = o3d.io.read_point_cloud('sample_data/house.pcd')
n_orig = len(pc.points)
pc_orig = pc.select_by_index([], invert=True)   # make a copy of original
planes = []     # list of result planes
fig = plt.figure(figsize=(23, 3))

i = 0
while i < n_planes:
    plane, inliers = pc.segment_plane(tol, n_points, iterations)
    angle = atan2(plane[2], hypot(plane[0], plane[1])) * 180. / pi  # angle of normal from horizontal
    plane_pc = pc.select_by_index(inliers)
    o3d.io.write_point_cloud('sample_data/plane_' + str(i) + '.pcd', plane_pc) # export planes
    ax = fig.add_subplot(1, n_planes, i+1, projection='3d')
    ax.set_title(f'{angle:.0f} degree {label(angle)}\n {len(inliers)} points')
    xyz = np.asarray(plane_pc.points)
    ax.scatter(xyz[:,0], xyz[:,1], xyz[:,2])
    #ch = plane_pc.compute_convex_hull()
    planes.append([plane, plane_pc])
    pc = pc.select_by_index(inliers, invert=True)   # rest of the point cloud
    i += 1
print(f'{len(pc.points)} points left from {n_orig}')

"""
Explanation: Find n_planes number of planes having most of the points.
End of explanation
"""


data = []
colors = ['red', 'blue', 'green', 'cyan', 'yellow', 'magenta', 'brown', 'orange']
for i, pp in enumerate(planes):
    xyz = np.asarray(pp[1].points)
    data.append(go.Scatter3d(x=xyz[:, 0], y=xyz[:, 1], z=xyz[:, 2], mode='markers',
                             marker=dict(size=1, color=colors[i % len(planes)])))
fig = go.Figure(data=data)
fig.update_layout(scene=dict(aspectmode='manual', aspectratio=dict(x=1, y=1, z=1),
                    xaxis_title='X', yaxis_title='Y', zaxis_title='Z'),
                    width=700, margin=dict(r=20, b=10, l=10, t=10))
fig.update_layout(showlegend=False)
fig.show()

"""
Explanation: Let's display the result!
End of explanation
"""


