import glob
import json
import postgresql as pg
from PIL import Image
from PIL import ImageDraw
from pathlib import Path
import numpy as np
from enum import Enum
from math import *
from libtiff import TIFF
from os import path
from shapely.geometry.polygon import Polygon
from shapely.geometry.point import Point
from shapely.geometry.multipoint import MultiPoint
from shapely.affinity import affine_transform
from functools import partial
from threading import Semaphore
from itertools import groupby
from postgresql.types.geometry import Box
from shapely.geometry import box
from itertools import repeat
import numpy.linalg as lin
from concurrent.futures import Future, ThreadPoolExecutor, ProcessPoolExecutor
from shapely.wkb import loads
from typing import List, Tuple
from collections import namedtuple
from argparse import ArgumentParser
from progressbar import Bar, ProgressBar
from pickle import dump, load
from functools import partial
from multiprocessing import Lock
import math
import cv2

BBox = namedtuple('BBox', ['detection_id', 'bbox', 'coverage'])


class StenCode(Enum):
    person = 1
    car = 2


def num_directories(tiffimage):
    count = 1
    TIFF.setdirectory(tiffimage, 0)
    while not TIFF.lastdirectory(tiffimage):
        TIFF.setdirectory(tiffimage, count)
        count += 1
    TIFF.setdirectory(tiffimage, 0)
    return count


def create_rot_matrix(euler: Point) -> np.matrix:
    x = np.radians(euler.x)
    y = np.radians(euler.y)
    z = np.radians(euler.z)
    Rx = np.array([[1, 0, 0],
                   [0, cos(x), -sin(x)],
                   [0, sin(x), cos(x)]], dtype=np.float)
    Ry = np.array([[cos(y), 0, sin(y)],
                   [0, 1, 0],
                   [-sin(y), 0, cos(y)]], dtype=np.float)
    Rz = np.array([[cos(z), -sin(z), 0],
                   [sin(z), cos(z), 0],
                   [0, 0, 1]], dtype=np.float)
    result = Rz @ Ry @ Rx
    return result


def create_affine(euler: Point, translation: Point) -> np.matrix:
    rotation = create_rot_matrix(euler)
    x = translation.x
    y = translation.y
    z = translation.z
    translation_mtx = np.array([[1, 0, 0, x],
                                [0, 1, 0, y],
                                [0, 0, 1, z],
                                [0, 0, 0, 1]], dtype=np.float64)
    rot_mtx = np.hstack((rotation, [[0], [0], [0]]))
    rot_mtx = np.vstack((rot_mtx, [0, 0, 0, 1]))
    result = translation_mtx @ rot_mtx
    return result


def draw_box(draw, box):
    draw.line(box[0] + box[1], fill=(255, 0, 0, 255), width=1)
    draw.line(box[1] + box[2], fill=(255, 0, 0, 255), width=1)
    draw.line(box[2] + box[3], fill=(255, 0, 0, 255), width=1)
    draw.line(box[3] + box[0], fill=(255, 0, 0, 255), width=1)
    draw.line(box[4] + box[5], fill=(255, 0, 0, 255), width=1)
    draw.line(box[5] + box[6], fill=(255, 0, 0, 255), width=1)
    draw.line(box[6] + box[7], fill=(255, 0, 0, 255), width=1)
    draw.line(box[7] + box[4], fill=(255, 0, 0, 255), width=1)
    draw.line(box[7] + box[3], fill=(255, 0, 0, 255), width=1)
    draw.line(box[4] + box[0], fill=(255, 0, 0, 255), width=1)
    draw.line(box[2] + box[-2], fill=(255, 0, 0, 255), width=1)
    draw.line(box[1] + box[-3], fill=(255, 0, 0, 255), width=1)


def to_world_space(image, depth, proj, view, w, h):
    world_space = np.empty((h, w, 4), dtype=np.float32)
    idx = np.indices((h, w), dtype=np.float32)
    world_space[:, :, 0] = idx[1]
    world_space[:, :, 1] = idx[0]
    world_space[:, :, 2] = depth
    world_space[:, :, 3] = 1
    world_space[:, :, 0] = ((2 * world_space[:, :, 0]) / w) - 1
    world_space[:, :, 1] = ((2 * (h - world_space[:, :, 1])) / h) - 1
    world_space = lin.inv(proj) @ world_space[:, :, :, np.newaxis]
    world_space = lin.inv(view) @ world_space[:, :, :]
    world_space = np.squeeze(world_space)
    world_space[:, :] = world_space[:, :] / world_space[:, :, 3, None]
    return np.squeeze(world_space)


