"""This module defines a player class with the OpenAI API on the main thread.
For a black-box implementation consider using the module env_player.
"""
from __future__ import annotations

import asyncio
import copy
import random
import time
from abc import ABC, abstractmethod
from logging import Logger
from typing import (
    Any,
    Awaitable,
    Callable,
    Dict,
    Generic,
    List,
    Optional,
    Tuple,
    TypeVar,
    Union,
)

from gym.core import ActType, Env, ObsType
from gym.spaces import Discrete, Space

from poke_env.concurrency import POKE_LOOP, create_in_poke_loop
from poke_env.environment.abstract_battle import AbstractBattle
from poke_env.player.battle_order import BattleOrder, ForfeitBattleOrder
from poke_env.player.player import Player
from poke_env.ps_client import AccountConfiguration
from poke_env.ps_client.server_configuration import (
    LocalhostServerConfiguration,
    ServerConfiguration,
)
from poke_env.teambuilder.teambuilder import Teambuilder

BattleType = TypeVar("BattleType", bound=AbstractBattle)

class _AsyncQueue:
    def __init__(self, queue: asyncio.Queue[Any]):
        self.queue = queue

    async def async_get(self):
        return await self.queue.get()

    def get(self):
        res = asyncio.run_coroutine_threadsafe(self.queue.get(), POKE_LOOP)
        return res.result()

    async def async_put(self, item: Any):
        await self.queue.put(item)

    def put(self, item: Any):
        task = asyncio.run_coroutine_threadsafe(self.queue.put(item), POKE_LOOP)
        task.result()

    def empty(self):
        return self.queue.empty()

    def join(self):
        task = asyncio.run_coroutine_threadsafe(self.queue.join(), POKE_LOOP)
        task.result()

    async def async_join(self):
        await self.queue.join()

class _AsyncPlayer(Generic[ObsType, ActType], Player):
    actions: _AsyncQueue
    observations: _AsyncQueue

    def __init__(
        self,
        user_funcs: OpenAIGymEnv[ObsType, ActType],
        username: str,
        **kwargs: Any,
    ):
        self.__class__.__name__ = username
        super().__init__(**kwargs)
        self.__class__.__name__ = "_AsyncPlayer"
        self.observations = _AsyncQueue(create_in_poke_loop(asyncio.Queue, 1))
        self.actions = _AsyncQueue(create_in_poke_loop(asyncio.Queue, 1))
        self.current_battle: Optional[AbstractBattle] = None
        self._user_funcs = user_funcs

    def choose_move(
        self, battle: AbstractBattle
    ) -> Union[BattleOrder, Awaitable[BattleOrder]]:
        return self._env_move(battle)

    async def _env_move(self, battle: AbstractBattle):
        if not self.current_battle or self.current_battle.finished:
            self.current_battle = battle
        if not self.current_battle == battle:
            raise RuntimeError("Using different battles for queues")
        battle_to_send = self._user_funcs.embed_battle(battle)
        await self.observations.async_put(battle_to_send)
        action = await self.actions.async_get()
        if action == -1:
            return ForfeitBattleOrder()
        return self._user_funcs.action_to_move(action, battle)

    def _battle_finished_callback(self, battle: AbstractBattle):
        to_put = self._user_funcs.embed_battle(battle)
        asyncio.run_coroutine_threadsafe(self.observations.async_put(to_put), POKE_LOOP)

class _ABCMetaclass(type(ABC)):
    pass

class _EnvMetaclass(type(Env)):
    pass

class _OpenAIGymEnvMetaclass(_EnvMetaclass, _ABCMetaclass):
    pass

