# -*- coding: utf-8 -*-
"""
Image Processing Actions
========================

A toolkit of image processing actions for segmentation.

author: Ed Beard
email: ejb207@cam.ac.uk, ed.beard94@gmail.com

"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import logging

import numpy as np
import os
from skimage.util import pad
from skimage.measure import regionprops

import itertools
import copy
from scipy import ndimage as ndi
# mmaser: add center_of_mass import for centroid recording
#from ndi.measurements import center_of_mass
from sklearn.cluster import KMeans
import osra_rgroup

from .model import Panel, Diagram, Label, Rect, Figure
from .io import imsave, imdel
from .clean import find_repeating_unit, clean_output
from .utils import crop, skeletonize, binarize, binary_close, binary_floodfill, merge_rect, merge_overlap

# Standard path to superatom dictionary file
parent_dir = os.path.dirname(os.path.abspath(__file__))
superatom_file = os.path.join(parent_dir, 'dict', 'superatom.txt')
spelling_file = os.path.join(parent_dir, 'dict', 'spelling.txt')

log = logging.getLogger(__name__)

def segment(fig):
    """ Segments image.

    :param fig: Input Figure
    :return panels: List of segmented Panel objects
    """

    bin_fig = binarize(fig)

    bbox = fig.get_bounding_box()
    skel_pixel_ratio = skeletonize_area_ratio(fig, bbox)

    log.debug(" The skeletonized pixel ratio is %s" % skel_pixel_ratio)

    # Choose kernel size according to skeletonized pixel ratio
    if skel_pixel_ratio > 0.025:
        kernel = 4
        closed_fig = binary_close(bin_fig, size=kernel)
        log.debug("Segmentation kernel size = %s" % kernel)

    elif 0.02 < skel_pixel_ratio <= 0.025:
        kernel = 6
        closed_fig = binary_close(bin_fig, size=kernel)
        log.debug("Segmentation kernel size = %s" % kernel)

    elif 0.015 < skel_pixel_ratio <= 0.02:
        kernel = 10
        closed_fig = binary_close(bin_fig, size=kernel)
        log.debug("Segmentation kernel size = %s" % kernel)

    else:
        kernel = 15
        closed_fig = binary_close(bin_fig, size=kernel)
        log.debug("Segmentation kernel size = %s" % kernel)

    # Using a binary floodfill to identify panel regions
    # mmaser: try adding .centroids here
    fill_img = binary_floodfill(closed_fig)
    # mmaser: tag_img now has .centroids property which should be list of
    # len num panels (tag_img is the fig, has .img (labels) too)
    tag_img = binary_tag(fill_img)
    # mmaser: each panel in panels now has panel.centroid as (row, col)
    panels = get_bounding_box(tag_img)

    # Removing relatively tiny pixel islands that are determined to be noise
    area_threshold = fig.get_bounding_box().area / 200
    width_threshold = fig.get_bounding_box().width / 150
    panels = [panel for panel in panels if panel.area > area_threshold or panel.width > width_threshold]
    return panels

def classify_kmeans(panels, fig, skel=True):
    """Takes input image and classifies through k means cluster of the panel area"""

    if len(panels) <= 1:
        raise Exception('Only one panel detected. Cannot cluster')
    return get_labels_and_diagrams_k_means_clustering(panels, fig, skel)

def get_labels_and_diagrams_k_means_clustering(panels, fig, skel=True):
    """ Splits into labels and diagrams using K-means clustering by the skeletonized area ratio or panel height.

    :param panels: List of Panel objects to be clustered
    :param fig: Input Figure
    :param skel: Boolean indication the clustering parameters to use
    :return Lists of Labels and Diagrams after clustering
    """

    cluster_params = []

    for panel in panels:
        if skel:
            cluster_params.append([skeletonize_area_ratio(fig, panel)])
        else:
            cluster_params.append([panel.height])

    all_params = np.array(cluster_params)

    km = KMeans(n_clusters=2)
    clusters = km.fit(all_params)

    group_1, group_2 = [], []

    for i, cluster in enumerate(clusters.labels_):
        if cluster == 0:
            group_1.append(panels[i])
        else:
            group_2.append(panels[i])

    if np.nanmean([panel.area for panel in group_1]) > np.nanmean([panel.area for panel in group_2]):
        diags = group_1
        labels = group_2
    else:
        diags = group_2
        labels = group_1

    # Convert to appropriate types
    # mmaser: up to here labels are still panels with .centroid property
    labels = [Label(label.left, label.right, label.top, label.bottom, label.tag) for label in labels]
    diags = [Diagram(diag.left, diag.right, diag.top, diag.bottom, diag.tag) for diag in diags]
    return labels, diags

def preprocessing(labels, diags, fig):
    """Pre-processing steps before final K-means classification
    :param labels: List of Label objects
    :param diags: List of Diagram objects
    :param fig: Figure object

    :return out_labels: List of Labels after merging and re-tagging
    :return out_diags: List of Diagrams after re-tagging
    """

    # Remove repeating unit indicators
    labels, diags = find_repeating_unit(labels, diags, fig)

    # Remove small pixel islands from diagrams
    diags = remove_diag_pixel_islands(diags, fig)

    # Merge labels together that are sufficiently local
    label_candidates_horizontally_merged = merge_label_horizontally(labels,fig)
    label_candidates_fully_merged = merge_labels_vertically(label_candidates_horizontally_merged)
    labels_converted = convert_panels_to_labels(label_candidates_fully_merged)

    # Re-tagging all diagrams and labels
    retagged_panels = retag_panels(labels_converted + diags)
    out_labels = retagged_panels[:len(labels_converted)]
    out_diags = retagged_panels[len(labels_converted):]

    return out_labels, out_diags

def label_diags(labels, diags, fig_bbox):
    """ Pair all Diagrams to Labels.

    :param labels: List of Label objects
    :param diags: List of Diagram objects
    :param fig_bbox: Co-ordinates of the bounding box of the entire figure

    :returns: List of Diagrams with assigned Labels

    """

    # Sort diagrams from largest to smallest
    # TODO mmaser change to sort by max(x_position) (actually have tags,
    # see comments throughtout)
    # YES should still be in order with x.tag
    #diags.sort(key=lambda x: x.area, reverse=True)
    # check what assign_label_to_diag does in next step before changing below
    # should be okay, this step actually just expands bbox of each diag until
    # it overlaps any label object, so ordering of diag really is what matters
    #
    diags.sort(key=lambda x: x.tag, reverse=False)
    initial_sorting = [assign_label_to_diag(diag, labels, fig_bbox) for diag in diags]

    # Identify failures by the presence of duplicate labels
    failed_diag_label = get_duplicate_labelling(initial_sorting)

    if len(failed_diag_label) == 0:
        return initial_sorting

    # Find average position of label relative to diagram for successful pairings (denoted with compass points: NSEW)
    successful_diag_label = [diag for diag in diags if diag not in failed_diag_label]

    # Where no sucessful pairings found, attempt looking 'South' for all diagrams (most common relative label position)
    if len(successful_diag_label) == 0:
        altered_sorting = [assign_label_to_diag_postprocessing(diag, labels, 'S', fig_bbox) for diag in failed_diag_label]
        if len(get_duplicate_labelling(altered_sorting)) != 0:
            altered_sorting = initial_sorting
            pass
        else:
            return altered_sorting
    else:
        # Get compass positions of labels relative to diagram
        diag_compass = [diag.compass_position(diag.label) for diag in successful_diag_label if diag.label]
        mode_compass = max(diag_compass, key=diag_compass.count)

        # Expand outwards in compass direction for all failures
        altered_sorting = [assign_label_to_diag_postprocessing(diag, labels, mode_compass, fig_bbox) for diag in failed_diag_label]

        # Check for duplicates after relabelling
        failed_diag_label = get_duplicate_labelling(altered_sorting + successful_diag_label)
        successful_diag_label = [diag for diag in successful_diag_label if diag not in failed_diag_label]

        # If no duplicates return all results
        if len(failed_diag_label) == 0:
            return altered_sorting + successful_diag_label

    # Add non duplicates to successes
    successful_diag_label.extend([diag for diag in altered_sorting if diag not in failed_diag_label])

    # Remove duplicate results
    diags_with_labels, diags_without_labels = remove_duplicates(failed_diag_label, fig_bbox)

    return diags_with_labels + successful_diag_label

def assign_label_to_diag(diag, labels, fig_bbox, rate=1):
    """ Iteratively expands the bounding box of diagram until it intersects a Label object

    :param diag: Input Diagram object to expand from
    :param labels: List of Label objects
    :param fig_bbox: Panel object representing the co-ordinates for the entire Figure
    :param rate: Number of pixels to expand by upon each iteration

    :return diag: Diagram with Label object assigned
    """

    probe_rect = Rect(diag.left, diag.right, diag.top, diag.bottom)
    found = False
    max_threshold_width = fig_bbox.width
    max_threshold_height = fig_bbox.height

    while found is False and (probe_rect.width < max_threshold_width or probe_rect.height < max_threshold_height):
        # Increase border value each loop
        probe_rect.right = probe_rect.right + rate
        probe_rect.bottom = probe_rect.bottom + rate
        probe_rect.left = probe_rect.left - rate
        probe_rect.top = probe_rect.top - rate

        for label in labels:
            if probe_rect.overlaps(label):
                found = True
                diag.label = label
    return diag

def assign_label_to_diag_postprocessing(diag, labels, direction, fig_bbox, rate=1):
    """ Iteratively expands the bounding box of diagram in the specified compass direction

    :param diag: Input Diagram object to expand from
    :param labels: List of Label objects
    :param direction: String representing determined compass direction (allowed values: 'E', 'S', 'W', 'N')
    :param fig_bbox: Panel object representing the co-ordinates for the entire Figure
    :param rate: Number of pixels to expand by upon each iteration
    """

    probe_rect = Rect(diag.left, diag.right, diag.top, diag.bottom)
    found = False

    def label_loop():

        for label in labels:
            # Only accepting labels in the average direction
            if diag.compass_position(label) != direction:
                pass
            elif probe_rect.overlaps(label):
                diag.label = label
                return True

        return False

    # Increase border value each loop
    if direction == 'E':
        while found is False and probe_rect.right < fig_bbox.right:
            probe_rect.right = probe_rect.right + rate
            found = label_loop()

    elif direction == 'S':
        while found is False and probe_rect.bottom < fig_bbox.bottom:
            probe_rect.bottom = probe_rect.bottom + rate
            found = label_loop()

    elif direction == 'W':
        while found is False and probe_rect.left > fig_bbox.left:
            probe_rect.left = probe_rect.left - rate
            found = label_loop()

    elif direction == 'N':
        while found is False and probe_rect.top > fig_bbox.top:
            probe_rect.top = probe_rect.top - rate
            found = label_loop()
    else:
        return diag

    return diag

def read_diagram_pyosra(diag, extension='jpg', debug=False, superatom_path=superatom_file, spelling_path=spelling_file):
    """ Converts a diagram to SMILES using pyosra

    :param diag: Diagram to be extracted
    :param extension: String file extension
    :param debug: Bool inicating debug mode

    :return smile: String of extracted chemical SMILE

    """

    # Add some padding to image to help resolve characters on the edge
    padded_img = pad(diag.fig.img, ((5, 5), (5, 5), (0, 0)), mode='constant', constant_values=1)

    # Save a temp image
    temp_img_fname = 'osra_temp.' + extension
    imsave(temp_img_fname, padded_img)

    # Run osra on temp image
    smile = osra_rgroup.read_diagram(temp_img_fname, debug=debug, superatom_file=superatom_path, spelling_file=spelling_path)

    if not smile:
        log.warning('No SMILES string was extracted for diagram %s' % diag.tag)

    if not debug:
        imdel(temp_img_fname)

    # mmaser addition to remove spaces from smiles outputs. clean_output
    # function in clean.py has this line commented out
    smile = smile.replace(' ', '')
    smile = clean_output(smile)
    return smile

def remove_diag_pixel_islands(diags, fig):
    """ Removes small pixel islands from the diagram

    :param diags: List of input Diagrams
    :param fig: Figure object

    :return diags: List of Diagrams with small pixel islands removed

    """

    for diag in diags:

        # Make a cleaned copy of image to be used when resolving diagrams
        clean_fig = copy.deepcopy(fig)

        diag_fig = Figure(crop(clean_fig.img, diag.left, diag.right, diag.top, diag.bottom))
        seg_fig = Figure(crop(clean_fig.img, diag.left, diag.right, diag.top, diag.bottom))
        sub_panels = segment(seg_fig)

        panel_areas = [panel.area for panel in sub_panels]
        diag_area = max(panel_areas)

        sub_panels = [panel for panel in sub_panels if panel.area != diag_area]

        sub_bbox = [(panel.left, panel.right, panel.top, panel.bottom) for panel in sub_panels]

        for bbox in sub_bbox:
            diag_fig.img[bbox[2]:bbox[3], bbox[0]:bbox[1]] = np.ones(3)

        diag.fig = diag_fig

    return diags

def pixel_ratio(fig, diag):
    """ Calculates the ratio of 'on' pixels to bounding box area for binary figure

    :param fig : Input binary Figure
    :param diag : Area to calculate pixel ratio

    :return ratio: Float detailing ('on' pixels / bounding box area)
    """

    cropped_img = crop(fig.img, diag.left, diag.right, diag.top, diag.bottom)
    ones = np.count_nonzero(cropped_img)
    all_pixels = np.size(cropped_img)
    ratio = ones / all_pixels
    return ratio

def binary_tag(fig):
    """ Tag connected regions with pixel value of 1

    :param fig: Input Figure
    :returns fig: Connected Figure
    """
    fig.img, no_tagged = ndi.label(fig.img)
    return fig

def get_bounding_box(fig):
    """ Gets the bounding box of each segment

    :param fig: Input Figure
    :returns panels: List of panel objects
    # mmaser note can supply fig as second arg to regionprops, so
    # regions = regionprops(fig.img, intensity_image=fig)
    #
    """
    panels = []
    regions = regionprops(fig.img)
    for region in regions:
        # mmaser y1, x1, y2, x2 = (top, left, bottom, right),
        y1, x1, y2, x2 = region.bbox
        # mmaser: Panel object now has .centroid property
        panels.append(Panel(x1, x2, y1, y2, region.label - 1))# Sets tags to start from 0
    return panels

def retag_panels(panels):
    """ Re-tag panels.

    :param panels: List of Panel objects
    :returns: List of re-tagged Panel objects
    """

    for i, panel in enumerate(panels):
        panel.tag = i
    return panels

def skeletonize_area_ratio(fig, panel):
    """ Calculates the ratio of skeletonized image pixels to total number of pixels

    :param fig: Input figure
    :param panel: Original panel object
    :return: Float : Ratio of skeletonized pixels to total area (see pixel_ratio)
    """

    skel_fig = skeletonize(fig)
    return pixel_ratio(skel_fig, panel)

def order_by_area(panels):
    """ Returns a list of panel objects ordered by area.

    :param panels: Input list of Panels
    :return panels: Output list of sorted Panels
    """

    def get_area(panel):
        return panel.area

    panels.sort(key=get_area)
    return panels

def merge_label_horizontally(merge_candidates, fig):
    """ Iteratively attempt to merge horizontally

    :param merge_candidates: Input list of Panels to be merged
    :return merge_candidates: List of Panels after merging
    """

    done = False

    # Identifies panels within horizontal merging criteria
    while done is False:
        ordered_panels = order_by_area(merge_candidates)
        merge_candidates, done = merge_loop_horizontal(ordered_panels, fig)

    merge_candidates, done = merge_all_overlaps(merge_candidates)
    return merge_candidates

def merge_labels_vertically(merge_candidates):
    """ Iteratively attempt to merge vertically

    :param merge_candidates: Input list of Panels to be merged
    :return merge_candidates: List of Panels after merging
    """

    # Identifies panels within horizontal merging criteria
    ordered_panels = order_by_area(merge_candidates)
    merge_candidates = merge_loop_vertical(ordered_panels)

    merge_candidates, done = merge_all_overlaps(merge_candidates)
    return merge_candidates

def merge_loop_horizontal(panels, fig_input):
    """ Iteratively merges panels by relative proximity to each other along the x axis.
        This is repeated until no panels are merged by the algorithm

    :param panels: List of Panels to be merged.

    :return output_panels: List of merged panels
    :return done: Bool indicating whether a merge occurred
    """

    output_panels = []
    blacklisted_panels = []
    done = True

    for a, b in itertools.combinations(panels, 2):

        # Check panels lie in roughly the same line, that they are of label size and similar height
        if abs(a.center[1] - b.center[1]) < 1.5 * a.height \
                and abs(a.height - b.height) < min(a.height, b.height):

            # Check that the distance between the edges of panels is not too large
            if (0 <= a.left - b.right < (min(a.height, b.height) * 2)) or (0 <= (b.left - a.right) < (min(a.height, b.height) * 2)):

                merged_rect = merge_rect(a, b)
                merged_panel = Panel(merged_rect.left, merged_rect.right, merged_rect.top, merged_rect.bottom, 0)
                output_panels.append(merged_panel)
                blacklisted_panels.extend([a, b])
                done = False

    log.debug('Length of blacklisted : %s' % len(blacklisted_panels))
    log.debug('Length of output panels : %s' % len(output_panels))

    for panel in panels:
        if panel not in blacklisted_panels:
            output_panels.append(panel)

    output_panels = retag_panels(output_panels)

    return output_panels, done

def merge_loop_vertical(panels):
    """ Iteratively merges panels by relative proximity to each other along the y axis.
        This is repeated until no panels are merged by the algorithm

    :param panels: List of Panels to be merged.

    :return output_panels: List of merged panels
    :return done: Bool indicating whether a merge occurred
    """

    output_panels = []
    blacklisted_panels = []

    # Merging labels that are in close proximity vertically
    for a, b in itertools.combinations(panels, 2):

        if (abs(a.left - b.left) < 3 * min(a.height, b.height) or abs(a.center[0] - b.center[0]) < 3 * min(a.height, b.height)) \
                and abs(a.center[1] - b.center[1]) < 3 * min(a.height, b.height) \
                and min(abs(a.top - b.bottom), abs(b.top - a.bottom)) < 2 * min(a.height, b.height):

            merged_rect = merge_rect(a, b)
            merged_panel = Panel(merged_rect.left, merged_rect.right, merged_rect.top, merged_rect.bottom, 0)
            output_panels.append(merged_panel)
            blacklisted_panels.extend([a, b])

    for panel in panels:
        if panel not in blacklisted_panels:
            output_panels.append(panel)

    output_panels = retag_panels(output_panels)

    return output_panels

def get_one_to_merge(all_combos, panels):
    """Merges the first overlapping set of panels found and an returns updated panel list

    :param all_combos: List of Tuple(Panel, Panel) objects of all possible combinations of the input 'panels' variable
    :param panels: List of input Panels

    :return panels: List of updated panels after one overlap is merged
    :return: Bool indicated whether all overlaps have been completed
    """

    for a, b in all_combos:

        overlap_panel = merge_overlap(a, b)
        if overlap_panel is not None:
            merged_panel = Panel(overlap_panel.left, overlap_panel.right, overlap_panel.top, overlap_panel.bottom, 0)
            panels.remove(a)
            panels.remove(b)
            panels.append(merged_panel)
            return panels, False

    return panels, True

def convert_panels_to_labels(panels):
    """ Converts a list of panels to a list of labels

    :param panels: Input list of Panels
    :return : List of Labels
    """

    return [Label(panel.left, panel.right, panel.top, panel.bottom, panel.tag) for panel in panels]

def merge_all_overlaps(panels):
    """ Merges all overlapping rectangles together

    :param panels : Input list of Panels
    :return output_panels: List of merged panels
    :return all_merged: Bool indicating whether all merges are completed
    """

    all_merged = False

    while all_merged is False:
        all_combos = list(itertools.combinations(panels, 2))
        panels, all_merged = get_one_to_merge(all_combos, panels)

    output_panels = retag_panels(panels)
    return output_panels, all_merged

def get_duplicate_labelling(labelled_diags):
    """ Returns diagrams sharing a Label object with other diagrams.

    :param labelled_diags: List of Diagrams with Label objects assigned
    :return failed_diag_label: List of Diagrams that share a Label object with another Diagram
    """

    failed_diag_label = set(diag for diag in labelled_diags if not diag.label)
    filtered_labelled_diags = [diag for diag in labelled_diags if diag not in failed_diag_label]

    # Identifying cases with the same label:
    for a, b in itertools.combinations(filtered_labelled_diags, 2):
        if a.label == b.label:
            failed_diag_label.add(a)
            failed_diag_label.add(b)

    return failed_diag_label

def remove_duplicates(diags, fig_bbox, rate=1):
    """
    Removes the least likely of the duplicate results.
    Likeliness is determined from the distance from the bounding box

    :param diags: All detected diagrams with assigned labels
    :param fig_bbox: Panel object representing the co-ordinates for the entire Figure
    :param rate: Number of pixels to expand by upon each iteration

    :return output_diags : List of Diagrams with Labels
    :return output_labelless_diags : List of Diagrams with Labels removed due to duplicates
    """

    output_diags = []
    output_labelless_diags = []

    # Unique labels
    unique_labels = set(diag.label for diag in diags if diag.label is not None)

    for label in unique_labels:

        diags_with_labels = [diag for diag in diags if diag.label is not None]
        diags_with_this_label = [diag for diag in diags_with_labels if diag.label.tag == label.tag]

        if len(diags_with_this_label) == 1:
            output_diags.append(diags_with_this_label[0])
            continue

        diag_and_displacement = [] # List of diag-distance tuples

        for diag in diags_with_this_label:

            probe_rect = Rect(diag.left, diag.right, diag.top, diag.bottom)
            found = False
            max_threshold_width = fig_bbox.width
            max_threshold_height = fig_bbox.height
            rate_counter = 0

            while found is False and (probe_rect.width < max_threshold_width or probe_rect.height < max_threshold_height):
                # Increase border value each loop
                probe_rect.right = probe_rect.right + rate
                probe_rect.bottom = probe_rect.bottom + rate
                probe_rect.left = probe_rect.left - rate
                probe_rect.top = probe_rect.top - rate

                rate_counter += rate

                if probe_rect.overlaps(label):
                    found = True
                    diag_and_displacement.append((diag, rate_counter))

        master_diag = min(diag_and_displacement, key=lambda x: x[1])[0]
        output_diags.append(master_diag)

        labelless_diags = [diag[0] for diag in diag_and_displacement if diag[0] is not master_diag]

        for diag in labelless_diags:
            diag.label = None

        output_labelless_diags.extend(labelless_diags)

    return output_diags, output_labelless_diags


# --------------------------------------------------------
# Sequence Level Semantics Aggregation
# Copyright (c) 2016 by Contributors
# Copyright (c) 2017 Microsoft
# Licensed under The Apache-2.0 License [see LICENSE for details]
# Modified by Yuwen Xiong
# --------------------------------------------------------
"""
Fast R-CNN:
data =
    {'data': [num_images, c, h, w],
    'rois': [num_rois, 5]}
