import os

import pytest
import lxml.etree

from aristo.core.knowledge_creator import KnowledgeCreator

test_data = [
    (["sinkhole"], 1, 1),
    (["56fsvhvnbv"], 2, 0),
    (["sinkhole", "cell"], 1, 2),
    (["sinkhole", "cell"], 5, 2)

]

@pytest.mark.parametrize("articles, batch_size, expected_answer", test_data)
def test_should_download_articles(tmpdir, articles, batch_size, expected_answer):
    output_file = str(tmpdir.mkdir("test_should_predict_the_correct_answer").join("wiki.xml"))
    sut = KnowledgeCreator()
    actual = 0;
    for file in sut.download_corpus(articles, output_file, batch_size):
        doc = lxml.etree.parse(file)
        actual = actual + doc.xpath('count(/*[local-name() = "mediawiki"]/*[local-name() = "page"])')
    assert actual == expected_answer


#!/usr/bin/env python3
# -*- author:xuanGe -*-
# -*- coding:utf-8 -*-
# -*- date:2020-2-13 -*-
import sys
import time
import asyncio
import libpcap as pcap
import ctypes as ct
from config import device
from log import info, success, debug, warning, error

async def collector(q=None):
    # open network interfaces
    errbuf = ct.create_string_buffer(pcap.PCAP_ERRBUF_SIZE + 1)
    pd = pcap.open_live(device, 65535, 1, 1000, errbuf)
    if not pd:
        error(f"c: pcap.open_live failed: {errbuf.value.decode('utf-8')}")
        debug(pd)
    if pcap.setnonblock(pd, 1, errbuf) == -1:
        error(f"c: pcap.setnonblock failed: {errbuf.value.decode('utf-8')}")
    # capture data paket
    try:
        while True:
            status = pcap.dispatch(pd, -1, handle_data, ct.cast(id(q), ct.POINTER(ct.c_ubyte)))
            await asyncio.sleep(0)
    except Exception as e:
        if str(e) != '':
            error(f"c: while filed: {e}")

@pcap.pcap_handler
def handle_data(q_id_pointer, hdr, pkt):
    q = ct.cast(q_id_pointer, ct.py_object).value
    debug(f"c: caplen: {hdr.contents.caplen}")
    debug(f"c: len: {hdr.contents.len}")
    data = ct.string_at(pkt, hdr.contents.caplen)
    #debug(f'data:{data}')
    if q:
        q.put_nowait(data)
        debug(f'c: q.put done: {q.qsize()}')

# start
if __name__ == "__main__":
    try:
        asyncio.run(collector(asyncio.Queue()))
    except KeyboardInterrupt:
        info('exit]')


# -*- coding: utf-8 -*-
"""
Created on Sun Jan 30 22:43:26 2022

@author: dvnny

DESCRIPTION

Practicing web scraping with BeautifulSoup4

Trying to pull coaching information from

"""

from bs4 import BeautifulSoup
import requests
import pandas as pd

# https://www.pro-football-reference.com/teams/nyj/coaches.htm#coaches_year

pfr_teams_site = "https://www.pro-football-reference.com/teams/"
suffix = "/coaches.htm#coaches_year"

team = "nyj"

r = requests.get(pfr_teams_site + team + suffix)

# Previewing info about the response
# print(r.headers)
# print(r.text)

soup = BeautifulSoup(r.text, "html.parser")
# print(soup.table)

# tables = soup.find_all('div', {'table'})
tables = soup.find_all('table')

# print(soup.prettify())

coach_data = tables[1]
print(coach_data)

df = pd.read_html(str(coach_data))[0]

print(df.head())
print(df.info())
print(df.columns)

df.columns = ["year","coach","regular_season_games","wins","losses","ties","win_percentage","playoff_games","playoff_wins","playoff_losses",\
                                  "playoff_result","coordinator_offense","coordinator_defense"]

#df.rename(columns={'Year':"year",'Coach':"coach",'G':"regular_season_games",'W':"win",'L':"loss",'T':"tie",'W-L%':"win_percentage"}, inplace=True)
                   # ,"playoff_games","playoff_win","playoff_loss","playoff_result","coordinator_offense","coordinator_defense"})
# df.astype({"year":"int64"}).dtypes

print(df.info())


from flask import session

from todo_app.data.Item import Item

import os, requests

def get_items():
    """
    Fetches all saved items from the session.

    Returns:
        list: The list of saved items.
    """

    BOARD_ID = os.getenv('CORNDEL_BOARD_ID')
    API = os.getenv('TRELLO_API')
    BOARD = os.getenv('TRELLO_BOARD').replace('{id}', BOARD_ID)

    # https://api.trello.com/1/board/{board_id}/cards
    request_url = API+BOARD+'?key='+os.getenv('TRELLO_API_KEY')+'&token='+os.getenv('TRELLO_TOKEN')
    trello_response = requests.get(request_url)

    return build_card_items(trello_response.json())

def build_card_items(trello_response):
    card_items = []

    for trello_list in trello_response:
        card_items.append(Item.from_trello_card(trello_list))

    return card_items

def get_item(id):
    """
    Fetches the saved item with the specified ID.

    Args:
        id: The ID of the item.

    Returns:
        item: The saved item, or None if no items match the specified ID.
    """
    items = get_items()
    return next((item for item in items if item.id == id), None)

def add_item(title, list_id):
    """
    Adds a new item with the specified title to the session.

    Args:
        title: The title of the card.
        list_id: The list to add the card to

    Returns:
        item: The saved item.
    """
    # Add the item to the list
    API = os.getenv('TRELLO_API')
    CARD = os.getenv('TRELLO_CARDS')
    request_url = API+CARD+'?key='+os.getenv('TRELLO_API_KEY')+'&token='+os.getenv('TRELLO_TOKEN')

    trello_response = requests.post(request_url, data={'name': title, 'idList': list_id})

    return trello_response

def save_item(card_id, list_id):
    """
    Updates an existing item in the session. If no existing item matches the ID of the specified item, nothing is saved.

    Args:
        card_id: The card_id to move to the new list.
        list_id: The list id to move the card to
    """

    API = os.getenv('TRELLO_API')
    CARD = os.getenv('TRELLO_CARD').replace('{id}', card_id)

    # https://api.trello.com/1/cards/{card_id}?key=....
    request_url = API+CARD+'?key='+os.getenv('TRELLO_API_KEY')+'&token='+os.getenv('TRELLO_TOKEN')

    trello_response = requests.put(request_url, data={'idList': list_id})

    return trello_response


# 11. 盛最多水的容器

"""
给你 n 个非负整数 a1，a2，...，an，每个数代表坐标中的一个点 (i, ai) 。在坐标内画 n 条垂直线，垂直线 i 的两个端点分别为 (i, ai) 和 (i, 0)。找出其中的两条线，使得它们与 x 轴共同构成的容器可以容纳最多的水。

说明：你不能倾斜容器，且 n 的值至少为 2。

 垂直线代表输入数组 [1,8,6,2,5,4,8,3,7]。在此情况下，容器能够容纳水（表示为蓝色部分）的最大值为 49。

示例：

输入：[1,8,6,2,5,4,8,3,7]
输出：49
"""
import json

class Solution:
    def maxArea(self, height: list) -> int:
        if not height:
            return 0
        left, right = 0, len(height) - 1
        res = 0
        while left < right:
            res = max(res, (right-left) * min(height[left], height[right]))
            if height[left] < height[right]:
                left += 1
            else:
                right -= 1
        return res

def stringToIntegerList(input):
    return json.loads(input)

def main():
    import sys
    import io
    def readlines():
        for line in io.TextIOWrapper(sys.stdin.buffer, encoding='utf-8'):
            yield line.strip('\n')

    lines = readlines()
    while True:
        try:
            line = next(lines)
            height = stringToIntegerList(line);

            ret = Solution().maxArea(height)

            out = str(ret);
            print(out)
        except StopIteration:
            break

if __name__ == '__main__':
    main()


from typing import Any
from absl.testing import parameterized

import tensorflow as tf
from tensorflow_gnn.models.contrastive_losses.deep_graph_infomax import layers

class DeepGraphInfomaxLogitsTest(tf.test.TestCase, parameterized.TestCase):

  @parameterized.named_parameters([
      dict(
          testcase_name="NonSequenceInput",
          inputs=tf.constant(range(8)),
          expected_error=r"Expected `Sequence` \(got .*\)",
      ),
      dict(
          testcase_name="InvalidSequenceInput",
          inputs=[tf.constant(range(8))] * 4,
          expected_error=r"Expected `Sequence` of length 2 \(got .*\)",
      ),
      dict(
          testcase_name="IncompatabileShapeInput",
          inputs=[tf.constant(range(8)), tf.constant(range(4))],
          expected_error=r"Shapes \(8,\) and \(4,\) are incompatible",
      ),
      dict(
          testcase_name="UndefinedInnerDimension",
          inputs=[tf.keras.Input((2, None)), tf.keras.Input((2, None))],
          expected_error=r"Expected a defined inner dimension \(got .*\)",
      ),
  ])
  def test_value_error(
      self,
      inputs: Any,
      expected_error: str):
    """Verifies invalid input raises `ValueError`."""
    with self.assertRaisesRegex(ValueError, expected_error):
      _ = layers.DeepGraphInfomaxLogits()(inputs)

  def test_deep_graph_infomax(self):
    """Verifies output logits shape."""
    x_clean = tf.random.uniform((1, 4))
    x_corrupted = tf.random.uniform((1, 4))
    logits = layers.DeepGraphInfomaxLogits()((x_clean, x_corrupted))

    # Clean and corrupted logits (i.e., shape (1, 2)).
    self.assertEqual(logits.shape, (1, 2))

if __name__ == "__main__":
  tf.test.main()


from fastapi.testclient import TestClient

from test.endpoints.sessions.conftest import Session

