# Copyright (c) 2018-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#

import numpy as np
import copy
from common.skeleton import Skeleton
from common.mocap_dataset import MocapDataset
from common.camera import normalize_screen_coordinates, image_coordinates, world_to_camera

h36m_skeleton = Skeleton(
    parents=[-1, 0, 1, 2, 3, 4, 0, 6, 7, 8, 9, 0, 11, 12, 13, 14, 12, 16, 17, 18, 19, 20, 19, 22, 12, 24, 25, 26, 27, 28, 27, 30],
    joints_left=[6, 7, 8, 9, 10, 16, 17, 18, 19, 20, 21, 22, 23],
    joints_right=[1, 2, 3, 4, 5, 24, 25, 26, 27, 28, 29, 30, 31])

h36m_cameras_intrinsic_params = [
    {
        'id': '54138969',
        'center': [512.54150390625, 515.4514770507812],
        'focal_length': [1145.0494384765625, 1143.7811279296875],
        'radial_distortion': [-0.20709891617298126, 0.24777518212795258, -0.0030751503072679043],
        'tangential_distortion': [-0.0009756988729350269, -0.00142447161488235],
        'res_w': 1000,
        'res_h': 1002,
        'azimuth': 70,  # Only used for visualization
    },
    {
        'id': '55011271',
        'center': [508.8486328125, 508.0649108886719],
        'focal_length': [1149.6756591796875, 1147.5916748046875],
        'radial_distortion': [-0.1942136287689209, 0.2404085397720337, 0.006819975562393665],
        'tangential_distortion': [-0.0016190266469493508, -0.0027408944442868233],
        'res_w': 1000,
        'res_h': 1000,
        'azimuth': -70,  # Only used for visualization
    },
    {
        'id': '58860488',
        'center': [519.8158569335938, 501.40264892578125],
        'focal_length': [1149.1407470703125, 1148.7989501953125],
        'radial_distortion': [-0.2083381861448288, 0.25548800826072693, -0.0024604974314570427],
        'tangential_distortion': [0.0014843869721516967, -0.0007599993259645998],
        'res_w': 1000,
        'res_h': 1000,
        'azimuth': 110,  # Only used for visualization
    },
    {
        'id': '60457274',
        'center': [514.9682006835938, 501.88201904296875],
        'focal_length': [1145.5113525390625, 1144.77392578125],
        'radial_distortion': [-0.198384091258049, 0.21832367777824402, -0.008947807364165783],
        'tangential_distortion': [-0.0005872055771760643, -0.0018133620033040643],
        'res_w': 1000,
        'res_h': 1002,
        'azimuth': -110,  # Only used for visualization
    },
]

h36m_cameras_extrinsic_params = {
    'S1': [
        {
            'orientation': [0.1407056450843811, -0.1500701755285263, -0.755240797996521, 0.6223280429840088],
            'translation': [1841.1070556640625, 4955.28466796875, 1563.4454345703125],
        },
        {
            'orientation': [0.6157187819480896, -0.764836311340332, -0.14833825826644897, 0.11794740706682205],
            'translation': [1761.278564453125, -5078.0068359375, 1606.2650146484375],
        },
        {
            'orientation': [0.14651472866535187, -0.14647851884365082, 0.7653023600578308, -0.6094175577163696],
            'translation': [-1846.7777099609375, 5215.04638671875, 1491.972412109375],
        },
        {
            'orientation': [0.5834008455276489, -0.7853162288665771, 0.14548823237419128, -0.14749594032764435],
            'translation': [-1794.7896728515625, -3722.698974609375, 1574.8927001953125],
        },
    ],
    'S2': [
        {},
        {},
        {},
        {},
    ],
    'S3': [
        {},
        {},
        {},
        {},
    ],
    'S4': [
        {},
        {},
        {},
        {},
    ],
    'S5': [
        {
            'orientation': [0.1467377245426178, -0.162370964884758, -0.7551892995834351, 0.6178938746452332],
            'translation': [2097.3916015625, 4880.94482421875, 1605.732421875],
        },
        {
            'orientation': [0.6159758567810059, -0.7626792192459106, -0.15728192031383514, 0.1189815029501915],
            'translation': [2031.7008056640625, -5167.93310546875, 1612.923095703125],
        },
        {
            'orientation': [0.14291371405124664, -0.12907841801643372, 0.7678384780883789, -0.6110143065452576],
            'translation': [-1620.5948486328125, 5171.65869140625, 1496.43701171875],
        },
        {
            'orientation': [0.5920479893684387, -0.7814217805862427, 0.1274748593568802, -0.15036417543888092],
            'translation': [-1637.1737060546875, -3867.3173828125, 1547.033203125],
        },
    ],
    'S6': [
        {
            'orientation': [0.1337897777557373, -0.15692396461963654, -0.7571090459823608, 0.6198879480361938],
            'translation': [1935.4517822265625, 4950.24560546875, 1618.0838623046875],
        },
        {
            'orientation': [0.6147197484970093, -0.7628812789916992, -0.16174767911434174, 0.11819244921207428],
            'translation': [1969.803955078125, -5128.73876953125, 1632.77880859375],
        },
        {
            'orientation': [0.1529948115348816, -0.13529130816459656, 0.7646096348762512, -0.6112781167030334],
            'translation': [-1769.596435546875, 5185.361328125, 1476.993408203125],
        },
        {
            'orientation': [0.5916101336479187, -0.7804774045944214, 0.12832270562648773, -0.1561593860387802],
            'translation': [-1721.668701171875, -3884.13134765625, 1540.4879150390625],
        },
    ],
    'S7': [
        {
            'orientation': [0.1435241848230362, -0.1631336808204651, -0.7548328638076782, 0.6188824772834778],
            'translation': [1974.512939453125, 4926.3544921875, 1597.8326416015625],
        },
        {
            'orientation': [0.6141672730445862, -0.7638262510299683, -0.1596645563840866, 0.1177929937839508],
            'translation': [1937.0584716796875, -5119.7900390625, 1631.5665283203125],
        },
        {
            'orientation': [0.14550060033798218, -0.12874816358089447, 0.7660516500473022, -0.6127139329910278],
            'translation': [-1741.8111572265625, 5208.24951171875, 1464.8245849609375],
        },
        {
            'orientation': [0.5912848114967346, -0.7821764349937439, 0.12445473670959473, -0.15196487307548523],
            'translation': [-1734.7105712890625, -3832.42138671875, 1548.5830078125],
        },
    ],
    'S8': [
        {
            'orientation': [0.14110587537288666, -0.15589867532253265, -0.7561917304992676, 0.619644045829773],
            'translation': [2150.65185546875, 4896.1611328125, 1611.9046630859375],
        },
        {
            'orientation': [0.6169601678848267, -0.7647668123245239, -0.14846350252628326, 0.11158157885074615],
            'translation': [2219.965576171875, -5148.453125, 1613.0440673828125],
        },
        {
            'orientation': [0.1471444070339203, -0.13377119600772858, 0.7670128345489502, -0.6100369691848755],
            'translation': [-1571.2215576171875, 5137.0185546875, 1498.1761474609375],
        },
        {
            'orientation': [0.5927824378013611, -0.7825870513916016, 0.12147816270589828, -0.14631995558738708],
            'translation': [-1476.913330078125, -3896.7412109375, 1547.97216796875],
        },
    ],
    'S9': [
        {
            'orientation': [0.15540587902069092, -0.15548215806484222, -0.7532095313072205, 0.6199594736099243],
            'translation': [2044.45849609375, 4935.1171875, 1481.2275390625],
        },
        {
            'orientation': [0.618784487247467, -0.7634735107421875, -0.14132238924503326, 0.11933968216180801],
            'translation': [1990.959716796875, -5123.810546875, 1568.8048095703125],
        },
        {
            'orientation': [0.13357827067375183, -0.1367100477218628, 0.7689454555511475, -0.6100738644599915],
            'translation': [-1670.9921875, 5211.98583984375, 1528.387939453125],
        },
        {
            'orientation': [0.5879399180412292, -0.7823407053947449, 0.1427614390850067, -0.14794869720935822],
            'translation': [-1696.04345703125, -3827.099853515625, 1591.4127197265625],
        },
    ],
    'S11': [
        {
            'orientation': [0.15232472121715546, -0.15442320704460144, -0.7547563314437866, 0.6191070079803467],
            'translation': [2098.440185546875, 4926.5546875, 1500.278564453125],
        },
        {
            'orientation': [0.6189449429512024, -0.7600917220115662, -0.15300633013248444, 0.1255258321762085],
            'translation': [2083.182373046875, -4912.1728515625, 1561.07861328125],
        },
        {
            'orientation': [0.14943228662014008, -0.15650227665901184, 0.7681233882904053, -0.6026304364204407],
            'translation': [-1609.8153076171875, 5177.3359375, 1537.896728515625],
        },
        {
            'orientation': [0.5894251465797424, -0.7818877100944519, 0.13991211354732513, -0.14715361595153809],
            'translation': [-1590.738037109375, -3854.1689453125, 1578.017578125],
        },
    ],
}

