<reponame>christiankuhl/VortexDynamics
"""
This module implements the N-vortex dynamics for an ellipse.
Insert code here as if it were inside a class definition statement.
"""

import numpy as np
from scipy.special import ellipk
from scipy.optimize import newton
from mpmath import ellipfun
from cmath import asin
from utils import vectorized

def __init__(self, a):
    """
    Instantiates an ellipse with half axes a and 1.
    """
    self.a = a
    self.xlim = (-a, a)
    self.ylim = (-1, 1)
    self._k = self.solve_for_k()
    self._K = ellipk(self._k)
    print(self._k)

def JGradG(self, z1, z2):
    z1, z2 = complex(*z1), complex(*z2)
    fz, fw, fprimez = self.F(z1), self.F(z2), self.Fprime(z1)
    igradG = -1.j / (2 * np.pi) * (fprimez * np.conj(fw) * (1 - fz * np.conj(fw))
                                    / np.abs((1 - fz * np.conj(fw)))**2
                                    - (fprimez * (fz - fw)) / np.abs((fz - fw))**2)
    return np.array([np.real(igradG), np.imag(igradG)])

@vectorized
def JGradh(self, z):
    # z = complex(*z)
    f, fprime, f2prime = self.F(z), self.Fprime(z), self.F2prime(z)
    # f, fprime, f2prime = z, 1, 0
    jhprime = 1.j / (2 * np.pi) * (2 * f * fprime / (1 - np.abs(f)**2)
                                     + fprime * f2prime / (np.abs(fprime)**2))
    return np.array([np.real(jhprime), np.imag(jhprime)])

def plot_me(self, t):
    return np.array([self.a*np.cos(2*np.pi*t), np.sin(2*np.pi*t)])

# Main interface implemented above, lots of helper functions below this point

def solve_for_k(self):
    def rhs(k):
        c = 2/np.pi * np.arcsinh(2 * self.a / (self.a**2 - 1))
        return ellipk(np.sqrt(1 - k**2)) / ellipk(k) - c
    return newton(rhs, .5)

@vectorized
def W(self, z):
    return 2 * self._K / np.pi * asin(z / np.sqrt(self.a**2 - 1))

def h(self, z):
    return 1/(2*np.pi)*np.log((1-np.abs(self.F(z))**2)/np.abs(self.Fprime(z)))

def F(self, z):
    return np.sqrt(self._k) * self.snW(z)

def Fprime(self, z):
    cn, dn = self.cnW(z), self.dnW(z)
    return (2 * np.sqrt(self._k) * self._K / (np.pi * np.sqrt(self.a**2 - 1 - z**2))
            * cn * dn)

def F2prime(self, z):
    sn, cn, dn = self.snW(z), self.cnW(z), self.dnW(z)
    return 2*np.sqrt(self._k)*self._K/(np.pi*(self.a**2-1-z**2)**(3/2))*cn*dn*z-4*np.sqrt(self._k)*self._K**2/(np.pi**2*(self.a**2-1-z**2))*sn*(dn**2+self._k**2*cn**2)

@vectorized
def sn(self, z):
    return ellipfun('sn', z, self._k**2).__complex__()

@vectorized
def cn(self, z):
    return ellipfun('cn', z, self._k**2).__complex__()

@vectorized
def dn(self, z):
    return ellipfun('dn', z, self._k**2).__complex__()

@vectorized
def snW(self, z):
    return self.sn(self.W(z))

@vectorized
def cnW(self, z):
    return self.cn(self.W(z))

@vectorized
def dnW(self, z):
    return self.dn(self.W(z))


import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.lines as lines
import matplotlib.gridspec as gridspec
import scipy.misc
from scipy import special
import paratf

plt.rcParams["font.family"] = "serif"
plt.rcParams["font.serif"] = "Times"

inches = 2
rows = 1
cols = 3
widths = [1,1,1]
heights = [1]
hspace = 0.05
wspace = 0.15

f = plt.figure(figsize=(inches*(np.sum(widths) + widths[0]*(wspace)*(cols - 1)),
                        inches*(np.sum(heights))))

outer_grid = gridspec.GridSpec(ncols=cols, nrows=rows,
                               width_ratios=widths, height_ratios=heights,
                               hspace=hspace, wspace=wspace)

def an(krda, n):
    return (n/np.pi)*((special.jn(n, krda)/(krda))**2)

def hdet(krd, theta, alpha):
    return an(krd*alpha, 1)*np.sin(theta)**2 + 0.5*an(krd*alpha, 2)*(alpha**2)*(np.cos(theta)**2)

alphas = [0.25, 0.5, 0.75]
alpha_strs = ['0.25', '0.5', '0.75']
for col in range(cols):
    alpha = alphas[col]
    ax = f.add_subplot(outer_grid[0, col])
    ax.set_xlim([0,2])
    ax.set_ylim([0,1])

    thetas = [0, np.pi/8, np.pi/4, 3*np.pi/8, np.pi/2]
    theta_strs = ['0', '$\pi/8$', '$\pi/4$', '$3\pi/8$', '$\pi/2$']
    for theta in thetas:
        nucr = np.linspace(1e-3, 2, 100)
        ax.plot(nucr, paratf.dpsf(nucr, theta, NA=alpha)/paratf.dpsf(1e-3, np.pi/2, NA=alpha), '-', c=[1 - theta/(np.pi/2), 0, 0], lw=0.5)

    # Monopole
    ax.plot(nucr, (paratf.jinc(0, nucr)/paratf.jinc(0, 0))**2, '-', c=[0, 1, 0], lw=1.5, zorder=0, alpha=1)

    ax.tick_params(direction='out', left=True, right=False, width=0.75)
    ax.yaxis.set_ticks([0, 0.25, 0.5, 0.75, 1.0])
    if col == 0:
        ax.set_yticklabels(['0','0.25','0.5','0.75','1'])
    else:
        ax.set_yticklabels([])
    ax.xaxis.set_ticks([0, 0.5, 1, 1.5, 2])
    ax.set_xticklabels(['0','0.5','1','1.5','2'])

    ax.text(0.5, 1.08, 'NA$/n_o ='+alpha_strs[col]+'$', ha='center', va='center', rotation=0, transform=ax.transAxes)
    if col == 0:
        ax.text(-0.4, 0.5, 'Renormalized\n dipole point spread function \n $h(r, \\vartheta)/h(0, \\pi/2)$', ha='center', va='center', rotation=90, transform=ax.transAxes)
        ax.text(0.5, -0.25, 'Scaled distance $\\nu_c r$', ha='center', va='center', transform=ax.transAxes)
        for i, theta_str in enumerate(theta_strs):
            if i == 4:
                theta_str = '$\\vartheta =' + theta_str[1:-1] + '$'
            ax.text(0.75, 0.4 + 0.1*i, theta_str, ha='right', va='center', transform=ax.transAxes)
            l1 = lines.Line2D([0.8,0.9], 2*[0.4 + 0.1*i], c=[1 - thetas[i]/(np.pi/2), 0, 0], lw=0.5, transform=ax.transAxes)
            f.lines.extend([l1])
        ax.text(0.75, 0.8 + 0.1*1, 'Monopole', ha='right', va='center', transform=ax.transAxes)
        l1 = lines.Line2D([0.8,0.9], 2*[0.8 + 0.1*1], c=[0, 1, 0], lw=1.5, transform=ax.transAxes)
        f.lines.extend([l1])

    if col != 0:
        ax.text(0.5, -0.25, '$\\nu_c r$', ha='center', va='center', transform=ax.transAxes)
f.savefig('dpsf.pdf', dpi=500, bbox_inches='tight')


"""
Project: RadarBook
File: rectangular_plate.py
Created by: <NAME>
On: 10/30/2018
Created with: PyCharm

Copyright (C) 2019 Artech House (<EMAIL>)
This file is part of Introduction to Radar Using Python and MATLAB
and can not be copied and/or distributed without the express permission of Artech House.
"""
from numpy import radians, sin, cos, sinc
from scipy.constants import c, pi

def radar_cross_section(frequency, width, length, incident_theta, observation_theta, observation_phi):
    """
    Calculate the bistatic radar cross section for a rectangular plate.
    :param frequency: The frequency of the incident energy (Hz).
    :param width: The width of the plate (m).
    :param length: The length of the plate (m).
    :param incident_theta: The incident angle theta (deg).
    :param observation_theta: The observation angle theta (deg).
    :param observation_phi: The observation angle phi (deg).
    :return: The bistatic radar cross section (m^2).
    """
    # Wavelength
    wavelength = c / frequency

    theta_i = radians(incident_theta)
    theta_o = radians(observation_theta)

    phi_o = radians(observation_phi)

    x = width / wavelength * sin(theta_o) * cos(phi_o)
    y = length / wavelength * (sin(theta_o) * sin(phi_o) - sin(theta_i))

    rcs_tm = 4.0 * pi * (length * width / wavelength) ** 2 * (cos(theta_i) ** 2 * (cos(theta_o) ** 2 * cos(phi_o) ** 2
                                                              + sin(phi_o) ** 2)) * sinc(x) ** 2 * sinc(y) ** 2

    rcs_te = 4.0 * pi * (length * width / wavelength) ** 2 * (cos(theta_o) ** 2 * sin(phi_o) ** 2
                                                              + cos(phi_o) ** 2) * sinc(x) ** 2 * sinc(y) ** 2

    return rcs_tm, rcs_te


import re
import os
import sys
import pandas as pd
from lxml import etree
from scipy import stats
import gzip
from sqlalchemy import create_engine

