faster-computer-control-lstm-unified-encoding.py

#!/usr/bin/env python
# coding: utf-8

# you can save a version of the notebook, then copy and edit the version.
# 
# # Useful Library:
# 
# [timm](https://github.com/huggingface/pytorch-image-models): pytorch image models

# you'd use einops to flatten the tensor.
# since you would be doing these operations separately.
# 
# to measure the overall loss, you would consider merging two types of tensor into one.

# this model does not need to train the ViT transformer, and does not need to decode the vector embedding of the ViT output.
# 
# mouse actions could be tokenized if send x and y coords sequentially.
# 
# make x and y coords continuous by squeezing it into 4 separate floating point numbers.
# 
# change the codebase and speed up training in next revision.

# # Code

# ## import some library and setup data filepaths

# In[1]:


# get_ipython().system('pip3 install einops')
import einops


# In[2]:


# try to train a model.
# first let's get the data.

# if you want to use the template generated by gpt4, you have to use kaggle.

##############################
## UPLOAD TO KAGGLE DATASET ##
##############################
import os

basePath = "/kaggle/input/agi-computer-control-test-dataset"
videoPath = os.path.join(basePath, "93755268.mp4")
# videoPath = "ffmpeg_output_test_hwaccel.mp4"
metadata_file_path = os.path.join(basePath,"screenshot_and_actions.json")

# keep it simple, just want to scratch the surface.
# we can run this on cpu.
# from PIL import Image
import numpy as np
import cv2
import ast
from pydantic import BaseModel, validator
from typing import Union
try:
    from typing import Literal
except:
    from typing_extensions import Literal # this is a failsafe.


# ## define some data models

# In[3]:


# import pynput
# no such dependency when training.


# change of the encoding? how to encode these stuff?
# train similar actions together? to train the encoder?

# sparse -> embedding layer -> dense over every output to decode?
# 
# when doing sparse decoding, you can rank probability and sample.
# when doing dense embedding, you can add noise and remove determinism.

# you can even add random noise to embedding input (before going into softmax?) to make network active (non-dead).

# fourier transform?
# you encode these things by addition, using fourier transform to decode.
#

# for mouse coordinates:
# you encode different part of info into different frequencies
# then you decode the info using some decoder
# you may share the decoder, to ensure its integrety.

# mouse_emb (elementwise_product) mouse_fft

# [special_emb]
# [vit_emb]
# [mouse_emb] <- compound emb: emb_1*compound_sin_1 + emb_2*compound_sin_2+ emb_3*compound_sin_3+ emb_4*compound_sin_4-> (use 2 dense layers, or transformer?) multiple decoder (dense)
# [keyboard_emb]

# mouse use continual spatial sparse encoding from HTM, then pass through Linear layer.
# keyboard use conventional embedding

# for continual embedding, you output a single value, instead of sparse.
# if choices are few, make it sparse.

# so:
# mouse - continual, squeezed to single bits.
# special bits - sparse (fft or ifft?)
# keys - sparse

# how to encode multiple mouse embedding?

class HIDActionBase:
    mouse_resolution: int = 1000
    keyboard_action_types = [
        "key_press",
        "key_release",
    ]
    mouse_action_types = [
        "mouse_move",
        "mouse_click",
        "mouse_scroll",
    ]
    action_types = [
        *keyboard_action_types,
        *mouse_action_types,
        # None,  # end of action
        # there is no such thing here. do it externally.
    ]
    mouse_buttons = [
        "Button.left",
        "Button.middle",
        "Button.right",
    ]
    keys = [
        """','""",
        """'.'""",
        """'/'""",
        """';'""",
        """\"'\"""",
        """'['""",
        """']'""",
        """'\\'""",
        """'='""",
        """'-'""",
        """'0'""",
        """'9'""",
        """'8'""",
        """'7'""",
        """'6'""",
        """'5'""",
        """'4'""",
        """'3'""",
        """'2'""",
        """'1'""",
        """'`'""",
        """'a'""",
        """'b'""",
        """'c'""",
        """'d'""",
        """'e'""",
        """'f'""",
        """'g'""",
        """'h'""",
        """'i'""",
        """'j'""",
        """'k'""",
        """'l'""",
        """'m'""",
        """'n'""",
        """'o'""",
        """'p'""",
        """'q'""",
        """'r'""",
        """'s'""",
        """'t'""",
        """'u'""",
        """'v'""",
        """'w'""",
        """'x'""",
        """'y'""",
        """'z'""",
        "Key.alt",
        "Key.alt",
        "Key.alt_r",
        "Key.alt_r",
        "Key.backspace",
        "Key.caps_lock",
        "Key.cmd",
        "Key.cmd",
        "Key.cmd_r",
        "Key.ctrl",
        "Key.ctrl",
        "Key.ctrl_r",
        "Key.delete",
        "Key.down",
        "Key.end",
        "Key.enter",
        "Key.esc",
        "Key.f1",
        "Key.f2",
        "Key.f3",
        "Key.f4",
        "Key.f5",
        "Key.f6",
        "Key.f7",
        "Key.f8",
        "Key.f9",
        "Key.f10",
        "Key.f11",
        "Key.f12",
        "Key.f13",
        "Key.f14",
        "Key.f15",
        "Key.f16",
        "Key.f17",
        "Key.f18",
        "Key.f19",
        "Key.f20",
        "Key.home",
        "Key.left",
        "Key.page_down",
        "Key.page_up",
        "Key.right",
        "Key.shift",
        "Key.shift",
        "Key.shift_r",
        "Key.space",
        "Key.tab",
        "Key.up",
        "Key.media_play_pause",
        "Key.media_volume_mute",
        "Key.media_volume_down",
        "Key.media_volume_up",
        "Key.media_previous",
        "Key.media_next",
    ]
    
    length = (len(action_types)
            + len(keys)
            + len(mouse_buttons)
            + 1 # mouse pressed
            + 4 * mouse_resolution) #,