class Human36mDataset(MocapDataset):
    def __init__(self, path3d, path2d, remove_static_joints=True, chunk_length=1, augment=False, pad=121, causal_shift=0, is_train=False):
        super().__init__(fps=50, skeleton=copy.deepcopy(h36m_skeleton))
        self.chunk_length = chunk_length
        self.augment = augment
        self.pad = pad
        self.causal_shift = causal_shift
        self.is_train = is_train
        self.train_subjects = ["S1", "S5", "S6", "S7", "S8"]
        self.test_subjects = ["S9", "S11"]

        self._cameras = copy.deepcopy(h36m_cameras_extrinsic_params)
        for cameras in self._cameras.values():
            for i, cam in enumerate(cameras):
                cam.update(h36m_cameras_intrinsic_params[i])
                for k, v in cam.items():
                    if k not in ['id', 'res_w', 'res_h']:
                        cam[k] = np.array(v, dtype='float32')

                # Normalize camera frame
                cam['center'] = normalize_screen_coordinates(cam['center'], w=cam['res_w'], h=cam['res_h']).astype('float32')
                cam['focal_length'] = cam['focal_length'] / cam['res_w'] * 2
                if 'translation' in cam:
                    cam['translation'] = cam['translation'] / 1000  # mm to meters

                # Add intrinsic parameters vector
                cam['intrinsic'] = np.concatenate((cam['focal_length'], cam['center'], cam['radial_distortion'], cam['tangential_distortion']))

        # Load serialized dataset
        self._data_3d = np.load(path3d, allow_pickle=True)['positions_3d'].item()

        if remove_static_joints:
            # Bring the skeleton to 17 joints instead of the original 32
            self.remove_joints([4, 5, 9, 10, 11, 16, 20, 21, 22, 23, 24, 28, 29, 30, 31])

            # Rewire shoulders to the correct parents
            self._skeleton._parents[11] = 8
            self._skeleton._parents[14] = 8

        data2d_file = np.load(path2d, allow_pickle=True)
        self._data_2d = data2d_file['positions_2d'].item()
        keypoints_symmetry = data2d_file['metadata'].item()['keypoints_symmetry']
        self.kps_left, self.kps_right = list(keypoints_symmetry[0]), list(keypoints_symmetry[1])

        self.process3DPose()
        self.normalize2DPose()
        # self._pairs = self._data_2d["S1"]["Directions"][0]
        self.getPairs()

        self.joints_left = self.skeleton().joints_left()
        self.joints_right = self.skeleton().joints_right()

    def supports_semi_supervised(self):
        return True

    def process3DPose(self):
        data_3d = self._data_3d
        self._data_3d = {}
        for subject, actions in data_3d.items():
            self._data_3d[subject] = {}
            for action_name, positions in actions.items():
                positions = positions[:, self.kept_joints]
                positions_3d = []
                for cam in self._cameras[subject]:
                    pos_3d = world_to_camera(positions, R=cam['orientation'], t=cam['translation'])
                    pos_3d[:, 1:] -= pos_3d[:, :1]  # Remove global offset, but keep trajectory in first position
                    positions_3d.append(pos_3d)
                self._data_3d[subject][action_name] = {'positions': positions, 'cameras': self._cameras[subject], 'positions_3d': positions_3d}

    def normalize2DPose(self):
        for subject in self._data_3d.keys():
            assert subject in self._data_2d, 'Subject {} is missing from the 2D detections dataset'.format(subject)
            for action in self._data_3d[subject].keys():
                assert action in self._data_2d[subject], 'Action {} of subject {} is missing from the 2D detections dataset'.format(action, subject)
                if 'positions_3d' not in self._data_3d[subject][action]:
                    continue
                for cam_idx in range(len(self._data_2d[subject][action])):
                    # We check for >= instead of == because some videos in H3.6M contain extra frames
                    mocap_length = self._data_3d[subject][action]['positions_3d'][cam_idx].shape[0]
                    assert self._data_2d[subject][action][cam_idx].shape[0] >= mocap_length

                    if self._data_2d[subject][action][cam_idx].shape[0] > mocap_length:
                        # Shorten sequence
                        self._data_2d[subject][action][cam_idx] = self._data_2d[subject][action][cam_idx][:mocap_length]

                    cam = self._cameras[subject][cam_idx]
                    kps = self._data_2d[subject][action][cam_idx]
                    kps[..., :2] = normalize_screen_coordinates(kps[..., :2], w=cam['res_w'], h=cam['res_h'])
                    self._data_2d[subject][action][cam_idx] = kps
                assert len(self._data_2d[subject][action]) == len(self._data_3d[subject][action]['positions_3d'])

    def getPairs(self):
        if self.is_train:
            subjects = self.train_subjects
        else:
            subjects = self.test_subjects

        for subject in subjects:
            assert subject in self._data_2d, 'Subject {} is missing from the 2D detections dataset'.format(subject)
            for action in self._data_3d[subject].keys():
                pose3dlist = self._data_3d[subject][action]["positions_3d"]
                pose2dlist = self._data_2d[subject][action]
                assert len(pose3dlist) == len(pose2dlist)
                for i in range(len(pose3dlist)):
                    assert pose3dlist[i].shape[0] == pose2dlist[i].shape[0]
                    n_chunks = (pose2dlist[i].shape[0] + self.chunk_length - 1) // self.chunk_length
                    offset = (n_chunks * self.chunk_length - pose2dlist[i].shape[0]) // 2
                    bounds = np.arange(n_chunks + 1) * self.chunk_length - offset
                    augment_vector = np.full(len(bounds - 1), False, dtype=bool)
                    self._pairs += zip(["{:s}_{:s}_{:d}".format(subject, action, i)] * len(bounds - 1), bounds[:-1], bounds[1:], augment_vector)
                    if self.augment:
                        self._pairs += zip(["{:s}_{:s}_{:d}".format(subject, action, i)] * len(bounds - 1), bounds[:-1], bounds[1:], ~augment_vector)

    def getDataFromPair(self, index):

        subject_action_camIndex, start_3d, end_3d, flip = self._pairs[index]

        subject, action, camIndex = subject_action_camIndex.split("_")
        seq_2d = self._data_2d[subject][action][int(camIndex)]

        start_2d = start_3d - self.pad - self.causal_shift
        end_2d = end_3d + self.pad - self.causal_shift
        low_2d = max(start_2d, 0)
        high_2d = min(end_2d, seq_2d.shape[0])
        pad_left_2d = low_2d - start_2d
        pad_right_2d = end_2d - high_2d
        if pad_left_2d != 0 or pad_right_2d != 0:
            group_2d = np.pad(copy.deepcopy(seq_2d[low_2d:high_2d]), ((pad_left_2d, pad_right_2d), (0, 0), (0, 0)), 'edge')
        else:
            group_2d = copy.deepcopy(seq_2d[low_2d:high_2d])

        if flip:
            # Flip 2D keypoints
            group_2d[:, :, 0] *= -1
            group_2d[:, self.kps_left + self.kps_right] = group_2d[:, self.kps_right + self.kps_left]

        seq_3d = self._data_3d[subject][action]["positions_3d"][int(camIndex)]

        low_3d = max(start_3d, 0)
        high_3d = min(end_3d, seq_3d.shape[0])
        pad_left_3d = low_3d - start_3d
        pad_right_3d = end_3d - high_3d
        if pad_left_3d != 0 or pad_right_3d != 0:
            group_3d = np.pad(copy.deepcopy(seq_3d[low_3d:high_3d]), ((pad_left_3d, pad_right_3d), (0, 0), (0, 0)), 'edge')
        else:
            group_3d = copy.deepcopy(seq_3d[low_3d:high_3d])

        if flip:
            # Flip 3D joints
            group_3d[:, :, 0] *= -1
            group_3d[:, self.joints_left + self.joints_right] = group_3d[:, self.joints_right + self.joints_left]

        group_cam = self._cameras[subject][int(camIndex)]['intrinsic']

        if flip:
            # Flip horizontal distortion coefficients
            group_cam[2] *= -1
            group_cam[7] *= -1

        group_3d[:,0,:] = 0

        return subject_action_camIndex, group_cam, group_3d, group_2d


from keras.models import Model
from keras.layers import Flatten, Convolution2D, Input, Dense
from keras.optimizers import Adam
import keras.backend as K
from PIL import Image
import numpy as np
from core import Env
from controller import sync, send_cmd, tetris_enum, p_wait, BTN_A, BTN_B, BTN_R, BTN_L, DPAD_U
import cv2
from gym import spaces
import pathlib
from statistics import mean
from random import shuffle, choice

RENDER = True

TRAINING_BUFFER = 20
INPUT_SHAPE = (90,160)
OUTPUT_ENV = 5
OUTPUT_REWARD = 3
OUTPUT_PRE = 2
ACTION_SPACE = len(tetris_enum)

CLASS_NAMES = np.array(['neutral', 'end_screen', 'error', 'title_a', 'title_b'])
CLASS_NAMES_PRE = np.array(['go', 'pregame'])
CLASS_NAMES_REWARD = np.array(['neutral', 'double', 'tetris'])

# Data paths
env_data = pathlib.Path('training/tetris_environment')
reward_data = pathlib.Path('training/tetris_line_clears')
pregame_data = pathlib.Path('training/tetris_pregame_opt')

# Model save paths
env_path = 'tetris/environment/env.ckpt'
reward_path = 'tetris/reward/reward.ckpt'
pregame_path = 'tetris/pregame/pregame.ckpt'

# Helper functions

def swap(a,b): return b, a

def one_hot(a, num_classes): return np.squeeze(np.eye(num_classes)[a.reshape(-1)])

def preprocess(im): return np.expand_dims(im.astype('float32') / 255.0, axis = 0)

