### --------------------------------------------------------------------------- ###
# loads monthly mean raw live stim and world cam luminance data files (saved at 4ms resolution)
# saves sanity check mean world cam video for each stimulus type
# measures display latency (lag between when bonsai tells a frame to display and when it actually displays)
# output display latency and sanity check plots/videos
# NOTE: this script uses ImageMagick to easily install ffmpeg onto Windows 10 (https://www.imagemagick.org/script/download.php)
# NOTE: in command line run with optional tags
#       1) '--a debug' to use only a subset of pupil location/size data
#       2) '--a vid_output' to generate sanity check mean world cam videos
#       3) '--a MOI' to find moments of interest in the mean world cam videos
#       4) '--loc *' to run with various root data locations (see first function below)
### --------------------------------------------------------------------------- ###
import logging
import pdb
import os
import glob
import datetime
import math
import sys
import itertools
import csv
import fnmatch
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from scipy import stats
import argparse
from IPython import embed
###################################
# SET CURRENT WORKING DIRECTORY
###################################
current_working_directory = os.getcwd()
###################################
# SCRIPT LOGGER
###################################
# grab today's date
now = datetime.datetime.now()
logging.basicConfig(filename="psa02_DisplayLatency_SanityChecks" + now.strftime("%Y-%m-%d_%H-%M-%S") + ".log", filemode='w', level=logging.INFO)
###################################
# FUNCTIONS
###################################

##########################################################
#### MODIFY THIS FIRST FUNCTION BASED ON THE LOCATIONS OF:
# 1) root_folder (parent folder with all intermediate data)
# AND
# 2) plots_folder (parent folder for all plots output from analysis scripts)
### Current default uses a debugging source dataset
##########################################################
def load_data(location='laptop'):
    if location == 'laptop':
        root_folder = r"C:\Users\taunsquared\Dropbox\SurprisingMinds\analysis\dataPythonWorkflows"
        plots_folder = r"C:\Users\taunsquared\Dropbox\SurprisingMinds\analysis\plots"
    elif location == 'office':
        root_folder = r"C:\Users\Kampff_Lab\Dropbox\SurprisingMinds\analysis\dataPythonWorkflows"
        plots_folder = r"C:\Users\Kampff_Lab\Dropbox\SurprisingMinds\analysis\plots"
    # monthly mean raw live and world cam luminances
    monthly_mean_lums_folders = fnmatch.filter(sorted(os.listdir(root_folder)), 'MeanStimuli_*')
    # display latency
    display_latency_folder = os.path.join(root_folder, 'displayLatency')
    # full dataset mean world cam output folder
    mean_world_cam_vids_folder = os.path.join(plots_folder, 'meanWorldCam')
    # sanity check
    raw_v_world_sanity_check_folder = os.path.join(plots_folder, 'rawVWorldSanityCheck')
    # Create output folders if they do not exist
    output_folders = [display_latency_folder, mean_world_cam_vids_folder, raw_v_world_sanity_check_folder]
    for output_folder in output_folders:
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
    return root_folder, plots_folder, monthly_mean_lums_folders, output_folders

##########################################################
def extract_stim_data(stim_data_array, world_or_raw, stim_data_dict):
    for row in stim_data_array:
        timebucket = row[0]
        weight = row[1]
        mean_data = row[2]
        this_tb_weighted_data = weight*mean_data
        if world_or_raw=='world':
            data_type = 'weighted keyframes'
        elif world_or_raw=='raw':
            data_type = 'weighted luminance'
        else:
            print('Invalid stimulus data type! Must be world or raw.')
            logging.warning('Invalid stimulus data type! Must be world or raw.')
        if timebucket in stim_data_dict.keys():
            stim_data_dict[timebucket][data_type].append(this_tb_weighted_data)
            stim_data_dict[timebucket]['weights'].append(weight)
        else:
            stim_data_dict[timebucket] = {data_type:[this_tb_weighted_data], 'weights':[weight]}

def calculate_weighted_sums(all_weighted_stim_data_dict, world_or_raw):
    weighted_sums_stim_dict = {}
    for stim in all_weighted_stim_data_dict.keys():
        weighted_sums_stim_dict[stim] = {}
        for timebucket in sorted(all_weighted_stim_data_dict[stim].keys()):
            if world_or_raw=='world':
                this_tb_weighted_sum = np.sum(np.array(all_weighted_stim_data_dict[stim][timebucket]['weighted keyframes']), axis=0)
                weighted_sum_key = 'keyframe, weighted sum'
            elif world_or_raw=='raw':
                this_tb_weighted_sum = np.sum(all_weighted_stim_data_dict[stim][timebucket]['weighted luminance'])
                weighted_sum_key = 'luminance, weighted sum'
            else:
                print('Invalid stimulus data type! Must be world or raw.')
                logging.warning('Invalid stimulus data type! Must be world or raw.')
            this_tb_weights_sum = np.sum(all_weighted_stim_data_dict[stim][timebucket]['weights'])
            weighted_sums_stim_dict[stim][timebucket] = {weighted_sum_key:this_tb_weighted_sum, 'summed weight':this_tb_weights_sum}
    return weighted_sums_stim_dict

def sanity_check_world_v_rawLive(world_dict, worldFull_or_worldCropped, raw_dict, timebucket_size, save_folder):
    for stim in world_dict.keys():
        if worldFull_or_worldCropped == 'full':
            world_label = 'world (full)'
            # figure path and title
            figPath = os.path.join(save_folder, 'Stim%d_worldFull_versus_raw_sanityCheck.png'%(stim))
            figTitle = 'Mean luminance of world cam (full) vs mean luminance of raw live stim during Stimulus %d'%(stim)
        elif worldFull_or_worldCropped == 'cropped':
            world_label = 'world (cropped)'
            # figure path and title
            figPath = os.path.join(save_folder, 'Stim%d_worldCropped_versus_raw_sanityCheck.png'%(stim))
            figTitle = 'Mean luminance of world cam (cropped to remove display latency) \n vs mean luminance of raw live stim during Stimulus %d'%(stim)
        plt.figure(figsize=(9, 9), dpi=200)
        plt.suptitle(figTitle, fontsize=12, y=0.98)
        plt.ylabel('Mean Luminance')
        plt.xlabel('Time (sec)')
        plt.xticks(np.arange(0,len(raw_dict[stim]),1000), np.arange(0,len(raw_dict[stim]),1000)*timebucket_size/1000, rotation=50)
        plt.plot(world_dict[stim], label=world_label)
        plt.plot(raw_dict[stim], label='raw')
        plt.legend()
        plt.savefig(figPath)
        plt.close()

def sanity_check_mean_world_vid(full_world_cam_dict, world_downsample_ms, original_sample_rate_ms, save_folder):
    fps_rate = int(1000/world_downsample_ms)
    world_cam_downsample_mult = int(world_downsample_ms/original_sample_rate_ms)
    for stim in full_world_cam_dict.keys():
        tbs_to_sample = np.arange(0, len(full_world_cam_dict[stim]['all frames, weighted sum']), world_cam_downsample_mult)
        downsampled_mean_frames = []
        for current_tb in tbs_to_sample:
            this_mean_frame = full_world_cam_dict[stim]['all frames, weighted sum'][current_tb]/full_world_cam_dict[stim]['weights'][current_tb]
            downsampled_mean_frames.append(this_mean_frame)
        # reshape into original world cam dimensions
        downsampled_mean_frames_reshaped = []
        for frame in downsampled_mean_frames:
            reshaped_frame = np.reshape(frame,(120,160))
            downsampled_mean_frames_reshaped.append(reshaped_frame)
        downsampled_mean_frames_reshaped = np.array(downsampled_mean_frames_reshaped)
        # save as mp4 video file
        write_path = os.path.join(save_folder, 'Stim%d_MeanWorldCam.mp4'%(stim))
        end_tbucket = len(downsampled_mean_frames_reshaped)
        # temporarily switch matplotlib backend in order to write video
        plt.switch_backend("Agg")
        # Set up formatting for the movie files
        Writer = animation.writers['ffmpeg']
        FF_writer = animation.FFMpegWriter(fps=fps_rate, codec='h264', metadata=dict(artist='Danbee Kim', album='Surprising Minds'))
        print('Drawing frames...')
        fig = plt.figure()
        i = 0
        im = plt.imshow(downsampled_mean_frames_reshaped[i], cmap='gray', animated=True)
        def updatefig(*args):
            global i
            if (i<end_tbucket-1):
                i += 1
            else:
                i = 0
            im.set_array(downsampled_mean_frames_reshaped[i])
            return im,
        ani = animation.FuncAnimation(fig, updatefig, frames=len(downsampled_mean_frames_reshaped), interval=world_downsample_ms, blit=True)
        print("Writing average world video frames to {path}...".format(path=write_path))
        ani.save(write_path, writer=FF_writer)
        plt.close(fig)
        print("Finished writing!")
        # restore default matplotlib backend
        plt.switch_backend('TkAgg')

