import numpy as np
import matplotlib.pyplot as plt
import random
from PIL import Image
from matplotlib import gridspec
from matplotlib.collections import PatchCollection
from matplotlib.patches import Rectangle, Circle
from matplotlib.pyplot import draw, figure, get_cmap, gray
from matplotlib.transforms import Affine2D
from numpy import array, uint32, round, sqrt, ceil, asarray, append
from numpy import spacing as eps
from hotspotter.other.AbstractPrintable import AbstractManager
from hotspotter.other.ConcretePrintable import Pref
from hotspotter.other.logger import logmsg, logdbg, logwarn
from warnings import catch_warnings, simplefilter
import colorsys
import os.path
import sys

class DrawManager(AbstractManager):
    def init_preferences(dm, default_bit=False):
        iom = dm.hs.iom
        if dm.draw_prefs == None:
            dm.draw_prefs = Pref(fpath=iom.get_prefs_fpath('draw_prefs'))
        dm.draw_prefs.bbox_bit       = True
        dm.draw_prefs.ellipse_bit    = False
        dm.draw_prefs.ellipse_alpha  = .6
        dm.draw_prefs.points_bit     = False
        dm.draw_prefs.result_view  = Pref(1, choices=['in_image', 'in_chip'])
        dm.draw_prefs.fignum         = 0
        dm.draw_prefs.num_result_cols = 3
        dm.draw_prefs.figsize        = (5,5)
        dm.draw_prefs.colormap       = Pref('hsv', hidden=True)
        dm.draw_prefs.in_qtc_bit     = Pref(False, hidden=True) #Draw in the Qt Console
        dm.draw_prefs.use_thumbnails = Pref(False, hidden=True)
        dm.draw_prefs.thumbnail_size = Pref(128, hidden=True)
        if not default_bit:
            dm.draw_prefs.load()
    # ---
    def show_splash(dm):
        splash_fname = os.path.join(os.path.realpath(os.path.dirname(__file__)), 'splash.png')
        if not os.path.exists(splash_fname):
            root_dir = os.path.realpath(os.path.dirname(__file__))
            while root_dir!=None:
                splash_fname = os.path.join(root_dir, "hotspotter", "front", 'splash.png')
                logdbg(splash_fname)
                exists_test = os.path.exists(splash_fname)
                logdbg('Exists:'+str(exists_test))
                if exists_test:
                    break
                tmp = os.path.dirname(root_dir)
                if tmp == root_dir:
                    root_dir = None
                else:
                    root_dir = tmp
        logdbg('Splash Fname: %r '% splash_fname)
        splash_img = asarray(Image.open(splash_fname))
        dm.add_images([splash_img],['Welcome to Hotspotter'])
        dm.end_draw()
    # ---
    def show_image(dm, gx):
        gm, cm = dm.hs.get_managers('gm','cm')
        gid        = gm.gx2_gid[gx]
        img_list   = gm.gx2_img_list(gx)
        title_list = ['gid='+str(gid)+'   gname='+gm.gx2_gname[gx]]
        dm.add_images(img_list, title_list)
        cx_list    = gm.gx2_cx_list[gx]
        if dm.draw_prefs.use_thumbnails is True:
            pass
        for cx in iter(cx_list):
            dm.draw_chiprep2(cx, axi=0, in_image_bit=True)
        dm.end_draw()
    # ---
    def show_chip(dm, cx, in_raw_chip=False, **kwargs):
        cm = dm.hs.cm
        cid, gname, chip = cm.cx2_(cx, 'cid', 'gname', 'chip')
        if in_raw_chip:
            chip = np.asarray(cm.cx2_pil_chip(cx, scaled=True,
                                              preprocessed=False, rotated=True,
                                             colored=True))
        if dm.draw_prefs.use_thumbnails is True:
            pass
        dm.add_images([chip], [gname])
        # Draw chiprep and return fsel incase rand is good
        fsel_ret = dm.draw_chiprep2(cx, axi=0, **kwargs)
        dm.end_draw()
        return fsel_ret
    # ---
    def show_query(dm, res, titleargs=None, enddraw=True):
        # Make sure draw is valid
        if res is None: dm.show_splash(); return
        # Get Chip Properties

        cm  = res.hs.cm
        qcm = res.qhs.cm

        titleargs =\
        ('cx', 'cid', 'nid', 'name', 'gname')
        ( qcx, qcid , qnid , qname , qgname ) = res.qcid2_(*titleargs)
        ( tcx, tcid , tnid , tname , tgname ) = res.tcid2_(*titleargs)

        (tcx, tscore, ) = res.tcid2_('cx','score')
        # Titles of the Plot

        #qtitle = 'gname: %s\nQUERY cid=%d, nid=%d' % (qgname, qcid, qnid)
        #ttile = ['cid=%d\n gname: %s\nrank/score=%d,%.2f' % tup for tup in zip(tcid, tgname, range(1,len(tscore)+1), tscore)]
        qtitle = 'gname: %s\nQUERY nid=%d' % (qgname, qnid)
        ttile = ['gname: %s\nrank/score=%d/%.2f' % tup for tup in zip(tgname, range(1,len(tscore)+1), tscore)]
        title_list = [qtitle] + ttile
        if dm.draw_prefs.use_thumbnails is True:
            pass

        # Add the images to draw
        in_image_bit = dm.draw_prefs.result_view == 'in_image'
        if in_image_bit:
            qimg = qcm.cx2_img(qcx)
            timg = cm.cx2_img_list(tcx)
            dm.add_images([qimg] + timg, title_list)
        elif dm.draw_prefs.result_view == 'in_chip':
            qchip = qcm.cx2_chip_list(qcx)
            tchip = cm.cx2_chip_list(tcx)
            dm.add_images(qchip + tchip, title_list)

        # Draw the Query Chiprep
        qaxi       = 0; qfsel      = []
        dm.draw_chiprep2(qcx, axi=qaxi, fsel=qfsel, qcm=qcm)
        # Draw the Top Result Chipreps
        for (tx, cx) in enumerate(tcx):
            fm    = res.rr.cx2_fm[cx]
            fs    = res.rr.cx2_fs[cx]
            axi   = tx+1
            if len(fs) == 0:
                qfsel = np.array([], np.uint32)
                fsel  = np.array([], np.uint32)
            else:
                qfsel = fm[fs > 0][:,0]
                fsel  = fm[fs > 0][:,1]
            dm.draw_chiprep2(cx,
                             axi=axi,
                             axi_color=axi,
                             fsel=fsel,
                             in_image_bit=in_image_bit)
            dm.draw_chiprep2(qcx,
                             axi=qaxi,
                             axi_color=axi,
                             fsel=qfsel,
                             in_image_bit=in_image_bit,
                             qcm=qcm)
        if enddraw:
            dm.end_draw()

    # ---

    def __init__(dm, hs):
        super( DrawManager, dm ).__init__( hs )
        dm.hs      =   hs
        dm.fignum  =    0
        dm.dpi     = 100 #72
        dm.draw_prefs = None
        dm.ax_list =   []
        dm.init_preferences()
    # ---
    def update_figsize(dm):
        fig = dm.get_figure()
        dm.draw_prefs.figsize = (fig.get_figheight(), fig.get_figwidth())
    # ---
    def get_figure(dm):
        guifig = dm.hs.uim.get_gui_figure()
        if guifig != None and dm.fignum == 0: # Check to see if we have access to the gui
            return guifig
        fig = figure(num=dm.fignum, figsize=dm.draw_prefs.figsize, dpi=dm.dpi, facecolor='w', edgecolor='k')
        return fig
    # ---
    def annotate_orientation(dm):
        logmsg('Please select an orientation of the torso (Click Two Points on the Image)')
        try:
            # Compute an angle from user interaction
            sys.stdout.flush()
            fig = dm.get_figure()
            pts = np.array(fig.ginput(2))
            logdbg('GInput Points are: '+str(pts))
            # Get reference point to origin
            refpt = pts[0] - pts[1]
            #theta = np.math.atan2(refpt[1], refpt[0])
            theta = np.math.atan(refpt[1]/refpt[0])
            logmsg('The angle in radians is: '+str(theta))
            return theta
        except Exception as ex:
            logmsg('Annotate Orientation Failed'+str(ex))
            return None
    def annotate_roi(dm):
        logmsg('Please Select a Rectangular Region of Interest (Click Two Points on the Image)')
        try:
            sys.stdout.flush()
            fig = dm.get_figure()
            pts = fig.ginput(2)
            logdbg('GInput Points are: '+str(pts))
            (x1, y1, x2, y2) = (pts[0][0], pts[0][1], pts[1][0], pts[1][1])
            xm = min(x1,x2)
            xM = max(x1,x2)
            ym = min(y1,y2)
            yM = max(y1,y2)
            (x, y, w, h) = (xm, ym, xM-xm, yM-ym)
            roi = array(round([x,y,w,h]),dtype=uint32)
            logmsg('The new ROI is: '+str(roi))
            return roi
        except Exception as ex:
            logmsg('Annotate ROI Failed'+str(ex))
            return None
    # ---
    def end_draw(dm):
        #gray()
        logdbg('Finalizing Draw with '+str(len(dm.ax_list))+' axes')
        fig = dm.get_figure()
        #fig.subplots_adjust(hspace=0.2, wspace=0.2)
        #fig.tight_layout(pad=.3, h_pad=None, w_pad=None)
        #fig.tight_layout()
        if dm.draw_prefs.in_qtc_bit:
            try:
                from IPython.back.display import display
                display(fig)
            except:
                logwarn('Cannot Draw in QTConsole')
        fig.show()
        dm.hs.uim.redraw_gui()
        fig.canvas.draw()
        #draw()
    # ---
    def save_fig(dm, save_file):
        dm.end_draw()
        fig = dm.get_figure()
        fig.savefig(save_file, format='png')
    # ---
    def add_images(dm, img_list, title_list=[]):
        fig = dm.get_figure(); fig.clf()
        num_images = len(img_list)
        #
        dm.ax_list     = [None]*num_images
        title_list     = title_list + ['NoTitle']*(num_images-len(title_list))
        # Fit Images into a grid
        max_columns = min(num_images, max(1,dm.draw_prefs.num_result_cols))
        if max_columns == 0: max_columns = 1
        nr = int( ceil( float(num_images)/max_columns) )
        nc = max_columns if num_images >= max_columns else 1
        #
        gs = gridspec.GridSpec( nr, nc )
        for i in xrange(num_images):
            #logdbg(' Adding the '+str(i)+'th Image')
            #logdbg('   * type(img_list[i]): %s'+str(type(img_list[i])))
            #logdbg('   * img_list[i].shape: %s'+str(img_list[i].shape))
            dm.ax_list[i] = fig.add_subplot(gs[i])
            imgplot = dm.ax_list[i].imshow(img_list[i])
            imgplot.set_cmap('gray')
            dm.ax_list[i].get_xaxis().set_ticks([])
            dm.ax_list[i].get_yaxis().set_ticks([])
            dm.ax_list[i].set_title(title_list[i])
            # transData: data coordinates -> display coordinates
            # transAxes: axes coordinates -> display coordinates
            # transLimits: data - > axes
        #
        #gs.tight_layout(fig)
        logdbg('Added '+str(num_images)+' images/axes')
    # ---
    def _get_fpt_ell_collection(dm, fpts, T_data, alpha, edgecolor):
        ell_patches = []
        for (x,y,a,c,d) in fpts: # Manually Calculated sqrtm(inv(A))
            with catch_warnings():
                simplefilter("ignore")
                aIS = 1/sqrt(a)
                cIS = (c/sqrt(a) - c/sqrt(d))/(a - d + eps(1))
                dIS = 1/sqrt(d)
            transEll = Affine2D([\
                    ( aIS,   0,   x),\
                    ( cIS, dIS,   y),\
                    (   0,   0,   1)])
            unitCirc1 = Circle((0,0),1,transform=transEll)
            ell_patches = [unitCirc1] + ell_patches
        ellipse_collection = PatchCollection(ell_patches)
        ellipse_collection.set_facecolor('none')
        ellipse_collection.set_transform(T_data)
        ellipse_collection.set_alpha(alpha)
        ellipse_collection.set_edgecolor(edgecolor)
        return ellipse_collection

    def draw_chiprep2(dm, cx, axi=0, fsel=None, in_image_bit=False, axi_color=0,
                      bbox_bit=None,
                      ell_alpha=None,
                      ell_bit=None,
                      xy_bit=None,
                      color=None,
                      qcm=None,
                      **kwargs):
        '''
        Draws a chip representation over an already drawn chip
        cx           - the chiprep to draw. Managed by the chip manager
        axi          - the axis index to draw it in
        #TODO: in_image_bit becomes data_coordinates
        in_image_bit - are the data coordinates image or rotated chip?
                       raw the chip by itself or in its original image
        axi_color    - use the color associated with axis index
                       (used for ploting queries)
        ---
        Others are preference overloads
        bbox_bit -
        ell_alpha
        ell_bit
        xy_bit
        ell_color
        '''
        # Allows display of cross database queries
        cm = dm.hs.cm if qcm is None else qcm
        # Grab Preferences
        xy_bit  = dm.draw_prefs.points_bit    if xy_bit    is None else xy_bit
        ell_bit = dm.draw_prefs.ellipse_bit   if ell_bit   is None else ell_bit
        bbox_bit  = dm.draw_prefs.bbox_bit      if bbox_bit  is None else bbox_bit
        ell_alpha = dm.draw_prefs.ellipse_alpha if ell_alpha is None else ell_alpha

        # Make sure alpha in range [0,1]
        if ell_alpha > 1: ell_alpha = 1.0
        if ell_alpha < 0: ell_alpha = 0.0

        # Get color from colormap or overloaded parameter
        if color is None:
            color   = plt.get_cmap('hsv')(float(axi_color)/len(dm.ax_list))[0:3]
            if axi_color == 0:
                color = [color[0], color[1]+.5, color[2]]

        # Axis We are drawing to.
        ax     = dm.ax_list[axi]
        T_data = ax.transData # data coordinates -> display coordinates
        # Data coordinates are chip coords

        if xy_bit or ell_bit or fsel != None:
            T_fpts = T_data if not in_image_bit else\
                    Affine2D(cm.cx2_T_chip2img(cx) ) + T_data
            fpts = cm.get_fpts(cx)
            if fsel is None: fsel = range(len(fpts))
            # ---DEVHACK---
            # Randomly sample the keypoints. (Except be sneaky)
            elif fsel == 'rand':
                # Get Relative Position
                minxy = fpts.min(0)[0:2]
                maxxy = fpts.max(0)[0:2]
                rel_pos = (fpts[:,0]-minxy[0])/(maxxy[0]-minxy[0])
                to_zero = 1 - np.abs(rel_pos - .5)/.5
                pdf = (to_zero / to_zero.sum())
                # Transform Relative Position to Probabilities
                # making it more likely to pick a centerpoint
                fsel = np.random.choice(xrange(len(fpts)), size=88, replace=False, p=pdf)
            # ---/DEVHACK---
            # Plot ellipses
            if ell_bit and len(fpts) > 0 and len(fsel) > 0:
                ells = dm._get_fpt_ell_collection(fpts[fsel,:],
                                                  T_fpts,
                                                  ell_alpha,
                                                  color)
                ax.add_collection(ells)
            # Plot xy points
            if xy_bit and len(fpts) > 0 and len(fsel) > 0:
                ax.plot(fpts[fsel,0], fpts[fsel,1], 'o',\
                        markeredgecolor=color,\
                        markerfacecolor=color,\
                        transform=T_fpts,\
                        markersize=2)
        # ===
        if bbox_bit:
            # Draw Bounding Rectangle in Image Coords
            [rx,ry,rw,rh] = cm.cx2_roi[cx]
            rxy = (rx,ry)
            # Convert to Chip Coords if needbe
            T_bbox = T_data if in_image_bit else\
                Affine2D( np.linalg.inv(cm.cx2_T_chip2img(cx)) ) + T_data
            bbox = Rectangle(rxy,rw,rh,transform=T_bbox)
            # Visual Properties
            bbox.set_fill(False)
            bbox.set_edgecolor(color)
            ax.add_patch(bbox)

            # Draw Text Annotation
            cid   = cm.cx2_cid[cx]
            name  = cm.cx2_name(cx)
            # Lower the value to .2 for the background color and set alpha=.7
            rgb_textFG = [1,1,1]
            hsv_textBG = colorsys.rgb_to_hsv(*color)[0:2]+(.2,)
            rgb_textBG = colorsys.hsv_to_rgb(*hsv_textBG)+(.7,)
            # Draw Orientation Backwards
            degrees = 0 if not in_image_bit else -cm.cx2_theta[cx]*180/np.pi
            txy = (0,rh) if not in_image_bit else (rx, ry+rh)

            chip_text =  'name='+name+'\n'+'cid='+str(cid)
            ax.text(txy[0]+1, txy[1]+1, chip_text,
                    horizontalalignment ='left',
                    verticalalignment   ='top',
                    transform           =T_data,
                    rotation            =degrees,
                    color               =rgb_textFG,
                    backgroundcolor     =rgb_textBG)
        return fsel


