import multiprocessing

from numpy.lib.type_check import common_type
from pysam import AlignmentFile
from scipy.sparse import hstack, vstack, csr_matrix, csc_matrix, coo_matrix, save_npz, load_npz
import numpy as np
from itertools import zip_longest
from Bio import SeqIO  # __init__ kicks in
import sys
# from numba import jit
try:
    from .tools import pathing
    from .cython_numpy.cython_np_array import to_array
except:
    # For debugging speed issues with pyinstrument
    # COMMAND: pyinstrument --show-all --outfile pileup.log pileup.py
    from DeepVCF import pathing, to_array

def get_ref_seq(reference_file: str, format: str = 'fasta'):
    reference_file = pathing(reference_file)
    seqrecords = list(SeqIO.parse(reference_file, format=format))
    if len(seqrecords) > 1:
        exit('ERROR :: Reference fasta must be a single complete sequence')
    return seqrecords[0].seq

def get_alignments(alignment_file, threads=multiprocessing.cpu_count()):
    return AlignmentFile(pathing(alignment_file), threads=threads)

class Pileup:
    """ """
    basepile = {
        'A':  [1, 0, 0, 0],
        'C':  [0, 1, 0, 0],
        'G':  [0, 0, 1, 0],
        'T':  [0, 0, 0, 1],
    }
    baseindex = {
        'A':  -4,
        'C':  -3,
        'G':  -2,
        'T':  -1,
    }
    base_index = {
        'A':  0,
        'C':  1,
        'G':  2,
        'T':  3,
    }
    base_insert_index = {
        'A':  4,
        'C':  5,
        'G':  6,
        'T':  7,
    }
    def __repr__():
        return 'Pileup 2D matrix - REF/deletion + ALT + ALT/insertion'

    def __init__(self,
                 reference_file,
                 alignment_file,
                 mimimum_alignment_coverage: float = .75,
                 save_pileup_to_destination: str = None,
                 use_saved_pileup: str = None,
                 minimum_coverage: int = 30,
                 heterozygous_threshold: float = .25,
                 minimum_variant_radius: int = 15,
                 overwrite_variant_calls: bool = False,
                 **kwargs,) -> None:
        # helper variables
        self.contig_start= None
        self.contig_end = None
        self.minimum_coverage = minimum_coverage
        self.heterozygous_threshold = heterozygous_threshold  # minimum difference a secondary alt needs to be for a variant to be called.
        self.minimum_variant_radius = minimum_variant_radius  # distance variants are allowed to be without being classified as complex
        self.mapping_coverages=[]
        # aligment parameters
        self.mimimum_alignment_coverage = mimimum_alignment_coverage
        # core data
        self.reference_file = pathing(reference_file)
        self.alignment_file = pathing(alignment_file)
        self.ref_seq = get_ref_seq(reference_file)
        # to speed-up debugging
        if use_saved_pileup:
            try:
                print('=== Using Input Pileup ===')
                self.pileup = str(pathing(use_saved_pileup))
                # open it up if its a npz file
                if isinstance(self.pileup, str):
                    self.pileup = np.load(self.pileup)
                for alignment in get_alignments(self.alignment_file):
                    self.contig_start = alignment.reference_start
                    break
                if overwrite_variant_calls:
                    print('=== Overwriting Variant Calls === ')
                    self.overwrite_variant_calls()
            except:
                print('=== Input Pileup Failed; Building Pileup From Scratch ===')
                self.pileup = self.get_contig_pileup()
        else:
            print('=== Building Pileup From Scratch ===')
            self.pileup = self.get_contig_pileup()
        if save_pileup_to_destination:
            self.save_pileup(save_pileup_to_destination)
        print('=== Pilup Complete ===')

    @staticmethod
    def _offset_sum(top, bot, left, right):
        top_len = len(top)
        bot_len = len(bot)
        if top.size == 0:
            return bot
        if right == 0 and left == 0:
            top += bot
        elif top_len <= left:
            padding = left - top_len
            return np.concatenate((np.pad(top, (0, padding), 'constant', constant_values=(0, 0)), bot))
        elif right > 0:
            top[left:] += bot[:top_len-left]
            top = np.concatenate((top, bot[-right:]))
        else:
            top[left:bot_len+left] += bot[:]
        return top

    def _get_contig_pileup_fragment(self, alignment):
        align_pairs = alignment.get_aligned_pairs()[alignment.query_alignment_start:alignment.query_alignment_end]
        # contig_fragment = [0]*(len(align_pairs)*8)
        contig_fragment = []
        # stepper = 0
        # step = 8
        for query_pos, ref_pos in align_pairs:
            if ref_pos != None:
                last_ref_pos = ref_pos # removes odd alignment edge-cases
            if ref_pos != None and query_pos != None:
                contig_fragment.extend(self.basepile[alignment.query_sequence[query_pos]][:] + [0, 0, 0, 0])
                # contig_fragment[self.base_index[alignment.query_sequence[query_pos]] + stepper] += 1
                # stepper += step
            elif ref_pos == None and query_pos != None and contig_fragment: # BUG in query_alignment_start; it seems its not always when it starts. b
                # contig_fragment = contig_fragment[:-8]
                # contig_fragment[self.base_insert_index[alignment.query_sequence[query_pos]] + stepper] += 1
                contig_fragment[self.baseindex[alignment.query_sequence[query_pos]]] += 1
            elif ref_pos != None and query_pos == None:
                # stepper += step
                contig_fragment.extend([0, 0, 0, 0, 0, 0, 0, 0])
            else:
                pass
                # print('SHIT')
                # contig_fragment = contig_fragment[:-8]  # no match on either side
        return contig_fragment, last_ref_pos

    def _get_contig_pileup(self):
        pileup_len = 8
        alignmentfile = get_alignments(self.alignment_file)
        prev_reference_end = 0
        prev_contig_fragment = np.empty(1)
        for i, alignment in enumerate(alignmentfile):
            # if i < 210000:
            #     continue
            # if i == 10000: break
            # If alignment less that 50% aligned, it's no good and will muddy up the pileup
            skip_base, total_aln_pos = 0, 0
            for code, adv in alignment.cigartuples:
                total_aln_pos += adv
                if code == 4:
                    skip_base += adv
            mapping_coverage = 1.0 - 1.0 * skip_base / (total_aln_pos+1)
            if mapping_coverage < self.mimimum_alignment_coverage:
                continue
            self.mapping_coverages.append(mapping_coverage)
            # initial start
            if i == 0 or not self.contig_start:
                self.contig_start = alignment.reference_start
                prev_reference_start = 0
                contig_pf, prev_reference_end = self._get_contig_pileup_fragment(alignment)
                prev_reference_end -= self.contig_start
                prev_contig_fragment = to_array(contig_pf)
                # prev_contig_fragment = contig_pf
                continue
            reference_start = alignment.reference_start - self.contig_start
            contig_pf, reference_end = self._get_contig_pileup_fragment(alignment)
            reference_end -= self.contig_start
            contig_fragment = to_array(contig_pf)
            # contig_fragment = contig_pf
            # arrays need to to be the same size
            left = reference_start - prev_reference_start
            right = reference_end - prev_reference_end  # can be negative offset for additional padding on incoming alignment
            left *= pileup_len
            right *= pileup_len
            # padding arrays to be the same size
            merged_contig_fragment = self._offset_sum(prev_contig_fragment, contig_fragment, left, right)
            # merged_contig_fragment = sum(self._offset_alignments(prev_contig_fragment, contig_fragment, right, left))

            # if not np.array_equal(merged_contig_fragment_, merged_contig_fragment):
            # if merged_contig_fragment_.shape != merged_contig_fragment_.shape:
            #     print(i, merged_contig_fragment_, merged_contig_fragment)
            #     import IPython
            #     IPython.embed()
            #     sys.exit()

            # alignments are sorted so the current reference slice is next
            prev_contig_fragment = merged_contig_fragment[(reference_start - prev_reference_start)*pileup_len:]
            # previous reference start to current reference start slice for vstack
            fragment = merged_contig_fragment[:(reference_start - prev_reference_start)*pileup_len]
            # update positions
            prev_reference_end = reference_end if reference_end >= prev_reference_end else prev_reference_end
            prev_reference_start = reference_start
            try:
                if fragment.size!=0: # repeating alignments at same ref start position
                    yield fragment.reshape(-1, pileup_len)
            except:
                print(fragment, merged_contig_fragment, i)
                sys.exit()
                # yield coo_matrix(fragment.reshape(-1, pileup_len), dtype=np.dtype('h'))
                # yield coo_matrix(fragment, dtype=np.dtype('h'))
        self.contig_end = prev_reference_end + self.contig_start
        if prev_contig_fragment.size!=0:
            yield prev_contig_fragment.reshape(-1, pileup_len)
            # yield coo_matrix(prev_contig_fragment.reshape(-1, pileup_len), dtype=np.dtype('h'))
            # yield coo_matrix(prev_contig_fragment, dtype=np.dtype('h'))
        alignmentfile.close()  # todo: do i need this?

    def _meets_filter(self, total_count: int, base_count: list, ref_base: str):
        """ If read nb is different or debatable enoough to be differnet (threshold) """
        base_count = zip('ACGT', base_count)
        base_count = sorted(base_count, key = lambda x:-x[1])
        if base_count[0][0] != ref_base:
            return 1
        p0 = 1.0 *  base_count[0][1] / total_count
        p1 = 1.0 *  base_count[1][1] / total_count
        if (p0 < 1.0 - self.heterozygous_threshold and p1 > self.heterozygous_threshold): # or base_count[0][0] != ref_base:
            return 1
        return 0

    #  TODO: There is a bottleneck somewhere in here. Memory issues are most likely the issue since time is linear with matrix depth & time allocating the actual memory was solved with cython.
    def get_contig_pileup(self):
        """
        Get pileup of reference+deletions, alt and alt+insertions for dynamic creation of tensors.

        Returns:
            csr_matrix: sparce column matrix to hold the complete segment of sequences in a tiny amount of memory.
        """
        contig_pileup = np.concatenate(list(self._get_contig_pileup()), axis=0)  # vstack: top+bot
        ref_sum = contig_pileup[:, :4].sum(axis=1)
        ref_seq_seg = self.ref_seq[self.contig_start:]
        ref_piles = []
        vpos = -float('inf')
        for pos, (nt, rs, base_count) in enumerate(zip(ref_seq_seg, ref_sum, contig_pileup[:, :4])):
            pile = [0] + self.basepile.get(nt, [0, 0, 0, 0])[:]
            try:  # TODO: give partial distributions for nts useing http://www.bioinformatics.org/sms/iupac.html
                # if pos + self.contig_start + 1 == 14513747:
                #     print(pos, nt, rs, base_count)
                pile[self.base_index[nt] + 1] = rs
            except:
                ref_piles.append(pile)
                continue
            if rs >= self.minimum_coverage:
                pile[0] = self._meets_filter(rs, base_count, nt)
            #     if variant_call:
            #         if abs(pos - vpos) < self.minimum_variant_radius:
            #             ref_piles[-1][0] = 0
            #             pile[0] = 0
            #         else:
            #             pile[0] = variant_call
            #         vpos = pos
            #     else:
            #         pile[0] = 0
            # else:
            #     pile[0] = 0
            ref_piles.append(pile)
        ref_pileup = np.array(ref_piles)
        return np.concatenate((ref_pileup, contig_pileup), axis=1)  # hstack: left+right
        # return vstack(stacks, dtype=np.dtype('h'))  # for sparse matrix if needed

    def overwrite_variant_calls(self):
        ref_sum = self.pileup[:, 1:5].sum(axis=1)
        ref_seq_seg = self.ref_seq[self.contig_start:]
        ref_piles = []
        vpos = -float('inf')
        for pos, (nt, rs, base_count) in enumerate(zip(ref_seq_seg, ref_sum, self.pileup[:, 5:9])):
            variant_call = [0]
            if rs >= self.minimum_coverage:
                variant_call = [self._meets_filter(rs, base_count, nt)]
                # if variant_call[0]:
                #     if abs(pos - vpos) < self.minimum_variant_radius:
                #         ref_piles[-1][0] = 0
                #     vpos = pos
            ref_piles.append(variant_call)
        ref_pileup = np.array(ref_piles)
        self.pileup = np.concatenate((ref_pileup, self.pileup[:, 1:]), axis=1)  # hstack: left+right

    def save_pileup(self, destination):
        """
        Save scipy array to a npz file so we don't have to build the array from scratch.

        Args:
            destination ([type]): [description]
        """
        np.save(pathing(destination, new=True, overwrite=True), self.pileup)
        # save_npz(pathing(destination, new=True, overwrite=True), self.pileup)

