# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# pylint: disable=no-member
"""
TridentNet Training Script.

This script is a simplified version of the training script in detectron2/tools.
"""
import math
from typing import List, Tuple, Union, Set, DefaultDict
from typing_extensions import Literal, TypedDict
import pickle
from pathlib import Path
from collections import defaultdict
import json
import os
import sys
from tqdm import tqdm
import torch
from torch.utils.data import Dataset, DataLoader
import torch.multiprocessing as mp
from torch.multiprocessing import set_start_method
import argtyped
import cv2
import lmdb
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.colors as mcolors
from PIL import Image
from apex import amp
import torch
from torch.nn import functional as F

sys.path.append("detectron2")

from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.engine import DefaultTrainer, default_setup
from detectron2.structures import Instances

from utils.extract_utils import get_image_blob
from models import add_config
from models.bua.box_regression import BUABoxes
from models.bua.layers.nms import batched_nms
# Only valid with fp32 inputs - give AMP the hint

Mode = Literal["bbox_feats", "roi_feats", "bboxes"]


class MatterportFeature(TypedDict):
    image_feat: List[Image.Image]
    view_ids: List[int]
    reverie_ids: List[int]
    image_h: int
    image_w: int
    fov: float
    boxes: List[List[int]]
    cls_probs: List[np.ndarray]


class Arguments(argtyped.Arguments):
    config_file: Path = Path("configs/bua-caffe/extract-bua-caffe-r101.yaml")
    part_ids: List[int] = [0]
    num_parts: int = 1
    extract_mode: Mode = "roi_feats"
    mode: str = 'caffe'
    overwrite: bool = False
    min_max_boxes: str = "min_max_default"
    output: Path = Path("output")
    bbox_dir: Path = Path("bbox_dir")
    connectivity: Path = Path("connectivity")
    visualize: bool = False
    preload: bool = False

    nms_thresh: float = 0.3
    conf_thresh: float = 0.4
    min_local_boxes: int = 5
    max_local_boxes: int = 20
    max_total_boxes: int = 100

    num_sweeps: int = 3
    views_per_sweep: int = 12
    # number of total views from one pano
    viewpoint_size: int = 36
    heading_inc: int = 30
    # margin of error for deciding if an object is closer to the centre of another view
    angle_margin: int=5
    # elevation on first sweep
    elevation_start: int =-30 
    # how much elevation increases each sweep
    elevation_inc: int =30 

    # DataLoader
    num_workers: int = 0
    batch_size: int = 1

    # Matterport LMDB containing view images (see pointer-reverie/scripts/matterport_images.py)
    matterport: Path = Path("matterport.lmdb")


def set_min_max_boxes(min_max_boxes):
    if min_max_boxes == "min_max_default":
        return []
    try:
        min_boxes = int(min_max_boxes.split(",")[0])
        max_boxes = int(min_max_boxes.split(",")[1])
    except:
        print("Illegal min-max boxes setting, using config default. ")
        return []
    cmd = [
        "MODEL.BUA.EXTRACTOR.MIN_BOXES",
        min_boxes,
        "MODEL.BUA.EXTRACTOR.MAX_BOXES",
        max_boxes,
    ]
    return cmd


def switch_extract_mode(mode: Mode):
    if mode == "roi_feats":
        switch_cmd = ["MODEL.BUA.EXTRACTOR.MODE", 1]
    elif mode == "bboxes":
        switch_cmd = ["MODEL.BUA.EXTRACTOR.MODE", 2]
    elif mode == "bbox_feats":
        switch_cmd = [
            "MODEL.BUA.EXTRACTOR.MODE",
            3,
            "MODEL.PROPOSAL_GENERATOR.NAME",
            "PrecomputedProposals",
        ]
    else:
        print("Wrong extract mode! ")
        exit()
    return switch_cmd