class Polite():
    """
    MALLET parameters used: 'output-topic-keys', 'output-doc-topics',
    'word-topic-counts-file', 'topic-word-weights-file',
    'xml-topic-report', 'xml-topic-phrase-report',
    'diagnostics-file', 'output-state'
    """

    class TableDef():
        def __init__(self, index=[], cols=[]):
            self.cols = cols
            self.index = index

    schema = dict(
        DOC = TableDef(['doc_id']),
        DOCTOPIC_NARROW = TableDef(['doc_id', 'topic_id']),
        DOCTOPIC = TableDef(['doc_id']),
        DOCWORD = TableDef(['doc_id', 'word_id']),
        PHRASE = TableDef(['phrase_str']),
        TOPIC = TableDef(['topic_id']),
        TOPICPHRASE = TableDef(['topic_id', 'topic_phrase']),
        TOPICWORD_DIAGS = TableDef(['topic_id', 'word_id']),
        TOPICWORD_NARROW = TableDef(['word_id', 'topic_id']),
        TOPICWORD_WEIGHTS = TableDef(['topic_id', 'word_str']),
        TOPICWORD = TableDef(['word_id']),
        VOCAB = TableDef(['word_id'])
    )

    def __init__(self, config_file, tables_dir='./', save_mode='csv'):
        """Initialize MALLET with trial name"""
        self.config_file = config_file
        self.tables_dir = tables_dir
        self._convert_config_file()
        self.save_mode = save_mode

        if self.save_mode == 'sql':
            engine = create_engine(f'sqlite:///{self.tables_dir}model.db', echo=True)
            self.db = engine.connect()

    def __del__(self):
        if self.save_mode == 'sql':
            self.db.close()

    def save_table(self, df, table_name):
        self.schema[table_name].cols = df.columns
        if self.save_mode == 'sql':
            df.to_sql(table_name, self.db, if_exists='replace', index=True)
        elif self.save_mode == 'csv':
            df.to_csv(self.tables_dir + f'{table_name}.csv')

    def get_table(self, table_name):
        index_cols = self.schema[table_name].index
        if self.save_mode == 'sql':
            df = pd.read_sql_table(table_name, self.db,
            index_col = index_cols)
        elif self.save_mode == 'csv':
            df = pd.read_csv(self.tables_dir + f'{table_name}.csv',
                index_col=index_cols)
        else:
            raise ValueError("No save method!")
        return df

    def _convert_config_file(self):
        """Converts the MALLLET config file into a Python dictionary."""
        self.config = {}
        with open(self.config_file, 'r') as cfg:
            for line in cfg.readlines():
                if not re.match(r'^#', line):
                    a, b = line.split()
                    b = b.strip()
                    if re.match(r'^\d+$', b):
                        b = int(b)
                    elif re.match(r'^\d+\.\d*$', b):
                        b = float(b)
                    elif re.match(r'^TRUE$', b, flags=re.IGNORECASE):
                        b = True
                    elif re.match(r'^FALSE$', b, flags=re.IGNORECASE):
                        b = False
                    self.config[a] = b

        # config = pd.DataFrame(self.config)

    def get_source_file(self, src_file_key):
        src_file = self.config[src_file_key]
        if not os.path.isfile(src_file):
            print(f"File {src_file} for {src_file_key} does not exist. Try running MALLET first.")
            sys.exit(1)
        else:
            return src_file

    def import_table_state(self):
        """Import the state file into docword table"""
        src_file = self.get_source_file('output-state')
        with gzip.open(src_file, 'rb') as f:
            docword = pd.DataFrame(
                [line.split() for line in f.readlines()[3:]],
                columns=['doc_id', 'src', 'word_pos', 'word_id', 'word_str', 'topic_id'])
        docword = docword[['doc_id', 'word_id', 'word_pos', 'topic_id']]
        docword = docword.astype('int')
        docword = docword.set_index(['doc_id', 'word_id'])

        # SAVE
        self.save_table(docword, 'DOCWORD')

    def import_table_topic(self):
        """Import data into topic table"""
        src_file = self.get_source_file('output-topic-keys')
        topic = pd.read_csv(src_file, sep='\t', header=None, index_col='topic_id',
            names=['topic_id', 'topic_alpha', 'topic_words'])
        topic['topic_alpha_zscore'] = stats.zscore(topic.topic_alpha)

        # SAVE
        self.save_table(topic, 'TOPIC')

    def import_tables_topicword_and_word(self):
        """Import data into topicword and word tables"""
        src_file = self.get_source_file('word-topic-counts-file')
        WORD = []
        TOPICWORD = []
        with open(src_file, 'r') as src:
            for line in src.readlines():
                row = line.strip().split()
                word_id, word_str = row[0:2]
                WORD.append((int(word_id), word_str))
                for item in row[2:]:
                    topic_id, word_count = item.split(':')
                    TOPICWORD.append((int(word_id), int(topic_id), int(word_count)))

        # May use schema for indexes
        word = pd.DataFrame(WORD, columns=['word_id', 'word_str']).set_index('word_id')
        topicword = pd.DataFrame(TOPICWORD, columns=['word_id', 'topic_id', 'word_count'])\
            .set_index(['word_id', 'topic_id'])
        topicword_wide = topicword.unstack(fill_value=0)
        topicword_wide.columns = topicword_wide.columns.droplevel(0)
        topicword_wide = topicword_wide / topicword_wide.sum()

        src_file2 = self.get_source_file('topic-word-weights-file')
        topicword_w = pd.read_csv(src_file2,  sep='\t', names=['topic_id','word_str','word_wgt'])\
            .set_index(['topic_id','word_str'])

        # COMBINE TOPICWORD_NARROW AND TOPICWORD_WEIGHTS
        # Note that word weights are just smoothed values
        # So we really should only save the smoothing beta parameter, e.g. .01
        # Get beta from self.config['beta']
        # Should have model table for this stuff ... import the config file
        # topicword_w = topicword_w.reset_index()
        # topicword_w['word_id'] = topicword_w.word_str.map(word.reset_index().set_index('word_str').word_id)
        # topicword_w = topicword_w.set_index(['topic_id','word_id'])
        # topicword['word_wgt'] = topicword_w.word_wgt

        # SAVE
        self.save_table(word, 'VOCAB')
        self.save_table(topicword, 'TOPICWORD_NARROW')
        self.save_table(topicword_wide, 'TOPICWORD')
        # self.save_table(topicword_w, 'TOPICWORD_WEIGHTS')

    def import_table_doctopic(self):
        """Import data into doctopic table"""
        src_file = self.get_source_file('output-doc-topics')
        doctopic = pd.read_csv(src_file, sep='\t', header=None)
        cols = ['doc_id', 'doc_tmp'] + [t for t in range(doctopic.shape[1]-2)]
        doctopic.columns = cols
        doctopic = doctopic.set_index('doc_id')
        doc = doctopic.doc_tmp.str.split(',', expand=True).iloc[:, :2]
        doc.columns = ['src_doc_id', 'doc_label']
        doc.index.name = 'doc_id'
        doctopic = doctopic.drop('doc_tmp', axis=1)
        doctopic_narrow = doctopic.unstack().to_frame('topic_weight')
        doctopic_narrow.index.names = ['doc_id', 'topic_id']
        doctopic_narrow['topic_weight_zscore'] = stats.zscore(doctopic_narrow.topic_weight)

        # SAVE
        self.save_table(doctopic, 'DOCTOPIC')
        self.save_table(doc, 'DOC')
        self.save_table(doctopic_narrow, 'DOCTOPIC_NARROW')

    def import_table_topicphrase(self):
        """Import data into topicphrase table"""
        src_file = self.get_source_file('xml-topic-phrase-report')
        TOPICPHRASE = []
        tree = etree.parse(src_file)
        for topic in tree.xpath('/topics/topic'):
            topic_id = int(topic.xpath('@id')[0])
            for phrase in topic.xpath('phrase'):
                phrase_weight = float(phrase.xpath('@weight')[0])
                phrase_count = int(phrase.xpath('@count')[0])
                topic_phrase = phrase.xpath('text()')[0]
                TOPICPHRASE.append((topic_id, topic_phrase, phrase_weight, phrase_count))
        topicphrase = pd.DataFrame(TOPICPHRASE,
            columns=['topic_id', 'topic_phrase', 'phrase_weight', 'phrase_count'])

        # Add phrase list to TOPIC
        # MOVE TO add_topic_glosses()
        topic = self.get_table('TOPIC')
        topic['phrases'] = topicphrase.groupby('topic_id').apply(lambda x: ', '.join(x.topic_phrase))

        topicphrase.set_index(['topic_id', 'topic_phrase'], inplace=True)

        phrase = topicphrase.value_counts('topic_phrase').to_frame('n_topics').sort_index()
        phrase['n_words'] = topicphrase.groupby('topic_phrase').phrase_count.sum().sort_index()
        M = topicphrase.sort_index().reset_index()
        M['topic_name'] = M.topic_id.astype('str').str.zfill(2)
        phrase['topic_list'] = M.groupby('topic_phrase').topic_name.apply(lambda x: ' '.join(x))\
            .to_frame('topic_list').sort_index()
        del(M)
        phrase['topic_weight_mean'] = topicphrase.groupby(['topic_phrase']).mean().phrase_weight.sort_index()

        topic['topic_label'] = topic.index.astype('str').str.zfill(2) + ': ' + topic.topic_words

        # SAVE
        self.save_table(phrase, 'PHRASE')
        self.save_table(topicphrase, 'TOPICPHRASE')
        self.save_table(topic, 'TOPIC')

    def add_topic_glosses(self):
        """Add glosses to topic table"""
        topicphrase = self.get_table('TOPICPHRASE')
        topic = self.get_table('TOPIC')
        topic['topic_gloss'] = topicphrase['phrase_weight'].unstack().idxmax(1)

        # SAVE
        self.save_table(topic, 'TOPIC')

    def add_diagnostics(self):
        """Add diagnostics data to topics and topicword_diags tables"""
        src_file = self.get_source_file('diagnostics-file')

        TOPIC = []
        TOPICWORD = []

        # Schema
        tkeys = ['id', 'tokens', 'document_entropy', 'word-length', 'coherence',
                 'uniform_dist', 'corpus_dist',
                 'eff_num_words', 'token-doc-diff', 'rank_1_docs',
                 'allocation_ratio', 'allocation_count',
                 'exclusivity']
        tints = ['id', 'tokens']
        wkeys = ['rank', 'count', 'prob', 'cumulative', 'docs', 'word-length', 'coherence',
                 'uniform_dist', 'corpus_dist', 'token-doc-diff', 'exclusivity']
        wints = ['rank', 'count', 'docs', 'word-length']

        tree = etree.parse(src_file)
        for topic in tree.xpath('/model/topic'):
            tvals = []
            for key in tkeys:
                xpath = f'@{key}'
                if key in tints:
                    tvals.append(int(float(topic.xpath(xpath)[0])))
                else:
                    tvals.append(float(topic.xpath(xpath)[0]))
            TOPIC.append(tvals)
            for word in topic.xpath('word'):
                wvals = []
                topic_id = tvals[0]  # Hopefully
                wvals.append(topic_id)
                word_str = word.xpath('text()')[0]
                wvals.append(word_str)
                for key in wkeys:
                    xpath = f'@{key}'
                    if key in wints:
                        wvals.append(int(float(word.xpath(xpath)[0])))
                    else:
                        wvals.append(float(word.xpath(xpath)[0]))
                TOPICWORD.append(wvals)

        tkeys = ['topic_{}'.format(re.sub('-', '_', k)) for k in tkeys]
        wkeys = ['topic_id', 'word_str'] + wkeys
        wkeys = [re.sub('-', '_', k) for k in wkeys]

        topic_diags = pd.DataFrame(TOPIC, columns=tkeys).set_index('topic_id')

        topics = self.get_table('TOPIC')
        topics = pd.concat([topics, topic_diags], axis=1)
        topicword_diags = pd.DataFrame(TOPICWORD, columns=wkeys).set_index(['topic_id', 'word_str'])

        word = self.get_table('VOCAB').reset_index().set_index('word_str')
        topicword_diags['word_id'] = topicword_diags.apply(lambda x: word.loc[x.name[1]].word_id, axis=1)
        topicword_diags = topicword_diags.reset_index().set_index(['topic_id', 'word_id'])

        # SAVE
        self.save_table(topics, 'TOPIC')
        self.save_table(topicword_diags, 'TOPICWORD_DIAGS')

    def do_all(self):
        """Run all importers and adders"""
        self.import_table_state()
        self.import_table_topic()
        self.import_tables_topicword_and_word()
        self.import_table_doctopic()
        self.import_table_topicphrase()
        self.add_diagnostics()
        self.add_topic_glosses()


import numpy as np
import random
import math
import zmq
import time
from scipy.spatial import distance
from tdw.output_data import IsOnNavMesh, Images, Bounds
from PIL import Image
import io
import os
from threading import Thread
from tdw.controller import Controller
from typing import List, Tuple, Dict, Optional, Union
import socket
from contextlib import closing
from tdw.librarian import ModelRecord
from pathlib import Path
import boto3
from botocore.exceptions import ProfileNotFound, ClientError
from subprocess import check_output, Popen, call
import re
from psutil import pid_exists
import base64