label =
    {'label': [num_rois],
    'bbox_target': [num_rois, 4 * num_classes],
    'bbox_weight': [num_rois, 4 * num_classes]}
roidb extended format [image_index]
    ['image', 'height', 'width', 'flipped',
     'boxes', 'gt_classes', 'gt_overlaps', 'max_classes', 'max_overlaps', 'bbox_targets']
"""

import numpy as np
import numpy.random as npr

from utils.image import get_image, tensor_vstack
from bbox.bbox_transform import bbox_overlaps, bbox_transform
from bbox.bbox_regression import expand_bbox_regression_targets

def get_rcnn_testbatch(roidb, cfg):
    """
    return a dict of testbatch
    :param roidb: ['image', 'flipped'] + ['boxes']
    :return: data, label, im_info
    """
    # assert len(roidb) == 1, 'Single batch only'
    imgs, roidb = get_image(roidb, cfg)
    im_array = imgs
    im_info = [np.array([roidb[i]['im_info']], dtype=np.float32) for i in range(len(roidb))]

    im_rois = [roidb[i]['boxes'] for i in range(len(roidb))]
    rois = im_rois
    rois_array = [np.hstack((0 * np.ones((rois[i].shape[0], 1)), rois[i])) for i in range(len(rois))]

    data = [{'data': im_array[i],
             'rois': rois_array[i]} for i in range(len(roidb))]
    label = {}

    return data, label, im_info

def get_rcnn_batch(roidb, cfg):
    """
    return a dict of multiple images
    :param roidb: a list of dict, whose length controls batch size
    ['images', 'flipped'] + ['gt_boxes', 'boxes', 'gt_overlap'] => ['bbox_targets']
    :return: data, label
    """
    num_images = len(roidb)
    imgs, roidb = get_image(roidb, cfg)
    im_array = tensor_vstack(imgs)

    assert cfg.TRAIN.BATCH_ROIS == -1 or cfg.TRAIN.BATCH_ROIS % cfg.TRAIN.BATCH_IMAGES == 0, \
        'BATCHIMAGES {} must divide BATCH_ROIS {}'.format(cfg.TRAIN.BATCH_IMAGES, cfg.TRAIN.BATCH_ROIS)

    if cfg.TRAIN.BATCH_ROIS == -1:
        rois_per_image = np.sum([iroidb['boxes'].shape[0] for iroidb in roidb])
        fg_rois_per_image = rois_per_image
    else:
        rois_per_image = cfg.TRAIN.BATCH_ROIS / cfg.TRAIN.BATCH_IMAGES
        fg_rois_per_image = np.round(cfg.TRAIN.FG_FRACTION * rois_per_image).astype(int)

    rois_array = list()
    labels_array = list()
    bbox_targets_array = list()
    bbox_weights_array = list()

    for im_i in range(num_images):
        roi_rec = roidb[im_i]

        # infer num_classes from gt_overlaps
        num_classes = roi_rec['gt_overlaps'].shape[1]

        # label = class RoI has max overlap with
        rois = roi_rec['boxes']
        labels = roi_rec['max_classes']
        overlaps = roi_rec['max_overlaps']
        bbox_targets = roi_rec['bbox_targets']

        im_rois, labels, bbox_targets, bbox_weights = \
            sample_rois(rois, fg_rois_per_image, rois_per_image, num_classes, cfg,
                        labels, overlaps, bbox_targets)

        # project im_rois
        # do not round roi
        rois = im_rois
        batch_index = im_i * np.ones((rois.shape[0], 1))
        rois_array_this_image = np.hstack((batch_index, rois))
        rois_array.append(rois_array_this_image)

        # add labels
        labels_array.append(labels)
        bbox_targets_array.append(bbox_targets)
        bbox_weights_array.append(bbox_weights)

    rois_array = np.array(rois_array)
    labels_array = np.array(labels_array)
    bbox_targets_array = np.array(bbox_targets_array)
    bbox_weights_array = np.array(bbox_weights_array)

    data = {'data': im_array,
            'rois': rois_array}
    label = {'label': labels_array,
             'bbox_target': bbox_targets_array,
             'bbox_weight': bbox_weights_array}

    return data, label

def sample_rois(rois, fg_rois_per_image, rois_per_image, num_classes, cfg,
                labels=None, overlaps=None, bbox_targets=None, gt_boxes=None, output_sort=False):
    """
    generate random sample of ROIs comprising foreground and background examples
    :param rois: all_rois [n, 4]; e2e: [n, 5] with batch_index
    :param fg_rois_per_image: foreground roi number
    :param rois_per_image: total roi number
    :param num_classes: number of classes
    :param labels: maybe precomputed
    :param overlaps: maybe precomputed (max_overlaps)
    :param bbox_targets: maybe precomputed
    :param gt_boxes: optional for e2e [n, 5] (x1, y1, x2, y2, cls)
    :return: (labels, rois, bbox_targets, bbox_weights)
    """
    if labels is None:
        overlaps = bbox_overlaps(rois[:, 1:].astype(np.float), gt_boxes[:, :4].astype(np.float))
        gt_assignment = overlaps.argmax(axis=1)
        overlaps = overlaps.max(axis=1)
        labels = gt_boxes[gt_assignment, 4]

    # foreground RoI with FG_THRESH overlap
    fg_indexes = np.where(overlaps >= cfg.TRAIN.FG_THRESH)[0]
    # guard against the case when an image has fewer than fg_rois_per_image foreground RoIs
    fg_rois_per_this_image = np.minimum(fg_rois_per_image, fg_indexes.size)
    # Sample foreground regions without replacement
    if len(fg_indexes) > fg_rois_per_this_image:
        fg_indexes = npr.choice(fg_indexes, size=fg_rois_per_this_image, replace=False)

    if output_sort:
        _overlaps = overlaps[fg_indexes]
        _order = np.argsort(_overlaps)[::-1]
        fg_indexes = fg_indexes[_order]

    # Select background RoIs as those within [BG_THRESH_LO, BG_THRESH_HI)
    bg_indexes = np.where((overlaps < cfg.TRAIN.BG_THRESH_HI) & (overlaps >= cfg.TRAIN.BG_THRESH_LO))[0]
    # Compute number of background RoIs to take from this image (guarding against there being fewer than desired)
    bg_rois_per_this_image = rois_per_image - fg_rois_per_this_image
    bg_rois_per_this_image = np.minimum(bg_rois_per_this_image, bg_indexes.size)
    # Sample foreground regions without replacement
    if len(bg_indexes) > bg_rois_per_this_image:
        bg_indexes = npr.choice(bg_indexes, size=bg_rois_per_this_image, replace=False)

    # indexes selected
    keep_indexes = np.append(fg_indexes, bg_indexes)

    # pad more to ensure a fixed minibatch size
    while keep_indexes.shape[0] < rois_per_image:
        gap = np.minimum(len(rois), rois_per_image - keep_indexes.shape[0])
        gap_indexes = npr.choice(range(len(rois)), size=gap, replace=False)
        keep_indexes = np.append(keep_indexes, gap_indexes)

    # select labels
    labels = labels[keep_indexes]
    # set labels of bg_rois to be 0
    labels[fg_rois_per_this_image:] = 0
    rois = rois[keep_indexes]

    # load or compute bbox_target
    if bbox_targets is not None:
        bbox_target_data = bbox_targets[keep_indexes, :]
    else:
        targets = bbox_transform(rois[:, 1:], gt_boxes[gt_assignment[keep_indexes], :4])
        if cfg.TRAIN.BBOX_NORMALIZATION_PRECOMPUTED:
            targets = ((targets - np.array(cfg.TRAIN.BBOX_MEANS))
                       / np.array(cfg.TRAIN.BBOX_STDS))
        bbox_target_data = np.hstack((labels[:, np.newaxis], targets))

    bbox_targets, bbox_weights = \
        expand_bbox_regression_targets(bbox_target_data, num_classes, cfg)

    return rois, labels, bbox_targets, bbox_weights


#-*- coding:utf-8 -*-
import os
import json
import threading
import numpy as np
from PIL import Image

import tensorflow as tf
from keras import losses
from keras import backend as K
from keras.utils import plot_model
from keras.preprocessing import image
from keras.preprocessing.sequence import pad_sequences
from keras.layers import Input, Dense, Flatten
from keras.layers.core import Reshape, Masking, Lambda, Permute
from keras.layers.recurrent import GRU, LSTM
from keras.layers.wrappers import Bidirectional, TimeDistributed
from keras.layers.normalization import BatchNormalization
from keras.layers.convolutional import Conv2D, MaxPooling2D, ZeroPadding2D
from keras.optimizers import SGD, Adam
from keras.models import Model
from keras.callbacks import EarlyStopping, ModelCheckpoint, LearningRateScheduler, TensorBoard

from imp import reload
import densenet
from keras.utils import multi_gpu_model

GPU_NUM = 2

img_h = 32
img_w = 280*2
batch_size = 64
maxlabellength = 10*2

def get_session(gpu_fraction=1.0):

    num_threads = os.environ.get('OMP_NUM_THREADS')
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_fraction)

    if num_threads:
        if GPU_NUM > 1:
            return tf.Session(config=tf.ConfigProto(
                gpu_options=gpu_options, intra_op_parallelism_threads=num_threads, allow_soft_placement=True))
        else:
            return tf.Session(config=tf.ConfigProto(
                gpu_options=gpu_options, intra_op_parallelism_threads=num_threads))
    else:
        if GPU_NUM > 1:
            return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
        else:
            return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, allow_soft_placement=True))

def readfile(filename):
    res = []
    with open(filename, 'r') as f:
        lines = f.readlines()
        for i in lines:
            res.append(i.strip())
    dic = {}
    for i in res:
        p = i.split(' ')
        dic[p[0]] = p[1:]
    return dic

class random_uniform_num():
    """
    均匀随机，确保每轮每个只出现一次
    """
    def __init__(self, total):
        self.total = total
        self.range = [i for i in range(total)]
        np.random.shuffle(self.range)
        self.index = 0
    def get(self, batchsize):
        r_n=[]
        if(self.index + batchsize > self.total):
            r_n_1 = self.range[self.index:self.total]
            np.random.shuffle(self.range)
            self.index = (self.index + batchsize) - self.total
            r_n_2 = self.range[0:self.index]
            r_n.extend(r_n_1)
            r_n.extend(r_n_2)
        else:
            r_n = self.range[self.index : self.index + batchsize]
            self.index = self.index + batchsize

        return r_n

def gen(data_file, image_path, batchsize=128, maxlabellength=10, imagesize=(32, 280)):
    image_label = readfile(data_file)
    _imagefile = [i for i, j in image_label.items()]
    x = np.zeros((batchsize, imagesize[0], imagesize[1], 1), dtype=np.float)
    labels = np.ones([batchsize, maxlabellength]) * 10000
    input_length = np.zeros([batchsize, 1])
    label_length = np.zeros([batchsize, 1])

    r_n = random_uniform_num(len(_imagefile))
    _imagefile = np.array(_imagefile)
    while 1:
        shufimagefile = _imagefile[r_n.get(batchsize)]
        for i, j in enumerate(shufimagefile):
            img1 = Image.open(os.path.join(image_path, j)).convert('L')
            img = np.array(img1, 'f') / 255.0 - 0.5

            x[i] = np.expand_dims(img, axis=2)
            # print('imag:shape', img.shape)
            str = image_label[j]
            label_length[i] = len(str)

            if(len(str) <= 0):
                print("len < 0", j)
            input_length[i] = imagesize[1] // 8
            labels[i, :len(str)] = [int(k) - 1 for k in str]

        inputs = {'the_input': x,
                'the_labels': labels,
                'input_length': input_length,
                'label_length': label_length,
                }
        outputs = {'ctc': np.zeros([batchsize])}
        yield (inputs, outputs)

def ctc_lambda_func(args):
    y_pred, labels, input_length, label_length = args
    return K.ctc_batch_cost(labels, y_pred, input_length, label_length)

def get_model(img_h, nclass):
    input = Input(shape=(img_h, None, 1), name='the_input')
    y_pred = densenet.dense_cnn(input, nclass)
    # y_pred = densenet.dense_blstm(input, nclass)

    basemodel = Model(inputs=input, outputs=y_pred)
    basemodel.summary()

    labels = Input(name='the_labels', shape=[None], dtype='float32')
    input_length = Input(name='input_length', shape=[1], dtype='int64')
    label_length = Input(name='label_length', shape=[1], dtype='int64')

    loss_out = Lambda(ctc_lambda_func, output_shape=(1,), name='ctc')([y_pred, labels, input_length, label_length])

    model = Model(inputs=[input, labels, input_length, label_length], outputs=loss_out)

    if GPU_NUM > 1:
        model = multi_gpu_model(model, gpus=GPU_NUM)

    model.compile(loss={'ctc': lambda y_true, y_pred: y_pred}, optimizer='adam', metrics=['accuracy'])

    return basemodel, model

if __name__ == '__main__':
    # char_set = open('char_std_5990.txt', 'r', encoding='utf-8').readlines()
    char_set = open('char_std_96_english_v1.txt', 'r', encoding='utf-8').readlines()
    char_set = ''.join([ch.strip('\n') if ch.strip('\n')!='' else ' ' for ch in char_set][1:] + ['卍'])
    print('char_set:',char_set)
    nclass = len(char_set)
    print('~~~~~~~~~~~~nclass:',nclass)

    K.set_session(get_session())
    reload(densenet)
    basemodel, model = get_model(img_h, nclass)

    modelPath = './models/weights_densenet_v2-08-2.77--.h5'
    if os.path.exists(modelPath):
        print("Loading model weights...")
        model.load_weights(modelPath, by_name=True, skip_mismatch=True)
        print('done!')

    basemodelPath = './models/pretrain_model/weights_densenet.h5'
    if os.path.exists(basemodelPath) and not os.path.exists(modelPath):
        print("Loading basemodelPath weights...")
        basemodel.load_weights(basemodelPath, by_name=True, skip_mismatch=True)
        print('done!')

    base_train_path = '/dockershare/hmbdata2/'
    base_test_path = '/dockershare/hmbdata3/'

    train_loader = gen(base_train_path + 'total_label_0-100.txt', base_train_path, batchsize=batch_size, maxlabellength=maxlabellength, imagesize=(img_h, img_w))
    test_loader = gen(base_test_path + 'total_label_0-100.txt', base_test_path, batchsize=batch_size, maxlabellength=maxlabellength, imagesize=(img_h, img_w))

    checkpoint = ModelCheckpoint(filepath='./models/weights_densenet_v3-{epoch:02d}-{val_loss:.2f}.h5', monitor='val_loss', save_best_only=False, save_weights_only=True)
    lr_schedule = lambda epoch: 0.0005 * 0.4**epoch
    learning_rate = np.array([lr_schedule(i) for i in range(10)])
    changelr = LearningRateScheduler(lambda epoch: float(learning_rate[epoch]))
    earlystop = EarlyStopping(monitor='val_loss', patience=2, verbose=1)
    tensorboard = TensorBoard(log_dir='./models/logs', write_graph=True)

    print('-----------Start training-----------')
    model.fit_generator(train_loader,
    	steps_per_epoch =  181999// batch_size,
    	epochs = 150,
    	initial_epoch = 8,
    	validation_data = test_loader,
    	validation_steps = 181999 // batch_size,
    	# callbacks = [checkpoint, changelr, tensorboard])
    callbacks = [checkpoint, earlystop, changelr, tensorboard])


from __future__ import with_statement
import numpy
import numpy as np
import flydra_core.reconstruct
import flydra_core._reconstruct_utils as ru
#import flydra_core.geom as geom
import flydra_core._fastgeom as geom
import time
from flydra_analysis.analysis.result_utils import get_caminfo_dicts, \
     get_fps, read_textlog_header
import tables
import tables as PT
import warnings
warnings.filterwarnings('ignore', category=tables.NaturalNameWarning)
import os, sys, pprint, tempfile, shutil
from flydra_core.kalman.flydra_tracker import Tracker
import flydra_core.kalman.flydra_kalman_utils as flydra_kalman_utils
from optparse import OptionParser
import flydra_core.kalman.dynamic_models as dynamic_models
import collections
import flydra_core.version
from flydra_core.data_descriptions import TextLogDescription
from flydra_core.reconstruct import do_3d_operations_on_2d_point
import flydra_analysis.a2.utils as utils
from flydra_analysis.a2.tables_tools import open_file_safe

# Not really "observations" but ML estimates
FilteredObservations = flydra_kalman_utils.FilteredObservations
convert_format = flydra_kalman_utils.convert_format
tmp = flydra_kalman_utils.ML_estimates_2d_idxs_type
ML_estimates_2d_idxs_type = tmp
del tmp

def process_frame(reconstructor,tracker,frame,frame_data,camn2cam_id,
                  debug=0):
    if debug is None:
        debug=0
    frame_data = tracker.calculate_a_posteriori_estimates(
        frame,frame_data,camn2cam_id,debug2=debug)

    frame_data = tracker.remove_duplicate_detections(frame,frame_data)
    max_err=tracker.kalman_model['hypothesis_test_max_acceptable_error']

    # Now, tracked objects have been updated (and their 2D data points
    # removed from consideration), so we can use old flydra
    # "hypothesis testing" algorithm on remaining data to see if there
    # are new objects.

    if debug>1:
        print 'for frame %d: data not gobbled:'%(frame,)
        pprint.pprint(dict(frame_data))
        print

    # Convert to format accepted by find_best_3d()
    found_data_dict,first_idx_by_cam_id = convert_format(
        frame_data,
        camn2cam_id,
        area_threshold=tracker.area_threshold)
    # found_data_dict contains undistorted points

    hypothesis_test_found_point = False
    # test to short-circuit rest of function
    if len(found_data_dict) >= 2:

        with_water = reconstructor.wateri is not None
        # Can only do 3D math with at least 2 cameras giving good
        # data.
        try:
            (this_observation_3d, this_observation_Lcoords, cam_ids_used,
             min_mean_dist) = ru.hypothesis_testing_algorithm__find_best_3d(
                reconstructor,
                found_data_dict,
                max_err,
                debug=debug,
                with_water=with_water)
        except ru.NoAcceptablePointFound:
            pass
        else:
            hypothesis_test_found_point = True

    if hypothesis_test_found_point:

        if debug > 5:
            print 'found new point using hypothesis testing:'
            print 'this_observation_3d',this_observation_3d
            print 'cam_ids_used',cam_ids_used
            print 'min_mean_dist',min_mean_dist

        believably_new = tracker.is_believably_new(
            this_observation_3d, debug=debug)
        if (debug > 5):
            print 'believably_new',believably_new

        if believably_new:
            assert min_mean_dist<max_err
            if debug > 5:
                print 'accepting point'

            # make mapping from cam_id to camn
            cam_id2camn = {}
            for camn in camn2cam_id:
                if camn not in frame_data:
                    continue # this camn not used this frame, ignore
                cam_id = camn2cam_id[camn]
                if cam_id in cam_id2camn:
                    print '*'*80
                    print """

