from __future__ import absolute_import
from __future__ import unicode_literals
from datetime import datetime

import pytz

from corehq.apps.domain.models import Domain
from corehq.apps.locations.models import SQLLocation
from corehq.apps.locations.permissions import user_can_edit_location, user_can_access_location_id, \
    user_can_view_location
from corehq.apps.reports.filters.base import BaseSingleOptionFilter
from corehq.apps.reports.filters.fixtures import AsyncLocationFilter
from corehq.apps.reports.filters.select import MonthFilter, YearFilter
from custom.common.filters import RestrictedAsyncLocationFilter
from memoized import memoized

def location_hierarchy_config(domain, location_types=None):
    location_types = location_types or ['state', 'district', 'block']
    return [
        (loc_type.name, [loc_type.parent_type.name if loc_type.parent_type else None])
        for loc_type in Domain.get_by_name(
            domain
        ).location_types if loc_type.code in location_types
    ]

# copy/paste from corehq.apps.location.utils
# added possibility to exclude test locations, test flag is custom added to the metadata in location object
def load_locs_json(domain, selected_loc_id=None, include_archived=False,
        user=None, only_administrative=False, show_test=False):
    """initialize a json location tree for drill-down controls on
    the client. tree is only partially initialized and branches
    will be filled in on the client via ajax.

    what is initialized:
    * all top level locs
    * if a 'selected' loc is provided, that loc and its complete
      ancestry

    only_administrative - if False get all locations
                          if True get only administrative locations
    """

    def loc_to_json(loc, project):
        ret = {
            'name': loc.name,
            'location_type': loc.location_type.name,  # todo: remove when types aren't optional
            'uuid': loc.location_id,
            'is_archived': loc.is_archived,
            'can_edit': True
        }
        if user:
            if user.has_permission(domain, 'access_all_locations'):
                ret['can_edit'] = user_can_edit_location(user, loc, project)
            else:
                ret['can_edit'] = user_can_access_location_id(domain, user, loc.location_id)
        return ret

    project = Domain.get_by_name(domain)

    locations = SQLLocation.root_locations(
        domain, include_archive_ancestors=include_archived
    )

    if only_administrative:
        locations = locations.filter(location_type__administrative=True)

    if not show_test:
        locations = [
            loc for loc in locations if loc.metadata.get('is_test_location', 'real') != 'test'
        ]

    loc_json = [
        loc_to_json(loc, project) for loc in locations
        if user is None or user_can_view_location(user, loc, project)
    ]

    # if a location is selected, we need to pre-populate its location hierarchy
    # so that the data is available client-side to pre-populate the drop-downs
    if selected_loc_id:
        selected = SQLLocation.objects.get(
            domain=domain,
            location_id=selected_loc_id
        )

        lineage = selected.get_ancestors()

        parent = {'children': loc_json}
        for loc in lineage:
            children = loc.child_locations(include_archive_ancestors=include_archived)
            if only_administrative:
                children = children.filter(location_type__administrative=True)

            # find existing entry in the json tree that corresponds to this loc
            try:
                this_loc = [k for k in parent['children'] if k['uuid'] == loc.location_id][0]
            except IndexError:
                # if we couldn't find this location the view just break out of the loop.
                # there are some instances in viewing archived locations where we don't actually
                # support drilling all the way down.
                break
            this_loc['children'] = [
                loc_to_json(loc, project) for loc in children
                if user is None or user_can_view_location(user, loc, project)
            ]
            parent = this_loc

    return loc_json

class ICDSTableauFilterMixin(object):
    def __init__(self, request, domain=None, timezone=pytz.utc, parent_report=None,
                 css_label=None, css_field=None):
        css_label = 'col-xs-4 col-md-4 col-lg-4 control-label'
        css_field = 'col-xs-8 col-md-8 col-lg-8'
        super(ICDSTableauFilterMixin, self).__init__(
            request, domain, timezone, parent_report, css_label, css_field
        )

class ICDSMonthFilter(ICDSTableauFilterMixin, MonthFilter):

    @property
    @memoized
    def selected(self):
        return self.get_value(self.request, self.domain) or "%02d" % datetime.now().month

class ICDSYearFilter(ICDSTableauFilterMixin, YearFilter):
    pass

class IcdsLocationFilter(AsyncLocationFilter):

    def load_locations_json(self, loc_id):
        show_test = self.request.GET.get('include_test', False)
        return load_locs_json(self.domain, loc_id, user=self.request.couch_user, show_test=show_test)

class IcdsRestrictedLocationFilter(AsyncLocationFilter):

    @property
    def location_hierarchy_config(self):
        return location_hierarchy_config(self.domain)

class TableauLocationFilter(ICDSTableauFilterMixin, RestrictedAsyncLocationFilter):

    auto_drill = False

    @property
    def location_hierarchy_config(self):
        return location_hierarchy_config(
            self.domain,
            location_types=['state', 'district', 'block', 'supervisor', 'awc']
        )

class CasteFilter(ICDSTableauFilterMixin, BaseSingleOptionFilter):
    slug = 'caste'
    label = 'Caste'

    @property
    @memoized
    def selected(self):
        return super(CasteFilter, self).selected or self.options[0][0]

    @property
    def options(self):
        return [
            ('All', 'All'),
            ('ST', 'ST'),
            ('SC', 'SC'),
            ('OBC', 'OBC'),
            ('Others', 'Others'),
        ]

class MinorityFilter(ICDSTableauFilterMixin, BaseSingleOptionFilter):
    slug = 'minority'
    label = 'Minority'

    @property
    @memoized
    def selected(self):
        return super(MinorityFilter, self).selected or self.options[0][0]

    @property
    def options(self):
        return [
            ('All', 'All'),
            ('Yes', 'Yes'),
            ('No', 'No')
        ]

class DisabledFilter(ICDSTableauFilterMixin, BaseSingleOptionFilter):
    slug = 'disabled'
    label = 'Disabled'

    @property
    @memoized
    def selected(self):
        return super(DisabledFilter, self).selected or self.options[0][0]

    @property
    def options(self):
        return [
            ('All', 'All'),
            ('Yes', 'Yes'),
            ('No', 'No')
        ]

class ResidentFilter(ICDSTableauFilterMixin, BaseSingleOptionFilter):
    slug = 'resident'
    label = 'Resident'

    @property
    @memoized
    def selected(self):
        return super(ResidentFilter, self).selected or self.options[0][0]

    @property
    def options(self):
        return [
            ('All', 'All'),
            ('Permanent', 'Permanent'),
            ('Resident', 'Resident')
        ]

class MaternalStatusFilter(ICDSTableauFilterMixin, BaseSingleOptionFilter):
    slug = 'ccs_status'
    label = 'Maternal status'

    @property
    @memoized
    def selected(self):
        return super(MaternalStatusFilter, self).selected or self.options[0][0]

    @property
    def options(self):
        return [
            ('All', 'All'),
            ('Pregnant', 'Pregnant'),
            ('Lactating', 'Lactating')
        ]

class ChildAgeFilter(ICDSTableauFilterMixin, BaseSingleOptionFilter):
    slug = 'child_age_tranche'
    label = 'Child age'

    @property
    @memoized
    def selected(self):
        return super(ChildAgeFilter, self).selected or self.options[0][0]

    @property
    def options(self):
        return [
            ('All', 'All'),
            ('0-28 days', '0-28 days'),
            ('28 days - 6mo', '28 days - 6mo'),
            ('1 yr', '1 yr'),
            ('2 yr', '2 yr'),
            ('3 yr', '3 yr'),
            ('4 yr', '4 yr'),
            ('5 yr', '5 yr'),
            ('6 yr', '6 yr'),
        ]

class THRBeneficiaryType(ICDSTableauFilterMixin, BaseSingleOptionFilter):
    slug = 'thr_beneficiary_type'
    label = 'THR Beneficiary Type'

    @property
    @memoized
    def selected(self):
        return super(THRBeneficiaryType, self).selected or self.options[0][0]

    @property
    def options(self):
        return [
            ('All', 'All'),
            ('Child', 'Child'),
            ('Pregnant', 'Pregnant'),
            ('Lactating', 'Lactating'),
        ]


"""Define word-based embedders."""

import abc
import collections
import os

import numpy as np
import tensorflow as tf

from google.protobuf import text_format
from tensorboard.plugins import projector

from opennmt import constants, tokenizers
from opennmt.data import dataset as dataset_util
from opennmt.data import text, vocab
from opennmt.inputters.inputter import Inputter
from opennmt.layers import common
from opennmt.utils import misc

def save_embeddings_metadata(
    log_dir, variable_name, vocabulary_file, num_oov_buckets=1
):
    """Registers an embedding variable for visualization in TensorBoard.

    This function registers :obj:`variable_name` in the ``projector_config.pbtxt``
    file and generates metadata from :obj:`vocabulary_file` to attach a label
    to each word ID.

    Args:
      log_dir: The active log directory.
      variable_name: The variable name in the checkpoint.
      vocabulary_file: The associated vocabulary file.
      num_oov_buckets: The number of additional unknown tokens.
    """
    # Assume it ends with /.ATTRIBUTES/VALUE
    filename = "%s.txt" % "_".join(variable_name.split("/")[:-2])
    metadata_path = os.path.join(log_dir, filename)

    with tf.io.gfile.GFile(vocabulary_file) as src, tf.io.gfile.GFile(
        metadata_path, mode="w"
    ) as dst:
        ws_index = 0
        for line in src:
            # The TensorBoard code checks line.trim().length == 0 when loading the
            # metadata file so make sure lines are not dropped.
            if not line.replace("\uFEFF", "").strip():
                dst.write("<whitespace%d>\n" % ws_index)
                ws_index += 1
            else:
                dst.write(line)
        if num_oov_buckets == 1:
            dst.write("<unk>\n")
        else:
            for i in range(num_oov_buckets):
                dst.write("<unk%d>\n" % i)

    config = projector.ProjectorConfig()

    # If the projector file exists, load it.
    config_path = os.path.join(log_dir, "projector_config.pbtxt")
    if tf.io.gfile.exists(config_path):
        with tf.io.gfile.GFile(config_path) as config_file:
            text_format.Merge(config_file.read(), config)

    # If this embedding is already registered, just update the metadata path.
    exists = False
    for meta in config.embeddings:
        if meta.tensor_name == variable_name:
            meta.metadata_path = filename
            exists = True
            break

    if not exists:
        embedding = config.embeddings.add()
        embedding.tensor_name = variable_name
        embedding.metadata_path = filename

    with tf.io.gfile.GFile(config_path, "w") as config_file:
        config_file.write(text_format.MessageToString(config))

