import cv2 as cv

if __name__=="__main__":
    pre = input("Prefix: ")
    print(f"Saving as {pre}_count.png")
    cam = cv.VideoCapture(-1)

    count = 0
    while(True):

        flag, frame = cam.read()

        if(not flag):
            print("Frame grab error")
            continue

        cv.imshow(f"{pre} Frame", frame)

        k = cv.waitKey(1)
        new_symbol = "0703_symbols/"
        if(k == ord('s')):
            cv.imwrite(f"{new_symbol}{pre}_{count}.png", frame)
            count+=1
        if(k == ord('q')):
            break

    cv.destroyAllWindows()


from beartype.typing import (
    Dict,
    Set,
)
from beartype._cave._cavemap import NoneTypeOr
from beartype._check.convert.convcoerce import clear_coerce_hint_caches
from beartype._conf.confcls import BeartypeConf
from beartype._data.cls.datacls import TYPES_BEARTYPEABLE
from beartype._data.hint.datahinttyping import (
    BeartypeableT,
    TypeStack,
)
from beartype._util.cls.utilclsset import set_type_attr
from beartype._util.module.utilmodget import get_object_module_name_or_none
from collections import defaultdict
from functools import wraps

# ....................{ DECORATORS ~ type                  }....................
def beartype_type(
    # Mandatory parameters.
    cls: BeartypeableT,
    conf: BeartypeConf,

    # Optional parameters.
    cls_stack: TypeStack = None,
) -> BeartypeableT:
    '''
    Decorate the passed class with dynamically generated type-checking.

    Parameters
    ----------
    cls : BeartypeableT
        Class to be decorated by :func:`beartype.beartype`.
    conf : BeartypeConf
        Beartype configuration configuring :func:`beartype.beartype` uniquely
        specific to this class.
    cls_stack : TypeStack, optional
        **Type stack** (i.e., either a tuple of the one or more
        :func:`beartype.beartype`-decorated classes lexically containing the
        class variable or method annotated by this hint *or* :data:`None`).
        Defaults to :data:`None`.

    Returns
    ----------
    BeartypeableT
        This class decorated by :func:`beartype.beartype`.
    '''
    assert isinstance(cls, type), f'{repr(cls)} not type.'
    assert isinstance(cls_stack, NoneTypeOr[tuple]), (
        f'{repr(cls_stack)} neither tuple nor "None".')
    # assert isinstance(conf, BeartypeConf), f'{repr(conf)} not configuration.'
    # print(f'Decorating type {repr(obj)}...')

    # ....................{ IMPORTS                        }....................
    # Avoid circular import dependencies.
    from beartype._decor.decorcore import beartype_object

    # ....................{ NOOP                           }....................
    # Original C-based __sizeof__() dunder method defined by this class, which
    # this decorator subsequently wraps with a pure-Python __sizeof__() dunder
    # method. Why? Tangential reasons that are obscure, profane, and have
    # absolutely *NOTHING* to do with the __sizeof__() dunder method itself.
    # Succinctly, @beartype needs a reasonably safe place to persist
    # @beartype-specific attributes pertaining to this class.
    #
    # Clearly, the obvious place would be this class itself. However, doing so
    # would fundamentally modify this class and thus *ALL* instances of this
    # class in an unexpected and thus possibly unsafe manner. Consider common
    # use cases like slots, introspection, pickling, and sizing. Clearly,
    # monkey-patching attributes into class dictionaries without the explicit
    # consent of class designers is an ill-advised approach.
    #
    # A less obvious but safer place is required. A method of this class would
    # be the ideal candidate; whereas everybody cares about object attributes
    # and thus class dictionaries, nobody cares about method attributes. This is
    # why @beartype safely monkey-patches attributes into @beartype-decorated
    # methods. However, which method? Most methods are *NOT* guaranteed to exist
    # across all possible classes. Adding a new method to this class would be no
    # better than adding a new attribute to this class; both modify class
    # dictionaries. Fortunately, Python currently guarantees *ALL* classes to
    # define at least 24 dunder methods as of Python 3.11. How? Via the root
    # "object" superclass. Unfortunately, *ALL* of these methods are C-based and
    # thus do *NOT* directly support monkey-patching: e.g.,
    #     >>> class AhMahGoddess(object): pass
    #     >>> AhMahGoddess.__init__.__beartyped_cls = AhMahGoddess
    #     AttributeError: 'wrapper_descriptor' object has no attribute
    #     '__beartyped_cls'
    #
    # Fortunately, *ALL* of these methods may be wrapped by pure-Python
    # equivalents whose implementations defer to their original C-based methods.
    # Unfortunately, doing so slightly reduces the efficiency of calling these
    # methods. Fortunately, a subset of these methods are rarely called under
    # production workloads; slightly reducing the efficiency of calling these
    # methods is irrelevant to almost all use cases. Of these, the most obscure,
    # largely useless, poorly documented, and single-use is the __sizeof__()
    # dunder method -- which is only ever called by the sys.getsizeof() utility
    # function, which itself is only ever called manually in a REPL or by
    # third-party object sizing packages. In short, __sizeof__() is perfect.
    cls_sizeof_old = cls.__sizeof__

    # True only if this decorator has already decorated this class, as indicated
    # by the @beartype-specific class variable "__beartyped_cls" monkey-patched
    # into a pure-Python __sizeof__() dunder method wrapper by a prior call to
    # this decorator passed this class.
    is_cls_beartyped = getattr(cls_sizeof_old, '__beartyped_cls', None)

    # If the value of this variable is that of this class, a prior call to this
    # decorator has already decorated this class. In this case, silently reduce
    # to a noop by returning this class as is.
    #
    # See where this variable is set below for further details.
    if is_cls_beartyped is cls:
        # print(f'Ignoring repeat decoration of {repr(cls)}...')
        return cls  # type: ignore[return-value]
    # Else, this decorator has yet to decorate this class.

    # ....................{ LOCALS                         }....................
    # Fully-qualified name of the module defining this class if this class is
    # defined by a module *OR* "None" otherwise (e.g., if this class is only
    # dynamically defined in-memory outside of any module structure).
    module_name = get_object_module_name_or_none(cls)

    # If this class is defined by a module...
    if module_name:
        # Unqualified basename of this class.
        type_name = cls.__name__

        # Set of the unqualified basenames of *ALL* classes in that module
        # previously decorated by this decorator.
        type_names_beartyped = _BEARTYPED_MODULE_TO_TYPE_NAME[module_name]

        # If a class with the same unqualified basename defined in a module with
        # the same fully-qualified name has already been marked as decorated by
        # this decorator, then either:
        # * That module has been externally reloaded. In this case, this class
        #   (along with the remainder of that module) has now been redefined.
        #   Common examples include:
        #   * Rerunning a Jupyter cell defining this class.
        #   * Refreshing a web app enabling hot reloading (i.e., automatic
        #     reloading of on-disk modules whose contents have been externally
        #     modified *AFTER* that app was initially run). Since most Python
        #     web app frameworks (e.g., Flask, Streamlit) support hot reloading,
        #     this is the common case.
        # * That module has internally redefined this class two or more times.
        #   This behaviour, while typically a bug, is also technically valid:
        #   e.g.,
        #       @beartype
        #       def MuhClass(object): ...
        #       @beartype
        #       def MuhClass(object): ...   # <-- this makes me squint
        #
        # In this case, clear *ALL* beartype-specific internal caches that have
        # been shown to fail when a class is redefined.
        if type_name in type_names_beartyped:
            #FIXME: Consider emitting a logging message instead if this branch
            #ever becomes computationally intensive, please.
            # print(f'@beartyped class "{module_name}.{type_name}" redefined!')

            # Clear the previously accessed set of the unqualified basenames of
            # *ALL* classes in that module previously decorated by this
            # decorator. Technically, this is optional. Pragmatically, this
            # *SHOULD* significantly improve the space and time constraints
            # associated with this class redefinition. Why? Because this class
            # being redefined implies that the module defining this class is
            # being redefined, which implies that all classes in that module are
            # being redefined as well. If we did *NOT* clear this set here, then
            # this set would continue to contain the unqualified basenames of
            # those other classes in that module; each @beartype-decorated
            # redefinition of those other classes would then unnecessarily clear
            # the same caches already cleared by the first @beartype-decorated
            # redefinition of a class in that module. Since doing so would be
            # overly aggressive and thus inefficient, avoiding doing so improves
            # efficiency in the common case of module redefinition.
            _BEARTYPED_MODULE_TO_TYPE_NAME.clear()

            # Clear *ALL* type hint coercion caches, which map from the
            # machine-readable representations of previously seen
            # non-self-cached type hints (e.g., "list[MuhClass]") to the first
            # seen instance of those hints (e.g., list[MuhClass]). Since this
            # class has been redefined, the first seen instance of those hints
            # could contain a reference to the first definition of this class;
            # if so, there now exists a discrepancy between the current
            # definition of this class and cached hints containing the prior
            # definition of this class. For safety, all caches possibly
            # containing those hints must now be assumed to be invalid. Failing
            # to clear these caches causes @beartype-decorated wrapper functions
            # to raise erroneous type-checking violations. See also:
            #     https://github.com/beartype/beartype/issues/288
            clear_coerce_hint_caches()
        # Else, this is the first decoration of this class by this decorator.
        # In this case...
        else:
            # Record that this class has now been decorated by this decorator.
            # Technically, this should (probably) be performed *AFTER* this
            # decorator has actually successfully decorated this class.
            # Pragmatically, doing so here is simply faster and... simpler.
            type_names_beartyped.add(type_name)
    # Else, this class is *NOT* defined by a module.

    # Replace the passed class stack with a new class stack appending this
    # decorated class to the top of this stack, reflecting the fact that this
    # decorated class is now the most deeply lexically nested class for the
    # currently recursive chain of @beartype-decorated classes.
    cls_stack = (
        # If the caller passed *NO* class stack, then this class is necessarily
        # the first decorated class being decorated directly by @beartype and
        # thus the root decorated class.
        #
        # Note this is the common case and thus tested first. Since nested
        # classes effectively do *NOT* exist in the wild, this comprises
        # 99.999% of all real-world cases.
        (cls,)
        if cls_stack is None else
        # Else, the caller passed a clack stack comprising at least a root
        # decorated class. Preserve that class as is to properly expose that
        # class elsewhere.
        cls_stack + (cls,)
    )

    # ....................{ DECORATION                     }....................
    # For the unqualified name and value of each direct (i.e., *NOT* indirectly
    # inherited) attribute of this class...
    for attr_name, attr_value in cls.__dict__.items():  # pyright: ignore[reportGeneralTypeIssues]
        # If this attribute is beartypeable...
        if isinstance(attr_value, TYPES_BEARTYPEABLE):
            # This attribute decorated with type-checking configured by this
            # configuration if *NOT* already decorated.
            attr_value_beartyped = beartype_object(
                obj=attr_value,
                conf=conf,
                cls_stack=cls_stack,
            )

            # Replace this undecorated attribute with this decorated attribute.
            set_type_attr(cls, attr_name, attr_value_beartyped)
        # Else, this attribute is *NOT* beartypeable. In this case, silently
        # ignore this attribute.

    # ....................{ MONKEY-PATCH                   }....................
    # Pure-Python __sizeof__() dunder method wrapping the original C-based
    # __sizeof__() dunder method declared by this class.
    @wraps(cls_sizeof_old)
    def cls_sizeof_new(self) -> int:
        return cls_sizeof_old(self)  # type: ignore[call-arg]

    # Monkey-patch a @beartype-specific instance variable into this wrapper,
    # recording that this decorator has now decorated this class.
    #
    # Note that we intentionally set this variable to this class rather than an
    # arbitrary value (e.g., "False", "None"). Why? Because subclasses of this
    # class will inherit this wrapper. If we simply set this variable to an
    # arbitrary value, we would be unable to decide above between the following
    # two cases:
    # * Whether this wrapper was inherited from its superclass, in which case
    #   this class has yet to be decorated by @beartype.
    # * Whether this wrapper was *NOT* inherited from its superclass, in which
    #   case this class has already been decorated by @beartype.
    cls_sizeof_new.__beartyped_cls = cls  # type: ignore[attr-defined]

    # Replace the original C-based __sizeof__() dunder method with this wrapper.
    # We intentionally call our set_type_attr() setter rather than attempting to
    # set this attribute directly. The latter approach efficiently succeeds for
    # standard pure-Python mutable classes but catastrophically fails for
    # non-standard C-based immutable classes (e.g., "enum.Enum" subclasses).
    set_type_attr(cls, '__sizeof__', cls_sizeof_new)

    # Return this class as is.
    return cls  # type: ignore[return-value]

