from datetime import datetime
from enum import Enum
from pathlib import Path
import re
from sys import stderr
from time import sleep
import typing

class ReleaseTypes(Enum):
    MAJOR = 0
    MINOR = 1
    PATCH = 2

valid_release_types = typing.Literal["major", "minor", "patch"]

class ReleaseStrings:
    """An easy way to pass the various flavours of release string between functions."""

    def __init__(self, input_tag: str):
        version_mask = r"v\d+\.\d+\.\d+\D*.*"
        regex_101 = "https://regex101.com/r/dLVaNH/1"
        if re.fullmatch(version_mask, input_tag) is None:
            message = (
                "Release tag does not match the input mask:\n"
                f"{version_mask}\n"
                f"({regex_101})"
            )
            raise ValueError(message)
        else:
            self.tag = input_tag  # v1.2.3rc0

        self.series = ".".join(self.tag.split(".")[:2])  # v1.2
        self.branch = self.series + ".x"  # v1.2.x
        self.release = self.tag[1:]  # 1.2.3rc0

class WhatsNewRsts:
    """An easy way to pass the paths of various What's New files between functions."""

    def __init__(self, release_strings: ReleaseStrings):
        src_dir = Path(__file__).parents[1] / "docs" / "src"
        whatsnew_dir = src_dir / "whatsnew"
        assert whatsnew_dir.is_dir()

        self.latest = whatsnew_dir / "latest.rst"
        self.release = whatsnew_dir / (release_strings.series[1:] + ".rst")
        self.index = whatsnew_dir / "index.rst"
        self.template = self.latest.with_suffix(".rst.template")

def _break_print(message: str):
    print()
    print(message)
    # Help with flow/visibility by waiting 1secs before proceeding.
    sleep(1)

def _mark_section(section_number: int):
    _break_print(f"SECTION {section_number} ...")

def _get_input(message: str, expected_inputs: str) -> str:
    _break_print(message)
    return input(expected_inputs + " : ")

def _wait_for_done(message: str):
    _break_print(message)
    done = False
    while not done:
        done = (
            input("Step complete? y / [n] : ").casefold() == "y".casefold()
        )

def _report_problem(message: str):
    print(message, file=stderr)
    # To ensure correct sequencing of messages.
    sleep(0.5)

def get_release_type() -> ReleaseTypes:
    release_type = None
    release_types_str = " ".join(
        [f"{m.name}={m.value}" for m in ReleaseTypes.__members__.values()]
    )
    message = "What type of release are you preparing?\nhttps://semver.org/"
    while release_type is None:
        input_type = _get_input(message, release_types_str)
        try:
            release_type = ReleaseTypes(int(input_type))
        except ValueError:
            _report_problem("Invalid release type. Please try again ...")
    _break_print(f"{release_type} confirmed.")
    return release_type

def get_release_tag() -> ReleaseStrings:
    # TODO: automate using setuptools_scm.
    release_strings = None
    message = (
        "Input the release tag you are creating today, including any release "
        "candidate suffix.\n"
        "https://semver.org/\n"
        "https://scitools-iris.readthedocs.io/en/latest/developers_guide/release.html?highlight=candidate#release-candidate"
    )
    while release_strings is None:
        input_tag = _get_input(message, "e.g. v1.2.3rc0")
        try:
            release_strings = ReleaseStrings(input_tag)
        except ValueError as err:
            _report_problem(str(err))
    return release_strings

def check_release_candidate(
    release_type: ReleaseTypes, release_strings: ReleaseStrings
) -> bool:
    is_release_candidate = "rc" in release_strings.tag

    message = "Checking tag for release candidate: "
    if is_release_candidate:
        message += "DETECTED\nThis IS a release candidate."
    else:
        message += "NOT DETECTED\nThis IS NOT a release candidate."
    _break_print(message)

    if release_type == ReleaseTypes.PATCH and is_release_candidate:
        message = (
            "Release candidates are not expected for PATCH releases. "
            "Are you sure you want to continue?"
        )
        if _get_input(message, "y / [n]").casefold() != "y".casefold():
            exit()
    return is_release_candidate

def check_first_in_series(
    release_type: ReleaseTypes,
    release_strings: ReleaseStrings,
    is_release_candidate: bool,
) -> bool:
    first_in_series = False
    if release_type != ReleaseTypes.PATCH:
        message = (
            "Have there been any prior releases in the "
            f"{release_strings.series} series, including release candidates?"
        )
        first_in_series = (
            _get_input(message, "[y] / n").casefold() == "n".casefold()
        )
        if first_in_series:
            _break_print("First in series confirmed.")
            if not is_release_candidate:
                message = (
                    "The first release in a series is expected to be a "
                    "release candidate, but this is not. Are you sure you "
                    "want to continue?"
                )
                if _get_input(message, "y / [n]").casefold() != "y".casefold():
                    exit()
        else:
            _break_print("Existing series confirmed.")
    return first_in_series

def update_standard_names(first_in_series: bool) -> None:
    if first_in_series:
        message = (
            "Update the file ``etc/cf-standard-name-table.xml`` to the latest CF "
            "standard names, via a new Pull Request.\n"
            "(This is used during build to automatically generate the sourcefile "
            "``lib/iris/std_names.py``).\n"
            "Latest standard names:\n"
            'wget "http://cfconventions.org/Data/cf-standard-names/current/src/cf-standard-name-table.xml";'
        )
        _wait_for_done(message)

def check_deprecations(release_type: ReleaseTypes) -> None:
    if release_type == ReleaseTypes.MAJOR:
        message = (
            "This is a MAJOR release - be sure to finalise all deprecations "
            "and FUTUREs from previous releases, via a new Pull Request.\n"
            "https://scitools-iris.readthedocs.io/en/latest/developers_guide/contributing_deprecations.html"
        )
        _wait_for_done(message)

def _delete_local_branch(branch_name: str):
    message = (
        "Before the next step, avoid a name clash by deleting any "
        "existing local branch, if one exists.\n"
        f"git branch -D {branch_name};\n"
        f"git push -d origin {branch_name};"
    )
    _wait_for_done(message)

def create_release_branch(
    release_strings: ReleaseStrings, first_in_series: bool
) -> None:
    # TODO: automate

    _break_print("Release branch management ...")

    if first_in_series:
        message = (
            "Visit https://github.com/SciTools/iris and create the"
            f"``{release_strings.branch}`` release branch from ``main``."
        )
        _wait_for_done(message)

    else:
        message = (
            "Cherry-pick any specific commits that are needed from ``main`` "
            f"onto {release_strings.branch} , to get the CI passing.\n"
            "E.g. a new dependency pin may have been introduced since "
            f"{release_strings.branch} was last updated from ``main``.\n"
            "DO NOT squash-merge - want to preserve the original commit SHA's."
        )
        _wait_for_done(message)

def finalise_whats_new(
    release_type: ReleaseTypes,
    release_strings: ReleaseStrings,
    is_release_candidate: bool,
    first_in_series: bool,
) -> WhatsNewRsts:
    _break_print("What's New finalisation ...")

    working_branch = release_strings.branch + ".updates"
    _delete_local_branch(working_branch)
    message = (
        f"Checkout a local branch from the official {release_strings.branch} branch.\n"
        "git fetch upstream;\n"
        f"git checkout upstream/{release_strings.branch} -b "
        f"{working_branch};"
    )
    _wait_for_done(message)

    rsts = WhatsNewRsts(release_strings)

    # TODO: automate
    if first_in_series:
        message = (
            "'Cut' the What's New for the release.\n"
            f"git mv {rsts.latest.absolute()} {rsts.release.absolute()};"
        )
        _wait_for_done(message)

        message = (
            f"In {rsts.index.absolute()}:\n"
            f"Replace references to {rsts.latest.name} with {rsts.release.name}"
        )
        _wait_for_done(message)

    _break_print(f"What's New file path = {rsts.release}")

    if not release_type == ReleaseTypes.PATCH:
        whatsnew_title = f"{release_strings.series} ({datetime.today().strftime('%d %b %Y')})"
        if is_release_candidate:
            whatsnew_title += " [release candidate]"
        # TODO: automate
        message = f"In {rsts.release.name}: set the page title to:\n{whatsnew_title}\n"
        if not is_release_candidate:
            message += (
                "\nBe sure to remove any existing mentions of release "
                "candidate from the title.\n"
            )
        _wait_for_done(message)

        message = (
            f"In {rsts.release.name}: ensure the page title underline is "
            "the exact same length as the page title text."
        )
        _wait_for_done(message)

        dropdown_title = f"\n{release_strings.series} Release Highlights\n"
        message = (
            f"In {rsts.release.name}: set the sphinx-design dropdown title to:{dropdown_title}"
        )
        _wait_for_done(message)

        message = (
            f"Review {rsts.release.name} to ensure it is a good reflection of "
            f"what is new in {release_strings.series}."
        )
        _wait_for_done(message)

        message = (
            "Work with the development team to populate the Release "
            f"Highlights dropdown section at the top of {rsts.release.name}."
        )
        _wait_for_done(message)

    else:
        message = (
            "Create a patch dropdown section at the top of "
            f"{rsts.release.name}.\n"
            f"See {rsts.template} for how this should be written."
        )
        _wait_for_done(message)

    if first_in_series:
        # TODO: automate
        message = (
            "Remove the What's New template file.\n"
            f"git rm {rsts.template.absolute()};"
        )
        _wait_for_done(message)

    message = (
        "Commit and push all the What's New changes.\n"
        f'git commit -am "Whats new updates for {release_strings.tag} .";\n'
        f"git push -u origin {working_branch};"
    )
    _wait_for_done(message)

    message = (
        f"Follow the Pull Request process to get {working_branch} "
        f"merged into upstream/{release_strings.branch} .\n"
        "Make sure the documentation is previewed during this process."
    )
    _wait_for_done(message)

    return rsts

def cut_release(
    release_strings: ReleaseStrings, is_release_candidate: bool
) -> None:
    _break_print("The release ...")

    message = (
        "Visit https://github.com/SciTools/iris/releases/new to open "
        "a blank new-release web page."
    )
    _wait_for_done(message)

    message = (
        f"Select {release_strings.branch} as the Target.\n"
        f"Input {release_strings.tag} as the new tag to create, and also as "
        "the Release title."
    )
    _wait_for_done(message)

    message = (
        "Craft an appropriate release description in the main text box.\n"
        "Be sure to mention the What's New entry, conda-forge and PyPI - you "
        "will need to return later to make these into links.\n"
        "Be careful to change the appropriate words if copying from a "
        "previous release description."
    )
    _wait_for_done(message)

    if is_release_candidate:
        message = (
            "This is a release candidate - include the following instructions "
            "for installing with conda or pip:\n"
            f"conda install -c conda-forge/label/rc_iris iris={release_strings.release}\n"
            f"pip install scitools-iris=={release_strings.release}"
        )
        _wait_for_done(message)

        message = (
            "This is a release candidate - tick the box to set this as a "
            "pre-release."
        )
        _wait_for_done(message)

    else:
        message = "Tick the box to set this as the latest release."
        _wait_for_done(message)

    message = "Click: Publish release !"
    _wait_for_done(message)

    message = (
        "The CI will now run against this new tag, including automatically "
        "publishing to PyPI."
    )
    _break_print(message)

    url = "https://github.com/SciTools/iris/actions/workflows/ci-wheels.yml"
    message = (
        f"Visit {url} to monitor the building, testing and publishing of "
        "the Iris sdist and binary wheel to PyPI."
    )
    _wait_for_done(message)

