# 1. Update the config file.
$jsonFilePath = "C:\tmp\MetaXRSimulator\MetaXRSimulator\config\sim_core_configuration.json"
$jsonContent = Get-Content $jsonFilePath | Out-String
$jsonObject = $jsonContent | ConvertFrom-Json
if (-not ($jsonObject.PSObject.Properties.Name -contains "session_capture")) {
    $jsonObject | Add-Member -MemberType NoteProperty -Name "session_capture" -Value "session_capture"
}
$jsonObject.session_capture = @{
    "delay_start_ms" = 10000
    "exec_state" = "replay"
    "quit_buffer_ms" = 1000
    "quit_when_complete" = $true
    "record_path" = "C:\tmp\test_recordings\recording.vrs"
    "replay_path" = "C:\tmp\test_recordings\replay_target.vrs"
}
$jsonObject | ConvertTo-Json -Depth 100 | Set-Content $jsonFilePath

# 2. Launch a bedroom for Synthetic Environment
Start-Process -FilePath "C:/tmp/MetaXRSimulator/MetaXRSimulator\.synth_env_server\synth_env_server.exe" -ArgumentList "Bedroom" -PassThru

# 3. Activate Meta XR Simulator
c:/tmp/MetaXRSimulator/MetaXRSimulator/activate_simulator.ps1

# 4. Launch the game to replay
Start-Process -FilePath "C:\tmp\output\game.exe" -Wait

# 5. Deactivate Meta XR Simulator
C:\tmp\MetaXRSimulator\MetaXRSimulator\deactivate_simulator.ps1

# 6. Stop Synthetic Environment
Get-Process -Name "synth_env_server" | Stop-Process

# 1. Update the config file.
$jsonFilePath = "C:\tmp\MetaXRSimulator\MetaXRSimulator\config\sim_core_configuration.json"
$jsonContent = Get-Content $jsonFilePath | Out-String
$jsonObject = $jsonContent | ConvertFrom-Json

if (-not ($jsonObject.PSObject.Properties.Name -contains "session_capture")) {
    $jsonObject | Add-Member -MemberType NoteProperty -Name "session_capture" -Value "session_capture"
}
$jsonObject.session_capture = @{
    "delay_start_ms" = 10000
    "exec_state" = "replay"
    "quit_buffer_ms" = 1000
    "quit_when_complete" = $true
    "record_path" = "C:\tmp\test_recordings\recording.vrs"
    "replay_path" = "C:\tmp\test_recordings\replay_new.vrs"
}
if (-not ($jsonObject.PSObject.Properties.Name -contains "use_batch_mode")) {
    $jsonObject | Add-Member -MemberType NoteProperty -Name "use_batch_mode" -Value $true
}
$jsonObject.use_batch_mode = $true

if (-not ($jsonObject.PSObject.Properties.Name -contains "log_location")) {
    $jsonObject | Add-Member -MemberType NoteProperty -Name "log_location" -Value "log_location"
}
$jsonObject.log_location = "C:\tmp\test_recordings\MetaXRSimulator.log"

$jsonObject | ConvertTo-Json -Depth 10 | Set-Content $jsonFilePath

# 2. Activate Meta XR Simulator
C:/tmp/MetaXRSimulator/MetaXRSimulator/activate_simulator.ps1

# 3. Launch a Synthetic Environment (skip this step if it is already running)

Start-Process -FilePath "C:/tmp/MetaXRSimulator/MetaXRSimulator\.synth_env_server\synth_env_server.exe" -ArgumentList "Bedroom" -PassThru

Start-Sleep -Seconds 5

# 4. Launch the game
Start-Process -FilePath "C:\tmp\output\game.exe" -Wait

# 5. Deactivate Meta XR Simulator
C:\tmp\MetaXRSimulator\MetaXRSimulator\deactivate_simulator.ps1
# Stop Synthetic Environment
Get-Process -Name "synth_env_server" | Stop-Process

# 6. Compare result

# Prepare requirements.txt for windows:
findstr /v "vrs==1.0.4" C:\tmp\MetaXRSimulator\MetaXRSimulator\scripts\requirements.txt > C:\tmp\tmp.txt
move C:\tmp\tmp.txt C:\tmp\requirements_windows.txt

# Download pyvrs
cd c:/tmp
git clone https://github.com/facebookresearch/pyvrs.git
cd pyvrs
git submodule sync --recursive
git submodule update --init --recursive

# Build and Run Pyvrs via Pixi
# Install pixi (details can be found: https://pixi.sh/latest/#__tabbed_1_2)

iwr -useb https://pixi.sh/install.ps1 | iex

pixi run install_pyvrs

pixi task add run_test "python C:/tmp/MetaXRSimulator/MetaXRSimulator/scripts/vrs_pixmatch.py C:/tmp/test_recordings/replay_new.vrs C:/tmp/test_recordings/replay_target.vrs --threshold 0.4 --best_match_pixels_diff_threshold 40000 --diffs_output_path C:/tmp/test_recordings"

pixi task add pip_install "pip install -r C:/tmp/requirements_windows.txt"
pixi run pip_install
pixi run run_test

# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
# -*- coding: utf-8 -*-

import argparse
from pathlib import Path
import numpy as np
import math

import pyvrs

import sys

from PIL import Image
from pixelmatch.contrib.PIL import pixelmatch
from matplotlib import pyplot as plt

class Stream:
    def __init__(self, streamId:str, startTime:int, filtered_reader:pyvrs.filter.FilteredVRSReader) -> None:
        self.streamId = streamId
        self.startTime = startTime
        self.filtered_reader = filtered_reader

# Display the image of the frames passed in for the specified camera angles
def displayFrames(record:  Image, replay: Image) -> None:
    plt.figure(figsize=(20, 10))

    # Create a plot to present each image and
    # adjust the image according to the frames value
    plt.subplot(1, 2, 1)
    image = np.asarray(record)
    plt.imshow(image, vmin=0, vmax=255, cmap="gray")
    plt.title("Record")
    plt.axis("off")

    plt.subplot(1, 2, 2)
    image = np.asarray(replay)
    plt.imshow(image, vmin=0, vmax=255, cmap="gray")
    plt.title("Replay")
    plt.axis("off")

    plt.show()

def getStreams(reader:pyvrs.SyncVRSReader, sId: str) -> [Stream]:
    ret = []

    streamIds = reader.stream_ids
    # print(streamIds)

    for streamId in streamIds:
        if not streamId.startswith(sId):
            continue

        streamIdx = streamId[len(sId)+1]

        # NOTE: This doesn't separate the streams by index, only by type.
        # TODO: Handle multiple image streams in the same VRS file by iterating over reader.stream_ids() to find the matches
        filtered_reader = reader.filtered_by_fields(
                            stream_ids = {streamId},
                            record_types = {'configuration'},
                        )
        startTime = filtered_reader[0].timestamp
        filtered_reader = reader.filtered_by_fields(
                            stream_ids = {streamId},
                            record_types = {'data'},
                        )
        ret.append( Stream(streamId, startTime, filtered_reader) )
        # print(streamId, streamIdx, startTime)

    return ret

def getSnapshots(stream:Stream):
    ret = []
    for record in stream.filtered_reader:
        # print(record)
        meta = record.metadata_blocks[0]
        # print(record.timestamp, meta['iValue'])
        numFrames = int(meta['iValue'])
        ret.append({'numFrames': numFrames, 'timestamp': record.timestamp-stream.startTime})

    return ret

def getTimestamp(stream:Stream, idx: int) -> float:
    # print(idx, len(stream.filtered_reader))
    assert idx < len(stream.filtered_reader)
    record = stream.filtered_reader[idx]
    return record.timestamp-stream.startTime

def getFrameIdx(stream:Stream, timestamp:float):
    nextIdx = 0
    record = None
    while nextIdx < len(stream.filtered_reader):
        record = stream.filtered_reader[nextIdx]
        if timestamp < record.timestamp-stream.startTime:
            return nextIdx
        nextIdx += 1
    # use the last record
    return len(stream.filtered_reader)-1

def getFrame(stream:Stream, frameIndex:int):
    if frameIndex < 0 or frameIndex >= len(stream.filtered_reader):
        return None

    record = stream.filtered_reader[frameIndex]
    # print("getFrame:", frameIndex, record)
    pixels = np.array(record.image_blocks[0])
    # print(record, f'shape={pixels.shape}')
    image = Image.fromarray(np.uint8(pixels)).convert('RGBA')
    # print(image.format, image.size, image.mode)
    return {'image':image, 'timestamp': record.timestamp}

def formatTime(time:float):
    return "{:.3f}sec".format(time)

def compareImages(args, sampleFrame, sampleIdx, matchFrame, matchIdx, prefix):
    img_diff = None
    if(args.diffs_output_path):
        img_diff = Image.new("RGBA", matchFrame['image'].size)

    numPixelsDiff = pixelmatch(sampleFrame['image'], matchFrame['image'], output=img_diff, threshold=args.threshold, includeAA=True)

    if(args.diffs_output_path):
        sampleFrame['image'].save(f"{args.diffs_output_path}/{prefix}-sampleIdx-{sampleIdx}.png")
        matchFrame['image'].save(f"{args.diffs_output_path}/{prefix}-matchIdx-{matchIdx}.png")
        img_diff.save(f"{args.diffs_output_path}/{prefix}-{numPixelsDiff}-diffs-sampleIdx-{sampleIdx}-matchIdx-{matchIdx}.png")

    return numPixelsDiff

