"""
This project produces "detections" for DeepSORT
in an attempt to trick the algorithm into moonwalking
over a crowd.

The framerate of rendering and detection can be distinct.
Also, all parameters (incl. framerate) can change along the way, 
thus positions cannot depend on that.
"""
from __future__ import annotations
from dataclasses import dataclass
import dataclasses
from typing import Optional
import time
import pyglet
import logging
import numpy as np
# from deep_sort_realtime.deepsort_tracker import DeepSort
from sort import Sort
from collections import defaultdict
import math



Interval = float # seconds

logger = logging.getLogger('moonwalk')

colorset = [(255, 0, 0),
 (0, 0, 255),
 (0, 255, 0),
 (0, 255, 255),
 (255, 0, 0),
 (255, 0, 255),
 (255, 255, 0)
 ]


VELOCITY_FACTOR = 30*.05 # a pixels/second velocity of 400 gives a frame-duration of .05

images_nrs = range(125, 162)
angles = [0,15,30,45]

walk_animations = {}
for a in angles:
    walk_images = []
    for nr in images_nrs:
        # smaller textures doesn't seem to matter much in speed
        fn = f'walk-animation/OUT/frame{nr:03d}_{a:02d}.png'
        pic: pyglet.image.ImageData = pyglet.image.load(fn)
        walk_images.append(pic)
    walk_animations[a] = pyglet.image.animation.Animation.from_image_sequence(walk_images, 1, True)

walk_animation = walk_animations[0]

# walk_image = pyglet.image.load('walk_bw.jpg') # TODO investigaet pyglet.resource: https://pyglet.readthedocs.io/en/latest/programming_guide/image.html#displaying-images
# image_grid = pyglet.image.ImageGrid(walk_image, rows=1, columns=4)
# # texture_grid = image_grid.get_texture_sequence()
# walk_animation = image_grid.get_animation(period=.05)
walk_animation_dim = {
    'y': walk_animation.get_max_height(),
    'x': walk_animation.get_max_width()
}


@dataclass
class Params:
    visual_variability: float = 0
    # video_fps: float = 60
    tracker_fps: float = 5.6
    # iou = None
    emitter_speed: float = 1.2 # objects per second
    object_velocity: float = 100 # pixels/second
    velocity_decay: float = 1 # make system a bit springy
    iou_threshold: float = 0 # SORT Intersection over union
    movement_angle: float = 0 # angle at which the blobs move (-45 - +45 degrees)

    def add_listener(self, attr: str, callback: callable):
        """
        Add a listener for a change, callback gets
        called with (attribute_name, old_value, new_value)
        """
        if not attr in self.get_listeners():
            self._listeners[attr] = []
        self._listeners[attr].append(callback)
    
    def get_listeners(self) -> dict[str, callable]:
        if not hasattr(self, "_listeners"):
            self._listeners = {}
        return self._listeners

    def __setattr__(self, attr, val):
        if attr == '_listeners':
            super().__setattr__(attr, val)
            return
    
        if attr == 'tracker_fps' and val < .1:
            # limit tracker fps
            val = .1
    
        if attr == 'emitter_speed' and val < .1:
            # limit tracker fps
            val = .1
        
        old_val = getattr(self, attr)
        super().__setattr__(attr, val)
        if attr in self.get_listeners():
            for listener in self._listeners[attr]:
                listener(attr, old_val, val)



        
class DetectedObject:
    def __init__(self, canvas: Canvas):
        self.canvas = canvas
        # TODO handle variability
        self.v = self.canvas.params.object_velocity
        self.w = 160 # width
        self.h = 320 # height
        self.l = -self.w # left
        self.t = (self.canvas.height - self.h)/2 - math.sin(self.canvas.params.movement_angle/360*math.pi*2) * self.canvas.width/2 # top

        # self.shape = pyglet.shapes.Rectangle(self.l, self.t, self.w, self.h, color=(255, 22, 20), batch=self.canvas.batch_figures)

        angle = min(angles, key=lambda x:abs(x-self.canvas.params.movement_angle))
        
        self.shape = pyglet.sprite.Sprite(img=walk_animations[angle],x=self.l, y=self.t, batch=self.canvas.batch_figures)
        self.shape.scale_x = self.w/walk_animation_dim['x']
        self.shape.scale_y = self.h/walk_animation_dim['y']
        # rectangle.opacity = 128
        # rectangle.rotation = 33
        #TODO renderer

    def update(self, dt: Interval):
        """
        Update position
        """
        self.l += dt * self.canvas.params.object_velocity * math.cos(self.canvas.params.movement_angle/360*2 * math.pi) 
        self.t += dt * self.canvas.params.object_velocity * math.sin(self.canvas.params.movement_angle/360*2 * math.pi)
        self.shape.x = self.l
        self.shape.y = self.t
        # TODO exponential decay with self.params.velocity_decay
 