class TDWUtils:
    """
    Utility functions for controllers.

    Usage:

    ```python
    from tdw.tdw_utils import TDWUtils
    ```
    """

    VECTOR3_ZERO = {"x": 0, "y": 0, "z": 0}

    # Cached values used during point cloud generation.
    __WIDTH: int = -1
    __HEIGHT: int = -1
    __CAM_TO_IMG_MAT: Optional[np.array] = None

    @staticmethod
    def vector3_to_array(vector3: Dict[str, float]) -> np.array:
        """
        Convert a Vector3 object to a numpy array.

        :param vector3: The Vector3 object, e.g. `{"x": 0, "y": 0, "z": 0}`

        :return A numpy array.
        """

        return np.array([vector3["x"], vector3["y"], vector3["z"]])

    @staticmethod
    def array_to_vector3(arr: np.array) -> Dict[str, float]:
        """
        Convert a numpy array to a Vector3.

        :param arr: The numpy array.

        :return A Vector3, e.g. `{"x": 0, "y": 0, "z": 0}`
        """

        return {"x": float(arr[0]), "y": float(arr[1]), "z": float(arr[2])}

    @staticmethod
    def vector4_to_array(vector4: Dict[str, float]) -> np.array:
        """
        Convert a Vector4 to a numpy array.

        :param vector4: The Vector4 object, e.g. `{"x": 0, "y": 0, "z": 0, "w": 0}`

        :return A numpy array.
        """

        return np.array([vector4["x"], vector4["y"], vector4["z"], vector4["w"]])

    @staticmethod
    def array_to_vector4(arr: np.array) -> Dict[str, float]:
        """
        Convert a numpy array to a Vector4.

        :param arr: The numpy array.

        :return A Vector4, e.g. `{"x": 0, "y": 0, "z": 0, "w": 0}`
        """

        return {"x": arr[0], "y": arr[1], "z": arr[2], "w": arr[3]}

    @staticmethod
    def color_to_array(color: Dict[str, float]) -> np.array:
        """
        Convert a RGB Color to a numpy array.

        :param color: The Color object, e.g. `{"r": 0, "g": 0, "b": 0, "a": 1}`

        :return A numpy array.
        """

        return np.array([round(color["r"] * 255), round(color["g"] * 255), round(color["b"] * 255)])

    @staticmethod
    def array_to_color(arr: np.array) -> Dict[str, float]:
        """
        Convert a numpy array to a RGBA Color. If no A value is supplied it will default to 1.

        :param arr: The array.

        :return A Color, e.g. `{"r": 0, "g": 0, "b": 0, "a": 1}`
        """

        return {"r": arr[0], "g": arr[1], "b": arr[2], "a": 1 if len(arr) == 3 else arr[3]}

    @staticmethod
    def get_random_point_in_circle(center: np.array, radius: float) -> np.array:
        """
        Get a random point in a circle, defined by a center and radius.

        :param center: The center of the circle.
        :param radius: The radius of the circle.

        :return A numpy array. The y value (`arr[1]`) is always 0.
        """

        alpha = 2 * math.pi * random.random()
        r = radius * math.sqrt(random.random())
        x = r * math.cos(alpha) + center[0]
        z = r * math.sin(alpha) + center[2]

        return np.array([x, 0, z])

    @staticmethod
    def get_magnitude(vector3: Dict[str, float]) -> float:
        """
        Get the magnitude of a Vector3.

        :param vector3: The Vector3 object, e.g. `{"x": 0, "y": 0, "z": 0}`

        :return The vector magnitude.
        """

        return np.linalg.norm(TDWUtils.vector3_to_array(vector3))

    @staticmethod
    def extend_line(p0: np.array, p1: np.array, d: float, clamp_y=True) -> np.array:
        """
        Extend the line defined by p0 to p1 by distance d. Clamps the y value to 0.

        :param p0: The origin.
        :param p1: The second point.
        :param d: The distance of which the line is to be extended.
        :param clamp_y: Clamp the y value to 0.

        :return: The position at distance d.
        """

        if clamp_y:
            p0[1] = 0
            p1[1] = 0

        # Get the distance between the two points.
        d0 = distance.euclidean(p0, p1)
        # Get the total distance.
        d_total = d0 + d

        return p1 + ((p1 - p0) * d_total)

    @staticmethod
    def get_distance(vector3_0: Dict[str, float], vector3_1: Dict[str, float]) -> float:
        """
        Calculate the distance between two Vector3 (e.g. `{"x": 0, "y": 0, "z": 0}`) objects.

        :param vector3_0: The first Vector3.
        :param vector3_1: The second Vector3.

        :return The distance.
        """

        return distance.euclidean(TDWUtils.vector3_to_array(vector3_0), TDWUtils.vector3_to_array(vector3_1))

    @staticmethod
    def get_box(width: int, length: int) -> List[Dict[str, int]]:
        """
        Returns a list of x,y positions that can be used to create a box with the `create_exterior_walls` command.
        :param width: The width of the box.
        :param length: The length of the box.

        :return The box as represented by a list of `{"x": x, "y": y}` dictionaries.
        """

        box = []
        for x in range(width):
            for y in range(length):
                if x == 0 or x == width - 1 or y == 0 or y == length - 1:
                    box.append({"x": x, "y": y})
        return box

    @staticmethod
    def get_vector3(x, y, z) -> Dict[str, float]:
        """
        :param x: The x value.
        :param y: The y value.
        :param z: The z value.

        :return: A Vector3: {"x": x, "y", y, "z": z}
        """

        return {"x": x, "y": y, "z": z}

    @staticmethod
    def create_empty_room(width: int, length: int) -> dict:
        """
        :param width: The width of the room.
        :param length: The length of the room.

        :return: A `create_exterior_walls` command that creates a box with dimensions (width, length).
        """

        return {"$type": "create_exterior_walls", "walls": TDWUtils.get_box(width, length)}

    @staticmethod
    def create_room_from_image(filepath: str, exterior_color=(255, 0, 0), interior_color=(0, 0, 0)) -> List[dict]:
        """
        Load a .png file from the disk and use it to create a room. Each pixel on the image is a grid point.

        :param filepath: The absolute filepath to the image.
        :param exterior_color: The color on the image marking exterior walls (default=red).
        :param interior_color: The color on the image marking interior walls (default=black).

        :return: A list of commands: The first creates the exterior walls, and the second creates the interior walls.
        """

        exterior_walls = []
        interior_walls = []

        # Read the image.
        img = Image.open(filepath)
        pixels = img.load()
        col, row = img.size

        # Read each pixel as a grid point.
        for i in range(row):
            for j in range(col):
                pixel = pixels[i, j]
                if len(pixel) == 4:
                    pixel = (pixel[0], pixel[1], pixel[2])
                if pixel == exterior_color:
                    exterior_walls.append({"x": i, "y": col - j})
                elif pixel == interior_color:
                    interior_walls.append({"x": i, "y": col - j})

        return [{"$type": "create_exterior_walls",
                 "walls": exterior_walls},
                {"$type": "create_interior_walls",
                 "walls": interior_walls}]

    @staticmethod
    def save_images(images: Images, filename: str, output_directory="dist", resize_to=None, append_pass: bool = True) -> None:
        """
        Save each image in the Images object.
        The name of the image will be: pass_filename.extension, e.g.: `"0000"` -> `depth_0000.png`
        The images object includes the pass and extension information.

        :param images: The Images object. Contains each capture pass plus metadata.
        :param output_directory: The directory to write images to.
        :param filename: The filename of each image, minus the extension. The image pass will be appended as a prefix.
        :param resize_to: Specify a (width, height) tuple to resize the images to. This is slower than saving as-is.
        :param append_pass: If false, the image pass will _not_ be appended to the filename as a prefix, e.g.: `"0000"`: -> "`0000.jpg"`
        """

        if not os.path.isdir(output_directory):
            os.makedirs(output_directory)

        for i in range(images.get_num_passes()):
            if append_pass:
                fi = images.get_pass_mask(i)[1:] + "_" + filename + "." + images.get_extension(i)
            else:
                fi = filename + "." + images.get_extension(i)

            if resize_to:
                TDWUtils.get_pil_image(images, i).resize((resize_to[0], resize_to[1]), Image.LANCZOS)\
                    .save(os.path.join(output_directory, fi))
            else:
                pass_mask = images.get_pass_mask(i)
                path = os.path.join(output_directory, fi)
                # The depth passes aren't png files, so we need to convert them.
                if pass_mask == "_depth" or pass_mask == "_depth_simple":
                    # Save the image.
                    Image.fromarray(TDWUtils.get_shaped_depth_pass(images=images, index=i)).save(path)
                # Every other pass can be saved directly to disk.
                else:
                    with open(path, "wb") as f:
                        f.write(images.get_image(i))

    @staticmethod
    def get_shaped_depth_pass(images: Images, index: int) -> np.array:
        """
        The `_depth` and `_depth_simple` passes are a 1D array of RGB values, as oppposed to a png or jpg like every other pass.
        This function reshapes the array into a 2D array of RGB values.

        :param images: The `Images` output data.
        :param index: The index in `Images` of the depth pass. See: `Images.get_pass_mask()`.

        :return: A reshaped depth pass. Shape is: `(height, width, 3)`.
        """

        return np.flip(np.reshape(images.get_image(index), (images.get_height(), images.get_width(), 3)), 0)

    @staticmethod
    def zero_padding(integer: int, width=4) -> str:
        """
        :param integer: The integer being converted.
        :param width: The total number of digits in the string. If integer == 3 and width == 4, output is: "0003".

        :return A string representation of an integer padded with zeroes, e.g. converts `3` to `"0003"`.
        """

        return str(integer).zfill(width)

    @staticmethod
    def get_pil_image(images: Images, index: int) -> Image:
        """
        Converts Images output data to a PIL Image object.
        Use this function to read and analyze an image in memory.
        Do NOT use this function to save image data to disk; `save_image` is much faster.

        :param images: Images data from the build.
        :param index: The index of the image in Images.get_image

        :return A PIL image.
        """

        return Image.open(io.BytesIO(images.get_image(index)))

    @staticmethod
    def get_random_position_on_nav_mesh(c: Controller, width: float, length: float, x_e=0, z_e=0, bake=True, rng=random.uniform) -> Tuple[float, float, float]:
        """
        Returns a random position on a NavMesh.

        :param c: The controller.
        :param width: The width of the environment.
        :param length: The length of the environment.
        :param bake: If true, send bake_nav_mesh.
        :param rng: Random number generator.
        :param x_e: The x position of the environment.
        :param z_e: The z position of the environment.

        :return The coordinates as a tuple `(x, y, z)`
        """

        if bake:
            c.communicate({'$type': 'bake_nav_mesh'})

        # Try to find a valid position on the NavMesh.
        is_on = False
        x, y, z = (0, 0, 0)
        while not is_on:
            # Get a random position.
            x = rng(-width / 2, width / 2) + x_e
            z = rng(-length / 2, length / 2) + z_e
            resp = c.communicate(
                {'$type': 'send_is_on_nav_mesh',
                 'position': {'x': x, 'y': 0, 'z': z},
                 'max_distance': 4.0
                 })
            answer = IsOnNavMesh(resp[0])
            is_on = answer.get_is_on()
            x, y, z = answer.get_position()
        return x, y, z

    @staticmethod
    def set_visual_material(c: Controller, substructure: List[dict], object_id: int, material: str, quality="med") -> List[dict]:
        """
        :param c: The controller.
        :param substructure: The metadata substructure of the object.
        :param object_id: The ID of the object in the scene.
        :param material: The name of the new material.
        :param quality: The quality of the material.

        :return A list of commands to set ALL visual materials on an object to a single material.
        """

        commands = []
        for sub_object in substructure:
            for i in range(len(sub_object["materials"])):
                commands.extend([c.get_add_material(material, library="materials_" + quality + ".json"),
                                 {"$type": "set_visual_material",
                                  "id": object_id,
                                  "material_name": material,
                                  "object_name": sub_object["name"],
                                  "material_index": i}])
        return commands

    @staticmethod
    def get_depth_values(image: np.array, depth_pass: str = "_depth", width: int = 256, height: int = 256, near_plane: float = 0.1, far_plane: float = 100) -> np.array:
        """
        Get the depth values of each pixel in a _depth image pass.
        The far plane is hardcoded as 100. The near plane is hardcoded as 0.1.
        (This is due to how the depth shader is implemented.)

        :param image: The image pass as a numpy array.
        :param depth_pass: The type of depth pass. This determines how the values are decoded. Options: `"_depth"`, `"_depth_simple"`.
        :param width: The width of the screen in pixels. See output data `Images.get_width()`.
        :param height: The height of the screen in pixels. See output data `Images.get_height()`.
        :param near_plane: The near clipping plane. See command `set_camera_clipping_planes`. The default value in this function is the default value of the near clipping plane.
        :param far_plane: The far clipping plane. See command `set_camera_clipping_planes`. The default value in this function is the default value of the far clipping plane.

        :return An array of depth values.
        """

        # Convert the image to a 2D image array.
        image = np.flip(np.reshape(image, (height, width, 3)), 0)
        if depth_pass == "_depth":
            depth_values = np.array((image[:, :, 0] + image[:, :, 1] / 256.0 + image[:, :, 2] / (256.0 ** 2)))
        elif depth_pass == "_depth_simple":
            depth_values = image[:, :, 0] / 256.0
        else:
            raise Exception(f"Invalid depth pass: {depth_pass}")
        # Un-normalize the depth values.
        return (depth_values * ((far_plane - near_plane) / 256.0)).astype(np.float32)

    @staticmethod
    def get_point_cloud(depth, camera_matrix: Union[np.array, tuple], vfov: float = 54.43222, filename: str = None, near_plane: float = 0.1, far_plane: float = 100) -> np.array:
        """
        Create a point cloud from an numpy array of depth values.

        :param depth: Depth values converted from a depth pass. See: `TDWUtils.get_depth_values()`
        :param camera_matrix: The camera matrix as a tuple or numpy array. See: [`send_camera_matrices`](https://github.com/threedworld-mit/tdw/blob/master/Documentation/api/command_api.md#send_camera_matrices).
        :param vfov: The field of view. See: [`set_field_of_view`](https://github.com/threedworld-mit/tdw/blob/master/Documentation/api/command_api.md#set_field_of_view)
        :param filename: If not None, the point cloud data will be written to this file.
        :param near_plane: The near clipping plane. See command `set_camera_clipping_planes`. The default value in this function is the default value of the near clipping plane.
        :param far_plane: The far clipping plane. See command `set_camera_clipping_planes`. The default value in this function is the default value of the far clipping plane.

        :return: An point cloud as a numpy array of `[x, y, z]` coordinates.
        """

        if isinstance(camera_matrix, tuple):
            camera_matrix = np.array(camera_matrix)
        camera_matrix = np.linalg.inv(camera_matrix.reshape((4, 4)))

        # Different from real-world camera coordinate system.
        # OpenGL uses negative z axis as the camera front direction.
        # x axes are same, hence y axis is reversed as well.
        # Source: https://learnopengl.com/Getting-started/Camera
        rot = np.array([[1, 0, 0, 0],
                        [0, -1, 0, 0],
                        [0, 0, -1, 0],
                        [0, 0, 0, 1]])
        camera_matrix = np.dot(camera_matrix, rot)

        # Cache some calculations we'll need to use every time.
        if TDWUtils.__HEIGHT != depth.shape[0] or TDWUtils.__WIDTH != depth.shape[1]:
            TDWUtils.__HEIGHT = depth.shape[0]
            TDWUtils.__WIDTH = depth.shape[1]

            img_pixs = np.mgrid[0: depth.shape[0], 0: depth.shape[1]].reshape(2, -1)
            # Swap (v, u) into (u, v).
            img_pixs[[0, 1], :] = img_pixs[[1, 0], :]
            img_pix_ones = np.concatenate((img_pixs, np.ones((1, img_pixs.shape[1]))))

            # Calculate the intrinsic matrix from vertical_fov.
            # Motice that hfov and vfov are different if height != width
            # We can also get the intrinsic matrix from opengl's perspective matrix.
            # http://kgeorge.github.io/2014/03/08/calculating-opengl-perspective-matrix-from-opencv-intrinsic-matrix
            vfov = vfov / 180.0 * np.pi
            tan_half_vfov = np.tan(vfov / 2.0)
            tan_half_hfov = tan_half_vfov * TDWUtils.__WIDTH / float(TDWUtils.__HEIGHT)
            fx = TDWUtils.__WIDTH / 2.0 / tan_half_hfov  # focal length in pixel space
            fy = TDWUtils.__HEIGHT / 2.0 / tan_half_vfov
            intrinsics = np.array([[fx, 0, TDWUtils.__WIDTH / 2.0],
                                   [0, fy, TDWUtils.__HEIGHT / 2.0],
                                   [0, 0, 1]])
            img_inv = np.linalg.inv(intrinsics[:3, :3])
            TDWUtils.__CAM_TO_IMG_MAT = np.dot(img_inv, img_pix_ones)

        points_in_cam = np.multiply(TDWUtils.__CAM_TO_IMG_MAT, depth.reshape(-1))
        points_in_cam = np.concatenate((points_in_cam, np.ones((1, points_in_cam.shape[1]))), axis=0)
        points_in_world = np.dot(camera_matrix, points_in_cam)
        points_in_world = points_in_world[:3, :].reshape(3, TDWUtils.__WIDTH, TDWUtils.__HEIGHT)
        points_in_cam = points_in_cam[:3, :].reshape(3, TDWUtils.__WIDTH, TDWUtils.__HEIGHT)
        if filename is not None:
            f = open(filename, 'w')
            for i in range(points_in_world.shape[1]):
                for j in range(points_in_world.shape[2]):
                    if points_in_cam[2, i, j] < (far_plane - near_plane):
                        f.write(f'{points_in_world[0, i, j]};{points_in_world[1, i, j]};{points_in_world[2, i, j]}\n')
        return points_in_world

    @staticmethod
    def create_avatar(avatar_type="A_Img_Caps_Kinematic", avatar_id="a", position=None, look_at=None) -> List[dict]:
        """
        This is a wrapper for `create_avatar` and, optionally, `teleport_avatar_to` and `look_at_position`.

        :param avatar_type: The type of avatar.
        :param avatar_id: The avatar ID.
        :param position: The position of the avatar. If this is None, the avatar won't teleport.
        :param look_at: If this isn't None, the avatar will look at this position.

        :return A list of commands to create theavatar.
        """

        # Create the avatar.
        commands = [{"$type": "create_avatar",
                     "type": avatar_type,
                     "id": avatar_id}]

        # Teleport the avatar.
        if position:
            commands.append({"$type": "teleport_avatar_to",
                             "avatar_id": avatar_id,
                             "position": position})
        if look_at:
            commands.append({"$type": "look_at_position",
                             "avatar_id": avatar_id,
                             "position": look_at})
        return commands

    @staticmethod
    def _send_start_build(socket, controller_address: str) -> dict:
        """
        This sends a command to the launch_binaries daemon running on a remote node
        to start a binary connected to the given controller address.

        :param socket: The zmq socket.
        :param controller_address: The host name or ip address of node running the controller.

        :return Build info dictionary containing build port.
        """
        request = {"type": "start_build",
                   "controller_address": controller_address}
        socket.send_json(request)
        build_info = socket.recv_json()
        return build_info

    @staticmethod
    def _send_keep_alive(socket, build_info: dict) -> dict:
        """
        This sends a command to the launch_binaries daemon running on a remote node
        to mark a given binary as still alive, preventing garbage collection.

        :param socket: The zmq socket.
        :param build_info: A diciontary containing the build_port.

        :return a heartbeat indicating build is still alive.
        """

        build_port = build_info["build_port"]
        request = {"type": "keep_alive", "build_port": build_port}
        socket.send_json(request)
        heartbeat = socket.recv_json()
        return heartbeat

    @staticmethod
    def _send_kill_build(socket, build_info: dict) -> dict:
        """
        This sends a command to the launch_binaries daemon running on a remote node to terminate a given binary.

        :param socket: The zmq socket.
        :param build_info: A diciontary containing the build_port.

        :return A kill_status indicating build has been succesfully terminated.
        """

        build_port = build_info["build_port"]
        request = {"type": "kill_build", "build_port": build_port}
        socket.send_json(request)
        kill_status = socket.recv_json()
        return kill_status

    @staticmethod
    def _keep_alive_thread(socket, build_info: dict) -> None:
        """
        This is a wrapper around the keep alive command to be executed in a separate thread.

        :param socket: The zmq socket.
        :param build_info: A diciontary containing the build_port.
        """
        while True:
            TDWUtils._send_keep_alive(socket, build_info)
            time.sleep(60)

    @staticmethod
    def launch_build(listener_port: int, build_address: str, controller_address: str) -> dict:
        """
        Connect to a remote binary_manager daemon and launch an instance of a TDW build.

        Returns the necessary information for a local controller to connect.
        Use this function to automatically launching binaries on remote (or local) nodes, and to
        automatically shut down the build after controller is finished. Call in the constructor
        of a controller and pass the build_port returned in build_info to the parent Controller class.

        :param listener_port: The port launch_binaries is listening on.
        :param build_address: Remote IP or hostname of node running launch_binaries.
        :param controller_address: IP or hostname of node running controller.

        :return The build_info dictionary containing build_port.
        """

        context = zmq.Context()
        socket = context.socket(zmq.REQ)
        socket.connect("tcp://" + build_address + ":%s" % listener_port)
        build_info = TDWUtils._send_start_build(socket, controller_address)
        thread = Thread(target=TDWUtils._keep_alive_thread,
                        args=(socket, build_info))
        thread.setDaemon(True)
        thread.start()
        return build_info

    @staticmethod
    def get_unity_args(arg_dict: dict) -> List[str]:
        """
        :param arg_dict: A dictionary of arguments. Key=The argument prefix (e.g. port) Value=Argument value.

        :return The formatted command line string that is accepted by unity arg parser.
        """

        formatted_args = []
        for key, value in arg_dict.items():
            prefix = "-" + key + "="
            if type(value) == list:
                prefix += ",".join([str(v) for v in value])
            else:
                prefix += str(value)
            formatted_args += [prefix]
        return formatted_args

    @staticmethod
    def find_free_port() -> int:
        """
        :return a free port.
        """

        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
            s.bind(("", 0))
            return int(s.getsockname()[1])

    @staticmethod
    def get_unit_scale(record: ModelRecord) -> float:
        """
        :param record: The model record.

        :return The scale factor required to scale a model to 1 meter "unit scale".
        """

        bounds = record.bounds

        # Get the "unit scale" of the object.
        s = 1 / max(
            bounds['top']['y'] - bounds['bottom']['y'],
            bounds['front']['z'] - bounds['back']['z'],
            bounds['right']['x'] - bounds['left']['x'])
        return s

    @staticmethod
    def validate_amazon_s3() -> bool:
        """
        Validate that your local Amazon S3 credentials are set up correctly.

        :return True if everything is OK.
        """

        config_path = Path.home().joinpath(".aws/config")
        new_config_path = not config_path.exists()
        # Generate a valid config file.
        if new_config_path:
            config_path.write_text("[default]\nregion = us-east-1\noutput = json")
            print(f"Generated a new config file: {config_path.resolve()}")
        try:
            session = boto3.Session(profile_name="tdw")
            s3 = session.resource("s3")
            s3.meta.client.head_object(Bucket='tdw-private', Key='models/windows/2018-2019.1/iron_box')
            return True
        except ProfileNotFound:
            print(f"ERROR! Your AWS credentials file is not set up correctly.")
            print("Your AWS credentials must have a [tdw] profile with valid keys.")
            return False
        except ClientError as e:
            print("ERROR! Could not access bucket tdw-private. Make sure you have the right permissions.")
            error_code = e.response['Error']['Code']
            print(e, error_code)
            return False

    @staticmethod
    def get_base64_flex_particle_forces(forces: list) -> str:
        """
        :param forces: The forces (see Flex documentation for how to arrange this array).

        :return: An array of Flex particle forces encoded in base64.
        """

        forces = np.array(forces, dtype=np.float32)
        return base64.b64encode(forces).decode()

    @staticmethod
    def color_to_hashable(color: Union[np.array, Tuple[int, int, int]]) -> int:
        """
        :param color: The color as an RGB array or tuple, where each value is between 0 and 255.

        :return: A hashable integer representation of the color array.
        """

        return (color[0] << 16) + (color[1] << 8) + color[2]

    @staticmethod
    def hashable_to_color(hashable: int) -> np.array:
        """
        :param hashable: A hashable integer representing an RGB color.

        :return: A color as a numpy array of integers between 0 and 255: `[r, g, b]`
        """

        return np.array([(hashable >> 16) & 255, (hashable >> 8) & 255, hashable & 255], dtype=int)

    @staticmethod
    def get_bounds_dict(bounds: Bounds, index: int) -> Dict[str, np.array]:
        """
        :param bounds: Bounds output data.
        :param index: The index in `bounds` of the target object.

        :return: A dictionary of the bounds. Key = the name of the position. Value = the position as a numpy array.
        """

        return {"top": np.array(bounds.get_top(index)),
                "bottom": np.array(bounds.get_bottom(index)),
                "left": np.array(bounds.get_left(index)),
                "right": np.array(bounds.get_right(index)),
                "front": np.array(bounds.get_front(index)),
                "back": np.array(bounds.get_back(index)),
                "center": np.array(bounds.get_center(index))}

    @staticmethod
    def get_bounds_extents(bounds: Bounds, index: int) -> np.array:
        """
        :param bounds: Bounds output data.
        :param index: The index in `bounds` of the target object.

        :return: The width (left to right), length (front to back), and height (top to bottom) of the bounds as a numpy array.
        """

        return np.array([np.linalg.norm(np.array(bounds.get_left(index)) - np.array(bounds.get_right(index))),
                         np.linalg.norm(np.array(bounds.get_front(index)) - np.array(bounds.get_back(index))),
                         np.linalg.norm(np.array(bounds.get_top(index)) - np.array(bounds.get_bottom(index)))])

    @staticmethod
    def get_closest_position_in_bounds(origin: np.array, bounds: Bounds, index: int) -> np.array:
        """
        :param origin: The origin from which the distance is calculated.
        :param bounds: Bounds output data.
        :param index: The index in `bounds` of the target object.

        :return: The position on the object bounds that is closest to `origin`.
        """

        object_bounds = TDWUtils.get_bounds_dict(bounds=bounds, index=index)

        # Get the closest point on the bounds.
        min_destination = ""
        min_distance = 10000
        for p in object_bounds:
            d = np.linalg.norm(origin - object_bounds[p])
            if d < min_distance:
                min_distance = d
                min_destination = p
        return object_bounds[min_destination]

    @staticmethod
    def get_angle(forward: np.array, origin: np.array, position: np.array) -> float:
        """
          :param position: The target position.
          :param origin: The origin position of the directional vector.
          :param forward: The forward directional vector.

          :return: The angle in degrees between `forward` and the direction vector from `origin` to `position`.
          """

        # Get the normalized directional vector to the target position.
        p0 = np.array([origin[0], origin[2]])
        p1 = np.array([position[0], position[2]])
        d = p1 - p0
        d = d / np.linalg.norm(d)
        f = np.array([forward[0], forward[2]])

        dot = f[0] * d[0] + f[1] * d[1]
        det = f[0] * d[1] - f[1] * d[0]
        angle = np.arctan2(det, dot)
        angle = np.rad2deg(angle)
        return angle

    @staticmethod
    def get_angle_between(v1: np.array, v2: np.array) -> float:
        """
        :param v1: The first directional vector.
        :param v2: The second directional vector.

        :return: The angle in degrees between two directional vectors.
        """

        ang1 = np.arctan2(v1[2], v1[0])
        ang2 = np.arctan2(v2[2], v2[0])

        return np.rad2deg((ang1 - ang2) % (2 * np.pi))

    @staticmethod
    def rotate_position_around(position: np.array, angle: float, origin: np.array = None) -> np.array:
        """
        Rotate a position by a given angle around a given origin.

        :param origin: The origin position.  If None, the origin is `[0, 0, 0]`
        :param position: The point being rotated.
        :param angle: The angle in degrees.

        :return: The rotated position.
        """

        if origin is None:
            origin = np.array([0, 0, 0])

        radians = np.deg2rad(angle)
        x, y = position[0], position[2]
        offset_x, offset_y = origin[0], origin[2]
        adjusted_x = (x - offset_x)
        adjusted_y = (y - offset_y)
        cos_rad = np.cos(radians)
        sin_rad = np.sin(radians)
        qx = offset_x + cos_rad * adjusted_x + sin_rad * adjusted_y
        qy = offset_y + -sin_rad * adjusted_x + cos_rad * adjusted_y

        return np.array([qx, position[1], qy])

    @staticmethod
    def euler_angles_to_rpy(euler_angles: np.array) -> np.array:
        """
        Convert Euler angles to ROS RPY angles.

        :param euler_angles: A numpy array: `[x, y, z]` Euler angles in degrees.

        :return: A numpy array: `[r, p, y]` angles in radians.
        """

        # Source: https://github.com/Unity-Technologies/URDF-Importer/blob/c41208565419b04907496baa93ad1b675d41dc20/com.unity.robotics.urdf-importer/Runtime/Extensions/TransformExtensions.cs#L85-L92
        return np.radians(np.array([-euler_angles[2], euler_angles[0], -euler_angles[1]]))

