from typing import Tuple, List, Union, Iterable, Callable, Optional, Dict

import re
import yaml
import json
import numpy as np
import tensorflow as tf
import tensorflow_gnn as tfgnn
from dataclasses import dataclass
from pathlib import Path

from graph2tac.loader.data_classes import DataConfig, GraphConstants, LoaderGraph, ProofstateMetadata, ProofstateContext, LoaderAction, LoaderActionSpec, LoaderProofstate, LoaderProofstateSpec, LoaderDefinition, LoaderDefinitionSpec
from graph2tac.loader.data_server import DataToTFGNN
from graph2tac.tfgnn.tasks import PredictionTask, DefinitionTask, GLOBAL_ARGUMENT_PREDICTION
from graph2tac.tfgnn.train import Trainer
from graph2tac.common import logger
from graph2tac.predict import Predict, predict_api_debugging, NUMPY_NDIM_LIMIT
from graph2tac.tfgnn.stack_graph_tensors import stack_graph_tensors

class BeamSearch(tf.keras.layers.Layer):
    """Beam search layer

    :param token_log_prob_fn: A tensorflow function which returns log probabilities for all the tokens.
        It also takes as input a cache dictionary which is stores hidden information for each partial sequence
        which can be used to compute the log probabilities, and it returns an updated version.  (The beam search
        will use gather to reorder this cache after each beam search step.)  Another dictionary static_data
        is used to pass in other tensorflow tensorflow objects which can be used in the computation.
        Inputs:  - partial sequences: [batch, beam_size, seq_length]
                 - beam_search_step: int32
                 - cache: dict of [batch, beam_size, ...]
                 - static_data: dict of tensors of any shape
        Outputs: - token_log_probs: [batch, beam_size, vocabulary]
                 - updated_cache: dict of [batch, beam_size, ...]
    :param stopping_fn: A tensorflow function which returns whether a beam ray is finished.
        It takes inputs and outputs similar to token_log_prob_fn.
        Inputs:  - partial sequences: [batch, beam_size, seq_length]
                 - beam_search_step: int32
                 - cache: dict of [batch, beam_size, ...]
                 - static_data: dict of tensors of any shape
        Outputs: - is_finished: [batch, beam_size, vocabulary]
                 - updated_cache: dict of [batch, beam_size, ...]
    """
    def __init__(
        self,
        token_log_prob_fn: Callable[
            [tf.Tensor, tf.Tensor, Dict[str, tf.Tensor], Dict[str, Union[tf.Tensor, tf.RaggedTensor]]],
            Tuple[tf.Tensor, Dict[str, tf.Tensor]]
        ],
        stopping_fn: Callable[
            [tf.Tensor, tf.Tensor, Dict[str, tf.Tensor], Dict[str, Union[tf.Tensor, tf.RaggedTensor]]],
            Tuple[tf.Tensor, Dict[str, tf.Tensor]]
        ],
    ):
        super().__init__()
        self.token_log_prob_fn = token_log_prob_fn
        self.stopping_fn = stopping_fn

    def one_beam_step(
        self,
        ids: tf.Tensor,  # [batch_size, beam_size0, seq_length]
        scores: tf.Tensor,  # [batch_size, beam_size0]
        cache: Dict[str, tf.Tensor],  # dict with values: [batch_size, beam_size0, ...]
        static_data: Dict[str, Union[tf.Tensor, tf.RaggedTensor]],
        i: tf.Tensor,  # int32
        max_beam_size: tf.Tensor,  # int32
    ) -> Tuple[
        tf.Tensor,  # ids: [batch_size, beam_size, seq_length+1]
        tf.Tensor,  # scores: [batch_size, beam_size]
        Dict[str, tf.Tensor],  # cache with values: [batch_size, beam_size, ...]
    ]:
        batch_size = tf.shape(ids)[0]
        beam_size0 = tf.shape(ids)[1]

        # [batch, beam_size0, vocabulary], dict of [batch_size, beam_size0, ...]
        token_scores, cache = self.token_log_prob_fn(ids, i, cache, static_data)
        vocabulary = tf.shape(token_scores)[2]

        # don't return full beam_size if not enough vocabulary to support it yet
        beam_size = tf.minimum(max_beam_size, beam_size0 * vocabulary)

        all_scores = tf.expand_dims(scores, axis=-1) + token_scores  # [batch_size, beam_size0, vocabulary]
        all_scores = tf.reshape(all_scores, shape=[batch_size, beam_size0*vocabulary])  # [batch_size, beam_size0*vocabulary]

        top_k = tf.math.top_k(all_scores, k=beam_size, sorted=True)
        top_k_scores = top_k.values  # [batch_size, beam_size]
        top_k_ixs = top_k.indices  # [batch_size, beam_size]
        top_k_beams = top_k_ixs // vocabulary  # [batch_size, beam_size]
        top_k_tokens = top_k_ixs % vocabulary  # [batch_size, beam_size]

        # reorder the beams and add on the new tokens
        ids = tf.gather(ids, top_k_beams, batch_dims=1)  # [batch_size, beam_size, seq_length]
        ids = tf.concat([ids, tf.expand_dims(top_k_tokens, axis=-1)], axis=-1)  # [batch_size, beam_size, seq_length+1]

        # reorder cache
        cache = {
            k: tf.gather(v, top_k_beams, batch_dims=1)  # [batch_size, beam_size, ...]}
            for k, v in cache.items()  # v: [batch_size, beam_size0, ...]
        }

        return ids, top_k_scores, cache

    def one_beam_step_with_early_stopping(
        self,
        is_finished: tf.Tensor, # bool
        ids: tf.Tensor,  # [batch_size, beam_size0, seq_length]
        scores: tf.Tensor,  # [batch_size, beam_size0]
        cache: Dict[str, tf.Tensor],  # dict with values: [batch_size, beam_size0, ...]
        static_data: Dict[str, Union[tf.Tensor, tf.RaggedTensor]],
        i: tf.Tensor,  # int32
        max_beam_size: tf.Tensor,  # int32
    ) -> Tuple[
        tf.Tensor,  # ids: [batch_size, beam_size, seq_length+1]
        tf.Tensor,  # scores: [batch_size, beam_size]
        Dict[str, tf.Tensor],  # cache with values: [batch_size, beam_size, ...]
    ]:
        if is_finished:
            return (is_finished, (ids, scores, cache))

        # check which are finished
        is_finished, cache = self.stopping_fn(ids, i, cache, static_data)  # [batch_size, beam_size0]
        is_finished = tf.reduce_all(is_finished)

        # wait until all are finished
        if is_finished:
            return (is_finished, (ids, scores, cache))

        return (is_finished, self.one_beam_step(ids, scores, cache, static_data, i, max_beam_size))

    def call(
        self,
        initial_ids: tf.Tensor,  # [batch_size, initial_seq_length]
        initial_cache: Dict[str, tf.Tensor],  # [batch_size, ...]
        static_data: Dict[str, Union[tf.Tensor, tf.RaggedTensor]],
        beam_size: tf.Tensor,  # int32
        max_decode_length: tf.Tensor,  # int32
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        batch_size = tf.shape(initial_ids)[0]

        # start with a beam_size of 1 for the input to the first step first step
        ids = tf.expand_dims(initial_ids, axis=1)  # [batch_size, beam_size0, initial_seq_length]
        cache = {
            k: tf.expand_dims(v, axis=1)  # [batch_size, beam_size0, ...]
            for k, v in initial_cache.items()  # v: [batch_size, ...]
        }
        scores = tf.expand_dims(tf.zeros(shape=[batch_size]), axis=-1)  # [batch_size, beam_size0]

        # iterate over the sequence
        is_finished = False
        # this loop is automatically converted into tf.while_loop
        shape_invariants = [
            (ids, tf.TensorShape([None, None, None])),
            (scores, tf.TensorShape([None, None])),
            (cache, {k: tf.TensorShape([None, None] + [None for i in range(len(tf.shape(v)) - 2)]) for k,v in cache.items()}),
        ]
        for i in range(max_decode_length):
            # the sizes of beam_size and seq_length can change during the loop and we need to let TF know this
            tf.autograph.experimental.set_loop_options(
                shape_invariants=shape_invariants
            )
            # ids: [batch_size, beam_size, initial_seq_length]
            # scores: [batch_size, beam_size]
            # cache: dict with values [batch_size, beam_size, ...]
            is_finished, (ids, scores, cache) = self.one_beam_step_with_early_stopping(is_finished, ids, scores, cache, static_data, i, beam_size)

        return ids, scores

class SelectBestResults(tf.keras.layers.Layer):
    """Select the best results from model.

    Returns best results given the model outputs.
    It takes as input the inference_output of the base inference model.
    It returns two things:
    - log_probabilities: the log probabilities of the returned results, shape [results]
    - results as a ragged tensor of shape [results, (sequence_length), 2]
      the sequence includes both the tactic id as the first element and the arguments
      the tactic is encoded as [tactic_id, tactic_id]
      the arguments are encoded as either [0, local_context_id] or [1, global_context_id]

    It may return less than search_expand_bound number of results,
    since it filters out results with probability 0.

    :param tactic_index_to_numargs: list tactic arg lengths for each tactic
    :param search_expand_bound: maximum number of results to return
    """
    def __init__(self, tactic_index_to_numargs: List[int], search_expand_bound: int):
        super().__init__()
        self.tactic_index_to_numargs = tf.cast(tf.constant(tactic_index_to_numargs), tf.int32)
        self.search_expand_bound = search_expand_bound

        self.beam_search = BeamSearch(
            token_log_prob_fn=self.token_log_prob_fn,
            stopping_fn=self.stopping_fn,
        )

    def token_log_prob_fn(
        self,
        partial_seqs: tf.Tensor,  # [batch, beam_size, seq_length] dtype: int
        pos: tf.Tensor,  # scalar dtype: int32,
        cache: Dict[str, tf.Tensor],  # value shapes: [batch, beam_size, ...]
        static_data: Dict[str, Union[tf.Tensor, tf.RaggedTensor]],
    ) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]:  # output shapes [batch, beam_size, vocab], dict of [batch, beam_size, ...]
        if pos == 0:
            # look up tactic logits
            batch_ix = cache["batch_ix"]  # [batch, beam_size]
            tactic_logits = static_data["tactic_logits"]  # [batch, vocab]
            tactic_logits = tf.gather(tactic_logits, batch_ix)  # [batch, beam_size, vocab]
            return tactic_logits, cache

        # look up argument logits
        arg_pos = tf.cast(pos - 1, tf.int32)
        arg_logits = static_data["arg_logits"]  # [batch, tactic, arg, vocab]
        batch_ix = cache["batch_ix"]  # [batch, beam_size]
        # first token is the tactic id.  It may be a stop token but that is accounted for
        tactic_id = self.get_tactic(partial_seqs)  # [batch, beam_size]
        indices = tf.stack([batch_ix, tactic_id, tf.ones_like(batch_ix) * arg_pos], axis=-1)  # [batch, beam_size]
        logits = tf.gather_nd(arg_logits, indices)
        return logits, cache

    @staticmethod
    def get_tactic(partial_seqs):
        if tf.size(partial_seqs) == 0:
            return tf.zeros(tf.shape(partial_seqs)[:2], tf.int32)
        else:
            return tf.reduce_max(partial_seqs[:, :, :1], axis=-1)

    def stopping_fn(
        self,
        partial_seqs: tf.Tensor,  # [batch, beam_size, seq_length] dtype: int
        pos: tf.Tensor,  # scalar dtype: int32,
        cache: Dict[str, tf.Tensor],  # value shapes: [batch, beam_size, ...]
        static_data: Dict[str, Union[tf.Tensor, tf.RaggedTensor]],
    ) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]:  # output shapes [batch, beam_size], dict of [batch, beam_size, ...]
        if pos == 0:
            is_finished = tf.zeros(shape=tf.shape(partial_seqs)[:2], dtype=bool)  # [batch, beam_size]
            return is_finished, cache
        else:
            arg_pos = tf.cast(pos - 1, tf.int32)
            # first token is the tactic id.  It may be a stop token but that is accounted for
            tactic_id = self.get_tactic(partial_seqs)  # [batch, beam_size]
            batch_ix = cache["batch_ix"]  # [batch, beam_size]
            arg_lengths = static_data["arg_lengths"]  # [batch, tactic, arg, vocab]
            indices = tf.stack([batch_ix, tactic_id], axis=-1)  # [batch, beam_size, 2]
            arg_lengths = tf.gather_nd(arg_lengths, indices)  # [batch, beam_size]

            is_finished = (arg_pos >= arg_lengths)
            return is_finished, cache

    def _select_best_results_search(
        self,
        tactic_logits: tf.Tensor,  # [batch, tactics]
        tactic_arg_counts: tf.Tensor,  # [batch, tactics] dtype:int32
        arg_logits: tf.RaggedTensor,  # [batch * tactics, None(args), small_cxt]
        beam_width: tf.Tensor,  # int32
    ) -> Tuple[
        tf.Tensor,  # tactic_token [batch-beam]
        tf.RaggedTensor,  # arg_token [batch-beam, None(args)]
        tf.Tensor,  # log_probs [batch-beam]
        tf.Tensor,  # batch_ix [batch-beam]
    ]:
        batch_size = tf.shape(tactic_logits)[0]

        # encode: compute sequences where
        # - the tactic id (or padding) is in the first position
        # - arguments (or padding) are in later positions
        # [batch, beam_size, sequence_length], [batch, beam_size]
        sequences, log_probs = self.beam_search(
            initial_ids=tf.zeros([batch_size, 0], tf.int32),  # [batch, 0]
            initial_cache={"batch_ix": tf.range(batch_size)},  # [batch],
            static_data={
                "arg_logits": arg_logits, # [batch, 1+tactic, max(args), vocab]
                "arg_lengths": tactic_arg_counts,  # [batch, 1+tactic]
                "tactic_logits": tactic_logits,  # [batch, 1+tactics]
            },
            beam_size=beam_width,
            max_decode_length=tf.reduce_max(tactic_arg_counts) + 1  # +1 for the tactic
        )

        # decode

        # make ragged to remove bad sequences
        # [batch, None(beam_size), sequence_length]
        sequences = tf.RaggedTensor.from_tensor(sequences, ragged_rank=1)
        log_probs = tf.RaggedTensor.from_tensor(log_probs, ragged_rank=1)  # [batch, None(beam)]

        # split tactics and args and record batch index
        tactic_token = tf.cast(sequences.values[:, 0], tf.int64)  # [batch-beam]
        arg_tokens = sequences.values[:, 1:]  # [batch-beam, max(args)]
        batch_ix = sequences.value_rowids()  # [batch-beam]

        # filter sequences for positive probability and a valid tactics
        is_good = (log_probs.values != -np.inf)  # [batch-beam]
        batch_ix = tf.boolean_mask(batch_ix, is_good)  # [batch-beam]
        log_probs = tf.boolean_mask(log_probs.values, is_good)  # [batch-beam]
        tactic_token = tf.boolean_mask(tactic_token, is_good)  # [batch-beam]
        arg_tokens = tf.boolean_mask(arg_tokens, is_good)  # [batch-beam, max(args)]

        # filter arguments for trailing padding.
        batch_tactic_token_nd_ix = tf.stack([batch_ix, tactic_token], axis=-1)  # [batch-beam, 2]
        arg_lengths = tf.gather_nd(tactic_arg_counts, batch_tactic_token_nd_ix)  # [batch-beam]
        arg_tokens = tf.RaggedTensor.from_tensor(arg_tokens, lengths=arg_lengths, ragged_rank=1)  # [batch-beam, None(args)]

        return tactic_token, arg_tokens, log_probs, batch_ix

    def _select_best_results_padding(
        self,
        tactic_logits: tf.Tensor,  # [batch, tactics]
        tactic_arg_counts: tf.Tensor,  # [batch, tactics] dtype:int32
        arg_logits: tf.RaggedTensor,  # [batch * tactics, None(args), small_cxt]
        beam_width: tf.Tensor,  # int32
    ) -> Tuple[
        tf.Tensor,  # tactic [batch-beam]
        tf.RaggedTensor,  # arg_ix [batch, None(args)]
        tf.Tensor,  # log_probs [batch-beam]
        tf.Tensor,  # batch_ix [batch-beam]
    ]:
        batch_size = tf.shape(tactic_logits)[0]
        tactic_cnt = tf.shape(tactic_logits)[1]

        PAD = 0

        # encode: add padding for args and tactics
        # for args padding:
        # if valid argument postion, the probability is 0 (log prob = -inf)
        # [batch*tactic, None(args), 1+cxt]
        arg_logits = tf.ragged.map_flat_values(tf.pad, arg_logits, paddings=[[0, 0], [1, 0]], constant_values=-np.inf)
        # if invalid argument postion, the probability is 1 (log prob = 0)
        log_prob_pad_token = tf.math.log(tf.one_hot(PAD, tf.shape(arg_logits)[-1]))  # [1+cxt]
        # [batch*tactic, max(args), 1+cxt]
        arg_logits = arg_logits.to_tensor(default_value=log_prob_pad_token)
        # pad tactic dimension as well.  the probability is always 0 (log prob = -inf) for the tactic pad token
        # [batch, tactic, max(args), 1+cxt]
        arg_logits = tf.reshape(arg_logits, shape=[batch_size, tactic_cnt, tf.shape(arg_logits)[1], tf.shape(arg_logits)[2]])
        # [batch, 1+tactic, max(args), 1+cxt]
        arg_logits = tf.pad(arg_logits, paddings=[[0,0], [1, 0], [0,0], [0,0]], constant_values=-np.inf)
        # [batch, 1+tactic]
        tactic_logits = tf.pad(tactic_logits, paddings=[[0,0], [1, 0]], constant_values=-np.inf)
        # [batch, 1+tactic]
        tactic_arg_counts = tf.pad(tactic_arg_counts, paddings=[[0,0], [1, 0]], constant_values = 0)

        # calculate results
        # [batch-beam], [batch-beam], [batch-beam], [batch-beam, None(args)]
        tactic_token, arg_tokens, log_probs, batch_ix = self._select_best_results_search(
            tactic_logits=tactic_logits,
            tactic_arg_counts=tactic_arg_counts,
            arg_logits=arg_logits,
            beam_width=beam_width
        )

        # decode
        # note, we have already removed all padding tokens
        tf.assert_greater(arg_tokens.values, tf.constant(PAD, tf.int32))
        tf.assert_greater(tactic_token, tf.constant(PAD, tf.int64))
        # so we can shift the index
        tactic_ix = tactic_token - 1  # [batch-beam]
        arg_ix = arg_tokens - 1  # [batch-beam, None(args)]

        return tactic_ix, arg_ix, log_probs, batch_ix

    def _select_best_results_top_k(
        self,
        tactic_logits: tf.Tensor,  # [batch, tactics]
        tactics: tf.Tensor,  # [batch, tactics] dtype:int32
        tactic_arg_counts: tf.Tensor,  # [batch, tactics] dtype:int32
        arg_logits: tf.RaggedTensor,  # [batch * tactics, None(args), cxt]
        beam_width: tf.Tensor,  # int32
    ) -> Tuple[
        tf.Tensor,  # tactic [batch-beam]
        tf.RaggedTensor,  # arg_ix [batch, None(args)]
        tf.Tensor,  # log_probs [batch-beam]
        tf.Tensor,  # batch_ix [batch-beam]
    ]:
        tactic_cnt = tf.shape(tactics)[1]

        # tactics are already restricted to top-k

        # arg encoding step: restrict to the top-k arguments
        # [batch-tactic-args, small_cxt], [batch-tactic-args, small_cxt]
        arg_logits_values, pre_topk_arg_indices = \
            tf.math.top_k(arg_logits.values, k=tf.minimum(beam_width, tf.shape(arg_logits)[-1]), sorted=False)
        pre_topk_arg_indices = arg_logits.with_values(pre_topk_arg_indices)  # [batch*tactic, None(args), small_cxt]
        arg_logits = arg_logits.with_values(arg_logits_values)  # [batch*tactic, None(args), small_cxt]

        # calculate results
        # tactic_ix and arg_ix are indices into tactic and pre_topk_arg_indices
        # [batch-beam], [batch-beam, None(args)], [batch-beam], [batch-beam]
        tactic_ix, arg_ix, log_probs, batch_ix  = self._select_best_results_padding(
            tactic_logits=tactic_logits,
            tactic_arg_counts=tactic_arg_counts,
            arg_logits=arg_logits,
            beam_width=beam_width
        )

        # decode
        # args (reindex back into the pre-top-k index)
        batch_tactic_ix = batch_ix * tf.cast(tactic_cnt, tf.int64) + tactic_ix  # [batch-beam]
        arg_indices = tf.gather(pre_topk_arg_indices, batch_tactic_ix)  # [batch-beam, None(args), small-cxt]
        arg_ix = tf.gather(arg_indices, arg_ix, batch_dims=2)  # [batch-beam, None(args)]

        # tactic (reindex back into tactics)
        batch_tactic_nd_ix = tf.stack([batch_ix, tactic_ix], axis=-1)  # [batch-beam, 2]
        final_tactic = tf.gather_nd(tactics, batch_tactic_nd_ix)  # [batch-beam]

        return final_tactic, arg_ix, log_probs, batch_ix

    def _select_best_results_local_global(
        self,
        tactic_logits: tf.Tensor,  # [batch, tactics]
        tactics: tf.Tensor,  # [batch, tactics] dtype:int32
        tactic_arg_counts: tf.Tensor,  # [batch, tactics] dtype:int32
        local_arg_logits: tf.RaggedTensor,  # [batch * tactics, None(args), global_cxt]
        global_arg_logits: tf.RaggedTensor,  # [batch * tactics, None(args), local_cxt]
        beam_width: tf.Tensor,  # int32
    ) -> Tuple[
        tf.RaggedTensor,  # log_probs [batch, None(options)]
        tf.RaggedTensor,  # log_probs [batch, None(options), None(tactic+args), 2]
    ]:
        # encode: combine the global and local_arguments together
        arg_logits = tf.concat([local_arg_logits, global_arg_logits], axis=-1)  # [batch*tactic, None(args), cxt]
        arg_logits = tf.ragged.map_flat_values(tf.math.log_softmax, arg_logits, axis=-1)  # [batch*tactic, None(args), cxt]
        global_arg_offset = tf.shape(local_arg_logits)[-1]

        # calculate results
        # [batch-beam], [batch-beam, None(args)], [batch-beam], [batch-beam]
        tactic, arg_ix, log_probs, batch_ix  = self._select_best_results_top_k(
            tactic_logits=tactic_logits,
            tactics=tactics,
            tactic_arg_counts=tactic_arg_counts,
            arg_logits=arg_logits,
            beam_width=beam_width
        )

        # decode
        is_local = arg_ix < global_arg_offset  # [batch-beam, None(args)]
        # 0 for local, 1 for global
        arg_type = tf.where(is_local, 0, 1)  # [batch-beam, None(args)]
        arg_index = tf.where(is_local, arg_ix, arg_ix - global_arg_offset)  # [batch-beam, None(args)]

        return tactic, arg_type, arg_index, log_probs, batch_ix

    @staticmethod
    def make_context_dense(
        logits: tf.RaggedTensor,  # [batch * top_k_tactics, None(args), None(cxt)]
        batch_size: tf.Tensor,  #int
    ) -> tf.Tensor:  # [batch * top_k_tactics, None(args), max(cxt)]
        # the cxt size is a function of the batch dim, so if the batch size is one, we can optimize this step
        if batch_size == 1:
            # the size of the cxt is constant so we can use tf.shape to make it a tensor
            logits_values_size = tf.cast(tf.shape(logits.values)[0], tf.int64)
            cxt_size = tf.maximum(tf.reduce_max(logits.values.row_lengths()), 0)
            # [1-top_k_tactics-args, cxt]
            logits_values = tf.reshape(logits.values.values, shape=[logits_values_size, cxt_size])
            # [1-top_k_tactics, None(args), cxt]
            logits = logits.with_values(
                logits_values
            )
        else:
            # [batch * top_k_tactics, None(args), cxt]
            logits = logits.with_values(logits.values.to_tensor(default_value=-np.inf))
        return logits

    def call(
        self,
        inference_output: Dict[str, Union[tf.Tensor, tf.RaggedTensor]],
    ):
        tactics = inference_output["tactic"]  # [batch, top_k_tactics]
        tactic_logits = inference_output["tactic_logits"]  # [batch, top_k_tactics]
        local_arguments_logits = inference_output["local_arguments_logits"]  # [batch * top_k_tactics, None(args), None(local_cxt)]
        global_arguments_logits = inference_output["global_arguments_logits"]  # [batch * top_k_tactics, None(args), None(global_cxt)]

        # find arg counts
        tactic_arg_counts = tf.gather(self.tactic_index_to_numargs, tactics) # [batch, top_k_tactics]

        # make last dimension dense
        batch_size = tf.shape(tactics)[0]
        # [batch * top_k_tactics, None(args), max(local_cxt)]
        local_arguments_logits = self.make_context_dense(local_arguments_logits, batch_size=batch_size)
        # [batch * top_k_tactics, None(args), max(global_cxt)]
        global_arguments_logits = self.make_context_dense(global_arguments_logits, batch_size=batch_size)

        # [batch-beam], [batch-beam, None(args)], [batch-beam, None(args)], [batch-beam]
        tactic, arg_type, arg_index, log_probs, batch_ix = self._select_best_results_local_global(
            tactic_logits=tactic_logits,
            tactics=tactics,
            tactic_arg_counts=tactic_arg_counts,
            global_arg_logits=global_arguments_logits,
            local_arg_logits=local_arguments_logits,
            beam_width=self.search_expand_bound
        )

        # return in desired form
        arg_decodings = tf.stack([arg_type, arg_index], axis=-1)  # [batch-beam, None(args), 2]
        tactic_decoding = tf.stack([tactic, tactic], axis=-1)  # [batch-beam, 2]
        tactic_decoding = tf.expand_dims(tactic_decoding, axis=1)  # [batch-beam, 1, 2]
        action_values = tf.concat([tactic_decoding, arg_decodings], axis=1)  # [batch-beam, None(1+args), 2]

        # [batch, None(beam), None(tactic+args), 2]
        actions = tf.RaggedTensor.from_value_rowids(
            values=action_values,  # [batch-beam, None(tactic+args), 2]
            value_rowids=batch_ix,  # [batch-beam]
            nrows=tf.cast(batch_size, tf.int64)
        )
        # [batch, None(beam)]
        log_probs = tf.RaggedTensor.from_value_rowids(
            values=log_probs,  # [batch-beam]
            value_rowids=batch_ix,  # [batch-beam]
            nrows=tf.cast(batch_size, tf.int64)
        )
        return log_probs, actions

