# -*- coding: utf-8 -*-

from functools import partial
from lib.cache_utils import *
from lib.clip import *
from lib.collection_utils import *
from lib.gpu_utils import *
from lib.image_utils import *
from lib.math_utils import *
from lib.processing_utils import *
from moviepy.editor import VideoFileClip
import multiprocessing
from multiprocessing import Pool
from multiprocessing.dummy import Pool as ThreadPool
import numpy as np
import os
from PIL import Image, ImageDraw, ImageFilter
from pprint import pprint
import subprocess
import sys

def addVideoArgs(parser):
    parser.add_argument('-in', dest="INPUT_FILE", default="tmp/samples.csv", help="Input file")
    parser.add_argument('-ss', dest="EXCERPT_START", type=float, default=-1, help="Excerpt start in seconds")
    parser.add_argument('-sd', dest="EXCERPT_DUR", type=float, default=-1, help="Excerpt duration in seconds")
    parser.add_argument('-db', dest="MASTER_DB", type=float, default=0.0, help="Master +/- decibels to be applied to final audio")
    parser.add_argument('-dir', dest="MEDIA_DIRECTORY", default="media/sample/", help="Input file")
    parser.add_argument('-width', dest="WIDTH", default=1920, type=int, help="Output video width")
    parser.add_argument('-height', dest="HEIGHT", default=1080, type=int, help="Output video height")
    parser.add_argument('-fps', dest="FPS", default=24, type=int, help="Output video frames per second")
    parser.add_argument('-outframe', dest="OUTPUT_FRAME", default="tmp/sample/frame.%s.png", help="Output frames pattern")
    parser.add_argument('-out', dest="OUTPUT_FILE", default="output/sample.mp4", help="Output media file")
    parser.add_argument('-threads', dest="THREADS", default=1, type=int, help="Amount of parallel frames to process (too many may result in too many open files)")
    parser.add_argument('-overwrite', dest="OVERWRITE", action="store_true", help="Overwrite existing frames?")
    parser.add_argument('-ao', dest="AUDIO_ONLY", action="store_true", help="Render audio only?")
    parser.add_argument('-vo', dest="VIDEO_ONLY", action="store_true", help="Render video only?")
    parser.add_argument('-cache', dest="CACHE_VIDEO", action="store_true", help="Cache video clips?")
    parser.add_argument('-cd', dest="CACHE_DIR", default="tmp/cache/", help="Dir for caching data")
    parser.add_argument('-ckey', dest="CACHE_KEY", default="sample", help="Key for caching data")
    parser.add_argument('-verifyc', dest="VERIFY_CACHE", action="store_true", help="Add a step for verifying existing cache data?")
    parser.add_argument('-recalc', dest="RECALC_CLIP_SIZE", action="store_true", help="Recalculate clip sizes?")
    parser.add_argument('-rand', dest="RANDOM_SEED", default=1, type=int, help="Random seed to use for pseudo-randomness")
    parser.add_argument('-pad0', dest="PAD_START", default=1000, type=int, help="Pad the beginning")
    parser.add_argument('-pad1', dest="PAD_END", default=3000, type=int, help="Pad the end")
    parser.add_argument('-rvb', dest="REVERB", default=80, type=int, help="Reverberence (0-100)")
    parser.add_argument('-debug', dest="DEBUG", action="store_true", help="Debug mode?")
    parser.add_argument('-mdb', dest="MATCH_DB", default=-16, type=int, help="Match decibels, -9999 for none")
    parser.add_argument('-maxdb', dest="MAX_DB", default=-16, type=int, help="Max decibels, -9999 for none")
    parser.add_argument('-precision', dest="PRECISION", default=3, type=int, help="Precision for position and size")
    parser.add_argument('-margin', dest="CLIP_MARGIN", default=0.06666666667, type=float, help="Margin between clips as a percentage of a clip's width")
    parser.add_argument('-alphar', dest="ALPHA_RANGE", default="0.33,1.0", help="Alpha range")
    parser.add_argument('-brightr', dest="BRIGHTNESS_RANGE", default="0.33,1.0", help="Brightness range")
    parser.add_argument('-mcd', dest="MIN_CLIP_DUR", default=1200, type=int, help="Minumum clip duration")
    parser.add_argument('-noise', dest="NOISE", default=0.0, type=float, help="Amount of pixel noise to add")
    parser.add_argument('-maxa', dest="MAX_AUDIO_CLIPS", default=-1, type=int, help="Maximum number of audio clips to play")
    parser.add_argument('-keep', dest="KEEP_FIRST_AUDIO_CLIPS", default=-1, type=int, help="Ensure the middle x audio files play")
    parser.add_argument('-inva', dest="INVERT_LOUDEST", action="store_true", help="Clips on the outer edge are louder?")
    parser.add_argument('-fa', dest="FRAME_ALPHA", default=1.0, type=float, help="For adding frame content on top of previous frames; must be 0 <= x < 1; lower number = slower fade of prev frames")
    parser.add_argument('-rmode', dest="RESIZE_MODE", default="fill", help="Mode for resizing frames: fill, contain, or warp")
    parser.add_argument('-io', dest="CLIP_INITIAL_OFFSET", default=0, type=int, help="Milliseconds to offset the initial clip state for continuity between compositions")
    parser.add_argument('-probe', dest="PROBE", action="store_true", help="Just spit out duration info?")
    parser.add_argument('-frame', dest="OUTPUT_SINGLE_FRAME", default=-1, type=int, help="Output only a single frame (indicated frame number)")
    parser.add_argument('-frange', dest="FRAME_RANGE", default="1,0", help="Frame range to render")

def clipsToFrame(p, clips, pixelData, precision=3, customClipToArrFunction=None, baseImage=None, gpuProgram=None, postProcessingFunction=None, preProcessingFunction=None, globalArgs={}):
    filename = p["filename"]
    width = p["width"]
    height = p["height"]

    frame = getValue(p, "frame", 1)
    ms = getValue(p, "ms", 0)
    overwrite = getValue(p, "overwrite", False)
    verbose = getValue(p, "verbose", False)
    debug = getValue(p, "debug", False)
    saveFrame = getValue(p, "saveFrame", filename)

    frameAlpha = getValue(globalArgs, "frameAlpha", None)
    isSequential = getValue(globalArgs, "isSequential", False)
    container = getValue(globalArgs, "container", None)

    im = None
    fileExists = filename and os.path.isfile(filename) and not overwrite
    clipArr = None
    returnValue = None

    globalArgsCopy = globalArgs.copy()
    globalArgsCopy["frame"] = frame
    globalArgsCopy["debug"] = debug

    if not fileExists and saveFrame or not saveFrame or isSequential:
        clipArr = clipsToNpArr(clips, ms, width, height, precision, customClipToArrFunction=customClipToArrFunction, globalArgs=globalArgsCopy)

    # frame does not exist, create frame image
    if not fileExists:
        im = Image.new(mode="RGBA", size=(width, height), color=(0, 0, 0, 255))
        if preProcessingFunction is not None:
            baseImage = preProcessingFunction(baseImage, ms, globalArgs=globalArgs)
        if pixelData is None:
            im = clipsToFrameOnTheFly(clips, clipArr, width, height, precision, baseImage=baseImage, globalArgs=globalArgsCopy)
        else:
            im = clipsToFrameGPU(clipArr, width, height, pixelData, precision, baseImage=baseImage, gpuProgram=gpuProgram, globalArgs=globalArgs)
        im = im.convert("RGB")
        # check to see if we're applying container-level effects
        if container is not None:
            # right now we only care about blur
            blur = container.vector.getBlur(ms)
            if blur > 0.0:
                im = blurImage(im, blur)
        # check to see if there's post-processing to be done
        if postProcessingFunction is not None:
            im = postProcessingFunction(im, ms, globalArgs=globalArgs)
        # save if necessary
        if saveFrame:
            im.save(filename)
            print("Saved frame %s" % filename)
        if frameAlpha is None:
            returnValue = im

    # frame has alpha, so darken for next frame
    if frameAlpha is not None and 0.0 <= frameAlpha < 1.0:
        if im is None:
            im = Image.open(filename)
        blackOverlay = Image.new(mode="RGB", size=im.size, color=(0, 0, 0))
        im = Image.blend(im, blackOverlay, frameAlpha)
        returnValue = im

    return returnValue

def clipsToFrameOnTheFly(clips, clipArrs, width, height, precision=3, baseImage=None, globalArgs={}):
    debug = getValue(globalArgs, "debug", False)
    vthreads = getValue(globalArgs, "vthreads", 1)
    precisionMultiplier = int(10 ** precision)

    validClips = []
    validClipArrs = []
    for i in range(len(clipArrs)):
        clip = clipArrToDict(clipArrs[i], precision)
        # only take clips that are visible
        if clip["width"] > 0.0 and clip["height"] > 0.0 and clip["alpha"] > 0.0:
            clip["filename"] = clips[i].props["filename"]
            clip["t"] = clips[i].props["start"] + clips[i].props["dur"] * clip["tn"]
            validClips.append(clip)
            validClipArrs.append(clipArrs[i])

    if debug:
        if baseImage is None:
            baseImage = Image.new(mode="RGB", size=(width, height), color=(0, 0, 0))
        draw = ImageDraw.Draw(baseImage)
        for clip in validClips:
            x = clip["x"]
            y = clip["y"]
            w = clip["width"]
            h = clip["height"]
            c = roundInt(clip["alpha"] * 255)
            color = (c, c, c)
            draw.rectangle([x, y, x+w, y+h], fill=color)
    else:
        validClips = addIndices(validClips, "_index")
        filenames = groupList(validClips, "filename")
        threads = getThreadCount(min(len(filenames), vthreads))
        pool = ThreadPool(threads)
        pixelData = pool.map(samplesToPixels, filenames)
        pool.close()
        pool.join()
        clipsPixelData = [[] for i in range(len(validClips))]
        for pd in pixelData:
            for i, pixels in pd:
                clipsPixelData[i] = [pixels]
        baseImage = clipsToFrameGPU(validClipArrs, width, height, clipsPixelData, precision, baseImage, globalArgs=globalArgs)

    return baseImage