class OpenAIGymEnv(
    Env[ObsType, ActType],
    ABC,
    metaclass=_OpenAIGymEnvMetaclass,
):
    """
    Base class implementing the OpenAI Gym API on the main thread.
    """

    _INIT_RETRIES = 100
    _TIME_BETWEEN_RETRIES = 0.5
    _SWITCH_CHALLENGE_TASK_RETRIES = 30
    _TIME_BETWEEN_SWITCH_RETIRES = 1

    def __init__(
        self,
        account_configuration: Optional[AccountConfiguration] = None,
        *,
        avatar: Optional[int] = None,
        battle_format: str = "gen8randombattle",
        log_level: Optional[int] = None,
        save_replays: Union[bool, str] = False,
        server_configuration: Optional[
            ServerConfiguration
        ] = LocalhostServerConfiguration,
        start_timer_on_battle_start: bool = False,
        start_listening: bool = True,
        ping_interval: Optional[float] = 20.0,
        ping_timeout: Optional[float] = 20.0,
        team: Optional[Union[str, Teambuilder]] = None,
        start_challenging: bool = False,
    ):
        """
        :param account_configuration: Player configuration. If empty, defaults to an
            automatically generated username with no password. This option must be set
            if the server configuration requires authentication.
        :type account_configuration: AccountConfiguration, optional
        :param avatar: Player avatar id. Optional.
        :type avatar: int, optional
        :param battle_format: Name of the battle format this player plays. Defaults to
            gen8randombattle.
        :type battle_format: Optional, str. Default to randombattles, with specifics
            varying per class.
        :param log_level: The player's logger level.
        :type log_level: int. Defaults to logging's default level.
        :param save_replays: Whether to save battle replays. Can be a boolean, where
            True will lead to replays being saved in a potentially new /replay folder,
            or a string representing a folder where replays will be saved.
        :type save_replays: bool or str
        :param server_configuration: Server configuration. Defaults to Localhost Server
            Configuration.
        :type server_configuration: ServerConfiguration, optional
        :param start_listening: Whether to start listening to the server. Defaults to
            True.
        :type start_listening: bool
        :param start_timer_on_battle_start: Whether to automatically start the battle
            timer on battle start. Defaults to False.
        :type start_timer_on_battle_start: bool
        :param ping_interval: How long between keepalive pings (Important for backend
            websockets). If None, disables keepalive entirely.
        :type ping_interval: float, optional
        :param ping_timeout: How long to wait for a timeout of a specific ping
            (important for backend websockets.
            Increase only if timeouts occur during runtime).
            If None pings will never time out.
        :type ping_timeout: float, optional
        :param team: The team to use for formats requiring a team. Can be a showdown
            team string, a showdown packed team string, of a ShowdownTeam object.
            Defaults to None.
        :type team: str or Teambuilder, optional
        :param start_challenging: Whether to automatically start the challenge loop or
            leave it inactive.
        :type start_challenging: bool
        """
        self.agent = _AsyncPlayer(
            self,
            username=self.__class__.__name__,  # type: ignore
            account_configuration=account_configuration,
            avatar=avatar,
            battle_format=battle_format,
            log_level=log_level,
            max_concurrent_battles=1,
            save_replays=save_replays,
            server_configuration=server_configuration,
            start_timer_on_battle_start=start_timer_on_battle_start,
            start_listening=start_listening,
            ping_interval=ping_interval,
            ping_timeout=ping_timeout,
            team=team,
        )
        self._actions = self.agent.actions
        self._observations = self.agent.observations
        self.action_space = Discrete(self.action_space_size())  # type: ignore
        self.observation_space = self.describe_embedding()
        self.current_battle: Optional[AbstractBattle] = None
        self.last_battle: Optional[AbstractBattle] = None
        self._keep_challenging: bool = False
        self._challenge_task = None
        self._seed_initialized: bool = False
        if start_challenging:
            self._keep_challenging = True
            self._challenge_task = asyncio.run_coroutine_threadsafe(
                self._challenge_loop(), POKE_LOOP
            )

    @abstractmethod
    def calc_reward(
        self, last_battle: AbstractBattle, current_battle: AbstractBattle
    ) -> float:
        """
        Returns the reward for the current battle state. The battle state in the previous
        turn is given as well and can be used for comparisons.

        :param last_battle: The battle state in the previous turn.
        :type last_battle: AbstractBattle
        :param current_battle: The current battle state.
        :type current_battle: AbstractBattle

        :return: The reward for current_battle.
        :rtype: float
        """
        pass

    @abstractmethod
    def action_to_move(self, action: int, battle: AbstractBattle) -> BattleOrder:
        """
        Returns the BattleOrder relative to the given action.

        :param action: The action to take.
        :type action: int
        :param battle: The current battle state
        :type battle: AbstractBattle

        :return: The battle order for the given action in context of the current battle.
        :rtype: BattleOrder
        """
        pass

    @abstractmethod
    def embed_battle(self, battle: AbstractBattle) -> ObsType:
        """
        Returns the embedding of the current battle state in a format compatible with
        the OpenAI gym API.

        :param battle: The current battle state.
        :type battle: AbstractBattle

        :return: The embedding of the current battle state.
        """
        pass

    @abstractmethod
    def describe_embedding(self) -> Space[ObsType]:
        """
        Returns the description of the embedding. It must return a Space specifying
        low bounds and high bounds.

        :return: The description of the embedding.
        :rtype: Space
        """
        pass

    @abstractmethod
    def action_space_size(self) -> int:
        """
        Returns the size of the action space. Given size x, the action space goes
        from 0 to x - 1.

        :return: The action space size.
        :rtype: int
        """
        pass

    @abstractmethod
    def get_opponent(
        self,
    ) -> Union[Player, str, List[Player], List[str]]:
        """
        Returns the opponent (or list of opponents) that will be challenged
        on the next iteration of the challenge loop. If a list is returned,
        a random element will be chosen at random during the challenge loop.

        :return: The opponent (or list of opponents).
        :rtype: Player or str or list(Player) or list(str)
        """
        pass

    def _get_opponent(self) -> Union[Player, str]:
        opponent = self.get_opponent()
        random_opponent = (
            random.choice(opponent) if isinstance(opponent, list) else opponent
        )
        return random_opponent

    def reset(
        self,
        *,
        seed: Optional[int] = None,
        return_info: bool = False,
        options: Optional[Dict[str, Any]] = None,
    ) -> Tuple[ObsType, Dict[str, Any]]:
        if seed is not None:
            super().reset(seed=seed)  # type: ignore
            self._seed_initialized = True
        elif not self._seed_initialized:
            super().reset(seed=int(time.time()))  # type: ignore
            self._seed_initialized = True
        if not self.agent.current_battle:
            count = self._INIT_RETRIES
            while not self.agent.current_battle:
                if count == 0:
                    raise RuntimeError("Agent is not challenging")
                count -= 1
                time.sleep(self._TIME_BETWEEN_RETRIES)
        if self.current_battle and not self.current_battle.finished:
            if self.current_battle == self.agent.current_battle:
                self._actions.put(-1)
                self._observations.get()
            else:
                raise RuntimeError(
                    "Environment and agent aren't synchronized. Try to restart"
                )
        while self.current_battle == self.agent.current_battle:
            time.sleep(0.01)
        self.current_battle = self.agent.current_battle
        battle = copy.copy(self.current_battle)
        battle.logger = None
        self.last_battle = copy.deepcopy(battle)
        return self._observations.get(), self.get_additional_info()

    def get_additional_info(self) -> Dict[str, Any]:
        """
        Returns additional info for the reset method.
        Override only if you really need it.

        :return: Additional information as a Dict
        :rtype: Dict
        """
        return {}

    def step(
        self, action: ActType
    ) -> Tuple[ObsType, float, bool, bool, Dict[str, Any]]:
        if not self.current_battle:
            obs, info = self.reset(return_info=True)
            return obs, 0.0, False, False, info
        if self.current_battle.finished:
            raise RuntimeError("Battle is already finished, call reset")
        battle = copy.copy(self.current_battle)
        battle.logger = None
        self.last_battle = copy.deepcopy(battle)
        self._actions.put(action)
        observation = self._observations.get()
        reward = self.calc_reward(self.last_battle, self.current_battle)
        terminated = False
        truncated = False
        if self.current_battle.finished:
            size = self.current_battle.team_size
            remaining_mons = size - len(
                [mon for mon in self.current_battle.team.values() if mon.fainted]
            )
            remaining_opponent_mons = size - len(
                [
                    mon
                    for mon in self.current_battle.opponent_team.values()
                    if mon.fainted
                ]
            )
            if (remaining_mons == 0) != (remaining_opponent_mons == 0):
                terminated = True
            else:
                truncated = True
        return observation, reward, terminated, truncated, self.get_additional_info()

    def render(self, mode: str = "human"):
        if self.current_battle is not None:
            print(
                "  Turn %4d. | [%s][%3d/%3dhp] %10.10s - %10.10s [%3d%%hp][%s]"
                % (
                    self.current_battle.turn,
                    "".join(
                        [
                            "⦻" if mon.fainted else "●"
                            for mon in self.current_battle.team.values()
                        ]
                    ),
                    self.current_battle.active_pokemon.current_hp or 0,
                    self.current_battle.active_pokemon.max_hp or 0,
                    self.current_battle.active_pokemon.species,
                    self.current_battle.opponent_active_pokemon.species,
                    self.current_battle.opponent_active_pokemon.current_hp or 0,
                    "".join(
                        [
                            "⦻" if mon.fainted else "●"
                            for mon in self.current_battle.opponent_team.values()
                        ]
                    ),
                ),
                end="\n" if self.current_battle.finished else "\r",
            )

    def close(self, purge: bool = True):
        if self.current_battle is None or self.current_battle.finished:
            time.sleep(1)
            if self.current_battle != self.agent.current_battle:
                self.current_battle = self.agent.current_battle
        closing_task = asyncio.run_coroutine_threadsafe(
            self._stop_challenge_loop(purge=purge), POKE_LOOP
        )
        closing_task.result()

    def seed(self, seed: Optional[int] = None):
        random.seed(seed)

    def background_send_challenge(self, username: str):
        """
        Sends a single challenge specified player. The function immediately returns
        to allow use of the OpenAI gym API.

        :param username: The username of the player to challenge.
        :type username: str
        """
        if self._challenge_task and not self._challenge_task.done():
            raise RuntimeError(
                "Agent is already challenging opponents with the challenging loop. "
                "Try to specify 'start_challenging=True' during instantiation or call "
                "'await agent.stop_challenge_loop()' to clear the task."
            )
        self._challenge_task = asyncio.run_coroutine_threadsafe(
            self.agent.send_challenges(username, 1), POKE_LOOP
        )

    def background_accept_challenge(self, username: str):
        """
        Accepts a single challenge specified player. The function immediately returns
        to allow use of the OpenAI gym API.

        :param username: The username of the player to challenge.
        :type username: str
        """
        if self._challenge_task and not self._challenge_task.done():
            raise RuntimeError(
                "Agent is already challenging opponents with the challenging loop. "
                "Try to specify 'start_challenging=True' during instantiation or call "
                "'await agent.stop_challenge_loop()' to clear the task."
            )
        self._challenge_task = asyncio.run_coroutine_threadsafe(
            self.agent.accept_challenges(username, 1, self.agent.next_team), POKE_LOOP
        )

    async def _challenge_loop(
        self,
        n_challenges: Optional[int] = None,
        callback: Optional[Callable[[AbstractBattle], None]] = None,
    ):
        if not n_challenges:
            while self._keep_challenging:
                opponent = self._get_opponent()
                if isinstance(opponent, Player):
                    await self.agent.battle_against(opponent, 1)
                else:
                    await self.agent.send_challenges(opponent, 1)
                if callback and self.current_battle is not None:
                    callback(copy.deepcopy(self.current_battle))
        elif n_challenges > 0:
            for _ in range(n_challenges):
                opponent = self._get_opponent()
                if isinstance(opponent, Player):
                    await self.agent.battle_against(opponent, 1)
                else:
                    await self.agent.send_challenges(opponent, 1)
                if callback and self.current_battle is not None:
                    callback(copy.deepcopy(self.current_battle))
        else:
            raise ValueError(f"Number of challenges must be > 0. Got {n_challenges}")

    def start_challenging(
        self,
        n_challenges: Optional[int] = None,
        callback: Optional[Callable[[AbstractBattle], None]] = None,
    ):
        """
        Starts the challenge loop.

        :param n_challenges: The number of challenges to send. If empty it will run until
            stopped.
        :type n_challenges: int, optional
        :param callback: The function to callback after each challenge with a copy of
            the final battle state.
        :type callback: Callable[[AbstractBattle], None], optional
        """
        if self._challenge_task and not self._challenge_task.done():
            count = self._SWITCH_CHALLENGE_TASK_RETRIES
            while not self._challenge_task.done():
                if count == 0:
                    raise RuntimeError("Agent is already challenging")
                count -= 1
                time.sleep(self._TIME_BETWEEN_SWITCH_RETIRES)
        if not n_challenges:
            self._keep_challenging = True
        self._challenge_task = asyncio.run_coroutine_threadsafe(
            self._challenge_loop(n_challenges, callback), POKE_LOOP
        )

    async def _ladder_loop(
        self,
        n_challenges: Optional[int] = None,
        callback: Optional[Callable[[AbstractBattle], None]] = None,
    ):
        if n_challenges:
            if n_challenges <= 0:
                raise ValueError(
                    f"Number of challenges must be > 0. Got {n_challenges}"
                )
            for _ in range(n_challenges):
                await self.agent.ladder(1)
                if callback and self.current_battle is not None:
                    callback(copy.deepcopy(self.current_battle))
        else:
            while self._keep_challenging:
                await self.agent.ladder(1)
                if callback and self.current_battle is not None:
                    callback(copy.deepcopy(self.current_battle))

    def start_laddering(
        self,
        n_challenges: Optional[int] = None,
        callback: Optional[Callable[[AbstractBattle], None]] = None,
    ):
        """
        Starts the laddering loop.

        :param n_challenges: The number of ladder games to play. If empty it
            will run until stopped.
        :type n_challenges: int, optional
        :param callback: The function to callback after each challenge with a
            copy of the final battle state.
        :type callback: Callable[[AbstractBattle], None], optional
        """
        if self._challenge_task and not self._challenge_task.done():
            count = self._SWITCH_CHALLENGE_TASK_RETRIES
            while not self._challenge_task.done():
                if count == 0:
                    raise RuntimeError("Agent is already challenging")
                count -= 1
                time.sleep(self._TIME_BETWEEN_SWITCH_RETIRES)
        if not n_challenges:
            self._keep_challenging = True
        self._challenge_task = asyncio.run_coroutine_threadsafe(
            self._ladder_loop(n_challenges, callback), POKE_LOOP
        )

    async def _stop_challenge_loop(
        self, force: bool = True, wait: bool = True, purge: bool = False
    ):
        self._keep_challenging = False

        if force:
            if self.current_battle and not self.current_battle.finished:
                if not self._actions.empty():
                    await asyncio.sleep(2)
                    if not self._actions.empty():
                        raise RuntimeError(
                            "The agent is still sending actions. "
                            "Use this method only when training or "
                            "evaluation are over."
                        )
                if not self._observations.empty():
                    await self._observations.async_get()
                await self._actions.async_put(-1)

        if wait and self._challenge_task:
            while not self._challenge_task.done():
                await asyncio.sleep(1)
            self._challenge_task.result()

        self._challenge_task = None
        self.current_battle = None
        self.agent.current_battle = None
        while not self._actions.empty():
            await self._actions.async_get()
        while not self._observations.empty():
            await self._observations.async_get()

        if purge:
            self.agent.reset_battles()

    def reset_battles(self):
        """Resets the player's inner battle tracker."""
        self.agent.reset_battles()

    def done(self, timeout: Optional[int] = None) -> bool:
        """
        Returns True if the task is done or is done after the timeout, false otherwise.

        :param timeout: The amount of time to wait for if the task is not already done.
            If empty it will wait until the task is done.
        :type timeout: int, optional

        :return: True if the task is done or if the task gets completed after the
            timeout.
        :rtype: bool
        """
        if self._challenge_task is None:
            return True
        if timeout is None:
            self._challenge_task.result()
            return True
        if self._challenge_task.done():
            return True
        time.sleep(timeout)
        return self._challenge_task.done()

    # Expose properties of Player class

    @property
    def battles(self) -> Dict[str, AbstractBattle]:
        return self.agent.battles

    @property
    def format(self) -> str:
        return self.agent.format

    @property
    def format_is_doubles(self) -> bool:
        return self.agent.format_is_doubles

    @property
    def n_finished_battles(self) -> int:
        return self.agent.n_finished_battles

    @property
    def n_lost_battles(self) -> int:
        return self.agent.n_lost_battles

    @property
    def n_tied_battles(self) -> int:
        return self.agent.n_tied_battles

    @property
    def n_won_battles(self) -> int:
        return self.agent.n_won_battles

    @property
    def win_rate(self) -> float:
        return self.agent.win_rate

    # Expose properties of Player Network Interface Class

    @property
    def logged_in(self) -> asyncio.Event:
        """Event object associated with user login.

        :return: The logged-in event
        :rtype: Event
        """
        return self.agent.ps_client.logged_in

    @property
    def logger(self) -> Logger:
        """Logger associated with the player.

        :return: The logger.
        :rtype: Logger
        """
        return self.agent.logger

    @property
    def username(self) -> str:
        """The player's username.

        :return: The player's username.
        :rtype: str
        """
        return self.agent.username

    @property
    def websocket_url(self) -> str:
        """The websocket url.

        It is derived from the server url.

        :return: The websocket url.
        :rtype: str
        """
        return self.agent.ps_client.websocket_url

    def __getattr__(self, item: str):
        return getattr(self.agent, item)

