#!/usr/bin/python
import pandas as pd
import numpy as np
from fuzzywuzzy import fuzz
import scipy
import multiprocessing as mp
from functools import partial
from sofa_print import *
from sofa_config import *
from sofa_common import *
from STree import *
from sklearn.cluster import KMeans

table_size = 0
iteration_timelines = []
iteration_table = []
blank_count = 0
def iter_profile(cfg, fields, df_gpu):
    elapsed_time = df_gpu['timestamp'].iat[len(df_gpu)-1] - df_gpu['timestamp'].iat[0]
    kernel_time = df_gpu['duration'].sum()
    payload = df_gpu['payload'].sum()
    return {'elapsed_time': elapsed_time, 'kernel_time': kernel_time, 'payload': payload}

def gpu_profile(logdir, cfg, df_gpu):
    total_kernel_time = 0.0
    total_gpu_time = 0.0
    if len(df_gpu) == 0:
        print_warning("No GPU traces to profile")
        return
    print_title("Task Time (MEMCPY included) for each Device (s)")
    grouped_df = df_gpu.groupby("deviceId")["duration"]
    total_tasktime = 0
    for key, item in grouped_df:
        print(("[%d]: %lf" % (int(float(key)), grouped_df.get_group(key).sum())))
        total_tasktime = total_tasktime + grouped_df.get_group(key).sum()
    n_devices = len(grouped_df)
    per_gpu_time = total_tasktime / n_devices
    print(("Averaged GPU time of devices: %.2lf" % per_gpu_time))

    print_title("Data Traffic (bidirection) for each Device (MB)")
    grouped_df = df_gpu.groupby("deviceId")["payload"]
    for key, item in grouped_df:
        print(("[%d]: %lf" % (key, grouped_df.get_group(key).sum() / 1000000.0)))

    grouped_df = df_gpu.groupby("copyKind")["duration"]
    for key, item in grouped_df:
        if key == 0:
            total_kernel_time = grouped_df.get_group(key).sum()

    print_title("All-reduce Time (s)")
    all_reduce_time = 0
    grouped_df = df_gpu.groupby("name")["duration"]
    for key, item in grouped_df:
        #print("[%s]: %lf" % (key, grouped_df.get_group(key).sum()))
        if key.find("AllReduce") != -1:
            all_reduce_time = all_reduce_time + grouped_df.get_group(key).sum()

    comm_profile(logdir, cfg, df_gpu)
    print(("MeasuredTotalKernelTime : %lf (s)" % total_kernel_time))

    print_title("Summary of Kernels")
    print(("MeasuredTotalKernelTime : %lf (s)" % total_kernel_time))
    print(("MeasuredAllReduceTime : %lf (s)" % all_reduce_time))
    get_top_k_events(df_gpu, 10)

def duration_sum(df_gpu):
    duration_us = int(sum(df_gpu[ df_gpu['duration'] > 1e-5 ]['duration'])/1e-6)
    return duration_us
def kernel_count(df_gpu):
    count = df_gpu[ df_gpu['duration'] > 1e-5 ].count()[0]
    return count

def main_string_generate_v1(df_gpu):
    wid = 0
    name_pre = ''
    name_table = {'':0}
    id_seq = []
    for i in range(len(df_gpu)):
        trace  = df_gpu.iloc[i]
        value = name_table.get(trace['name'])
        if value == None:
            wid = wid + 1
            name_table[trace['name']] = wid
            value = wid
        id_seq.append(str(value))
        #id_seq.append(str(int(trace['duration']/1e-5)))
    main_string = ','.join(id_seq)
    return main_string

def main_string_generate_v2(df_gpu):
    tick = 0
    tick_next = 0
    tick_begin = 0
    tick_end = int( round( ( t_df_end - t_df_begin ) / time_interval ) )
    event_names = get_top_k_events(df_gpu,10)
    events = event_names[:]
    pattern_table = [] #np.array(np.zeros((tick_end,len(event_names)+2),dtype=int))

    iteration_pattern_count = 1
    while(tick < tick_end):
        tick_next = tick + 1
        tick_event = (df_gpu.loc[:, "timestamp"] - t_df_begin) / time_interval
        #slice trace to block by time interval
        df_block = df_gpu[( tick_event < tick_next) & ( tick_event >= tick)]
        #Create vector and count
        vector = np.array(np.zeros(len(event_names)+2))
        #print(df_block.shape)
        if df_block.shape[0] > 0:
            for i in range(len(event_names)):
                count = event_count('name', event_names[i], df_block)
                vector[i] = count

            vector[len(event_names)] = 0;#duration_sum(df_block)
            vector[len(event_names)+1] = 0;#kernel_count(df_block)
            matched_index = pattern_matching(pattern_table, events, vector, threshold)
            if matched_index != -1:
                iteration_timelines.append(str(matched_index))
            else:
                iteration_timelines.append('0')
        else:
            iteration_timelines.append('0')
        tick += 1
    main_string = ",".join(iteration_timelines)

def get_top_k_events(df, topk):
    topk_events=[]
    print("Top %d Events:"%topk)
    gby = df.groupby(['name'])
    df_agg = gby.aggregate(np.sum)
    df_agg_sorted = df_agg.sort_values(by=['duration'],ascending=False)
    #memcpy = ['copyKind_1_','copyKind_2_','copyKind_8_']
    eventName = df_agg_sorted[df_agg_sorted.columns[0:0]].head(topk).index.values.tolist()
    #eventName.extend(memcpy)

    for i in range(len(eventName)):
        print('[%d] %s'%(i,eventName[i]))
    return eventName

def get_memcpyHtoD(df):

    df_sorted = df.sort_values(by=['payload'],ascending=False)
    HtoD = df_sorted[df_sorted['name'].str.contains('copyKind_1_', na=False)]['name'].head(1).values

    return HtoD

def select_pattern(candidate_pattern):
    #print("candidate_pattern:", candidate_pattern)
    candidate_pattern_filtered = []
    for cp in candidate_pattern:
        if len(cp)>1:
            if cp.count(cp[0])+cp.count(cp[1]) != len(cp):
               # if cp.count('0') < 2:
                candidate_pattern_filtered.append(cp)
                    #print('filtered cp = '+cp)

    if len(candidate_pattern_filtered) > 0:
        pattern = max(candidate_pattern_filtered, key = len)
        print("pattern selected:", pattern)
        return pattern
    else:
        print_warning('all candidate patterns are not selected')

def iter_detect(logdir, cfg, df_gpu, time_interval, threshold, iteration_times):
    global iteration_timelines, iteration_table, blank_count, \
            iteration_table
    df_gpu.to_csv('dfgpu.csv')
    t_df_begin = df_gpu.iloc[0]['timestamp']
    t_df_end = df_gpu.iloc[-1]['timestamp']
    candidate_pattern=[]
    main_string = main_string_generate_v1(df_gpu)
    st = STree(main_string)
    st.find_repeat_pattern(candidate_pattern, iteration_times)
    #print("main_string:", main_string)
    pattern = select_pattern(candidate_pattern)
    if pattern:
        print('original string length of main_string = %d' % len(main_string))
        wid_seq = main_string.split(',')
        pat_seq = pattern.split(',')
        total_length = num_wids = len(wid_seq)
        print('context string length = %d' % num_wids )
        print('pattern length = %d' % len(pat_seq) )
        #block_size = pattern.count(',') + 1
        print(pat_seq)
        block_size = len(pat_seq)

        block_begin = 0
        block_end = block_size
        step = 1
        #step = int(block_size/8)
        iteration_count = 0
        fuzzyRatioTable = []
        while block_begin <= (total_length - block_size):
            blockString = ",".join(wid_seq[block_begin:block_end])
            #print('A(%d): %s' % (len(blockString),blockString))
            #print('B(%d): %s' % (len(pattern), pattern))
            fuzz_ratio = fuzz.ratio(blockString,pattern)

            fuzzyRatioTable.append((block_begin,fuzz_ratio))
            #print(fuzz_ratio)
            block_begin = block_begin + step
            block_end = block_begin + block_size
        #find largest fuzzy ratio n blocks (n = iteration_times)
        ind = []
        begTable = []
        endTable = []
        beg = 0
        end = 0
        fw_threshold = 80
        for i in range(len(fuzzyRatioTable)):
            if fuzzyRatioTable[i][1] > fw_threshold:
                ind.append(fuzzyRatioTable[i][0])

        if len(ind) > 0:
            #plot event by event
            for i in ind:
                iteration_table.append((df_gpu.iloc[i]['timestamp'],df_gpu.iloc[i+block_size-1]['timestamp']))
            begTable.append(ind[0])
            beg = ind[0]
            end = ind[0] + block_size
            #plot window by window
            #for index in ind:
            #    iteration_table.append((float(index*step) * time_interval + t_df_begin, float(index*step + block_size) * time_interval + t_df_begin))
            #begTable.append(ind[0])
            #beg = ind[0]
            #end = ind[0] + block_size
        else:
            print_warning("No matched strings by fuzzywuzzy of threshold %d."%fw_threshold)
    else:
        print('No iteration patterns detected.')

def event_count(column, eventName, df):
        #get the count of rows that contain eventName
        count = 0
        for i in range(df.shape[0]):
            if df.iloc[i][column].find(eventName) != -1:
                count = count + 1
        return count

def pattern_matching(pattern_table, events, vector, threshold):
    global table_size
    if np.sum(vector) == 0:
        return -1
    for i in range(table_size):
        pvector = pattern_table[i]
        if similar(pvector, vector, threshold):
            return i
    pattern_table.append(vector)

    return table_size

def similar(a, b, threshold):
    result = 1.0 - scipy.spatial.distance.cosine(a, b)
    if result > threshold:
        langa = np.linalg.norm(a)
        langb = np.linalg.norm(b) + 1e-9
        ratio = np.abs(langa/langb)
        if ratio < 1.25 and ratio > 0.75:
            return True
    return False

def traces_to_json(path):
    global iteration_table
    sofa_traces = None
    with open(path, 'r') as f:
        sofa_traces = f.readlines()
    with open(path, 'w') as f:
        f.writelines([line for line in sofa_traces[:-1]])
        f.write("\n\n")
        f.write("iteration_detection_begin = ")
        f.write('{"name": "iteration_begin", "color": "rgba(241,156,162,1)", "data": [')
        for (IT_beg, IT_end) in iteration_table:
            #print("begin:%f end:%f duration:%f"%(IT_beg, IT_end, IT_end-IT_beg))
            for i in range(10):
                dot_y = np.power(10,float(i-5))
                f.write('{"name": "iteration_begin", "x": ' + str(IT_beg) + ', "y": ' + str(dot_y) + '}, ')
        f.write(']}\n')
        f.write("iteration_detection_end = ")
        f.write('{"name": "iteration_end", "color": "rgba(241,15,162,1)", "data": [')
        for (IT_beg, IT_end) in iteration_table:
            #print("begin:%f end:%f duration:%f"%(IT_beg, IT_end, IT_end-IT_beg))
            for i in range(10):
                dot_y = np.power(10,float(i-5))
                f.write('{"name": "iteration_end", "x": ' + str(IT_end) + ', "y": ' + str(dot_y) + '}, ')
        f.write(']}\n')

        for s in sofa_traces[-1].split(','):
            if s.find(']') == -1:
                f.write(s+',')
        f.write("iteration_detection_begin,iteration_detection_end]")

