'''узел программы генерации точек возле поверхности мэшей'''
import pickle
import random

import rclpy
from rclpy.node import Node
from visualization_msgs.msg import MarkerArray, Marker

from gen_locations_on_mesh_msgs.msg import LocationsOnMesh
from gen_locations_on_mesh_msgs.srv import GenerateLocation, RemoveLocation

import numpy as np

import trimesh
# Библиотека работы с мэшами.
# Мэш должен состоять только из треугольников.
# Совпадение координат вершин не допускается.
# доступ к частям мэша:
# mesh = trimesh.load_mesh(filename)
# mesh.edges  # (ne,2)
# mesh.faces  # (nf,3)
# mesh.vertices  # (nv,3)


def get_all_triangles_with_correct_orientation(vs, fs, max_a):
    '''создать список всех треугольных граней, ориентированных
    одинаково: при виде сверху точки проходятся по часовой стрелке.
    Грани с углом наклона к вертикали больше max_a радиан не рассматриваются'''
    min_cos = np.cos(max_a)
    res_faces = []  # [(ind0, ind1, ind2, normal)]
    # для каждой грани
    for i in range(fs.shape[0]):
        # извлечь индексы вершин
        ind0 = fs[i,0]
        ind1 = fs[i,1]
        ind2 = fs[i,2]
        # извлечь координаты вершин
        v0 = vs[ind0,:]
        v1 = vs[ind1,:]
        v2 = vs[ind2,:]
        # сформировать векторы сторон
        sd1 = v1 - v0
        sd2 = v2 - v0
        # рассчитать нормаль
        nrm = np.cross(sd2,sd1)
        nrm /= np.linalg.norm(nrm)
        # если грань слишком наклонена, она не рассматривается
        if abs(nrm[2]) < min_cos:
            continue
        # определяем порядок вершин
        # по направлению нормали
        if nrm[2] > 0:
            res = (ind0, ind1, ind2, nrm)
        else:
            res = (ind0, ind2, ind1, -nrm)
        res_faces.append(res)
        #if i % 10000 == 0:
        #    print(f'{i/fs.shape[0]*100:3.1f}%')
    return res_faces

def calc_connections(fs):
    '''рассчитать для каждой грани её соединения
    --- список граней, с которой она связана'''
    connections = {}  # face_id -> [face_id]
    unknown_conns = {}  # (v1, v2) -> face_id
    for i,(ind0, ind1, ind2, _) in enumerate(fs):
        # для каждого ребра: если оно уже встречалось,
        # добавим связь; если нет, вставим в список ожидания
        for edg in [(ind1, ind0), (ind2, ind1), (ind0, ind2)]:
            if edg in unknown_conns:
                # добавить связь между i и unknown_conns[edg]
                if i in connections:
                    connections[i].append(unknown_conns[edg])
                else:
                    connections[i] = [unknown_conns[edg]]
                if unknown_conns[edg] in connections:
                    connections[unknown_conns[edg]].append(i)
                else:
                    connections[unknown_conns[edg]] = [i]
                # убрать ребро
                del unknown_conns[edg]
            else:  # связи нет, добавляем её в список ожидания
                unknown_conns[(edg[1], edg[0])] = i
    return connections

def calc_pieces(faces, conns):
    '''найти непрерывные куски поверхности'''
    part_used = [False] * len(faces)
    parts = []
    curr_index = 0
    # обходим все грани
    while curr_index < len(faces):
        if part_used[curr_index]:
            # пропустить уже использованную грань
            curr_index += 1
            continue
        # начинаем новый кусок с текущей грани
        new_part = set([curr_index])
        part_used[curr_index] = True
        # поместим в очередь все грани, которые надо искать на текущем шаге
        queue = [curr_index]
        # пока очередь не закончится, будем искать грани
        while len(queue) > 0:
            curr_f = queue[0]
            del queue[0]
            # пропустить грань без соседей
            if curr_f not in conns:
                continue
            # для всех соседей:
            for nb in conns[curr_f]:
                if not part_used[nb]:
                    # добавляем грань в кусок
                    part_used[nb] = True
                    # ставим её в очередь
                    queue.append(nb)
                    new_part.add(nb)
        # очередь закончена, кусок сформирован
        parts.append(new_part)
        curr_index += 1

    return parts

