import sys
import gc
from numpy.random import seed
seed(7)
import numpy as np
import glob
from utils import get_classes
from sklearn.model_selection import StratifiedShuffleSplit
import cv2
from keras.utils import to_categorical
import scipy.io as sio
from scipy.misc import imresize

def load_inputs(mode, test_subject, parameters, val_index=None):
    data_files_folder = parameters['data_files_folder']
    images_folder = parameters['images_folder']
    of_folder = parameters['of_folder']
    val_size = parameters['percentage_of_train_for_validation']
    classes_file = parameters['classes_file']
    test_evaluation_method = parameters['test_evaluation_method']
    L = parameters['L']
    of_mean = parameters['of_mean']
    image_mean = sio.loadmat(parameters['image_mean'])['image_mean']
    image_shape = (parameters['width'], parameters['height'])
    augment_data = parameters['apply_data_augmentation']
    _width, _height = parameters['dag_resize_shape']

    # Load the .txt where train/test partition are saved
    if mode == 'train' or mode == 'val':
        data_file = data_files_folder + '/train_{}.txt'.format(test_subject)
    elif mode == 'test':
        data_file = data_files_folder + '/test_{}.txt'.format(test_subject)
    else:
        print('No valid mode to load data. Options: train, val, test.')
        sys.exit()

    # Load video folders and their respective classes to memory
    folders, labels = [], []
    class_dict = get_classes(classes_file)
    with open(data_file, 'r') as f:
        content = f.readlines()
        for i in range(len(content)):
            folder, label = content[i].strip().split(' ')
            folders.append(folder)
            class_name = folder[folder.find('/')+1:folder.rfind('/')]
            labels.append(class_dict[class_name])

    # Stratify the dataset in train to get a validation partition:
    # i.e., get a subset of the training set with the same class distsribution
    if mode == 'train':
        sss = StratifiedShuffleSplit(
            n_splits=1, test_size=val_size, random_state=0
        )
        indices = sss.split(folders, labels)
        val_index = indices.next()[1]

    X, Y = [], []
    for i in range(len(folders)):
        folder = folders[i]
        class_name = folder[folder.find('/')+1:folder.rfind('/')]
        # Only include the elements with index inside val_index
        # for the validation set
        if mode == 'val':
            if i in val_index:
                X.append(folder)
                Y.append(int(labels[i]))
        # Do not include elements of the validation set in the training set
        elif mode == 'train':
            if not i in val_index:
                X.append(folder)
                Y.append(int(labels[i]))
        elif mode == 'test':
            X.append(folder)
            Y.append(int(labels[i]))
    del folders, labels
    gc.collect()

    folders_in_class = dict()
    class_of_folder = dict()
    for folder, label in zip(X, Y):
        temp = folder[:folder.rfind('/')]
        class_name = temp[temp.rfind('/')+1:]

        # Store the folders by class
        if not folders_in_class.has_key(class_name):
            folders_in_class[class_name] = []
        folders_in_class[class_name].append(folder)
        class_of_folder[folder] = label

    perm = np.random.permutation(len(folders_in_class.keys()))
    (batches_images, batches_stacks, batches_labels,
     inputs_per_video, video_names) = [], [], [], [],[]
    #batch_labels_video = []
    #video_durations = []
    nb_videos = 0

    # p contains a class index, randomized by perm
    for p in np.asarray(folders_in_class.keys())[list(perm)]:
        folders = folders_in_class[p]
        for element in folders:
            video_names.append(element)
            inputs_per_video.append(nb_videos)

            # Load images:
            frames = glob.glob(images_folder + element + '/frame*')
            # In case of less than L images, replicate them
            if len(frames) < L:
                reps = int(L/len(frames))
                temp = frames
                frames = []
                for i in range(len(temp)):
                    for _ in xrange(reps):
                        frames.append(temp[i])
            temp = []
            indices = []
            if test_evaluation_method != 0:
                for i in xrange(len(frames)):
                    img = cv2.imread(frames[i]) - image_mean
                    # subtract the image mean (used in the pre-training)
                    if mode == 'train' and augment_data:
                        img = imresize(img, (_width, _height), interp='bilinear')
                    temp.append(img)
            else:
                indices = list(np.asarray(np.linspace(0,len(frames)-1,25),dtype=np.int))
                for i in indices:
                    img = cv2.imread(frames[i]) - image_mean
                    img = imresize(img, (_width, _height), interp='bilinear')
                    temp.append(img)
            batches_images.append(temp)

            # Load optical flow and stack L horizontal and vertical images
            x_frames = glob.glob(of_folder + element + '/flow_x*')
            y_frames = glob.glob(of_folder + element + '/flow_y*')
            if len(frames) < L:
                reps = int(L/len(x_frames))
                temp_x, temp_y = x_frames, y_frames
                x_frames, y_frames = [], []
                for i in range(len(temp_x)):
                    for _ in xrange(reps):
                        x_frames.append(temp_x[i])
                        y_frames.append(temp_y[i])

            rest = len(x_frames) % L
            add = 0
            # In case of spare optical flow images, add one more stack
            # taking the last L images (there is some overlap with the
            # previous stack)
            if rest > 0:
                add = 1
            temp = []
            if test_evaluation_method != 0:
                for r in xrange((len(x_frames) // L) + add):
                    low, high = r*L, (r+1)*L
                    if high > len(x_frames):
                        low, high = -L, len(x_frames)
                    i = 0
                    if mode == 'train' and augment_data:
                        flow = np.zeros(shape=(_width, _height, 2*L,),
                                        dtype=np.float32)
                    else:
                        flow = np.zeros(shape=image_shape + (2*L,),
                                        dtype=np.float32)
                    for flow_x_file, flow_y_file in zip(
                            x_frames[low:high],y_frames[low:high]):
                        img_x = cv2.imread(flow_x_file, cv2.IMREAD_GRAYSCALE)
                        img_y = cv2.imread(flow_y_file, cv2.IMREAD_GRAYSCALE)
                        # Resize to larger size in order to crop them later
                        if mode == 'train' and augment_data:
                            img_x = imresize(img_x, (_width, _height),
                                             interp='bilinear')
                            img_y = imresize(img_y, (_width, _height),
                                             interp='bilinear')
                        flow[:,:,2*i] = img_x
                        flow[:,:,2*i+1] = img_y
                        i += 1
                    # Subtract the Optical Flow mean (used in the pre-training)
                    temp.append(flow - of_mean)
                    nb_videos += 1
                assert len(temp) == ((len(x_frames) // L)+add)
            else:
                for i in indices:
                    flow = np.zeros(shape=(_width, _height, 2*L,),
                                        dtype=np.float32)
                    n = 0
                    for j in range(i-5,i+4):
                        k = max(min(j,0),len(x_frames)-1)
                        img_x = cv2.imread(x_frames[k], cv2.IMREAD_GRAYSCALE)
                        img_y = cv2.imread(y_frames[k], cv2.IMREAD_GRAYSCALE)
                        img_x = imresize(img_x, (_width, _height),
                                             interp='bilinear')
                        img_y = imresize(img_y, (_width, _height),
                                         interp='bilinear')
                        flow[:,:,2*n] = img_x
                        flow[:,:,2*n+1] = img_y
                        n += 1
                    temp.append(flow - of_mean)
                    nb_videos += 1
            batches_stacks.append(temp)
            batches_labels.append(class_of_folder[element])
    if mode != 'test':
        return {'images': batches_images, 'stacks': batches_stacks,
                'labels': batches_labels, 'inputs_per_video': inputs_per_video,
                'val_index': val_index, 'video_names': video_names}
    else:
        return {'images': batches_images, 'stacks': batches_stacks,
                'labels': batches_labels, 'inputs': nb_videos}

def load_train_val_image_dataset(parameters, test_subject):
    training_set = load_inputs('train', test_subject, parameters)
    validation_set = load_inputs(
        'val', test_subject, parameters, val_index=training_set['val_index']
    )
    assert len(training_set['images']) == len(training_set['labels']), (
        'Number of inputs should be equal to the number of labels in '
        'training set but found: {} inputs and {} labels'.format(
            len(training_set['images']), len(training_set['labels'])
        )
    )
    assert len(validation_set['images']) == len(validation_set['labels']), (
        'Number of inputs should be equal to the number of labels in '
        'validation set but found: {} inputs and {} labels'.format(
        len(validation_set['images']), len(validation_set['labels']))
    )
    return training_set, validation_set

def load_test_image_dataset(parameters, test_subject):
    test_set = load_inputs('test', test_subject, parameters)
    assert len(test_set['images']) == len(test_set['labels']), (
        'Number of inputs should be equal to the number of labels in '
        'test set but found: {} inputs and {} labels'.format(
        len(test_set['images']), len(test_set['labels']))
    )
    return test_set

# This function was done by @rubensancor
def load_one_test_image(parameters, test_subject):
    data_files_folder = parameters['data_files_folder']
    images_folder = parameters['images_folder']
    of_folder = parameters['of_folder']
    L = parameters['L']
    of_mean = parameters['of_mean']
    image_mean = sio.loadmat(parameters['image_mean'])['image_mean']
    image_shape = (parameters['width'], parameters['height'])

    data_file = data_files_folder + '/test_{}.txt'.format(test_subject)
    with open(data_file, "r") as f:
        content = f.readlines()
        path = content[0]
        folder, label = path.strip().split(' ')

    img = cv2.imread(glob.glob(images_folder + folder + "/frame*")[5]) - image_mean

    of_path_x = glob.glob(of_folder + folder + "/flow_x*")[:10]
    of_path_y = glob.glob(of_folder + folder + "/flow_y*")[:10]
    i = 0
    flow = np.zeros(shape=image_shape + (2*L,), dtype=np.float32)
    for flow_x_file, flow_y_file in zip( of_path_x, of_path_y):
        img_x = cv2.imread(flow_x_file, cv2.IMREAD_GRAYSCALE)
        img_y = cv2.imread(flow_y_file, cv2.IMREAD_GRAYSCALE)
        flow[:,:,2*i] = img_x
        flow[:,:,2*i+1] = img_y
        i += 1
    of = flow - of_mean

    return img, of, label

def batch_generator(mode, parameters, dataset):
    test_evaluation_method = parameters['test_evaluation_method']
    nb_classes = parameters['nb_classes']
    L = parameters['L']
    batch_size = parameters['batch_size']
    augment_data = parameters['apply_data_augmentation']
    width, height = parameters['width'], parameters['height']
    _width, _height = parameters['dag_resize_shape']
    width_dif, height_dif = _width-width, _height-height
    evens = [n for n in range(20) if n % 2 == 0]

    while True:
        if mode == 'test':
            nb_inputs = dataset['inputs']
            # TEST METHOD 1: Get 25 images + optical flow stacks,
            # crop 4 corneres and center and get their mirrored version.
            # Majority voting with the 10*25 inputs obtained to classify the
            # video
            if test_evaluation_method == 0:
                for i in xrange(len(dataset['images'])):
                    batch_images, batch_ofs = [], []
                    for j in range(len(dataset['images'][i])):
                        crops = [[0,width,0,height],
                                 [0,width,height_dif,height_dif+height],
                                 [width_dif, width_dif+width,
                                  height_dif,height_dif+height],
                                 [width_dif,width_dif+width,0,height],
                                 [width_dif/2,width_dif/2+width,
                                  height_dif/2,height_dif/2+height]
                                ]
                        for crop in crops:
                            img = dataset['images'][i][j][
                                                crop[0]:crop[1], crop[2]:crop[3],:
                                             ]
                            stack = dataset['stacks'][i][j][
                                                crop[0]:crop[1], crop[2]:crop[3],:
                                            ]
                            for mirroring in [False, True]:
                                if mirroring:
                                    img = np.fliplr(img)
                                    stack[...,evens] = 255 - stack[...,evens]
                                batch_images.append(img)
                                batch_ofs.append(stack)

                    yield (np.asarray(batch_images),
                           np.asarray(batch_ofs),
                           np.asarray(
                               to_categorical(dataset['labels'][i], nb_classes)
                               )
                          )
            # TEST METHOD 2: Get all the inputs of the video and do
            # majority voting with all of them
            elif test_evaluation_method == 1:
                for i in xrange(len(dataset['stacks'])):
                    batch_images, batch_ofs = [], []
                    for j in range(len(dataset['stacks'][i])):
                        if j*L+L > len(dataset['images'][i]):
                            pos_image = -(L/2)
                        else:
                            pos_image = j*L + L/2
                        #print(len(dataset['stacks'][i]), len(dataset['images'][i]),
                        #      j, pos_image)
                        batch_images.append(dataset['images'][i][pos_image])
                        batch_ofs.append(dataset['stacks'][i][j])
                    yield (np.asarray(batch_images),
                           np.asarray(batch_ofs),
                           np.asarray(
                               to_categorical(dataset['labels'][i], nb_classes)
                               )
                          )
        # Train and val
        else:
            nb_inputs = dataset['inputs_per_video'][-1]
            perm = np.random.permutation(nb_inputs)
            batch_images, batch_ofs, batch_labels = [], [], []
            for i in xrange(nb_inputs):
                pos = np.searchsorted(
                    dataset['inputs_per_video'], perm[i], side='left'
                )
                if dataset['inputs_per_video'][pos] > perm[i]:
                    pos -= 1
                pos_stack = perm[i]-dataset['inputs_per_video'][pos]

                if pos_stack*L+L >= len(dataset['images'][pos]):
                    if mode == 'train':
                        inner_pos = (len(dataset['images'][pos]) - 1 -
                                    np.random.randint(0, L))
                    else:
                        inner_pos = len(dataset['images'][pos]) - 1 - 5
                else:
                    if mode == 'train':
                        inner_pos = np.random.randint(pos_stack*L,pos_stack*L+L)
                    else:
                        inner_pos = pos_stack*L + L/2

                # Data augmentation
                if mode == 'train' and augment_data:
                    # Random Cropping
                    dx = np.random.randint(0, _width-width)
                    dy = np.random.randint(0, _height-height)
                    img = dataset['images'][pos][inner_pos][
                                    dx:dx+width, dy:dy+height
                                ]
                    stack = dataset['stacks'][pos][pos_stack][
                                        dx:dx+width, dy:dy+height,:
                                    ]
                    # Random Mirroring
                    rand = np.random.rand(1)
                    if rand > 0.5:
                        img = np.fliplr(img)
                        stack[...,evens] = 255 - stack[...,evens]
                    batch_images.append(img)
                    batch_ofs.append(stack)
                    batch_labels.append(
                        to_categorical(dataset['labels'][pos], nb_classes)
                    )
                else:
                    batch_images.append(dataset['images'][pos][inner_pos])
                    batch_ofs.append(dataset['stacks'][pos][pos_stack])
                    batch_labels.append(
                        to_categorical(dataset['labels'][pos], nb_classes)
                    )
                if len(batch_images) == batch_size:
                    yield (np.asarray(batch_images), np.asarray(batch_ofs),
                           np.asarray(batch_labels))
                    batch_images, batch_ofs, batch_labels = [], [], []
            if mode == 'val' or mode == 'test':
                if len(batch_images) > 0:
                    yield (np.asarray(batch_images), np.asarray(batch_ofs),
                           np.asarray(batch_labels))
                    batch_images, batch_ofs, batch_labels = [], [], []
            elif mode == 'train' and nb_inputs % batch_size > 0:
                indices = np.random.choice(
                    range(nb_inputs), nb_inputs % batch_size
                )
                for i in indices:
                    pos = np.searchsorted(
                        dataset['inputs_per_video'], i, side='left'
                    )
                    if dataset['inputs_per_video'][pos] > i:
                        pos -= 1
                    pos_stack = i-dataset['inputs_per_video'][pos]
                    if pos_stack*L+L >= len(dataset['images'][pos]):
                        if mode == 'train':
                            inner_pos = (len(dataset['images'][pos]) - 1 -
                                        np.random.randint(0, L))
                        else:
                            inner_pos = len(dataset['images'][pos]) - 1 - 5
                    else:
                        if mode == 'train':
                            inner_pos = np.random.randint(
                                pos_stack*L,pos_stack*L+L
                            )
                        else:
                            inner_pos = pos_stack*L + L/2

                    # Data augmentation
                    if mode == 'train' and augment_data:
                        # Random Cropping
                        dx = np.random.randint(0, _width-width)
                        dy = np.random.randint(0, _height-height)
                        img = dataset['images'][pos][inner_pos][
                                        dx:dx+width, dy:dy+height
                                    ]
                        stack = dataset['stacks'][pos][pos_stack][
                                            dx:dx+width, dy:dy+height,:
                                        ]
                        # Random Mirroring
                        rand = np.random.rand(1)
                        if rand > 0.5:
                            img = np.fliplr(img)
                            stack[...,evens] = 255 - stack[...,evens]

                        batch_images.append(img)
                        batch_ofs.append(stack)
                        batch_labels.append(
                            to_categorical(dataset['labels'][pos], nb_classes)
                        )

                    else:
                        batch_images.append(dataset['images'][pos][inner_pos])
                        batch_ofs.append(dataset['stacks'][pos][pos_stack])
                        batch_labels.append(
                            to_categorical(dataset['labels'][pos], nb_classes)
                        )
                    if len(batch_images) == batch_size:
                        yield (np.asarray(batch_images), np.asarray(batch_ofs),
                               np.asarray(batch_labels))
                        batch_images, batch_ofs, batch_labels = [], [], []


from functools import partial
from kivy.uix.anchorlayout import AnchorLayout
from kivy.clock import Clock
from constants import green_color, red_color, maroon_color
from application.application_utils import create_grid_layout, create_label
from application.screens.games.base_game import BaseGame
from application.widgets.image_button import ImageButton

class NumberGapFillingGameScreen(BaseGame):

    def __init__(self, screen_name, sign_visual_dict, parent_function, rng=None, **kwargs):
        super().__init__(screen_name, sign_visual_dict, parent_function, rng, 'NumberGapFilling', **kwargs)
        self.current_index = 0
        self.__create_question()
        self.__build_screen()

    def __create_question(self):
        self.first_operand = self.rng.integers(0, 10)
        self.second_operand = self.rng.integers(1, 10)
        self.result = str(self.first_operand * self.second_operand)
        self.hidden_word = ' '.join(['_'] * len(self.result))

    def __build_screen(self):
        self.grid_layout = create_grid_layout(rows=2, spacing='5dp')
        self.layout.add_widget(self.grid_layout)
        self.__build_instructions()
        self.__build_buttons()

    def __build_instructions(self):
        self.grid_layout_text = create_grid_layout(rows=2, size_hint=(1, 0.1))
        self.grid_layout.add_widget(self.grid_layout_text)

        instructions = ('Select the signs in the correct order to complete the ' +
            f'result of: {self.first_operand} x {self.second_operand}')
        instructions_label = create_label(text=instructions, size_hint=(1, 0.4))
        self.grid_layout_text.add_widget(instructions_label)

        self.word_label = create_label(text=self.hidden_word, size_hint=(1, 0.6))
        self.grid_layout_text.add_widget(self.word_label)

    def __build_buttons(self):
        anchor_layout = AnchorLayout()
        self.grid_layout.add_widget(anchor_layout)

        self.grid_layout_signs = create_grid_layout(rows=2, spacing='7dp',
            cols=len(self.sign_visual_dict) // 2, size_hint=(1, 0.9))
        anchor_layout.add_widget(self.grid_layout_signs)

        options = list(self.sign_visual_dict.keys())
        self.rng.shuffle(options)
        for sign in options:
            self.__add_button(sign)

    def __add_button(self, sign):
        visual_path = self.sign_visual_dict[sign]
        button = ImageButton(on_press=self.button_callback, image_source=visual_path, size_hint_y=0.4)
        button.id = sign
        self.grid_layout_signs.add_widget(button)

    def button_callback(self, instance):
        if self.__correct_answer(instance.id):
            self.__trigger_color_change(instance, green_color)
            self.__update_hidden_word()
            if self.__is_game_finished():
                self.finish_game()
        else:
            self.num_mistakes += 1
            self.__trigger_color_change(instance, red_color)

    def __correct_answer(self, selected_sign):
        return selected_sign == self.result[self.current_index]

    def __update_hidden_word(self):
        index = 2 * self.current_index
        self.hidden_word = self.hidden_word[:index] + self.result[self.current_index] + self.hidden_word[index+1:]
        self.word_label.text = self.hidden_word
        self.current_index += 1

    def set_match_color_to_default(self, button_pressed, *args):
        button_pressed.border_color = maroon_color

    def __trigger_color_change(self, button_pressed, color):
        button_pressed.border_color = color
        partial_func = partial(self.set_match_color_to_default, button_pressed)
        Clock.schedule_once(partial_func, 0.2)

    def __is_game_finished(self):
        return self.current_index >= len(self.result)


from config import dino_pretrain_path
from .resnet import ResNet
from .vision_transformer import *
import torch

def get_backbone(args, **kwargs):
    if "resnet" in args.arch:
        return _get_resnet_backbone(args)
    elif "vit" in args.arch:
        return _get_vit_backbone(args, **kwargs)
    else:
        raise NotImplementedError("The arch has not implemented.")

def _get_resnet_backbone(args):
    backbone = ResNet(args.arch, args.low_res)
    return backbone

def _get_vit_backbone(args, drop_path_rate=0):
    vit_backbone = vit_base(drop_path_rate=drop_path_rate)
    # vit_backbone = vit_small(drop_path_rate=drop_path_rate)
    try:
        state_dict = torch.load(dino_pretrain_path, map_location='cpu')
        vit_backbone.load_state_dict(state_dict)
    except RuntimeError:
        print("Noting you are failed to load the pretrained model, we will load the dino pretrained model.")
    for m in vit_backbone.parameters():
        m.requires_grad = False

    # Only finetune layers from block 'args.grad_from_block' onwards
    for name, m in vit_backbone.named_parameters():
        if 'block' in name:
            block_num = int(name.split('.')[1])
            if block_num >= args.grad_from_block:
                m.requires_grad = True

    return vit_backbone


from construct import Struct, Int32ul, Int16ul, Int8ul, Nibble, Const, IfThenElse, this, Bytes, BitStruct

from formats.common import FlagsEnumCumulative

'''
typedef enum FLG_CONTIGUOUS_MIP {
    FLG_7ZIP = 0x40000000,
    FLG_ADDRU_MASK = 0xf,
    FLG_ADDRV_MASK = 0xf0,
    FLG_ARRTEX = 0x200000,
    FLG_COMPR_MASK = 0xe0000000,
    FLG_CONTIGUOUS_MIP = 0x100,
    FLG_CUBTEX = 0x800,
    FLG_GAMMA_EQ_1 = 0x8000,
    FLG_GENMIP_BOX = 0x2000,
    FLG_GENMIP_KAIZER = 0x4000,
    FLG_GLES3_TC_FMT = 0x100000,
    FLG_HASBORDER = 0x400,
    FLG_HOLD_SYSMEM_COPY = 0x10000,
    FLG_HQ_PART = 0x80000,
    FLG_NEED_PAIRED_BASETEX = 0x20000,
    FLG_NONPACKED = 0x200,
    FLG_OODLE = 0x60000000,
    FLG_REV_MIP_ORDER = 0x40000,
    FLG_VOLTEX = 0x1000,
    FLG_ZLIB = 0x80000000,
    FLG_ZSTD = 0x20000000
} FLG_CONTIGUOUS_MIP;
'''
ddsx_flags_enum = FlagsEnumCumulative(
    Int32ul,
    FLG_7ZIP=0x40_000_000,
    FLG_ADDRU_MASK=0xf,
    FLG_ADDRV_MASK=0xf0,
    FLG_ARRTEX=0x200_000,
    FLG_COMPR_MASK=0xe0_000_000,
    FLG_CONTIGUOUS_MIP=0x100,
    FLG_CUBTEX=0x800,
    FLG_GAMMA_EQ_1=0x8_000,
    FLG_GENMIP_BOX=0x2_000,
    FLG_GENMIP_KAIZER=0x4_000,
    FLG_GLES3_TC_FMT=0x100_000,
    FLG_HASBORDER=0x400,
    FLG_HOLD_SYSMEM_COPY=0x10_000,
    FLG_HQ_PART=0x80_000,
    FLG_NEED_PAIRED_BASETEX=0x20_000,
    FLG_NONPACKED=0x200,
    FLG_OODLE=0x60_000_000,
    FLG_REV_MIP_ORDER=0x40_000,
    FLG_VOLTEX=0x1_000,
    FLG_ZLIB=0x80_000_000,
    FLG_ZSTD=0x20_000_000
)

'''
struct Header {
    uint label;
    uint d3dFormat;
    uint flags;
    ushort w;
    ushort h;
    uchar levels;
    uchar hqPartLevels;
    ushort depth;
    ushort bitsPerPixel;
    uchar lQmip:4;
    uchar mQmip:4;
    uchar dxtShift:4;
    uchar uQmip:4;
    uint memSz;
    uint packedSz;
};
'''
ddsx_header = Struct(
    "label" / Const(b"DDSx"),
    "d3dFormat" / Bytes(4),
    # there error in parsing flags, fixed somewhere on future of construct, but not on 2.9.24
    "flags" / ddsx_flags_enum,
    "w" / Int16ul,
    "h" / Int16ul,
    "levels" / Int8ul,
    "hqPartLevels" / Int8ul,
    "depth" / Int16ul,
    "bitsPerPixel" / Int16ul,
    # this field added, from Embedded(BitStruct constrict package updating
    "bits_" / BitStruct(
        "lQmip" / Nibble,
        "mQmip" / Nibble,
        "dxtShift" / Nibble,
        "uQmip" / Nibble,
    ),
    "memSz" / Int32ul,
    "packedSz" / Int32ul,
)

ddsx = Struct(
    "header" / ddsx_header,
    "body" / IfThenElse(
        this.header.packedSz != 0,
        Bytes(this.header.packedSz),
        Bytes(this.header.memSz)
    )
)


import hashlib
import json
import os
import string
import sys

import typing

sys.path.append(os.getcwd())
import disk_set

def truncated_sha1(pwd):  # type: (typing.Text) -> bytes
    return hashlib.sha1(pwd.encode("utf8")).digest()[:8]

def brute_force(length, strings):
    # type: (int, typing.Sequence[typing.Text]) -> typing.Iterator[typing.Text]
    stack = [
        [s] for s in reversed(strings)
    ]  # type: typing.List[typing.List[typing.Text]]
    while stack:
        cur = stack.pop()  # type: typing.List[typing.Text]
        yield "".join(cur)
        if len(cur) < length:
            stack.extend(cur + [s] for s in reversed(strings))

with open("hash_set_truncated_size_847223402", "rb") as f:
    N = 847223402  # type: int
    with disk_set.DiskSet(f, N, 8) as ds:
        t = tuple(
            p
            for p in brute_force(3, string.ascii_lowercase)
            if truncated_sha1(p) not in ds
        )  # type: typing.Tuple[typing.Text, ...]
        print(json.dumps(t, separators=(",", ":")))


from typing import cast
from unittest.mock import AsyncMock, patch

import pytest

from starknet_py.contract import Contract
from starknet_py.hash.address import compute_address
from starknet_py.hash.selector import get_selector_from_name
from starknet_py.net.account.account import Account
from starknet_py.net.account.base_account import BaseAccount
from starknet_py.net.client_errors import ClientError
from starknet_py.net.client_models import (
    Call,
    DeployAccountTransaction,
    DeployAccountTransactionResponse,
    EstimatedFee,
    SierraContractClass,
    TransactionStatus,
)
from starknet_py.net.full_node_client import FullNodeClient
from starknet_py.net.gateway_client import GatewayClient
from starknet_py.net.models import StarknetChainId
from starknet_py.net.models.transaction import Declare, DeclareV2
from starknet_py.net.signer.stark_curve_signer import KeyPair
from starknet_py.net.udc_deployer.deployer import Deployer
from starknet_py.tests.e2e.fixtures.constants import MAX_FEE

@pytest.mark.run_on_devnet
@pytest.mark.asyncio
@pytest.mark.parametrize(
    "client",
    [GatewayClient(net="custom.net"), FullNodeClient(node_url="custom.net/rpc")],
)
async def test_get_balance_throws_when_token_not_specified(account, client):
    modified_account = Account(
        address=account.address,
        client=client,
        key_pair=KeyPair(1, 2),
        chain=cast(StarknetChainId, 1),
    )
    with pytest.raises(
        ValueError,
        match="Argument token_address must be specified when using a custom network.",
    ):
        await modified_account.get_balance()

@pytest.mark.asyncio
async def test_balance_when_token_specified(account, erc20_contract):
    balance = await account.get_balance(erc20_contract.address)

    assert balance == 200

@pytest.mark.asyncio
async def test_estimated_fee_greater_than_zero(erc20_contract, account):
    erc20_contract = Contract(
        address=erc20_contract.address, abi=erc20_contract.data.abi, provider=account
    )

    estimated_fee = (
        await erc20_contract.functions["balanceOf"]
        .prepare("1234", max_fee=0)
        .estimate_fee(block_hash="latest")
    )

    assert estimated_fee.overall_fee > 0
    assert (
        estimated_fee.gas_price * estimated_fee.gas_usage == estimated_fee.overall_fee
    )

@pytest.mark.asyncio
async def test_estimate_fee_for_declare_transaction(account, map_compiled_contract):
    declare_tx = await account.sign_declare_transaction(
        compiled_contract=map_compiled_contract, max_fee=MAX_FEE
    )

    estimated_fee = await account.client.estimate_fee(tx=declare_tx)

    assert isinstance(estimated_fee.overall_fee, int)
    assert estimated_fee.overall_fee > 0
    assert (
        estimated_fee.gas_usage * estimated_fee.gas_price == estimated_fee.overall_fee
    )

@pytest.mark.asyncio
@pytest.mark.parametrize("key, val", [(20, 20), (30, 30)])
async def test_sending_multicall(account, map_contract, key, val):
    calls = [
        map_contract.functions["put"].prepare(key=10, value=10),
        map_contract.functions["put"].prepare(key=key, value=val),
    ]

    res = await account.execute(calls=calls, max_fee=int(1e20))
    await account.client.wait_for_tx(res.transaction_hash)

    (value,) = await map_contract.functions["get"].call(key=key)

    assert value == val

@pytest.mark.run_on_devnet
@pytest.mark.asyncio
async def test_get_block_traces(gateway_account):
    traces = await gateway_account.client.trace_block_transactions(block_number=2)

    assert traces.traces != []

@pytest.mark.asyncio
async def test_rejection_reason_in_transaction_receipt(map_contract):
    with pytest.raises(ClientError, match=r".*INSUFFICIENT_MAX_FEE.*"):
        _ = await map_contract.functions["put"].invoke(key=10, value=20, max_fee=1)

def test_sign_and_verify_offchain_message_fail(account, typed_data):
    signature = account.sign_message(typed_data)
    signature = [signature[0] + 1, signature[1]]
    result = account.verify_message(typed_data, signature)

    assert result is False

def test_sign_and_verify_offchain_message(account, typed_data):
    signature = account.sign_message(typed_data)
    result = account.verify_message(typed_data, signature)

    assert result is True

@pytest.mark.asyncio
async def test_get_class_hash_at(map_contract, account):
    class_hash = await account.client.get_class_hash_at(
        map_contract.address, block_hash="latest"
    )

    assert class_hash != 0

@pytest.mark.asyncio()
async def test_get_nonce(account, map_contract):
    nonce = await account.get_nonce()
    address = map_contract.address
    block = await account.client.get_block(block_number="latest")

    tx = await account.execute(
        Call(
            to_addr=address, selector=get_selector_from_name("put"), calldata=[10, 20]
        ),
        max_fee=MAX_FEE,
    )
    await account.client.wait_for_tx(tx.transaction_hash)

    new_nonce = await account.get_nonce()
    new_nonce_latest_block = await account.get_nonce(block_number="latest")

    old_nonce = await account.get_nonce(block_number=block.block_number)

    assert isinstance(nonce, int) and isinstance(new_nonce, int)
    assert new_nonce == nonce + 1

    assert old_nonce == nonce
    assert new_nonce_latest_block == new_nonce

@pytest.mark.asyncio
@pytest.mark.parametrize(
    "calls", [[Call(10, 20, [30])], [Call(10, 20, [30]), Call(40, 50, [60])]]
)
async def test_sign_invoke_transaction(account, calls):
    signed_tx = await account.sign_invoke_transaction(calls, max_fee=MAX_FEE)

    assert isinstance(signed_tx.signature, list)
    assert len(signed_tx.signature) > 0
    assert signed_tx.max_fee == MAX_FEE

@pytest.mark.asyncio
async def test_sign_invoke_transaction_auto_estimate(account, map_contract):
    signed_tx = await account.sign_invoke_transaction(
        Call(map_contract.address, get_selector_from_name("put"), [3, 4]),
        auto_estimate=True,
    )

    assert isinstance(signed_tx.signature, list)
    assert len(signed_tx.signature) > 0
    assert signed_tx.max_fee > 0

@pytest.mark.asyncio
async def test_sign_declare_transaction(account, map_compiled_contract):
    signed_tx = await account.sign_declare_transaction(
        map_compiled_contract, max_fee=MAX_FEE
    )

    assert isinstance(signed_tx, Declare)
    assert signed_tx.version == 1
    assert isinstance(signed_tx.signature, list)
    assert len(signed_tx.signature) > 0
    assert signed_tx.max_fee == MAX_FEE

@pytest.mark.asyncio
async def test_sign_declare_transaction_auto_estimate(account, map_compiled_contract):
    signed_tx = await account.sign_declare_transaction(
        map_compiled_contract, auto_estimate=True
    )

    assert isinstance(signed_tx, Declare)
    assert signed_tx.version == 1
    assert isinstance(signed_tx.signature, list)
    assert len(signed_tx.signature) > 0
    assert signed_tx.max_fee > 0

@pytest.mark.asyncio
async def test_sign_declare_v2_transaction(
    account, sierra_minimal_compiled_contract_and_class_hash
):
    (
        compiled_contract,
        compiled_class_hash,
    ) = sierra_minimal_compiled_contract_and_class_hash

    signed_tx = await account.sign_declare_v2_transaction(
        compiled_contract,
        compiled_class_hash=compiled_class_hash,
        max_fee=MAX_FEE,
    )

    assert isinstance(signed_tx, DeclareV2)
    assert signed_tx.version == 2
    assert isinstance(signed_tx.signature, list)
    assert len(signed_tx.signature) > 0
    assert signed_tx.max_fee == MAX_FEE

@pytest.mark.asyncio
async def test_sign_declare_v2_transaction_auto_estimate(
    account, sierra_minimal_compiled_contract_and_class_hash
):
    (
        compiled_contract,
        compiled_class_hash,
    ) = sierra_minimal_compiled_contract_and_class_hash

    signed_tx = await account.sign_declare_v2_transaction(
        compiled_contract,
        compiled_class_hash=compiled_class_hash,
        auto_estimate=True,
    )

    assert isinstance(signed_tx, DeclareV2)
    assert signed_tx.version == 2
    assert isinstance(signed_tx.signature, list)
    assert len(signed_tx.signature) > 0
    assert signed_tx.max_fee > 0

@pytest.mark.asyncio
async def test_declare_contract_raises_on_sierra_contract_without_compiled_class_hash(
    sierra_minimal_compiled_contract_and_class_hash, account
):
    compiled_contract, _ = sierra_minimal_compiled_contract_and_class_hash
    with pytest.raises(
        ValueError,
        match="Signing sierra contracts requires using `sign_declare_v2_transaction` method.",
    ):
        await account.sign_declare_transaction(compiled_contract=compiled_contract)

@pytest.mark.asyncio
async def test_sign_deploy_account_transaction(account):
    class_hash = 0x1234
    salt = 0x123
    calldata = [1, 2, 3]
    signed_tx = await account.sign_deploy_account_transaction(
        class_hash, salt, calldata, max_fee=MAX_FEE
    )

    assert isinstance(signed_tx.signature, list)
    assert len(signed_tx.signature) > 0
    assert signed_tx.max_fee == MAX_FEE
    assert signed_tx.class_hash == class_hash
    assert signed_tx.contract_address_salt == salt
    assert signed_tx.constructor_calldata == calldata

@pytest.mark.asyncio
async def test_sign_deploy_account_transaction_auto_estimate(
    account, account_with_validate_deploy_class_hash
):
    class_hash = account_with_validate_deploy_class_hash
    salt = 0x1234
    calldata = [account.signer.public_key]
    signed_tx = await account.sign_deploy_account_transaction(
        class_hash, salt, calldata, auto_estimate=True
    )

    assert isinstance(signed_tx.signature, list)
    assert len(signed_tx.signature) > 0
    assert signed_tx.max_fee > 0
    assert signed_tx.class_hash == class_hash
    assert signed_tx.contract_address_salt == salt
    assert signed_tx.constructor_calldata == calldata

@pytest.mark.skip("No 0.12.1 devnet")
@pytest.mark.asyncio
async def test_deploy_account(client, deploy_account_details_factory, map_contract):
    address, key_pair, salt, class_hash = await deploy_account_details_factory.get()

    deploy_result = await Account.deploy_account(
        address=address,
        class_hash=class_hash,
        salt=salt,
        key_pair=key_pair,
        client=client,
        chain=StarknetChainId.TESTNET,
        max_fee=int(1e16),
    )
    await deploy_result.wait_for_acceptance()

    account = deploy_result.account

    assert isinstance(account, BaseAccount)
    assert account.address == address

    # Test making a tx
    res = await account.execute(
        calls=Call(
            to_addr=map_contract.address,
            selector=get_selector_from_name("put"),
            calldata=[30, 40],
        ),
        max_fee=MAX_FEE,
    )
    tx_receipt = await account.client.wait_for_tx(res.transaction_hash)

    assert tx_receipt.finality_status in (
        TransactionStatus.ACCEPTED_ON_L1,
        TransactionStatus.ACCEPTED_ON_L2,
    )

@pytest.mark.asyncio
async def test_deploy_account_raises_on_incorrect_address(
    client, deploy_account_details_factory
):
    address, key_pair, salt, class_hash = await deploy_account_details_factory.get()

    with pytest.raises(
        ValueError,
        match=f"Provided address {hex(0x111)} is different than computed address {hex(address)}",
    ):
        await Account.deploy_account(
            address=0x111,
            class_hash=class_hash,
            salt=salt,
            key_pair=key_pair,
            client=client,
            chain=StarknetChainId.TESTNET,
            max_fee=MAX_FEE,
        )

@pytest.mark.asyncio
@pytest.mark.parametrize(
    "call_contract, client",
    [
        (
            "starknet_py.net.gateway_client.GatewayClient.call_contract",
            "gateway_client",
        ),
        (
            "starknet_py.net.full_node_client.FullNodeClient.call_contract",
            "full_node_client",
        ),
    ],
)
async def test_deploy_account_raises_on_no_enough_funds(
    deploy_account_details_factory, call_contract, client, request
):
    client = request.getfixturevalue(client)
    address, key_pair, salt, class_hash = await deploy_account_details_factory.get()

    with patch(call_contract, AsyncMock()) as mocked_balance:
        mocked_balance.return_value = (0, 0)

        with pytest.raises(
            ValueError,
            match="Not enough tokens at the specified address to cover deployment costs",
        ):
            await Account.deploy_account(
                address=address,
                class_hash=class_hash,
                salt=salt,
                key_pair=key_pair,
                client=client,
                chain=StarknetChainId.TESTNET,
                max_fee=MAX_FEE,
            )

@pytest.mark.asyncio
@pytest.mark.parametrize(
    "client_method_path, client",
    [
        (
            "starknet_py.net.gateway_client.GatewayClient",
            "gateway_client",
        ),
        (
            "starknet_py.net.full_node_client.FullNodeClient",
            "full_node_client",
        ),
    ],
)
async def test_deploy_account_passes_on_enough_funds(
    deploy_account_details_factory, client_method_path, client, request
):
    address, key_pair, salt, class_hash = await deploy_account_details_factory.get()
    client = request.getfixturevalue(client)

    with patch(
        client_method_path + ".call_contract", AsyncMock()
    ) as mocked_balance, patch(
        client_method_path + ".deploy_account", AsyncMock()
    ) as mocked_deploy:
        mocked_balance.return_value = (0, 100)
        mocked_deploy.return_value = DeployAccountTransactionResponse(
            transaction_hash=0x1
        )

        await Account.deploy_account(
            address=address,
            class_hash=class_hash,
            salt=salt,
            key_pair=key_pair,
            client=client,
            chain=StarknetChainId.TESTNET,
            max_fee=MAX_FEE,
        )

# TODO (#1056): change this test to braavos account
@pytest.mark.skip(
    reason="'__validate_execute__' doesn't allow any other calldata than in the constructor"
)
@pytest.mark.asyncio
async def test_deploy_account_uses_custom_calldata(
    client, deploy_account_details_factory, fee_contract
):
    _, key_pair, salt, class_hash = await deploy_account_details_factory.get()
    calldata = [1, 2, 3, 4]
    address = compute_address(
        salt=salt,
        class_hash=class_hash,
        constructor_calldata=calldata,
        deployer_address=0,
    )

    res = await fee_contract.functions["transfer"].invoke(
        recipient=address, amount=int(1e16), max_fee=MAX_FEE
    )
    await res.wait_for_acceptance()

    deploy_result = await Account.deploy_account(
        address=address,
        class_hash=class_hash,
        salt=salt,
        key_pair=key_pair,
        client=client,
        chain=StarknetChainId.TESTNET,
        constructor_calldata=calldata,
        max_fee=int(1e16),
    )

    tx = await client.get_transaction(deploy_result.hash)
    assert isinstance(tx, DeployAccountTransaction)
    assert tx.constructor_calldata == calldata

@pytest.mark.asyncio
async def test_sign_invoke_tx_for_fee_estimation(account, map_contract):
    call = map_contract.functions["put"].prepare(key=40, value=50)
    transaction = await account.sign_invoke_transaction(calls=call, max_fee=MAX_FEE)

    estimate_fee_transaction = await account.sign_for_fee_estimate(transaction)

    estimation = await account.client.estimate_fee(estimate_fee_transaction)
    assert estimation.overall_fee > 0

    # Verify that the transaction signed for fee estimation cannot be sent
    with pytest.raises(ClientError):
        await account.client.send_transaction(estimate_fee_transaction)

    # Verify that original transaction can be sent
    result = await account.client.send_transaction(transaction)
    await account.client.wait_for_tx(result.transaction_hash)

@pytest.mark.asyncio
async def test_sign_declare_tx_for_fee_estimation(account, map_compiled_contract):
    transaction = await account.sign_declare_transaction(
        compiled_contract=map_compiled_contract, max_fee=MAX_FEE
    )

    estimate_fee_transaction = await account.sign_for_fee_estimate(transaction)

    estimation = await account.client.estimate_fee(estimate_fee_transaction)
    assert estimation.overall_fee > 0

    # Verify that the transaction signed for fee estimation cannot be sent
    with pytest.raises(ClientError):
        await account.client.declare(estimate_fee_transaction)

    # Verify that original transaction can be sent
    result = await account.client.declare(transaction)
    await account.client.wait_for_tx(result.transaction_hash)

@pytest.mark.asyncio
async def test_sign_deploy_account_tx_for_fee_estimation(
    client, deploy_account_details_factory
):
    address, key_pair, salt, class_hash = await deploy_account_details_factory.get()

    account = Account(
        address=address,
        client=client,
        key_pair=key_pair,
        chain=StarknetChainId.TESTNET,
    )

    transaction = await account.sign_deploy_account_transaction(
        class_hash=class_hash,
        contract_address_salt=salt,
        constructor_calldata=[key_pair.public_key],
        max_fee=MAX_FEE,
    )

    estimate_fee_transaction = await account.sign_for_fee_estimate(transaction)

    estimation = await account.client.estimate_fee(transaction)
    assert isinstance(estimation, EstimatedFee)
    assert estimation.overall_fee > 0

    # Verify that the transaction signed for fee estimation cannot be sent
    with pytest.raises(ClientError):
        await account.client.deploy_account(estimate_fee_transaction)

    # Verify that original transaction can be sent
    result = await account.client.deploy_account(transaction)
    await account.client.wait_for_tx(result.transaction_hash)

@pytest.mark.asyncio
async def test_sign_transaction_custom_nonce(account, cairo1_hello_starknet_class_hash):
    deployment = Deployer().create_contract_deployment(cairo1_hello_starknet_class_hash)
    deploy_tx = await account.sign_invoke_transaction(deployment.call, max_fee=MAX_FEE)

    new_balance = 30
    invoke_tx = await account.sign_invoke_transaction(
        Call(
            deployment.address,
            get_selector_from_name("increase_balance"),
            [new_balance],
        ),
        nonce=deploy_tx.nonce + 1,
        max_fee=MAX_FEE,
    )

    deploy_res = await account.client.send_transaction(deploy_tx)
    invoke_res = await account.client.send_transaction(invoke_tx)

    await account.client.wait_for_tx(deploy_res.transaction_hash)
    await account.client.wait_for_tx(invoke_res.transaction_hash)

    result = await account.client.call_contract(
        Call(deployment.address, get_selector_from_name("get_balance"), [])
    )

    assert invoke_tx.nonce == deploy_tx.nonce + 1
    assert result == [new_balance]

@pytest.mark.asyncio
async def test_argent_cairo1_account_deploy(
    full_node_client,
    argent_cairo1_account_class_hash,
    deploy_account_details_factory,
):
    address, key_pair, salt, class_hash = await deploy_account_details_factory.get(
        class_hash=argent_cairo1_account_class_hash, argent_calldata=True
    )

    deploy_result = await Account.deploy_account(
        address=address,
        class_hash=class_hash,
        salt=salt,
        key_pair=key_pair,
        client=full_node_client,
        constructor_calldata=[key_pair.public_key, 0],
        chain=StarknetChainId.TESTNET,
        max_fee=int(1e16),
    )
    await deploy_result.wait_for_acceptance()
    account = deploy_result.account

    assert isinstance(account, BaseAccount)
    assert await account.cairo_version == 1

    account_contract_class = await full_node_client.get_class_at(
        contract_address=account.address, block_number="latest"
    )

    assert isinstance(account_contract_class, SierraContractClass)

@pytest.mark.asyncio
async def test_argent_cairo1_account_execute(
    deployed_balance_contract,
    argent_cairo1_account: BaseAccount,
):
    # verify that initial balance is 0
    get_balance_call = Call(
        to_addr=deployed_balance_contract.address,
        selector=get_selector_from_name("get_balance"),
        calldata=[],
    )
    get_balance = await argent_cairo1_account.client.call_contract(
        call=get_balance_call, block_number="latest"
    )

    assert get_balance[0] == 0

    value = 20
    increase_balance_by_20_call = Call(
        to_addr=deployed_balance_contract.address,
        selector=get_selector_from_name("increase_balance"),
        calldata=[value],
    )
    execute = await argent_cairo1_account.execute(
        calls=increase_balance_by_20_call, max_fee=int(1e16)
    )
    await argent_cairo1_account.client.wait_for_tx(tx_hash=execute.transaction_hash)
    receipt = await argent_cairo1_account.client.get_transaction_receipt(
        tx_hash=execute.transaction_hash
    )

    # TODO (#1179): devnet 0.3.0 still return STATUS instead of FINALITY_STATUS
    assert receipt.status == TransactionStatus.ACCEPTED_ON_L2

    # verify that the previous call was executed
    get_balance_call = Call(
        to_addr=deployed_balance_contract.address,
        selector=get_selector_from_name("get_balance"),
        calldata=[],
    )
    get_balance = await argent_cairo1_account.client.call_contract(
        call=get_balance_call, block_number="latest"
    )

    assert get_balance[0] == value

# TODO (#1184): remove that
@pytest.mark.asyncio
async def test_cairo1_account_deprecations(
    deployed_balance_contract,
    argent_cairo1_account: BaseAccount,
):
    call = Call(
        to_addr=deployed_balance_contract.address,
        selector=get_selector_from_name("increase_balance"),
        calldata=[20],
    )
    with pytest.warns(
        DeprecationWarning,
        match="Parameter 'cairo_version' has been deprecated. It is calculated automatically based on your account's "
        "contract class.",
    ):
        _ = await argent_cairo1_account.execute(
            calls=call, max_fee=int(1e16), cairo_version=1
        )
        _ = await argent_cairo1_account.sign_invoke_transaction(
            calls=call, max_fee=int(1e16), cairo_version=1
        )


"""raytracing helper functions for deferred rendering using pyembree as a backend

Currently, meshes are represented as simple dictionaries with vertices, faces, uvs and optionally additional per-vertex
attributes (e.g., tangents, bitangents and normals). This is due to a limitation in the trimesh.Trimesh class, namely
that texture coordinates can only be specified per vertex and not per face, preventing proper unwrapping. Inputs of type
trimesh.Trimesh or trimesh.Scene are therefore converted to dictionaries or lists of dictionaries with the above keys.

The most important function is embree_render_deferred():

embree_render_deferred() produces inputs for deferred shading. It takes a list of meshes and a camera dict, as well as a
point light position as input and returns np.ndarray buffers with all relevant geometric quantities per intersected
pixel: 3D intersection points, interpolated vertex normals and tangents, normalized local and unnormalized global light
and view directions.

For additional light sources, get_local_dirs() can be used with the buffers returned from embree_render_deferred() and
a light position.

embree_render_deferred() calls the following functions: embree_create_scene(), embree_intersect_scene(),
interpolate_vertex_attributes() and get_local_dirs().
- embree_create_scene() constructs an EmbreeScene object, given the list of meshes. Camera parameters are either user-
specified, or will be set automatically (random camera position facing the scene center). Unless specified, the camera
focal length is automatically set so that the scene tightly fits onto the camera sensor. Unless explicitly disabled, it
computes per-vertex tangent frames, using each mesh's texture coordinates.
- embree_intersect_scene() performs the actual ray tracing and returns a pixel mask with the camera's resolution
indicating which pixels are hit, as well as buffers for all intersected pixels with the following attributes:
geometry and triangle IDs for each intersection, intersection depth (distance from camera), 3D intersection point,
barycentric coordinates within each triangle
- interpolate_vertex_attributes() computes per pixel texture coordinates and tangent frames by interpolating with the
barycentric coordinates returned from the embree_intersect_scene()

missing features:
- tracing shadow rays
- camera distortion model
- fallback tangent computing for meshes without texture coordinates
"""
from copy import deepcopy
import numpy as np
import trimesh.visual.texture
from warnings import warn
from pyembree.rtcore_scene import EmbreeScene
from pyembree.mesh_construction import TriangleMesh

from pysmtb.image import assign_masked
from pysmtb.utils import Dct, find_dim, dims_execpt, safe_divide

try:
    # we only optionally depend on trimesh, a lot of functionality also works without
    from trimesh import Scene, Trimesh
except ModuleNotFoundError:
    from types import NoneType
    Scene = NoneType
    Trimesh = NoneType

def normalize(vec, axis=0):
    """normalize input array along specified dimension"""
    return safe_divide(vec, np.linalg.norm(vec, axis=axis, keepdims=True))

def get_vertices(mesh, join_geometries=False):
    """given an input mesh (trimesh.Trimesh, trimesh.Scene, dict), return the vertices as NV x 3 array"""
    if isinstance(mesh, Trimesh):
        vertices = mesh.vertices
    elif isinstance(mesh, Scene):
        if not join_geometries and len(mesh.geometry) > 1:
            raise Exception('mesh has multiple geometries, set join_geometries=True')
        vertices = np.concatenate([g.vertices for g in mesh.geometry.values()], axis=0)
    elif isinstance(mesh, dict):
        if not hasattr(mesh, 'vertices'):
            raise Exception('mesh dict should have field vertices')
        vertices = mesh['vertices']
    else:
        raise Exception('unsupported input type: ' + str(type(mesh)))
    return vertices

def get_normals(mesh, join_geometries=False):
    """given an input mesh (trimesh.Trimesh, trimesh.Scene, dict), return the vertex normals as NV x 3 array"""
    if isinstance(mesh, Trimesh):
        normals = mesh.normals
    elif isinstance(mesh, Scene):
        if not join_geometries and len(mesh.geometry) > 1:
            raise Exception('mesh has multiple geometries, set join_geometries=True')
        normals = np.concatenate([g.normals for g in mesh.geometry.values()], axis=0)
    elif isinstance(mesh, dict):
        if not hasattr(mesh, 'normals'):
            raise Exception('mesh dict should have field normals')
        normals = mesh['normals']
    else:
        raise Exception('unsupported input type: ' + str(type(mesh)))
    return normals

def get_vertex_normals(mesh):
    """return per-vertex normal vectors, which are usually obtained by weighted averaging of each vertex's surrounding
    face normals"""
    if isinstance(mesh, Trimesh):
        vertex_normals = mesh.vertex_normals
    elif isinstance(mesh, Scene):
        vertex_normals = np.concatenate([g.vertex_normals for g in mesh.geometry.values()], axis=0)
    elif isinstance(mesh, dict):
        if not hasattr(mesh, 'vertex_normals'):
            raise Exception('mesh dict should have field vertex_normals')
        vertex_normals = mesh['vertex_normals']
    else:
        raise Exception('unsupported input type: ' + str(type(mesh)))
    return vertex_normals

def get_uv_coords(mesh, join_geometries=False):
    """given an input mesh (trimesh.Trimesh, trimesh.Scene, dict), return the texture coordinates as NV x 2 array"""
    if isinstance(mesh, Trimesh) and hasattr(mesh, 'visual') and hasattr(mesh.visual, 'uv'):
        uvs = mesh.visual.uv
    elif isinstance(mesh, Trimesh) and 'ply_raw' in mesh.metadata.keys():
        props = mesh.metadata['ply_raw']['vertex']['properties'].keys()
        u_key = 'u' if 'u' in props else 's'
        v_key = 'v' if 'v' in props else 't'
        u = mesh.metadata['ply_raw']['vertex']['data'][u_key]
        v = mesh.metadata['ply_raw']['vertex']['data'][v_key]
        uvs = np.stack((u, v), axis=1)
    elif isinstance(mesh, Scene):
        if not join_geometries and len(mesh.geometry) > 1:
            raise Exception('mesh has multiple geometries, set join_geometries=True')
        uvs = np.concatenate([g.visual.uv for g in mesh.geometry.values()], axis=0)
    elif isinstance(mesh, dict):
        if not hasattr(mesh, 'uvs'):
            raise Exception('mesh dict should have field uvs')
        uvs = mesh['uvs']
    else:
        raise Exception('unsupported input type: ' + str(type(mesh)))
    return uvs

def get_faces(mesh, join_geometries=False):
    """given an input mesh (trimesh.Trimesh, trimesh.Scene, dict), return the faces as NF x 3 array"""
    if isinstance(mesh, Trimesh):
        faces = mesh.faces
    elif isinstance(mesh, Scene):
        if not join_geometries and len(mesh.geometry) > 1:
            raise Exception('mesh has multiple geometries, set join_geometries=True')
        raise NotImplementedError('join_geometries is not supported for faces')
    elif isinstance(mesh, dict):
        if not hasattr(mesh, 'faces'):
            raise Exception('mesh dict should have field faces')
        faces = mesh['faces']
    else:
        raise Exception('unsupported input type: ' + str(type(mesh)))

    if faces.ndim == 1:
        faces = np.stack([face[1] for face in faces], axis=0)
    if faces.shape[1] == 4:
        faces = np.concatenate((faces[:, :3], faces[:, [0, 2, 3]]), axis=0)
    return faces

def get_tangent_frames(mesh):
    """compute tangent frames per vertex, returns the mesh as dict, even if input was a Trimesh"""
    # algorithm:
    # source: Eric Lengyel, Foundations of Game Engine Development - Volume 2: Rendering
    # originally under http://www.terathon.com/code/tangent.html
    # https://web.archive.org/web/20190211192552/http://www.terathon.com/code/tangent.html
    # given triangle vertices p_0, p_1, p_2 (counterclockwise) with corresponding texture coordinates (u_i, v_i)
    # it holds: p_i − p_j = (u_i − u_j) * t + (v_i − v_j) * b
    #
    # define:
    # e_1 = p_1 - p_0,   (x_1, y_1) = (u_1 - u_0, v_1 - v_0)
    # e_2 = p_2 - p_0,   (x_2, y_2) = (u_2 - u_0, v_2 - v_0)
    #
    # we solve for tangent and bitangent vectors t and b using a system of equations:
    # e_1 = x_1 * t + y_1 * b
    # e_2 = x_2 * t + y_2 * b
    #
    # [t, b] = 1 / (x_1 * y_2 - x_2 * y_1) * [e_1, e_2] * [[y_2, -x_2], [-y_1, x_1]]
    #
    # once solved, compute per-vertex tangents and bitangents by summing up the vectors from each triangle and
    # re-orthonormalizing, given per-vertex normal vectors:
    # t' = normalize(t - dot(t, n) * n)
    # b' = normalize(b - dot(b, n) * n - dot(b, t') * t')

    # convert trimesh inputs to dicts
    if isinstance(mesh, Trimesh):
        mesh = trimesh_to_dict(mesh)

    if isinstance(mesh, Scene):
        meshes = trimesh_scene_to_dicts(mesh)
        for mesh_id in range(len(meshes)):
            meshes[mesh_id] = get_tangent_frames(meshes[mesh_id])
        return meshes

    if hasattr(mesh, 'vertex_tangent_frames'):
        # tangent space already computed
        return mesh

    vertices = get_vertices(mesh)
    faces = get_faces(mesh)
    uvs = get_uv_coords(mesh)
    vertex_tangents = np.zeros_like(vertices)
    vertex_bitangents = np.zeros_like(vertices)
    vertex_normals = mesh.vertex_normals

    face_verts = vertices[faces.T, :]
    face_uvs = uvs[faces.T, :]

    e1 = face_verts[1] - face_verts[0]
    e2 = face_verts[2] - face_verts[0]

    x1 = face_uvs[1, :, 0:1] - face_uvs[0, :, 0:1]
    y1 = face_uvs[1, :, 1:2] - face_uvs[0, :, 1:2]
    x2 = face_uvs[2, :, 0:1] - face_uvs[0, :, 0:1]
    y2 = face_uvs[2, :, 1:2] - face_uvs[0, :, 1:2]

    r = safe_divide(1, (x1 * y2 - x2 * y1))
    t = (e1 * y2 - e2 * y1) * r
    b = (e2 * x1 - e1 * x2) * r

    # sum (bi)tangents that are associated to a vertex from its neighboring triangles
    vertex_tangents[faces[:, 0], :] += t
    vertex_tangents[faces[:, 1], :] += t
    vertex_tangents[faces[:, 2], :] += t
    vertex_bitangents[faces[:, 0], :] += b
    vertex_bitangents[faces[:, 1], :] += b
    vertex_bitangents[faces[:, 2], :] += b

    # normalize summed vectors
    vertex_tangents = normalize(vertex_tangents, axis=1)
    vertex_bitangents = normalize(vertex_bitangents, axis=1)

    # re-orthonormalization
    vertex_tangents = normalize(vertex_tangents - np.sum(vertex_normals * vertex_tangents, axis=1, keepdims=True) * vertex_normals, axis=1)
    vertex_bitangents = normalize(vertex_bitangents - np.sum(vertex_bitangents * vertex_normals, axis=1, keepdims=True) * vertex_normals
                                  - np.sum(vertex_bitangents * vertex_tangents, axis=1, keepdims=True) * vertex_tangents, axis=1)

    # determine handedness of tangent frame
    sgn_mask = np.sum(np.cross(vertex_tangents, vertex_bitangents, axis=1) * vertex_normals, axis=1, keepdims=True) > 0
    det_sgn = np.ones_like(sgn_mask, dtype=np.float32)
    det_sgn[~sgn_mask] = -1.

    # NV x 3 x 3 array (tangent, bitangent & normal stacked in rows)
    matrix_to_tangent_space = np.stack((vertex_tangents, det_sgn * vertex_bitangents, vertex_normals), axis=1)

    mesh.vertex_tangent_frames = matrix_to_tangent_space
    mesh.det_sgn = det_sgn
    return mesh

def create_unit_sphere(res=100):
    """create sphere centered on origin with radius 1 and res many segments in azimuth axis"""
    thetas = np.linspace(0, np.pi, res // 2)
    phis = np.linspace(0, 2 * np.pi, res)
    phis, thetas = np.meshgrid(phis, thetas)
    cp = np.cos(phis)
    sp = np.sin(phis)
    ct = np.cos(thetas)
    st = np.sin(thetas)
    xyz = np.stack((cp * st, sp * st, ct), axis=2).reshape((-1, 3))
    normals = xyz
    uvs = np.stack((phis / (2 * np.pi), thetas / np.pi), axis=2).reshape((-1, 2))

    nv = res * (res // 2)
    faces = np.r_[:nv].reshape(res // 2, res, order='C')

    # # faces:
    # # 0, 1, 2,
    # # 3, 4, 5,
    # # 6, 7, 8
    # faces = np.stack((
    #     faces[:-1],  # 0, 1, 2,
    #     np.roll(faces, -1, axis=0)[:-1],  # 3, 4, 5,
    #     np.roll(faces, -1, axis=1)[:-1],  # 1, 2, 0
    #     np.roll(np.roll(faces, -1, axis=0)[:-1], -1, axis=1),  # 4, 5, 3
    #     np.roll(faces, -1, axis=1)[:-1],  # 1, 2, 0
    #     np.roll(faces, -1, axis=0)[:-1],  # 3, 4, 5
    # ), axis=2).reshape((-1, 3))

    # faces:
    # 0, 1, 2,
    # 3, 4, 5,
    # 6, 7, 8
    faces = np.stack((
        faces[:-1],  # 0, 1, 2,
        np.roll(faces, -1, axis=0)[:-1],  # 3, 4, 5,
        np.roll(faces, -1, axis=1)[:-1],  # 1, 2, 0
        np.roll(np.roll(faces, -1, axis=0)[:-1], -1, axis=1),  # 4, 5, 3
        np.roll(faces, -1, axis=1)[:-1],  # 1, 2, 0
        np.roll(faces, -1, axis=0)[:-1],  # 3, 4, 5
    ), axis=2).reshape((-1, 3))

    # mesh = Trimesh(vertices=xyz.reshape((-1, 3)), faces=faces.reshape((-1, 3)))
    # mesh.show()

    return Dct(vertices=xyz, uvs=uvs, faces=faces, vertex_normals=normals)

def create_unit_cube():
    """create cube centered on origin with side length 1"""
    # x: side, y: front-back, z: up-down
    xyz = np.array([[0, 0, 0],  # 0, bottom
                    [1, 0, 0],  # 1
                    [1, 1, 0],  # 2
                    [0, 1, 0],  # 3
                    [0, 0, 0],  # 4 front
                    [1, 0, 0],  # 5
                    [1, 0, 1],  # 6
                    [0, 0, 1],  # 7
                    [1, 0, 0],  # 8 right
                    [1, 1, 0],  # 9
                    [1, 1, 1],  # 10
                    [1, 0, 1],  # 11
                    [1, 1, 0],  # 12 back
                    [0, 1, 0],  # 13
                    [0, 1, 1],  # 14
                    [1, 1, 1],  # 15
                    [0, 1, 0],  # 16 left
                    [0, 0, 0],  # 17
                    [0, 0, 1],  # 18
                    [0, 1, 1],  # 19
                    [0, 0, 1],  # 20 top
                    [1, 0, 1],  # 21
                    [1, 1, 1],  # 22
                    [0, 1, 1],  # 23
                    ], dtype=np.float32)
    normals = np.array(([[ 0,  0, -1]] * 4) +  # bottom
                       ([[ 0, -1,  0]] * 4) +  # front
                       ([[ 1,  0,  0]] * 4) +  # right
                       ([[ 0,  1,  0]] * 4) +  # back
                       ([[-1,  0,  0]] * 4) +  # left
                       ([[ 0,  0,  1]] * 4)    # top
                       )
    uvs = np.array([[0.25, 0.375],  # 0, bottom
                    [0.50, 0.375],  # 1
                    [0.50, 0.125],  # 2
                    [0.25, 0.125],  # 3
                    [0.25, 0.375],  # 4 front
                    [0.50, 0.375],  # 5
                    [0.50, 0.625],  # 6
                    [0.25, 0.625],  # 7
                    [0.50, 0.375],  # 8 right
                    [0.75, 0.375],  # 9
                    [0.75, 0.625],  # 10
                    [0.50, 0.625],  # 11
                    [0.75, 0.375],  # 12 back
                    [1.00, 0.375],  # 13
                    [1.00, 0.625],  # 14
                    [0.75, 0.625],  # 15
                    [0.00, 0.375],  # 16 left
                    [0.25, 0.375],  # 17
                    [0.25, 0.625],  # 18
                    [0.00, 0.625],  # 19
                    [0.25, 0.625],  # 20 top
                    [0.50, 0.625],  # 21
                    [0.50, 0.875],  # 22
                    [0.25, 0.875],  # 23
                    ])

    # from pysmtb.plotting import text3
    # th = text3(xyz)
    xyz = xyz - np.r_[0.5, 0.5, 0.5][None]

    faces = np.array([
        [0, 2, 1],  # bottom
        [2, 0, 3],
        [4, 5, 6],  # front
        [6, 7, 4],
        [8, 9, 10],  # right
        [10, 11, 8],
        [12, 13, 14],  # back
        [14, 15, 12],
        [16, 17, 18],  # left
        [18, 19, 16],
        [20, 21, 22],  # top
        [22, 23, 20],
                      ])
    # mesh = Trimesh(vertices=xyz.reshape((-1, 3)), faces=faces.reshape((-1, 3)))
    # mesh.show()

    return Dct(vertices=xyz, uvs=uvs, faces=faces, vertex_normals=normals)

def get_bbox(inp, dims: tuple = None):
    """get bounding box of vertices"""
    bbox = None
    if isinstance(inp, Scene):
        return inp.bounds

    if isinstance(inp, list):
        for mesh in inp:
            bbox_ = get_bbox(mesh)
            bbox = bbox_ if bbox is None else  np.array([np.minimum(bbox[0, :], bbox_[0, :]),
                                                         np.maximum(bbox[1, :], bbox_[1, :])])
        return bbox

    if isinstance(inp, Trimesh):
        inp = inp.vertices
    elif isinstance(inp, dict):
        inp = inp['vertices']
    elif not isinstance(inp, np.ndarray):
        raise Exception('input must be trimesh.Trimesh, trimesh.Scene or list of dicts with vertices key')

    if dims is None:
        dims = dims_execpt(inp, find_dim(inp, size=3))
    lower = np.min(inp, axis=dims).squeeze()
    upper = np.max(inp, axis=dims).squeeze()
    bbox = np.r_[lower[None], upper[None]]
    return bbox

def get_bbox_corners(bbox):
    """construct 8 corner points given extreme diagonal points of AABB"""
    xl, yl, zl = bbox[0]
    xu, yu, zu = bbox[1]
    corners = np.r_[np.r_[xl, yl, zl][None],
                    np.r_[xu, yl, zl][None],
                    np.r_[xu, yu, zl][None],
                    np.r_[xl, yu, zl][None],
                    np.r_[xl, yl, zu][None],
                    np.r_[xu, yl, zu][None],
                    np.r_[xu, yu, zu][None],
                    np.r_[xl, yu, zu][None]
                    ]
    return corners

def trimesh_scene_to_dicts(scene: Scene):
    """convert trimesh.Scene into list of dicts with vertices, faces and uv coordinates"""
    output = []
    for mesh in list(scene.geometry.values()):
        output.append(trimesh_to_dict(mesh))
    return output

def trimesh_to_dict(mesh: Trimesh):
    """extract vertices, faces and uv coordinates from trimesh.Trimesh and return them in dict"""
    return Dct(vertices=get_vertices(mesh),
               faces=get_faces(mesh),
               uvs=get_uv_coords(mesh),
               vertex_normals=get_vertex_normals(mesh))

def cam_extrinsics(cam):
    """compute camera's extrinsics matrix given position, lookat and up vectors in the input dict"""
    if not hasattr(cam, 'forward'):
        cam.forward = normalize(cam.lookat - cam.position)

    if not hasattr(cam, 'up'):
        cam.up = np.r_[0., 1., 0.]

    # re-orthonormalization
    cam_side = normalize(np.cross(cam.up, cam.forward))
    cam.up = normalize(np.cross(cam.forward, cam_side))

    # set up extrinsics matrix
    cam.extrinsics = np.r_[np.r_[np.r_[cam_side[None], cam.up[None], cam.forward[None]].T, np.r_[0., 0., 0.][None]].T, np.r_[cam.position, 1.][None]].T

    return cam

def set_tight_cam(vertices, cam: dict = None, use_bbox: bool = True, visualize: bool = False):
    # backward compatibility
    warn('set_tight_cam() is deprecated and should be replaced by cam_auto_zoom()')
    return cam_auto_zoom(vertices=vertices, cam=cam, use_bbox=use_bbox, visualize=visualize)

def cam_auto_zoom(vertices, cam: dict = None, use_bbox: bool = True, visualize: bool = False, debug: bool = False):
    """automatically set camera focal length to enclose a scene (either all vertices, or their axis aligned bounding
    box) in the view frustum"""
    if not isinstance(vertices, np.ndarray) and not isinstance(vertices, list):
        raise Exception('vertices must be nv x 3 np.ndarray or list of such arrays')

    bbox = get_bbox(vertices)
    bbox_diam = np.linalg.norm(np.diff(bbox, axis=0))
    scene_center = np.mean(bbox, axis=0)

    if use_bbox:
        # project only scene bounding box corners, not the vertices themselves to estimate the camera's FOV
        # (less accurate but faster)
        vertices = get_bbox_corners(bbox)

    if isinstance(vertices, list):
        vertices = np.concatenate(vertices, axis=0)

    # set up camera if not provided: lookat: scene center, position: random direction & distance from lookat
    if cam is None:
        cam = Dct(res_x=512, res_y=256)
    cam = Dct(cam)

    if not hasattr(cam, 'cx'):
        cam.cx = cam.res_x / 2.
        cam.cy = cam.res_y / 2.

    if not hasattr(cam, 'extrinsics'):
        if not hasattr(cam, 'position'):
            random_dir = normalize(np.random.rand(3) - 0.5)
            random_dist = (bbox_diam / 2 + 10. * np.random.rand())
            cam.position = scene_center + random_dist * random_dir

        if not hasattr(cam, 'lookat'):
            cam.lookat = scene_center

        if not hasattr(cam, 'up'):
            cam.up = np.r_[0., 1., 0.]

        # compute extrinsics matrix from position, lookat and up vectors
        cam = cam_extrinsics(cam)

    # bring bbox corners into camera space
    corners_cam = (np.linalg.inv(cam.extrinsics) @ np.r_[vertices.T, np.ones((vertices.shape[0], 1)).T]).T[:, :3]
    if np.abs(corners_cam[:, 2]).min() == 0.0:
        raise Exception('camera inside object, cannot automatically set fov / focal length')

    # perspective projection
    corners_projected = corners_cam / corners_cam[:, 2:3]
    # compute bbox of projected points in image plane
    bbox_projected = np.r_[np.min(corners_projected, axis=0)[None], np.max(corners_projected, axis=0)[None]]

    # at unit focal length (image plane) compute sensor size necessary to enclose projected scene bbox
    min_sensor_width = 2 * np.maximum(np.abs(bbox_projected[0, 0]), bbox_projected[1, 0])
    min_sensor_height = 2 * np.maximum(np.abs(bbox_projected[0, 1]), bbox_projected[1, 1])
    xy_ratio = min_sensor_width / min_sensor_height

    # account for physical aspect ratio
    film_xy_ratio = cam.res_x / cam.res_y

    # adjust sensor dimension with its actual aspect ratio
    if xy_ratio > film_xy_ratio:
        # x too wide --> increase y
        min_sensor_height = min_sensor_width / film_xy_ratio
    else:
        # y too high --> increase x
        min_sensor_width = film_xy_ratio * min_sensor_height
    cam.sensor_width = min_sensor_width
    cam.sensor_height = min_sensor_height

    cam.fov_x_deg = np.rad2deg(2 * np.arctan(cam.sensor_width / 2))
    cam.fov_y_deg = np.rad2deg(2 * np.arctan(cam.sensor_height / 2))

    # now compute focal length for a 36mm sensor (along its larger axis)
    if cam.res_x > cam.res_y:
        film_width = 1  # 36
        film_height = film_width * cam.res_y / cam.res_x
    else:
        film_height = 1  # 36
        film_width = film_height * cam.res_x / cam.res_y
    cam.pix_per_mm_x = cam.res_x / film_width
    cam.pix_per_mm_y = cam.res_y / film_height

    cam.fx_mm = film_width / (2 * np.tan(np.deg2rad(cam.fov_x_deg) / 2))
    cam.fy_mm = film_height / (2 * np.tan(np.deg2rad(cam.fov_y_deg) / 2))
    cam.fx = cam.pix_per_mm_x * cam.fx_mm
    cam.fy = cam.pix_per_mm_y * cam.fy_mm

    cam.intrinsics = np.array([[cam.fx, 0., cam.cx], [0., cam.fy, cam.cy], [0., 0., 1.]])

    if debug:
        output = Dct(cam=cam, corners_projected=corners_projected, bbox_projected=bbox_projected)
    else:
        output = cam

    if visualize:
        corners_bbox_projected_world = (cam.extrinsics @ np.r_[get_bbox_corners(bbox_projected).T, np.ones((8, 1)).T]).T[:, :3]
        ray_origins, ray_dirs = create_rays(cam)
        corners_projected_world = (cam.extrinsics @ np.r_[corners_projected.T, np.ones((corners_projected.shape[0], 1)).T])[:3].T
        from pysmtb.plotting import plot3, quiver3, scatter3
        sh = scatter3(cam.position, marker='s', color=np.r_[0, 0, 0])
        ax = sh.axes
        quiver3(cam.position[None], cam.up[None], color=np.r_[0.6, 0.6, 0.6], axes=ax)
        quiver3(cam.position[None], scene_center[None] - cam.position[None], color=np.r_[0.2, 0.2, 0.2], axes=ax)
        plot3(corners_bbox_projected_world[np.r_[np.r_[:4], 0]].T, axes=ax, color=np.r_[0.5, 0.4, 0.0])
        if use_bbox:
            for ci in range(8):
                plot3(np.r_[vertices[ci:ci+1, :], corners_projected_world[ci:ci+1, :], cam.position[None]].T, axes=ax, color=0.85 * np.r_[1, 1, 1])
            scatter3(vertices, axes=ax, color=np.r_[0.1, 0.4, 0.85])
            plot3(vertices[np.r_[np.r_[:4], 0, np.r_[4:8], 4, 5, 1, 2, 6, 7, 3]].T, axes=ax, color=np.r_[0.1, 0.4, 0.85])
            scatter3(corners_projected_world, axes=ax, color=np.r_[0.85, 0.6, 0.1])
            plot3(corners_projected_world[np.r_[np.r_[:4], 0, np.r_[4:8], 4, 5, 1, 2, 6, 7, 3]].T, axes=ax, color=np.r_[0.85, 0.6, 0.1])
        else:
            scatter3(vertices[::123, :], marker='.', axes=ax, color=np.r_[0.1, 0.4, 0.85])
            scatter3(corners_projected_world[::123, :], marker='.', axes=ax, color=np.r_[0.85, 0.6, 0.1])

        scatter3(ray_origins[::3*123] + ray_dirs[::3*123], marker='.', axes=ax)
        # quiver3(ray_origins[::123], ray_dirs[::123], axes=ax)
    return output

def cam_auto_zoom_trajectory(vertices, cam_positions, cam: dict = None, use_bbox: bool = True, visualize: bool = False):
    """given scene vertices and an N x 3 array of camera positions, as well as a camera dict with initial intrisics
    (resolution, principal point, ...), set the camera's focal length to ensure the projected scene is always contained
    in the view frustum; returns cam dict with fx and fy set to the mininum over all focal lengths, as well as a list of
    extrinsics matrices corresponding to the camera positions"""
    cam_orig = deepcopy(cam)
    lookat = np.mean(get_bbox(vertices), axis=0)
    fx = []
    fy = []
    extrinsics = []
    centers = []
    corners_projected = []
    bbox_projected = []
    for cam_pos in cam_positions:
        cam = deepcopy(cam_orig)
        cam.position = cam_pos
        cam.lookat = lookat
        cam = cam_extrinsics(cam)

        debug = cam_auto_zoom(vertices=vertices, cam=cam, use_bbox=use_bbox, debug=True)
        corners_projected.append(debug.corners_projected)
        bbox_projected.append(debug.bbox_projected)
        cam = debug.cam
        centers.append(cam.lookat)
        fx.append(cam.fx)
        fy.append(cam.fy)
        extrinsics.append(cam.extrinsics)
    # final focal length is the minimum (largest field of view) over each position's focal length
    cam.fx = np.min(fx)
    cam.fy = np.min(fy)

    if visualize:
        from pysmtb.plotting import scatter3
        sh = scatter3(vertices, '.')
        ax = sh.axes
        scatter3(cam_positions, 's', axes=ax)
        scatter3(np.array([(ex @ np.concatenate((bbp, np.ones((bbp.shape[0], 1))), axis=1).T).T[:, :3] for ex, bbp in
                           zip(extrinsics, bbox_projected)]).reshape((-1, 3)), '.', axes=ax)

    return cam, extrinsics

def create_rays(cam, repeat_origin=True, xs: np.ndarray = None, ys: np.ndarray = None):
    """given camera specifications (resolution in x and y, extrinsics and intrinsics), compute ray origins and
    directions for the full pixel raster (or optionally a patch specified by x and y ranges)"""
    if xs is None:
        xs = np.r_[0.5: float(cam.res_x) + 0.5]
    if ys is None:
        ys = np.r_[0.5: float(cam.res_y) + 0.5]
    xs, ys = np.meshgrid(xs, ys)
    pts_pixels = np.stack((xs.ravel(), ys.ravel(), np.ones(xs.size)), axis=0).astype(np.float32)

    # construct ray origins and directions
    ray_origins = np.array(cam.extrinsics @ np.r_[0, 0, 0, 1])[None, :3].astype(np.float32)
    if repeat_origin:
        ray_origins = ray_origins.repeat(xs.size, axis=0)
    ray_dirs = (np.linalg.inv(cam.intrinsics) @ pts_pixels).T
    ray_dirs = normalize(ray_dirs, axis=1)
    ray_dirs = (cam.extrinsics[:3, :3] @ ray_dirs.T).astype(np.float32).T
    return ray_origins, ray_dirs

def embree_create_scene(meshes,
                        with_tangent_frames: bool = True,
                        scene: EmbreeScene = None,
                        cam: dict = None,
                        auto_cam: bool = True,
                        auto_cam_bbox: bool = True,
                        auto_cam_visualize: bool = False):
    """given a scene as a list of meshes, construct a pyembree scene, optionally adjust the camera to automatically
    focus the scene, returns embree scene and camera"""

    if scene is None:
        # create new scene
        scene = EmbreeScene()

    if isinstance(meshes, Trimesh):
        meshes = meshes.scene()

    if isinstance(meshes, Scene):
        meshes = trimesh_scene_to_dicts(meshes)

    if isinstance(meshes, dict):
        meshes = [meshes]

    if not isinstance(meshes, list):
        raise Exception('input of type ' + str(type(meshes)) + ' is not supported')

    # precompute per-vertex tangent frames
    if with_tangent_frames:
        for mesh_id in range(len(meshes)):
            meshes[mesh_id] = get_tangent_frames(meshes[mesh_id])

    if auto_cam and auto_cam_bbox:
        bbox = get_bbox(meshes)
    all_vertices = []
    for mesh in meshes:
        TriangleMesh(scene, mesh['vertices'], mesh['faces'])
        if auto_cam and not auto_cam_bbox:
            all_vertices.append(mesh.vertices)

    if auto_cam:
        cam = cam_auto_zoom(bbox if auto_cam_bbox else all_vertices, cam=cam, use_bbox=auto_cam_bbox, visualize=auto_cam_visualize)
    return scene, cam

def embree_intersect_scene(scene: EmbreeScene,
                           cam: dict,
                           return_as_rasters: bool = False,
                           xs: np.ndarray = None,
                           ys: np.ndarray = None,
                           ):
    """
    Given a pyembree scene and a camera dict, shoot rays and intersect the scene, returning a dict with either
    N x C or H x W x C arrays if return_as_rasters == True. Optionally, a range of pixels can be specified as separate
    x- and y-coordinate arrays to only trace a region of interest. Such pixel coordinates should be in the range from 0
    to cam.res_x - 1 and 0 to cam.res_y - 1 respectively.

    The returned dict has keys:
    hit: binary mask indicating hit / miss (always H x W x 1)
    tfar: depth in camera space, 1e37 for missed pixels
    Ng: (unnormalized!) geometric (per-face) normal vectors
    u, v: barycentric coordinates for each intersected triangle
    primID: triangle ID
    geomID: sub-mesh ID
    """
    cam = Dct(cam)
    res_x = cam.res_x
    res_y = cam.res_y
    intrinsics = np.array([[cam.fx, 0, cam.cx], [0, cam.fy, cam.cy], [0, 0, 1]])

    if hasattr(cam, 'extrinsics'):
        extrinsics = cam.extrinsics
    else:
        extrinsics = np.eye(4)

    # generate pixel array
    if xs is None:
        xs = np.r_[0.5: float(res_x) + 0.5]
    elif np.remainder(xs[0], 1) == 0.:
        xs += 0.5
    if ys is None:
        ys = np.r_[float(res_y) + 0.5: 0.5: -1]
    elif np.remainder(ys[0], 1) == 0.:
        ys += 0.5
    xs, ys = np.meshgrid(xs, ys)
    pts_pixels = np.stack((xs.ravel(), ys.ravel(), np.ones(xs.size)), axis=0).astype(np.float32)
    num_pix = res_y * res_x

    # construct ray origins and directions
    ray_origins = np.array(extrinsics @ np.r_[0, 0, 0, 1])[None, :3].astype(np.float32).repeat(xs.size, axis=0)
    ray_dirs = (np.linalg.inv(intrinsics) @ pts_pixels).T
    ray_dirs = normalize(ray_dirs, axis=1)
    ray_dirs = (extrinsics[:3, :3] @ ray_dirs.T).astype(np.float32).T

    # trace rays
    its = scene.run(ray_origins, ray_dirs, output=1)
    its = Dct(its)

    # compute hit mask
    hit = its.tfar < 1e37
    # always reshape hit mask to bitmap resolution
    its.hit = hit.reshape((res_y, res_x, 1))

    # normalize normals
    its.Ng[hit, :] = normalize(its.Ng[hit, :], axis=1)

    # reshape all fields to image size, optionally multiply with hit mask to set missed pixels to 0
    # geomIDs are 0-based so we cannot overwrite their background with 0
    for key in its.keys():
        if return_as_rasters:
            its[key] = its[key].reshape((res_y, res_x, -1))
        elif key not in ['hit']:
            its[key] = its[key].reshape((num_pix, -1))
            its[key] = its[key][hit, :]

    # compute intersection points
    its['points'] = ray_origins[hit, :] + its['tfar'] * ray_dirs[hit, :]
    return its

def interpolate_vertex_attributes(its, meshes, return_as_rasters: bool = False):
    """given intersection dict with barycentric coordinates ('u', 'v'), geometry IDs and triangle IDs, as well as the
    corresponding list of meshes, compute texture coordinates coordinates for each intersected pixel"""

    if isinstance(meshes, Scene):
        meshes = list(meshes.geometry.values())
    elif not isinstance(meshes, list):
        raise Exception('meshes should be list of trimesh.Trimesh objects or one trimesh.Scene object')

    # compute uv coordinates
    hit = its['hit']
    num_hit = np.count_nonzero(hit)
    its['uvs'] = np.zeros((num_hit, 2), dtype=np.float32)
    its['tangent_frames'] = np.zeros((num_hit, 3, 3), dtype=np.float32)
    for mesh_id in range(len(meshes)):
        # select only those intersections corresponding to the current mesh
        mesh_mask = its['geomID'] == mesh_id

        # barycentric coordinates are used to interpolate vertex attributes
        u = its['u'][mesh_mask]
        v = its['v'][mesh_mask]
        uvw = np.stack([1 - u - v, u, v], axis=0)

        # get per-triangle-vertex attributes (currently uv coordinates and tangent frames)
        triangle_ids = its['primID'][mesh_mask]
        its_vertex_ids = get_faces(meshes[mesh_id])[triangle_ids, :].T
        face_vertex_uvs = get_uv_coords(meshes[mesh_id])[its_vertex_ids, :]
        meshes[mesh_id] = get_tangent_frames(meshes[mesh_id])
        face_vertex_tangent_frames = meshes[mesh_id].vertex_tangent_frames[its_vertex_ids, :, :]  # 3 x (3 * NF) x 3 x 3

        # interpolate per pixel texture coordinates
        its['uvs'][mesh_mask.flatten(), :] = np.sum(face_vertex_uvs * uvw[..., None], axis=0)
        # interpolate per pixel tangent frames
        its['tangent_frames'][mesh_mask.flatten(), :, :] = np.sum(face_vertex_tangent_frames * uvw[:, :, None, None], axis=0)

    if return_as_rasters:
        # reshape to image size
        its['uvs'] = its['uvs'].reshape(its['u'].shape[:2] + (2, ))
        its['tangent_frames'] = its['tangent_frames'].reshape(its['u'].shape[:2] + (3, 3))
    return its

def get_local_dirs(its, global_position, normalized=True):
    """given intersection dict with per-pixel interpolated tangent frames and intersection points, as well as a global
    positions (e.g. light or view), compute per pixel local directions and unnormalized global directions (e.g., for
    computing light attenuation)"""
    dirs_global = global_position[None, :] - its['points']
    frames = its['tangent_frames']
    dirs_local = np.einsum('nij,nj->ni', frames, dirs_global)
    if normalized:
        dirs_local = normalize(dirs_local, axis=-1)
    return Dct(dirs_local=dirs_local, dirs_global=dirs_global)

def embree_render_deferred(meshes = None,
                           with_tangent_frames: bool = True,
                           cam: dict = None,
                           auto_cam: bool = True,
                           auto_cam_bbox: bool = False,
                           light_position: tuple = (0., 10., 0.),
                           return_as_rasters: bool = False):
    """given a list of meshes, optionally a camera and a point light position, render the scene via ray tracing;
    returns pixel buffers with:
    - 3D intersection points
    -
    - view and light directions in tangent space
    - """

    if meshes is None:
        meshes = []
        meshes.append(create_unit_sphere())
        meshes.append(create_unit_cube())

    bbox = get_bbox(meshes)
    bbox_diam = np.linalg.norm(np.diff(bbox, axis=0))

    if cam is None:
        cam = Dct(res_x=384, res_y=256)
    if not hasattr(cam, 'cx'):
        cam.cx = cam.res_x / 2.
        cam.cy = cam.res_y / 2.

    if not hasattr(cam, 'position'):
        cam.position = np.mean(bbox, axis=0) + (bbox_diam / 2 + 10. * np.random.rand()) * normalize(np.random.rand(3) - 0.5)

    # scene & camera setup
    scene, cam = embree_create_scene(meshes=meshes, with_tangent_frames=with_tangent_frames, cam=cam, auto_cam=auto_cam,
                                     auto_cam_bbox=auto_cam_bbox, auto_cam_visualize=False)

    # ray tracing
    buffers = embree_intersect_scene(scene=scene, cam=cam, return_as_rasters=return_as_rasters)

    # per pixel tangent frames
    buffers = interpolate_vertex_attributes(buffers, meshes, return_as_rasters=return_as_rasters)
    h, w = cam.res_y, cam.res_x

    # local light & view directions
    view_dirs = get_local_dirs(buffers, cam.position, normalized=True)
    light_dirs = get_local_dirs(buffers, light_position, normalized=True)

    if return_as_rasters:
        mask = buffers['hit']
        buffers['view_dirs_local'] = assign_masked(mask, view_dirs.pop('dirs_local'))
        buffers['view_dirs_global'] = assign_masked(mask, view_dirs.pop('dirs_global'))
        buffers['light_dirs_local'] = assign_masked(mask, light_dirs.pop('dirs_local'))
        buffers['light_dirs_global'] = assign_masked(mask, light_dirs.pop('dirs_global'))
    else:
        buffers['view_dirs_local'] = view_dirs.pop('dirs_local')
        buffers['view_dirs_global'] = view_dirs.pop('dirs_global')
        buffers['light_dirs_local'] = light_dirs.pop('dirs_local')
        buffers['light_dirs_global'] = light_dirs.pop('dirs_global')

    return buffers