def trace_timeline(path):
    global iteration_timelines, blank_count
    i = 0
    with open(path, 'w') as f:
        for x in range(blank_count):
            f.write(str(i) + "," + str(0))
            i += 1
        for index in iteration_timelines:
            f.write(str(i) + "," + str(index))
            i += 1

def sofa_aisi(logdir, cfg, df_cpu, df_gpu):
    global iteration_table
    a = 0
    try:
        iter_detect(logdir, cfg, df_gpu, 0.01, 0.8, cfg.iterations)
        traces_to_json(logdir + 'report.js')
        xlist = []
        for t in iteration_table:
            xlist.append((t[0],0))
        X = np.array(xlist)
        kmeans = KMeans(n_clusters=cfg.iterations, random_state=0).fit(X)
        times = []
        for c in kmeans.cluster_centers_:
            times.append(c[0])
        times = np.sort(times)
        trace_timeline(logdir + 'iteration_timeline.txt')

        iter_summary_fields = ['elapsed_time', 'kernel_time', 'fw_time', 'bw_time', 'comm_time', 'cpu_time', 'bw_h2d', 'bw_d2h', 'bw_p2p', 'bw_d2h_big', 'bw_d2h_big', 'bw_p2p_big', 'payload']
        iter_list = []
        for i in range(1,len(times)):
            #print_title("Perormance analyze of iteration-%d"%(i))
            overlapness = 0.0
            cond1 = (df_gpu['timestamp'] >= times[i-1])
            cond2 = (df_gpu['timestamp'] <  times[i])
            df_gpu_iteration = df_gpu[ cond1 & cond2]
            df_cpu_iteration = df_cpu[ cond1 & cond2]
            if len(df_gpu_iteration) > 0:
                iter_list.append(iter_profile(cfg, iter_summary_fields, df_gpu_iteration))
        iter_summary = pd.DataFrame( iter_list, columns=iter_summary_fields )

        print_title('Per-iteration Performance Summary')
        mean_step_time = iter_summary.loc[1:,'elapsed_time'].mean()
        mean_kernel_time = iter_summary.loc[1:,'kernel_time'].mean()
        print("Elapsed time of initial iteration (s): ", iter_summary.loc[0,'elapsed_time'])
        print("Averaged elapsed time of iterations excluding initial one (s): ", iter_summary.loc[1:,'elapsed_time'].mean())
        print("Total CUDA kernel time (s): ", iter_summary.loc[1:,'kernel_time'].mean())
        print("Total CUDA  payload (B): ", iter_summary.loc[1:,'payload'].mean())
        print_title('Performance Optimization Hints')
        if mean_kernel_time / mean_step_time > 0.8:
            print("The profiled program is a compute-bound workload, try increasing # of GPUs to improve throughput")
        elif mean_kernel_time / mean_step_time < 0.5:
            print("The profiled program is a communication-bound workload, %d bytes are monitored on PCIe bus"%payload)
        print('\n\n')
    except IOError:
        print_warning(
            "gputrace.csv is not found. If there is no need to profile GPU, just ignore it.")


import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import random
import numpy as np
import seaborn as sns
from extract_arrays import extract_arrays, get_refseq

def run(args):
    sns.set_style('ticks')

    # read in data
    refseq = get_refseq(args.ref)
    arrays = extract_arrays(args.bam, refseq, copy_filter='hard').arrays

    switch_dict = {1:0, 0:1}
    # dictionary of "switch probabilities" (i.e., error rates)
    # for T-C or C-T changes at K3L-His47Arg
    switch_probs = { 'r7':  {1: 0.015, 0: 0.023},
                     'r9':  {1: 0.024, 0: 0.023},
                     'r94': {1: 0.026, 0: 0.005} }
    # use the switch probabilities for the particular sequencing
    # chemistry used to make the input BAM
    switch_probs = switch_probs[args.chemistry]
    all_homogeneous = []
    true_mixed, true_homogeneous = 0, 0
    # loop over every array in the population,
    # and convert mixed arrays into homogenous arrays
    for array in arrays:
        # single-copy arrays can't be "mixed"
        if len(array) < 2: continue
        # if the array is mixed, convert it to a homogenous array
        if len(set(array)) == 2:
            true_mixed += 1
            # if the array contains equal numbers of K3L-His47Arg
            # and K3L-WT copies, randomly decide whether it gets converted
            # into a homogenous WT or His47Arg array
            if array.count(1) - array.count(0) == 0:
                switch = random.random() < 0.5
                if switch:
                    all_homogeneous.append([0 for a in range(len(array))])
                else:
                    all_homogeneous.append([1 for a in range(len(array))])
            # else, assign the array to the homogeneity of its
            # predominant allele
            elif array.count(1) - array.count(0) > 0:
                all_homogeneous.append([1 for a in range(len(array))])
            else:
                all_homogeneous.append([0 for a in range(len(array))])
        # if the array is already homogeneous, keep it as such
        else:
            true_homogeneous += 1
            all_homogeneous.append(array)
    simulation_results = []
    sims = 1000
    for sim in range(sims):
        # loop over every one of the converted homogenous arrays,
        # and introduce "switches" (i.e., T-C or C-T) changes to
        # mimic sequencing errors at the K3L-His47Arg site
        new_arrays = []
        for array in all_homogeneous:
            new_array = []
            for copy in array:
                if random.random() < switch_probs[copy]:
                    new_array.append(switch_dict[copy])
                else: new_array.append(copy)
            new_arrays.append(new_array)
        # record the numbers of mixed and homogenous arrays in the
        # randomly "switched" population of arrays
        sim_mixed, sim_homogeneous = 0, 0
        for array in new_arrays:
            if len(set(array)) == 2:
                sim_mixed += 1
            else: sim_homogeneous += 1
        # record the fraction of mixed arrays in the simulated population
        simulation_results.append(sim_mixed / float(sim_mixed + sim_homogeneous))

    sns.kdeplot(np.array(simulation_results), label="Expected fraction of mixed arrays")
    plt.axvline(true_mixed / float(true_mixed + true_homogeneous), color='r',
                    label="Observed fraction of mixed arrays", ls='--')
    plt.ylabel("Kernel density")
    plt.xlabel("Fraction of mixed arrays in population")
    plt.legend()
    sns.despine(trim=True)
    if args.png:
        plt.savefig(args.o + '_' + args.chemistry + '.png', dpi=200, bbox_inches='tight')
    else:
        plt.savefig(args.o + '_' + args.chemistry + '.eps', bbox_inches='tight')

def main(argv):
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("--bam", required=True, help='Path to sorted BAM file.')
    p.add_argument("--ref", required=True, help='Path to FASTA reference genome.')
    p.add_argument("--chemistry", required=True, help='Flowcell chemistry used to generate reads in BAM. One of "r94", "r9", or "r7."')
    p.add_argument("-o", help='Name of output plot.', default='sim-errors')
    p.add_argument("-png", help='Output as PNG.', action='store_true')
    run(p.parse_args(argv))

if __name__ == "__main__":
    import sys
    main(sys.argv[1:])


# Copyright (c) Facebook, Inc. and its affiliates
# Copyright (c) MTRF authors

import collections
from enum import Enum
from gym import utils
import numpy as np
import os
from pathlib import Path
import pickle
from typing import Any, Dict, Optional, NewType, Sequence, Union, Tuple

from r3l.r3l_envs.base_env.base_env_deltapos import SawyerDhandDeltaBaseEnv
from r3l.utils.quatmath import quat2euler, euler2quat
from r3l.utils.circle_math import circle_distance, circle_distance_mod
from r3l.utils.range import get_range_from_params
from r3l.robot.object import ObjectState
from r3l.robot.default_configs import (
    DEFAULT_DHAND_ROBOT_CONFIG,
    DEFAULT_SAWYER_ROBOT_CONFIG,
    DEFAULT_OBJECT_CONFIG,
)

class ObjectType(Enum):
    Valve3 = 1
    Rod = 2
    Dodecahedron = 3
    DodecahedronBasket = 4
    DodecahedronBulb = 5
    Dumbbell = 6
    Mug = 7
    Pipe = 8

DHAND_INHAND_ASSET_PATH = Path(__file__).absolute().parent / "assets"
OBJECT_TYPE_TO_MODEL_PATH = {
    ObjectType.Valve3: str(DHAND_INHAND_ASSET_PATH / "sawyer_dhand_inhand_valve.xml"),
    ObjectType.Rod: str(DHAND_INHAND_ASSET_PATH / "sawyer_dhand_inhand_rod.xml"),
    ObjectType.Dodecahedron: str(DHAND_INHAND_ASSET_PATH / "sawyer_dhand_inhand_dodecahedron.xml"),
    ObjectType.DodecahedronBasket: str(DHAND_INHAND_ASSET_PATH / "sawyer_dhand_inhand_dodecahedron_basket.xml"),
    ObjectType.DodecahedronBulb: str(DHAND_INHAND_ASSET_PATH / "sawyer_dhand_inhand_dodecahedron_bulb.xml"),
    ObjectType.Dumbbell: str(DHAND_INHAND_ASSET_PATH / "sawyer_dhand_inhand_dumbbell.xml"),
    ObjectType.Mug: str(DHAND_INHAND_ASSET_PATH / "sawyer_dhand_inhand_mug.xml"),
    ObjectType.Pipe: str(DHAND_INHAND_ASSET_PATH / "sawyer_dhand_inhand_pipe.xml"),
}

