#!/usr/bin/env python3

import rospy
from tauv_msgs.msg import SonarPulse, FeatureDetection, FeatureDetections
from tauv_msgs.srv import SonarControl, FeatureDetectionsSync
from geometry_msgs.msg import Point
import numpy as np
from scipy.ndimage import gaussian_filter
import skimage
from tauv_util.transforms import tf2_transform_to_homogeneous
import tf2_ros as tf2
import cv2 as cv

#for testing
import matplotlib
import matplotlib.pyplot as plt

POINT_DIM = 3

class SonarGraph:
    def __init__(self):
        rospy.init_node('sonar_map',anonymous=True)

        self._load_config()
   
        x = rospy.wait_for_message("/sonar/pulse",SonarPulse)

        self.cur_angle = x.angle
        self.data_len = x.data_length

        #map used for cluster processing
        self.heat_map = np.zeros((self.HEAT_MAP_DIM, self.HEAT_MAP_DIM), dtype=np.float64)

        #use below to visualize sonar processing
        # self.heat = np.zeros_like(self.heat_map)

        self.global_origin = [self.global_range/2.0, self.global_range/2.0] #yx origin
        self.square_size = self.global_range/self.HEAT_MAP_DIM

        #self._tf_buffer: tf2.Buffer = tf2.Buffer()
        #self._tf_listener: tf2.TransformListener = tf2.TransformListener(self._tf_buffer)

        #for to-world transform operations
        self.RELATIVE_POINT_MAP = np.linspace(0, self.range, num=self.num_distances, endpoint=True)
        self.rot = np.zeros((3,3))
        self.positions = np.empty((self.num_distances, 3, 4))
        self.positions[:,:,3] = 1

        #subscribes to sonar driver publishings
        self.pulse_reader = rospy.Subscriber("/sonar/pulse",SonarPulse, self.read)
        #starts scanning immediately, might wanna start from false
        self.survey_mode = True

        #communicates directly w global map to send survey points
        rospy.wait_for_service("/global_map/sync_detections")
        self.global_map_com = rospy.ServiceProxy("/global_map/sync_detections", FeatureDetectionsSync)

        #survey/start/stop/reset service for mission communications
        self.serve = rospy.Service("/sonar/map", SonarControl, self.handle_sig)

        #used for testing. Publishes all sonar detections
        # self.pub = rospy.Publisher("/feature_detections_sonar", FeatureDetections, queue_size=10)

    def _load_config(self):
        #params for data reading
        self.range = rospy.get_param('sonar_map/device/range')
        self.gradians = rospy.get_param('sonar_map/map/degrees') #not rad OR degrees
        self.num_distances=rospy.get_param('sonar_map/device/num_samples') 
        self.global_range = rospy.get_param('sonar_map/map/global_range') 
        self.residue_factor = rospy.get_param('sonar_map/map/residue_factor') 
        self.confidence_increment = rospy.get_param('sonar_map/map/confidence_increment') 
        self.ignore_radius = rospy.get_param('sonar_map/device/ignore_radius') 
        self.HEAT_MAP_DIM = rospy.get_param('sonar_map/map/map_dimension') 

        #params for heat map processing
        self.blur_sigma = rospy.get_param('sonar_map/map/blur_sigma')
        self.pre_filter_thresh = rospy.get_param('sonar_map/map/pre_filter_thresh')
        self.post_filter_thresh = rospy.get_param('sonar_map/map/post_filter_thresh')
        self.closing_size = rospy.get_param('sonar_map/map/closing_size')
        self.min_object_size = rospy.get_param('sonar_map/map/min_object_size')
        self.line_of_sight_weight = rospy.get_param('sonar_map/map/line_of_sight_weight')

        #@theo please help xP
        #self._tf_namespace: str = rospy.get_param('tf_namespace')
        #self.line_of_sight = rospy.get_param('sonar_map/device/line_of_sight')

    #handles the enumerated interactions w sonar map
    def handle_sig(self, req):
        if(req.op==0): #stop reading in data
            self.survey_mode = False
            return True
        elif(req.op==1): #start reading in data
            self.survey_mode = True
            return True
        elif(req.op==2): #survey, pass detections to map
            res = self.survey ()
            return res
        elif(req.op==3): #reset map
            return (self.__reset ())
        return False

    #resets the map
    def __reset(self):
        self.heat_map = np.zeros((self.HEAT_MAP_DIM, self.HEAT_MAP_DIM), dtype=np.float64)
        return True

    def __to_rad(self,angle):
        return (angle/400.0)*2*np.pi

    #transforms a point to a 2D heat map coordinate
    def __get_map_square(self, point): #y x z
        row = np.round((point[0]+self.global_origin[0])/self.square_size)
        col = np.round((point[1]+self.global_origin[1])/self.square_size)
        if(row>=0 and row<self.HEAT_MAP_DIM and col>=0 and col<self.HEAT_MAP_DIM):
            return [(int)(row), (int)(col)]
        else:
            raise Exception("Point out of map bounds")

    #service callback. Processes current heat map and sends detections to global maps
    def survey(self):
        #clears decaying detections that pile-up over-time
        self.heat_map = np.where(self.heat_map>self.pre_filter_thresh, self.heat_map, 0)

        #blur
        blurred = gaussian_filter(self.heat_map, sigma = self.blur_sigma)

        #threshold filter for detections
        threshed = (blurred>self.post_filter_thresh).astype(np.uint8)

        #closing to connect openings
        closed = cv.morphologyEx(threshed, cv.MORPH_CLOSE, cv.getStructuringElement(cv.MORPH_RECT,(self.closing_size, self.closing_size)))

        #find connected components
        numDets, self.labels, stats, centroids = cv.connectedComponentsWithStats(closed)

        #create a feature detections message to pass to map
        detections = FeatureDetections()
        detections.detector_tag = "sonar"
        detections.detections = []
        status = True

        #for visualization
        # self.heat = np.copy(self.heat_map)

        #add center point to buckets
        for contour in range(numDets):
            if(stats[contour,cv.CC_STAT_AREA]>self.min_object_size):
                #add centroid translated by global origin to bucket
                detection = FeatureDetection()
                detection.position = Point(self.square_size*(centroids[contour][0] - self.global_origin[0]),self.square_size*(centroids[contour][1] - self.global_origin[1]),-50)
                detection.orientation = Point(0,0,0)
                detection.tag = "unknown"
                detection.confidence = 0.5
                detection.SE2 = True

                #for detection visualization
                # self.heat = cv.circle(self.heat, (np.round(centroids[contour][0]).astype(int), np.round(centroids[contour][1]).astype(int)), 10, 155, thickness=5)

                #add all detections at once (sonar does clustering)
                #this initially leads to redundancy problems since map only clusters across frames. So if sonar
                #detections close together but separate, map will not cluster them together
                # detections.detections.append(detection)

                #OR

                #add detections individually (leave it up to map to cluster)
                detections.detections = [detection]
                status = status and self.global_map_com(detections).success

                #for detection visualization
                # self.pub.publish(detections)
         
        #uncomment below as well if sonars clustering
        # status = self.global_map_com(detections).success

        return status

    #ray filtering: performs max/line-of-sight filtering and return a single detection per ray
    def __get_ray(self, data):
        sig = np.frombuffer(data, dtype=np.uint8).copy()
        sig[0:self.ignore_radius]=0

        #sort
        sorted = np.argsort(-sig)

        #calculate scores
        scores = np.array(sig)[sorted]-self.line_of_sight_weight*sorted

        #max arg over values - args
        max_ind = sorted[np.argmax(scores)]
        ray = np.zeros_like(sig)
        ray[max_ind] = 1
        
        #arg 1 else 0
        return sig*ray

    def read(self, data):
        #if we are in survey mode
        if not self.survey_mode:
            return

        #transform data to world coordinates
        coords = self.__transform_to_world(self.__to_rad(data.angle-1), self.__to_rad(data.angle), self.__to_rad(data.angle+1))

        #gets the signal data
        signal = self.__get_ray(data.data)

        #retainment of previous map data, scales down to stay within data range
        new_map = self.heat_map*(1-self.residue_factor)

        try:
            prev_points = np.apply_along_axis(self.__get_map_square, 1, (coords[self.ignore_radius]))

            #iterates over each point in the ray, transforms it to map coordinate
            for i in range(self.ignore_radius, self.data_len-1):
                cur_points = np.apply_along_axis(self.__get_map_square, 1, (coords[i+1]))
                #fits hexagon to approximate sonar ray arc between previous and current distance
                hex = np.concatenate((prev_points, cur_points))
                #fills in the approximated area with the signal detection
                new_map = cv.fillPoly(new_map, [hex], signal[i]*(1-self.residue_factor+self.confidence_increment))
                prev_points = cur_points

        #ends loop at first out-of-bounds transform
        except Exception as err:
            rospy.logwarn(err)
        
        self.heat_map = new_map+(self.heat_map*self.residue_factor)

        #for visualization
        # show = cv.applyColorMap(np.round(self.heat).astype(np.uint8), cv.COLORMAP_JET)
        # cv.imshow("Color-map",show)
        # cv.waitKey(1)

    #TODO the static TFs are quite untested and need some fixing :P
    def __transform_to_world(self, low, mid, high):
        #assuming northwards 0 degrees and NED
        #TODO TODO TODO THIS TRANSFORM WILL NEED TO BE RECALCULATED AT COMP
        self.rot[:,0:2] = [[np.cos(low), -np.sin(low)],[np.cos(mid), -np.sin(mid)],[np.cos(high), -np.sin(high)]]
        self.positions[:,:,0:3] = np.multiply.outer(self.RELATIVE_POINT_MAP,self.rot)

        #TODO FIX
        #world_frame = f'{self._tf_namespace}/vehicle' #change to world
        #sonar_frame = f'{self._tf_namespace}/sonar'

        try:
            #transform = self._tf_buffer.lookup_transform(
            #world_frame,
            #sonar_frame,
            #rospy.Time.now())

            #H = tf2_transform_to_homogeneous(transform)
            #TODO rn its just doing a silly little identity mutiplication, change this to TFs
            H = np.eye(4,4)

            return ((np.matmul (self.positions, H))[:,:,0:3])

        except (tf2.LookupException, tf2.ConnectivityException, tf2.ExtrapolationException) as e:
            rospy.logwarn(f'Could not get transform from {world_frame} to {sonar_frame}: {e}')


def main():
    sonar = SonarGraph()
    rospy.spin()

if __name__ == '__main__':
    main()