import os
import io
import socket

from enum import Enum
from queue import Queue
from typing import List, Tuple, Optional, NewType, Union

import torch

import persia.env as env

from persia.logger import get_default_logger
from persia.sparse.optim import Optimizer
from persia.prelude import PyPersiaCommonContext, PyPersiaBatchData, PyTensor
from persia.distributed import DistributedBaseOption, get_default_distributed_option

_CURRENT_CXT = None

_logger = get_default_logger()

PythonTrainBatch = NewType("PythonTrainBatch", object)

def _check_finite(tensors: List[torch.Tensor]) -> bool:
    """Check if all tensors in the input list contain only finite elements.

    Arguments:
        tensors (List[torch.Tensor]): List of tensor to be checked.

    Returns:
        bool: ``True`` if all elements in ``tensors`` are finite or None.
    """
    return all([torch.isfinite(t).all() if t is not None else True for t in tensors])

def _cast_dlpack2torch_tensor(
    pytensor: PyTensor, requires_grad: bool = False
) -> torch.Tensor:
    """Convert the DLPack PythonCapsule to torch tensor

    Arguments:
        pytensor (PyTensor): PersiaTensor wrapper that contains dlpack information.
        requires_grad (bool, optional): Whether current tensor requires grad or not.
    Returns: pytorch tensor
    """

    import torch.utils.dlpack as dlpack

    tensor = dlpack.from_dlpack(pytensor.dlpack)
    tensor.requires_grad = requires_grad
    return tensor

class PreprocessMode(Enum):
    r"""Mode of preprocessing.

    Used by ``EmbeddingCtx.prepare_features`` to generate features of different datatypes.

    When set to ``TRAIN``, ``prepare_features`` will return a torch tensor with ``requires_grad`` attribute set to ``True``.
    When set to ``EVAL``, ``prepare_features`` will return a torch tensor with ``requires_grad`` attribute set to ``False``.
    ``INFERENCE`` behaves almost identical to ``PreprocessMode.EVAL``, except that ``INFERENCE`` allows ""EmbeddingCtx`` to process the ``PythonTrainBatch`` without a target tensor.
    """
    TRAIN = 1
    EVAL = 2
    INFERENCE = 3

class BaseCtx:
    r"""Initializes a common context for other persia context, e.g. `DataCtx`, `EmbeddingCtx` and `TrainCtx`.
    This class should not be instantiated directly.
    """

    def __init__(
        self, threadpool_worker_size: int = 10, device_id: Optional[int] = None
    ):
        """
        Arguments:
            threadpool_worker_size (int): Rpc threadpool worker size.
            device_id (int, optional): The CUDA device to use for this process.
        """
        self.origin_context = None

        if device_id is not None and device_id >= 0:
            assert torch.cuda.is_available() and (
                0 <= device_id < torch.cuda.device_count()
            ), f"device_id: {device_id} invalid!"

            torch.cuda.set_device(device_id)
        else:
            device_id = None

        self.device_id = device_id

        replica_index = (
            env.get_rank() if env.get_rank() != -1 else env.get_replica_index()
        )
        replica_size = (
            env.get_world_size()
            if env.get_world_size() != -1
            else env.get_replica_size()
        )

        self.common_context = PyPersiaCommonContext(
            threadpool_worker_size, replica_index, replica_size, device_id
        )
        _logger.info(
            f"init persia context, replica_size: {replica_size} replica_index: {replica_index}"
        )

    def _enter(self):
        """Hook when enter the context"""
        ...

    def _exit(self):
        """Hook when exit the context"""
        ...

    def __enter__(self):
        self._enter()

        global _CURRENT_CXT
        self.origin_context = _CURRENT_CXT
        _CURRENT_CXT = self

        return self

    def __exit__(self, exc_type, value, trace):
        self._exit()

        global _CURRENT_CXT
        _CURRENT_CXT = self.origin_context

        if exc_type:
            import traceback

            _logger.error("\n" + traceback.format_exc())

class DataCtx(BaseCtx):
    r"""Provides the communicate ability for data generator component to send the PersiaBatchData
    to the trainer and embedding middleware.

    Example:
        >>> from persia.prelude import PyPersiaBatchData
        >>> loader = make_simple_loader()
        >>> with DataCtx() as ctx:
        >>>     for (dense, batch_sparse_ids, target) in loader:
        >>>         batch_data = PyPersiaBatchData()
        >>>         batch_data.add_dense([dense])
        >>>         batch_data.add_sparse(batch_sparse_ids)
        >>>         batch_data.add_target(target)
        >>>         ctx.send_data(batch_data)
    """

    def __init__(
        self,
        *args,
        **kwargs,
    ):
        super(DataCtx, self).__init__(*args, **kwargs)
        self.prepare()
        _logger.info("Data ctx prepare done.")

    def prepare(self):
        """Do some preparation to init `DataCtx`."""

        self.common_context.init_nats_publisher(None)
        self.common_context.wait_servers_ready()

    def send_sparse_to_middleware(self, data: PyPersiaBatchData):
        """Send PersiaBatchData from data compose to middleware side.

        Arguments:
            data (PyPersiaBatchData): PersiaBatchData that haven't been process.
        """
        self.common_context.send_sparse_to_middleware(data)

    def send_dense_to_trainer(self, data: PyPersiaBatchData):
        """Send PersiaBatchData from data compose to trainer side.

        Arguments:
            data (PyPersiaBatchData): PersiaBatchData that have been sent to middleware.
        """
        self.common_context.send_dense_to_trainer(data)

    def send_data(self, data: PyPersiaBatchData):
        """Send PersiaBatchData from data compose to trainer and middleware side.

        Arguments:
            data (PyPersiaBatchData): PersiaBatchData that haven't been process.
        """
        self.send_sparse_to_middleware(data)
        self.send_dense_to_trainer(data)

class EmbeddingConfig:
    r"""Embedding hyperparameters, argument of ``EmbeddingCtx``."""

    def __init__(
        self,
        emb_initialization: Tuple[float, float] = (-0.01, 0.01),
        admit_probability: float = 1.0,
        weight_bound: float = 10,
    ):
        """
        Arguments:
            emb_initialization (Tuple[float, float], optional): Lower and upper bound of embedding uniform initialization.
            admit_probability (float, optional): The probability (0<=, <=1) of admitting a new embedding.
            weight_bound (float, optional): Restrict each element value of an embedding in [-weight_bound, weight_bound].
        """
        self.emb_initialization = emb_initialization
        self.admit_probability = admit_probability
        self.weight_bound = weight_bound

def _get_default_embedding_config():
    """Get default embedding configuration"""
    return EmbeddingConfig()