class TestDeleteSession:
    def test_delete_session(
        self,
        single_user_session,
        test_app: TestClient,
        login_and_get_token,
    ):
        test_response = test_app.delete(
            f"/sessions/{single_user_session['_id']}",
            headers={"Authorization": f"Bearer {login_and_get_token}"},
        )

        assert test_response.status_code == 204

        test_session = Session()
        deleted_session = test_session.find_one(single_user_session["_id"])

        assert deleted_session is None


from events import EventClass
from PySide6.QtWidgets import QFileDialog
from UI.Templates import customFileDialog
import os

class MenuHandler(EventClass):
    def __init__(self, window):
        self.handler = window.handler
        self.window = window
        self.connectActions()

    def connectActions(self):
        handler, window = (self.handler, self.window)
        mb = window.menuBar()

        # FILE
        File = mb.addMenu("&File")
        File.addAction("Save", self.onSave, "Ctrl+s")
        File.addAction("Load", self.onLoad, "Ctrl+l")

        File.addAction("Load Dataset", self.onDatasetLoad, "Ctrl+d")
        File.addAction("Load Model", self.onModelLoad, "Ctrl+e")

        File.addAction("Load Zero Model", self.loadZeroModel, "Ctrl+0")
        File.addAction("Load Prediction", self.loadPrepredictedModel)

        # File.addAction("Preferences", self.onPreferences)
        # File.addAction("Exit", self.onExit)

        # LOUPE
        Loupe = mb.addMenu("&Loupe")
        Loupe.addAction("New", self.newLoupe, "Ctrl+n")

    def onSave(self):
        (path, _) = QFileDialog.getSaveFileName(self.handler.window)
        if path is None or path.strip() == "":
            return

        self.handler.env.newTask(
            self.handler.env.save,
            args=(path,),
            visual=True,
            name=f"Saving at {os.path.basename(path)}",
            threaded=True,
        )

    def onLoad(self):
        path = QFileDialog.getExistingDirectory(self.handler.window)
        if path is None or path.strip() == "":
            return

        self.handler.env.newTask(
            self.handler.env.load,
            args=(path,),
            visual=True,
            name=f"Loading {os.path.basename(path)}",
            threaded=True,
        )

    def onPreferences(self):
        pass

    def onExit(self):
        self.eventPush("QUIT_EVENT")

    def onDatasetLoad(self):
        env = self.handler.env
        fileTypes = list(env.datasetTypes.keys())
        extensions = [
            env.datasetTypes[x].datasetFileExtension for x in fileTypes
        ]
        path, typ = customFileDialog(
            self.handler.window, fileTypes=fileTypes, extensions=extensions
        )

        env.taskLoadDataset(path, typ)

    def onModelLoad(self):
        env = self.handler.env
        fileTypes = list(env.modelTypes.keys())
        extensions = [env.modelTypes[x].modelFileExtension for x in fileTypes]
        path, typ = customFileDialog(
            self.handler.window, fileTypes=fileTypes, extensions=extensions
        )

        env.taskLoadModel(path, typ)

    def loadPrepredictedModel(self):
        env = self.handler.env
        names = [x.getName() for x in env.getAllDatasets(excludeSubs=True)]
        keys = [x.fingerprint for x in env.getAllDatasets(excludeSubs=True)]
        extensions = ["*.npz"] * len(names)

        path, typ = customFileDialog(
            self.handler.window, fileTypes=names, extensions=extensions
        )
        idx = names.index(typ)
        env.loadPrepredictedDataset(path, keys[idx])

    def newLoupe(self):
        self.handler.newLoupe()

    def loadZeroModel(self):
        env = self.handler.env
        env.loadZeroModel()


import csv
import json
from symbol import except_clause
from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.support.ui import WebDriverWait
from webdriver_manager.chrome import ChromeDriverManager
from selenium.webdriver.common.by import By

import time

browser = webdriver.Chrome(ChromeDriverManager().install())
links = ['https://www.zomato.com/ncr/delivery-in-south-delhi', 'https://www.zomato.com/ncr/delivery-in-north-delhi',
         'https://www.zomato.com/ncr/delivery-in-east-delhi', 'https://www.zomato.com/ncr/delivery-in-west-delhi', 'https://www.zomato.com/ncr/delivery-in-central-delhi']
restaurants = {}
for link in links:
    browser.get(link)
    time.sleep(2)
    restaurants[link] = set()

    rows = browser.find_elements(By.CSS_SELECTOR, ".sc-doWzTn")
    while len(rows) < 17:
        browser.execute_script(
            "window.scrollTo(0, document.body.scrollHeight);")
        time.sleep(2)
        rows = browser.find_elements(By.CSS_SELECTOR, ".sc-1mo3ldo-0")
        print(len(rows))

    # find all links hrefs

    a_links = browser.find_elements_by_tag_name('a')
    # for a in a_links:

    print(len(a_links))
    # filter links having /order in them
    iteration = 0
    for a in a_links:
        print(iteration)
        href = a.get_attribute('href')
        try:
            if 'order' in href:
                restaurants[link].add(href)
        except:
            pass
        iteration += 1
    restaurants[link] = list(restaurants[link])
    # for row in rows:
    #     # print all elements in row
    #     try:
    #         # print(row.find_elements_by_tag_name("a"))
    #         restaurants[link].add(row.find_element_by_tag_name("a").get_attribute("href"))
    #     except:
    #         print("Error")
    #         pass
        # try:
        #     jumbo_tracker = row.find_element(By.CSS_SELECTOR, ".jumbo-tracker")
        #     if jumbo_tracker:
        #         a_link = jumbo_tracker.find_element(By.HTML, "a")
        #         restaurants[link].append(a_link.get_attribute("href"))
        #     else:
        #         print("No jumbo tracker")
        # except:
        #     print("No jumbo tracker")
        #     pass
    # break

print(restaurants)
print(len(restaurants))
# add to json file named zomato.json
with open('zomato.json', 'w') as f:
    json.dump(json.dumps(restaurants), f)

# - Created a web scraper using python and selenium to scrape data from zomato website.
# - The script is able to go to five different locations and scrape the data from the page.
# - Challenges:
#     - The page is dynamic and the data is loaded using lazy loading. We have to scroll down to load more data.
#     - The data is loaded in a grid format and we have to extract the data from each grid.
#     - The scraper needs to visit multiple pages and extract data from each page. This requires a lot of time

with open('zomato.json') as f:
    data = json.load(f)
    print(data)

browser = webdriver.Chrome(ChromeDriverManager().install())
restaurants = {}
for link in data:
    restaurants[link] = []
    for href in data[link]:
    # href = data[link][0]
        browser.get(href)
        time.sleep(2)
        try:
            name = browser.find_element(by=By.XPATH, value='/html/body/div[1]/div/main/div/section[3]/section/section/div/div/div/h1').text
        except:
            name = ""
        try:
            rating = browser.find_element(by=By.XPATH, value='/html/body/div[1]/div/main/div/section[3]/section/section/div/div/div/section/div[3]/div[1]/div/div/div[1]').text
        except:
            rating = "0"
        browser.find_element(by=By.XPATH, value='/html/body/div[1]/div/main/div/article/div/section/section/div[1]/h2/a').click()
        time.sleep(2)
        try:
            address = browser.find_element(by=By.XPATH, value='/html/body/div[1]/div/main/div/section[4]/section/article/section/p').text
        except:
            address = ""
        try:
            phone = browser.find_element(by=By.XPATH, value='/html/body/div[1]/div/main/div/section[4]/section/article/p').text.split('#')[0]
        except:
            phone = ""
        try:
            popularDishes = browser.find_element(by = By.XPATH, value='/html/body/div[1]/div/main/div/section[4]/section/section/article[1]/section[1]/p[1]').text
            popularDishes = popularDishes.split(', ')
        except:
            popularDishes = []
        try:
            knownFor = browser.find_element(by = By.XPATH, value='/html/body/div[1]/div/main/div/section[4]/section/section/article[1]/section[1]/p[2]').text
            knownFor = knownFor.split(', ')
        except:
            knownFor = []
        try:
            avgCost = browser.find_element(by = By.XPATH, value='/html/body/div[1]/div/main/div/section[4]/section/section/article[1]/section[1]/p[3]').text.split(' ')[0].split('₹')[1]
        except:
            avgCost = ""
        try:
            cuisineSection = browser.find_element(by = By.XPATH, value='/html/body/div[1]/div/main/div/section[4]/section/section/article[1]/section[1]/section[3]')
            cuisine = cuisineSection.find_elements(by = By.TAG_NAME, value='a')
            cuisine = [c.text for c in cuisine]
        except:
            cuisine = []
        try:
            reviewSection = browser.find_element(by=By.XPATH, value='/html/body/div[1]/div/main/div/section[4]/section/section/section[2]/div')
            reviewHighlights = reviewSection.find_elements(by=By.TAG_NAME, value='span')
            reviewHighlights = [r.text for r in reviewHighlights]
            if 'Write a Review' in reviewHighlights:
                reviewHighlights.remove('Write a Review')
        except:
            reviewHighlights = []
        print(name, rating, address, phone, popularDishes, knownFor, avgCost, cuisine, reviewHighlights)
        # write to json file
        restaurants[link].append({
            'name': name,
            'rating': rating,
            'address': address,
            'phone': phone,
            'popularDishes': popularDishes,
            'knownFor': knownFor,
            'avgCost': avgCost,
            'cuisine': cuisine,
            'reviewHighlights': reviewHighlights
        })
        with open('zomatoData.json', 'w') as f:
            json.dump(restaurants, f)
        # write to zomato.csv
    # with open('zomato.csv', 'a', encoding='utf-8') as f:
    #     writer = csv.writer(f)
    #     writer.writerow([name, rating, address, phone, popularDishes, knownFor, avgCost, cuisine, reviewHighlights])


from time import time
from collections import defaultdict
from collections import deque
from collections import namedtuple
from os.path import basename
from os.path import splitext
from logging import getLogger

from . import pathmaker
from .archive import ArchiveFile
from .archive import FileNotFoundError
from .exception import BadArchiveFile
from .exception import UnsupportedArchiveFile

