# Aziz Nechi
# If you have any trouble with submitting to S3, please contact me. This part is on me.
# As for the conversion, the fetching from url and the sound editing, I have made sure
# of its veracity.
# Please make sure the url you provide is correct, and doesn't block downloads from
# unknown sources.

# Your S3 bucket name here:
BUCKET = ''

# To read and write MP3 files
import pydub

# Math
import numpy as np

# Get requests to download mp3 or wav files
import requests

# Used to replace file operations with in memory operations (with ByteIO)
import io

# In memory copy of objects
from copy import copy

# Imports related to the use of S3
import logging
import boto3
from botocore.exceptions import ClientError

class AudioProcessing():
    def __init__(self, url=None, bytes_object=None):
        """
        In memory simple audio operation on wav files. Supports mp3 to wav conversion
        """
        if(url):
            # create file from url
            r = requests.get(url)
            bytes_object = r.content

        self.read(bytes_object)
        return

    def add_echo(self, delay):
        '''Applies an echo that is 0...<input audio duration in seconds> seconds from the beginning'''
        output_audio = copy(self.audio_data)
        output_delay = delay * self.freq

        for count in range(len(self.audio_data)):
            e = self.audio_data[count]
            output_audio[count] = e + self.audio_data[count - int(output_delay)]

        self.audio_data = output_audio

    def set_audio_speed(self, speed_factor):
        '''Sets the speed of the audio by a floating-point factor'''
        sound_index = np.round(np.arange(0, len(self.audio_data), speed_factor))
        self.audio_data = self.audio_data[sound_index[sound_index < len(self.audio_data)].astype(int)]

    def filter_frequency(self, threshold):
        output_audio = copy(self.audio_data)
        for count in range(len(self.audio_data)):
            e = self.audio_data[count]
            if ( e < threshold ):
                output_audio[count] = threshold
            else :
                output_audio[count] = e
        self.audio_data = output_audio

    def custom_filter(self, threshold):
        output_audio = copy(self.audio_data)
        for count in range(len(self.audio_data)):
            e = self.audio_data[count]
            output_audio[count] = (e / threshold)

        self.audio_data = output_audio

    def set_volume(self, level):
        '''Sets the overall volume of the data via floating-point factor'''
        output_audio = copy(self.audio_data)
        for count in range(len(self.audio_data)):
            e = self.audio_data[count]
            output_audio[count] = (e * level)

        self.audio_data = output_audio

    def deepen(self, factor=0.85):
        self.freq = int(self.freq * factor)

    def pitch(self, factor=0.9):
        self.freq = int(self.freq / factor)

    def hide_voice (self):
        self.set_volume(1.3)
        #self.custom_filter(1.3)
        self.add_echo(delay=0.02)
        self.deepen(factor=0.8)
        output = io.BytesIO()
        return self.write(output)

    def to_bytes(self):
        """
        Returns io.BytesIO stream of the edited mp3
        """
        output = io.BytesIO()
        return self.write(output)

    def read(self, f, normalized=False):
        """MP3 to numpy array"""
        a = pydub.AudioSegment.from_mp3(io.BytesIO(f))
        y = np.array(a.get_array_of_samples())
        if a.channels == 2:
            y = y.reshape((-1, 2))
        if normalized:
            self.audio_data =  np.float32(y) / 2**15
            self.freq = a.frame_rate
            return a.frame_rate, np.float32(y) / 2**15
        else:
            self.audio_data = y
            self.freq = a.frame_rate
            return a.frame_rate, y

    def write(self, f, normalized=False):
        """numpy array to MP3"""
        sr = self.freq
        x = self.audio_data
        channels = 2 if (x.ndim == 2 and x.shape[1] == 2) else 1
        if normalized:  # normalized array - each item should be a float in [-1, 1)
            y = np.int16(x * 2 ** 15)
        else:
            y = np.int16(x)
        song = pydub.AudioSegment(y.tobytes(), frame_rate=sr, sample_width=2, channels=channels)
        song.export(f, format="mp3", bitrate="320k")
        return f

def upload_file_to_S3(file_name, bucket, object_name=None):
    """Upload a file to an S3 bucket

    :param file_name: File to upload
    :param bucket: Bucket to upload to
    :param object_name: S3 object name. If not specified then file_name is used
    :return: True if file was uploaded, else False
    """

    # If S3 object_name was not specified, use file_name
    if object_name is None:
        object_name = file_name

    # Upload the file
    s3_client = boto3.client('s3')
    try:
        response = s3_client.upload_file(file_name, bucket, object_name)
        return response
    except ClientError as e:
        print("S3 error occured "+ e)
        return None

if __name__ == '__main__':

    # print("URL from of mp3 file: ", end="")
    # YOUR_URL = str(input())

    YOUR_URL = debug_url
    x = AudioProcessing(bytes_object="who-are-you.mp3")
    # x.set_audio_speed(1.1)

    # Deepened voice :
    x.deepen()

    # Pitched voice :
    # x.pitch()

    result = x.to_bytes()
    # upload_file_to_S3(result, BUCKET, "voice_changed.mp3")

    with open("slim.mp3", "wb") as f:
        f.write(result.getbuffer())


#! /usr/bin/env python
# -*- coding: utf-8 -*-