def clipsToFrameGPU(clips, width, height, clipsPixelData, precision=3, baseImage=None, gpuProgram=None, globalArgs={}):
    c = getValue(globalArgs, "colors", 3)
    offset = 0
    maxScaleFactor = 2.0
    precisionMultiplier = int(10 ** precision)

    # filter out clips with no pixels, or zero [width, height, alpha]
    # keep track of how many pixels we'll need
    indices = []
    pixelCount = 0
    for i in range(len(clips)):
        clip = clipArrToDict(clips[i], precision)
        frameCount = len(clipsPixelData[i])
        # only take clips that are visible
        if frameCount > 0 and clip["width"] > 0.0 and clip["height"] > 0.0 and clip["alpha"] > 0.0:
            indices.append(i)
            tn = clip["tn"]
            h, w, _ = clipsPixelData[i][roundInt(tn * (frameCount-1))].shape
            # we want to resample if scaled too much
            scaleFactor = 1.0 * w / clip["width"]
            if scaleFactor > maxScaleFactor:
                w = roundInt(clip["width"])
                h = roundInt(clip["height"])
            # we need to resize if blurred or rotated
            if clip["blur"] > 0.0 or clip["rotation"] % 360.0 > 0.0:
                _x, _y, newW, newH = bboxRotate(0, 0, roundInt(clip["width"]), roundInt(clip["height"]), angle=45.0)
                w = roundInt(newW)
                h = roundInt(newH)
            pixelCount += int(h*w*c)

    validCount = len(indices)
    propertyCount = Clip.gpuPropertyCount
    properties = np.zeros((validCount, propertyCount), dtype=np.int32)
    pixelData = np.zeros(pixelCount, dtype=np.uint8)
    for i, clipIndex in enumerate(indices):
        clipArr = clips[clipIndex]
        x, y, tw, th, alpha, t, zindex, rotation, blur, brightness = tuple(clipArr)
        clip = clipArrToDict(clipArr, precision)
        clipPixelData = clipsPixelData[clipIndex]
        frameCount = len(clipPixelData)
        tn = clip["tn"]
        pixels = clipPixelData[roundInt(tn * (frameCount-1))]
        h, w, _c = pixels.shape
        # we want to resample if scaled too much
        scaleFactor = 1.0 * w / clip["width"]
        if scaleFactor > maxScaleFactor or clip["blur"] > 0.0 or clip["rotation"] % 360.0 > 0.0:
            rw = roundInt(clip["width"])
            rh = roundInt(clip["height"])
            im = None
            # assume we are debugging if single pixel
            if h==1 and w==1:
                newPixels = np.zeros((rh, rw, _c), dtype=np.uint8)
                newPixels[:,:] = pixels[0,0]
                im = Image.fromarray(newPixels, mode="RGB")
            else:
                im = Image.fromarray(pixels, mode="RGB")

            # apply effects before resize for better quality
            if clip["blur"] > 0.0 or clip["rotation"] % 360.0 > 0.0:
                im, _x, _y = applyEffects(im, 0, 0, clip["rotation"], clip["blur"], colors=c)
                # retrieve new coordinates based on target/resized size
                newX, newY, newW, newH = bboxRotate(clip["x"], clip["y"], roundInt(clip["width"]), roundInt(clip["height"]), angle=45.0)
                rw = roundInt(newW)
                rh = roundInt(newH)
                # x, y, tw, th changes if we rotate or blur
                x = roundInt(newX * precisionMultiplier)
                y = roundInt(newY * precisionMultiplier)
                tw = roundInt(newW * precisionMultiplier)
                th = roundInt(newH * precisionMultiplier)

            # resize image
            imW, imH = im.size
            if imW != rw or imH != rh:
                resampleType = Image.LANCZOS if imW > rw else Image.NEAREST
                im = im.resize((rw, rh), resample=resampleType)

            # im.save("output/test_pil.png")
            pixels = np.array(im)
            h, w, _c = pixels.shape
        # pixels are size 3, but need size 4
        if c > _c:
            fillVals = np.full((h, w, 1), 255, dtype='uint8')
            pixels = np.concatenate((pixels, fillVals), axis=2)
        properties[i] = np.array([offset, x, y, w, h, tw, th, alpha, zindex, brightness])
        px0 = offset
        px1 = px0 + int(h*w*c)
        pixelData[px0:px1] = pixels.reshape(-1)
        offset += int(h*w*c)

    pixels = clipsToImageGPU(width, height, pixelData, properties, c, precision, gpuProgram=gpuProgram, baseImage=baseImage)
    return Image.fromarray(pixels, mode="RGB")

def compileFrames(infile, fps, outfile, padZeros, audioFile=None, quality="high"):
    print("Compiling frames...")
    padStr = '%0'+str(padZeros)+'d'

    # https://trac.ffmpeg.org/wiki/Encode/H.264
    # presets: veryfast, faster, fast, medium, slow, slower, veryslow
    #   slower = better quality
    # crf: 0 is lossless, 23 is the default, and 51 is worst possible quality
    #   17 or 18 to be visually lossless or nearly so
    preset = "veryslow"
    crf = "18"
    if quality=="medium":
        preset = "medium"
        crf = "23"
    elif quality=="low":
        preset = "medium"
        crf = "28"

    if audioFile:
        command = ['ffmpeg','-y',
                    '-framerate',str(fps)+'/1',
                    '-i',infile % padStr,
                    '-i',audioFile,
                    '-c:v','libx264',
                    '-preset', preset,
                    '-crf', crf,
                    '-r',str(fps),
                    '-pix_fmt','yuv420p',
                    '-c:a','aac',
                    # '-q:v','1',
                    '-b:a', '192k',
                    # '-shortest',
                    outfile]
    else:
        command = ['ffmpeg','-y',
                    '-framerate',str(fps)+'/1',
                    '-i',infile % padStr,
                    '-c:v','libx264',
                    '-preset', preset,
                    '-crf', crf,
                    '-r',str(fps),
                    '-pix_fmt','yuv420p',
                    # '-q:v','1',
                    outfile]
    print(" ".join(command))
    finished = subprocess.check_call(command)
    print("Done.")

# resize video using fill method
def fillVideo(video, w, h):
    ratio = 1.0 * w / h
    vw, vh = video.size
    vratio = 1.0 * vw / vh

    # first, resize video
    newW = w
    newH = h
    if vratio > ratio:
        newW = h * vratio
    else:
        newH = w / vratio
    resized = video.resize((roundInt(newW), roundInt(newH)))

    # and then crop
    x = 0
    y = 0
    if vratio > ratio:
        x = roundInt((newW - w) * 0.5)
    else:
        y = roundInt((newH - h) * 0.5)
    cropped = resized.crop(x, y, w, h)

    return cropped

def frameToMs(frame, fps, roundResult=True):
    result = (1.0 * frame / fps) * 1000.0
    if roundResult:
        result = roundInt(result)
    return result

def getAlpha(clip):
    alpha = clip["alpha"] if "alpha" in clip and clip["alpha"] < 1.0 else 1.0
    return roundInt(alpha*255)

def getCenter(clip):
    cx = clip["x"] + clip["width"] * 0.5
    cy = clip["y"] + clip["height"] * 0.5
    return (cx, cy)

def getDurationFromFile(filename, accurate=False):
    result = 0
    if os.path.isfile(filename):
        if accurate:
            try:
                video = VideoFileClip(filename, audio=False)
                result = video.duration
                video.reader.close()
                del video
            except UnicodeDecodeError:
                print("Unicode decode error for %s" % filename)
                result = 0
            except IOError:
                print("I/O error for %s" % filename)
                result = 0
        else:
            command = ['ffprobe', '-v', 'error', '-show_entries', 'format=duration', '-of', 'default=noprint_wrappers=1:nokey=1', filename]
            try:
                result = subprocess.check_output(command).strip()
            except subprocess.CalledProcessError:
                result = 0
    return float(result)

def getEmptyVideoClipImage(clip):
    cw = roundInt(clip["width"])
    ch = roundInt(clip["height"])
    videoPixels = np.zeros((ch, cw, 4), dtype='uint8')
    clipImg = Image.fromarray(videoPixels, mode="RGBA")
    return clipImg

# e.g. returns ['audio', 'video'] for a/v files
def getMediaTypes(filename):
    result = []
    if os.path.isfile(filename):
        command = ['ffprobe', '-loglevel', 'error', '-show_entries', 'stream=codec_type', '-of', 'csv=p=0', filename]
        # print(" ".join(command))
        try:
            result = [line.decode("utf-8") for line in subprocess.check_output(command).splitlines()]
        except subprocess.CalledProcessError:
            result = []
    return result

def getSolidPixels(color, width=100, height=100):
    dim = len(color)
    pixels = np.zeros((height, width, dim))
    pixels[:,:] = color
    return pixels

def getVideoClipImage(video, videoDur, clip=None, t=None, resizeMode="fill", resampleType="default"):
    videoT = clip["t"] / 1000.0 if t is None else t / 1000.0
    cw = None
    ch = None
    if clip is not None:
        cw = roundInt(clip["width"])
        ch = roundInt(clip["height"])
    delta = videoDur - videoT
    # check if we need to loop video clip
    if delta < 0:
        videoT = videoT % videoDur
    # hack: ffmpeg sometimes has trouble reading the very end of the video; choose 500ms from end
    elif delta < 0.5:
        videoT = videoDur - 0.5
    # a numpy array representing the RGB picture of the clip
    try:
        videoPixels = video.get_frame(videoT)
    except IOError:
        print("Could not read pixels for %s at time %s" % (video.filename, videoT))
        if ch is None or cw is None:
            cw = 1280
            ch = 720
        videoPixels = np.zeros((ch, cw, 3), dtype='uint8')
    clipImg = Image.fromarray(videoPixels, mode="RGB")
    if cw is not None and ch is not None:
        clipImg = resizeImage(clipImg, cw, ch, resizeMode, resampleType)
    return clipImg

def getVideoClipImageFromFile(fn, nt=0.5):
    video = VideoFileClip(fn, audio=False)
    videoDur = video.duration
    t = nt * videoDur * 1000

    image = getVideoClipImage(video, videoDur, clip=None, t=t)
    video.reader.close()
    del video
    return image

def getRotation(clip):
    rotation = clip["rotation"] if "rotation" in clip else 0.0
    angle = normalizeAngle(rotation)
    return angle

def hasAudio(filename):
    return ("audio" in getMediaTypes(filename))

def isVideoExtension(ext):
    return (ext in ['.mp4', '.mov', '.avi', '.wmv'])

def loadVideoPixelData(clips, fps, cacheDir="tmp/", width=None, height=None, verifyData=True, cache=True, resizeMode="fill"):
    # load videos
    filenames = list(set([clip.props["filename"] for clip in clips]))
    fileCount = len(filenames)
    msStep = frameToMs(1, fps, False)
    clipsPixelData = [None] * len(clips)

    for i, clip in enumerate(clips):
        if "maxWidth" in clip.props and "maxHeight" in clip.props:
            break
        if i <= 0:
            print("Warning: max width/height not set, setting it to given width/height")
        if "maxWidth" not in clip.props:
            clip.setProp("maxWidth", clip.props["width"])
        if "maxHeight" not in clip.props:
            clip.setProp("maxHeight", clip.props["height"])

    # only open one video at a time
    for i, fn in enumerate(filenames):
        # check for cache fors filename
        cacheFn = cacheDir + os.path.basename(fn) + ".p"
        loaded = False
        fileCacheData = []
        if cache:
            loaded, fileCacheData = loadCacheFile(cacheFn)
        vclips = [c for c in clips if fn==c.props["filename"]]
        valid = True

        # Verify loaded data
        if loaded and verifyData:
            print("Verifying cache data for %s..." % cacheFn)
            clipTimesSet = set(fileCacheData[0])
            for clip in vclips:
                start = clip.props["start"]
                end = start + clip.props["dur"]
                ms = start
                while ms < end:
                    t = roundInt(ms)
                    ms += msStep
                    if t not in clipTimesSet:
                        print("%s not found in %s. Resetting cache data" % (t, cacheFn))
                        loaded = False
                        break
                    index = fileCacheData[0].index(t)
                    clipH, clipW, _ = fileCacheData[1][index].shape
                    if roundInt(clip.props["maxWidth"]) > clipW:
                        print("Clip width is too small (%s > %s) for %s at %s. Resetting cache data" % (clip.props["maxWidth"], clipW, cacheFn, t))
                        loaded = False
                        break
                    # TODO: Add check for aspect ratio?
                if not loaded:
                    break
            print("Verified cache data for %s" % cacheFn)

        if not loaded:
            print("No cache for %s, rebuilding..." % fn)
            video = VideoFileClip(fn, audio=False)
            videoDur = video.duration
            clipTimes = []
            clipPixels = []

            # Already loaded cache data; append to it rather than overwrite it
            if len(fileCacheData) > 0:
                clipTimes, clipPixels = fileCacheData

            # extract frames from videos
            vclipCount = len(vclips)
            for j, clip in enumerate(vclips):
                start = clip.props["start"]
                end = start + clip.props["dur"]
                ms = start
                while ms < end:
                    fclip = clip.props.copy()
                    fclip["width"] = fclip["maxWidth"]
                    fclip["height"] = fclip["maxHeight"]
                    t = roundInt(ms)
                    ms += msStep
                    # already exists, check size
                    existIndex = None
                    if t in set(clipTimes):
                        existIndex = clipTimes.index(t)
                        clipH, clipW, _ = clipPixels[existIndex].shape
                        if roundInt(fclip["width"]) <= clipW:
                            continue
                    clipResizeMode = getValue(clip.props, "resizeMode", resizeMode)
                    clipImg = getVideoClipImage(video, videoDur, fclip, t, clipResizeMode)
                    if existIndex is not None:
                        clipPixels[existIndex] = np.array(clipImg, dtype=np.uint8)
                    else:
                        clipTimes.append(t)
                        clipPixels.append(np.array(clipImg, dtype=np.uint8))
                printProgress(j+1, vclipCount)

            # save cache data
            fileCacheData = (clipTimes, clipPixels)

            if cache:
                saveCacheFile(cacheFn, fileCacheData, overwrite=True)

            # close video to free up memory
            video.reader.close()
            del video

        # assign pixel data to clips
        for clip in vclips:
            start = clip.props["start"]
            end = start + clip.props["dur"]
            ms = start
            pixelData = []
            while ms < end:
                t = roundInt(ms)
                ms += msStep
                index = fileCacheData[0].index(t)
                pixelData.append(fileCacheData[1][index])
            clipsPixelData[clip.props["index"]] = pixelData
            # clip.setProp("framePixelData", pixelData)

        printProgress(i+1, fileCount)

    print("Finished loading pixel data.")
    return clipsPixelData

