import sys
import os.path
from enum import Enum
from typing import TYPE_CHECKING, Optional, Union

import pygame
from pygame import gfxdraw
from pygame.locals import *

from code.omok3.animations import SpaceAnimation, StoneDeployedAnimation, StoneIdleAnimation, StoneReservedAnimation
from code.network.udp.client_udp import OmokUDP

class Color(Enum):
    BLUE = (0, 0, 255)
    RED = (255, 0, 0)
    GREEN = (0, 255, 0)
    BLACK = (0, 0, 0)
    WHITE = (255, 255, 255)

class Team(Enum):
    NONE = 0
    BLACK = 1
    WHITE = 2

    def color(self) -> tuple[int, int, int]:
        if self == Team.NONE:
            return (127, 127, 127)
        elif self == Team.BLACK:
            return (0, 0, 0)
        elif self == Team.WHITE:
            return (255, 255, 255)
        else:
            raise ValueError()

class State(Enum):
    EMPTY = 0
    BLACK = 1
    WHITE = 2
    BLACK_RESERVED = 3
    WHITE_RESERVED = 4

class Validity(Enum):
    INVALID = 0
    RESERVE = 1
    DENY = 2
    CONFIRM = 3

class Game:
    FPS = 60

    def __init__(self, multiplay=False, udp: OmokUDP = None) -> None:
        self.multiplay = multiplay
        if self.multiplay:
            self.udp = udp
            self.udp.rx_event = lambda data: self.decode_udp_and_update(data)
            self.udp.start_listen()
            self.opponent_move: tuple[bool, tuple[int, int]] = (False, (0, 0))
            self.opponent_chat: str = ""
            self.encode_udp_and_send("HELLO")

        pygame.init()
        self.clock = pygame.time.Clock()

        pygame.mouse.set_visible(False)
        self.mouse_clicking = False

        pygame.display.set_caption("Real-time LiveOmok3")  # 창 제목 설정
        self.displaysurf = pygame.display.set_mode((875, 1075))
        root = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))  # LiveOmok/
        img_path = os.path.join(root, "res", "board-PIL-step.png")  # LiveOmok/res/board-PIL-step.png
        self.background = pygame.image.load(img_path)
        self.dotum = pygame.font.SysFont("malgun gothic", 30)
        # self.dotum = pygame.font.SysFont("dotumche", 30)
        self.black_heuk = self.dotum.render("흑", True, (0, 0, 0))
        self.white_baek = self.dotum.render("백", True, (255, 255, 255))

        self.board = Board(self)
        self.black_gauge = 0.0
        self.white_gauge = 0.0

        self.displaysurf.blit(self.background, (0, 0))

    def event_lclick(self, click_pos: tuple[int, int]):
        target = self.find_mouse_pointed_space(click_pos)
        if (target is not None) and (target.valid(Team.BLACK) != Validity.INVALID):
            target.click(Team.BLACK)
            if self.multiplay:
                self.encode_udp_and_send(click_pos)

    def event_rclick(self, click_pos: tuple[int, int]):
        if self.multiplay:
            pass
        else:
            target = self.find_mouse_pointed_space(click_pos)
            if (target is not None) and (target.valid(Team.WHITE) != Validity.INVALID):
                target.click(Team.WHITE)

    def loop(self) -> str:
        mouse_over_shadow = pygame.Surface((51, 51), pygame.SRCALPHA)
        mouse_over_shadow.fill((255, 255, 255, 0))  # RGBA : totally transparent
        mouse_over_shadow.set_alpha(128)
        cursor_with_gauge = pygame.Surface((41, 41), pygame.SRCALPHA)

        while True:
            # 고정된 그림: 백 점수판, 배경, 흑 점수판
            pygame.draw.rect(
                self.displaysurf,
                Color.BLACK.value,
                [(0, 0), (875, 100)]
            )
            self.displaysurf.blit(self.background, (0, 100))
            pygame.draw.rect(
                self.displaysurf,
                Color.WHITE.value,
                [(0, 975), (875, 100)]
            )

            # 마우스 이벤트
            for event in pygame.event.get():
                left_clicked = pygame.mouse.get_pressed()[0]
                if event.type == pygame.MOUSEBUTTONDOWN and left_clicked:
                    if self.mouse_clicking:  # mouse holding
                        continue
                    self.mouse_clicking = True
                    pos = pygame.mouse.get_pos()
                    self.event_lclick(pos)
                elif event.type == pygame.MOUSEBUTTONDOWN and pygame.mouse.get_pressed()[2]:
                    pos = pygame.mouse.get_pos()
                    self.event_rclick(pos)
                elif event.type == pygame.MOUSEBUTTONUP and not left_clicked:
                    self.mouse_clicking = False
                if event.type == QUIT:
                    pygame.quit()
                    sys.exit()

            # UDP 이벤트(넷 플레이에서만)
            if self.multiplay:
                if self.opponent_move[0] is True:
                    target = self.find_mouse_pointed_space(self.opponent_move[1])
                    if target is not None:
                        target.click(Team.WHITE)
                    self.opponent_move = (False, (0, 0))

            # 마우스 hover 시 그림자 표시
            pos = pygame.mouse.get_pos()
            target = self.find_mouse_pointed_space(pos)
            if target is not None:
                if target.state == State.EMPTY:
                    pygame.draw.circle(
                        mouse_over_shadow,
                        (127, 127, 127),
                        mouse_over_shadow.get_rect().center,
                        17
                    )
                    self.displaysurf.blit(mouse_over_shadow, target.rect.topleft)

            # 점수판 표시
            black_5, white_5 = self.board.count_5_connected()
            black_score = self.dotum.render(f"{black_5}/3", True, (0, 0, 0))
            white_score = self.dotum.render(f"{white_5}/3", True, (255, 255, 255))
            self.displaysurf.blit(self.black_heuk, (75, 1010))
            self.displaysurf.blit(self.white_baek, (700, 35))
            self.displaysurf.blit(black_score, (125, 1010))
            self.displaysurf.blit(white_score, (750, 35))

            # 게이지 표시
            if self.black_gauge < 3.0:
                self.black_gauge += 1/180  # 3초에 하나씩 채워진다
            if self.white_gauge < 3.0:
                self.white_gauge += 1/180
            self.draw_gauge()

            if black_5 >= 3:
                pygame.mouse.set_visible(True)
                return "Black wins!"
            elif white_5 >= 3:
                return "White wins!"

            self.board.update()

            # 마우스 커서 표시
            cursor_with_gauge.fill((255, 255, 255, 0))  # RGBA : totally transparent
            gfxdraw.aatrigon(cursor_with_gauge, 20, 22, 14, 4, 26, 4, (0, 0, 255))
            gfxdraw.aatrigon(cursor_with_gauge, 20, 22, 38, 27, 33, 36, (0, 0, 255))
            gfxdraw.aatrigon(cursor_with_gauge, 20, 22, 7, 36, 2, 27, (0, 0, 255))
            if 0.0 <= self.black_gauge < 1.0:
                gfxdraw.filled_trigon(
                    cursor_with_gauge,
                    20, 22,
                    round(20 - 6 * self.black_gauge), round(22 - 18 * self.black_gauge),
                    round(20 + 6 * self.black_gauge), round(22 - 18 * self.black_gauge),
                    (0, 0, 255)
                )
            elif 1.0 <= self.black_gauge < 2.0:
                gfxdraw.filled_trigon(cursor_with_gauge, 20, 22, 14, 4, 26, 4, (0, 0, 255))
                gfxdraw.filled_trigon(
                    cursor_with_gauge,
                    20, 22,
                    round(20 + 18 * (self.black_gauge-1)), round(22 + 5 * (self.black_gauge-1)),
                    round(20 + 13 * (self.black_gauge-1)), round(22 + 14 * (self.black_gauge-1)),
                    (0, 0, 255)
                )
            elif 2.0 <= self.black_gauge < 3.0:
                gfxdraw.filled_trigon(cursor_with_gauge, 20, 22, 14, 4, 26, 4, (0, 0, 255))
                gfxdraw.filled_trigon(cursor_with_gauge, 20, 22, 38, 27, 33, 36, (0, 0, 255))
                gfxdraw.filled_trigon(
                    cursor_with_gauge,
                    20, 22,
                    round(20 - 13 * (self.black_gauge-2)), round(22 + 14 * (self.black_gauge-2)),
                    round(20 - 18 * (self.black_gauge-2)), round(22 + 5 * (self.black_gauge-2)),
                    (0, 0, 255)
                )
            else:  # self.black_gauge >= 3.0
                gfxdraw.filled_trigon(cursor_with_gauge, 20, 22, 14, 4, 26, 4, (0, 0, 255))
                gfxdraw.filled_trigon(cursor_with_gauge, 20, 22, 38, 27, 33, 36, (0, 0, 255))
                gfxdraw.filled_trigon(cursor_with_gauge, 20, 22, 7, 36, 2, 27, (0, 0, 255))
            self.displaysurf.blit(cursor_with_gauge, (pos[0]-21, pos[1]-21))

            pygame.display.update()
            self.clock.tick(Game.FPS)

    def draw_gauge(self):
        # 기본 돌 3개 (새치기 시 돌려받는 걸 대비해 3개 더 그린다)
        for i in range(6):
            pygame.draw.circle(
                self.displaysurf,
                (0, 0, 0),
                (250 + 40*i, 1025),
                20
            )
            pygame.draw.circle(
                self.displaysurf,
                (255, 255, 255),
                (400 + 40*i, 50),
                20
            )
        # 가리개를 통해 원의 일부를 가려서 활꼴을 그린다
        pygame.draw.rect(
            self.displaysurf,
            (255, 255, 255),
            (230 + self.black_gauge * 40, 1000, (3.0 - self.black_gauge) * 40 + 130, 50)
        )
        pygame.draw.rect(
            self.displaysurf,
            (0, 0, 0),
            (380 + self.white_gauge * 40, 25, (3.0 - self.white_gauge) * 40 + 130, 50)
        )

    def decode_udp_and_update(self, msg: str) -> None:
        print("Opponent - " + msg)
        if msg.startswith("Move:"):
            # msg에는 Move 헤더와 클릭한 좌표가 들어온다. "Move:530,230"
            msg = msg[5:]
            x, y, *other = msg.split(',')
            if other:  # 좌표가 3개 이상
                return
            self.opponent_move = (True, (int(x), int(y)))

        elif msg.startswith("Chat:"):
            # msg에는 Chat 헤더와 채팅 내용이 들어온다. "Chat:Ez LOL"
            self.opponent_chat = msg[5:]
            print(self.opponent_chat)

    def encode_udp_and_send(self, pos_or_msg: Union[tuple[int, int], str]) -> None:
        if isinstance(pos_or_msg, tuple):
            msg = f"Move:{pos_or_msg[0]},{pos_or_msg[1]}"
        elif isinstance(pos_or_msg, str):
            msg = f"Chat:{pos_or_msg}"
        print("Me - " + msg)
        self.udp.send(msg)

    def find_mouse_pointed_space(self, pos: tuple[int, int]) -> Optional['Space']:
        target: Optional['Space'] = None
        for row in self.board.spaces:
            for space in row:
                if space.rect.inflate(-10, -10).collidepoint(pos):  # Taxi distance
                # if pygame.Vector2(space.rect.center).distance_to(pos) < 25:  # Euclidean distance
                    target = space
                    break
            if target is not None:
                break
        return target