def check_rtd(
    release_strings: ReleaseStrings, is_release_candidate: bool
) -> None:
    _break_print("Read the Docs checks ...")

    message = (
        "Visit https://readthedocs.org/projects/scitools-iris/versions/ and "
        "make sure you are logged in."
    )
    _wait_for_done(message)

    message = f"Set {release_strings.tag} to Active, un-Hidden."
    _wait_for_done(message)

    message = f"Set {release_strings.branch} to Active, Hidden."
    _wait_for_done(message)

    message = (
        "Keep the latest 2 branch doc builds active - those formatted 0.0.x - "
        "deactivate older ones."
    )
    _wait_for_done(message)

    message = (
        f"Visit https://scitools-iris.readthedocs.io/en/{release_strings.tag} "
        "to confirm:\n\n"
        "- The docs have rendered.\n"
        f"- The version badge in the top left reads: {release_strings.tag} .\n"
        "   (this demonstrates that setuptools_scm has worked correctly).\n"
        "- The What's New looks correct.\n"
        f"- {release_strings.tag} is available in RTD's version switcher.\n\n"
        "NOTE: the docs can take several minutes to finish building."
    )
    if not is_release_candidate:
        message += (
            "- Selecting 'stable' in the version switcher also brings up the "
            f"{release_strings.tag} render."
        )
    _wait_for_done(message)

    message = (
        f"Visit https://scitools-iris.readthedocs.io/en/{release_strings.branch} "
        "to confirm:\n\n"
        "- The docs have rendered\n"
        f"- The version badge in the top left includes: {release_strings.branch} .\n"
        f"- {release_strings.branch} is NOT available in RTD's version switcher.\n\n"
        "NOTE: the docs can take several minutes to finish building."
    )
    _wait_for_done(message)

def check_pypi(
    release_strings: ReleaseStrings, is_release_candidate: bool
) -> str:
    _break_print("PyPI checks ...")
    _break_print("If anything goes wrong, manual steps are in the documentation.")

    message = (
        "Confirm that the following URL is correctly populated:\n"
        f"https://pypi.org/project/scitools-iris/{release_strings.release}/"
    )
    _wait_for_done(message)

    message = (
        f"Confirm that {release_strings.release} is at the top of this page:\n"
        "https://pypi.org/project/scitools-iris/#history"
    )
    _wait_for_done(message)

    if is_release_candidate:
        message = (
            f"Confirm that {release_strings.release} is marked as a "
            f"pre-release on this page:\n"
            "https://pypi.org/project/scitools-iris/#history"
        )
    else:
        message = (
            f"Confirm that {release_strings.release} is the tag shown on the "
            "scitools-iris PyPI homepage:\n"
            "https://pypi.org/project/scitools-iris/"
        )
    _wait_for_done(message)

    message = (
        f"Visit the below and click `view hashes` for the Source Distribution"
        f"(`.tar.gz`):\n"
        f"https://pypi.org/project/scitools-iris/{release_strings.release}#files\n"
    )
    sha256 = _get_input(message, "Input the SHA256 hash")

    message = (
        "Confirm that pip install works as expected:\n"
        f"pip install scitools-iris=={release_strings.release};"
    )
    _wait_for_done(message)

    return sha256

def update_conda_forge(
    release_strings: ReleaseStrings, is_release_candidate: bool, sha256: str
) -> None:
    _break_print("conda-forge updates ...")

    if not is_release_candidate:
        message = (
            "NOTE: after several hours conda-forge automation will create a "
            "Pull Request against conda-forge/iris-feedstock (via the "
            "regro-cf-autotick-bot). Quicker to sort it now, manually ..."
        )
        _break_print(message)

    message = (
        "Make sure you have a GitHub fork of:\n"
        "https://github.com/conda-forge/iris-feedstock"
    )
    _wait_for_done(message)

    message = (
        "Make sure you have a local clone of your iris-feedstock fork.\n"
        "`cd` into your clone."
    )
    _wait_for_done(message)

    if is_release_candidate:
        upstream_branch = "release-candidate"
    else:
        upstream_branch = "main"

    # TODO: automate
    message = (
        "Checkout a new branch for the conda-forge changes for this release:\n"
        "git fetch upstream;\n"
        f"git checkout upstream/{upstream_branch} -b {release_strings.tag};\n"
    )
    _wait_for_done(message)

    message = (
        "Update ./recipe/meta.yaml:\n\n"
        f"- The version at the very top of the file: {release_strings.release}\n"
        f"- The sha256 hash: {sha256}\n"
        "- Requirements: align the packages and pins with those in the Iris repo\n"
        "- Maintainers: update with any changes to the dev team\n"
        "- MAKE SURE everything else is correct - plenty of other things "
        "might need one-off changes.\n"
    )
    _wait_for_done(message)

    # TODO: automate
    message = (
        "No other file normally needs changing in iris-feedstock, so push up "
        "the changes to prepare for a Pull Request:\n"
        f'git commit -am "Recipe updates for {release_strings.tag} .";\n'
        f"git push -u origin {release_strings.tag};"
    )
    _wait_for_done(message)

    message = (
        f"Follow the Pull Request process to get {release_strings.tag} branch "
        f"merged into upstream/{upstream_branch} .\n"
        "Specific conda-forge guidance will be automatically given once the "
        "PR is created."
    )
    _wait_for_done(message)

    message = (
        f"Confirm that {release_strings.release} appears in this list:\n"
        "https://anaconda.org/conda-forge/iris/files"
    )
    _wait_for_done(message)

    if not is_release_candidate:
        message = (
            f"Confirm that {release_strings.release} is displayed on this "
            "page as the latest available:\n"
            "https://anaconda.org/conda-forge/iris"
        )
        _wait_for_done(message)

    if is_release_candidate:
        channel_command = " -c conda-forge/label/rc_iris "
    else:
        channel_command = " "
    message = (
        "Confirm that conda (or mamba) install works as expected:\n"
        f"conda create -n tmp_iris{channel_command}iris={release_strings.release};\n"
        f"conda remove -n tmp_iris --all;"
    )
    _wait_for_done(message)

def update_links(release_strings: ReleaseStrings) -> None:
    _break_print("Link updates ...")

    message = (
        "Revisit the GitHub release:\n"
        f"https://github.com/SciTools/iris/releases/tag/{release_strings.tag}\n"
        "You have confirmed that Read the Docs, PyPI and conda-forge have all "
        "updated correctly. Include the following links in the release "
        "notes:\n\n"
        f"https://scitools-iris.readthedocs.io/en/{release_strings.tag}/\n"
        f"https://pypi.org/project/scitools-iris/{release_strings.release}/\n"
        f"https://anaconda.org/conda-forge/iris?version={release_strings.release}\n"
    )
    _wait_for_done(message)

    message = (
        "Update the release page in GitHub discussions, with the above links "
        "and anything else appropriate.\n"
        "https://github.com/SciTools/iris/discussions"
    )
    _wait_for_done(message)

def twitter_announce(
    release_strings: ReleaseStrings, first_in_series: bool
) -> None:
    message = (
        "Announce the release via https://twitter.com/scitools_iris, and any "
        "other appropriate message boards (e.g. Yammer).\n"
        "Any content used for the announcement should be stored in the "
        "SciTools/twitter-scitools-iris GitHub repo.\n"
    )
    if not first_in_series:
        message += (
            f"Consider replying within an existing {release_strings.series} "
            "announcement thread, if appropriate."
        )
    _wait_for_done(message)

def update_citation(
    release_strings: ReleaseStrings, is_release_candidate: bool
) -> None:
    if not is_release_candidate:
        src_dir = Path(__file__).parents[1] / "docs" / "src"
        citation_rst = src_dir / "userguide" / "citation.rst"
        assert citation_rst.is_file()
        message = (
            f"Follow the Pull Request process to update {citation_rst.name} "
            "with the correct dates, DOI and version string for "
            f"{release_strings.tag}.\n"
            f"{citation_rst.absolute()}\n\n"
            f"The PR should target {release_strings.branch} (prior to merge-back)."
        )
        _wait_for_done(message)

def merge_back(
    release_strings: ReleaseStrings, first_in_series: bool, rsts: WhatsNewRsts
) -> None:
    _break_print("Branch merge-back ...")

    merge_commit = (
        "BE SURE TO MERGE VIA A MERGE-COMMIT (not a squash-commit), to "
        "preserve the commit SHA's."
    )

    if first_in_series:
        # TODO: automate

        working_branch = release_strings.branch + ".mergeback"
        _delete_local_branch(working_branch)
        message = (
            "Checkout a local branch from the official ``main`` branch.\n"
            "git fetch upstream;\n"
            f"git checkout upstream/main -b {working_branch};"
        )
        _wait_for_done(message)

        message = (
            f"Merge in the commits from {release_strings.branch}.\n"
            f"{merge_commit}\n"
            f"git merge upstream/{release_strings.branch} --no-ff "
            '-m "Merging release branch into main";'
        )
        _wait_for_done(message)

        message = (
            "Recreate the following files, which are present in ``main``, but "
            f"are currently deleted from {working_branch}:\n"
            f"{rsts.latest.absolute()}\n"
            f"{rsts.template.absolute()}\n"
            "THEN:\n"
            f"git add {rsts.latest.absolute()};\n"
            f"git add {rsts.template.absolute()};\n"
        )
        _wait_for_done(message)

        message = (
            f"In {rsts.index.absolute()}:\n"
            f"Add {rsts.latest.name} to the top of the list of .rst files, "
            f"and set the top include:: to be {rsts.latest.name} ."
        )
        _wait_for_done(message)

        message = (
            "Commit and push all the What's New changes.\n"
            'git commit -am "Restore latest Whats New files.";\n'
            f"git push -u origin {working_branch};"
        )
        _wait_for_done(message)

        message = (
            "Follow the Pull Request process to get "
            f"{working_branch} merged into upstream/main .\n"
            "Make sure the documentation is previewed during this process.\n"
            f"{merge_commit}"
        )
        _wait_for_done(message)

    else:
        message = (
            f"Propose a merge-back from {release_strings.branch} into ``main`` by "
            f"visiting this URL and clicking `Create pull request`:\n"
            f"https://github.com/SciTools/iris/compare/main...{release_strings.branch}\n"
            f"{merge_commit}"
        )
        _wait_for_done(message)
        message = (
            f"Once the pull request is merged ensure that the {release_strings.branch} "
            "release branch is restored.\n"
            "GitHub automation rules may have automatically deleted the release branch."
        )
        _wait_for_done(message)

