# ---Dependencies---
import torch
import torch.nn as nn

# ---Custom---
from utils.ops import extract_kpt_vectors

Default = {'margin': 0.5, 'alpha': None, 'beta': None, 'n_neg': 10}

# noinspection PyArgumentList
class PixelwiseContrastiveLoss(nn.Module):
    """
    Implementation of "pixel-wise" contrastive loss. Contrastive loss typically compares two whole images.
            L = (Y) * (1/2 * d**2) + (1 - Y) * (1/2 * max(0, margin - d)**2)

    In this instance, we instead compare pairs of features within those images.
    Positive matches are given by ground truth correspondences between images.
    Negative matches are generated on-the-fly based on provided parameters.

    Attributes:
        margin (float): Target margin distance between positives and negatives
        alpha (int): Minimum distance from original positive KeyPoint
        beta (int): Maximum distance from original positive KeyPoint
        n-neg (int): Number of negative samples to generate

    Methods:
        forward: Compute pixel-wise contrastive loss
        forward_eval: Detailed forward pass for logging
    """
    def __init__(self, margin=0.5, alpha=None, beta=None, n_neg=10):
        super().__init__()
        self.margin = margin
        self.alpha = alpha
        self.beta = beta
        self.n_neg = n_neg

        self._dist = nn.PairwiseDistance()

    def __repr__(self):
        return f'{self.__class__.__qualname__}({self.margin}, {self.alpha}, {self.beta}, {self.n_neg})'

    def __str__(self):
        return f'Min{self.alpha or 0}_Max{self.beta or "Inf"}'

    @staticmethod
    def create_parser(parser):
        parser.add_argument('--margin', default=0.5, help='Target distance between negative feature embeddings.')
        parser.add_argument('--alpha', default=None, type=float, help='Minimum distance from positive KeyPoint')
        parser.add_argument('--beta', default=None, type=float, help='Maximum distance from positive KeyPoint')
        parser.add_argument('--n_neg', default=10, help='Number of negative samples to generate')

    def forward(self, features, labels):
        """ Compute pixel-wise contrastive loss.
        :param features: Vertically stacked feature maps (b, n-dim, h*2, w)
        :param labels: Horizontally stacked correspondence KeyPoints (b, n-kpts, 4) -> (x1, y1, x2, y2)
        :return: Loss
        """
        source, target = torch.chunk(features, 2, dim=-2)
        source_kpts, target_kpts = torch.chunk(labels, 2, dim=-1)

        loss = self._positive_loss(source, target, source_kpts, target_kpts)[0]
        loss += self._negative_loss(source, target, source_kpts, target_kpts)[0]
        return loss

    def forward_eval(self, features, labels):
        source, target = torch.chunk(features, 2, dim=-2)
        source_kpts, target_kpts = torch.chunk(labels, 2, dim=-1)

        pos_loss, pos_dist = self._positive_loss(source, target, source_kpts, target_kpts)
        neg_loss, neg_dist = self._negative_loss(source, target, source_kpts, target_kpts)

        loss = (pos_loss+neg_loss).item()
        output = {
            'scalars': {
                'loss': loss,
                'positive': pos_dist.mean().item(),
                'negative': neg_dist.mean().item(),
            },
            'histograms': {
                'hist_positive': pos_dist.detach().cpu(),
                'hist_negative': neg_dist.detach().cpu(),
            }
        }
        return loss, output

    def _calc_distance(self, source, target, source_kpts, target_kpts):
        source_descriptors = extract_kpt_vectors(source, source_kpts).permute([0, 2, 1])
        target_descriptors = extract_kpt_vectors(target, target_kpts).permute([0, 2, 1])
        return self._dist(source_descriptors, target_descriptors)

    def _positive_loss(self, source, target, source_kpts, target_kpts):
        dist = self._calc_distance(source, target, source_kpts, target_kpts)
        loss = (dist**2).mean() / 2
        return loss, dist

    def _negative_loss(self, source, target, source_kpts, target_kpts):
        dsource_kpts, dtarget_kpts = self._generate_negative_like(source, source_kpts, target_kpts)

        dist = self._calc_distance(source, target, dsource_kpts, dtarget_kpts)
        margin_dist = (self.margin - dist).clamp(min=0.0)
        loss = (margin_dist ** 2).mean() / 2
        return loss, dist

    def _generate_negative_like(self, other, source_kpts, target_kpts):
        # Source points remain the same
        source_kpts = source_kpts.repeat([1, self.n_neg, 1])

        # Target points + offset according to method
        target_kpts = target_kpts.repeat([1, self.n_neg, 1])
        target_kpts = self._permute_negatives(target_kpts, other.shape)
        return source_kpts, target_kpts

    def _permute_negatives(self, kpts, shape):
        h, w = shape[-2:]
        # (max(h, w) - low) means that even after getting the remainder points will be further away than low
        low = self.alpha if self.alpha else 0
        high = self.beta if self.beta else (max(h, w) - low)

        # Generate random shift for each KeyPoint
        shift = torch.randint_like(kpts, low=low, high=high)
        shift *= torch.sign(torch.rand_like(shift, dtype=torch.float)-0.5).short()  # Random + or - shift

        # Initial shift to satisfy max distance
        new_kpts = kpts + shift
        new_kpts %= torch.tensor((w, h), dtype=torch.short, device=new_kpts.device)

        # Shift to satisfy min distance
        diffs = new_kpts - kpts
        diff_clamp = torch.clamp(diffs, min=-high, max=high)
        new_kpts += (diff_clamp - diffs)

        return new_kpts


from StringIO import StringIO
from datetime import datetime
from lxml import etree, objectify
import os
import pytz
import re
import sys

class ObjectifiedDateTime(objectify.ObjectifiedDataElement):
    @property
    def pyval(self):
        result = re.match('(\d{4})-(\d{2})-(\d{2})T(\d{2}):(\d{2}):(\d{2})\.?(\d{1,6})?Z?', str(self)).groups()
        if result[6]:
            values = map(int, result)
        else:
            values = map(int, result[:6])

        values.append(pytz.timezone("UTC"))

        return datetime(*values)

    @staticmethod
    def check(string):
        if not re.match("\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.?\d{0,6}Z?", string):
            raise ValueError

xml_datetime_type = objectify.PyType('dateTime', ObjectifiedDateTime.check, ObjectifiedDateTime)
xml_datetime_type.xmlSchemaTypes = ('dateTime',)
xml_datetime_type.register()

class QMLParser(object):
    def __init__(self):
        # Patch XSD provided by QuakeML website to fit our needs
        xsd_path = os.path.join(os.path.dirname(__file__), 'QuakeML-BED-1.2.xsd')
        xsd_root = etree.parse(xsd_path).getroot()
        self.xsd_root = self.__patch_xsd(xsd_root)

        xsd_stringio =  StringIO(etree.tostring(self.xsd_root))
        self.parser = objectify.makeparser(schema=etree.XMLSchema(file=xsd_stringio))

    def __apply_attributes(self, obj):
        for attrname, value in obj.attrib.items():
            if attrname in ('channelCode', 'networkCode', 'publicID', 'stationCode'):
                setattr(obj, attrname, str(value))

        for obj in obj.__dict__.values():
            if isinstance(obj, objectify.ObjectifiedElement):
                map(self.__apply_attributes, obj)

    def __patch_xsd(self, root):
        del(root.attrib['targetNamespace'])

        # Remove any references to "bed" in "base" and "type" attributes
        #
        # attr.attrname not available with lxml 2.2
        # Can't use this code :
        # for attr in root.xpath('//@type | //@base'):
        #     if attr.startswith('bed:'):
        #         parent.attrib[attr.attrname] = attr[4:]
        for elem in root.xpath('//*[@type] | //*[@base]'):
            for attrname, attrvalue in elem.items():
                if attrvalue.startswith('bed:'):
                    elem.attrib[attrname] = attrvalue[4:]

        # Replace the first expected element (EventParameters) by Event element
        xsd_first_element = root.xpath("//*[@name='eventParameters']")[0]
        xsd_first_element.attrib['name'] = 'event'
        xsd_first_element.attrib['type'] = 'Event'

        # Remove publicID from Arrival
        arrival = root.xpath("//*[@name='Arrival']")[0]
        arrival_publicid = arrival.xpath(".//*[@name='publicID']")[0]
        arrival.remove(arrival_publicid)

        # Remove ResourceReference from WaveformstreamID
        waveformstreamid = root.xpath("//*[@name='WaveformStreamID']")[0]
        for elem in waveformstreamid.xpath(".//*[@name]"):
            waveformstreamid.append(elem)

        simplecontent = waveformstreamid.xpath("./*")[0]
        waveformstreamid.remove(simplecontent)

        # Remove any references to Namespaces in root element
        patched_root = etree.Element(root.tag, root.attrib)
        map(patched_root.append, root.getchildren())

        return patched_root

    def events(self, filename):
        for xml_event, elem in etree.iterparse(filename, tag='event'):
            xml_stringio = StringIO(etree.tostring(elem))
            elem.clear()

            obj = objectify.parse(xml_stringio, self.parser).getroot()
            self.__apply_attributes(obj)
            yield obj


import abc
import os
import re
import unicodedata
from typing import List, Union

import numpy as np
import tensorflow as tf
import tensorflow_datasets as tds
from utils.character import en, jp

import frontends.wordpiece as wordpiece

def preprocess_paths(paths: Union[List, str]):
    if isinstance(paths, list):
        return [
            path
            if path.startswith("gs://")
            else os.path.abspath(os.path.expanduser(path))
            for path in paths
        ]
    elif isinstance(paths, str):
        return (
            paths
            if paths.startswith("gs://")
            else os.path.abspath(os.path.expanduser(paths))
        )
    else:
        return None

class TextFeaturizer(metaclass=abc.ABCMeta):
    def __init__(self):
        self.blank = None
        self.tokens2indices = {}
        self.tokens = []
        self.num_classes = None

    def preprocess_text(self, text):
        text = unicodedata.normalize("NFC", text.lower())
        return text.strip("\n")  # remove trailing newline

    def normalize_indices(self, indices: tf.Tensor) -> tf.Tensor:
        """Remove -1 in indices by replacing them with blanks
        Args:
            indices (tf.Tensor): shape any

        Returns:
            (tf.Tensor): normalized indices with shape same as indices
        """
        minus_one = -1 * tf.ones_like(indices, dtype=tf.int32)
        blank_like = self.blank * tf.ones_like(indices, dtype=tf.int32)
        return tf.where(indices == minus_one, blank_like, indices)

    def prepand_blank(self, text: tf.Tensor) -> tf.Tensor:
        """Prepand blank index for transducer models"""
        return tf.concat([[self.blank], text], axis=0)

    @abc.abstractclassmethod
    def init_upoints(self):
        """Create a dictionary of upoints that map to target indicies"""
        raise NotImplementedError()

    @abc.abstractclassmethod
    def string2Indices(self, text):
        """Converts a string to a list of indices

        Args:
            text (string): sequence of characters

        Returns:
            (tf.Tensor): sequence of ints
        """
        raise NotImplementedError()

    @abc.abstractclassmethod
    def indices2String(self, indices):
        """Directly converts a list of indices to string

        Note: this function is intended for only debugging purpose

        Args:
            indices (tf.Tensor): [B, None]

        Returns:
            (tf.Tensor): transcripts of dtype tf.string with dim [B]
        """
        raise NotImplementedError()

    @abc.abstractclassmethod
    def indices2upoints(self, indices):
        """Transforms predicted indices to unicode code points

        Args:
            indices (tf.Tensor): indices in shape [None].

        Returns:
            (tf.Tensor): unicode code points transcript with dtype tf.int32 and shape
            [None]

        """
        raise NotImplementedError()