class Board:

    def __init__(self, game: Game) -> None:
        self.game = game
        self.spaces = [[Space(self, j, i) for j in range(15)] for i in range(15)]

    def update(self) -> None:
        for row in self.spaces:
            for space in row:
                try:
                    next(space.animation.play())
                except StopIteration as e:
                    if isinstance(space.animation, StoneReservedAnimation):
                        space.animation = StoneDeployedAnimation(space)
                    elif isinstance(space.animation, StoneDeployedAnimation):
                        space.animation = StoneIdleAnimation(space)

    def count_5_connected(self) -> tuple[int, int]:
        # 참고로 6목은 5목 2개로 본다. 7목은 5목 3개로 보기 때문에 바로 이긴다.
        black_5 = 0
        white_5 = 0
        for row in self.spaces:
            for space in row:
                for team in [space.is_horizontal_5(), space.is_vertical_5(), space.is_slash_5(), space.is_backslash_5()]:
                    if team == Team.BLACK: black_5 += 1
                    elif team == Team.WHITE: white_5 += 1
        return black_5, white_5

class Space:

    def __init__(self, board: Board, x, y) -> None:
        self.board = board
        self.x = x  # 0-14
        self.y = y  # 0-14
        self.team = Team.NONE
        # self.team = Team.BLACK
        self.animation = SpaceAnimation(self)
        # self.animation = StoneIdleAnimation(self)

    def __repr__(self) -> str:
        row = str(15 - self.y)
        col = "ABCDEFGHIJKLMNOP"[self.x]
        return f"Space('{col}{row}')"

    @property
    def rect(self) -> Rect:
        return pygame.Rect(56 + 51 * self.x, 156 + 51 * self.y, 51, 51)

    @property
    def state(self) -> State:
        if isinstance(self.animation, SpaceAnimation):
            return State.EMPTY
        elif isinstance(self.animation, (StoneIdleAnimation, StoneDeployedAnimation)):
            if self.team == Team.BLACK:
                return State.BLACK
            elif self.team == Team.WHITE:
                return State.WHITE
            else:
                raise ValueError()
        elif isinstance(self.animation, StoneReservedAnimation):
            if self.team == Team.BLACK:
                return State.BLACK_RESERVED
            elif self.team == Team.WHITE:
                return State.WHITE_RESERVED
            else:
                raise ValueError()

    def valid(self, clicker_team: Team) -> Validity:
        if self.state == State.EMPTY:
            if (clicker_team == Team.BLACK and self.board.game.black_gauge >= 1.0) \
                or (clicker_team == Team.WHITE and self.board.game.white_gauge >= 1.0):
                return Validity.RESERVE
        elif self.state == State.BLACK_RESERVED:
            if (clicker_team == Team.BLACK and self.board.game.black_gauge >= 1.0):
                return Validity.CONFIRM
            if (clicker_team == Team.WHITE and self.board.game.white_gauge >= 1.5):
                return Validity.DENY
        elif self.state == State.WHITE_RESERVED:
            if (clicker_team == Team.BLACK and self.board.game.black_gauge >= 1.5):
                return Validity.DENY
            if (clicker_team == Team.WHITE and self.board.game.white_gauge >= 1.0):
                return Validity.CONFIRM
        else: return Validity.INVALID

    def click(self, clicker_team: Team):
        if self.state == State.EMPTY:
            if clicker_team == Team.BLACK:
                if self.board.game.black_gauge >= 1.0:
                    self.team = Team.BLACK
                    self.animation = StoneReservedAnimation(self)
                    self.board.game.black_gauge -= 1.0
            elif clicker_team == Team.WHITE:
                if self.board.game.white_gauge >= 1.0:
                    self.team = Team.WHITE
                    self.animation = StoneReservedAnimation(self)
                    self.board.game.white_gauge -= 1.0

        elif self.state == State.BLACK_RESERVED:
            if clicker_team == Team.BLACK:  # 흑 확정
                if self.board.game.black_gauge >= 1.0:
                    self.animation = StoneDeployedAnimation(self)
                    self.board.game.black_gauge -= 1.0
            elif clicker_team == Team.WHITE:  # 백 새치기
                if self.board.game.white_gauge >= 1.5:
                    if self.animation.current_frame > 15:
                        self.team = clicker_team
                        self.animation = StoneDeployedAnimation(self)
                        self.board.game.white_gauge -= 1.5
                        # self.board.game.black_gauge = min(3.0, self.board.game.black_gauge + 1.0)
                        self.board.game.black_gauge += 1.0

        elif self.state == State.WHITE_RESERVED:
            if clicker_team == Team.WHITE:  # 백 확정
                if self.board.game.white_gauge >= 1.0:
                    self.animation = StoneDeployedAnimation(self)
                    self.board.game.white_gauge -= 1.0
            elif clicker_team == Team.BLACK:  # 흑 새치기
                if self.board.game.black_gauge >= 1.5:
                    if self.animation.current_frame > 15:
                        self.team = clicker_team
                        self.animation = StoneDeployedAnimation(self)
                        self.board.game.black_gauge -= 1.5
                        # self.board.game.white_gauge = min(3.0, self.board.game.white_gauge + 1.0)
                        self.board.game.white_gauge += 1.0

    def is_horizontal_5(self) -> Team:
        if not 2 <= self.x <= 12:
            return Team.NONE
        if self.state not in [State.BLACK, State.WHITE]:
            return Team.NONE
        if self.board.spaces[self.y][self.x - 2].state \
            == self.board.spaces[self.y][self.x - 1].state \
            == self.state \
            == self.board.spaces[self.y][self.x + 1].state \
            == self.board.spaces[self.y][self.x + 2].state:
            return self.team
        else:
            return Team.NONE

    def is_vertical_5(self) -> Team:
        if not 2 <= self.y <= 12:
            return Team.NONE
        if self.state not in [State.BLACK, State.WHITE]:
            return Team.NONE
        if self.board.spaces[self.y - 2][self.x].state \
            == self.board.spaces[self.y - 1][self.x].state \
            == self.state \
            == self.board.spaces[self.y + 1][self.x].state \
            == self.board.spaces[self.y + 2][self.x].state:
            return self.team
        else:
            return Team.NONE

    def is_slash_5(self) -> Team:
        if (not 2 <= self.x <= 12) or (not 2 <= self.y <= 12):
            return Team.NONE
        if self.state not in [State.BLACK, State.WHITE]:
            return Team.NONE
        if self.board.spaces[self.y - 2][self.x + 2].state \
            == self.board.spaces[self.y - 1][self.x + 1].state \
            == self.state \
            == self.board.spaces[self.y + 1][self.x - 1].state \
            == self.board.spaces[self.y + 2][self.x - 2].state:
            return self.team
        else:
            return Team.NONE

    def is_backslash_5(self) -> Team:
        if (not 2 <= self.x <= 12) or (not 2 <= self.y <= 12):
            return Team.NONE
        if self.state not in [State.BLACK, State.WHITE]:
            return Team.NONE
        if self.board.spaces[self.y - 2][self.x - 2].state \
            == self.board.spaces[self.y - 1][self.x - 1].state \
            == self.state \
            == self.board.spaces[self.y + 1][self.x + 1].state \
            == self.board.spaces[self.y + 2][self.x + 2].state:
            return self.team
        else:
            return Team.NONE

if __name__ == '__main__':
    Game().loop()


from django.shortcuts import redirect, render, HttpResponse
from pedidos.serializers import PedidosSerializer
from produtos.models import Produtos
from django.contrib import messages
from .forms import FormPedido, FormStatusPedido
from .models import Pedidos
from django.db.models import Q
from django.contrib.auth.decorators import login_required

@login_required(login_url="/usuario/login/")
def index(request):
    pedidos = Pedidos.objects.filter(status_de_pagamento='pendente', status_do_pedido='pronto')
    user = request.user

    return render(request, 'html/pedidos.html', {'pedidos': pedidos, 'user': user})

@login_required(login_url="/usuario/login/")
def pedido_detalhe(request, pk):
    pedido_detalhe = Pedidos.objects.get(pk=pk)
    list_pedidos = pedido_detalhe.pedidos.all()
    # print(pedido_detalhe.status_de_pagamento[0:].capitalize())

    if request.method == "POST":
        form = FormPedido(request.POST, instance=pedido_detalhe)
        if form.is_valid():
            form.save()
            return redirect(F'/pedidos-detalhe/{pk}')
    else:
        form = FormPedido()
    return render(request, 'html/pedido-detalhe.html', { 'pedido_detalhe': pedido_detalhe, 'form': form, 'list_pedidos': list_pedidos })

def historico_de_pedidos(request):
    all_pedidos = Pedidos.objects.filter(status_de_pagamento='pago')
    return render(request, 'html/historico-de-pedidos.html', {'all_pedidos': all_pedidos})

def pedidos_da_cozinha(request):
    return render(request, 'html/pedidos-da-cozinha.html')

