import random
import sys

import pygame
from pygame.locals import *

import Bird
import Pipe

pygame.init()

WIDTH = 288
HEIGHT = 512

FPS = 30
CLOCK = pygame.time.Clock()

output = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Flappy Bird OOP game")
pygame.display.set_icon(pygame.image.load('game.ico'))

game_back = [pygame.image.load("content/bg.png")]

digits = [pygame.image.load('content/0.png'),
		  pygame.image.load('content/1.png'),
		  pygame.image.load('content/2.png'),
		  pygame.image.load('content/3.png'),
		  pygame.image.load('content/4.png'),
		  pygame.image.load('content/5.png'),
		  pygame.image.load('content/6.png'),
		  pygame.image.load('content/7.png'),
		  pygame.image.load('content/8.png'),
		  pygame.image.load('content/9.png')]

SCORE = 0

SOUNDS = {
	'die': pygame.mixer.Sound('content/die.wav'),
	'hit': pygame.mixer.Sound('content/hit.wav'),
	'point': pygame.mixer.Sound('content/point.wav'),
	'wing': pygame.mixer.Sound('content/wing.wav')
}

def getPipesReady():
	lower_new_pipe_1 = Pipe.LowerPipe()
	upper_new_pipe_1 = Pipe.UpperPipe()

	Pipe.PIPE_LOWER = random.randrange(Pipe.PIPE_GAP + 30, Pipe.MIN_Y - 30)
	Pipe.PIPE_UPPER = Pipe.PIPE_LOWER - Pipe.PIPE_GAP - Pipe.PIPE_HEIGHT

	lower_new_pipe_2 = Pipe.LowerPipe()
	lower_new_pipe_2.x += (WIDTH * 0.5)

	global lower_pipe_list
	lower_pipe_list = [lower_new_pipe_1, lower_new_pipe_2]

	upper_new_pipe_2 = Pipe.UpperPipe()
	upper_new_pipe_2.x += (WIDTH * 0.5)

	global upper_pipe_list
	upper_pipe_list = [upper_new_pipe_1, upper_new_pipe_2]

def generateBackground():
	global background
	background = game_back

def setupBirds():
	global bird_ind
	bird_ind = random.randrange(0, 3)

	global bird
	bird = Bird.Bird(bird_ind)

def remapObjsOnScreen():
	output.blit(random.choice(background), (0, 0))

	for pipe in lower_pipe_list:
		pipe.draw(output)

	for pipe in upper_pipe_list:
		pipe.draw(output)

	bird.draw(output)
	show_score(SCORE)

	pygame.display.update()

def show_score(score):
	score_digits = [int(x) for x in list(str(score))]
	total_width = 0

	for digit in score_digits:
		total_width += digits[digit].get_width()

	x_offset = (WIDTH - total_width) / 2

	for digit in score_digits:
		output.blit(digits[digit], (x_offset, HEIGHT * 0.1))
		x_offset += digits[digit].get_width()

# --------- UPDATER FUNCTIONS ---------

def update_pipe_list():
	Pipe.PIPE_LOWER = random.randrange(Pipe.PIPE_GAP + 30, Pipe.MIN_Y - 30)
	Pipe.PIPE_UPPER = Pipe.PIPE_LOWER - Pipe.PIPE_GAP - Pipe.PIPE_HEIGHT

	if 0 < lower_pipe_list[0].x < 3:
		new_pipe = Pipe.LowerPipe(pipe_col)
		lower_pipe_list.append(new_pipe)

	if lower_pipe_list[0].x < 0 - Pipe.PIPE_WIDTH:
		lower_pipe_list.pop(0)

	if 0 < upper_pipe_list[0].x < 3:
		new_pipe = Pipe.UpperPipe(pipe_col)
		upper_pipe_list.append(new_pipe)

	if upper_pipe_list[0].x < 0 - Pipe.PIPE_WIDTH:
		upper_pipe_list.pop(0)

# What if collision is found?

def detect_collision():
	if (bird.y + Bird.bird_pics[0][1].get_height()) >= 400 - 5:
		if not bird.endOfLife:
			SOUNDS['hit'].play()
		bird.endOfLife = True
		return True
	else:
		for upper_pipe, lower_pipe in zip(upper_pipe_list, lower_pipe_list):
			if (pygame.sprite.collide_mask(bird, upper_pipe) is not None) or \
					(pygame.sprite.collide_mask(bird, lower_pipe) is not None):
				if not bird.endOfLife:
					SOUNDS['hit'].play()
					SOUNDS['die'].play()
				bird.endOfLife = True
				return True
	return False

def check_for_score():
	player_mid_pos = bird.x + Bird.bird_pics[0][1].get_width() / 2

	for pipe in upper_pipe_list:
		pipe_mid_pos = pipe.x + Pipe.PIPE_IMAGES[0][0].get_width() / 2
		if pipe_mid_pos <= player_mid_pos < pipe_mid_pos + 4 and not bird.endOfLife:
			global SCORE
			SCORE += 1
			SOUNDS['point'].play()

# --------- MAIN GAME FUNCTIONS ---------

def start():
	generateBackground()
	setupBirds()

	while True:
		for event in pygame.event.get():
			if event.type == QUIT or (event.type == KEYDOWN and event.key == K_ESCAPE):
				pygame.quit()
				sys.exit()
			elif event.type == KEYDOWN and (event.key == K_UP or event.key == K_SPACE): return

		output.blit(random.choice(background), (0, 0))
		bird.draw(output)
		bird.update()
		pygame.display.update()
		CLOCK.tick(FPS)

def main_loop():
	global SCORE
	SCORE = 0

	start()
	getPipesReady()
	while True:
		for event in pygame.event.get():
			if event.type == QUIT or (event.type == KEYDOWN and event.key == K_ESCAPE):
				pygame.quit()
				sys.exit()
			elif (event.type == MOUSEBUTTONDOWN or (event.type == KEYDOWN and
														(event.key == K_UP or event.key == K_SPACE))) and \
					(not bird.y <= 0) and not bird.endOfLife:
				bird.jump()
				SOUNDS['wing'].play()

			elif bird.endOfLife and (event.type == KEYDOWN and (event.key == K_UP or event.key == K_SPACE) or
											 event.type == MOUSEBUTTONDOWN):
				main_loop()

		update_pipe_list()
		check_for_score()

		for pipe in lower_pipe_list:
			pipe.update()

		for pipe in upper_pipe_list:
			pipe.update()

		bird.update()

		if detect_collision():
			for upper_pipe, lower_pipe in zip(upper_pipe_list, lower_pipe_list):
				upper_pipe.x_vel = 0
				lower_pipe.x_vel = 0

		remapObjsOnScreen()

		CLOCK.tick(FPS)

main_loop()  # GAME INITIALIZER


# --------------------------------------------------------
# Adapted from  https://github.com/microsoft/unilm/tree/master/beit
# --------------------------------------------------------
import math
import os
from functools import partial
import logging
import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import drop_path, to_2tuple, trunc_normal_

#from .transformer import PatchDropout
from .layers.rope import VisionRotaryEmbedding, VisionRotaryEmbeddingFast

import numpy as np

if os.getenv('ENV_TYPE') == 'deepspeed':
    from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint
else:
    from torch.utils.checkpoint import checkpoint
'''
try:
    import xformers.ops as xops
except ImportError:
    xops = None
    print("Please 'pip install xformers'")
'''

class PatchDropout(nn.Module):
    """
    https://arxiv.org/abs/2212.00794
    """

    def __init__(self, prob, exclude_first_token=True):
        super().__init__()
        assert 0 <= prob < 1.
        self.prob = prob
        self.exclude_first_token = exclude_first_token  # exclude CLS token
        logging.info(f"os.getenv('RoPE')={os.getenv('RoPE')}")

    def forward(self, x):
        if not self.training or self.prob == 0.:
            return x

        if self.exclude_first_token:
            cls_tokens, x = x[:, :1], x[:, 1:]
        else:
            cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])

        batch = x.size()[0]
        num_tokens = x.size()[1]

        batch_indices = torch.arange(batch)
        batch_indices = batch_indices[..., None]

        keep_prob = 1 - self.prob
        num_patches_keep = max(1, int(num_tokens * keep_prob))

        rand = torch.randn(batch, num_tokens)
        patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices

        x = x[batch_indices, patch_indices_keep]

        if self.exclude_first_token:
            x = torch.cat((cls_tokens, x), dim=1)

        if self.training and os.getenv('RoPE') == '1':
            return x, patch_indices_keep

        return x