class SawyerDhandInHandObjectBaseEnv(SawyerDhandDeltaBaseEnv):
    VALVE_DEFAULT_OBSERVATION_KEYS = (
        "dhand_qpos", # "sawyer_arm_qpos",
        "dhand_qvel", # "sawyer_arm_qvel",
        "mocap_pos", "mocap_euler",

        # Object info
        "object_xyz", # Not including object euler and qvel
        # object "z orientation" accounts for flips
        "object_top_angle_cos",
        "object_top_angle_sin",
        "target_xyz",
        # target "z orientation" accounting for flips
        "target_top_angle_cos",
        "target_top_angle_sin",

        "object_to_hand_xyz",
        # These all account for site symmetry
        "object_site1_to_target_site1_xyz_err",
        "object_site2_to_target_site2_xyz_err",
        "object_site3_to_target_site3_xyz_err"
    )

    ROD_DEFAULT_OBSERVATION_KEYS = (
        "dhand_qpos", "sawyer_arm_qpos",
        "mocap_pos", "mocap_euler",

        # Object info
        "object_xyz", # Not including object euler and qvel
        "object_z_orientation_cos", "object_z_orientation_sin",

        "target_xyz",
        "target_z_orientation_cos", "target_z_orientation_sin",

        "object_to_hand_xyz",
    )

    PIPE_DEFAULT_OBSERVATION_KEYS = (
        "dhand_qpos", "sawyer_arm_qpos",
        "mocap_pos", "mocap_euler",

        # Object info
        "object_xyz", # Not including object euler and qvel
        # "object_z_orientation_cos", "object_z_orientation_sin",

        "target_xyz",
        # "target_z_orientation_cos", "target_z_orientation_sin",

        "object_to_hand_xyz",
        "object_normal_to_target_normal_err",
        "object_parallel_to_target_parallel_err",
    )

    DODECAHEDRON_DEFAULT_OBSERVATION_KEYS = (
        "dhand_qpos", "sawyer_arm_qpos",
        "mocap_pos", "mocap_euler",

        # Object info
        "object_xyz", # Not including object euler and qvel
        "target_xyz",

        "object_to_hand_xyz",
        "object_to_target_xyz",
    )

    def __init__(
            self,
            model_path=None,
            sim=None,
            object_type: ObjectType = ObjectType.Rod,
            observation_keys: list = ROD_DEFAULT_OBSERVATION_KEYS,
            collect_bin_counts: bool = False,
            n_bins: int = 32,
            reset_policy_directories: list = [],
            reset_state_pkl_path: str = None,
            reset_offset: np.ndarray = None,
            save_reset_videos: bool = False,
            **kwargs
    ):
        self.object_type = object_type

        # Discrete count rewards
        self.x_range = (0.72 - 0.4, 0.72 + 0.4)
        self.y_range = (0.15 - 0.4, 0.15 + 0.4)
        self._collect_bin_counts = collect_bin_counts
        self.n_bins = n_bins
        self.x_bins = np.linspace(self.x_range[0], self.x_range[1], self.n_bins)
        self.y_bins = np.linspace(self.y_range[0], self.y_range[1], self.n_bins)
        self.bin_counts = np.ones((self.n_bins + 1, self.n_bins + 1))

        # Loaded policy resets
        self._reset_dist = np.zeros(1)
        self._reset_policies = []
        self._wrapped_reset_policies = []
        self._reset_horizons = []
        self._reset_envs = []
        self._reset_imgs = []
        self._save_reset_videos = save_reset_videos

        # Track goal position
        self._last_target_xyz = np.zeros(3)
        self._last_target_euler = np.zeros(3)
        self._pose_target = 0.7 * np.array([
            1, 1, 1, 1,
            -1, -1, -1, -1,
            1, 1, 1, 1,
            -1, -1, -1, -1,
        ])

        # Initialize observation keys depending on object type
        if self.object_type in (
                ObjectType.Dodecahedron,
                ObjectType.DodecahedronBasket,
                ObjectType.DodecahedronBulb):
            observation_keys = self.DODECAHEDRON_DEFAULT_OBSERVATION_KEYS
        elif self.object_type in (ObjectType.Valve3,):
            observation_keys = self.VALVE_DEFAULT_OBSERVATION_KEYS
        elif self.object_type in (ObjectType.Rod,):
            observation_keys = self.ROD_DEFAULT_OBSERVATION_KEYS
        elif self.object_type in (ObjectType.Pipe,):
            observation_keys = self.PIPE_DEFAULT_OBSERVATION_KEYS

        try:
            if self.script_motion:
                observation_keys = observation_keys + ("step_counter",)
        except AttributeError:
            pass

        # Init the environment + sim
        env_params = dict(
            observation_keys=observation_keys,
            dhand_config=DEFAULT_DHAND_ROBOT_CONFIG,
            sawyer_config=DEFAULT_SAWYER_ROBOT_CONFIG,
            object_config=DEFAULT_OBJECT_CONFIG,
        )
        env_params.update(kwargs)
        if sim:
            super().__init__(sim=sim, **env_params)
        else:
            assert object_type in OBJECT_TYPE_TO_MODEL_PATH
            object_model_path = OBJECT_TYPE_TO_MODEL_PATH.get(object_type)
            # If a model_path is already specified, make sure that it matches the object type.
            if model_path is not None:
                assert model_path == object_model_path
            else:
                model_path = object_model_path
            super().__init__(
                model_path=model_path,
                observation_keys=observation_keys,
                **kwargs
            )

            # Boost damping in the arms to avoid high accelerations
            # TODO(justinvyu): Is this also the case on the hardware?
            self.sim.model.dof_damping[:7] *= 10
            self.sim.model.dof_armature[:7] *= 10

        # Get model ids
        self.target_bid = self.sim.model.body_name2id('target')
        self.grasp_left = self.sim.model.site_name2id('grasp_left')
        self.grasp_right = self.sim.model.site_name2id('grasp_right')
        self.grasp_id = self.sim.model.site_name2id('grasp')
        if self.object_type in (ObjectType.Rod, ObjectType.Pipe):
            self.rod_center = self.sim.model.site_name2id('rod_center')
            self.rod_left = self.sim.model.site_name2id('rod_left')
            self.rod_right = self.sim.model.site_name2id('rod_right')
            self.rod_normal = self.sim.model.site_name2id('rod_normal')
            self.rod_parallel = self.sim.model.site_name2id('rod_parallel')

            self.target_center = self.sim.model.site_name2id('target_center')
            self.target_left = self.sim.model.site_name2id('target_left')
            self.target_right = self.sim.model.site_name2id('target_right')
            self.target_normal = self.sim.model.site_name2id('target_normal')
            self.target_parallel = self.sim.model.site_name2id('target_parallel')
        elif self.object_type == ObjectType.Valve3:
            self.valve_1 = self.sim.model.site_name2id('valve3_1')
            self.valve_2 = self.sim.model.site_name2id('valve3_2')
            self.valve_3 = self.sim.model.site_name2id('valve3_3')
            self.target_1 = self.sim.model.site_name2id('valve3_target_1')
            self.target_2 = self.sim.model.site_name2id('valve3_target_2')
            self.target_3 = self.sim.model.site_name2id('valve3_target_3')

        # Track goal position
        self._last_target_xyz = self.sim.data.get_body_xpos("target").copy()
        self._last_target_euler = quat2euler(self.sim.data.get_body_xquat("target").copy())

        # Object offset = the offset relative to the object that the hand is reset to
        self.obj_offset = reset_offset
        if self.obj_offset is None:
            self.obj_offset = np.array([-0.15, 0, 0.175])

        if reset_state_pkl_path:
            self._use_reset_state = True
            with open(reset_state_pkl_path, 'rb') as f:
                reset_states = pickle.load(f)
                if isinstance(reset_states, dict):
                    assert "sawyer" in reset_states
                    assert "dhand" in reset_states
                    assert "object" in reset_states
                    self._reset_sawyer_state = reset_states["sawyer"]
                    self._reset_dhand_state = reset_states["dhand"]
                    self._reset_object_state = reset_states["object"]
                else:
                    self._reset_sawyer_state = reset_states[0]
                    self._reset_dhand_state = reset_states[1]
                    self._reset_object_state = reset_states[2]
        else:
            self._use_reset_state = False
            reset_policy_directories = reset_policy_directories or []
            for directory in reset_policy_directories:
                self._load_policy(directory)

    def _load_policy(self, saved_policy_dir, phase_num=None):
        try:
            import softlearning
        except ModuleNotFoundError:
            print('Error: softlearning package not found. Unable to load reset policy.')
            return None
        from softlearning.policies.utils import get_policy_from_variant
        from softlearning.environments.utils import get_environment_from_params

        policy_dir = Path(saved_policy_dir)
        variant_path = policy_dir / "params.pkl"
        policy_weights_path = policy_dir / "policy_params.pkl"
        assert variant_path.exists() and policy_weights_path.exists(), (
            "Error loading policy and variant: we expect a file at:\n"
            + str(variant_path) + "\n" + str(policy_weights_path)
        )
        with open(variant_path, 'rb') as f:
            variant = pickle.load(f)
        with open(policy_weights_path, 'rb') as f:
            policy_weights = pickle.load(f)

        if phase_num is not None:
            policy_weights = policy_weights[phase_num]

        from softlearning.environments.adapters.gym_adapter import GymAdapter
        from softlearning.models.utils import flatten_input_structure
        reset_env_params = variant["environment_params"]["evaluation"]
        reset_env_kwargs = reset_env_params.pop("kwargs", {})
        reset_env_kwargs["sim"] = self.sim
        reset_env_params["kwargs"] = reset_env_kwargs
        reset_env = get_environment_from_params(reset_env_params)
        # Create environment as softlearning expects it for policy initialization
        reset_env = GymAdapter(None, None, env=reset_env)
        self._reset_envs.append(reset_env)

        reset_policy = get_policy_from_variant(variant, reset_env)
        reset_policy.set_weights(policy_weights)
        self._reset_policies.append(reset_policy)

        # Save some time by taking a max reset horizon of 50 steps
        horizon = min(
            50,
            variant.get('sampler_params', {}).get('kwargs', {}).get('max_path_length', 50)
        )
        self._reset_horizons.append(horizon)

        def wrapped_policy(obs_dict):
            feed_dict = {
                key: obs_dict[key][None, ...]
                for key in reset_policy.observation_keys
            }
            observation = flatten_input_structure(feed_dict)
            with reset_policy.set_deterministic(True):
                action = reset_policy.actions_np(observation)[0]
            return action
        self._wrapped_reset_policies.append(wrapped_policy)
        return wrapped_policy

    def run_reset_policy(self):
        if self._use_reset_state:
            # Reset angles so that mocap doesn't go crazy after many iters
            self.sawyer_robot.reset(command_angles=True)
            self.sawyer_robot.set_state(self._reset_sawyer_state, command_angles=False)
            self.do_simulation(self.act_mid - 0.75 * self.act_rng, 300) # wait for arm to be stable
            self.dhand_robot.set_state(self._reset_dhand_state)
            self.object.set_state(self._reset_object_state)
        else:
            self._reset_imgs = []
            for reset_policy, reset_env, horizon in zip(self._wrapped_reset_policies, self._reset_envs, self._reset_horizons):
                reset_env.reset_robot()
                for _ in range(horizon):
                    action = reset_policy(reset_env.get_obs_dict())
                    reset_env.step(action)
                    if self._save_reset_videos:
                        self._reset_imgs.append(self.render(width=480, height=480, mode="rgb_array"))
                self._reset_dist = self.get_obs_dict()["object_to_target_xy_distance"]

    def get_obs_dict(self) -> Dict[str, np.ndarray]:
        obs_dict = super().get_obs_dict()
        # Object data
        object_qpos, object_qvel = obs_dict["object_qpos"], obs_dict["object_qvel"]
        object_xyz, object_quat = object_qpos[:3], object_qpos[3:]
        object_euler = quat2euler(object_quat)
        object_z_orientation = object_euler[2]

        # Target data
        # NOTE: Targets need to be set in the OptiTrack coordinate space for hardware
        target_xyz = self._last_target_xyz.copy()
        target_euler = self._last_target_euler.copy()
        target_z_orientation = target_euler[2]

        # TODO(justinvyu): Add an offset to the end of the Sawyer xyz
        grasp_xyz = self.sim.data.get_site_xpos("grasp").copy()
        if self.is_hardware:
            grasp_xyz = obs_dict["mocap_pos"][:3]

        relative_xyz = object_xyz - target_xyz
        circle_dist = circle_distance(object_euler, target_euler)
        hand_to_obj_dist = np.linalg.norm(grasp_xyz - object_xyz)
        hand_to_target_dist = np.linalg.norm(grasp_xyz - target_xyz)

        x, y, z = object_xyz
        x_d, y_d = np.digitize(x, self.x_bins), np.digitize(y, self.y_bins)

        hand_pose_dist = np.linalg.norm(obs_dict["dhand_qpos"] - self._pose_target)

        circle_dist_mod_180 = np.array([
            circle_distance_mod(object_z_orientation, target_z_orientation, mod=np.pi)
        ])

        obs_dict.update(collections.OrderedDict((
            # Object info
            ("object_xyz", object_xyz),
            ("object_xy_discrete", np.array([x_d, y_d])),
            ("object_quat", object_quat),
            ("object_euler", object_euler),
            ("object_z_orientation", object_z_orientation),
            ("object_z_orientation_cos", np.array([np.cos(object_z_orientation)])),
            ("object_z_orientation_sin", np.array([np.sin(object_z_orientation)])),
            ("object_qvel", object_qvel),

            # Target info
            ("target_xyz", target_xyz),
            ("target_euler", target_euler),
            ("target_quat", euler2quat(target_euler)),
            ("target_z_orientation", target_z_orientation),
            ("target_z_orientation_cos", np.array([np.cos(target_z_orientation)])),
            ("target_z_orientation_sin", np.array([np.sin(target_z_orientation)])),

            # Distances
            ("object_to_target_xyz_distance", np.linalg.norm(relative_xyz)),
            ("object_to_target_xy_distance", np.linalg.norm(relative_xyz[:2])),
            ("object_to_target_x_distance", np.abs(relative_xyz[0])),
            ("object_to_target_y_distance", np.abs(relative_xyz[1])),
            ("object_to_target_z_distance", np.abs(relative_xyz[2])),
            ("object_to_target_circle_distances", circle_dist),
            ("object_to_target_circle_distance", np.linalg.norm(circle_dist)),
            ("object_to_target_mod_120_circle_distance", np.array([0])),
            ("object_to_target_mod_180_circle_distance", circle_dist_mod_180),
            ("object_to_hand_xyz_distance", hand_to_obj_dist),
            ("target_to_hand_xyz_distance", hand_to_target_dist),
            ("pose_dist", hand_pose_dist),

            # Relative vectors
            ("object_to_target_xyz", relative_xyz),
            ("object_to_hand_xyz", grasp_xyz - object_xyz),
            ("grasp_xyz", grasp_xyz),
        )))
        if self._wrapped_reset_policies:
            obs_dict["reset_policy_xy_distance"] = self._reset_dist.copy()

        # Add these as placeholders (to be populated for each object specifically)

        # Valve specific observations
        obs_dict["object_site1_to_target_site1_xyz_err"] = np.zeros(3)
        obs_dict["object_site2_to_target_site2_xyz_err"] = np.zeros(3)
        obs_dict["object_site3_to_target_site3_xyz_err"] = np.zeros(3)
        obs_dict["object_top_angle_cos"] = np.array([0])
        obs_dict["object_top_angle_sin"] = np.array([0])
        obs_dict["object_top_angle"] = np.array([0])
        obs_dict["target_top_angle_cos"] = np.array([0])
        obs_dict["target_top_angle_sin"] = np.array([0])
        obs_dict["target_top_angle"] = np.array([0])

        # Rod/pipe specific observations
        obs_dict["object_normal_to_target_normal_err"] = np.zeros(3)
        obs_dict["object_parallel_to_target_parallel_err"] = np.zeros(3)
        obs_dict["object_normal_to_target_normal_distance"] = np.zeros(1)
        obs_dict["object_parallel_to_target_parallel_distance"] = np.zeros(1)

        obs_dict['step_counter'] = np.array([0])

        if not self.initializing and self.object_type == ObjectType.Valve3:
            target1_xpos = self.sim.data.site_xpos[self.target_1].copy()
            valve1_xpos = self.sim.data.site_xpos[self.valve_1].copy()

            valve_dir = np.array(valve1_xpos - object_xyz)
            # Project onto xy plane
            valve_dir[2] = 0
            # Unit vector in direction of the red prong
            valve_dir /= np.linalg.norm(valve_dir)

            target_dir = np.array(target1_xpos - target_xyz)
            # Project onto xy plane
            target_dir[2] = 0
            # Unit vector in direction of the red prong
            target_dir /= np.linalg.norm(target_dir)

            object_angle = np.arctan2(valve_dir[1], valve_dir[0])
            obs_dict["object_top_angle"] = np.array([object_angle])
            obs_dict["object_top_angle_cos"] = np.array([valve_dir[0]])
            obs_dict["object_top_angle_sin"] = np.array([valve_dir[1]])
            target_angle = np.arctan2(target_dir[1], target_dir[0])
            obs_dict["target_top_angle"] = np.array([target_angle])
            obs_dict["target_top_angle_cos"] = np.array([target_dir[0]])
            obs_dict["target_top_angle_sin"] = np.array([target_dir[1]])

            obs_dict["object_to_target_mod_120_circle_distance"] = np.array([
                circle_distance_mod(object_angle, target_angle, mod=(2 * np.pi / 3))
            ])

            # Calculating site distances (of closest matching sites)
            perm_1_dist = np.mean([
                self.get_site_distance(self.target_1, self.valve_1),
                self.get_site_distance(self.target_2, self.valve_2),
                self.get_site_distance(self.target_3, self.valve_3)
            ])
            perm_2_dist = np.mean([
                self.get_site_distance(self.target_1, self.valve_2),
                self.get_site_distance(self.target_2, self.valve_3),
                self.get_site_distance(self.target_3, self.valve_1)
            ])
            perm_3_dist = np.mean([
                self.get_site_distance(self.target_1, self.valve_3),
                self.get_site_distance(self.target_2, self.valve_1),
                self.get_site_distance(self.target_3, self.valve_2)
            ])

            if perm_1_dist <= perm_2_dist and perm_1_dist <= perm_3_dist:
                # Use the first permutation because it is the closest
                valve_1 = self.valve_1
                valve_2 = self.valve_2
                valve_3 = self.valve_3
                object_to_target_dist = perm_1_dist
            elif perm_2_dist <= perm_1_dist and perm_2_dist <= perm_3_dist:
                # Use the 2nd
                valve_1 = self.valve_2
                valve_2 = self.valve_3
                valve_3 = self.valve_1
                object_to_target_dist = perm_2_dist
            elif perm_3_dist <= perm_1_dist and perm_3_dist <= perm_2_dist:
                # Use the 3rd
                valve_1 = self.valve_3
                valve_2 = self.valve_1
                valve_3 = self.valve_2
                object_to_target_dist = perm_3_dist

            obs_dict["object_sites_to_target_sites_distance"] = object_to_target_dist
            obs_dict["object_site1_to_target_site1_xyz_err"] = self.sim.data.site_xpos[self.target_1] - self.sim.data.site_xpos[valve_1]
            obs_dict["object_site2_to_target_site2_xyz_err"] = self.sim.data.site_xpos[self.target_2] - self.sim.data.site_xpos[valve_2]
            obs_dict["object_site3_to_target_site3_xyz_err"] = self.sim.data.site_xpos[self.target_3] - self.sim.data.site_xpos[valve_3]

        if not self.initializing and self.object_type in (ObjectType.Rod, ObjectType.Pipe):
            object_center = self.sim.data.site_xpos[self.rod_center]
            target_center = self.sim.data.site_xpos[self.target_center]
            object_normal = self.sim.data.site_xpos[self.rod_normal]
            object_parallel = self.sim.data.site_xpos[self.rod_parallel]
            target_normal = self.sim.data.site_xpos[self.target_normal]
            target_parallel = self.sim.data.site_xpos[self.target_parallel]

            object_normal_vec = obs_dict["object_normal_vector"] = object_normal - object_center
            target_normal_vec = obs_dict["target_normal_vector"] = target_normal - target_center
            object_parallel_vec = obs_dict["object_parallel_vector"] = object_parallel - object_center
            target_parallel_vec = obs_dict["target_parallel_vector"] = target_parallel - target_center

            normal_err = obs_dict["object_normal_to_target_normal_err"] = object_normal_vec - target_normal_vec
            parallel_err = obs_dict["object_parallel_to_target_parallel_err"] = object_parallel_vec - target_parallel_vec
            obs_dict["object_normal_to_target_normal_distance"] = np.linalg.norm(normal_err)
            obs_dict["object_parallel_to_target_parallel_distance"] = np.linalg.norm(parallel_err)

        return obs_dict

    def get_site_distance(self, site1, site2):
        return np.linalg.norm(self.sim.data.site_xpos[site1] - self.sim.data.site_xpos[site2])

    def step(self, action):
        obs, rew, done, info = super().step(action)
        if self._collect_bin_counts:
            x_d, y_d = self.last_obs_dict["object_xy_discrete"]
            self.bin_counts[x_d, y_d] += 1
        return obs, rew, done, info

    def get_reward_dict(
            self,
            action: np.ndarray,
            obs_dict: Dict[str, np.ndarray]
    ) -> Dict[str, np.ndarray]:
        xy_d = obs_dict["object_xy_discrete"].squeeze()
        x_d, y_d = xy_d[0], xy_d[1]
        circle_dist = obs_dict["object_to_target_circle_distance"]
        circle_dist_mod_180 = obs_dict["object_to_target_mod_180_circle_distance"]
        xy_dist = obs_dict["object_to_target_xy_distance"]
        x_dist = obs_dict["object_to_target_x_distance"]
        y_dist = obs_dict["object_to_target_y_distance"]
        z_dist = obs_dict["object_to_target_z_distance"]
        xyz_dist = obs_dict["object_to_target_xyz_distance"]
        hand_to_obj_distance = obs_dict["object_to_hand_xyz_distance"]
        hand_to_target_distance = obs_dict["target_to_hand_xyz_distance"]

        reward_dict = collections.OrderedDict((
            # Negative average distance btwn ends of object and hand
            ('object_to_hand_xyz_distance_reward', -hand_to_obj_distance),
            ('target_to_hand_xyz_distance_reward', -hand_to_target_distance),
            # Negative distance between center of object to desired location
            ("object_to_target_xy_distance_reward", -xy_dist),
            ("object_to_target_xyz_distance_reward", -xyz_dist),
            ("object_to_target_x_distance_reward", -x_dist),
            ("object_to_target_y_distance_reward", -y_dist),
            ("object_to_target_z_distance_reward", -z_dist),

            ("object_to_target_circle_distance_reward", -circle_dist),
            ("object_to_target_mod_120_circle_distance_reward", -obs_dict["object_to_target_mod_120_circle_distance"]),
            ("object_to_target_mod_180_circle_distance_reward", -circle_dist_mod_180),
            ("object_to_target_circle_distance_normalized_reward", -circle_dist / np.pi),
            # Measuring spread of the Dhand fingers compared to desired grip
            ('span_dist', np.array([0.])),
            ('hand_pose_dist_reward', -obs_dict["pose_dist"]),

            ('xy_discrete_sqrt_count_reward', 1 / np.sqrt(self.bin_counts[x_d, y_d])),
            # ('xy_discrete_invlog_count_reward', 1 / np.log(1 + self.bin_counts[x_d, y_d])),

            ("object_normal_to_target_normal_distance_reward", -obs_dict["object_normal_to_target_normal_distance"]),
            ("object_parallel_to_target_parallel_distance_reward", -obs_dict["object_parallel_to_target_parallel_distance"]),
        ))

        if not self.initializing:
            reward_dict["span_dist"] = -.03 * np.linalg.norm(
                self.sim.data.qpos[7:23] - self.act_mid - 0.5 * self.act_rng)

        return reward_dict

    def get_score_dict(
            self,
            obs_dict: Dict[str, np.ndarray],
            reward_dict: Dict[str, np.ndarray]
    ) -> Dict[str, np.ndarray]:
        return collections.OrderedDict()

    def _discretize_xys(self, xys):
        assert isinstance(xys, np.ndarray) and xys.ndim == 2 and xys.shape[1] == 2
        x_d = np.expand_dims(np.digitize(xys[:, 0], self.x_bins), 1)
        y_d = np.expand_dims(np.digitize(xys[:, 1], self.y_bins), 1)
        return np.concatenate((x_d, y_d), axis=1)

    def get_count_bonuses(self, xys):
        """
        Parameters
        ----------
        xys : np.ndarray (batch_size x 2)
            Must have batch_size as its first dimension, even if it's only 1
        """
        xy_d = self._discretize_xys(xys)
        count_bonuses = 1 / np.sqrt(self.bin_counts[xy_d[:, 0], xy_d[:, 1]])
        return count_bonuses.reshape((-1, 1))

    def reset_target(self):
        pass

    def set_goal(self, goal_index):
        """Required to implement this method for softlearning MultiSAC."""
        pass