# ....................{ PRIVATE ~ globals                  }....................
_BEARTYPED_MODULE_TO_TYPE_NAME: Dict[str, Set[str]] = defaultdict(set)
'''
**Decorated classname registry (i.e., dictionary mapping from the
fully-qualified name of each module defining one or more classes decorated by
the :func:`beartype.beartype` decorator to the set of the unqualified basenames
of all classes in that module decorated by that decorator).
'''


from typing import Any

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch import Tensor

class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.lstm = nn.LSTM(4, 64)
        self.fc = nn.Linear(64, 1)

    def __call__(self, *input, **kwargs) -> Any:
        return super().__call__(*input, **kwargs)

    def forward(self, input: Tensor):
        x, (_, _) = self.lstm(input)  # Output of LSTM is of the structure (outs, (h_n, c_n)), we ignore the two latter
        x = x.view(-1, 64)
        x = self.fc(x)

        return x

def custom_loss(inputs, targets):
    return torch.sqrt(torch.sum(torch.pow(torch.log(inputs + 1) - torch.log(targets + 1), 2)) / inputs.shape[0])

if __name__ == '__main__':
    model = Model()
    data = np.array([[1, 2, 3, 4], [2, 3, 4, 5], [3, 4, 5, 6]])
    tensor = torch.from_numpy(data).view(-1, 1, 4).float()

    loss_function = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=0.1)

    print(model(tensor))

    for epoch in range(1000):
        model.zero_grad()

        output = model(tensor)
        output = output.view(-1)

        loss = custom_loss(output, torch.tensor([10.0, 10.0, 10.0]))
        loss.backward()
        optimizer.step()

    print(model(tensor))


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import KDEFunction
import time
import plot_map
import math
from SecondtoTime import sec2time

# 读取GPS数据
# 数据读取time：57~58s
data = pd.read_csv(r'../Data Process/RawGPSData.csv', header=None)
data.columns = ['VehicleNum', 'Stime', 'Lng', 'Lat']
# print(data.shape)
# print(data.head(5))

# 获取出租车数据经纬度最值
# maxLng = data['Lng'].max()
# maxLat = data['Lat'].max()
# minLng = data['Lng'].min()
# minLat = data['Lat'].min()

# bounds
# print(maxLng, minLng, maxLat, minLat)
# bounds = [117.9963607788086, 115.20406341552734, 40.87900924682617, 39.00091552734375]
# bounds = [115.20406341552734, 39.00091552734375, 117.9963607788086, 40.87900924682617]

# 北京市经纬度边界【最小经度，最小纬度，最大经度，最大纬度】
# bounds = [115.7, 39.4, 117.4, 41.6]
bounds = [116.0, 39.6, 116.8, 40.2]
MaxX = bounds[2]
MinX = bounds[0]
MaxY = bounds[3]
MinY = bounds[1]

# 坐标采样个数,即Sample_Num*Sample_Num
Sample_Num = 100

# 计算采样间隔
Sample_Space_X = (MaxX - MinX) / Sample_Num
Sample_Space_Y = (MaxY - MinY) / Sample_Num

# 采样矩阵，Sample_Num*Sample_Num
Sample_X, Sample_Y = np.mgrid[MinX:MaxX:complex(0, Sample_Num), MinY:MaxY:complex(0, Sample_Num)]  # 步长实数为间隔，虚数为个数
# Sample_X, Sample_Y = np.mgrid[MinX:MaxX:Sample_Space_X, MinY:MaxY:Sample_Space_Y]  # 步长实数为间隔，虚数为个数
Sample_Metric = np.transpose(np.vstack((Sample_X.ravel(), Sample_Y.ravel())))
# print(Sample_Metric)
print(Sample_Metric.shape)

# 计算100秒时的热点数据
test_sec = 28800
data_10 = data[data['Stime'] == test_sec]
# print(data_10)
Point_Num = len(data_10)
X = np.array(data_10[['Lng', 'Lat']])

# 计时
start = time.process_time()

# # 计算地球经纬度与公里数转换 1度多少千米
deltaLng_km = (MaxX-MinX)*(2 * math.pi * 6371 * math.cos((MinY + MaxY) * math.pi/360))/360
deltaLat_km = (MaxY-MinY)*(2 * math.pi * 6371)/360
# print(deltaLng_km, deltaLat_km)

# 公里数
bandwidth_list = [0.1, 0.25, 0.5, 0.75, 1, 2, 3, 6, 10]
# 不同带宽分析
bandwidth1 = np.array(bandwidth_list)/np.sqrt(deltaLat_km*deltaLng_km)

# 0.25-0.5km效果最佳

print(bandwidth1)

# 用来正常显示中文标签
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['font.size'] = 3.5
# plt.rcParams['font.size'] = 4
fig = plt.figure(dpi=300)

for i in range(len(bandwidth1)):

    # KDEFunction
    kde, dense = KDEFunction.KDE(X, Sample_Metric, bandwidth1[i])

    # 打印进度
    print(i, '/8', dense.max(), dense.min())

    # 保存特征数据num*3 经度|纬度|核密度估计值
    # KDEFunction.SaveNPZ(Sample_Metric, dense, "fdata")

    # 画核密度等高线图
    # KDEFunction.PlotContour(MinX, MaxX, MinY, MaxY, Sample_Space_X, Sample_Space_Y, kde, plotbool=True, savebool=False, picname="KDE_Random")

    # 北京市经纬度边界
    # fig = plt.figure(dpi=300)
    # ax = plt.subplot(111)
    # plt.sca(ax)
    # fig.tight_layout(rect=(0.05, 0.1, 1, 0.9))
    # plt.subplot(3, 3, i+1)

    # 散点图分布
    # plt.scatter(data_10['Lng'], data_10['Lat'], s=1, alpha=0.5)

    # 画核密度等高线图
    # 作图
    Contour_X = Sample_X
    Contour_Y = Sample_Y
    Contour_Z = np.reshape(dense, Contour_X.shape)

    # axis = fig.gca()
    axis = fig.add_subplot(3, 3, i+1)

    # 背景
    plot_map.plot_map(plt, bounds, zoom=12, style=3)
    # plt.axis('off')
    plt.axis('on')
    axis.set_xlim(MinX, MaxX)
    axis.set_ylim(MinY, MaxY)

    contourf_set = axis.contourf(Contour_X, Contour_Y, Contour_Z, levels=np.linspace(0, Contour_Z.max(), 12), alpha=0.7, cmap='Reds')
    # contourf_set = axis.contourf(Contour_X, Contour_Y, Contour_Z, alpha=0.6, cmap='Reds')

    # 设置colorbar
    cbar = fig.colorbar(contourf_set, fraction=0.04, pad=0.04)
    cbar.set_clim(Contour_Z.min(), Contour_Z.max())

    # contour_set = axis.contour(Contour_X, Contour_Y, Contour_Z, colors='r')
    # axis.clabel(contour_set, inline=1, fontsize=5)

    # axis.set_xlabel('Longitude')
    # axis.set_ylabel('Latitude')
    plt.title("带宽" + str(bandwidth_list[i]) + "km时" + "出租车热点分布（8点）", fontsize='7')

# axis on
plt.subplots_adjust(top=0.92, bottom=0.08, right=0.96, left=0.04, hspace=0.32, wspace=0.25)
# axis off
# plt.subplots_adjust(top=0.92, bottom=0.08, right=0.96, left=0.04, hspace=0.4, wspace=0.15)

end = time.process_time()
print("Data Process 2 Step time:", end - start)

# 图片存储在HotSpotPlot目录
plt_name = './KDEbandAnaysis/HotSpot_band_list_update_7'
plt.savefig(plt_name+'.png')
plt.savefig(plt_name+'.svg')

plt.show()


from __future__ import annotations

import pickle
from typing import Type, Union

import pytest
from commonmodel.field_types import (
    Binary,
    Boolean,
    DateTime,
    Decimal,
    FieldType,
    Float,
    Integer,
    Interval,
    Json,
    LongBinary,
    LongText,
    Text,
    str_to_field_type,
    ensure_field_type,
)

