import csv
import os
import numpy as numpy
from bs4 import BeautifulSoup
from requests import get
import lxml
from multiprocessing import Process, Pool
import pandas as pd
def getname(line):
	"""
	Deciphers and parses the name of a player given a line from a file within
	./data_loader/data/early_applicants/txt
	"""
	parts = line.replace('\n', '').replace('|', ',').split(',')
	fullname = parts[0]
	split_name = fullname.split(' ')
	firstname = split_name[0]
	lastname = ' '.join(split_name[1:])
	return firstname, lastname

def early_applicants_txt2csv():
	"""
	Some of the nba early_applicant website formats are difficult to webscrape, so assuming
	we have the directory ./data_loader/data/early_applicants/txt which has a handful of
	txt files containing pasted content from the site, we will generate csv files
	of those players.
	"""
	rootdir = './data_loader/data/early_applicants/txt'
	filenames = os.listdir(rootdir)
	filenames = [os.path.join(rootdir, filename) for filename in filenames]

	for filename in filenames:
		data = []
		with open(filename, 'r') as file:
			with open(filename.replace('txt', 'csv'), 'w') as csvfile:
				writer = csv.writer(csvfile)
				writer.writerow(['first_name', 'last_name'])
				line = file.readline()
				while line:
					writer.writerow(getname(line))
					line = file.readline()

def get_year_csv(args):
	"""
	Generates the year csv given some url to web scrape and the year you want to scrape on.
	Used to generate the early_applicants csv validation file
	"""
	url, year = args
	content = get(url).content
	soup = BeautifulSoup(content, 'lxml')
	table = soup.findAll('table')[1]
	players = table.findAll('tr')
	with open('./data_loader/data/early_applicants/csv/{0}.csv'.format(year), 'w') as csvfile:
		writer = csv.writer(csvfile)
		for i, player in enumerate(players):
			if (i == 0):
				writer.writerow(['first_name', 'last_name'])
			else:
				name = player.find('td').getText().strip().split(' ')
				first_name = name[0]
				last_name = ' '.join(name[1:])
				writer.writerow([first_name, last_name])

def get_nba_data(soup):
	"""
	Given some soup object, parses and returns nba player information. Used when
	generating draftpick data

	@param soup: BeautifulSoup object of a table row
	"""
	player = soup.find('td', {'class':'player'}).getText()
	team = soup.find('td', {'class':'team'}).getText()
	affiliation = soup.find('td', {'class':'first text'}).getText()
	year, num, pick, overrall = soup.findAll('td', {'class':''})
	return [player, team, affiliation, year.getText(), num.getText(), pick.getText(),
	        overrall.getText()]

def nba_draftpicks():
	"""
	Generates nba draftpicks using pasted html source code inside the directory
	./data_loader/data/draftpicks/html
	"""
	url = 'https://stats.nba.com/draft/history/?Season='
	html_rootdir = './data_loader/data/draftpicks/html'
	csv_rootdir = './data_loader/data/draftpicks/csv'

	raw_files = [os.path.join(html_rootdir, filename) for filename in os.listdir(html_rootdir)]
	html_files = []
	for filename in raw_files:
	    with open(filename, 'r') as file:
	        html_files.append(file.read())

	header = ['Player', 'Team', 'Affiliation',
		'Year', 'Round Number', 'Round Pick',
		'Overall Pick']

	for i, date in enumerate(os.listdir(html_rootdir)):
		filename = date.replace('.htm', '')
		csv_filepath = os.path.join(csv_rootdir, '{0}.csv'.format(filename))
		with open(csv_filepath, 'w') as csv_file:
		    soup = BeautifulSoup(html_files[i], 'lxml')
		    table = soup.find('table')
		    writer = csv.writer(csv_file)
		    writer.writerow(header)
		    rows = table.findAll('tr')
		    for k, row in enumerate(rows):
		        if k != 0:
		            writer.writerow(get_nba_data(row))

def create_player_stats():
	player_box = pd.read_csv("./data_loader/data/ncaa/player_box.csv")
	player_info = pd.read_csv("./data_loader/data/ncaa/player_info.csv")

	# Filter out unwanted values
	player_box = player_box[player_box['player_id']!=-101]
	player_box = player_box[~player_box['player_id'].isnull()]
	player_box = player_box[player_box['season']!=None]
	player_box = player_box[player_box['season']!=2018]

	player_info = player_info[player_info['player_id']!=-101]
	player_info = player_info[~player_info['player_id'].isnull()]
	player_info = player_info[player_info['season']!=None]
	player_info = player_info[player_info['season']!=2018]
	player_info = player_info[~player_info['first_name'].isnull()]
	player_info = player_info[~player_info['last_name'].isnull()]

	# Combine the two datasets
	player_box["player_id_AND_season"] = player_box["player_id"].map(int).map(str) + "_AND_" + player_box["season"].map(int).map(str)
	player_info["player_id_AND_season"] = player_info["player_id"].map(int).map(str) + "_AND_" + player_info["season"].map(int).map(str)
	player_stats = pd.merge(player_box, player_info, on='player_id_AND_season', how='outer')

	player_stats = player_stats[['pts', 'fga', 'fga3', 'fgm', 'fgm3', 'fta', 'ftm', 'ast', 'blk', 'stl', 'dreb', 'oreb', 'reb', 'pf', 'tf', 'tov', 'mins_played', 'grade_level', 'division', 'first_name', 'last_name', 'height_in',
	'school', 'conference', 'position', "season_x"]]
	player_stats = player_stats.dropna()

	# TODO optimize, takes too long
	# Appends categorical argument of whether player applied to the nba
	applied_stats = applied.join(player_stats)

	# Appends categorical argument of whether player was accepted into the nba
	draft_stats = draft.join(applied_stats)

	df = draft_stats.copy()
	test = df.dropna()
	test.to_csv('./data_loader/data/ncaa/player_stats.csv', index=False)

def getall():
	"""
	Generates all custom made datasets
	"""
	args = (
		('http://www.nba.com/2016/news/04/26/early-entry-candidates-2016-draft/', '2016'),
		('http://www.nba.com/2015/news/04/28/early-entry-candidates-for-2015-draft/', '2015')
	)
	pool = Pool()
	pool.map(get_year_csv, args)
	pool.apply(nba_draftpicks)
	pool.apply(early_applicants_txt2csv)
	pool.close()
	pool.join()

def get_classifier_data_feature():
	players = pd.read_csv('./data_loader/data/ncaa/player_stats.csv')
	#used categorical cols: 'school', 'conference', 'position',
	X = players.drop(columns=['drafted', 'first_name', 'last_name', 'xml_name', 'class_desc',
		 'player_id_AND_season', 'player_id_x', 'game_date', 'game_id', 'season_x', 'team_code_x'])
	X.to_csv("player_stats_raw_features.csv")
	y = players.drafted
	y.to_csv("player_stats_labels.csv")

if __name__ == '__main__':
	getall()


import argparse
import dxfwrite
from math import sqrt, acosh, cosh, ceil

def main():
    parser = argparse.ArgumentParser(description='Catenary generator', add_help=True)
    parser.add_argument('-o', '--output', nargs='?', type=str, metavar='PATH',
                        required=True,
                        help='Output file path')
    parser.add_argument('-s', '--size', nargs='?', type=float, metavar='DIM',
                        required=True, help='Square dimension')
    parser.add_argument('-r', '--repeats', nargs='?', type=float, metavar='N',
                        required=False, default=1, help='Number of repeats')
    parser.add_argument('--resolution', nargs='?', type=int, metavar='N',
                        required=False, default=100, help='Number of segments per repeat')
    args = parser.parse_args()

    dxf = dxfwrite.DXFEngine()
    drawing = dxf.drawing()
    a = args.size / 2
    d = a * sqrt(2)
    b = a * acosh(sqrt(2))
    def f(x):
        while x < 0:
            x += 2 * b
        while x >= 2 * b:
            x -= 2 * b
        return d - a * cosh((x - b) / a)
    x = [2 * b * i / args.resolution for i in range(int(ceil(args.resolution * args.repeats)))]
    y = list(map(f, x))
    for i in range(1, len(x)):
        drawing.add(dxf.line((x[i-1], y[i-1]), (x[i], y[i])))
    drawing.saveas(args.output)

if __name__ == "__main__":
    main()


from math import sin, cos, pi

from core import SimpleSection, ComplexSection, Dimensions, cached_property

# ==============================================================================
# S I M P L E   S E C T I O N S
# ==============================================================================

class Rectangle(SimpleSection):
    dimensions = Dimensions(a=None, b=None)

    def check_dimensions(self, dims):
        if dims.a <= 0:
            raise ValueError("Invalid dimensions: a <= 0")
        if dims.b <= 0:
            raise ValueError("Invalid dimensions: b <= 0")

    @cached_property
    def A(self):
        return self.density * self.a * self.b

    @cached_property
    def _cog(self):
        return 0.0, 0.0

    @cached_property
    def _I0(self):
        _I11 = self.a * self.b**3 / 12.
        _I22 = self.b * self.a**3 / 12.
        _I12 = 0.0
        return tuple(self.density * i for i in (_I11, _I22, _I12))

class CircularSector(SimpleSection):
    dimensions = Dimensions(ro=None, ri=None, phi=None)

    def check_dimensions(self, dims):
        if dims.ri < 0:
            raise ValueError("Invalid dimensions: ri < 0")
        if dims.ro <= dims.ri:
            raise ValueError("Invalid dimensions: ro < ri")
        if dims.phi <= 0:
            raise ValueError("Invalid dimensions: phi <= 0")
        if dims.phi > 2*pi:
            raise ValueError("Invalid dimensions: phi > 2*pi")

    @cached_property
    def A(self):
        ro  = self.ro
        ri  = self.ri
        phi = self.phi
        A   = 0.5 * (ro**2 - ri**2) * self.phi
        return self.density * A

    @cached_property
    def _cog(self):
        ro  = self.ro
        ri  = self.ri
        phi = self.phi
        A   = self.A / self.density
        S2 = 2./3. * (ro**3 - ri**3) * sin(0.5*phi)
        _e1 = S2 / A
        _e2 = 0.0
        return _e1, _e2

    @cached_property
    def _I0(self):
        ro  = self.ro
        ri  = self.ri
        phi = self.phi
        _e1, _e2 = self._cog
        _I11 = self.density * 0.125 * (ro**4 - ri**4) * (phi - sin(phi))
        _I22 = self.density * 0.125 * (ro**4 - ri**4) * (phi + sin(phi))
        _I12 = self.density * 0.0
        return self.parallel_axis((_I11, _I22, _I12), self._cog, reverse=True)

class CircularSegment(SimpleSection):
    dimensions = Dimensions(r=None, phi=None)

    def check_dimensions(self, dims):
        if dims.r <= 0:
            raise ValueError("Invalid dimensions: r <= 0")
        if dims.phi <= 0:
            raise ValueError("Invalid dimensions: phi <= 0")
        if dims.phi > 2*pi:
            raise ValueError("Invalid dimensions: phi > 0")

    @cached_property
    def _cog(self):
        r, phi = self.r, self.phi
        _e1 = 4 * sin(0.5*phi)**3 * r / (3 * (phi - sin(phi)) )
        _e2 = 0.0
        return _e1, _e2

    @cached_property
    def A(self):
        r, phi = self.r, self.phi
        return self.density * 0.5 * r**2 * (phi - sin(phi))

    @cached_property
    def _I0(self):
        r   = self.r
        phi = self.phi
        A  = self.A / self.density
        e1, e2 = self._cog
        I11 = 1. / 48. * r**4 * (6*phi - 8*sin(phi) + sin(2*phi))
        I22 = 0.125 * r**4 * (phi - sin(phi)*cos(phi)) - A * e1**2
        I12 = 0.0
        return tuple(self.density * i for i in (I11, I22, I12))

