from datetime import datetime
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
plt.switch_backend('agg')
import io
from torchvision import transforms as trans
from data.data_pipe import de_preprocess
import torch
from model import l2_norm
import pdb
import cv2

from face_detection.accuracy_evaluation import predict
from face_detection.config_farm import configuration_10_320_20L_5scales_v2 as cfg
import mxnet as mx
import numpy as np
def separate_bn_paras(modules):
    if not isinstance(modules, list):
        modules = [*modules.modules()]
    paras_only_bn = []
    paras_wo_bn = []
    for layer in modules:
        if 'model' in str(layer.__class__):
            continue
        if 'container' in str(layer.__class__):
            continue
        else:
            if 'batchnorm' in str(layer.__class__):
                paras_only_bn.extend([*layer.parameters()])
            else:
                paras_wo_bn.extend([*layer.parameters()])
    return paras_only_bn, paras_wo_bn

def prepare_facebank(conf, model, mtcnn, tta = True):
    model.eval()
    ctx = mx.gpu(0)
    symbol_file_path = 'face_detection/symbol_farm/symbol_10_320_20L_5scales_v2_deploy.json'
    model_file_path = 'face_detection/saved_model/configuration_10_320_20L_5scales_v2/train_10_320_20L_5scales_v2_iter_1800000.params'
    face_detector = predict.Predict(mxnet=mx,
                             symbol_file_path=symbol_file_path,
                             model_file_path=model_file_path,
                             ctx=ctx,
                             receptive_field_list=cfg.param_receptive_field_list,
                             receptive_field_stride=cfg.param_receptive_field_stride,
                             bbox_small_list=cfg.param_bbox_small_list,
                             bbox_large_list=cfg.param_bbox_large_list,
                             receptive_field_center_start=cfg.param_receptive_field_center_start,
                             num_output_scales=cfg.param_num_output_scales)
    embeddings =  []
    names = ['Unknown']
    for path in conf.facebank_path.iterdir():
        if path.is_file():
            continue
        else:
            embs = []
            for filename in path.iterdir():
                if not filename.is_file():
                    continue
                else:
                    try:
                        print(filename)
                        image = Image.open(filename)
                        img = image
                        # img = np.array(image)
                        # img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

                        # faces, infer_time = face_detector.predict(img, resize_scale=0.5, score_threshold=0.4, top_k=10000, \
                        #                     NMS_threshold=0.2, NMS_flag=True, skip_scale_branch_list=[])
                        # img_size = 112
                        # print(len(faces))
                        # margin = 0

                        # img_h, img_w, _ = np.shape(image)

                        # for i, bbox in enumerate(faces):
                        #     x1, y1, x2, y2= bbox[0], bbox[1], bbox[2] ,bbox[3]
                        #     xw1 = max(int(x1 - margin ), 0)
                        #     yw1 = max(int(y1 - margin ), 0)
                        #     xw2 = min(int(x2 + margin ), img_w - 1)
                        #     yw2 = min(int(y2 + margin ), img_h - 1)
                        #     face =  cv2.resize(img[yw1:yw2 + 1, xw1:xw2 + 1], (img_size, img_size))
                        #     # img = Image.fromarray(face[...,::-1])
                        #     img = face
                        #     break
                    except Exception  as e:
                        print(e)
                        continue

                    if img.size != (112, 112):
                        img = mtcnn.align(img)
                    print(type(img))
                    # cv2.imshow('window', img)
                    # img.show()
                    # if cv2.waitKey() == ord('q'):
                        # break
                    with torch.no_grad():
                        if tta:
                            img = trans.functional.to_grayscale(img, num_output_channels=3)
                            mirror = trans.functional.hflip(img)
                            emb = model(conf.test_transform(img).to(conf.device).unsqueeze(0))
                            emb_mirror = model(conf.test_transform(mirror).to(conf.device).unsqueeze(0))

                            v_mirror = trans.functional.vflip(mirror)
                            v_emb_mirror = model(conf.test_transform(v_mirror).to(conf.device).unsqueeze(0))

                            v_img = trans.functional.vflip(img)
                            v_img_mirror = model(conf.test_transform(v_img).to(conf.device).unsqueeze(0))

                            embs.append(l2_norm(emb + emb_mirror))
                            # embs.append(l2_norm(emb + emb_mirror + v_emb_mirror + v_img_mirror))
                            # embs.append(emb)
                            # embs.append(emb_mirror)
                            # embs.append(v_emb_mirror)
                            # embs.append(v_img_mirror)
                        else:
                            embs.append(model(conf.test_transform(img).to(conf.device).unsqueeze(0)))
        if len(embs) == 0:
            continue
        embedding = torch.cat(embs).mean(0,keepdim=True)
        embeddings.append(embedding)
        names.append(path.name)
    embeddings = torch.cat(embeddings)
    names = np.array(names)
    torch.save(embeddings, conf.facebank_path/'facebank.pth')
    np.save(conf.facebank_path/'names', names)
    return embeddings, names

def load_facebank(conf):
    embeddings = torch.load(conf.facebank_path/'facebank.pth')
    names = np.load(conf.facebank_path/'names.npy')
    return embeddings, names

def face_reader(conf, conn, flag, boxes_arr, result_arr, learner, mtcnn, targets, tta):
    while True:
        try:
            image = conn.recv()
        except:
            continue
        try:
            bboxes, faces = mtcnn.align_multi(image, limit=conf.face_limit)
        except:
            bboxes = []

        results = learner.infer(conf, faces, targets, tta)

        if len(bboxes) > 0:
            print('bboxes in reader : {}'.format(bboxes))
            bboxes = bboxes[:,:-1] #shape:[10,4],only keep 10 highest possibiity faces
            bboxes = bboxes.astype(int)
            bboxes = bboxes + [-1,-1,1,1] # personal choice
            assert bboxes.shape[0] == results.shape[0],'bbox and faces number not same'
            bboxes = bboxes.reshape([-1])
            for i in range(len(boxes_arr)):
                if i < len(bboxes):
                    boxes_arr[i] = bboxes[i]
                else:
                    boxes_arr[i] = 0
            for i in range(len(result_arr)):
                if i < len(results):
                    result_arr[i] = results[i]
                else:
                    result_arr[i] = -1
        else:
            for i in range(len(boxes_arr)):
                boxes_arr[i] = 0 # by default,it's all 0
            for i in range(len(result_arr)):
                result_arr[i] = -1 # by default,it's all -1
        print('boxes_arr ： {}'.format(boxes_arr[:4]))
        print('result_arr ： {}'.format(result_arr[:4]))
        flag.value = 0

hflip = trans.Compose([
            de_preprocess,
            trans.ToPILImage(),
            trans.functional.hflip,
            trans.ToTensor(),
            trans.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])

def hflip_batch(imgs_tensor):
    hfliped_imgs = torch.empty_like(imgs_tensor)
    for i, img_ten in enumerate(imgs_tensor):
        hfliped_imgs[i] = hflip(img_ten)
    return hfliped_imgs

def get_time():
    return (str(datetime.now())[:-10]).replace(' ','-').replace(':','-')

def gen_plot(fpr, tpr):
    """Create a pyplot plot and save to buffer."""
    plt.figure()
    plt.xlabel("FPR", fontsize=14)
    plt.ylabel("TPR", fontsize=14)
    plt.title("ROC Curve", fontsize=14)
    plot = plt.plot(fpr, tpr, linewidth=2)
    buf = io.BytesIO()
    plt.savefig(buf, format='jpeg')
    buf.seek(0)
    plt.close()
    return buf

def draw_box_name(bbox,name,frame):
    frame = cv2.rectangle(frame,(bbox[0],bbox[1]),(bbox[2],bbox[3]),(0,0,255),6)
    frame = cv2.putText(frame,
                    name,
                    (bbox[0],bbox[1]),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    2,
                    (0,255,0),
                    3,
                    cv2.LINE_AA)
    return frame


import os.path as op
from textwrap import dedent
import glob
import os

import nibabel
from nibabel import Nifti1Image, MGHImage
import numpy as np
from fsl.data.image import Image as FSLImage

from regtricks.image_space import ImageSpace
from regtricks import x5_interface as x5
from regtricks import application_helpers as apply
from regtricks import multiplication as multiply
from regtricks.transforms.transform import Transform

class Registration(Transform):
    """
    Affine (4x4) transformation between two images.

    Args:
        src2ref (str/np.ndarray): path to text-like file to load or np.ndarray
    """

    def __init__(self, src2ref):
        Transform.__init__(self)

        if isinstance(src2ref, str):
            src2ref = np.loadtxt(src2ref)

        if (src2ref.shape != (4,4)
            or (np.abs(src2ref[3,:] - [0,0,0,1]) > 1e-9).any()):
            raise RuntimeError("src2ref must be a 4x4 affine matrix, where ",
                               "the last row is [0,0,0,1].")

        self.__src2ref = src2ref

    @classmethod
    def from_flirt(cls, src2ref, src, ref):
        """
        Load an affine (4x4) transformation between two images directly from
        FLIRT's -omat output.

        Args:
            src2ref (str/np.ndarray): path to text-like file to load or np.ndarray
            src: the source of the transform
            ref: the reference (or target) of the transform

        Returns:
            Registration object
        """

        if isinstance(src2ref, str):
            src2ref = np.loadtxt(src2ref)

        if not isinstance(src, ImageSpace):
            src = ImageSpace(src)
        src_spc = src
        if not isinstance(ref, ImageSpace):
            ref = ImageSpace(ref)
        ref_spc = ref

        src2ref = (ref_spc.FSL2world @ src2ref @ src_spc.world2FSL)
        return Registration(src2ref)

    def __len__(self):
        return 1

    def __repr__(self):

        formatter = "{:8.3f}".format
        with np.printoptions(precision=3, formatter={'all': formatter}):
            text = (f"""\
                Registration (linear) with properties:
                src2ref:       {self.src2ref[0,:]}
                               {self.src2ref[1,:]}
                               {self.src2ref[2,:]}
                               {self.src2ref[3,:]}""")
        return dedent(text)

    @property
    def ref2src(self):
        return np.linalg.inv(self.__src2ref)

    @property
    def src2ref(self):
        return self.__src2ref

    @classmethod
    def identity(cls):
        return Registration(np.eye(4))

    def inverse(self):
        """Self inverse"""
        constructor = type(self)
        return constructor(self.ref2src)

    def to_fsl(self, src, ref):
        """
        Return transformation in FSL convention, for given src and ref,
        as np.array. This will be 3D in the case of MotionCorrections
        """

        if not isinstance(src, ImageSpace):
            src = ImageSpace(src)
        if not isinstance(ref, ImageSpace):
            ref = ImageSpace(ref)

        return ref.world2FSL @ self.src2ref @ src.FSL2world

    def to_flirt(self, src, ref):
        """Alias for self.to_fsl()"""
        return self.to_fsl(src, ref)

    def save_txt(self, path):
        """Save as textfile at path"""
        np.savetxt(path, self.src2ref)

    def save_fsl(self, path, src, ref):
        """Save in FSL convention as textfile at path"""
        m = self.to_fsl(src, ref)
        np.savetxt(path, m)

    def prepare_cache(self, ref):
        """
        Cache re-useable data before interpolate_and_scale. Just the voxel
        index grid of the reference space is stored
        """

        self.cache = ref.ijk_grid('ij').reshape(-1, 3)

    def resolve(self, src, ref, *unused):
        """
        Return a coordinate array and scale factor that maps reference voxels
        into source voxels, including the transform. Uses cached values, if
        available.

        Args:
            src (ImageSpace): in which data currently exists and interpolation
                will be performed
            ref (ImageSpace): in which data needs to be expressed

        Returns:
            (np.ndarray, 1) coordinates on which to interpolate and identity
                scale factor
        """

        # Array of all voxel indices in the reference grid
        # Map them into world coordinates, apply the transform
        # and then into source voxel coordinates for the interpolation
        ref2src_vox = (src.world2vox @ self.ref2src @ ref.vox2world)
        ijk = apply.aff_trans(ref2src_vox, self.cache).T
        scale = 1
        return (ijk, scale)