##########################################################
# BEGIN SCRIPT
##########################################################
if __name__=='__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--a", nargs='?', default="no_vid_output")
    parser.add_argument("--loc", nargs='?', default='laptop')
    args = parser.parse_args()
    ###################################
    # SOURCE DATA AND OUTPUT FILE LOCATIONS
    ###################################
    root_folder, plots_folder, monthly_mean_lums_folders, output_folders = load_data(args.loc)
    display_latency_folder = output_folders[0]
    mean_world_cam_vids_folder = output_folders[1]
    raw_v_world_sanity_check_folder = output_folders[2]
    logging.info('\n MAIN DATA FOLDER: %s \n PLOTS FOLDER: %s \n MONTHLY MEAN RAW LIVE AND WORLD CAM STIMULI DATA FOLDER: %s' % (root_folder, plots_folder, monthly_mean_lums_folders))
    print('MAIN DATA FOLDER: %s \n PLOTS FOLDER: %s \n MONTHLY MEAN RAW LIVE AND WORLD CAM STIMULI DATA FOLDER: %s' % (root_folder, plots_folder, monthly_mean_lums_folders))
    ###################################
    # TIMING/SAMPLING VARIABLES FOR DATA EXTRACTION
    ###################################
    # downsample = collect data from every 40ms or other multiples of 20
    downsampled_bucket_size_ms = 40
    original_bucket_size_in_ms = 4
    max_length_of_stim_vid = 60000 # milliseconds
    no_of_time_buckets = max_length_of_stim_vid/original_bucket_size_in_ms
    downsampled_no_of_time_buckets = max_length_of_stim_vid/downsampled_bucket_size_ms
    downsample_multiplier = int(downsampled_bucket_size_ms/original_bucket_size_in_ms)
    ###################################
    # STIMULI VID INFO
    ###################################
    stim_vids = [24.0, 25.0, 26.0, 27.0, 28.0, 29.0]
    stim_name_to_float = {"Stim24": 24.0, "Stim25": 25.0, "Stim26": 26.0, "Stim27": 27.0, "Stim28": 28.0, "Stim29": 29.0}
    stim_float_to_name = {24.0: "Stim24", 25.0: "Stim25", 26.0: "Stim26", 27.0: "Stim27", 28.0: "Stim28", 29.0: "Stim29"}
    phase_names = ['calib', 'octo', 'unique1', 'unique2', 'unique3', 'unique4', 'unique5', 'unique6']
    ###############################################################
    # Load mean monthly raw live stim and world cam luminance files
    # extract monthly mean world cam and raw live stim data (4ms resolution)
    ###############################################################
    # world cam data structure
    all_weighted_world_keyFrames = {key:{} for key in stim_vids}
    # raw live stim data structure
    all_weighted_rawLive_timebuckets = {key:{} for key in stim_vids}
    # collect length of each stimulus type in 4ms resolution
    supersampled_length_all_stims = {key:[] for key in stim_vids}
    # extract raw live stim and world cam data
    for monthly_mean_folder in monthly_mean_lums_folders:
        raw_live_stim_files = glob.glob(root_folder + os.sep + monthly_mean_folder + os.sep + '*_meanRawLiveStim_*.npy')
        world_cam_files = glob.glob(root_folder + os.sep + monthly_mean_folder + os.sep + '*_meanWorldCam_*.npy')
        # extract world cam data
        for world_cam in world_cam_files:
            stim_type = stim_name_to_float[os.path.basename(world_cam).split('_')[1]]
            vid_count = float(os.path.basename(world_cam).split('_')[-1][:-8])
            world_cam_frames = np.load(world_cam)
            extract_stim_data(world_cam_frames, 'world', all_weighted_world_keyFrames[stim_type])
        # extract raw live data
        for raw_live_stim in raw_live_stim_files:
            stim_type = stim_name_to_float[os.path.basename(raw_live_stim).split('_')[1]]
            vid_count = float(os.path.basename(raw_live_stim).split('_')[-1][:-8])
            raw_live_array = np.load(raw_live_stim)
            supersampled_length_all_stims[stim_type].append(len(raw_live_array))
            extract_stim_data(raw_live_array, 'raw', all_weighted_rawLive_timebuckets[stim_type])
    ########################################################
    # Calculate full dataset mean world cam for each stimulus
    ########################################################
    logging.info('Calculating full dataset mean world camera for each unique stimulus...')
    print('Calculating full dataset mean world camera for each unique stimulus...')
    weighted_sums_world_keyFrames = calculate_weighted_sums(all_weighted_world_keyFrames, 'world')
    # Fill in gaps between keyframes in world cam
    weighted_sums_world_all_frames = {key:{} for key in stim_vids}
    for stim in weighted_sums_world_keyFrames.keys():
        ordered_keyframes = sorted(weighted_sums_world_keyFrames[stim].keys())
        this_stim_all_weighted_frames = []
        this_stim_all_weights = []
        for i, keyframe in enumerate(ordered_keyframes):
            if i==0 and keyframe==0:
                this_stim_all_weighted_frames.append(weighted_sums_world_keyFrames[stim][keyframe]['keyframe, weighted sum'])
                this_stim_all_weights.append(weighted_sums_world_keyFrames[stim][keyframe]['summed weight'])
            elif i==0 and keyframe!=0:
                for frame in range(keyframe-1):
                    this_stim_all_weighted_frames.append(np.nan)
                    this_stim_all_weights.append(np.nan)
                this_stim_all_weighted_frames.append(weighted_sums_world_keyFrames[stim][keyframe]['keyframe, weighted sum'])
                this_stim_all_weights.append(weighted_sums_world_keyFrames[stim][keyframe]['summed weight'])
            else:
                prev_keyframe = ordered_keyframes[i-1]
                if keyframe - prev_keyframe > 1:
                    for frame in range(prev_keyframe, keyframe-1):
                        this_stim_all_weighted_frames.append(weighted_sums_world_keyFrames[stim][prev_keyframe]['keyframe, weighted sum'])
                        this_stim_all_weights.append(weighted_sums_world_keyFrames[stim][prev_keyframe]['summed weight'])
                this_stim_all_weighted_frames.append(weighted_sums_world_keyFrames[stim][keyframe]['keyframe, weighted sum'])
                this_stim_all_weights.append(weighted_sums_world_keyFrames[stim][keyframe]['summed weight'])
        full_length_this_stim = np.min(supersampled_length_all_stims[stim])
        if ordered_keyframes[-1] < full_length_this_stim:
            for frame in range(ordered_keyframes[-1], full_length_this_stim):
                this_stim_all_weighted_frames.append(weighted_sums_world_keyFrames[stim][ordered_keyframes[-1]]['keyframe, weighted sum'])
                this_stim_all_weights.append(weighted_sums_world_keyFrames[stim][ordered_keyframes[-1]]['summed weight'])
        weighted_sums_world_all_frames[stim] = {'all frames, weighted sum':this_stim_all_weighted_frames, 'weights':this_stim_all_weights}
    # Calculate weighted mean frame and weighted mean luminance for each world cam timebucket
    world_all_weighted_mean_frames = {key:None for key in stim_vids}
    world_all_weighted_mean_luminance = {key:None for key in stim_vids}
    for stim in weighted_sums_world_all_frames.keys():
        this_stim_weighted_mean_frames = []
        this_stim_weighted_mean_lum = []
        for tb, summed_frame in enumerate(weighted_sums_world_all_frames[stim]['all frames, weighted sum']):
            this_tb_weighted_mean_frame = summed_frame/weighted_sums_world_all_frames[stim]['weights'][tb]
            this_stim_weighted_mean_frames.append(this_tb_weighted_mean_frame)
            this_tb_weighted_mean_luminance = np.sum(this_tb_weighted_mean_frame)
            this_stim_weighted_mean_lum.append(this_tb_weighted_mean_luminance)
        world_all_weighted_mean_frames[stim] = np.array(this_stim_weighted_mean_frames)
        world_all_weighted_mean_luminance[stim] = np.array(this_stim_weighted_mean_lum)
    ########################################################
    # Find timebuckets marking start and end of each phase
    # UNDER CONSTRUCTION
    ########################################################
    if args.a == 'MOI':
        reshaped_world_all_weighted_mean_frames = {key:None for key in stim_vids}
        for stim in world_all_weighted_mean_frames.keys():
            reshaped_frames = []
            for i, frame in enumerate(world_all_weighted_mean_frames[stim]):
                reshaped_frame = np.reshape(frame,(120,160))
                reshaped_frames.append(reshaped_frame)
                # figure path and title
                figPath = os.path.join(mean_world_cam_vids_folder, 'Stim%d_meanWorldCamSanityCheck_tb%06d_4msResolution.png'%(stim, i))
                figTitle = 'Stim%d: mean world cam sanity check \n timebucket: %06d'%(stim, i)
                plt.figure(figsize=(9, 9), dpi=200)
                plt.suptitle(figTitle, fontsize=12, y=0.98)
                plt.imshow(reshaped_frame)
                plt.savefig(figPath)
                plt.close()

            reshaped_world_all_weighted_mean_frames[stim] = np.array(reshaped_frames)
        stim_to_check = input('Which unique stimulus would you like to check for moments of interest?')
        print('Checking %s'%(stim_to_check))
        go_to_timebucket = input('Jump to timebucket:')
        embed()

def display_mean_world_vid_frame(world_vid_frames_dict, stim_num, timebucket):
    plt.imshow(world_vid_frames_dict[stim_num][timebucket])
    plt.show()

    ########################################################
    # Calculate full dataset raw live for each stimulus
    ########################################################
    logging.info('Calculating full dataset raw live vid for each unique stimulus...')
    print('Calculating full dataset raw live vid for each unique stimulus...')
    weighted_sums_raw_timebuckets = calculate_weighted_sums(all_weighted_rawLive_timebuckets, 'raw')
    # Calculate weighted mean luminance for each raw live stim timebucket
    raw_all_weighted_mean_luminance = {key:None for key in stim_vids}
    for stim in weighted_sums_raw_timebuckets.keys():
        this_stim_weighted_mean_lum = []
        for tb in sorted(weighted_sums_raw_timebuckets[stim].keys()):
            this_tb_weighted_mean_lum = weighted_sums_raw_timebuckets[stim][tb]['luminance, weighted sum']/weighted_sums_raw_timebuckets[stim][tb]['summed weight']
            this_stim_weighted_mean_lum.append(this_tb_weighted_mean_lum)
        raw_all_weighted_mean_luminance[stim] = np.array(this_stim_weighted_mean_lum)
    ########################################################
    # Calculate display latency
    # this should be when the mean luminance in world cam drops more than 400,000
    # save as binary file output
    ########################################################
    all_true_start_tb = []
    all_true_start_tb_dict = {key:None for key in stim_vids}
    for stim in world_all_weighted_mean_luminance.keys():
        true_calib_start = None
        for tb, mean_lum in enumerate(world_all_weighted_mean_luminance[stim]):
            if true_calib_start==None:
                true_calib_start = [tb, mean_lum]
            elif abs(true_calib_start[1]-mean_lum)>400000:
                true_calib_start = [tb, mean_lum]
                break
        all_true_start_tb.append([stim, true_calib_start[0]])
        all_true_start_tb_dict[stim] = true_calib_start[0]
        logging.info('Display latency for stim %d: %d'%(stim, true_calib_start[0]))
        print('Display latency for stim %d: %d'%(stim, true_calib_start[0]))
    all_true_start_tb_array = np.array(all_true_start_tb)
    logging.info('Writing display latencies to binary file...')
    print('Writing display latencies to binary file...')
    display_latency_output = display_latency_folder + os.sep + 'displayLatencies.npy'
    np.save(display_latency_output, all_true_start_tb_array)
    ########################################################
    # sanity check: plot mean world cam (cropped to "real start") + full raw live for each stimulus to visually check display latency
    # save plots
    ########################################################
    # CROP LUMINANCE ARRAYS BASED ON all_true_start_tb
    world_all_weighted_mean_luminance_cropped = {key:None for key in stim_vids}
    for stim in all_true_start_tb_dict.keys():
        true_start = all_true_start_tb_dict[stim]
        cropped_lum_array = world_all_weighted_mean_luminance[stim][true_start:]
        world_all_weighted_mean_luminance_cropped[stim] = cropped_lum_array
    # plot sanity check
    logging.info('Saving sanity check plots of mean luminance for world cam versus raw live stim...')
    print('Saving sanity check plots of mean luminance for world cam versus raw live stim...')
    sanity_check_world_v_rawLive(world_all_weighted_mean_luminance_cropped, 'cropped', raw_all_weighted_mean_luminance, original_bucket_size_in_ms, raw_v_world_sanity_check_folder)
    ########################################################
    # save full dataset mean world cam video
    # downsample mean world cam video for 25 fps (one frame every 40 ms)
    ########################################################
    if args.a == 'no_vid_output':
        logging.info('No mean world cam video output saved. To save mean world cam video, run with optional input --a vid_output.')
        print('No mean world cam video output saved. To save mean world cam video, run with optional input --a vid_output.')
    elif args.a == 'vid_output':
        logging.info('Saving sanity check videos of mean luminance for world cam...')
        print('Saving sanity check videos of mean luminance for world cam...')
        sanity_check_mean_world_vid(weighted_sums_world_all_frames, downsampled_bucket_size_ms, original_bucket_size_in_ms, mean_world_cam_vids_folder)
    else:
        logging.warning('%s is not a valid optional input to this script! \n Completing script without generating mean world cam video output...' % (args.a))
        print('%s is not a valid optional input to this script! \n Completing script without generating mean world cam video output...' % (args.a))
# FIN


from keras.layers import Conv1D,Dense, Activation
from keras.models import Sequential, load_model
from keras.optimizers import Adam
import tensorflow as tf
import numpy as np
import os
import matplotlib.pyplot as plt
import pandas as pd
from joblib import dump, load
import sys

def setEvaluationFilesPath(video):
    return '../evaluationFiles/'+video+'/rep_10/'

MODEL_NAME='GreenABR'
if len(sys.argv)<2 :
    print('Please run the script with one of the test videos (tos,bbb,doc). i.e. python evaluate.py tos')
    quit()
else:
    VIDEO=sys.argv[1]
    if VIDEO not in ['tos','bbb','doc']:
        print('You must use one of the test videos, [tos,bbb,doc]')
        quit()
S_DIM=7
A_DIM = 10

VIDEO_BIT_RATE = [150,300, 750, 1200, 1850, 2850, 4300,6000,9000,12000] # Kbps
BUFFER_NORM_FACTOR = 10.0
CHUNK_TIL_VIDEO_END_CAP = 44
M_IN_K = 1000.0
M_IN_N = 1000000.0
REBUF_PENALTY = 4.3  # 1 sec rebuffering -> 3 Mbps
SMOOTH_PENALTY = 1
DEFAULT_QUALITY = 1  # default video quality without agent
RANDOM_SEED = 42
RAND_RANGE = 1000

REWARD_FOLDER='reward_logs/'
TEST_LOG_FOLDER = 'test_results/'
TEST_TRACES = '../test_sim_traces/'
EVAL_FILE_PATH=setEvaluationFilesPath(VIDEO)
PHONE_VMAF = pd.read_csv(EVAL_FILE_PATH+'vmaf_phone.csv')
REGULAR_VMAF = pd.read_csv(EVAL_FILE_PATH+'vmaf_phone.csv')
POWER_ATTRIBUTES= pd.read_csv('../power_attributes.csv')
POWER_MES=pd.read_csv(EVAL_FILE_PATH+'power_measurements.csv')