class LegacyOpenAIGymEnv(OpenAIGymEnv[ObsType, ActType], ABC):
    """
    Subclass of OpenAIGymEnv compatible with the old gym API.
    If you need compatibility with the old gym API you should use the
    `wrap_for_old_gym_api` function.
    """

    def reset(
        self,
        *,
        seed: Optional[int] = None,
        return_info: bool = False,
        options: Optional[Dict[str, Any]] = None,
    ) -> Tuple[ObsType, Dict[str, Any]]:
        obs, info = super().reset(seed=seed, return_info=True, options=options)
        return obs, info

    def step(
        self, action: ActType
    ) -> Tuple[ObsType, float, bool, bool, Dict[str, Any]]:
        obs, reward, terminated, truncated, info = super().step(action)
        return obs, reward, terminated, truncated, info

class _OpenAIGymEnvWrapper(LegacyOpenAIGymEnv[ObsType, ActType]):
    def __init__(self, environment: OpenAIGymEnv[ObsType, ActType]):
        self._wrapped = environment
        self.step = super().step.__get__(self._wrapped, self._wrapped.__class__)
        self.reset = super().reset.__get__(self._wrapped, self._wrapped.__class__)
        self._instantiated = True

    def calc_reward(
        self, last_battle: AbstractBattle, current_battle: AbstractBattle
    ) -> float:
        return self._wrapped.calc_reward(last_battle, current_battle)

    def action_to_move(self, action: int, battle: AbstractBattle) -> BattleOrder:
        return self._wrapped.action_to_move(action, battle)

    def embed_battle(self, battle: AbstractBattle) -> ObsType:
        return self._wrapped.embed_battle(battle)

    def describe_embedding(self) -> Space[ObsType]:
        return self._wrapped.describe_embedding()

    def action_space_size(self) -> int:
        return self._wrapped.action_space_size()

    def get_opponent(self) -> Union[Player, str, List[Player], List[str]]:
        return self._wrapped.get_opponent()

    def __getattr__(self, item: str):
        if item == "_instantiated":
            return False
        return getattr(self._wrapped, item)

    def __setattr__(self, key: str, value: Any):
        if not self._instantiated:
            return super().__setattr__(key, value)
        return setattr(self._wrapped, key, value)

def wrap_for_old_gym_api(
    env: OpenAIGymEnv[ObsType, ActType]
) -> LegacyOpenAIGymEnv[ObsType, ActType]:
    """
    Wraps an OpenAIGymEnv in order to support the old gym API.

    :param env: the environment to wrap.
    :type env: OpenAIGymEnv

    :return: The wrapped environment
    :rtype: OpenAIGymEnv
    """
    return _OpenAIGymEnvWrapper(env)


# 方法:让算法共用,方便重复使用
# def 方法名(参数,参数)
# 加法
# def add(s1,s2):
#     sum=s1+s2
#     # 将sum返回调用处
#     return sum
# a= add(1,2)
# print(a)

# 减法
def sub():
 s1=int(input('输入被减数'))
 s2=int(input('输入减数'))
 c=s1-s2
 return c


class Student:

    def __init__(self, m1, m2):
        self.m1 = m1
        self.m2 = m2

    def sum(self, a=None, b=None, c=None):
        add = 0
        if a!=None and b!=None and c!=None:
            add = a + b + c

        elif a!=None and b!=None:
            add = a + b

        else:
            add = a

        return add

s1 = Student(55, 96)

print(s1.sum(5))

class A:

    def show(self):
        print('A is Show')

class B(A):
    pass

class C(A):
    def show(self):
        print('C is Show')

a = A()
a.show()

b = B()
b.show()

c = C()
c.show()


import os
from typing import List
from matplotlib.axes import Axes
from matplotlib.ticker import FormatStrFormatter

figure_style = {
    # tight layout used for better visualisation
    'tight_layout': [0.05, 0.05, 0.95, 0.95],
    # xlimits used for the visualisation of non-expectation values
    'xlimits': (None, None),
    # ylimits used for the visualisation of non-expecation values
    'ylimits': (None, None),
    # format string for xticks
    'xtick_format': "%.0f",
    # format string for yticks
    'ytick_format': "%.1f",
    # scaling factor used in QuantDisplace2
    'scaling': 1
}

def scaling():
    if "FIGURE_SCALING" in os.environ:
        return float(os.environ["FIGURE_SCALING"])
    else:
        return 1

def set_scaling(scale: float):
    os.environ["FIGURE_SCALING"] = str(scale)

def set_axis_limits(axes: List[Axes], style: dict, direction: str):
    """
    Set new axes limits based on the 'xlimits' or 'ylimits' key given in
    the style dictionary for all axis elements passed in the
    collection of axis elements.

    Parameters
    ------------
        axes: List[matplotlib.Axes]
            Collection of matplotlib.Axes instances
        style: dict
            Style dictionary containing 'xlim' keyword
        direction: str
            Can either be 'x' or 'y'
    """
    assert(direction in ['x', 'y'], "Direction must be 'x' or 'y'")
    ax_min, ax_max = style['xlimits'] if direction == 'x' else style['ylimits']

    for ax in axes:
        if direction == 'x':
            ax.set_xlim(ax_min, ax_max)
        else:
            ax.set_ylim(ax_min, ax_max)

def axis_tick_formatting(axes: List[Axes], style: dict, direction: str):
    """
    Set the string formatting of the axis ticks.
    """
    assert(direction in ['x', 'y'], "Direction must be 'x' or 'y'")

    format, attr = (style['xtick_format'], 'xaxis') if direction == 'x' else (style['ytick_format'], 'yaxis')

    for ax in axes: getattr(ax, attr).set_major_formatter(FormatStrFormatter(format))

def apply_styles(axes: List[Axes], style: dict):
    style_keys = style.keys()

    if "xlimits" in style_keys:
        if not (None, None) ==  style["xlimits"]:
            set_axis_limits(axes, style, 'x')

    if "ylimits" in style_keys:
        if not (None, None) == style["ylimits"]:
            set_axis_limits(axes, style, 'y')

    if "xtick_format" in style_keys:
        if style["xtick_format"] is not None:
            axis_tick_formatting(axes, style, 'x')

    if "ytick_format" in style_keys:
        if style["ytick_format"] is not None:
            axis_tick_formatting(axes, style, 'y')

    if "scaling" in style_keys:
        if style["scaling"] is not None:
            set_scaling(style['scaling'])


#!/usr/bin/env python3

import pandas as pd
import numpy as np
import lightkurve as lk

#=================================
# EXOPLANET DATA FORMAT
# koi_period, koi_time0 (BKJD), koo_dur, koi_depth, koi_ror, koi_prad,
# COLUMN kepid:          KepID
# COLUMN koi_period:     Orbital Period [days]
# COLUMN koi_time0bk:    Transit Epoch [BKJD]
# COLUMN koi_duration:   Transit Duration [hrs]
# COLUMN koi_depth:      Transit Depth [ppm]
# COLUMN koi_ror:        Planet-Star Radius Ratio
# COLUMN koi_prad:       Planetary Radius [Earth radii]
# COLUMN koi_quarters:   Quarters
# COLUMN koi_disposition: Exoplanet Archive Disposition
#=================================
#https://exoplanetarchive.ipac.caltech.edu/cgi-bin/TblView/nph-tblView?app=ExoTbls&config=cumulative
# Description of the columns in the list of all known KOI - cumulative.tab

def threshold_positive(x):
    """Find outliers and replace them with samples from a Gaussian"""
    data = np.copy(x)
    std = np.nanstd(data)
    absdiff = np.abs(np.diff(data, prepend=data[0]))
    thresh = 3 * std

    outlier = np.zeros(len(data), dtype=bool)
    outlier[data > thresh] = True
    outlier[absdiff > thresh] = True

    # replace outliers with randomly sampled values
    std_inliers = np.nanstd(data[~outlier])
    data[outlier] = np.random.normal(0, std_inliers, size=outlier.sum())

    return data

def get_lightcurve(test=False):
    """Return a lightcurve and its associated data

    Args:
        test (bool): always return the same lightcurve so it can be cached by
            lightkurve

    Returns:
        lc (pandas DataFrame): phase-folded and cleaned lightcurve in a pandas dataframe.
            The relevant column is 'flux_clean'
        lc_info (pandas Series): contains information about the specific
            lightcurve such as Kepler ID ('kepid'), orbital period ('koi_period'),
            transit duration ('koi_duration'), planet/star ratio ('koi_ror'),
    """

    lc_sample_time = 29.424
    BKJD_start_time = 2454833 # offset of Kepler start time in BJD
    data = pd.read_csv('cumulative_2023.02.25_10.21.50.tab', sep='\t', dtype={'koi_quarters': str})
    data = data.dropna()
    data['num_quarters'] = data.koi_quarters.apply(lambda x: sum(map(int, list(x))))

    # filter out bad lightcurve candidates
    data = data[
        (data.koi_period > 4) &
        (data.koi_period < 12) &
        (data.koi_prad > 3) &
        (data.num_quarters>12)
    ]

    # choose a lightcurve
    if test:
        lc_info = pd.Series(data.iloc[0])
    else:
        lc_info = data.sample().squeeze()

    # download data for lightcurve and optionally cache it
    lc_files = lk.search_lightcurve(
        f"KIC {lc_info.kepid}", author='Kepler', cadence = "long"
    ).download_all(download_dir='.' if test else None)

    # insert extra info from lightkurve into lc_info before converting
    # lc_files to pandas tables
    lc_info['ra'] = lc_files[0].ra_obj
    lc_info['dec'] = lc_files[0].dec_obj
    lc_info['teff'] = lc_files[0].teff

    # import ipdb
    # ipdb.set_trace()

    lc_chunks = []
    for lcf in lc_files:
        lc_chunk = lcf.to_pandas()
        median = np.nanmedian(lc_chunk.flux)
        lc_chunk['flux_clean'] = (threshold_positive(lc_chunk.flux - median) + median) / median
        lc_chunks.append(lc_chunk)

    # stitch lightcurve data and filter out low quality data points
    lc = pd.concat(lc_chunks).query('quality == 0')

    # phase fold data points
    lc['folded_time'] = (
        lc.index - lc_info.koi_time0bk - lc_info.koi_period / 2
    ) % lc_info.koi_period
    lc = lc.sort_values('folded_time')

    # hours from middle of transit
    lc['transit_hours'] = (lc.folded_time - lc_info.koi_period / 2) * 24

    return lc, lc_info