def pedido_cozinha_detalhe(request, pk):
    pedido = Pedidos.objects.get(pk=pk)
    list_pedidos = pedido.pedidos.all()
    if request.method == "POST":
        form = FormStatusPedido(request.POST, instance=pedido)

        if form.is_valid():
            form.save()
            if form.data['status_do_pedido'] == 'pronto':
                return redirect('pedidos-cozinha')
            else:
                messages.success(request, 'PEDIDO ALTERADO COM SUCESSO!')
            # return redirect(f'pedidos-detalhe-cozinha/{pk}')
        else:
            # Adicionar mensagem de erro
            messages.error(request, 'Erro ao enviar informações')
    else:
        form = FormStatusPedido(instance=pedido)
    return render(request, 'html/detalhe-pedido-cozinha.html', {'pedido': pedido, 'form': form, 'list_pedidos':list_pedidos})


import pygame, sys, math, random
pygame.init()
screen = pygame.display.set_mode([1680,1000])
dirt = pygame.image.load("dirt.jpg")
mountian = pygame.image.load("mountian.jpg")
water = pygame.image.load("water.jpg")

terrains = [{"sprite": dirt}, {"sprite": mountian}, {"sprite": water}]

themap = [[-1 for x in range(150)] for y in range (150)]

for x in range (0,150):
    for y in range (0,150):
        themap[x][y] = terrains[random.randint(0,2)]

while True:
    pygame.time.delay(50)

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            sys.exit()
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                for x in range (0,150):
                    for y in range (0,150):
                        themap[x][y] = terrains[random.randint(0,2)]

    for x in range (0, 150):
        for y in range (0, 150):
            screen.blit(themap[x][y]["sprite"], [x * 16, y * 16])

    pygame.display.flip()


import urllib.request
from bs4 import BeautifulSoup

base_url = 'http://www.domofond.ru/prodazha-kvartiry-magnitogorsk-c2293'

def get_html(url):
    response = urllib.request.urlopen(url)
    return response.read()

def pagination(html):
    soup = BeautifulSoup(html)
    number = soup.find_all('ul', {'class': 'pagination'})[-1].text
    number = re.sub("\D", "", number)
    return number

def parse(html):

    soup = BeautifulSoup(html)
    table = soup.find('div', {'id': 'listingResults'})
    rows = table.find_all('div', {'class': 'pull-left df_listingTileContent '})

    flats = []

    for row in rows:
        #Находим цену квартиры в теге <p> и преобразовываем в текст
        price = row.find('p').text
        #Преобразовываем весь найденный текст в число /отбрасывая все буквы, оставляя лишь числа/
        price = re.sub("\D", "", price)

        #Находим цену квадрата квартиры в теге <span>
        square_price = row.find('span').text
        square_price = re.sub("\D", "", square_price)

        square = row.find('li').text
        square = re.sub("\D", "", square)

        floor = (row.find('li')).find_next('li').text
        floor = re.findall('(\d+)', floor)

        address = row.find('span', {'itemprop': 'address'}).text

        #добавляем квартиры в словарь
        flats.append({
                'price': price,
                'square_price': square_price,
                'square': square,
                'floor': floor[0],
                'floor_all': floor[-1],
                'address': address
        })

    return flats

def main():
    page_count = int(pagination(get_html(base_url)))

    parse(get_html('http://www.domofond.ru/prodazha-kvartiry-magnitogorsk-c2293'))

    flats = []

    for page in range(1, page_count):
        #print('Parsim %d%%' %(page / page_count * 10))
        flats.extend(parse(get_html(base_url + '?Page=%d' % page)))

    for flat in flats:
        print(flat)

if __name__ == '__main__':
    main()


from collections import deque
from queue import Queue

class Graph:
    def __init__(self, v):
        self.V = v
        # 每个node的入度
        self.in_degree = [0 for _ in range(self.V)]
        # 入度为0的集合
        self.dq = deque([])
        # 邻接表
        self.adj = [[] for _ in range(self.V)]

    def add_edge(self, v1, v2):
        self.adj[v1].append(v2)
        self.in_degree[v2]+=1

    def topo_sort(self):
        for v in range(self.V):
            if self.in_degree[v]==0:
                self.dq.append(v)

        # print(self.dq)
        count = 0
        while len(self.dq) != 0:
            v = self.dq.popleft()
            print(v)
            count += 1
            for w in self.adj[v]:
                self.in_degree[w]-=1
                if self.in_degree[w]==0:
                    self.dq.append(w)
            # print(self.dq)

        if count<self.V:
            return False
        else:
            return True

if __name__ == "__main__":
    g = Graph(6)
    g.add_edge(5, 2)
    g.add_edge(5, 0)
    g.add_edge(4, 0)
    g.add_edge(4, 1)
    g.add_edge(2, 3)
    g.add_edge(3, 1)
    print(g.in_degree)
    print(g.adj)
    print(g.topo_sort())


from dataclasses import dataclass

import pytest

from jinja2_component.environment import ComponentEnvironment

@pytest.fixture
def env():
    return ComponentEnvironment()

@pytest.fixture
def hello_component():
    @dataclass
    class Hello:
        template_string: str = '
Hello
'

    return Hello

@pytest.fixture
def helloworld_component():
    @dataclass
    class HelloWorld:
        name: str = 'World'
        template_string: str = '
Hello {{name}}
'

    return HelloWorld

def test_import():
    assert 'ComponentEnvironment' == ComponentEnvironment.__name__

def test_construction(env):
    k = 'jinja2_component.extension.ComponentExtension'
    ext = env.extensions
    assert 'ComponentEnvironment' == env.__class__.__name__
    assert {} == env.components
    assert k in ext.keys()
    assert 1 == len(ext.keys())

def test_register_components(env, hello_component):
    assert {} == env.components
    env.register_components([hello_component])
    tag_name = hello_component.__name__
    extensions = list(env.extensions.values())
    assert tag_name in extensions[0].tags

def test_load_template(env, hello_component):
    env.register_components([hello_component])
    hc = hello_component()
    env.load_template(hc)
    assert hello_component.__name__ in env.templates

def test_render_component(helloworld_component):
    env = ComponentEnvironment([helloworld_component])
    result = env.render_component(helloworld_component)
    assert '
Hello World
' == result

def test_render_component_extra(helloworld_component):
    env = ComponentEnvironment([helloworld_component])
    extra_context = dict(name='extra')
    result = env.render_component(helloworld_component,
                                  extra_context=extra_context)
    assert '
Hello extra
' == result

def test_render_string(env, helloworld_component):
    env.register_components([helloworld_component])
    ts = '{% HelloWorld %}'
    result = env.render_string(ts)
    assert '
Hello World
' == result

# This is the problematic one
# def test_render_string_context(env, helloworld_component):
#     env.register_components([helloworld_component])
#     ts = '{% HelloWorld name=context_name %}'
#     context = dict(context_name='CONTEXTBABY')
#     result = env.render_string(ts, context=context)
#     assert '
Hello CONTEXTBABY
' == result


import numpy as np
import itertools

# # define strength of cards
strength = {f'{i}': i for i in range(1, 11)}
strength['Walet'] = 11
strength['Dama'] = 12
strength['Król'] = 13
strength['As'] = 14

# create and mix cards based on given colors and numbers
def mix_cards(colors, numbers):
    # Create all pack
    allCards = []
    for col, num in itertools.product(colors, numbers):
        allCards.append(num + ' ' + col)
    # shuffled cards
    np.random.shuffle(allCards)
    return allCards

def deal_cards(allCards, n_players):
    game = []
    for i in range(n_players):
        cards = []
        for j in range(len(allCards)):
            if j % n_players == i:
                cards.append(allCards[j])
        game.append(cards)
    return game

# remove player from game
def remove_player(game, players):
    # if player has no more cards is removed from game
    players_to_remove = []
    for card, player in zip(game, players):
        if card == []:
            players_to_remove.append(player)
    for item in players_to_remove:
        players.remove(item)
        game.remove([])
    return game, players

# function which give back cards to owners (used after war ended with dead-heat)
def cards_come_back(game, players, cards_to_take):
    n = len(players)
    for i in range(n):
        for j in range(len(cards_to_take)):
            if j % n == i:
                game[i].append(cards_to_take[j])
    return game

# one round
def one_round(game, players, war):
    """
    :param game: list containing lists of cards for each player
    :param players: list of players - length must be the same as length of game
    :param war: bool - True if current round is the war
    :return: winner, cards which got winner, players, game
    """
    # cards to win
    cards_to_take = []
    round_ = []
    # print players and their cards
    for card, player in zip(game, players):
        print(player, card)

    if war:
        # if round is the war, first card doesn't matter so it goes strictly to cards to win
        for card, player in zip(game, players):
            cards_to_take.append(card[0])
            card.pop(0)
        # if one of players had only one card, remove him from game
        game2, players2 = remove_player(game.copy(), players.copy())

        # if all players were removed, then it is dead-heat
        if len(players2) == 0:
            game = cards_come_back(game, players, cards_to_take)
            return None, cards_to_take, players, game
        # update game and players to state after removing players
        game = game2
        players = players2

    # if only one player stayed, it's winner
    if len(players) == 1:
        return 0, cards_to_take, players, game

    else:
        # choose first card of each player and remove it from cards
        for card, player in zip(game, players):
            round_.append(card[0])
            card.pop(0)
        # scores for cards
        scores = list(map(lambda x: strength[x.split()[0]], round_))
        scores = np.array(scores)
        # maximum score
        max_val = np.max(scores)
        # add cards to cards for winner
        cards_to_take += round_

        # if there is only one max score, we have winner! Otherwise we have war
        if np.sum(scores == max_val) == 1:
            winner = np.argmax(scores)
        else:
            print('wojna')

            # choose players for war and their cards
            players_in_war = [i for i in range(len(scores)) if scores[i] == max_val]
            players_for_war = [players[i] for i in players_in_war]
            print('w wojnie biorą udział gracze', players_for_war)
            game_for_war = [game[i] for i in players_in_war]

            # if any players for war has no more cards, remove him from war
            if [] in game_for_war:
                game_for_war2, players_for_war2 = remove_player(game_for_war, players_for_war)
                # if there is no more players then it is dead-heat, otherwise update players for war and go into war
                if players_for_war2 == []:
                    print('remis - karty wracają do właścicieli')
                    game = cards_come_back(game, players, cards_to_take)
                    return None, [], players, game
                else:
                    game_for_war, players_for_war = game_for_war2, players_for_war2
            winner, cards_to_take_after_war, players_for_war, game_for_war = one_round(
                game_for_war,
                players_for_war,
                war=True)

            # if winer after war is None, it is dead-heat
            if winner is None:
                game = cards_come_back(game, players, cards_to_take)
                return None, [], players, game

            winner = players_in_war[winner]
            cards_to_take += cards_to_take_after_war

        # give cards to winner
        if not war:
            print('zwycięzca to', players[winner])
            print('wygrał', cards_to_take)
            game[winner] += cards_to_take

        # check if there is player with no more cards and remove him/her
        if [] in game:
            game, players = remove_player(game, players)
        return winner, cards_to_take, players, game