#rodar com pyton 3.8(base)

import numpy as np
from PIL import Image
from mtcnn.mtcnn import MTCNN
from tensorflow.keras.models import load_model
import cv2

pessoas = ["MATHEUS", "MAYKI"]
num_classes = len(pessoas)

cap = cv2.VideoCapture(0)#0

detector = MTCNN()
facenet = load_model("facenet_keras.h5")
model = load_model("modelo_faces.h5")

def extract_face(image, box, required_size=(160,160)):

    pixels = np.asarray(image)

    x1, y1, width, height = box
    x2, y2 = x1 + width, y1 + height

    face = pixels[y1:y2, x1:x2]

    image = Image.fromarray(face)
    image = image.resize(required_size)
    return np.asarray(image)

def get_embedding(facenet, face_pixels):
    face_pixels = face_pixels.astype('float32')

    mean, std = face_pixels.mean(), face_pixels.std()
    face_pixels = (face_pixels - mean) / std

    samples = np.expand_dims(face_pixels, axis=0)

    yhat = facenet.predict(samples)

    return yhat[0]

while True:
    _, frame = cap.read()
    faces = detector.detect_faces(frame)

    for face in faces:
        confidence = face['confidence'] * 100

        if confidence >= 98:
            x1, y1, w, h = face['box']
            face = extract_face(frame, face['box'])

            face = face.astype("float32") / 255
            emb = get_embedding(facenet, face)
            tensor = np.expand_dims(emb, axis=0)

            classe = model.predict_classes(tensor)[0]
            prob = model.predict_proba(tensor)
            prob = prob[0][classe] * 100
            user = str(pessoas[classe]).upper()

            color = (192, 255, 119)#BGR

            cv2.rectangle(frame, (x1, y1), (x1+w, y1+h), color, 2)

            font = cv2.FONT_HERSHEY_SIMPLEX
            font_scale = 0.5

            cv2.putText(frame, user, (x1, y1-10), font, fontScale=font_scale, color=color, thickness=1)

    cv2.imshow("FACE RECOGNITION", frame)
    key = cv2.waitKey(1)

    if key==27:
        break

cap.release()
cv2.destroyAllWindows()


from shutil import rmtree
import json
import numpy as np
import traceback
import pydub
from pydub.playback import play
from os import listdir
import pygame
from pythonosc import udp_client

PORT_NUMBER = 3333
IP_NUM = '192.168.43.144'
client = udp_client.SimpleUDPClient(IP_NUM, PORT_NUMBER)

folder = 'guitar-electric'

# initialize the sounds
sounds = []

pygame.mixer.init()

# iterate over each of the sounds
for file in listdir('./../sound/' + folder):

    # create the new sound and save it
    sound = pygame.mixer.Sound('./../sound/' + folder + '/' + file)

    # save the sound
    sounds.append(sound)

# create the key points
KEY_POINTS = [
            "Nose",
            "Neck",
            "RShoulder",
            "RElbow",
            "RWrist",
            "LShoulder",
            "LElbow",
            "LWrist",
            "MidHip",
            "RHip",
            "RKnee",
            "RAnkle",
            "LHip",
            "LKnee",
            "LAnkle",
            "REye",
            "LEye",
            "REar",
            "LEar",
            "LBigToe",
            "LSmallToe",
            "LHeel",
            "RBigToe",
            "RSmallToe",
            "RHeel",
            "Background"
        ]

# try to erase the previous data
try: rmtree('./../data/live')
except Exception as e: pass

# initialize count and debounce
count = 0
debounce = 10000

# create a list for right wrist points
RWrists = []

# continuously search for new json files
while True:
    try:
        # try to open the next json file as f
        with open('./../data/live/' + str(count).zfill(12) + '_' + 'keypoints.json') as f:

            # read the poeple from the json object
            people = json.load(f)['people']

            # get the pose key points
            data = people[0]['pose_keypoints_2d'] if people else [0 for _ in range(75)]

            # get the indices for left hip and wrist and right wrist
            lh_idx, lw_idx, rw_idx = KEY_POINTS.index('LHip'), KEY_POINTS.index('RWrist'), KEY_POINTS.index('LWrist')

            # get the current left hip and left wrist points
            LHip = np.array(data[lh_idx * 3: lh_idx * 3 + 2])
            LWrist = np.array(data[lw_idx * 3: lw_idx * 3 + 2])

            # get the current right wrist
            RWrists.append(np.array(data[rw_idx * 3: rw_idx * 3 + 2]))

            # if the length of wrists is more than 5
            if len(RWrists) > 200:

                # remove the first element
                RWrists.pop(0)

            # calculate the wrist velocity for the last half second
            RWristVelocity = np.array([0.0, 0.0]) if len(RWrists) < 5 else (RWrists[-1] - RWrists[0]) / 0.4

            # calculate the left wrist difference from the hip
            handDistance = np.linalg.norm(LHip - LWrist)

            # if the right wrist's downward velocity is more than 0.2
            if RWristVelocity[1] > 0.05 and debounce < 0 and handDistance > 0.05:

                i_prev = 0.05

                # iterate over the sounds
                for sound, j, i, note in zip(sounds, range(len(sounds)), np.linspace(0.2, 0.8, len(sounds)), listdir('./../sound/' + folder)):

                    # if the hand distance is in the right range
                    if i_prev < handDistance < i:

                        # set the sound volume based on strum speed
                        # sound.set_volume(np.linalg.norm(RWristVelocity))

                        client.send_message('/hit', j + 1)

                        # play the sound
                        sound.play()

                        # print the strum
                        print('Strum:', note.split('.')[0], np.linalg.norm(RWristVelocity))

                    i_prev = i

                # reset the debounce
                debounce = 1000

            # increment count and decrement debounce
            debounce -= 1
            count += 1

    except FileNotFoundError:
        if count > 0:
            count -= 1

    except json.JSONDecodeError as e:
        pass

    except Exception as e:
        traceback.print_exc()