class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """
    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)

    def extra_repr(self) -> str:
        return 'p={}'.format(self.drop_prob)

class Mlp(nn.Module):
    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        drop=0.,
        subln=False,

        ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()

        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()

        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        # x = self.drop(x)
        # commit this for the orignal BERT implement
        x = self.ffn_ln(x)

        x = self.fc2(x)
        x = self.drop(x)
        return x

class SwiGLU(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.,
                norm_layer=nn.LayerNorm, subln=False):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features

        self.w1 = nn.Linear(in_features, hidden_features)
        self.w2 = nn.Linear(in_features, hidden_features)

        self.act = act_layer()
        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
        self.w3 = nn.Linear(hidden_features, out_features)

        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x1 = self.w1(x)
        x2 = self.w2(x)
        hidden = self.act(x1) * x2
        x = self.ffn_ln(hidden)
        x = self.w3(x)
        x = self.drop(x)
        return x

class SwiGLU_sep(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.,
                norm_layer=nn.LayerNorm, subln=False, num_det_tokens=100):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features

        self.w1_patch = nn.Linear(in_features, hidden_features)
        self.w1_det = nn.Linear(in_features, hidden_features)
        self.w2_patch = nn.Linear(in_features, hidden_features)
        self.w2_det = nn.Linear(in_features, hidden_features)

        self.act = act_layer()
        self.ffn_ln_patch = norm_layer(hidden_features) if subln else nn.Identity()
        self.ffn_ln_det = norm_layer(hidden_features) if subln else nn.Identity()
        self.w3_patch = nn.Linear(hidden_features, out_features)
        self.w3_det = nn.Linear(hidden_features, out_features)

        self.drop = nn.Dropout(drop)

        self.num_det_tokens = num_det_tokens

    def forward(self, x):
        x_patch = x[:,:-self.num_det_tokens,:]
        x_det = x[:,-self.num_det_tokens:,:]
        # x1 = self.w1(x)
        x1_patch = self.w1_patch(x_patch)
        x1_det = self.w1_det(x_det)
        # x2 = self.w2(x)
        x2_patch = self.w2_patch(x_patch)
        x2_det = self.w2_det(x_det)
        # hidden = self.act(x1) * x2
        hidden_patch = self.act(x1_patch) * x2_patch
        hidden_det = self.act(x1_det) * x2_det
        # x = self.ffn_ln(hidden)
        x_patch = self.ffn_ln_patch(hidden_patch)
        x_det = self.ffn_ln_det(hidden_det)
        # x = self.w3(x)
        x_patch = self.w3_patch(x_patch)
        x_det = self.w3_det(x_det)
        x = torch.cat((x_patch, x_det), dim=1)
        x = self.drop(x)
        return x

def memory_efficient_attention_pytorch(query, key, value, attn_bias=None, p=0., scale=None, attn_mask=None):
    # query     [batch, seq_len, n_head, head_dim]
    # key       [batch, seq_len, n_head, head_dim]
    # value     [batch, seq_len, n_head, head_dim]
    # attn_bias [batch, n_head, seq_len, seq_len]

    if scale is None:
        scale = 1 / query.shape[-1] ** 0.5

    # BLHC -> BHLC
    #query = query.transpose(1, 2)
    #key = key.transpose(1, 2)
    #value = value.transpose(1, 2)

    query = query * scale
    # BHLC @ BHCL -> BHLL
    attn = query @ key.transpose(-2, -1)
    if attn_bias is not None:
        attn = attn + attn_bias

    if attn_mask is not None:
        attn_mask = attn_mask.bool()
        attn = attn.masked_fill(~attn_mask, float("-inf"))

    attn = attn.softmax(-1)
    attn = F.dropout(attn, p)
    # BHLL @ BHLC -> BHLC
    out = attn @ value
    # BHLC -> BLHC
    out = out.transpose(1, 2)
    return out

class Attention(nn.Module):
    def __init__(
            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,proj_drop=0., window_size=None,
            attn_head_dim=None, xattn=False, rope=None, subln=False, norm_layer=nn.LayerNorm, num_det_tokens=100):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.num_det_tokens = num_det_tokens
        if attn_head_dim is not None:
            head_dim = attn_head_dim
        all_head_dim = head_dim * self.num_heads
        self.scale = qk_scale or head_dim ** -0.5

        self.subln = subln
        if self.subln:
            self.q_proj = nn.Linear(dim, all_head_dim, bias=False)
            self.k_proj = nn.Linear(dim, all_head_dim, bias=False)
            self.v_proj = nn.Linear(dim, all_head_dim, bias=False)
        else:
            self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)

        if qkv_bias:
            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
        else:
            self.q_bias = None
            self.v_bias = None

        if window_size:
            self.window_size = window_size
            self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
            self.relative_position_bias_table = nn.Parameter(
                torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
            # cls to token & token 2 cls & cls to cls

            # get pair-wise relative position index for each token inside the window
            coords_h = torch.arange(window_size[0])
            coords_w = torch.arange(window_size[1])
            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
            coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
            relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
            relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
            relative_coords[:, :, 1] += window_size[1] - 1
            relative_coords[:, :, 0] *= 2 * window_size[1] - 1
            relative_position_index = \
                torch.zeros(size=(window_size[0] * window_size[1] + 1, ) * 2, dtype=relative_coords.dtype)
            relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
            relative_position_index[0, 0:] = self.num_relative_distance - 3
            relative_position_index[0:, 0] = self.num_relative_distance - 2
            relative_position_index[0, 0] = self.num_relative_distance - 1

            self.register_buffer("relative_position_index", relative_position_index)
        else:
            self.window_size = None
            self.relative_position_bias_table = None
            self.relative_position_index = None

        self.attn_drop = nn.Dropout(attn_drop)
        self.inner_attn_ln = norm_layer(all_head_dim) if subln else nn.Identity()
        # self.proj = nn.Linear(all_head_dim, all_head_dim)
        self.proj = nn.Linear(all_head_dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)
        self.xattn = xattn
        self.xattn_drop = attn_drop

        self.rope = rope

    def forward(self, x, rel_pos_bias=None, attn_mask=None):
        B, N, C = x.shape
        if self.subln:
            q = F.linear(input=x, weight=self.q_proj.weight, bias=self.q_bias)
            # np.savetxt('q.csv',q[0].cpu().detach().numpy(),fmt='%.4f',delimiter=',')
            k = F.linear(input=x, weight=self.k_proj.weight, bias=None)
            v = F.linear(input=x, weight=self.v_proj.weight, bias=self.v_bias)

            q = q.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)     # B, num_heads, N, C
            k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
            v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
        else:

            qkv_bias = None
            if self.q_bias is not None:
                qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))

            qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
            qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)   # 3, B, num_heads, N, C
            q, k, v = qkv[0], qkv[1], qkv[2]

        if self.rope:
            # slightly fast impl
            q_t = q[:, :, 1:-self.num_det_tokens, :]
            ro_q_t = self.rope(q_t)
            q = torch.cat((q[:, :, :1, :], ro_q_t, q[:, :, -self.num_det_tokens:, :]), -2).type_as(v)

            k_t = k[:, :, 1:-self.num_det_tokens, :]
            ro_k_t = self.rope(k_t)
            k = torch.cat((k[:, :, :1, :], ro_k_t, k[:, :, -self.num_det_tokens:, :]), -2).type_as(v)

        if self.xattn:
            #q = q.permute(0, 2, 1, 3)   # B, num_heads, N, C -> B, N, num_heads, C
            #k = k.permute(0, 2, 1, 3)
            #v = v.permute(0, 2, 1, 3)

            #x = xops.memory_efficient_attention(
            #    q, k, v,
            #    p=self.xattn_drop,
            #    scale=self.scale,
            #    )
            x = memory_efficient_attention_pytorch(q, k, v, p=self.xattn_drop, scale=self.scale, attn_mask=attn_mask)

            x = x.reshape(B, N, -1)
            x = self.inner_attn_ln(x)
            x = self.proj(x)
            x = self.proj_drop(x)
        else:
            q = q * self.scale
            attn = (q @ k.transpose(-2, -1))

            if self.relative_position_bias_table is not None:
                relative_position_bias = \
                    self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
                        self.window_size[0] * self.window_size[1] + 1,
                        self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
                relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
                attn = attn + relative_position_bias.unsqueeze(0).type_as(attn)

            if rel_pos_bias is not None:
                attn = attn + rel_pos_bias.type_as(attn)

            if attn_mask is not None:
                attn_mask = attn_mask.bool()
                attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf"))

            attn = attn.softmax(dim=-1)
            attn = self.attn_drop(attn)

            x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
            x = self.inner_attn_ln(x)
            x = self.proj(x)
            x = self.proj_drop(x)
        return x

class Attention_sep(nn.Module):
    def __init__(
            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,proj_drop=0., window_size=None,
            attn_head_dim=None, xattn=False, rope=None, subln=False, norm_layer=nn.LayerNorm, num_det_tokens=100):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.num_det_tokens = num_det_tokens
        if attn_head_dim is not None:
            head_dim = attn_head_dim
        all_head_dim = head_dim * self.num_heads
        self.scale = qk_scale or head_dim ** -0.5

        self.subln = subln
        if self.subln:
            self.q_proj_patch = nn.Linear(dim, all_head_dim, bias=False)
            self.k_proj_patch = nn.Linear(dim, all_head_dim, bias=False)
            self.v_proj_patch = nn.Linear(dim, all_head_dim, bias=False)
            self.q_proj_det = nn.Linear(dim, all_head_dim, bias=False)
            self.k_proj_det = nn.Linear(dim, all_head_dim, bias=False)
            self.v_proj_det = nn.Linear(dim, all_head_dim, bias=False)
        else:
            self.qkv_patch = nn.Linear(dim, all_head_dim * 3, bias=False)
            self.qkv_det = nn.Linear(dim, all_head_dim * 3, bias=False)

        if qkv_bias:
            self.q_bias_patch = nn.Parameter(torch.zeros(all_head_dim))
            self.v_bias_patch = nn.Parameter(torch.zeros(all_head_dim))
            self.q_bias_det = nn.Parameter(torch.zeros(all_head_dim))
            self.v_bias_det = nn.Parameter(torch.zeros(all_head_dim))
        else:
            self.q_bias_patch = None
            self.v_bias_patch = None
            self.q_bias_det = None
            self.v_bias_det = None

        if window_size:
            raise NotImplementedError
        else:
            self.window_size = None
            self.relative_position_bias_table = None
            self.relative_position_index = None

        self.attn_drop = nn.Dropout(attn_drop)
        self.inner_attn_ln = norm_layer(all_head_dim) if subln else nn.Identity()
        # self.proj = nn.Linear(all_head_dim, all_head_dim)
        self.proj_patch = nn.Linear(all_head_dim, dim)
        self.proj_det = nn.Linear(all_head_dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)
        self.xattn = xattn
        self.xattn_drop = attn_drop

        self.rope = rope

    def forward(self, x, rel_pos_bias=None, attn_mask=None):
        B, N, C = x.shape
        if self.subln:
            q_patch = F.linear(input=x[:,:-self.num_det_tokens,:], weight=self.q_proj_patch.weight, bias=self.q_bias_patch)
            k_patch = F.linear(input=x[:,:-self.num_det_tokens,:], weight=self.k_proj_patch.weight, bias=None)
            v_patch = F.linear(input=x[:,:-self.num_det_tokens,:], weight=self.v_proj_patch.weight, bias=self.v_bias_patch)

            q_det = F.linear(input=x[:,-self.num_det_tokens:,:], weight=self.q_proj_det.weight, bias=self.q_bias_det)
            k_det = F.linear(input=x[:,-self.num_det_tokens:,:], weight=self.k_proj_det.weight, bias=None)
            v_det = F.linear(input=x[:,-self.num_det_tokens:,:], weight=self.v_proj_det.weight, bias=self.v_bias_det)

            q = torch.cat((q_patch, q_det), dim=1)
            k = torch.cat((k_patch, k_det), dim=1)
            v = torch.cat((v_patch, v_det), dim=1)

            q = q.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)     # B, num_heads, N, C
            k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
            v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
        else:
            qkv_bias_patch = None
            qkv_bias_det = None
            if self.q_bias_patch is not None:
                qkv_bias_patch = torch.cat((self.q_bias_patch, torch.zeros_like(self.v_bias_patch, requires_grad=False), self.v_bias_patch))
            if self.q_bias_det is not None:
                qkv_bias_det = torch.cat((self.q_bias_det, torch.zeros_like(self.v_bias_det, requires_grad=False), self.v_bias_det))

            qkv_patch = F.linear(input=x[:,:-self.num_det_tokens,:], weight=self.qkv_patch.weight, bias=qkv_bias_patch)
            qkv_det = F.linear(input=x[:,-self.num_det_tokens:,:], weight=self.qkv_det.weight, bias=qkv_bias_det)
            qkv = torch.cat((qkv_patch, qkv_det), dim=1)
            qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)   # 3, B, num_heads, N, C
            q, k, v = qkv[0], qkv[1], qkv[2]

        if self.rope:
            # slightly fast impl
            q_t = q[:, :, 1:-self.num_det_tokens, :]
            ro_q_t = self.rope(q_t)
            q = torch.cat((q[:, :, :1, :], ro_q_t, q[:, :, -self.num_det_tokens:, :]), -2).type_as(v)

            k_t = k[:, :, 1:-self.num_det_tokens, :]
            ro_k_t = self.rope(k_t)
            k = torch.cat((k[:, :, :1, :], ro_k_t, k[:, :, -self.num_det_tokens:, :]), -2).type_as(v)

        if self.xattn:
            #q = q.permute(0, 2, 1, 3)   # B, num_heads, N, C -> B, N, num_heads, C
            #k = k.permute(0, 2, 1, 3)
            #v = v.permute(0, 2, 1, 3)

            #x = xops.memory_efficient_attention(
            #    q, k, v,
            #    p=self.xattn_drop,
            #    scale=self.scale,
            #    )
            x = memory_efficient_attention_pytorch(q, k, v, p=self.xattn_drop, scale=self.scale, attn_mask=attn_mask)

            x = x.reshape(B, N, -1)
            x = self.inner_attn_ln(x)
            x_patch = self.proj_patch(x[:,:-self.num_det_tokens,:])
            x_det = self.proj_det(x[:,-self.num_det_tokens:,:])
            x = torch.cat((x_patch, x_det), dim=1)
            x = self.proj_drop(x)
        else:
            q = q * self.scale
            attn = (q @ k.transpose(-2, -1))

            if self.relative_position_bias_table is not None:
                relative_position_bias = \
                    self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
                        self.window_size[0] * self.window_size[1] + 1,
                        self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
                relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
                attn = attn + relative_position_bias.unsqueeze(0).type_as(attn)

            if rel_pos_bias is not None:
                attn = attn + rel_pos_bias.type_as(attn)

            if attn_mask is not None:
                attn_mask = attn_mask.bool()
                attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf"))

            attn = attn.softmax(dim=-1)
            attn = self.attn_drop(attn)

            x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
            x = self.inner_attn_ln(x)
            x = self.proj(x)
            x = self.proj_drop(x)
        return x

class Block(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
                 window_size=None, attn_head_dim=None, xattn=False, rope=None, postnorm=False,
                 subln=False, naiveswiglu=False, num_det_tokens=100, partial_finetune=False, partial_finetune_attn=False):
        super().__init__()
        self.norm1 = norm_layer(dim)
        if partial_finetune_attn:
            self.attn = Attention_sep(
                dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
                attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim,
                xattn=xattn, rope=rope, subln=subln, norm_layer=norm_layer, num_det_tokens=num_det_tokens)
        else:
            self.attn = Attention(
                dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
                attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim,
                xattn=xattn, rope=rope, subln=subln, norm_layer=norm_layer, num_det_tokens=num_det_tokens)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)

        if naiveswiglu:
            if partial_finetune:
                self.mlp = SwiGLU_sep(
                    in_features=dim,
                    hidden_features=mlp_hidden_dim,
                    subln=subln,
                    norm_layer=norm_layer,
                    num_det_tokens=num_det_tokens,
                )
            else:
                self.mlp = SwiGLU(
                    in_features=dim,
                    hidden_features=mlp_hidden_dim,
                    subln=subln,
                    norm_layer=norm_layer,
                )
        else:
            self.mlp = Mlp(
                in_features=dim,
                hidden_features=mlp_hidden_dim,
                act_layer=act_layer,
                subln=subln,
                drop=drop
            )

        if init_values is not None and init_values > 0:
            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
        else:
            self.gamma_1, self.gamma_2 = None, None

        self.postnorm = postnorm

    def forward(self, x, rel_pos_bias=None, attn_mask=None):
        if self.gamma_1 is None:
            if self.postnorm:
                x = x + self.drop_path(self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
                x = x + self.drop_path(self.norm2(self.mlp(x)))
            else:
                x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
                x = x + self.drop_path(self.mlp(self.norm2(x)))
        else:
            if self.postnorm:
                x = x + self.drop_path(self.gamma_1 * self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
                x = x + self.drop_path(self.gamma_2 * self.norm2(self.mlp(x)))
            else:
                x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
                x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
        return x

class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x, **kwargs):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraints
        #assert H == self.img_size[0] and W == self.img_size[1], \
        #    f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x).flatten(2).transpose(1, 2)  # [1, 14*14, vision_cfg.width:768]
        return x

class RelativePositionBias(nn.Module):

    def __init__(self, window_size, num_heads):
        super().__init__()
        self.window_size = window_size
        self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
        self.relative_position_bias_table = nn.Parameter(
            torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
        # cls to token & token 2 cls & cls to cls

        # get pair-wise relative position index for each token inside the window
        coords_h = torch.arange(window_size[0])
        coords_w = torch.arange(window_size[1])
        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
        relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
        relative_coords[:, :, 1] += window_size[1] - 1
        relative_coords[:, :, 0] *= 2 * window_size[1] - 1
        relative_position_index = \
            torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
        relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
        relative_position_index[0, 0:] = self.num_relative_distance - 3
        relative_position_index[0:, 0] = self.num_relative_distance - 2
        relative_position_index[0, 0] = self.num_relative_distance - 1

        self.register_buffer("relative_position_index", relative_position_index)

    def forward(self):
        relative_position_bias = \
            self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
                self.window_size[0] * self.window_size[1] + 1,
                self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
        return relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww

class EVAVisionTransformer(nn.Module):
    """ Vision Transformer with support for patch or hybrid CNN input stage
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, num_det_tokens=100, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,drop_path_rate=0.,
                 norm_layer=nn.LayerNorm, init_values=None, patch_dropout=0.,use_abs_pos_emb=True, use_rel_pos_bias=False,
                 use_shared_rel_pos_bias=False, rope=False,use_mean_pooling=True, init_scale=0.001, grad_checkpointing=False,
                 xattn=False, postnorm=False,pt_hw_seq_len=16, intp_freq=False, naiveswiglu=False, subln=False, attn_mask=False,
                 partial_finetune=False, partial_finetune_attn=False):
        super().__init__()
        #self.img_size = img_size
        if isinstance(img_size,tuple):
            self.img_size = img_size
        else:
            self.img_size = to_2tuple(img_size)
        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models

        self.num_det_tokens = num_det_tokens

        self.patch_embed = PatchEmbed(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        self.patch_size = patch_size
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        # self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        if use_abs_pos_emb:
            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        else:
            self.pos_embed = None
        self.pos_drop = nn.Dropout(p=drop_rate)

        if use_shared_rel_pos_bias:
            self.rel_pos_bias = RelativePositionBias(window_size=self.patch_embed.patch_shape, num_heads=num_heads)
        else:
            self.rel_pos_bias = None

        if rope:
            half_head_dim = embed_dim // num_heads // 2
            hw_seq_len = img_size // patch_size
            self.rope = VisionRotaryEmbeddingFast(
                dim=half_head_dim,
                pt_seq_len=pt_hw_seq_len,
                ft_seq_len=hw_seq_len if intp_freq else None,
                # patch_dropout=patch_dropout
            )
        else:
            self.rope = None

        self.naiveswiglu = naiveswiglu
        self.partial_finetune = partial_finetune
        self.partial_finetune_attn = partial_finetune_attn

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        self.use_rel_pos_bias = use_rel_pos_bias
        self.blocks = nn.ModuleList([
            Block(
                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
                init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None,
                xattn=xattn, rope=self.rope, postnorm=postnorm, subln=subln, naiveswiglu=naiveswiglu,
                num_det_tokens = self.num_det_tokens, partial_finetune=partial_finetune, partial_finetune_attn = partial_finetune_attn)
            for i in range(depth)])
        self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
        self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None
        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()

        if self.pos_embed is not None:
            trunc_normal_(self.pos_embed, std=.02)

        trunc_normal_(self.cls_token, std=.02)
        # trunc_normal_(self.mask_token, std=.02)

        self.apply(self._init_weights)
        self.fix_init_weight()

        if isinstance(self.head, nn.Linear):
            trunc_normal_(self.head.weight, std=.02)
            self.head.weight.data.mul_(init_scale)
            self.head.bias.data.mul_(init_scale)

        # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
        self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()

        self.grad_checkpointing = grad_checkpointing
        self.has_mid_pe = False

        self.attn_mask = attn_mask

    def fix_init_weight(self):
        def rescale(param, layer_id):
            param.div_(math.sqrt(2.0 * layer_id))

        for layer_id, layer in enumerate(self.blocks):
            if self.partial_finetune_attn:
                rescale(layer.attn.proj_patch.weight.data, layer_id + 1)
                rescale(layer.attn.proj_det.weight.data, layer_id + 1)
            else:
                rescale(layer.attn.proj.weight.data, layer_id + 1)
            if self.naiveswiglu:
                if self.partial_finetune:
                    rescale(layer.mlp.w3_patch.weight.data, layer_id + 1)
                    rescale(layer.mlp.w3_det.weight.data, layer_id + 1)
                else:
                    rescale(layer.mlp.w3.weight.data, layer_id + 1)
            else:
                rescale(layer.mlp.fc2.weight.data, layer_id + 1)

    def get_cast_dtype(self) -> torch.dtype:
        return self.blocks[0].mlp.fc2.weight.dtype

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def get_num_layers(self):
        return len(self.blocks)

    def lock(self, unlocked_groups=0):
        assert unlocked_groups == 0, 'partial locking not currently supported for this model'
        for param in self.parameters():
            param.requires_grad = False

    @torch.jit.ignore
    def set_grad_checkpointing(self, enable=True):
        self.grad_checkpointing = enable

    @torch.jit.ignore
    def no_weight_decay(self):
        return {'pos_embed', 'cls_token'}

    def get_classifier(self):
        return self.head

    def reset_classifier(self, num_classes, global_pool=''):
        self.num_classes = num_classes
        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()

    def finetune_det(self, img_size=[800, 1344], mid_pe_size=None, use_checkpoint=False):
        # import pdb;pdb.set_trace()

        import math
        g = math.pow(self.pos_embed.size(1) - 1, 0.5)
        if int(g) - g != 0:
            self.pos_embed = torch.nn.Parameter(self.pos_embed[:, 1:, :])

        # self.det_token_num = det_token_num
        self.det_token = nn.Parameter(torch.zeros(1, self.num_det_tokens, self.embed_dim))
        self.det_token = trunc_normal_(self.det_token, std=.02)
        cls_pos_embed = self.pos_embed[:, 0, :]
        cls_pos_embed = cls_pos_embed[:,None]
        det_pos_embed = torch.zeros(1, self.num_det_tokens, self.embed_dim)
        det_pos_embed = trunc_normal_(det_pos_embed, std=.02)
        patch_pos_embed = self.pos_embed[:, 1:, :]
        patch_pos_embed = patch_pos_embed.transpose(1,2)
        B, E, Q = patch_pos_embed.shape
        P_H, P_W = self.img_size[0] // self.patch_size, self.img_size[1] // self.patch_size
        patch_pos_embed = patch_pos_embed.view(B, E, P_H, P_W)
        H, W = img_size
        new_P_H, new_P_W = H//self.patch_size, W//self.patch_size
        # torch.Size([1, 768, 50, 84])
        patch_pos_embed = nn.functional.interpolate(patch_pos_embed, size=(new_P_H,new_P_W), mode='bicubic', align_corners=False)
        patch_pos_embed = patch_pos_embed.flatten(2).transpose(1,2) # torch.Size([1, 4200, 768])
        self.pos_embed = torch.nn.Parameter(torch.cat((cls_pos_embed, patch_pos_embed), dim=1), requires_grad=not (self.partial_finetune or self.partial_finetune_attn))
        self.det_pos_embed = torch.nn.Parameter(det_pos_embed)

        _, head_dim = self.rope.freqs_cos.shape
        freqs_cos = self.rope.freqs_cos.transpose(0,1).view(B, head_dim, P_H, P_W)
        freqs_sin = self.rope.freqs_sin.transpose(0,1).view(B, head_dim, P_H, P_W)
        freqs_cos = nn.functional.interpolate(freqs_cos, size=(new_P_H,new_P_W), mode='bicubic', align_corners=False) #torch.Size([1, 64, 50, 84])
        freqs_sin = nn.functional.interpolate(freqs_sin, size=(new_P_H,new_P_W), mode='bicubic', align_corners=False)
        self.register_buffer("freqs_cos", freqs_cos.squeeze().flatten(1).transpose(0,1))
        self.register_buffer("freqs_sin", freqs_sin.squeeze().flatten(1).transpose(0,1))

        self.img_size = img_size
        if mid_pe_size == None:
            self.has_mid_pe = False
            print('No mid pe')
        else:
            print('Has mid pe')
            self.mid_pos_embed = nn.Parameter(torch.zeros(self.depth - 1, 1, 1 + (mid_pe_size[0] * mid_pe_size[1] // self.patch_size ** 2) + self.num_det_tokens, self.embed_dim))
            trunc_normal_(self.mid_pos_embed, std=.02)
            self.has_mid_pe = True
            self.mid_pe_size = mid_pe_size
        self.use_checkpoint=use_checkpoint

    def InterpolateInitPosEmbed(self, pos_embed, img_size=(800, 1344)):
        # import pdb;pdb.set_trace()
        _, head_dim = self.rope.freqs_cos.shape

        cls_pos_embed = pos_embed[:, 0, :]
        cls_pos_embed = cls_pos_embed[:,None]
        patch_pos_embed = pos_embed[:, 1:, :] # [1,4200,768]
        patch_pos_embed = patch_pos_embed.transpose(1,2)
        freqs_cos = self.freqs_cos.transpose(0,1)
        freqs_sin = self.freqs_sin.transpose(0,1)
        B, E, Q = patch_pos_embed.shape

        P_H, P_W = self.img_size[0] // self.patch_size, self.img_size[1] // self.patch_size
        patch_pos_embed = patch_pos_embed.view(B, E, P_H, P_W)
        freqs_cos = freqs_cos.view(B, head_dim, P_H, P_W)
        freqs_sin = freqs_sin.view(B, head_dim, P_H, P_W)

        # P_H, P_W = self.img_size[0] // self.patch_size, self.img_size[1] // self.patch_size
        # patch_pos_embed = patch_pos_embed.view(B, E, P_H, P_W)

        H, W = img_size
        new_P_H, new_P_W = H//self.patch_size, W//self.patch_size
        patch_pos_embed = nn.functional.interpolate(patch_pos_embed, size=(new_P_H,new_P_W), mode='bicubic', align_corners=False)
        patch_pos_embed = patch_pos_embed.flatten(2).transpose(1, 2)
        scale_pos_embed = torch.cat((cls_pos_embed, patch_pos_embed), dim=1)

        freqs_cos = nn.functional.interpolate(freqs_cos, size=(new_P_H,new_P_W), mode='bicubic', align_corners=False)
        freqs_sin = nn.functional.interpolate(freqs_sin, size=(new_P_H,new_P_W), mode='bicubic', align_corners=False)
        self.rope.freqs_cos = freqs_cos.squeeze().flatten(1).transpose(0,1)
        self.rope.freqs_sin = freqs_sin.squeeze().flatten(1).transpose(0,1)

        return scale_pos_embed

    def get_rect_mask(self, patch_seq_len=196):

        total_seq_len = 1 + patch_seq_len + self.num_det_tokens
        mask = torch.ones(size=[total_seq_len, total_seq_len], device='cuda')
        mask[:-self.num_det_tokens, -self.num_det_tokens:] = 0

        return mask

    def forward_features(self, x, return_all_features=False):

        B, H, W = x.shape[0], x.shape[2], x.shape[3]

        x = self.patch_embed(x)  # [B,3,H,W] -> torch.Size([B, n_patch, embed_dim=768])
        batch_size, seq_len, _ = x.size()

        # interpolate init pe
        if (self.pos_embed.shape[1] - 1) != seq_len:
            temp_pos_embed = self.InterpolateInitPosEmbed(self.pos_embed, img_size=(H,W))
        else:
            temp_pos_embed = self.pos_embed
            self.rope.freqs_cos = self.freqs_cos
            self.rope.freqs_sin = self.freqs_sin

        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        det_token = self.det_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x, det_token), dim=1)  # [B,1+n_patch+100, embed_dim=768]
        if self.pos_embed is not None:
            x = x + torch.cat((temp_pos_embed, self.det_pos_embed), dim=1)
        x = self.pos_drop(x)

        # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
        if os.getenv('RoPE') == '1':
            if self.training and not isinstance(self.patch_dropout, nn.Identity):
                x, patch_indices_keep = self.patch_dropout(x)
                self.rope.forward = partial(self.rope.forward, patch_indices_keep=patch_indices_keep)
            else:
                self.rope.forward = partial(self.rope.forward, patch_indices_keep=None)
                x = self.patch_dropout(x)  # Identity()
        else:
            x = self.patch_dropout(x)

        rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None  # None
        attn_mask = self.get_rect_mask(patch_seq_len=seq_len) if self.attn_mask else None
        for blk in self.blocks:
            if self.grad_checkpointing:
                x = checkpoint(blk, x, (rel_pos_bias, attn_mask))
            else:
                x = blk(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)

        if not return_all_features:  # x:[B, 1+n_patch+100, embed_dim=768]
            x = self.norm(x)
            if self.fc_norm is not None:
                return self.fc_norm(x.mean(1))
            else:
                return x[:, -self.num_det_tokens:, :]  # torch.Size([1, 100, 768])
        return x

    def forward(self, x, return_all_features=False):
        if return_all_features:
            return self.forward_features(x, return_all_features)
        x = self.forward_features(x)
        # x = self.head(x) # [1,512]
        return x


"""Collection of path, file, and directory utilities.

Intended to add a self containing, consistent, and robust layer over the
standard os module, and thus `os` should not be needed outside this module.
There are four main areas of file/directory functionality covered:
introspection, generic removal/creation, validation, and path parsing.

Notes
-----
* All path parameters for every function are normalized (made path independent).
* Functions with names that start with 'ensure' are meant to be used prior to
  file input/output, to ensure failures are caught ahead of time. This is very
  important since virtually nothing in the API will work correctly if data is not
  retrieved correctly.
* Minimal usage of neighboring functions in this module to reduce call-stack
  complexity, thus os.path functions are used repetitively.
"""
import os
import errno
import codecs
import shutil
import inspect
import importlib
import contextlib
from typing import Sequence

# --------------------------------- file and directory inspection ----------------------------------
def is_existent_file(file_path: str) -> bool:
    """Return True if given file exists, else False.

    See Also
    --------
    * Function `ensure_file_exists` for the exception raising counterpart
    """
    try:
        ensure_file_exists(file_path)
        return True
    except FileNotFoundError:
        return False

def is_existent_directory(dir_path: str) -> bool:
    """Return True if given path is dir exists, else False.

    See Also
    --------
    * Function `ensure_directory_exists` for the exception raising counterpart
    """
    try:
        ensure_directory_exists(dir_path)
        return True
    except FileNotFoundError:
        return False

def is_empty_file(file_path: str) -> bool:
    """Return True if file is empty, else False. Raises if file not present."""
    ensure_file_exists(file_path)
    return os.stat(os.path.normpath(file_path.strip(' '))).st_size == 0

def get_path_of_python_source(obj: object) -> str:
    """Return path to the python source code of given object.

    Raises
    ------
    ValueError
        * If inspected path is not a python file
    FileNotFoundError
        * If inspected path does not exist (extremely unlikely)

    Notes
    -----
    * Given object is a package if and only if its name is same as its parent
      directory's, and its source file is an __init__.py. In that case, we
      return the parent directory of the source file (exact package path)
    """
    file_path = os.path.normpath(inspect.getfile(object=obj))
    if os.path.splitext(file_path)[-1].lower() != '.py':
        raise ValueError('Only python source objects can be inspected for path.')
    if not os.path.isfile(file_path):
        raise FileNotFoundError(errno.ENOENT, 'Python source cannot be found ', file_path)

    if file_path.endswith('__init__.py') and os.path.dirname(file_path).endswith(obj.__name__):
        return os.path.normpath(os.path.dirname(file_path))
    return file_path

# -------------------------------- file and directory manipulation ---------------------------------
def create_directory(dir_path: str) -> None:
    """Make single directory (no intermediates) at given path if creatable.

    See Also
    --------
    * For requirements of a successful directory creation, and the corresponding
      exceptions raised, visit function `ensure_directory_is_creatable`.
    """
    dir_path_ = os.path.normpath(dir_path)
    ensure_directory_is_creatable(dir_path_)
    os.mkdir(path=dir_path_)

def create_temp_file():
    """Create temp file."""
    pass

def remove_directory(dir_path: str, ignore_errors=False) -> None:
    """Remove directory and all children at given path, with an option of ignoring any Errors."""
    dir_path_ = os.path.normpath(dir_path.strip(' '))
    if ignore_errors:
        with contextlib.suppress(OSError):
            shutil.rmtree(path=dir_path_, ignore_errors=True)
    else:
        try:
            ensure_directory_exists(dir_path_)
            shutil.rmtree(path=dir_path_, ignore_errors=False)
        except Exception:
            raise OSError(f'Directory \'{dir_path_}\' failed to be removed.')

def remove_file(file_path: str, ignore_errors=False) -> None:
    """Remove file at given path, with an option of ignoring any Errors."""
    file_path_ = os.path.normpath(file_path.strip(' '))
    if ignore_errors:
        with contextlib.suppress(OSError):
            os.remove(file_path_)
    else:
        try:
            ensure_file_exists(file_path_)
            os.remove(file_path_)
        except Exception:
            raise OSError(f'File \'{file_path_}\' failed to be removed.')

@contextlib.contextmanager
def change_directory(dir_path: str) -> str:
    """Temporarily change directory to given path, similar to 'cd' in terminal.

    Examples
    --------
    >>> import os
    >>> from stem_center_analytics import PROJECT_DIR

    >>> previous_dir = os.getcwd()
    >>> with change_directory(PROJECT_DIR): print(os.getcwd() == PROJECT_DIR)
    True
    >>> previous_dir == os.getcwd()
    True
    """
    new_cwd = os.path.normpath(dir_path)
    old_cwd = os.path.normpath(os.getcwd())
    ensure_directory_exists(new_cwd)
    try:
        os.chdir(new_cwd)
        yield new_cwd
    finally:
        os.chdir(old_cwd)

# ---------------------------------------- error checking ------------------------------------------
def ensure_successful_imports(path: str, names: Sequence[str]) -> None:
    """Valid if all object names are successfully imported at given path location.

    Parameters
    ----------
    path : str
        Path at which to import the given names. If file path given,
        names are imported from its parent directory
    names : array-like of str
        Names of objects to import in the order given. This can be the name of
        any importable object, but to ensure successful imports favor
        full, absolute import names such as 'package.module.func', over '.func'

    Raises
    ------
    ImportError
        * If any module in names failed to be imported correctly

    Notes
    -----
    * Imports are lost (goes out of scope) once the function returns
    """
    path_ = os.path.normpath(path.strip(' '))
    dir_to_import_at = path_ if os.path.isdir(path_) else os.path.normpath(os.path.dirname(path_))

    # import given names from the inferred directory location
    unsuccessful_imports = []
    with change_directory(dir_to_import_at):
        for name in names:
            try:
                importlib.import_module(name)
            except ImportError:
                unsuccessful_imports.append(name)
    if unsuccessful_imports:
        raise ImportError(f'Failed to import modules - {tuple(unsuccessful_imports)}.')

def ensure_file_exists(file_path: str) -> None:
    """Valid if path exists as a file.

    Raises
    ------
    FileNotFoundError
        * If path is not an existing file
    ValueError
        * If file does not end with one of `file_type` extensions
    """
    file_path_ = os.path.normpath(file_path.strip(' '))
    if not os.path.isfile(file_path_):
        raise FileNotFoundError(errno.ENOENT, 'No such file', file_path_)

def ensure_directory_exists(dir_path: str) -> None:
    """Valid if path exists as a directory.

    Raises
    ------
    FileNotFoundError
        * If path is a child of a non-existent parent directory
    """
    dir_path_ = os.path.normpath(dir_path.strip(' '))
    if not os.path.isdir(dir_path_):
        raise FileNotFoundError(errno.ENOENT, 'No such directory', dir_path_)

def ensure_file_is_creatable(file_path: str) -> None:
    """Creatable if vacant path in existing directory..

    Raises
    ------
    FileExistsError
        * If file path already exists
    FileNotFoundError
        * If file path is a child of a non-existent parent directory
    ValueError
        * If file does not end with one of `file_type` extensions
    """
    file_path_ = os.path.normpath(file_path.strip(' '))
    parent_dir_path_ = os.path.normpath(os.path.dirname(file_path_))
    if os.path.exists(file_path_):
        raise FileExistsError(errno.ENOENT, 'Cannot create file at non-vacant location', file_path_)
    if not os.path.isdir(parent_dir_path_):
        raise FileNotFoundError(errno.ENOENT, 'Cannot create file in non-existent parent directory',
                                parent_dir_path_)

def ensure_directory_is_creatable(dir_path: str) -> None:
    """Directory is creatable if path is vacant in existing directory.

    Raises
    ------
    FileExistsError
        * If directory path already exists
    FileNotFoundError
        * If path is a child of a non-existent parent directory
    """
    dir_path_ = os.path.normpath(dir_path.strip(' '))
    parent_dir_path_ = os.path.normpath(os.path.dirname(dir_path_))
    if os.path.exists(dir_path_):
        raise FileExistsError(errno.ENOENT, 'Cannot create directory at non-vacant location',
                              dir_path_)
    if not os.path.isdir(parent_dir_path_):
        raise FileNotFoundError(errno.ENOENT,
                                'Cannot create a sub directory in a non-existent parent directory',
                                parent_dir_path_)

def ensure_path_is_absolute(file_path: str) -> None:
    """Valid if file path is absolute.

    Raises
    ------
    ValueError
        * If file path is not absolute

    Notes
    -----
    * In contrast to most other functions in `os_lib` starting with 'ensure',
      given path is NOT checked for existence
    """
    if not os.path.isabs(os.path.normpath(file_path.strip(' '))):
        raise ValueError('File path \'{}\' is not absolute.')

def ensure_valid_file_type(file_path: str, file_type: str) -> None:
    """Valid if file is one of the given file types.

    Raises
    ------
    ValueError
        * If file name has more than one extension
        * If file name doesn't have extension of `file_type`

    Notes
    -----
    * In contrast to most other functions in `os_lib` starting with 'ensure',
      given path is NOT checked for existence
    """
    file_type_ = file_type.lstrip('.')
    file_name, file_extension = get_basename(file_path), get_extension(file_path).lstrip('.')
    if file_name.count('.') != 1:
        raise ValueError(f'File \'{file_name}\' is invalid - file name must have a single extension.')
    if file_extension != file_type_:
        raise ValueError(f'File \'{file_name}\' is invalid - file must be of type \'{file_type_}\'.')

def ensure_valid_file_encoding(file_path: str, encoding: str) -> None:
    """Valid if file at given path can be successfully decoded with given encoding."""
    file_path_ = os.path.normpath(file_path.strip(' '))
    encoding_ = encoding.strip().lower()

    try:
        with codecs.open(file_path_, mode='r', encoding=encoding_, errors='strict'): pass
    except LookupError:
        raise LookupError(f'Unknown encoding - \'{encoding_}\'.') from None
    except UnicodeError:
        raise UnicodeDecodeError(f'File \'{file_path_}\' cannot be decoded as \'{encoding_}\'.')

# ------------------------------------- path name manipulation -------------------------------------
def join_path(path: str, *args: str) -> str:
    """Return path joined with given items.

    Notes
    -----
    * Will break if items contain separators
    """
    base_path_ = os.path.normpath(path.strip(' '))
    path_components = [arg.strip(' ') for arg in args]
    return os.path.normpath(os.path.join(base_path_, *path_components))

def normalize_path(path: str) -> str:
    """Return platform independent path with leading/trailing whitespace removed."""
    return os.path.normpath(path.strip(' '))

def get_extension(path: str) -> str:
    """Return path's (last) extension, including leading '.' separator."""
    return os.path.splitext(os.path.normpath(path.strip(' ')))[-1].lower()

def get_parent_dir(path: str) -> str:
    """Return immediate parent directory of given path."""
    return os.path.normpath(os.path.dirname(os.path.normpath(path.strip(' '))))

def get_basename(path: str, include_extension: bool=True) -> str:
    """Return unix-style basename (eg: '/foo/bar/' => 'bar')."""
    # normalize path, and return the last component
    name = os.path.split(os.path.normpath(path.strip(' ')))[-1]
    return name if include_extension else os.path.splitext(name)[0]


import sys
import os

from mock import patch, Mock, call
from tank_test.tank_test_base import *

# import the test base class
test_python_path = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "..", "python")
)
sys.path.append(test_python_path)
from base_test import TestShotgunUtilsFramework