def main():
    # how many round with no winner to stop game
    n = 10000
    # how many times repeat that
    m = int(input('Ile razy chcesz powtórzyć rozgrywke? '))
    end = 0
    # min no. of players
    min_players = 2
    # max no. of players
    max_players = 17
    # number of players
    n_players = int(input('Podaj liczbę graczy: '))
    while (n_players < min_players) or (n_players > max_players):
        n_players = int(input(f'Podałeś złą liczbę. Liczba graczy musi być pomiędzy {min_players} a {max_players}: '))
    for j in range(m):
        # define colors and numbers in cards, create and mix them
        colors = ['Kier', 'Karo', 'Trefl', 'Pik']
        numbers = [str(i) for i in range(2, 11)] + ['Walet', 'Dama', 'Król', 'As']
        allCards = mix_cards(colors, numbers)

        # Names of players
        players = [f'Gracz {i + 1}' for i in range(n_players)]

        # dealing cards
        game = deal_cards(allCards, n_players)

        # one game go on till someone win or end if after n rounds there is no winner
        i = 0
        while len(players) > 1 and i < n:
            i += 1

            print('liczba graczy to', len(players))
            winner, cards_to_take, players, game = one_round(game, players, war=False)

        if i == n:
            print('nie skonczyli')
        else:
            print('skonczyli')
            end += 1

    print('Grę skończyli', end, 'razy na', m)

if __name__ == '__main__':
    main()


from pathlib import Path
from typing import Any, Optional

import numpy as np
import torch
import torchvision.transforms.functional as TF
from PIL import Image
from torchvision import transforms
from yacs.config import CfgNode as CN

from .base_dataset import BaseDataset, Ingr, Metadata

class Nutrition5kDataset(BaseDataset):
    """
    Nutrition 5k Dataset Class

    Args:
        imgs_dir (Path): Path to images directory
        metadatas_path (Path): Path to metadatas (csv file)
        splits (list[str]): list of image names in current split
        transform (Optional[transforms.Compose]): RGB image transforms (default=None)
        transform_depth (Optional[transforms.Compose]): depth image transforms (default=None)
        normalize_depth (Optional[transforms.Compose]): dpeth normalization (default=None)
        rotate_flip (Optional[bool]): Whether to perform rotation and horizontal flip data augmentation (default=True)
        w_depth (Optional[bool]): Whether to load depth data (default=True)
        w_mask (Optional[bool]): Whether to load top-1 mask (default=True)
        w_mask_all (Optional[bool]): Whether to load all masks (default=True)
    """
    def __init__(self, *args, **kwargs) -> None:
        super(Nutrition5kDataset, self).__init__(*args, **kwargs)
        self.metadatas_dict: dict[str, Metadata] = {}
        self.init_metadatas()

    @staticmethod
    def transform_(imgs: dict[str, Any]) -> dict[str, torch.Tensor]:
        # rotate
        params = transforms.RandomRotation.get_params([-180, 180])
        imgs = {k: TF.rotate(img, params) for k, img in imgs.items()}

        if np.random.rand() < 0.5:
            imgs = {k: TF.hflip(img) for k, img in imgs.items()}
        return imgs

    def init_metadatas(self):
        """
        Extract metadatas for all dishes and normalize
        """
        mean_metadata = self.mean_metadata
        std_metadata = self.std_metadata
        for line in open(self.metadatas_path, "r").readlines():
            line = line.rstrip()
            data_list = line.split(",")
            data_dict = Metadata()
            dish_id = data_list[0]
            data_dict.dish_id = dish_id
            data_dict.cal = (float(data_list[1]) - mean_metadata.cal) / std_metadata.cal
            data_dict.mass = (
                float(data_list[2]) - mean_metadata.mass
            ) / std_metadata.mass
            data_dict.fat = (float(data_list[3]) - mean_metadata.fat) / std_metadata.fat
            data_dict.carb = (
                float(data_list[4]) - mean_metadata.carb
            ) / std_metadata.carb
            data_dict.protein = (
                float(data_list[5]) - mean_metadata.protein
            ) / std_metadata.protein
            data_dict.ingrs = []
            for id, name, grams, cal, fat, carb, protein in zip(*[data_list[x::7] for x in range(6, 13)]):  # type: ignore
                ingr = Ingr(
                    **{
                        "id": id,
                        "name": name,
                        "grams": grams,
                        "cal": cal,
                        "fat": fat,
                        "carb": carb,
                        "protein": protein,
                    }
                )
                data_dict.ingrs.append(ingr)
            assert len(data_list) == len(data_dict.ingrs) * 7 + 6
            self.metadatas_dict[dish_id] = data_dict

    def __getitem__(self, index: int) -> dict[str, Any]:
        metadata = self.metadatas_dict[self.splits[index]]
        img_dir = Path.joinpath(self.imgs_dir, self.splits[index])
        rgb_path = Path.joinpath(img_dir, "rgb.png")
        depth_path = Path.joinpath(img_dir, "depth_raw.png")
        mask_path = Path.joinpath(img_dir, "mask.pt")
        mask_all_path = Path.joinpath(img_dir, "mask_all.pt")
        sample = {
            "metadata": metadata,
            "rgb_path": str(rgb_path),
            "depth_path": str(depth_path),
        }
        imgs = {}
        rgb_img = Image.open(rgb_path)
        rgb_img = self.transform(rgb_img)
        imgs["rgb_img"] = rgb_img
        if self.w_depth:
            depth_img = Image.open(depth_path)
            depth_img = self.transform_depth(depth_img).float()  # type: ignore
            depth_img = self.normalize_depth(depth_img)
            imgs["depth_img"] = depth_img
        if self.w_mask:
            mask = torch.load(mask_path)
            mask = mask.unsqueeze(0)
            imgs["mask"] = mask
        if self.w_mask_all:
            imgs["mask_all"] = torch.load(mask_all_path)

        if self.rotate_flip:
            imgs = self.transform_(imgs)
        if self.w_mask:
            imgs["mask"] = imgs["mask"].squeeze(0)
        sample.update(imgs)
        return sample

def make_dataset(
    config: Optional[CN],
    imgs_dir: str = ".",
    metadatas_path: str = ".",
    splits_train_path: str = ".",
    splits_test_path: str = ".",
    unnormalized_int_tensor: bool = False,
    **kwargs,
) -> dict[str, BaseDataset]:
    if isinstance(config, CN):
        imgs_dir_p = Path(config.DATA.IMGS_DIR)
        metadatas_path_p = Path(config.DATA.METADATAS_PATH)
        splits_train_path_p = Path(config.DATA.SPLITS_TRAIN_PATH)
        splits_test_path_p = Path(config.DATA.SPLITS_TEST_PATH)
        if config.TRAIN.TRAINER != "multi_mask":
            kwargs.update(w_mask_all=False)
        if config.TRAIN.TRAINER == "normal" or config.TRAIN.TRAINER == "noisy":
            kwargs.update(w_mask=False)
    else:
        imgs_dir_p = Path(imgs_dir)
        metadatas_path_p = Path(metadatas_path)
        splits_train_path_p = Path(splits_train_path)
        splits_test_path_p = Path(splits_test_path)

    ### dishes with obvious mislabels
    removed = {
        "dish_1564159636",
        "dish_1551138237",
        "dish_1551232973",
        "dish_1551381990",
        "dish_1551389458",
        "dish_1551389551",
        "dish_1551389588",
        "dish_1551567508",
        "dish_1551567573",
        "dish_1551567604",
        "dish_1560974769",
    }
    splits_train = [
        line.rstrip() for line in open(splits_train_path_p, "r").readlines()
    ]
    splits_train = list(
        filter(
            lambda t: imgs_dir_p.joinpath(t, "rgb.png").is_file()
            and imgs_dir_p.joinpath(t, "depth_raw.png").is_file()
            and t not in removed,
            splits_train,
        )
    )
    splits_test = [line.rstrip() for line in open(splits_test_path_p, "r").readlines()]
    splits_test = list(
        filter(
            lambda t: imgs_dir_p.joinpath(t, "rgb.png").is_file()
            and imgs_dir_p.joinpath(t, "depth_raw.png").is_file()
            and t not in removed,
            splits_test,
        )
    )

    model_name = config.MODEL.NAME if config is not None else ""

    if unnormalized_int_tensor or model_name == "openseed":
        print("unnormalized")
        transform = transforms.Compose([transforms.PILToTensor()])
        normalize_depth = transforms.Normalize(3091, 5.361)
    else:
        transform = None
        normalize_depth = None
    train_dataset = Nutrition5kDataset(
        imgs_dir_p,
        metadatas_path_p,
        splits_train,
        transform=transform,
        normalize_depth=normalize_depth,
        **kwargs,
    )
    test_dataset = Nutrition5kDataset(
        imgs_dir_p,
        metadatas_path_p,
        splits_test,
        transform=transform,
        normalize_depth=normalize_depth,
        **kwargs,
    )
    return {"train": train_dataset, "test": test_dataset}


from uuid import uuid4
from hyperion.utils.logger import ProjectLogger
from hyperion.audio.io.source.in_device import InDevice

import os
import argparse
import numpy as np
import soundfile as sf