class EmbeddingCtx(BaseCtx):
    r"""Provides the embedding-related functionality. EmbeddingCtx can run offline test or online inference
    depending on different preprocess_mode. The simplest way to get this context is by using ``persia.ctx.eval_ctx()`` or
    ``persia.ctx.inference_ctx`` to get the ``EmbeddingCtx`` instance.

    Example:
        >>> from persia.prelude import PyPersiaBatchData
        >>> model = get_dnn_model()
        >>> loader = make_dataloader()
        >>> embedding_config = EmbeddingConfig()
        >>> with EmbeddingCtx(
        ...     model=model,
        ...     PreprocessMode.EVAL,
        ...     embedding_config
        ... ) as ctx:
        >>>     for (dense, batch_sparse_ids, target) in loader:
        >>>         batch_data = PyPersiaBatchData()
        >>>         batch_data.add_dense([dense])
        >>>         batch_data.add_sparse(batch_sparse_ids)
        >>>         batch_data.add_target(target)
        >>>         python_train_batch = ctx.get_embedding_from_data(batch_data)
        >>>         (output, target) = ctx.forward(python_train_batch)
    """

    def __init__(
        self,
        preprocess_mode: PreprocessMode,
        model: Optional[torch.nn.Module] = None,
        embedding_config: Optional[EmbeddingConfig] = None,
        *args,
        **kwargs,
    ):
        """
        Arguments:
            preprocess_mode (PreprocessMode): Different preprocess mode effect the behavior of ``prepare_features``.
            model (torch.nn.Module): Torch model matched with embeddings in this context.
            embedding_config (EmbeddingConfig, optional): The embedding configuration that will be sent to the embedding server.
        """
        super(EmbeddingCtx, self).__init__(*args, **kwargs)
        self.preprocess_mode = preprocess_mode
        self.model = model
        self.embedding_config = embedding_config or _get_default_embedding_config()

        self.current_batch = None

    def _enter(self):
        if self.embedding_config is not None:
            self.configure_embedding_servers(self.embedding_config)

    def configure_embedding_servers(
        self,
        embedding_config: EmbeddingConfig,
    ):
        """Apply Embedding config to embedding servers.
        Arguments:
            embedding_config (EmbeddingConfig): The embedding configuration that will be sent to the embedding server.
        """
        self.common_context.configure_embedding_servers(
            embedding_config.emb_initialization[0],
            embedding_config.emb_initialization[1],
            embedding_config.admit_probability,
            embedding_config.weight_bound > 0,
            embedding_config.weight_bound,
        )

    def forward(
        self, batch: PythonTrainBatch
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """Call `prepare_features` and then do a forward step of the model in context.

        Arguments:
            batch (PythonTrainBatch): Training data provided by PersiaML upstream including
                dense, target, sparse data and meta info.

        Returns:
            the tuple of output data and target data.
        """
        assert self.model is not None, "model not found, please init context with model"
        dense, sparse, target = self.prepare_features(batch)
        output = self.model(dense, sparse)
        return (output, target)

    def prepare_features(
        self, batch: PythonTrainBatch
    ) -> Tuple[torch.Tensor, List[torch.Tensor], Optional[torch.Tensor]]:
        """Converts the dense, sparse and target raw data in``PythonTrainBatch`` to `torch.Tensor``.

        Arguments:
            batch (PythonTrainBatch): Training data provided by PersiaML upstream including
                dense, target, sparse data and meta info.

        Returns:
            the tuple of dense data, list of sparse data and target data.
        """
        if self.preprocess_mode == PreprocessMode.INFERENCE:
            batch.target_tensor = None
        else:
            # pytype: disable=attribute-error
            batch.target = batch.consume_all_targets()
            # pytype: enable=attribute-error
            assert len(batch.target) == 1
            batch.target = batch.target[0]
            batch.target_tensor = _cast_dlpack2torch_tensor(batch.target)

        is_training = self.preprocess_mode == PreprocessMode.TRAIN  # cache property

        # pytype: disable=attribute-error
        batch.dense = batch.consume_all_dense_features()
        # pytype: enable=attribute-error
        batch.dense = batch.dense[0]
        batch.dense_tensor = _cast_dlpack2torch_tensor(batch.dense)

        # pytype: disable=attribute-error
        batch.emb = batch.consume_all_sparse_features()
        # pytype: enable=attribute-error

        batch.emb_slot = []
        # sparse embedding processing
        emb_tensors, forward_tensors = [], []

        for emb in batch.emb:
            if emb.is_raw_embedding():
                # no duplicate id in raw_id_tensor
                (
                    raw_embedding,
                    index,
                    non_empty_index,
                    sample_id_num,
                ) = emb.get_raw_embedding()

                batch.emb_slot.append([raw_embedding, index, non_empty_index])
                distinct_id_tensor = _cast_dlpack2torch_tensor(raw_embedding)
                index_tensor = _cast_dlpack2torch_tensor(
                    index
                )  # tensor shape (1, batch_size * sample_fixed_size)
                max_index = index_tensor.max()
                size_of_distinct_id_tensor = distinct_id_tensor.shape[0]

                assert (
                    max_index < size_of_distinct_id_tensor
                ), "raw embedding select index larger than tensor"

                non_empty_index_tensor = _cast_dlpack2torch_tensor(
                    non_empty_index
                )  # tensor shape (-1), variable length

                batch_size = len(sample_id_num)
                dim = distinct_id_tensor.shape[-1]
                sample_fixed_size = index_tensor.shape[1] // batch_size
                index_select_raw_tensor = distinct_id_tensor.index_select(
                    0, index_tensor.view(-1)
                )
                index_select_raw_tensor.requires_grad = is_training

                raw_fixed_size_tensor = index_select_raw_tensor.view(
                    -1, sample_fixed_size, dim
                )
                mask = (
                    index_tensor.view(batch_size, sample_fixed_size, 1) != 0
                ).half()  # generate mask
                raw_fixed_size_tensor_with_mask = torch.cat(
                    [raw_fixed_size_tensor, mask], dim=2
                )
                emb_tensors.append(
                    (
                        raw_embedding.name,
                        distinct_id_tensor,
                        index_tensor,
                        non_empty_index_tensor,
                        index_select_raw_tensor,
                    )
                )
                forward_tensors.append(raw_fixed_size_tensor_with_mask)
            else:
                emb = emb.get_sum_embedding()
                batch.emb_slot.append([emb])
                sum_tensor = _cast_dlpack2torch_tensor(emb, requires_grad=is_training)
                forward_tensors.append(sum_tensor)
                emb_tensors.append((emb.name, None, None, None, sum_tensor))

        batch.forward_tensors = forward_tensors
        batch.emb_tensors = emb_tensors
        self.current_batch = batch

        return batch.dense_tensor, batch.forward_tensors, batch.target_tensor

    def dump_checkpoint(
        self,
        dst_dir: str,
        dense_filename: str = "dense.pt",
        jit_dense_filename: str = "jit_dense.pt",
        blocking: bool = True,
        with_jit_model: bool = False,
    ):
        """Dump the dense and sparse checkpoint to destination directory.

        Arguments:
            dst_dir (str): Destination directory.
            dense_filename (str, optional): Dense checkpoint filename.
            jit_dense_filename (str, optional): Jit dense checkpoint filename.
            blocking (bool, optional): Dump embedding checkpoint in blocking mode or not.
            with_jit_model (bool, optional): Dump jit script dense checkpoint or not.
        """
        assert self.model is not None, "model not found, please init context with model"

        if with_jit_model:
            self.dump_dense(self.model, dst_dir, jit_dense_filename, True)
        self.dump_dense(self.model, dst_dir, dense_filename)

        self.dump_embedding(dst_dir, blocking=blocking)

    def load_checkpoint(
        self,
        src_dir: str,
        map_location: Optional[str] = None,
        dense_filename: str = "dense.pt",
        blocking: bool = True,
    ):
        """Load the dense and sparse checkpoint from source directory.

        Arguments:
            src_dir (str): Source directory.
            map_location (str, optional): Load the dense checkpoint to specific device.
            dense_filename (str, optional): Dense checkpoint filename.
            blocking (bool, optional): Dump embedding checkpoint in blocking mode or not.
        """
        assert self.model is not None, "model not found, please init context with model"

        dense_model_filepath = os.path.join(src_dir, dense_filename)
        if os.path.exists(dense_model_filepath):
            self.load_dense(self.model, dense_model_filepath, map_location=map_location)

        self.load_embedding(src_dir, blocking=blocking)

    def dump_embedding(self, dst_dir: str, blocking: bool = True):
        """Dump embeddings to ``dst_dir``. Use ``TrainCtx.wait_for_dump_embedding`` to wait until finished
        if ``blocking=False``.

        Arguments:
            dst_dir (str): Destination directory.
            blocking (bool, optional): Dump embedding in blocking mode or not.
        """
        self.common_context.dump(dst_dir)
        if blocking:
            self.wait_for_dump_embedding()

    def load_embedding(self, src_dir: str, blocking: bool = True):
        """Load embeddings from ``src_dir``. Use ``TrainCtx.wait_for_load_embedding`` to wait until finished
        if ``blocking=False``.

        Arguments:
            src_dir (str): Directory to load embeddings.
            blocking (bool, optional): Dump embedding in blocking mode or not.
        """
        self.common_context.load(src_dir)
        if blocking:
            self.wait_for_load_embedding()

    def dump_dense(
        self,
        dense: Union[torch.nn.Module, torch.optim.Optimizer],
        dst_dir: str,
        file_name: str,
        is_jit: bool = False,
    ):
        """Dump torch model or optimizer to ``dst_dir`` as ``file_name``.

        Arguments:
            dense (torch.nn.Module or torch.optim.Optimizer): dense model or optimizer to be dumped.
            dst_dir (str): Destination directory.
            file_name (str): Destination filename.
            is_jit (bool, optional): whether to dump model as jit script.
        """
        buffer = io.BytesIO()
        if not is_jit:
            torch.save(dense.state_dict(), buffer)
        else:
            assert isinstance(
                dense, torch.nn.Module
            ), "saving an optimizer as jit script"
            jit_model = torch.jit.script(dense)
            torch.jit.save(jit_model, buffer)
        bytes_model = buffer.getvalue()
        self.common_context.dump_to_file(bytes_model, dst_dir, file_name)

    def load_dense(
        self,
        dense: Union[torch.nn.Module, torch.optim.Optimizer],
        src_filepath: str,
        map_location: Optional[str] = None,
    ):
        """Load the torch state dict from source file path.

        Arguments:
            dense (torch.nn.Module or torch.optim.Optimizer): dense model or optimizer to restore.
            src_filepath (str): Source file path.
            map_location (str, optional): Load the dense checkpoint to specific device.
        """
        dense_bytes = self.common_context.read_from_file(src_filepath)
        buffer = io.BytesIO(dense_bytes)
        buffer.seek(0)
        state_dict = torch.load(buffer, map_location=map_location)
        dense.load_state_dict(state_dict)

    def wait_for_dump_embedding(self):
        """Wait for the embedding dump process."""
        self.common_context.wait_for_emb_dumping()

    def wait_for_load_embedding(self):
        """Wait for the embedding load process."""
        self.common_context.wait_for_emb_loading()

    def get_embedding_size(self) -> List[int]:
        """Get number of ids on all embedding servers."""
        return self.common_context.get_embedding_size()

    def clear_embeddings(self):
        """Clear all embeddings on all embedding servers."""
        self.common_context.clear_embeddings()

    def get_embedding_from_data(
        self, data: PyPersiaBatchData, device_id: Optional[int] = None
    ) -> PythonTrainBatch:
        """Get embeddings of the input batch data.

        Arguments:
            data (PyPersiaBatchData): Input data without embeddings.
            device_id (int, optional): The CUDA device to use for this process.

        Returns:
            Input data with embeddings.
        """
        return self.common_context.get_embedding_from_data(data, device_id)

    def get_embedding_from_bytes(
        self, data: bytes, device_id: Optional[int] = None
    ) -> PythonTrainBatch:
        """Get embeddings of the serialized input batch data.

        Arguments:
            data (PyPersiaBatchData): Serialized input data without embeddings.
            device_id (int, optional): The CUDA device to use for this process.

        Returns:
            Input data with embeddings.
        """
        return self.common_context.get_embedding_from_bytes(data, device_id)

class TrainCtx(EmbeddingCtx):
    r"""Subclass of ``EmbeddingCtx`` that provide the backward ability to update the sparse embedding.

    Example:
        >>> import torch
        >>> model = get_dnn_model()
        >>> sparse_optimizer = persia.sparse.optim.SGD(lr=1e-3)
        >>> dense_optimizer = torch.optim.SGD(lr=1e-3)
        >>> loss_fn = torch.nn.BCELoss(reduction="mean")
        >>> with TrainCtx(
        >>>     sparse_optimizer,
        >>>     dense_optimizer,
        >>> ) as ctx:
        >>>     parallel_model = ctx.model
        >>>     for batch_data in dataloder:
        >>>         dense, sparse, target = ctx.prepare_features(data)
        >>>         output = parallel_model(dense, sparse)
        >>>         loss = loss_fn(output, target)
        >>>         scaled_loss = ctx.backward(loss)
    """

    def __init__(
        self,
        sparse_optimizer: Optimizer,
        dense_optimizer: torch.optim.Optimizer,
        grad_scalar_update_factor: float = 4,
        backward_buffer_size: int = 10,
        backward_workers_size: int = 8,
        grad_update_buffer_size: int = 60,
        lookup_emb_directly: bool = True,
        mixed_precision: bool = True,
        distributed_option: Optional[DistributedBaseOption] = None,
        *args,
        **kwargs,
    ):
        """
        Arguments:
            sparse_optimizer (persia.sparse.optim.Optimizer): Optimizer for the embeddings.
            dense_optimizer (torch.optim.Optimizer): Optimizer for dense parameters.
            grad_scalar_update_factor (float, optional): Update factor of ``Gradscalar`` to ensure loss scale finitely if set ``mixed_precision=True``.
            backward_buffer_size (int, optional): Max number of not updated gradients queued.
            backward_workers_size (int, optional): Number of workers sending embedding gradients in parallel.
            grad_update_buffer_size (int, optional): Number of reference cache , hold the gradient tensor reference to avoid
                meet dangle data in gradient backward phase.
            lookup_emb_directly (bool, optional): Lookup embedding directly without isolation data compose.
            mixed_precision (bool): Enable mixed_precision or not.
            distributed_option (DistributedBaseOption, optional): DistributedOption to converted model to dataparallel model.
        """
        super(TrainCtx, self).__init__(PreprocessMode.TRAIN, *args, **kwargs)

        assert (
            sparse_optimizer is not None
        ), "Sparse_optimizer should not be none in train context"
        assert grad_scalar_update_factor > 0, "grad scalar should greater than zero"
        assert (
            self.model is not None
        ), "Model not found, please init context with pytorch model"

        assert grad_scalar_update_factor > 0, "grad scalar should greater than zero"

        world_size = env.get_world_size()
        assert world_size != -1, "WORLD_SIZE not set"
        rank_id = env.get_rank()
        assert rank_id != -1, "RANK not set"

        self.world_size = world_size
        self.rank_id = rank_id

        assert not mixed_precision or (
            mixed_precision and torch.cuda.is_available()
        ), "Mixed precision training only support on cuda device."
        self.mixed_precision = mixed_precision

        if self.mixed_precision:
            self.grad_scalar_update_factor = grad_scalar_update_factor
            self.grad_scaler = torch.cuda.amp.GradScaler()
            self.update_times = 0

        if world_size > 1:
            distributed_option = distributed_option or get_default_distributed_option()
            not_env_file = not distributed_option.init_with_env_file()
            not_exists_master_addr = not distributed_option.master_addr
            if not_env_file and not_exists_master_addr:
                master_addr = self._get_master_addr()
            else:
                master_addr = None

            model, dense_optimizer = distributed_option.convert2distributed_model(
                self.model,
                world_size,
                rank_id,
                self.device_id,
                master_addr=master_addr,
                optimizer=dense_optimizer,
            )
            self.model = model
            _logger.info("Distributed training context init done.")
        else:
            _logger.info("SingleMachine training context init done.")

        self.dense_optimizer = dense_optimizer
        self.sparse_optimizer = sparse_optimizer

        self.wait_servers_ready()

        if lookup_emb_directly:
            init_rpc_client_num = self._init_middlewrae_rpc_client()
            _logger.info(f"Successfully init {init_rpc_client_num} rpc client")

        self.backward_workers_size = backward_workers_size

        from persia.prelude import PyBackward

        self.grad_queue = Queue(grad_update_buffer_size)
        self.backward_engine = PyBackward(backward_buffer_size)

    def _enter(self):
        super()._enter()

        self.sparse_optimizer.apply()
        self.backward_engine.launch(self.backward_workers_size)

    def _exit(self):
        super()._exit()

        self.backward_engine.shutdown()

    def _get_master_addr(self) -> str:
        """Get leader(rank 0) ip address."""
        if self.rank_id == 0:
            master_addr = socket.gethostbyname(socket.gethostname())
            self.common_context.init_master_discovery_service(master_addr)
            _logger.info(f"init addr is {master_addr}")
        else:
            self.common_context.init_master_discovery_service(None)
            master_addr = self.common_context.master_addr
            _logger.info(f"master addr is {master_addr}")
        return master_addr

    def _init_middlewrae_rpc_client(self) -> int:
        middleware_addr_list = self.common_context.get_middleware_addr_list()
        assert len(middleware_addr_list) > 0, "Not available middleware."
        for middleware_addr in middleware_addr_list:
            self.common_context.init_rpc_client_with_addr(middleware_addr)
        return len(middleware_addr_list)

    def wait_servers_ready(self):
        """query embedding server to check if servers are ready"""

        self.common_context.init_nats_publisher(self.world_size)
        self.common_context.wait_servers_ready()

    def backward(
        self, loss: torch.Tensor, embedding_gradient_check_frequency: int = 20
    ) -> torch.Tensor:
        """Compute the gradient of current dense and sparse tensors.

        This method supports mixed precision training. Depending on whether the current embedding gradient is finite or not, ``GradScalar``
        can update the scale automatically to restrict to parameters in finite range.

        Arguments:
            loss (torch.Tensor): Loss of current batch.
            embedding_gradient_check_frequency (int, optional): The frequency to check gradient finite or not for current embedding.
        """
        if self.mixed_precision:
            loss = self.grad_scaler.scale(loss)
            scale = self.grad_scaler.get_scale()
        else:
            scale = 1  # Always equal to 1 when disable mixed_precision training

        loss.backward()

        finite = self._on_backward(scale, embedding_gradient_check_frequency)

        if self.mixed_precision:
            self.grad_scaler.step(self.dense_optimizer)

            if finite:
                self.grad_scaler.update()
            else:
                self.grad_scaler.update(scale / self.grad_scalar_update_factor)
        else:
            self.dense_optimizer.step()

        self.dense_optimizer.zero_grad()
        return loss

    def _on_backward(self, loss_scale: float, embedding_gradient_check_frequency: int):
        """Update the embeddings' gradients

        Arguments:
            loss_scale (float): The loss that scaled by GradScalar, loss_scale always equal to 1 for cpu training scenes.
            embedding_gradient_check_frequency (int): The frequency to check gradient finite or not for current embedding.
        """
        if self.grad_queue.full():
            self.grad_queue.get()

        finite = True

        if (
            self.mixed_precision
            and self.update_times % embedding_gradient_check_frequency == 0
        ):
            finite = _check_finite(
                [emb[-1].grad for emb in self.current_batch.emb_tensors]
            )
            self.update_times += 1

        grad_slots, empty_grads = [], []
        gradient_batch = self.current_batch.create_gradient_batch()

        for (
            emb_name,
            distinct_id_tensor,
            index,
            non_zero_index,
            emb_tensor,
        ) in self.current_batch.emb_tensors:
            if emb_tensor.grad is None:
                gradient_batch.add_skipped_gradient(emb_name)
                empty_grads.append(emb_name)
            else:
                if distinct_id_tensor is not None:
                    if distinct_id_tensor.shape[0] > 1:
                        grad = torch.zeros_like(distinct_id_tensor, dtype=torch.float32)
                        non_zero_grad = emb_tensor.grad.index_select(
                            0, non_zero_index.view(-1)
                        ).float()
                        non_zero_index = index.view(-1)[non_zero_index.view(-1)]
                        grad.index_add_(0, non_zero_index, non_zero_grad)
                        grad = grad[1:, :]
                        is_f16_gradient = False
                    else:
                        grad = None
                else:
                    grad = emb_tensor.grad  # type: torch.Tensor
                    is_f16_gradient = True

                if grad is not None:

                    grad_slots.append(grad)
                    gradient_batch.add_gradient(
                        emb_name,
                        grad.data_ptr(),
                        grad.shape,
                        is_f16_gradient,
                        loss_scale,
                    )

        if self.device_id is not None:
            torch.cuda.synchronize()

        self.backward_engine.update_sparse_gradient_batched(gradient_batch)
        self.grad_queue.put(grad_slots)

        if len(empty_grads) > 0:
            _logger.warning(
                f"Current batch exists empty gradient tensors, num: {len(empty_grads)}, {empty_grads}"
            )
        return finite

    def dump_checkpoint(
        self,
        dst_dir: str,
        dense_filename: str = "dense.pt",
        jit_dense_filename: str = "jit_dense.pt",
        opt_filename: str = "opt.pt",
        blocking: bool = True,
        with_jit_model: bool = False,
    ):
        """Dump the dense and sparse checkpoint to destination directory.

        Arguments:
            dst_dir (str): Destination directory.
            dense_filename (str, optional): Dense checkpoint filename.
            jit_dense_filename (str, optional): Jit dense checkpoint filename.
            opt_filename (str, optional): Optimizer checkpoint filename.
            blocking (bool, optional): Dump embedding checkpoint in blocking mode or not.
            with_jit_model (bool, optional): Dump dense checkpoint as jit script or not.
        """
        super().dump_checkpoint(
            dst_dir,
            dense_filename=dense_filename,
            jit_dense_filename=jit_dense_filename,
            blocking=blocking,
            with_jit_model=with_jit_model,
        )

        self.dump_dense(self.dense_optimizer, dst_dir, opt_filename)

    def load_checkpoint(
        self,
        src_dir: str,
        map_location: Optional[str] = None,
        dense_filename: str = "dense.pt",
        opt_filename: str = "opt.pt",
        blocking: bool = True,
    ):
        """Load the dense and sparse checkpoint from source directory.

        Arguments:
            src_dir (str): Source directory.
            map_location (str, optional): Load the dense checkpoint to specific device.
            dense_filename (str, optional): Dense checkpoint filename.
            opt_filename (str, optional): Optimizer checkpoint filename.
            blocking (bool, optional): Dump embedding checkpoint in blocking mode or not.
        """
        super().load_checkpoint(
            src_dir,
            map_location=map_location,
            dense_filename=dense_filename,
            blocking=blocking,
        )

        optimizer_filepath = os.path.join(src_dir, opt_filename)
        if os.path.exists(optimizer_filepath):
            self.load_dense(self.dense_optimizer, optimizer_filepath)

def cnt_ctx() -> Optional[BaseCtx]:
    """Get the BaseCtx recently entered."""
    return _CURRENT_CXT

def eval_ctx(*args, **kwargs) -> EmbeddingCtx:
    """Get the ``EmbeddingCtx`` with the ``PreprocessMode.EVAL`` mode."""
    return EmbeddingCtx(PreprocessMode.EVAL, *args, **kwargs)

class InferCtx(EmbeddingCtx):
    r"""Subclass of ``EmbeddingCtx`` that provide the forward ability without nats servers.

    Example:
        >>> from persia.ctx import InferCtx
        >>> persia_context = InferCtx()
        >>> batch = persia_context.get_embedding_from_bytes(batch, device_id)
        >>> model_input = persia_context.prepare_features(batch)
    """

    def __init__(
        self,
        middleware_addrs: List[str],
        *args,
        **kwargs,
    ):
        """
        Arguments:
            middleware_addrs (List[str]): middleware address(ip:port) list.
        """
        super(InferCtx, self).__init__(PreprocessMode.INFERENCE, *args, **kwargs)

        for addr in middleware_addrs:
            self.common_context.init_rpc_client_with_addr(addr)

    r"""Wait for middleware and embedding server ready for serving."""

    def wait_for_serving(self):
        self.common_context.wait_for_serving()


from __future__ import unicode_literals
from django.db import models
import facebook
import logging
import datetime
import stomp
import json

import moneyed
from djmoney.models.fields import MoneyField
from djmoney.models.managers import money_manager
# Create your models here.
logger = logging.getLogger(__name__)

STOMP_SERVER_URL = '35.162.74.72'
STOMP_PORT = 61613

class BdayAppUsermanager(models.Manager):

    def create_user(self,email_id=None):
        user = self.create(email_id=email_id,creation_date=datetime.datetime.now())
        return user

class BdayAppUser(models.Model):
    email_id = models.EmailField(null=True)
    creation_date = models.DateTimeField()
    objects = BdayAppUsermanager()

    class Meta:
        ordering = ['-creation_date']

    def update_email_id(self,email_id):
        self.email_id = email_id
        self.save()

class EventManager(models.Manager):

    def create_event(self,admin,name,start_date,end_date=None,picture=None,members=None,type='DEFAULT'):
        event = self.create(name=name,start_date=start_date,end_date=end_date,picture=picture,type=type,creation_date=datetime.datetime.now())
        logger.info(event.id)
        event.admin.add(admin)
        if members is not None:
            for member in members:
                event.members.add(member)
        event.save()
        UnreadChatBuffer.objects.create(event=event, user=admin, last_read_chat=None)
        return event

class Event(models.Model):
    # id = models.BigAutoField()
    picture = models.ImageField(upload_to='images/%Y/%m/%d', null=True)
    admin = models.ManyToManyField(BdayAppUser, related_name='admin',null=True,blank=True)
    members = models.ManyToManyField(BdayAppUser, related_name='members', null=True,blank=True)
    name = models.CharField(max_length=100)
    type = models.CharField(max_length=100, default='DEFAULT')
    creation_date = models.DateTimeField()
    start_date = models.DateTimeField()
    end_date = models.DateTimeField(null=True,blank=True)

    objects = EventManager()

    class Meta:
        ordering = ['-creation_date', 'name']

    def __update_picture(self,picture):
        if picture is not None:
            self.picture = picture
            self.save()

    def __update_admins(self, users):
        if users is not None:
            for user in users:
                self.admin.add(user)
            self.save()
    def __remove_admin(self,users):
        if users is not None:
            for user in users:
                self.admin.add(user)
            self.save()
    # Users is a list
    def __update_members(self,users):
        if users is not None:
            for user in users:
                self.members.add(user)
            self.save()
    def __remove_members(self,users):
        if users is not None:
            for user in users:
                self.members.remove(user)
            self.save()
    def __update_name(self,name):
        if name is not None:
            self.name = name
            self.save()

    def __update_start_date(self,start_date):
        if start_date is not None:
            self.start_date = start_date
            self.save()

    def __update_end_date(self,end_date):
        if end_date is not None:
            self.end_date = end_date
            self.save()

    def __remove_picture(self):
        self.picture = None
        self.save()

    def update(self,picture=None,admins=None,members=None,name=None,start_date=None,end_date=None):
        self.__update_picture(picture)
        self.__update_admins(admins)
        self.__update_members(members)
        self.__update_name(name)
        self.__update_start_date(start_date)
        self.__update_end_date(end_date)

class ReminderManager(models.Manager):

    def create_reminder(self,name,date,user,event_associated_flag=False,associated_event=None,picture=None,type=None,reminder_for=None,picture_file=None):
        reminder = self.create(name=name,reminder_date=date,user=user,event_associated_flag=event_associated_flag,associated_event=associated_event,picture=picture,type=type,reminder_for=reminder_for,picture_file=picture_file,creation_date = datetime.datetime.now())
        return reminder

class Reminder(models.Model):
    picture = models.URLField(null=True,blank=True)
    picture_file = models.ImageField(upload_to='images/%Y/%m/%d', null=True)
    reminder_date = models.DateField()
    name = models.CharField(max_length=50)
    event_associated_flag = models.BooleanField(default=False)
    associated_event = models.OneToOneField(Event, default=None, null=True, blank=True)
    user = models.ForeignKey(BdayAppUser)
    creation_date = models.DateTimeField()
    type = models.CharField(max_length=50, default='CUSTOM')
    reminder_for = models.ForeignKey(BdayAppUser, default=None, null=True, blank=True,related_name='reminder_for')
    objects = ReminderManager()

    class Meta:
        ordering = ['-creation_date', 'name']

    def __update_picture(self,picture):
        if picture is not None:
            self.picture = picture
            self.save()

    def __update_reminder_date(self,date):
        if date is not None:
            self.reminder_date = date
            self.save()

    def __update_name(self,name):
        if name is not None:
            self.name = name
            self.save()

    def __update_event_associated_flag(self,flag):
        if flag is not None:
            self.event_associated_flag = flag
            self.save()

    def __update_associated_event(self,event):
        if event is not None:
            self.associated_event = event
            self.save()

    def update_reminder(self,picture=None,date=None,name=None,event_associated_flag=None,associated_event=None):
        self.__update_picture(picture)
        self.__update_reminder_date(date)
        self.__update_name(name)
        self.__update_event_associated_flag(event_associated_flag)
        self.__update_associated_event(associated_event)

class UserProfileManager(models.Manager):

    def create_user_profile(self,access_token,profile_type,user):
        try:
            birthday, profile_id, user_name, picture, first_name, last_name, gender = None, None,None,None, None, None, None
            #total_number_of_friends_on_app = 0
            if profile_type == "FB":
                graph = facebook.GraphAPI(access_token=access_token, version='2.8')
                profile_id = graph.get_object(id='me')['id']
                #user_name = graph.get_object(id='me')['name']
                picture = graph.get_object(id='me', fields='picture')['picture']['data']['url']
                response_string_ = graph.get_object(id='me',fields = 'birthday,email,first_name,last_name,gender,name')
                if 'birthday' in response_string_:
                    birthday = response_string_['birthday']
                if 'email' in response_string_:
                    if user.email_id is None:
                        user.email_id = response_string_['email']
                        user.save()
                if 'name' in response_string_:
                    user_name = response_string_['name']

                if 'first_name' in response_string_:
                    first_name = response_string_['first_name']

                if 'last_name' in response_string_:
                    last_name = response_string_['last_name']

                if 'gender' in response_string_:
                    gender = response_string_['gender']

                user_profile = self.create(access_token=access_token,profile_id=profile_id,birthday=birthday,picture=picture,profile_type=profile_type,user=user,user_name=user_name,first_name=first_name,last_name=last_name,gender=gender,creation_date = datetime.datetime.now(),phone_number=None)
                user_profile.save()
                friends = graph.get_connections('me', "friends")
                while True:
                    try:
                        for friend in friends['data']:
                            id = friend['id'].encode('utf-8')
                            try:
                                if UserProfile.objects.get(profile_id=id) is not None:
                                    friends_profile = UserProfile.objects.get(profile_id=id)
                                    user_profile.app_friends.add(friends_profile.user)
                                    friends_profile.app_friends.add(user)
                                    if friends_profile.birthday is not None:
                                        friends_bday_reminder = Reminder.objects.create_reminder(friends_profile.user_name+"'s Bday",
                                                                         datetime.datetime.strptime(
                                                                            friends_profile.birthday, '%m/%d/%Y'),
                                                                         user_profile.user, False, None, friends_profile.picture,
                                                                         type='FACEBOOK',
                                                                         reminder_for=friends_profile.user
                                                                         )
                                        friends_bday_reminder.save()
                                    if user_profile.birthday is not None:
                                        self_bday_reminder_for_friends = Reminder.objects.create_reminder(user_profile.user_name+"'s Bday",
                                                                                datetime.datetime.strptime(
                                                                                user_profile.birthday,
                                                                                '%m/%d/%Y'),
                                                                                friends_profile.user,
                                                                                False, None,
                                                                                user_profile.picture,
                                                                                type='FACEBOOK',
                                                                                reminder_for=user_profile.user
                                                                                )
                                        self_bday_reminder_for_friends.save()
                            except UserProfile.DoesNotExist:
                                continue
                        # Attempt to make a request to the next page of data, if it exists.
                        friends = graph.get_connections("me", "friends", after=friends['paging']['cursors']['after'])
                    except KeyError:
                        # When there are no more pages (['paging']['next']), break from the
                        # loop and end the script.
                        break
                logger.debug("facebook type user profile created successfully")
                user_profile.save()
                return user_profile
        except Exception as e:
            logger.error(e)
            logger.error("could not create user profile")

class UserProfile(models.Model):
    # id = models.BigAutoField()
    first_name = models.CharField(max_length=50, null=True)
    last_name = models.CharField(max_length=50, null=True)
    phone_number = models.DecimalField(max_digits=10, null=True, decimal_places=0)
    gender = models.CharField(max_length=15, null=True)
    user_name = models.CharField(max_length=50)
    profile_id = models.CharField(max_length=50, unique=True)
    picture = models.URLField()
    access_token = models.CharField(max_length=150)
    user = models.OneToOneField(BdayAppUser, null=True, blank=True, related_name='user_id')
    app_friends = models.ManyToManyField(BdayAppUser, null=True, symmetrical=True, related_name='friends')
    profile_type = models.CharField(max_length=2, choices=(
        ('FB', 'FACEBOOK'),
        ('GP', 'GOOGLEPLUS')
    ))
    birthday = models.CharField(max_length=10,null=True,blank=True)
    street_address = models.CharField(max_length=10, null=True, blank=True)
    address_line_2 = models.CharField(max_length=10, null=True, blank=True)
    city = models.CharField(max_length=10, null=True, blank=True)
    state = models.CharField(max_length=10, null=True, blank=True)
    pincode = models.CharField(max_length=10, null=True, blank=True)
    #address = models.CharField(max_length=100,null=True,blank=True)
    contact_privacy_field = models.CharField(max_length=10, choices=(
        ('PUBLIC', 'PB'),
        ('PRIVATE', 'PV')
    ), default='PV')
    creation_date = models.DateTimeField()
    objects = UserProfileManager()
    class Meta:
        ordering = ['-creation_date']

    def __get_user_name(self):
        try:
            if self.profile_type is "FB":
                graph = facebook.GraphAPI(access_token=self.access_token, version='2.8')
                self.user_name = graph.get_object(id='me')['name']
            logger.info("user_name stored successfully")
        except Exception as e:
            logger.error(e)
    def __get_profile_id(self):
        try:
            if self.profile_type is "FB":
                graph = facebook.GraphAPI(access_token=self.access_token, version='2.8')
                self.profile_id = graph.get_object(id='me')['id']
            logger.info("profile_id stored successfully")
        except Exception as e:
            logger.error(e)
    def __get_picture(self):
        try:
            if self.profile_type is "FB":
                graph = facebook.GraphAPI(access_token=self.access_token, version='2.8')
                self.picture = graph.get_object(id='me', fields='picture')['picture']['data']['url']
            logger.info("profile picture updated successfully")
        except Exception as e:
            logger.error(e)
    def __get_birthday(self):
        try:
            if self.profile_type is "FB":
                graph = facebook.GraphAPI(access_token=self.access_token, version='2.8')
                response_string_ = graph.get_object(id='me', fields='birthday')
                if 'birthday' in response_string_:
                    self.birthday = response_string_['birthday']
                    logger.info("birthday is " +self.birthday)
            logger.info("birthday stored successfully")
        except Exception as e:
            logger.error(e)
    def __get_email(self):
        try:
            if self.profile_type is "FB":
                graph = facebook.GraphAPI(access_token=self.access_token, version='2.8')
                response_string_ = graph.get_object(id='me', fields='email')
                if 'email' in response_string_:
                    logger.info("email is : " + response_string_['email'])
                    return response_string_['email']
            logger.info("email_id stored successfully")
        except Exception as e:
            logger.error(e)
    def __get_friends(self):
        friends_list = []
        try:
            if self.profile_type is "FB":
                graph = facebook.GraphAPI(access_token=self.access_token, version='2.8')
                friends = graph.get_connections('me', "friends")
                while True:
                    try:
                        for friend in friends['data']:
                            id = friend['id'].encode('utf-8')
                            friends_list.append(id)
                        # Attempt to make a request to the next page of data, if it exists.
                        friends = graph.get_connections("me", "friends", after=friends['paging']['cursors']['after'])
                    except KeyError:
                        # When there are no more pages (['paging']['next']), break from the
                        # loop and end the script.
                        break
                return friends_list
        except Exception as e:
            logger.error(e)
            return friends_list
    def update_user_name(self,user_name):
        try:
            if user_name is not None:
                self.user_name = user_name
                self.save()
                logger.info("user_name updated")
        except Exception as e:
            logger.error(e)
    def update_picture(self,picture):
        try:
            if picture is not None:
                self.picture = picture
                self.save()
                logger.info("picture updated successfully")
        except Exception as e:
            logger.error(e)
    def update_access_token(self,access_token):
        try:
            if self.access_token is not None:
                self.access_token = access_token
                logger.info("access_token updated successfully")
        except Exception as e:
            logger.error(e)
    def update_picture_from_social_network(self):
        try:
            if self.profile_type is "FB":
                self.__get_picture()
                self.save()
                logger.info("picture updated from social network")
        except Exception as e:
            logger.error(e)
    def update_app_friends(self):
        try:
            users = self.__get_friends()
            if users is not None:
                for profile_id in users:
                    try:
                        if UserProfile.objects.get(profile_id=profile_id) is not None:
                            friends_profile = UserProfile.objects.get(profile_id=profile_id)
                            self.app_friends.add(friends_profile.user)
                            if friends_profile.birthday is not None:
                                friends_bday_reminder = Reminder.objects.create_reminder(friends_profile.user_name,
                                                                                datetime.datetime.strptime(
                                                                                friends_profile.birthday,
                                                                                '%m/%d/%Y'),
                                                                                self.user, False, None,
                                                                                friends_profile.picture
                                                                                )
                                friends_bday_reminder.save()
                            if self.birthday is not None:
                                self_bday_reminder_for_friends = Reminder.objects.create_reminder(self.user_name,
                                                                                datetime.datetime.strptime(
                                                                                self.birthday,
                                                                                '%m/%d/%Y'),
                                                                                friends_profile.user, False, None,
                                                                                self.picture
                                                                                )
                                self_bday_reminder_for_friends.save()
                    except UserProfile.DoesNotExist:
                        continue
            self.save()
            logger.info("friend list updated successfully")
        except Exception as e:
            logger.error(e)

class MessageManager(models.Manager):
    def create_message(self,message,to_,from_,is_read=False,read_date=None):
        message = self.create(message=message,to_field=to_,from_field=from_,creation_date=datetime.datetime.now(),is_read=False,read_date=None)
        return message

class Message(models.Model):
    message = models.CharField(max_length=1000)
    to_field = models.ForeignKey(BdayAppUser, related_name='to_')
    from_field = models.ForeignKey(BdayAppUser, related_name='from_')
    creation_date = models.DateTimeField()
    is_read = models.BooleanField(default=False)
    read_date = models.DateTimeField(null=True,blank=True)
    objects = MessageManager()

    class Meta:
        ordering = ['-creation_date']

class NotificationManager(models.Manager):
    def create_notification(self,message,associated_user,url,is_read=False,read_date=None,type=None,event_reminder_type=None,event_for_user=None,event=None):
        notification = self.create(message=message,associated_user=associated_user,creation_date=datetime.datetime.now(),url=url,is_read=False,type=type,event_reminder_type=event_reminder_type,event_for_user=event_for_user,event=event)
        return notification

class Notification(models.Model):
    message = models.CharField(max_length=1000)
    associated_user = models.ForeignKey(BdayAppUser)
    type = models.CharField(max_length=500,null=True,blank=True)
    event_reminder_type = models.CharField(max_length=500,null=True,blank=True)
    event_for_user = models.ForeignKey(BdayAppUser, related_name='event_for_user', null=True, blank=True)
    event = models.ForeignKey(Event, null=True, blank=True)
    creation_date = models.DateTimeField()
    is_read = models.BooleanField(default=False)
    url = models.URLField(null=True,blank=True)
    read_date = models.DateTimeField(null=True, blank=True)
    objects = NotificationManager()
    class Meta:
        ordering = ['-creation_date']

class WishManager(models.Manager):
    def create_wish(self,url,name,website_name,website_url,price,picture,user=None,event=None):
        wish = self.create(url=url,name=name,website_name=website_name,website_url=website_url,
                           price=price,picture=picture,user=user,event=event,creation_date=datetime.datetime.now())
        return wish

class Wish(models.Model):
    url = models.URLField()
    name = models.CharField(max_length=500)
    website_name = models.CharField(max_length=500)
    website_url = models.URLField()
    price = models.FloatField()
    picture = models.URLField()
    creation_date = models.DateTimeField()
    user = models.ForeignKey(BdayAppUser,null=True)
    event = models.ForeignKey(Event, null=True)
    objects = WishManager()
    class Meta:
        ordering = ['-creation_date']

class ChatManager(models.Manager):
    def create_event_chat(self, event, user,wish = None,message_field=None, url_field=None, file_field = None, type='LR',creation_date = datetime.datetime.now()):
        if message_field is None and url_field is None and file_field is None:
            pass
        else:
            event_chat = self.create(event=event,wish = wish,message_field = message_field, url_field = url_field, file_field = file_field, user=user,type=type,creation_date = creation_date)
            return event_chat

class EventChat(models.Model):
    event = models.ForeignKey(Event)
    message_field = models.CharField(max_length=100,null=True)
    creation_date = models.DateTimeField()
    type = models.CharField(max_length=100)
    url_field = models.URLField(null=True)
    file_field = models.ImageField(upload_to='images/%Y/%m/%d', null=True)
    wish = models.ForeignKey(Wish,null=True)
    user = models.ForeignKey(BdayAppUser)
    objects = ChatManager()

    class Meta:
        ordering = ['-id']

class WalletManager(models.Manager):
    def create_wallet(self, minimum_balance, maximum_balance, type, default_currency, associated_user, associated_event, creation_date=datetime.datetime.now()):
        initial_balance = moneyed.Money(0, 'INR')
        wallet = self.create(balance=initial_balance,minimum_balance=minimum_balance, maximum_balance=maximum_balance,type=type,default_currency=default_currency,associated_user=associated_user,associated_event=associated_event, creation_date=creation_date)
        return wallet

class Wallet(models.Model):
    balance = MoneyField(max_digits=10, decimal_places=2, default_currency='INR')
    minimum_balance = MoneyField(max_digits=10, decimal_places=2, default_currency='INR')
    maximum_balance = MoneyField(max_digits=10, decimal_places=2, default_currency='INR')
    type = models.CharField(max_length=20, choices=(
        ('U','USER'),
        ('E','EVENT')
    ))
    default_currency = models.CharField(max_length=20)
    associated_user = models.OneToOneField(BdayAppUser, related_name='user_wallet', null=True)
    associated_event = models.OneToOneField(Event, related_name='event_wallet', null=True)
    creation_date = models.DateTimeField()
    objects = money_manager(WalletManager())

# class GroupieWallet(models.Model):
#     balance = MoneyField(max_digits=10, decimal_places=2, default_currency='INR')
#     minimum_balance = MoneyField(max_digits=10, decimal_places=2, default_currency='INR')
#     maximum_balance = MoneyField(max_digits=10, decimal_places=2, default_currency='INR')
#     type = models.CharField(max_length=20, choices=(
#         ('U','USER'),
#         ('E','EVENT')
#     ))
#     default_currency = models.CharField(max_length=20)
#     associated_user = models.OneToOneField(BdayAppUser, related_name='user_wallet', null=True)
#     associated_event = models.OneToOneField(Event, related_name='event_wallet', null=True)
#     creation_date = models.DateTimeField()
#     objects = money_manager(WalletManager())

class Transaction(models.Model):
    from_wallet = models.ForeignKey(Wallet, related_name='from_wallet', null=True)
    to_wallet = models.ForeignKey(Wallet, related_name='to_wallet', null=True)
    type = models.CharField(max_length=20, choices=(
        ('D','DEBIT'),
        ('C','CREDIT')
    ))
    amount = MoneyField(max_digits=10, decimal_places=2, default_currency='INR')
    default_currency = models.CharField(max_length=10)
    status = models.CharField(max_length=20, choices=(
        ('S', 'TXN_SUCCESS'),
        ('P', 'PENDING'),
        ('F', 'TXN_FAILURE'),
        ('O', 'OPEN')
    ))
    order_id = models.CharField(max_length=1000, null=True)
    external_transaction_flag = models.BooleanField()
    external_subscription_id = models.BigIntegerField(null=True)
    external_transaction_id = models.BigIntegerField(null=True)
    bank_transaction_id = models.CharField(max_length=100,null=True)
    # external_response_code = models.CharField(max_length=100, null=True)
    # external_response_message = models.CharField(max_length=100, null=True)
    transaction_date = models.DateTimeField()
    gateway_name = models.CharField(max_length=100, null=True)
    bank_name = models.CharField(max_length=100, null=True)
    payment_mode = models.CharField(max_length=100, null=True)
    promo_camp_id = models.CharField(max_length=100, null=True)
    promo_status = models.CharField(max_length=100, null=True)
    promo_response_code = models.CharField(max_length=100, null=True)
    #checksum_hash = models.CharField(max_length=1000, null=True)
    objects = money_manager(models.Manager())

class UnreadChatBuffer(models.Model):
    event = models.ForeignKey(Event)
    user = models.ForeignKey(BdayAppUser)
    last_read_chat = models.ForeignKey(EventChat,null=True,blank=True)

class Category(models.Model):
    name = models.CharField(max_length=100)
    logo = models.ImageField(upload_to='images/%Y/%m/%d')

class SubCategory(models.Model):
    name = models.CharField(max_length=100)
    logo = models.ImageField(upload_to='images/%Y/%m/%d')
    category = models.ForeignKey(Category)

class GiftStore(models.Model):
    name = models.CharField(max_length=100)
    link = models.URLField()
    affiliated_link = models.URLField()
    logo = models.URLField()
    category = models.ForeignKey(Category)
    subcategory = models.ForeignKey(SubCategory)

class OtherLifeEvents(models.Model):
    name = models.CharField(max_length=100)
    date_of_event = models.DateTimeField()
    created_date = models.DateTimeField(default=datetime.datetime.now)
    type = models.CharField(max_length=100, default='DEFAULT')
    user = models.ForeignKey(BdayAppUser, related_name='other_life_events', on_delete=models.CASCADE)

    class Meta:
        ordering = ['-created_date']

class CashoutCard(models.Model):
    type = models.CharField(max_length=100)
    terms_and_conditions = models.CharField(max_length=3000, null=True, blank=True)
    shipping_and_handling = models.CharField(max_length=3000, null=True, blank=True)
    brand_details = models.CharField(max_length=1000, null=True, blank=True)
    coupon_code = models.CharField(max_length=100)
    amount = MoneyField(max_digits=10, decimal_places=2, default_currency='INR')
    user = models.ForeignKey(BdayAppUser, related_name='cards', on_delete=models.CASCADE)
    created_date = models.DateTimeField(default=datetime.datetime.now)
    cashout_transaction_id = models.ForeignKey(Transaction, null=True, blank=True)
    associated_event = models.ForeignKey(Event, null=True, blank=True, related_name="associated_event")
    objects = money_manager(models.Manager())

class CashoutMessage(models.Model):
    text = models.CharField(max_length=400)
    user = models.ForeignKey(BdayAppUser, on_delete=models.CASCADE)
    associated_event = models.ForeignKey(Event, null=True, blank=True, related_name="associated_event")

class GroupyWallet(models.Model):
    balance = MoneyField(max_digits=10, decimal_places=2, default_currency='INR')
    minimum_balance = MoneyField(max_digits=10, decimal_places=2, default_currency='INR')
    maximum_balance = MoneyField(max_digits=10, decimal_places=2, default_currency='INR')
    default_currency = models.CharField(max_length=20)
    associated_user = models.OneToOneField(BdayAppUser, related_name='user_wallet', null=True)
    creation_date = models.DateTimeField()
    objects = money_manager(models.Manager())
    deduct_percentage = models.DecimalField(max_digits=2)


# Wrap a gamestream session around a launcher program with configurable background tasks and resolution switching
import win32api
import win32.lib.win32con as win32con
import win32gui
import pywintypes
import pyautogui
import configparser
import subprocess
import psutil
import os
import sys
from time import sleep

# Window enumeration handler function per https://www.blog.pythonlibrary.org/2014/10/20/pywin32-how-to-bring-a-window-to-front/
def windowEnumerationHandler(hwnd, top_windows):
    top_windows.append((hwnd, win32gui.GetWindowText(hwnd)))

def set_resolution(gamestream_width, gamestream_height):
    print("Switching resolution to {0}x{1}".format(gamestream_width, gamestream_height))
    devmode = pywintypes.DEVMODEType()
    devmode.PelsWidth = int(gamestream_width)
    devmode.PelsHeight = int(gamestream_height)
    devmode.Fields = win32con.DM_PELSWIDTH | win32con.DM_PELSHEIGHT
    win32api.ChangeDisplaySettings(devmode, 0)

def get_process_name(p):
    # If there are permission errors reading a process name, it's probably not the one we want, so skip it.
    try:
        p_name = p.name()
    except (PermissionError, psutil.AccessDenied):
        p_name = ""
        pass
    return p_name

def reset_launcher_resolution(gamestream_width, gamestream_height, launcher_window_name):
    # Check to ensure desired GSLP resolution is still set whenever the launcher is in focus in case it didn't reset when exiting a game
    focused_window = win32gui.GetWindowText(win32gui.GetForegroundWindow()).lstrip()
    #print("Trying to match", launcher_window_name, focused_window)
    if focused_window.startswith(launcher_window_name):
        #print("Matched")
        current_width = str(win32api.GetSystemMetrics(0))
        current_height = str(win32api.GetSystemMetrics(1))
        if current_width != gamestream_width and current_height != gamestream_height:
            print("Resolutions don't match, changing from", current_width, current_height, "to", gamestream_width, gamestream_height)
            set_resolution(gamestream_width, gamestream_height)

# Define a default config file to write if we're missing one
config_filename = 'gamestream_playnite.ini'
default_config = """[LAUNCHER]
# The path to your Playnite.FullscreenApp.exe
launcher_path = %%LOCALAPPDATA%%\Playnite\Playnite.FullscreenApp.exe
# Name of the window to watch to close the session when it's gone
launcher_window_name = Playnite

[BACKGROUND]
# List as many exe's or bat's as you want here. They will run at the start of the GameStream session and be killed at the end.
# background_exe_1 = C:\Program Files (x86)\JoyToKey\JoyToKey.exe
# background_exe_2 = C:\WINDOWS\system32\mspaint.exe

[SETTINGS]
# Set debug = 1 to leave a window running after gamestream to see error messages from GSLP
debug = 0
# Set sleep_on_exit to 1 to put the computer to sleep after the session
sleep_on_exit = 0
# Set close_watch_method to "process" to wait for the launcher process to totally die to exit (can be problematic if it closes to tray), or "window" to just wait the the window to close
close_watch_method = window
"""

# Write the default config
if not os.path.exists(config_filename):
    with open(config_filename, 'w') as out_file:
        out_file.write(default_config)

# Target resolution for gamestream environment
try:
    gamestream_width = sys.argv[1]
    gamestream_height = sys.argv[2]
    # If there's a 3rd argument after the .py/.exe, use it as a custom launcher path
    if len(sys.argv) == 4:
        config_filename = sys.argv[3]
except IndexError:
    print("Error parsing arguments. Did you mean to run one of the .bat launcher scripts?")
    print("Usage: gamestream_launchpad.exe 1920 1080 [config.ini]")
    input("Press Enter to exit.")
    sys.exit(1)

# Parse the config file and assume defaults otherwise
config = configparser.ConfigParser()
config.read(config_filename)
cfg_launcher_path = config['LAUNCHER'].get('launcher_path', r'%LOCALAPPDATA%\Playnite\Playnite.FullscreenApp.exe')
cfg_launcher_window_name = config['LAUNCHER'].get('launcher_window_name', 'Playnite')
cfg_bg_paths = config['BACKGROUND']
debug = config['SETTINGS'].get('debug', '0')
sleep_on_exit = config['SETTINGS'].get('sleep_on_exit', '0')
close_watch_method = config['SETTINGS'].get('close_watch_method', 'window')

launcher_exec_name = os.path.basename(cfg_launcher_path)

# Set resolution to target
set_resolution(gamestream_width, gamestream_height)

# Start background programs, if they're available
for path in cfg_bg_paths:
    expanded_path = os.path.expandvars(cfg_bg_paths[path])
    if os.path.exists(expanded_path):
        print("Launching", expanded_path)
        exec_name = os.path.basename(expanded_path)
        # Kill the process first if it's already running
        if exec_name in (get_process_name(p) for p in psutil.process_iter()):
            os.system('taskkill /f /im ' + exec_name)
        # Start the process
        subprocess.Popen(expanded_path)

# A launcher value of false will create a wait inside of the console instead watching a program
if cfg_launcher_path.lower() == "false":
    input('Press enter to end the GameStream session.')
else:
    # Minimize all windows
    print("Minimizing windows")
    pyautogui.hotkey('winleft', 'd')

    # Playnite has to be killed before it will start in fullscreen mode
    if "Playnite" in launcher_exec_name:
        if "Playnite.FullscreenApp.exe" in (get_process_name(p) for p in psutil.process_iter()):
            os.system('taskkill /f /im ' + "Playnite.FullscreenApp.exe")
        if "Playnite.DesktopApp.exe" in (get_process_name(p) for p in psutil.process_iter()):
            os.system('taskkill /f /im ' + "Playnite.DesktopApp.exe")

        # Move mouse cursor into the lower-right corner to pseudo-hide it because sticks out in playnite fullscreen
        pyautogui.FAILSAFE = False
        pyautogui.moveTo(9999, 9999, duration = 0)

    # Start game launcher
    print("Starting game launcher")
    launcher_exe = os.path.expandvars(cfg_launcher_path)
    subprocess.Popen(launcher_exe)

    # Focus launcher in the foreground and maximize
    results = []
    top_windows = []
    launcher_focused = False
    while not launcher_focused:
        win32gui.EnumWindows(windowEnumerationHandler, top_windows)
        for i in top_windows:
            if cfg_launcher_window_name in i[1]:
                if not 'Fullscreen' in cfg_launcher_path:
                    print("Maximizing", cfg_launcher_window_name)
                    win32gui.ShowWindow(i[0], 3)
                print("Focusing", cfg_launcher_window_name)
                win32gui.SetForegroundWindow(i[0])
                launcher_focused = True
                launcher_window_handle = i[0]
                break
        sleep(1)

    # Watch for closing the launcher window to return to the system's original configuration
    if close_watch_method == "window":
        print("Watching for launcher window to close")
        while win32gui.IsWindowVisible(launcher_window_handle):
            #print("Visible:", launcher_window_handle)
            sleep(2)
            reset_launcher_resolution(gamestream_width, gamestream_height, cfg_launcher_window_name)
    # Alternative method that waits for the process to die (can be problematic if it minimizes to system tray)
    elif close_watch_method == "process":
        print("Watching for launcher process to die")
        while True:
            if launcher_exec_name in (get_process_name(p) for p in psutil.process_iter()):
                sleep(2)
                reset_launcher_resolution(gamestream_width, gamestream_height, cfg_launcher_window_name)
            else:
                break
    else:
        print("No valid close_watch_method in the config. Press Enter when you're done.")
        input()

# Terminate background programs, if they're available
for path in cfg_bg_paths:
    expanded_path = os.path.expandvars(cfg_bg_paths[path])
    if os.path.exists(expanded_path):
        exec_name = os.path.basename(expanded_path)
        print("Terminating", exec_name)
        if exec_name in (get_process_name(p) for p in psutil.process_iter()):
            os.system('taskkill /f /im ' + exec_name)

# Kill gamestream
print("Terminating GameStream session.")
if "nvstreamer.exe" in (get_process_name(p) for p in psutil.process_iter()):
    os.system('taskkill /f /im nvstreamer.exe')

# Restore original resolution
print('Restoring original resolution.')
win32api.ChangeDisplaySettings(None, 0)

if sleep_on_exit == '1':
# Put computer to sleep
    print("Going to sleep")
    os.system("rundll32.exe powrprof.dll,SetSuspendState 0,1,0")

if debug == '1':
    # Leave window open for debugging
    input("Paused for debug review. Press Enter key to close.")


"""
    花园里有 N 个花盆，每个花盆里都有一朵花。这 N 朵花会在 N 天内依次开放，每天有且仅有一朵花会开放并且会一直盛开下去。
    给定一个数组 flowers 包含从 1 到 N 的数字，每个数字表示在那一天开放的花所在的花盆编号。

    例如， flowers[i] = x 表示在第 i 天盛开的花在第 x 个花盆中，i 和 x 都在 1 到 N 的范围内。
    给你一个整数 k，请你输出在哪一天恰好有两朵盛开的花，他们中间间隔了 k 朵花并且都没有开放。
    如果不存在，输出 -1。

    样例 1:
        输入:
            flowers: [1,3,2]
            k: 1
        输出: 2
        解释: 在第二天，第一朵和第三朵花都盛开了。

    样例 2:
        输入:
            flowers: [1,2,3]
            k: 1
        输出: -1
"""
from typing import List

class Solution:
    def kEmptySlots(self, bulbs: List[int], K: int) -> int:
        pass

    @classmethod
    def solve_1(cls, bulbs: List[int], K: int) -> int:
        # 第i盆开花的时间
        days = [0] * (len(bulbs) + 1)
        for index, row in enumerate(bulbs):
            days[row] = index + 1

        result = float("inf")
        left = 1
        right = left + K + 1

        while right < len(days):

            for i in range(left + 1, right):
                if days[i] <= days[left] or days[i] <= days[right]:
                    left = i
                    right = left + K + 1
                    break
            else:
                cur_result = max(days[left], days[right])
                result = min(result, cur_result)
                left = right
                right = left + K + 1
        return result if result != float("-inf") else 1


def reset_weather():
	# Remove weather multiplier if possible - done after accuracy displaying if restarting
	try:
		global weather_mult
		del(weather_mult)
	except:
		pass

def get_weather_advanced():
        global weather_mult
        while True:
                # Manually input weather accuracy multiplier. Make it 100% if entered value is invalid.
                try:
                        weather_mult = float(input("Weather visibility percentage (blank for 100%): ")) / 100
                except:
                        weather_mult = 1
                        break
                # Break the loop if the weather multipler is greater than 0 and less than or equal to 1.
                if (weather_mult > 0 and weather_mult <= 1):
                        break

def get_weather_normal():
	# Dictionary of weather keys and multiplier values.
	WEATHER_DICT = {
	"": 1,
	"clear": 1,
	"drythunderstorm": 1,
	"rain": 0.8,
	"snow": 0.8,
	"rainythunderstorm": 0.8,
	"fog": 0.5,
	"foggyrain": 0.5
	}
	weather_list = tuple(WEATHER_DICT.keys())
	while True:
		# Prompt weather key entry
		weather = input("Weather ('list' for valid inputs, or blank for clear): ")
		if (weather in weather_list):
			break
		if weather == "list":
			count_max = len(weather_list)
			for x in range(1, count_max):
				print("    -",weather_list[x])

	# Weather multiplier matches up with multipler to the suitable key.
	global weather_mult
	weather_mult = WEATHER_DICT[weather]

def get_weat_mode(m):
        while True:
                if (m == "normal" or m == "n"):
                        get_weather_normal()
                        break
                elif (m == "advanced" or m == "a"):
                        get_weather_advanced()
                        break
                else:
                        m = str(input("Input mode for weather ('modes' for modes, blank to skip): "))
                        if (m == ""):
                                break
                        elif (m == "modes"):
                                print("    - normal or n\n    - advanced or a")


from languageModel import LanguageModel
from unigram import Unigram
from bigram import Bigram
from trigram import Trigram
import random

class Interpolation(LanguageModel):
    def __init__(self):
        self.unigram_model = Unigram()
        self.bigram_model = Bigram()
        self.trigram_model = Trigram()
        self.unigram_lambda = .25
        self.bigram_lambda = .25
        self.trigram_lambda = .5

    def train(self, trainingSentences):
        self.unigram_model.train(trainingSentences)
        self.bigram_model.train(trainingSentences)
        self.trigram_model.train(trainingSentences)

    #Arbitrary lambdas.
    def getWordProbability(self, sentence, index):
        return (self.trigram_lambda * self.trigram_model.getWordProbability(sentence, index)) \
               + (self.bigram_lambda * self.bigram_model.getWordProbability(sentence, index)) \
               + (self.unigram_lambda * self.unigram_model.getWordProbability(sentence, index))

    #Doesn't matter which model we use here- vocabulary is the same
    def getVocabulary(self, context):
        return self.trigram_model.getVocabulary(context)

    #What does generating a sentence in an interpolation model look like?
    #I don't know, so what I've done is generate a word using trigram, bigram, and
    #unigram model some of the time, using the same values in getWordProbability
    def generateSentence(self):
        sentence = []
        prev_previous = LanguageModel.START
        previous = random.choice(list(self.trigram_model.word_count.keys()))
        for i in range(20):
            model_choice = random.random()
            if model_choice <= self.trigram_lambda:
                word = self.trigram_model.generateWord(prev_previous, previous)
            elif model_choice > self.trigram_lambda and model_choice <= self.trigram_lambda + self.bigram_lambda:
                word = self.bigram_model.generate_word(previous)
            else:
                word = self.unigram_model.generateWord()
            sentence.append(word)
            prev_previous = previous
            previous = word
            if word == LanguageModel.STOP:
                break
        return sentence


from modules.logging import logging_decorator
from telegram.ext import CommandHandler
import requests

def module_init(gd):
    global subscription_key
    commands = gd.config["commands"]
    subscription_key = gd.config['subscription_key']
    for command in commands:
        gd.dp.add_handler(CommandHandler(command, web_search, pass_args=True))

@logging_decorator('g')
def web_search(bot, update, args):
    query = ' '.join(args)
    if not query:
        update.message.reply_text("You need a query to search!")
        return
    try:
        result_link = get_first_result(query)
    except:
        update.message.reply_text("Sorry, случилась какая-то жопа!")
        return
    if result_link is None:
        update.message.reply_text("Nothing found!")
        return
    msg_text = "[link]({})".format(result_link)
    update.message.reply_text(msg_text, parse_mode="Markdown")

def get_first_result(query):
    search_url = "https://api.cognitive.microsoft.com/bing/v7.0/search"
    headers = {"Ocp-Apim-Subscription-Key": subscription_key}
    params = {
        "q": query,
        "textDecorations": False,
        "textFormat": "Raw",
    }
    response = requests.get(search_url, headers=headers, params=params)
    response.raise_for_status()
    search_results = response.json()
    if 'webPages' not in search_results:
        return None
    values = search_results['webPages']['value']
    link = values[0]['url']
    return link


from datetime import datetime, date, timedelta
from bs4 import BeautifulSoup
from requests import Session

class KrepSession(Session):
    def __init__(self, sso_id, sso_pw, uuid):
        Session.__init__(self)
        self.get(f'https://krep.kaist.ac.kr/krep/ctf/mng/cdc/CtfMngCdc001M51.do?mobile_sn_no={uuid}')
        self.post('https://krep.kaist.ac.kr/krep/m_login.do', data={'b001': sso_id, 'b002': sso_pw})

    def get_qrcode(self):
        r = self.get('https://krep.kaist.ac.kr/krep/ctf/mng/cdc/CtfMngCdc001M11.do')
        return r.content

    def get_status(self):
        r = self.get('https://krep.kaist.ac.kr/krep/ctf/mng/cdc/CtfMngCdc001M53.do')
        s = BeautifulSoup(r.content)
        t = s.get_text().split()

        curr = t[12]
        next = t[14]
        defer = datetime.strptime(f'{t[16]} {t[17]}', '%Y-%m-%d %H:%M:%S')
        due = defer + timedelta(minutes=150) if next == '출근' \
            else datetime.strptime(f'{t[19]} {t[20]}', '%Y-%m-%d %H:%M:%S')

        return KrepStatus(curr, next, defer, due)

class KrepStatus:
    def __init__(self, curr, next, defer, due):
        self.curr = curr
        self.next = next
        self.defer = defer
        self.due = due

    def __str__(self):
        defer = self.defer.strftime('%H:%M') if self.defer.date() == date.today() else self.defer.strftime('%Y-%m-%d')
        return f'{self.curr} / {defer}부터 {self.next} 가능'

    @property
    def alarm_text(self):
        return self.next + '도장 찍기'


import cv2
import numpy as np
import matplotlib.pyplot as plt

def HW4(bonus=False):
    """

    :return:
    """
    MIN_MATCH_COUNT = 10
    # part 1
    # 1.Load the "stones" image and convert it to grayscale.
    img1 = cv2.imread('stones.png')
    gray = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    # 2.Detect SIFT features and compute descriptors for them using sift.detectAndCompute().
    sift = cv2.xfeatures2d.SIFT_create()
    kp1, des_stone = sift.detectAndCompute(gray, None)
    img1 = cv2.drawKeypoints(img1,kp1,img1)
    # 3.Show the SIFT features on the image.
    # cv2.imshow('sift', draw_keypts)
    # 4.Iterate through the frames of the video and detect SIFT features & descriptors on each frame
    cap = cv2.VideoCapture('input.mov')
    # Part 2, 1.Create a "brute force" matcher (BFMatcher) object.
    bf = cv2.BFMatcher()
    i = 0
    inlier_ = outlier_ = 0
    frame_size = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    filename = 'overlay_Chan.avi' if bonus else 'boarder_Chan.avi'
    video = cv2.VideoWriter(filename=filename, fourcc=cv2.VideoWriter_fourcc('M', 'P', 'E', 'G'), fps=30,
                            frameSize=frame_size)

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        sift = cv2.xfeatures2d.SIFT_create()
        kp2, des_frame = sift.detectAndCompute(gray, None)
        # frames.append(cv2.drawKeypoints(frame, kp2, frame))
        # part 2
        # 2.At each video frame, compute the first two nearest-neighbor matches for each keypoint in the video frame
        # (matching to the target "stones" image).
        matches = bf.knnMatch(des_stone, des_frame, k=2)
        # 3.Apply the "ratio test" to filter out unreliable matches, using a threshold of 0.7
        # store all the good matches as per Lowe's ratio test.
        good = []
        for m, n in matches:
            if m.distance < 0.7 * n.distance:
                good.append(m)
        # 4. Draw the matches between the video frame and the target image using cv.drawMatches()
        if len(good) > MIN_MATCH_COUNT:
            src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
            dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
            # part 3
            # 1.Fit a homography to the matches using findHomography() with the RANSAC option.
            # Note that you might need to invert the resulting homography matrix using np.linalg.inv().
            H, inlier = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
            matchesMask = inlier.ravel().tolist()
            print(matchesMask)
            h, w, d = img1.shape
            pts = np.float32([[0, 0], [0, h - 1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2)
            dst = cv2.perspectiveTransform(pts, H)
            # Draw polygon around the tracking target in the video frame using cv.polylines()
            frame = cv2.polylines(frame, [np.int32(dst)], True, 255, 3, cv2.LINE_AA)
        else:
            print("Not enough matches are found - {}/{}".format(len(good), MIN_MATCH_COUNT))
            matchesMask = None
        # 4.Draw the inlier matches between the video frame and the target image using cv.drawMatches()

        draw_params = dict(matchColor=(0, 255, 0),  # draw matches in green color
                           singlePointColor=None,
                           matchesMask=matchesMask,  # draw only inliers
                           flags=2)
        inlier_matches = cv2.drawMatches(img1, kp1, frame, kp2, good, None, **draw_params)
        # cv2.imwrite('inlier_matches'+str(i)+'.png', inlier_matches)
        if not i:
            inlier_h, inlier_w, inlier_d = inlier_matches.shape
            video2 = cv2.VideoWriter(filename='matches.avi', fourcc=cv2.VideoWriter_fourcc('M', 'P', 'E', 'G'),
                                     fps=15,
                                     frameSize=(inlier_w, inlier_h))
        video2.write(inlier_matches)
        if bonus:
            # bonus
            # Use the calculate homography to warp the overlay image (see cv.perspectiveWarp() ).
            overlay = cv2.imread('overlay.png')
            # To put the overlay on the image, you need to create a mask which has ones over the warped overlay and zeros
            # outside of it.  To do this, create an image of all ones and
            # warp it using perspectiveWarp().
            mask = np.ones(overlay.shape)
            mask = cv2.warpPerspective(mask, H, frame_size)
            overlay = cv2.warpPerspective(overlay, H, frame_size)
            # Calculate the composite image as (1-mask) * frame + mask * overlay.
            comp_img = np.multiply(1-mask, frame) + np.multiply(mask, overlay)
            frame = np.uint8(comp_img)
        video.write(frame)
        i += 1
    # part 3  5. Calculate the average percentage of inliers.
    for i in range(len(good)):
        if matchesMask[i]:
            inlier_ += 1
        else:
            outlier_ += 1
    percent_inlier = inlier_ / (inlier_ + outlier_) * 100
    percent_outlier = 100-percent_inlier
    print("percent inlier = %.6f" % percent_inlier)
    print("percent outlier = %.6f" % percent_outlier)
    video.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    # True to generate output with overlay
    # False to generate output with border
    HW4(True)


# Definition for a undirected graph node
class UndirectedGraphNode:
    def __init__(self, x):
        self.label = x
        self.neighbors = []

class Solution:
    # @param node, a undirected graph node
    # @return a undirected graph node
    def cloneGraph(self, node):
        if node is None:
            return

        dictionary = {node.label: UndirectedGraphNode(node.label)}
        left = set()
        left.add(node)

        while left:
            cur = next(iter(left))
            for nb in cur.neighbors:
                if nb.label not in list(dictionary.keys()):
                    dictionary[nb.label] = UndirectedGraphNode(nb.label)
                    left.add(nb)
                dictionary[cur.label].neighbors.append(dictionary[nb.label])
            left.remove(cur)
        return dictionary[node.label]

    # Time Limit Exceeded
    def cloneGraph2(self, node):
        if node is None:
            return
        m = {}
        return self.cloneNode2(node, m)

    def cloneNode2(self, node, m):
        res = UndirectedGraphNode(node.label)
        m[node.label] = res
        for nb in node.neighbors:
            if nb.label in list(m.keys()):
                res.neighbors.append(m[nb.label])
            else:
                res.neighbors.append(self.cloneNode2(nb, m))
        return res

n0 = UndirectedGraphNode(0)
n1 = UndirectedGraphNode(1)
n2 = UndirectedGraphNode(2)
n3 = UndirectedGraphNode(3)
n4 = UndirectedGraphNode(4)
n5 = UndirectedGraphNode(5)

n0.neighbors.append(n1)
n0.neighbors.append(n5)

n1.neighbors.append(n2)
n1.neighbors.append(n5)

n2.neighbors.append(n3)

n3.neighbors.append(n4)
n3.neighbors.append(n4)

n4.neighbors.append(n5)
n4.neighbors.append(n5)

res = Solution().cloneGraph(n0)


import numpy as np
import pickle
class Model():

    def __init__(self, feature_count, n, gamma, alpha, parameter_vectors):
        '''
        feature_count: number of features

        n: number of steps in temporal difference Q-learning
        gamma: discounting factor
        alpha: learning rate

        parameter_vectors: start vectors for training
        '''

        self.feature_count = feature_count
        self.parameter_vectors = parameter_vectors

        self.n = n
        self.gamma = gamma
        self.alpha = alpha

        self.discounts = (np.ones((n+1))*gamma)**np.linspace(0,n,n+1) #discounting factors

        self.batches = [[] for _ in range(6)]  #six batches for the six moves
        self.rewards = [[]]
        self.states = [[]]

        self.TAU = 0
        self.T = 400 #maximum 400 time steps per episode

    def add_step(self, tau, t, state, action, reward):
        '''
        tau: episode
        t: time step in episode
        action: action that is executed
        reward: reward that is earned after the action
        '''
        if tau > self.TAU: #new episode
            self.TAU = tau
            self.rewards.append([])
            self.states.append([])

        self.batches[action].append((tau, t))
        self.rewards[tau].append(reward)
        self.states[tau].append(state)

    def update_parameter_vectors(self):
        #print(self.parameter_vectors)
        #calculate The Ys
        Y = np.empty((self.TAU+1, self.T))
        Y.fill(np.nan)
        for tau in range(self.TAU+1):
            steps = len(self.states[tau])
            for t in range(steps-self.n):
                rewards = np.array(self.rewards[tau][t:t+self.n]) #rewards for next n steps
                Y[tau][t] = np.nansum(rewards*self.discounts[0:-1])
                Y[tau][t] += np.max(np.dot(self.states[tau][t+self.n],self.parameter_vectors.T))*self.discounts[-1]
            for t in range(steps-self.n,steps):
                rewards = np.array(self.rewards[tau][t:]) #rewards till end of the game
                discounts = np.array(self.discounts[0:len(self.rewards[tau][t:])]) #discounts for these rewards
                Y[tau][t] = np.nansum(rewards*discounts)

        #update the parameter vectors
        for i,beta in enumerate(self.parameter_vectors):
            batch = self.batches[i]
            sum_ = np.zeros(len(beta))
            #print(beta)
            for tau, t in batch:
                if np.isnan(Y[tau][t]):
                    print(tau, t)
                sum_ += self.states[tau][t] * (Y[tau][t] - np.dot(self.states[tau][t], beta))
            #print(sum_)
            self.parameter_vectors[i] += (self.alpha / len(batch)) * sum_

        #reset for next training cicle
        self.batches = [[] for _ in range(6)]  #six batches for the six moves
        self.rewards = [[]]
        self.states = [[]]

        self.TAU = 0

        #print("***\nUpdated:")
        #print(self.parameter_vectors)

        # Store the model
        with open("my-saved-model.pt", "wb") as file:
            pickle.dump(self.parameter_vectors, file)

    def predict_action(self, state):
        #return softmax of the actions for given state
        Q = np.dot(state, self.parameter_vectors.T)
        #print("***")
        #print(Q)
        #print("softmax")
        #print(np.exp(Q)/sum(np.exp(Q)))
        Q /= np.max(np.abs(Q))
        return Q, np.exp(Q)/sum(np.exp(Q)), np.argmax(Q) #Q, softmax, max action index


#!/usr/bin/env python

import pytest

import netmiko
from DEVICE_CREDS import *

def setup_module(module):

    module.EXPECTED_RESPONSES = {
        'enable_prompt' : 'RP/0/0/CPU0:XRv-1#',
        'base_prompt'   : 'RP/0/0/CPU0:XRv-1',
        'interface_ip'  : '169.254.254.181',
        'config_mode'   : '(config)'
    }

    show_ver_command = 'show version'
    commit_history_cmd = 'show configuration commit list'
    module.basic_command = 'show ipv4 int brief'

    SSHClass = netmiko.ssh_dispatcher(cisco_xr['device_type'])
    net_connect = SSHClass(**cisco_xr)

    module.show_version = net_connect.send_command(show_ver_command)
    module.show_ip = net_connect.send_command(module.basic_command)

    net_connect.enable()
    module.enable_prompt = net_connect.find_prompt()

    current_commit_history = net_connect.send_command(commit_history_cmd)

    # get the current 10 digit commit Label ID
    module.current_commit = current_commit_history.split('\n')[4].split()[1]
    module.config_mode = net_connect.config_mode()
    config_commands = ['logging monitor warning', 'logging buffered 30720100', 'logging console errors']
    net_connect.send_config_set(config_commands, commit=True)

    new_commit_history = net_connect.send_command(commit_history_cmd)

    # get the new 10 digit commit Label ID
    module.new_commit = new_commit_history.split('\n')[4].split()[1]

    module.config_commands_output = net_connect.send_command("show run | inc logging")

    net_connect.disconnect()

def test_enable_mode():
    assert enable_prompt == EXPECTED_RESPONSES['enable_prompt']

def test_config_mode():
    assert EXPECTED_RESPONSES['config_mode'] in config_mode

def test_command_set():
    assert 'logging buffered 30720100' in config_commands_output
    assert 'logging console errors' in config_commands_output

def test_commit():
    assert int(current_commit) < int(new_commit)


#!/usr/bin/env python

import subprocess as sp
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import vplot as vpl
import os
import sys
import matplotlib.patches as mpatches

#Variable List
dest = "KDwarfStatic"
num = 100
L_sun = 3.846e26
a_earth = 1
listf = "list_KDwarfStatic"

# if the list file exists, extract data for plotting

lum0, obliq0, semi0, inst, snowball, northCapL, northCapS, southCapL, southCapS, icebeltL, icebeltS, iceFree, tGlobal = np.loadtxt(
    listf, unpack=True)

PolarCaps = np.zeros(num*num)
MoistGreen = np.zeros(num*num)
for i in np.arange(num*num):

    if (
    northCapL[i] == 1 and southCapL[i] == 1 and icebeltS[i] == 0 and icebeltL[i] == 0 or
    northCapS[i] == 1 and northCapL[i] == 1 and southCapS[i] == 1 and southCapL[i] == 1 and icebeltS[i] == 0 and icebeltL[i] == 0 or
    northCapS[i] == 1 and southCapS[i] == 1 and icebeltS[i] == 0 and icebeltL[i] == 0 or
    northCapL[i] == 1 and southCapS[i] == 1 and icebeltS[i] == 0 and icebeltL[i] == 0 or
    northCapS[i] == 1 and southCapL[i] == 1 and icebeltS[i] == 0 and icebeltL[i] == 0
    ):
        PolarCaps[i] = 1

    if tGlobal[i] >= 70 and iceFree[i] == 1:
        MoistGreen[i] = 1

lum0 = np.reshape(lum0, (num, num))
obliq0 = np.reshape(obliq0, (num, num)) * 180 / np.pi
semi0 = np.reshape(semi0, (num, num)) / 1.49598e11
inst = np.reshape(inst, (num, num)) / 1350
snowball = np.reshape(snowball, (num, num))
northCapL = np.reshape(northCapL, (num, num))
northCapS = np.reshape(northCapS, (num, num))
southCapL = np.reshape(southCapL, (num, num))
southCapS = np.reshape(southCapS, (num, num))
icebeltL = np.reshape(icebeltL, (num, num))
icebeltS = np.reshape(icebeltS, (num, num))
iceFree = np.reshape(iceFree, (num, num))
tGlobal = np.reshape(tGlobal, (num, num))

PolarCaps = np.reshape(PolarCaps, (num, num))
MoistGreen = np.reshape(MoistGreen, (num, num))

plt.figure(figsize=(9,6.5))
plt.ylabel("Instellation [Earth]", fontsize=16)
plt.xlabel("Obliquity [$^\circ$]", fontsize=16)
plt.ylim(0.867,1.211)
plt.xlim(0,90)

iFF = plt.contourf(obliq0, inst, iceFree, [0, 1], colors = vpl.colors.dark_blue)
sNF = plt.contourf(obliq0, inst, snowball, [0.5, 1], colors = '#efefef')
PcF = plt.contourf(obliq0,inst,PolarCaps, [0.5, 1], colors = vpl.colors.purple)
icF = plt.contourf(obliq0,inst,icebeltL, [0.5, 1], colors = vpl.colors.pale_blue)

h1, _ = iFF.legend_elements()
h2, _ = icF.legend_elements()
h3, _ = sNF.legend_elements()
h4, _ = PcF.legend_elements()

plt.legend([h1[0], h2[0], h3[0], h4[0]], [ 'Ice Free', 'Ice Belt', 'Snowball','Polar Ice Caps'],loc = 'upper left', bbox_to_anchor=(0, 1.02, 1, 0.102),ncol=4, mode="expand", borderaxespad=0)

if (sys.argv[1] == 'pdf'):
    plt.savefig(dest + '.pdf', dpi=300)
if (sys.argv[1] == 'png'):
    plt.savefig(dest + '.png', dpi=300)

plt.show()
plt.close()