class TestShotgunNavDataHandler(TestShotgunUtilsFramework):
    """
    Tests for the data handler low level io
    """

    def setUp(self):
        """
        Fixtures setup
        """
        super(TestShotgunNavDataHandler, self).setUp()
        self.shotgun_model = self.framework.import_module("shotgun_model")

    def test_generate_data_request(self):
        """
        Test generate_data_request
        """
        test_path = os.path.join(self.tank_temp, "test_nav_handler.pickle")

        dh = self.shotgun_model.data_handler_nav.ShotgunNavDataHandler(
            root_path="/",
            seed_entity_field="Version.entity",
            entity_fields=None,
            cache_path=test_path,
        )

        dh.load_cache()

        # first let the data handler perform its request
        # and make sure it registers it in an expected way
        mock_data_retriever = Mock()
        mock_data_retriever.execute_nav_expand = Mock(return_value=1234)

        request_id = dh.generate_data_request(mock_data_retriever, "/")

        mock_data_retriever.execute_nav_expand.assert_called_once_with(
            "/", "Version.entity", None
        )
        self.assertEqual(request_id, 1234)

    def test_updates(self):
        """
        Test generate_data_request
        """
        test_path = os.path.join(self.tank_temp, "test_nav_handler.pickle")

        dh = self.shotgun_model.data_handler_nav.ShotgunNavDataHandler(
            root_path="/",
            seed_entity_field="Version.entity",
            entity_fields=None,
            cache_path=test_path,
        )

        dh.load_cache()

        # apply first data chunk
        sg_data = {
            "path": "/",
            "children": [
                {"label": "foo", "has_children": True, "path": "/foo"},
                {"label": "bar", "has_children": True, "path": "/bar"},
                {"label": "baz", "has_children": False, "path": "/baz"},
            ],
        }

        diff = dh.update_data(sg_data)

        # we are getting a diff of two added nodes back.
        self.assertEqual(len(diff), 3)

        self.assertEqual(diff[0]["mode"], dh.ADDED)
        self.assertEqual(diff[1]["mode"], dh.ADDED)
        self.assertEqual(diff[2]["mode"], dh.ADDED)

        foo = diff[0]["data"]
        bar = diff[1]["data"]
        baz = diff[1]["data"]

        self.assertEqual(foo.unique_id, "/foo")
        self.assertEqual(foo.field, None)
        self.assertEqual(
            foo.shotgun_data, {"has_children": True, "label": "foo", "path": "/foo"}
        )
        self.assertEqual(foo.parent, None)
        self.assertEqual(foo.is_leaf(), False)

        self.assertEqual(dh.get_data_item_from_uid("/foo"), foo)

        # now apply an update
        sg_data = {
            "path": "/",
            "children": [
                {"label": "foo2", "has_children": True, "path": "/foo2"},
                {"label": "bar", "has_children": True, "path": "/bar"},
                {"label": "baz", "has_children": False, "path": "/baz"},
            ],
        }
        diff = dh.update_data(sg_data)

        self.assertEqual(len(diff), 2)
        self.assertEqual(diff[0]["mode"], dh.ADDED)
        self.assertEqual(diff[1]["mode"], dh.DELETED)

        dh.save_cache()
        dh.unload_cache()
        self.assertFalse(dh.is_cache_loaded())
        dh.load_cache()

        self.assertEqual(dh.get_data_item_from_uid("/foo2"), diff[0]["data"])
        self.assertEqual(dh.get_data_item_from_uid("/Prop"), None)

    def test_generate_child_nodes(self):
        """
        Tests child node generation from cache
        """

        test_path = os.path.join(self.tank_temp, "test_nav_handler.pickle")

        dh = self.shotgun_model.data_handler_nav.ShotgunNavDataHandler(
            root_path="/",
            seed_entity_field="Version.entity",
            entity_fields=None,
            cache_path=test_path,
        )

        dh.load_cache()

        # apply first data chunk
        sg_data = {
            "path": "/",
            "children": [
                {"label": "foo", "has_children": True, "path": "/foo"},
                {"label": "bar", "has_children": True, "path": "/bar"},
                {"label": "baz", "has_children": False, "path": "/baz"},
            ],
        }

        dh.update_data(sg_data)

        sg_data = {
            "path": "/foo",
            "children": [
                {"label": "smith", "has_children": False, "path": "/foo/smith"},
                {"label": "jones", "has_children": False, "path": "/foo/jones"},
                {"label": "brown", "has_children": False, "path": "/foo/brown"},
            ],
        }

        dh.update_data(sg_data)

        callback = Mock()

        dh.generate_child_nodes(None, None, callback)

        calls = [
            call(None, dh.get_data_item_from_uid("/foo")),
            call(None, dh.get_data_item_from_uid("/bar")),
            call(None, dh.get_data_item_from_uid("/baz")),
        ]

        # Python 2 and 3 have different ordering internally
        # for their keys, so we can't expect the items
        # be generated in the same order.
        callback.assert_has_calls(calls, any_order=True)

        # and check children
        callback = Mock()
        dh.generate_child_nodes("/foo", None, callback)
        calls = [
            call(None, dh.get_data_item_from_uid("/foo/smith")),
            call(None, dh.get_data_item_from_uid("/foo/brown")),
            call(None, dh.get_data_item_from_uid("/foo/jones")),
        ]
        # Python 2 and 3 have different ordering internally
        # for their keys, so we can't expect the items
        # be generated in the same order.
        callback.assert_has_calls(calls, any_order=True)