INPRINT_DURATION = 5
SAMPLING_RATE = 16000
CHUNK_DURATION_MS = 512
CHUNK_SIZE = int(SAMPLING_RATE * CHUNK_DURATION_MS / 1000)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Record voice inprint for a new speaker')
    parser.add_argument('name', type=str, help='Speaker name.')
    parser.add_argument('samples_dir', default='~/.hyperion/speakers_samples', type=str, help='Speaker samples directory.')
    args = parser.parse_args()

    args.samples_dir = os.path.expanduser(args.samples_dir)

    mic = InDevice(-2, 16000, db=60)
    mic.open()

    buffer = None
    recording_size = INPRINT_DURATION * SAMPLING_RATE
    stream = mic()
    for chunk in stream:
        ProjectLogger().info('Sound chunk received.')
        if buffer is not None and len(buffer) >= recording_size:
            break
        if chunk is None:
            continue
        buffer = chunk if buffer is None else np.concatenate([buffer, chunk], axis=0)

    mic.close()

    outdir = os.path.join(args.samples_dir, args.name)
    os.makedirs(outdir, exist_ok=True)
    sf.write(os.path.join(outdir, f'{uuid4()}.wav'), buffer[:recording_size, ...], SAMPLING_RATE)


import re

from pygments.style import Style
from pygments.lexer import RegexLexer, bygroups
from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
    Number, Punctuation

from pygments.token import Keyword, Name, Comment, String, Error, \
    Number, Operator, Generic

class MyStyle(Style):
    default_style = ""
    styles = {
        Comment:                'italic #888',
        Keyword:                'bold #005',
        Name:                   '#f00',
        Name.Function:          '#0f0',
        Name.Class:             'bold #0f0',
        String:                 'bg:#eee #111',
        Keyword.Reserved:       'bold ',
    }

class DynaLexer(RegexLexer):
    """
    Lexer for Dyna files.
    """
    name = 'Prolog'
    aliases = ['prolog']
    filenames = ['*.dyna']
    mimetypes = ['text/x-dyna']

    flags = re.UNICODE | re.MULTILINE

    tokens = {
        'root': [
            (r'^#.*', Comment.Single),
            (r'/\*', Comment.Multiline, 'nested-comment'),
            (r'%.*', Comment.Single),

            # aggregator literal
            (r'(max|min|bag|set|[:+*&|?])=|:-|->', Keyword),

            # REPL commands
            ('(clear_memos|retract_rule|help|tickle|check_invariants|compare|obligated_children|rules|state|flush|agenda|query|retract_rule|state|memos|flush|agenda|run_agenda|subscribe|checkpoint|load|sol|verbose)',
             Keyword.Reserved),

            # FLOAT:
            ('-?[0-9]*((\.[0-9]+(e[\-+]?[0-9]+)?)|(e[\-+]?[0-9]+))', Number),
            # INT
            ('-?[0-9]+', Number),

            # VAR:
            ("[$]?[A-ZΑΒΓΔΕΖΗΘΙΚΛΜΝΞΟΠΡΣΤΥΦΧΨΩ][a-zA-Z0-9'_αβγδεζηθικλμνξοπρστυφχψωΑΒΓΔΕΖΗΘΙΚΛΜΝΞΟΠΡΣΤΥΦΧΨΩ]*", Keyword),

            #EQ: "="

            # FUNCTOR1:
            ("([$]?[a-zαβγδεζηθικλμνξοπρστυφχψω]['a-zA-Z0-9_αβγδεζηθικλμνξοπρστυφχψωΑΒΓΔΕΖΗΘΙΚΛΜΝΞΟΠΡΣΤΥΦΧΨΩ]*)", Name.Function),

            # FUNCTOR2:
            ("('[^']+')", Name.Function),

            #COMMENT: "%" /[^\n\r]*/ (NEWLINE|/$/)

            # EOR:
            ('\.\s+|\.\s*$|\.(?=[^\d])', Punctuation),
            ('[()]', Punctuation),

            # Strings
            (r'"((\\"|[^"])*?)"', String),

#            # character literal
#            (r'0\'.', String.Char),
#            (r'0b[01]+', Number.Bin),
#            (r'0o[0-7]+', Number.Oct),
#            (r'0x[0-9a-fA-F]+', Number.Hex),
#            # literal with prepended base
#            (r'\d\d?\'[a-zA-Z0-9]+', Number.Integer),
#            (r'(\d+\.\d*|\d*\.\d+)([eE][+-]?[0-9]+)?', Number.Float),
#            (r'\d+', Number.Integer),
#            (r'[\[\](){}|.,;!]', Punctuation),
#            (r':-|-->', Punctuation),
#            (r'"(?:\\x[0-9a-fA-F]+\\|\\u[0-9a-fA-F]{4}|\\U[0-9a-fA-F]{8}|'
#             r'\\[0-7]+\\|\\["\nabcefnrstv]|[^\\"])*"', String.Double),
#            (r"'(?:''|[^'])*'", String.Atom),  # quoted atom
#            # Needs to not be followed by an atom.
#            # (r'=(?=\s|[a-zA-Z\[])', Operator),
            (r'is\b', Operator),
            (r'(<|>|=<|>=|==|=:=|=|/|//|\*|\+|-)(?=\s|[a-zA-Z0-9\[])', Operator),
#            (r'(mod|div|not)\b', Operator),
            (r'_', Keyword),  # The don't-care variable
#            (r'([a-z]+)(:)', bygroups(Name.Namespace, Punctuation)),
#            (u'([a-z\u00c0-\u1fff\u3040-\ud7ff\ue000-\uffef]'
#             u'[\w$\u00c0-\u1fff\u3040-\ud7ff\ue000-\uffef]*)'
#             u'(\\s*)(:-|-->)',
#             bygroups(Name.Function, Text, Operator)),  # function defn
#            (u'([a-z\u00c0-\u1fff\u3040-\ud7ff\ue000-\uffef]'
#             u'[\w$\u00c0-\u1fff\u3040-\ud7ff\ue000-\uffef]*)'
#             u'(\\s*)(\\()',
#             bygroups(Name.Function, Text, Punctuation)),
#            (u'[a-z\u00c0-\u1fff\u3040-\ud7ff\ue000-\uffef]'
#             u'[\w$\u00c0-\u1fff\u3040-\ud7ff\ue000-\uffef]*',
#             String.Atom),  # atom, characters
#            # This one includes !
#            (u'[#&*+\\-./:<=>?@\\\\^~\u00a1-\u00bf\u2010-\u303f]+',
#             String.Atom),  # atom, graphics
#            (r'[A-Z_]\w*', Name.Variable),
#            (u'\\s+|[\u2000-\u200f\ufff0-\ufffe\uffef]', Text),
        ],
        'nested-comment': [
            (r'\*/', Comment.Multiline, '#pop'),
            (r'/\*', Comment.Multiline, '#push'),
            (r'[^*/]+', Comment.Multiline),
            (r'[*/]', Comment.Multiline),
        ],
    }

    def analyse_text(text):
        return True


"""
任意文件下载漏洞批量下载脚本
"""
import os
import threading

import requests
import threadpool

from config import *

total_urls = 0
complete_num = 0
dir_lock = threading.Lock()
progress_lock = threading.Lock()

def request_handler(urls):
    """
    urls 已经去重去参数，只保留相对路径，假设urls.txt中的链接为 http://xx.xx.com/xx/xx.php,
    执行到这一步就变成了 xx/xx.php。
    接下来可以用自定义函数处理它，比如利用文件包含漏洞下载源码的时候需要用php伪协议。
    @params:
        urls: 形如['xx/xx.php','xx/xxx/xxx.php']储存相对路径的列表
    """
    return urls

def response_handler(response):
    """
    用于处理下载后的页面，有些情况下服务器会返回多余的数据，这时候可以实现这个函数把多余数据删掉。
    @params:
        response: requests.Response 类，服务器响应包，详情参考requests文档。
    @rtype bytes
    """
    return response.content

def printProgress(iteration, total, prefix='', suffix='', decimals=1, barLength=100):
    """
    Call in a loop to create a terminal progress bar
    @params:
        iteration   - Required  : current iteration (Int)
        total       - Required  : total iterations (Int)
        prefix      - Optional  : prefix string (Str)
        suffix      - Optional  : suffix string (Str)
        decimals    - Optional  : positive number of decimals in percent complete (Int)
        barLength   - Optional  : character length of bar (Int)
    """
    import sys
    formatStr = "{0:." + str(decimals) + "f}"
    percent = formatStr.format(100 * (iteration / float(total)))
    filledLength = int(round(barLength * iteration / float(total)))
    bar = '#' * filledLength + '-' * (barLength - filledLength)
    sys.stdout.write('\r%s |%s| %s%s %s  ' %
                     (prefix, bar, percent, '%', suffix)),
    if iteration == total:
        sys.stdout.write('\n')
    sys.stdout.flush()

def data_clean(urls):
    slice_start = urls[0].index(SITE)+len(SITE)+1
    urls = set(map(lambda url: url.partition('?')[0].strip(), urls))
    urls = filter(lambda url: url.rpartition('.')[-1] in EXT_NAME, urls)
    urls = list(map(lambda url: url[slice_start:], urls))
    urls = request_handler(urls)
    return urls

def real_down(url):
    file_path = os.path.join('..', SITE, url)
    dir_path = file_path.rpartition('/')[0]
    dir_lock.acquire()
    if not os.path.exists(dir_path) and dir_path != '':
        os.makedirs(dir_path)
    dir_lock.release()
    # 参数 urls 经过去重，因此 file_path 具有唯一性，不存在多个线程同时操作同一 file_path 的情况，所以下面的操作无需上锁
    if not os.path.exists(file_path):
        response = requests.get(LINK.format(url))
        download_content = response_handler(response)
        with open(file_path, 'wb') as f:
            f.write(download_content)
    global complete_num
    progress_lock.acquire()
    complete_num += 1
    printProgress(complete_num, total_urls, 'Progress:', url, barLength=30)
    progress_lock.release()