BITRATE_MAX=12000.0
FILE_SIZE_MAX=1775324.0
QUALITY_MAX=100.0
MOTION_MAX=20.15
PIXEL_RATE_MAX=2073600.0
POWER_MAX=1690.0

def load_trace(cooked_trace_folder=TEST_TRACES):
    cooked_files = os.listdir(cooked_trace_folder)
    if '.DS_Store' in cooked_files:
        cooked_files.remove('.DS_Store')
    all_cooked_time = []
    all_cooked_bw = []
    all_file_names = []
    for cooked_file in cooked_files:
        file_path = cooked_trace_folder + cooked_file
        cooked_time = []
        cooked_bw = []
        # print file_path
        with open(file_path, 'r') as f:
            for line in f:
                parse = line.split()
                if len(parse)>0:
#                     print(parse[0], parse[1])
                    cooked_time.append(float(parse[0]))
                    cooked_bw.append(float(parse[1]))
        all_cooked_time.append(cooked_time)
        all_cooked_bw.append(cooked_bw)
        all_file_names.append(cooked_file)

    return all_cooked_time, all_cooked_bw, all_file_names

# power model parameters for network power
p_alpha = 210
p_betha = 28

SEGMENT_SIZE = 4.0
power_threshold = 2500
byte_to_KB = 1000
KB_to_MB=1000.0

def Estimate_Network_Power_Consumption(thr, chunk_file_size):
    return (chunk_file_size * (p_alpha*1/thr+p_betha))

MIN_ENERGY=1000
NORMALIZATION_SCALAR=1000
SCALING_FACTOR=-1.2
def Calculate_Energy_Penalty(energy_chunk):
    penalty= ((energy_chunk-MIN_ENERGY)/NORMALIZATION_SCALAR)**2/SCALING_FACTOR
    return min(0,penalty)

def baseline_model():
    # create model
    model = Sequential()
    model.add(Dense(7, input_dim=5, kernel_initializer='normal', activation='relu'))
    model.add(Dense(20, kernel_initializer='normal', activation='relu'))
    model.add(Dense(1, kernel_initializer='normal'))
    # Compile model
    model.compile(loss='mean_squared_error', optimizer='adam')
    return model

class Environment:
    def __init__(self, all_cooked_time, all_cooked_bw, random_seed=RANDOM_SEED):
        self.MILLISECONDS_IN_SECOND = 1000.0
        self.B_IN_MB = 1000000.0
        self.BITS_IN_BYTE = 8.0
        self.RANDOM_SEED = 42
        self.VIDEO_CHUNCK_LEN = 4000.0  # millisec, every time add this amount to buffer
        self.BITRATE_LEVELS = A_DIM
        self.TOTAL_VIDEO_CHUNCK = 44
        self.BUFFER_THRESH = 60000.0  # millisec, max buffer limit
        self.DRAIN_BUFFER_SLEEP_TIME = 500.0  # millisec
        self.PACKET_PAYLOAD_PORTION = 0.95
        self.LINK_RTT = 80  # millisec
        self.PACKET_SIZE = 1500  # bytes
        self.VIDEO_SIZE_FILE = EVAL_FILE_PATH+'video_size_'
        assert len(all_cooked_time) == len(all_cooked_bw)

        self.local_power_model=load('../power_model.joblib')

        np.random.seed(random_seed)

        self.all_cooked_time = all_cooked_time
        self.all_cooked_bw = all_cooked_bw

        self.video_chunk_counter = 0
        self.buffer_size = 0

        # pick a random trace file
        self.trace_idx = 0
        self.cooked_time = self.all_cooked_time[self.trace_idx]
        self.cooked_bw = self.all_cooked_bw[self.trace_idx]

        # randomize the start point of the trace
        # note: trace file starts with time 0
        self.mahimahi_start_ptr = 1
        self.mahimahi_ptr = 1
        self.last_mahimahi_time = self.cooked_time[self.mahimahi_ptr - 1]

        self.video_size = {}  # in bytes
        for bitrate in range(self.BITRATE_LEVELS):
            self.video_size[bitrate] = []
            with open(self.VIDEO_SIZE_FILE + str(bitrate)) as f:
                for line in f:
                    self.video_size[bitrate].append(int(line.split()[0]))

    def reset(self):
        self.video_chunk_counter = 0
        self.buffer_size = 0

        # pick a random trace file
        # self.trace_idx = 0
        self.cooked_time = self.all_cooked_time[self.trace_idx]
        self.cooked_bw = self.all_cooked_bw[self.trace_idx]

        # randomize the start point of the trace
        # note: trace file starts with time 0
        self.mahimahi_ptr = 1
        self.last_mahimahi_time = self.cooked_time[self.mahimahi_ptr - 1]

        self.video_size = {}  # in bytes
        for bitrate in range(self.BITRATE_LEVELS):
            self.video_size[bitrate] = []
            with open(self.VIDEO_SIZE_FILE + str(bitrate)) as f:
                for line in f:
                    self.video_size[bitrate].append(int(line.split()[0]))

        return np.zeros(S_DIM)

    def get_video_chunk(self, quality):

        assert quality >= 0
        assert quality < self.BITRATE_LEVELS

        video_chunk_size = self.video_size[quality][self.video_chunk_counter]

        # use the delivery opportunity in mahimahi
        delay = 0.0  # in ms
        video_chunk_counter_sent = 0  # in bytes

        while True:  # download video chunk over mahimahi
            throughput = self.cooked_bw[self.mahimahi_ptr] * self.B_IN_MB / self.BITS_IN_BYTE
            duration = self.cooked_time[self.mahimahi_ptr] - self.last_mahimahi_time

            packet_payload = throughput * duration * self.PACKET_PAYLOAD_PORTION

            if video_chunk_counter_sent + packet_payload > video_chunk_size:

                fractional_time = (video_chunk_size - video_chunk_counter_sent) /throughput / self.PACKET_PAYLOAD_PORTION
                delay += fractional_time
                self.last_mahimahi_time += fractional_time
                assert(self.last_mahimahi_time <= self.cooked_time[self.mahimahi_ptr])
                break

            video_chunk_counter_sent += packet_payload
            delay += duration
            self.last_mahimahi_time = self.cooked_time[self.mahimahi_ptr]
            self.mahimahi_ptr += 1

            if self.mahimahi_ptr >= len(self.cooked_bw):
            # loop back in the beginning
            # note: trace file starts with time 0
                self.mahimahi_ptr = 1
                self.last_mahimahi_time = 0

        delay *= self.MILLISECONDS_IN_SECOND
        delay += self.LINK_RTT

    # add a multiplicative noise to the delay
        # delay *= np.random.uniform(self.NOISE_LOW, self.NOISE_HIGH)

        # rebuffer time
        rebuf = np.maximum(delay - self.buffer_size, 0.0)

        # update the buffer
        self.buffer_size = np.maximum(self.buffer_size - delay, 0.0)

        # add in the new chunk
        self.buffer_size += self.VIDEO_CHUNCK_LEN

        # sleep if buffer gets too large
        sleep_time = 0
        if self.buffer_size > self.BUFFER_THRESH:
          # exceed the buffer limit
          # we need to skip some network bandwidth here
          # but do not add up the delay
            drain_buffer_time = self.buffer_size - self.BUFFER_THRESH
            sleep_time = np.ceil(drain_buffer_time / self.DRAIN_BUFFER_SLEEP_TIME) * self.DRAIN_BUFFER_SLEEP_TIME
            self.buffer_size -= sleep_time

            while True:
                duration = self.cooked_time[self.mahimahi_ptr] - self.last_mahimahi_time
                if duration > sleep_time / self.MILLISECONDS_IN_SECOND:
                    self.last_mahimahi_time += sleep_time / self.MILLISECONDS_IN_SECOND
                    break
                sleep_time -= duration * self.MILLISECONDS_IN_SECOND
                self.last_mahimahi_time = self.cooked_time[self.mahimahi_ptr]
                self.mahimahi_ptr += 1

                if self.mahimahi_ptr >= len(self.cooked_bw):
                  # loop back in the beginning
                  # note: trace file starts with time 0
                    self.mahimahi_ptr = 1
                    self.last_mahimahi_time = 0

        # the "last buffer size" return to the controller
        # Note: in old version of dash the lowest buffer is 0.
        # In the new version the buffer always have at least
        # one chunk of video
        return_buffer_size = self.buffer_size

        self.video_chunk_counter += 1
        video_chunk_remain = self.TOTAL_VIDEO_CHUNCK - self.video_chunk_counter

        end_of_video = False
        if self.video_chunk_counter >= self.TOTAL_VIDEO_CHUNCK:
            end_of_video = True
            self.buffer_size = 0
            self.video_chunk_counter = 0

            # pick a random trace file
            self.trace_idx += 1
            if self.trace_idx >= len(self.all_cooked_time):
                self.trace_idx = 0
            self.cooked_time = self.all_cooked_time[self.trace_idx]
            self.cooked_bw = self.all_cooked_bw[self.trace_idx]

            # randomize the start point of the video
            # note: trace file starts with time 0
            self.mahimahi_ptr = self.mahimahi_start_ptr
            self.last_mahimahi_time = self.cooked_time[self.mahimahi_ptr - 1]

        next_video_chunk_sizes = []
        for i in range(self.BITRATE_LEVELS):
            next_video_chunk_sizes.append(self.video_size[i][self.video_chunk_counter])

        return self.video_chunk_counter,delay,sleep_time,return_buffer_size / self.MILLISECONDS_IN_SECOND, rebuf / self.MILLISECONDS_IN_SECOND, video_chunk_size,next_video_chunk_sizes, end_of_video, video_chunk_remain
    def normalize_parameters(self,bitrate,t_sec):
        time=self.video_chunk_counter*SEGMENT_SIZE+t_sec
        b_n= VIDEO_BIT_RATE[bitrate] / BITRATE_MAX
        d=POWER_ATTRIBUTES[(POWER_ATTRIBUTES['Bitrate']==VIDEO_BIT_RATE[bitrate]) &
                         (POWER_ATTRIBUTES['Time']==time)]
        f_n=d['FileSize']/FILE_SIZE_MAX
        q_n= d['Quality']/QUALITY_MAX
        m_n=d['Motion']/MOTION_MAX
        p_n=d['PixelRate']/PIXEL_RATE_MAX
        return np.reshape(np.array([b_n,f_n,q_n,m_n,p_n]),(1,5))

    def calculate_local_energy(self, bitrate):
        total_energy=0
        for i in range(int(SEGMENT_SIZE)):
            pars=self.normalize_parameters(bitrate,i)
            pred=self.local_power_model.predict(pars)
            power=pred[0]*POWER_MAX # equal to energy as it is for 1 second
            total_energy+=power
        return total_energy

    def step(self, bit_rate, last_bitrate):
        video_chunk_counter,delay, sleep_time, buffer_size, rebuf, video_chunk_size, next_video_chunk_sizes, end_of_video, video_chunk_remain = self.get_video_chunk(bit_rate)

        throughput=float(video_chunk_size) / float(delay) / M_IN_K
        new_state= np.zeros(S_DIM)
        new_state[0]=VIDEO_BIT_RATE[bit_rate] / float(np.max(VIDEO_BIT_RATE))
        new_state[1]=buffer_size / BUFFER_NORM_FACTOR
        new_state[2]=throughput
        new_state[3]=float(delay) / M_IN_K / BUFFER_NORM_FACTOR
        new_state[4]=np.minimum(video_chunk_remain,CHUNK_TIL_VIDEO_END_CAP) / float(CHUNK_TIL_VIDEO_END_CAP)

        log=np.zeros(6)
        log[0]=self.TOTAL_VIDEO_CHUNCK if video_chunk_counter==0 else video_chunk_counter
        log[1]=delay
        log[2]=sleep_time
        log[3]=buffer_size
        log[4]=rebuf
        log[5]=video_chunk_size

        quality=(PHONE_VMAF['VMAF_' + str(bit_rate+1)][log[0]-1])
        new_state[6]=quality

        estimated_energy=(POWER_MES['P_' + str(bit_rate+1)][log[0]-1])
        new_state[5]= estimated_energy

        reward=0
        quality=quality/20.0
        quality_reward=quality

        if log[0]==1:
            rebuffer_penalty=0.0
        else:
            rebuffer_penalty= REBUF_PENALTY*rebuf

        if log[0]==1:
            smooth_penalty=0.0
        else:
            smooth_penalty=SMOOTH_PENALTY* np.abs((PHONE_VMAF['VMAF_' + str(bit_rate+1)][log[0]-1]/20.0)- (PHONE_VMAF['VMAF_' + str(last_bitrate+1)][log[0]-2]/20.0))

        reward=quality_reward - rebuffer_penalty - smooth_penalty

        energy_penalty=Calculate_Energy_Penalty(estimated_energy)

        return new_state, reward, end_of_video, estimated_energy, video_chunk_size, quality, rebuffer_penalty, smooth_penalty,energy_penalty, log