def paint_pixels(world_space, world, box: Tuple[np.array, np.array], val):
    model_space = lin.inv(world) @ world_space[:, :, :, np.newaxis]
    model_space = np.squeeze(model_space)
    min = np.hstack((box[0], [0]))
    max = np.hstack((box[1], [2]))
    gt = model_space[:, :, :] > min
    lt = model_space[:, :, :] < max
    allgt = np.all(gt, axis=-1)
    alllt = np.all(lt, axis=-1)
    all = np.logical_and(allgt, alllt)
    # all = alllt
    return (val * all).astype(np.uint32)


def draw_bboxes(image, depth, proj, view, boxes: List[MultiPoint], xforms: List[np.matrix]):
    transformed_boxes = [affine_transform(x, list(y[0:3, 0:3].flat) + list(y[0:3, 3].flat)) for x, y in
                         zip(boxes, xforms)]
    boxes_matrix = [[np.matrix([[p.x], [p.y], [p.z], [1]]) for p in x.geoms] for x in boxes]
    boxes_world = [[y @ x for x in box] for (y, box) in zip(xforms, boxes_matrix)]
    boxes_view = [[view @ x for x in box] for box in boxes_world]
    boxes_proj = [[proj @ x for x in box] for box in boxes_view]
    boxes_ndc = [[x / x[3] for x in box] for box in boxes_proj]
    boxes_image = [[[((vec.A1[0] + 1) / 2) * 1920, (1 - ((vec.A1[1] + 1) / 2)) * 1080] for vec in box] for box in
                   boxes_ndc]

    # transformed_boxes = [affine_transform(x, view.A1[:12]) for x in transformed_boxes]
    # transformed_boxes = [affine_transform(x, proj.A1[:12]) for x in transformed_boxes]
    positions = [y @ np.matrix([[0], [0], [0], [1]]) for y in xforms]
    positions_view = [view @ x for x in positions]
    positions_proj = [proj @ x for x in positions_view]
    positions_ndc = [x / x[3] for x in positions_proj]
    positions_image = [[((vec.A1[0] + 1) / 2) * 1920, (1 - ((vec.A1[1] + 1) / 2)) * 1080] for vec in positions_ndc]
    # depths = [depth[y,x] for x, y in positions_image]
    # pdepths = [pos.A1[2] for pos in positions_ndc]
    # depths_diff = [abs(a-b) for a, b in zip(depths, pdepths)]
    pilimage = Image.fromarray(image, mode="RGBA")
    draw = ImageDraw.Draw(pilimage)
    for x, y in [pos for sublist in boxes_image for pos in sublist]:
        if x > 1920 or x < 0 or y > 1080 or y < 0:
            continue
        draw.ellipse([x - 4, y - 4, x + 4, y + 4], fill=(0, 255, 0, 200))
    for x, y in positions_image:
        if x > 1920 or x < 0 or y > 1080 or y < 0:
            continue
        draw.ellipse([x - 4, y - 4, x + 4, y + 4], fill=(255, 0, 0, 200))
    for box in boxes_image:
        # draw.line(box[0] + box[-1], fill=(255,0,0,255), width=10)
        draw_box(draw, box)

    del draw
    pilimage.show("Test")
    pilimage.save("test.bmp")
    # print(transformed_boxes)