def test_instantiation():
    Boolean()
    Integer()
    Float()
    Decimal()
    Decimal(12)
    Decimal(12, 2)
    Decimal(scale=12, precision=2)
    Binary()
    LongBinary()
    Text()
    Text(length=255)
    LongText()
    Interval()

def test_repr():
    assert repr(Boolean()) == "Boolean"
    assert repr(Text(length=255)) == "Text(length=255)"
    assert repr(Decimal()) == "Decimal"
    assert repr(Decimal(10)) == "Decimal(precision=10)"
    assert repr(Decimal(16, 2)) == "Decimal(precision=16, scale=2)"

def test_repr_round_trip():
    assert ensure_field_type(repr(Text())) == Text()
    assert ensure_field_type(repr(Boolean())) == Boolean()
    assert ensure_field_type(repr(Decimal())) == Decimal()
    assert ensure_field_type(repr(Decimal(scale=16, precision=2))) == Decimal(2, 16)

def test_pickle():
    assert pickle.loads(pickle.dumps(Boolean())) == Boolean()
    assert pickle.loads(pickle.dumps(Decimal())) == Decimal()
    assert pickle.loads(pickle.dumps(Decimal(16, 2))) == Decimal(16, 2)
    assert pickle.loads(pickle.dumps(Text(length=255))) == Text(length=255)

@pytest.mark.parametrize(
    "s,expected",
    [
        ("Text", Text()),
        ("Boolean", Boolean()),
        ("Text(length=3)", Text(length=3)),
        ("Text(3)", Text(length=3)),
        ("DateTime", DateTime()),
        ("Interval", Interval()),
        ("DateTime(timezone=False)", DateTime(timezone=False)),
        ("Decimal(precision=16, scale=2)", Decimal(16, 2)),
        ("Decimal(16,2)", Decimal(16, 2)),
        ("Decimal()", Decimal()),
        ("Json(astext_type=Text())", Json()),
    ],
)
def test_str_to_field_types(s: str, expected: Union[FieldType, Type[FieldType]]):
    assert str_to_field_type(s) == expected


import hashlib
import os
import os.path
import re
import struct
import subprocess
import sys

from conftest import need_to_install_package_err

from elftools.elf.elffile import ELFFile

import pytest

TEST_DIR = os.path.join(os.path.abspath(os.path.dirname(__file__)), "elf2image")

try:
    import esptool
except ImportError:
    need_to_install_package_err()

def try_delete(path):
    try:
        os.remove(path)
    except OSError:
        pass

def segment_matches_section(segment, section):
    """segment is an ImageSegment from an esptool binary.
    section is an elftools ELF section

    Returns True if they match
    """
    sh_size = (section.header.sh_size + 0x3) & ~3  # pad length of ELF sections
    return section.header.sh_addr == segment.addr and sh_size == len(segment.data)

@pytest.mark.host_test
class BaseTestCase:
    @classmethod
    def setup_class(self):
        # Save the current working directory to be resotred later
        self.stored_dir = os.getcwd()
        os.chdir(TEST_DIR)

    @classmethod
    def teardown_class(self):
        # Restore the stored working directory
        os.chdir(self.stored_dir)

    def assertEqualHex(self, expected, actual, message=None):
        try:
            expected = hex(expected)
        except TypeError:  # if expected is character
            expected = hex(ord(expected))
        try:
            actual = hex(actual)
        except TypeError:  # if actual is character
            actual = hex(ord(actual))
        assert expected == actual, message

    def assertImageDoesNotContainSection(self, image, elf, section_name):
        """
        Assert an esptool binary image object does not
        contain the data for a particular ELF section.
        """
        with open(elf, "rb") as f:
            e = ELFFile(f)
            section = e.get_section_by_name(section_name)
            assert section, f"{section_name} should be in the ELF"
            sh_addr = section.header.sh_addr
            data = section.data()
            # no section should start at the same address as the ELF section.
            for seg in sorted(image.segments, key=lambda s: s.addr):
                print(
                    f"comparing seg {seg.addr:#x} sec {sh_addr:#x} len {len(data):#x}"
                )
                assert (
                    seg.addr != sh_addr
                ), f"{section_name} should not be in the binary image"

    def assertImageContainsSection(self, image, elf, section_name):
        """
        Assert an esptool binary image object contains
        the data for a particular ELF section.
        """
        with open(elf, "rb") as f:
            e = ELFFile(f)
            section = e.get_section_by_name(section_name)
            assert section, f"{section_name} should be in the ELF"
            sh_addr = section.header.sh_addr
            data = section.data()
            # section contents may be smeared across multiple image segments,
            # so look through each segment and remove it from ELF section 'data'
            # as we find it in the image segments. When we're done 'data' should
            # all be accounted for
            for seg in sorted(image.segments, key=lambda s: s.addr):
                print(
                    f"comparing seg {seg.addr:#x} sec {sh_addr:#x} len {len(data):#x}"
                )
                if seg.addr == sh_addr:
                    overlap_len = min(len(seg.data), len(data))
                    assert (
                        data[:overlap_len] == seg.data[:overlap_len]
                    ), f"ELF '{section_name}' section has mis-matching bin image data"
                    sh_addr += overlap_len
                    data = data[overlap_len:]

            # no bytes in 'data' should be left unmatched
            assert len(data) == 0, (
                f"ELF {elf} section '{section_name}' has no encompassing"
                f" segment(s) in bin image (image segments: {image.segments})"
            )

    def assertImageInfo(self, binpath, chip="esp8266", assert_sha=False):
        """
        Run esptool.py image_info on a binary file,
        assert no red flags about contents.
        """
        cmd = [sys.executable, "-m", "esptool", "--chip", chip, "image_info", binpath]
        try:
            output = subprocess.check_output(cmd)
            output = output.decode("utf-8")
            print(output)
        except subprocess.CalledProcessError as e:
            print(e.output)
            raise
        assert re.search(
            r"Checksum: [a-fA-F0-9]{2} \(valid\)", output
        ), "Checksum calculation should be valid"
        if assert_sha:
            assert re.search(
                r"Validation Hash: [a-fA-F0-9]{64} \(valid\)", output
            ), "SHA256 should be valid"
        assert (
            "warning" not in output.lower()
        ), "Should be no warnings in image_info output"

    def run_elf2image(self, chip, elf_path, version=None, extra_args=[]):
        """Run elf2image on elf_path"""
        cmd = [sys.executable, "-m", "esptool", "--chip", chip, "elf2image"]
        if version is not None:
            cmd += ["--version", str(version)]
        cmd += [elf_path] + extra_args
        print("\nExecuting {}".format(" ".join(cmd)))
        try:
            output = subprocess.check_output(cmd)
            output = output.decode("utf-8")
            print(output)
            assert (
                "warning" not in output.lower()
            ), "elf2image should not output warnings"
        except subprocess.CalledProcessError as e:
            print(e.output)
            raise

class TestESP8266V1Image(BaseTestCase):
    ELF = "esp8266-nonosssdk20-iotdemo.elf"
    BIN_LOAD = "esp8266-nonosssdk20-iotdemo.elf-0x00000.bin"
    BIN_IROM = "esp8266-nonosssdk20-iotdemo.elf-0x10000.bin"

    @classmethod
    def setup_class(self):
        super(TestESP8266V1Image, self).setup_class()
        self.run_elf2image(self, "esp8266", self.ELF, 1)

    @classmethod
    def teardown_class(self):
        super(TestESP8266V1Image, self).teardown_class()
        try_delete(self.BIN_LOAD)
        try_delete(self.BIN_IROM)

    def test_irom_bin(self):
        with open(self.ELF, "rb") as f:
            e = ELFFile(f)
            irom_section = e.get_section_by_name(".irom0.text")
            assert (
                irom_section.header.sh_size == os.stat(self.BIN_IROM).st_size
            ), "IROM raw binary file should be same length as .irom0.text section"

    def test_loaded_sections(self):
        image = esptool.bin_image.LoadFirmwareImage("esp8266", self.BIN_LOAD)
        # Adjacent sections are now merged, len(image.segments) should
        # equal 2 (instead of 3).
        assert len(image.segments) == 2
        self.assertImageContainsSection(image, self.ELF, ".data")
        self.assertImageContainsSection(image, self.ELF, ".text")
        # Section .rodata is merged in the binary with the previous one,
        # so it won't be found in the binary image.
        self.assertImageDoesNotContainSection(image, self.ELF, ".rodata")

class TestESP8266V12SectionHeaderNotAtEnd(BaseTestCase):
    """Ref https://github.com/espressif/esptool/issues/197 -
    this ELF image has the section header not at the end of the file"""

    ELF = "esp8266-nonossdkv12-example.elf"
    BIN_LOAD = ELF + "-0x00000.bin"
    BIN_IROM = ELF + "-0x40000.bin"

    @classmethod
    def teardown_class(self):
        try_delete(self.BIN_LOAD)
        try_delete(self.BIN_IROM)

    def test_elf_section_header_not_at_end(self):
        self.run_elf2image("esp8266", self.ELF)
        image = esptool.bin_image.LoadFirmwareImage("esp8266", self.BIN_LOAD)
        assert len(image.segments) == 3
        self.assertImageContainsSection(image, self.ELF, ".data")
        self.assertImageContainsSection(image, self.ELF, ".text")
        self.assertImageContainsSection(image, self.ELF, ".rodata")