def main():
    _mark_section(1)
    release_type = get_release_type()

    _mark_section(2)
    release_strings = get_release_tag()

    _mark_section(3)
    is_release_candidate = check_release_candidate(
        release_type,
        release_strings,
    )

    _mark_section(4)
    is_first_in_series = check_first_in_series(
        release_type,
        release_strings,
        is_release_candidate,
    )

    _mark_section(5)
    update_standard_names(
        is_first_in_series,
    )

    _mark_section(6)
    check_deprecations(
        release_type,
    )

    _mark_section(7)
    create_release_branch(
        release_strings,
        is_first_in_series,
    )

    _mark_section(8)
    whats_new_rsts = finalise_whats_new(
        release_type,
        release_strings,
        is_release_candidate,
        is_first_in_series,
    )

    _mark_section(9)
    cut_release(
        release_strings,
        is_release_candidate,
    )

    _mark_section(10)
    check_rtd(
        release_strings,
        is_release_candidate,
    )

    _mark_section(11)
    sha256 = check_pypi(
        release_strings,
        is_release_candidate,
    )

    _mark_section(12)
    update_conda_forge(
        release_strings,
        is_release_candidate,
        sha256,
    )

    _mark_section(13)
    update_links(
        release_strings,
    )

    _mark_section(14)
    twitter_announce(
        release_strings,
        is_first_in_series,
    )

    _mark_section(15)
    update_citation(
        release_strings,
        is_release_candidate,
    )

    _mark_section(16)
    merge_back(
        release_strings,
        is_first_in_series,
        whats_new_rsts,
    )

    _break_print("RELEASE COMPLETE. Congratulations! ðŸŽ‰")

if __name__ == "__main__":
    main()


from django.db import models
from django.contrib.auth.models import AbstractUser
from datetime import date

class CustomUser(AbstractUser):
    OFFICE_MANAGER = 1
    EMPLOYEE = 2

    ROLE_CHOICES = (
        (OFFICE_MANAGER, 'Office manager'),
        (EMPLOYEE, 'Employee'),
    )

    role = models.PositiveSmallIntegerField(choices=ROLE_CHOICES,
                                            blank=True, null=True)

    class Meta(AbstractUser.Meta):
        swappable = 'AUTH_USER_MODEL'

    def __str__(self):
        return self.username

class Room(models.Model):
    name = models.CharField(max_length=100)
    chairs_number = models.IntegerField()
    projector = models.BooleanField()
    marker_board = models.BooleanField()

    def __str__(self):
        return self.name

class Booking(models.Model):
    room = models.ForeignKey(Room, on_delete=models.CASCADE)
    user = models.ForeignKey(CustomUser, on_delete=models.CASCADE)
    date_from = models.DateField(default=date.today)
    time_from = models.TimeField(default="00:00")
    date_to = models.DateField(default=date.today,)
    time_to = models.TimeField(default="00:00")


from sys import stdin

def A():
    t = int(input())
    while t:
        t -= 1
        s = input()
        s1 = input()
        dp = [0] * 26
        for i in range(0, len(s)):
            dp[ord(s[i]) - ord('a')] = i
        ans = 0
        for i in range(0, len(s1) - 1):
            ans += abs(dp[ord(s1[i+1]) - ord('a')] - dp[ord(s1[i]) - ord('a')])
        print(ans)
A()
def B():
    t = int(input())
    while t:
        t -= 1
        x, n = [int(x) for x in stdin.readline().split()]
        k = n // 4 * 4 + 1
        for i in range(k, n+1):
            if x & 1:
                x += i
            else:
                x -= i
        print(x)

def C():
    t = int(input())
    while t:
        t -= 1
        n = int(input())
        x = [int(x) for x in stdin.readline().split()]
        x.sort()
        ans = x[0]
        for i in range(0, n-1):
            ans = max(ans, x[i+1] - x[i])
        print(ans)

def D():
    t = int(input())
    while t:
        t -= 1
        n = int(input())
        x = [int(x) for x in stdin.readline().split()]
        s = input()
        l, r = [], []
        for i in range(0, n):
            if s[i] == 'B':
                l.append(x[i])
            else:
                r.append(x[i])
        l.sort()
        r.sort()
        r.reverse()
        ans = True
        for i in range(0, len(l)):
            if l[i] < i + 1:
                ans = False
        for i in range(0, len(r)):
            if r[i] > n - i:
                ans = False
        if ans:
            print("yes")
        else:
            print("no")
D()

CP problems/Codeforces/Codeforces Round #753 (Div. 3).py


import sys
import string
import numpy as np

from tabulate import tabulate

def computePrecisionRecall(f, header):

    results = np.zeros((10, 10))
    table_header = [str(x) for x in range(0, 9+1)]
    table_header.insert(0, ' ')
    #print(table_header)
    #print(results)
    # old stuff
    print(header)
    #print(tabulate(results, tablefmt="latex", floatfmt=".2f"))
    true_positive = dict()
    false_positive = dict()

    total_predicted = dict()
    total_actual = dict()

    for i in range(0, 9+1):
        true_positive[i] = 0
        false_positive[i] = 0
        total_predicted[i] = 0
        total_actual[i] = 0

    found_prediction = False
    for line in f:
        pieces = line.rstrip('\n').split(' ')
        if "Prediction:" == pieces[0]:

            # needed to compute precision and recall
            found_prediction = True
            pred = int(pieces[1])
            actual = int(pieces[3])
            total_predicted[pred] += 1
            total_actual[actual] += 1
            if pred == actual:
                true_positive[pred] += 1
            else:
                false_positive[pred] += 1
            #new stuff
            results[pred][actual] += 1
        else:
            if found_prediction:
                break

    pr_header = ["Digit", "Precision", "Recall"]
    pr_data = []
    for i in range(0, 9+1):
        # compute precision and recall for one digit at a time
        recall = true_positive[i] / float(total_actual[i])
        precision = true_positive[i] / float(total_predicted[i])

        entry = [str(i), precision, recall]
        pr_data.append(entry)
        #print("Digit: " + str(i) + " Recall: " + str(recall) + " Precision: " + str(precision))
    print("Result Counts")
    proc_results = results.tolist()
    for i in range(0, 9+1):
        proc_results[i].insert(0, str(i))
    print(tabulate(proc_results, headers=table_header, tablefmt="latex"))
    print("Precision and Recall Values")
    print(tabulate(pr_data, headers=pr_header, tablefmt="latex"))

if __name__ == "__main__":
    with open(sys.argv[1]) as f:
        computePrecisionRecall(f, "Without k-d tree")
        computePrecisionRecall(f, "With k-d tree")


#!/usr/bin/python3

from pprint import pprint
from math import atan2, pi

asteroids = set()
with open("inputs/input10.txt", "r") as file:
    y = 0
    for line in (line.rstrip() for line in file):
        x = 0
        for char in line:
            if char != '.':
                asteroids.add( (x,y) )
            x += 1
        y += 1

angles = dict()
distances = dict()
for a in asteroids:
    for b in asteroids:
        if a==b:
            continue
        if a not in angles:
            angles[a] = dict()
        if a not in distances:
            distances[a] = dict()
        distances[a][b] = ( (a[0]-b[0])**2 + (a[1]-b[1])**2 ) ** 0.5
        angle = atan2(b[0]-a[0], a[1]-b[1])
        if angle < 0:
            angle += 2*pi
        angles[a][b] = angle

asteroids_by_angle = dict()
for a in asteroids:
    aa = angles[a]
    aba = dict()
    for b in aa.keys():
        int_angle = int(aa[b] * 180 / pi) % 360
        if int_angle not in aba:
            aba[int_angle] = set()
        aba[int_angle].add(b)
    asteroids_by_angle[a] = aba

def visible_from(a):
    visible = set()
    for b in asteroids:
        if a == b:
            continue
        b_angle = angles[a][b]
        b_distance = distances[a][b]
        int_ang = int(b_angle * 180 / pi) % 360
        possible_blockers = set()
        for ang in [int_ang, int_ang-1 % 360, int_ang+1 % 360]:
            if ang in asteroids_by_angle[a]:
                possible_blockers.update(asteroids_by_angle[a][ang])
        for c in possible_blockers:
            if a == c or b == c:
                continue
            if distances[a][c] < b_distance and abs(angles[a][c] - b_angle) < 0.00000001:
                break
        else:
            visible.add(b)
    return visible

max_visible = 0
best_location = None
for a in asteroids:
    visible = len(visible_from(a))
    if visible > max_visible:
        max_visible = visible
        best_location = a

print("Part 1:", best_location, max_visible)

asteroids.remove(best_location)
laser_angle = 0
count = 0
angles_bl = angles[best_location]
while count < 200 and asteroids:
    visible = visible_from(best_location)
    best_angle = laser_angle + 4*pi
    best_v = None
    for v in visible:
        angle = angles_bl[v]
        while (angle < laser_angle):
            angle += 2*pi
        angle = angle-laser_angle
        if angle < best_angle:
            best_angle = angle
            best_v = v
    asteroids.remove(best_v)
    laser_angle += best_angle + 0.00000001
    while laser_angle > 2*pi:
        laser_angle -= 2*pi
    count += 1

print("Part 2:", best_v, best_v[0] * 100 + best_v[1])


class LargestPrimeFactorial(object):
	def __init__(self):
		pass

	def _isPrime(self, n):
		i = 2
		while i <= (n / 2):
			if n % i == 0:
				return False
			i += 1
		return True

	def _calcPrimesTo(self, n):
		returnVal = []
		i = 3
		while i < n:
			if self._isPrime(i):
				returnVal.append(i)
			i += 1
		print('Got all the primes')
		return returnVal

	def _calcFactorsOf(self, thisList, factor):
		returnVal = []
		for item in thisList:
			if factor % item == 0:
				returnVal.append(item)
		return returnVal

	def calc(self, n):
		allFactorialPrimes = self._calcFactorsOf(self._calcPrimesTo(n), n)
		print(allFactorialPrimes)
		return max(allFactorialPrimes)

def main():
	largestPrimeFactorial = LargestPrimeFactorial()
	largestPrimeFactorial.calc(600851475143)

if __name__ == '__main__':
	main()


from urllib.request import urlopen
import json
import pickle
import os
import csv

def dict_to_list_of_dicts(dict_list, relevant_properties):
    a = list()
    for i in range(0, len(dict_list)):
        a.append({j:dict_list[i][j] for j in relevant_properties})
    return a

def read_urls(path):
    urls = list()
    file = open(path, "r")
    for line in file:
        split_line = line.strip()
        split_line = split_line.split(";")
        urls.append((split_line[0], split_line[1], split_line[2]))
    return urls