logger = getLogger(__name__)

# XXX the i prefix here means archive internal, not for mapper.
FileEntry = namedtuple(
    'FileEntry', ['archive_path', 'ifilename', 'ifile_size'])

class DefaultMapper(object):
    """
    Mapper that tracks the nested structure within archive files.
    """

    def __init__(self, path=None, pathmaker_name='default', _pathmaker=None,
            overwrite=False, include_arcname=False, splitext_arcname=False):
        """
        Initialize the mapping, optionally with a path to an archive
        file.

        Mapping dict keys are names of file or directory, values are
        either a tuple that represent a file, or a dict to represent a
        directory.
        """

        self.include_arcname = include_arcname
        self.splitext_arcname = splitext_arcname
        self.overwrite = overwrite
        if _pathmaker:
            self.pathmaker = _pathmaker
        else:
            self.pathmaker = getattr(pathmaker, pathmaker_name)()

        # The actual filesystem mapping
        self.mapping = {}
        # a mapping with keys of generated paths against source archive.
        # its keys includes a map of directory to its source archive.
        self.reverse_mapping = defaultdict(deque)
        # Tracked file added timestamps (see _load_infolist)
        self.archives = {}
        # A flattened mapping of archive to its list of internal entries
        # including directory entries.
        self.archive_ifilenames = {}

        if path:
            self.load_archive(path)

    def mkdir(self, path_fragments):
        """
        Creates the dir entries identified by path if not already exists
        and return the complete directory.
        """

        # set current to root node
        current = self.mapping

        for c, frag in enumerate(path_fragments):
            if frag in current:
                current = current[frag]
                if not isinstance(current, dict):
                    raise ValueError(
                        'cannot create directory `%(filename)s` at '
                        '`%(path)s/`: file entry exists.' % {
                            'filename': frag,
                            'path': '/'.join(path_fragments[:c]),
                        }
                    )
            else:
                # create directory dict entry and set current.
                current[frag] = current = {}

        return current

    def traverse(self, path):
        """
        Traverse to path, or return the entry identified by path.
        """

        path_fragments = path and path.split('/') or []
        return self._traverse(path_fragments)

    def _traverse(self, path_fragments):
        current = self.mapping

        for frag in path_fragments:
            if not isinstance(current, dict) or frag not in current:
                # No such frag in dir.
                return None
            current = current[frag]

        return current

    def _load_infolist(self, archive_path, infolist):
        self.archives[archive_path] = time()
        archive_name = (
            splitext(basename(archive_path))[0]
            if self.splitext_arcname else
            basename(archive_path)
        ) + '/'
        self.archive_ifilenames[archive_path] = i_filenames = []

        for info in infolist:
            if self.include_arcname:
                raw_filename = archive_name + info.filename
            else:
                raw_filename = info.filename
            frags, filename = self.pathmaker(raw_filename)

            # the internal filename.
            ifilename = '/'.join([] + frags + [filename])

            fentry = FileEntry(archive_path, info.filename, info.file_size)

            # These may appear redundant compared to the ones below, but
            # do note these are always added despite conflict that will
            # appear at the filesystem presentation level.  This is used
            # for locating current zip files for regeneration if an
            # archive was removed from this mapping.
            self.reverse_mapping[ifilename].append(fentry)
            i_filenames.append(ifilename)

            try:
                target = self.mkdir(frags)
            except ValueError as e:
                # using info.filename rather than raw_filename because
                # the message is provided by the exception generated by
                # self.mkdir.
                logger.warning(
                    '`%s` could not be created: %s', info.filename, e.args[0])
                continue

            if not filename:
                # was a directory entry
                continue

            if filename in target:
                if not self.overwrite:
                    logger.info('`%s` already exists; ignoring', info.filename)
                    continue
            target[filename] = fentry

    def _unload_functions(self):
        """
        Return a tuple of needed functions for unloading an archive out
        of the mappings.
        """

        if self.overwrite:
            # check the rightmost (newest) item, pop from right
            return -1, deque.pop
        else:
            # check the leftmost (oldest) item, pop from left
            return 0, deque.popleft

    def _unload_infolist(self, archive_path):
        # pop this out right away to mark this as to be pruned off.
        ifilenames = self.archive_ifilenames.pop(archive_path)
        index, pop = self._unload_functions()
        all_frags = set()
        for ifilename in ifilenames:
            # lookup via the reverse mapping to see that this ifilename
            # is the active check that the current active
            fentries = self.reverse_mapping[ifilename]

            if fentries[index].archive_path != archive_path:
                # leave the reverse mapping in place.
                continue

            original_filename = fentries[index].ifilename

            archive_ifilenames = self.archive_ifilenames
            # We have a match, time to clean up to the latest fresh and
            # active entry that other previous calls to this function
            # have left behind.
            fentry_replacement = None
            try:
                while fentries[index].archive_path not in archive_ifilenames:
                    pop(fentries)
                else:
                    # Successfully left the loop, the fileentry is now
                    # the replacement.
                    fentry_replacement = fentries[index]
            except IndexError:
                # This no longer exists in any active archive.
                self.reverse_mapping.pop(ifilename)

            # We have an ifilename created in _load_infolist, invert
            # operation to derive filename and frags.
            raw = ifilename.split('/')
            filename = raw[-1]
            frags = raw[:-1]

            if frags:
                all_frags.add(tuple(frags))

            info = self._traverse(frags)
            if not info is None:
                # The file's directory may not have been added to
                # self.mapping, if its creation may have been blocked
                # by another file.
                if filename in info:
                    # Only remove, for the same reason as above.
                    if (original_filename.endswith('/') ==
                            isinstance(info[filename], dict)):
                        # Only if the types are the same type should
                        # this entry be removed (file for file, dir for
                        # dir).
                        info.pop(filename)

                    if fentry_replacement:
                        # Restore the fileentry with the replacement
                        # from above.
                        info[filename] = fentry_replacement

                # XXX filename could be empty, does it mean the
                # directory can be removed if empty?

            # if ifilename_paths is not empty, it may be possible to
            # restore that entry with the newer (or previous) "version"
            # from the archive, if it also resolves to a file, but this
            # will need more thought to do, given that files and dirs
            # are two different types.

        # finally, purge all empty directories.  Yes this includes
        # directories that may not be wholly owned by this archive.
        for frags in sorted(all_frags, key=lambda x: len(x)):
            target = self._traverse(frags[:-1])
            if target.get(frags[-1]) == {}:
                target.pop(frags[-1])

        # discard the date associated with this archive path too.
        self.archives.pop(archive_path)

    def load_archive(self, archive_path):
        """
        Load an archive file identified by archive_path into the
        mapping.
        """

        try:
            with ArchiveFile(archive_path) as af:
                self._load_infolist(archive_path, af.infolist())
            logger.info('loaded `%s`', archive_path)
            return True
        except BadArchiveFile:
            logger.warning(
                '`%s` appears to be an invalid archive file', archive_path)
        except UnsupportedArchiveFile as e:
            logger.warning('`%s` %s', archive_path, str(e))
        except FileNotFoundError:
            logger.warning(
                '`%s` does not exist.', archive_path)
        except:
            logger.exception('Exception')
        return False

    def unload_archive(self, archive_path):
        self._unload_infolist(archive_path)
        logger.info('unloaded `%s`', archive_path)

    def open(self, path):
        info = self.traverse(path)
        if info is None:
            return
        archive_path, filename, _ = info
        # it is possible to return those values, but given that the
        # underlying files can change, or that new stack comes in, it's
        # best not to directly expose this.
        try:
            with ArchiveFile(archive_path) as zf:
                # zf does remain open because of this.
                return (id(info), zf.open(filename))
        except BadArchiveFile:  # pragma: no cover
            logger.warning(
                '`%s` became an invalid archive file', archive_path)
        except FileNotFoundError:  # pragma: no cover
            # shoud unload the archive?
            logger.warning(
                '`%s` no longer exists.', archive_path)
        except:  # pragma: no cover
            logger.exception('Exception')
        return False

    def readfile(self, path):
        """
        Return the complete file with information contained in path.
        """

        # This is an alternative implementation that returns the entire
        # contents.  Should include similar exception handling as above,
        # but that pattern should really be generalized.

        info = self.traverse(path)

        archive_filename, filename, _ = info
        with ArchiveFile(archive_filename) as af:
            with af.open(filename) as f:
                return f.read()

    def readdir(self, path):
        """
        Return a listing of all files in a directory
        """

        info = self.traverse(path)
        if not isinstance(info, dict):
            return []
        return list(info.keys())


"""Train lib function."""
import datetime
import functools
import math
import re
import time
from contextlib import nullcontext

# pylint: disable=fixme,import-error,import-outside-toplevel,invalid-name,no-name-in-module,wrong-import-order
import numpy as np
import torch
import torch.distributed as dist
import torch.sagemaker as tsm
import torch.utils.data
import transformers
from accelerate import init_empty_weights
from checkpoints import (
    _CHECKPOINT_DIR_REGEX,
    _DEFAULT_STATE_DICT_TYPE,
    CheckpointingMethod,
    get_coordinator_rank,
    is_action_rank,
    load_checkpoint,
    save_checkpoint,
)
from data.pipelines import GPTDataPipeline, create_data_pipeline
from fsdp_utils import get_backward_fetch_policy, get_sharding_strategy, get_transformer_layer
from logging_utils import (
    create_args_table,
    get_logger,
    log_and_write_eval_metrics,
    log_train_metrics,
    show_env_vars,
    write_metrics_train_step,
)
from memory_tracker import memory_status, memory_status_cpu
from packaging import version as pversion
from torch import optim
from torch.distributed.elastic.multiprocessing.errors import record
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy
from torch.sagemaker import transform
from torch.sagemaker.delayed_param import DelayedParamIniter
from torch.sagemaker.distributed.fsdp.fully_sharded_data_parallel import (
    register_tp_processing_pre_comm_hook,
)
from torch.sagemaker.grad_norm import clip_grad_norm_
from torch.sagemaker.utils import utils as tsm_utils  # pylint: disable=no-name-in-module
from train_utils import (
    apply_activation_checkpoint,
    compute_num_params,
    compute_tflops,
    create_model,
    get_learning_rate_scheduler,
    get_model_config,
    get_param_groups_by_weight_decay,
    patch_neox_rope,
)
from transformers import set_seed