def find_largest_piece(pieces):
    '''найти самый большой кусок мэшей'''
    # TODO учитывать не количество, а площадь
    if len(pieces) == 0:
        raise ValueError('No pieces found')
    ls = [len(v) for v in pieces]
    max_len = max(ls)
    return pieces[ls.index(max_len)]

def gen_new_point(piece, trs, vs, existing_ps, dist_limit, vert_shift,
                  max_z_coord, num_generation_trials):
    '''создать новую точку случайным образом;
    вернуть словарь точек (с новой)'''
    # будем снова и снова пытаться сделать точку, удалённую от других;
    # если после множества попыток это не получится, смиримся
    # и примем ту, что есть
    curr_trial = 0
    while curr_trial < num_generation_trials:
        # выбрать случайный треугольник
        tr_ind = random.randrange(len(piece))
        # сформировать координаты случайной точки в нём
        alpha = random.random()
        beta = random.random()
        if alpha + beta > 1:
            kappa = (alpha + beta - 1) * 0.5
            alpha -= 2 * kappa
            beta -= 2 * kappa
        # получить её абсолютные координаты
        ind0, ind1, ind2, nrm = trs[tr_ind]
        c0 = vs[ind0,:]
        c1 = vs[ind1,:]
        c2 = vs[ind2,:]
        p = c0 + (c1-c0)*alpha + (c2-c0)*beta
        p[2] += vert_shift
        if p[2] > max_z_coord:
            continue
        # проверить совпадение координат с существующими точками
        succ = True
        for ep in existing_ps.values():
            if np.linalg.norm(p - ep) < dist_limit:
                succ = False
                break
        if succ or curr_trial == num_generation_trials - 1:
            break
    # выдать найденную точку
    if len(existing_ps) > 0:
        new_ind = max(existing_ps.keys()) + 1
    else:
        new_ind = 1
    existing_ps[new_ind] = p
    return new_ind