if __name__ == '__main__':
    with open('urls.txt', 'r', encoding='utf8') as f:
        relative_urls = data_clean(f.readlines())
    total_urls = len(relative_urls)
    pool = threadpool.ThreadPool(THREADS)
    tasks = threadpool.makeRequests(real_down, relative_urls)
    [pool.putRequest(task) for task in tasks]
    pool.wait()


from PIL import Image
from absl import app
from options import FLAGS as opts
import data
import datetime
import functools
import glob
import losses
import networks
import numpy as np
import options
import os.path as osp
import random
import skimage.measure
import staged_model
import tensorflow as tf
import time
import utils

def build_model_fn(use_exponential_moving_average=True):
  """Builds and returns the model function for an estimator.

  Args:
    use_exponential_moving_average: bool. If true, the exponential moving
    average will be used.

  Returns:
    function, the model_fn function typically required by an estimator.
  """
  arch_type = opts.arch_type
  use_appearance = opts.use_appearance
  def model_fn(features, labels, mode, params):
    """An estimator build_fn."""
    del labels, params
    if mode == tf.estimator.ModeKeys.TRAIN:
      step = tf.train.get_global_step()

      x_in = features['conditional_input']
      x_gt = features['expected_output']  # ground truth output
      x_app = features['peek_input']

      if opts.training_pipeline == 'staged':
        ops = staged_model.create_computation_graph(x_in, x_gt, x_app=x_app,
                                                    arch_type=opts.arch_type)
        op_increment_step = tf.assign_add(step, 1)
        train_disc_op = ops['train_disc_op']
        train_renderer_op = ops['train_renderer_op']
        train_op = tf.group(train_disc_op, train_renderer_op, op_increment_step)

        utils.HookReport.log_tensor(ops['total_loss_d'], 'total_loss_d')
        utils.HookReport.log_tensor(ops['loss_d_real'], 'loss_d_real')
        utils.HookReport.log_tensor(ops['loss_d_fake'], 'loss_d_fake')
        utils.HookReport.log_tensor(ops['total_loss_g'], 'total_loss_g')
        utils.HookReport.log_tensor(ops['loss_g_gan'], 'loss_g_gan')
        utils.HookReport.log_tensor(ops['loss_g_recon'], 'loss_g_recon')
        utils.HookReport.log_tensor(step, 'global_step')

        return tf.estimator.EstimatorSpec(
            mode=mode, loss=ops['total_loss_d'] + ops['total_loss_g'],
            train_op=train_op)
      else:
        raise NotImplementedError('%s training is not implemented.' %
                                  opts.training_pipeline)
    elif mode == tf.estimator.ModeKeys.EVAL:
      raise NotImplementedError('Eval is not implemented.')
    else:  # all below modes are for difference inference tasks.
      # Build network and initialize inference variables.
      g_func = networks.RenderingModel(arch_type, use_appearance)
      if use_appearance:
        app_func = g_func.get_appearance_encoder()
      if use_exponential_moving_average:
        ema = tf.train.ExponentialMovingAverage(decay=0.999)
        var_dict = ema.variables_to_restore()
        tf.train.init_from_checkpoint(osp.join(opts.train_dir), var_dict)

      if mode == tf.estimator.ModeKeys.PREDICT:
        x_in = features['conditional_input']
        if use_appearance:
          x_app = features['peek_input']
          x_app_embedding, _, _ = app_func(x_app)
        else:
          x_app_embedding = None
        y = g_func(x_in, x_app_embedding)
        tf.logging.info('DBG: shape of y during prediction %s.' % str(y.shape))
        return tf.estimator.EstimatorSpec(mode=mode, predictions=y)

      # 'eval_subset' mode is same as PREDICT but it concatenates the output to
      # the input render, semantic map and ground truth for easy comparison.
      elif mode == 'eval_subset':
        x_in = features['conditional_input']
        x_gt = features['expected_output']
        if use_appearance:
          x_app = features['peek_input']
          x_app_embedding, _, _ = app_func(x_app)
        else:
          x_app_embedding = None
        y = g_func(x_in, x_app_embedding)
        tf.logging.info('DBG: shape of y during prediction %s.' % str(y.shape))
        x_in_rgb = tf.slice(x_in, [0, 0, 0, 0], [-1, -1, -1, 3])
        if opts.use_semantic:
          x_in_semantic = tf.slice(x_in, [0, 0, 0, 4], [-1, -1, -1, 3])
          output_tuple = tf.concat([x_in_rgb, x_in_semantic, y, x_gt], axis=2)
        else:
          output_tuple = tf.concat([x_in_rgb, y, x_gt], axis=2)
        return tf.estimator.EstimatorSpec(mode=mode, predictions=output_tuple)

      # 'compute_appearance' mode computes and returns the latent z vector.
      elif mode == 'compute_appearance':
        assert use_appearance, 'use_appearance is set to False!'
        x_app_in = features['peek_input']
        # NOTE the following line is a temporary hack (which is
        # specially bad for inputs smaller than 512x512).
        x_app_in = tf.image.resize_image_with_crop_or_pad(x_app_in, 512, 512)
        app_embedding, _, _ = app_func(x_app_in)
        return tf.estimator.EstimatorSpec(mode=mode, predictions=app_embedding)

      # 'interpolate_appearance' mode expects an already computed latent z
      # vector as input passed a value to the dict key 'appearance_embedding'.
      elif mode == 'interpolate_appearance':
        assert use_appearance, 'use_appearance is set to False!'
        x_in = features['conditional_input']
        x_app_embedding = features['appearance_embedding']
        y = g_func(x_in, x_app_embedding)
        tf.logging.info('DBG: shape of y during prediction %s.' % str(y.shape))
        return tf.estimator.EstimatorSpec(mode=mode, predictions=y)
      else:
        raise ValueError('Unsupported mode: ' + mode)

  return model_fn

def make_sample_grid_and_save(est, dataset_name, dataset_parent_dir, grid_dims,
                              output_dir, cur_nimg):
  """Evaluate a fixed set of validation images and save output.

  Args:
    est: tf,estimator.Estimator, TF estimator to run the predictions.
    dataset_name: basename for the validation tfrecord from which to load
      validation images.
    dataset_parent_dir: path to a directory containing the validation tfrecord.
    grid_dims: 2-tuple int for the grid size (1 unit = 1 image).
    output_dir: string, where to save image samples.
    cur_nimg: int, current number of images seen by training.

  Returns:
    None.
  """
  num_examples = grid_dims[0] * grid_dims[1]
  def input_val_fn():
    dict_inp = data.provide_data(
        dataset_name=dataset_name, parent_dir=dataset_parent_dir, subset='val',
        batch_size=1, crop_flag=True, crop_size=opts.train_resolution,
        seeds=[0], max_examples=num_examples,
        use_appearance=opts.use_appearance, shuffle=0)
    x_in = dict_inp['conditional_input']
    x_gt = dict_inp['expected_output']  # ground truth output
    x_app = dict_inp['peek_input']
    return x_in, x_gt, x_app

  def est_input_val_fn():
    x_in, _, x_app = input_val_fn()
    features = {'conditional_input': x_in, 'peek_input': x_app}
    return features

  images = [x for x in est.predict(est_input_val_fn)]
  images = np.array(images, 'f')
  images = images.reshape(grid_dims + images.shape[1:])
  utils.save_images(utils.to_png(utils.images_to_grid(images)), output_dir,
                    cur_nimg)

def visualize_image_sequence(est, dataset_name, dataset_parent_dir,
                             input_sequence_name, app_base_path, output_dir):
  """Generates an image sequence as a video and stores it to disk."""
  batch_sz = opts.batch_size
  def input_seq_fn():
    dict_inp = data.provide_data(
        dataset_name=dataset_name, parent_dir=dataset_parent_dir,
        subset=input_sequence_name, batch_size=batch_sz, crop_flag=False,
        seeds=None, use_appearance=False, shuffle=0)
    x_in = dict_inp['conditional_input']
    return x_in

  # Compute appearance embedding only once and use it for all input frames.
  app_rgb_path = app_base_path + '_reference.png'
  app_rendered_path = app_base_path + '_color.png'
  app_depth_path = app_base_path + '_depth.png'
  app_sem_path = app_base_path + '_seg_rgb.png'
  x_app = _load_and_concatenate_image_channels(
      app_rgb_path, app_rendered_path, app_depth_path, app_sem_path)
  def seq_with_single_appearance_inp_fn():
    """input frames with a fixed latent appearance vector."""
    x_in_op = input_seq_fn()
    x_app_op = tf.convert_to_tensor(x_app)
    x_app_tiled_op = tf.tile(x_app_op, [tf.shape(x_in_op)[0], 1, 1, 1])
    return {'conditional_input': x_in_op,
            'peek_input': x_app_tiled_op}

  images = [x for x in est.predict(seq_with_single_appearance_inp_fn)]
  for i, gen_img in enumerate(images):
    output_file_path = osp.join(output_dir, 'out_%04d.png' % i)
    print('Saving frame #%d to %s' % (i, output_file_path))
    with tf.gfile.Open(output_file_path, 'wb') as f:
      f.write(utils.to_png(gen_img))