class Polygon(SimpleSection, list):

    @cached_property
    def A(self):
        if len(self) < 3:
            raise ValueError("Cannot calculate A: Polygon must have at least three vertices")

        #Area will be positive if vertices are ordered counter-clockwise
        x1, x2 = self.__looped_vertices()
        n = len(self)
        A = 0.5 * sum(x1[i]*x2[i+1] - x1[i+1]*x2[i] for i in range(n))
        return self.density * A

    @cached_property
    def _cog(self):
        if len(self) < 3:
            raise ValueError("Cannot calculate _cog: Polygon must have at least three vertices")

        x1, x2 = self.__looped_vertices()
        n = len(self)
        A = self.A / self.density
        _e1 = 1. / (6 * A) * sum( (x1[i] + x1[i+1]) * (x1[i]*x2[i+1] - x1[i+1]*x2[i]) for i in range(n) )
        _e2 = 1. / (6 * A) * sum( (x2[i] + x2[i+1]) * (x1[i]*x2[i+1] - x1[i+1]*x2[i]) for i in range(n) )
        return _e1, _e2

    @cached_property
    def _I0(self):
        if len(self) < 3:
            raise ValueError("Cannot calculate _I0: Polygon must have at least three vertices")

        x1, x2 = self.__looped_vertices()
        n = len(self)
        _I11 = 1./12. * sum( (x2[i]**2 + x2[i]*x2[i+1] + x2[i+1]**2) * (x1[i]*x2[i+1] - x1[i+1]*x2[i]) for i in range(n))
        _I22 = 1./12. * sum( (x1[i]**2 + x1[i]*x1[i+1] + x1[i+1]**2) * (x1[i]*x2[i+1] - x1[i+1]*x2[i]) for i in range(n))
        _I12 = 1./24. * sum( (x1[i]*x2[i+1] + 2*x1[i]*x2[i] + 2*x1[i+1]*x2[i+1] + x1[i+1]*x2[i] )*(x1[i]*x2[i+1] - x1[i+1]*x2[i]) for i in range(n))

        _I = tuple(self.density * i for i in (_I11, _I22, _I12))
        return self.parallel_axis(_I, self._cog, reverse=True)

    def __looped_vertices(self):
        n = len(self)
        x1 = [self[i%n][0] for i in range(n + 1)]
        x2 = [self[i%n][1] for i in range(n + 1)]
        return x1, x2

    # Override list methods which add new items to the list
    # Only allow to add items consisting of two values which can be
    # convered to float.
    # Any change of vertices must call self.reset_cached_properties
    # =============================================================

    def append(self, vertex):
        vertex = self.convert_to_vertices(vertex)[0]
        list.append(self, vertex)
        self.reset_cached_properties()

    def extend(self, vertices):
        vertices = self.convert_to_vertices(*vertices)
        list.extend(self, vertices)
        self.reset_cached_properties()

    def insert(self, i, vertex):
        vertex = self.convert_to_vertices(vertex)[0]
        list.insert(self, i, vertex)
        self.reset_cached_properties()

    def __setitem__(self, i, vertex):
        vertex = self.convert_to_vertices(vertex)[0]
        list.__setitem__(self, i, vertex)
        self.reset_cached_properties()

    def __setslice__(self, i, j, vertices):
        vertices = self.convert_to_vertices(*vertices)
        list.__setslice__(self, i, j, vertices)
        self.reset_cached_properties()

    @staticmethod
    def convert_to_vertices(*items):
        return [(float(x), float(y)) for x, y in items]

    # =============================================================

class Triangle(Polygon):

    # Override list methods which add new items to the list
    # Only allow to add items consisting of two values which can be
    # convered to float.
    # Any change of vertices must call self.reset_cached_properties
    # =============================================================

    def append(self, vertex):
        if len(self) == 3:
            raise IndexError("Triangle cannot have more than 3 vertices")
        super(Triangle, self).append(vertex)

    def extend(self, vertices):
        vertices = self.convert_to_vertices(*vertices)
        if len(self) + len(vertices) > 3:
            raise IndexError("Triangle cannot have more than 3 vertices")
        super(Triangle, self).extend(vertices)

    def insert(self, i, vertex):
        if len(self) == 3:
            raise IndexError("Triangle cannot have more than 3 vertices")
        super(Triangle, self).insert(i, vertex)

    def __setslice__(self, i, j, vertices):
        vertices = self.convert_to_vertices(*vertices)
        resulting = self[:]
        resulting.__setslice__(i, j, vertices)
        if len(resulting) > 3:
            raise IndexError("Triangle cannot have more than 3 vertices")
        super(Triangle, self).__setslice__(i, j, vertices)

    # =============================================================

    def reset_cached_properties(self):

        if len(self) == 3:
            x, y = zip(*self)
            v1 = x[1] - x[0], y[1] - y[0]
            v2 = x[2] - x[1], y[2] - y[1]
            sin = v1[0] * v2[1] - v1[1] * v2[0]
            if sin < 0:
                self[:] = self[0], self[2], self[1]
        super(Triangle, self).reset_cached_properties()

# ==============================================================================
# C O M P L E X   S E C T I O N S
# ==============================================================================

class Circle(ComplexSection):
    dimensions = Dimensions(r=None)
    sections   = [CircularSector]

    def update_sections(self):
        self.sections[0].set_dimensions(ri=0, ro=self.r, phi=2*pi)

class Box(ComplexSection):
    dimensions = Dimensions(a=None, b=None, ta=None, tb=None)
    sections = [Polygon]

    def check_dimensions(self, dims):
        if dims.a <= 0:
            raise ValueError("Invalid dimensions: a <= 0")
        if dims.b <= 0:
            raise ValueError("Invalid dimensions: b <= 0")
        if dims.ta <= 0:
            raise ValueError("Invalid dimensions: ta <= 0")
        if dims.tb <= 0:
            raise ValueError("Invalid dimensions: tb <= 0")
        if dims.a <= 2*dims.tb:
            raise ValueError("Invalid dimensions: a <= 2*tb")
        if dims.b <= 2*dims.ta:
            raise ValueError("Invalid dimensions: b <= 2*ta")

    def update_sections(self):
        ao = 0.5 * self.a
        ai = 0.5 * (self.a - 2*self.tb)
        bo = 0.5 * self.b
        bi = 0.5 * (self.b - 2*self.ta)

        polygon = self.sections[0]
        polygon[:] = [
            (-ao, -bo),
            ( ao, -bo),
            ( ao,  bo),
            (-ao,  bo),
            (-ao, -bi),
            (-ai, -bi),
            (-ai,  bi),
            ( ai,  bi),
            ( ai, -bi),
            (-ao, -bi)]

class Ring(ComplexSection):
    sections   = [CircularSector]
    dimensions = Dimensions(ro=None, ri=None)

    def update_sections(self):
        self.sections[0].set_dimensions(ro=self.ro, ri=self.ri, phi=2*pi)

    def check_dimensions(self, dims):
        super(Ring, self).check_dimensions(dims)

class Wedge(ComplexSection):
    sections   = [CircularSector]
    dimensions = Dimensions(r=None, phi=None)

    def update_sections(self):
        self.sections[0].set_dimensions(ro=self.r, ri=0, phi=self.phi)

class WedgeRing(CircularSector):
    pass

class BaseFillet(ComplexSection):
    sections   = [Triangle, CircularSegment]
    dimensions = Dimensions(r=None, phi=None)
    densities  = [1.0, -1.0]

    def check_dimensions(self, dims):
        if dims.r <= 0:
            raise ValueError("Invalid dimensions: r <= 0")
        if dims.phi <= 0:
            raise ValueError("Invalid dimensions: phi <= 0")
        if dims.phi == pi:
            raise ValueError("Invalid dimensions: phi = pi")
        if dims.phi >= 2*pi:
            raise ValueError("Invalid dimensions: phi >= 2*pi")

    def update_sections(self):
        def sign(x):
            return x / abs(x)

        alpha = self.phi/2
        beta  = abs(pi - self.phi)
        theta = pi * (self.phi < pi)
        a = self.r * cos(alpha) / sin(alpha)
        b = self.r * cos(alpha)**2 / sin(alpha) * sign(a)
        c = self.r * cos(alpha)
        d = self.r / sin(alpha) * sign(a)

        triangle = self.sections[0]
        triangle[:] = [
            (0,  0),
            (b,  c),
            (b, -c)]

        segment = self.sections[1]
        segment.set_dimensions(r=self.r, phi=beta)
        segment.set_position(d1=d, d2=0, theta=theta)

        if self.phi > pi:
            self.densities = [-d for d in self.__class__.densities]
        else:
            self.densities = self.__class__.densities[:]
        self.set_density(self.density)

class Fillet(ComplexSection):
    sections   = [BaseFillet]
    dimensions = Dimensions(r=None, phi0=None, phi1=None)

    def check_dimensions(self, dims):
        if dims.r <= 0:
            raise ValueError("Invalid dimensions: r <= 0")
        if dims.phi1 <= dims.phi0:
            raise ValueError("Invalid dimensions: phi1 <= phi0")
        if dims.phi1 - dims.phi0 >= 2*pi:
            raise ValueError("Invalid dimensions: phi1 - phi0 >= 2*pi")

    def update_sections(self):
        phi = self.phi1 - self.phi0
        theta = 0.5 * (self.phi0 + self.phi1)
        self.sections[0].set_dimensions(r=self.r, phi=phi)
        self.sections[0].set_position(d1=0, d2=0, theta=theta)


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

import os

extensions = [
    'sphinx.ext.autodoc',
    'sphinx.ext.autosummary',
    'sphinx.ext.coverage',
    'sphinx.ext.doctest',
    'sphinx.ext.extlinks',
    'sphinx.ext.ifconfig',
    'sphinx.ext.napoleon',
    'sphinx.ext.todo',
    'sphinx.ext.intersphinx',
    'sphinx_click.ext',
    'nbsphinx',
    'sphinx_copybutton',
    'sphinx.ext.mathjax',
    'recommonmark',
    'IPython.sphinxext.ipython_console_highlighting',
]

if os.getenv('SPELLCHECK'):
    extensions += 'sphinxcontrib.spelling',
    spelling_show_suggestions = True
    spelling_lang = 'en_US'

# to exclude traditional Python prompts from your copied code
copybutton_prompt_text = ">>> "
# All lines of the code blocks will be copied after the prompts are stripped.
copybutton_only_copy_prompt_lines = False

nbsphinx_execute = 'never'

autodoc_mock_imports = ['torch', 'tensorflow']
autosummary_generate = True

source_suffix = {
    '.rst': 'restructuredtext',
    '.txt': 'markdown',
    '.md': 'markdown',
}
master_doc = 'index'
project = 'Hangar'
year = '2019-2020'
author = 'Richard Izzo'
copyright = '{0}, {1}'.format(year, author)
version = release = '0.5.2'

pygments_style = 'default'
pygments_lexer = 'PythonConsoleLexer'
highlight_options = {
    'python3': True
}
templates_path = ['.']
exclude_patterns = ['_build', '**.ipynb_checkpoints']
extlinks = {
    'issue': ('https://github.com/tensorwerk/hangar-py/issues/%s', '#'),
    'pr': ('https://github.com/tensorwerk/hangar-py/pull/%s', 'PR #'),
}
intersphinx_mapping = {
    'python': ('https://docs.python.org/3', None),
    'torch': ('https://pytorch.org/docs/master', None),
    'numpy': ('http://docs.scipy.org/doc/numpy', None),
}

# Regular expressions that match URIs that should not be checked
# when doing a linkcheck build
linkcheck_ignore = [
    r'http://localhost:\d+/?', 'http://localhost/',
    'https://github.com/tensorwerk/hangar-py',
    r'https://github.com/tensorwerk/hangar-py/.*',
    r'http://tensorwerk.com/hangar-benchmarks/',
    r'https://tensorwerk.com/hangar-benchmarks',
]
linkcheck_retries = 3

# on_rtd is whether we are on readthedocs.org
on_rtd = os.environ.get('READTHEDOCS', None) == 'True'

# if not on_rtd:  # only set the theme if we're building docs locally
# html_theme = 'sphinx_rtd_theme'
html_theme = 'sphinx_material'

html_sidebars = {
    "**": ["logo-text.html", "globaltoc.html", "localtoc.html", "searchbox.html"]
}

html_short_title = '%s-%s' % (project, version)

napoleon_use_ivar = True
napoleon_use_rtype = True
napoleon_use_param = True
napoleon_include_init_with_doc = True

add_module_names = False
doctest_test_doctest_blocks = None
autoclass_content = 'class'

# Material theme options (see theme.conf for more information)
html_theme_options = {

    # Set the name of the project to appear in the navigation.
    'nav_title': 'Hangar',

    # Set the color and the accent color
    'color_primary': 'deep-purple',
    'color_accent': 'blue',

    # Set the repo location to get a badge with stats
    'repo_url': 'https://github.com/tensorwerk/hangar-py/',
    'repo_name': 'Hangar',
    'repo_type': 'github',

    # Visible levels of the global TOC; -1 means unlimited
    'globaltoc_depth': -1,
    # If False, expand all TOC entries
    'globaltoc_collapse': True,
    # If True, show hidden TOC entries
    'globaltoc_includehidden': True,
}


from __future__ import annotations
from csv import DictReader
from os import path
from typing import TYPE_CHECKING
from warnings import warn

if TYPE_CHECKING:
    from ..DhsArticle import DhsArticle

script_folder = path.dirname(__file__)