logger = get_logger()

def finetune_with_pretrained_weights_check(args):
    # returns True for start of finetuning only
    return args.hf_pretrained_model_name_or_dir is not None and args.resume_from_checkpoint is None

def finetune_check(args):
    # returns True for start of finetuning as well as resuming
    return args.hf_pretrained_model_name_or_dir is not None

def eval_model(model, data_pipeline, num_batches):
    """Eval step."""
    model = model.eval()
    n_batches = 0
    loss = 0.0

    with torch.no_grad():
        for batch_idx, input_data in enumerate(data_pipeline.val_dataloader):
            input_ids, mask = data_pipeline.get_val_batch(input_data)

            if batch_idx >= num_batches:
                break

            loss += model(input_ids=input_ids, attention_mask=None, labels=input_ids)["loss"]
            n_batches += 1

    if n_batches > 0:
        detached_loss = loss.detach()
        torch.distributed.all_reduce(detached_loss)
        loss = detached_loss.item() / dist.get_world_size()
        loss /= n_batches
        ppl = math.exp(loss)
    else:
        loss = -1.0
        ppl = -1.0

    return loss, ppl

def reduce_loss(loss):
    loss_detached = loss.detach()
    dist.all_reduce(loss_detached)
    loss_scalar = loss_detached.item() / dist.get_world_size()
    return loss_scalar

def train_step(
    args, batch_idx, nvtx_warmup_iters, data_pipeline, input_data, model, optimizer, lr_scheduler
):
    if batch_idx >= nvtx_warmup_iters:
        torch.cuda.nvtx.range_push(f"iteration{batch_idx}")

    input_ids, mask, labels = data_pipeline.get_batch(input_data)

    if batch_idx == 0:
        # checking only on batch 0 to reduce checks during runtime
        assert (
            input_ids.shape[1] == args.max_context_width
        ), f"Input data passed {input_ids.shape} does not respect max_context_width set. Note that this is not strictly necessary, but added to prevent mistakes. If you intend to do this, please remove this check."
        assert (
            input_ids.shape[1] <= args.max_context_width
        ), "Input data passed is larger than max_context_width for model. You need to change max_context_width so model can expect larger sequences"

    optimizer.zero_grad(set_to_none=True)

    torch.cuda.synchronize()
    step_start = time.time()

    if batch_idx >= nvtx_warmup_iters:
        torch.cuda.nvtx.range_push("forward")

    # uses default causal mask
    loss = model(input_ids=input_ids, attention_mask=None, labels=labels)["loss"]

    if batch_idx >= nvtx_warmup_iters:
        # for forward
        torch.cuda.nvtx.range_pop()

    if args.enable_memory_profiling > 0 and batch_idx < 5:
        memory_status_cpu("After_forward_cpu")
        memory_status(msg="After_forward")

    if batch_idx >= nvtx_warmup_iters:
        torch.cuda.nvtx.range_push("backward")

    loss.backward()

    if batch_idx >= nvtx_warmup_iters:
        # for backward
        torch.cuda.nvtx.range_pop()

    if args.enable_memory_profiling > 0 and batch_idx < 5:
        memory_status_cpu("After_train_step_cpu")
        memory_status(msg="After_train_step")

    if batch_idx >= nvtx_warmup_iters:
        torch.cuda.nvtx.range_push("opt_step")

    grad_norm = clip_grad_norm_(model, args.grad_clip)
    optimizer.step()
    lr_scheduler.step()

    if batch_idx >= nvtx_warmup_iters:
        # for opt step
        torch.cuda.nvtx.range_pop()

    if args.clean_cache > 0:
        # empty the cache to avoid OOM
        torch.cuda.empty_cache()

    if batch_idx >= nvtx_warmup_iters:
        # for step
        torch.cuda.nvtx.range_pop()

    torch.cuda.synchronize()
    step_time = time.time() - step_start

    if args.enable_memory_profiling > 0 and batch_idx < 5:
        memory_status(msg="After_opt_step")

    batch_num_sequences = input_ids.shape[0]
    batch_seqlen = input_ids.shape[1]
    return loss, step_time, batch_num_sequences, batch_seqlen, grad_norm

# pylint: disable=no-member,too-many-arguments,too-many-branches,too-many-locals,too-many-statements
def train(
    model,
    optimizer,
    lr_scheduler,
    writers,
    model_config,
    start_epoch,
    start_train_path_index,
    resume_from_sequence_number,
    num_params,
    total_steps,
    args,
    global_rank,
    world_size,
    checkpointing_pg_metadata,
):
    """Train."""
    if args.enable_memory_profiling > 0:
        memory_status_cpu(msg="before train step")

    model.train()
    dp_rank = global_rank
    dp_size = world_size

    if args.tensor_parallel_degree > 1:
        dp_rank //= args.tensor_parallel_degree
        dp_size //= args.tensor_parallel_degree

    if global_rank == 0:
        logger.info("Creating train dataloader")

    throughputs = []
    # Set the same seed for computation
    set_seed(args.seed)

    data_pipeline = create_data_pipeline(
        args, start_train_path_index, resume_from_sequence_number, dp_rank, dp_size
    )
    cur_seq_index = resume_from_sequence_number
    epoch = start_epoch
    while total_steps < args.max_steps:
        nvtx_warmup_iters = 3
        if global_rank == 0:
            logger.info(f"Starting training with epoch {epoch}")

        # additional loop around is for GPTDataset as there can be multiple dataloaders
        if isinstance(data_pipeline, GPTDataPipeline):
            # with new path if incremented at the end of this for loop
            data_pipeline.create_train_dataset()

        for batch_idx, input_data in enumerate(data_pipeline.train_dataloader):
            if total_steps >= args.max_steps:
                break

            if args.profile_nsys > 0 and batch_idx == nvtx_warmup_iters:
                torch.cuda.cudart().cudaProfilerStart()

            loss, step_time, batch_num_sequences, batch_seqlen, grad_norm = train_step(
                args,
                batch_idx,
                nvtx_warmup_iters,
                data_pipeline,
                input_data,
                model,
                optimizer,
                lr_scheduler,
            )
            total_steps += 1
            cur_seq_index += batch_num_sequences
            sample_processed = batch_num_sequences * dp_size
            throughput = sample_processed / step_time
            throughputs.append(throughput)

            tflops_per_gpu = compute_tflops(throughput, num_params, world_size, batch_seqlen)

            if not total_steps % args.logging_freq and args.log_reduced_training_loss > 0:
                loss_scalar = reduce_loss(loss)
            else:
                loss_scalar = loss.item()

            current_lr = lr_scheduler.get_lr()
            display_step = total_steps - 1
            if global_rank == 0:
                write_metrics_train_step(
                    writers,
                    display_step,
                    loss_scalar,
                    throughput,
                    tflops_per_gpu,
                    current_lr,
                    grad_norm,
                )
                if not total_steps % args.logging_freq:
                    log_train_metrics(
                        args,
                        total_steps,
                        display_step,
                        loss_scalar,
                        throughput,
                        tflops_per_gpu,
                        current_lr,
                        grad_norm,
                        throughputs,
                        num_params,
                        dp_size,
                        batch_seqlen,
                    )

            # evaluate on validation
            if args.validation_freq and not total_steps % args.validation_freq:
                cur_state = np.random.get_state()
                torch.cuda.empty_cache()
                val_loss, val_ppl = eval_model(model, data_pipeline, args.validation_batches)
                if global_rank == 0:
                    log_and_write_eval_metrics(writers, display_step, val_loss, val_ppl)
                model = model.train()
                if args.preserve_np_state > 0:
                    np.random.set_state(cur_state)

            # checkpoint
            if not total_steps % args.checkpoint_freq[0]:

                if isinstance(data_pipeline, GPTDataPipeline):
                    save_train_path_index = data_pipeline.cur_train_path
                else:
                    save_train_path_index = 0
                save_train_seq_index = cur_seq_index
                # technically we have processed save_train_seq_index sequences in this file
                # and so index to start from is save_train_seq_index
                user_content = {
                    "cli_args": args.__dict__,
                    "model_config": model_config,
                    "num_params": num_params,
                    "total_steps": total_steps,
                    "epoch": epoch,
                    "start_train_path_index": save_train_path_index,
                    "resume_from_sequence_number": save_train_seq_index,
                }

                subdir = f"{args.model_type}-{total_steps}steps"
                if global_rank == 0 and not re.match(_CHECKPOINT_DIR_REGEX, subdir):
                    raise ValueError(
                        f"Please double check hard-coded checkpoint subdir pattern: `{subdir}` "
                        f"not matching `{_CHECKPOINT_DIR_REGEX}`."
                    )

                if args.enable_memory_profiling > 0:
                    msg = f"({_DEFAULT_STATE_DICT_TYPE})"
                    memory_status(msg=f"Before ckpt @{display_step} {msg}")
                save_checkpoint(
                    model,
                    optimizer,
                    lr_scheduler,
                    user_content,
                    get_sharding_strategy(args.sharding_strategy),
                    args.checkpoint_dir[0],
                    subdir,
                    args.num_kept_checkpoints[0],
                    checkpointing_pg_metadata,
                    tensor_parallel_degree=int(args.tensor_parallel_degree),
                    checkpoint_type=args.checkpoint_type,
                )
                if args.enable_memory_profiling > 0:
                    msg = f"({_DEFAULT_STATE_DICT_TYPE})"
                    memory_status(msg=f"After ckpt  @{display_step} {msg}")

        if isinstance(data_pipeline, GPTDataPipeline):
            incremented_in_epoch = data_pipeline.increment_path_in_epoch()
            if not incremented_in_epoch:
                # path index set to 0
                epoch += 1
        else:
            epoch += 1
    # Using median throughput across all steps, could be more robust.
    return total_steps, np.median(throughputs) if throughputs else 0