class ReplayBuffer(object):
    def __init__(self, max_size, input_shape, n_actions, discrete=False):
        self.mem_size = max_size
        self.mem_cntr = 0
        self.discrete = discrete
        self.state_memory = np.zeros((self.mem_size, input_shape))
        self.new_state_memory = np.zeros((self.mem_size, input_shape))
        dtype = np.int8 if self.discrete else np.float32
        self.action_memory = np.zeros((self.mem_size, n_actions), dtype=np.int8)
        self.reward_memory = np.zeros(self.mem_size)
        self.terminal_memory = np.zeros(self.mem_size, dtype=np.float32)

    def store_transition(self, state, action, reward, state_, done):
        index = self.mem_cntr % self.mem_size
        self.state_memory[index] = state
        self.new_state_memory[index] = state_
        # store one hot encoding of actions, if appropriate
        if self.discrete:
            actions = np.zeros(self.action_memory.shape[1])
            actions[action] = 1.0
            self.action_memory[index] = actions
        else:
            self.action_memory[index] = action
        self.reward_memory[index] = reward
        self.terminal_memory[index] = 1 - done
        self.mem_cntr += 1

    def sample_buffer(self, batch_size):
        max_mem = min(self.mem_cntr, self.mem_size)
        batch = np.random.choice(max_mem, batch_size)

        states = self.state_memory[batch]
        actions = self.action_memory[batch]
        rewards = self.reward_memory[batch]
        states_ = self.new_state_memory[batch]
        terminal = self.terminal_memory[batch]

        return states, actions, rewards, states_, terminal

def build_dqn(lr, n_actions, input_dims, fc1_dims, fc2_dims):
    model = Sequential([
                Dense(fc1_dims, input_shape=(input_dims,)),
                Activation('relu'),
                Dense(fc2_dims),
                Activation('relu'),
                Dense(128),
                Activation('relu'),
                Dense(64),
                Activation('relu'),
                Dense(n_actions)])

    model.compile(optimizer=Adam(lr=lr), loss='mse')
    return model

class DDQNAgent(object):
    def __init__(self, alpha, gamma, n_actions, epsilon, batch_size,
                 input_dims,epsilon_dec=0.996,  epsilon_end=0.01,
                 mem_size=50000, fname=MODEL_NAME, replace_target=100):
        self.action_space = [i for i in range(n_actions)]
        self.n_actions = n_actions
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_dec = epsilon_dec
        self.epsilon_min = epsilon_end
        self.batch_size = batch_size
        self.model_file = MODEL_NAME+".h5"
        self.replace_target = replace_target
        self.memory = ReplayBuffer(mem_size, input_dims, n_actions,
                                   discrete=True)
        self.q_eval = build_dqn(alpha, n_actions, input_dims,256, 256)
        self.q_target = build_dqn(alpha, n_actions, input_dims, 256,256)

    def remember(self, state, action, reward, new_state, done):
        self.memory.store_transition(state, action, reward, new_state, done)

    def choose_action(self, state):
        state = state[np.newaxis, :]
        rand = np.random.random()
        if rand < self.epsilon:
            action = np.random.choice(self.action_space)
        else:
            actions = self.q_eval.predict(state)
            action = np.argmax(actions)

        return action

    def learn(self):
        if self.memory.mem_cntr > self.batch_size:
            state, action, reward, new_state, done = self.memory.sample_buffer(self.batch_size)

            action_values = np.array(self.action_space, dtype=np.int8)
            action_indices = np.dot(action, action_values)

            q_next = self.q_target.predict(new_state)
            q_eval = self.q_eval.predict(new_state)
            q_pred = self.q_eval.predict(state)

            max_actions = np.argmax(q_eval, axis=1)

            q_target = q_pred

            batch_index = np.arange(self.batch_size, dtype=np.int32)

            q_target[batch_index, action_indices] = reward + self.gamma*q_next[batch_index, max_actions.astype(int)]*done

            _ = self.q_eval.fit(state, q_target, verbose=0)

            self.epsilon = self.epsilon*self.epsilon_dec if self.epsilon > self.epsilon_min else self.epsilon_min
            if self.memory.mem_cntr % self.replace_target == 0:
                self.update_network_parameters()

    def update_network_parameters(self):
        self.q_target.set_weights(self.q_eval.get_weights())

    def save_model(self):
        self.q_eval.save(self.model_file)

    def load_model(self):
        self.q_eval = load_model(self.model_file)
        # if we are in evaluation mode we want to use the best weights for
        # q_target
        if self.epsilon == 0.0:
            self.update_network_parameters()

# In[175]:

def main():
    REWARD_MODEL=MODEL_NAME

    all_cooked_time, all_cooked_bw, all_file_names = load_trace(TEST_TRACES)

    input_dims=7

    env=Environment(all_cooked_time, all_cooked_bw, RANDOM_SEED)
    ddqn_agent = DDQNAgent(alpha=0.001, gamma=0.99, n_actions=A_DIM, epsilon=0.0,batch_size=64,input_dims=input_dims)

    ddqn_agent.load_model()

    log_path = TEST_LOG_FOLDER+VIDEO+'/log_'+REWARD_MODEL+'_reward_' + all_file_names[env.trace_idx]
    os.makedirs(os.path.dirname(log_path), exist_ok=True)
    log_file = open(log_path, 'w')

    reward_log_path = 'reward_logs/'+VIDEO+'/log_'+REWARD_MODEL+'_reward_' + all_file_names[env.trace_idx]
    os.makedirs(os.path.dirname(reward_log_path), exist_ok=True)
    reward_log_file = open(reward_log_path, 'w')

    trace_log_path='reward_logs/'+VIDEO+"_log_"+REWARD_MODEL
    os.makedirs(os.path.dirname(trace_log_path), exist_ok=True)
    trace_log=open(trace_log_path, 'w')

    video_count=0
    for video_count in range(len(all_file_names)):
        done = False
        score = 0
        total_energy=0
        total_data=0
        last_bit_rate=1
        total_quality=0
        total_rebuffer_pen=0
        total_rebuffer_time=0
        total_smooth_pen=0
        total_smooth_time=0
        total_energy_penalty=0
        time_stamp=0
        first_chunk=True
        log_path =  TEST_LOG_FOLDER+VIDEO+'/log_'+REWARD_MODEL+'_reward_' + all_file_names[env.trace_idx]
        log_file = open(log_path, 'w')
        log_file.write('video_chunk' + '\t' +
                           'bitrate' + '\t' +
                           'buffer_size' + '\t' +
                           'rebuf' + '\t' +
                           'video_chunk_size' + '\t' +
                           'delay' + '\t' +
                           'phone_vmaf' + '\t' +
                           'regular_vmaf' + '\t' +
                           'energy' + '\t' +
                           'reward' + '\n'
                           )

        observation=env.reset()
        while not done:
            if first_chunk:
                action=0
            else :
                action = ddqn_agent.choose_action(observation)
            observation_, reward, done, energy, data,quality, rebuffer_penalty, smooth_penalty, energy_penalty, log= env.step(action,last_bit_rate)
            first_chunk=False
            score += reward
            total_energy+=energy
            total_data+=data
            total_quality+=quality
            total_rebuffer_pen=+rebuffer_penalty
            if rebuffer_penalty > 0:
                total_rebuffer_time+=1
            total_smooth_pen+=smooth_penalty
            if smooth_penalty > 0:
                total_smooth_time +=1
            total_energy_penalty+=energy_penalty

            time_stamp+=log[1]
            time_stamp+=log[2]
            # print("the chunk counter is ", str(log[0]))
            log_file.write(str(log[0]-1) + '\t' +
                           str(VIDEO_BIT_RATE[action]) + '\t' +
                           str(log[3]) + '\t' +
                           str(log[4]) + '\t' +
                           str(log[5]) + '\t' +
                           str(log[1]) + '\t' +
                           str(PHONE_VMAF['VMAF_' + str(action+1)][log[0]-1]) + '\t' +
                           str(REGULAR_VMAF['VMAF_' + str(action+1)][log[0]-1]) + '\t' +
                           str(energy) + '\t' +
                           str(reward) + '\n'
                           )

            log_file.flush()

            reward_log_file.write(str(log[0]) + '\t' +
                           str(time_stamp) + '\t' +
                           str(VIDEO_BIT_RATE[action]) + '\t' +
                           str(quality) + '\t' +
                           str(rebuffer_penalty) + '\t' +
                           str(smooth_penalty) + '\t' +
                           str(energy_penalty) + '\t' +
                           str(PHONE_VMAF['VMAF_' + str(action+1)][log[0]-1]) + '\t' +
                           str(REGULAR_VMAF['VMAF_' + str(action+1)][log[0]-1]) + '\t' +
                           str(energy) + '\t' +
                           str(reward) + '\n'
                           )

            reward_log_file.flush()
            observation = observation_
            last_bit_rate=action

        log_file.write('\n')
        log_file.close()
        reward_log_file.write('\n')
        reward_log_file.close()
        print('Completed trace ', all_file_names[env.trace_idx])
        print('Trace: ', video_count,'score: %.2f' % score)
        trace_log.write(str(video_count+1)+'\t'+
                        str(all_file_names[env.trace_idx]+'\t')+
                        str(score)+'\t'+
                        str(total_energy)+'\t'+
                        str(total_data)+'\t'+
                        str(total_quality)+'\t'+
                        str(total_rebuffer_pen)+'\t'+
                        str(total_rebuffer_time)+'\t'+
                        str(total_smooth_pen)+'\t'+
                        str(total_smooth_time)+'\t'+
                        str(total_energy_penalty)+'\n')
        trace_log.flush()
        video_count+=1
        reward_log_path = 'reward_logs/'+VIDEO+'/log_'+REWARD_MODEL+'_reward_' + all_file_names[env.trace_idx]
        reward_log_file = open(reward_log_path, 'w')

    trace_log.write('\n')
    trace_log.close()

if __name__ == "__main__":
    main()


# -*- coding: utf-8 -*-
from pathlib import Path
from collections import defaultdict
import numpy as np
import attr
import matplotlib
import Bio.PDB as PDB

from typing import Optional, Tuple, Dict, Set
Residue = PDB.Residue
from Bio.PDB.vectors import Vector
from geometry import project_on_plane

