text:
'''
keras generator for videos. contains functions for loading and processing videos for
multiple architectures, currently mainly focused on i3d from DeepMind.
'''

__author__ = "Camilo Fosco"
__email__ = "cfosco@nvidia.com"

import numpy as np
import random
import os
# from skimage.color.adapt_rgb import adapt_rgb, each_channel, hsv_value
# from skimage import filters
# from skimage.filters import roberts, sobel, scharr, prewitt
import pandas as pd
from timeit import default_timer as timer
from vid_utils import load_video_opencv, load_video_opencv_fixed_frames, load_video_imageio, load_video_skvideo, plot_frames
from vid_utils import time_crop, video_spatial_crop, video_resize, video_rescale, video_truncate_and_rescale, video_resample
from PIL import Image
import matplotlib.pyplot as plt
import cv2
import keras
import sys
sys.path.append('../')

# from imgaug import augmenters as iaa
# from imgaug import parameters as iap
import i3d_config as cfg
from keras.preprocessing import image
from IPython.display import clear_output

def build_str2label(csv_path):
    lines = pd.read_csv(csv_path, header=None, sep=None).values
    label_dictionary = {}
    for line in lines:
        label_dictionary[line[0]] = line[1]

    return label_dictionary

def build_label2str(csv_path):
    lines = pd.read_csv(csv_path, header=None, sep=None).values
    label_dictionary = {}
    for line in lines:
        label_dictionary[line[1]] = line[0]

    return label_dictionary

def preprocess_spoof(vid_list, is_train=True):
    pass

def preprocess_frames_and_bboxes(batch, is_train=True):
    '''Preprocess function for ROI pooling. Preprocesses both input frames and bboxes accordingly.'''

    N= len(batch)
    channels = batch[0].shape[3]

    vid_array = []

    for i in range(N):
        # start = timer()
        ## Resize, smallest dim 256 pixels, bilinear interpolation
        resi_frames, scale_x, scale_y = video_resize_fixed(vid_list[i])

        bbox[i, :, 0] = int(bbox[i, :, 0]*scale_x)
        bbox[i, :, 1] = int(bbox[i, :, 1]*scale_y)
        bbox[i, :, 2] = int(bbox[i, :, 2]*scale_x)
        bbox[i, :, 3] = int(bbox[i, :, 3]*scale_y)

        # t1=timer()
        ## During training, randomly select a 224x224 crop. During test, take center 224x224 crop.
        # crop_frames = video_spatial_crop(resi_frames, size=cfg._IMAGE_SIZE, randomize=is_train)
        # t2=timer()
        ## Rescale between -1 and 1
        resc_frames = video_rescale(vid_list[i], maxi=1, mini=0)
        # t3=timer()
        ## Append frames to array
        vid_array.append( resc_frames )
        # t4=timer()
        # print('resizing: %.3f, sp_crop: %.3f, rescaling: %.3f, assigning to array: %.3f' % (t1-start, t2-t1,t3-t2,t4-t3))
    return vid_array

def preprocess_i3d_rgb(vid_list, is_train=True):
    ''' Preprocess a list of raw videos. The videos must be in numpy array format.'''

    N= len(vid_list)
    # print('len(vid_list)',len(vid_list))
    # print('vid_listshape',vid_list[0].shape)
    channels = vid_list[0].shape[3]
    vid_array = np.zeros((N, cfg._NUM_FRAMES, cfg._IMAGE_SIZE, cfg._IMAGE_SIZE, channels), dtype=np.float32)

    for i,vid in enumerate(vid_list):
        # Sample at 25 fps.
        # resa_frames = video_resample(vid, fps=25)

        # During training, randomly select a sequence of N frames. During test, take the middle N frames.
        resa_frames = time_crop(vid, N=cfg._NUM_FRAMES, is_train= is_train)

        # Resize, smallest dim 256 pixels, bilinear interpolation
        resi_frames = video_resize(resa_frames)

        # During training, randomly select a 224x224 crop. During test, take center 224x224 crop.
        crop_frames = video_spatial_crop(resi_frames, size=cfg._IMAGE_SIZE, randomize=is_train)

        # Rescale between -1 and 1
        resc_frames = video_rescale(crop_frames)

        # making sure that we're using float32
        vid_array[i] = resc_frames.astype(np.float32)

    return vid_array

def preprocess_rgb_fast_with_caption(X, is_train=True):
    return [preprocess_rgb_fast(X[0], is_train=is_train), X[1]]

def preprocess_flow_fast_with_caption(X, is_train=True):
    return [preprocess_flow_fast(X[0], is_train=is_train), X[1]]

def preprocess_rgb_fast(vid_list, is_train=True):
    '''Preprocess input frames faster than preprocess_i3d. for this to work,
    The vid_list has to contain videos that are either numpy arrays or lists of frames,
    frames being numpy arrays of the same width, height and channel. The frames must contain values from 0 to 255.
    Also, most importantly, the videos must already be time cropped,
    which means they must have exactly _NUM_FRAMES frames.

    If planning on feeding videos of more or less than cfg._NUM_FRAMES (for testing on untrimmed or unlooped videos, for
    example), use the preprocess_rgb_fast_var_input function '''

    N= len(vid_list)
    channels = vid_list[0].shape[3]

    vid_array = np.zeros((N, cfg._NUM_FRAMES, cfg._IMAGE_SIZE, cfg._IMAGE_SIZE, channels), dtype=np.float32)

    full_time=0
    for i in range(N):
#         start = timer()
        ## Resize, smallest dim 256 pixels, bilinear interpolation
#         resi_frames = video_resize(vid_list[i])

#         t1=timer()
        ## During training, randomly select a 224x224 crop. During test, take center 224x224 crop.
        crop_frames = video_spatial_crop(vid_list[i], size=cfg._IMAGE_SIZE, randomize=is_train)

        ##DEBUG
        # crop_frames/=255.

#         t2=timer()
        ## Rescale between -1 and 1 - THIS HAPPENS INSIDE THE NETWORK NOW
        # resc_frames = video_rescale(crop_frames)
#         t3=timer()

        ## Append frames to array
        vid_array[i] = crop_frames.astype(np.float32)
#         t4=timer()
#         print('resizing: %.4f, sp_crop: %.4f, rescaling: %.4f, assigning to array: %.4f' % (t1-start, t2-t1,t3-t2,t4-t3))
#         full_time+=t4-start

#     print('Full time on this batch of %d:' % N, full_time)
    return vid_array

def preprocess_rgb_fast_var_input(vid_list, is_train=True):
    '''Preprocess input frames faster than preprocess_i3d. for this to work,
    The vid_list has to contain videos that are either numpy arrays or lists of frames,
    frames being numpy arrays of the same width, height and channel. The frames must contain values from 0 to 255.

    This function doesn't assume a fixed number of input frames. Used for feeding
    batches where the videos have different number of frames (e.g. testing situation)
    '''

    N= len(vid_list)
    channels = vid_list[0].shape[3]

    vid_array = []

    for i in range(N):
        # Resize, smallest dim 256 pixels, bilinear interpolation
        resi_frames = video_resize(vid_list[i])

        # During training, randomly select a 224x224 crop. During test, take center 224x224 crop.
        crop_frames = video_spatial_crop(resi_frames, size=cfg._IMAGE_SIZE, randomize=is_train)

        # Rescale between -1 and 1
        resc_frames = video_rescale(crop_frames)

        vid_array.append(resc_frames.astype(np.float32))

    return vid_array

def preprocess_i3d_flow(vid_list, is_train=False):

    # print('len(vid_list)',len(vid_list))
    # print('vid_listshape',vid_list[0].shape)
    channels = vid_list[0].shape[3]
    vid_array = np.empty((len(vid_list), cfg._NUM_FRAMES, cfg._IMAGE_SIZE, cfg._IMAGE_SIZE, channels), dtype=np.float32)

    for i,vid in enumerate(vid_list):
        # Sample at 25 fps.
        # resa_frames = video_resample(vid, fps=25)

        # During training, randomly select a sequence of N frames. During test, take the middle N frames.
        resa_frames = time_crop(vid, N=cfg._NUM_FRAMES, is_train= is_train)

        # Resize, smallest dim 256 pixels, bilinear interpolation
        resi_frames = video_resize(resa_frames)

        # During training, randomly select a 224x224 crop. During test, take center 224x224 crop.
        crop_frames = video_spatial_crop(resi_frames, size=cfg._IMAGE_SIZE, randomize=is_train)

        # Rescale between -1 and 1
        resc_frames = video_truncate_and_rescale(crop_frames)

        # making sure that we're using float32
        vid_array[i] = resc_frames.astype(np.float32)

    return vid_array

def preprocess_flow_fast(vid_list, is_train=False):

    # print('len(vid_list)',len(vid_list))
    # print('vid_listshape',vid_list[0].shape)
    channels = vid_list[0].shape[3]
    vid_array = np.empty((len(vid_list), cfg._NUM_FRAMES, cfg._IMAGE_SIZE, cfg._IMAGE_SIZE, channels), dtype=np.float32)

    for i,vid in enumerate(vid_list):
        # Sample at 25 fps.
        # resa_frames = video_resample(vid, fps=25)

        # During training, randomly select a sequence of N frames. During test, take the middle N frames.
#         resa_frames = time_crop(vid, N=cfg._NUM_FRAMES, is_train= is_train)

        # Resize, smallest dim 256 pixels, bilinear interpolation - done offline
#         resi_frames = video_resize(vid)

        # During training, randomly select a 224x224 crop. During test, take center 224x224 crop.
        crop_frames = video_spatial_crop(vid, size=cfg._IMAGE_SIZE, randomize=is_train)

        # Rescale between -1 and 1 - done inside the net
#         resc_frames = video_truncate_and_rescale(crop_frames)

        # making sure that we're using float32
        vid_array[i] = crop_frames.astype(np.float32)

    return vid_array

def load_hmdb_npy(filenames, type='rgb', hmdb_npy_path = '../data/HMDB51', is_train=False, add_suffix=False):
    '''Npy file loader for the HMDB dataset. Returns rgb or flow npy files.'''

    # Change this if loading anything other than the preprocessed npy files found in ../data/HMDB51

    if type == 'rgb':
        channels = 3
        suffix = '_rgb'
    elif type =='flow':
        channels = 2
        suffix = '_flow'
    else:
        raise ValueError('Unrecognized type of image to load')

    if not add_suffix: suffix=''

    N = len(filenames)
    loaded_npy = np.zeros((N, cfg._NUM_FRAMES, 256, 256, channels))

    i=0
    for file in filenames:
        npy_name = file.split('.')[0]+suffix+'.npy'

#         print(os.path.join(hmdb_npy_path, npy_name))
        try:
            npy_arr = np.load(os.path.join(hmdb_npy_path, npy_name))
            loaded_npy[i] = time_crop(npy_arr,cfg._NUM_FRAMES)
        except Exception as err:
            print(err)
            print('Exception caught with video:', npy_name)
        i+=1

    return loaded_npy

def load_vids_from_path(filenames, path, load_func='opencv', is_train=False):
    '''Raw video loader for files in the string array filenames. each string in the list
    Must be a valid path of a .avi or .mpeg file'''

    if isinstance(load_func, str):
        if load_func == 'opencv':
            load_func = load_video_opencv_fixed_frames
        elif load_func == 'opencv_full':
            load_func = load_video_opencv
        elif load_func == 'imageio':
            load_func = load_video_imageio
        elif load_func == 'skvideo':
            load_func = load_video_skvideo
        else:
            raise ValueError('Unknown load function')

    videos = []
    for file in filenames:
        # print('loading file:',os.path.join(path, file))
        # start=timer()
        vid = np.array(load_func(os.path.join(path, file), is_train=is_train))
        # end=timer()
        # print(file,' - time loading:',end-start)
        if len(vid.shape) is not 4:
            raise ValueError('Video '+file+' has wrong shape: '+str(vid.shape))
        # print(file+': vid.shape in load_vids is', vid.shape)

        videos.append(vid)

    return videos

