from __future__ import annotations
from enum import auto
from typing import Optional, TYPE_CHECKING

from base_enum import BaseEnum
from monster_base import MonsterBase
from random_gen import RandomGen
from helpers import get_all_monsters

from data_structures.referential_array import ArrayR
from data_structures.queue_adt import CircularQueue
from data_structures.stack_adt import ArrayStack
from data_structures.array_sorted_list import ArraySortedList
from data_structures.sorted_list_adt import ListItem

if TYPE_CHECKING:
    from battle import Battle

class MonsterTeam:

    class TeamMode(BaseEnum):

        FRONT = auto()
        BACK = auto()
        OPTIMISE = auto()

    class SelectionMode(BaseEnum):

        RANDOM = auto()
        MANUAL = auto()
        PROVIDED = auto()

    class SortMode(BaseEnum):

        HP = auto()
        ATTACK = auto()
        DEFENSE = auto()
        SPEED = auto()
        LEVEL = auto()

    TEAM_LIMIT = 6

    def mapping(self, monster : MonsterBase, sort_mode : SortMode):
        """
        Returns a value which will be used as the key in the ListItem objecct

        :param monster: The monster from which we want to get the value from
        :para sort_mode: The attribute that the monsters will be sorted with reference to

        :complexity: O(1)
        """

        if sort_mode == self.SortMode.HP:
            return monster.get_hp()
        elif sort_mode == self.SortMode.ATTACK:
            return monster.get_attack()
        elif sort_mode == self.SortMode.DEFENSE:
            return monster.get_defense()
        elif sort_mode == self.SortMode.SPEED:
            return monster.get_speed()
        elif sort_mode == self.SortMode.LEVEL:
            return monster.get_level()

    def __init__(self, team_mode: TeamMode, selection_mode, **kwargs) -> None:
        """
        The function initializes a team of monsters based on the team mode and selection mode provided.

        :implementation:
            We use a different data structure for each team mode.
            For the front team mode we use a stack as the monster last added is the one that is retrieved first
            For the back team mode we use a queue here as the monster first added is the one that is retrieved first
            For the optimise team mode we use a sorted list data structure since it is able to sort monsters when they are being added

        :param team_mode: The `team_mode` parameter determines the mode in which the team is organized.
        :param selection_mode: The `selection_mode` parameter determines how the monsters are selected
        for the team.

        :complexity:
        All best cases are considered in the context of using the select provided to create the teams
        All worst cases are considered in the context of using the random generator to create the teams

        FRONT:
                O(n+m) best case
                O(n+ a * m) worst case
        BACK:
                O(n+m) best case
                O(n + a*m) worst case
        OPTIMISE:
                O(n + m) best case
                O(a*n*log(m!))  worst case

        where n is the size of the team,
        a is the number of monsters in the game
        and m is the number of monsters to be added

        """

        self.team_mode = team_mode
        self.lives = 1

        if team_mode == self.TeamMode.FRONT:
            self.team = ArrayStack[MonsterBase](self.TEAM_LIMIT)
        elif team_mode == self.TeamMode.BACK:
            self.team = CircularQueue[MonsterBase](self.TEAM_LIMIT)
        elif team_mode == self.TeamMode.OPTIMISE:
            self.sort_mode = kwargs.get('sort_key')
            self.descending = True
            self.team = ArraySortedList(self.TEAM_LIMIT)
        else:
            raise ValueError(f"team_mode {team_mode} not supported.")

        self.monsters = CircularQueue[MonsterBase](self.TEAM_LIMIT) #Keeping a track of the team for revival.

        if selection_mode == self.SelectionMode.RANDOM:
            self.select_randomly(**kwargs)
        elif selection_mode == self.SelectionMode.MANUAL:
            self.select_manually(**kwargs)
        elif selection_mode == self.SelectionMode.PROVIDED:
            self.select_provided(**kwargs)
        else:
            raise ValueError(f"selection_mode {selection_mode} not supported.")

    def add_to_team(self, monster: MonsterBase):
        """
        Adds a new monster to the team

        :implementation:
            According to which team mode has been chosen, we must change how we add to the team
            since each team mode uses a different data structure.
            For the optimise team mode, we need to additionally check if we need to sort in a descending or ascending order
            since the special for optimise changes the order that must be maintained.

        :param monster: The monster to be added to the team

        :complexity:

        FRONT:
            Best case: O(1)
            Worst case: O(1)
        BACK:
            Best case: O(1)
            Worst case: O(1)
        OPTIMISE:
            Best case: O(log(n))
            Worst case: O(n)

            where n is the number of monsters in the team
        """
        if self.team_mode == self.TeamMode.FRONT:
            self.team.push(monster)
        elif self.team_mode == self.TeamMode.BACK:
            self.team.append(monster)
        elif self.team_mode == self.TeamMode.OPTIMISE:
            self.optimise_add(self.team, monster)

    def optimise_add(self, team : ArraySortedList, monster : MonsterBase) -> None:
        """
        The function optimise_add adds a monster to a team with a specific sorting mode.

        :param team: The "team" parameter is an object representing a team.
        :param monster: The "monster" parameter is an object representing a monster.
        """
        if self.descending:
            item = ListItem(monster, -1 * self.mapping(monster, self.sort_mode))
            team.add(item)
        else:
            item = ListItem(monster, self.mapping(monster, self.sort_mode))
            team.add(item)

    def retrieve_from_team(self) -> MonsterBase:
        """
        The function retrieves a monster object from a team based on the team mode.

        :implementation:
            Based on the team mode, the function will return a monster from the team differently.

        :returns: A monster object
        :complexity:

        FRONT:
            Best case: O(1)
            Worst case: O(1)
        BACK:
            Best case: O(1)
            Worst case: O(1)
        OPTIMISE:
            Best case: O(n)
            Worst case: O(n)

            This is because when an element is deleted, all the elements after that one will need to be shifted to the left.

        where n is the number of monsters in the team
        """
        if self.team_mode == self.TeamMode.FRONT:
            return self.team.pop()
        elif self.team_mode == self.TeamMode.BACK:
            return self.team.serve()
        elif self.team_mode == self.TeamMode.OPTIMISE:
            return self.team.delete_at_index(0).value

    def special(self) -> None:
        """
        The special function rearranges the team based on the selected team mode.

        :implementation:
            Based on which team mode is selected, special will be executed differently.
            FRONT: The first 3 monsters in the team will be appended into a queue and then pushed back onto the stack.
            This will effectively reverse the order of the first 3 monsters in the team
            BACK: The first half of the team will be appended into a queue and the second half will be pushed onto the stack.
            Then the second half of the team will be appended back into the queue and then the first half of the team will be
            appended back into the queue.
            OPTIMISE: The descending boolean flag will be set to its opposite value and then each monster will be reinserted back into
            the team.

        :complexity:
        FRONT:
            Best case: O(1)
            Worst case: O(1)

            The worst case is O(1) since in the long run as the team grows arbitrarily large,
            the largest amount that will need to be reversed is 3 which is a constant
        BACK:
            Best case: O(n)
            Worst case: O(n)
        OPTIMISE:
            Best case: O(nlog(n))
            Worst case: O(n^2)

            The best case occurs when we switch from ascending to descending and the worst case occurs when we switch
            from descending to ascending

        where n is the number of monsters in the team.
        """

        if self.team_mode == self.TeamMode.FRONT:
            temp_team_holder = CircularQueue[MonsterBase](min(3, len(self.team)))
            for _ in range(min(3, len(self.team))):
                temp_team_holder.append(self.team.pop())
            for _ in range(len(temp_team_holder)):
                self.team.push(temp_team_holder.serve())

        elif self.team_mode == self.TeamMode.BACK:
            first_half_size = len(self.team)//2
            first_half_team = CircularQueue[MonsterBase](first_half_size)
            second_half_team = ArrayStack[MonsterBase](len(self.team) - first_half_size)
            for _ in range(first_half_size):
                first_half_team.append(self.team.serve())
            for _ in range(len(self.team)):
                second_half_team.push(self.team.serve())
            for _ in range(len(second_half_team)):
                self.team.append(second_half_team.pop())
            for _ in range(len(first_half_team)):
                self.team.append(first_half_team.serve())

        elif self.team_mode == self.TeamMode.OPTIMISE:
            temp = ArraySortedList[MonsterBase](len(self.team))
            self.descending = not self.descending
            for _ in range(len(self.team)):
                self.optimise_add(temp, self.retrieve_from_team())
            self.team = temp

    def regenerate_team(self) -> None:
        """
        The `regenerate_team` function resets the HP of all monsters in the team and adds them back to
        the team.

        :implementation:
            The team will be reset. This will be done differently depending on which team mode has been selected.
            Then, using the queue which stored all the monsters that were orginally in the team will be readded back into the team.
            They will also will be maintained within the queue so the team can be regenerated again if necessary.

        :complexity:
        FRONT:
            Best case: O(n)
            Worst case: O(n)
        BACK:
            Best case: O(n)
            Worst case: O(n)
        OPTIMISE:
            Best case: O(n)
            Worst case: O(nlog(n!))

        where n is the number of monsters.
        """

        if self.team_mode == self.TeamMode.FRONT or self.team_mode == self.TeamMode.BACK:
            self.team.clear()

        elif self.team_mode == self.TeamMode.OPTIMISE:
            self.team.reset()
            self.descending = True

        for _ in range(len(self.monsters)):
                monster = self.monsters.serve()
                monster.set_hp(monster.get_max_hp())
                self.add_to_team(monster)
                self.monsters.append(monster)

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

    def select_randomly(self):
        """"
        Creates a random team of monsters

        :complexity:
        Best case: O(m*n)
        Worst Case: O(m*n)

        where m is the number of monsters and n is the team size

        """
        team_size = RandomGen.randint(1, self.TEAM_LIMIT)
        monsters = get_all_monsters()
        n_spawnable = 0
        for x in range(len(monsters)):
            if monsters[x].can_be_spawned():
                n_spawnable += 1

        for _ in range(team_size):
            spawner_index = RandomGen.randint(0, n_spawnable-1)
            cur_index = -1
            for x in range(len(monsters)):
                if monsters[x].can_be_spawned():
                    cur_index += 1
                    if cur_index == spawner_index:
                        # Spawn this monster
                        self.add_to_team(monsters[x]())
                        self.monsters.append(monsters[x]())
                        break
            else:
                raise ValueError("Spawning logic failed.")

    def select_manually(self):
        """
        Prompt the user for input on selecting the team.
        Any invalid input should have the code prompt the user again.

        First input: Team size. Single integer
        For _ in range(team size):
            Next input: Prompt selection of a Monster class.
                * Should take a single input, asking for an integer.
                    This integer corresponds to an index (1-indexed) of the helpers method
                    get_all_monsters()
                * If invalid of monster is not spawnable, should ask again.

        Add these monsters to the team in the same order input was provided. Example interaction:

        How many monsters are there? 2
        MONSTERS Are:
        1: Flamikin [✔]
        2: Infernoth [❌]
        3: Infernox [❌]
        4: Aquariuma [✔]
        5: Marititan [❌]
        6: Leviatitan [❌]
        7: Vineon [✔]
        8: Treetower [❌]
        9: Treemendous [❌]
        10: Rockodile [✔]
        11: Stonemountain [❌]
        12: Gustwing [✔]
        13: Stormeagle [❌]
        14: Frostbite [✔]
        15: Blizzarus [❌]
        16: Thundrake [✔]
        17: Thunderdrake [❌]
        18: Shadowcat [✔]
        19: Nightpanther [❌]
        20: Mystifly [✔]
        21: Telekite [❌]
        22: Metalhorn [✔]
        23: Ironclad [❌]
        24: Normake [❌]
        25: Strikeon [✔]
        26: Venomcoil [✔]
        27: Pythondra [✔]
        28: Constriclaw [✔]print
        33: Leafadder [✔]
        34: Iceviper [✔]
        35: Rockpython [✔]
        36: Soundcobra [✔]
        37: Psychosnake [✔]
        38: Groundviper [✔]
        39: Faeboa [✔]
        40: Bugrattler [✔]
        41: Darkadder [✔]
        Which monster are you spawning? 38
        MONSTERS Are:
        1: Flamikin [✔]
        2: Infernoth [❌]
        3: Infernox [❌]
        4: Aquariuma [✔]
        5: Marititan [❌]
        6: Leviatitan [❌]
        7: Vineon [✔]
        8: Treetower [❌]
        9: Treemendous [❌]
        10: Rockodile [✔]
        11: Stonemountain [❌]
        12: Gustwing [✔]
        13: Stormeagle [❌]
        14: Frostbite [✔]
        15: Blizzarus [❌]
        16: Thundrake [✔]
        17: Thunderdrake [❌]
        18: Shadowcat [✔]
        19: Nightpanther [❌]
        20: Mystifly [✔]
        21: Telekite [❌]
        22: Metalhorn [✔]
        23: Ironclad [❌]
        24: Normake [❌]
        25: Strikeon [✔]()
        27: Pythondra [✔]
        28: Constriclaw [✔]
        29: Shockserpent [✔]
        30: Driftsnake [✔]
        31: Aquanake [✔]
        32: Flameserpent [✔]
        33: Leafadder [✔]
        34: Iceviper [✔]
        35: Rockpython [✔]
        36: Soundcobra [✔]
        37: Psychosnake [✔]
        38: Groundviper [✔]
        39: Faeboa [✔]
        40: Bugrattler [✔]
        41: Darkadder [✔]
        Which monster are you spawning? 2
        This monster cannot be spawned.
        Which monster are you spawning? 1

        :complexity: O(n) in both best/worst case where n is the number of monsters to be spawned
        """
        while True:
            try:
                team_size = int(input("How many monsters are there (pick between 1 and 6): "))
                if 1 <= team_size <= self.TEAM_LIMIT:
                    break
                print("Please enter a number between 1 and 6.")
            except ValueError:
                print("Please enter a number. It must be between 1 and 6.")
        monsters = get_all_monsters()
        while team_size > 0:
            try:
                monster_index = int(input("Enter the index of the monster you would like on your team (between 1 and 41): "))
                if 1 <= monster_index <= 41 and monsters[monster_index - 1].can_be_spawned():
                    self.add_to_team(monsters[monster_index- 1]())
                    self.monsters.append(monsters[monster_index- 1]())
                    team_size -= 1
                else:
                    print("Sorry, a monster with that index does not exist or cannot be spawned. Please enter another monster index between 1 and 41.")
            except ValueError:
                print("Sorry, please enter a number. It must be between 1 and 41")

    def select_provided(self, provided_monsters:Optional[ArrayR[type[MonsterBase]]]=None, **kwargs):
        """
        Generates a team based on a list of already provided monster classes.

        While the type hint imples the argument can be none, this method should never be None

        :complexity: O(n) in both best/worst case where n is the number of monsters provided

        Example team if in TeamMode.FRONT:
        [Gustwing Instance, Aquariuma Instance, Flamikin Instance]
        """
        for monster in provided_monsters:
            if monster.can_be_spawned() and not self.team.is_full():
                self.add_to_team(monster())
                self.monsters.append(monster())
            else:
                raise ValueError("Too many monsters or a monster cannot be spawned")

    def choose_action(self, currently_out: MonsterBase, enemy: MonsterBase) -> Battle.Action:
        # This is just a placeholder function that doesn't matter much for testing.
        from battle import Battle
        if currently_out.get_speed() >= enemy.get_speed() or currently_out.get_hp() >= enemy.get_hp():
            return Battle.Action.ATTACK
        if self.team.is_empty():
            return Battle.Action.ATTACK
        return Battle.Action.SWAP


