import asyncio
import base64
import datetime as dt
import glob
import json
import logging
import math
import os
from abc import abstractmethod
from dataclasses import dataclass
from typing import List, Tuple

import ffmpeg
import numpy as np
import requests
import websockets
from cv2 import cv2
from shapely.geometry import Point, Polygon
from sqlalchemy import orm, Column, Integer, SmallInteger, VARCHAR, ForeignKey, Float, Boolean, String
from sqlalchemy.dialects.mysql import DATETIME, TEXT
from sqlalchemy.orm import relationship, backref

from server.database import Base, db_session
from server.instance.config import PROCESSORS_PREVIEW_DIR, FACE_DETECTOR_URL, FACE_WS_ADDRESS, FACE_WS_PORT

DT_FORMAT = '%Y-%m-%d %H:%M:%S.%f'


@dataclass
class DetectedObject:
    """
    Storage class for detected object infos.

    """
    # relative to w or h of a frame values
    x_min: float
    y_min: float
    x_max: float
    y_max: float
    prob: float
    cls: int = -1

    def x_center(self) -> float:
        return (self.x_min + self.x_max) / 2

    def y_center(self) -> float:
        return (self.y_min + self.y_max) / 2

    @property
    def h(self) -> float:
        return self.y_max - self.y_min

    @property
    def w(self) -> float:
        return self.x_max - self.x_min

    def point(self, x_offset: int = 0, y_offset: int = 0) -> Point:
        """
        Returns a Point inside the object. (0,0) gives a center; (-1, 0) gives x_min, y_center; etc
        :param x_offset: [-1, 1]
        :param y_offset: [-1, 1]
        :return:
        """
        if -1 <= x_offset <= 1 and -1 <= y_offset <= 1:
            x = self.x_min + self.w / 2 * (1 + x_offset)
            y = self.y_min + self.h / 2 * (1 + y_offset)
            return Point(x, y)
        else:
            logging.error('Offsets are out of [-1,1] {} {}'.format(x_offset, y_offset))
            raise ValueError


class Frame:
    """
    Contains image with timestamp and list of detected objects on it (to be filled after detection)
    """

    def __init__(self, img: np.ndarray, ts: dt.datetime):
        self.image = img
        self.ts = ts
        self.objects: List[DetectedObject] = []


class Track:
    """
    Used for track counter only. It stores sequence of some unique object images as list of DetectedObjects.
    Allows to predict next position of an object based on his history.
    """

    def __init__(self, first_obj: DetectedObject, ts: dt.datetime) -> None:
        self._objs = [first_obj]
        self.last_frame_ts = ts

    @property
    def objs(self) -> List[DetectedObject]:
        return self._objs

    @property
    def length(self) -> int:
        return len(self._objs)

    def add_obj(self, obj: DetectedObject, ts: dt.datetime):
        self._objs.append(obj)
        self.last_frame_ts = ts

    def last_obj(self) -> DetectedObject:
        return self._objs[-1]

    def best_prediction(self) -> Point:
        if self.length > 1:
            x = 2 * self._objs[-1].point().x - self._objs[-2].point().x
            y = 2 * self._objs[-1].point().y - self._objs[-2].point().y
            prediction = Point(x, y)
        else:
            prediction = self.last_obj().point()
        return prediction

    def drop_old_objs(self, max_track_size: int):
        self._objs = self._objs[max_track_size * -1:]


class Scene:
    """
    Used for track counter only. Stores all found Tracks.
    """
    tracks: List[Track]

    def __init__(self) -> None:
        self.tracks = []


class Camera(Base):
    """
    Camera class. Connects camera info on connection and processing
    """
    __tablename__ = 'camera'
    id = Column(Integer, primary_key=True)
    watch_fps = Column(Integer, default=1)  # #FPS to read from camera
    watch_rows = Column(Integer, default=1)  # rows to combine frames for detection
    watch_cols = Column(Integer, default=1)  # cols to combine frames for detection
    tz = Column(SmallInteger)  # camera time zone offset (for storage of results)
    stream_url = Column(VARCHAR(length=100), unique=True, nullable=False)  # connection URL

    @property
    def grid_size(self) -> [int, int]:
        """
        Mosaic grid size to be passed for object detector
        :return:
        """
        return self.watch_rows, self.watch_cols