class PDBFile:
    def __init__(self, fname: str, id: str = None):
        """
        Helper class to work with the pdb file. Takes a file name to a pdb
        file.
        """
        if fname[-3:] == 'pdb':
            self.PDB = PDB.PDBParser()
        elif fname[-3:] == 'cif':
            self.PDB = PDB.FastMMCIFParser()
        if id is None:
            id = fname

        self.struct = self.PDB.get_structure(
            id, fname)  #.child_list[0].child_dict['A']
        self.res = PDB.Selection.unfold_entities(self.struct, "R")
        self.kd = PDB.NeighborSearch(list(self.struct.get_atoms()))

    @classmethod
    def from_pdb_id(cls, pdb_id: str, save_dir=None):
        """Download PDB file from ID, will saved in the homedir"""
        if save_dir is None:
            pdbl = PDB.PDBList()
            save_dir = str(Path.home())
        fname = pdbl.retrieve_pdb_file(pdb_id,
                                       file_format="mmCif",
                                       pdir=save_dir)
        return cls(fname)

    def get_res(self, res_name: str):
        match_res = filter(lambda r: r.get_resname() == res_name, self.res)
        return next(match_res)

    def find_near(self,
                  res_name: str,
                  radius: float = 5,
                  res_type=None,
                  atom_names=None,
                  exclude_self=False):
        """Find residues near given residue, currently assumes the residues is
        unique"""
        match_res = filter(lambda r: r.get_resname() == res_name, self.res)
        center_res = next(match_res)
        nearby_res: Set[PDB.Residue] = set()
        for atom in center_res:
            if atom_names is not None:
                if atom.name not in atom_names:
                    continue
            nearby_res |= set(self.kd.search(atom.coord, radius, level="R"))
        if exclude_self:
            out = list(nearby_res - {center_res})
        else:
            out = list(nearby_res)
        if res_type is not None:
            out = [i for i in out if i.get_resname() == res_type]
        return out

    def calculate_angles(self,
                         residue: PDB.Residue,
                         other_tdm: PDB.Vector = None,
                         do_project: bool = True):
        res = ResModes(residue)
        if other_tdm is not None:
            for v in res.modes.values():
                v.angle = np.rad2deg(v.tdm.angle(other_tdm))
                assert v.angle is not None
                if do_project:
                    if v.angle > 90.:
                        v.angle = 180. - v.angle
        return res

    def atom_vec_angles(self,
                        a: PDB.Atom,
                        b: PDB.Atom,
                        tdm: PDB.Atom,
                        do_project: bool = True):
        vec = b.get_vector() - a.get_vector()
        angle = np.rad2deg(tdm.angle(vec))

        if do_project:
            if angle > 90.:
                angle = 180. - angle
        return angle

def make_coord_array(res: PDB.Residue) -> np.ndarray:
    out = np.empty((3, len(res.child_list)))
    for i, a in enumerate(res.get_atoms()):
        out[:, i] = a.get_coord()
    return out

@attr.s(auto_attribs=True)
class Vibration:
    """Single Vibration"""
    name: str
    freq_range: Tuple[float, float] = attr.ib()
    strength: float = attr.ib(0)
    tdm: Vector = attr.ib(None)
    origin: Optional[Vector] = attr.ib(None)
    angle: Optional[float] = attr.ib(None)

WATER_NAMES = ("HOH", "WAT", "TIP3W", "TIP", "LYS")

@attr.s(auto_attribs=True)
class ResModes:
    residue: PDB.Residue
    modes: Dict[str, Vibration] = attr.Factory(dict)
    use_protonation: bool = False

    def vec(self, a: Optional[str], b=None) -> PDB.Vector:
        "Helper function to extract atom-vectors and atom-atom vectors"
        if b is not None:
            return self.residue[a].get_vector() - self.residue[b].get_vector()
        else:
            return self.residue[a].get_vector()

    def get_modes(self, kind: str):
        iter = filter(lambda x: x.startswith(kind), self.modes)
        return list(iter)

    def __attrs_post_init__(self):
        "Calculate all known modes"
        resname = self.residue.get_resname()
        #if resname in WATER_NAMES:  # Dont handle water
        #    return
        if not self.residue.id[0] == ' ':
            return
        self.amide_one()
        if resname in ["ASP", "GLU"]:
            is_protonated = any(k in res for k in ("HE1", "HE2", "HD1", "HD2"))
            if not is_protonated or not self.use_protonation:
                self.carboxylate_antisym(resname)
            if is_protonated or not self.use_protonation:
                self.carbonyl_stretch(resname)

    def amide_one(self):
        """Estimates the direction of the Amide I vibration by taking
        the C=O vector of the petide group and rotate it accordingly."""
        own_id = self.residue.id[1]
        next_N = self.residue.parent[self.residue.id[1] + 1]['N'].get_vector()
        vec_cn = -(self.vec('C') - next_N)
        vec_co = self.vec('C', 'O')
        normal = (vec_cn**vec_co).normalized()
        for angle, name in zip((0, 10, 20), ("No rot", "D2O", "H2O")):
            rot_mat = PDB.rotaxis(np.deg2rad(-angle), normal)
            vec = vec_co.left_multiply(rot_mat)
            vib = Vibration(name='Amide 1 %s' % name,
                            freq_range=(1640, 1660),
                            strength=3,
                            tdm=vec,
                            angle=np.nan,
                            origin=self.vec('C'))
            self.modes[vib.name] = vib

    def carboxylate_antisym(self, resname="ASP"):
        """Returns an estimate of the antisymetric COO- vibration by
        returning the O-O vector."""
        if resname == "ASP":
            vec = self.vec('OD1', 'OD2')
            freq_range = (1565, 1600)
            origin = self.vec('OD1')
        elif resname == "GLU":
            vec = self.vec('OE1', 'OE2')
            freq_range = (1555, 1600)
            origin = self.vec('OE1')

        vib = Vibration(
            name="COO- AS ",
            freq_range=freq_range,
            strength=4,
            tdm=vec,
            angle=np.nan,
            origin=origin,
        )
        self.modes[vib.name] = vib

    def carbonyl_stretch(self, resname):
        """Returns the carbonyl stretching by returning the C=O vector"""

        if resname == "ASP":
            origin = self.vec('CG')
            freq_range = (1690, 1740)
            if "HD1" in res or not self.use_protonation:
                vec = self.vec('OD2', 'CG')
                vib = Vibration(
                    name="CO (OD2)",
                    freq_range=freq_range,
                    tdm=vec,
                    origin=origin,
                )
                self.modes[vib.name] = vib
            if "HD2" in res or not self.use_protonation:
                vec = self.vec('OD1', 'CG')
                vib = Vibration(
                    name="CO (OD1)",
                    freq_range=freq_range,
                    tdm=vec,
                    origin=origin,
                )
                self.modes[vib.name] = vib

        if resname == "GLU":
            origin = self.vec('CD')
            freq_range = (1690, 1740)
            if "HE1" in res or not self.use_protonation:
                vec = self.vec('OE2', 'CD')
                vib = Vibration(
                    name="CO (OE2)",
                    freq_range=freq_range,
                    tdm=vec,
                    origin=origin,
                )
                self.modes[vib.name] = vib
            if "HE2" in res or not self.use_protonation:
                vec = self.vec('OE1', 'CD')
                vib = Vibration(name="CO (OE1)",
                                freq_range=freq_range,
                                origin=origin,
                                tdm=vec)
                self.modes[vib.name] = vib

    def ring_breathmode(self, res_name):
        if res_name == 'TRP':
            v1 = self.residue["CD1"].get_vector()
            vec = v1 - self.residue["CG"].get_vector()
            freq_range = (1615, 1625)
            origin = self.vec("CG")

        if res_name == 'TYR':
            vec = self.vec("CZ", "CG")
            freq_range = (1490, 1510)
            origin = self.vec("CG")

        vib = Vibration(
            name="CC",
            freq_range=freq_range,
            strength=4,
            tdm=vec,
            origin=origin,
        )
        self.modes[vib.name] = vib

if __name__ == "__main__":

    struct = PDBFile.from_pdb_id('5ZIM')
    ns = struct.find_near("RET", 10, atom_names=('C15', 'C13'))

    retinal = struct.get_res('RET')
    vec = retinal['C14'].get_vector() - retinal['C10'].get_vector()

    by_type: Dict[str, Dict[str, float]] = defaultdict(dict)
    by_res = {}
    for res in ns:
        r = struct.calculate_angles(res, vec)
        resname = res.get_resname() + str(res.id[1])
        resname = resname[:1] + resname[1:3].lower() + resname[3:]
        for i in r.modes:
            by_type[i][resname] = r.modes[i].angle
        by_res[resname] = r

    waters = 404, 406, 407
    chain = struct.struct[0]['A']
    chain_dict = {res.id[1]: res for res in chain}
    print("Water angles\n-----------------")
    for i in waters:
        for j in waters:
            if i == j:
                pass
            else:
                a = chain_dict[i]['O']
                b = chain_dict[j]['O']
                print(i, j,  a-b, struct.atom_vec_angles(a, b, vec))

    print("CO stretch angles\n-----------------")
    for r in by_res:
        for m in by_res[r].get_modes('CO'):
            print(f"{r} {m} {by_res[r].modes[m].angle: .1f}")


import yaml
from util.prepare_dataset import load_dataset,create_dataloader
from util.functions import log_parser,batch_iterator, collapse_phn
from model.las_model import Listener,Speller
import numpy as np
from torch.autograd import Variable
import torch
import torch.nn as nn
import time
import os.path as path
import os as os
import matplotlib.pyplot as plt

# Load example config file for experiment
config_path = 'config/las_example_config.yaml'
conf = yaml.load(open(config_path,'r'))

listener_model_path = path.join("checkpoint3"+"listener")
speller_model_path = path.join("checkpoint3"+"speller")

import editdistance as ed
from glob import glob
import re as re
import os as os
import json as _json
import os.path as path
import soundfile as sf
import utils as utils
import random as _random
import utils.hangul as _hangul

from random import shuffle as shuffle
from importlib import import_module as import_module
from functools import wraps
from python_speech_features import mfcc
import scipy.io.wavfile as wav
from six.moves import xrange as range

from abc import ABC, abstractmethod

import pandas as pd
import random as random
import tensorflow as tf
from tqdm import tqdm

import pickle as _pickle
import os.path as _path
from random import shuffle as _shuffle
import numpy as _np

learning_rate=0.0001
n_classes=56

batch_size=4
feat_len =13
max_label_len=26   #max_label_len+2해서 sos eos 추가
max_time_step=448
pre_trained=True

ctc_labels = _hangul.JAMOS
labels = [" "] + ctc_labels

def load_data(data_dir):

  assert _path.isdir(data_dir)
  load_rawdata = getattr(import_module("utils.data"),"opendict")
  preprocess_audio = getattr(import_module("utils.audio"), "mfcc")
  preprocess_text = getattr(import_module("utils.hangul"), "jamo_token")
  cache_path = _path.join(data_dir, "opendict" + "_cache.pkl")

  print("Loading cached data.")
  with open(cache_path, "rb") as f:
      cache = _pickle.load(f)
  return [[audio, preprocess_text(text)] for (audio, text) in cache]

data=load_data('/home/suji93/data/')
data = list(filter(lambda x: x[0] is not None,data))

#max_len=max(list(map(lambda x:len(x[1]), data)))+2

def batch_generator(data, batch_size,idx=0, num_workers=1, bin_size=100, bin_num=10):

  tmp = []
  bin_idx = 0
  for _ in range(bin_num):
    tmp.append(list(filter(lambda x: bin_idx <= len(x) and len(x) < bin_idx+bin_size,data)))
    bin_idx += bin_size
  order = list(range(0, len(data), batch_size))
  while True:
    shuffle(order)
    for i in range(bin_num):
      shuffle(tmp[i])
    data = sum(tmp, [])
    for i in order:
      if i % num_workers == idx:
        yield list(zip(*data[i:(i+batch_size)]))

_random.seed(29520)
_random.shuffle(data)

n_train = int(len(data) * 0.85/batch_size) * batch_size
n_valid=n_train+int(len(data)*0.05)
n_test=n_valid+int(len(data)*0.1)
test_set=data[n_valid:n_test]
valid_set=data[n_train:n_valid]
batchs = {
      "train": batch_generator(data[:n_train], batch_size),
      "valid": batch_generator(data[n_train:n_valid], batch_size),
      "test": batch_generator(data[n_valid:n_test], batch_size)
  }

def OneHotEncode(Y,max_label_len,max_idx):
    new_y = np.zeros((len(Y),max_label_len,max_idx))
    for idx,label_seq in enumerate(Y):
        cnt = 1
        for label in label_seq:
            new_y[idx,cnt,label] = 1.0
            cnt += 1
        new_y[idx,cnt,1] = 1.0 # <eos>
    return new_y

#######training set##############

ctc_labels = _hangul.JAMOS
labels = [" "] + ctc_labels
jamo2index = {k:(v+2) for v,k in enumerate(labels)}
index2jamo = {(v+2):k for v,k in enumerate(labels)}

def encode(string):
    return list(map(lambda c: jamo2index[c], string))

def decode(string):
    return list(map(lambda c: index2jamo[c], string))

if pretrained:
    listener=torch.load(listener_model_path)
    speller=torch.load(speller_model_path)