@record
def main(args):
    """Main function to train GPT."""
    # Sanity check for args.
    # - Checkpoints.
    # TODO(sliuxl): Supporting one single checkpoint dir now, and multiple dirs support is missing.
    ckpt_lens = (
        len(args.checkpoint_dir),
        len(args.checkpoint_freq),
        len(args.num_kept_checkpoints),
    )
    if len(set(ckpt_lens)) != 1:
        raise ValueError(f"Len mismtach for checkpoint dir, freq vs num to keep:  {ckpt_lens}.")

    if args.distributed_backend == "smddp":
        import smdistributed.dataparallel.torch.torch_smddp  # pylint: disable=unused-import

    dist.init_process_group(args.distributed_backend, timeout=datetime.timedelta(seconds=7200))
    global_rank = dist.get_rank()
    device = global_rank % torch.cuda.device_count()
    world_size = dist.get_world_size()

    if args.tensorboard_dir and global_rank == 0:
        from torch.utils.tensorboard import SummaryWriter

        logger.info("Writing metrics for tensorboard to %s.", args.tensorboard_dir)
        writers = tuple(SummaryWriter(log_dir=tb_dir) for tb_dir in args.tensorboard_dir)
        table_str = create_args_table(args.__dict__)
        for writer in writers:
            writer.add_text("Arguments", table_str)
    else:
        writers = ()

    smp_config_dict = {
        "activation_loading_horizon": args.activation_loading_horizon,
        "sm_activation_offloading": args.offload_activations > 0,
    }
    if args.shard_degree is not None:
        smp_config_dict["hybrid_shard_degree"] = args.shard_degree

    smp_config_dict["tensor_parallel_degree"] = args.tensor_parallel_degree
    if args.tensor_parallel_degree > 1 and args.use_smp_implementation < 1:
        args.use_smp_implementation = 1
        if global_rank == 0:
            logger.info(
                f"Tensor parallelism (TP) is enabled as tensor_parallel_degree is set to {args.tensor_parallel_degree} (>1). Switching use_smp_implementation to 1 so we can use SMP optimized implementation for TP."
            )
    if args.use_smp_implementation:
        # For our Mem usage fix to TE, this needs to be True
        args.use_orig_params = 1

    if args.use_synthetic_data and args.validation_freq is not None:
        # Overriding validation freq to None as synthetic data
        args.validation_freq = None

    tsm.init(smp_config_dict)
    show_env_vars(0)

    if global_rank == 0:
        for index, (key, value) in enumerate(sorted(args.__dict__.items()), 1):
            logger.info("Arguments [%03d/%03d] %-30s: %s", index, len(args.__dict__), key, value)
        logger.info("Transformers version: %s", transformers.__version__)
        logger.info("World size = %d: # nodes = %d.", world_size, world_size / 8)

        gbs = (
            world_size
            * args.max_context_width
            * args.train_batch_size
            / args.tensor_parallel_degree
        )
        logger.info("Global batch size in tokens: %10d (%5.2fM).", gbs, gbs / 1e6)

    set_seed(args.seed)

    if args.enable_memory_profiling > 0:
        memory_status_cpu(msg="before model creation")

    if args.bf16:
        dtype = torch.bfloat16
    else:
        dtype = torch.get_default_dtype()

    if finetune_check(args):
        from transformers import AutoConfig

        # Using config for finetune mode, else uses args to create model
        model_config = AutoConfig.from_pretrained(args.hf_pretrained_model_name_or_dir)
    else:
        model_config = get_model_config(args)

    if pversion.parse(transformers.__version__) >= pversion.parse("4.34.0"):
        model_config._flash_attn_2_enabled = True

    delayed_param_initer = None
    with tsm_utils.timeit(True, "Model creation", global_rank):
        if args.delayed_param:
            if finetune_with_pretrained_weights_check(args) and dist.get_rank() == 0:
                # create model with pretrained weights on one rank even if we want to use
                # delayed param, param init on other ranks will still be delayed
                model = create_model(
                    args,
                    model_config=model_config,
                    dtype=dtype,
                    pretrained_model_weights=args.hf_pretrained_model_name_or_dir
                    if finetune_with_pretrained_weights_check(args)
                    else None,
                )
                num_params = compute_num_params(model)
            else:
                with init_empty_weights():
                    model = create_model(
                        args,
                        model_config=model_config,
                        dtype=dtype,
                    )
                num_params = compute_num_params(model)
            if finetune_check(args):
                dist.barrier()
        else:
            model = create_model(
                args,
                model_config=model_config,
                dtype=dtype,
                pretrained_model_weights=args.hf_pretrained_model_name_or_dir
                if finetune_with_pretrained_weights_check(args) and dist.get_rank() == 0
                else None,
            )
            num_params = compute_num_params(model)

        if args.use_smp_implementation:
            load_state_dict_from_rank0 = (
                True if finetune_with_pretrained_weights_check(args) else False
            )
            model = transform(model, load_state_dict_from_rank0=load_state_dict_from_rank0)

        if args.delayed_param:
            # param init fn for delayed param creation
            if finetune_check(args):
                if dist.get_rank() != 0:
                    delayed_param_initer = DelayedParamIniter(model)
            else:
                delayed_param_initer = DelayedParamIniter(model)

    assert set(x.dtype for x in model.parameters()) == set(
        [torch.float32]
    ), "Model parameters should be in fp32 for FSDP mixed precision"

    start = time.time()

    if global_rank == 0:
        logger.info(
            "Created model with total parameters: %d (%.2f B)", num_params, num_params * 1e-9
        )

    transformer_layer = get_transformer_layer(args.model_type, args.use_smp_implementation)

    if args.auto_wrap_policy == "transformer_auto_wrap_policy":
        gpt_auto_wrap_policy = functools.partial(
            transformer_auto_wrap_policy,
            transformer_layer_cls={
                transformer_layer,
            },
        )
    elif args.auto_wrap_policy == "size_based_auto_wrap_policy":
        gpt_auto_wrap_policy = functools.partial(
            size_based_auto_wrap_policy,
        )

    torch.cuda.set_device(device)
    if args.bf16:
        # buffer set to fp32 as some models in HF such as llama hard code buffers to fp32
        # to be similar with that we set this to fp32
        buffer_dtype = torch.float32 if args.use_smp_implementation else dtype
        mixed_precision_policy = MixedPrecision(
            param_dtype=dtype, reduce_dtype=dtype, buffer_dtype=buffer_dtype
        )
    else:
        mixed_precision_policy = None

    if args.enable_memory_profiling > 0:
        memory_status_cpu(msg="before fsdp wrapper")

    sharding_strategy = get_sharding_strategy(args.sharding_strategy)

    with (
        delayed_param_initer.validate_params_and_buffers_inited()
        if (delayed_param_initer and not finetune_with_pretrained_weights_check(args))
        else nullcontext(),
        tsm_utils.timeit(True, "FSDP constructor", global_rank),
    ):
        model = FSDP(  # pylint: disable=unexpected-keyword-arg
            model,
            auto_wrap_policy=gpt_auto_wrap_policy,
            mixed_precision=mixed_precision_policy,
            sharding_strategy=sharding_strategy,
            backward_prefetch=get_backward_fetch_policy(args.backward_fetch_policy),
            forward_prefetch=args.forward_prefetch,
            limit_all_gathers=args.limit_all_gathers,
            device_id=torch.cuda.current_device(),
            use_orig_params=args.use_orig_params > 0,
            param_init_fn=delayed_param_initer.get_param_init_fn()
            if delayed_param_initer
            else None,
            post_param_init_fn=delayed_param_initer.get_post_param_init_fn()
            if delayed_param_initer
            else None,
            sync_module_states=True if finetune_with_pretrained_weights_check(args) else False,
        )
    # Barrier is a workaround to reduce extra memory usage with SMDDP backend
    # after the broadcast that happens when we use sync_module_states
    # This can be removed once the SMDDP issue is fixed
    dist.barrier()
    memory_status(msg="After FSDP")

    if global_rank == 0:
        logger.info("Wrapped model with FSDP")

    if args.enable_memory_profiling > 0:
        memory_status(msg="after fsdp wrapper")

    if args.activation_checkpointing > 0:
        apply_activation_checkpoint(args, model=model)

    if args.offload_activations > 0:
        from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import offload_wrapper

        model = offload_wrapper(model)

        # Patch RoPE for GPT NEoX where they are created on Host to move them to Device
        if args.model_type == "gpt_neox" and args.patch_neox_rope > 0:
            patch_neox_rope(model)

    param_groups = get_param_groups_by_weight_decay(model)

    optimizer = optim.AdamW(
        param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
    )

    if global_rank == 0:
        logger.info("Created optimizer")

    lr_scheduler = get_learning_rate_scheduler(optimizer, args)

    checkpointing_pg_metadata = (
        model.process_group,
        get_coordinator_rank(model.process_group),
        is_action_rank(global_rank),
    )

    if args.resume_from_checkpoint:
        (
            model,
            optimizer,
            lr_scheduler,
            epoch,
            total_steps,
            start_train_path_index,
            resume_from_sequence_number,
        ) = load_checkpoint(
            args,
            model,
            optimizer,
            lr_scheduler,
            args.resume_from_checkpoint,
            sharding_strategy,
            checkpointing_pg_metadata,
            tensor_parallel_degree=int(args.tensor_parallel_degree),
            checkpoint_type=args.checkpoint_type,
        )

    else:
        total_steps = 0
        epoch = 0
        start_train_path_index = 0
        resume_from_sequence_number = 0

    start = time.time()
    # total_steps, throughput, loss
    total_steps, _ = train(
        model,
        optimizer,
        lr_scheduler,
        writers,
        model_config,
        epoch,
        start_train_path_index,
        resume_from_sequence_number,
        num_params,
        total_steps,
        args,
        global_rank,
        world_size,
        checkpointing_pg_metadata,
    )
    time_to_train = time.time() - start

    dist.barrier()

    if args.save_final_model:
        save_checkpoint(
            model,
            None,
            None,
            {"model_config": model_config},
            None,
            args.model_dir if args.model_dir is not None else args.checkpoint_dir[0],
            "" if args.model_dir is not None else "model",
            1,
            None,
            int(args.tensor_parallel_degree),
            checkpoint_type=CheckpointingMethod.FULL,
        )

    if global_rank == 0:
        logger.info(
            "FSDP training finished successfully %fs (%fmin).", time_to_train, time_to_train / 60.0
        )
    dist.destroy_process_group()