#             1)

    @staticmethod
    def unshift_keycode(keycode: str) -> Union[str, None]:
        unshift_keycodes = {
            "!": "1",
            "@": "2",
            "#": "3",
            "$": "4",
            "%": "5",
            "^": "6",
            "&": "7",
            "*": "8",
            "(": "9",
            ")": "0",
            "_": "-",
            "+": "=",
            "{": "[",
            "}": "]",
            "|": "\\",
            ":": ";",
            '"': "'",
            "<": ",",
            ">": ".",
            "?": "/",
            "~": "`",
        }
        ctrl_keycodes = {
            "\x01": "a",
            "\x02": "b",
            "\x03": "c",
            "\x04": "d",
            "\x05": "e",
            "\x06": "f",
            "\x07": "g",
            "\x08": "h",
            "\t": "i",
            "\n": "j",
            "\x0b": "k",
            "\x0c": "l",
            "\r": "m",
            "\x0e": "n",
            "\x0f": "o",
            "\x10": "p",
            "\x11": "q",
            "\x12": "r",
            "\x13": "s",
            "\x14": "t",
            "\x15": "u",
            "\x16": "v",
            "\x17": "w",
            "\x18": "x",
            "\x19": "y",
            "\x1a": "z",
            "<219>": "[",
            "<221>": "]",
            "<189>": "-",
            "<187>": "=",
            "<192>": "`",
            "<48>": "0",
            "<49>": "1",
            "<50>": "2",
            "<51>": "3",
            "<52>": "4",
            "<53>": "5",
            "<54>": "6",
            "<55>": "7",
            "<56>": "8",
            "<57>": "9",
            "<220>": "\\",
            "<186>": ";",
            "<222>": "'",
            "<188>": ",",
            "<190>": ".",
            "<191>": "/",
        }
        keycode = unshift_keycodes.get(keycode, ctrl_keycodes.get(keycode, keycode))
        # still, this is something out of concern.
        if keycode.startswith("<") and keycode.endswith(">"):
            print("Discarding unconvertable keycode: %s" % keycode)
            # keycode = pynput.keyboard.KeyCode(int(keycode[1:-1]))
            return
        return keycode

    @staticmethod
    def uncover_keycode(keycode: str) -> Union[str, None]:
        if not keycode.startswith("Key."):
            keycode_converted = HIDActionBase.unshift_keycode(
                keycode
                if keycode.startswith("<") and keycode.endswith(">")
                else ast.literal_eval(keycode)
            )
            return keycode_converted
            # this could be None.
            # when this is None, simply skip this code. do not end the conversion. skip it.
        else:
            return keycode