class AudioUtils:
    """
    Utility class for recording audio in TDW using [fmedia](https://stsaz.github.io/fmedia/).

    Usage:

    ```python
    from tdw.tdw_utils import AudioUtils
    from tdw.controller import Controller

    c = Controller()

    initialize_trial()  # Your code here.

    # Begin recording audio. Automatically stop recording at 10 seconds.
    AudioUtils.start(output_path="path/to/file.wav", until=(0, 10))

    do_trial()  # Your code here.

    # Stop recording.
    AudioUtils.stop()
    ```
    """

    # The process ID of the audio recorder.
    RECORDER_PID: Optional[int] = None
    # The audio capture device.
    DEVICE: Optional[str] = None

    @staticmethod
    def get_system_audio_device() -> str:
        """
        :return: The audio device that can be used to capture system audio.
        """

        devices = check_output(["fmedia", "--list-dev"]).decode("utf-8").split("Capture:")[1]
        dev_search = re.search("device #(.*): Stereo Mix", devices, flags=re.MULTILINE)
        assert dev_search is not None, "No suitable audio capture device found:\n" + devices
        return dev_search.group(1)

    @staticmethod
    def start(output_path: Union[str, Path], until: Optional[Tuple[int, int]] = None) -> None:
        """
        Start recording audio.

        :param output_path: The path to the output file.
        :param until: If not None, fmedia will record until `minutes:seconds`. The value must be a tuple of 2 integers. If None, fmedia will record until you send `AudioUtils.stop()`.
        """

        if isinstance(output_path, str):
            p = Path(output_path).resolve()
        else:
            p = output_path

        # Create the directory.
        if not p.parent.exists():
            p.parent.mkdir(parents=True)

        # Set the capture device.
        if AudioUtils.DEVICE is None:
            AudioUtils.DEVICE = AudioUtils.get_system_audio_device()
        fmedia_call = ["fmedia",
                       "--record",
                       f"--dev-capture={AudioUtils.DEVICE}",
                       f"--out={str(p.resolve())}",
                       "--globcmd=listen"]
        # Automatically stop recording.
        if until is not None:
            fmedia_call.append(f"--until={TDWUtils.zero_padding(until[0], 2)}:{TDWUtils.zero_padding(until[1], 2)}")
        with open(os.devnull, "w+") as f:
            AudioUtils.RECORDER_PID = Popen(fmedia_call,
                                            stderr=f).pid

    @staticmethod
    def stop() -> None:
        """
        Stop recording audio (if any fmedia process is running).
        """

        if AudioUtils.RECORDER_PID is not None:
            with open(os.devnull, "w+") as f:
                call(['fmedia', '--globcmd=quit'], stderr=f, stdout=f)
            AudioUtils.RECORDER_PID = None

    @staticmethod
    def is_recording() -> bool:
        """
        :return: True if the fmedia recording process still exists.
        """

        return AudioUtils.RECORDER_PID is not None and pid_exists(AudioUtils.RECORDER_PID)