# This file is a wrapper for an OpenAI gym environment
class Tetris(Env):
    # -1 so that the bot cannot choose BTN_NONE
    action_space = spaces.Discrete(ACTION_SPACE - 1)

    def __init__(self, train_reward=False, train_env=False, train_pregame=False, capture_device=0, debug=False):
        self.DEBUG = debug
        self.clock = 0
        self.reward = self.model_reward()
        self.env = self.model_environment()
        self.pregame = self.model_pregame()
        # Sync to the switch so that we can send controller commands
        sync()
        if self.DEBUG:
            print('Successfully synced to MCU')
        # Connect controller
        send_cmd(BTN_L + BTN_R)
        p_wait(0.5)
        send_cmd()
        p_wait(0.1)
        send_cmd(BTN_A)
        p_wait(0.5)
        send_cmd()
        if self.DEBUG:
            print('Successfully paired controller')
        # Capture switch screen
        try:
            self.switch_screen = cv2.VideoCapture(capture_device)
        except:
            print('Unable to open capture device. Is it plugged in?')
            exit()
        if train_env:
            print('Starting environment training')
            self.fit_environment()
            print('Finished environment training - saving')
            print('You can now set train_environment = False')
            self.env.save_weights(env_path)
        if train_pregame:
            print('Starting pregame optimization training')
            self.fit_pregame()
            print('Finished optimization training - saving')
            print('You can now set train_pregame = False')
            self.pregame.save_weights(pregame_path)
        if train_reward:
            print('Starting reward training')
            self.fit_reward()
            print('Finished reward training - saving')
            print('You can now set train_reward = False')
            self.reward.save_weights(reward_path)

    # A model that determines what state the game is in
    # (IN_GAME, RESULT_SCREEN, ERROR_SCREEN, TITLE_A, TITLE_B)
    def model_environment(self):
        frame = Input(shape=(*INPUT_SHAPE, 3))
        x = Convolution2D(32, kernel_size=(8,8), strides=4, activation='relu')(frame)
        x = Convolution2D(64, kernel_size=(4,4), strides=2, activation='relu')(x)
        x = Convolution2D(64, kernel_size=(3,3), strides=1, activation='relu')(x)
        x = Flatten()(x)
        x = Dense(512, activation='relu')(x)
        output = Dense(OUTPUT_ENV, activation='softmax')(x)
        m = Model(inputs = frame, outputs = output)
        m.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
        try:
            m.load_weights(env_path)
        except:
            print('Unable to load environment network')
        return m

    # A model for determining reward
    # (neutral [no reward], junk sent [.5 point], tetris [1 point])
    def model_reward(self):
        frame = Input(shape=(*INPUT_SHAPE, 3))
        x = Convolution2D(32, kernel_size=(8,8), strides=4, activation='relu')(frame)
        x = Convolution2D(64, kernel_size=(4,4), strides=2, activation='relu')(x)
        x = Convolution2D(64, kernel_size=(3,3), strides=1, activation='relu')(x)
        x = Flatten()(x)
        x = Dense(512, activation='relu')(x)
        output = Dense(OUTPUT_REWARD, activation='softmax')(x)
        m = Model(inputs = frame, outputs = output)
        m.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
        try:
            m.load_weights(reward_path)
        except:
            print('Unable to load reward network')
        return m

    # A model for waiting until the GO screen
    # (GO screen, pregame waiting screen)
    def model_pregame(self):
        frame = Input(shape=(*INPUT_SHAPE, 3))
        x = Convolution2D(32, kernel_size=(8,8), strides=4, activation='relu')(frame)
        x = Convolution2D(64, kernel_size=(4,4), strides=2, activation='relu')(x)
        x = Convolution2D(64, kernel_size=(3,3), strides=1, activation='relu')(x)
        x = Flatten()(x)
        x = Dense(512, activation='relu')(x)
        output = Dense(OUTPUT_PRE, activation='softmax')(x)
        m = Model(inputs = frame, outputs = output)
        m.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
        try:
            m.load_weights(pregame_path)
        except:
            print('Unable to load pregame-optimization network')
        return m

    # A model for determining reward
    # (neutral [no reward], junk sent [.5 point], tetris [1 point])
    def fit_reward(self, epochs = 20):
        for i in range(epochs):
            batch = np.zeros((10, *INPUT_SHAPE, 3))
            labels = np.zeros((10, OUTPUT_REWARD))
            x = 0
            paths = list(reward_data.glob('*/*.png'))
            shuffle(paths)
            accuracy = 0
            total = 0
            for path in paths:
                path = str(path)
                # Already resized
                try:
                    im = cv2.imread(path)
                except:
                    print('Unable to read in image')
                    continue
                batch[x] = im / 255.0
                # Split by windows path divider
                label = path.split('\\')[2]
                found = np.where(CLASS_NAMES_REWARD == label)
                labels[x] = one_hot(found[0][0], CLASS_NAMES_REWARD.size)
                x += 1
                if x == 10:
                    accuracy += self.reward.train_on_batch(x = batch, y = labels)[1]
                    total += 1
                    x = 0
            if i == epochs - 1:
                print('Final accuracy reward:', accuracy/total)

    # Called at the beggining to fit network to training labels (pre-recorded)
    def fit_environment(self, epochs = 3):
        for _ in range(epochs):
            batch = np.zeros((10, *INPUT_SHAPE, 3))
            labels = np.zeros((10, OUTPUT_ENV))
            x = 0
            paths = list(env_data.glob('*/*.png'))
            shuffle(paths)
            for path in paths:
                path = str(path)
                # Already resized
                try:
                    im = cv2.imread(path)
                except:
                    print('Unable to read in image')
                    continue
                batch[x] = im / 255.0
                # Split by windows path divider
                label = path.split('\\')[2]
                found = np.where(CLASS_NAMES == label)
                labels[x] = one_hot(found[0][0], CLASS_NAMES.size)
                x += 1
                if x == 10:
                    self.env.train_on_batch(x = batch, y = labels)
                    x = 0

    # Called at the beggining to fit network to training labels (pre-recorded)
    def fit_pregame(self, epochs = 3):
        for _ in range(epochs):
            batch = np.zeros((10, *INPUT_SHAPE, 3))
            labels = np.zeros((10, OUTPUT_PRE))
            x = 0
            paths = list(pregame_data.glob('*/*.png'))
            shuffle(paths)
            for path in paths:
                path = str(path)
                # Already resized
                try:
                    im = cv2.imread(path)
                except:
                    print('Unable to read in image')
                    continue
                batch[x] = im / 255.0
                # Split by windows path divider
                label = path.split('\\')[2]
                found = np.where(CLASS_NAMES_PRE == label)
                labels[x] = one_hot(found[0][0], CLASS_NAMES_PRE.size)
                x += 1
                if x == 10:
                    self.pregame.train_on_batch(x = batch, y = labels)
                    x = 0

    def step(self, action):
        # Send the chosen button press to switch
        send_cmd(tetris_enum[action])
        p_wait(.05)
        # Read in the resulting frame
        ret_val, self.current_frame = self.switch_screen.read()
        if RENDER:
            self.render()
        if ret_val == False:
            print('Error reading from switch screen. Switch disconnected?')
            print('Exiting')
            exit()

        img_resize = cv2.resize(self.current_frame, swap(*INPUT_SHAPE))
        im = Image.fromarray(img_resize)
        # Unsure if these 2 lines are necessary
        im.save('temp.png')
        im_import = cv2.imread('temp.png')

        preprocessed = preprocess(im_import)

        # Evaluate on environment network
        env_out = np.argmax(self.env.predict_on_batch(preprocessed))
        if self.DEBUG and self.clock % 10 == 0:
            print('Current environment evaluation:', CLASS_NAMES[env_out])
        done = False
        if env_out == 1:
            # Endgame screen
            reward = -1.0
            done = True
        elif env_out != 0:
            # Not in-game
            reward = 0.0
            done = True
        else:
            reward_func = np.argmax(self.reward.predict_on_batch(preprocessed))
            if reward_func == 0: # neutral
                reward = 0.0
            elif reward_func == 1: # line clear
                # Drop this value once agent learns how to do normal line clears
                reward = 1.0
            else:
                reward = 1.0 # tetris
            if self.DEBUG:
                print('Current reward:', reward)
        self.clock += 1
        # Crop to get only area of tetris
        cropped = self.current_frame[:, 480:-480]
        #cv2.imshow('cropped', cropped)
        #cv2.waitKey(1)
        return cropped, reward, done, {}

    def reset(self):
        self.clock = 0
        # While we are not in-game
        while True:
            # Read in current state of switch screen
            ret_val, self.current_frame = self.switch_screen.read()
            if ret_val == False:
                print('Error reading from switch screen. Switch disconnected?')
                print('Exiting')
                exit()
            if RENDER:
                self.render()

            img_resize = cv2.resize(self.current_frame, swap(*INPUT_SHAPE))
            im = Image.fromarray(img_resize)
            im.save('temp.png')
            im_import = cv2.imread('temp.png')

            preprocessed = preprocess(im_import)
            # Evaluate on environment
            env_out = np.argmax(self.env.predict_on_batch(preprocessed))
            if self.DEBUG:
                print('Environment:', CLASS_NAMES[env_out])
            if env_out == 0:
                # Check if in pregame or GO screen
                pregame_out = np.argmax(self.pregame.predict_on_batch(preprocessed))
                if CLASS_NAMES_PRE[pregame_out] == 'pregame':
                    # Still on pregame, wait
                    if self.DEBUG:
                        print('Waiting on pregame')
                    p_wait(.05)
                    continue
                else:
                    # We have successfully started a new game, return GO screen
                    if self.DEBUG:
                        print('GO')
                    cropped = self.current_frame[:, 480:-480]
                    return cropped
            if env_out == 1 or env_out == 2 or env_out == 3:
                # End game screen, title screen, or error screen
                send_cmd(BTN_A)
            if env_out == 4:
                # Screen for pressing B button
                send_cmd(BTN_B)
            # Send none to clear button press
            p_wait(3.0)
            send_cmd()
            p_wait(3.0)

    def render(self):
        try:
            resized_frame = cv2.resize(self.current_frame, (1280, 720))
        except NameError:
            print('No frames found')
            return
        cv2.imshow('Switch Screen', resized_frame)
        cv2.waitKey(1)

    def close(self):
        cv2.destroyAllWindows()
        self.switch_screen.release()


import math
import threading
import cv2
import numpy as np

class IMAGE_DATALOADER:
    def __init__(self, img_path_list, batch_size, resized_width, resized_height, num_cam):
        self.data = img_path_list
        self.batch_size = batch_size
        self.resized_width, self.resized_height = resized_width, resized_height
        self.num_cam = num_cam

    def __len__(self):
        return math.ceil(len(self.data) / self.batch_size)

    def __getitem__(self, idx):
        # return  [[imgA_cam1, imgA_cam2, ...], [imgB_cam1, imgB_cam2, ...], ...]
        indices = self.indices[idx*self.batch_size:(idx+1)*self.batch_size]

        paths_batch = [list(self.data[i].values()) for i in indices] #[[cam1_path, cam2_path, ...], [cam1_path, cam2_path, ...], ...]
        img_dict = {} #{0:{}, 1:{}, ...}
        for batch_idx in range(self.batch_size) :
            img_dict[batch_idx] = {key : None for key in range(self.num_cam)}

        threads = []
        for batch_idx, paths in enumerate(paths_batch):
            for cam_idx, path in enumerate(paths) :
                t = threading.Thread(target=self.load_img, args=(path, batch_idx, cam_idx, img_dict))
                threads.append(t)
                t.start()
                #self.load_img(path, batch_idx, cam_idx, img_dict)
        for t in threads :
            t.join()

        #img_dict : {0:{'1':imgA_cam1, '2':imgA_cam2, ....}, 1:{'1':imgB_cam1, '2':imgB_cam2, ....}, ...}

        batch = []
        for images_in_one_batch in list(img_dict.values()):
            batch.append(list(images_in_one_batch.values()))

        batch_np = np.array(batch).transpose(1, 0, 2, 3, 4)
        batch_list = [cam_batch for cam_batch in batch_np]
        return batch_list

    def load_img(self, path, batch_idx, cam_idx, img_dict):
        img = cv2.imread(path)
        #print(path)
        img = cv2.resize(img, (self.resized_width, self.resized_height), interpolation=cv2.INTER_CUBIC)
        img_dict[batch_idx][cam_idx] = img

    def set_indices(self, indices):
        self.indices = indices