class HIDAction(BaseModel, HIDActionBase):
    # static method: from_action
    # static method: from_ndarray
    # instance method: to_ndarray
    # instance method: to_action
    max_x: int
    max_y: int
    action_type: Union[
        Literal["key_press"],  # ["key_press", "'w'"]
        Literal["key_release"],  # ["key_release", "'r'"]
        Literal[
            "mouse_move"
        ],  # ["mouse_move", [176.7734375, 580.40625]], "timeStamp": 1680247557.125498}
        Literal[
            "mouse_click"
        ],  # ["mouse_click", [176.7734375, 580.40625, "Button.left", true]]
        Literal["mouse_scroll"],  # ["mouse_scroll", [938.76171875, 318.75, 0, 0]]
#         None,  # end_of_action
    ] # you need to specify this.
    key: Union[
        Literal["""','"""],
        Literal["""'.'"""],
        Literal["""'/'"""],
        Literal["""';'"""],
        Literal["""\"'\""""],
        Literal["""'['"""],
        Literal["""']'"""],
        Literal["""'\\'"""],
        Literal["""'='"""],
        Literal["""'-'"""],
        Literal["""'0'"""],
        Literal["""'9'"""],
        Literal["""'8'"""],
        Literal["""'7'"""],
        Literal["""'6'"""],
        Literal["""'5'"""],
        Literal["""'4'"""],
        Literal["""'3'"""],
        Literal["""'2'"""],
        Literal["""'1'"""],
        Literal["""'`'"""],
        Literal["""'a'"""],
        Literal["""'b'"""],
        Literal["""'c'"""],
        Literal["""'d'"""],
        Literal["""'e'"""],
        Literal["""'f'"""],
        Literal["""'g'"""],
        Literal["""'h'"""],
        Literal["""'i'"""],
        Literal["""'j'"""],
        Literal["""'k'"""],
        Literal["""'l'"""],
        Literal["""'m'"""],
        Literal["""'n'"""],
        Literal["""'o'"""],
        Literal["""'p'"""],
        Literal["""'q'"""],
        Literal["""'r'"""],
        Literal["""'s'"""],
        Literal["""'t'"""],
        Literal["""'u'"""],
        Literal["""'v'"""],
        Literal["""'w'"""],
        Literal["""'x'"""],
        Literal["""'y'"""],
        Literal["""'z'"""],
        Literal["Key.alt"],
        Literal["Key.alt"],
        Literal["Key.alt_r"],
        Literal["Key.alt_r"],
        Literal["Key.backspace"],
        Literal["Key.caps_lock"],
        Literal["Key.cmd"],
        Literal["Key.cmd"],
        Literal["Key.cmd_r"],
        Literal["Key.ctrl"],
        Literal["Key.ctrl"],
        Literal["Key.ctrl_r"],
        Literal["Key.delete"],
        Literal["Key.down"],
        Literal["Key.end"],
        Literal["Key.enter"],
        Literal["Key.esc"],
        Literal["Key.f1"],
        Literal["Key.f2"],
        Literal["Key.f3"],
        Literal["Key.f4"],
        Literal["Key.f5"],
        Literal["Key.f6"],
        Literal["Key.f7"],
        Literal["Key.f8"],
        Literal["Key.f9"],
        Literal["Key.f10"],
        Literal["Key.f11"],
        Literal["Key.f12"],
        Literal["Key.f13"],
        Literal["Key.f14"],
        Literal["Key.f15"],
        Literal["Key.f16"],
        Literal["Key.f17"],
        Literal["Key.f18"],
        Literal["Key.f19"],
        Literal["Key.f20"],
        Literal["Key.home"],
        Literal["Key.left"],
        Literal["Key.page_down"],
        Literal["Key.page_up"],
        Literal["Key.right"],
        Literal["Key.shift"],
        Literal["Key.shift"],
        Literal["Key.shift_r"],
        Literal["Key.space"],
        Literal["Key.tab"],
        Literal["Key.up"],
        Literal["Key.media_play_pause"],
        Literal["Key.media_volume_mute"],
        Literal["Key.media_volume_down"],
        Literal["Key.media_volume_up"],
        Literal["Key.media_previous"],
        Literal["Key.media_next"],
        None,
    ] = None

    mouse_button: Union[
        Literal["Button.left"], Literal["Button.middle"], Literal["Button.right"], None
    ] = None
    mouse_pressed: Union[bool, None] = None
    x: Union[float, None] = None
    y: Union[float, None] = None
    dx: Union[float, None] = None
    dy: Union[float, None] = None

    @validator("max_x", "max_y")
    def greater_than_zero(cls, v):
        assert type(v) == int
        assert v > 0
        return v

    @validator("action_type")
    def action_type_within_action_types(cls, v):
        if v:
            assert v in HIDActionBase.action_types
        return v

    @validator("key")
    def key_within_keys(cls, v):
        if v:
            assert v in HIDActionBase.keys
        return v

    @validator("mouse_button")
    def mouse_button_within_mouse_buttons(cls, v):
        if v:
            assert v in HIDActionBase.mouse_buttons
        return v

    @validator("mouse_pressed")
    def mouse_pressed_type_check(cls, v):
        if v:
            assert type(v) == bool
        return v

    @staticmethod
    def from_action_json(action_json: list, max_x: int, max_y: int):
        action_type = action_json[0]
        action_args = action_json[1]

        construct_args = dict(max_x=max_x, max_y=max_y, action_type=action_type)

        if action_type == "key_press":
            assert action_args in HIDActionBase.keys

            construct_args.update(dict(key=action_args))
        elif action_type == "key_release":
            assert action_args in HIDActionBase.keys

            construct_args.update(dict(key=action_args))
        elif action_type == "mouse_move":
            assert action_args[0] >= 0 and action_args[0] <= max_x
            assert action_args[1] >= 0 and action_args[1] <= max_y

            construct_args.update(dict(x=action_args[0], y=action_args[1]))
        elif action_type == "mouse_click":
            assert action_args[0] >= 0 and action_args[0] <= max_x
            assert action_args[1] >= 0 and action_args[1] <= max_y
            assert action_args[2] in HIDActionBase.mouse_buttons
            assert type(action_args[3]) == bool

            construct_args.update(
                dict(
                    x=action_args[0],
                    y=action_args[1],
                    mouse_button=action_args[2],
                    mouse_pressed=action_args[3],
                )
            )
        elif action_type == "mouse_scroll":
            assert action_args[0] >= 0 and action_args[0] <= max_x
            assert action_args[1] >= 0 and action_args[1] <= max_y
            assert action_args[2] >= -max_x and action_args[2] <= max_x
            assert action_args[3] >= -max_y and action_args[3] <= max_y

            construct_args.update(
                dict(
                    x=action_args[0],
                    y=action_args[1],
                    dx=action_args[2],
                    dy=action_args[3],
                )
            )
        else:
            raise Exception(
                "Unknown action type: %s\naction args: %s" % (action_type, action_args)
            )

        mHIDAction = HIDAction(**construct_args)
        return mHIDAction

    @staticmethod
    def from_ndarray(ndarray: np.ndarray, max_x: int, max_y: int):
        assert ndarray.shape == (
            HIDActionBase.length,
        )
        cursor = 0

        action_type_ndarray = ndarray[cursor : cursor + len(HIDActionBase.action_types)]
        cursor += len(HIDActionBase.action_types)
        action_type_index = np.argmax(action_type_ndarray)
        action_type = HIDActionBase.action_types[action_type_index]
        del action_type_ndarray
        del action_type_index

        construct_args = dict(max_x=max_x, max_y=max_y, action_type=action_type)

        if action_type:
            key_ndarray = ndarray[cursor : cursor + len(HIDActionBase.keys)]
            cursor += len(HIDActionBase.keys)
            key_index = np.argmax(key_ndarray)
            key = HIDActionBase.keys[key_index]
            del key_ndarray
            del key_index

            mouse_button_ndarray = ndarray[
                cursor : cursor + len(HIDActionBase.mouse_buttons)
            ]
            cursor += len(HIDActionBase.mouse_buttons)
            mouse_button_index = np.argmax(mouse_button_ndarray)
            mouse_button = HIDActionBase.mouse_buttons[mouse_button_index]
            del mouse_button_ndarray
            del mouse_button_index

            mouse_pressed_ndarray = ndarray[cursor : cursor + 1]
            cursor += 1
            mouse_pressed = bool(mouse_pressed_ndarray[0][0])
            del mouse_pressed_ndarray

            x_ndarray = ndarray[cursor : cursor + HIDActionBase.mouse_resolution]
            cursor += HIDActionBase.mouse_resolution
            x_index = np.argmax(x_ndarray)
            x = (x_index / HIDActionBase.mouse_resolution) * max_x
            del x_ndarray
            del x_index

            y_ndarray = ndarray[cursor : cursor + HIDActionBase.mouse_resolution]
            cursor += HIDActionBase.mouse_resolution
            y_index = np.argmax(y_ndarray)
            y = (y_index / HIDActionBase.mouse_resolution) * max_y
            del y_ndarray
            del y_index

            dx_ndarray = ndarray[cursor : cursor + HIDActionBase.mouse_resolution]
            cursor += HIDActionBase.mouse_resolution
            dx_index = np.argmax(dx_ndarray)
            dx = (dx_index / HIDActionBase.mouse_resolution) * 2 * max_x - max_x
            del dx_ndarray
            del dx_index

            dy_ndarray = ndarray[cursor : cursor + HIDActionBase.mouse_resolution]
            cursor += HIDActionBase.mouse_resolution
            dy_index = np.argmax(dy_ndarray)
            dy = (dy_index / HIDActionBase.mouse_resolution) * 2 * max_y - max_y
            del dy_ndarray
            del dy_index

            if action_type == "key_press":
                construct_args.update(dict(key=key))
            elif action_type == "key_release":
                construct_args.update(dict(key=key))
            elif action_type == "mouse_move":
                construct_args.update(dict(x=x, y=y))
            elif action_type == "mouse_click":
                construct_args.update(
                    dict(
                        x=x, y=y, mouse_button=mouse_button, mouse_pressed=mouse_pressed
                    )
                )
            elif action_type == "mouse_scroll":
                construct_args.update(dict(x=x, y=y, dx=dx, dy=dy))
        else:
            pass

        del cursor

        mHIDAction = HIDAction(**construct_args)
        return mHIDAction

    def round_within(self, number: Union[int, float], number_name: str) -> int:
        result = round(number)
        if result > self.mouse_resolution:
            print(f"Warning: {number_name} overflow")
            print(f"Value {result} greater than {self.mouse_resolution}")
            return self.mouse_resolution
        elif result < 0:
            print(f"Warning: {number_name} overflow")
            print(f"Value {result} smaller than 0")
            return 0
        return result

    def to_ndarray(
        self,
    ) -> np.ndarray:
        action_type_ndarray = np.zeros((len(self.action_types), 1))
        action_type_ndarray[self.action_types.index(self.action_type)] = 1

        key_ndarray = np.zeros((len(self.keys), 1))
        if self.key:
            key_ndarray[self.keys.index(self.key)] = 1

        mouse_button_ndarray = np.zeros((len(self.mouse_buttons), 1))
        if self.mouse_button:
            mouse_button_ndarray[self.mouse_buttons.index(self.mouse_button)] = 1

        mouse_pressed_array = np.zeros((1, 1))
        if self.mouse_pressed:
            mouse_pressed_array[0] = 1

        x_ndarray = np.zeros((self.mouse_resolution, 1))
        if self.x:
            x_ndarray[
                self.round_within(self.mouse_resolution * self.x / self.max_x, "X")
            ] = 1

        y_ndarray = np.zeros((self.mouse_resolution, 1))
        if self.y:
            y_ndarray[
                self.round_within(self.mouse_resolution * self.y / self.max_y, "Y")
            ] = 1

        dx_ndarray = np.zeros((self.mouse_resolution, 1))
        if self.dx:
            dx_ndarray[
                self.round_within(
                    self.mouse_resolution * (self.dx + self.max_x) / (2 * self.max_x),
                    "DX",
                )
            ] = 1

        dy_ndarray = np.zeros((self.mouse_resolution, 1))
        if self.dy:
            dy_ndarray[
                self.round_within(
                    self.mouse_resolution * (self.dy + self.max_y) / (2 * self.max_y),
                    "DY",
                )
            ] = 1

        ndarray = np.concatenate(
            [
                action_type_ndarray,
                key_ndarray,
                mouse_button_ndarray,
                mouse_pressed_array,
                x_ndarray,
                y_ndarray,
                dx_ndarray,
                dy_ndarray,
            ]
        )
        return ndarray

    def to_action_json(
        self,
    ) -> Union[list, None]:
        action_type = self.action_type
        if action_type:
            if action_type == "key_press":
                assert self.key in self.keys

                action_args = self.key
            elif action_type == "key_release":
                assert self.key in self.keys

                action_args = self.key
            elif action_type == "mouse_move":
                assert self.x >= 0 and self.x <= self.max_x
                assert self.y >= 0 and self.y <= self.max_y

                action_args = [self.x, self.y]
            elif action_type == "mouse_click":
                assert self.x >= 0 and self.x <= self.max_x
                assert self.y >= 0 and self.y <= self.max_y
                assert self.mouse_button in self.mouse_buttons
                assert type(self.mouse_pressed) == bool

                action_args = [self.x, self.y, self.mouse_button, self.mouse_pressed]
            elif action_type == "mouse_scroll":
                assert self.x >= 0 and self.x <= self.max_x
                assert self.y >= 0 and self.y <= self.max_y
                assert self.dx >= -self.max_x and self.dx <= self.max_x
                assert self.dy >= -self.max_y and self.dy <= self.max_y

                action_args = [self.x, self.y, self.dx, self.dy]
            else:
                raise Exception("Unknown action_type: %s" % action_type)
            action_json = [action_type, action_args]
        else:
            action_json = None
        return action_json