#################################
#
#   MAIN
#
#################################
if  __name__ == "__main__":

    import argparse
    import numpy as np

    from pysedm import get_sedmcube, io

    # ================= #
    #   Options         #
    # ================= #
    parser = argparse.ArgumentParser(
        description=""" run the interactive plotting of a given cube""",
        formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument('infile', type=str, default="None",
                        help='cube filepath')

    # // AUTOMATIC EXTRACTION
    parser.add_argument('--inputcube',   action="store_true", default=False,
                        help='Is the input a cube ?')

    parser.add_argument('--reducer',  type=str, default=io.SEDM_REDUCER,
                        help='What is your name? This will be added in the header [default: $SEDM_REDUCER]')

    #  which extraction
    parser.add_argument('--aperture',  type=str, default=None,
                        help='Do you want a simple aperture extraction ?')

    parser.add_argument('--ap_bkgdscale',  type=float, default=1.4,
                        help='Default width scale of the outter background ring [default 1.4]')

    # Auto PSF
    parser.add_argument('--auto',  type=str, default=None,
                        help='Shall this run an automatic PSF extraction')

    parser.add_argument('--autorange',  type=str, default="4000,8000",
                        help='Wavelength range [in Angstrom] for measuring the metaslice PSF')

    parser.add_argument('--autobins',  type=int, default=7,
                        help='Number of bins within the wavelength range (see --autorange)')

    parser.add_argument('--buffer',  type=float, default=10,
                        help='Radius [in spaxels] of the aperture used for the PSF fit. (see --centroid for aperture center)')

    parser.add_argument('--psfmodel',  type=str, default="NormalMoffatTilted",
                        help='PSF model used for the PSF fit: NormalMoffat{Flat/Tilted/Curved}')

    parser.add_argument('--seeing',  type=float, default=2,
                        help='NOT READY YET Expected seeing (not quite in arcsec). This helps the PSF model to converge.')

    # Contour separation with contsep.py
    parser.add_argument('--contsep',  action="store_true", default=False,
                        help='Shall this run contsep to avoid non-target spaxels within the fit.')

    parser.add_argument('--contsep_offset', nargs=2, type=float, default=[0, 0],
                        help='Offset between reference image and cube image for the contsep module.')

    parser.add_argument('--contsep_fakemag', type=float, default=16.0,
                        help="Fake magnitude of the target for a reference image.")

    parser.add_argument('--contsep_forcedmag', type=float, default=0.0,
                        help="Addition to contsep_mag to select more host spaxels.")

    # Cosmic ray removal with byecr.py
    parser.add_argument('--byecr', action="store_true", default=False,
                         help="Shall this run byecr to remove cosmic rays at a specific spaxel at a specific wavelengh.")

    parser.add_argument('--byecr_lbda', type=float, default=None,
                         help="To check cosmic rays at a given wavelength.")

    parser.add_argument('--byecr_cut', type=float, default=5.0,
                         help="Cut criteria for byecr module.")

    parser.add_argument('--byecr_wspectral', action="store_true", default=False,
                         help="Also use a spectral filtering for byecr.")

    parser.add_argument('--byecr_showcube', action="store_true", default=False,
                         help="Show a cube with detected cosmic rays.")

    # Centroid
    parser.add_argument('--centroid',  type=str, default="auto", nargs="+",
                        help='Where is the point source expected to be ?'+
                            "\nYou have 3 potential cases:"+
                            "\n   * --centroid auto: Use astrometry from guider if possible or falls back to the 'brightest' mode"+
                            "\n   * --centroid brightest: Use the position of the bightest spaxels "+
                            "\n   * --centroid x0,y0: Use the given x0 y0 position (separated by a coma)."
                            )

    parser.add_argument('--centroiderr',  type=str, default="None",nargs=2,
                        help='Where is the point source location guess boundaries ?'+
                        "\nYou have 2 potential cases:"+
                            "\n   * --centroiderr None: [3 3] if position from the guider, [5 5] if from the 'brightest' mode "+
                            "\n   * --centroiderr [X Y]: use the given values")

    # wavelength step
    parser.add_argument('--display',  action="store_true", default=False,
                        help='Select the area to fit using the display function.')

    parser.add_argument('--vmin',  type=str, default="2",
                        help='Data Percentage used for imshow "vmin" when using the --display mode')

    parser.add_argument('--vmax',  type=str, default="98",
                        help='Data Percentage used for imshow "vmax" when using the --display mode')

    # Flux calibration
    parser.add_argument('--nofluxcal',  action="store_true", default=False,
                        help='No Flux calibration')

    parser.add_argument('--fluxcalsource',  type=str, default="None",
                        help="Path to a 'fluxcal' .fits file. This file will be used for the flux calibration"+
                        "\nIf nothing given, the nearest (in time and within the night) fluxcal will be used.")

    # wavelength step
    parser.add_argument('--lstep',  type=int, default=1,
                        help='Slice width in lbda step: default is 1, use 2 for fainter source and maybe 3 for really faint target')

    # I/O options
    parser.add_argument('--tag',  type=str, default="None",
                        help='Add the tag on output filename.')

    parser.add_argument('--nofig',    action="store_true", default=False,
                        help='Do not create any plot')

    parser.add_argument('--notroncation',  action="store_true", default=False,
                        help='Do not apply the blue and red end edge troncation on flux calibrated spectra')

    # - Standard Star object
    parser.add_argument('--std',  action="store_true", default=False,
                        help='Set this to True to tell the program you what to build a calibration spectrum from this object')

    args = parser.parse_args()

    # ================= #
    #   The Scripts     #
    # ================= #

    # --------- #
    #  Date     #
    # --------- #
    if not args.inputcube:
        date = args.infile

    # --------- #
    #  Parsing  #
    # --------- #
    if len(args.centroid)==1:
        args.centroid = args.centroid[0]
    elif len(args.centroid)>2:
        raise ValueError("--centroid should be 'auto', 'brightness' or X Y")

    # ------------- #
    #               #
    # APERTURE      #
    #               #
    # ------------- #
    if args.aperture is not None and len(args.aperture) >0:
        raise NotImplementedError("--aperture option not implemented yet. Use --auto")
        if args.std:
            raise NotImplementedError("--std option cannot be used (yet) with --aperture. Use --auto")

    # ------------- #
    #               #
    # 3D PSF        #
    #               #
    # ------------- #
    if args.auto is not None and len(args.auto) >0:

        for target in args.auto.split(","):

            filecubes = io.get_night_files(date, "cube.*", target=target.replace(".fits","")) if not args.inputcube else [target]
            print("cube file from which the spectra will be extracted [auto]: "+ ", ".join(filecubes))

            # - loop over the file cube
            for filecube in filecubes:
                print("Automatic extraction of target %s, file: %s"%(target, filecube))
                # ===================== #
                # Running ExtractStar   #
                # ===================== #
                #
                # DATA IN
                #
                cube = get_sedmcube(filecube)
                cube.header['REDUCER'] = args.reducer
                #
                # OPTIONS
                #
                es_options = {
                    # Display mode ?
                    "display":args.display, "displayprop":{"vmin":args.vmin, "vmax":args.vmax},

                    # Metaslices
                    # -- How many slices
                    "step1range": np.asarray(args.autorange.split(","), dtype="float"),
                    "step1bins": args.autobins,
                    # -- Centroid guess
                    "centroid":args.centroid, "prop_position":{"centroiderr":args.centroid},
                    # Sub IFU
                    "spaxelbuffer": args.buffer,
                    # PSF Model
                    "psfmodel": args.psfmodel,
                    "spaxels_to_use":None, # not ready yet, provide here directly the spaxels to be used
                    "spaxels_to_avoid":None, # not ready yet, provide here directly the spaxels to be used
                    "fwhm_guess": args.seeing,
                    # Spectral options
                    "slice_width": int(args.lstep),
                    #
                    #
                    "verbose":True,
                    }

                # This means that you want to remove non-target spaxel contamination
                if args.contsep:
                    from pysedm import contsep
                    print(" Starting contsep non-target spaxel removal ".center(50, "-"))
                    targetid = io.filename_to_id(filecube)
                    cont = contsep.get_spaxels_from_constsep(date, targetid,
                                                             offset=args.contsep_offset,
                                                             fake_mag=args.contsep_fakemag,
                                                             forced_mag=args.contsep_forcedmag)
                    es_options["spaxels_to_avoid"] =  list( cont.get_others_spaxels(spaxels_id=False) )

                # Remove cosmic ray at a specific spaxel at a given wavelength.
                if args.byecr:
                    from pysedm import byecr
                    print(" Starting byecr cosmic ray removal ".center(50, "-"))
                    targetid = io.filename_to_id(filecube)
                    byecrclass = byecr.get_cr_spaxels_from_byecr(date, targetid)
                    cr_df = byecrclass.get_cr_spaxel_info(lbda_index=args.byecr_lbda,
                                                          cut_criteria=args.byecr_cut,
                                                          wspectral=args.byecr_wspectral)

                    cube.data[cr_df["cr_lbda_index"], cr_df["cr_spaxel_index"]] = np.nan

                    print("Number of detected cosmic rays = %i" %len(cr_df))
                    print("Number of cosmic ray affected spaxels = %i" %len(np.unique(cr_df["cr_spaxel_index"])))

                # ===================== #
                # SOURCE EXTRACTION     #
                # ===================== #
                print(" Starting extract_pointsource ".center(50, "-"))
                try:
                    es_out = cube.extract_pointsource(**es_options)
                except IOError:
                    # Astrometry was not solved
                    print("ERROR: Astrometry file not found, "
                          "using brightest spaxel instead")
                    es_options["centroid"] = "brightest"
                    es_out = cube.extract_pointsource(**es_options)
                except AttributeError:
                    # Astrometry solved, but target outside IFU
                    print("ERROR: Astrometry places target outside IFU, "
                          "using brightest spaxel instead")
                    es_options["centroid"] = "brightest"
                    es_out = cube.extract_pointsource(**es_options)
                    cube.header['QUALITY'] = 3
                # -> The output object
                es_object = cube.extractstar

                if args.contsep:
                    es_object.raw_spectrum.header.set("CONTMAG", cont.target_contsep_mag,
                                                          "contsep splitting magnitude")
                    es_object.raw_spectrum.header.set("NCONTSPX",
                                                len(es_options["spaxels_to_avoid"]),
                                                "total number of non-target spaxels from contsep")

                if args.byecr:
                    es_object.raw_spectrum.header.set("NCR", len(cr_df),
                                                "total number of detected cosmic-rays from byecr")
                    es_object.raw_spectrum.header.set("NCRSPX",
                                                len(np.unique(cr_df["cr_spaxel_index"])),
                                                "total number of cosmic-ray affected spaxels")

                # -
                # - SAVING
                # -
                add_tag = "_%s" % args.tag if args.tag is not None and args.tag not in ["None", ""] else ""
                add_info_spec = "_notfluxcal" if not es_object.is_spectrum_fluxcalibrated() else ""
                spec_info = "_lstep%s" % es_options["slice_width"] + add_info_spec

                print("Tag added", add_tag)
                es_object.writeto(basename=None, add_tag="auto" + add_tag, add_info=None)

                # -
                # - PLOTTING
                # -
                plot_tag = "_auto" + add_tag + spec_info + "_"
                if not args.nofig:
                    print(" Saving figures ".center(50, "-"))
                    es_object.show_adr(            savefile=es_object.basename.replace("{placeholder}","adr_fit" + plot_tag))
                    es_object.show_metaslices(     savefile=es_object.basename.replace("{placeholder}","metaslices" + plot_tag))
                    es_object.show_extracted_spec( savefile=es_object.basename.replace("{placeholder}","extracted_spec" + plot_tag))
                    es_object.show_mla(            savefile=es_object.basename.replace("{placeholder}","ifu_spaxels_source" + plot_tag))
                    es_object.show_psf(            savefile=es_object.basename.replace("{placeholder}","psfprofile" + plot_tag), sliceid=2)
                    es_object.spectrum.show(       savefile=es_object.basename.replace("{placeholder}","spec" + plot_tag))
                    if args.contsep:
                        cont.show_ifudata( wcontour=False, wtargetspaxel=True, wotherspaxel=True,
                        savefile=es_object.basename.replace("{placeholder}","contsep" + plot_tag) )
                    if args.byecr_showcube:
                        byecrclass.show_cr_spaxels( wcube=args.byecr_showcube, lbda_index=args.byecr_lbda, cut_criteria=args.byecr_cut, wspectral=args.byecr_wspectral,
                        savefile=es_object.basename.replace("{placeholder}","byecr" + plot_tag) )

                # -
                # - Standard Specific
                # -
                if args.std:
                    print("Standard Star object, starting estimation of flux calibration")
                    #
                    # Flux Calibrator
                    #
                    speccal, fl = es_object.get_fluxcalibrator(persecond=True, troncate_edges=False)
                    if speccal is None:
                        print("WARNING: Getting the Flux Calibrator Failed.")
                    else:
                        speccal.writeto(speccal.filename.replace("fluxcal", "fluxcal" + plot_tag))

                    #
                    # PLOTTING
                    #
                    if not args.nofig:
                        if cube.header['IMGTYPE'].lower() in ['standard'] and es_object.is_spectrum_fluxcalibrated():
                            from pysedm.fluxcalibration import show_fluxcalibrated_standard
                            show_fluxcalibrated_standard(es_object.spectrum, savefile=es_object.basename.replace("{placeholder}","calibcheck_spec" + plot_tag)+".pdf")
                            show_fluxcalibrated_standard(es_object.spectrum, savefile=es_object.basename.replace("{placeholder}","calibcheck_spec" + plot_tag)+".png")
                        if fl is not None:
                            fl.show(savefile=speccal.filename.replace(".fits",".pdf"),
                                        show=False)
                            fl.show(savefile=speccal.filename.replace(".fits",".png"),
                                        show=False)


import signal
import sys
from random import randint

import numpy as np
import pyaudio
from yeelight import *
from time import sleep

FORMAT = pyaudio.paInt16
CHANNELS = 2
RATE = 44100
CHUNK = 1024

def signal_handler(sig, frame):
    bulbs[2].stop_music()
    bulbs[1].stop_music()
    bulbs[20].stop_music()
    sys.exit(0)

signal.signal(signal.SIGINT, signal_handler)

bulbs = [Bulb('192.168.1.100'), Bulb('192.168.1.101'), Bulb('192.168.1.102')]

maxValue = 2 ** 16
bars = 35
p = pyaudio.PyAudio()
# start Recording
stream = p.open(format=FORMAT, channels=CHANNELS,
                rate=RATE, input=True,
                frames_per_buffer=CHUNK)
bulbs[2].turn_on()
bulbs[2].start_music()

bulbs[1].turn_on()
bulbs[1].start_music()

bulbs[0].turn_on()
bulbs[0].start_music()

bulbs[0].set_brightness(255)
bulbs[1].set_brightness(255)
bulbs[2].set_brightness(255)

i = 0
while True:
    i = i + 1
    data = np.fromstring(stream.read(1024), dtype=np.int16)
    dataL = data[0::2]
    dataR = data[1::2]
    peakL = np.abs(np.max(dataL) - np.min(dataL)) / maxValue
    peakR = np.abs(np.max(dataR) - np.min(dataR)) / maxValue
    lString = "#" * int(peakL * bars) + "-" * int(bars - peakL * bars)
    rString = "#" * int(peakR * bars) + "-" * int(bars - peakR * bars)
    bulbs[0].set_rgb(255, 0, 128)
    bulbs[1].set_rgb(255, 0, 128)
    bulbs[2].set_rgb(255, 0, 128)

    rndm = randint(1, 2)

    if (peakL * 1000 / 2 > 70):
        bulbs[0].set_rgb(0, 255, 0)
        bulbs[1].set_rgb(0, 255, 0)
        bulbs[2].set_rgb(0, 255, 0)
        bulbs[0].set_brightness(peakL * 1000 / 2)
        bulbs[1].set_brightness(peakL * 1000 / 2)
        bulbs[2].set_brightness(peakL * 1000 / 2)
        i = 0
    elif (peakL * 1000 / 2 > 100):

        bulbs[0].set_rgb(255, 255, 255)
        bulbs[1].set_rgb(255, 255, 255)
        bulbs[2].set_rgb(255, 255, 255)

    print(peakL * 1000 / 2)
    #bulbs[0].set_brightness(peakL * 1000 / 2)
    #bulbs[1].set_brightness(peakL * 1000 / 2)
    #bulbs[2].set_brightness(peakL * 1000 / 2)
    i = 0


#!/usr/bin/env python3

from __future__ import print_function
import sys
import os
from array import array
import cv2
import numpy as np
import time
import openncc as ncc
import struct
import math

min_score = 0.8
media_head = 'iII13I'

def coordinate_is_valid(x1, y1, x2, y2):
    if ((x1 < 0) or (x1 > 1)):
        return False
    if ((y1 < 0) or (y1 > 1)):
        return False
    if ((x2 < 0) or (x2 > 1)):
        return False
    if ((y2 < 0) or (y2 > 1)):
        return False
    if ((x1 >= x2) or (y1 >= y2)):
        return False
    # print("coordinate_is_valid true")
    return True

def print_hex(bytes):
    l = [hex(int(i)) for i in bytes]
    print(" ".join(l))

def get_header_info(spec):
    frames = struct.unpack(media_head, bytes(spec))
    return frames[0], frames[1], frames[2]

def calc_similar(face_vec, standardemb):
	dot_out = 0
	vec_sum = 0
	for i in range(256):
		vec_sum += face_vec[i] * face_vec[i];

	vec_sum = math.sqrt(vec_sum);
	for i in range(256):
		face_vec[i] = face_vec[i] / vec_sum;

	dot_out = 0;
	for i in range(256):
		dot_out += face_vec[i] * standardemb[i]

	return dot_out

faceLib = [
-0.054260,
0.968262,
-0.499023,
1.067383,
1.729492,
-1.036133,
0.794922,
-0.025360,
-0.373535,
1.203125,
0.597656,
-0.909668,
0.729980,
-2.687500,
1.304688,
-0.976562,
0.211182,
-0.124268,
-0.927734,
-0.179321,
-1.613281,
-0.815918,
-1.208008,
0.757812,
0.097229,
-1.232422,
1.401367,
0.678223,
0.511230,
-2.101562,
0.104736,
0.359131,
-1.122070,
-0.223022,
-0.830078,
0.831055,
0.697266,
-0.326172,
-1.934570,
1.367188,
-0.572754,
-0.148193,
-1.887695,
0.146118,
0.627930,
-2.746094,
-1.091797,
-0.540039,
-1.600586,
-0.174561,
-1.806641,
1.315430,
0.875977,
2.404297,
-1.711914,
-0.583984,
-0.020615,
-1.338867,
0.907227,
-1.554688,
1.161133,
1.010742,
-0.660645,
0.286865,
-0.551758,
-0.156006,
0.717285,
0.762695,
0.582031,
-0.655273,
1.076172,
-0.012634,
-1.116211,
-0.598633,
-0.724121,
-0.318848,
-0.865234,
-0.186768,
-1.458008,
0.981445,
0.033752,
0.625000,
1.940430,
-0.059143,
0.468262,
1.219727,
0.780762,
-1.144531,
0.183228,
0.107300,
0.169678,
1.334961,
-0.533691,
0.596680,
1.090820,
0.819824,
-1.291992,
-0.210083,
-1.312500,
1.328125,
0.074402,
1.840820,
1.423828,
-0.952148,
-0.435547,
-1.138672,
1.110352,
0.064575,
1.963867,
0.922852,
1.352539,
0.506836,
0.667969,
0.379883,
0.270996,
0.404297,
1.996094,
0.586426,
-0.577148,
-0.348633,
-0.183472,
2.597656,
1.267578,
0.114990,
-0.166992,
-0.837891,
0.866699,
1.536133,
-0.417236,
-0.618652,
-0.160522,
0.309082,
1.122070,
0.682129,
1.201172,
0.149536,
-0.988281,
-0.074097,
-0.661133,
-1.432617,
0.964355,
2.667969,
0.714844,
0.419434,
-0.346191,
-0.835938,
-0.358154,
0.773926,
0.934082,
1.106445,
-0.321289,
0.235962,
0.033875,
0.051910,
-0.373291,
0.143188,
0.878418,
0.281494,
-0.912598,
0.174072,
-0.836426,
0.053772,
-0.829590,
-0.257080,
0.883789,
-0.397949,
-0.842285,
-0.030838,
0.337158,
-0.343262,
-0.094788,
0.832031,
-0.095398,
-0.973633,
-1.229492,
1.606445,
-0.383301,
-0.528809,
-0.077332,
-0.737305,
-1.667969,
-0.333252,
1.216797,
-0.499756,
1.037109,
-0.812012,
2.025391,
-0.271240,
0.761719,
0.614746,
0.728027,
0.122437,
0.248779,
-0.465820,
0.947266,
-0.069214,
-0.401611,
-0.135864,
-0.512695,
1.193359,
-0.024048,
-0.530273,
1.072266,
0.209839,
1.353516,
0.825684,
-0.003265,
1.353516,
0.270020,
-0.050385,
-1.509766,
0.148193,
-0.539551,
0.088074,
0.062073,
-0.111877,
-1.545898,
-1.327148,
-1.704102,
0.880859,
-1.027344,
-0.336182,
-0.052551,
1.624023,
-0.321045,
0.809082,
0.913574,
-0.441162,
0.329590,
0.252197,
0.077759,
-0.316406,
-1.588867,
-2.382812,
-0.415283,
-0.748535,
1.193359,
-0.019989,
0.984375,
-1.216797,
1.408203,
-1.426758,
-1.748047,
-1.165039,
-0.017349,
1.116211,
0.912598,
1.005859,
0.542480,
-0.289551,
0.886230,
-0.002472,
-1.326172,
-1.346680,
0.936035,
-1.043945,
]
SIMI_VAL=7.0
def main():
    global  faceLib
    res = ncc.load_fw("./moviUsbBoot", "fw/flicRefApp.mvcmd")
    if res < 0:
        printf('load firmware error!')
        sys.exit(1)

    print("get usb %d sdk versin %s" % (ncc.get_usb_version(), ncc.get_sdk_version()))

    print("get fw version: %s and ncc id %s" % (ncc.camera_get_fw_version(),
                                                ncc.camera_get_ncc_id()))

    sensors = ncc.CameraSensor()
    sensor1 = ncc.SensorModesConfig()
    if sensors.GetFirstSensor(sensor1) == 0:
        print("camera: %s, %dX%d@%dfps, AFmode:%d, maxEXP:%dus,gain[%d, %d]\n" % (
            sensor1.moduleName, sensor1.camWidth, sensor1.camHeight, sensor1.camFps,
            sensor1.AFmode, sensor1.maxEXP, sensor1.minGain, sensor1.maxGain))

    sensor2 = ncc.SensorModesConfig()
    while sensors.GetNextSensor(sensor2) == 0:
        print("camera: %s, %dX%d@%dfps, AFmode:%d, maxEXP:%dus,gain[%d, %d]\n" % (
            sensor2.moduleName, sensor2.camWidth, sensor2.camHeight, sensor2.camFps,
            sensor2.AFmode, sensor2.maxEXP, sensor2.minGain, sensor2.maxGain))

    ncc.camera_select_sensor(0)  # 0 1080p 1 4k
    cameraCfg = sensor1

    net1_info = ncc.Network1Par()
    net1_info.inputFormat = ncc.IMG_FORMAT_BGR_PLANAR
    net1_info.meanValue = [0.0, 0.0, 0.0]
    net1_info.stdValue = 1

    net1_info.isOutputYUV = 1
    net1_info.isOutputH26X = 0
    net1_info.isOutputJPEG = 0

    net1_info.imageWidth = cameraCfg.camWidth
    net1_info.imageHeight = cameraCfg.camHeight
    net1_info.startX = 0
    net1_info.startY = 0
    net1_info.endX = cameraCfg.camWidth
    net1_info.endY = cameraCfg.camHeight
    net1_info.inputDimWidth = 300
    net1_info.inputDimHeight = 300

    # extInputs = np.zeros(ncc.MAX_EXTINPUT_SIZE,dtype = np.uint8)
    # print('input size {}'.format(extInputs.size))
    net1_info.extInputs = [0] * ncc.MAX_EXTINPUT_SIZE  # tobytes()
    net1_info.modelCascade = 1
    net1_info.inferenceACC = 0

    net2_info = ncc.Network2Par()
    net2_info.startXAdj = -5
    net2_info.startYAdj = -5
    net2_info.endXAdj = 5
    net2_info.endYAdj = 5

    labelMask = [0]*ncc.MAX_LABEL_SIZE
    labelMask[1]=1
    net2_info.labelMask = labelMask
    net2_info.minConf = 0.7

    net2_info.inputDimWidth = 128
    net2_info.inputDimHeight = 128
    net2_info.inputFormat = ncc.IMG_FORMAT_BGR_PLANAR

    net2_info.meanValue = [0.0, 0.0, 0.0]
    net2_info.stdValue = 1

    count=ncc.MAX_EXTINPUT_SIZE//2
    extInputs = array('h',[0]*count)
    net2_info.extInputs = extInputs.tobytes()
    net2_info.modelCascade = 0
    #print("input  0={} 1={} type={}".format(int(ncc.f32Tof16(1.0)),extInputs[1],type(extInputs[1])));
    blob1 = "./blob/2020.3/face-detection-retail-0004/face-detection-retail-0004.blob";
    blob2 = "./blob/2020.3/face-reidentification-retail-0095/face-reidentification-retail-0095.blob";
    ret = ncc.sdk_net2_init(None, None, \
                            blob1, net1_info, 0, \
                            blob2, net2_info, 0)
    metasize = 3 * 1024 * 1024
    print("xlink_init ret=%d  %d" % (ret, metasize))
    if (ret < 0):
        return

    oft_x = net1_info.startX
    oft_y = net1_info.startY
    dis_w = net1_info.endX - net1_info.startX
    dis_h = net1_info.endY - net1_info.startY

    offset = struct.calcsize(media_head)  # 64
    size = cameraCfg.camWidth * cameraCfg.camHeight * 2
    yuvbuf = bytearray(size + offset)
    metabuf = bytearray(metasize)

    ncc.camera_video_out(ncc.YUV420p, ncc.VIDEO_OUT_CONTINUOUS)
    bmeta=False
    bsave_face=False
    if os.access("face.txt", os.F_OK):
    	faceLib=np.loadtxt('face.txt', dtype=float).tolist()
    	print(faceLib)

    while (True):
        size = ncc.GetYuvData(yuvbuf)
        if (size <= 0):
            time.sleep(0.1)  # 0.1 second
            continue

        numarry = np.array(yuvbuf[offset:size])  # skip head frameSpecOut 64 bytes
        # print("buf   len=%d/%d" % (numarry.size,size))
        yuv = numarry.reshape((int(cameraCfg.camHeight * 3 / 2), cameraCfg.camWidth))
        bgr = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR_I420, 3)

        size, outsize0,outsize1,item_num = ncc.GetMetaDataExt(metabuf)
        if (size > 0)and (item_num>0):
            bmeta=True
            print("ret=%d size0=%d size1=%d,num=%d" %(size,outsize0,outsize1,item_num ))
            # spec=metabuf[0:offset]
            #type1,seq,m_size=get_header_info(metabuf[0:offset])
            #print("meta type=%d seq=%d  size=%d" % (type1,seq,m_size))

            #f = open("sample.txt", "wb")
            #f.write(metabuf)
            #f.close()

            from_i=offset + ncc.OUTPUT_INDEX_SIZE
            count=outsize0//2
            detMetadata = struct.unpack('H'*count,metabuf[from_i:from_i+outsize0])

            from_i+=outsize0   #skip output 0
            count=outsize1//2
            print("out2 count %d from %d to %d" % (count,from_i,size))
            secondMeta = struct.unpack('H'*count*item_num,metabuf[from_i:size])

            for i in range(item_num):
                image_id = int(ncc.f16Tof32(detMetadata[i * 7 + 0]))
                if (image_id < 0):
                    break

                label = int(ncc.f16Tof32(detMetadata[i * 7 + 1]))
                score = ncc.f16Tof32(detMetadata[i * 7 + 2])
                x0 = ncc.f16Tof32(detMetadata[i * 7 + 3])
                y0 = ncc.f16Tof32(detMetadata[i * 7 + 4])
                x1 = ncc.f16Tof32(detMetadata[i * 7 + 5])
                y1 = ncc.f16Tof32(detMetadata[i * 7 + 6])
               # print("item sise=%d score:%.2f<->min:%.2f  rec:(%.3f,%.3f)<->(%.3f,%.3f) " %(item_num,score,min_score,x0,y0,x1,y1))
                if ((not coordinate_is_valid(x0, y0, x1, y1)) or (score < min_score) or (labelMask[label]==0)):
                    continue

                x = int(x0 * dis_w + oft_x)
                y = int(y0 * dis_h + oft_y)
                w = int((x1 - x0) * dis_w)
                h = int((y1 - y0) * dis_h)

                cv2.rectangle(bgr, (x, y), (x + w, y + h), (255, 128, 128), 2)
                regMetadata=secondMeta[i*count:i*count+count]

                regRet = array('f')

                for j in range(256):
                	regRet.append(ncc.f16Tof32(regMetadata[j]))

                if 	bsave_face:
                		print("!!!!save face!!!!")
                		cv2.imwrite('face.png', bgr)
                		nf=np.array(regRet)
                		np.savetxt('face.txt', np.array(nf), fmt='%.2f')
                		bsave_face=False

                similarVal = calc_similar(regRet, faceLib)
                result = ''
                print("get similar {:.2f}".format(similarVal))
                if similarVal > SIMI_VAL:
                	result = 'RET{:2.2f}, OK'.format(similarVal);
                else:
                	result = 'RET{:2.2f}, NG'.format(similarVal);
                cv2.putText(bgr, result, (x, y - 20), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 1)

        img_scaled = cv2.resize(bgr, None, fx=0.7, fy=0.7, interpolation=cv2.INTER_CUBIC)
        cv2.namedWindow('openncc', cv2.WINDOW_AUTOSIZE)
        cv2.imshow('openncc', img_scaled)
        c=cv2.waitKey(10)
        if c==115:    #s
        	bsave_face=True
        elif ( c== 27):
            break
    ncc.sdk_uninit()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    sys.exit(main() or 0)