def loadVideoPixelDataDebug(clipCount):
    clipsPixelData = np.zeros((clipCount, 1, 1, 1, 3))
    for i in range(clipCount):
        clipsPixelData[i, 0, 0, 0] = getRandomColor(i)
    return clipsPixelData

def loadVideoPixelDataFromFrames(frames, clips, containerW, containerH, fps, cacheDir="tmp/", cacheKey="sample", verifyData=True, cache=True, debug=False, precision=3, customClipToArrFunction=None, customClipToArrCalcFunction=None, globalArgs={}):
    frameCount = len(frames)
    clipCount = len(clips)
    precisionMultiplier = int(10 ** precision)
    cacheFile = cacheKey + "_maxes.p"
    resizeMode = getValue(globalArgs, "resizeMode", "fill")
    recalculateClipSizes = getValue(globalArgs, "recalculateClipSizes", False)

    if debug:
        clipsPixelData = loadVideoPixelDataDebug(clipCount)
        return clipsPixelData

    loaded = False
    clipWidthMaxes = None
    if cache:
        loaded, clipWidthMaxes = loadCacheFile(cacheDir+cacheFile)

    if not loaded or len(clipWidthMaxes) != clipCount or recalculateClipSizes:
        print("Calculating clip size/position from frame sequence...")
        clipWidthMaxes = np.zeros(clipCount, dtype=np.int32)
        clipCompare = np.zeros((2, clipCount), dtype=np.int32)
        ccfunction = customClipToArrFunction
        # override custom clip arr function for calcuation
        if customClipToArrCalcFunction is "default":
            ccfunction = None
        elif customClipToArrCalcFunction is not None:
            ccfunction = customClipToArrCalcFunction
        for i, frame in enumerate(frames):
            ms = frame["ms"]
            # frameClips = clipsToDictsGPU(clips, ms, container, precision)
            frameClips = clipsToNpArr(clips, ms, containerW, containerH, precision, customClipToArrFunction=ccfunction, globalArgs=globalArgs)
            clipCompare[0] = clipWidthMaxes
            clipCompare[1] = frameClips[:,2] # just take the width (2)
            clipWidthMaxes = np.amax(clipCompare, axis=0)
            printProgress(i+1, frameCount)
        if cache:
            saveCacheFile(cacheDir+cacheFile, clipWidthMaxes, overwrite=True)

    # visualize the result
    # from matplotlib import pyplot as plt
    # totalMax = np.amax(clipWidthMaxes)
    # colors = clipWidthMaxes.astype(np.float32) / totalMax * 255.0
    # colors = colors.astype(np.uint8)
    # im = Image.fromarray(colors, mode="L")
    # plt.imshow(np.asarray(im))
    # plt.show()
    # sys.exit()

    # update clips with max width/height
    for i, clip in enumerate(clips):
        width = max(1.0 * clipWidthMaxes[clip.props["index"]] / precisionMultiplier, clip.props["width"])
        height = width / clip.vector.aspectRatio
        clip.setProp("maxWidth", width)
        clip.setProp("maxHeight", height)
        # print("%s, %s" % (clip.props["width"], clip.props["height"]))

    clipsPixelData = loadVideoPixelData(clips, fps, cacheDir=cacheDir, verifyData=verifyData, cache=cache, resizeMode=resizeMode)

    return clipsPixelData

def msToFrame(ms, fps):
    return roundInt((ms / 1000.0) * fps)

def parseVideoArgs(args):
    d = vars(args)
    d["THREADS"] = min(args.THREADS, multiprocessing.cpu_count()) if args.THREADS > 0 else multiprocessing.cpu_count()
    d["AUDIO_OUTPUT_FILE"] = args.OUTPUT_FILE.replace(".mp4", ".mp3")
    d["MS_PER_FRAME"] = frameToMs(1, args.FPS, False)
    d["CACHE_VIDEO"] = args.CACHE_VIDEO
    d["VOLUME_RANGE"] = tuple([float(v) for v in args.VOLUME_RANGE.strip().split(",")]) if "VOLUME_RANGE" in d else (0.0, 1.0)
    d["ALPHA_RANGE"] =  tuple([float(v) for v in args.ALPHA_RANGE.strip().split(",")])
    d["BRIGHTNESS_RANGE"] =  tuple([float(v) for v in args.BRIGHTNESS_RANGE.strip().split(",")])
    d["FRAME_RANGE"] =  tuple([int(v) for v in args.FRAME_RANGE.strip().split(",")])
    d["VIDEO_THREADS"] = args.VIDEO_THREADS if "VIDEO_THREADS" in d else 1
    if args.OUTPUT_SINGLE_FRAME > 0:
        d["VIDEO_ONLY"] = True

def processFrames(params, clips, clipsPixelData, threads=1, precision=3, verbose=True, customClipToArrFunction=None, postProcessingFunction=None, preProcessingFunction=None, globalArgs={}):
    if len(params) < 1:
        return

    count = len(params)
    print("Processing %s frames" % count)
    threads = getThreadCount(threads)

    frameAlpha = getValue(globalArgs, "frameAlpha", 1.0)
    isSequential = getValue(globalArgs, "isSequential", False)
    baseImage = getValue(globalArgs, "baseImage", None)
    propagateFrames = (0.0 <= frameAlpha < 1.0)
    if propagateFrames:
        isSequential = True

    # load gpu program
    p0 = params[0]
    colorDimensions = getValue(globalArgs, "colors", 3)
    pcount = Clip.gpuPropertyCount
    gpuProgram = loadMakeImageProgram(p0["width"], p0["height"], pcount, colorDimensions, precision)

    if threads > 1 and not isSequential:
        pool = ThreadPool(threads)
        pclipsToFrame = partial(clipsToFrame, clips=clips, pixelData=clipsPixelData, precision=precision, customClipToArrFunction=customClipToArrFunction, baseImage=baseImage, gpuProgram=gpuProgram, postProcessingFunction=postProcessingFunction, preProcessingFunction=preProcessingFunction, globalArgs=globalArgs)
        pool.map(pclipsToFrame, params)
        pool.close()
        pool.join()
    else:
        prevImage = None
        for i, p in enumerate(params):
            baseImage = prevImage if propagateFrames else baseImage
            prevImage = clipsToFrame(p, clips=clips, pixelData=clipsPixelData, precision=precision, customClipToArrFunction=customClipToArrFunction, baseImage=baseImage, gpuProgram=gpuProgram, postProcessingFunction=postProcessingFunction, preProcessingFunction=preProcessingFunction, globalArgs=globalArgs)
            if verbose:
                printProgress(i+1, count)

def samplesToPixels(f):
    video = VideoFileClip(f["filename"], audio=False)
    videoDur = video.duration
    fclips = f["items"]
    pixelData = [0 for i in range(len(fclips))]
    for i, fclip in enumerate(fclips):
        clipImg = getVideoClipImage(video, videoDur, {"width": ceilInt(fclip["width"]), "height": ceilInt(fclip["height"]), "t": fclip["t"]})
        pixels = np.array(clipImg)
        pixelData[i] = (fclip["_index"], pixels)
    video.reader.close()
    del video
    return pixelData


import bpy
from imutils import face_utils
import dlib
import cv2
import time
import numpy
from bpy.props import FloatProperty