# Loading the libraries we need
import numpy as np
import cv2
import time

#Creating a VideoCapture object to read video from the primary camera
cap=cv2.VideoCapture(0)

#Creating a VideoWriterObject to save the output video
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter('wizard_smaran.avi' , fourcc, 20.0, (640,480))

#allow the system to sleep for 3 sec before webcam starts
time.sleep(2)

# Capture the background in range of 30 or 60 without you!
background = 0
for i in range(30):
    ret, background = cap.read()#capturing image

#Now we capture you in real time!
while(cap.isOpened()):
    ret, img = cap.read()

    if not ret:
        break

    # Converting the color space from BGR to HSV as BGR is more sensitive to light
    hsv=cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    #Generating mask for red range 1(0-10)
    lower_red = np.array([0,120,70])
    upper_red = np.array([10,255,255])
    mask1 = cv2.inRange(hsv , lower_red , upper_red)

    #Generating mask for red range 1(170-180)
    lower_red = np.array([170,120,70])
    upper_red = np.array([180,255,255])
    mask2 = cv2.inRange(hsv , lower_red , upper_red)

    #Combining the masks obtained for both the ranges
    mask1 = mask1 + mask2

    mask1=cv2.morphologyEx(mask1, cv2.MORPH_OPEN ,np.ones((3,3) , np.uint8) , iterations=2)
    mask2=cv2.morphologyEx(mask1, cv2.MORPH_DILATE ,np.ones((3,3) , np.uint8) , iterations=1)

    # Segmenting out cloth color
    mask2 = cv2.bitwise_not(mask1)

    # Segment the red color part out of the frame using bitwise and with the inverted mask
    layer1 = cv2.bitwise_and(background, background, mask=mask1)
    # Create image showing static background frame pixels only for the masked region
    layer2 = cv2.bitwise_and(img, img, mask=mask2)

    final_output = cv2.addWeighted(layer1 , 1, layer2 , 1, 0)

    cv2.imshow('Invisible Smaran' , final_output)
    k=cv2.waitKey(10)

    #Keyboard Interupt
    if k==27:
        break

cap.release()
cv2.destroyAllWindows()


import pathlib
import re
import argparse
import yaml
from Bio.PDB import PDBParser
from Bio import SeqIO
import os
import numpy as np
import pandas as pd
import helpers as nhp
from LatticeModelComparison import LatticeModelComparison
from ParallelTempering import ParallelTempering

pdb_parser = PDBParser()
__location__ = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__)))

parser = argparse.ArgumentParser(description='Generate cubic lattice models for a (list of) pdb names or AA seqs')
# --- IO ---
parser.add_argument('--in-dir', type=str,
                    help='entity/fasta file or directory containing entity/fasta files of which to generate models')
parser.add_argument('--cm-pdb-dir', type=str, default=None,
                    help='center-of-mass pdbs directory, for finetuning of starting structures. Necessary if --finetune-structure is on')
parser.add_argument('--finetune-structure', action='store_true',
                    help='Minimize lattice model RMSD w.r.t. center-of-mass structure provided through cm-pdb-dir by the same ID')
# --- model parameters ---
parser.add_argument('--temp-range', type=float, nargs=2,
                    help='Use a range of equally spaced temperatures between given values in a parallel tempering search.')
parser.add_argument('--tagged-resn', type=str, default='',
                    help='Define residue(s) with 1-letter-code that should be replaced with a very hydrophillic tag.')
parser.add_argument('--pairs-mat', type=str, default=f'{__location__}/potential_matrices/aa_water2_abeln2011.txt',
                    help='Matrix with pair potentials for interactions between residues and residues vs solvents.')
parser.add_argument('--labeling-model', type=str, default='perfect',
                    help='Yaml file containing labeling probability for each residue, for each labeling chemistry '
                         '(no value == no labeling). Supply "standard" for regular model or "perfect" for '
                         'no mislabeling or path to yaml [default: perfect].')
parser.add_argument('--experimental-mode', type=int, default=0,
                    help='Several experimental settings under 1 switch, numeric.')
# --- lattice properties ---
parser.add_argument('--lattice-type', type=str, default='cubic', choices=['cubic', 'bcc', 'bcc_old'],
                    help='Set type of lattice to use. Choices: cubic, bcc [ default: cubic]')
parser.add_argument('--starting-structure', type=str,
                    choices=['free_random', 'anchored_random', 'stretched', 'serrated'], default='stretched',
                    help='Choose how to initialize structure [default: stretched]')

# --- model iterations parameters ---
parser.add_argument('--nb-steps', type=int, default=1,
                    help='number of mutations to perform at each MC iteration.')
parser.add_argument('--accomodate-tags', action='store_true',
                    help='Run MC iterations until tag penalty is 0.')
parser.add_argument('--iters-per-tempswap', default=1000, type=int,
                    help='If using parallel tempering, define per how many rounds a temperature swap should be performed')
parser.add_argument('--mc-iters', type=int, default=1000,
                    help='number Monte Carlo iterations to perform for each model.')
parser.add_argument('--nb-models', type=int, default=5,
                    help='number of models to create per AA sequence.')
parser.add_argument('--nb-processes', type=int, default=4,
                    help='Define how many processes to engage at once in parallel tempering.')
parser.add_argument('--early-stopping-rounds', type=int, default=-1,
                    help='Number of consecutive temperature swap rounds without improvement in energy'
                         'before training is stopped. Supply -1 to disable early stopping.')
parser.add_argument('--restart', action='store_true',
                    help='Restart previous run; if folder with structures carrying same id exists, start with'
                         'next model index. If not, start with 0.')
parser.add_argument('--start-idx', type=int,
                    help='Models should be generated starting with this index. Overrides --restart!')
parser.add_argument('--no-regularization', action='store_true',
                    help='Do not add regularization term to energy function.')
parser.add_argument('--free-sampling', action='store_true',
                    help='Do not optimize structure during snapshot generation, accept all')

# --- Result saving options ---
parser.add_argument('--out-dir', type=str, required=True,
                    help='Location where model pdb files are stored.')
parser.add_argument('--store-rg', type=str, choices=['off', 'tswap', 'full'], default='tswap', required=False,
                    help='Store the estimated radius of gyration at each training step in the produced pdb.')
parser.add_argument('--save-all-pdbs', action='store_true',  # todo superfluous, remove
                    help='Save all pdbs at intermediate steps as well, instead of keeping only the eventual best pdb')
parser.add_argument('--snapshots', nargs=2, type=int, default=[0,0],
                    help='If given [n,s], saves n snapshots with s steps in between after convergence/end of run.')
parser.add_argument('--save-intermediate-structures', action='store_true',
                    help='Save structure after temperature swaps in this pdb file')
parser.add_argument('--store-energies', action='store_true',
                    help='Store base energy and individual contributions to energy in tsv file')

# --- Temporary debugging options ---
parser.add_argument('--equidist-temps', action='store_true',
                    help='Temperatures equidistance spaced in T range (i.o. beta range).')

args = parser.parse_args()

if args.out_dir[-1] != '/': args.out_dir += '/'

pairs_mat = nhp.get_pairs_mat(args.pairs_mat)
ent_list = nhp.parse_input_dir(args.in_dir, '*.npz')
# ent_list = nhp.parse_input_dir(args.in_dir, '*.ent')
# if len(ent_list) == 0: ent_list = nhp.parse_input_dir(args.in_dir, '*.fasta')
# if len(ent_list) == 0: ent_list = nhp.parse_input_dir(args.in_dir, '*.pdb')
# if len(ent_list) == 0: ent_list = nhp.parse_input_dir(args.in_dir, '*.npz')
if len(ent_list) == 0: raise ValueError('No npz files found')

# Load a labeling model, describing the probability of labeling a given residue for a given labeling chemistry
# Default: perfect labeling
labeling_model = args.labeling_model
if labeling_model == 'standard':
    labeling_model = f'{__location__}/data/labeling_models/standard.yml'
elif labeling_model == 'perfect':
    labeling_model = f'{__location__}/data/labeling_models/perfect.yml'
with open(labeling_model, 'r') as fh: labeling_model = yaml.load(fh, yaml.SafeLoader)

# Load cm pdbs, if available
cm_dict = {}
if args.cm_pdb_dir is not None:
    cm_list = nhp.parse_input_dir(args.cm_pdb_dir, '*.pdb')
    cm_dict = {os.path.splitext(os.path.basename(cmd))[0].split('_')[0]:
                   np.vstack([atm.coord for atm in pdb_parser.get_structure('cm', cmd).get_atoms() if atm.name == 'CA'])
               for cmd in cm_list}
elif args.finetune_structure:
    raise ValueError('Must provide --cm-pdb-dir if finetuning structure.')