# ## read data from path

# In[4]:


import json

with open(metadata_file_path, "r") as f:
    data = json.loads(f.read())

# print(data.keys())
# ['screenshot_and_actions', 'perspective_size', 'timestep']

perspective_width, perspective_height = data["perspective_size"]
timestep = data["timestep"]

print("PERSPECTIVE:", perspective_width, perspective_height)
print("TIMESTEP:", timestep)

import re

class VideoCaptureContextManager:
    def __init__(self, videoPath):
        self.videoPath = videoPath
        
    def __enter__(self):
        print("Entering the context...")
        self.cap = cv2.VideoCapture(self.videoPath)
        return self.cap

    def __exit__(self, exc_type, exc_value, exc_tb):
        try:
            self.cap.release()
        finally:
            import gc
            gc.collect()
            print("Leaving the context...")
        #  print(exc_type, exc_value, exc_tb, sep="\n")


# ## test and define the model

# In[5]:


import torch
import torchvision


# In[6]:


vit_model = torchvision.models.vit_b_16(pretrained=True)


# In[7]:


# use einops.pack and einops.unpack

# def size_splits(tensor, split_sizes, dim=0):
#     """Splits the tensor according to chunks of split_sizes.
    
#     Arguments:
#         tensor (Tensor): tensor to split.
#         split_sizes (list(int)): sizes of chunks
#         dim (int): dimension along which to split the tensor.
#     """
#     if dim < 0:
#         dim += tensor.dim()
    
#     dim_size = tensor.size(dim)
#     if dim_size != torch.sum(torch.Tensor(split_sizes)):
#         raise KeyError("Sum of split sizes exceeds tensor dim")
    
#     splits = torch.cumsum(torch.Tensor([0] + split_sizes), dim=0)[:-1]

#     return tuple(tensor.narrow(int(dim), int(start), int(length)) 
#         for start, length in zip(splits, split_sizes))


# In[8]:


# get_ipython().system('pip3 install pydantic-numpy')
from pydantic_numpy import NDArray


# In[9]:


import numpy as np
from typing import Mapping, List

class ConsciousBase:
    data_types = ['image', 'HIDAction']
    special_tokens = ['image_newline', 'image_end', 'action_end', None]
#     vector_size = 1+2+1000+4110 # visual things are pre-encoded. no raw image here!

    image_dim = 224
    image_channels = 3
    
    data_type_length = len(data_types)
    special_token_length = len(special_tokens)
    image_length = image_dim*image_dim*image_channels
    
    # FIX 1: plus to colon.
    split_sizes = [ len(data_types),
    len(special_tokens) ,
    image_length , # FIX 9: change to flattened image bits count
    HIDActionBase.length  # 4110?
                  ]
    length = sum(split_sizes)
    
    # you cannot easily revert this compression by argmax or something else.
    # so you need the decoder.
    