class ProcessorEvent(Base):
    """
    Processing result class. It is generated by Processors. For each of them this record means different things.
    For example for ObjectCounter value = number of objects in frame at this ts;
    for traffic counter value = number of objects left the zone at this ts
    """
    __tablename__ = 'processor_event'
    id = Column(Integer, primary_key=True)
    processor_id = Column(Integer, ForeignKey('processor.id'), nullable=False)  # processor id
    processor = relationship('Processor')
    ts = Column(DATETIME(fsp=6), nullable=False)  # timestamp of event detection
    value = Column(SmallInteger)  # detected value


class Processor(Base):
    """
    Base class. Must not be barely initialized. Only used as an abstract class for ORM.
    """
    __tablename__ = 'processor'
    id = Column(Integer, primary_key=True)
    camera_id = Column(Integer, ForeignKey('camera.id'), nullable=False)
    camera = relationship('Camera', backref=backref('processors', lazy=True))
    zones_str = Column(TEXT, nullable=False, default='[]')  # detection zones as a list of lists of polygon vertices
    threshold = Column(Float, default=0)  # minimal probability value to count result
    enabled = Column(Boolean, default=False)  # true to enable it; false to skip
    output_hls = Column(Boolean, default=False)  # true to save output m3u8 hlp

    type = Column(String(50), nullable=False)
    __mapper_args__ = {
        'polymorphic_identity': 'processor',
        'polymorphic_on': type
    }

    def __init__(self, **kwargs) -> None:
        super().__init__(kwargs)
        self.video_builder = None

    @orm.reconstructor
    def init_on_load(self):
        self.video_builder = None

    @abstractmethod
    def process(self, frames: List[Frame]):
        pass

    @property
    def zones(self) -> List[List[List]]:
        # parse zones from config string
        return json.loads(self.zones_str)

    @property
    def polygons(self) -> List[Polygon]:
        # make polygons from zones
        return [Polygon(zone) for zone in self.zones]

    def zones_mask(self, h: int, w: int) -> np.ndarray:
        base = np.zeros((h, w), dtype=np.uint8)
        scaled_zones = []
        for zone in self.zones:
            scaled_zones.append([(p[0] * w, p[1] * h) for p in zone])
        return cv2.fillPoly(base, np.array(scaled_zones).astype(int), (255, 255, 255))

    def update_video_builder(self, h: int, w: int):
        # enable or disable video builder
        if self.output_hls:

            if self.video_builder is None:
                file_mask = 'processed_stream'
                out_dir = os.path.join(PROCESSORS_PREVIEW_DIR, str(self.id))

                # prepare output folder
                os.makedirs(out_dir, exist_ok=True)
                fileList = glob.glob(os.path.join(out_dir, file_mask + '*'))
                for filePath in fileList:
                    try:
                        os.remove(filePath)
                    except Exception as e:
                        logging.warning("Error while cleaning output folder:", filePath)
                        logging.warning(str(e))
                # init video builder
                manifest_fp = os.path.join(out_dir, file_mask + '.m3u8')
                self.video_builder = (ffmpeg.input('pipe:', format='rawvideo', pix_fmt='bgr24', s='{}x{}'.format(w, h),
                                                   framerate=self.camera.watch_fps)
                                      .output(manifest_fp, **{'c:v': 'libx264', 'b:v': 880000}, pix_fmt='yuv420p',
                                              f='hls', hls_time=10, hls_list_size=3, start_number=1,
                                              hls_flags='delete_segments+program_date_time')
                                      .global_args('-loglevel', 'error')
                                      .overwrite_output()
                                      .run_async(pipe_stdin=True))
        else:
            if self.video_builder:
                self.video_builder.stdin.close()
                self.video_builder.wait()
                self.video_builder = None

    @staticmethod
    def draw_zones(img: np.ndarray, zones_mask: np.ndarray, color: Tuple[int, int, int] = (66, 183, 42)):
        # draw zones
        base = np.zeros(img.shape, img.dtype)
        base[:, :] = color
        mask = cv2.bitwise_and(base, base, mask=zones_mask)
        cv2.addWeighted(mask, 1, img, 1, 0, img)