def load_vids_opencv(filenames, path, is_train=False):
    return load_vids_from_path(filenames, path, load_func='opencv', is_train=is_train)

def load_vids_skvideo(filenames, path, is_train=False):
    return load_vids_from_path(filenames, path, load_func='skvideo', is_train=is_train)

def load_vids_imageio(filenames, path, is_train=False):
    return load_vids_from_path(filenames, path, load_func='imageio', is_train=is_train)

def load_hmdb_npy_rgb(filenames, path, is_train=False):
    return load_hmdb_npy(filenames, type='rgb', hmdb_npy_path = path, is_train=is_train)

def load_hmdb_npy_flow(filenames, path, is_train=False):
    return load_hmdb_npy(filenames, type='flow', hmdb_npy_path = path, is_train=is_train)

def load_single_frame_as_video(folders_with_frames, path,
                is_train=False, frames_to_get=cfg._NUM_FRAMES,
                return_index = False):
    '''Loads a single frame from folder with frames and expands it to frames_to_get frames'''

    first_pass=True
    for i,f in enumerate(folders_with_frames):
        frames = np.load(os.path.join(path, f))
        # print('frames.shape',frames.shape)
        if first_pass:
            h = frames.shape[1]
            w = frames.shape[2]
            ch = frames.shape[3]
            batch = np.zeros((len(folders_with_frames), frames_to_get, h, w, ch))
            first_pass=False
        batch[i]= np.repeat(frames[0][np.newaxis, :,:,:], frames_to_get, axis=0)

    return batch

def load_frames(folders_with_frames, path,
                is_train=False, frames_to_get=cfg._NUM_FRAMES,
                return_fixed_frames=True, return_index=False,
                remove_suffix=False, ex_on_empty_folder=True, stride=2):
    '''Loads frames, considering that all frames in a given folder have the same width and height.'''

    batch=[]
    start_frames=[]

    if remove_suffix:
        if folders_with_frames[0].endswith('.mp4'):
            folders_with_frames = [f[:-4] for f in folders_with_frames]

    for j in range(0, len(folders_with_frames), stride):
        dirpath = os.path.join(path,folders_with_frames[j])
        if not os.path.isdir(dirpath):
            if ex_on_empty_folder:
                raise ValueError('Folder '+os.path.join(path,f)+' does not exist')
            else:
                batch.append(np.zeros((frames_to_get, cfg._IMAGE_SIZE, cfg._IMAGE_SIZE, 3)))
                continue

        first_pass=True
        frames = sorted([n for n in os.listdir(dirpath) if (n.endswith('.jpg') or n.endswith('.png') or n.endswith('.jpeg'))])
        if not frames:
            if ex_on_empty_folder:
                raise ValueError('Folder '+dirpath+' has zero frames in it')
            else:
                batch.append(np.zeros((frames_to_get, cfg._IMAGE_SIZE, cfg._IMAGE_SIZE, 3)))
                continue

        n_frames = len(frames)

        if not frames_to_get:
            frames_to_get = n_frames

        if n_frames <= frames_to_get*stride:

            for i in range(0, n_frames, stride):
                img = image.load_img(os.path.join(path, f, frames[i]))
                img = image.img_to_array(img)
                if first_pass:
                    rows,cols,channels = img.shape
                    imgs = np.zeros((n_frames, rows, cols, channels))
                    first_pass=False
                imgs[i]=img

            if is_train or return_fixed_frames:
                imgs = time_crop(imgs,frames_to_get)

            start_frames.append(0)
            batch.append(imgs)

        else:
            if is_train:
                start_f = np.random.randint(n_frames-frames_to_get)
            else:
                start_f = n_frames//2-(frames_to_get*stride)//2

            for i in range(0, frames_to_get*stride, stride):
                idx = start_f+i
                img = image.load_img(os.path.join(path, f, frames[idx]))
                img = image.img_to_array(img)
                if first_pass:
                    rows,cols,channels = img.shape
                    imgs = np.zeros((frames_to_get, rows, cols, channels))
                    first_pass=False
                imgs[i]=img

            start_frames.append(start_f)
            batch.append(imgs)
    if return_index:
        return batch, start_frames
    return batch

def load_frames_and_bbox(folders_with_frames, path,
                is_train=False, frames_to_get=cfg._NUM_FRAMES,
                return_fixed_frames=True):
    ''' Loads both frames and bboxes, and returns botheleents in a list.
    Made for the RoiPooling model, where a ground truth bbox is now used
    to know what slice of the feature map to take.

    Inputs:
    -------
    folders_with_frames: list of str
        List of all folders to be processed where frames are located.
        usually has size of a batch. Example: ['climb/vid1', 'run/vid2']
    path: str
        path to dataset
    is_train; bool
        Whether we are in training phase or not
    frames_to_get: int
        Number of frames to get from the video, or to reloop the vid if smaller.
    return_fixed_frames: bool
        Whether to always return frames-to-get frames or to be more lenient (smaller vids are temporally untouched)

    Outputs:
    --------
    batch: list of numpy arrays corresponding to the videos in folders_with_frames

    '''

    frame_batch, start_frames = load_frames(folders_with_frames, path,
                    is_train, frames_to_get,
                    return_fixed_frames, return_index=True)

    bbox_batch = load_bboxes(folders_with_frames,path,start_frames,
                            frames_to_get, return_fixed_frames)

    return [frame_batch, bbox_batch]

def load_bboxes(video_names,path,start_frames,
                frames_to_get=cfg._NUM_FRAMES,
                return_fixed_frames=True):

    for i,v in enumerate(video_names):
        bbox=np.load(os.path.join(path, v, 'bboxes.npy'))
        bbox[:,2] = bbox[:,2]-bbox[:,0]
        bbox[:,3] = bbox[:,3]-bbox[:,1]
        n_frames = bbox.shape[0]
        new_bbox = np.zeros(frames_to_get, 4)
        if bbox.shape[0] <= frames_to_get:
            if not return_fixed_frames:
                batch.append(bbox)
                continue
            for j in frames_to_get:
                new_bbox[j] = bbox[j%n_frames]
        else:
            new_bbox = bbox[start_frames[i]:start_frames[i]+frames_to_get]

        batch.append(new_bbox)

    return batch

def load_crops_raw(folders_with_frames, path, is_train=False):
    '''Loads frames, assuming that they're crops. Crops might not have the same width and height from one frame to the other.'''

    DESIRED_SIZE = cfg._IMAGE_SIZE
    CHANNELS = 3
    vids=[]

    for f in folders_with_frames:
        crops = sorted([n for n in os.listdir(os.path.join(path,f)) if (n.endswith('.jpg') or n.endswith('.png') or n.endswith('.jpeg'))])
        if not crops:
            raise ValueError('Folder '+os.path.join(path,f)+' has zero frames in it')
        frames=np.zeros((len(crops), DESIRED_SIZE, DESIRED_SIZE, CHANNELS), dtype=np.uint8)
        for i,c in enumerate(crops):
            frames[i]=resize_crop(os.path.join(path,f,c), DESIRED_SIZE)
        vids.append(frames)
    return vids

def load_crops_fixed(folders_with_frames, path,
                    is_train=False, frames_to_get=cfg._NUM_FRAMES,
                    return_fixed_frames=True):
    '''load crops, but always the amount indicated in frames-to-get, not more or less.

    Inputs:
    -------
    folders_with_frames: list of str
        List of all folders to be processed where frames are located.
        usually has size of a batch. Example: ['climb/vid1', 'run/vid2']
    path: str
        path to dataset
    is_train; bool
        Whether we are in training phase or not
    frames_to_get: int
        Number of frames to get from the video, or to reloop the vid if smaller.
    return_fixed_frames: bool
        Whether to always return frames-to-get frames or to be more lenient (smaller vids are temporally untouched)

    Returns:
    --------
    batch: list of numpy arrays
        Batch of sequnces composed of crops (which have not been necessarily resized).

    '''

    DESIRED_SIZE = cfg._IMAGE_SIZE
    CHANNELS = 3
    batch=[]

    for f in folders_with_frames:
        crops = sorted([n for n in os.listdir(os.path.join(path,f)) if (n.endswith('.jpg') or n.endswith('.png') or n.endswith('.jpeg'))])
        if not crops:
            raise ValueError('Folder '+os.path.join(path,f)+' has zero frames in it')
        n_frames = len(crops)

        if n_frames <= frames_to_get:
            frames=np.zeros((n_frames, DESIRED_SIZE, DESIRED_SIZE, CHANNELS), dtype=np.uint8)
            for i in range(n_frames):
                frames[i]=resize_crop(os.path.join(path,f,crops[i]), DESIRED_SIZE)

            if is_train or return_fixed_frames:
                frames = time_crop(frames,frames_to_get)

        else:
            if is_train:
                start_f = np.random.randint(n_frames-frames_to_get)
            else:
                start_f = n_frames//2-frames_to_get//2
            frames=np.zeros((frames_to_get, DESIRED_SIZE, DESIRED_SIZE, CHANNELS), dtype=np.uint8)
            for i in range(frames_to_get):
                idx = start_f+i
                frames[i]=resize_crop(os.path.join(path,f,crops[idx]), DESIRED_SIZE)

        batch.append(frames)

    return batch

def load_crops_in_streams(stream_folders, path, streams_to_concat=4, is_train=True):

    CROP_SIZE = cfg._IMAGE_SIZE
    FRAMES_PER_STREAM = 16
    CHANNELS = 3
    EXTENSION = 'extend10'

    vids = []
    for f in stream_folders:
        streams = sorted(os.listdir(os.path.join(path, f, EXTENSION)))
        # print('len streams:',len(streams))
        init = 0
        if len(streams)<=streams_to_concat:
            streams_to_concat=len(streams)
        elif is_train:
            init = np.random.randint(len(streams)-streams_to_concat)

        frames=np.zeros((FRAMES_PER_STREAM*streams_to_concat, CROP_SIZE, CROP_SIZE, CHANNELS))
        i=0
        for s in streams[init:init+streams_to_concat]:
            crops = sorted(os.listdir(os.path.join(path,f,EXTENSION,s)))
            if len(crops)>16:
                crops = crops[:16]
            for c in crops:
                frames[i]=resize_crop(os.path.join(path,f,EXTENSION,s,c), CROP_SIZE)
                i+=1
        vids.append(frames)

    return vids

def load_crops_raw_and_stream(folders, path, streams_to_concat=4, is_train=True):

    if not folders:
        raise ValueError('No folders to load received in load_crops_raw_and_stream function')

    raw_folders = []
    stream_folders = []
    for f in folders:
        # print('folder:',f)
        is_stream = False
        files = os.listdir(os.path.join(path, f))
        for file in files:
            if 'extend' in file:
                is_stream = True
                break
            elif '.jpg' in file or '.png' in file:
                break
        if is_stream:
            stream_folders.append(f)
            # print('appended to stream')
        else:
            raw_folders.append(f)
            # print('appended to raw')

    if len(stream_folders)+len(raw_folders)==0:
        raise ValueError('Folder length sums to zero. INFO DUMP\
        \nFolders: %s\nstream_folders: %s\nraw_folders: %s\n'
        % (str(folders), str(raw_folders), str(stream_folders)))

    # print('length stream, raw folder names:', len(stream_folders), len(raw_folders))
    vids_raw = load_crops_fixed(raw_folders, path=path, is_train=is_train)
    vids_stream = load_crops_in_streams(stream_folders, path=path, streams_to_concat=streams_to_concat, is_train=is_train)
    vids_concat = vids_raw+vids_stream
    # print('len(vids_raw, vids_stream)',len(vids_raw), len(vids_stream))

    if len(vids_raw)+ len(vids_stream)==0:
        raise ValueError('Loaded video sizes sums to zero. INFO DUMP\
        \nFolders: %s\nstream_folders: %s\nraw_folders: %s \nlen(vids_raw): %s\nlen(vids_stream): %s\
        \nvids_raw: '+str(vids_raw)+'\nvids_stream:'+str(vids_stream)
        % (str(folders), str(raw_folders), str(stream_folders), str(len(vids_raw)),str(len(vids_stream))))

    return vids_concat