# Loop over protein files
for ent in ent_list:
    pdb_id, ext = os.path.splitext(os.path.basename(ent))
    out_dir_pdb = f'{args.out_dir}{pdb_id}/'
    pathlib.Path(out_dir_pdb).mkdir(parents=True, exist_ok=True)
    model_idx_start = 0
    nb_models = args.nb_models  # For each protein [nb_models] model runs are performed

    if args.finetune_structure:
        if pdb_id.split('_')[0] not in cm_dict:
            print(f'Cannot find cm structure {pdb_id} for polishing, skipping...')

    if args.start_idx:  # Start counting at a different index than 0 - only used for snakemake implementation
        model_idx_start = args.start_idx
        nb_models = model_idx_start + nb_models
    elif args.restart:  # Allow restarting an interrupted run, by checking final output pdb for a given index
        existing_fn_list = nhp.parse_input_dir(out_dir_pdb, pattern='*.pdb')
        highest_mod_idx = -1
        for pdb_fn in existing_fn_list:
            re_obj = re.search(f'(?<={pdb_id}_)[0-9]+(?=.pdb)', pdb_fn)
            if re_obj: highest_mod_idx = max(highest_mod_idx, int(re_obj.group(0)))
        if highest_mod_idx + 1 >= nb_models:
            print(f'Restart for {pdb_id} not required: model {model_idx_start} of {nb_models} was already produced.')
            continue
        model_idx_start = highest_mod_idx + 1

    tagged_resn = '' if args.tagged_resn == 'None' else args.tagged_resn
    with np.load(ent) as fh:
        aa_seq = fh['sequence']
        coords = fh['coords'].astype(int)
        ss_df = pd.DataFrame(fh['secondary_structure'], columns=['H', 'S', 'L'])

    if args.accomodate_tags and os.path.exists(f'{out_dir_pdb}{pdb_id}_{args.start_idx}_unoptimizedTags.npz'):
        is_retry = True
        print(f'Reloading previous attempt at tag accomodation')
        with np.load(f'{out_dir_pdb}{pdb_id}_{args.start_idx}_unoptimizedTags.npz', allow_pickle=True) as fh:
            # aa_seq_original = aa_seq
            aa_seq = fh['sequence']
            coords = fh['coords'].astype(int)
            ss_df = pd.DataFrame(fh['secondary_structure'], columns=['H', 'S', 'L'])
            tagged_resi = fh['tagged_resi'][()]
            # tagged_resi = {tr:[] for tr in fh['tagged_resn']}
            # for ri, resn in enumerate(aa_seq):
            #     if resn == 'TAG': tagged_resi[aa_seq_original[ri]].append(ri)
                # tagged_resi = {: [ri for ri, resn in enumerate(aa_seq) if resn == 'TAG']}
    else:
        is_retry = False

    # Set temperatures at which parallel-tempered models are run
    if args.equidist_temps:
        # option 1: temperatures are linearly spaced as temperatures and then converted to beta-values
        temp_list = np.linspace(min(args.temp_range), max(args.temp_range), num=args.nb_processes)
        beta_list = 0.01 / temp_list
    else:
        # option 2: temperatures are linearly spaced as beta-values
        beta_min = 0.01 / max(args.temp_range)
        beta_max = 0.01 / min(args.temp_range)
        if args.experimental_mode == 8:
            beta_list = np.concatenate((np.linspace(beta_min, beta_max, num=args.nb_processes-1)[::-1], [np.nan]))
        else:
            beta_list = np.linspace(beta_min, beta_max, num=args.nb_processes)[::-1]

    # Construct model
    for ns in range(model_idx_start, nb_models):
        print(f'generating new model at {out_dir_pdb}{pdb_id}_{ns}.pdb ')
        if os.path.exists(f'{out_dir_pdb}{pdb_id}_{ns}.pdb'):
            print(f'skipping {pdb_id}_{ns}, already exists')
            continue

        print(f'{nhp.print_timestamp()} Generating model {ns+1} for {pdb_id} ')

        # Determine which residues are labeled. Here labeling model is applied (i.e. incomplete/mis-labeling takes place)
        if not is_retry:
            tagged_resi = nhp.get_tagged_resi(labeling_model, tagged_resn, aa_seq)

        if args.nb_processes > 1:  # Run parallel tempering with [args.nb_processes] chains
            pt = ParallelTempering(
                # --- model ID args ---
                pdb_id=pdb_id,  # unique id for model
                model_nb=ns,    # Idx of this model run
                # --- model run params ---
                aa_sequence=aa_seq,                                # Amino acid sequence in 1-letter array
                tagged_resn=tagged_resn,
                beta_list=beta_list,                               # beta values at which to run parallel tempering chains (i.e. 0.01/T)
                nb_processes=args.nb_processes,                    # Number of processes/cores to engage at once (== number of parallel tempering chains)

                lattice_type=args.lattice_type,                    # Type of lattice to adhere to (e.g. cubic or bcc)
                coords=coords,                                     # lattice CA starting coordinates, [L x 3] int np array, or None if starting coordinates should be generated
                starting_structure=args.starting_structure,        # If no starting coords are provided, indicate how starting coords should be generated (e.g. 'stretched' or 'random')

                nb_steps=args.nb_steps,                            # Number of mutations to apply at each MC iteration
                nb_iters=args.mc_iters,                            # Total number of MC iterations
                iters_per_tempswap=args.iters_per_tempswap,        # Number of MC iterations to perform before swapping models between parallel tempering chains

                early_stopping_rounds=args.early_stopping_rounds,  # number of rounds in which no better model is accepted before stopping early, -1 to disable

                experimental_mode=args.experimental_mode,
                cm_coords=cm_dict.get(pdb_id.split('_')[0], None),
                finetune_structure=args.finetune_structure,
                # -- E modifiers ---
                pairs_mat=pairs_mat,        # pandas df containing energy modifier for each combination of residue types + water interactions
                secondary_structure=ss_df,  # secondary structure energy modifiers, [L x 3] pandas df with E modifier for H(elix)/S(trand)/L(oop) resp. in columns
                tagged_resi=tagged_resi,    # tagged residue indices
                # --- output args ---
                save_dir=out_dir_pdb,                                           # directory to save results to
                store_rg=args.store_rg,                                         # Store radius of gyration at each temp swap ('tswap'), each MC iteration ('full') or not at all ('off')
                snapshots=args.snapshots,                                       # Tuple [N, T] denoting to make N snapshots of final structure spaced by T MC iterations
                save_intermediate_structures=args.save_intermediate_structures, # if True, save structure for each temp swap in separate file
                store_energies=args.store_energies,                              # if True, store energy for each temp swap in separate file
                no_regularization=args.no_regularization,

                accomodate_tags=args.accomodate_tags,
                free_sampling=args.free_sampling
            )
            pt.do_mc_parallel()
        else:  # If one process available, run single chain at lowest temperature
            beta = 0.01 / args.temp_range[0]
            intermediate_fn = None
            if args.save_intermediate_structures:
                intermediate_fn = f'{out_dir_pdb}{pdb_id}_{ns}_intermediates.pdb'
            lmc = LatticeModelComparison(
                pdb_id=pdb_id,
                # --- model ID args ---
                mod_id=ns, # Idx for this model run
                # --- model run params ---
                aa_sequence=aa_seq,  # Amino acid sequence in 1-letter array
                beta=beta,           # Beta value at which to run chain (i.e. 0.01/T)

                lattice_type=args.lattice_type,             # Type of lattice to adhere to (e.g. cubic or bcc)
                coords=coords,                              # Lattice CA starting coordinates, [L x 3] int np array, or None if starting coordinates should be generated
                starting_structure=args.starting_structure, # If no starting coords are provided, indicate how starting coords should be generated (e.g. 'stretched' or 'random')

                nb_steps=args.nb_steps,  # Number of mutations to apply at each MC iteration
                experimental_mode=args.experimental_mode,
                # --- E modifiers ---
                pairs_mat=pairs_mat,        # pandas df containing energy modifier for each combination of residue types + water interactions
                secondary_structure=ss_df,  # secondary structure energy modifiers, [L x 3] pandas df with E modifier for H(elix)/S(trand)/L(oop) resp. in columns
                tagged_resi=tagged_resi,    # tagged residue indices
                # --- output args ---
                store_rg=args.store_rg,  # Store radius of gyration at each temp swap ('tswap'), each MC iteration ('full') or not at all ('off')
                stepwise_intermediates=intermediate_fn,
                no_regularization=args.no_regularization
            )
            lmc.do_mc(args.mc_iters)
            lmc.make_snapshots(args.snapshots[0], args.snapshots[1], [lmc.best_model.rg],
                               f'{out_dir_pdb}{pdb_id}_{ns}.pdb', free_sampling=args.free_sampling)


"""
Executes the parallelization of Stochastic DiNT through MPI4Py
"""

import os
import os.path
import sys
import math
import numpy as np

#take the current main() function/subroutine and move that code
#   into a main_run_DiNT() function/subroutine that main() now calls.
#Then, one can create a libDiNT.a library that a new lightweight MPI code
#   could link in where each MPI rank spawned would call main_run_DiNT()
#   to launch its own instance of a DiNT calculation.
#Making sure each instance of DiNT had unique input and output is the
#   next step, and then it’s a matter of optimizing performance,
#   which in this case is usually turning off all unnecessary writing of
#   output and/or refactoring how output gets written so that it’s
#   efficient when 100,000s of DiNTs are all running (e.g. could capture
#   important output in memory and write single output file in parallel).

from mpi4py import MPI
#from libDiNT.a import dint
#or use f2py?
import stoch_driver.collisions as collisions
import stoch_driver.zonz as zonz
import stoch_driver.constants as constants

class Trajectory:
    """
    A representation of a stochastic trajectory.
    =========================== ================================================
    Attribute                   Description
    =========================== ================================================
    `rank`                      MPI process rank assigned to this trajectory
    `A`                         A list of the reactants present in the system
    `X`                         A list of the reactive collider species in the system
    `M`                         A list of the nonreactive collider species in the system
    `natoms`                    An array of the number of atoms in associated A species
    `temp`                      Temperature of the system in Kelvin
    `xA`                        Number of molar equivalents of each A species
    `xX`                        Number of molar equivalents of each X species
    `xM`                        Number of molar equivalents of each M species
    `steps`                     Maximum number of Monte Carlo steps the simulation should take
    `Ei_list`                   A list of the initial energies for each step in the trajectory
    `Ef_list`                   A list of the final energies for each step in the trajectory
    `Ji_list`                   A list of the initial J for each step in the trajectory
    `Jf_list`                   A list of the final J for each step in the trajectory
    --------------------------- ----------------------------------------------------
    `xAn`                       Normalized mole fraction of A species
    `xXn`                       Normalized mole fraction of A species
    `xMn`                       Normalized mole fraction of A species
    =========================== ================================================
    """
    def __init__(self, label, A, X, M, natomsA, natomsX, natomsM,
                 xA, xX, xM, temperature, steps):
        self.label = label
        self.A = A
        self.X = X
        self.M = M
        self.natomsA = natomsA
        self.natomsX = natomsX
        self.natomsM = natomsM
        self.temperature = temperature
        self.xA = xA
        self.xX = xX
        self.xM = xM
        self.steps = steps

        self.outputDirectory = join(os.path.abspath(outputDirectory),"b",rank)

        self.Ei_list = []
        self.Ji_list = []
        self.Ef_list = []
        self.Jf_list = []

#    def species(self):

        #for i in range(len(A)):
        #    self.A[i] = A[i]
        #for i in range(len(X)):
        #    self.X[i] = X[i]
        #for i in range(len(M)):
        #    self.M[i] = M[i]

#        return A, X, M

#    def moleFracs(self):

        #for i in range(len(xA)):
        #    self.xA[i] = xA[i]
        #for i in range(len(xA)):
        #    self.xX[i] = xX[i]
        #for i in range(len(xA)):
        #    self.xM[i] = xM[i]

#        return xA, xX, xM

    def initEJ(self, ei, ji):
        self.Ei_list.append(ei)
        self.Ji_list.append(ji)

    def updateEJ(self, ef, jf):
        self.Ef_list.append(ef)
        self.Jf_list.append(jf)

    def executeStochasticDynamics(self, rates):
        # Set path
        DRIVE_PATH = os.getcwd()

        xA = self.xA
        xX = self.xX
        xM = self.xM
        zhsA = rates.zhsA
        zhsX = rates.zhsX
        zhsM = rates.zhsM
        zljA = rates.zljA
        zljX = rates.zljX
        zljM = rates.zljM
        bmaxA = rates.bmaxA
        bmaxX = rates.bmaxX
        bmaxM = rates.bmaxM

        # Initialize trajectory objects
        collisions.normMoleFracs(xA, xX, xM, zhsA, zhsX, zhsM,
                                 zljA, zljX, zljM, bmaxA, bmaxX, bmaxM)

        # Initialize E, J values
        zonz.pEJ(temp)

        ef = ee*constants.cmtokcalmol + 4.5*constants.evtokcalmol
        self.initEJ(ei = ef, ji = jf)
        self.Ei_list.append(ei)
        self.Ji_list.append(ji)

        # Continue stochastic dynamics until reaction occurs or max steps reached:
        while nn[0] < maxsteps:
            x = random.randrange(0,1)
            it = 0
            test = pp/ps
            while (x > test):
                it += 1
                test += pp[it]/ps
            nn[it] += 1

            if (it == 0): # Counter
                # print to specific traj file
                print(nn[0],ji,ei,jf,ef)
            elif (it == 1): # A + M
                collisions.prepAMTraj(ef,jf,temp)    # creates new input file for A+M
                print("Run dint-alkT here")
