import colorsys
import logging
import math
import numpy as np
from enum import Enum, unique
import cv2
import matplotlib as mpl
import matplotlib.colors as mplc
import matplotlib.figure as mplfigure
import pdb
import torch
from matplotlib.backends.backend_agg import FigureCanvasAgg

from mseg.utils.fb_colormap import random_color, fixed_color
from mseg.utils.mask_utils import get_mask_from_polygon, get_most_populous_class

logger = logging.getLogger(__name__)

__all__ = ["ColorMode", "VisImage", "Visualizer"]

_SMALL_OBJECT_AREA_THRESH = 1000
_LARGE_MASK_AREA_THRESH = 120000
_OFF_WHITE = (1.0, 1.0, 240.0 / 255)
_BLACK = (0, 0, 0)
_RED = (1.0, 0, 0)

@unique
class ColorMode(Enum):
    """
    Enum of different color modes to use for instance visualizations.

    Attributes:
        IMAGE: Picks a random color for every instance and overlay segmentations with low opacity.
        SEGMENTATION: Let instances of the same category have similar colors, and overlay them with
            high opacity. This provides more attention on the quality of segmentation.
        IMAGE_BW: same as IMAGE, but convert all areas without masks to gray-scale.
            Only available for drawing per-instance mask predictions.
    """

    IMAGE = 0
    SEGMENTATION = 1
    IMAGE_BW = 2

class GenericMask:
    """
    Attribute:
        polygons (list[ndarray]): list[ndarray]: polygons for this mask.
            Each ndarray has format [x, y, x, y, ...]
        mask (ndarray): a binary mask
    """

    def __init__(self, mask_or_polygons, height, width):
        self._mask = self._polygons = self._has_holes = None
        self.height = height
        self.width = width

        m = mask_or_polygons
        if isinstance(m, dict):
            # RLEs
            assert "counts" in m and "size" in m
            if isinstance(m["counts"], list):  # uncompressed RLEs
                h, w = m["size"]
                assert h == height and w == width
                m = mask_util.frPyObjects(m, h, w)
            self._mask = mask_util.decode(m)[:, :]
            return

        if isinstance(m, list):  # list[ndarray]
            self._polygons = [np.asarray(x).reshape(-1) for x in m]
            return

        if isinstance(m, np.ndarray):  # assumed to be a binary mask
            assert m.shape[1] != 2, m.shape
            assert m.shape == (height, width), m.shape
            self._mask = m.astype("uint8")
            return

        raise ValueError("GenericMask cannot handle object {} of type '{}'".format(m, type(m)))

    @property
    def mask(self):
        if self._mask is None:
            self._mask = self.polygons_to_mask(self._polygons)
        return self._mask

    @property
    def polygons(self):
        if self._polygons is None:
            self._polygons, self._has_holes = self.mask_to_polygons(self._mask)
        return self._polygons

    @property
    def has_holes(self):
        if self._has_holes is None:
            if self._mask is not None:
                self._polygons, self._has_holes = self.mask_to_polygons(self._mask)
            else:
                self._has_holes = False  # if original format is polygon, does not have holes
        return self._has_holes

    def mask_to_polygons(self, mask):
        # cv2.RETR_CCOMP flag retrieves all the contours and arranges them to a 2-level
        # hierarchy. External contours (boundary) of the object are placed in hierarchy-1.
        # Internal contours (holes) are placed in hierarchy-2.
        # cv2.CHAIN_APPROX_NONE flag gets vertices of polygons from contours.
        res, hierarchy = cv2.findContours(mask.astype("uint8"), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)

        if hierarchy is None:  # empty mask
            return [], False
        has_holes = (hierarchy.reshape(-1, 4)[:, 3] >= 0).sum() > 0

        res = [x.squeeze() for x in res]
        res = [x for x in res if x.size >= 6]  # should have at least 3 vertices to be valid
        return res, has_holes

    def polygons_to_mask(self, polygons):
        rle = mask_util.frPyObjects(polygons, self.height, self.width)
        rle = mask_util.merge(rle)
        return mask_util.decode(rle)[:, :]

    def area(self):
        return self.mask.sum()

    def bbox(self):
        p = mask_util.frPyObjects(self.polygons, self.height, self.width)
        p = mask_util.merge(p)
        bbox = mask_util.toBbox(p)
        bbox[2] += bbox[0]
        bbox[3] += bbox[1]
        return bbox

class _PanopticPrediction:
    def __init__(self, panoptic_seg, segments_info):
        self._seg = panoptic_seg

        self._sinfo = {s["id"]: s for s in segments_info}  # seg id -> seg info
        segment_ids, areas = torch.unique(panoptic_seg, sorted=True, return_counts=True)
        areas = areas.numpy()
        sorted_idxs = np.argsort(-areas)
        self._seg_ids, self._seg_areas = segment_ids[sorted_idxs], areas[sorted_idxs]
        self._seg_ids = self._seg_ids.tolist()
        for sid, area in zip(self._seg_ids, self._seg_areas):
            if sid in self._sinfo:
                self._sinfo[sid]["area"] = float(area)

    def non_empty_mask(self):
        """
        Returns:
            (H, W) array, a mask for all pixels that have a prediction
        """
        empty_ids = []
        for id in self._seg_ids:
            if id not in self._sinfo:
                empty_ids.append(id)
        if len(empty_ids) == 0:
            return np.zeros(self._seg.shape, dtype=np.uint8)
        assert len(empty_ids) == 1, ">1 ids corresponds to no labels. This is currently not supported"
        return (self._seg != empty_ids[0]).numpy().astype(np.bool)

    def semantic_masks(self):
        for sid in self._seg_ids:
            sinfo = self._sinfo.get(sid)
            if sinfo is None or sinfo["isthing"]:
                # Some pixels (e.g. id 0 in PanopticFPN) have no instance or semantic predictions.
                continue
            yield (self._seg == sid).numpy().astype(np.bool), sinfo

    def instance_masks(self):
        for sid in self._seg_ids:
            sinfo = self._sinfo.get(sid)
            if sinfo is None or not sinfo["isthing"]:
                continue
            mask = (self._seg == sid).numpy().astype(np.bool)
            if mask.sum() > 0:
                yield mask, sinfo

def _create_text_labels(classes, scores, class_names):
    """
    Args:
        classes (list[int] or None):
        scores (list[float] or None):
        class_names (list[str] or None):

    Returns:
        list[str] or None
    """
    labels = None
    if classes is not None and class_names is not None and len(class_names) > 1:
        labels = [class_names[i] for i in classes]
    if scores is not None:
        if labels is None:
            labels = ["{:.0f}%".format(s * 100) for s in scores]
        else:
            labels = ["{} {:.0f}%".format(l, s * 100) for l, s in zip(labels, scores)]
    return labels

class VisImage:
    def __init__(self, img, scale=1.0):
        """
        Args:
            img (ndarray): an RGB image of shape (H, W, 3).
            scale (float): scale the input image
        """
        self.img = img
        self.scale = scale
        self.width, self.height = img.shape[1], img.shape[0]
        self._setup_figure(img)

    def _setup_figure(self, img):
        """
        Args:
            Same as in :meth:`__init__()`.

        Returns:
            fig (matplotlib.pyplot.figure): top level container for all the image plot elements.
            ax (matplotlib.pyplot.Axes): contains figure elements and sets the coordinate system.
        """
        fig = mplfigure.Figure(frameon=False)
        self.dpi = fig.get_dpi()
        # add a small 1e-2 to avoid precision lost due to matplotlib's truncation
        # (https://github.com/matplotlib/matplotlib/issues/15363)
        fig.set_size_inches(
            (self.width * self.scale + 1e-2) / self.dpi,
            (self.height * self.scale + 1e-2) / self.dpi,
        )
        self.canvas = FigureCanvasAgg(fig)
        # self.canvas = mpl.backends.backend_cairo.FigureCanvasCairo(fig)
        ax = fig.add_axes([0.0, 0.0, 1.0, 1.0])
        ax.axis("off")
        ax.set_xlim(0.0, self.width)
        ax.set_ylim(self.height)

        self.fig = fig
        self.ax = ax

    def save(self, filepath):
        """
        Args:
            filepath (str): a string that contains the absolute path, including the file name, where
                the visualized image will be saved.
        """
        if filepath.endswith(".jpg") or filepath.endswith(".png"):
            # faster than matplotlib's imshow
            cv2.imwrite(filepath, self.get_image()[:, :, ::-1])
        else:
            # support general formats (e.g. pdf)
            self.ax.imshow(self.img, interpolation="nearest")
            self.fig.savefig(filepath)

    def get_image(self):
        """
        Returns:
            ndarray: the visualized image of shape (H, W, 3) (RGB) in uint8 type.
              The shape is scaled w.r.t the input image using the given `scale` argument.
        """
        canvas = self.canvas
        s, (width, height) = canvas.print_to_buffer()
        if (self.width, self.height) != (width, height):
            img = cv2.resize(self.img, (width, height))
        else:
            img = self.img

        # buf = io.BytesIO()  # works for cairo backend
        # canvas.print_rgba(buf)
        # width, height = self.width, self.height
        # s = buf.getvalue()

        buffer = np.frombuffer(s, dtype="uint8")

        # imshow is slow. blend manually (still quite slow)
        img_rgba = buffer.reshape(height, width, 4)
        rgb, alpha = np.split(img_rgba, [3], axis=2)

        try:
            import numexpr as ne  # fuse them with numexpr

            visualized_image = ne.evaluate("img * (1 - alpha / 255.0) + rgb * (alpha / 255.0)")
        except ImportError:
            alpha = alpha.astype("float32") / 255.0
            visualized_image = img * (1 - alpha) + rgb * alpha

        visualized_image = visualized_image.astype("uint8")

        return visualized_image