import pygame
from projectil import projectile
import animation
# CREATION CLASSE JOUER
class player (animation.AnimateSprite):

    def __init__(self,game):
        super().__init__('player')
        self.game=game
        self.vie = 150
        self.max_vie = 150
        self.attack = 20
        self.vitesse = 5
        self.son = pygame.mixer.Sound("assets/sounds/tir.ogg")
        self.all_projectiles=pygame.sprite.Group()

        self.rect = self.image.get_rect()
        self.rect.x= 500
        self.rect.y = 600
    def launch_projectile(self):
        self.start_animation()
        self.son.play()

        self.all_projectiles.add(projectile(self))

    def deplacement_droite(self):

        if not self.game.check_collision(self,self.game.all_monstres):
            self.rect.x += self.vitesse
    def deplacement_gauche(self):
        self.rect.x -= self.vitesse

    def char_bar_vie(self, surface):

        pygame.draw.rect(surface, (63, 60, 60), [self.rect.x + 20, self.rect.y + 10, self.max_vie, 7])
        pygame.draw.rect(surface, (111, 210, 45),[self.rect.x + 20, self.rect.y +10, self.vie, 7])
    def domage(self,amount):

        if self.vie - amount > amount:
            self.vie -= amount
        else:
            self.game.game_over()
    def update_animation(self):
        self.animat()


from django import forms
from apps.presentacion.models import Presentacion

class PresentacionForm(forms.ModelForm):

    class Meta:
        model = Presentacion

        fields = "__all__"

        labels = {
            'cod_presentacion': 'Codigo de la presetacion',
            'nom_presentacion': 'Nombre de la presentacion',
        }

    def __init__(self, *args, **kwargs):
        super(PresentacionForm, self).__init__(*args, **kwargs)

        for field in self.fields:
            self.fields[field].widget.attrs.update({'class': 'form-control'})


"""
Implementaion for Beta VAE
"""
from distutils import log
from torch.utils import data
from carla_disentanglement.datasets.ground_truth_data import GroundTruthDataset
from carla_disentanglement.metrics.factor_vae_score import factor_vae_score
import os
import subprocess
from typing import Callable
from collections.abc import Sequence
from datetime import datetime
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.utils import save_image
from PIL import Image
from PIL.Image import Image as PILImage
import numpy as np
from tqdm import tqdm
from carla_disentanglement.metrics.beta_vae_score import beta_vae_score
from torch.utils.tensorboard import SummaryWriter
from disentanglement_lib.evaluation.metrics import beta_vae
from disentanglement_lib.evaluation.metrics import factor_vae
from disentanglement_lib.evaluation.metrics import dci
from disentanglement_lib.evaluation.metrics import mig
from disentanglement_lib.evaluation.metrics import sap_score

import gin.tf
gin.parse_config_file("metrics.gin")

class VAEModule(nn.Module):
    """
    pytorch Module for Variational Auto Encoders
    """

    def __init__(self, encoder: nn.Module, decoder: nn.Module, activation=nn.Sigmoid(), loss_activation: nn.Module = None):
        super().__init__()

        self.encoder = encoder
        self.decoder = decoder
        self.activation = activation
        self.use_model_regulariser = False
        self.loss_activation = loss_activation if loss_activation is not None else self.activation

    def encode(self, x: torch.Tensor):
        """Encodeds Input to latents

        Args:
            x (Tensor): input

        Returns:
            Tensor: latent encoding
        """
        return self.encoder(x)

    def decode(self, z):
        """Decondes latent z into output (with (sigmoid) activation)

        Args:
            z (Tensor): latent

        Returns:
            Tensor: output
        """
        return self.activation(self.decoder(z))

    def reparametrize(self, mu, logvar):
        """Reparamterization trick for training

        Args:
            mu (Tensor): mean of latents
            logvar (Tensor): log variance of latents

        Returns:
            Tensor: latent sample
        """
        std = logvar.mul(0.5).exp_()
        eps = std.data.new(std.size()).normal_()
        return eps.mul(std).add_(mu)

    def forward(self, x):
        """Forward pass from input to reconstruction (including reparamterization trick).
        Omits (sigmoid) activation for the use of BCEWithLogitsLoss

        Args:
            x (Tensor): input

        Returns:
            Tensor, (Tensor, Tensor): reconstruction, (mu, log_var), z
        """
        mu, logvar = self.encode(x)
        z = self.reparametrize(mu, logvar)
        # NOTE: we use decodeR(!) here not decode (which includes sigmoid activation)
        return self.loss_activation(self.decoder(z)), (mu, logvar), z

    def reconstruct(self, x):
        """Forward pass from input to reconstruction, without reparamterization trick.
        Only uses mu, ignores logvar

        Args:
            x (Tensor): input

        Returns:
            Tensor, (Tensor, Tensor): reconstruction, (mu, log_var)
        """
        mu, logvar = self.encode(x)
        return self.activation(self.decoder(mu)), (mu, logvar)

    def eval_forward(self, x):
        """Forward pass from input to reconstruction, without reparamterization trick.
        Only uses mu, ignores logvar
        Omits (sigmoid) activation for the use of BCEWithLogitsLoss

        Args:
            x (Tensor): input

        Returns:
            Tensor, (Tensor, Tensor): reconstruction, (mu, log_var)
        """
        mu, logvar = self.encode(x)
        return self.loss_activation(self.decoder(mu)), (mu, logvar)

    def model_regulariser(self):
        return 0