class OpenCVAnimOperator(bpy.types.Operator):
    """Operator which runs its self from a timer"""
    bl_idname = "wm.opencv_operator"
    bl_label = "OpenCV Animation Operator"

    # p = our pre-treined model directory
    #p = "D:\\Users\\jason\\shape_predictor_68_face_landmarks.dat" # Windows
    #p = "/Users/jason/Downloads/shape_predictor_68_face_landmarks.dat" # macOS
    p = "/home/jason/Downloads/shape_predictor_68_face_landmarks.dat" # linux
    detector = dlib.get_frontal_face_detector()
    predictor = dlib.shape_predictor(p)
    # rig_name - take it from your scene collection tree
    rig_name = "RIG-Vincent"

    _timer = None
    _cap  = None

    width = 800
    height = 600

    stop :bpy.props.BoolProperty()

    # 3D model points.
    model_points = numpy.array([
                                (0.0, 0.0, 0.0),             # Nose tip
                                (0.0, -330.0, -65.0),        # Chin
                                (-225.0, 170.0, -135.0),     # Left eye left corner
                                (225.0, 170.0, -135.0),      # Right eye right corne
                                (-150.0, -150.0, -125.0),    # Left Mouth corner
                                (150.0, -150.0, -125.0)      # Right mouth corner
                            ], dtype = numpy.float32)
    # Camera internals
    camera_matrix = numpy.array(
                            [[height, 0.0, width/2],
                            [0.0, height, height/2],
                            [0.0, 0.0, 1.0]], dtype = numpy.float32
                            )

    # Keeps a moving average of given length
    def smooth_value(self, name, length, value):
        if not hasattr(self, 'smooth'):
            self.smooth = {}
        if not name in self.smooth:
            self.smooth[name] = numpy.array([value])
        else:
            self.smooth[name] = numpy.insert(arr=self.smooth[name], obj=0, values=value)
            if self.smooth[name].size > length:
                self.smooth[name] = numpy.delete(self.smooth[name], self.smooth[name].size-1, 0)
        sum = 0
        for val in self.smooth[name]:
            sum += val
        return sum / self.smooth[name].size

    # Keeps min and max values, then returns the value in a ranve 0 - 1
    def get_range(self, name, value):
        if not hasattr(self, 'range'):
            self.range = {}
        if not name in self.range:
            self.range[name] = numpy.array([value, value])
        else:
            self.range[name] = numpy.array([min(value, self.range[name][0]), max(value, self.range[name][1])] )
        val_range = self.range[name][1] - self.range[name][0]
        if val_range != 0:
            return (value - self.range[name][0]) / val_range
        else:
            return 0

    def modal(self, context, event):

        if (event.type in {'RIGHTMOUSE', 'ESC'}) or self.stop == True:
            self.cancel(context)
            return {'CANCELLED'}

        if event.type == 'TIMER':
            self.init_camera()
            _, image = self._cap.read()
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            rects = self.detector(gray, 0)
            # bpy.context.scene.frame_set(frame_num)

            # For each detected face, find the landmark.
            for (i, rect) in enumerate(rects):
                shape = self.predictor(gray, rect)
                shape = face_utils.shape_to_np(shape)

                #2D image points. If you change the image, you need to change vector
                image_points = numpy.array([shape[30],     # Nose tip - 31
                                            shape[8],      # Chin - 9
                                            shape[36],     # Left eye left corner - 37
                                            shape[45],     # Right eye right corne - 46
                                            shape[48],     # Left Mouth corner - 49
                                            shape[54]      # Right mouth corner - 55
                                        ], dtype = numpy.float32)

                dist_coeffs = numpy.zeros((4,1)) # Assuming no lens distortion

                if hasattr(self, 'rotation_vector'):
                    (success, self.rotation_vector, self.translation_vector) = cv2.solvePnP(self.model_points,
                        image_points, self.camera_matrix, dist_coeffs, flags=cv2.SOLVEPNP_ITERATIVE,
                        rvec=self.rotation_vector, tvec=self.translation_vector,
                        useExtrinsicGuess=True)
                else:
                    (success, self.rotation_vector, self.translation_vector) = cv2.solvePnP(self.model_points,
                        image_points, self.camera_matrix, dist_coeffs, flags=cv2.SOLVEPNP_ITERATIVE,
                        useExtrinsicGuess=False)

                if not hasattr(self, 'first_angle'):
                    self.first_angle = numpy.copy(self.rotation_vector)

                bones = bpy.data.objects[self.rig_name].pose.bones

                bones["head_fk"].rotation_euler[0] = self.smooth_value("h_x", 3, (self.rotation_vector[0] - self.first_angle[0])) / 1   # Up/Down
                bones["head_fk"].rotation_euler[2] = self.smooth_value("h_y", 3, -(self.rotation_vector[1] - self.first_angle[1])) / 1.5  # Rotate
                bones["head_fk"].rotation_euler[1] = self.smooth_value("h_z", 3, (self.rotation_vector[2] - self.first_angle[2])) / 1.3   # Left/Right

                bones["mouth_ctrl"].location[2] = self.smooth_value("m_h", 2, -self.get_range("mouth_height", numpy.linalg.norm(shape[62] - shape[66])) * 0.06 )
                bones["mouth_ctrl"].location[0] = self.smooth_value("m_w", 2, (self.get_range("mouth_width", numpy.linalg.norm(shape[54] - shape[48])) - 0.5) * -0.04)
                bones["brow_ctrl_L"].location[2] = self.smooth_value("b_l", 3, (self.get_range("brow_left", numpy.linalg.norm(shape[19] - shape[27])) -0.5) * 0.04)
                bones["brow_ctrl_R"].location[2] = self.smooth_value("b_r", 3, (self.get_range("brow_right", numpy.linalg.norm(shape[24] - shape[27])) -0.5) * 0.04)

                bones["head_fk"].keyframe_insert(data_path="rotation_euler", index=-1)
                bones["mouth_ctrl"].keyframe_insert(data_path="location", index=-1)
                bones["brow_ctrl_L"].keyframe_insert(data_path="location", index=2)
                bones["brow_ctrl_R"].keyframe_insert(data_path="location", index=2)

                for (x, y) in shape:
                    cv2.circle(image, (x, y), 2, (0, 255, 255), -1)

            cv2.imshow("Output", image)
            cv2.waitKey(1)

        return {'PASS_THROUGH'}

    def init_camera(self):
        if self._cap == None:
            self._cap = cv2.VideoCapture(0)
            self._cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
            self._cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
            self._cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            time.sleep(0.5)

    def stop_playback(self, scene):
        print(format(scene.frame_current) + " / " + format(scene.frame_end))
        if scene.frame_current == scene.frame_end:
            bpy.ops.screen.animation_cancel(restore_frame=False)

    def execute(self, context):
        bpy.app.handlers.frame_change_pre.append(self.stop_playback)

        wm = context.window_manager
        self._timer = wm.event_timer_add(0.02, window=context.window)
        wm.modal_handler_add(self)
        return {'RUNNING_MODAL'}

    def cancel(self, context):
        wm = context.window_manager
        wm.event_timer_remove(self._timer)
        cv2.destroyAllWindows()
        self._cap.release()
        self._cap = None

def register():
    bpy.utils.register_class(OpenCVAnimOperator)

def unregister():
    bpy.utils.unregister_class(OpenCVAnimOperator)

if __name__ == "__main__":
    register()

    # test call
    #bpy.ops.wm.opencv_operator()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # 强制cpu计算
import tensorflow as tf
import mxnet as mx
from sklearn.metrics import roc_curve, auc
import prettytable as pt
from tqdm import tqdm
import numpy as np
import cv2
import pickle
import gc
from verification import evaluate
#
PB_DIR = "./mobilenet_arcface_optimized.pb"
VALID_DIR = "./valid_dataset"
class FACENET_PB:
    def __init__(self, model_path):
        graph = tf.Graph()
        with graph.as_default():
            with open(model_path, 'rb') as f:
                graph_def = tf.GraphDef.FromString(f.read())
                tf.import_graph_def(graph_def, name='')

        config = tf.ConfigProto(
            allow_soft_placement=True)
        self.sess = tf.Session(graph=graph, config=config)

    def get_embedding(self, images):
        feeds = {
            self.sess.graph.get_tensor_by_name("input_1:0") : images,     # 人脸抠图 4D float tensor with format [faceimg1,faceimg2 ....] 单张人脸faceimg1 112X112 像素,归一化[-1,1]之间,RGB通道
        }
        fetches = self.sess.graph.get_tensor_by_name("embeddings/l2_normalize:0")    # 抽取特征 2D float tensorflow [[1X512],....] 每行为单张人脸特征, 1X512

        embedding = self.sess.run(fetches, feed_dict=feeds)
        return embedding

    def free(self):
        self.sess.close()

#  model = FACENET_PB(PB_DIR)
#  img = cv2.imread('./img_0.jpg')
#  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
#img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
#  img = img - 127.5
#  img = img * 0.0078125
#  img = np.expand_dims(img,axis=0)
#  print("img:",img.shape,img)
#  emb = model.get_embedding(img)
#  print("emb:",emb)
#
def predict_evaluate_data_PB(db_name, image_size , model, embedsize = 512):
    batch = 100
    bins, issame_list = pickle.load(open(os.path.join(db_name), 'rb'), encoding='bytes')
    datasets = np.zeros( (batch , image_size[0], image_size[1], 3))
    embedings = np.zeros((len(issame_list)*2, embedsize))
    count = 0

    for i in tqdm(range(len(issame_list)*2),desc="evalutaing %s ..."%db_name) :
        _bin = bins[i]
        img = mx.image.imdecode(_bin).asnumpy()
        del _bin
        # img = cv2.imdecode(np.fromstring(_bin, np.uint8), -1)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        #img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        img = img - 127.5
        img = img * 0.0078125
        datasets[count, ...] = img
        count += 1
        i = i+1
        if i % batch  == 0:
            count = 0
            #  print('loading bin', i)
            bz_data = datasets[0:batch,...]
            bz_embedings = np.zeros([len(bz_data),512])
            for ii in range(len(bz_data)):
                em = model.get_embedding(np.expand_dims(bz_data[ii],axis = 0))
                bz_embedings[ii] = em
            embedings[i-batch:i,...] = bz_embedings
    else:
        ll = i % batch
        bz_embedings = np.zeros([ll,512])
        if ll !=0:
            ll_in = i // batch
            bz_data = datasets[0:ll,...]
            for ii in range(len(bz_data)):
                em = model.get_embedding(np.expand_dims(bz_data[ii],axis = 0))
                bz_embedings[ii] = em
            embedings[ll_in*batch : ll_in*batch+ll: ,...] = bz_embedings
    del bins
    gc.collect()
    return embedings, issame_list

def pb_model_test(model_dir):
    model = FACENET_PB(model_dir)
    img_size = 112
    valid_dir = VALID_DIR
    datadirs = os.listdir(valid_dir)
    embedding_size = 512

    result = pt.PrettyTable( ["data set", "AUC", "ACC" ,"VR @ FAR ", "dist max", "dist min"])
    val_all = []
    auc_all = []
    acc_all = []
    dist_all = []
    for i in range(len(datadirs)) :
        path = valid_dir + "/" + datadirs[i]
        embeddings, issamelab = predict_evaluate_data_PB(path, [img_size,img_size], model, embedding_size)

        embeddings1 = embeddings[0::2]
        embeddings2 = embeddings[1::2]
        diff = np.subtract(embeddings1, embeddings2)
        dist = np.sum(np.square(diff), 1)
        del embeddings1,embeddings2
        fpr, tpr,ths = roc_curve(np.asarray(issamelab).astype('int'),dist, pos_label=0 )
        auc_score = auc(fpr, tpr)
        if 0:
            plt.figure()
            lw = 2
            plt.plot(fpr, tpr, color='darkorange',
                              lw=lw, label='ROC curve (area = %0.2f)' % auc_score)
            plt.plot([0, 1], [0, 1], color='navy', lw=lw, linestyle='--')
            plt.xlim([0.0, 1.0])
            plt.ylim([0.0, 1.05])
            plt.xlabel('False Positive Rate')
            plt.ylabel('True Positive Rate')
            plt.title('Receiver operating characteristic example')
            plt.legend(loc="lower right")
            #  plt.show()
            plt.savefig("roc.png")

        #  print('-----10 folds------')
        tpr, fpr, accuracy, val, val_std, far = evaluate(embeddings, issamelab)
        del embeddings,issamelab
        gc.collect()
        result.add_row([datadirs[i] ,
                         round(auc_score,2),
                         "%1.3f+-%1.3f" % (np.mean(accuracy), np.std(accuracy)),
                         "%2.5f+-%2.5f @ FAR=%2.5f" % (val, val_std, far),
                         round(np.max(dist),2),
                         round(np.min(dist),2),
                             ])
        val_all.append(val)
        acc_all.append(accuracy)
        auc_all.append(auc_score)
        dist_all.append([np.max(dist),np.min(dist)])

    print(result)

pb_model_test(PB_DIR)


#!/usr/bin/env python
#

interactive = False

import sys, os
datadir = os.path.join(os.path.dirname(__file__), "../data")
sys.path.insert(0, datadir)
here = os.path.dirname(__file__)

import unittest, warnings
import numpy as np, histogram.hdf as hh, histogram as H
from multiphonon.backward import sqe2dos
from dos import loadDOS

class TestCase(unittest.TestCase):

    def test2a(self):
        "sqe2dos: V exp"
        iqehist = hh.load(os.path.join(datadir, 'XYZ2-iqe-Ei_20.h5'))
        initdos = hh.load(os.path.join(datadir, 'XYZ2-initdos-Ei_80.h5'))
        iterdos = sqe2dos.sqe2dos(
            iqehist, T=300,
            Ecutoff=17.1, elastic_E_cutoff=(-3., 1.1), M=79.452,
            C_ms=.05, Ei=20., workdir='work-XYZ2',
            initdos = initdos)
        list(iterdos)
        if interactive:
            dos = hh.load('work-XYZ2/final-dos.h5')
            pylab.errorbar(dos.E, dos.I, dos.E2**.5, label='final')
            pylab.legend()
            pylab.show()
        return

    pass  # end of TestCase

if __name__ == "__main__":
    interactive = True
    import pylab
    unittest.main()

# End of file


import os
import numpy as np
import torch
from torch.optim import RMSprop
from torch.nn.utils import clip_grad_norm_

from .ppo import PPOAgent
from wappo.network import CriticNetwork