DEFAULT_WIKIDATA_LINKS_FILE = path.join(script_folder, "wikidata_dhs_wikipedia_articles_gndid_instanceof.csv")
WIKIDATA_QUERY_FILE = path.join(script_folder, "wikidata_dhs_wikipedia_articles_gndid.sparql")

LOADED_WIKIDATA_LINKS_CSVS = set()
WIKIDATA_LINKS = dict()

WIKIDATA_URL_KEY = "item"

SPARQL_DOWNLOAD_DISCLAIMER = \
    f"A prerequisite is to have manually downloaded the result of the sparql query in file '{WIKIDATA_QUERY_FILE}' at 'https://query.wikidata.org/' " + \
    f"as a csv at this location '{DEFAULT_WIKIDATA_LINKS_FILE}' (or provide to this function the location of your csv file as function argument)."

def load_wikidata_links(wikidata_links_file = DEFAULT_WIKIDATA_LINKS_FILE):
    """Loads the csv links in a dictionary of the form dhsid->list(linked wikidata_wikipedia entities)\n\n""" + SPARQL_DOWNLOAD_DISCLAIMER
    if wikidata_links_file not in LOADED_WIKIDATA_LINKS_CSVS:
        if not path.exists(wikidata_links_file):
            raise Exception(
                f"dhs_scraper.wikidata.load_wikidata_links() wikidata_links_file at location '{wikidata_links_file}' not found.\n"+
                SPARQL_DOWNLOAD_DISCLAIMER
            )
        with open(wikidata_links_file) as f:
            reader = DictReader(f)
            for r in reader:
                links_list = WIKIDATA_LINKS.get(r["dhsid"])
                if links_list is not None:
                    links_list.append(r)
                else:
                    WIKIDATA_LINKS[r["dhsid"]] = [r]
        LOADED_WIKIDATA_LINKS_CSVS.add(wikidata_links_file)
    return WIKIDATA_LINKS

def get_wikidata_links_from_dhs_id(dhs_id, wikidata_links_file = DEFAULT_WIKIDATA_LINKS_FILE):
    load_wikidata_links(wikidata_links_file)
    wiki_links = WIKIDATA_LINKS.get(dhs_id)
    if wiki_links is not None:
        return wiki_links
    else:
        return []

def get_wikidata_main_link_from_dhs_id(dhs_id, language:str, wikidata_links_file = DEFAULT_WIKIDATA_LINKS_FILE):
    load_wikidata_links(wikidata_links_file)
    wiki_links = get_wikidata_links_from_dhs_id(dhs_id)
    wikipedia_page_title_key = "name"+language
    wd_url_wk_title = (None,None)
    for l in wiki_links:
        if l[wikipedia_page_title_key] is not None and l[wikipedia_page_title_key] not in ["", "null"]:
            wd_url_wk_title = (l[WIKIDATA_URL_KEY], l[wikipedia_page_title_key])
            break
        else:
            wd_url_wk_title = (l[WIKIDATA_URL_KEY], None)
    return wd_url_wk_title

def add_wikidata_wikipedia_to_text_links(dhs_article:DhsArticle, wikidata_links_file = DEFAULT_WIKIDATA_LINKS_FILE):
    """adds "wiki_links" attribute to dhs_article.text_links, also adds one as main "wikidata_url" and "wikipedia_page_title".

    Adds all the wikidata entities that have a wikidata P902 property pointing to the given dhs_id to "wiki_links".
    For the "wikidata_url" and "wikipedia_page_title" attributes, takes the first entries that has a wikipedia_page_title in given language.
    """
    load_wikidata_links(wikidata_links_file)
    if "text_links" in dhs_article.__dict__:
        for block_links in dhs_article.text_links:
            for link in block_links:
                lng, dhsid, v = dhs_article.get_language_id_version_from_url(link["href"])
                wiki_links = get_wikidata_links_from_dhs_id(dhsid)
                link["wiki_links"] = wiki_links
                wikidata_url, wikipedia_page_title = get_wikidata_main_link_from_dhs_id(dhsid, dhs_article.language)
                link["wikidata_url"] = wikidata_url
                link["wikipedia_page_title"] = wikipedia_page_title

    else:
        warn("dhs_scraper.wikidata.dhs_article_add_wikidata_wikipedia_links() trying to add links to a DhsArticle having no text_links: skipping.")
    return dhs_article


#encoding=utf-8
import torch
import csv
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torch.optim as optim
import torch.utils.data as Data
import matplotlib.pyplot as plt
import numpy as np
class mydataset:
    def __init__(self, tens,lab):
        self.dataset =tens
        self.label=lab
    def __len__(self):
        return len(self.dataset)
    def __getitem__(self, idx):
        text = self.dataset[idx]
        label = self.label[idx]
        sample = {"text": text, "label": label}
        return sample
class Net(nn.Module):  # 定义网络，继承torch.nn.Module
    def __init__(self):
        super(Net, self).__init__()
        #self.pool = nn.AvgPool2d(2, 2)  # 池化层
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels=1,
                            out_channels=4,
                            kernel_size=3,
                            stride=2,
                            padding=1),
            torch.nn.BatchNorm2d(4),
            torch.nn.ReLU()
        )
        self.conv2 = torch.nn.Sequential(
            torch.nn.Conv2d(4, 8, 3, 2, 1),
            torch.nn.BatchNorm2d(8),
            torch.nn.ReLU()
        )
        self.conv3 = torch.nn.Sequential(
            torch.nn.Conv2d(8, 16, 3, 2, 1),
            torch.nn.BatchNorm2d(16),
 	    torch.nn.ReLU()
        )
        self.fc1 = nn.Linear(480,250)  # 全连接层
        self.fc2 = nn.Linear(250,120)
        self.fc3 = nn.Linear(120,6)

    def forward(self, x):  # 前向传播

        #x = self.pool(F.relu(self.conv1(x)))  # F就是torch.nn.functional
        #x = self.pool(F.relu(self.conv2(x)))
        x=self.conv1(x)
        x=self.conv2(x)
        x = self.conv3(x)
        x = self.fc1(x.view(x.size(0), -1))
        x = self.fc2(x) # .view( )是一个tensor的方法，使得tensor改变size但是元素的总数是不变的。
        x = self.fc3(x)
 # 从卷基层到全连接层的维度转换
        return x

def trainandsave():
    # 神经网络结构
    net = Net()
      # 学习率为0.001
    optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)
    criterion = nn.CrossEntropyLoss()  # 损失函数也可以自己定义，我们这里用的交叉熵损失函数
    # 训练部分
    for epoch in range(10):  # 训练的数据量为5个epoch，每个epoch为一个循环
        # 每个epoch要训练所有的图片，每训练完成200张便打印一下训练的效果（loss值）
        running_loss = 0.0  # 定义一个变量方便我们对loss进行输出
        for i, data in enumerate(trainloader):  # 这里我们遇到了第一步中出现的trailoader，代码传入数据
            # enumerate是python的内置函数，既获得索引也获得数据
            # get the inputs
            inputs= data["text"] # data是从enumerate返回的data，包含数据和标签信息，分别赋值给inputs和labels
    #        print(inputs.shape)
            labels=data["label"]
            # wrap them in Variable
            inputs, labels = Variable(inputs), Variable(labels)  # 转换数据格式用Variable
            optimizer.zero_grad()  # 梯度置零，因为反向传播过程中梯度会累加上一次循环的梯度
            # forward + backward + optimize
            inputs=inputs.float()
            labels=labels.float()
            outputs = net(inputs)  # 把数据输进CNN网络net
            loss = criterion(outputs, labels.long())  # 计算损失值
            loss.backward()  # loss反向传播
            optimizer.step()  # 反向传播后参数更新
            running_loss += loss.item()  # loss累加
            if i % 2000 == 0:
                print('[%d, %5d] loss: %.3f' %
                      (epoch + 1, i + 1, running_loss / 2000))  # 然后再除以200，就得到这两百次的平均损失值
                running_loss = 0.0  # 这一个200次结束后，就把running_loss归零，下一个200次继续使用
    print('Finished Training')
    # 保存神经网络
    return net
if __name__ =='__main__':
    ttens = []
    tlab=[]
    etens = []
    elab=[]
    i=0
    count2=0
    with open("/home/zc/lyd/train20-40b.csv") as f:
         row=csv.reader(f,delimiter = ',')
         next(row)
         for r in row:
             keyt=int(r[1])
             lab=int(r[2][:-1])
            # if((lab==1)and(count2<=20000)):
            #    count2+=1
            # elif((lab==1)and(count2>20000)):
            #    continue
            # if((lab==1) or (lab==2)):
            #     lab=2
            # if((lab==2)and(count2<=25000)):
            #     count2+=1
            # elif((lab==2)and(count2>25000)):
            #     continue
            # if(lab>4):
            #     continue
             if(keyt==9):
                if(i%5!=0):
                    i=i+1
                    tlab.append(lab)
                    temp=r[0][1:-1].split('|')
                    tempten=[]
                    for j in temp:
                        line=j.split('/')
                        tempten.append(list(map(int,line)))
                    ttens.append(torch.tensor([tempten]))
                else:
                    i=i+1
                    elab.append(lab)
                    temp = r[0][1:-1].split('|')
                    tempten = []
                    for j in temp:
                        line = j.split('/')
                        tempten.append(list(map(int, line)))
                    etens.append(torch.tensor([tempten]))
    td=mydataset(ttens,torch.tensor(tlab))
    test=mydataset(etens,torch.tensor(elab))
    BATCH_SIZE = 8
    trainloader = Data.DataLoader(
        dataset=td,
        batch_size=BATCH_SIZE,
        num_workers=0,
        drop_last=True
    )
    testloader = Data.DataLoader(
        dataset=test,
        batch_size=BATCH_SIZE,
        num_workers=0,
        drop_last=True
    )

    model = trainandsave()
    accuracy_sum = []
    for i,data in enumerate(testloader):
        test_x = Variable(data["text"]).float()
        test_y = Variable(data["label"]).float()
    #    print(test_x.shape)
        out = model(test_x)
        # print('test_out:\t',torch.max(out,1)[1])
        # print('test_y:\t',test_y)
        accuracy = torch.max(out,1)[1].numpy() == test_y.numpy()
        accuracy_sum.append(accuracy.mean())
       # print(torch.max(out,1)[1].numpy(),test_y.numpy())
    print('总准确率：\t',sum(accuracy_sum)/len(accuracy_sum))


import timeit
import random
import matplotlib.pyplot as plt
import pandas as pd
import os

from ordered_linked_list import OrderedLinkedList
from binary_search_tree import BinarySearchTree
from red_black_tree import RedBlackTree

n = 5006
step = 50
test_per_iteration = 30

def random_array(n):
    array = []
    for i in range(0, n):
        array.append(i)
    random.shuffle(array)
    return array

# Creo le funzioni di test
# Funzione di test dell'inserimento di un valore in una struttura dati
def measure_insert_test(insert_function, array):
    times = []
    t = 0
    for j in range(len(array)):
        # Misuro il tempo di esecuzione di una singola iterazione
        times.append(timeit.timeit(stmt=lambda: insert_function(array[j]), number=1))
        t += times[j]
    return (t / len(array)) * 1000  # tempo in ms

# Funzione di test della ricerca del k-esimo elemento più piccolo
def measure_os_test(os_function, start, length):
    k = length // 2  # Cerco l'elemento più o meno centrale
    return timeit.timeit(stmt=lambda: os_function(start, k),
                         number=test_per_iteration) / test_per_iteration * 1000  # tempo in ms

# Funzione di test della ricerca del rank di un elemento
def measure_rank_test(rank_function, ric):
    return timeit.timeit(stmt=lambda: rank_function(ric),
                         number=test_per_iteration) / test_per_iteration * 1000  # tempo in ms)