# can it be consciousnessless?

class ConsciousBlock(BaseModel):
    data_type:Union[Literal["image"],
                   Literal["HIDAction"]] # 2 bits, required
    special_token:Union[Literal["image_newline"],
                       Literal["image_end"],
                       Literal["action_end"], # change some of these bits into -torch.inf, so you won't have paradox like results.
                       None] = None # 4 bits
    image_data: Union[None, NDArray] = None # what is the shape of this image data? assume to be [3,224,224] (c h w) flattened
    action_data: Union[None, NDArray] = None # assume to be: (4110, )
    # [1,1000] -> [3,1000,1000] -> [3,224,224]
    #    einsum.repeat       conv2d
    
    # so, maybe you still need some ViT decode layer.
    
    @staticmethod 
    def from_json(data:Mapping):
        mConsciousBlock = ConsciousBlock(**data)
        return mConsciousBlock

    def to_json(self) -> Mapping:
        mJson = self.dict()
        return mJson
    
    @staticmethod
    def from_tensor(tensor:torch.Tensor):
        # check its shape.
        assert tensor.shape == (ConsciousBlock.length,)
        split_sizes = ConsciousBase.split_sizes
        data_bits, special_bits, image_data, action_data = einops.unpack(tensor, [[s] for s in split_sizes], "*")
#         data_bits, special_bits, image_data, action_data = size_splits(tensor, split_sizes)
        
        data_type = ConsciousBase.data_types[int(torch.argmax(data_bits))]
        if data_type == "image":
            special_bits[2] = -torch.inf
            special_token = ConsciousBase.special_tokens[int(torch.argmax(special_bits))]
            
            mConsciousBlock = ConsciousBlock(data_type=data_type, special_token=special_token, image_data=image_data)
        elif data_type == "HIDAction":
            special_bits[0:2] = -torch.inf
            special_token = ConsciousBase.special_tokens[int(torch.argmax(special_bits))]
            
            mConsciousBlock = ConsciousBlock(data_type=data_type, special_token=special_token, action_data=action_data)
        else:
            raise Exception("Unknown data_type:", data_type)
        return mConsciousBlock
    
    def to_tensor(self) -> torch.Tensor:
        mTensorBase = torch.zeros((ConsciousBase.length))
        # BUG 1: not enough values to unpack
#         print(ConsciousBase.length)
#         print(ConsciousBase.split_sizes)
        data_bits, special_bits, image_data, action_data = einops.unpack(mTensorBase, [[s] for s in ConsciousBase.split_sizes], '*')
#         data_bits, special_bits, image_data, action_data = size_splits(mTensorBase, ConsciousBase.split_sizes)
        data_bits[ConsciousBase.data_types.index(self.data_type)] = 1
        if self.data_type == 'image':
            # BUG 2: comparing ndarray to None
            # FIX 2: change "!=" into "is not"
            assert self.image_data is not None
            assert self.image_data.shape == (ConsciousBase.image_length, )
            assert self.special_token!="action_end"
            
            image_data = torch.Tensor(self.image_data)
            special_bits[ConsciousBase.special_tokens.index(self.special_token)] = 1
            
        elif self.data_type == 'HIDAction':
            assert self.action_data is not None
            assert self.action_data.shape == (HIDActionBase.length, )
            assert self.special_token not in ['image_newline', 'image_end']
            
            action_data = torch.Tensor(self.action_data)
            special_bits[ConsciousBase.special_tokens.index(self.special_token)] = 1
            
        mTensor, _ = einops.pack((data_bits, special_bits, image_data, action_data), "*")
#         mTensor = torch.concat((data_bits, special_bits, image_data, action_data))
        del mTensorBase
        return mTensor
    
class ConsciousFlow(BaseModel, ConsciousBase):
    consciousBlocks:List[ConsciousBlock]
        
    @staticmethod
    def from_json(data:List[Mapping]):
        mList = [ConsciousBlock.from_json(j) for j in data]
        mConsciousFlow = ConsciousFlow(mList)
        return mConsciousFlow

    def to_json(self) -> List[Mapping]:
        mJson = [c.to_json() for c in self.consciousBlocks]
        return mJson
    
    @staticmethod
    def from_tensor(tensor:torch.Tensor):
        consciousBlockCount, vector_length = tensor.shape
        assert vector_length ==  ConsciousBase.vector_length
        mConsciousBlocks = []
        for i in range(consciousBlockCount):
            arr = tensor[i,:]
            mConsciousBlock = ConsciousBlock.from_tensor(arr)
            mConsciousBlocks.append(mConsciousBlock)
        mConsciousFlow = ConsciousFlow(mConsciousBlocks)
        return mConsciousFlow
    
    def to_tensor(self) -> torch.Tensor:
        mTensor, _ = einops.pack([c.to_tensor() for c in self.consciousBlocks],'* d')
#         mTensor = torch.Tensor([c.to_tensor() for c in self.consciousBlocks])
        return mTensor


# In[10]:


# notice: when in online mode only image will be backpropagated.
# like using some upside down mirror.


class CustomModel(torch.nn.Module):
    def __init__(self, vit_model, hidden_size_vit=1000, vit_block_size=228):
#     def __init__(self, rwkv_model, vit_model, tokenizer, hidden_size_rwkv, hidden_size_vit, output_size, vit_times = 4, vit_block_size=228):
        super(CustomModel, self).__init__()
#         self.rwkv_model = rwkv_model # processing language, generate actions.
        self.vit_model = vit_model
        self.vit_block_size = vit_block_size # this is default.
#         self.vit_times = vit_times
        
        # seq2seq alike.