class BetaVAE(object):
    """
    Beta VAE model
    https://openreview.net/forum?id=Sy2fzU9gl
    """

    def __init__(self, model: VAEModule, beta=1.0, reconstruction='bce', tag=None, model_regulariser_weight=0.01, tag_suffix: str = None):

        self.model = model
        if reconstruction == 'mse_no_activation':
            self.reconstruction_loss = nn.MSELoss(reduction='sum')
            self.model.loss_activation = nn.Identity()
        elif reconstruction == 'mse':
            self.reconstruction_loss = nn.MSELoss(reduction='sum')
        else:
            self.reconstruction_loss = nn.BCEWithLogitsLoss(reduction='sum')
            self.model.loss_activation = nn.Identity()

        self.num_channels = model.encoder._num_channels
        self.z_dim = model.encoder._latent_dim
        self.image_size = model.encoder._image_size

        self.beta = beta
        self.model_regulariser_weight = model_regulariser_weight

        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print(self.device)
        self.model.to(self.device)

        tag = str(tag)+'_' if tag else ''
        self.full_tag = tag+str(self.__class__.__name__)+self.getParamString()
        if tag_suffix is not None:
            self.full_tag += "_" + tag_suffix
        print(self.full_tag, tag_suffix)
        self.num_log_values = 4

        self.initOptimizer()

    def getParamString(self):
        return "_beta"+str(self.beta)

    def gauss_vae_regulariser(self, mu: torch.Tensor, log_var: torch.Tensor):
        """
        Compute the regularisation factor for training a VAE
        with a Gaussian latent space.

        Parameters
        ----------
        mu: Tensor
            Mean predictions from variational encoder.
        log_var: Tensor
            log_var predictions from variational encoder.

        Returns
        -------
        kl_div : Tensor
            The Kullback-Leibler divergence between a Gaussian distribution
            with the given parameters and the standard normal distribution.
        log info : None
        """
        return ((torch.exp(log_var) + mu*mu - log_var - 1)/2).sum(1).mean(0)

    def _visualize_reconstruction(self, ds: torch.utils.data.Dataset, count=10, indices=None):
        """Create visualatin of VAE reconstruction

        Args:
            ds (torch.utils.data.Dataset): dataset
            count (int, optional): Number of sampels to be visualized. Defaults to 10.

        Returns:
            torch.Tensor: image
        """

        if indices != None:
            ds = torch.utils.data.Subset(ds, indices)
            count = len(indices)

        data_loader = DataLoader(ds, batch_size=count)
        xs = next(iter(data_loader))[0]

        with torch.no_grad():
            # get reconstructions
            preds = self.model.reconstruct(xs.to(self.device))[0].cpu()

            # concatenate all data
            x_cat = torch.cat(
                [torch.cat([x, torch.zeros(self.num_channels, self.image_size, 1)], dim=-1) for x in xs], dim=-1)
            pred_cat = torch.cat(
                [torch.cat([x, torch.zeros(self.num_channels, self.image_size, 1)], dim=-1) for x in preds], dim=-1)

            return torch.cat([x_cat, pred_cat], dim=1)

    def visualize_reconstruction(self, ds: torch.utils.data.Dataset, count=10, indices=None):
        """Create visualatin of VAE reconstruction
        How to use: display(image, metadata={'width': '100%'})

        Args:
            ds (torch.utils.data.Dataset): dataset
            count (int, optional): Number of sampels to be visualized. Defaults to 10.

        Returns:
            PIL.Image: image
        """

        image = self._visualize_reconstruction(ds, count, indices)

        to_image = transforms.Compose([
            # inverts normalization of image # implement unnormalize when we use normalization
            # transforms.Normalize(-means / stds, 1. / stds),
            transforms.Lambda(lambda x: torch.clamp(x, 0, 1)),
            transforms.ToPILImage()
        ])

        return to_image(image)

    def visualize_latents_static(self, ds: torch.utils.data.Dataset, index=0, limit=3):
        image = self._visualize_latents_static(ds, index, limit=limit)

        to_image = transforms.Compose([
            # inverts normalization of image # implement unnormalize when we use normalization
            # transforms.Normalize(-means / stds, 1. / stds),
            transforms.Lambda(lambda x: torch.clamp(x, 0, 1)),
            transforms.ToPILImage()
        ])

        return to_image(image)

    def _visualize_latents_static(self, ds: torch.utils.data.Dataset, index=0, limit=3):
        inter = 2/3
        interpolation = torch.arange(-limit, limit+0.1, inter)

        img = ds.__getitem__(index)[0].cuda().unsqueeze(0)
        z_org = self.model.encode(img)[0]
        z_dim = z_org.shape[-1]
        if len(z_org.shape) < 2:
            z_org = z_org.unsqueeze(0)

        samples = []
        for z_index in range(z_dim):
            z = z_org.clone()
            inters = []
            for val in interpolation:
                z[:, z_index] = val
                sample = self.model.decode(z).data
                inters.append(sample)
                inters.append(torch.full((1, sample.shape[1], sample.shape[2], 1), .5).cuda())

            samples.append(torch.cat(inters, -1))
            samples.append(torch.full((1, sample.shape[1], 1, samples[-1].shape[-1]), .5).cuda())

        samples = torch.cat(samples, -2).cpu()
        return samples.squeeze()

    def visualize_latents(self, ds: torch.utils.data.Dataset, index=0, output_dir='./viz/', delay=100):
        limit = 3
        inter = 2/3
        interpolation = torch.arange(-limit, limit+0.1, inter)

        img = ds.__getitem__(index)[0].cuda().unsqueeze(0)
        z_org = self.model.encode(img)[0]
        z_dim = z_org.shape[-1]
        if len(z_org.shape) < 2:
            z_org = z_org.unsqueeze(0)

        samples = []
        for z_index in range(z_dim):
            z = z_org.clone()
            for val in interpolation:
                z[:, z_index] = val
                sample = self.model.decode(z).data
                samples.append(sample)

        samples = torch.cat(samples).cpu()
        gif = samples.view(z_dim, len(interpolation), self.num_channels, self.image_size, self.image_size).transpose(0, 1)

        for j in range(gif.shape[0]):
            save_image(tensor=gif[j].cpu(),
                       fp=os.path.join(output_dir, '{}.png'.format(j)),
                       nrow=z_dim, pad_value=1)

        image_str = os.path.join(output_dir, '*.png')
        output_gif = os.path.join(output_dir, self.full_tag+'_'+datetime.now().strftime("%Y-%m-%d_%H-%M-%S")+'.gif')
        str1 = 'magick convert -delay '+str(delay)+' -loop 0 ' + image_str + ' ' + output_gif
        subprocess.call(str1, shell=True)

        return output_gif

    def initOptimizer(self, lr=1e-4, betas=None):
        """Initialize adam optimizer

        Args:
            lr (float, optional): Learning rate. Defaults to 1e-4.
            betas ((float,float), optional): betas param for Adam.
        """
        if betas is not None:
            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr, betas=betas)
        else:
            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)

    def beforeTrain(self):
        # keep track of training
        self.writer = SummaryWriter(
            "runs/"+self.full_tag+"_"+datetime.now().strftime('%b%d_%H-%M-%S'))

    def afterTrainStep(self, epoch_data, epoch_step):
        pass

    def train(self, ds: torch.utils.data.Dataset,
              batch_size=256,
              epochs=100,
              eval_every=10,
              score_every=10,
              reconstruction_every=10,
              reconstruct_indices: Sequence = None,
              save_every=20,
              save_path="./checkpoints/",
              num_workers=8,
              score_ds: GroundTruthDataset = None,
              max_iter=None,
              log_deadNeurons=False,
              log_stepdata=False):
        """Start training

        Args:
            ds (torch.utils.data.Dataset): Dataset
            batch_size (int, optional): batch size. Defaults to 256.
            epochs (int, optional): number of epoches. Defaults to 100.
            score (int, optional): after how many epochs score metrics will be calculated. Defaults to 10.
            eval_every (int, optional): after how many epochs evaluation will be perforemd. Defaults to 10.
            reconstruction_every (int, optional): after how many epochs reconstruction_callback will be called. Defaults to 10.
            save_every (int, optional): after how many epochs model will be saved. Defaults to 20.
            save_path (str, optional): path for model state dict. Defaults to "./checkpoints/".
        """

        self.beforeTrain()

        train_size = int(0.75 * len(ds))
        test_size = len(ds) - train_size
        train_dataset, test_dataset = torch.utils.data.random_split(ds, [train_size, test_size], generator=torch.Generator().manual_seed(0))

        data_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=num_workers)
        self.global_iter = 0
        self.model.train()

        last_save = None

        if log_deadNeurons:
            dn_weights = []
            dn_update_matix = []
            for name, param in self.model.named_parameters():
                if 'weight' in name:
                    dn_weights.append(name+" "+str(param.shape)[11:-1])
                    # two update maps: num updates, last update

                    dn_update_matix.append([torch.zeros(param.shape), torch.zeros(param.shape), torch.zeros([epochs] + list(param.shape),dtype=torch.bool)])

        for epoch in range(epochs):
            self.epoch_start = self.global_iter
            self.train_iter = epoch
            epoch_data = torch.empty((len(data_loader), self.num_log_values), dtype=torch.float32, device=self.device)

            if log_deadNeurons:
                dn_epoch_data = torch.empty((len(data_loader), len(dn_weights)), dtype=torch.float32, device=self.device)

            with tqdm(data_loader) as bar:
                for i, (x_true, _) in enumerate(bar):
                    self.global_iter += 1

                    x_true = x_true.to(self.device)

                    loss = self.calc_loss(x_true, batch_size, epoch_data, i)

                    self.optimizer.zero_grad(set_to_none=True)
                    loss.backward(retain_graph=False)
                    if log_deadNeurons:
                        dn_i=0
                        for name, param in self.model.named_parameters():
                            if 'weight' in name:
                                dn_filter = (param.grad == 0)
                                dn_epoch_data[i, dn_i] = dn_filter.sum().div(param.grad.numel())
                                if log_stepdata:
                                    self.writer.add_scalar('Dead Neurons_/weight.'+str(dn_i)+'.'+name+" "+str(param.shape)[11:-1], dn_epoch_data[i, dn_i].cpu().numpy(), self.global_iter)

                                dn_filter_inv = ~dn_filter
                                update_map, iteration_map, deadneuron_map_epochs = dn_update_matix[dn_i]
                                update_map[dn_filter_inv] +=1 #update_map
                                iteration_map[dn_filter_inv] = self.global_iter #iteration_map
                                deadneuron_map_epochs[epoch, dn_filter] = True

                                dn_i += 1

                    self.optimizer.step()

                    self.afterTrainStep(epoch_data, i)

                avg_loss = self.log_data(epoch, epoch_data, 'train')
                if log_deadNeurons:
                    dn_epoch_data = dn_epoch_data.detach().mean(axis=0).cpu().numpy()
                    for dn_i, name in enumerate(dn_weights):
                        self.writer.add_scalar('Dead Neurons/weight.'+str(dn_i)+'.'+name, dn_epoch_data[dn_i], epoch)
                        temp = ((dn_update_matix[dn_i][1] < self.epoch_start)*1.).mean().detach().cpu().numpy()
                        self.writer.add_scalar('Dead Neurons/weight.'+str(dn_i)+'.'+name+"_", temp, epoch)

                        self.writer.add_scalar('Dead_Neurons/weight.'+str(dn_i)+'.'+name, (dn_update_matix[dn_i][2][epoch]*1.0).mean(), epoch)

                if epoch % save_every == 0 and epoch > 0:
                    last_save = self.full_tag + "_" + str(epoch) + '_' + \
                        datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + "_"+str(avg_loss)+".save"
                    torch.save(self.model.state_dict(), save_path+last_save)

                print({"epoch": epoch, "step": self.global_iter, "loss": avg_loss, "last_save": last_save})
                # bar.set_postfix(epoch_stat=epoch, global_iter=self.global_iter, loss=avg_loss, last_save=last_save)

            do_eval = epoch % eval_every == 0
            do_recon = epoch % reconstruction_every == 0
            do_score = epoch % score_every == 0
            if do_eval or do_recon or do_score:
                self.model.eval()

            if do_eval:
                eval_loss = self.eval(test_dataset, batch_size=batch_size)
                self.log_data(epoch, eval_loss, 'eval')

            if do_recon:
                img = self._visualize_reconstruction(ds, indices=reconstruct_indices)
                self.writer.add_image('Reconstruction', torch.clamp(img, 0, 1), epoch)

            if do_score:
                scores = self.score(score_ds if score_ds is not None else ds)
                print(scores)
                for score_name, score_values in scores.items():
                    # self.writer.add_scalars(score_name, score_values, global_step=epoch)
                    for score_value_name, value in score_values.items():
                        self.writer.add_scalar(score_name+"/"+score_value_name, value, global_step=epoch)

            if do_eval or do_recon or do_score:
                self.model.train()

            if max_iter and self.global_iter >= max_iter:
                break

        # if we didn't calc scores in last iteration do it now
        if not do_score:
            scores = self.score(score_ds if score_ds is not None else ds)
            print(scores)
            for score_name, score_values in scores.items():
                # self.writer.add_scalars(score_name, score_values, global_step=epoch)
                for score_value_name, value in score_values.items():
                    self.writer.add_scalar(score_name+"/"+score_value_name, value, global_step=epoch)

        last_save = self.full_tag + "_" + \
            datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + "_"+str(avg_loss)+".save"
        torch.save(self.model.state_dict(), save_path+last_save)
        print("saved "+last_save)

        return_dict = {}
        if log_deadNeurons:
            return_dict['deadNeurondata'] = dn_update_matix
            return_dict['weights'] = dn_weights

        return return_dict

    def score(self, ds):
        random_state = np.random.RandomState(0)
        beta_vae_score = beta_vae.compute_beta_vae_sklearn(
            ds, make_representor(self),
            random_state, batch_size=64, num_train=10000, num_eval=5000)

        random_state = np.random.RandomState(0)
        factor_vae_score = factor_vae.compute_factor_vae(
            ds, make_representor(self),
            random_state, num_variance_estimate=10000)

        random_state = np.random.RandomState(0)
        mig_score = mig.compute_mig(ds, make_representor(self), random_state)

        # random_state = np.random.RandomState(0)
        # dci_score = dci.compute_dci(ds, make_representor(self), random_state)

        random_state = np.random.RandomState(0)
        sap = sap_score.compute_sap(ds, make_representor(self), random_state)

        return {
            "beta_vae_score": beta_vae_score,
            "factor_vae_score": factor_vae_score,
            "mig_score": mig_score,
            # "dci_score": dci_score,
            "sap_score": sap,
        }

    def calc_loss(self, x_true, batch_size, epoch_data, epoch_step):
        y, (mu, log_var), _ = self.model.forward(x_true)
        kld = self.gauss_vae_regulariser(mu, log_var)
        reg_loss = self.beta * kld
        recon_loss = self.reconstruction_loss(y, x_true).div(batch_size)

        loss = recon_loss + reg_loss

        if self.model.use_model_regulariser:
            loss += self.model_regulariser_weight * self.model.model_regulariser()

        epoch_data[epoch_step, 0] = recon_loss.detach()
        epoch_data[epoch_step, 1] = reg_loss.detach()
        epoch_data[epoch_step, 2] = kld.detach()

        with torch.no_grad():
            mse = nn.functional.mse_loss(
                self.model.activation(y) if isinstance(self.model.loss_activation, nn.Identity) else y,
                x_true,
                reduction='sum').div(batch_size)
            epoch_data[epoch_step, 3] = mse.detach()

        return loss

    def log_data(self, epoch, epoch_data, type='train'):
        epoch_data = epoch_data.detach().mean(axis=0).cpu().numpy()
        l1, l2 = epoch_data[0], epoch_data[1]

        self.writer.add_scalar('Loss reconstruction/'+type, l1, epoch)
        self.writer.add_scalar('Loss regularisation/'+type, l2, epoch)
        self.writer.add_scalar('KL divergence/'+type, epoch_data[2], epoch)
        self.writer.add_scalar('MSE reconstruction/'+type, epoch_data[3], epoch)

        return l1 + l2

    def eval(self, ds, batch_size=256):
        self.model.eval()
        data_loader = DataLoader(ds, batch_size=batch_size, shuffle=True)

        with torch.no_grad():
            epoch_loss = torch.empty((len(data_loader), self.num_log_values), dtype=torch.float32, device=self.device)
            for i, (x_true, _) in enumerate(tqdm(data_loader)):
                x_true = x_true.to(self.device)

                self.calc_loss(x_true, batch_size, epoch_loss, i)

            return epoch_loss

    def loadModel(self, checkpoint, folder="./checkpoints/"):
        """Loads model checkpint

        Args:
            checkpoint (str): Filename of checkpoint
            folder (str, optional): dir where checkpoint is located. Defaults to "./checkpoints/".
        """
        self.model.load_state_dict(torch.load(folder + checkpoint))
        self.model.eval()

    def encode(self, x):
        return self.model.encode(x.to(self.device))[0]