import functools
import warnings

import numpy as np
import sparse
import scipy
import scipy.sparse as ssp
import torch
import torch.sparse as tsparse

def _ensure_tensor(x):
    if isinstance(x, torch.Tensor):
        if x.ndim == 0:
            return x.item()
        if x.is_sparse:
            return STensor.from_torch(x)
        return DTensor(x)
    elif isinstance(x, torch.Size):
        return DTensor(torch.as_tensor(x))
    else:
        return x

def _wrap_ret(func):
    """Wrap return value of func to spartan tensor types.
    """
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        ret = func(*args, **kwargs)
        try:
            ret = torch.as_tensor(ret).cuda()
        except Exception:
            pass
        return _ensure_tensor(ret)
    return wrapper

class TorchArithmetic:
    # Unary Opeartors
    @_wrap_ret
    def __neg__(self):
        return -self._data

    @_wrap_ret
    def __pos__(self):
        return +self._data

    @_wrap_ret
    def __invert__(self):
        return ~self._data

    # Binary Operators
    @_wrap_ret
    def __add__(self, other):
        try:
            return self._data + other._data
        except AttributeError:
            return self._data + other

    @_wrap_ret
    def __mul__(self, other):
        try:
            return self._data * other._data
        except AttributeError:
            return self._data * other

    @_wrap_ret
    def __truediv__(self, other):
        try:
            return self._data / other._data
        except AttributeError:
            return self._data / other

    @_wrap_ret
    def __floordiv__(self, other):
        try:
            return self._data // other._data
        except AttributeError:
            return self._data // other

    @_wrap_ret
    def __mod__(self, other):
        try:
            return self._data % other._data
        except AttributeError:
            return self._data % other

    @_wrap_ret
    def __pow__(self, other, modulo=None):
        try:
            return self._data ** other._data
        except AttributeError:
            return self._data ** other

    @_wrap_ret
    def __matmul__(self, other):
        try:
            return self._data @ other._data
        except AttributeError:
            return self._data @ other

    @_wrap_ret
    def __and__(self, other):
        try:
            return self._data & other._data
        except AttributeError:
            return self._data & other

    @_wrap_ret
    def __or__(self, other):
        try:
            return self._data | other._data
        except AttributeError:
            return self._data | other

    @_wrap_ret
    def xor(self, other):
        try:
            return self._data ^ other._data
        except AttributeError:
            return self._data ^ other

    @_wrap_ret
    def __lshift__(self, other):
        return self._data << other

    @_wrap_ret
    def __rshift__(self, other):
        return self._data >> other

    # Comparison Operators
    @_wrap_ret
    def __eq__(self, other):
        try:
            return self._data == other._data
        except AttributeError:
            return self._data == other

    @_wrap_ret
    def __gt__(self, other):
        try:
            return self._data > other._data
        except AttributeError:
            return self._data > other

    @_wrap_ret
    def __lt__(self, other):
        try:
            return self._data < other._data
        except AttributeError:
            return self._data < other

    @_wrap_ret
    def __ge__(self, other):
        try:
            return self._data >= other._data
        except AttributeError:
            return self._data >= other

    @_wrap_ret
    def __le__(self, other):
        try:
            return self._data <= other._data
        except AttributeError:
            return self._data <= other