# spector.py
# ALS 2017/06/01

import sys
import numpy as np
import astropy.table as at
import matplotlib.pyplot as plt
import astropy.units as u
from astropy.io import fits
import astropy.constants as const
import modelBC03
import os

from ..obsobj import Operator
from .. import filters
from . import getconti
from . import extrap
from . import linelist
from . import lineflux

class Spector(Operator):

	def __init__(self, **kwargs):
		"""
		Spector, an operator on spectrum

		Params
		----------
		Operator params:
			/either
				obj (object of class obsobj): with attributes ra, dec, dir_obj
			/or
				ra (float)
				dec (float)
				/either
					dir_obj (string)
				/or
					dir_parent (string): attr dir_obj is set to dir_parent+'SDSSJXXXX+XXXX/'

		survey (str):
			survey of the photometric system
			if not provided, use self.obj.survey. Raise exception if self.obj.survey does not exist.
		z (float):
			redshift, if not provided, use obj.z or obj.sdss.z
		survey_spec (str):
			survey of the spectrum
			if not provided, use the instrument of sdss_xid.csv
		decompose_method = 'modelBC03' (str)
			'modelBC03' or 'running_median'
			the method for decomposing spectrum into continuum and lines, and extrapolate the continum.

		Attributes
		----------
		Operator Attributes:
			obj (instance of objObj)
			ra (float)
			dec (float)
			dir_obj (string)

		survey (str): e.g., 'hsc'
			survey of the photometric system
		survey_spec (str): e.g., 'sdss' or 'boss'
			survey of the spectrum
		bands (list):
			e.g., ['g', 'r', 'i', 'z', 'y'] for survey = 'hsc'
		z (float):
		decompose_method (str)
			'modelBC03' or 'running_median'
		conti_model
			if decompose_method is modelBC03 then this is an instance of modelBC03
		"""

		super(Spector, self).__init__(**kwargs)

		# set survey
		if hasattr(self.obj, 'survey'):
			default_survey = self.obj.survey
			self.survey = kwargs.pop('survey', default_survey)
		else:
			self.survey = kwargs.pop('survey')

		# set survey_spec
		self.survey_spec = kwargs.pop('survey_spec', 'auto')

		if self.survey_spec in ['sdss', 'boss', 'eboss', 'auto']:
			if self.survey_spec == 'auto':
				self.obj.add_sdss(toload_photoobj=False)
				self.survey_spec = self.obj.sdss.instrument.lower()

			# # sanity check - spec name consistent
			# if (self.survey_spec != self.obj.sdss.instrument.lower()):
			# 	raise Exception("[spector] survey_spec inconsistent with sdss_xid.csv:instrument")

		# set z
		if 'z' in kwargs:
			self.z = kwargs.pop('z')

		elif hasattr(self.obj, 'z'):
			self.z = self.obj.z

		elif self.survey_spec in ['sdss', 'boss', 'eboss', 'auto']:
			self.obj.add_sdss(toload_photoobj=False)
			self.z = kwargs.pop('z', self.obj.sdss.z)

		# set self.decompose_method
		self.decompose_method = kwargs.pop('decompose_method', 'modelBC03')
		self.conti_model = None

		# set others
		self.bands = filters.filtertools.surveybands[self.survey]
		self.waverange = filters.filtertools.waverange[self.survey]

		# define paths
		self.fp_spec = self.dir_obj+'spec.fits'
		self.fp_spec_decomposed = self.dir_obj+'spec_decomposed.ecsv'
		self.fp_spec_contextrp = self.dir_obj+'spec_contextrp.ecsv'
		self.fp_spec_mag = self.dir_obj+'spec_mag.csv'
		self.fp_spec_lineflux = self.dir_obj+'spec_lineflux.csv'
		self.fp_spec_linefrac = self.dir_obj+'spec_linefrac.csv'

		self.spec, self.ws = self.get_spec_ws()
		self.u_spec = 1.e-17*u.Unit('erg / (Angstrom cm2 s)') # default unit of spec
		self.u_ws = u.AA # default unit of ws

	def get_spec_ws(self, forceload_from_fits=False):
		"""
		read spec and ws, either from spec_decomposed.ecsv or spec.fits (if forced or spec_decomposed.ecsv does not exist)

		Param
		-----
		forceload_from_fits=False:
			if true, then load from fits.

		Return
		------
		spec (astropy table column with unit)
		ws (astropy table column with unit)

		Default units
		-------------
		u_spec = 1.e-17*u.Unit('erg / (Angstrom cm2 s)')
		u_ws = u.AA
		"""

		fn = self.fp_spec_decomposed

		if not os.path.isfile(fn) or forceload_from_fits:
			spec, ws, __ = self.__read_spec_ws_ivar_from_fits()
			return spec, ws
		else:
			tab = at.Table.read(fn, format='ascii.ecsv')
			return tab['spec'], tab['ws']

	def get_spec_ws_from_spectab(self, component, fn=None):
		"""
		read certain spec table.ecsv to get spec of a given component
		if fn not specified then use the default file that would contain the specified component,
		e.g.:
			'spec_decomposed.ecsv' for component = "all", "cont", "line"
			'spec_contextrp.ecsv'    for component = "contextrp"

		Param
		-----
		component:
			either ['all', 'line', 'cont', 'contextrp']
		fn='spec_decomposed.ecsv'

		Return
		------
		spec (astropy table column with unit)
		ws (astropy table column with unit)

		Default units
		-------------
		u_spec = 1.e-17*u.Unit('erg / (Angstrom cm2 s)')
		u_ws = u.AA
		"""
		if fn == None:
			if component in ['all', 'cont', 'line']:
				fn = self.fp_spec_decomposed
				self.make_spec_decomposed_ecsv(overwrite=False)
			elif component in ['contextrp']:
				fn = self.fp_spec_contextrp
				self.make_spec_contextrp_ecsv(overwrite=False)
			else:
				raise Exception("[spector] component not recognized")

		tab = at.Table.read(fn, format='ascii.ecsv')
		col = self.__get_spectab_colname(component)

		return tab[col], tab['ws']

	def make_spec_decomposed_ecsv(self, overwrite=False):
		"""
		saving seperated continuum and line spectrum in csv file

		Params
		------
		self
		overwrite=False

		Return
		------
		status

		note: this part of the code can be refactorized better.

		"""

		fn = self.fp_spec_decomposed

		if (not os.path.isfile(fn)) or overwrite:

			spec, ws = self.get_spec_ws()

			ws_uless = np.array((ws/self.u_ws).to(u.dimensionless_unscaled))
			spec_uless = np.array((spec/self.u_spec).to(u.dimensionless_unscaled))

			iscon, speccont, specline, __, model = getconti.decompose_cont_line_t2AGN(spec_uless, ws_uless, self.z, method=self.decompose_method)

			self.conti_model = model # modelBC03 instance if set method='modelBC03', otherwise None.

			tab = at.Table([ws_uless, spec_uless, speccont, specline, iscon], names=['ws', 'spec', 'speccont', 'specline', 'iscon'])
			tab['ws'].unit = self.u_ws
			tab['spec'].unit = self.u_spec
			tab['speccont'].unit = self.u_spec
			tab['specline'].unit = self.u_spec

			tab.write(fn, format='ascii.ecsv', overwrite=overwrite)

			# sanity check: units are identical
			units = [tab[col].unit for col in ['spec', 'speccont', 'specline']]
			if len(set(units)) > 1:
				raise Exception("[spector] units in table spec_decomposed are not identical")

		status = os.path.isfile(fn)
		return status

	def make_spec_contextrp_ecsv(self, overwrite=False, refit=False):
		"""
		extrapolate continuum to cover all of the wavelength range of filters
		there are two methods:
			for self.conti_model modelBC03: use the bestfit
			for running_median: polynomial fit

		"""
		fn = self.fp_spec_contextrp

		if (not os.path.isfile(fn)) or overwrite:
			speccont, ws = self.get_spec_ws_from_spectab(component='cont')
			ws_uless = np.array((ws/self.u_ws).to(u.dimensionless_unscaled))
			speccont_uless = np.array((speccont/self.u_spec).to(u.dimensionless_unscaled))

			l0, l1 = self.waverange

			if self.decompose_method == 'modelBC03':
				if self.conti_model is None or refit:
					m = modelBC03.modelBC03(extinction_law='none')
					m.fit(ws=ws_uless, spec=speccont_uless, z=self.z)
				else:
					m = self.conti_model # reuse
				speccon_ext = m.bestfit
				ws_ext = m.ws_bestfit

			elif self.decompose_method == 'running_median':
				speccon_ext, ws_ext = extrap.extrap_to_ends(ys=speccont_uless, xs=ws_uless, x_end0=l0, x_end1=l1, polydeg=1, extbase_length=2000.)
			else:
				raise Exception("method is not recognized")

			assert (np.min(ws_ext*self.u_ws) < l0) & (np.max(ws_ext*self.u_ws) > l1)
			col_contextrp = self.__get_spectab_colname('contextrp')
			tab = at.Table([ws_ext, speccon_ext], names=['ws', col_contextrp])
			tab['ws'].unit = self.u_ws
			tab[col_contextrp].unit = self.u_spec
			tab.write(fn, format='ascii.ecsv', overwrite=overwrite)

		status = os.path.isfile(fn)
		return status

	def make_spec_mag(self, overwrite=False):
		"""
		make table spec_mag.csv that contains the convolved spectral magnitude and fnu in each band

		Params
		------
		self
		overwrite=False

		Return
		------
		status
		"""
		#==========================================================================

		fn = self.fp_spec_mag

		self.make_spec_decomposed_ecsv(overwrite=False)

		if not os.path.isfile(fn) or overwrite:
			print("[spector] making spec_mag")
			tabmag = at.Table()
			tabfnu = at.Table()

			for component in ['all', 'cont', 'line', 'contextrp']:
				for band in self.bands:
					colfnu = self.__get_specmag_colname(band, component=component, fluxquantity='fnu')
					colmag = self.__get_specmag_colname(band, component=component, fluxquantity='mag')
					try:
						fnu = self._calc_Fnu_in_band(band=band, component=component)
					except KeyboardInterrupt:
						sys.exit(0)
					except:
						print(("[spector] skip calculating fnu of {} in band {}".format(component, band)))
					else:
						mag = fnu.to(u.ABmag)
						fnu_nm = fnu.to(u.nanomaggy)
						tabmag[colmag] = [mag.value]
						tabfnu[colfnu] = [fnu_nm.value]

			tab = at.hstack([tabmag, tabfnu])
			tab.meta['comments'] = [
									"survey_photo: {}".format(self.survey),
									"survey_spec: {}".format(self.survey_spec),
									"unit_mag: ABmag",
									"unit_fnu: nanomaggy",
									]

			tab.write(fn, comment='#', format='ascii.csv', overwrite=overwrite)
		else:
			print("[spector] skip making spec_mag as file exists")

		status = os.path.isfile(fn)
		return status

	def get_spec_mag_tab(self):
		""" return spec_mag table"""
		self.make_spec_mag(overwrite=False)
		return at.Table.read(self.fp_spec_mag, comment='#', format='ascii.csv')

	def get_spec_mag_value(self, band, component='all', fluxquantity='mag'):
		"""
		return requested value, reading from spec_mag.csv

		Params
		------
		component='line':
			or 'all', 'cont', 'contextrp'
		fluxquantity='mag'
		band='i'

		Return
		------
		x (float): the value
		"""
		tab = self.get_spec_mag_tab()
		col = self.__get_specmag_colname(band=band, component=component, fluxquantity=fluxquantity)

		return tab[col][0]

	def get_fnu_ratio_band1_over_band2(self, band1, band2, component='all'):
		"""
		return fnu_band1/fnu_band2

		Params
		------
		band1, band2, component='all'

		Return
		------
		x (float): the ratio
		"""
		fluxquantity = 'fnu'
		tab = self.get_spec_mag_tab()
		col1 = self.__get_specmag_colname(band=band1, component=component, fluxquantity=fluxquantity)
		col2 = self.__get_specmag_colname(band=band2, component=component, fluxquantity=fluxquantity)

		ratio = tab[col1][0]/tab[col2][0]

		return ratio

	def plot_spec(self, wfilters=True, wconti=True, wcontextrp=True, wline=True, wspec=True, overwrite=True):
		"""
		Params
		------
		self
		wfilters=True
		wconti=False

		Return
		------
		status (bool)
		"""
		fn = self.dir_obj+'spec.pdf'

		if not os.path.isfile(fn) or overwrite:
			plt.close('all')
			plt.figure(1, figsize=(12, 6))
			plt.clf()

			plt.plot(0., 0., ls='', color='black', label='z='+'%.3f'%self.z) # show z in legend

			if wfilters:
				for band in self.bands:
					trans, ws_trans = self._get_norm_trans_func(band=band)
					plt.plot(ws_trans, trans/max(trans), label=band)

			if wspec:
				spec, ws = self.get_spec_ws()
				norm = max(spec)
				plt.plot(ws, spec/norm, color='0.3', lw=1.5, label='__nolabel__')

			if wline:
				speccont, ws = self.get_spec_ws_from_spectab(component='cont')
				specline, ws = self.get_spec_ws_from_spectab(component='line')
				specplot = specline+speccont
				specplot[specline==0] = np.nan
				plt.plot(ws, specplot/norm, color='black', lw=1.5, label='line')

			if wconti:
				speccont, ws = self.get_spec_ws_from_spectab(component='cont')
				plt.plot(ws, speccont/norm, color='cyan', lw=2, label='continuum')

			if wcontextrp:
				speccextrp, ws = self.get_spec_ws_from_spectab(component='contextrp')
				plt.plot(ws, speccextrp/norm, color='grey', lw=0.5, label='extrapolated conti')

			plt.legend(loc='upper right')
			if self.survey_spec == 'sdss':
				plt.xlim(2980.0, 11230.0)
			elif self.survey_spec == 'boss':
				plt.xlim(3500.0, 12000.0)
			plt.ylim(0., 1.)
			plt.savefig(fn, overwrite=overwrite)

		status = os.path.isfile(fn)
		return status

	def calc_fline_over_fnuband(self, band, line):
		"""
		calculate the conversion ratio r = (flux_line / fnu_band), to convert continuum subtracted image to line intensity map

		by multiplying this ratio to the band image (in fnu [erg/s/cm^2/Hz]) one can obtain the observed flux of the line (in f [erg/s/cm^2]). Notice that this is flux in observed frame, for measurements one still needs to transform it to the rest frame (depending on z).

		r = (c / (T(w_k) * w_k)) * frac_k  := dnu * frac_k

			where
				dnu    =  (c / (T(w_k) * w_k))
				frac_k = (f_k * T(w_k) * w_k) / sum(f_i * T(w_i) * w_i)

		Params
		------
		self
		band (str)
		line (str)

		Return
		------
		ratio (quantity in unit of Hz)
		"""
		f, __ = self._get_line_flux(line=line, wunit=False)
		w = self._get_line_obs_wave(line=line, wunit=False)
		T = self._get_norm_trans(wavelength=w, band=band, bounds_error=True)

		dnu = const.c / (T * w * u.AA)

		frac = self._get_line_frac(band=band, line=line)

		# sanity check: all strong lines are considered
		tab_linefrac = at.Table.read(self.fp_spec_linefrac, format='ascii.csv', comment='#')
		stronglines = self._list_stronglines_in_band(band=band)
		for line in stronglines:
			if 'frac_{}'.format(line) not in tab_linefrac.colnames:
				raise Exception("[spector] strong line {} in band is not contained in spec_linefrac.csv".format(line))

		r = (dnu * frac).to(u.Hz)
		return r

	def _get_line_frac(self, band, line='OIII5008'):
		"""
		read line flux from file spec_linefrac.csv.

		Params
		------
		band, line='OIII5008', wunit=False

		Return
		------
		frac (float)
		"""
		fn = self.fp_spec_linefrac

		if not os.path.isfile(fn):
			self.make_linefrac(band=band, overwrite=False)

		tab = at.Table.read(fn, format='ascii.csv', comment='#')

		if tab['lineband'][0] != band:
			raise Exception("[spector] _get_line_frac the band required is not provided by the current spec_linefrac.csv")

		linetag = 'frac_{}'.format(line)
		frac = tab[linetag][0]

		return frac

	def make_linefrac(self, band, lines=['NeIII3870', 'NeIII3969', 'Hg', 'Hb', 'OIII4960', 'OIII5008', 'OI6302', 'OI6366'], tofixOIIIratio=True, overwrite=False):
		"""
		make file spec_linefrac.csv that contains the fraction each of the strong lines have in a specific band.
		Columns: f_{band}_{line}, T_{band}_{line}, w_{band}_{line}, frac_{band}_{line}

		The fraction is based on the f*T*w of the line. Only the strong lines are listed.

		If tofixOIIIratio = True, then the ratio between OIII5008 and OIII4960 is fixed to the theoretical ratio of 2.98, see
		Storey + 2000. http://adsabs.harvard.edu/abs/2000MNRAS.312..813S.

		Params
		------
		band
		lines=['NeIII3870', 'NeIII3969', 'Hg', 'Hb', 'OIII4960', 'OIII5008', 'OI6302', 'OI6366']
		tofixOIIIratio=True
		overwrite=False

		Return
		------
		status
		"""
		fn = self.fp_spec_linefrac

		self.make_lineflux(overwrite=overwrite)

		if not os.path.isfile(fn) or overwrite:
			print("[spector] making spec_linefrac")

			tab = at.Table([[band]], names=['lineband'])

			fwt_sum = 0.
			for line in lines:
				f, __ = self._get_line_flux(line=line, wunit=False)
				w = self._get_line_obs_wave(line=line, wunit=False)
				T = self._get_norm_trans(wavelength=w, band=band, bounds_error=False)

				fwt = max(f*w*T, 0)
				fwt_sum = fwt_sum + fwt

				col_new = at.Table([[f], [w], [T], [fwt]], names=['f_{}'.format(line), 'w_{}'.format(line), 't_{}'.format(line), 'fwt_{}'.format(line)])
				tab = at.hstack([tab, col_new])

			for line in lines:
				frac = tab['fwt_{}'.format(line)][0] / fwt_sum
				col_new = at.Table([[frac]], names=['frac_{}'.format(line)])
				tab = at.hstack([tab, col_new])

			if tofixOIIIratio:
				r = 2.98
				frac_OIIItotal = tab['frac_OIII4960'] + tab['frac_OIII5008']
				tab['frac_OIII5008'] = frac_OIIItotal * r / (1.+r)
				tab['frac_OIII4960'] = frac_OIIItotal * 1. / (1.+r)

			tab.write(fn, format='ascii.csv', overwrite=overwrite)

		else:
			print("[spector] skip making spec_linefrac as file exists")

		status = os.path.isfile(fn)
		return status

	def make_lineflux(self, lines=['NeIII3870', 'NeIII3969', 'Hg', 'Hb', 'OIII4960', 'OIII5008', 'OI6302', 'OI6366'], u_flux=u.Unit("1E-17 erg cm-2 s-1"), overwrite=False):
		"""
		make file spec_lineflux.csv that contains the flux of the specified lines. The fluxes are calculated by integrating the line component of the spectrum over a window of +/- 1400 km/s.

		WARNING: currently only Hb, and OIII lines are supported. For lines that are overlapped, e.g., Ha and NII, the current implemenation will double count the flux.

		Params
		------
		lines=['Hb', 'OIII4960', 'OIII5008']
		u_flux=u.Unit("1E-17 erg cm-2 s-1")
			the unit of the output
		overwrite=False

		Return
		------
		status (bool)
		"""
		fn = self.fp_spec_lineflux

		self.make_spec_decomposed_ecsv(overwrite=False)

		if not os.path.isfile(fn) or overwrite:
			print("[spector] making spec_lineflux")
			tab = at.Table()

			for line in lines:
				f, ferr = self._calc_line_flux(line=line, u_flux=u_flux, wunit=False)

				col_new = at.Table([[f], [ferr]], names=['f_{}'.format(line), 'ferr_{}'.format(line)])
				tab = at.hstack([tab, col_new])

			tab.meta['comments'] = ["unit_flux: {}".format(u_flux.to_string()),]

			tab.write(fn, comment='#', format='ascii.csv', overwrite=overwrite)

		else:
			print("[spector] skip making spec_lineflux as file exists")

		status = os.path.isfile(fn)
		return status

	def _calc_line_flux(self, line='OIII5008', dv=1400*u.km/u.s, u_flux=u.Unit("1E-17 erg cm-2 s-1"), wunit=False):
		"""
		calculate the flux of the line by trpz integrating the decomposed emission line component of the spectrum over a range of +/- dv

		WARNING: currently only Hb, and OIII lines are supported. For lines that are overlapped, e.g., Ha and NII, the current implemenation will double count the flux.

		WARNING: the errors of the flux is propogated from the variance ('ivar' column) of the sdss spectrum. However, sdss's noise could be correlated between neighbors but such a covariance is not quantified, which will result in 10-20% in the error estimates.  We here artifically boost the flux error by 20% to incoporate such uncertainty on the error, see:
		http://www.sdss.org/dr12/spectro/caveats/

		Params
		------
		line='OIII5008' (str)

		dv=1400*u.km/u.s (quantity)
			half of the width over which the spectrum is integrated to calculate flux

		u_flux=u.Unit("1E-17 erg cm-2 s-1")
			the unit of the output

		wunit=False
			whether the output to come with unit or not

		Return
		------
		f (float or quantity)
			unit is in the dimension of erg cm-2 s-1
		ferr (float or quantity)
			same unit as f
		"""

		# sanity check
		if line not in ['NeIII3870', 'NeIII3969', 'Hg', 'Hb', 'OIII4960', 'OIII5008', 'OI6302', 'OI6366']:
			raise Exception("[spector] _calc_line_flux does not support lines other than Hb and OIII as those are not tested. ")

		# get w range
		beta = (dv/const.c).to_value(u.dimensionless_unscaled)
		w = self._get_line_obs_wave(line=line, wunit=False)
		w0 = w*(1-beta)
		w1 = w*(1+beta)

		# get spectrum
		spec, ws = self.get_spec_ws_from_spectab(component='line')
		__, __, ivar = self.__read_spec_ws_ivar_from_fits()

		f, ferr = lineflux.calc_line_flux(spec, ws, ivar, w0, w1, u_flux)

		# artificially boost the error to account for pixel covariance
		ferr = ferr * 1.2

		if wunit:
			return f, ferr
		else:
			return f.to_value(u_flux), ferr.to_value(u_flux)

	def _get_line_flux(self, line='OIII5008', wunit=False):
		"""
		read line flux from file spec_lineflux.csv. For details, see make_spec_lineflux().

		Params
		------
		line='OIII5008', wunit=False

		Return
		------
		f, ferr
		"""
		fn = self.fp_spec_lineflux

		if not os.path.isfile(fn):
			self.make_lineflux(overwrite=False)

		tab = at.Table.read(fn, format='ascii.csv', comment='#')

		f = tab['f_{}'.format(line)][0]

		ferr = tab['ferr_{}'.format(line)][0]

		if not wunit:
			return f, ferr

		else:
			u_flux = u.Quantity(tab.meta['comments'][0].split(': ')[1])
			return f*u_flux, ferr*u_flux

	def _get_line_flux_sdss(self, line='OIII5008', wunit=False):
		"""
		Flux of line from SDSS spec.fits higher extensions
		It is measured by Gaussian fit, for details, see:
		http://classic.sdss.org/dr7/dm/flatFiles/spZline.html

		PARAMS
		------
		line = 'OIII5008' (str)
		wunit = False

		Return
		------
		f (float):
			flux
			or astropy quantify with units u.Unit("1E-17 erg cm-2 s-1")
		ferr (float)
			error on flux
			or astropy quantify with units u.Unit("1E-17 erg cm-2 s-1")
		"""

		# get sdss linename
		linename = linelist.sdssLINENAME[line]

		# read flux
		table = self.obj.sdss.get_spec()[3].data
		i = [table['LINENAME']==linename]
		f = table[i]['LINEAREA'][0]
		ferr = table[i]['LINEAREA_ERR'][0]

		if not wunit:
			return f, ferr
		else:
			return f*u.Unit("1E-17 erg cm-2 s-1"), ferr*u.Unit("1E-17 erg cm-2 s-1")

	def _get_line_obs_wave(self, line='OIII5008', wunit=False):
		"""
		PARAMS
		------
		line = 'OIII5008' (str)
		wunit = False

		Return
		------
		w (float):
			redshifted wavelength of line
			or astropy quantify with units u.Unit("AA")
		"""
		w = filters.getllambda(ion=line, vacuum=True) * (1. + self.z)

		try: len(w)
		except:
			pass
		else:
			if len(w) == 1:
				w = w[0]
			else:
				raise Exception('got more than one wavelength')

		if not wunit:
			return w
		else:
			return w*u.Unit("AA")

	def _list_stronglines_in_band(self, band, threshold=0.01):
		"""
		return a list of strong lines in band where the filter transmission function is higher than threshold (in fraction)

		Params
		------
		self
		band (str)
		threshold (float):
			the fractional threshold (relative to the filter peak) that defines the wavelength boundary of a band
		wunit (bool)

		Return
		------
		llist (array of str)
			e.g., ['OIII5008', 'OIII4960', 'Hb', ...]
		"""
		w1, w2 = filters.filtertools.getFilterBoundaries(threshold=threshold, band=band, survey=self.survey, withunit=False)

		llist = []
		for line in linelist.strongline:
			w = self._get_line_obs_wave(line=line, wunit=False)
			if (w > w1) & (w < w2):
				llist += [line]

		return llist

	def __read_spec_ws_ivar_from_fits(self, u_spec=1.e-17*u.Unit('erg / (Angstrom cm2 s)'), u_ws=u.AA, wunit=True):
		"""
		Params
		------
		self
		u_spec=1.e-17*u.Unit('erg / (Angstrom cm2 s)')
		u_ws=u.AA
		wunit=True

		Return
		------
		spec (nparray)
		ws (nparray)
		ivar (nparray)

		Default units
		-------------
		u_spec = 1.e-17*u.Unit('erg / (Angstrom cm2 s)')
		u_ws = u.AA
		"""
		fn = self.fp_spec
		if self.survey_spec in ['sdss', 'boss', 'boss']:
			if os.path.isfile(fn):
				hdus = fits.open(fn)
			else:
				raise IOError("[Spector] spec fits file does not exist")

			hdus[0].header
			spectable = hdus[1].data
			spec, ws, ivar = spectable['flux'], 10.**spectable['loglam'], spectable['ivar']

			if wunit:
				spec = spec*u_spec
				ws = ws*u_ws
				ivar = ivar / (u_spec**2)

			# instrument_header = at.Table(hdus[2].data)['INSTRUMENT'][0].lower()
			# if self.survey_spec != instrument_header:
			# 	self.survey_spec = instrument_header
			# 	print("[Spector] updating survey_spec to reflect instrument in spec.fits header -- {}".format(instrument_header))

			return at.Column(spec, name=['spec']), at.Column(ws, name=['ws']), at.Column(ivar, name=['ivar'])
		else:
			raise NameError("[Spector] survey_spec not recognized")

	def _get_norm_trans_func(self, band='i'):
		"""
		return normalized transmission function and its wavelength coordinate
		the normalization is such that int{ trans{l} * dlnl} = 1.

		Params
		------
		self
		band='i'

		Return
		------
		trans  (array)
		ws_trans (array)
		"""
		trans, ws_trans = filters.filtertools.getNormTransFunc(band=band, survey=self.survey)
		return trans, ws_trans*u.AA

	def _get_norm_trans(self, wavelength, band='i', bounds_error=False):
		"""
		return normalized transmission function at a specific wavelenth, see filters.getNormTrans()
		the normalization is such that int{ trans{l} * dlnl} = 1.

		Params
		------
		self
		wave (float): wavelength to evaluate the function at
		band='i'
		bounds_error (bool):
			whether to raise error when interpolation outside of array is attempted

		Return
		------
		trans  (float)
		"""
		trans = filters.filtertools.getNormTrans(wavelength, band=band, survey=self.survey, bounds_error=bounds_error)
		return trans

	def _calc_Fnu_in_band(self, band, component='all'):
		"""
		Params
		------
		band=band
		component='all': from ['all', 'cont', 'line']
			which spectral component to operate on

		Return
		------
		Fnu (quantity in units "erg s-1 cm-2 Hz-1")
		"""

		spec, ws = self.get_spec_ws_from_spectab(component=component)

		trans, ws_trans = self._get_norm_trans_func(band=band)

		Fnu = filters.inttools.calc_Fnu_in_band_from_fl(fl=spec, ws=ws, trans=trans, ws_trans=ws_trans, isnormed=True)
		return Fnu

	def _calc_mAB_in_band(self, band, component='all'):
		Fnu = self._calc_Fnu_in_band(band=band, component='all')
		return Fnu.to(u.ABmag)

	def __get_specmag_colname(self, band, component, fluxquantity):
		"""
		band : e.g. 'i'
		component: from ['all', 'cont', 'line']
		fluxquantity: from ['fnu', 'mag']
		"""
		tag_component = {'all': '', 'cont': 'cont', 'line': 'line', 'contextrp': 'contextrp'}
		tag_quantity = {'fnu': 'Fnu', 'mag': 'Mag'}
		return 'spec{0}{1}_{2}'.format(tag_component[component], tag_quantity[fluxquantity], band)

	def __get_spectab_colname(self, component):
		"""
		band : e.g. 'i'
		component: from ['all', 'cont', 'line']
		fluxquantity: from ['fnu', 'mag']
		"""
		if component == 'ws':
			return 'ws'
		else:
			tag_component = {'all': '', 'cont': 'cont', 'line': 'line', 'contextrp': 'contextrp'}
			return 'spec{0}'.format(tag_component[component])