import os
import json
import requests
import datetime
import sys

def read_file(path):
    f = open(path, "r")
    string = f.read()
    f.close()
    return string

def write_file(path, string, timestamp=False):
    time = datetime.datetime.now()
    path_time = path + "_" + time.strftime("%Y-%m-%d_%H") + "h.json"

    if os.path.exists(path_time):
        print("WRITE: Not written, same data from same hour.")
        return read_file(path_time)

    # Always overwrite default
    f = open(path + ".json", "w")
    f.write(string)
    f.close()

    # Creates times save if needed.
    if timestamp:
        f = open(path_time, "w")
        f.write(string)
        f.close()

    return string

def create_access_token(client_id, client_secret, region='eu'):
    data = {
        'grant_type': 'client_credentials'
    }

    response = requests.post(f'https://{region}.battle.net/oauth/token',
                             data=data, auth=(client_id, client_secret))

    return response

# reads the 'settings' file
settings = json.loads(read_file("settings"))
# print(settings)

region = settings["region"]
realm_name = settings["realm"]
client_id = settings["client_id"]
client_secret = settings["client_secret"]

# try creating a token
try:
    token = create_access_token(client_id, client_secret).json()[
        "access_token"]
except:
    print("Token was not given by Blizz API")
    exit()

# Gets the whole auction house for a realm
r = requests.get(
    f"https://{region}.api.blizzard.com/data/wow/realm/{realm_name}?namespace=dynamic-{region}&locale=en_GB&access_token={token}")

# Get the base request url
try:
    base_url = r.json()["connected_realm"]["href"].split("?")[0]
except:
    print("json was bad")
    print(r.text)
    exit()
# print(base_url)

r = requests.get(
    f"{base_url}/auctions?namespace=dynamic-{region}&locale=en_US&access_token={token}")

# saving the request to region-server and region-server-date
auctions = write_file(f"{region}-{realm_name}",
                      r.text, True)

auctions_json = json.loads(auctions)
ids = [int(x) for x in sys.argv[1:]]
act_list = []
# Parse all the AH for the items we have
for auction in auctions_json["auctions"]:
    if auction["item"]["id"] in ids:
        try:
            price = auction["buyout"] / 10000
        except:
            price = auction["unit_price"] / 10000

        act_list += [(auction["item"]["id"], price, auction["quantity"])]

act_list.sort()
# print(act_list)
prev = 0
for auction in act_list:
    if (auction[0] != prev):
        prev = auction[0]
        print("=============================================")
        r = requests.get(
            f"https://{region}.api.blizzard.com/data/wow/item/{prev}?namespace=static-{region}&locale=en_US&access_token={token}")
        print(r.json()["name"])

    print(auction[0], end='')

    mid_space = 20 - len(str(auction[0])) - len(str(auction[1]))
    print(' ' * mid_space, end='')

    print(auction[1], end='')

    mid_space = 10 - len(str(auction[2]))
    print(' ' * mid_space, end='')

    print(auction[2], end='\n')

# r = requests.get(f"https://{region}.battle.net/oauth/token")
# print(r.json())

# https://{region}.api.blizzard.com/data/wow/realm/{realm_name}?namespace=dynamic-{region}&locale=en_GB&access_token={token}

# https://{region}.api.blizzard.com/data/wow/connected-realm/{realm_id}/auctions?namespace=dynamic-{region}&locale=en_US&access_token={token}
# {base_url}/auctions?namespace=dynamic-{region}&locale=en_US&access_token={token}

# https://{region}.api.blizzard.com/data/wow/item/{item_id}?namespace=static-{region}&locale=en_US&access_token={token}


import random
from os import path

import e2e_tests_utils
from demisto_client.demisto_api.rest import ApiException

from demisto_sdk.commands.download.downloader import Downloader
from demisto_sdk.commands.format.format_module import format_manager
from demisto_sdk.commands.generate_docs import generate_playbook_doc
from demisto_sdk.commands.upload.uploader import Uploader
from demisto_sdk.commands.validate.validate_manager import ValidateManager