class TestESP8266V2Image(BaseTestCase):
    def _test_elf2image(self, elfpath, binpath, mergedsections=[]):
        try:
            self.run_elf2image("esp8266", elfpath, 2)
            image = esptool.bin_image.LoadFirmwareImage("esp8266", binpath)
            print("In test_elf2image", len(image.segments))
            assert 4 - len(mergedsections) == len(image.segments)
            sections = [".data", ".text", ".rodata"]
            # Remove the merged sections from the `sections` list
            sections = [sec for sec in sections if sec not in mergedsections]
            for sec in sections:
                self.assertImageContainsSection(image, elfpath, sec)
            for sec in mergedsections:
                self.assertImageDoesNotContainSection(image, elfpath, sec)

            irom_segment = image.segments[0]
            assert irom_segment.addr == 0, "IROM segment 'load address' should be zero"
            with open(elfpath, "rb") as f:
                e = ELFFile(f)
                sh_size = (
                    e.get_section_by_name(".irom0.text").header.sh_size + 15
                ) & ~15
                assert len(irom_segment.data) == sh_size, (
                    f"irom segment ({len(irom_segment.data):#x}) should be same size "
                    f"(16 padded) as .irom0.text section ({sh_size:#x})"
                )

            # check V2 CRC (for ESP8266 SDK bootloader)
            with open(binpath, "rb") as f:
                f.seek(-4, os.SEEK_END)
                image_len = f.tell()
                crc_stored = struct.unpack("<I", f.read(4))[0]
                f.seek(0)
                crc_calc = esptool.bin_image.esp8266_crc32(f.read(image_len))
                assert crc_stored == crc_calc

            # test imageinfo doesn't fail
            self.assertImageInfo(binpath)

        finally:
            try_delete(binpath)

    def test_nonossdkimage(self):
        ELF = "esp8266-nonossdkv20-at-v2.elf"
        BIN = "esp8266-nonossdkv20-at-v2-0x01000.bin"
        self._test_elf2image(ELF, BIN)

    def test_espopenrtosimage(self):
        ELF = "esp8266-openrtos-blink-v2.elf"
        BIN = "esp8266-openrtos-blink-v2-0x02000.bin"
        # .rodata section is merged with the previous one: .data
        self._test_elf2image(ELF, BIN, [".rodata"])

class TestESP32Image(BaseTestCase):
    def _test_elf2image(self, elfpath, binpath, extra_args=[]):
        try:
            self.run_elf2image("esp32", elfpath, extra_args=extra_args)
            image = esptool.bin_image.LoadFirmwareImage("esp32", binpath)
            self.assertImageInfo(
                binpath,
                "esp32",
                True if "--ram-only-header" not in extra_args else False,
            )
            return image
        finally:
            try_delete(binpath)

    def test_bootloader(self):
        ELF = "esp32-bootloader.elf"
        BIN = "esp32-bootloader.bin"
        image = self._test_elf2image(ELF, BIN)
        assert len(image.segments) == 3
        for section in [".iram1.text", ".iram_pool_1.text", ".dram0.rodata"]:
            self.assertImageContainsSection(image, ELF, section)

    def test_app_template(self):
        ELF = "esp32-app-template.elf"
        BIN = "esp32-app-template.bin"
        image = self._test_elf2image(ELF, BIN)
        # Adjacent sections are now merged, len(image.segments) should
        # equal 5 (instead of 6).
        assert len(image.segments) == 5
        # the other segment is a padding or merged segment
        for section in [
            ".iram0.vectors",
            ".dram0.data",
            ".flash.rodata",
            ".flash.text",
        ]:
            self.assertImageContainsSection(image, ELF, section)
        # check that merged sections are not in the binary image
        for mergedsection in [".iram0.text"]:
            self.assertImageDoesNotContainSection(image, ELF, mergedsection)

    def test_too_many_sections(self, capsys):
        ELF = "esp32-too-many-sections.elf"
        BIN = "esp32-too-many-sections.bin"
        with pytest.raises(subprocess.CalledProcessError):
            self._test_elf2image(ELF, BIN)
        output = capsys.readouterr().out
        assert "max 16" in output
        assert "linker script" in output

    def test_use_segments(self):
        ELF = "esp32-zephyr.elf"
        BIN = "esp32-zephyr.bin"
        # default behaviour uses ELF sections,
        # this ELF will produce 8 segments in the bin
        image = self._test_elf2image(ELF, BIN)
        # Adjacent sections are now merged, len(image.segments) should
        # equal 4 (instead of 8).
        assert len(image.segments) == 4

        # --use_segments uses ELF segments(phdrs), produces just 2 segments in the bin
        image = self._test_elf2image(ELF, BIN, ["--use_segments"])
        assert len(image.segments) == 2

    def test_ram_only_header(self):
        ELF = "esp32-app-template.elf"
        BIN = "esp32-app-template.bin"
        # --ram-only-header produces just 2 visible segments in the bin
        image = self._test_elf2image(ELF, BIN, ["--ram-only-header"])
        assert len(image.segments) == 2

class TestESP8266FlashHeader(BaseTestCase):
    def test_2mb(self):
        ELF = "esp8266-nonossdkv20-at-v2.elf"
        BIN = "esp8266-nonossdkv20-at-v2-0x01000.bin"
        try:
            self.run_elf2image(
                "esp8266",
                ELF,
                version=2,
                extra_args=["--flash_size", "2MB", "--flash_mode", "dio"],
            )
            with open(BIN, "rb") as f:
                header = f.read(4)
                print(f"header {header}")
                self.assertEqualHex(0xEA, header[0])
                self.assertEqualHex(0x02, header[2])
                self.assertEqualHex(0x30, header[3])
        finally:
            try_delete(BIN)

class TestESP32FlashHeader(BaseTestCase):
    def test_16mb(self):
        ELF = "esp32-app-template.elf"
        BIN = "esp32-app-template.bin"
        try:
            self.run_elf2image(
                "esp32",
                ELF,
                extra_args=[
                    "--flash_size",
                    "16MB",
                    "--flash_mode",
                    "dio",
                    "--min-rev",
                    "1",
                ],
            )
            with open(BIN, "rb") as f:
                header = f.read(24)
                self.assertEqualHex(0xE9, header[0])
                self.assertEqualHex(0x02, header[2])
                self.assertEqualHex(0x40, header[3])
                self.assertEqualHex(0x01, header[14])  # chip revision
        finally:
            try_delete(BIN)

class TestELFSHA256(BaseTestCase):
    ELF = "esp32-app-cust-ver-info.elf"
    SHA_OFFS = 0xB0  # absolute offset of the SHA in the .bin file
    BIN = "esp32-app-cust-ver-info.bin"

    """
    esp32-app-cust-ver-info.elf was built with the following application version info:

    const __attribute__((section(".rodata_desc"))) esp_app_desc_t esp_app_desc = {
        .magic_word = 0xffffffff,
        .secure_version = 0xffffffff,
        .reserv1 = {0xffffffff, 0xffffffff},
        .version = "_______________________________",
        .project_name = "-------------------------------",
        .time = "xxxxxxxxxxxxxxx",
        .date = "yyyyyyyyyyyyyyy",
        .idf_ver = "zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz",
        .app_elf_sha256 =
        {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
        .reserv2 = {0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff,
                    0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff,
                    0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff,
                    0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff},
    };

    This leaves zeroes only for the fiels of SHA-256 and the test will fail
    if the placement of zeroes are tested at the wrong place.

    00000000: e907 0020 780f 0840 ee00 0000 0000 0000  ... x..@........
    00000010: 0000 0000 0000 0001 2000 403f 605a 0000  ........ .@?`Z..
    00000020: ffff ffff ffff ffff ffff ffff ffff ffff  ................
    00000030: 5f5f 5f5f 5f5f 5f5f 5f5f 5f5f 5f5f 5f5f  ________________
    00000040: 5f5f 5f5f 5f5f 5f5f 5f5f 5f5f 5f5f 5f00  _______________.
    00000050: 2d2d 2d2d 2d2d 2d2d 2d2d 2d2d 2d2d 2d2d  ----------------
    00000060: 2d2d 2d2d 2d2d 2d2d 2d2d 2d2d 2d2d 2d00  ---------------.
    00000070: 7878 7878 7878 7878 7878 7878 7878 7800  xxxxxxxxxxxxxxx.
    00000080: 7979 7979 7979 7979 7979 7979 7979 7900  yyyyyyyyyyyyyyy.
    00000090: 7a7a 7a7a 7a7a 7a7a 7a7a 7a7a 7a7a 7a7a  zzzzzzzzzzzzzzzz
    000000a0: 7a7a 7a7a 7a7a 7a7a 7a7a 7a7a 7a7a 7a00  zzzzzzzzzzzzzzz.
    000000b0: 0000 0000 0000 0000 0000 0000 0000 0000  ................    SHA-256 here
    000000c0: 0000 0000 0000 0000 0000 0000 0000 0000  ................
    000000d0: ffff ffff ffff ffff ffff ffff ffff ffff  ................
    000000e0: ffff ffff ffff ffff ffff ffff ffff ffff  ................
    000000f0: ffff ffff ffff ffff ffff ffff ffff ffff  ................
    00000100: ffff ffff ffff ffff ffff ffff ffff ffff  ................
    00000110: ffff ffff ffff ffff ffff ffff ffff ffff  ................
    00000120: 6370 755f 7374 6172 7400 0000 1b5b 303b  cpu_start....[0;

    """

    def test_binary_patched(self):
        try:
            self.run_elf2image(
                "esp32",
                self.ELF,
                extra_args=["--elf-sha256-offset", f"{self.SHA_OFFS:#x}"],
            )
            image = esptool.bin_image.LoadFirmwareImage("esp32", self.BIN)
            rodata_segment = image.segments[0]
            bin_sha256 = rodata_segment.data[
                self.SHA_OFFS - 0x20 : self.SHA_OFFS - 0x20 + 32
            ]  # subtract 0x20 byte header here

            with open(self.ELF, "rb") as f:
                elf_computed_sha256 = hashlib.sha256(f.read()).digest()

            with open(self.BIN, "rb") as f:
                f.seek(self.SHA_OFFS)
                bin_sha256_raw = f.read(len(elf_computed_sha256))

            assert elf_computed_sha256 == bin_sha256
            assert elf_computed_sha256 == bin_sha256_raw
        finally:
            try_delete(self.BIN)

    def test_no_overwrite_data(self, capsys):
        with pytest.raises(subprocess.CalledProcessError):
            self.run_elf2image(
                "esp32",
                "esp32-bootloader.elf",
                extra_args=["--elf-sha256-offset", "0xb0"],
            )
        output = capsys.readouterr().out
        assert "SHA256" in output
        assert "zero" in output