#add dictionary to outide for dataset name
# def pickle_locations(url_location, path):
#     urls = read_urls(url_location)
#     PATH = path
#     if os.path.exists(PATH):
#         os.remove(PATH)
#     a = {}
#     file = open(PATH, "wb+")
#     for i in urls:
#         request = urlopen(i[0])
#         data = json.loads(request.read())
#         json_path = i[1].strip()
#         json_path = json_path.split(",")
#         id = data[json_path[0]][0]["id"]
#         id = id.split(".")[0]
#         request = urlopen(i[0])
#         data = json.loads(request.read())
#         json_path = i[1].strip()
#         json_path = json_path.split(",")
#         relevant_properties = i[2].strip()
#         relevant_properties = relevant_properties.split(",")
#         a.update({id: dict_to_list_of_dicts([j[json_path[1]] for j in data[json_path[0]]], relevant_properties)})
#         for j in a[id]:
#             if relevant_properties[0] == "coordinates":
#                 if type(j["coordinates"][0]) is list:
#                     if len(j["coordinates"]) == 1:
#                         j["coordinates"] = j["coordinates"][0]
#                     else:
#                         for k in range(1, len(j["coordinates"])):
#                             a[id].append({"coordinates": j["coordinates"][k]})
#                         j["coordinates"] = j["coordinates"][0]
#         print(str(id) + " " + str(a[id][:100]))
#     pickle.dump(a, file)
#     file.close()

def pickle_suburb_data(url_location, path):
    urls = read_urls(url_location)
    PATH = path
    if os.path.exists(PATH):
        os.remove(PATH)
    a = list()
    file = open(PATH, "wb+")
    for i in urls:
        request = urlopen(i[0])
        data = json.loads(request.read())
        json_path = i[1].strip()
        json_path = json_path.split(",")
        relevant_properties = i[2].strip()
        relevant_properties = relevant_properties.split(",")
        for j in range(0, len(data[json_path[0]])):
            j_temp = data[json_path[0]][j]
            for k in dict_to_list_of_dicts([j_temp[json_path[1]]], relevant_properties):
                if len(a) > j:
                    a[j].update({l: j_temp[json_path[1]][l] for l in k})
                else:
                    a.append({l: j_temp[json_path[1]][l] for l in k})
    pickle.dump(a, file)
    file.close()

def append_internet_location_data_pickle(a, csv_path, path, suburb_list):
    PATH = path
    if os.path.exists(PATH):
        os.remove(PATH)
    num_internets = dict()
    file = open(PATH, "wb+")
    with open(csv_path) as int_loc:
        reader = csv.reader(int_loc, delimiter = ",")
        for row in reader:
            if not row[3] in num_internets:
                num_internets[row[3]] = 1
            else:
                num_internets[row[3]] +=1
    for i in range(0, len(suburb_list)):
        if suburb_list[i] in num_internets:
            a[i]["number_internet_locations"] = num_internets[suburb_list[i]]
        else:
            a[i]["number_internet_locations"] = 0
    pickle.dump(a, file)
    file.close()

def get_suburb_list():
    a = list()
    request = urlopen("https://myvic-app-dev-gis.beta.vic.gov.au/geoserver/myvic/ows?service=WFS&version=1.0.0&request=GetFeature&outputFormat=application%2Fjson&typeName=myvic%3Aeducation_suburb")
    data = json.loads(request.read())
    for j in range(0, len(data["features"])):
        j_temp = data["features"][j]
        a.append(j_temp["properties"]["ssc_name"])
    return a

if __name__ == "__main__":
    PATH = "prepared_data.pkl"
    pickle_suburb_data("urls", PATH)
    file = open(PATH, "rb")
    a = pickle.load(file)
    file.close()
    suburb_list = get_suburb_list()
    append_internet_location_data_pickle(a, "public_internet_locations_vic.csv",PATH, suburb_list)


import numpy
import torch.utils.data
import numpy as np
import pandas as pd
from database_access import psycopg2_connection_to_database, psycopg2_connection_to_remote_database, \
    sqlalchemy_connection_to_database, sqlalchemy_connection_to_remote_database
import time
import shutil
import os
from sklearn.preprocessing import StandardScaler

def standardization(x, stat_len):
    """x = x.numpy()
    shape = np.shape(x)
    x = np.reshape(x, (-1, stat_len))
    scaler = StandardScaler()
    scaler.fit(x)
    x_stand = scaler.transform(x)
    x_stand = torch.from_numpy(np.reshape(x_stand,shape))"""
    m = x.view(-1, stat_len).mean(-2, keepdim=True)
    s = x.view(-1, stat_len).std(-2, unbiased=False, keepdim=True)
    x_stand = (x - m) / s
    safe_tensor = torch.where(torch.isnan(x_stand), torch.zeros_like(x_stand), x_stand)
    return safe_tensor

def get_column_names(table_name, location, run_location):
    if location == "local":
        con = psycopg2_connection_to_database()
    elif location == "remote":
        con, server = psycopg2_connection_to_remote_database(run_location)
    column_names = []
    with con.cursor() as cur:
        cur.execute("SELECT column_name FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = '%s'" % (table_name))
        columns = cur.fetchall()
        for tuple in columns:
            column_names += [tuple[0]]
        con.close()
    if location == "remote":
        server.stop()
    return column_names

def load_data(batch_size, num_of_train_year, window_size, end_year, granularity, database_location, run_location,
              epochs):
    engine, server = sqlalchemy_connection_to_remote_database()
    column_names = get_column_names()
    table_name = ("cumulative_team_stats", "game_stats")
    with engine.connect() as con:
        query = 'SELECT g."teamIdHome",g."teamIdAway",g."goalsHome",g."goalsAway",%s FROM %s AS t \
                     JOIN %s AS g ON t."gameId" = g."gameId" \
                     WHERE t."gameId" >= %s AND t."gameId" < %s \
                     ORDER BY t."timeDate" ASC, t."gameId" ASC' \
                % (",".join(f't."{name}"' for name in column_names), table_name[0], table_name[1],
                   int(str(start_year) + "020000"), int(str(end_year + 1) + "020000"))
        df = pd.read_sql(query, con)

    column_name = get_column_names('cumulative_player_stats_4')
    START_YEAR = 2000
    cur.execute(
        'SELECT %s FROM cumulative_player_stats_4 AS p JOIN game_stats AS g ON p.game = g."gameId" WHERE p.game >= %s ORDER BY g."timeDate"' % (
            ",".join(f'p.{name}' for name in column_name), int(str(START_YEAR) + "020000")))
    data = cur.fetchall()
    trn_data = []
    val_data = []
    highest_id = 0
    for match in data:
        if match[1] > highest_id:
            highest_id = match[1]
        if match[2] > highest_id:
            highest_id = match[2]
        if int(str(match[0])[:4]) < END_YEAR:
            trn_data.append(match)
        else:
            val_data.append(match)
    print("trn")
    trn_dataset = Dataset(trn_data)
    print("val")
    print("highest", highest_id)
    val_dataset = Dataset(val_data)
    trn_loader = torch.utils.data.DataLoader(trn_dataset, batch_size=batch_size, shuffle=False)
    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=32)
    cur.close()
    con.close()
    # print("Connection to database is CLOSED.")
    return trn_loader, val_loader