#                dint.dint-alkT ### call appropriate dint exe
            elif (it == 2): # A + X
                collisions.prepAXTraj(nmodes,ef,jf,temp,bmax)   # creates new input file for A+X
                print("Run dint-ch4hALL here")
#                dint.dint-ch4hALL ### call appropriate dint exe
            else:
                print("Error: Iteration not found")
                break

            with open('fort.31') as f:
                for line in f:
                    pass    # grab last line from fort.31
                f31 = split(line)

            ji = f31[12]
            ei = (f31[13]+f31[14]+f31[15])*constants.eVtokcalmol
            jf = f31[24]
            ef = (f31[25]+f31[26]+f31[27])*constants.eVtokcalmol

            self.updateEJ(ef = ef, jf = jf)
            self.Ef_list.append(ef)
            self.Jf_list.append(jf)

#            print("Traj b",i,": Timestep ",nn[0]," , M collision ",nn[it]," ( ",ji," , ",ei," ) --> ( ",jf," , ",ef" )")

class Reactants:
    def __init__(self, A, natomsA, xA, B):
        self.A = A
        self.natomsA = natomsA
        self.xA = xA
        self.B = B

class xColliders:
    def __init__(self, X, natomsX, xX):
        self.X = X
        self.natomsX = natomsX
        self.xX = xX

class mColliders:
    def __init__(self, M, natomsM, xM):
        self.M = M
        self.natomsM = natomsM
        self.xM = xM


import torch
from torch import nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
from dql.helper import Transition
from dql.helper import get_latest_model_path
# from torch.utils.tensorboard import SummaryWriter
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class ReplayMemory:
    def __init__(self, capacity, batch_size):
        self.capacity = capacity
        self.memory = []
        self.batch_size = batch_size
        self.position = 0

    def push(self, *args):
        if len(self.memory) < self.capacity:
            self.memory.append(None)
        self.memory[self.position] = Transition(*args)
        self.position = (self.position + 1) % self.capacity
    # did find a torch equivalent to np.random.choice :(

    def sample(self):
        random_idx = np.random.choice(
            np.arange(len(self.memory)), size=self.batch_size)
        res = [self.memory[i] for i in random_idx]
        return res

    def __len__(self):
        return len(self.memory)

# the architecture for both policy net and target net

class DQN(nn.Module):
    def __init__(self, h, w, n_actions, frame_length):
        super(DQN, self).__init__()
        self.n_actions = n_actions
        self.conv1 = nn.Conv2d(frame_length, 16, kernel_size=5, stride=2)
        self.bn1 = nn.BatchNorm2d(16)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=5, stride=2)
        self.bn2 = nn.BatchNorm2d(32)
        self.conv3 = nn.Conv2d(32, 32, kernel_size=5, stride=2)
        self.bn3 = nn.BatchNorm2d(32)

        def conv2d_size_out(size, kernel_size=5, stride=2):
            return (size - kernel_size) // stride + 1
        convw = conv2d_size_out(conv2d_size_out(conv2d_size_out(w)))
        convh = conv2d_size_out(conv2d_size_out(conv2d_size_out(h)))
        linear_input_size = convw * convh * 32
        self.output_layer = nn.Linear(linear_input_size, n_actions)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        return self.output_layer(x.view(x.size(0), -1))

    def eps_greedy_policy(self, state, eps):
        if np.random.sample() > eps:
            with torch.no_grad():
                return self.greedy_policy(state)
        else:
            return torch.tensor([[np.random.choice(np.arange(self.n_actions))]], device=device, dtype=torch.long)
    def greedy_policy(self, state):
        return self(state).argmax(1)

class Model:
    def __init__(self, env, hps):
        self.hps = hps
        # self.summary_writter = SummaryWriter(
        #     log_dir='runs/{}'.format(self.hps['envname']))
        self.policy_net = DQN(
            *self.hps['image_size'], env.action_space.n, self.hps['frame_length'])
        self.target_net = DQN(
            *self.hps['image_size'], env.action_space.n, self.hps['frame_length'])
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.loss = torch.empty(1)
        self.optimizer = optim.RMSprop(self.policy_net.parameters())
        self.target_net.eval()
        self.policy_net.train()
        return

    def optimize_model(self, batch):
        state_batch = torch.stack(batch.state)
        action_batch = torch.stack(batch.action)
        reward_batch = torch.stack(batch.reward)
        next_state_batch = torch.stack(batch.next_state)
        done_batch = torch.tensor(batch.done, device=device)
        current_q = self.policy_net(state_batch).gather(
            1, action_batch.unsqueeze(1)).squeeze()
        next_q = self.target_net(next_state_batch)
        TD_target_batch = reward_batch + \
            (~done_batch).type(torch.float32) * \
            self.hps['gamma'] * next_q.max(dim=1).values
        self.loss = F.smooth_l1_loss(current_q, TD_target_batch)
        self.optimizer.zero_grad()
        self.loss.backward()
        self.optimizer.step()
        return self.loss

    def update_target(self):
        self.target_net.load_state_dict(self.policy_net.state_dict())
        return
    def load_checkpoint(self):
        model_ckpt_path, latest_episode_num = get_latest_model_path(self.hps)
        checkpoint = torch.load(model_ckpt_path)
        self.optimizer.load_state_dict(checkpoint['optimizer'])
        self.policy_net.load_state_dict(checkpoint['policynet'])
        self.update_target()
        # self.target_net.load_state_dict(checkpoint['target_net'])
        return latest_episode_num


#!/usr/bin/python
#Halfling

import numpy as np
import pandas as pd
import random

def LuckyHalf():
    Name = "Lucky (Halfling)"
    Feature = "When you roll a 1 on an attack roll, ability check, or saving throw, you can reroll the die and must use the new roll."

	return Name, Feature

def Brave():
    Name = "Brave"
    Feature = "You have advantage on saving throws against being frightened."

	return Name, Feature

def NimbleHalf():
    Name = "Halfling Nimbleness"
    Feature = "You can move through the space of any creature that is of a size larger than yours."

	return Name, Feature

def NaturalStealth():
    Name = "Naturally Stealthy"
    Feature = "You can attempt to hide even when you are obscured only by a creature that is at least one size larger than you."

	return Name, Feature

def StoutResil():
    Name = "Stout Resilience"
    Feature = "You have advantage on saving throws against poison, and you have resistance against poison damage."

	return Name, Feature


import os, math, tempfile, gc, requests
import librosa
import numpy as np
from flask import jsonify

# Google AI
from googleapiclient import discovery

ALLOWED_EXTENSIONS = ['aiff', 'wav', 'mp3']
GCP_PROJECT = 'voices-to-emotions'
GCP_MODEL = 'emotionrecognition'
MFCC_FEATURES = 12
MFCC_LENGTH = 200
SILENCE_TRESHOLD = 50
AI_VERSION = 'v2'

emotion_dict = {
    0: 'neutral',
    1: 'calm',
    2: 'happy',
    3: 'sad',
    4: 'angry',
    5: 'fearful',
    6: 'disgust',
    7: 'surprised'
}

def save_from_uri(uri):
    req = requests.get(uri)
    if req.status_code == 200:
        temp_dir = tempfile.gettempdir()
        extension = uri.rsplit('.', 1)[1].lower()
        file_path = os.path.join(temp_dir, 'temp.' + extension)

        print(req.encoding)
        print(req.content)
        if req.encoding == 'UTF-8':
            raise Exception('Invalid file type')

        with open(file_path, 'wb') as f:
            f.write(req.content)

        return file_path, extension

    raise Exception('Invalid file type')

def map_predictions(predictions):
    res = []
    for p in predictions:
        classes = p["classes"]
        mapped_classes = {}

        for i in range(0, len(classes)):
            mapped_classes[emotion_dict[i]] = round(classes[i] * 10000) / 10000

        res.append(mapped_classes)

    return res

def get_predictions(instances):
    service = discovery.build('ml', 'v1')
    name = 'projects/{}/models/{}/versions/{}'.format(GCP_PROJECT, GCP_MODEL, AI_VERSION)

    response = service.projects().predict(
        name=name,
        body={'instances': instances}
    ).execute()

    if 'error' in response:
        raise RuntimeError(response['error'])

    return response['predictions']

def load_audio_data(file_path):
    # 16000 Hz = VoIP
    wave, sr = librosa.load(file_path, mono=True, sr=16000)
    wave_frag_offsets = librosa.effects.split(wave, top_db=SILENCE_TRESHOLD)

    results = []
    timestamps = []
    for offsets in wave_frag_offsets:
        start_sec = offsets[0] / sr
        last_secs = start_sec
        frag_duration = (offsets[1] - offsets[0]) / sr

        wave_fragment = wave[offsets[0]:offsets[1]]
        mfcc = librosa.feature.mfcc(wave_fragment, sr, n_mfcc=MFCC_FEATURES)
        _, y_size = mfcc.shape

        splitted_mfcc = np.array_split(mfcc, math.ceil(y_size / MFCC_LENGTH), axis=1)

        one_mfcc_secs = frag_duration / y_size

        for short_mfcc in splitted_mfcc:
            mfcc_len = len(short_mfcc[0])
            short_mfcc = np.pad(short_mfcc, ((0, 0), (0, MFCC_LENGTH - len(short_mfcc[0]))), mode='constant',
                                constant_values=0)
            results.append(np.array(short_mfcc).tolist())

            duration = one_mfcc_secs * mfcc_len
            timestamps.append([last_secs, last_secs + duration])
            last_secs += duration

    return results, timestamps

def mfcc(request):
    if not request.method == 'GET':
        return jsonify({
            "type": 'error',
            "message": "Unknown request"
        })

    uri = request.args['uri']
    if uri:
        file_path, extension = save_from_uri(uri)

        if extension in ALLOWED_EXTENSIONS:
            mfcc, timestamps = load_audio_data(file_path)
            gc.collect()

            predictions = map_predictions(get_predictions(mfcc))
            gc.collect()

            # Remove the file, KEEP THIS AT THE END!
            os.remove(file_path)

            return jsonify({
                "type": 'success',
                "data": {
                    'timestamps': timestamps,
                    'emotions': predictions
                }
            })

        else:
            return jsonify({
                "type": 'error',
                "message": "Invalid filetype"
            })

    else:
        return jsonify({
            "type": 'error',
            "message": "Please provide a file"
        })


import pygame
import numpy as np
from Source import UI_functions as UI
from Source import battleships_functions_bot as bfb
from Source import battleships_functions_check as bfc