class TFGNNPredict(Predict):
    def __init__(self,
                 log_dir: Path,
                 tactic_expand_bound: int,
                 search_expand_bound: int,
                 debug_dir: Optional[Path] = None,
                 checkpoint_number: Optional[int] = None,
                 exclude_tactics: Optional[List[str]] = None,
                 allocation_reserve: float = 0.5,
                 numpy_output: bool = True,
                 ):
        """
        @param log_dir: the directory for the checkpoint that is to be loaded (as passed to the Trainer class)
        @param tactic_expand_bound: the number of top base tactics to consider
        @param search_expand_bound: the max number of results to return
        @param debug_dir: set to a directory to dump pickle files for every API call that is made
        @param checkpoint_number: the checkpoint number we want to load (use `None` for the latest checkpoint)
        @param exclude_tactics: a list of tactic names to exclude from all predictions
        @param allocation_reserve: proportional size of extra allocated space when resizing the nodes embedding array
        @param numpy_output: set to True to return the predictions as a tuple of numpy arrays (for evaluation purposes)
        """

        self._exporter = DataToTFGNN()
        self._allocation_reserve = allocation_reserve

        # create dummy dataset for pre-processing purposes
        graph_constants_filepath_yaml = log_dir / 'config' / 'graph_constants.yaml'
        graph_constants_filepath_json = log_dir / 'config' / 'graph_constants.json'
        if not graph_constants_filepath_json.exists():
            logger.info(f'no json graph_constants file, trying to load yaml')
            with graph_constants_filepath_yaml.open('r') as yml_file:
                graph_constants_d = yaml.load(yml_file, Loader=yaml.SafeLoader)
            if "global_context" in graph_constants_d:
                global_context = graph_constants_d["global_context"]
                assert global_context == list(range(len(global_context)))
                del graph_constants_d["global_context"]
            with graph_constants_filepath_json.open('w') as json_file:
                json.dump(graph_constants_d, json_file)
        else:
            logger.info(f'loading json graph_constants file')
            with graph_constants_filepath_json.open('r') as json_file:
                graph_constants_d = json.load(json_file)

        graph_constants_d['data_config'] = DataConfig(**graph_constants_d['data_config'])
        graph_constants = GraphConstants(**graph_constants_d)

        # call to parent constructor to defines self.graph_constants
        super().__init__(
            graph_constants=graph_constants,
            tactic_expand_bound=tactic_expand_bound,
            search_expand_bound=search_expand_bound,
            debug_dir=debug_dir,
        )

        # to build dummy proofstates we will need to use a tactic taking no arguments
        self._dummy_tactic_id = tf.argmin(graph_constants.tactic_index_to_numargs)  # num_arguments == 0

        # the decoding mechanism currently does not support tactics with more than NUMPY_NDIM_LIMIT
        self.fixed_tactic_mask = tf.constant(np.array(graph_constants.tactic_index_to_numargs) < NUMPY_NDIM_LIMIT)

        # mask tactics explicitly excluded from predictions
        if exclude_tactics is not None:
            exclude_tactics = set(exclude_tactics)
            self.fixed_tactic_mask &= tf.constant([(tactic_name not in exclude_tactics) for tactic_name in graph_constants.tactic_index_to_string])

        # create prediction task
        prediction_yaml_filepath = log_dir / 'config' / 'prediction.yaml'
        self.prediction_task = PredictionTask.from_yaml_config(graph_constants=graph_constants,
                                                               yaml_filepath=prediction_yaml_filepath)
        self.prediction_task_type = self.prediction_task.get_config()['prediction_task_type']

        # create task to select best results from prediction task
        self.select_best_results_task = SelectBestResults(
            tactic_index_to_numargs=self.graph_constants.tactic_index_to_numargs,
            search_expand_bound=self._search_expand_bound
        )

        # create definition task
        definition_yaml_filepath = log_dir / 'config' / 'definition.yaml'
        if definition_yaml_filepath.is_file():
            self.definition_task = DefinitionTask.from_yaml_config(graph_embedding=self.prediction_task.graph_embedding,
                                                                   gnn=self.prediction_task.gnn,
                                                                   yaml_filepath=definition_yaml_filepath)
        else:
            self.definition_task = None

        # load training checkpoint
        checkpoint = tf.train.Checkpoint(prediction_task=self.prediction_task.checkpoint)
        if self.definition_task is not None:
            checkpoint.definition_task = self.definition_task.get_checkpoint()

        checkpoints_path = log_dir / 'ckpt'
        available_checkpoints = {int(re.search(r'ckpt-(\d+).index', str(ckpt)).group(1)): ckpt.with_suffix('')
                                 for ckpt in checkpoints_path.glob('*.index')}
        if checkpoint_number is None:
            checkpoint_number = max(available_checkpoints.keys())
            logger.info(f'no checkpoint number specified, using latest available checkpoint #{checkpoint_number}')
        elif checkpoint_number not in available_checkpoints.keys():
            logger.error(f'checkpoint #{checkpoint_number} is not available')
            raise ValueError(f'checkpoint number {checkpoint_number} not found')

        try:
            load_status = checkpoint.restore(save_path=str(available_checkpoints[checkpoint_number]))
        except tf.errors.OpError as error:
            logger.error(f'unable to restore checkpoint #{checkpoint_number}!')
            raise error
        else:
            load_status.expect_partial().assert_nontrivial_match().assert_existing_objects_matched().run_restore_ops()
            logger.info(f'restored checkpoint #{checkpoint_number}!')

        node_label_num = self.graph_constants.node_label_num
        extra_label_num = round(self._allocation_reserve*node_label_num)
        if extra_label_num > 0: self._allocate_definitions(node_label_num + extra_label_num)

        self._compile_network()

    def _allocate_definitions(self, new_node_label_num) -> None: # explicit change of the network array

        logger.info(f'extending global context from {self.graph_constants.node_label_num} to {new_node_label_num} elements')
        new_graph_emb_layer = self.prediction_task.graph_embedding.extend_embeddings(new_node_label_num)
        self.prediction_task.graph_embedding = new_graph_emb_layer

        if self.definition_task is not None:
            self.definition_task._graph_embedding = new_graph_emb_layer
        self.graph_constants.node_label_num = new_node_label_num
        self.prediction_task.global_arguments_logits.update_embedding_matrix(
            embedding_matrix=self.prediction_task.graph_embedding.get_node_embeddings()
        )

    @predict_api_debugging
    def allocate_definitions(self, new_node_label_num : int) -> None:
        if self.prediction_task_type != GLOBAL_ARGUMENT_PREDICTION:
            # no need to update anything if we are not going to use the global context
            return

        if new_node_label_num <= self.graph_constants.node_label_num:
            # already have sufficient array
            return

        new_node_label_num += round(self._allocation_reserve*new_node_label_num)

        self._allocate_definitions(new_node_label_num)
        self._compile_network()

    @predict_api_debugging
    def compute_new_definitions(self, new_cluster_subgraphs: List[LoaderDefinition]) -> None:
        if self.definition_task is None:
            raise RuntimeError('cannot update definitions when a definition task is not present')

        assert len(new_cluster_subgraphs) == 1
        self._compute_and_replace_definition_embs(new_cluster_subgraphs[0])

    @tf.function(input_signature = (LoaderProofstateSpec,))
    def _make_proofstate_graph_tensor(self, state : LoaderProofstate):
        action = LoaderAction(
            self._dummy_tactic_id,
            tf.zeros(shape=(0), dtype=tf.int64),
            tf.zeros(shape=(0), dtype=tf.int64),
        )
        graph_id = tf.constant(-1, dtype=tf.int64)
        x = DataToTFGNN.proofstate_to_graph_tensor(state, action, graph_id)
        return x

    def _compile_network(self):
        @tf.function(input_signature = (LoaderDefinitionSpec,))
        def compute_and_replace_definition_embs(loader_definition):
            graph_tensor = self._exporter.definition_to_graph_tensor(loader_definition)
            definition_graph = stack_graph_tensors([graph_tensor])

            scalar_definition_graph = definition_graph.merge_batch_to_components()
            definition_embeddings = self.definition_task(scalar_definition_graph).flat_values
            defined_labels = Trainer._get_defined_labels(definition_graph).flat_values

            self.prediction_task.graph_embedding.update_node_embeddings(
                embeddings=definition_embeddings,
                indices=defined_labels
            )
        self._compute_and_replace_definition_embs = compute_and_replace_definition_embs

        inference_model_bare = self.prediction_task.create_inference_model(
            tactic_expand_bound=self._tactic_expand_bound,
            graph_constants=self.graph_constants
        )
        allowed_model_tactics_spec = tf.TensorSpec(shape=(None,), dtype=tf.int32)
        @tf.function(input_signature = (LoaderProofstateSpec, allowed_model_tactics_spec))
        def inference_model(state, allowed_model_tactics):
            tactic_mask = tf.scatter_nd(
                indices = tf.expand_dims(allowed_model_tactics, axis = 1),
                updates = tf.ones_like(allowed_model_tactics, dtype=bool),
                shape = [self.graph_constants.tactic_num]
            )
            graph_tensor_single = self._make_proofstate_graph_tensor(state)
            graph_tensor_stacked = stack_graph_tensors([graph_tensor_single])
            inference_output = inference_model_bare({
                self.prediction_task.PROOFSTATE_GRAPH: graph_tensor_stacked,
                self.prediction_task.TACTIC_MASK: tf.expand_dims(tactic_mask, axis=0),
            })
            return self.select_best_results_task(inference_output)
        self._inference_model = inference_model

        dummy_graph = LoaderGraph(
            nodes = np.zeros([1], dtype = int),
            edges = np.zeros([0,2], dtype = int),
            edge_labels = np.zeros([0], dtype = int),
            edge_offsets = np.zeros([0], dtype = int),
        )
        dummy_loader_proofstate = LoaderProofstate(
            graph = dummy_graph,
            root = 0,
            context = ProofstateContext(
                local_context = np.zeros([0], dtype = int),
                global_context = np.zeros([1], dtype = int), # make sure at least one element in global context
            ),
            metadata = ProofstateMetadata(
                name = "",
                step = 0,
                is_faithful = False,
            ),
        )

        # run on a dummy input to force precompilation
        self._inference_model(dummy_loader_proofstate, np.zeros([1], dtype = int))

    # Currently not used
    def _make_proofstate_batch(self, datapoints : Iterable[LoaderProofstate]):
        return stack_graph_tensors([
            self._make_proofstate_graph_tensor(x)
            for x in datapoints
        ])

    @predict_api_debugging
    def ranked_predictions(self,
                           state: LoaderProofstate,
                           available_global: Optional[np.ndarray] = None,
                           allowed_model_tactics: Optional[Iterable[int]] = None
                           ) -> Tuple[np.ndarray, np.ndarray]:  #
        """
        Produces predictions for a single proof-state.
        """
        if available_global is not None:
            raise NotImplementedError('available_global is not supported yet')

        best_log_probs, best_tactics = self._inference_model(state, allowed_model_tactics)
        return (best_tactics.numpy()[0], best_log_probs.numpy()[0])

    # (!) NOT MAINTAINED
    def _evaluate(self,
                  proofstate_graph_dataset: tf.data.Dataset,
                  batch_size: int,
                  tactic_expand_bound: int,
                  total_expand_bound: int,
                  search_expand_bound: Optional[int] = None,
                  allowed_model_tactics: Optional[Iterable[int]] = None
                  ) -> Tuple[float, float]:  # per_proofstate_passrate, per_lemma_passrate

        raise NotImplemented("Running unmaintained code. Delete this line at your own risk.")