"""Training script for the DeepLab-LargeFOV network on the PASCAL VOC dataset
   for semantic image segmentation.

This script trains the model using augmented PASCAL VOC dataset,
which contains approximately 10000 images for training and 1500 images for validation.
"""

from __future__ import print_function

import argparse
import os
import time
import math
import scipy.misc
import matplotlib

matplotlib.use('Agg')
import matplotlib.pyplot as plt

import tensorflow as tf
import numpy as np
import scipy.misc

from nets.large_fov.image_reader import ImageReader
from nets.large_fov.model import DeepLabLFOVModel
from nets.large_fov.utils import decode_labels

BATCH_SIZE = 10
DATA_DIRECTORY = '/home/automan/Data/Pascal/VOC2012'
DATA_LIST_PATH = './dataset/val.txt'
INPUT_SIZE = '321,321'

MEAN_IMG = tf.Variable(np.array((104.00698793,116.66876762,122.67891434)), trainable=False, dtype=tf.float32)
RANDOM_SCALE = False
RESTORE_FROM = 'snapshots/'
SAVE_DIR = 'images_val/'
WEIGHTS_PATH = None

IMG_MEAN = np.array((104.00698793,116.66876762,122.67891434), dtype=np.float32)

def get_arguments():
    """Parse all the arguments provided from the CLI.

    Returns:
      A list of parsed arguments.
    """
    parser = argparse.ArgumentParser(description="DeepLabLFOV Network")
    parser.add_argument("--batch_size", type=int, default=BATCH_SIZE,
                        help="Number of images sent to the network in one step.")
    parser.add_argument("--data_dir", type=str, default=DATA_DIRECTORY,
                        help="Path to the directory containing the PASCAL VOC dataset.")
    parser.add_argument("--data_list", type=str, default=DATA_LIST_PATH,
                        help="Path to the file listing the images in the dataset.")
    parser.add_argument("--input_size", type=str, default=INPUT_SIZE,
                        help="Comma-separated string with height and width of images.")
    parser.add_argument("--restore_from", type=str, default=RESTORE_FROM,
                        help="Where restore checkpoint from.")
    parser.add_argument("--weights_path", type=str, default=WEIGHTS_PATH,
                        help="Path to the file with caffemodel weights. "
                            "If not set, all the variables are initialised randomly.")
    parser.add_argument("--save_dir", type=str, required=True,
                        help="Where to save figures with predictions.")
    return parser.parse_args()