class CharFeaturizer(TextFeaturizer):
    def __init__(
        self,
        vocab_path: str = None,
        blank_at_zero: bool = True,
        lang: str = "en",
        normalize: bool = False,
    ):
        super().__init__()

        self.vocab_path = vocab_path
        self.blank_at_zero = blank_at_zero
        self.lang = lang
        self.normalize = normalize

        self.init_vocabulary()
        self.init_upoints()

    def init_vocabulary(self):
        lines = []
        if self.vocab_path is not None:
            with open(self.vocab_path, "r", encoding="utf-8") as fin:
                lines.extend(fin.readlines())
        else:
            if self.lang == "en":
                lines = en
            elif self.lang == "jp":
                lines = jp
            else:
                raise ValueError("This language is not supported.")

        self.blank = 0 if self.blank_at_zero else None
        self.tokens2indices = {}
        self.tokens = []

        index = 1 if self.blank == 0 else 0
        for line in lines:
            if self.normalize:
                line = self.preprocess_text(line)

            self.tokens2indices[line[0]] = index
            self.tokens.append(line[0])
            index += 1

        if self.blank is None:
            self.blank = len(self.tokens)  # blank not at zero

        self.vocab_array = self.tokens.copy()
        self.tokens.insert(self.blank, "")  # add blank token to tokens
        self.num_classes = len(self.tokens)
        self.tokens = tf.convert_to_tensor(self.tokens, dtype=tf.string)

    def init_upoints(self):
        self.upoints = tf.strings.unicode_decode(self.tokens, "UTF-8").to_tensor(
            shape=[None, 1]
        )

    def string2Indices(self, text: str) -> tf.Tensor:
        if self.normalize:
            text = self.preprocess_text(text)

        text = list(text.strip())  # remove trailing space
        indices = [self.tokens2indices[token] for token in text]
        return tf.convert_to_tensor(indices, dtype=tf.int32)

    def indices2String(self, indices: tf.Tensor) -> tf.Tensor:
        indices = self.normalize_indices(indices)
        tokens = tf.gather_nd(self.tokens, tf.expand_dims(indices, axis=-1))
        with tf.device("/CPU:0"):  # string data is not supported on GPU
            tokens = tf.strings.reduce_join(tokens, axis=-1)
        return tokens

    @tf.function(input_signature=[tf.TensorSpec([None], dtype=tf.int32)])
    def indices2upoints(self, indices: tf.Tensor) -> tf.Tensor:
        indices = self.normalize_indices(indices)
        upoints = tf.gather_nd(self.upoints, tf.expand_dims(indices, axis=-1))
        return tf.gather_nd(upoints, tf.where(tf.not_equal(upoints, 0)))

    @tf.function(input_signature=[tf.TensorSpec([None, None], dtype=tf.int32)])
    def indices2upointsBatch(self, indices: tf.Tensor) -> tf.Tensor:
        """Transforms predicted indices to unicode code points in batch

        Note: Unlike `indices2upoints`, to keep batch dimension, this function does
        not remove blank tokens.

        L: length of predicted tokens

        Args:
            indices: [B, L]

        """
        indices = self.normalize_indices(indices)
        # [B, L, 1]
        upoints = tf.gather_nd(self.upoints, tf.expand_dims(indices, axis=-1))
        # [B, L]
        upoints = tf.squeeze(upoints, axis=-1)
        return upoints

class WordpieceFeaturizer(TextFeaturizer):
    PAD_TOKEN = "<pad>"  # also used as blank token

    def __init__(
        self,
        subwords_path: str,
        subwords_corpus: str = None,
        target_vocab_size: int = 1024,
        max_token_length: int = 4,
        normalize: bool = False,
    ):
        """Text featurizer based on Wordpiece

        TODO: Since Japanese is the top priority as of 2021/3 and japanese asr performs
        better in character (see refs), wordpiece in production is not supported,
        althought it is just fine in training. This needs to be exportable to TFLite!!
        ref: https://github.com/espnet/espnet/pull/326
        ref: http://www.interspeech2020.org/uploadfile/pdf/Mon-2-2-2.pdf

        """
        super().__init__()

        self.target_vocab_size = target_vocab_size
        self.max_token_length = max_token_length
        self.max_corpus_chars = None
        self.reserved_tokens = [self.PAD_TOKEN]

        self.normalize = normalize

        if subwords_path and os.path.exists(subwords_path):
            self.subwords = self.load_from_file(subwords_path)
        else:
            self.subwords = self.build_from_corpus(subwords_corpus, subwords_path)

        self.blank = 0  # subword treats blank as 0
        self.num_classes = self.subwords.vocab_size

        self.init_upoints()

    def init_upoints(self):
        text = [""]  # first one is reserved for blank token
        # start iteration after blank
        for idx in np.arange(1, self.num_classes, dtype=np.int32):
            token = self.subwords.detokenize(tf.constant([[idx]], tf.int32))
            token = token.values.numpy()[0].decode("utf-8")

            # `##_` in wordpiece means, a beginging space should be trimmed, otherwise
            # prepend a space
            if "#" in token:
                token = re.sub("#", "", token)
            else:
                token = " " + token

            text.append(token)
        self.upoints = tf.strings.unicode_decode(text, "UTF-8")
        self.upoints = self.upoints.to_tensor()  # [num_classes, max_token_length]

    def build_from_corpus(self, corpus_files: list = None, output_file: str = None):
        filename = preprocess_paths(output_file)

        def corpus_generator():
            for file in corpus_files:
                with open(file, "r", encoding="utf-8") as f:
                    lines = f.read().splitlines()
                    lines = lines[1:]
                for line in lines:
                    line = line.split("\t")
                    yield line[-1]

        wordpiece.build_from_corpus(
            corpus_generator(),
            output_file_path=filename,
            target_vocab_size=self.target_vocab_size,
            max_token_length=self.max_token_length,
            max_corpus_chars=self.max_corpus_chars,
            reserved_tokens=self.reserved_tokens,
        )

        subwords = wordpiece.WordpieceTokenizer(filename, token_out_type=tf.int32)
        return subwords

    def load_from_file(self, filename: str = None):
        filename = preprocess_paths(filename)
        subwords = wordpiece.WordpieceTokenizer(filename, token_out_type=tf.int32)
        return subwords

    def string2Indices(self, text: str) -> tf.Tensor:
        if self.normalize:
            text = self.preprocess_text(text)

        # remove trailing space
        text = text.strip()
        splitted_text = text.split()
        indices = self.subwords.tokenize(splitted_text)
        # `tokenize` returns ragged tensor
        indices = indices.merge_dims(0, -1)
        return tf.cast(indices, dtype=tf.int32)

    def indices2String(self, indices: tf.Tensor) -> tf.Tensor:
        with tf.device("/CPU:0"):  # string data is not supported on GPU
            indices = self.normalize_indices(indices)
            # returns ragged tensor with one element corresponds to one word
            # eg) ["i'm", "<pad>", "thirsty", "<pad>", ...]
            text = self.subwords.detokenize(indices)
            # ["i'm", "", "thirsty", ""]
            text = tf.strings.regex_replace(text, self.PAD_TOKEN, "")
            # ["i'm  thirsty "] <- double spaces & trailing space
            text = tf.strings.reduce_join(text, separator=" ", axis=-1)
            # ["i'm thirsty "] <- trailing space
            text = tf.strings.regex_replace(text, " +", " ")
            # ["i'm thirsty"]
            text = tf.strings.strip(text)
            return text

    @tf.function(input_signature=[tf.TensorSpec([None], dtype=tf.int32)])
    def indices2upoints(self, indices: tf.Tensor) -> tf.Tensor:
        indices = self.normalize_indices(indices)
        upoints = tf.gather_nd(self.upoints, tf.expand_dims(indices, axis=-1))
        return tf.gather_nd(upoints, tf.where(tf.not_equal(upoints, 0)))

class DeprecatedWordpieceFeaturizer(TextFeaturizer):
    def __init__(
        self,
        subwords_path: str,
        subwords_corpus: str = None,
        target_vocab_size: int = 1024,
        max_token_length: int = 4,
        normalize: bool = False,
    ):
        super().__init__()

        self.target_vocab_size = target_vocab_size
        self.max_token_length = max_token_length
        self.max_corpus_chars = None
        # unlike undeprecated version, tds.deprecated.text.SubwordTextEncoder
        # automatically register blank token to zero index
        self.reserved_tokens = None

        self.normalize = normalize

        if subwords_path and os.path.exists(subwords_path):
            self.subwords = self.load_from_file(subwords_path)
        else:
            self.subwords = self.build_from_corpus(subwords_corpus)
            self.save_to_file(subwords_path)

        self.blank = 0  # subword treats blank as 0
        self.num_classes = self.subwords.vocab_size

        self.init_upoints()

    def init_upoints(self):
        text = [""]  # first one is reserved for blank token
        for idx in np.arange(1, self.num_classes, dtype=np.int32):
            token = self.subwords.decode([idx])
            text.append(token)
        self.upoints = tf.strings.unicode_decode(text, "UTF-8")
        self.upoints = self.upoints.to_tensor()  # [num_classes, max_token_length]

    def build_from_corpus(self, corpus_files: list = None):
        def corpus_generator():
            for file in corpus_files:
                with open(file, "r", encoding="utf-8") as f:
                    lines = f.read().splitlines()
                    lines = lines[1:]
                for line in lines:
                    line = line.split("\t")
                    yield line[-1]

        subwords = tds.deprecated.text.SubwordTextEncoder.build_from_corpus(
            corpus_generator(),
            self.target_vocab_size,
            self.max_token_length,
            self.max_corpus_chars,
            self.reserved_tokens,
        )
        return subwords

    def load_from_file(self, filename: str = None):
        filename = preprocess_paths(filename)
        filename_prefix = os.path.splitext(filename)[0]
        subwords = tds.deprecated.text.SubwordTextEncoder.load_from_file(
            filename_prefix
        )
        return subwords

    def save_to_file(self, filename: str = None):
        filename = preprocess_paths(filename)
        filename_prefix = os.path.splitext(filename)[0]
        return self.subwords.save_to_file(filename_prefix)

    def string2Indices(self, text: str) -> tf.Tensor:
        if self.normalize:
            text = self.preprocess_text(text)

        text = text.strip()  # remove trailing space
        indices = self.subwords.encode(text)
        return tf.convert_to_tensor(indices, dtype=tf.int32)

    def indices2String(self, indices: tf.Tensor) -> tf.Tensor:
        with tf.device("/CPU:0"):  # string data is not supported on GPU
            total = tf.shape(indices)[0]
            batch = tf.constant(0, dtype=tf.int32)
            transcripts = tf.TensorArray(
                dtype=tf.string,
                size=total,
                dynamic_size=False,
                infer_shape=False,
                clear_after_read=False,
                element_shape=tf.TensorShape([]),
            )

            def cond(batch, total, _):
                return tf.less(batch, total)

            def body(batch, total, transcripts):
                upoints = self.indices2upoints(indices[batch])
                transcripts = transcripts.write(
                    batch, tf.strings.unicode_encode(upoints, "UTF-8")
                )
                return batch + 1, total, transcripts

            _, _, transcripts = tf.while_loop(
                cond, body, loop_vars=[batch, total, transcripts]
            )

            return transcripts.stack()

    @tf.function(input_signature=[tf.TensorSpec([None], dtype=tf.int32)])
    def indices2upoints(self, indices: tf.Tensor) -> tf.Tensor:
        indices = self.normalize_indices(indices)
        upoints = tf.gather_nd(self.upoints, tf.expand_dims(indices, axis=-1))
        return tf.gather_nd(upoints, tf.where(tf.not_equal(upoints, 0)))


#python standard library
from time import time as now, sleep

#apetools
from basepollster import BasePollster, CSV_JOIN, ZERO

class DevicePollerEnum(object):
    """
    A Holder of rssi-poller constants
    """
    __slots__ = ()
    rssi = 'RSSI'
    devicepoller = 'devicepoller'
# end class DevicePollerEnum