#         self.hidden_size = hidden_size_rwkv+hidden_size_vit
        self.hidden_size = hidden_size_vit
    
        self.HIDEncoder = torch.nn.Linear(HIDActionBase.length, 1000)
        self.HIDDecoder = torch.nn.Linear(1000, HIDActionBase.length) # use torch.where or something 
        # sparse matrix?
        
        # some magic going elsewhere.
        self.ViTDecoder = torch.nn.Conv2d(in_channels=9, out_channels=3, kernel_size=22, stride=4, dilation=5) # n c h w
    
        # FIX 13: change 1000 to ConsciousBase.data_type_length+ConsciousBase.special_token_length+1000
        io_h_size=ConsciousBase.data_type_length+ConsciousBase.special_token_length+1000
        self.rnn = torch.nn.LSTM(input_size=io_h_size, hidden_size = io_h_size, batch_first=True)

        # use tensor.
    def forward(self, conscious_stream:torch.Tensor, target_output:Union[None, torch.Tensor]=None):
        # if have target, and our datatype bits are wrong, we only optimize the datatype bits, making image/action bits identical
        # otherwise just calculate all bits from the model.
        
        # conscious_stream: [batch_size, (ConsciousBase.length) data_type+special_tokens+image_bits+action_bits]
        # you need another pydantic model for it. 
        
        # BUG 9: input shape error
        # FIX 9: check input and output shape
#         print(conscious_stream.shape, ConsciousBase.length, )
#         if target_output is not None:
#             print(target_output.shape)
        batch_size, seq_length, dim_block = conscious_stream.shape
        assert dim_block == ConsciousBase.length
        
        conscious_stream_reshaped = einops.rearrange(conscious_stream, "b s d -> (b s) d")
        
        datatype_bits, special_bits, image_bits, action_bits = einops.unpack(conscious_stream_reshaped, [[s] for s in ConsciousBase.split_sizes],"b *")
#         datatype_bits, special_bits, image_bits, action_bits = size_splits(conscious_stream_reshaped, ConsciousBase.split_sizes,dim=1)
#         desired_size = self.vit_block_size*self.vit_times

        datatypes = torch.argmax(datatype_bits, dim=1)
        image_indexs = datatypes == 0
        action_indexs = datatypes == 1
        
        special_tokens = torch.argmax(special_bits, dim=1)
        image_newline_indexs = special_tokens == 0
        image_end_indexs = special_tokens == 1
        action_end_indexs = special_tokens == 2
        nop_indexs = special_tokens == 3
        
#         4+2+1000
        
        # you are gonna take actions.
        # prepare some zeros.
        
        # BUG 10: len() on int
        # FIX 10: find and fix incorrect len() calls
        batched_rnn_input = torch.zeros((batch_size*seq_length, ConsciousBase.data_type_length+ConsciousBase.special_token_length+1000))
        
        batched_rnn_input[image_indexs, 0] = 1
        batched_rnn_input[action_indexs, 1] = 1
        
        # BUG 11: indexs not defined
        # FIX 11: indexes -> indexs
        batched_rnn_input[image_newline_indexs, 2] = 1
        batched_rnn_input[image_end_indexs, 3] = 1
        batched_rnn_input[action_end_indexs, 4] = 1
        batched_rnn_input[nop_indexs, 5] = 1
        
        # process this.
        datatype_and_special_token_length = ConsciousBase.data_type_length + ConsciousBase.special_token_length
        
        # BUG 12: unannotated unknown axes in einops.rearrange
        # FIX 12: annotate these axes
        selected_image_bits = image_bits[image_indexs, :]
        transformed_image_bits = einops.rearrange(selected_image_bits, "b (c h w) -> b c h w", h = ConsciousBase.image_dim, w = ConsciousBase.image_dim, c = ConsciousBase.image_channels)
        processed_image_bits = self.vit_model(transformed_image_bits)
        batched_rnn_input[image_indexs, datatype_and_special_token_length:] = processed_image_bits
        
        selected_action_bits = action_bits[action_indexs, :]
        processed_action_bits = self.HIDEncoder(selected_action_bits)
        batched_rnn_input[action_indexs, datatype_and_special_token_length:] = processed_action_bits
        
        batched_rnn_input_reshaped = einops.rearrange(batched_rnn_input, "(b s) d -> b s d", b = batch_size, s = seq_length)
        
        # BUG 13: mismatched shape for rnn
        _, (h1, c1) = self.rnn(batched_rnn_input_reshaped)
        
        # shape of h1: [batch_size, dim_block]
        
        if target_output is not None:
            target_datatype_bits, target_special_bits, target_image_bits, target_action_bits = einops.unpack(target_output, [[s] for s in ConsciousBase.split_sizes], "b *")
        
        output_datatype_bits, output_special_bits, output_data_bits = einops.unpack(einops.rearrange(h1, "b s d -> (b s) d"), [[s] for s in [ConsciousBase.data_type_length, ConsciousBase.special_token_length, 1000]] , "b *")
#         output_datatype_bits, output_special_bits, output_data_bits = size_splits(h1, [ConsciousBase.data_type_length, ConsciousBase.special_token_length, 1000] ,dim=1)
        output_datatypes = torch.argmax(output_datatype_bits, dim=1)
        
        output_image_indexs = output_datatypes == 0
        output_action_indexs = output_datatypes == 1
        
        if target_output is not None:
            target_output_datatypes = torch.argmax(target_datatype_bits, dim=1)
            
            target_output_image_indexs = target_output_datatypes == 0
            target_output_action_indexs = target_output_datatypes == 1
            
            common_output_image_indexs = torch.logical_and(target_output_image_indexs, output_image_indexs)
            common_output_action_indexs = torch.logical_and(target_output_action_indexs, output_action_indexs)
            
            common_output_indexs = torch.logical_or(common_output_image_indexs, common_output_action_indexs)
            
            target_exclusive_output_image_indexs = torch.logical_and(target_output_image_indexs, torch.logical_not(output_image_indexs))
            target_exclusive_output_action_indexs = torch.logical_and(target_output_action_indexs, torch.logical_not(output_action_indexs))
            
            target_exclusive_output_indexs = torch.logical_or(target_exclusive_output_image_indexs, target_exclusive_output_action_indexs)
        
        if target_output is not None:
            selected_output_image_bits = output_data_bits[common_output_image_indexs, :]
        else:
            selected_output_image_bits = output_data_bits[output_image_indexs, :]
        processed_output_image_bits = einops.repeat(selected_output_image_bits, "b d -> b d 9")
        processed_output_image_bits = einops.einsum(processed_output_image_bits, processed_output_image_bits,"b h c, b w c -> b c h w")
        processed_output_image_bits = self.ViTDecoder(processed_output_image_bits)
        processed_output_image_bits = einops.rearrange(processed_output_image_bits, "b c h w -> b (c h w)")
        
        if target_output is not None:
            selected_output_action_bits = output_data_bits[common_output_action_indexs,:]
        else:
            selected_output_action_bits = output_data_bits[output_action_indexs,:]
        processed_output_action_bits = self.HIDDecoder(selected_output_action_bits)
        
        # preparing blank output
        output_0 = torch.zeros((batch_size, ConsciousBase.length))
        
        output_datatype_bits_0, output_special_token_bits_0, output_image_bits_0, output_action_bits_0 = einops.unpack(output_0, [[s] for s in ConsciousBase.split_sizes], "b *")