def test_e2e_demisto_sdk_flow_playbook_testsuite(tmpdir):
    # Importing TestSuite classes from Demisto-SDK, as they are excluded when pip installing the SDK.
    e2e_tests_utils.cli(f"mkdir {tmpdir}/git")
    e2e_tests_utils.git_clone_demisto_sdk(
        destination_folder=f"{tmpdir}/git/demisto-sdk", sdk_git_branch="master"
    )
    from TestSuite.playbook import Playbook
    from TestSuite.repo import Repo

    repo = Repo(tmpdir)

    unique_id = random.randint(1, 1000)
    pack_name = "foo_" + str(unique_id)
    pack = repo.create_pack(name=pack_name)
    playbook_name = "pb_" + pack_name
    playbook: Playbook = pack.create_playbook(name=playbook_name)
    playbook.create_default_playbook(name=playbook_name)
    assert path.exists(f"{tmpdir}/Packs/{pack_name}/Playbooks/{playbook_name}.yml")

    print(
        f"Trying to upload playbook from {tmpdir}/Packs/{pack_name}/Playbooks/{playbook_name}.yml"
    )
    Uploader(
        input=f"{tmpdir}/Packs/{pack_name}/Playbooks/{playbook_name}.yml", insecure=True
    ).upload()

    # Preparing updated pack folder
    e2e_tests_utils.cli(f"mkdir {tmpdir}/Packs/{pack_name}_testsuite")

    print(
        f"Trying to download the updated playbook from {playbook_name} to {tmpdir}/Packs/{pack_name}_testsuite/Playbooks"
    )
    Downloader(
        output=f"{tmpdir}/Packs/{pack_name}_testsuite",
        input=[playbook_name],
        insecure=True,
    ).download()
    assert path.exists(
        f"{tmpdir}/Packs/{pack_name}_testsuite/Playbooks/{playbook_name}.yml"
    )

    print(
        "Generating docs (creating a readme file)"
        f" for the playbook {tmpdir}/Packs/{pack_name}_testsuite/Playbooks/{playbook_name}.yml"
    )
    generate_playbook_doc.generate_playbook_doc(
        input_path=f"{tmpdir}/Packs/{pack_name}_testsuite/Playbooks/{playbook_name}.yml"
    )
    assert path.exists(
        f"{tmpdir}/Packs/{pack_name}_testsuite/Playbooks/{playbook_name}_README.md"
    )

    print(
        f"Formating playbook {tmpdir}/Packs/{pack_name}_testsuite/Playbooks/{playbook_name}.yml"
    )
    format_manager(
        input=f"{tmpdir}/Packs/{pack_name}_testsuite/Playbooks/{playbook_name}.yml",
        assume_yes=True,
    )
    print(
        f"Validating playbook {tmpdir}/Packs/{pack_name}_testsuite/Playbooks/{playbook_name}.yml"
    )
    ValidateManager(
        file_path=f"{tmpdir}/Packs/{pack_name}_testsuite/Playbooks/{playbook_name}.yml"
    ).run_validation()

    print(
        f"Uploading updated playbook {tmpdir}/Packs/{pack_name}_testsuite/Playbooks/{playbook_name}.yml"
    )
    Uploader(
        input=f"{tmpdir}/Packs/{pack_name}/Playbooks/{playbook_name}.yml", insecure=True
    ).upload()

def test_e2e_demisto_sdk_flow_playbook_client(tmpdir, insecure: bool = True):
    unique_id = random.randint(1, 1000)
    pack_name = "foo_" + str(unique_id)
    playbook_name = "pb_" + str(unique_id)

    demisto_client = e2e_tests_utils.connect_to_server(insecure=insecure)
    body = [
        {
            "name": playbook_name,
            "propagationLabels": ["all"],
            "tasks": {
                "0": {
                    "id": "0",
                    "unqiueId": "0",
                    "type": "start",
                    "nextTasks": None,
                    "task": {},
                }
            },
        }
    ]

    header_params = {
        "Accept": "application/json",  # noqa: E501
        "Accept-Encoding": "gzip, deflate, br",
        "Content-Type": "application/json",
    }

    try:
        demisto_client.api_client.call_api(
            resource_path="/playbook/save",
            method="POST",
            header_params=header_params,
            body=body,
        )
    except ApiException as ae:
        print(f"*** Failed to create playbook {playbook_name}, reason: {ae}")
        assert False

    # Preparing updated pack folder
    e2e_tests_utils.cli(f"mkdir -p {tmpdir}/Packs/{pack_name}_client")

    print(
        f"Trying to download the updated playbook from {playbook_name} to {tmpdir}/Packs/{pack_name}_client/Playbooks"
    )
    Downloader(
        output=f"{tmpdir}/Packs/{pack_name}_client",
        input=[playbook_name],
        insecure=True,
    ).download()
    assert path.exists(
        f"{tmpdir}/Packs/{pack_name}_client/Playbooks/{playbook_name}.yml"
    )

    print(
        "Generating docs (creating a readme file)"
        f" for the playbook {tmpdir}/Packs/{pack_name}_client/Playbooks/{playbook_name}.yml"
    )
    generate_playbook_doc.generate_playbook_doc(
        input_path=f"{tmpdir}/Packs/{pack_name}_client/Playbooks/{playbook_name}.yml"
    )
    assert path.exists(
        f"{tmpdir}/Packs/{pack_name}_client/Playbooks/{playbook_name}_README.md"
    )

    print(
        f"Formating playbook {tmpdir}/Packs/{pack_name}_client/Playbooks/{playbook_name}.yml"
    )
    format_manager(
        input=f"{tmpdir}/Packs/{pack_name}_client/Playbooks/{playbook_name}.yml",
        assume_yes=True,
    )
    print(
        f"Validating playbook {tmpdir}/Packs/{pack_name}_client/Playbooks/{playbook_name}.yml"
    )
    ValidateManager(
        file_path=f"{tmpdir}/Packs/{pack_name}_client/Playbooks/{playbook_name}.yml"
    ).run_validation()

    print(
        f"Uploading updated playbook {tmpdir}/Packs/{pack_name}_client/Playbooks/{playbook_name}.yml"
    )
    Uploader(
        input=f"{tmpdir}/Packs/{pack_name}/Playbooks/{playbook_name}.yml", insecure=True
    ).upload()


import numpy as np
import matplotlib.pyplot as plt
from matplotlib import patches
import argparse
import shutil
import subprocess
from tqdm import tqdm

import os

def draw_frame(state, filename):
    plt.figure(figsize=(8,8))

    length = 9.0
    radius = 1.0
    theta = state[0]

    xpos = length * np.sin(theta)
    ypos = -length * np.cos(theta)

    plt.plot([0, xpos], [0, ypos], 'k-', linewidth=2.0)

    circle = plt.Circle((xpos, ypos), radius, color='springgreen')
    plt.gca().add_patch(circle)

    circle = plt.Circle((0, 0), 0.3, color='black')
    plt.gca().add_patch(circle)

    plt.axis('equal')
    plt.xlim([-12, 12])
    plt.ylim([-12, 12])

    plt.savefig(filename)
    plt.close()

def dump_images(x_trj, temp_dir):
    T = x_trj.shape[0]
    for t in tqdm(range(T)):
        filename = "image_{:05d}".format(t)
        draw_frame(x_trj[t,:], os.path.join(temp_dir, filename))

def save_video(temp_dir, file_dir, video_name, framerate=30):
    image_name = os.path.join(temp_dir, "image_%5d.png")
    video_name = os.path.join(file_dir, video_name)
    subprocess.call([
        'ffmpeg', '-framerate', str(framerate), '-i',
        image_name, '-r', str(framerate), '-pix_fmt', 'yuv420p',
        video_name
    ])

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("filename", help=".npy trajectory file to animate.")
    args = parser.parse_args()
    file_path = args.filename

    x_trj = np.load(file_path)
    file_dir = os.path.split(file_path)[0]
    file_name = os.path.split(file_path)[1]
    temp_dir = os.path.join(file_dir, "temp")

    videofile_name = file_name[:-4] + ".mp4"

    os.mkdir(temp_dir)

    dump_images(x_trj, temp_dir)
    save_video(temp_dir, file_dir, videofile_name)

    shutil.rmtree(temp_dir)


import pandas as pd
import numpy as np
import yaml, time, os
import xarray as xr
import datetime
import subprocess
import shutil
import os

def saveRunLog(config,code_dir,outFile):
    """
    function to add the run specs to a csv log file
    :param config: [dict] the current config dictionary
    :param code_dir: [str] path to river-dl directory
    :param outFile: [str] the filename for the output
    """

    #check if the log already exists
    newLog = not os.path.exists(outFile)

    #add a default run description if needed
    if not "runDescription" in config.keys():
        config['runDescription']=""

    #replace commas with semi-colons for csv readability
    config['runDescription']=config['runDescription'].replace(",",";")

    with open(outFile,"a+") as f:
        if newLog:
            f.write("'Date','Directory','Description'\n")
        f.write("%s,'%s','%s'\n"%(datetime.date.today().strftime("%m/%d/%y"),os.path.join(os.path.relpath(os.getcwd(),code_dir),config['out_dir']),config['runDescription']))

def asRunConfig(config, code_dir, outFile):
    """
    function to save the as-run config settings to a text file
    :param config: [dict] the current config dictionary
    :param code_dir: [str] path to river-dl directory
    :param outFile: [str] the filename for the output
    """
    #store some run parameters
    config['runDate']=datetime.date.today().strftime("%m/%d/%y")
    with open(os.path.join(code_dir,".git/HEAD"),'r') as head:
        ref = head.readline().split(' ')[-1].strip()
        branch = ref.split("/")[-1]
    with open(os.path.join(code_dir,'.git/', ref),'r') as git_hash:
        commit = git_hash.readline().strip()
    status = str(subprocess.Popen(['git status'], cwd = None if code_dir=="" else code_dir,shell=True,stdout=subprocess.PIPE).communicate()[0]).split("\\n")
    modifiedFiles = [x.split()[1].strip() for x in status if "modified" in x]
    newFiles = [x.split()[1].strip() for x in status if "new file" in x]
    config['gitStatus']= 'unknown' if len(status)==1 else 'dirty' if len(modifiedFiles)>0 or len(newFiles)>0 else 'clean'
    #if the repo is dirty, make a zipped archive of the code directory
    if config['gitStatus']=='dirty':
        shutil.make_archive(os.path.join(os.path.dirname(outFile),"river_dl"),"zip",os.path.join(code_dir,"river_dl"))
    config['gitModified']=modifiedFiles
    config['gitNew']=newFiles
    config['gitBranch']=branch
    config['gitCommit'] = commit
    #and the file info for the input files
    config['input_file_info']={config[x]:{'file_size':os.stat(config[x]).st_size,'file_date':time.strftime("%m/%d/%Y %I:%M:%S %p",time.localtime(os.stat(config[x]).st_ctime))} for x in config.keys() if "file" in x and x!="input_file_info"}
    with open(outFile,'w') as f:
        yaml.dump(config, f, default_flow_style=False)

    #add the log entry
    saveRunLog(config,code_dir,os.path.join(code_dir,"runLog.csv"))

def scale(dataset, std=None, mean=None):
    """
    scale the data so it has a standard deviation of 1 and a mean of zero
    :param dataset: [xr dataset] input or output data
    :param std: [xr dataset] standard deviation if scaling test data with dims
    :param mean: [xr dataset] mean if scaling test data with dims
    :return: scaled data with original dims
    """
    if not isinstance(std, xr.Dataset) or not isinstance(mean, xr.Dataset):
        std = dataset.std(skipna=True)
        mean = dataset.mean(skipna=True)
    # adding small number in case there is a std of zero
    scaled = (dataset - mean) / (std + 1e-10)
    check_if_finite(std)
    check_if_finite(mean)
    return scaled, std, mean