class DTensor(TorchArithmetic):
    """A dense multi-deimensional tensor on GPU (based on PyTorch).

    Parameters
    ----------
    value : np.ndarray or array_like
        Data of dense tensor.
    dtype : data-type, optional
        Data type of the tensor, by default None(inferred automatically)

    Examples
    --------
    You can create a `DTensor` in many ways. For example, `numpy.ndarray`
    `torch.Tensor` and `list of list`.

    >>> x = np.random.rand(3, 4)
    >>> A = st.DTensor.from_numpy(x)
    >>> A
    DTensor(
    tensor([[0.9604, 0.9878, 0.4680, 0.0225],
            [0.5677, 0.3412, 0.6834, 0.1586],
            [0.6883, 0.5613, 0.8819, 0.2513]], device='cuda:0',
        dtype=torch.float64)
    )
    >>> A = st.DTensor(x)
    >>> A
    DTensor(
    tensor([[0.9604, 0.9878, 0.4680, 0.0225],
            [0.5677, 0.3412, 0.6834, 0.1586],
            [0.6883, 0.5613, 0.8819, 0.2513]], device='cuda:0',
        dtype=torch.float64)
    )
    >>> x = torch.rand(3, 4)
    >>> A = st.DTensor(x)
    >>> A
    DTensor(
    tensor([[0.8155, 0.3988, 0.6346, 0.5665],
            [0.2955, 0.8131, 0.0774, 0.7335],
            [0.0919, 0.8163, 0.5959, 0.9401]], device='cuda:0')
    )
    >>> A = st.DTensor([[1, 2, 3, 4], [2, 1, 4, 3], [1, 3, 2, 4]])
    >>> A
    DTensor(
    tensor([[1, 2, 3, 4],
            [2, 1, 4, 3],
            [1, 3, 2, 4]], device='cuda:0')
    )
    """

    def __init__(self, data, dtype=None):
        if isinstance(data, DTensor):
            self._data = data._data
        elif isinstance(data, STensor):
            self._data = data._data.to_dense()
        else:
            self._data = torch.as_tensor(data, dtype=dtype)
        self._data: torch.Tensor = self._data.cuda()

    def __repr__(self):
        return '%s(\n%r\n)' % (type(self).__name__, self._data)

    def __len__(self):
        return self._data.__len__()

    def __copy__(self):
        return self._data.__copy__()

    def __iter__(self):
        return iter(self._data)

    # Slice
    @_wrap_ret
    def __getitem__(self, index):
        return self._data.__getitem__(index)

    def __setitem__(self, index, value):
        self._data.__setitem__(index, value)

    def __delitem__(self, index):
        self._data.__delitem__(index)

    @_wrap_ret
    def __getattr__(self, name):
        if name.startswith('_'):
            return super().__getattr__(name)
        ret = getattr(self._data, name)
        if callable(ret):
            raise AttributeError(
                f"{type(self).__name__} doesn't have attribute {name}")
        return ret

    def __setattr__(self, name, value):
        if name.startswith('_'):
            super().__setattr__(name, value)
        else:
            setattr(self._data, name, value)

    def __delattr__(self, name):
        if name.startswith('_'):
            super().__delattr__(name)
        else:
            delattr(self._data, name)

    @_wrap_ret
    def all(self, axis=None, keepdims=False):
        """Wrapper of `torch.all`

        Parameters
        ----------
        axis : None or int or tuple of ints , optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        DTensor or bool
            Result of `all` operation as a DTensor or a boolean value.
        """
        if axis is None:
            return torch.all(self._data)
        else:
            return torch.all(self._data, dim=axis, keepdim=keepdims)

    @_wrap_ret
    def any(self, axis=None, keepdims=False):
        """Wrapper of `torch.any`

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        DTensor or bool
            Result of `any` operation as a DTensor or a boolean value.
        """
        if axis is None:
            return torch.any(self._data)
        return torch.any(self._data, dim=axis, keepdim=keepdims)

    @_wrap_ret
    def min(self, axis=None, keepdims=False):
        """Wrapper of `torch.Tensor.min`

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        DTensor or scalar
            Result of `min` operation as a DTensor or a scalar value.
        """
        return self._data.min(dim=axis, keepdim=keepdims).values

    @_wrap_ret
    def max(self, axis=None, keepdims=False):
        """Wrapper of `torch.Tensor.max`

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        DTensor or scalar
            Result of `max` operation as a DTensor or a scalar value.
        """
        return self._data.max(dim=axis, keepdim=keepdims).values

    @_wrap_ret
    def sum(self, axis=None, keepdims=False):
        """Wrapper of `torch.Tensor.sum`

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        DTensor or scalar
            Result of `sum` operation as a DTensor or a scalar value.
        """
        if axis is None:
            return torch.sum(self._data)
        return self._data.sum(dim=axis, keepdim=keepdims)

    @_wrap_ret
    def prod(self, axis=None, keepdims=False):
        """Wrapper of `torch.Tensor.prod`

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        DTensor or scalar
            Result of `prod` operation as a DTensor or a scalar value.
        """
        if axis is None:
            return torch.prod(self._data)
        return self._data.prod(dim=axis, keepdim=keepdims)

    @_wrap_ret
    def mean(self, axis=None, keepdims=False):
        """Wrapper of `torch.Tensor.mean`

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        DTensor or scalar
            Result of `mean` operation as a DTensor or a scalar value.
        """
        if axis is None:
            return torch.mean(self._data)
        return self._data.sum(dim=axis, keepdim=keepdims)

    @_wrap_ret
    def var(self, axis=None, keepdims=False):
        """Wrapper of `torch.Tensor.var`, note that result is biased(divisor is `n`).

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are \
                left to size one, by default False

        Returns
        -------
        DTensor or scalar
            Result of `var` operation as a DTensor or a scalar value.
        """
        if axis is None:
            return torch.var(self._data)
        return self._data.var(dim=axis, keepdim=keepdims)

    @_wrap_ret
    def std(self, axis=None, keepdims=False):
        """Wrapper of `torch.Tensor.std`

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        DTensor or scalar
            Result of `std` operation as a DTensor or a scalar value.
        """
        if axis is None:
            return torch.std(self._data)
        return self._data.std(dim=axis, keepdim=keepdims)

    @_wrap_ret
    def dot(self, other):
        """Wrapper of `torch.Tensor.dot`. Perform tensor dot operation with
        another tensor, including vector inner product, matrix multiplication
        and general tensor dot.

        Parameters
        ----------
        other : DTensor
            The second operand of dot operation.

        Returns
        -------
        DTensor or scalar
            Tensor dot result as a DTensor or a scalar value.
        """
        if other.ndim == 1:
            return self._data.dot(other._data)
        else:
            return self._data.matmul(other._data)

    @_wrap_ret
    def reshape(self, shape):
        """Wrapper of `torch.Tensor.reshape`. Gives a new dense tensor with
        new shape and current data.

        Parameters
        ----------
        shape : int or tuple of ints
            Shape of new tensor.

        Returns
        -------
        DTensor
            New dense tensor with given shape.
        """
        return self._data.reshape(shape)

    @_wrap_ret
    def nonzero(self):
        """Wrapper of `torch.Tensor.nonzero`. Find the indices of non-zero elements.

        Returns
        -------
        tuple of DTensor
            Indices of non-zero elements.
        """
        return tuple([DTensor(d) for d in self._data.nonzero()])

    @_wrap_ret
    def astype(self, dtype):
        """Wrapper of `torch.Tensor.astype`. Convert the array to a new type.

        Parameters
        ----------
        dtype : data-type
            Type of new tensor.

        Returns
        -------
        DTensor
            New dense tensor with given data type.
        """
        return self._data.type(dtype)

    def to_numpy(self):
        """Return the tensor as `numpy.ndarray`.

        Returns
        -------
        numpy.ndarray
            Returned numpy array.
        """
        return self._data.cpu().numpy()

    @classmethod
    def from_numpy(cls, x):
        """Construct a `DTensor` from a `numpy.ndarray`.

        Parameters
        ----------
        x : numpy.ndarray
            Input array.

        Returns
        -------
        DTensor
            Dense tensor constructed from given array.

        Raises
        ------
        TypeError
            Raise if the parameter is not `numpy.ndarray`.
        """
        if not isinstance(x, np.ndarray):
            raise TypeError(
                f"Argument type should be `numpy.ndarray`, got {type(x)}")
        t = cls.__new__(cls)
        t._data = torch.from_numpy(x).cuda()
        return t