if __name__ == '__main__':
    pileup = Pileup(
        reference_file='/home/tmsincomb/Desktop/SupplementaryDataset1/SIMULATED_SNP_CONTAINING_GENOMES/neisseria/FDAARGOS_205/FDAARGOS_205_simulated.fasta',
        alignment_file='/home/tmsincomb/Desktop/SupplementaryDataset1/SIMULATED_SNP_CONTAINING_GENOMES/neisseria/FDAARGOS_205/dwgsim-default.mapped.bam',
    )
    print(pileup.pileup[:10])
    # import IPython
    # IPython.embed()
    # print(pileup.pileup[:10])

    vcf = pathing('~/Dropbox/thesis/VariantNET/testing_data/chr21/chr21.vcf')
    fasta = pathing('~/Dropbox/thesis/VariantNET/testing_data/chr21/chr21.fa')
    bam = pathing('~/Dropbox/thesis/VariantNET/testing_data/chr21/hg38.NA12878-WashU_chr21-14069662-46411975.bam')
    bami = pathing('~/Dropbox/thesis/VariantNET/testing_data/chr21/hg38.NA12878-WashU_chr21-14069662-46411975.bam.bai')
    bed = pathing('~/Dropbox/thesis/VariantNET/testing_data/chr21/CHROM21_v.3.3.2_highconf_noinconsistent.bed')
    # p = Pileup(fasta, bam)

import numpy,from numpy,from scipy,from numba


# Copyright (c) 2020 by Fraunhofer Institute for Energy Economics
# and Energy System Technology (IEE), Kassel. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be found in the LICENSE file.

import numpy as np
import pandapipes as ppipes
import pandapower as pp
import pandas as pd
from pandapipes.control.run_control import prepare_run_ctrl as prepare_run_ctrl_ppipes
from pandapower.control.run_control import prepare_run_ctrl as prepare_run_ctrl_pp, \
    net_initialization, get_recycle, control_initialization, control_finalization, \
    _evaluate_net as _evaluate_net, control_implementation, get_controller_order, NetCalculationNotConverged

try:
    import pandaplan.core.pplog as logging
except ImportError:
    import logging

logger = logging.getLogger(__name__)

def _evaluate_multinet(multinet, levelorder, ctrl_variables, **kwargs):
    """
    Within a control loop after all controllers applied their their action "_evaluate_multinet" checks if all
    nets affectd in one level did converge or not

    :param multinet: multinet with multinet controllers, net distinct controllers and several pandapipes/pandapower nets
    :type multinet: pandapipes.Multinet
    :param levelorder: list of tuples given the correct order of the different controllers within one level
    :type levelorder: list
    :param ctrl_variables: contains all relevant information and boundaries required for a successful control run
    :type ctrl_variables: dict
    :param kwargs: additional keyword arguments handed to each run function
    :type kwargs: dict
    :return: as the ctrl_variables are adapted they are returned
    :rtype: dict
    """
    levelorder = np.array(levelorder)
    multinet_converged = []
    rel_nets = _relevant_nets(multinet, levelorder)
    for net_name in multinet['nets'].keys():
        net = multinet['nets'][net_name]
        rel_levelorder = levelorder[rel_nets[net_name]]
        ctrl_variables[net_name] = _evaluate_net(net, rel_levelorder, ctrl_variables[net_name], **kwargs) if np.any(
            rel_nets[net_name]) else ctrl_variables[net_name]
        multinet_converged += [ctrl_variables[net_name]['converged']]
    ctrl_variables['converged'] = np.all(multinet_converged)
    return ctrl_variables

def _relevant_nets(multinet, levelorder):
    """
    This function determines the relevant nets in each level, i.e. only the nets affected in each level are investigated
    and checked.

    :param multinet: multinet with multinet controllers, net distinct controllers and several pandapipes/pandapower nets
    :type multinet: pandapipes.Multinet
    :param levelorder: list of tuples given the correct order of the different controllers within one level
    :type levelorder: list
    :return: dictionary of all nets in multinet.nets. Entries are booleans. True means net is in the
    corresponding level, False means it is not.
    :rtype: dict
    """
    net_names = dict()

    nns = []
    levelorder = np.array(levelorder)
    rel_levelorder_multi = levelorder[:, 1].__eq__(multinet)
    controller = levelorder[rel_levelorder_multi, 0]
    nns += [ctrl.get_all_net_names() for ctrl in controller]

    for net_name in multinet['nets'].keys():
        net = multinet['nets'][net_name]
        rel_levelorder = levelorder[:, 1].__eq__(net)
        level_excl = [False if net_name not in nn else True for nn in nns]
        rel_levelorder_multi[rel_levelorder_multi] = level_excl
        net_names[net_name] = np.maximum(rel_levelorder_multi, rel_levelorder)

    return net_names

def net_initialization_multinet(multinet, ctrl_variables, **kwargs):
    """
    If one controller affecting a net requires an initial_run, a loadflow/pipeflow is conducted.

    :param multinet: multinet with multinet controllers, net distinct controllers and several pandapipes/pandapower nets
    :type multinet: pandapipes.Multinet
    :param ctrl_variables: contains all relevant information and boundaries required for a successful control run
    :type ctrl_variables: dict
    :param kwargs: additional keyword arguments handed to each run function
    :type kwargs: dict
    :return: as the ctrl_variables are adapted they are returned
    :rtype: dict
    """
    ctrl_variables['converged'] = False

    for net_name in multinet['nets'].keys():
        net = multinet['nets'][net_name]
        kwargs['recycle'], kwargs['only_v_results'] = \
            ctrl_variables[net_name]['recycle'], ctrl_variables[net_name]['only_v_results']
        ctrl_variables[net_name] = net_initialization(net, ctrl_variables[net_name], **kwargs)
        ctrl_variables['converged'] = max(ctrl_variables['converged'], ctrl_variables[net_name]['converged'])
    return ctrl_variables

def run_control(multinet, ctrl_variables=None, max_iter=30, **kwargs):
    """
    Main function to call a multnet with controllers.

    Function is running control loops for the controllers specified in net.controller
    Runs controller until each one converged or max_iter is hit.

    1. Call initialize_control() on each controller
    2. Calculate an inital run (if it is enabled, i.e. setting the initial_run veriable to True)
    3. Repeats the following steps in ascending order of controller_order until total convergence of all
       controllers for each level:
       a) Evaluate individual convergence for all controllers in the level
       b) Call control_step() for all controllers in the level on diverged controllers
       c) Fire run function (or optionally another function like runopf or whatever you defined)
    4. Call finalize_control() on each controller

    :param multinet: multinet with multinet controllers, net distinct controllers and several pandapipes/pandapower nets
    :type multinet: pandapipes.Multinet
    :param ctrl_variables: contains all relevant information and boundaries required for a successful control run. To \
           define ctrl_variables yourself, following entries for each net are required:\n
           - level (list): gives a list of levels to be investigated \n
           - controller_order (list): nested list of tuples given the correct order of the
             different controllers within one level\
           - run (funct, e.g. pandapower.runpp, pandapipes.pipeflow): function to be used to
             conduct a loadflow/pipeflow \n
           - initial_run (boolean): Is a initial_run for a net required or not\n
           - continue_on_divergence (boolean): What to do if loadflow/pipeflow is not converging, fires control_repair
    :type ctrl_variables: dict, default: None
    :param max_iter: number of iterations for each controller to converge
    :type max_iter: int, default: 30
    :param kwargs: additional keyword arguments handed to each run function
    :type kwargs: dict
    :return: runs an entire control loop
    :rtype: None
    """
    ctrl_variables = prepare_run_ctrl(multinet, ctrl_variables)
    controller_order = ctrl_variables['controller_order']

    # initialize each controller prior to the first power flow
    control_initialization(controller_order)

    # initial run (takes time, but is not needed for every kind of controller)
    ctrl_variables = net_initialization_multinet(multinet, ctrl_variables, **kwargs)

    # run each controller step in given controller order
    control_implementation(multinet, controller_order, ctrl_variables, max_iter,
                           evaluate_net_fct=_evaluate_multinet, **kwargs)

    # call finalize function of each controller
    control_finalization(controller_order)