def setup(args: Arguments):
    """
    Create configs and perform basic setups.
    """
    cfg = get_cfg()
    add_config(args, cfg)
    cfg.merge_from_file(str(args.config_file))
    # cfg.merge_from_list(args.opts)
    cfg.merge_from_list(switch_extract_mode(args.extract_mode))
    cfg.merge_from_list(set_min_max_boxes(args.min_max_boxes))
    cfg.freeze()
    default_setup(cfg, args)
    return cfg




def get_detections_from_im( im: np.ndarray, model, cfg, fov: float, boxes = None
        ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    im = np.array(im)
    dataset_dict = get_image_blob(im, cfg.MODEL.PIXEL_MEAN)

    # extract bbox only
    if boxes is not None:
        boxes = boxes * dataset_dict["im_scale"]
        proposals = Instances(dataset_dict["image"].shape[-2:])
        proposals.proposal_boxes = BUABoxes(boxes)
        dataset_dict["proposals"] = proposals

    with torch.no_grad():
        if cfg.MODEL.BUA.ATTRIBUTE_ON:
            boxes, cls_scores, features_pooled, attr_scores = model(
                [dataset_dict]
            )
        else:
            boxes, cls_scores, features_pooled = model([dataset_dict])

    return boxes[0].tensor, cls_scores[0], features_pooled[0]


def filter_bbox(boxes, prob, view_id: int, width: int, height: int, fov: float, args: Arguments):
    num_boxes, num_cls = boxes.shape

    # Keep only the best detections
    max_conf = torch.zeros((num_boxes)).to(prob.device)
    for cls_ind in range(1, num_cls):
        scores = prob[:, cls_ind]
        keep = batched_nms(boxes, scores, torch.arange(num_boxes), args.nms_thresh)
        mask = scores[keep] > max_conf[keep]
        max_conf[keep][mask] = scores[keep][mask]

    keep_boxes = max_conf >= args.conf_thresh
    if len(keep_boxes) < args.min_local_boxes:
        top_boxes = torch.topk(max_conf, args.min_local_boxes).indices
        keep_boxes = torch.zeros(num_boxes).bool().to(prob.device)
        keep_boxes[top_boxes] = True
    elif len(keep_boxes) > args.max_local_boxes:
        top_boxes = torch.topk(max_conf, args.max_local_boxes).indices
        keep_boxes = torch.zeros(num_boxes).bool().to(prob.device)
        keep_boxes[top_boxes] = True

    # Discard any box that would be better centered in another image
    # threshold for pixel distance from center of image
    focal_length = (height/2)/math.tan(math.radians(fov/2)) # focal length
    hor_thresh = focal_length * math.tan(math.radians(args.heading_inc/2 + args.angle_margin))
    vert_thresh = focal_length * math.tan(math.radians(args.elevation_inc/2 + args.angle_margin))
    center_x = 0.5*(boxes[:,0]+boxes[:,2])
    center_y = 0.5*(boxes[:,1]+boxes[:,3])
    reject = (center_x < width/2-hor_thresh) | (center_x > width/2+hor_thresh)

    if view_id >= args.views_per_sweep: # not lowest sweep
        reject |= (center_y > height/2+vert_thresh)
    if view_id < args.viewpoint_size - args.views_per_sweep: # not highest sweep
        reject |= (center_y < height/2-vert_thresh)
    keep_boxes = keep_boxes & ~reject

    return keep_boxes


def get_ft_head_elev(bbox: torch.Tensor, view_id: torch.Tensor, width: int, height: int, fov: float):
    # Calculate the heading and elevation of the center of each observation
    center_x = 0.5 * (bbox[:, 0] + bbox[:, 2])
    center_y = 0.5 * (bbox[:, 1] + bbox[:, 3])
    focal_length = torch.tensor((height / 2) / math.tan(math.radians(fov / 2))).to(bbox.device)

    heading = torch.deg2rad((view_id % 12) * 30)
    ft_heading = heading + torch.atan2(center_x - width / 2, focal_length)
    # normalize featureHeading
    ft_heading = ft_heading % (math.pi * 2)
    assert (0 <= ft_heading).all() and (ft_heading <= math.pi * 2).all()
    # force it to be the positive remainder, so that 0 <= angle < 360
    more_than_pi = ft_heading > math.pi
    ft_heading[more_than_pi] = (ft_heading - math.pi * 2)[more_than_pi]
    assert (-math.pi <= ft_heading).all() and (ft_heading <= math.pi).all()

    elevation = torch.deg2rad((view_id // 12) * 30)
    ft_elevation = elevation + torch.atan2(
        -center_y + height / 2, focal_length
    )

    return ft_heading, ft_elevation


def filter_panorama(boxes: torch.Tensor, probs: torch.Tensor, features: torch.Tensor, view_ids: torch.Tensor, max_boxes: int, width: int, height: int, fov: float) -> torch.Tensor:
    ft_heading, ft_elevation = get_ft_head_elev(boxes, view_ids, width, height, fov)
    # Remove the most redundant features (that have similar heading, elevation and 
    # are close together to an existing feature in cosine distance)
    ext_features = features.unsqueeze(2)
    feat_dist = F.cosine_similarity(ext_features, ext_features.transpose(0, 2), dim=1)

    indices = torch.triu_indices(*feat_dist.shape, 1)

    heading_diff_tri = F.pdist(ft_heading.unsqueeze(1), 2)
    heading_diff_tri = torch.minimum(heading_diff_tri, 2*math.pi - heading_diff_tri)
    heading_diff = torch.zeros_like(feat_dist)
    heading_diff[indices[0], indices[1]] = heading_diff_tri

    elevation_diff_tri = F.pdist(ft_elevation.unsqueeze(1), 2)
    elevation_diff = torch.zeros_like(feat_dist)
    elevation_diff[indices[0], indices[1]] = elevation_diff_tri

    total_dist = feat_dist + heading_diff + elevation_diff # Could add weights

    # Discard diagonal and upper triangle by setting large distance
    dist = total_dist[indices[0], indices[1]]
    arg_ind = torch.argsort(dist)

    # Remove indices of the most similar features (in appearance and orientation)
    keep = set(range(feat_dist.shape[0]))
    ix = 0
    while len(keep) > max_boxes:
        min_ind = arg_ind[ix]
        i, j = indices[:, min_ind].tolist()

        if i not in keep or j not in keep:
            ix += 1
            continue

        if probs[i,1:].max() > probs[j,1:].max():
            keep.remove(j)
        else:
            keep.remove(i)
        ix += 1

    return torch.Tensor(list(keep)).long().to(boxes.device)


def extract_feat(worker_id, viewpoint_lists, args: Arguments):
    part_id = args.part_ids[worker_id]
    viewpoint_list = viewpoint_lists[part_id]
    num_viewpoints = len(viewpoint_list)

    cfg = setup(args)
    print("Number of viewpoints on split{}: {}.".format(part_id, num_viewpoints))

    num_gpus = torch.cuda.device_count()
    assert num_gpus != 0
    device_id = part_id % num_gpus
    torch.cuda.set_device(device_id)

    if args.num_parts == 1:
        lmdb_path = args.output
    else:
        lmdb_path = args.output.parent / f"{args.output.stem}-{part_id}.lmdb"
    writer = LMDBWriter(str(lmdb_path), map_size=int(1e12), buffer_size=300, overwrite=args.overwrite)

    model = DefaultTrainer.build_model(cfg)
    DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
        cfg.MODEL.WEIGHTS, resume=False
    )

    model.eval()
    done = set(tuple(bkey.decode().split("_")) for bkey in  writer._keys)
    print('Done', len(done))
    viewpoint_list = [v for v in viewpoint_list if v not in done]

    dataset = MatterportDataset(viewpoint_list, args)
    dataloader = DataLoader(
        dataset,
        num_workers=args.num_workers,
        batch_size=10,
        collate_fn=lambda x: x,
    )

    batch: List[Tuple[MatterportFeature, DefaultDict, str, str]]
    for batch in tqdm(dataloader):
        for feats, bboxes_per_view_id, scan, viewpoint in batch:
            all_boxes = []
            all_features = []
            all_probs = []
            all_view_ids = []
            all_reverie_ids = []

            for view_id, im in zip(feats['view_ids'], feats['image_feat']):

                if args.extract_mode == "bbox_feats":
                    boxes = torch.Tensor([b["bbox"] for b in bboxes_per_view_id[view_id]])
                    all_reverie_ids += [b["obj_id"] for b in bboxes_per_view_id[view_id]]
                else:
                    boxes = None

                detected_boxes, probs, features = get_detections_from_im(np.array(im), model, cfg, feats['fov'], boxes)

                if args.extract_mode != 'bbox_feats':
                    mask_keep = filter_bbox(
                        detected_boxes,
                        probs,
                        view_id,
                        feats['image_w'],
                        feats['image_h'],
                        feats['fov'],
                        args
                    )
                    boxes = detected_boxes[mask_keep]
                    probs = probs[mask_keep]
                    features = features[mask_keep]

                all_boxes.append(boxes)
                all_probs.append(probs)
                all_features.append(features)
                num_bbox = boxes.shape[0]
                all_view_ids += [view_id] * num_bbox


            boxes = torch.cat(all_boxes)
            keep_ind = filter_panorama(
                boxes,
                torch.cat(all_probs),
                torch.cat(all_features),
                torch.Tensor(all_view_ids).to(boxes.device),
                args.max_total_boxes, 
                feats['image_w'],
                feats['image_h'],
                feats['fov'],
            )

            key = f"{scan}_{viewpoint}"
            image_feat = torch.cat(all_features)[keep_ind].cpu()
            bbox = torch.cat(all_boxes)[keep_ind].cpu()
            probs = torch.cat(all_probs)[keep_ind].cpu()
            view_ids = torch.Tensor(all_view_ids)[keep_ind].cpu()
            reverie_ids = (
                keep_ind.cpu()
                if all_reverie_ids == []
                else torch.Tensor(all_reverie_ids)
            )

            data = {
                "image_feat": image_feat,
                "boxes": bbox,
                "image_h": feats['image_h'],
                "image_w": feats['image_w'],
                "fov": feats['fov'],
                "view_ids": view_ids,
                "reverie_ids": reverie_ids,
                "cls_probs": probs,

            }
            writer.put(key.encode('ascii'), pickle.dumps(data))

        writer.flush()


def main():
    args = Arguments()

    # Extract features.
    viewpoints = load_viewpoints(args)
    viewpoints_lists = [viewpoints[i::args.num_parts] for i in range(args.num_parts)]
    assert len(viewpoints) == len(sum(viewpoints_lists, []))
    
    if args.part_ids == []:
        raise ValueError()
    elif len(args.part_ids) == 1:
        extract_feat(0, viewpoints_lists, args)
    else:
        parts = (viewpoints_lists, args)
        print('Spawn ', len(args.part_ids), 'procs')
        mp.spawn(extract_feat, nprocs=len(args.part_ids), args=parts)


def load_viewpoints(args: Arguments) -> List[Tuple[str, str]]:
    viewpoint_ids = set()

    with open(args.connectivity / "scans.txt") as f:
        scans = set(x.strip() for x in f)

    for scan in scans:
        with open(args.connectivity / f"{scan}_connectivity.json") as f:
            data = json.load(f)
            viewpoint_ids |= set((scan, x["image_id"]) for x in data if x["included"])
    try:
         set_start_method('spawn', force=True)
    except RuntimeError:
        pass


    print("Loaded: %d viewpoints" % len(viewpoint_ids))

    # Remove viewpoints related to empty JSON
    if args.extract_mode == 'bbox_feats':
        empty = set(
            (scan, vp)
            for scan, vp in viewpoint_ids
            if not has_bbox(scan, vp, args.bbox_dir)
        )
        print("Empty: %d viewpoints" % len(empty))
        viewpoint_ids = viewpoint_ids - empty
        print("Total: %d viewpoints" % len(viewpoint_ids))

    return sorted(viewpoint_ids)


def has_bbox(scan: str, vp: str, bbox_dir: Path) -> bool:
    with open(bbox_dir / f"{scan}_{vp}.json") as fid:
        data = json.load(fid)
    for obj in data[vp].values():
        if obj["bbox2d"] != []:
            return True
    return False


def load_json(path):
    with open(path, "r") as fid:
        data = json.load(fid)
    return data


def load_boxes(boxes_path: Path):
    boxes_per_vp = {}
    for json_file in boxes_path.glob("*.json"):
        boxes_per_vp.update(load_json(json_file))
    return boxes_per_vp


class LMDBWriter:
    def __init__(self, path: Union[Path, str], map_size: int, buffer_size: int, overwrite: bool):
        self._env = lmdb.open(str(path), map_size=map_size, max_readers=512)
        self._buffer: List[Tuple[bytes, bytes]] = []
        self._buffer_size = buffer_size
        self._overwrite = overwrite

        with self._env.begin(write=False) as txn:
            bkeys = txn.get("__keys__".encode("ascii"))
            if bkeys is None:
                self._keys: Set[bytes] = set()
            else:
                self._keys = set(pickle.loads(bkeys))

    def put(self, bkey: bytes, value: bytes):
        if not self._overwrite and bkey in self._keys:
            print(bkey, 'is already present')
            return
        self._buffer.append((bkey, value))
        self._keys.add(bkey)
        if len(self._buffer) == self._buffer_size:
            self.flush()

    def get_keys(self) -> List[str]:
        return [k.decode("ascii") for k in self._keys]

    def flush(self):
        with self._env.begin(write=True) as txn:
            for bkey, value in self._buffer:
                txn.put(bkey, value)
            txn.put("__keys__".encode("ascii"), pickle.dumps(self._keys))

        self._buffer = []


class MatterportDataset(Dataset):
    def __init__(self, viewpoints: List[Tuple[str, str]], args: Arguments):
        self.viewpoints = viewpoints
        self.args = args
        self.counter = 0
        self.bbox = load_boxes(args.bbox_dir)
        self.env = lmdb.open(str(args.matterport))
        self.txn = self.env.begin(write=False)

    def __len__(self):
        return len(self.viewpoints)

    def __iter__(self):
        self.counter = 0
        return self

    def __next__(self):
        self.counter += 1
        if self.counter == len(self):
            raise StopIteration()
        return self[self.counter]

    def __getitem__(self, index: int) -> Tuple[MatterportFeature, DefaultDict, str, str]:
        scan, viewpoint = self.viewpoints[index]

        key = f"{scan}_{viewpoint}"
        feats = self.txn.get(key.encode('ascii'))
        if feats is None:
            raise RuntimeError(f"Can't find the key {key}")
        feats = pickle.loads(feats)

        bboxes_per_view_id = defaultdict(list)
        for obj_id, item in self.bbox[viewpoint].items():
            for view_id, bbox in zip(item["visible_pos"], item["bbox2d"]):
                xywh2xyxy(bbox)
                bboxes_per_view_id[view_id].append({"bbox": bbox, "obj_id": obj_id})

        return feats, bboxes_per_view_id, scan, viewpoint


def plot_viewpoint(im, boxes, filename: Path):
    plt.imshow(im[:, :, ::-1])
    ax = plt.gca()

    for b, color in zip(boxes, mcolors.BASE_COLORS):
        bbox = b["bbox"]
        h = bbox[3] - bbox[1]
        w = bbox[2] - bbox[0]
        pos = (bbox[0], bbox[1])
        rect = patches.Rectangle(
            (pos), w, h, linewidth=1, edgecolor=color, facecolor="none"
        )
        ax.add_patch(rect)

    plt.savefig(filename)


def xywh2xyxy(bbox):
    bbox[2] += bbox[0]
    bbox[3] += bbox[1]


def transform_img(im):
    """ Prep opencv BGR 3 channel image for the network """
    blob = np.array(im, copy=True)
    return blob


if __name__ == "__main__":
    try:
         set_start_method('spawn', force=True)
    except RuntimeError:
        pass

    main()
