import argparse
import cv2
import json
import math
import numpy as np
# from pycocotools.coco import COCO
# from pycocotools.cocoeval import COCOeval

import torch

# from datasets.coco import CocoValDataset
from pose2d_models.with_mobilenet import PoseEstimationWithMobileNet
from modules.keypoints import extract_keypoints, group_keypoints
from modules.load_state import load_state

def run_coco_eval(gt_file_path, dt_file_path):
    annotation_type = 'keypoints'
    print('Running test for {} results.'.format(annotation_type))

    coco_gt = COCO(gt_file_path)
    coco_dt = coco_gt.loadRes(dt_file_path)

    result = COCOeval(coco_gt, coco_dt, annotation_type)
    result.evaluate()
    result.accumulate()
    result.summarize()

def normalize(img, img_mean, img_scale):
    img = np.array(img, dtype=np.float32)
    img = (img - img_mean) * img_scale
    return img

def pad_width(img, stride, pad_value, min_dims):
    h, w, _ = img.shape
    h = min(min_dims[0], h)
    min_dims[0] = math.ceil(min_dims[0] / float(stride)) * stride
    min_dims[1] = max(min_dims[1], w)
    min_dims[1] = math.ceil(min_dims[1] / float(stride)) * stride
    pad = []
    pad.append(int(math.floor((min_dims[0] - h) / 2.0)))
    pad.append(int(math.floor((min_dims[1] - w) / 2.0)))
    pad.append(int(min_dims[0] - h - pad[0]))
    pad.append(int(min_dims[1] - w - pad[1]))
    padded_img = cv2.copyMakeBorder(img, pad[0], pad[2], pad[1], pad[3],
                                    cv2.BORDER_CONSTANT, value=pad_value)
    return padded_img, pad

def convert_to_coco_format(pose_entries, all_keypoints):
    coco_keypoints = []
    scores = []
    for n in range(len(pose_entries)):
        if len(pose_entries[n]) == 0:
            continue
        keypoints = [0] * 17 * 3
        to_coco_map = [0, -1, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3]
        person_score = pose_entries[n][-2]
        position_id = -1
        for keypoint_id in pose_entries[n][:-2]:
            position_id += 1
            if position_id == 1:  # no 'neck' in COCO
                continue

            cx, cy, score, visibility = 0, 0, 0, 0  # keypoint not found
            if keypoint_id != -1:
                cx, cy, score = all_keypoints[int(keypoint_id), 0:3]
                cx = cx + 0.5
                cy = cy + 0.5
                visibility = 1
            keypoints[to_coco_map[position_id] * 3 + 0] = cx
            keypoints[to_coco_map[position_id] * 3 + 1] = cy
            keypoints[to_coco_map[position_id] * 3 + 2] = visibility
        coco_keypoints.append(keypoints)
        scores.append(person_score * max(0, (pose_entries[n][-1] - 1)))  # -1 for 'neck'
    return coco_keypoints, scores

