from ch02_math.solutions.ex04_primes import calc_primes_up_to

def calc_prime_factors(value):
    all_primes = calc_primes_up_to(value)

    prime_factors = []
    remaining_value = value

    # solange gerade, immer wieder durch 2 teilen
    while remaining_value % 2 == 0 and remaining_value >= 2:
        remaining_value = remaining_value // 2
        prime_factors.append(2)

    # prüfe Rest auf Primzahl
    if is_prime(all_primes, remaining_value):
        prime_factors.append(remaining_value)
    else:
        # Rest ist keine Primzahl, weitere Prüfung
        while remaining_value > 1:
            for current_prime in all_primes:
                if remaining_value % current_prime == 0:
                    remaining_value = remaining_value // current_prime
                    prime_factors.append(current_prime)

                    # starte nochmal von vorne, weil jeder Teiler
                    # mehrfach vorkommen kann
                    break

    return prime_factors

def is_prime(all_primes, n):
    return n in all_primes

def calc_prime_factors_optimized(value):
    all_primes = calc_primes_up_to(value)

    prime_factors = []

    remaining_value = value
    while remaining_value > 1:
        for current_prime in all_primes:
            while remaining_value % current_prime == 0:
                remaining_value = remaining_value // current_prime
                prime_factors.append(current_prime)

    return prime_factors

def main():
    print(calc_prime_factors(28))
    print(calc_prime_factors(123))
    print(calc_prime_factors(7225))
    print(calc_prime_factors(7227))
    print("---------------")
    print(calc_prime_factors_optimized(28))
    print(calc_prime_factors_optimized(123))
    print(calc_prime_factors_optimized(7225))
    print(calc_prime_factors_optimized(7227))

if __name__ == "__main__":
    main()


#
# @lc app=leetcode.cn id=71 lang=python3
#
# [71] 简化路径
#

# @lc code=start
class Solution:
    def simplifyPath(self, path: str) -> str:
        ps = path.split('/')
        x = []
        for p in ps:
            if p == '' or p == '.':
                continue
            elif p == '..':
                if len(x) > 0:
                    x.pop()
            else:
                x.append(p)

        return '/' + '/'.join(x)
# @lc code=end


import cv2
import os
import argparse
import numpy as np
import torch
import matplotlib.pyplot as plt
import torchvision.transforms as transforms

from lib.multi_depth_model_woauxi import RelDepthModel
from lib.net_tools import load_ckpt

def parse_args():
    parser = argparse.ArgumentParser(
        description='Configs for LeReS')
    parser.add_argument('--load_ckpt', default='./res50.pth', help='Checkpoint path to load')
    parser.add_argument('--backbone', default='resnext101', help='Checkpoint path to load')

    args = parser.parse_args()
    return args

def scale_torch(img):
    """
    Scale the image and output it in torch.tensor.
    :param img: input rgb is in shape [H, W, C], input depth/disp is in shape [H, W]
    :param scale: the scale factor. float
    :return: img. [C, H, W]
    """
    if len(img.shape) == 2:
        img = img[np.newaxis, :, :]
    if img.shape[2] == 3:
        transform = transforms.Compose([transforms.ToTensor(),
		                                transforms.Normalize((0.485, 0.456, 0.406) , (0.229, 0.224, 0.225) )])
        img = transform(img)
    else:
        img = img.astype(np.float32)
        img = torch.from_numpy(img)
    return img