class TestHashAppend(BaseTestCase):
    ELF = "esp32-bootloader.elf"
    BIN = "esp32-bootloader.bin"

    # 15th byte of the extended header after the 8-byte image header
    HASH_APPEND_OFFSET = 15 + 8

    @classmethod
    def teardown_class(self):
        try_delete(self.BIN)

    def test_hash_append(self):
        self.run_elf2image(
            "esp32",
            self.ELF,
            extra_args=["-o", self.BIN],
        )
        with open(self.BIN, "rb") as f:
            bin_with_hash = f.read()

        assert bin_with_hash[self.HASH_APPEND_OFFSET] == 1

        # drop the last 32 bytes (SHA256 digest)
        expected_bin_without_hash = bytearray(bin_with_hash[:-32])
        # disable the hash append byte in the file header
        expected_bin_without_hash[self.HASH_APPEND_OFFSET] = 0

        try_delete(self.BIN)
        self.run_elf2image(
            "esp32",
            self.ELF,
            extra_args=["--dont-append-digest", "-o", self.BIN],
        )

        with open(self.BIN, "rb") as f:
            bin_without_hash = f.read()

        assert bin_without_hash[self.HASH_APPEND_OFFSET] == 0
        assert bytes(expected_bin_without_hash) == bin_without_hash


from pathlib import Path

# ########################
# DIRECTORIES
# ########################
''' Define new directory variables according to following rules:
    1) Use CAPITAL LETTERS
    2) Start with phrase 'DIR'
    3) Use '_' as delimiter
'''
# DIR_SRC points to the folder source code is being held
DIR_SRC = Path(__file__).parent.resolve()
DIR_ROOT = DIR_SRC.parent  # Points to Root directory of the project

# ########################
# SINGLE FILES
# ########################
''' Define new file variables according to following rules:
    1) Use CAPITAL LETTERS
    2) Start with phrase 'FILE'
    3) Use '_' as delimiter
'''
FILE_REQS = DIR_SRC / "requirements.txt"  # Points to requirements.txt
FILE_FLASK = DIR_SRC / "parkingmanagement.py"
FILE_DATABASE = DIR_SRC / 'parkinglot.db'

# ########################
# SETTINGS
# ########################
DEMO = True
VERBOSE = True
DEBUG = True
CLEAN_DB = True
SILENT_INSTALL = False

# ########################
# CLASS (for Flask)
# ########################
class Config(object):
    SECRET_KEY = 'fission-liturgy9-ru+f_ous-comely-burp'
    """docstring for Config"""
    def __init__(self, arg):
        super(Config, self).__init__()
        self.arg = arg


#!/usr/bin/env python3
# Copyright (C) 2018 RDA Technologies Limited and/or its affiliates("RDA").
# All rights reserved.
#
# This software is supplied "AS IS" without any warranties.
# RDA assumes no responsibility or liability for the use of the software,
# conveys no license or title under any patent, copyright, or mask work
# right to the product. RDA reserves the right to make changes in the
# software without notification.  RDA also make no representation or
# warranty that such application will be suitable for the specified use
# without further testing or modification.

import argparse
import sys
import binascii

DESCRIPTION = """
Calculate CRC32 of files
"""

def main():
    parser = argparse.ArgumentParser(description=DESCRIPTION)

    parser.add_argument(
        "fnames",
        metavar="FILE_NAME",
        nargs="+",
        help="file names to calculate CRC32")

    args = parser.parse_args()

    for fname in args.fnames:
        with open(fname, 'rb') as fh:
            data = fh.read()

        crc = binascii.crc32(data)
        print('%08x %s' % (crc, fname))

if __name__ == "__main__":
    main()


"""
Downloading input files from URIs, with plugin modules for different URI schemes

Download URI plugins are installed & registered using the setuptools entry point group
"miniwdl.plugin.file_download", with name equal to the URI scheme (e.g. "gs" or "s3").

The plugin entry point should be a context manager, which the runtime keeps open for the duration of
the download operation. Given the desired URI, it should quickly yield a tuple with:
    1. source code of a WDL 1.0 task to perform the download
    2. dict of Cromwell-style JSON inputs to give to the task
miniwdl then executes this specified operation, expecting it to produce an output "File file" with
the downloaded file. By doing the heavy lifting in a WDL task, the operation gets to inherit all
the functionality of miniwdl's task runtime, e.g. pulling docker image with binary dependencies,
resource scheduling & isolation, logging, error/signal handling, retry, etc.
The Python context manager itself might be used to obtain and manage the lifetime of any needed
security credentials.
"""
import os
import logging
import traceback
import tempfile
import hashlib
import shlex
from contextlib import ExitStack
from typing import Optional, Generator, Dict, Any, Tuple, Callable
from . import config
from .cache import CallCache
from .._util import compose_coroutines
from .._util import StructuredLogMessage as _

def _load(cfg: config.Loader):
    if getattr(cfg, "_downloaders", None):
        return

    # default public URI downloaders
    file_downloaders = {
        "https": aria2c_downloader,
        "http": aria2c_downloader,
        "ftp": aria2c_downloader,
    }
    directory_downloaders = {}

    # plugins
    for plugin_name, plugin_fn in config.load_plugins(cfg, "file_download"):
        file_downloaders[plugin_name] = plugin_fn
    for plugin_name, plugin_fn in config.load_plugins(cfg, "directory_download"):
        directory_downloaders[plugin_name] = plugin_fn

    setattr(cfg, "_downloaders", (file_downloaders, directory_downloaders))

def _downloader(
    cfg: config.Loader, uri: str, directory: bool = False
) -> Optional[Callable[..., Generator[Dict[str, Any], Dict[str, Any], None]]]:
    _load(cfg)
    colon = uri.find(":")
    if colon <= 0:
        return None
    scheme = uri[:colon]
    return getattr(cfg, "_downloaders")[1 if directory else 0].get(scheme, None)

def able(cfg: config.Loader, uri: Optional[str], directory: bool = False) -> bool:
    """
    Returns True if uri appears to be a URI we know how to download
    """
    return bool(uri and _downloader(cfg, uri, directory=directory) is not None)

def run(
    cfg: config.Loader, logger: logging.Logger, uri: str, directory: bool = False, **kwargs
) -> str:
    """
    Download the URI and return the local filename.

    kwargs are passed through to ``run_local_task``, so ``run_dir`` and ``logger_prefix`` may be
    useful in particular.
    """

    from .error import RunFailed, DownloadFailed, Terminated, error_json
    from .task import run_local_task
    from .. import parse_document, values_from_json, values_to_json, Walker

    gen = _downloader(cfg, uri, directory=directory)
    assert gen
    try:
        logger.info(_(f"start {'directory ' if directory else ''}download", uri=uri))
        with compose_coroutines([lambda kwargs: gen(cfg, logger, **kwargs)], {"uri": uri}) as cor:
            recv = next(cor)

            if "task_wdl" in recv:
                task_wdl, inputs = (recv[k] for k in ["task_wdl", "inputs"])

                doc = parse_document(task_wdl, version="development")  # pyre-ignore
                assert len(doc.tasks) == 1 and not doc.workflow
                doc.typecheck()
                Walker.SetParents()(doc)
                task = doc.tasks[0]
                inputs = values_from_json(inputs, task.available_inputs)  # pyre-ignore
                subdir, outputs_env = run_local_task(
                    cfg, task, inputs, run_id=("download-" + task.name), **kwargs
                )

                recv = cor.send(
                    {"outputs": values_to_json(outputs_env), "dir": subdir}  # pyre-ignore
                )

            ans = recv["outputs"]["directory" if directory else "file"]
            assert isinstance(ans, str) and os.path.exists(ans)
            logger.info(_(f"downloaded{' directory' if directory else ' file'}", uri=uri, file=ans))
            return ans

    except RunFailed as exn:
        if isinstance(exn.__cause__, Terminated):
            raise exn.__cause__ from None
        raise DownloadFailed(uri) from exn.__cause__
    except Exception as exn:
        logger.debug(traceback.format_exc())
        logger.error(_("downloader error", uri=uri, **error_json(exn)))
        raise DownloadFailed(uri) from exn

def run_cached(
    cfg,
    logger: logging.Logger,
    cache: CallCache,
    uri: str,
    run_dir: str,
    directory: bool = False,
    **kwargs,
) -> Tuple[bool, str]:
    """
    Cached download logic: returns the file from the cache if available; otherwise, runs the
    download and puts it into the cache before returning
    """
    cached = cache.get_download(uri, directory=directory, logger=logger)
    if cached:
        return True, cached
    cache_path = cache.download_cacheable(uri, directory=directory)
    cache_path_preexists = cache_path and os.path.exists(cache_path)
    if cache_path and not cache_path_preexists:
        # run the download in a holding area under the cache directory, so that it can later be
        # moved atomically into its final cache location
        run_dir = os.path.join(
            cfg["file_io"]["root"], os.path.join(cfg["download_cache"]["dir"], "ops")
        )
    filename = run(cfg, logger, uri, directory=directory, run_dir=run_dir, **kwargs)
    if cache_path_preexists:
        # a cache entry had already existed, but we didn't use it (--no-cache).
        # FIXME: it'd be better to replace the old copy...but what if another workflow is using it?
        logger.warning(
            _(
                "ignored a previously-cached download, which remains in the cache",
                uri=uri,
                downloaded=filename,
                cache_path=cache_path,
            )
        )
        # use the newly downloaded copy in the current run directory
        cache.memo_download(uri, filename, directory=directory)
        return False, filename
    return False, cache.put_download(
        uri, os.path.realpath(filename), directory=directory, logger=logger
    )

# WDL tasks for downloading a file based on its URI scheme

def aria2c_downloader(
    cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
) -> Generator[Dict[str, Any], Dict[str, Any], None]:
    wdl = r"""
    task aria2c {
        input {
            String uri
            String docker
            Int connections = 10
        }
        command <<<
            set -euxo pipefail
            mkdir __out
            cd __out
            aria2c -x ~{connections} -s ~{connections} \
                --file-allocation=none --retry-wait=2 --stderr=true --enable-color=false \
                "~{uri}"
        >>>
        output {
            File file = glob("__out/*")[0]
        }
        runtime {
            docker: docker
        }
    }
    """
    recv = yield {
        "task_wdl": wdl,
        "inputs": {"uri": uri, "docker": cfg["download_aria2c"]["docker"]},
    }
    yield recv  # pyre-ignore