def save(saver, sess, logdir, step):
    model_name = 'model.ckpt'
    checkpoint_path = os.path.join(logdir, model_name)

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

    saver.save(sess, checkpoint_path, global_step=step)
    print('The checkpoint has been created.')

def load(loader, sess, ckpt_path):
    '''Load trained weights.

    Args:
      loader: TensorFlow saver object.
      sess: TensorFlow session.
      ckpt_path: path to checkpoint file with parameters.
    '''
    loader.restore(sess, ckpt_path)
    print("Restored model parameters from {}".format(ckpt_path))

def main():
    """Create the model and start the training."""
    args = get_arguments()
    """
    h, w = map(int, args.input_size.split(','))
    input_size = None  # (h, w)
    """
    # Create queue coordinator.
    coord = tf.train.Coordinator()

    # Load reader.
    with tf.name_scope("create_inputs"):
        reader = ImageReader(
            args.data_dir,
            args.data_list,
            None,
            RANDOM_SCALE,
            coord)
        image_batch, label_batch, shape_batch = reader.dequeue(args.batch_size)

    # Create network.
    net = DeepLabLFOVModel(args.weights_path)

    pred = net.preds(image_batch)

    # Set up tf session and initialize variables.
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)
    init = tf.global_variables_initializer()
    sess.run(init)

    saver = tf.train.Saver(var_list=tf.trainable_variables())

    ckpt = tf.train.get_checkpoint_state(args.restore_from)
    load(saver, sess, ckpt.model_checkpoint_path)

    # Start queue threads.
    threads = tf.train.start_queue_runners(coord=coord, sess=sess)

    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)
    if not os.path.exists(args.save_dir + 'mask/'):
        os.makedirs(args.save_dir + 'mask/')
    if not os.path.exists(args.save_dir + 'pred/'):
        os.makedirs(args.save_dir + 'pred/')

    # Iterate over training steps.
    for step in range(int(math.ceil(reader.image_num / args.batch_size))):
        start_time = time.time()
        images, labels, shapes, preds = sess.run([image_batch, label_batch, shape_batch, pred])
        for i in range(len(preds)):
            shape = shapes[i]
            label = (labels[i])[:shape[0], :shape[1], :]
            prediction = (preds[i])[:shape[0], :shape[1], :]
            scipy.misc.imsave(args.save_dir + 'mask/' + str(step * args.batch_size + i) + '.png',
                              np.squeeze(label, axis=2))
            scipy.misc.imsave(args.save_dir + 'pred/' + str(step * args.batch_size + i) + '.png',
                              np.squeeze(prediction, axis=2))

            # write readable images into /output/
            vimage = (images[i])[:shape[0], :shape[1], :]
            vlabel = decode_labels(labels[i, :, :, 0])[:shape[0], :shape[1], :]
            vprediction = decode_labels(preds[i, :, :, 0])[:shape[0],:shape[1],:]
            scipy.misc.imsave(args.save_dir + 'output/' + str(step * args.batch_size + i) + '_i.png',
                              vimage)
            scipy.misc.imsave(args.save_dir + 'output/' + str(step * args.batch_size + i) + '_m.png',
                              vlabel)
            scipy.misc.imsave(args.save_dir + 'output/' + str(step * args.batch_size + i) + '_p.png',
                              vprediction)

        duration = time.time() - start_time
        print('step {:d} \t ({:.3f} sec/step)'.format(step, duration))

    coord.request_stop()
    coord.join(threads)

if __name__ == '__main__':
    main()


from __future__ import absolute_import, unicode_literals

import base64
import hashlib
import json
import logging
import os
import pickle
import re
from builtins import str

import numpy as np
from boto.s3.connection import Key, S3Connection

from .constants import *

from btb import ParamTypes

from future import standard_library  # isort:skip
standard_library.install_aliases()

import urllib  # isort:skip

# global variable storing this machine's public IP address
# (so we only have to fetch it once)
public_ip = None

# URL which should give us our public-facing IP address
PUBLIC_IP_URL = 'http://ip.42.pl/raw'

logger = logging.getLogger('atm')

def hash_dict(dictionary, ignored_keys=None):
    """
    Hash a python dictionary to a hexadecimal string.
    http://stackoverflow.com/questions/5884066/hashing-a-python-dictionary
    """
    dictionary = dict(dictionary)  # copy dictionary
    for key in (ignored_keys or []):
        del dictionary[key]
    return hashlib.md5(repr(sorted(dictionary.items())).encode('utf8')).hexdigest()

def hash_nested_tuple(tup):
    """ Hash a nested tuple to hexadecimal """
    return hashlib.md5(repr(sorted(tup)).encode('utf8')).hexdigest()

def hash_string(s):
    """ Hash a string to hexadecimal """
    return hashlib.md5(str(s).encode('utf8')).hexdigest()

def ensure_directory(directory):
    """ Create directory if it doesn't exist. """
    if not os.path.exists(directory):
        os.makedirs(directory)

def get_public_ip():
    """
    Get the public IP address of this machine. If the request times out,
    return "localhost".
    """
    global public_ip
    if public_ip is None:
        try:
            response = urllib.request.urlopen(PUBLIC_IP_URL, timeout=2)
            data = str(response.read().strip())
            # pull an ip-looking set of numbers from the response
            match = re.search('\d+\.\d+\.\d+\.\d+', data)
            if match:
                public_ip = match.group()
        except Exception as e:  # any exception, doesn't matter what
            logger.error('could not get public IP: %s' % e)
            public_ip = 'localhost'

    return public_ip

def object_to_base_64(obj):
    """ Pickle and base64-encode an object. """
    pickled = pickle.dumps(obj)
    return base64.b64encode(pickled)

def base_64_to_object(b64str):
    """
    Inverse of object_to_base_64.
    Decode base64-encoded string and then unpickle it.
    """
    decoded = base64.b64decode(b64str)
    return pickle.loads(decoded)

def obj_has_method(obj, method):
    """http://stackoverflow.com/questions/34439/finding-what-methods-an-object-has"""
    return hasattr(obj, method) and callable(getattr(obj, method))

## Converting hyperparameters to and from BTB-compatible formats