ERROR: It appears that you have >1 camn for a cam_id at a certain
frame. This almost certainly means that you are using a data file
recorded with an older version of flydra.MainBrain and that the
cameras were re-synchronized during the saving of a data file. You
will have to manually find out which camns to ignore (use
flydra_analysis_print_camera_summary) and then use the --exclude-camns
option to this program.

"""
                    print '*'*80
                    print
                    print 'frame',frame
                    print 'camn',camn
                    print 'frame_data',frame_data
                    print
                    print 'cam_id2camn',cam_id2camn
                    print 'camn2cam_id',camn2cam_id
                    print
                    raise ValueError('cam_id already in dict')
                cam_id2camn[cam_id]=camn

            # find camns
            this_observation_camns = [
                cam_id2camn[cam_id] for cam_id in cam_ids_used]

            this_observation_idxs = [
                first_idx_by_cam_id[cam_id] for cam_id in cam_ids_used]

            if debug>5:
                print 'this_observation_camns',this_observation_camns
                print 'this_observation_idxs',this_observation_idxs

                print 'camn','raw 2d data','reprojected 3d->2d'
                for camn in this_observation_camns:
                    cam_id = camn2cam_id[camn]
                    repro=reconstructor.find2d(
                        cam_id, this_observation_3d )
                    print camn,frame_data[camn][0][0][:2],repro

            ####################################
            #  Now join found point into Tracker
            tracker.join_new_obj( frame,
                                  this_observation_3d,
                                  this_observation_Lcoords,
                                  this_observation_camns,
                                  this_observation_idxs,
                                  debug=debug,
                                  )
    if debug > 5:
        print
        print 'At end of frame %d, all live tracked objects:'%frame
        tracker.debug_info(level=debug)
        print
        print '-'*80
    elif debug > 2:
        print 'At end of frame %d, all live tracked objects:'%frame
        tracker.debug_info(level=debug)
        print

class KalmanSaver:
    def __init__(self,
                 h5file,
                 reconstructor,
                 cam_id2camns=None,
                 min_observations_to_save=0,
                 textlog_save_lines = None,
                 dynamic_model_name=None,
                 dynamic_model=None,
                 fake_timestamp=None,
                 debug=False):
        self.cam_id2camns = cam_id2camns
        self.min_observations_to_save = min_observations_to_save
        self.debug = debug

        self.kalman_saver_info_instance = flydra_kalman_utils.KalmanSaveInfo(
            name=dynamic_model_name)
        kalman_estimates_description = (
            self.kalman_saver_info_instance.get_description())

        filters = tables.Filters(1, complib='lzo') # compress

        self.h5file = h5file
        reconstructor.save_to_h5file(self.h5file)
        self.h5_xhat = self.h5file.create_table(
            self.h5file.root,'kalman_estimates',
            kalman_estimates_description,
            "Kalman a posteriori estimates of tracked object",
            filters=filters)
        self.h5_xhat.attrs.dynamic_model_name = dynamic_model_name
        self.h5_xhat.attrs.dynamic_model = dynamic_model

        self.h5_obs = self.h5file.create_table(
            self.h5file.root,'ML_estimates', FilteredObservations,
            "observations of tracked object",filters=filters)

        self.h5_2d_obs_next_idx = 0

        # Note that ML_estimates_2d_idxs_type() should
        # match dtype with tro.observations_2d.

        self.h5_2d_obs = self.h5file.create_vlarray(
            self.h5file.root,
            'ML_estimates_2d_idxs',
            ML_estimates_2d_idxs_type(),
            "camns and idxs")

        self.obj_id = 0

        self.h5textlog = self.h5file.create_table(
            self.h5file.root,'textlog',TextLogDescription,'text log')

        if 1:
            textlog_row = self.h5textlog.row
            cam_id = 'mainbrain'
            if fake_timestamp is None:
                timestamp = time.time()
            else:
                timestamp = fake_timestamp

            list_of_textlog_data = [
                (timestamp,cam_id,timestamp,text)
                for text in textlog_save_lines ]
            for textlog_data in list_of_textlog_data:
                (mainbrain_timestamp,cam_id,host_timestamp,message
                 ) = textlog_data
                textlog_row['mainbrain_timestamp'] = mainbrain_timestamp
                textlog_row['cam_id'] = cam_id
                textlog_row['host_timestamp'] = host_timestamp
                textlog_row['message'] = message
                textlog_row.append()

            self.h5textlog.flush()

        self.h5_xhat_names = PT.Description(
            kalman_estimates_description().columns)._v_names
        self.h5_obs_names = PT.Description(
            FilteredObservations().columns)._v_names
        self.all_kalman_calibration_data = []

    def close(self):
        pass

    def save_tro(self,tro,force_obj_id=None):
        if len(tro.observations_frames) < self.min_observations_to_save:
            # not enough data to bother saving
            return

        self.obj_id += 1

        if force_obj_id is not None:
            save_obj_id = force_obj_id
        else:
            save_obj_id = self.obj_id

        if self.debug:
            print 'saving %s as obj_id %d'%(repr(self), save_obj_id)

        # save observation 2d data indexes
        debugADS=False

        if debugADS:
            print '2D indices: ----------------'

        this_idxs = []
        for camns_and_idxs in tro.observations_2d:
            this_idxs.append( self.h5_2d_obs_next_idx )
            self.h5_2d_obs.append( camns_and_idxs )

            if debugADS:
                print ' %d: %s'%(self.h5_2d_obs_next_idx,str(camns_and_idxs))
            self.h5_2d_obs_next_idx += 1
        self.h5_2d_obs.flush()

        if debugADS:
            print

        # becomes obs_2d_idx (index into 'ML_estimates_2d_idxs')
        this_idxs = numpy.array( this_idxs, dtype=numpy.uint64 )

        # save observations ####################################
        observations_frames = numpy.array(
            tro.observations_frames, dtype=numpy.uint64)
        obj_id_array = numpy.empty(
            observations_frames.shape, dtype=numpy.uint32)
        obj_id_array.fill(save_obj_id)
        MLE_position = numpy.array(
            tro.MLE_position, dtype=numpy.float32)
        MLE_Lcoords = numpy.array(
            tro.MLE_Lcoords, dtype=numpy.float32)
        list_of_pos = [MLE_position[:,i]
                       for i in range(MLE_position.shape[1])]
        list_of_lines = [MLE_Lcoords[:,i]
                         for i in range(MLE_Lcoords.shape[1])]
        array_list = ([obj_id_array,observations_frames]+
                      list_of_pos+[this_idxs]+list_of_lines)
        obs_recarray = numpy.rec.fromarrays(
            array_list, names = self.h5_obs_names)
        if 0:
            # End tracking at last non-nan observation (must be > 1
            # camera for final points).
            idx = numpy.nonzero(~numpy.isnan(MLE_position))[0][-1]
            last_observation_frame = observations_frames[idx]
        else:
            # End tracking at last observation (can be 1 camera for
            # final points).
            last_observation_frame = observations_frames[-1]

        if debugADS:
            print 'kalman observations: --------------'
            for row in obs_recarray:
                print row['frame'], row['obs_2d_idx']

        self.h5_obs.append(obs_recarray)
        self.h5_obs.flush()

        # save xhat info (kalman estimates) ##################

        frames = numpy.array(tro.frames, dtype=numpy.uint64)
        xhat_data = numpy.array(tro.xhats, dtype=numpy.float32)
        timestamps = numpy.array(tro.timestamps, dtype=numpy.float64)
        P_data_full = numpy.array(tro.Ps, dtype=numpy.float32)

        # don't guess after last observation
        cond = frames <= last_observation_frame
        frames = frames[cond]
        xhat_data = xhat_data[cond]
        timestamps = timestamps[cond]
        P_data_full = P_data_full[cond]

        obj_id_array = numpy.empty(frames.shape, dtype=numpy.uint32)
        obj_id_array.fill(save_obj_id)

        # one list entry per column
        list_of_xhats = [xhat_data[:,i]
                         for i in range(xhat_data.shape[1])]

        tmp = self.kalman_saver_info_instance.covar_mats_to_covar_entries
        list_of_Ps = tmp(P_data_full)
        del tmp

        xhats_recarray = numpy.rec.fromarrays(
            [obj_id_array,frames,timestamps]+list_of_xhats+list_of_Ps,
            names = self.h5_xhat_names)

        self.h5_xhat.append(xhats_recarray)
        self.h5_xhat.flush()

def kalmanize(src_filename,
              do_full_kalmanization=True,
              dest_filename=None,
              reconstructor=None,
              reconstructor_filename=None,
              start_frame=None,
              stop_frame=None,
              exclude_cam_ids=None,
              exclude_camns=None,
              dynamic_model_name=None,
              debug=False,
              frames_per_second=None,
              area_threshold=0,
              min_observations_to_save=0,
              options=None,
              ):
    if options is None:
        # get default options
        parser = get_parser()
        (options, args) = parser.parse_args([])

    if debug:
        numpy.set_printoptions(precision=3,linewidth=120,suppress=False)

    if exclude_cam_ids is None:
        exclude_cam_ids = []

    if exclude_camns is None:
        exclude_camns = []

    use_existing_filename = True

    if reconstructor is not None:
        assert isinstance(reconstructor, flydra_core.reconstruct.Reconstructor )
        assert reconstructor_filename is None

    with open_file_safe(src_filename,mode='r') as results:
        camn2cam_id, cam_id2camns = get_caminfo_dicts(results)

        if do_full_kalmanization:
            if dynamic_model_name is None:
                if hasattr(results.root,'kalman_estimates'):
                    if hasattr(results.root.kalman_estimates.attrs,'dynamic_model_name'):
                        dynamic_model_name = (
                            results.root.kalman_estimates.attrs.dynamic_model_name)
                        warnings.warn('dynamic model not specified. '
                                      'using "%s"'%dynamic_model_name)
            if dynamic_model_name is None:
                dynamic_model_name = 'EKF mamarama, units: mm'
                warnings.warn('dynamic model not specified. '
                              'using "%s"'%dynamic_model_name)
            else:
                print 'using dynamic model "%s"'%dynamic_model_name

            if reconstructor_filename is not None:
                if reconstructor_filename.endswith('h5'):
                    with PT.open_file(reconstructor_filename,mode='r') as fd:
                        reconstructor = flydra_core.reconstruct.Reconstructor(
                            fd,
                            minimum_eccentricity=options.force_minimum_eccentricity)
                else:
                    reconstructor = flydra_core.reconstruct.Reconstructor(
                        reconstructor_filename,
                        minimum_eccentricity=options.force_minimum_eccentricity)
            else:
                # reconstructor_filename is None
                if reconstructor is None:
                    reconstructor = flydra_core.reconstruct.Reconstructor(
                        results,
                        minimum_eccentricity=options.force_minimum_eccentricity)

            if options.force_minimum_eccentricity is not None:
                if (reconstructor.minimum_eccentricity !=
                    options.force_minimum_eccentricity):
                    raise ValueError('could not force minimum_eccentricity')

            if dest_filename is None:
                dest_filename = os.path.splitext(
                    results.filename)[0]+'.kalmanized.h5'
        else:
            use_existing_filename = False
            dest_filename = tempfile.mktemp(suffix='.h5')

        if reconstructor is not None and reconstructor.cal_source_type == 'pytables':
            save_reconstructor_filename = reconstructor.cal_source.filename
        else:
            warnings.warn('unable to determine reconstructor source '
                          'filename for %r'%reconstructor.cal_source_type)
            save_reconstructor_filename = None

        if frames_per_second is None:
            frames_per_second = get_fps(results)
            if do_full_kalmanization:
                print 'read frames_per_second from file', frames_per_second

        dt = 1.0/frames_per_second

        if options.sync_error_threshold_msec is None:
            # default is IFI/2
            sync_error_threshold = (0.5*dt)
        else:
            sync_error_threshold=options.sync_error_threshold_msec/1000.0

        if os.path.exists(dest_filename):
            if use_existing_filename:
                raise ValueError('%s already exists. Will not '
                                'overwrite.'%dest_filename)
            else:
                os.unlink(dest_filename)

        with open_file_safe(dest_filename, mode="w", title="tracked Flydra data file",
                          delete_on_error=True) as h5file:

            if 'experiment_info' in results.root:
                results.root.experiment_info._f_copy(h5file.root,recursive=True)

            if do_full_kalmanization:
                parsed = read_textlog_header(results)
                if 'trigger_CS3' not in parsed:
                    parsed['trigger_CS3'] = 'unknown'
                textlog_save_lines = [
                    'kalmanize running at %s fps, (top %s, trigger_CS3 %s, flydra_version %s)'%(
                    str(frames_per_second),str(parsed.get('top','unknown')),
                    str(parsed['trigger_CS3']),flydra_core.version.__version__),
                    'original file: %s'%(src_filename,),
                    'dynamic model: %s'%(dynamic_model_name,),
                    'reconstructor file: %s'%(save_reconstructor_filename,),
                    ]

                kalman_model = dynamic_models.get_kalman_model(
                    name=dynamic_model_name, dt=dt )

                h5saver = KalmanSaver(h5file,
                                      reconstructor,
                                      cam_id2camns=cam_id2camns,
                                      min_observations_to_save=min_observations_to_save,
                                      textlog_save_lines=textlog_save_lines,
                                      dynamic_model_name=dynamic_model_name,
                                      dynamic_model=kalman_model,
                                      debug=debug,
                                      fake_timestamp = options.fake_timestamp,
                                      )

                tracker = Tracker(
                    reconstructor,
                    kalman_model=kalman_model,
                    save_all_data=True,
                    area_threshold=area_threshold,
                    area_threshold_for_orientation=options.area_threshold_for_orientation,
                    disable_image_stat_gating=options.disable_image_stat_gating,
                    orientation_consensus=options.orientation_consensus,
                    fake_timestamp=options.fake_timestamp,
                    )

                tracker.set_killed_tracker_callback( h5saver.save_tro )

                # copy timestamp data into newly created kalmanized file
                if hasattr(results.root,'trigger_clock_info'):
                    results.root.trigger_clock_info._f_copy(h5file.root)

            data2d = results.root.data2d_distorted

            frame_count = 0
            last_frame = None
            frame_data = collections.defaultdict(list)
            time_frame_all_cam_timestamps = []
            time_frame_all_camns = []

            if 1:
                time1 = time.time()
                if do_full_kalmanization:
                    print 'loading all frame numbers...'
                frames_array = numpy.asarray(data2d.read(field='frame'))
                time2 = time.time()
                if do_full_kalmanization:
                    print 'done in %.1f sec'%(time2-time1)
                    if (not options.disable_image_stat_gating and
                        'cur_val' in data2d.colnames):
                        warnings.warn('No pre-filtering of data based on zero '
                                      'probability -- more data association work is '
                                      'being done than necessary')

            if len(frames_array)==0:
                # no data
                print 'No 2D data. Nothing to do.'
                return

            if do_full_kalmanization:
                print '2D data range: approximately %d<frame<%d'%(
                    frames_array[0], frames_array[-1])

            if do_full_kalmanization:
                accum_frame_spread = None
            else:
                accum_frame_spread = []
                accum_frame_spread_fno = []
                accum_frame_all_timestamps = []
                accum_frame_all_camns = []

            max_all_check_times = -np.inf

            for row_start, row_stop in utils.iter_non_overlapping_chunk_start_stops(
                frames_array, min_chunk_size=500000, size_increment=1000, status_fd=sys.stdout):

                print 'Doing initial scan of approx frame range %d-%d.'%(
                    frames_array[row_start],frames_array[row_stop-1])

                this_frames_array = frames_array[row_start:row_stop]
                if start_frame is not None:
                    if this_frames_array.max() < start_frame:
                        continue
                if stop_frame is not None:
                    if this_frames_array.min() > stop_frame:
                        continue

                data2d_recarray = data2d.read(start=row_start, stop=row_stop)
                this_frames = data2d_recarray['frame']
                print 'Examining frames %d-%d in detail.'%(this_frames[0],this_frames[-1])
                this_row_idxs = np.argsort( this_frames )
                for ii in range(len(this_row_idxs)+1):

                    if ii >= len(this_row_idxs):
                        finish_frame=True
                    else:
                        finish_frame = False

                        this_row_idx = this_row_idxs[ii]

                        row = data2d_recarray[this_row_idx]

                        new_frame = row['frame']

                        if start_frame is not None:
                            if new_frame < start_frame:
                                continue
                        if stop_frame is not None:
                            if new_frame > stop_frame:
                                continue

                        if last_frame != new_frame:
                            if new_frame < last_frame:
                                print 'new_frame',new_frame
                                print 'last_frame',last_frame
                                raise RuntimeError("expected continuously increasing "
                                                   "frame numbers")
                            finish_frame =True

                    if finish_frame:
                        # new frame
                        ########################################
                        # Data for this frame is complete
                        if last_frame is not None:

                            this_frame_spread = 0.0
                            if len(time_frame_all_cam_timestamps) > 1:
                                check_times = np.array(time_frame_all_cam_timestamps)
                                check_times -= check_times.min()
                                this_frame_spread = check_times.max()
                                if accum_frame_spread is not None:
                                    accum_frame_spread.append( this_frame_spread )
                                    accum_frame_spread_fno.append( last_frame )

                                    accum_frame_all_timestamps.append(
                                        time_frame_all_cam_timestamps )
                                    accum_frame_all_camns.append(
                                        time_frame_all_camns )

                                max_all_check_times = max( this_frame_spread,
                                                           max_all_check_times)
                                if this_frame_spread > sync_error_threshold:
                                    if this_frame_spread==max_all_check_times:
                                        print '%s frame %d: sync diff: %.1f msec'%(
                                            os.path.split(results.filename)[-1],
                                            last_frame,
                                            this_frame_spread*1000.0)

                            if debug > 5:
                                print
                                print 'frame_data for frame %d'%(last_frame,)
                                pprint.pprint(dict(frame_data))
                                print
                            if do_full_kalmanization:
                                if this_frame_spread > sync_error_threshold:
                                    if debug > 5:
                                        print (
                                            'frame sync error (spread %.1f msec), '
                                            'skipping'%(this_frame_spread*1e3,))
                                        print
                                    warnings.warn('Synchronization error detected, '
                                                  'but continuing analysis without '
                                                  'potentially bad data.')
                                else:
                                    process_frame(reconstructor,tracker,
                                                  last_frame,frame_data,camn2cam_id,
                                                  debug=debug)
                            frame_count += 1
                            if do_full_kalmanization and frame_count%1000==0:
                                time2 = time.time()
                                dur = time2-time1
                                fps = frame_count/dur
                                print 'frame % 10d, kalmanization/data association speed: % 8.1f fps'%(
                                    last_frame,fps)
                                time1 = time2
                                frame_count = 0

                        ########################################
                        frame_data = collections.defaultdict(list)
                        time_frame_all_cam_timestamps = [] # clear values
                        time_frame_all_camns = [] # clear values
                        last_frame = new_frame

                    camn = row['camn']
                    try:
                        cam_id = camn2cam_id[camn]
                    except KeyError:
                        # This will happen if cameras were re-synchronized (and
                        # thus gain new cam_ids) immediately before saving was
                        # turned on in MainBrain. The reason is that the network
                        # buffers are still full of old data coming in from the
                        # cameras.
                        warnings.warn('WARNING: no cam_id for camn '
                                      '%d, skipping this row of data'%camn)
                        continue

                    if cam_id in exclude_cam_ids:
                        # exclude this camera
                        continue

                    if camn in exclude_camns:
                        # exclude this camera
                        continue

                    time_frame_all_cam_timestamps.append( row['timestamp'] )
                    time_frame_all_camns.append( row['camn'] )

                    if do_full_kalmanization:

                        x_distorted = row['x']
                        if numpy.isnan(x_distorted):
                            # drop point -- not found
                            continue
                        y_distorted = row['y']

                        (x_undistorted,y_undistorted) = reconstructor.undistort(
                            cam_id,(x_distorted,y_distorted))

                        (area,slope,eccentricity,frame_pt_idx) = (row['area'],
                                                                  row['slope'],
                                                                  row['eccentricity'],
                                                                  row['frame_pt_idx'])

                        if 'cur_val' in row.dtype.fields: cur_val = row['cur_val']
                        else: cur_val = None
                        if 'mean_val' in row.dtype.fields: mean_val = row['mean_val']
                        else: mean_val = None
                        if 'sumsqf_val' in row.dtype.fields: sumsqf_val = row['sumsqf_val']
                        else: sumsqf_val = None

                        # FIXME: cache this stuff?
                        pmat_inv = reconstructor.get_pmat_inv(cam_id)
                        camera_center = reconstructor.get_camera_center(cam_id)
                        camera_center = numpy.hstack((camera_center[:,0],[1]))
                        camera_center_meters = reconstructor.get_camera_center(
                            cam_id)
                        camera_center_meters = numpy.hstack((camera_center_meters[:,0],[1]))
                        helper = reconstructor.get_reconstruct_helper_dict()[cam_id]
                        rise=slope
                        run=1.0
                        if np.isinf(rise):
                            if rise > 0:
                                rise=1.0
                                run=0.0
                            else:
                                rise=-1.0
                                run=0.0

                        (p1, p2, p3, p4,
                         ray0, ray1, ray2, ray3, ray4, ray5) = do_3d_operations_on_2d_point(
                            helper, x_undistorted,y_undistorted,
                            pmat_inv,
                            camera_center,
                            x_distorted,y_distorted,
                            rise, run)
                        line_found = not numpy.isnan(p1)
                        pluecker_hz_meters = (ray0, ray1, ray2, ray3, ray4, ray5)

                        # Keep in sync with kalmanize.py and data_descriptions.py
                        pt_undistorted = (x_undistorted,y_undistorted,
                                          area,slope,eccentricity,
                                          p1,p2,p3,p4, line_found,
                                          frame_pt_idx, cur_val, mean_val, sumsqf_val)

                        projected_line_meters=geom.line_from_HZline(pluecker_hz_meters)

                        frame_data[camn].append((pt_undistorted,projected_line_meters))

            if do_full_kalmanization:
                tracker.kill_all_trackers() # done tracking

        if not do_full_kalmanization:
            os.unlink(dest_filename)

    if accum_frame_spread is not None:
        # save spread data to file for analysis
        accum_frame_spread = np.array(accum_frame_spread)
        accum_frame_spread_fno = np.array(accum_frame_spread_fno)
        if options.dest_file is not None:
            accum_frame_spread_filename = options.dest_file
        else:
            accum_frame_spread_filename = src_filename + '.spreadh5'

        cam_ids = cam_id2camns.keys()
        cam_ids.sort()
        camn_order = []
        for cam_id in cam_ids:
            camn_order.extend( cam_id2camns[cam_id] )

        camn_order = np.array(camn_order)
        cam_id_array = np.array(cam_ids)

        N_cams = len(camn_order)
        N_frames = len(accum_frame_spread_fno)

        all_timestamps=np.empty( ( N_frames, N_cams ), dtype=np.float)
        all_timestamps.fill(np.nan)
        for i,(timestamps,camns) in enumerate(zip(
            accum_frame_all_timestamps,
            accum_frame_all_camns)):

            for j,camn in enumerate(camn_order):
                try:
                    idx = camns.index(camn)
                except ValueError:
                    continue # not found, skip
                timestamp = timestamps[idx]
                all_timestamps[i,j]=timestamp

        h5 = tables.open_file( accum_frame_spread_filename, mode='w')
        h5.create_array( h5.root, 'spread', accum_frame_spread,
                        'frame timestamp spreads (sec)')
        h5.create_array( h5.root, 'framenumber', accum_frame_spread_fno,
                        'frame number')
        h5.create_array( h5.root, 'all_timestamps', all_timestamps,
                        'all timestamps')
        h5.create_array( h5.root, 'camn_order', camn_order,
                        'camn_order')
        h5.create_array( h5.root, 'cam_id_array', cam_id_array,
                        'cam_id_array')
        h5.close()
        print 'saved %s'%accum_frame_spread_filename

    if max_all_check_times > sync_error_threshold:
        if not options.keep_sync_errors:
            if do_full_kalmanization:
                print 'max_all_check_times %.2f msec'%(
                    max_all_check_times*1000.0)
                handle, target = tempfile.mkstemp(
                    os.path.split(dest_filename)[1])
                os.unlink(target) # remove original file there
                shutil.move( dest_filename, target)

                raise ValueError(
                    'Synchonization errors exist in the data. Moved result file'
                    ' to ensure it is not confused with valid data. The new '
                    'location is: %s'%(target,))

            else:
                sys.exit(1) # sync error
    else:
        if not do_full_kalmanization:
            print '%s no sync differences greater than %.1f msec'%(
                os.path.split(src_filename)[-1],
                sync_error_threshold*1000.0,
                )

def check_sync():
    usage = """%prog FILE [options]