class STensor(TorchArithmetic):
    """A sparse multi-dimensional tensor on GPU (based on PyTorch).

    Parameters
    ----------
    data : Tuple of indices (ndim x nnz) and values (1 x nnz) or STenosr
        or DTensor
        Data of sparse tensor.
    shape : tuple, optional
        Shape of sparse tensor, by default None (inferred automatically)
    """

    def __init__(self, data, shape=None):
        if type(data) is tuple:
            indices, values = data
            self._data = tsparse.FloatTensor(indices, values, shape=shape)
        elif isinstance(data, DTensor):
            self._data = data._data.to_sparse()
        elif isinstance(data, STensor):
            self._data = data._data
        else:
            self._data = tsparse.FloatTensor(data, shape=shape)
        self._data = self._data.cuda().coalesce()

    def __repr__(self):
        return '%s(\n%r\n)' % (type(self).__name__, self._data)

    def __len__(self):
        return self._data.__len__()

    def __copy__(self):
        return self._data.__copy__()

    def __iter__(self):
        return iter(self._data)

    # Notice that Pytorch's GPU sparse tensor doesn't support slice fully yet!
    # Slice
    @_wrap_ret
    def __getitem__(self, index):
        return self._data.__getitem__(index)

    def __setitem__(self, index, value):
        self._data.__setitem__(index, value)

    def __delitem__(self, index):
        self._data.__delitem__(index)

    @_wrap_ret
    def __getattr__(self, name):
        if name.startswith('_'):
            return super().__getattr__(name)
        ret = getattr(self._data, name)
        if callable(ret):
            raise AttributeError(
                f"{type(self).__name__} doesn't have attribute {name}")
        return ret

    def __setattr__(self, name, value):
        if name.startswith('_'):
            super().__setattr__(name, value)
        else:
            setattr(self._data, name, value)

    def __delattr__(self, name):
        if name.startswith('_'):
            super().__delattr__(name)
        else:
            delattr(self._data, name)

    @_wrap_ret
    def all(self, axis=None, keepdims=False):
        """Wrapper of `torch.all`

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        STensor or bool
            Result of `all` operation as a STensor or a boolean value.
        """
        if keepdims:
            warnings.warn(f"The keepdims parameter has no effect yet",
                          UserWarning)
        if axis is None:
            return torch.all(self._data)
        return torch.all(self._data, axis=axis, keepdim=keepdims)

    @_wrap_ret
    def any(self, axis=None, keepdims=False):
        """Wrapper of `torch.any`

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        STensor or bool
            Result of `any` operation as a STensor or a boolean value.
        """
        if keepdims:
            warnings.warn(f"The keepdims parameter has no effect yet",
                          UserWarning)
        if axis is None:
            return torch.any(self._data)
        return torch.any(self._data, axis=axis, keepdim=keepdims)

    @_wrap_ret
    def min(self, axis=None, keepdims=False):
        """Not implmented yet.

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        STensor or scalar
            Result of `min` operation as a STensor or a scalar value.
        """
        raise NotImplementedError

    @_wrap_ret
    def max(self, axis=None, keepdims=False):
        """Not implemented yet.

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        STensor or scalar
            Result of `max` operation as a STensor or a scalar value.
        """
        raise NotImplementedError

    @_wrap_ret
    def sum(self, axis=None, keepdims=False):
        """Wrapper of `torch.sparse.sum`

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        STensor or scalar
            Result of `sum` operation as a STensor or a scalar value.
        """
        if keepdims:
            warnings.warn(f"The keepdims parameter has no effect yet",
                          UserWarning)
        if axis is None:
            return tsparse.sum(self._data)
        return tsparse.sum(self._data, dim=axis)

    @_wrap_ret
    def prod(self, axis=None, keepdims=False):
        """Not implemented yet.

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        STensor or scalar
            Result of `prod` operation as a STensor or a scalar value.
        """
        raise NotImplementedError

    @_wrap_ret
    def mean(self, axis=None, keepdims=False):
        """Not implemented yet.

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        STensor or scalar
            Result of `mean` operation as a STensor or a scalar value.
        """
        raise NotImplementedError

    @_wrap_ret
    def var(self, axis=None, keepdims=False):
        """Not implemented yet.

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        STensor or scalar
            Result of `var` operation as a STensor or a scalar value.
        """
        raise NotImplementedError

    @_wrap_ret
    def std(self, axis=None, keepdims=False):
        """Not implemented yet.

        Parameters
        ----------
        axis : None or int or tuple of ints, optional
            Axis or axes to operate on, by default
            None
        keepdims : bool, optional
            If true, the axes along which the operation performed are left
            to size one, by default False

        Returns
        -------
        STensor or scalar
            Result of `std` operation as a STensor or a scalar value.
        """
        raise NotImplementedError

    @_wrap_ret
    def dot(self, other):
        """Wrapper of `torch.sparse.mm`. Perform tensor dot operation with
        another tensor, including vector inner product, matrix multiplication
        and general tensor dot.

        Parameters
        ----------
        other : DTensor or STensor
            The second operand of dot operation.

        Returns
        -------
        DTensor or STensor
            Tensor dot result as a DTensor or a STensor or a scalar value.
        """
        return tsparse.mm(self._data, other._data)

    @_wrap_ret
    def todense(self):
        """Convert to a dense tensor.

        Returns
        -------
        DTensor
            Dense tensor coverted from the sparse tensor.
        """
        return self._data.to_dense()

    @_wrap_ret
    def reshape(self, shape):
        """Not implemented yet.

        Parameters
        ----------
        shape : int or tuple of ints
            Shape of new tensor.

        Returns
        -------
        STensor
            New sparse tensor with given shape.
        """
        raise NotImplementedError

    @_wrap_ret
    def nonzero(self):
        """Wrapper of `torch.Tensor.indices`. Find the indices of non-zero elements.

        Returns
        -------
        tuple of DTensor
            Indices of non-zero elements.
        """
        return tuple([x for x in self._data.indices()])

    @_wrap_ret
    def astype(self, dtype):
        """Not implemented yet.

        Parameters
        ----------
        dtype : data-type
            Type of new tensor.

        Returns
        -------
        STensor
            New sparse tensor with given data type.
        """
        raise NotImplementedError

    def to_scipy(self, format='coo'):
        """Return the sparse tensor as a `scipy.sparse.spmatrix`.

        Parameters
        ----------
        format : str, optional, {'coo', 'csr', 'csc', 'lil', 'dok'}
            Format of scipy sparse matrix, by default 'coo'

        Returns
        -------
        scipy.sparse.spmatrix
            Returned scipy sparse matrix.
        """
        if self._data.ndim != 2:
            raise ValueError(
                f"Only 2-dim sparse tensor can be converted to scipy sparse\
                matrix, got {self._data.ndim}")
        indices = self._data.indices().cpu().numpy()
        row, col = indices[0], indices[1]
        vals = self._data.values().cpu().numpy()
        shape = tuple(self._data.shape)
        return ssp.coo_matrix((vals, (row, col)), shape=shape).asformat(format)

    def to_sparse_array(self):
        """Return the sparse tensor as a `sparse.COO`.

        Returns
        -------
        sparse.COO
            Returned sparse array.
        """
        indices = self._data.indices().cpu().numpy()
        vals = self._data.values().cpu().numpy()
        shape = tuple(self._data.shape)
        return sparse.COO(indices, vals, shape=shape)

    @classmethod
    def from_numpy(cls, x: np.ndarray):
        """Construct a `STensor` from a `numpy.ndarray`.

        Parameters
        ----------
        x : numpy.ndarray
            Input array.

        Returns
        -------
        STensor
            Sparse tensor constructed from given array.

        Raises
        ------
        TypeError
            Raise if the parameter is not `numpy.ndarray`.
        """
        if not isinstance(x, np.ndarray):
            raise TypeError(
                f"Argument type should be `numpy.ndarray`, got {type(x)}")
        t = cls.__new__(cls)
        t._data = torch.from_numpy(x).to_sparse().cuda()
        t._data = t._data.coalesce()
        return t

    @classmethod
    def from_scipy_sparse(cls, x: scipy.sparse.spmatrix):
        """Construct a `STensor` from a `scipy.sparse.spmatrix`.

        Parameters
        ----------
        x : scipy.sparse.spmatrix
            Input sparse matrix.

        Returns
        -------
        STensor
            Sparse tensor constructed from given sparse matrix.

        Raises
        ------
        TypeError
            Raise if the parameter is not `scipy.sparse.spmatrix`.
        """
        if not isinstance(x, scipy.sparse.spmatrix):
            raise TypeError(
                f"Argument type should be `scipy.sparse.spmatrix`, \
                got {type(x)}")
        x = x.tocoo()
        indices = torch.from_numpy(np.vstack([x.row, x.col])).type(torch.long)
        values = torch.from_numpy(x.data)
        shape = torch.Size(x.shape, dtype=torch.long)
        t = cls.__new__(cls)
        t._data = tsparse.FloatTensor(indices, values, shape).cuda()
        t._data = t._data.coalesce()
        return t

    @classmethod
    def from_sparse_array(cls, x: sparse.COO):
        """Construct a sparse tensor from a `sparse.COO`.

        Parameters
        ----------
        x : sparse.COO
            Input COO sparse array.

        Returns
        -------
        STensor
            Sparse tensor constructed from given sparse array.

        Raises
        ------
        TypeError
            Raise if the parameter is not `sparse.COO`.
        """
        if not isinstance(x, sparse.COO):
            raise TypeError(
                f"Argument type should be `sparse.SparseArray`, got {type(x)}")
        indices = torch.from_numpy(x.coords).type(torch.long)
        values = torch.from_numpy(x.data)
        shape = torch.Size(x.shape, dtype=torch.long)
        t = cls.__new__(cls)
        t._data = tsparse.FloatTensor(indices, values, shape).cuda()
        t._data = t._data.coalesce()
        return t

    @classmethod
    def from_torch(cls, x: torch.Tensor):
        """Construct a `STensor` from a `torch.Tensor`.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor.

        Returns
        -------
        STensor
            Sparse tensor constructed from given array.

        Raises
        ------
        TypeError
            Raise if the parameter is not `torch.Tensor`.
        """
        if not isinstance(x, torch.Tensor):
            raise TypeError(
                f"Argument type should be `torch.Tensor`, got {type(x)}")
        t = cls.__new__(cls)
        if not x.is_sparse:
            x = x.to_sparse()
        t._data = x.cuda()
        t._data = t._data.coalesce()
        return t