import numpy,import astropy,from astropy


#!/usr/bin/env python3
import rospy
rospy.init_node('dino_detect', log_level=rospy.DEBUG)
import sys
rospy.loginfo(sys.path)

import torch
import torchvision
import torch.nn.functional as F
import time
import datetime
import numpy as np
import cv2
import torchvision.transforms as transforms
import PIL.Image
import boto3
import os
import logging
logging.basicConfig(level=logging.DEBUG)
from jetbot import Camera

steering_model_path = rospy.get_param("ml_model_directory")+rospy.get_param("steering_model")
dino_detect_model_path = rospy.get_param("ml_model_directory")+rospy.get_param("dino_detect_model")

try:
    model_roadfollow = torchvision.models.resnet18(pretrained=False)
    model_roadfollow.fc = torch.nn.Linear(512, 2)
    model_roadfollow.load_state_dict(torch.load(steering_model_path))

    model_dinodet = torchvision.models.resnet18(pretrained=False)
    model_dinodet.fc = torch.nn.Linear(512, 6)
    model_dinodet.load_state_dict(torch.load(dino_detect_model_path))

    device = torch.device('cuda')

    model = model_roadfollow.to(device)
    model = model_roadfollow.eval().half()

    model_dinodet = model_dinodet.to(device)
    model_dinodet = model_dinodet.eval()

    camera = Camera.instance(width=int(rospy.get_param("image_width")), height=int(rospy.get_param("image_height")))

    from jetbot import Robot
    robot = Robot()
    speed_gain_slider = float(rospy.get_param("speed"))
    steering_gain_slider = float(rospy.get_param("steering_gain"))
    steering_dgain_slider = float(rospy.get_param("steering_d_gain"))
    steering_bias_slider = float(rospy.get_param("steering_bias"))
    angle = 0.0
    angle_last = 0.0
    prev_class = -1
