#!/usr/bin/env python2

import os
import yaml
import utm

import cv2
import numpy as np

import rospy

import std_msgs
from std_msgs.msg import Float64, Bool
from sensor_msgs.msg import Imu
from geometry_msgs.msg import TwistWithCovarianceStamped, PoseWithCovariance, TwistWithCovariance, Pose, Point, Quaternion, Twist, Vector3
from rosflight_msgs.msg import Command, RCRaw, Status
from nav_msgs.msg import Odometry, GridCells

# NED frame: X is North, Y east, Z down.
# Latitude: NS, Longitude: WE


def roundi(v):
    return int(round(v))


def ll2ned(ll):
    u = utm.from_latlon(ll[0], ll[1])
    return (u[1], u[0], 0)


def scale_keepout(img, scale, target):
    r = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
                                  (roundi(1 / scale / 2 / target), roundi(1 / scale / 2 / target)))
    r = cv2.dilate(r, k)
    r = cv2.resize(r,
                   (roundi(img.shape[0] * scale / target), roundi(img.shape[1] * scale / target)))
    _, r = cv2.threshold(r, 128, 255, cv2.THRESH_BINARY)

    k = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    d = cv2.dilate(r, k)
    r = cv2.max(r, d // 2)
    r[:, 0] = 255
    r[:, -1] = 255
    r[0, :] = 255
    r[-1, :] = 255

    return r


class Map:
    def __init__(self, base, spec, scale=1):
        self._spec = spec
        self._display = cv2.imread(os.path.join(base, spec['display']))

        self._calib = np.array(spec['calibration'])
        self._ll_base = np.min(self._calib[:, 0:2], 0)
        self._ned_base = np.array(ll2ned(self._ll_base))

        ll_max = np.max(self._calib[:, 0:2], 0)
        ned_max = np.array(ll2ned(ll_max))
        self._ned_scale = (ned_max - self._ned_base)[:2] / (ll_max - self._ll_base)

        keepout = cv2.imread(os.path.join(base, spec['keepout']))

        lim = np.zeros((2, 2))
        px = np.polyfit(self._calib[:, 2], self._calib[:, 0], 1)
        lim[:, 0] = np.polyval(px, (0, keepout.shape[0]))
        py = np.polyfit(self._calib[:, 3], self._calib[:, 1], 1)
        lim[:, 1] = np.polyval(py, (0, keepout.shape[1]))
        lim -= self._ll_base
        lim *= self._ned_scale

        self._size = np.max(lim, 0) - np.min(lim, 0)
        self._resolution = self._size / keepout.shape[:2]

        self._scale = np.mean(self._resolution)
        self._keepout = scale_keepout(keepout, np.mean(self._scale), 1)

    def run(self):
        rospy.init_node('revosim', anonymous=True)

        self._keepout_pub = rospy.Publisher('keepout', GridCells, queue_size=10)

        r = rospy.Rate(1)

        while not rospy.is_shutdown():
            h = std_msgs.msg.Header(stamp=rospy.Time.now(), frame_id='map')
            c = []
            s = self._keepout.shape
            for y, row in enumerate(self._keepout):
                for x, v in enumerate(row):
                    if v >= 200:
                        c.append(Point(x=x, y=s[0] - y - 1))
            g = GridCells(header=h, cell_width=1, cell_height=1, cells=c)
            self._keepout_pub.publish(g)
            r.sleep()


def main():
    base = '/home/michaelh/projects/nppilot-v2/src/nppilot/'
    with open(base + 'resources/site1/site1.yaml') as f:
        spec = yaml.load(f)
    m = Map(base + 'resources/site1', spec)
    m.run()


if __name__ == '__main__':
    main()