def make_representor(model: BetaVAE):
    """
    Encloses the pytorch ScriptModule in a callable that can be used by `disentanglement_lib`.
    Parameters
    ----------
    model : torch.nn.Module or torch.jit.ScriptModule
        The Pytorch model.
    cuda : bool
        Whether to use CUDA for inference. Defaults to the return value of the `use_cuda`
        function defined above.
    Returns
    -------
    callable
        A callable function (`representation_function` in dlib code)
    """
    # Deepcopy doesn't work on ScriptModule objects yet:
    # https://github.com/pytorch/pytorch/issues/18106
    # model = deepcopy(model)
    # cuda = use_cuda() if cuda is None else cuda
    # model = model.cuda() if cuda else model.cpu()

    # Define the representation function
    def _represent(x):
        # assert isinstance(x, np.ndarray), \
        #     "Input to the representation function must be a ndarray."
        # assert x.ndim == 4, \
        #     "Input to the representation function must be a four dimensional NHWC tensor."
        # # Convert from NHWC to NCHW
        # x = np.moveaxis(x, 3, 1)
        # # Convert to torch tensor and evaluate
        # x = torch.from_numpy(x).float().to('cuda' if cuda else 'cpu')
        with torch.no_grad():
            y = model.encode(x)
        y = y.cpu().numpy()
        assert y.ndim == 2, \
            "The returned output from the representor must be two dimensional (NC)."
        return y

    return _represent


'''
A student class, instantiated 1 per user

'''

class QuestionSetProgress:
    '''
    data is the history of the Student on this particular question
    '''
    def __init__(self, questionDictionary, numberTestSessions, testData, ):
        self.questions = questionDictionary
        self.timesTookTest = numberTestSessions
        self.data = testData

'''

class Student:
    def __init__(self, PastInfo ):
        self.studentMind = dataTable

        self.

'''


import sys, threading
import django.db

import config
from . import currentUser, api, login, dump
from .lib import db, util, rpc, logging #@UnresolvedImport
from .lib.error import Error, UserError, InternalError

def logCall(function, args, kwargs):
	logging.log(category="api", method=function.__name__, args=args, kwargs=kwargs, user=currentUser().name if currentUser() else None)

class Wrapper:
	def __init__(self):
		self.semaphore = threading.Semaphore(config.MAX_REQUESTS)
	def __enter__(self):
		self.semaphore.acquire()
	def __exit__(self, exc_type, exc_val, exc_tb):
		self.semaphore.release()
		if django.db.transaction.is_dirty():
			django.db.transaction.commit()
		django.db.connection.close()

@db.commit_after
def handleError(error, function, args, kwargs):
	if not isinstance(error, Error):
		if isinstance(error, TypeError) and function.__name__ in str(error):
			error = UserError.wrap(error, data={"function": function.__name__, "args": args, "kwargs": kwargs})
		else:
			error = InternalError.wrap(error, data={"function": function.__name__, "args": args, "kwargs": kwargs})
	logging.logException()
	if isinstance(error, InternalError):
		dump.dumpException()
	return error

@db.commit_after
def afterCall(*args, **kwargs):
	pass

def runServer(server):
	try:
		server.serve_forever()
	except KeyboardInterrupt:
		pass

servers = []

def wrapError(error, func, args, kwargs):
	error = handleError(error, func, args, kwargs)
	if isinstance(error, rpc.Fault):
		return error
	assert isinstance(error, Error)
	if error.code == UserError.NOT_LOGGED_IN:
		return rpc.xmlrpc.ErrorUnauthorized()
	return rpc.Fault(999, error.rawstr)

def start():
	print >>sys.stderr, "Starting RPC servers"
	global servers
	del servers[:]
	for settings in config.SERVER:
		server_address = ('', settings["PORT"])
		sslOpts = None
		if settings["SSL"]:
			sslOpts = rpc.SSLOpts(private_key=settings["SSL_OPTS"]["key_file"], certificate=settings["SSL_OPTS"]["cert_file"], client_certs=None)
		server = rpc.xmlrpc.XMLRPCServerIntrospection(server_address, sslOpts=sslOpts, loginFunc=login, wrapper=Wrapper(), beforeExecute=logCall, afterExecute=afterCall, onError=wrapError)
		server.register(api)
		print >>sys.stderr, " - %s:%d, SSL: %s" % (server_address[0], server_address[1], bool(sslOpts))
		util.start_thread(server.serve_forever)
		servers.append(server)

def stop():
	for server in servers:
		server.shutdown()


#!/usr/bin/python

def person(name,age,**kw):
    print('name:',name,'age:',age,'other:',kw)

person('Bob',35,city='Beijing')

from flask import Flask

app = Flask(__name__)

@app.route('/')
def hello_world():
    return 'Hello World!'

if __name__ == '__main__':
    app.run(host="0.0.0.0")


from __future__ import division

import s2s
import argparse
import torch
import torch.nn as nn
from torch import cuda
from torch.autograd import Variable
import math
import time
import logging
import random
import numpy as np

try:
    import ipdb
except ImportError:
    pass
from nltk.translate import bleu_score
from s2s.xinit import xavier_normal, xavier_uniform
import os
import xargs
torch.backends.cudnn.enabled=False

parser = argparse.ArgumentParser(description='train.py')
xargs.add_data_options(parser)
xargs.add_model_options(parser)
xargs.add_train_options(parser)

opt = parser.parse_args()

logging.basicConfig(format='%(asctime)s [%(levelname)s:%(name)s]: %(message)s', level=logging.INFO)
log_file_name = time.strftime("%Y%m%d-%H%M%S") + '.log.txt'
if opt.log_home:
    log_file_name = os.path.join(opt.log_home, log_file_name)
file_handler = logging.FileHandler(log_file_name, encoding='utf-8')
file_handler.setFormatter(logging.Formatter('%(asctime)s [%(levelname)-5.5s:%(name)s] %(message)s'))
logging.root.addHandler(file_handler)
logger = logging.getLogger(__name__)

logger.info('My PID is {0}'.format(os.getpid()))
logger.info('PyTorch version: {0}'.format(str(torch.__version__)))
logger.info(opt)

if torch.cuda.is_available() and not opt.gpus:
    logger.info("WARNING: You have a CUDA device, so you should probably run with -gpus 0")

if opt.seed > 0:
    torch.manual_seed(opt.seed)

if opt.gpus:
    if opt.cuda_seed > 0:
        torch.cuda.manual_seed(opt.cuda_seed)
    cuda.set_device(opt.gpus[0])

logger.info('My seed is {0}'.format(torch.initial_seed()))
logger.info('My cuda seed is {0}'.format(torch.cuda.initial_seed()))

def NMTCriterion(vocabSize):
    weight = torch.ones(vocabSize)
    weight[s2s.Constants.PAD] = 0

    crit = nn.NLLLoss(weight, size_average=False)
    if opt.gpus:
        crit.cuda()
    return crit

def reverse_NMTCriterion(vocabSize):
    weight = torch.zeros(vocabSize)
    weight[s2s.Constants.PAD] = 0

    crit = nn.NLLLoss(weight, size_average=False)
    if opt.gpus:
        crit.cuda()
    return crit

def loss_function(g_outputs, g_targets, generator, crit, eval=False):
    batch_size = g_outputs.size(1)

    g_out_t = g_outputs.view(-1, g_outputs.size(2))
    g_prob_t = generator(g_out_t)

    g_loss = crit(g_prob_t, g_targets.view(-1))
    total_loss = g_loss
    report_loss = total_loss.item()
    return total_loss, report_loss, 0