def get_controller_order_multinet(multinet):
    """
    Defining the controller order per level.

    Takes the order and level columns from net.controller.
    If levels are specified, the levels and orders are executed in ascending order.

    :param multinet: multinet with multinet controllers, net distinct controllers and several pandapipes/pandapower nets
    :type multinet: pandapipes.Multinet
    :return: nested list of tuples given the correct order of the controllers, respectively for each level
    :rtype: list
    """

    net_list = []
    controller_list = []

    if hasattr(multinet, "controller") and len(multinet.controller[multinet.controller.in_service]) != 0:
        # if no controllers are in the net, we have no levels and no order lists
        multinets = [multinet] * len(multinet.controller)
        net_list += multinets
        controller_list += [multinet.controller.values]

    for net_name in multinet['nets'].keys():
        net = multinet['nets'][net_name]
        if not (hasattr(net, 'controller') and len(net.controller[net.controller.in_service]) != 0):
            # if no controllers are in the net, we have no levels and no order lists
            continue
        nets = [net] * len(net.controller)
        net_list += nets
        controller_list += [net.controller.values]

    if not len(controller_list):
        # if no controllers are in the net, we have no levels and no order lists
        return [0], [[]]
    else:
        controller_list = pd.DataFrame(np.concatenate(controller_list), columns=multinet.controller.columns)
        controller_list = controller_list.astype(multinet.controller.dtypes)
        return get_controller_order(net_list, controller_list)

def prepare_run_ctrl(multinet, ctrl_variables, **kwargs):
    """
    Prepares run control functions.

    Internal variables needed:
        - level (list): gives a list of levels to be investigated
        - controller_order (list): nested list of tuples given the correct order of the different controllers
        within one level
        - run (funct, e.g. pandapower.runpp, pandapipes.pipeflow): function to be used to conduct a loadflow/pipeflow
        - initial_run (boolean): Is a initial_run for a net required or not
        - continue_on_divergence (boolean): What to do if loadflow/pipeflow is not converging, fires control_repair

    You don't need to define it for each net. If one net is not defined, the default settings are used.

    :param multinet: multinet with multinet controllers, net distinct controllers and several pandapipes/pandapower nets
    :type multinet: pandapipes.Multinet
    :param ctrl_variables: contains all relevant information and boundaries required for a successful control run.
    :type ctrl_variables: dict
    :return: adapted ctrl_variables for all nets with all required boundary informaiton
    :rtype: dict
    """

    # sort controller_order by order if not already done
    if ctrl_variables is None:
        ctrl_variables = dict()

    excl_net = []

    if hasattr(multinet, "controller") and len(multinet.controller[multinet.controller.in_service]) != 0:
        for _, c in multinet['controller'].iterrows():
            net_names = c.object.get_all_net_names()
            for net_name in net_names:
                if net_name not in ctrl_variables.keys():
                    ctrl_variables[net_name] = {'run': None, 'initial_run': None,
                                                'continue_on_divergence': None, 'errors': ()}
                net = multinet['nets'][net_name]
                if isinstance(net, ppipes.pandapipesNet):
                    ctrl_variables_net = prepare_run_ctrl_ppipes(net, None, **kwargs)
                elif isinstance(net, pp.pandapowerNet):
                    ctrl_variables_net = prepare_run_ctrl_pp(net, None, **kwargs)
                else:
                    raise ValueError('the given nets are neither pandapipes nor pandapower nets')

                ctrl_variables[net_name]['run'] = ctrl_variables_net['run']
                ctrl_variables[net_name]['errors'] = ctrl_variables_net['errors']
                ctrl_variables[net_name]['initial_run'] = ctrl_variables[net_name]['initial_run'] if \
                    ctrl_variables[net_name]['initial_run'] is not None else ctrl_variables_net['initial_run']
                ctrl_variables[net_name]['only_v_results'], ctrl_variables[net_name]['recycle'] = \
                    get_recycle(ctrl_variables_net)
                ctrl_variables[net_name]['continue_on_divergence'] = \
                    ctrl_variables[net_name]['continue_on_divergence'] if \
                    ctrl_variables[net_name]['continue_on_divergence'] is not None else \
                    ctrl_variables_net['continue_on_divergence']
                excl_net += [net_name]

    for net_name in multinet['nets'].keys():
        if net_name in excl_net:
            continue
        if net_name not in ctrl_variables.keys():
            ctrl_variables[net_name] = {'run': None, 'initial_run': False,
                                        'continue_on_divergence': None, 'errors': ()}
        net = multinet['nets'][net_name]
        if isinstance(net, ppipes.pandapipesNet):
            ctrl_variables_net = prepare_run_ctrl_ppipes(net, None, **kwargs)
        elif isinstance(net, pp.pandapowerNet):
            ctrl_variables_net = prepare_run_ctrl_pp(net, None, **kwargs)
        else:
            raise ValueError('the given nets are neither pandapipes nor pandapower nets')
        ctrl_variables[net_name]['run'] = ctrl_variables_net['run']
        ctrl_variables[net_name]['errors'] = ctrl_variables_net['errors']
        ctrl_variables[net_name]['initial_run'] = ctrl_variables[net_name]['initial_run'] if \
            ctrl_variables[net_name]['initial_run'] is not None else ctrl_variables_net['initial_run']
        ctrl_variables[net_name]['only_v_results'], ctrl_variables[net_name]['recycle'] = \
            get_recycle(ctrl_variables_net)
        ctrl_variables[net_name]['continue_on_divergence'] = \
            ctrl_variables[net_name]['continue_on_divergence'] if \
            ctrl_variables[net_name]['continue_on_divergence'] is not None else \
            ctrl_variables_net['continue_on_divergence']

    if ('check_each_level') in kwargs:
        check = kwargs.pop('check_each_level')
        ctrl_variables['check_each_level'] = check
    else:
        ctrl_variables['check_each_level'] = True

    ctrl_variables['errors'] = (NetCalculationNotConverged,)

    ctrl_variables['level'], ctrl_variables['controller_order'] = get_controller_order_multinet(multinet)

    return ctrl_variables


"""
IRAC imaging of the cluster field taken before and after the SN discovery
epoch

Nothing obviously there, though difference imaging doesn't work very well for
the rotated IRAC PSF

"""

import os
import glob

import numpy as np

import astropy.io.fits as pyfits
import astropy.wcs as pywcs

AORS = ['r58289152', 'r58785536']

def sync():
    for aor in AORS:
        os.system(f'aws s3 sync s3://grizli-v1/IRAC/AORS/{aor}/ ./{aor}/')

def mosaic():

    from grizli import utils
    from drizzlepac.astrodrizzle import ablot

    head, outputwcs = utils.make_wcsheader(ra=24.51568758, dec=-21.92554133, size=240, pixscale=1.0, get_hdu=False, theta=0)

    files = {}
    drz = {}

    for ch in ['ch1', 'ch2']:
        files[ch] = {}
        drz[ch] = {}

        for aor in AORS:
            files[ch][aor] = glob.glob(f'{aor}/{ch}/*/*fits.gz')
            files[ch][aor].sort()

            sci_list = [pyfits.open(file)['CBCD'].data for file in files[ch][aor]]
            unc_list = [pyfits.open(file)['CBUNC'].data for file in files[ch][aor]]
            N = len(sci_list)
            wht_list = [1/unc**2 for unc in unc_list]
            bcd_headers = [pyfits.open(file)['CBCD'].header for file in files[ch][aor]]

            headers = [pyfits.open(file)['WCS'].header for file in files[ch][aor]]
            wcs_list = [pywcs.WCS(head, relax=True) for head in headers]

            for i in range(N):
                mask = (unc_list[i] == 0) | (~np.isfinite(sci_list[i]))
                mask |= (~np.isfinite(unc_list[i]))

                sci_list[i] -= headers[i]['PEDESTAL']
                sci_list[i][mask] = 0
                wht_list[i][mask] = 0
                wcs_list[i].pscale = utils.get_wcs_pscale(wcs_list[i])

            _drz = utils.drizzle_array_groups(sci_list, wht_list, wcs_list, outputwcs=outputwcs, kernel='square', pixfrac=0., verbose=True)

            for iter in range(2):
                cr_wht = [wht*1 for wht in wht_list]

                exptime = 0
                ndrizim = 0

                # CR rejection
                for i in range(N):
                    _blt = ablot.do_blot(_drz[0], _drz[4], wcs_list[i], 1., coeffs=True, interp='poly5')
                    cr = (sci_list[i] - _blt)*np.sqrt(wht_list[i]) > 12
                    cr &= (_blt != 0)
                    cr_wht[i] *= (cr == 0)

                    if (_blt != 0).sum() > 10:
                        exptime += bcd_headers[9]['EXPTIME']
                        ndrizim += 1

                _drz = utils.drizzle_array_groups(sci_list, cr_wht, wcs_list, outputwcs=outputwcs, kernel='square', pixfrac=0., verbose=True)

            _drz[3]['EXPTIME'] = exptime
            _drz[3]['NDRIZIM'] = ndrizim

            for k in ['DATE_OBS', 'MJD_OBS', 'INSTRUME', 'CHNLNUM']:
                _drz[3][k] = bcd_headers[0][k]

            pyfits.writeto(f'{aor}-{ch}_drz_sci.fits', data=_drz[0], header=_drz[3], overwrite=True)
            pyfits.writeto(f'{aor}-{ch}_drz_wht.fits', data=_drz[1], header=_drz[3], overwrite=True)
            drz[ch][aor] = _drz


# PyGame template.

# Import standard modules.
import sys

# Import non-standard modules.
import pygame
from pygame.locals import *
import numpy as np
import math

BLACK = (0, 0, 0)
WHITE = (255, 255, 255)

class Player():
    def __init__(self, location, isRight, controllable=False):
        self.isRightt = isRight
        self.loc = location

    def draw(self, screen):
        pass

    def update(self):
        pass

class Ball():
    def __init__(self, loc):
        self.location = loc
        self.dir = np.random.uniform(0, 2*np.pi)
        # print(self.dir)
        self.speed = 200

    def draw(self, screen):
        pygame.draw.circle(screen, WHITE, self.getLoc(), 30, 1)

    def getLoc(self):
        loc = (int(np.round(self.location[0])), int(np.round(self.location[1])))
        return loc

    def update(self, dt):
        # print(dt)
        dx = np.cos(self.dir)*self.speed*dt/1000
        dy = np.sin(self.dir)*self.speed*dt/1000
        # print(dy, dx)
        # print(dx, dy)
        self.location = (self.location[0] + dx, self.location[1] + dy)