def load_player_data(num_of_train_year, window_size, end_year, granularity, database_location, run_location, file_name):
    """
    if location == "local":
        con = psycopg2_connection_to_database()
        table_name = 'cumulative_player_stats_6'
    elif location == "remote":
        con, server = psycopg2_connection_to_remote_database(run_location)
        table_name = 'cumulative_player_stats'
    column_names = get_column_names(table_name, location,run_location)
    max_length = 0
    game_numbers = []
    start_train_year = end_year - num_of_test_year - num_of_val_year - num_of_train_year
    start_val_year = end_year - num_of_test_year - num_of_val_year
    start_test_year = end_year - num_of_test_year
    with con.cursor() as cur:
    """
    start_all = time.perf_counter()
    start_cpu = time.process_time()

    if database_location == "local":
        con = psycopg2_connection_to_database()
        table_name = ("cumulative_player_stats_6", "game_stats")
    elif database_location == "remote":
        con, server = psycopg2_connection_to_remote_database(run_location)
        table_name = ("cumulative_player_stats", "game_stats")
    column_names = get_column_names(table_name[0], database_location, run_location)
    game_numbers = {}
    max_length = [0, 0]
    # print(num_of_train_year)
    start_year = end_year - num_of_train_year
    if start_year <= 2004:
        start_year -= 1  # -1 protoze v roce 2004 byla vyluka a nehral se jediny zapas
    print(start_year)
    with con.cursor() as cur:
        cur.execute('SELECT pi."primaryPosition",%s FROM %s AS p \
                     JOIN %s AS g ON p.game = g."gameId" \
                     JOIN player_info AS pi ON p.player = pi.id \
                     WHERE p.game >= %s AND p.game < %s \
                     ORDER BY g."timeDate" ASC , p.game ASC, p.side ASC, p.player ASC' # ORDER BY g."timeDate" ASC , p.game ASC, p.side ASC, p.player ASC'
                    % (",".join(f'p."{name}"' for name in column_names), table_name[0], table_name[1],
                       int(str(start_year) + "020000"),
                       int(str(end_year + 1) + "020000")))
        data = cur.fetchall()
    con.close()
    if database_location == "remote":
        server.stop()
    # print(len(data),len(data[0]))
    grouped_data = []
    labels = []
    game_idx = column_names.index('game') + 1
    team_idx = column_names.index('team') + 1
    player_idx = column_names.index('player') + 1
    result_idx = column_names.index('results') + 1
    side_idx = column_names.index('side') + 1
    goals_per_minute_idx = column_names.index('goalsPerMinute') + 1
    goalie_stats = ['timeOnIce', 'pim', 'shortHandedShotsAgainst', 'powerPlayShotsAgainst', 'evenShotsAgainst', 'saves',
                    'shots', 'assists', 'goals', 'shortHandedSaves', 'powerPlaySaves', 'evenSaves', 'goalsPerMinute',
                    'shotsPerMinute', 'assistsPerMinute', 'pimPerMinute', 'goalsPerGame', 'shotsPerGame',
                    'assistsPerGame', 'pimPerGame', 'timeOnIcePerGame', 'savesPerMinute', 'savesPerGame',
                    'savesPercentage', 'powerPlaySavesPerGame', 'powerPlaySavesPercentage', 'shortHandedSavesPerGame',
                    'shortHandedSavesPercentage', 'evenSavesPerGame', 'evenSavesPercentage']
    skater_stats = ['timeOnIce', 'takeaways', 'faceoffTaken', 'shortHandedAssists', 'faceOffWins', 'powerPlayTimeOnIce',
                    'shots', 'assists', 'goals', 'plusMinus', 'hits', 'shortHandedGoals', 'evenTimeOnIce', 'blocked',
                    'goalsPerMinute', 'shotsPerMinute', 'assistsPerMinute', 'powerPlayGoals', 'powerPlayAssists',
                    'goalsPerGame', 'shotsPerGame', 'assistsPerGame', 'timeOnIcePerGame', 'shortHandedTimeOnIce',
                    'giveaways', 'hitsPerMinute', 'goalsPercentage', 'takeawaysPerMinute', 'giveawaysPerMinute',
                    'blockedPerMinute', 'powerPlayGoalsPerMinute', 'powerPlayAssistsPerMinute',
                    'shortHandedGoalsPerMinute', 'shortHandedAssistsPerMinute', 'evenGoalsPerMinute',
                    'evenAssistsPerMinute', 'hitsPerGame', 'takeawaysPerGame', 'giveawaysPerGame', 'blockedPerGame',
                    'powerPlayGoalsPerGame', 'powerPlayAssistsPerGame', 'shortHandedGoalsPerGame',
                    'shortHandedAssistsPerGame', 'evenGoalsPerGame', 'evenAssistsPerGame', 'faceoffPercentage',
                    'powerPlayTimeOnIcePerGame', 'shortHandedTimeOnIcePerGame', 'teamGoalsPercentage',
                    'teamShotsPercentage']  # , 'teamBlockedPercentage', 'teamTakeawaysPercentage', 'teamHitsPercentage']
    if granularity == "Medium":
        temp = []
        for stat in goalie_stats:
            if "Per" in stat:
                continue
            else:
                temp.append(stat)
        goalie_stats = temp
        temp = []
        for stat in skater_stats:
            if "Per" in stat:
                continue
            else:
                temp.append(stat)
        skater_stats = temp
    temp = []
    for stat in skater_stats:
        if "hits" in stat:
            continue
        elif "takeaways" in stat:
            continue
        elif "giveaways" in stat:
            continue
        elif "blocked" in stat:
            continue
        else:
            temp.append(stat)
    skater_stats = temp
    print(skater_stats)
    col_names = get_column_names(table_name[0], database_location, run_location)
    goalie_stats_index = np.where(np.isin(col_names, goalie_stats))[0] + 1
    skater_stats_index = np.where(np.isin(col_names, skater_stats))[0] + 1
    player_stats_index = (skater_stats_index, goalie_stats_index)
    # print(goalie_stats_index)
    last_side = None
    if __name__ != '__main__':
        shutil.move(os.path.join("scripts", "loading", file_name), os.path.join("scripts", "loaded", file_name))

    """
    data1 = data
    start = time.process_time()

    data = np.array(data)
    keys = np.unique(data[:, 4])
    reindex_player_id = dict(zip(keys, [i for i in range(len(keys))]))
    """
    print("CPU loading: ", time.process_time() - start_cpu)
    print("ALL loading: ", time.perf_counter() - start_all)
    start_all = time.perf_counter()
    start_cpu = time.process_time()

    reindex_player_id = {}
    i = 1
    for item in data:
        player_id = item[player_idx]
        if player_id not in reindex_player_id:
            reindex_player_id[player_id] = i
            i += 1
    # print("CPU reindex: ", time.process_time() - start_cpu)
    # print("ALL reindex: ", time.perf_counter() - start_all)

    # print(reindex_player_id)
    average_player = []
    i = -1
    for item in data:
        # print(item[decision_stat_idx])
        if item[0] == "Goalie":
            goalie_flag = 1
        else:
            goalie_flag = 0

        if item[game_idx] not in game_numbers:
            i += 1
            game_numbers[item[game_idx]] = i
            #print(game_numbers)
            # game_numbers.append(item[game_idx])
            # game_numbers.append(i)

            if item[side_idx] == 0:
                labels.append(item[result_idx])
            if granularity == "Low":
                grouped_data.append([[],
                                     []])
            else:
                grouped_data.append([[[], []],  #
                                     [[], []]])  # +6s
                average_player.append([[[0] * len(player_stats_index[0]), [0] * len(player_stats_index[1])],  #
                                       [[0] * len(player_stats_index[0]), [0] * len(player_stats_index[1])]])  #

        if granularity == "Low":
            grouped_data[-1][item[side_idx]].append(reindex_player_id[item[player_idx]])
            # print(item[player_idx])
        elif granularity == "Medium":
            # print(list(item[player_stats_index[goalie_flag]].astype(np.float32)))
            stats = [item[index] for index in player_stats_index[goalie_flag]]  # +2s

            grouped_data[-1][item[side_idx]][goalie_flag].append(stats)  # +1s
            # print(type(item[player_stats_index[goalie_flag]].astype(np.float32).tolist()))
            # print(item[game_idx],item[side_idx])
            # print(average_player[-1][int(item[side_idx])][goalie_flag])
            # print(list(item[player_stats_index[goalie_flag]].astype(np.float32)))
            # print(stats)
            # print(item[game_idx],list(game_numbers.keys())[list(game_numbers.values()).index(i)])
            # print(average_player[-1][int(item[side_idx])][goalie_flag])
            # print(stats)
            average_player[-1][item[side_idx]][goalie_flag] = [a + b for a, b in
                                                               zip(average_player[-1][item[side_idx]][goalie_flag],
                                                                   stats)]
            # print(average_player[-1][int(item[side_idx])][goalie_flag])
            # print(average_player[-1][int(item[side_idx])][goalie_flag])
        elif granularity == "High":
            stats = [item[index] for index in player_stats_index[goalie_flag]]  # +2s
            #stats[0] = int(str(item[game_idx])[2:4]+str(item[game_idx])[-4:])
            grouped_data[-1][item[side_idx]][goalie_flag].append(stats)  # +1s
            average_player[-1][item[side_idx]][goalie_flag] = [a + b for a, b in
                                                               zip(average_player[-1][item[side_idx]][goalie_flag],
                                                                   stats)]

            # grouped_data[-1][int(item[side_idx])][goalie_flag].append(item[player_stats_index[goalie_flag]].astype(np.float32).tolist())
            # average_player[-1][int(item[side_idx])][goalie_flag] = [a + b for a, b in zip(average_player[-1][int(item[side_idx])][goalie_flag],item[player_stats_index[goalie_flag]].astype(np.float32).tolist())]

        if granularity == "Low":
            if max_length[0] < len(grouped_data[-1][item[side_idx]]):
                max_length[0] = len(grouped_data[-1][item[side_idx]])
        else:
            if max_length[goalie_flag] < len(grouped_data[-1][item[side_idx]][goalie_flag]):
                max_length[goalie_flag] = len(grouped_data[-1][item[side_idx]][goalie_flag])

    print("CPU grouped: ", time.process_time() - start_cpu)
    print("ALL grouped: ", time.perf_counter() - start_all)
    start_all = time.perf_counter()
    start_cpu = time.process_time()

    # print(grouped_data[-1])
    # print(len(grouped_data),len(grouped_data[-1]),len(grouped_data[-1][item[side_idx]]))

    # print(type(grouped_data[0]), type(grouped_data[0][0]), type(grouped_data[0][0][0]), type(grouped_data[0][0][0][0]), type(grouped_data[0][0][0][0][0]))
    # print(self.max_length)
    # stat_len = len(grouped_data[0][0][0])
    # print(max_length)
    none_zero = []
    for g, game in enumerate(grouped_data):
        none_zero.append([])
        for t, team in enumerate(game):
            none_zero[g].append([])
            if granularity == "Low":
                num_of_players = len(grouped_data[g][t])
                none_zero[g][t] = [num_of_players]
                for m in range(max_length[0] - num_of_players):
                    grouped_data[g][t].append(0)
            else:
                for s, stat in enumerate(team):
                    # print(g,t,s)
                    num_of_players = len(grouped_data[g][t][s])
                    for m in range(max_length[s] - num_of_players):
                        grouped_data[g][t][s].append([stat / num_of_players for stat in average_player[g][t][s]])

                # print(len(grouped_data[g][t][s]))
                # pass
        # for j in range(max_length - len(grouped_data[i][0])):
        #    grouped_data[i][0].append([0] * stat_len)
        # for j in range(max_length - len(grouped_data[i][1])):
        #    grouped_data[i][1].append([0] * stat_len)
        # grouped_data[i][0] += [0] * (max_length - len(grouped_data[i][0]))
        # grouped_data[i][1] += [0] * (max_length - len(grouped_data[i][1]))
        # if len(grouped_data[i][0]) < max_length or len(grouped_data[i][1]) < max_length:
        #    print(len(grouped_data[i][0]), len(grouped_data[i][1]))
        # print(grouped_data[i][0])
        # print(game, len(grouped_data[i][0]),len(grouped_data[i][1]),len(grouped_data[i][0][29]),len(grouped_data[i][1][29]))
        # print(len(grouped_data), len(grouped_data[i]), len(grouped_data[i][0]),len(grouped_data[i][0]))

    print("CPU padded: ", time.process_time() - start_cpu)
    print("ALL padded: ", time.perf_counter() - start_all)
    start_all = time.perf_counter()
    start_cpu = time.process_time()

    if granularity == "Low":
        stat_len = (max(reindex_player_id.values()), max(reindex_player_id.values()))
    else:
        stat_len = (len(grouped_data[0][0][0][0]), len(grouped_data[0][0][1][0]))
    trn_label = []
    trn_data = []
    trn_game = []
    trn_none_zero = []
    tst_label = []
    tst_data = []
    tst_game = []
    tst_none_zero = []

    for game_num, i in game_numbers.items():
        if int(str(game_num)[:4]) < end_year: #> start_year: #
            if not trn_game:
                trn_data.append([])
                trn_label.append([])
                trn_game.append([])
                trn_none_zero.append([])
            elif int(str(game_num)[:4]) != int(str(trn_game[-1][0])[:4]):
                trn_data.append([])
                trn_label.append([])
                trn_game.append([])
                trn_none_zero.append([])
            #if int(str(game_num)[:4]) == end_year and int(str(game_num)[-4:]) > 1082:
            #    continue
            trn_data[-1].append(grouped_data[i])
            trn_label[-1].append(labels[i] + 1)
            trn_game[-1].append(game_num)
            if granularity == "Low":
                trn_none_zero[-1].append(none_zero[i])
            # print(len(trn_data),len(trn_data[-1]))
        elif int(str(game_num)[:4]) == end_year and int(str(game_num)[-4:]) <= 1082:
        #elif int(str(game_num)[:4]) == start_year:# and int(str(game_num)[-4:]) <= 1082:
            if not tst_game:
                tst_data.append([])
                tst_label.append([])
                tst_game.append([])
                tst_none_zero.append([])
            tst_data[-1].append(grouped_data[i])
            tst_label[-1].append(labels[i] + 1)
            tst_game[-1].append(game_num)
            if granularity == "Low":
                tst_none_zero[-1].append(none_zero[i])
    for i, year in enumerate(trn_game):
        last_game = None
        for game in year:
            if str(game)[:4] != last_game:
                print(game)
            last_game = str(game)[:4]
        print(i, len(year))

    # pokus = np.array(trn_data, dtype=object)
    # print(pokus.shape)
    # skaters = np.array(pokus[:, :, 0].tolist(), dtype=float)
    # goalies = np.array(pokus[:, :, 1].tolist(), dtype=float)
    # print(skaters.shape, goalies.shape)
    tst_dataset = {}
    trn_dataset = {}
    trn_dataset["data"] = trn_data
    trn_dataset["labels"] = trn_label
    if granularity == "Low":
        trn_dataset["none_zero"] = trn_none_zero
    trn_loader = Loader(trn_dataset, window_size, test=False)
    tst_dataset["data"] = tst_data
    tst_dataset["labels"] = tst_label
    if granularity == "Low":
        tst_dataset["none_zero"] = tst_none_zero
    tst_loader = Loader(tst_dataset, None, test=True)
    print("CPU ready: ", time.process_time() - start_cpu)
    print("ALL ready: ", time.perf_counter() - start_all)
    start_all = time.perf_counter()
    start_cpu = time.process_time()
    return trn_loader, tst_loader, stat_len