except Exception as e:
    rospy.loginfo("Error occured:", str(e))

dino_names = rospy.get_param("dinosaurs")

mean = 255.0 * np.array([0.485, 0.456, 0.406])
stdev = 255.0 * np.array([0.229, 0.224, 0.225])

mean_roadfollow = torch.Tensor([0.485, 0.456, 0.406]).cuda().half()
std_roadfollow = torch.Tensor([0.229, 0.224, 0.225]).cuda().half()

normalize = torchvision.transforms.Normalize(mean, stdev)

def push_to_s3(img):
    try:
        bucket_name = "dinobot-roscon-unknown"
        timestamp = int(time.time())
        now = datetime.datetime.now()
        key = "unknown_dino/{}_{}_{}_{}_{}.jpg".format(now.month, now.day, now.hour, now.minute, timestamp)

        #s3 = boto3.client('s3')

        #encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 90]
        #_, jpg_data = cv2.imencode('.jpg', img, encode_param)
        #response = s3.put_object(ACL='public-read',
        #                         Body=jpg_data.tostring(),
        #                         Bucket=bucket_name,
        #                         Key=key)
        msg = key
        print("pub to s3:"+key)
    except Exception as e:
        msg = "Pushing to S3 failed: " + str(e)
    return msg

def publish_to_iot(msg):
    print("publish to iot: "+msg)