def infer(net, img, scales, base_height, stride, pad_value=(0, 0, 0), img_mean=(128, 128, 128), img_scale=1/256):
    normed_img = normalize(img, img_mean, img_scale)
    height, width, _ = normed_img.shape
    scales_ratios = [scale * base_height / float(height) for scale in scales]
    avg_heatmaps = np.zeros((height, width, 19), dtype=np.float32)
    avg_pafs = np.zeros((height, width, 38), dtype=np.float32)

    for ratio in scales_ratios:
        scaled_img = cv2.resize(normed_img, (0, 0), fx=ratio, fy=ratio, interpolation=cv2.INTER_CUBIC)
        min_dims = [base_height, max(scaled_img.shape[1], base_height)]
        padded_img, pad = pad_width(scaled_img, stride, pad_value, min_dims)

        tensor_img = torch.from_numpy(padded_img).permute(2, 0, 1).unsqueeze(0).float().cuda()
        stages_output = net(tensor_img)

        stage2_heatmaps = stages_output[-2]
        heatmaps = np.transpose(stage2_heatmaps.squeeze().cpu().data.numpy(), (1, 2, 0))
        heatmaps = cv2.resize(heatmaps, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
        heatmaps = heatmaps[pad[0]:heatmaps.shape[0] - pad[2], pad[1]:heatmaps.shape[1] - pad[3]:, :]
        heatmaps = cv2.resize(heatmaps, (width, height), interpolation=cv2.INTER_CUBIC)
        avg_heatmaps = avg_heatmaps + heatmaps / len(scales_ratios)

        stage2_pafs = stages_output[-1]
        pafs = np.transpose(stage2_pafs.squeeze().cpu().data.numpy(), (1, 2, 0))
        pafs = cv2.resize(pafs, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
        pafs = pafs[pad[0]:pafs.shape[0] - pad[2], pad[1]:pafs.shape[1] - pad[3], :]
        pafs = cv2.resize(pafs, (width, height), interpolation=cv2.INTER_CUBIC)
        avg_pafs = avg_pafs + pafs / len(scales_ratios)

    return avg_heatmaps, avg_pafs

def evaluate(labels, output_name, images_folder, net, multiscale=False, visualize=False):
    net = net.cuda().eval()
    base_height = 368
    scales = [1]
    if multiscale:
        scales = [0.5, 1.0, 1.5, 2.0]
    stride = 8

    dataset = CocoValDataset(labels, images_folder)
    coco_result = []
    for sample in dataset:
        file_name = sample['file_name']
        img = sample['img']

        avg_heatmaps, avg_pafs = infer(net, img, scales, base_height, stride)

        total_keypoints_num = 0
        all_keypoints_by_type = []
        for kpt_idx in range(18):  # 19th for bg
            total_keypoints_num += extract_keypoints(avg_heatmaps[:, :, kpt_idx], all_keypoints_by_type, total_keypoints_num)

        pose_entries, all_keypoints = group_keypoints(all_keypoints_by_type, avg_pafs)

        coco_keypoints, scores = convert_to_coco_format(pose_entries, all_keypoints)

        image_id = int(file_name[0:file_name.rfind('.')])
        for idx in range(len(coco_keypoints)):
            coco_result.append({
                'image_id': image_id,
                'category_id': 1,  # person
                'keypoints': coco_keypoints[idx],
                'score': scores[idx]
            })

        if visualize:
            for keypoints in coco_keypoints:
                for idx in range(len(keypoints) // 3):
                    cv2.circle(img, (int(keypoints[idx * 3]), int(keypoints[idx * 3 + 1])),
                               3, (255, 0, 255), -1)
            cv2.imshow('keypoints', img)
            key = cv2.waitKey()
            if key == 27:  # esc
                return

    with open(output_name, 'w') as f:
        json.dump(coco_result, f, indent=4)

    run_coco_eval(labels, output_name)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--labels', type=str, required=True, help='path to json with keypoints val labels')
    parser.add_argument('--output-name', type=str, default='detections.json',
                        help='name of output json file with detected keypoints')
    parser.add_argument('--images-folder', type=str, required=True, help='path to COCO val images folder')
    parser.add_argument('--checkpoint-path', type=str, required=True, help='path to the checkpoint')
    parser.add_argument('--multiscale', action='store_true', help='average inference results over multiple scales')
    parser.add_argument('--visualize', action='store_true', help='show keypoints')
    args = parser.parse_args()

    net = PoseEstimationWithMobileNet()
    checkpoint = torch.load(args.checkpoint_path)
    load_state(net, checkpoint)

    evaluate(args.labels, args.output_name, args.images_folder, net, args.multiscale, args.visualize)


# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Model an instrument response for spectroscopic simulations.

An instrument model is usually initialized from a configuration used to create
a simulator and then accessible via its ``instrument`` attribute, for example:

    >>> import specsim.simulator
    >>> simulator = specsim.simulator.Simulator('test')  # doctest: +IGNORE_OUTPUT
    >>> print(np.round(simulator.instrument.fiber_diameter, 1))
    107.0 um

See :doc:`/api` for examples of changing model parameters defined in the
configuration. No attributes can be changed after a simulator has
been created.  File a github issue if you would like to change this.

An :class:`Instrument` includes one or more
:class:`Cameras <specsim.camera.Camera>`.
"""
from __future__ import print_function, division

import numpy as np
import os.path
import scipy.interpolate
import scipy.integrate

import astropy.constants
import astropy.units as u

import specsim.camera
import specsim.fastfiberacceptance
import specsim.config

class Instrument(object):
    """Model the instrument response of a fiber spectrograph.

    A spectrograph can have multiple :mod:`cameras <specsim.camera>` with
    different wavelength coverages. Objects representing each camera are
    contained in a list accessible from our ``cameras`` attribute, which will
    be in order of increasing effective wavelength.

    No instrument attributes can be changed after an instrument has been
    created. Create a github issue if you would like to change this.

    Parameters
    ----------
    name : str
        Descriptive name of this instrument.
    wavelength : astropy.units.Quantity
        Array of wavelength bin centers where the instrument response is
        calculated, with units.
    fiberloss_method : str
        Must be "table" or "galsim" or "fastsim".  Specifies how fiber acceptance fractions
        will be loaded or calculated.
    fiber_acceptance_dict : dict or None
        Dictionary of fiber acceptance fractions tabulated for different
        source models, with keys corresponding to source model names.
        Ignored when fiberloss_method is "galsim".
    fast_fiber_acceptance : specsim.fastfiberacceptance.FastFiberAcceptance or None
        Initialized instance to use when fiberloss_method is "fastsim".
        Ignored for other values of fiberloss_method.
    fiberloss_num_wlen : int
        Number of wavelengths where the fiberloss fraction should be tabulated
        for interpolation.  Ignored when fiberloss_method is not "galsim".
    fiberloss_num_pixels : int
        Number of pixels used to subdivide the fiber diameter for
        numerical convolution and integration calculations.
        Ignored when fiberloss_method  is not "galsim".
    blur_function : callable
        Function of field angle and wavelength that returns the corresponding
        RMS blur in length units (e.g., microns).
    offset_function : callable
        Function of focal-plane position (x,y) in angular units and wavelength
        that returns the corresponding radial centroid offset in length
        units (e.g., microns).
    cameras : list
        List of :class:`specsim.camera.Camera` instances representing the
        camera(s) of this instrument.
    primary_mirror_diameter : astropy.units.Quantity
        Diameter of the primary mirror, with units.
    obscuration_diameter : astropy.units.Quantity
        Diameter of a central obscuration of the primary mirror, with units.
    support_width : astropy.units.Quantity
        Width of the obscuring supports, with units.
    fiber_diameter : astropy.units.Quantity
        Physical diameter of the simulated fibers, with units of length.
        Converted to an on-sky diameter using the plate scale.
    field_radius : astropy.units.Quantity
        Maximum radius of the field of view in length units measured at
        the focal plane. Converted to an angular field of view using the
        plate scale.
    radial_scale : callable
        Callable function that returns the plate scale in the radial
        (meridional) direction (with appropriate units) as a function of
        focal-plane distance (with length units) from the boresight.
    azimuthal_scale : callable
        Callable function that returns the plate scale in the azimuthal
        (sagittal) direction (with appropriate units) as a function of
        focal-plane distance (with length units) from the boresight.
    """
    def __init__(self, name, wavelength, fiberloss_method,
                 fiber_acceptance_dict, fast_fiber_acceptance, fiberloss_num_wlen,
                 fiberloss_num_pixels, blur_function, offset_function, cameras,
                 primary_mirror_diameter, obscuration_diameter, support_width,
                 fiber_diameter, field_radius, radial_scale, azimuthal_scale):
        self.name = name
        self._wavelength = wavelength
        self.fiber_acceptance_dict = fiber_acceptance_dict
        self.fast_fiber_acceptance = fast_fiber_acceptance
        # Both fiber_acceptance_dict and fast_fiber_acceptance must be initialized
        # before assigning to fiberloss_method (since its setter checks their values).
        self.fiberloss_method = fiberloss_method
        self.fiberloss_num_wlen = fiberloss_num_wlen
        self.fiberloss_num_pixels = fiberloss_num_pixels
        self._blur_function = blur_function
        self._offset_function = offset_function
        self.cameras = cameras
        self.primary_mirror_diameter = primary_mirror_diameter
        self.obscuration_diameter = obscuration_diameter
        self.support_width = support_width
        self.fiber_diameter = fiber_diameter
        self.field_radius = field_radius
        self.radial_scale = radial_scale
        self.azimuthal_scale = azimuthal_scale

        # Calculate the effective area of the primary mirror.
        D = self.primary_mirror_diameter
        obs = self.obscuration_diameter
        support_area = 0.5*(D - obs) * self.support_width
        self.effective_area = (
            np.pi * ((0.5 * D) ** 2 - (0.5 * obs) ** 2) - 4 * support_area)

        # Tabulate the mapping between focal plane radius and boresight
        # opening angle by integrating the radial plate scale.
        # Use mm and radians as the canonical units.
        self._radius_unit, self._angle_unit = u.mm, u.rad
        radius = np.linspace(
            0., self.field_radius.to(self._radius_unit).value, 1000)
        dradius_dangle = self.radial_scale(radius * self._radius_unit).to(
            self._radius_unit / self._angle_unit).value
        angle = scipy.integrate.cumtrapz(
            1. / dradius_dangle, radius, initial=0.)

        # Record the maximum field angle corresponding to our field radius.
        self.field_angle = angle[-1] * self._angle_unit

        # Build dimensionless linear interpolating functions of the
        # radius <-> angle map using the canonical units.
        self._radius_to_angle = scipy.interpolate.interp1d(
            radius, angle, kind='linear', copy=True, bounds_error=True)
        self._angle_to_radius = scipy.interpolate.interp1d(
            angle, radius, kind='linear', copy=True, bounds_error=True)

        # Calculate the energy per photon at each wavelength.
        hc = astropy.constants.h * astropy.constants.c
        energy_per_photon = (hc / self._wavelength).to(u.erg)

        # Calculate the rate of photons incident on the focal plane per
        # wavelength bin per unit spectral flux density. The fiber acceptance
        # fraction is not included in this calculation.
        wavelength_bin_size = np.gradient(self._wavelength)
        self.photons_per_bin = (
            self.effective_area * wavelength_bin_size / energy_per_photon
            ).to((u.cm**2 * u.Angstrom) / u.erg)

        wave_mid = []
        for i, camera in enumerate(self.cameras):
            wave_min, wave_max = camera.wavelength_min, camera.wavelength_max
            wave_mid.append(0.5 * (wave_min + wave_max))
            if i == 0:
                self.wavelength_min = wave_min
                self.wavelength_max = wave_max
            else:
                self.wavelength_min = min(self.wavelength_min, wave_min)
                self.wavelength_max = max(self.wavelength_max, wave_max)

        # Sort cameras in order of increasing wavelength.
        self.cameras = [x for (y, x) in sorted(zip(wave_mid, self.cameras))]

    @property
    def fiberloss_method(self):
        """The current method used to calculate fiber acceptance fractions.
        """
        return self._fiberloss_method

    @fiberloss_method.setter
    def fiberloss_method(self, fiberloss_method):
        """Set the method used to calculate fiber acceptance fractions.

        Must be one of "table" or "galsim" or "fastsim".
        """
        if fiberloss_method not in ('table', 'galsim', 'fastsim' ):
            raise ValueError(
                'fiberloss_method must be "table" or "galsim" or "fastsim".')
        if fiberloss_method == 'table' and self.fiber_acceptance_dict is None:
            raise ValueError('Missing required instrument.fiberloss.table.')
        if fiberloss_method == 'fastsim' and self.fast_fiber_acceptance is None:
            raise ValueError(
                'Missing required instrument.fiberloss.fast_fiber_acceptance_path.')
        if fiberloss_method == 'galsim':
            try:
                import galsim
            except ImportError:
                raise ValueError('The galsim package is not installed.')
        self._fiberloss_method = fiberloss_method

    def field_radius_to_angle(self, radius):
        """Convert focal plane radius to an angle relative to the boresight.

        The mapping is derived from the radial (meridional) plate scale
        function :math:`dr/d\\theta(r)` via the integral:

        .. math::

            \\theta(r) = \\int_0^{r} \\frac{dr}{dr/d\\theta(r')}\\, dr'

        The input values must be within the field of view.
        Use :meth:`field_angle_to_radius` for the inverse transform.

        Parameters
        ----------
        radius : astropy.units.Quantity
            One or more radius values where the angle should be calculated.
            Values must be between 0 and ``field radius``.

        Returns
        -------
        astropy.units.Quantity
            Opening angle(s) relative to the boresight corresponding to
            the input radius value(s).

        Raises
        ------
        ValueError
            One or more input values are outside the allowed range.
        """
        return self._radius_to_angle(
            radius.to(self._radius_unit)) * self._angle_unit

    def field_angle_to_radius(self, angle):
        """Convert focal plane radius to an angle relative to the boresight.

        The mapping :math:`r(\\theta)` is calculated by numerically inverting
        the function :math:`\\theta(r)`.

        The input values must be within the field of view.
        Use :meth:`field_radius_to_angle` for the inverse transform.

        Parameters
        ----------
        angle : astropy.units.Quantity
            One or more angle values where the radius should be calculated.
            Values must be between 0 and ``field_angle``.

        Returns
        -------
        astropy.units.Quantity
            Radial coordinate(s) in the focal plane corresponding to the
            input angle value(s).

        Raises
        ------
        ValueError
            One or more input values are outside the allowed range.
        """
        return self._angle_to_radius(
            angle.to(self._angle_unit)) * self._radius_unit

    def get_blur_rms(self, wavelength, angle):
        """Get the instrument PSF blur at the specified field angle.

        Parameters
        ----------
        wavelength : astropy.units.Quantity
            Wavelength where the blur should be calculated.
        angle : astropy.units.Quantity
            Angular separation from the field center.

        Returns
        -------
        astropy.units.Quantity
            RMS blur of the instrument at this wavelength and field radius
            in length units.
        """
        return self._blur_function(angle, wavelength)

    def get_centroid_offset(self, angle_x, angle_y, wavelength):
        """Get the instrument centroid offset at the specified field angles.

        This method does not make any assumptions about how the x and y
        axes are defined, as long as (0, 0) is the field center.

        Note that the focal-plane position is input as angles relative to
        the field center, while the offsets are returned as lengths relative
        to the nominal fiber center.

        Parameters
        ----------
        angle_x : astropy.units.Quantity
            Angular separation from the field center along x.
        angle_y : astropy.units.Quantity
            Angular separation from the field center along y.
        wavelength : astropy.units.Quantity
            Wavelength where the blur should be calculated.

        Returns
        -------
        tuple
            Tuple (dx, dy) of astropy quantities giving the spot centroid
            offset components at this wavelength and position in the focal
            plane.  Offsets are given in length units, e.g., microns.
        """
        return self._offset_function(angle_x, angle_y, wavelength)

    def get_focal_plane_optics(self, focal_x, focal_y, wlen_grid):
        """Calculate the optical parameters at a set of focal-plane positions.

        Uses :meth:`get_centroid_offset`, :meth:`get_blur_rms`, and
        :meth:`field_radius_to_angle` to calculate the optics at each focal
        plane location.

        This method does not make any assumptions about how the x and y
        axes are defined, as long as (0, 0) is the field center.  However
        radial symmetry is broken by the (dx, dy) offsets calculated by
        :meth:`get_centroid_offset`.

        Note that units are required for the input arrays and included with
        the returned arrays.

        Parameters
        ----------
        focal_x : :class:`astropy.units.Quantity`
            1D array of X coordinates in the focal plane relative to the
            boresight, with length units.
        focal_y : :class:`astropy.units.Quantity`
            1D array of Y coordinates in the focal plane relative to the
            boresight, with length units.
        wlen_grid : :class:`astropy.units.Quantity`
            1D array of wavelengths where parameters should be tabulated,
            with length units.

        Returns
        -------
        tuple
            Tuple of arrays scale, blur, offset with shapes (N,2), (N,M) and
            (N,M,2) where N is the size of the 1D input (x,y) arrays, M is
            the size of the input wavelength grid, and axes of length 2
            correspond to radial and azimuthal axes (not the input x,y!).
            All output arrays have units.
        """
        # Check for valid units on the input arrays.
        try:
            focal_x_mm = focal_x.to(u.mm).value
            focal_y_mm = focal_y.to(u.mm).value
            wlen_grid_ang = wlen_grid.to(u.Angstrom).value
        except astropy.units.UnitConversionError:
            raise ValueError('Input arrays have invalid units.')
        except AttributeError:
            raise ValueError('Input arrays are missing required units.')

        # Check for expected input array shapes.
        if len(focal_x_mm.shape) != 1 or len(wlen_grid_ang.shape) != 1:
            raise ValueError('Input arrays must be 1D.')
        if focal_x_mm.shape != focal_y_mm.shape:
            raise ValueError('Input (x,y) arrays have different shapes.')

        # Allocate output arrays.
        n_xy = len(focal_x_mm)
        n_wlen = len(wlen_grid_ang)
        scale = np.empty((n_xy, 2))
        blur = np.empty((n_xy, n_wlen))
        offset = np.empty((n_xy, n_wlen, 2))

        # Convert x, y offsets in length units to field angles.
        focal_r = np.sqrt(focal_x**2+focal_y**2)
        angle_r = self.field_radius_to_angle(focal_r)
        angle_x = np.zeros(focal_x.shape) * angle_r.unit
        angle_y = np.zeros(focal_y.shape) * angle_r.unit
        positive_radius = focal_r>0
        angle_x[positive_radius] = (
            angle_r[positive_radius] / focal_r[positive_radius]
            ) * focal_x[positive_radius]
        angle_y[positive_radius] = (
            angle_r[positive_radius] / focal_r[positive_radius]
            ) * focal_y[positive_radius]

        # Calculate the radial and azimuthal plate scales at each location.
        scale[:, 0] = self.radial_scale(focal_r).to(u.um / u.arcsec).value
        scale[:, 1] = self.azimuthal_scale(focal_r).to(u.um / u.arcsec).value

        # Calculate the transformations between polar and Cartesian coordinates.
        phi = np.arctan2(focal_y_mm, focal_x_mm)
        cos_phi = np.cos(phi)
        sin_phi = np.sin(phi)

        # Lookup the instrument blur and centroid offset at each
        # wavelength for this focal-plane position.
        for i, wlen in enumerate(wlen_grid):
            # Lookup the RMS blurs in focal-plane microns.
            blur[:, i] = self.get_blur_rms(wlen, angle_r).to(u.um).value
            # Lookup the radial centroid offsets in focal-plane microns.
            dx, dy = self.get_centroid_offset(angle_x, angle_y, wlen)
            dx_um = dx.to(u.um).value
            dy_um = dy.to(u.um).value
            # Rotate to polar coordinates.
            offset[:, i, 0] =  cos_phi * dx_um + sin_phi * dy_um
            offset[:, i, 1] =  -sin_phi * dx_um + cos_phi * dy_um

        return scale * (u.um / u.arcsec), blur * u.um, offset * u.um

    def plot_field_distortion(self):
        """Plot focal plane distortions over the field of view.

        Requires that the matplotlib package is installed.
        """
        import matplotlib.pyplot as plt

        # Tabulate the field radius - angle mapping.
        radius = np.linspace(0., self.field_radius.to(u.mm).value, 500) * u.mm
        angle = self.field_radius_to_angle(radius).to(u.deg)

        # Calculate the r**2 weighted mean inverse radial scale by minimizing
        # angle - mean_inv_radial_scale * radius with respect to
        # mean_inv_radial_scale.
        mean_inv_radial_scale = (
            np.sum(radius ** 3 * angle) / np.sum(radius ** 4))
        mean_radial_scale = (1. / mean_inv_radial_scale).to(u.um / u.arcsec)

        # Calculate the angular distortion relative to the mean radial scale.
        distortion = (angle - radius * mean_inv_radial_scale).to(u.arcsec)

        # Eliminate round off error so that the zero distortion case is
        # correctly recovered.
        distortion = np.round(distortion, decimals=5)

        # Calculate the fiber area as a function of radius.
        radial_size = (
            0.5 * self.fiber_diameter / self.radial_scale(radius))
        azimuthal_size = (
            0.5 * self.fiber_diameter / self.azimuthal_scale(radius))
        fiber_area = (np.pi * radial_size * azimuthal_size).to(u.arcsec ** 2)

        # Calculate the r**2 weighted mean fiber area.
        mean_fiber_area = np.sum(radius ** 2 * fiber_area) / np.sum(radius ** 2)

        # Calculate the dimensionless fiber area ratio.
        fiber_area_ratio = (fiber_area / mean_fiber_area).si.value

        # Calculate the dimensionless ratio of azimuthal / radial plate scales
        # which is the ratio of the on-sky radial / azimuthal extends.
        shape_ratio = (self.azimuthal_scale(radius) /
                       self.radial_scale(radius)).si.value

        # Make the plots.
        fig, (ax1, ax2) = plt.subplots(2, sharex=True, figsize=(8, 8))

        ax1.plot(angle, distortion, 'b-', lw=2)
        ax1.set_ylabel('Field angle distortion [arcsec]', fontsize='large')
        ax1.set_xlim(0., self.field_angle.to(u.deg).value)
        ax1.grid()

        ax1.axhline(0., color='r')
        xy = 0.5 * self.field_angle.to(u.deg).value, 0.
        label = '{0:.1f}'.format(mean_radial_scale)
        ax1.annotate(label, xy, xy, color='r', horizontalalignment='center',
                     verticalalignment='bottom', fontsize='large')

        ax2.plot(angle, fiber_area_ratio, 'b', lw=2, label='Area ratio')
        ax2.plot(angle, shape_ratio, 'k', lw=2, ls='--',
                 label='Radial/azimuthal')
        ax2.set_ylabel('Fiber sky area and shape ratios', fontsize='large')
        ax2.grid()
        ax2.legend(loc='upper right')

        ax2.axhline(1., color='r')
        xy = 0.5 * self.field_angle.to(u.deg).value, 1.
        label = '{0:.3f}'.format(mean_fiber_area)
        ax2.annotate(label, xy, xy, color='r', horizontalalignment='center',
                     verticalalignment='bottom', fontsize='large')

        ax2.set_xlabel('Field angle [deg]', fontsize='large')
        plt.subplots_adjust(
            left=0.10, right=0.98, bottom=0.07, top=0.97, hspace=0.05)

    def plot(self, flux=1e-17 * u.erg / (u.cm**2 * u.s * u.Angstrom),
             exposure_time=1000 * u.s, cmap='nipy_spectral'):
        """Plot a summary of this instrument's model.

        Requires that the matplotlib package is installed.

        Parameters
        ----------
        flux : astropy.units.Quantity
            Constant source flux to use for displaying the instrument response.
        exposure_time : astropy.units.Quantity
            Exposure time to use for displaying the instrument response.
        cmap : str or matplotlib.colors.Colormap
            Matplotlib colormap name or instance to use for displaying the
            instrument response.  Colors are selected for each camera
            according to its central wavelength, so a spectral color map
            will give reasonably intuitive results.
        """
        import matplotlib.pyplot as plt
        import matplotlib.cm as cm

        fig, (ax1, ax2) = plt.subplots(2, sharex=True, figsize=(8, 8))
        ax1_rhs = ax1.twinx()
        ax2_rhs = ax2.twinx()
        cmap = cm.get_cmap(cmap)

        wave = self._wavelength.value
        wave_unit = self._wavelength.unit
        dwave = np.gradient(wave)

        if self.fiber_acceptance_dict:
            for source_type in self.fiber_acceptance_dict:
                # Plot fiber acceptance fractions without labels.
                ax1.plot(wave, self.fiber_acceptance_dict[source_type], 'k--')
        for camera in self.cameras:
            cwave = camera._wavelength

            # Use an approximate spectral color for each band.
            mid_wave = 0.5 * (camera.wavelength_min + camera.wavelength_max)
            color = cmap(
                (mid_wave - self.wavelength_min) /
                (self.wavelength_max - self.wavelength_min))

            # Calculate number of photons with perfect fiber acceptance.
            nphot = (flux * self.photons_per_bin * exposure_time *
                     camera.throughput / dwave)
            dark_noise = np.sqrt(
                (camera.dark_current_per_bin * exposure_time).value)
            total_noise = np.sqrt(
                dark_noise ** 2 + camera.read_noise_per_bin.value ** 2)

            ax1.plot(cwave, camera.throughput, ls='-', color=color)

            ax1_rhs.plot(cwave, nphot.value, ls=':', color=color)
            ax1_rhs.fill_between(
                cwave, total_noise / dwave, lw=0, color=color, alpha=0.2)
            ax1_rhs.fill_between(
                cwave, dark_noise / dwave, lw=0, color=color, alpha=0.2)
            ax1_rhs.plot(cwave, total_noise / dwave, ls='-.', color=color)

            ax2.plot(
                cwave, camera.rms_resolution.to(wave_unit).value,
                ls='-', color=color)
            ax2.plot(
                cwave, camera.row_size.to(wave_unit / u.pixel).value,
                ls='--', color=color)

            ax2_rhs.plot(
                cwave, camera.neff_spatial.to(u.pixel), ls=':', color=color)

        ax1.plot([], [], 'k--', label='Fiber Acceptance')
        ax1.plot([], [], 'k-', label='Camera Throughput')
        ax1.plot([], [], 'k:', label='{0}'.format(flux))
        ax1.plot([], [], 'k-.', label='Dark + Readout Noise')
        ax1.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
                   ncol=2, mode="expand", borderaxespad=0.)

        ax2.plot([], [], 'k-', label='RMS Resolution')
        ax2.plot([], [], 'k--', label='Row Size')
        ax2.plot([], [], 'k:', label='Column Size')
        ax2.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
                   ncol=3, mode="expand", borderaxespad=0.)

        ax1.set_ylim(0., None)
        ax1.set_ylabel('Fiber, Camera Throughput')
        ax1_rhs.set_ylim(0., None)
        ax1_rhs.set_ylabel(
            'Photons, Electrons / Exposure / {0}'.format(wave_unit))
        ax2.set_ylim(0., None)
        ax2.set_ylabel('RMS Resolution, Row Size [{0}]'.format(wave_unit))
        ax2_rhs.set_ylim(0., None)
        ax2_rhs.set_ylabel('Effective Column Size [pixels]')
        ax2.set_xlabel('Wavelength [{0}]'.format(wave_unit))
        ax2.set_xlim(wave[0], wave[-1])

def initialize(config, camera_output=True):
    """Initialize the instrument model from configuration parameters.

    This method is responsible for creating a new :class:`Instrument` as
    well as the :class:`Cameras <specsim.camera.Camera>` it includes.

    Parameters
    ----------
    config : :class:`specsim.config.Configuration`
        The configuration parameters to use.
    camera_output : bool
        Initialize support for resolution convolution and downsampling for
        each camera when True.

    Returns
    -------
    Instrument
        An initialized instrument model including one or more
        :class:`cameras <specsim.camera.Camera>`.
    """
    name = config.instrument.name
    cameras = config.instrument.cameras
    camera_names = cameras.keys()
    initialized_cameras = []
    for camera_name in camera_names:
        camera = getattr(cameras, camera_name)
        ccd = config.load_table(
            camera.ccd, ['row_size', 'fwhm_resolution', 'neff_spatial'])
        throughput = config.load_table(camera.throughput, 'throughput')
        constants = config.get_constants(camera,
            ['read_noise', 'dark_current', 'gain', 'num_sigmas_clip',
             'output_pixel_size'])
        initialized_cameras.append(specsim.camera.Camera(
            camera_name, config.wavelength, throughput,
            ccd['row_size'], ccd['fwhm_resolution'],
            ccd['neff_spatial'], constants['read_noise'],
            constants['dark_current'], constants['gain'],
            constants['num_sigmas_clip'], constants['output_pixel_size'],
            allow_convolution=camera_output))

    constants = config.get_constants(
        config.instrument,
        ['primary_mirror_diameter', 'obscuration_diameter',
         'support_width', 'fiber_diameter', 'field_radius'])

    try:
        # Try to read a tabulated plate scale first.
        plate_scale = config.load_table(
            config.instrument.plate_scale,
            ['radius', 'radial_scale', 'azimuthal_scale'], interpolate=False)
        r_vec = plate_scale['radius']
        sr_vec = plate_scale['radial_scale']
        sa_vec = plate_scale['azimuthal_scale']
        # Build dimensionless linear interpolators for the radial and azimuthal
        # scales using the native units from the tabulated data.
        sr_interpolate = scipy.interpolate.interp1d(
            r_vec.value, sr_vec.value, kind='linear', copy=True)
        sa_interpolate = scipy.interpolate.interp1d(
            r_vec.value, sa_vec.value, kind='linear', copy=True)
        # Wrap interpolators in lambdas that take care of units.
        radial_scale = lambda r: (
            sr_interpolate(r.to(r_vec.unit).value) * sr_vec.unit)
        azimuthal_scale = lambda r: (
            sa_interpolate(r.to(r_vec.unit).value) * sa_vec.unit)
    except AttributeError:
        # Fall back to a constant value.
        plate_scale_constant = config.get_constants(
            config.instrument.plate_scale, ['value'])
        value = plate_scale_constant['value']
        # Create lambdas that return the constant plate scale with units.
        # Use np.ones_like to ensure correct broadcasting.
        radial_scale = lambda r: value * np.ones_like(r.value)
        azimuthal_scale = lambda r: value * np.ones_like(r.value)

    # Initialize for both fiberloss methods so that method can be changed
    # at run time.
    fiberloss_method = config.instrument.fiberloss.method
    fiberloss_num_wlen = config.instrument.fiberloss.num_wlen
    fiberloss_num_pixels = config.instrument.fiberloss.num_pixels
    if hasattr(config.instrument.fiberloss, 'table'):
        fiber_acceptance_dict = config.load_table(
            config.instrument.fiberloss, 'fiber_acceptance', as_dict=True)
    else:
        fiber_acceptance_dict = None
    if hasattr(config.instrument.fiberloss, 'fast_fiber_acceptance_path'):
        filename = os.path.join(
            config.abs_base_path,
            config.instrument.fiberloss.fast_fiber_acceptance_path)
        if not os.path.isfile(filename) :
            raise RuntimeError(
            'Cannot find file {}. May need to update desimodel svn ?'
            .format(filename))
        fast_fiber_acceptance = specsim.fastfiberacceptance.FastFiberAcceptance(
            filename)
    else:
        fast_fiber_acceptance = None

    blur_value = getattr(config.instrument.blur, 'value', None)
    if blur_value:
        blur_value = specsim.config.parse_quantity(blur_value, u.micron)
        blur_function = lambda angle, wlen: blur_value
    else:
        blur_function = config.load_table2d(
            config.instrument.blur, 'wavelength', 'r=')

    offset_value = getattr(config.instrument.offset, 'value', None)
    if offset_value:
        offset_value = specsim.config.parse_quantity(offset_value, u.micron)
        offset_function = (
            lambda angle_x, angle_y, wlen: (offset_value, 0 * u.um))
    else:
        # Build an interpolator in (r, wlen) of radial chromatic offsets.
        radial_offset_function = config.load_table2d(
            config.instrument.offset, 'wavelength', 'r=')
        # Look for an optional file of achromatic offsets.
        # Static achromatic term, correlated in focal plane
        if hasattr(config.instrument.offset, 'static'):
            # Build an interpolator in (x, y).
            interpolators = config.load_fits2d(
                config.instrument.offset.static, xy_unit=u.deg,
                dx='XOFFSET', dy='YOFFSET')
            static_fx = interpolators['dx']
            static_fy = interpolators['dy']
        else:
            static_fx  = None
            static_fy  = None
        # Random uncorrelated achromatic term
        if hasattr(config.instrument.offset, 'sigma1d'):
            if hasattr(config.instrument.offset, 'seed'):
                offset_gen = np.random.RandomState(config.instrument.offset.seed)
            else:
                offset_gen = np.random.RandomState()
            sigma1d=specsim.config.parse_quantity(config.instrument.offset.sigma1d)
            random_fx = lambda angle_x, angle_y: offset_gen.normal(size=angle_x.shape) * sigma1d
            random_fy = lambda angle_x, angle_y: offset_gen.normal(size=angle_x.shape) * sigma1d
        else :
            random_fx = None
            random_fy = None
        # Combine the interpolators into a function of (x, y, wlen) that
        # returns (dx, dy).  Use default parameter values to capture the
        # necessary state in the inner function's closure.
        def offset_function(angle_x, angle_y, wlen,
                            fr=radial_offset_function,
                            static_fx=static_fx,
                            static_fy=static_fy,
                            random_fx=random_fx,
                            random_fy=random_fy) :
            angle_r = np.sqrt(angle_x ** 2 + angle_y ** 2)
            dr = fr(angle_r, wlen)
            # Special handling of the origin.
            not_at_origin = (angle_r > 0.)
            ux = np.ones(shape=dr.shape, dtype=float)
            uy = np.ones(shape=dr.shape, dtype=float)
            ux[not_at_origin] = angle_x / angle_r
            uy[not_at_origin] = angle_y / angle_r
            dx = dr * ux
            dy = dr * uy
            # Add interpolated offsets if any.
            if static_fx is not None : dx += static_fx(angle_x, angle_y)
            if static_fy is not None : dy += static_fy(angle_x, angle_y)
            if random_fx is not None : dx += random_fx(angle_x, angle_y)
            if random_fy is not None : dy += random_fy(angle_x, angle_y)
            return dx , dy

    instrument = Instrument(
        name, config.wavelength, fiberloss_method, fiber_acceptance_dict,
        fast_fiber_acceptance, fiberloss_num_wlen, fiberloss_num_pixels,
        blur_function, offset_function, initialized_cameras,
        constants['primary_mirror_diameter'], constants['obscuration_diameter'],
        constants['support_width'], constants['fiber_diameter'],
        constants['field_radius'], radial_scale, azimuthal_scale)

    if config.verbose:
        # Print some derived quantities.
        print('Telescope effective area: {0:.3f}'
              .format(instrument.effective_area))
        print('Field of view diameter: {0:.1f} = {1:.2f}.'
              .format(2 * instrument.field_radius.to(u.mm),
                      2 * instrument.field_angle.to(u.deg)))
        if fiber_acceptance_dict is not None:
            print('Fiberloss source types: {0}.'
                  .format(instrument.fiber_acceptance_dict.keys()))

    return instrument

import numpy,import scipy,import astropy


import numpy as np
import dlib
import pickle

predictor_path = "models/shape_predictor_5_face_landmarks.dat"
face_rec_model_path = "models/dlib_face_recognition_resnet_model_v1.dat"
faces_path = "models/faces.p"

labels_mapping = {
    1: "Tal",
    2: "Dror"
}

TRY_IDENTIFY = False

class Instance:
    _instances_count = 0
    face_detector = dlib.get_frontal_face_detector()
    sp = dlib.shape_predictor(predictor_path)
    facerec = dlib.face_recognition_model_v1(face_rec_model_path)
    faces = pickle.load(open(faces_path, "rb"))
    faces_dist_thresh = 0.55

    def __init__(self, event, frame=None):
        self._features = event._features
        self._events = [event]
        self._stability = 1
        self._bbox = event._bbox
        self._label = event._label
        self._identified = False
        self._score = event._score
        self._com = event._com

        self._estimate_object_distance()

        self._name = "%s_%d" % (self._label, Instance._instances_count)
        Instance._instances_count += 1

        if self._label == 'person' and self._identified is False:
            self._identify(frame)

    @property
    def stability(self):
        return self._stability

    @property
    def features(self):
        return self._features

    def _identify(self, frame):
        if TRY_IDENTIFY is False:
            return

        if self.stability % 10 > 0:
            return

        dets = self.face_detector(frame, 1)

        if len(dets) == 0:
            return

        d = dets[0] # assume only one face if any
        shape = self.sp(frame, d)

        face_descriptor = self.facerec.compute_face_descriptor(frame, shape)
        face_descriptor_str = str(face_descriptor).replace('\n', ',')
        face_descriptor_arr = face_descriptor_str.split(',')
        pred = self.faces.predict([face_descriptor_arr])
        dist, _ = self.faces.kneighbors([face_descriptor_arr])
        mean_dist = np.mean(dist)

        if mean_dist < self.faces_dist_thresh:
            rec = pred[0]

            self._identified = True
            self._name = labels_mapping[rec]
        else:
            print("detect face but unknown")

    def update(self, ev, partial=False, frame=None):
        self._events.append(ev)
        self._features = ev._features
        self._label = ev._label
        self._score = ev._score
        self._com = ev._com
        self._bbox = ev._bbox
        self._estimate_object_distance()
        self._stability = min(self._stability + 1, 10)

        if self._label == 'person' and self._identified == False:
            self._identify(frame)

    def _estimate_object_distance(self, far_threshold=3.0, near_threshold=1.5):

        _person_width_in_meters = 0.6
        _focal_length = 530.0

        width_in_pix = float(self._bbox[3]) - float(self._bbox[1])

        dist = (_person_width_in_meters * _focal_length) / width_in_pix

        self._distance = dist

        if dist < near_threshold:
            self._distance_level = 'Near'
        elif dist > far_threshold:
            self._distance_level = 'Far'
        else:
            self._distance_level = 'Mid'

class ObjectsTracker:
    def __init__(self):
        self._instances = []
        self._appearance_factor = 0.5
        self._euclidean_factor = 0.5

        self._threshold = 30.0

    @property
    def instances(self):
        return self._instances

    @staticmethod
    def L2(a, b):
        return np.linalg.norm(a-b)

    def match_events_to_instances(self, events, frame=None):
        if len(self._instances) == 0:
            return events

        total_dist = np.zeros((len(events), len(self._instances)))

        for i, inst in enumerate(self._instances):
            for j, ev in enumerate(events):
                total_dist[j, i] = self._appearance_factor * self.L2(ev._features, inst._features) + \
                                   self._euclidean_factor * self.L2(ev._com, inst._com) \
                                   + 1000 * float(ev._label != inst._label)

        unused_events = []
        matched_insts = []
        for j, ev in enumerate(events):
            i = np.argmin(total_dist[j, :])
            val = total_dist[j, i]
            if val < self._threshold:
                inst = self._instances[i]
                inst.update(ev, frame=frame)
                total_dist[:, i] = 9999999.0
                matched_insts.append(inst)
            else:
                unused_events.append(ev)

        for inst in self._instances:
            if inst in matched_insts:
                continue
            inst._stability = 0
            # inst._stability -= 1

        return unused_events

    def initiate_new_instances(self, events, frame=None):
        for event in events:
            self._instances.append(Instance(event, frame))

    def filter_events(self, events):
        return events

    def track(self, events, frame=None):
        events = self.filter_events(events)
        unmatched_events = self.match_events_to_instances(events, frame)
        self.initiate_new_instances(unmatched_events, frame)

        self.remove_old_instances()

        return self._instances

    def remove_old_instances(self):
        keep_instances = []
        for inst in self._instances:
            if inst.stability > 0:
                keep_instances.append(inst)

        self._instances = keep_instances

    def current_proposals(self):
        return []


import os
import time
import numpy as np
import pandas as pd
from tqdm import tqdm
import matplotlib.pyplot as plt
from matplotlib import patches
from utils_tf_record.read_dataset_utils import parse_camera_tfrecord_example, write_serialized_string
from waymo_open_dataset.protos import metrics_pb2
from waymo_open_dataset import label_pb2
from absl import app
from absl import flags
import tensorflow as tf

os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
tf.enable_eager_execution()

FLAGS = flags.FLAGS
flags.DEFINE_list('detections_files',
                    ["predictions_final_sample/aofinal_frcnn_500_256_extrafeat_red_softmax_bikes_higherlr/validation_detections.tfrecord",
                    "predictions_final_sample/aofinal_frcnn_500_256_extrafeat/validation_detections.tfrecord",
                     "predictions_final_sample/aofinal_frcnn_500_256_extrafeat_redsoftmax/validation_detections.tfrecord",
                     "predictions_final_sample/aofinal_frcnn_500_256_extrafeat_weights/validation_detections.tfrecord"
                     ],
                    'TFRecord detections files separated by comma')
flags.DEFINE_list('ignore_vehicles', [], 'detection_files of models which should ignore vehicle predictions.')
flags.DEFINE_list('ignore_pedestrians', [], 'detection_files of models which should ignore pedestrian predictions.')
flags.DEFINE_list('ignore_cyclists', [], 'detection_files of models which should ignore cyclist predictions.')
flags.DEFINE_string('output_path', "predictions_final_sample/ensemble4",
                    'Path to output serialized .bin files with metrics.Objects proto format')
flags.DEFINE_float("iou_threshold", 0.7, "IOU Threshold used to consider two predictions the same predictions")
flags.DEFINE_string("ensemble_type", 'NMS', 'Ensemble method to use: NMS (non-maximum supression) or AVG (Average the position between same predictions)')

def create_object(frame_features, x_max, x_min, y_max, y_min, label, difficulty=None, score=None):
    obj = metrics_pb2.Object()

    def create_label(img_width, img_height, x_max, x_min, y_max, y_min, label_type, difficulty=None):
        lab = label_pb2.Label()
        lab.box.center_x = (x_max + x_min) / 2 * img_width
        lab.box.center_y = (y_max + y_min) / 2 * img_height
        lab.box.length = (x_max - x_min) * img_width
        lab.box.width = (y_max - y_min) * img_height

        lab.type = 4 if label_type == 3 else label_type  # Revert cyclist label to 4
        if difficulty is not None:
            lab.detection_difficulty_level = 1 if difficulty == 0 else difficulty
        return lab

    label = create_label(frame_features['image/width'], frame_features['image/height'],
                         x_max, x_min, y_max, y_min, label, difficulty)
    obj.object.MergeFrom(label)
    if score:
        obj.score = score
    obj.context_name = frame_features["image/context_name"]
    obj.frame_timestamp_micros = frame_features["image/frame_timestamp_micros"]
    obj.camera_name = frame_features["image/camera_name"]

    return obj

def to_numpy(frame, features,  predictions, model=None, ignore_vehicles=False, ignore_pedestrians=False, ignore_cyclists=False):
    image_id = frame['image/context_name'].numpy().decode() + '_' + str(int(frame['image/frame_timestamp_micros'])) + '_' + str(int(frame['image/camera_name']))
    frame_features = {
        "image/context_name": frame['image/context_name'].numpy().decode(),
        "image/frame_timestamp_micros": int(frame['image/frame_timestamp_micros']),
        "image/camera_name": int(frame['image/camera_name']),
        "image/width": int(frame['image/width']),
        "image/height": int(frame['image/height']),
        'predictions': [],
        'predictions_1': [],
        'predictions_2': [],
        'predictions_3': [],
    }

    if image_id not in features:
        features[image_id] = frame_features

    # Predictions
    zip_pd_boxes = zip(frame['image/detection/bbox/xmax'], frame['image/detection/bbox/xmin'],
                       frame['image/detection/bbox/ymax'], frame['image/detection/bbox/ymin'],
                       frame['image/detection/label'], frame["image/detection/score"])

    npreds = len(predictions)

    i = npreds
    for (x_max, x_min, y_max, y_min, label, score) in zip_pd_boxes:
        if int(label.numpy()) == 1 and ignore_vehicles:
            continue
        if int(label.numpy()) == 2 and ignore_pedestrians:
            continue
        if int(label.numpy()) == 3 and ignore_cyclists:
            continue
        row = [
            y_min.numpy(),
            x_min.numpy(),
            y_max.numpy(),
            x_max.numpy(),
            label.numpy(),
            score.numpy()
        ]
        predictions.append(row)
        features[image_id]['predictions'] += [i]
        features[image_id]['predictions_{}'.format(int(label.numpy()))] += [i]
        i += 1

    return features, predictions

def AvgEnsemble(dets, iou_threshold=0.5):
    """
    Args:
        dets: list of [ymin, xmin, ymax, xmax, class, score]
        iou_threshold: float

    Returns: new detections like a list of [ymin, xmin, ymax, xmax, class, score]
    """

    def computeIOU(box1, box2):
        """
        Args:
            box1:  [ymin, xmin, ymax, xmax]
            box2:  [ymin, xmin, ymax, xmax]
        Returns: iou between box1 and box2
        """
        y11, x11, y12, x12 = box1
        y21, x21, y22, x22 = box2

        x_left = max(x11, x21)
        y_top = max(y11, y21)
        x_right = min(x12, x22)
        y_bottom = min(y12, y22)

        if x_right < x_left or y_bottom < y_top:
            return 0.0

        intersect_area = (x_right - x_left) * (y_bottom - y_top)
        box1_area = (x12 - x11) * (y12 - y11)
        box2_area = (x22 - x21) * (y22 - y21)

        iou = intersect_area / (box1_area + box2_area - intersect_area)
        return iou

    keep = []
    used = []
    for box1 in dets:
        if box1 in used:
            continue
        used.append(box1)

        found = [box1]
        for box2 in dets:
            if box2[4] != box1[4] or box2 in used :
                # Skip if box already used or box of different class
                continue
            iou = computeIOU(box1[:4], box2[:4])
            if iou > iou_threshold:
                used.append(box2)
                found.append(box2)

        y1_sum = 0.
        x1_sum = 0.
        y2_sum = 0.
        x2_sum = 0.
        weight_sum = 0.
        score = 1.

        for b in found:
            y1_sum += b[0] * b[5]
            x1_sum += b[1] * b[5]
            y2_sum += b[2] * b[5]
            x2_sum += b[3] * b[5]
            weight_sum += b[5]
            score *= (1 - b[5])

        y1 = y1_sum/weight_sum
        x1 = x1_sum/weight_sum
        y2 = y2_sum/weight_sum
        x2 = x2_sum/weight_sum

        score = 1-score

        label = int(box1[4])

        new_box = [y1, x1, y2, x2, label, score]
        keep.append(new_box)
    return keep

def main(_):
    detections_files = FLAGS.detections_files
    IOU_THRESHOLD = FLAGS.iou_threshold
    ENSEMBLE_TYPE = FLAGS.ensemble_type
    IGNORE_VEHICLES = FLAGS.ignore_vehicles
    IGNORE_CYCLISTS = FLAGS.ignore_cyclists
    IGNORE_PEDESTRIANS = FLAGS.ignore_pedestrians
    assert all(ignore_file in detections_files for ignore_file in IGNORE_VEHICLES + IGNORE_PEDESTRIANS + IGNORE_CYCLISTS), 'error while defining ignore class predictions'
    print("Num models: {}\nNum models ignoring vehicles: {}\nNum models ignoring pedestrians: {}\nNum models ignoring cyclists: {}\nIOU Threshold: {}\nEnsemble type: {}\n".format(len(detections_files), len(IGNORE_VEHICLES), len(IGNORE_PEDESTRIANS), len(IGNORE_CYCLISTS), IOU_THRESHOLD, ENSEMBLE_TYPE))

    # Reading predictions from models
    t0 = time.time()
    detections_datasets = [(detections_file, tf.data.TFRecordDataset(detections_file.strip(), compression_type='') \
                               .map(lambda x: parse_camera_tfrecord_example(x, with_detections=True))) \
                           for detections_file in detections_files]

    detections = []
    features = {}

    for j, (detection_file, detections_dataset) in enumerate(detections_datasets):
        ignore_vehicles = detection_file in IGNORE_VEHICLES
        ignore_pedestrians = detection_file in IGNORE_PEDESTRIANS
        ignore_cyclists = detection_file in IGNORE_CYCLISTS
        for i, frame in tqdm(enumerate(detections_dataset), desc="Reading predictions from model #{}".format(j)):
            features, detections = to_numpy(frame, features, detections, model=str(j), ignore_vehicles=ignore_vehicles,
                                            ignore_pedestrians=ignore_pedestrians, ignore_cyclists=ignore_cyclists)

    detections = np.array(detections)

    num_original_detections = len(detections)

    image_ids = []
    # Merging predictions for each images
    if ENSEMBLE_TYPE == 'NMS':
        keep = []
        for image_id in tqdm(features, desc="Merging predictions for each images (NMS method)"):
            image_features = features[image_id]

            labels = [1, 2, 3]
            for label in labels:
                indices = image_features['predictions_{}'.format(label)]
                image_detection_label = detections[image_features['predictions_{}'.format(label)]]

                boxes = tf.convert_to_tensor(image_detection_label[:, :4], tf.float32)
                scores = tf.convert_to_tensor(image_detection_label[:, 5], tf.float32)
                max_output_size = tf.convert_to_tensor(len(image_detection_label), tf.int32)
                iou_threshold = IOU_THRESHOLD

                filtered_image_detections = tf.image.non_max_suppression(boxes, scores, max_output_size, iou_threshold)
                filtered_image_detections = [indices[x] for x in filtered_image_detections.numpy()]
                keep += filtered_image_detections
                image_ids += [image_id for _ in range(len(filtered_image_detections))]

        detections = detections[keep]

    elif ENSEMBLE_TYPE == 'AVG':
        new_detections = []
        for image_id in tqdm(features, desc="Merging predictions for each images (NMS method)"):
            image_features = features[image_id]
            image_detections = detections[image_features['predictions']]

            context_name, frame_timestamp_micros, camera_name, width, height = image_id[0], image_id[1], image_id[2], image_id[3], image_id[4]

            boxes = [np_box.tolist() for np_box in image_detections]
            new_dets = AvgEnsemble(boxes, IOU_THRESHOLD)

            new_detections += new_dets
            image_ids += [image_id for _ in range(len(new_dets))]

        detections = np.array(new_detections)

    num_detections = len(detections)

    # Creating prediction objects
    predictions = metrics_pb2.Objects()
    for index in tqdm(list(range(len(detections))), desc='Creating prediction objects'):
        image_features = features[image_ids[index]]
        y_min, x_min, y_max,  x_max, label, score = detections[index]
        frame_features = {
            "image/context_name": image_features["image/context_name"],
            "image/frame_timestamp_micros": image_features["image/frame_timestamp_micros"],
            "image/camera_name": image_features["image/camera_name"],
            "image/width": image_features["image/width"],
            "image/height": image_features["image/height"]
        }
        obj = create_object(frame_features, x_max, x_min, y_max, y_min, int(label), difficulty=None,
                            score=score)
        predictions.objects.extend([obj])

    # Write to file
    output_path = FLAGS.output_path
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    output_path = os.path.join(output_path, 'ensemble_{}_predictions.bin'.format(ENSEMBLE_TYPE.lower()))
    write_serialized_string(output_path, predictions)

    t1 = time.time()
    print("Detections remaining: {}/{} ({:.2f})".format(num_detections, num_original_detections, num_detections/num_original_detections))
    print("Processing time: {:.2f} s".format(t1-t0))

    return 0

if __name__ == '__main__':
    app.run(main)


from agent_dir.agent import Agent
import scipy
import numpy as np

import os
import keras
import tensorflow as tf
from keras.models import Sequential,load_model
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2D
from keras.optimizers import Adam, Adamax, RMSprop
from keras import backend as K

MAX_TIMESTEP = 1000
MAX_EP = 3000
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.1
set_session(tf.Session(config=config))

np.random.seed(1)
tf.set_random_seed(1)

# Reference: https://github.com/mkturkcan/Keras-Pong/blob/master/keras_pong.py
def discount_rewards(r):
    gamma=0.99
    discounted_r = np.zeros_like(r)
    running_add = 0
    for t in reversed(range(0, r.size)):
        running_add = running_add * gamma + r[t]
        discounted_r[t] = running_add
    return discounted_r

def categorical_crossentropy(target, output):
    _epsilon =  tf.convert_to_tensor(10e-8, dtype=output.dtype.base_dtype)
    output = tf.clip_by_value(output, _epsilon, 1. - _epsilon)
    return - tf.reduce_sum(target * tf.log(output),axis=len(output.get_shape()) - 1)

class Agent_PGC(Agent):
    def __init__(self, env, args):
        super(Agent_PGC,self).__init__(env)

        self.log_path = './pgc.log'
        self.model_path = './pong_model_cart_checkpoint.h5'
        if args.test_pg:
            print('loading trained model')
            self.model = load_model('./pong_model_cart_checkpoint.h5')
        else:
            self.env = env
            self.actions_avialbe = env.action_space.n
            self.feature_dim = env.observation_space.shape[0]
            self.learning_rate  = 1e-3

            model = Sequential()
            model.add(Dense(20,activation='relu',input_shape=(self.feature_dim,)))
            model.add(Dense(self.actions_avialbe,activation='softmax'))

            opt = Adam(lr=self.learning_rate)
            model.compile(loss=categorical_crossentropy, optimizer=opt)
            self.model = model

    def init_game_setting(self):
        self.prev_x = None

    def train(self):
        # Init
        log = open(self.log_path,'w')
        log.write('reward,avg_reward\n')
        batch_size = 1
        frames, prob_actions, dlogps, drs =[], [], [], []
        tr_x, tr_y = [],[]
        avg_reward = []
        reward_sum = 0
        t = 0
        ep_number = 0
        prev_x = None
        observation = self.env.reset()
        # Training progress
        while True:

            aprob = self.model.predict(observation.reshape((1,self.feature_dim)), batch_size=1).flatten()
            frames.append(observation)
            prob_actions.append(aprob)
            action = np.random.choice(self.actions_avialbe, 1, p=aprob.reshape((self.actions_avialbe)))[0]
            y = np.zeros([self.actions_avialbe])
            y[action] = 1
            observation, reward, done, info = self.env.step(action)

            reward_sum += reward
            drs.append(reward)
            dlogps.append(np.array(y).astype('float32') - aprob)

            t += 1

            if done or t> 1000:
                ep_number +=1
                ep_x = np.vstack(frames)
                ep_dlogp = np.vstack(dlogps)
                ep_reward = np.vstack(drs)
                # Discount and normalize rewards
                discounted_ep_reward = discount_rewards(ep_reward)
                discounted_ep_reward -= np.mean(discounted_ep_reward)
                discounted_ep_reward /= np.std(discounted_ep_reward)
                ep_dlogp *= discounted_ep_reward

                # Store current episode into training batch
                tr_x.append(ep_x)
                tr_y.append(ep_dlogp)
                frames, dlogps, drs =[], [], []
                if ep_number % batch_size == 0:
                    input_tr_y = prob_actions + self.learning_rate * np.squeeze(np.vstack(tr_y))
                    self.model.train_on_batch(np.vstack(tr_x).reshape(-1,self.feature_dim), input_tr_y)
                    tr_x,tr_y,prob_actions = [],[],[]
                    # Checkpoint
                    os.remove(self.model_path) if os.path.exists(self.model_path) else None
                    self.model.save(self.model_path)

                avg_reward.append(float(reward_sum))
                if len(avg_reward)>30: avg_reward.pop(0)
                print('Epsidoe {:} reward {:.2f}, Last 30ep Avg. rewards {:.2f}.'.format(ep_number,reward_sum,np.mean(avg_reward)))
                print('{:.4f},{:.4f}'.format(reward_sum,np.mean(avg_reward)),end='\n',file=log,flush=True)
                reward_sum = 0
                t = 0
                observation = self.env.reset()
            if ep_number>3000:
                break

    def make_action(self, observation, test=True):
        """

        Input:
            observation: np.array
                current RGB screen of game, shape: (210, 160, 3)

        Return:
            action: int
                the predicted action from trained model
        """
        cur_x = prepro(observation)
        x = cur_x - self.prev_x if self.prev_x is not None else np.zeros(cur_x.shape)
        self.prev_x = cur_x
        aprob = self.model.predict(x.reshape((1,80,80,1)), batch_size=1).flatten()

        return np.argmax(aprob)


import numpy as np
from falass import dataformat
import matplotlib.pyplot as plt
import MDAnalysis as mda

class Files:
    """File parsing.

    Definition of the files to be used in the analysis; such as the .pdb file, .lgt file (is one exists), .dat file
    (is one exists). Also defined at this time are the resolution (this is ignored if the .dat file has 4 columns)
    and the percentage error in the intensity (this is also ignored if the .dat file has >2 columns), and whether
    the simulation cell should be flipped in the xy-plane.

    Parameters
    ----------
    pdbfile: str
        Path and name of the .pdb file to be analysed.
    lgtfile: str, optional
        Path and name of the .lgt file (which contains the scattering lengths of each of the atom types in the
        pdbfile). If a lgtfile is not defined falass will help in the creation of one.
        Currently the .lgt file style that is supported is a 3 column space separated txt file where the columns are
        atom_type, real_scattering_length, and imaginary_scattering_length respectively.
    datfile: str, optional
        Path and name of the .dat file (from which the analysis q vectors are drawn and also for subsequent
        comparison between theory and experiment). If a datfile is not defined falass will allow the user to
        define a range of q vectors to calculate the reflectometry over.
        Currently the .dat file style that is supported is a 2, 3, and 4 column space separated txt files where the
        columns are q, i, di, and dq respectively.
    resolution: float, optional
        Percentage of the q vector for the width of the resolution Gaussian function to be used in the data
        smearing, if a 4 column datfile is given this is ignored.
    ierror: float, optional
        Percentage error of the intensity to be assumed, if a >2 column datfile is used this is ignored.
    flip: bool, optional
        False if the system should be read as is, true is the simulation cell should be rotated through the
        xy-plane -- note that falass treats the first side that the neutron or X-ray interacts with as that at z=0.
    xray: bool, optional
        True if the scattering length of the particles should be scaled by the classical radius of an electron.
    """
    def __init__(self, pdbfile=None, lgtfile=None, datfile=None, resolution=5., ierror=5., flip=False, xray=False):
        self.pdbfile = pdbfile
        self.cell = []
        self.atoms = []
        self.number_of_timesteps = 0
        self.times = []
        self.lgtfile = lgtfile
        self.scat_lens = []
        self.datfile = datfile
        self.expdata = []
        self.ierror = ierror
        self.resolution = resolution
        self.flip = flip
        self.xray = xray
        return

    def set_file(self, pdbfile=None, lgtfile=None, datfile=None):
        """Edits files.

        Let the subsequent definition, or redefinition of the pdbfile, lgtfile or datfile.

        Parameters
        ----------
        pdbfile: str
            Path and name of the .pdb file to be analysed.
        lgtfile: str, optional
            Path and name of the .lgt file (which contains the scattering lengths of each of the atom types in the
            pdbfile). If a lgtfile is not defined falass will help in the creation of one.
            Currently the .lgt file style that is supported is a 3 column space separated txt file where the columns are
            atom_type, real_scattering_length, and imaginary_scattering_length respectively.
        datfile: str, optional
            Path and name of the .dat file (from which the /home/arm61/progs/refnxanalysis q vectors are drawn and also for subsequent
            comparison between theory and experiment). If a datfile is not defined falass will allow the user to
            define a range of q vectors to calculate the reflectometry over.
            Currently the .dat file style that is supported is a 2, 3, and 4 column space separated txt files where the
            columns are q, i, di, and dq respectively.
        """
        if pdbfile:
            self.pdbfile = pdbfile
        if lgtfile:
            self.lgtfile = lgtfile
        if datfile:
            self.datfile = datfile
        return

    def read_pdb(self):
        """Parse .pdb.

        Reads the .pdb file into memory. Currently the atoms must have the title 'ATOM', the timestep time needs to
        be the last text in the 'TITLE' line, and the cell dimensions are taken from the 'CRYST1' line, and assumed to
        be orthorhomic. Non-orthorhomic cells are not necessarily supported.
        """
        print("Reading PDB file")
        self.u = u = mda.Universe(self.pdbfile)

        self.cell = []
        self.atoms = []

        for ts in u.trajectory:
            self.cell.append(u.dimensions[:3].copy())
            # grab z positions
            pos = u.atoms.positions[:, 2]
            # flip?
            if self.flip:
                pos = flip_zpos(u.dimensions[2], pos)
            self.atoms.append(
                [dataformat.AtomPositions(at.name, p) for at, p in zip(u.atoms, pos)]
            )

        lines = line_count(self.pdbfile)
        with open(self.pdbfile, 'r') as f:
            for i, line in enumerate(f):
                if "TITLE  " in line:
                    self.number_of_timesteps, new_time = iterate_time(self.number_of_timesteps, line)
                    self.times.append(new_time)

        return

    def read_lgt(self):
        """Parses .lgt.

        Parses the lgtfile. If no lgtfile is defined falass will help the user to build one by working through the
        atom types in the pdb file and requesting input of the real and imaginary scattering lengths. This will also
        occur if a atom type if found in the pdbfile but not in the given lgts file. falass will write the lgtfile
        to disk if atom types do not feature in the given lgtfile or one is written from scratch.
        """
        if self.lgtfile:
            lines = line_count(self.lgtfile)
            print("Reading LGT file")
            percentage = 0
            print_update(percentage)
            file = open(self.lgtfile, 'r')
            for i, line in enumerate(file):
                percentage_new = np.floor(i / lines * 100)
                percentage = check_update(percentage, percentage_new)
                line_list = line.split()
                duplicate = check_duplicates(self.scat_lens, line_list[0])
                if not duplicate:
                    i = 1
                    if self.xray:
                        i *= 2.817940
                    self.scat_lens.append(dataformat.ScatLens(line_list[0], float(line_list[1])*i, float(line_list[2])*i))
            print_update(100)
            file.close()
        else:
            raise ValueError("No lgtfile has been defined.")
        return

    def read_dat(self):
        """Parses .dat.

        Parses the .dat file, supporting 2, 3, and 4 column files consisting of q, i, di, and dq with comments in lines
        where the first character is a '#'. If there is no .dat file the get_qs() function should be used to generate
        q vectors to allow for the calculation of the reflectometry profile.
        """
        self.expdata=[]
        if self.datfile:
            lines = line_count(self.datfile)
            print("Reading DAT file")
            percentage = 0
            print_update(percentage)
            file = open(self.datfile, 'r')
            for i, line in enumerate(file):
                percentage_new = np.floor(i / lines * 100)
                percentage = check_update(percentage, percentage_new)
                if line[0] != '#':
                    line_list = line.split()
                    if len(line_list) == 2:
                        self.expdata.append(dataformat.QData(float(line_list[0]), float(line_list[1]),
                                                             float(line_list[1]) * (self.ierror / 100),
                                                             float(line_list[0]) * (self.resolution / 100)))
                    if len(line_list) == 3:
                        self.expdata.append(dataformat.QData(float(line_list[0]), float(line_list[1]),
                                                             float(line_list[2]),
                                                             float(line_list[0]) * (self.resolution / 100)))
                    if len(line_list) == 4:
                        self.expdata.append(dataformat.QData(float(line_list[0]), float(line_list[1]),
                                                             float(line_list[2]), float(line_list[3])))
            print_update(100)
            file.close()
        else:
            print('No DAT file has been given, therefore no comparison will be conducted, please use the get_qs '
                  'function. Alternatively the DAT file can be added using the setFile function.')
        return

    def get_qs(self, start=0.005, end=0.5, number=50):
        """Make custom q-vectors.

        If no datfile exists this function should be used to generate a linear-spaced range of q-vector for the
        calculation of the reflectometry over.

        Parameters
        ----------
        start: float, optional
            The first q-vector for which the reflectometry should be calculated.
        end: float, optional
            The last q-vector for which the reflectometry should be calculated.
        number: int, optional
            The number of q-vectors.
        """
        q_values = np.linspace(start, end, number)
        for i in range(0, len(q_values)):
            self.expdata.append(dataformat.QData(q_values[i], None, None, q_values[i] * (self.resolution / 100)))
        return

    def plot_dat(self, rq4=True):
        """Plot the experimental data file that has been read in.

        Parameters
        ----------
        rq4: bool, optional
            Should the plot be created with a y-axis of Rq^4
        """
        if self.datfile:
            x = []
            y = []
            dy = []
            plt.rc('text')
            plt.rc('font', family='serif')
            if rq4:
                for i in range(0, len(self.expdata)):
                    x.append(self.expdata[i].q)
                    y.append(np.log10(self.expdata[i].i * self.expdata[i].q ** 4))
                    dy.append((self.expdata[i].di * self.expdata[i].q ** 4) /
                              (self.expdata[i].i * np.log(10)))
                    plt.ylabel('log($Rq^4$) (\AA$^4$)')
            else:
                for i in range(0, len(self.expdata)):
                    x.append(self.expdata[i].q)
                    y.append(np.log10(self.expdata[i].i))
                    dy.append(self.expdata[i].di / (self.expdata[i].i * np.log(10)))
                    plt.ylabel('log($R$)')
            x = np.asarray(x)
            y = np.asarray(y)
            dy = np.asarray(dy)
            plt.errorbar(x, y, yerr=dy, marker='o', ls='')
            plt.xlabel('$q$ (\AA)')
            plt.show()
        else:
            raise ValueError("No data file is defined.")
        return plt

def check_duplicates(array, check):
    """Stops duplicate atom types.

    Checks if an atom type has already been added to an array.

    Parameters
    ----------
    array: array-type ScatLens
        The array to check.
    check: str
        The atom type to try and find.

    Returns
    -------
    bool
        True if the atom type is already present in the scatlen type array, false if not.
    """
    for i in range(0, len(array)):
        if array[i].atom == check:
            return True
    return False

def line_count(filename):
    """File length.

    Quickly counts the number of lines in a file

    Parameters
    ----------
    filename: str
        Name of the file that the number of lines is desired for.

    Returns
    -------
    int
        Number of lines in the file
    """
    i = 0
    with open(filename) as f:
        for i, l in enumerate(f):
            pass
    return i + 1

def flip_zpos(cell, zpos):
    """Flip the z-position.

    Flips the z-position through the xy-plane.

    Parameters
    ----------
    cell: float
        z-cell dimension.
    zpos: float
        z-position for an atom.

    Returns
    -------
    float
        z-position after the flipping
    """
    return np.sqrt(np.square(cell - zpos))

def iterate_time(number_of_timesteps, line):
    """Increases the number of timesteps

    Iterates the number of timesteps and finds the new timestep value.

    Parameters
    ----------
    number_of_timesteps: int
        Number of timesteps found so far.
    line: str
        'TITLE' line from the pdb file which has timestep information.

    Returns
    -------
    int
        Number of timesteps iterated by 1.
    float
        Time of newest timestep.
    """
    number_of_timesteps += 1
    new_time = float(line.split()[-1])
    return number_of_timesteps, new_time

def print_update(percentage):
    """Print percentage read-in.

    Prints a percentage of how much as been read in at a given time.

    Parameters
    ----------
    string: int
        Percentage read-in complete.
    """
    print("[{} {} % ]".format('#' * int(percentage / 10), int(percentage / 10) * 10))

def check_update(percentage, percentage_new):
    """Check if an update string should be printed.

    Assess if an update string should be printed to the screen.

    Parameters
    ----------
    string: int
        Percentage of way through reading.

    Returns
    -------
    int
        New percentage of way through reading.
    """
    if percentage_new > percentage + 9:
        percentage = percentage_new
        print_update(percentage)
    return percentage