def sel_partition_data(dataset, time_idx_name, start_dates, end_dates):
    """
    select the data from a date range or a set of date ranges
    :param dataset: [xr dataset] input or output data with date dimension
    :param time_idx_name: [str] name of column that is used for temporal index
        (usually 'time')
    :param start_dates: [str or list] fmt: "YYYY-MM-DD"; date(s) to start period
    (can have multiple discontinuos periods)
    :param end_dates: [str or list] fmt: "YYYY-MM-DD"; date(s) to end period
    (can have multiple discontinuos periods)
    :return: dataset of just those dates
    """
   # if it just one date range
    if isinstance(start_dates, str):
        if isinstance(end_dates, str):
            return dataset.sel({time_idx_name: slice(start_dates, end_dates)})
        else:
            raise ValueError("start_dates is str but not end_date")
    # if it's a list of date ranges
    elif isinstance(start_dates, list) or isinstance(start_dates, tuple):
        if len(start_dates) == len(end_dates):
            data_list = []
            for i in range(len(start_dates)):
                date_slice = slice(start_dates[i], end_dates[i])
                data_list.append(dataset.sel({time_idx_name: date_slice}))
            return xr.concat(data_list, dim=time_idx_name)
        else:
            raise ValueError("start_dates and end_dates must have same length")
    else:
        raise ValueError("start_dates must be either str, list, or tuple")

def separate_trn_tst(
    dataset,
    time_idx_name,
    train_start_date,
    train_end_date,
    val_start_date=None,
    val_end_date=None,
    test_start_date=None,
    test_end_date=None,
):
    """
    separate the train data from the test data according to the start and end
    dates. This assumes your training data is in one continuous block. Be aware,
    if your train/test/val partitions are discontinuous (composed of multiple
    periods), depending on your sequence length and how the data line up, you
    could end up with sequences starting in one period and ending in another.
    The breaking up of sequences would happen in the `convert_batch_reshape`
    function
    :param dataset: [xr dataset] input or output data with dims
    :param time_idx_name: [str] name of column that is used for temporal index
        (usually 'time')
    :param train_start_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to start
    train period (can have multiple discontinuous periods)
    :param train_end_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to end train
     period (can have multiple discontinuous periods)
    :param val_start_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to start
     validation period (can have multiple discontinuous periods)
    :param val_end_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to end
    validation period (can have multiple discontinuous periods)
    :param test_start_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to start
    test period (can have multiple discontinuous periods)
    :param test_end_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to end test
    period (can have multiple discontinuous periods)
    :return: [tuple] separated data
    """
    train = sel_partition_data(
        dataset, time_idx_name, train_start_date, train_end_date
    )

    if val_start_date and val_end_date:
        val = sel_partition_data(
            dataset, time_idx_name, val_start_date, val_end_date
        )

    elif val_start_date and not val_end_date:
        raise ValueError("With a val_start_date a val_end_date must be given")
    elif val_end_date and not val_start_date:
        raise ValueError("With a val_end_date a val_start_date must be given")
    else:
        val = None

    if test_start_date and test_end_date:
        test = sel_partition_data(
            dataset, time_idx_name, test_start_date, test_end_date
        )
    elif test_start_date and not test_end_date:
        raise ValueError("With a test_start_date a test_end_date must be given")
    elif test_end_date and not test_start_date:
        raise ValueError("With a test_end_date a test_start_date must be given")
    else:
        test = None

    return train, val, test

def split_into_batches(data_array, seq_len=365, offset=1.0):
    """
    split training data into batches with size of batch_size
    :param data_array: [numpy array] array of training data with dims [nseg,
    ndates, nfeat]
    :param seq_len: [int] length of sequences (e.g., 365)
    :param offset: [float] How to offset the batches. Values < 1 are taken as fractions, (e.g., 0.5 means that
    the first batch will be 0-365 and the second will be 182-547), values > 1 are used as a constant number of
    observations to offset by.
    :return: [numpy array] batched data with dims [nbatches, nseg, seq_len
    (batch_size), nfeat]
    """
    if offset>1:
        period = int(offset)
    else:
        period = int(offset*seq_len)
    num_batches = data_array.shape[1]//period
    combined=[]
    for i in range(num_batches+1):
        idx = int(period*i)
        batch = data_array[:,idx:idx+seq_len,...]
        combined.append(batch)
    combined = [b for b in combined if b.shape[1]==seq_len]
    combined = np.asarray(combined)
    return combined

def read_obs(obs_file, y_vars, x_data):
    """
    read and format multiple observation files. we read in the pretrain data to
    make sure we have the same indexing.
    :param x_data: [xr.Dataset] xarray dataset used to match spatial and
    temporal domain
    :param y_vars: [list of str] which variables_to_log to prepare data for
    :param obs_file: [list] filenames of observation file
    :return: [xr dataset] the observations in the same time
    """
    ds = xr.open_zarr(obs_file, consolidated=False)
    obs = xr.merge([x_data, ds], join="left")
    obs = obs[y_vars]
    return obs

def join_catch_properties(x_data_ts, catch_props):
    """
    append the catchment properties to the x time series data
    :param x_data_ts: [xr dataset] timeseries x-data
    :param catch_props: [xr dataset] catchment properties data
    :return: [xr dataset] the merged datasets
    """
    # broadcast the catchment properties on the ts data so that there is a value
    # for each date
    _, ds_catch = xr.broadcast(x_data_ts, catch_props)
    return xr.merge([x_data_ts, ds_catch], join="left")

def prep_catch_props(x_data_ts, catch_prop_file, catch_prop_vars, spatial_idx_name, replace_nan_with_mean=True):
    """
    read catch property file and join with ts data
    :param x_data_ts: [xr dataset] timeseries x-data
    :param catch_prop_file: [str] the feather file of catchment attributes
    :param catch_prop_vars: [list of str] the catchment attributes to use, if None, all attributes will be kept
    :param spatial_idx_name: [str] name of column that is used for spatial
        index (e.g., 'seg_id_nat')
    :param replace_nan_with_mean: [bool] if true, any nan will be replaced with
    the mean of that variable
    :return: [xr dataset] merged datasets
    """
    df_catch_props = pd.read_feather(catch_prop_file)

    #keep only the requested variables
    if catch_prop_vars:
        catch_prop_vars.append(spatial_idx_name)
        df_catch_props = df_catch_props[catch_prop_vars]

    # replace nans with column means
    if replace_nan_with_mean:
        df_catch_props = df_catch_props.apply(
            lambda x: x.fillna(x.mean()), axis=0
        )
    ds_catch_props = df_catch_props.set_index(spatial_idx_name).to_xarray()

    return join_catch_properties(x_data_ts, ds_catch_props)

def reshape_for_training(data):
    """
    reshape the data for training
    :param data: training data (either x or y_dataset or mask) dims: [nbatch, nseg,
    len_seq, nfeat/nout]
    :return: reshaped data [nbatch * nseg, len_seq, nfeat/nout]
    """
    n_batch, n_seg, seq_len, n_feat = data.shape
    return np.reshape(data, [n_batch * n_seg, seq_len, n_feat])

def get_exclude_start_end(exclude_grp):
    """
    get the start and end dates for the exclude group
    :param exclude_grp: [dict] dictionary representing the exclude group from
    the exclude yml file
    :return: [tuple of datetime objects] start date, end date
    """
    start = exclude_grp.get("start_date")
    if start:
        start = datetime.datetime.strptime(start, "%Y-%m-%d")

    end = exclude_grp.get("end_date")
    if end:
        end = datetime.datetime.strptime(end, "%Y-%m-%d")
    return start, end

def get_exclude_vars(exclude_grp):
    """
    get the variables_to_log to exclude for the exclude group
    :param exclude_grp: [dict] dictionary representing the exclude group from
    the exclude yml file
    :return: [list] variables_to_log to exclude
    """
    variable = exclude_grp.get("variable")
    if not variable or variable == "both":
        return ["seg_tave_water", "seg_outflow"]
    elif variable == "temp":
        return ["seg_tave_water"]
    elif variable == "flow":
        return ["seg_outflow"]
    else:
        raise ValueError("exclude variable must be flow, temp, or both")

def get_exclude_seg_ids(exclude_grp, all_segs):
    """
    get the segments to exclude
    :param exclude_grp: [dict] dictionary representing the exclude group from
    the exclude yml file
    :param all_segs: [array] all of the segments. this is needed if we are doing
    a reverse exclusion
    :return: [list like] the segments to exclude
    """
    # ex_segs are the sites to exclude
    if "seg_id_nats_ex" in exclude_grp.keys():
        ex_segs = exclude_grp["seg_id_nats_ex"]
    # exclude all *but* the "seg_id_nats_in"
    elif "seg_id_nats_in" in exclude_grp.keys():
        ex_mask = ~all_segs.isin(exclude_grp["seg_id_nats_in"])
        ex_segs = all_segs[ex_mask]
    else:
        ex_segs = all_segs
    return ex_segs

def exclude_segments(y_data, exclude_segs):
    """
    exclude segments from being trained on by setting their weights as zero
    :param y_data:[xr dataset] y_dataset data. this is used to get the dimensions
    :param exclude_segs: [list] list of segments to exclude in the loss
    calculation
    :return:
    """
    weights = initialize_weights(y_data, 1)
    for seg_grp in exclude_segs:
        # get the start and end dates is present
        start, end = get_exclude_start_end(seg_grp)
        exclude_vars = get_exclude_vars(seg_grp)
        segs_to_exclude = get_exclude_seg_ids(seg_grp, weights.seg_id_nat)

        # loop through the data_vars
        for v in exclude_vars:
            # set those weights to zero
            weights[v].load()
            weights[v].loc[
                dict(date=slice(start, end), seg_id_nat=segs_to_exclude)
            ] = 0
    return weights

def initialize_weights(y_data, initial_val=1):
    """
    initialize all weights with a value.
    :param y_data:[xr dataset] y_dataset data. this is used to get the dimensions
    :param initial_val: [num] a number to initialize the weights with. should
    be between 0 and 1 (inclusive)
    :return: [xr dataset] dataset weights initialized with a uniform value
    """
    weights = y_data.copy(deep=True)
    for v in y_data.data_vars:
        weights[v].load()
        weights[v].loc[:, :] = initial_val
    return weights