#    try:
#        client = boto3.client('iot-data')
#        response = client.publish(topic ='dino-detect' , qos =0, payload=msg)
#    except Exception as e:
#        msg = "Publish to IoT failed: " + str(e)

def preprocess(camera_value):
    global device, normalize
    x = camera_value
    x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)
    x = x.transpose((2, 0, 1))
    x = torch.from_numpy(x).float()
    x = normalize(x)
    x = x.to(device)
    x = x[None, ...]
    return x

def preprocess_roadfollow(image):
    image = PIL.Image.fromarray(image)
    image = transforms.functional.to_tensor(image).to(device).half()
    image.sub_(mean_roadfollow[:, None, None]).div_(std_roadfollow[:, None, None])
    return image[None, ...]

def find_dino(change):
    x = change['new']
    x = preprocess(x)
    y = model_dinodet(x)
    y_dino = F.softmax(y, dim=1)
    topk = y_dino.cpu().topk(1)

    return (e.data.numpy().squeeze().tolist() for e in topk)

def move_bot(image, robot_stop):
    global angle, angle_last
    if robot_stop:
        robot.stop()
        robot.left_motor.value=0
        robot.left_motor.value=0
        time.sleep(2)
        robot_stop = False
    else:
        xy = model_roadfollow(preprocess_roadfollow(image)).detach().float().cpu().numpy().flatten()
        x = xy[0]
        y = (0.5 - xy[1]) / 2.0
        speed_slider = speed_gain_slider
        angle = np.arctan2(x, y)
        pid = angle * steering_gain_slider + (angle - angle_last) * steering_dgain_slider
        angle_last = angle
        steering_slider = pid + steering_bias_slider
        robot.left_motor.value = max(min(speed_slider + steering_slider, 1.0), 0.0)
        robot.right_motor.value = max(min(speed_slider - steering_slider, 1.0), 0.0)

while True:
    try:
        img = camera.value
        robot_stop = False
        probs, classes = find_dino({'new': img})
        msg = "Start..."
        s3url = ""
        if probs > 0.6 and prev_class != classes:
            prev_class = classes
            robot_stop = True
            if classes == 5:
                s3url = push_to_s3(img)
                msg = '{"dinosaur":"unknown"' + ',"confidence":"' + str(probs) +'","url":"'+ s3url +'"}'
            else:
                msg = '{"dinosaur":"' + dino_names[classes] + '"' + ',"confidence":"' + str(probs) +'"}'
            print(msg)
            publish_to_iot(msg)
        move_bot(img, robot_stop)
    except Exception as e:
        rospy.loginfo("Error occured:", str(e))


import numpy as np
import argparse
import torch
from torch.backends import cudnn
import torch.nn.functional as tfunc
from torch.utils.data import DataLoader
from math import sqrt
import json
import os
from .utils.dataloader import RSNADataSet
from .utils.score import compute_auroc
from .utils.model import DenseNet121, DenseNet121Eff, load_checkpoint
import onnx
import onnxruntime
from PIL import Image
from torchvision import transforms
from openvino.inference_engine import IECore

class RSNAInference():
    def __init__(self, model, data_loader_test, class_count, checkpoint, class_names, device):
        self.device = device
        self.model = model.to(self.device)
        self.checkpoint = checkpoint
        load_checkpoint(self.model, self.checkpoint)
        self.data_loader_test = data_loader_test
        self.class_count = class_count
        self.class_names = class_names
        self.model.eval()

    def test_onnx(self, img_path, onnx_checkpoint):
        onnx_model = onnx.load(onnx_checkpoint)
        onnx.checker.check_model(onnx_model)
        # The validity of the ONNX graph is verified by checking the model’s version,
        # the graph’s structure, as well as the nodes and their inputs and outputs.
        ort_session = onnxruntime.InferenceSession(onnx_checkpoint)
        sample_image = Image.open(img_path).convert('RGB')
        to_tensor = transforms.ToTensor()
        sample_image = to_tensor(sample_image)
        sample_image.unsqueeze_(0)
        ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(sample_image)}
        ort_outs = ort_session.run(None, ort_inputs)
        torch_out = self.model(sample_image.cuda())
        np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)

    def load_inference_model(self, run_type, onnx_checkpoint):
        if run_type == 'pytorch':
            model = self.model
        elif run_type == 'onnx':
            model = onnxruntime.InferenceSession(onnx_checkpoint)
        else:
            ie = IECore()
            model_xml = os.path.splitext(onnx_checkpoint)[0] + ".xml"
            model_bin = os.path.splitext(model_xml)[0] + ".bin"
            model_temp = ie.read_network(model_xml, model_bin)
            model = ie.load_network(network=model_temp, device_name='CPU')
        return model

    def validate_models(self, run_type, onnx_checkpoint =''):
        cudnn.benchmark = True
        out_gt = torch.FloatTensor()
        out_pred = torch.FloatTensor()
        model = self.load_inference_model(run_type, onnx_checkpoint)

        with torch.no_grad():
            for var_input, var_target in self.data_loader_test:
                if run_type in ('pytorch', 'onnx'):
                    var_target = var_target.to(self.device)
                    var_input = var_input.to(self.device)
                    out_gt = out_gt.to(self.device)
                    out_pred = out_pred.to(self.device)
                    out_gt = torch.cat((out_gt, var_target), 0)
                else:
                    out_gt = torch.cat((out_gt, var_target), 0)
                _, c, h, w = var_input.size()
                var_input = var_input.view(-1, c, h, w)
                to_tensor = transforms.ToTensor()
                if run_type == 'pytorch':
                    out = model(var_input)
                elif run_type == 'onnx':
                    ort_inputs = {model.get_inputs()[0].name: to_numpy(var_input)}
                    out = model.run(None, ort_inputs)
                    out = np.array(out)
                    out = to_tensor(out).squeeze(1).transpose(dim0=1, dim1=0).to(self.device)
                else:
                    out = model.infer(inputs={'input': var_input})['output']
                    out = to_tensor(out).squeeze(1)

                out_pred = torch.cat((out_pred, out), 0)
                one_hot_gt = tfunc.one_hot(out_gt.squeeze(1).long()).float()
                auroc_individual = compute_auroc(one_hot_gt, out_pred, self.class_count)
        auroc_mean = np.array(auroc_individual).mean()
        return auroc_mean

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

def main(args):
    checkpoint= args.checkpoint
    class_count = 3

    class_names = ['Lung Opacity', 'Normal', 'No Lung Opacity / Not Normal']
    dpath = args.dpath
    img_pth = os.path.join(args.dpath, 'processed_data/')
    numpy_path = os.path.join(args.dpath, 'data_split/')
    with open(os.path.join(dpath, 'rsna_annotation.json')) as lab_file:
        labels = json.load(lab_file)
    test_list = np.load(os.path.join(numpy_path, 'test_list.npy')).tolist()

    dataset_test = RSNADataSet(test_list, labels, img_pth, transform=True)
    data_loader_test = DataLoader(
        dataset=dataset_test,
        batch_size=1,
        shuffle=False,
        num_workers=4,
        pin_memory=False)

    if args.optimised:
        alpha = args.alpha
        phi = args.phi
        beta = args.beta
        if beta is None:
            beta = round(sqrt(2 / alpha), 3)
        alpha = alpha ** phi
        beta = beta ** phi
        model = DenseNet121Eff(alpha, beta, class_count)
    else:
        model = DenseNet121(class_count)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # use gpu if available

    rsna_inference = RSNAInference(model, data_loader_test, class_count, checkpoint, class_names, device)

    test_auroc = rsna_inference.validate_models(run_type='pytorch')
    print(f"Test AUROC is {test_auroc}")