#         output_datatype_bits_0, output_special_token_bits_0, output_image_bits_0, output_action_bits_0 = size_splits(output, ConsciousBase.split_sizes, dim=1)
        
        if target_output is not None:
            output_datatype_bits_0[target_exclusive_output_indexs,:] = target_datatype_bits[target_exclusive_output_indexs,:]
            output_datatype_bits_0[common_output_indexs,:] = output_datatype_bits[common_output_indexs,:]
            
            output_special_token_bits_0[target_exclusive_output_indexs,:] = output_special_bits[target_exclusive_output_indexs,:]
            output_special_token_bits_0[common_output_indexs,:] = output_special_bits[common_output_indexs,:]
            
            output_special_token_bits_0[common_output_image_indexs,2] = 0
            output_special_token_bits_0[common_output_action_indexs,:2] = 0
            
            output_image_bits_0[target_exclusive_output_image_indexs,:] = target_image_bits[target_exclusive_output_image_indexs,:]
            output_action_bits_0[target_exclusive_output_action_indexs, :] = target_action_bits[target_exclusive_output_action_indexs,:]
        
            output_image_bits_0[common_output_image_indexs,:] = processed_output_image_bits
            output_action_bits_0[common_output_action_indexs, :] = processed_output_action_bits
        else:
            output_datatype_bits_0 = output_datatype_bits
            
            output_special_bits_0 = output_special_bits
            
            output_special_bits_0[output_image_indexs, 2] = 0
            output_special_bits_0[output_action_indexs, :2] = 0
        
            output_image_bits_0[output_image_indexs, :] = processed_output_image_bits
            output_action_bits_0[output_action_indexs, :] = processed_output_action_bits
        
        output, _ = einops.pack((output_datatype_bits_0, output_special_token_bits_0, output_image_bits_0, output_action_bits_0), "b *")
#         output = torch.concat((output_datatype_bits_0, output_special_token_bits_0, output_image_bits_0, output_action_bits_0), dim=1)
        
        return output


# ## training

# In[11]:


# train the model.
rnn_hidden_state = None # because this is needed for every session per video.

# hidden_size_rwkv = hidden_size_vit = 1024
# output_size = # output size shall be identical to ConsciousBase.length

model = CustomModel(vit_model)


# In[12]:


from torch.optim import Adam
# from torch.nn import BCELoss
# from torch.nn import MSELoss
from torch.nn import CrossEntropyLoss
lr = 0.00001
optimizer= Adam(model.parameters(), lr=lr)

# loss_fn = BCELoss(reduction='none')
# loss_fn = MSELoss()
# loss_fn = CrossEntropyLoss(reduction='sum')
loss_fn = CrossEntropyLoss(reduction='mean')


# In[13]:


import cv2

def resizeImage(im, desired_size):
    # im = cv2.imread(im_pth)
    old_size = im.shape[:2] # old_size is in (height, width) format

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

    # new_size should be in (width, height) format

    im = cv2.resize(im, (new_size[1], new_size[0]))

    delta_w = desired_size - new_size[1]
    delta_h = desired_size - new_size[0]
    top, bottom = delta_h//2, delta_h-(delta_h//2)
    left, right = delta_w//2, delta_w-(delta_w//2)

    color = [0, 0, 0]
    new_im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT,
        value=color)
    return new_im


# In[ ]:


# just for test.
# shall we do some backward?

# consider merging multiple ViT converted images to batch them?
# from PIL import Image

desired_size = 224*4

# for every image it adds 16 conscious block in total.
# it could be a huge array

# so we would append into some queue class, when hit the critical point, it will train.
from typing import Callable

import gc

class Trainer:
    def __init__(self, model, loss_fn, optimizer):
        self.model = model
        self.loss_fn = loss_fn
        self.optimizer = optimizer # shall be registered to model parameters.
        
    def step(self, batched_input, batched_output=None):
        # BUG 8: model forward keyword error
        # FIX 8: fixing keyword, adding default keyword argument
        model_output = self.model.forward(batched_input, target_output=batched_output)
        loss = self.loss_fn(model_output, batched_output)
        print("LOSS?")
        print(loss)
        
        # this loss is incorrect. shall use some argmax stuff.
        # to ensure that this thing is the thing that we want.
        loss.backward()
        self.optimizer.step()
        self.optimizer.zero_grad()
        


class SequentialTrainingQueue:
    def __init__(self, context_length:int, batch_size:int, trainer:Trainer):
        self.context_length = context_length
        self.batch_size = batch_size
        
        self.trainer = trainer
        
        self.max_critical_length = self.context_length+self.batch_size
        self.min_critical_length = self.context_length+1
        
        self.consciousVectors = []
        
    def enqueue(self, consciousBlock:ConsciousBlock, clear:bool=False):
        # change that into some tensor first.
        
        # BUG 3: consciousBlock has tensor output but not numpy ndarray
        # FIX 3: find and replace all "consciousBlock.to_nparray()" with "consciousBlock.to_tensor()"
#         print(consciousBlock)
#         print(type(consciousBlock))
        consciousVector = consciousBlock.to_tensor()
        self.consciousVectors.append(consciousVector)
        
        if not clear:
            # BUG 5: no "max_critical_point"
            # FIX 5: replace it with "max_critical_length"
            if len(self.consciousVectors) == self.max_critical_length:
                self.train()
        else:
            self.clear()
    
    def train(self, clear:bool=False):
        # train the model and clear the queue.
        # size of queue before: self.context_length+self.batch_size (should be? at least geq to self.length+1)
        # size of queue after training: self.context_length
        
        if len(self.consciousVectors) >= self.min_critical_length:
            batch_size = len(self.consciousVectors) - self.context_length
            # BUG 6: missing self.
            # FIX 6: adding self.
            
            # BUG 7: torch.Tensor conversion error
            # FIX 7: change "torch.Tensor([])" into einops.pack
            