import moviepy.editor as mp
import os
#import matplotlib.pyplot as plt
import numpy as np
import librosa
import panns_inference
from panns_inference import AudioTagging, SoundEventDetection, labels
import time
import json
from json.decoder import JSONDecodeError
from collections import defaultdict

def get_audio_tagging_result(clipwise_output, number_of_classes=10, classes_set=set(), threshold=0.2):
    """Visualization of audio tagging result.

    Args:
      clipwise_output: (classes_num,)
    """
    sorted_indexes = np.argsort(clipwise_output)[::-1]
    result = {}
    # Print audio tagging top probabilities
    for k in range(number_of_classes):
        if np.array(labels)[sorted_indexes[k]] in classes_set and clipwise_output[sorted_indexes[k]] > threshold:
            print('{}: {:.3f}'.format(np.array(labels)[sorted_indexes[k]], clipwise_output[sorted_indexes[k]]))
            result[np.array(labels)[sorted_indexes[k]]] = [round(clipwise_output[sorted_indexes[k]], 2)]
    return result

def store_result(new_result={}):
    with open('/home/nvr/airesults/ser.json', 'r') as f:
        try:
            prev = json.load(f)
        except JSONDecodeError:
            prev = {}
    result = defaultdict(list)

    for d in (prev, new_result):
        for key, value in d.items():
            result[key].extend(value)

    with open('/home/nvr/airesults/ser.json', 'w') as f:
        json.dump(dict(result), f)

'''
def plot_sound_event_detection_result(framewise_output):
    """Visualization of sound event detection result.

    Args:
      framewise_output: (time_steps, classes_num)
    """
    out_fig_path = 'results/sed_result.png'
    os.makedirs(os.path.dirname(out_fig_path), exist_ok=True)

    classwise_output = np.max(framewise_output, axis=0) # (classes_num,)

    idxes = np.argsort(classwise_output)[::-1]
    idxes = idxes[0:5]

    ix_to_lb = {i : label for i, label in enumerate(labels)}
    lines = []
    for idx in idxes:
        line, = plt.plot(framewise_output[:, idx], label=ix_to_lb[idx])
        lines.append(line)

    plt.legend(handles=lines)
    plt.xlabel('Frames')
    plt.ylabel('Probability')
    plt.ylim(0, 1.)
    plt.savefig(out_fig_path)
    print('Save fig to {}'.format(out_fig_path))
'''

if __name__ == '__main__':
    """Example of using panns_inferece for audio tagging and sound evetn detection.
    """
    device = 'cpu'  # 'cuda' | 'cpu'
    at = AudioTagging(checkpoint_path=None, device=device)
    audio_path = "/home/nvr/converted_audio_files/"
    with open('/opt/iotistic-mnvr/config/default.json', 'r') as f:
        config = json.load(f)['SERParameters']
    classes = set(config['enabledclasses'].split('|'))
    threshold = config['threshold']

    while True:
        audio_files = os.listdir(audio_path)
        if not audio_files:
            time.sleep(0.1)
            continue
        t0 = time.time()
        # audio_path = 'resources/R9_ZSCveAHg_7s.wav'
        (audio, _) = librosa.core.load(audio_path + audio_files[0], sr=32000, mono=True)
        os.remove(audio_path + audio_files[0])
        audio = audio[None, :]  # (batch_size, segment_samples)

        print('------ Audio tagging ------')
        (clipwise_output, embedding) = at.inference(audio)
        """clipwise_output: (batch_size, classes_num), embedding: (batch_size, embedding_size)"""
        result = get_audio_tagging_result(clipwise_output[0], 10, classes, threshold)
        store_result(result)
        # print('------ Sound event detection ------')
        # sed = SoundEventDetection(checkpoint_path=None, device=device)
        # framewise_output = sed.inference(audio)
        # """(batch_size, time_steps, classes_num)"""

        # plot_sound_event_detection_result(framewise_output[0])

        print('done file within ', time.time() - t0)

        ########### run store_result on another thread?


import argparse
from functools import partial
import hashlib
import json
import logging
from pathlib import Path
import pickle
from typing import Callable, List, Tuple
import warnings

import faiss
import numpy as np
import spacy
from tqdm import tqdm

from text_to_uri import english_filter, replace_numbers

CACHE_DIR = Path('.cache/')

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
SPACY_MODEL='en_core_web_lg'

def init_cache():
    if not CACHE_DIR.exists():
        logger.debug(f'Creating cache dir at: {CACHE_DIR}')
        CACHE_DIR.mkdir(parents=True)

def _cache_path(fn, args, kwargs):
    fn_name = fn.__name__
    args_string = ','.join(str(arg) for arg in args)
    kwargs_string = json.dumps(kwargs)
    byte_string = fn_name + args_string + kwargs_string
    hash_object = hashlib.sha1(byte_string.encode())
    return CACHE_DIR / hash_object.hexdigest()

def cache():
    def decorator(fn):
        def load_cached_if_available(*args, **kwargs):
            path = _cache_path(fn, args, kwargs)
            if path.exists():
                logger.debug(f'Loading `{fn.__name__}` output from cache')
                with open(path, 'rb') as f:
                    return pickle.load(f)
            output = fn(*args, **kwargs)
            with open(path, 'wb') as f:
                pickle.dump(output, f)
            return output
        return load_cached_if_available
    return decorator

class Vocab:
    def __init__(self, words) -> None:
        self.idx_to_word = words
        self.word_to_idx = {word: idx for idx, word in enumerate(words)}

@cache()
def read_embedding_file(embedding_file: Path) -> Tuple[Vocab, np.ndarray]:

    logger.debug(f'Reading embeddings from {embedding_file}')

    with open(embedding_file, 'r') as f:
        info = next(f)
        shape = tuple(int(x) for x in info.split())
        embeddings = np.zeros(shape, dtype=np.float32)

        words = []
        for i, line in tqdm(enumerate(f), total=shape[0]):
            word, *embedding = line.split()
            embedding = np.array([float(x) for x in embedding])
            words.append(word)
            embeddings[i] = embedding

    vocab = Vocab(words)

    return vocab, embeddings

def build_index(metric: str, embeddings: np.ndarray) -> faiss.swigfaiss_avx2.Index:

    logger.debug(f'Building search index')

    if metric == 'cosine':
        index = faiss.IndexFlatIP(embeddings.shape[-1])
    elif metric == 'l2':
        index = faiss.IndexFlatL2(embeddings.shape[-1])
    else:
        raise ValueError(f'Bad metric: {metric}')

    index.add(embeddings)

    return index

def generate_instances(dataset: Path):
    with open(dataset, 'r') as f:
        for line in f:
            yield(json.loads(line))

def get_extraction_fn(extraction_strategy: str,
                      ngram_length: int) -> Callable[[List[str], Vocab], List[str]]:
    if extraction_strategy == 'exhaustive':
        return partial(exhaustive_extraction, ngram_length=ngram_length)
    elif extraction_strategy == 'greedy':
        return partial(greedy_extraction, ngram_length=ngram_length)
    elif extraction_strategy == 'root':
        return partial(root_extraction, ngram_length=ngram_length)
    else:
        raise ValueError(f'Bad extraction strategy: {extraction_strategy}')

def exhaustive_extraction(phrase: List[str],
                          vocab: Vocab,
                          ngram_length: int) -> List[str]:
    tokens = english_filter([x.lower() for x in phrase])
    num_tokens = len(tokens)
    out = []
    for n in range(1, ngram_length):
        for i in range(num_tokens - n + 1):
            concept = replace_numbers('_'.join(tokens[i: i+n]))
            if concept in vocab.word_to_idx:
                out.append(concept)
    return out

def greedy_extraction(phrase: List[str],
                      vocab: Vocab,
                      ngram_length: int) -> List[str]:
    tokens = english_filter([x.lower() for x in phrase])
    out = []
    while len(tokens) > 0:
        for n in range(ngram_length + 1, 0, -1):
            concept = replace_numbers('_'.join(tokens[:n]))
            if concept in vocab.word_to_idx:
                out.append(concept)
                tokens = tokens[n:]
                break
            elif n == 1:
                tokens = tokens[n:]
    return out

def root_extraction(phrase: List[str],
                    vocab: Vocab,
                    ngram_length: int) -> List[str]:
    doc = nlp(' '.join(phrase))

    # Logic for noun phrases
    for chunk in doc.noun_chunks:
        if chunk.root.dep_ == 'ROOT':

            # Return entire chunk if in vocab
            tokens = english_filter([token.text.lower() for token in chunk])
            concept = replace_numbers('_'.join(tokens))
            if concept in vocab.word_to_idx:
                return [concept]

            # Otherwise return the just the root token
            root = replace_numbers(chunk.root.text.lower())
            if root in vocab.word_to_idx:
                return [root]

    # Logic for other types of roots
    for token in doc:
        if token.dep_ == 'ROOT':
            concept = replace_numbers(token.text.lower())
            if concept in vocab.word_to_idx:
                return [concept]

    return []