def load_pretrained_embeddings(
    embedding_file,
    vocabulary_file,
    num_oov_buckets=1,
    with_header=True,
    case_insensitive_embeddings=True,
):
    """Returns pretrained embeddings relative to the vocabulary.

    The :obj:`embedding_file` must have the following format:

    .. code-block:: text

        N M
        word1 val1 val2 ... valM
        word2 val1 val2 ... valM
        ...
        wordN val1 val2 ... valM

    or if :obj:`with_header` is ``False``:

    .. code-block:: text

        word1 val1 val2 ... valM
        word2 val1 val2 ... valM
        ...
        wordN val1 val2 ... valM

    This function will iterate on each embedding in :obj:`embedding_file` and
    assign the pretrained vector to the associated word in :obj:`vocabulary_file`
    if found. Otherwise, the embedding is ignored.

    If :obj:`case_insensitive_embeddings` is ``True``, word embeddings are assumed
    to be trained on lowercase data. In that case, word alignments are case
    insensitive meaning the pretrained word embedding for "the" will be assigned
    to "the", "The", "THE", or any other case variants included in
    :obj:`vocabulary_file`.

    Args:
      embedding_file: Path the embedding file. Entries will be matched against
        :obj:`vocabulary_file`.
      vocabulary_file: The vocabulary file containing one word per line.
      num_oov_buckets: The number of additional unknown tokens.
      with_header: ``True`` if the embedding file starts with a header line like
        in GloVe embedding files.
      case_insensitive_embeddings: ``True`` if embeddings are trained on lowercase
        data.

    Returns:
      A Numpy array of shape ``[vocabulary_size + num_oov_buckets, embedding_size]``.
    """
    # Map words to ids from the vocabulary.
    word_to_id = collections.defaultdict(list)
    with tf.io.gfile.GFile(vocabulary_file) as vocabulary:
        count = 0
        for word in vocabulary:
            word = word.strip()
            if case_insensitive_embeddings:
                word = word.lower()
            word_to_id[word].append(count)
            count += 1

    # Fill pretrained embedding matrix.
    with tf.io.gfile.GFile(embedding_file) as embedding:
        pretrained = None

        if with_header:
            next(embedding)

        for line in embedding:
            fields = line.strip().split()
            word = fields[0]

            if pretrained is None:
                pretrained = np.random.normal(
                    size=(count + num_oov_buckets, len(fields) - 1)
                )

            # Lookup word in the vocabulary.
            if word in word_to_id:
                ids = word_to_id[word]
                for index in ids:
                    pretrained[index] = np.asarray(fields[1:])

    return pretrained

def add_sequence_controls(
    ids,
    length=None,
    start_id=None,
    end_id=None,
    pad_id=constants.PADDING_ID,
):
    """Adds sequence control tokens.

    Args:
      ids: Sequence of token ids as a 1D ``tf.Tensor``, a 2D ``tf.Tensor``, or a
        2D ``tf.RaggedTensor``.
      length: Length of each sequence in the batch as a 1D ``tf.Tensor``.
        Required if :obj:`ids` is a 2D ``tf.Tensor``.
      start_id: Token id to prepend to the sequences (set ``None`` to disable).
      end_id: Token id to append to the sequences (set ``None`` to disable).
      pad_id: Padding token ID.

    Returns:
      The updated token ids, and optionally the updated length if it was initially
      passed as argument.

    Raises:
      ValueError: if :obj:`ids` is a dense 2D ``tf.Tensor``, but :obj:`length` is not set.
      ValueError: if the rank of :obj:`ids` is not 1 or 2.
    """
    input_is_ragged = isinstance(ids, tf.RaggedTensor)
    if start_id is None and end_id is None:
        return ids if input_is_ragged or length is None else (ids, length)

    batch_size = None
    rank = ids.shape.rank
    if rank == 2:
        if not input_is_ragged:
            if length is None:
                raise ValueError(
                    "length argument is required when passing a dense batch of sequences"
                )
            ids = tf.RaggedTensor.from_tensor(ids, lengths=length)
        batch_size = ids.nrows()
    elif rank != 1:
        raise ValueError("Unsupported rank %d (expected 1 or 2)" % rank)

    def _make_column(value):
        value = tf.constant(value, dtype=ids.dtype)
        if batch_size is not None:
            value = tf.fill([batch_size], value)
        return tf.expand_dims(value, -1)

    concat_inputs = []
    if start_id is not None:
        concat_inputs.append(_make_column(start_id))
    concat_inputs.append(ids)
    if end_id is not None:
        concat_inputs.append(_make_column(end_id))

    ids = tf.concat(concat_inputs, axis=-1)

    if not input_is_ragged:
        if batch_size is not None:
            ids = ids.to_tensor(default_value=pad_id)
        if length is not None:
            length += len(concat_inputs) - 1
            return ids, length

    return ids

class TextInputter(Inputter):
    """An abstract inputter that processes text."""

    def __init__(self, num_oov_buckets=1, **kwargs):
        super().__init__(**kwargs)
        self.num_oov_buckets = num_oov_buckets
        self.noiser = None
        self.in_place_noise = False
        self.noise_probability = None
        self.vocabulary_file = None
        self.vocabulary_size = None
        self.tokens_to_ids = None
        self.ids_to_tokens = None
        self.tokenizer = None

    def initialize(self, data_config):
        self.vocabulary_file = data_config["vocabulary"]
        (
            self.vocabulary_size,
            self.tokens_to_ids,
            self.ids_to_tokens,
        ) = vocab.create_lookup_tables(
            self.vocabulary_file,
            self.num_oov_buckets,
            as_asset=data_config.get("export_vocabulary_assets", True),
        )
        tokenizer_config = data_config.get("tokenization")
        self.tokenizer = tokenizers.make_tokenizer(tokenizer_config)

    def set_noise(self, noiser, in_place=True, probability=None):
        """Enables noise to be applied to the input features.

        Args:
          noiser: A :class:`opennmt.data.WordNoiser` instance.
          in_place: If ``False``, the noisy version of the input will be stored as
            a separate feature prefixed with ``noisy_``.
          probability: When :obj:`in_place` is enabled, the probability to apply the
            noise.

        Raises:
          ValueError: if :obj:`in_place` is enabled but a :obj:`probability` is not
            set.
        """
        if in_place and probability is None:
            raise ValueError("In-place noise requires a probability")
        self.noiser = noiser
        self.in_place_noise = in_place
        self.noise_probability = probability

    def export_assets(self, asset_dir):
        self._assert_is_initialized()
        return self.tokenizer.export_assets(asset_dir, asset_prefix=self.asset_prefix)

    def make_dataset(self, data_file, training=None):
        return dataset_util.make_datasets(tf.data.TextLineDataset, data_file)

    def get_dataset_size(self, data_file):
        if isinstance(data_file, list):
            return list(map(misc.count_lines, data_file))
        return misc.count_lines(data_file)

    def has_prepare_step(self):
        # For performance reasons, we apply external tokenizers on a batch of
        # dataset elements during the preparation step.
        return not self.tokenizer.in_graph and not isinstance(
            self.tokenizer, tokenizers.SpaceTokenizer
        )

    def prepare_elements(self, elements, training=None):
        return {"tokens": self.tokenizer.tokenize(elements, training=training)}

    def make_features(self, element=None, features=None, training=None):
        """Tokenizes raw text."""
        self._assert_is_initialized()
        if features is None:
            features = {}
        if "tokens" in features:
            return features

        element = features.pop("text", element)
        if isinstance(element, dict):
            tokens = element["tokens"]
        else:
            element = tf.convert_to_tensor(element, dtype=tf.string)
            tokens = self.tokenizer.tokenize(element, training=training)

        if isinstance(tokens, tf.RaggedTensor):
            length = tokens.row_lengths()
            tokens = tokens.to_tensor(default_value=constants.PADDING_TOKEN)
        else:
            length = tf.shape(tokens)[0]
        if training and self.noiser is not None:
            noisy_tokens, noisy_length = self.noiser(
                tokens,
                sequence_length=length,
                probability=self.noise_probability if self.in_place_noise else None,
            )
            if self.in_place_noise:
                tokens = noisy_tokens
                length = noisy_length
            else:
                # Call make_features again to fill the remaining noisy features.
                noisy_features = dict(tokens=noisy_tokens, length=noisy_length)
                noisy_features = self.make_features(
                    features=noisy_features, training=training
                )
                for key, value in noisy_features.items():
                    features["noisy_%s" % key] = value
        features["length"] = length
        features["tokens"] = tokens
        return features

    def input_signature(self):
        self._assert_is_initialized()
        if self.tokenizer.in_graph:
            return {"text": tf.TensorSpec([None], tf.string)}
        else:
            return {
                "tokens": tf.TensorSpec([None, None], tf.string),
                "length": tf.TensorSpec([None], tf.int32),
            }

    def _assert_is_initialized(self):
        if self.tokenizer is None:
            raise RuntimeError(
                "The input layer is not initialized. You should initialize "
                "the model by calling model.initialize(data_config)."
            )

def _format_bool(b):
    return "yes" if b else "no"