class DevicePoller(BasePollster):
    """
    A DevicePoller
    """
    def __init__(self, *args, **kwargs):
        super(DevicePoller, self).__init__(*args, **kwargs)
        return

    @property
    def name(self):
        """
        :return: the name to look for in the logs
        """
        if self._name is None:
            self._name = DevicePollerEnum.devicepoller
        return self._name

    @property
    def expression(self):
        """
        This isn't used by the device poller right now

        :return: None
        """
        return self._expression

    def run(self):
        """
        :postcondition: the poller is sending rssi values to the output
        """
        interval = self.interval
        #if self.use_header:
        self.output.write("timestamp,rssi,noise,bitrate\n")
        while True:
            if self.event is not None:
                self.event.wait()
            start_time = now()
            datum = CSV_JOIN.format(self.timestamp(), self.device.poll())
            self.output.writeline(datum)
            self.logger.debug(datum)
            sleep(max(interval - now() + start_time, ZERO))
        return
# end DevicePoller


from typing import Any, Callable, List, Optional, Tuple

import sqlparse  # type:ignore
from sqlparse.sql import Token  # type:ignore
from sqlparse.tokens import Name, Newline, Punctuation  # type:ignore
from sqlparse.utils import remove_quotes  # type:ignore

class LayerSqlFunction:
    """
    A parsed Layer SQL statement
    """

    SUPPORTED_FUNCTION_TRAIN = "train"
    SUPPORTED_FUNCTION_PREDICT = "predict"
    SUPPORTED_FUNCTION_AUTOML = "automl"
    SUPPORTED_FUNCTION_TYPES = [SUPPORTED_FUNCTION_TRAIN, SUPPORTED_FUNCTION_PREDICT, SUPPORTED_FUNCTION_AUTOML]

    def __init__(self, function_type: str, source_name: str, target_name: str) -> None:
        self.function_type = function_type
        self.source_name = source_name
        self.target_name = target_name

    def __repr__(self) -> str:
        return (
            f"<LayerSQL function_type:{self.function_type}"
            + f" source_name:{self.source_name}"
            + f" target_name:{self.target_name}>"
        )

class LayerPredictFunction(LayerSqlFunction):
    def __init__(
        self,
        source_name: str,
        target_name: str,
        model_name: str,
        select_columns: List[str],
        predict_columns: List[str],
        all_columns: List[str],
        prediction_alias: str,
        sql: str,
    ) -> None:
        super().__init__(
            function_type=self.SUPPORTED_FUNCTION_PREDICT, source_name=source_name, target_name=target_name
        )
        self.model_name = model_name
        self.select_columns = select_columns
        self.predict_columns = predict_columns
        self.all_columns = all_columns
        self.prediction_alias = prediction_alias
        self.sql = sql

class LayerTrainFunction(LayerSqlFunction):
    def __init__(
        self,
        source_name: str,
        target_name: str,
        train_columns: List[str],
    ) -> None:
        super().__init__(function_type=self.SUPPORTED_FUNCTION_TRAIN, source_name=source_name, target_name=target_name)
        self.train_columns = train_columns

class LayerAutoMLFunction(LayerSqlFunction):
    def __init__(
        self,
        source_name: str,
        target_name: str,
        model_type: str,
        feature_columns: List[str],
        target_column: str,
        sql: str,
    ) -> None:
        super().__init__(function_type=self.SUPPORTED_FUNCTION_AUTOML, source_name=source_name, target_name=target_name)
        self.feature_columns = feature_columns
        self.target_column = target_column
        self.model_type = model_type
        self.sql = sql

def find_from_token(tokens: List[Token]) -> List[Token]:
    return expect_tokens(tokens, [keyword("from"), lambda x: isinstance(x, sqlparse.sql.Identifier)])

def get_from_where_clause(select_tokens: List[Token]) -> Tuple[str, str]:
    from_token = find_from_token(select_tokens)
    if not from_token:
        raise ValueError("Invalid sql. Missing 'from' clause.")
    source = from_token[0].value
    where_statement = " ".join((x.value for x in from_token[1:]))
    return source, where_statement

def get_cols_from_container(brace_container_token: Token) -> List[str]:
    _, cols_token, _ = brace_container_token.tokens
    if isinstance(cols_token, sqlparse.sql.IdentifierList):
        return [x.value for x in clean_separators(cols_token.tokens)]
    return [cols_token.value]

def build_sql(cols: List[str], source: str, where_statement: str) -> str:
    col_str = ", ".join(cols)
    if len(where_statement) > 0:
        return f"select {col_str} from {source} {where_statement}"
    else:
        return f"select {col_str} from {source}"

class LayerSQLParser:

    ALL_COLUMNS_WILDCARD = "*"

    def parse(self, sql: str) -> Optional[LayerSqlFunction]:
        """
        returns None if not a layer SQL statement
        returns an instance of LayerSQL if a valid layer SQL statement
        """
        parsed = sqlparse.parse(sql)
        if not parsed:
            return None

        tokens = parsed[0].tokens
        layer_func = next((x for x in find_functions(tokens) if is_layer_function(x)), None)
        if not layer_func:
            return None

        target_name_group = next(
            (x for x in expect_tokens(tokens, [keyword("create or replace"), keyword("table"), group()])), None
        )
        if not target_name_group:
            return None
        target_name = self._get_target_name_from_group(target_name_group)

        if is_automl_function(layer_func):
            return self.parse_automl(layer_func, target_name)
        elif is_predict_function(layer_func):
            return self.parse_predict(layer_func, target_name)
        elif is_train_function(layer_func):
            return self.parse_train(layer_func, target_name)
        else:
            raise ValueError(f"Unsupported function: {layer_func.get_name()}")

    def _get_target_name_from_group(self, token: Token) -> str:
        clean_inner_tokens = remove_tokens(token.flatten(), newline())
        target_name_tokens = slice_between_tokens(clean_inner_tokens, name(), whitespace())
        if target_name_tokens:
            return "".join([x.value for x in target_name_tokens])
        else:
            raise ValueError("Invalid sql syntax for Layer function")

    def parse_predict(self, layer_func_token: Token, target: str) -> LayerPredictFunction:

        # We need the parent `select` statement that contains the function
        # to get access to the selected columns and the source relation
        select_sttmt = find_parent(layer_func_token, lambda x: isinstance(x, sqlparse.sql.Parenthesis))
        clean_select_sttmt = []
        if not select_sttmt:
            raise ValueError("SQL syntax error")
        else:
            clean_select_sttmt = clean_separators(select_sttmt.tokens)

        # sqlparse doesn't seem to parse correctly the inner contents of this parenthesis.
        # Here, we rebuild the sql and use parse it again to get the relevant tokens
        inner_sql_text = " ".join(x.value for x in clean_select_sttmt)
        inner_sql = sqlparse.parse(inner_sql_text)[0]
        clean_inner_sql = remove_tokens(inner_sql.tokens, lambda x: x.is_whitespace)

        source, where_statement = get_from_where_clause(clean_inner_sql)

        # extract the selected columns in the query
        columns_incl_layer = []
        select_columns_list = find_token(clean_inner_sql, lambda x: isinstance(x, sqlparse.sql.IdentifierList))
        if select_columns_list is None:
            raise ValueError("SQL syntax error. Select columns are missing.")
        else:
            columns_incl_layer = find_tokens(
                select_columns_list.tokens, lambda x: isinstance(x, sqlparse.sql.Identifier)
            )
        columns = remove_tokens(columns_incl_layer, lambda t: find_layer_function(t) is not None)
        select_columns = [t.value for t in columns]

        # extract the layer prediction function
        clean_func_tokens = clean_separators(layer_func_token[1].tokens)
        if len(clean_func_tokens) < 3:
            invalid_func = " ".join(t.value for t in clean_func_tokens)
            raise ValueError(f"Invalid predict function syntax {invalid_func}")
        model_name, _, bracket_container = clean_func_tokens
        predict_model = remove_quotes(model_name.value)
        predict_columns = get_cols_from_container(bracket_container)

        all_columns = select_columns + list(set(predict_columns) - set(select_columns))
        sql_text = build_sql(all_columns, source, where_statement)
        sql = sqlparse.format(sql_text, keyword_case="lower", strip_whitespace=True)

        prediction_alias = layer_func_token.parent.get_alias() or "prediction"

        return LayerPredictFunction(
            source,
            target,
            predict_model,
            select_columns,
            predict_columns,
            all_columns,
            prediction_alias,
            sql,
        )

    def parse_automl(self, layer_func_token: Token, target: str) -> LayerAutoMLFunction:

        # We need the parent `select` statement that contains the function
        # to get access to the selected columns and the source relation
        select_sttmt = find_parent(layer_func_token, lambda x: isinstance(x, sqlparse.sql.Parenthesis))
        clean_select_sttmt = []
        if not select_sttmt:
            raise ValueError("SQL syntax error")
        else:
            clean_select_sttmt = clean_separators(select_sttmt.tokens)

        # sqlparse doesn't seem to parse correctly the inner contents of this parenthesis.
        # Here, we rebuild the sql and use parse it again to get the relevant tokens
        inner_sql_text = " ".join(x.value for x in clean_select_sttmt)
        inner_sql = sqlparse.parse(inner_sql_text)[0]
        clean_inner_sql = remove_tokens(inner_sql.tokens, lambda x: x.is_whitespace)

        source, where_statement = get_from_where_clause(clean_inner_sql)

        # extract the layer prediction function
        clean_func_tokens = clean_separators(layer_func_token[1].tokens)
        if len(clean_func_tokens) < 4:
            invalid_func = " ".join(t.value for t in clean_func_tokens)
            raise ValueError(f"Invalid automl function syntax {invalid_func}")
        model_type_token, _, bracket_container, target_column_token = clean_func_tokens
        model_type = remove_quotes(model_type_token.value)
        feature_columns = get_cols_from_container(bracket_container)
        target_column = remove_quotes(target_column_token.value)
        all_columns = list(dict.fromkeys((feature_columns + [target_column])))
        sql_text = build_sql(all_columns, source, where_statement)
        sql = sqlparse.format(sql_text, keyword_case="lower", strip_whitespace=True)

        return LayerAutoMLFunction(source, target, model_type, feature_columns, target_column, sql)

    def parse_train(self, layer_func_token: Token, target: str) -> LayerTrainFunction:
        select = find_parent(layer_func_token, lambda x: isinstance(x, sqlparse.sql.Identifier))
        if not select:
            raise ValueError("SQL syntax error")
        else:
            clean_select = clean_separators(select)
            source, _ = get_from_where_clause(layer_func_token.parent.parent.tokens)
            _, train_func = clean_select
            if len(train_func.tokens) < 2:
                invalid_func = "".join(x.value for x in select.flatten())
                raise ValueError(f"Invalid train function syntax {invalid_func}")
            _, parenthesis_group = train_func.tokens
            cols = self.extract_columns(clean_separators(parenthesis_group.tokens))

            return LayerTrainFunction(source, target, cols)

    def extract_columns(self, content_tokens: List[Token]) -> List[str]:
        if len(content_tokens) < 1:
            return [self.ALL_COLUMNS_WILDCARD]
        if content_tokens[0].value.lower() == "array":
            if len(content_tokens) < 2:
                raise ValueError("Invalid train function syntax")
            identifiers = clean_separators(content_tokens[1].tokens)[0]
            return [x.value for x in identifiers.get_identifiers()]
        else:
            return [content_tokens[0].value]

# Functions to assist the SQL Parsing

TokenPredicate = Callable[[Token], bool]

def get_predict_function_name(layer_column_token: Token) -> Optional[str]:
    if layer_column_token.has_alias():
        return layer_column_token.get_name()
    return "predict"

def find_token(tokens: List[Token], pred: TokenPredicate) -> Optional[Token]:
    for token in tokens:
        if pred(token):
            return token
    return None

def find_tokens(tokens: List[Token], pred: TokenPredicate) -> List[Token]:
    return [x for x in tokens if pred(x)]

def remove_tokens(tokens: List[Token], pred: TokenPredicate) -> List[Token]:
    return [x for x in tokens if not pred(x)]

def clean_separators(tokens: List[Token]) -> List[Token]:
    return remove_tokens(tokens, lambda x: x.is_whitespace or x.ttype == Punctuation)