class Game:
    def __init__(self, ballpos, player1pos, player2pos):
        speed = 1
        self.ball = Ball(loc=ballpos)
        self.p1 = Player(player1pos, isRight=True, controllable=True)
        self.p2 = Player(player2pos, isRight=False, controllable=True)

    def draw(self, screen):
        screen.fill(BLACK)  # Fill the screen with black.
        self.p1.draw(screen)
        self.p2.draw(screen)
        self.ball.draw(screen)

        # self.ball.draw(screen)
        # Flip the display so that the things we drew actually show up.
        pygame.display.flip()

    def update(self, dt):
        """
        Update game. Called once per frame.
        dt is the amount of time passed since last frame.
        If you want to have constant apparent movement no matter your framerate,
        what you can do is something like

        x += v * dt

        and this will scale your velocity based on time. Extend as necessary."""
        self.ball.update(dt)
        # Go through events that are passed to the script by the window.
        for event in pygame.event.get():
            # We need to handle these events. Initially the only one you'll want to care
            # about is the QUIT event, because if you don't handle it, your game will crash
            # whenever someone tries to exit.
            if event.type == QUIT:
                pygame.quit()  # Opposite of pygame.init
                sys.exit()  # Not including this line crashes the script on Windows. Possibly
                # on other operating systems too, but I don't know for sure.
            # Handle other events as you wish.

def drawScreen(screen):
    """
    Draw things to the window. Called once per frame.
    """
    screen.fill(BLACK)  # Fill the screen with black.

    # Flip the display so that the things we drew actually show up.
    pygame.display.flip()