class QuaternionUtils:
    """
    Helper functions for using quaternions.

    Quaternions are always numpy arrays in the following order: `[x, y, z, w]`.
    This is the order returned in all Output Data objects.

    Vectors are always numpy arrays in the following order: `[x, y, z]`.
    """

    """:class_var
    The global up directional vector.
    """
    UP = np.array([0, 1, 0])
    """:class_var
    The global forward directional vector.
    """
    FORWARD: np.array = np.array([0, 0, 1])
    """:class_var
    The quaternion identity rotation.
    """
    IDENTITY = np.array([0, 0, 0, 1])

    @staticmethod
    def get_inverse(q: np.array) -> np.array:
        """
        Source: https://referencesource.microsoft.com/#System.Numerics/System/Numerics/Quaternion.cs

        :param q: The quaternion.

        :return: The inverse of the quaternion.
        """

        x = q[0]
        y = q[1]
        z = q[2]
        w = q[3]

        ls = x * x + y * y + z * z + w * w
        inv = 1.0 / ls

        return np.array([-x * inv, -y * inv, -z * inv, w * inv])

    @staticmethod
    def multiply(q1: np.array, q2: np.array) -> np.array:
        """
        Multiply two quaternions.
        Source: https://stackoverflow.com/questions/4870393/rotating-coordinate-system-via-a-quaternion

        :param q1: The first quaternion.
        :param q2: The second quaternion.
        :return: The multiplied quaternion: `q1 * q2`
        """

        x1 = q1[0]
        y1 = q1[1]
        z1 = q1[2]
        w1 = q1[3]

        x2 = q2[0]
        y2 = q2[1]
        z2 = q2[2]
        w2 = q2[3]

        w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
        x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
        y = w1 * y2 + y1 * w2 + z1 * x2 - x1 * z2
        z = w1 * z2 + z1 * w2 + x1 * y2 - y1 * x2
        return np.array([x, y, z, w])

    @staticmethod
    def get_conjugate(q: np.array) -> np.array:
        """
        Source: https://stackoverflow.com/questions/4870393/rotating-coordinate-system-via-a-quaternion

        :param q: The quaternion.

        :return: The conjugate of the quaternion: `[-x, -y, -z, w]`
        """

        x = q[0]
        y = q[1]
        z = q[2]
        w = q[3]

        return np.array([-x, -y, -z, w])

    @staticmethod
    def multiply_by_vector(q: np.array, v: np.array) -> np.array:
        """
        Multiply a quaternion by a vector.
        Source: https://stackoverflow.com/questions/4870393/rotating-coordinate-system-via-a-quaternion

        :param q: The quaternion.
        :param v: The vector.

        :return: A directional vector calculated from: `q * v`
        """

        q2 = (v[0], v[1], v[2], 0.0)
        return QuaternionUtils.multiply(QuaternionUtils.multiply(q, q2), QuaternionUtils.get_conjugate(q))[:-1]

    @staticmethod
    def world_to_local_vector(position: np.array, origin: np.array, rotation: np.array) -> np.array:
        """
        Convert a vector position in absolute world coordinates to relative local coordinates.
        Source: https://answers.unity.com/questions/601062/what-inversetransformpoint-does-need-explanation-p.html

        :param position: The position vector in world coordinates.
        :param origin: The origin vector of the local space in world coordinates.
        :param rotation: The rotation quaternion of the local coordinate space.

        :return: `position` in local coordinates.
        """

        return QuaternionUtils.multiply_by_vector(q=QuaternionUtils.get_inverse(q=rotation), v=position - origin)

    @staticmethod
    def get_up_direction(q: np.array) -> np.array:
        """
        :param q: The rotation as a quaternion.

        :return: A directional vector corresponding to the "up" direction from the quaternion.
        """

        return QuaternionUtils.multiply_by_vector(q, QuaternionUtils.UP)

    @staticmethod
    def euler_angles_to_quaternion(euler: np.array) -> np.array:
        """
        Convert Euler angles to a quaternion.

        :param euler: The Euler angles vector.

        :return: The quaternion representation of the Euler angles.
        """

        roll = euler[0]
        pitch = euler[1]
        yaw = euler[2]
        cy = np.cos(yaw * 0.5)
        sy = np.sin(yaw * 0.5)
        cp = np.cos(pitch * 0.5)
        sp = np.sin(pitch * 0.5)
        cr = np.cos(roll * 0.5)
        sr = np.sin(roll * 0.5)

        w = cy * cp * cr + sy * sp * sr
        x = cy * cp * sr - sy * sp * cr
        y = sy * cp * sr + cy * sp * cr
        z = sy * cp * cr - cy * sp * sr
        return np.array([x, y, z, w])

    @staticmethod
    def quaternion_to_euler_angles(quaternion: np.array) -> np.array:
        """
        Convert a quaternion to Euler angles.

        :param quaternion: A quaternion as a nump array.

        :return: The Euler angles representation of the quaternion.
        """

        x = quaternion[0]
        y = quaternion[1]
        z = quaternion[2]
        w = quaternion[3]
        ysqr = y * y

        t0 = +2.0 * (w * x + y * z)
        t1 = +1.0 - 2.0 * (x * x + ysqr)
        ex = np.degrees(np.arctan2(t0, t1))

        t2 = +2.0 * (w * y - z * x)
        t2 = np.where(t2 > +1.0, +1.0, t2)

        t2 = np.where(t2 < -1.0, -1.0, t2)
        ey = np.degrees(np.arcsin(t2))

        t3 = +2.0 * (w * z + x * y)
        t4 = +1.0 - 2.0 * (ysqr + z * z)
        ez = np.degrees(np.arctan2(t3, t4))

        return np.array([ex, ey, ez])

    @staticmethod
    def get_y_angle(q1: np.array, q2: np.array) -> float:
        """
        Source: https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles

        :param q1: The first quaternion.
        :param q2: The second quaternion.

        :return: The angle between the two quaternions in degrees around the y axis.
        """

        qd = QuaternionUtils.multiply(QuaternionUtils.get_conjugate(q1), q2)

        return np.rad2deg(2 * np.arcsin(qd[1]))