def reduce_training_data_random(
    data_file,
    train_start_date="1980-10-01",
    train_end_date="2004-09-30",
    reduce_amount=0,
    out_file=None,
    segs=None,
):
    """
    artificially reduce the amount of training data in the training dataset
    :param train_start_date: [str] date (fmt YYYY-MM-DD) for when training data
    starts
    :param train_end_date: [str] date (fmt YYYY-MM-DD) for when training data
    ends
    :param data_file: [str] path to the observations data file
    :param reduce_amount: [float] fraction to reduce the training data by.
    For example, if 0.9, a random 90% of the training data will be set to nan
    :param out_file: [str] file to which the reduced dataset will be written
    :param segs: [array-like] segments to reduce data of
    :return: [xarray dataset] updated weights (nan where reduced)
    """
    # read in an convert to dataframe
    ds = xr.open_zarr(data_file)
    df = ds.to_dataframe()
    idx = pd.IndexSlice
    df_trn = df.loc[idx[train_start_date:train_end_date, :], :]
    if segs:
        df_trn = df_trn.loc[idx[:, segs], :]
    non_null = df_trn.dropna()
    reduce_idx = non_null.sample(frac=reduce_amount).index
    df.loc[reduce_idx] = np.nan
    reduced_ds = df.to_xarray()
    if out_file:
        reduced_ds.to_zarr(out_file)
    return reduced_ds

def filter_reduce_dates(df, start_date, end_date, reduce_between=False):
    df_filt = df.copy()
    df_filt = df_filt.reset_index()
    if reduce_between:
        df_filt = df_filt[
            (df_filt["date"] > start_date) & (df_filt["date"] < end_date)
        ]
    else:
        df_filt = df_filt[
            (df_filt["date"] < start_date) | (df_filt["date"] > end_date)
        ]
    return df_filt.set_index(df.index.names)

def reduce_training_data_continuous(
    data_file,
    reduce_start="1980-10-01",
    reduce_end="2004-09-30",
    train_start=None,
    train_end=None,
    segs=None,
    reduce_between=True,
    out_file=None,
):
    """
    reduce the amount of data in the training dataset by replacing a section
    or the inverse of that section with nan
    :param data_file: [str] path to the observations data file
    :param reduce_start: [str] date (fmt YYYY-MM-DD) for when reduction data
    starts
    :param reduce_end: [str] date (fmt YYYY-MM-DD) for when reduction data
    ends
    :param train_start: [str] data (fmt YYYY-MM-DD) the start of the training
    data. If left None (default) the reduction will be done for all of the
    dates. This is only relevant if reduce_between == False, because in that
    case, the inverse of the range (reduce_start, reduce_end) is used and it may
    be necessary to limit making the nans to the training time period
    :param train_end: [str] data (fmt YYYY-MM-DD) the end of the training
    data. If left None (default) the reduction will be done for all of the
    dates. This is only relevant if reduce_between == False, because in that
    case, the inverse of the range (reduce_start, reduce_end) is used and it may
    be necessary to limit making the nans to the training time period
    :param segs: [list-like] segment id's for which the data should be reduced
    :param reduce_between: [bool] if True the data *in* the range (reduce_start,
    reduce_end) will be made nan. if False, the data *outside* of that range
    will be made nan
    :param out_file: [str] file to which the reduced dataset will be written
    :return: [xarray dataset] updated weights (nan where reduced)
    """
    # read in an convert to dataframe
    ds = xr.open_zarr(data_file)
    df = ds.to_dataframe()
    idx = pd.IndexSlice
    df_red = df.copy()
    df_red = df_red.loc[idx[train_start:train_end, :]]
    if segs:
        df_red = df_red.loc[idx[:, segs], :]
    df_red = filter_reduce_dates(
        df_red, reduce_start, reduce_end, reduce_between
    )
    df.loc[df_red.index] = np.nan
    reduced_ds = df.to_xarray()
    if out_file:
        reduced_ds.to_zarr(out_file)
    return reduced_ds

def convert_batch_reshape(
    dataset,
    spatial_idx_name="seg_id_nat",
    time_idx_name="date",
    seq_len=365,
    offset=1.0
):
    """
    convert xarray dataset into numpy array, swap the axes, batch the array and
    reshape for training
    :param dataset: [xr dataset] data to be batched
    :param spatial_idx_name: [str] name of column that is used for spatial
        index (e.g., 'seg_id_nat')
    :param time_idx_name: [str] name of column that is used for temporal index
        (usually 'time')
    :param seq_len: [int] length of sequences (e.g., 365)
    :param offset: [float] 0-1, how to offset the batches (e.g., 0.5 means that
    the first batch will be 0-365 and the second will be 182-547)
    :return: [numpy array] batched and reshaped dataset
    """
    # If there is no dataset (like if a test or validation set is not supplied)
    # just return None
    if not dataset:
        return None

    # convert xr.dataset to numpy array
    dataset = dataset.transpose(spatial_idx_name, time_idx_name)

    arr = dataset.to_array().values

    # if the dataset is empty, just return it as is
    if dataset[time_idx_name].size == 0:
        return arr

    # before [nfeat, nseg, ndates]; after [nseg, ndates, nfeat]
    # this is the order that the split into batches expects
    arr = np.moveaxis(arr, 0, -1)

    # batch the data
    # after [nbatch, nseg, seq_len, nfeat]
    batched = split_into_batches(arr, seq_len=seq_len, offset=offset)

    # reshape data
    # after [nbatch * nseg, seq_len, nfeat]
    reshaped = reshape_for_training(batched)
    return reshaped

def coord_as_reshaped_array(
    dataset,
    coord_name,
    spatial_idx_name="seg_id_nat",
    time_idx_name="date",
    seq_len=365,
    offset=1.0,
):
    """
    convert an xarray coordinate to an xarray data array and reshape that array
    :param dataset:
    :param coord_name: [str] the name of the coordinate to convert/reshape
    :param spatial_idx_name: [str] name of column that is used for spatial
        index (e.g., 'seg_id_nat')
    :param time_idx_name: [str] name of column that is used for temporal index
        (usually 'time')
    :param seq_len: [int] length of sequences (e.g., 365)
    :param offset: [float] 0-1, how to offset the batches (e.g., 0.5 means that
    the first batch will be 0-365 and the second will be 182-547)
    :return:
    """
    # If there is no dataset (like if a test or validation set is not supplied)
    # just return None
    if not dataset:
        return None

    # I need one variable name. It can be any in the dataset, but I'll use the
    # first
    first_var = next(iter(dataset.data_vars.keys()))
    coord_array = xr.broadcast(dataset[coord_name], dataset[first_var])[0]
    new_var_name = coord_name + "1"
    dataset[new_var_name] = coord_array
    reshaped_np_arr = convert_batch_reshape(
        dataset[[new_var_name]],
        spatial_idx_name,
        time_idx_name,
        seq_len=seq_len,
        offset=offset,
    )
    return reshaped_np_arr

def check_if_finite(xarr):
    assert np.isfinite(xarr.to_array().values).all()

def log_variables(y_dataset, variables_to_log):
    """
    take the log of given variables
    :param variables_to_log: [list of str] variables to take the log of
    :param y_dataset: [xr dataset] the y data
    :return: [xr dataset] the data logged
    """
    for v in variables_to_log:
        y_dataset[v].load()
        y_dataset[v].loc[:, :] = y_dataset[v] + 1e-6
        y_dataset[v].loc[:, :] = xr.ufuncs.log(y_dataset[v])
    return y_dataset