else:
    listener = Listener(**conf['model_parameter'])
    speller = Speller(**conf['model_parameter'])
optimizer = torch.optim.Adam([{'params':listener.parameters()}, {'params':speller.parameters()}],
                              lr=conf['training_parameter']['learning_rate'])

def preprocess(data):
    audio=data[0]
    label=data[1]

    #zero padding
    input_len_feed = np.asarray(list(map(len, audio)))
    inputs_feed = np.zeros((len(audio), max_time_step, feat_len), np.float32)

    #one-hot
    preprocess_text=[encode(label[i]) for i in range(len(label))]
    label=OneHotEncode(preprocess_text,max_label_len,max_idx=n_classes)
    for i, l in enumerate(input_len_feed):
      inputs_feed[i, :l, :] = audio[i]
    return (inputs_feed, label)

def LetterErrorRate2(pred_y,true_y):
    ed_accumalate = []
    for p,t in zip(pred_y,true_y):
        compressed_t = [w for w in t if (w!=1 and w!=0)]
        compressed_p = []
        for p_w in p:
            if p_w == 0:
                continue
            if p_w == 1:
                break
            compressed_p.append(p_w)

        ed_accumalate.append(ed.eval(compressed_p[1:-1],compressed_t[1:-1])/len(compressed_t))
    return ed_accumalate

def model(batch_audio, batch_text, is_training):

     batch_input = Variable(torch.from_numpy(batch_audio)).type(torch.FloatTensor).cuda()
     batch_label = Variable(torch.from_numpy(np.array(batch_text)), requires_grad=False).cuda()
     objective = nn.CrossEntropyLoss(ignore_index=0).cuda()
     #objective = nn.NLLLoss(ignore_index=0)

     optimizer.zero_grad()
     listner_feature = listener(batch_input)

     if is_training:
         raw_pred_seq, attention_record = speller(listner_feature,ground_truth=batch_label,teacher_force_rate=tf_rate)

     else:
         raw_pred_seq, attention_record = speller(listner_feature,ground_truth=None,teacher_force_rate=0)

     #output_class_dim=56
     pred_y = torch.cat([torch.unsqueeze(each_y,1) for each_y in raw_pred_seq],1).view(-1,n_classes+1)
     true_y = torch.max(batch_label,dim=2)[1].view(-1)

     loss = objective(pred_y,true_y)

     if is_training:
         loss.backward()
         optimizer.step()

     batch_loss = loss.cpu().data.numpy()
     batch_ler = LetterErrorRate(torch.max(pred_y,dim=1)[1].cpu().data.numpy().reshape(batch_size,max_label_len),
                                    true_y.cpu().data.numpy().reshape(batch_size,max_label_len))

     pred=torch.max(pred_y,dim=1)[1].cpu().data.numpy().reshape(batch_size,max_label_len)
     true=true_y.cpu().data.numpy().reshape(batch_size,max_label_len)

     batch_ler=sum(batch_ler)/len(batch_ler)
     return batch_loss,batch_ler, pred, true

###############training##################

best_ler = 1.0
traing_log = open('log'+'train.log','w')

num_epochs=10
for epoch in range(num_epochs):

    epoch_head = time.time()
    loss = 0.0
    ler = 0.0
    val_loss = 0.0
    val_ler = 0.0

    # Teacher forcing rate linearly decay
    tf_rate = 0.8 - (0.8-0.0)

    for _ in range(int(n_train/batch_size)):

        audio_train,label_train=preprocess(list(next(batchs["train"])))
        train_loss,train_ler,tr_pred_y, tr_true_y=model(audio_train, label_train,True)

        loss += train_loss
        ler+= train_ler

    tr_loss=loss/int(n_train/batch_size)
    tr_ler=ler/int(n_train/batch_size)
    training_time = float(time.time()-epoch_head)

    print('Train Error Rate',tr_ler)
    print('Train loss',tr_loss)

    audio_valid,label_valid=preprocess(list(next(batchs["valid"])))
    val_loss,val_ler,val_pred_y, val_true_y=model(audio_valid, label_valid,False)

    # Checkpoint

    print('valid ler',val_ler)
    print('valid loss',val_loss)
    print('val prediction', val_pred_y)
    print('val true', val_true_y)
    torch.save(listener, listener_model_path)
    torch.save(speller, speller_model_path)

###################test####################################

dfa = {0:[-1,-1, 1, 1,-1, 1,-1,-1, 1, 1, 1,-1,-1,-1,-1,-1,-1,-1, 1, 1, 1,-1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 0],
  1:[-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1],
  2:[-1, 3, 3, 3, 3, 3, 3, 3, 3,-1, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,-1, 3, 3, 3, 3, 3,-1, 3, 3, 3, 3, 3,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 2],
  3:[ 0,-1, 1, 1,-1, 1,-1,-1, 1, 1, 1,-1,-1,-1,-1,-1,-1,-1, 1, 1, 1,-1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 3]}

def accepts(transitions,initial,accepting,s):
    state = initial
    for c in s:
        state = transitions[state][c]
        if state==-1:
            break
    return state in accepting

class WER:
    def __init__(self, true_data, pred_data):
        self.data1 = true_data
        self.data2 = pred_data

    def compress(self,label):
        i=0
        val=[]
        while label[i]!=1:
            val.append(label[i])
            i+=1
        return val[1:]

    def sep(self,text):
        mo=[]
        new=[]
        for i in text:
            if i!=2:
                new.append(i)
            else:
                mo.append(new)
                new=[]
        mo.append(new)
        return mo

    def count2(self, text):
        count=0
        for i in text:
            if i==2:
                count+=1
        return count

    def wer(self):
        text1=self.compress(self.data1)
        text2=self.compress(self.data2)

        min_num=min(self.count2(text1),self.count2(text2))

        score=0
        i=0
        #2의 개수
        while i<=min_num:

            if self.sep(text1)[i]!=self.sep(text2)[i]:
                score+=1
            else:
                score+=0
            i+=1
        total_wer=score/len(text1)
        return total_wer

ler=0.0
for _ in range(int(n_test/batch_size)):

    audio_test,label_test=preprocess(list(next(batchs["test"])))
    te_ler,te_pred_y, te_true_y=model(audio_test, label_test,False)

    ler+= te_ler

    wer=0
    for i in range(len(te_pred_y)):
        if accepts(dfa,0,{3},te_pred_y[i])==True:
            wer+=WER(te_true_y[i],te_pred_y[i])

test_wer=wer/int(n_test/batch_size)
print('test_wer',wer)

test_ler=ler/int(n_test/batch_size)
print('Testing Error Rate',test_ler)


##########################################################
# pytorch-kaldi v.0.1
# Mirco Ravanelli, Titouan Parcollet
# Mila, University of Montreal
# October 2018
##########################################################

import sys
import configparser
import os
from utils import is_sequential_dict, model_init, optimizer_init, forward_model, progress
from data_io import load_counts
import numpy as np
import random
import torch
from distutils.util import strtobool
import time
import threading
import torch

from data_io import read_lab_fea, open_or_fd, write_mat
from utils import shift

def read_next_chunk_into_shared_list_with_subprocess(
    read_lab_fea, shared_list, cfg_file, is_production, output_folder, wait_for_process
):
    p = threading.Thread(target=read_lab_fea, args=(cfg_file, is_production, shared_list, output_folder))
    p.start()
    if wait_for_process:
        p.join()
        return None
    else:
        return p

def extract_data_from_shared_list(shared_list):
    data_name = shared_list[0]
    data_end_index_fea = shared_list[1]
    data_end_index_lab = shared_list[2]
    fea_dict = shared_list[3]
    lab_dict = shared_list[4]
    arch_dict = shared_list[5]
    data_set = shared_list[6]
    return data_name, data_end_index_fea, data_end_index_lab, fea_dict, lab_dict, arch_dict, data_set

def convert_numpy_to_torch(data_set_dict, save_gpumem, use_cuda):
    if not (save_gpumem) and use_cuda:
        data_set_inp = torch.from_numpy(data_set_dict["input"]).float().cuda()
        data_set_ref = torch.from_numpy(data_set_dict["ref"]).float().cuda()
    else:
        data_set_inp = torch.from_numpy(data_set_dict["input"]).float()
        data_set_ref = torch.from_numpy(data_set_dict["ref"]).float()
    data_set_ref = data_set_ref.view((data_set_ref.shape[0], 1))
    return data_set_inp, data_set_ref