def stencil_cull(segmentations, stencil, cls=StenCode.car, w=1920, h=1080):
    mask = np.full((h, w), 0x7, dtype=np.uint8)
    stencil_unmask = np.bitwise_and(stencil, mask)
    thresh = cv2.compare(stencil_unmask, cls.value, cv2.CMP_EQ)
    thresh = (thresh // 255).astype(np.uint32)
    # person_thresh = cv2.compare(stencil_unmask, StenCode.person.value, cv2.CMP_EQ)
    # person_thresh = (person_thresh // 255).astype(np.uint32)
    return segmentations * thresh


def get_bboxes(img, detections, w=1920, h=1080):
    result = []
    for val in detections:
        a = np.where(img == val['handle'])
        if len(a[0]) == 0:
            continue
        bbox = Box(((np.min(a[1]) / w, np.min(a[0]) / h), (np.max(a[1]) / w, np.max(a[0]) / h)))
        count = np.count_nonzero(a)
        normw = np.min(a[1]) - np.max(a[1])
        normh = np.min(a[0]) - np.max(a[0])
        npix = normw * normh
        cov = count / npix
        result.append(BBox(val['detection_id'], bbox, cov))
    return result


def process_detections(base_data_dir, detections):
    # detections = list(detections)
    # print('\n' + str(len(detections)) + '\n')
    # sem.acquire()
    # imgpath = Path(base_data_dir) / Path(detections[0]['runguid']) / Path(detections[0]['imagepath'], mode='r')
    # filename = 'info-{}.tiff'
    filename = '{}.tiff'
    depth_filename = '{}-depth.tiff'
    stencil_filename = '{}-stencil.tiff'
    imgpath = path.join(base_data_dir, filename.format(detections[0]['imagepath']))
    depth_imgpath = path.join(base_data_dir, depth_filename.format(detections[0]['imagepath']))
    stencil_imgpath = path.join(base_data_dir, stencil_filename.format(detections[0]['imagepath']))
    if not path.exists(imgpath):
        print(imgpath)
        return (None, None)
    tiffimg = TIFF.open(str(imgpath))
    depth_tiffimg = TIFF.open(str(depth_imgpath))
    stencil_tiffimg = TIFF.open(str(stencil_imgpath))
    img = Image.open(str(imgpath))
    w = img.width
    h = img.height
    # del img   # this causes segfault when running sequentially, not in parallel
    image = np.empty((h, w, 4), dtype=np.uint8)
    depth = np.empty((h, w), dtype=np.float32)
    stencil = np.empty((h, w), dtype=np.uint8)
    TIFF.setdirectory(tiffimg, 0)
    TIFF.readencodedstrip(tiffimg, 0, image.ctypes.data, -1)
    lastdir = num_directories(tiffimg) - 1
    # TIFF.setdirectory(tiffimg, lastdir - 1)
    TIFF.readencodedstrip(depth_tiffimg, 0, depth.ctypes.data, -1)
    # TIFF.setdirectory(tiffimg, lastdir)
    TIFF.readencodedstrip(stencil_tiffimg, 0, stencil.ctypes.data, -1)
    affines = [create_affine(loads(x['rot']), loads(x['pos'])) for x in detections]
    bboxes_geom = [loads(x['fullbox']) for x in detections]
    bboxes = [(np.array(loads(x['bbox3d_min'])), np.array(loads(x['bbox3d_max']))) for x in detections]
    handles = [x['handle'] for x in detections]
    view = np.array(detections[0]['view_matrix'], dtype=np.float64)
    proj = np.array(detections[0]['proj_matrix'], dtype=np.float64)
    world_space = to_world_space(image, depth, proj, view, w, h)
    output = np.zeros((h, w), dtype=np.uint32)
    max = math.pow(2, 24) - 1
    step = max // len(handles)
    i = step
    for box, xform, val in zip(bboxes, affines, handles):
        output += paint_pixels(world_space, xform, box, val)
        i += step
    output_cars = stencil_cull(output, stencil, StenCode.car, w, h)
    # output_peds = stencil_cull(output, stencil, StenCode.person, w, h)
    # draw_bboxes(image, depth, proj, view, bboxes_geom,  affines)
    # colorout = cv2.applyColorMap(output.astype(np.uint8), cv2.COLORMAP_JET)
    output_col = output_cars.view(np.uint8).reshape(output_cars.shape+(4,))[...,1:4]
    #cv2.imshow("color", output_col)
    #cv2.waitKey(0)
    result = get_bboxes(output_cars, detections, w, h)
    # result += get_bboxes(output_peds, detections)
    return (result, output_cars)


def get_conn():
    """
    :rtype: postgresql.api.Connection
    """
    with open('config.json', 'r', encoding='utf-8') as file:
        credentials = json.load(file)
    db_user = credentials['db_user']
    db_name = credentials['db_name']
    db_password = credentials['db_password']
    db_host = credentials['db_host']
    db_port = credentials['db_port']
    # db = postgresql.open('pq://user:password@host:port/database')
    conn = pg.open("pq://{}:{}@{}:{}/{}".format(db_user, db_password, db_host, db_port, db_name))
    return conn


def process(pixel_path, base_data_dir, run):
    conn = get_conn()
    print("Query Images.....")
    # query for real extraction
    # prep_stmt = conn.query(
    #     "SELECT snapshot_id, detection_id, runguid::text, imagepath, view_matrix, proj_matrix, handle, pos::bytea, rot::bytea, bbox,"
    #     "ngv_box3dpolygon(bbox3d)::bytea as fullbox,"
    #     "ST_MakePoint(ST_XMin(bbox3d), ST_YMin(bbox3d), ST_ZMin(bbox3d))::bytea as bbox3d_min,"
    #     "ST_MakePoint(ST_XMax(bbox3d), ST_YMax(bbox3d), ST_ZMax(bbox3d))::bytea as bbox3d_max FROM detections JOIN snapshots USING (snapshot_id) JOIN runs USING (run_id) JOIN sessions USING(session_id)"
    #     "WHERE session_id=$1 and processed=false and camera_pos <-> pos < 200 order by snapshot_id desc", session)

    # copeid and modified query for sample data
    sample_names = glob.glob(base_data_dir + '/*[0-9].tiff', recursive=False)
    sample_names = [path.splitext(path.basename(s))[0].replace('info-', '') for s in sample_names]

    prep_stmt = conn.query((
        "SELECT snapshot_id, detection_id, runguid::text, imagepath, view_matrix, proj_matrix, handle, pos::bytea, rot::bytea, bbox,"+
        "ngv_box3dpolygon(bbox3d)::bytea as fullbox,"+
        "ST_MakePoint(ST_XMin(bbox3d), ST_YMin(bbox3d), ST_ZMin(bbox3d))::bytea as bbox3d_min,"+
        "ST_MakePoint(ST_XMax(bbox3d), ST_YMax(bbox3d), ST_ZMax(bbox3d))::bytea as bbox3d_max FROM detections "+
        "JOIN snapshots USING (snapshot_id) JOIN runs USING (run_id) "+
        "WHERE run_id={} "+
        "AND imagepath IN ({}) "+
        "order by snapshot_id desc").format(run, ', '.join(['\'{}\''.format(s) for s in sample_names])))

    pbar = ProgressBar(max_value=len(prep_stmt)).start()
    i = 0
    sem = Semaphore(100)
    pool = ProcessPoolExecutor(100)
    results = []
    conn.close()
    lck = Lock()
    def on_done(snapshot_id, x):
        result = x.result()
        sem.release()
        with lck:
            nonlocal i
            nonlocal results
            pbar.update(i)
            i += 1
            #if result is None: return
            #results.append((snapshot_id, x.result()[0], x.result()[1]))
            upload([(snapshot_id, result[0], result[1])], Path(pixel_path))

    last_id = 0
    for snapshot_id, detections in groupby(prep_stmt, key=lambda x: x['snapshot_id']):
        sem.acquire()
        detections = list(detections)
        last_id = snapshot_id

        # serial
        # on_done(snapshot_id, process_detections(base_data_dir, detections))
        # parallel
        result = pool.submit(process_detections, base_data_dir, detections)
        result.add_done_callback(partial(on_done, snapshot_id))
    pool.shutdown(wait=True)
    pbar.finish()

    conn = get_conn()
    conn.query("UPDATE snapshots set processed=false where snapshot_id=$1", last_id)
    conn.close()
    # print(results)

    return results

    # original db
    # conn = pg.open("pq://sim-group:XXXXXXX@database.ngvlab.org/sim_annotations")


def upload(results, pixel_path: Path):
    conn = get_conn()
    update_query = conn.prepare("UPDATE detections SET best_bbox=$1 WHERE detection_id = $2")
    done_query = conn.prepare("UPDATE snapshots set processed=true where snapshot_id=$1")
    #pbar2 = ProgressBar()
    print("Uploading Results")
    print("Snapshot: {}".format(results[0][0]))
    for snapshot_id, detections, stencil in results:
        done_query(snapshot_id)
        if detections is None: continue
        print("Found: {} detections".format(len(detections)))
        for item in detections:
            update_query(item.bbox, item.detection_id)
        if stencil is not None:
            np.savez_compressed(str(pixel_path / (str(snapshot_id) + ".npz")), stencil)
            #cv2.imwrite(str(pixel_path / (str(snapshot_id) + ".png")), stencil)
    conn.close()


if __name__ == "__main__":
    parser = ArgumentParser(
        description="Process a GTA session using data from the stencil buffer as well as the camera parameters")
    parser.add_argument("--run", dest='run', required=True, type=int, help="the run to process")
    parser.add_argument("--dataroot", dest='dataroot', required=True, type=str, help="Location of the data")
    parser.add_argument("--pixel_path", dest="pixel_path", required=True, type=str,
                        help="Location to output pixel annotations")
    parser.add_argument("--resume", dest='resume', required=False, type=str, help="resume from a file")
    save_parser = parser.add_mutually_exclusive_group(required=False)
    save_parser.add_argument('--save', dest='save', action='store_true')
    save_parser.add_argument('--no-save', dest='save', action='store_false')
    parser.set_defaults(feature=True)
    args = parser.parse_args()
    results = None
    if (args.resume is None):
        results = process(args.pixel_path, args.dataroot, args.run)
        print("dumping results")
        with open('results.pkl', 'wb') as f:
            dump(results, f)

    else:
        print("loading results")
        with open(args.resume, 'rb') as f:
            results = load(f)

    #if args.save:
        #upload(results, Path(args.pixel_path))