import numpy as np
import pandas as pd
import hashlib
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from sklearn.feature_selection import SelectKBest, f_classif
from xgboost import XGBClassifier
from sklearn.utils.class_weight import compute_class_weight
from imblearn.under_sampling import CondensedNearestNeighbour
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import ConfusionMatrixDisplay
from sklearn import metrics, random
from sklearn.ensemble import RandomForestClassifier, StackingClassifier
from imblearn.over_sampling import SMOTE
from sklearn.metrics import roc_curve, roc_auc_score, auc
from sklearn.preprocessing import label_binarize
from sklearn.metrics import f1_score
from sklearn.metrics import matthews_corrcoef
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_curve, auc, balanced_accuracy_score
from sklearn.preprocessing import label_binarize
from scipy import interp
from itertools import cycle
from sklearn.metrics import precision_recall_fscore_support
import shap

def make_seed(seedStr):
    return int(hashlib.md5(seedStr.encode("utf-8")).hexdigest()[24:], 16)

def get_actual_class(tumor_code, actual_classes):
    return actual_classes.get(tumor_code)
def optimize_thresholds(y_test, y_pred_proba):
    best_mcc = -1
    best_thresholds = (0, 0)

    for t_low in np.linspace(0, 1, num=100):
        for t_high in np.linspace(t_low, 1, num=100):
            y_pred = np.zeros_like(y_test)
            y_pred[y_pred_proba[:, 0] > t_low] = 0
            y_pred[y_pred_proba[:, 1] > t_high] = 1

            mcc = matthews_corrcoef(y_test, y_pred)

            if mcc > best_mcc:
                best_mcc = mcc
                best_thresholds = (t_low, t_high)
                print (best_mcc)
    return best_thresholds