import os
import numpy as np
import matplotlib.pyplot as plt
from scipy.io import loadmat
from svm import SVM

#  Exercise 6 | Support Vector Machines

def linearKernel(x1, x2):
    return x1.T @ x2

def gaussianKernel(x1, x2, sigma):
    diff = x1 - x2
    return np.exp(-(diff.T @ diff)/(2*sigma*sigma))

def plotData(x, y):
    pos = np.where(y==1)
    neg = np.where(y==0)
    plt.plot(x[pos[0],0], x[pos[0],1], 'b+', lw = 1.5, ms = 4)
    plt.plot(x[neg[0],0], x[neg[0],1], 'ko', mfc = 'y', ms = 4)

def visualizeBoundaryLinear(X, y, model):
    w = model.w
    b = model.b
    xp = np.linspace(min(X[:,0]), max(X[:,0]), 100)
    yp = - (w[0] * xp + b)/w[1]
    plotData(X, y)
    plt.plot(xp, yp, '-b')

def visualizeBoundary(X, y, model):
    #VISUALIZEBOUNDARY plots a non-linear decision boundary learned by the SVM

    # Plot the training data on top of the boundary
    plotData(X, y)

    # Make classification predictions over a grid of values
    x1plot = np.linspace(min(X[:,0]), max(X[:,0]), 100)
    x2plot = np.linspace(min(X[:,1]), max(X[:,1]), 100)
    X1, X2 = np.meshgrid(x1plot, x2plot)
    vals = np.zeros(X1.shape)
    for i in range(X1.shape[1]):
        this_X = np.c_[X1[:, i], X2[:, i]]
        vals[:, i] = model.svmPredict(this_X)

    # Plot the SVM boundary
    plt.contour(X1, X2, vals, [0.5])

def dataset3Params(X, y, Xval, yval):
    Cs = np.array([0.01, 0.03, 0.1, 0.3, 1, 3, 10, 30])
    sigmas = Cs
    model = SVM()

    costs = np.zeros((Cs.size, sigmas.size))
    for i in range(Cs.size):
        for j in range(sigmas.size):
            sigma = sigmas[j]
            #need to define gaussianKernelLambda lambda here again to capture sigma
            gaussianKernelLambda = lambda x1, x2: (gaussianKernel(x1, x2, sigma))
            gaussianKernelLambda.__name__ = 'gaussianKernel'

            model.svmTrain(X, y.astype(float), Cs[i], gaussianKernelLambda)
            predictions = model.svmPredict(Xval)
            costs[i, j] = np.mean(predictions != yval)
    i, j = np.unravel_index(np.argmin(costs), costs.shape)
    return (Cs[i], sigmas[j])

## =============== Part 1: Loading and Visualizing Data ================
print('Loading and Visualizing Data ...')

# Load from ex6data1:
# You will have X, y in your environment
scriptdir = os.path.dirname(os.path.realpath(__file__))
data = loadmat(scriptdir + '//ex6data1.mat')
X = data['X']
y = data['y'].ravel()
# Plot training data
plotData(X, y)
plt.show()

input('Program paused. Press enter to continue.')

## ==================== Part 2: Training Linear SVM ====================
#  The following code will train a linear SVM on the dataset and plot the
#  decision boundary learned.

print('\nTraining Linear SVM ...')

# You should try to change the C value below and see how the decision
# boundary varies (e.g., try C = 1000)
C = 1
model = SVM()
#we passed y.astype(float), firstly to be sure y -1 can be assigned to singed int.
#secondly we want to preserve original y to 0 and 1
model.svmTrain(X, y.astype(float), C, linearKernel, 1e-3, 20)

visualizeBoundaryLinear(X, y, model)
plt.show()
input('Program paused. Press enter to continue.\n')

## =============== Part 3: Implementing Gaussian Kernel ===============

print('\nEvaluating the Gaussian Kernel ...')

x1 = np.array([1, 2, 1])
x2 = np.array([0, 4, -1])
sigma = 2

#use lambda to be able to pass sigma
gaussianKernelLambda = lambda x1, x2: (gaussianKernel(x1, x2, sigma))
gaussianKernelLambda.__name__ = 'gaussianKernel'

sim = gaussianKernelLambda(x1, x2)

print(f'Gaussian Kernel between x1 = [1c 2, 1], x2 = [0, 4, -1], sigma = {sigma} :')
print(f'\t{sim}\n(for sigma = 2, this value should be about 0.324652)\n')

input('Program paused. Press enter to continue.\n')

## =============== Part 4: Visualizing Dataset 2 ================

print('Loading and Visualizing Data ...')

# Load from ex6data2:
# You will have X, y in your environment
data = loadmat(scriptdir + '//ex6data2.mat')
X = data['X']
y = data['y'].ravel()

# Plot training data
plotData(X, y)
plt.show()

input('Program paused. Press enter to continue.\n')

## ========== Part 5: Training SVM with RBF Kernel (Dataset 2) ==========

print('\nTraining SVM with RBF Kernel (this may take 1 to 2 minutes) ...')

# SVM Parameters
C = 1; sigma = 0.1

# We set the tolerance and max_passes lower here so that the code will run
# faster. However, in practi ce, you will want to run the training to
# convergence.
model.svmTrain(X, y.astype(float), C, gaussianKernelLambda)
visualizeBoundary(X, y, model)
plt.show()
input('Program paused. Press enter to continue.')

## =============== Part 6: Visualizing Dataset 3 ================

print('Loading and Visualizing Data ...')

# Load from ex6data3:
# You will have X, y in your environment
data = loadmat(scriptdir + '//ex6data3.mat')

X = data['X']
y = data['y'].ravel()
Xval = data['Xval']
yval = data['yval'].ravel()

# Plot training data
plotData(X, y)
plt.show()

input('Program paused. Press enter to continue.\n')

## ========== Part 7: Training SVM with RBF Kernel (Dataset 3) ==========

#  This is a different dataset that you can use to experiment with. Try
#  different values of C and sigma here.

# Try different SVM Parameters here
C, sigma = dataset3Params(X, y, Xval, yval)
# Train the SVM
model.svmTrain(X, y.astype(float), C, gaussianKernelLambda)
visualizeBoundary(X, y, model)
plt.show()
input('Program paused. Press enter to continue.\n')


import os
import cv2
import glob
import h5py
import json
import joblib
import argparse
import numpy as np
from tqdm import tqdm
import os.path as osp
import scipy.io as sio

import sys
sys.path.append('.')

from lib.models import spin
from lib.core.config import TCMR_DB_DIR, BASE_DATA_DIR
from lib.utils.utils import tqdm_enumerate
from lib.data_utils._kp_utils import convert_kps
from lib.data_utils._img_utils import get_bbox_from_kp2d
from lib.data_utils._feature_extractor import extract_features

from lib.data_utils._occ_utils import load_occluders
from lib.models.smpl import H36M_TO_J14, SMPL_MODEL_DIR, SMPL
from lib.utils.smooth_bbox import get_smooth_bbox_params, get_all_bbox_params
from lib.utils.vis import draw_skeleton

VIS_THRESH = 0.3

def cam2pixel(cam_coord, f, c):
    x = cam_coord[:,0] / cam_coord[:,2] * f[0] + c[0]
    y = cam_coord[:,1] / cam_coord[:,2] * f[1] + c[1]
    z = cam_coord[:,2]
    return np.stack((x,y,z),1)