def vector_to_params(vector, tunables, categoricals, constants):
    """
    Converts a numpy vector to a dictionary mapping keys to named parameters.

    vector: single example to convert

    Examples of the format for SVM sigmoid hyperpartition:

    tunables = [('C', HyperParameter(type='float_exp', range=(1e-05, 1e5))),
                ('degree', HyperParameter(type='int', range=(2, 4))),
                ('gamma', HyperParameter(type='float_exp', range=(1e-05, 1e5)))]

    categoricals = (('kernel', 'poly'),
                    ('probability', True),
                    ('_scale', True))

    constants = [('cache_size', 15000)]
    """
    params = {}

    # add the tunables
    for i, elt in enumerate(vector):
        key, struct = tunables[i]
        if struct.type in [ParamTypes.INT, ParamTypes.INT_EXP]:
            params[key] = int(elt)
        elif struct.type in [ParamTypes.FLOAT, ParamTypes.FLOAT_EXP]:
            params[key] = float(elt)
        else:
            raise ValueError('Unknown data type: {}'.format(struct.type))

    # add the fixed categorical settings and fixed constant values
    for key, value in categoricals + constants:
        params[key] = value

    return params

def params_to_vectors(params, tunables):
    """
    Converts a list of parameter vectors (with metadata) into a numpy array
    ready for BTB tuning.

    Args:
        params: list of hyperparameter vectors. Each vector is a dict mapping
            the names of parameters to those parameters' values.
        tunables: list of HyperParameter metadata structures describing all
            the optimizable hyperparameters that should be in each vector. e.g.

        tunables = [('C',  HyperParameter(type='float_exp', range=(1e-5, 1e5))),
                    ('degree',  HyperParameter('int', (2, 4))),
                    ('gamma',  HyperParameter('float_exp', (1e-05, 1e5)))]

    Returns:
        vectors: np.array of parameter vectors ready to be optimized by a
            Gaussian Process (or what have you).
            vectors.shape = (len(params), len(tunables))
    """
    # make sure params is iterable
    if not isinstance(params, (list, np.ndarray)):
        params = [params]

    keys = [k[0] for k in tunables]
    vectors = np.zeros((len(params), len(keys)))
    for i, p in enumerate(params):
        for j, k in enumerate(keys):
            vectors[i, j] = p[k]
    return vectors

## Serializing and deserializing data on disk

def _make_save_path_old(dir, classifier, suffix):
    """
    Generate the base save path for a classifier's model and metrics files,
    based on the classifier's dataset name and hyperparameters.
    """
    run_hash = hash_string(classifier.datarun.dataset.name)
    params_hash = hash_dict(classifier.hyperparameter_values)
    filename = "%s-%s-%s.%s" % (run_hash, params_hash,
                                classifier.datarun.description, suffix)
    return os.path.join(dir, filename)

def make_save_path(dir, classifier, suffix):
    """
    Generate the base save path for a classifier's model and metrics files,
    based on the classifier's dataset name and hyperparameters.
    """
    run_name = "".join([c for c in classifier.datarun.dataset.name
                        if c.isalnum() or c in (' ', '-', '_')]).rstrip()
    params_hash = hash_dict(classifier.hyperparameter_values)[:8]
    filename = "%s-%s.%s" % (run_name, params_hash, suffix)
    return os.path.join(dir, filename)

def save_model(classifier, model_dir, model):
    """
    Save a serialized version of a Model object for a particular classifier.
    The object will be stored at a path generated from the classifier's
    attributes.
    """
    path = make_save_path(model_dir, classifier, 'model')
    logger.info('Saving model in: %s' % path)
    with open(path, 'wb') as f:
        pickle.dump(model, f, protocol=pickle.HIGHEST_PROTOCOL)
    return path

def save_metrics(classifier, metric_dir, metrics):
    """
    Save a JSON-serialized version of a set of performance metrics for a
    particular classifier. The metrics will be stored at a path generated from
    the classifier's attributes.
    """
    path = make_save_path(metric_dir, classifier, 'metric')
    logger.info('Saving metrics in: %s' % path)
    with open(path, 'w') as f:
        json.dump(metrics, f)
    return path

def load_model(classifier, model_dir):
    """ Load the Model object for a particular classifier """
    path = make_save_path(model_dir, classifier, 'model')
    with open(path, 'rb') as f:
        return pickle.load(f)

def load_metrics(classifier, metric_dir):
    """ Load the performance metrics for a particular classifier """
    path = make_save_path(metric_dir, classifier, 'metric')
    with open(path) as f:
        return json.load(f)

## Downloading data from the web

def get_local_data_path(data_path):
    """
    given a data path of the form "s3://..." or "http://...", return the local
    path where the file should be stored once it's downloaded.
    """
    if data_path is None:
        return None, None

    m = re.match(S3_PREFIX, data_path)
    if m:
        path = data_path[len(m.group()):].split('/')
        return os.path.join(DATA_DL_PATH, path[-1]), FileType.S3

    m = re.match(HTTP_PREFIX, data_path)
    if m:
        path = data_path[len(m.group()):].split('/')
        return os.path.join(DATA_DL_PATH, path[-1]), FileType.HTTP

    return data_path, FileType.LOCAL

def download_file_s3(aws_path, aws_config, local_folder=DATA_DL_PATH):
    """ Download a file from an S3 bucket and save it in the local folder. """
    # remove the prefix and extract the S3 bucket, folder, and file name
    m = re.match(S3_PREFIX, aws_path)
    split = aws_path[len(m.group()):].split('/')
    s3_bucket = split.pop(0)
    s3_folder = '/'.join(split[:-1])
    keyname = split[-1]

    # create the local folder if necessary
    if local_folder is not None:
        ensure_directory(local_folder)
        path = os.path.join(local_folder, keyname)
    else:
        path = keyname

    if os.path.isfile(path):
        logger.warning('file %s already exists!' % path)
        return path

    conn = S3Connection(aws_config.access_key, aws_config.secret_key)
    bucket = conn.get_bucket(s3_bucket)

    if s3_folder:
        aws_keyname = os.path.join(s3_folder, keyname)
    else:
        aws_keyname = keyname

    logger.debug('downloading data from S3...')
    s3key = Key(bucket)
    s3key.key = aws_keyname
    s3key.get_contents_to_filename(path)
    logger.info('file saved at %s' % path)

    return path

def download_file_http(url, local_folder=DATA_DL_PATH):
    """ Download a file from a public URL and save it locally. """
    filename = url.split('/')[-1]
    if local_folder is not None:
        ensure_directory(local_folder)
        path = os.path.join(local_folder, filename)
    else:
        path = filename

    if os.path.isfile(path):
        logger.warning('file %s already exists!' % path)
        return path

    logger.debug('downloading data from %s...' % url)
    f = urllib.request.urlopen(url)
    data = f.read()
    with open(path, 'wb') as outfile:
        outfile.write(data)
    logger.info('file saved at %s' % path)

    return path

def download_data(train_path, test_path=None, aws_config=None):
    """
    Download a set of train/test data from AWS (if necessary) and return the
    path to the local data.
    """
    local_train_path, train_type = get_local_data_path(train_path)
    local_test_path, test_type = get_local_data_path(test_path)

    # if the data are not present locally, try to download them from the
    # internet (either an S3 bucket or a http connection)
    if not os.path.isfile(local_train_path):
        if train_type == FileType.HTTP:
            assert (download_file_http(train_path) == local_train_path)
        elif train_type == FileType.S3:
            assert (download_file_s3(train_path, aws_config=aws_config) ==
                    local_train_path)

    if local_test_path and not os.path.isfile(local_test_path):
        if test_type == FileType.HTTP:
            assert (download_file_http(test_path) == local_test_path)
        elif test_type == FileType.S3:
            assert (download_file_s3(test_path, aws_config=aws_config) ==
                    local_test_path)

    return local_train_path, local_test_path


######################################################
#
# PyRAI2MD 2 module for MOLCAS interface
#
# Author Jingbai Li
# Sep 20 2021
#
######################################################

import os, sys, shutil
import numpy as np

from PyRAI2MD.Quantum_Chemistry.qc_molcas import MOLCAS
from PyRAI2MD.Utils.coordinates import S2F, MolcasCoord