class WAPPOAgent(PPOAgent):

    def __init__(self, venv_source, venv_target, log_dir, device,
                 num_steps=10**6, lr_ppo=5e-4, gamma=0.999,
                 rollout_length=16, num_minibatches=8, epochs_ppo=3,
                 clip_range_ppo=0.2, coef_value=0.5, coef_ent=0.01,
                 lambd=0.95, max_grad_norm=0.5, num_initial_blocks=1,
                 lr_critic=5.e-4, epochs_critic=5, coef_conf=10.0,
                 clip_range_critic=0.01):
        super().__init__(
            venv_source, venv_target, log_dir, device, num_steps, lr_ppo,
            gamma, rollout_length, num_minibatches, epochs_ppo, clip_range_ppo,
            coef_value, coef_ent, lambd, max_grad_norm, num_initial_blocks)

        # Adversarial network.
        self.network_critic = CriticNetwork(
            feature_dim=self.network_ppo.feature_dim).to(device)

        # Optimizers.
        self.optim_critic = RMSprop(
            self.network_critic.parameters(), lr=lr_critic)

        self.epochs_critic = epochs_critic
        self.coef_conf = coef_conf
        self.clip_range_critic = clip_range_critic

    def update(self):
        loss_policies = []
        loss_values = []
        loss_critics = []
        loss_confs = []

        for samples in self.storage_source.iter(self.batch_size):
            states_target = self.storage_target.sample(self.batch_size)

            loss_policy, loss_value, loss_conf = \
                self.update_ppo(*samples, states_target)
            loss_critic = self.update_critic(samples[0], states_target)

            self.update_steps += 1
            loss_policies.append(loss_policy)
            loss_values.append(loss_value)
            loss_confs.append(loss_conf)
            loss_critics.append(loss_critic)

        self.writer.add_scalar(
            'loss/policy', np.mean(loss_policies), self.steps)
        self.writer.add_scalar(
            'loss/value', np.mean(loss_values), self.steps)
        self.writer.add_scalar(
            'loss/critic', np.mean(loss_critics), self.steps)
        self.writer.add_scalar(
            'loss/conf', np.mean(loss_confs), self.steps)

    def update_ppo(self, states_source, actions, values_old, value_targets,
                   log_probs_old, advantages, states_target):
        # Include confusion loss every 5 updates.
        if_update_conf = float(self.update_steps % self.epochs_critic == 0)

        values, log_probs, mean_entropy, source_features = \
            self.network_ppo.evaluate_actions(states_source, actions)

        # >>> Value >>> #
        values_clipped = values_old + (values - values_old).clamp(
            -self.clip_range_ppo, self.clip_range_ppo)
        loss_value1 = (values - value_targets).pow(2)
        loss_value2 = (values_clipped - value_targets).pow(2)
        loss_value = torch.max(loss_value1, loss_value2).mean()
        # >>> Value >>> #

        # >>> Policy >>> #
        ratio = torch.exp(log_probs - log_probs_old)
        loss_policy1 = -ratio * advantages
        loss_policy2 = -torch.clamp(
            ratio,
            1.0 - self.clip_range_ppo,
            1.0 + self.clip_range_ppo
        ) * advantages
        loss_policy = torch.max(loss_policy1, loss_policy2).mean()
        # >>> Policy >>> #

        # >>> Confusion >>> #
        target_features = \
            self.network_ppo.calculate_features(states_target)
        source_preds = self.network_critic(source_features)
        target_preds = self.network_critic(target_features)
        loss_conf = torch.mean(source_preds) - torch.mean(target_preds)
        # >>> Confusion >>> #

        # >>> Total >>> #
        loss = \
            loss_policy \
            - self.coef_ent * mean_entropy \
            + self.coef_value * loss_value \
            + if_update_conf * self.coef_conf * loss_conf
        # >>> Total >>> #

        self.optim_ppo.zero_grad()
        loss.backward()
        clip_grad_norm_(self.network_ppo.parameters(), self.max_grad_norm)
        self.optim_ppo.step()

        return loss_policy.detach().item(), loss_value.detach().item(), \
            loss_conf.detach().item()

    def update_critic(self, states_source, states_target):
        with torch.no_grad():
            source_features = \
                self.network_ppo.calculate_features(states_source)
            target_features = \
                self.network_ppo.calculate_features(states_target)

        source_preds = self.network_critic(source_features)
        target_preds = self.network_critic(target_features)

        loss_critic = - torch.mean(source_preds) + torch.mean(target_preds)

        self.optim_critic.zero_grad()
        loss_critic.backward()
        clip_grad_norm_(self.network_critic.parameters(), self.max_grad_norm)
        self.optim_critic.step()

        for p in self.network_critic.parameters():
            p.data.clamp_(-self.clip_range_critic, self.clip_range_critic)

        return loss_critic.detach().item()

    def save_models(self, save_dir):
        super().save_models(save_dir)
        torch.save(
            self.network_critic.state_dict(),
            os.path.join(save_dir, 'network_critic.pth'))

    def load_models(self, save_dir):
        super().load_models(save_dir)
        self.network_critic.load_state_dict(
            torch.load(os.path.join(save_dir, 'network_critic.pth')))


import time
import weakref
import logging

import numpy as np
import wgpu.backends.rs

from .. import Renderer
from ...linalg import Matrix4, Vector3
from ...objects import (
    id_provider,
    KeyboardEvent,
    RootEventHandler,
    PointerEvent,
    WheelEvent,
    WindowEvent,
    WorldObject,
)
from ...cameras import Camera
from ...resources import Buffer, Texture, TextureView
from ...utils import array_from_shadertype, Color

from . import _blender as blender_module
from ._flusher import RenderFlusher
from ._pipelinebuilder import ensure_pipeline
from ._update import update_buffer, update_texture, update_texture_view

logger = logging.getLogger("pygfx")

# Definition uniform struct with standard info related to transforms,
# provided to each shader as uniform at slot 0.
# todo: a combined transform would be nice too, for performance
# todo: same for ndc_to_world transform (combined inv transforms)
stdinfo_uniform_type = dict(
    cam_transform="4x4xf4",
    cam_transform_inv="4x4xf4",
    projection_transform="4x4xf4",
    projection_transform_inv="4x4xf4",
    physical_size="2xf4",
    logical_size="2xf4",
    flipped_winding="i4",  # A bool, really
)

def _get_sort_function(camera: Camera):
    """Given a scene object, get a function to sort wobject-tuples"""

    def sort_func(wobject_tuple: WorldObject):
        wobject = wobject_tuple[0]
        z = (
            Vector3()
            .set_from_matrix_position(wobject.matrix_world)
            .apply_matrix4(proj_screen_matrix)
            .z
        )
        return wobject.render_order, z

    proj_screen_matrix = Matrix4().multiply_matrices(
        camera.projection_matrix, camera.matrix_world_inverse
    )

    return sort_func

class SharedData:
    """An object to store global data to share between multiple wgpu renderers.

    Since renderers don't render simultaneously, they can share certain
    resources. This safes memory, but more importantly, resources that
    get used in wobject pipelines should be shared to avoid having to
    constantly recompose the pipelines of wobjects that are rendered by
    multiple renderers.
    """

    def __init__(self, canvas):

        # Create adapter and device objects - there should be just one per canvas.
        # Having a global device provides the benefit that we can draw any object
        # anywhere.
        # We could pass the canvas to request_adapter(), so we get an adapter that is
        # at least compatible with the first canvas that a renderer is create for.
        # However, passing the object has been shown to prevent the creation of
        # a canvas (on Linux + wx), so, we never pass it for now.
        self.adapter = wgpu.request_adapter(
            canvas=None, power_preference="high-performance"
        )
        self.device = self.adapter.request_device(
            required_features=[], required_limits={}
        )

        # Create a uniform buffer for std info
        self.stdinfo_buffer = Buffer(array_from_shadertype(stdinfo_uniform_type))
        self.stdinfo_buffer._wgpu_usage |= wgpu.BufferUsage.UNIFORM

        # A cache for shader objects
        self.shader_cache = {}