def load_team_data(num_of_train_year, window_size, end_year, granularity, database_location, run_location, file_name):
    if database_location == "local":
        con = psycopg2_connection_to_database()
        table_name = ("cumulative_team_stats", "game_stats")
    elif database_location == "remote":
        con, server = psycopg2_connection_to_remote_database(run_location)
        table_name = ("cumulative_team_stats", "game_stats")
    column_names = get_column_names(table_name[0], database_location, run_location)
    print(column_names)
    game_numbers = []
    # print(num_of_train_year)
    start_year = end_year - num_of_train_year
    # print(start_year)
    if start_year <= 2004:
        start_year -= 1  # -1 protoze v roce 2004 byla vyluka a nehral se jediny zapas
    with con.cursor() as cur:
        cur.execute('SELECT g."teamIdHome",g."teamIdAway",g."goalsHome",g."goalsAway",%s FROM %s AS t \
                     JOIN %s AS g ON t."gameId" = g."gameId" \
                     WHERE t."gameId" >= %s AND t."gameId" < %s \
                     ORDER BY t."timeDate" ASC, t."gameId" ASC'
                    % (",".join(f't."{name}"' for name in column_names), table_name[0], table_name[1],
                       int(str(start_year) + "020000"), int(str(end_year + 1) + "020000")))
        data = cur.fetchall()
    con.close()
    if database_location == "remote":
        server.stop()
    shutil.move(os.path.join("scripts", "loading", file_name), os.path.join("scripts", "loaded", file_name))
    # print(len(data), len(data[0]))
    print(column_names)
    print(data[0])
    game_idx = column_names.index('gameId') + 4
    home_team_idx = column_names.index('teamId0') + 4
    away_team_idx = column_names.index('teamId1') + 4
    home_total_idx = column_names.index('goalsForAllTimeAverageHome0') + 4
    away_total_idx = column_names.index('goalsForAllTimeAverageHome1') + 4
    grouped_data = []
    labels = []

    for item in data:
        game_numbers.append(item[game_idx])
        labels.append(np.sign(item[2] - item[3]))
        if granularity == "Low":
            grouped_data.append(list(item[:2]))
        elif granularity == "Medium":
            grouped_data.append(list(item[home_team_idx + 1:home_total_idx] + item[away_team_idx + 1:away_total_idx]))
            # print(len(grouped_data[-1]))
            # print(grouped_data[-1])
        elif granularity == "High":
            #test = list(item[home_team_idx + 1:away_team_idx] + item[away_team_idx + 1:])
            #test.insert(0,int(str(item[game_idx])[2:4]+str(item[game_idx])[-4:]))
            grouped_data.append(list(item[home_team_idx + 1:away_team_idx] + item[away_team_idx + 1:]))
            # print(len(grouped_data[-1]))
            # print(grouped_data[-1])
    # print(np.unique(np.array(grouped_data).flatten()))
    print(grouped_data[0])
    if granularity == "Low":
        stat_len = np.amax(np.array(grouped_data)) + 1
        # print(stat_len)
    else:
        stat_len = len(grouped_data[0])
    trn_label = []
    trn_data = []
    trn_game = []
    tst_label = []
    tst_data = []
    tst_game = []
    for i, game_num in enumerate(game_numbers):
        if int(str(game_num)[:4]) < end_year:
            if not trn_game:
                trn_data.append([])
                trn_label.append([])
                trn_game.append([])
            elif int(str(game_num)[:4]) != int(str(trn_game[-1][0])[:4]):
                trn_data.append([])
                trn_label.append([])
                trn_game.append([])
            trn_data[-1].append(grouped_data[i])
            trn_label[-1].append(labels[i] + 1)
            trn_game[-1].append(game_num)
            # print(len(trn_data),len(trn_data[-1]))
        elif int(str(game_num)[:4]) == end_year and int(str(game_num)[-4:]) <= 1082:
            if not tst_game:
                tst_data.append([])
                tst_label.append([])
                tst_game.append([])
            tst_data[-1].append(grouped_data[i])
            tst_label[-1].append(labels[i] + 1)
            tst_game[-1].append(game_num)

    tst_dataset = {}
    trn_dataset = {}
    trn_dataset["data"] = trn_data
    trn_dataset["labels"] = trn_label
    # print(len(trn_data), len(trn_data[0]))
    trn_loader = Loader(trn_dataset, window_size, test=False)
    tst_dataset["data"] = tst_data
    tst_dataset["labels"] = tst_label
    # print(len(tst_data), len(tst_data[0]))
    tst_loader = Loader(tst_dataset, None, test=True)
    """
    for i in range(3):
        print("epoch",i)
        for (trn, val) in trn_loader:
            print(trn['data'].shape)
            print(val['data'].shape)
    """
    # trn_loader = torch.utils.data.DataLoader(trn_dataset, batch_size=batch_size, shuffle=False)
    # tst_loader = torch.utils.data.DataLoader(tst_dataset, batch_size=batch_size)
    return trn_loader, tst_loader, stat_len

class Loader():
    def __init__(self, dataset, sliding_window, test):
        self.test = test
        # self.data = [np.array(year, dtype='f4') for year in dataset["data"]]
        # print(len(dataset["data"]),len(dataset["data"][0]),len(dataset["data"][0][0]),len(dataset["data"][0][0][0]),len(dataset["data"][0][0][0][1]))
        try:
            if "none_zero" in dataset:
                self.none_zero = [np.array(year, dtype='l') for year in dataset["none_zero"]]
                self.data = [np.array(year, dtype=np.int64) for year in dataset["data"]]
            else:
                self.data = [np.array(year, dtype='f4') for year in dataset["data"]]
                self.none_zero = None
        except:

            for year in dataset["data"]:
                np_year = np.array(year, dtype=object)
                skater_year = np_year[:, :, 0].tolist()
                goalie_year = np_year[:, :, 1].tolist()
            self.data = {"skaters": [np.array(np.array(year, dtype=object)[:, :, 0].tolist(), dtype='f4') for year in
                                     dataset["data"]],
                         "goalies": [np.array(np.array(year, dtype=object)[:, :, 1].tolist(), dtype='f4') for year in
                                     dataset["data"]]}

        # print(self.data[0].shape)
        # print(type(dataset["data"][0][0]))
        self.labels = dataset["labels"]
        # print(len(self.labels))
        self.sliding_window = sliding_window
        self.index = 0
        self.window = bool(sliding_window)

    def __len__(self):
        if self.window:
            return len(self.labels) - self.sliding_window
        else:
            return len(self.labels) - 1

    def __iter__(self):
        self.index = 0
        if not self.window:
            self.sliding_window = 0
        return self

    def __next__(self):
        if not self.test:
            trn_batch = {}
            val_batch = {}
            if not self.window:
                self.sliding_window += 1
            if self.index + self.sliding_window + 1 > len(self.labels):
                raise StopIteration
            for i in range(self.sliding_window + 1): #,-1,-1):
                # print("index", i, "sliding window", self.sliding_window)
                # print(trn_batch)
                if not trn_batch:
                    trn_batch["labels"] = torch.tensor(self.labels[self.index + i], dtype=torch.long)
                    try:
                        trn_batch["data"] = torch.from_numpy(self.data[self.index + i])
                        if self.none_zero is not None:
                            trn_batch["none_zero"] = torch.from_numpy(self.none_zero[self.index + i])
                    except:
                        trn_batch["data"] = {"skaters": torch.from_numpy(self.data["skaters"][self.index + i]),
                                             "goalies": torch.from_numpy(self.data["goalies"][self.index + i])}
                elif i < self.sliding_window/2: # i != 0:
                    print("trn", i)
                    trn_batch["labels"] = torch.cat(
                        (trn_batch["labels"], torch.tensor(self.labels[self.index + i], dtype=torch.long)), 0)
                    try:
                        trn_batch["data"] = torch.cat((trn_batch["data"], torch.from_numpy(self.data[self.index + i])),
                                                      0)
                        if self.none_zero is not None:
                            trn_batch["none_zero"] = torch.cat(
                                (trn_batch["none_zero"], torch.from_numpy(self.none_zero[self.index + i])), 0)
                    except:
                        trn_batch["data"]["skaters"] = torch.cat(
                            (trn_batch["data"]["skaters"], torch.from_numpy(self.data["skaters"][self.index + i])), 0)
                        trn_batch["data"]["goalies"] = torch.cat(
                            (trn_batch["data"]["goalies"], torch.from_numpy(self.data["goalies"][self.index + i])), 0)
                else:
                    print("val", i)
                    if not val_batch:
                        val_batch["labels"] = torch.tensor(self.labels[self.index + i], dtype=torch.long)
                        try:
                            val_batch["data"] = torch.from_numpy(self.data[self.index + i])
                            if self.none_zero is not None:
                                val_batch["none_zero"] = torch.from_numpy(self.none_zero[self.index + i])
                        except:
                            val_batch["data"] = {"skaters": torch.from_numpy(self.data["skaters"][self.index + i]),
                                                 "goalies": torch.from_numpy(self.data["goalies"][self.index + i])}
                    else:
                        val_batch["labels"] = torch.cat(
                            (val_batch["labels"], torch.tensor(self.labels[self.index + i], dtype=torch.long)), 0)
                        try:
                            val_batch["data"] = torch.cat(
                                (val_batch["data"], torch.from_numpy(self.data[self.index + i])),
                                0)
                            if self.none_zero is not None:
                                val_batch["none_zero"] = torch.cat(
                                    (val_batch["none_zero"], torch.from_numpy(self.none_zero[self.index + i])), 0)
                        except:
                            val_batch["data"]["skaters"] = torch.cat(
                                (val_batch["data"]["skaters"], torch.from_numpy(self.data["skaters"][self.index + i])),
                                0)
                            val_batch["data"]["goalies"] = torch.cat(
                                (val_batch["data"]["goalies"], torch.from_numpy(self.data["goalies"][self.index + i])),
                                0)

                # print(val_batch["data"].shape)
                # print(trn_batch['data'].shape)
            if self.window:
                self.index += 1
            if val_batch:
                return (trn_batch, val_batch)
            else:
                return (trn_batch)
        else:
            trn_batch = {}
            if self.index == len(self.labels):
                raise StopIteration
            trn_batch["labels"] = torch.tensor(self.labels[0], dtype=torch.long)
            try:
                trn_batch["data"] = torch.from_numpy(self.data[0])
                if self.none_zero is not None:
                    trn_batch["none_zero"] = torch.from_numpy(self.none_zero[0])
            except:
                trn_batch["data"] = {"skaters": torch.from_numpy(self.data["skaters"][0]),
                                     "goalies": torch.from_numpy(self.data["goalies"][0])}
            self.index += 1
            return (trn_batch)