import mock
import mongomock

from app import create_app
from app.crawl.house_crawler import HouseCrawler
from app.mongodb.mongodb_manager import MongoDbManager

@mock.patch('app.mongodb.mongodb_manager.MongoDbManager.get_client')
@mock.patch('app.crawl.house_crawler.HouseCrawler._create_payloads')
def test_should_insertion_operate_properly(mock_create_payloads, mock_get_client):
    mock_create_payloads.return_value = [
        {'is_new_list': '1',
          'type': '1',
          'kind': '0',
          'searchtype': 1,
          'regionid': '1',
          'firstRow': 30,
          'totalRows': 12503},
         {'is_new_list': '1',
          'type': '1',
          'kind': '0',
          'searchtype': 1,
          'regionid': '1',
          'firstRow': 30,
          'totalRows': 12503},
         {'is_new_list': '1',
          'type': '1',
          'kind': '0',
          'searchtype': 1,
          'regionid': '1',
          'firstRow': 60,
          'totalRows': 12503},
    ]

    mock_get_client.return_value = mongomock.MongoClient()
    app = create_app()
    with app.app_context():
        crawl_mgr = HouseCrawler.get_instance()
        mongo_mgr = MongoDbManager.get_instance()
        result = crawl_mgr.run()
        houses = mongo_mgr.get_client().five9one.houses.find()

    assert result == 'finished'
    assert houses.count() == 30 * 3 - crawl_mgr.duplicate_count