class Visualizer:
    def __init__(self, img_rgb, metadata, scale=1.0, instance_mode=ColorMode.IMAGE):
        """
        Args:
            img_rgb: a numpy array of shape (H, W, C), where H and W correspond to
                the height and width of the image respectively. C is the number of
                color channels. The image is required to be in RGB format since that
                is a requirement of the Matplotlib library. The image is also expected
                to be in the range [0, 255].
            metadata (MetadataCatalog): image metadata.
        """
        self.img = np.asarray(img_rgb).clip(0, 255).astype(np.uint8)
        self.metadata = metadata
        self.output = VisImage(self.img, scale=scale)
        self.cpu_device = torch.device("cpu")

        # too small texts are useless, therefore clamp to 9
        self._default_font_size = max(np.sqrt(self.output.height * self.output.width) // 90, 10 // scale)
        self._instance_mode = instance_mode

    def draw_sem_seg(self, sem_seg, area_threshold=None, alpha=0.8):
        """
        Draw semantic segmentation predictions/labels.

        Args:
            sem_seg (Tensor or ndarray): the segmentation of shape (H, W).
            area_threshold (int): segments with less than `area_threshold` are not drawn.
            alpha (float): the larger it is, the more opaque the segmentations are.

        Returns:
            output (VisImage): image object with visualizations.
        """
        if isinstance(sem_seg, torch.Tensor):
            sem_seg = sem_seg.numpy()
        labels, areas = np.unique(sem_seg, return_counts=True)
        sorted_idxs = np.argsort(-areas).tolist()
        labels = labels[sorted_idxs]
        for label in labels:
            try:
                mask_color = [x / 255 for x in self.metadata.stuff_colors[label]]
            except (AttributeError, IndexError):
                mask_color = None

            binary_mask = (sem_seg == label).astype(np.uint8)
            text = self.metadata.stuff_classes[label]
            self.draw_binary_mask(
                binary_mask,
                color=mask_color,
                edge_color=_OFF_WHITE,
                text=text,
                alpha=alpha,
                area_threshold=area_threshold,
            )
        return self.output

    def draw_panoptic_seg_predictions(self, panoptic_seg, segments_info, area_threshold=None, alpha=0.7):
        """
        Draw panoptic prediction results on an image.

        Args:
            panoptic_seg (Tensor): of shape (height, width) where the values are ids for each
                segment.
            segments_info (list[dict]): Describe each segment in `panoptic_seg`.
                Each dict contains keys "id", "category_id", "isthing".
            area_threshold (int): stuff segments with less than `area_threshold` are not drawn.

        Returns:
            output (VisImage): image object with visualizations.
        """
        pred = _PanopticPrediction(panoptic_seg, segments_info)

        if self._instance_mode == ColorMode.IMAGE_BW:
            self.output.img = self._create_grayscale_image(pred.non_empty_mask())

        # draw mask for all semantic segments first i.e. "stuff"
        for mask, sinfo in pred.semantic_masks():
            category_idx = sinfo["category_id"]
            try:
                mask_color = [x / 255 for x in self.metadata.stuff_colors[category_idx]]
            except AttributeError:
                mask_color = None

            text = self.metadata.stuff_classes[category_idx]
            self.draw_binary_mask(
                mask,
                color=mask_color,
                edge_color=_OFF_WHITE,
                text=text,
                alpha=alpha,
                area_threshold=area_threshold,
            )

        # draw mask for all instances second
        all_instances = list(pred.instance_masks())
        if len(all_instances) == 0:
            return self.output
        masks, sinfo = list(zip(*all_instances))
        category_ids = [x["category_id"] for x in sinfo]

        try:
            scores = [x["score"] for x in sinfo]
        except KeyError:
            scores = None
        labels = _create_text_labels(category_ids, scores, self.metadata.thing_classes)

        try:
            colors = [random_color(rgb=True, maximum=1) for k in category_ids]
        except AttributeError:
            colors = None
        self.overlay_instances(masks=masks, labels=labels, assigned_colors=colors, alpha=alpha)

        return self.output

    def overlay_instances(self, label_map, id_to_class_name_map, alpha=0.5):
        """
        Args:
            boxes (Boxes, RotatedBoxes or ndarray): either a :class:`Boxes`,
                or an Nx4 numpy array of XYXY_ABS format for the N objects in a single image,
                or a :class:`RotatedBoxes`,
                or an Nx5 numpy array of (x_center, y_center, width, height, angle_degrees) format
                for the N objects in a single image,
            labels (list[str]): the text to be displayed for each instance.
            masks (masks-like object): Supported types are:

                * `structures.masks.PolygonMasks`, `structures.masks.BitMasks`.
                * list[list[ndarray]]: contains the segmentation masks for all objects in one image.
                    The first level of the list corresponds to individual instances. The second
                    level to all the polygon that compose the instance, and the third level
                    to the polygon coordinates. The third level should have the format of
                    [x0, y0, x1, y1, ..., xn, yn] (n >= 3).
                * list[ndarray]: each ndarray is a binary mask of shape (H, W).
                * list[dict]: each dict is a COCO-style RLE.
            keypoints (Keypoint or array like): an array-like object of shape (N, K, 3),
                where the N is the number of instances and K is the number of keypoints.
                The last dimension corresponds to (x, y, visibility or score).
            assigned_colors (list[matplotlib.colors]): a list of colors, where each color
                corresponds to each mask or box in the image. Refer to 'matplotlib.colors'
                for full list of formats that the colors are accepted in.

        Returns:
            output (VisImage): image object with visualizations.
        """
        for label_idx in np.unique(label_map):
            color = fixed_color(rgb=True, maximum=1, idx=label_idx)
            text = id_to_class_name_map[label_idx]

            segment = label_map == label_idx

            # import pdb; pdb.set_trace()
            mask_obj = GenericMask(segment, self.output.height, self.output.width)
            polygons, _ = mask_obj.mask_to_polygons(segment)

            for polygon_verts in polygons:
                # print('Polygon verts: ', polygon_verts)

                # rasterize this specific polygon
                segment_mask = get_mask_from_polygon(polygon_verts, self.output.height, self.output.width)
                class_mode_idx = get_most_populous_class(segment_mask, label_map)
                if class_mode_idx != label_idx:
                    continue
                self.draw_polygon(polygon_verts, color, alpha=alpha)

                # import matplotlib.pyplot as plt
                # plt.imshow(segment_mask)
                # plt.show()

                # draw text in the center (defined by median) when box is not drawn
                # median is less sensitive to outliers.
                text_pos = np.median(segment_mask.nonzero(), axis=1)[::-1]
                horiz_align = "center"

                x0 = np.amin(polygon_verts[:, 0])
                x1 = np.amax(polygon_verts[:, 0])

                y0 = np.amin(polygon_verts[:, 1])
                y1 = np.amax(polygon_verts[:, 1])

                # else:
                #     continue  # drawing the box confidence for keypoints isn't very useful.
                # # for small objects, draw text at the side to avoid occlusion
                # instance_area = (y1 - y0) * (x1 - x0)
                # if (
                #     instance_area < _SMALL_OBJECT_AREA_THRESH * self.output.scale
                #     or y1 - y0 < 40 * self.output.scale
                # ):
                #     if y1 >= self.output.height - 5:
                #         text_pos = (x1, y0)
                #     else:
                #         text_pos = (x0, y1)

                height_ratio = (y1 - y0) / np.sqrt(self.output.height * self.output.width)
                lighter_color = self._change_color_brightness(color, brightness_factor=0.7)
                font_size = np.clip((height_ratio - 0.02) / 0.08 + 1, 1.2, 2) * 0.5 * self._default_font_size
                self.draw_text(
                    text,
                    text_pos,
                    color=lighter_color,
                    horizontal_alignment=horiz_align,
                    font_size=font_size,
                )
                # import matplotlib.pyplot as plt
                # plt.imshow(self.output.get_image())
                # plt.show()

        return self.output.get_image()

    """
    Primitive drawing functions:
    """

    def draw_text(self, text, position, *, font_size=None, color="g", horizontal_alignment="center", rotation=0):
        """
        Args:
            text (str): class label
            position (tuple): a tuple of the x and y coordinates to place text on image.
            font_size (int, optional): font of the text. If not provided, a font size
                proportional to the image width is calculated and used.
            color: color of the text. Refer to `matplotlib.colors` for full list
                of formats that are accepted.
            horizontal_alignment (str): see `matplotlib.text.Text`
            rotation: rotation angle in degrees CCW

        Returns:
            output (VisImage): image object with text drawn.
        """
        if not font_size:
            font_size = self._default_font_size

        # since the text background is dark, we don't want the text to be dark
        # color = np.maximum(list(mplc.to_rgb(color)), 0.2)
        # color[np.argmax(color)] = max(0.8, np.max(color))

        # force text color to be white!
        color = np.array([1.0, 1.0, 1.0])

        x, y = position
        self.output.ax.text(
            x,
            y,
            text,
            size=font_size * self.output.scale,
            family="sans-serif",
            bbox={"facecolor": "black", "alpha": 0.8, "pad": 0.7, "edgecolor": "none"},
            verticalalignment="top",
            horizontalalignment=horizontal_alignment,
            color=color,
            zorder=10,
            rotation=rotation,
        )
        return self.output

    def draw_binary_mask(self, binary_mask, color=None, *, edge_color=None, text=None, alpha=0.5, area_threshold=4096):
        """
        Args:
            binary_mask (ndarray): numpy array of shape (H, W), where H is the image height and
                W is the image width. Each value in the array is either a 0 or 1 value of uint8
                type.
            color: color of the mask. Refer to `matplotlib.colors` for a full list of
                formats that are accepted. If None, will pick a random color.
            edge_color: color of the polygon edges. Refer to `matplotlib.colors` for a
                full list of formats that are accepted.
            text (str): if None, will be drawn in the object's center of mass.
            alpha (float): blending efficient. Smaller values lead to more transparent masks.
            area_threshold (float): a connected component small than this will not be shown.

        Returns:
            output (VisImage): image object with mask drawn.
        """
        if color is None:
            color = random_color(rgb=True, maximum=1)
        if area_threshold is None:
            area_threshold = 4096

        has_valid_segment = False
        binary_mask = binary_mask.astype("uint8")  # opencv needs uint8
        mask = GenericMask(binary_mask, self.output.height, self.output.width)
        shape2d = (binary_mask.shape[0], binary_mask.shape[1])

        if not mask.has_holes:
            # draw polygons for regular masks
            for segment in mask.polygons:
                area = mask_util.area(mask_util.frPyObjects([segment], shape2d[0], shape2d[1]))
                if area < area_threshold:
                    continue
                has_valid_segment = True
                segment = segment.reshape(-1, 2)
                self.draw_polygon(segment, color=color, edge_color=edge_color, alpha=alpha)
        else:
            rgba = np.zeros(shape2d + (4,), dtype="float32")
            rgba[:, :, :3] = color
            rgba[:, :, 3] = (mask.mask == 1).astype("float32") * alpha
            has_valid_segment = True
            self.output.ax.imshow(rgba)

        if text is not None and has_valid_segment:
            # TODO sometimes drawn on wrong objects. the heuristics here can improve.
            lighter_color = self._change_color_brightness(color, brightness_factor=0.7)
            _num_cc, cc_labels, stats, centroids = cv2.connectedComponentsWithStats(binary_mask, 8)
            largest_component_id = np.argmax(stats[1:, -1]) + 1

            # draw text on the largest component, as well as other very large components.
            for cid in range(1, _num_cc):
                if cid == largest_component_id or stats[cid, -1] > _LARGE_MASK_AREA_THRESH:
                    # median is more stable than centroid
                    # center = centroids[largest_component_id]
                    center = np.median((cc_labels == cid).nonzero(), axis=1)[::-1]
                    self.draw_text(text, center, color=lighter_color)
        return self.output

    def draw_polygon(self, segment, color, edge_color=None, alpha=0.5):
        """
        Args:
            segment: numpy array of shape Nx2, containing all the points in the polygon.
            color: color of the polygon. Refer to `matplotlib.colors` for a full list of
                formats that are accepted.
            edge_color: color of the polygon edges. Refer to `matplotlib.colors` for a
                full list of formats that are accepted. If not provided, a darker shade
                of the polygon color will be used instead.
            alpha (float): blending efficient. Smaller values lead to more transparent masks.

        Returns:
            output (VisImage): image object with polygon drawn.
        """
        if edge_color is None:
            # make edge color darker than the polygon color
            if alpha > 0.8:
                edge_color = self._change_color_brightness(color, brightness_factor=-0.7)
            else:
                edge_color = color
        edge_color = mplc.to_rgb(edge_color) + (1,)

        polygon = mpl.patches.Polygon(
            segment,
            fill=True,
            facecolor=mplc.to_rgb(color) + (alpha,),
            edgecolor=edge_color,
            linewidth=max(self._default_font_size // 15 * self.output.scale, 1),
        )
        self.output.ax.add_patch(polygon)
        return self.output

    """
    Internal methods:
    """

    def _jitter(self, color):
        """
        Randomly modifies given color to produce a slightly different color than the color given.

        Args:
            color (tuple[double]): a tuple of 3 elements, containing the RGB values of the color
                picked. The values in the list are in the [0.0, 1.0] range.

        Returns:
            jittered_color (tuple[double]): a tuple of 3 elements, containing the RGB values of the
                color after being jittered. The values in the list are in the [0.0, 1.0] range.
        """
        color = mplc.to_rgb(color)
        vec = np.random.rand(3)
        # better to do it in another color space
        vec = vec / np.linalg.norm(vec) * 0.5
        res = np.clip(vec + color, 0, 1)
        return tuple(res)

    def _create_grayscale_image(self, mask=None):
        """
        Create a grayscale version of the original image.
        The colors in masked area, if given, will be kept.
        """
        img_bw = self.img.astype("f4").mean(axis=2)
        img_bw = np.stack([img_bw] * 3, axis=2)
        if mask is not None:
            img_bw[mask] = self.img[mask]
        return img_bw

    def _change_color_brightness(self, color, brightness_factor):
        """
        Depending on the brightness_factor, gives a lighter or darker color i.e. a color with
        less or more saturation than the original color.

        Args:
            color: color of the polygon. Refer to `matplotlib.colors` for a full list of
                formats that are accepted.
            brightness_factor (float): a value in [-1.0, 1.0] range. A lightness factor of
                0 will correspond to no change, a factor in [-1.0, 0) range will result in
                a darker color and a factor in (0, 1.0] range will result in a lighter color.

        Returns:
            modified_color (tuple[double]): a tuple containing the RGB values of the
                modified color. Each value in the tuple is in the [0.0, 1.0] range.
        """
        assert brightness_factor >= -1.0 and brightness_factor <= 1.0
        color = mplc.to_rgb(color)
        polygon_color = colorsys.rgb_to_hls(*mplc.to_rgb(color))
        modified_lightness = polygon_color[1] + (brightness_factor * polygon_color[1])
        modified_lightness = 0.0 if modified_lightness < 0.0 else modified_lightness
        modified_lightness = 1.0 if modified_lightness > 1.0 else modified_lightness
        modified_color = colorsys.hls_to_rgb(polygon_color[0], modified_lightness, polygon_color[2])
        return modified_color

    def get_output(self):
        """
        Returns:
            output (VisImage): the image output containing the visualizations added
            to the image.
        """
        return self.output


import requests
import re
from grepgithub import Hits, fetch_grep_app
import time
import random
from crawlab import save_item
from datetime import datetime

SOURCEGRAPH = "Sourcegraph"
GREPAPP = "Grep.app"

def extract_onion_addresses(text):
    onion_address_pattern = re.compile(r'[a-z2-7]{55}d\.onion')
    matches = re.findall(onion_address_pattern, text)
    return set(matches)

def search_sourcegraph(domain_name):
    headers =   {'User-Agent': 'Mozilla/5.0 (Windows NT 6.1; Win64; x64; rv:47.0) Gecko/20100101 Firefox/47.0'}
    url = f"https://sourcegraph.com/search/stream?q=context%3Aglobal%20{domain_name}%20count%3A2000%20fork%3Ayes%20&v=V3&t=standard&sm=0&dl=0&dk=html&dc=1&display=1500&cm=t"
    response = requests.get(url, headers=headers)
    addresses = extract_onion_addresses(response.text)
    onions = {}
    for address in addresses:
        onions[address] = "Sourcegraph"
    return onions

def search_grepapp(domain_name):
    next_page, hits, count = fetch_grep_app(page=1, query="d.onion")

    while next_page and next_page < 101:    # Does not paginate after 100
        time.sleep(random.uniform(1, 2))
        next_page, page_hits, _ = fetch_grep_app(page=next_page, query="d.onion")
        hits.merge(page_hits)
    print()

    onion_addresses = {}
    for repo, path_data in hits.hits.items():
        for path, lines in path_data.items():
            for line_num, line in lines.items():
                for address in extract_onion_addresses(line):
                    if address not in onion_addresses.keys():
                        onion_addresses[address] = set()
                    onion_addresses[address].add(repo)

    onions = {}
    for onion, repos in onion_addresses.items():
        repos = list(repos)
        repos_string = repos[0]
        if len(repos) > 1:
            for repo in repos[1:]:
                repos_string =  repos_string + ', ' + repo
        repos_string = 'Grep.app (' + repos_string + ')'
        onions[onion] = repos_string

    return onions

for source in [GREPAPP, SOURCEGRAPH]:

    print(f"Discovering onions in {source} ...\n")

    if source == SOURCEGRAPH:
        addresses = search_sourcegraph("d.onion")
    elif source == GREPAPP:
        addresses = search_grepapp("d.onion")

    for onion, advertiser in addresses.items():
        print(onion, advertiser)
        today = datetime.today().timestamp()
        save_item({
            'advertiser': advertiser,
            'onion': onion,
            'timestamp': today
        })

    print()
    print(f"Discovered [{len(addresses)}] in {source}!!!\n")

    print()
    print("#####################################################")
    print()


from PyQt5.QtCore import QPoint, QSettings, QSignalMapper, QSize, Qt, QFileInfo
from PyQt5.QtGui import QKeySequence
from PyQt5.QtWidgets import QAction, QApplication, QFileDialog, QMainWindow, QMdiArea, QMessageBox, qApp, QStyle

from plain_text_import_dialog import PlainTextImportDialog
from pyqtchild import MDIChildPlot

class MainWindow(QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()

        self.mdiArea = QMdiArea()
        self.mdiArea.setHorizontalScrollBarPolicy(Qt.ScrollBarAsNeeded)
        self.mdiArea.setVerticalScrollBarPolicy(Qt.ScrollBarAsNeeded)
        self.setCentralWidget(self.mdiArea)

        self.mdiArea.subWindowActivated.connect(self.update_menus)
        self.windowMapper = QSignalMapper(self)
        self.windowMapper.mapped.connect(self.set_active_sub_window)

        self.newAct = QAction(self)
        self.newAct.setIcon(self.style().standardIcon(QStyle.SP_FileIcon))
        # self.newAct.setIcon(QIcon(':/images/new.png'))
        self.newAct.setIconText('New')
        self.newAct.setShortcut(QKeySequence.New)
        self.newAct.setStatusTip('Create a new file')
        self.newAct.triggered.connect(self.new_file)

        self.openAct = QAction(self)
        self.openAct.setIcon(self.style().standardIcon(QStyle.SP_DirOpenIcon))
        self.openAct.setIconText('Open...')
        self.openAct.setShortcut(QKeySequence.Open)
        self.openAct.setStatusTip('Open an existing file')
        self.openAct.triggered.connect(self.open)

        self.saveAct = QAction(self)
        self.saveAct.setIcon(self.style().standardIcon(QStyle.SP_DialogSaveButton))
        self.saveAct.setIconText('Save')
        self.saveAct.setShortcut(QKeySequence.Save)
        self.saveAct.setStatusTip('Save the document to disk')
        self.saveAct.triggered.connect(self.save)

        self.saveAsAct = QAction(self)
        self.saveAsAct.setIconText('Save As...')
        self.saveAsAct.setShortcut(QKeySequence.SaveAs)
        self.saveAsAct.setStatusTip('Save the document under a new name')
        self.saveAsAct.triggered.connect(self.save_as)

        self.exitAct = QAction(self)
        self.exitAct.setIconText('Exit')
        self.exitAct.setShortcut(QKeySequence.Quit)
        self.exitAct.setStatusTip('Exit the application')
        self.exitAct.triggered.connect(qApp.closeAllWindows)

        self.cutAct = QAction(self)
        # self.cutAct.setIcon(QIcon(':/images/cut.png'))
        self.cutAct.setIconText('Cut')
        self.cutAct.setShortcut(QKeySequence.Cut)
        self.cutAct.setStatusTip("Cut the current selection's contents to the clipboard")
        self.cutAct.triggered.connect(self.cut)

        self.copyAct = QAction(self)
        # self.copyAct.setIcon(QIcon(':/images/copy.png'))
        self.copyAct.setIconText('Copy')
        self.copyAct.setShortcut(QKeySequence.Copy)
        self.copyAct.setStatusTip("Copy the current selection's contents to the clipboard")
        self.copyAct.triggered.connect(self.copy)

        self.pasteAct = QAction(self)
        # self.pasteAct.setIcon(QIcon(':/images/paste.png'))
        self.pasteAct.setIconText('Paste')
        self.pasteAct.setShortcut(QKeySequence.Paste)
        self.pasteAct.setStatusTip("Paste the clipboard's contents into the current selection")
        self.pasteAct.triggered.connect(self.paste)

        self.closeAct = QAction(self)
        self.closeAct.setIconText('Close')
        self.closeAct.setShortcut('Ctrl+W')
        self.closeAct.setStatusTip("Close the active window")
        self.closeAct.triggered.connect(self.mdiArea.closeActiveSubWindow)

        self.closeAllAct = QAction(self)
        self.closeAllAct.setIconText('Close All')
        self.closeAllAct.setStatusTip('Close all the windows')
        self.closeAllAct.triggered.connect(self.mdiArea.closeAllSubWindows)

        self.tileAct = QAction(self)
        self.tileAct.setIconText('Tile')
        self.tileAct.setStatusTip('Tile the windows')
        self.tileAct.triggered.connect(self.mdiArea.tileSubWindows)

        self.cascadeAct = QAction(self)
        self.cascadeAct.setIconText('Cascade')
        self.cascadeAct.setStatusTip('Cascade the windows')
        self.cascadeAct.triggered.connect(self.mdiArea.cascadeSubWindows)

        self.nextAct = QAction(self)
        self.nextAct.setIconText('Next')
        self.nextAct.setShortcut(QKeySequence.NextChild)
        self.nextAct.setStatusTip('Move the focus to the next window')
        self.nextAct.triggered.connect(self.mdiArea.activateNextSubWindow)

        self.previousAct = QAction(self)
        self.previousAct.setIconText('Previous')
        self.previousAct.setShortcut(QKeySequence.PreviousChild)
        self.previousAct.setStatusTip('Move the focus to the previous window')
        self.previousAct.triggered.connect(self.mdiArea.activatePreviousSubWindow)

        self.separatorAct = QAction(self)
        self.separatorAct.setSeparator(True)

        self.aboutAct = QAction(self)
        self.aboutAct.setIconText('About')
        self.aboutAct.setStatusTip("Show the application's About box")
        self.aboutAct.triggered.connect(self.about)

        self.aboutQtAct = QAction(self)
        self.aboutQtAct.setIconText('About Qt')
        self.aboutQtAct.setStatusTip("Show the Qt library's About box")
        self.aboutQtAct.triggered.connect(qApp.aboutQt)

        self.fileMenu = self.menuBar().addMenu('File')
        self.fileMenu.addAction(self.newAct)
        self.fileMenu.addAction(self.openAct)
        self.fileMenu.addAction(self.saveAct)
        self.fileMenu.addAction(self.saveAsAct)
        self.fileMenu.addSeparator()
        action = self.fileMenu.addAction('Switch layout direction')
        action.triggered.connect(self.switch_layout_direction)
        self.fileMenu.addAction(self.exitAct)

        self.editMenu = self.menuBar().addMenu('Edit')
        self.editMenu.addAction(self.cutAct)
        self.editMenu.addAction(self.copyAct)
        self.editMenu.addAction(self.pasteAct)

        self.windowMenu = self.menuBar().addMenu('Window')
        self.update_window_menu()
        self.windowMenu.aboutToShow.connect(self.update_window_menu)

        self.menuBar().addSeparator()

        self.helpMenu = self.menuBar().addMenu('Help')
        self.helpMenu.addAction(self.aboutAct)
        self.helpMenu.addAction(self.aboutQtAct)

        self.file_tool_bar = self.addToolBar('File')
        self.file_tool_bar.addAction(self.newAct)
        self.file_tool_bar.addAction(self.openAct)
        self.file_tool_bar.addAction(self.saveAct)

        self.edit_tool_bar = self.addToolBar('Edit')
        self.edit_tool_bar.addAction(self.cutAct)
        self.edit_tool_bar.addAction(self.copyAct)
        self.edit_tool_bar.addAction(self.pasteAct)

        self.statusBar().showMessage('Ready')

        self.update_menus()

        self.settings = QSettings('SavSoft', 'Combiner')

        self.last_directory: str = ''
        self.read_settings()

        self.setWindowTitle('MDI')
        self.setUnifiedTitleAndToolBarOnMac(True)

    def closeEvent(self, event):
        self.mdiArea.closeAllSubWindows()
        if self.active_mdi_child():
            event.ignore()
        else:
            self.write_settings()
            event.accept()

    def new_file(self):
        child = self.create_mdi_child()
        child.new_file()
        child.show()

    def open(self):
        filters = ['IRTECON files (*.grd)', 'Plain text files (*.csv *.tsv *.dat *.txt)']
        file_name, _ = QFileDialog.getOpenFileName(self, filter=';;'.join(filters),
                                                   directory=self.last_directory,
                                                   options=QFileDialog.DontUseNativeDialog)
        if file_name:
            child = self.create_mdi_child()
            if QFileInfo(file_name).suffix() == 'grd':
                if child.load_irtecon_file(file_name):
                    self.statusBar().showMessage('File loaded', 2000)
                    self.last_directory = QFileInfo(file_name).dir().absolutePath()
                    child.show()
                else:
                    child.close()
            else:
                print(PlainTextImportDialog(file_name).exec())
                raise NotImplementedError

    def save(self):
        if self.active_mdi_child() and self.active_mdi_child().save():
            self.statusBar().showMessage('File saved', 2000)

    def save_as(self):
        if self.active_mdi_child() and self.active_mdi_child().save_as():
            self.statusBar().showMessage('File saved', 2000)

    def cut(self):
        if self.active_mdi_child():
            self.active_mdi_child().cut()

    def copy(self):
        if self.active_mdi_child():
            self.active_mdi_child().copy()

    def paste(self):
        if self.active_mdi_child():
            self.active_mdi_child().paste()

    def about(self):
        QMessageBox.about(self, 'About MDI',
                          'The <b>MDI</b> example demonstrates how to write multiple '
                          'document interface applications using Qt.')

    def update_menus(self):
        has_mdi_child = (self.active_mdi_child() is not None)
        self.saveAct.setEnabled(has_mdi_child)
        self.saveAsAct.setEnabled(has_mdi_child)
        self.pasteAct.setEnabled(has_mdi_child)
        self.closeAct.setEnabled(has_mdi_child)
        self.closeAllAct.setEnabled(has_mdi_child)
        self.tileAct.setEnabled(has_mdi_child)
        self.cascadeAct.setEnabled(has_mdi_child)
        self.nextAct.setEnabled(has_mdi_child)
        self.previousAct.setEnabled(has_mdi_child)
        self.separatorAct.setVisible(has_mdi_child)

    def update_window_menu(self):
        self.windowMenu.clear()
        self.windowMenu.addAction(self.closeAct)
        self.windowMenu.addAction(self.closeAllAct)
        self.windowMenu.addSeparator()
        self.windowMenu.addAction(self.tileAct)
        self.windowMenu.addAction(self.cascadeAct)
        self.windowMenu.addSeparator()
        self.windowMenu.addAction(self.nextAct)
        self.windowMenu.addAction(self.previousAct)
        self.windowMenu.addAction(self.separatorAct)

        windows = self.mdiArea.subWindowList()
        self.separatorAct.setVisible(len(windows) != 0)

        for i, window in enumerate(windows):
            child = window.widget()

            text = f'{i + 1:d} {child.user_friendly_current_file()}'
            if i < 9:
                text = '' + text

            action = self.windowMenu.addAction(text)
            action.setCheckable(True)
            action.setChecked(child == self.active_mdi_child())
            action.triggered.connect(self.windowMapper.map)
            self.windowMapper.setMapping(action, window)

    def create_mdi_child(self):
        child = MDIChildPlot()
        self.mdiArea.addSubWindow(child)
        return child

    def read_settings(self):
        pos = self.settings.value('pos', QPoint(200, 200))
        self.move(pos)
        size = self.settings.value('size', QSize(400, 400))
        self.resize(size)
        self.last_directory = self.settings.value('directory', '')

    def write_settings(self):
        self.settings.setValue('pos', self.pos())
        self.settings.setValue('size', self.size())
        self.settings.setValue('directory', self.last_directory)

    def active_mdi_child(self):
        active_sub_window = self.mdiArea.activeSubWindow()
        if active_sub_window:
            return active_sub_window.widget()
        return None

    def switch_layout_direction(self):
        if self.layoutDirection() == Qt.LeftToRight:
            qApp.setLayoutDirection(Qt.RightToLeft)
        else:
            qApp.setLayoutDirection(Qt.LeftToRight)

    def set_active_sub_window(self, window):
        if window:
            self.mdiArea.setActiveSubWindow(window)

if __name__ == '__main__':
    import sys

    app = QApplication(sys.argv)
    mainWin = MainWindow()
    mainWin.show()
    sys.exit(app.exec_())


# -*- coding: utf-8 -*-
from flask import Flask
from flask import jsonify, request, make_response
from flask_sqlalchemy import SQLAlchemy

from data import db_session

# Importing all models to work correctly in app.py
from data.transport_type import TransportTypes
from data.regions_to_couriers import CouriersToRegions
from data.orders import Orders
from data.regions import Regions
from data.delivery import Delivery
from data.couriers import *

# Other packages
import os
from pprint import pprint
import datetime

app = Flask(__name__)

@app.route('/couriers', methods=['POST'])
def post_couriers():
    '''
    1) POST /couriers
    '''
    try:
        data = request.get_json()
        couriers_list = data['data']
    except BaseException as be:
        return make_response(
            jsonify({'validation_error': 'wrong json data'}), 400)

    db_sess = db_session.create_session()

    validation_error = {"validation_error": {
        "couriers": []
    }
    }

    valid_couriers = []
    for courier_json in couriers_list:

        if Couriers.validate_courier_json(courier_json, logger=app.logger):
            courier_json['working_hours'] = convert_wh_hours_to_str(
                courier_json['working_hours'])
            courier_json['courier_type'] = db_sess.query(
                TransportTypes).filter(
                    TransportTypes.type_name
                    == courier_json['courier_type']).first()
        else:
            validation_error["validation_error"]['couriers'].append(
                {'id': courier_json['courier_id']})
            continue

        added_regions = []
        regions = []
        for region_name in courier_json['regions']:
            if db_sess.query(Regions).filter(Regions.region_code ==
                                             region_name).first(
            ) is None and region_name not in added_regions:
                reg = Regions(region_name)
                db_sess.add(reg)
                # db_sess.commit()
                added_regions.append(reg)
            else:
                reg = db_sess.query(Regions).filter(
                    Regions.region_code == region_name).first()
                regions.append(reg)

        courier = Couriers()
        courier.courier_id = courier_json['courier_id']
        courier.courier_type_id = courier_json['courier_type'].type_id
        courier.regions = added_regions + regions
        courier.working_hours = courier_json['working_hours']
        valid_couriers.append(courier)

    if len(validation_error["validation_error"]['couriers']) != 0:
        return make_response(jsonify(validation_error), '400 Bad Request')

    response = {'couriers': []}
    app.logger.info(f'Adding the following Couriers:')
    for cour in valid_couriers:
        app.logger.info(cour)
        db_sess.add(cour)
        response['couriers'].append({"id": cour.courier_id})
        db_sess.commit()

    return make_response(jsonify(response), '201 Created')

@app.route('/couriers/<int:courier_id>', methods=['PATCH'])
def patch_couriers(courier_id):
    '''
    2) PATCH /couriers/$courier_id
    '''
    if isinstance(courier_id, int):
        if courier_id > 0:
            data = request.get_json()
            db_sess = db_session.create_session()

            if Couriers.validate_patch(data, logger=app.logger):

                cour = db_sess.query(Couriers).filter(
                    Couriers.courier_id.like(courier_id)).first()

                if cour:
                    app.logger.info(
                        f'Changing {cour} with following params:\n{data.keys()}')
                    for key in data:
                        # TODO Bad perfomance, change with getattr/setattr(?) or
                        # property function
                        if isinstance(data[key], str):

                            new_type = db_sess.query(TransportTypes).filter(
                                TransportTypes.type_name == data[key]).first()
                            if new_type:
                                cour.change_cour_type(new_type, db_sess)

                        elif isinstance(data[key][0], int):

                            new_regions = data[key]
                            cour.change_cour_regions(new_regions, db_sess)

                        elif isinstance(data[key][0], str):

                            new_working_hours = data[key]
                            cour.change_cour_work_hours(
                                new_working_hours, db_sess)

                    return make_response(
                        jsonify(
                            Couriers.make_courier_response(
                                cour)),
                        '200 OK')
                else:
                    app.logger.info(f'No courier with id {courier_id}')
                    return '', '404 Not found'

    app.logger.info(f'Requets courier with id {courier_id} is bad')
    return '', '400 Bad request'

@ app.route('/orders', methods=['POST'])
def set_orders():
    '''
    3) POST /orders
    '''
    try:
        data = request.get_json()
        orders_list = data['data']
    except BaseException as be:
        app.logger.exception(f'Wrong JSON data passed: {be}')
        return make_response('400 Bad request')

    db_sess = db_session.create_session()
    ve = False
    validation_error = {"validation_error": {
        "orders": []
    }
    }

    valid_orders = []
    for order_json in orders_list:

        if Orders.validate_order_json(order_json, logger=app.logger):
            if ve:
                continue
            order = Orders()
            order.order_id = order_json['order_id']
            order.weight = order_json['weight']

            query_region = db_sess.query(Regions).filter(
                Regions.region_code == order_json['region']).first()
            if query_region is None:
                reg = Regions(region_code=order_json['region'])
            else:
                reg = query_region

            db_sess.add(reg)

            order.region = reg
            order.delivery_time = convert_wh_hours_to_str(
                order_json['delivery_hours'])

            valid_orders.append(order)
        else:
            try:
                validation_error["validation_error"]['orders'].append(
                    {'id': order_json['order_id']}
                )
            except KeyError as ke:
                validation_error["validation_error"]['schema_error'] = order_json
            db_sess.rollback()
            ve = True
            continue

    if ve:
        app.logger.info(f'Passes not valid data: {validation_error}')
        return make_response(jsonify(validation_error), '400 Bad request')
    else:
        response = {'orders': []}

        for order in valid_orders:
            db_sess.add(order)
            response['orders'].append({"id": order.order_id})
            db_sess.commit()
        return make_response(jsonify(response), '201 Created')

@ app.route('/orders/assign', methods=['POST'])
def assign_orders():
    '''
    4) POST /orders/assign
    '''
    data = request.get_json()
    db_sess = db_session.create_session()

    if data is not None:
        if Couriers.validate_assigment(data):
            courier_id = data['courier_id']

            cour = db_sess.query(Couriers).filter(
                Couriers.courier_id == courier_id).first()

            if cour is not None:
                if not cour.is_working():

                    regions = [reg.region_id for reg in cour.regions]

                    orders = db_sess.query(Orders).filter(
                        Orders.order_complete_time == None).filter(
                            Orders.delivery_id == None).filter(
                                Orders.weight <= cour.courier_type.type_weight,
                                Orders.region_id.in_(regions)).all()

                    for_deliver = []
                    for order in orders:
                        if cour.could_he_take(order):
                            for_deliver.append(order)

                    if len(for_deliver) > 0:
                        response = {
                            "orders": [],
                            "assign_time": "2021-01-10T09:32:14.42Z"
                        }

                        delivery = Delivery()
                        for order in for_deliver:
                            response['orders'].append({'id': order.order_id})
                            # --- DB ---

                        delivery.delivery_courier = cour
                        delivery.assigned_courier_type_id = cour.courier_type_id
                        delivery.assigned_courier_type = cour.courier_type
                        assign_time = datetime.datetime.now()
                        delivery.assign_time = assign_time
                        delivery.orders_in_delivery = for_deliver
                        app.logger.info(
                            f'{cour} got {delivery} with {for_deliver}')
                        db_sess.add(delivery)
                        db_sess.commit()
                        # --- END ---

                        response['assign_time'] = assign_time.isoformat()[
                            :22] + "Z"
                        return make_response(jsonify(response), '200 OK')
                    return make_response(jsonify({'orders': []}), '200 OK')
                else:
                    try:
                        current_delivery = cour.get_current_delivery()
                        app.logger.info(
                            f'Courier {courier_id} have {current_delivery}')
                        orders_list = db_sess.query(
                            Orders.order_id).filter(
                                Orders.delivery_id
                                == current_delivery.delivery_id).filter(
                                    Orders.order_complete_time == None
                        ).all()
                        response = {
                            'orders':
                            [{'id': order.order_id} for order in orders_list],
                            'assign_time': current_delivery.get_str_assign_time()}
                        return make_response(jsonify(response), '200 OK')
                    except BaseException:
                        app.logger.exception(
                            '/assign can`t show any information.')
                        return '', '400 Bad request'
    return '', '400 Bad request'

@ app.route('/orders/complete', methods=['POST'])
def complete_order():
    '''
    5) POST /orders/complete
    '''
    data = request.get_json()
    db_sess = db_session.create_session()
    if data:
        if Orders.validate_complete(data, db_sess, logger=app.logger):
            order = db_sess.query(Orders).filter(
                Orders.order_id == data['order_id']).first()

            response = {'order_id': order.order_id}
            complete_time = datetime.datetime.fromisoformat(
                data['complete_time'][:22] + '0')

            if not order.is_completed():
                order.order_complete_time = complete_time
            else:
                return make_response(jsonify(response), '200 OK')

            db_sess.add(order)

            delivery = order.delivery
            if delivery.is_completed():
                delivery.delivery_complete_time = complete_time

            db_sess.add(delivery)
            db_sess.commit()
            return make_response(jsonify(response), '200 OK')
    return '', '400 Bad request'

@ app.route('/couriers/<int:courier_id>', methods=['GET'])
def get_courier_info(courier_id):
    '''
    6) GET /couriers/$courier_id
    '''
    db_sess = db_session.create_session()
    if courier_id >= 1:

        courier = db_sess.query(Couriers).filter(
            Couriers.courier_id == courier_id).first()

        if courier:
            delivery = db_sess.query(Delivery)\
                .filter(Delivery.delivery_complete_time is not None)\
                .filter(Delivery.courier_id == courier_id).all()

            # Counting Earnings of courier
            earnings = 0
            for d in delivery:
                earnings += d.count_earning()
                courier.earnings = earnings

            add_response = {'earnings': earnings}
            # Counting rationg of courier
            if len(delivery) > 0:
                regions_avg = Regions.count_avg_time_from_orders(delivery)
                if len(regions_avg) != 0:
                    rating = Couriers.count_rating_from_regions_avg(
                        regions_avg)

                    courier.rating = rating
                    add_response['rating'] = rating

            db_sess.add(courier)
            response = Couriers.make_courier_response(
                courier, **add_response)
            db_sess.commit()
            return make_response(jsonify(response), 200)
        else:
            app.logger.info(f'No courier with id {courier_id}')

    return '', '404 Not found'

def main(debug=False):
    # Preparing db and run app
    db_session.global_init_sqlite('db.sqlite')
    app.config.from_envvar('API_CONFIG')
    app.run(host='0.0.0.0', port=8080)

if __name__ == '__main__':
    main(debug=True)


import sys
import os
from PyQt5.QtCore import QObject, QUrl
from PyQt5.QtGui import QDesktopServices
from typing import List, cast

from QD.Event import CallFunctionEvent
from QD.FlameProfiler import pyqtSlot
from QD.Math.Vector import Vector
from QD.Scene.Selection import Selection
from QD.Scene.Iterator.BreadthFirstIterator import BreadthFirstIterator
from QD.Operations.GroupedOperation import GroupedOperation
from QD.Operations.RemoveSceneNodeOperation import RemoveSceneNodeOperation
from QD.Operations.TranslateOperation import TranslateOperation
from QD.Resources import Resources

import qidi.QIDIApplication
from qidi.Operations.SetParentOperation import SetParentOperation
from qidi.MultiplyObjectsJob import MultiplyObjectsJob
from qidi.Settings.SetObjectExtruderOperation import SetObjectExtruderOperation
from qidi.Settings.ExtruderManager import ExtruderManager

from qidi.Operations.SetBuildPlateNumberOperation import SetBuildPlateNumberOperation

from QD.Logger import Logger
from QD.Scene.SceneNode import SceneNode

class QIDIActions(QObject):
    def __init__(self, parent: QObject = None) -> None:
        super().__init__(parent)

    @pyqtSlot()
    def openDocumentation(self) -> None:
        # Starting a web browser from a signal handler connected to a menu will crash on windows.
        # So instead, defer the call to the next run of the event loop, since that does work.
        # Note that weirdly enough, only signal handlers that open a web browser fail like that.
        event = CallFunctionEvent(self._openUrl, [QUrl("https://qidi.com/en/resources/manuals/software")], {})
        qidi.QIDIApplication.QIDIApplication.getInstance().functionEvent(event)

    @pyqtSlot()
    def openBugReportPage(self) -> None:
        event = CallFunctionEvent(self._openUrl, [QUrl("https://github.com/QIDI/QIDI/issues/new/choose")], {})
        qidi.QIDIApplication.QIDIApplication.getInstance().functionEvent(event)

    @pyqtSlot()
    def factorySetting(self):
        Resources.factoryReset()
        #sys.exit(1)
        os._exit(0)

    @pyqtSlot()
    def homeCamera(self) -> None:
        """Reset camera position and direction to default"""

        scene = qidi.QIDIApplication.QIDIApplication.getInstance().getController().getScene()
        camera = scene.getActiveCamera()
        if camera:
            diagonal_size = qidi.QIDIApplication.QIDIApplication.getInstance().getBuildVolume().getDiagonalSize()
            camera.setPosition(Vector(-80, 250, 700) * diagonal_size / 375)
            camera.setPerspective(True)
            camera.lookAt(Vector(0, 0, 0))

    @pyqtSlot()
    def centerSelection(self) -> None:
        """Center all objects in the selection"""

        operation = GroupedOperation()
        for node in Selection.getAllSelectedObjects():
            current_node = node
            parent_node = current_node.getParent()
            while parent_node and parent_node.callDecoration("isGroup"):
                current_node = parent_node
                parent_node = current_node.getParent()

            #   This was formerly done with SetTransformOperation but because of
            #   unpredictable matrix deconstruction it was possible that mirrors
            #   could manifest as rotations. Centering is therefore done by
            #   moving the node to negative whatever its position is:
            center_operation = TranslateOperation(current_node, -current_node._position)
            operation.addOperation(center_operation)
        operation.push()

    @pyqtSlot(int,int,bool)
    def multiplySelection(self, count: int,spacing: int,checked:bool) -> None:
        """Multiply all objects in the selection

        :param count: The number of times to multiply the selection.
        """
        min_offset = qidi.QIDIApplication.QIDIApplication.getInstance().getBuildVolume().getEdgeDisallowedSize() + 2  # Allow for some rounding errors
        job = MultiplyObjectsJob(Selection.getAllSelectedObjects(), count,spacing,checked, min_offset = max(min_offset, 8))
        job.start()

    @pyqtSlot()
    def deleteSelection(self) -> None:
        """Delete all selected objects."""

        if not qidi.QIDIApplication.QIDIApplication.getInstance().getController().getToolsEnabled():
            return

        removed_group_nodes = [] #type: List[SceneNode]
        op = GroupedOperation()
        nodes = Selection.getAllSelectedObjects()
        for node in nodes:
            op.addOperation(RemoveSceneNodeOperation(node))
            group_node = node.getParent()
            if group_node and group_node.callDecoration("isGroup") and group_node not in removed_group_nodes:
                remaining_nodes_in_group = list(set(group_node.getChildren()) - set(nodes))
                if len(remaining_nodes_in_group) == 1:
                    removed_group_nodes.append(group_node)
                    op.addOperation(SetParentOperation(remaining_nodes_in_group[0], group_node.getParent()))
                    op.addOperation(RemoveSceneNodeOperation(group_node))

            # Reset the print information
            qidi.QIDIApplication.QIDIApplication.getInstance().getController().getScene().sceneChanged.emit(node)

        op.push()

    @pyqtSlot(str)
    def setExtruderForSelection(self, extruder_id: str) -> None:
        """Set the extruder that should be used to print the selection.

        :param extruder_id: The ID of the extruder stack to use for the selected objects.
        """

        operation = GroupedOperation()

        nodes_to_change = []
        for node in Selection.getAllSelectedObjects():
            # If the node is a group, apply the active extruder to all children of the group.
            if node.callDecoration("isGroup"):
                for grouped_node in BreadthFirstIterator(node): #type: ignore #Ignore type error because iter() should get called automatically by Python syntax.
                    if grouped_node.callDecoration("getActiveExtruder") == extruder_id:
                        continue

                    if grouped_node.callDecoration("isGroup"):
                        continue

                    nodes_to_change.append(grouped_node)
                continue

            # Do not change any nodes that already have the right extruder set.
            if node.callDecoration("getActiveExtruder") == extruder_id:
                continue

            nodes_to_change.append(node)

        if not nodes_to_change:
            # If there are no changes to make, we still need to reset the selected extruders.
            # This is a workaround for checked menu items being deselected while still being
            # selected.
            ExtruderManager.getInstance().resetSelectedObjectExtruders()
            return

        for node in nodes_to_change:
            operation.addOperation(SetObjectExtruderOperation(node, extruder_id))
        operation.push()

    @pyqtSlot(int)
    def setBuildPlateForSelection(self, build_plate_nr: int) -> None:
        Logger.log("d", "Setting build plate number... %d" % build_plate_nr)
        operation = GroupedOperation()

        root = qidi.QIDIApplication.QIDIApplication.getInstance().getController().getScene().getRoot()

        nodes_to_change = []  # type: List[SceneNode]
        for node in Selection.getAllSelectedObjects():
            parent_node = node  # Find the parent node to change instead
            while parent_node.getParent() != root:
                parent_node = cast(SceneNode, parent_node.getParent())

            for single_node in BreadthFirstIterator(parent_node):  # type: ignore #Ignore type error because iter() should get called automatically by Python syntax.
                nodes_to_change.append(single_node)

        if not nodes_to_change:
            Logger.log("d", "Nothing to change.")
            return

        for node in nodes_to_change:
            operation.addOperation(SetBuildPlateNumberOperation(node, build_plate_nr))
        operation.push()

        Selection.clear()

    def _openUrl(self, url: QUrl) -> None:
        QDesktopServices.openUrl(url)


"""Download dataset : https://www.robots.ox.ac.uk/~vgg/data/pets/
"""

import os
from typing import List, Tuple, Union

import albumentations as A
import numpy as np
import torch
from albumentations.pytorch import ToTensorV2
from PIL import Image
from torch.utils.data import Dataset

def make_transform(image_set: str) -> A.Compose:
    """Creates a transformation pipeline based on the given image set ("train" or "test").

    Args:
    - image_set (str): The set of images for which transformations need to be made ("train" or otherwise).

    Returns:
    - A.Compose: A composition of various augmentations.
    """
    # Mean and Standard deviation for normalization
    mean = (0.485, 0.456, 0.406)
    std = (0.229, 0.224, 0.225)

    # Augmentations for training set
    if image_set == "train":
        return A.Compose(
            [
                A.Rotate(limit=35),  # Rotation augmentation
                A.HorizontalFlip(),  # Horizontal flip augmentation
                A.VerticalFlip(),  # Vertical flip augmentation
                A.Resize(height=256, width=256),  # Resizing the images
                A.Normalize(mean=mean, std=std, max_pixel_value=255.0),  # Normalizing images
                ToTensorV2(),  # Converting images to PyTorch tensors
            ],
        )
    # Augmentations for validation or test set
    else:
        return A.Compose(
            [
                A.Resize(height=256, width=256),  # Resizing the images
                A.Normalize(mean=mean, std=std, max_pixel_value=255.0),  # Normalizing images
                ToTensorV2(),  # Converting images to PyTorch tensors
            ],
        )

class OxfordIiitPets(Dataset):
    """PyTorch Dataset for Oxford-IIIT Pet Dataset.

    Args:
    - image_dir (str): Path to the directory containing images.
    - images_list (List[str]): List of paths to the images.
    - mask_dir (str): Path to the directory containing masks.
    - transforms (Union[A.Compose, None]): Transformations to apply on the images and masks.
    """

    def __init__(
        self,
        image_dir: str,
        images_list: List[str],
        mask_dir: str,
        transforms: Union[A.Compose, None] = None,
    ):
        self.image_dir = image_dir
        self.images_list = images_list
        self.mask_dir = mask_dir
        self.transforms = transforms

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

    def __getitem__(self, idx) -> Tuple[torch.Tensor, torch.Tensor]:
        """Fetches an image and its corresponding mask based on the index provided.

        Args:
        - idx (int): Index of the image and mask pair to fetch.

        Returns:
        - tuple: A tuple containing the transformed image and mask.
        """
        ip = os.path.join(self.image_dir, self.images_list[idx])
        mp = os.path.join(self.mask_dir, self.images_list[idx].replace("jpg", "png"))

        image = Image.open(ip).convert("RGB")
        mask = Image.open(mp)

        mask = np.array(mask)
        mask[mask == 3] = 0  # background
        mask[mask == 2] = 2  # border
        mask[mask == 1] = 1  # pet

        if self.transforms:
            transformed = self.transforms(image=np.array(image), mask=mask)
            image = transformed["image"]
            mask = transformed["mask"]

        return image, mask

# if __name__=="__main__":
#     import glob

#     import matplotlib.pyplot as plt
#     from torch.utils.data import DataLoader

#     image_dir = "sample_dataset/oxford_iiit_pets/images"
#     mask_dir = "sample_dataset/oxford_iiit_pets/annotations/trimaps"

#     image_dir = r"C:\Users\anant\Downloads\images.tar\images"
#     mask_dir = r"C:\Users\anant\Downloads\oxford_iiit_pets\annotations.tar\annotations\trimaps"
#     images_list = []
#     for ext in [".jpg", ".jpeg", ".png"]: # filtering images
#         files = glob.glob(f"{image_dir}/*{ext}")
#         images_list += [os.path.basename(f) for f in files]

#     dset = OxfordIiitPets(
#         mask_dir=mask_dir,
#         image_dir=image_dir,
#         images_list=images_list,
#         transforms=make_transform(image_set="val"),
#     )

#     loader = DataLoader(dset, batch_size=1)
#     for x, y in loader:
#         print(x.shape, y.shape)
#         plt.subplot(1, 2, 1)
#         plt.imshow(x[0].permute(1, 2, 0))

#         plt.subplot(1, 2, 2)
#         plt.imshow(y[0])

#         plt.show()


import numpy as np
import matplotlib.pyplot as plt
from tri_functions import *
import os
from matplotlib import animation
import time
from mpl_toolkits.axes_grid1 import Divider, Size

@jit(nopython=True)
def get_e_stack(vi,vj):
    e_stack = np.zeros((vi.shape[0],3,2,2))
    for i in range(vi.shape[0]):
        for j in range(3):
            e_stack[i,j] = np.row_stack((vi[i,j],vj[i,j]))
    return e_stack

def plot_mesh(ax,v,neigh,Lx,Ly,color="black",alpha=1):
    v = v[~np.isnan(v[:,0])]
    neigh = neigh[neigh[:,0]!=-1]
    v_neigh = tri_call3(v, neigh)
    v_CCW = roll3(v_neigh, 1)
    vn = get_vn(v, v_CCW, Lx, Ly)
    vi = np.expand_dims(v,1)*np.ones_like(vn)
    vj = vi-vn
    e_stack = get_e_stack(vi,vj)
    ax.plot(e_stack[:,:,:,0].reshape(-1,2).T,e_stack[:,:,:,1].reshape(-1,2).T,color=color,alpha=alpha)
    # for i in range(vi.shape[0]):
    #     for j in range(3):
    #         e = np.row_stack((vi[i,j],vj[i,j]))
    #         ax.plot(e[:,0],e[:,1],color="black")
    ax.set(xlim=(0,Lx),ylim=(0,Ly),aspect=1)

def plot_tri(ax,tcnt,Lx,Ly,color="black"):
    tcnt = tcnt[~np.isnan(tcnt[:,0])]
    for i in range(3):
        cnts = np.array((tcnt[:, i], tcnt[:, np.mod(i + 1, 3)]))
        cnts[1] = cnts[0] + per(cnts[1]-cnts[0],Lx,Ly)
        ax.plot(cnts[:,:,0],cnts[:,:,1],color=color)

def fix_axes_size_incm(axew, axeh,fig,ax,pad=0.8):
    axew = axew*pad
    axeh = axeh*pad

    #lets use the tight layout function to get a good padding size for our axes labels.
    # fig = plt.gcf()
    # ax = plt.gca()
    # fig.tight_layout()
    #obtain the current ratio values for padding and fix size
    oldw, oldh = fig.get_size_inches()
    l = ax.figure.subplotpars.left
    r = ax.figure.subplotpars.right
    t = ax.figure.subplotpars.top
    b = ax.figure.subplotpars.bottom

    #work out what the new  ratio values for padding are, and the new fig size.
    neww = axew+oldw*(1-r+l)
    newh = axeh+oldh*(1-t+b)
    newr = r*oldw/neww
    newl = l*oldw/neww
    newt = t*oldh/newh
    newb = b*oldh/newh

    #right(top) padding, fixed axes size, left(bottom) pading
    hori = [Size.Scaled(newr), Size.Fixed(axew), Size.Scaled(newl)]
    vert = [Size.Scaled(newt), Size.Fixed(axeh), Size.Scaled(newb)]

    divider = Divider(fig, (0.0, 0.0, 1., 1.), hori, vert, aspect=False)
    # the width and height of the rectangle is ignored.

    ax.set_axes_locator(divider.new_locator(nx=1, ny=1))

    #we need to resize the figure now, as we have may have made our axes bigger than in.
    fig.set_size_inches(neww,newh)

def animate_mesh(v_save,neigh_save,L,n_frames=100, file_name=None, dir_name="plots",scale=10):
    if L.shape == L.size:
        L = np.ones((v_save.shape[0],1))*np.expand_dims(L,0)

    if not os.path.exists(dir_name):
        os.makedirs(dir_name)

    fig = plt.figure(figsize=(scale,scale*L[:,1].max()/L[:,0].max()))
    sf = scale/L[:,0].max()
    ax1 = fig.add_subplot(1, 1, 1)

    skip = int((v_save.shape[0]) / n_frames)

    def animate(i):

        ax1.cla()

        v = v_save[i*skip]
        neigh = neigh_save[i*skip]
        Lx,Ly = L[i*skip]
        plot_mesh(ax1,v,neigh,Lx,Ly)
        xlim = (0, Lx)
        ylim = (0, Ly)
        # ax1.set_facecolor('black')

        ax1.set(aspect=1, xlim=xlim, ylim=ylim)
        fix_axes_size_incm(Lx*sf,Ly*sf,fig,ax1)

    Writer = animation.writers["ffmpeg"]
    writer = Writer(fps=20, bitrate=500)
    if file_name is None:
        file_name = "animation %d" % time.time()
    an = animation.FuncAnimation(fig, animate, frames=n_frames, interval=200)
    an.save("%s/%s.mp4" % (dir_name, file_name), writer=writer, dpi=264)


import sys
from skimage.transform.radon_transform import convert_to_float
sys.path.append('/content/drive/MyDrive/FSBD')
from utils.config import conf
sys.path.append('/content/drive/MyDrive/FSBD/data')

import util
from cutout import Cutout
from tr_dataset import TrBboxDataset

import torch as t
from skimage import transform as sktsf
from torchvision import transforms as tvtsf
from torch.utils.data import DataLoader
import torch
import torchvision.transforms as T
from random import randint
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from PIL import Image,ImageDraw
import cv2

def inverse_normalize(img):
  if conf.caffe_pretrain:
      img = img + (np.array([122.7717, 115.9465, 102.9801]).reshape(3, 1, 1))
      return img[::-1, :, :]
  # approximate un-normalize for visualize
  return (img * 0.225 + 0.45).clip(min=0, max=1) * 255

def pytorch_normalze(img):

  normalize = tvtsf.Normalize(mean=[0.485, 0.456, 0.406],
                              std=[0.229, 0.224, 0.225])
  img = normalize(t.from_numpy(img))
  return img.numpy()

def caffe_normalize(img):

  img = img[[2, 1, 0], :, :]  # RGB-BGR
  img = img * 255
  mean = np.array([122.7717, 115.9465, 102.9801]).reshape(3, 1, 1)
  img = (img - mean).astype(np.float32, copy=True)
  return img

def preprocess(img, min_size=600, max_size=1000):

  C, H, W = img.shape
  scale1 = min_size / min(H, W)
  scale2 = max_size / max(H, W)
  scale = min(scale1, scale2)
  img = img / 255.
  img = sktsf.resize(img, (C, H * scale, W * scale), mode='reflect',anti_aliasing=False)
  # both the longer and shorter should be less than
  # max_size and min_size
  if conf.caffe_pretrain:
      normalize = caffe_normalize
  else:
      normalize = pytorch_normalze
  return normalize(img)

class Transform(object):

  def __init__(self, min_size=600, max_size=1000):
    self.min_size = min_size
    self.max_size = max_size

  def __call__(self, in_data):
    img, bbox, label = in_data
    _, H, W = img.shape
    img = preprocess(img, self.min_size, self.max_size)
    _, o_H, o_W = img.shape
    scale = o_H / H
    bbox = util.resize_bbox(bbox, (H, W), (o_H, o_W))

    # horizontally flip
    img, params = util.random_flip(
        img, x_random=True, return_param=True)
    bbox = util.flip_bbox(
        bbox, (o_H, o_W), x_flip=params['x_flip'])

    return img, bbox, label, scale

class Dataset:

  def __init__(self, split = 'beta'):
    self.db = TrBboxDataset( split = split)
    self.tsf = Transform(conf.min_size, conf.max_size)

  def __getitem__(self, idx):
    ori_img, bbox, label, fname = self.db.__getitem__(idx)

    img, bbox, label, scale = self.tsf((ori_img, bbox, label))

    return img.copy(), bbox.copy(), label.copy(), scale, fname

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

class TestDataset:
  def __init__(self,split = 'beta'):
      self.db = TrBboxDataset(split = split)

  def __getitem__(self, idx):
      ori_img, bbox, label, fname = self.db.__getitem__(idx)
      img = preprocess(ori_img)
      return img, ori_img.shape[1:], bbox, label, fname

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

class CDDataset:

  def __init__(self, split = 'beta', transforms = False):
    self.db = TrBboxDataset(split = split)
    self.tsf = Transform(conf.min_size, conf.max_size)
    self.transforms = transforms

  def __getitem__(self, idx):
    ori_img, bbox, label, fname = self.db.__getitem__(idx)

    img, bbox, label, scale = self.tsf((ori_img, bbox, label))

    if self.transforms:

      if randint(0, 1000) % 2 == 1 or True :

        img = img.transpose((1,2,0))
        objtotensor = T.ToTensor()
        objgaus = T.GaussianBlur(kernel_size=(5, 9), sigma=(0.1, 5))
        objjit = T.ColorJitter(brightness=.5, hue=.3)
        objco = Cutout(1,8)
        img = objtotensor(img.copy())
        img = objgaus(img)
        img = objjit(img)

        img = objco(img)

        img = img.numpy()

    return img.copy(), bbox.copy(), label.copy(), scale, fname

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

class FSBDDataset2x:

  def __init__(self, split = 'beta'):
    self.db = TrBboxDataset( split = split)
    self.tsf = Transform(conf.min_size, conf.max_size)

  def __getitem__(self, idx):
    real_idx = idx//2
    ori_img, bbox, label, fname = self.db.__getitem__(real_idx)

    img, bbox, label, scale = self.tsf((ori_img, bbox, label))

    img = img.transpose((1,2,0))
    objtotensor = T.ToTensor()
    objgaus = T.GaussianBlur(kernel_size=(5, 9), sigma=(0.1, 5))
    objjit = T.ColorJitter(brightness=.5, hue=.3)
    objco = Cutout(1,8)
    img = objtotensor(img.copy())

    if idx % 2 != 0 :
      img = objgaus(img)
      img = objjit(img)
      img = objco(img)

    img = img.numpy()

    return img.copy(), bbox.copy(), label.copy(), scale, fname

  def __len__(self):
    return (len(self.db)*2)

class FSBDDataset4x:

  def __init__(self, split = 'beta'):
    self.db = TrBboxDataset( split = split)
    self.tsf = Transform(conf.min_size, conf.max_size)

  def __getitem__(self, idx):
    real_idx = idx//4
    ori_img, bbox, label, fname = self.db.__getitem__(real_idx)

    img, bbox, label, scale = self.tsf((ori_img, bbox, label))

    img = img.transpose((1,2,0))
    objtotensor = T.ToTensor()
    objgaus = T.GaussianBlur(kernel_size=(5, 9), sigma=(0.1, 5))
    objjit = T.ColorJitter(brightness=.5, hue=.3)
    objco = Cutout(1,8)
    img = objtotensor(img.copy())

    if idx % 4 != 0 :
      img = objgaus(img)
      img = objjit(img)
      img = objco(img)

    img = img.numpy()

    return img.copy(), bbox.copy(), label.copy(), scale, fname

  def __len__(self):
    return (len(self.db)*4)

if __name__ == '__main__':

  print("Inside dataset")

  conf.data_dir = conf.ten_shot_tn
  conf.image_path = conf.ten_shot_image
  conf.annot_path = conf.ten_shot_annot

  db = FSBDDataset2x(split='beta')
  print(db.__len__())

  for ii in range(db.__len__()):
    db.__getitem__(ii)

  img,bbox,label,scale,fname = db.__getitem__(0)
  print("Displaying img shape : ",img.shape)

  img = img.transpose((1,2,0))

  fig, ax = plt.subplots()

  for bb in bbox:
    h = bb[2]-bb[0]
    w = bb[3]-bb[1]
    rect = patches.Rectangle((bb[1], bb[0]), w, h, linewidth=1, edgecolor='r', facecolor='none')
    ax.add_patch(rect)

  plt.imshow(img)
  plt.savefig('abc.png')