from collections import deque
from aiomeasures.events import Event

class Collector(deque):
    """Caped list of metrics
    """

    def flush(self, rate=None, formatter=None):
        while True:
            try:
                metric = self.popleft()
                if isinstance(metric, Event):
                    yield formatter(metric)
                    continue
                if metric.value is None:
                    continue
                if rate and metric.rate and rate < metric.rate:
                    continue
                if formatter:
                    try:
                        yield formatter(metric)
                    except ValueError:
                        continue
                else:
                    yield metric
            except IndexError:
                raise StopIteration


"""
SuperScript network module
===========

Classes:
    SocialNetwork
        Class for tracking social network of collaborations
"""

import networkx as nx
from itertools import combinations
from mesa.space import NetworkGrid

from .config import HISTORICAL_SUCCESS_RATIO_THRESHOLD

class SocialNetwork(NetworkGrid):
    """Social network class.

    Is a subclass of mesa.space.NetworkGrid

    Edges in this network count the number of successful collaborations
    between each pair of workers.
    ...

    Attributes:
        model: model.SuperScriptModel
            Reference to main model.
        G: nx.Graph
            Networkx x graph object, created in main model.
        success_threshold: float
            Fraction of team members who need to have worked before on
            successful project for get_team_historical_success_flag to
            return True.
    """

    def __init__(self, model, G,
                 success_threshold=HISTORICAL_SUCCESS_RATIO_THRESHOLD):

        self.model = model
        self.G = G
        self.success_threshold = success_threshold

    def initialise(self):
        """Create the network.

        Adds all workers as nodes in the network.
        Calls base class constructor and calls base class method
        to place each worker on their node.
        """

        for worker in self.model.schedule.agents:
            self.G.add_node(
                worker.worker_id,
                department=worker.department.dept_id
            )
        super().__init__(self.G)

        for worker, node in zip(self.model.schedule.agents, self.G.nodes()):
            self.place_agent(worker, node)

    def replace_worker(self, old_worker, new_worker):
        """Replaces worker in the network (called when a worker is
        replaced due to inactivity.

        Args:
            old_worker: worker.Worker
                Worker being replace.
            new_worker: worker.Worker
                Worker replacing old worker.
        """

        self._remove_agent(old_worker, old_worker.worker_id)
        self.G = nx.relabel_nodes(
            self.G,
            {old_worker.worker_id: new_worker.worker_id},
            copy=False
        )
        self.place_agent(new_worker, new_worker.worker_id)

    def add_team_edges(self, team):
        """Adds all edges between pairs of workers in this team.

        If an edge already exists, its weight is incremented by 1.

        Args:
            team: organisation.Team
                Provides list of members that have successfully
                collaborated on a project and are to be added
                to the network.
        """

        pairs = list(combinations(team.members.keys(), 2))
        for pair in pairs:

            if (pair[0], pair[1]) not in self.G.edges():
                self.G.add_edge(pair[0], pair[1], weight=1)
            else:
                self.G[pair[0]][pair[1]]['weight'] += 1

    def get_team_historical_success_flag(self, team):
        """Determine is team meets required threshold for previous
        successful collaboration.

        Args:
            team: organisation.Team

        Returns:
            bool: True if more than self.success_threshold of possible
                  edges are realised in this network.
        """

        pairs = list(combinations(team.members.keys(), 2))
        success_ratio = 0
        for pair in pairs:
            if (pair[0], pair[1]) in self.G.edges():
                success_ratio += 1

        success_ratio = (
            success_ratio / len(pairs) if len(pairs) > 0 else 0.0
        )
        if success_ratio >= self.success_threshold:
            return True
        else:
            return False

    def save(self):
        """Save network for later analysis."""
        nx.write_gpickle(
            self.G,
            self.model.io_dir
            + '/network_timestep_%d.gpickle' % self.model.time
        )


from django import template
from brandapp.models import *

from django.contrib.auth.models import User
register = template.Library()

@register.filter
def check_temp(val):
	try:
		temp=emailTemplate.objects.get(category=val)
		print("False")
		return False
	except Exception as e:
		print("True")
		return True


import numpy as np
from keras.datasets import mnist