def awscli_downloader(
    cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
) -> Generator[Dict[str, Any], Dict[str, Any], None]:
    inputs: Dict[str, Any] = {"uri": uri, "docker": cfg["download_awscli"]["docker"]}
    with ExitStack() as cleanup:
        inputs["aws_credentials"] = prepare_aws_credentials(cfg, logger, cleanup)

        wdl = r"""
        task aws_s3_cp {
            input {
                String uri
                String docker
                File? aws_credentials
            }

            command <<<
                set -euo pipefail
                if [ -n "~{aws_credentials}" ]; then
                    source "~{aws_credentials}"
                fi
                export AWS_RETRY_MODE=standard
                export AWS_MAX_ATTEMPTS=5
                set -x
                mkdir __out
                if ! aws s3 cp "~{uri}" __out/ ; then
                    # Retry with --no-sign-request in case the object is public. Without this flag,
                    # the previous invocation could have failed either because (i) no AWS
                    # credentials are available or (ii) the IAM policy restricts accessible S3
                    # buckets regardless of whether the desired object is public.
                    rm -f __out/*
                    >&2 echo 'Retrying with --no-sign-request in case the object is public.' \
                         ' If the overall operation fails, the real error may precede this message.'
                    aws s3 cp --no-sign-request "~{uri}" __out/
                fi
            >>>

            output {
                File file = glob("__out/*")[0]
            }

            runtime {
                docker: docker
            }
        }
        """
        recv = yield {"task_wdl": wdl, "inputs": inputs}
    yield recv  # pyre-ignore

def awscli_directory_downloader(
    cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
) -> Generator[Dict[str, Any], Dict[str, Any], None]:
    inputs: Dict[str, Any] = {"uri": uri, "docker": cfg["download_awscli"]["docker"]}
    with ExitStack() as cleanup:
        inputs["aws_credentials"] = prepare_aws_credentials(cfg, logger, cleanup)

        wdl = r"""
        task aws_s3_cp_directory {
            input {
                String uri
                String docker
                File? aws_credentials
            }

            String dnm = basename(uri, "/")

            command <<<
                set -euo pipefail
                if [ -n "~{aws_credentials}" ]; then
                    source "~{aws_credentials}"
                fi
                export AWS_RETRY_MODE=standard
                export AWS_MAX_ATTEMPTS=5
                set -x
                mkdir -p "__out/~{dnm}/"
                if ! aws s3 cp --recursive "~{uri}" "__out/~{dnm}/" ; then
                    # Retry with --no-sign-request in case the object is public. Without this flag,
                    # the previous invocation could have failed either because (i) no AWS
                    # credentials are available or (ii) the IAM policy restricts accessible S3
                    # buckets regardless of whether the desired object is public.
                    rm -f "__out/~{dnm}/*"
                    >&2 echo 'Retrying with --no-sign-request in case the folder is public.' \
                        ' If the overall operation fails, the real error may precede this message.'
                    aws s3 cp --recursive --no-sign-request "~{uri}" "__out/~{dnm}/"
                fi
            >>>

            output {
                Directory directory = "__out/" + dnm
            }

            runtime {
                docker: docker
            }
        }
        """
        recv = yield {"task_wdl": wdl, "inputs": inputs}
    yield recv  # pyre-ignore

def prepare_aws_credentials(
    cfg: config.Loader, logger: logging.Logger, cleanup: ExitStack
) -> Optional[str]:
    host_aws_credentials = {}
    if "AWS_EC2_METADATA_DISABLED" in os.environ:
        # https://github.com/aws/aws-cli/issues/5623
        host_aws_credentials["AWS_EC2_METADATA_DISABLED"] = os.environ["AWS_EC2_METADATA_DISABLED"]
    # get AWS credentials from boto3 (unless prevented by configuration)
    if cfg["download_awscli"].get_bool("host_credentials"):
        import boto3  # pyre-fixme

        try:
            b3creds = boto3.session.Session().get_credentials()
            host_aws_credentials["AWS_ACCESS_KEY_ID"] = b3creds.access_key
            host_aws_credentials["AWS_SECRET_ACCESS_KEY"] = b3creds.secret_key
            host_aws_credentials["AWS_SESSION_TOKEN"] = b3creds.token
        except Exception:
            pass

    if host_aws_credentials:
        # write credentials to temp file that'll self-destruct afterwards
        host_aws_credentials = (
            "\n".join(f"export {k}={shlex.quote(v)}" for (k, v) in host_aws_credentials.items())
            + "\n"
        )
        aws_credentials_file = cleanup.enter_context(
            tempfile.NamedTemporaryFile(
                prefix=hashlib.sha256(host_aws_credentials.encode()).hexdigest(),
                delete=True,
                mode="w",
            )
        )
        print(host_aws_credentials, file=aws_credentials_file, flush=True)
        # make file group-readable to ensure it'll be usable if the docker image runs as non-root
        os.chmod(aws_credentials_file.name, os.stat(aws_credentials_file.name).st_mode | 0o40)
        logger.getChild("awscli_downloader").info("loaded host AWS credentials")
        return aws_credentials_file.name
    else:
        logger.getChild("awscli_downloader").info(
            "no AWS credentials available via host awscli/boto3; if needed, "
            "configure them and set [download_awscli] host_credentials=true. "
            "(On EC2: awscli might still assume role from instance metadata "
            "service.)"
        )
        return None

def gsutil_downloader(
    cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
) -> Generator[Dict[str, Any], Dict[str, Any], None]:
    """
    Built-in downloader plugin for public gs:// URIs; registered by setup.cfg entry_points section

    TODO: adopt security credentials from runtime environment
    """
    if uri == "gs://8675309":
        # hook for test coverage of exception handler
        raise RuntimeError("don't change your number")
    wdl = r"""
    task gsutil_cp {
        input {
            String uri
            String docker
        }
        command <<<
            set -euxo pipefail
            mkdir __out/
            gsutil -q cp "~{uri}" __out/
        >>>
        output {
            File file = glob("__out/*")[0]
        }
        runtime {
            docker: docker
        }
    }
    """
    yield (  # pyre-ignore
        yield {"task_wdl": wdl, "inputs": {"uri": uri, "docker": cfg["download_gsutil"]["docker"]}}
    )

def gsutil_directory_downloader(
    cfg: config.Loader, logger: logging.Logger, uri: str, **kwargs
) -> Generator[Dict[str, Any], Dict[str, Any], None]:
    """
    Built-in downloader plugin for public gs:// URIs; registered by setup.cfg entry_points section

    TODO: adopt security credentials from runtime environment
    """
    wdl = r"""
    task gsutil_cp {
        input {
            String uri
            String docker
        }

        String dnm = basename(uri, "/")

        command <<<
            set -euxo pipefail
            mkdir __out/
            gsutil -q -m cp -r "~{uri}" __out/
        >>>
        output {
            Directory directory = "__out/" + dnm
        }
        runtime {
            docker: docker
        }
    }
    """
    yield (  # pyre-ignore
        yield {"task_wdl": wdl, "inputs": {"uri": uri, "docker": cfg["download_gsutil"]["docker"]}}
    )


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as an
import re

def get_timeblock(colvalue):
    '''Determines the time block (1,2,3,4) of the given timestamp value based on its hour. Each day is divided into four
    6-hour blocks. Block 1: (0-6hr), Block 2: (6-12hr), Block3: (12-18hr), Block4: (18-24hr)'''
    hour_ofthe_day = colvalue.hour #hour value of the date
    block = 0
    if hour_ofthe_day < 6: #assigning block value based on the hour value.
        block = 1
    elif 6 <= hour_ofthe_day < 12:
        block = 2
    elif 12 <= hour_ofthe_day < 18:
        block = 3
    else:
        block = 4
    return block

def get_activity_type(oldvalue):
    '''Determines the activity type based on the inout value and categorises it into one of these 6 values.
    {Create, Fork, Delete, Commit, PullRequest, Other}'''
    categories = ['Create', 'Fork', 'Delete', 'Commit', 'PullRequest', 'Other']
    expr = "^"+"|".join(categories) #regex to search for the category type.
    p = re.compile(expr)
    if p.search(oldvalue) is not None:
        newvalue = p.search(oldvalue).group(0) #assigning the matched category.
    else:
        newvalue = 'Other' #If not matched, assign 'Other'
    return newvalue

def get_country2(location):
    '''Determines the country of the location as USA or Other. If location has one of us states or cities,
    it is marked as USA, else it is markes as Other'''
    us_states = ["AL", "AK", "AZ", "AR", "CA", "CO", "CT", "DE", "FL", "GA",
                 "HI", "ID", "IL", "IN", "IA", "KS", "KY", "LA", "ME", "MD",
                 "MA", "MI", "MN", "MS", "MO", "MT", "NE", "NV", "NH", "NJ",
                 "NM", "NY", "NC", "ND", "OH", "OK", "OR", "PA", "RI", "SC",
                 "SD", "TN", "TX", "UT", "VT", "VA", "WA", "WV", "WI", "WY",
                 "Alabama", "Alaska", "Arizona", "Arkansas", "California", "Colorado",
                 "Connecticut", "Delaware", "Florida", "Georgia", "Hawaii", "Idaho", "Illinois",
                 "Indiana", "Iowa", "Kansas", "Kentucky", "Louisiana", "Maine", "Maryland",
                 "Massachusetts", "Michigan", "Minnesota", "Mississippi", "Missouri", "Montana",
                 "Nebraska", "Nevada", "New Hampshire", "New Jersey", "New Mexico", "New York",
                 "North Carolina", "North Dakota", "Ohio", "Oklahoma", "Oregon", "Pennsylvania",
                 "Rhode Island", "South Carolina", "South Dakota", "Tennessee", "Texas", "Utah",
                 "Vermont", "Virginia", "Washington", "West Virginia", "Wisconsin", "Wyoming"]
    us_cities = ["Chicago", "San Francisco", "Seattle"]
    expr = "|".join(us_states + us_cities) #regex to match us locations
    expr += "|USA|U.S."
    match = re.compile(expr)
    if match.search(location) is not None: #If location in us, assign 'USA'
        country = 'USA'
    else: #If no match found, assign 'Other'
        country = 'Other'
    return country

file_names = ['github_timeline_data000000000000.csv', 'github_timeline_data000000000001.csv',
              'github_timeline_data000000000002.csv', 'github_timeline_data000000000003.csv',
              'github_timeline_data000000000004.csv', 'github_timeline_data000000000005.csv',
              'github_timeline_data000000000006.csv', 'github_timeline_data000000000007.csv']
df_list = []
for file in file_names: #reading csv files
    df_temp = pd.read_csv(file, usecols=['repository_url', 'created_at', 'type', 'actor_attributes_location'])
    df_list.append(df_temp)
df = pd.concat(df_list, axis=0, ignore_index=True) #concatenated dataframe