if __name__ == '__main__':

    args = parse_args()

    # create depth model
    depth_model = RelDepthModel(backbone=args.backbone)
    depth_model.eval()

    # load checkpoint
    load_ckpt(args, depth_model, None, None)
    depth_model.cuda()

    image_dir = os.path.dirname(os.path.dirname(__file__)) + '/test_images/'
    imgs_list = os.listdir(image_dir)
    imgs_list.sort()
    imgs_path = [os.path.join(image_dir, i) for i in imgs_list if i != 'outputs']
    image_dir_out = image_dir + '/outputs'
    os.makedirs(image_dir_out, exist_ok=True)

    for i, v in enumerate(imgs_path):
        print('processing (%04d)-th image... %s' % (i, v))
        rgb = cv2.imread(v)
        rgb_c = rgb[:, :, ::-1].copy()
        gt_depth = None
        A_resize = cv2.resize(rgb_c, (448, 448))
        rgb_half = cv2.resize(rgb, (rgb.shape[1]//2, rgb.shape[0]//2), interpolation=cv2.INTER_LINEAR)

        img_torch = scale_torch(A_resize)[None, :, :, :]
        pred_depth = depth_model.inference(img_torch).cpu().numpy().squeeze()
        pred_depth_ori = cv2.resize(pred_depth, (rgb.shape[1], rgb.shape[0]))

        # if GT depth is available, uncomment the following part to recover the metric depth
        #pred_depth_metric = recover_metric_depth(pred_depth_ori, gt_depth)

        img_name = v.split('/')[-1]
        cv2.imwrite(os.path.join(image_dir_out, img_name), rgb)
        # save depth
        plt.imsave(os.path.join(image_dir_out, img_name[:-4]+'-depth.png'), pred_depth_ori, cmap='rainbow')
        cv2.imwrite(os.path.join(image_dir_out, img_name[:-4]+'-depth_raw.png'), (pred_depth_ori/pred_depth_ori.max() * 60000).astype(np.uint16))


import curses
import datetime

from curses import wrapper

def prev_value(x, field):
    limit = 24 if field == 'hh' else 60
    return (x + limit - 1) % limit

def next_value(x, field):
    limit = 24 if field == 'hh' else 60
    return (x + 1) % limit

def render(last_key, stdscr, offsets, now, focus):
    offsety = offsets[0]
    offsetx = offsets[0]

    fmt = '{hh:0>2}   {mm:0>2}   {ss:0>2}'

    lines = [
            fmt.format(
                hh=prev_value(now['hh'], 'hh') if focus == 'hh' else '  ',
                mm=prev_value(now['mm'], 'mm') if focus == 'mm' else '  ',
                ss=prev_value(now['ss'], 'ss') if focus == 'ss' else '  ',
                ),
            fmt.format(
                hh=now['hh'],
                mm=now['mm'],
                ss=now['ss'],
                ),
            fmt.format(
                hh=next_value(now['hh'], 'hh') if focus == 'hh' else '  ',
                mm=next_value(now['mm'], 'mm') if focus == 'mm' else '  ',
                ss=next_value(now['ss'], 'ss') if focus == 'ss' else '  ',
                ),
            ]

    for idx, line in enumerate(lines):
        if idx == 1:
            attr = curses.color_pair(0)
        else:
            attr = curses.color_pair(1) | curses.A_BOLD

        stdscr.addstr(idx + offsety, offsetx + 1, line, attr)

    focus_l, focus_r = None, None
    if focus == 'hh':
        focus_l, focus_r = 0, 3
    elif focus == 'mm':
        focus_l, focus_r = 5, 8
    elif focus == 'ss':
        focus_l, focus_r = 10, 13

    if focus_l is not None and focus_r is not None:
        stdscr.addstr(offsety + 1, offsetx + focus_l, '[')
        stdscr.addstr(offsety + 1, offsetx + focus_r, ']')

    stdscr.addstr(offsety + 1, offsetx + 4, ':')
    stdscr.addstr(offsety + 1, offsetx + 9, ':')

def main(stdscr):
    stdscr.clear()

    curses.curs_set(False)

    curses.init_pair(1, curses.COLOR_BLACK, curses.COLOR_BLACK)

    lines = [
            '┌' + '─' * 14 + '┐',
            '│' + ' ' * 14 + '│',
            '│' + ' ' * 14 + '│',
            '│' + ' ' * 14 + '│',
            '└' + '─' * 14 + '┘',
            ]
    for idx, line in enumerate(lines):
        attr = curses.color_pair(1) | curses.A_BOLD
        stdscr.addstr(idx, 0, line, attr)

    fields = ['hh', 'mm', 'ss']
    now = datetime.datetime.now()
    now = {'hh': now.hour, 'mm': now.minute, 'ss': now.second}

    focus = 0
    key = 'none'
    while True:
        render(key, stdscr, (1, 1), now, fields[focus])
        key = stdscr.getkey()

        if key == 'q':
            break

        elif key in ('h', 'KEY_LEFT'):
            focus = max(0, focus - 1)

        elif key in ('l', 'KEY_RIGHT'):
            focus = min(len(fields) - 1, focus + 1)

        elif key in ('k', 'KEY_UP'):
            now[fields[focus]] = prev_value(now[fields[focus]], fields[focus])

        elif key in ('j', 'KEY_DOWN'):
            now[fields[focus]] = next_value(now[fields[focus]], fields[focus])

        elif key == '0':
            now[fields[focus]] = 0

        elif key == '$':
            now[fields[focus]] = prev_value(0, fields[focus])

    return now['hh'], now['mm'], now['ss']

if __name__ == '__main__':
    print('Selected time: {:0>2}:{:0>2}:{:0>2}'.format(*wrapper(main)))


import os

from cpp2py import Config

from tools import TESTCASES_PATH, cpp2py_tester

def test_external_library():
    """set LD_LIBRARY_PATH before test to find the shared liarary"""
    EXTERNAL_DIR = "anotherincludedir"
    library_dir = os.path.abspath(os.path.join(TESTCASES_PATH, EXTERNAL_DIR))

    cwd = os.getcwd()
    try:
        os.chdir(library_dir)
        os.system("make")
    finally:
        os.chdir(cwd)

    config = Config()
    config.add_library_dir(library_dir)

    @cpp2py_tester("withincludedir.hpp", config=config, incdirs=[EXTERNAL_DIR])
    def run():
        from withincludedir import length

        assert length(3.0, 4.0) == 5.0

    run()

@cpp2py_tester("withincludedir.hpp", incdirs=["anotherincludedir"])
def test_another_include_dir():
    from withincludedir import length

    assert length(3.0, 4.0) == 5.0


#!/usr/bin/env python

def largest_multiple_of_three(ints):
    # Sort numbers in descending order
    ints.sort(reverse=True)

    # Function to calculate the sum of digits
    sum_digits = sum(ints)

    # Function to remove the smallest element with a specified remainder
    def remove_smallest(remainder):
        for i in range(len(ints) - 1, -1, -1):
            if ints[i] % 3 == remainder:
                del ints[i]
                return True
        return False

    # Adjust the list to make the sum of its digits divisible by 3
    if sum_digits % 3 == 1:
        # Try to remove one number with a remainder of 1, else remove two with a remainder of 2
        if not remove_smallest(1):
            remove_smallest(2) and remove_smallest(2)
    elif sum_digits % 3 == 2:
        # Try to remove one number with a remainder of 2, else remove two with a remainder of 1
        if not remove_smallest(2):
            remove_smallest(1) and remove_smallest(1)

    return int("".join(map(str, ints))) if ints and sum(ints) % 3 == 0 else -1

# Test Cases
test1 = largest_multiple_of_three([8, 1, 9]) == 981
test2 = largest_multiple_of_three([8, 6, 7, 1, 0]) == 8760
test3 = largest_multiple_of_three([1]) == -1

test1, test2, test3


# Gabrielle Mehltretter 20065730 April 5th 2019
# This program sorts and searches data from the Toronto Pearson Airport.
# Calculates mean temperature, max temperature, amounts of rain and snow
# giving dates as well.

# This given function creates a list of dictionaries from the file.
# Each dictionary has the data for a certain month of a certain year.

def readData(filename):
    fileIn = open(filename, 'r')
    allData = []
    line = fileIn.readline()
    while line != "":
        line = fileIn.readline().strip()
        if line != "":
            values = line.split(',')
            monthData = {}
            monthData['year'] = int(values[0])
            monthData['month'] = int(values[1])
            monthData['yearmonth'] = int(values[0]) * 100 + int(values[1])
            monthData['meanT'] = float(values[2])
            monthData['maxT'] = float(values[3])
            monthData['minT'] = float(values[4])
            monthData['rain'] = float(values[5])
            monthData['snow'] = float(values[6])
            allData.append(monthData)
    fileIn.close()
    return allData

# This function prints the beginning and end portions of the
# supplied list of dictionaries.

def showSome(allData):
    for i in range(10):
        print(allData[i])
    print("<snip>")
    for i in range(-10, 0):
        print(allData[i])

# This function ensures that given integers fall within
# proper parameters.

def getInt(prompt, lowLimit=None, highLimit=None):
    numberOK = False
    while not numberOK:
        try:
            userNum = int(input(prompt))
            if lowLimit != None and userNum < lowLimit:
                print("The number must be higher than", lowLimit)
                print("Please try again.")
            elif highLimit != None and userNum > highLimit:
                print("The number must be lower than", highLimit)
                print("Please try again.")
            else:
                numberOK = True
        except ValueError:
            print("Your entry is not a valid integer, please try again.")
    return userNum

# This insertion sort function sorts the file in increasing order
# by "yearmonth". Starting with the first index decides if
# the value should be before or after the zero index. Continues
# through every index, switching where needed.

def insertionSort(allData, key):
    for i in range(1, len(allData)):
        temp = allData[i]
        j = i - 1
        while j >= 0 and temp[key] < allData[j][key]:
            allData[j + 1] = allData[j]
            j -= 1
        allData[j + 1] = temp

# This function uses a binary search to find the amount of rain
# from a certain date. Binary search finds the midpoint in the data
# it then cuts the data in half and searches the half with the
# point it is looking for. It continues to cut the data in half
# until it finds the data it is looking for. The function assumes the
# data has been sorted in increasing order by date. The function will
# raise a ValueError exception if the year and month target do not
# exist in the data.

def findRain(allData, target):
    low = 0
    high = len(allData) - 1
    found = False
    while low <= high and not found:
        mid = (low + high) // 2
        if target == allData[mid]["yearmonth"]:
            found = True
        else:
            if target < allData[mid]["yearmonth"]:
                high = mid - 1
            else:
                low = mid + 1
    return allData[mid]["rain"]

    raise ValueError("Target rainfall not found!")

# This function finds the max value in the data for a certain key.

def findMax(allData, key):
    val = 0
    for i in range(1, len(allData)):
        if allData[i][key] > allData[val][key]:
            val = i
    return allData[val]

# This function finds the min value in the data for a certain key.

def findMin(allData, key):
    val = 0
    for i in range(1, len(allData)):
        if allData[i][key] < allData[val][key]:
            val = i
    return allData[val]

# This function returns a list of dictionaries that say the total amount
# of snow for every year in the data.

def getAnnualSnow(allData):
    snowData = []
    i = 0
    while i < len(allData):
        snowSum = 0
        currentYear = allData[i]["year"]
        while i < len(allData) and allData[i]["year"] == currentYear:
            snowSum = snowSum + allData[i]["snow"]
            i = i + 1

        yearOfSnow = {}
        yearOfSnow["year"] = allData[i-1]["year"]
        yearOfSnow["totalsnow"] = snowSum

        snowData.append(yearOfSnow)
    return snowData

# This function calculates the mean temperatures for a year in the data.

def saveAnnualMeanTemp(allData, filename):
    file = open(filename, "w+")
    i = 0
    while i < len(allData):
        tempSum = 0
        currentYear = allData[i]["year"]
        while i < len(allData) and allData[i]["year"] == currentYear:
            tempSum = tempSum + allData[i]["meanT"]
            i = i + 1
        tempSum = tempSum / 12
    file.write(str(tempSum) + "\n")
    file.close()

# The main function was given. It first opens the file and then prints both the
# unsorted and sorted list of dictionaries made from the file. It then prompts
# for a certian month and year and prints the rainfall for that month. It will
# raise a ValueError exception if there is no data for that month. It then will
# print the max and min for temperature as well as the max snowfall and rainfall.
# Also prints the highest, lowest and median snowfall. Finally it will print the
# highest, lowest and median mean temperature for the data.

def main():
    db = readData("TorontoWeatherData.csv")
    unsortedDb = readData("TorontoWeatherData.csv")
    print("Before sorting, as read from file:")
    showSome(db)
    insertionSort(db, 'yearmonth')
    print("\nAfter sorting by date:")
    showSome(db)

    searchYear = getInt("Enter year for rainfall search: ", 1938, 2018)
    searchMonth = getInt("Enter month for rainfall search: ", 1, 12)
    searchYearMonth = 100 * searchYear + searchMonth
    try:
        rainfall = findRain(db, searchYearMonth)
        print(f"Rainfall was {rainfall} mm.")
    except ValueError as message:
        print(message)

    maxR = findMax(db, 'maxT')
    print(f"\nHighest temperature {maxR['maxT']} deg C, in month {maxR['month']}, {maxR['year']}.")
    minR = findMin(db, 'minT')
    print(f"Lowest temperature {minR['minT']} deg C, in month {minR['month']}, {minR['year']}.")
    maxR = findMax(db, 'rain')
    print(f"Highest rainfall {maxR['rain']} mm, in month {maxR['month']}, {maxR['year']}.")
    maxR = findMax(db, 'snow')
    print(f"Highest snowfall {maxR['snow']} cm, in month {maxR['month']}, {maxR['year']}.")

    annualSnow = getAnnualSnow(db)
    insertionSort(annualSnow, 'totalsnow')
    minR = annualSnow[0]
    print(f"\nLowest annual snowfall {minR['totalsnow']} cm, in {minR['year']}.")
    medR = annualSnow[len(annualSnow) // 2]
    print(f"Median annual snowfall {medR['totalsnow']} cm.")
    maxR = annualSnow[len(annualSnow) - 1]
    print(f"Highest annual snowfall {maxR['totalsnow']} cm, in {maxR['year']}.")

    insertionSort(db, 'meanT')
    minR = db[0]
    print(f"\nLowest mean temperature {minR['meanT']} deg C, in month {minR['month']}, {minR['year']}.")
    medR = db[len(db) // 2]
    print(f"Median mean temperature {medR['meanT']} deg C.")
    maxR = db[-1]
    print(f"Highest mean temperature {maxR['meanT']} deg C, in month {maxR['month']}, {maxR['year']}.")

    saveAnnualMeanTemp(db, "YearMeans.txt")

main()


'''TkOS New 1.0源代码
作者：bilibili -小唐玩电脑-
相关视频只在B站发布。
本程序使用MIT许可证。'''

import tkinter
from tkinter import messagebox
from tkinter import ttk
from tkinter import filedialog
import os
import sys
import platform
import webbrowser

def shutdown():
    que = messagebox.askokcancel("提示",'你确定要关闭TkOS吗？')
    if que == True:
        sys.exit()

def shutdown_windows():
    que = messagebox.askokcancel("提示","你确信要关闭系统吗？\n注意这个真的会让Windows关机！")
    if que == True:
        os.system("shutdown -s -t 10 -c 因为TkOS用户要关机，所以将在10秒后关机。")

def cantclose():
    messagebox.showerror("系统","请在系统菜单关闭此程序。")

def go_to_uploader():
    webbrowser.open("https://space.bilibili.com/1473734923")

def cmd():
    def sent_command():
        get_command = input_command.get()
        if get_command == "version":
            messagebox.showinfo(message = "TkOS New 1.0\n开发版本\n不是最终效果，请等待正式版\n0.0.114514\n本程序已使用MIT证书。\nsystem:{}".format(platform.system()))
        elif get_command == "sysinfo":
            messagebox.showinfo(message="{}".format(platform.system()))
        elif get_command == 'goto windows cmd.exe':
            print("TkOS:Open app\"cmd.exe\"is succcessfully.\nIf you want to exit the cmd.exe,please insert the\"exit\"command to exit the cmd.exe")
            os.system('cmd')
        elif get_command == 'run textedit':
            print("open textedit")
            text_edit()
        else:
            print("sorry, your command is not a vaild TkOS command.")
        input_command.delete(0,tkinter.END)

    command = tkinter.Toplevel()
    command.title("终端")
    tkinter.Label(command, text = '请输入终端命令：').grid(row = 0, column=0)
    input_command = tkinter.Entry(command)
    input_command.grid(row = 0, column = 1)
    ttk.Button(command, text = "OK", command = sent_command).grid(row = 0,column=2)

def lltang():
    def ok():
        get_answer = your_input.get()
        who_are_you = ['你叫什么名字？', '你是谁？']
        fucks = ['卧槽','他妈的', '屮', '艹','傻逼', '草泥马', '操你妈', '我日你仙人',\
           '我日你先人', 'Fuck', 'fuck','Shit','shit','bitch', 'Bitch']
        one = who_are_you in get_answer
        two = fucks in get_answer
        if one:
            messagebox.showinfo(message = "我是小小唐。")
        elif two:
            messagebox.showinfo(message = '说脏话是一种不对的行为。')
        else:
            messagebox.showinfo("NiShuoSha?",'你说啥？')

    '''我是小小唐，文字助手'''
    lltang = tkinter.Toplevel()
    lltang.title("小小唐")
    tkinter.Label(lltang, text = "哈喽啊，你想说些什么？").grid(row = 0, column = 0)
    your_input = tkinter.Entry(lltang)
    your_input.grid(row = 0, column = 1)
    ttk.Button(lltang, text = '提交', command = ok).gri(row = 0, column = 2)

def update_log():
    log = tkinter.Toplevel()
    log.title("版本更新日志")
    tkinter.Label(log, text = '当前版本：开发版本3\n内部版本0.0.114514.2').pack()
    update = tkinter.Text(log)
    update.pack()
    update_text = '版本更新日志（按版本编号从大到小排列）\n\
0.0.114514.3:\n\
完善部分UI。\n\
完善文本编辑。\n\
新增程序打开器。\n\
0.0.114514.2：\n\
更新了更新日志的查看。\n\
新增文本编辑。\n\
0.0.114514.1：\n\
新增终端。\n\
新增文字助手小小唐（目前有bug ）。\n\
更新关于页面。\n\
0.0.114514：\n\
无更新。'
    update.insert(tkinter.END, update_text)

def text_edit():
    def open():
        file_open = filedialog.askopenfile(title = '打开文件', parent = windows, filetypes = [("文本文档", '.txt', \
            '所有文档', '.*')])
        with open(file_open, 'r') as file:
            text.insert(tkinter.END, file)
            file.close()
    def save():
        file_text = text.get(1.0, tkinter.END)
        file_name = filedialog.asksaveasfile(title = '保存/另存为文件...', parent = windows, filetypes = \
            [("文本文档", '.txt'), ("所有文件", '.*')])
        if  not file_name:
            return
        with open(file_name, 'w') as file:
            file.write(file_text)
            file.close()

    def new():
        if text.get(1.0, tkinter.END) == '':
            text.delete(1.0, tkinter.END)
        else:
            want = messagebox.askquestion(message="是否要保存当前的文件？")
            if want == True:
                save()
    windows = tkinter.Toplevel()
    windows.title('TkOS文本编辑器')
    text = tkinter.Text(windows)
    text.pack()
    text_menu = tkinter.Menu(windows)
    file_menu = tkinter.Menu(text_menu, tearoff = 0)
    help_menu = tkinter.Menu(text_menu, tearoff = 0)
    text_menu.add_cascade(label = '文件(F)', menu = file_menu)
    file_menu.add_command(label = '新建', command = new)
    file_menu.add_command(label = '保存', command = save)
    file_menu.add_command(label = '另存为')
    file_menu.add_command(label = '打开')
    file_menu.add_separator()
    help_menu.add_command(label = '关于')
    windows.config(menu = text_menu)

def open_app():
    def ok():
        app_will_open = var.get()
        if app_will_open == "文本编辑":
            text_edit()
        elif app_will_open == "终端":
            cmd()

    app = tkinter.Toplevel()
    app.title("打开应用程序")
    tkinter.Label(app, text = '请选择应用程序：').pack()
    var = tkinter.StringVar()
    app_select = ttk.Combobox(app, textvariable=var)
    app_select['value'] = ('文本编辑', '终端')
    app_select.pack()
    ttk.Button(app ,text = '确定', command = ok).pack()

def about():
    about_windows = tkinter.Toplevel()
    about_windows.title("关于TkOS")
    tkinter.Label(about_windows, text = 'TkOS New Version 1.0\n版本0.0.114514.1（开发版本2）\n制作者bilibili -小唐玩电脑-\n本程序使用MIT许可证。').pack()
    ttk.Button(about_windows, text = '官方网站', command = webbrowser.open("https://zqtang10/github.io/TkOS_New_1.0")).pack()

root = tkinter.Tk()
root.title('主界面')
root.protocol("WM_DELETE_WINDOW", cantclose)
menu = tkinter.Menu()
system = tkinter.Menu(menu, tearoff = 0)
menu.add_cascade(label = '系统(S)', menu = system)
system.add_command(label = '关机系统', command = shutdown)
system.add_command(label = "关机（注意这个真的会关机）", command = shutdown_windows)
system.add_separator()
system.add_command(label = '终端', command = cmd)
system.add_command(label = '打开应用程序', command = open_app)
system.add_separator()
system.add_command(label = '更新日志', command = update_log)
system.add_command(label = '关于',command = about)
root.config(menu = menu)
tkinter.Label(root, text = 'TkOS New 1.0 开发版\n制作者：bilibili up主-小唐玩电脑-\n三连再看，养成习惯').grid(row = 0,column=0)
ttk.Button(root, text = '去看看up主首页', command = go_to_uploader).grid(row = 1,column=0)
ttk.Button(root, text = '我是小小唐', command = lltang).grid(row = 1, column = 1)
root.mainloop()


import cv2

import raytracing
import open3d
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from dataset.database import parse_database_name, get_database_split, BaseDatabase
from network.field import SDFNetwork, SingleVarianceNetwork, NeRFNetwork, AppShadingNetwork, get_intersection, \
    extract_geometry, sample_pdf, MCShadingNetwork
from utils.base_utils import color_map_forward, downsample_gaussian_blur
from utils.raw_utils import linear_to_srgb

def build_imgs_info(database:BaseDatabase, img_ids):
    images = [database.get_image(img_id) for img_id in img_ids]
    poses = [database.get_pose(img_id) for img_id in img_ids]
    Ks = [database.get_K(img_id) for img_id in img_ids]

    images = np.stack(images, 0)
    images = color_map_forward(images).astype(np.float32)
    Ks = np.stack(Ks, 0).astype(np.float32)
    poses = np.stack(poses, 0).astype(np.float32)
    return {'imgs': images, 'Ks': Ks, 'poses': poses}

def imgs_info_to_torch(imgs_info, device='cpu'):
    for k, v in imgs_info.items():
        v = torch.from_numpy(v)
        if k.startswith('imgs'): v = v.permute(0,3,1,2)
        imgs_info[k] = v.to(device)
    return imgs_info

def imgs_info_slice(imgs_info, idxs):
    new_imgs_info={}
    for k, v in imgs_info.items():
        new_imgs_info[k]=v[idxs]
    return new_imgs_info

def imgs_info_to_cuda(imgs_info):
    for k, v in imgs_info.items():
        imgs_info[k]=v.cuda()
    return imgs_info

def imgs_info_downsample(imgs_info, ratio):
    b, _, h, w = imgs_info['imgs'].shape
    dh, dw = int(ratio*h), int(ratio*w)
    imgs_info_copy = {k:v for k,v in imgs_info.items()}
    imgs_info_copy['imgs'], imgs_info_copy['Ks'] = [], []
    for bi in range(b):
        img = imgs_info['imgs'][bi].cpu().numpy().transpose([1,2,0])
        img = downsample_gaussian_blur(img, ratio)
        img = cv2.resize(img, (dw,dh), interpolation=cv2.INTER_LINEAR)
        imgs_info_copy['imgs'].append(torch.from_numpy(img).permute(2,0,1))
        K = torch.from_numpy(np.diag([dw / w, dh / h, 1]).astype(np.float32)) @ imgs_info['Ks'][bi]
        imgs_info_copy['Ks'].append(K)

    imgs_info_copy['imgs'] = torch.stack(imgs_info_copy['imgs'], 0)
    imgs_info_copy['Ks'] = torch.stack(imgs_info_copy['Ks'], 0)
    return imgs_info_copy

class NeROShapeRenderer(nn.Module):
    default_cfg = {
        # standard deviation for opacity density
        'std_net': 'default',
        'std_act': 'exp',
        'inv_s_init': 0.3,
        'freeze_inv_s_step': None,

        # geometry network
        'sdf_net': 'default',
        'sdf_activation': 'none',
        'sdf_bias': 0.5,
        'sdf_n_layers': 8,
        'sdf_freq': 6,
        'sdf_d_out': 257,
        'geometry_init': True,

        # shader network
        'shader_config': {},

        # sampling strategy
        'n_samples': 64,
        'n_bg_samples': 32,
        'inf_far': 1000.0,
        'n_importance': 64,
        'up_sample_steps': 4,  # 1 for simple coarse-to-fine sampling
        'perturb': 1.0,
        'anneal_end': 50000,
        'train_ray_num': 512,
        'test_ray_num': 1024,
        'clip_sample_variance': True,

        # dataset
        'database_name': 'nerf_synthetic/lego/black_800',

        # validation
        'test_downsample_ratio': True,
        'downsample_ratio': 0.25,
        'val_geometry': False,

        # losses
        'rgb_loss': 'charbonier',
        'apply_occ_loss': True,
        'occ_loss_step': 20000,
        'occ_loss_max_pn': 2048,
        'occ_sdf_thresh': 0.01,

        "fixed_camera": False,
    }

    def __init__(self, cfg, training=True):
        super().__init__()
        self.cfg = {**self.default_cfg, **cfg}

        self.sdf_network = SDFNetwork(d_out=self.cfg['sdf_d_out'], d_in=3, d_hidden=256,
                                      n_layers=self.cfg['sdf_n_layers'],
                                      skip_in=[self.cfg['sdf_n_layers'] // 2], multires=self.cfg['sdf_freq'],
                                      bias=self.cfg['sdf_bias'], scale=1.0,
                                      geometric_init=self.cfg['geometry_init'],
                                      weight_norm=True, sdf_activation=self.cfg['sdf_activation'])

        self.deviation_network = SingleVarianceNetwork(init_val=self.cfg['inv_s_init'], activation=self.cfg['std_act'])

        # background nerf is a nerf++ model (this is outside the unit bounding sphere, so we call it outer nerf)
        self.outer_nerf = NeRFNetwork(D=8, d_in=4, d_in_view=3, W=256, multires=10, multires_view=4, output_ch=4, skips=[4], use_viewdirs=True)
        nn.init.constant_(self.outer_nerf.rgb_linear.bias, np.log(0.5))

        self.color_network = AppShadingNetwork(self.cfg['shader_config'])
        self.sdf_inter_fun = lambda x: self.sdf_network.sdf(x)

        if training:
            self._init_dataset()

    def _init_dataset(self):
        # train/test split
        self.database = parse_database_name(self.cfg['database_name'])
        self.train_ids, self.test_ids = get_database_split(self.database)
        self.train_ids = np.asarray(self.train_ids)

        self.train_imgs_info = build_imgs_info(self.database, self.train_ids)
        self.train_imgs_info = imgs_info_to_torch(self.train_imgs_info, 'cpu')
        b, _, h, w = self.train_imgs_info['imgs'].shape
        print(f'training size {h} {w} ...')
        self.train_num = len(self.train_ids)

        self.test_imgs_info = build_imgs_info(self.database, self.test_ids)
        self.test_imgs_info = imgs_info_to_torch(self.test_imgs_info, 'cpu')
        self.test_num = len(self.test_ids)

        # clean the data if we already have
        if hasattr(self, 'train_batch'):
            del self.train_batch

        self.train_batch, self.train_poses, self.tbn, _, _ = self._construct_ray_batch(self.train_imgs_info)
        self.train_poses = self.train_poses.float().cuda()

        self._shuffle_train_batch()

    def _shuffle_train_batch(self):
        self.train_batch_i = 0
        shuffle_idxs = torch.randperm(self.tbn, device='cpu')  # shuffle
        for k, v in self.train_batch.items():
            self.train_batch[k] = v[shuffle_idxs]

    def _construct_ray_batch(self, imgs_info, device='cpu'):
        imn, _, h, w = imgs_info['imgs'].shape
        coords = torch.stack(torch.meshgrid(torch.arange(h), torch.arange(w)), -1)[:, :, (1, 0)]  # h,w,2
        coords = coords.to(device)
        coords = coords.float()[None, :, :, :].repeat(imn, 1, 1, 1)  # imn,h,w,2
        coords = coords.reshape(imn, h * w, 2)
        coords = torch.cat([coords + 0.5, torch.ones(imn, h * w, 1, dtype=torch.float32, device=device)], 2)  # imn,h*w,3

        # imn,h*w,3 @ imn,3,3 => imn,h*w,3
        dirs = coords @ torch.inverse(imgs_info['Ks']).permute(0, 2, 1)
        imgs = imgs_info['imgs'].permute(0, 2, 3, 1).reshape(imn, h * w, 3)  # imn,h*w,3
        idxs = torch.arange(imn, dtype=torch.int64, device=device)[:, None, None].repeat(1, h * w, 1)  # imn,h*w,1
        poses = imgs_info['poses']  # imn,3,4

        rn = imn * h * w
        ray_batch = {
            'dirs': dirs.float().reshape(rn, 3).to(device),
            'rgbs': imgs.float().reshape(rn, 3).to(device),
            'idxs': idxs.long().reshape(rn, 1).to(device),
        }
        return ray_batch, poses, rn, h, w

    def nvs(self, pose, K, h, w):
        device = 'cuda'
        K = torch.from_numpy(K.astype(np.float32)).unsqueeze(0).to(device)
        pose = torch.from_numpy(pose.astype(np.float32)).unsqueeze(0).to(device)

        coords = torch.stack(torch.meshgrid(torch.arange(h), torch.arange(w)), -1)[:, :, (1, 0)]  # h,w,2
        coords = coords.to(device)
        coords = coords.float()[None, :, :, :].repeat(1, 1, 1, 1)  # 1,h,w,2
        coords = coords.reshape(1, h * w, 2)
        coords = torch.cat([coords + 0.5, torch.ones(1, h * w, 1, dtype=torch.float32, device=device)], 2)  # 1,h*w,3

        # 1,h*w,3 @ imn,3,3 => 1,h*w,3
        dirs = coords @ torch.inverse(K).permute(0, 2, 1)
        idxs = torch.arange(1, dtype=torch.int64, device=device)[:, None, None].repeat(1, h * w, 1)  # 1,h*w,1

        rn = h * w
        ray_batch = {
            'dirs': dirs.float().reshape(rn, 3).to(device),
            'idxs': idxs.long().reshape(rn, 1).to(device),
        }

        trn = 1024
        output_color = []
        for ri in range(0, rn, trn):
            cur_ray_batch = {}
            for k, v in ray_batch.items(): cur_ray_batch[k] = v[ri:ri + trn]

            with torch.no_grad():
                rays_o, rays_d, near, far, human_poses = self._process_ray_batch(cur_ray_batch, pose)
                cur_outputs = self.render(rays_o, rays_d, near, far, human_poses, 0, 0, is_train=False, step=300000)
                output_color.append(cur_outputs['ray_rgb'].detach().cpu().numpy())

        output_color = np.reshape(np.concatenate(output_color, 0), [h, w, 3])
        return output_color

    def get_anneal_val(self, step):
        if self.cfg['anneal_end'] < 0:
            return 1.0
        else:
            return np.min([1.0, step / self.cfg['anneal_end']])

    @staticmethod
    def near_far_from_sphere(rays_o, rays_d):
        a = torch.sum(rays_d ** 2, dim=-1, keepdim=True)
        b = 2.0 * torch.sum(rays_o * rays_d, dim=-1, keepdim=True)
        mid = 0.5 * (-b) / a
        near = mid - 1.0
        far = mid + 1.0
        near = torch.clamp(near, min=1e-3)
        return near, far

    def get_human_coordinate_poses(self, poses):
        pn = poses.shape[0]
        cam_cen = (-poses[:, :, :3].permute(0, 2, 1) @ poses[:, :, 3:])[..., 0]  # pn,3
        if self.cfg['fixed_camera']:
            pass
        else:
            cam_cen[..., 2] = 0

        Y = torch.zeros([1, 3]).expand(pn, 3)
        Y[:, 2] = -1.0
        Z = torch.clone(poses[:, 2, :3])  # pn, 3
        Z[:, 2] = 0
        Z = F.normalize(Z, dim=-1)
        X = torch.cross(Y, Z)  # pn, 3
        R = torch.stack([X, Y, Z], 1)  # pn,3,3
        t = -R @ cam_cen[:, :, None]  # pn,3,1
        return torch.cat([R, t], -1)

    def _process_ray_batch(self, ray_batch, poses):
        rays_d = ray_batch['dirs']  # rn,3
        idxs = ray_batch['idxs'][..., 0]  # rn

        rays_o = poses[:, :, :3].permute(0, 2, 1) @ -poses[:, :, 3:]  # trn,3,1
        rays_o = rays_o[idxs, :, 0]  # rn,3
        rays_d = poses[idxs, :, :3].permute(0, 2, 1) @ rays_d.unsqueeze(-1)
        rays_d = rays_d[..., 0]  # rn,3

        rays_o = rays_o
        rays_d = F.normalize(rays_d, dim=-1)
        near, far = self.near_far_from_sphere(rays_o, rays_d)

        human_poses = self.get_human_coordinate_poses(poses)
        return rays_o, rays_d, near, far, human_poses[idxs]  # rn, 3, 4

    def test_step(self, index, step,):
        target_imgs_info, target_img_ids = self.test_imgs_info, self.test_ids
        imgs_info = imgs_info_slice(target_imgs_info, torch.from_numpy(np.asarray([index], np.int64)))
        gt_depth, gt_mask = self.database.get_depth(target_img_ids[index])  # used in evaluation
        if self.cfg['test_downsample_ratio']:
            imgs_info = imgs_info_downsample(imgs_info, self.cfg['downsample_ratio'])
            h, w = gt_depth.shape
            dh, dw = int(self.cfg['downsample_ratio'] * h), int(self.cfg['downsample_ratio'] * w)
            gt_depth, gt_mask = cv2.resize(gt_depth, (dw, dh), interpolation=cv2.INTER_NEAREST), \
                cv2.resize(gt_mask.astype(np.uint8), (dw, dh), interpolation=cv2.INTER_NEAREST)
        gt_depth, gt_mask = torch.from_numpy(gt_depth), torch.from_numpy(gt_mask.astype(np.int32))
        ray_batch, input_poses, rn, h, w = self._construct_ray_batch(imgs_info)

        input_poses = input_poses.float().cuda()
        for k, v in ray_batch.items(): ray_batch[k] = v.cuda()

        trn = self.cfg['test_ray_num']
        outputs_keys = ['ray_rgb', 'gradient_error', 'normal', 'depth']
        outputs_keys += [
            'diffuse_albedo', 'diffuse_light', 'diffuse_color',
            'specular_albedo', 'specular_light', 'specular_color', 'specular_ref',
            'metallic', 'roughness', 'occ_prob', 'indirect_light', 'occ_prob_gt',
        ]
        if self.color_network.cfg['human_light']:
            outputs_keys += ['human_light']

        outputs = {k: [] for k in outputs_keys}
        for ri in range(0, rn, trn):
            cur_ray_batch= {k:v[ri:ri+trn] for k, v in ray_batch.items()}
            rays_o, rays_d, near, far, human_poses = self._process_ray_batch(cur_ray_batch, input_poses)
            cur_outputs = self.render(rays_o, rays_d, near, far, human_poses, 0, 0, is_train=False, step=step)
            for k in outputs_keys: outputs[k].append(cur_outputs[k].detach())

        for k in outputs_keys: outputs[k] = torch.cat(outputs[k], 0)
        outputs['loss_rgb'] = self.compute_rgb_loss(outputs['ray_rgb'], ray_batch['rgbs'])
        outputs['gt_rgb'] = ray_batch['rgbs'].reshape(h, w, 3)
        outputs['ray_rgb'] = outputs['ray_rgb'].reshape(h, w, 3)

        # used in evaluation
        outputs['gt_depth'] = gt_depth.unsqueeze(-1)
        outputs['gt_mask'] = gt_mask.unsqueeze(-1)

        self.zero_grad()
        return outputs

    def train_step(self, step):
        rn = self.cfg['train_ray_num']
        # fetch to gpu
        train_ray_batch = {k: v[self.train_batch_i:self.train_batch_i + rn].cuda() for k, v in self.train_batch.items()}
        self.train_batch_i += rn
        if self.train_batch_i + rn >= self.tbn: self._shuffle_train_batch()
        train_poses = self.train_poses.cuda()
        rays_o, rays_d, near, far, human_poses = self._process_ray_batch(train_ray_batch, train_poses)

        outputs = self.render(rays_o, rays_d, near, far, human_poses, -1, self.get_anneal_val(step), is_train=True, step=step)
        outputs['loss_rgb'] = self.compute_rgb_loss(outputs['ray_rgb'], train_ray_batch['rgbs'])  # ray_loss
        return outputs

    def compute_rgb_loss(self, rgb_pr, rgb_gt):
        if self.cfg['rgb_loss'] == 'l2':
            rgb_loss = torch.sum((rgb_pr - rgb_gt) ** 2, -1)
        elif self.cfg['rgb_loss'] == 'l1':
            rgb_loss = torch.sum(F.l1_loss(rgb_pr, rgb_gt, reduction='none'), -1)
        elif self.cfg['rgb_loss'] == 'smooth_l1':
            rgb_loss = torch.sum(F.smooth_l1_loss(rgb_pr, rgb_gt, reduction='none', beta=0.25), -1)
        elif self.cfg['rgb_loss'] == 'charbonier':
            epsilon = 0.001
            rgb_loss = torch.sqrt(torch.sum((rgb_gt - rgb_pr) ** 2, dim=-1) + epsilon)
        else:
            raise NotImplementedError
        return rgb_loss

    def density_activation(self, density, dists):
        return 1.0 - torch.exp(-F.softplus(density) * dists)

    def compute_density(self, points):
        points_norm = torch.norm(points, dim=-1, keepdim=True)
        points_norm = torch.clamp(points_norm, min=1e-3)
        sigma = self.outer_nerf.density(torch.cat([points / points_norm, 1.0 / points_norm], -1))[..., 0]
        return sigma

    @staticmethod
    def upsample(rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

        dist = (next_z_vals - prev_z_vals)
        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
        return z_samples

    def cat_z_vals(self, rays_o, rays_d, z_vals, new_z_vals, sdf, last=False):
        batch_size, n_samples = z_vals.shape
        _, n_importance = new_z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * new_z_vals[..., :, None]
        z_vals = torch.cat([z_vals, new_z_vals], dim=-1)
        z_vals, index = torch.sort(z_vals, dim=-1)

        if not last:
            new_sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
            sdf = torch.cat([sdf, new_sdf], dim=-1)
            xx = torch.arange(batch_size)[:, None].expand(batch_size, n_samples + n_importance).reshape(-1)
            index = index.reshape(-1)
            sdf = sdf[(xx, index)].reshape(batch_size, n_samples + n_importance)

        return z_vals, sdf

    def sample_ray(self, rays_o, rays_d, near, far, perturb):
        n_samples = self.cfg['n_samples']
        n_bg_samples = self.cfg['n_bg_samples']
        n_importance = self.cfg['n_importance']
        up_sample_steps = self.cfg['up_sample_steps']

        # sample points
        batch_size = len(rays_o)
        z_vals = torch.linspace(0.0, 1.0, n_samples)  # sn
        z_vals = near + (far - near) * z_vals[None, :]  # rn,sn
        z_vals_outside = torch.linspace(1e-3, 1.0 - 1.0 / (n_bg_samples + 1.0), n_bg_samples)

        if perturb > 0:
            t_rand = (torch.rand([batch_size, 1]) - 0.5)
            z_vals = z_vals + t_rand * 2.0 / n_samples

            mids = .5 * (z_vals_outside[..., 1:] + z_vals_outside[..., :-1])
            upper = torch.cat([mids, z_vals_outside[..., -1:]], -1)
            lower = torch.cat([z_vals_outside[..., :1], mids], -1)
            t_rand = torch.rand([batch_size, z_vals_outside.shape[-1]])
            z_vals_outside = lower[None, :] + (upper - lower)[None, :] * t_rand

        z_vals_outside = far / torch.flip(z_vals_outside, dims=[-1]) + 1.0 / n_bg_samples

        # Up sample
        with torch.no_grad():
            pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
            sdf = self.sdf_network.sdf(pts).reshape(batch_size, n_samples)

            for i in range(up_sample_steps):
                rn, sn = z_vals.shape
                if self.cfg['clip_sample_variance']:
                    inv_s = self.deviation_network(torch.empty([1, 3])).expand(rn, sn - 1)
                    inv_s = torch.clamp(inv_s, max=64 * 2 ** i)  # prevent too large inv_s
                else:
                    inv_s = torch.ones(rn, sn - 1) * 64 * 2 ** i
                new_z_vals = self.upsample(rays_o, rays_d, z_vals, sdf, n_importance // up_sample_steps, inv_s)
                z_vals, sdf = self.cat_z_vals(rays_o, rays_d, z_vals, new_z_vals, sdf, last=(i + 1 == up_sample_steps))

        z_vals = torch.cat([z_vals, z_vals_outside], -1)
        return z_vals

    def render(self, rays_o, rays_d, near, far, human_poses, perturb_overwrite=-1, cos_anneal_ratio=0.0, is_train=True, step=None):
        """
        :param rays_o: rn,3
        :param rays_d: rn,3
        :param near:   rn,1
        :param far:    rn,1
        :param human_poses:     rn,3,4
        :param perturb_overwrite: set 0 for inference
        :param cos_anneal_ratio:
        :param is_train:
        :param step:
        :return:
        """
        perturb = self.cfg['perturb']
        if perturb_overwrite >= 0:
            perturb = perturb_overwrite
        z_vals = self.sample_ray(rays_o, rays_d, near, far, perturb)
        ret = self.render_core(rays_o, rays_d, z_vals, human_poses, cos_anneal_ratio=cos_anneal_ratio, step=step, is_train=is_train)
        return ret

    def compute_validation_info(self, z_vals, rays_o, rays_d, weights, human_poses, step):
        depth = torch.sum(weights * z_vals, -1, keepdim=True)  # rn,
        points = depth * rays_d + rays_o  # rn,3
        gradients = self.sdf_network.gradient(points)  # rn,3
        inner_mask = torch.norm(points, dim=-1, keepdim=True) <= 1.0
        outputs = {
            'depth': depth,  # rn,1
            'normal': ((F.normalize(gradients, dim=-1) + 1.0) * 0.5) * inner_mask,  # rn,3
        }

        feature_vector = self.sdf_network(points)[..., 1:]  # rn,f
        _, occ_info, inter_results = self.color_network(points, gradients, -F.normalize(rays_d, dim=-1), feature_vector, human_poses, inter_results=True, step=step)
        _, occ_prob, _ = get_intersection(self.sdf_inter_fun, self.deviation_network, points, occ_info['reflective'], sn0=128, sn1=9)  # pn,sn-1
        occ_prob_gt = torch.sum(occ_prob, dim=-1, keepdim=True)
        outputs['occ_prob_gt'] = occ_prob_gt
        for k, v in inter_results.items(): inter_results[k] = v * inner_mask
        outputs.update(inter_results)
        return outputs

    def compute_sdf_alpha(self, points, dists, dirs, cos_anneal_ratio, step):
        # points [...,3] dists [...] dirs[...,3]
        sdf_nn_output = self.sdf_network(points)
        sdf = sdf_nn_output[..., 0]
        feature_vector = sdf_nn_output[..., 1:]

        gradients = self.sdf_network.gradient(points)  # ...,3
        inv_s = self.deviation_network(points).clip(1e-6, 1e6)  # ...,1
        inv_s = inv_s[..., 0]

        if self.cfg['freeze_inv_s_step'] is not None and step < self.cfg['freeze_inv_s_step']:
            inv_s = inv_s.detach()

        true_cos = (dirs * gradients).sum(-1)  # [...]
        iter_cos = -(F.relu(-true_cos * 0.5 + 0.5) * (1.0 - cos_anneal_ratio) +
                     F.relu(-true_cos) * cos_anneal_ratio)  # always non-positive

        # Estimate signed distances at section points
        estimated_next_sdf = sdf + iter_cos * dists * 0.5
        estimated_prev_sdf = sdf - iter_cos * dists * 0.5

        prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)
        next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)

        p = prev_cdf - next_cdf
        c = prev_cdf

        alpha = ((p + 1e-5) / (c + 1e-5)).clip(0.0, 1.0)  # [...]
        return alpha, gradients, feature_vector, inv_s, sdf

    def compute_density_alpha(self, points, dists, dirs, nerf):
        norm = torch.norm(points, dim=-1, keepdim=True)
        points = torch.cat([points / norm, 1.0 / norm], -1)
        density, color = nerf(points, dirs)  # [...,1] [...,3]
        alpha = self.density_activation(density[..., 0], dists)
        color = linear_to_srgb(torch.exp(torch.clamp(color, max=5.0)))
        return alpha, color

    def compute_occ_loss(self, occ_info, points, sdf, gradients, dirs, step):
        if step < self.cfg['occ_loss_step']: return torch.zeros(1)

        occ_prob = occ_info['occ_prob']
        reflective = occ_info['reflective']

        # select a subset for occ loss
        # note we only apply occ loss on the surface
        inner_mask = torch.norm(points, dim=-1) < 0.999  # pn
        sdf_mask = torch.abs(sdf) < self.cfg['occ_sdf_thresh']
        normal_mask = torch.sum(gradients * dirs, -1) < 0  # pn
        mask = (inner_mask & normal_mask & sdf_mask)

        if torch.sum(mask) > self.cfg['occ_loss_max_pn']:
            indices = torch.nonzero(mask)[:, 0]  # npn
            idx = torch.randperm(indices.shape[0], device='cuda')  # npn
            indices = indices[idx[:self.cfg['occ_loss_max_pn']]]  # max_pn
            mask_new = torch.zeros_like(mask)
            mask_new[indices] = 1
            mask = mask_new

        if torch.sum(mask) > 0:
            inter_dist, inter_prob, inter_sdf = get_intersection(self.sdf_inter_fun, self.deviation_network, points[mask], reflective[mask], sn0=64, sn1=16)  # pn,sn-1
            occ_prob_gt = torch.sum(inter_prob, -1, keepdim=True)
            return F.l1_loss(occ_prob[mask], occ_prob_gt)
        else:
            return torch.zeros(1)

    def render_core(self, rays_o, rays_d, z_vals, human_poses, cos_anneal_ratio=0.0, step=None, is_train=True):
        batch_size, n_samples = z_vals.shape

        # section length in original space
        dists = z_vals[..., 1:] - z_vals[..., :-1]  # rn,sn-1
        dists = torch.cat([dists, dists[..., -1:]], -1)  # rn,sn
        mid_z_vals = z_vals + dists * 0.5

        points = rays_o.unsqueeze(-2) + rays_d.unsqueeze(-2) * mid_z_vals.unsqueeze(-1)
        inner_mask = torch.norm(points, dim=-1) <= 1.0
        outer_mask = ~inner_mask

        dirs = rays_d.unsqueeze(-2).expand(batch_size, n_samples, 3)
        human_poses_pt = human_poses.unsqueeze(-3).expand(batch_size, n_samples, 3, 4)
        dirs = F.normalize(dirs, dim=-1)
        alpha, sampled_color = torch.zeros(batch_size, n_samples), torch.zeros(batch_size, n_samples, 3)

        if torch.sum(outer_mask) > 0:
            alpha[outer_mask], sampled_color[outer_mask] = self.compute_density_alpha(points[outer_mask], dists[outer_mask], -dirs[outer_mask], self.outer_nerf)

        if torch.sum(inner_mask) > 0:
            alpha[inner_mask], gradients, feature_vector, inv_s, sdf = self.compute_sdf_alpha(points[inner_mask], dists[inner_mask], dirs[inner_mask], cos_anneal_ratio, step)
            sampled_color[inner_mask], occ_info = self.color_network(points[inner_mask], gradients, -dirs[inner_mask], feature_vector, human_poses_pt[inner_mask], step=step)
            # Eikonal loss
            gradient_error = (torch.linalg.norm(gradients, ord=2, dim=-1) - 1.0) ** 2
        else:
            gradient_error = torch.zeros(1)

        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[..., :-1]  # rn,sn
        color = (sampled_color * weights[..., None]).sum(dim=1)

        outputs = {
            'ray_rgb': color,  # rn,3
            'gradient_error': gradient_error,  # rn
        }

        if torch.sum(inner_mask) > 0:
            outputs['std'] = torch.mean(1 / inv_s)
        else:
            outputs['std'] = torch.zeros(1)

        if step < 1000:
            mask = torch.norm(points, dim=-1) < 1.2
            outputs['sdf_pts'] = points[mask]
            outputs['sdf_vals'] = self.sdf_network.sdf(points[mask])[..., 0]

        if self.cfg['apply_occ_loss']:
            # occlusion loss
            if torch.sum(inner_mask) > 0:
                outputs['loss_occ'] = self.compute_occ_loss(occ_info, points[inner_mask], sdf, gradients, dirs[inner_mask], step)
            else:
                outputs['loss_occ'] = torch.zeros(1)

        if not is_train:
            outputs.update(self.compute_validation_info(z_vals, rays_o, rays_d, weights, human_poses, step))

        return outputs

    def forward(self, data):
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
        is_train = 'eval' not in data
        step = data['step']

        if is_train:
            outputs = self.train_step(step)
        else:
            index = data['index']
            outputs = self.test_step(index, step=step)

            if index == 0 and self.cfg['val_geometry']:
                bbox_min = -torch.ones(3)
                bbox_max = torch.ones(3)
                vertices, triangles = extract_geometry(bbox_min, bbox_max, 128, 0, lambda x: self.sdf_network.sdf(x))
                outputs['vertices'] = vertices
                outputs['triangles'] = triangles

        torch.set_default_tensor_type('torch.FloatTensor')
        return outputs

    def predict_materials(self):
        name = self.cfg['name']
        mesh = open3d.io.read_triangle_mesh(f'data/meshes/{name}-300000.ply')
        xyz = np.asarray(mesh.vertices)
        xyz = torch.from_numpy(xyz.astype(np.float32)).cuda()
        torch.set_default_tensor_type('torch.cuda.FloatTensor')

        metallic, roughness, albedo = [], [], []
        batch_size = 8192
        for vi in range(0, xyz.shape[0], batch_size):
            feature_vectors = self.sdf_network(xyz[vi:vi + batch_size])[:, 1:]
            m, r, a = self.color_network.predict_materials(xyz[vi:vi + batch_size], feature_vectors)
            metallic.append(m.cpu().numpy())
            roughness.append(r.cpu().numpy())
            albedo.append(a.cpu().numpy())

        return {'metallic': np.concatenate(metallic, 0),
                'roughness': np.concatenate(roughness, 0),
                'albedo': np.concatenate(albedo, 0)}

class NeROMaterialRenderer(nn.Module):
    default_cfg={
        'train_ray_num': 512,
        'test_ray_num': 1024,

        'database_name': 'real/bear/raw_1024',
        'rgb_loss': 'charbonier',

        'mesh': 'data/meshes/bear_shape-300000.ply',

        'shader_cfg': {},

        'reg_mat': True,
        'reg_diffuse_light': True,
        'reg_diffuse_light_lambda': 0.1,
        'fixed_camera': False,
    }
    def __init__(self, cfg, is_train=True):
        self.cfg = {**self.default_cfg, **cfg}
        super().__init__()
        self.warned_normal = False
        self._init_geometry()
        self._init_dataset(is_train)
        self._init_shader()

    def _init_geometry(self):
        self.mesh = open3d.io.read_triangle_mesh(self.cfg['mesh'])
        self.ray_tracer = raytracing.RayTracer(np.asarray(self.mesh.vertices), np.asarray(self.mesh.triangles))

    def _init_dataset(self, is_train):
        # train/test split
        self.database = parse_database_name(self.cfg['database_name'])
        self.train_ids, self.test_ids = get_database_split(self.database, 'validation')
        self.train_ids = np.asarray(self.train_ids)

        if is_train:
            self.train_imgs_info = build_imgs_info(self.database, self.train_ids)
            self.train_imgs_info = imgs_info_to_torch(self.train_imgs_info, 'cpu')
            self.train_num = len(self.train_ids)

            self.test_imgs_info = build_imgs_info(self.database, self.test_ids)
            self.test_imgs_info = imgs_info_to_torch(self.test_imgs_info, 'cpu')
            self.test_num = len(self.test_ids)

            self.train_batch = self._construct_ray_batch(self.train_imgs_info)
            self.tbn = self.train_batch['rays_o'].shape[0]
            self._shuffle_train_batch()

    def _init_shader(self):
        self.cfg['shader_cfg']['is_real'] = self.cfg['database_name'].startswith('real')
        self.shader_network = MCShadingNetwork(self.cfg['shader_cfg'], lambda o,d: self.trace(o,d))

    def trace_in_batch(self, rays_o, rays_d, batch_size=1024**2, cpu=False):
        inters, normals, depth, hit_mask = [], [], [], []
        rn = rays_o.shape[0]
        for ri in range(0, rn, batch_size):
            inters_cur, normals_cur, depth_cur, hit_mask_cur = self.trace(rays_o[ri:ri+batch_size], rays_d[ri:ri+batch_size])
            if cpu:
                inters_cur = inters_cur.cpu()
                normals_cur = normals_cur.cpu()
                depth_cur = depth_cur.cpu()
                hit_mask_cur = hit_mask_cur.cpu()
            inters.append(inters_cur)
            normals.append(normals_cur)
            depth.append(depth_cur)
            hit_mask.append(hit_mask_cur)
        return torch.cat(inters, 0), torch.cat(normals, 0), torch.cat(depth, 0), torch.cat(hit_mask, 0)

    def trace(self, rays_o, rays_d):
        inters, normals, depth = self.ray_tracer.trace(rays_o, rays_d)
        depth = depth.reshape(*depth.shape, 1)
        normals = -normals
        normals = F.normalize(normals, dim=-1)
        if not self.warned_normal:
            print('warn!!! the normals are flipped in NeuS by default. You may flip the normal according to your mesh!')
            self.warned_normal=True
        miss_mask = depth >= 10
        hit_mask = ~miss_mask
        return inters, normals, depth, hit_mask

    def _warn_ray_tracing(self, centers):
        centers = centers.reshape([-1,3])
        distance = torch.norm(centers,dim=-1) + 1.0
        max_dist = torch.max(distance).cpu().numpy()
        if max_dist>10.0:
            print(f'warning!!! the max distance from the camera is {max_dist:.4f}, which is beyond 10.0 for the ray tracer')

    def get_human_coordinate_poses(self, poses):
        pn = poses.shape[0]
        cam_cen = (-poses[:, :, :3].permute(0,2,1) @ poses[:, :, 3:])[..., 0]  # pn,3
        if self.cfg['fixed_camera']:
            pass
        else:
            cam_cen[..., 2] = 0

        Y = torch.zeros([1, 3], device='cpu').expand(pn, 3)
        Y[:, 2] = -1.0
        Z = torch.clone(poses[:, 2, :3])  # pn, 3
        Z[:, 2] = 0
        Z = F.normalize(Z, dim=-1)
        X = torch.cross(Y, Z)  # pn, 3
        R = torch.stack([X, Y, Z], 1)  # pn,3,3
        t = -R @ cam_cen[:, :, None]  # pn,3,1
        return torch.cat([R, t], -1)

    def _construct_ray_batch(self, imgs_info, device='cpu', is_train=True):
        imn, _, h, w = imgs_info['imgs'].shape
        coords = torch.stack(torch.meshgrid(torch.arange(h),torch.arange(w)),-1)[:,:,(1,0)] # h,w,2
        coords = coords.to('cpu')
        coords = coords.float()[None,:,:,:].repeat(imn,1,1,1) # imn,h,w,2
        coords = coords.reshape(imn,h*w,2)
        coords = torch.cat([coords+0.5, torch.ones(imn, h*w, 1, dtype=torch.float32, device='cpu')], 2) # imn,h*w,3

        # imn,h*w,3 @ imn,3,3 => imn,h*w,3
        rays_d = coords @ torch.inverse(imgs_info['Ks']).permute(0, 2, 1)
        poses = imgs_info['poses']  # imn,3,4
        R, t = poses[:, :, :3], poses[:, :, 3:]
        rays_d = rays_d @ R
        rays_d = F.normalize(rays_d, dim=-1)
        rays_o = -R.permute(0,2,1) @ t # imn,3,3 @ imn,3,1
        self._warn_ray_tracing(rays_o)
        rays_o = rays_o.permute(0, 2, 1).repeat(1, h*w, 1) # imn,h*w,3
        inters, normals, depth, hit_mask = self.trace_in_batch(rays_o.reshape(-1,3), rays_d.reshape(-1,3), cpu=True) # imn
        inters, normals, depth, hit_mask = inters.reshape(imn,h*w,3), normals.reshape(imn,h*w,3), depth.reshape(imn,h*w,1), hit_mask.reshape(imn, h*w)

        human_poses = self.get_human_coordinate_poses(poses) # imn,3,4
        human_poses = human_poses.unsqueeze(1).repeat(1,h*w,1,1) # imn,h*w,3,4
        rgb = imgs_info['imgs'].reshape(imn,3,h*w).permute(0,2,1) # imn,h*w,3

        if is_train:
            ray_batch={
                'rays_o': rays_o[hit_mask].to(device),
                'rays_d': rays_d[hit_mask].to(device),
                'inters': inters[hit_mask].to(device),
                'normals': normals[hit_mask].to(device),
                'depth': depth[hit_mask].to(device),
                'human_poses': human_poses[hit_mask].to(device),
                'rgb': rgb[hit_mask].to(device),
            }
        else:
            assert imn==1
            ray_batch={
                'rays_o': rays_o[0].to(device),
                'rays_d': rays_d[0].to(device),
                'inters': inters[0].to(device),
                'normals': normals[0].to(device),
                'depth': depth[0].to(device),
                'human_poses': human_poses[0].to(device),
                'rgb': rgb[0].to(device),
                'hit_mask': hit_mask[0].to(device),
            }
        return ray_batch

    def _shuffle_train_batch(self):
        self.train_batch_i = 0
        shuffle_idxs = torch.randperm(self.tbn, device='cpu') # shuffle
        for k, v in self.train_batch.items():
            self.train_batch[k] = v[shuffle_idxs]

    def shade(self, pts, view_dirs, normals, human_poses, is_train, step=None):
        rgb_pr, outputs = self.shader_network(pts, view_dirs, normals, human_poses, step, is_train)
        outputs['rgb_pr'] = rgb_pr
        return outputs

    def compute_rgb_loss(self, rgb_pr, rgb_gt):
        if self.cfg['rgb_loss'] == 'l1':
            rgb_loss = torch.sum(F.l1_loss(rgb_pr, rgb_gt, reduction='none'), -1)
        elif self.cfg['rgb_loss'] == 'charbonier':
            epsilon = 0.001
            rgb_loss = torch.sqrt(torch.sum((rgb_gt - rgb_pr) ** 2, dim=-1) + epsilon)
        else:
            raise NotImplementedError
        return rgb_loss

    def compute_diffuse_light_regularization(self, diffuse_lights):
        diffuse_white_reg = torch.sum(torch.abs(diffuse_lights - torch.mean(diffuse_lights, dim=-1, keepdim=True)), dim=-1)
        return diffuse_white_reg * self.cfg['reg_diffuse_light_lambda']

    def train_step(self, step):
        rn = self.cfg['train_ray_num']
        pts = self.train_batch['inters'][self.train_batch_i:self.train_batch_i + rn].cuda()
        view_dirs = -self.train_batch['rays_d'][self.train_batch_i:self.train_batch_i + rn].cuda() # view directions are opposite to ray directions
        normals = self.train_batch['normals'][self.train_batch_i:self.train_batch_i + rn].cuda()
        rgb_gt = self.train_batch['rgb'][self.train_batch_i:self.train_batch_i + rn].cuda()
        human_poses = self.train_batch['human_poses'][self.train_batch_i:self.train_batch_i + rn].cuda()

        shade_outputs = self.shade(pts, view_dirs, normals, human_poses, True, step)
        shade_outputs['rgb_gt'] = rgb_gt
        shade_outputs['loss_rgb'] = self.compute_rgb_loss(shade_outputs['rgb_pr'],shade_outputs['rgb_gt'])
        if self.cfg['reg_mat']:
            shade_outputs['loss_mat_reg'] = self.shader_network.material_regularization(
                pts, normals, shade_outputs['metallic'], shade_outputs['roughness'], shade_outputs['albedo'], step)
        if self.cfg['reg_diffuse_light']:
            shade_outputs['loss_diffuse_light'] = self.compute_diffuse_light_regularization(shade_outputs['diffuse_light'])

        self.train_batch_i += rn
        if self.train_batch_i + rn >= self.tbn: self._shuffle_train_batch()
        return shade_outputs

    def test_step(self, index):
        test_imgs_info = imgs_info_slice(self.test_imgs_info, torch.from_numpy(np.asarray([index],np.int64)))
        _, _, h, w = test_imgs_info['imgs'].shape
        ray_batch = self._construct_ray_batch(test_imgs_info, 'cuda', False)
        trn = self.cfg['test_ray_num']

        output_keys = {'rgb_gt':3, 'rgb_pr':3, 'specular_light':3, 'specular_color':3, 'diffuse_light':3, 'diffuse_color':3, 'albedo':3, 'metallic':1, 'roughness':1}
        outputs = {k:[] for k in output_keys.keys()}
        rn = ray_batch['rays_o'].shape[0]
        for ri in range(0, rn, trn):
            hit_mask = ray_batch['hit_mask'][ri:ri+trn]
            outputs_cur = {k: torch.zeros(hit_mask.shape[0], d) for k, d in output_keys.items()}
            if torch.sum(hit_mask)>0:
                pts = ray_batch['inters'][ri:ri+trn][hit_mask]
                view_dirs = -ray_batch['rays_d'][ri:ri+trn][hit_mask]
                normals = ray_batch['normals'][ri:ri+trn][hit_mask]
                rgb_gt = ray_batch['rgb'][ri:ri+trn][hit_mask]
                human_poses = ray_batch['human_poses'][ri:ri+trn][hit_mask]

                shade_outputs = self.shade(pts, view_dirs, normals, human_poses, False)

                outputs_cur['rgb_pr'][hit_mask] = shade_outputs['rgb_pr']
                outputs_cur['rgb_gt'][hit_mask] = rgb_gt
                outputs_cur['specular_light'][hit_mask] = shade_outputs['specular_light']
                outputs_cur['specular_color'][hit_mask] = shade_outputs['specular_color']
                outputs_cur['diffuse_color'][hit_mask] = shade_outputs['diffuse_color']
                outputs_cur['diffuse_light'][hit_mask] = shade_outputs['diffuse_light']
                outputs_cur['albedo'][hit_mask] = shade_outputs['albedo']
                outputs_cur['metallic'][hit_mask] = shade_outputs['metallic']
                outputs_cur['roughness'][hit_mask] = torch.sqrt(shade_outputs['roughness']) # note: we assume predictions are roughness squared

            for k in output_keys.keys():
                outputs[k].append(outputs_cur[k])

        for k in output_keys.keys():
            outputs[k] = torch.cat(outputs[k], 0).reshape(h, w, -1)

        return outputs

    def forward(self, data):
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
        is_train = 'eval' not in data
        step = data['step']

        if is_train:
            outputs = self.train_step(step)
        else:
            index = data['index']
            outputs = self.test_step(index)

        torch.set_default_tensor_type('torch.FloatTensor')
        return outputs

    def predict_materials(self, batch_size=8192):
        verts = torch.from_numpy(np.asarray(self.mesh.vertices,np.float32)).cuda().float()
        metallic, roughness, albedo = [], [], []
        for vi in range(0, verts.shape[0], batch_size):
            m, r, a = self.shader_network.predict_materials(verts[vi:vi+batch_size])
            r = torch.sqrt(torch.clamp(r, min=1e-7)) # note: we assume predictions are squared roughness!!!
            metallic.append(m.cpu().numpy())
            roughness.append(r.cpu().numpy())
            albedo.append(a.cpu().numpy())

        return {'metallic': np.concatenate(metallic, 0),
                'roughness': np.concatenate(roughness, 0),
                'albedo': np.concatenate(albedo, 0)}

name2renderer={
    'shape': NeROShapeRenderer,
    'material': NeROMaterialRenderer,
}


a = int(input('Um número: '))
b = int(input('Outro número: '))
#resto a = a % 2 primeiro IF
resto_a = a % 2
restO_b = b % 2
#Primeiro IF
#if resto == 0:
#    print('O número {} é par!')
#else:
#    print('Esse número é impar!')

if resto_a == 0 or restO_b == 0:
    print('Foi digitado um número par.')
else:
    print('Não foi digitado nenhum número par.')


"""Validate the various output formats"""

import os
import tempfile
from typing import Any, Callable, Dict, Iterable, List, Tuple

import numpy
import numpy.typing
import npyosmium
import npyosmium.io
import npyosmium.osm
import pytest

from pyhgtmap.hgt.tile import TileContours
from pyhgtmap.output import make_elev_classifier, o5mUtil, osmUtil, pbfUtil

class OSMDecoder(npyosmium.SimpleHandler):
    """Basic OSM file decoder, relying on official osmium library."""

    def __init__(self) -> None:
        super().__init__()
        self.nodes: Dict[int, Any] = {}
        self.ways: Dict[int, Any] = {}

    def node(self, n: npyosmium.osm.Node) -> None:
        try:
            self.nodes[n.id] = (n.location.lat, n.location.lon)
        except Exception:
            pass

    def way(self, w: npyosmium.osm.Way) -> None:
        self.ways[w.id] = ([node.ref for node in w.nodes], [tag for tag in w.tags])

def arrays_from_lists(
    coordinates_lists: List[List[Tuple[int, int]]]
) -> List[numpy.typing.NDArray]:
    """Helper to convert list of lists into array of arrays."""
    return [
        numpy.array(
            [numpy.array(coordinates, dtype=numpy.float64) for coordinates in way]
        )
        for way in coordinates_lists
    ]

@pytest.fixture
def tile_contours() -> TileContours:
    return TileContours(
        nb_nodes=8,
        nb_ways=3,
        contours={
            # Elevation 0
            0: arrays_from_lists(
                [
                    # Closed loop
                    [(1, 1), (1, 2), (2, 2), (2, 1), (1, 1)],
                    # Open one
                    [(3, 1), (3, 2)],
                ]
            ),
            # Elevation 50
            50: arrays_from_lists([[(4, 1), (4, 2)]]),
            # Elevation 100
            100: [],
            # Elevation 150
            150: [],
        },
    )

@pytest.fixture
def bounding_box() -> Tuple[float, float, float, float]:
    """Bounding box of all fake data nodes"""
    return (1, 1, 4, 2)

@pytest.fixture
def elev_classifier() -> Callable[[int], str]:
    return make_elev_classifier(100, 50)

@pytest.fixture
def elevations() -> Iterable[int]:
    return [0, 50, 100, 150]

def check_osmium_result(osm_file_name: str) -> None:
    """Check results using osmium. Data corresponding to shared fixtures."""
    osm_decoder = OSMDecoder()
    reader = npyosmium.io.Reader(osm_file_name)
    header = reader.header()
    # Check header attributes
    assert header.box().bottom_left == npyosmium.osm.Location(1, 1)
    assert header.box().top_right == npyosmium.osm.Location(4, 2)
    if osm_file_name[-4:] != ".o5m":
        # Not implemented for o5m format
        assert header.get("generator") == "pyhgtmap 123"
    osm_decoder.apply_file(osm_file_name)

    # Check nodes
    assert osm_decoder.nodes == {
        1000: (1.0, 1.0),
        1001: (2.0, 1.0),
        1002: (2.0, 2.0),
        1003: (1.0, 2.0),
        1004: (1.0, 3.0),
        1005: (2.0, 3.0),
        1006: (1.0, 4.0),
        1007: (2.0, 4.0),
    }

    # Check ways
    assert osm_decoder.ways == {
        2000: (
            # Loop must be closed, re-using first ID
            [1000, 1001, 1002, 1003, 1000],
            [
                npyosmium.osm.Tag(k="ele", v="0"),
                npyosmium.osm.Tag(k="contour", v="elevation"),
                npyosmium.osm.Tag(k="contour_ext", v="elevation_major"),
            ],
        ),
        2001: (
            [1004, 1005],
            [
                npyosmium.osm.Tag(k="ele", v="0"),
                npyosmium.osm.Tag(k="contour", v="elevation"),
                npyosmium.osm.Tag(k="contour_ext", v="elevation_major"),
            ],
        ),
        2002: (
            [1006, 1007],
            [
                npyosmium.osm.Tag(k="ele", v="50"),
                npyosmium.osm.Tag(k="contour", v="elevation"),
                # Medium elevation expected
                npyosmium.osm.Tag(k="contour_ext", v="elevation_medium"),
            ],
        ),
    }

class TestOutputOsm:
    @staticmethod
    def test_produce_osm(tile_contours: TileContours, elev_classifier) -> None:
        """Generate OSM file out of mocked data and check content."""
        with tempfile.TemporaryDirectory() as tempdir:
            osm_file_name = os.path.join(tempdir, "output.osm")
            osm_output = osmUtil.Output(
                osm_file_name,
                osmVersion=0.6,
                pyhgtmap_version="123",
                boundsTag='<bounds minlat="1" minlon="1" maxlat="2" maxlon="4"/>',
                gzip=0,
                elevClassifier=elev_classifier,
                timestamp=False,
            )

            # Write OSM file
            next_node_id, ways = osm_output.write_nodes(
                tile_contours, ' time="some time"', 1000, 0.6
            )
            assert next_node_id == 1008
            osm_output.write_ways(ways, 2000)
            osm_output.done()

            # Check file output
            contents: str = open(osm_file_name).read()
            assert (
                contents
                == """<?xml version="1.0" encoding="utf-8"?>
<osm version="0.6" generator="pyhgtmap 123">
<bounds minlat="1" minlon="1" maxlat="2" maxlon="4"/>
<node id="1000" lat="1.0000000" lon="1.0000000" version="1" time="some time"/>
<node id="1001" lat="2.0000000" lon="1.0000000" version="1" time="some time"/>
<node id="1002" lat="2.0000000" lon="2.0000000" version="1" time="some time"/>
<node id="1003" lat="1.0000000" lon="2.0000000" version="1" time="some time"/>
<node id="1004" lat="1.0000000" lon="3.0000000" version="1" time="some time"/>
<node id="1005" lat="2.0000000" lon="3.0000000" version="1" time="some time"/>
<node id="1006" lat="1.0000000" lon="4.0000000" version="1" time="some time"/>
<node id="1007" lat="2.0000000" lon="4.0000000" version="1" time="some time"/>
<way id="2000" version="1"><nd ref="1000"/>
<nd ref="1001"/>
<nd ref="1002"/>
<nd ref="1003"/>
<nd ref="1000"/>
<tag k="ele" v="0"/><tag k="contour" v="elevation"/><tag k="contour_ext" v="elevation_major"/></way>
<way id="2001" version="1"><nd ref="1004"/>
<nd ref="1005"/>
<tag k="ele" v="0"/><tag k="contour" v="elevation"/><tag k="contour_ext" v="elevation_major"/></way>
<way id="2002" version="1"><nd ref="1006"/>
<nd ref="1007"/>
<tag k="ele" v="50"/><tag k="contour" v="elevation"/><tag k="contour_ext" v="elevation_medium"/></way>
</osm>
"""
            )

            # Check file with osmium
            check_osmium_result(osm_file_name)

class TestOutputPbf:
    @staticmethod
    def test_produce_pbf(
        tile_contours: TileContours,
        elev_classifier,
        bounding_box: Tuple[float, float, float, float],
    ) -> None:
        """Generate PBF file out of mocked data and check content."""
        with tempfile.TemporaryDirectory() as tempdir:
            osm_file_name = os.path.join(tempdir, "output.osm.pbf")
            osm_output = pbfUtil.Output(
                osm_file_name,
                osmVersion=0.6,
                pyhgtmap_version="123",
                bbox=bounding_box,
                elevClassifier=elev_classifier,
            )

            # Write OSM file
            next_node_id, ways = osm_output.write_nodes(
                tile_contours, ' time="some time"', 1000, 0.6
            )
            assert next_node_id == 1008
            osm_output.write_ways(ways, 2000)
            osm_output.done()

            # Check file content with osmium
            check_osmium_result(osm_file_name)

            # Check file size to ensure there's no major drop in efficiency (compression, dense encoding)
            assert os.stat(osm_file_name).st_size < 420

    @staticmethod
    def test_node_id_overflow(
        tile_contours: TileContours,
        elev_classifier,
        bounding_box: Tuple[float, float, float, float],
    ) -> None:
        with tempfile.TemporaryDirectory() as tempdir:
            osm_file_name = os.path.join(tempdir, "output.osm.pbf")
            osm_output = pbfUtil.Output(
                osm_file_name,
                osmVersion=0.6,
                pyhgtmap_version="123",
                bbox=bounding_box,
                elevClassifier=elev_classifier,
            )
            start = 2147483647
            next_node_id, ways = osm_output.write_nodes(
                tile_contours, ' time="some time"', start, 0.6
            )

            # Check int32 boundary has been passed without issue
            assert next_node_id == 2147483655

class TestOutputO5m:
    @staticmethod
    def test_produce_o5m(
        tile_contours: TileContours,
        elev_classifier,
        bounding_box: Tuple[float, float, float, float],
    ) -> None:
        """Generate PBF file out of mocked data and check content."""
        with tempfile.TemporaryDirectory() as tempdir:
            osm_file_name = os.path.join(tempdir, "output.osm.o5m")
            osm_output = o5mUtil.Output(
                osm_file_name,
                osmVersion=0.6,
                pyhgtmap_version="123",
                bbox=bounding_box,
                elevClassifier=elev_classifier,
            )

            # Write OSM file
            next_node_id, ways = osm_output.write_nodes(
                tile_contours, ' time="some time"', 1000, 0.6
            )
            assert next_node_id == 1008
            osm_output.write_ways(ways, 2000)
            osm_output.done()

            # Check file with osmium
            check_osmium_result(osm_file_name)


from collections import defaultdict

def is_member(words, query):
    for word in words:
        if is_match(query, word):
            return True

    return False

def is_match(query, word):
    if len(query) != len(word):
        return False

    for query_char, word_char in zip(query, word):
        if query_char != word_char and query_char != '*':
            return False

    return True

print(is_member(["foo", "bar", "baz"], "foo"))

class Dictionary:
    def __init__(self, words):
        self.trie = defaultdict(dict)
        [ self.add(word) for word in words ]

    def is_match(self, query):
        def word_match(node, idx=0):
            if idx == len(query):
                return "." in node

            letter = query[idx]
            if query[idx] == "*":
                return True in [ word_match(child, idx+1) for child in node.values() ]

            if letter not in node:
                return False

            return word_match(node[letter], idx+1)

        return word_match(self.trie)

    def add(self, word):
        curr = self.trie
        for idx, letter in enumerate(word):
            if letter not in curr:
                curr[letter] = {}

            curr = curr[letter]

        curr["."] = {}

dictionary = Dictionary(["foo", "bar", "baz"])
print(dictionary.is_match("foo"))
print(dictionary.is_match("f**"))
print(dictionary.is_match("f*o"))
print(dictionary.is_match("***"))
print(dictionary.is_match("*z*"))


import os
import pdb
import sys
import time
import torch
import shutil
import random
from barbar import Bar
from datetime import datetime
from torchsummary import summary
from sklearn.metrics import accuracy_score

class SGPU_TrnOnly:
    def __init__(self,
                 params,
                 net,
                 optimizer,
                 criterion,
                 trainloader,
                 cuda_device_id=0,
                 seed=random.randint(101, 201)):
        """ Trains a model on one GPU.

        Parameters
        ----------
        params: Dict
        A dictionary having training parameters.
        Ex:
        ```
        {
            "max_epochs": 100,
        }
        ```
        net: Custom Network Instance
            Instance of pytorch custom network we will be training
        optimizer: Optimizer instance
        criterion: Loss instance
        trainloader: Instance of Custom DataLoader
            Training data loader instance
        cuda_device_id: int, optional
            CUDA device ID. Defaults to 0.
        seed: int, optional
            sets random seed for gpu and cpu
        """
        # Set random seed
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed) # --> not needed. Just in case.

        self.params = params

        # Load network into GPU
        self.net = net
        self.optimizer = optimizer
        self.criterion = criterion
        self.trainloader = trainloader

        # Load network into GPU device
        self.device = self._get_cuda_device(cuda_device_id)
        self.net.to(self.device)

        # Create checkpoint directory if it does not exist
        self.work_dir = self.params['work_dir']
        if os.path.isdir(self.work_dir):
            shutil.rmtree(self.work_dir)
        print(f"INFO: Creating work directory {self.params['work_dir']}")
        os.makedirs(self.work_dir)

        # Create log json file
        self.log = open(self.params['log_pth'], "w")

        # Epochs for which we need to save checkpoints
        self.max_epochs = self.params['max_epochs']
        ckpt_save_interval = self.params['ckpt_save_interval']
        self.ckpt_save_epochs = list(
            range(0, self.max_epochs, ckpt_save_interval))

        # Save model diagram in model.txt
        model_stats = summary(self.net, params['input_shape'])
        summary_str = str(model_stats)
        model_floc = f"{os.path.dirname(self.params['log_pth'])}/model.txt"
        model_f = open(model_floc, "w")
        model_f.write(summary_str)
        model_f.close()

    def _get_cuda_device(self, cuda_device_id):
        """ Sets NVIDIA device with id `cuda_device_id` to
        be used for training.

        Parameters
        ----------
        net: Custom Network Instance
            Instance of pytorch custom network we will be training
        cuda_device_id: int
            CUDA device ID
        """
        if not torch.cuda.is_available():
            raise Exception("ERROR: Cuda is not found in this environment")

        num_cuda_devices = torch.cuda.device_count()
        cuda_devices = list(range(0,
                                  num_cuda_devices))  # cuda index start from 0

        if not (cuda_device_id <= num_cuda_devices - 1):
            raise Exception(
                f"ERROR: Cuda device {cuda_device_id} is not found.\n"
                f"ERROR: Found {num_cuda_devices}("
                f"{cuda_devices}), Cuda devices")

        # If no errors load network onto cuda device
        print("INFO: Network sucessfully loaded into "
              f"CUDA device {cuda_device_id}")
        return f"cuda:{cuda_device_id}"

    def train(self):
        """ Trains a network with specified dataloaser.

        Parameters
        ----------

        """
        print(f"INFO: Starting training for {self.max_epochs} epochs")
        now = datetime.now().strftime("%m/%d/%Y, %H:%M:%S")
        self.log.write(f'{{"mode": "info", "start_time":"{now}" }}')

        # Epoch loop
        for epoch in range(self.max_epochs):
            print(
                f"\n*****************Epoch {epoch}**************************")

            # Training
            print("Trn:")
            self.net.train()
            trnloss, trnaccu, trntt = self._train_and_get_loss_and_accuracy()

            # Save every n epochs
            if epoch in self.ckpt_save_epochs:
                ckpt_loc = f"{self.work_dir}/epoch_{epoch}.pth"
                self._save_model(epoch, trnloss, ckpt_loc)

            # Print useful information per epoch
            trnstr = (f'{{"mode": "train", "epoch":{epoch}, "acc":{trnaccu}, '
                      f'"loss":{trnloss}}}')

            # Write log file (this will slow down trainning <<< How much?)
            self.log.write(f"\n{trnstr}")
            self.log.flush()
            os.fsync(self.log.fileno())

            # Print to console
            print(f"Trn: {trnstr}")

        # Recording end time
        now = datetime.now().strftime("%m/%d/%Y, %H:%M:%S")
        self.log.write(f'\n{{"mode": "info", "end_time":"{now}" }}')

        self.log.close()

    def _train_and_get_loss_and_accuracy(self):
        """ Trains and returns avrage training loss and accuracy
        """
        # Training Iteration loop
        trngt_lst = []
        trntt_lst = []
        trnloss_lst = []
        trnpred_lst = []
        for idx, data in enumerate(Bar(self.trainloader)):

            # Train for current iteration
            trnloss, trnpred, trngt, tt = self._train_cur_iter(data)

            # Collect training loss and prediction and time taken
            trntt_lst += [tt]
            trngt_lst += trngt
            trnloss_lst += [trnloss]
            trnpred_lst += trnpred

        # Training loss and prediction accuracy per epoch
        trngt_lst = [round(x) for x in trngt_lst]
        trnpred_lst = [round(x) for x in trnpred_lst]
        tot_tt = round(sum(trntt_lst), 5)
        trnaccu = round(accuracy_score(trngt_lst, trnpred_lst), 2)
        avg_trnloss = round(sum(trnloss_lst) / len(trnloss_lst), 2)

        return avg_trnloss, trnaccu, tot_tt

    def _train_cur_iter(self, data):
        """ Trains f
        """
        # Label and input tensors
        labels, inputs = (data[0].to(self.device, non_blocking=True),
                          data[1].to(self.device, non_blocking=True))
        labels = torch.reshape(labels.float(), (-1, 1))

        # Training Core
        st = time.time()  # Start time
        self.optimizer.zero_grad()  # Zeroing the gradient
        outputs = self.net(inputs)  # Prediction
        loss = self.criterion(outputs, labels)  # Calculate loss
        loss.backward()  # Back propagation
        self.optimizer.step()  # Update the gradients
        et = time.time()  # End time
        tt = round(et - st, 5)  # time taken

        # Collecting information to cpu memory and to list
        ipred = outputs.data.clone()
        ipred = ipred.to("cpu").numpy().flatten().tolist()
        iloss = loss.data.clone()
        iloss = iloss.to("cpu").numpy().tolist()
        ilabels = labels.data.clone()
        ilabels = ilabels.to("cpu").numpy().flatten().tolist()

        return iloss, ipred, ilabels, tt

    def _save_model(self, epoch, trnloss, ckpt_loc):
        """ Saves model as pth file
        """
        torch.save(
            {
                'epoch': epoch,
                'model_state_dict': self.net.state_dict(),
                'optimizer_state_dict': self.optimizer.state_dict(),
                'loss': trnloss,
            }, ckpt_loc)


# SHIV's code for activity selection
'''
# sample test case

2
6
1 3 2 5 8 5
2 4 6 7 9 9
10
41 13 4 70 10 58 61 34 100 79
68 39 12 97 13 66 82 38 120 99
'''

tn = int(input())

# sort by end times and then pick out those whose start times were greater
# than prev ends
def process(club):
	club.sort(key = lambda x: x[1])
	prev, tasks = 0, 0
	for s, e in club:
		if s >= prev:
			tasks+=1
			prev=e
	return tasks

for i in range(tn):
	tasks = int(input())
	starts = list(map(int, input().split(' ')))
	ends = list(map(int, input().split(' ')))
	club = list(map(list, zip(starts, ends)))
	print(process(club))


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

import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

def save_result_img(args, in_list, name):
        plt.figure()
        plt.plot(range(len(in_list)), in_list)
        plt.ylabel(name)
        plt.savefig('./save/fed_{}_{}_{}_C{}_iid{}_{}.png'.format(args.dataset, args.model, args.epochs, args.frac, args.iid, name))

def test_img(net_g, datatest, args):
    net_g.eval()
    # testing
    test_loss = 0
    correct = 0

    correct_list = []
    loss_list = []

    data_loader = DataLoader(datatest, batch_size=args.bs)
    l = len(data_loader)
    for idx, (data, target) in enumerate(data_loader):
        if args.gpu != -1:
            data, target = data.cuda(), target.cuda()
        log_probs = net_g(data)
        # sum up batch loss
        test_loss += F.cross_entropy(log_probs, target, reduction='sum').item()
        # get the index of the max log-probability
        y_pred = log_probs.data.max(1, keepdim=True)[1]
        correct += y_pred.eq(target.data.view_as(y_pred)).long().cpu().sum()

        # store correct list for report
        correct_list.append(100.00 * correct / len(data_loader.dataset))
        # loss_list.append(test_loss)

    test_loss /= len(data_loader.dataset)
    accuracy = 100.00 * correct / len(data_loader.dataset)
    # save_result_img(args, correct_list, name="test accurcy")
    if args.verbose:
        print('\nTest set: Average loss: {:.4f} \nAccuracy: {}/{} ({:.2f}%)\n'.format(
            test_loss, correct, len(data_loader.dataset), accuracy))
    return accuracy, test_loss

def state_dict_to_xferable( state_dict, lr=None, ep=None ):
    import json
    xfer_dict = {}
    if lr is not None:
        xfer_dict['lr'] = lr
    if ep is not None:
        xfer_dict['ep'] = ep
    for k in state_dict:
        xfer_dict[k] = state_dict[k].cpu().tolist()
    xfer_string = json.dumps(xfer_dict)
    return xfer_string

def xferable_to_state_dict(xfer_string):
    import json
    import collections
    json_decode_dict = json.loads(xfer_string)
    recovered_dict = collections.OrderedDict()
    for k in json_decode_dict:
        if k != 'lr' and k != 'ep':
            recovered_dict[k] = torch.Tensor(json_decode_dict[k])
    print("recovered_dict: ", type(recovered_dict))
    return recovered_dict, json_decode_dict['lr'], json_decode_dict['ep']


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

# (c) 2013, Chatham Financial <oss@chathamfinancial.com>
# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)

from __future__ import absolute_import, division, print_function
__metaclass__ = type

ANSIBLE_METADATA = {'metadata_version': '1.1',
                    'status': ['preview'],
                    'supported_by': 'community'}

DOCUMENTATION = '''
---
module: rabbitmq_vhost
short_description: Manage the state of a virtual host in RabbitMQ
description:
  - Manage the state of a virtual host in RabbitMQ
version_added: "1.1"
author: '"Chris Hoffman (@choffman)"'
options:
  name:
    description:
      - The name of the vhost to manage
    required: true
    aliases: [vhost]
  node:
    description:
      - erlang node name of the rabbit we wish to configure
    default: rabbit
    version_added: "1.2"
  tracing:
    description:
      - Enable/disable tracing for a vhost
    type: bool
    default: 'no'
    aliases: [trace]
  state:
    description:
      - The state of vhost
    default: present
    choices: [present, absent]
'''

EXAMPLES = '''
# Ensure that the vhost /test exists.
- rabbitmq_vhost:
    name: /test
    state: present
'''

from ansible.module_utils.basic import AnsibleModule

class RabbitMqVhost(object):
    def __init__(self, module, name, tracing, node):
        self.module = module
        self.name = name
        self.tracing = tracing
        self.node = node

        self._tracing = False
        self._rabbitmqctl = module.get_bin_path('rabbitmqctl', True)

    def _exec(self, args, run_in_check_mode=False):
        if not self.module.check_mode or (self.module.check_mode and run_in_check_mode):
            cmd = [self._rabbitmqctl, '-q', '-n', self.node]
            rc, out, err = self.module.run_command(cmd + args, check_rc=True)
            return out.splitlines()
        return list()

    def get(self):
        vhosts = self._exec(['list_vhosts', 'name', 'tracing'], True)

        for vhost in vhosts:
            if '\t' not in vhost:
                continue

            name, tracing = vhost.split('\t')
            if name == self.name:
                self._tracing = self.module.boolean(tracing)
                return True
        return False

    def add(self):
        return self._exec(['add_vhost', self.name])

    def delete(self):
        return self._exec(['delete_vhost', self.name])

    def set_tracing(self):
        if self.tracing != self._tracing:
            if self.tracing:
                self._enable_tracing()
            else:
                self._disable_tracing()
            return True
        return False

    def _enable_tracing(self):
        return self._exec(['trace_on', '-p', self.name])

    def _disable_tracing(self):
        return self._exec(['trace_off', '-p', self.name])

def main():
    arg_spec = dict(
        name=dict(required=True, aliases=['vhost']),
        tracing=dict(default='off', aliases=['trace'], type='bool'),
        state=dict(default='present', choices=['present', 'absent']),
        node=dict(default='rabbit'),
    )

    module = AnsibleModule(
        argument_spec=arg_spec,
        supports_check_mode=True
    )

    name = module.params['name']
    tracing = module.params['tracing']
    state = module.params['state']
    node = module.params['node']
    result = dict(changed=False, name=name, state=state)
    rabbitmq_vhost = RabbitMqVhost(module, name, tracing, node)

    if rabbitmq_vhost.get():
        if state == 'absent':
            rabbitmq_vhost.delete()
            result['changed'] = True
        else:
            if rabbitmq_vhost.set_tracing():
                result['changed'] = True
    elif state == 'present':
        rabbitmq_vhost.add()
        rabbitmq_vhost.set_tracing()
        result['changed'] = True

    module.exit_json(**result)

if __name__ == '__main__':
    main()


import collections

class Node:
    def __init__(self, val=0, neighbors=None):
        self.val = val
        self.neighbors = neighbors

    def __repr__(self):
        return "node value:" + str(self.val)

def dfs(node: Node) -> Node:
    c = {node: Node(node.val, [])}
    stack = [node]

    while stack:
        n = stack.pop()
        for neighbor in n.neighbors:
            if neighbor not in c:
                c[neighbor] = Node(neighbor.val, [])
                stack.append(neighbor)
            c[n].neighbors.append(c[neighbor])

    return c[node]

def bfs(node: Node) -> Node:
    c = {node: Node(node.val, [])}
    deque = collections.deque([node])

    while deque:
        d = deque.popleft()
        for neighbor in d.neighbors:
            if neighbor not in c:
                d.append(neighbor)
                c[neighbor] = Node(neighbor.val)
            c[d].neighbors.append(c[neighbor])

    return c[node]

if __name__ == "__main__":
    a = Node(1)
    b = Node(2)
    c = Node(3)
    d = Node(4)
    a.neighbors = [b, d]
    b.neighbors = [a, c]
    c.neighbors = [b, d]
    d.neighbors = [a, c]

    res = dfs(a)


from ctypes import *

# #define PINK_MAX_RANDOM_ROWS   (30)
PINK_MAX_RANDOM_ROWS = 30

# #define PINK_RANDOM_BITS       (24)
PINK_RANDOM_BITS = 24

# #define PINK_RANDOM_SHIFT      ((sizeof(long)*8)-PINK_RANDOM_BITS)
PINK_RANDOM_SHIFT = (sizeof(c_long) * 8) - PINK_RANDOM_BITS

# type of `pink_rows[]` member of `pink_noise_t` struct
PINK_ROWS_TYPE = c_long * PINK_MAX_RANDOM_ROWS

# Struct used to track the pink noise generator state
class Pink_noise_t(Structure):
    """
    typedef struct
    {
      long      pink_rows[PINK_MAX_RANDOM_ROWS];
      long      pink_running_sum;   /* Used to optimize summing of generators. */
      int       pink_index;        /* Incremented each sample. */
      int       pink_index_mask;    /* Index wrapped by ANDing with this mask. */
      float     pink_scalar;       /* Used to scale within range of -1.0 to +1.0 */
    } pink_noise_t;
    """
    _fields_ = (("pink_rows", PINK_ROWS_TYPE),
                ("pink_running_sum", c_long),
                ("pink_index", c_int),
                ("pink_index_mask", c_int),
                ("pink_scalar", c_float))


import pandas as pd
import numpy as np

from torch.utils.data import DataLoader
import warnings,transformers,logging,torch
from transformers import TrainingArguments,Trainer
from transformers import AutoModelForSequenceClassification,AutoTokenizer, BertTokenizer
import datasets
from datasets import load_dataset, Dataset, DatasetDict
from sklearn.metrics import log_loss
import torch.nn.functional as F

df = pd.read_csv('data/train.csv')

warnings.simplefilter('ignore')
logging.disable(logging.WARNING)

model_nm = 'microsoft/deberta-v3-base'
tokz = AutoTokenizer.from_pretrained(model_nm)
sep = tokz.sep_token
df['inputs'] = df.discourse_type + sep +df.discourse_text
new_label = {"discourse_effectiveness": {"Ineffective": 0, "Adequate": 1, "Effective": 2}}
df = df.replace(new_label)
df = df.rename(columns = {"discourse_effectiveness": "label"})
ds = Dataset.from_pandas(df)
def tok_func(x): return tokz(x["inputs"], truncation=True)
tok_func(ds[0])

inps = "discourse_text","discourse_type"
tok_ds = ds.map(tok_func, batched=True, remove_columns=inps+('inputs','discourse_id','essay_id'))

essay_ids = df.essay_id.unique()
np.random.seed(42)
np.random.shuffle(essay_ids)

val_prop = 0.2
val_sz = int(len(essay_ids)*val_prop)
val_essay_ids = essay_ids[:val_sz]

is_val = np.isin(df.essay_id, val_essay_ids)
idxs = np.arange(len(df))
val_idxs = idxs[ is_val]
trn_idxs = idxs[~is_val]

dds = DatasetDict({"train":tok_ds.select(trn_idxs),
             "test": tok_ds.select(val_idxs)})

def get_dds(df, train=True):
    ds = Dataset.from_pandas(df)
    to_remove = ['discourse_text','discourse_type','inputs','discourse_id','essay_id']
    tok_ds = ds.map(tok_func, batched=True, remove_columns=to_remove)
    if train:
        return DatasetDict({"train":tok_ds.select(trn_idxs), "test": tok_ds.select(val_idxs)})
    else:
        return tok_ds

lr,bs = 8e-5,16
wd,epochs = 0.01,4

def score(preds): return {'log loss': log_loss(preds.label_ids, F.softmax(torch.Tensor(preds.predictions)))}
def get_trainer(dds):
    args = TrainingArguments('outputs', learning_rate=lr, warmup_ratio=0.1, lr_scheduler_type='cosine',
        evaluation_strategy="epoch", per_device_train_batch_size=bs, per_device_eval_batch_size=bs*2,
        num_train_epochs=epochs, weight_decay=wd, report_to='none')
    model = AutoModelForSequenceClassification.from_pretrained(model_nm, num_labels=3)
    return Trainer(model, args, train_dataset=dds['train'], eval_dataset=dds['test'],
                   tokenizer=tokz, compute_metrics=score)

trainer = get_trainer(dds)

trainer.train()


#Comma code

def myListBreaker(myList):
    myString = ''
    for i in range(len(myList)-1):
            myString += (str(myList[i]) + ',')
    myString += 'and ' + str(myList[-1])
    print(myString)

newList = [1, 'two', 3, 'four']
myListBreaker(newList)


# """
# Day 5: Hydrothermal Venture
# https://adventofcode.com/2021/day/5
#
# Usage:
# $ python solution.py input.txt
# """

import sys
from typing import Iterator, TypeAlias

import numpy as np

Line: TypeAlias = list[tuple[int, int]]

def main(path: str) -> None:
    assert 6267 == solve(load(path))
    assert 20196 == solve(load(path, use_diagonals=True))

def solve(lines: list[Line]) -> int:
    """Return the number of overlapping points."""
    diagram = make_diagram(lines)
    for line in lines:
        for point in line:
            diagram[point] += 1
    return len(diagram[diagram > 1])

def make_diagram(lines: list[Line]) -> np.ndarray:
    """Return a square matrix initialized with zeros."""
    size = np.concatenate([np.array(a).flatten() for a in lines]).max()
    return np.zeros(shape=[size + 1] * 2, dtype=int)

def load(path: str, use_diagonals: bool = False) -> list[Line]:
    """Return a list of lines comprised of points."""

    def load_lazy() -> Iterator:
        with open(path, encoding="utf-8") as file:
            for line in file:
                x1, y1, x2, y2 = np.fromstring(
                    line.replace("->", ","), dtype=int, sep=","
                )
                if x1 == x2:
                    yield vertical(x1, y1, y2)
                elif y1 == y2:
                    yield horizontal(y1, x1, x2)
                elif abs(x2 - x1) == (abs(y2 - y1)):
                    if use_diagonals:
                        yield diagonal(x1, y1, x2, y2)

    return list(load_lazy())

def vertical(x1: int, y1: int, y2: int) -> Line:
    """Return points of a vertical line."""
    y = np.linspace(y1, y2, dtype=int, num=abs(y2 - y1) + 1)
    x = np.repeat(x1, len(y))
    return list(zip(y, x))

def horizontal(y1: int, x1: int, x2: int) -> Line:
    """Return points of a horizontal line."""
    x = np.linspace(x1, x2, dtype=int, num=abs(x2 - x1) + 1)
    y = np.repeat(y1, len(x))
    return list(zip(y, x))

def diagonal(x1: int, y1: int, x2: int, y2: int) -> Line:
    """Return points of a diagonal line."""
    x = np.linspace(x1, x2, dtype=int, num=abs(x2 - x1) + 1)
    y = np.linspace(y1, y2, dtype=int, num=abs(y2 - y1) + 1)
    return list(zip(y, x))

if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "input.txt")