class TrafficCounter(Processor):
    """
    Traffic counting processor. Takes series of frames and tries to find tracks finishing in one of zones.
    Detection is made by process method. If a new list of frames will be passed,
    it will try to append new objects to the existing tracks.
    """
    __tablename__ = 'traffic_counter'
    id = Column(Integer, ForeignKey('processor.id'), primary_key=True)

    def __init__(self, camera_id: int, zones_str: str, threshold: float):
        super().__init__(camera_id=camera_id, zones_str=zones_str, threshold=threshold)
        self.scene = Scene()  # storage of existing tracks

    @orm.reconstructor
    def init_on_load(self):
        super().init_on_load()
        self.scene = Scene()

    __mapper_args__ = {
        'polymorphic_identity': 'traffic',
    }

    def process(self, frames: List[Frame]):
        """
        Assuming that frames is not empty and sorted by ts.
        Will count number ob track entered any of selected zone
        :param frames:
        :return:
        """
        # prepare video builder
        h, w, _ = frames[0].image.shape
        self.update_video_builder(h, w)
        # generate zones mask if video builder is required
        if self.video_builder:
            zones_mask = self.zones_mask(h, w)
        else:
            zones_mask = None

        logging.debug('Scene tracks: {}'.format(len(self.scene.tracks)))
        min_dt = frames[0].ts.strftime(DT_FORMAT)
        max_dt = frames[-1].ts.strftime(DT_FORMAT)
        logging.debug('TS range: {} - {}'.format(min_dt, max_dt))

        # calculation
        finished_tracks_counter = 0
        for frame in frames:
            finished_tracks = self._analyze_frame(frame)
            if finished_tracks:
                finished_tracks_counter += len(finished_tracks)
                e = ProcessorEvent(processor_id=self.id, ts=frame.ts, value=len(finished_tracks))
                db_session.add(e)
            if self.video_builder:
                frame = self._visualize(frame.image, finished_tracks_counter, finished_tracks, zones_mask)
                self.video_builder.stdin.write(frame.astype(np.uint8).tobytes())
        db_session.commit()

    def _analyze_frame(self, frame: Frame, x_weight=1.0, y_weight=1.0, max_frames_gap=5, max_next_point_dst=0.1,
                       min_track_size=3,
                       max_track_size=10) -> List[Track]:
        """
        Map objects to tracks and returns finished tracks
        """
        new_objects = frame.objects.copy()
        logging.debug('New objects: {}'.format(len(new_objects)))
        # if no objects were found then just leave this procedure
        if not new_objects:
            return []

        # if not try to extend each track at one object; unused objects create new tracks
        for track in self.scene.tracks:
            track_obj_point = track.best_prediction()
            new_objects.sort(key=lambda o: distance(o.point(), track_obj_point, x_weight, y_weight))
            next_obj = new_objects[0]
            if distance(next_obj.point(), track_obj_point, x_weight, y_weight) < max_next_point_dst:
                track.add_obj(next_obj, frame.ts)
                new_objects.remove(next_obj)
                logging.debug('Appended object to track')
            if not new_objects:
                break

        # delete old tracks (that were not updated for max_frames_gap seconds)
        self.scene.tracks = [t for t in self.scene.tracks if
                             frame.ts - t.last_frame_ts < dt.timedelta(seconds=max_frames_gap)]
        logging.debug('Scene tracks (alive): {}'.format(len(self.scene.tracks)))
        # find finished tracks of objects
        finished_tracks = []
        for i, track in enumerate(self.scene.tracks):
            if at_roe(track.last_obj(), self.polygons) and track.length >= min_track_size:
                finished_tracks.append(track)
                self.scene.tracks.remove(track)

        # leave only last max_track_size points in each track
        for track in self.scene.tracks:
            track.drop_old_objs(max_track_size)

        # create new tracks from unused objects if they are not in roe
        for obj in new_objects:
            if not at_roe(obj, self.polygons):
                self.scene.tracks.append(Track(obj, frame.ts))
        logging.debug('Scene tracks (with new): {}'.format(len(self.scene.tracks)))
        logging.debug('Finished tracks: {}'.format(len(finished_tracks)))
        return finished_tracks

    def _visualize(self, frame: np.ndarray, found: int, finished_tracks: List[Track],
                   zones_mask: np.ndarray) -> np.ndarray:
        """
        Draws scene tracks at frame. Takes non finished tracks from scene and finished from argument.
        Changes frame!
        :param finished_tracks:
        :return:
        """
        img = frame.copy()
        h, w, _ = img.shape
        color_active_track = (0, 0, 255)
        color_finish_track = (0, 255, 0)
        color_text = (0, 0, 0)
        draw_tasks = [{'tracks': self.scene.tracks, 'color': color_active_track},
                      {'tracks': finished_tracks, 'color': color_finish_track}]

        # draw detection zones
        self.draw_zones(img, zones_mask)

        # draw tracks each with own color
        for task in draw_tasks:
            tracks = task['tracks']
            color = task['color']
            for track in tracks:
                points = [(int(obj.point().x * w), int(obj.point().y * h)) for obj in track.objs]
                cv2.polylines(img, [np.int32(points)], False, color)
                for obj in track.objs:
                    cv2.circle(img, (int(obj.point().x * w), int(obj.point().y * h)), 2, color)
                obj = track.last_obj()
                p1 = (int(round(obj.x_min * w)), int(round(obj.y_min * h)))
                p2 = (int(round(obj.x_max * w)), int(round(obj.y_max * h)))
                cv2.rectangle(img, p1, p2, color, 3)
                cv2.putText(img, '{:.2f} %'.format(obj.prob * 100), (p1[0], p1[1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 1.0,
                            color_text, lineType=cv2.LINE_AA)

        # draw number of found objects
        text = 'Found: {}'.format(found)
        cv2.putText(img, text, (20, 30), cv2.FONT_HERSHEY_SIMPLEX, 1.0, color_text)
        return img


class ObjectsCounter(Processor):
    """
    Counts objects at the frame.
    Saves result to DB only for the first frame in sequence and for the frame with different number of objects.
    Counts only in zone objects.
    """
    __tablename__ = 'objects_counter'
    id = Column(Integer, ForeignKey('processor.id'), primary_key=True)

    __mapper_args__ = {
        'polymorphic_identity': 'object',
    }

    def __init__(self, camera_id: int, zones_str: str, threshold: float):
        super().__init__(camera_id=camera_id, zones_str=zones_str, threshold=threshold)

    @orm.reconstructor
    def init_on_load(self):
        super().init_on_load()

    def process(self, frames: List[Frame]):
        """
        Assuming that frames is not empty and sorted by ts.
        Will count number of detected objects in each frame and write its value to DB when changed
        :param frames:
        :return:
        """
        # prepare video builder
        h, w, _ = frames[0].image.shape
        self.update_video_builder(h, w)
        # generate zones mask if video builder is required
        if self.video_builder:
            zones_mask = self.zones_mask(h, w)
        else:
            zones_mask = None
        # first frame always goes to DB
        prev_val = sum([at_roe(obj, self.polygons) for obj in frames[0].objects])
        e = ProcessorEvent(processor_id=self.id, ts=frames[0].ts, value=prev_val)
        db_session.add(e)
        for frame_id, frame in enumerate(frames):
            objs_in_zone = [obj for obj in frame.objects if at_roe(obj, self.polygons)]
            cur_val = len(objs_in_zone)
            if prev_val != cur_val:
                e = ProcessorEvent(processor_id=self.id, ts=frame.ts, value=cur_val)
                db_session.add(e)
                prev_val = cur_val
            if self.video_builder:
                frame_img = self._visualize(frame, objs_in_zone, zones_mask)
                self.video_builder.stdin.write(frame_img.astype(np.uint8).tobytes())
        db_session.commit()

    def _visualize(self, frame: Frame, objs_in_zone: List[DetectedObject], zones_mask: np.ndarray) -> np.ndarray:
        """
        Draws objects at frame.
        Does not change frame.image
        :return:
        """
        color_not_in_zone = (0, 0, 255)
        color_in_zone = (0, 255, 0)
        color_text = (0, 0, 0)

        img = frame.image.copy()
        h, w, _ = img.shape
        # draw detection zones
        self.draw_zones(img, zones_mask)
        # draw detected objects
        for obj in frame.objects:
            p1 = (int(round(obj.x_min * w)), int(round(obj.y_min * h)))
            p2 = (int(round(obj.x_max * w)), int(round(obj.y_max * h)))
            if obj in objs_in_zone:
                color = color_in_zone
            else:
                color = color_not_in_zone
            cv2.rectangle(img, p1, p2, color, 3)
            cv2.putText(img, '{:.2f} %'.format(obj.prob * 100), (p1[0], p1[1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 1.0,
                        color_text, lineType=cv2.LINE_AA)
        # draw number of found objects
        text = 'Found: {}'.format(len(objs_in_zone))
        cv2.putText(img, text, (20, 30), cv2.FONT_HERSHEY_SIMPLEX, 1.0, color_text)
        return img


class FaceDetector(Processor):
    """
    Detects faces of in zone objects. Scans only upper left square of a detected object.
    Sends faces to FACE_WS server.
    """
    __tablename__ = 'face_detector'
    id = Column(Integer, ForeignKey('processor.id'), primary_key=True)
    __mapper_args__ = {
        'polymorphic_identity': 'face',
    }

    def __init__(self, camera_id: int, zones_str: str, threshold: float):
        super().__init__(camera_id=camera_id, zones_str=zones_str, threshold=threshold)
        self.container = 'jpg'

    @orm.reconstructor
    def init_on_load(self):
        super().init_on_load()
        self.container = 'jpg'

    def process(self, frames: List[Frame]):
        # todo: scan not only top left square but the whole image anf if no face => several squares
        #  OR create some better NN
        # prepare video builder
        h, w, _ = frames[0].image.shape
        self.update_video_builder(h, w)
        # generate zones mask if video builder is required
        if self.video_builder:
            zones_mask = self.zones_mask(h, w)
        else:
            zones_mask = None
        content_type = 'image/jpeg'
        headers = {'content-type': content_type}
        for frame in frames:
            good_faces = []
            # scan upper square of each detected object to find a face
            frame_faces = []
            for obj in frame.objects:
                if at_roe(obj, self.polygons):
                    square_size = int(min(obj.w * w, obj.h * h))
                    square_y_min = int(obj.y_min * h)
                    square_y_max = square_y_min + square_size
                    square_x_min = int(obj.x_min * w)
                    square_x_max = square_x_min + square_size
                    upper_square = frame.image[square_y_min: square_y_max, square_x_min: square_x_max, :]
                    if upper_square.size > 0:
                        _, buf = cv2.imencode('.jpg', upper_square)
                        r = requests.post(FACE_DETECTOR_URL, data=buf.tostring(), headers=headers)
                        faces = json.loads(r.text)
                        boxes = faces['boxes']
                        conf = faces['conf']
                        if boxes:
                            for box, conf in zip(boxes, conf):
                                frame_faces.append(DetectedObject(obj.x_min + box[0] / w, obj.y_min + box[1] / h,
                                                                  obj.x_min + box[2] / w, obj.y_min + box[3] / h,
                                                                  conf))
                                if conf >= self.threshold:
                                    face = upper_square[int(box[1]):int(box[3]), int(box[0]):int(box[2]), :]
                                    if face.size > 0:
                                        _, face_buffer = cv2.imencode('.' + self.container, face)
                                        jpg_as_text = base64.b64encode(face_buffer).decode()
                                        good_faces.append({'shape': [int(v) for v in face.shape],
                                                           'bbox': [int(v) for v in box],
                                                           'conf': float(conf),
                                                           'img': jpg_as_text})
            # do not send empty faces
            if good_faces:
                data = {'frame_ts': frame.ts.strftime(DT_FORMAT),
                        'container': self.container,
                        'camera_id': self.camera_id,
                        'camera_url': self.camera.stream_url,
                        'faces': good_faces}
                asyncio.get_event_loop().run_until_complete(self._send_to_ws(json.dumps(data)))

            if self.video_builder:
                # draw faces with zones and save to video file
                frame_img = self._visualize(frame, frame_faces, zones_mask)
                self.video_builder.stdin.write(frame_img.astype(np.uint8).tobytes())

    def _visualize(self, frame: Frame, faces: List[DetectedObject], zones_mask: np.ndarray) -> np.ndarray:
        """
        Draws faces at frame.
        Does not change frame.image
        :return:
        """
        color_low_prob = (0, 0, 255)
        color_high_prob = (0, 255, 0)
        color_text = (0, 0, 0)

        img = frame.image.copy()
        h, w, _ = img.shape
        # draw detection zones
        self.draw_zones(img, zones_mask)
        # draw detected objects
        for face in faces:
            p1 = (int(round(face.x_min * w)), int(round(face.y_min * h)))
            p2 = (int(round(face.x_max * w)), int(round(face.y_max * h)))
            if face.prob >= self.threshold:
                color = color_high_prob
            else:
                color = color_low_prob
            cv2.rectangle(img, p1, p2, color, 3)
            cv2.putText(img, '{:.2f} %'.format(face.prob * 100), (p1[0], p1[1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 1.0,
                        color_text, lineType=cv2.LINE_AA)
        # draw number of found objects
        text = 'Found: {}'.format(len(faces))
        cv2.putText(img, text, (20, 30), cv2.FONT_HERSHEY_SIMPLEX, 1.0, color_text)
        return img

    @staticmethod
    async def _send_to_ws(data: str):
        address = FACE_WS_ADDRESS
        if address == '0.0.0.0':
            address = 'localhost'
        port = FACE_WS_PORT
        uri = "ws://{}:{}".format(address, port)
        try:
            async with websockets.connect(uri) as websocket:
                await websocket.send('detector')
                if await websocket.recv() == 'OK':
                    await websocket.send(data)
                else:
                    logging.error('WS server did not allow to send message')
                await websocket.close()
        except Exception as e:
            logging.error('Failed to send faces to WS from detector')
            logging.error(str(e))


def distance(p1: Point, p2: Point, x_weight=1.0, y_weight=1.0, metric='euclidean') -> float:
    """
    Returns weighted by x and y Euclidean distance.
    Weighting can be useful when perspective is known.

    :param p1:
    :param p2:
    :param x_weight:
    :param y_weight:
    :param metric:
    :return:
    """
    if metric == 'euclidean':
        return math.sqrt(float((p1.x - p2.x) ** 2) / x_weight + float((p1.y - p2.y) ** 2) / y_weight)


def at_roe(obj: DetectedObject, polygons: List[Polygon]) -> bool:
    """
    Checks if obj is in at least one polygon.
    Checks only one point (obj.point()) which is usually a center of an object.
    :param obj:
    :param polygons:
    :return:
    """
    if polygons:
        for pol in polygons:
            if pol.contains(obj.point()):
                return True
        return False
    else:
        return True