def generate_copy_loss_function(g_outputs, c_outputs, g_targets,
                                c_switch, c_targets, c_gate_values,
                                generator, nlu_generator, classifier, crit, copyCrit,
                                reverse_crit, guideCrit, nluCrit, guides, g_hiddens, nlu_outputs,type_emb):

    # g_out is output [(TargetLength, BatchSize, Hidden)]
    batch_size = g_outputs.size(1)
    targetLength = g_outputs.size(0)

    guides = guides.transpose(1,0)

    # transform typed target with the special <ent> symbol to mention level target
    m_targets = g_targets.clone().detach()
    m_targets = m_targets.transpose(1,0)
    for i in range(batch_size):
        count = 0
        for j in range(targetLength):
            if m_targets[i][j] == s2s.Constants.SS:
                m_targets[i][j] = guides[i][count]
                count += 1
            else:
                m_targets[i][j] = s2s.Constants.PAD

    m_targets = m_targets.transpose(1,0)

    # calculate nlu prediction distribution
    nlu_out_t = nlu_outputs.view(-1, nlu_outputs.size(2))
    nlu_prob_t = nlu_generator(nlu_out_t)
    nlu_prob_t = nlu_prob_t.view(-1, batch_size, nlu_prob_t.size(1))

    # calculate generation prediction distribution
    g_out_t = g_outputs.view(-1, g_outputs.size(2))
    g_prob_t = generator(g_out_t)
    g_prob_t = g_prob_t.view(-1, batch_size, g_prob_t.size(1))
    g_hiddens = g_hiddens.squeeze(1).transpose(1,0)

    # calculate the mention predictor loss
    cl_loss = 0
    type_emb = type_emb.transpose(1,0) # injected type embedding from source
    g_prob_t_trans = g_prob_t.transpose(1,0)

    g_targets = g_targets.transpose(1,0)
    guide_probs = []
    for i in range(batch_size):
        guide_prob = []
        guide = []
        count = 0
        for j in range(targetLength):
            if g_targets[i][j] == s2s.Constants.SS:
                injected = torch.cat([g_hiddens[i][j],type_emb[i][count]],dim=0)

                guide_prob.append(torch.log(classifier(injected.unsqueeze(0))))
                guide.append(guides[i][count])
                count += 1
        guide_prob = torch.stack(guide_prob).squeeze(1)
        guide = torch.stack(guide).view(-1)
        if i == 0:
            cl_loss = guideCrit(guide_prob,guide)
        else:
            cl_loss += guideCrit(guide_prob,guide)
    g_targets = g_targets.transpose(1,0)

    g_output_prob = g_prob_t + 1e-8

    g_output_prob_log = torch.log(g_output_prob)
    nlu_prob_log = torch.log(nlu_prob_t)
    g_output_prob_log = g_output_prob_log * ((1 - c_switch).unsqueeze(2).expand_as(g_output_prob_log))

    g_output_prob_log = g_output_prob_log.view(-1, g_output_prob_log.size(2))
    nlu_prob_log = nlu_prob_log.view(-1, nlu_prob_log.size(2))

    # g_output_prob_log [(TargetLength * BatchSize, vocab_size)]
    g_loss = crit(g_output_prob_log, g_targets.view(-1))
    nlu_loss = nluCrit(nlu_prob_log, m_targets.view(-1))

    total_loss = g_loss + nlu_loss * 1.5 + cl_loss * 2.5
    report_loss = total_loss.item()
    return total_loss, report_loss, 0, g_loss, cl_loss, nlu_loss

def addPair(f1, f2, f3):
    for x, y1, z in zip(f1, f2, f3):
        yield (x, y1, z)
    yield (None, None, None)

def load_dev_data(translator, src_file, bio_file, feat_files, tgt_file, guide_file):
    dataset, raw = [], []
    srcF = open(src_file, encoding='utf-8')
    guideF = open(guide_file, encoding='utf-8')
    tgtF = open(tgt_file, encoding='utf-8')
    bioF = open(bio_file, encoding='utf-8')
    featFs = [open(x, encoding='utf-8') for x in feat_files]

    src_batch, tgt_batch = [], []
    bio_batch, feats_batch = [], []
    guide_batch = []
    for line, tgt, guide in addPair(srcF, tgtF, guideF):
        if (line is not None) and (tgt is not None) and (guide is not None):
            src_tokens = line.strip().split(' ')
            src_batch += [src_tokens]
            guide_tokens = guide.strip().split(' ')
            guide_batch += [guide_tokens]
            tgt_tokens = tgt.strip().split(' ')
            tgt_batch += [tgt_tokens]
            bio_tokens = bioF.readline().strip().split(' ')
            bio_batch += [bio_tokens]
            feats_tokens = [reader.readline().strip().split((' ')) for reader in featFs]
            feats_batch += [feats_tokens]

            if len(src_batch) < opt.batch_size:
                continue
        else:
            # at the end of file, check last batch
            if len(src_batch) == 0:
                break
        data = translator.buildData(src_batch, bio_batch, feats_batch, tgt_batch, guide_batch)
        dataset.append(data)
        raw.append((src_batch, tgt_batch))
        src_batch, tgt_batch = [], []
        bio_batch, feats_batch = [], []
        guide_batch = []
    srcF.close()
    bioF.close()
    for f in featFs:
        f.close()
    tgtF.close()
    guideF.close()
    return (dataset, raw)

evalModelCount = 0
totalBatchCount = 0

def evalModel(model, translator, evalData):
    global evalModelCount
    evalModelCount += 1
    ofn_prefix = opt.ofn_prefix
    if evalModelCount % 2 == 0:
        ofn_prefix += '_test'
    else:
        ofn_prefix += '_dev'
    ofn = ofn_prefix + '.{0}'.format(evalModelCount)
    if opt.save_path:
        ofn = os.path.join(opt.save_path, ofn)

    predict, gold = [], []
    processed_data, raw_data = evalData
    for batch, raw_batch in zip(processed_data, raw_data):
        #  (1) read src and tgt (bio and feats are part of the legacy code and are never used in this project)
        src, bio, feats, tgt, guide, indices = batch[0]
        src_batch, tgt_batch = raw_batch

        #  (2) translate
        pred, predScore, predIsCopy, predCopyPosition, attn, _ = translator.translateBatch(src, bio, feats, tgt, guide)
        pred, predScore, predIsCopy, predCopyPosition, attn = list(zip(
            *sorted(zip(pred, predScore, predIsCopy, predCopyPosition, attn, indices),
                    key=lambda x: x[-1])))[:-1]

        #  (3) convert indexes to words
        predBatch = []
        for b in range(src[0].size(1)):
            n = 0
            predBatch.append(
                translator.buildTargetTokens(pred[b][n], src_batch[b],
                                             predIsCopy[b][n], predCopyPosition[b][n], attn[b][n])
            )
        # nltk BLEU evaluator needs tokenized sentences
        gold += [[r] for r in tgt_batch]
        predict += predBatch
        for i, j in zip([[r] for r in tgt_batch], predBatch):
            if random.random() > 1:
                print('{}\n{}\n{}\n'.format(i, j, '*'*50))

    no_copy_mark_predict = [[word.replace('[[', '').replace(']]', '') for word in sent] for sent in predict]
    bleu = bleu_score.corpus_bleu(gold, no_copy_mark_predict)
    report_metric = bleu

    # (4) save generation result
    with open(ofn, 'w', encoding='utf-8') as of:
        for p in predict:
            of.write(' '.join(p) + '\n')
    print(ofn)
    return report_metric