def prep_y_data(
    y_data_file,
    y_vars,
    x_data,
    train_start_date,
    train_end_date,
    val_start_date=None,
    val_end_date=None,
    test_start_date=None,
    test_end_date=None,
    val_sites=None,
    test_sites=None,
    spatial_idx_name="seg_id_nat",
    time_idx_name="date",
    seq_len=365,
    log_vars=None,
    exclude_file=None,
    normalize_y=True,
    y_type="obs",
    y_std=None,
    y_mean=None,
    trn_offset = 1.0,
    tst_val_offset = 1.0
):
    """
    prepare y_dataset data

    :param y_data_file: [str] temperature observations file
    :param y_vars: [str or list of str] target variable(s)
    :param x_data: [xr.Dataset] xarray dataset used to match spatial and
    temporal domain
    :param spatial_idx_name: [str] name of column that is used for spatial
        index (e.g., 'seg_id_nat')
    :param time_idx_name: [str] name of column that is used for temporal index
        (usually 'time')
    :param train_start_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to start
    train period (can have multiple discontinuous periods)
    :param train_end_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to end train
    period (can have multiple discontinuous periods)
    :param val_start_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to start
    validation period (can have multiple discontinuous periods)
    :param val_end_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to end
    validation period (can have multiple discontinuous periods)
    :param test_start_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to start
    test period (can have multiple discontinuous periods)
    :param test_end_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to end test
    period (can have multiple discontinuous periods)
    :param val_sites: [list of site_ids] sites to retain for validation. These
    sites will be witheld from training
    :param test_sites: [list of site_ids] sites to retain for testing. These
    sites will be witheld from training and validation
    :param seq_len: [int] length of sequences (e.g., 365)
    :param log_vars: [list-like] which variables_to_log (if any) to take log of
    :param exclude_file: [str] path to exclude file
    :param normalize_y: [bool] whether or not to normalize the y_dataset values
    :param y_type: [str] "obs" if observations or "pre" if pretraining
    :param y_std: [array-like] standard deviations of y_dataset variables_to_log
    :param y_mean: [array-like] means of y_dataset variables_to_log
    :param trn_offset: [str] value for the training offset
    :param tst_val_offset: [str] value for the testing and validation offset
    :returns: training and testing data along with the means and standard
    deviations of the training input and output data
    """
    # I assume that if `y_vars` is a string only one variable has been passed
    # so I put that in a list which is what the rest of the functions expect
    if isinstance(y_vars, str):
        y_vars = [y_vars]

    # when specifying mean and std, they get passed as an np.ndarray where we need xr.Datasets
    if isinstance(y_mean, np.ndarray):
        y_mean = xr.Dataset(dict(zip(y_vars,y_mean)))
        y_std = xr.Dataset(dict(zip(y_vars,y_std)))

    y_data = read_obs(y_data_file, y_vars, x_data)

    y_trn, y_val, y_tst = separate_trn_tst(
        y_data,
        time_idx_name,
        train_start_date,
        train_end_date,
        val_start_date,
        val_end_date,
        test_start_date,
        test_end_date,
    )

    # replace validation sites' (and test sites') data with np.nan
    if val_sites:
        y_trn = y_trn.where(~y_trn[spatial_idx_name].isin(val_sites))

    if test_sites:
        y_trn = y_trn.where(~y_trn[spatial_idx_name].isin(test_sites))
        y_val = y_val.where(~y_val[spatial_idx_name].isin(test_sites))

    if log_vars:
        y_trn = log_variables(y_trn, log_vars)

    # filter pretrain/finetune y_dataset
    if exclude_file:
        exclude_segs = read_exclude_segs_file(exclude_file)
        y_wgts = exclude_segments(y_trn, exclude_segs=exclude_segs)
    else:
        y_wgts = initialize_weights(y_trn)
    # scale y_dataset training data and get the mean and std
    # scale the validation partition to benchmark epoch performance
    if normalize_y:
    # check if mean and std are already calculated/exist
        if not isinstance(y_std, xr.Dataset) or not isinstance(y_mean, xr.Dataset):
            y_trn, y_std, y_mean = scale(y_trn)
        else:
            y_trn, _, _ = scale(y_trn,y_std,y_mean)

        if y_val:
            y_val, _, _ = scale(y_val, y_std, y_mean)

        if y_tst:
            y_tst, _, _ = scale(y_tst, y_std, y_mean)
    else:
        _, y_std, y_mean = scale(y_trn)
        y_std = y_std/y_std
        y_mean = y_mean * 0

    if y_type == 'obs':
        data = {
            "y_obs_trn": convert_batch_reshape(
                y_trn, spatial_idx_name, time_idx_name, offset=trn_offset, seq_len=seq_len
            ),
            "y_obs_wgts": convert_batch_reshape(
                y_wgts, spatial_idx_name, time_idx_name, offset=trn_offset, seq_len=seq_len
            ),
            "y_obs_val": convert_batch_reshape(
                y_val, spatial_idx_name, time_idx_name, offset=tst_val_offset, seq_len=seq_len
            ),
            "y_obs_tst": convert_batch_reshape(
                y_tst, spatial_idx_name, time_idx_name, offset=tst_val_offset, seq_len=seq_len
            ),
            "y_std": y_std.to_array().values,
            "y_mean": y_mean.to_array().values,
            "y_obs_vars": y_vars,
        }
    elif y_type == 'pre':
        if normalize_y:
            if not isinstance(y_std, xr.Dataset) or not isinstance(y_mean, xr.Dataset):
                y_trn, y_std, y_mean = scale(y_trn)
            else:
                y_data, _, _ = scale(y_data, y_std, y_mean)

        data = {
            "y_pre_full": convert_batch_reshape(
                y_data, spatial_idx_name, time_idx_name, offset=trn_offset, seq_len=seq_len
            ),
            "y_pre_trn": convert_batch_reshape(
                y_trn, spatial_idx_name, time_idx_name, offset=trn_offset, seq_len=seq_len
            ),
            "y_pre_vars": y_vars,
        }
    return data

def prep_all_data(
    x_data_file,
    y_data_file,
    x_vars,
    train_start_date,
    train_end_date,
    val_start_date=None,
    val_end_date=None,
    test_start_date=None,
    test_end_date=None,
    val_sites=None,
    test_sites=None,
    y_vars_finetune=None,
    y_vars_pretrain=None,
    spatial_idx_name="seg_id_nat",
    time_idx_name="date",
    seq_len=365,
    pretrain_file=None,
    distfile=None,
    dist_idx_name="rowcolnames",
    dist_type="updown",
    catch_prop_file=None,
    catch_prop_vars=None,
    exclude_file=None,
    log_y_vars=False,
    out_file=None,
    segs=None,
    normalize_y=True,
    trn_offset = 1.0,
    tst_val_offset = 1.0
):
    """
    prepare input and output data for DL model training read in and process
    data into training and testing datasets. the training and testing data are
    scaled to have a std of 1 and a mean of zero
    :param x_data_file: [str] path to Zarr file with x data. Data should have
    a spatial coordinate and a time coordinate that are specified in the
    `spatial_idx_name` and `time_idx_name` arguments
    :param y_data_file: [str] observations Zarr file. Data should have a spatial
    coordinate and a time coordinate that are specified in the
    spatial_idx_name` and `time_idx_name` arguments
    :param train_start_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to start
    train period (can have multiple discontinuous periods)
    :param train_end_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to end train
    period (can have multiple discontinuous periods)
    :param val_start_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to start
    validation period (can have multiple discontinuous periods)
    :param val_end_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to end
    validation period (can have multiple discontinuous periods)
    :param test_start_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to start
    test period (can have multiple discontinuous periods)
    :param test_end_date: [str or list] fmt: "YYYY-MM-DD"; date(s) to end test
    period (can have multiple discontinuous periods)
    :param val_sites: [list of site_ids] sites to retain for validation. These
    sites will be witheld from training
    :param test_sites: [list of site_ids] sites to retain for testing. These
    sites will be witheld from training and validation
    :param spatial_idx_name: [str] name of column that is used for spatial
    index (e.g., 'seg_id_nat')
    :param time_idx_name: [str] name of column that is used for temporal index
    (usually 'time')
    :param x_vars: [list] variables_to_log that should be used as input. If None, all
    of the variables_to_log will be used
    :param y_vars_finetune: [str or list of str] finetune target variable(s)
    :param y_vars_pretrain: [str or list of str] pretrain target variable(s)
    :param seq_len: [int] length of sequences (e.g., 365)
    :param pretrain_file: [str] Zarr file with the pretraining data. Should have
    a spatial coordinate and a time coordinate that are specified in the
    `spatial_idx_name` and `time_idx_name` arguments
    :param distfile: [str] path to the distance matrix .npz file
    :param dist_idx_name: [str] name of index to sort dist_matrix by. This is
    the name of an array in the distance matrix .npz file
    :param dist_type: [str] type of distance matrix ("upstream", "downstream" or
    "updown")
    :param catch_prop_file: [str] the path to the catchment properties file. If
    left unfilled, the catchment properties will not be included as predictors
    :param catch_prop_vars: [list of str] list of catchment properties to use. If
    left unfilled and a catchment property file is supplied all variables will be used.
    :param exclude_file: [str] path to exclude file
    :param log_y_vars: [bool] whether or not to take the log of discharge in
    training
    :param segs: [list-like] which segments to prepare the data for
    :param normalize_y: [bool] whether or not to normalize the y_dataset values
    :param out_file: [str] file to where the values will be written
    :param trn_offset: [str] value for the training offset
    :param tst_val_offset: [str] value for the testing and validation offset
    :returns: training and testing data along with the means and standard
    deviations of the training input and output data
            "x_trn": x training data
            "x_val": x validation data
            "x_tst": x test data
            "x_std": x standard deviations
            "x_mean": x means
            "x_cols": x column names
            "ids_trn": segment ids of the training data
            "times_trn": dates of the training data
            "ids_val": segment ids of the validation data
            "times_val": dates of the validation data
            "ids_tst": segment ids of the test data
            "times_tst": dates of the test data
            'y_pre_trn': y_dataset pretrain data for train set
            'y_obs_trn': y_dataset observations for train set
            "y_pre_wgts": y_dataset weights for pretrain data
            "y_obs_wgts": weights for y_dataset observations
            "y_obs_val": y_dataset observations for validation set
            "y_obs_tst": y_dataset observations for train set
            "y_std": standard deviations of y_dataset data
            "y_mean": means of y_dataset data
            "y_obs_vars": y_dataset observation variable names
            'y_pre_val': y_dataset pretrain data for validation data
            'y_pre_tst': y_dataset pretrain data for test data
            'y_pre_vars':  y_dataset pretrain data variable names
            "dist_matrix": prepared adjacency matrix
    """
    if pretrain_file and not y_vars_pretrain:
        raise ValueError("included pretrain file but no pretrain vars")

    x_data = xr.open_zarr(x_data_file,consolidated=False)
    x_data = x_data.sortby([spatial_idx_name, time_idx_name])

    if segs:
        x_data = x_data.sel({spatial_idx_name: segs})

    x_data = x_data[x_vars]

    if catch_prop_file:
        x_data = prep_catch_props(x_data, catch_prop_file, catch_prop_vars, spatial_idx_name)
        #update the list of x_vars
        x_vars = list(x_data.data_vars)
    # make sure we don't have any weird or missing input values
    check_if_finite(x_data)
    x_trn, x_val, x_tst = separate_trn_tst(
        x_data,
        time_idx_name,
        train_start_date,
        train_end_date,
        val_start_date,
        val_end_date,
        test_start_date,
        test_end_date,
    )

    x_trn_scl, x_std, x_mean = scale(x_trn)

    x_scl, _, _ = scale(x_data,std=x_std,mean=x_mean)

    if x_val:
        x_val_scl, _, _ = scale(x_val, std=x_std, mean=x_mean)
    else:
        x_val_scl = None

    if x_tst:
        x_tst_scl, _, _ = scale(x_tst, std=x_std, mean=x_mean)
    else:
        x_tst_scl = None

    # read, filter observations for finetuning

    x_data_dict = {
        "x_pre_full": convert_batch_reshape(
            x_scl,spatial_idx_name, time_idx_name, seq_len=seq_len,
            offset=trn_offset,
        ),
        "x_trn": convert_batch_reshape(
            x_trn_scl, spatial_idx_name, time_idx_name, seq_len=seq_len,
            offset=trn_offset,
        ),
        "x_val": convert_batch_reshape(
            x_val_scl,
            spatial_idx_name,
            time_idx_name,
            offset=tst_val_offset,
            seq_len=seq_len,
        ),
        "x_tst": convert_batch_reshape(
            x_tst_scl,
            spatial_idx_name,
            time_idx_name,
            offset=tst_val_offset,
            seq_len=seq_len,
        ),
        "x_std": x_std.to_array().values,
        "x_mean": x_mean.to_array().values,
        "x_vars": np.array(x_vars),
        "ids_trn": coord_as_reshaped_array(
            x_trn,
            spatial_idx_name,
            spatial_idx_name,
            time_idx_name,
            offset=trn_offset,
            seq_len=seq_len,
        ),
        "times_trn": coord_as_reshaped_array(
            x_trn,
            time_idx_name,
            spatial_idx_name,
            time_idx_name,
            offset=trn_offset,
            seq_len=seq_len,
        ),
        "ids_val": coord_as_reshaped_array(
            x_val,
            spatial_idx_name,
            spatial_idx_name,
            time_idx_name,
            offset=tst_val_offset,
            seq_len=seq_len,
        ),
        "times_val": coord_as_reshaped_array(
            x_val,
            time_idx_name,
            spatial_idx_name,
            time_idx_name,
            offset=tst_val_offset,
            seq_len=seq_len,
        ),
        "ids_tst": coord_as_reshaped_array(
            x_tst,
            spatial_idx_name,
            spatial_idx_name,
            time_idx_name,
            offset=tst_val_offset,
            seq_len=seq_len,
        ),
        "times_tst": coord_as_reshaped_array(
            x_tst,
            time_idx_name,
            spatial_idx_name,
            time_idx_name,
            offset=tst_val_offset,
            seq_len=seq_len,
        ),
    }
    if distfile:
        x_data_dict["dist_matrix"] = prep_adj_matrix(
            infile=distfile,
            dist_type=dist_type,
            dist_idx_name=dist_idx_name,
            segs=segs,
        )

    y_obs_data = {}
    y_pre_data = {}
    if y_data_file:
        y_obs_data = prep_y_data(
            y_data_file=y_data_file,
            y_vars=y_vars_finetune,
            x_data=x_data,
            train_start_date=train_start_date,
            train_end_date=train_end_date,
            val_start_date=val_start_date,
            val_end_date=val_end_date,
            test_start_date=test_start_date,
            test_end_date=test_end_date,
            val_sites=val_sites,
            test_sites=test_sites,
            spatial_idx_name=spatial_idx_name,
            time_idx_name=time_idx_name,
            seq_len=seq_len,
            log_vars=log_y_vars,
            exclude_file=exclude_file,
            normalize_y=normalize_y,
            y_type="obs",
            trn_offset = trn_offset,
            tst_val_offset = tst_val_offset
        )
        # if there is a y_data_file and a pretrain file, use the observation
        # mean and standard deviation to do the scaling/centering
        if pretrain_file:
            y_pre_data = prep_y_data(
                y_data_file=pretrain_file,
                y_vars=y_vars_pretrain,
                x_data=x_data,
                train_start_date=train_start_date,
                train_end_date=train_end_date,
                val_start_date=val_start_date,
                val_end_date=val_end_date,
                test_start_date=test_start_date,
                test_end_date=test_end_date,
                spatial_idx_name=spatial_idx_name,
                time_idx_name=time_idx_name,
                seq_len=seq_len,
                log_vars=log_y_vars,
                exclude_file=exclude_file,
                normalize_y=normalize_y,
                y_type="pre",
                y_std=y_obs_data["y_std"],
                y_mean=y_obs_data["y_mean"],
                trn_offset = trn_offset,
                tst_val_offset = tst_val_offset
            )
    # if there is no observation file, use the pretrain mean and standard dev
    # to do the scaling/centering
    elif pretrain_file and not y_obs_data:
        y_pre_data = prep_y_data(
            y_data_file=pretrain_file,
            y_vars=y_vars_pretrain,
            x_data=x_data,
            train_start_date=train_start_date,
            train_end_date=train_end_date,
            val_start_date=val_start_date,
            val_end_date=val_end_date,
            test_start_date=test_start_date,
            test_end_date=test_end_date,
            spatial_idx_name=spatial_idx_name,
            time_idx_name=time_idx_name,
            seq_len=seq_len,
            log_vars=log_y_vars,
            exclude_file=exclude_file,
            normalize_y=normalize_y,
            y_type="pre",
            trn_offset = trn_offset,
            tst_val_offset = tst_val_offset
        )
    else:
        raise Warning("No y_dataset data was provided")

    all_data = {**x_data_dict, **y_obs_data, **y_pre_data}
    if out_file:
        np.savez_compressed(out_file, **all_data)
    return all_data