class WgpuRenderer(RootEventHandler, Renderer):
    """Object used to render scenes using wgpu.

    The purpose of a renderer is to render (i.e. draw) a scene to a
    canvas or texture. It also provides picking, defines the
    anti-aliasing parameters, and any post processing effects.

    A renderer is directly associated with its target and can only render
    to that target. Different renderers can render to the same target though.

    It provides a ``.render()`` method that can be called one or more
    times to render scenes. This creates a visual representation that
    is stored internally, and is finally rendered into its render target
    (the canvas or texture).
                                  __________
                                 | blender  |
        [scenes] -- render() --> |  state   | -- flush() --> [target]
                                 |__________|

    The internal representation is managed by the blender object. The
    internal render textures are typically at a higher resolution to
    reduce aliasing (SSAA). The blender has auxilary buffers such as a
    depth buffer, pick buffer, and buffers for transparent fragments.
    Depending on the blend mode, a single render call may consist of
    multiple passes (to deal with semi-transparent fragments).

    The flush-step resolves the internal representation into the target
    texture or canvas, averaging neighbouring fragments for anti-aliasing.

    Parameters:
        target (WgpuCanvas or Texture): The target to render to, and what
            determines the size of the render buffer.
        pixel_ratio (float, optional): How large the physical size of the render
            buffer is in relation to the target's physical size, for antialiasing.
            See the corresponding property for details.
        show_fps (bool): Whether to display the frames per second. Beware that
            depending on the GUI toolkit, the canvas may impose a frame rate limit.
    """

    _shared = None

    _wobject_pipelines_collection = weakref.WeakValueDictionary()

    def __init__(
        self,
        target,
        *args,
        pixel_ratio=None,
        show_fps=False,
        blend_mode="default",
        sort_objects=False,
        enable_events=True,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)

        # Check and normalize inputs
        if not isinstance(target, (Texture, TextureView, wgpu.gui.WgpuCanvasBase)):
            raise TypeError(
                f"Render target must be a canvas or texture (view), not a {target.__class__.__name__}"
            )
        self._target = target

        # Process other inputs
        self.pixel_ratio = pixel_ratio
        self._show_fps = bool(show_fps)

        # Make sure we have a shared object (the first renderer create it)
        canvas = target if isinstance(target, wgpu.gui.WgpuCanvasBase) else None
        if WgpuRenderer._shared is None:
            WgpuRenderer._shared = SharedData(canvas)

        # Init counter to auto-clear
        self._renders_since_last_flush = 0

        # Get target format
        if isinstance(target, wgpu.gui.WgpuCanvasBase):
            self._canvas_context = self._target.get_context()
            self._target_tex_format = self._canvas_context.get_preferred_format(
                self._shared.adapter
            )
            # Also configure the canvas
            self._canvas_context.configure(
                device=self._shared.device,
                format=self._target_tex_format,
                usage=wgpu.TextureUsage.RENDER_ATTACHMENT,
            )
        else:
            self._target_tex_format = self._target.format
            # Also enable the texture for render and display usage
            self._target._wgpu_usage |= wgpu.TextureUsage.RENDER_ATTACHMENT
            self._target._wgpu_usage |= wgpu.TextureUsage.TEXTURE_BINDING

        # Prepare render targets.
        self.blend_mode = blend_mode
        self.sort_objects = sort_objects

        # Prepare object that performs the final render step into a texture
        self._flusher = RenderFlusher(self._shared.device)

        # Initialize a small buffer to read pixel info into
        # Make it 256 bytes just in case (for bytes_per_row)
        self._pixel_info_buffer = self._shared.device.create_buffer(
            size=16,
            usage=wgpu.BufferUsage.COPY_DST | wgpu.BufferUsage.MAP_READ,
        )

        if enable_events:
            self.enable_events()

    @property
    def device(self):
        """A reference to the used wgpu device."""
        return self._shared.device

    @property
    def target(self):
        """The render target. Can be a canvas, texture or texture view."""
        return self._target

    @property
    def pixel_ratio(self):
        """The ratio between the number of internal pixels versus the logical pixels on the canvas.

        This can be used to configure the size of the render texture
        relative to the canvas' logical size. By default (value is None) the
        used pixel ratio follows the screens pixel ratio on high-res
        displays, and is 2 otherwise.

        If the used pixel ratio causes the render texture to be larger
        than the physical size of the canvas, SSAA is applied, resulting
        in a smoother final image with less jagged edges. Alternatively,
        this value can be set to e.g. 0.5 to lower* the resolution (e.g.
        for performance during interaction).
        """
        return self._pixel_ratio

    @pixel_ratio.setter
    def pixel_ratio(self, value):
        if value is None:
            self._pixel_ratio = None
        elif isinstance(value, (int, float)):
            self._pixel_ratio = None if value <= 0 else float(value)
        else:
            raise TypeError(
                f"Rendered.pixel_ratio expected None or number, not {value}"
            )

    @property
    def rect(self):
        """The rectangular viewport for the renderer area."""
        return (0, 0) + self.logical_size

    @property
    def logical_size(self):
        """The size of the render target in logical pixels."""
        target = self._target
        if isinstance(target, wgpu.gui.WgpuCanvasBase):
            return target.get_logical_size()
        elif isinstance(target, Texture):
            return target.size[:2]  # assuming pixel-ratio 1
        elif isinstance(target, TextureView):
            return target.texture.size[:2]  # assuming pixel-ratio 1
        else:
            raise TypeError(f"Unexpected render target {target.__class__.__name__}")

    @property
    def physical_size(self):
        """The physical size of the internal render texture."""

        # Get physical size of the target
        target = self._target
        if isinstance(target, wgpu.gui.WgpuCanvasBase):
            target_psize = target.get_physical_size()
        elif isinstance(target, Texture):
            target_psize = target.size[:2]
        elif isinstance(target, TextureView):
            target_psize = target.texture.size[:2]
        else:
            raise TypeError(f"Unexpected render target {target.__class__.__name__}")

        # Determine the pixel ratio of the render texture
        if self._pixel_ratio:
            pixel_ratio = self._pixel_ratio
        else:
            target_lsize = self.logical_size
            pixel_ratio = target_psize[0] / target_lsize[0]
            if pixel_ratio <= 1:
                pixel_ratio = 2.0  # use 2 on non-hidpi displays

        # Determine the physical size of the internal render textures
        return tuple(max(1, int(pixel_ratio * x)) for x in target_lsize)

    @property
    def blend_mode(self):
        """The method for handling transparency:

        * "default" or None: Select the default: currently this is "ordered2".
        * "opaque": single-pass approach that consider every fragment opaque.
        * "ordered1": single-pass approach that blends fragments (using alpha blending).
          Can only produce correct results if fragments are drawn from back to front.
        * "ordered2": two-pass approach that first processes all opaque fragments and then
          blends transparent fragments (using alpha blending) with depth-write disabled. The
          visual results are usually better than ordered1, but still depend on the drawing order.
        * "weighted": two-pass approach that for order independent transparency,
          using alpha weights.
        * "weighted_depth": two-pass approach for order independent transparency,
          with weights based on alpha and depth (McGuire 2013). Note that the depth
          range affects the (quality of the) visual result.
        * "weighted_plus": three-pass approach for order independent transparency,
          in wich the front-most transparent layer is rendered correctly, while
          transparent layers behind it are blended using alpha weights.
        """
        return self._blend_mode

    @blend_mode.setter
    def blend_mode(self, value):
        # Massage and check the input
        if value is None:
            value = "default"
        value = value.lower()
        if value == "default":
            value = "ordered2"
        # Map string input to a class
        m = {
            "opaque": blender_module.OpaqueFragmentBlender,
            "ordered1": blender_module.Ordered1FragmentBlender,
            "ordered2": blender_module.Ordered2FragmentBlender,
            "weighted": blender_module.WeightedFragmentBlender,
            "weighted_depth": blender_module.WeightedDepthFragmentBlender,
            "weighted_plus": blender_module.WeightedPlusFragmentBlender,
        }
        if value not in m:
            raise ValueError(
                f"Unknown blend_mode '{value}', use any of {set(m.keys())}"
            )
        # Set blender object
        self._blend_mode = value
        self._blender = m[value]()
        # If the blend mode has changed, we may need a new _wobject_pipelines
        self._set_wobject_pipelines()
        # If our target is a canvas, request a new draw
        if isinstance(self._target, wgpu.gui.WgpuCanvasBase):
            self._target.request_draw()

    @property
    def sort_objects(self):
        """Whether to sort world objects before rendering. Default False.

        * ``True``: the render order is defined by 1) the object's ``render_order``
          property; 2) the object's distance to the camera; 3) the position object
          in the scene graph (based on a depth-first search).
        * ``False``: don't sort, the render order is defined by the scene graph alone.
        """
        return self._sort_objects

    @sort_objects.setter
    def sort_objects(self, value):
        self._sort_objects = bool(value)

    def _set_wobject_pipelines(self):
        # Each WorldObject has associated with it a wobject_pipeline:
        # a dict that contains the wgpu pipeline objects. This
        # wobject_pipeline is also associated with the blend_mode,
        # because the blend mode affects the pipelines.
        #
        # Each renderer has ._wobject_pipelines, a dict that maps
        # wobject -> wobject_pipeline. This dict is a WeakKeyDictionary -
        # when the wobject is destroyed, the associated pipeline is
        # collected as well.
        #
        # Renderers with the same blend mode can safely share these
        # wobject_pipeline dicts. Therefore, we make use of a global
        # collection. Since this global collection is a
        # WeakValueDictionary, if all renderes stop using a certain
        # blend mode, the associated pipelines are removed as well.
        #
        # In a diagram:
        #
        # _wobject_pipelines_collection -> _wobject_pipelines -> wobject_pipeline
        #        global                         renderer              wobject
        #   WeakValueDictionary              WeakKeyDictionary         dict

        # Below we set this renderer's _wobject_pipelines. Note that if the
        # blending has changed, we automatically invalidate all "our" pipelines.
        self._wobject_pipelines = WgpuRenderer._wobject_pipelines_collection.setdefault(
            self.blend_mode, weakref.WeakKeyDictionary()
        )

    def render(
        self,
        scene: WorldObject,
        camera: Camera,
        *,
        rect=None,
        viewport=None,
        clear_color=None,
        flush=True,
    ):
        """Render a scene with the specified camera as the viewpoint.

        Parameters:
            scene (WorldObject): The scene to render, a WorldObject that
                optionally has child objects.
            camera (Camera): The camera object to use, which defines the
                viewpoint and view transform.
            rect (tuple, optional): The rectangular region to draw into,
                expressed in logical pixels, a.k.a. the viewport.
            clear_color (bool, optional): Whether to clear the color buffer
                before rendering. By default this is True on the first
                call to ``render()`` after a flush, and False otherwise.
            flush (bool, optional): Whether to flush the rendered result into
                the target (texture or canvas). Default True.
        """
        device = self.device

        now = time.perf_counter()  # noqa
        if self._show_fps:
            if not hasattr(self, "_fps"):
                self._fps = now, now, 1
            elif now > self._fps[0] + 1:
                print(f"FPS: {self._fps[2]/(now - self._fps[0]):0.1f}")
                self._fps = now, now, 1
            else:
                self._fps = self._fps[0], now, self._fps[2] + 1

        # Define whether to clear color.
        if clear_color is None:
            clear_color = self._renders_since_last_flush == 0
        clear_color = bool(clear_color)
        self._renders_since_last_flush += 1

        # We always clear the depth, because each render() should be "self-contained".
        # Any use-cases where you normally would control depth-clearing should
        # be covered by the blender. Also, this way the blender can better re-use internal
        # buffers. The only rule is that the color buffer behaves correctly on multiple renders.

        # todo: also note that the fragment shader is (should be) optional
        #      (e.g. depth only passes like shadow mapping or z prepass)

        # Get size for the render textures
        logical_size = self.logical_size
        physical_size = self.physical_size
        if not all(i > 0 for i in logical_size):
            return
        pixel_ratio = physical_size[1] / logical_size[1]

        # Update the render targets
        self._blender.ensure_target_size(device, physical_size)

        # todo: 11-04-2022 - remove viewport arg and this backwards compat logic in the next release
        if viewport is not None:
            if not hasattr(self, "_warned_for_viewport_arg"):
                self._warned_for_viewport_arg = True
                logger.warning(
                    "render(viewport=xx) is deprecated, use rect=xx instead."
                )
            if not rect:
                rect = viewport

        # Get viewport in physical pixels
        if not rect:
            scene_lsize = logical_size
            scene_psize = physical_size
            physical_viewport = 0, 0, physical_size[0], physical_size[1], 0, 1
        elif len(rect) == 4:
            scene_lsize = rect[2], rect[3]
            physical_viewport = [int(i * pixel_ratio + 0.4999) for i in rect]
            physical_viewport = tuple(physical_viewport) + (0, 1)
            scene_psize = physical_viewport[2], physical_viewport[3]
        else:
            raise ValueError(
                "The viewport rect must be None or 4 elements (x, y, w, h)."
            )

        # Ensure that matrices are up-to-date
        scene.update_matrix_world()
        camera.set_view_size(*scene_lsize)
        camera.update_matrix_world()  # camera may not be a member of the scene
        camera.update_projection_matrix()

        # Update stdinfo uniform buffer object that we'll use during this render call
        self._update_stdinfo_buffer(camera, scene_psize, scene_lsize)

        # Get the list of objects to render, as they appear in the scene graph
        wobject_list = []
        scene.traverse(wobject_list.append, True)

        # Ensure each wobject has pipeline info, and filter objects that we cannot render
        wobject_tuples = []
        any_has_changed = False
        for wobject in wobject_list:
            if not wobject.material:
                continue
            wobject_pipeline, has_changed = ensure_pipeline(self, wobject)
            if wobject_pipeline:
                any_has_changed |= has_changed
                wobject_tuples.append((wobject, wobject_pipeline))

        # Command buffers cannot be reused. If we want some sort of re-use we should
        # look into render bundles. See https://github.com/gfx-rs/wgpu-native/issues/154
        # If we do get this to work, we should trigger a new recording
        # when the wobject's children, visibile, render_order, or render_pass changes.

        # Sort objects
        if self.sort_objects:
            sort_func = _get_sort_function(camera)
            wobject_tuples.sort(key=sort_func)

        # Record the rendering of all world objects, or re-use previous recording
        command_buffers = []
        command_buffers += self._render_recording(
            wobject_tuples, physical_viewport, clear_color
        )
        command_buffers += self._blender.perform_combine_pass(self._shared.device)
        command_buffers

        # Collect commands and submit
        device.queue.submit(command_buffers)

        if flush:
            self.flush()

    def flush(self):
        """Render the result into the target texture view. This method is
        called automatically unless you use ``.render(..., flush=False)``.
        """

        # Note: we could, in theory, allow specifying a custom target here.

        if isinstance(self._target, wgpu.gui.WgpuCanvasBase):
            raw_texture_view = self._canvas_context.get_current_texture()
        else:
            if isinstance(self._target, Texture):
                texture_view = self._target.get_view()
            elif isinstance(self._target, TextureView):
                texture_view = self._target
            update_texture(self._shared.device, texture_view.texture)
            update_texture_view(self._shared.device, texture_view)
            raw_texture_view = texture_view._wgpu_texture_view[1]

        # Reset counter (so we can auto-clear the first next draw)
        self._renders_since_last_flush = 0

        command_buffers = self._flusher.render(
            self._blender.color_view,
            None,
            raw_texture_view,
            self._target_tex_format,
        )
        self.device.queue.submit(command_buffers)

    def _render_recording(
        self,
        wobject_tuples,
        physical_viewport,
        clear_color,
    ):

        # You might think that this is slow for large number of world
        # object. But it is actually pretty good. It does iterate over
        # all world objects, and over stuff in each object. But that's
        # it, really.
        # todo: we may be able to speed this up with render bundles though

        command_encoder = self.device.create_command_encoder()
        blender = self._blender

        # ----- compute pipelines

        compute_pass = command_encoder.begin_compute_pass()

        for wobject, wobject_pipeline in wobject_tuples:
            for pinfo in wobject_pipeline.get("compute_pipelines", ()):
                compute_pass.set_pipeline(pinfo["pipeline"])
                for bind_group_id, bind_group in enumerate(pinfo["bind_groups"]):
                    compute_pass.set_bind_group(
                        bind_group_id, bind_group, [], 0, 999999
                    )
                compute_pass.dispatch_workgroups(*pinfo["index_args"])

        compute_pass.end()

        # ----- render pipelines

        for pass_index in range(blender.get_pass_count()):

            color_attachments = blender.get_color_attachments(pass_index, clear_color)
            depth_attachment = blender.get_depth_attachment(pass_index)
            render_mask = blender.passes[pass_index].render_mask
            if not color_attachments:
                continue

            render_pass = command_encoder.begin_render_pass(
                color_attachments=color_attachments,
                depth_stencil_attachment={
                    **depth_attachment,
                    "stencil_load_op": wgpu.LoadOp.load,
                    "stencil_store_op": wgpu.StoreOp.store,
                },
                occlusion_query_set=None,
            )
            render_pass.set_viewport(*physical_viewport)

            for wobject, wobject_pipeline in wobject_tuples:
                if not (render_mask & wobject_pipeline["render_mask"]):
                    continue
                for pinfo in wobject_pipeline["render_pipelines"]:
                    render_pass.set_pipeline(pinfo["pipelines"][pass_index])
                    for slot, vbuffer in pinfo["vertex_buffers"].items():
                        render_pass.set_vertex_buffer(
                            slot,
                            vbuffer._wgpu_buffer[1],
                            vbuffer.vertex_byte_range[0],
                            vbuffer.vertex_byte_range[1],
                        )
                    for bind_group_id, bind_group in enumerate(pinfo["bind_groups"]):
                        render_pass.set_bind_group(bind_group_id, bind_group, [], 0, 99)
                    # Draw with or without index buffer
                    if pinfo["index_buffer"] is not None:
                        ibuffer = pinfo["index_buffer"]
                        render_pass.set_index_buffer(ibuffer, 0, ibuffer.size)
                        render_pass.draw_indexed(*pinfo["index_args"])
                    else:
                        render_pass.draw(*pinfo["index_args"])

            render_pass.end()

        return [command_encoder.finish()]

    def _update_stdinfo_buffer(self, camera, physical_size, logical_size):
        # Update the stdinfo buffer's data
        stdinfo_data = self._shared.stdinfo_buffer.data
        stdinfo_data["cam_transform"].flat = camera.matrix_world_inverse.elements
        stdinfo_data["cam_transform_inv"].flat = camera.matrix_world.elements
        stdinfo_data["projection_transform"].flat = camera.projection_matrix.elements
        stdinfo_data[
            "projection_transform_inv"
        ].flat = camera.projection_matrix_inverse.elements
        # stdinfo_data["ndc_to_world"].flat = np.linalg.inv(stdinfo_data["cam_transform"] @ stdinfo_data["projection_transform"])
        stdinfo_data["physical_size"] = physical_size
        stdinfo_data["logical_size"] = logical_size
        stdinfo_data["flipped_winding"] = camera.flips_winding
        # Upload to GPU
        self._shared.stdinfo_buffer.update_range(0, 1)
        update_buffer(self._shared.device, self._shared.stdinfo_buffer)

    # Picking

    def get_pick_info(self, pos):
        """Get information about the given window location. The given
        pos is a 2D point in logical pixels (with the origin at the
        top-left). Returns a dict with fields:

        * "ndc": The position in normalized device coordinates, the 3d element
            being the depth (0..1). Can be translated to the position
            in world coordinates using the camera transforms.
        * "rgba": The value in the color buffer. All zero's when rendering
          directly to the screen (bypassing post-processing).
        * "world_object": the object at that location (provided that
          the object supports picking).
        * Additional pick info may be available, depending on the type of
          object and its material. See the world-object classes for details.
        """

        # Make pos 0..1, so we can scale it to the render texture
        logical_size = self.logical_size
        float_pos = pos[0] / logical_size[0], pos[1] / logical_size[1]

        # Prevent out of range and picking before first draw
        out_of_range = not (0 <= float_pos[0] <= 1 and 0 <= float_pos[1] <= 1)
        blender_zero_size = self._blender.size == (0, 0)
        if out_of_range or blender_zero_size:
            return {"rgba": Color(0, 0, 0, 0), "world_object": None}

        # Sample
        encoder = self.device.create_command_encoder()
        self._copy_pixel(encoder, self._blender.color_tex, float_pos, 0)
        self._copy_pixel(encoder, self._blender.pick_tex, float_pos, 8)
        queue = self.device.queue
        queue.submit([encoder.finish()])

        # Collect data from the buffer
        data = self._pixel_info_buffer.map_read()
        color = Color(x / 255 for x in tuple(data[0:4].cast("B")))
        pick_value = tuple(data[8:16].cast("Q"))[0]
        wobject_id = pick_value & 1048575  # 2**20-1
        wobject = id_provider.get_object_from_id(wobject_id)
        # Note: the position in world coordinates is not included because
        # it depends on the camera, but we don't "own" the camera.

        info = {
            "rgba": color,
            "world_object": wobject,
        }

        if wobject:
            pick_info = wobject._wgpu_get_pick_info(pick_value)
            info.update(pick_info)
        return info

    def _copy_pixel(self, encoder, render_texture, float_pos, buf_offset):

        # Map position to the texture index
        w, h, d = render_texture.size
        x = max(0, min(w - 1, int(float_pos[0] * w)))
        y = max(0, min(h - 1, int(float_pos[1] * h)))

        # Note: bytes_per_row must be a multiple of 256.
        encoder.copy_texture_to_buffer(
            {
                "texture": render_texture,
                "mip_level": 0,
                "origin": (x, y, 0),
            },
            {
                "buffer": self._pixel_info_buffer,
                "offset": buf_offset,
                "bytes_per_row": 256,  # render_texture.bytes_per_pixel,
                "rows_per_image": 1,
            },
            copy_size=(1, 1, 1),
        )

    def snapshot(self):
        """Create a snapshot of the currently rendered image."""

        # Prepare
        device = self._shared.device
        texture = self._blender.color_tex
        size = texture.size
        bytes_per_pixel = 4

        # Note, with queue.read_texture the bytes_per_row limitation does not apply.
        data = device.queue.read_texture(
            {
                "texture": texture,
                "mip_level": 0,
                "origin": (0, 0, 0),
            },
            {
                "offset": 0,
                "bytes_per_row": bytes_per_pixel * size[0],
                "rows_per_image": size[1],
            },
            size,
        )

        return np.frombuffer(data, np.uint8).reshape(size[1], size[0], 4)

    def request_draw(self, draw_function=None):
        """Forwards a request_draw call to the target canvas. If the renderer's
        target is not a canvas (e.g. a texture) this function does
        nothing.
        """
        request_draw = getattr(self.target, "request_draw", None)
        if request_draw:
            request_draw(draw_function)

    def enable_events(self):
        """Add event handlers for a specific list of events that are generated
        by the canvas. The handler is the ``convert_event`` method in order to
        convert the Wgpu event dicts into Pygfx event objects."""

        # Check for ``add_event_handler`` attribute. Silently 'fail' as it is
        # perfectly fine to pass a texture as a target to the renderer
        if hasattr(self._target, "add_event_handler"):
            self._target.add_event_handler(self.convert_event, *EVENTS_TO_CONVERT)

    def disable_events(self):
        """Remove the event handlers from the canvas."""
        if hasattr(self._target, "remove_event_handler"):
            self._target.remove_event_handler(self.convert_event, *EVENTS_TO_CONVERT)

    def convert_event(self, event: dict):
        """Converts Wgpu event (dict following jupyter_rfb spec) to Pygfx Event object,
        adds picking info and then dispatches the event in the Pygfx event system.
        """
        event_type = event["event_type"]
        target = None
        if "x" in event and "y" in event:
            info = self.get_pick_info((event["x"], event["y"]))
            target = info["world_object"]
            event["pick_info"] = info

        ev = EVENT_TYPE_MAP[event_type](
            type=event_type, **event, target=target, root=self
        )
        self.dispatch_event(ev)