def train(dataset_name, dataset_parent_dir, load_pretrained_app_encoder,
          load_trained_fixed_app, save_samples_kimg=50):
  """Main training procedure.

  The trained model is saved in opts.train_dir, the function itself does not
   return anything.

  Args:
    save_samples_kimg: int, period (in KiB) to save sample images.

  Returns:
    None.
  """
  image_dir = osp.join(opts.train_dir, 'images')  # to save validation images.
  tf.gfile.MakeDirs(image_dir)
  config = tf.estimator.RunConfig(
      save_summary_steps=(1 << 10) // opts.batch_size,
      save_checkpoints_steps=(save_samples_kimg << 10) // opts.batch_size,
      keep_checkpoint_max=5,
      log_step_count_steps=1 << 30)
  model_dir = opts.train_dir
  if (opts.use_appearance and load_trained_fixed_app and
      not tf.train.latest_checkpoint(model_dir)):
    tf.logging.warning('***** Loading resume_step from %s!' %
                       opts.fixed_appearance_train_dir)
    resume_step = utils.load_global_step_from_checkpoint_dir(
        opts.fixed_appearance_train_dir)
  else:
    tf.logging.warning('***** Loading resume_step (if any) from %s!' %
                       model_dir)
    resume_step = utils.load_global_step_from_checkpoint_dir(model_dir)
  if resume_step != 0:
    tf.logging.warning('****** Resuming training at %d!' % resume_step)

  model_fn = build_model_fn()  # model function for TFEstimator.

  hooks = [utils.HookReport(1 << 12, opts.batch_size)]

  if opts.use_appearance and load_pretrained_app_encoder:
    tf.logging.warning('***** will warm-start from %s!' %
                       opts.appearance_pretrain_dir)
    ws = tf.estimator.WarmStartSettings(
        ckpt_to_initialize_from=opts.appearance_pretrain_dir,
        vars_to_warm_start='appearance_net/.*')
  elif opts.use_appearance and load_trained_fixed_app:
    tf.logging.warning('****** finetuning will warm-start from %s!' %
                       opts.fixed_appearance_train_dir)
    ws = tf.estimator.WarmStartSettings(
        ckpt_to_initialize_from=opts.fixed_appearance_train_dir,
        vars_to_warm_start='.*')
  else:
    ws = None
    tf.logging.warning('****** No warm-starting; using random initialization!')

  est = tf.estimator.Estimator(model_fn, model_dir, config, params={},
                               warm_start_from=ws)

  for next_kimg in range(opts.save_samples_kimg, opts.total_kimg + 1,
                         opts.save_samples_kimg):
    next_step = (next_kimg << 10) // opts.batch_size
    if opts.num_crops == -1:  # use random crops
      crop_seeds = None
    else:
      crop_seeds = list(100 * np.arange(opts.num_crops))
    input_train_fn = functools.partial(
        data.provide_data, dataset_name=dataset_name,
        parent_dir=dataset_parent_dir, subset='train',
        batch_size=opts.batch_size, crop_flag=True,
        crop_size=opts.train_resolution, seeds=crop_seeds,
        use_appearance=opts.use_appearance)
    est.train(input_train_fn, max_steps=next_step, hooks=hooks)
    tf.logging.info('DBG: kimg=%d, cur_step=%d' % (next_kimg, next_step))
    tf.logging.info('DBG: Saving a validation grid image %06d to %s' % (
        next_kimg, image_dir))
    make_sample_grid_and_save(est, dataset_name, dataset_parent_dir, (3, 3),
                              image_dir, next_kimg << 10)

def _build_inference_estimator(model_dir):
  model_fn = build_model_fn()
  est = tf.estimator.Estimator(model_fn, model_dir)
  return est

def evaluate_sequence(dataset_name, dataset_parent_dir, virtual_seq_name,
                      app_base_path):
  output_dir = osp.join(opts.train_dir, 'seq_output_%s' % virtual_seq_name)
  tf.gfile.MakeDirs(output_dir)
  est = _build_inference_estimator(opts.train_dir)
  visualize_image_sequence(est, dataset_name, dataset_parent_dir,
                           virtual_seq_name, app_base_path, output_dir)

def evaluate_image_set(dataset_name, dataset_parent_dir, subset_suffix,
                       output_dir=None, batch_size=6):
  if output_dir is None:
    output_dir = osp.join(opts.train_dir, 'validation_output_%s' % subset_suffix)
  tf.gfile.MakeDirs(output_dir)
  model_fn_old = build_model_fn()
  def model_fn_wrapper(features, labels, mode, params):
    del mode
    return model_fn_old(features, labels, 'eval_subset', params)
  model_dir = opts.train_dir
  est = tf.estimator.Estimator(model_fn_wrapper, model_dir)
  est_inp_fn = functools.partial(
      data.provide_data, dataset_name=dataset_name,
      parent_dir=dataset_parent_dir, subset=subset_suffix,
      batch_size=batch_size, use_appearance=opts.use_appearance, shuffle=0)

  print('Evaluating images for subset %s' % subset_suffix)
  images = [x for x in est.predict(est_inp_fn)]
  print('Evaluated %d images' % len(images))
  for i, img in enumerate(images):
    output_file_path = osp.join(output_dir, 'out_%04d.png' % i)
    print('Saving file #%d: %s' % (i, output_file_path))
    with tf.gfile.Open(output_file_path, 'wb') as f:
      f.write(utils.to_png(img))

def _load_and_concatenate_image_channels(rgb_path=None, rendered_path=None,
                                         depth_path=None, seg_path=None,
                                         size_multiple=64):
  """Prepares a single input for the network."""
  if (rgb_path is None and rendered_path is None and depth_path is None and
      seg_path is None):
    raise ValueError('At least one of the inputs has to be not None')

  channels = ()
  if rgb_path is not None:
    rgb_img = np.array(Image.open(rgb_path)).astype(np.float32)
    rgb_img = utils.crop_to_multiple(rgb_img, size_multiple)
    channels = channels + (rgb_img,)
  if rendered_path is not None:
    rendered_img = np.array(Image.open(rendered_path)).astype(np.float32)
    if not opts.use_alpha:
      rendered_img = rendered_img[:, :, :3]  # drop the alpha channel
    rendered_img = utils.crop_to_multiple(rendered_img, size_multiple)
    channels = channels + (rendered_img,)
  if depth_path is not None:
    depth_img = np.array(Image.open(depth_path))
    depth_img = depth_img.astype(np.float32)
    depth_img = utils.crop_to_multiple(depth_img[:, :, np.newaxis],
                                       size_multiple)
    channels = channels + (depth_img,)
    # depth_img = depth_img * (2.0 / 255) - 1.0
  if seg_path is not None:
    seg_img = np.array(Image.open(seg_path)).astype(np.float32)
    seg_img = utils.crop_to_multiple(seg_img, size_multiple)
    channels = channels + (seg_img,)
  # Concatenate and normalize channels
  img = np.dstack(channels)
  img = np.expand_dims(img, axis=0)
  img = img * (2.0 / 255) - 1.0
  return img

def infer_dir(model_dir, input_dir, output_dir):
  tf.gfile.MakeDirs(output_dir)
  est = _build_inference_estimator(opts.train_dir)

  def read_image(base_path, is_appearance=False):
    if is_appearance:
      ref_img_path = base_path + '_reference.png'
    else:
      ref_img_path = None
    rendered_img_path = base_path + '_color.png'
    depth_img_path = base_path + '_depth.png'
    seg_img_path = base_path + '_seg_rgb.png'
    img = _load_and_concatenate_image_channels(
        rgb_path=ref_img_path, rendered_path=rendered_img_path,
        depth_path=depth_img_path, seg_path=seg_img_path)
    return img

  def get_inference_input_fn(base_path, app_base_path):
    x_in = read_image(base_path, False)
    x_app_in = read_image(app_base_path, True)
    def infer_input_fn():
      return {'conditional_input': x_in, 'peek_input': x_app_in}
    return infer_input_fn

  file_paths = sorted(glob.glob(osp.join(input_dir, '*_depth.png')))
  base_paths = [x[:-10] for x in file_paths]  # remove the '_depth.png' suffix
  for inp_base_path in base_paths:
    est_inp_fn = get_inference_input_fn(inp_base_path, inp_base_path)
    img = next(est.predict(est_inp_fn))
    basename = osp.basename(inp_base_path)
    output_img_path = osp.join(output_dir, basename + '_out.png')
    print('Saving generated image to %s' % output_img_path)
    with tf.gfile.Open(output_img_path, 'wb') as f:
      f.write(utils.to_png(img))

def joint_interpolation(model_dir, app_input_dir, st_app_basename,
                        end_app_basename, camera_path_dir):
  """
  Interpolates both viewpoint and appearance between two input images.
  """
  # Create output direcotry
  output_dir = osp.join(model_dir, 'joint_interpolation_out')
  tf.gfile.MakeDirs(output_dir)

  # Build estimator
  model_fn_old = build_model_fn()
  def model_fn_wrapper(features, labels, mode, params):
    del mode
    return model_fn_old(features, labels, 'interpolate_appearance', params)
  def appearance_model_fn(features, labels, mode, params):
    del mode
    return model_fn_old(features, labels, 'compute_appearance', params)
  config = tf.estimator.RunConfig(
      save_summary_steps=1000, save_checkpoints_steps=50000,
      keep_checkpoint_max=50, log_step_count_steps=1 << 30)
  model_dir = model_dir
  est = tf.estimator.Estimator(model_fn_wrapper, model_dir, config, params={})
  est_app = tf.estimator.Estimator(appearance_model_fn, model_dir, config,
                                   params={})

  # Compute appearance embeddings for the two input appearance images.
  app_inputs = []
  for app_basename in [st_app_basename, end_app_basename]:
    app_rgb_path = osp.join(app_input_dir, app_basename + '_reference.png')
    app_rendered_path = osp.join(app_input_dir, app_basename + '_color.png')
    app_depth_path = osp.join(app_input_dir, app_basename + '_depth.png')
    app_seg_path = osp.join(app_input_dir, app_basename + '_seg_rgb.png')
    app_in = _load_and_concatenate_image_channels(
        rgb_path=app_rgb_path, rendered_path=app_rendered_path,
        depth_path=app_depth_path, seg_path=app_seg_path)
    # app_inputs.append(tf.convert_to_tensor(app_in))
    app_inputs.append(app_in)

  embedding1 = next(est_app.predict(
      lambda: {'peek_input': app_inputs[0]}))
  embedding1 = np.expand_dims(embedding1, axis=0)
  embedding2 = next(est_app.predict(
      lambda: {'peek_input': app_inputs[1]}))
  embedding2 = np.expand_dims(embedding2, axis=0)

  file_paths = sorted(glob.glob(osp.join(camera_path_dir, '*_depth.png')))
  base_paths = [x[:-10] for x in file_paths]  # remove the '_depth.png' suffix

  # Compute interpolated appearance embeddings
  num_interpolations = len(base_paths)
  interpolated_embeddings = []
  delta_vec = (embedding2 - embedding1) / (num_interpolations - 1)
  for delta_iter in range(num_interpolations):
    x_app_embedding = embedding1 + delta_iter * delta_vec
    interpolated_embeddings.append(x_app_embedding)

  # Generate and save interpolated images
  for frame_idx, embedding in enumerate(interpolated_embeddings):
    # Read in input frame
    frame_render_path = osp.join(base_paths[frame_idx] + '_color.png')
    frame_depth_path = osp.join(base_paths[frame_idx] + '_depth.png')
    frame_seg_path = osp.join(base_paths[frame_idx] + '_seg_rgb.png')
    x_in = _load_and_concatenate_image_channels(
        rgb_path=None, rendered_path=frame_render_path,
        depth_path=frame_depth_path, seg_path=frame_seg_path)

    img = next(est.predict(
        lambda: {'conditional_input': tf.convert_to_tensor(x_in),
                 'appearance_embedding': tf.convert_to_tensor(embedding)}))
    output_img_name = '%s_%s_%03d.png' % (st_app_basename, end_app_basename,
                                          frame_idx)
    output_img_path = osp.join(output_dir, output_img_name)
    print('Saving interpolated image to %s' % output_img_path)
    with tf.gfile.Open(output_img_path, 'wb') as f:
      f.write(utils.to_png(img))