def expect_tokens(tokens: List[Token], predicates: List[TokenPredicate]) -> List[Token]:
    if not tokens or not predicates:
        return tokens
    pred = predicates[0]
    for i, token in enumerate(tokens):
        if pred(token):
            return expect_tokens(tokens[i:], predicates[1:])
    return []

def slice_between_tokens(tokens: List[Token], start: TokenPredicate, end: TokenPredicate) -> Optional[List[Token]]:
    if not tokens:
        return None
    start_index = -1
    end_index = -1
    start_seen = False
    for i, token in enumerate(tokens):
        if not start_seen and start(token):
            start_index = i
            start_seen = True
        if start_seen and end(token):
            end_index = i
            break
    if start_index >= 0 and end_index > 0:
        return tokens[start_index:end_index]
    else:
        return None

def find_functions(elem: Any) -> List[Token]:
    from itertools import chain

    funcs: List[List[Token]] = [[]]
    if isinstance(elem, List):
        funcs = [find_functions(x) for x in elem]
    elif isinstance(elem, sqlparse.sql.Statement):
        funcs = [find_functions(x) for x in elem.tokens]
    elif isinstance(elem, sqlparse.sql.IdentifierList):
        funcs = [find_functions(x) for x in elem.get_identifiers()]
    elif isinstance(elem, sqlparse.sql.Identifier):
        funcs = [find_functions(x) for x in elem.tokens]
    elif isinstance(elem, sqlparse.sql.Parenthesis):
        funcs = [find_functions(x) for x in elem.tokens]
    elif isinstance(elem, sqlparse.sql.Function):
        return [elem]
    return list(chain.from_iterable(funcs))

def find_parent(token: Token, pred: TokenPredicate) -> Optional[Token]:
    parent = token.parent
    if not parent:
        return None
    elif pred(parent):
        return parent
    else:
        return find_parent(parent, pred)

def is_layer_function(func: sqlparse.sql.Function) -> bool:
    parent = func.parent
    return isinstance(parent, sqlparse.sql.Identifier) and parent.tokens[0].value == "layer"

def find_layer_function(tokens: List[Token]) -> Optional[Token]:
    return next((x for x in find_functions(tokens) if is_layer_function(x)), None)

def is_predict_function(func: sqlparse.sql.Function) -> bool:
    return func[0].value.lower() == LayerSqlFunction.SUPPORTED_FUNCTION_PREDICT

def is_automl_function(func: sqlparse.sql.Function) -> bool:
    return func[0].value.lower() == LayerSqlFunction.SUPPORTED_FUNCTION_AUTOML

def is_train_function(func: sqlparse.sql.Function) -> bool:
    return func[0].value.lower() == LayerSqlFunction.SUPPORTED_FUNCTION_TRAIN

def keyword(value: str) -> TokenPredicate:
    def check_token(token: Token) -> bool:
        return token.is_keyword and token.value.lower() == value.lower()

    return check_token

def whitespace() -> Callable[[Token], bool]:
    return lambda x: x.is_whitespace

def name() -> Callable[[Token], bool]:
    return lambda x: x.ttype is Name

def newline() -> Callable[[Token], bool]:
    return lambda x: x.ttype is Newline

def group() -> Callable[[Token], bool]:
    return lambda x: x.is_group


import requests
import hashlib
import subprocess
import os

def main():

    # Get the expected SHA-256 hash value of the VLC installer
    expected_sha256 = get_expected_sha256()

    # Download (but don't save) the VLC installer from the VLC website
    installer_data = download_installer()

    # Verify the integrity of the downloaded VLC installer by comparing the
    # expected and computed SHA-256 hash values
    if installer_ok(installer_data, expected_sha256):

        # Save the downloaded VLC installer to disk
        installer_path = save_installer(installer_data)

        # Silently run the VLC installer
        run_installer(installer_path)

        # Delete the VLC installer from disk
        delete_installer(installer_path)

def get_expected_sha256():
    file_url = 'https://download.videolan.org/pub/videolan/vlc/3.0.17.4/win64/vlc-3.0.17.4-win64.exe.sha256'
    msg = requests.get(file_url)
    # Check whether the download was successful
    if msg.status_code == requests.codes.ok:
    # Extract text file content from response message
        filecontents = msg.text.split()
        hash = filecontents[0]
        return hash

def download_installer():
    file_url = 'https://download.videolan.org/pub/videolan/vlc/3.0.17.4/win64/vlc-3.0.17.4-win64.exe'
    resp_msg = requests.get(file_url)
    # Check whether the download was successful
    if resp_msg.status_code == requests.codes.ok:
    # Extract binary file content from response message body
        execontent = resp_msg.content
    # Calculate SHA-256 hash value
    return execontent

def installer_ok(installer_data, expected_sha256):
    if expected_sha256 == hashlib.sha256(installer_data).hexdigest():
        return True
    else:
        return False

def save_installer(installer_data):
    path = r'C:\Users\TBif\Desktop\vlcsetup.exe'
    with open(path, 'wb') as file:
        file.write(installer_data)
    return path

def run_installer(installer_path):
    subprocess.run([installer_path, '/L=1033', '/S'])
    return

def delete_installer(installer_path):
    os.remove(installer_path)
    return

if __name__ == '__main__':
    main()


from django.conf.urls import patterns, url
from contributors import views

urlpatterns = patterns('',
    url(r'^$', views.main, name='main'),
    url(r'^article/(?P<article_id>\d+)', views.article, name='article'),
    url(r'^contributors', views.contributors, name='contirbutors'),
    url(r'^contributor/(?P<contributor_id>.+)', views.contributor, name='contirbutors'),
    url(r'^library', views.archive, name='archive'),
    url(r'^about', views.about, name='about'),
    url(r'^submitted/', views.submitted, name='submitted'),
    url(r'^mailing/', views.mailinglist, name='mailinglist'),
    url(r'^contact/', views.contact, name='contact'),
)


#Package Imports
import json
import multiprocessing as mp
import os
import random
import glob
import copy
import math
from os.path import abspath, dirname, exists, join
import sys
sys.path.append("./")
from pathlib import Path
from sklearn.model_selection import train_test_split

import ase
from ase import Atoms, Atom
from ase.visualize import view
from ase.io import read, write
from ase.io.trajectory import Trajectory
from ase.visualize.plot import plot_atoms
from ase.constraints import dict2constraint
from ase.calculators.singlepoint import SinglePointCalculator
from ase.db import connect
from ase.neighborlist import NeighborList, natural_cutoffs
from ase.io import extxyz

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

import seaborn as sns
import networkx as nx
import mendeleev
import itertools
import pickle
import shutil
from datetime import datetime

import torch
from torch_geometric.loader import DataLoader
from torch_geometric.utils import from_networkx
from torch_geometric.data import Data
from torch_geometric.utils import to_networkx

from pymatgen.io.ase import AseAtomsAdaptor
from scipy.spatial import distance_matrix

class GaussianDistance():
    """
    Expands the distance by Gaussian basis.
    Unit: angstrom
    """

    def __init__(self, dmin=0.0, dmax=6, step=0.2, var=None):
        """
        Args:
            dmin (float): Minimum interatomic distance
            dmax (float): Maximum interatomic distance
            step (float): Step size for the Gaussian filter
            var (float, optional): Variance of Gaussian basis. Defaults to step if not given
        """
        assert dmin < dmax
        assert dmax - dmin > step

        self.filter = np.arange(dmin, dmax + step, step)
        self.embedding_size = len(self.filter)

        if var is None:
            var = step

        self.var = var

    def expand(self, distances):
        """Apply Gaussian distance filter to a numpy distance array
        Args:
            distances (ArrayLike): A distance matrix of any shape
        Returns:
            Expanded distance matrix with the last dimension of length
            len(self.filter)
        """
        distances = np.array(distances)

        return np.exp(-((distances[..., np.newaxis] - self.filter) ** 2) / self.var ** 2)

    def get_layer(self,slb, thresh = 1):
        """
        * slb (atoms object): The system, WITHOUT adsorbate attached
        * thresh (float): The threshold in Angstroms of the comparison between z heights
        * layer_remapped (list of floats): An ordered list of the layer heights identified in the system
        * num_layers (int): # of layers in system
        * num_atoms (int): # of atoms in system
        * layer_ID_Indices (list of lists): A list of lists of atom ID #s corresponding to the atoms placed in each layer classification.
          NOTE: layer_ID_Indices[0] contains the ID for atoms in the ASE system that were identified as belonging to the height in layer_remapped[0]

        """

        layer_ID_Indices = []
        similar_idx = []
        num_atoms = len(slb)
        slab_z_list= slb.positions[:,2].tolist()
        unique_z = np.unique(slab_z_list).tolist()

        # For the unique heights in the system...
        # loop, and see if any other entries are close
        # If they are, save the indices of these redundant heights so we can eliminate them
        for h1 in unique_z:
            for h2 in unique_z:
                if h1 != h2 and unique_z.index(h2) > unique_z.index(h1) and math.isclose(h1, h2, abs_tol = thresh):
                    similar_idx.append(h2)

        # This gets all the distinct z values together by removing the entries in the similar list
        layer_z = [unique_z[i] for i in range(len(unique_z)) if unique_z[i] not in similar_idx]
        # This is how many layers we think we have
        num_layers = len(layer_z)
        # This is the set of layer heights we found
        layer_heights = layer_z

        # Find the indices present in each layer we have found
        slotted_indices = []
        for i in range(num_layers):
            layer_ID_Indices.append([])
            for j in range(num_atoms):
                if math.isclose(slab_z_list[j], layer_heights[i], abs_tol = 1) == True and j not in slotted_indices:
                    layer_ID_Indices[i].append(j)
                    slotted_indices.append(j)
        # This next section organizes the indices so that we have each list of indices for a layer ordered in terms of layer ascending/descending instead of randomly
        layer_remapped = copy.deepcopy(layer_heights)
        layer_remapped.sort()
        ordering = [layer_heights.index(layers) for layers in layer_remapped]
        #print("Ordering", ordering)
        layer_ID_Indices = [layer_ID_Indices[i] for i in ordering]

        return layer_remapped, ordering, layer_ID_Indices