def findTheMatchingFrame(frame_stream, frame_name, target_frame_stream, target_frame_name, sampleIdx, max_test_frames, args, minPixelsDiff, bestMatch):
    sampleTS = getTimestamp(frame_stream, sampleIdx)
    sampleFrame = getFrame(frame_stream, sampleIdx)
    matchStartIdx = getFrameIdx(target_frame_stream, sampleTS)

    # search around the frame of interest until the closest match is found

    # find the matching frame in the target_frame_stream
    for matchIdx in range(matchStartIdx, matchStartIdx+max_test_frames):
        matchFrame = getFrame(target_frame_stream, matchIdx)
        if not matchFrame:
            break

        numPixelsDiff = compareImages(args, sampleFrame, sampleIdx, matchFrame, matchIdx, f"{target_frame_name}-{frame_name}")

        print("  compared ", frame_name, sampleIdx, " and ", target_frame_name, matchIdx, ": #px diff=", numPixelsDiff)
        if numPixelsDiff >= minPixelsDiff:
            continue

        minPixelsDiff = numPixelsDiff
        bestMatch = {'sampleStream':frame_name, 'minPixelsDiff':numPixelsDiff, 'sampleFrame':sampleFrame, 'sampleIdx':sampleIdx, 'matchStream':target_frame_name, 'matchFrame':matchFrame, 'matchIdx':matchIdx}

        if(minPixelsDiff == 0):
            break
    return (minPixelsDiff, bestMatch)

def main():

    parser = argparse.ArgumentParser(description='use pixelmatch to diff vrs files')
    parser.add_argument('record', type=str, help="VRS reference recording file to compare against")
    parser.add_argument('replay', type=str, help="VRS replay to check")
    # parser.add_argument('record frames')
    parser.add_argument('--threshold', type=float, required=False, default=0.1, help="threshold over which the pixels are considered to be different")
    parser.add_argument('--sample_location', type=float, required=False, default=0.5, help="where in the group to test: 0=beginning, 0.5=middle, 1=end")
    parser.add_argument('--max_test_frames', type=int, required=False, default=3, help="max number of matches to perform in each direction")
    parser.add_argument('--show_matches', type=bool, required=False, default=False, help="display the corresponding screenshots side-by-side once a match is identified")
    parser.add_argument('--diffs_output_path', type=str, required=False, default=None, help="output folder for image diffs")
    parser.add_argument('--best_match_pixels_diff_threshold', type=int, required=False, default=100, help="threshold on the pixels diff of the best match")

    args = parser.parse_args()

    record_path = Path(args.record)
    record_reader = pyvrs.SyncVRSReader(record_path)

    replay_path = Path(args.replay)
    replay_reader = pyvrs.SyncVRSReader(replay_path)

    # NOTE: only use the first snapshot stream
    record_snapshot_stream = getStreams(record_reader, "400")[0]
    replay_snapshot_stream = getStreams(replay_reader, "400")[0]

    # print(f'snapshots for {record_path}')
    record_groups = getSnapshots(record_snapshot_stream)
    # print(f'snapshots for {replay_path}')
    replay_groups = getSnapshots(replay_snapshot_stream)

    assert len(record_groups) == len(replay_groups)

    print("Comparing", record_path, "with", replay_path, ": num groups=", len(record_groups))

    # NOTE: only use the first frame stream for now
    record_frame_stream = getStreams(record_reader, "8003")[0]
    replay_frame_stream = getStreams(replay_reader, "8003")[0]

    groupOffset = 0
    for groupIdx in range(0, len(record_groups)):
        assert record_groups[groupIdx]['numFrames'] == replay_groups[groupIdx]['numFrames']
        numFrames = record_groups[groupIdx]['numFrames']

        sample_location = args.sample_location
        sampleIdx = groupOffset+int(numFrames*sample_location)
        max_test_frames = args.max_test_frames

        print("Group", groupIdx, ":numFrames=", numFrames,"sampleIndex=", sampleIdx)

        minPixelsDiff = 100000000000
        bestMatch = None

        # first check the replay against the record
        (minPixelsDiff, bestMatch) = findTheMatchingFrame(record_frame_stream, "record", replay_frame_stream, "replay", sampleIdx, max_test_frames, args, minPixelsDiff, bestMatch)

        if minPixelsDiff > args.best_match_pixels_diff_threshold:
            print("best match diff is > ", args.best_match_pixels_diff_threshold, " so matching record against replay")
            # then check the record against the replay
            (minPixelsDiff, bestMatch) = findTheMatchingFrame(replay_frame_stream, "replay", record_frame_stream, "record", sampleIdx, max_test_frames, args, minPixelsDiff, bestMatch)

        # report the best match
        if bestMatch:
            print(f" Best sample is {bestMatch['sampleStream']}#{bestMatch['sampleIdx']} ts={formatTime(bestMatch['sampleFrame']['timestamp'])}, best match is {bestMatch['matchStream']}#{bestMatch['matchIdx']} ts={formatTime(bestMatch['matchFrame']['timestamp'])}, #px diff={bestMatch['minPixelsDiff']}")
            if(args.show_matches):
                displayFrames(bestMatch['sampleFrame']['image'], bestMatch['matchFrame']['image'])
            if(bestMatch['minPixelsDiff']> args.best_match_pixels_diff_threshold):
                print(" The best match's pixels difference exceeds the threshold of ", args.best_match_pixels_diff_threshold, ". So this is a mismatch")
                sys.exit(-1)
            else:
                print(" The best match's pixels difference is less than the threshold of ", args.best_match_pixels_diff_threshold, ". So this is a match")
        else:
            print(" no match found!")
            sys.exit(-2)

        groupOffset += numFrames

    return 0

if __name__ == "__main__":
    main()