EVENT_TYPE_MAP = {
    "resize": WindowEvent,
    "close": WindowEvent,
    "pointer_down": PointerEvent,
    "pointer_up": PointerEvent,
    "pointer_move": PointerEvent,
    "double_click": PointerEvent,
    "wheel": WheelEvent,
    "key_down": KeyboardEvent,
    "key_up": KeyboardEvent,
}

EVENTS_TO_CONVERT = (
    "key_down",
    "key_up",
    "pointer_down",
    "pointer_move",
    "pointer_up",
    "wheel",
    "close",
    "resize",
)


# -*- coding: utf-8 -*-
# Fielname = dmu380_offline_sim.py

"""
IMU fusion.
Created on 2018-03-15
@author: dongxiaoguang
"""

# import
import os
import struct
import platform
import math
import numpy as np
from ctypes import *
from ctypes import wintypes

# globals
VERSION = '1.0'

R2D = 180.0 / math.pi

class SIM_COMFIG(Structure):
    '''
    read config params from file, and put params in this structure
    '''
    _fields_ = [("_inputDataRate", c_bool),
                ("inputDataRate", c_uint8),
                ("_outputDataRate", c_bool),
                ("outputDataRate", c_uint8),
                ("_hasMags", c_bool),
                ("hasMags", c_bool),
                ("_useMags", c_bool),
                ("useMags", c_bool),
                ("_hasGps", c_bool),
                ("hasGps", c_bool),
                ("_useGps", c_bool),
                ("useGps", c_bool),
                ("_freeIntegrate", c_bool),
                ("freeIntegrate", c_bool),
                ("_dynamicMotion", c_bool),
                ("dynamicMotion", c_bool),
                ("_stationaryLockYaw", c_bool),
                ("stationaryLockYaw", c_bool),
                ("_turnSwitchThreshold", c_bool),
                ("turnSwitchThreshold", c_float),
                ("_hardIron_X", c_bool),
                ("hardIron_X", c_float),
                ("_hardIron_Y", c_bool),
                ("hardIron_Y", c_float),
                ("_softIronScaleRatio", c_bool),
                ("softIronScaleRatio", c_float),
                ("_softIronAngle", c_bool),
                ("softIronAngle", c_float),
                ("_headingTrackOffset", c_bool),
                ("headingTrackOffset", c_ushort),
                ("_accelLPFType", c_bool),
                ("accelLPFTypeStr", c_char * 64),
                ("_accelSwitch", c_bool),
                ("accelSwitch", c_float),
                ("_linAccelSwitchDelay", c_bool),
                ("linAccelSwitchDelay", c_float),
                ("_Free_Integration_Cntr", c_bool),
                ("Free_Integration_Cntr", c_float),
                ("_Stabilize_System", c_bool),
                ("Stabilize_System", c_float),
                ("_Initialize_Attitude", c_bool),
                ("Initialize_Attitude", c_float),
                ("_High_Gain_AHRS", c_bool),
                ("High_Gain_AHRS", c_float),
                ("_Low_Gain_AHRS", c_bool),
                ("Low_Gain_AHRS", c_float),
                ("_maxGpsDropTime", c_bool),
                ("maxGpsDropTime", c_int32),
                ("_maxReliableDRTime", c_bool),
                ("maxReliableDRTime", c_int32),
                ("_procCovarMult_rateBias", c_bool),
                ("procCovarMult_rateBias", c_float),
                ("_procCovarMult_attitude", c_bool),
                ("procCovarMult_attitude", c_float),
                ("_measCovarMult_roll", c_bool),
                ("measCovarMult_roll", c_float),
                ("_measCovarMult_pitch", c_bool),
                ("measCovarMult_pitch", c_float)
                ]