dataDel = pd.read_csv("Delayed1.csv")
dataPre = pd.read_csv("Precontrast1.csv")
dataArt = pd.read_csv("Arterial1.csv")
dataVen = pd.read_csv("Venous1.csv")

data = pd.concat([dataDel, dataPre, dataArt, dataVen], axis=1)

labelclassnum2 = data['Grade']
code_data = data[['tumor and ID']].copy()
data = data.drop(['RPCID', 'tumorKey', 'patientId', 'Grade', 'tumor and ID'], axis=1)

datafeatures = np.array(data)

accuracy = np.zeros([100])
output_file = "all_reports_stack.txt"
with open(output_file, "w") as f:
    f.write("")
print ('Starting the Code ...')
misclassified_count = {}
actual_classes = {}
for i in range(100):

    seed1 = make_seed(f"growth1{i}")
    seed2 = make_seed(f"growth2{i}")
    seed3 = make_seed(f"growth3{i}")
    seed4 = make_seed(f"growth4{i}")

    np.random.seed(seed3)
    random.seed(seed4)

    x_train, x_test, y_train, y_test, train_index, test_index = train_test_split(datafeatures, labelclassnum2, data.index,
                                                                                test_size=0.15, random_state=seed1,
                                                                                stratify=labelclassnum2)

    cnn = CondensedNearestNeighbour(random_state=42)
    x_train, y_train = cnn.fit_resample(x_train, y_train)

    smote = SMOTE(random_state=42)
    x_train, y_train = smote.fit_resample(x_train, y_train)

    class_weights = compute_class_weight(class_weight='balanced', classes=np.unique(y_train), y=y_train)
    weights = dict(zip(np.unique(y_train), class_weights))
    scale_pos_weight = weights[1] / weights[0]

    scaler = StandardScaler()
    x_train = scaler.fit_transform(x_train)
    x_test = scaler.transform(x_test)

    k_best = SelectKBest(score_func=f_classif, k=20)
    x_train = k_best.fit_transform(x_train, y_train)
    x_test = k_best.transform(x_test)
    print ('DONE')
    # Grid search for XGBoost
    xgb_model = XGBClassifier(tree_method='gpu_hist', gpu_id=0)
    xgb_param_grid = {
        'scale_pos_weight': [scale_pos_weight],
        'max_depth': [3, 4, 5, 6],
        'min_child_weight': [1, 2, 3],
        'gamma': [0, 0.1, 0.2],
        'colsample_bytree': [0.6, 0.8, 1.0],
        'subsample': [0.6, 0.8, 1.0],
    }

    xgb_grid_search = GridSearchCV(xgb_model, xgb_param_grid, scoring='f1', cv=5, n_jobs=-1)
    xgb_grid_search.fit(x_train, y_train)
    print ('DONE XGB')
    # Grid search for Random Forest
    rf_model = RandomForestClassifier()
    rf_param_grid = {
        'n_estimators': [100, 200, 300],
        'max_depth': [None, 3, 4, 5],
        'min_samples_split': [2, 5, 10],
        'min_samples_leaf': [1, 2, 4],
        'class_weight': ['balanced', 'balanced_subsample'],
    }
    print ('DONE RF')
    rf_grid_search = GridSearchCV(rf_model, rf_param_grid, scoring='f1', cv=5, n_jobs=-1)
    rf_grid_search.fit(x_train, y_train)

    best_xgb_model = xgb_grid_search.best_estimator_
    best_xgb_model.fit(x_train, y_train)

    best_rf_model = rf_grid_search.best_estimator_
    best_rf_model.fit(x_train, y_train)
    '''
    xgb_y_pred_proba = best_xgb_model.predict_proba(x_test)
    rf_y_pred_proba = best_rf_model.predict_proba(x_test)
    avg_y_pred_proba = (xgb_y_pred_proba + rf_y_pred_proba) / 2
    avg_y_pred = np.argmax(avg_y_pred_proba, axis=1)

    # ENSEMBLING: Evaluate the performance of the ensemble model
    ensemble_accuracy = accuracy_score(y_test, avg_y_pred)
    print("Ensemble Accuracy: %.2f%%" % (ensemble_accuracy * 100.0))
    '''
    stacking_clf = StackingClassifier(
        estimators=[('xgb', best_xgb_model), ('rf', best_rf_model)],
        final_estimator=LogisticRegression(),
        cv=5
    )
    stacking_clf.fit(x_train, y_train)

    stacked_y_pred = stacking_clf.predict(x_test)
    stacked_y_pred_proba = stacking_clf.predict_proba(x_test)
    stacked_accuracy = accuracy_score(y_test, stacked_y_pred)
    print("Stacking Classifier Accuracy: %.2f%%" % (stacked_accuracy * 100.0))

    test_code_data = code_data.iloc[test_index]
    pred_data = list(zip(test_code_data['tumor and ID'], stacked_y_pred, y_test))

    for code, pred, actual in pred_data:
        if code not in misclassified_count:
            misclassified_count[code] = 0
            actual_classes[code] = actual
        if pred != actual:
            misclassified_count[code] += 1

    ensemble_report = classification_report(y_test, stacked_y_pred)
    print(ensemble_report)

    y_pred_proba_low_grade = stacked_y_pred_proba[:, 0]
    y_pred_proba_high_grade = stacked_y_pred_proba[:, 1]

    optimal_threshold_low, optimal_threshold_high = optimize_thresholds(y_test, stacked_y_pred_proba)

    y_pred_optimal = np.zeros_like(y_test)
    y_pred_optimal = np.where(stacked_y_pred_proba[:, 1] > optimal_threshold_high, 1, 0)
    optimal_classification_report = classification_report(y_test, y_pred_optimal)
    print(optimal_classification_report)

    optimal_y_pred_proba_pos = stacked_y_pred_proba[:, 1]

    optimal_fpr, optimal_tpr, _ = roc_curve(y_test, optimal_y_pred_proba_pos)
    optimal_roc_auc = auc(optimal_fpr, optimal_tpr)

    lw = 2
    plt.figure()
    plt.plot(optimal_fpr, optimal_tpr, color='darkorange', lw=lw,
             label=f'ROC curve (area = {optimal_roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], 'k--', lw=lw)
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic (ROC)')
    plt.legend(loc="lower right")
    plt.savefig(f'roc_curve_optimal_stack{i}.png')
    plt.close()

    optimal_conf_mat = confusion_matrix(y_test, y_pred_optimal)
    disp = ConfusionMatrixDisplay(confusion_matrix=optimal_conf_mat, display_labels=['LowGrade', 'HighGrade'])
    disp.plot(cmap='Blues')
    plt.savefig(f'optimal_confusion_matrix_stack{i}.png')
    '''
    explainer = shap.Explainer(stacking_clf.named_estimators_['xgb'], x_train)

    # Calculate SHAP values for the x_test dataset
    shap_values = explainer(x_test)

    # Create a SHAP force plot and save it to a file
    test_sample_idx = 0  # Choose the index of the test sample you want to visualize
    shap.plots.force(shap_values[test_sample_idx], matplotlib=True, show=False)
    plt.savefig(f'shap_force_plot_optimal_iter_{i+1}.png', bbox_inches='tight')
    plt.close()
    '''
    with open(output_file, "a") as f:
        f.write(f"Iteration {i+1} - Optimal Thresholds\n")
        f.write(f"Low Grade: {optimal_threshold_low}, High Grade: {optimal_threshold_high}\n")
        f.write(optimal_classification_report)
        f.write("\n")

threshold = max(20, 0.25 * 100)
repeatedly_wrong_codes = {code: count for code, count in misclassified_count.items() if count >= threshold}

with open("repeatedly_wrong_codes_stack.csv", "w") as f:
    f.write("Code,Count,Actual\n")
    for code, count in repeatedly_wrong_codes.items():
        actual_class = get_actual_class(code, actual_classes)
        f.write(f"{code},{count},{actual_class}\n")

average_accuracy = np.mean(accuracy)
print("Average accuracy: %.2f%%" % (average_accuracy * 100.0))


import time
import os
import pip

if os.name == 'nt':
    clear = lambda: os.system('cls')
else:
    clear = lambda: os.system('clear')

def import_or_install():
    try:
        from pydrive.auth import GoogleAuth
        from pydrive.drive import GoogleDrive
        import pyperclip
    except ImportError:
        pip.main(['install', 'pydrive'])
        pip.main(['install', 'pyperclip'])

        clear()
        print('pydrive has been installed!')
import_or_install()

from requests import get
from pydrive.auth import GoogleAuth
from pydrive.drive import GoogleDrive
from requests.sessions import extract_cookies_to_jar

os.chdir('../credentials')
gauth = GoogleAuth()
gauth.LocalWebserverAuth()
drive = GoogleDrive(gauth)
os.chdir('../data')
externalIPAddress = ''

class UpdateIP:

    def DriveRootFolderQuery():
        global externalIPAddress
        #print('externalip in DriveRootFolderQuery is')
        #print(externalIPAddress)
        externalIPAddress = get('https://api.ipify.org').text
        file_list = drive.ListFile({'q': "'root' in parents and trashed=false"}).GetList()
        for file1 in file_list:
            if  'IPUpdate.txt' in file1['title']:
                #print('title: %s, id: %s' % (file1['title'], file1['id']))
                UpdateIP.UpdateFileContent(file1)
        return

    def UpdateFileContent(file1):
        global externalIPAddress
        if externalIPAddress != file1.GetContentString():
            print('The ip does not match. Trying to set the content to the global ip in string.')

            file1.SetContentString(externalIPAddress) # Set content of the file from given string.
            file1.Upload()
        else:

            print('The ip matches. Sleeping for 300 seconds')
            time.sleep(300)

while True:
    UpdateIP.DriveRootFolderQuery()


import abc
from tempfile import TemporaryDirectory
from typing import Any, Callable, Dict, List, Tuple, Type, Union

from nebullvm.config import ACTIVATION_METRIC_DROP_THS
from nebullvm.core.models import (
    OptimizedModel,
    OptimizationTime,
    ModelParams,
    ModelCompiler,
    ModelCompressor,
    DeepLearningFramework,
    DeviceType,
    QuantizationType,
)
from nebullvm.operations.base import Operation
from nebullvm.operations.inference_learners.base import (
    BuildInferenceLearner,
)
from nebullvm.operations.inference_learners.builders import (
    DeepSparseBuildInferenceLearner,
    FasterTransformerBuildInferenceLearner,
    IntelNeuralCompressorBuildInferenceLearner,
    ONNXApacheTVMBuildInferenceLearner,
    ONNXBuildInferenceLearner,
    ONNXTensorRTBuildInferenceLearner,
    OpenVINOBuildInferenceLearner,
    PyTorchApacheTVMBuildInferenceLearner,
    PyTorchTensorRTBuildInferenceLearner,
    TensorflowBuildInferenceLearner,
    TFLiteBuildInferenceLearner,
    TorchNeuronBuildInferenceLearner,
    TorchXLABuildInferenceLearner,
    TorchDynamoBuildInferenceLearner,
    TorchScriptBuildInferenceLearner,
)
from nebullvm.operations.measures.measures import MetricDropMeasure
from nebullvm.operations.measures.utils import (
    compute_optimized_running_time,
    compute_relative_difference,
)
from nebullvm.operations.optimizations.compilers.base import Compiler
from nebullvm.operations.optimizations.compilers.deepsparse import (
    DeepSparseCompiler,
)
from nebullvm.operations.optimizations.compilers.faster_transformer import (
    FasterTransformerCompiler,
)
from nebullvm.operations.optimizations.compilers.intel_neural_compressor import (  # noqa: E501
    IntelNeuralCompressorCompiler,
)
from nebullvm.operations.optimizations.compilers.onnxruntime import (
    ONNXCompiler,
)
from nebullvm.operations.optimizations.compilers.openvino import (
    OpenVINOCompiler,
)
from nebullvm.operations.optimizations.compilers.tensor_rt import (
    ONNXTensorRTCompiler,
    PyTorchTensorRTCompiler,
)
from nebullvm.operations.optimizations.compilers.tensorflow import (
    TensorflowBackendCompiler,
    TFLiteBackendCompiler,
)
from nebullvm.operations.optimizations.compilers.torch_dynamo import (
    TorchDynamoCompiler,
)
from nebullvm.operations.optimizations.compilers.torch_neuron import (
    TorchNeuronCompiler,
)
from nebullvm.operations.optimizations.compilers.torch_xla import (
    TorchXLACompiler,
)
from nebullvm.operations.optimizations.compilers.torchscript import (
    TorchScriptCompiler,
)
from nebullvm.operations.optimizations.compilers.tvm import (
    ONNXApacheTVMCompiler,
    PyTorchApacheTVMCompiler,
)
from nebullvm.optional_modules.tensorflow import tensorflow as tf
from nebullvm.optional_modules.torch import torch
from nebullvm.tools.data import DataManager
from nebullvm.tools.transformations import MultiStageTransformation
from nebullvm.tools.utils import get_throughput

class Optimizer(Operation, abc.ABC):
    def __init__(self):
        super().__init__()
        self.optimized_models = []
        self.source_dl_framework = None
        self.pipeline_dl_framework = None
        self.compiler_ops = {}
        self.build_inference_learner_ops = {}
        self.validity_check_op = MetricDropMeasure()

    def execute(
        self,
        model: Any,
        input_data: DataManager,
        optimization_time: OptimizationTime,
        metric_drop_ths: float,
        metric: Callable,
        model_params: ModelParams,
        model_outputs: List[Tuple[Any, ...]],
        ignore_compilers: List[ModelCompiler],
        ignore_compressors: List[ModelCompressor],
        source_dl_framework: DeepLearningFramework,
        is_diffusion: bool = False,
    ) -> List[OptimizedModel]:
        self.source_dl_framework = source_dl_framework

        # TODO: implement and select compressors from hardware

        compilers = self._select_compilers_from_hardware()

        remove_compiler_list = []
        add_compiler_list = []
        for compiler in ignore_compilers:
            if compiler in MULTI_FRAMEWORK_COMPILERS:
                add_compiler_list += MULTI_FRAMEWORK_COMPILERS[compiler]
                remove_compiler_list.append(compiler)

        for c in remove_compiler_list:
            ignore_compilers.remove(c)

        ignore_compilers += add_compiler_list

        (
            self.compiler_ops,
            self.build_inference_learner_ops,
        ) = self._load_compilers(
            ignore_compilers=ignore_compilers,
            compilers=compilers,
        )
        self._optimize(
            model=model,
            input_data=input_data,
            optimization_time=optimization_time,
            metric_drop_ths=metric_drop_ths,
            metric=metric,
            model_params=model_params,
            model_outputs=model_outputs,
            ignore_compilers=ignore_compilers,
            is_diffusion=is_diffusion,
        )

        return self.optimized_models

    @abc.abstractmethod
    def _select_compilers_from_hardware(self):
        raise NotImplementedError()

    @staticmethod
    def _load_compilers(
        ignore_compilers: List[ModelCompiler],
        compilers: List[ModelCompiler],
    ):
        compiler_ops = {
            compiler: COMPILER_TO_OPTIMIZER_MAP[compiler]()
            for compiler in compilers
            if compiler not in ignore_compilers
            and compiler in COMPILER_TO_OPTIMIZER_MAP
        }
        build_inference_learner_ops = {
            compiler: COMPILER_TO_INFERENCE_LEARNER_MAP[compiler]()
            for compiler in compilers
            if compiler not in ignore_compilers
            and compiler in COMPILER_TO_OPTIMIZER_MAP
        }

        return compiler_ops, build_inference_learner_ops

    def free_model_gpu(self, model: Any):
        # Free gpu memory
        if self.device.type is DeviceType.GPU:
            try:
                model.cpu()
            except Exception:
                pass
            try:
                with torch.cuda.device(self.device.to_torch_format()):
                    torch.cuda.empty_cache()
            except Exception:
                pass

    def _optimize(
        self,
        model: Union[torch.nn.Module, tf.Module, str],
        input_data: DataManager,
        optimization_time: OptimizationTime,
        metric_drop_ths: float,
        metric: Callable,
        model_params: ModelParams,
        model_outputs: List[Tuple[Any, ...]],
        ignore_compilers: List[ModelCompiler],
        is_diffusion: bool = False,
    ):

        if metric_drop_ths is not None:
            q_types = [
                None,
            ]
            if metric_drop_ths > 0:
                q_types.append(QuantizationType.HALF)
            if metric_drop_ths > ACTIVATION_METRIC_DROP_THS:
                q_types.append(QuantizationType.DYNAMIC)
                if input_data is not None:
                    q_types.append(QuantizationType.STATIC)
        else:
            q_types = [None]

        optimization_info = []
        for compiler, compiler_op, build_inference_learner_op in zip(
            self.compiler_ops.keys(),
            self.compiler_ops.values(),
            self.build_inference_learner_ops.values(),
        ):
            for q_type in q_types:
                input_tfms = MultiStageTransformation([])

                self.free_model_gpu(model)

                with TemporaryDirectory() as tmp_dir:
                    try:
                        compiler_op.to(self.device).execute(
                            model=model,
                            input_data=input_data,
                            model_params=model_params,
                            metric_drop_ths=metric_drop_ths
                            if q_type is not None
                            else None,
                            quantization_type=q_type,
                            input_tfms=input_tfms,
                            onnx_output_path=tmp_dir,
                            is_diffusion=is_diffusion,
                        )

                        compiled_model = compiler_op.get_result()
                        if compiled_model is not None:
                            build_inference_learner_op.to(self.device).execute(
                                model=compiled_model,
                                model_orig=compiler_op.model_orig
                                if hasattr(compiler_op, "model_orig")
                                else None,
                                model_params=model_params,
                                input_tfms=input_tfms,
                                source_dl_framework=self.source_dl_framework,
                                quantization_type=q_type,
                            )
                            inference_learner = (
                                build_inference_learner_op.get_result()
                            )

                            if inference_learner is not None:
                                test_input_data, ys = input_data.get_split(
                                    "test"
                                ).get_list(with_ys=True)

                                self.validity_check_op.execute(
                                    inference_learner,
                                    test_input_data,
                                    model_outputs,
                                    metric_drop_ths,
                                    metric_func=metric
                                    if q_type is not None
                                    else compute_relative_difference,
                                    ys=ys,
                                )

                                if self.validity_check_op.valid:
                                    latency = compute_optimized_running_time(
                                        inference_learner, input_data
                                    )
                                    self.logger.info(
                                        f"Optimized model latency: {latency} "
                                        f"sec/iter"
                                    )

                                    if (
                                        compiler not in ignore_compilers
                                        and optimization_time
                                        is OptimizationTime.CONSTRAINED
                                    ):
                                        ignore_compilers.append(compiler)

                                    self.optimized_models.append(
                                        OptimizedModel(
                                            inference_learner=inference_learner,  # noqa: E501
                                            metric_drop=self.validity_check_op.measure_result,  # noqa: E501
                                            compiler=compiler,
                                            technique=q_type.name
                                            if q_type is not None
                                            else "None",
                                            latency_seconds=latency,
                                            throughput=get_throughput(
                                                latency,
                                                # Normal models have batch
                                                # size B, diffusion models
                                                # have batch size 2B
                                                model_params.batch_size
                                                if not is_diffusion
                                                else model_params.batch_size
                                                / 2,
                                            ),
                                            size_mb=inference_learner.get_size()  # noqa: E501
                                            / 1e6,
                                        )
                                    )

                                    opt_info_dict = {
                                        "compiler": f"{self.pipeline_dl_framework.value}_{compiler.value}",  # noqa: E501
                                        "technique": q_type.value
                                        if q_type
                                        else "none",
                                        "latency": latency,
                                    }
                                    if (
                                        metric_drop_ths is not None
                                        and q_type is not None
                                    ):
                                        opt_info_dict[
                                            "metric_loss"
                                        ] = (
                                            self.validity_check_op.measure_result  # noqa: E501
                                        )
                                        opt_info_dict[
                                            "metric"
                                        ] = metric.__name__
                                    optimization_info.append(opt_info_dict)
                                else:
                                    self.logger.warning(
                                        "The optimized model will be "
                                        "discarded due to poor results "
                                        "obtained with the given metric."
                                    )

                                if self.device.type in [
                                    DeviceType.GPU,
                                    DeviceType.TPU,
                                ]:
                                    inference_learner.free_gpu_memory()
                    except Exception as ex:
                        self.logger.warning(
                            f"Optimization failed with "
                            f"{self.pipeline_dl_framework} "
                            f"interface of {compiler}. Got error {ex}. "
                            f"If possible the compilation will be re-scheduled"
                            f" with another interface. Please consult the "
                            f"documentation for further info or open an issue "
                            f"on GitHub for receiving assistance."
                        )
                        optimization_info.append(
                            {
                                "compiler": compiler.value,
                                "technique": q_type.value
                                if q_type
                                else "none",
                                "latency": -1,
                            }
                        )
        if self.feedback_collector is not None:
            self.feedback_collector.store_info(
                key="optimizations",
                value=optimization_info,
            )

MULTI_FRAMEWORK_COMPILERS = {
    ModelCompiler.TENSOR_RT: [
        ModelCompiler.TENSOR_RT_TORCH,
        ModelCompiler.TENSOR_RT_ONNX,
    ],
    ModelCompiler.APACHE_TVM: [
        ModelCompiler.APACHE_TVM_TORCH,
        ModelCompiler.APACHE_TVM_ONNX,
    ],
}

COMPILER_TO_OPTIMIZER_MAP: Dict[ModelCompiler, Type[Compiler]] = {
    ModelCompiler.TORCHSCRIPT: TorchScriptCompiler,
    ModelCompiler.DEEPSPARSE: DeepSparseCompiler,
    ModelCompiler.INTEL_NEURAL_COMPRESSOR: IntelNeuralCompressorCompiler,
    ModelCompiler.TENSOR_RT_TORCH: PyTorchTensorRTCompiler,
    ModelCompiler.TENSOR_RT_ONNX: ONNXTensorRTCompiler,
    ModelCompiler.APACHE_TVM_TORCH: PyTorchApacheTVMCompiler,
    ModelCompiler.APACHE_TVM_ONNX: ONNXApacheTVMCompiler,
    ModelCompiler.ONNX_RUNTIME: ONNXCompiler,
    ModelCompiler.OPENVINO: OpenVINOCompiler,
    ModelCompiler.TFLITE: TFLiteBackendCompiler,
    ModelCompiler.XLA: TensorflowBackendCompiler,
    ModelCompiler.TORCH_NEURON: TorchNeuronCompiler,
    ModelCompiler.TORCH_XLA: TorchXLACompiler,
    ModelCompiler.TORCH_DYNAMO: TorchDynamoCompiler,
    ModelCompiler.FASTER_TRANSFORMER: FasterTransformerCompiler,
}

COMPILER_TO_INFERENCE_LEARNER_MAP: Dict[
    ModelCompiler, Type[BuildInferenceLearner]
] = {
    ModelCompiler.TORCHSCRIPT: TorchScriptBuildInferenceLearner,
    ModelCompiler.DEEPSPARSE: DeepSparseBuildInferenceLearner,
    ModelCompiler.INTEL_NEURAL_COMPRESSOR: IntelNeuralCompressorBuildInferenceLearner,  # noqa: E501
    ModelCompiler.TENSOR_RT_TORCH: PyTorchTensorRTBuildInferenceLearner,
    ModelCompiler.TENSOR_RT_ONNX: ONNXTensorRTBuildInferenceLearner,
    ModelCompiler.APACHE_TVM_TORCH: PyTorchApacheTVMBuildInferenceLearner,
    ModelCompiler.APACHE_TVM_ONNX: ONNXApacheTVMBuildInferenceLearner,
    ModelCompiler.ONNX_RUNTIME: ONNXBuildInferenceLearner,
    ModelCompiler.OPENVINO: OpenVINOBuildInferenceLearner,
    ModelCompiler.TFLITE: TFLiteBuildInferenceLearner,
    ModelCompiler.XLA: TensorflowBuildInferenceLearner,
    ModelCompiler.TORCH_NEURON: TorchNeuronBuildInferenceLearner,
    ModelCompiler.TORCH_XLA: TorchXLABuildInferenceLearner,
    ModelCompiler.TORCH_DYNAMO: TorchDynamoBuildInferenceLearner,
    ModelCompiler.FASTER_TRANSFORMER: FasterTransformerBuildInferenceLearner,
}


import time
from time import sleep
from selenium import webdriver
from selenium.webdriver import ActionChains

def move(distance):
    """
    :param distance: 距离 int
    :return: 距离总和为distance的数组 tracks[]
    """
    v = 0
    t = 0.2
    x = distance
    tracks = []
    current = 0
    # x = x+10
    while current < x:
        a = 10
        v0 = v
        # 单位时间内位移公式
        s = v0 * t + 0.5 * a * (t ** 2)
        # 当前位移
        current = current + s
        tracks.append(round(s))
        v = v0 + a * t
    return tracks

def main():
    # 点击登录
    driver = webdriver.Chrome(executable_path='F:/pythonProject/pubilcVenv/seleniumDriver/chrome 91/chromedriver.exe')
    driver.get('https://www.qcc.com/')

    # wait
    sleep(5)

    longin = driver.find_element_by_xpath(r"//a[@onclick='showLoginModal();']/span")
    longin.click()
    sleep(3)

    #sliding_block
    sliding_block = driver.find_element_by_xpath(r"//span[@class='nc_iconfont btn_slide']")
    ActionChains(driver).click_and_hold(sliding_block).perform()
    for x in move(distance=500):
        ActionChains(driver).move_by_offset(xoffset=x, yoffset=0).perform()
    ActionChains(driver).release(sliding_block).perform()
    time.sleep(3)

    pass

if __name__ == '__main__':
    main()


from decimal import Decimal

from pydantic import Field, SecretStr

from hummingbot.client.config.config_data_types import BaseConnectorConfigMap, ClientFieldData
from hummingbot.core.data_type.trade_fee import TradeFeeSchema

DEFAULT_FEES = TradeFeeSchema(
    maker_percent_fee_decimal=Decimal("0.0002"),
    taker_percent_fee_decimal=Decimal("0.0004"),
    buy_percent_fee_deducted_from_returns=True
)

CENTRALIZED = True

EXAMPLE_PAIR = "BTC-USDT"

BROKER_ID = "x-3QreWesy"

class BinancePerpetualConfigMap(BaseConnectorConfigMap):
    connector: str = Field(default="binance_perpetual", client_data=None)
    binance_perpetual_api_key: SecretStr = Field(
        default=...,
        client_data=ClientFieldData(
            prompt=lambda cm: "Enter your Binance Perpetual API key",
            is_secure=True,
            is_connect_key=True,
            prompt_on_new=True,
        )
    )
    binance_perpetual_api_secret: SecretStr = Field(
        default=...,
        client_data=ClientFieldData(
            prompt=lambda cm: "Enter your Binance Perpetual API secret",
            is_secure=True,
            is_connect_key=True,
            prompt_on_new=True,
        )
    )

KEYS = BinancePerpetualConfigMap.construct()

OTHER_DOMAINS = ["binance_perpetual_testnet"]
OTHER_DOMAINS_PARAMETER = {"binance_perpetual_testnet": "binance_perpetual_testnet"}
OTHER_DOMAINS_EXAMPLE_PAIR = {"binance_perpetual_testnet": "BTC-USDT"}
OTHER_DOMAINS_DEFAULT_FEES = {"binance_perpetual_testnet": [0.02, 0.04]}

class BinancePerpetualTestnetConfigMap(BaseConnectorConfigMap):
    connector: str = Field(default="binance_perpetual_testnet", client_data=None)
    binance_perpetual_testnet_api_key: SecretStr = Field(
        default=...,
        client_data=ClientFieldData(
            prompt=lambda cm: "Enter your Binance Perpetual testnet API key",
            is_secure=True,
            is_connect_key=True,
            prompt_on_new=True,
        )
    )
    binance_perpetual_testnet_api_secret: SecretStr = Field(
        default=...,
        client_data=ClientFieldData(
            prompt=lambda cm: "Enter your Binance Perpetual testnet API secret",
            is_secure=True,
            is_connect_key=True,
            prompt_on_new=True,
        )
    )

    class Config:
        title = "binance_perpetual"

OTHER_DOMAINS_KEYS = {"binance_perpetual_testnet": BinancePerpetualTestnetConfigMap.construct()}


from flask import Flask, render_template, request
import base64
from scipy.misc import imsave, imread, imresize
import numpy as np
import pdb
import sys
import os
from load_DL_data import *
app = Flask(__name__)

global model, graph
model,graph = init()

@app.route('/draw_canvas')
def index():
   return render_template('canvas.html')

@app.route('/predict', methods=['POST', 'GET'])
def precict_data():
   print("under predict")
   enc_img = request.get_data()
   # pdb.Pdb(stdout=sys.__stdout__).set_trace()
   enc_img = enc_img[22:]
   with open('canvas_img.png', 'wb') as img :
         img.write(base64.b64decode(enc_img))

   # img_name = 'canvas_img.png'  # I assume you have a way of picking unique filenames
   dec_img = imread('canvas_img.png', mode='L')
   dec_img_swap = np.invert(dec_img)
   dec_img_swap = imresize(dec_img_swap, (28,28))
   dec_img_swap = dec_img_swap.reshape(1, 784)

   with graph.as_default():
       out = model.predict(dec_img_swap)
       response = np.array_str(np.argmax(out,axis=1))
       return response

if __name__ == '__main__':
   app.run(host='0.0.0.0',debug = True)


import asyncio

def format_response(given_command_type: str, command_dict: dict, execution_result: str) -> dict:
    """Format the response to be sent back to the client."""
    math_expression = command_dict.get("expression", None)
    output = {
        given_command_type: command_dict.get('command_name', math_expression) + " " + " ".join(
            command_dict.get("parameters", [])),
        "result": execution_result
    }
    return output

async def systemic_processing(command_name: str, parameters: list) -> str:
    """Perform some systemic processing on the command."""
    proc = await asyncio.create_subprocess_exec(command_name, *parameters, stdout=asyncio.subprocess.PIPE,
                                                stderr=asyncio.subprocess.PIPE)
    await proc.wait()
    result = await proc.stdout.read()

    # If the result is empty, then there was an error
    if result == b'':
        result = await proc.stderr.read()

    return result.decode()

async def computational_command_processing(expression: str) -> str:
    result = eval(expression)
    return str(result)


import sys
from PyQt6.QtWidgets import QWidget, QApplication
from PyQt6.QtGui import QPainter, QColor, QFont
from PyQt6.QtCore import Qt

# We begin with drawing some Unicode text on the client area of a window.
class Example(QWidget):

    def __init__(self):
        super().__init__()

        self.text = None
        self.initUI()

    def initUI(self):
        self.text = "Shaxzodbek Sobitov Ravshan o'g'li\nInagamova Shoira"

        self.setGeometry(300, 300, 350, 300)
        self.setWindowTitle('Drawing text')
        self.show()

    # In our example, we draw some text in Cylliric.
    # The text is vertically and horizontally aligned.
    def paintEvent(self, event):
        # Drawing is done within the paint event.
        qp = QPainter()
        qp.begin(self)

        """
        The QPainter class is responsible for all the low-level painting.
        All the painting methods go between begin and end methods.
        The actual painting is delegated to the drawText method.
        """
        self.drawText(event, qp)
        qp.end()

    def drawText(self, event, qp):
        # Here we define a pen and a font which are used to draw the text.
        qp.setPen(QColor(168, 34, 3))
        qp.setFont(QFont('Decorative', 10))

        """
        The drawText method draws text on the window.
        The rect method of the paint event returns the rectangle that needs to be updated.
        With the Qt.Alignment.AlignCenter we align the text in both dimensions.
        """
        qp.drawText(event.rect(), Qt.AlignmentFlag.AlignCenter, self.text)

def main():
    app = QApplication(sys.argv)
    ex = Example()
    sys.exit(app.exec())

if __name__ == '__main__':
    main()


# This file contains code for finding occurrences of techniques in a hobby training log
#
# Author: Josh McIntyre
#

import re
import csv

import config

# Build a regular expression from technique name variations
def build_regex(technique_variations):

    regex_str = "|".join(technique_variations) if len(technique_variations) > 1 else technique_variations[0]
    regex = re.compile(regex_str, re.IGNORECASE)

    return regex

# Find all instances of the specified regex
def count_instances(data, regex):

    count = 0
    for row in data:
        if regex.search(row[config.HEADER]):
            count += 1

    return count

# Parse the CSV log
def parse_log(log):

    data = []
    with open(log) as f:
        dr = csv.DictReader(f)
        for row in dr:
            data.append(row)

    return data

# The main entry point for the program
def main():

    data = parse_log(config.LOGFILE)
    for variation_list in config.TECHNIQUE_VARIATIONS:
        regex = build_regex(variation_list)
        name = variation_list[0]
        count = count_instances(data, regex)
        print(f"Instances of technique {name}: {count}")

if __name__ == "__main__":
    main()


# import torchvision.models as models
import torch
import torch.nn as nn
from transformers import BertModel, BertConfig

class TextEncoder(nn.Module):
    def __init__(self, pretrained=True):
        super(TextEncoder, self).__init__()
        if pretrained:  # if use pretrained scibert model
            self.main_model = BertModel.from_pretrained('bert_pretrained/')
        else:
            config = BertConfig(vocab_size=31090, )
            self.main_model = BertModel(config)

        self.dropout = nn.Dropout(0.1)
        # self.hidden_size = self.main_model.config.hidden_size

    def forward(self, input_ids, attention_mask):
        device = input_ids.device
        typ = torch.zeros(input_ids.shape).long().to(device)
        output = self.main_model(input_ids, token_type_ids=typ, attention_mask=attention_mask)['pooler_output']  # b,d
        logits = self.dropout(output)
        return logits

if __name__ == '__main__':
    model = TextEncoder()
    for name, param in model.named_parameters():
        print(name)


from datetime import date
import os

from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from preprocessors import ToDatetime
from transformers.gamestate import (
    ComputeDaysSinceStart,
    ComputeNetScore,
    EncodeOnBaseOccupancy,
    Identity,
    EncodeHandedness,
    EncodeInningTopBot,
)
from transformers.movingaverage import MovingAverage
from transformers.headtohead import HeadToHead
from transformers.parkfactor import ParkFactor
from constants import IDENTITY_COLS, ROOT_DIR

def make_data_pipeline(
    training_data_start_date: date,
    training_data_end_date: date,
    ma365_start_date: date,
    ma30_start_date: date
):

    preprocessors = Pipeline(steps=[
        ("game_date_to_datetime", ToDatetime(input_col="game_date")),
    ])

    feature_transformers = ColumnTransformer(
        [
            (
                "identity",
                Identity(),
                IDENTITY_COLS,
            ),
            (
                "encode_on_base_occupancy",
                EncodeOnBaseOccupancy(output_cols=["on_1b", "on_2b", "on_3b"]),
                ["on_1b", "on_2b", "on_3b", ],
            ),
            (
                "compute_net_score",
                ComputeNetScore(output_col="net_score"),
                ["bat_score", "fld_score", ],
            ),
            (
                "compute_days_since_start",
                ComputeDaysSinceStart(output_col="days_since_start"),
                ["game_date", ],
            ),
            (
                "encode_handedness",
                EncodeHandedness(output_cols=["batter_stand", "pitcher_throws"]),
                ["stand", "p_throws", ],
            ),
            (
                "encode_inning",
                EncodeInningTopBot(output_col="topbot"),
                ["inning_topbot", ],
            ),
            (
                "batter_365_days_ma",
                MovingAverage(
                    output_cols=["PA", "1B", "2B", "3B", "HR", "BB", "SO", "DP", "FO", "HBP", "SF", "SH", "wOBA", "mEV", "aEV", "pwOBA", "pPA"],
                    player_type="batter",
                    ma_days=365,
                    stats_to_compute=["PA", "1B", "2B", "3B", "HR", "BB", "SO", "DP", "FO", "HBP", "SF", "SH", "wOBA", "mEV", "aEV", "pwOBA", "pPA"],
                    training_data_start_date=training_data_start_date, #date(2008, 2, 27),
                    training_data_end_date=training_data_end_date, #date(2023, 5, 3),
                    ma_start_date=ma365_start_date, #date(2009, 2, 26), # datetime(2022, 4, 7) + timedelta(365)
                ),
                ["game_date", "batter", "launch_speed", "p_throws", ],
            ),
            (
                "pitcher_365_days_ma",
                MovingAverage(
                    output_cols=["PA", "1B", "2B", "3B", "HR", "BB", "SO", "DP", "FO", "HBP", "SF", "SH", "wOBA", "mEV", "aEV", "pwOBA", "pPA"],
                    player_type="pitcher",
                    ma_days=365,
                    stats_to_compute=["PA", "1B", "2B", "3B", "HR", "BB", "SO", "DP", "FO", "HBP", "SF", "SH", "wOBA", "mEV", "aEV", "pwOBA", "pPA"],
                    training_data_start_date=training_data_start_date, #date(2008, 2, 27),
                    training_data_end_date=training_data_end_date, #date(2023, 5, 3),
                    ma_start_date=ma365_start_date, #date(2009, 2, 26), # datetime(2022, 4, 7) + timedelta(365)
                ),
                ["game_date", "pitcher", "launch_speed", "stand", ],
            ),
            (
                "batter_30_days_ma",
                MovingAverage(
                    output_cols=["PA", "1B", "2B", "3B", "HR", "BB", "SO", "wOBA"],
                    player_type="batter",
                    ma_days=30,
                    stats_to_compute=["PA", "1B", "2B", "3B", "HR", "BB", "SO", "wOBA"],
                    training_data_start_date=training_data_start_date, #date(2008, 2, 27),
                    training_data_end_date=training_data_end_date, #date(2023, 5, 3),
                    ma_start_date=ma30_start_date, #date(2008, 3, 28), # datetime(2022, 4, 7) + timedelta(30)
                ),
                ["game_date", "batter", "launch_speed", "p_throws", ],
            ),
            # (
            #     "pitcher_30_days_ma",
            #     MovingAverage(
            #         output_cols=["PA", "1B", "2B", "3B", "HR", "BB", "SO", "wOBA"],
            #         player_type="pitcher",
            #         ma_days=30,
            #         stats_to_compute=["PA", "1B", "2B", "3B", "HR", "BB", "SO", "wOBA"],
            #         training_data_start_date=training_data_start_date, #date(2008, 2, 27),
            #         training_data_end_date=training_data_end_date, #date(2023, 5, 3),
            #         ma_start_date=ma30_start_date, #date(2008, 3, 28), # datetime(2022, 4, 7) + timedelta(30)
            #     ),
            #     ["game_date", "pitcher", "launch_speed", "stand", ],
            # ),
            (
                "head_to_head",
                HeadToHead(
                    output_cols=["1B", "2B", "HR", "BB", "SO", "PA", "wOBA"],
                    stats_to_compute=["1B", "2B", "HR", "BB", "SO", "PA", "wOBA"],
                ),
                ["batter", "pitcher", ],
            ),
            (
                "ball_park",
                ParkFactor(
                    output_cols=["1B", "2B", "3B", "HR", "BB"],
                    stats_to_compute=["1B", "2B", "3B", "HR", "BB"],
                    park_factors_file_path=os.path.join(ROOT_DIR, "intermediate/park_factors.json"),
                ),
                ["home_team", "game_year", ],
            ),
        ],
        remainder='drop',
    )

    data_pipeline = Pipeline(
        [
            ("preprocessors", preprocessors),
            ("feature_transformers", feature_transformers),
        ]
    )
    return data_pipeline


#!/usr/bin/env python

import rospy
from std_msgs.msg import Float64
import scipy.io as sio
import dynamic_reconfigure.client as reconfig

def gui_callback(config):
    global control_onoff
    global reference_path
    global step_time

    if config['enable_control']:
        control_onoff = True
    else:
        control_onoff = False

    reference_path = config['curve_path']
    step_time = config['step_time']

def get_curve():
    global control_onoff
    global reference_path
    global step_time

    rospy.init_node('curve_node', anonymous=True)

    # published topics
    pub = {}
    pub['calibration'] = rospy.Publisher('calibration_channel', Float64, queue_size=10)

    # communicate with the dynamic server
    dyn_params = reconfig.Client('server', config_callback=gui_callback)

    step_time = 60  # se eu alterar aqui eu preciso alterar no client_node step_t = 6
    length_reference = 12000
    new_freq = length_reference/step_time

    rate = rospy.Rate(new_freq)

    control_onoff = False
    count = 0
    loaded = 0
    number_steps = 0

    while not rospy.is_shutdown():

        if control_onoff:
            # load reference
            if loaded == 0:
                try:
                    reference_data = sio.loadmat(reference_path)
                    pw_ref = reference_data['pw_ref'][0]
                    loaded = 1
                except:
                    print("Ooops, the file you tried to load is not in the folder.")
                    pw_ref = [0, 0, 0, 0]

            if count > len(pw_ref) - 1:
                count = 0
                number_steps = number_steps + 1

            # send messages
            if number_steps == 0:
                pub['calibration'].publish(pw_ref[count])
            else:
                pub['calibration'].publish(0)

            # update counter
            count = count + 1

        else:
            count = 0
            number_steps = 0
            loaded = 0
            pub['calibration'].publish(0)

        rate.sleep()

if __name__ == '__main__':
    try:
        get_curve()
    except rospy.ROSInterruptException:
        pass


import numpy as np
import os
from PIL import Image
import pandas as pd
import argparse
import thread, time
from colorama import init
from termcolor import colored
import logging

values=[]
compared=[]

def compare(file1, path2):
  im = [None, None]
  im[0] = (np.array(Image.open(file1))).astype(np.int)
  for file2 in os.listdir(path2):
    compared.append(file1 + ' - ' + file2)
    try:
      im[1] = (np.array(Image.open(path2+'/'+file2))).astype(np.int)
      c=np.abs(im[0] - im[1]).sum()
      if c <= 500:
        values.append((file1.split("/")[-1],file2,c))
        logging.warning(colored("\tPositive (<treshold): "+ file1 + " " + file2 + "Treshold: "+str(c), "yellow", "on_blue"))
    except:
      logging.warning(colored("\tComparisson error: "+file1+ " vs "+ file2, "yellow", "on_blue"))

def thumbnails(path):
  for file1 in os.listdir(path):
    try:
      img=Image.open(path + "/" + file1).convert('L').resize((32,32), resample=Image.BICUBIC)
      img.save(path + "../thumbX/" + file1)
    except:
      pass

if __name__=='__main__':
  import sys
  init()
  logging.warning(colored('image comparer. Identify equal or similar images in two folders', "green"))
  parser = argparse.ArgumentParser(description='image comparer. Identify equal or similar images in two folders.')
  parser.add_argument('-p1','--path1', help='path for the first folder containing only images', required=True)
  parser.add_argument('-p2','--path2', help='path for the second folder containing only images', required=True)
  parser.add_argument('-c','--case', help='casename, for results filename', required=True)
  args = parser.parse_args()
  argsdict = vars(args)
  path1 = argsdict['path1']
  path2 = argsdict['path2']
  logging.warning(colored("Comparing images using resize(32x32), Gray Scale and verifying pizel by pixel. (Treshold < 500)","yellow"))
  logging.warning(colored("\tpaths: "+ path1 + " " + path2, "yellow"))
  values.append((path1,path2,"compare"))
  logging.warning("\tCreating thumbnails for images on paths ")

  for path in values[0][:2]:
    if not os.path.exists(path + "../thumbX/"):
      os.makedirs(path + "../thumbX/")
    logging.warning(colored("\t - Working on path "+path, "yellow"))
    thumbnails(path)

  logging.warning("\tComparing thumbails ...")
  num_files=len(os.listdir(path1 + "../thumbX/"))*len(os.listdir(path2 + "../thumbX/"))
  logging.warning(colored("\n - Total images to compare: "+ str(num_files) + "\n", "yellow"))
  for file1 in os.listdir(path1 + "../thumbX/"):
     thread.start_new_thread(compare, (path1+"../thumbX/"+file1, path2+"../thumbX/"))
  p=0
  while len(compared) != num_files:
    p=p+1
    if p%2==0:
      logging.warning("\tCompared files: "+str(len(compared)))
    time.sleep(5)

  with open(argsdict['case']+'.txt', 'w') as f:
      for img in values:
        f.write("%s\n" %str(img) )

  logging.warning(colored("Comparison finished.","green"))


import numpy as np
import pandas as pd
from PIL import Image
import cv2
from flask import Flask, jsonify, request
from sklearn.model_selection import train_test_split
X = np.load('image.npz')
X = X['arr_0']
y = pd.read_csv("labels.csv")
y = y["labels"]

classes = ["A","B" ,"C" ,"D" ,"E" ,"F" ,"G" ,"H" ,"I" ,"J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z"]
n_classes = len(classes)

X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=9, train_size=3500, test_size=500)