def trainModel(model, translator, trainData, validData, testData, dataset, optim):
    logger.info(model)
    model.train()
    torch.cuda.empty_cache()
    for name,para in model.named_parameters():
        if para.requires_grad:
            print(name,para.data.shape)
    logger.warning("Set model to {0} mode".format('train' if model.decoder.dropout.training else 'eval'))

    # define criterion of each GPU
    criterion = NMTCriterion(dataset['dicts']['tgt'].size())
    guideCriterion = NMTCriterion(dataset['dicts']['guide_src'].size())
    nluCriterion = NMTCriterion(dataset['dicts']['guide_src'].size())
    reverse_criterion = reverse_NMTCriterion(dataset['dicts']['tgt'].size())
    copyLossF = nn.NLLLoss(size_average=False)

    start_time = time.time()

    def saveModel(metric=None):
        model_state_dict = model.module.state_dict() if len(opt.gpus) > 1 else model.state_dict()
        model_state_dict = {k: v for k, v in model_state_dict.items() if 'generator' not in k}
        generator_state_dict = model.generator.module.state_dict() if len(
            opt.gpus) > 1 else model.generator.state_dict()
        #  (5) drop a checkpoint
        checkpoint = {
            'model': model_state_dict,
            'generator': generator_state_dict,
            'dicts': dataset['dicts'],
            'opt': opt,
            'epoch': epoch,
            'optim': optim
        }
        save_model_path = 'model'
        if opt.save_path:
            if not os.path.exists(opt.save_path):
                os.makedirs(opt.save_path)
            save_model_path = opt.save_path + os.path.sep + save_model_path
        if metric is not None:
            torch.save(checkpoint, '{0}_dev_metric_{1}_e{2}.pt'.format(save_model_path, round(metric, 4), epoch))
        else:
            torch.save(checkpoint, '{0}_e{1}.pt'.format(save_model_path, epoch))

    def trainEpoch(epoch):
        if opt.extra_shuffle and epoch > opt.curriculum:
            logger.info('Shuffling...')
            trainData.shuffle()

        # shuffle mini batch order
        batchOrder = torch.randperm(len(trainData))

        total_loss, total_words, total_num_correct = 0, 0, 0
        report_loss, report_tgt_words, report_src_words, report_num_correct = 0, 0, 0, 0
        start = time.time()
        epoch_loss = 0
        epoch_cl_loss = 0
        epoch_nlu_loss = 0
        for i in range(len(trainData)):
            global totalBatchCount
            totalBatchCount += 1
            """
            [0] ([0][0] wrap(srcBatch), [0][1] lengths), [1] ([1][0] wrap(bioBatch), [1][1] lengths), [2] ([2][0] (wrap(x) for x in featBatches), [2][1] lengths),
            [3] ([3][0] wrap(tgtBatch), [3][1] wrap(copySwitchBatch), [3][2] wrap(copyTgtBatch)),
            [4] [4] indices (not used) # original indices will be excluded !
            """
            batchIdx = batchOrder[i] if epoch > opt.curriculum else i
            batch = trainData[batchIdx][:-1]  # exclude original indices

            model.zero_grad()
            '''
            # g_out is output [(TargetLength, BatchSize, Hidden)]
            # c_out is attn_out [(TargetLength, BatchSize, InputLength)]
            # c_gate is copy_prob [(TargetLength, BatchSize, 1)]
            '''
            g_outputs, c_outputs, c_gate_values, g_hiddens, nlu_outputs, type_emb = model(batch, False)
            targets = batch[3][0][1:]  # exclude <s> from targets
            copy_switch = batch[3][1][1:]
            c_targets = batch[3][2][1:]
            g_guides = batch[4][0]

            loss, res_loss, num_correct, g_loss, cl_loss, nlu_loss = generate_copy_loss_function(
                g_outputs, c_outputs, targets, copy_switch, c_targets, c_gate_values, model.generator, model.nlu_generator, model.classifier,
                criterion, copyLossF, reverse_criterion, guideCriterion, nluCriterion, g_guides, g_hiddens, nlu_outputs, type_emb)
            torch.cuda.empty_cache()
            epoch_loss += g_loss
            epoch_cl_loss += cl_loss
            epoch_nlu_loss += nlu_loss

            if math.isnan(res_loss) or res_loss > 1e20:
                logger.info('catch NaN')
                ipdb.set_trace()

            # update the parameters
            loss.backward()
            optim.step()

            num_words = targets.data.ne(s2s.Constants.PAD).sum().item()
            report_loss += res_loss
            report_num_correct += num_correct
            report_tgt_words += num_words
            report_src_words += batch[0][-1].data.sum()
            total_loss += res_loss
            total_num_correct += num_correct
            total_words += num_words
            if i % opt.log_interval == 0:
                logger.info(
                    "Epoch %2d, %6d/%5d/%5d; acc: %6.2f; loss: %6.2f; words: %5d; ppl: %6.2f; %3.0f src tok/s; %3.0f tgt tok/s; %6.0f s elapsed" %
                    (epoch, totalBatchCount, i + 1, len(trainData),
                     report_num_correct / report_tgt_words * 100,
                     report_loss,
                     report_tgt_words,
                     math.exp(min((report_loss / report_tgt_words), 16)),
                     report_src_words / max((time.time() - start), 1.0),
                     report_tgt_words / max((time.time() - start), 1.0),
                     time.time() - start))

                report_loss = report_tgt_words = report_src_words = report_num_correct = 0
                start = time.time()

            if validData is not None and totalBatchCount % opt.eval_per_batch == 0 and totalBatchCount >= opt.start_eval_batch:
                model.eval()
                logger.warning("Set model to {0} mode".format('train' if model.decoder.dropout.training else 'eval'))
                valid_bleu = evalModel(model, translator, validData)
                model.train()
                logger.warning("Set model to {0} mode".format('train' if model.decoder.dropout.training else 'eval'))
                model.decoder.attn.mask = None
                logger.info('Validation Score: %g' % (valid_bleu * 100))
                if valid_bleu >= optim.best_metric:
                    saveModel(valid_bleu)
                optim.updateLearningRate(valid_bleu, epoch)

            if testData is not None and totalBatchCount % opt.eval_per_batch == 0 and totalBatchCount >= opt.start_eval_batch:
                model.eval()
                logger.warning("Set model to {0} mode".format('train' if model.decoder.dropout.training else 'test'))
                test_bleu = evalModel(model, translator, testData)
                logger.info('Test Score: %g' % (test_bleu * 100))
        print('epoch g loss: ')
        print(epoch_loss)
        print('epoch cl loss: ')
        print(epoch_cl_loss)
        print('epoch nlu loss: ')
        print(epoch_nlu_loss)

        return total_loss / total_words, total_num_correct / total_words

    for epoch in range(opt.start_epoch, opt.epochs + 1):
        logger.info('')
        # train for one epoch on the training set
        train_loss, train_acc = trainEpoch(epoch)
        # calculate metric for each epoch training
        train_ppl = math.exp(min(train_loss, 100))
        logger.info('Train perplexity: %g' % train_ppl)
        logger.info('Train accuracy: %g' % (train_acc * 100))
        logger.info('Saving checkpoint for epoch {0}...'.format(epoch))
        saveModel()

def main():

    import onlinePreprocess
    onlinePreprocess.lower = opt.lower_input
    onlinePreprocess.seq_length = opt.max_sent_length
    onlinePreprocess.shuffle = 1 if opt.process_shuffle else 0
    from onlinePreprocess import prepare_data_online

    # opt.train_src (source file of sequence) 'it is a replica of the grotto at lourdes , france where the virgin mary reputedly appeared to saint bernadette soubirous in 1858 .'
    # opt.src_vocab (source file of vocab) 'the(word) 4(index) 256272(frequency) 0.06749202214022335'
    # opt.train_bio (answer position embedding) 'O O O O O O O O O O O O O O O O O O B I I O O O'
    # opt.bio_vocab (source file of answer position vocab) 'O(bio) 4(index) 2525015(frequency) 0.8958601572376024'
    # opt.train_feats (source file of postag/ner/case) 'PERSON/UPCASE/NN ...' (3 different embeddings)
    # opt.feat_vocab (source file of answer feat vocab)
    # opt.train_tgt (source file of question) 'to whom did the virgin mary allegedly appear in 1858 in lourdes france ?'
    # opt.tgt_vocab (source file of vocab) same file with opt.src_vocab !!
    dataset = prepare_data_online(opt.train_src, opt.src_vocab, opt.train_bio, opt.bio_vocab, opt.train_feats,
                                  opt.feat_vocab, opt.train_tgt, opt.tgt_vocab, opt.train_guide_src, opt.guide_src_vocab)

    trainData = s2s.Dataset(dataset['train']['src'], dataset['train']['bio'], dataset['train']['feats'],
                            dataset['train']['tgt'], dataset['train']['switch'], dataset['train']['c_tgt'], opt.batch_size, opt.gpus, dataset['train']['guide_src'])

    dicts = dataset['dicts']
    logger.info(' * vocabulary size. source = %d; target = %d' % (dicts['src'].size(), dicts['tgt'].size()))
    logger.info(' * number of training sentences. %d' % len(dataset['train']['src']))
    logger.info(' * maximum batch size. %d' % opt.batch_size)

    logger.info('Building Model ...')
    encoder = s2s.Models.Encoder(opt, dicts['src'], dicts['guide_src'])
    decoder = s2s.Models.Decoder(opt, dicts['tgt'])
    decIniter = s2s.Models.DecInit(opt)

    ''' generator map output embedding to vocab size vector then softmax'''
    generator = nn.Sequential(
        nn.Linear(opt.dec_rnn_size // opt.maxout_pool_size, dicts['tgt'].size()),
        nn.Softmax(dim=1)
    )
    classifier = nn.Sequential(
        nn.Linear(opt.dec_rnn_size + 300, dicts['guide_src'].size()),
        nn.Softmax(dim=1)
    )
    nlu_generator = nn.Sequential(
        nn.Linear(opt.dec_rnn_size * 2, dicts['guide_src'].size()),
        nn.Softmax(dim=1)
    )

    model = s2s.Models.NMTModel(encoder, decoder, decIniter)
    model.generator = generator
    model.classifier = classifier
    model.nlu_generator = nlu_generator
    translator = s2s.Translator(opt, model, dataset)

    if len(opt.gpus) >= 1:
        model.cuda()
        generator.cuda()
        classifier.cuda()
        nlu_generator.cuda()
    else:
        model.cpu()
        generator.cpu()
        classifier.cpu()
        nlu_generator.cpu()

    # if len(opt.gpus) > 1:
    #     model = nn.DataParallel(model, device_ids=opt.gpus, dim=1)
    #     generator = nn.DataParallel(generator, device_ids=opt.gpus, dim=0)

    for pr_name, p in model.named_parameters():
        logger.info(pr_name)
        # p.data.uniform_(-opt.param_init, opt.param_init)
        if p.dim() == 1:
            # p.data.zero_()
            p.data.normal_(0, math.sqrt(6 / (1 + p.size(0))))
        else:
            nn.init.xavier_normal_(p, math.sqrt(3))

    encoder.load_pretrained_vectors(opt)
    decoder.load_pretrained_vectors(opt)

    optim = s2s.Optim(
        opt.optim, opt.learning_rate,
        max_grad_norm=opt.max_grad_norm,
        max_weight_value=opt.max_weight_value,
        lr_decay=opt.learning_rate_decay,
        start_decay_at=opt.start_decay_at,
        decay_bad_count=opt.halve_lr_bad_count
    )
    optim.set_parameters(model.parameters())

    validData = None
    if opt.dev_input_src and opt.dev_ref:
        validData = load_dev_data(translator, opt.dev_input_src, opt.dev_bio, opt.dev_feats, opt.dev_ref, opt.dev_guide_src)

    testData = None
    if opt.test_input_src and opt.test_ref:
        testData = load_dev_data(translator, opt.test_input_src, opt.test_bio, opt.test_feats, opt.test_ref, opt.test_guide_src)

    trainModel(model, translator, trainData, validData, testData, dataset, optim)

if __name__ == "__main__":
    main()


#!/usr/bin/env python3

'''
Given an unsorted array nums, reorder it in-place such that nums[0] < nums[1] > nums[2] < nums[3].... Example:

Input: nums = [3,5,2,1,6,4]
Output: One possible answer is [3,5,1,6,2,4]
'''

from typing import List

class Solution:
	def wiggleSort(self, nums: List[int]) -> None:
		nums.sort()
		nums1 = nums[:]

		j = int((len(nums)-1)/2)
		k = len(nums)-1
		jj = j
		i = 0
		while i < len(nums)-1:
			if j > -1:
				nums[i] = nums1[j]
				i += 1
				j -= 1
			if k > jj:
				nums[i] = nums1[k]
				i += 1
				k -= 1

def main():
	sol = Solution()
	nums = [1,4,2,4,4,6]

	sol.wiggleSort(nums)
	print(nums)

if __name__ == "__main__":
	main()

1. 1-D Array (19:19)/314. Wiggle Sort II/wiggleSort.py