class GPS_DATA(Structure):
    '''
    Input GPS data structure
    '''
    _fields_ = [("gpsFixType", c_uint8),
                ("gpsUpdate", c_uint8),
                ("numSatellites", c_uint8),
                ("itow", c_uint),
                ("latitude", c_double),
                ("longitude", c_double),
                ("vNed", c_double*3),
                ("trueCourse", c_double),
                ("rawGroundSpeed", c_double),
                ("altitude", c_double),
                ("GPSSecondFraction", c_double),
                ("GPSmonth", c_uint8),
                ("GPSday", c_uint8),
                ("GPSyear", c_uint8),
                ("GPSHour", c_char),
                ("GPSMinute", c_char),
                ("GPSSecond", c_char),
                ("GPSHorizAcc", c_float),
                ("GPSVertAcc", c_float),
                ("HDOP", c_float),
                ("geoidAboveEllipsoid", c_float),
                ]

class ODO_DATA(Structure):
    '''
    Input odometer data structure
    '''
    _fields_ = [("odoUpdate", c_uint8),
                ("v", c_float),
                ]

class EKF_STATE(Structure):
    '''
    Return EFK state in this structure
    '''
    _fields_ = [("timeStep", c_uint32),
                ("kfPosN", c_double*3),
                ("kfVelN", c_float*3),
                ("kfQuat", c_float*4),
                ("kfRateBias", c_float*3),
                ("kfAccelBias", c_float*3),
                ("kfCorrectedRateB", c_float*3),
                ("kfCorrectedAccelB", c_float*3),
                ("algoFilteredYawRate", c_float),
                ("kfEulerAngles", c_float*3),
                ("algoState", c_int),
                ("algoTurnSwitch", c_ushort),
                ("algoLinAccelSwitch", c_uint8),]

class DMU380Sim(object):
    '''
    A wrapper form DMU380 algorithm offline simulation.
    '''
    def __init__(self, config_file):
        '''
        Args:
            config_file: a configuration file
        '''
        self.ext = '.so'
        # platform
        if platform.system() == 'Windows':
            self.ext = '.dll'
            if struct.calcsize("P") == 8:
                self.ext = '-x64' + self.ext
            else:
                self.ext = '-x86' + self.ext
        else:
            raise OSError('Only support windows.')
        # algorithm description
        self.input = ['fs', 'gyro', 'accel', 'gps', 'gps_visibility', 'time', 'gps_time', 'odo']
        self.output = ['algo_time', 'pos', 'vel', 'att_euler', 'wb', 'ab']
        self.batch = True
        self.results = None
        # algorithm vars
        this_dir = os.path.dirname(__file__)
        self.config_lib = os.path.join(this_dir, 'dmu380_sim_lib/libsim_utilities' + self.ext)
        self.sim_lib = os.path.join(this_dir, 'dmu380_sim_lib/aceinna_ins' + self.ext)
        if not (os.path.exists(self.config_lib) and os.path.exists(self.sim_lib)):
            if not self.build_lib():
                raise OSError('Shared libs not found.')
        self.parse_config = cdll.LoadLibrary(self.config_lib)
        self.sim_engine = cdll.LoadLibrary(self.sim_lib)
        # initialize algorithm
        self.sim_config = SIM_COMFIG()
        self.parse_config.parseConfigFile(c_char_p(config_file.encode('utf-8')),\
                                          pointer(self.sim_config))
        self.sim_engine.SimInitialize(pointer(self.sim_config))
        # if mag required?
        if self.sim_config.hasMags and self.sim_config.useMags:
            self.input.append('mag')

    def run(self, set_of_input):
        '''
        main procedure of the algorithm
        Args:
            set_of_input is a tuple or list consistent with self.input
                fs: sample frequency, Hz
                gyro: numpy array of size (n,3), rad/s
                accel: numpy array of size (n,3), m/s/s
                mag: numpy array of size (n,3), Gauss
        '''
        # get input
        idx = 0
        fs = set_of_input[idx]
        idx += 1
        gyro = set_of_input[idx]
        idx += 1
        accel = set_of_input[idx]
        idx += 1
        gps = set_of_input[idx]
        idx += 1
        gps_visibility = set_of_input[idx]
        idx += 1
        time = set_of_input[idx]
        idx += 1
        gps_time = set_of_input[idx]
        idx += 1
        odometer = set_of_input[idx]
        idx += 1
        if 'mag' in self.input:
            mag = set_of_input[idx]
        n = accel.shape[0]
        # algo output
        time_step = np.zeros((n,))
        pos = np.zeros((n, 3))
        vel = np.zeros((n, 3))
        euler_angles = np.zeros((n, 3))
        rate_bias = np.zeros((n, 3))
        accel_bias = np.zeros((n, 3))
        # run
        g = GPS_DATA()
        odo = ODO_DATA()
        ekf_state = EKF_STATE()
        output_len = 0
        idx_gps = 0
        for i in range(0, n):
            a = np.zeros((3,))
            w = np.zeros((3,))
            m = np.zeros((3,))
            a = accel[i, :] / 9.80665
            w = gyro[i, :]
            if 'mag' in self.input:
                m = mag[i, :] / 100.0
            aptr = a.ctypes.data_as(POINTER(c_double))
            wptr = w.ctypes.data_as(POINTER(c_double))
            mptr = m.ctypes.data_as(POINTER(c_double))
            # fill in GPS data
            g.gpsUpdate = 0
            if gps_time[idx_gps] == time[i]:
                # there is a GPS update
                g.gpsUpdate = 1
                g.gpsFixType = int(gps_visibility[idx_gps])
                g.latitude = gps[idx_gps, 0] * R2D
                g.longitude = gps[idx_gps, 1] * R2D
                g.altitude = gps[idx_gps, 2]
                g.vNed[0] = gps[idx_gps, 3]
                g.vNed[1] = gps[idx_gps, 4]
                g.vNed[2] = gps[idx_gps, 5]
                g.rawGroundSpeed = math.sqrt(g.vNed[0]*g.vNed[0] + g.vNed[1]*g.vNed[1])
                g.trueCourse = math.atan2(g.vNed[1], g.vNed[0]) * R2D
                # print(g.trueCourse, g.vNed[1], g.vNed[0])
                g.HDOP = 1.0
                g.GPSHorizAcc = g.HDOP * 3.0
                g.GPSVertAcc = g.GPSHorizAcc * 1.5
                g.itow = int(gps_time[idx_gps] * 1000)
                idx_gps += 1
                if idx_gps == gps_time.shape[0]:
                    idx_gps = gps_time.shape[0] - 1
            # fill in odo data
            odo.odoUpdate = 1
            odo.v = odometer[i]
            self.sim_engine.SimRun(aptr, wptr, mptr, pointer(g), pointer(odo))
            # get output
            self.sim_engine.GetEKF_STATES(pointer(ekf_state))
            # time_step[output_len] = ekf_state.timeStep / fs
            time_step[output_len] = i / fs
            # Euler angles order is [roll pitch yaw] in the algo
            # We use yaw [pitch roll yaw] order in the simulation
            pos[output_len, 0] = ekf_state.kfPosN[0] / R2D
            pos[output_len, 1] = ekf_state.kfPosN[1] / R2D
            pos[output_len, 2] = ekf_state.kfPosN[2]
            vel[output_len, 0] = ekf_state.kfVelN[0]
            vel[output_len, 1] = ekf_state.kfVelN[1]
            vel[output_len, 2] = ekf_state.kfVelN[2]
            euler_angles[output_len, 0] = ekf_state.kfEulerAngles[2]
            euler_angles[output_len, 1] = ekf_state.kfEulerAngles[1]
            euler_angles[output_len, 2] = ekf_state.kfEulerAngles[0]
            rate_bias[output_len, 0] = ekf_state.kfRateBias[0]
            rate_bias[output_len, 1] = ekf_state.kfRateBias[1]
            rate_bias[output_len, 2] = ekf_state.kfRateBias[2]
            accel_bias[output_len, 0] = ekf_state.kfAccelBias[0]
            accel_bias[output_len, 1] = ekf_state.kfAccelBias[1]
            accel_bias[output_len, 2] = ekf_state.kfAccelBias[2]
            output_len += 1
        # results
        self.results = [time_step[0:output_len],\
                        pos[0:output_len, :],\
                        vel[0:output_len, :],\
                        euler_angles[0:output_len, :],\
                        rate_bias[0:output_len, :],\
                        accel_bias[0:output_len, :]]

    def update(self, gyro, acc, mag=np.array([0.0, 0.0, 0.0])):
        '''
        Mahony filter for gyro, acc and mag.
        Args:
        Returns:
        '''
        pass

    def get_results(self):
        '''
        Returns:
            algorithm results as specified in self.output
                algorithm time step, sec
                Euler angles [yaw pitch roll], rad
        '''
        return self.results

    def reset(self):
        '''
        Reset the fusion process to uninitialized state.
        '''
        windll.kernel32.FreeLibrary.argtypes = [wintypes.HMODULE]
        windll.kernel32.FreeLibrary(self.sim_engine._handle)
        self.sim_engine = cdll.LoadLibrary(self.sim_lib)
        self.sim_engine.SimInitialize(pointer(self.sim_config))

    def build_lib(self, dst_dir=None, src_dir=None):
        '''
        Build shared lib
        Args:
            dst_dir: dir to put the built libs in.
            src_dir: dir containing the source code.
        Returns:
            True if success, False if error.
        '''
        if self.ext == '.dll':
            print("Automatic generation of .dll is not supported.")
            return False
        this_dir = os.path.dirname(__file__)
        # get dir containing the source code
        if src_dir is None:
            src_dir = os.path.join(this_dir, '//home//dong//c_projects//dmu380_sim_src//')
        if not os.path.exists(src_dir):
            print('Source code directory ' + src_dir + ' does not exist.')
            return False
        # get dir to put the libs in
        if dst_dir is None:
            dst_dir = os.path.join(this_dir, './/dmu380_sim_lib//')
        if not os.path.exists(dst_dir):
            os.mkdir(dst_dir)

        algo_lib = 'libdmu380_algo_sim.so'
        sim_utilities_lib = 'libsim_utilities.so'
        # get current workding dir
        cwd = os.getcwd()
        # create the cmake dir
        cmake_dir = src_dir + '//cmake//'
        if not os.path.exists(cmake_dir):
            os.mkdir(cmake_dir)
        else:
            os.system("rm -rf " + cmake_dir + "*")
        # call cmake and make to build the libs
        os.chdir(cmake_dir)
        ret = os.system("cmake ..")
        ret = os.system("make")
        algo_lib = cmake_dir + 'algo//' + algo_lib
        sim_utilities_lib = cmake_dir + 'SimUtilities//' + sim_utilities_lib
        if os.path.exists(algo_lib) and os.path.exists(sim_utilities_lib):
            os.system("mv " + algo_lib + " " + dst_dir)
            os.system("mv " + sim_utilities_lib + " " + dst_dir)

        # restore working dir
        os.chdir(cwd)
        return True