class MotionCorrection(Registration):
    """
    A sequence of Registration objects, one for each volume in a timeseries.

    Args:
        mats: a path to a directory containing transformation matrices, in
            name order (all files will be loaded), or a list of individual
            filenames, or a list of np.arrays
    """

    def __init__(self, mats):
        Transform.__init__(self)

        if isinstance(mats, str):
            if op.isdir(mats):
                mats = sorted(glob.glob(op.join(mats, '*')))
                if not mats:
                    raise RuntimeError("Did not find any matrices in %s" % mats)
            else:
                mat = np.loadtxt(mats)
                if mat.shape[0] % 4:
                    raise ValueError("Matrix loaded from %s " % mats,
                                     "should be sized (4xN) x 4.")
                mats = [ mat[i*4:(i+1)*4,:] for i in range(mat.shape[0] // 4) ]

        self.__transforms = []
        for mat in mats:
            if isinstance(mat, (np.ndarray, str)):
                m = Registration(mat)
            else:
                m = mat
            self.__transforms.append(m)

    def from_flirt(self, *args):
        raise NotImplementedError("Use the MotionCorrection.from_mcflirt() method")

    @classmethod
    def from_mcflirt(cls, mats, src, ref):
        """
        Load a MotionCorrection object directly from MCFLIRT's -omat directory.
        Note that for within-timeseries registration, the src and ref arguments
        should take the same value.

        Args:
            mats: a path to a directory containing transformation matrices, in
                name order (all files will be loaded), or a list of individual
                filenames, or a list of np.arrays
            src: source of the transforms (ie, the image being corrected)
            ref: the target of the transforms (normally same as src)

        Returns:
            MotionCorrection
        """

        if isinstance(mats, str):
            mats = sorted(glob.glob(op.join(mats, '*')))
            if not mats:
                raise RuntimeError("Did not find any matrices in %s" % mats)

        if isinstance(mats, np.ndarray):
            if mats.ndim == 3:
                if not mats.shape[:2] ==(4,4):
                    raise ValueError("A 3D stack of matrices should have size "
                            "(4,4) in first two dimensions")
            if mats.ndim == 2:
                if ((mats.shape[0] % 4) or (mats.shape[1] != 4)):
                    raise ValueError("A 2D array should have size Nx4, "
                        "where N is divisible by 4")
                mats = [ mats[4*m : 4*(m+1),:] for m in range(mats.shape[0] // 4) ]

        return MotionCorrection([ Registration.from_flirt(m, src, ref) for m in mats ])

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

    def __repr__(self):
        t = self[0]

        formatter = "{:8.3f}".format
        with np.printoptions(precision=3, formatter={'all': formatter}):
            text = (f"""\
                MotionCorrection (linear) with properties:
                series length:   {len(self)}
                src2ref_0:       {t.src2ref[0,:]}
                                 {t.src2ref[1,:]}
                                 {t.src2ref[2,:]}
                                 {t.src2ref[3,:]}""")
        return dedent(text)

    def __getitem__(self, idx):
        """Access individual Registration objects from within series"""
        return self.__transforms[idx]

    @classmethod
    def identity(cls, length):
        return MotionCorrection([Registration.identity()] * length)

    @classmethod
    def from_registration(cls, reg, length):
        """
        Produce a MotionCorrection by repeating a Registration object
        n times (eg, 10 copies of a single transform)
        """

        return MotionCorrection([reg.src2ref] * length)

    @property
    def transforms(self):
        """List of Registration objects representing each volume of transform"""
        return self.__transforms

    @property
    def src2ref(self):
        """List of src to ref transformation matrices"""
        return [ t.src2ref for t in self.transforms ]

    @property
    def ref2src(self):
        """List of ref to src transformation matrices"""
        return [ t.ref2src for t in self.transforms ]

    def to_fsl(self, src, ref):
        """Transformation matrices in FSL terms"""
        return [ t.to_fsl(src, ref) for t in self.transforms ]

    def save_txt(self, outdir, prefix="MAT_"):
        """
        Save individual transformation matrices in text format
        in outdir. Matrices will be named prefix_001...

        Args:
            outdir: directory in which to save
            src: (optional) path to image, or ImageSpace, source space of
                transformation
            ref: as above, for reference space of transformation
            convention: "world" or "fsl", if fsl then src/ref must be given
            prefix: prefix for naming each matrix
        """

        os.makedirs(outdir, exist_ok=True)
        for idx, r in enumerate(self.transforms):
            p = op.join(outdir, "{}{:04d}.txt".format(prefix, idx))
            r.save_txt(p)

    def save_fsl(self, outdir, src, ref, prefix="MAT_"):
        """Save in FSL convention as textfiles at path"""

        os.makedirs(outdir, exist_ok=True)
        for idx, r in enumerate(self.transforms):
            p = op.join(outdir, "{}{:04d}.txt".format(prefix, idx))
            m = r.to_fsl(src, ref)
            np.savetxt(p, m)

    def resolve(self, src, ref, at_idx):
        """
        Return a coordinate array and scale factor that maps reference voxels
        into source voxels, including the transform. Uses cached values, if
        available.

        Args:
            src (ImageSpace): in which data currently exists and interpolation
                will be performed
            ref (ImageSpace): in which data needs to be expressed
            at_idx (int): index number within series of transforms to apply

        Returns:
            (np.ndarray, 1) coordinates on which to interpolate and identity
                scale factor
        """

        # Array of all voxel indices in the reference grid
        # Map them into world coordinates, apply the transform
        # and then into source voxel coordinates for the interpolation
        ref2src_vox = (src.world2vox
                       @ self.ref2src[at_idx]
                       @ ref.vox2world)
        ijk = apply.aff_trans(ref2src_vox, self.cache).T
        scale = 1
        return ijk, scale


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

import os
#os.environ['CUDA_VISIBLE_DEVICES'] = ''

import numpy as np
#import cv2
#import matplotlib.pyplot as plt
import  sys
sys.path.append('/home/mydisk/code/matting/Unet4')
sys.path.append('/home/mydisk/code/matting/Unet4/segmentation_models.pytorch-master')
import segmentation_models_pytorch as smp
import albumentations as albu
import torch
import time
from torch.autograd import Variable

import torch
print(torch.version.__version__)

#384x512

#pth_path = '/home/pt/Desktop/code/matting/traindata/unet4/yxl_1/res18_yxl.pth'
pth_path = '/home/mydisk/code/matting/traindata/unet4/best_model.pth'

model = torch.load(pth_path)
model=model.cuda()

#dummy_input = Variable(torch.randn(1, 3, 224 , 224)).cpu()
#dummy_input = torch.randn(1, 4, 320, 320, requires_grad=True)
model=model.cpu()
print(model)

#model = torchvision.models.resnet18(pretrained=True)
model.eval()
example = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("/home/mydisk/code/androidA/android-demo-app/HelloWorldApp/app/src/main/assets/my_model.pt")


import numpy as np
from keras.models import load_model
from dataload import loadTestData
from rouge import Rouge
import nltk

def dummy_rouge(sentence, summary, alpha=0.5):
    if sentence in summary:
        return 1
    return 0

def dummy_loadTestData():
    testing_data = [ [ np.array(["This sentence is important for doc0." ,
                                 "Such a sentence is irrelevent for doc 0."]),
                       np.random.rand(2,5,300),
                       np.array(["This sentence is important for doc0."]) ],
                     [ np.array(["Lol that sentence is awesome for do1." ,
                                 "No way, this is irrelevent"]),
                       np.random.rand(2,5,300),
                                np.array(["Lol that sentence is awesome for do1."]) ] ]
    return testing_data

def test(model, testing_data, batch_size = 128, upper_bound = 100, threshold = 1, metric = "ROUGE1"):
    """
        Build the actual summaries for test data and evaluate them
        To do:
            - load the actual x_test (embed test sentences) and y_test (compute rouge score)

        Parameters:
            testing_data           - np.array
                                        ex: [ doc1, doc2, ... , docn]
                                         where doci = [sentences, x_test, summary]
                                             where sentences = np.array of string
                                                   x_test = np.array of matrices (embedded sentences)
                                                   summaries = np.array of sentences

        Returns:
            eval                   - float between 0 and 1.
    """
    rouge = Rouge()
    r1evals = []
    r2evals = []
    summaries = []

    r1bestPred = []
    r1bestTrue = []
    r1worstPred = []
    r1worstTrue = []
    r1bestScore = -1.0
    r1worstScore = 1.1

    r2bestPred = []
    r2bestTrue = []
    r2worstPred = []
    r2worstTrue = []
    r2bestScore = -1.0
    r2worstScore = 1.1

    for doc in testing_data:
        sentences = doc[0]

        x_test_old = doc[1]
        s1 = x_test_old.shape[0]
        (s3,s4) = x_test_old[0].shape
        print(s1,s3,s4)
        x_test = np.random.rand(s1,1,190,s4)
        for i in range(s1) :
            x_test[i] = np.array( [ np.pad(x_test_old[i], ((190-s3,0),(0,0)), 'constant') ] )

        true_summary = doc[2]

        predicted_scores = model.predict(x_test, batch_size=batch_size)
        #argsorted_scores= np.argsort(predicted_scores)
        argsorted_scores = np.argpartition(np.transpose(predicted_scores)[0], 1)

        predicted_summary = []
        summary_length = 0

        i = 0

        while i < len(sentences) and summary_length < upper_bound:
            sentence = sentences[argsorted_scores[i]]
            #if ( dummy_rouge( sentence , predicted_summary ) < threshold ):
            sentence = np.array([sentence])
            #print(sentence, predicted_summary)
            predicted_summary.append(sentence)
            summary_length += len(nltk.word_tokenize(sentence[0]))

            i+=1

        #print(predicted_summary)

        #if metric == "ROUGE1" :
        #    N = 1
        #elif metric == "ROUGE2":
        #    N = 0

        #evals.append(dummy_rouge( predicted_summary, true_summary, alpha = N))
        r1score = rouge.saliency(predicted_summary, true_summary, alpha=1)
        r2score = rouge.saliency(predicted_summary, true_summary, alpha=0)

        if r1score > r1bestScore:
            r1bestScore = r1score
            r1bestPred = predicted_summary
            r1bestTrue = true_summary
        if r1score < r1worstScore:
            r1worstScore = r1score
            r1worstPred = predicted_summary
            r1worstTrue = true_summary

        if r2score > r2bestScore:
            r2bestScore = r2score
            r2bestPred = predicted_summary
            r2bestTrue = true_summary
        if r2score < r2worstScore:
            r2worstScore = r2score
            r2worstPred = predicted_summary
            r2worstTrue = true_summary

        r1evals.append(r1score)
        r2evals.append(r2score)

        #evals.append(rouge.saliency(predicted_summary, true_summary, alpha=N))
        summaries.append((predicted_summary, true_summary))

    print("&& PRINTING BEST AND WORST SUMMARIES && ")

    print("ROUGE 1 --")
    print("BEST:")
    print(" score:", r1bestScore)
    print(" predicted:", r1bestPred)
    print(" true:", r1bestTrue)

    print()
    print("WORST:")
    print(" score:", r1worstScore)
    print(" predicted:", r1worstPred)
    print(" true:", r1worstTrue)
    print()

    print("ROUGE 2 --")
    print("BEST:")
    print(" score:", r2bestScore)
    print(" predicted:", r2bestPred)
    print(" true:", r2bestTrue)

    print()
    print("WORST:")
    print(" score:", r2worstScore)
    print(" predicted:", r2worstPred)
    print(" true:", r2worstTrue)

    print(" *--*--*--*--*--*--*--*")

    return np.mean(r1evals), np.mean(r2evals)

def main():
    model = load_model('../model-nfilt-200.h5')
    #testing_data = dummy_loadTestData()
    #testing_data = loadTestData("../data/DUC2002_Summarization_Documents")
    testing_data = loadTestData("../data/test_subset")
    print(testing_data)

    rouge1_score, rouge2_score = test(model, testing_data, upper_bound=100, metric = "ROUGE1")
    #rouge2_score = test(model, testing_data, upper_bound=100, metric = "ROUGE2")
    print("ROUGE1:",rouge1_score)
    print("ROUGE2:",rouge2_score)

if __name__ == "__main__":
    main()


import cv2 as cv
import numpy as np
from tensorflow.keras.models import load_model
model=load_model('facial.h5')
cv.ocl.setUseOpenCL(False)
dic={0:"Angry",1:"Disgusted",2:"Fearful",3:"Happy",4:"Neutral",5:"Sad",6:"Surprised"}
cap=cv.VideoCapture(0)
while True:
    ret,frame=cap.read()
    if not ret:
        break
    bounding_box=cv.CascadeClassifier(r'C:\Users\naman\anaconda3\Lib\site-packages\cv2\data\haarcascade_frontalface_default.xml')
    grey_frame=cv.cvtColor(frame,cv.COLOR_BGR2GRAY)
    faces=bounding_box.detectMultiScale(grey_frame,scaleFactor=1.3,minNeighbors=5)
    for (x,y,w,h) in faces:
        cv.rectangle(frame, (x,y-50),(x+w, y+h+10),(255,0,0),2)
        roi=grey_frame[y:y+h,x:x+w]
        cropped=np.expand_dims(np.expand_dims(cv.resize(roi,(48,48)),-1),0)
        pred=model.predict(cropped)
        maxindex=int(np.argmax(pred))
        cv.putText(frame,dic[maxindex],(x+20,y-60),cv.FONT_HERSHEY_SIMPLEX,1,(255,255,255),2,cv.LINE_AA)
    cv.imshow('Video',cv.resize(frame,(1200,860),interpolation=cv.INTER_CUBIC))
    if cv.waitKey(1)&0xFF==ord('q'):
        cap.release()
        cv.destroyAllWindows()
        break


import sys
import pydart2 as pydart
import numpy as np

from skate_cma.hpcma import HpCma
from skate_cma.PenaltyType import PenaltyType

def objective(i, penalty_option0, penalty_option1, penalty_weight):
    """
    0 : com h, left foot, #left foot com, max y velocity end, max y angvel end
    1 : am, axis, height com end, max y diff end
    2 : am, axis
    3 : right foot end, right foot end com
    4 : com height, right foot, right foot com, right foot com end, right foot end
    """

    penalty_weight[PenaltyType.MAX_Y_ANGULAR_MOMENTUM] = 0.1
    penalty_weight[PenaltyType.MAX_Y_VELOCITY_END] = 0.01
    penalty_weight[PenaltyType.MAX_Y_ANGULAR_MOMENTUM_END] = 0.01
    penalty_weight[PenaltyType.MAX_Y_POS_DIFF_END] = 0.1
    penalty_weight[PenaltyType.TORQUE] = 0.

    if i == 0:
        # jump ready
        penalty_option0[PenaltyType.COM_HEIGHT] = 0.64
        penalty_option0[PenaltyType.LEFT_FOOT_CONTACT] = True
        penalty_option0[PenaltyType.MAX_Y_VELOCITY_END] = True
        penalty_option0[PenaltyType.MAX_Y_ANGULAR_MOMENTUM_END] = True

        # jump max
        penalty_option1[PenaltyType.MAX_Y_ANGULAR_MOMENTUM] = True
        penalty_option1[PenaltyType.COM_IN_LEFT_FOOT] = True
        penalty_option1[PenaltyType.COM_IN_RIGHT_FOOT] = True
        penalty_option1[PenaltyType.COM_IN_HEAD] = True
        penalty_option1[PenaltyType.MAX_Y_POS_DIFF_END] = True
        penalty_option1[PenaltyType.COM_HEIGHT_END] = 1.45

    elif i == 1:
        # jump max
        penalty_option0[PenaltyType.MAX_Y_ANGULAR_MOMENTUM] = True
        penalty_option0[PenaltyType.COM_IN_LEFT_FOOT] = True
        penalty_option0[PenaltyType.COM_IN_RIGHT_FOOT] = True
        penalty_option0[PenaltyType.COM_IN_HEAD] = True
        penalty_option0[PenaltyType.MAX_Y_POS_DIFF_END] = True
        penalty_option0[PenaltyType.COM_HEIGHT_END] = 1.45

        # jump down
        penalty_option1[PenaltyType.MAX_Y_ANGULAR_MOMENTUM] = True
        penalty_option1[PenaltyType.COM_IN_RIGHT_FOOT] = True
        penalty_option1[PenaltyType.COM_IN_LEFT_FOOT] = True
        penalty_option1[PenaltyType.COM_IN_HEAD] = True

    elif i == 2:
        # jump down
        penalty_option0[PenaltyType.MAX_Y_ANGULAR_MOMENTUM] = True
        penalty_option0[PenaltyType.COM_IN_LEFT_FOOT] = True
        penalty_option0[PenaltyType.COM_IN_RIGHT_FOOT] = True
        penalty_option0[PenaltyType.COM_IN_HEAD] = True

        # prepare landing
        penalty_option1[PenaltyType.COM_IN_RIGHT_FOOT_END] = True
        penalty_option1[PenaltyType.RIGHT_FOOT_CONTACT_END] = True

    elif i == 3:
        # prepare landing
        penalty_option0[PenaltyType.COM_IN_RIGHT_FOOT_END] = True
        penalty_option0[PenaltyType.RIGHT_FOOT_CONTACT_END] = True

        # landed
        penalty_option1[PenaltyType.COM_IN_RIGHT_FOOT] = True
        penalty_option1[PenaltyType.RIGHT_FOOT_CONTACT] = True
        penalty_option1[PenaltyType.COM_HEIGHT] = 0.75
        penalty_option1[PenaltyType.COM_IN_RIGHT_FOOT_END] = True
        penalty_option1[PenaltyType.RIGHT_FOOT_CONTACT_END] = True
    elif i == 4:
        # landed
        penalty_option0[PenaltyType.COM_IN_RIGHT_FOOT] = True
        penalty_option0[PenaltyType.RIGHT_FOOT_CONTACT] = True
        penalty_option0[PenaltyType.COM_HEIGHT] = 0.75
        penalty_option0[PenaltyType.COM_IN_RIGHT_FOOT_END] = True
        penalty_option0[PenaltyType.RIGHT_FOOT_CONTACT_END] = True

if __name__ == '__main__':
    pydart.init()
    if len(sys.argv) == 1:
        cma = HpCma('jump0507_2', 4, sigma=0.1, max_time=2., cma_timeout=1)
        # cma = HpCma('jump0507_2', 4, sigma=0.1, max_time=2., start_state_num=4, start_state_sol_dir='jump0507_2_model_201905192040/', cma_timeout=1)
    elif len(sys.argv) == 5:
        cma = HpCma(sys.argv[1], int(sys.argv[2]), sigma=float(sys.argv[3]), max_time=float(sys.argv[4]))
    elif len(sys.argv) == 8:
        cma = HpCma(sys.argv[1], int(sys.argv[2]), sigma=float(sys.argv[3]), max_time=float(sys.argv[4]), start_state_num=int(sys.argv[5]), start_state_sol_dir=sys.argv[6], cma_timeout=int(sys.argv[7]))
    else:
        cma = HpCma('jump0507_2', 1)

    dq = np.zeros(57)
    dq[3] = -1.5
    cma.set_init_dq(dq)
    cma.objective = objective
    cma.run()


"""Hourglass Network for Keras.
# Reference paper
- [Objects as Points]
  (https://arxiv.org/pdf/1904.07850.pdf)
# Reference implementation
- [PyTorch CenterNet]
  (https://github.com/xingyizhou/CenterNet/blob/master/src/lib/models/networks/large_hourglass.py)
- [Keras Stacked_Hourglass_Network_Keras]
  (https://github.com/yuanyuanli85/Stacked_Hourglass_Network_Keras/blob/master/src/net/hourglass.py)
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

from keras.models import Model
from keras.layers import Conv2D, Input, Activation, BatchNormalization, Add, UpSampling2D, ZeroPadding2D
from keras.utils import get_file
import keras.backend as K
import numpy as np

CTDET_COCO_WEIGHTS_PATH = (
  'https://github.com/see--/keras-centernet/'
  'releases/download/0.1.0/ctdet_coco_hg.hdf5')

HPDET_COCO_WEIGHTS_PATH = (
  'https://github.com/see--/keras-centernet/'
  'releases/download/0.1.0/hpdet_coco_hg.hdf5')

def normalize_image(image):
  """Normalize the image for the Hourglass network.
  # Arguments
    image: BGR uint8
  # Returns
    float32 image with the same shape as the input
  """
  mean = [0.40789655, 0.44719303, 0.47026116]
  std = [0.2886383, 0.27408165, 0.27809834]
  return ((np.float32(image) / 255.) - mean) / std

def HourglassNetwork(heads, num_stacks, cnv_dim=256, inres=(512, 512), weights='ctdet_coco',
                     dims=[256, 384, 384, 384, 512]):
  """Instantiates the Hourglass architecture.
  Optionally loads weights pre-trained on COCO.
  Note that the data format convention used by the model is
  the one specified in your Keras config at `~/.keras/keras.json`.
  # Arguments
      num_stacks: number of hourglass modules.
      cnv_dim: number of filters after the resolution is decreased.
      inres: network input shape, should be a multiple of 128.
      weights: one of `None` (random initialization),
            'ctdet_coco' (pre-training on COCO for 2D object detection),
            'hpdet_coco' (pre-training on COCO for human pose detection),
            or the path to the weights file to be loaded.
      dims: numbers of channels in the hourglass blocks.
  # Returns
      A Keras model instance.
  # Raises
      ValueError: in case of invalid argument for `weights`,
          or invalid input shape.
  """
  if not (weights in {'ctdet_coco', 'hpdet_coco', None} or os.path.exists(weights)):
    raise ValueError('The `weights` argument should be either '
                     '`None` (random initialization), `ctdet_coco` '
                     '(pre-trained on COCO), `hpdet_coco` (pre-trained on COCO) '
                     'or the path to the weights file to be loaded.')
  input_layer = Input(shape=(inres[0], inres[1], 3), name='HGInput')
  inter = pre(input_layer, cnv_dim)
  prev_inter = None
  outputs = []
  for i in range(num_stacks):
    prev_inter = inter
    _heads, inter = hourglass_module(heads, inter, cnv_dim, i, dims)
    outputs.extend(_heads)
    if i < num_stacks - 1:
      inter_ = Conv2D(cnv_dim, 1, use_bias=False, name='inters_.%d.0' % i)(prev_inter)
      inter_ = BatchNormalization(epsilon=1e-5, name='inters_.%d.1' % i)(inter_)

      cnv_ = Conv2D(cnv_dim, 1, use_bias=False, name='cnvs_.%d.0' % i)(inter)
      cnv_ = BatchNormalization(epsilon=1e-5, name='cnvs_.%d.1' % i)(cnv_)

      inter = Add(name='inters.%d.inters.add' % i)([inter_, cnv_])
      inter = Activation('relu', name='inters.%d.inters.relu' % i)(inter)
      inter = residual(inter, cnv_dim, 'inters.%d' % i)

  model = Model(inputs=input_layer, outputs=outputs)
  if weights == 'ctdet_coco':
    weights_path = get_file(
      '%s_hg.hdf5' % weights,
      CTDET_COCO_WEIGHTS_PATH,
      cache_subdir='models',
      file_hash='ce01e92f75b533e3ff8e396c76d55d97ff3ec27e99b1bdac1d7b0d6dcf5d90eb')
    model.load_weights(weights_path)
  elif weights == 'hpdet_coco':
    weights_path = get_file(
      '%s_hg.hdf5' % weights,
      HPDET_COCO_WEIGHTS_PATH,
      cache_subdir='models',
      file_hash='5c562ee22dc383080629dae975f269d62de3a41da6fd0c821085fbee183d555d')
    model.load_weights(weights_path)
  elif weights is not None:
    model.load_weights(weights,by_name=True)

  return model

def hourglass_module(heads, bottom, cnv_dim, hgid, dims):
  # create left features , f1, f2, f4, f8, f16 and f32
  lfs = left_features(bottom, hgid, dims)

  # create right features, connect with left features
  rf1 = right_features(lfs, hgid, dims)
  rf1 = convolution(rf1, 3, cnv_dim, name='cnvs.%d' % hgid)

  # add 1x1 conv with two heads, inter is sent to next stage
  # head_parts is used for intermediate supervision
  heads = create_heads(heads, rf1, hgid)
  return heads, rf1

def convolution(_x, k, out_dim, name, stride=1):
  padding = (k - 1) // 2
  _x = ZeroPadding2D(padding=padding, name=name + '.pad')(_x)
  _x = Conv2D(out_dim, k, strides=stride, use_bias=False, name=name + '.conv')(_x)
  _x = BatchNormalization(epsilon=1e-5, name=name + '.bn')(_x)
  _x = Activation('relu', name=name + '.relu')(_x)
  return _x

def residual(_x, out_dim, name, stride=1):
  shortcut = _x
  num_channels = K.int_shape(shortcut)[-1]
  _x = ZeroPadding2D(padding=1, name=name + '.pad1')(_x)
  _x = Conv2D(out_dim, 3, strides=stride, use_bias=False, name=name + '.conv1')(_x)
  _x = BatchNormalization(epsilon=1e-5, name=name + '.bn1')(_x)
  _x = Activation('relu', name=name + '.relu1')(_x)

  _x = Conv2D(out_dim, 3, padding='same', use_bias=False, name=name + '.conv2')(_x)
  _x = BatchNormalization(epsilon=1e-5, name=name + '.bn2')(_x)

  if num_channels != out_dim or stride != 1:
    shortcut = Conv2D(out_dim, 1, strides=stride, use_bias=False, name=name + '.skip.0')(
        shortcut)
    shortcut = BatchNormalization(epsilon=1e-5, name=name + '.skip.1')(shortcut)

  _x = Add(name=name + '.add')([_x, shortcut])
  _x = Activation('relu', name=name + '.relu')(_x)
  return _x

def pre(_x, num_channels):
  # front module, input to 1/4 resolution
  _x = convolution(_x, 7, 128, name='pre.0', stride=2)
  _x = residual(_x, num_channels, name='pre.1', stride=2)
  return _x

def left_features(bottom, hgid, dims):
  # create left half blocks for hourglass module
  # f1, f2, f4 , f8, f16, f32 : 1, 1/2, 1/4 1/8, 1/16, 1/32 resolution
  # 5 times reduce/increase: (256, 384, 384, 384, 512)
  features = [bottom]
  for kk, nh in enumerate(dims):
    pow_str = ''
    for _ in range(kk):
      pow_str += '.low2'
    _x = residual(features[-1], nh, name='kps.%d%s.low1.0' % (hgid, pow_str), stride=2)
    _x = residual(_x, nh, name='kps.%d%s.low1.1' % (hgid, pow_str))
    features.append(_x)
  return features

def connect_left_right(left, right, num_channels, num_channels_next, name):
  # left: 2 residual modules
  left = residual(left, num_channels_next, name=name + 'up1.0')
  left = residual(left, num_channels_next, name=name + 'up1.1')

  # up: 2 times residual & nearest neighbour
  out = residual(right, num_channels, name=name + 'low3.0')
  out = residual(out, num_channels_next, name=name + 'low3.1')
  out = UpSampling2D(name=name + 'low3.upsampleNN')(out)
  out = Add(name=name + 'low3.add')([left, out])
  return out

def bottleneck_layer(_x, num_channels, hgid):
  # 4 residual blocks with 512 channels in the middle
  pow_str = 'low2.' * 5
  _x = residual(_x, num_channels, name='kps.%d.%s0' % (hgid, pow_str))
  _x = residual(_x, num_channels, name='kps.%d.%s1' % (hgid, pow_str))
  _x = residual(_x, num_channels, name='kps.%d.%s2' % (hgid, pow_str))
  _x = residual(_x, num_channels, name='kps.%d.%s3' % (hgid, pow_str))
  return _x

def right_features(leftfeatures, hgid, dims):
  rf = bottleneck_layer(leftfeatures[-1], dims[-1], hgid)
  for kk in reversed(range(len(dims))):
    pow_str = ''
    for _ in range(kk):
      pow_str += 'low2.'
    rf = connect_left_right(leftfeatures[kk], rf, dims[kk], dims[max(kk - 1, 0)], name='kps.%d.%s' % (hgid, pow_str))
  return rf

def create_heads(heads, rf1, hgid):
  _heads = []
  for head in sorted(heads):
    num_channels = heads[head]
    _x = Conv2D(256, 3, use_bias=True, padding='same', name=head + '.%d.0.conv' % hgid)(rf1)
    _x = Activation('relu', name=head + '.%d.0.relu' % hgid)(_x)
    _x = Conv2D(num_channels, 1, use_bias=True, name=head + '.%d.1' % hgid)(_x)
    _heads.append(_x)
  return _heads

if __name__ == '__main__':
  kwargs = {
    'num_stacks': 2,
    'cnv_dim': 256,
    'inres': (512, 512),
  }
  heads = {
    'hm': 80,
    'reg': 2,
    'wh': 2
  }
  model = HourglassNetwork(heads=heads, **kwargs)
  print(model.summary(line_length=200))
  # from IPython import embed; embed()


import torch, math, time, argparse, os
import random, dataset, utils, losses, net
import numpy as np

from dataset.Inshop import Inshop_Dataset
from net.resnet import *
from net.googlenet import *
from net.bn_inception import *
from dataset import sampler
from torch.utils.data.sampler import BatchSampler
from torch.utils.data.dataloader import default_collate

from tqdm import *
import wandb

seed = 1
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)  # set random seed for all gpus

parser = argparse.ArgumentParser(description=
                                 'Official implementation of `Proxy Anchor Loss for Deep Metric Learning`'
                                 + 'Our code is modified from `https://github.com/dichotomies/proxy-nca`'
                                 )
# export directory, training and val datasets, test datasets

parser.add_argument('--workers', default=4, type=int, dest='nb_workers', help='Number of workers for dataloader.')
parser.add_argument('--l2-norm', default=1, type=int, help='L2 normlization')
parser.add_argument('--bn-freeze', default = 1, type = int,help = 'Batch normalization parameter freeze')
parser.add_argument('--remark', default='', help='Any reamrk')

parser.add_argument('--LOG_DIR', default='../logs/cars_512', help='Path to log folder')
parser.add_argument('--batch-size', default=32, type=int, dest='sz_batch', help='Number of samples per batch.')
parser.add_argument('--dataset', default='cars', help='Training dataset, e.g. cub, cars, SOP, Inshop, logo2k')
parser.add_argument('--data_root', default='/home/ruofan/PycharmProjects/ProxyNCA-/mnt/datasets/')
parser.add_argument('--embedding-size', default=512, type=int, dest='sz_embedding',
                    help='Size of embedding that is appended to backbone model.')
parser.add_argument('--model', default='resnet50', help='Model for training')
parser.add_argument('--loss', default='MS', help='Criterion for training')
parser.add_argument('--gpu-id', default=1, type=int, help='ID of GPU that is used for training. -1 means use all')
parser.add_argument('--lr', default=1e-5, type=float, help='Learning rate setting')

args = parser.parse_args()

if args.gpu_id != -1:
    torch.cuda.set_device(args.gpu_id)

# Directory for Log
LOG_DIR = args.LOG_DIR
data_root = args.data_root

if args.dataset != 'Inshop':
    ev_dataset = dataset.load(
        name=args.dataset,
        root=data_root,
        mode='eval',
        transform=dataset.utils.make_transform(
            is_train=False,
            is_inception=(args.model == 'bn_inception')
        ))

    dl_ev = torch.utils.data.DataLoader(
        ev_dataset,
        batch_size=args.sz_batch,
        shuffle=False,
        num_workers=args.nb_workers,
        pin_memory=True
    )

else:
    query_dataset = Inshop_Dataset(
        root=data_root,
        mode='query',
        transform=dataset.utils.make_transform(
            is_train=False,
            is_inception=(args.model == 'bn_inception')
        ))

    dl_query = torch.utils.data.DataLoader(
        query_dataset,
        batch_size=args.sz_batch,
        shuffle=False,
        num_workers=args.nb_workers,
        pin_memory=True
    )

    gallery_dataset = Inshop_Dataset(
        root=data_root,
        mode='gallery',
        transform=dataset.utils.make_transform(
            is_train=False,
            is_inception=(args.model == 'bn_inception')
        ))

    dl_gallery = torch.utils.data.DataLoader(
        gallery_dataset,
        batch_size=args.sz_batch,
        shuffle=False,
        num_workers=args.nb_workers,
        pin_memory=True
    )

# Backbone Model
if args.model.find('googlenet') + 1:
    model = googlenet(embedding_size=args.sz_embedding, pretrained=True, is_norm=args.l2_norm, bn_freeze=args.bn_freeze)
elif args.model.find('bn_inception') + 1:
    model = bn_inception(embedding_size=args.sz_embedding, pretrained=True, is_norm=args.l2_norm,
                         bn_freeze=args.bn_freeze)
elif args.model.find('resnet18') + 1:
    model = Resnet18(embedding_size=args.sz_embedding, pretrained=True, is_norm=args.l2_norm, bn_freeze=args.bn_freeze)
elif args.model.find('resnet50') + 1:
    model = Resnet50(embedding_size=args.sz_embedding, pretrained=True, is_norm=args.l2_norm, bn_freeze=args.bn_freeze)
elif args.model.find('resnet101') + 1:
    model = Resnet101(embedding_size=args.sz_embedding, pretrained=True, is_norm=args.l2_norm, bn_freeze=args.bn_freeze)
model = model.cuda()
model.load_state_dict(torch.load('{}/{}_{}_{}_best.pth'.format(LOG_DIR, args.dataset, args.model, args.loss))['model_state_dict'])

print("**Evaluating...**")
if args.dataset == 'Inshop':
    utils.evaluate_cos_Inshop(model, dl_query, dl_gallery)
elif args.dataset != 'SOP':
    utils.evaluate_cos(model, dl_ev)


# This code is auto-generated.
import http.client as http_client
import logging
import os

import numpy as np
from joblib import load
from scipy import sparse

from sagemaker_containers.beta.framework import encoders
from sagemaker_containers.beta.framework import worker
from sagemaker_sklearn_extension.externals import read_csv_data

def _is_inverse_label_transform():
    """Returns True if if it's running in inverse label transform."""
    return os.getenv('AUTOML_TRANSFORM_MODE') == 'inverse-label-transform'

def _is_feature_transform():
    """Returns True if it's running in feature transform mode."""
    return os.getenv('AUTOML_TRANSFORM_MODE') == 'feature-transform'

def _sparsify_if_needed(x):
    """Returns a sparse matrix if the needed for encoding to sparse recordio protobuf."""
    if os.getenv('AUTOML_SPARSE_ENCODE_RECORDIO_PROTOBUF') == '1' \
            and not sparse.issparse(x):
        return sparse.csr_matrix(x)
    return x

def _split_features_target(x):
    """Returns the features and target by splitting the input array."""
    if os.getenv('AUTOML_TRANSFORM_MODE') == 'feature-transform':
        return _sparsify_if_needed(x), None

    if sparse.issparse(x):
        return x[:, 1:], x[:, 0].toarray()
    return _sparsify_if_needed(x[:, 1:]), np.ravel(x[:, 0])

def model_fn(model_dir):
    """Loads the model.

    The SageMaker Scikit-learn model server loads model by invoking this method.

    Parameters
    ----------
    model_dir: str
        the directory where the model files reside

    Returns
    -------
    : AutoMLTransformer
        deserialized model object that can be used for model serving

    """
    return load(filename=os.path.join(model_dir, 'model.joblib'))

def predict_fn(input_object, model):
    """Generates prediction for the input_object based on the model.

    The SageMaker Scikit-learn model server invokes this method with the return value of the input_fn.
    This method invokes the loaded model's transform method, if it's expected to transform the input data
    and invokes the loaded model's inverse_label_transform, if the task is to perform
    inverse label transformation.

    Parameters
    ----------
    input_object : array-like
        the object returned from input_fn

    model : AutoMLTransformer
        the model returned from model_fn

    Returns
    -------
    : ndarray
        transformed input data or inverse transformed label data

    """
    if isinstance(input_object, worker.Response):
        return input_object

    if _is_inverse_label_transform():
        return model.inverse_label_transform(
            input_object.ravel().astype(np.float).astype(np.int)
        )
    try:
        return model.transform(input_object)
    except ValueError as e:
        return worker.Response(
            response='{}'.format(str(e) or 'Unknown error.'),
            status=http_client.BAD_REQUEST
        )

def input_fn(request_body, request_content_type):
    """Decodes request body to 2D numpy array.

    The SageMaker Scikit-learn model server invokes this method to deserialize the request data into an object
    for prediction.

    Parameters
    ----------
    request_body : str
        the request body
    request_content_type : str
        the media type for the request body

    Returns
    -------
    : array-like
        decoded data as 2D numpy array

    """
    content_type = request_content_type.lower(
    ) if request_content_type else "text/csv"
    content_type = content_type.split(";")[0].strip()

    if content_type == 'text/csv':
        if isinstance(request_body, str):
            byte_buffer = request_body.encode()
        else:
            byte_buffer = request_body
        val = read_csv_data(source=byte_buffer)
        logging.info(f"Shape of the requested data: '{val.shape}'")
        return val

    return worker.Response(
        response=f"'{request_content_type}' is an unsupported content type.",
        status=http_client.UNSUPPORTED_MEDIA_TYPE
    )

def output_fn(prediction, accept_type):
    """Encodes prediction to accept type.

    The SageMaker Scikit-learn model server invokes this method with the result of prediction and
    serializes this according to the response MIME type. It expects the input to be numpy array and encodes
    to requested response MIME type.

    Parameters
    ----------
    prediction : array-like
        the object returned from predict_fn

    accept_type : str
        the expected MIME type of the response

    Returns
    -------
    : Response obj
        serialized predictions in accept type

    """
    if isinstance(prediction, worker.Response):
        return prediction

    if _is_inverse_label_transform():
        if accept_type == 'text/csv':
            return worker.Response(
                response=encoders.encode(prediction, accept_type),
                status=http_client.OK,
                mimetype=accept_type
            )
        else:
            return worker.Response(
                response=f"Accept type '{accept_type}' is not supported "
                         f"during inverse label transformation.",
                status=413
            )

    if isinstance(prediction, tuple):
        X, y = prediction
    else:
        X, y = _split_features_target(prediction)

    if accept_type == 'application/x-recordio-protobuf':
        return worker.Response(
            response=encoders.array_to_recordio_protobuf(
                _sparsify_if_needed(X).astype('float32'),
                y.astype('float32') if y is not None else y
            ),
            status=http_client.OK,
            mimetype=accept_type
        )

    if accept_type == 'text/csv':
        if y is not None:
            X = np.column_stack(
                (np.ravel(y), X.todense() if sparse.issparse(X) else X)
            )

        return worker.Response(
            response=encoders.encode(X, accept_type),
            status=http_client.OK,
            mimetype=accept_type
        )
    return worker.Response(
        response=f"Accept type '{accept_type}' is not supported.",
        status=http_client.NOT_ACCEPTABLE
    )

def execution_parameters_fn():
    """Return the response for execution-parameters request used by SageMaker Batch transform.

    The SageMaker Scikit-learn model server invokes when execution-parameters endpoint is called.
    For the models generated by AutoML Jobs, returns the MaxPayloadInMB to be 1MB for feature transform
    used during inference and defaults to 6MB otherwise.
    """
    if _is_feature_transform():
        return worker.Response(
            response='{"MaxPayloadInMB":1}',
            status=http_client.OK,
            mimetype="application/json"
        )
    return worker.Response(
        response='{"MaxPayloadInMB":6}',
        status=http_client.OK,
        mimetype="application/json"
    )


from typing import List, Dict, Any, Tuple
from pathlib import Path
import re
import numpy as np
import pandas as pd
from xml.etree.ElementTree import Element  # noqa: S405
import defusedxml.ElementTree as ET
from resistics.time import ChanMetadata
from resistics.calibrate import SensorCalibrationReader, CalibrationData
from resistics.spectra import SpectraMetadata

class SensorCalibration_RSP_RSPX_Base(SensorCalibrationReader):
    """Base class for RSP and RSPX calibration data readers"""

    file_str: str = "Metronix_Coil-----TYPE-$sensor_$chopper-ID-$serial$extension"
    """The file string to search for. Various parameters will be substituted"""

    def _get_path(self, dir_path: Path, metadata: SpectraMetadata, chan: str) -> Path:
        """
        Get the path to the calibration file

        Parameters
        ----------
        dir_path : Path
            The directory  path to look for calibration files
        metadata : SpectraMetadata
            SpectraMetadata with data information
        chan : str
            The channel to calibrate

        Returns
        -------
        Path
            The path to the calibration file
        """
        chan_metadata = metadata.chans_metadata[chan]
        chopper_str = "LF" if chan_metadata.chopper else "HF"
        sensor_str = re.sub("[^0-9]", "", chan_metadata.sensor)
        sensor_str = f"{int(sensor_str):03d}"
        serial_str = f"{int(chan_metadata.serial):06d}"
        file_name = self.file_str.replace("$sensor", sensor_str)
        file_name = file_name.replace("$serial", serial_str)
        file_name = file_name.replace("$chopper", chopper_str)
        file_name = file_name.replace("$extension", self.extension)
        return dir_path / file_name

    def _get_chopper(self, file_path: Path) -> bool:
        """Get whether the calibration is chopper on or off"""
        if "LF" in file_path.stem or "BB" in file_path.stem:
            return True
        return False

class SensorCalibrationRSP(SensorCalibration_RSP_RSPX_Base):
    """
    Reader for RSP calibration files

    RSP data is in units:

    - F [Hz]
    - Magnitude [mv/nT]
    - Phase [deg]

    Data is returned with units:

    - F [Hz]
    - Magnitude [mV/nT]
    - Phase [radians]

    The static gain for RSP files is applied to the magnitude as it is read in
    """

    extension: str = ".RSP"

    def read_calibration_data(
        self, file_path: Path, chan_metadata: ChanMetadata
    ) -> CalibrationData:
        """
        Read data from a RSP calibration file

        Parameters
        ----------
        file_path : Path
            The file path of the calibration file
        chan_metadata : ChanMetadata
            The channel metadata for the channel to be calibrated

        Returns
        -------
        CalibrationData
            The calibration data
        """
        with file_path.open("r") as f:
            lines = f.readlines()
        lines = [x.strip() for x in lines]
        data_dict = self._read_metadata(lines)
        data_dict["chopper"] = self._get_chopper(file_path)
        df = self._read_data(lines)
        df["magnitude"] = df["magnitude"] * data_dict["static_gain"]
        df["phase"] = df["phase"] * (np.pi / 180)
        data_dict["frequency"] = df.index.values.tolist()
        data_dict["magnitude"] = df["magnitude"].values.tolist()
        data_dict["phase"] = df["phase"].values.tolist()
        data_dict["file_path"] = file_path
        return CalibrationData(**data_dict)

    def _read_metadata(self, lines: List[str]) -> Dict[str, Any]:
        """Read the calibration file metadata"""
        serial, sensor = self._get_sensor_details(lines)
        static_gain = self._get_static_gain(lines)
        return {
            "serial": serial,
            "sensor": sensor,
            "static_gain": static_gain,
            "magnitude_unit": "mV/nT",
            "phase_unit": "radians",
        }

    def _get_sensor_details(self, lines: List[str]) -> Tuple[int, str]:
        """Get sensor details from the file"""
        serial: int = 1
        sensor: str = ""
        for line in lines:
            if "induction coil no" in line:
                split1 = line.split(":")[1]
                serial = int(split1.split("-")[0].strip())
            if "SensorType" in line:
                sensor = line.split()[1]
        return serial, sensor

    def _get_static_gain(self, lines: List[str]) -> float:
        static_gain: float = 1.0
        for line in lines:
            if "StaticGain" in line:
                static_gain = float(line.split()[1])
                return static_gain
        return static_gain

    def _read_data(self, lines: List[str]) -> pd.DataFrame:
        """Read data from calibration file"""
        read_from = self._get_read_from(lines)
        data_lines = self._get_data_lines(lines, read_from)
        data = np.array([x.split() for x in data_lines], dtype=np.float32)
        df = pd.DataFrame(data=data, columns=["frequency", "magnitude", "phase"])
        return df.set_index("frequency").sort_index()

    def _get_read_from(self, lines: List[str]) -> int:
        """Get the line number to read from"""
        for idx, line in enumerate(lines):
            if "FREQUENCY" in line:
                return idx + 2
        raise ValueError("Unable to determine location of data in file")

    def _get_data_lines(self, lines: List[str], idx: int) -> List[str]:
        """Get the data lines out of the file"""
        data_lines: List[str] = []
        while idx < len(lines) and lines[idx] != "":
            data_lines.append(lines[idx])
            idx += 1
        return data_lines

class SensorCalibrationRSPX(SensorCalibration_RSP_RSPX_Base):
    """
    Read data from RSPX calibration file

    RSPX data is in units:

    - F [Hz]
    - Magnitude [mv/nT]
    - Phase [deg]

    Data is returned with units:

    - F [Hz]
    - Magnitude [mV/nT]
    - Phase [radians]

    Static gain is applied to the magnitude
    """

    extension: str = ".RSPX"

    def read_calibration_data(
        self, file_path: Path, chan_metadata: ChanMetadata
    ) -> CalibrationData:
        """
        Read RSPX file

        Parameters
        ----------
        file_path : Path
            The file path of the calibration file
        chan_metadata : ChanMetadata
            The channel metadata for the channel to be calibrated

        Returns
        -------
        CalibrationData
            The calibration data
        """
        root = ET.parse(file_path).getroot()
        data_dict = self._read_metadata(root)
        data_dict["chopper"] = self._get_chopper(file_path)
        df = self._read_data(root)
        df["magnitude"] = df["magnitude"] * data_dict["static_gain"]
        df["phase"] = df["phase"] * (np.pi / 180)
        data_dict["frequency"] = df.index.values.tolist()
        data_dict["magnitude"] = df["magnitude"].values.tolist()
        data_dict["phase"] = df["phase"].values.tolist()
        data_dict["file_path"] = file_path
        return CalibrationData(**data_dict)

    def _read_metadata(self, root: Element) -> Dict[str, Any]:
        """Read the calibration file metadata"""
        serial, sensor = self._get_sensor_details(root)
        static_gain = self._get_static_gain(root)
        return {
            "serial": serial,
            "sensor": sensor,
            "static_gain": static_gain,
            "magnitude_unit": "mV/nT",
            "phase_unit": "radians",
        }

    def _get_sensor_details(self, root: Element) -> Tuple[int, str]:
        """Get sensor details"""
        serial: int = 1
        if root.find("SensorId") is not None:
            serial = int(root.find("SensorId").text)
        sensor: str = ""
        if root.find("SensorSpecification") is not None:
            sensor = root.find("SensorSpecification").text
        return serial, sensor

    def _get_static_gain(self, root) -> float:
        """Get the static gain"""
        static_gain: float = 1.0
        if root.find("StaticGain") is not None:
            static_gain = float(root.find("StaticGain").text)
        return static_gain

    def _read_data(self, root: Element) -> pd.DataFrame:
        """Get data in a DataFrame"""
        data = []
        for resp in root.findall("ResponseData"):
            data.append(
                [
                    np.float32(resp.get("Frequency")),
                    np.float32(resp.get("Magnitude")),
                    np.float32(resp.get("Phase")),
                ]
            )
        df = pd.DataFrame(data=data, columns=["frequency", "magnitude", "phase"])
        return df.set_index("frequency").sort_index()


import os
import torch
import torch.nn as nn

from ResNet import DKU_ResNet

def predict(model_checkpoint, dir_data):
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    model = DKU_ResNet(2).to(device)

    state = torch.load(model_checkpoint)
    model.load_state_dict(state)
    model.eval()

    l_utt = []
    test_li = []
    feature_type = 'GDgram'
    from extract_feats import extract_logspec,extract_GDgram
    if feature_type == 'logspec':
        extract_func = extract_logspec
    elif feature_type == 'GDgram':
        extract_func = extract_GDgram
    for r, ds, fs in os.walk(dir_data):
        for f in fs:
            # if os.path.splitext(f)[1] != '.wav': continue
            test_li += [f]
            l_utt.append('/'.join([r, f.replace('\\', '/')]))
    feats = extract_func(l_utt,ret=True)
    out = []
    y_score = []
    res = []
    like=[]
    import numpy as np
    with torch.no_grad():
        for sample in feats:
            sample = np.expand_dims(sample,axis=0)
            sample = torch.from_numpy(sample).to(device)
            output = model(sample)
            pred = output.argmax(dim=1, keepdim=True)
            loglikelihood = nn.LogSoftmax(dim=1)(output).cpu().numpy()
            score = loglikelihood[:, 1] - loglikelihood[:, 0]
            y_score+=[score]
            like+=[loglikelihood]
            out += [output.cpu().numpy()]
            res += [pred.cpu().numpy()]
    for i in range(len(test_li)):
        print(test_li[i], out[i], like[i], y_score[i],res[i])

if __name__ == '__main__':
    model_ckpath = 'results/models/best-eer-ep28-0.724280.pt'
    audio_dir = 'audio'
    predict(model_ckpath, audio_dir)


from astropy.io import fits
from astropy.table import Table
from astropy.time import Time
import astropy.units as u
import os
import numpy as np
from srttools.io import mkdir_p, locations, read_data_fitszilla, \
    get_chan_columns, classify_chan_columns
from srttools.utils import scantype, force_move_file, minmax, median_diff
from srttools.fit import detrend_spectroscopic_data
import warnings
from astropy import log

def default_scan_info_table():
    return Table(names=['scan_id', 'start', 'stop',
                        'ra_min', 'ra_max', 'ra_d',
                        'dec_min', 'dec_max', 'dec_d',
                        'az_min', 'az_max', 'az_d',
                        'el_min', 'el_max', 'el_d',
                        'glon_min', 'glon_max', 'glon_d',
                        'glat_min', 'glat_max', 'glat_d',
                        'is_skydip', 'kind', 'direction'],

                 dtype=[int, float, float,
                        float, float, float, float, float, float,
                        float, float, float, float, float, float,
                        float, float, float, float, float, float,
                        bool, 'S10', 'S5'])

def get_subscan_info(subscan):
    info = default_scan_info_table()
    scan_id = subscan.meta['SubScanID']
    start, stop = minmax(subscan['time'])
    ramin, ramax = minmax(subscan['ra'])
    decmin, decmax = minmax(subscan['dec'])
    azmin, azmax = minmax(subscan['az'])
    elmin, elmax = minmax(subscan['el'])
    is_skydip = subscan.meta['is_skydip']

    d_ra = median_diff(subscan['ra'])
    d_dec = median_diff(subscan['dec'])
    d_az = median_diff(subscan['az'])
    d_el = median_diff(subscan['el'])

    ravar = (ramax - ramin) * np.cos(np.mean((decmin, decmax)))
    decvar = decmax - decmin
    azvar = (azmax - azmin) * np.cos(np.mean((elmin, elmax)))
    elvar = elmax - elmin

    tot_eq = np.sqrt(ravar ** 2 + decvar ** 2)
    tot_hor = np.sqrt(elvar ** 2 + azvar ** 2)
    ravar /= tot_eq
    decvar /= tot_hor

    directions = np.array(['ra', 'dec', 'az', 'el'])
    allvars = np.array([ravar, decvar, azvar, elvar])

    if tot_eq > 2 * tot_hor:
        kind = 'point'
        direction = ''
    else:
        kind = 'line'
        direction = directions[np.argmax(allvars)]

    info.add_row([scan_id, start, stop,
                  ramin, ramax, d_ra, decmin, decmax, d_dec,
                  azmin, azmax, d_az, elmin, elmax, d_el,
                  0, 0, 0, 0, 0, 0, is_skydip, kind, direction])

    return info

def format_direction(direction):
    """
    Examples
    --------
    >>> format_direction('ra')
    'ra'
    >>> format_direction('el')
    'alat'
    >>> format_direction('az')
    'alon'
    """
    lowerdir = direction.lower()
    if lowerdir == 'el':
        return 'alat'
    elif lowerdir == 'az':
        return 'alon'
    return direction

def get_observing_strategy_from_subscan_info(info):
    """Get observing strategy from subscan information."""
    kinds = info['kind']
    skydips = info['is_skydip']
    lines = info[kinds == 'line']
    points = info[kinds == 'point']
    ctype = 'RA/DEC'
    durations = (info['stop'] - info['start']) * 86400

    xspc, yspc = (0, 0)

    zigzag = False

    stype = 'MAP'
    direction = 'Unkn'
    length = 0

    if np.all(skydips):
        stype = 'SKYDIP'
        mode = 'OTF'
        geom = 'LINE'
        direction = 'ALAT'
    elif len(lines) > len(points):
        mode = 'OTF'
        ra_lines = lines[lines['direction'] == 'ra']
        dec_lines = lines[lines['direction'] == 'dec']
        az_lines = lines[lines['direction'] == 'az']
        el_lines = lines[lines['direction'] == 'el']

        directions = np.array(['ra', 'dec', 'az', 'el'])
        nsub = np.array([len(lines[lines['direction'] == d])
                         for d in directions])

        direction = directions[np.argmax(nsub)]
        if direction in ['ra', 'dec']:
            lon_lines, dlon, lat_lines, dlat = ra_lines, 'ra', dec_lines, 'dec'
        elif direction in ['az', 'el']:
            lon_lines, dlon, lat_lines, dlat = az_lines, 'az', el_lines, 'el'
        else:
            raise ValueError('Unknown scan direction')

        ctype = format_direction(dlon) + '/' + format_direction(dlat)
        sample_dist_lon = lon_lines[dlon + '_d']
        sample_dist_lat = lon_lines[dlat + '_d']

        if len(lon_lines) == len(lat_lines):
            geom = 'CROSS'
            zigzag = True
            length = \
                np.median(lon_lines[dlon + '_max'] - lon_lines[dlon + '_min'])
        elif len(lon_lines) > len(lat_lines):
            geom = 'LINE'
            # if we see an inversion of direction, set zigzag to True
            zigzag = np.any(sample_dist_lon[:-1] * sample_dist_lon[1:] < 0)
            length = \
                np.median(lon_lines[dlon + '_max'] - lon_lines[dlon + '_min'])
            direction = format_direction(dlon)
            xspc = 0
            yspc = median_diff(info[dlat + '_min'], sorting=True)
        else:
            geom = 'LINE'
            zigzag = np.any(sample_dist_lat[:-1] * sample_dist_lat[1:] < 0)
            length = \
                np.median(lat_lines[dlat + '_max'] - lat_lines[dlat + '_min'])
            direction = format_direction(dlat)
            yspc = 0
            xspc = median_diff(info[dlon + '_min'], sorting=True)

    else:
        mode = 'RASTER'
        geom = 'SINGLE'

    results = type('results', (), {})()
    results.mode = mode
    results.geom = geom
    results.sep = (xspc, yspc)
    results.zigzag = zigzag
    results.length = length
    results.type = stype
    results.ctype = ctype
    results.stype = stype
    results.scanvel = length / np.median(durations)
    results.direction = direction
    results.nobs = len(info['scan_id'])
    results.scantime = np.median(durations)
    return results

def _copy_hdu_and_adapt_length(hdu, length):
    data = hdu.data
    columns = []
    for col in data.columns:
        newvals = [data[col.name][0]] * length
        newcol = fits.Column(name=col.name, array=newvals,
                             format=col.format)
        columns.append(newcol)
    newhdu = fits.BinTableHDU.from_columns(columns)
    newhdu.header = hdu.header
    return newhdu

keywords_to_reset = [
    '11CD2F', '11CD2I', '11CD2J', '11CD2R', '11CD2S',
    '1CRPX2F', '1CRPX2I', '1CRPX2J', '1CRPX2R', '1CRPX2S', '1CRVL2F',
    '1CRVL2I', '1CRVL2J', '1CRVL2R', '1CRVL2S', '1CTYP2F', '1CTYP2I',
    '1CTYP2J', '1CTYP2R', '1CTYP2S', '1CUNI2F', '1CUNI2I', '1CUNI2J',
    '1CUNI2R', '1CUNI2S', '1SOBS2F', '1SOBS2I', '1SOBS2J', '1SOBS2R',
    '1SOBS2S', '1SPEC2F', '1SPEC2I', '1SPEC2J', '1SPEC2R', '1SPEC2S',
    '1VSOU2R', 'AN', 'ANRX', 'AW', 'AWRX', 'BANDWID', 'BLATOBJ', 'BLONGOBJ',
    'CA', 'CARX', 'DEWCABIN', 'DEWRTMOD', 'DEWUSER', 'DEWZERO', 'DISTANCE',
    'ECCENTR', 'FDELTACA', 'FDELTAIA', 'FDELTAIE', 'FDELTAX', 'FDELTAXT',
    'FDELTAY', 'FDELTAYT', 'FDELTAZ', 'FDELTAZT', 'FDTYPCOD', 'FEBEBAND',
    'FEBEFEED', 'FEGAIN', 'FREQRES', 'FRTHRWHI', 'FRTHRWLO', 'GRPID1',
    'GRPLC1', 'HACA', 'HACA2', 'HACA2RX', 'HACA3', 'HACA3RX', 'HACARX',
    'HASA', 'HASA2', 'HASA2RX', 'HASARX', 'HECA2', 'HECA2RX', 'HECA3',
    'HECA3RX', 'HECE', 'HECE2', 'HECE2RX', 'HECE6', 'HECE6RX', 'HECERX',
    'HESA', 'HESA2', 'HESA2RX', 'HESA3', 'HESA3RX', 'HESA4', 'HESA4RX',
    'HESA5', 'HESA5RX', 'HESARX', 'HESE', 'HESERX', 'HSCA', 'HSCA2',
    'HSCA2RX', 'HSCA5', 'HSCA5RX', 'HSCARX', 'HSSA3', 'HSSA3RX', 'IA', 'IARX',
    'IE', 'IERX', 'INCLINAT', 'LATOBJ', 'LONGASC', 'LONGOBJ', 'LONGSTRN',
    'NFEBE', 'NOPTREFL', 'NPAE', 'NPAERX', 'NPHASES', 'NRX', 'NRXRX', 'NRY',
    'NRYRX', 'NUSEBAND', 'OMEGA', 'OPTPATH', 'ORBEPOCH', 'ORBEQNOX', 'PATLAT',
    'PATLONG', 'PDELTACA', 'PDELTAIA', 'PDELTAIE', 'PERIDATE', 'PERIDIST',
    'REFOFFX', 'REFOFFY', 'REF_ONLN', 'REF_POL', 'RESTFREQ', 'SBSEP',
    'SCANLEN', 'SCANLINE', 'SCANNUM', 'SCANPAR1', 'SCANPAR2', 'SCANROT',
    'SCANRPTS', 'SCANSKEW', 'SCANTIME', 'SCANXSPC', 'SCANXVEL', 'SCANYSPC',
    'SIDEBAND', 'SIG_ONLN', 'SIG_POL', 'SKYFREQ', 'SWTCHMOD', 'TBLANK',
    'TRANSITI', 'TSYNC', 'WCSNM2F', 'WCSNM2I', 'WCSNM2J', 'WCSNM2R',
    'WCSNM2S', 'WOBTHROW', 'WOBUSED']

def pack_data(scan, polar_dict, detrend=False):
    """Pack data into MBFITS-ready format

    Examples
    --------
    >>> scan = {'Feed0_LCP': np.arange(4), 'Feed0_RCP': np.arange(4, 8)}
    >>> polar = {'LCP': 'Feed0_LCP', 'RCP': 'Feed0_RCP'}
    >>> res = pack_data(scan, polar)
    >>> np.allclose(res, [[0, 4], [1, 5], [2, 6], [3, 7]])
    True
    >>> scan = {'Feed0_LCP': np.arange(2), 'Feed0_RCP': np.arange(2, 4),
    ...         'Feed0_Q': np.arange(4, 6), 'Feed0_U': np.arange(6, 8)}
    >>> polar = {'LCP': 'Feed0_LCP', 'RCP': 'Feed0_RCP', 'Q': 'Feed0_Q',
    ...          'U': 'Feed0_U'}
    >>> res = pack_data(scan, polar)
    >>> np.allclose(res, [[0, 2, 4, 6], [1, 3, 5, 7]])
    True
    >>> scan = {'Feed0_LCP': np.ones((2, 4)), 'Feed0_RCP': np.zeros((2, 4))}
    >>> polar = {'LCP': 'Feed0_LCP', 'RCP': 'Feed0_RCP'}
    >>> res = pack_data(scan, polar)
    >>> np.allclose(res, [[[ 1.,  1.,  1.,  1.], [ 0.,  0.,  0.,  0.]],
    ...                   [[ 1.,  1.,  1.,  1.], [ 0.,  0.,  0.,  0.]]])
    True
    """

    polar_list = list(polar_dict.keys())
    if 'LCP' in polar_list:
        data = [scan[polar_dict['LCP']], scan[polar_dict['RCP']]]
        try:
            data.append(scan[polar_dict['Q']])
            data.append(scan[polar_dict['U']])
        except KeyError:
            pass
    else:  # pragma: no cover
        raise ValueError('Polarization kind not implemented yet')

    if detrend:
        new_data = []
        for d in data:
            detr, _ = detrend_spectroscopic_data(0, d, 'als')
            new_data.append(detr)
        data = new_data
    return np.stack(data, axis=1)

def reset_all_keywords(header):
    """Set a specific list of keywords to zero or empty string.

    Examples
    --------
    >>> from astropy.io.fits import Header
    >>> h = Header({'SCANNUM': 5, 'OPTPATH': 'dafafa', 'a': 'blabla'})
    >>> h2 = reset_all_keywords(h)
    >>> h2['SCANNUM']
    0
    >>> h2['OPTPATH']
    ''
    >>> # This is not in the list of keywords to eliminate
    >>> h2['a']
    'blabla'
    """
    import six
    for key in keywords_to_reset:
        if key in header:
            if isinstance(header[key], six.string_types):
                header[key] = ''
            else:
                header[key] = type(header[key])(0)
    return header

class MBFITS_creator():
    def __init__(self, dirname, test=False):
        self.dirname = dirname
        self.test = test
        mkdir_p(dirname)
        curdir = os.path.dirname(__file__)
        datadir = os.path.join(curdir, '..', 'data')
        self.template_dir = os.path.join(datadir, 'mbfits_template')

        self.FEBE = {}

        self.GROUPING = 'GROUPING.fits'
        with fits.open(os.path.join(self.template_dir,
                                    'GROUPING.fits'),
                       memmap=False) as grouping_template:
            grouping_template[1].data = grouping_template[1].data[:1]

            grouping_template.writeto(
                os.path.join(self.dirname, self.GROUPING), overwrite=True)

        self.SCAN = 'SCAN.fits'
        with fits.open(os.path.join(self.template_dir,
                                    'SCAN.fits'),
                       memmap=False) as scan_template:
            scan_template[1].data['FEBE'][0] = 'EMPTY'

            scan_template.writeto(os.path.join(self.dirname, self.SCAN),
                                  overwrite=True)
        self.date_obs = Time.now()
        self.scan_info = default_scan_info_table()
        self.nfeeds = None
        self.ra = 0
        self.dec = 0
        self.site = None
        self.lst = 1e32

    def fill_in_summary(self, summaryfile):
        log.info('Loading {}'.format(summaryfile))
        with fits.open(summaryfile, memmap=False) as hdul:
            header = hdul[0].header
            hdudict = dict(header.items())

        self.ra = np.degrees(hdudict['RightAscension'])
        self.dec = np.degrees(hdudict['Declination'])
        self.restfreq = None
        if 'RESTFREQ1' in hdudict:
            self.resfreq = hdudict['RESTFREQ1']

        try:
            self.date_obs = Time(hdudict['DATE-OBS'])
        except KeyError:
            self.date_obs = Time(hdudict['DATE'])
        try:
            self.obsid = int(hdudict['OBSID'])
        except (KeyError, ValueError):
            self.obsid = 9999

        with fits.open(os.path.join(self.dirname, self.GROUPING),
                       memmap=False) as grouphdul:
            groupheader = grouphdul[0].header
            groupdict = dict(groupheader.items())
            for key in hdudict.keys():
                if key in groupdict:
                    groupheader[key] = hdudict[key]
            groupheader['RA'] = self.ra
            groupheader['DEC'] = self.dec
            groupheader['DATE-OBS'] = self.date_obs.value
            groupheader['MJD-OBS'] = self.date_obs.mjd
            groupheader['SCANNUM'] = self.obsid
            grouphdul.writeto('tmp.fits', overwrite=True)

        force_move_file('tmp.fits', os.path.join(self.dirname, self.GROUPING))

        with fits.open(os.path.join(self.dirname, self.SCAN),
                       memmap=False) as scanhdul:
            scanheader = reset_all_keywords(scanhdul[1].header)
            scandict = dict(scanheader.items())
            for key in hdudict.keys():
                if key[:5] in ['NAXIS', 'PGCOU', 'GCOUN']:
                    continue
                if key in scandict:
                    scanheader[key] = hdudict[key]
            # Todo: update with correct keywords
            scanheader['DATE-OBS'] = self.date_obs.value
            scanheader['MJD'] = self.date_obs.mjd
            scanheader['SCANNUM'] = self.obsid
            scanhdul.writeto('tmp.fits', overwrite=True)

        force_move_file('tmp.fits', os.path.join(self.dirname, self.SCAN))

    def add_subscan(self, scanfile, detrend=False):
        log.info('Loading {}'.format(scanfile))

        subscan = read_data_fitszilla(scanfile)
        subscan_info = get_subscan_info(subscan)

        self.scan_info.add_row(subscan_info[0])

        time = Time(subscan['time'] * u.day, scale='utc', format='mjd')
        if self.date_obs.mjd > time[0].mjd:
            self.date_obs = time[0]
        if self.site is None:
            self.site = subscan.meta['site']

        chans = get_chan_columns(subscan)

        combinations = classify_chan_columns(chans)

        if self.nfeeds is None:
            self.nfeeds = len(combinations.keys())

        for feed in combinations:
            felabel = subscan.meta['receiver'] + '{}'.format(feed)
            febe = felabel + '-' + subscan.meta['backend']

            datapar = os.path.join(self.template_dir, '1',
                                   'FLASH460L-XFFTS-DATAPAR.fits')
            with fits.open(datapar, memmap=False) as subs_par_template:
                n = len(subscan)
                # ------------- Update DATAPAR --------------
                subs_par_template[1] = \
                    _copy_hdu_and_adapt_length(subs_par_template[1], n)

                newtable = Table(subs_par_template[1].data)
                newtable['MJD'] = subscan['time']
                newtable['LST'][:] = \
                    time.sidereal_time('apparent',
                                       locations[subscan.meta['site']].lon
                                       ).value
                if newtable['LST'][0] < self.lst:
                    self.lst = newtable['LST'][0]
                newtable['INTEGTIM'][:] = \
                    subscan['Feed0_LCP'].meta['sample_rate']
                newtable['RA'] = subscan['ra'].to(u.deg)
                newtable['DEC'] = subscan['dec'].to(u.deg)
                newtable['AZIMUTH'] = subscan['az'].to(u.deg)
                newtable['ELEVATIO'] = subscan['el'].to(u.deg)
                _, direction = scantype(subscan['ra'], subscan['dec'],
                                        el=subscan['el'], az=subscan['az'])

                direction_cut = \
                    direction.replace('<', '').replace('>', '').lower()
                if direction_cut in ['ra', 'dec']:
                    baslon = subscan['ra'].to(u.deg)
                    baslat = subscan['dec'].to(u.deg)

                    yoff = baslat.value - self.dec
                    # GLS projection
                    xoff = \
                        (baslon.value - self.ra)
                    newtable['LONGOFF'] = xoff * np.cos(np.radians(self.dec))
                    newtable['LATOFF'] = yoff
                elif direction_cut in ['el', 'az']:
                    warnings.warn('AltAz projection not implemented properly')
                    baslon, baslat = \
                        subscan['az'].to(u.deg), subscan['el'].to(u.deg)
                    newtable['LONGOFF'] = 0 * u.deg
                    newtable['LATOFF'] = 0 * u.deg
                else:
                    raise ValueError('Unknown coordinates')

                newtable['CBASLONG'] = baslon
                newtable['CBASLAT'] = baslat
                newtable['BASLONG'] = baslon
                newtable['BASLAT'] = baslat

                newhdu = fits.table_to_hdu(newtable)
                subs_par_template[1].data = newhdu.data
                subs_par_template[1].header['DATE-OBS'] = \
                    time[0].fits.replace('(UTC)', '')
                subs_par_template[1].header['LST'] = newtable['LST'][0]
                subs_par_template[1].header['FEBE'] = febe
                subs_par_template[1].header['SCANDIR'] = \
                    format_direction(direction_cut).upper()
                subs_par_template[1].header['SCANNUM'] = self.obsid

                outdir = str(subscan.meta['SubScanID'])
                mkdir_p(os.path.join(self.dirname, outdir))
                new_datapar = os.path.join(outdir,
                                           febe + '-DATAPAR.fits')
                subs_par_template.writeto('tmp.fits', overwrite=True)

            force_move_file('tmp.fits',
                            os.path.join(self.dirname, new_datapar))

            arraydata = os.path.join(self.template_dir, '1',
                                     'FLASH460L-XFFTS-ARRAYDATA-1.fits')

            new_arraydata_rows = []
            bands = list(combinations[feed].keys())
            for baseband in combinations[feed]:
                nbands = np.max(bands)
                ch = list(combinations[feed][baseband].values())[0]

                packed_data = pack_data(subscan, combinations[feed][baseband],
                                        detrend=detrend)
                # ------------- Update ARRAYDATA -------------
                with fits.open(arraydata, memmap=False) as subs_template:
                    subs_template[1] = \
                        _copy_hdu_and_adapt_length(subs_template[1], n)

                    new_header = \
                        reset_all_keywords(subs_template[1].header)

                    new_header['SCANNUM'] = self.obsid
                    new_header['SUBSNUM'] = subscan.meta['SubScanID']
                    new_header['DATE-OBS'] = self.date_obs.fits
                    new_header['FEBE'] = febe
                    new_header['BASEBAND'] = baseband
                    new_header['NUSEBAND'] = nbands
                    new_header['CHANNELS'] = subscan.meta['channels']
                    new_header['SKYFREQ'] = \
                        subscan[ch].meta['frequency'].to('Hz').value
                    if self.restfreq is not None:
                        new_header['RESTFREQ'] = self.restfreq
                    else:
                        new_header['RESTFREQ'] = new_header['SKYFREQ']
                    bandwidth = subscan[ch].meta['bandwidth'].to('Hz').value
                    new_header['BANDWID'] = bandwidth

                    new_header['FREQRES'] = bandwidth / new_header['CHANNELS']

                    # Todo: check sideband
                    new_header['SIDEBAND'] = 'USB'
                    # Todo: check all these strange keywords. These are
                    # probably NOT the rest frequencies!
                    new_header['1CRVL2F'] = new_header['RESTFREQ']
                    new_header['1CRVL2S'] = new_header['RESTFREQ']
                    for i in ['1CRPX2S', '1CRPX2R', '1CRPX2F', '1CRPX2J']:
                        new_header[i] = (new_header['CHANNELS'] + 1) // 2

                    subs_template[1].header = new_header
                    newtable = Table(subs_template[1].data)
                    newtable['MJD'] = subscan['time']
                    newtable['DATA'] = packed_data
                    newhdu = fits.table_to_hdu(newtable)
                    subs_template[1].data = newhdu.data

                    subname = febe + '-ARRAYDATA-{}.fits'.format(baseband)
                    new_sub = \
                        os.path.join(outdir, subname)
                    subs_template.writeto('tmp.fits', overwrite=True)

                    new_arraydata_rows.append([2, new_sub, 'URL',
                                               'ARRAYDATA-MBFITS',
                                               subscan.meta['SubScanID'], febe,
                                               baseband])

                force_move_file('tmp.fits',
                                os.path.join(self.dirname, new_sub))

            # Finally, update GROUPING file
            with fits.open(os.path.join(self.dirname,
                                        self.GROUPING),
                           memmap=False) as grouping:
                newtable = Table(grouping[1].data)
                if febe not in self.FEBE:

                    nfebe = len(list(self.FEBE.keys()))
                    new_febe = self.add_febe(febe, combinations, feed,
                                             subscan[ch].meta,
                                             bands=bands)

                    grouping[0].header['FEBE{}'.format(nfebe)] = febe
                    grouping[0].header['FREQ{}'.format(nfebe)] = \
                        subscan[ch].meta['frequency'].to('Hz').value
                    grouping[0].header['BWID{}'.format(nfebe)] = \
                        subscan[ch].meta['bandwidth'].to('Hz').value
                    grouping[0].header['LINE{}'.format(nfebe)] = ''

                    newtable.add_row([2, new_febe, 'URL', 'FEBEPAR-MBFITS',
                                      -999, febe, -999])
                    self.FEBE[febe] = new_febe

                newtable.add_row([2, new_datapar, 'URL', 'DATAPAR-MBFITS',
                                  -999, febe, -999])

                for row in new_arraydata_rows:
                    newtable.add_row(row)
                new_hdu = fits.table_to_hdu(newtable)
                grouping[1].data = new_hdu.data
                grouping[0].header['INSTRUME'] = subscan[ch].meta['backend']
                grouping[0].header['TELESCOP'] = self.site

                grouping.writeto('tmp.fits', overwrite=True)

            force_move_file('tmp.fits',
                            os.path.join(self.dirname, self.GROUPING))

            if self.test:
                break

    def add_febe(self, febe, feed_info, feed, meta, bands=None):
        if bands is None:
            bands = [1]
        polar = 'N'
        polar_code = polar[0]

        febe_name = febe + '-FEBEPAR.fits'

        with fits.open(
                os.path.join(self.template_dir,
                             'FLASH460L-XFFTS-FEBEPAR.fits'),
                memmap=False) as febe_template:

            febe_template[1].header = \
                reset_all_keywords(febe_template[1].header)

            febedata = Table(febe_template[1].data)
            # FEBEFEED stores the total number of feeds for the receiver in
            # use.  A receiver outputting two polarisations counts as two
            # feeds.  For an array, count the total no.  of pixels, even if
            # not all in use.
            febedata['USEBAND'] = np.array([bands])
            febedata['NUSEFEED'] = np.array([[2]])
            febedata['USEFEED'] = \
                np.array([[feed * 2 + 1, feed * 2 + 2,
                           feed * 2 + 1, feed * 2 + 2]])
            febedata['BESECTS'] = np.array([[0]])
            febedata['FEEDTYPE'] = np.array([[1, 2, 3, 4]])
            febedata['POLTY'][:] = np.array([polar_code])
            febedata['POLA'][:] = np.array([[0., 0.]])
            new_hdu = fits.table_to_hdu(febedata)

            febe_template[1].data = new_hdu.data
            # TODO: fill in the information given in the subscan[ch]

            new_febe = os.path.join(self.dirname, febe_name)

            febe_template[1].header['DATE-OBS'] = self.date_obs.fits
            febe_template[1].header['FEBE'] = febe
            febe_template[1].header['FEBEFEED'] = self.nfeeds * 2
            febe_template[1].header['NUSEBAND'] = max(bands)
            febe_template[1].header['NPHASES'] = 1
            febe_template[1].header['SWTCHMOD'] = 'NONE'
            febe_template[1].header['SCANNUM'] = self.obsid

            if 'Q' in feed_info[feed][bands[0]].keys():
                febe_template[1].header['FDTYPCOD'] = '1:L, 2:R, 3:Q, 4:U'
            else:
                febe_template[1].header['FDTYPCOD'] = '1:L, 2:R'
            febe_template.writeto('tmp.fits', overwrite=True)
        force_move_file('tmp.fits', new_febe)

        with fits.open(os.path.join(self.dirname, self.SCAN),
                       memmap=False) as scan:
            newtable = Table(scan[1].data)

            if newtable['FEBE'][0].strip() == 'EMPTY':
                newtable['FEBE'][0] = febe
            else:
                newtable.add_row([febe])

            new_hdu = fits.table_to_hdu(newtable)
            scan[1].data = new_hdu.data
            scanheader = scan[1].header
            scanheader['SITELONG'] = np.degrees(meta['SiteLongitude'])
            scanheader['SITELAT'] = np.degrees(meta['SiteLatitude'])
            scanheader['SITEELEV'] = meta['SiteHeight']
            diameter = 64. if meta['site'].lower().strip() == 'srt' else 32.
            scanheader['DIAMETER'] = diameter
            scanheader['PROJID'] = meta['Project_Name']

            scan.writeto('tmp.fits', overwrite=True)
        force_move_file('tmp.fits', os.path.join(self.dirname, self.SCAN))

        return febe_name

    def update_scan_info(self):
        info = \
            get_observing_strategy_from_subscan_info(self.scan_info)

        with fits.open(os.path.join(self.dirname, self.SCAN),
                       memmap=False) as scanhdul:
            scanheader = scanhdul[1].header
            # Todo: update with correct keywords
            scanheader['CTYPE'] = info.ctype
            scanheader['CTYPE1'] = 'RA---GLS'
            scanheader['CTYPE2'] = 'DEC--GLS'
            scanheader['CRVAL1'] = self.ra
            scanheader['CRVAL2'] = self.dec
            scanheader['BLONGOBJ'] = self.ra
            scanheader['BLATOBJ'] = self.dec
            scanheader['LONGOBJ'] = self.ra if not info.ctype[0] == 'A' else 0
            scanheader['LATOBJ'] = self.dec if not info.ctype[0] == 'A' else 0
            scanheader['EQUINOX'] = 2000.
            scanheader['GRPLC1'] = 'GROUPING.fits'
            scanheader['LST'] = self.lst
            scanheader['LATPOLE'] = 90.
            scanheader['LONPOLE'] = 0.
            scanheader['PATLONG'] = 0
            scanheader['MOVEFRAM'] = False
            if info.ctype == 'ALON/ALAT':
                scanheader['WCSNAME'] = 'Absolute horizontal'
            scanheader['SCANTYPE'] = info.stype.upper()
            scanheader['SCANDIR'] = info.direction.upper()
            scanheader['SCANXVEL'] = info.scanvel
            scanheader['SCANTIME'] = info.scantime
            scanheader['SCANMODE'] = info.mode.upper()
            scanheader['SCANGEOM'] = info.geom.upper()
            scanheader['SCANLINE'] = 1
            scanheader['SCANLEN'] = np.degrees(info.length)

            scanheader['SCANYSPC'] = np.degrees(info.sep[1])
            scanheader['SCANXSPC'] = np.degrees(info.sep[0])
            scanheader['SCANPAR1'] = -999
            scanheader['SCANPAR2'] = -999
            scanheader['ZIGZAG'] = info.zigzag
            scanheader['PHASE1'] = 'sig'
            scanheader['PHASE2'] = 'sig'
            scanheader['NOBS'] = info.nobs
            scanheader['NSUBS'] = info.nobs
            scanheader['WOBCYCLE'] = 0.
            scanheader['WOBDIR'] = 'NONE'
            scanheader['WOBMODE'] = 'NONE'
            scanheader['WOBPATT'] = 'NONE'
            scanhdul.writeto('tmp.fits', overwrite=True)
        force_move_file('tmp.fits', os.path.join(self.dirname, self.SCAN))

    def wrap_up_file(self):
        import copy
        prihdu = fits.PrimaryHDU()
        with fits.open(os.path.join(self.dirname, self.GROUPING),
                       memmap=False) as grouhdl:
            prihdu.header = copy.deepcopy(grouhdl[0].header)
            file_list = list(zip(grouhdl[1].data['MEMBER_LOCATION'],
                                 grouhdl[1].data['EXTNAME'],
                                 grouhdl[1].data['FEBE']))

        hdulists = {}
        for febe in self.FEBE.keys():
            hdulists[febe] = fits.HDUList([prihdu])

            with fits.open(os.path.join(self.dirname, self.SCAN),
                           memmap=False) as scanhdul:
                scanhdul[1].data['FEBE'] = [febe]
                newhdu = type(scanhdul[1])()
                newhdu.data = scanhdul[1].data
                newhdu.header = scanhdul[1].header
                hdulists[febe].append(newhdu)

        for fname, ext, febe in file_list:
            if febe == '':
                continue
            with fits.open(os.path.join(self.dirname, fname),
                           memmap=False) as hl:
                newhdu = type(hl[ext])()
                newhdu.data = hl[ext].data
                newhdu.header = hl[ext].header
                hdulists[febe].append(newhdu)

        fnames = {}
        for febe, hdulist in hdulists.items():
            fname = self.dirname + '.' + febe + '.fits'
            hdulist.writeto(fname, overwrite=True)
            hdulist.close()
            fnames[febe] = fname

        return fnames

import numpy,import astropy,from astropy


import cv2 as cv
import numpy as np
import os
import pyautogui
from time import time, sleep
from screencapture import WindowCapture
from computervision import ComputerVision
from hsvfilter import HsvFilter
from threading import Thread
from tkinter import *
import entry
import win32api

#TKINTER ----------------------------------------------------------------------------------
def myClick():
    global inp
    inp = e.get()

def myClick2():
    global SLEEP_TIMER
    SLEEP_TIMER = e.get()

root = Tk()
root.title('Bot by Nikolas Achatz')
e = Entry(root, width= 60)
e.pack()
inp = None
SLEEP_TIMER = 0.2
Button(root, text='Enter client name (Leave empty for desktop view)', command=myClick).pack()
Button(root, text='Sleep Timer for actions', command=myClick2).pack()
Button(root, text="Send", command=root.destroy).pack()
root.mainloop()

#TKINTER END ----------------------------------------------------------------------------------

print("Client Window " + inp)
print("Ticker " + SLEEP_TIMER)

inp = "League of Legends (TM) Client"

#This is the window we want to capture and work with.
wincap = entry.attemptDataSend(inp)

#THIS IS WHAT WE WANT TO FIND in our window
if inp == "Runescape":
    cv_find = ComputerVision('runescapet.JPG')
elif inp == "League":
    cv_find = ComputerVision('minion.JPG')
else:
    cv_find = ComputerVision('runescapet.JPG')

# else:
#     win32api.MessageBox(0, 'You did not attach to a good client, we have nothing to search for.', 'Error')
#     exit(0)

#This is to find our HSV values for new objects.
#cv_find.init_control_gui()

#Determines if we are already doing a task.
is_in_action = False

#This is our preset HSV values for processed searches
if inp == "Runescape":
    hsv_filter = HsvFilter(0, 137, 0, 20, 221, 96, 52, 0, 0, 0)
else:
    hsv_filter = HsvFilter(0,74,33,8,255,255,0,0,0,0)

# Using a different thread we can run the bot actions in parallel
def bot_action(rectangles):
    if len(rectangles) > 0:
        targets = cv_find.get_click_points(rectangles)
        target = wincap.get_screen_position(targets[0])
        pyautogui.moveTo(x=target[0], y=target[1])
        pyautogui.click(button='right')
        sleep(int(SLEEP_TIMER))
        global is_in_action
        is_in_action = False

#Used for FPS monitoring. Not needed
#loop_time = time()

while(True):

    # Continuously grabs the image on the screen using windows OS
    screenshot = wincap.get_screenshot()

    # This is our normal image being modified by the HSV values set above
    processed_img = cv_find.apply_hsv_filter(screenshot, hsv_filter)

    # FINDS INSTANCES OF THE OBJECT WE WANT TO FIND, second value is adjusted to change acceptance
    rectangles = cv_find.find(processed_img, 0.28)

    # Draw the rectangles of positive matches
    output_img = cv_find.draw_rectangles(screenshot, rectangles)

    cv.imshow("Processed", processed_img)
    cv.imshow("Matches", screenshot)

    # Bot action thread, if we aren't in a thread, we can send a new task out.
    if not is_in_action:
        is_in_action = True
        t = Thread(target=bot_action, args=(rectangles,))
        t.start()

    #print('FPS {} '.format(1 /(time() - loop_time)))
    #loop_time = time()

    key = cv.waitKey(1)
    if key == ord('q'):
        cv.destroyAllWindows()
        break

print("done")