class MOLCAS_TINKER(MOLCAS):
    """ MOLCAS/TINKER QMMMM calculation interface

        Parameters:          Type:
            keywords         dict         keywords dict
            id               int          calculation index

        Attribute:           Type:
            natom            int          number of active qmmm atoms without link atoms
            qmmm_key         str          tinker key file
            tinker           str          path to tinker

        Functions:           Returns:
            evaluate         self         run qmmm calculation
    """

    def __init__(self, keywords = None, id = None, runtype = 'qm'):
        super().__init__(keywords = keywords, id = id, runtype = runtype)
        self.qmmm_key       = keywords['molecule']['qmmm_key']
        self.tinker         = keywords['molcas']['tinker']

        ## add tinker to runscript
        self.runscript = 'export TINKER=%s\n%s' % (self.tinker, self.runscript)

    def _write_coord(self, traj):
        ## write tinker xyz file and copy tinker key

        xyz = self._update_txyz(traj)
        with open('%s/%s.xyz' % (self.calcdir, self.project), 'w') as out:
            out.write(xyz)

        shutil.copy2(self.qmmm_key, '%s/%s.key' % (self.calcdir, self.project))

    def _update_txyz(self, traj):
        ## prepare new tinker xyz file

        traj = traj.applyqmmm() # compute the H cap coordinates

        active = traj.active
        link   = traj.link
        xyz    = traj.coord
        la     = traj.Hcap_coord
        txyz   = traj.txyz

        ## check qmmm xyz info completeness
        if len(txyz) == 0:
            sys.exit('\n  VariabError\n PyRAI2MD: QMMM xyz information is missing,\
 please use tinker xyz or set the path to qmmm_xyz in &molecule or change qm to molcas in &control')

        out = '%s\n' % (txyz[0])

        for line in txyz[1:]:
            index = int(line[0]) - 1
            part1 = '%6s  %-3s' % (line[0], line[1])
            part3 = ''.join(['%6s' % x for x in line[5:]])
            if index in active:
                x, y, z = xyz[active.index(index)]
                part2 = '%12.6f%12.6f%12.6f' % (x, y, z)
            elif index in link:
                x, y, z = la[link.index(index)]
                part2 = '%12.6f%12.6f%12.6f' % (x, y, z)
            else:
                part2 = '%12s%12s%12s' % (line[2], line[3], line[4])
            out += '%s%s%s\n' % (part1, part2, part3)

        return out

    def _read_data(self, natom):
        ## read molcas logfile and pack data

        with open('%s/%s.log' % (self.calcdir, self.project), 'r') as out:
            log = out.read().splitlines()
        spin       = -1
        coord      = []
        casscf     = []
        gradient   = []
        nac        = []
        soc        = []
        soc_mtx    = []
        soc_state  = 0
        sin_state  = 0
        tri_state  = 0
        flag       = 'grad'
        for i, line in enumerate(log):
            if   """Cartesian coordinates in Angstrom""" in line:
                coord = log[i + 4: i + 4 + natom]
                coord = MolcasCoord(coord)

            elif """Final state energy(ies)""" in line:
                spin += 1
                if   """::    RASSCF root number""" in log[i+3]:
                    shift_line = 3  # normal energy output format
                    en_col = -1
                else:
                    shift_line = 5  # relativistic energy output format
                    en_col = 1
                e = [float(x.split()[en_col]) for x in log[i + shift_line: i + shift_line + self.ci[spin]]]
                casscf += e

            elif """Molecular gradients """ in line:
                flag = 'grad'

            elif """CI derivative coupling""" in line:
                flag = 'nac'

            elif """Molecular gradients, after ESPF""" in line and flag == 'grad':
                g = log[i + 8: i + 8 + natom]
                g = S2F(g)
                gradient.append(g)

            elif """Molecular gradients, after ESPF""" in line and flag == 'nac':
                n = log[i + 8: i + 8 + natom]
                n = S2F(n)
                nac.append(n)

            elif """Nr of states""" in line:
                soc_state = int(line.split()[-1])

            elif """Root nr:""" in line:
                tri_state = int(line.split()[-1])
                sin_state = soc_state - tri_state

            elif """Spin-orbit section""" in line:
                soc_dim = int(sin_state * self.mult[0] + tri_state * self.mult[1])
                soc_urt = int(soc_dim * (soc_dim + 1) / 2)
                soc_sfs = np.zeros([soc_dim, soc_dim])
                soc_mtx = np.zeros([soc_state, soc_state])

                # form soc matrix by spin free eigenstates
                for so_el in log[i+11:i+11+soc_urt]:
                    i1, s1, ms1, i2, s2, ms2, real_part, imag_part, absolute = so_el.split()
                    i1 = int(i1) - 1
                    i2 = int(i2) - 1
                    va = float(absolute)
                    soc_sfs[i1, i2] = va
                    soc_sfs[i2, i1] = va

                # reduce soc matrix intto configuration state
                for s1 in range(sin_state):
                    for s2 in range(tri_state):
                        p2 = sin_state + s2
                        first_col = int(sin_state + s2 * self.mult[1])
                        final_col = int(sin_state + (s2 + 1) * self.mult[1])
                        soc_mtx[s1, p2] = np.sum(soc_sfs[s1, first_col: final_col]**2)**0.5
                        soc_mtx[p2, s1] = soc_mtx[s1, p2]

        ## extract soc matrix elements
        if len(self.soc_coupling) > 0 and len(soc_mtx) > 0:
            for pair in self.soc_coupling:
                s1, s2 = pair
                socme = np.array(soc_mtx[s1, s2 - self.ci[0] + sin_state]).reshape(1)
                soc.append(socme)

        ## pack data
        energy   = np.array(casscf)
        gradient = np.array(gradient)
        nac      = np.array(nac)
        soc      = np.array(soc)

        return coord, energy, gradient, nac, soc

    def evaluate(self, traj):
        ## main function to run Molcas calculation and communicate with other PyRAIMD modules

        ## load trajctory info
        self.ci = traj.ci
        self.mult = traj.mult
        self.soc_coupling = traj.soc_coupling
        self.nsoc = traj.nsoc
        self.nnac = traj.nnac
        self.nstate = traj.nstate

        ## compute properties
        energy = []
        gradient = []
        nac = []
        soc = []
        completion = 0

        ## setup Molcas calculation
        self._setup_molcas(traj)

        ## run Molcas calculation
        self._run_molcas()

        ## read Molcas output files
        coord, energy, gradient, nac, soc = self._read_data(traj.natom)

        if  len(energy) >= self.nstate and\
            len(gradient) >= self.nstate and\
            len(nac) >= self.nnac and\
            len(soc) >= self.nsoc:
            completion = 1

        ## clean up
        if self.keep_tmp == 0:
            shutil.rmtree(self.calcdir)

        # update trajectory
        traj.energy = np.copy(energy)
        traj.grad = np.copy(gradient)
        traj.nac = np.copy(nac)
        traj.soc = np.copy(soc)
        traj.err_energy = None
        traj.err_grad = None
        traj.err_nac = None
        traj.err_soc = None
        traj.status = completion

        return traj


import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torchvision.models as models
from torch.autograd import Variable
import numpy as np
from model.utils.config import cfg
from model.rpn.rpn import _RPN
from model.roi_pooling.modules.roi_pool import _RoIPooling
from model.roi_crop.modules.roi_crop import _RoICrop
from model.roi_align.modules.roi_align import RoIAlignAvg
from model.rpn.proposal_target_layer_cascade import _ProposalTargetLayer
from model.rpn.bbox_transform import bbox_overlaps_batch, bbox_transform_batch
import time
import pdb
from model.utils.net_utils import _smooth_l1_loss, _crop_pool_layer, _affine_grid_gen, _affine_theta

class _fasterRCNN(nn.Module):
    """ faster RCNN """
    def __init__(self, classes, class_agnostic):
        super(_fasterRCNN, self).__init__()
        self.classes = classes
        self.n_classes = len(classes)
        self.class_agnostic = class_agnostic
        # loss
        self.RCNN_loss_cls = 0
        self.RCNN_loss_bbox = 0

        # define rpn
        self.RCNN_rpn = _RPN(self.dout_base_model)
        self.RCNN_proposal_target = _ProposalTargetLayer(self.n_classes)
        self.RCNN_roi_pool = _RoIPooling(cfg.POOLING_SIZE, cfg.POOLING_SIZE, 1.0/16.0)
        self.RCNN_roi_align = RoIAlignAvg(cfg.POOLING_SIZE, cfg.POOLING_SIZE, 1.0/16.0)

        self.grid_size = cfg.POOLING_SIZE * 2 if cfg.CROP_RESIZE_WITH_MAX_POOL else cfg.POOLING_SIZE
        self.RCNN_roi_crop = _RoICrop()

    def forward(self, im_data, im_info, gt_boxes, num_boxes):
        batch_size = im_data.size(0)

        im_info = im_info.data
        gt_boxes = gt_boxes.data
        num_boxes = num_boxes.data

        # feed image data to base model to obtain base feature map
        base_feat = self.RCNN_base(im_data)

        # feed base feature map tp RPN to obtain rois
        rois, rpn_loss_cls, rpn_loss_bbox = self.RCNN_rpn(base_feat, im_info, gt_boxes, num_boxes)

        # if it is training phrase, then use ground trubut bboxes for refining
        if self.training:
            roi_data = self.RCNN_proposal_target(rois, gt_boxes, num_boxes)
            rois, rois_label, rois_target, rois_inside_ws, rois_outside_ws = roi_data

            rois_label = Variable(rois_label.view(-1).long())
            rois_target = Variable(rois_target.view(-1, rois_target.size(2)))
            rois_inside_ws = Variable(rois_inside_ws.view(-1, rois_inside_ws.size(2)))
            rois_outside_ws = Variable(rois_outside_ws.view(-1, rois_outside_ws.size(2)))
        else:
            rois_label = None
            rois_target = None
            rois_inside_ws = None
            rois_outside_ws = None
            rpn_loss_cls = 0
            rpn_loss_bbox = 0

        rois = Variable(rois)

        # update 20191026: get the index of nodes in graph for rois (default: batch_size = 1)
        # if we want to change batch_size, we should consider to change roi2gt_assignment[0]
        # roi_part_match[0] and  roi_part_match_overlap[0] and so on

        if self.training:

            part_threshold\
                = 0.25

            # first, calculate the overlaps among rois and gt, get the max roi for each gt (node_cls)
            overlaps = bbox_overlaps_batch(rois, gt_boxes)
            max_overlaps_rois2gt, roi2gt_assignment = torch.max(overlaps, 1)

            # second, calculate the overlaps among rois and rois_select,
            # using threshold to select roi for each rois_select (node_part)

            rois_cls_tmp = rois[:, roi2gt_assignment[0], :]
            rois_cls_num = np.argwhere(gt_boxes[:, :, 4].cpu().data.numpy()[0] != 0).shape[0]
            rois_cls_tmp = rois_cls_tmp[:,:rois_cls_num, :]
            rois_cls = rois_cls_tmp.new(rois_cls_tmp.size(0), rois_cls_tmp.size(1), 5).zero_()
            rois_cls[:, :, :4] = rois_cls_tmp[:, :, 1:5]
            rois_cls[:, :, 4] = rois_cls_tmp[:, :, 0]

            # rois_cls_idx_list is the idx related from rois_cls to rois
            roi_cls_idx_list = roi2gt_assignment[0][:rois_cls_num]

            overlaps = bbox_overlaps_batch(rois, rois_cls)
            max_overlaps_rois2cls, roi2cls_assignment = torch.max(overlaps, 2)

            roi_part_match_overlap = max_overlaps_rois2cls.cpu().data.numpy()
            roi_part_match = roi2cls_assignment.cpu().data.numpy()

            # roi_part_idx_list is the idx related from rois_part to rois
            roi_part_idx_list = []
            roi_part_match_idx = np.unique(roi_part_match[0])
            for roi_cls_idx in roi_part_match_idx:
                match_idx_tmp = np.transpose(np.argwhere(roi_part_match[0] == roi_cls_idx))[0]
                match_overlap_tmp = roi_part_match_overlap[0][match_idx_tmp]
                # use threshold to select rois_part
                match_idx_tmp_select = np.transpose(np.argwhere(match_overlap_tmp > part_threshold))[0]
                match_idx_tmp = match_idx_tmp[match_idx_tmp_select]
                roi_part_idx_list.append(torch.from_numpy(match_idx_tmp))

        # do roi pooling based on predicted rois
        if cfg.POOLING_MODE == 'crop':
            # pdb.set_trace()
            # pooled_feat_anchor = _crop_pool_layer(base_feat, rois.view(-1, 5))
            grid_xy = _affine_grid_gen(rois.view(-1, 5), base_feat.size()[2:], self.grid_size)
            grid_yx = torch.stack([grid_xy.data[:,:,:,1], grid_xy.data[:,:,:,0]], 3).contiguous()
            pooled_feat = self.RCNN_roi_crop(base_feat, Variable(grid_yx).detach())
            if cfg.CROP_RESIZE_WITH_MAX_POOL:
                pooled_feat = F.max_pool2d(pooled_feat, 2, 2)
        elif cfg.POOLING_MODE == 'align':
            pooled_feat = self.RCNN_roi_align(base_feat, rois.view(-1, 5))
        elif cfg.POOLING_MODE == 'pool':
            pooled_feat = self.RCNN_roi_pool(base_feat, rois.view(-1,5))

        # feed pooled features to top model
        pooled_feat = self._head_to_tail(pooled_feat)

        # update 20191027: build graph for rois based on index (default: batch_size = 1)
        if self.training:
            adj_jud = np.zeros((0))
            adj_rois = torch.zeros(0).cuda().long()
            for i in range(roi_cls_idx_list.shape[0]):
                adj_jud = np.concatenate((adj_jud, [1]))
                adj_rois = torch.cat((adj_rois, roi_cls_idx_list[i:i+1]))
                try:
                    adj_jud = np.concatenate((adj_jud, np.zeros((roi_part_idx_list[i].shape[0]))))
                    adj_rois = torch.cat((adj_rois, roi_part_idx_list[i].cuda()))
                except IndexError:
                    print ('IndexError happen, continue')
                    continue

            node_cls_idx = np.transpose(np.argwhere(adj_jud == 1))[0]

            adj_matrix_bin = np.zeros((len(adj_jud), len(adj_jud)))

            # link edges for node_cls to node_cls
            for k in range(len(node_cls_idx)-1):
                idx_node_cls_1 = node_cls_idx[k]
                idx_node_cls_2 = node_cls_idx[k + 1]
                adj_matrix_bin[idx_node_cls_1, idx_node_cls_2] = 1
                adj_matrix_bin[idx_node_cls_2, idx_node_cls_1] = 1

            # link edges for node_cls to related node_part
            for k in range(len(node_cls_idx)-1):
                idx_start = node_cls_idx[k]
                idx_end = node_cls_idx[k + 1]
                for s in range(idx_start, idx_end):
                    for t in range(idx_start, idx_end):
                        if s == t:
                            adj_matrix_bin[s, t] = 0
                        else:
                            adj_matrix_bin[s, t] = 1

            # calculate the adj_mat based on adj_matrix_bin, the weights on edges are the cosine distance between nodes
            adj_matrix = np.zeros((len(adj_jud), len(adj_jud)))

            cos = torch.nn.CosineSimilarity(dim=0, eps=1e-6)

            for s in range(len(adj_jud)):
                for t in range(len(adj_jud)):
                    if adj_matrix_bin[s, t] == 1:
                        node_feat_s = pooled_feat[adj_rois[s], :]
                        node_feat_t = pooled_feat[adj_rois[t], :]
                        adj_matrix[s, t] = cos(node_feat_s, node_feat_t)
                    else:
                        adj_matrix[s, t] = 0

            adj_matrix = torch.from_numpy(adj_matrix).float().cuda()

            pooled_feat[adj_rois, :] = F.relu(self.gcn1(pooled_feat[adj_rois, :], adj_matrix))
            pooled_feat[adj_rois, :] = F.relu(self.gcn2(pooled_feat[adj_rois, :], adj_matrix))

        # compute bbox offset
        bbox_pred = self.RCNN_bbox_pred(pooled_feat)
        if self.training and not self.class_agnostic:
            # select the corresponding columns according to roi labels
            bbox_pred_view = bbox_pred.view(bbox_pred.size(0), int(bbox_pred.size(1) / 4), 4)
            bbox_pred_select = torch.gather(bbox_pred_view, 1, rois_label.view(rois_label.size(0), 1, 1).expand(rois_label.size(0), 1, 4))
            bbox_pred = bbox_pred_select.squeeze(1)

        # compute object classification probability
        cls_score = self.RCNN_cls_score(pooled_feat)
        cls_prob = F.softmax(cls_score, 1)

        RCNN_loss_cls = 0
        RCNN_loss_bbox = 0

        if self.training:
            # classification loss
            RCNN_loss_cls = F.cross_entropy(cls_score, rois_label)

            # bounding box regression L1 loss
            RCNN_loss_bbox = _smooth_l1_loss(bbox_pred, rois_target, rois_inside_ws, rois_outside_ws)

        cls_prob = cls_prob.view(batch_size, rois.size(1), -1)
        bbox_pred = bbox_pred.view(batch_size, rois.size(1), -1)

        # update 2019-6-17:fix the bug for dimension specified as 0...
        if self.training:
            rpn_loss_cls = torch.unsqueeze(rpn_loss_cls, 0)
            rpn_loss_bbox = torch.unsqueeze(rpn_loss_bbox, 0)
            RCNN_loss_cls = torch.unsqueeze(RCNN_loss_cls, 0)
            RCNN_loss_bbox = torch.unsqueeze(RCNN_loss_bbox, 0)

        return rois, cls_prob, bbox_pred, rpn_loss_cls, rpn_loss_bbox, RCNN_loss_cls, RCNN_loss_bbox, rois_label

    def _init_weights(self):
        def normal_init(m, mean, stddev, truncated=False):
            """
            weight initalizer: truncated normal and random normal.
            """
            # x is a parameter
            if truncated:
                m.weight.data.normal_().fmod_(2).mul_(stddev).add_(mean) # not a perfect approximation
            else:
                m.weight.data.normal_(mean, stddev)
                m.bias.data.zero_()

        normal_init(self.RCNN_rpn.RPN_Conv, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_rpn.RPN_cls_score, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_rpn.RPN_bbox_pred, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_cls_score, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_bbox_pred, 0, 0.001, cfg.TRAIN.TRUNCATED)

    def create_architecture(self):
        self._init_modules()
        self._init_weights()