#This class handle model/data imports
class ASE_To_Graph():
    def __init__(self,db_file,embedding):
        '''
        inputs:
        db_file (str) - filename/path of the ASE db file
        embedding (str) - name of the embedding you want to implement
        '''

        self.db_filename = db_file
        self.embedding = embedding

    def import_db(self):
        '''
        self explanatory function
        '''

        with connect(self.db_filename) as db:
            ase_db = db
        self.db = ase_db
        return ase_db

    #Not used in algorithm, just helps with data inspection
    def bond_symbol(self,atoms, a1, a2):
        '''
        self explanatory function
        '''

        return "{}{}".format(*sorted((atoms[a1].symbol, atoms[a2].symbol)))

    #Get electronegativity (EN)
    def get_EN(self,atoms, a1, a2):
        '''
        self explanatory function
        '''

        try:
            avg = (mendeleev.element(atoms[a1].symbol).en_pauling + mendeleev.element(atoms[a2].symbol).en_pauling)/2
        except:
            avg = 0.0
        try:
            dif = abs(mendeleev.element(atoms[a1].symbol).en_pauling - mendeleev.element(atoms[a2].symbol).en_pauling)
        except:
            dif = 0
        return np.array([avg, dif])

    def node_symbol(self,atom):
        '''
        self explanatory function
        '''

        return "".format(atom.symbol)

    def add_atoms_node(self,graph, atoms, a1):
        '''
        self explanatory function
        '''

        graph.add_node(self.node_symbol(atoms[a1]), index=a1)

    def add_atoms_edge(self,graph, atoms, a1, a2):
        '''
        self explanatory function
        '''

        EN = self.get_EN(atoms, a1, a2)
        distance = atoms.get_distances(a1,a2)
        graph.add_edge(a1,
                       a2,
                       EN= EN,
                       distance = distance)

    #Main surface analysis automation function
    def analyse_surface(self,atoms, adsorbate,radii_multiplier=1.1, skin = 0.25, **kwargs):
        '''
        inputs:
        atoms (ASE atoms) - atom to analyze
        adsorbate (str) - adsorbate to analyze
        radii_multiplier (float) - atomic radius factor
        skin (float) - skin factor

        outputs:
        binding_atoms (list) - binding atom index
        adsorbate_atoms (list) - adsorbate atom index
        bonds (list) - connections
        '''

        #Extract nearest neighbor indexes
        nl = NeighborList(natural_cutoffs(atoms, radii_multiplier), self_interaction=False,
                              bothways=True, skin=skin)

        nl.update(atoms)
        #Extract index of adsorbate atoms
        adsorbate_atoms = [atom.index for atom in atoms if atom.symbol in adsorbate]

        # Add all atoms to graph
        distances = atoms.get_all_distances(mic=True)

        binding_atoms = []
        bonds = []
        for index, atom in enumerate(atoms):
            neighbors, offsets = nl.get_neighbors(index)
            for neighbor, offset in zip(neighbors, offsets):
                if np.all(offset==np.zeros(3)):
                    if distances[index][neighbor] > 2.5 and (bool(index in adsorbate_atoms) ^ bool(neighbor in adsorbate_atoms)):

                        continue
                    if (bool(index not in adsorbate_atoms) ^ bool(neighbor not in adsorbate_atoms)) and (index not in adsorbate_atoms):
                        binding_atoms.append(index)

                bonds.append((index, neighbor))

        return binding_atoms, adsorbate_atoms, bonds

    def extract_ads_geom(self, atms, adsorbate, elem_features, plot=False):
        '''
        inputs:
        atms (ASE Atoms) - slab to analyze
        adsorbate (str) - adsorbate to analyze
        elem_features (str) - embedding features initialized with the class
        plot (bool) - plot the exactracted structure

        outputs:
        surf_n (ASE atoms) - extracted surface
        full (nx graph) - resulting graph
        '''

        #Extract neighbors
        binding_atoms, adsorbate_atoms, bonds = self.analyse_surface(atms, adsorbate=adsorbate)
        nl = NeighborList(natural_cutoffs(atms, 1.1), self_interaction=False,
                                  bothways=True, skin=0.25)
        nl.update(atms)
        poss = [atms.positions[i] for i in binding_atoms]
        symbols = [atms[i].symbol for i in binding_atoms]
        distances =atms.get_all_distances(mic=True)
        #extract coordiantes of atoms
        for atm in  binding_atoms:
            neighbors, offsets = nl.get_neighbors(atm)
            for i, offset in zip(neighbors, offsets):
                poss.append(atms.positions[i] + np.dot(offset, atms.get_cell()))
                symbols.append(atms[i].symbol)
        #extract ads atoms coords
        for atm in  adsorbate_atoms:
            poss.append(atms.positions[atm])
            symbols.append(atms[atm].symbol)

        coords = np.vstack(poss)

        #generate extracted neighbors structure
        surf_n = Atoms(symbols,positions=coords )
        ase.geometry.get_duplicate_atoms(surf_n, cutoff=0.1, delete=True)
        if plot:
            fig, ax = plt.subplots(1,2)
            plot_atoms(surf_n,ax[0],radii=0.9, rotation=('0x,0y,0z'))
            plot_atoms(surf_n,ax[1],radii=0.9, rotation=('-90x,45y,0z'))

        params = {"binding_atoms":binding_atoms,
                "adsorbate_atoms":adsorbate_atoms,
                "bonds":bonds}

        adsorbate_atms = [atom.index for atom in surf_n if atom.symbol in adsorbate]

        #Generate Graph
        full = nx.Graph()
        full.add_nodes_from(range(len(surf_n)))
        distances = surf_n.get_all_distances(mic=True)

        numbers = surf_n.get_atomic_numbers().tolist()
        symbols = surf_n.get_chemical_symbols()

        for i, feat_dict in full.nodes(data=True):
            feat_dict.update({'symbol':  symbols[i]})
            feat_dict.update({'atomic_number':  numbers[i]})
            feat_dict.update({'node_attr':  np.array(elem_features[symbols[i]]).astype(np.float32)})
        comb = list(itertools.combinations([i for i in range(len(surf_n))],2))

        for atm1, atm2 in comb:
            if atm1==atm2:
                continue
            if distances[atm1][atm2] > 2.8 :
                continue

            if distances[atm1][atm2] > 2.3 and (bool(atm1 in adsorbate_atms) ^ bool(atm2 in adsorbate_atms)):
                continue
            if distances[atm1][atm2] > 2.3 and (bool(atm1 in adsorbate_atms) and bool(atm2 in adsorbate_atms)):
                continue
            self.add_atoms_edge(full, surf_n,  atm1, atm2)
        return surf_n,  full #params

    #convert a ASE atoms object from an ASE db into a graph
    def get_data(self,row, atom_features,map_warning=False):
        '''
        inputs:
        row (ASE rows) - row to convert into a torch graph
        atom_features (str) - embedding features
        map_warning (bool) - print warnings about failed mappings
        '''

        gdf = GaussianDistance()
        adsorbate=row.Adsorbate
        atm = row.toatoms()
        slabId = row.Hea_ID
        ads_energy = row.Energy
        #embed graph object with features
        try:
            atms, G = self.extract_ads_geom(atms=atm, adsorbate=list(set(list(adsorbate))), elem_features=atom_features)
            data =from_networkx(G)

            data.x = data.node_attr
            edge_attr = data.EN[:,0]*data.distance[:,0]
            edge_attr.unsqueeze(-1)
            edge_attr = gdf.expand(edge_attr)
            edge_attr = torch.tensor(edge_attr)
            data.edge_attr = edge_attr
            data.__setattr__('idx',slabId)
            data.__setattr__('adsorbate',adsorbate)
            data.__setattr__('postions',torch.from_numpy(atms.positions))
            data.y = torch.tensor([ads_energy])

            try:
                atm1 = copy.deepcopy(atms)
                del(atm1[-len(list(adsorbate)):])
                atm1.center()
                _,_, layers = gdf.get_layer( atm1, thresh = 1)
                surface = layers[-1]
                subsurface = layers[-2]
                data.__setattr__('surface',[ atm1.symbols[i] for i in surface])
                data.__setattr__('subsurface',[ atm1.symbols[i] for i in subsurface])
            #Throw error
            except Exception as er1:
                if map_warning:
                    print(f" Warnn ing!!! surface and subsurface atoms not found in ID[{slabId}]")
                data.__setattr__('surface',[])
                data.__setattr__('subsurface',[])
            return data
        except Exception as err:
            print(f"{err}")
            return None

    def get_data_from_asedb(self, db,  adsorbate, emb="cgcnn92", ncpus=1):
        '''
        inputs:
        db (ASE db) - db to fully extract
        adsorbate (str) - adsorbate to analyze/extract
        emb (str) - embedding features
        ncpus (int) - core cpus
        '''

        gdf = GaussianDistance()
        rows = [row for row in db.select('Energy', Adsorbate=adsorbate) ]

        elem_emb = join(os.path.dirname(os.path.realpath("__file__")), f"element/{emb}.json")
        with open(elem_emb) as f:
                atom_features = json.load(f)

        if len(rows) > 100:
            print(f'Warning, {len(rows)} ASE atoms will likely take {len(rows)/10}+ min to complete')

        results = []
        for i in range(len(rows)):
            if i==0:
                results.append(self.get_data( rows[i],atom_features))
            elif i%20==0:
                print(f'{i}/{len(rows)} Graphs generated')
            else:
                results.append(self.get_data( rows[i],atom_features))

        datas = results

        with open(f"data_{adsorbate}_{emb}.pickle", 'wb') as handle:
            pickle.dump(datas, handle)

        self.adsorbate = adsorbate
        self.emb = emb
        self.pickle_name = f"data_{adsorbate}_{emb}.pickle"

    #Provide a sample visual of AGRA extraction
    def sample_visual(self):
        with open('element/cgcnn92.json') as f:
            atom_features = json.load(f)

        slab_sample_visual = self.db.get_atoms(1)
        extraction_sample_visual, g = self.extract_ads_geom(slab_sample_visual,
                                                           adsorbate=['O','C'],
                                                           elem_features=atom_features)

        ocp_a2g = OCP_AtomsToGraphs(
            max_neigh=50,
            radius=6,
            r_energy=False,
            r_forces=False,
            r_distances=True,
            r_edges=True,
            r_fixed=True,)

        ocp_slab = ocp_a2g.convert(slab_sample_visual)
        ocp_graph = to_networkx(ocp_slab)

        fig, ax = plt.subplots(2,2,figsize=(8,6),subplot_kw={'aspect': 'equal'})

        gs = gridspec.GridSpec(2, 2, width_ratios=[1, 1], height_ratios=[1, 1])
        ax[0, 0] = plt.subplot(gs[0, 0])
        ax[0, 1] = plt.subplot(gs[0, 1])
        ax[1, 0] = plt.subplot(gs[1, 0])
        ax[1, 1] = plt.subplot(gs[1, 1])

        plot_atoms(slab_sample_visual,ax[0,0],radii=0.9, rotation=('0x,0y,0z'))
        ax[0,0].set_title("ASE db Extracted Structure")
        ax[0,0].set_xticks([])
        ax[0,0].set_yticks([])

        plot_atoms(extraction_sample_visual,ax[0,1],radii=0.9, rotation=('0x,0y,0z'))
        ax[0,1].set_title("AGRA Refined Structure")
        ax[0,1].set_xticks([])
        ax[0,1].set_yticks([])

        nx.draw_networkx(g,with_labels=True,ax=ax[1,1])
        ax[1,1].set_title("Resulting AGRA Graph")

        nx.draw_networkx(ocp_graph,with_labels=True,ax=ax[1,0])
        ax[1,0].set_title("Resulting OCP Graph")

        plt.tight_layout()
        return None

    #split generated graph data into train/test data
    def split_data(self,split=0.1):
        '''
        inputs:
        split (float) - train/test split

        outputs:
        train/test (np.array) - train test datasets
        '''

        with open(self.pickle_name, 'rb') as handle:
            dataset= pickle.load(handle)

        train,test,ytrain,ytest = train_test_split(dataset,dataset,test_size=split)

        with open(f"data_{self.adsorbate}_Train_{self.emb}.pickle", 'wb') as handle:
            pickle.dump(train, handle)
        with open(f"data_{self.adsorbate}_Test_{self.emb}.pickle", 'wb') as handle:
            pickle.dump(test, handle)
        return train,test