"""
class Loader():
    def __init__(self, dataset, sliding_window, start_year, num_of_val_year, epochs):
        self.dataset = dataset
        self.window = sliding_window
        self.start_year = start_year
        print(start_year)
        print(dataset[0]["game"])
        self.unique_years = []
        last = 0
        for i in dataset:
            if int(str(i["game"])[:4]) != last:
                self.unique_years.append(int(str(i["game"])[:4]))
                last = int(str(i["game"])[:4])
        self.num_of_unique_years = len(self.unique_years)
        self.num_of_val_year = num_of_val_year
        self.num_of_epochs = epochs
        self.num_of_slides = self.num_of_unique_years - (self.window + self.num_of_val_year) + 1
        self.window_years = self.unique_years[self.cur_slide:self.cur_slide+self.window+self.num_of_val_year]
        self.window_index = 0
        print(self.window+self.num_of_val_year, self.num_of_unique_years,self.num_of_slides,self.num_of_epochs//self.num_of_slides)
        self.cur_epoch = 0
        self.cur_slide = 0
        self.trn_batch = {}
        self.val_batch = {}
        self.cur_year = self.window_years[self.window_index]

    def __iter__(self):
        self.cur_epoch += 1
        if self.cur_epoch % (self.num_of_epochs//self.num_of_slides) == 0:
            self.cur_slide += 1
            self.window_years = self.unique_years[self.cur_slide:self.cur_slide + self.window + self.num_of_val_year]
            print("changin")
        return self

    def __next__(self):
        for match in self.dataset:
            #print(self.cur_slide,self.window,self.unique_years[self.cur_slide:self.cur_slide+self.window])
            if int(str(match["game"])[:4]) > self.cur_year: #and int(str(match["game"])[:4]) <= window_years[-1]:
                print(batch["data"].size(),batch["labels"].size())
                self.window_index += 1
                self.cur_year = self.window_years[self.window_index]

            if int(str(match["game"])[:4]) == self.cur_year:
                if not batch:
                    batch["labels"] = torch.tensor([match["labels"]])
                    batch["data"] = torch.unsqueeze(torch.from_numpy(match["data"]), 0)
                else:
                    batch["labels"] = torch.cat((batch["labels"], torch.tensor([match["labels"]])), 0)
                    batch["data"] = torch.cat((batch["data"], torch.unsqueeze(torch.from_numpy(match["data"]), 0)), 0)
                #print(batch["labels"].size(),batch["data"].size())
            elif int(str(match["game"])[:4]) > window_years[-1]:
                raise StopIteration
"""

class Dataset(torch.utils.data.Dataset):
    def __init__(self, labels, data, game):
        self.labels = np.array(labels)
        self.data = np.array(data)
        self.game = np.array(game)
        # print(self.labels.shape)
        # print(self.data.shape)
        """
        self.start_year = START
        self.end_year = END
        self.column_names = get_column_names('cumulative_player_stats_4')
        self.game_numbers = []
        #print(self.column_names)
        with con.cursor() as cur:
            cur.execute('SELECT %s FROM cumulative_player_stats_4 AS p \
                         JOIN game_stats AS g ON p.game = g."gameId" \
                         WHERE p.game >= %s AND p.game < %s \
                         ORDER BY g."timeDate" ASC, p.game ASC, p.side ASC, p.player ASC'
                         % (",".join(f'p."{name}"' for name in self.column_names), int(str(self.start_year) + "020000"), int(str(self.end_year + 1) + "020000")) )
            self.data = cur.fetchall()
        self.grouped_data = []
        self.labels = []
        game_idx = self.column_names.index('game')
        team_idx = self.column_names.index('team')
        player_idx = self.column_names.index('player')
        result_idx = self.column_names.index('results')
        side_idx = self.column_names.index('side')
        self.max_length = 0
        for item in self.data:
            if item[game_idx] not in self.game_numbers:
                self.game_numbers.append(item[game_idx])
                if item[side_idx] == 0:
                    self.labels.append(item[result_idx])
                else:
                    self.labels.append(-item[result_idx])
                self.grouped_data.append([[], []])
            if item[side_idx] == 0:
                self.grouped_data[-1][0] += list(item[result_idx + 1:])
            else:
                self.grouped_data[-1][1] += list(item[result_idx + 1:])
            if self.max_length < len(self.grouped_data[-1][0]):
                self.max_length = len(self.grouped_data[-1][0])
            elif self.max_length < len(self.grouped_data[-1][1]):
                self.max_length = len(self.grouped_data[-1][1])
        #print(self.max_length)
        for i, game in enumerate(self.game_numbers):
            #print(game, self.labels[i], len(self.grouped_data[i][0]), len(self.grouped_data[i][1]))
            self.grouped_data[i][0] += [0] * (self.max_length - len(self.grouped_data[i][0]))
            self.grouped_data[i][1] += [0] * (self.max_length - len(self.grouped_data[i][1]))
            if len(self.grouped_data[i][0]) < self.max_length or len(self.grouped_data[i][1]) < self.max_length:
                print(len(self.grouped_data[i][0]),len(self.grouped_data[i][1]))
        #print(len(self.grouped_data),len(self.grouped_data[0]),len(self.grouped_data[0][0]))
        self.grouped_data = np.array(self.grouped_data)
        #print(self.grouped_data.shape)
        self.labels = np.array(self.labels)
        #print(self.labels.shape)
        """
        """
            if item[game_idx] not in self.game_numbers:
                self.game_numbers.append(item[game_idx])
            if item[game_idx] not in self.grouped_data:
                self.grouped_data[item[game_idx]] = {}
            if item[team_idx] not in self.grouped_data[item[game_idx]]:
                self.grouped_data[item[game_idx]][item[team_idx]] = []
            self.grouped_data[item[game_idx]][item[team_idx]] += list(item[result_idx+1:])
            if item[game_idx] not in self.labels:
                self.labels[item[game_idx]] = {'results': {}, 'side': {}}
            if item[team_idx] not in self.labels[item[game_idx]]['results']:
                self.labels[item[game_idx]]['results'][item[team_idx]] = item[result_idx]
                self.labels[item[game_idx]]['side'][item[team_idx]] = item[side_idx]
        for name in self.column_names:
            print(name, self.data[0][self.column_names.index(name)])
        for game in self.grouped_data:
            for team in self.grouped_data[game]:
                #print(len(self.grouped_data[game][team]))
                if len(self.grouped_data[game][team]) > self.max_length:
                    self.max_length = len(self.grouped_data[game][team])

        print(self.grouped_data[item[game_idx]])
        print(self.labels[item[game_idx]])
        #print(self.max_length)
        """
        # print(self.game_numbers)
        # values = len(np.unique(list(zip(*self.data))[1]))
        # print(len(self.game_numbers),len(values))
        """
            data_array = []
        for teamId in data:
            data_vector = [0] * 54
            #print(teamId[1],teamId[2])
            data_vector[teamId[1]-1] = 1
            data_vector[teamId[2]-1] = 1
            data_array.append(data_vector)
        #print(data_array)
        #self.data = [np.array(teamId[1:3]) for teamId in data]
        self.data = data_array
        #print(data[0])
        #print(self.data[0])
        goal_diff = [goals[3]-goals[4] for goals in data]
        self.labels = np.sign(goal_diff)+1
        weights = [len(goal_diff) / np.count_nonzero(self.labels == 0), len(goal_diff) / np.count_nonzero(self.labels == 1),
                   len(goal_diff) / np.count_nonzero(self.labels == 2)]
        print(weights)

        #print(self.labels)
        """

    def __getitem__(self, i):
        """
        return_dict = {'labels': [], 'data': []}
        game_number = self.game_numbers[i]
        teams = list(self.labels[game_number]['side'].keys())
        if self.labels[game_number]['side'][teams[0]] == 0:
            home_team = teams[0]
            away_team = teams[1]
        elif self.labels[game_number]['side'][teams[1]] == 0:
            home_team = teams[1]
            away_team = teams[0]
        else:
            print('something is wrong')
        return_dict['labels'] = self.labels[game_number]['results'][home_team]

        return_dict['data'] = self.grouped_data[game_number][home_team] + [0]*(self.max_length-len(self.grouped_data[game_number][home_team])) + self.grouped_data[game_number][away_team] + [0]*(self.max_length-len(self.grouped_data[game_number][away_team]))
        """
        return {
            'game': self.game[i],
            'labels': self.labels[i].astype('i8'),
            # torch wants labels to be of type LongTensor, in order to compute losses
            'data': self.data[i].astype('f4')  # First retype to float32 (default dtype for torch)
        }

    def __len__(self):
        # return len(np.unique(list(zip(*self.data))[1]))
        return self.labels.size

if __name__ == '__main__':
    goalie_stats = ['timeOnIce', 'pim', 'shortHandedShotsAgainst', 'powerPlayShotsAgainst', 'evenShotsAgainst', 'saves',
                    'shots', 'assists', 'goals', 'shortHandedSaves', 'powerPlaySaves', 'evenSaves', 'goalsPerMinute',
                    'shotsPerMinute', 'assistsPerMinute', 'pimPerMinute', 'goalsPerGame', 'shotsPerGame',
                    'assistsPerGame',
                    'pimPerGame', 'timeOnIcePerGame', 'savesPerMinute', 'savesPerGame', 'savesPercentage',
                    'powerPlaySavesPerGame', 'powerPlaySavesPercentage', 'shortHandedSavesPerGame',
                    'shortHandedSavesPercentage', 'evenSavesPerGame', 'evenSavesPercentage']
    skater_stats = ['timeOnIce', 'takeaways', 'faceoffTaken', 'shortHandedAssists', 'faceOffWins', 'powerPlayTimeOnIce',
                    'shots', 'assists', 'goals', 'plusMinus', 'hits', 'shortHandedGoals', 'evenTimeOnIce', 'blocked',
                    'goalsPerMinute', 'shotsPerMinute', 'assistsPerMinute', 'powerPlayGoals', 'powerPlayAssists',
                    'goalsPerGame', 'shotsPerGame', 'assistsPerGame', 'timeOnIcePerGame', 'shortHandedTimeOnIce',
                    'giveaways', 'hitsPerMinute', 'goalsPercentage', 'takeawaysPerMinute', 'giveawaysPerMinute',
                    'blockedPerMinute', 'powerPlayGoalsPerMinute', 'powerPlayAssistsPerMinute',
                    'shortHandedGoalsPerMinute',
                    'shortHandedAssistsPerMinute', 'evenGoalsPerMinute', 'evenAssistsPerMinute', 'hitsPerGame',
                    'takeawaysPerGame', 'giveawaysPerGame', 'blockedPerGame', 'powerPlayGoalsPerGame',
                    'powerPlayAssistsPerGame', 'shortHandedGoalsPerGame', 'shortHandedAssistsPerGame',
                    'evenGoalsPerGame',
                    'evenAssistsPerGame', 'faceoffPercentage', 'powerPlayTimeOnIcePerGame',
                    'shortHandedTimeOnIcePerGame',
                    'teamGoalsPercentage', 'teamShotsPercentage', 'teamBlockedPercentage', 'teamTakeawaysPercentage',
                    'teamHitsPercentage']
    col_names = get_column_names("cumulative_player_stats_6", "local", "local")

    print(col_names)
    # print(np.array(col_names)[skater_stats_index])
    # print(np.array(col_names)[np.where(np.isin(col_names,goalie_stats+skater_stats,invert=True))])

    trn, tst, stat_len = load_player_data(3, 2, 2019, "High", "local", "local", "b.txt")

    print(len(trn))
    for i, (trn_b, val_b) in enumerate(trn):
        # print((trn_b["data"]))
        print((trn_b["data"]["skaters"].size()))
        print((trn_b["data"]["goalies"].size()))
        print((val_b["data"]["skaters"].size()))
        print((val_b["data"]["goalies"].size()))
        print(trn_b["data"]["skaters"][0][0][0][:6])
        print(val_b["data"]["skaters"][0][0][0][:6])
        print(trn_b["data"]["skaters"][4960][0][0][:6])
        print(val_b["data"]["skaters"][1270][0][0][:6])
        # print(len(trn_b["data"]))
        # print(type(trn_b["data"]))
        # print(type(trn_b))

        pass
    """
    data = Dataset(psycopg2_connection_to_database(),2000, 2002)
    print(len(data))
    for i in range(len(data)):
        print(data[i])
        if i == len(data)-1:
            break
        if len(data[i]) != len(data[i+1]):
            print(len(data[i]), len(data[i+1]))
        else:
            print(i,'ok')
    """