"""
    jupylet/sprite.py

    Copyright (c) 2020, Nir Aides - nir@winpdb.org

    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions are met:

    1. Redistributions of source code must retain the above copyright notice, this
       list of conditions and the following disclaimer.
    2. Redistributions in binary form must reproduce the above copyright notice,
       this list of conditions and the following disclaimer in the documentation
       and/or other materials provided with the distribution.

    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
    ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
    WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
    DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
    ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
    (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
    LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
    ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
    SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""

import webcolors
import moderngl
import pathlib
import math
import glm

import PIL.Image

import moderngl_window as mglw
import numpy as np

from .collision import trbl, hitmap_and_outline_from_alpha, compute_collisions
from .resource import load_texture, pil_from_texture, get_shader_2d
from .utils import glm_dumps, glm_loads
from .color import c2v
from .state import State
from .node import Node, aa2q, q2aa
from .lru import SPRITE_TEXTURE_UNIT

_empty_array = np.array([])

class Sprite(Node):

    def __init__(
        self,
        img,
        x=0,
        y=0,
        scale=1.0,
        angle=0.0,
        anchor_x='center',
        anchor_y='center',
        color='white',
        flip=True,
        mipmap=True,
        autocrop=False,
        anisotropy=8.0,
        height=None,
        width=None,
        name=None,
        collisions=True,
    ):

        super(Sprite, self).__init__(
            name,
            rotation=aa2q(glm.radians(angle)),
            scale=None,
            position=glm.vec3(x, y, 0),
        )

        self.mipmap = mipmap
        self.autocrop = autocrop
        self.anisotropy = anisotropy

        self.geometry = mglw.geometry.quad_2d(
            size=(1.0, 1.0),
            pos=(0.5, 0.5)
        )

        self.texture = load_texture(
            img,
            anisotropy=anisotropy,
            autocrop=autocrop,
            mipmap=mipmap,
            flip=False,
        )
        self.texture.repeat_x = False
        self.texture.repeat_y = False

        self.baseline = 0
        self.components = self.texture.components
        self.color4 = glm.vec4(1., 1., 1., 1.)
        self.flip = flip

        self._collisions = collisions

        if self._collisions:
            self.hitmap, self.outline = hitmap_and_outline_from_alpha(self.image)

        self.scale = scale

        if width:
            self.width = width

        elif height:
            self.height = height

        self.set_anchor(anchor_x, anchor_y)
        self.color = color

    def update(self, shader):
        pass

    def draw(self, shader=None):
        return self.render(shader)

    def render(self, shader=None):

        shader = shader or get_shader_2d()

        if self._dirty:
            self.update(shader)

        shader['components'] = self.components
        shader['color'].write(self.color4)
        shader['flip'] = self.flip

        shader['model'].write(self.matrix)
        shader['texture0'] = SPRITE_TEXTURE_UNIT

        self.texture.use(location=SPRITE_TEXTURE_UNIT)
        self.geometry.render(shader)

    @property
    def scale(self):
        return self.scale0.x / self.texture.width

    @scale.setter
    def scale(self, scale):
        self.scale0 = scale * glm.vec3(self.texture.width, self.texture.height, 1)

    @property
    def x(self):
        return self.position.x

    @x.setter
    def x(self, value):
        self.position.x = value

    @property
    def y(self):
        return self.position.y

    @y.setter
    def y(self, value):
        self.position.y = value

    @property
    def angle(self):
        angle, axis = q2aa(self.rotation)
        return round(glm.degrees(angle * glm.sign(axis.z)), 4)

    @angle.setter
    def angle(self, angle):
        self.rotation = aa2q(glm.radians(angle))

    def set_anchor(self, ax=None, ay=None):

        self._ax = ax
        self._ay = ay

        if ax == 'left':
            self.anchor.x = 0
        elif ax == 'center':
            self.anchor.x = 0.5
        elif ax == 'right':
            self.anchor.x = 1.
        elif type(ax) in (int, float):
            self.anchor.x = ax / self.width

        if ay == 'bottom':
            self.anchor.y = 0
        if ay == 'baseline':
            self.anchor.y = self.baseline
        elif ay == 'center':
            self.anchor.y = 0.5
        elif ay == 'top':
            self.anchor.y = 1.
        elif type(ay) in (int, float):
            self.anchor.y = ay / self.width

    @property
    def width(self):
        return self.scale0.x

    @width.setter
    def width(self, width):
        self.scale0 = glm.vec3(self.texture.width, self.texture.height, 1) * width / self.texture.width

    @property
    def height(self):
        return self.scale0.y

    @height.setter
    def height(self, height):
        self.scale0 = glm.vec3(self.texture.width, self.texture.height, 1) * height / self.texture.height

    @property
    def image(self):
        return pil_from_texture(self.texture)

    @image.setter
    def image(self, img):

        scale = self.scale

        self.texture.release()
        self.texture = load_texture(
            img,
            anisotropy=self.anisotropy,
            autocrop=self.autocrop,
            mipmap=self.mipmap,
            flip=False,
        )
        self.texture.repeat_x = False
        self.texture.repeat_y = False

        self.scale = scale

        if self._collisions:
            self.hitmap, self.outline = hitmap_and_outline_from_alpha(self.image)

    def collisions_with(self, o, debug=False):

        #if self.distance_to(o) > self.radius + o.radius:
        #    return

        x0, y0 = self.position.xy
        x1, y1 = o.position.xy

        t0, r0, b0, l0 = self._trbl()
        t1, r1, b1, l1 = o._trbl()

        if t0 + y0 < b1 + y1 or t1 + y1 < b0 + y0:
            return _empty_array[:0]

        if r0 + x0 < l1 + x1 or r1 + x1 < l0 + x0:
            return _empty_array[:0]

        return compute_collisions(o, self, debug=debug)

    def distance_to(self, o=None, pos=None):

        x, y = pos or (o.position.x, o.position.y)

        dx = x - self.position.x
        dy = y - self.position.y

        return (dx ** 2 + dy ** 2) ** 0.5

    def angle_to(self, o=None, pos=None):

        qd = {
            (True, True): 0,
            (True, False): 180,
            (False, False): 180,
            (False, True): 360,
        }

        x, y = pos or (o.position.x, o.position.y)

        dx = x - self.position.x
        dy = y - self.position.y

        a0 = math.atan(dy / (dx or 1e-7)) / math.pi * 180 + qd[(dy >= 0, dx >= 0)]

        return -a0

    @property
    def top(self):
        t, r, b, l = self._trbl()
        return self.position.y + t

    @property
    def right(self):
        t, r, b, l = self._trbl()
        return self.position.x + r

    @property
    def bottom(self):
        t, r, b, l = self._trbl()
        return self.position.y + b

    @property
    def left(self):
        t, r, b, l = self._trbl()
        return self.position.x + l

    @property
    def radius(self):
        t, r, b, l = self._trbl()
        rs = max(t, b) ** 2 + max(r, l) ** 2
        return rs ** .5

    def _trbl(self):
        tx = self.texture
        return trbl(
            tx.width,
            tx.height,
            self.anchor.x * tx.width,
            self.anchor.y * tx.height,
            self.angle,
            self.scale
        )

    def wrap_position(self, width, height, margin=50):
        self.position.x = (self.position.x + margin) % (width + 2 * margin) - margin
        self.position.y = (self.position.y + margin) % (height + 2 * margin) - margin

    def clip_position(self, width, height, margin=0):
        self.position.x = max(-margin, min(margin + width, self.position.x))
        self.position.y = max(-margin, min(margin + height, self.position.y))

    @property
    def opacity(self):
        return self.color4.a

    @opacity.setter
    def opacity(self, opacity):
        self.color4.a = opacity

    @property
    def color(self):
        return self.color4

    @color.setter
    def color(self, color):
        self.color4 = c2v(color, self.color4.a)

    def get_state(self):
        return dict(
            node = super(Sprite, self).get_state(),
            color4 = glm_dumps(glm.vec4(self.color4)),
            mipmap = self.mipmap,
            autocrop = self.autocrop,
            anisotropy = self.anisotropy,
            baseline = self.baseline,
            components = self.components,
            flip = self.flip,
            image = self.image,
        )

    def set_state(self, s):

        for k, v in s.items():
            if k == 'node':
                super(Sprite, self).set_state(v)
            else:
                setattr(self, k, glm_loads(v))


from collections import deque
from typing import Dict, List, Tuple, Union

import numpy as np

from utils import utils

class ReplayBuffer:
    def __init__(
        self, buffer_length: int = 10000, batch_size: int = 32, n_steps: int = 3, gamma: float = 0.99, num_envs: int = 1
    ):
        self.num_envs = num_envs
        self.n_steps = n_steps
        self.buffer_length = buffer_length // num_envs
        self.batch_size = batch_size
        self.n_step_buffer = {env_id: deque([], maxlen=n_steps) for env_id in range(num_envs)}
        self.buffer = {env_id: deque([], maxlen=int(buffer_length // num_envs) + 1) for env_id in range(num_envs)}
        self.beta = 1
        self.gamma = gamma

    def step(self):
        pass

    def clear_n_step_buffer(self, env_id: int):
        for trans_id in range(1, len(self.n_step_buffer[env_id])):
            self.add_new_priority(env_id)
            n_step_transition = utils.Transition(
                *(
                    self.n_step_buffer[env_id][trans_id].state,
                    self.n_step_buffer[env_id][trans_id].action,
                    self.n_step_buffer[env_id][trans_id].next_state,
                    utils.compute_cumulated_return(list(self.n_step_buffer[env_id])[trans_id:], self.gamma),
                    self.n_step_buffer[env_id][-1].done,
                )
            )
            self.buffer[env_id].append(n_step_transition)

        self.n_step_buffer[env_id] = deque([], maxlen=self.n_steps)

    def add_new_priority(self, env_id: int):
        pass

    def push(self, env_id: int, *args):
        one_step_transition = utils.Transition(*args)
        self.n_step_buffer[env_id].append(one_step_transition)
        if len(self.n_step_buffer[env_id]) < self.n_steps:
            return

        n_step_transition = utils.Transition(
            *(
                self.n_step_buffer[env_id][0].state,
                self.n_step_buffer[env_id][0].action,
                one_step_transition.next_state,
                utils.compute_cumulated_return(self.n_step_buffer[env_id], self.gamma),
                one_step_transition.done,
            )
        )

        self.add_new_priority(env_id)
        self.buffer[env_id].append(n_step_transition)

        if one_step_transition.done:
            self.clear_n_step_buffer(env_id)

    def sample(self) -> Tuple[List, np.array, np.array]:
        raise NotImplementedError("Replay buffer does not implement 'sample()' method!")

    def sample_curl(self, crop_size: int):
        raise NotImplementedError("Replay buffer does not implement 'sample_curl()' method!")

    def update(self, indices: np.array, priorities: np.array):
        pass

    @property
    def buffer_flattened(self) -> Dict:
        raise NotImplementedError("Replay buffer does not implement 'buffer_flattened()' method!")

    def __len__(self):
        len_total = 0
        for env_id in range(self.num_envs):
            len_total += len(self.buffer[env_id])

        return len_total

class ExperienceReplay(ReplayBuffer):
    def __init__(
        self, buffer_length: int = 10000, batch_size: int = 32, n_steps: int = 3, gamma: float = 0.99, num_envs: int = 1
    ):
        super().__init__(buffer_length, batch_size, n_steps, gamma, num_envs)

    def sample(self) -> Tuple[List, np.array, np.array]:
        sample_indices = np.random.choice(list(self.buffer_flattened.keys()), size=self.batch_size)
        sample_transitions = [self.buffer_flattened[idx][2] for idx in sample_indices]
        return sample_transitions, sample_indices, np.ones(self.batch_size)

    def sample_curl(self, crop_size: int) -> Tuple[List, np.array, np.array]:
        sample_indices = np.random.choice(list(self.buffer_flattened.keys()), size=self.batch_size)
        sample_transitions = [self.buffer_flattened[idx][2] for idx in sample_indices]
        for i, transition in enumerate(sample_transitions):
            curl_transition = utils.TransitionCurl(
                *(
                    utils.random_crop(transition.state, crop_size),
                    utils.random_crop(transition.state, crop_size),
                    utils.random_crop(transition.state, crop_size),
                    transition.action,
                    utils.random_crop(transition.next_state, crop_size),
                    transition.reward,
                    transition.done,
                )
            )
            sample_transitions[i] = curl_transition

        return sample_transitions, sample_indices, np.ones(self.batch_size)

    @property
    def buffer_flattened(self) -> Dict:
        flattened_buffer = {}
        idx = 0
        for env_id in range(self.num_envs):
            for t_id in range(len(self.buffer[env_id])):
                flattened_buffer[idx] = (env_id, t_id, self.buffer[env_id][t_id])
                idx += 1

        return flattened_buffer

class PrioritizedExperienceReplay(ReplayBuffer):
    def __init__(
        self,
        buffer_length: int = 10000,
        batch_size: int = 32,
        n_steps: int = 3,
        alpha: float = 0.75,
        beta0: float = 0.4,
        replay_delay: int = 1000,
        gamma: float = 0.99,
        num_envs: int = 1,
    ):
        super().__init__(buffer_length, batch_size, n_steps, gamma, num_envs)

        self.alpha = alpha
        self.beta0 = beta0
        self.beta = beta0
        self.priorities = {env_id: deque([], maxlen=buffer_length) for env_id in range(num_envs)}
        self.total_steps = (buffer_length // batch_size) * replay_delay

    def step(self):
        self.beta = np.minimum(self.beta + (1 - self.beta0) / self.total_steps, 1)

    def sample(self) -> Tuple[List, np.array, np.array]:
        flattened_buffer = self.buffer_flattened
        priorities_tmp = np.array([entry[3] for _, entry in flattened_buffer.items()])
        probabilities = priorities_tmp ** self.alpha
        probabilities /= probabilities.sum()

        sample_indices = np.random.choice(
            list(flattened_buffer.keys()), size=self.batch_size, p=probabilities, replace=False
        )
        sample_transitions = [flattened_buffer[idx][2] for idx in sample_indices]

        weights = (probabilities[sample_indices] * len(self.buffer)) ** (-self.beta)
        weights = np.array(weights / weights.max(), dtype=np.float32)

        return sample_transitions, sample_indices, weights

    def sample_curl(self, crop_size: int) -> Tuple[List, np.array, np.array]:
        flattened_buffer = self.buffer_flattened
        priorities_tmp = np.array([entry[3] for _, entry in flattened_buffer.items()])
        probabilities = priorities_tmp ** self.alpha
        probabilities /= probabilities.sum()

        sample_indices = np.random.choice(
            list(flattened_buffer.keys()), size=self.batch_size, p=probabilities, replace=False
        )
        sample_transitions = [flattened_buffer[idx][2] for idx in sample_indices]

        for i, transition in enumerate(sample_transitions):
            curl_transition = utils.TransitionCurl(
                *(
                    utils.random_crop(transition.state, crop_size),
                    utils.random_crop(transition.state, crop_size),
                    utils.random_crop(transition.state, crop_size),
                    transition.action,
                    utils.random_crop(transition.next_state, crop_size),
                    transition.reward,
                    transition.done,
                )
            )
            sample_transitions[i] = curl_transition

        weights = (probabilities[sample_indices] * len(self.buffer)) ** (-self.beta)
        weights = np.array(weights / weights.max(), dtype=np.float32)

        return sample_transitions, sample_indices, weights

    def add_new_priority(self, env_id: int):
        if len(self.buffer[env_id]) == 0:
            self.priorities[env_id].append(1.0)
        else:
            self.priorities[env_id].append(np.max(np.array([entry[3] for _, entry in self.buffer_flattened.items()])))
            priorities_tmp = np.array(self.priorities[env_id])
            self.priorities[env_id] = deque(priorities_tmp, maxlen=self.buffer_length)

    def update(self, indices: np.array, priorities: np.array):
        buf_flattened = self.buffer_flattened
        for i, flattened_idx in enumerate(indices):
            env_id, t_id, _, _ = buf_flattened[flattened_idx]
            self.priorities[env_id][t_id] = np.abs(priorities[i])

    @property
    def buffer_flattened(self) -> Dict:
        flattened_buffer = {}
        idx = 0
        for env_id in range(self.num_envs):
            for t_id in range(len(self.buffer[env_id])):
                flattened_buffer[idx] = (env_id, t_id, self.buffer[env_id][t_id], self.priorities[env_id][t_id])
                idx += 1

        return flattened_buffer


#!/usr/local/bin/python
import sys
import subprocess
import os
import numpy as np

if len(sys.argv) < 3:
    print 'usage:' + sys.argv[0] + ' <max> <out file> <in bam 1> <in bam 2> <..>'
    exit(1)

max_c = int(sys.argv[1])
out_file = sys.argv[2]

o = open('.exclude.tmp','w')

for i in range(3,len(sys.argv)):
    bam_file = sys.argv[i]
    coverage_file = bam_file + '.coverage'

    f = open(coverage_file, 'r')
    for l in f:
        a = l.rstrip().split('\t')
        if (int(a[3]) > max_c):
            o.write(l)
    f.close()
o.close()

p = subprocess.Popen(\
        'cat .exclude.tmp | ' \
        'sort -S 20G -k1,1 -k2,2n | ' \
        'bedtools merge -i stdin -d 50 ' \
        '> ' + out_file, shell=True)
os.waitpid(p.pid, 0)
os.unlink('.exclude.tmp')