def Play_Game(screen, bg, cfg):

    #Init
    screen, bg = UI.Update_Screen_Values(screen, bg)
    pygame.time.Clock().tick(cfg["Basic"].getint("FPS"))

    #Resources (Images, Icons, Fonts)
    font = pygame.font.Font("Assets/Font/impact.ttf", 26)
    font_b = pygame.font.Font("Assets/Font/impact.ttf", 40)
    font_c = pygame.font.SysFont('segoeuisymbol', 34)
    square = pygame.image.load("Assets/Images/Square.png")
    rectangle = pygame.image.load("Assets/Images/Rectangle.png")
    grid = pygame.image.load("Assets/Images/WhiteGrid.png")

    #Initial Values
    RUNNING = True
    SHOW_HAHA = True
    CANT_GENERATE1 = False
    CANT_GENERATE2 = False
    CLICK = False
    bfb.load_config_file(cfg)
    Bmap1 = np.zeros((cfg["Rules"].getint("Y_RANGE"),cfg["Rules"].getint("X_RANGE")), dtype = np.int32)
    Bmap2 = np.zeros((cfg["Rules"].getint("Y_RANGE"),cfg["Rules"].getint("X_RANGE")), dtype = np.int32)
    Bmap1, CANT_GENERATE1 = bfb.generate_bot_ships(Bmap1)
    Bmap2, CANT_GENERATE2 = bfb.generate_bot_ships(Bmap2)
    rects, images_pos, text_pos = UI.Rect_AI_AI_Set()
    rect_map = UI.Rect_Player_AI_Map()
    images = [rectangle, square, rectangle, rectangle]
    texts = [font.render("PLAY", True, (52, 52, 54)),
             font.render("χ", True, (52, 52, 54)),
             font.render("AI 1", True, (52, 52, 54)),
             font.render("AI 2", True, (52, 52, 54)),
             font_c.render("⟳", True, (52, 52, 54)),
             font_c.render("⟳", True, (52, 52, 54)),
             font_b.render(cfg["Text"]["AI1"], True, (255, 255, 255)),
             font_b.render(cfg["Text"]["AI2"], True, (255, 255, 255)),
             font_b.render(cfg["Text"]["SCORE"], True, (255, 255, 255)),
             font_b.render(str(cfg["Points"].getint("AI1_PTS")) + " - " + str(cfg["Points"].getint("AI2_PTS")), True, (255, 255, 255))]

    #InGame
    while RUNNING:

        #Screen properties per update
        mx, my = pygame.mouse.get_pos()
        screen.blit(bg,(0,0))

        #Draw functions
        UI.Draw_Left_Map_Set(screen, Bmap1, grid)
        UI.Draw_Right_Map_Set(screen, Bmap2, grid)
        UI.Draw_Pos(screen, images, images_pos)
        UI.Draw_Pos(screen, texts, text_pos)

        #Clickable buttons
        if CLICK:
            if rects[0].collidepoint((mx,my)):
                return Bmap1, Bmap2, True

        if CLICK:
            if rects[1].collidepoint((mx,my)):
                return None, None, False

            if rects[2].collidepoint((mx,my)):
                Bmap1 = np.zeros((cfg["Rules"].getint("Y_RANGE"),cfg["Rules"].getint("X_RANGE")), dtype = np.int32)
                Bmap1, CANT_GENERATE1 = bfb.generate_bot_ships(Bmap1)

            if rects[3].collidepoint((mx,my)):
                Bmap2 = np.zeros((cfg["Rules"].getint("Y_RANGE"),cfg["Rules"].getint("X_RANGE")), dtype = np.int32)
                Bmap2, CANT_GENERATE2 = bfb.generate_bot_ships(Bmap2)

        #Events and update
        pygame.display.update()
        CLICK = False

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                RUNNING = False
                pygame.quit()
            if event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 1:
                    CLICK = True


from scarv import *

import pybedtools
import pyranges as pr
from keras.models import load_model
import keras_genomics
import joblib
import numpy as np
import pandas as pd
import sys

chrom = sys.argv[1]
chrom_base = "chrX" if "X" in chrom else chrom

window_size = int(sys.argv[2])
ancestry = sys.argv[3]
pop_size = int(sys.argv[4])
outFile = sys.argv[5]

k = 25
genome = pybedtools.genome_registry.hg38
reference_fasta = "/rds/project/who1000-1/rds-who1000-cbrc/ref/UCSC/hg38/hg38.fa"

snvs = scarv_queries.load_variants("/rds/project/who1000-1/rds-who1000-cbrc/user/jwt44/scarv_pipeline_gnomad_hg38/" + ancestry +"/variants/snvs_by_chr/pass_snvs_" + chrom_base + ".bed")

deletions = scarv_queries.load_variants("/rds/project/who1000-1/rds-who1000-cbrc/user/jwt44/scarv_pipeline_gnomad_hg38/" + ancestry +"/variants/pass_deletions.bed")[chrom_base]
deletions.ref = deletions.ref.str.slice(1)
deletions.alt = "X"
deletions = deletions[deletions.ref.str.len() == 1]
deletions.Start += 1
deletions = deletions[deletions.lengths() == 1]

snvs_and_deletions = pr.PyRanges(pd.concat([snvs.as_df(), deletions.as_df()]))

insertions = scarv_queries.load_variants("/rds/project/who1000-1/rds-who1000-cbrc/user/jwt44/scarv_pipeline_gnomad_hg38/" + ancestry +"/variants/pass_insertions.bed")[chrom_base]
insertions.ref = "X"
insertions.alt = insertions.alt.str.slice(1)
insertions = insertions[insertions.alt.str.len() == 1]

reliable_sites = pr.read_bed("/rds/project/who1000-1/rds-who1000-cbrc/user/jwt44/scarv_pipeline_gnomad_hg38/" + ancestry + "/quality_filtering/reliable_sites_by_chr/reliable_sites_" + chrom + ".bed")

preds_by_kmer_file = "/rds/project/who1000-1/rds-who1000-cbrc/user/jwt44/scarv_pipeline_gnomad_hg38/" + ancestry + "/preds_by_kmer.pkl"
preds_by_kmer = pd.read_pickle(preds_by_kmer_file)

chrom_gr = pr.from_dict({'Chromosome': [chrom_base], 'Start': [0], 'End': [genome[chrom_base][1]]})
chrom_seq = pr.get_fasta(chrom_gr, reference_fasta)[0].upper()
chrom_seq_padded = "X".join(chrom_seq)
chrom_seq_ohe = np.array(pd.get_dummies(pd.Series(list(chrom_seq_padded)))[['A','C','X','G','T']])