def resize_crop(crop, desired_size=cfg._IMAGE_SIZE):

    # Open
    im = Image.open(crop)
    old_size = im.size  # old_size[0] is in (width, height) format

    # Generate image ratio
    ratio = float(desired_size)/max(old_size)
    new_size = tuple([int(x*ratio) for x in old_size])

    # Resize
    im = im.resize(new_size, Image.ANTIALIAS)

    # Create a new image and paste the resized on it
    new_im = Image.new("RGB", (desired_size, desired_size))
    new_im.paste(im, ((desired_size-new_size[0])//2,
                        (desired_size-new_size[1])//2))

    return new_im

# previously called load_labels
def load_labels_video(filenames, label_dictionary, dataset_path=None, label_pos_in_string=-2):

    labels = np.zeros((len(filenames),len(label_dictionary)), dtype=np.int16)
    i=0
    for file in filenames:
        classname = file.split('/')[label_pos_in_string]
        labels[i, label_dictionary[classname] ] = 1
        i+=1

    return labels

def load_labels_mem(filenames, str2label_dict, dataset_path=None, label_pos_in_string=None, remove_slash=True):

    labels = np.zeros(len(filenames))
    for i,file in enumerate(filenames):
        if '/' in file and remove_slash:
            file = file.split('/')[-1]
        # print('FILE IN LOAD LABELS MEM:',file)
        labels[i] = str2label_dict[file] if isinstance(str2label_dict[file], float) else str2label_dict[file][0]

    return labels

def load_labels_mem_alpha(filenames, str2label_dict, dataset_path=None, label_pos_in_string=None):

#     print('filenames in load labels mem alpha',filenames)
    labels = np.zeros((len(filenames), 2))
    for i,file in enumerate(filenames):
        # print('FILE IN LOAD LABELS MEM:',file)
        labels[i] = str2label_dict[file]

    return labels

def load_labels_as_list(filenames, str2label_dict, dataset_path=None, label_pos_in_string=None):

    labels = []
    for i, file in enumerate(filenames):
        labels.append( str2label_dict[file] )
    return np.array(labels)

def load_labels_frame(filenames, str2label_dict, dataset_path=None, label_pos_in_string=None):
    labels=[]
    for file in filenames:
        # The numpy file loaded here must be a matrix of frames x n_classes.
        # For each frame, there has to be a corresponding one-hot encoded vector indicating the correct class
        labels.append(np.load(os.path.join(dataset_path+'_aux', 'frame_labels', file+'.npy')))
    return labels

def augment_imgaug(X, augment):
    seq = get_imgaug_sequence(augment)
    return augment_video_batch(X, seq)

def augment_raw(X, augment):
    return augment_video_batch(X, augment)

#### THREADSAFE
import threading
class threadsafe_iter:
    def __init__(self, it):
        self.it = it
        self.lock = threading.Lock()
    def __iter__(self):
        return self
    def __next__(self):
        with self.lock:
            return next(self.it)

def threadsafe_generator(f):
    def g(*a, **kw):
        return threadsafe_iter(f(*a, **kw))
    return g

#######################
### GENERATOR CLASS ###
#######################

class VideoSeqGenerator(keras.utils.Sequence):

    '''Generates video data for Keras models. Extends keras Sequence class,
    making it work nicely with fit_generator and evaluate_generator, as well as
    multiprocessing.

    '''

    def __init__(self,
                txt_file=None,
                batch_size=1,
                files=None,
                label_csv="../../../nfs_share/datasets/HMDB51/videos_aux/labels.csv",
                str2label_dict=None,
                augment=None,
                load_labels_func = 'per_video',
                load_func='hmdb_npy_rgb',
                preprocess_func=None,
                augment_func=None,
                dataset_path=None,
                shuffle=True,
                is_train=True,
                verbose=1,
                use_imgaug=True,
                return_labels=True,
                remove_excess_files=True,
                has_label_col=False,
                it_per_epoch=None,
                sample_weights=None,
                input_captions=None,
                len_vocab=cfg.LEN_VOCAB):

        'Initialization'

        self.txt_file = txt_file
        self.batch_size = batch_size
        self.files = files
        self.augment = augment
        self.load_func = load_func
        self.preprocess_func = preprocess_func
        self.augment_func = augment_func
        self.load_labels_func = load_labels_func
        self.dataset_path = dataset_path
        self.label_csv = label_csv
        self.shuffle = shuffle
        self.is_train = is_train
        self.verbose = verbose
        self.use_imgaug = use_imgaug
        self.return_labels = return_labels
        self.times=[timer()]
        self.sample_weights=sample_weights
        self.input_captions=input_captions

        self.it_per_epoch = it_per_epoch
        self.epoch_counter = 0
        self.count = 0
        self.len_vocab = len_vocab

        load_func_dict = {  'vids_opencv': load_vids_opencv,
                            'vids_imageio': load_vids_imageio,
                            'vids_skvideo': load_vids_skvideo,
                            'npy_rgb': load_hmdb_npy_rgb,
                            'npy_flow': load_hmdb_npy_flow,
                            'crops_raw': load_crops_raw,
                            'crops_fixed': load_crops_fixed,
                            'crops_stream': load_crops_in_streams,
                            'crops_mixed': load_crops_raw_and_stream,
                            'frames_raw': load_frames,
                            'single_frame': load_single_frame_as_video}

        preprocess_func_dict = { 'i3d_rgb': preprocess_i3d_rgb,
                                 'i3d_flow': preprocess_i3d_flow,
                                 'fast_rgb': preprocess_rgb_fast,
                                 'fast_flow': preprocess_flow_fast,
                                 'cap_fast_rgb': preprocess_rgb_fast_with_caption,
                                 'cap_fast_flow': preprocess_flow_fast_with_caption,
                                  }

        augment_func_dict = {'augment_imgaug': augment_imgaug,
                            'augment_raw': augment_raw}

        load_labels_dict = {'per_video': load_labels_video,
                            'per_frame': load_labels_frame,
                            'mem': load_labels_mem,
                            'mem_alpha': load_labels_mem_alpha,
                            'as_list': load_labels_as_list}

        ## Reading files to load from txt file
        # start=timer()
        if has_label_col:
            if self.files is None:
                self.files, self.label_array = read_videofile_txt(self.txt_file, ret_label_array=has_label_col)
            else:
                self.label_array = [f[1] for f in self.files]
                self.files = [f[0] for f in self.files]
        else:
            if self.files is None:
                self.files = read_videofile_txt(self.txt_file)
            self.label_array = None

        self.full_file_list = self.files
        self.num_files = len(self.full_file_list)

        ## If limiting number of iterations per epoch
        if self.it_per_epoch and self.it_per_epoch*self.batch_size > self.num_files:
            print('Caution, it_per_epoch is bigger than the amount of iterations a normal epoch takes. \
                    Disregarding parameter - epochs will have %d iterations, as normal.' % self.num_files//self.batch_size)
            self.it_per_epoch = None

        ## Checking if batch size is multiple of num files
        excess_files = self.num_files % self.batch_size
        if excess_files != 0:
            print('Caution: batch_size (%d) is not multiple of given dataset length (%d). This can cause issues with multigpu training.' % (self.batch_size, self.num_files) )
            if remove_excess_files:
                print('Truncating given files: removing last %d elements' % excess_files)
                self.files = self.files[:-excess_files]
                self.num_files = len(self.files)
                print('New num_files:', self.num_files)

        ## Defining load function to use
        if isinstance(self.load_func, str):
            self.load_func = load_func_dict[self.load_func]

        ## Defining preprocess function to use
        if isinstance(self.preprocess_func, str):
            self.preprocess_func = preprocess_func_dict[self.preprocess_func]

        ## Defining augment function to use
        if isinstance(self.augment_func, str):
            self.augment_func = augment_func_dict[self.augment_func]

        ## Defining label function to use
        if isinstance(self.load_labels_func, str):
            self.load_labels_func = load_labels_dict[self.load_labels_func]

        ## Shuffling files and setting list of files to use if it_per_epoch is not None
        self.on_epoch_end()

        ## Build dictionary to switch from string label to numbered label
        if self.return_labels:
            if str2label_dict is None:
                print('Building stringtolabel')
                self.str2label_dict = build_str2label(label_csv)
            else:
                self.str2label_dict = str2label_dict
            self._original_str2label_dict = self.str2label_dict.copy()

        if self.verbose:
            print('Calling VideoSeqGenerator. Batch size: ',self.batch_size,
            '. Number of files received:',self.num_files,'. Augmentation: ',self.augment)

    def __len__(self):
        'Denotes the number of batches per epoch'
        if self.it_per_epoch:
            return self.it_per_epoch
#         if not self.is_train:
#             return np.ceil(self.num_files / self.batch_size)
#         else:
        return self.num_files // self.batch_size

    def __getitem__(self, index):
        'Generate one batch of data'

        start = timer()

        ###DEBUG
#         print('\n\nIndex: '+str(index)+' - Files/folders to load in this call:',self.files[index*self.batch_size:(index+1)*self.batch_size])

        files_to_use = self.files[index*self.batch_size:(index+1)*self.batch_size]

        if len(files_to_use)==0:
            raise ValueError('Slice of files has no elements. INFO DUMP:\nself.batch_size: %d\
            \nindex: %d\nlen(self.files): %d\nself.files[index*self.batch_size]: %s' % (self.batch_size, index, len(self.files), '-'))

        kwargs = {}
        if self.input_captions is not None:
            kwargs['idx_cap'] = np.random.randint(len(self.input_captions[files_to_use[0]]))
            kwargs['idx_seq'] = {n:np.random.randint(len(self.input_captions[n][kwargs['idx_cap']])) for n in files_to_use}
            kwargs['input_captions'] = self.input_captions
            kwargs['len_vocab'] = self.len_vocab

        ## LOAD DATA
        X = self.load_func(files_to_use,
                           self.dataset_path,
                           is_train=self.is_train,
                           **kwargs)

        if self.return_labels:
            Y = self.load_labels_func(files_to_use,
                                        self.str2label_dict,
                                        self.label_array,
                                        **kwargs)

        if self.sample_weights is not None:
            W = self.sample_weights[index*self.batch_size:(index+1)*self.batch_size]

        t1=timer()

        ## PREPROCESS DATA
        if self.preprocess_func is not None:
            X = self.preprocess_func(X, is_train=self.is_train)

        t2=timer()

        ## AUGMENT DATA
        if self.augment_func is None:
            if self.augment is not None:
                if self.use_imgaug:
                    X = augment_imgaug(X, self.augment)
                else:
                    X = augment_raw(X, self.augment)

        else:
            X = self.augment_func(X, self.augment)

        # Always do random flipping.
        if self.is_train:
            if np.random.choice([False, True]):
                if type(X)==list and len(X)==2:
                    X[0] = np.flip(X[0], axis=3)
                else:
                    X = np.flip(X, axis=3)

        t3=timer()

        if not (self.count % cfg._FREQ_OF_TIME_PRINTS) and self.verbose:
            print('\nTotal time on this batch:', t3-start,
            ' - it:',self.count, '- idx:',index)
            print('Time loading:', t1-start)
            print('Time preprocessing:', t2-t1)
            print('Time augmenting:', t3-t2)

        ## DEBUG
#         plot_frames(X[0], is_01image=True, frames_to_show=4)
#         print('Y',Y)

        self.count+=1
        if self.return_labels:
            if self.sample_weights is not None:
                return X, Y, [W]+[np.ones(len(W))]*(len(Y)-1)
            else:
                return X, Y
        else:
            return X

    def on_epoch_end(self):
        'Updates indexes after each epoch'

        if self.it_per_epoch:
            if self.epoch_counter == 0:
                if self.shuffle:
                    random.shuffle(self.full_file_list)
            virt_idx = (self.epoch_counter+1)*self.it_per_epoch*self.batch_size
            if virt_idx >= self.num_files:
                self.files = self.full_file_list[(self.epoch_counter*self.it_per_epoch*self.batch_size)%self.num_files:]
                self.files.extend(self.full_file_list[:virt_idx%self.num_files])
                self.epoch_counter = 0
            else:
                self.files = self.full_file_list[self.epoch_counter*self.it_per_epoch*self.batch_size:virt_idx]
                self.epoch_counter+=1
        else:
            if self.shuffle:
                if self.sample_weights is None:
                    random.shuffle(self.files)
                else:
                    idxs = list(range(len(self.files)))
                    np.random.shuffle(idxs)
                    self.files = self.files[idxs]
                    self.sample_weights = self.sample_weights[idxs]

        # print(f'Files to see in this epoch: {len(self.files)} \
        # (it_per_epoch = {self.it_per_epoch}, batch_size = {self.batch_size}, epoch_counter = {self.epoch_counter})')

class VideoCapGenerator(VideoSeqGenerator):
    pass

#####################
### AUGMENTATION ####
#####################

def get_imgaug_sequence(augment):
    # TODO: Modify raw imgaug code to play nicely with videos, mainly take batches
    # of vids and process them all at once in matrix manner. Might improve speed.

    seq = iaa.Sequential()

    AUG_NUM = 11
    act_array = np.random.rand(AUG_NUM) >0.5
    p = np.random.rand(AUG_NUM)

    for a in augment:
        if a == 'ch_shift' and act_array[0]:
#             print(a, p[0])
            seq.add(iaa.Multiply((0.5+p[0], 0.65+p[0]), per_channel=1))
        elif a == 'br_shift' and act_array[1]:
#             print(a, p[1])
            seq.add(iaa.Multiply((0.8+p[1]*.6, 0.9+p[1]*.6)))
        elif a == 'edge' and act_array[2]:
#             print(a, p[2])
            seq.add(iaa.Sharpen(alpha=p[2]/2+.1, lightness=p[2]/2+.6))
        elif a == 'dropout' and act_array[3]:
#             print(a, p[3])
            seq.add(iaa.CoarseDropout(np.random.rand()*.15, size_percent=(p[3]*.2+.005)))
        elif a == 'crop' and act_array[4]:
#             print(a, p[4])
            seq.add(iaa.Crop(percent=(0.05, .15))) # crop images from each side by 0 to 16px (randomly chosen)
        elif a == 'flip' and act_array[5]:
#             print(a, p[5])
            seq.add(iaa.Fliplr(1.0))
        elif a == 'stretch' and act_array[6] :
#             print(a, p[6])
            seq.add(iaa.Affine(scale={'x':(1.2,1.5),'y':(1.2,1.5)}))
        elif a == 'crop_pad' and act_array[7]:
#             print(a, p[7])
            seq.add(iaa.CropAndPad(
            # This is: top, right, bot, left. Negative = crop, pos = pad
            percent=((0, -0.1), (p[7]*0.3+0.1, p[7]*0.3+0.05) , (0, -0.1), (p[7]*0.3+0.1, p[7]*0.3+0.05))))
        elif a == 'shear' and act_array[8]:
#             print(a, p[8])
            seq.add(iaa.Affine(shear=(p[8]*-10, p[8]*10)))
        elif a == 'cutout' and act_array[9]:
#             print(a, p[9])
            seq.add(iaa.Lambda(cutout_img, heatmap_func, keypoint_func))
        elif a == 'simplex' and act_array[10]:
#             print(a, p[10])
            seq.add(iaa.SimplexNoiseAlpha(
                first=iaa.Multiply(0), sigmoid_thresh=iap.Normal(8.0, 3.0),
                per_channel=False))

    return seq

def augment_video_batch(batch, augment):
    if isinstance(augment, iaa.Augmenter):
        for i in range(len(batch)):
            if np.random.choice([True,False]):
                batch[i] = augment.augment_images(batch[i])
    else:
        for i in range(len(batch)):
            if 'crop' in augment:
                batch[i] = spatial_crop(batch[i])
            if 'flip' in augment:
                batch[i] = random_flip(batch[i])
            if 'edge' in augment:
                batch[i] = edge_enhancement(batch[i])
            if 'stretch' in augment:
                batch[i] = stretch_video(batch[i])
            if 'ch_shift' in augment:
                batch[i] = channel_shift(batch[i])
            if 'br_shift' in augment:
                batch[i] = brightness_shift(batch[i])
            if 'cutout' in augment:
                batch[i] = cutout_video(batch[i])
            if 'rotate' in augment:
                batch[i] = rotate_video(batch[i])
            if 'shear' in augment:
                batch[i] = shear(batch[i])

    return batch

def stretch_video():
    pass

def heatmap_func(hm_on_images, random_state, parents, hooks):
    return hm_on_images

def keypoint_func(keypoints_on_images, random_state, parents, hooks):
    return keypoints_on_images

def cutout_img(images, random_state, parents, hooks):
    fill_value = 0 #if np.random.rand()<0.5 else -1

    for img in images:
        H,W,C = img.shape
        cutout_size = np.random.randint(H*0.18)+int(H*0.1)
        x = np.random.randint(W-cutout_size+1)
        y = np.random.randint(H-cutout_size+1)
        img[y:y+cutout_size,x:x+cutout_size,:] = fill_value
    return images

def cutout_video(vid, cutout_percent_height=0.35, prob = 0.5):

    r = np.random.rand()
    if r<prob:
        # print('Cutout triggered')
        F,H,W,C = vid.shape
        cutout_size = int(cutout_percent_height*H)
        if cutout_size >= W:
            cutout_size = W//2

        fill_value = 0 #if np.random.rand()<0.5 else -1

        # In 50% of the cases, generate uniform cutout accross frames.
        if r<prob/2:
            x = np.random.randint(W-cutout_size+1)
            y = np.random.randint(H-cutout_size+1)

            vid[:,y:y+cutout_size,x:x+cutout_size,:] = fill_value

        # In the other 50%, randomize the cutout for every frame
        else:
            for i in range(F):
                x = np.random.randint(W-cutout_size+1)
                y = np.random.randint(H-cutout_size+1)

                vid[i,y:y+cutout_size,x:x+cutout_size,:] = fill_value

    return vid

def channel_shift(vid, filter):
    raise NotImplementedError()

def brightness_shift():
    raise NotImplementedError()

def random_flip(vid, prob=0.5):
    if np.random.rand()<prob:
        # print('Flip triggered')
        F,H,W,C = vid.shape
        for i in range(F):
            for j in range(C):
                vid[i,:,:,j] = np.fliplr(vid[i,:,:,j])

    return vid

def edge_enhancement(vid, prob=0.5, method = 'sobel'):

    if np.random.rand()<prob:
        # print('Edge enhancement triggered')
        if method =='sobel':
            @adapt_rgb(each_channel)
            def sobel_each(image):
                if isinstance(image[0][0], np.uint8):
                    return image//2+sobel(image)
                return np.clip((image*0.5+sobel(image)+0.5)/2, -1, 1)
            edge_filt = sobel_each
        elif method == 'roberts':
            edge_filt = roberts
        elif method == 'scharr':
            edge_filt = scharr
        elif method == 'prewitt':
            edge_filt = prewitt
        else:
            raise ValueError('Unknown edge enhancemenet method: '+method)

        for i in range(len(vid)):
            vid[i] = edge_filt(vid[i])

    return vid

def pca_enhancement():
    raise NotImplementedError()

def build_hmdb_textfiles(hmdb_splitfiles_dir,  output_dir='', split_n=1):
    """ Builds textfiles tailored to HMDB dataset required for the keras generator video_generator.

        Args
        ----
        hmdb_splitfiles_dir: string
            Directory where the splitfiles are located (downloaded from hmdb website)
        output_dir: string
            directory to save the txt files in
        split_n: int
            split number to generate files for (has to be 1,2 or 3)

        Raises
        -----
        ValueError if split_n is not 1 2 or 3

    """
    if split_n >=4 or split_n <=0:
        raise ValueError('HMDB51 has only 3 splits - split_n must be 1, 2 or 3.')

    txt_names = os.listdir(hmdb_splitfiles_dir)
    train_vids=[]
    test_vids=[]
    for name in txt_names:
        if 'split'+str(split_n) in name:
            print('looking in', name)
            classname = name.split('_test')[0]
            with open(os.path.join(hmdb_splitfiles_dir,name),'r') as f:
                lines=f.readlines()
            for l in lines:
                if l[-3]=='1':
                    train_vids.append(classname+'/'+l[:-4])
                elif l[-3]=='2':
                    test_vids.append(classname+'/'+l[:-4])

            print('len(train_vids)',len(train_vids))
            print('len(test_vids)',len(test_vids))

    # Saving txts
    with open(os.path.join(output_dir, 'train_split0'+str(split_n)+'.txt'), 'w') as f:
        for v in train_vids:
            f.write("%s\n" % v)

    # Saving test txts
    with open(os.path.join(output_dir, 'test_split0'+str(split_n)+'.txt'), 'w') as f:
        for v in test_vids:
            f.write("%s\n" % v)

def read_videofile_txt(txt_path, ret_label_array=False):
    files=[]
    if ret_label_array:
        label_array=[]
    with open(txt_path, 'r') as f:
        for line in f:
            # print(line)
            filename = line[:-1].split(',')[0].split(' ')[0]
            files.append(filename)
            if ret_label_array:
                label_array.append(line.split(',')[1])
    if ret_label_array:
        return files, label_array
    else:
        return files

# THE ABOVE FUNCTION COULD BE REPLACED BY SIMPLY [l[:-1] for l in open(txt_path, 'r')]
def read_txt_oneliner(txt_path):
    return [l[:-1] for l in open(txt_path, 'r')]

def show_N_results(model, txt_file, N=5, show_top_5=True, idxs=[0,16,32,48,-1], l2s_dict={}, augment = None,
                    dataset_path=None, csv_path=None, prep_func=None, load_func='hmdb_npy_rgb',
                    func_gen = True, shuffle=False, use_imgaug=True):
    if func_gen:
        vid_gen = video_generator(txt_file, dataset_path=dataset_path,
                                batch_size=1, augment = augment, shuffle=shuffle, is_train=False,
                                label_csv=csv_path, load_func=load_func, preprocess_func=prep_func,
                                use_imgaug=use_imgaug)
    else:
        vid_gen = VideoSeqGenerator(txt_file, dataset_path=dataset_path,
                                batch_size=1, augment = augment, shuffle=shuffle, is_train=False,
                                label_csv=csv_path, load_func=load_func, preprocess_func=prep_func,
                                use_imgaug=use_imgaug)

    if N is None:
        N = len([l for l in open(txt_file)])
    y_true=[]
    y_pred=[]
    is_in_top_5=0
    for i in range(N):
        if func_gen:
            batch = next(vid_gen)
        else:
            batch = vid_gen[i]
        y_true.append(np.argmax(batch[1],axis=1)[0])
        prediction = model.predict(batch[0])
        y_pred.append(np.argmax(prediction, axis=1)[0])
        top_5_preds = np.argsort(prediction[0])[-5:]
        if y_true[i] in top_5_preds:
            is_in_top_5+=1

        if l2s_dict:
            str_pred = l2s_dict[y_pred[i]]
            str_true = l2s_dict[y_true[i]]
        else:
            str_pred=str_true='-'
        caption = 'Predicted: '+str_pred+' / '+str(y_pred[i])+' - True: '+str_true+' / '+str(y_true[i])

        # print('np.min(batch[0][0])',np.min(batch[0][0]))
        # print('np.max(batch[0][0])',np.max(batch[0][0]))

        plot_frames(batch[0][0], title=caption, idxs =idxs)
        if show_top_5:
            print('Top 5 predictions:')
            for idx in top_5_preds[::-1]:
                print(l2s_dict[idx],' with prob %.3f' % prediction[0][idx])
            print()

    print('Accuracy in results shown:')
    print('TOP 1:', np.mean(np.equal(y_true,y_pred)))
    print('TOP 5:', is_in_top_5/N )

### VIDEO GENERATOR AS FUNCTION - DEPRECATED
# @threadsafe_generator
# def video_generator(txt_file=None,
#                     batch_size=1,
#                     files=None,
#                     label_csv="../../../datasets/HMDB51/videos_aux/labels.csv",
#                     augment=None,
#                     load_func='hmdb_npy_rgb',
#                     preprocess_func=None,
#                     augment_func=None,
#                     label_func='per_video',
#                     dataset_path=None,
#                     shuffle=True,
#                     is_train=True,
#                     verbose=1, use_imgaug=False,
#                     return_labels=True):

#     '''
#     DEPRECATED! Use videoSeqGenerator instead

#     Keras Generator for videos. Example flow:
#             # 0. Turn to numpy array
#             # 1. Resize
#             # 2. Rescale
#             # 3. Resample fps (optional)
#             # 4. Random time crop
#             # 5. Random space crop
#             # 6. Augmentation

#     The first lines define the function dictionaries for generator generalization.
#     If using a new dataset or architecture, add the corresponding load_func for the dataset
#     and the preprocess_func for the architecture to the dictionaries.

#     Args:
#     -----
#     txt_file: string
#         file containing the names of the vids to load, in format "classname/videoname.avi".
#         The list will be shuffled or not, depending on the flag.
#     batch_size: int
#         number of videos to load per batch.
#     label_csv: string
#         The path to the csv file containing tuples (class_name, numeric_label) for the current dataset
#     type: string
#         One of flow or rgb. selects which type of data should be loaded.
#         This assumes that flow data augmentationdata is available.
#     shuffle: boolean
#         wether to shuffle the incoming data list or not.
#     augment: list of strings
#         Array of augmentation strings. Currently supported: 'cutout', 'flip', 'edge'

#     Yields:
#     -------
#     (X,Y): tuple with batch of videos and corresponding labels.
#     '''

#     load_func_dict = {  'vids_opencv':load_vids_imageio,
#                         'vids_imageio': load_vids_imageio,
#                         'vids_skvideo': load_vids_skvideo,
#                         'npy_rgb': load_hmdb_npy_rgb,
#                         'npy_flow': load_hmdb_npy_flow,
#                         'crops_raw': load_crops_raw,
#                         'crops_fixed': load_crops_fixed,
#                         'crops_stream': load_crops_in_streams,
#                         'crops_mixed': load_crops_raw_and_stream,
#                         'frames_raw': load_frames }

#     preprocess_func_dict = { 'i3d_rgb': preprocess_i3d_rgb,
#                              'i3d_flow': preprocess_i3d_flow,
#                              'fast_rgb': preprocess_fast}

#     augment_func_dict = {'augment_imgaug': augment_imgaug,
#                         'augment_raw': augment_raw}

#     load_labels_dict = {'per_video': load_labels_video,
#                         'per_frame': load_labels_frame}

#     ## Reading files to load from txt file
#     # start=timer()
#     if files == None:
#         files=read_videofile_txt(txt_file)

#     ## Defining load function to use
#     if isinstance(load_func, str):
#         load_func = load_func_dict[load_func]

#     ## Defining preprocess function to use
#     if isinstance(preprocess_func, str):
#         preprocess_func = preprocess_func_dict[preprocess_func]

#     ## Defining augment function to use
#     if isinstance(augment_func, str):
#         augment_func = augment_func_dict[augment_func]

#     ## Defining label function to use
#     if isinstance(load_labels_func, str):
#         load_labels_func = load_labels_dict[load_labels_func]

#     ## Shuffling files
#     L = len(files)
#     if shuffle:
#         random.shuffle(files)

#     ## This is called once at the beggining of the training, so the shuffle happens once, at that point.
#     if verbose:
#         print('Calling video_generator. Batch size: ',batch_size,'. Number of files received:',L,'. Augmentation: ',augment)

#     ## Build dictionary to switch from string label to numbered label
#     if return_labels:
#         str2label_dict = build_str2label(label_csv)

#     ## This line is just to make the generator infinite, keras needs that
#     while True:

#         ## Define starting idx for batch slicing
#         batch_start = 0
#         batch_end = batch_size
#         c=1

#         ## Loop over the txt file while there are enough unseen files for one more batch
#         while batch_start < L:
#             ## LOAD DATA
#             limit = min(batch_end, L)
#             # print('STEP',c,' - yielding', limit-batch_start,'videos.')
#             start = timer()
#             ## load_func can return either a numpy array of shape (batch, frames, h, w, c)
#             ## or a list of videos with different n_frames, widths and heights. preprocess_func has to work accordingly.
#             if dataset_path is not None:
#                 X = load_func(files[batch_start:limit], dataset_path, is_train=is_train)
#             else:
#                 X = load_func(files[batch_start:limit], is_train=is_train)
#             if return_labels:
#                 Y = load_labels_func(files[batch_start:limit], str2label_dict, dataset_path)

#             ## AUGMENT DATA
#             if augment_func is None:
#                 if augment is not None:
#                     if use_imgaug:
#                         X = augment_imgaug(X, augment)
#                     else:
#                         X = augment_raw(X, augment)
#             else:
#                 X = augment_func(X, augment)

#             ## PREPROCESS DATA
#             if preprocess_func is not None:
#                 X = preprocess_func(X, is_train=is_train)

#             t1 = timer()
#             if not c % cfg._FREQ_OF_TIME_PRINTS:
#                 print('Time elapsed loading, augmenting and preprocessing:', t1-start, ' - it:',c)

#             ## YIELD DATA
#             if return_labels:
#                 yield X,Y #a tuple with two numpy arrays with batch_size samples
#             else:
#                 yield X
#             ## Increasing idxs for next batch
#             batch_start += batch_size
#             batch_end += batch_size
#             c+=1


text:
import numpy as np
import torch
from torch.optim import Adam
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import gym
import time
from tqdm import tqdm
import os, sys
from json import JSONEncoder
import json
import spinup.algos.pytorch.vpg.core as core
from spinup.algos.pytorch.ppo.broil_ppo import *
from spinup.utils.logx import EpochLogger
from spinup.envs.pointbot_const import *
from spinup.utils.mpi_pytorch import setup_pytorch_for_mpi, sync_params, mpi_avg_grads
from spinup.utils.mpi_tools import mpi_fork, mpi_avg, proc_id, mpi_statistics_scalar, num_procs
from spinup.examples.pytorch.broil_rtg_pg_v2.pointbot_reward_utils import PointBotReward
from spinup.examples.pytorch.broil_rtg_pg_v2.cartpole_reward_utils import CartPoleReward
# from spinup.examples.pytorch.broil_rtg_pg_v2.cheetah_reward_utils import CheetahReward
from spinup.examples.pytorch.broil_rtg_pg_v2.reacher_reward_utils import ReacherReward
from spinup.examples.pytorch.broil_rtg_pg_v2.shelf_reward_utils import ShelfReward
from spinup.examples.pytorch.broil_rtg_pg_v2.cvar_utils import cvar_enumerate_pg
import dmc2gym

def visualize_policy(env, num_rollouts, ac, num_rew_fns, std_r, mean_r, reward_dist, local_steps_per_epoch, broil_alpha, file_data, args, max_ep_len, t, broil_lambda, training=False):

    start_time = time.time()
    o, ep_ret, ep_len = env.reset(), 0, 0

    cvar_list = []
    wc_ret_list = []
    bc_ret_list = []
    ret_list = []
    obstacle_list = []
    trajectories_x = []
    trajectories_y = []
    trash_trajectories = []
    num_trashes = []
    obs_times = []
    count = 0
    # Main loop: collect experience in env and update/log each epoch
    for i in range(num_rollouts):
        print("Rollout: ", i)
        total_reward_dist = np.zeros(num_rew_fns)
        running_ret = 0
        num_runs = 1
        obstacles = 0
        num_constraint_violations = 0
        num_episodes = 0
        done = False
        count += 1
        constraint_violated = False
        j = 0
        while not done:
            a, v, logp = ac.step(torch.as_tensor(o, dtype=torch.float32))

            # TODO: Test unnormalizing the values
            v = (v * std_r.numpy()) + mean_r.numpy()
            j += 1
            next_o, r, d, info = env.step(a)

            rew_dist = reward_dist.get_reward_distribution(env, next_o)
            total_reward_dist += rew_dist.flatten()
            running_ret += r
            ep_ret += r
            ep_len += 1
            if args.env == 'PointBot-v0':
                obstacles += int(env.obstacle(next_o))
            # Update obs (critical!)
            o = next_o

            timeout = ep_len == max_ep_len
            terminal = d or timeout
            epoch_ended = t==local_steps_per_epoch-1

            done = done or j == env._max_episode_steps

        # Store stuff for saving data

        posterior_reward_weights = reward_dist.posterior

        #calculate sigma and find either the conditional value at risk or entropic risk measure given the current policy
            #Calculate policy gradient for conditional value at risk
        sigma, cvar = cvar_enumerate_pg(total_reward_dist, posterior_reward_weights, broil_alpha)

        ret_list.append(running_ret / float(num_runs))
        wc_ret_list.append(np.min(total_reward_dist) / float(num_runs))
        bc_ret_list.append(np.max(total_reward_dist / float(num_runs)))
        cvar_list.append(cvar)
        obstacle_list.append(obstacles / float(num_runs))
        obs_times.append(env.feature[0])
        if TRASH:
            num_trashes.append(env.feature[2])
        print('True returns:', ret_list)
        print('Cvar: ', cvar_list)
        print('Worst case:', wc_ret_list)

        if args.env == 'PointBot-v0':
            plt.ylim((env.grid[2], env.grid[3]))
            plt.xlim((env.grid[0], env.grid[1]))
            plt.scatter([env.hist[0][0]],[env.hist[0][2]],  [8], '#00FF00', zorder=11)
            x, y = [env.hist[i][0] for i in range(1, len(env.hist))], [env.hist[i][2] for i in range(1, len(env.hist))]
            plt.scatter(x, y, [6]*len(x), zorder =9)
            if TRASH:
                for j in env.current_trash_taken:
                    plt.scatter([j[0]],[j[1]], [25], zorder = 10, color = '#000000')
                plt.scatter([env.next_trash[0]],[env.next_trash[1]], [25], zorder = 10, color = '#000000')
            x_bounds = [obstacle.boundsx for obstacle in env.obstacle.obs]
            y_bounds = [obstacle.boundsy for obstacle in env.obstacle.obs]
            for i in range(len(x_bounds)):
                plt.gca().add_patch(patches.Rectangle((x_bounds[i][0], y_bounds[i][0]), width=x_bounds[i][1] - x_bounds[i][0], height=y_bounds[i][1] - y_bounds[i][0], fill=True, alpha=.5, linewidth=1, zorder = 0, edgecolor='#d3d3d3',facecolor='#d3d3d3'))

            if not os.path.exists(file_data + 'Rollouts'):
                os.makedirs(file_data + 'Rollouts')
            plt.savefig(file_data + 'Rollouts'+ '/rollout_' + str(count) + '.png')
            plt.clf()

        o, ep_ret, ep_len = env.reset(), 0, 0

    experiment_name = args.env + '_alpha_' + str(broil_alpha) + '_lambda_' + str(broil_lambda) + '_rollouts'

    metrics = {"conditional value at risk": ('_cvar', cvar_list),
            "true_return": ('_true_return', ret_list),
            "worst case return": ('_worst_case_return', wc_ret_list),
            "best case return": ('_best_case_return', bc_ret_list),
            "obstacle_collision": ('_obstacles', obstacle_list)}

    for metric, result in metrics.items():
        file_metric_description, results = result
        file_path = file_data + 'Rollouts/' + experiment_name + file_metric_description + '.txt'
        #assert not os.path.isfile(file_path)  # make sure we are making a new file and not overwriting
        with open(file_path, 'w') as f:
            for item in results:
                f.write("%s\n" % item)

    f = open(file_data + '/rolloutAverages.txt', "w")
    f.write("Trash: " + str(num_trashes))
    f.write("\nAverage Trashes: "+ str(np.average(num_trashes[-100:])))
    f.write("\nAverage Trashes Stnd Dev: "+ str(np.std(num_trashes[-100:])))
    f.write("\nObs times: " + str(obs_times))
    f.write("\nAverage Obstacle Time: "+ str(np.average(obs_times[-100:])))
    f.write("\nAverage Obstacle Time Stnd Dev: "+ str(np.std(obs_times[-100:])))
    f.close()

ext: py

max_issues_repo_name: zaynahjaved/pg-broil

max_forks_repo_name: zaynahjaved/pg-broil

reason: import numpy


text:
################################################## KE LIANG ############################################################
import torch
from torch import nn
from torchvision import transforms
from torchvision.utils import make_grid
import numpy as np
from PIL import Image
from config import _C as cfg
from model import get_model

model = get_model(model_name=cfg.MODEL.ARCH)
resume_path = 'checkpoint/Modified_Residual_Transfer_Model_apparent_2.pth'##modify the path of the model here
# resume_path = 'checkpoint/Modified_Residual_Transfer_Model_real_2.pth'
checkpoint = torch.load(resume_path, map_location="cpu")
model.load_state_dict(checkpoint['state_dict'])
image = Image.open('visual_layers/000096.jpg') ##modify the path of the model here

transform = transforms.Compose([transforms.Resize((224, 224)),
                         transforms.ToTensor(),
                         transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                              std=[0.229, 0.224, 0.225])])
img = transform(image)
img = img.unsqueeze(0)

def save_img(tensor, name):
    tensor = tensor.permute((1, 0, 2, 3))
    im = make_grid(tensor, normalize=True, scale_each=True, nrow=8, padding=2).permute((1, 2, 0))
    im = (im.data.numpy() * 255.).astype(np.uint8)
    Image.fromarray(im).save(name + '.jpg')

new_model = nn.Sequential(*list(model.children())[:0])
f1 = new_model(img)
# save_img(f1, 'visual_layers/Real/layer1')
save_img(f1, 'visual_layers/Apparent/layer1')##modify the path of output

new_model = nn.Sequential(*list(model.children())[:1])
f2 = new_model(img)
# save_img(f2, 'visual_layers/Real/layer2')
save_img(f2, 'visual_layers/Apparent/layer2')##modify the path of output

new_model = nn.Sequential(*list(model.children())[:2])
f3 = new_model(img)
# save_img(f3, 'visual_layers/Real/layer3')
save_img(f3, 'visual_layers/Apparent/layer3')##modify the path of output

new_model = nn.Sequential(*list(model.children())[:3])
f4 = new_model(img)
# save_img(f4, 'visual_layers/Real/layer4')
save_img(f4, 'visual_layers/Apparent/layer4')##modify the path of output

new_model = nn.Sequential(*list(model.children())[:4])
f5 = new_model(img)
# save_img(f5, 'visual_layers/Real/layer5')
save_img(f5, 'visual_layers/Apparent/layer5')##modify the path of output

new_model = nn.Sequential(*list(model.children())[:5])
f6 = new_model(img)
# save_img(f6, 'visual_layers/Real/layer6')
save_img(f6, 'visual_layers/Apparent/layer6')##modify the path of output

new_model = nn.Sequential(*list(model.children())[:6])
f7 = new_model(img)
print(f7.shape)
# save_img(f7, 'visual_layers/Real/layer7')
save_img(f7, 'visual_layers/Apparent/layer7')##modify the path of output

new_model = nn.Sequential(*list(model.children())[:7])
f8 = new_model(img)
print(f8.shape)
# save_img(f8, 'visual_layers/Real/layer8')
save_img(f7, 'visual_layers/Apparent/layer8')##modify the path of output

new_model = nn.Sequential(*list(model.children())[:8])
f9 = new_model(img)
print(f9.shape)
# save_img(f9, 'visual_layers/Real/layer9')
save_img(f9, 'visual_layers/Apparent/layer9')##modify the path of output
########################################################################################################################


text:
# !/usr/bin/env python
# encoding: utf-8
"""
Session.py

Created by Tomas HJ Knapen on 2009-11-26.
Copyright (c) 2009 TK. All rights reserved.
"""

# import python packages:
from __future__ import division
from sklearn.linear_model import Ridge
from IPython import embed as shell
import numpy as np
import pylab as pl
from scipy.stats import pearsonr,spearmanr,kendalltau
from scipy.signal import fftconvolve
from scipy import ndimage
from scipy import interpolate
from lmfit import minimize, Minimizer, Parameters, Parameter, report_fit, report_errors
from skimage.morphology import disk
from matplotlib import cm
import seaborn as sn
from sympy.solvers import solve
from sympy import Symbol, exp
import hrf_estimation as he
import time as t
import os
sn.set(style="ticks")

class gpf(object):
	def __init__(self, design_matrix, max_eccentricity, n_pixel_elements, ssr, rtime,slice_no):
		self.design_matrix = design_matrix
		self.max_eccentricity = max_eccentricity
		self.n_pixel_elements = n_pixel_elements
		self.ssr = ssr
		self.rtime = rtime
		self.slice = slice_no

		X = np.linspace(-max_eccentricity, max_eccentricity, n_pixel_elements)
		Y = np.linspace(-max_eccentricity, max_eccentricity, n_pixel_elements)
		self.MG = np.meshgrid(X, Y)

	#define model function and pass independent variables x and y as a list
	def twoD_Gaussian(self,  xo, yo, sigma):
		(x,y) = self.MG
		theta=0
		a = (np.cos(theta)**2)/(2*sigma**2) + (np.sin(theta)**2)/(2*sigma**2)
		b = -(np.sin(2*theta))/(4*sigma**2) + (np.sin(2*theta))/(4*sigma**2)
		c = (np.sin(theta)**2)/(2*sigma**2) + (np.cos(theta)**2)/(2*sigma**2)
		gauss = np.exp( - (a*((x-xo)**2) + 2*b*(x-xo)*(y-yo) + c*((y-yo)**2)))
		gauss[disk((self.n_pixel_elements-1)/2)==0] = 0
		return gauss

	def raw_model_prediction(self,  xo, yo, sigma):
		g = self.twoD_Gaussian(xo, yo, sigma).reshape(self.n_pixel_elements**2)
		return np.dot(self.design_matrix.reshape(-1,self.n_pixel_elements**2), g)

	def hrf_model_prediction(self, xo, yo, sigma,hrf_params):

		rmp = self.raw_model_prediction( xo, yo, sigma)
		original_size = len(rmp)
		rmp = np.repeat(rmp, self.ssr, axis=0)
		xx = np.arange(0,32,self.rtime/float(self.ssr))
		self.hrf_kernel = hrf_params[0] * he.hrf.spmt(xx) +hrf_params[1]* he.hrf.dspmt(xx) +hrf_params[2] * he.hrf.ddspmt(xx)
		if self.hrf_kernel.shape[0] % 2 == 1:
			self.hrf_kernel = np.r_[self.hrf_kernel, 0]
		self.hrf_kernel /= np.abs(self.hrf_kernel).sum()

		# add slice timing correction to fit
		convolved_mp = fftconvolve( rmp, self.hrf_kernel, 'full' )[int(self.slice)::int(self.ssr)][:int(original_size)]

		return convolved_mp, self.hrf_kernel

def fitRidge_for_Dumoulin(design_matrix, timeseries, n_iter = 50, compute_score = False, verbose = True,valid_regressors=[],n_pixel_elements=[], alpha = 1.0):
	"""fitRidge fits a design matrix to a given timeseries.
	It computes the coefficients and returns these coefficients
	plus the correlation between the model fit and timeseries.
	"""
	br = Ridge(alpha = alpha)
	br.fit(design_matrix, timeseries)
	predicted_signal = br.coef_ * design_matrix
	srp = list(spearmanr(timeseries, predicted_signal.sum(axis = 1)))
	srp = [srp[0], -np.log10(srp[1])]

	PRF = np.zeros(n_pixel_elements**2)
	PRF[valid_regressors] = br.coef_
	PRF = np.reshape(PRF,(n_pixel_elements,n_pixel_elements))
	maximum = ndimage.measurements.maximum_position(PRF)

	start_params = {}
	start_params['xo'], start_params['yo'] = maximum[1]/float(n_pixel_elements)*2-1, maximum[0]/float(n_pixel_elements)*2-1

	return start_params, PRF, predicted_signal.sum(axis = 1)

def fit_PRF_on_concatenated_data(data_shared,voxels_in_this_slice,n_TRs,n_slices,fit_on_all_data,plotbool,raw_design_matrices, dm_for_BR,
	valid_regressors, n_pixel_elements_convolved, n_pixel_elements_raw,plotdir,voxno,slice_no,randint,roi,TR,model,hrf_params_shared,all_results_shared,conditions,
	results_frames,	postFix=[],max_eccentricity=1,max_xy = 5,orientations=['0','45','90','135','180','225','270','315','X'],stim_radius = 7.5,
	nuisance_regressors = []):
	"""
	"""

	# grab data for this fit procedure from shared memory
	time_course = np.array(data_shared[:,voxels_in_this_slice][:,voxno])
	hrf_params = np.array(hrf_params_shared[:,voxels_in_this_slice][:,voxno])

	n_orientations = len(orientations)

	# already initialize the final PRF dict
	PRFs = {}

	if fit_on_all_data:

		#########################################################################################################################################################################################################################
		#### Instantiate parameters
		#########################################################################################################################################################################################################################

		## initiate search space with Ridge prefit
		Ridge_start_params, PRFs['Ridge'], BR_predicted = fitRidge_for_Dumoulin(dm_for_BR, time_course, valid_regressors=valid_regressors, n_pixel_elements=n_pixel_elements_convolved, alpha=1e14)

		## initiate parameters:
		params = Parameters()

		# one baseline parameter
		params.add('baseline',value=0.0)

		# two location parameters
		params.add('xo_%s'%conditions[0], value = Ridge_start_params['xo'])
		params.add('yo_%s'%conditions[0], value = Ridge_start_params['yo'])

		params.add('sigma_center_%s'%conditions[0],value=0.1,min=1e-20)#min=1e-201 # this means initialization at 0.1 * 7.5 = 0.75 degrees, with minimum of 0.075 degrees
		params.add('amp_center_%s'%conditions[0],value=0.05,min=1e-20)#min=1e-201 # this is initialized at 0.001

		# surround parameters
		params.add('sigma_surround_%s'%conditions[0],value=0.3,expr='sigma_center_%s+delta_sigma_%s'%(conditions[0],conditions[0])) # surround size should roughly be 5 times that of the center
		params.add('delta_sigma_%s'%conditions[0],value=0.4,min=1e-20) # this difference parameter ensures that the surround is always larger than the center#,min=1e-20000000001
		params.add('amp_surround_%s'%conditions[0],value=-0.005,max=1e-20,expr='-amp_center_%s+delta_amplitude_%s'%(conditions[0],conditions[0])) # initialized at 10% of center amplitude #max=-0.0000000001,
		params.add('delta_amplitude_%s'%conditions[0],value=0.045,min=1e-20) # this difference parameter ensures that the surround is never deeper than the center is high,min=0.0000000001

		# when fitting an OG model, set all surround and delta parameters to 0 and to not vary and set the expression to None, otherwise it will start to vary anyway
		if model == 'OG':
			params['amp_surround_%s'%conditions[0]].value,params['amp_surround_%s'%conditions[0]].vary,params['amp_surround_%s'%conditions[0]].expr = 0, False, None
			params['delta_sigma_%s'%conditions[0]].vary,params['sigma_surround_%s'%conditions[0]].vary =  False, False
			params['delta_amplitude_%s'%conditions[0]].vary = False
	else:

		#########################################################################################################################################################################################################################
		#### INITIATING PARAMETERS with all results
		#########################################################################################################################################################################################################################

		# grab data for this fit procedure from shared memory
		all_results = np.array(all_results_shared[:,voxels_in_this_slice][:,voxno])

		## initiate parameters:
		params = Parameters()

		# shared baseline param:
		params.add('baseline', value = all_results[results_frames['baseline']])

		# location parameters
		for condition in conditions:
			params.add('xo_%s'%condition, value = all_results[results_frames['xo']])
			params.add('yo_%s'%condition, value = all_results[results_frames['yo']])

			# center parameters:
			params.add('sigma_center_%s'%condition,value=all_results[results_frames['sigma_center']]/stim_radius,min=1e-20) # this means initialization at 0.05/2 * 15 = 1.5 degrees, ,min=0.0084
			params.add('amp_center_%s'%condition,value=all_results[results_frames['amp_center']],min=1e-20) # this is initialized at 0.001 ,min=0.0000000001

			# surround parameters
			params.add('sigma_surround_%s'%condition,value=all_results[results_frames['sigma_surround']]/stim_radius,expr='sigma_center_%s+delta_sigma_%s'%(condition,condition)) # surround size should roughly be 5 times that of the center
			params.add('amp_surround_%s'%condition,value=all_results[results_frames['amp_surround']],max=-1e-20,expr='-amp_center_%s+delta_amplitude_%s'%(condition,condition)) # initialized at 10% of center amplitudemax=-0.0000000001
			params.add('delta_sigma_%s'%condition,value=all_results[results_frames['delta_sigma']],min=1e-20) # this difference parameter ensures that the surround is always larger than the centermin=0.0000000001
			params.add('delta_amplitude_%s'%condition,value=all_results[results_frames['delta_amplitude']],min=1e-20) # this difference parameter ensures that the surround is never deeper than the center is highmin=0.0000000001

			# when fitting an OG model, set all surround and delta parameters to 0 and to not vary and set the expression to None, otherwise it will start to vary anyway
			if model == 'OG':
				params['amp_surround_%s'%condition].value,params['amp_surround_%s'%condition].vary,params['amp_surround_%s'%condition].expr = 0, False, None
				params['delta_sigma_%s'%condition].vary,params['sigma_surround_%s'%condition].vary = False, False
				params['delta_amplitude_%s'%condition].vary=False

		g = gpf(design_matrix = raw_design_matrices[conditions[0]], max_eccentricity = max_eccentricity, n_pixel_elements = n_pixel_elements_raw, rtime = TR, ssr = 1,slice_no=slice_no)

		# recreate PRFs
		this_surround_PRF = g.twoD_Gaussian(all_results[results_frames['xo']],all_results[results_frames['yo']],
			all_results[results_frames['sigma_surround']]/stim_radius) * all_results[results_frames['amp_surround']]
		this_center_PRF = g.twoD_Gaussian(all_results[results_frames['xo']], all_results[results_frames['yo']],
			all_results[results_frames['sigma_center']]/stim_radius) * all_results[results_frames['amp_center']]
		PRFs['All_fit'] = this_center_PRF + this_surround_PRF

	#########################################################################################################################################################################################################################
	#### Prepare fit object and function
	#########################################################################################################################################################################################################################

	# initiate model prediction object
	ssr = np.round(1/(TR/float(n_slices)))

	gpfs = {}
	for condition in conditions:
		gpfs[condition] = gpf(design_matrix = raw_design_matrices[condition], max_eccentricity = max_eccentricity, n_pixel_elements = n_pixel_elements_raw, rtime = TR, ssr = ssr,slice_no=slice_no)

	def residual(params,recreate=False):

		# combine all stimulus regressors
		combined_model_prediction = np.ones_like(time_course)*params['baseline'].value
		for ci,condition in enumerate(conditions):
			combined_model_prediction += gpfs[condition].hrf_model_prediction(params['xo_%s'%condition].value, params['yo_%s'%condition].value,
				params['sigma_center_%s'%condition].value,hrf_params)[0] * params['amp_center_%s'%condition].value
			combined_model_prediction += gpfs[condition].hrf_model_prediction(params['xo_%s'%condition].value, params['yo_%s'%condition].value,
				params['sigma_surround_%s'%condition].value, hrf_params)[0] * params['amp_surround_%s'%condition].value

		return time_course - combined_model_prediction

	#########################################################################################################################################################################################################################
	#### evalute fit
	#########################################################################################################################################################################################################################

	# optimize parameters
	minimize(residual, params, args=(), kws={},method='powell')

	#########################################################################################################################################################################################################################
	#### Recreate resulting predictions and PRFs with optimized parameters
	#########################################################################################################################################################################################################################

	# initiate model prediction at baseline value
	combined_model_prediction = np.ones_like(time_course) * params['baseline'].value

	# now loop over conditions, create prediction and add to total prediction
	model_predictions = {}
	for ci,condition in enumerate(conditions):
		this_center_model_prediction = gpfs[condition].hrf_model_prediction(params['xo_%s'%condition].value, params['yo_%s'%condition].value,
			params['sigma_center_%s'%condition].value,hrf_params)[0] * params['amp_center_%s'%condition].value
		this_surround_model_prediction = gpfs[condition].hrf_model_prediction(params['xo_%s'%condition].value, params['yo_%s'%condition].value,
			params['sigma_surround_%s'%condition].value, hrf_params)[0] * params['amp_surround_%s'%condition].value
		model_predictions[condition] = this_center_model_prediction + this_surround_model_prediction
		combined_model_prediction += model_predictions[condition]

		# recreate PRFs
		this_center_PRF = gpfs[condition].twoD_Gaussian(params['xo_%s'%condition].value, params['yo_%s'%condition].value,
			params['sigma_center_%s'%condition].value) * params['amp_center_%s'%condition].value
		this_surround_PRF = gpfs[condition].twoD_Gaussian(params['xo_%s'%condition].value, params['yo_%s'%condition].value,
			params['sigma_surround_%s'%condition].value) * params['amp_surround_%s'%condition].value
		PRFs[condition] = this_center_PRF + this_surround_PRF

	#########################################################################################################################################################################################################################
	#### Get fit diagnostics
	#########################################################################################################################################################################################################################

	# add fwhm, necessary when fitting DoG
	reconstruction_radius = 10
	this_ssr = 1000
	t = np.linspace(-reconstruction_radius,reconstruction_radius,this_ssr*reconstruction_radius)

	fwhms = {}
	surround_sizes = {}
	for condition in conditions:
		PRF_2D = params['amp_center_%s'%condition].value * np.exp(-t**2/(2*params['sigma_center_%s'%condition].value**2)) + params['amp_surround_%s'%condition].value * np.exp(-t**2/(2*(params['sigma_surround_%s'%condition].value)**2))
		## then, we fit a spline through this line, and get the roots (the fwhm points) of the spline:
		spline=interpolate.UnivariateSpline(range(len(PRF_2D)),PRF_2D-np.max(PRF_2D)/2,s=0)
		## and compute the distance between them
		try:
			fwhms[condition] = ((np.diff(spline.roots())/len(t)*reconstruction_radius) * stim_radius)[0]
		except:
			## when this procedure fails, set fwhm to 0:
			fwhms[condition] = 0

		## now find the surround size in the same way
		if (model == 'OG') + (params['amp_surround_%s'%condition].value == 0):
			surround_sizes[condition] = 0
		else:
			spline = interpolate.UnivariateSpline(range(len(PRF_2D)),PRF_2D+np.min(PRF_2D),s=0)
			surround_sizes[condition] = ((np.diff(spline.roots())/len(t)*reconstruction_radius) * stim_radius)[0]

	## EVALUATE OVERALL MODEL FIT QUALITY
	stats = {}
	stats['spearman'] = spearmanr(time_course,combined_model_prediction)[0]
	stats['pearson'] = pearsonr(time_course,combined_model_prediction)[0]
	stats['RSS'] = np.sum((time_course - combined_model_prediction)**2)
	stats['r_squared'] = 1 - stats['RSS']/np.sum((time_course - np.mean(time_course)) ** 2)
	stats['kendalls_tau'] = kendalltau(time_course,combined_model_prediction)[0]

	## CREATE SEPERATE RESULTS DICT PER CONDITION
	results = {}
	for condition in conditions:
		results[condition] = {}
		results[condition]['baseline'] = params['baseline'].value
		# params from fit
		for key in params.keys():
			if condition in key:
				if condition in key:
					# leave out the condition in the keys (as the results frames are identical across conditions)
					new_key = key[:-len(condition)-1]
				else:
					new_key = key
				results[condition][new_key] = params[key].value

		results[condition]['ecc'] = np.linalg.norm([params['xo_%s'%condition].value,params['yo_%s'%condition].value]) * stim_radius
		results[condition]['sigma_center'] *= stim_radius
		results[condition]['sigma_surround'] *= stim_radius

		# derived params
		results[condition]['polar'] = np.arctan2(params['yo_%s'%condition].value,params['xo_%s'%condition].value)
		results[condition]['fwhm'] = fwhms[condition]
		results[condition]['surround_size'] = surround_sizes[condition]
		results[condition]['SI'] = ((params['amp_surround_%s'%condition].value * (params['sigma_surround_%s'%condition].value**2) )
			/ (params['amp_center_%s'%condition].value * (params['sigma_center_%s'%condition].value**2) ))

		# if the resulting PRF falls outside of the stimulus radius,
		# set the multiplier here to 0 so that it falls off the retmaps
		if results[condition]['ecc'] < (stim_radius):
			multiplier = stats['r_squared']
		else:
			multiplier = 0.001

		# here for only voxels within stim region:
		results[condition]['real_polar_stim_region'] = np.cos(results[condition]['polar'])*np.arctanh(multiplier)
		results[condition]['imag_polar_stim_region'] = np.sin(results[condition]['polar'])*np.arctanh(multiplier)

		# and for all voxels:
		results[condition]['real_polar'] = np.cos(results[condition]['polar'])*np.arctanh(stats['r_squared'])
		results[condition]['imag_polar'] = np.sin(results[condition]['polar'])*np.arctanh(stats['r_squared'])

	#########################################################################################################################################################################################################################
	#### Plot results
	#########################################################################################################################################################################################################################

	if plotbool * (stats['r_squared']>0.0):# (np.random.randint(10)<10):#* (stats['r_squared']>0.1):#(stats['r_squared']>0.1):# * :# :#* (results['ecc'] < 3) :#:# * * randint ) #* :#* )

		n_TRs = n_TRs[0]
		n_runs = int(len(time_course) / n_TRs)
		if fit_on_all_data:
			plot_conditions = ['Ridge','All']
		else:
			plot_conditions = conditions + ['All_fit']
		plot_dir = os.path.join(plotdir, '%s'%roi)
		if not os.path.isdir(plot_dir): os.mkdir(plot_dir)

		f=pl.figure(figsize=(20,8)); rowi = (n_runs+4)

		import colorsys
		colors = np.array([colorsys.hsv_to_rgb(c,0.6,0.9) for c in np.linspace(0,1,3+1)])[:-1]

		for runi in range(n_runs):
			s = f.add_subplot(rowi,1,runi+1)
			pl.plot(time_course[n_TRs*runi:n_TRs*(runi+1)],'-ok',linewidth=0.75,markersize=2.5)#,label='data'
			if not fit_on_all_data:
				for ci, condition in enumerate(conditions):
					pl.plot(model_predictions[condition][n_TRs*runi:n_TRs*(runi+1)]+params['baseline'].value,color=colors[ci],label='%s model'%condition,linewidth=2)
				pl.plot([0,n_TRs],[params['baseline'].value,params['baseline'].value],color=colors[0],linewidth=1)
			else:
				pl.plot(combined_model_prediction[n_TRs*runi:n_TRs*(runi+1)],color=colors[0],label='model',linewidth=2)
			sn.despine(offset=10)
			pl.xlim(0,850)
			if runi == (n_runs-1):
				pl.xlabel('TRs')
			else:
				pl.xticks([])
			if runi == (n_runs/2):
				pl.legend(loc='best',fontsize=8)
				if 'psc' in postFix:
					pl.ylabel('% signal change')
				else:
					pl.ylabel('unkown unit')
			pl.yticks([int(np.min(time_course)),0,int(np.max(time_course))])
			pl.ylim([int(np.min(time_course)),int(np.max(time_course))])

		rowi = (n_runs+2)/2
		k = 0
		for ci, condition in enumerate(plot_conditions):
			k+= 1
			s = f.add_subplot(rowi,len(plot_conditions)*2,(rowi-1)*len(plot_conditions)*2+k,aspect='equal')
			pl.imshow(PRFs[condition],origin='lowerleft',interpolation='nearest',cmap=cm.coolwarm)

			pl.axis('off')
			s.set_title('%s PRF'%condition)

			k+= 1
			if not (condition == 'Ridge') + (condition == 'All_fit'):
				s = f.add_subplot(rowi,len(plot_conditions)*2,(rowi-1)*len(plot_conditions)*2+k)
				pl.imshow(np.ones((n_pixel_elements_raw,n_pixel_elements_raw)),cmap='gray')
				pl.clim(0,1)
				if model == 'OG':
					s.text(n_pixel_elements_raw/2,n_pixel_elements_raw/2, "\n%s PARAMETERS: \n\nbaseline: %.2f\nsize: %.2f\namplitude: %.6f\n\n\nDERIVED QUANTIFICATIONS: \n\nr-squared: %.2f\necc: %.2f\nFWHM: %.2f"%
						(condition,results[condition]['baseline'],results[condition]['sigma_center'],results[condition]['amp_center'],
							stats['r_squared'],results[condition]['ecc'],results[condition]['fwhm']),
						horizontalalignment='center',verticalalignment='center',fontsize=10,bbox={'facecolor':'white', 'alpha':1, 'pad':10})
				elif model == 'DoG':
					s.text(n_pixel_elements_raw/2,n_pixel_elements_raw/2, "\n%s PARAMETERS: \n\nbaseline: %.2f\nsd center: %.2f\nsd surround: %.2f\namp center: %.6f\namp surround: %.6f\n\nDERIVED QUANTIFICATIONS: \n\nr squared: %.2f\necc: %.2f\nFWHM: %.2f\nsurround size: %.2f\nsupression index: %.2f"
						%(condition,results[condition]['baseline'],results[condition]['sigma_center'],results[condition]['sigma_surround'],results[condition]['amp_center'],
						results[condition]['amp_surround'],stats['r_squared'],results[condition]['ecc'],results[condition]['fwhm'],results[condition]['surround_size'],
						results[condition]['SI']),horizontalalignment='center',verticalalignment='center',fontsize=10,bbox={'facecolor':'white', 'alpha':1, 'pad':10})
				pl.axis('off')

		# pl.tight_layout()
		pl.savefig(os.path.join(plot_dir, 'vox_%d_%d_%d.pdf'%(slice_no,voxno,n_pixel_elements_raw)))
		pl.close()

	return results, stats

reason: import numpy,from scipy,from sympy


text:
import os
import torch
import numpy as np
from torch.distributions.normal import Normal
from torch.nn.functional import mse_loss
from copy import deepcopy
from rlil.environments import State, action_decorator, Action
from rlil.initializer import get_writer, get_device, get_replay_buffer
from rlil import nn
from tqdm import tqdm
from .base import Agent, LazyAgent

class BRAC(Agent):
    """
    Behavior Regularized Actor Critic (BRAC)

    BRAC is a general framework for evaluating proposed offline RL methods
    as well as a number of simple baselines across a variety of offline
    continuous control tasks.
    https://arxiv.org/abs/1911.11361

    There are four main options for the behavior regularized actor critic framework:
    1 How to estimate the Q value (e.g. the BEAR uses weighted mixture of Q-values).
     According to the BRAC paper, double Q learning like TD3 works well.
    2 Which divergence to use. The paper says that KL divergence works well.
    divergence work well.
    3 Whether to learn lagrange multipliers alpha adaptively.
     The paper says that a fixed alpha is superior than an adaptive alpha.
    4 Whether to use a value penalty in the Q function objective or just use policy regularization.
    The paper says that the performance imploves slightly when using a value
     penalty in the Q function objective.

    Then, this implementation uses following tricks:
    1. double Q learning like TD3, 2. use KL divergence,
    3. fixed alpha, 4. Q function objective penalty.

    The behavior policy is estimated using simple bc with SoftDeterministicPolicy.
    Args:
        q_1 (QContinuous): An Approximation of the continuous action Q-function.
        q_2 (QContinuous): An Approximation of the continuous action Q-function.
        policy (SoftDeterministicPolicy): An Approximation of a soft deterministic policy.
        behavior_policy (SoftDeterministicPolicy): An approximation of the behavior policy.
        bc_iters (int): Number of training steps for behavior cloning.
        alpha (float): Value of lagrange multipliers. Trick 3.
        n_div_samples (int): Number of samples used for computing divergence.
        discount_factor (float): Discount factor for future rewards.
        minibatch_size (int): The number of experiences to sample in each training update.
    """

    def __init__(self,
                 q_1,
                 q_2,
                 policy,
                 behavior_policy,
                 bc_iters=5000,
                 alpha=0.1,
                 n_div_samples=4,
                 discount_factor=0.99,
                 minibatch_size=100,
                 ):
        # objects
        self.q_1 = q_1
        self.q_2 = q_2
        self.policy = policy
        self.behavior_policy = behavior_policy
        self.replay_buffer = get_replay_buffer()
        self.device = get_device()
        self.writer = get_writer()
        # hyperparameters
        self.alpha = alpha
        self.bc_iters = bc_iters
        self.n_div_samples = n_div_samples
        self.minibatch_size = minibatch_size
        self.discount_factor = discount_factor
        # private
        self._train_count = 0

    def act(self, states, reward=None):
        actions = Action(self.policy.no_grad(
            states.to(self.device))[0]).to("cpu")
        return actions

    def train(self):
        if self._train_count == 0:
            self.train_bc()

        self._train_count += 1

        # sample transitions from buffer
        (states, actions, rewards,
         next_states, _, _) = self.replay_buffer.sample(self.minibatch_size)

        # Trick 2: KL divergence regularization
        policy_mean, policy_logvar = self.policy.mean_logvar(states)
        behavior_mean, behavior_logvar = self.behavior_policy.mean_logvar(
            states)
        kl = nn.kl_gaussian(policy_mean, policy_logvar,
                            behavior_mean, behavior_logvar)

        # Trick 1: clipped double Q learning
        next_actions, _ = self.policy.target(next_states)
        q_values = torch.min(self.q_1.target(next_states, Action(next_actions)),
                             self.q_2.target(next_states, Action(next_actions)))

        # Trick 4: Q target with divergence penalty
        q_targets = rewards + self.discount_factor * \
            (q_values - self.alpha * kl.detach())

        self.q_1.reinforce(
            mse_loss(self.q_1(states, actions), q_targets))
        self.q_2.reinforce(
            mse_loss(self.q_2(states, actions), q_targets))

        # Update policy with a warmstart
        policy_loss = self.alpha * kl
        if self._train_count >= 5000:
            policy_actions, _ = self.policy(states)
            policy_loss -= self.q_1(states, Action(policy_actions))
        self.policy.reinforce(policy_loss.mean())

        self.writer.add_scalar('q_targets/mean', q_targets.detach().mean())
        self.writer.add_scalar('loss/kl', kl.detach().mean())
        self.writer.train_steps += 1

    def train_bc(self):
        for _ in tqdm(range(self.bc_iters)):
            (states, actions, _, _, _, _) = self.replay_buffer.sample(
                self.minibatch_size)
            bc_actions = Action(self.behavior_policy(states)[0])
            loss = mse_loss(bc_actions.features, actions.features)
            self.behavior_policy.reinforce(loss)

        # freeze behavior policy
        for param in self.behavior_policy.model.parameters():
            param.requires_grad = False

    def should_train(self):
        return True

    def make_lazy_agent(self, evaluation=False, store_samples=True):
        policy_model = deepcopy(self.policy.model)
        return BracLazyAgent(policy_model.to("cpu"),
                             evaluation=evaluation,
                             store_samples=store_samples)

    def load(self, dirname):
        for filename in os.listdir(dirname):
            if filename == 'policy.pt':
                self.policy.model = torch.load(os.path.join(
                    dirname, filename), map_location=self.device)
            if filename in ('q_1.pt'):
                self.q_1.model = torch.load(os.path.join(dirname, filename),
                                            map_location=self.device)
            if filename in ('q_2.pt'):
                self.q_1.model = torch.load(os.path.join(dirname, filename),
                                            map_location=self.device)
            if filename in ('behavior_policy.pt'):
                self.behavior_policy.model = torch.load(os.path.join(dirname, filename),
                                                        map_location=self.device)

class BracLazyAgent(LazyAgent):
    """
    Agent class for sampler.
    """

    def __init__(self, policy_model, *args, **kwargs):
        self._policy_model = policy_model
        super().__init__(*args, **kwargs)
        if self._evaluation:
            self._policy_model.eval()

    def act(self, states, reward):
        super().act(states, reward)
        self._states = states
        with torch.no_grad():
            outputs = self._policy_model(states, return_mean=True)
            self._actions = Action(outputs).to("cpu")
        return self._actions