# A program to create several scatter plots of the random walk class

import matplotlib.pyplot as plt
from random_walk import Random_Walk

# Generating a new random walk for use in plotting
rw = Random_Walk()
rw.fill_walk()

plt.scatter(rw.x_location, rw.y_location, s=10)
plt.show()

# Creating a colormap that darkens as the steps progress

step_list = list(range(rw.steps))
plt.scatter(rw.x_location, rw.y_location, c=step_list, cmap=plt.cm.Purples,
            edgecolor='none', s=10)
plt.show()

# Making the start(green) and end(red) points more visible
plt.scatter(rw.x_location, rw.y_location, c=step_list, cmap=plt.cm.Purples,
            edgecolor='none', s=10)
plt.scatter(0, 0, color='green', s=30)
plt.scatter(rw.x_location[-1], rw.y_location[-1], color='red', s=30)
plt.show()

# Plotting the random walk as a line rather than a scatter plot
plt.plot(rw.x_location, rw.y_location, linewidth=1)
plt.show()


from flask import Flask, send_from_directory, Response
import cv2
import sys
import numpy
import time
import threading
import subprocess
import asyncio
import websockets
import random

frame_time = 1/12
app = Flask(__name__)
commonVideo = []
yesVideo = []
hereVideo = []

oldFrameNum = 0
currentVideo = commonVideo
frameNum = 1
ws_switch = ""
inc = True
is_off = False