def test(oll_i, oll_os, oll_r, bst_i, bst_os, bst_r, rbt_i, rbt_os, rbt_r, n, step):
    # Creo le strutture dati
    oll = OrderedLinkedList()
    bst = BinarySearchTree()
    rbt = RedBlackTree()

    # Esecuzione test
    for i in range(5, n, step):
        # Creo un array con i elementi random tra 0 e n
        arr = random_array(i)

        # svolgo i test sull'inserimento e mi salvo i risultati espressi in ms
        oll_i.append(measure_insert_test(oll.insert, arr))
        bst_i.append(measure_insert_test(bst.insert, arr))
        rbt_i.append(measure_insert_test(rbt.insert, arr))

        # Svolgo i test sulla ricerca del k-esimo elemento più piccolo e mi salvo i risultati espressi in ms
        oll_os.append(measure_os_test(oll.order_statistic, oll.head, i))
        bst_os.append(measure_os_test(bst.get_kesimo, bst.root, i))
        rbt_os.append(measure_os_test(rbt.os_select, rbt.root, i))

        # Svolgo i test sulla ricerca del rango di un elemento casuale dell'array
        oll_r.append(measure_rank_test(oll.oll_rank, oll.search(oll.head, i // 2)))
        bst_r.append(measure_rank_test(bst.get_rank, bst.search(bst.root, i // 2)))
        rbt_r.append(measure_rank_test(rbt.os_rank, rbt.search(rbt.root, i // 2)))

        # Svuoto le liste
        oll = OrderedLinkedList()
        bst = BinarySearchTree()
        rbt = RedBlackTree()

    return oll_i, oll_os, oll_r, bst_i, bst_os, bst_r, rbt_i, rbt_os, rbt_r

if __name__ == '__main__':

    # Creo le liste in cui salvare i risultati dei test
    assex = []

    # Serve per i grafici
    for i in range(5, n, step):
        assex.append(i)

    oll_insert_times = [0.0] * len(assex)
    bst_insert_times = [0.0] * len(assex)
    rbt_insert_times = [0.0] * len(assex)
    oll_order_statistic_times = [0.0] * len(assex)
    bst_order_statistic_times = [0.0] * len(assex)
    rbt_order_statistic_times = [0.0] * len(assex)
    oll_rank_times = [0.0] * len(assex)
    bst_rank_times = [0.0] * len(assex)
    rbt_rank_times = [0.0] * len(assex)

    print("Parto con l'esecuzione dei test")

    # Sommo tutti i valori dei vari test
    for j in range(test_per_iteration):
        print("Eseguo il test numero: ", (j+1), "/", test_per_iteration)
        oll_temp_i = []
        bst_temp_i = []
        rbt_temp_i = []
        oll_temp_os = []
        bst_temp_os = []
        rbt_temp_os = []
        oll_temp_r = []
        bst_temp_r = []
        rbt_temp_r = []
        oll_temp_i, oll_temp_os, oll_temp_r, bst_temp_i, bst_temp_os, bst_temp_r, rbt_temp_i, rbt_temp_os, rbt_temp_r = test(
            oll_temp_i, oll_temp_os, oll_temp_r,
            bst_temp_i, bst_temp_os, bst_temp_r,
            rbt_temp_i, rbt_temp_os, rbt_temp_r, n,
            step)

        for k in range(len(oll_temp_i)):
            # Inserimento
            oll_insert_times[k] += oll_temp_i[k]
            bst_insert_times[k] += bst_temp_i[k]
            rbt_insert_times[k] += rbt_temp_i[k]

            # Statistica D'Ordine
            oll_order_statistic_times[k] += oll_temp_os[k]
            bst_order_statistic_times[k] += bst_temp_os[k]
            rbt_order_statistic_times[k] += rbt_temp_os[k]

            # Rango
            oll_rank_times[k] += oll_temp_r[k]
            bst_rank_times[k] += bst_temp_r[k]
            rbt_rank_times[k] += rbt_temp_r[k]

    print("Faccio le medie")

    # Faccio le medie
    for s in range(0, len(assex)):
        # Inserimento
        oll_insert_times[s] = oll_insert_times[s] / test_per_iteration
        bst_insert_times[s] = bst_insert_times[s] / test_per_iteration
        rbt_insert_times[s] = rbt_insert_times[s] / test_per_iteration

        # Statistica d'ordine
        oll_order_statistic_times[s] = oll_order_statistic_times[s] / test_per_iteration
        bst_order_statistic_times[s] = bst_order_statistic_times[s] / test_per_iteration
        rbt_order_statistic_times[s] = rbt_order_statistic_times[s] / test_per_iteration

        # Rango
        oll_rank_times[s] = oll_rank_times[s] / test_per_iteration
        bst_rank_times[s] = bst_rank_times[s] / test_per_iteration
        rbt_rank_times[s] = rbt_rank_times[s] / test_per_iteration

    # Funzione per svuotare una cartella
    def svuota_cartella(cartella):
        file_lista = os.listdir(cartella)
        for file in file_lista:
            file_path = os.path.join(cartella, file)
            os.remove(file_path)

    # Svuota la cartella "tabelle"
    svuota_cartella("tabelle")

    # Svuota la cartella "immagini"
    svuota_cartella("immagini")

    # Grafici test inserimento
    print("Creo i grafici")

    # Grafico inserimento lista linkata ordinata
    plt.clf()
    plt.plot(assex, oll_insert_times, color='blue', label='Inserimento lista linkata ordinata')
    plt.xlabel('Dimensione dell\'array')
    plt.ylabel('Tempo di esecuzione (ms)')
    plt.title('Tempi inserimento lista ordinata linkata')
    plt.legend()
    # Salvo i grafdici nella cartella immagini
    plt.savefig('immagini/oll_ins.png')
    plt.show()

    # Grafico inserimento albero binario
    plt.clf()
    plt.plot(assex, bst_insert_times, color='green', label='Inserimento albero binario')
    plt.xlabel('Dimensione dell\'array')
    plt.ylabel('Tempo di esecuzione (ms)')
    plt.title('Tempi inserimento albero binario')
    plt.legend()
    plt.savefig('immagini/bst_ins.png')
    plt.show()

    # Grafico inserimento albero rosso-nero
    plt.clf()
    plt.plot(assex, rbt_insert_times, color='red', label='Inserimento albero rosso-nero')
    plt.xlabel('Dimensione dell\'array')
    plt.ylabel('Tempo di esecuzione (ms)')
    plt.title('Tempi inserimento albero rosso-nero')
    plt.legend()
    plt.savefig('immagini/rbt_ins.png')
    plt.show()

    # Grafico confronto tra i tre tempi di inserimento
    plt.clf()
    plt.plot(assex, oll_insert_times, color='blue', label='Inserimento lista linkata ordinata')
    plt.plot(assex, bst_insert_times, color='green', label='Inserimento albero binario')
    plt.plot(assex, rbt_insert_times, color='red', label='Inserimento albero rosso-nero')
    plt.xlabel('Dimensione dell\'array')
    plt.ylabel('Tempo di esecuzione (ms)')
    plt.title('Confronto tempi inserimento')
    plt.legend()
    plt.savefig('immagini/conf_ins.png')
    plt.show()

    # Grafici ricerca del k-esimo elemento più piccolo

    # Grafico ricerca del k-esimo elemento più piccolo in una lista linkata ordinata
    plt.clf()
    plt.plot(assex, oll_order_statistic_times, color='blue', label='Ricerca k-esimo lista linkata ordinata')
    plt.xlabel('Dimensione dell\'array')
    plt.ylabel('Tempo di esecuzione (ms)')
    plt.title('Tempi ricerca k-esimo lista linkata ordinata')
    plt.legend()
    plt.savefig('immagini/oll_os.png')
    plt.show()

    # Grafico ricerca del k-esimo elemento più piccolo in un albero binario
    plt.clf()
    plt.plot(assex, bst_order_statistic_times, color='green', label='Ricerca k-esimo albero binario')
    plt.xlabel('Dimensione dell\'array')
    plt.ylabel('Tempo di esecuzione (ms)')
    plt.title('Tempi ricerca k-esimo albero binario')
    plt.legend()
    plt.savefig('immagini/bst_os.png')
    plt.show()

    # Grafico ricerca del k-esimo elemento più piccolo in un albero rosso-nero
    plt.clf()
    plt.plot(assex, rbt_order_statistic_times, color='red', label='Ricerca k-esimo albero rosso-nero')
    plt.xlabel('Dimensione dell\'array')
    plt.ylabel('Tempo di esecuzione (ms)')
    plt.title('Tempi ricerca k-esimo albero rosso-nero')
    plt.legend()
    plt.savefig('immagini/rbt_os.png')
    plt.show()

    # Grafico confronto tra le ricerche dei k-esimi elementi più piccoli
    plt.clf()
    plt.plot(assex, oll_order_statistic_times, color='blue', label='Ricerca k-esimo lista linkata ordinata')
    plt.plot(assex, bst_order_statistic_times, color='green', label='Ricerca k-esimo albero binario')
    plt.plot(assex, rbt_order_statistic_times, color='red', label='Ricerca k-esimo albero rosso-nero')
    plt.xlabel('Dimensione dell\'array')
    plt.ylabel('Tempo di esecuzione (ms)')
    plt.title('Tempi ricerca k-esimo lista linkata ordinata')
    plt.legend()
    plt.savefig('immagini/conf_os.png')
    plt.show()

    # Grafici ricerca rank di un elemento

    # Grafico ricerca rank di un elemento in una lista ordinata
    plt.clf()
    plt.plot(assex, oll_rank_times, color='blue', label='Ricerca rank di un elemento lista linkata ordinata')
    plt.xlabel('Dimensione dell\'array')
    plt.ylabel('Tempo di esecuzione (ms)')
    plt.title('Tempi ricerca rank di un elemento lista linkata ordinata')
    plt.legend()
    plt.savefig('immagini/oll_rank.png')
    plt.show()

    # Grafico ricerca rank di un elemento in un albero binario
    plt.clf()
    plt.plot(assex, bst_rank_times, color='green', label='Ricerca rank di un elemento in un albero binario')
    plt.xlabel('Dimensione dell\'array')
    plt.ylabel('Tempo di esecuzione (ms)')
    plt.title('Tempi ricerca rank di un elemento in un albero binario')
    plt.legend()
    plt.savefig('immagini/bst_rank.png')
    plt.show()

    # Grafico ricerca rank di un elemento in un albero rosso-nero
    plt.clf()
    plt.plot(assex, rbt_rank_times, color='red', label='Ricerca rank di un elemento in un albero rosso-nero')
    plt.xlabel('Dimensione dell\'array')
    plt.ylabel('Tempo di esecuzione (ms)')
    plt.title('Tempi ricerca rank di un elemento in un albero rosso-nero')
    plt.legend()
    plt.savefig('immagini/rbt_rank.png')
    plt.show()

    # Confronto grafici ricerca rank di un elemento
    plt.clf()
    plt.plot(assex, oll_rank_times, color='blue', label='Ricerca rank di un elemento lista linkata ordinata')
    plt.plot(assex, bst_rank_times, color='green', label='Ricerca rank di un elemento in un albero binario')
    plt.plot(assex, rbt_rank_times, color='red', label='Ricerca rank di un elemento in un albero rosso-nero')
    plt.xlabel('Dimensione dell\'array')
    plt.ylabel('Tempo di esecuzione (ms)')
    plt.title('Tempi ricerca rank di un elemento in un albero rosso-nero')
    plt.legend()
    plt.savefig('immagini/conf_rank.png')
    plt.show()

    # Creazione delle tabelle

    # Creazione delle liste dei dati per inserimento, ricerca k-esima statistica d'ordine e rank
    data_inserimento = {'Dimensione dell\'array': assex, 'Lista linkata ordinata': oll_insert_times,
                        'Albero binario': bst_insert_times, 'Albero rosso-nero': rbt_insert_times}
    df_ins = pd.DataFrame(data_inserimento)
    data_order_statistic = {'Dimensione dell\'array': assex, 'Lista linkata ordinata': oll_order_statistic_times,
                            'Albero binario': bst_order_statistic_times,
                            'Albero rosso-nero': rbt_order_statistic_times}
    df_os = pd.DataFrame(data_order_statistic)
    data_rank = {'Dimensione dell\'array': assex, 'Lista linkata ordinata': oll_rank_times,
                 'Albero binario': bst_rank_times, 'Albero rosso-nero': rbt_rank_times}
    df_rank = pd.DataFrame(data_rank)

    # Approssimazione dei tempi di esecuzione
    df_ins = df_ins.round(4)
    df_os = df_os.round(4)
    df_rank = df_rank.round(4)

    plt.clf()

    # Visualizzazione delle tabelle come immagini separate
    plt.figure(figsize=(10, 6))  # Imposta la dimensione dell'immagine

    # Tabella 1: Inserimento
    plt.title("Inserimento")
    plt.axis('off')  # Rimuove gli assi
    table1 = plt.table(cellText=df_ins.values, colLabels=df_ins.columns, loc='center')
    table1.auto_set_font_size(False)
    table1.set_fontsize(10)
    plt.savefig('tabelle/tabella_inserimento.png')  # Salva l'immagine della tabella di inserimento come file PNG
    plt.show()

    # Creazione di una nuova figura per la tabella di ricerca della k-esima statistica d'ordine
    plt.clf()
    plt.figure(figsize=(10, 6))
    plt.title("Ricerca K-esima Statistica d'Ordine")
    plt.axis('off')
    table2 = plt.table(cellText=df_os.values, colLabels=df_os.columns, loc='center')
    table2.auto_set_font_size(False)
    table2.set_fontsize(10)
    plt.savefig('tabelle/tabella_kesima_statistica_ordine.png')
    plt.show()

    # Creazione di una nuova figura per la tabella di ricerca rank
    plt.clf()
    plt.figure(figsize=(10, 6))
    plt.title("Ricerca Rank")
    plt.axis('off')
    table3 = plt.table(cellText=df_rank.values, colLabels=df_rank.columns, loc='center')
    table3.auto_set_font_size(False)
    table3.set_fontsize(10)
    plt.savefig('tabelle/tabella_ricerca_rank.png')
    plt.show()


'''
******* 문제 푼 후 느낀 것 *********
처음에 그냥 comnbination 을 구해서 처리했더니, 메모리 초과가 발생하였다.
100개의 아이템들이 있을때 50개를 구한다하면 100C50 으로 제한된 메모리 10,000,000 을 쉽게 넘어버린다.

따라서 이는 다른 방법으로 풀도록 해야한다. 아니면 메모리를 적게 사용하도록 dfs 로 조합을 직접 구현?
dp를 이용해서 구현해야한다. 정말 어려운 문제 같다. 이런 문제를 보고 바로 dp 를 생각하는 것이 어려운 것 같다.

문제 조건의 N 의 수 등을 잘 확인해야할 것 같다...
'''

'''
2022/07/30
boj.kr/7579
앱

문제 :
우리는 스마트폰을 사용하면서 여러 가지 앱(App)을 실행하게 된다. 대개의 경우 화면에 보이는 ‘실행 중’인 앱은 하나뿐이지만 보이지 않는 상태로 많은 앱이 '활성화'되어 있다.
앱들이 활성화 되어 있다는 것은 화면에 보이지 않더라도 메인 메모리에 직전의 상태가 기록되어 있는 것을 말한다.
현재 실행 중이 아니더라도 이렇게 메모리에 남겨두는 이유는 사용자가 이전에 실행하던 앱을 다시 불러올 때에 직전의 상태를 메인 메모리로부터 읽어 들여 실행 준비를 빠르게 마치기 위해서이다.

하지만 스마트폰의 메모리는 제한적이기 때문에 한번이라도 실행했던 모든 앱을 활성화된 채로 메인 메모리에 남겨두다 보면 메모리 부족 상태가 오기 쉽다.
새로운 앱을 실행시키기 위해 필요한 메모리가 부족해지면 스마트폰의 운영체제는 활성화 되어 있는 앱들 중 몇 개를 선택하여 메모리로부터 삭제하는 수밖에 없다.
이러한 과정을 앱의 ‘비활성화’라고 한다.

메모리 부족 상황에서 활성화 되어 있는 앱들을 무작위로 필요한 메모리만큼 비활성화 하는 것은 좋은 방법이 아니다.
비활성화된 앱들을 재실행할 경우 그만큼 시간이 더 필요하기 때문이다. 여러분은 이러한 앱의 비활성화 문제를 스마트하게 해결하기 위한 프로그램을 작성해야 한다

현재 N개의 앱, A1, ..., AN이 활성화 되어 있다고 가정하자. 이들 앱 Ai는 각각 mi 바이트만큼의 메모리를 사용하고 있다.
또한, 앱 Ai를 비활성화한 후에 다시 실행하고자 할 경우, 추가적으로 들어가는 비용(시간 등)을 수치화 한 것을 ci 라고 하자.
이러한 상황에서 사용자가 새로운 앱 B를 실행하고자 하여, 추가로 M 바이트의 메모리가 필요하다고 하자.
즉, 현재 활성화 되어 있는 앱 A1, ..., AN 중에서 몇 개를 비활성화 하여 M 바이트 이상의 메모리를 추가로 확보해야 하는 것이다.
여러분은 그 중에서 비활성화 했을 경우의 비용 ci의 합을 최소화하여 필요한 메모리 M 바이트를 확보하는 방법을 찾아야 한다.

입력:
입력은 3줄로 이루어져 있다. 첫 줄에는 정수 N과 M이 공백문자로 구분되어 주어지며, 둘째 줄과 셋째 줄에는 각각 N개의 정수가 공백문자로 구분되어 주어진다.
둘째 줄의 N개의 정수는 현재 활성화 되어 있는 앱 A1, ..., AN이 사용 중인 메모리의 바이트 수인 m1, ..., mN을 의미하며, 셋째 줄의 정수는 각 앱을 비활성화 했을 경우의 비용 c1, ..., cN을 의미한다

단, 1 ≤ N ≤ 100, 1 ≤ M ≤ 10,000,000이며, 1 ≤ m1, ..., mN ≤ 10,000,000을 만족한다. 또한, 0 ≤ c1, ..., cN ≤ 100이고, M ≤ m1 + m2 + ... + mN이다.

출력:
필요한 메모리 M 바이트를 확보하기 위한 앱 비활성화의 최소의 비용을 계산하여 한 줄에 출력해야 한다.
'''
import sys
from itertools import combinations
input = sys.stdin.readline

N, M = map(int, input().split())

memory_cost_list =  [0] + list(map(int, input().split()))
off_cost_list = [0] + list(map(int, input().split()))

max_cost = sum(off_cost_list)

# 메모리 비용을 기준으로 DP 를 하게 되면, 최대 10,000,000 * 100 개 까지 생성될 수 있기 때문에
# 너무 많은 메모리를 차지하게 된다.
# off 비용을 기준으로 하게되면 100 * 100 으로 10,000 이기 때문에 이것을 이용하도록 한다.
dp = [[0 for _ in range(max_cost+1)] for _ in range(N+1)]

answer = sys.maxsize

for i in range(1, N+1):
    for j in range(1, max_cost+1):
        # 현재 확인하는 비용이 현재 확인하는 메모리공간의 교체 비용보다 작다면 그냥 i-1에 있는 것을 사용한다
        if j < off_cost_list[i]:
            dp[i][j] = dp[i-1][j]
        else :
            dp[i][j] = max(dp[i-1][j], dp[i-1][j-off_cost_list[i]] + memory_cost_list[i])
            if dp[i][j] >= M :
                answer = min(answer, j)
                break

print(answer)


from django.contrib.auth import logout
from django.contrib.auth.decorators import login_required, user_passes_test
from django.contrib.auth.models import User
from django.core.exceptions import ObjectDoesNotExist, PermissionDenied
from django.db.models import Count
from django.forms import (
    ModelForm,
    ModelMultipleChoiceField,
    modelformset_factory,
)
from django.shortcuts import get_object_or_404, render
from django.urls import reverse
from django.utils.translation import gettext_lazy as _
from django.views.decorators.http import require_http_methods
from django.views.generic.edit import UpdateView

from ..mixins import LoginRequiredMixin, NoStudentsMixin, student_check
from ..models import (
    Answer,
    AnswerAnnotation,
    Assignment,
    Discipline,
    Question,
    QuestionFlag,
    QuestionFlagReason,
    ShownRationale,
)

@login_required
@user_passes_test(student_check, login_url="/access_denied_and_logout/")
def research_index(request):
    template = "peerinst/research/index.html"
    context = {"disciplines": Discipline.objects.all()}
    return render(request, template, context)

def get_question_annotation_counts(discipline_title, annotator, assignment_id):
    """
    Returns:
    ========
    list of dicts, one for each question in discipline, and the counts of
    Answers and AnswerAnnotations for each
    """
    if discipline_title:
        questions_qs = Question.objects.filter(
            discipline__title=discipline_title
        )
    elif assignment_id:
        questions_qs = Assignment.objects.get(
            identifier=assignment_id
        ).questions.order_by("assignmentquestions__rank")

    # FIXME:
    translation_table = str.maketrans("ABCDEFG", "1234567")

    question_annotation_counts = []
    for q in questions_qs:
        d1 = {}
        d1["question"] = q

        d1["question_expert_answers"] = q.answer_set.filter(expert=True)

        # need at least one sample answer that is not marked as expert for each
        # answer choice
        d1["enough_sample_answers"] = 0 not in list(
            q.get_frequency(all_rationales=True)["first_choice"].values()
        )

        d1["total_annotations"] = AnswerAnnotation.objects.filter(
            score__isnull=False, answer__question_id=q.pk
        ).count()
        d1["total_annotations_by_user"] = AnswerAnnotation.objects.filter(
            score__isnull=False, answer__question_id=q.pk, annotator=annotator
        ).count()
        flagged_questions = Question.flagged_objects.all()
        flagged_by_user = QuestionFlag.objects.filter(
            user=annotator
        ).values_list("question", flat=True)
        if q.pk in flagged_by_user:
            d1["flag_color_code"] = "red"
        elif q in flagged_questions:
            d1["flag_color_code"] = "goldenrod"
            d1["flagged_reasons"] = "; ".join(
                [
                    " - ".join(e)
                    for e in QuestionFlag.objects.filter(
                        flag=True, question=q
                    ).values_list("user__username", "flag_reason__title")
                ]
            )
        else:
            d1["flag_color_code"] = None

        answer_frequencies = q.get_frequency_json("first_choice")
        for d2 in answer_frequencies:
            a_choice = str(d2["answer_label"][0]).translate(translation_table)
            d2.update(
                {
                    "annotation_count": AnswerAnnotation.objects.filter(
                        score__isnull=False,
                        answer__question_id=q.pk,
                        answer__first_answer_choice=a_choice,
                    ).count(),
                    "annotation_count_by_user": AnswerAnnotation.objects.filter(  # noqa
                        score__isnull=False,
                        answer__question_id=q.pk,
                        answer__first_answer_choice=a_choice,
                        annotator=annotator,
                    ).count(),
                }
            )
        d1["answerchoices"] = answer_frequencies
        question_annotation_counts.append(d1)

    return question_annotation_counts

@login_required
@user_passes_test(student_check, login_url="/access_denied_and_logout/")
def research_discipline_question_index(
    request, discipline_title=None, assignment_id=None
):
    template = "peerinst/research/question_index.html"

    annotator = get_object_or_404(User, username=request.user)

    question_annotation_counts = get_question_annotation_counts(
        discipline_title=discipline_title,
        annotator=annotator,
        assignment_id=assignment_id,
    )
    request.session["assignment_id"] = assignment_id

    context = {
        "questions": question_annotation_counts,
        "discipline_title": discipline_title,
        "assignment_id": assignment_id,
        "annotations_count": AnswerAnnotation.objects.filter(
            annotator=annotator, score__isnull=False
        ).count(),
        "annotator": annotator,
    }
    return render(request, template, context)

@require_http_methods(["GET", "POST"])
@login_required
@user_passes_test(student_check, login_url="/access_denied_and_logout/")
def research_question_answer_list(
    request,
    question_pk,
    answerchoice_value,
    discipline_title=None,
    assignment_id=None,
):
    template = "peerinst/research/answer_list.html"
    if discipline_title:
        reverse_url_name = "research-question-answer-list-by-discipline"
    elif assignment_id:
        reverse_url_name = "research-question-answer-list-by-assignment"

    annotator = get_object_or_404(User, username=request.user)
    if not annotator:
        logout(request)
        raise PermissionDenied

    # FIXME
    if answerchoice_value.isdigit():
        answerchoice_id = answerchoice_value
    else:
        answerchoice_id = {"A": 1, "B": 2, "C": 3, "D": 4, "E": 5, "F": 6}[
            answerchoice_value
        ]

    answer_qs = Question.objects.get(pk=question_pk).answer_set.filter(
        question_id=question_pk, first_answer_choice=answerchoice_id
    )
    for a in answer_qs:
        annotation, created = AnswerAnnotation.objects.get_or_create(
            answer=a,
            annotator=annotator,
        )

    # only need two expert scores per rationale,
    # and those marked never show even by one person can be excluded
    already_scored = (
        AnswerAnnotation.objects.filter(
            answer__question_id=question_pk,
            score__isnull=False,
            answer__first_answer_choice=answerchoice_id,
        )
        .values("answer")
        .order_by("answer")
        .annotate(times_scored=Count("answer"))
        .filter(times_scored__gte=3)
        .values_list("answer__id", flat=True)
    )

    queryset = (
        AnswerAnnotation.objects.filter(
            answer__question_id=question_pk,
            annotator=annotator,
            answer__first_answer_choice=answerchoice_id,
        )
        .exclude(answer__id__in=already_scored)
        .annotate(times_shown=Count("answer__shown_answer"))
        .order_by("-times_shown")
    )

    AnswerAnnotationFormset = modelformset_factory(
        AnswerAnnotation, fields=("score",), extra=0
    )

    formset = AnswerAnnotationFormset(queryset=queryset)

    context = {
        "formset": formset,
        "question": Question.objects.get(id=question_pk),
        "discipline_title": discipline_title,
        "question_pk": question_pk,
        "assignment_id": assignment_id,
        "annotations_count": AnswerAnnotation.objects.filter(
            annotator=annotator,
            score__isnull=False,
            answer__question_id=question_pk,
            answer__first_answer_choice=answerchoice_id,
        ).count(),
        "annotator": annotator,
    }

    if request.method == "POST":
        formset = AnswerAnnotationFormset(request.POST)
        if formset.is_valid():
            instances = formset.save()
            context.update(message=_("Scores updated"))

    return render(request, template, context)

def research_all_annotations_for_question(
    request, question_pk, discipline_title=None, assignment_id=None
):
    """
    Returns:
    ====================
    All answer annotations for a question, ordered for
    comparison by researchers
    """
    template = "peerinst/research/all_annotations.html"
    question = Question.objects.get(id=question_pk)
    context = {
        "question": question,
        "question_pk": question_pk,
        "discipline_title": discipline_title,
        "assignment_id": assignment_id,
    }

    all_annotations = []
    for label, _answerchoice_text in question.get_choices():
        d1 = {}
        d1["answerchoice"] = label
        d1["annotations"] = []

        # FIXME
        if label.isdigit():
            answerchoice_id = label
        else:
            answerchoice_id = {"A": 1, "B": 2, "C": 3, "D": 4, "E": 5, "F": 6}[
                label
            ]

        a_list = list(
            set(
                AnswerAnnotation.objects.filter(
                    score__isnull=False,
                    answer__question_id=question_pk,
                    answer__first_answer_choice=answerchoice_id,
                ).values_list("answer", flat=True)
            )
        )

        a_list.sort()

        for a in a_list:
            d2 = {}
            d2["answer"] = Answer.objects.get(pk=a)
            d2["scores"] = AnswerAnnotation.objects.filter(
                answer=a, score__isnull=False
            ).values("score", "annotator__username")
            d2["times_shown"] = ShownRationale.objects.filter(
                shown_answer=a
            ).count()
            d2["times_chosen"] = Answer.objects.filter(
                chosen_rationale_id=a
            ).count()
            d1["annotations"].append(d2)

        all_annotations.append(d1)

    context["all_annotations"] = all_annotations

    return render(request, template, context)

class QuestionFlagForm(ModelForm):
    flag_reason = ModelMultipleChoiceField(
        queryset=QuestionFlagReason.objects.all()
    )

    class Meta:
        model = QuestionFlag
        fields = ["flag", "flag_reason", "comment"]

@require_http_methods(["GET", "POST"])
@login_required
@user_passes_test(student_check, login_url="/access_denied_and_logout/")
def flag_question_form(
    request, question_pk, discipline_title=None, assignment_id=None
):
    """
    Get or Create QuestionFlag object for user, and allow edit
    """
    template = "peerinst/research/flag_question.html"
    question = get_object_or_404(Question, pk=question_pk)
    message = None

    try:
        question_flag = QuestionFlag.objects.get(
            user=request.user, question=question
        )
    except ObjectDoesNotExist:
        question_flag = None

    if request.method == "GET":
        form = QuestionFlagForm(instance=question_flag)

    if request.method == "POST":
        form = QuestionFlagForm(request.POST, instance=question_flag)

        if form.is_valid():
            form.instance.user = request.user
            form.instance.question = question
            instance = form.save()
            if instance.flag:
                message = _(
                    """
                    You have flagged this question, and your input
                    has been forwarded to a myDALITE content moderator.
                    """
                )
            else:
                message = _(
                    """
                    You have un-flagged this question, and thus it will be
                    taken off the list of potentially problematic questions.
                    """
                )

    context = {
        "form": form,
        "question": question,
        "message": message,
        "discipline_title": discipline_title,
        "assignment_id": assignment_id,
        "expert_answers": question.answer_set.filter(expert=True),
    }

    return render(request, template, context)

class AnswerExpertUpdateView(LoginRequiredMixin, NoStudentsMixin, UpdateView):
    model = Answer
    fields = ["expert"]
    template_name = "peerinst/research/answer-expert-update.html"

    def get_context_data(self, **kwargs):
        context = super().get_context_data(**kwargs)
        context["teacher"] = self.request.user.teacher
        context["question"] = Question.objects.get(pk=self.object.question_id)
        return context

    def get_success_url(self):
        return reverse(
            "research-fix-expert-rationale",
            kwargs={"question_id": self.object.question_id},
        )


#!/usr/bin/env python3

import argparse
import subprocess
import os

class DriverOpenStack:
    @staticmethod
    def cleanup(args):
        subprocess.check_call(
            [
                "ansible-playbook",
                "openstack/playbooks/clean.yaml",
                "-e",
                "prefix=%s" % args.prefix,
            ]
        )

    @staticmethod
    def purge(args):
        subprocess.check_call(
            [
                "ansible-playbook",
                "openstack/playbooks/purge.yaml",
                "-e",
                "prefix=%s" % args.prefix,
            ]
        )

    @staticmethod
    def deploy(args):
        subprocess.check_call(
            [
                "ansible-playbook",
                "-vv",
                "openstack/playbooks/deploy.yaml",
                "-i",
                "openstack/inventory/openstack.yaml",
                "-e",
                "prefix=%s" % args.prefix,
            ]
        )
        subprocess.check_call(
            [
                "ansible-playbook",
                "playbooks/prepare_datastore.yaml",
                "-i",
                "openstack/inventory/openstack.yaml",
                "-vv",
            ]
        )
        subprocess.check_call(
            [
                "ansible-playbook",
                "playbooks/write_inventory.yaml",
                "-i",
                "openstack/inventory/openstack.yaml",
                "-vv",
            ]
        )
        if args.run_test:
            subprocess.check_call(
                [
                    "ansible-playbook",
                    "-vv",
                    "playbooks/run_test.yaml",
                    "-i",
                    "openstack/inventory/openstack.yaml",
                    "-e",
                    "prefix=%s" % args.prefix,
                ]
            )

class DriverLibvirt:
    @staticmethod
    def cleanup(args):
        for i in ["esxi1", "esxi2", "vcenter", "datastore"]:
            subprocess.check_call(["vl", "stop", i])

    @staticmethod
    def purge(args):
        subprocess.check_call(["vl", "down"])

    @staticmethod
    def deploy(args):
        # prepare virt lightning configuration
        cmd = []
        cmd.append("ansible-playbook")
        cmd.append("libvirt/playbooks/prepare_virt_lightning.yaml")
        cmd.extend(
            [
                '-e "{}_distro={}"'.format(k, vars(args)[k])
                for k in ("datastore", "esxi", "vcenter")
                if vars(args)[k] is not None
            ]
        )
        cmd.append("-vv")
        subprocess.check_call(cmd)
        # create virt lightning vm
        subprocess.check_call(
            ["vl", "up", "--virt-lightning-yaml", "libvirt/virt-lightning.yaml"]
        )
        # download zuul useful roles
        cmd = []
        cmd.append("ansible-playbook")
        cmd.append("playbooks/download_zuul_ci_roles.yaml")
        cmd.append("-vv")
        subprocess.check_call(cmd)
        # create ansible inventory
        subprocess.check_call("vl ansible_inventory>inventory", shell=True)
        # unset any ANSIBLE_ROLES_PATH variable to read the ansible.cfg file
        if "ANSIBLE_ROLES_PATH" in os.environ:
            del os.environ["ANSIBLE_ROLES_PATH"]
        # prepare datastore
        subprocess.check_call(
            [
                "ansible-playbook",
                "playbooks/prepare_datastore.yaml",
                "-i",
                "inventory",
                "-vv",
            ]
        )
        subprocess.check_call(
            [
                "ansible-playbook",
                "playbooks/write_inventory.yaml",
                "-i",
                "inventory",
                "-vv",
            ]
        )
        if args.run_test:
            subprocess.check_call(
                [
                    "ansible-playbook",
                    "-vv",
                    "playbooks/run_test.yaml",
                    "-i",
                    "inventory",
                    "-e",
                    "prefix=%s" % args.prefix,
                ]
            )

parser = argparse.ArgumentParser(description="Process some integers.")

actions = {
    "deploy": "Deploy the environment.",
    "cleanup": "Remove the ESXi VMs and the VCenter, preserver the datastore.",
    "purge": "Remove all the resources.",
}

help_actions = ", ".join(["ðŸ”¥{k}: {v}".format(k=k, v=v) for k, v in actions.items()])
parser.add_argument(
    "action", type=str, choices=actions.keys(), metavar=None, help=help_actions
)
parser.add_argument(
    "--driver",
    dest="driver",
    type=str,
    choices=["libvirt"],  # the OpenStack support is currently broken
    default="libvirt",
    help="Driver to use",
)
parser.add_argument("--prefix", type=str, default="")
parser.add_argument("--run-test", type=bool, default=False)
parser.add_argument("--datastore-distro", type=str, dest="datastore")
parser.add_argument("--esxi-distro", type=str, dest="esxi")
parser.add_argument("--vcenter-distro", type=str, dest="vcenter")

args = parser.parse_args()

if args.driver == "openstack":
    driver = DriverOpenStack
elif args.driver == "libvirt":
    driver = DriverLibvirt

getattr(driver, args.action)(args)


import json

with open('szwb-data\json\szwb-dict-list.json', encoding='UTF-8') as f:

    dest_dict = {}
    source_dict = json.load(f)
    for date in source_dict:
        index_in_date = 0
        for record in source_dict[date]:
            index_all = record['date'].replace('-', '') + format(index_in_date, '03d')
            dest_dict[index_all] = record
            index_in_date = index_in_date + 1

with open('szwb-data\json\szwb-delist.json', 'w', encoding='UTF-8') as df:
    json.dump(dest_dict, df, ensure_ascii=False, indent=4, separators=(',', ':'))


from scipy.stats import pearsonr
import os
import sys
from time import sleep
import numpy as np
import torch
import torch.nn as nn
from torch.nn import PoissonNLLLoss, MSELoss
import torch.nn.functional as F
import os.path as path
from torchdeepretina.utils import get_cuda_info, save_checkpoint, GaussRegularizer
from torchdeepretina.datas import loadexpt, DataContainer, DataDistributor
from torchdeepretina.models import *
import torchdeepretina.analysis as analysis
import time
from tqdm import tqdm
import math
import torch.multiprocessing as mp
from queue import Queue
from collections import deque
import psutil
import gc
import resource
import json

class Trainer:
    def __init__(self, run_q=None, return_q=None, early_stopping=10, stop_tolerance=0.01):
        self.run_q = run_q
        self.ret_q = return_q
        self.early_stopping = early_stopping
        self.tolerance = stop_tolerance
        self.prev_acc = None

    def stop_early(self, acc):
        if self.early_stopping <= 0:
            return False # use 0 as way to cancel early stopping
        if self.prev_acc is None:
            self.prev_acc = acc
            self.stop_count = 0
            return False
        if acc-self.prev_acc < self.tolerance:
            self.stop_count += 1
            if self.stop_count >= self.early_stopping:
                return True
        self.stop_count = 0
        self.prev_acc = acc
        return False

    def loop_training(self):
        train_args = self.run_q.get()
        while True:
            try:
                results = self.train(*train_args)
                self.ret_q.put([results])
                train_args = self.run_q.get()
            except Exception as e:
                print("Caught error",e,"on", train_args[0]['exp_num'], "will retry in 100 seconds...")
                sleep(100)

    def get_model_and_distr(self, hyps, model_hyps, train_data):
        """
        hyps: dict
            dict of relevant hyperparameters
        model_hyps: dict
            dict of relevant hyperparameters
        train_data: DataContainer
            a DataContainer of the training data as returned by get_data
        """
        model = hyps['model_class'](**model_hyps)
        model = model.to(hyps['device'])
        num_val = 10000
        seq_len = 1 if not model.recurrent else hyps['recur_seq_len']
        shift_labels = False if 'shift_labels' not in hyps else hyps['shift_labels']
        data_distr = DataDistributor(train_data, num_val, batch_size=hyps['batch_size'], shuffle=hyps['shuffle'],
                                                                    recurrent=model.recurrent, seq_len=seq_len,
                                                                    shift_labels=shift_labels)
        data_distr.torch()
        return model, data_distr

    def record_session(self, hyps, model):
        """
        hyps: dict
            dict of relevant hyperparameters
        model: torch nn.Module
            the model to be trained
        """
        if not os.path.exists(hyps['save_folder']):
            os.mkdir(hyps['save_folder'])
        with open(os.path.join(hyps['save_folder'],"hyperparams.txt"),'w') as f:
            f.write(str(model)+'\n')
            for k in sorted(hyps.keys()):
                f.write(str(k) + ": " + str(hyps[k]) + "\n")
        with open(os.path.join(hyps['save_folder'],"hyperparams.json"),'w') as f:
            temp_hyps = {k:v for k,v in hyps.items()}
            del temp_hyps['model_class']
            json.dump(temp_hyps, f)

    def get_hs(self, hyps, model, batch_size=None):
        """
        hyps: dict
            dict of relevant hyperparameters
        model: torch nn.Module
            the model to be trained
        batch_size: int
        """
        hs = None
        device = hyps['device']
        if model.recurrent:
            if batch_size is None:
                batch_size = hyps['batch_size']
            hs = [torch.zeros(batch_size, *h).to(device) for h in model.h_shapes]
            if model.kinetic:
                hs[0][:,0] = hyps['intl_pop']
                hs[1] = deque([],maxlen=hyps['recur_seq_len'])
                for i in range(hyps['recur_seq_len']):
                    hs[1].append(torch.zeros(batch_size, *model.h_shapes[1]).to(device))
        return hs

    def print_train_update(self, error, grade, l1, model, n_loops, i):
        loss = error + grade + l1
        s = "Loss: {:.5e} – ".format(loss.item())
        if grade.item() > 0:
            s += "Error: {:.5e} – Grade: {:.5e} – ".format(error.item(), grade.item())
        if model.kinetic:
            ps = model.kinetics.named_parameters()
            s += " – ".join([str(name)+":"+str(round(p.data.item(),4)) for name,p in list(ps)]) + " – "
        print(s, i, "/", n_loops, end="       \r")

    def validate_recurrent(self, hyps, model, data_distr, hs, loss_fn, verbose=False):
        """
        hyps: dict
            dict of relevant hyperparameters
        model: torch nn.Module
            the model to be trained
        data_distr: DataDistributor
            the data distribution object as obtained through self.get_model_and_distr
        hs: list
            the hidden states of the recurrent model as obtained through self.get_hs
        loss_fn: loss function
            the loss function
        """
        val_preds = []
        val_targs = []
        val_loss = 0
        batch_size = hyps['batch_size']
        n_loops = data_distr.val_shape[0]*batch_size
        step_size = 1
        device = hyps['device']
        for b in range(batch_size):
            for i, (val_x, val_y) in enumerate(data_distr.val_sample(step_size)):
                # val_x (batch_size, seq_len, depth, height, width)
                # val_y (batch_size, seq_len, n_ganlion)
                outs, hs = model(val_x[b:b+1,0].to(device), hs)
                val_y = val_y[:,0]
                val_preds.append(outs[:,None,:].cpu().detach().numpy())
                val_targs.append(val_y[b:b+1,None,:].cpu().detach().numpy())
                val_loss += loss_fn(outs, val_y.to(device)).item()
                if hyps['l1'] > 0:
                    val_loss += (hyps['l1'] * torch.norm(outs, 1).float()/outs.shape[0]).item()
                if verbose and i%(n_loops//10) == 0:
                    print("{}/{}".format(b*data_distr.val_shape[0] + i,data_distr.val_shape[0]*batch_size), end="     \r")
        return val_loss, val_preds, val_targs

    def validate_static(self, hyps, model, data_distr, loss_fn, step_size=500, verbose=False):
        """
        hyps: dict
            dict of relevant hyperparameters
        model: torch nn.Module
            the model to be trained
        data_distr: DataDistributor
            the data distribution object as obtained through self.get_model_and_distr
        loss_fn: loss function
            the loss function
        step_size: int
            optional size of batches when evaluating validation set
        """
        val_preds = []
        val_targs = []
        val_loss = 0
        n_loops = data_distr.val_shape[0]//step_size
        device = hyps['device']
        for i, (val_x, val_y) in enumerate(data_distr.val_sample(step_size)):
            outs = model(val_x.to(device)).detach()
            val_preds.append(outs.cpu().detach().numpy())
            val_targs.append(val_y.cpu().detach().numpy())
            val_loss += loss_fn(outs, val_y.to(device)).item()
            if hyps['l1'] > 0:
                val_loss += (hyps['l1'] * torch.norm(outs, 1).float()/outs.shape[0]).item()
            if verbose and i%(n_loops//10) == 0:
                print("{}/{}".format(i*step_size,data_distr.val_y.shape[0]), end="     \r")
        return val_loss, val_preds, val_targs

    def train(self, hyps, model_hyps, device, verbose=False):
        """
        hyps: dict
            dict of relevant hyperparameters
        model_hyps: dict
            dict of relevant hyperparameters
        model: torch nn.Module
            the model to be trained
        train_data: DataContainer
            a DataContainer of the training data as returned by get_data
        """
        # Initialize miscellaneous parameters
        torch.cuda.empty_cache()
        hyps['device'] = device
        batch_size = hyps['batch_size']
        if 'skip_nums' in hyps and hyps['skip_nums'] is not None and len(hyps['skip_nums']) > 0 and hyps['exp_num'] in hyps['skip_nums']:
            print("Skipping", hyps['save_folder'])
            results = {"save_folder":hyps['save_folder'], "Loss":None, "ValAcc":None, "ValLoss":None, "TestPearson":None}
            return results

        # Get Data, Make Model, Record Initial Hyps and Model
        train_data, test_data = get_data(hyps)
        model_hyps["n_units"] = train_data.y.shape[-1]
        model_hyps['centers'] = train_data.centers
        model, data_distr = self.get_model_and_distr(hyps, model_hyps, train_data)
        print("train shape:", data_distr.train_shape)
        print("val shape:", data_distr.val_shape)
        self.record_session(hyps, model)
        teacher = None
        if 'teacher' in hyps and hyps['teacher'] is not None:
            teacher = analysis.read_model_file(hyps['teacher'])
            teacher.to(device)
            teacher.eval()
            #if hyps['teacher_layers'] is not None:

        # Make optimization objects (lossfxn, optimizer, scheduler)
        optimizer, scheduler, loss_fn = get_optim_objs(hyps, model, train_data.centers)
        if 'gauss_reg' in hyps and hyps['gauss_reg'] > 0:
            gauss_reg = GaussRegularizer(model, [0,6], std=hyps['gauss_reg'])

        # Training
        for epoch in range(hyps['n_epochs']):
            print("Beginning Epoch", epoch, " -- ", hyps['save_folder'])
            print()
            n_loops = data_distr.n_loops
            hs = self.get_hs(hyps, model)
            model.train(mode=True)
            epoch_loss = 0
            stats_string = 'Epoch ' + str(epoch) + " -- " + hyps['save_folder'] + "\n"
            starttime = time.time()

            # Train Loop
            for i,(x,label) in enumerate(data_distr.train_sample()):
                optimizer.zero_grad()
                label = label.float().to(device)

                # Error Evaluation
                if model.recurrent:
                    y,error,grade,hs = recurrent_eval(hyps, x, label, model, hs, loss_fn, teacher=teacher)
                else:
                    y,error,grade = static_eval(hyps, x, label, model, loss_fn, teacher=teacher)
                activity_l1 = torch.zeros(1).to(device) if hyps['l1']<=0 else hyps['l1'] * torch.norm(y, 1).float().mean()
                if 'gauss_reg' in hyps and hyps['gauss_reg'] > 0:
                    activity_l1 += hyps['gauss_loss_coef']*gauss_reg.get_loss()

                # Backwards Pass
                loss = error + activity_l1 + grade
                loss.backward(retain_graph=True)
                optimizer.step()

                epoch_loss += loss.item()
                if verbose:
                    self.print_train_update(error, grade, activity_l1, model, n_loops, i)
                if math.isnan(epoch_loss) or math.isinf(epoch_loss) or hyps['exp_name']=="test":
                    break

            # Clean Up Train Loop
            avg_loss = epoch_loss/n_loops
            stats_string += 'Avg Loss: ' + str(avg_loss) + " -- exec time:"+ str(time.time() - starttime) +"\n"
            del x
            del y
            del label

            # Validation
            model.eval()
            with torch.no_grad():

                # Miscellaneous Initializations
                if model.recurrent:
                    step_size = 1
                    n_loops = data_distr.val_shape[0]*batch_size
                    hs = self.get_hs(hyps, model, batch_size=step_size)
                else:
                    step_size = 500
                    n_loops = data_distr.val_shape[0]//step_size
                if verbose:
                    print()
                    print("Validating")

                # Validation Block
                if model.recurrent:
                    val_loss, val_preds, val_targs = self.validate_recurrent(hyps, model, data_distr, hs,
                                                                                loss_fn, verbose=verbose)
                else:
                    val_loss, val_preds, val_targs = self.validate_static(hyps, model, data_distr, loss_fn,
                                                                       step_size=step_size, verbose=verbose)

                # Validation Evaluation
                val_loss = val_loss/n_loops
                n_units = data_distr.val_y.shape[-1]
                if model.recurrent:
                    val_preds = np.concatenate(val_preds, axis=1).reshape((-1,n_units))
                    val_targs = np.concatenate(val_targs, axis=1).reshape((-1,n_units))
                else:
                    val_preds = np.concatenate(val_preds, axis=0)
                    val_targs = np.concatenate(val_targs, axis=0)
                pearsons = []
                for cell in range(val_preds.shape[-1]):
                    pearsons.append(pearsonr(val_preds[:, cell], val_targs[:,cell])[0])
                stats_string += "Val Cell Pearsons:" + " - ".join([str(p) for p in pearsons])+'\n'
                val_acc = np.mean(pearsons)
                stop = self.stop_early(val_acc)

                # Exp Validation
                exp_val_acc = None
                if hyps["log_poisson"] and hyps['softplus'] and not model.infr_exp:
                    val_preds = np.exp(val_preds)
                    for cell in range(val_preds.shape[-1]):
                        pearsons.append(pearsonr(val_preds[:, cell], val_targs[:,cell])[0])
                    exp_val_acc = np.mean(pearsons)

                # Clean Up
                stats_string += "Avg Val Pearson: {} -- Val Loss: {} -- Exp Val: {}\n".format(val_acc, val_loss, exp_val_acc)
                scheduler.step(val_loss)
                del val_preds

                # Validation on Test Subset (Static Only)
                avg_pearson = 0
                if not model.recurrent and test_data is not None:
                    test_x = torch.from_numpy(test_data.X)
                    test_obs = model(test_x.to(device)).cpu().detach().numpy()
                    rng = range(test_obs.shape[-1])
                    if verbose:
                        rng = tqdm(rng)
                    for cell in rng:
                        obs = test_obs[:,cell]
                        lab = test_data.y[:,cell]
                        r,p = pearsonr(obs,lab)
                        avg_pearson += r
                        stats_string += 'Cell ' + str(cell) + ': ' + str(r)+"\n"
                    avg_pearson = avg_pearson / float(test_obs.shape[-1])
                    stats_string += "Avg Test Pearson: "+ str(avg_pearson) + "\n"
                    del test_obs

            # Save Model Snapshot
            optimizer.zero_grad()
            save_dict = {
                "model_type": hyps['model_type'],
                "model_hyps": model_hyps,
                "model_state_dict":model.state_dict(),
                "optim_state_dict":optimizer.state_dict(),
                "loss": avg_loss,
                "epoch":epoch,
                "val_loss":val_loss,
                "val_acc":val_acc,
                "exp_val_acc":exp_val_acc,
                "test_pearson":avg_pearson,
                "norm_stats":train_data.stats,
            }
            for k in hyps.keys():
                if k not in save_dict:
                    save_dict[k] = hyps[k]
            save_checkpoint(save_dict, hyps['save_folder'], 'test', del_prev=True)

            # Print Epoch Stats
            gc.collect()
            max_mem_used = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
            stats_string += "Memory Used: {:.2f} mb".format(max_mem_used / 1024)+"\n"
            print(stats_string)
            # If loss is nan, training is futile
            if math.isnan(avg_loss) or math.isinf(avg_loss) or stop:
                break

        # Final save
        results = {"save_folder":hyps['save_folder'], "Loss":avg_loss, "ValAcc":val_acc, "ValLoss":val_loss, "TestPearson":avg_pearson}
        with open(hyps['save_folder'] + "/hyperparams.txt",'a') as f:
            f.write("\n" + " ".join([str(k)+":"+str(results[k]) for k in sorted(results.keys())]) + '\n')
        return results

def recurrent_eval(hyps, x, label, model, hs, loss_fn, teacher=None):
    """
    hyps: dict
        dict of relevant hyperparameters
    x: torch FloatTensor
        a batch of the training data
    label: torch FloatTensor
        a batch of the training labels
    model: torch nn.Module
        the model to be trained
    hs: list
        the hidden states of the recurrent model as obtained through self.get_hs
    """
    hs_out = hs
    batch_size = hyps['batch_size']
    ys = []
    answers = []
    device = hyps['device']
    teacher_device = device
    if teacher is None:
        teacher = lambda x: None
        teacher_device = torch.device("cpu")
    for ri in range(hyps['recur_seq_len']):
        ins = x[:,ri]
        y, hs_out = model(ins.to(device), hs_out)
        ys.append(y.view(batch_size, 1, label.shape[-1]))
        with torch.no_grad():
            ans = teacher(x.to(teacher_device))
        answers.append(ans)
        if ri == 0 and not hyps['reset_hs']:
            if model.kinetic:
                hs[0] = hs_out[0].data.clone()
                hs[1] = deque([h.data.clone() for h in hs[1]], maxlen=hyps['recur_seq_len'])
            else:
                hs = [h.data.clone() for h in hs_out]
    y = torch.cat(ys, dim=1)
    error = loss_fn(y,label)/hyps['recur_seq_len']

    # Teacher
    if answers[0] is not None:
        answers = torch.cat(answers, dim=1)
        grade = hyps['teacher_coef']*F.mse_loss(y,answers.data)/hyps['recur_seq_len']
        hyps['teacher_coef'] *= hyps['teacher_decay']
    else:
        grade = torch.zeros(1).to(device)

    # Kinetics
    if hyps['model_type'] == "KineticsModel" and model.kinetics.kfr + model.kinetics.ksi > 1:
        error += ((1-(model.kinetics.kfr+model.kinetics.ksi))**2).mean()

    return y, error, grade, hs

def static_eval(hyps, x, label, model, loss_fn, teacher=None):
    """
    hyps: dict
        dict of relevant hyperparameters
    x: torch FloatTensor
        a batch of the training data
    label: torch FloatTensor
        a batch of the training labels
    model: torch nn.Module
        the model to be trained
    teacher: torch nn.Module
        optional teacher network
    """
    device = hyps['device']
    y = model(x.to(device))
    error = loss_fn(y,label)
    if teacher is not None:
        with torch.no_grad():
            ans = teacher(x.to(device))
        grade = F.mse_loss(y,ans.data)
    else:
        grade = torch.zeros(1).to(device)
    return y,error,grade

def get_data(hyps):
    """
    hyps: dict
        dict of relevant hyperparameters
    """
    img_depth, img_height, img_width = hyps['img_shape']
    train_data = DataContainer(loadexpt(hyps['dataset'],hyps['cells'], hyps['stim_type'],'train',img_depth,0))
    norm_stats = [train_data.stats['mean'], train_data.stats['std']]

    try:
        test_data = DataContainer(loadexpt(hyps['dataset'],hyps['cells'],hyps['stim_type'],'test',img_depth,0,
                                                                                norm_stats=norm_stats))
        test_data.X = test_data.X[:500]
        test_data.y = test_data.y[:500]
    except:
        test_data = None
    return train_data, test_data

def get_optim_objs(hyps, model, centers=None):
    """
    hyps: dict
        dict of relevant hyperparameters
    model: torch nn.Module
        the model to be trained
    centers: list of tuples or lists, shape: (n_cells, 2)
        the centers of each ganglion cell in terms of image coordinates
        if None centers is ignored
    """
    if 'lossfxn' not in hyps:
        hyps['lossfxn'] = "PoissonNLLLoss"
    if hyps['lossfxn'] == "PoissonNLLLoss" and 'log_poisson' in hyps:
        loss_fn = globals()[hyps['lossfxn']](log_input=hyps['log_poisson'])
    else:
        loss_fn = globals()[hyps['lossfxn']]()
    optimizer = torch.optim.Adam(model.parameters(), lr = hyps['lr'], weight_decay = hyps['l2'])
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor = 0.1)
    return optimizer, scheduler, loss_fn

def fill_hyper_q(hyps, hyp_ranges, keys, hyper_q, idx=0):
    """
    Recursive function to load each of the hyperparameter combinations
    onto a queue.

    hyps - dict of hyperparameters created by a HyperParameters object
        type: dict
        keys: name of hyperparameter
        values: value of hyperparameter
    hyp_ranges - dict of ranges for hyperparameters to take over the search
        type: dict
        keys: name of hyperparameters to be searched over
        values: list of values to search over for that hyperparameter
    keys - keys of the hyperparameters to be searched over. Used to
            specify order of keys to search
    train - method that handles training of model. Should return a dict of results.
    hyper_q - Queue to hold all parameter sets
    idx - the index of the current key to be searched over
    """
    # Base call, runs the training and saves the result
    if idx >= len(keys):
        if 'exp_num' not in hyps:
            if 'starting_exp_num' not in hyps: hyps['starting_exp_num'] = 0
            hyps['exp_num'] = hyps['starting_exp_num']
        hyps['save_folder'] = hyps['exp_name'] + "/" + hyps['exp_name'] +"_"+ str(hyps['exp_num'])
        for k in keys:
            hyps['save_folder'] += "_" + str(k)+str(hyps[k])

        test_hyps = get_model_hyps(hyps)

        # Make model hyps
        hyps['model_class'] = globals()[hyps['model_type']]
        model_hyps = {k:v for k,v in hyps.items()}
        fn_args = set(hyps['model_class'].__init__.__code__.co_varnames)
        if "kwargs" in fn_args:
            fn_args = fn_args | set(TDRModel.__init__.__code__.co_varnames)
        keys = list(model_hyps.keys())
        for k in keys:
            if k not in fn_args:
                del model_hyps[k]

        for k in test_hyps.keys():
            assert test_hyps[k] == model_hyps[k]

        # Load q
        hyper_q.put([{k:v for k,v in hyps.items()}, {k:v for k,v in model_hyps.items()}])
        hyps['exp_num'] += 1

    # Non-base call. Sets a hyperparameter to a new search value and passes down the dict.
    else:
        key = keys[idx]
        for param in hyp_ranges[key]:
            hyps[key] = param
            hyper_q = fill_hyper_q(hyps, hyp_ranges, keys, hyper_q, idx+1)
    return hyper_q

def get_device(visible_devices, cuda_buffer=3000):
    info = get_cuda_info()
    for i,mem_dict in enumerate(info):
        if i in visible_devices and mem_dict['remaining_mem'] >= cuda_buffer:
            return i
    return -1

def mp_hyper_search(hyps, hyp_ranges, keys, n_workers=4, visible_devices={0,1,2,3,4,5}, cuda_buffer=3000,
                                                    ram_buffer=6000, early_stopping=10, stop_tolerance=.01):
    starttime = time.time()
    # Make results file
    if not os.path.exists(hyps['exp_name']):
        os.mkdir(hyps['exp_name'])
    results_file = hyps['exp_name']+"/results.txt"
    with open(results_file,'a') as f:
        f.write("Hyperparameters:\n")
        for k in hyps.keys():
            if k not in hyp_ranges:
                f.write(str(k) + ": " + str(hyps[k]) + '\n')
        f.write("\nHyperranges:\n")
        for k in hyp_ranges.keys():
            f.write(str(k) + ": [" + ",".join([str(v) for v in hyp_ranges[k]])+']\n')
        f.write('\n')

    hyper_q = mp.Queue()

    hyper_q = fill_hyper_q(hyps, hyp_ranges, keys, hyper_q, idx=0)
    total_searches = hyper_q.qsize()
    print("n_searches:", total_searches)

    n_workers = min(total_searches, n_workers) # No need to waste resources
    run_q = mp.Queue(n_workers)
    return_q = mp.Queue(n_workers+1)
    workers = []
    procs = []
    print("Initializing workers")
    for i in range(n_workers):
        worker = Trainer(run_q, return_q, early_stopping=early_stopping, stop_tolerance=stop_tolerance)
        workers.append(worker)
        proc = mp.Process(target=worker.loop_training)
        proc.start()
        procs.append(proc)

    result_count = 0
    print("Starting Hyperloop")
    while result_count < total_searches:
        print("Running Time:", time.time()-starttime)
        device = get_device(visible_devices, cuda_buffer)
        enough_ram = psutil.virtual_memory().free//1028**2 > ram_buffer
        # must be careful not to threadlock here
        if (not enough_ram or device <= -1) and hyper_q.qsize() >= total_searches:
            print("RAM shortage or no devices available, sleeping for 20s")
            time.sleep(20)
        elif not hyper_q.empty() and not run_q.full():
            hyperset = hyper_q.get()
            hyperset.append(device)
            print("Loading hyperset...")
            run_q.put(hyperset)
            time.sleep(5) # Timer to ensure ram measurements are completed appropriately
            print("Loaded", hyperset[0]["exp_num"])
        else:
            print("Waiting...")
            results = return_q.get()[0]
            print("Collected", results['save_folder'])
            with open(results_file,'a') as f:
                results = " -- ".join([str(k)+":"+str(results[k]) for k in sorted(results.keys())])
                f.write("\n"+results+"\n")
            result_count += 1

    for proc in procs:
        proc.terminate()
        proc.join(timeout=1.0)

def hyper_search(hyps, hyp_ranges, keys, device, early_stopping=10, stop_tolerance=.01):
    starttime = time.time()
    # Make results file
    if not os.path.exists(hyps['exp_name']):
        os.mkdir(hyps['exp_name'])
    results_file = hyps['exp_name']+"/results.txt"
    with open(results_file,'a') as f:
        f.write("Hyperparameters:\n")
        for k in hyps.keys():
            if k not in hyp_ranges:
                f.write(str(k) + ": " + str(hyps[k]) + '\n')
        f.write("\nHyperranges:\n")
        for k in hyp_ranges.keys():
            f.write(str(k) + ": [" + ",".join([str(v) for v in hyp_ranges[k]])+']\n')
        f.write('\n')

    hyper_q = Queue()

    hyper_q = fill_hyper_q(hyps, hyp_ranges, keys, hyper_q, idx=0)
    total_searches = hyper_q.qsize()
    print("n_searches:", total_searches)

    trainer = Trainer(early_stopping=early_stopping, stop_tolerance=stop_tolerance)

    result_count = 0
    print("Starting Hyperloop")
    while not hyper_q.empty():
        print("Searches left:", hyper_q.qsize(),"-- Running Time:", time.time()-starttime)
        hyperset = hyper_q.get()
        hyperset.append(device)
        results = trainer.train(*hyperset, verbose=True)
        with open(results_file,'a') as f:
            results = " -- ".join([str(k)+":"+str(results[k]) for k in sorted(results.keys())])
            f.write("\n"+results+"\n")

def get_model_hyps(hyps):
    hyps = {k:v for k,v in hyps.items()}
    hyps['model_class'] = globals()[hyps['model_type']]
    model_hyps = {k:v for k,v in hyps.items()}

    fn_args = set(hyps['model_class'].__init__.__code__.co_varnames)
    if "kwargs" in fn_args:
        fn_args = fn_args | set(TDRModel.__init__.__code__.co_varnames)
    keys = list(model_hyps.keys())
    for k in keys:
        if k not in fn_args:
            del model_hyps[k]
    return model_hyps

class ModulePackage:
    def __init__(self, model, layer_keys):
        self.model = model
        self.layer_keys = set(layer_keys)
        self.layers = dict()
        self.handles = []
        for name, mod in model.named_modules():
            if name in self.layer_keys:
                self.handles.append(mod.register_forward_hook(self.get_hook(name)))

    def get_hook(self, layer_key):
        def hook(module, inp, out):
            self.layers[layer_key] = out
        return hook

    def remove_hooks(self):
        for h in self.handles:
            h.remove()
        self.handles = []

    def detach_all(self):
        self.remove_hooks()
        for k in self.layers.keys():
            try:
                self.layers[k].detach().cpu()
                del self.layers[k]
            except:
                pass

    def __call__(self, x):
        return self.model(x)

class LossFxnWrapper:
    def __init__(self, loss_fn, hyps, centers):
        self.centers = centers
        self.loss_fn = loss_fn
        self.coords = self.convert_coords(hyps, centers)

    def convert_coords(self, hyps, centers):
        """
        Assumes a stride of 1 with 0 padding in each layer.
        """
        # Each quantity is even, thus the final half_effective_ksize is odd
        half_effective_ksize = (hyps['ksizes'][0]-1) + (hyps['ksizes'][1]-1) + (hyps['ksizes'][2]//2-1) + 1
        coords = []
        for center in centers:
            row = min(max(0,center[0]-half_effective_ksize), hyps['img_shape'][1]-2*(half_effective_ksize-1))
            col = min(max(0,center[1]-half_effective_ksize), hyps['img_shape'][2]-2*(half_effective_ksize-1))
            coords.append([row,col])
        return torch.LongTensor(coords)

    def __call__(self, x, y):
        idxs = torch.arange(x.shape[-4]).long()
        chans = torch.arange(x.shape[-3]).long()
        units = x[...,:,chans,self.coords[:,0],self.coords[:,1]]
        return self.loss_fn(units, y)