def interpolate_appearance(model_dir, input_dir, target_img_basename,
                           appearance_img1_basename, appearance_img2_basename):
  # Create output direcotry
  output_dir = osp.join(model_dir, 'interpolate_appearance_out')
  tf.gfile.MakeDirs(output_dir)

  # Build estimator
  model_fn_old = build_model_fn()
  def model_fn_wrapper(features, labels, mode, params):
    del mode
    return model_fn_old(features, labels, 'interpolate_appearance', params)
  def appearance_model_fn(features, labels, mode, params):
    del mode
    return model_fn_old(features, labels, 'compute_appearance', params)
  config = tf.estimator.RunConfig(
      save_summary_steps=1000, save_checkpoints_steps=50000,
      keep_checkpoint_max=50, log_step_count_steps=1 << 30)
  model_dir = model_dir
  est = tf.estimator.Estimator(model_fn_wrapper, model_dir, config, params={})
  est_app = tf.estimator.Estimator(appearance_model_fn, model_dir, config,
                                   params={})

  # Compute appearance embeddings for the two input appearance images.
  app_inputs = []
  for app_basename in [appearance_img1_basename, appearance_img2_basename]:
    app_rgb_path = osp.join(input_dir, app_basename + '_reference.png')
    app_rendered_path = osp.join(input_dir, app_basename + '_color.png')
    app_depth_path = osp.join(input_dir, app_basename + '_depth.png')
    app_seg_path = osp.join(input_dir, app_basename + '_seg_rgb.png')
    app_in = _load_and_concatenate_image_channels(
        rgb_path=app_rgb_path, rendered_path=app_rendered_path,
        depth_path=app_depth_path, seg_path=app_seg_path)
    # app_inputs.append(tf.convert_to_tensor(app_in))
    app_inputs.append(app_in)

  embedding1 = next(est_app.predict(
      lambda: {'peek_input': app_inputs[0]}))
  embedding2 = next(est_app.predict(
      lambda: {'peek_input': app_inputs[1]}))
  embedding1 = np.expand_dims(embedding1, axis=0)
  embedding2 = np.expand_dims(embedding2, axis=0)

  # Compute interpolated appearance embeddings
  num_interpolations = 10
  interpolated_embeddings = []
  delta_vec = (embedding2 - embedding1) / num_interpolations
  for delta_iter in range(num_interpolations + 1):
    x_app_embedding = embedding1 + delta_iter * delta_vec
    interpolated_embeddings.append(x_app_embedding)

  # Read in the generator input for the target image to render
  rendered_img_path = osp.join(input_dir, target_img_basename + '_color.png')
  depth_img_path = osp.join(input_dir, target_img_basename + '_depth.png')
  seg_img_path = osp.join(input_dir, target_img_basename + '_seg_rgb.png')
  x_in = _load_and_concatenate_image_channels(
      rgb_path=None, rendered_path=rendered_img_path,
      depth_path=depth_img_path, seg_path=seg_img_path)

  # Generate and save interpolated images
  for interpolate_iter, embedding in enumerate(interpolated_embeddings):
    img = next(est.predict(
        lambda: {'conditional_input': tf.convert_to_tensor(x_in),
                 'appearance_embedding': tf.convert_to_tensor(embedding)}))
    output_img_name = 'interpolate_%s_%s_%s_%03d.png' % (
        target_img_basename, appearance_img1_basename, appearance_img2_basename,
        interpolate_iter)
    output_img_path = osp.join(output_dir, output_img_name)
    print('Saving interpolated image to %s' % output_img_path)
    with tf.gfile.Open(output_img_path, 'wb') as f:
      f.write(utils.to_png(img))

def main(argv):
  del argv
  configs_str = options.list_options()
  tf.gfile.MakeDirs(opts.train_dir)
  with tf.gfile.Open(osp.join(opts.train_dir, 'configs.txt'), 'wb') as f:
    f.write(configs_str)
  tf.logging.info('Local configs\n%s' % configs_str)

  if opts.run_mode == 'train':
    dataset_name = opts.dataset_name
    dataset_parent_dir = opts.dataset_parent_dir
    load_pretrained_app_encoder = opts.load_pretrained_app_encoder
    load_trained_fixed_app = opts.load_from_another_ckpt
    batch_size = opts.batch_size
    train(dataset_name, dataset_parent_dir, load_pretrained_app_encoder,
          load_trained_fixed_app)
  elif opts.run_mode == 'eval':  # generate a camera path output sequence from TFRecord inputs.
    dataset_name = opts.dataset_name
    dataset_parent_dir = opts.dataset_parent_dir
    virtual_seq_name = opts.virtual_seq_name
    inp_app_img_base_path = opts.inp_app_img_base_path
    evaluate_sequence(dataset_name, dataset_parent_dir, virtual_seq_name,
                      inp_app_img_base_path)
  elif opts.run_mode == 'eval_subset':  # generate output for validation set (encoded as TFRecords)
    dataset_name = opts.dataset_name
    dataset_parent_dir = opts.dataset_parent_dir
    virtual_seq_name = opts.virtual_seq_name
    evaluate_image_set(dataset_name, dataset_parent_dir, virtual_seq_name,
                       opts.output_validation_dir, opts.batch_size)
  elif opts.run_mode == 'eval_dir':  # evaluate output for a directory with input images
    input_dir = opts.inference_input_path
    output_dir = opts.inference_output_dir
    model_dir = opts.train_dir
    infer_dir(model_dir, input_dir, output_dir)
  elif opts.run_mode == 'interpolate_appearance':  # interpolate appearance only between two images.
    model_dir = opts.train_dir
    input_dir = opts.inference_input_path
    target_img_basename = opts.target_img_basename
    app_img1_basename = opts.appearance_img1_basename
    app_img2_basename = opts.appearance_img2_basename
    interpolate_appearance(model_dir, input_dir, target_img_basename,
                           app_img1_basename, app_img2_basename)
  elif opts.run_mode == 'joint_interpolation':  # interpolate viewpoint and appearance between two images
    model_dir = opts.train_dir
    app_input_dir = opts.inference_input_path
    st_app_basename = opts.appearance_img1_basename
    end_app_basename = opts.appearance_img2_basename
    frames_dir = opts.frames_dir
    joint_interpolation(model_dir, app_input_dir, st_app_basename,
                        end_app_basename, frames_dir)
  else:
    raise ValueError('Unsupported --run_mode %s' % opts.run_mode)

if __name__ == '__main__':
  app.run(main)


import os
import json

import torch
import numpy as np
import random

from asd import BEST_PRETRAIN_MODEL_NAME, create_agent
from dataset import ASDPatientInteractionSimulator, SimPaDataset

EVAL_FILE = "Metrics.json"

def write_json(data, fp):
    with open(fp, "w") as outfile:
        json.dump(data, outfile, indent=4)

def initialize_seed(seed):
    """Method for seeding the random generators.

    Parameters
    ----------
    seed: int
        the seed to be used.

    Returns
    -------
    None

    """
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)

def load_model(model, path):
    data = torch.load(path)
    model.load_state_dict(data["agent_state_dict"])
    model.eval()
    return model

def asd_classifier_test(agent, eval_dl, args):
    with torch.no_grad():
        agent.train(False)
        results = eval_dl.evaluate(
            agent,
            args.interaction_length,
            args.num_valid_trajs,
            None,
            args.seed,
            args.cuda_idx,
            args.thres,
            args.masked_inquired_actions,
            args.compute_metrics_flag,
            args.eval_batch_size,
        )
    out_path = f"{args.output}/{args.eval_prefix}{EVAL_FILE}"
    write_json(results, out_path)
    return results

def asd_test(args, params=None):
    """Method for pretraining an agent against a gym environment.

    Parameters
    ----------
    args: dict
        The arguments as provided in the command line.
    params: dict
        The parameters as provided in the configuration file.
        Default: None

    Return
    ------
    None

    """

    if params is None:
        params = {}

    assert args.batch_size > 0

    if not os.path.exists(args.output):
        os.makedirs(args.output)

    if args.seed is None:
        args.seed = 1234

    print(f"Asd testing: Using seed [{args.seed}] for this process.")

    args.metric = "accuracy"

    # initialize the random generator
    initialize_seed(args.seed)

    # create environment
    env = ASDPatientInteractionSimulator(args, patient_filepath=args.data, train=False)
    print(f"Test environment size: {env.sample_size}")

    args.inp_size = env.state_size + 5
    args.symptom_prediction_size = env.symptom_size
    args.action_state_size = 1
    args.patho_size = env.diag_size if args.patho_size is None else args.patho_size

    # import pdb;pdb.set_trace()
    # instantiate the agent
    saved_model_path = f"{args.output}/{BEST_PRETRAIN_MODEL_NAME}"
    agent = create_agent(args)

    args.device = torch.device(
        f"cuda:{args.cuda_idx}"
        if torch.cuda.is_available() and args.cuda_idx is not None
        else "cpu"
    )

    agent = load_model(agent, saved_model_path)
    agent.to(args.device)

    # create the datasets
    ds_valid = SimPaDataset(env, args.interaction_length)

    # infer
    asd_classifier_test(agent, ds_valid, args)