scaled_X_train = X_train/255.0
scaled_X_test = X_test/255.0

def get_alphabet(image):

    #opening
    #image
    im_pil = Image.open(image)
    alph_image = im_pil.convert('L')

    alph_image_resized = alph_image.resize((28,28), Image.ANTIALIAS)

    #setting minimum and max pixel
    pixel_filter = 20
    min_pixel = np.percentile(alph_image_resized, pixel_filter)
    alph_image_resized_inverted_scaled = np.clip(alph_image_resized-min_pixel, 0, 255)
    max_pixel = np.max(alph_image_resized)

    # inverting for proper recognition
    alph_image_resized_inverted_scaled = np.asarray(alph_image_resized_inverted_scaled)/max_pixel

    # converting 2 array and predicting the alphabet
    test_sample = np.array(alph_image_resized_inverted_scaled).reshape(1,660)
    test_pred = train_test_split.predict(test_sample)
    return test_pred[0]

def predict_data():
  image = cv2.imdecode(np.fromstring(request.files.get("alphabet").read(), np.uint8), cv2.IMREAD_UNCHANGED)
  image = request.files.get("alphabet")
  alphabet = get_alphabet(image)
  return jsonify({
    "alphabet_predicted": alphabet
  }), 200

predict_data()


import numpy as np
import networkx
from mealpy.utils.problem import Problem
import tsplib95
class TSPProblem(Problem):
    def __init__(self, lb, ub,minmax, distance_matrix,name="TSPProblem",**kwargs):
        self.distance_matrix = distance_matrix
        super().__init__(lb, ub, minmax,**kwargs)
        self.name = name

    def generate_position(self,lb=None, ub=None):
        return np.random.permutation(len(lb))

    def amend_position(self,solution,lb=None, ub=None):
        solution = np.clip(solution, lb,ub)
        # solution_int = solution.astype(np.int64)
        solution_set = set(list(range(0, len(solution))))
        solution_done = np.array([-1, ] * len(solution))
        solution_int = solution.astype(np.int32)
        city_unique, city_counts = np.unique(solution_int, return_counts=True)

        for idx, city in enumerate(solution_int):
            if solution_done[idx] != -1:
                continue
            if city in city_unique:
                solution_done[idx] = city
                city_unique = np.where(city_unique == city, -1, city_unique)
            else:
                list_cities_left = list(solution_set - set(city_unique) - set(solution_done))
                solution_done[idx] = list_cities_left[0]
        # print(f"Final solution: {solution_done}")
        # solution_done = np.clip(solution_done, lb,ub)
        return solution_done
        # return solution_int

    def fit_func(self, solution):
        solution = np.clip(solution, self.lb,self.ub)
        dist = self.distance_matrix[solution]
        sum =0.
        for i in range(len(dist)-1):
            sum += dist[i,i+1]
        sum += dist[-1,0]
        return sum

def returnMatrix(problem):
    if problem.edge_weight_format == "FULL_MATRIX":
        g = problem.get_graph()
        distance = networkx.to_numpy_matrix(g)
    else:
        ## euclidian distance
        g = problem.get_graph()
        distance = networkx.to_numpy_matrix(g)
        # for i in range(len(distance)):
        #     for j in range(len(distance)):
        #         if i != j:
        #             distance[i,j] = np.linalg.norm(np.array(g.nodes[i]['coord']) - np.array(g.nodes[j]['coord']))
    return distance

def returnProblem(path, type):
    if type == 'atsp':
        problem = tsplib95.load(path)
        # distance = returnMatrix(problem)
        return problem
    elif type == 'tsp':
        problem = tsplib95.load(path,special=tsplib95.distances.euclidean)
        # distance = returnMatrix(problem)
        return problem