def run_nn_refac01(
    data_name, data_set, data_end_index, fea_dict, lab_dict, arch_dict, cfg_file, processed_first, next_config_file
):
    def _read_chunk_specific_config(cfg_file):
        if not (os.path.exists(cfg_file)):
            sys.stderr.write("ERROR: The config file %s does not exist!\n" % (cfg_file))
            sys.exit(0)
        else:
            config = configparser.ConfigParser()
            config.read(cfg_file)
        return config

    def _get_batch_size_from_config(config, to_do):
        if to_do == "train":
            batch_size = int(config["batches"]["batch_size_train"])
        elif to_do == "valid":
            batch_size = int(config["batches"]["batch_size_valid"])
        elif to_do == "forward":
            batch_size = 1
        return batch_size

    def _initialize_random_seed(config):
        seed = int(config["exp"]["seed"])
        torch.manual_seed(seed)
        random.seed(seed)
        np.random.seed(seed)

    def _load_model_and_optimizer(fea_dict, model, config, arch_dict, use_cuda, multi_gpu, to_do):
        inp_out_dict = fea_dict
        nns, costs = model_init(inp_out_dict, model, config, arch_dict, use_cuda, multi_gpu, to_do)
        optimizers = optimizer_init(nns, config, arch_dict)
        for net in nns.keys():
            pt_file_arch = config[arch_dict[net][0]]["arch_pretrain_file"]
            if pt_file_arch != "none":
                if use_cuda:
                    checkpoint_load = torch.load(pt_file_arch)
                else:
                    checkpoint_load = torch.load(pt_file_arch, map_location="cpu")
                nns[net].load_state_dict(checkpoint_load["model_par"])
                if net in optimizers:
                    optimizers[net].load_state_dict(checkpoint_load["optimizer_par"])
                    optimizers[net].param_groups[0]["lr"] = float(
                        config[arch_dict[net][0]]["arch_lr"]
                    )  # loading lr of the cfg file for pt
            if multi_gpu:
                nns[net] = torch.nn.DataParallel(nns[net])
        return nns, costs, optimizers, inp_out_dict

    def _open_forward_output_files_and_get_file_handles(forward_outs, require_decodings, info_file, output_folder):
        post_file = {}
        for out_id in range(len(forward_outs)):
            if require_decodings[out_id]:
                out_file = info_file.replace(".info", "_" + forward_outs[out_id] + "_to_decode.ark")
            else:
                out_file = info_file.replace(".info", "_" + forward_outs[out_id] + ".ark")
            post_file[forward_outs[out_id]] = open_or_fd(out_file, output_folder, "wb")
        return post_file

    def _get_batch_config(data_set_input, seq_model, to_do, data_name, batch_size):
        N_snt = None
        N_ex_tr = None
        N_batches = None
        if seq_model or to_do == "forward":
            N_snt = len(data_name)
            N_batches = int(N_snt / batch_size)
        else:
            N_ex_tr = data_set_input.shape[0]
            N_batches = int(N_ex_tr / batch_size)
        return N_snt, N_ex_tr, N_batches

    def _prepare_input(
        snt_index,
        batch_size,
        inp_dim,
        ref_dim,
        beg_snt_fea,
        beg_snt_lab,
        data_end_index_fea,
        data_end_index_lab,
        beg_batch,
        end_batch,
        seq_model,
        arr_snt_len_fea,
        arr_snt_len_lab,
        data_set_inp,
        data_set_ref,
        use_cuda,
    ):
        def _zero_padding(
            inp,
            ref,
            max_len_fea,
            max_len_lab,
            data_end_index_fea,
            data_end_index_lab,
            data_set_inp,
            data_set_ref,
            beg_snt_fea,
            beg_snt_lab,
            snt_index,
            k,
        ):
            def _input_and_ref_have_same_time_dimension(N_zeros_fea, N_zeros_lab):
                if N_zeros_fea == N_zeros_lab:
                    return True
                return False

            snt_len_fea = data_end_index_fea[snt_index] - beg_snt_fea
            snt_len_lab = data_end_index_lab[snt_index] - beg_snt_lab
            N_zeros_fea = max_len_fea - snt_len_fea
            N_zeros_lab = max_len_lab - snt_len_lab
            if _input_and_ref_have_same_time_dimension(N_zeros_fea, N_zeros_lab):
                N_zeros_fea_left = random.randint(0, N_zeros_fea)
                N_zeros_lab_left = N_zeros_fea_left
            else:
                N_zeros_fea_left = 0
                N_zeros_lab_left = 0
            inp[N_zeros_fea_left : N_zeros_fea_left + snt_len_fea, k, :] = data_set_inp[
                beg_snt_fea : beg_snt_fea + snt_len_fea, :
            ]
            ref[N_zeros_lab_left : N_zeros_lab_left + snt_len_lab, k, :] = data_set_ref[
                beg_snt_lab : beg_snt_lab + snt_len_lab, :
            ]
            return inp, ref, snt_len_fea, snt_len_lab

        if len(data_set_ref.shape) == 1:
            data_set_ref = data_set_ref.shape.view((data_set_ref.shape[0], 1))
        max_len = 0
        if seq_model:
            max_len_fea = int(max(arr_snt_len_fea[snt_index : snt_index + batch_size]))
            max_len_lab = int(max(arr_snt_len_lab[snt_index : snt_index + batch_size]))
            inp = torch.zeros(max_len_fea, batch_size, inp_dim).contiguous()
            ref = torch.zeros(max_len_lab, batch_size, ref_dim).contiguous()
            for k in range(batch_size):
                inp, ref, snt_len_fea, snt_len_lab = _zero_padding(
                    inp,
                    ref,
                    max_len_fea,
                    max_len_lab,
                    data_end_index_fea,
                    data_end_index_lab,
                    data_set_inp,
                    data_set_ref,
                    beg_snt_fea,
                    beg_snt_lab,
                    snt_index,
                    k,
                )
                beg_snt_fea = data_end_index_fea[snt_index]
                beg_snt_lab = data_end_index_lab[snt_index]
                snt_index = snt_index + 1
        else:
            if to_do != "forward":
                inp = data_set[beg_batch:end_batch, :].contiguous()
            else:
                snt_len_fea = data_end_index_fea[snt_index] - beg_snt_fea
                snt_len_lab = data_end_index_lab[snt_index] - beg_snt_lab
                inp = data_set_inp[beg_snt_fea : beg_snt_fea + snt_len_fea, :].contiguous()
                ref = data_set_ref[beg_snt_lab : beg_snt_lab + snt_len_lab, :].contiguous()
                beg_snt_fea = data_end_index_fea[snt_index]
                beg_snt_lab = data_end_index_lab[snt_index]
                snt_index = snt_index + 1
        if use_cuda:
            inp = inp.cuda()
            ref = ref.cuda()
        return inp, ref, max_len_fea, max_len_lab, snt_len_fea, snt_len_lab, beg_snt_fea, beg_snt_lab, snt_index

    def _optimization_step(optimizers, outs_dict, config, arch_dict):
        for opt in optimizers.keys():
            optimizers[opt].zero_grad()
        outs_dict["loss_final"].backward()
        for opt in optimizers.keys():
            if not (strtobool(config[arch_dict[opt][0]]["arch_freeze"])):
                optimizers[opt].step()

    def _update_progress_bar(to_do, i, N_batches, loss_sum):
        if to_do == "train":
            status_string = (
                "Training | (Batch "
                + str(i + 1)
                + "/"
                + str(N_batches)
                + ")"
                + " | L:"
                + str(round(loss_sum.cpu().item() / (i + 1), 3))
            )
            if i == N_batches - 1:
                status_string = "Training | (Batch " + str(i + 1) + "/" + str(N_batches) + ")"
        if to_do == "valid":
            status_string = "Validating | (Batch " + str(i + 1) + "/" + str(N_batches) + ")"
        if to_do == "forward":
            status_string = "Forwarding | (Batch " + str(i + 1) + "/" + str(N_batches) + ")"
        progress(i, N_batches, status=status_string)

    def _write_info_file(info_file, to_do, loss_tot, err_tot, elapsed_time_chunk):
        with open(info_file, "w") as text_file:
            text_file.write("[results]\n")
            if to_do != "forward":
                text_file.write("loss=%s\n" % loss_tot.cpu().numpy())
                text_file.write("err=%s\n" % err_tot.cpu().numpy())
            text_file.write("elapsed_time_chunk=%f\n" % elapsed_time_chunk)
        text_file.close()

    def _save_model(to_do, nns, multi_gpu, optimizers, info_file, arch_dict):
        if to_do == "train":
            for net in nns.keys():
                checkpoint = {}
                if multi_gpu:
                    checkpoint["model_par"] = nns[net].module.state_dict()
                else:
                    checkpoint["model_par"] = nns[net].state_dict()
                if net in optimizers:
                    checkpoint["optimizer_par"] = optimizers[net].state_dict()
                else:
                    checkpoint["optimizer_par"] = dict()
                out_file = info_file.replace(".info", "_" + arch_dict[net][0] + ".pkl")
                torch.save(checkpoint, out_file)

    def _get_dim_from_data_set(data_set_inp, data_set_ref):
        inp_dim = data_set_inp.shape[1]
        ref_dim = 1
        if len(data_set_ref.shape) > 1:
            ref_dim = data_set_ref.shape[1]
        return inp_dim, ref_dim

    from data_io import read_lab_fea_refac01 as read_lab_fea
    from utils import forward_model_refac01 as forward_model

    config = _read_chunk_specific_config(cfg_file)
    _initialize_random_seed(config)

    output_folder = config["exp"]["out_folder"]
    use_cuda = strtobool(config["exp"]["use_cuda"])
    multi_gpu = strtobool(config["exp"]["multi_gpu"])
    to_do = config["exp"]["to_do"]
    info_file = config["exp"]["out_info"]
    model = config["model"]["model"].split("\n")
    forward_outs = config["forward"]["forward_out"].split(",")
    forward_normalize_post = list(map(strtobool, config["forward"]["normalize_posteriors"].split(",")))
    forward_count_files = config["forward"]["normalize_with_counts_from"].split(",")
    require_decodings = list(map(strtobool, config["forward"]["require_decoding"].split(",")))
    save_gpumem = strtobool(config["exp"]["save_gpumem"])
    is_production = strtobool(config["exp"]["production"])
    batch_size = _get_batch_size_from_config(config, to_do)

    if processed_first:
        shared_list = list()
        p = read_next_chunk_into_shared_list_with_subprocess(
            read_lab_fea, shared_list, cfg_file, is_production, output_folder, wait_for_process=True
        )
        data_name, data_end_index_fea, data_end_index_lab, fea_dict, lab_dict, arch_dict, data_set_dict = extract_data_from_shared_list(
            shared_list
        )
        data_set_inp, data_set_ref = convert_numpy_to_torch(data_set_dict, save_gpumem, use_cuda)
    else:
        data_set_inp = data_set["input"]
        data_set_ref = data_set["ref"]
        data_end_index_fea = data_end_index["fea"]
        data_end_index_lab = data_end_index["lab"]
    shared_list = list()
    data_loading_process = None
    if not next_config_file is None:
        data_loading_process = read_next_chunk_into_shared_list_with_subprocess(
            read_lab_fea, shared_list, next_config_file, is_production, output_folder, wait_for_process=False
        )
    nns, costs, optimizers, inp_out_dict = _load_model_and_optimizer(
        fea_dict, model, config, arch_dict, use_cuda, multi_gpu, to_do
    )
    if to_do == "forward":
        post_file = _open_forward_output_files_and_get_file_handles(
            forward_outs, require_decodings, info_file, output_folder
        )

    seq_model = is_sequential_dict(config, arch_dict)
    N_snt, N_ex_tr, N_batches = _get_batch_config(data_set_inp, seq_model, to_do, data_name, batch_size)
    beg_batch = 0
    end_batch = batch_size
    snt_index = 0
    beg_snt_fea = 0
    beg_snt_lab = 0
    arr_snt_len_fea = shift(shift(data_end_index_fea, -1, 0) - data_end_index_fea, 1, 0)
    arr_snt_len_lab = shift(shift(data_end_index_lab, -1, 0) - data_end_index_lab, 1, 0)
    arr_snt_len_fea[0] = data_end_index_fea[0]
    arr_snt_len_lab[0] = data_end_index_lab[0]
    data_set_inp_dim, data_set_ref_dim = _get_dim_from_data_set(data_set_inp, data_set_ref)
    inp_dim = data_set_inp_dim + data_set_ref_dim
    loss_sum = 0
    err_sum = 0
    start_time = time.time()
    for i in range(N_batches):
        inp, ref, max_len_fea, max_len_lab, snt_len_fea, snt_len_lab, beg_snt_fea, beg_snt_lab, snt_index = _prepare_input(
            snt_index,
            batch_size,
            data_set_inp_dim,
            data_set_ref_dim,
            beg_snt_fea,
            beg_snt_lab,
            data_end_index_fea,
            data_end_index_lab,
            beg_batch,
            end_batch,
            seq_model,
            arr_snt_len_fea,
            arr_snt_len_lab,
            data_set_inp,
            data_set_ref,
            use_cuda,
        )
        if to_do == "train":
            outs_dict = forward_model(
                fea_dict,
                lab_dict,
                arch_dict,
                model,
                nns,
                costs,
                inp,
                ref,
                inp_out_dict,
                max_len_fea,
                max_len_lab,
                batch_size,
                to_do,
                forward_outs,
            )
            _optimization_step(optimizers, outs_dict, config, arch_dict)
        else:
            with torch.no_grad():
                outs_dict = forward_model(
                    fea_dict,
                    lab_dict,
                    arch_dict,
                    model,
                    nns,
                    costs,
                    inp,
                    ref,
                    inp_out_dict,
                    max_len_fea,
                    max_len_lab,
                    batch_size,
                    to_do,
                    forward_outs,
                )
        if to_do == "forward":
            for out_id in range(len(forward_outs)):
                out_save = outs_dict[forward_outs[out_id]].data.cpu().numpy()
                if forward_normalize_post[out_id]:
                    counts = load_counts(forward_count_files[out_id])
                    out_save = out_save - np.log(counts / np.sum(counts))
                write_mat(output_folder, post_file[forward_outs[out_id]], out_save, data_name[i])
        else:
            loss_sum = loss_sum + outs_dict["loss_final"].detach()
            err_sum = err_sum + outs_dict["err_final"].detach()
        beg_batch = end_batch
        end_batch = beg_batch + batch_size
        _update_progress_bar(to_do, i, N_batches, loss_sum)
    elapsed_time_chunk = time.time() - start_time
    loss_tot = loss_sum / N_batches
    err_tot = err_sum / N_batches
    del inp, ref, outs_dict, data_set_inp_dim, data_set_ref_dim
    _save_model(to_do, nns, multi_gpu, optimizers, info_file, arch_dict)
    if to_do == "forward":
        for out_name in forward_outs:
            post_file[out_name].close()
    _write_info_file(info_file, to_do, loss_tot, err_tot, elapsed_time_chunk)
    if not data_loading_process is None:
        data_loading_process.join()
        data_name, data_end_index_fea, data_end_index_lab, fea_dict, lab_dict, arch_dict, data_set_dict = extract_data_from_shared_list(
            shared_list
        )
        data_set_inp, data_set_ref = convert_numpy_to_torch(data_set_dict, save_gpumem, use_cuda)
        data_set = {"input": data_set_inp, "ref": data_set_ref}
        data_end_index = {"fea": data_end_index_fea, "lab": data_end_index_lab}
        return [data_name, data_set, data_end_index, fea_dict, lab_dict, arch_dict]
    else:
        return [None, None, None, None, None, None]