def process_file(filename, arr):
    print("Processing " + filename)
    file = cv2.VideoCapture(filename)
    while True:
        retval, im = file.read()
        if (retval == False):
            break
        imgencode = cv2.imencode('.jpg', im)[1]
        bytesData = imgencode.tobytes()
        arr.append(bytesData)
    del(file)
    print("Processed " + filename)

def gen_frames():
    global currentVideo
    global is_off
    while True:
        time.sleep(frame_time)
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + get_next_frame() + b'\r\n')  # concat frame one by one and show result
        if is_off:
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + bytearray() + b'\r\n')  # concat frame one by one and show result
            is_off = False
            currentVideo = commonVideo
            return ""

def get_next_frame():
    global frameNum
    if is_off:
        return bytearray()
    if (currentVideo == commonVideo or currentVideo == hereVideo):
        out = currentVideo[frameNum]
        inc_frame_num()
        return out

    xrand = random.randint(0, 10)
    if xrand == 0:
        xrand = random.random()
        frame_reduction = int(xrand * (1/frame_time))
        if (frameNum <= len(currentVideo) - 1 - frame_reduction):   # dont go out of bounds
            time.sleep(xrand)
            frameNum = frameNum + frame_reduction
            return currentVideo[frameNum]

    out = currentVideo[frameNum]
    # # randomize bytes on image
    # for i in range(0, len(currentVideo)-1):
    #     delta = random.randint(-currFrame[i], 255 - currFrame[i])
    #     out.append(currFrame[i] + delta)
    inc_frame_num()
    return out

def inc_frame_num():
    global inc
    global frameNum
    global oldFrameNum
    global is_off
    global currentVideo
    global commonVideo

    if inc:
        frameNum = frameNum + 1

        if frameNum >= len(currentVideo) - 4:
            if currentVideo == hereVideo:
                frameNum = oldFrameNum
                currentVideo = commonVideo
                return
            if currentVideo != commonVideo:
                is_off = True
                print("sleepy boi")
                time.sleep(3)
                return
            inc = False
    else:
        frameNum = frameNum - 1
        if frameNum <= 2:
            inc = True

@ app.route('/yes')
def switch_to_yes():
    global currentVideo
    global frameNum
    global ws_switch
    currentVideo = yesVideo
    frameNum = 1
    ws_switch = "yes"
    return Response()

@ app.route('/here')
def switch_to_here():
    global currentVideo
    global frameNum
    global ws_switch
    global oldFrameNum
    oldFrameNum = frameNum
    currentVideo = hereVideo
    frameNum = 1
    ws_switch = "here"
    return Response()

@app.route("/off")
def off():
    global is_off
    is_off = True