def sort_dist_matrix(mat, row_col_names, segs=None):
    """
    sort the distance matrix by id
    :return:
    """
    if segs is not None:
        row_col_names = row_col_names.astype(type(segs[0]))
    df = pd.DataFrame(mat, columns=row_col_names, index=row_col_names)
    if segs:
        df = df[segs]
        df = df.loc[segs]
    df = df.sort_index(axis=0)
    df = df.sort_index(axis=1)
    return df

def prep_adj_matrix(infile, dist_type, dist_idx_name, segs=None, out_file=None):
    """
    process adj matrix.
    **The resulting matrix is sorted by id **
    :param infile: [str] path to the distance matrix .npz file
    :param dist_type: [str] type of distance matrix ("upstream", "downstream" or
    "updown")
    :param dist_idx_name: [str] name of index to sort dist_matrix by. This is
    the name of an array in the distance matrix .npz file
    :param segs: [list-like] which segments to prepare the data for
    :param out_file: [str] path to save the .npz file to
    :return: [numpy array] processed adjacency matrix
    """
    adj_matrices = np.load(infile)
    adj = adj_matrices[dist_type]
    adj = sort_dist_matrix(adj, adj_matrices[dist_idx_name], segs=segs)
    adj = np.where(np.isinf(adj), 0, adj)
    adj = -adj
    mean_adj = np.mean(adj[adj != 0])
    std_adj = np.std(adj[adj != 0])
    adj[adj != 0] = adj[adj != 0] - mean_adj
    adj[adj != 0] = adj[adj != 0] / std_adj
    adj[adj != 0] = 1 / (1 + np.exp(-adj[adj != 0]))

    I = np.eye(adj.shape[0])
    A_hat = adj.copy() + I
    D = np.sum(A_hat, axis=1)
    D_inv = D ** -1.0
    D_inv = np.diag(D_inv)
    A_hat = np.matmul(D_inv, A_hat)
    if out_file:
        np.savez_compressed(out_file, dist_matrix=A_hat)
    return A_hat

def read_exclude_segs_file(exclude_file):
    """
    read the exclude segs file. should be a yml file with start_date and list of
    segments to exclude
    --
    example exclude file:

    group_after_2017:
        start_date: "2017-10-01"
        variable: "temp"
        seg_id_nats_ex:
            - 1556
            - 1569
    group_2018_water_year:
        start_date: "2017-10-01"
        end_date: "2018-10-01"
        seg_id_nats_ex:
            - 1653
    group_all_time:
        seg_id_nats_in:
            - 1806
            - 2030

    --
    :param exclude_file: [str] exclude segs file
    :return: [list] list of dictionaries of segments to exclude. dict keys must
    have 'seg_id_nats' and may also have 'start_date' and 'end_date'
    """
    with open(exclude_file, "r") as s:
        d = yaml.safe_load(s)
    return [val for key, val in d.items()]


import logging.config

import hydra
import numpy as np
import pandas as pd
from omegaconf import DictConfig, OmegaConf

from data import read_data
from entities import PredictionPipelineParams, read_prediction_pipeline_params
from features import make_features, load_transformer
from models import (
    predict_model,
    load_model,
)
from utils import save_config

logger = logging.getLogger(__name__)

def save_prediction(predictions: np.ndarray, output_path: str) -> str:
    pd.Series(predictions).to_csv(output_path, index=False)
    return output_path

def predict_pipeline(params: PredictionPipelineParams):
    logger.info(f"Start prediction.")

    data = read_data(params.data_path)
    logger.info(f"Data loaded. Raw data shape: {data.shape}")

    pipeline = load_transformer(params.transformer_path)
    logger.info(f"Transformer loaded: {pipeline}")

    model = load_model(params.model_path)
    logger.info(f"Model loaded: {model}")

    train_features = make_features(pipeline, data)
    logger.info(f"Test features shape: {train_features.shape}")

    predictions = predict_model(train_features, model)
    predictions_path = save_prediction(predictions, params.output_path)
    logger.info(f"Predictions saved in {predictions_path}")

@hydra.main(config_path="../configs/", config_name="config_pred.yaml")
def predict_pipeline_command(cfg: DictConfig) -> None:
    cfg_yaml = OmegaConf.to_yaml(cfg)
    save_config(cfg_yaml)
    logger.info(f"Config params:")
    logger.info(cfg_yaml)

    params = read_prediction_pipeline_params(cfg)
    predict_pipeline(params)

if __name__ == "__main__":
    predict_pipeline_command()


import json
from os import listdir
from os.path import isfile, join, realpath, dirname
from models import Episode, Season, Serie

def getSerieInfo(file_path):
	try:
		file = open((file_path), 'r')
		data = json.loads(file.read())
		file.close()
		if data["video"]["type"] != "show":
			return False
		ret_json = {
			'serie_name': data['video']['title'].upper(),
			'seasons': data['video']['seasons']
		}
		return ret_json
	except (IOError, KeyError, ValueError) as e:
		return False

#dirpath = join(dirname(realpath(__file__)), 'series')
#series = [f for f in listdir(dirpath) if isfile(join(dirpath, f))]

#for i in xrange(0, len(series)):

#print counting
#print len(series)


import os
import logging
import sys
import json
from typing import Tuple, List

import pandas as pd
import torch
from sklearn import metrics
from minio import Minio
from minio.error import S3Error
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

from net import model, optimizer

LOG_FORMAT = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
FORMATTER = logging.Formatter(LOG_FORMAT)

logger = logging.getLogger()
logger.setLevel("INFO")
stdout_handler = logging.StreamHandler(sys.stdout)
stdout_handler.setFormatter(FORMATTER)
file_handler = logging.FileHandler("container.log")
file_handler.setFormatter(FORMATTER)
logger.addHandler(stdout_handler)
logger.addHandler(file_handler)

def download_and_prepare_data() -> Tuple[torch.Tensor, torch.Tensor]:
    """Preparation function

    :return: test data and labels
    :rtype: Tuple[torch.Tensor, torch.Tensor]
    """

    client = Minio(
        endpoint=os.getenv("MINIO_ADDRESS"),
        secure=False,
        access_key=os.getenv("MINIO_ROOT_USER"),
        secret_key=os.getenv("MINIO_ROOT_PASSWORD"),
    )
    logger.info("Connected to Minio object storage.")
    try:
        client.bucket_exists(os.getenv("RUN_ID"))
    except S3Error as exc:
        logger.error("Bucket does not exist!")
        raise exc

    client.fget_object(
        os.getenv("RUN_ID"),
        "data_prep/artifacts/test/test.csv",
        "./data/test.csv",
    )
    client.fget_object(
        os.getenv("RUN_ID"),
        "train-eval/artifacts/train/model.pth",
        "./data/model.pth",
    )
    dataset = pd.read_csv("./data/test.csv")
    x = dataset.drop(["species", "labels"], axis=1).values
    y = dataset["labels"].values
    x = torch.FloatTensor(x)
    y = torch.LongTensor(y)

    return x, y