def link(graphs: List,
         output: Path = None,
         embedding_file: Path = Path('../numberbatch-en-19.08.txt'),
         metric: str = 'cosine',
         extraction_strategy: str = 'greedy',
         ngram_length: int = 3,
         num_candidates: int = 5,
         debug: bool = False) -> List:
    """
    Browse the top-k conceptnet candidates for a node.

    Parameters
    ==========
    input : List
        A list containing parsed alpha NLI graphs.
    output : Path
        Jsonl file to serialize output to.
    embedding_file : Path
        A txt file containing the embeddings.
    metric: str
        Similarity metric. One of: 'cosine', 'l2'
    extraction_strategy: str
        Approach for extracting concepts from mentions. One of: 'exhaustive', 'greedy'
    ngram_length: int
        Max length of n-grams to consider during concept extraction.
    num_candidates : int
        Number of candidates to return.
    """
    assert metric in {'cosine', 'l2'}
    assert extraction_strategy in {'exhaustive', 'greedy', 'root'}

    if debug:
        logger.setLevel(logging.DEBUG)

    global nlp
    nlp = spacy.load(SPACY_MODEL)

    init_cache()
    vocab, embeddings = read_embedding_file(embedding_file)
    index = build_index(metric, embeddings)
    extraction_fn = get_extraction_fn(extraction_strategy, ngram_length)

    if output:
        output_file = open(output, 'w')
    output_instances=[]
    for instance in graphs:
        output_instance = instance.copy()
        for uri, node in instance['nodes'].items():
            # Extract concept tokens from phrase
            phrase = node['phrase']
            concepts = extraction_fn(phrase, vocab)
            concept_ids = np.array([vocab.word_to_idx[concept] for concept in concepts])

            if len(concept_ids) > 0:
                if len(concept_ids) > 1:
                    mean_embedding = np.mean(embeddings[concept_ids], axis=0, keepdims=True)
                    query = np.concatenate((embeddings[concept_ids], mean_embedding), axis=0)
                else:
                    query = embeddings[concept_ids]

                scores, candidate_ids = index.search(query, num_candidates)

                # Convert from cosine similarity to distance
                if metric == 'cosine':
                    scores = 1 - scores

                # If phrase contains k concepts then the search returns k * num_candidates results.
                # Reduces to top-k
                scores = scores.flatten()
                candidate_ids = candidate_ids.flatten()
                top_k_indices = np.argsort(scores)[:num_candidates]
                scores = scores[top_k_indices]
                candidate_ids = candidate_ids[top_k_indices]
                output_instance['nodes'][uri]['candidates'] = []
                for candidate_id, score in zip(np.nditer(candidate_ids), np.nditer(scores)):
                    candidate = vocab.idx_to_word[candidate_id]
                    if '#' in candidate:
                        logger.warning(f'Encountered candidate URI: "{candidate}". Due to preprocessing steps '
										'used to produce the ConceptNet Numberbatch embeddings this is '
										'likely a bad link, and will be skipped.')
                        continue
                    output_instance['nodes'][uri]['candidates'].append({
						'uri': '/c/en/' + candidate,  # TODO: Support other KBs
						'score': score.item()
                    })
            else:
                output_instance['nodes'][uri]['candidates'] = []
        output_instances.append(output_instance)
        if output:
            output_file.write(json.dumps(output_instance) + '\n')
    if output:
        output_file.close()
    return output_instances

def main():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--input', type=Path, required=True, help='JSON-Lines file containing the text graphs.')
    parser.add_argument('--output', type=Path, required=False, help='Output JSON-Lines file.')
    parser.add_argument('--embedding_file', type=Path, required=False,
                        help='KG node embeddings. Currently only ConceptNet Numberbatch supported.')
    parser.add_argument('--metric', type=str, default='cosine',
                        help='Distance metric used for nearest neighbor search. One of: "cosine", "l2".')
    parser.add_argument('--extraction_strategy', type=str, default='greedy',
                        help='Approach for extracting concepts from text. One of: "exhaustive", "greedy", "root".'
                             'Exhaustive extraction produces all viable n-grams, e.g., "high school" -> {"high", "school", "high_school"}. '
                             'Greedy search produces only the largest viable n-grams, e.g., "high school" -> {"high_school"}. '
                             'Root extraction only produces the root of the dependency parse, e.g., "his car" -> {"car"}.')
    parser.add_argument('--ngram_length', type=int, default=3,
                        help='Maximum length of ngrams to consider when converting text to KG nodes (i.e. ConceptNet concepts).')
    parser.add_argument('--num_candidates', type=int, default=5, help='Number of candidates to return.')
    parser.add_argument('--debug', action='store_true', help='Enables debug statements.')
    args = parser.parse_args()

    link(graphs=generate_instances(args.input),
         output=args.output,
         embedding_file=args.embedding_file,
         metric=args.metric,
         extraction_strategy=args.extraction_strategy,
         ngram_length=args.ngram_length,
         num_candidates=args.num_candidates,
         debug=args.debug)

if __name__ == '__main__':
    main()


#from batchprocess import Batch_data, pretraining_batch
import tensorflow as tf
import pandas as pd
import itertools
import os
import logging
import datetime
import argparse
import numpy as np
from collections import Counter
import time
from nets import model_softmax as model
from batch import TrainData,MultiTaskBatchManager
from tools import get_available_gpus,average_gradients
import multiprocessing as mp

parser = argparse.ArgumentParser(description='Model Training')
parser.add_argument('--chkpt', type=str, help='checkpoint dir name')
parser.add_argument('--chkpt_load_dir', type=str, help='pre trained checkpoint dir name',default=' ')
parser.add_argument('--E_train', type=int, help='number of training epochs',default=5)
parser.add_argument('--lr', type=float, help='learning rate',default=0.0001)
parser.add_argument('--bs', type=int, help='training BatchSize',default=64)
parser.add_argument('--csv', type=str, help='training Data CSV file',default='ted_train.csv')

args = parser.parse_args()
parameters='--chkpt '+args.chkpt+'--chkpt_load_dir '+args.chkpt_load_dir
parameters+='--E_train '+str(args.E_train)+'--lr '+str(args.lr)+'--bs '+str(args.bs)

if args.chkpt_load_dir==' ':
    chkpt_load_dir=args.chkpt
else:
    chkpt_load_dir=args.chkpt_load_dir
chkpt_dir = args.chkpt
if not os.path.exists(chkpt_dir):
    os.mkdir(chkpt_dir)

train_df = pd.read_csv('ted_train.csv')
total_speakers = max(train_df['speaker_id'].unique())
batch_size=args.bs

logging.basicConfig(
    filename=chkpt_dir + '/' + datetime.datetime.now().strftime('training_log_%d_%m_%Y_%H_%M.log'),
    level=logging.DEBUG)
logging.info(datetime.datetime.now().strftime('%d_%m_%Y_%H_%M_') + 'Training Parameters : %s' %parameters)
chkpt_file=chkpt_dir + '/model.ckpt'

T_graph=tf.Graph()
num_epoch=args.E_train

def train():
    BatchLoader=MultiTaskBatchManager(Data=train_df,batch_size=batch_size,Ntasks=5,Nepochs=num_epoch)
    print('epoch size : = %s'%(len(train_df)/batch_size))
    with T_graph.as_default():
        tower_grads=[]
        global_step = tf.Variable(0, name='global_step', trainable=False)
        x=tf.placeholder('float32',shape=[batch_size,None])
        y=tf.placeholder('int32',shape=[batch_size])
        optimizer = tf.train.AdamOptimizer(learning_rate=args.lr)
        available_gpus=get_available_gpus()
        num_clones=len(available_gpus)
        #x,y,flag=train_data.GetBatch()
        #x=tf.convert_to_tensor(x, dtype=tf.float32)
        #y=tf.convert_to_tensor(y, dtype=tf.int32)
        #batch_queue = tf.contrib.slim.prefetch_queue.prefetch_queue([x, y], capacity=2 * num_clones)
        print('Number of clones = %d'%num_clones)
        with tf.variable_scope(tf.get_variable_scope()):
            for i in range(num_clones):
                with tf.device(available_gpus[i]):
                    #batch_x, batch_y = batch_queue.dequeue()
                    # Distribute data among all clones equally.
                    step = int(batch_size / float(num_clones))
                    # Network outputs.
                    prediction= model(x[i * step: (i + 1) * step],step,total_speakers)
                    label_onehot = tf.one_hot(y[i * step : (i + 1) * step], depth=total_speakers+1)
                    #prediction= model(batch_x,batch_size,total_speakers)
                    #label_onehot = tf.one_hot(batch_y, depth=total_speakers+1)
                    SOFTMAX = tf.nn.softmax_cross_entropy_with_logits(logits=prediction, labels=label_onehot)
                    with tf.name_scope('loss'):
                        loss = tf.reduce_mean(SOFTMAX)
                    # Reuse variables for the next tower.
                    tf.get_variable_scope().reuse_variables()
                    # Calculate the gradients for the batch of data on this tower.
                    grads = optimizer.compute_gradients(loss)
                    tower_grads.append(grads)
        grads = average_gradients(tower_grads)
        # Apply the gradients to adjust the shared variables.
        apply_gradient_op = optimizer.apply_gradients(grads, global_step=global_step)
        # Track the moving averages of all trainable variables.
        MOVING_AVERAGE_DECAY = 0.9999
        variable_averages = tf.train.ExponentialMovingAverage(MOVING_AVERAGE_DECAY, global_step)
        variables_averages_op = variable_averages.apply(tf.trainable_variables())
        train_op = tf.group(apply_gradient_op, variables_averages_op)

        summaries=set(tf.get_collection(tf.GraphKeys.SUMMARIES))
        summaries.add(tf.summary.scalar('loss',loss))
        summary_op=tf.summary.merge(list(summaries))

        with tf.Session(graph=T_graph,config=tf.ConfigProto(allow_soft_placement=True,log_device_placement=False)) as sess:
            saver = tf.train.Saver()
            summary_writer=tf.summary.FileWriter(chkpt_dir,graph=T_graph)
            sess.run(tf.global_variables_initializer())
            tf.train.start_queue_runners(sess=sess)
            if os.path.exists(chkpt_dir + '/checkpoint'):
                print('restoring !!')
                saver.restore(sess, chkpt_file)
            elif not os.path.exists(chkpt_dir):
                os.mkdir(chkpt_dir)
            print('Training Started !!')
            isrunning=True
            stepcount=0
            steploss=0
            epoch_loss=0
            nep=0
            while isrunning:
                stepcount+=1
                #st=time.time()
                batch_x, batch_y, flag,isrunning=BatchLoader.next_batch()
                #print('data loading time %s'%(time.time()-st))
                if isrunning:
                    #st=time.time()
                    _,c,summary,g= sess.run([apply_gradient_op,loss,summary_op,global_step],feed_dict={x:batch_x,y: batch_y})
                    #print('training time %s'%(time.time()-st))
                    steploss+=c
                    epoch_loss+=c
                    if stepcount%100==0:
                        epoch_loss+=steploss
                        save_path = saver.save(sess, chkpt_file)
                        print('step_loss : %s '%steploss)
                        logging.info('step_loss : %s'%(steploss))
                        steploss=0
                    if stepcount%(len(train_df)/batch_size)==0:
                        logging.info('Epoch %d loss : %f'%(nep,epoch_loss))
                        print('Epoch %d loss : %f'%(nep,epoch_loss))
                        epoch_loss=0
                        nep+=1
            logging.info('last Batch %d loss : %f'%(nep,epoch_loss))
            print('last Batch %d loss : %f'%(nep,epoch_loss))
            BatchLoader.close()
if __name__ == '__main__':
    train()


#!/usr/bin/python3
import mido
import pigpio
from numpy import interp
import board
import neopixel

pixels = neopixel.NeoPixel(board.D18, 300)

port = mido.open_input('f_midi')  # open USB port
red = 0
green = 0
blue = 0
while True:
    try:
        for msg in port.iter_pending():

            if(msg.type == 'note_on'):
                out = interp(msg.velocity, [0,127],[0,255])
                if(msg.note == 60):
                    red = out
                if(msg.note == 61):
                    green = out
                if(msg.note == 62):
                    blue = out
                if(msg.note == 65):
                    red = out
                    green = out
                    blue = out
                if(msg.note == 66):
                    red = out
                    green = out
                    blue = out
            if(msg.type == 'note_off'):
                if(msg.note == 60):
                    red = 0
                if(msg.note == 61):
                    green = 0
                if(msg.note == 62):
                    blue = 0
                if(msg.note == 65):
                    red = 0
                    green = 0
                    blue = 0
            pixels.fill((red, green, blue))
    except AttributeError as error:
        print("Error excepted")
    pass