#This entire class is copied over from the OCP github solely for the purpose of graph generation comparison (https://github.com/Open-Catalyst-Project/ocp)
class OCP_AtomsToGraphs:
    def __init__(
        self,
        max_neigh=200,
        radius=6,
        r_energy=False,
        r_forces=False,
        r_distances=False,
        r_edges=True,
        r_fixed=True,
        r_pbc=False,
    ):
        self.max_neigh = max_neigh
        self.radius = radius
        self.r_energy = r_energy
        self.r_forces = r_forces
        self.r_distances = r_distances
        self.r_fixed = r_fixed
        self.r_edges = r_edges
        self.r_pbc = r_pbc

    def _get_neighbors_pymatgen(self, atoms):
        """Preforms nearest neighbor search and returns edge index, distances,
        and cell offsets"""
        struct = AseAtomsAdaptor.get_structure(atoms)
        _c_index, _n_index, _offsets, n_distance = struct.get_neighbor_list(
            r=self.radius, numerical_tol=0, exclude_self=True
        )

        _nonmax_idx = []
        for i in range(len(atoms)):
            idx_i = (_c_index == i).nonzero()[0]
            # sort neighbors by distance, remove edges larger than max_neighbors
            idx_sorted = np.argsort(n_distance[idx_i])[: self.max_neigh]
            _nonmax_idx.append(idx_i[idx_sorted])
        _nonmax_idx = np.concatenate(_nonmax_idx)

        _c_index = _c_index[_nonmax_idx]
        _n_index = _n_index[_nonmax_idx]
        n_distance = n_distance[_nonmax_idx]
        _offsets = _offsets[_nonmax_idx]

        return _c_index, _n_index, n_distance, _offsets

    def _reshape_features(self, c_index, n_index, n_distance, offsets):
        """Stack center and neighbor index and reshapes distances,
        takes in np.arrays and returns torch tensors"""
        edge_index = torch.LongTensor(np.vstack((n_index, c_index)))
        edge_distances = torch.FloatTensor(n_distance)
        cell_offsets = torch.LongTensor(offsets)

        # remove distances smaller than a tolerance ~ 0. The small tolerance is
        # needed to correct for pymatgen's neighbor_list returning self atoms
        # in a few edge cases.
        nonzero = torch.where(edge_distances >= 1e-8)[0]
        edge_index = edge_index[:, nonzero]
        edge_distances = edge_distances[nonzero]
        cell_offsets = cell_offsets[nonzero]

        return edge_index, edge_distances, cell_offsets

    def convert(
        self,
        atoms,
    ):
        """Convert a single atomic stucture to a graph.
        Args:
            atoms (ase.atoms.Atoms): An ASE atoms object.
        Returns:
            data (torch_geometric.data.Data): A torch geometic data object with positions, atomic_numbers, tags,
            and optionally, energy, forces, distances, edges, and periodic boundary conditions.
            Optional properties can included by setting r_property=True when constructing the class.
        """

        # set the atomic numbers, positions, and cell
        atomic_numbers = torch.Tensor(atoms.get_atomic_numbers())
        positions = torch.Tensor(atoms.get_positions())
        cell = torch.Tensor(atoms.get_cell()).view(1, 3, 3)
        natoms = positions.shape[0]
        # initialized to torch.zeros(natoms) if tags missing.
        # https://wiki.fysik.dtu.dk/ase/_modules/ase/atoms.html#Atoms.get_tags
        tags = torch.Tensor(atoms.get_tags())

        # put the minimum data in torch geometric data object
        data = Data(
            cell=cell,
            pos=positions,
            atomic_numbers=atomic_numbers,
            natoms=natoms,
            tags=tags,
        )

        # optionally include other properties
        if self.r_edges:
            # run internal functions to get padded indices and distances
            split_idx_dist = self._get_neighbors_pymatgen(atoms)
            edge_index, edge_distances, cell_offsets = self._reshape_features(
                *split_idx_dist
            )

            data.edge_index = edge_index
            data.cell_offsets = cell_offsets
        if self.r_energy:
            energy = atoms.get_potential_energy(apply_constraint=False)
            data.y = energy
        if self.r_forces:
            forces = torch.Tensor(atoms.get_forces(apply_constraint=False))
            data.force = forces
        if self.r_distances and self.r_edges:
            data.distances = edge_distances
        if self.r_fixed:
            fixed_idx = torch.zeros(natoms)
            if hasattr(atoms, "constraints"):
                from ase.constraints import FixAtoms

                for constraint in atoms.constraints:
                    if isinstance(constraint, FixAtoms):
                        fixed_idx[constraint.index] = 1
            data.fixed = fixed_idx
        if self.r_pbc:
            data.pbc = torch.tensor(atoms.pbc)

        return data

    def convert_all(
        self,
        atoms_collection,
        processed_file_path=None,
        collate_and_save=False,
        disable_tqdm=False,
    ):
        """Convert all atoms objects in a list or in an ase.db to graphs.
        Args:
            atoms_collection (list of ase.atoms.Atoms or ase.db.sqlite.SQLite3Database):
            Either a list of ASE atoms objects or an ASE database.
            processed_file_path (str):
            A string of the path to where the processed file will be written. Default is None.
            collate_and_save (bool): A boolean to collate and save or not. Default is False, so will not write a file.
        Returns:
            data_list (list of torch_geometric.data.Data):
            A list of torch geometric data objects containing molecular graph info and properties.
        """

        # list for all data
        data_list = []
        if isinstance(atoms_collection, list):
            atoms_iter = atoms_collection

        elif isinstance(atoms_collection, ase.db.sqlite.SQLite3Database):
            atoms_iter = atoms_collection.select()

        elif isinstance(
            atoms_collection, ase.io.trajectory.SlicedTrajectory
        ) or isinstance(atoms_collection, ase.io.trajectory.TrajectoryReader):
            atoms_iter = atoms_collection
        else:
            raise NotImplementedError

        for atoms in tqdm(
            atoms_iter,
            desc="converting ASE atoms collection to graphs",
            total=len(atoms_collection),
            unit=" systems",
            disable=disable_tqdm,
        ):
            # check if atoms is an ASE Atoms object this for the ase.db case
            if not isinstance(atoms, ase.atoms.Atoms):
                atoms = atoms.toatoms()
            data = self.convert(atoms)
            data_list.append(data)

        if collate_and_save:
            data, slices = collate(data_list)
            torch.save((data, slices), processed_file_path)

        return data_list


# [ 분석 시 고려사항 5. 최적의 매개변수 조합 (grid search) : train/val/test ]
# - 변수의 최적의 조합을 찾는 과정
# - 중첩 for문으로 구현 가능, grid search 기법으로 간단히 구현 가능
# - train/validation/test set으로 분리
# - 매개변수의 선택은 validation set으로 평가

# 예제) grid search - knn iris data
# 1) data loading
# 2) data split
trainval_x, test_x, trainval_y, test_y = train_test_split(df_iris.data,
                                                          df_iris.target,
                                                          random_state=0)

train_x, val_x, train_y, val_y = train_test_split(trainval_x,
                                                  trainval_y,
                                                  random_state=0)

# 3) 모델 학습 및 매개변수 튜닝
best_score = 0
for i in range(1,11) :
    m_knn = knn_c(i)
    m_knn.fit(train_x, train_y)
    vscore = m_knn.score(val_x, val_y)

    if vscore > best_score :
        best_score = vscore    # best_score의 갱신
        best_params = i        # best parameter의 저장

# 4) 매개변수 고정 후 재학습 및 평가
m_knn = knn_c(best_params)
m_knn.fit(trainval_x, trainval_y)
m_knn.score(test_x, test_y)    # 0.973, best parameter = 5
# --------------------------------------------------------------------------- #

# -------------------------------- 연 습 문 제 -------------------------------- #
# 연습문제 6. grid search - random forest cancer data set
# 1) data loading
from sklearn.datasets import load_breast_cancer
df_cancer = load_breast_cancer()
df_cancer.data.shape    # (569, 30)

# 2) data split
trainval_x, test_x, trainval_y, test_y = train_test_split(df_cancer.data,
                                                          df_cancer.target,
                                                          random_state = 0)

train_x, val_x, train_y, val_y = train_test_split(trainval_x,
                                                  trainval_y,
                                                  random_state = 0)

# 3) 모델 학습 및 매개변수 튜닝 (min_samples_split, max_features)
# min_samples_split : 2 ~ 10
# max_features : 1 ~ 30
# 3-1) 교차 검증 수행 x => 9 * 30
best_score = 0
for i in range(2, 11) :
    for j in range(1, 31) :
        m_rf = rf_c(min_samples_split = i, max_features = j, random_state = 0)
        m_rf.fit(train_x, train_y)
        vscore = m_rf.score(val_x, val_y)

        if vscore > best_score :
            best_score = vscore
            best_params = {'min_samples_split' : i,
                           'max_features' : j}

m_rf = rf_c(**best_params, random_state = 0)    # ** 사용하면 키워드 인자 넣을 수 있음
m_rf.fit(trainval_x, trainval_y)
m_rf.score(test_x, test_y)    # 0.972

# 3-2) 교차 검증 수행 (매개변수 튜닝 시) => 9 * 30 * 5
best_score = 0
for i in range(2, 11) :
    for j in range(1, 31) :
        m_rf = rf_c(min_samples_split = i, max_features = j, random_state = 0)
        m_rf.fit(train_x, train_y)
        cr_vscore = cross_val_score(m_rf, trainval_x, trainval_y, cv = 5)

        vscore = cr_vscore.mean()
        if vscore > best_score :
            best_score = vscore
            best_params = {'min_samples_split' : i,
                           'max_features' : j}

m_rf = rf_c(**best_params, random_state = 0)    # ** 사용하면 키워드 인자 넣을 수 있음
m_rf.fit(trainval_x, trainval_y)
m_rf.score(test_x, test_y)    # 0.944

# grid search 기법
# - 위의 중첩 for문을 사용한 매개변수의 조합을 찾는 과정을 함수화
# - CV 기법을 포함시켜 validation data set에 대한 반복 test를 수행

from sklearn.model_selection import GridSearchCV

# 1) 모델 생성
m_rf = rf_c()

# 2) 그리드 서치 기법을 통한 매개변수 조합 찾기
# 2-1) 매개변수 조합 생성
v_params = {'min_samples_split' : np.arange(2, 11),
            'max_features' : np.arange(1, 31)}

# 2-2) 그리드 서치 모델 생성
m_grid = GridSearchCV(m_rf,           # 적용 모델
                      v_params,       # 매개변수 조합 (딕셔너리)
                      cv = 5)

# 2-3) 그리드 서치에 의한 모델 학습
m_grid.fit(trainval_x, trainval_y)

# 2-4) 결과 확인
dir(m_grid)
m_grid.best_score_    # 베스트 매개변수 값을 갖는 평가 점수
m_grid.best_params_   # {'max_features': 1, 'min_samples_split': 2}

df_result = DataFrame(m_grid.cv_results_)
df_result.T    # 여기서 test_score들은 우리가 아는 test score가 아니고, validation score라고 생각하면 됨

df_result.T.iloc[:, 0]    # 첫 번째 매개변수 셋 결과

# cv = 5이기 때문에 아래처럼 5회 진행
# split0_test_score                                              0.94186
# split1_test_score                                             0.952941
# split2_test_score                                             0.929412
# split3_test_score                                             0.988235
# split4_test_score                                             0.988235

# 2-5) 최종 평가
m_grid.score(test_x, test_y)    # 0.951 => 모델 자체에 best params가 저장되어 있음

# 2-6) 그리드 서치 결과 시각화
df_result.mean_test_score    # 교차 검증의 결과 (5개의 점수에 대한 평균)
arr_score = np.array(df_result.mean_test_score).reshape(30, 9)
# reshape(max_features, min_samples) => reshape은 행우선 순위로 배치 되기 때문에
df_result.loc[:, ['params', 'mean_test_score']]

# =============================================================================
# # [ 참고 : df_result의 결과 reshape 시 배치 순서 ]
# #                             min_samples_split
# #                         2          3          4          5      ...     10
# # 'max_features': 1   0.960137   0.953133   0.960109   0.955486
# # 'max_features': 2
# #                     ....
# # 'max_features': 30
# =============================================================================

import mglearn
plt.rc('figure', figsize = (10, 10))
plt.rc('font', size = 7)

mglearn.tools.heatmap(value,        # 숫자 배열
                      xlabel,       # x축 이름
                      ylabel,       # y축 이름
                      xticklabels,  # x축 눈금
                      yticklabels)  # y축 눈금

mglearn.tools.heatmap(arr_score,
                      'min_samples_split',
                      'max_features',
                      v_params['min_samples_split'],
                      v_params['max_features'],
                      cmap = 'viridis')
# --------------------------------------------------------------------------- #


def divisors(a):
    a = abs(a)
    n = a
    if n == 1:
        return {1}
    res = {1, n}
    i = 2
    while i < n:
        if a % i == 0:
            res.add(i)
            temp = a // i
            res.add(temp)
            n = temp
        i += 1
    return sorted(res)

def isprime(a):
    return len(a) == 2

# Expensive
def primefact(a):
    res = dict()
    n = a
    while True:
        div = divisors(n)
        if isprime(div):
            if res.get(n) is not None:
                res[n] += 1
            else:
                res[n] = 1
            break
        for i in div:
            if isprime(divisors(i)):  # i is prime
                if res.get(i) is not None:
                    res[i] += 1
                else:
                    res[i] = 1
                n = n // i
        if n <= 1:
            break
    return res

# bullet point 1
print(divisors(147))
print(divisors(13195))

# bullet point 2
print(primefact(147))
print(primefact(13195))


# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def lowestCommonAncestor(self, root: 'TreeNode', p: 'TreeNode', q: 'TreeNode') -> 'TreeNode':
        curr = root
        # if p.val < root.val and q.val < root.val
        # root needs to go one level further
        while curr:
            if p.val < curr.val and q.val < curr.val:
                curr = curr.left
            elif p.val > curr.val and q.val > curr.val:
                curr = curr.right
            else:
                break
        return curr


from algorithm import Algorithm
from function import Function

def ask_for_number(question):
    try:
        point = int(input(question))
    except:
        print("Czy aby na pewno podałeś numer?")

    return point

def ask_for_float(question):
    try:
        point = float(input(question))
    except:
        print("Czy aby na pewno podałeś numer?")

    return point

def ask_for_function():
    print("Najpierw podaj wspolczynniki funkcji.")
    a = ask_for_number("a = ")
    b = ask_for_number("b = ")
    c = ask_for_number("c = ")
    d = ask_for_number("d = ")
    function = Function(a, b, c, d)
    return function

def main():
    print("Witaj w algorytmie genetycznym!")
    function = ask_for_function()
    chromosomes = ask_for_number("Podaj liczbe chromosomow: ")
    iterations = ask_for_number("Podaj liczbe iteracji do przeprowadzenia: ")
    crossing_rate = ask_for_float("Podaj wspolczynnik krzyzowania: ")
    mutation_rate = ask_for_float("Podaj wspolczynnik mutacji: ")
    algorithm = Algorithm(chromosomes, iterations, function, crossing_rate, mutation_rate)
    algorithm.generate()

if __name__ == "__main__":
    main()


import torch.nn as nn

class Discriminator(nn.Module):
    def __init__(self, img_dim):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(in_features=img_dim, out_features=128),
            nn.LeakyReLU(0.01),
            nn.Linear(128, 1),
            nn.Sigmoid(),
        )

    def forward(self, x):
        return self.layers.forward(x)

class Generator(nn.Module):
    def __init__(self, noise_dim, img_dim):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(noise_dim, 256),
            nn.LeakyReLU(0.01),
            nn.Linear(256, img_dim),
            nn.Tanh(),
        )

    def forward(self, x):
        return self.layers.forward(x)


from unittest.mock import patch

import pytest
from urllib3.response import HTTPResponse

from stravabot.clients import mapbox

@patch("stravabot.clients.mapbox.MAPBOX_ACCESS_TOKEN", "mapbox-access-token")
def test_static_image(requests_mock):
    expected_response = HTTPResponse(body="the-raw-response", status=200)
    requests_mock.get(
        "https://api.mapbox.com/styles/v1/mapbox/light-v10/static/the-overlay/auto/123x456",
        raw=expected_response,
    )
    response = mapbox.static_image("the-overlay", 123, 456)

    assert requests_mock.call_count == 1
    assert requests_mock.last_request.qs == {"access_token": ["mapbox-access-token"]}
    assert response == expected_response

@patch("stravabot.clients.mapbox.MAPBOX_ACCESS_TOKEN", "mapbox-access-token")
def test_static_image_raises_on_non_200(requests_mock):
    expected_response = HTTPResponse(body="the-raw-response", status=500)
    requests_mock.get(
        "https://api.mapbox.com/styles/v1/mapbox/light-v10/static/the-overlay/auto/123x456",
        raw=expected_response,
    )
    with pytest.raises(mapbox.MapboxError):
        mapbox.static_image("the-overlay", 123, 456)


import argparse
import socket
import time

parser = argparse.ArgumentParser(description="UDP Client")
parser.add_argument(
    "-l",
    "--listen",
    type=str,
    default="127.0.0.1",
    help="IP to check connectivity",
)
parser.add_argument(
    "-p",
    "--port",
    type=int,
    default=1234,
    help="Port to check connectivity",
)

args = parser.parse_args()
ip = args.listen
port = args.port

print(f"Check udp on {ip}:{port}")

while True:
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    client_socket.settimeout(1.0)
    message = b"PING"
    addr = (ip, port)

    start = time.time()
    client_socket.sendto(message, addr)
    try:
        data, server = client_socket.recvfrom(1024)
        end = time.time()
        elapsed = end - start
        print(f"{data} {elapsed:.6f}s")
    except socket.timeout:
        print("REQUEST TIMED OUT")
    time.sleep(1)


#!/usr/bin/env python
#

from cliff.helpers.system import System
import scipy
from scipy import stats
from scipy.spatial.distance import pdist, cdist, squareform
import numpy as np
import logging
import pickle
import cliff.helpers.constants as constants
import cliff.helpers.utils as utils
import math
import os
import copy
import time
import glob

import cliff.tests as t
testpath = os.path.abspath(t.__file__).split('__init__')[0]
# qml
import qml
from qml.representations import get_slatm_mbtypes

class Multipole:
    '''
    Multipole class. Predicts multipoles from machine learning.
    No local axis system. Instead, basis set expansion along the pairwise vectors.
    '''

    def __init__(self, options, ref = None):

        name = options.name
        # Set logger
        self.logger = options.logger
        self.ref = ref
        self.multipoles = None
        self.descr_train  = {'H':[], 'C':[], 'O':[], 'N':[], 'S':[], 'Cl':[], 'F':[], 'Br':[]}
        self.target_train = {'H':[], 'C':[], 'O':[], 'N':[], 'S':[], 'Cl':[], 'F':[], 'Br':[]}
        # support vector regression
        self.clf = None
        # alpha_train has size 1,3,9
        self.max_coeffs = [1, 3, 9]
        self.offset_mtp = [0, 1, 4]
        self.alpha_train = {'H':None, 'C':None, 'O':None, 'N':None, 'S':None, 'Cl':None, 'F':None, 'Br':None}
        self.ml_method  = options.multipole_ml_method
        self.kernel     = options.multipole_kernel
        self.krr_sigma  = options.multipole_krr_sigma
        self.krr_lambda = options.multipole_krr_lambda
        # Normalization of the target data - mean and std for each MTP component
        self.norm_tgt_mean = {'H':np.zeros((3)),'C':np.zeros((3)),'O':np.zeros((3)), 'N':np.zeros((3)), 'S':np.zeros((3)), 'Cl':np.zeros((3)), 'F':np.zeros((3)), 'Br':np.zeros((3))}
        self.norm_tgt_std  = {'H':np.ones((3)), 'C':np.ones((3)), 'O':np.ones((3)), 'N':np.ones((3)), 'S':np.ones((3)), 'Cl':np.ones((1)), 'F':np.zeros((3)), 'Br':np.zeros((3))}
        self.num_mols_train = {'H':0, 'C':0, 'O':0, 'N':0, 'S':0, 'Cl':0, 'F':0, 'Br':0}
        # Set of qml mols. Used for SLATM
        self.qml_mols = []
        self.qml_filter_ele = []
        self.mbtypes = None
        self.ref_mtp = options.multipole_ref_mtp

        self.cutoff = options.multipole_rcut

        self.ref_path = ""
        if self.ref_mtp:
            self.ref_path = options.multipole_ref_path
        if options.test_mode:
            self.ref_path = testpath + self.ref_path
        self.training_dir = options.multipole_training

        self.correct_charge = options.multipole_correct_charge

        self.nn_mtp = None

    def set_ref_path(self, path):
        self.ref_path = path

    def load_ml(self, load_file=None):
        '''Load machine learning model'''
        # Try many atoms and see which atoms we find
        load_files = []
        if load_file != None:
            load_files.append(load_file)
        else:
            load_files = glob.glob(self.training_dir + "/*.pkl")
            self.logger.info(
                "    Loading Multipole training from %s" % self.training_dir)

            #print("MTP models: ", self.training_dir)
        for mtp_file in load_files:
            try:
                with open(mtp_file, 'rb') as f:
                    descr_train_at, alpha_train, norm_tgt_mean, \
                    norm_tgt_std, mbtypes = pickle.load(f,encoding="ISO-8859-1")
                    #norm_tgt_std, mbtypes = pickle.load(f) #try for old pickles
                    #norm_tgt_std, mbtypes = pickle.load(f, encoding='latin1') #try for old pickles
                    for e in self.descr_train.keys():
                        if e in descr_train_at.keys() and len(descr_train_at[e]) > 0:
                            # Update
                            self.descr_train[e] = descr_train_at[e]
                            self.alpha_train[e] = alpha_train[e]
                            self.norm_tgt_mean[e] = norm_tgt_mean[e]
                            self.norm_tgt_std[e] = norm_tgt_std[e]
                    self.mbtypes = mbtypes
            except:
                print("Could not load model ", mtp_file)
        return None

    def save_ml(self, save_file):
        '''Save machine learning model'''
        self.logger.info("   Saving multipole machine learning model to %s" %
            save_file)
        with open(save_file, 'wb') as f:
            pickle.dump([self.descr_train, self.alpha_train,
                         self.norm_tgt_mean, self.norm_tgt_std, self.mbtypes], f, protocol=2)
        return None

    def train_mol(self):
        '''Train machine learning model of multipole rank mtp_rank and
        basis set expansion coefficient coeff.'''
        # SLATM: First compute mbtypes and the representation
        # Reinitialize descriptor
        for key in self.descr_train.keys():
            self.descr_train[key] = []
        # Check that mbtypes exists
        if self.mbtypes is None:
            raise ValueError("mbtypes missing")
        for i,mol in enumerate(self.qml_mols):
            mol.generate_slatm(self.mbtypes, rcut = self.cutoff, local=True)
            for j,at in enumerate(self.qml_filter_ele[i]):
                if at == 1:
                    # Do include the descriptor
                    self.descr_train[mol.atomtypes[j]].append(mol.representation[j])
        for e in self.descr_train.keys():
            size_training = len(self.target_train[e])
            # self.normalize(e)
           # print(e, size_training, len(self.descr_train[e]))
           # print(self.alpha_train)
           # exit()
            if len(self.descr_train[e]) > 0:
                self.logger.info("Training set size: %d atoms; %d molecules" % (size_training,
                    self.num_mols_train[e]))
                #print("Training set size: %d atoms; %d molecules" % (size_training,
                #    self.num_mols_train[e]))
                tgt_prop = [[self.target_train[e][i][mtp_rank][coeff]
                            for mtp_rank in range(3)
                            for coeff in range(self.max_coeffs[mtp_rank])]
                            for i in range(len(self.target_train[e]))]
                pairwise_dists = squareform(pdist(self.descr_train[e],
                    constants.ml_metric[self.kernel]))
                self.logger.info("building kernel matrix of size (%d,%d); %7.4f Gbytes" \
                    % (size_training, size_training, 8*size_training**2/1e9))
                #print("building kernel matrix of size (%d,%d); %7.4f Gbytes" \
                #    % (size_training, size_training, 8*size_training**2/1e9))
                power  = constants.ml_power[self.kernel]
                prefac = constants.ml_prefactor[self.kernel]
                kmat = np.exp(- pairwise_dists**power / (prefac*self.krr_sigma**power))
                #kmat += self.krr_lambda*np.identity(len(self.target_train[e]))
                kmat.flat[::len(self.target_train[e])+1] += self.krr_lambda

                self.alpha_train[e] = np.linalg.solve(kmat,tgt_prop)
        print(self.alpha_train)
        self.logger.info("training of multipoles finished.")
        return None

    def predict_mol(self, _system, charge=0, xyz=None, force_predict = False):
        '''Predict multipoles in local reference frame given descriptors.'''
        tp = time.time()
        _system.initialize_multipoles()
        _system.compute_basis()

        _system.build_slatm(self.mbtypes,self.cutoff)
        power  = constants.ml_power[self.kernel]
        prefac = constants.ml_prefactor[self.kernel]
        for e in self.alpha_train.keys():
            if self.alpha_train[e] is not None:
                # slowest part of the whole project
                pairwise_dists = cdist(_system.slatm, \
                    self.descr_train[e], constants.ml_metric[self.kernel])
                kmat = np.exp(- pairwise_dists**power / (prefac*self.krr_sigma**power))
                pred = np.dot(kmat,self.alpha_train[e])
                for i in range(_system.num_atoms):
                    if _system.elements[i] == e:
                        _system.mtp_expansion[i] = pred[i]
        # Revert normalization
        # self.rev_normalize(_system)
        # Correct to get integer charge
        if self.correct_charge:
            # Weigh by ML error
            mol_mu = sum([constants.ml_chg_correct_error[ele]
                            for ele in _system.elements])
            totalcharge = sum([mtp[0] for mtp in _system.mtp_expansion])
            excess_chg = totalcharge - float(charge)
            if mol_mu > 0.:
                for i,mtp_i in enumerate(_system.mtp_expansion):
                    w_i = constants.ml_chg_correct_error[_system.elements[i]]
                    mtp_i[0] += -1.*excess_chg * (w_i/mol_mu)
        # Compute multipoles from basis set expansion
        _system.expand_multipoles()

        self.logger.debug("Predicted multipole expansion for %s" % ( _system.xyz[0]))

       # print("    Time spent predicting multipoles:                     %8.3f s" % (time.time() - tp))
        return None

    def add_mol_to_training(self, new_system, pun, atom=None, xyz=None):
        'Add molecule to training set'
        new_system.initialize_multipoles()

        # Don't build SLATM yet, only add information to mbtypes
        mol = None
        if new_system.xyz[0] is None:
            if xyz is not None:
                mol = qml.Compound(xyz)
            else:
                raise ValueError("Missing xyz file")
        else:
            mol = qml.Compound(new_system.xyz[0])
        self.qml_mols.append(mol)
        if atom is None:
            self.qml_filter_ele.append([1 for i in range(mol.natoms)])
        else:
            self.qml_filter_ele.append([1 if (str(mol.atomtypes[i]) == atom) else 0
                                for i in range(mol.natoms)])
        new_system.multipoles = pun
        # Read in multipole moments from txt file
        if len(new_system.multipoles) != new_system.num_atoms:
            raise Exception("Wrong number of charges in %s" % (len(pun)))

        for i in range(len(new_system.elements)):
            ele_i = new_system.elements[i]
            if (ele_i == atom) or atom is None:
                if ele_i not in self.target_train.keys():
                    self.target_train[ele_i] = []
                    self.descr_train[ele_i] = []
                    self.num_mols_train[ele_i] = 0
                new_target_train = []
                # Rotate system until atom pairs point in all x,y,z directions
                vec_all_dir = new_system.compute_basis()
                # charge
                new_target_train.append([new_system.multipoles[i][0]])
                # dipole
                new_target_train.append(np.dot(new_system.multipoles[i][1:4],
                                    new_system.basis[i].T))
                # quadrupole
                tmp = np.dot(np.dot(new_system.basis[i],
                    utils.spher_to_cart(new_system.multipoles[i][4:9])),
                            new_system.basis[i].T).reshape((9,))
                new_target_train.append(tmp)
                self.target_train[ele_i].append(new_target_train)
            if atom in new_system.elements or atom is None:
                self.num_mols_train[ele_i] += 1
        self.logger.info("Added file to training set: %s" % new_system)
        return None