@ app.route('/vid')
def vid():
    return Response(gen_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')

@ app.route('/aud')
def aud():
    def generate():
        with open('videos/test.mp4', 'rb') as file:
            file.seek(0)
            for chunk in iter(lambda: file.read(4096), b''):
                yield chunk
    return Response(generate(), mimetype='video/mp4')

@ app.route('/<path:path>')
def send_js(path):
    return send_from_directory('../videos', path)

async def echo(websocket, path):
    global ws_switch
    print("WS connected")
    while True:
        if (ws_switch):
            await websocket.send(ws_switch)
            ws_switch = ""
        time.sleep(0.1)

def start_ws():
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    asyncio.get_event_loop().run_until_complete(
        websockets.serve(echo, '0.0.0.0', 5001))
    asyncio.get_event_loop().run_forever()

def start():
    # process loop video
    t = threading.Thread(target=process_file, args=(
        "videos/common.mp4", commonVideo))
    t.start()

    # process respond message
    t = threading.Thread(target=process_file, args=(
        "videos/test.mp4", yesVideo))
    t.start()

    # process respond message
    t = threading.Thread(target=process_file, args=(
        "videos/here.mp4", hereVideo))
    t.start()

    t = threading.Thread(target=start_ws)
    t.start()

    app.run(host="0.0.0.0", port=5000, debug=False)


#!/usr/bin/env python3
"""
This code holds the solution for part 2 of day 3 of the Advent of Code for 2022.
"""

from collections import defaultdict
import sys

def process_group(group):
    badge_score = 0
    item_map = defaultdict(int)

    for rucksack in group:
        items = set([x for x in rucksack])

        for i in items:
            item_map[i] += 1

    for key, count in item_map.items():
        if count == len(group):
            if key > 'a':
                badge_score = ord(key) - ord('a') + 1
            else:
                badge_score = ord(key.lower()) - ord('a') + 27
            break

    return badge_score

def calculate_solution(items):
    result = 0

    group=[]
    for item in items:
        group.append(item)

        if len(group) == 3:
            result += process_group(group)
            group=[]

    return result

def run_test(test_input, groups, priorities, expected_solution):
    """
    Helper method for running some unit tests whilst minimising repetative code.
    """

    for items, priority in zip(groups, priorities):
        result=process_group(items)
        if result != priority:
            print(
                f'Test for {items} FAILED. Got a result of {result}, not {priority}')
            sys.exit(-1)

    result=calculate_solution(test_input)

    if result != expected_solution:
        print(
            f'Test for {test_input} FAILED. Got a result of {result}, not {expected_solution}')
        sys.exit(-1)

    print(f'Test for {test_input} passed.')

    return result

# Run any tests that we've defined to help validate our code prior to
# trying to solve the puzzle.

test_list = ["vJrwpWtwJgWrhcsFMMfFFhFp", "jqHRNqRjqzjGDLGLrsFMfFZSrLrFZsSL", "PmmdzqPrVvPwwTWBwg",
             "wMqvLMZHhHMvwLHjbvcjnnSBnvTQFn", "ttgJtRGJQctTZtZT", "CrZsJsPPZsGzwwsLwLmpwMDw"]
groups = [
    ["vJrwpWtwJgWrhcsFMMfFFhFp", "jqHRNqRjqzjGDLGLrsFMfFZSrLrFZsSL", "PmmdzqPrVvPwwTWBwg"],
    ["wMqvLMZHhHMvwLHjbvcjnnSBnvTQFn", "ttgJtRGJQctTZtZT", "CrZsJsPPZsGzwwsLwLmpwMDw"]
]

result = run_test(test_list, groups, [18, 52], 70)

print('')
print('-----------------')
print('All Tests PASSED.')
print('-----------------')
print('')

# Ok, so if we reach here, then we can be reasonably sure that the code
# above is working correctly. Let's use the actual captcha now.

with open('input.txt', 'r') as f:
    input_data = [line.strip() for line in f]
    answer = calculate_solution(input_data)

    print(f'Solution is {answer}')


import sys

import itertools

import numpy as np
import matplotlib.pyplot as plt

from sklearn.metrics import mean_squared_error
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import (
    LinearRegression,
    Ridge,
    Lasso,
    LassoCV,
    RidgeCV,
    ElasticNet,
)

# The `PlotManager` class provides methods to initialize a plot, plot regression models, and display
# the plot with a legend.
class PlotManager:
    def __init__(self):
        """
        The function initializes a plot with a title, x and y labels, and a grid, and plots a line with a
        specific range and style.
        """
        x = np.linspace(-10, 10, 1000)

        plt.title = "Regression models comparison"
        plt.xlabel("Predicted values")
        plt.ylabel("Real values")

        plt.plot(x, x, color="grey", linestyle="--", label="Y = x")
        plt.ylim(-5, 5)
        plt.xlim(-5, 5)
        plt.grid(True)

    def plot_regression(
        self, y_real=[], y_pred=[], model_name="No model name", color="blue"
    ):
        """
        The function `plot_regression` plots the real values (`y_real`) against the predicted values
        (`y_pred`) using a scatter plot with a specified color and model name as labels.

        :param y_real: The actual values of the dependent variable (y) in the regression model
        :param y_pred: The predicted values of the regression model
        :param model_name: The model_name parameter is a string that represents the name of the regression
        model being plotted. It is used as a label for the plot, defaults to No model name (optional)
        :param color: The color parameter is used to specify the color of the plotted points. It can be any
        valid color name or a hexadecimal color code, defaults to blue (optional)
        """
        plt.plot(y_real, y_pred, "o", color=color, label=model_name)

        for i in range(len(y_real)):
            plt.plot([y_real[i], y_real[i]], [y_real[i], y_pred[i]], color=color, linestyle='--')
            plt.text(y_real[i], y_pred[i], str(i), fontsize=10, color='red')

    def show(self):
        """
        The function displays a legend on a matplotlib plot and shows the plot.
        """
        plt.legend(loc="upper left")
        plt.show()

# The `Regression_Model_Tester` class is used to test a regression model by performing
# cross-validation and calculating the mean error of the model.
class Regression_Model_Tester:
    def __init__(self, X, Y, used_model, used_model_name, plot, color, parameters = None):
        """
        The above function is a constructor that initializes the attributes of an object, including the
        input data, the used model, and the model name, as well as empty lists for predicted values, real
        values, and prediction errors.

        :param X: The X parameter represents the input data for the model. It could be a matrix or an array
        containing the features or independent variables used for prediction
        :param Y: The Y parameter represents the target variable or the dependent variable in a machine
        learning model. It is the variable that we are trying to predict or estimate based on the input
        variables X
        :param used_model: The `used_model` parameter is the machine learning model that will be used for
        prediction. It could be any model such as linear regression, decision tree, random forest, etc. The
        specific model will be passed as an argument when creating an instance of this class
        :param used_model_name: The name of the machine learning model that is being used
        """
        self.X = X
        self.Y = Y

        self.used_model = used_model
        self.used_model_name = used_model_name

        self.y_pred = []
        self.y_real = []
        self.prediction_error = []
        self.plot = plot
        self.color = color

        self.parameters = parameters

        self.column_removed = []

    @property
    def best_error_removing_column(self):
        return self._previous_best_error, self._previous_best_combination

    @property
    def errors(self) -> np.array:
        """
        The function calculates the error of the model by taking the mean of the prediction error.
        :return: the mean of the `prediction_error` array.
        """
        if len(self.prediction_error) == 0:
            return sys.maxsize
        return np.mean(self.prediction_error)

    @errors.setter
    def errors_setter(self, errors):
        self._previous_best_error = errors

    def _train_model(self, X_set, Y_set, i ):
        """
        The `_train_model` function trains a model on a training set, makes predictions on a test set, and
        calculates the mean squared error between the predicted and actual values.

        :param X_set: X_set is a numpy array containing the input features for the training data. Each row
        of X_set represents a single training example, and each column represents a different feature
        :param Y_set: Y_set is the set of target values or labels for the training data. It represents the
        true values that the model is trying to predict
        :param i: The parameter "i" represents the index of the data point that is being used as the test
        set. It is used to split the data into training and testing sets by excluding the data point at
        index "i" from the training set and using it as the test set
        :return: the mean squared error between the predicted values (y_pred) and the actual values
        (y_test).
        """

        x_train_set_cpy = np.delete(np.copy(X_set), i, axis=0)
        y_train_set_cpy = np.delete(np.copy(Y_set), i, axis=0)

        x_test = X_set[i : i + 1]
        y_test = Y_set[i : i + 1]

        self.used_model.fit(x_train_set_cpy, y_train_set_cpy)

        # Test prediction with the part of the training set
        y_pred = self.used_model.predict(x_test)

        y_pred, y_test

        return(mean_squared_error(y_test, y_pred), y_pred, y_test)

    def _cross_validation_with_all_columns(self) -> None:
        """
        The function performs cross-validation by training a model on all columns of the input data and
        calculating the prediction error for each iteration.
        :return: None
        """
        for i in range(len(self.X)):
            SE, y_pred, y_real = self._train_model(np.copy(self.X), np.copy(self.Y), i)

            self.y_pred.append(y_pred)
            self.y_real.append(y_real)

            self.prediction_error.append(SE)

        return None

    def _cross_validation_removing_colums(self, number_of_column_to_remove = 1) ->  None:
        """
        The function performs cross-validation by removing a specified number of columns from a dataset and
        training a model on the modified dataset.

        :param number_of_column_to_remove: The parameter `number_of_column_to_remove` is an integer that
        specifies the maximum number of columns to remove from the dataset. It determines the length of the
        combinations of columns that will be generated. For example, if `number_of_column_to_remove` is set
        to 2, the code will generate, defaults to 1 (optional)
        :return: None.
        """

        column_to_remove = []
        see_vector = []
        y_pred_vector= []
        y_real_vector  =[]

        for length in range (1, number_of_column_to_remove + 1):
            for combo in itertools.combinations(range(10), length):
                column_to_remove.append(combo)

        #print(f"number_of colums to remove {(number_of_column_to_remove)} column_to_remove {column_to_remove}")

        for i in column_to_remove:
            x_set = np.delete(np.copy(self.X), i, axis=1)

            for  j in range (len(x_set)):
                model_results = self._train_model(x_set, np.copy(self.Y), j)
                see_vector.append(model_results[0])
                y_pred_vector.append(model_results[1])
                y_real_vector.append(model_results[2])

            if np.mean(see_vector) < self.errors or np.isnan(self.errors):
                self.column_removed = i
                self.prediction_error = see_vector
                self.y_pred = y_pred_vector
                self.y_real = y_real_vector

            see_vector = []
            y_pred_vector = []
            y_real_vector = []

        return None

    def run_validation(self, number_of_column_to_remove = 0) -> np.ndarray:
        """
        The `run_validation` function performs cross-validation and optionally removes columns and plots the
        model, and then prints the mean error and the best combination of removed columns.

        :param number_of_column_to_remove: The parameter "number_of_column_to_remove" is an optional
        parameter that specifies the number of columns to remove during cross-validation. If the value is 0,
        the method will perform cross-validation without removing any columns. If the value is greater than
        0, the method will perform cross-validation by removing, defaults to 0 (optional)
        :param plot_model: A boolean flag indicating whether or not to plot the model. If set to True, the
        model will be plotted; if set to False, the model will not be plotted, defaults to False (optional)
        :return: None
        """
        if  number_of_column_to_remove == 0 :
            self._cross_validation_with_all_columns()
        else:
            self._cross_validation_removing_colums(number_of_column_to_remove)

        return self.errors

    def plot_model(self) -> None:
        """
        The function `plot_model` plots the predicted values against the real values.
        :return: None.
        """
        self.plot.plot_regression(
            np.array(self.y_real).reshape(-1),
            np.array(self.y_pred).reshape(-1),
            f"{self.used_model_name} : {self.errors}",
            color=self.color,
        )
        return None

    def model_logger(self):
        print(f"Current Model: {self.used_model_name} current special parameters : {self.parameters}  current_error : {self.errors} current columns removed : {self.column_removed}")

def elastic_net(X_train, Y_train, plot, color, alphas=[0.09775999999999777], l1_ratios= [0.89], number_of_columns_to_remove = 0, plot_model=False) -> None:
    """
    The function ElasticNet performs Elastic Net regression on the given training data and returns the
    errors of the model.

    :param X_train: The training data for the independent variables (features)
    :param Y_train: The target variable for training the model
    :param plot: A boolean value indicating whether or not to plot the validation results
    :param color: The "color" parameter is used to specify the color of the plot when visualizing the
    results. It can be any valid color value, such as "red", "blue", "#FF0000" (hexadecimal color code),
    or "rgb(255, 0, 0)" (
    :param alpha: The alpha parameter controls the regularization strength of the Elastic Net model. It
    determines the balance between the L1 and L2 penalties. A higher alpha value leads to stronger
    regularization and can help prevent overfitting
    :param l1_ratio: The l1_ratio parameter in ElasticNet is a value between 0 and 1 that determines the
    balance between L1 and L2 regularization
    :return: the errors from the Elastic Net model.
    """

    best_elastic_model = None

    for alpha in alphas:
        for l1_ratio in l1_ratios:
            parameters = {"alpha": alpha, "l1_ratio": l1_ratio}

            elastic_net = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, max_iter=10000)

            elastic_net_model = Regression_Model_Tester(
                X_train,
                Y_train,
                elastic_net,
                "Elastic Net",
                plot,
                color,
                parameters
            )

            if best_elastic_model is None:
                best_elastic_model = elastic_net_model

            current_sse = elastic_net_model.run_validation(number_of_column_to_remove=number_of_columns_to_remove)

            if (current_sse < best_elastic_model.errors):
                best_elastic_model = elastic_net_model

    best_elastic_model.model_logger()

    if plot_model is True:
        best_elastic_model.plot_model()

    return None