def run_nn(
    data_name, data_set, data_end_index, fea_dict, lab_dict, arch_dict, cfg_file, processed_first, next_config_file
):

    # This function processes the current chunk using the information in cfg_file. In parallel, the next chunk is load into the CPU memory

    # Reading chunk-specific cfg file (first argument-mandatory file)
    if not (os.path.exists(cfg_file)):
        sys.stderr.write("ERROR: The config file %s does not exist!\n" % (cfg_file))
        sys.exit(0)
    else:
        config = configparser.ConfigParser()
        config.read(cfg_file)

    # Setting torch seed
    seed = int(config["exp"]["seed"])
    torch.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)

    # Reading config parameters
    output_folder = config["exp"]["out_folder"]
    use_cuda = strtobool(config["exp"]["use_cuda"])
    multi_gpu = strtobool(config["exp"]["multi_gpu"])

    to_do = config["exp"]["to_do"]
    info_file = config["exp"]["out_info"]

    model = config["model"]["model"].split("\n")

    forward_outs = config["forward"]["forward_out"].split(",")
    forward_normalize_post = list(map(strtobool, config["forward"]["normalize_posteriors"].split(",")))
    forward_count_files = config["forward"]["normalize_with_counts_from"].split(",")
    require_decodings = list(map(strtobool, config["forward"]["require_decoding"].split(",")))

    use_cuda = strtobool(config["exp"]["use_cuda"])
    save_gpumem = strtobool(config["exp"]["save_gpumem"])
    is_production = strtobool(config["exp"]["production"])

    if to_do == "train":
        batch_size = int(config["batches"]["batch_size_train"])

    if to_do == "valid":
        batch_size = int(config["batches"]["batch_size_valid"])

    if to_do == "forward":
        batch_size = 1

    # ***** Reading the Data********
    if processed_first:

        # Reading all the features and labels for this chunk
        shared_list = []

        p = threading.Thread(target=read_lab_fea, args=(cfg_file, is_production, shared_list, output_folder))
        p.start()
        p.join()

        data_name = shared_list[0]
        data_end_index = shared_list[1]
        fea_dict = shared_list[2]
        lab_dict = shared_list[3]
        arch_dict = shared_list[4]
        data_set = shared_list[5]

        # converting numpy tensors into pytorch tensors and put them on GPUs if specified
        if not (save_gpumem) and use_cuda:
            data_set = torch.from_numpy(data_set).float().cuda()
        else:
            data_set = torch.from_numpy(data_set).float()

    # Reading all the features and labels for the next chunk
    shared_list = []
    p = threading.Thread(target=read_lab_fea, args=(next_config_file, is_production, shared_list, output_folder))
    p.start()

    # Reading model and initialize networks
    inp_out_dict = fea_dict

    [nns, costs] = model_init(inp_out_dict, model, config, arch_dict, use_cuda, multi_gpu, to_do)

    # optimizers initialization
    optimizers = optimizer_init(nns, config, arch_dict)

    # pre-training and multi-gpu init
    for net in nns.keys():
        pt_file_arch = config[arch_dict[net][0]]["arch_pretrain_file"]

        if pt_file_arch != "none":
            if use_cuda:
                checkpoint_load = torch.load(pt_file_arch)
            else:
                checkpoint_load = torch.load(pt_file_arch, map_location="cpu")
            nns[net].load_state_dict(checkpoint_load["model_par"])
            optimizers[net].load_state_dict(checkpoint_load["optimizer_par"])
            optimizers[net].param_groups[0]["lr"] = float(
                config[arch_dict[net][0]]["arch_lr"]
            )  # loading lr of the cfg file for pt

        if multi_gpu:
            nns[net] = torch.nn.DataParallel(nns[net])

#         print('--------- before ----------')
#         print(nns)
#
#         num_bits = 4
#         threshold = 1
#         nns[net].module.fsq_x.update_num_bits_and_threshold(num_bits, threshold)
#         for each in nns[net].module.fsq_hs:
#             each.update_num_bits_and_threshold(num_bits, threshold)
#
#         print('--------- after ----------')
#         print(nns)

    if to_do == "forward":

        post_file = {}
        for out_id in range(len(forward_outs)):
            if require_decodings[out_id]:
                out_file = info_file.replace(".info", "_" + forward_outs[out_id] + "_to_decode.ark")
            else:
                out_file = info_file.replace(".info", "_" + forward_outs[out_id] + ".ark")
            post_file[forward_outs[out_id]] = open_or_fd(out_file, output_folder, "wb")

    # check automatically if the model is sequential
    seq_model = is_sequential_dict(config, arch_dict)

    # ***** Minibatch Processing loop********
    if seq_model or to_do == "forward":
        N_snt = len(data_name)
        N_batches = int(N_snt / batch_size)
    else:
        N_ex_tr = data_set.shape[0]
        N_batches = int(N_ex_tr / batch_size)

    beg_batch = 0
    end_batch = batch_size

    snt_index = 0
    beg_snt = 0

    start_time = time.time()

    # array of sentence lengths
    arr_snt_len = shift(shift(data_end_index, -1, 0) - data_end_index, 1, 0)
    arr_snt_len[0] = data_end_index[0]

    loss_sum = 0
    err_sum = 0

    if to_do == "forward":
        num_zeroact_inps = 0
        num_act_inps = 0
        num_zeroact_rcrs = 0
        num_act_rcrs = 0
        num_ops = 0

    inp_dim = data_set.shape[1]
    for i in range(N_batches):

        max_len = 0

        if seq_model:

            max_len = int(max(arr_snt_len[snt_index : snt_index + batch_size]))
            inp = torch.zeros(max_len, batch_size, inp_dim).contiguous()

            for k in range(batch_size):

                snt_len = data_end_index[snt_index] - beg_snt
                N_zeros = max_len - snt_len

                # Appending a random number of initial zeros, tge others are at the end.
                N_zeros_left = random.randint(0, N_zeros)

                # randomizing could have a regularization effect
                inp[N_zeros_left : N_zeros_left + snt_len, k, :] = data_set[beg_snt : beg_snt + snt_len, :]

                beg_snt = data_end_index[snt_index]
                snt_index = snt_index + 1

        else:
            # features and labels for batch i
            if to_do != "forward":
                inp = data_set[beg_batch:end_batch, :].contiguous()
            else:
                snt_len = data_end_index[snt_index] - beg_snt
                inp = data_set[beg_snt : beg_snt + snt_len, :].contiguous()
                beg_snt = data_end_index[snt_index]
                snt_index = snt_index + 1

        # use cuda
        if use_cuda:
            inp = inp.cuda()

        if to_do == "train":
            # Forward input, with autograd graph active
            outs_dict = forward_model(
                fea_dict,
                lab_dict,
                arch_dict,
                model,
                nns,
                costs,
                inp,
                inp_out_dict,
                max_len,
                batch_size,
                to_do,
                forward_outs,
            )

            for opt in optimizers.keys():
                optimizers[opt].zero_grad()

            outs_dict["loss_final"].backward()

            # # Gradient Clipping (th 0.1)
            # for net in nns.keys():
            #    torch.nn.utils.clip_grad_norm_(nns[net].parameters(), 0.1)

            for opt in optimizers.keys():
                if not (strtobool(config[arch_dict[opt][0]]["arch_freeze"])):
                    optimizers[opt].step()

            # Clip decay rate to (0,1)
            for net in nns.keys():
                for name, param in nns[net].named_parameters():
                    if 'taus' in name:
                        w = param.data
                        w = w.clamp(0, 1)

        else:
            with torch.no_grad():  # Forward input without autograd graph (save memory)
                if to_do == "forward":
                    outs_dict, stat_zeroact_inps, stat_act_inps, stat_zeroact_rcrs, stat_act_rcrs, stat_ops = forward_model(
                        fea_dict,
                        lab_dict,
                        arch_dict,
                        model,
                        nns,
                        costs,
                        inp,
                        inp_out_dict,
                        max_len,
                        batch_size,
                        to_do,
                        forward_outs,
                    )
                    for key in stat_zeroact_inps.keys():
                        num_zeroact_inps += int(torch.sum(stat_zeroact_inps[key]))
                        num_act_inps += int(torch.sum(stat_act_inps[key]))
                        num_zeroact_rcrs += int(torch.sum(stat_zeroact_rcrs[key]))
                        num_act_rcrs += int(torch.sum(stat_act_rcrs[key]))
                        num_ops += int(torch.sum(stat_ops[key]))
                else:
                    outs_dict = forward_model(
                        fea_dict,
                        lab_dict,
                        arch_dict,
                        model,
                        nns,
                        costs,
                        inp,
                        inp_out_dict,
                        max_len,
                        batch_size,
                        to_do,
                        forward_outs,
                    )

        if to_do == "forward":
            for out_id in range(len(forward_outs)):

                out_save = outs_dict[forward_outs[out_id]].data.cpu().numpy()

                if forward_normalize_post[out_id]:
                    # read the config file
                    counts = load_counts(forward_count_files[out_id])
                    out_save = out_save - np.log(counts / np.sum(counts))

                # save the output
                write_mat(output_folder, post_file[forward_outs[out_id]], out_save, data_name[i])
        else:
            loss_sum = loss_sum + outs_dict["loss_final"].detach()
            err_sum = err_sum + outs_dict["err_final"].detach()

        # update it to the next batch
        beg_batch = end_batch
        end_batch = beg_batch + batch_size

        # Progress bar
        if to_do == "train":
            status_string = (
                "Training | (Batch "
                + str(i + 1)
                + "/"
                + str(N_batches)
                + ")"
                + " | L:"
                + str(round(loss_sum.cpu().item() / (i + 1), 3))
            )
            if i == N_batches - 1:
                status_string = "Training | (Batch " + str(i + 1) + "/" + str(N_batches) + ")"

        if to_do == "valid":
            status_string = "Validating | (Batch " + str(i + 1) + "/" + str(N_batches) + ")"
        if to_do == "forward":
            status_string = "Forwarding | (Batch " + str(i + 1) + "/" + str(N_batches) + ")"

        progress(i, N_batches, status=status_string)

    elapsed_time_chunk = time.time() - start_time

    loss_tot = loss_sum / N_batches
    err_tot = err_sum / N_batches

    # clearing memory
    del inp, outs_dict, data_set

    # save the model
    if to_do == "train":

        for net in nns.keys():
            checkpoint = {}
            if multi_gpu:
                checkpoint["model_par"] = nns[net].module.state_dict()
            else:
                checkpoint["model_par"] = nns[net].state_dict()

            checkpoint["optimizer_par"] = optimizers[net].state_dict()

            out_file = info_file.replace(".info", "_" + arch_dict[net][0] + ".pkl")
            torch.save(checkpoint, out_file)

    if to_do == "forward":
        for out_name in forward_outs:
            post_file[out_name].close()
        print("num_zeroact_inps={}".format(num_zeroact_inps))
        print("num_act_inps={}".format(num_act_inps))
        print("num_zeroact_rcrs={}".format(num_zeroact_rcrs))
        print("num_act_rcrs={}".format(num_act_rcrs))
        print("avg_ops={}".format(num_ops/N_batches))

    # Write info file
    with open(info_file, "w") as text_file:
        text_file.write("[results]\n")
        if to_do != "forward":
            text_file.write("loss=%s\n" % loss_tot.cpu().numpy())
            text_file.write("err=%s\n" % err_tot.cpu().numpy())
        else:
            text_file.write("num_zeroact_inps={}\n".format(num_zeroact_inps))
            text_file.write("num_act_inps={}\n".format(num_act_inps))
            text_file.write("num_zeroact_rcrs={}\n".format(num_zeroact_rcrs))
            text_file.write("num_act_rcrs={}\n".format(num_act_rcrs))
            text_file.write("avg_ops={}\n".format(num_ops/N_batches))
        text_file.write("elapsed_time_chunk=%f\n" % elapsed_time_chunk)

    text_file.close()

    # Getting the data for the next chunk (read in parallel)
    p.join()
    data_name = shared_list[0]
    data_end_index = shared_list[1]
    fea_dict = shared_list[2]
    lab_dict = shared_list[3]
    arch_dict = shared_list[4]
    data_set = shared_list[5]

    # converting numpy tensors into pytorch tensors and put them on GPUs if specified
    if not (save_gpumem) and use_cuda:
        data_set = torch.from_numpy(data_set).float().cuda()
    else:
        data_set = torch.from_numpy(data_set).float()

    return [data_name, data_set, data_end_index, fea_dict, lab_dict, arch_dict]