#importing modules
from tkinter import *
from tkinter.colorchooser import askcolor
from tkinter import *
from tkinter.colorchooser import askcolor
class Paint(object):

    pen_size = 5.0
    ink_color = 'black'

    def __init__(self):
        self.root = Tk()#initializing the Tkinter
        #pen
        self.pen = Button(self.root, text='pen', command=self.use_pen)#setup the Button for PEN
        self.pen.grid(row=0, column=0)#set up button grid
        #brush
        self.brush = Button(self.root, text='brush', command=self.use_brush)#setup Brush button
        self.brush.grid(row=0, column=1)#set up button grid
        #color-chooser
        self.color_button = Button(self.root, text='color', command=self.choose_color)#setup Color button
        self.color_button.grid(row=0, column=2)#set up button grid
        #eraser
        self.eraser = Button(self.root, text='eraser', command=self.use_eraser)#setup eraser button
        self.eraser.grid(row=0, column=3)  #set up button grid
        #choose-size
        self.choose_size = Scale(self.root, from_=1, to=10, orient=HORIZONTAL)#setup size scale
        self.choose_size.grid(row=0, column=4)#set up button grid
        #paint window
        self.draw_canvas = Canvas(self.root, bg='white', width=600, height=600)#setup canvas
        self.draw_canvas.grid(row=1, columnspan=5)#set up button grid

        self.setup()
        self.root.mainloop()

    def setup(self):
        self.old_x = None
        self.old_y = None
        self.line_width = self.choose_size.get()
        self.color = self.ink_color
        self.eraser_on = True
        self.active_button = self.pen
        self.draw_canvas.bind('<B1-Motion>', self.paint)
        self.draw_canvas.bind('<ButtonRelease-1>', self.reset)

    def use_pen(self):
        self.activate_button(self.pen)

    def use_brush(self):
        self.activate_button(self.brush)

    def choose_color(self):
        self.eraser_on = False
        self.color = askcolor(color=self.color)[1]

    def use_eraser(self):
        self.activate_button(self.eraser, eraser_mode=True)

    def activate_button(self, button, eraser_mode=False):
        self.active_button.config(relief=RAISED)
        button.config(relief=SUNKEN)
        self.active_button = button
        self.eraser_on = eraser_mode

    def paint(self, event):
        self.line_width = self.choose_size.get()
        paint_color = 'white' if self.eraser_on else self.color
        if self.old_x and self.old_y:
            self.draw_canvas.create_line(self.old_x, self.old_y, event.x, event.y,
                               width=self.line_width, fill=paint_color,
                               capstyle=ROUND, smooth=TRUE, splinesteps=36)
        self.old_x = event.x
        self.old_y = event.y

    def reset(self, event):
        self.old_x, self.old_y = None, None

if __name__ == '__main__':
    Paint()


def divisibleSumPairs(n, k, ar):

    pares = 0
    for i in range(n):
        for j in range(i + 1, n):
            if (ar[i] + ar[j]) % k == 0:
                pares += 1
    return pares

c = 6
d = 3
arr = [1, 3, 2, 6, 1, 2]
res = divisibleSumPairs(c, d, arr)
print(res)


# -*- coding: utf-8 -*-
# encoding: utf-8
"""
Created on Mon May 29 17:29:24 2017

@author: sky
"""
import pandas as pd
import numpy as np
from sklearn.utils import shuffle
from sklearn import preprocessing

fileName1 = '..\data\dsjtzs_txfz_training.txt'
fileName2 = '..\data\dsjtzs_txfz_test1.txt'
#min_max_scaler = preprocessing.MinMaxScaler()
#X_train = min_max_scaler.fit_transform(X_train)
#X_test = min_max_scaler.fit_trasform(X_test)
def giveValue(df,X):
    for i, s in enumerate(df['data']):
        for j,x in enumerate(s.split(';')[:-1]):
            X[i][j] = list(map(int,x.split(',')))
    return X

def getData(filename1,filename2):
    df1 = pd.read_csv(filename1,sep=' ',header=None)
    df1.columns = ['id','data','target','label']
    df1 = shuffle(df1)
    n1 = len(df1)
    df2 = pd.read_csv(filename2,sep=' ',header=None)
    df2.columns = ['id','data','target']
    n2 = len(df2)
    X_train = np.zeros((n1,300,3),dtype=np.int32)
    y = np.zeros((n1,1))
    y[:,0] = df1['label']

    X_test = np.zeros((n2,300,3),dtype=np.int32)

    X_train = giveValue(df1,X_train)
    X_test = giveValue(df2,X_test)
    return X_train,y,X_test

X_train, Y_train, X_test = getData(fileName1,fileName2)

from keras.models import Sequential
from keras.layers import Conv1D, MaxPooling1D
from keras.layers import Activation, Dropout, Flatten, Dense
from keras import backend as K
from keras.utils import np_utils,plot_model
from keras.optimizers import SGD
from keras import initializers

epochs = 50
batch_size = 300

#Y_train = np_utils.to_categorical(Y_train,2)

model = Sequential()
model.add(Conv1D(16, kernel_initializer='glorot_uniform',kernel_size=3,input_shape=(300,3)))
model.add(Activation('relu'))
model.add(Conv1D(32, kernel_size=3))
model.add(Activation('relu'))
model.add(MaxPooling1D(pool_size=2))
model.add(Conv1D(32, kernel_size=3))
model.add(Activation('relu'))
model.add(Conv1D(64, kernel_size=2))
model.add(Activation('relu'))
model.add(MaxPooling1D(pool_size=2))
model.add(Conv1D(64, kernel_size=3))
model.add(Activation('relu'))
model.add(Conv1D(64, kernel_size=3))
model.add(Activation('relu'))
model.add(MaxPooling1D(pool_size=2))

'''
model.add(Conv1D(16, kernel_size=3))
model.add(Activation('relu'))
model.add(Conv1D(16, kernel_size=3))
model.add(Activation('relu'))
model.add(Conv1D(16, kernel_size=2))
model.add(Activation('relu'))
model.add(MaxPooling1D(pool_size=2))
model.add(Conv1D(16, kernel_size=3))
model.add(Activation('relu'))
model.add(Conv1D(16, kernel_size=3))
model.add(Activation('relu'))
model.add(Conv1D(16, kernel_size=2))
model.add(Activation('relu'))
model.add(MaxPooling1D(pool_size=2))
model.add(Dropout(0.25))
model.add(Dense(16))
model.add(Activation('relu'))
'''
model.add(Flatten())
model.add(Dense(1, activation='sigmoid'))

sgd = SGD(lr=0.1, decay=1e-6, momentum=0.9, nesterov=True)
model.compile(optimizer=sgd, loss='mean_squared_error', metrics=['accuracy'])
model.fit(X_train, Y_train, batch_size=batch_size, epochs=epochs,validation_split=0.3)
model.save_weights('cnn3.h5',overwrite=True)

'''
df = model.predict_classes(X_test, batch_size=batch_size)
df = pd.DataFrame(df)
df.columns = ['behavior']
df2 = df[df['behavior'] == 0]
df2=df2.reset_index()
df2 = df2.drop(['behavior'],axis=1)
df2 = df2['index'] + 1
#df2.to_csv('cnn1.txt',index=False)
'''


#%%
from pathlib import Path

import matplotlib.pyplot as plt
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

from dataset import YearDataset
from utils import year_indices, get_valid_years

DEVICE='cuda:0'
#%%
model = AutoModelForCausalLM.from_pretrained('gpt2')
model.to(DEVICE)
tokenizer = AutoTokenizer.from_pretrained('gpt2')
years_to_sample_from = get_valid_years(tokenizer, 1000, 1900)
N = 400
ds = YearDataset(years_to_sample_from, N, Path("cache/potential_nouns.txt"), tokenizer, balanced=False, device=DEVICE, eos=True)

# %%
N_prior = 200
start_years = years_to_sample_from[torch.randint(years_to_sample_from.size(0), (N_prior,))]
years_XX = start_years // 100
years_XX00 = (start_years // 100) * 100
years = []
for XX00 in years_XX00:
    sample_space = years_to_sample_from[(years_to_sample_from >= XX00) & (years_to_sample_from < XX00+100)]
    years.append(sample_space[torch.randint(sample_space.size(0), (5,))])
years = torch.stack(years)
years_YY = years % 100

years_strings = [f'{str(y.tolist())[1:-1]}, {XX}' for y, XX in zip(years, years_XX)]
years_tokens = tokenizer(years_strings, return_tensors="pt")['input_ids'].to(DEVICE)

with torch.inference_mode():
    logits = model(years_tokens).logits[:, -1]
probs = torch.softmax(logits, dim=-1)
year_probs = probs[:, year_indices]
topk = torch.topk(probs, k=5)
topk_tokens = [[tokenizer._convert_id_to_token(top) for top in ex] for ex in topk.indices]

def comp_prob(probs, years_YY, gt=True):
    comps = []
    for prob, year in zip(probs, years_YY[:, -1]):
        if gt:
            comps.append(prob[year+1:].sum())
        else:
            comps.append(prob[:year+1].sum())
    return torch.stack(comps)
print(year_probs[torch.arange(year_probs.size(0)), years_YY[:, -1]+1].mean())
print(comp_prob(year_probs, years_YY).mean())
print(comp_prob(year_probs, years_YY, False).mean())
# %%
i = 4
plt.plot(year_probs[i].cpu())
plt.title(f"GPT-2 Probabilities when YY={years_YY[i].tolist()}")
plt.xlabel(f"Predicted Year")
plt.ylabel(f"probability")
plt.show()
# %%