"""
Genotype a set of SV calls in a given set of BAMs.
"""
import argparse
from collections import defaultdict
import csv
import intervaltree
import logging
import numpy as np
import operator
import pandas as pd
import pybedtools
import pysam
from sklearn import mixture
np.random.seed(1)

logging.basicConfig(filename="genotyper.log",level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Input file looks like this with SV coordinates in local assembly at columns 6-8.
# chr1    350793  350794  insertion       40      chr1-350756-362784|utg7180000000000|merged      37      77
CHROMOSOME=0
START=1
END=2
EVENT_TYPE=3
EVENT_LENGTH=4
CONTIG_NAME=5
CONTIG_START=6
CONTIG_END=7

BAM_CMATCH = 0
BAM_CSOFT_CLIP = 4
DEVIATIONS_FOR_THRESHOLD = 1.5
# Allow at most 2 mismatches in a 100 bp read.
MAX_ERROR_RATE = 0.02
MAX_LIKELIHOOD = 100

def has_perfect_mapping(read):
    """
    Returns True if the given pysam read maps perfectly.

    A "perfect" mapping is either a mapping quality greater than zero and no
    more than 2% mismatches or a full-length alignment of the read without any
    mismatches. Also consider soft-clipping and other non-match alignment
    features in the CIGAR string by summing all non-match bases and comparing
    against max errors allowed.
    """
    max_errors = np.ceil(MAX_ERROR_RATE * read.qlen)
    return (
        not read.is_unmapped and
        ((read.mapq > 0 and dict(read.tags)["NM"] <= max_errors and sum([i[1] for i in read.cigar if i[0] == BAM_CSOFT_CLIP]) <= max_errors) or
         (len(read.cigar) == 1 and read.cigar[0][0] == BAM_CMATCH and read.cigar[0][1] == read.qlen and dict(read.tags)["NM"] == 0))
    )

def spans_region(read, region):
    """
    Returns True if the given pysam read spans the given pybedtools.Interval,
    ``region``.
    """
    return has_perfect_mapping(read) and read.reference_start <= region.start and read.reference_end >= region.end

def has_gaps_in_region(read, region):
    """
    Returns True if the given pysam read spans the given pybedtools.Interval,
    ``region``.
    """
    # If the given read has gaps in its alignment to the reference inside the
    # given interval (more than one block inside the SV event itself), there are
    # gaps inside the SV.
    tree = intervaltree.IntervalTree()
    for block in read.get_blocks():
        tree[block[0]:block[1]] = block

    return len(tree[region.start:region.end]) > 1

def pair_spans_regions(read_pair, regions):
    """
    Returns True if the given pysam reads spans the given pybedtools.Interval
    list, ``regions``. Read pairs can span the regions even if either read
    overlaps the region breakpoints.
    """
    return len(read_pair) == 2 and read_pair[0].reference_start < regions[0].start and read_pair[1].reference_end > regions[-1].end

def pair_has_proper_orientation(read_pair):
    """
    Returns True if the given pysam reads map towards each other (i.e., read
    with lowest coordinate is in forward orientation and other read is in
    reverse).
    """
    return len(read_pair) == 2 and not read_pair[0].is_reverse and read_pair[1].is_reverse

def soft_clips_at_breakpoint(read, region):
    """
    Returns True if the given pysam read maps up to the edge of the given
    pybedtools.Interval, ``region`` and soft clips at that edge.

    Two cases are possible:

    1. The end of the read soft clips such that the read's reference end stops
    at the breakpoint start and the end of the read is softclipped.

    2. The beginning of the read soft clips such that the read's reference start
    begins at the breakpoint end and the beginning of the read is softclipped.

    In both cases, the read should be mapped in its best location in the
    reference ("perfect mapping").
    """
    return (
        (read.reference_end == region.start and read.cigar[-1][0] == BAM_CSOFT_CLIP) or
        (read.reference_start == region.end + 1 and read.cigar[0][0] == BAM_CSOFT_CLIP)
    )

def maps_outside_regions(read, regions):
    """
    Returns True if the given pysam read maps outside the range of the given
    regions.
    """
    return (
        (read.reference_start < regions[0].start and read.reference_end < regions[0].start) or
        (read.reference_start > regions[-1].end and read.reference_end > regions[-1].end)
    )

def is_proper_pair(read_pair, lower_insert_size_threshold, upper_insert_size_threshold):
    """
    Returns True if both reads in the given pair map perfectly, map within the
    expected insert size, and map with the proper orientation.
    """
    return (
        lower_insert_size_threshold <= np.abs(read_pair[0].isize) <= upper_insert_size_threshold and
        pair_has_proper_orientation(read_pair) and
        any(map(has_perfect_mapping, read_pair))
    )

def get_insert_sizes_for_region(bam, region):
    """
    Return insert sizes reads in the given pysam.AlignmentFile, ``bam``, and the
    given pybedtools.Interval, ``region``.
    """
    return [read.tlen for read in bam.fetch(region.chrom, region.start, region.end)
            if has_perfect_mapping(read) and not read.mate_is_unmapped and read.tlen >= 0 and read.tlen <= 1000]

def get_depth_for_region(bam_fields, regions, breakpoints, region_type="control"):
    """
    Return read depth for concordant reads in the given dictionary
    ``bam_fields`` with a pysam.AlignmentFile in the field "bam" and the given
    pybedtools.BedTool of regions, ``regions``.

    The depth calculated across the given regions depends on the type of region
    which is "control" by default. Other acceptable values for the region type
    are "insertion" or "deletion". If the region type is not "control" then the
    coordinates given by ``breakpoints`` are used to evaluate concordant and
    discordant read pairs.
    """
    bam = bam_fields["file"]
    lower_insert_threshold = bam_fields["lower_insert_threshold"]
    upper_insert_threshold = bam_fields["upper_insert_threshold"]

    # Convert breakpoints to a list for easier indexing.
    breakpoints = list(breakpoints)

    # Get all distinct reads in given regions.
    reads = set()
    for region in regions:
        for read in bam.fetch(region.chrom, region.start, region.end):
            # Exclude secondary and supplementary alignments (flags 0x100 and 0x800).
            if not read.is_secondary and (read.flag & 0x800 == 0):
                reads.add(read)

    # Group reads by read name.
    reads_by_name = defaultdict(list)
    for read in reads:
        reads_by_name[read.qname].append(read)

    logger.debug("Found %i potential read pairs for %s", len(reads_by_name), str(region).strip())

    concordant_pairs = []

    for read_name, read_pair in reads_by_name.iteritems():
        # Sort reads by position to get reads "1" and "2" in expected order.
        read_pair = sorted(read_pair, key=operator.attrgetter("pos"))

        is_concordant = False
        if region_type == "insertion":
            # Reads are in a proper pair and either of the reads maps across
            # either of the breakpoints.
            if is_proper_pair(read_pair, lower_insert_threshold, upper_insert_threshold) and any([spans_region(read, breakpoint) for read in read_pair for breakpoint in breakpoints]):
                is_concordant = True
        elif region_type == "deletion":
            # Reads are in a proper pair and either of the reads maps across the
            # deletion breakpoint.
            if is_proper_pair(read_pair, lower_insert_threshold, upper_insert_threshold) and any([spans_region(read, breakpoint) for read in read_pair for breakpoint in breakpoints]):
                is_concordant = True
        else:
            # Control regions look for all properly paired reads with proper
            # orientation.
            if is_proper_pair(read_pair, lower_insert_threshold, upper_insert_threshold):
                is_concordant = True

        if is_concordant:
            concordant_pairs.append(read_pair)

    return concordant_pairs

def genotype_call_with_read_pair(concordant, discordant, std_depth):
    """
    Genotype call based on total depth of concordant and discordant reads a la
    Hormozdiari et al. 2010 (Genome Research).
    """
    #homozygous_deletion_threshold = std_depth
    homozygous_deletion_threshold = 5

    if concordant < homozygous_deletion_threshold and discordant < homozygous_deletion_threshold:
        genotype = "./."
        discordant_genotype_likelihood = np.power(2, discordant) / np.power(2, homozygous_deletion_threshold)
        concordant_genotype_likelihood = np.power(2, concordant) / np.power(2, homozygous_deletion_threshold)
        shared_likelihood = np.sqrt(np.square(concordant_genotype_likelihood) + np.square(discordant_genotype_likelihood))
        genotype_likelihood = np.floor(-10 * np.log10(shared_likelihood))
    else:
        expected_discordant_lower_bound = concordant * 0.25
        expected_discordant_upper_bound = concordant * 4

        if discordant < expected_discordant_lower_bound:
            genotype = "1/1"
            # Calculate likelihood of homozygous alternate genotype as
            # Phred-scaled proportion of distance between the observed
            # discordant depth and expected depth for the corresponding
            # concordant depth.
            genotype_likelihood = np.floor(-10 * np.log10(np.power(2, discordant) / np.power(2, expected_discordant_lower_bound)))
        elif expected_discordant_lower_bound <= discordant < expected_discordant_upper_bound:
            # Calculate likelihood of heterozygous genotype as Phred-scaled
            # proportion of distance between the observed discordant depth and
            # expected depth for the corresponding concordant depth.
            genotype_ratio = np.power(2, np.abs(discordant - expected_discordant_upper_bound)) / np.power(2, expected_discordant_upper_bound)
            genotype_likelihood = np.floor(-10 * np.log10(1 - min(genotype_ratio, 1)))
            genotype = "1/0"
        else:
            # Calculate likelihood of homozygous "reference" genotype as
            # Phred-scaled proportion of distance between the observed
            # discordant depth and expected depth for the corresponding
            # concordant depth.
            genotype_ratio = np.power(2, np.abs(discordant - expected_discordant_upper_bound)) / np.power(2, expected_discordant_upper_bound)
            genotype_likelihood = np.floor(-10 * np.log10(1 - min(genotype_ratio, 1)))
            genotype = "0/0"

    return genotype, genotype_likelihood

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("sv_calls", help="BED file of SV calls with reference coordinates in columns 1-3, SV type in 4, length in 5, and contig coordinates in 6-8")
    parser.add_argument("control_regions", help="BED file of control regions with empirical copy number in column 4")
    parser.add_argument("bams", nargs="+", help="one or more BAMs per sample to genotype; read groups must be present in BAM")
    parser.add_argument("output", help="file to emit VCF output to; use /dev/stdout for piping")
    args = parser.parse_args()

    # Load control regions with copy number in column 4. Use copy number 2
    # regions.
    control_regions = pybedtools.BedTool(args.control_regions)
    logger.debug("Loaded %i control regions", len(control_regions))
    copy_2_regions = [region for region in control_regions if region.name == "2"]
    logger.debug("Found %i copy 2 regions", len(copy_2_regions))

    chromosome_sizes = None
    bams_by_name = defaultdict(dict)
    for bam in args.bams:
        bam_file = pysam.AlignmentFile(bam, "rb")

        # Prepare chromosome lengths based on one BAM's header.
        if chromosome_sizes is None:
            chromosome_sizes = dict(zip(bam_file.references, [(0, length) for length in bam_file.lengths]))
            logger.debug("Got %i chromosome sizes", len(chromosome_sizes))

        insert_sizes = []
        for control_region in copy_2_regions:
            #logging.debug("Get insert size distribution from copy 2 region %s", control_region)
            insert_sizes.extend(get_insert_sizes_for_region(bam_file, control_region))

        mean_insert_size = np.median(insert_sizes)
        std_insert_size = np.std(insert_sizes)
        lower_insert_threshold = int(mean_insert_size - (DEVIATIONS_FOR_THRESHOLD * std_insert_size))
        upper_insert_threshold = int(mean_insert_size + (DEVIATIONS_FOR_THRESHOLD * std_insert_size))

        logger.debug("Mean insert size at copy 2 regions for %s is %s", bam, mean_insert_size)
        logger.debug("Standard deviation of insert size at copy 2 regions for %s is %s", bam, std_insert_size)

        bams_by_name[bam]["file"] = bam_file
        bams_by_name[bam]["mean_insert_size"] = mean_insert_size
        bams_by_name[bam]["std_insert_size"] = std_insert_size
        bams_by_name[bam]["lower_insert_threshold"] = lower_insert_threshold
        bams_by_name[bam]["upper_insert_threshold"] = upper_insert_threshold

        # Calculate concordant support across all copy 2 regions.
        logger.debug("Calculating depth at copy 2 regions for %s", bam)
        # control_region_depths = []
        # for control_region in copy_2_regions:
        #     # Keep only the depth for concordant reads at the edge of each
        #     # control region plus or minus the mean insert size of this BAM's
        #     # library. This provides a baseline for SV breakpoints.
        #     breakpoint_intervals = [
        #         (control_region.chrom, control_region.start, control_region.start)
        #     ]

        #     breakpoints = pybedtools.BedTool(breakpoint_intervals).set_chromsizes(chromosome_sizes).slop(b=int(mean_insert_size)).merge()
        #     logging.debug("Testing copy 2 region %s", breakpoints)
        #     depth = get_depth_for_region(bams_by_name[bam], breakpoints, breakpoints)[0]
        #     control_region_depths.append(depth)

        # mean_depth = np.mean(control_region_depths)
        # std_depth = np.std(control_region_depths)
        # logger.debug("Mean depth at copy 2 regions for %s is %s", bam, mean_depth)
        # logger.debug("Standard deviation of depth at copy 2 regions for %s is %s", bam, std_depth)

        # bams_by_name[bam]["mean_depth"] = mean_depth
        # bams_by_name[bam]["std_depth"] = std_depth
        bams_by_name[bam]["std_depth"] = 5

        # # # Calculate mean read depth across all control regions.
        # # logger.debug("Calculating depth at control regions for %s", bam)
        # # #windowed_control_regions = control_regions.window_maker(b=control_regions, w=5000, i="src")
        # # control_region_depths = []
        # # for control_region in control_regions:
        # #     logging.debug("Testing control region %s", control_region)
        # #     depth = get_depth_for_region(bam_file, control_region)
        # #     control_region_depths.append((control_region.name, "%.2f" % depth))

        # # with open("control_depths.tab", "w") as oh:
        # #     control_writer = csv.writer(oh, delimiter="\t", lineterminator="\n")
        # #     control_writer.writerow(("copy", "depth"))
        # #     for depth in control_region_depths:
        # #         control_writer.writerow(depth)

    concordant_reads = pysam.AlignmentFile("concordant_reads.bam", "wb", header=bams_by_name.values()[0]["file"].header)

    #with open(args.output, "w") as output_fh:
        #writer = csv.writer(output_fh, delimiter="\t", lineterminator="\n")
    columns = ("chr", "start", "end", "sv_call", "concordant", "discordant")

    sv_calls = pybedtools.BedTool(args.sv_calls)

    results = []
    for sv_call in sv_calls:
        if sv_call.name == "deletion":
            breakpoint_intervals = [
                (sv_call[5], int(sv_call[6]), int(sv_call[7])),
            ]
            reference_intervals = [
                (sv_call[0], int(sv_call[1]), int(sv_call[1]) + 1),
                (sv_call[0], int(sv_call[2]) - 1, int(sv_call[2]))
            ]
            reference_call_type = "insertion"
        else:
            breakpoint_intervals = [
                (sv_call[5], int(sv_call[6]), int(sv_call[6]) + 1),
                (sv_call[5], int(sv_call[7]) - 1, int(sv_call[7]))
            ]
            reference_intervals = [
                (sv_call[0], int(sv_call[1]), int(sv_call[1]) + 1),
            ]
            reference_call_type = "deletion"

        logger.debug("Breakpoint intervals: %s", breakpoint_intervals)
        breakpoint_intervals = pybedtools.BedTool(breakpoint_intervals)
        reference_intervals = pybedtools.BedTool(reference_intervals)

        for bam_name, bam in bams_by_name.iteritems():
            # Inspect either side of the SV breakpoint(s) by adding the mean
            # insert size of the read pairs to either side of each
            # breakpoint. Merge overlapping breakpoints to avoid double
            # counting support across adjacent breakpoints.
            breakpoint_regions = breakpoint_intervals.set_chromsizes(chromosome_sizes).slop(b=int(bam["mean_insert_size"])).merge()

            # Get concordant read pairs across the breakpoint intervals of
            # this SV.
            breakpoint_concordant = get_depth_for_region(bam, breakpoint_regions, breakpoint_intervals, sv_call.name)

            # Inspect either side of the corresponding SV breakpoint(s) in
            # the reference by adding the mean insert size of the read pairs
            # to either side of each breakpoint. Merge overlapping
            # breakpoints to avoid double counting support across adjacent
            # breakpoints.
            reference_regions = reference_intervals.set_chromsizes(chromosome_sizes).slop(b=int(bam["mean_insert_size"])).merge()

            # Get concordant read pairs across the breakpoint_intervals of
            # the corresponding region of the reference. These pairs are
            # discordant with respect to the SV haplotype.
            breakpoint_discordant = get_depth_for_region(bam, reference_regions, reference_intervals, reference_call_type)

            # Save concordant and discordant read pairs to their respective
            # BAMs.
            for read_pair in breakpoint_concordant:
                for read in read_pair:
                    concordant_reads.write(read)

            for read_pair in breakpoint_discordant:
                for read in read_pair:
                    concordant_reads.write(read)

            breakpoint_concordant_depth, breakpoint_discordant_depth = len(breakpoint_concordant), len(breakpoint_discordant)

            # Calculate the size of the regions inspected and use the ratio
            # to adjust read depths for insertions.
            breakpoint_regions_size = sum([region.end - region.start for region in breakpoint_regions])
            reference_regions_size = sum([region.end - region.start for region in reference_regions])
            logger.debug("Breakpoint regions size: %s", breakpoint_regions_size)
            logger.debug("Reference regions size: %s", reference_regions_size)

            if sv_call.name == "insertion":
                breakpoint_concordant_depth = breakpoint_concordant_depth / (breakpoint_regions_size / float(reference_regions_size))
            else:
                breakpoint_discordant_depth = breakpoint_discordant_depth / (reference_regions_size / float(breakpoint_regions_size))

            logger.debug("Found concordant depth for %s: %s", sv_call.name, breakpoint_concordant_depth)
            logger.debug("Found discordant depth for %s: %s", sv_call.name, breakpoint_discordant_depth)

            #genotype, genotype_likelihood = genotype_call_with_read_pair(breakpoint_concordant_depth, breakpoint_discordant_depth, bam["std_depth"])
            #logger.debug("%s had %s concordant, %s discordant, %s stddev for genotype of %s (GL: %s)", sv_call.name, breakpoint_concordant_depth, breakpoint_discordant_depth, bam["std_depth"], genotype, genotype_likelihood)

            results.append((sv_call[-3:] + [sv_call.name, breakpoint_concordant_depth, breakpoint_discordant_depth]))

        # variant_id = "_".join(row[:EVENT_LENGTH + 1])
        # reference_allele = "N"
        # alternate_allele = "<%s>" % row[EVENT_TYPE]
        # variant_quality = 30
        # filter_status = "PASS"
        # info_pairs = (("END", row[CONTIG_END]),)
        # info = ";".join(["=".join(pair) for pair in info_pairs])
        # format = "GT"

        # output_row = (
        #     row[CONTIG_NAME],
        #     row[CONTIG_START],
        #     variant_id,
        #     reference_allele,
        #     alternate_allele,
        #     variant_quality,
        #     filter_status,
        #     info,
        #     format
        # )

        #writer.writerow(output_row)

    logger.debug("Genotyping...")
    df = pd.DataFrame(results, columns=columns)
    haplotype_ratio = np.log10(np.array((df["concordant"] + 1) / (df["discordant"] + 1)))
    g = mixture.GMM(n_components=3)
    g.fit(haplotype_ratio)

    # Create a lookup table of genotypes for the classes assigned by the GMM
    # based on the assumption that there will be three genotypes for which the
    # lowest mean will be homozygous reference, the next mean will be
    # heterozygous, and the highest mean will be homozygous alternate.
    genotype_by_class = dict(
        zip(
            np.array(sorted(list(enumerate(g.means_.flatten())), key=operator.itemgetter(1)), dtype=np.int8)[:, 0],
            ("0/0", "1/0", "1/1")
        )
    )

    # Calculate genotypes.
    genotypes = g.predict(haplotype_ratio)

    # Assign a standard genotype value for each class calculated by the GMM.
    df["genotype"] = [genotype_by_class.get(genotype, "./.") for genotype in genotypes]

    # Calculate a Phred-scale genotype likelihood from the probability that the
    # assigned genotype from the GMM is correct.
    probabilities = g.predict_proba(haplotype_ratio)
    likelihoods = np.rint(-10 * np.log10(1 - np.apply_along_axis(max, 1, probabilities)))
    likelihoods[np.isinf(likelihoods)] = MAX_LIKELIHOOD
    df["genotype_likelihood"] = likelihoods.astype(int)

    # Convert genotypes for low coverage regions to unknowns and set likelihoods
    # to 0.
    df.loc[(df["concordant"] < 5) & (df["discordant"] < 5), "genotype"] = "./."
    df.loc[(df["concordant"] < 5) & (df["discordant"] < 5), "genotype_likelihood"] = 0

    # Save genotypes.
    df.to_csv(args.output, sep="\t", index=False)

    concordant_reads.close()


import csim
from gym import spaces
from gym.utils import seeding
from time import localtime, strftime
import numpy as np
import copy
import sys
import matplotlib.pyplot as plt
from norm import Normalizer
from collections import deque

# SCENARIO = "Basic"
# SCENARIO = "Passing"
# SCENARIO = "Dot"
SCENARIO = "Hallway"

class Env:
	def __init__(self):
		self.AGENT_NUM = 1
		self.OBSTACLE_NUM = 6
		self.Parser = csim.Parser(SCENARIO, self.AGENT_NUM, self.OBSTACLE_NUM)

		body_dim = 18
		vision_dim = 36
		vel_dim = 36
		dim = body_dim + vision_dim + vel_dim

		ob_low = 0.0 * np.ones(dim) # 14 + 45 + 45 = 104
		ob_high =1.0 * np.ones(dim)

		self.normalizer_ac = Normalizer(
			real_val_max=1.0*np.ones(3),
			real_val_min=-1.0*np.ones(3),
			norm_val_max=3*np.ones(3),
			norm_val_min=-3*np.ones(3)
		)

		for i in range(2):
			ob_low[i] = -1.0

		ob_low[3] = -1.0
		ob_high[3] = 1.0
		ob_low[4] = -0.5
		ob_high[4] = 2.0
		ob_low[4] = -1.5
		ob_high[4] = 1.5

		# vel observation
		for i in range(body_dim+vision_dim, dim):
			ob_low[i] = -3.0
			ob_high[i] = 3.0

		self.observation_space = spaces.Box(
			low = ob_low,
			high = ob_high
		)

		self.action_space = spaces.Box(
			low = self.normalizer_ac.norm_val_min,
			high = self.normalizer_ac.norm_val_max
		)

		self.seed()
		self.reward_range = (-float('inf'), float('inf'))
		self.metadata = {'render.modes' : []}
		self.spec = None

		self.cur_step = 0
		self.initPlot()

# ================plot=============================
	def initPlot(self):
		self.fig = plt.figure()
		self.ax1 = self.fig.add_subplot(111)
		self.ax1.plot([], label='reward', color='black')
		self.ax1.plot([], label='target', color='red')
		self.ax1.plot([], label='pref', color='green')
		self.ax1.plot([], label='smooth', color='blue')
		self.ax1.plot([], label='bubble', color = 'orange')
		self.ax1.plot([], label='col', color='purple')
		self.ax1.plot([], label='dir', color='cyan')
		self.ax1.legend()

		self.resetRewardCum()

		self.reward_queue = deque(maxlen=100)
		self.target_queue = deque(maxlen=100)
		self.prefV_queue = deque(maxlen=100)
		self.smooth_queue = deque(maxlen=100)
		self.col_queue = deque(maxlen=100)
		self.bubble_queue = deque(maxlen=100)
		self.dir_queue = deque(maxlen=100)

		self.reward_plot = []
		self.target_plot = []
		self.prefV_plot = []
		self.smooth_plot = []
		self.col_plot = []
		self.bubble_plot = []
		self.dir_plot = []

	def resetRewardCum(self):
		self.reward_cum = 0
		self.reward_target_cum = 0
		self.reward_prefV_cum = 0
		self.reward_smooth_cum = 0
		self.reward_col_cum = 0
		self.reward_bubble_cum = 0
		self.reward_dir_cum = 0

	def setRewardPlot(self, reward, rewards):
		self.reward_target_cum += rewards[0]
		self.reward_prefV_cum += rewards[1]
		self.reward_smooth_cum += rewards[2]
		self.reward_col_cum += rewards[3]
		self.reward_bubble_cum += rewards[4]
		self.reward_dir_cum += rewards[5]

		self.reward_cum += rewards[0]+rewards[1]+rewards[2]+rewards[3]+rewards[4] + rewards[5]

	def appendPlot(self):
		self.reward_queue.append(self.reward_cum)
		self.target_queue.append(self.reward_target_cum)
		self.prefV_queue.append(self.reward_prefV_cum)
		self.smooth_queue.append(self.reward_smooth_cum)
		self.col_queue.append(self.reward_col_cum)
		self.bubble_queue.append(self.reward_bubble_cum)
		self.dir_queue.append(self.reward_dir_cum)

	def drawPlot(self):
		self.reward_plot.append(np.mean(self.reward_queue))
		self.target_plot.append(np.mean(self.target_queue))
		self.prefV_plot.append(np.mean(self.prefV_queue))
		self.smooth_plot.append(np.mean(self.smooth_queue))
		self.bubble_plot.append(np.mean(self.bubble_queue))
		self.col_plot.append(np.mean(self.col_queue))
		self.dir_plot.append(np.mean(self.dir_queue))

		plt.ion()
		self.ax1.plot(self.reward_plot, label='reward', color='black')
		self.ax1.plot(self.target_plot, label='target', color='red')
		self.ax1.plot(self.prefV_plot, label='pref', color='green')
		self.ax1.plot(self.smooth_plot, label='smooth', color='blue')
		self.ax1.plot(self.bubble_plot, label='col', color='orange')
		self.ax1.plot(self.col_plot, label='col', color='purple')
		self.ax1.plot(self.dir_plot, label='dir', color='cyan')
		plt.draw()
		plt.pause(0.1)
		plt.show()

	def plotSave(self):
		plt.savefig("../data/ckpt/network/graph"+str(strftime("%Y%m%d_%I%M", localtime())))

# =================================================

	def step(self, a):
		memory = self.Parser.Step(self.actions, False)

		obs = memory['obs']
		obs = self.convert_to_numpy(obs)

		body_state = obs['agent'][0]['body_state']
		sensor_state = obs['agent'][0]['sensor_state']
		velocity_state = obs['agent'][0]['velocity_state']
		state = np.concatenate((body_state, sensor_state, velocity_state), axis=None)
		reward = memory['reward']
		isTerm = memory['isTerm']

		if np.isnan(state).any():
			print("state nan")
			sys.exit(1)

		rewards = memory['reward_sep']
		self.cur_step += 1

		self.setRewardPlot(reward, rewards)
		if isTerm:
			self.appendPlot()
			self.resetRewardCum()

		return state, reward, isTerm, {}

	def steps(self, a_):
		a_norm = []
		for ac in a_:
			a_norm.append(self.normalizer_ac.norm_to_real(ac))

		self.actions =np.array(a_norm, dtype=np.float32)

	def step_render(self, a_):
		a_norm = []
		for ac in a_:
			a_norm.append(self.normalizer_ac.norm_to_real(ac))
			# if ac[0] > 3.0 or ac[1] > 3.0 or ac[2] > 3.0:
			# print("ac : ", ac)
			# print("ac_ : ", self.normalizer_ac.norm_to_real(ac))

		memory = self.Parser.Step(np.array(a_norm, dtype=np.float32), True)

		obs = memory['obs']
		obs = self.convert_to_numpy(obs)

		# print("action : ", a_norm)
		# print("state : ", obs['agent'][0]['body_state'])

		state = obs
		reward = memory['reward']
		isTerm = memory['isTerm']

		return state, reward, isTerm, {}

	def states(self):
		obs = self.Parser.Observe()
		obs = self.convert_to_numpy(obs)

		states = []
		for i in range(len(obs['agent'])):
			body_state = obs['agent'][i]['body_state']
			sensor_state = obs['agent'][i]['sensor_state']
			velocity_state = obs['agent'][i]['velocity_state']
			state = np.concatenate((body_state, sensor_state, velocity_state), axis=None)
			states.append(state)

		return states

	def seed(self, seed=None):
		self.np_random, seed = seeding.np_random(seed)
		return [seed]

	def reset(self):
		self.Parser.Reset(-1, self.AGENT_NUM, self.OBSTACLE_NUM)
		obs = self.Parser.Observe()
		obs = self.convert_to_numpy(obs)

		body_state = obs['agent'][0]['body_state']
		sensor_state = obs['agent'][0]['sensor_state']
		velocity_state = obs['agent'][0]['velocity_state']
		state = np.concatenate((body_state, sensor_state, velocity_state), axis=None)

		return state

	def observe(self):
		obs = self.Parser.Observe()
		obs = self.convert_to_numpy(obs)

		return obs

	def convert_to_numpy(self, obs):
		for i in range(self.AGENT_NUM): # r, p, d, front, color
			r_data = obs['agent'][i]['render_data']
			obs['agent'][i]['r'] = np.array([r_data[0], r_data[1]])
			obs['agent'][i]['p'] = np.array([r_data[2], r_data[3]])
			obs['agent'][i]['d'] = np.array([r_data[4], r_data[5]])
			obs['agent'][i]['front'] = r_data[6]
			obs['agent'][i]['color'] = np.array([r_data[7], r_data[8], r_data[9]])
			obs['agent'][i]['body_state'] = np.array(obs['agent'][i]['body_state'])
			obs['agent'][i]['sensor_state'] = np.array(obs['agent'][i]['sensor_state'])
			obs['agent'][i]['velocity_state'] = np.array(obs['agent'][i]['velocity_state'])
			obs['agent'][i]['offset'] = np.array(obs['agent'][i]['offset_data'])

		obstacle_num = len(obs['obstacle'])
		for i in range(obstacle_num):
			obs['obstacle'][i]['p'] = np.array(obs['obstacle'][i]['p'])
			obs['obstacle'][i]['r'] = np.array(obs['obstacle'][i]['r'])
			obs['obstacle'][i]['front'] = np.array(obs['obstacle'][i]['front'])

		wall_num = len(obs['wall'])
		for i in range(wall_num):
			obs['wall'][i]['p'] = np.array(obs['wall'][i]['p'])
			obs['wall'][i]['w'] = obs['wall'][i]['w']
			obs['wall'][i]['h'] = obs['wall'][i]['h']

		return obs


import retro
from gym.envs.classic_control import rendering
import numpy as np
import cv2
import neat
import pickle
import sys
import win32api

viewer = rendering.SimpleImageViewer()
env = retro.make('SuperMarioWorld-Snes', 'YoshiIsland2Yellow')

imgarray = []

xpos_end = 0

def kill():
    mx, my = win32api.GetCursorPos()
    if mx==0 and my==0:
        print("over")
        cv2.destroyAllWindows()
        env.reset()
        env.viewer = None
        env.close()
        env.render(close=True)
        sys.exit()

def repeat_upsample(rgb_array, k=1, l=1, err=[]):
    # repeat kinda crashes if k/l are zero
    if k <= 0 or l <= 0:
        if not err:
            print ("Number of repeats must be larger than 0, k: {}, l: {}, returning default array!".format(k, l))
            err.append('logged')
        return rgb_array

    # repeat the pixels k times along the y axis and l times along the x axis
    # if the input image is of shape (m,n,3), the output image will be of shape (k*m, l*n, 3)

    return np.repeat(np.repeat(rgb_array, k, axis=0), l, axis=1)

def eval_genomes(genomes, config):

    for genome_id, genome in genomes:
        ob = env.reset()
        ac = env.action_space.sample()

        inx, iny, inc = env.observation_space.shape

        inx = int(inx/6)
        iny = int(iny/6)

        net = neat.nn.recurrent.RecurrentNetwork.create(genome, config)

        current_max_fitness = 0
        fitness_current = 0
        frame = 0
        counter = 0
        xMax = 0

        done = False
        cv2.namedWindow("main", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("main", (256*2,240*2))
        cv2.namedWindow("ob", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("ob", (256*2,240*2))

        while not done:
            kill()
            rgb = env.render('rgb_array')
            upscaled=repeat_upsample(rgb,3, 3)
            viewer.imshow(upscaled)

            frame += 1
            scaledimg = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
            scaledimg = cv2.resize(scaledimg, (inx, iny))
            #ob = cv2.resize(ob, (inx, iny))
            ob = cv2.cvtColor(scaledimg, cv2.COLOR_RGB2GRAY)
            #ob = np.reshape(ob, (inx,iny))

            cv2.imshow('main', scaledimg)
            cv2.imshow('ob', ob)
            #cv2.waitKey(1)

            imgarray = np.ndarray.flatten(ob)

            nnOutput = net.activate(imgarray)
            #print(nnOutput)
            ob, rew, done, info = env.step(nnOutput)

            #print(info['lives'])

            #xpos = info['x']
            #xpos_end = info['screen_x_end']

            #make sure he goes right
            if xMax < info['x']:
                xMax = info['x']
            fitness_current += rew

            if fitness_current > current_max_fitness:
                current_max_fitness = fitness_current
                counter = 0
            else:
                counter += 1

            if done or counter == 350 or info['lives'] == 3:
                done = True
                fitness_current += info['score']
                print(genome_id, fitness_current)

            genome.fitness = fitness_current

config = neat.Config(neat.DefaultGenome, neat.DefaultReproduction,
                     neat.DefaultSpeciesSet, neat.DefaultStagnation,
                     'config-feedforward-mw')

p = neat.Population(config)
#p = neat.Checkpointer.restore_checkpoint("neat-checkpoint-260")

p.add_reporter(neat.StdOutReporter(True))
stats = neat.StatisticsReporter()
p.add_reporter(stats)
p.add_reporter(neat.Checkpointer(10))

winner = p.run(eval_genomes)

with open('winner.pkl', 'wb') as output:
    pickle.dump(winner, output, 1)


# import the necessary packages
from collections import OrderedDict
import numpy as np
import cv2
import dlib
import imutils
from pynput.mouse import Button, Controller
import array as arr
webcam = cv2.VideoCapture("./foo.jpg")

facial_features_cordinates = {}
classes = None

with open("yolov3.txt", 'r') as f:
    classes = [line.strip() for line in f.readlines()]
mouse=Controller()

change = 0;
prevch = 0;
changed = False
def get_output_layers(net):

    layer_names = net.getLayerNames()

    output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]

    return output_layers

def draw_prediction(img, class_id, confidence, x, y, x_plus_w, y_plus_h):
    global changed
    if class_id==0:
        Width = img.shape[1]
        Height = img.shape[0]
        label = str(classes[class_id])

        color = COLORS[class_id]

        cv2.rectangle(img, (x,y), (x_plus_w,y_plus_h), color, 2)
        cv2.circle(image, (int(Width/2),int(Height/2)), 5, (0,255,255))
        midx = int(Width/2)
        midy = int(Height/2)
        avx = int((x+x_plus_w)/2)
        avy = int((y+y_plus_h)/2)
        cv2.line(img,(avx,avy),(midx+5,midy),(0,0,255))
        cv2.line(img,(avx,avy),(midx-5,midy),(0,0,255))
        print("Distance from centre to general location:{0}".format(midx-avx))
        change = abs(midx-avx)
        #true = spin, false = stop
        if change<150:
            if (changed):
                print("stop rotate")
                mouse.click(Button.left, 1)
                changed = False
        else:
            if not changed:
                print("rotate")
                mouse.click(Button.left, 1)
                changed = True
        cv2.putText(img, label, (x-10,y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

COLORS = np.random.uniform(0, 255, size=(len(classes), 3))

net = cv2.dnn.readNet("yolov3.weights", "yolov3.cfg")
while True:
    webcam = cv2.VideoCapture("./foo.jpg")
    res, frame = webcam.read()
    if res:

        image = frame
        #print('Original Dimensions : ',image.shape)
        #image = imutils.resize(image, width=500)
        Width = image.shape[1]
        Height = image.shape[0]
        scale = 0.00392
        blob = cv2.dnn.blobFromImage(image, scale, (416,416), (0,0,0), True, crop=False)
        net.setInput(blob)
        outs = net.forward(get_output_layers(net))
        class_ids = []
        confidences = []
        boxes = []
        conf_threshold = 0.5
        nms_threshold = 0.4
        for out in outs:
            for detection in out:
                scores = detection[5:]
                class_id = np.argmax(scores)
                confidence = scores[class_id]
                cv2.circle(image, (int(Width/2),int(Height/2)), 5, (0,255,255))
                if confidence > 0.5:
                    center_x = int(detection[0] * Width)
                    center_y = int(detection[1] * Height)
                    w = int(detection[2] * Width)
                    h = int(detection[3] * Height)
                    #print(Width/2)

                    x = center_x - w / 2
                    y = center_y - h / 2
                    class_ids.append(class_id)
                    confidences.append(float(confidence))
                    boxes.append([x, y, w, h])

        indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)
        if not len(indices):
            if not changed:
                print("rotate")
                mouse.click(Button.left, 1)
                changed = True
        for i in indices:
            i = i[0]
            box = boxes[i]
            x = box[0]
            y = box[1]
            w = box[2]
            h = box[3]
            draw_prediction(image, class_ids[i], confidences[i], round(x), round(y), round(x+w), round(y+h))

#image = cv2.imread("images/image_1.jpg")

webcam.release()
cv2.destroyAllWindows()

# detect faces in the grayscale image

# loop over the face detections


#!/usr/bin/env python

# multiplexContours.py #################
#
# A "multiplexer" to combine many contours (like for different SRs) using the expected CLs as the figure of merit
#
# By: Larry Lee - Mar 2018
# Updated in June 2020 by Jonathan Long

import pickle
import ROOT
import math,sys,os,argparse
import itertools
import copy

ROOT.gROOT.SetBatch()

parser = argparse.ArgumentParser()
parser.add_argument('--inputFiles','-i', type=str, nargs='+', help='input json files', required=True)
parser.add_argument("--outputFile","-o", type = str, help="output ROOT file", default = "multiplexedContours.root")
parser.add_argument("--debug","-d",      help = "print extra debugging info", action="store_true", default=False)
parser.add_argument("--skipTheory", action="store_true", help="Skip theory variations")
parser.add_argument("--plotUsedContours", "-p", action="store_true", help="Generate a plot of the used contours.")
parser.add_argument("--ignoreInteriorLines", action="store_true", help="Ignore subRegions that are contained in the total expected. ")
parser.add_argument("--noSubRegionMerge", action="store_false", help="Do not merge split subregions.")
parser.add_argument("--showPastLimit", action="store_true", help="Show the best expected region past the limit when using --plotUsedContours.")
parser.add_argument("--minArea", type=float, default=-1, help="Remove subregions below this area.  Helps to remove very small subRegions by setting this to 1.")
parser.add_argument("--doEntirePlane", action="store_true", help="Pick best contour for the entire plane, i.e., well out side of the obs/exp contours.  This will be slower, but you will need this to show upper limit values outside of the exclusion.")

args = parser.parse_args()

## If I need to use scipy, please let me have scipy. I'll even help you out!
try:
    import matplotlib as mpl
    mpl.use('Agg')
    import matplotlib.pyplot as plt
    import matplotlib.patches as mpatches
    from matplotlib.lines import Line2D
    import numpy as np
    import scipy.interpolate

    from shapely.ops import cascaded_union, polygonize, split, unary_union
    from shapely.geometry import Polygon,MultiPolygon
    from shapely.geometry import LineString, Point
    from shapely.geometry.collection import GeometryCollection
    from shapely import affinity
    from shapely.geometry.base import BaseMultipartGeometry

except:
    print (">>> You need scipy/matplotlib and shapely to run this. And you had to have run harvestToContours in scipy mode [default]")
    print (">>> In an ATLAS environment, you can...")
    print ('>>> > lsetup "views LCG_94 x86_64-centos7-gcc62-opt" # or replace with a newer LCG release or more appropriate arch')
    print (">>> ")
    print (">>> Try that and then run this again!")
    sys.exit(1)

def main():

    print (">>> Launching multiplexContours.py! ")

    dict_TFiles = {}
    dict_Surfaces = {}
    dict_rbf_Surf = {}
    dict_Exp = {}
    dict_Obs = {}
    dict_UL = {}

    dict_Exp_u1s = {}
    dict_Exp_d1s = {}

    dict_Obs_u1s = {}
    dict_Obs_d1s = {}

    isoExpectedContours = {}

    listOfInputFiles = [tmp.translate({None:", "}) for tmp in args.inputFiles]

    print (">>> Grabbing input information from input files:")
    print (">>> >>> " + " ".join(listOfInputFiles))
    print (">>> ")
    print (">>> FYI, I'm also going to grab the *.expectedSurface.pkl files that should be in the same location")
    print (">>> P.S. using some smoothing during the harvest stage will speed up the contour slicing step and generally provide in a more satisfactory, and less wonky, result.")

    for inputFileName in listOfInputFiles:
        dict_TFiles[inputFileName] = ROOT.TFile(inputFileName)
        try:
            dict_Surfaces[inputFileName] = pickle.load( open( inputFileName+".expectedSurface.pkl", "rb" ) )
        except FileNotFoundError:
            print (">>> I need those *.expectedSurface.pkl files that are produced from harvestToContours! Exiting...")
            sys.exit(1)

        dict_Exp[inputFileName] = dict_TFiles[inputFileName].Get("Exp_0").Clone(inputFileName+"_Exp_0")
        dict_Obs[inputFileName] = dict_TFiles[inputFileName].Get("Obs_0").Clone(inputFileName+"_Obs_0")

        dict_Exp_u1s[inputFileName] = dict_TFiles[inputFileName].Get("SubGraphs/clsu1s_Contour_0").Clone(inputFileName+"_Exp_u1s")
        dict_Exp_d1s[inputFileName] = dict_TFiles[inputFileName].Get("SubGraphs/clsd1s_Contour_0").Clone(inputFileName+"_Exp_d1s")

        if not args.skipTheory:
            dict_Obs_u1s[inputFileName] = dict_TFiles[inputFileName].Get("Obs_0_Up").Clone(inputFileName+"_Obs_u1s")
            dict_Obs_d1s[inputFileName] = dict_TFiles[inputFileName].Get("Obs_0_Down").Clone(inputFileName+"_Obs_d1s")

        dict_UL[inputFileName] = dict_TFiles[inputFileName].Get("upperLimit_gr").Clone(inputFileName+"upperLimit_gr")

    print (">>> Creating output ROOT file: %s"%args.outputFile)

    outputFile = ROOT.TFile(args.outputFile,"RECREATE")

    fig, ax = plt.subplots()

    print (">>> Creating iso expected contours for all pairwise combinations of SRs")

    # Loop over combinations of input files and grab coordinates from surfances in pkl files
    # Thanks to Luigi Longo for the more comprehensive input checks
    for (region1, region2) in itertools.combinations(listOfInputFiles, 2):
        #x,y = dict_Surfaces[region1]["x"], dict_Surfaces[region1]["y"]
        #z1,z2 = dict_Surfaces[region1]["z"], dict_Surfaces[region2]["z"]
        #cs = ax.contour(x,y,z1-z2, ":", levels=[0] ) # this creates contours where the two surfaces have the same Z value
        #
        #isoExpectedContours[(region1,region2)] = cs.allsegs[0] # grab list of arrays of coords of intersection contour

        x,y = dict_Surfaces[region1]["x"], dict_Surfaces[region1]["y"]

        # Used to check that regions have same grid
        x_ref = dict_Surfaces[region2]['x']
        y_ref = dict_Surfaces[region2]['y']

        # Input checks
        if len(x) == len(x_ref) and len(y) == len(y_ref):
            if len(x)>0 and len(y)>0:
                if len(x[0]) == len(x_ref[0]) and len(y[0]) == len(y_ref[0]):
                    if ( np.isclose(x_ref,x).all() and np.isclose(y_ref,y).all() ):
                        z1,z2 = dict_Surfaces[region1]["z"], dict_Surfaces[region2]["z"]
                        cs = ax.contour(x,y,z1-z2, ":", levels=[0] ) # this creates contours where the two surfaces have the same Z value

                        isoExpectedContours[(region1,region2)] = cs.allsegs[0] # grab list of arrays of coords of intersection contour
                    else:
                        if not np.isclose(x_ref,x).all():
                            print(f"X for {region1} from: {x[0][0]:f} to {y[-1][0]:f}")
                            print(f"X for {region2} from: {x_ref[0][0]:f} to {x_ref[-1][0]:f}")
                        if not np.isclose(y_ref,y).all():
                            print(f"Y for {region1} from: {y[0][0]:f} to {y[-1][0]:f}")
                            print(f"Y for {region2} from: {y_ref[0][0]:f} to {y_ref[-1][0]:f}")

                        print(">>> Error, inputs are not close.  %s is different"%(region1))
                        sys.exit(1)

                else:
                    print(">>> Error, inputs must use same x and y limits and resolution in harvest contours.  %s is different"%(region1))
                    sys.exit(1)

    ax.cla()

    uncutRegions = {} # dictionary of all contours for each region

    bestRegions = {}

    # Retrieve and convert tgraphs of various contours into polygons
    for inputFileName in listOfInputFiles:
        region = {}

        region["Exp"]     = tGraphToPolygon(dict_Exp[inputFileName])
        region["Exp_u1s"] = tGraphToPolygon(dict_Exp_u1s[inputFileName])
        region["Exp_d1s"] = tGraphToPolygon(dict_Exp_d1s[inputFileName])
        region["Obs"]     = tGraphToPolygon(dict_Obs[inputFileName])
        if not args.skipTheory:
            region["Obs_u1s"] = tGraphToPolygon(dict_Obs_u1s[inputFileName])
            region["Obs_d1s"] = tGraphToPolygon(dict_Obs_d1s[inputFileName])

        uncutRegions[inputFileName] = region

        bestRegions[inputFileName] = list()

    sumOfExpecteds = cascaded_union([polygons["Exp"] for regionName,polygons in uncutRegions.items()]) # geometric union of all expected contours

    sumOfContours = copy.deepcopy(sumOfExpecteds)
    for key in ["Exp_u1s", "Exp_d1s", "Obs"] + (["Obs_u1s","Obs_d1s"] if not args.skipTheory else []):

        sumOfContours = cascaded_union([polygons[key] for regionName,polygons in uncutRegions.items()] + [sumOfContours]) # geometric union of all contours.  This defines the region where contours are chopped up

    if args.doEntirePlane:
        # make rectangle bound
        x_r = [x[0][0], x[0][-1], x[-1][-1], x[-1][0]]
        y_r = [y[0][0], y[0][-1], y[-1][-1], y[-1][0]]

        p_r = tGraphToPolygon(convertArraysToTGraph(x_r,y_r))
        sumOfContours = cascaded_union([sumOfContours, p_r])

    sumOfExpectedsArea = sumOfExpecteds.area

    if args.debug:
        print(">>> Number of regions being added to the best expected contour: %d"%len(uncutRegions) )
        print(">>> Integral of the best expected contour: %d"%sumOfExpecteds.area)

        x,y = sumOfExpecteds.exterior.coords.xy
        convertArraysToTGraph(x,y).Write("debug_sumOfExpecteds")
        ax.cla()
        ax.plot(x,y,alpha=0.5)
        fig.savefig("debug_sumOfExpecteds.pdf")

        for i,singleExpectedCurve in enumerate([v["Exp"] for k,v in uncutRegions.items()]):
            x,y = singleExpectedCurve.exterior.coords.xy
            ax.cla()
            ax.plot(x,y,alpha=0.5)
            fig.savefig("debug_singleExpected_%d.pdf"%i)

    allIsoExpectedContours = []
    allIsoExpectedContoursLineStrings = []
    allIsoExpectedContours_regionName = []
    for two_region_tuple in isoExpectedContours:

        for intersectionContour in isoExpectedContours[two_region_tuple]:
            # JDL: couldn't this be a list of arrays?

            ## Second clause is better handling of bad geometries like splitting overlapping lines.  Not sure it's always needed, but did fix a crash on a complicated test
            try:
                tmpLineString = LineString(intersectionContour)
                if args.ignoreInteriorLines:
                    if tmpLineString.within(sumOfExpecteds): # check if intersection contour is completely within sum of expecteds
                        continue
                allIsoExpectedContours.append( intersectionContour )
                allIsoExpectedContoursLineStrings.append( tmpLineString )
            except ValueError:
                try:
                    if Polygon(intersectionContour).is_valid:
                        tmpLineString = LineString(intersectionContour)
                        if args.ignoreInteriorLines:
                            if tmpLineString.within(sumOfExpecteds): # check if intersection contour is completely within sum of expecteds
                                continue

                        allIsoExpectedContours.append( intersectionContour )
                        allIsoExpectedContoursLineStrings.append( tmpLineString )

                    else:

                        ls_temp = LineString(intersectionContour)
                        lr = LineString(ls_temp.coords[:]+ls_temp.coords[0:1])

                        mls = unary_union(lr)

                        for ls in mls:
                            allIsoExpectedContours.append(intersectionContour)
                            allIsoExpectedContoursLineStrings.append(ls)

                except ValueError:
                    print(">>> Failed to handle an intersection contour")
                    continue

    # Add cut lines of the input expectected contours.  This helps with irregularities in the final expectected contour seen in a test with EWK comb contours
    for regionName,polygons in uncutRegions.items():
        allIsoExpectedContoursLineStrings.append( LineString(polygons["Exp"].exterior.coords) )

    if args.debug:
        ax.cla()
        x,y = sumOfExpecteds.exterior.xy
        ax.plot(x,y,linewidth=1,color='k')
        for ls in allIsoExpectedContoursLineStrings:
            x,y = ls.coords.xy
            ax.plot(x,y,linewidth=2,color=np.random.rand(3,),linestyle=":")
        fig.savefig("debug_allEICs.pdf")

    print(">>> Total number of isoExpectedContours: %d"%len(allIsoExpectedContoursLineStrings))

    listOfBestCurves = {}

    for tmpKey in uncutRegions[inputFileName]:
        listOfBestCurves[tmpKey] = []

    print (">>> Cutting space up into regions based on iso expected contours")

    # Cut up the union of contours into individual parts
    allIsoExpectedContoursLineStrings = sorted(allIsoExpectedContoursLineStrings, key=lambda x: x.length, reverse=True)

    # Cut up a the sum of all contours
    subRegions = [sumOfContours]

    for cutLine in allIsoExpectedContoursLineStrings: # loop over all cuts, i.e. iso contours
        tmpSubRegions = []
        for subRegion in subRegions: # apply cut to each subRegion

            # special case of interior circle that doesn't split properly.  Only really needed for looking at usedContour figure
            if subRegion.contains(cutLine) and not cutLine.is_ring and len(cutLine.coords) > 2:
                x_start, y_start = cutLine.coords[0:1][0][0], cutLine.coords[0:1][0][1]
                x_end, y_end = cutLine.coords[-2:-1][0][0], cutLine.coords[-2:-1][0][1]
                x_end2, y_end2 = cutLine.coords[-3:-2][0][0], cutLine.coords[-3:-2][0][1]

                dist_close = math.sqrt((x_start-x_end)**2 + (y_start- y_end)**2 )
                dist_neighbor = math.sqrt((x_end2-x_end)**2 + (y_end2- y_end)**2 )

                # if closing the line requires a jump not greater than 1.5x the distance in the 2nd to last point, try it
                if abs(dist_close-dist_neighbor)/dist_neighbor < 1.5:
                    cutLine = LineString(cutLine.coords[:]+cutLine.coords[0:1])
                else:
                    if args.debug:
                        print("Debug: cutline distance too big to close")

            tmpGeoms = split(subRegion, cutLine)
            if len(tmpGeoms) == 1:
                # Shapely is supposed to return the original poly if no split is made, but have seen edge cases where it does not.
                tmpGeoms=[subRegion]

            tmpGeomsClean = []
            # Clean up some regions that might cause issues later
            for tg in tmpGeoms:
                # Get rid of very small regions
                if tg.area < 1e-10:
                    continue

                # Another corner case.  This removes stray lines coming off of polygons.
                tmp = tg.buffer(-1e-1).buffer(5e-1).intersection(tg)
                if type(tmp) == MultiPolygon:
                    for t in tmp:
                        if t.area>0:
                            tmpGeomsClean.append(t)
                elif type(tmp) == GeometryCollection:
                    for t in tmp:
                        if type(t) != LineString and type(t) != Point:
                            if t.area>0:
                                tmpGeomsClean.append(t)
                else:
                    tmpGeomsClean.append( tmp )

            tmpSubRegions.extend(tmpGeomsClean)

        subRegions = tmpSubRegions

    if args.debug:
        print(">>> Number of regions in this plane: %d"%len(subRegions))

    if args.debug: # JDL
        ax.cla()
        for i,poly in enumerate(subRegions):
            x,y = poly.exterior.xy
            ax.plot(x,y,linewidth=2,color=np.random.rand(3,),linestyle=":")
        x,y = sumOfExpecteds.exterior.xy
        ax.plot(x,y,linewidth=1,color='k')
        fig.savefig("debug_CutUpExpected.pdf")

    # Gather points to do interpretation only once
    ZI_reg_dict = {}
    interiorPoints = []
    for subRegion in subRegions:
        x_cent,y_cent = subRegion.representative_point().xy
        interiorPoints.append([x_cent[0], y_cent[0]])

    for regionName,region in uncutRegions.items():
        x,y = dict_Surfaces[regionName]["x"], dict_Surfaces[regionName]["y"]
        z = dict_Surfaces[regionName]["z"]
        ZI_reg_dict[regionName] = scipy.interpolate.griddata( (x.flatten(), y.flatten()), z.flatten(), interiorPoints, method="linear" )

    # Fill list of best SR per subRegion
    subRegion_SR_dict = {}
    for iSubRegion,subRegion in enumerate(subRegions):
        bestSR = ""
        bestZvalue = -100

        # See if the test point is inside of an expected curve
        interiorPoint = Point(interiorPoints[iSubRegion])
        isInExpecteds = sumOfExpecteds.contains(interiorPoint)

        if isInExpecteds: # only consider SRs where point is in expected
            for regionName,region in uncutRegions.items():
                if region["Exp"].contains(interiorPoint):
                    ZI = ZI_reg_dict[regionName][iSubRegion]

                    if ZI > bestZvalue:
                        bestZvalue = ZI
                        bestSR = regionName
        else: # just go through and find the best point
            for regionName,region in uncutRegions.items():
                ZI = ZI_reg_dict[regionName][iSubRegion]

                if ZI > bestZvalue:
                    bestZvalue = ZI
                    bestSR = regionName

        subRegion_SR_dict[subRegion.area]=bestSR

    nSubRegionsPre = len(subRegions)

    # Merge subregions, speeds up later steps
    if args.noSubRegionMerge:
        i = 0
        while i < len(subRegions):
            subRegion = subRegions[i]

            match = False
            for j,sR in enumerate(subRegions):
                if i==j: continue
                if  (doPolygonsTouch(subRegion,sR) or subRegion.contains(sR) or sR.contains(subRegion)) and subRegion_SR_dict[subRegion.area] == subRegion_SR_dict[sR.area]: # do Polygons touch or does one contain the other and have the same assigned SR
                    newPoly = cascaded_union([subRegion, sR])

                    if type(newPoly) == MultiPolygon: continue # Even though they shouldn't, don't merge things that don't form a single polygon

                    subRegions.remove(subRegion)
                    subRegions.remove(sR)
                    subRegions.append(newPoly)

                    subRegion_SR_dict[newPoly.area] = subRegion_SR_dict[subRegion.area]
                    match=True
                    break
            if not match:
                i+=1
        print(">>> Number of subRegions reduced from %d to %d"%(nSubRegionsPre,len(subRegions)))
    else:
        print(">>> Number of subRegions %d"%(len(subRegions)))

    if args.minArea>0:
        areaRemoved = 0.
        for sR in subRegions:
            if sR.area < args.minArea:
                areaRemoved += sR.area
                subRegions.remove(sR)
        print(">>>Area removed over total %f / %f"%(areaRemoved,sumOfContours.area))
        print(">>> New number of subRegions %d"%(len(subRegions)))

    if args.debug: # JDL
        ax.cla()
        for i,poly in enumerate(subRegions):
            x,y = poly.exterior.xy
            ax.plot(x,y,linewidth=2,color=np.random.rand(3,),linestyle=":")
        x,y = sumOfExpecteds.exterior.xy
        ax.plot(x,y,linewidth=1,color='k')
        fig.savefig("debug_CutUpExpected_Merged.pdf")

    # Now go through subcurves and dice them up
    for iSubRegion,subRegion in enumerate(subRegions): # subRegions are the cut up space associated with the best SR in each region
        if args.debug:
            print(">>> >>> Loop through sub regions: %d"%iSubRegion)

        # this subRegion -- Figure out which SR it corresponds to

        # Pick Best SR
        bestSR = subRegion_SR_dict[subRegion.area]

        if bestSR=="":
            allIsoExpectedContours_regionName.append("none")
            print(">>> Skipping a subRegion because I couldn't find it's SR")
            continue
        else:
            allIsoExpectedContours_regionName.append(bestSR)

        # now I know that this subRegion should go with bestSR
        # so now I need to loop over all other types of curves associated with best SR
        # and cut them up with all the IECs (e.g. obs, exp+1sig, etc)

        if args.debug:
            print(">>> >>> Identified the best SR in this region as %s"%bestSR)

        if args.showPastLimit:
            bestRegions[bestSR].append(subRegion)

        # Now loop through types of curves and cookie cut out shapes based on subRegions
        debugCounter = 0
        for typeOfCurve, srCurve in uncutRegions[bestSR].items(): #typeOfCurve is 'Exp', 'Obs' etc, srCurve is then that curves polygon, i.e. a single SR

            if subRegion.intersects( srCurve):
                cookieCut = subRegion.intersection(srCurve)
                if type(cookieCut) == Polygon:
                    listOfBestCurves[typeOfCurve].append(cookieCut)
                    if not args.showPastLimit and typeOfCurve == 'Exp': bestRegions[bestSR].append(cookieCut) # keep best exp curves for each SR
                elif type(cookieCut) == MultiPolygon:
                    for poly in cookieCut:
                        listOfBestCurves[typeOfCurve].append(poly)
                        if not args.showPastLimit and typeOfCurve == 'Exp': bestRegions[bestSR].append(poly) # keep best exp curves for each SR
                elif type(cookieCut) == GeometryCollection:
                    for poly in cookieCut:
                        if type(poly) == Polygon:
                            listOfBestCurves[typeOfCurve].append(poly)
                            if not args.showPastLimit and typeOfCurve == 'Exp': bestRegions[bestSR].append(poly) # keep best exp curves for each SR
                        else:
                            if args.debug:
                                print(">>> Possibly skipping something: geom obj type: %s"%(type(poly)))
                                ax.cla()
                                x,y = subRegion.exterior.xy
                                ax.plot(x,y,linewidth=2,color='k',linestyle=":")
                                x,y = srCurve.exterior.xy
                                ax.plot(x,y,linewidth=2,color='r',linestyle=":")

                                x,y = poly.coords.xy
                                ax.plot(x,y,linewidth=2,color='b',linestyle="-")

                                fig.savefig(f"debug_nonpolygeocolcut_{debugCounter}.pdf")

                else:
                    if args.debug:
                        print(">>> Possibly skipping something: cookieCut type: %s"%(type(cookieCut)))
                        ax.cla()
                        x,y = subRegion.exterior.xy
                        ax.plot(x,y,linewidth=2,color='k',linestyle=":")
                        x,y = srCurve.exterior.xy
                        ax.plot(x,y,linewidth=2,color='r',linestyle=":")

                        x,y = cookieCut.coords.xy
                        ax.plot(x,y,linewidth=2,color='b',linestyle="-")
                        fig.savefig(f"debug_nonpolycut_{debugCounter}.pdf")

                debugCounter +=1
    if args.debug:
        for typeOfCurve in listOfBestCurves:
            print("nCurves to Sum for %s = %d"%(typeOfCurve,len(listOfBestCurves[typeOfCurve])))
            ax.cla()
            for poly in listOfBestCurves[typeOfCurve]:
                x,y = poly.exterior.xy
                ax.plot(x,y,linewidth=2,color=np.random.rand(3,),linestyle=":")
            x,y = sumOfExpecteds.exterior.xy
            ax.plot(x,y,linewidth=1,color='k')
            fig.savefig("debug_polyToSum_%s.pdf"%(typeOfCurve))

    ax.cla()

    for iIEC,IEC in enumerate(allIsoExpectedContours):
        convertArraysToTGraph(IEC[:,0],IEC[:,1]).Write("isoExpectedContour_%d"%iIEC)
        ax.plot(IEC[:,0],IEC[:,1],alpha=0.5,linewidth=0.1)

    print (">>> Creating summed curves and writing to output file")
    summedCurves = {}

    expCurve = ()
    for typeOfCurve in listOfBestCurves:
        print(">>> Adding together: %s"%typeOfCurve)

        summedCurves[typeOfCurve]   = cascaded_union( listOfBestCurves[typeOfCurve] )

        # Try various things to get one good polygon
        if type(summedCurves[typeOfCurve])==MultiPolygon:
            # Try to join multipoly by expanding everything a tinybit
            summedCurves[typeOfCurve]   = cascaded_union([c.buffer(1e-1) for c in listOfBestCurves[typeOfCurve]])

            # For some reason polygons inside another don't always behave well
            #if type(summedCurves[typeOfCurve])==MultiPolygon:
            #
            #    for poly in listOfBestCurves[typeOfCurve]:
            #        # if poly is contained in another, remove it
            #        for polyTest in listOfBestCurves[typeOfCurve]:
            #            if poly==polyTest: continue
            #            if polyTest.contains(poly):
            #                print("removed poly")
            #                listOfBestCurves[typeOfCurve].remove(poly)
            #
            #    summedCurves[typeOfCurve]   = cascaded_union( listOfBestCurves[typeOfCurve] )

            # Last resort, convex hull
            if type(summedCurves[typeOfCurve])==MultiPolygon and not ("Exp_u1s" == typeOfCurve or "Exp_d1s" == typeOfCurve):
                print(">>> Warning, I've made a convex_hull of %s.  It won't have all of the expected inward wiggles."%(typeOfCurve))
                summedCurves[typeOfCurve] = summedCurves[typeOfCurve].convex_hull

        if "Exp_u1s" == typeOfCurve or "Exp_d1s" == typeOfCurve:
            continue

        x,y =  summedCurves[typeOfCurve].exterior.coords.xy
        ax.plot(x,y,alpha=0.5)

        if "Exp" == typeOfCurve:
            expCurve = (x,y)

        convertArraysToTGraph(x,y).Write("%s"%typeOfCurve)

    print (">>> Creating Exp +/-1 sigma band and writing to output file")

    if "Exp_u1s" in summedCurves and "Exp_d1s" in summedCurves:

        summedCurves["Exp_u1s"] = summedCurves["Exp_u1s"].buffer(1e-1) # this fixed some bad behavior where the band was the whole area
        band = summedCurves["Exp_d1s"].difference(summedCurves["Exp_u1s"])

        if type(band)==MultiPolygon:
            band = sorted(band, key=lambda x: x.area, reverse=True)
            for i,thing in enumerate(band):
                x,y = thing.exterior.coords.xy
                ax.plot(x,y,alpha=0.9)
                convertArraysToTGraph(x,y).Write("ExpectedBand_%d"%i)
        else:
            x,y = band.exterior.coords.xy
            ax.plot(x,y,alpha=0.9)

            convertArraysToTGraph(x,y).Write("ExpectedBand")

    # Upper Limits (little grey numbers)
    makeBestValueGraph(bestRegions, dict_UL).Write("upperLimit_gr")

    if args.debug:
        print (">>> Saving debugging plot: debug.pdf")
        fig.savefig("debug.pdf")

    outputFile.Write()
    outputFile.Close()

    for name in listOfInputFiles:
        if name not in allIsoExpectedContours_regionName:
            print("Warning! Region " + name + " was not used.  For best results, remove it from the input files and run again.")

    return (bestRegions, expCurve)

def tGraphToPolygon(myGraph, pinPoint=None):
    x,y = convertTGraphToArrays(myGraph)
    if pinPoint==None:
        x = np.append(x,[min(x)])
        y = np.append(y,[min(y)])
    if not Polygon(list(zip(x,y))).is_valid:
        x,y = x[:-1],y[:-1]

    return Polygon(list(zip(x,y)))

def convertTGraphToArrays(mygraph):
    size = mygraph.GetN()
    x1buffer = mygraph.GetX()
    y1buffer = mygraph.GetY()

    x1 = np.array([ x1buffer[i]  for i in range(size) ])
    y1 = np.array([ y1buffer[i]  for i in range(size) ])

    return (x1,y1)

def convertArraysToTGraph(x,y):
    mygraph = ROOT.TGraph(len(x))
    for iPoint,(x,y) in enumerate(zip(x,y)):
        mygraph.SetPoint(iPoint,x,y)
    return mygraph

def convertTGraph2DToArrays(mygraph):
    size = mygraph.GetN()
    x1buffer = mygraph.GetX()
    y1buffer = mygraph.GetY()
    z1buffer = mygraph.GetZ()

    x1 = np.array([ x1buffer[i]  for i in range(size) ])
    y1 = np.array([ y1buffer[i]  for i in range(size) ])
    z1 = np.array([ z1buffer[i]  for i in range(size) ])

    return (x1,y1,z1)

def convertArraysToTGraph2D(x,y,z):
    mygraph = ROOT.TGraph2D(len(x))
    for iPoint,(x,y,z) in enumerate(zip(x,y,z)):
        mygraph.SetPoint(iPoint,x,y,z)
    return mygraph

def doPolygonsTouch(poly1, poly2):

    return len(set(poly1.exterior.coords) & set(poly2.exterior.coords) ) > 2

def plotBestRegion(bestRegions, expCurve):

    fout = ROOT.TFile("usedExpectedContours.root","recreate")

    fig, ax = plt.subplots()

    colors = ['b','g','r','c','m','y','k']
    for i in range(10):
        colors.append(np.random.rand(3,))

    # Plot pieces
    for reg in list(bestRegions.keys()):
        for curve in bestRegions[reg]:

            try:
                x,y = curve.exterior.xy
                c = colors[ (list(bestRegions.keys())).index(reg) ]
                ax.fill(x,y,alpha=0.3,linewidth=0.2,color=c)

                tg = convertArraysToTGraph(x,y)
                name = reg.replace(".root","")+"_"+str(i)
                tg.SetNameTitle(name, name)
                tg.Write()
            except:
                print(">>> Skipped a subRegion.")
                pass

    # Plot combined expected
    ax.plot(expCurve[0],expCurve[1],linewidth=1,color='k')
    tg = convertArraysToTGraph(expCurve[0],expCurve[1])
    tg.SetNameTitle("comb_expected","comb_expected")
    tg.Write()

    # Make legend
    patches=[]
    for i,reg in enumerate(bestRegions.keys()):
        patches.append(mpatches.Patch(color=colors[i], label=reg))

    patches.append(Line2D([0], [0], color='k', lw=1, label="Combined Expected"))

    x1,x2,y1,y2 = plt.axis()
    plt.axis((x1,x2,y1,1.7*y2))

    ax.legend(handles=patches, loc=1)

    print (">>> Saving used expected contours: usedExpectedContours.pdf")
    fig.savefig("usedExpectedContours.pdf")
    fout.Close()

def makeBestValueGraph(bestRegions, inputGraphDict):

    outpoints = []

    # Gather points
    valuesDict = {fname:[] for fname in inputGraphDict.keys()}
    points = []
    for fname in inputGraphDict.keys():
        (x,y,z) = convertTGraph2DToArrays(inputGraphDict[fname])
        points+=zip(x,y)
        valuesDict[fname]+= zip(x,y,z)

    # unique set of points to loop over
    #print(points)
    points = list(set(points))

    # Loop over points, get best region, retrieve UL, and save.
    for p in points:
        x = p[0]
        y = p[1]
        z = -1
        test_point = Point( x, y )

        # Loop over best regions
        for reg in bestRegions:
            for curve in bestRegions[reg]:

                # Buffer solves issues when point is on the boundary
                if curve.buffer(1e-10).contains(test_point):

                    for v in valuesDict[reg]:
                        x_test = v[0]
                        y_test = v[1]
                        z_test = v[2]

                        if x == x_test and y == y_test:
                            z = z_test
                            outpoints.append( (x,y,z) )
                            break
                    break

    x_out = [o[0] for o in outpoints]
    y_out = [o[1] for o in outpoints]
    z_out = [o[2] for o in outpoints]

    return convertArraysToTGraph2D(x_out, y_out, z_out)

if __name__ == "__main__":

    bestRegions, expCurve = main()

    if args.plotUsedContours:
        plotBestRegion(bestRegions, expCurve)


import numpy as np
import torch
from torch.nn.modules.utils import _pair

def mask_target(pos_proposals_list, pos_assigned_gt_inds_list, gt_masks_list,
                cfg):
    """Compute mask target for positive proposals in multiple images.

    Args:
        pos_proposals_list (list[Tensor]): Positive proposals in multiple
            images.
        pos_assigned_gt_inds_list (list[Tensor]): Assigned GT indices for each
            positive proposals.
        gt_masks_list (list[:obj:`BaseInstanceMasks`]): Ground truth masks of
            each image.
        cfg (dict): Config dict that specifies the mask size.

    Returns:
        list[Tensor]: Mask target of each image.
    """
    cfg_list = [cfg for _ in range(len(pos_proposals_list))]
    mask_targets = map(mask_target_single, pos_proposals_list,
                       pos_assigned_gt_inds_list, gt_masks_list, cfg_list)
    mask_targets = list(mask_targets)
    if len(mask_targets) > 0:
        mask_targets = torch.cat(mask_targets)
    return mask_targets

def mask_target_single(pos_proposals, pos_assigned_gt_inds, gt_masks, cfg):
    """Compute mask target for each positive proposal in the image.

    Args:
        pos_proposals (Tensor): Positive proposals.
        pos_assigned_gt_inds (Tensor): Assigned GT inds of positive proposals.
        gt_masks (:obj:`BaseInstanceMasks`): GT masks in the format of Bitmap
            or Polygon.
        cfg (dict): Config dict that indicate the mask size.

    Returns:
        Tensor: Mask target of each positive proposals in the image.
    """
    device = pos_proposals.device
    mask_size = _pair(cfg.mask_size)
    num_pos = pos_proposals.size(0)
    if num_pos > 0:
        proposals_np = pos_proposals.cpu().numpy()
        # print("before")
        # print(proposals_np)

        maxh, maxw = gt_masks.height, gt_masks.width
        proposals_np[:, [0, 2]] = np.clip(proposals_np[:, [0, 2]], 0, maxw)
        proposals_np[:, [1, 3]] = np.clip(proposals_np[:, [1, 3]], 0, maxh)
        # print("after")
        # print(proposals_np)
        pos_assigned_gt_inds = pos_assigned_gt_inds.cpu().numpy()
        # print(pos_assigned_gt_inds)
        # print(gt_masks)
        mask_targets = gt_masks.crop_and_resize(
            proposals_np, mask_size, device=device,
            inds=pos_assigned_gt_inds).to_ndarray()

        mask_targets = torch.from_numpy(mask_targets).float().to(device)
    else:
        mask_targets = pos_proposals.new_zeros((0, ) + mask_size)

    return mask_targets


import time
import cv2 as cv
import numpy as np
from libs.centroid_object_tracker import CentroidTracker
from scipy.spatial import distance as dist

class Distancing:

    def __init__(self, config):
        self.config = config
        self.ui = None
        self.detector = None
        self.device = self.config.get_section_dict('Detector')['Device']
        self.running_video = False
        self.tracker = CentroidTracker(
            maxDisappeared=int(self.config.get_section_dict("PostProcessor")["MaxTrackFrame"]))
        if self.device == 'Jetson':
            from libs.detectors.jetson.Detector import Detector
            self.detector = Detector(self.config)
        elif self.device == 'EdgeTPU':
            from libs.detectors.edgetpu.Detector import Detector
            self.detector = Detector(self.config)
        elif self.device == 'Dummy':
            self.detector = None

        self.image_size = [int(i) for i in self.config.get_section_dict('Detector')['ImageSize'].split(',')]

        if self.device != 'Dummy':
            print('Device is: ', self.device)
            print('Detector is: ', self.detector.name)
            print('image size: ', self.image_size)

    def set_ui(self, ui):
        self.ui = ui

    def __process(self, cv_image):
        """
        return object_list list of  dict for each obj,
        obj["bbox"] is normalized coordinations for [x0, y0, x1, y1] of box
        """
        if self.device == 'Dummy':
            return cv_image, [], None

        # Resize input image to resolution
        resolution = [int(i) for i in self.config.get_section_dict('App')['Resolution'].split(',')]
        cv_image = cv.resize(cv_image, tuple(resolution))

        resized_image = cv.resize(cv_image, tuple(self.image_size[:2]))
        rgb_resized_image = cv.cvtColor(resized_image, cv.COLOR_BGR2RGB)
        tmp_objects_list = self.detector.inference(rgb_resized_image)
        for obj in tmp_objects_list:
            box = obj["bbox"]
            x0 = box[1]
            y0 = box[0]
            x1 = box[3]
            y1 = box[2]
            obj["centroid"] = [(x0 + x1) / 2, (y0 + y1) / 2, x1 - x0, y1 - y0]
            obj["bbox"] = [x0, y0, x1, y1]

        objects_list, distancings = self.calculate_distancing(tmp_objects_list)
        return cv_image, objects_list, distancings

    def process_video(self, video_uri):
        input_cap = cv.VideoCapture(video_uri)

        if (input_cap.isOpened()):
            print('opened video ', video_uri)
        else:
            print('failed to load video ', video_uri)
            return

        self.running_video = True
        while input_cap.isOpened() and self.running_video:
            _, cv_image = input_cap.read()
            cv_image, objects, distancings = self.__process(cv_image)
            self.ui.update(cv_image, objects, distancings)

        input_cap.release()
        self.running_video = False

    def process_image(self, image_path):
        cv_image = cv.imread(image_path)
        cv_image, objects, distancings = self.__process(cv_image)
        self.ui.update(cv_image, objects, distancings)

    def calculate_distancing(self, objects_list):
        """
        this function post-process the raw boxes of object detector and calculate a distance matrix
        for detected bounding boxes.
        post processing is consist of:
        1. omitting large boxes by filtering boxes which are biger than the 1/4 of the size the image.
        2. omitting duplicated boxes by applying an auxilary non-maximum-suppression.
        3. apply a simple object tracker to make the detection more robust.

        params:
        object_list: a list of dictionaries. each dictionary has attributes of a detected object such as
        "id", "centroid" (a tuple of the normalized centroid coordinates (cx,cy,w,h) of the box) and "bbox" (a tuple
        of the normalized (xmin,ymin,xmax,ymax) coordinate of the box)

        returns:
        object_list: the post processed version of the input
        distances: a NxN ndarray which i,j element is distance between i-th and l-th bounding box

        """
        new_objects_list = self.ignore_large_boxes(objects_list)
        new_objects_list = self.non_max_suppression_fast(new_objects_list,
                                                         float(self.config.get_section_dict("PostProcessor")[
                                                                   "NMSThreshold"]))
        tracked_boxes = self.tracker.update(new_objects_list)
        new_objects_list = [tracked_boxes[i] for i in tracked_boxes.keys()]
        for i, item in enumerate(new_objects_list):
            item["id"] = item["id"].split("-")[0] + "-" + str(i)

        centroids = np.array( [obj["centroid"] for obj in new_objects_list] )
        distances = dist.cdist(centroids, centroids)
        return new_objects_list, distances

    @staticmethod
    def ignore_large_boxes(object_list):

        """
        filtering boxes which are biger than the 1/4 of the size the image
        params:
            object_list: a list of dictionaries. each dictionary has attributes of a detected object such as
            "id", "centroid" (a tuple of the normalized centroid coordinates (cx,cy,w,h) of the box) and "bbox" (a tuple
            of the normalized (xmin,ymin,xmax,ymax) coordinate of the box)
        returns:
        object_list: input object list without large boxes
        """
        large_boxes = []
        for i in range(len(object_list)):
            if (object_list[i]["centroid"][2] * object_list[i]["centroid"][3]) > 0.25:
                large_boxes.append(i)
        updated_object_list = [j for i, j in enumerate(object_list) if i not in large_boxes]
        return updated_object_list

    @staticmethod
    def non_max_suppression_fast(object_list, overlapThresh):

        """
        omitting duplicated boxes by applying an auxilary non-maximum-suppression.
        params:
        object_list: a list of dictionaries. each dictionary has attributes of a detected object such
        "id", "centroid" (a tuple of the normalized centroid coordinates (cx,cy,w,h) of the box) and "bbox" (a tuple
        of the normalized (xmin,ymin,xmax,ymax) coordinate of the box)

        overlapThresh: threshold of minimum IoU of to detect two box as duplicated.

        returns:
        object_list: input object list without duplicated boxes
        """
        # if there are no boxes, return an empty list
        boxes = np.array([item["centroid"] for item in object_list])
        corners = np.array([item["bbox"] for item in object_list])
        if len(boxes) == 0:
            return []
        if boxes.dtype.kind == "i":
            boxes = boxes.astype("float")
        # initialize the list of picked indexes
        pick = []
        cy = boxes[:, 1]
        cx = boxes[:, 0]
        h = boxes[:, 3]
        w = boxes[:, 2]
        x1 = corners[:, 0]
        x2 = corners[:, 2]
        y1 = corners[:, 1]
        y2 = corners[:, 3]
        area = (h + 1) * (w + 1)
        idxs = np.argsort(cy + (h / 2))
        while len(idxs) > 0:
            last = len(idxs) - 1
            i = idxs[last]
            pick.append(i)
            xx1 = np.maximum(x1[i], x1[idxs[:last]])
            yy1 = np.maximum(y1[i], y1[idxs[:last]])
            xx2 = np.minimum(x2[i], x2[idxs[:last]])
            yy2 = np.minimum(y2[i], y2[idxs[:last]])

            w = np.maximum(0, xx2 - xx1 + 1)
            h = np.maximum(0, yy2 - yy1 + 1)
            # compute the ratio of overlap
            overlap = (w * h) / area[idxs[:last]]
            # delete all indexes from the index list that have
            idxs = np.delete(idxs, np.concatenate(([last],
                                                  np.where(overlap > overlapThresh)[0])))
        updated_object_list = [j for i, j in enumerate(object_list) if i in pick]
        return updated_object_list


import os
import numpy as np
import pandas as pd

#             0, 1,  2,  3,    4,    5,    6,  7,     8,         9,
orig_seg = [( 1, 1,  9,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg1"),
            ( 2, 1, 10,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg2"),
            ( 3, 1,  9,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg3"),
            ( 4, 1, 11,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg4"),
            ( 5, 1, 15,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg5"),
            ( 6, 1, 13,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg6"),
            ( 7, 1, 12,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg7"),
            ( 8, 1, 14,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg8"),
            ( 9, 1, 10,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg9"),
            (10, 1, 11,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg10"),
            (11, 1, 12,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg11"),
            (12, 1, 13,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg12"),
            (13, 1, 14,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg13"),
            (14, 1, 15,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg14"),
            (15, 1,  0,  0, 0.00,  0.0,  0.0,  0,  0.04, "origSeg15")]

orig_rch = [(1, 39, 15,  1,  1),
            (1, 39, 16,  1,  2),
            (1, 39, 17,  1,  3),
            (1, 38, 17,  1,  4),
            (1, 38, 18,  1,  5),
            (1, 38, 19,  1,  6),
            (1, 38, 20,  1,  7),
            (1, 38, 21,  1,  8),
            (1, 38, 22,  1,  9),
            (1, 38, 23,  1, 10),
            (1, 37, 24,  1, 11),
            (1, 37, 25,  1, 12),
            (1, 36, 25,  1, 13),
            (1, 36, 26,  1, 14),
            (1, 35, 26,  1, 15),
            (1, 35, 27,  1, 16),
            (1, 36, 28,  1, 17),
            (1, 36, 29,  1, 18),
            (1, 36, 30,  1, 19),
            (1, 36, 31,  1, 20),
            (1, 36, 32,  1, 21),
            (1, 35, 32,  1, 22),
            (1, 35, 33,  1, 23),
            (1, 34, 33,  1, 24),
            (1, 34, 34,  1, 25),
            (1, 33, 34,  1, 26),
            (1, 33, 35,  1, 27),
            (1, 32, 35,  1, 28),
            (1, 32, 36,  1, 29),
            (1, 32, 37,  1, 30),
            (1, 53, 37,  2,  1),
            (1, 52, 37,  2,  2),
            (1, 52, 38,  2,  3),
            (1, 51, 38,  2,  4),
            (1, 50, 38,  2,  5),
            (1, 50, 39,  2,  6),
            (1, 49, 39,  2,  7),
            (1, 48, 39,  2,  8),
            (1, 47, 39,  2,  9),
            (1, 47, 40,  2, 10),
            (1, 46, 40,  2, 11),
            (1, 45, 40,  2, 12),
            (1, 45, 39,  2, 13),
            (1, 44, 39,  2, 14),
            (1, 43, 38,  2, 15),
            (1, 42, 38,  2, 16),
            (1, 41, 38,  2, 17),
            (1, 40, 38,  2, 18),
            (1, 40, 39,  2, 19),
            (1, 39, 39,  2, 20),
            (1, 38, 39,  2, 21),
            (1, 37, 40,  2, 22),
            (1, 36, 40,  2, 23),
            (1, 36, 41,  2, 24),
            (1, 35, 41,  2, 25),
            (1, 34, 41,  2, 26),
            (1, 33, 41,  2, 27),
            (1, 32, 41,  2, 28),
            (1, 31, 42,  2, 29),
            (1, 31, 33,  3,  1),
            (1, 31, 34,  3,  2),
            (1, 31, 35,  3,  3),
            (1, 31, 36,  3,  4),
            (1, 31, 37,  3,  5),
            (1, 48, 48,  4,  1),
            (1, 47, 48,  4,  2),
            (1, 46, 48,  4,  3),
            (1, 46, 47,  4,  4),
            (1, 45, 47,  4,  5),
            (1, 44, 47,  4,  6),
            (1, 43, 47,  4,  7),
            (1, 42, 47,  4,  8),
            (1, 41, 47,  4,  9),
            (1, 41, 48,  4, 10),
            (1, 40, 48,  4, 11),
            (1, 39, 48,  4, 12),
            (1, 38, 47,  4, 13),
            (1, 37, 47,  4, 14),
            (1, 36, 48,  4, 15),
            (1, 35, 48,  4, 16),
            (1, 35, 49,  4, 17),
            (1, 34, 49,  4, 18),
            (1, 34, 50,  4, 19),
            (1, 33, 50,  4, 20),
            (1, 55, 72,  5,  1),
            (1, 54, 72,  5,  2),
            (1, 53, 72,  5,  3),
            (1, 52, 72,  5,  4),
            (1, 51, 72,  5,  5),
            (1, 50, 73,  5,  6),
            (1, 49, 73,  5,  7),
            (1, 48, 73,  5,  8),
            (1, 48, 74,  5,  9),
            (1, 47, 74,  5, 10),
            (1, 46, 75,  5, 11),
            (1, 45, 75,  5, 12),
            (1, 45, 76,  5, 13),
            (1, 44, 76,  5, 14),
            (1, 45, 62,  6,  1),
            (1, 44, 62,  6,  2),
            (1, 43, 62,  6,  3),
            (1, 43, 63,  6,  4),
            (1, 42, 63,  6,  5),
            (1, 41, 63,  6,  6),
            (1, 40, 63,  6,  7),
            (1, 24, 55,  7,  1),
            (1, 25, 55,  7,  2),
            (1, 25, 56,  7,  3),
            (1, 26, 56,  7,  4),
            (1, 27, 56,  7,  5),
            (1, 28, 57,  7,  6),
            (1, 29, 57,  7,  7),
            (1, 30, 57,  7,  8),
            (1, 31, 57,  7,  9),
            (1, 32, 57,  7, 10),
            (1, 33, 57,  7, 11),
            (1, 33, 58,  7, 12),
            (1, 34, 58,  7, 13),
            (1, 34, 59,  7, 14),
            (1, 35, 59,  7, 15),
            (1, 36, 59,  7, 16),
            (1, 37, 60,  7, 17),
            (1, 23, 71,  8,  1),
            (1, 24, 71,  8,  2),
            (1, 25, 71,  8,  3),
            (1, 26, 71,  8,  4),
            (1, 27, 72,  8,  5),
            (1, 27, 73,  8,  6),
            (1, 28, 73,  8,  7),
            (1, 29, 73,  8,  8),
            (1, 30, 73,  8,  9),
            (1, 31, 73,  8, 10),
            (1, 32, 73,  8, 11),
            (1, 33, 73,  8, 12),
            (1, 34, 73,  8, 13),
            (1, 34, 74,  8, 14),
            (1, 35, 74,  8, 15),
            (1, 36, 74,  8, 16),
            (1, 36, 73,  8, 17),
            (1, 37, 73,  8, 18),
            (1, 38, 72,  8, 19),
            (1, 39, 72,  8, 20),
            (1, 40, 72,  8, 21),
            (1, 41, 72,  8, 22),
            (1, 42, 72,  8, 23),
            (1, 42, 73,  8, 24),
            (1, 31, 38,  9,  1),
            (1, 31, 39,  9,  2),
            (1, 31, 40,  9,  3),
            (1, 31, 41,  9,  4),
            (1, 31, 42,  9,  5),
            (1, 30, 42,  9,  6),
            (1, 30, 43, 10,  1),
            (1, 30, 44, 10,  2),
            (1, 29, 44, 10,  3),
            (1, 29, 45, 10,  4),
            (1, 29, 46, 10,  5),
            (1, 29, 47, 10,  6),
            (1, 30, 47, 10,  7),
            (1, 30, 48, 10,  8),
            (1, 31, 49, 10,  9),
            (1, 32, 50, 10, 10),
            (1, 32, 51, 11,  1),
            (1, 33, 52, 11,  2),
            (1, 33, 53, 11,  3),
            (1, 34, 53, 11,  4),
            (1, 34, 54, 11,  5),
            (1, 35, 54, 11,  6),
            (1, 35, 55, 11,  7),
            (1, 35, 56, 11,  8),
            (1, 36, 57, 11,  9),
            (1, 36, 58, 11, 10),
            (1, 36, 59, 11, 11),
            (1, 37, 59, 11, 12),
            (1, 37, 60, 11, 13),
            (1, 38, 60, 11, 14),
            (1, 38, 61, 12,  1),
            (1, 38, 62, 12,  2),
            (1, 38, 63, 12,  3),
            (1, 39, 63, 12,  4),
            (1, 39, 64, 13,  1),
            (1, 39, 65, 13,  2),
            (1, 40, 65, 13,  3),
            (1, 40, 66, 13,  4),
            (1, 40, 67, 13,  5),
            (1, 40, 68, 13,  6),
            (1, 41, 69, 13,  7),
            (1, 41, 70, 13,  8),
            (1, 42, 71, 13,  9),
            (1, 42, 72, 13, 10),
            (1, 42, 73, 13, 11),
            (1, 42, 73, 14,  1),
            (1, 43, 73, 14,  2),
            (1, 43, 74, 14,  3),
            (1, 43, 75, 14,  4),
            (1, 44, 75, 14,  5),
            (1, 44, 76, 14,  6),
            (1, 44, 77, 15,  1),
            (1, 44, 78, 15,  2),
            (1, 44, 79, 15,  3),
            (1, 45, 79, 15,  4)
]

def gen_mf6_sfr_connections():
    conns = []
    for i in np.arange(0, len(orig_seg)):
        tup = orig_seg[i]
        segid = tup[0]
        ioutseg = tup[2]
        iupseg = tup[3]

        # Get all reaches associated with segment
        # Find an element in a list of tuples
        allrchs = [item for item in orig_rch if item[3] == segid]

        # Loop through allrchs and generate list of connections
        for rchx in allrchs:
            # rchx will be a tuple
            upconn = []
            dnconn = []

            if rchx[4] == 1:      # checks if first rch of segment
                # Collect all segs that dump to the current one (there may not be any)
                dumpersegs = [item for item in orig_seg if item[2] == segid]
                # For every seg that outflows to current, set last reach of it as
                # an upstream connection
                for dumper in dumpersegs:
                    dumper_seg_id = dumper[0]
                    rch_cnt = len([item for item in orig_rch if item[3] == dumper_seg_id])
                    lastrch = [item for item in orig_rch if item[3] == dumper_seg_id and item[4] == rch_cnt]
                    idx = orig_rch.index(lastrch[0])
                    upconn.append(int(idx))

                # Current reach is the most upstream reach for current segment
                if iupseg == 0:
                    pass
                elif iupseg > 0:  # Lake connections, signified with negative numbers, aren't handled here
                    iupseg_rchs = [item for item in orig_rch if item[3] == iupseg]
                    # Get the index of the last reach of the segement that was the upstream segment in the orig sfr file
                    idx = orig_rch.index(iupseg_rchs[len(iupseg_rchs)-1])   # From: https://stackoverflow.com/questions/20239312/find-an-exact-tuple-match-in-a-list-of-tuples-and-return-its-index
                    upconn.append(idx)

                # Even if the first reach of a segement, it will have an outlet,
                # either the next reach in the segment, or first reach of outseg,
                # which should be taken care of below
                if len(allrchs) > 1:
                    idx = orig_rch.index(rchx)
                    # adjust idx for 0-based and increment to next item in list
                    dnconn.append(int(idx + 1) * -1)

            elif rchx[4] > 1 and not rchx[4] == len(allrchs):
                # Current reach is 'interior' on the original segment and therefore
                # should only have 1 upstream & 1 downstream segement
                idx = orig_rch.index(rchx)
                # B/c 0-based, idx will already be incremented by -1
                upconn.append(int(idx - 1))
                # adjust idx for 0-based and increment to next item in list
                dnconn.append(int(idx + 1) * -1)  # all downstream connections are negative in MF6

            if rchx[4] == len(allrchs):
                # If the last reach in a multi-reach segment, always need to account
                # for the reach immediately upstream (single reach segs dealt with
                # above), unless of course we're dealing with a single reach segment
                # like in the case of a spillway from a lake
                if len(allrchs) != 1:
                    idx = orig_rch.index(rchx)
                    # B/c 0-based, idx will already be incremented by -1
                    upconn.append(int(idx - 1))

                # Current reach is last reach in segment and may have multiple
                # downstream connections, particular when dealing with diversions.
                if ioutseg == 0:
                    pass
                elif ioutseg > 0:       # Lake connections, signified with negative numbers, aren't handled here
                    idnseg_rchs = [item for item in orig_rch if item[3] == ioutseg and item[4] == 1]
                    idx = orig_rch.index(idnseg_rchs[0])
                    # adjust idx for 0-based and increment to next item in list
                    dnconn.append(int(idx) * -1)

                # In addition to ioutseg, look for all segments that may have the
                # current segment as their iupseg
                possible_divs = [item for item in orig_seg if item[3] == rchx[3]]
                for segx in possible_divs:
                    # Next, peel out all first reach for any segments listed in possible_divs
                    first_rchs = [item for item in orig_rch if item[3] == segx[0] and item[4] == 1]
                    for firstx in first_rchs:
                        idx = orig_rch.index(firstx)
                        # adjust idx for 0-based and increment to next item in list
                        dnconn.append(int(idx) * -1)

            # Append the collection of upconn & dnconn as an entry in a list
            idx = orig_rch.index(rchx)
            # Adjust current index for 0-based
            conns.append([idx] + upconn + dnconn)

    return conns


import numpy as np
import os
import cv2
import time

data_path = "C:/Users/mabin/Dropbox/SKaggle/generative-dog-images/all-dogs/"
# data_path = "C:/Users/mabin/Dropbox/SKaggle/generative-dog-images/resized/"
data_name = os.listdir(data_path)
data_size = []
for i in data_name:
    size = os.stat(data_path + i).st_size
    data_size.append(size)

maxx = np.max(data_size)
maxxw = np.where(maxx == data_size)
minn = np.min(data_size)
tmpimg = cv2.imread(data_path + data_name[int(maxxw[0])])
img_height, img_width, img_channels = tmpimg.shape[:3]
if img_height >= img_width:
    img_size = img_height
else:
    img_size = img_width
img_size = 32
s1 = time.time()
for i in data_name:
    tmpimg = cv2.imread(data_path + i)
    tmpimg = cv2.resize(tmpimg, (img_size, img_size))
    cv2.imwrite(data_path + "../resized32/" + i, tmpimg)
    print(np.array(tmpimg).size)
print(time.time() - s1)


import sys
sys.path.append("../../")
import unittest
import paddle
import numpy as np
from paddleslim import UnstructuredPruner
from paddle.vision.models import mobilenet_v1
import paddle.vision.transforms as T
import paddle.fluid as fluid
from paddle.static import InputSpec as Input
import paddle.nn.functional as F

class TestStaticMasks(unittest.TestCase):
    def __init__(self, *args, **kwargs):
        super(TestStaticMasks, self).__init__(*args, **kwargs)
        paddle.disable_static()
        transform = T.Compose([T.Transpose(), T.Normalize([127.5], [127.5])])
        self.train_dataset = paddle.vision.datasets.MNIST(
            mode="train", backend="cv2", transform=transform)
        self.train_loader = paddle.io.DataLoader(
            self.train_dataset,
            places=paddle.set_device('cpu'),
            return_list=True)

        def _reader():
            for data in self.val_dataset:
                yield data

        self.val_reader = _reader

    def _update_masks(self, pruner, t):
        for name, sub_layer in pruner.model.named_sublayers():
            for param in sub_layer.parameters(include_sublayers=False):
                mask = pruner.masks.get(param.name)
                bool_tmp = (paddle.abs(param) < t)
                paddle.assign(bool_tmp, output=mask)

    def runTest(self):
        with fluid.unique_name.guard():
            net = paddle.vision.models.LeNet()
            optimizer = paddle.optimizer.Adam(
                learning_rate=0.001, parameters=net.parameters())
            inputs = [Input([None, 1, 28, 28], 'float32', name='image')]
            labels = [Input([None, 1], 'int64', name='label')]
        pruner = UnstructuredPruner(net, mode='ratio', ratio=0.55)
        net.train()
        self._update_masks(pruner, 0.0)
        pruner.update_params()
        self._update_masks(pruner, 1.0)
        pruner.set_static_masks()
        sparsity_0 = UnstructuredPruner.total_sparse(net)
        for i, data in enumerate(self.train_loader):
            x_data = data[0]
            y_data = paddle.to_tensor(data[1])
            logits = net(x_data)
            loss = F.cross_entropy(logits, y_data)
            loss.backward()
            optimizer.step()
            optimizer.clear_grad()
            if i == 10: break
        sparsity_1 = UnstructuredPruner.total_sparse(net)
        pruner.update_params()
        sparsity_2 = UnstructuredPruner.total_sparse(net)
        print(sparsity_0, sparsity_1, sparsity_2)
        self.assertEqual(sparsity_0, 1.0)
        self.assertEqual(sparsity_2, 1.0)
        self.assertLess(sparsity_1, 1.0)

if __name__ == "__main__":
    unittest.main()