if __name__=="__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint",
        required=False,
        help="start training from a checkpoint model weight",
        default= None,
        type = str)
    parser.add_argument("--dpath",
        required=True,
        help="Path to dataset",
        type =str)
    parser.add_argument("--optimised",
        required=False,
        default=False,
        help="enable flag for eff model",
        action='store_true')
    parser.add_argument("--alpha",
        required=False,
        help="alpha for the model",
        default=(11 / 6),
        type=float)
    parser.add_argument("--phi",
        required=False,
        help="Phi for the model.",
        default=1.0,
        type=float)
    parser.add_argument("--beta",
        required=False,
        help="Beta for the model.",
        default=None,
        type=float)
    custom_args = parser.parse_args()
    main(custom_args)


from base64 import b64decode
from contextlib import contextmanager
from errno import EEXIST
from json import dumps
from os import path, makedirs, getcwd
from sys import version_info
from typing import overload, Iterable, Any, Union

import numpy as np
from cv2 import cv2
from pytest import skip
from selenium.common.exceptions import WebDriverException
from selenium.webdriver.common.by import By
from skimage.metrics import structural_similarity as ssim

if version_info >= (3, 8):
    from functools import singledispatchmethod

class SnapSizeError(AssertionError):
    """ Sizes of provided images are not equal. """

class SnapDifferenceError(AssertionError):
    """ Contents of provided images are not equal. """

class SnapTypeError(TypeError):
    """ Invalid data type given. """

class SnapInvalidCallWarning(Warning):
    """ Invalid attributes given. """

class SnapshotComparator(object):
    """
    Class that calculates the difference between two provided images and creates difference image.
    """

    def __init__(self, output_snap, reference_snap):
        self.output_snap = self._read_snap(output_snap)
        self.reference_snap = self._read_snap(reference_snap)

        self._ssim_equality = None
        self._difference_image = None

    @property
    def difference(self):
        return round(((1 - self._ssim_equality) * 100), 4)

    @property
    def difference_image(self):
        return cv2.imencode('.png', self._draw_contours())[1].tobytes()

    def get_equality_and_diff_image(self):
        self._ssim_equality, self._difference_image = ssim(*self._get_grayscale(), full=True, gaussian_weights=True)

    def _get_grayscale(self):
        return cv2.cvtColor(self.reference_snap, cv2.COLOR_BGR2GRAY), cv2.cvtColor(self.output_snap, cv2.COLOR_BGR2GRAY)

    def _draw_contours(self):
        contours = cv2.findContours(
            cv2.threshold(
                (self._difference_image * 255).astype('uint8'), 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU
            )[1],
            cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_SIMPLE
        )
        contours = contours[0] if len(contours) == 2 else contours[1]

        alpha = 0.6
        color = (0, 255, 0, 125)
        return cv2.addWeighted(
            cv2.drawContours(self.output_snap.copy(), contours, -1, color, -1), alpha, self.output_snap, 1 - alpha, 0
        )

    if version_info >= (3, 8):
        @singledispatchmethod
        def _read_snap(self, snap):
            raise SnapTypeError(f'{type(snap)} is an unsupported type, expected - {bytes}, {str}')

        @_read_snap.register
        def _(self, snap: str):
            return cv2.imread(snap)

        @_read_snap.register
        def _(self, snap: bytes):
            return cv2.imdecode(np.asarray(bytearray(snap), dtype=np.uint8), cv2.IMREAD_COLOR)

    else:
        @overload
        def _read_snap(self, snap: str) -> np.ndarray: ...
        @overload
        def _read_snap(self, snap: bytes) -> np.ndarray: ...

        def _read_snap(self, snap):
            if isinstance(snap, str):
                return cv2.imread(snap)
            elif isinstance(snap, bytes):
                return cv2.imdecode(np.asarray(bytearray(snap), dtype=np.uint8), cv2.IMREAD_COLOR)
            else:
                raise SnapTypeError(f'{type(snap)} is an unsupported type, expected - {bytes}, {str}')

class Asserter(SnapshotComparator):
    """
    Class asserting that the difference between two provided images is under the difference_limit.
    """
    def __init__(self, output_snap, reference_snap, difference_limit=0):
        self.difference_limit = difference_limit

        super().__init__(output_snap, reference_snap)

        self.assert_snap_sizes()
        self.get_equality_and_diff_image()

    def assert_snap(self):
        if self.difference > self.difference_limit:
            raise SnapDifferenceError(
                f'The difference between new screenshot and reference ({self.difference}) '
                f'exceeds acceptable limit ({self.difference_limit})')

    def assert_snap_sizes(self):
        if self.output_snap.shape != self.reference_snap.shape:
            raise SnapSizeError(
                f'Screenshot sizes do not match.\n'
                f'New screenshot size - {self.output_snap.shape[0]} на {self.output_snap.shape[1]},'
                f'reference size - {self.reference_snap.shape[0]} на {self.reference_snap.shape[1]}.'
            )

class Snappy(object):
    """
    Main pytest-snappy class that implements comparing methods used in test functions.
    """
    def __init__(self, driver, refresh_reference):
        """
        Snappy object initialization method.

        :param driver: Instance of Selenium WebDriver
        :param refresh_reference: Boolean value specified on the command line
        by '--refresh_references' parameter (False by default)
        """
        self.driver = driver
        self.refresh_reference = refresh_reference

        self.fullpage = True
        self.mask_locators = None
        self.locator = None

        self.filename = None

        self.output_snap = None
        self.difference_image = None

        self.reference_directory = path.realpath(path.join(getcwd(), 'snap_references'))
        try:
            makedirs(self.reference_directory)
        except OSError as error:
            if error.errno == EEXIST:
                pass
            else:
                raise error

        self.driver.maximize_window()

    def assert_snapshots(self, difference_limit=0) -> None:
        """
        Assert-variant of compare_snapshots method for situations when you dont need context manager functionality.
        For example, when all operations before the screenshot is taken are guaranteed to complete.

        :param difference_limit: Difference threshold, upon exceeding which an exception will be raised (0 by default).
        """
        with self.compare_snapshots(difference_limit):
            pass

    @contextmanager
    def compare_snapshots(self, difference_limit=0) -> None:
        """
        Context manager for performing screenshot comparing (using the Asserter class),
        reference saving and errors raising in cases when screenshots contents or sizes are different.

        :param difference_limit: Difference threshold, upon exceeding which an exception will be raised (0 by default).
        """
        yield

        reference_file = path.join(self.reference_directory, f'{self.filename}.png')

        if self.locator:
            self.output_snap = self._get_element_screenshot_as_bytes(*self.locator)
        elif self.fullpage:
            if self.mask_locators:
                self._mask_elements(self.mask_locators)
            self.output_snap = self._get_fullpage_screenshot_as_bytes()
        else:
            raise SnapInvalidCallWarning('Invalid arguments given, fullpage should be True or locator should be given')

        if not path.isfile(reference_file) or self.refresh_reference:
            with open(reference_file, 'wb') as file:
                file.write(self.output_snap)
            skip(
                'Reference snapshot is missing or out of date. '
                'The current screenshot is saved as a reference.'
            )
        else:
            asserter = Asserter(self.output_snap, reference_file, difference_limit)
            try:
                asserter.assert_snap()
            except SnapDifferenceError as error:
                self.difference_image = asserter.difference_image
                raise error
            except SnapSizeError as error:
                self.difference_image = self.output_snap
                raise error

    def _mask_elements(self, mask_locators: Iterable[Iterable]) -> None:
        """
        Masks (makes transparent) elements found by locators in self.mask_locators
        (in order to hide dynamic elements on page).

        :param mask_locators: Iterable object that contains iterable locators (may be tuples, lists, etc)
        """
        for locator in mask_locators:
            by, locator, *_ = locator
            for element in self.driver.find_elements(by=by, value=locator):
                try:
                    self.driver.execute_script('arguments[0].setAttribute("style", "opacity:0;");', element)
                except WebDriverException as error:
                    print('Error : ', str(error))
                    raise error

    def _get_fullpage_screenshot_as_bytes(self) -> bytes:
        """
        Makes screenshot of full page and returns it as bytes (may not work in non-chromium browsers).

        :return: Byte type screenshot of full page
        """
        def send(cmd, params):
            resource = f'/session/{self.driver.session_id}/chromium/send_command_and_get_result'
            url = self.driver.command_executor._url + resource
            body = dumps({'cmd': cmd, 'params': params})
            response = self.driver.command_executor._request('POST', url, body)
            return response.get('value')

        def evaluate(script):
            response = send('Runtime.evaluate', {'returnByValue': True, 'expression': script})
            return response['result']['value']

        metrics = evaluate("""
            ({
                width: Math.max(window.innerWidth, document.body.scrollWidth, document.documentElement.scrollWidth)|0,
                height: Math.max(innerHeight, document.body.scrollHeight, document.documentElement.scrollHeight)|0,
                deviceScaleFactor: window.devicePixelRatio || 1,
                mobile: typeof window.orientation !== "undefined"
            })
        """)
        send('Emulation.setDeviceMetricsOverride', metrics)
        screenshot = send('Page.captureScreenshot', {'format': 'png', 'fromSurface': True})
        send('Emulation.clearDeviceMetricsOverride', {})

        return b64decode(screenshot['data'])

    def _get_element_screenshot_as_bytes(self, by: Union[By, str], locator: str, *_: Any) -> bytes:
        """
        Makes screenshot of element found by self.locator.

        :param by: WebDriver By object or str e.g. 'css selector' or 'xpath'
        :param locator: Element locator
        :param _: Unused parameters for case when locator tuple contains more than two values
        :return: Byte type screenshot of element found by given locator
        """
        return self.driver.find_element(by=by, value=locator).screenshot_as_png