import sys
'''
Section 1: Collect customer input
'''

# Rental rates are assigned to three different variables
chargeBudget = 40.00
chargeDaily = 60.00
chargeWeekly = 190.00

# Rental rate variables are assigned to codes - B, D, W
B = chargeBudget
D = chargeDaily
W = chargeWeekly

rentalCode = input("(B)ugdet, (D)aily, or (W)eekly rental?\n").upper()  # The .upper method is included here to avoid error should the user put in a lowercase letter, it will be converted to uppercase
print("You have selected the " + rentalCode + " rate.") # This line is used to ensure the correct rental code has been selected

# Conditional statements based on the rental code input of the user

if rentalCode == 'B':
  print("Enter the number of days Rented:") # Asking user for number of days rented
  rentalPeriod = int(input()) # Accept user input for number of days rented
  print("\nNumber of Days rented: ", rentalPeriod, "\nDaily Budget rate: $", "%.2f" % chargeBudget) # Printing confirmation of number of days rented at the correct rate.  This is to ensure all data is correct and is debugged
elif rentalCode == 'D':
  print("Enter the number of days Rented:") # Asking user for number of days rented
  rentalPeriod = int(input()) # Accept user input for number of days rented
  print("\nNumber of Days rented: ", rentalPeriod, "\nDaily rate: $", "%.2f" % chargeDaily) # Printing confirmation of number of days rented at the correct rate.  This is to ensure all data is correct and is debugged
else: # else statement to account for the selection of weekly rate
  print("Enter the number of weeks Rented:")  # Asking user for number of weeks rented
  rentalPeriod = int(input()) # Accept user input for number of days rented
  print("\nNumber of Weeks rented: ", rentalPeriod, "\nWeekly rate: $", "%.2f" % chargeWeekly) # Printing confirmation of number of days rented at the correct rate.  This is to ensure all data is correct and is debugged

# Using print function to verify data and debug if necessary
if rentalCode == 'B':
  baseCharge = rentalPeriod * chargeBudget
  print("Total rental base charge: $", baseCharge, "\n")
elif rentalCode == 'D':
  baseCharge = rentalPeriod * chargeDaily
  print("Total rental base charge: $", baseCharge, "\n")
else:
  baseCharge = rentalPeriod * chargeWeekly
  print("Total rental base charge: $", baseCharge, "\n")

# Milage calculations and charges
print("Starting Odometer Reading:")
odoStart = int(input()) # Accept user inpur variable as an intege
print("Ending Odometer Reading:")
odoEnd = int(input()) # Accept user inpur variable as an intege

#Calculate total miles driven
totalMiles = odoEnd - odoStart  # Set varibable 'totalMiles' (ending - starting)

print("\nMileage summary\n\nStarting mileage:", odoStart,  "\nEnding mileage: ", odoEnd, "\nTotal mileage: ", totalMiles) # Printing mileage summarty in three lines to verify data and debug if necessary, and to check calculation of totalMiles

# Set three variables which will be used in the conditional statement
# It is important to read and understand these variables, so as to not create repitation in comments below
averageDayMiles = totalMiles / rentalPeriod # Calculates the average DAILY miles driven for the rental period
extraDayMiles = int(averageDayMiles) - 100  # Customer will be charged per mile, if they exceed an average of 100 miles a day for their rental period
averageWeekMiles = totalMiles / rentalPeriod  # Calculates the average WEEKLY miles driven for the rental period
extraWeekMiles = int(averageWeekMiles) - 900  # Customer will be charged per mile, if they exceed an average of 100 miles a day for their rental period

#Conditional statements with nesting
if rentalCode == 'B': # Conditional statement for BUDGET
  mileCharge = totalMiles * .25 # 25 cents is charged for ever mile driven
  print("Budget milage charge total: $", mileCharge)  # Prints the total charged for BUDGET rental to verify data, and debug if necessary
elif rentalCode == 'D': # Or if rental code is DAILY >>>> branched into another conditional statement for different milage calculations
  if averageDayMiles <= 100:  # If average daily miles driven was 100 or less
    mileCharge = 0  # There is no charge on the miles driven
    print("Daily milage charge total: $", mileCharge) # Prints the total charged for DAILY rental to verify data, and debug if necessary
  else: # Or else, if the average daily miles is greater than 100
    mileCharge = extraDayMiles * .25 # Usiing the variable 'extraDailyMiles', $00.25 is charged on every mile over the average 100 miles driven.
    print("Daily milage charge total: $", mileCharge) # Prints the total charged for DAILY rental to verify data, and debug if necessary
else: # Or if rental code is WEEKLY >>>> branched into another conditional statement for different milage calculations
  if averageWeekMiles <= 900: # If average weekly miles driven was 900 or less
    mileCharge = 0 # There is no charge on the miles driven
    print("Weekly milage charge total: $", mileCharge)  # Prints the total charged for WEEKLY rental to verify data, and debug if necessary
  else: # However if the average weekly miles driven exceeds 900 miles >>>> charges below apply
    mileCharge = rentalPeriod * 100.00  # Usiing the variable 'extraWeekMiles', $00.25 is charged on every mile over the average 100 miles driven.
    print("Weekly milage charge total: $", mileCharge)  # Prints the total charged for BUDGET rental to verify data, and debug if necessary

#   Final calculations

amtDue = baseCharge + mileCharge    #   Assigning variable 'amtDue' >>> we gater the data from the baseCharge and add the mileCharge

print("\nRental Summary\n") #   Print 'Rental Summary'
#   This summary prints all the information clearly, using '\n' to break lines
print("Rental Code: ", rentalCode, "\nRental Period: ", rentalPeriod, "\nStarting Odometer: ", odoStart, "\nEnding Odometer: ", odoEnd, "\nMiles Driven: ", totalMiles, "\nAmount Due: $", amtDue)


import time
import sys
sys.path.append("../..")
import DFRobot_BC20

RED = 'R'
BLUE = 'B'
GREEN = 'G'
YELLOW = 'Y'
PURPLE = 'P'
CYAN = 'C'
WHITE = 'W'

'''For general controllers. Communicate by IIC
   Connect Instructions
      Raspberry Pi   |    Module(BC20)
      SDA(GPIO3)     |       D/T
      SCL(GPIO5)     |       C/R
      GND(GPIO6)     |       GND
      VCC(GPIO2)     |       VCC

   IIC address(A0,A1)
     0x30:(A0=0,A1=0)
     0x31:(A0=0,A1=1)
     0x32:(A0=1,A1=0)
     0x33:(A0=1,A1=1) default'''

'''For general controllers. Communicate by Serial
   Connect Instructions
      Raspberry Pi   |    Module(BC20)
      RXD(GPIO10)    |       D/T
      TXD(GPIO8)     |       C/R
      GND(GPIO6)     |       GND
      VCC(GPIO2)     |       VCC'''

BC20=DFRobot_BC20.DFRobot_BC20_IIC(0x33)

#BC20=DFRobot_BC20.DFRobot_BC20_Serial()

print("Starting the BC20.Please wait...")
BC20.changeColor(RED)
while BC20.powerOn()==False:
  print(".")
  BC20.ledON()
  time.sleep(0.5)
  BC20.ledOFF()
  time.sleep(0.5)
print("BC20 started successfully!\r\n")

BC20.configSleepMode(0)

print("Turning on GNSS ... ")
BC20.setQGNSSC(1)
BC20.changeColor(YELLOW)
if BC20.getQGNSSC() == 0:
  print(".")
  BC20.ledON()
  time.sleep(0.5)
  BC20.ledOFF()
  time.sleep(0.5)
print("GNSS is ON.")
BC20.changeColor(CYAN)

while True:
  BC20.getQGNSSRD("\"NMEA/GLL\"")

#UTC time, format: hhmmss.ss, ex. 162436.54 => 16:24:36.54
#h - hours
#m - minutes
#s - seconds
  print("UTC Time: "),
  print(BC20.sGLL.UTC_Time)

#Latitude, format: ddmm.mmmmm, ex. 3150.7820 => 31deg 50.7820min
#d - degrees
#m - minutes
  print("Latitude: "),
  print(BC20.sGLL.LatitudeVal),

#Latitude north or south
#N - North
#S - South
  print(BC20.sGLL.LatitudeDir)

#Longitude, format: dddmm.mmmmm, ex. 12135.6794 = 121deg 35.6794min
#d - degrees
#m - minutes
  print("Longitude: "),
  print(BC20.sGLL.LongitudeVal),

#Longitude wast or west
#E - East
#W - West
  print(BC20.sGLL.LongitudeDir)

#DataStatus
#V - Invalid
#A - Valid

  print("Data Status: "),
  print(BC20.sGLL.DataStatus)

#Positioning Mode
#N - No fix
#A - Autonomous GPS fix
#D - Differential GPS fix

  print("Positioning Mode: "),
  print(BC20.sGLL.PositionMode)
  print("")

  BC20.ledON()
  time.sleep(0.5)
  BC20.ledOFF()
  time.sleep(3)