def pixel2cam(pixel_coord, f, c):
    x = (pixel_coord[:,0] - c[0]) / f[0] * pixel_coord[:,2]
    y = (pixel_coord[:,1] - c[1]) / f[1] * pixel_coord[:,2]
    z = pixel_coord[:,2]
    return np.stack((x,y,z),1)

def world2cam(world_coord, R, t):
    cam_coord = np.dot(R, world_coord.transpose(1,0)).transpose(1,0) + t.reshape(1,3)
    return cam_coord

def cam2world(cam_coord, R, t):
    world_coord = np.dot(np.linalg.inv(R), (cam_coord - t.reshape(1,3)).transpose(1,0)).transpose(1,0)
    return world_coord

def read_data_train(dataset_path, set='train', debug=False):
    dataset = {
        'vid_name': [],
        'frame_id': [],
        'joints3D': [],
        'joints2D': [],
        'shape': [],
        'pose': [],
        'bbox': [],
        'img_name': [],
        'features': [],
    }

    # occluders = load_occluders('./data/VOC2012')

    model = spin.get_pretrained_hmr()

    if set == 'train':
        subjects = [1,5,6,7,8]
    else:
        subjects= [9, 11]
    for subject in subjects:
        annot_path = osp.join(dataset_path, 'annotations')
        # camera load
        with open(osp.join(annot_path, 'Human36M_subject' + str(subject) + '_camera.json'), 'r') as f:
            cameras = json.load(f)
        # joint coordinate load
        with open(osp.join(annot_path, 'Human36M_subject' + str(subject) + '_joint_3d.json'), 'r') as f:
            joints = json.load(f)
        # SMPL parameters obtained by NeuralAnnot will be released (https://arxiv.org/abs/2011.11232) after publication
        # # smpl parameter load
        # with open(osp.join(annot_path, 'Human36M_subject' + str(subject) + '_SMPL_NeuralAnnot.json'), 'r') as f:
        #     smpl_params = json.load(f)

        seq_list = sorted(glob.glob(dataset_path + f'/images/s_{subject:02d}*'))
        for seq in tqdm(seq_list):
            seq_name = seq.split('/')[-1]
            act = str(int(seq_name.split('_act_')[-1][0:2]))
            subact = str(int(seq_name.split('_subact_')[-1][0:2]))
            cam = str(int(seq_name.split('_ca_')[-1][0:2]))
            # if cam != '4':  # front camera (Table 6)
            #     continue
            print("seq name: ", seq)

            img_paths = sorted(glob.glob(seq + '/*.jpg'))
            num_frames = len(img_paths)
            if num_frames < 1:
                continue
            # camera parameter
            cam_param = cameras[cam]
            R, t, f, c = np.array(cam_param['R'], dtype=np.float32), np.array(cam_param['t'], dtype=np.float32), np.array(
                cam_param['f'], dtype=np.float32), np.array(cam_param['c'], dtype=np.float32)

            # img starts from index 1, and annot starts from index 0
            poses = np.zeros((num_frames, 72), dtype=np.float32)
            shapes = np.zeros((num_frames, 10), dtype=np.float32)
            j3ds = np.zeros((num_frames, 49, 3), dtype=np.float32)
            j2ds = np.zeros((num_frames, 49, 3), dtype=np.float32)

            for img_i in tqdm(range(num_frames)):
                # smpl_param = smpl_params[act][subact][str(img_i)][cam]
                # pose = np.array(smpl_param['pose'], dtype=np.float32)
                # shape = np.array(smpl_param['shape'], dtype=np.float32)

                joint_world = np.array(joints[act][subact][str(img_i)], dtype=np.float32)
                # match right, left
                match = [[1, 4], [2, 5], [3, 6]]
                for m in match:
                    l, r = m
                    joint_world[l], joint_world[r] = joint_world[r].copy(), joint_world[l].copy()
                joint_cam = world2cam(joint_world, R, t)
                joint_img = cam2pixel(joint_cam, f, c)

                j3d = convert_kps(joint_cam[None, :, :] / 1000, "h36m", "spin").reshape((-1, 3))
                j3d = j3d - j3d[39]  # 4 is the root

                joint_img[:, 2] = 1
                j2d = convert_kps(joint_img[None, :, :], "h36m", "spin").reshape((-1,3))

                # poses[img_i] = pose
                # shapes[img_i] = shape
                j3ds[img_i] = j3d
                j2ds[img_i] = j2d

                """
                import torch
                smpl = SMPL(SMPL_MODEL_DIR, batch_size=1, create_transl=False)

                p = torch.from_numpy(pose).float().reshape(1,-1,3)
                s = torch.from_numpy(shape).float().reshape(1,-1)
                J_regressor = torch.from_numpy(np.load(osp.join(TCMR_DATA_DIR, 'J_regressor_h36m.npy'))).float()
                output = smpl(betas=s, body_pose=p[:, 3:], global_orient=p[:, :3])
                vertices = output.vertices
                J_regressor_batch = J_regressor[None, :].expand(vertices.shape[0], -1, -1).to(vertices.device)
                temp_j3d = torch.matmul(J_regressor_batch, vertices) * 1000
                # temp_j3d = temp_j3d - temp_j3d[:, 0, :]
                temp_j3d = temp_j3d[0, H36M_TO_J14, :]

                gt_j3d = joint_cam - joint_cam[0, :]
                gt_j3d = gt_j3d[H36M_TO_J14, :]

                print("CHECK: ", (temp_j3d-gt_j3d))
                """

            bbox_params, time_pt1, time_pt2 = get_smooth_bbox_params(j2ds, vis_thresh=VIS_THRESH, sigma=8)
            # bbox_params, time_pt1, time_pt2 = get_all_bbox_params(j2ds, vis_thresh=VIS_THRESH)

            """
            img = cv2.imread(img_paths[0])
            temp = draw_skeleton(img, j2ds[0], dataset='spin', unnormalize=False, thickness=2)
            cv2.imshow('img', temp)
            cv2.waitKey(0)
            cv2.destroyAllWindows()
            cv2.waitKey(1)
            """

            # process bbox_params
            c_x = bbox_params[:, 0]
            c_y = bbox_params[:, 1]
            scale = bbox_params[:, 2]

            w = h = 150. / scale
            w = h = h * 0.9  # 1.1 for h36m_train_25fps_occ_db.pt
            bbox = np.vstack([c_x, c_y, w, h]).T

            img_paths_array = np.array(img_paths)[time_pt1:time_pt2][::2]
            bbox = bbox[::2]
            # subsample frame to 25 fps

            dataset['vid_name'].append(np.array([f'{seq}_{subject}'] * num_frames)[time_pt1:time_pt2][::2])
            dataset['frame_id'].append(np.arange(0, num_frames)[time_pt1:time_pt2][::2])
            dataset['joints3D'].append(j3ds[time_pt1:time_pt2][::2])
            dataset['joints2D'].append(j2ds[time_pt1:time_pt2][::2])
            dataset['shape'].append(shapes[time_pt1:time_pt2][::2])
            dataset['pose'].append(poses[time_pt1:time_pt2][::2])

            dataset['img_name'].append(img_paths_array)
            dataset['bbox'].append(bbox)

            features = extract_features(model, None, img_paths_array, bbox,
                                        kp_2d=j2ds[time_pt1:time_pt2][::2], debug=debug, dataset='h36m', scale=1.0)  # 1.2 for h36m_train_25fps_occ_db.pt

            dataset['features'].append(features)

    for k in dataset.keys():
        dataset[k] = np.concatenate(dataset[k])
        print(k, dataset[k].shape)

    return dataset

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--dir', type=str, help='dataset directory', default='data/h36m')
    parser.add_argument('--set', type=str, help='select train/test set', default='train')

    args = parser.parse_args()

    # import torch
    # torch.set_num_threads(8)

    dataset = read_data_train(args.dir, args.set)
    joblib.dump(dataset, osp.join(TCMR_DB_DIR, f'h36m_{args.set}_25fps_tight_db.pt'))  # h36m_train_25fps_occ_db.pt


import numpy as np
from scipy.optimize import fsolve
# from sympy import symbols, exp, solve
# from sympy import Symbol

# mu_m, lam, t = symbols('mu_m, lam, t')
v = - 1/3
A = 281.6 * 1000

INITIAL_WEIGHT = 10
ONE_YEAR_WEIGHT = 35

def rechards_equations(A, v, mu_m, lam, t):
    temp1 = (mu_m / A) * ((1+v) ** (1 + 1/v)) * (lam - t)
    temp2 = np.exp(temp1)
    temp3 = v * np.exp(1 + v) * temp2
    temp4 = A * (1 + temp3)**(-1/v)

    return temp4

def equations(p):
    mu_m, lam = p
    return (rechards_equations(A, v, mu_m, lam, 0) - INITIAL_WEIGHT, rechards_equations(A, v, mu_m, lam, 1)-ONE_YEAR_WEIGHT)

# print(equations((mu_m, lam)))

# (15.954190216465834, -11.657871577866423)

error = 10000
best_mu_0 = 0
best_lam_0 = 0
mu_res = 0
lam_res = 0
for mu_m_0 in range(1, 1000, 10):
    for lam_0 in range(1, 1000, 10):
        print('=======')
        print('mu_m_0 = ', mu_m_0)
        print('lam_0 = ', lam_0)
        mu_m, lam = fsolve(equations, (mu_m_0, lam_0))
        temp1 = rechards_equations(A, v, mu_m, lam, 0) - INITIAL_WEIGHT
        temp2 = rechards_equations(A, v, mu_m, lam, 1) - ONE_YEAR_WEIGHT
        print(temp1)
        print(temp2)
        if np.abs(temp1) + np.abs(temp2) < error:
            error = np.abs(temp1) + np.abs(temp2)
            best_lam_0 = lam_0
            best_mu_0 = mu_m_0
            mu_res = mu_m
            lam_res = lam
        print()

print('best_mu_0 = ', best_mu_0)
print('best_lam_0 = ', best_lam_0)
print('error = ', error)

print('mu_res = ', mu_res)
print('lam_res = ', lam_res)


import sys
import os
import subprocess
import json
from scipy import stats
from argparse import ArgumentParser

parser = ArgumentParser(description="Get Ideogram.js annotations for an SRR")
parser.add_argument("--acc", required=True, help="SRR accession")
args = parser.parse_args()

acc = args.acc
out = acc + "_counts"

os.chdir("counting")

subprocess.call(["python", "counter.py", "--inp", acc, "--out", out])

os.chdir("../formatter")
subprocess.call([
    "python", "formatter.py", "--type", "srr", "--lookup", "gene_lookup_GRCh37.tsv",
    "--inp", "../counting/" + out + "_norm.tsv", "--out", acc + ".json"
])


'''
    @author <NAME>.
    @version 2019.3.19
    @function
        Use 'rdpcap()' to read packet file captured by Wireshark
        Use 'pkt.time' to get packet receiving time in ms
        Use 'pkt.len' to get RadioTap dummy length
        Use 'len(pkt)' to get whole frame length
        Use dtw to calculate distance between time series
'''

from scapy.utils import rdpcap
from scapy.layers.dot11 import RadioTap
import numpy as np

all_packets = rdpcap('data/evil-twin-20190318-part1.pcapng')
client_packets = []
ap_packets = []

for i in range(0, len(all_packets)):
    time = all_packets[i].time
    length = len(all_packets[i]) - all_packets[i].len
    point = [time, length]
    if '50:5b:c2:d8:d1:21' in all_packets[i].addr2:
        client_packets.append(point)
    else:
        ap_packets.append(point)