class GenLocationsOnMeshNode(Node):

    def __init__(self):
        super().__init__('gen_locations_on_mesh_node')

        # состояние
        self.vs = None
        self.trs = None
        self.the_piece = None
        self.existing_points = {}
        self.dropped_points = {}

        self.declare_parameter('mesh_filename', '')
        self.declare_parameter('init_number_of_points', 10)
        self.declare_parameter('max_z_coord', 2.)
        self.declare_parameter('min_distance_between_points', 2.)
        self.declare_parameter('height_upon_surface', 0.1)
        self.declare_parameter('num_generation_trials', 1000)

        self._load_data()

        # создать инфраструктуру ROS
        self.state_pub_msg = LocationsOnMesh()
        self.state_pub = self.create_publisher(LocationsOnMesh, 'locations', 1)
        self.markers_msg = MarkerArray()
        self.markers_dict = {}
        self.markers_pub = self.create_publisher(MarkerArray, 'locations_markers', 1)
        self.generate_location_srv = self.create_service(GenerateLocation,
                                                         'generate_location',
                                                         self.generate_location_cb)
        self.remove_location_srv = self.create_service(RemoveLocation,
                                                       'remove_location',
                                                       self.remove_location_cb)
        self.pub_timer = self.create_timer(0.1, self.pub_timer_cb)
        # сгенерировать нужное количество точек
        num_ps = self.get_parameter('init_number_of_points').get_parameter_value().integer_value
        for i in range(num_ps):
            self.gen_new_point()

    def pub_timer_cb(self):
        '''таймер, периодически публикующий координаты точек'''
        self.publish_unity_coords()
        self.publish_real_coords_markers()

    def publish_unity_coords(self):
        '''опубликовать результаты в левостроронней системе координат Unity'''
        # собрать в сообщение все данные про объекты
        self.state_pub_msg.location_id = []
        self.state_pub_msg.x = []
        self.state_pub_msg.y = []
        self.state_pub_msg.z = []
        for l_id, p in self.existing_points.items():
            self.state_pub_msg.location_id.append(l_id)
            # self.state_pub_msg.x.append(p[0])
            # self.state_pub_msg.y.append(p[1])
            # self.state_pub_msg.z.append(p[2])
            self.state_pub_msg.x.append(-p[1] - 400.)
            self.state_pub_msg.y.append(p[2])
            self.state_pub_msg.z.append(p[0] + 400.)
        # опубликовать их
        self.state_pub.publish(self.state_pub_msg)

    def publish_real_coords_markers(self):
        '''опубликовать данные в правосторонней системе координат реального мира'''
        curr_timestamp = self.get_clock().now().to_msg()
        for m in self.markers_dict.values():
            m.header.stamp = curr_timestamp
        self.markers_msg.markers = list(self.markers_dict.values())
        self.markers_pub.publish(self.markers_msg)

    def _load_data(self):
        '''загрузка мэша и его предобработка'''
        # сначала пытаемся загрузить кэшированный вариант
        cache_filename = self.get_parameter('mesh_filename').get_parameter_value().string_value + '.pkl'
        try:
            with open(cache_filename, 'rb') as f:
                dt = pickle.load(f)
                self.vs, self.trs, self.the_piece = dt
        except (pickle.PickleError,
                FileNotFoundError):
            # загрузка кэша не удалась, обрабатываем файлы полностью
            self.get_logger().info('Loading data from mesh file in gen_locations_on_mesh')
            mesh_filename = self.get_parameter('mesh_filename').get_parameter_value().string_value
            mesh = trimesh.load_mesh(mesh_filename)
            self.vs = mesh.vertices
            self.trs = get_all_triangles_with_correct_orientation(mesh.vertices,
                                                                mesh.faces,
                                                                1.5)
            #conns = calc_connections(self.trs)
            #pieces = calc_pieces(self.trs, conns)
            #self.the_piece = find_largest_piece(pieces)
            self.the_piece = list(range(len(self.trs)))
            # сохраняем кэш
            try:
                with open(cache_filename, 'wb') as f:
                    dt = (self.vs, self.trs, self.the_piece)
                    pickle.dump(dt, f)
                self.get_logger().info('Cache for gen_locations_on_mesh saved successfully')
            except pickle.PickleError:
                self.get_logger().warn('Error while saving cache in gen_locations_on_mesh')


    def generate_location_cb(self, request, response):
        '''callback for generation service'''
        ind = self.gen_new_point()
        response.location_id = ind
        response.x = self.existing_points[ind][0]
        response.y = self.existing_points[ind][1]
        response.z = self.existing_points[ind][2]
        return response

    def gen_new_point(self):
        '''create point, return its index'''
        max_z_coord = self.get_parameter('max_z_coord').get_parameter_value().double_value
        min_distance_between_points = self.get_parameter('min_distance_between_points').get_parameter_value().double_value
        height_upon_surface = self.get_parameter('height_upon_surface').get_parameter_value().double_value
        num_generation_trials = self.get_parameter('num_generation_trials').get_parameter_value().integer_value
        ind = gen_new_point(self.the_piece,
                            self.trs,
                            self.vs,
                            self.existing_points,
                            min_distance_between_points,
                            height_upon_surface,
                            max_z_coord,
                            num_generation_trials)
        # create marker for the new point
        m = Marker()
        m.header.frame_id = 'map'
        m.id = ind
        m.type = Marker.SPHERE
        m.action = Marker.ADD
        m.pose.position.x = self.existing_points[ind][0] + 400
        m.pose.position.y = self.existing_points[ind][1] + 400
        m.pose.position.z = self.existing_points[ind][2]
        m.pose.orientation.w = 1.
        m.scale.x = 1.
        m.scale.y = 1.
        m.scale.z = 1.
        m.color.r = 1.
        m.color.a = 1.
        self.markers_dict[ind] = m
        return ind

    def remove_location_cb(self, request, response):
        #self.get_logger().info('Incoming request RM')
        if request.location_id in self.existing_points:
            self.dropped_points[request.location_id] = self.existing_points[request.location_id]
            del self.existing_points[request.location_id]
            # update marker
            m = self.markers_dict[request.location_id]
            m.color.r = 0.
            m.color.g = 0.5
            m.color.b = 1.
            response.success = True
        else:
            response.success = False
        return response


def main(args=None):
    rclpy.init(args=args)
    gen_node = GenLocationsOnMeshNode()
    try:
        rclpy.spin(gen_node)
    except KeyboardInterrupt:
        pass

if __name__ == '__main__':
    main()