def polynomial_model(X_train, Y_train, plot, color, degree=2, number_of_columns_to_remove = 0, plot_model=False) -> None:
    """
    The function `polynomial_model` fits a polynomial regression model of a specified degree to the
    given training data and evaluates its performance using a regression model tester.

    :param X_train: The training data for the independent variable(s) (features)
    :param Y_train: The parameter Y_train represents the target variable or the dependent variable in
    your dataset. It is the variable that you are trying to predict or model using the independent
    variables (X_train)
    :param plot: The "plot" parameter is a boolean value that determines whether or not to plot the
    regression line and data points. If set to True, the regression line and data points will be
    plotted. If set to False, no plot will be generated
    :param color: The "color" parameter is used to specify the color of the plot in the polynomial_model
    function. It can be any valid color value, such as "red", "blue", "green", etc
    :param degree: The degree parameter determines the degree of the polynomial features to be used in
    the polynomial regression model. It specifies the maximum power of the independent variable(s) in
    the polynomial equation. For example, if degree=2, the polynomial regression model will include
    features with powers 0, 1, and, defaults to 2 (optional)
    :return: None.
    """

    polynomial_features = PolynomialFeatures(degree=degree, include_bias=True)

    x_poly = polynomial_features.fit_transform(X_train)

    polynomial_regression_model = Regression_Model_Tester(
        x_poly,
        Y_train,
        LinearRegression(),
        "Polynomial Regression",
        plot,
        color
    )
    polynomial_regression_model.run_validation(number_of_column_to_remove=number_of_columns_to_remove)

    polynomial_regression_model.model_logger()

    if plot_model is True:
        polynomial_regression_model.plot_model()

    return None

def ridge_model(X_train, Y_train, plot, color, alphas=[2.08710000000000001], number_of_columns_to_remove = 0, plot_model=False) -> None:
    """
    The function `ridge_model` performs ridge regression on the given training data and prints the
    validation results.

    :param X_train: The X_train parameter is the training data for the independent variables. It should
    be a matrix or dataframe with shape (n_samples, n_features), where n_samples is the number of
    samples or observations and n_features is the number of independent variables or features
    :param Y_train: The parameter Y_train represents the target variable or the dependent variable in
    your training dataset. It is the variable that you are trying to predict or model using the
    independent variables (X_train)
    :return: None.
    """

    best_ridge_model = None

    for alpha in alphas:

        parameters = {"alpha": alpha}

        ridge_regression_model = Regression_Model_Tester(
            X_train,
            Y_train,
            Ridge(alpha=alpha, max_iter=10000),
            "Ridge Regression",
            plot,
            color,
            parameters
        )

        if best_ridge_model is None:
            best_ridge_model = ridge_regression_model

        current_ridge_model_error = ridge_regression_model.run_validation(number_of_column_to_remove=number_of_columns_to_remove)

        if (current_ridge_model_error < best_ridge_model.errors):
            best_ridge_model = ridge_regression_model

    if (plot_model is True):
        best_ridge_model.plot_model()

    best_ridge_model.model_logger()

    return None

def lasso_model(X_train, Y_train, plot, color, alphas=[0.08710000000000001], number_of_columns_to_remove = 0, plot_model=False) -> None:
    """
    The function `lasso_model` trains and tests a Lasso regression model using the provided training
    data.

    :param X_train: The parameter X_train is the training data for the independent variables. It should
    be a matrix or dataframe with shape (n_samples, n_features), where n_samples is the number of
    samples or observations and n_features is the number of independent variables or features
    :param Y_train: The parameter Y_train represents the target variable or the dependent variable in
    your dataset. It is the variable that you are trying to predict or model using the independent
    variables (X_train)
    :return: None.
    """

    best_lasso_model = None

    for alpha in alphas:

        parameters = {"alpha": alpha}

        lasso_regression_model = Regression_Model_Tester(
            X_train,
            Y_train,
            Lasso(alpha=alpha, max_iter=10000),
            "Lasso Regression",
            plot,
            color,
            parameters
        )

        if best_lasso_model is None:
            best_lasso_model = lasso_regression_model

        current_lasso_model_error = lasso_regression_model.run_validation(number_of_column_to_remove=number_of_columns_to_remove)

        if (current_lasso_model_error < best_lasso_model.errors):
            best_lasso_model = lasso_regression_model

    if (plot_model is True):
        best_lasso_model.plot_model()

    best_lasso_model.model_logger()

    return None

def linear_model(X_train, Y_train, plot, color, number_of_columns_to_remove = 0, plot_model=False) -> None:
    """
    The function `linear_model` trains and tests a linear regression model using the given training
    data.

    :param X_train: The X_train parameter is the training data for the independent variables in your
    linear regression model. It should be a 2-dimensional array or dataframe where each row represents a
    sample and each column represents a feature
    :param Y_train: The parameter Y_train represents the target variable or the dependent variable in
    your dataset. It is the variable that you are trying to predict or model using the independent
    variables (X_train)
    :return: None.
    """

    linear_regression_model = Regression_Model_Tester(
        X_train,
        Y_train,
        LinearRegression(),
        "Linear Regression",
        plot,
        color
    )
    linear_regression_model.run_validation(number_of_column_to_remove=number_of_columns_to_remove)

    linear_regression_model.model_logger()

    if (plot_model is True):
        linear_regression_model.plot_model()

    return None

def main():
    """
    The main function loads training data and calls three different models: linear_model, ridge_model,
    and lasso_model.
    """

    plot = PlotManager()

    ridge_alpha = np.arange(2.0, 3.1, 0.1).tolist()
    l_ratio = [0.89]
    elastic_net_alpha = [0.09775999999999777]
    #ridge_alpha = [2.08710000000000001]
    lasso_alpha = [0.08710000000000001]

    f = open("errors.txt", "w")
    f.write("Errors for regression models:\n")

    x_train_set = np.load("input_files/1_exercise/X_train_regression1.npy")
    y_train_set = np.load("input_files/1_exercise/y_train_regression1.npy")

    #linear_model(X_train=x_train_set,
    #             Y_train=y_train_set,
    #             plot= plot,
    #             color="blue",
    #             number_of_columns_to_remove=0,
#                 plot_model=True)

    #polynomial_model(X_train=x_train_set,
    #                 Y_train=y_train_set,
    #                 plot=plot,
    #                 color="cyan",
    #                 degree=2)

    ridge_model(X_train=x_train_set,
                Y_train=y_train_set,
                plot=plot,
                color="green",
                alphas=ridge_alpha,
                number_of_columns_to_remove = 9, plot_model=True)

    #lasso_model(X_train=x_train_set,
    #            Y_train=y_train_set,
    #            plot=plot,
    #            color="red",
    #            alphas=lasso_alpha)

    #elastic_net(X_train=x_train_set,
    #            Y_train=y_train_set,
    #            plot=plot,
    #            color="purple",
    #            alphas=elastic_net_alpha,
    #            l1_ratios=l_ratio)

    plot.show()
    f.close()

if __name__ == "__main__":
    main()