df2 = df[['type', 'created_at', 'actor_attributes_location']]
df2['created_at'] = pd.to_datetime(df2['created_at']) #pandas timestamps
df2['timeblocks'] = df2['created_at'].apply(get_timeblock) #assigning timeblocks based on the timestamp
df3 = df2.groupby('timeblocks').count() #Counting # of events(/dev activities) for each timeblock
df3.rename(columns={'type': 'dev_activity'}, inplace=True)
#print(df3.head())

df3['dev_activity'].plot.bar(rot=0) #bar plot of development activities broken down by time blocks.
plt.title('3a) Development activity over time of the day')
plt.ylabel('# of activities')
plt.xticks([0, 1, 2, 3], ['12am-6am', '6am-12pm', '12pm-6pm', '6pm-12am'])
plt.show()

df2['activity_type'] = df2['type'].apply(get_activity_type) #categorizing activities into (Create, Delete, Fork, ...etc.,)
df4 = df2.pivot_table('created_at', index='timeblocks', columns='activity_type', aggfunc='count') # pivot table with
# timeblocks as index. For each of the timeblocks(1,2,3,4), the corr. row has # of events for different categories (Create, Delete, Fork, ...etc.,)
#print(df4)

df4[['Create', 'Fork', 'Delete', 'Commit', 'PullRequest', 'Other']].plot(kind='bar', stacked=True, rot=0) #stacked bar graph
# with different activity types.
plt.title('3b) Development activities over time of the day')
plt.ylabel('# of activities')
plt.xticks([0, 1, 2, 3], ['12am-6am', '6am-12pm', '12pm-6pm', '6pm-12am'])
plt.show()

df5 = df2.dropna() #dropping rows with no locations
df5['COUNTRY'] = df5['actor_attributes_location'].apply(get_country2) #assigning country
df5 = df5.pivot_table('created_at', index='timeblocks', columns='COUNTRY', aggfunc='count') #pivot table with timeblocks and countries.
# For each timeblock, the corr. row has # of events/activities for countries (USA, Other).
#print(df5)

df5[['USA', 'Other']].plot(kind='bar', stacked=True, rot=0) #stacked bar graph for dev activities analysed
# furthur by country (USA, Other)
plt.title('3c) Development activities in USA vs. rest of the world')
plt.ylabel('# of activities')
plt.xticks([0, 1, 2, 3], ['12am-6am', '6am-12pm', '12pm-6pm', '6pm-12am'])
plt.show()

daysoftheweek = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday']
df6 = df2[['type', 'created_at']]
df6['dayofweek'] = df6['created_at'].dt.dayofweek #get day of the week based on the timestamp
df6['weekofyear'] = df6['created_at'].dt.weekofyear #get week of the year based on the timestamp
df7 = df6.groupby('dayofweek', as_index=False)['type'].count() #Obtaining # of activities for each day of the week.
df7.sort_values('type', ascending=False, inplace=True)
most_active_day = df7.iloc[0, 0] #Most active day based on the highest # of activities.
print('Part 3d)')
print('Most active development day of the week in the time window:', daysoftheweek[most_active_day])

df6 = df6.pivot_table('type', index='weekofyear', columns='dayofweek',aggfunc='count').fillna(0) #pivot table with
# weekofyear and dayofweek. For each weekofyear, corr. row in the pivot table has # of activities for all days of the
# week (Mon, Tue, Wed, Thu, Fri, Sat, Sun)
max_value = df6.max().max() #maximum # of activities in the pivot table. Needed for y limit below.
#print(df6)

fig = plt.figure()
plt.title('3d) Animation of development activities over weeks')
plt.xticks([0, 1, 2, 3, 4, 5, 6], ['Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat', 'Sun'])
plt.xlabel('Day of the week')
plt.ylabel('# of activities')
plt.ylim(0, max_value+100)
xdata = np.arange(0, 7)
ydata = np.array(df6.iloc[0]) #first bar plot in the animation.
barplot = plt.bar(xdata, ydata, width=0.5, color="rmygbck")

def figinit():
    pass

def redraw(framenum):
    ydata = np.array(df6.iloc[framenum]) #next bar plot values.
    # Here, each row in df6 corresponds to a week in the year 2012.
    index = 0
    for bar in barplot:
        bar.set_height(ydata[index])
        index += 1
    return barplot,

anim = an.FuncAnimation(fig, redraw, frames=len(df6.index), interval=500, init_func=figinit, blit=False, repeat=False)
plt.show()


from functools import partial
from math import floor
from random import randrange
import pygame

import Engine.SceneManager as SceneManager
from Engine.Scenes.Fight.Fighter import Fighter
from Engine.Scenes.Fight.Options import Options

from Engine.Event import Event
from Engine.Tools.Message import Message
from Engine.Tools.Colors import black, type_colors
from Engine.SceneAbstract import SceneAbstract
from monster import fight

class Fight(SceneAbstract):
    def __init__(self):
        self.players = []
        self.screen_size = pygame.display.get_surface().get_size()

    def attack(self):
        print("activating attacks")
        self.active.deactivate()
        self.active = self.make_attacks()

    def turn_attacks(self, move):
        if self.challenger.speed > self.foe.speed:
            self.foe.defend(self.challenger.use(self.foe, move.id - 1))
            if not self.foe.is_alive():
                print("Turn end !\n")
                return
            self.challenger.defend(self.foe.use(self.challenger, int(floor(randrange(0, 3)))))
        else:
            self.challenger.defend(self.foe.use(self.challenger, int(floor(randrange(0, 3)))))
            if not self.challenger.is_alive():
                print("Turn end !\n")
                return
            self.foe.defend(self.challenger.use(self.foe, move.id - 1))

    def move(self, move):
        print(self.challenger.name, "uses", move.name)
        self.turn_attacks(move)
        self.active.deactivate()
        self.active = self.make_options()

    def heal(self):
        print(self.challenger.hp)

    def give_up(self):
        # SceneManager.SceneManager.get_instance().set_scene(3, data={"won": False})
        SceneManager.SceneManager.get_instance().set_scene(0)

    def fight_end(self):
        won = True
        if self.challenger.hp <= 0:
            won = False
        SceneManager.SceneManager.get_instance().set_scene(3, data={"won": won})

    def make_objects(self):
        objects = Options((3*self.screen_size[0]/5, 3*self.screen_size[1]/5), (self.screen_size[0]/3, self.screen_size[1]/3))
        return objects

    def make_options(self):
        options = Options((3*self.screen_size[0]/5, 3*self.screen_size[1]/5), (self.screen_size[0]/3, self.screen_size[1]/3))
        options.add_option("attack", (200,200,200), self.attack)
        options.add_option("heal", (200,200,200), self.heal)
        options.add_option("give up", (100,100,100), self.give_up)
        return options

    def make_attacks(self):
        attacks = Options((3*self.screen_size[0]/5, 3*self.screen_size[1]/5), (self.screen_size[0]/3, self.screen_size[1]/3))
        for move in self.challenger.moves:
            attack = partial(self.move, move)
            attacks.add_option(move.name, type_colors[move.type], attack)
        return attacks

    def setup(self, data):
        self.players = data["players"]
        print(self.players[0])
        self.challenger = Fighter(self.players[0], (self.screen_size[0]/10, 3*self.screen_size[1]/5), (self.screen_size[0]/3, self.screen_size[1]/3))
        self.foe = Fighter(self.players[1], (3*self.screen_size[0]/5, self.screen_size[1]/10), (self.screen_size[0]/3, self.screen_size[1]/3))
        self.message = Message((10, 10), (self.screen_size[0]/2, self.screen_size[1]/3))
        self.message.add_content("Make your move")
        self.message.next()
        self.active = self.make_options()
        Event.get_instance().debug()

    def update(self, screen):
        screen.fill(black)
        self.challenger.update()
        self.challenger.display(screen)
        self.foe.update()
        self.foe.display(screen)
        self.message.update(screen)
        self.active.update(screen)
        if self.challenger.hp < 0 or self.foe.hp < 0:
            self.fight_end()


###################
# Libs
###################

import torch
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter
import time
from argparse import ArgumentParser

import os
import datetime

from model import TSP_net
from utils import compute_tour_length

# visualization
from IPython.display import set_matplotlib_formats, clear_output
set_matplotlib_formats('png2x','pdf')
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
try:
    import networkx as nx
    from scipy.spatial.distance import pdist, squareform
    from concorde.tsp import TSPSolver # !pip install -e pyconcorde
except:
    pass
import warnings
warnings.filterwarnings("ignore", category=UserWarning)

###################
# Hyper-parameters
###################
embedding_choices = ['linear', 'conv_same_padding', 'conv', 'conv_linear']  # conv_XY, conv_same_padding, conv_linear
parser = ArgumentParser()
parser.add_argument('--debug', action='store_true', default=False)
parser.add_argument('--exp_name', type=str, required=True)
parser.add_argument('--exp_detail', type=str, default='')
parser.add_argument('-n', '--nb_nodes', type=int, choices=[20, 50, 100], required=True)
parser.add_argument('--bsz', type=int, required=True)
parser.add_argument('--gpu_id', type=str, required=True)

parser.add_argument('--embedding', type=str, choices=embedding_choices, required=True)
parser.add_argument('--nb_neighbors', type=int, default=None)  # For CNN
parser.add_argument('--kernel_size', type=int, default=None)  # For CNN

parser.add_argument('--segm_len', type=int, default=None)
parser.add_argument('--batchnorm', action='store_true', default=False)

parser.add_argument('--resume_training', action='store_true', default=False)
parser.add_argument('--resume_file', type=str, default=None)

parser.add_argument('--dim_emb', type=int, default=128)
parser.add_argument('--dim_ff', type=int, default=512)
parser.add_argument('--dim_input_nodes', type=int, default=2)
parser.add_argument('--nb_layers_encoder', type=int, default=6)
parser.add_argument('--nb_layers_decoder', type=int, default=2)
parser.add_argument('--nb_heads', type=int, default=8)
parser.add_argument('--nb_epochs', type=int, default=100)
parser.add_argument('--nb_batch_per_epoch', type=int, default=2500)
parser.add_argument('--nb_batch_eval', type=int, default=20)
parser.add_argument('--lr', type=int, default=1e-4)
parser.add_argument('--tol', type=int, default=1e-3)
parser.add_argument('--max_len_PE', type=int, default=1000)
args = parser.parse_args()
if args.embedding == 'conv':
    assert not args.nb_neighbors is None
    assert not args.kernel_size is None
    assert args.nb_neighbors == args.kernel_size - 1
