# This is a CamiTK python action
#
# Testing MeshComponent in python: compute the mesh barycenter and use matplotlib
# to show a 3D plots of a sample of the points and the barycenter.

import numpy as np
import camitk
import matplotlib
matplotlib.use('Qt5Agg')
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # Needed to activate 3D projection
from PySide2.QtWidgets import *

def init(self:camitk.Action):
    camitk.info("CamiTK version: " + camitk.__version__)
    # # camitk.Application.open(camitk.Core.get_test_data_dir() + "/brain.mha")
    # camitk.Application.open(camitk.Core.getTestDataDir() + "/sinus_skin.vtk")
    # for c in camitk.Application.getTopLevelComponents():
    #     print(c.getName()) # print will print to CamiTK console without timestamp, level and location in code

def compute_barycenter(self:camitk.Action):
    points = self.getTargets()[-1].getPointSetAsNumpy()
    barycenter = np.mean(points, axis=0)
    self.setParameterValue("Barycenter", "QVector3D(" + str(barycenter[0]) + "," + str(barycenter[1]) + "," + str(barycenter[2]) +")")
    return points, barycenter

def targetDefined(self:camitk.Action):
    compute_barycenter(self)

def subsample_points(points, max_points=1000):
    n = points.shape[0]
    if n <= max_points:
        return points  # nothing to do
    indices = np.random.choice(n, size=max_points, replace=False)
    return points[indices]

def process(self:camitk.Action):
    points, barycenter = compute_barycenter(self)

    # 3D matplot is too slow for big mesh, only takes 1 every 100 points
    
    subsampled = subsample_points(points)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    plt.get_current_fig_manager().set_window_title(self.getTargets()[-1].getName())

    # Plot all points
    ax.scatter(subsampled[:,0], subsampled[:,1], subsampled[:,2], c='blue', marker='o', label='Points')

    # Plot barycenter
    ax.scatter(barycenter[0], barycenter[1], barycenter[2],
            c='red', marker='^', s=100, label='Barycenter')

    # Labels
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    # Legend
    ax.legend()

    # Equal aspect ratio (optional but looks better)
    ax.set_box_aspect([np.ptp(subsampled[:,0]), np.ptp(subsampled[:,1]), np.ptp(subsampled[:,2])])

    plt.show()

    msg = QMessageBox()
    msg.setIcon(QMessageBox.Information)

    # setting message for Message Box
    msg.setText("You can now check the plot...")
    
    # setting Message box window title
    msg.setWindowTitle("Plot Information")
    
    # declaring buttons on Message Box
    msg.setStandardButtons(QMessageBox.Ok | QMessageBox.Cancel)
    msg.exec_()
    
    return