class WordEmbedder(TextInputter):
    """Simple word embedder."""

    def __init__(self, embedding_size=None, dropout=0.0, **kwargs):
        """Initializes the parameters of the word embedder.

        Args:
          embedding_size: The size of the resulting embedding.
            If ``None``, an embedding file must be provided.
          dropout: The probability to drop units in the embedding.
          **kwargs: Additional layer keyword arguments.
        """
        super().__init__(**kwargs)
        self.embedding_size = embedding_size
        self.embedding_file = None
        self.dropout = dropout
        self.decoder_mode = False
        self.mark_start = None
        self.mark_end = None

    def set_decoder_mode(self, enable=True, mark_start=None, mark_end=None):
        """Make this inputter produce sequences for a decoder.

        In this mode, the returned "ids_out" feature is the decoder output sequence
        and "ids" is the decoder input sequence.

        Args:
          enable: Enable the decoder mode.
          mark_start: Mark the sequence start. If ``None``, keep the current value.
          mark_end: Mark the sequence end. If ``None``, keep the current value.
        """
        self.decoder_mode = enable
        if mark_start is not None:
            self.mark_start = mark_start
        if mark_end is not None:
            self.mark_end = mark_end

    def get_length(self, features, ignore_special_tokens=False):
        length = features["length"]
        if ignore_special_tokens:
            # Decoder mode shifts the sequences by one timesteps.
            num_special_tokens = -1 if self.decoder_mode else 0
            if self.mark_start:
                num_special_tokens += 1
            if self.mark_end:
                num_special_tokens += 1
            length -= num_special_tokens
        return length

    def get_oov_tokens(self, features):
        tokens, ids = features["tokens"], features["ids"]
        if self.mark_start:
            ids = ids[1:] if ids.shape.rank == 1 else ids[:, 1:]
        oov_tokens_pos = tf.where(ids >= (self.vocabulary_size - self.num_oov_buckets))
        return tf.gather_nd(tokens, oov_tokens_pos)

    def initialize(self, data_config):
        super().initialize(data_config)
        embedding = data_config.get("embedding")
        if embedding is None and self.embedding_size is None:
            raise ValueError("embedding_size must be set")
        if embedding is not None:
            self.embedding_file = embedding["path"]
            self.trainable = embedding.get("trainable", True)
            self.embedding_file_with_header = embedding.get("with_header", True)
            self.case_insensitive_embeddings = embedding.get("case_insensitive", True)
        sequence_controls = data_config.get("sequence_controls")
        if sequence_controls:
            self.mark_start = sequence_controls.get("start")
            self.mark_end = sequence_controls.get("end")

        tf.get_logger().info(
            "Initialized %s input layer:", self.asset_prefix.strip("_")
        )
        tf.get_logger().info(" - vocabulary size: %d", self.vocabulary_size)
        tf.get_logger().info(
            " - special tokens: BOS=%s, EOS=%s",
            _format_bool(self.mark_start),
            _format_bool(self.mark_end),
        )

    def make_features(self, element=None, features=None, training=None):
        """Converts words tokens to ids."""
        features = super().make_features(
            element=element, features=features, training=training
        )
        if "ids" not in features:
            features["ids"] = self.tokens_to_ids.lookup(features["tokens"])
            if self.mark_start or self.mark_end:
                features["ids"], features["length"] = add_sequence_controls(
                    features["ids"],
                    features["length"],
                    start_id=constants.START_OF_SENTENCE_ID
                    if self.mark_start
                    else None,
                    end_id=constants.END_OF_SENTENCE_ID if self.mark_end else None,
                )
        if self.decoder_mode:
            ids = features["ids"]
            length = features["length"]
            if ids.shape.rank == 2:
                mask = tf.sequence_mask(
                    length - 1, maxlen=tf.shape(ids)[1], dtype=ids.dtype
                )
                features["ids"] = (ids * mask)[:, :-1]
                features["ids_out"] = ids[:, 1:]
            else:
                features["ids"] = ids[:-1]
                features["ids_out"] = ids[1:]
            features["length"] -= 1
        return features

    def build(self, input_shape):
        if self.embedding_file:
            pretrained = load_pretrained_embeddings(
                self.embedding_file,
                self.vocabulary_file,
                num_oov_buckets=self.num_oov_buckets,
                with_header=self.embedding_file_with_header,
                case_insensitive_embeddings=self.case_insensitive_embeddings,
            )
            self.embedding_size = pretrained.shape[-1]
            initializer = tf.constant_initializer(value=pretrained.astype(self.dtype))
        else:
            initializer = None
        self.embedding = self.add_weight(
            "embedding",
            [self.vocabulary_size, self.embedding_size],
            initializer=initializer,
            trainable=self.trainable,
        )
        super().build(input_shape)

    def call(self, features, training=None):
        ids = features["ids"]

        if getattr(self, "_tflite_mode", False):
            # Workaround for TensorFlow issue #42410.
            ids = tf.expand_dims(ids, axis=-1)
            return tf.gather_nd(self.embedding, ids)

        outputs = tf.nn.embedding_lookup(self.embedding, ids)
        outputs = common.dropout(outputs, self.dropout, training=training)
        return outputs

    def visualize(self, model_root, log_dir):
        save_embeddings_metadata(
            log_dir,
            misc.get_variable_name(self.embedding, model_root),
            self.vocabulary_file,
            num_oov_buckets=self.num_oov_buckets,
        )

    def map_v1_weights(self, weights):
        return [(self.embedding, weights["w_embs"])]

class CharEmbedder(TextInputter):
    """Base class for character-aware inputters."""

    def __init__(self, embedding_size, dropout=0.0, **kwargs):
        """Initializes the parameters of the character embedder.

        Args:
          embedding_size: The size of the character embedding.
          dropout: The probability to drop units in the embedding.
          **kwargs: Additional layer keyword arguments.
        """
        super().__init__(**kwargs)
        self.embedding_size = embedding_size
        self.embedding = None
        self.dropout = dropout

    def make_features(self, element=None, features=None, training=None):
        """Converts words to characters."""
        if features is None:
            features = {}
        if "char_ids" in features:
            return features
        if "chars" in features:
            chars = features["chars"]
        else:
            features = super().make_features(
                element=element, features=features, training=training
            )
            chars = text.tokens_to_chars(features["tokens"])
            chars = chars.to_tensor(default_value=constants.PADDING_TOKEN)
        features["char_ids"] = self.tokens_to_ids.lookup(chars)
        return features

    def build(self, input_shape):
        self.embedding = self.add_weight(
            "char_embedding", [self.vocabulary_size, self.embedding_size]
        )
        super().build(input_shape)

    @abc.abstractmethod
    def call(self, features, training=None):
        raise NotImplementedError()

    def visualize(self, model_root, log_dir):
        save_embeddings_metadata(
            log_dir,
            misc.get_variable_name(self.embedding, model_root),
            self.vocabulary_file,
            num_oov_buckets=self.num_oov_buckets,
        )

    def _embed(self, inputs, training):
        mask = tf.math.not_equal(inputs, 0)
        outputs = tf.nn.embedding_lookup(self.embedding, inputs)
        outputs = common.dropout(outputs, self.dropout, training=training)
        return outputs, mask

class CharConvEmbedder(CharEmbedder):
    """An inputter that applies a convolution on characters embeddings."""

    def __init__(
        self,
        embedding_size,
        num_outputs,
        kernel_size=5,
        stride=3,
        dropout=0.0,
        **kwargs
    ):
        """Initializes the parameters of the character convolution embedder.

        Args:
          embedding_size: The size of the character embedding.
          num_outputs: The dimension of the convolution output space.
          kernel_size: Length of the convolution window.
          stride: Length of the convolution stride.
          dropout: The probability to drop units in the embedding.
          **kwargs: Additional layer keyword arguments.
        """
        super().__init__(embedding_size, dropout=dropout, **kwargs)
        self.output_size = num_outputs
        self.conv = tf.keras.layers.Conv1D(
            num_outputs, kernel_size, strides=stride, padding="same"
        )

    def call(self, features, training=None):
        inputs = features["char_ids"]
        flat_inputs = tf.reshape(inputs, [-1, tf.shape(inputs)[-1]])
        outputs, _ = self._embed(flat_inputs, training)
        outputs = self.conv(outputs)
        outputs = tf.reduce_max(outputs, axis=1)
        outputs = tf.reshape(outputs, [-1, tf.shape(inputs)[1], self.output_size])
        return outputs

class CharRNNEmbedder(CharEmbedder):
    """An inputter that runs a single RNN layer over character embeddings."""

    def __init__(
        self, embedding_size, num_units, dropout=0.2, cell_class=None, **kwargs
    ):
        """Initializes the parameters of the character RNN embedder.

        Args:
          embedding_size: The size of the character embedding.
          num_units: The number of units in the RNN layer.
          dropout: The probability to drop units in the embedding and the RNN
            outputs.
          cell_class: The inner cell class or a callable taking :obj:`num_units` as
            argument and returning a cell. Defaults to a LSTM cell.
          **kwargs: Additional layer keyword arguments.

        Raises:
          ValueError: if :obj:`encoding` is invalid.
        """
        super().__init__(embedding_size, dropout=dropout, **kwargs)
        if cell_class is None:
            cell_class = tf.keras.layers.LSTMCell
        self.rnn = tf.keras.layers.RNN(cell_class(num_units))
        self.num_units = num_units

    def call(self, features, training=None):
        inputs = features["char_ids"]
        flat_inputs = tf.reshape(inputs, [-1, tf.shape(inputs)[-1]])
        embeddings, mask = self._embed(flat_inputs, training)
        outputs = self.rnn(embeddings, mask=mask, training=training)
        outputs = tf.reshape(outputs, [-1, tf.shape(inputs)[1], self.num_units])
        return outputs


import os.path

import numpy as np
import torch
import cv2
from torch import tensor
import librosa
import librosa.display
import random
import shutil