if args.resume_training:
    assert not args.resume_file is None

print(args)
if not args.debug:
    if not os.path.exists('./logs'):
        os.makedirs('./logs')
    writer = SummaryWriter('./logs/'+args.exp_name)
else:
    print("########### RUNNING ON DEBUG MODE ############")

###################
# Hardware : CPU / GPU(s)
###################

device = torch.device("cpu"); gpu_id = -1 # select CPU

gpu_id = args.gpu_id
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
if torch.cuda.is_available():
    device = torch.device("cuda")
    print('GPU name: {:s}, gpu_id: {:s}'.format(torch.cuda.get_device_name(0),gpu_id))

# if args.debug:
#     device = torch.device("cpu"); gpu_id = -1 # select CPU
print(device)

###################
# Small test set for quick algorithm comparison
# Note : this can be removed
###################

save_1000tsp = True
save_1000tsp = False
if save_1000tsp:
    bsz = 1000
    x = torch.rand(bsz, args.nb_nodes, args.dim_input_nodes, device='cpu')
    print(x.size(),x[0])
    data_dir = os.path.join("data")
    if not os.path.exists(data_dir):
        os.makedirs(data_dir)
    if args.nb_nodes==20 : torch.save({ 'x': x, }, '{}.pkl'.format(data_dir + "/1000tsp20"))
    if args.nb_nodes==50 : torch.save({ 'x': x, }, '{}.pkl'.format(data_dir + "/1000tsp50"))
    if args.nb_nodes==100 : torch.save({ 'x': x, }, '{}.pkl'.format(data_dir + "/1000tsp100"))

checkpoint = None
if args.nb_nodes==20 : checkpoint = torch.load("data/1000tsp20.pkl")
if args.nb_nodes==50 : checkpoint = torch.load("data/1000tsp50.pkl")
if args.nb_nodes==100 : checkpoint = torch.load("data/1000tsp100.pkl")
if checkpoint is not None:
    x_1000tsp = checkpoint['x'].to(device)
    n = x_1000tsp.size(1)
    print('nb of nodes :',n)
else:
    x_1000tsp = torch.rand(1000, args.nb_nodes, args.dim_input_nodes, device='cpu')
    n = x_1000tsp.size(1)
    print('nb of nodes :',n)

###################
# Instantiate a training network and a baseline network
###################
try:
    del model_train # remove existing model
    del model_baseline # remove existing model
except:
    pass

model_train = TSP_net(args.embedding, args.nb_neighbors, args.kernel_size,
                      args.dim_input_nodes, args.dim_emb, args.dim_ff, args.nb_layers_encoder, args.nb_layers_decoder,
                      args.nb_heads, args.max_len_PE, args.segm_len, batchnorm=args.batchnorm)

model_baseline = TSP_net(args.embedding, args.nb_neighbors, args.kernel_size,
                      args.dim_input_nodes, args.dim_emb, args.dim_ff, args.nb_layers_encoder, args.nb_layers_decoder,
                      args.nb_heads, args.max_len_PE, args.segm_len, batchnorm=args.batchnorm)

# uncomment these lines if trained with multiple GPUs
print(torch.cuda.device_count())
if torch.cuda.device_count()>1:
    model_train = nn.DataParallel(model_train)
    model_baseline = nn.DataParallel(model_baseline)
# uncomment these lines if trained with multiple GPUs

optimizer = torch.optim.Adam( model_train.parameters() , lr = args.lr )

model_train = model_train.to(device)
model_baseline = model_baseline.to(device)
model_baseline.eval()

print(args); print('')

# Logs
if not args.debug:
    time_stamp=datetime.datetime.now().strftime("%y-%m-%d--%H-%M-%S")
    # file_name = 'logs'+'/'+time_stamp + "-n{}".format(args.nb_nodes) + "-gpu{}".format(args.gpu_id) + ".txt"
    file_name = f'logs/{args.exp_name}.txt'
    file = open(file_name,"w",1)
    file.write(time_stamp+'\n\n')
    file.write(args.exp_detail+'\n\n')
    for arg in vars(args):
        file.write(arg)
        hyper_param_val="={}".format(getattr(args, arg))
        file.write(hyper_param_val)
        file.write('\n')
    file.write('\n\n')
else:
    torch.autograd.set_detect_anomaly(True)
plot_performance_train = []
plot_performance_baseline = []
epoch_ckpt = 0
tot_time_ckpt = 0

# Uncomment these lines to re-start training with saved checkpoint
checkpoint_dir = 'checkpoint'
if args.resume_training:
    checkpoint_file = os.path.join(checkpoint_dir, args.resume_file)
    checkpoint = torch.load(checkpoint_file, map_location=device)
    epoch_ckpt = checkpoint['epoch'] + 1
    tot_time_ckpt = checkpoint['tot_time']
    plot_performance_train = checkpoint['plot_performance_train']
    plot_performance_baseline = checkpoint['plot_performance_baseline']
    model_baseline.load_state_dict(checkpoint['model_baseline'])
    model_train.load_state_dict(checkpoint['model_train'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    print('Re-start training with saved checkpoint file={:s}\n  Checkpoint at epoch= {:d} and time={:.3f}min\n'.format(checkpoint_file,epoch_ckpt-1,tot_time_ckpt/60))
    del checkpoint
# Uncomment these lines to re-start training with saved checkpoint

###################
# Main training loop
###################
start_training_time = time.time()

for epoch in range(0,args.nb_epochs):

    # re-start training with saved checkpoint
    epoch += epoch_ckpt
    if epoch == args.nb_epochs:
        break

    ###################
    # Train model for one epoch
    ###################
    start = time.time()
    model_train.train()

    for step in range(1,args.nb_batch_per_epoch+1):

        # generate a batch of random TSP instances
        x = torch.rand(args.bsz, args.nb_nodes, args.dim_input_nodes, device=device) # size(x)=(bsz, nb_nodes, dim_input_nodes)

        # compute tours for model
        tour_train, sumLogProbOfActions = model_train(x, deterministic=False) # size(tour_train)=(bsz, nb_nodes), size(sumLogProbOfActions)=(bsz)

        # compute tours for baseline
        with torch.no_grad():
            tour_baseline, _ = model_baseline(x, deterministic=True)

        # get the lengths of the tours
        L_train = compute_tour_length(x, tour_train) # size(L_train)=(bsz)
        L_baseline = compute_tour_length(x, tour_baseline) # size(L_baseline)=(bsz)

        # backprop
        loss = torch.mean( (L_train - L_baseline)* sumLogProbOfActions )
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    time_one_epoch = time.time()-start
    time_tot = time.time()-start_training_time + tot_time_ckpt

    ###################
    # Evaluate train model and baseline on 10k random TSP instances
    ###################
    model_train.eval()
    mean_tour_length_train = 0
    mean_tour_length_baseline = 0
    for step in range(0,args.nb_batch_eval):

        # generate a batch of random tsp instances
        x = torch.rand(args.bsz, args.nb_nodes, args.dim_input_nodes, device=device)

        # compute tour for model and baseline
        with torch.no_grad():
            tour_train, _ = model_train(x, deterministic=True)
            tour_baseline, _ = model_baseline(x, deterministic=True)

        # get the lengths of the tours
        L_train = compute_tour_length(x, tour_train)
        L_baseline = compute_tour_length(x, tour_baseline)

        # L_tr and L_bl are tensors of shape (bsz,). Compute the mean tour length
        mean_tour_length_train += L_train.mean().item()
        mean_tour_length_baseline += L_baseline.mean().item()

    mean_tour_length_train =  mean_tour_length_train/ args.nb_batch_eval
    mean_tour_length_baseline =  mean_tour_length_baseline/ args.nb_batch_eval

    # evaluate train model and baseline and update if train model is better
    update_baseline = mean_tour_length_train+args.tol < mean_tour_length_baseline
    if update_baseline:
        model_baseline.load_state_dict( model_train.state_dict() )

    # Compute TSPs for small test set
    # Note : this can be removed
    with torch.no_grad():
        tour_baseline, _ = model_baseline(x_1000tsp, deterministic=True)
    mean_tour_length_test = compute_tour_length(x_1000tsp, tour_baseline).mean().item()

    # For checkpoint
    plot_performance_train.append([ (epoch+1), mean_tour_length_train])
    plot_performance_baseline.append([ (epoch+1), mean_tour_length_baseline])

    # Compute optimality gap
    if args.nb_nodes==50: gap_train = mean_tour_length_train/5.692- 1.0
    elif args.nb_nodes==100: gap_train = mean_tour_length_train/7.765- 1.0
    else: gap_train = -1.0

    # Print and save in txt file
    mystring_min = 'Epoch: {:d}, epoch time: {:.3f}min, tot time: {:.3f}day, L_train: {:.3f}, L_base: {:.3f}, L_test: {:.3f}, gap_train(%): {:.3f}, update: {}'.format(
        epoch, time_one_epoch/60, time_tot/86400, mean_tour_length_train, mean_tour_length_baseline, mean_tour_length_test, 100*gap_train, update_baseline)
    print(mystring_min) # Comment if plot display
    if not args.debug:
        writer.add_scalar('mean_tour_len_train', mean_tour_length_train, epoch)
        writer.add_scalar('mean_tour_len_baseline', mean_tour_length_baseline, epoch)

        file.write(mystring_min+'\n')

        # Saving checkpoint
        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
        filename = f'{checkpoint_dir}/{args.exp_name}.pkl'
        if (epoch+1) % 10 == 0:
            filename = f'{checkpoint_dir}/{args.exp_name}_e{epoch}.pkl'
        torch.save({
            'epoch': epoch,
            'time': time_one_epoch,
            'tot_time': time_tot,
            'loss': loss.item(),
            'TSP_length': [torch.mean(L_train).item(), torch.mean(L_baseline).item(), mean_tour_length_test],
            'plot_performance_train': plot_performance_train,
            'plot_performance_baseline': plot_performance_baseline,
            'mean_tour_length_test': mean_tour_length_test,
            'model_baseline': model_baseline.state_dict(),
            'model_train': model_train.state_dict(),
            'whole_model': model_baseline,
            'optimizer': optimizer.state_dict(),
            }, filename)

if not args.debug:
    writer.close()