class ObjectEmitter:
    """
    Emit detectable objects
    """
    def __init__(self, params: Params, canvas: Canvas):
        self.lastEmit = 0
        self.params = params
        self.canvas = canvas

    def emit(self, dt: Interval) -> list[DetectedObject]:
        self.lastEmit += dt
        if self.lastEmit is None or self.lastEmit >= 1/self.params.emitter_speed:
            logger.info('emit!')
            obj = DetectedObject(self.canvas)
            self.lastEmit = 0
            return [obj]
        return []       
    
class MissingDict[T](dict):
    """
    collections.defaultdict does not accept arguments, but this is what we want/need.
    This implementation should do the trick
    """
    def __init__(self, factory: callable, values={}):
        self.update(values)
        self.factory = factory

    def __missing__(self, key) -> T:
        self[key] = self.factory(key)
        return self[key]



class Canvas:
    """
    A canvas with moving objects
    """
    def __init__(self, params: Params):
        self.width = 1920
        self.height = 1080
        self.objects: list[DetectedObject] = []
        self.lastSnapshot: Optional[float] = None
        self.params = params
        self.emitter = ObjectEmitter(self.params, self)
        

        self.hide_stats = False
        
        config = pyglet.gl.Config(sample_buffers=1, samples=4, double_buffer=True)
        # , fullscreen=self.config.render_window
        self.window = pyglet.window.Window(width=self.width, height=self.height, config=config, fullscreen=False)
        self.window.set_handler('on_draw', self.on_draw)
        self.window.set_handler('on_key_press', self.on_key_press)
        self.window.set_handler('on_mouse_scroll', self.on_mouse_scroll)

        self.window.set_handler('on_refresh', self.on_refresh)

        # self.window.set_handler('on_refresh', self.on_refresh)
        # self.window.set_handler('on_close', self.on_close)

        # Purple background color:
        pyglet.gl.glClearColor(230/255,230/255,230/255,230/255)
        self.draw_stats = True
        self.fps_display = pyglet.window.FPSDisplay(window=self.window, color=(255,255,255,255), samples=100)
        self.fps_display.label.x = self.window.width - 150
        self.fps_display.label.y = self.window.height - 17
        self.fps_display.label.bold = False
        self.fps_display.label.font_size = 10
        
        self.batch_figures = pyglet.graphics.Batch()
        self.batch_bounding_boxes = pyglet.graphics.Batch()
        self.batch_info = pyglet.graphics.Batch()

        self.tracks = []

        self.labels = {
            'objects': pyglet.text.Label("", x=20, y=30, color=(255,255,255,255), batch=self.batch_info),
            'tracks': pyglet.text.Label("", x=120, y=30, color=(255,255,255,255), batch=self.batch_info),
        }
        for i, field in enumerate(dataclasses.fields(self.params)):
            self.labels[field.name] = pyglet.text.Label(f"{field.name}: {field.default}", x=20, y=30 + 20*(i+1), color=(255,255,255,255), batch=self.batch_info)
        
        # TODO: Add a number next to the box using pyglet.graphics.Group()
        self.track_shapes: MissingDict[np.float64, tuple[pyglet.shapes.Box, pyglet.text.Label]] = MissingDict(self.getTrackBboxShapes)

        self.tracker = Sort(max_age=5, min_hits=1, iou_threshold=0) #DeepSort(max_age=5)

        pyglet.clock.schedule_interval(self.on_track, 1/self.params.tracker_fps)

        self.interval_items: list[pyglet.clock._ScheduledIntervalItem] = [i for i in pyglet.clock._default._schedule_interval_items if i.func == self.on_track]
        self.params.add_listener('tracker_fps', self.on_tracker_fps_change)
        self.params.add_listener('iou_threshold', self.on_iou_threshold_change)
        self.params.add_listener('object_velocity', self.on_object_velocity_change)
        self.params.add_listener('movement_angle', self.on_movement_angle_change)
        # trigger first time around
        self.on_object_velocity_change('object_velocity', None, self.params.object_velocity)
    
    def getTrackBboxShapes(self, track_id) -> tuple[pyglet.shapes.Box, pyglet.text.Label]:
        color = colorset[int(track_id) % len(colorset)]
        return (
            pyglet.shapes.Box(0,0,0,0,color=color,thickness=2, batch=self.batch_bounding_boxes),
            pyglet.text.Label(f"{track_id:.0f}", x=0, y=0, anchor_y='top', color=color, batch=self.batch_bounding_boxes),
            # pyglet.shapes.Lab(0,0,0,0,color=(0,255,0),thickness=2, batch=self.batch_bounding_boxes)
        )
        
            
    def on_tracker_fps_change(self, field, old_value, new_value):
        """
        Param dataclass listener for changes to tracker_fps
        """
        for ii in self.interval_items:
            ii.interval = 1/new_value
            logger.debug(f"Set tracker interval to {ii.interval}")
    
    def on_object_velocity_change(self, field, old_value, new_value):
        """
        Param dataclass listener for changes to object_velocity as to change walk animation
        """
        duration = max(.001, VELOCITY_FACTOR / new_value)
        logger.debug(f"set frame duration to {duration=} (for {new_value} p/s, factor: {VELOCITY_FACTOR})")
        for anim in walk_animations.values():
            for frame in anim.frames:
                    # a velocity of 
                    frame.duration = duration

        # for ii in self.interval_items:
        #     ii.interval = 1/new_value
            # logger.debug(f"Set tracker interval to {ii.interval}")
    
    def on_movement_angle_change(self, field, old_value, new_value):
        """
        Param dataclass listener for changes to movement-angle
        """
        angle = min(angles, key=lambda x:abs(x-new_value))

        if hasattr(self, 'angle') and self.angle == angle:
            return
        
        self.angle = angle

        for obj in self.objects:
            idx = obj.shape._frame_index
            obj.shape.image = walk_animations[self.angle]
            obj.shape._frame_index = idx
    
    def on_iou_threshold_change(self, field, old_value, new_value):
        """
        Param dataclass listener for changes to iou_threshold
        """
        self.tracker.iou_threshold = new_value

    def run(self):

        self.event_loop = pyglet.app.EventLoop()
        # pyglet.clock.schedule_interval(self.check_running, 0.1)
        # pyglet.clock.schedule(self.check_frames)
        # pyglet.clock.schedule(self.track)
        self.event_loop.run()

    def on_draw(self):
        self.window.clear()

        self.batch_figures.draw()
        self.batch_bounding_boxes.draw()
        self.batch_info.draw()
        if self.draw_stats:
            self.fps_display.draw()

    def on_close(self):
        logger.info('closing')
        pass


    def on_key_press(self, symbol, modifiers):
        if symbol == pyglet.window.key.Q or symbol == pyglet.window.key.ESCAPE:
            self.window.close()
            exit()
        
        if symbol == pyglet.window.key.V:
            level = logging.INFO if logger.getEffectiveLevel() == logging.DEBUG else logging.DEBUG
            logger.setLevel(level)
            logger.info(f"set log level: {level}")
        if symbol == pyglet.window.key.S:
            self.draw_stats = not self.draw_stats
            logger.info(f"rendering stats: {self.draw_stats}")
        if symbol == pyglet.window.key.UP:
            logger.debug('up')
            self.params.object_velocity += (10 if pyglet.window.key.MOD_SHIFT & modifiers else 1)
        if symbol == pyglet.window.key.DOWN:
            logger.debug('down')
            self.params.object_velocity -= (10 if pyglet.window.key.MOD_SHIFT & modifiers else 1)

    def on_mouse_scroll(self, x, y, scroll_x, scroll_y):
        # determine xy position to select var to change,
        # then change according to scroll_y
        for param_name, param_value in dataclasses.asdict(self.params).items():
            if x >= self.labels[param_name].x and \
                x <= (self.labels[param_name].x + self.labels[param_name].content_width) and \
                y >= self.labels[param_name].y and \
                y <= (self.labels[param_name].y + self.labels[param_name].content_height):
                    if param_value == 0:
                        if scroll_y < 1:
                            param_value -= .1
                        else:
                            param_value += .1

                    if scroll_y < 0:
                        param_value /= (-scroll_y * 1.25)
                    else:
                        param_value *= (scroll_y * 1.25)
                    setattr(self.params, param_name, param_value)


    def on_track(self, dt):
        # bbs = object_detector.detect(frame) 
        objects = self.snapshot()
        # TODO chipper & embedder
        # bbs = [([o.l, o.t, o.w, o.h], 1, 1) for o in objects]
        # DEEP SORT: self.tracks = self.tracker.update_tracks(bbs, frame=np.zeros([1280,720])) # bbs expected to be a list of detections, each in tuples of ( [left,top,w,h], confidence, detection_class )
        
        bbs = np.array([[o.l, o.t, o.l+ o.w, o.t+o.h, 1, 1] for o in objects])
        self.tracks = self.tracker.update(bbs) # a numpy array of detections in the format [[x1,y1,x2,y2,score],[x1,y1,x2,y2,score],...]

        # self.tracks is a np array where each row contains a valid bounding box and track_id (last column)

        # remove old shapes
        ids = [track[4] for track in self.tracks]
        for k in list(self.track_shapes.keys()):
            if k not in ids:
                self.track_shapes.pop(k)
                logger.debug(f"shape removed {k}" )
        # print([t[4] for t in self.tracks])

    
    def set_tracker_fps(self, fps: float):
        self.params.tracker_fps = fps
        for interval in self.interval_items:
            interval.interval = 1/fps


    def prune(self):
        """
        Loop over objects remove those out of the frame
        """
        for i, object in enumerate(self.objects.copy()):
            if object.l > self.width:
                logging.info(f'Delete {i}')
                el = self.objects.pop(i)
                el.shape.delete() # clear the attached shape
        

    def snapshot(self) -> list[DetectedObject]:
        """
        Update all object positions base on dt = now - lastSnapshot
        """
        now = time.monotonic()
        if self.lastSnapshot is None:
            self.lastSnapshot = now

        dt = now - self.lastSnapshot

        self.objects.extend(self.emitter.emit(dt))
        for object in self.objects:
            object.update(dt)
        self.prune()

        self.lastSnapshot = now
        # print(f"duration: {time.monotonic()-now)
        return self.objects
    

    
    def on_refresh(self, dt: float):
        objects = self.snapshot()

        self.labels['objects'].text = f"Objects: {len(objects)}"
        self.labels['tracks'].text = f"Tracks: {len(self.tracks)}"
        # self.labels['velocity'].text = f"Velocity: {self.params.object_velocity}"
        # self.labels['tracker_fps'].text = f"Tracker FPS: {self.params.tracker_fps}"

        for name, value in dataclasses.asdict(self.params).items():
            self.labels[name].text = f"{name}: {value}"

        # BIG impact 
        for track in self.tracks:
            nr = track[4]
            box: pyglet.shape.Box = self.track_shapes[nr][0]
            label: pyglet.text.Label = self.track_shapes[nr][1]
            if box.x == 0 and box.y == 0:
                # set initial position
                box.x = track[0]
                box.y = track[1]
                label.x = track[0]
                label.y = track[1]
                box.width = track[2] - track[0]
                box.height = track[3] - track[1]
            else: # not really impact
                #  exponential decay
                label.x = box.x = exponentialDecay(box.x, track[0], 12, dt)
                label.y = box.y = exponentialDecay(box.y, track[1], 12, dt)
                # setting width and height on label is not needed _and_ makes it super slow
                box.width = exponentialDecay(box.width, track[2] - track[0], 12, dt)
                box.height = exponentialDecay(box.height, track[3] - track[1], 12, dt)

            
        # TODO: shape in DetectedObject 
        # rectangle = shapes.Rectangle(250, 300, 400, 200, color=(255, 22, 20), batch=batch)
        # rectangle.opacity = 128
        # rectangle.rotation = 33
        # print(objects)
        # id(objects)


def exponentialDecay(a, b, decay, dt):
    """Exponential decay as alternative to Lerp
    Introduced by Freya Holmér: https://www.youtube.com/watch?v=LSNQuFEDOyQ
    """
    return b + (a-b) * math.exp(-decay * dt)

        
        
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    params = Params()
    canvas = Canvas(params)
    canvas.run()