def test(x_test: torch.Tensor, y_test: torch.Tensor) -> Tuple[List[float], List[float]]:
    """Train function

    :param x_test: test data
    :type x_test: torch.Tensor
    :param y_test: test labels
    :type y_test: torch.Tensor
    :return truth and predicted labels
    :rtype: Tuple[List[float], List[float]]
    """

    y_hat_list = []
    y_list = []
    checkpoint = torch.load("/container_files/data/model.pth")
    model.load_state_dict(checkpoint["model_state_dict"])
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    model.eval()
    with torch.no_grad():
        for val, y in zip(x_test, y_test):
            y_hat = model(val)
            y_list.append(y.item())
            y_hat_list.append(y_hat.argmax().item())
    logger.debug("Predictions: %s", y_hat_list)
    logger.debug("Truth: %s", y_list)

    return y_list, y_hat_list

def compute_metrics(
    conf: dict, true_labels: List[float], predicted_labels: List[float]
) -> None:
    """Calculate evaluation metrics (accuracy, precision, recall, F1-score, confusion matrix)

    :param conf: full configuration
    :type conf: dict
    :param true_labels: list of ground truth labels
    :type true_labels: List[float]
    :param predicted_labels: list of predicted labels
    :type predicted_labels: List[float]
    """

    labels_to_classes_mapping = {
        v: k for k, v in conf["data_prep"].get("mapping").items()
    }
    true_classes = [labels_to_classes_mapping[item] for item in true_labels]
    predicted_classes = [labels_to_classes_mapping[item] for item in predicted_labels]
    logger.info("Computing classification metrics.")
    accuracy = metrics.accuracy_score(true_classes, predicted_classes)
    macro_averaged_precision = metrics.precision_score(
        true_classes, predicted_classes, average="macro"
    )
    macro_averaged_recall = metrics.recall_score(
        true_classes, predicted_classes, average="macro"
    )
    macro_averaged_f1_score = metrics.f1_score(
        true_classes, predicted_classes, average="macro"
    )
    metrics_dict = {
        "Accuracy": accuracy,
        "Precision": macro_averaged_precision,
        "Recall": macro_averaged_recall,
        "F1-score": macro_averaged_f1_score,
    }
    metrics_df = pd.DataFrame(metrics_dict, index=[0])
    metrics_df.to_csv(
        os.path.join(conf["train_eval"].get("metrics_dir"), "metrics.csv"),
        encoding="utf-8",
        index=False,
    )
    plt.figure(figsize=(18, 8))
    sns.heatmap(
        metrics.confusion_matrix(true_classes, predicted_classes),
        annot=True,
        xticklabels=np.unique(true_classes),
        yticklabels=np.unique(true_classes),
        cmap="summer",
    )
    plt.xlabel("Predicted Classes")
    plt.ylabel("True Classes")
    plt.savefig(os.path.join(conf["train_eval"].get("metrics_dir"), "CM.png"))

def push_data_to_minio(conf: dict) -> None:
    """Push checkpoints and the container log to Minio.

    :param conf: configuration of the train and eval stages
    :type conf: dict
    """

    client = Minio(
        endpoint=os.getenv("MINIO_ADDRESS"),
        secure=False,
        access_key=os.getenv("MINIO_ROOT_USER"),
        secret_key=os.getenv("MINIO_ROOT_PASSWORD"),
    )
    logger.info("Connected to Minio object storage.")
    try:
        client.bucket_exists(os.getenv("RUN_ID"))
    except S3Error as exc:
        logger.error("Bucket does not exist!")
        raise exc

    logger.info("Uploading new files to Minio...")
    for file in os.listdir(conf.get("metrics_dir")):
        client.fput_object(
            os.getenv("RUN_ID"),
            f"train-eval/artifacts/test/{file}",
            f"{os.path.join(conf.get('metrics_dir'), file)}",
        )
    client.fput_object(
        os.getenv("RUN_ID"),
        "train-eval/logs/test/container.log",
        "container.log",
    )
    logger.info("Upload successful!")

def main(conf: dict) -> None:
    """Training main function

    :param conf: configuration of the train and eval stages
    :type conf: dict
    """

    torch.manual_seed(42)
    if not os.path.exists(conf["train_eval"].get("metrics_dir")):
        os.makedirs(conf["train_eval"].get("metrics_dir"), exist_ok=True)
    x_test, y_test = download_and_prepare_data()
    y_list, y_hat_list = test(x_test, y_test)
    compute_metrics(conf, y_list, y_hat_list)
    push_data_to_minio(conf["train_eval"])

if __name__ == "__main__":
    if os.getenv("CONFIG_FILE_PATH") is None:
        logger.error("Configuration file path is missing!")
        sys.exit(0)
    else:
        with open(os.getenv("CONFIG_FILE_PATH"), "r") as json_file:
            config = json.load(json_file)
        main(config)


# -*- coding: utf-8 -*-
"""
Created on Mon Sep 28 19:31:23 2020

@author: Gustavo Kildegaard
@mail: gustavo.kildegaard@gmail.com
"""
# %%
import datetime
# %%
fecha_hora = datetime.datetime.now()
print(fecha_hora)
# %%
fecha = datetime.date.today()
print(fecha)
# %%
nacim_gus = datetime.date(year=1988, day=2, month=6)
nacim_agus = datetime.date(year=1990, day=18, month=3)
hoy = datetime.date.today()
dif1 = hoy - nacim_gus
dif2 = hoy - nacim_agus
años1 = (dif1 / 365.2425).days
meses1 = (dif1.days / 365.2425 - (dif1 / 365.2425).days) * 12
años2 = (dif2 / 365.2425).days
meses2 = (dif2.days / 365.2425 - (dif2 / 365.2425).days) * 12
print(f'Gus: {años1} años y {meses1:.0f} meses.')
print(f'Agus: {años2} años y {meses2:.0f} meses.')
# %%

def edad(dia: int, mes: int, anio: int) -> list:
    import datetime
    nacim = datetime.date(year=anio, day=dia, month=mes)
    hoy = datetime.date.today()
    dif = hoy - nacim
    edad = []
    edad.append((dif / 365.2425).days)
    edad.append(int((dif.days / 365.2425 - (dif / 365.2425).days) * 12))
    return edad

def print_edad(nombre: str, edad: list):
    print(f'{nombre}: {edad[0]} años y {edad[1]} meses.')

# %%
hoy = datetime.datetime.now()
fecha_random = datetime.datetime(1990, 3, 18, 15, 0, 0)
fecha_random - hoy
# %%
from datetime import *

a = time(0, 30, 0)
b = time(0, 45, 0)


import weakref
from . import workers
import time
import json
import logging
import cherrypy
from . import messagebus, pages, auth, widgets
from .unitsofmeasure import strftime
from .config import config

mlogger = logging.getLogger("system.msgbuslog")

logger = logging.getLogger("system.notifications")
ilogger = logging.getLogger("system.notifications.important")

notificationslog = []

class API(widgets.APIWidget):
    def onNewSubscriber(self, user, cid, **kw):
        self.send(['all', notificationslog])

api = API()
api.require("/admin/mainpage.view")

toolbarapi=widgets.APIWidget()
toolbarapi.require("/admin/mainpage.view")
toolbarapi.echo=False

def f(u,v,id):
    if v[0]=='countsince':
        toolbarapi.sendTo(json.dumps(countnew(v[1])),id)

toolbarapi.attach2(f)

def countnew(since):
    normal = 0
    errors = 0
    warnings = 0
    total = 0
    x = list(notificationslog)
    x.reverse()
    for i in x:
        if not i[0] > since:
            break
        else:
            if 'warning' in i[1]:
                warnings += 1
            elif 'error' in i[1]:
                errors += 1
            else:
                normal += 1
            total += 1
    return [total, normal, warnings, errors]

class WI():
    @cherrypy.expose
    def countnew(self, **kwargs):
        pages.require('/admin/mainpage.view')
        return json.dumps(countnew(float(kwargs['since'])))

    @cherrypy.expose
    def mostrecent(self, **kwargs):
        pages.require('/admin/mainpage.view')
        return json.dumps(notificationslog[-int(kwargs['count']):])

def doPlyer(t,m):
    try:
        try:
            from plyer import notification
        except ImportError:
            return
        n=notification

        n.notify(title='Kaithem '+t, message=m[:140], ticker='')
    except Exception:
        logger.exception("Could not do the notification")

epochAndRemaining = [0,15]

def subscriber(topic, message):
    global notificationslog
    notificationslog.append((time.time(), topic, message))
    # Delete all but the most recent N notifications, where N is from the config file.
    notificationslog = notificationslog[-config['notifications-to-keep']:]

    # TODO:
    # Not threadsafe. But it is still better than the old polling based system.
    try:
        toolbarapi.send(["newmsg"])
    except:
        logging.exception("Error pushing notifications")

    api.send(['notification', [time.time(), topic, message]])

    if 'error' in topic or 'warning' in topic:

        # Add allowed notifications at a rate of  under 1 per miniute up to 15 "stored credits"
        epochAndRemaining[1] = max((time.monotonic()-epochAndRemaining[0])/240 + epochAndRemaining[1], 15)
        epochAndRemaining[0]=time.monotonic()

        if epochAndRemaining[1] > 1:
            epochAndRemaining[1] -= 1

            def f():
                doPlyer(topic.split("/")[-1],message)
            workers.do(f)

messagebus.subscribe('/system/notifications/#', subscriber)

def printer(t, m):
    if 'error' in t:
        logger.error(t+':'+m)
    elif 'warning' in t:
        logger.warning(t+':'+m)
    elif 'important' in t:
        ilogger.info(t+':'+m)
    else:
        logger.info(t+':'+m)

messagebus.subscribe('/system/notifications/#', printer)

def mprinter(t, m):
    if 'error' in t:
        mlogger.error(t+':'+m)
    elif 'warning' in t:
        mlogger.warning(t+':'+m)
    elif 'important' in t:
        mlogger.info(t+':'+m)
    else:
        mlogger.info(t+':'+m)

for i in config['print-topics']:
    messagebus.subscribe(i, mprinter)