def get_img_tensor_from_sequence(parent: str) -> tensor:
    """
    函数用途：将image sequence转为对应的tensor
    参数举例：
    parent:"../data/original_data/train/061_foam_brick/1"
    返回值:60*440*440的tensor(取20张序列中间的图片用于后续的训练)
    """

    img_list = []
    for jpg_name in os.listdir(os.path.join(parent, "rgb")):
        png_name = jpg_name.split(".")[0] + ".png"

        rgb_file_path = os.path.join(parent, "rgb", jpg_name)
        mask_file_path = os.path.join(parent, "mask", png_name)

        rgb = cv2.imread(rgb_file_path)
        rgb = rgb[20:rgb.shape[0] - 20, 100:rgb.shape[1] - 100, :]
        mask = cv2.imread(mask_file_path, cv2.IMREAD_GRAYSCALE)

        img = cv2.bitwise_and(src1=rgb, src2=rgb, mask=mask)

        img = img.transpose(2, 0, 1)

        img_list.append(img)

    if len(img_list) < 20:
        img_list += [img_list[-1]] * (20 - len(img_list))
    elif len(img_list) > 20:
        reduced_len = len(img_list) - 20
        img_list = img_list[reduced_len // 2:]
        img_list = img_list[:20]

    assert len(img_list) == 20

    result = np.concatenate(img_list, axis=0)

    result = (result - 4.04) / 16.63

    return torch.tensor(result, dtype=torch.float)

def audio_to_features(audio_src_path: str):
    """
    函数用途：将audio数据转为梅尔频谱，返回对应的tensor
    """

    data = np.load(audio_src_path, allow_pickle=True)

    mel_array = np.zeros((4, 64, 256))
    mfcc_array = np.zeros((4, 64, 256))
    for index in range(4):
        mel_spec = librosa.feature.melspectrogram(y=data['audio'][:, index], sr=data['audio_samplerate'], n_mels=64)
        mel_spec_db = librosa.amplitude_to_db(mel_spec, ref=np.max)

        mfcc = librosa.feature.mfcc(y=data['audio'][:, index], sr=data['audio_samplerate'], n_mfcc=64)

        mel_array[index, :, :] = cv2.resize(mel_spec_db, (256, 64))
        mfcc_array[index, :, :] = cv2.resize(mfcc, (256, 64))

    mel_array = (mel_array - (-77.36)) / 8.26
    mfcc_array = (mfcc_array - (-5.15)) / 73.61

    feature = np.concatenate((mel_array, mfcc_array), axis=0)

    return torch.tensor(feature, dtype=torch.float)

def preprocess():
    """
    函数用途：将image sequence和audio转为tensor存储起来
    """

    parent = "../data/original_data/train"
    dst = "../data/processed_data"

    if not os.path.exists(dst):
        os.mkdir(dst)

    for class_dir in os.listdir(parent):
        if not os.path.exists(os.path.join(dst, class_dir)):
            os.mkdir(os.path.join(dst, class_dir))

        class_dir_path = os.path.join(parent, class_dir)

        for index in os.listdir(class_dir_path):
            index_path = os.path.join(class_dir_path, index)

            print(index_path)

            sequence_img_tensor = get_img_tensor_from_sequence(index_path)
            torch.save(obj=sequence_img_tensor, f=os.path.join(dst, class_dir, "image_" + str(index) + ".pth"))

            sequence_audio_tensor = audio_to_features(os.path.join(index_path, "audio_data.pkl"))
            torch.save(obj=sequence_audio_tensor, f=os.path.join(dst, class_dir, "audio_" + str(index) + ".pth"))

def generate_audio_or_image_data(choice="audio"):
    """
    函数用途：为 audio classify 或者 image classify 产生需要的数据
    """

    parent = "../data/processed_data"
    dst = os.path.join("../data", choice)

    if not os.path.exists(dst):
        os.mkdir(dst)

    for item in ["train", "val", "test"]:
        if not os.path.exists(os.path.join(dst, item)):
            os.mkdir(os.path.join(dst, item))

    for class_dir in os.listdir(parent):

        for item in ["train", "val", "test"]:
            if not os.path.exists(os.path.join(dst, item, class_dir)):
                os.mkdir(os.path.join(dst, item, class_dir))

        class_path = os.path.join(parent, class_dir)

        data_path_list = [os.path.join(class_path, item) for item in os.listdir(class_path) if choice in item]

        for data_path in data_path_list:
            p = random.random()

            if 0 <= p < 0.8:
                shutil.copy(src=data_path, dst=os.path.join(dst, "train", class_dir))
                print(os.path.join(dst, "train", class_dir))
            elif 0.8 <= p < 0.9:
                shutil.copy(src=data_path, dst=os.path.join(dst, "val", class_dir))
                print(os.path.join(dst, "val", class_dir))
            else:
                shutil.copy(src=data_path, dst=os.path.join(dst, "test", class_dir))
                print(os.path.join(dst, "test", class_dir))

if __name__ == "__main__":
    preprocess()
    generate_audio_or_image_data(choice="audio")
    generate_audio_or_image_data(choice="image")


#import PyQt5

import matplotlib
matplotlib.use('Qt5Agg')
import matplotlib.pyplot as plt
import cv2

# from detectron2.data import MetadataCatalog
# from detectron2.structures import Instances
# from detectron2.utils.visualizer import Visualizer, VisImage
# import numpy as np

# Salmon outline detector
#
# This program detects the outline of a swimming salmon from experiments as part of the Master's thesis of
# Eirik Ruben Grimholt Søvik at Marine Technological Center, NTNU, Trondheim.
#
# Relevant variables:
#       - OUTPUT_FOLDER
#       -

#import argparse

import numpy as np
#import re

#from detectron2.config import get_cfg, CfgNode
from detectron2.data import MetadataCatalog
#from detectron2.engine import DefaultPredictor
from detectron2.structures import Instances
from detectron2.utils.visualizer import  VisImage #,Visualizer

import math
#from detectron2.engine import DefaultTrainer
#from detectron2 import model_zoo
from PIL import Image
import torch #, torchvision

print(torch.__version__, torch.cuda.is_available())

from detectron2.utils.logger import setup_logger
#import os, cv2, random
import cv2
import os

os.environ.pop("QT_QPA_PLATFORM_PLUGIN_PATH") # compatibility with PyQt5
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
#from detectron2.data import DatasetCatalog

#from detectron2.data.datasets import register_coco_instances
#from detectron2.config import LazyConfig
#from detectron2.engine import DefaultTrainer
#from detectron2.utils.visualizer import ColorMode
from detectron2.checkpoint import DetectionCheckpointer
from detectron2 import modeling

setup_logger()

import time, os #, fnmatch, shutil

def calculateDistance(x1, y1, x2, y2):
    dist = math.hypot(x2 - x1, y2 - y1)
    return dist

def get_cropped_leaf(img, predictor, return_mapping=False, resize=None):
    # convert to numpy
    img = np.array(img)[:, :, ::-1]
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

    # get prediction
    outputs = predictor(img)

    # get boxes
    ins = outputs["instances"]
    pred_masks = ins.get_fields()["pred_masks"]
    boxes = ins.get_fields()["pred_boxes"]

    # get main leaf mask if the area is >= the mean area of boxes and is closes to the centre
    masker = pred_masks[np.argmin(
        [calculateDistance(x[0], x[1], int(img.shape[1] / 2), int(img.shape[0] / 2)) for i, x in
         enumerate(boxes.get_centers()) if (boxes[i].area() >= torch.mean(boxes.area()).to("cpu")).item()])].to(
        "cpu").numpy().astype(np.uint8)

    # mask image
    mask_out = cv2.bitwise_and(img, img, mask=masker)

    # find contours and boxes
    contours, hierarchy = cv2.findContours(masker.copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    contour = contours[np.argmax([cv2.contourArea(x) for x in contours])]
    rotrect = cv2.minAreaRect(contour)
    box = cv2.boxPoints(rotrect)
    box = np.int0(box)

    # crop image
    cropped = get_cropped(rotrect, box, mask_out)

    # resize
    rotated = MakeLandscape()(Image.fromarray(cropped))

    if not resize == None:
        resized = ResizeMe((resize[0], resize[0], resize[1]))(rotated)
    else:
        resized = rotated

    if return_mapping:
        img = cv2.drawContours(img, [box], 0, (0, 0, 255), 10)
        img = cv2.drawContours(img, contours, -1, (255, 150,), 10)
        return resized, ResizeMe((int(resize[0]), int(resize[1])))(Image.fromarray(img))

    return resized, mask_out, contours, contour

def get_cropped(rotrect, box, image):
    width = int(rotrect[1][0])
    height = int(rotrect[1][1])

    src_pts = box.astype("float32")
    # coordinate of the points in box points after the rectangle has been straightened
    dst_pts = np.array([[0, height - 1],
                        [0, 0],
                        [width - 1, 0],
                        [width - 1, height - 1]], dtype="float32")

    # the perspective transformation matrix
    M = cv2.getPerspectiveTransform(src_pts, dst_pts)

    # directly warp the rotated rectangle to get the straightened rectangle
    warped = cv2.warpPerspective(image, M, (width, height))
    return warped

# image manipulations

class ResizeMe(object):
    # resize and center image in desired size
    def __init__(self, desired_size):

        self.desired_size = desired_size

    def __call__(self, img):

        img = np.array(img).astype(np.uint8)

        desired_ratio = self.desired_size[1] / self.desired_size[0]
        actual_ratio = img.shape[0] / img.shape[1]

        desired_ratio1 = self.desired_size[0] / self.desired_size[1]
        actual_ratio1 = img.shape[1] / img.shape[0]

        if desired_ratio < actual_ratio:
            img = cv2.resize(img, (int(self.desired_size[1] * actual_ratio1), self.desired_size[1]), None,
                             interpolation=cv2.INTER_AREA)
        elif desired_ratio > actual_ratio:
            img = cv2.resize(img, (self.desired_size[0], int(self.desired_size[0] * actual_ratio)), None,
                             interpolation=cv2.INTER_AREA)
        else:
            img = cv2.resize(img, (self.desired_size[0], self.desired_size[1]), None, interpolation=cv2.INTER_AREA)

        h, w, _ = img.shape

        new_img = np.zeros((self.desired_size[1], self.desired_size[0], 3))

        hh, ww, _ = new_img.shape

        yoff = int((hh - h) / 2)
        xoff = int((ww - w) / 2)

        new_img[yoff:yoff + h, xoff:xoff + w, :] = img

        return Image.fromarray(new_img.astype(np.uint8))

class MakeLandscape():
    # flip if needed
    def __init__(self):
        pass

    def __call__(self, img):
        if img.height > img.width:
            img = np.rot90(np.array(img))
            img = Image.fromarray(img)
        return img

def cv2_imshow(im):
    im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
    plt.figure(), plt.imshow(im), plt.axis('off')
    plt.show() #maybe turn off?
    print("cv2_imshow function used!")

def load_dataset(DATASETNAME):
    # MODELNAME = "/" + DATASETNAME + "_model.pth"

    DATASETNAME = "inputs/salmon1015to1020"

    print("Setting config...")
    cfg = get_cfg()
    cfg.MODEL.DEVICE = "cuda"
    cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
    cfg.DATASETS.TRAIN = ("salmon_train",)
    cfg.DATASETS.TEST = ()
    cfg.DATALOADER.NUM_WORKERS = 8  # previously 2
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
        "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  # Let training initialize from model zoo
    cfg.SOLVER.IMS_PER_BATCH = 6
    cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR
    cfg.SOLVER.MAX_ITER = 400  # 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset
    cfg.SOLVER.STEPS = []  # do not decay learning rate
    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128  # faster, and good enough for this toy dataset (default: 512)
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # only has one class (ballon). (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets)
    # NOTE: this config means the number of classes, but a few popular unofficial tutorials incorrect uses num_classes+1 here.
    # DefaultTrainer.auto_scale_workers(cfg,8)
    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)

    # Inference should use the config with parameters that are used in training
    # cfg now already contains everything we've set previously. We changed it a little bit for inference:

    cfg.MODEL.WEIGHTS = os.path.join("output/Models/30des2021model.pth")  # path to model previously trained
    print(cfg.MODEL.WEIGHTS)
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7  # set a custom testing threshold

    print("Loading default predictor with custom config")
    predictor = DefaultPredictor(cfg)
    model = modeling.build_model(cfg)

    DetectionCheckpointer(model).load(cfg.OUTPUT_DIR + "/model_final.pth")

    print("Testing to see if we got anything useful. Let us load our new dataset and print some of the")
    print("predictions to see how they are!")
    #data = DatasetCatalog.get("salmon_train")

    #return data, cfg, model, predictor
    return cfg, model, predictor

from scipy.interpolate import UnivariateSpline

def smoother(mask, plot=False, k=1, s=2, N=400):
    x = mask[:, 0]
    y = mask[:, 1]

    x_new = np.zeros([len(x)+1])
    y_new = np.zeros([len(y)+1])

    x_new[0:len(x)] = x
    x_new[-1] = x[0]
    y_new[0:len(y)] = y
    y_new[-1] = y[0]

    x = x_new
    y = y_new

    points = np.vstack((x, y)).T

    # Linear length along the line:
    distance = np.cumsum(np.sqrt(np.sum(np.diff(points, axis=0) ** 2, axis=1)))
    distance = np.insert(distance, 0, 0) / distance[-1]

    # for s in range(25,50,5):
    # Build a list of the spline function, one for each dimension:
    # k = 1
    # s = 50
    splines = [UnivariateSpline(distance, coords, k=k, s=s) for coords in points.T]

    # Computed the spline for the asked distances:
    alpha = np.linspace(0, 1, N)
    points_fitted = np.vstack(spl(alpha) for spl in splines).T

    # for spl in splines:
    #    points_fitted = np.vstack(spl(alpha)).T

    if plot == True:
        # Graph:
        # plt.plot(*points.T, 'ok', label='original points')
        # plt.plot(mask[:, 0], mask[:, 1])
        plt.plot(points_fitted[:, 0], points_fitted[:, 1], 'ro', label='fitted spline k=' + str(k) + ', s=' + str(s))
        plt.plot(x,y,'bx')
        plt.plot(x[0],y[0],'go',x[-1],y[-1],'go')
        plt.axis('equal')
        plt.legend()
        plt.xlabel('x')
        plt.ylabel('y')
        # plt.show()
        # time.sleep(1)
        # plt.pause(10)

        plt.show()
        wait = input('Press any key to go to the next iteration')
    return points_fitted

# import mask_smoother

if __name__ == "__main__":

    save = False
    plotting = True

    t = time.localtime()
    timestamp = time.strftime('%b-%d-%Y_%H%M', t)
    OUTPUT_FOLDER = ("mask_output_" + timestamp)
    # DATASETNAME = "30des2021coco"
    DATASETNAME = "inputs/salmon1015to1020"

    OUTPUT_FOLDER = ("mask_output_april_" + timestamp)
    # DATASETNAME = "inputs/salmon1015to1020"
    cfg, model, predictor = load_dataset(DATASETNAME)

    print("Loading dataset: " + DATASETNAME)
    print("Output folder: " + OUTPUT_FOLDER)
    print("Saving set to: " + str(save))

    count = 1
    dirlist = os.listdir(DATASETNAME)
    print("Sequencing images now. In total " + str(len(dirlist)) + " images")
    image_file: str
    if save==True:
        if not os.path.exists(OUTPUT_FOLDER):
            os.makedirs(OUTPUT_FOLDER)
            os.makedirs(OUTPUT_FOLDER + "/images")
            os.makedirs(OUTPUT_FOLDER + "/masks")

    for image_file in dirlist:

        print("Processing image " + str(count) + " of " + str(len(dirlist)))
        count += 1
        # image_file = "salmon_5_1_extract.jpg"
        # img: np.ndarray = cv2.imread(image_file)

        img: np.ndarray = cv2.imread(DATASETNAME + "/" + image_file)
        try:
            output: Instances = predictor(img)["instances"]
        except AttributeError: # skip image if there is an error
            print("Image skipped! Image " + image_file + " encountered an attribute error and had to be skipped")
            continue

        v = Visualizer(img[:, :, ::-1],
                       MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
                       scale=1.0)
        result: VisImage = v.draw_instance_predictions(output.to("cpu"))
        result_image: np.ndarray = result.get_image()[:, :, ::-1]
        cv2_imshow(result_image)

        # get file name without extension, -1 to remove "." at the end
        # out_file_name: str = re.search(r"(.*)\.", image_file).group(0)[:-1]
        # out_file_name += "_processed.png"

        resized, mask_out, contours, contour = get_cropped_leaf(img, predictor)  # returntype true

        data_in_array = np.array(contour)
        transposed = data_in_array.T
        x, y = transposed

        contour_2D = np.vstack((x, y)).T
        #mask, plot = False, k = 1, s = 2, N = 400
        smooth_contour_2D = smoother(contour_2D, plot=plotting, k=1, s=3, N=400)

        #plotting=False
        if plotting:
            fig, ax = plt.subplots()

            ax.plot(contour_2D[:, 0], contour_2D[:, 1], 'o')
            #ax.plot(smooth_contour_2D[:, 0], smooth_contour_2D[:, 1])
            ax.legend(["Contour", "Smoothed countour"])
            ax.set_title("Contour generated by detectron2")
            ax.set_aspect('equal', 'box')
            ax.set_xlabel("x-axis")
            ax.set_ylabel("y-axis")
            plt.show()

        # fig, ax = plt.subplots(1,1)
        # ax.plot(contour_3D[:, 0], contour_3D[:, 1], 'ro')
        # ax.axis('equal')
        # plt.title("contour_3D")
        # fig.show()

        # plt.plot(x_val, y_val)
        # plt.show()

        # save numpy array as npy file
        from numpy import asarray
        from pathlib import Path
        if save==True:

            #cv2.imwrite(OUTPUT_FOLDER+"/images/"+image_file, result_image)

            path_filename = Path(image_file)
            image_file_nosuffix = str(path_filename.with_suffix(''))

            #os.makedirs(OUTPUT_FOLDER + "/masks/" + image_file )#os.path.basename(image_file["file_name"]))
            os.makedirs(OUTPUT_FOLDER + "/masks/" + image_file_nosuffix)
            #cv2.imwrite(OUTPUT_FOLDER + "/images/" + os.path.basename(image_file["file_name"]) + ".jpg", result_image)

            np.save(OUTPUT_FOLDER + "/masks/" + image_file_nosuffix + "/" + image_file_nosuffix, smooth_contour_2D)

            print("Masks saved!")

    print("script finished!")


from flask import Flask
import pika

app = Flask(__name__)

@app.route("/sendMsg/<string:msg>")
def send(msg):
    parameters = pika.URLParameters('amqp://imaya:imaya@rabbitmq:5672/%2F')
    connection = pika.BlockingConnection(parameters)
    channel = connection.channel()
    channel.queue_declare(queue="request", durable=True)
    channel.basic_publish(exchange='', routing_key='request', body=msg, properties=pika.BasicProperties(delivery_mode = 2))
    connection.close()

    html = "
Hello !
" \
       "Message Sent
" \
       "Again"
    return html.format()

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


from coinapi_service import coinapi_get_exchange_rates_extended
import json
from pathlib import Path
from datetime import datetime, date

def get_json_rates(rates_data):
    return json.dumps([{ "date":d['time_period_start'][0:10], "Value":int(d['rate_close']) } for d in rates_data])

def save_data_to_json_file(json_data, filename):
    with open(filename, 'w') as file:
        file.write(json_data)
        file.close()

def load_json_data_from_file(data_filename_to_load):
    with open(data_filename_to_load, 'r') as file:
        try:
            return json.load(file)
        except json.JSONDecodeError:
            print("Le fichier est vide ou contient des données JSON non valides.")
            return []

def get_and_manage_rates_data(assets,date_start_str,date_end_str):
    pass

assets = "BTC/EUR"
data_filename = assets.replace("/","_")+".json"
date_start_str = date(2023, 4, 10)
date_end_str = date.today()
rates= []
if Path(data_filename).is_file():
    #Json File exist
    rates= load_json_data_from_file(data_filename)
    if rates:
        save_data_date_start_str = datetime.strptime(rates[0]["date"], "%Y-%m-%d").date()
        save_data_date_end_str = datetime.strptime(rates[-1]["date"], "%Y-%m-%d").date()
        print(save_data_date_end_str, save_data_date_start_str)
        print("save_data_date_start_str<date_start_str: ",save_data_date_start_str<date_start_str)
        if save_data_date_start_str<date_start_str:
            rates = coinapi_get_exchange_rates_extended(assets,save_data_date_start_str,date_start_str) + rates
        if save_data_date_end_str > date_end_str :
            rates+=coinapi_get_exchange_rates_extended(assets,date_end_str,save_data_date_end_str)
        print("save_data_date_end_str > date_end_str :",save_data_date_end_str > date_end_str)
        rates = get_json_rates(rates)
        save_data_to_json_file(rates, data_filename)
else:
    #Json File don't exist
    rates = coinapi_get_exchange_rates_extended(assets, date_start_str,date_end_str)
    print()
    if rates:
        json_data= get_json_rates(rates)
        print(assets, "nombres de cours : ", len(rates))
        save_data_to_json_file(json_data, data_filename)


# MATH 3670 Article Clustering
# Authors: Akash Patel, Nabeel Mahmood, Taylor Fields

import statistics as stat
from nltk.corpus import reuters
import os
import numpy as np
np.set_printoptions(threshold=np.inf)

from scipy.cluster.vq import kmeans, vq

### Get complete word list

os.listdir('/home/apatel435/nltk_data/corpora/reuters/training/')

wordlist = [];
num_articles = 0
artnum = 1000
j = 0
for filename in os.listdir('/home/apatel435/nltk_data/corpora/reuters/training/'):
    if j < artnum:
        words = reuters.words('training/' + str(filename))
        wordlist.append(words)
        num_articles += 1
    j += 1

# Flatten
wordlist = [item for sublist in wordlist for item in sublist]
# Lowercase
wordlist = [x.lower() for x in wordlist]
# Filter repeats
wordlist = list(set(wordlist))
# Sort
wordlist.sort()

### Filter common words and special characters
# take out numbers
wordlist = [ x for x in wordlist if x.isalpha() ]

# Filter common words
common = ["a", "about", "above", "above", "across", "after", "afterwards",
"again", "against", "all", "almost", "alone", "along", "already",
"also","although","always","am","among", "amongst", "amongst", "amount",  "an",
"and", "another", "any","anyhow","anyone","anything","anyway", "anywhere",
"are", "around", "as",  "at", "back","be","became",
"because","become","becomes", "becoming", "been", "before", "beforehand",
"behind", "being", "below", "beside", "besides", "between", "beyond", "bill",
"both", "bottom","but", "by", "call", "can", "cannot", "cant", "co", "con",
"could", "couldnt", "do", "done", "down",
"due", "during", "each", "e.g.", "eight", "either", "eleven","else",
"elsewhere", "empty", "enough", "etc", "even", "ever", "every", "everyone",
"everything", "everywhere", "except", "few", "fifteen", "fifty", "fill", "find",
"fire", "first", "five", "for", "former", "formerly", "forty", "found", "four",
"from", "front", "full", "further", "get", "give", "go", "had", "has", "hasnt",
"have", "he", "hence", "her", "here", "hereafter", "hereby", "herein",
"hereupon", "hers", "herself", "him", "himself", "his", "how", "however",
"hundred", "ie", "if", "in", "inc", "indeed", "interest", "into", "is", "it",
"its", "itself", "keep", "last", "latter", "latterly", "least", "less", "ltd",
"made", "many", "may", "me", "meanwhile", "might", "mill", "mine", "more",
"moreover", "most", "mostly", "move", "much", "must", "my", "myself", "name",
"namely", "neither", "never", "nevertheless", "next", "nine", "no", "nobody",
"none", "noone", "nor", "not", "nothing", "now", "nowhere", "of", "off",
"often", "on", "once", "one", "only", "onto", "or", "other", "others",
"otherwise", "our", "ours", "ourselves", "out", "over", "own","part", "per",
"perhaps", "please", "put", "rather", "re", "same", "see", "seem", "seemed",
"seeming", "seems", "serious", "several", "she", "should", "show", "side",
"since", "sincere", "six", "sixty", "so", "some", "somehow", "someone",
"something", "sometime", "sometimes", "somewhere", "still", "such",
"take", "ten", "than", "that", "the", "their", "them", "themselves", "then",
"thence", "there", "thereafter", "thereby", "therefore", "therein", "thereupon",
"these", "they", "thick", "thin", "third", "this", "those", "though", "three",
"through", "throughout", "thru", "thus", "to", "together", "too", "top",
"toward", "towards", "twelve", "twenty", "two", "un", "under", "until", "up",
"upon", "us", "very", "via", "was", "we", "well", "were", "what", "whatever",
"when", "whence", "whenever", "where", "whereafter", "whereas", "whereby",
"wherein", "whereupon", "wherever", "whether", "which", "while", "whither",
"who", "whoever", "whole", "whom", "whose", "why", "will", "with", "within",
"without", "would", "yet", "you", "your", "yours", "yourself", "yourselves",
"the"]

wordlist = [x for x in wordlist if x not in common]

num_wordlist = len(wordlist)

### Continue with matrix building
# Add a 1 if the article contains the word in wordlist
arr = np.zeros((num_wordlist,num_articles))

i = 0

for filename in os.listdir('/home/apatel435/nltk_data/corpora/reuters/training/'):

    if i < artnum:
        words = reuters.words('training/' + str(filename))
        # Lowercase
        words = [x.lower() for x in words]
        # Filter repeats
        words = list(set(words))

        # Add the one if word is in the wordlist
        for x in range(0, num_wordlist):
            if wordlist[x] in words:
                arr[x,i] = 1

    i += 1

stackarr = arr[:,0]

for x in range(1,artnum):
    stackarr = np.vstack((stackarr,arr[:,x]))

# Covariance matrix
cov_arr = np.cov(stackarr)

### Perform k - means algo clustering
# Cluster size chosen to be 5 due to relatively large dataset (1000x1000)
clusters = 5

centroids,_= kmeans(cov_arr,clusters)
idx,_ = vq(cov_arr,centroids)

# Sort matrix by clusters
# ie vectors that are in cluster one are adjacent to one another and so on
# for each cluster

index = idx.argsort();
cov_arrsort = cov_arr[index];

# Transpose matrix to sort the rows
# Matrix is symmetric so transposing it produces the same result
cov_arrtrans = cov_arrsort.transpose();
cluster_cov_arr = cov_arrtrans[index];
print(cluster_cov_arr)


from sklearn.gaussian_process import GaussianProcessRegressor, kernels
from sklearn.datasets import make_regression
import numpy as np
import matplotlib.pyplot as plt
import scipy.stats as ss

'''
def log_marginal_likelihood(gp, theta):
    return gp.log_marginal_likelihood(theta)

def log_prior_parameters(theta):
    return np.sum(np.log([ss.norm(0, 1).pdf(theta_k) for theta_k in theta]))

def log_joint_unnorm(gp, theta):
    return log_marginal_likelihood(gp, theta) + log_prior_parameters(theta)
'''

class Gaussian_Process():

    def __init__(self, kernel, mode):
        self.kernel = kernel
        if mode == 'OPT':
            self.gp = GaussianProcessRegressor(kernel=kernel,
                                               n_restarts_optimizer=10,
                                               normalize_y=True)
        elif mode == 'MAP':
            self.gp = GaussianProcessRegressor(kernel=kernel,
                                               optimizer = None,
                                               n_restarts_optimizer=10,
                                               normalize_y=True)
        else:
            raise Exception('Wrong GP initial mode!!!')
    '''
    def optimise__with_splice_sampling(self, initial_theta, num_iters, sigma, burn_in,  X, y):
        self.fit(X, y)
        slice_sampler = Slice_sampler(num_iters = num_iters,
                                      sigma = sigma,
                                      burn_in = burn_in,
                                      gp = self)
        samples = slice_sampler.sample(init = initial_theta)
        theta_opt = [np.mean(samples[0]),np.mean(samples[1])]

        self.gp.set_params(**{"kernel__k1__length_scale": np.mean(samples[0]),
                            "kernel__k2__noise_level": np.power(np.mean(samples[0]),2)})
    '''

    def set_params(self, theta_set):
        self.gp.set_params(**{"kernel__k1__noise_level": np.abs(theta_set[0]),
                        "kernel__k2__k1__constant_value": np.abs(theta_set[1]),
                        "kernel__k2__k2__length_scale": theta_set[2:]})

    def log_marginal_likelihood(self,theta):
        return self.gp.log_marginal_likelihood(theta)

    def log_prior_parameters(self, theta):
        return np.sum(np.log([ss.norm(0, 1).pdf(theta_k) for theta_k in theta]))

    def log_joint_unnorm(self, theta):
        return self.log_marginal_likelihood(theta) + self.log_prior_parameters(theta)

    def fit(self, X, y):
        self.gp.fit(X, y)

    def predict(self, X):
        mu, sigma = self.gp.predict(X, return_std =True)
        return mu, sigma

class Slice_sampler():

    def __init__(self, num_iters, sigma, burn_in):
        self.num_iters = num_iters # specify number of samples we want to draw
        self.sigma = sigma
        self.burn_in = burn_in

    def sample(self, init, gp, step_out=True):

        D = len(init)
        samples = np.zeros((D, self.num_iters))

        xx = init.copy()

        for i in range(self.num_iters + self.burn_in):
            perm = list(range(D))
            np.random.shuffle(perm)
            last_llh = gp.log_joint_unnorm(xx)

            for d in perm:
                llh0 = last_llh + np.log(np.random.rand())
                rr = np.random.rand(1)
                x_l = xx.copy()
                x_l[d] = x_l[d] - rr * self.sigma[d]
                x_r = xx.copy()
                x_r[d] = x_r[d] + (1 - rr) * self.sigma[d]

                if step_out:
                    llh_l = gp.log_joint_unnorm(x_l)
                    while llh_l > llh0:
                        x_l[d] = x_l[d] - self.sigma[d]
                        llh_l = gp.log_joint_unnorm(x_l)
                    llh_r = gp.log_joint_unnorm(x_r)
                    while llh_r > llh0:
                        x_r[d] = x_r[d] + self.sigma[d]
                        llh_r = gp.log_joint_unnorm(x_r)

                x_cur = xx.copy()
                while True:
                    xd = np.array(np.random.rand() * (x_r[d] - x_l[d]) + x_l[d])
                    x_cur[d] = xd.copy()
                    last_llh = gp.log_joint_unnorm(x_cur)
                    if last_llh > llh0:
                        xx[d] = xd.copy()
                        break
                    elif xd > xx[d]:
                        x_r[d] = xd
                    elif xd < xx[d]:
                        x_l[d] = xd
                    else:
                        raise RuntimeError('Slice sampler shrank too far.')

            if i == 0:
                pass
                #print ("burn-in")
            elif i > self.burn_in and i % 100 == 0:
                pass
                #print ('iteration', i - self.burn_in)

            if i >= self.burn_in:
                samples[:, i-self.burn_in] = xx.copy().ravel() # index starting from 0

        #plt.hist(samples[0])

        return samples


from __future__ import print_function
import simpy
import random
from statistics import mean
from pifo_wrapper import SkipListWrapper

def test(env):
    NUM_SKIP_LISTS = 6
    PERIOD = 1
    MAX_NODES = 512
    OUTREG_WIDTH = 16
    ENQ_FIFO_DEPTH = 16
    CLK_FREQ = 200 # MHz
    PKT_RATE = 14.8*4 # MPkts/sec
    PKT_INTERVAL = int(CLK_FREQ/PKT_RATE)
    NumRuns = 1
    NumOps = 128
    enq_nclks_list = []
    deq_nclks_list = []

    enq_in_pipe  = simpy.Store(env)
    enq_out_pipe = simpy.Store(env)
    deq_in_pipe  = simpy.Store(env)
    deq_out_pipe = simpy.Store(env)

    print ('@ {:04d} - starting skip list init'.format(env.now))

    slw = SkipListWrapper(env, enq_in_pipe, enq_out_pipe, deq_in_pipe, deq_out_pipe, num_sl=NUM_SKIP_LISTS, period=PERIOD, size=MAX_NODES, outreg_width=OUTREG_WIDTH, enq_fifo_depth=ENQ_FIFO_DEPTH, rd_latency=1, wr_latency=1)

    for i in range(NUM_SKIP_LISTS):
        while slw.sl[i].busy == 1:
            yield env.timeout(PERIOD)

    print ('@ {:04d} - done skip list init'.format(env.now))

    for j in range(NumRuns):
        print ("Run:", j)
        # Enqueue some values and print skip list
        for i in range (NumOps):
            val = random.randint(0,100)
            hsp = mdp = -1
            t1 = env.now
            slw.enq_in_pipe.put((val, hsp, mdp))
            yield slw.enq_out_pipe.get()
            enq_nclks = env.now - t1
            enq_nclks_list.append(enq_nclks)
            print ('enq: {} - {} clks'.format(val, enq_nclks))
            yield env.timeout(PKT_INTERVAL)

        for i in range(NUM_SKIP_LISTS):
            while slw.sl[i].enq_fifo.fill_level() != 0 or slw.sl[i].busy == 1:
                yield env.timeout(PERIOD)

        # Dequeue all values and print skip list
        while slw.num_entries > 0:
            t1 = env.now
            slw.deq_in_pipe.put(True)
            (val, hsp, mdp, deq_nclks) = yield slw.deq_out_pipe.get()
            deq_nclks = env.now - t1
            print ('deq: {} - {} clks'.format(val, deq_nclks))
            deq_nclks_list.append(deq_nclks)
            yield env.timeout(PKT_INTERVAL)

    # Wait until skip lists are done
    for i in range(NUM_SKIP_LISTS):
        while slw.sl[i].busy == 1:
            yield env.timeout(PERIOD)
        # Stop deq_sl processes
        slw.sl[i].enq_sl_proc.interrupt('Done')
        slw.sl[i].deq_sl_proc.interrupt('Done')

    print ("Time measurements (min, avg, max) for {} runs of {} enq/deq ops".format(NumRuns, NumOps))
    print ("Enq: {:5.2f} {:5.2f} {:5.2f}".format(min(enq_nclks_list), mean(enq_nclks_list), max(enq_nclks_list)))
    print ("Deq: {:5.2f} {:5.2f} {:5.2f}".format(min(deq_nclks_list), mean(deq_nclks_list), max(deq_nclks_list)))

# Main
env = simpy.Environment()
env.process(test(env))
env.run()


import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions.categorical import Categorical
import torch.nn.functional as F
import numpy as np

class PPOmemory:
    def __init__(self, mini_batch_size, cfg):
        self.gamma = cfg.gamma  # 折扣因子
        self.gae_lambda = cfg.gae_lambda  # GAE参数
        self.states = []  # 状态
        self.actions = []  # 实际采取的动作
        self.probs = []  # 动作概率
        self.vals = []  # critic输出的状态值
        self.rewards = []  # 奖励
        self.dones = []  # 结束标志
        self.available_actions = []

        self.mini_batch_size = mini_batch_size  # minibatch的大小

    def gae_adv(self,):

        # 计算GAE
        vals_arr, rewards_arr, dones_arr = np.array(self.vals), np.array(self.rewards), np.array(self.dones)
        returns = list()
        G = 0

        for r in rewards_arr[::-1]:
            G = r + self.gamma * G
            returns.insert(0, G)

        self.advantage = np.array(returns)

    def sample(self):
        n_states = len(self.states)  # memory记录数量=20
        self.mini_batch_size = n_states
        batch_start = np.arange(0, n_states, self.mini_batch_size)  # 每个batch开始的位置[0,5,10,15]
        indices = np.arange(n_states, dtype=np.int64)  # 记录编号[0,1,2....19]
        np.random.shuffle(indices)  # 打乱编号顺序[3,1,9,11....18]
        mini_batches = [indices[i:i + self.mini_batch_size] for i in batch_start]  # 生成4个minibatch，每个minibatch记录乱序且不重复

        return np.array(self.states), np.array(self.actions), np.array(self.probs), \
               np.array(self.vals), np.array(self.rewards), np.array(self.dones), self.advantage, np.array(self.available_actions), mini_batches

    # 每一步都存储trace到memory
    def push(self, state, action, prob, val, reward, done, available_action):
        self.states.append(state)
        self.actions.append(action)
        self.probs.append(prob)
        self.vals.append(val)
        self.rewards.append(reward)
        self.dones.append(done)
        self.available_actions.append(available_action)

    # 固定步长更新完网络后清空memory
    def clear(self):
        self.states = []
        self.actions = []
        self.probs = []
        self.vals = []
        self.rewards = []
        self.dones = []

class Actor(nn.Module):
    def __init__(self, n_states, n_actions, hidden_dim):
        super(Actor, self).__init__()
        self.base1 = nn.Sequential(
            nn.Linear(n_states, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),)
        self.base2 = nn.Sequential(
            nn.Linear(n_states, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),)

        self.out1 = nn.Linear(hidden_dim, n_actions)
        self.out2 = nn.Linear(hidden_dim, n_actions)

    def forward(self, state, available_action, test=False):
        mask = (available_action == 0)
        x1 = self.base1(state)
        x1 = self.out1(x1)
        x2 = self.base2(state)
        x2 = self.out2(x2)
        x1[mask[:,0]] = -1e10
        x2[mask[:,1]] = -1e10
        probs_in = F.softmax(x1, dim=1)
        probs_out = F.softmax(x2, dim=1)
        dist_in = Categorical(probs_in)
        dist_out = Categorical(probs_out)
        if test:
            return [probs_in, probs_out]
        else:
            return {'dist_in':dist_in, 'dist_out':dist_out}

class Critic(nn.Module):
    def __init__(self, n_states, hidden_dim):
        super(Critic, self).__init__()
        self.critic = nn.Sequential(
            nn.Linear(n_states, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1))

    def forward(self, state):
        value = self.critic(state)
        return value

class Agent:
    def __init__(self, n_states, n_actions, cfg):
        # 训练参数
        self.gamma = cfg.gamma  # 折扣因子
        self.n_epochs = cfg.n_epochs  # 每次更新重复次数
        self.gae_lambda = cfg.gae_lambda  # GAE参数
        self.policy_clip = cfg.policy_clip  # clip参数
        self.device = cfg.device  # 运行设备

        # AC网络及优化器
        self.actor = Actor(n_states, n_actions, cfg.hidden_dim).to(cfg.device)
        self.critic = Critic(n_states, cfg.hidden_dim).to(cfg.device)
        self.actor_optim = optim.Adam(self.actor.parameters(), lr=cfg.actor_lr)
        self.critic_optim = optim.Adam(self.critic.parameters(), lr=cfg.critic_lr)

        # 经验池
        self.memory = PPOmemory(cfg.mini_batch_size,cfg)

    def choose_action(self, state, available_action):
        state = torch.FloatTensor(state).unsqueeze(0).to(self.device)  # 数组变成张量
        available_action = torch.LongTensor(available_action).unsqueeze(0).to(self.device)  # 数组变成张量

        dist = self.actor(state, available_action)  # action分布
        dist_in, dist_out = dist['dist_in'], dist['dist_out']
        value = self.critic(state)  # state value值

        action_in = dist_in.sample()  # 随机选择action
        action_out = dist_out.sample()  # 随机选择action

        prob_in = torch.squeeze(dist_in.log_prob(action_in)).item()  # action对数概率
        prob_out = torch.squeeze(dist_out.log_prob(action_out)).item()  # action对数概率

        action_in = torch.squeeze(action_in).item()
        action_out = torch.squeeze(action_out).item()
        value = torch.squeeze(value).item()
        return [action_in, action_out], [prob_in, prob_out], value

    def test_action(self, state, available_action):
        state = torch.FloatTensor(state).unsqueeze(0).to(self.device)  # 数组变成张量
        available_action = torch.LongTensor(available_action).unsqueeze(0).to(self.device)  # 数组变成张量

        probs = self.actor(state, available_action, test=True)  # action分布
        probs_in, probs_out = probs[0], probs[1]

        action_in = probs_in.argmax()  # 随机选择action
        action_out = probs_out.argmax()  # 随机选择action

        action_in = torch.squeeze(action_in).item()
        action_out = torch.squeeze(action_out).item()
        return [action_in, action_out]

    def learn(self):
        self.memory.gae_adv()
        for _ in range(self.n_epochs):
            # memory中的trace以及处理后的mini_batches，mini_batches只是trace索引而非真正的数据
            states_arr, actions_arr, old_probs_arr, vals_arr,\
                rewards_arr, dones_arr, Gt, available_actions, mini_batches = self.memory.sample()

            # mini batch 更新网络
            values = vals_arr[:]
            Gt = torch.FloatTensor(Gt).to(self.device)
            values = torch.FloatTensor(values).to(self.device)
            for batch in mini_batches:
                states = torch.FloatTensor(states_arr[batch]).to(self.device)
                old_probs = torch.FloatTensor(old_probs_arr[batch]).to(self.device)
                actions = torch.LongTensor(actions_arr[batch]).to(self.device)
                active_mask = torch.FloatTensor(available_actions[batch].sum(axis=2) > 0).to(self.device)
                value_mask = torch.FloatTensor(available_actions[batch].sum(axis=(1,2)) > 0).to(self.device)
                available_actions = torch.LongTensor(available_actions[batch]).to(self.device)  # 数组变成张量
                # mini batch 更新一次critic和actor的网络参数就会变化
                # 需要重新计算新的dist,values,probs得到ratio,即重要性采样中的新旧策略比值
                critic_value = torch.squeeze(self.critic(states))

                dist = self.actor(states, available_actions)
                dist_in, dist_out = dist['dist_in'], dist['dist_out']
                new_probs_in = dist_in.log_prob(actions[:, 0]) # in action
                new_probs_out = dist_out.log_prob(actions[:, 1]) # out action
                prob_ratio_in = new_probs_in.exp() / old_probs[:, 0].exp() # in ratio
                prob_ratio_out = new_probs_out.exp() / old_probs[:, 1].exp() # in ratio

                advantage = Gt[batch] - critic_value.detach()

                # in
                weighted_probs_in = advantage * prob_ratio_in
                weighted_clip_probs_in = torch.clamp(prob_ratio_in, 1 - self.policy_clip,\
                     1 + self.policy_clip) * advantage
                actor_loss_in = -torch.sum(torch.min(weighted_probs_in, weighted_clip_probs_in) * active_mask[:,0]) / active_mask[:,0].sum()

                # in
                weighted_probs_out = advantage * prob_ratio_out
                weighted_clip_probs_out = torch.clamp(prob_ratio_out, 1 - self.policy_clip,\
                     1 + self.policy_clip) * advantage
                actor_loss_our = -torch.sum(torch.min(weighted_probs_out, weighted_clip_probs_out) * active_mask[:,1]) / active_mask[:,1].sum()

                actor_loss = actor_loss_in + actor_loss_our
                self.actor_optim.zero_grad()
                actor_loss.backward()
                self.actor_optim.step()

                # critic loss
                critic_loss = torch.sum(((Gt[batch] - critic_value) ** 2) * value_mask) / value_mask.sum()
                critic_loss = critic_loss.mean()
                # total_loss

                # 更新
                self.critic_optim.zero_grad()
                critic_loss.backward()
                self.critic_optim.step()
        self.memory.clear()


import pygame as pg
from settings import *
import net
from worrior import *

def make_text(text, font, color, surface, x, y):
    textobj = font.render(text, 1, color)
    textrect = textobj.get_rect()
    textrect.topleft = (x, y)
    surface.blit(textobj, textrect)

#def is_in_DB(nick):
#    return False

screen = pg.display.set_mode((1280, 720))
font = pg.font.Font(None, 32)
font2 = pg.font.Font(None, 64)

pg.display.set_caption('Main Menu')
clock = pg.time.Clock()
pg.init()
click = False
done = False
hero = Worrior()

def main_menu(player_list):

    input_box = pg.Rect(520, 300, 140, 32)
    color_inactive = pg.Color('lightskyblue3')
    color_active = pg.Color('dodgerblue2')
    color = color_inactive
    active = False
    text = ''
    done = False
    nick = ''

    print(f'player list: {player_list}')

    while not done:

        for event in pg.event.get():
            if event.type == pg.QUIT:
                done = True
            if event.type == pg.MOUSEBUTTONDOWN:

                if input_box.collidepoint(event.pos):

                    active = not active

                else:
                    active = False

                color = color_active if active else color_inactive
            if event.type == pg.KEYDOWN:
                if active:
                    if event.key == pg.K_RETURN:
                        nick = text
                        print(text)
                        text = ''
                        print(net.is_in_DB(nick))
                        playerData = next((x for x in player_list if x['name'] == nick), None)

                        if playerData == None:

                            print('wchodzi do pentli i ')
                            print(nick)
                            create_champ(nick)
                            done = True
                        else:
                            print(playerData)
                            global hero
                            hero = Worrior(ad=playerData['ad'], arm=playerData['arm'], curr_quests=playerData['curr_quests'], dir=playerData['dir'], exp=playerData['exp'], gold=playerData['gold'], hp=playerData['hp'], map=playerData['map'], moving=playerData['moving'], name=playerData['name'], pa=playerData['pa'], position_x=playerData['position_x'], position_y=playerData['position_y'], prev_map=playerData['prev_map'], prev_pos=playerData['prev_pos'], quests_completed=playerData['quests_completed'], sprite=playerData['sprite'], velocity=playerData['velocity'])
                            return hero

                    elif event.key == pg.K_BACKSPACE:
                        text = text[:-1]
                    else:
                        text += event.unicode

        screen.fill((30, 30, 30))

        txt_surface = font.render(text, True, color)

        width = max(200, txt_surface.get_width()+10)
        input_box.w = width

        screen.blit(txt_surface, (input_box.x+5, input_box.y+5))

        pg.draw.rect(screen, color, input_box, 2)

        make_text('Welcome in our game!', font2, color_inactive, screen, 375, 50)
        make_text('Insert your nickname to start the game', font, color_inactive, screen, 400, 200)
        make_text('Your Nick:', font, color_inactive, screen, 560, 260)
        pg.display.flip()

def create_champ(nick):
    global click
    running = True
    available_points = 20
    strength = 0
    health = 0
    armor = 0
    action = 0

    while running:
        screen.fill(darkgray)
        make_text('Creating your champion', font, black, screen, 20, 20)
        mx, my = pygame.mouse.get_pos()
        make_text(f'Remainings available points: {available_points}', font, black, screen, 900, 20)

        #strength --------------------
        button_strength_add = pygame.Rect(350, 100, 76, 76)
        button_strength_reduce = pygame.Rect(450, 100, 76, 76)

        make_text(f'Strength points: {strength}', font, black, screen, 20, 125)

        if button_strength_add.collidepoint((mx, my)):
            if click:
                if available_points > 0:
                    available_points -= 1
                    strength += 1

        if button_strength_reduce.collidepoint((mx, my)):
            if click:
                if strength > 0:
                    strength -= 1
                    available_points += 1

        #health-----------------------
        button_health_add = pygame.Rect(350, 200, 76, 76)
        button_health_reduce = pygame.Rect(450, 200, 76, 76)

        make_text(f'Health points: {health}', font, black, screen, 20, 225)

        if button_health_add.collidepoint((mx, my)):
            if click:
                if available_points > 0:
                    available_points -= 1
                    health += 1

        if button_health_reduce.collidepoint((mx, my)):
            if click:
                if health > 0:
                    health -= 1
                    available_points += 1

        #armor ----------------------
        button_armor_add = pygame.Rect(350, 300, 76, 76)
        button_armor_reduce = pygame.Rect(450, 300, 76, 76)

        make_text(f'Armor points: {armor}', font, black, screen, 20, 325)

        if button_armor_add.collidepoint((mx, my)):
            if click:
                if available_points > 0:
                    available_points -= 1
                    armor += 1

        if button_armor_reduce.collidepoint((mx, my)):
            if click:
                if armor > 0:
                    armor -= 1
                    available_points += 1

        #action ----------------------
        button_action_add = pygame.Rect(350, 400, 76, 76)
        button_action_reduce = pygame.Rect(450, 400, 76, 76)

        make_text(f'Action points: {action}', font, black, screen, 20, 425)

        if button_action_add.collidepoint((mx, my)):
            if click:
                if available_points > 0:
                    available_points -= 1
                    action += 1

        if button_action_reduce.collidepoint((mx, my)):
            if click:
                if action > 0:
                    action -= 1
                    available_points += 1

        # submit
        button_submit = pygame.Rect(800, 270, 256, 256)

        if button_submit.collidepoint((mx, my)):
            global hero, done
            if click:
                if available_points == 0:
                    action += 10
                    print("zmieniam hirka")
                    hero = Worrior(name=nick, sprite='./BULBA64.png', position_x = 0, position_y = 0, hp = (health * 50), ad = (strength * 25), arm = (armor * 10), pa = action)
                    print(f'moje statystyki {hero.name} \n {hero.sprite}   \n {hero.hp}  \n {hero.arm}')
                    done = True
                    running = False
                    return hero

        #generating buttons ------------
        pygame.draw.rect(screen, lightskyblue3, button_strength_add)
        pygame.draw.rect(screen, lightskyblue3, button_strength_reduce)
        screen.blit(plus2, button_strength_add)
        screen.blit(minus2, button_strength_reduce)

        pygame.draw.rect(screen, lightskyblue3, button_health_add)
        pygame.draw.rect(screen, lightskyblue3, button_health_reduce)
        screen.blit(plus2, button_health_add)
        screen.blit(minus2, button_health_reduce)

        pygame.draw.rect(screen, lightskyblue3, button_armor_add)
        pygame.draw.rect(screen, lightskyblue3, button_armor_reduce)
        screen.blit(plus2, button_armor_add)
        screen.blit(minus2, button_armor_reduce)

        pygame.draw.rect(screen, lightskyblue3, button_action_add)
        pygame.draw.rect(screen, lightskyblue3, button_action_reduce)
        screen.blit(plus2, button_action_add)
        screen.blit(minus2, button_action_reduce)

        pygame.draw.rect(screen, darkgray, button_submit)
        screen.blit(start3, button_submit)

        click = False

        for event in pg.event.get():
            if event.type == pg.QUIT:
                pygame.quit()
                sys.exit()
            if event.type == pg.KEYDOWN:
                if event.key == pg.K_ESCAPE:
                    running = False
            if event.type == pg.MOUSEBUTTONDOWN:
                if event.button == 1:
                    click = True

        pygame.display.update()
        clock.tick(60)

def get_hero():
    global hero
    return hero