x = np.array(client_packets)
y = np.array(ap_packets)

from scipy.spatial.distance import euclidean
from fastdtw import fastdtw

distance, path = fastdtw(x, y, dist=euclidean)
print(distance)


import numpy as np
from scipy import linalg
import pathlib, sys
file_path = pathlib.Path(__file__).parent.absolute()

from pressio4py import ode, rom
from pressio4py.apps.burgers1d import Burgers1d

gold = np.array(
  [5.0081542681376, 5.016629490569,
   5.025433912557,  5.0345792953115,
   5.0440827179355, 5.0539568295087,
   5.0642107801363, 5.074857742734,
   5.0859146000515, 5.0974001265619,
   5.1093302710968, 5.1217197481536,
   5.1345846667293, 5.1479436063682,
   5.1618137609004, 5.1762071980595,
   5.1911395190849, 5.2066322357211,
   5.222706587389,  5.2393822195142,
   5.2566784890019, 5.274617970535,
   5.2932246323729, 5.3125186218141,
   5.3325236627322, 5.3532729201416,
   5.3747971779128, 5.3971189932731,
   5.4202577535351, 5.4442348269811,
   5.469078757402,  5.4948202159561,
   5.5214859714822, 5.5491009348394,
   5.5776911098501, 5.6072849195866,
   5.6379131952825, 5.6696069037791,
   5.7023980878343, 5.7363239274031,
   5.7714263431002, 5.807744410524,
   5.8453128737429, 5.884168702448,
   5.9243510856362, 5.9658923478856,
   6.0088164545724, 6.0531503069487,
   6.0989210765093, 6.1461565470309])

#----------------------------------------
class OdeObserver:
  def __init__(self, fomRec):
    self.fomRec = fomRec

  def __call__(self, timeStep, time, state):
    fs = self.fomRec(state)
    assert(fs.shape[0]==50)

#----------------------------------------
def test_euler():
  meshSize = 50
  romSize  = 20
  Nsteps   = 3500
  dt       = 0.01

  # create fom object
  appObj = Burgers1d(meshSize)
  # reference state
  yRef = np.ones(meshSize)

  # I have to make phi a column-major array to ensure
  # pressio does not make a copy of this
  basisFile = str(file_path) + "/basis_euler.txt"
  phi = np.copy(np.loadtxt(basisFile), order='F')
  decoder = rom.Decoder(phi)

  # create rom state
  yRom = np.zeros(romSize)
  # create problem
  scheme = ode.stepscheme.ForwardEuler
  problem = rom.galerkin.DefaultExplicitProblem(scheme, appObj, decoder, yRom, yRef)

  fomRecon = problem.fomStateReconstructor()
  # the observer is called to monitor evolution of rom_state and
  # uses the reconstructor object to reconstruct FOM state
  myObs = OdeObserver(fomRecon)
  ode.advance_n_steps_and_observe(problem, yRom, 0., dt, Nsteps, myObs)

  # reconstruct full state at the end
  yFomFinal = fomRecon(yRom)
  print(yFomFinal)

  # check solution is right
  for y1,y2 in zip(gold, yFomFinal):
    assert( np.abs(y1-y2) < 1e-12)


<gh_stars>1-10
from astropy.io import fits
from astropy.convolution import convolve, Box1DKernel

import scipy as sp
import matplotlib
import matplotlib.pyplot as plt

import glob

'''
O 436
B 582
A 745
F 766
G 596
K 759
M 306
'''

'''
O 476, 8773, 9818
B 96, 378, 462, 489, 492
A 17, 114, 120, 136
F 52, 158
G 25, 27, 30, 85
K 61, 65
M 256, 291, 300
'''

i = [476, 378, 17, 158, 30, 61, 256]

c = ['O', 'B', 'A', 'F', 'G', 'K', 'M'][::-1]

loc = 5891

files = [glob.glob('/data2/cpb405/Training_2/*.fits')[j] for j in i][::-1]

fig, ax = plt.subplots(figsize = (5,0.9*5*sp.sqrt(2)))

ax.axvline(6565, c = 'r', alpha = 0.1)
ax.text(6600, 7, 'Ha', color = 'r')
ax.axvline(4862, c = 'r', alpha = 0.1)
ax.text(4900, 7, 'Hb', color = 'r')
ax.axvline(4342, c = 'r', alpha = 0.1)
ax.text(4400, 7, 'Hg', color = 'r')

for idx in range(len(files)):

    with fits.open(files[idx]) as hdulist:
        flux = hdulist[0].data[0]
        init = hdulist[0].header['COEFF0']
        disp = hdulist[0].header['COEFF1']
        CLS = hdulist[0].header['CLASS']
        SCLS = hdulist[0].header['SUBCLASS'][0]
        #print('{}, {}, {}'.format(idx, CLS, SCLS))

    wavelength = 10**sp.arange(init, init+disp*(len(flux)-0.9), disp)

    wavelength = wavelength[:-100]
    flux = flux[:-100]

    flux = sp.array(flux)

    wi = sp.searchsorted(wavelength, loc)

    #wi = -1

    flux = flux/sp.amax(flux)

    ax.plot(wavelength, flux + idx, label = c[idx], c = '#1f77b4')

    ax.annotate(c[idx], xy = (wavelength[sp.argmax(flux)]-75, idx+1.03))

ax.set_title('Stellar Spectra')
ax.set_xlabel('Wavelength \ Angstroms')
ax.set_ylabel('Normalised Flux')
plt.yticks([]," ")
#ax.set_yticklabels([])
#ax.get_yaxis().set_visible(False)
plt.tight_layout()
plt.savefig('MK.pdf')
plt.show()


from TCP_monitor import TCP_monitor
from OnlineHeart import OnlineHeart
from Silver import Silver
from LotteryResult import LotteryResult
from Tasks import Tasks
from connect import connect
from rafflehandler import Rafflehandler
import asyncio
from login import login
from printer import Printer
from statistics import Statistics
from bilibili import bilibili
import biliconsole
from pkLottery import PKLottery
from guardLottery import GuardLottery
from schedule import Schedule
import configloader
import os

fileDir = os.path.dirname(os.path.realpath('__file__'))
file_user = fileDir + "/conf/user.conf"
dic_user = configloader.load_user(file_user)

loop = asyncio.get_event_loop()
printer = Printer(dic_user['thoroughly_log']['on'])
bilibili()
Statistics()
rafflehandler = Rafflehandler()
biliconsole.Biliconsole()

task = OnlineHeart()  # 在线心跳
task1 = Silver()  # 领取银瓜子
task2 = Tasks()  # 获取每日包裹奖励，签到功能，领取每日任务奖励，应援团签到，过期礼物处理，银瓜子兑换硬币，硬币换瓜子，将当前佩戴的勋章亲密度送满，
task3 = LotteryResult()  # 广播抽奖检测
task4 = connect()  # 新的战疫分区直播间实际上没有弹幕区???
task5 = PKLottery()  # 大乱斗抽奖？
task6 = GuardLottery()  # 上船奖励？

tasks1 = [
    login().login_new()
]
loop.run_until_complete(asyncio.wait(tasks1))

# 任务
# import threading
# console_thread = threading.Thread(target=biliconsole.controler)
# console_thread.start()

tasks = [
    task.run(),
    task1.run(),
    task2.run(),
    biliconsole.Biliconsole().run(),#？？？？？
    task4.create(),
    task3.query(),
    rafflehandler.run(),
    task5.run(),
    task6.run()
]

if dic_user['monitoy_server']['on']:#监控服务器
    monitor = TCP_monitor()
    task_tcp_conn = monitor.connectServer(
        dic_user['monitoy_server']['host'], dic_user['monitoy_server']['port'], dic_user['monitoy_server']['key'])
    task_tcp_heart = monitor.HeartbeatLoop()
    tasks.append(task_tcp_conn)
    tasks.append(task_tcp_heart)

schedule = Schedule()
if dic_user['regular_sleep']['on']:
    tasks.append(schedule.run(dic_user['regular_sleep']['schedule']))
    Schedule().scheduled_sleep = True

tasks = list(map(asyncio.ensure_future, tasks))
loop.run_until_complete(asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION))
Printer().printer('\n'.join(map(repr, asyncio.Task.all_tasks())), "Info", "green")
for task in tasks:
    Printer().printer(repr(task._state), "Info", "green")
    if task._state == 'FINISHED':
        Printer().printer(f"Task err: {repr(task.exception())}", "Error", "red")
loop.close()

# console_thread.join()


"""This module contains auxiliary functions for the robust optimization problem."""
from functools import partial

import numpy as np
from scipy.optimize import fminbound

from robupy.tests.resources.pre_numba.checks import pre_numba_checks
from robupy.tests.resources.pre_numba.config import EPS_FLOAT
from robupy.tests.resources.pre_numba.config import MAX_FLOAT

def pre_numba_criterion_full(v, q, beta, lambda_):
    """This is the criterion function for ..."""
    pre_numba_checks("criterion_full_in", v, q, beta, lambda_)

    # We want to rule out an infinite logarithm.
    arg_ = np.clip(np.sum(q * np.exp(v / lambda_)), EPS_FLOAT, None)

    rslt = lambda_ * np.log(arg_) + lambda_ * beta

    pre_numba_checks("criterion_full_out", rslt)

    return rslt

def pre_numba_calculate_p(v, q, lambda_):
    """This function return the optimal ..."""
    pre_numba_checks("calculate_p_in", v, q, lambda_)

    p = q * np.clip(np.exp(v / lambda_), None, MAX_FLOAT)
    p = p / np.sum(p)

    pre_numba_checks("calculate_p_out", p)

    return p

def pre_numba_get_worst_case_probs(v, q, beta, is_cost=True):
    """This function return the worst case measure."""
    pre_numba_checks("get_worst_case_in", v, q, beta)

    if beta == 0.0:
        return q.copy()

    # We can use this function to determine the worst case if we pass in costs or if we pass in
    # utility.
    if not is_cost:
        v_intern = -np.array(v).copy()
    else:
        v_intern = np.array(v)

    # We scale the value function to avoid too large evaluations of the exponential function in
    # the calculate_p() function.
    v_scaled = v_intern / max(abs(v_intern))

    upper = np.clip(
        (max(v_scaled) - np.matmul(q, v_scaled)) / beta, 2 * EPS_FLOAT, None
    )
    lower = EPS_FLOAT

    criterion = partial(pre_numba_criterion_full, v_scaled, q, beta)

    rslt = fminbound(criterion, lower, upper, xtol=EPS_FLOAT, full_output=True)
    p = pre_numba_calculate_p(v_scaled, q, rslt[0])

    pre_numba_checks("get_worst_case_out", p, q, beta, rslt)

    return p

def pre_numba_get_worst_case_outcome(v, q, beta, is_cost=True):
    """This function calculates the worst case outcome."""
    pre_numba_checks("get_worst_case_outcome_in", v, q, beta)

    # We want to handle two cases explicitly. First we deal with the case that there is no
    # ambiguity in the transition probabilities. Second, we look at the case where the all mass
    # assigned to the worst-case realization is inside the feasible set.
    if beta == 0:
        return np.matmul(q, v)
    elif beta >= max(-np.log(q)):
        if is_cost:
            return max(v)
        else:
            return min(v)

    p = pre_numba_get_worst_case_probs(v, q, beta, is_cost)
    rslt = np.matmul(p, v)

    pre_numba_checks("get_worst_case_outcome_out", v, q, beta, is_cost, rslt)

    return rslt