(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

train_images = train_images.reshape(train_images.shape[0], -1)  # (60000, 784)
train_images = train_images.astype('float32') / 255

test_images = test_images.reshape(test_images.shape[0], -1)
test_images = test_images.astype('float32') / 255

# 将标签数据转为int32 并且形状为(60000,1)
train_labels = train_labels.astype(np.int32)
test_labels = test_labels.astype(np.int32)
train_labels = train_labels.reshape(-1, 1)
test_labels = test_labels.reshape(-1, 1) #(10000,1)

test_num = [0,0,0,0,0,0,0,0,0,0]
for n in range(0,10):
    test_num[n] = np.sum(test_labels == n)
print(test_num)
print(np.sum(test_num))


import numpy as np
from scipy.ndimage import rotate
from scipy.signal import convolve2d
import math
from .dataReader import DataReader
import imutils
from PIL import Image
from math import ceil
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA

def conv2(x, y, mode='same'):
    """
    Convoles two images similar to MATLAB's convolution'
    Args:
        x (_numpy_arr_): _First arr to convolve_
        y (_numpy_arr_): _Second array to convolve with_
        mode (str, optional): Defaults to 'same'.

    Returns:
        _type_: _description_
    """
    return np.rot90(convolve2d(np.rot90(x, 2), np.rot90(y, 2), mode=mode), 2)

class OptimizationFunction:
    """
    Optimization reated functions to align the coarse image exactly with fine image
    """
    def optima(self, inputs, coarse_image, fine_image, edge, N):
        """
        RMSE cost Function to be used for optimization and obtain parameters
        Args:
            inputs (_list_): _Swarm positions for optimization_
            coarse_image (_numpy_array_): _Band from coarse image_
            fine_image (_numpy array_): _Band from fine image_
            edge (_int_): _Difference between the size of coarse and fine image_
            N (_int_): _Size of fine image_

        Returns:
            _list_: _Cost from each swarm position_
        """
        inputs = np.array(inputs)
        cost = []

        # Process for each particle position
        for i  in range(len(inputs)):

            # Apply FFT on image and use RMSE for optimization
            filtered_CoarseImage = self.FFT(inputs[i], coarse_image, edge, N)
            cost.append(np.sqrt(np.mean(np.power(filtered_CoarseImage - fine_image, 2))))

        cost = np.array(cost).T
        return cost

    def FFT(self, parameters, image, edge, N):
        """
        FFT function which applies PSF and finds appropraites portion from complete coarse image
        Args:
            parameters (_list_): _Parameters for applying PSF and shifting_
            image (_numpy_arr_): _Numpy array of image_
            edge (_int_): _Difference between the coarse and fine image_
            N (_int_): _Size of the fine image_

        Returns:
            _numpy_arr_: _coarse image which aligns to fine image_
        """
        fwhm_x = parameters[0]
        fwhm_y = parameters[1]
        sh_x = parameters[2]
        sh_y = parameters[3]
        angle = parameters[4]

        # Calculate SMF for given parameters
        psf_cal = self.SMF(fwhm_x, fwhm_y, angle)

        # convolve image with PSF
        image_fft = conv2(image, psf_cal, 'same')
        ends = list(image_fft.shape)

        # Cut the coarse image such that it aligns with the fine image exactly
        image_fft = image_fft[int(edge + sh_y): int(ends[0] - edge + sh_y), int(edge + sh_x): int(ends[1]-edge + sh_x)]
        return image_fft.reshape(1, (N * N))

    def SMF(self, fwhm_x, fwhm_y, angle):
        """
        SMF obtains the PSF and creates the appropraite mask and rotates for it to align exactly
        Args:
            fwhm_x (_int_): _accomodates for PSF in X direction_
            fwhm_y (_int_): _accomodates for PSF in Y direction_
            angle (_int_): _rotates the PSF_

        Returns:
            _numpy_arr_: _rotated PSF array_
        """

        # Obtain PSF
        psf_gaussian = self.PSF(fwhm_x, fwhm_y)
        psf_mask = np.zeros_like(psf_gaussian)

        # PSF Mask
        psf_mask[int(ceil(fwhm_y/2)) : int(ceil(2*fwhm_y-(fwhm_y/2))) + 1, int(ceil(fwhm_x/2)) : int(ceil(2 * fwhm_x - (fwhm_x/2))) +1] = 1

        # Rotate MASK and the gaussian in counter clockwise direction
        rot_psf_mask = np.array(Image.fromarray(psf_mask).rotate(angle = angle, expand = True))
        psf_gaussian = np.array(Image.fromarray(psf_gaussian).rotate(angle = angle, expand = True))

        psf_gaussian[np.where(rot_psf_mask == 0)] = 0
        rowpix, colpix = np.where(rot_psf_mask == 1)

        # Cut the required portion of PSF which actually has an impact
        psf = psf_gaussian[np.min(rowpix):np.max(rowpix) + 1, np.min(colpix):np.max(colpix) + 1]
        psf = psf / np.sum(psf)

        # Normalize and return the final PSF
        return np.round(psf, 4)

    def PSF(self, fwhm_x, fwhm_y):
        """
        PSF for accomodating change in pixels because of sensor observations
        Args:
            fwhm_x (_int_): _Useful for defining x side of pixel change effects_
            fwhm_y (_int_): _Useful for defining y side of pixel change effects_

        Returns:
            _numpy_arr_: _PSF array by modelling using gaussian function_
        """
        # TODO : Fix and try to check if it's opposite
        sd_x = fwhm_x / (2 * math.sqrt(2 * math.log(2)))
        sd_y = fwhm_y / (2 * math.sqrt(2 * math.log(2)))

        x_max = fwhm_x * 2
        y_max = fwhm_y * 2

        # Gaussian Function for both x and y
        gaussian_x = np.matrix(self.gaussian(x_max + 1, sd_x))
        gaussian_y = np.matrix(self.gaussian(y_max + 1, sd_y))
        gaussian_xy = np.matmul(gaussian_y.getH(), gaussian_x)

        # Combination of both gaussians
        psf = gaussian_xy / (np.sum(gaussian_xy))
        psf = np.round(psf, 4)
        return psf

    def gaussian(self, M, std):
        """
        Gaussian Function used to account for PSF
        Args:
            M (_float_): _Mean of the Gaussian distribution_
            std (_float_): _Standard Deviation of the Gaussian distribution_

        Returns:
            _numpy_arr_
        """
        Mn = (M - 1)/2
        n = np.array([i for i in range(-int(Mn), int(Mn) + 1, 1)])
        gwin = np.exp(-0.5*np.power((n/std), 2))
        return gwin


import pytest

import pickle

import pandas as pd

from daskperiment.backend import MongoBackend, RedisBackend
from daskperiment.backend.mongo import MongoKey

class NoSQLBase(object):

    @classmethod
    def init_backend(cls):
        raise NotImplementedError

    def make_key(self, key):
        """
        Create key to backend
        """
        return key

    def make_expected(self, exp):
        """
        Convert expected input to DB internal repr
        """
        return exp

    @classmethod
    def teardown_class(cls):
        backend = cls.init_backend()
        backend._delete_cache()

    def test_set_get(self):
        backend = self.init_backend()

        key = self.make_key('test_key')
        backend.set(key, 0)
        assert backend.get(key) == self.make_expected(0)

    def test_get_not_exist(self):
        backend = self.init_backend()
        key = self.make_key('no_key')
        assert backend.get(key) is None

    @pytest.mark.parametrize('key_string', ['list_key', 'dot_list_key.child',
                                            'num_list_key.series.1'])
    def test_append_list(self, key_string):
        backend = self.init_backend()

        key = self.make_key(key_string)
        backend.append_list(key, 10)
        assert backend.get_list(key) == self.make_expected([10])

        backend.append_list(key, 20)
        assert backend.get_list(key) == self.make_expected([10, 20])

        backend.append_list(key, 30)
        backend.append_list(key, 40)
        assert backend.get_list(key) == self.make_expected([10, 20, 30, 40])

        new_key = self.make_key(key_string + '2')
        backend.append_list(new_key, 1)
        assert backend.get_list(new_key) == self.make_expected([1])
        assert backend.get_list(key) == self.make_expected([10, 20, 30, 40])

    def test_no_list_key(self):
        backend = self.init_backend()
        key = self.make_key('no_list')
        assert backend.get_list(key) == []
        assert backend.get(key) is None

    def test_increment(self):
        backend = self.init_backend()

        key = self.make_key('incr_key')
        assert backend.increment(key) == 1
        assert backend.increment(key) == 2
        assert backend.increment(key) == 3
        assert backend.increment(key) == 4

        new_key = self.make_key('new_incr_key')
        assert backend.increment(new_key) == 1
        assert backend.increment(key) == 5

    def test_save_load_text(self):
        backend = self.init_backend()

        key = self.make_key('text_key')
        backend.save_text(key, 'aaa')
        assert backend.load_text(key) == 'aaa'

        backend.save_text(key, 'bbb')
        assert backend.load_text(key) == 'bbb'

    @pytest.mark.parametrize('key_string', ['obj_key', 'num_obj_key.series.1'])
    def test_save_load_object(self, key_string):
        backend = self.init_backend()
        obj = {'a': 1, 'b': pd.Timestamp('2011-01-01')}

        key = self.make_key(key_string)
        backend.save_object(key, obj)
        assert backend.load_object(key) == obj

    def test_pickle_roundtrip_save_load_object(self):
        backend = self.init_backend()

        pickle_obj_key = self.make_key('pickle_obj_key')
        obj = {'a': 1, 'b': pd.Timestamp('2011-01-01')}
        backend.save_object(pickle_obj_key, obj)
        assert backend.load_object(pickle_obj_key) == obj

        pickle_text_key = self.make_key('pickle_text_key')
        backend.save_object(pickle_text_key, 'xxxx')
        assert backend.load_object(pickle_text_key) == 'xxxx'

        new_backend = pickle.loads(pickle.dumps(backend))
        assert new_backend.load_object(pickle_obj_key) == obj
        assert new_backend.load_object(pickle_text_key) == 'xxxx'

        assert backend.load_object(pickle_obj_key) == obj
        assert backend.load_object(pickle_text_key) == 'xxxx'

class TestRedisBackend(NoSQLBase):

    @classmethod
    def init_backend(cls):
        uri = 'redis://localhost:6379/0'
        return RedisBackend('redis_backend', uri)

    def make_expected(self, exp):
        if isinstance(exp, int):
            return str(exp).encode()
        elif isinstance(exp, list):
            return [self.make_expected(e) for e in exp]
        else:
            raise NotImplementedError

class TestMongoBackend(NoSQLBase):

    @classmethod
    def init_backend(cls):
        uri = 'mongodb://localhost:27017/test_db'
        return MongoBackend('mongo_backend', uri)

    def make_key(self, key):
        return MongoKey({'test_mongo_key': key})

    def test_mongo_key(self):
        backend = self.init_backend()

        key = MongoKey({'experiment_id': 'test_mongo_key',
                        'trial_id': 1})

        backend.set(key, 100)
        assert backend.get(key) == 100
        backend.set(key, 200)
        assert backend.get(key) == 200

        new_key = MongoKey({'experiment_id': 'test_mongo_key',
                            'trial_id': 2})

        backend.set(new_key, 300)
        assert backend.get(new_key) == 300
        assert backend.get(key) == 200

    def test_mongo_key_internal(self):
        backend = self.init_backend()

        key = MongoKey({'experiment_id': 'test_mongo_key_internal',
                        'trial_id': 1354,
                        'internal_unique_key': 'xxxxx'})

        backend.set(key, 100)
        assert backend.get(key) == 100

        key = MongoKey({'internal_unique_key': 'xxxxx'},
                       field_name='experiment_id')
        assert backend.get(key) == 'test_mongo_key_internal'

        key = MongoKey({'internal_unique_key': 'xxxxx'},
                       field_name='trial_id')
        assert backend.get(key) == 1354

    def test_append_list_mongo_key(self):
        backend = self.init_backend()

        key = MongoKey({'experiment_id': 'test_append_list_mongo_key',
                        'trial_id': 1})

        backend.append_list(key, 10)
        assert backend.get_list(key) == [10]

        backend.append_list(key, 20)
        assert backend.get_list(key) == [10, 20]

    def test_increment_mongo_key(self):
        backend = self.init_backend()

        key = MongoKey({'experiment_id': 'test_increment_mongo_key',
                        'trial_id': 1})

        assert backend.increment(key) == 1
        assert backend.increment(key) == 2
        assert backend.increment(key) == 3
        assert backend.increment(key) == 4

        new_key = MongoKey({'experiment_id': 'test_increment_mongo_key',
                            'trial_id': 2})
        assert backend.increment(new_key) == 1
        assert backend.increment(key) == 5

class TestMongoKey(object):

    def test_eq(self):
        key1 = MongoKey({'doc': 1, 'id': 2})
        key2 = MongoKey({'doc': 1, 'id': 2})
        assert key1 == key2

        key3 = MongoKey({'doc': 1, 'id': 3})
        assert key1 != key3

        key4 = MongoKey({'doc': 1, 'id': 2}, field_name='xxx')
        assert key1 != key4
        assert key2 != key4

    def test_repr(self):
        key = MongoKey({'doc': 1, 'id': 2})
        exp = "MongoKey(meta={'doc': 1, 'id': 2}, field_name=value)"
        assert repr(key) == exp

    def test_wildcard(self):
        key = MongoKey({'key1': 1, 'key2': 2, 'key3': '*'})
        assert key.document_meta == {'key1': 1, 'key2': 2}


#
# @lc app=leetcode.cn id=1534 lang=python3
#
# [1534] 统计好三元组
#
# https://leetcode-cn.com/problems/count-good-triplets/description/
#
# algorithms
# Easy (78.21%)
# Likes:    5
# Dislikes: 0
# Total Accepted:    7.9K
# Total Submissions: 10K
# Testcase Example:  '[3,0,1,1,9,7]\n7\n2\n3'
#
# 给你一个整数数组 arr ，以及 a、b 、c 三个整数。请你统计其中好三元组的数量。
#
# 如果三元组 (arr[i], arr[j], arr[k]) 满足下列全部条件，则认为它是一个 好三元组 。
#
#
# 0 <= i < j < k < arr.length
# |arr[i] - arr[j]| <= a
# |arr[j] - arr[k]| <= b
# |arr[i] - arr[k]| <= c
#
#
# 其中 |x| 表示 x 的绝对值。
#
# 返回 好三元组的数量 。
#
#
#
# 示例 1：
#
# 输入：arr = [3,0,1,1,9,7], a = 7, b = 2, c = 3
# 输出：4
# 解释：一共有 4 个好三元组：[(3,0,1), (3,0,1), (3,1,1), (0,1,1)] 。
#
#
# 示例 2：
#
# 输入：arr = [1,1,2,2,3], a = 0, b = 0, c = 1
# 输出：0
# 解释：不存在满足所有条件的三元组。
#
#
#
#
# 提示：
#
#
# 3 <= arr.length <= 100
# 0 <= arr[i] <= 1000
# 0 <= a, b, c <= 1000
#
#
#

from typing import List

# @lc code=start
class Solution:
    def countGoodTriplets(self, arr: List[int], a: int, b: int, c: int) -> int:
        res = 0
        for i in range(len(arr)):
            for j in range(i + 1, len(arr)):
                for k in range(j + 1, len(arr)):
                    if abs(arr[i] - arr[j]) <= a and \
                       abs(arr[j] - arr[k]) <= b and \
                       abs(arr[i] - arr[k]) <= c:
                        res += 1
        return res

# @lc code=end


import os
import json
import yaml
import scipy
import pickle
import numpy as np
from copy import deepcopy
from torch.utils.data import Dataset

from utils import read_nifti

class Geometry(object):
    def __init__(self, config):
        self.v_res = config['nVoxel'][0]    # ct scan
        self.p_res = config['nDetector'][0] # projections
        self.v_spacing = np.array(config['dVoxel'])[0]    # mm
        self.p_spacing = np.array(config['dDetector'])[0] # mm

        self.DSO = config['DSO'] # mm
        self.DSD = config['DSD'] # mm

    def project(self, points, angle):
        # points: [N, 3] ranging from [0, 1]
        # d_points: [N, 2] ranging from [-1, 1]

        points = deepcopy(points).astype(float)
        points[:, :2] -= 0.5 # [-0.5, 0.5]
        points[:, 2] = 0.5 - points[:, 2] # [-0.5, 0.5]
        points *= self.v_res * self.v_spacing # mm

        angle = -1 * angle # inverse direction
        rot_M = np.array([
            [np.cos(angle), -np.sin(angle), 0],
            [np.sin(angle),  np.cos(angle), 0],
            [            0,              0, 1]
        ])
        points = points @ rot_M.T

        d1 = self.DSO
        d2 = self.DSD

        coeff = (d2) / (d1 - points[:, 0]) # N,
        d_points = points[:, [2, 1]] * coeff[:, None] # [N, 2] float
        d_points /= (self.p_res * self.p_spacing)
        d_points *= 2 # NOTE: some points may fall outside [-1, 1]

        return d_points

class Mixed_CBCT_dataset(Dataset):
    def __init__(self, dst_list, **kwargs) -> None:
        super().__init__()
        print('mixed_dataset:', dst_list)
        self.name_list = dst_list
        self.datasets = []
        for dst_name in self.name_list:
            self.datasets.append(CBCT_dataset(dst_name, **kwargs))

    def __len__(self):
        dst_len = [len(d) for d in self.datasets]
        return np.sum(dst_len)

    @property
    def num_dst(self):
        return len(self.datasets)

    def find_dst(self, index):
        for i, dst in enumerate(self.datasets):
            if index >= len(dst):
                index -= len(dst)
            else:
                return i, index

    def __getitem__(self, index):
        dst_idx, index = self.find_dst(index)

        dst_vec = np.zeros(self.num_dst)
        dst_vec[dst_idx] = 1.

        data_dict = self.datasets[dst_idx][index]
        data_dict['dst_vec'] = dst_vec
        data_dict['dst_name'] = self.name_list[dst_idx]
        return data_dict

class CBCT_dataset(Dataset):
    def __init__(
            self,
            dst_name,
            split='train',
            num_views=10,
            npoint=5000,
            out_res=256,
            random_views=False,
            view_offset=0
        ):
        super().__init__()
        dst_root = './data'

        # load dataset info
        if dst_name in ['knee_cbct']:
            data_root = os.path.join(dst_root, dst_name)
            with open(os.path.join(data_root, 'info.json'), 'r') as f:
                cfg = json.load(f)
                name_list = sorted(cfg[split])
                print('CBCT_dataset, name: {}, split: {}, len: {}.'.format(dst_name, split, len(name_list)))
        else:
            raise ValueError(dst_name)

        # load projection config
        with open(os.path.join(data_root, cfg['projection_config']), 'r') as f:
            proj_cfg = yaml.safe_load(f)
            self.geo = Geometry(proj_cfg)

        # prepare points
        if split == 'train':
            # load blocks' coordinates [train only]
            self.blocks = np.load(os.path.join(data_root, cfg['blocks']))
        else:
            # prepare sampling points
            points = np.mgrid[:out_res, :out_res, :out_res]
            points = points.astype(float) / (out_res - 1)
            points = points.reshape(3, -1)
            self.points = points.transpose(1, 0) # N, 3

        # other parameters
        self.out_res = out_res
        self.data_root = data_root
        self.cfg = cfg

        self.name_list = name_list
        self.npoint = npoint
        self.is_train = (split == 'train')
        self.num_views = num_views
        self.random_views = random_views
        self.view_offset = view_offset

    def __len__(self):
        return len(self.name_list)

    def sample_projections(self, name):
        # -- load projections
        with open(os.path.join(self.data_root, self.cfg['projections'].format(name)), 'rb') as f:
            data = pickle.load(f)
            projs = data['projections'] # K, 1, res^2
            angles = data['angles']     # K,

        # -- sample projections
        views = np.linspace(0, len(projs), self.num_views, endpoint=False).astype(int)
        offset = np.random.randint(len(projs) - views[-1]) if self.random_views else self.view_offset
        views += offset
        projs = projs[views].astype(float) / 255.
        projs = projs[:, None, ...]
        angles = angles[views]

        # -- de-normalization [required for mixed dataset]
        # projs = projs * self.cfg['projection_norm'] / 0.2

        return projs, angles

    def load_ct(self, name):
        image = read_nifti(os.path.join(self.data_root, self.cfg['image'].format(name)))
        image = image.astype(np.float32) / 255.
        if self.out_res == 128:
            image = scipy.ndimage.zoom(image, 0.5, order=3, prefilter=False)
        elif self.out_res != 256:
            raise ValueError
        return image

    def load_block(self, name, b_idx):
        path = os.path.join(self.data_root, self.cfg['image_block'].format(name, b_idx))
        return np.load(path)

    def sample_points(self, points, values=None):
        choice = np.random.choice(len(points), size=self.npoint, replace=False)
        points = points[choice]
        if values is not None:
            values = values[choice]
            values = values.astype(float) / 255.
            return points, values
        else: return points

    def __getitem__(self, index):
        name = self.name_list[index]

        # -- load projections
        projs, angles = self.sample_projections(name)

        # -- load sampling points
        if not self.is_train:
            points = self.points
            image = self.load_ct(name)
            p_gt = np.zeros(len(points))
        else:
            b_idx = np.random.randint(len(self.blocks))
            block_values = self.load_block(name, b_idx)
            block_coords = self.blocks[b_idx] # N, 3
            points, p_gt = self.sample_points(block_coords, block_values)

        # -- project points and view direction
        proj_points = []
        for a in angles:
            p = self.geo.project(points, a)
            proj_points.append(p)
        proj_points = np.stack(proj_points, axis=0) # M, N, 2

        # -- normalize points
        points = deepcopy(points) # ~[0, 1]
        points[:, :2] -= 0.5 # [-0.5, 0.5]
        points[:, 2] = 0.5 - points[:, 2] # [-0.5, 0.5]
        points *= 2 # => [-1, 1]

        # -- normalize viewing angles
        angles = np.array(angles, dtype=float) # ~[0, +pi]
        angles = angles / np.pi * 2 - 1 # => [-1, 1]

        # -- collect data
        ret_dict = {
            'name': name,
            'points': points,           # 3D points
            'angles': angles[:, None],  # angles
            'proj_points': proj_points, # projected points
            'projs': projs,             # 2D projections
            'p_gt': p_gt[None, :]       # labels
        }
        if not self.is_train:
            ret_dict['image'] = image

        return ret_dict

if __name__ == '__main__':
    dst = CBCT_dataset(dst_name='knee_zhao', random_views=True, num_views=10)
    item = dst[0]
    import pdb; pdb.set_trace()


import json
from pprint import pprint
from src.environment import ReacherEnvironment
from src.mappo import MAPPO
from src.models import SimpleMAPPOAgent
import torch
import numpy as np
import datetime
import matplotlib.pyplot as plt

def train(*args, **kwargs):
    print(kwargs)

    device = 'cpu'
    kwargs['mappo']['device'] = device

    env = ReacherEnvironment(**kwargs['env'])
    env.reset(train_mode=True)

    state_dim = env.get_state_dim()
    action_dim = env.get_action_dim()
    kwargs['agent']['state_dim'] = state_dim
    kwargs['agent']['action_dim'] = action_dim
    kwargs['agent']['num_agents'] = env.get_num_agents()

    agent = SimpleMAPPOAgent(**kwargs['agent']).to(device)
    alg = MAPPO(agent=agent, **kwargs['mappo'])
    scores = alg.train(env, 5000)

    dt = str(datetime.datetime.now().strftime("%m_%d_%Y_%I_%M_%p"))
    model_fname = "../models/mappo_tennis_{}.pt".format(dt)
    torch.save(agent, model_fname)

    scores_fname = "../reports/mappo_tennis_{}".format(dt)
    np.save(scores_fname, np.asarray(scores))

    plt.plot(scores)
    plt.plot(np.convolve(scores, np.ones(100)/100)[:-100])
    fig_name = "../reports/mappo_tennis_{}.png".format(dt)
    plt.savefig(fig_name)

if __name__ == '__main__':
    with open('../config.json') as data_file:
        kwargs = json.load(data_file)
    pprint(kwargs)
    train(**kwargs)


from dataclasses import asdict
import json
from pathlib import Path
from genetic_algo.structures import Dataset, DatasetItem
from genetic_algo.dataset import load_from_file

def test_load_dataset(tmpdir):
    dataset_items = [
        DatasetItem(
            name="Item1",
            fitness=1,
            weight=0.5,
        )
    ]

    dataset = Dataset(
        expected_weight=10,
        items=dataset_items,
    )

    dataset_file = tmpdir / Path("foo.json")

    with open(dataset_file, "w") as fp:
        json.dump(asdict(dataset), fp)

    result = load_from_file(dataset_file)

    assert result == dataset