#             print(self.consciousVectors)
            batched_input, _ = einops.pack([
#             batched_input = torch.Tensor([
                    einops.pack(self.consciousVectors[i:i+self.context_length],'* d')[0] for i in range(batch_size)
#                     torch.Tensor(self.consciousVectors[i:i+self.context_length]) for i in range(batch_size)
                ],'* s d')
            batched_output, _ = einops.pack([
#             batched_output = torch.Tensor([
                    self.consciousVectors[self.context_length+i] for i in range(batch_size)
                ], '* d')
            self.trainer.step(batched_input, batched_output)
            
            if not clear:
                # now remove some elements.
                self.consciousVectors = self.consciousVectors[-self.context_length:]
            else:
                self.consciousVectors = []
        
    def clear(self):
        # check if anything left in queue. call at the end of queue.
        self.train(clear=True)

# technically we would "shuffle" the dataset.
# but now, we just read it in sequence.

# make it small.
context_length = 10
batch_size = 3

trainer = Trainer(model, loss_fn, optimizer)

sequentialTrainingQueue = SequentialTrainingQueue(context_length, batch_size, trainer)

with VideoCaptureContextManager(videoPath) as cap:
    last_index = -1
    for screenshot_and_actions in data["screenshot_and_actions"]:
        screenshot = screenshot_and_actions["screenshot"]
        actions = screenshot_and_actions["actions"]
        print()
        print("SCREENSHOT:", screenshot)
        image_path = screenshot["imagePath"]
        image_size = screenshot["imageSize"]
        index = re.findall(r"\d+", image_path)[0]
        index = int(index)
        print("IMAGE INDEX:", index)
        
        encoded_actions = []

        for action in actions:
            action_type, action_args = action["HIDEvent"]
            if action_type in HIDActionBase.keyboard_action_types:
                action_args = HIDActionBase.uncover_keycode(action_args)
                if action_args is None:
                    print("Skipping")
                    continue
            mHIDAction = HIDAction.from_action_json(
                [action_type, action_args],
                max_x=perspective_width,
                max_y=perspective_height,
            )  # related to mouse coordinates.
            mHIDActionNDArray = mHIDAction.to_ndarray()
            print(mHIDActionNDArray.shape)


        if index != last_index:
            ret, image = cap.read()
#             print(image.shape) # 400, 640, 3
# shall be the lowest quality.
            # loop through all actions in this area.
            image_resized = resizeImage(image, desired_size)
            image_reshaped = einops.rearrange(image_resized, "h w c -> c h w")
#             image_reshaped = np.rollaxis(image_resized, 2, 0) # (3, 896, 896)
            image_sliced = [
                image_reshaped[:,\
                                                   x*224:(x+1)*224,\
                                                   y*224:(y+1)*224
                                                  ] for x in range(4) for y in range(4)
            ] #) # c h w
            
            # IMAGE RESHAPED: (896, 3, 896)?
            # IMAGE RESHAPED: (896, 896, 3)
#             print('IMAGE RESHAPED:', image_reshaped.shape)
#             print('IMAGE SLICED:', image_sliced.shape)
#     (16, 3, 224, 224)
# hell?
            for index, im in enumerate(image_sliced):
                im = einops.rearrange(im, "c h w -> (c h w)")
                st = None
                if index == 15:
                    st = "image_end"
                elif index !=0 and (index+1)%4 == 0:
                    st = "image_newline"
                    
                # BUG 4: tuple
                # FIX 4: remove .to_tensor() method call
                consciousBlock = ConsciousBlock(data_type='image', special_token=st, image_data=im)
#                 print(consciousBlock)
                sequentialTrainingQueue.enqueue(consciousBlock)
#             last_output = torch.zeros(1, output_size)
            for index, EA in enumerate(encoded_actions):
                st = None
                if index+1 == len(encoded_actions):
                    st = "action_end"
                consciousBlock = ConsciousBlock(data_type = "HIDAction", special_token=st, action_data = EA)
                sequentialTrainingQueue.enqueue(consciousBlock)
                
#                 print('ACTION SHAPE?', EA.to_ndarray().shape) # (4111, 1)
#                 if last_output is not None:
#                     last_output = last_output.detach()
#                 if rnn_hidden_state is not None:
#                     rnn_hidden_state = (rnn_hidden_state[0].detach(), rnn_hidden_state[1].detach())
#                 last_output, rnn_hidden_state = model.forward(torch.tensor(image_sliced).float(), last_output, rnn_hidden_state)
# #                 output, rnn_hidden_state = model.forward(instruction, pil_image, rnn_hidden_state)
#                 print("OUTPUT:", last_output.shape)
#                 print("HIDDEN_STATE:", rnn_hidden_state[0].shape, rnn_hidden_state[1].shape) # used for next action, if possible.
#                 target = torch.tensor(EA.to_ndarray().reshape(1,-1)).float()
#                 print("TARGET SHAPE:", target.shape)
#                 print("OUTPUT SHAPE:", last_output.shape)
#                 # what the heck?
# #                 TARGET SHAPE: torch.Size([1, 4111])
# #                 OUTPUT SHAPE: torch.Size([1, 4111])
#                 print(target)
#                 print(last_output)
#                 with torch.enable_grad():
                
            del image
            del image_reshaped
            del image_resized
            last_index = index
        else:
            pass
        
        # no need to jump to index. it is monotonic.
        # image = read_image_from_video_at_given_index_as_ndarray(index, cap)
        # SCREENSHOT: {'timeStamp': 1680247598.852561, 'imagePath': 'screenshots/498.raw', 'imageSize': [2560, 1600]}
        # print(image.shape)  # (1600, 2560, 3)
        # width, height, channel
        # breakpoint()
        
        print("ACTIONS:", actions)
        # keyboard action types: ['key_press', 'key_release']
        # mouse action types: ['mouse_move', 'mouse_click', 'mouse_scroll']
        # cv2.imshow("IMG", image)
        # cv2.waitKey(0)        
        # del image
        
        # not deleting the image. keep it!

        # now the data reader is complete. focus on the network design.
        # shall we?
sequentialTrainingQueue.clear()