This command will exit with a non-zero exit code if there are sync errors.
"""

    parser = OptionParser(usage)
    parser.add_option("--start", type="int",
                      help="first frame",
                      metavar="START")

    parser.add_option("--stop", type="int",
                      help="last frame",
                      metavar="STOP")

    parser.add_option("--dest-file", type="string",
                      help=("filename of .spreadh5 file (otherwise defaults "
                            "to FILE.spreadh5)"))

    parser.add_option("--sync-error-threshold-msec", type="float", default=None)

    parser.add_option(
        "--exclude-cam-ids", type='string',
        help="camera ids to exclude from reconstruction (space separated)",
        metavar="EXCLUDE_CAM_IDS")

    parser.add_option(
        "--exclude-camns", type='string',
        help="camera numbers to exclude from reconstruction (space separated)",
        metavar="EXCLUDE_CAMNS")

    parser.add_option("--debug", type="int",
                      metavar="DEBUG")

    (options, args) = parser.parse_args()
    if options.exclude_cam_ids is not None:
        options.exclude_cam_ids = options.exclude_cam_ids.split()

    if options.exclude_camns is not None:
        options.exclude_camns = [int(camn)
                                 for camn in options.exclude_camns.split()]

    src_filename = args[0]

    kalmanize(src_filename,
              do_full_kalmanization=False, # only check for sync errors
              start_frame=options.start,
              stop_frame=options.stop,
              exclude_cam_ids=options.exclude_cam_ids,
              exclude_camns=options.exclude_camns,
              debug = options.debug,
              options=options,
              )

def get_parser():
    usage = '%prog FILE [options]'

    parser = OptionParser(usage)

    parser.add_option("-d", "--dest-file", dest="dest_filename", type='string',
                      help="save to hdf5 file (append if already present)",
                      metavar="DESTFILE")

    parser.add_option(
        "-r", "--reconstructor", dest="reconstructor_path",
        type='string',
        help=("calibration/reconstructor path (if not specified, "
              "defaults to FILE)"),
        metavar="RECONSTRUCTOR")

    parser.add_option("--sync-error-threshold-msec", type="float", default=None)

    parser.add_option(
        "--save-cal-dir", type='string',
        help="directory name in which to save new calibration data",
        default=None,
        )

    parser.add_option(
        "--fake-timestamp", type='float',
        help="value of timestamp to use when saving logfile",
        default=None,
        )

    parser.add_option("--fps", dest='fps', type='float',
                      help="frames per second (used for Kalman filtering)")

    parser.add_option(
        "--exclude-cam-ids", type='string',
        help="camera ids to exclude from reconstruction (space separated)",
        metavar="EXCLUDE_CAM_IDS")

    parser.add_option(
        "--exclude-camns", type='string',
        help="camera numbers to exclude from reconstruction (space separated)",
        metavar="EXCLUDE_CAMNS")

    parser.add_option("--dynamic-model", dest="dynamic_model", type='string')

    parser.add_option("--start", type="int",
                      help="first frame",
                      metavar="START")

    parser.add_option("--stop", type="int",
                      help="last frame",
                      metavar="STOP")

    parser.add_option("--debug", type="int",
                      metavar="DEBUG")

    parser.add_option('--keep-sync-errors', action='store_true',
                      default=False, help='keep files with sync errors')

    parser.add_option("--area-threshold", type='float',
                      default=0.0,
                      help="area threshold (used to filter incoming 2d points)")

    parser.add_option("--area-threshold-for-orientation", type='float',
                      default=0.0,
                      help="minimum area to compute orientation")

    parser.add_option(
        "--min-observations-to-save", type='int', default=2,
        help=('minimum number of observations required for a kalman object '
              'to be saved'))

    parser.add_option("--orientation-consensus", type='int', default=0)

    parser.add_option("--force-minimum-eccentricity", type='float',
                      default=None)

    parser.add_option("--disable-image-stat-gating", action='store_true',
                      help="disable gating the data based on image statistics",
                      default=False)
    return parser

def main():
    parser = get_parser()
    (options, args) = parser.parse_args()
    if options.exclude_cam_ids is not None:
        options.exclude_cam_ids = options.exclude_cam_ids.split()

    if options.exclude_camns is not None:
        options.exclude_camns = [int(camn)
                                 for camn in options.exclude_camns.split()]

    if len(args)>1:
        print 'args',args
        print >> sys.stderr, ("arguments interpreted as FILE supplied more "
                              "than once")
        parser.print_help()
        sys.exit(1)

    if len(args)<1:
        parser.print_help()
        sys.exit(1)

    src_filename = args[0]

    args = (src_filename,)
    kwargs = dict(
              dest_filename=options.dest_filename,
              reconstructor_filename=options.reconstructor_path,
              start_frame=options.start,
              stop_frame=options.stop,
              exclude_cam_ids=options.exclude_cam_ids,
              exclude_camns=options.exclude_camns,
              dynamic_model_name = options.dynamic_model,
              debug = options.debug,
              frames_per_second = options.fps,
              area_threshold=options.area_threshold,
              min_observations_to_save=options.min_observations_to_save,
              options=options,
              )

    if int(os.environ.get('PROFILE','0')):
        import cProfile
        import lsprofcalltree
        p = cProfile.Profile()
        print 'running kalmanize in profile mode'
        p.runctx('kalmanize(*args, **kwargs)',globals(),locals())
        k = lsprofcalltree.KCacheGrind(p)
        data = open(os.path.expanduser('~/kalmanize.kgrind'), 'w+')
        k.output(data)
        data.close()
    else:
        kalmanize(*args, **kwargs)

if __name__=='__main__':
    main()


"""
Generate speaker embeddings and metadata for training
"""
import os
import pickle
import regex
import hashlib

from model_bl import D_VECTOR
from collections import OrderedDict
import numpy as np
from numpy.random import RandomState
import torch

import argparse

from utils import *

python_object_filter = ['npy', 'pkl']
nat_sort = lambda s: [int(t) if t.isdigit() else t.lower() for t in regex.split(r'(\d+)', s)]

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Make metadata.')
    # parser.add_argument('-r', '--root-dir', default='./spmel/wavs', help='root directory of audio dataset', type=str)
    parser.add_argument('-r', '--root-dir', default='./spmel/test-clean', help='root directory of audio dataset', type=str)
    parser.add_argument('-l', '--tag-level', default=0, help='tag level', type=int)
    parser.add_argument('-u', '--num-uttrs', default=999, help='number of utterances', type=int)
    parser.add_argument('-c', '--len-crop', default=128, help='crop length', type=int)

    args = parser.parse_args()

    C = D_VECTOR(dim_input=80, dim_cell=768, dim_emb=256).eval().cuda()
    c_checkpoint = torch.load('pretrained/3000000-BL.ckpt')
    new_state_dict = OrderedDict()
    for key, val in c_checkpoint['model_b'].items():
        new_key = key[7:]
        new_state_dict[new_key] = val
    C.load_state_dict(new_state_dict)
    num_uttrs = args.num_uttrs
    len_crop = args.len_crop

    # Directory containing mel-spectrograms
    root_dir = args.root_dir

    tag_level = args.tag_level
    if tag_level == 0:
        _, speaker_id, _ = next(os.walk(root_dir))

        speaker_id = sorted(speaker_id, key=nat_sort)
        # print(speaker_id)
        speakers = []
        for speaker in speaker_id:
            # print('Processing speaker: %s' % speaker)
            utterances = []
            utterances.append(speaker)

            speaker_root = os.path.join(root_dir, speaker)
            all_file_list = []
            for dir_name, subdir_list, file_list in os.walk(speaker_root):
                for file_name in file_list:
                    all_file_list.append(os.path.join(dir_name, file_name))

            rngsd = int(hashlib.sha256(speaker_root.encode('utf-8')).hexdigest() , 16) % 10**8
            prng = RandomState(rngsd)

            real_num_uttrs = num_uttrs
            if len(all_file_list) < real_num_uttrs:
                real_num_uttrs = len(all_file_list)

            # new
            idx_uttrs = prng.choice(len(all_file_list), size=real_num_uttrs, replace=False)
            embs = []
            for i in range(real_num_uttrs):
                target_file = all_file_list[idx_uttrs[i]]
                tmp = np.load(target_file)
                # print(target_file, tmp.shape)
                if tmp.shape[0] > 0 and tmp.shape[0] < len_crop:
                    repeat_num = np.ceil(len_crop/tmp.shape[0]).astype(int)
                    res_mel = np.concatenate([tmp for i in range(repeat_num)])
                else:
                    res_mel = tmp
                if res_mel.shape[0] == len_crop:
                    left = 0
                else:
                    left = prng.randint(0, res_mel.shape[0]-len_crop)
                melsp = torch.from_numpy(res_mel[np.newaxis, left:left+len_crop, :]).cuda()
                emb = C(melsp)
                embs.append(emb.detach().squeeze().cpu().numpy())

            if len(embs) > 0:
                utterances.append(np.mean(embs, axis=0))
                for file_name in sorted(all_file_list, key=nat_sort):
                    utterances.append(os.path.relpath(file_name, root_dir))
                speakers.append(utterances)

            print('Make metadata completed. ID: {}. Sampled embeddings number: {}.'.format(speaker, len(embs)))

            # old
            '''
            idx_uttrs = prng.choice(len(all_file_list), size=real_num_uttrs, replace=False)

            # print(idx_uttrs, end=' ')
            embs = []
            for i in range(real_num_uttrs):
                target_file = all_file_list[idx_uttrs[i]]
                tmp = np.load(target_file)
                # print(tmp.shape)
                candidates = np.delete(np.arange(len(all_file_list)), idx_uttrs)

                insufficient = False

                # search for long enough spect
                while tmp.shape[0] <= len_crop:
                    if len(candidates) == 0:
                        insufficient = True
                        break
                    idx_alt = prng.choice(candidates)
                    target_file_alt = all_file_list[idx_alt]
                    tmp = np.load(target_file_alt)
                    candidates = np.delete(candidates, np.argwhere(candidates==idx_alt))
                if insufficient:
                    break
                left = prng.randint(0, tmp.shape[0]-len_crop)
                melsp = torch.from_numpy(tmp[np.newaxis, left:left+len_crop, :]).cuda()
                emb = C(melsp)
                embs.append(emb.detach().squeeze().cpu().numpy())

            if len(embs) > 0:
                utterances.append(np.mean(embs, axis=0))

                for file_name in sorted(all_file_list, key=nat_sort):
                    utterances.append(os.path.relpath(file_name, root_dir))
                speakers.append(utterances)
            '''

        with open(os.path.join(root_dir, 'train.pkl'), 'wb') as write_pkl:
            pickle.dump(speakers, write_pkl)
    elif tag_level == 1:
        pass
    else:
        pass


import os
import numpy as np
import tensorflow as tf
import time
import cv2
import datetime

from graph_utils import load_graph
from PIL import ImageDraw, Image
from styx_msgs.msg import TrafficLight

# to save monitor image for debug. if set True, uncomment "import visualization_utils as vis_util" please.
# if "import visualization_utils" occurs error, for detail:https://github.com/MarkBroerkens/CarND-Capstone/blob/master/README.md#prepare-environment
SAVE_MONITOR_IMAGE = False # to save monitor image for debug.
#import visualization_utils as vis_util

# to show labeled image, for test only, if submit code to udacity, should be set SHOW_MONITOR_IMAGE = False and comment below "imort" lines.
# take a show: open a new terminal ,and run "rosrun image_view image_view image:=/clssifier_monitor_image"
SHOW_MONITOR_IMAGE = False
#import visualization_utils as vis_util
#import rospy
#from sensor_msgs.msg import Image as Image_msg
#from cv_bridge import CvBridge

class TLClassifier(object):
    def __init__(self, is_site):
        #TODO load classifier

        sess = None

        if is_site:
            sess, _ = load_graph('models/real_model.pb')
        else:
            sess, _ = load_graph('models/sim_model.pb')

        self.sess = sess
        self.sess_graph = self.sess.graph
        # Definite input and output Tensors for sess
        self.image_tensor = self.sess_graph.get_tensor_by_name('image_tensor:0')
        # Each box represents a part of the image where a particular object was detected.
        self.detection_boxes = self.sess_graph.get_tensor_by_name('detection_boxes:0')
        # Each score represent how level of confidence for each of the objects.
        # Score is shown on the result image, together with the class label.
        self.detection_scores = self.sess_graph.get_tensor_by_name('detection_scores:0')
        self.detection_classes = self.sess_graph.get_tensor_by_name('detection_classes:0')
        self.num_classes = 3

        self.category_index = {1: {'id': 1, 'name': u'red'}, 2: {'id': 2, 'name': u'yellow'}, 3: {'id': 3, 'name': u'green'}}

        self.image_count = 0
        self.last_pred = TrafficLight.UNKNOWN
        self.pub_tl_clssifier_monitor = None
        self.bridge = None

        if SHOW_MONITOR_IMAGE:
            self.pub_tl_clssifier_monitor = rospy.Publisher('/clssifier_monitor_image', Image_msg, queue_size=2)
            self.bridge = CvBridge()

    def get_classification(self, image, wp = 0):
        """Determines the color of the traffic light in the image

        Args:
            image (cv::Mat): image containing the traffic light

        Returns:
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)

        """

        cv2_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image_np_expanded = np.expand_dims(cv2_image, axis=0)

        (boxes, scores, classes) = self.sess.run([self.detection_boxes, self.detection_scores, self.detection_classes], feed_dict={self.image_tensor: image_np_expanded})

        prediction = 4
        min_score_thresh=.6
        sq_boxes = np.squeeze(boxes)
        sq_classes = np.squeeze(classes).astype(np.int32)
        sq_scores = np.squeeze(scores)

        for i in range(sq_boxes.shape[0]):
            if sq_scores is None or sq_scores[i] > min_score_thresh:
                if sq_classes[i] in self.category_index.keys():
                    prediction = sq_classes[i]
                    print("Found traffic light: {ID:%s  color:%s  pred_score:%.4f}"%(prediction, str(self.category_index[sq_classes[i]]['name']), sq_scores[i]))
                    min_score_thresh = sq_scores[i]

        if SAVE_MONITOR_IMAGE:

            dt_str = str(time.time()).replace('.','')
            file_name = "Monitor_IMG_" + dt_str + '_'+ str(int(prediction - 1)) +'.jpg'
            output_filename = "light_classification/IMGS/tl_origin/" + file_name
            cv2.imwrite(output_filename, image)

            vis_image = vis_util.visualize_boxes_and_labels_on_image_array(
                     image,
                     sq_boxes,
                     sq_classes,
                     sq_scores,
                     self.category_index,
                     use_normalized_coordinates=True,
                     line_thickness=1
                     #min_score_thresh = min_score_thresh
                     )
            image_pil = Image.fromarray(cv2.cvtColor(vis_image,cv2.COLOR_BGR2RGB))
            file_name = "Monitor_IMG_" + dt_str + '_'+ str(int(prediction - 1)) +'.png'
            output_filename = "light_classification/IMGS/tl_predict_output/" + file_name
            image_pil.save(output_filename, 'PNG')

        if SHOW_MONITOR_IMAGE:
            vis_image = vis_util.visualize_boxes_and_labels_on_image_array(
                     image,
                     sq_boxes,
                     sq_classes,
                     sq_scores,
                     self.category_index,
                     use_normalized_coordinates=True,
                     line_thickness=1
                     )

            image_message = self.bridge.cv2_to_imgmsg(cv2.cvtColor(vis_image,cv2.COLOR_BGR2RGB), encoding="rgb8")
            self.pub_tl_clssifier_monitor.publish(image_message)

        rtn = TrafficLight.UNKNOWN

        if prediction == 1:
            rtn = TrafficLight.RED
        elif prediction == 2:
            rtn = TrafficLight.YELLOW
        elif prediction == 3:
            rtn = TrafficLight.GREEN

        self.last_pred = rtn
        self.image_count += 1

        return rtn


import gym
import numpy as np
import torch
from ray.rllib import SampleBatch
from ray.rllib.agents.impala import ImpalaTrainer
from ray.rllib.agents.impala.vtrace_torch_policy import VTraceTorchPolicy, VTraceLoss, make_time_major
from ray.rllib.models.torch.torch_action_dist import TorchCategorical
from ray.rllib.policy.torch_policy import LearningRateSchedule, EntropyCoeffSchedule
from ray.rllib.utils.torch_ops import sequence_mask

from griddly.util.rllib.torch.conditional_actions.conditional_action_mixin import ConditionalActionMixin

def build_invalid_masking_vtrace_loss(policy, model, dist_class, train_batch):

    model_out, _ = model.from_batch(train_batch)

    if isinstance(policy.action_space, gym.spaces.Discrete):
        is_multidiscrete = False
        output_hidden_shape = [policy.action_space.n]
    elif isinstance(policy.action_space, gym.spaces.MultiDiscrete):
        is_multidiscrete = True
        output_hidden_shape = policy.action_space.nvec.astype(np.int32)
    else:
        is_multidiscrete = False
        output_hidden_shape = 1

    def _make_time_major(*args, **kw):
        return make_time_major(policy, train_batch.get("seq_lens"), *args,
                               **kw)

    actions = train_batch[SampleBatch.ACTIONS]
    dones = train_batch[SampleBatch.DONES]
    rewards = train_batch[SampleBatch.REWARDS]
    behaviour_action_logp = train_batch[SampleBatch.ACTION_LOGP]
    behaviour_logits = train_batch[SampleBatch.ACTION_DIST_INPUTS]

    invalid_action_mask = train_batch['invalid_action_mask']

    if 'seq_lens' in train_batch:
        max_seq_len = policy.config['rollout_fragment_length']
        mask_orig = sequence_mask(train_batch["seq_lens"], max_seq_len)
        mask = torch.reshape(mask_orig, [-1])
    else:
        mask = torch.ones_like(rewards)

    model_out += torch.maximum(torch.tensor(torch.finfo().min), torch.log(invalid_action_mask))
    action_dist = dist_class(model_out, model)

    if isinstance(output_hidden_shape, (list, tuple, np.ndarray)):
        unpacked_behaviour_logits = torch.split(
            behaviour_logits, list(output_hidden_shape), dim=1)
        unpacked_outputs = torch.split(
            model_out, list(output_hidden_shape), dim=1)
    else:
        unpacked_behaviour_logits = torch.chunk(
            behaviour_logits, output_hidden_shape, dim=1)
        unpacked_outputs = torch.chunk(model_out, output_hidden_shape, dim=1)
    values = model.value_function()

    # Prepare actions for loss.
    loss_actions = actions if is_multidiscrete else torch.unsqueeze(
        actions, dim=1)

    # Inputs are reshaped from [B * T] => [T - 1, B] for V-trace calc.
    policy.loss = VTraceLoss(
        actions=_make_time_major(loss_actions, drop_last=True),
        actions_logp=_make_time_major(
            action_dist.logp(actions), drop_last=True),
        actions_entropy=_make_time_major(
            action_dist.entropy(), drop_last=True),
        dones=_make_time_major(dones, drop_last=True),
        behaviour_action_logp=_make_time_major(
            behaviour_action_logp, drop_last=True),
        behaviour_logits=_make_time_major(
            unpacked_behaviour_logits, drop_last=True),
        target_logits=_make_time_major(unpacked_outputs, drop_last=True),
        discount=policy.config["gamma"],
        rewards=_make_time_major(rewards, drop_last=True),
        values=_make_time_major(values, drop_last=True),
        bootstrap_value=_make_time_major(values)[-1],
        dist_class=TorchCategorical if is_multidiscrete else dist_class,
        model=model,
        valid_mask=_make_time_major(mask, drop_last=True),
        config=policy.config,
        vf_loss_coeff=policy.config["vf_loss_coeff"],
        entropy_coeff=policy.entropy_coeff,
        clip_rho_threshold=policy.config["vtrace_clip_rho_threshold"],
        clip_pg_rho_threshold=policy.config["vtrace_clip_pg_rho_threshold"])

    return policy.loss.total_loss

def setup_mixins(policy, obs_space, action_space, config):
    ConditionalActionMixin.__init__(policy)
    EntropyCoeffSchedule.__init__(policy, config["entropy_coeff"],
                                  config["entropy_coeff_schedule"])
    LearningRateSchedule.__init__(policy, config["lr"], config["lr_schedule"])

ConditionalActionVTraceTorchPolicy = VTraceTorchPolicy.with_updates(
    name="ConditionalActionVTraceTorchPolicy",
    loss_fn=build_invalid_masking_vtrace_loss,
    before_init=setup_mixins,
    mixins=[LearningRateSchedule, EntropyCoeffSchedule, ConditionalActionMixin]
)

def get_vtrace_policy_class(config):
    if config['framework'] == 'torch':
        return ConditionalActionVTraceTorchPolicy
    else:
        raise NotImplementedError('Tensorflow not supported')

ConditionalActionImpalaTrainer = ImpalaTrainer.with_updates(name="ConditionalActionImpalaTrainer",
                                                            default_policy=ConditionalActionVTraceTorchPolicy,
                                                            get_policy_class=get_vtrace_policy_class)