# Compute chromosome-wide predictions
k_mers = pd.Series([chrom_seq_padded[(pos-k//2):(pos+k//2+1)] for pos in range(len(chrom_seq_padded)-k//2)])

preds = np.empty(shape=(2*genome[chrom_base][1]-1, 5))
preds[:] = np.nan

reliable_indices = np.sort(np.concatenate([2*reliable_sites.tile(1).Start.to_numpy(), 2*reliable_sites.tile(1).Start.to_numpy()+1]))
identified_nucleotides = np.where(["N" not in seq for seq in k_mers])[0]
query_indices = np.intersect1d(reliable_indices, identified_nucleotides)

preds[query_indices] = preds_by_kmer.loc[k_mers[query_indices]].to_numpy()

# Compute chromosome-wide allele frequencies
allele_counts = 2 * pop_size * chrom_seq_ohe.astype('int32')

allele_counts = scarv_assess.process_variants (insertions, allele_counts, insertion=True)
allele_counts = scarv_assess.process_variants (snvs_and_deletions, allele_counts, insertion=False)

allele_counts[np.setdiff1d(np.arange(allele_counts.shape[0]), query_indices)] = 0   # remove contribution from unreliable or unidentified sequences
allele_frequencies = allele_counts / allele_counts.sum(axis=1)[:, np.newaxis]

# Generate score
observed_entropy = scarv_assess.get_entropy(allele_frequencies)
observed_entropy_sliding = np.convolve(observed_entropy, [1] * 2 * window_size, 'same')[::2]

predicted_entropy = scarv_assess.get_entropy(preds)
predicted_entropy_sliding = np.convolve(predicted_entropy, [1] * 2 * window_size, 'same')[::2]

coverage_indicator = np.zeros(2*genome[chrom_base][1]-1, dtype=np.uint8)
coverage_indicator[query_indices] = 1
coverage_sliding = np.convolve(coverage_indicator, [1] * 2 * window_size, 'same')[::2]

# Print to file
output = pd.DataFrame({'chrom': chrom_base, 'start': np.arange(genome[chrom_base][1]), 'end': np.arange(genome[chrom_base][1]) + 1,
    'obs': observed_entropy_sliding, 'exp': predicted_entropy_sliding, 'coverage': coverage_sliding})

output.to_csv(outFile, sep='\t', header=False, index=False, compression='gzip')


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Mar 13 14:30:12 2019

@author:Soumya
"""
import os.path
import random
import torchvision.transforms as transforms
import torch
from data.base_dataset import BaseDataset
#from data.image_folder import make_dataset
from PIL import Image
import pandas
#from natsort import natsorted, ns
#import glob
#import numpy as np
#import sys
#sys.path.append('/home/esa/Downloads/face-alignment')

class AlignedDataset(BaseDataset):
    def initialize(self, opt):
        self.opt = opt
        self.root = opt.dataroot
        self.dir_AB = os.path.join(opt.dataroot)
        self.ref_path=opt.which_ref
        self.csv_path=opt.csv_path

    def __getitem__(self, index):

        SI_path= self.ref_path

        df2 = pandas.read_csv(self.csv_path)
        a=list(range(296,299))+ list(range(679,696))#+ list(range(299,435))
        refdf = df2[df2.columns[a]]

        f=0
        A1 = Image.open(SI_path).convert('RGB')
        A1 = A1.resize((self.opt.loadSize, self.opt.loadSize), Image.BICUBIC)
        A1 = transforms.ToTensor()(A1)
        A = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(A1)

        for i in range(0,len(refdf)):

            paramA=torch.tensor(refdf.values[i]).view(1,-1)
            paramA[0,0:3]=(paramA[0,0:3]-(-0.70))/1.4
            paramA[0,3:20]=paramA[0,3:20]/5

            if f==0:
                P1=paramA
            else:
                P1=torch.cat([P1,paramA],dim=0)
            f=f+1

        if (not self.opt.no_flip) and random.random() < 0.5:
            idx = [i for i in range(A.size(2) - 1, -1, -1)]
            idx = torch.LongTensor(idx)
            A = A.index_select(2, idx)

        return {'A': A, 'PB': P1,
                'A_paths': SI_path}#SAB_path[0]

    def __len__(self):
        return len(self.ref_path)

    def name(self):
        return 'AlignedDataset'


from music21 import converter,instrument,note,chord,stream
import tensorflow as tf
from keras.layers import Dense, Conv2D, Dropout, Flatten, MaxPooling2D,LSTM,Activation
from keras.models import Sequential
import keras.models
import numpy as np
from keras.utils import np_utils
from keras.callbacks import ModelCheckpoint
import glob
import pickle

def Predict(model):

    with open('objs.pkl','rb') as f:
        networkInput, pitchNames, numNotes, nVocab, normalNetworkInput = pickle.load(f)

    #randomly select a start point in our network inputs
    start = np.random.randint(0, len(networkInput)-1)
    #create a dictonary that allows us to convert back from the ltsm output to note names.
    intToNote = dict((number, note) for number, note in enumerate(pitchNames))
    #the prediction pattern at the randomly generated start point.
    pattern = networkInput[start]
    #array that will store the output predictions
    predictionOutput = []
    # generate numNotes notes
    for note_index in range(numNotes):
        #converts the prediction input to a format that can be used with the lstm predict function
        predictionInput = np.reshape(pattern, (1, len(pattern), 1))
        predictionInput = predictionInput / float(nVocab)
        #using this converted input we predict the next note in the sequence.
        prediction = model.predict(predictionInput, verbose=0)
        #this prediciton variable holds the probabilitys of each note to be played next
        #we use the argmax value to get the one with the bigest probability
        index = np.argmax(prediction)
        #we convert the most likely result to a note
        result = intToNote[index]
        #finanly we add this note to the predicition output list
        predictionOutput.append(result)

        #we add this note to the end of the list
        pattern.append(index)
        #we remove the first note in the list
        pattern = pattern[1:len(pattern)]
        #this allows us to now make our next prediciton bassed on the last 100 notes (witch now includes our predicted note at the end and our the first note being removed)

    #parseing this output to a midi file.
    #offset represents how far from the start of the midi file the note should be played.
    offset = 0
    offSetRepeat = 0
    #output notes will be the final string of notes with offset values to be converted to midi
    outputNotes=[]
    #sets note and chord objects based on models predicitions
    for pattern in predictionOutput:
        #if statement checking weather the next value in prediciton output is a chord.
        #if so we seperate the chord into an array of notes
        #create a new note object for these notes
        #assign default iinstrument as piano
        #then we create a chord object from these notes and assign the current offset.
        #finaly we appened the the outputnotes list.
        if('.' in pattern) or pattern.isdigit():
            notesInChord = pattern.split('.')
            cNotes = []
            for currentNote in notesInChord:
                newNote = note.Note(int(currentNote))
                newNote.storedInstrument = instrument.Piano()
                cNotes.append(newNote)
            newChord = chord.Chord(cNotes)
            newChord.offset=offset
            outputNotes.append(newChord)
        #if note a chord its simply a single note in this case we create a new note object assign the current offset assign default instrument and append to the output notes
        else:
                newNote = note.Note(pattern)
                newNote.offset = offset
                newNote.storedInstrument = instrument.Piano()
                outputNotes.append(newNote)
        #we increase the offset after each generated note has been added
        #if statement that generates 1 bar or four beat rythms

        #adds offset amount to total ofset of generated peice and reduces the offset repeat amount by 1
        offset += 0.5

        #create a stream from our output notes
        #write the output stream to the file OUTPUT.mid
    MIDIStream = stream.Stream(outputNotes)
    return MIDIStream


#	PHOTOMETRY CODE FOR PYTHON 3.X
#	2019.06.20	CREATED BY	Gregory S.H. Paek
#	2019.10.06	MODIFIED BY	Gregory S.H. Paek
#============================================================
import os, glob
import numpy as np
import matplotlib.pyplot as plt
from astropy.table import Table, vstack
from astropy.io import ascii
from astropy.io import fits
from astropy.time import Time
from astropy.coordinates import SkyCoord
from astropy import units as u
from astropy.wcs import WCS
from imsng import phot
import time
#============================================================
#	FUNCTION
#============================================================
def phot_routine(inim, refcatname, phottype, tra, tdec, path_base='./', aperture='MAG_APER_7', detsig=3.0, frac=0.95):
	#------------------------------------------------------------
	#	HEADER INFO
	hdr			= fits.getheader(inim)
	xcent, ycent= hdr['NAXIS1']/2., hdr['NAXIS2']/2.
	try:
		w			= WCS(inim)
		radeg, dedeg= w.all_pix2world(xcent, ycent, 1)
		radeg, dedeg= np.asscalar(radeg), np.asscalar(dedeg)
	except:
		print('BAD WCS INFORMATION?')
		radeg,dedeg	= hdr['CRVAL1'], hdr['CRVAL2']
	#xcent, ycent= w.all_world2pix(radeg, dedeg, 1)
	#------------------------------------------------------------
	try:
		date_obs	= hdr['date-obs']
		jd			= round(Time(date_obs, format='isot', scale='utc').jd, 3)
	except:
		date_obs	= None
		jd			= None
	#------------------------------------------------------------
	#	NAME INFO
	obs, obj		= 'KMTNET', inim[:-7]
	refmagkey		= 'R'
	refmagerkey 	= refmagkey+'err'
	indx_obs		= np.where(obstbl['obs']==obs)
	gain, pixscale	= obstbl[indx_obs]['gain'][0], obstbl[indx_obs]['pixelscale'][0]
	#------------------------------------------------------------
	#	REF. CATALOG QUERY
	#------------------------------------------------------------
	refcatlist	= glob.glob(path_refcat+'/*.cat')
	#------------------------------------------------------------
	if		refcatname	== 'PS1':
		if path_refcat+'/ps1-'+obj+'.cat' not in refcatlist:
			querytbl	= phot.ps1_query(obj, radeg, dedeg, path_refcat, radius=3.0)
		else:
			querytbl	= ascii.read(path_refcat+'/ps1-'+obj+'.cat')
		reftbl, refcat  = phot.ps1_Tonry(querytbl, obj)
	#------------------------------------------------------------
	elif 	refcatname	== 'SDSS':
		if path_refcat+'/sdss-'+obj+'.cat' not in refcatlist:
			querytbl        = phot.sdss_query(obj, radeg, dedeg, path_refcat)
		else:
			querytbl        = ascii.read(path_refcat+'/sdss-'+obj+'.cat')
		reftbl, refcat  = phot.sdss_Blaton(querytbl, obj)
	#------------------------------------------------------------
	elif	refcatname	== 'APASS':
		if path_refcat+'/apass-'+obj+'.cat' not in refcatlist:
			querytbl        = phot.apass_query(obj, radeg, dedeg, path_refcat, radius=2.0)
		else:
			querytbl        = ascii.read(path_refcat+'/apass-'+obj+'.cat')
		reftbl, refcat  = phot.apass_Blaton(querytbl, obj)
	#------------------------------------------------------------
	elif	refcatname	== '2MASS':
		if path_refcat+'/2mass-'+obj+'.cat' not in refcatlist:
			querytbl        = phot.twomass_query(obj, radeg, dedeg, path_refcat, band=refmagkey, radius=1.0)
		else:
			querytbl        = ascii.read(path_refcat+'/2mass-'+obj+'.cat')
		reftbl, refcat  = querytbl, '2mass-'+obj+'.cat'
	#------------------------------------------------------------
	#	SourceEXtractor
	#------------------------------------------------------------
	peeing, seeing	= phot.psfex(inim, pixscale)
	param_secom	= dict(	inim=inim,
						gain=gain, pixscale=pixscale, seeing=seeing,
						det_sigma=detsig,
						backsize=str(64), backfiltersize=str(3),
						psf=True, check=False)
	# intbl0, incat	= phot.secom(**param_secom)
	intbl, incat	= phot.secom(**param_secom)
	'''
	#	CENTER POS. & DIST CUT
	deldist		= phot.sqsum((xcent-intbl0['X_IMAGE']), (ycent-intbl0['Y_IMAGE']))
	indx_dist	= np.where(deldist < np.sqrt(frac)*(xcent+ycent)/2.)
	intbl		= intbl0[indx_dist]
	intbl.write(incat, format='ascii', overwrite=True)
	'''
	#	MATCHING
	param_match	= dict(	intbl=intbl, reftbl=reftbl,
						inra=intbl['ALPHA_J2000'], indec=intbl['DELTA_J2000'],
						refra=reftbl['ra'], refdec=reftbl['dec'], sep=2.0)
	mtbl		= phot.matching(**param_match)
	#------------------------------------------------------------
	#	ZEROPOINT CALCULATION
	#------------------------------------------------------------
	inmagkey	= aperture
	inmagerkey	= '_'.join(['MAGERR', inmagkey.split('_')[1], inmagkey.split('_')[2]])
	param_st4zp	= dict(	intbl=mtbl,
						inmagerkey=aperture,
						refmagkey=refmagkey,
						refmagerkey=refmagerkey,
						refmaglower=10,
						refmagupper=16,
						refmagerupper=0.1,
						inmagerupper=0.1)
	param_zpcal	= dict(	intbl=phot.star4zp(**param_st4zp),
						inmagkey=inmagkey, inmagerkey=inmagerkey,
						refmagkey=refmagkey, refmagerkey=refmagerkey,
						sigma=2.0)
	zp, zper, otbl, xtbl	= phot.zpcal(**param_zpcal)

	#------------------------------------------------------------
	#	ZEROPOINT PLOT
	#------------------------------------------------------------
	outname	= path_base+'/{0}.{1}.zpcal.png'.format(inim[:-5], inmagkey)
	phot.zpplot(	outname=outname,
					otbl=otbl,xtbl=xtbl,
					inmagkey=inmagkey, inmagerkey=inmagerkey,
					refmagkey=refmagkey, refmagerkey=refmagerkey,
					zp=zp, zper=zper)
	param_plot	= dict(	inim		= inim,
						numb_list	= otbl['NUMBER'],
						xim_list	= otbl['X_IMAGE'],
						yim_list	= otbl['Y_IMAGE'],
						add			= True,
						numb_addlist= xtbl['NUMBER'],
						xim_addlist	= xtbl['X_IMAGE'],
						yim_addlist	= xtbl['Y_IMAGE'])
	try:
		phot.plotshow(**param_plot)
	except:
		print('FAIL TO DRAW ZEROPOINT GRAPH')
		pass
	#------------------------------------------------------------
	#	TARGET PHOTOMETRY
	#------------------------------------------------------------
	skymean, skymed, skysig		= phot.bkgest_mask(inim)
	aper	= 2*peeing
	ul		= phot.limitmag(detsig, zp, aper, skysig)
	#------------------------------------------------------------
	#	ADD HEADER INFO
	#------------------------------------------------------------
	phot.puthdr(inim, 'SEEING',	round(seeing, 3),		hdrcomment='SEEING [arcsec]')
	phot.puthdr(inim, 'PEEING',	round(peeing, 3),		hdrcomment='SEEING [pixel]')
	phot.puthdr(inim, 'SKYSIG',	round(skysig, 3),		hdrcomment='SKY SIGMA VALUE')
	phot.puthdr(inim, 'SKYVAL',	round(skymed, 3),		hdrcomment='SKY MEDIAN VALUE')
	phot.puthdr(inim, 'OPTZP',	round(zp, 3),			hdrcomment='2*SEEING DIAMETER')
	phot.puthdr(inim, 'OPTZPERR',round(zper, 3),		hdrcomment='2*SEEING DIAMETER')
	phot.puthdr(inim, 'OPTUL',	round(ul, 3),			hdrcomment='2*SEEING 3 sigma limit mag')
	phot.puthdr(inim, 'STDNUMB',len(otbl),				hdrcomment='# OF STD STARS')
	#------------------------------------------------------------
	#	NORMAL PHOTOMETRY
	#------------------------------------------------------------
	if phottype == 'normal':
		intbl['REAL_'+inmagkey]		= zp + intbl[inmagkey]
		intbl['REAL_'+inmagerkey]	= phot.sqsum(zper, intbl[inmagerkey])
		indx_targ	= phot.targetfind(tra, tdec, intbl['ALPHA_J2000'], intbl['DELTA_J2000'], sep=seeing)
		if indx_targ != None:
			mag, mager	= intbl[indx_targ]['REAL_'+inmagkey], intbl[indx_targ]['REAL_'+inmagerkey]
		else:
			mag, mager	= -99, -99
	#------------------------------------------------------------
	#	SUBTRACTION PHOTOMETRY
	#------------------------------------------------------------
	elif phottype == 'subt':
		subim	= 'hd'+inim
		phot.puthdr(subim, 'SEEING',	round(seeing, 3),		hdrcomment='SEEING [arcsec]')
		phot.puthdr(subim, 'PEEING',	round(peeing, 3),		hdrcomment='SEEING [pixel]')
		phot.puthdr(subim, 'SKYSIG',	round(skysig, 3),		hdrcomment='SKY SIGMA VALUE')
		phot.puthdr(subim, 'SKYVAL',	round(skymed, 3),		hdrcomment='SKY MEDIAN VALUE')
		phot.puthdr(subim, 'OPTZP',	round(zp, 3),				hdrcomment='2*SEEING DIAMETER')
		phot.puthdr(subim, 'OPTZPERR',round(zper, 3),			hdrcomment='2*SEEING DIAMETER')
		phot.puthdr(subim, 'OPTUL',	round(ul, 3),				hdrcomment='2*SEEING 3 sigma limit mag')
		phot.puthdr(subim, 'STDNUMB',len(otbl),					hdrcomment='# OF STD STARS')
		os.system('cp {} {}'.format(inim[:-5]+'.psf', subim[:-5]+'.psf'))
		param_subcom	= dict(	inim=subim,
								gain=gain, pixscale=pixscale, seeing=seeing,
								det_sigma=3,
								backsize=str(64), backfiltersize=str(3),
								psf=True, check=False)
		subtbl, subcat	= phot.secom(**param_subcom)
		subtbl['REAL_'+inmagkey]	= zp + subtbl[inmagkey]
		subtbl['REAL_'+inmagerkey]	= phot.sqsum(zper, subtbl[inmagerkey])
		indx_targ	= phot.targetfind(tra, tdec, subtbl['ALPHA_J2000'], subtbl['DELTA_J2000'], sep=seeing)
		if indx_targ != None:
			mag, mager	= subtbl[indx_targ]['REAL_'+inmagkey], subtbl[indx_targ]['REAL_'+inmagerkey]
		else:
			mag, mager	= -99, -99
	#------------------------------------------------------------
	#	CALC. DEPTH
	#------------------------------------------------------------
	elif phottype == 'depth':
		mag, mager	= -99, -99

	onetbl	= Table([[inim], [obs], [obj], [round(radeg, 3)], [round(dedeg, 3)], [date_obs], [jd], [refmagkey], [len(otbl)], [round(zp, 3)], [round(zper, 3)], [round(seeing, 3)], [round(skymed, 3)], [round(skysig, 3)], [round(ul, 3)], [mag], [mager]],
					names=('image', 'obs', 'obj', 'ra', 'dec', 'date-obs', 'jd', 'filter', 'stdnumb', 'zp', 'zper', 'seeing', 'skyval', 'skysig', 'ul', 'mag', 'magerr'))
	return onetbl
#============================================================
#	USER SETTING
#============================================================
path_base	= './'
path_obs	= '/home/sonic/Research/table'
path_refcat	= '/home/sonic/Research/cat/refcat'
#------------------------------------------------------------
obstbl		= ascii.read(path_obs+'/obs.txt')
#	TARGET COORD.	[deg]
#------------------------------------------------------------
#	IMAGES TO PHOTOMETRY
#	INPUT FORMAT	: Calib-[OBS]-[TARGET]-[DATE]-[TIME]-[BAND]*.fits
#------------------------------------------------------------
kmtnetkey	= 'nn'
imlist		= glob.glob('a*{}-4*.fits'.format(kmtnetkey))
# imlist.sort()
for inim in glob.glob('a*{}-*.fits'.format(kmtnetkey)):
	if '-4' not in inim:
		imlist.append(inim)

for img in imlist: print(img)

photlist	= []
refcatname	= 'APASS'
phottype	= 'depth'
starttime	= time.time()
#============================================================
#	MAIN COMMAND
#============================================================
for inim in imlist:
	try:
		plt.ioff()
		param_phot	= dict(	inim=inim, refcatname=refcatname, phottype=phottype,
							tra=0, tdec=0, path_base='./', aperture='MAG_APER_7',
							detsig=3.0, frac=0.9)
		photlist.append(phot_routine(**param_phot))
		os.system('rm psf*.fits snap*.fits *.xml seg.fits')
	except:
		pass
#------------------------------------------------------------
#	FINISH
#------------------------------------------------------------
if len(photlist) == 0:
	print('PHOTOMETRY FAILED!')
else:
	photbl		= vstack(photlist)
	if 'phot-{}.dat'.format(kmtnetkey) in glob.glob(path_base+'/phot-{}.dat'.format(kmtnetkey)):
		os.system('mv {} {}'.format(path_base+'/phot-{}.dat'.format(kmtnetkey), path_base+'/phot-{}.dat.bkg'.format(kmtnetkey)))
	photbl.write(path_base+'/phot-{}.dat'.format(kmtnetkey), format='ascii', overwrite=True)
	deltime		= time.time() - starttime
	print('All PROCESS IS DONE.\t('+str(round(deltime, 1))+' sec)')


import cv2
import dlib
import time
import glob
import numpy as np

imgs, heights, widths = [], [], []
# read imgs
for f in glob.glob("pics/afterWork/*.jpg"):
    if 'YSL' in f:
        img = cv2.imread(f, -1)  # 参数-1表示返回原图
        h, w = img.shape[:2]  # 切片
        heights.append(h)
        widths.append(w)
        imgs.append(img)
# set minimum heights and weights
min_height = min(heights)
min_width = min(widths)
# resize imgs
for i, x in enumerate(imgs):
    # i为每个图像的序号, x为每个图像的多维像素矩阵
    imgs[i] = x[:min_height:2, :min_width:2]  # 切片 以步长为3
# concatenation
img0 = np.concatenate(imgs[:2], 1)  # 横着拼三个
# resize logo
imgLogo = cv2.imread('/Users/ryan/PycharmProjects/detect_mouth/venv/pics/logo/ysl.jpg',cv2.IMREAD_UNCHANGED)
h, w = img0.shape[:2]
dim = (w, int(h/3))
imgLogo = cv2.resize(imgLogo, dim, interpolation = cv2.INTER_AREA)
cv2.imshow('logo',imgLogo)
img9 = np.concatenate([img0,imgLogo], 0)# 竖着拼起来
#resize final pic
print(min_width,min_height)
#img9 = cv2.resize(img9,(1080,1439), interpolation = cv2.INTER_AREA)
h,w = img9.shape[:2]
if int(h*0.75) != w:
    print('not 4:3')
    x = int(w*1.25)-h
    img9 = cv2.copyMakeBorder(img9,int(x/2), int(x/2), 0, 0, cv2.BORDER_CONSTANT, value=(0,0,0))
cv2.imwrite("pics/final/result_1x1.jpg", img9)


import os
import sys
import glob
import copy
import math
import random
import csv
import numpy as np

# from torch.utils.data import Dataset
import torch
import cv2
from tqdm import tqdm

# Semantic Segmentation classes
classes = {
    0: [0, 0, 0],         # None
    1: [70, 70, 70],      # Buildings
    2: [190, 153, 153],   # Fences
    3: [72, 0, 90],       # Other
    4: [220, 20, 60],     # Pedestrians
    5: [153, 153, 153],   # Poles
    6: [157, 234, 50],    # RoadLines
    7: [128, 64, 128],    # Roads
    8: [244, 35, 232],    # Sidewalks
    9: [107, 142, 35],    # Vegetation
    10: [0, 0, 255],      # Vehicles
    11: [102, 102, 156],  # Walls
    12: [220, 220, 0]     # TrafficSigns
}

class CarlaDataset(torch.utils.data.Dataset):
    def __init__(self, data_dir=None, train=True, model_name=None, split_size=0.88):
        self.data_dir = data_dir
        self.train = train
        self.img_size = (88, 200)
        self.split_size = split_size
        self.model_name = model_name

        if self.train:
            self.dir = os.path.join(self.data_dir, 'train')
        else:
            self.dir = os.path.join(self.data_dir, 'test')

        self.img_data, self.target_data = self._process(self.dir)   # type : list

    def __len__(self):
        return len(self.img_data)

    def __getitem__(self, index):
        # img --> RGB image, measurements --> [steer, throttle, speed]
        try:
            measurement = {
                'img': self.img_data[index],
                'speed': np.array([self.target_data[index][2]]),
                'direction': np.array([self.target_data[index][3]]),
                'target_action': self.target_data[index][:2]
            }

        except IndexError:
            print("Blank Image")
            measurement = {
                'img': np.zeros((17, self.img_size[0], self.img_size[1])),
                'speed': np.array([0.0]),
                'direction': np.array([-1.0]),
                'target_action': np.array([[0.0, 0.0]])
            }
        return measurement

    @staticmethod
    def get_measurements(path):
        measurements = []
        with open(glob.glob(os.path.join(path, 'measurements*'))[0], 'r') as ofd:
            w = csv.DictReader(ofd)

            for row in w:
                steer = float(row['steer'])
                if float(row['throttle']) >= 0.0:
                    throttle = float(row['throttle'])
                else:
                    throttle =  float(row['brake'])

                command = float(row['direction'])

                if 'speed' in row.keys():
                    speed = float(row['speed'])
                else:
                    speed = float((row[None])[0])

                measurements.append(np.array([steer, throttle, speed, command]))

        return measurements

    def _process(self, dir_path):
        episode_lst = glob.glob(os.path.join(dir_path, 'episode_*'))
        if len(episode_lst) == 0:
            return None, None
        random.shuffle(episode_lst)

        imgs, measurements = [], []
        append = imgs.append
        for episode in tqdm(episode_lst):
            c_rgb_lst = glob.glob(os.path.join(episode, 'rgb_*'))
            c_semseg_lst = glob.glob(os.path.join(episode,'semseg_*'))
            c_depth_lst = glob.glob(os.path.join(episode, 'depth_*'))
            measurements = self.get_measurements(episode)

            for rgb, semseg, depth in zip(c_rgb_lst, c_semseg_lst, c_depth_lst):

                rgb = cv2.imread(rgb)
                rgb = rgb[:, :, ::-1]
                rgb = rgb.astype(np.float32)
                rgb /= 255.0

                semseg = cv2.imread(semseg)
                semseg = semseg[:, :, ::-1]
                semseg = self.cityscape_to_index(semseg)
                semseg = semseg.astype(np.float32)
                semseg /= 255.0

                depth = cv2.imread(depth)
                depth = depth[:, :, ::-1]
                depth = depth[:, :, 0]

                out = np.dstack((rgb, semseg, depth))
                out = np.transpose((2, 0, 1))
                append(out)
        return imgs, measurements

    @staticmethod
    def cityscape_to_index(image):
        result = np.zeros((image.shape[0], image.shape[1], 13))
        where = np.where
        items = classes.items()
        for key, value in items:
            x, y, z = where(image == value)
            result[x, y, key] = 1
        return result


# Copyright 2021 Zilliz. All rights reserved.
#
# 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.

from typing import NamedTuple
import torch
import numpy

from towhee.operator import Operator

class PytorchTransformerVecOperator(Operator):
    """
    Pytorch transformer vector operator, the output is embedding instead of prediction
        Args:
            model(nn.Module):
                Model object
    """

    def __init__(self, model) -> None:
        super().__init__()
        self._model = model

    def __call__(self, img_tensor: torch.Tensor) -> NamedTuple('Outputs', [('embedding', numpy.ndarray)]):
        self._model.eval()
        with torch.no_grad():
            outputs = self._model(img_tensor).squeeze(0)
        Outputs = NamedTuple('Outputs', [('embedding', numpy.ndarray)])
        return Outputs(outputs)