def runPyGame():
    # Initialise PyGame.
    pygame.init()

    # Set up the clock. This will tick every frame and thus maintain a relatively constant framerate. Hopefully.
    fps = 60.0
    fpsClock = pygame.time.Clock()

    # Set up the window.
    width, height = 640, 480
    screen = pygame.display.set_mode((width, height))

    # screen is the surface representing the window.
    # PyGame surfaces can be thought of as screen sections that you can draw onto.
    # You can also draw surfaces onto other surfaces, rotate surfaces, and transform surfaces.
    ballpos = (width//2, height//2)
    p1pos = (0, height//2)
    p2pos = (width, height//2)

    game = Game(ballpos, p1pos, p2pos)
    # Main game loop.
    dt = 1/fps  # dt is the time since last frame.
    while True:  # Loop forever!
        game.draw(screen)
        game.update(dt)
        # You can update/draw here, I've just moved the code for neatness.

        # update(dt)
        # draw(screen)

        dt = fpsClock.tick(fps)

if __name__ == '__main__':
    runPyGame()


import numpy as np
from collections import OrderedDict

from robolearn_envs.pybullet.core.bullet_env import BulletEnv
from gym.spaces import Box

from robolearn_envs.pybullet.planarmanipulators.planarmanipulator3dof \
    import PlanarManipulator3Dof
from robolearn_envs.pybullet.planarmanipulators.planarmanipulator4dof \
    import PlanarManipulator4Dof

from robolearn_envs.pybullet.common import GoalCircle
from robolearn_envs.pybullet.common import Plane
from robolearn_envs.pybullet.common import PusherBorder
from robolearn_envs.pybullet.common import Cylinder2d

from robolearn_envs.utils.transformations import euler_from_quat
from robolearn_envs.utils.transformations import create_quat

from robolearn_envs.reward_functions.state_rewards \
    import exponential_distance_reward

# GOAL_RANGE_X = (0.55, 0.7)
# GOAL_RANGE_Y = (0.55, 0.7)
# HALF_TGT_RANGE_X = (0.40, 0.70)
# HALF_TGT_RANGE_Y = (-0.2, 0.70)
GOAL_RANGE_X = (0.63, 0.68)  # Paper
GOAL_RANGE_Y = (0.33, 0.38)  # Paper
HALF_TGT_RANGE_X = (0.50, 0.55)  # Paper
HALF_TGT_RANGE_Y = (0.1, 0.20)  # Paper
FULL_TGT_RANGE_X = (-0.10, 0.70)
FULL_TGT_RANGE_Y = (-0.7, 0.70)

GOAL_RANGE_X = (0.53, 0.78)  # Prueba 04/03/19
GOAL_RANGE_Y = (0.23, 0.48)  # Prueba 04/03/19
TGT_RANGE_X = (0.40, 0.65)  # Prueba 04/03/19
TGT_RANGE_Y = (0.0, 0.30)  # Prueba 04/03/19

MAX_GOAL_NORM = 0.8
MAX_CYLINDER_POS_NORM = 0.65
MIN_CYLINDER_POS_NORM = 0.45
ROBOT_CONFIG_RANGE = [
    (-1.35, 1.35),
    (-0.90, 0.9),
    (-1.5, 1.5),
    (-1.5, 1.5),  # Only used for 4dof robot
]

class Pusher2DBulletEnv(BulletEnv):
    def __init__(
            self,
            robot_dof=3,
            is_render=False,
            obs_distances=False,
            only_position=True,  # TODO: Scale yaw cost if only_position=False
            rdn_goal_pose=True,
            goal_pose=None,
            tgt_pose=None,
            rdn_tgt_object_pose=True,
            robot_config=None,
            rdn_robot_config=True,
            tgt_cost_weight=1.0,
            goal_cost_weight=3.0,
            ctrl_cost_weight=1.0e-2,
            goal_tolerance=0.05,
            max_time=None,  # seconds
            half_env=True,
            sim_timestep=1/240.,
            frame_skip=1,
            seed=None,
    ):
        """Planar3DoF robot seeking to push a cylinder to a target position.
        """
        super(Pusher2DBulletEnv, self).__init__(
            sim_timestep=sim_timestep, frameskip=frame_skip,
            is_render=is_render, seed=seed,
        )

        # Plane
        # self.plane = PlaneObject(plane_type='plane_simple')
        self.plane = Plane(plane_type='simple', color=None)
        self.add_to_sim(self.plane)

        # Border
        self._half_env = half_env
        self.border = PusherBorder(color='gray', half_env=self._half_env)
        self.add_to_sim(self.border)

        # Robot
        self._robot_config_is_rdn = rdn_robot_config
        if robot_config is None:
            self._robot_config_mean = \
                [np.deg2rad(-100), np.deg2rad(45), np.deg2rad(20), np.deg2rad(5)][:robot_dof]
        else:
            self._robot_config_mean = robot_config
        self._robot_config_var = [0.05, 0.2, 0.3, 0.3]
        if robot_dof == 3:
            self.robot = PlanarManipulator3Dof(
                self_collision=True,
                init_config=self._robot_config_mean,
                robot_model=None,
                control_mode='joint_torque',
            )
        elif robot_dof == 4:
            self.robot = PlanarManipulator4Dof(
                self_collision=False,
                init_config=self._robot_config_mean,
                robot_model=None,
                control_mode='joint_torque',
            )
        else:
            raise NotImplementedError("Only planar manipulators with 3 and 4"
                                      "DoFs have been implemented!")

        self.add_to_sim(self.robot)
        self._init_robot_config = self.robot.initial_configuration

        # Goal
        self._goal_pos_is_rdn = rdn_goal_pose
        self._goal_state = np.inf * np.ones(3)
        if tgt_pose is None:
            if self._half_env:
                tgt_range_x = HALF_TGT_RANGE_X
                tgt_range_y = HALF_TGT_RANGE_Y
            else:
                tgt_range_x = FULL_TGT_RANGE_X
                tgt_range_y = FULL_TGT_RANGE_Y

            while (np.linalg.norm(self._goal_state[:2]) > MAX_CYLINDER_POS_NORM) or \
                    (np.linalg.norm(self._goal_state[:2]) < MIN_CYLINDER_POS_NORM):
                self._goal_state[0] = \
                    np.random.uniform(tgt_range_x[0], tgt_range_x[1])
                self._goal_state[1] = \
                    np.random.uniform(tgt_range_y[0], tgt_range_y[1])
        else:
            self._goal_state = tgt_pose
        self.goal = GoalCircle(init_pos=(0, 0, -0.0055), color='green')
        self.add_to_sim(self.goal)

        # Cylinder Object
        self.obj_pose_is_rdn = rdn_tgt_object_pose
        if tgt_pose is None:
            self._obj_state_mean = [0.6, 0.25, 1.4660]
        else:
            self._obj_state_mean = tgt_pose
        # self.goal_state_mean = [[0.8+0*0.1+0*0.105,  0*0.105,   0.],
        #                        ]
        self._obj_state_var = [0.05 * 1, 0.2 * .1, 0.2]
        self._obj_height = 0.055

        self.cylinder = Cylinder2d(init_pos=[self._obj_state_mean[0],
                                             self._obj_state_mean[1],
                                             self._obj_height],
                                   color='yellow')
        self.add_to_sim(self.cylinder)

        # Reset environment so we can get info from pybullet
        self.set_rendering(False)
        self._is_env_instantiation_complete = False
        self.reset()
        self._is_render = is_render

        # ACTUATION
        self.action_space = Box(
            self.robot.low_action_bounds,
            self.robot.high_action_bounds,
            dtype=np.float32
        )

        # OBSERVATION
        robot_state_dim = self.robot.observation_dim  # joint pos/vel
        self._use_obs_distances = obs_distances
        if only_position:
            self._pose_dof = 2   # [x, y]
        else:
            self._pose_dof = 3   # [x, y, theta]
        if obs_distances:
            optitrack_dim = 2*self._pose_dof
        else:
            optitrack_dim = 3*self._pose_dof
        obs_dim = robot_state_dim + optitrack_dim
        self.observation_space = Box(
            low=-np.inf, high=np.inf, shape=(obs_dim,),
            dtype=np.float32
        )
        self._observation = OrderedDict()
        self._observation['robot_state'] = np.zeros(robot_state_dim)
        if obs_distances:
            self._observation['obj_ee_diff'] = np.zeros(self._pose_dof)
            self._observation['goal_obj_diff'] = np.zeros(self._pose_dof)
        else:
            self._observation['cylinder_pose'] = np.zeros(self._pose_dof)
            self._observation['goal_pose'] = np.zeros(self._pose_dof)
            self._observation['ee_pose'] = np.zeros(self._pose_dof)

        # STATE
        state_dim = robot_state_dim + 3*3  # robot_state + (XYyaw)*(EE, OBJ, GOAL)
        self.state_space = Box(
            low=-np.inf, high=np.inf, shape=(state_dim,),
            dtype=np.float32
        )
        self._state = OrderedDict(
            robot_state=np.zeros(robot_state_dim),
            object_pose=np.zeros(3),
            goal_pose=np.zeros(3),
            ee_pose=np.zeros(3),
        )

        # Environment settings
        self.max_time = max_time
        if self._half_env:
            # self.set_visualizer_data(distance=1.1, yaw=-30, pitch=-90,
            self.set_visualizer_data(distance=1.1, yaw=0.0, pitch=-89.95,
                                     target_pos=[0.35, 0.0, 0.0])
        else:
            self.set_visualizer_data(distance=1.1, yaw=-30, pitch=-90,
                                     target_pos=None)
        self.set_rendering(self._is_render)

        # Initial conditions
        self._initial_conditions = []
        self._initial_obs_conditions = []
        self._current_init_cond = np.zeros(self.state_dim)

        # Costs
        self._tgt_cost_weight = tgt_cost_weight
        self._goal_cost_weight = goal_cost_weight
        self._goal_tolerance = goal_tolerance
        self._ctrl_cost_weight = ctrl_cost_weight
        # Max reward
        self._max_rewards = np.zeros(1)
        self._is_env_instantiation_complete = True

    def _reset_env(self, *args, **kwargs):
        """Reset the robot and the objects in the environment.

        Function called by self.reset()

        Args:
            *args:
            **kwargs:

        Returns:
            np.ndarray: Current environment observation

        """
        # Reset Simulation
        self.reset_simulation()

        # Reset ground in simulation
        self.plane.reset()

        # Border
        self.border.reset()

        # Robot
        condition = kwargs.pop('condition', None)
        robot_state = self.robot.reset()
        if self._is_env_instantiation_complete:
            if condition is None:
                if self._robot_config_is_rdn:
                    wrong_config = True
                    while wrong_config:
                        init_robot_config = np.clip(
                            [
                                np.random.uniform(ROBOT_CONFIG_RANGE[0][0],
                                                  ROBOT_CONFIG_RANGE[0][1]),
                                np.random.uniform(ROBOT_CONFIG_RANGE[1][0],
                                                  ROBOT_CONFIG_RANGE[1][1]),
                                np.random.uniform(ROBOT_CONFIG_RANGE[2][0],
                                                  ROBOT_CONFIG_RANGE[2][1]),
                                np.random.uniform(ROBOT_CONFIG_RANGE[3][0],
                                                  ROBOT_CONFIG_RANGE[3][1]),
                            ][:self.robot.dof],
                            self.min_robot_config,
                            self.max_robot_config
                        )
                        self.robot.set_state(np.concatenate((
                            init_robot_config,
                            np.zeros_like(init_robot_config)))
                        )
                        x_lim = -0.10 if self._half_env else -0.5
                        wrong_config = self.get_ee_pose()[0] <= x_lim
                else:
                    init_robot_config = self.init_robot_config

            else:
                if isinstance(condition, int):
                    init_robot_config = self._init_robot_configs[condition]
                else:
                    raise NotImplementedError
                    # init_robot_config = np.clip(condition,
                    #                             self.min_robot_config,
                    #                             self.max_robot_config)
            self.robot.initial_configuration = init_robot_config

        # Goal
        if self._is_env_instantiation_complete:
            if condition is None:
                if self._goal_pos_is_rdn:
                    self._goal_state = np.array([np.inf, np.inf, 0])
                    if self._half_env:
                        tgt_range_x = HALF_TGT_RANGE_X
                        tgt_range_y = HALF_TGT_RANGE_Y
                    else:
                        tgt_range_x = FULL_TGT_RANGE_X
                        tgt_range_y = FULL_TGT_RANGE_Y

                    while (np.linalg.norm(self._goal_state) > MAX_CYLINDER_POS_NORM) or \
                            (np.linalg.norm(self._goal_state) < MIN_CYLINDER_POS_NORM):
                        self._goal_state[0] = \
                            np.random.uniform(tgt_range_x[0], tgt_range_x[1])
                        self._goal_state[1] = \
                            np.random.uniform(tgt_range_y[0], tgt_range_y[1])
            else:
                self._goal_state = self._init_tgt_states[condition]

        self.goal.reset(pose=(self._goal_state[0], self._goal_state[1], -0.0055))

        # Target Object(s)
        self.cylinder.reset()
        if self._is_env_instantiation_complete:
            xy_offset = np.array([np.inf, np.inf])
            if self.obj_pose_is_rdn:
                wrong_obj_pose = True
                while (np.linalg.norm(self._obj_state_mean[:2] + xy_offset) >
                       MAX_CYLINDER_POS_NORM) or \
                      (np.linalg.norm(self._obj_state_mean[:2] + xy_offset) <
                       MIN_CYLINDER_POS_NORM) or wrong_obj_pose:
                    xy_offset = \
                        self.np_random.randn(2)*np.sqrt(self._obj_state_var[:2])

                    # Set XY position
                    des_pos = np.zeros(3)
                    des_pos[:2] = self._obj_state_mean[:2] + xy_offset
                    des_pos[0] = np.clip(des_pos[0], TGT_RANGE_X[0], TGT_RANGE_X[1])
                    des_pos[1] = np.clip(des_pos[1], TGT_RANGE_Y[0], TGT_RANGE_Y[1])
                    # Set the height (Z)
                    des_pos[2] = self._obj_height
                    # Set the orientation (Yaw)
                    yaw_offset = 0
                    des_yaw = self._obj_state_mean[2] + yaw_offset
                    # Set new cylinder pose
                    self.cylinder.set_pose(position=des_pos,
                                           orientation=create_quat(rot_yaw=des_yaw))
                    wrong_obj_pose = self.is_robot_touching_cylinder()

                yaw_offset = \
                    self.np_random.randn(1)*np.sqrt(self._obj_state_var[2])
            else:
                raise NotImplementedError
                # xy_offset = np.zeros(2)
                # yaw_offset = 0

        # Update Environment State and Observation
        state = self._update_env_state()
        observation = self._update_env_obs()

        # Replace init_cond with current state
        if self._is_env_instantiation_complete:
            self._current_init_cond = np.copy(self._state)

        # Update max reward
        if self._is_env_instantiation_complete:
            self._max_reward = np.array([0.])
            # self._max_reward = self._calc_max_reward()

        # Visualization (if applicable)
        self.enable_vis()

        return observation

    # ############## #
    # ACTION related #
    # ############## #

    def _set_action(self, action):
        self.robot.apply_action(action)

        # Simulation step
        self.sim_step()

    # ############# #
    # STATE related #
    # ############# #
    def _update_env_state(self):
        if self._is_env_instantiation_complete:
            # Robot state
            self._state['robot_state'][:] = \
                self.robot.get_state(only_controlled_joints=True)

            self._state['object_pose'][:] = self.get_object_pose()
            self._state['goal_pose'][:] = self.get_goal_pose()
            self._state['ee_pose'][:] = self.get_ee_pose()

        return self.get_state()

    def get_state(self):
        if self._is_env_instantiation_complete:
            return np.concatenate(
                [state.flatten() for state in self._state.values()]
            )

    def set_state(self, state):
        """
        Set the state of the environment.
        Args:
            state (float[Sdim]): Desired state of the environment.

        Returns:

        """
        if not self._is_env_instantiation_complete:
            raise ValueError("Not possible to set state!")
        if len(state) != self.state_dim:
            raise ValueError("Desired state size does not correpond to "
                             "the environment state. (%d != %d)" %
                             (len(state), self.state_dim))
        init_idx = 0
        for key, val in self._state.items():
            dim = np.prod(val.shape)
            self._state[key] = state[init_idx:init_idx+dim]
            # Set robot state
            if key == 'robot_state':
                self.robot.set_state(self._state[key])
            # elif key.endswith('_pose'):

            init_idx += dim

        self._update_env_state()
        self._update_env_obs()

    # ################### #
    # OBSERVATION related #
    # ################### #

    def _update_env_obs(self):
        if self._is_env_instantiation_complete:
            # Robot state
            self._observation['robot_state'][:] = \
                self.robot.get_state(only_controlled_joints=True)

            obj_state = self.get_object_pose()[:self._pose_dof]
            goal_state = self.get_goal_pose()[:self._pose_dof]
            ee_state = self.get_ee_pose()[:self._pose_dof]

            if self._use_obs_distances:
                self._observation['obj_ee_diff'][:] = obj_state - ee_state
                self._observation['goal_obj_diff'][:] = goal_state - obj_state
            else:
                self._observation['cylinder_pose'][:] = obj_state
                self._observation['goal_pose'][:] = goal_state
                self._observation['ee_pose'][:] = ee_state

        return self.get_observation()

    def get_observation(self):
        if self._is_env_instantiation_complete:
            return np.concatenate(
                [obs.flatten() for obs in self._observation.values()]
            )

    def set_observation(self, observation):
        if not self._is_env_instantiation_complete:
            raise ValueError("Not possible to set observation!")
        if len(observation) != self.obs_dim:
            raise ValueError("Desired observation size does not correpond to "
                             "the environment observation dim. (%d != %d)" %
                             (len(observation), self.obs_dim))
        init_idx = 0
        for key, val in self._observation:
            dim = np.prod(self._observation[val].shape)
            self._observation[val] = observation[init_idx:init_idx+dim]
            init_idx += dim

    # ############ #
    # STEP related #
    # ############ #

    def _check_termination(self):
        if self.max_time is None:
            ee_pos = self.get_ee_pose()[:2]

            x_tol = -0.10 if self._half_env else -0.50
            done = ee_pos[0] <= x_tol
            if done:
                done_trigger = 'failure'
            else:
                done_trigger = False

        else:
            if self.sim_time >= self.max_time:
                done = True
                done_trigger = 'time'
            else:
                done = False
                done_trigger = False

        info = {
            'done_trigger': done_trigger,
        }

        return done, info

    def _compute_reward(self, state, action, next_state=None):
        goal_pose = self.get_goal_pose()[:self._pose_dof]
        object_pose = self.get_object_pose()[:self._pose_dof]
        ee_pose = self.get_ee_pose()[:self._pose_dof]

        # Distance object to goal
        goal_obj_reward, _ = exponential_distance_reward(
            goal_pose - object_pose, temperature=0.1)
        goal_obj_reward = self._goal_cost_weight * goal_obj_reward

        # Give distance object to goal reward only if touching it
        goal_obj_reward *= (self.is_robot_touching_cylinder(only_with_ee=True)
                            and self.is_cylinder_inside_gripper())

        # Distance ee to object
        obj_ee_reward = exponential_distance_reward(object_pose - ee_pose,
                                                    temperature=0.6)[0].item()
        ee_obj_reward = self._tgt_cost_weight * obj_ee_reward

        # Control
        ctrl_cost = self._ctrl_cost_weight * np.square(action).sum()

        reward_composition = np.array([
            ee_obj_reward,
            goal_obj_reward,
            -ctrl_cost,
        ])

        reward = reward_composition.sum().item()

        info = {
            'reward_composition': reward_composition
        }

        return reward, info

    def _calc_max_reward(self):
        max_reward = self._compute_reward(np.zeros(self.state_dim),
                                          np.zeros(self.action_dim)
                                          )
        return max_reward

    def set_tgt_pos(self, tgt_pose):
        self._tgt_state = tgt_pose

    def add_initial_condition(self, state):
        self.set_rendering(False)
        self.reset(-1)
        self.set_rendering(self._is_render)
        full_state = self.get_state()
        full_obs = self.get_observation()

        self._initial_conditions.append(full_state)
        self._initial_obs_conditions.append(full_obs)

    def clear_initial_condition(self, idx=None):
        if idx is None:
            self._initial_conditions.clear()
        else:
            self._initial_conditions.pop(idx)

    def get_conditions(self, cond=None):
        if cond is None:
            return list(self._initial_conditions)
        else:
            return self._initial_conditions[cond]

    def get_total_joints(self):
        return self.robot.total_joints

    def get_ee_pose(self):
        """ Returns pose of the end-effector, defined by (X, Y, yaw) """
        gripper_pose = self.robot.links['gripper_center'].get_pose()
        return self.convert_2d_pose(gripper_pose)

    def get_object_pose(self):
        """ Returns pose of the object object, defined by (X, Y, yaw) """
        tgt_pose = self.cylinder.get_pose()
        return self.convert_2d_pose(tgt_pose)

    def get_goal_pose(self):
        """ Returns pose of the goal, defined by (X, Y, yaw) """
        goal_pose = self.goal.get_pose()
        return self.convert_2d_pose(goal_pose)

    def is_robot_touching_cylinder(self, only_with_ee=False):
        if only_with_ee:
            robot_body = self.robot.links['gripper']
        else:
            robot_body = None

        is_touching = len(self.get_contacts(self.robot, self.cylinder,
                                            robot_body)) > 0
        return is_touching

    def is_cylinder_inside_gripper(self):
        is_inside_gripper = np.linalg.norm(
            self.get_object_pose()[:2] - self.get_ee_pose()[:2]) < 0.03
        return is_inside_gripper

    @staticmethod
    def convert_2d_pose(pose):
        """ It converts a pose defined by (x,y,z,ox,oy,oz,ow) to (x,z,theta)"""
        xy = pose[:2]
        ori = euler_from_quat(pose[3:])[2]
        # ori = normalize_angle(ori, range='pi')
        return np.array([xy[0], xy[1], ori])

    def set_robot_config(self, joint_config):
        self.robot.set_state(joint_config, np.zeros_like(joint_config))

    @property
    def initial_conditions(self):
        return self._initial_conditions

    @property
    def initial_obs_conditions(self):
        return self._initial_obs_conditions

    @property
    def min_robot_config(self):
        limits_robot_config = self.robot.get_joint_limits()
        return [limit[0] for limit in limits_robot_config]

    @property
    def max_robot_config(self):
        limits_robot_config = self.robot.get_joint_limits()
        return [limit[1] for limit in limits_robot_config]

    @property
    def init_robot_config(self):
        return self.robot.initial_configuration

if __name__ == "__main__":
    render = True
    # render = False

    H = 10000

    env = Pusher2DBulletEnv(robot_dof=4,
                            is_render=render)
    print('env_reset')
    env.reset()
    if render:
        env.render()

    for tt in range(H):
        print('env_step: %02d' % tt)
        # next_obs, rew, done, info = env.step(env.action_space.sample())
        next_obs, rew, done, info = env.step(np.zeros_like(env.action_space.sample()))

        if done:
            print("The environment is done!")
            # break

        if render:
            env.render()

    print("Correct!")


# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

import time
import cv2
import numpy as np
from PIL import Image
from mindspore import nn
from mindspore.ops import operations as ops
from mindspore.train.callback import Callback
from mindspore.common.tensor import Tensor

class UnetEval(nn.Cell):
    """
    Add Unet evaluation activation.
    """
    def __init__(self, net, need_slice=False):
        super(UnetEval, self).__init__()
        self.net = net
        self.need_slice = need_slice
        self.transpose = ops.Transpose()
        self.softmax = ops.Softmax(axis=-1)
        self.argmax = ops.Argmax(axis=-1)
        self.squeeze = ops.Squeeze(axis=0)

    def construct(self, x):
        out = self.net(x)
        if self.need_slice:
            out = self.squeeze(out[-1:])
        out = self.transpose(out, (0, 2, 3, 1))
        softmax_out = self.softmax(out)
        argmax_out = self.argmax(out)
        return (softmax_out, argmax_out)

class TempLoss(nn.Cell):
    """A temp loss cell."""
    def __init__(self):
        super(TempLoss, self).__init__()
        self.identity = ops.identity()
    def construct(self, logits, label):
        return self.identity(logits)

def apply_eval(eval_param_dict):
    """run Evaluation"""
    model = eval_param_dict["model"]
    dataset = eval_param_dict["dataset"]
    metrics_name = eval_param_dict["metrics_name"]
    index = 0 if metrics_name == "dice_coeff" else 1
    eval_score = model.eval(dataset, dataset_sink_mode=False)["dice_coeff"][index]
    return eval_score

class dice_coeff(nn.Metric):
    """Unet Metric, return dice coefficient and IOU."""
    def __init__(self, cfg_unet, print_res=True):
        super(dice_coeff, self).__init__()
        self.clear()
        self.cfg_unet = cfg_unet
        self.print_res = print_res

    def clear(self):
        self._dice_coeff_sum = 0
        self._iou_sum = 0
        self._samples_num = 0

    def update(self, *inputs):
        if len(inputs) != 2:
            raise ValueError('Need 2 inputs ((y_softmax, y_argmax), y), but got {}'.format(len(inputs)))
        y = self._convert_data(inputs[1])
        self._samples_num += y.shape[0]
        y = y.transpose(0, 2, 3, 1)
        b, h, w, c = y.shape
        if b != 1:
            raise ValueError('Batch size should be 1 when in evaluation.')
        y = y.reshape((h, w, c))
        if self.cfg_unet["eval_activate"].lower() == "softmax":
            y_softmax = np.squeeze(self._convert_data(inputs[0][0]), axis=0)
            if self.cfg_unet["eval_resize"]:
                y_pred = []
                for i in range(self.cfg_unet["num_classes"]):
                    y_pred.append(cv2.resize(np.uint8(y_softmax[:, :, i] * 255), (w, h)) / 255)
                y_pred = np.stack(y_pred, axis=-1)
            else:
                y_pred = y_softmax
        elif self.cfg_unet["eval_activate"].lower() == "argmax":
            y_argmax = np.squeeze(self._convert_data(inputs[0][1]), axis=0)
            y_pred = []
            for i in range(self.cfg_unet["num_classes"]):
                if self.cfg_unet["eval_resize"]:
                    y_pred.append(cv2.resize(np.uint8(y_argmax == i), (w, h), interpolation=cv2.INTER_NEAREST))
                else:
                    y_pred.append(np.float32(y_argmax == i))
            y_pred = np.stack(y_pred, axis=-1)
        else:
            raise ValueError('config eval_activate should be softmax or argmax.')
        y_pred = y_pred.astype(np.float32)
        inter = np.dot(y_pred.flatten(), y.flatten())
        union = np.dot(y_pred.flatten(), y_pred.flatten()) + np.dot(y.flatten(), y.flatten())

        single_dice_coeff = 2 * float(inter) / float(union+1e-6)
        single_iou = single_dice_coeff / (2 - single_dice_coeff)
        if self.print_res:
            print("single dice coeff is: {}, IOU is: {}".format(single_dice_coeff, single_iou))
        self._dice_coeff_sum += single_dice_coeff
        self._iou_sum += single_iou

    def eval(self):
        if self._samples_num == 0:
            raise RuntimeError('Total samples num must not be 0.')
        return (self._dice_coeff_sum / float(self._samples_num), self._iou_sum / float(self._samples_num))

class StepLossTimeMonitor(Callback):

    def __init__(self, batch_size, per_print_times=1):
        super(StepLossTimeMonitor, self).__init__()
        if not isinstance(per_print_times, int) or per_print_times < 0:
            raise ValueError("print_step must be int and >= 0.")
        self._per_print_times = per_print_times
        self.batch_size = batch_size

    def step_begin(self, run_context):
        self.step_time = time.time()

    def step_end(self, run_context):

        step_seconds = time.time() - self.step_time
        step_fps = self.batch_size*1.0/step_seconds

        cb_params = run_context.original_args()
        loss = cb_params.net_outputs

        if isinstance(loss, (tuple, list)):
            if isinstance(loss[0], Tensor) and isinstance(loss[0].asnumpy(), np.ndarray):
                loss = loss[0]

        if isinstance(loss, Tensor) and isinstance(loss.asnumpy(), np.ndarray):
            loss = np.mean(loss.asnumpy())

        cur_step_in_epoch = (cb_params.cur_step_num - 1) % cb_params.batch_num + 1

        if isinstance(loss, float) and (np.isnan(loss) or np.isinf(loss)):
            raise ValueError("epoch: {} step: {}. Invalid loss, terminating training.".format(
                cb_params.cur_epoch_num, cur_step_in_epoch))
        self.losses.append(loss)
        if self._per_print_times != 0 and cb_params.cur_step_num % self._per_print_times == 0:
            # TEST
            print("step: %s, loss is %s, fps is %s" % (cur_step_in_epoch, loss, step_fps), flush=True)

    def epoch_begin(self, run_context):
        self.epoch_start = time.time()
        self.losses = []

    def epoch_end(self, run_context):
        cb_params = run_context.original_args()
        epoch_cost = time.time() - self.epoch_start
        step_in_epoch = (cb_params.cur_step_num - 1) % cb_params.batch_num + 1
        step_fps = self.batch_size * 1.0 * step_in_epoch / epoch_cost
        print("epoch: {:3d}, avg loss:{:.4f}, total cost: {:.3f} s, per step fps:{:5.3f}".format(
            cb_params.cur_epoch_num, np.mean(self.losses), epoch_cost, step_fps), flush=True)

def mask_to_image(mask):
    return Image.fromarray((mask * 255).astype(np.uint8))

def filter_checkpoint_parameter_by_list(param_dict, filter_list):
    """remove useless parameters according to filter_list"""
    for key in list(param_dict.keys()):
        for name in filter_list:
            if name in key:
                print("Delete parameter from checkpoint: ", key)
                del param_dict[key]
                break


import eel
import io
import cv2
import numpy as np
from PIL import ImageGrab
import os
import time

"""
button with picture of each champion.
pressing button adds and removes that champion from the list.
when the champion is found, play a sound

use js to edit list
then send list to python to check screenshot, if champion is on screen
then output sound
"""

eel.init('web')

#print(os.listdir("champions"))
#Champions we are looking for in screenshots
#champions = ["darius.png", "garen.png", "vayne.png", "lucian.png"]

def screenGrab():
    sc = ImageGrab.grab()
    screenshotname = "screenshot"+str(int(time.time()))+".png"
    sc.save(os.getcwd() + "\\livescreenshots\\" +screenshotname , "PNG")
    return screenshotname

@eel.expose
def checkCurrentScreen(champions):
    print(champions)
    onScreenChampions = [] #No champions found on screen yet

    screenshot_name = screenGrab() #Get name of screenshot and read it into cv2 as currentScreen
    currentScreen = cv2.imread("livescreenshots\\"+screenshot_name)

    w_screenshot, h_screenshot = currentScreen.shape[:-1] #Get dimensions of screenshot

    #Define the dimensions we want the champion icons to be, based on the dimensions of the screenshot
    w = round(w_screenshot/(1080/185))
    h = round(h_screenshot/(1920/139))

    #Where each champion is checked
    for champion in champions:
        query = cv2.resize(cv2.imread("../champions/"+champion), (w, h)) #Resizes the champion icon accordingly and reads it into the variable query

        result = cv2.matchTemplate(currentScreen, query, cv2.TM_CCOEFF_NORMED)
        threshold = 0.9
        loc = np.where(result >= threshold)

        #Labels currentScreen where champion is found, prints which champion is found and adds them to the list of champions on the screen
        for pt in zip(*loc[::-1]):
            cv2.rectangle(currentScreen, pt, (pt[0] + w, pt[1] + h),(0,255,255),2)
            onScreenChampions.append(champion)
            print("Found", champion)

    if onScreenChampions: #If something is on screen, print the list
        print(list(set(onScreenChampions)))
        eel.playSound()
        #If this is running, user should be notified

    eel.runCheckScreen(champions)
    #checkCurrentScreen(champions)

eel.start('index.html', size=(1280, 720))
print("started")

Champion Detection/_GUI/oldmain for champion tracking.py


import time
import math as m
import redis
import struct
import numpy as np
from adafruit_servokit import ServoKit

r = redis.Redis(host='localhost', port=6379, db=0)
kit = ServoKit(channels=16)

#controllable variables
def rget_and_float(name, default = None):
    output = r.get(name)
    if output == None:
        return default
    else:
        return float(output)

rear_diff_locked = int(rget_and_float('rear_diff_locked', 1))
front_diff_locked = int(rget_and_float('front_diff_locked', 1))
gear = int(rget_and_float('gear', 1))
low_battery_voltage = rget_and_float('low_battery_voltage', 3.5)

#----

speed_cap = 45 #percentage of max speed
#steering angle 30 - 150

throttle_stop = 72
throttle_full_forward = 180
throttle_full_reverse = 0

steering_pin = 15
esc_pin = 14

frontdiff_pin = 11
reardiff_pin = 13
gearbox_pin = 12

gear_servo_pos = [0, 60, 110]

rear_diff_servo_pos = [78, 15] #0 locked, 1 open

front_diff_servo_pos = [120, 55] #0 locked, 1 open

def steering_angle(angle):
    if angle > 55:
        angle = 55
    if angle < -55:
        angle = -55
    kit.servo[steering_pin].angle = -angle + 88

def driving_speed_signal(speed):
    if speed > 100:
        speed = 100
    if speed < -72:
        speed = -72
    kit.servo[esc_pin].angle = speed * speed_cap / 100 + 72

driving = True
in_motion_start = time.time()
while driving:
    rear_diff_locked = int(rget_and_float('rear_diff_locked', 1))
    front_diff_locked = int(rget_and_float('front_diff_locked', 1))
    gear = int(rget_and_float('gear', 1))

    kit.servo[gearbox_pin].angle = gear_servo_pos[gear]
    kit.servo[reardiff_pin].angle = rear_diff_servo_pos[rear_diff_locked]
    kit.servo[frontdiff_pin].angle = front_diff_servo_pos[front_diff_locked]

    low_battery_voltage = rget_and_float('low_battery_voltage', 3.5)
    voltages_received = r.get('voltages')
    if voltages_received is None:
        print("no battery info")
        break
    else:
        voltages = np.array(struct.unpack('%sf' %2, voltages_received))
    if voltages.min() < low_battery_voltage:
        print(voltages.min())
        break

    target_speed = r.get('target_speed')
    current_speed_received = r.get('current_speed')

    if current_speed_received is not None:
        current_speed = float(current_speed_received)
        #print(current_speed)

    if target_speed is None:
        #print("no driving input received")
        driving_speed_signal(0)
        in_motion_start = time.time()
    else:
        target_speed = float(target_speed)
        if target_speed > 0:
            if current_speed < 0.05 and time.time() - in_motion_start > 2:
                driving_speed_signal(target_speed * 1.5)
                #print("driving faster")
            else:
                driving_speed_signal(target_speed * 1)
                #print("driving normal speed")
        else:
            driving_speed_signal(0)
            #print("stopped")
            in_motion_start = time.time()

    angle_received = r.get('angle')
    if angle_received is None:
        #print("no steering input received")
        steering_angle(0)
    else:
        steering_angle(float(angle_received))
    r.psetex('log_driving_running', 1000, "on")
    time.sleep(0.03) # ???

print("stopping")
driving_speed_signal(0)
steering_angle(-20)
time.sleep(1)
steering_angle(20)
time.sleep(1)
steering_angle(-20)
time.sleep(1)
steering_angle(0)


import cv2 as cv
import tensorflow as tf
import numpy as np

class Predictor:

    # Darknet and CNN Parameters
    confidence_threshold = 0.1  # Confidence threshold
    nms_threshold = 0.6  # Non-maximum suppression threshold

    yolo_net_width = 416
    yolo_net_height = 416

    # Load all models and configs
    plates_yolo_config = "dl_prediction/config/yolov3_plates.cfg"
    plates_yolo_weights = "dl_prediction/models/yolov3_plates_final.weights"
    plates_classes = ['Plate']

    chars_classes = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', 'G',
                     'H', 'I', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z']
    chars_yolo_config = "dl_prediction/config/yolov3_chars.cfg"
    chars_yolo_weights = "dl_prediction/models/yolov3_chars_final.weights"

    def __init__(self):
        self.plates_yolo_net = self.get_yolo_net(self.plates_yolo_config, self.plates_yolo_weights)
        self.chars_yolo_net = self.get_yolo_net(self.chars_yolo_config, self.chars_yolo_weights)
        self.cnn_chars_model = tf.keras.models.load_model('dl_prediction/models/cnn_chars_recognition.h5')

    @staticmethod
    def get_yolo_net(config, weights):
        yolo_net = cv.dnn.readNetFromDarknet(config, weights)
        yolo_net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
        yolo_net.setPreferableTarget(cv.dnn.DNN_TARGET_CPU)
        return yolo_net

    @staticmethod
    def process_license_plate(license_plate):
        gray = cv.cvtColor(license_plate, cv.COLOR_BGR2GRAY)
        ret, thresh = cv.threshold(gray, 127, 255, 0)

        contours, _ = cv.findContours(thresh, cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
        areas = [cv.contourArea(c) for c in contours]

        if len(areas) != 0:
            max_index = np.argmax(areas)
            cnt = contours[max_index]
            x, y, w, h = cv.boundingRect(cnt)
            cv.rectangle(license_plate, (x, y), (x + w, y + h), (0, 255, 0), 2)
            processed_license_plate = license_plate[y: y + h, x: x + w]
        else:
            processed_license_plate = license_plate

        return processed_license_plate

    @staticmethod
    def draw_pred(frame, name, conf, left, top, right, bottom, color=(0, 255, 0)):
        cv.rectangle(frame, (left, top), (right, bottom), color, 3)
        # label = '{}:{}'.format(name, '%.2f' % conf)
        label = name
        label_size, base_line = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        top = max(top, label_size[1])
        cv.rectangle(frame, (left, top - round(1.5 * label_size[1])),
                     (left + round(1.5 * label_size[0]), top + base_line), (0, 0, 255), cv.FILLED)
        cv.putText(frame, label, (left, top), cv.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 0), 1)

    def predict_boxes(self, frame, yolo_outputs, is_license_plate=True):
        classes = []
        confidences = []
        boxes = []

        max_confidence = 0.0
        for output in yolo_outputs:
            for prediction in output:
                scores = prediction[5:]
                class_id = np.argmax(scores)
                confidence = scores[class_id]
                max_confidence = max(confidence, max_confidence)
                if confidence > self.confidence_threshold:
                    center_x = int(prediction[0] * frame.shape[1])
                    center_y = int(prediction[1] * frame.shape[0])

                    width = int(prediction[2] * frame.shape[1])
                    height = int(prediction[3] * frame.shape[0])
                    left = int(center_x - width / 2)
                    top = int(center_y - height / 2)

                    classes.append(class_id)
                    confidences.append(float(confidence))
                    boxes.append([left, top, width, height])
        indices = cv.dnn.NMSBoxes(boxes, confidences, self.confidence_threshold, self.nms_threshold)

        positions = []
        chars = []

        for index in indices:
            index = index[0]
            box = boxes[index]
            left = box[0]
            top = box[1]
            width = box[2]
            height = box[3]

            positions.append(left)

            if is_license_plate and max_confidence == confidences[index]:
                # Draw prediction rectangle for License Plate

                license_plate = frame[top: top + height, left: left + width]
                cv.rectangle(frame, (left, top), (left + width, top + height), (0, 255, 0), 3)
                # cv.imshow('License Plate', license_plate)
                # cv.imwrite('./test_bha.jpg', license_plate.astype(np.uint8))

                # Process Licence plate to cover to Gray and to enhance contours
                processed_license_plate = self.process_license_plate(license_plate)
                # cv.imshow('Processed License Plate', license_plate)
                # cv.imwrite('./test_bha2.jpg', processed_license_plate.astype(np.uint8))

                self.draw_pred(frame, self.plates_classes[0], confidences[index], left, top, left + width, top + height)
                return "", processed_license_plate
            else:
                char = self.chars_classes[classes[index]]
                chars.append(char)
                self.draw_pred(frame, char, confidences[index], left, top, left + width, top + height, color=(90, 0, 255))

        sorted_chars = [x for _, x in sorted(zip(positions, chars))]
        return "".join(sorted_chars), frame

    def cnn_char_recognition(self, image):
        gray_char = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
        gray_char = cv.resize(gray_char, (75, 100))
        image = gray_char.reshape((1, 100, 75, 1))
        image = image / 255.0
        predictions = self.cnn_chars_model.predict(image)
        max_confidence_index = np.argmax(predictions)
        return self.chars_classes[max_confidence_index]

    def canny(self, image, sigma=0.33):
        lower = int(max(0, (1.0 - sigma) * np.median(image)))
        upper = int(min(255, (1.0 + sigma) * np.median(image)))
        edges = cv.Canny(image, lower, upper)
        return edges

    def cnn_recognize_plate(self, frame):
        gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
        thresh_inv = cv.adaptiveThreshold(gray, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY_INV, 39, 1)
        edges = self.canny(thresh_inv)

        contours, _ = cv.findContours(edges.copy(), cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)

        sorted_ctrs = sorted(contours, key=lambda x: cv.boundingRect(x)[0])
        area = frame.shape[0] * frame.shape[1]

        chars = []
        for i, ctr in enumerate(sorted_ctrs):
            x, y, w, h = cv.boundingRect(ctr)
            roi_area = w * h
            non_max_sup = roi_area / area

            if (non_max_sup >= 0.015) and (non_max_sup < 0.09):
                if (h > 1.2 * w) and (3 * w >= h):
                    char = frame[y:y + h, x:x + w]
                    chars.append(self.cnn_char_recognition(char))
                    cv.rectangle(frame, (x, y), (x + w, y + h), (90, 0, 255), 2)
        licensePlate = "".join(chars)
        return licensePlate

    @staticmethod
    def resize_license_plate(license_plate):
        scale_percent = 300  # percent of original size
        width = int(license_plate.shape[1] * scale_percent / 100)
        height = int(license_plate.shape[0] * scale_percent / 100)
        return cv.resize(license_plate, (width, height), interpolation=cv.INTER_AREA)

    def get_image_blob(self, image):
        return cv.dnn.blobFromImage(image, 1 / 255, (self.yolo_net_width, self.yolo_net_height), [0, 0, 0], 1,
                                    crop=False)

    def predict(self, input_path, output_car_path, output_license_path_original, output_license_path, video_path=None, is_cnn=False, is_image=True):
        vc = cv.VideoCapture(input_path)

        FPS = 2
        # Limit number of frames per sec
        vc.set(cv.CAP_PROP_FPS, FPS)

        if not is_image:
            vid_writer = cv.VideoWriter(video_path, cv.VideoWriter_fourcc('H','2','6','4'), 30, (
            round(vc.get(cv.CAP_PROP_FRAME_WIDTH)), round(vc.get(cv.CAP_PROP_FRAME_HEIGHT))))

        # Read Frames from the file if video, else read first frame from image
        while cv.waitKey(1) < 0:
            exists, frame = vc.read()
            if not exists:
                cv.waitKey(2000)
                print("End of frames")
                vid_writer.release()
                break

            # DETECT LICENSE PLATE

            car_image_blob = self.get_image_blob(frame)

            # Feed the input image to the Yolo Network
            self.plates_yolo_net.setInput(car_image_blob)

            # Get All Unconnected Yolo layers
            plates_yolo_layers = [self.plates_yolo_net.getLayerNames()[i[0] - 1] for i in
                                  self.plates_yolo_net.getUnconnectedOutLayers()]

            # Forward pass the input to yolov3 net and get outputs
            plates_output = self.plates_yolo_net.forward(plates_yolo_layers)

            # Remove the bounding boxes with low confidence and draw box for license plate
            license_num, processed_license_plate = self.predict_boxes(frame, plates_output)
            if is_image:
                cv.imwrite(output_license_path_original, processed_license_plate.astype(np.uint8))

            if not is_cnn:
                # IDENTIFY LICENSE PLATE NUMBER USING YOLOV3

                # Resize the license plate so we can feed it to the second trained yolo network
                resized_license_plate = self.resize_license_plate(processed_license_plate)

                license_plate_image_blob = self.get_image_blob(resized_license_plate)
                # license_plate_image_blob = np.reshape(license_plate_image_blob, (1, 3, yolo_net_width,
                # yolo_net_height))

                # Feed the input image to the Yolo Network
                self.chars_yolo_net.setInput(license_plate_image_blob)

                # Get All Unconnected Yolo layers
                chars_yolo_layers = [self.chars_yolo_net.getLayerNames()[i[0] - 1] for i in
                                     self.chars_yolo_net.getUnconnectedOutLayers()]

                # Forward pass the input to yolov3 net and get outputs
                chars_output = self.chars_yolo_net.forward(chars_yolo_layers)

                license_number, processed_license_plate = self.predict_boxes(processed_license_plate, chars_output, is_license_plate=False)
                print(license_number)

            elif is_cnn:
                # IDENTIFY LICENSE PLATE NUMBER USING CNN
                license_number = self.cnn_recognize_plate(processed_license_plate)

            if is_image:
                cv.imwrite(output_license_path, processed_license_plate.astype(np.uint8))
                cv.imwrite(output_car_path, frame.astype(np.uint8))
                return license_number
            else:
                vid_writer.write(frame.astype(np.uint8))

if __name__ == "__main__":
    Predictor()

# license = Predictor().predict('/Users/moni/projects/LicensePlateDetector/m657_1.jpg', False)
# print(license)


"""
Created on Mar 25, 2013

@author: Markovtsev Vadim <v.markovtsev@samsung.com>

███████████████████████████████████████████████████████████████████████████████

Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied.  See the License for the
specific language governing permissions and limitations
under the License.

███████████████████████████████████████████████████████████████████████████████
"""

import logging
import numpy
from .library import Library
from .formatters import Formatters
from .explorer import Explorer

class SetupFeaturesFailedException(Exception):
    """Raised when setup_features_extraction() return null.
    """
    pass

class ExtractionFailedException(Exception):
    """Raised when extract_sound_features() return status code different
    from 0.
    """
    pass

class Extractor(object):
    """
    Audio feature extractor.
    """

    logger = logging.getLogger("sfm.Extractor")

    def __init__(self, features, buffer_size, sampling_rate, channels=1):
        self._config = None
        self.features = features
        self.features_dict = {f.name: f for f in self.features}
        self.buffer_size = buffer_size
        self.sampling_rate = sampling_rate
        self.channels = channels
        flen = len(self.features)
        fstrs = Library().new("char*[]", flen)
        # prevent from garbage collecting fstrs contents
        fstrs_ref = [None] * flen
        for i, f in enumerate(self.features):
            fstrs[i] = fstrs_ref[i] = Library().new("char[]", f.description(
                {"sampling_rate": sampling_rate,
                 "channels": channels}).encode())
        self._config = Library().setup_features_extraction(
            fstrs, len(self.features), buffer_size, sampling_rate)
        # fstrs_ref is still alive at this point
        del fstrs_ref
        if self._config:
            self.logger.debug("Successfully set up %d features (config %s)",
                              flen, self._config)
        else:
            self.logger.error("Failed to set up features")
            raise SetupFeaturesFailedException()

    def __del__(self):
        if not self._config:
            return
        Library().destroy_features_configuration(self._config)
        self.logger.debug("Destroyed config %s", self._config)

    RAW_KEY_NAME = "RAW"

    def calculate_raw(self, buffer):
        """
        Calculates the audio features.
        """
        if not self._config:
            self.logger.error("Unable to calculate features")
            return None
        fnames = Library().new("char***")
        results = Library().new("void***")
        rlengths = Library().new("int**")
        status = Library().extract_sound_features(
            self._config, Library().cast(
                "int16_t*", buffer.__array_interface__["data"][0]),
            fnames, results, rlengths)
        self.logger.debug("extract_speech_features() returned status %d "
                          "(results = %s)", status, results[0])
        if status != 0:
            raise ExtractionFailedException()
        ret = {}
        for i in range(len(self.features)):
            length = rlengths[0][i]
            fname = Library().string(fnames[0][i]).decode()
            feature = self.features_dict[fname]
            self.logger.debug(feature.name + " yielded %d bytes", length)
            buffer = Library().buffer(results[0][i], length)
            format_name = feature.transforms[-1].output_format
            if format_name == "":
                format_name = Explorer().transforms[
                    feature.transforms[-1].name].output_format
            ret[fname] = Formatters.parse(numpy.frombuffer(
                buffer, dtype=numpy.byte, count=length), format_name)
        ret[Extractor.RAW_KEY_NAME] = results[0]
        Library().free_results(len(self.features), fnames[0],
                               Library().NULL, rlengths[0])
        self.logger.debug("Freed the service information")
        return ret

    def free_results(self, results):
        """
        Disposes the resources allocated for calculation results.
        """
        if results:
            Library().free_results(len(self.features), Library().NULL,
                                   results[Extractor.RAW_KEY_NAME],
                                   Library().NULL)
            self.logger.debug("Freed extraction results %s", results["RAW"])

    def calculate(self, buffer):
        """
        Default is to copy the data from calculate_raw().
        """
        results = self.calculate_raw(buffer)
        for name in results:
            if name != Extractor.RAW_KEY_NAME:
                results[name] = numpy.copy(results[name])
        self.free_results(results)
        del(results[Extractor.RAW_KEY_NAME])
        return results

    def report(self, file_name):
        """
        Saves the extraction report graph.
        """
        Library().report_extraction_graph(
            self._config, Library().new("char[]", file_name))