import pandas as pd
import pyranges as pr
import re

def extract_dna_id(filename):
    pattern = "genomics\/3_vcf\/.*\/(.*)\/.*"
    dna_id = re.search(pattern, filename).group(1)
    return dna_id

df_anno = pd.read_csv(snakemake.input['sample_anno'], sep='\t')

_df_anno = df_anno[
    ~df_anno['DNA_VCF_FILE'].isna()
]
_df_anno['VCF_ID'] = _df_anno.apply(lambda x: extract_dna_id(x['DNA_VCF_FILE']), axis=1)
df_anno = df_anno.set_index('RNA_ID').join(_df_anno.set_index('RNA_ID')['VCF_ID']).reset_index()
df_anno.to_csv(snakemake.output['sample_anno_updated'], sep='\t')


#! /usr/bin/env python
# -*- coding: utf-8 -*-

"""
Module that contains functions related with Maya tag functionality for ueGear.
"""

from __future__ import print_function, division, absolute_import

import maya.cmds as cmds
import maya.api.OpenMaya as OpenMaya

from mgear.uegear import utils, log

logger = log.uegear_logger

TAG_ASSET_GUID_ATTR_NAME = "ueGearAssetGuid"
TAG_ASSET_TYPE_ATTR_NAME = "ueGearAssetType"
TAG_ASSET_NAME_ATTR_NAME = "ueGearAssetName"
TAG_ASSET_PATH_ATTR_NAME = "ueGearAssetPath"
TAG_ACTOR_NAME_ATTR_NAME = "ueGearActorName"
ALL_TAGS_ATTR_NAMES = [
    TAG_ASSET_GUID_ATTR_NAME,
    TAG_ASSET_TYPE_ATTR_NAME,
    TAG_ASSET_NAME_ATTR_NAME,
    TAG_ASSET_PATH_ATTR_NAME,
    TAG_ACTOR_NAME_ATTR_NAME,
]

class TagTypes(object):
    """
    Class that holds all available tag types.
    """

    Skeleton = "skeleton"
    StaticMesh = "staticmesh"
    SkeletalMesh = "skeletalmesh"
    Camera = "camera"
    Alembic = "alembic"
    MetahumanBody = "metahumanbody"
    MetahumanFace = "metahumanface"
    Sequence = "sequence"

def auto_tag(node=None, remove=False):
    """
    Automatically tags given (or current selected nodes) so ueGear exporter can identify how to export the specific
    nodes.

    :param str or list(str) or None node: node/s to tag.
    :param bool remove: if True tag will be removed.
    """

    nodes = utils.force_list(node or cmds.ls(sl=True, long=True))

    for node in nodes:
        found_skin_clusters = utils.get_skin_clusters_for_node(node)
        if found_skin_clusters and cmds.objectType(node) == "joint":
            remove_tag(node) if remove else apply_tag(
                node, attribute_value=TagTypes.SkeletalMesh
            )
        else:
            shapes = cmds.listRelatives(node, fullPath=True, shapes=True)
            if not shapes:
                continue
            first_shape = utils.get_first_in_list(shapes)
            if not first_shape:
                continue
            object_type = cmds.objectType(first_shape)
            if object_type == "mesh":
                found_skin_clusters = utils.get_skin_clusters_for_node(
                    first_shape
                )
                if found_skin_clusters:
                    remove_tag(node) if remove else apply_tag(
                        node, attribute_value=TagTypes.Skeleton
                    )
                else:
                    remove_tag(node) if remove else apply_tag(
                        node, attribute_value=TagTypes.StaticMesh
                    )
            elif object_type == "camera":
                remove_tag(node) if remove else apply_tag(
                    node, attribute_value=TagTypes.Camera
                )

def apply_tag(
    node=None, attribute_name=TAG_ASSET_TYPE_ATTR_NAME, attribute_value=""
):
    """
    Creates a new tag attribute with given value into given node/s (or selected nodes).

    :param str or list(str) or None node: nodes to apply tag to.
    :param str attribute_name: tag attribute value to use. By default, TAG_ASSET_TYPE_ATTR_NAME will be used.
    :param str attribute_value: value to set tag to.
    """

    nodes = utils.force_list(node or cmds.ls(sl=True))
    attribute_value = str(attribute_value)

    for node in nodes:
        if not cmds.attributeQuery(attribute_name, node=node, exists=True):
            cmds.addAttr(node, longName=attribute_name, dataType="string")
        cmds.setAttr(
            "{}.{}".format(node, attribute_name),
            attribute_value,
            type="string",
        )
        if attribute_value:
            logger.info(
                'Tagged "{}.{}" as {}.'.format(
                    node, attribute_name, attribute_value
                )
            )
        else:
            logger.info(
                'Tagged "{}.{}" as empty.'.format(node, attribute_name)
            )

def remove_tag(node=None, attribute_name=TAG_ASSET_TYPE_ATTR_NAME):
    """
    Removes tag attribute from the given node.

    :param str or list(str) or None node: nodes to remove tag from.
    :param str attribute_name: tag attribute value to remove. By default, TAG_ASSET_TYPE_ATTR_NAME will be used.
    """

    nodes = utils.force_list(node or cmds.ls(sl=True))

    for node in nodes:
        if not cmds.attributeQuery(attribute_name, node=node, exists=True):
            continue
        cmds.deleteAttr("{}.{}".format(node, attribute_name))
        logger.info(
            'Removed attribute {} from "{}"'.format(attribute_name, node)
        )

def remove_all_tags(node=None):
    """
    Removes all ueGear tags from the given node.

    :param str or list(str) or None node: nodes to remove tags from.
    """

    nodes = utils.force_list(node or cmds.ls(sl=True))

    for attribute_name in ALL_TAGS_ATTR_NAMES:
        remove_tag(nodes, attribute_name=attribute_name)

def apply_alembic_tag(node=None, remove=False):
    """
    Applies alembic tag to given node/s (or selected nodes).

    :param str or list(str) or None node: node/s to tag.
    :param bool remove: if True tag will be removed.
    """

    remove_tag(node=node) if remove else apply_tag(
        node=node, attribute_value=TagTypes.Alembic
    )

def find_tagged_nodes(
    tag_name=TAG_ASSET_TYPE_ATTR_NAME, nodes=None, tag_value=None
):
    """
    Returns a list with all nodes that are tagged with the given tag name and has a value set.

    :param str tag_name: name of the tag to search. By default, TAG_ATTR_NAME will be used.
    :param str or list(str) or None nodes: list of nodes to find tags of, if not given all nodes in the scene will be
            checked.
    :param str tag_value: if given only tag with given value will be returned.
    :return: list of found tagged nodes.
    :rtype: list(str)
    """

    found_tagged_nodes = list()
    nodes = utils.force_list(nodes or cmds.ls())
    for node in nodes:
        if not cmds.attributeQuery(tag_name, node=node, exists=True):
            continue
        found_tag_value = cmds.getAttr("{}.{}".format(node, tag_name))
        if not found_tag_value or (
            tag_value is not None and found_tag_value != tag_value
        ):
            continue
        found_tagged_nodes.append(node)

    return found_tagged_nodes

def find_tagged_selected_nodes(tag_name):
    """
    Returns a list with all selected nodes that are tagged with the given tag name and has a value set.

    :param str tag_name: name of the tag to search. By default, TAG_ATTR_NAME will be used.
    :return: list of found tagged nodes.
    :rtype: list(str)
    """

    return find_tagged_nodes(nodes=cmds.ls(sl=True))

def find_tagged_node_attributes(tag_name=TAG_ASSET_TYPE_ATTR_NAME, nodes=None):
    """
    Returns a list with all node attributes that are tagged with the given tag name and has a value set.

    :param str tag_name: name of the tag to search. By default, TAG_ATTR_NAME will be used.
    :param str or list(str) or None nodes: list of nodes to find tags of, if not given all nodes in the scene will be
            checked.
    :return: list of found tagged nodes.
    :rtype: list(str)
    """

    found_tagged_node_attributes = list()
    nodes = utils.force_list(nodes or cmds.ls(long=True))
    for node in nodes:
        if not cmds.attributeQuery(tag_name, node=node, exists=True):
            continue
        if not cmds.getAttr("{}.{}".format(node, tag_name)):
            continue
        found_tagged_node_attributes.append("{}.{}".format(node, tag_name))

    return found_tagged_node_attributes

def find_tagged_selected_node_attributes(tag_name):
    """
    Returns a list with all selected node attributes that are tagged with the given tag name and has a value set.

    :param str tag_name: name of the tag to search. By default, TAG_ATTR_NAME will be used.
    :return: list of found tagged nodes.
    :rtype: list(str)
    """

    return find_tagged_node_attributes(nodes=cmds.ls(sl=True))

def tag_values(tag_name=TAG_ASSET_TYPE_ATTR_NAME, nodes=None):
    """
    Returns a list with all node attribute values that are tagged with the given tag name.

    :param str tag_name:name of the tag to search value of.
    :param str or list(str) nodes: list of nodes to find tags of, if not given all nodes in the scene will be checked.
    :return: list of tagged node values.
    :rtype: list(object)
    """

    found_tag_values = list()
    nodes = utils.force_list(nodes or cmds.ls(long=True))
    for node in nodes:
        if not cmds.attributeQuery(tag_name, node=node, exists=True):
            found_tag_values.append(None)
            continue
        found_tag_values.append(cmds.getAttr("{}.{}".format(node, tag_name)))

    return found_tag_values

def tag_match(dag_path, tag_value, tag):
    """
    Validates if the object specified by its dag path, has the same tag and value
    assigned to it.

    :param OpenMaya.DagPath dag_path: The object you want to validate has the
           following tag and data assigned.
    :param str tag_value: value assigned to the tag.
    :param str tag: tag to correlate with.
    :return: True if the object has matching tag and the values are the same.
    :rtype: bool
    """
    dag_node = OpenMaya.MFnDagNode(dag_path)
    attr = dag_node.attribute(tag)
    plug = dag_node.findPlug(attr, False)
    plug_value = plug.asString()
    return plug_value == tag_value


from __future__ import annotations

from unittest import mock

from kombu.utils.objects import cached_property

class test_cached_property:

    def test_deleting(self):

        class X:
            xx = False

            @cached_property
            def foo(self):
                return 42

            @foo.deleter
            def foo(self, value):
                self.xx = value

        x = X()
        del x.foo
        assert not x.xx
        x.__dict__['foo'] = 'here'
        del x.foo
        assert x.xx == 'here'

    def test_when_access_from_class(self):

        class X:
            xx = None

            @cached_property
            def foo(self):
                return 42

            @foo.setter
            def foo(self, value):
                self.xx = 10

        desc = X.__dict__['foo']
        assert X.foo is desc

        assert desc.__get__(None) is desc
        assert desc.__set__(None, 1) is desc
        assert desc.__delete__(None) is desc
        assert desc.setter(1)

        x = X()
        x.foo = 30
        assert x.xx == 10

        del x.foo

    def test_locks_on_access(self):

        class X:
            @cached_property
            def foo(self):
                return 42

        x = X()

        # Getting the value acquires the lock, and may do so recursively
        # on Python < 3.12 because the superclass acquires it.
        with mock.patch.object(X.foo, 'lock') as mock_lock:
            assert x.foo == 42
        mock_lock.__enter__.assert_called()
        mock_lock.__exit__.assert_called()

        # Setting a value also acquires the lock.
        with mock.patch.object(X.foo, 'lock') as mock_lock:
            x.foo = 314
        assert x.foo == 314
        mock_lock.__enter__.assert_called_once()
        mock_lock.__exit__.assert_called_once()

        # .. as does clearing the cached value to recompute it.
        with mock.patch.object(X.foo, 'lock') as mock_lock:
            del x.foo
        assert x.foo == 42
        mock_lock.__enter__.assert_called_once()
        mock_lock.__exit__.assert_called_once()


import cv2

def draw_boxes(im, boxes, class_names=None, scores=None, colors=None):
    scores = [None] * len(boxes) if scores is None else scores
    colors = [None] * len(boxes) if colors is None else colors
    class_names = [None] * len(boxes) if class_names is None else class_names

    for params in zip(boxes, class_names, scores, colors):
        _draw_box(im, *params)

    return im

def _draw_box(im, box, class_name=None, score=None, color=None):
    x1, y1, x2, y2 = box
    color = color if color is not None else (0, 255, 0)
    msg = class_name.capitalize() if class_name else None
    if msg is not None and score is not None:
        msg += f' [{int(score * 100)}]'

    cv2.rectangle(im, (x1, y1), (x2, y2), color=color, thickness=2)
    if msg is not None:
        cv2.rectangle(im, (x1 - 1, y1 - 20), (x2 + 1, y1), color, -1)
        cv2.putText(im, msg, (x1 + 10, y1 - 8), cv2.FONT_HERSHEY_SIMPLEX ,
                    .5, (0, 0, 0), 2, cv2.LINE_AA)

    return im


# -*- coding: utf-8 -*-

from django.db import models
from django.utils.translation import gettext_lazy as _
from django.utils import timezone
from django.db import IntegrityError, transaction

from .managers import TopicNotificationQuerySet
from spirit.core.conf import settings

class TopicNotification(models.Model):
    UNDEFINED, MENTION, COMMENT = range(3)
    ACTION_CHOICES = (
        (UNDEFINED, _("Undefined")),
        (MENTION, _("Mention")),
        (COMMENT, _("Comment")))

    user = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        related_name='st_topic_notifications',
        on_delete=models.CASCADE)
    topic = models.ForeignKey(
        'spirit_topic.Topic',
        on_delete=models.CASCADE)
    comment = models.ForeignKey(
        'spirit_comment.Comment',
        on_delete=models.CASCADE)

    date = models.DateTimeField(default=timezone.now)
    action = models.IntegerField(choices=ACTION_CHOICES, default=UNDEFINED)
    is_read = models.BooleanField(default=False)
    is_active = models.BooleanField(default=False)

    objects = TopicNotificationQuerySet.as_manager()

    class Meta:
        unique_together = ('user', 'topic')
        ordering = ['-date', '-pk']
        verbose_name = _("topic notification")
        verbose_name_plural = _("topics notification")

    def get_absolute_url(self):
        if self.topic_id != self.comment.topic_id:
            # Out of sync
            return self.topic.get_absolute_url()
        return self.comment.get_absolute_url()

    @property
    def text_action(self):
        return self.ACTION_CHOICES[self.action][1]

    @property
    def is_mention(self):
        return self.action == self.MENTION

    @property
    def is_comment(self):
        return self.action == self.COMMENT

    @classmethod
    def mark_as_read(cls, user, topic):
        if not user.is_authenticated:
            return

        (cls.objects
         .filter(user=user, topic=topic)
         .update(is_read=True))

    @classmethod
    def create_maybe(cls, user, comment, is_read=True, action=COMMENT):
        # Create a dummy notification
        return cls.objects.get_or_create(
            user=user,
            topic=comment.topic,
            defaults={
                'comment': comment,
                'action': action,
                'is_read': is_read,
                'is_active': True})

    @classmethod
    def notify_new_comment(cls, comment):
        (cls.objects
         .filter(topic=comment.topic, is_active=True, is_read=True)
         .exclude(user=comment.user)
         .update(
            comment=comment,
            is_read=False,
            action=cls.COMMENT,
            date=timezone.now()))

    @classmethod
    def notify_new_mentions(cls, comment, mentions):
        if not mentions:
            return

        # TODO: refactor
        for user in mentions.values():
            try:
                with transaction.atomic():
                    cls.objects.create(
                        user=user,
                        topic=comment.topic,
                        comment=comment,
                        action=cls.MENTION,
                        is_active=True)
            except IntegrityError:
                pass

        (cls.objects
         .filter(
            user__in=tuple(mentions.values()),
            topic=comment.topic,
            is_read=True)
         .update(
            comment=comment,
            is_read=False,
            action=cls.MENTION,
            date=timezone.now()))

    @classmethod
    def bulk_create(cls, users, comment):
        return cls.objects.bulk_create([
            cls(user=user,
                topic=comment.topic,
                comment=comment,
                action=cls.COMMENT,
                is_active=True)
            for user in users])

    # XXX add tests
    # XXX fix with migration (see issue #237)
    @classmethod
    def sync(cls, comment, topic):
        # Notifications can go out of sync
        # when the comment is no longer
        # within the topic (i.e moved).
        # User is subscribed to the topic,
        # not the comment, so we either update
        # it to a newer comment or set it as undefined
        if comment.topic_id == topic.pk:
            return
        next_comment = (
            topic.comment_set
                .filter(date__gt=comment.date)
                .order_by('date')
                .first())
        if next_comment is None:
            (cls.objects
             .filter(comment=comment, topic=topic)
             .update(is_read=True, action=cls.UNDEFINED))
            return
        (cls.objects
         .filter(comment=comment, topic=topic)
         .update(comment=next_comment, action=cls.COMMENT))


import os
from dotenv import load_dotenv
from configparser import ConfigParser

conf = ConfigParser()
conf.read('model.conf')

load_dotenv('.env')

def _getenv(key, default): return type(default)(os.getenv(key)) if os.getenv(key) else default

SERVER_IP = _getenv('SERVER_IP', '0.0.0.0')  # Service IP
SERVER_PORT = _getenv('SERVER_PORT', '6002')  # Service IP

REGISTER = _getenv('REGISTER', 0) # register to the management service

MANAGER_IP = _getenv('MANAGER_IP', '127.0.0.1') # Management server address
MANAGER_PORT = _getenv('MANAGER_PORT', 5005) # Management server address
MANAGER_INTERFACE_REGISTER = _getenv('MANAGER_INTERFACE_REGISTER', '/model/register')
MANAGER_INTERFACE_CANCEL = _getenv('MANAGER_INTERFACE_CANCEL', '/model/cancel')

MODEL_TYPE = _getenv('MODEL_TYPE', conf.get('model', 'model_type', fallback='')) # Service type
MODEL_VERSION = _getenv('MODEL_VERSION', 1) # Service version number

ENGINE_FILE_PATH = _getenv('ENGINE_FILE_PATH', conf.get('model', 'engine_file_path', fallback=''))
CLASS_NUM = _getenv('CLASS_NUM', int(conf.get('model', 'class_num', fallback='0')))
CLASS_NAMES = [name.strip() for name in _getenv('CLASS_NAMES', conf.get('model', 'class_names')).split(',')]

KEY = _getenv('KEY', 'LONGYUAN')


import tensorflow as tf
import pandas as pd
from sklearn.metrics import multilabel_confusion_matrix, confusion_matrix, precision_score, recall_score, f1_score

def calculate_output(model, actual_classes, session, feed_dict):
    actuals = tf.argmax(actual_classes, 1)
    predictions = tf.argmax(model, 1)
    actuals = session.run(actuals, feed_dict)
    predictions = session.run(predictions, feed_dict)
    return actuals, predictions

def tf_confusion_metrics(model, actual_classes, session, feed_dict):
  import numpy as np
  cat = 5

  actuals, predictions = calculate_output(model, actual_classes, session, feed_dict)

  lbls = [*range(cat)]
  mcm = multilabel_confusion_matrix(actuals, predictions, labels=lbls)
  tp = mcm[:, 1, 1]
  tn = mcm[:, 0, 0]
  fn = mcm[:, 1, 0]
  fp = mcm[:, 0, 1]

  cm = confusion_matrix(actuals, predictions, labels=lbls, sample_weight=None)
  tp = np.mean(tp)
  tn = np.mean(tn)
  fp = np.mean(fp)
  fn = np.mean(fn)
  try:
      tpr = float(tp)/(float(tp) + float(fn))

      accuracy = (float(tp) + float(tn))/(float(tp) + float(fp) + float(fn) + float(tn))

      recall = tpr
      if((fp+tp)!=0):
        precision = float(tp)/(float(tp) + float(fp))
        f1_score = (2 * (precision * recall)) / (precision + recall)
      else:
        precision=0
        f1_score=0

      fp_rate=float(fp)/(float(fp)+float(tn))
      fn_rate=float(fn)/(float(fn)+float(tp))

      # return precision, recall, f1_score, accuracy, fp_rate, fn_rate
      PR = str(round(precision * 100, 2))
      RC = str(round(recall * 100, 2))
      F1 = str(round(f1_score * 100, 2))
      ACC = str(round(accuracy * 100, 2))
      FPR = str(round(fp_rate * 100, 2))
      FNR = str(round(fn_rate * 100, 2))

      data_pd=[['PR',PR],['RC', RC],['F1', F1],['ACC', ACC],['FPR', FPR], ['FNR', FNR],['tp', tp],['tn', tn],['fp', fp], ['fn', fn]]

      df = pd.DataFrame(data_pd, columns=['Measure', 'Percentage'])

  except Exception as e:
    print(e)
    data_pd = [['PR', 'Err'], ['RC', 'Err'], ['F1', 'Err'], ['ACC', 'Err'], ['FPR', 'Err'], ['FNR', 'Err']]
    df = pd.DataFrame(data_pd, columns=['Measure', 'Percentage'])

  return df

def tf_confusion_metrics_2(model, actual_classes, session, feed_dict):
    actuals, predictions = calculate_output(model, actual_classes, session, feed_dict)
    cm = tf.confusion_matrix(actuals, predictions)

    print("Confusion Matrix")

    return session.run(cm, feed_dict)

def Macro_calculate_measures_tf(y_true, y_pred, session, feed_dict):
    y_true, y_pred = calculate_output(y_pred, y_true, session, feed_dict)
    pr=  precision_score(y_true, y_pred, average='macro')
    rc = recall_score(y_true, y_pred, average='macro')
    f1 = f1_score(y_true, y_pred, average='macro')
    print("pr, rc, f1:" ,str(pr)+ str(rc)+str(f1))
    return pr, rc, f1


import gc
import os
from pathlib import Path
from typing import Any, Dict, cast

import mlflow
import numpy as np
import onnx
import torch
import transformers
from pytorch_lightning import Trainer, seed_everything
from pytorch_lightning.callbacks import ModelCheckpoint
from transformers.modeling_utils import PreTrainedModel
from transformers.onnx import FeaturesManager, export, validate_model_outputs

from crypto_sentiment_demo_app.models.train.base import IModelTrain, TrainRegistry

from .dataset import build_dataloaders, split_train_val
from .pipeline import MetricTracker, SentimentPipeline

transformers.logging.set_verbosity_error()

os.environ["TOKENIZERS_PARALLELISM"] = "false"

@TrainRegistry.register("bert")
class Bert(IModelTrain):
    """Bert model. Wrapper for hugging face models.

    :param cfg: model config
    """

    def __init__(self, cfg: Dict[str, Any]):
        """Init model."""
        super().__init__(cfg)

        self.model_cfg = self.cfg["model"]
        self.class_names = self.cfg["data"]["class_names"]

        if self.model_cfg["device"] == "gpu" and not torch.cuda.is_available():
            self.device = torch.device("cpu")
        else:
            self.device = torch.device(self.model_cfg["device"])

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        """Fit model.

        :param X: train data
        :param y: train labels
        """
        seed_everything(self.model_cfg["seed"])

        train_data, val_data, train_labels, val_labels = split_train_val(X, y)

        train_dataloader, val_dataloader = build_dataloaders(
            self.model_cfg, train_data, train_labels, val_data, val_labels
        )

        self.model = SentimentPipeline(self.model_cfg)

        metric_tracker = MetricTracker()

        checkpoint_path = Path(self.model_cfg["checkpoint_path"]).parent
        checkpoint_filename = Path(self.model_cfg["checkpoint_path"]).stem

        checkpoint_callback = ModelCheckpoint(
            save_top_k=1,
            monitor="val_acc",
            mode="max",
            dirpath=checkpoint_path,
            filename=checkpoint_filename,
        )

        gpus = 1 if self.device.type == "cuda" and torch.cuda.is_available() else 0

        self.trainer = Trainer(
            max_epochs=self.model_cfg["epochs"],
            gpus=gpus,
            callbacks=[metric_tracker, checkpoint_callback],
            num_sanity_val_steps=0,
            enable_checkpointing=True,
            logger=False,
        )

        self.trainer.fit(
            self.model,
            train_dataloaders=train_dataloader,
            val_dataloaders=val_dataloader,
        )

    def save(self) -> None:
        """Save model."""
        save_dir = Path(self.model_cfg["path_to_model"]).parent
        filename = Path(self.model_cfg["path_to_model"]).stem

        pt_path = save_dir / f"{filename}.pt"
        onnx_path = save_dir / f"{filename}.onnx"

        self._onnx_export(onnx_path)
        onnx_model = onnx.load_model(onnx_path)
        mlflow.onnx.log_model(onnx_model=onnx_model, artifact_path="bert", registered_model_name="bert")
        del onnx_model
        gc.collect()

        self.model = SentimentPipeline.load_from_checkpoint(self.model_cfg["checkpoint_path"], cfg=self.model_cfg)
        cast(PreTrainedModel, self.model.model).eval()
        cast(PreTrainedModel, self.model.tokenizer).save_pretrained(pt_path)
        cast(PreTrainedModel, self.model.model).save_pretrained(pt_path)

    def load(self) -> None:
        """Load model checkpoint."""
        self.model = SentimentPipeline.load_from_checkpoint(self.model_cfg["checkpoint_path"], cfg=self.model_cfg)

    def _onnx_export(self, path: Path):
        model_kind, model_onnx_config = FeaturesManager.check_supported_model_or_raise(
            self.model.model, feature="sequence-classification"
        )
        onnx_config = model_onnx_config(self.model.model.config)
        onnx_inputs, onnx_outputs = export(
            self.model.tokenizer, self.model.model, onnx_config, onnx_config.default_onnx_opset, path
        )
        validate_model_outputs(
            onnx_config, self.model.tokenizer, self.model.model, path, onnx_outputs, onnx_config.atol_for_validation
        )

    def enable_mlflow_logging(self) -> None:
        mlflow.set_experiment("bert")
        mlflow.pytorch.autolog()


#!/usr/bin/env python3

user_input = str(input("Please enter a phrase (only characters A-Z): "))

phrase = user_input.split()
result = " "

for i in phrase:
	result += str(i[0]).upper()

print (result)


"""Rotate Image

You are given an n x n 2D matrix representing an image, rotate the image by 90 degrees (clockwise).

You have to rotate the image in-place, which means you have to modify the input 2D matrix directly. DO NOT allocate another 2D matrix and do the rotation.

Input: matrix = [[1,2,3],[4,5,6],[7,8,9]]
Output: [[7,4,1],[8,5,2],[9,6,3]]

Input: matrix = [[5,1,9,11],[2,4,8,10],[13,3,6,7],[15,14,12,16]]
Output: [[15,13,2,5],[14,3,4,1],[12,6,8,9],[16,7,10,11]]
"""

from typing import List
import unittest

def rotate(matrix: List[List[int]]) -> None:
    """
    Do not return anything, modify matrix in-place instead.
    """
    n = len(matrix)
    for row in range((n+1)//2):
        for col in range(n//2):
            temp = matrix[col][n-1-row]
            matrix[col][n-1-row] = matrix[row][col]
            matrix[row][col] = matrix[n-1-col][row]
            matrix[n-1-col][row] = matrix[n-1-row][n-1-col]
            matrix[n-1-row][n-1-col] = temp

class TestProblems(unittest.TestCase):
    def test_rotate_image(self):
        actual = rotate([[1,2,3],[4,5,6],[7,8,9]])
        expected = [[7,4,1],[8,5,2],[9,6,3]]
        self.assertCountEqual(actual, expected)

        actual_1 = rotate([[5,1,9,11],[2,4,8,10],[13,3,6,7],[15,14,12,16]])
        expected_1 = [[15,13,2,5],[14,3,4,1],[12,6,8,9],[16,7,10,11]]
        self.assertCountEqual(actual_1, expected_1)

if __name__ == '__main__':
    unittest.main()


import frappe
from frappe.model.document import Document

class Sales(Document):

	def before_save(self):
		total_amount = 0

		for item in self.item:
			item.amount = item.product_price * item.quantity
			total_amount += item.amount

			product = frappe.get_doc('Product', item.product_name)
			# Decrease the product quantity from Product DocType
			product.quantity -= item.quantity
			product.save()

		self.total_due    = total_amount - self.receive_amount
		self.total_amount = total_amount - self.discount

		customer         = frappe.get_doc('Customer', self.customer)

		# Get the previous values and save it to Customer DocType
		previous_receive = customer.total_receive
		previous_dues    = customer.total_due
		previous_amount  = customer.total_amount

		if customer.total_receive == 0 or customer.total_due == 0:
			customer.total_receive = self.receive_amount
			customer.total_due     = self.total_due
			customer.toal_amount   = self.total_amount
		else:
			customer.total_receive = previous_receive + self.receive_amount
			customer.total_due     = previous_dues + self.total_due
			customer.total_amount  = previous_amount + self.total_amount

		customer.save()


import bottle
import json
import random

from . import DatabaseManager
from .product import Product
import recommender.vector.arithmetic
import recommender.rocchio.algorithm

@bottle.route('/product/get/<doc_id:int>')
def product_get(doc_id):
    d = product_manager.get_product(doc_id).as_dictionary()
    result = {'result': d}
    return result

@bottle.route('/product/remove/<doc_id:int>', method='DELETE')
def product_remove(doc_id):
    try:
        product_manager.remove_document(doc_id)
    except:
        return {'result': False}
    return {'result': True}

@bottle.route('/product/all')
def product_get_all():
    l = [ p.as_dictionary() for p in product_manager.get_all_products()]
    result = {'result': l}
    #bottle.response.content_type = 'application/json'
    return result

@bottle.route('/product/random/<count:int>')
def product_random(count):
    products = product_manager.get_all_products()
    rands = []
    while len(rands) < count:
        index = random_generator.randint(0, len(products)-1)

        rands.append(products[index].as_dictionary())

        products.remove(products[index])
        pass
    result = {'result': rands};
    return result

@bottle.route('/product/insert', method='POST')
def product_insert():
    """
    curl -X POST -d "product={'image_name':'img.jpg','terms':{'a':1,'b':3}}"
    """
    try:
        product_json = bottle.request.forms.get('product')
        product_dict = json.loads(product_json)

        p = Product()
        p.image_name = product_dict['image_name']
        p.terms = product_dict['terms']

        product_manager.add_document(p)
    except:
        return {'result': False}
    return {'result': True}

@bottle.route('/vector/default/<doc_id:int>')
def vector_default(doc_id):
    d = (product_vector_manager
        .get_vector_for_document_id(doc_id)
        .as_dictionary()
        )
    result = {'result': d}
    return result

@bottle.route('/vector/df')
def vector_df():
    d = (
        product_vector_manager
        .get_document_frequency_vector()
        .as_dictionary()
    )
    result = {'result': d}
    return result

@bottle.route('/vector/idf')
def vector_idf():
    d = (
        product_vector_manager
        .get_inverse_document_frequency_vector()
        .as_dictionary()
    )
    result = {'result': d}
    return result

@bottle.route('/vector/tf/<doc_id:int>')
def vector_tf(doc_id):
    d = (
        product_vector_manager
        .get_term_frequency_vector(doc_id)
        .as_dictionary()
    )
    result = {'result': d}
    return result

@bottle.route('/vector/tfidf/<doc_id:int>')
def vector_tfidf(doc_id):
    d = (
        product_vector_manager
        .get_tfidf_vector(doc_id)
        .as_dictionary()
    )
    result = {'result': d}
    return result

@bottle.route('/vector/user/<user_id:int>')
def vector_user_by_id(user_id):
    d = (
        user_vector_manager
            .get_user_vector_for_id(user_id)
            .as_dictionary()
    )
    result = {'result': d}
    return result

@bottle.route('/vector/user/<user_name>')
def vector_user_by_name(user_name):
    d = (
        user_vector_manager
            .get_user_vector_for_name(user_name)
            .as_dictionary()
    )
    result = {'result': d}
    return result

@bottle.route('/user/all')
def get_all_users():
    user_list = user_vector_manager.get_all_users_by_name()
    result = {'result': user_list}
    return result

@bottle.route('/user/create/<user_name>')
def create_user_by_name(user_name):
    user_vector_manager.create_user(user_name)
    return {'result': True}

@bottle.route('/user/exists/<user_name>')
def exists_user_by_name(user_name):
    d = {}
    d['exists'] = user_vector_manager.has_user_with_name(user_name)
    result = {'result': d}
    return result

@bottle.route('/user/remove/<user_name>', method='DELETE')
def remove_user_by_name(user_name):
    try:
        user_id = user_vector_manager.get_user_id_for_name(user_name)
        user_vector_manager.remove_user(user_id)
    except:
        return {'result': False}
    return {'result': True}

@bottle.route('/user/createifnotexist/<user_name>')
def create_user_if_not_exists(user_name):
    if not user_vector_manager.has_user_with_name(user_name):
        create_user_by_name(user_name)
    return {'result': True}

@bottle.route('/user/setpreference/<user_name>/<product_id:int>')
def add_preference_to_user(user_name, product_id):
    user_id = user_vector_manager.get_user_id_for_name(user_name)
    user_vector_manager.set_user_preference(user_id, product_id, True)
    return {'result': True}

@bottle.route('/user/setnopreference/<user_name>/<product_id:int>')
def add_preference_to_user(user_name, product_id):
    user_id = user_vector_manager.get_user_id_for_name(user_name)
    user_vector_manager.set_user_preference(user_id, product_id, False)
    return {'result': True}

@bottle.route('/user/update/<user_name>')
def get_user_update(user_name):
    user_id = user_vector_manager.get_user_id_for_name(user_name)

    weights = recommender.rocchio.default_weights()

    update_user(user_id, weights)

    return {'result': True}

@bottle.route('/user/update/<user_name>/<alpha:int>/<beta:int>/<gamma:int>')
def get_user_update(user_name, alpha, beta, gamma):
    user_id = user_vector_manager.get_user_id_for_name(user_name)

    if alpha < 0:
        alpha = 0
    elif alpha > 100:
        alpha = 100;
    if beta < 0:
        beta = 0
    elif beta > 100:
        beta = 100
    if gamma < 0:
        gamma = 0
    elif gamma > 100:
        gamma = 100

    weights = alpha / 100, beta / 100, gamma / 100

    update_user(user_id, weights)
    return {'result': True}

@bottle.route('/user/relevant/<user_name>')
def get_user_preference(user_name):
    user_id = user_vector_manager.get_user_id_for_name(user_name)
    relevant_vectors = user_vector_manager.get_relevant_document_vector_list(user_id)
    relevant_products = [
        product_manager.get_product(v.document_id).as_dictionary()
        for v in relevant_vectors
    ]

    result = {'result': relevant_products}
    return result

@bottle.route('/user/nonrelevant/<user_name>')
def get_user_no_preference(user_name):
    user_id = user_vector_manager.get_user_id_for_name(user_name)
    non_relevant_vectors = user_vector_manager.get_non_relevant_document_vector_list(user_id)
    non_relevant_products = [
        product_manager.get_product(v.document_id).as_dictionary()
        for v in non_relevant_vectors
    ]

    result = {'result': non_relevant_products}
    return result

@bottle.route('/recommendations/<user_name>/<k:int>')
def get_recommendation(user_name, k):
    vector = user_vector_manager.get_user_vector_for_name(user_name)
    others = product_vector_manager.get_all_vectors()

    #distance_function = recommender.vector.arithmetic.hamming_distance
    #distance_function = recommender.vector.arithmetic.euclidean_distance

    recommendations = vector_arithmetic.k_nearest_neighbours(k, vector, others)
    products = [
        product_manager.get_product(vector.document_id).as_dictionary()
        for _, vector in recommendations
    ]
    result = {'result': products}
    return result

database_manager = None
product_manager = None
product_vector_manager = None
document_manager = None
user_vector_manager = None
term_manager = None

random_generator = None

vector_arithmetic = recommender.vector.arithmetic

def run(database_path, host, port):
    _init(database_path)
    bottle.run(host=host, port=port, debug=True)

def _init(database_path):
    global database_manager
    global product_manager
    global product_vector_manager
    global document_manager
    global user_vector_manager
    global term_manager

    global random_generator

    database_manager = DatabaseManager(database_path)
    product_manager = database_manager.get_product_manager()
    product_vector_manager = database_manager.get_product_vector_manager()
    document_manager = database_manager.get_document_manager()
    user_vector_manager  = database_manager.get_user_vector_manager()
    term_manager = database_manager.get_term_manager()

    random_generator = random.Random()

def update_user(user_id, weights):
    user_vector = user_vector_manager.get_user_vector_for_id(user_id)
    relevant = user_vector_manager.get_relevant_document_vector_list(user_id)
    non_relevant = user_vector_manager.get_non_relevant_document_vector_list(user_id)

    uvector = recommender.rocchio.algorithm.calculate(user_vector, relevant, non_relevant, weights)
    user_vector_manager.update_user_vector(user_id, uvector);
    pass


import math
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

#Dane początkowe
k1 = 1
m = 1
h = 0.05

x0 = 10
vx0 = 0

w1 = math.sqrt(k1/m)
A1 = math.sqrt((vx0*vx0)/(w1*w1) + (x0*x0))

iloscPunktow = 1000

#oś XY
setXl = 0
setXr = 55
setYl = 49.95
setYr = 50.04

if(vx0 <= 0):
   fi1 = math.acos(x0/A1) * 180/math.pi
else:
   fi1 = -math.acos(x0/A1) * 180/math.pi

#Wypisanie danych poczatkowych
print("\nk1 = {0}\nm = {1}\nh = {2}\nx0 = {3}\nvx0 = {4}\nw1 = {5}\nA1 = {6}\nfi1 = {7}"
      .format(k1, m, h, x0, vx0, w1, A1, fi1))
print("\nIlosć punktów = {0}".format(iloscPunktow))

#Czas
time = []
for i in range(0, iloscPunktow+1):
    time.append(round((i*h), 2))
    #print(time[i])

#rozwiazanie dokladne listy oraz wartosci poczatkowe
dokladneX = []
dokladneVX = []
dokladneX.append(x0)
dokladneVX.append(vx0)
dokladneE = []

#metoda eulera listy oraz wartosci poczatkowe
eulerX = []
eulerVX = []
eulerX.append(x0)
eulerVX.append(vx0)
eulerE = []

#metoda punktu posredniego listy oraz wartosci poczatkowe
posredniaX = []
posredniaVX = []
posredniaX.append(x0)
posredniaVX.append(vx0)
posredniaE = []

#metoda verleta listy oraz wartosci poczatkowe
verletX = []
verletVX = []
verletX.append(x0)
verletVX.append(vx0)
verletE = []

#metoda beemana listy oraz wartosci poczatkowe
beemanX = []
beemanVX = []
beemanX.append(x0)
beemanVX.append(vx0)
beemanE = []

#uzupelnianie list
for i in range(1, iloscPunktow+1):
    #dokladna
    dokX = A1 * math.cos(w1 * time[i] + fi1 * math.pi / 180)
    dokVX = -A1 * w1 * math.sin(w1 * time[i] + fi1 * math.pi/180)
    dokladneX.append(dokX)
    dokladneVX.append(dokVX)
    #euler
    eulX = eulerX[i - 1] + eulerVX[i - 1] * h
    eulVX = eulerVX[i - 1] - (w1 * w1) * eulerX[i - 1] * h
    eulerX.append(eulX)
    eulerVX.append(eulVX)
    #posrednia
    posX = posredniaX[i-1] + posredniaVX[i-1] * h - 0.5 * (w1 * w1) * posredniaX[i-1] * (h * h)
    posVX = posredniaVX[i-1] - (w1 * w1) * posredniaX[i-1] * h
    posredniaX.append(posX)
    posredniaVX.append(posVX)
    #verlet
    verX = verletX[i - 1] + verletVX[i - 1] * h - 0.5 * (w1 * w1) * verletX[i - 1] * (h * h)
    verletX.append(verX)
    verVX = verletVX[i - 1] - 0.5 * (w1 * w1) * (verletX[i - 1] + verletX[i]) * h
    verletVX.append(verVX)
    #beeman
    # z verleta liczone
    if(i == 1):
        beemanX.append(verletX[1])
        beemanVX.append(verletVX[1])
    else:
        bemX = beemanX[i - 1] + beemanVX[i - 1] * h + (w1 * w1) * (beemanX[i - 2] - 4 * beemanX[i - 1]) * (h * h)/6
        beemanX.append(bemX)
        bemVX = beemanVX[i - 1] + (w1 * w1) * (beemanX[i - 2] - 5 * beemanX[i - 1] - 2 * beemanX[i]) * h/6
        beemanVX.append(bemVX)

#energia

for i in range(0, iloscPunktow+1):
    dokE = 0.5 * k1 * (A1*A1)
    dokladneE.append(dokE)
    eulE = m * (eulerVX[i] * eulerVX[i])/2 + k1 * (eulerX[i] * eulerX[i]/2)
    eulerE.append(eulE)
    posE = m * (posredniaVX[i] * posredniaVX[i])/2 + k1 * (posredniaX[i] * posredniaX[i]/2)
    posredniaE.append(posE)
    verE = m * (verletVX[i] * verletVX[i])/2 + k1 * (verletX[i] * verletX[i]/2)
    verletE.append(verE)
    bemE = m * (beemanVX[i] * beemanVX[i])/2 + k1 * (beemanX[i] * beemanX[i]/2)
    beemanE.append(bemE)

#Animacja
xdata = []
ydata = []

xdata2 = []
ydata2 = []

xdata3 = []
ydata3 = []

font1 = {'family': 'serif', 'color': 'blue', 'size': 20}
font2 = {'family': 'serif', 'color': 'darkred', 'size': 15}
fig, ax = plt.subplots()
ax.set_xlim(setXl, setXr)
ax.set_ylim(setYl, setYr)
plt.title("Energia całkowita oscylatora", fontdict=font1)
plt.xlabel("t", fontdict = font2)
plt.ylabel("E", fontdict = font2)

line, = ax.plot(0, 0, '.') #niebieski
line2, = ax.plot(0, 0, 'r.') #czerwony
line3, = ax.plot(0, 0, 'g.') #zielony

line.set_label('rozwiązanie Dokładne')
line2.set_label('metoda Verleta')
line3.set_label('metoda Beemana')
ax.legend()

def animation_frame(i):
    xdata.append(time[i])
    ydata.append(dokladneE[i])

    xdata2.append(time[i])
    ydata2.append(verletE[i])

    xdata3.append(time[i])
    ydata3.append(beemanE[i])

    line.set_xdata(xdata)
    line.set_ydata(ydata)

    line2.set_xdata(xdata2)
    line2.set_ydata(ydata2)

    line3.set_xdata(xdata3)
    line3.set_ydata(ydata3)

    return line, line2, line3,

animation = FuncAnimation(fig, func = animation_frame, frames = np.arange(0, iloscPunktow + 1, 1), interval = 5)
plt.show()


from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

from .conv_tasnet import TCN, GatedTCN
from .lobe.activation import get_activation
from .lobe.norm import get_norm
from .lobe.rnn import FSMN, ConditionFSMN

class Unet(nn.Module):
    """
    Generic_Args:
        input_type: Real or RI(real+image)
        input_dim: input feature dimension
        activation_type: activation function
        norm_type: normalization function
        dropout: if not 0, add dropout in down-CNN layers

    Unet_Args:
        channels: controlled input/output channel for Unet
        kernel_t: kernel size in time axis for each down cnn layer
        kernel_f: kernel size in freq axis for each down/up cnn layer
        stride_t: stride size in time axis for each down cnn layer
        stride_f: stride size in freq axis for each down/up cnn layer
        dilation_t: dilation size in time axis for each down cnn layer
        dilation_f: dilation size in freq axis for each down/up cnn layer
        delay: add lookahead frames in each down cnn layers, if 0 means causal cnn operation
        transpose_t_size: the kernel size of ConvTranspose2d's time axis for up cnn layer
        skip_conv
    """

    def __init__(
        self,
        input_type: str = "RI",
        input_dim: int = 512,
        activation_type: str = "PReLU",
        norm_type: str = "bN2d",
        dropout: float = 0.05,
        channels: Tuple = (1, 1, 8, 8, 16, 16),
        transpose_t_size: int = 2,
        skip_conv: bool = False,
        kernel_t: Tuple = (5, 1, 9, 1, 1),
        stride_t: Tuple = (1, 1, 1, 1, 1),
        dilation_t: Tuple = (1, 1, 1, 1, 1),
        kernel_f: Tuple = (1, 5, 1, 5, 1),
        stride_f: Tuple = (1, 4, 1, 4, 1),
        dilation_f: Tuple = (1, 1, 1, 1, 1),
        delay: Tuple = (0, 0, 1, 0, 0),
        multi_output: int = 1,
    ):
        super().__init__()
        assert (
            len(kernel_t)
            == len(kernel_f)
            == len(stride_t)
            == len(stride_f)
            == len(dilation_t)
            == len(dilation_f)
        )
        self.input_type = input_type
        self.input_dim = input_dim
        self.multi_output = multi_output
        self.activation_type = activation_type
        self.norm_type = norm_type
        self.dropout = dropout
        self.skip_conv = skip_conv

        # Structure information
        self.kernel_t = kernel_t
        self.kernel_f = kernel_f
        self.stride_t = stride_t
        self.stride_f = stride_f
        self.dilation_t = dilation_t
        self.dilation_f = dilation_f
        self.transpose_t_size = transpose_t_size

        active_cls = get_activation(activation_type.lower())
        norm_cls = get_norm(norm_type)

        self.n_cnn = len(kernel_t)
        self.channels = list(channels)
        self.kernel = list(
            zip(kernel_f, kernel_t)
        )  # each layer's kernel size (freq, time)
        self.delay = delay  # how much delay for each layer
        self.dilation = list(zip(dilation_f, dilation_t))
        self.stride = list(zip(stride_f, stride_t))
        self.t_kernel = transpose_t_size

        # Check relationship between feature-type and input-channel
        if input_type.lower() == "ri":
            self.num_freq = input_dim // 2
            self.channels[0] = self.channels[0] * 2  # will expand RI channel

        elif input_type.lower() == "real":
            self.num_freq = input_dim

        else:
            raise TypeError("Input feature type should be RI-concate, RI-stack or Real")

        # CNN-down, downsample in frequency axis
        self.cnn_down = nn.ModuleList()
        for i in range(self.n_cnn):
            encode = []
            freq_pad = (
                self.kernel[i][0] // 2,
                self.kernel[i][0] // 2,
            )  # center padding in frequency axis
            time_pad = (self.kernel[i][1] - self.delay[i] - 1, self.delay[i])
            encode += [
                nn.ZeroPad2d(time_pad + freq_pad),  # (left, right, top, down)
                nn.Conv2d(
                    self.channels[i],
                    self.channels[i + 1],
                    kernel_size=self.kernel[i],
                    stride=self.stride[i],
                    dilation=self.dilation[i],
                ),
                norm_cls(self.channels[i + 1]),
                active_cls(),
                nn.Dropout(self.dropout),
            ]

            self.cnn_down.append(nn.Sequential(*encode))

        # CNN-up, upsample in frequency axis
        self.cnn_up = nn.ModuleList()
        skip_double = 2 if not skip_conv else 1
        skip_double = [skip_double] * self.n_cnn

        for i in reversed(range(self.n_cnn)):
            s, _ = self.stride[i]
            k = self.kernel[i][0]
            p = k // 2
            op = s - k + 2 * p
            encode = []
            if i != 0:
                encode += [
                    nn.ConvTranspose2d(
                        self.channels[i + 1] * skip_double[i],
                        self.channels[i],
                        kernel_size=(k, self.t_kernel),
                        stride=self.stride[i],
                        dilation=self.dilation[i],
                        padding=(p, 0),
                        output_padding=(op, 0),
                    ),
                    norm_cls(self.channels[i]),
                    active_cls(),
                ]

            else:
                # linear output
                encode += [
                    nn.ConvTranspose2d(
                        self.channels[i + 1] * skip_double[i],
                        self.channels[i] * self.multi_output,
                        kernel_size=(k, self.t_kernel),
                        stride=self.stride[i],
                        dilation=self.dilation[i],
                        padding=(p, 0),
                        output_padding=(op, 0),
                    )
                ]

            self.cnn_up.append(nn.Sequential(*encode))

        if skip_conv:
            self.skip_cnn = nn.ModuleList()
            for i in reversed(range(self.n_cnn)):
                encode = []
                encode += [
                    nn.Conv2d(
                        self.channels[i + 1],
                        self.channels[i + 1],
                        kernel_size=(1, 1),
                        stride=1,
                    ),
                    active_cls(),
                ]

                self.skip_cnn.append(nn.Sequential(*encode))

    def shape_info(self):
        # input_shape = [N, ch, C, T]
        # conv-transpose output size is:
        #   (freq): (input_shape[2] -1) * stride[0] - 2*padding[0] + dilation[0] * (kernel_size[0]-1) + output_padding[0] + 1
        #   (time): (input_shape[2] -1) * stride[1] - 2*padding[1] + dilation[1] * (kernel_size[1]-1) + output_padding[1] + 1
        down_shape = [self.num_freq]
        for i in range(self.n_cnn):
            stride, _ = self.stride[i]
            if down_shape[i - 1] % stride == 0:
                _f = down_shape[-1] // stride
            else:
                _f = down_shape[-1] // stride
                _f += 1
            down_shape.append(_f)

        up_shape = [_f]
        for i in range(self.n_cnn):
            stride, _ = self.stride[-i - 1]
            kernel_size = self.kernel[-i - 1][0]
            padding = kernel_size // 2
            output_padding = stride - kernel_size + 2 * padding
            _f = (
                (up_shape[-1] - 1) * stride
                - 2 * padding
                + self.dilation[-i - 1][0] * (kernel_size - 1)
                + output_padding
                + 1
            )
            up_shape.append(_f)

        return down_shape, up_shape

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: input tensor shape [N, C, T]

        Returns:
            output tensor has shape [N, C, T]
        """
        if self.input_type.lower() == "ri":
            _re, _im = torch.chunk(x, 2, dim=-2)
            x = torch.stack([_re, _im], dim=1)  # [N, C, T] -> [N, 2, C, T]
        else:
            if x.dim() == 3:
                x = x.unsqueeze(1)  # [N, 1, C, T]

        skip = [x.clone()]

        # forward CNN-down layers
        for cnn_layer in self.cnn_down:
            x = cnn_layer(x)  # [N, ch, C, T]
            skip.append(x)

        # forward CNN-up layers
        for i, cnn_layer in enumerate(self.cnn_up):
            if self.skip_conv:
                x += self.skip_cnn[i](skip[-i - 1])
            else:
                x = torch.cat([x, skip[-i - 1]], dim=1)

            x = cnn_layer(x)
            if self.t_kernel != 1:
                x = x[
                    ..., : -(self.t_kernel - 1)
                ]  # transpose-conv with t-kernel size would increase (t-1) length

        if self.multi_output != 1:
            batch, ch, fdim, tdim = x.shape
            x = x.reshape(batch, self.multi_output, -1, fdim, tdim)
            if self.input_type.lower() == "ri":
                _re = x[:, :, 0, :, :]
                _im = x[:, :, 1, :, :]
                x = torch.cat([_re, _im], dim=2)

            else:
                x = x.squeeze(2)  # [N, M, 1, C, T] -> [N, C, T]

        else:
            if self.input_type.lower() == "ri":
                _re = x[:, 0, :, :]
                _im = x[:, 1, :, :]
                x = torch.cat([_re, _im], dim=1)

            else:
                x = x.squeeze(1)  # [N, 1, C, T] -> [N, C, T]

        return x

    @property
    def get_args(self) -> Dict:
        return {
            "input_type": self.input_type,
            "input_dim": self.input_dim,
            "activation_type": self.activation_type,
            "norm_type": self.norm_type,
            "dropout": self.dropout,
            "channels": self.channels,
            "transpose_t_size": self.transpose_t_size,
            "skip_conv": self.skip_conv,
            "kernel_t": self.kernel_t,
            "stride_t": self.stride_t,
            "dilation_t": self.dilation_t,
            "kernel_f": self.kernel_f,
            "stride_f": self.stride_f,
            "dilation_f": self.dilation_f,
            "delay": self.delay,
            "multi_output": self.multi_output,
        }

class UnetTcn(Unet):
    """
    Improve temporal modeling ability by inserting a TCN inside an Unet model.

    Args:
        embed_dim: Embedding feature dimension.
        embed_norm: If True, applies the 2-norm on the input embedding.
    """

    def __init__(
        self,
        embed_dim: int = 0,
        embed_norm: bool = False,
        input_type: str = "RI",
        input_dim: int = 512,
        activation_type: str = "PReLU",
        norm_type: str = "bN2d",
        dropout: float = 0.05,
        channels: Tuple = (1, 1, 8, 8, 16, 16),
        transpose_t_size: int = 2,
        transpose_delay: bool = False,
        skip_conv: bool = False,
        kernel_t: Tuple = (5, 1, 9, 1, 1),
        stride_t: Tuple = (1, 1, 1, 1, 1),
        dilation_t: Tuple = (1, 1, 1, 1, 1),
        kernel_f: Tuple = (1, 5, 1, 5, 1),
        stride_f: Tuple = (1, 4, 1, 4, 1),
        dilation_f: Tuple = (1, 1, 1, 1, 1),
        delay: Tuple = (0, 0, 1, 0, 0),
        tcn_layer: str = "normal",
        tcn_kernel: int = 3,
        tcn_dim: int = 256,
        tcn_dilated_basic: int = 2,
        per_tcn_stack: int = 5,
        repeat_tcn: int = 4,
        tcn_with_embed: List = [1, 0, 0, 0, 0],
        tcn_use_film: bool = False,
        tcn_norm: str = "gLN",
        dconv_norm: str = "gGN",
        causal: bool = False,
    ):
        super().__init__(
            input_type,
            input_dim,
            activation_type,
            norm_type,
            dropout,
            channels,
            transpose_t_size,
            skip_conv,
            kernel_t,
            stride_t,
            dilation_t,
            kernel_f,
            stride_f,
            dilation_f,
            delay,
        )

        self.embed_dim = embed_dim
        self.embed_norm = embed_norm
        self.tcn_layer = tcn_layer
        self.tcn_dim = tcn_dim
        self.tcn_kernel = tcn_kernel
        self.per_tcn_stack = per_tcn_stack
        self.repeat_tcn = repeat_tcn
        self.tcn_dilated_basic = tcn_dilated_basic
        self.tcn_with_embed = tcn_with_embed
        self.tcn_norm = tcn_norm
        self.dconv_norm = dconv_norm
        self.tcn_use_film = tcn_use_film
        self.causal = causal
        self.transpose_delay = transpose_delay

        # TCN module's
        temporal_input_dim = self.num_freq
        for stride, _ in self.stride:
            if temporal_input_dim % stride == 0:
                temporal_input_dim //= stride
            else:
                temporal_input_dim //= stride
                temporal_input_dim += 1

        temporal_input_dim *= self.channels[-1]  # extend by channel size

        if self.tcn_layer.lower() == "normal":
            tcn_cls = TCN
        elif self.tcn_layer.lower() == "gated":
            print("GatedTCN would ignore dconv_norm configuration.")
            tcn_cls = GatedTCN
        else:
            raise NameError

        assert per_tcn_stack == len(tcn_with_embed)
        self.tcn_list = nn.ModuleList()
        for _ in range(repeat_tcn):
            _tcn = []

            for i in range(per_tcn_stack):
                if tcn_with_embed[i]:
                    if self.tcn_layer.lower() == "normal":
                        _tcn.append(
                            tcn_cls(
                                temporal_input_dim,
                                tcn_dim,
                                kernel=tcn_kernel,
                                dilation=tcn_dilated_basic ** i,
                                emb_dim=embed_dim,
                                causal=causal,
                                tcn_norm=tcn_norm,
                                dconv_norm=dconv_norm,
                            )
                        )
                    else:
                        _tcn.append(
                            tcn_cls(
                                temporal_input_dim,
                                tcn_dim,
                                kernel=tcn_kernel,
                                dilation=tcn_dilated_basic ** i,
                                emb_dim=embed_dim,
                                causal=causal,
                                tcn_norm=tcn_norm,
                                use_film=tcn_use_film,
                            )
                        )
                else:
                    if self.tcn_layer.lower() == "normal":
                        _tcn.append(
                            tcn_cls(
                                temporal_input_dim,
                                tcn_dim,
                                kernel=tcn_kernel,
                                dilation=tcn_dilated_basic ** i,
                                emb_dim=0,
                                causal=causal,
                                tcn_norm=tcn_norm,
                                dconv_norm=dconv_norm,
                            )
                        )
                    else:
                        _tcn.append(
                            tcn_cls(
                                temporal_input_dim,
                                tcn_dim,
                                kernel=tcn_kernel,
                                dilation=tcn_dilated_basic ** i,
                                emb_dim=0,
                                causal=causal,
                                tcn_norm=tcn_norm,
                                use_film=False,
                            )
                        )

            self.tcn_list.append(nn.ModuleList(_tcn))

    def forward(
        self, x: torch.Tensor, dvec: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        Args:
            x: input tensor shape [N, C, T]
            dvec: conditional tensor shape [N, C]

        Returns:
            output tensor has shape [N, C, T]
        """
        # normalize
        if self.embed_norm and dvec is not None:
            dvec = F.normalize(dvec, p=2, dim=1)

        if self.input_type.lower() == "ri":
            _re, _im = torch.chunk(x, 2, dim=-2)
            x = torch.stack([_re, _im], dim=1)  # [N, C, T] -> [N, 2, C, T]
        else:
            if x.dim() == 3:
                x = x.unsqueeze(1)  # [N, 1, C, T]

        skip = [x.clone()]

        # forward CNN-down layers
        for cnn_layer in self.cnn_down:
            x = cnn_layer(x)  # [N, ch, C, T]
            skip.append(x)

        # forward TCN block
        N_ori, ch, C_ori, T = x.shape
        x = x.reshape(N_ori, ch * C_ori, T)

        for r in range(self.repeat_tcn):
            for i in range(len(self.tcn_list[r])):
                if self.tcn_with_embed[i]:
                    x = self.tcn_list[r][i](x, dvec)
                else:
                    x = self.tcn_list[r][i](x)

        x = x.reshape(N_ori, ch, C_ori, T)

        # forward CNN-up layers
        for i, cnn_layer in enumerate(self.cnn_up):
            if self.skip_conv:
                x += self.skip_cnn[i](skip[-i - 1])
            else:
                x = torch.cat([x, skip[-i - 1]], dim=1)

            x = cnn_layer(x)
            if self.t_kernel != 1:
                if self.transpose_delay:
                    x = x[
                        ..., (self.t_kernel - 1) :
                    ]  # transpose-conv with t-kernel size would increase (t-1) length
                else:
                    x = x[
                        ..., : -(self.t_kernel - 1)
                    ]  # transpose-conv with t-kernel size would increase (t-1) length

        if self.input_type.lower() == "ri":
            _re = x[:, 0, :, :]
            _im = x[:, 1, :, :]
            x = torch.cat([_re, _im], dim=1)

        else:
            x = x.squeeze(1)  # [N, 1, C, T] -> [N, C, T]

        return x

    @property
    def get_args(self) -> Dict:
        return {
            "input_type": self.input_type,
            "input_dim": self.input_dim,
            "activation_type": self.activation_type,
            "norm_type": self.norm_type,
            "dropout": self.dropout,
            "channels": self.channels,
            "transpose_t_size": self.transpose_t_size,
            "transpose_delay": self.transpose_delay,
            "skip_conv": self.skip_conv,
            "kernel_t": self.kernel_t,
            "stride_t": self.stride_t,
            "dilation_t": self.dilation_t,
            "kernel_f": self.kernel_f,
            "stride_f": self.stride_f,
            "dilation_f": self.dilation_f,
            "delay": self.delay,
            "embed_dim": self.embed_dim,
            "embed_norm": self.embed_norm,
            "tcn_norm": self.tcn_norm,
            "dconv_norm": self.dconv_norm,
            "tcn_layer": self.tcn_layer,
            "tcn_dim": self.tcn_dim,
            "tcn_kernel": self.tcn_kernel,
            "tcn_dilated_basic": self.tcn_dilated_basic,
            "repeat_tcn": self.repeat_tcn,
            "per_tcn_stack": self.per_tcn_stack,
            "tcn_with_embed": self.tcn_with_embed,
            "tcn_use_film": self.tcn_use_film,
            "causal": self.causal,
        }

class UnetFsmn(Unet):
    """
    Improve temporal modeling ability by inserting a FSMN inside an Unet model.

    Args:
        embed_dim: Embedding feature dimension.
        embed_norm: If True, applies the 2-norm on the input embedding.
    """

    def __init__(
        self,
        embed_dim: int = 0,
        embed_norm: bool = False,
        input_type: str = "RI",
        input_dim: int = 512,
        activation_type: str = "PReLU",
        norm_type: str = "bN2d",
        dropout: float = 0.05,
        channels: Tuple = (1, 1, 8, 8, 16, 16),
        transpose_t_size: int = 2,
        transpose_delay: bool = False,
        skip_conv: bool = False,
        kernel_t: Tuple = (5, 1, 9, 1, 1),
        stride_t: Tuple = (1, 1, 1, 1, 1),
        dilation_t: Tuple = (1, 1, 1, 1, 1),
        kernel_f: Tuple = (1, 5, 1, 5, 1),
        stride_f: Tuple = (1, 4, 1, 4, 1),
        dilation_f: Tuple = (1, 1, 1, 1, 1),
        delay: Tuple = (0, 0, 1, 0, 0),
        fsmn_l_context: int = 3,
        fsmn_r_context: int = 0,
        fsmn_dim: int = 256,
        num_fsmn: int = 8,
        fsmn_with_embed: List = [1, 1, 1, 1, 1, 1, 1, 1],
        fsmn_norm: str = "gLN",
        use_film: bool = True,
    ):
        super().__init__(
            input_type,
            input_dim,
            activation_type,
            norm_type,
            dropout,
            channels,
            transpose_t_size,
            skip_conv,
            kernel_t,
            stride_t,
            dilation_t,
            kernel_f,
            stride_f,
            dilation_f,
            delay,
        )

        self.transpose_delay = transpose_delay
        self.embed_dim = embed_dim
        self.embed_norm = embed_norm
        self.fsmn_l_context = fsmn_l_context
        self.fsmn_r_context = fsmn_r_context
        self.fsmn_dim = fsmn_dim
        self.num_fsmn = num_fsmn
        self.fsmn_with_embed = fsmn_with_embed
        self.fsmn_norm = fsmn_norm
        self.use_film = use_film

        # FSMN module's
        temporal_input_dim = self.num_freq
        for stride, _ in self.stride:
            if temporal_input_dim % stride == 0:
                temporal_input_dim //= stride
            else:
                temporal_input_dim //= stride
                temporal_input_dim += 1

        temporal_input_dim *= self.channels[-1]  # extend by channel size

        assert num_fsmn == len(fsmn_with_embed)
        self.fsmn_list = nn.ModuleList()

        for i in range(num_fsmn):
            if fsmn_with_embed[i]:
                self.fsmn_list.append(
                    ConditionFSMN(
                        temporal_input_dim,
                        temporal_input_dim,
                        fsmn_dim,
                        embed_dim,
                        fsmn_l_context,
                        fsmn_r_context,
                        norm_type=fsmn_norm,
                        use_film=use_film,
                    )
                )

            else:
                self.fsmn_list.append(
                    FSMN(
                        temporal_input_dim,
                        temporal_input_dim,
                        fsmn_dim,
                        fsmn_l_context,
                        fsmn_r_context,
                        norm_type=fsmn_norm,
                    )
                )

    def forward(
        self, x: torch.Tensor, dvec: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        Args:
            x: input tensor shape [N, C, T]
            dvec: conditional tensor shape [N, C]

        Returns:
            output tensor has shape [N, C, T]
        """
        # normalize
        if self.embed_norm and dvec is not None:
            dvec = F.normalize(dvec, p=2, dim=1)

        if self.input_type.lower() == "ri":
            _re, _im = torch.chunk(x, 2, dim=-2)
            x = torch.stack([_re, _im], dim=1)  # [N, C, T] -> [N, 2, C, T]
        else:
            if x.dim() == 3:
                x = x.unsqueeze(1)  # [N, 1, C, T]

        skip = [x.clone()]

        # forward CNN-down layers
        for cnn_layer in self.cnn_down:
            x = cnn_layer(x)  # [N, ch, C, T]
            skip.append(x)

        # forward FSMN block
        N_ori, ch, C_ori, T = x.shape
        x = x.reshape(N_ori, ch * C_ori, T)
        memory = None
        for i in range(len(self.fsmn_list)):
            if self.fsmn_with_embed[i]:
                x, memory = self.fsmn_list[i](x, dvec, memory)
            else:
                x, memory = self.fsmn_list[i](x, memory)

        x = x.reshape(N_ori, ch, C_ori, T)

        # forward CNN-up layers
        for i, cnn_layer in enumerate(self.cnn_up):
            if self.skip_conv:
                x += self.skip_cnn[i](skip[-i - 1])
            else:
                x = torch.cat([x, skip[-i - 1]], dim=1)

            x = cnn_layer(x)
            if self.t_kernel != 1:
                if self.transpose_delay:
                    x = x[
                        ..., (self.t_kernel - 1) :
                    ]  # transpose-conv with t-kernel size would increase (t-1) length
                else:
                    x = x[
                        ..., : -(self.t_kernel - 1)
                    ]  # transpose-conv with t-kernel size would increase (t-1) length

        if self.input_type.lower() == "ri":
            _re = x[:, 0, :, :]
            _im = x[:, 1, :, :]
            x = torch.cat([_re, _im], dim=1)

        else:
            x = x.squeeze(1)  # [N, 1, C, T] -> [N, C, T]

        return x

    @property
    def get_args(self) -> Dict:
        return {
            "input_type": self.input_type,
            "input_dim": self.input_dim,
            "activation_type": self.activation_type,
            "norm_type": self.norm_type,
            "dropout": self.dropout,
            "channels": self.channels,
            "transpose_t_size": self.transpose_t_size,
            "transpose_delay": self.transpose_delay,
            "skip_conv": self.skip_conv,
            "kernel_t": self.kernel_t,
            "stride_t": self.stride_t,
            "dilation_t": self.dilation_t,
            "kernel_f": self.kernel_f,
            "stride_f": self.stride_f,
            "dilation_f": self.dilation_f,
            "delay": self.delay,
            "embed_dim": self.embed_dim,
            "embed_norm": self.embed_norm,
            "fsmn_l_context": self.fsmn_l_context,
            "fsmn_r_context": self.fsmn_r_context,
            "fsmn_dim": self.fsmn_dim,
            "num_fsmn": self.num_fsmn,
            "fsmn_with_embed": self.fsmn_with_embed,
            "fsmn_norm": self.fsmn_norm,
            "use_film": self.use_film,
        }


import os
import sys

if not "DEVITO_OPENMP" in os.environ or os.environ["DEVITO_OPENMP"] != "1":
    print("*** WARNING: Devito OpenMP environment variable has not been set ***", file=sys.stderr)

import numpy as np
from sympy import Matrix, Eq, solve
import progressbar
from devito import TimeData, Operator, t, x, y, z, logger as devito_logger, parameters as devito_parameters

from . import sim

devito_logger.set_log_level('WARNING')

def vector_laplacian(u):
    return Matrix([u[0].dx2 + u[0].dy2 + u[0].dz2,
                   u[1].dx2 + u[1].dy2 + u[1].dz2,
                   u[2].dx2 + u[2].dy2 + u[2].dz2])

def vector_gradient(u):
    return u[0].dx**2 + u[0].dy**2 + u[0].dz**2 + u[1].dx**2 + u[1].dy**2 + u[1].dz**2 + u[2].dx**2 + u[2].dy**2 + u[2].dz**2

def curl(u):
    return Matrix([u[2].dy - u[1].dz,
                   u[0].dz - u[2].dx,
                   u[1].dx - u[0].dy])

expression_cache = {}

class Sim(sim.Sim):

    framework_name = "Devito"

    @property
    def data_shape(self):
        # Devito doesn't like numpy types for the grid dimensions, and it needs to be a tuple, so shape needs to be converted
        return tuple(int(i) for i in self.grid_params.n)

    def data_matrix(self, settings):
        return Matrix([TimeData(name='m_x', **settings),
            TimeData(name='m_y', **settings),
            TimeData(name='m_z', **settings)])

    def generate_step_kernel(self):
        settings = {"shape":self.buffer_dims, "space_order":2}
        m = self.data_matrix(settings)
        c = 2 / (self.mu0 * self.sim_params.Ms)
        zeeman = Matrix(self.sim_params.H)
        exchange = self.sim_params.A * c * vector_laplacian(m)
        e = Matrix(self.sim_params.e)
        anisotropy = self.sim_params.K * c * m.dot(e) * e
        dmi = self.sim_params.D * c * curl(m)
        heff = zeeman + exchange + anisotropy + dmi
        crossHeff = m.cross(heff)
        dmdt_rhs = -self.gamma0 / (1 + self.sim_params.alpha**2) * (crossHeff + self.sim_params.alpha * m.cross(crossHeff))
        dmdt_lhs = Matrix([TimeData(name='dmdt_x', **settings),
            TimeData(name='dmdt_y', **settings),
            TimeData(name='dmdt_z', **settings)])
        dmdt_correction = self.correction * dmdt_lhs.dot(dmdt_lhs)**0.5 * (1 - m.dot(m)) * m
        string_llg = str(dmdt_rhs) + str(dmdt_correction)

        if string_llg in expression_cache:
            update = expression_cache[string_llg]
        else:
            update = []
            if self.correction > 0:
                # if using correction solve in 2 steps; calculate dmdt, then calculate m[t+1] = dmdt + correction
                for i, dmdti in enumerate(dmdt_lhs):
                    update.append(Eq(dmdti, dmdt_rhs[i]))
                llg_eqn = Matrix([mi.dt for mi in m]) - (dmdt_lhs + dmdt_correction)
            else:
                # if not using correction; m[t+1] = dmdt
                llg_eqn = Matrix([mi.dt for mi in m]) - dmdt_rhs

            print("Solving LLG Sympy expressions ...", file=sys.stderr)
            with progressbar.ProgressBar(max_value=len(m)) as bar:
                for i, mi in enumerate(m):
                    update.append(Eq(mi.forward, solve(llg_eqn[i], mi.forward)[0]))
                    bar.update(i + 1)
            expression_cache[string_llg] = update

        bcs = []
        nx, ny, nz = self.buffer_dims
        if self.periodic_boundary:
            for mi in m:
                bcs += [Eq(mi.indexed[t, x, y, 0], mi.indexed[t, x, y, nz - 2])]
                bcs += [Eq(mi.indexed[t, x, y, nz - 1], mi.indexed[t, x, y, 1])]
                bcs += [Eq(mi.indexed[t, x, 0, z], mi.indexed[t, x, ny - 2, z])]
                bcs += [Eq(mi.indexed[t, x, ny - 1, z], mi.indexed[t, x, 1, z])]
                bcs += [Eq(mi.indexed[t, 0, y, z], mi.indexed[t, nx - 2, y, z])]
                bcs += [Eq(mi.indexed[t, nx - 1, y, z], mi.indexed[t, 1, y, z])]
        else:
            for mi in m:
                bcs += [Eq(mi.indexed[t, x, y, 0], 0.)]
                bcs += [Eq(mi.indexed[t, x, y, nz - 1], 0.)]
                bcs += [Eq(mi.indexed[t, x, 0, z], 0.)]
                bcs += [Eq(mi.indexed[t, x, ny - 1, z], 0.)]
                bcs += [Eq(mi.indexed[t, 0, y, z], 0.)]
                bcs += [Eq(mi.indexed[t, nx - 1, y, z], 0.)]

        dx, dy, dz = self.grid_params.d
        dt = self.time_params.d
        subs = {x.spacing: dx, y.spacing: dy, z.spacing: dz, t.spacing: dt}
        op = Operator(bcs + update, subs=subs)

        # Call op trigger compilation
        op(time=1)

        def step(f, t):
            for i, mi in enumerate(m):
                mi.data[(0, ) + self.buffer_slice] = f[i]
            op(time=self.save_every + 1)
            for i, mi in enumerate(m):
                t[i] = mi.data[(self.save_every % 2, ) + self.buffer_slice]

        return step

"""
    def energy_expr(self, m):
        dV = self.grid_params.prod_d
        e = Matrix(self.sim_params.e)
        H = Matrix(self.sim_params.H)
        Kc = dV * -self.sim_params.K
        Ac = dV * self.sim_params.A
        Dc = dV * -self.sim_params.D
        Hc = dV * -self.mu0 * self.sim_params.Ms

        return {"Zeeman":Hc * m.dot(H),
            "Exchange":Ac * vector_gradient(m),
            "Anisotropy":Kc * (m.dot(e))**2,
            "DMI":Dc * m.dot(curl(m))}
    def generate_energy_kernel(self):
        settings = {"shape":self.buffer_dims, "space_order":2}
        m = self.data_matrix(settings)
        energy_expr = self.energy_expr(m)
        E = TimeData(name='E', **settings)
        eqn = Eq(E, sum(energy_expr.values()))
        dx, dy, dz = self.grid_params.d
        subs = {x.spacing: dx, y.spacing: dy, z.spacing: dz}
        # turn dle off because some eqns are 1st and some are 2nd order, requiring different bounds.
        op = Operator(eqn, subs=subs, dle=False)

        # Call op trigger compilation
        op()

        def energy(d):
            for i, mi in enumerate(m):
                mi.data[0] = d[i]
            op(time=1)
            return E.data[0]

        return energy

    def generate_detailed_energy_kernel(self, terms):
        def energy(d):
            settings = {"shape":self.buffer_dims, "space_order":2, "time_dim":len(d), "save":True}
            m = self.data_matrix(settings)
            energy_expr = self.energy_expr(m)

            names = [k for k in terms if k in energy_expr]
            symbols = []
            eqns = []
            for key in names:
                symbol = TimeData(name='E_{}'.format(key), **settings)
                symbols.append(symbol)
                eqns.append(Eq(symbol, energy_expr[key]))

            dx, dy, dz = self.grid_params.d
            subs = {x.spacing: dx, y.spacing: dy, z.spacing: dz}
            # turn dle off because some eqns are 1st and some are 2nd order, requiring different bounds.
            op = Operator(eqns, subs=subs, dle=False)

            for i, mi in enumerate(m):
                for j, dj in enumerate(d):
                    mi.data[j] = dj[i]
            op()
            ret = {}
            for i, name in enumerate(names):
                ret[name] = []
                for dj in symbols[i].data:
                    ret[name].append(dj)
            return ret

        return energy
"""


#! /usr/bin/env python
import sys
import csv
import screed
import random
import argparse

import sourmash

import sequtils                 # local import

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('genome')
    parser.add_argument('-e', '--error-rate', type=float, default=.01)
    parser.add_argument('-r', '--read-length', type=int, default=100,
                        help="Length of reads to generate")
    parser.add_argument("-S", "--seed", dest="seed", help="Random seed", type=int,
                        default=1)
    parser.add_argument("-k", "--ksize", default=31, help="k-mer size")
    parser.add_argument("-o", "--output", required=True,
                        help="CSV output of detection curve")

    args = parser.parse_args()

    READLEN=args.read_length
    ERROR_RATE=args.error_rate
    NUM_FRACMINHASH = 5

    random.seed(args.seed)                  # make this reproducible, please.

    records = list(screed.open(args.genome))
    assert len(records) == 1
    record = records[0]

    genome = record.sequence
    len_genome = len(genome)

    total_mh = sourmash.MinHash(0, args.ksize, scaled=1)
    total_mh.add_sequence(genome)
    all_hashes = set(total_mh.hashes)

    # make NUM_FRACMINHASH minhashes each with different mmh seeds
    all_hashes_list = []
    scaled_mh_list = []
    for i in range(NUM_FRACMINHASH):
        smh = sourmash.MinHash(0, args.ksize, scaled=1000, seed=i + 42)
        all_hashes_i = smh.copy_and_clear()
        all_hashes_i.add_sequence(genome)
        scaled_mh_list.append(smh)
        all_hashes_list.append(all_hashes_i)

    print('genome size:', len_genome, file=sys.stderr)
    print('readlen:', READLEN, file=sys.stderr)
    print('error rate:', ERROR_RATE, file=sys.stderr)
    print('num k-mers:', len(total_mh))

    reads_mut = 0
    total_mut = 0

    print(f"Read in template genome {0} of length {1} from {2}".format(record["name"], len_genome, args.genome), file=sys.stderr)
    print(f"Generating reads of length {READLEN} with an error rate of 1 in {ERROR_RATE}", file=sys.stderr)

    it = sequtils.generate_mutated_reads(genome, READLEN, ERROR_RATE)
    it = iter(it)

    fp = open(args.output, 'w', newline="")
    csv_w = csv.writer(fp)
    headers = ['num_reads', 'coverage', 'n_detected', 'f_detected']
    for i in range(NUM_FRACMINHASH):
        headers.append(f"smash_count_{i}")
    csv_w.writerow(headers)
    csv_w.writerow([0, 0, 0, 0] + [0]*NUM_FRACMINHASH)

    n_reads = 0
    total_bp_in_reads = 0
    f01 = len(all_hashes) * 0.1

    remaining_hashes = set(all_hashes)
    while len(remaining_hashes) > f01:
        start, read, read_mutations = next(it)

        if read_mutations:
            reads_mut += 1
            total_mut += read_mutations

        n_reads += 1
        total_bp_in_reads += len(read)

        # first, track _all_ hashes for actual k-mer detection
        mh = total_mh.copy_and_clear()
        mh.add_sequence(read)
        remaining_hashes -= set(mh.hashes)

        n_detected = len(all_hashes) - len(remaining_hashes)
        f_detected = n_detected / len(all_hashes)
        coverage = total_bp_in_reads / len_genome

        # now, track sourmash detection & intersect with legit hashes:
        smash_detection = []
        for smh, all_hashes_i in zip(scaled_mh_list, all_hashes_list):
            smh.add_sequence(read)
            smh_hashes = set(smh.hashes)
            smh_hashes.intersection_update(all_hashes_i.hashes)
            smash_detection.append(len(smh_hashes))

        csv_w.writerow([n_reads, f"{coverage:.4f}", n_detected, f"{f_detected:.4f}"] + smash_detection)

        sys.stdout.write(u'\r\033[K')
        sys.stdout.write(f"...{n_reads} reads, {len(all_hashes)} missing k-mers, {total_bp_in_reads / len_genome:.2f} coverage")
        sys.stdout.flush()

    fp.close()

if __name__ == '__main__':
    sys.exit(main())


# suhaarslan.com
from random import randbytes
class Client:
    def keyControl(self, x):
        # takes first 32 bytes
        return x[:self.__byte]
    # Control Funtions
    def __init__(self, a, b):
        # a, b 32 bytes public key / a, b string
        self.__byte = 128
        self.public_key_1 = bytes(self.keyControl(a.encode()))
        self.public_key_2 = bytes(self.keyControl(b.encode()))
        self.__createPrivateKey()
    def __createPrivateKey(self):
        # private key 32 bytes
        self.__private_key = randbytes(self.__byte)
    def getPrivateKey(self):
        return self.__private_key
    # Key Functions
    def Make(self, onc = 0):
        if onc == 0:
            self.local = hex((int(self.public_key_1.hex(), 16)^int(self.__private_key.hex(), 16))&int(self.public_key_2.hex(), 16)-1)
            return self.local
        else:
            self.l_onc = hex((int(onc, 16)^int(self.__private_key.hex(), 16))&int(self.public_key_2.hex(), 16)-1)
            return self.l_onc
    # Calcs

k1 = """Ua){jk2#N^=yShan.]}:+#'TZL6s!F!WG8A=&-ML{gJ(B>5$xC=X/]H'[6gyNn6*B`4:UB,~)et[">$9:d#9F6nQjcp,!pm5FPP(=VGTXe6U=Ypta&JjrRfE}"/j~g"/"""
k2 = """rt$}Lu9Gdsu:^&>8[2>waMC}g+q[=g~KJ=ymp5"`=:&M-XUDQ&SB3Yc_B-V/5b@_kt(:[=r`98C(r2rE@wA#c_T8k+D>EMqrG5$\_xUaDx)Tr4_J"b{vud+X<9'N<:sB"""

alpha = Client(k1, k2)
alphaCalc = alpha.Make()

beta = Client(k1, k2)
betaCalc = beta.Make()

print(alpha.Make(onc=betaCalc))
print(beta.Make(onc=alphaCalc))


import logging
import pathlib
import sqlite3

logger = logging.getLogger(__name__)

def is_database_exists(db_path):
    return pathlib.Path(db_path).exists()

def open_connection(db_path):
    if is_database_exists(db_path):
        logger.debug(f"Connecting to {db_path}")
        try:
            return sqlite3.connect(db_path)
        except Exception:
            logger.exception(f"Failed to connect to {db_path}")
            raise
    else:
        raise RuntimeError(f"Databse {db_path} doesn't exist")

def close_connection(connection):
    assert connection is not None
    logger.debug("Closing connection")
    connection.close()

def create_database(db_path):
    logger.info(f"Creating empty database at {db_path}")
    if not is_database_exists(db_path):
        try:
            connection = sqlite3.connect(db_path)
        except Exception:
            logging.exception("Failed to create database")
            raise
        else:
            close_connection(connection)
    else:
        raise RuntimeError(f"Database {db_path} already exists")

class DatabaseIO:
    def __init__(self, db_path):
        self._path = db_path
        self._connection = None

    def __enter__(self):
        self._connection = open_connection(self._path)
        return self

    def __exit__(self, exc_type, exc_val, exc_traceback):
        del exc_type
        del exc_val
        del exc_traceback
        close_connection(self._connection)
        self._connection = None
        return False


# __author__ = 'heyin'
# __date__ = '2019/2/14 16:03'
# google翻译rpc服务端代码
import sys

sys.path.append('../')
import json

import grpc
import time
from concurrent import futures
from rpc_server.fanyi import fanyi_pb2, fanyi_pb2_grpc
from rpc_conf import HOST, PORT, ONE_DAY_IN_SECONDS
from core import google

js = google.Py4Js()

class Translate(fanyi_pb2_grpc.TranslateServicer):
    def DoTranslate(self, request, context):
        args = request.text
        args = json.loads(args)
        src = args.get('src')
        dest = args.get('dest')
        cookies = args.get('cookies')

        # 下边内容为谷歌的翻译操作
        ret = google.translate(js, args.get('content'), src, dest, cookies)
        return fanyi_pb2.Data(text=ret)

def serve():
    grpcServer = grpc.server(futures.ThreadPoolExecutor(max_workers=4))
    fanyi_pb2_grpc.add_TranslateServicer_to_server(Translate(), grpcServer)
    grpcServer.add_insecure_port(HOST + ':' + PORT)
    grpcServer.start()
    try:
        while True:
            time.sleep(ONE_DAY_IN_SECONDS)
    except KeyboardInterrupt:
        grpcServer.stop(0)

if __name__ == '__main__':
    serve()


import logging
import paddle.fluid as fluid
import paddle.fluid.dygraph.nn as nn
from utils import build_norm_layer, build_conv_layer, Sequential

class BasicBlock(fluid.dygraph.Layer):
    expansion = 1

    def __init__(self,
                 inplanes,
                 planes,
                 stride=1,
                 dilation=1,
                 downsample=None,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN')):
        super(BasicBlock, self).__init__()

        self.norm1_name, norm1 = build_norm_layer(norm_cfg, planes, postfix=1)
        self.norm2_name, norm2 = build_norm_layer(norm_cfg, planes, postfix=2)

        self.conv1 = build_conv_layer(
            conv_cfg,
            inplanes,
            planes,
            3,
            stride=stride,
            padding=dilation,
            dilation=dilation,
            bias_attr=False)
        self.add_sublayer(self.norm1_name, norm1)

        self.conv2 = build_conv_layer(
            conv_cfg, planes, planes, 3, padding=1, bias_attr=False)
        self.add_sublayer(self.norm2_name, norm2)

        self.relu = fluid.layers.relu
        self.downsample = downsample
        self.stride = stride
        self.dilation = dilation

    @property
    def norm1(self):
        return getattr(self, self.norm1_name)

    @property
    def norm2(self):
        return getattr(self, self.norm2_name)

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.norm1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.norm2(out)

        if self.downsample is not None:
            identity = self.downsample(x)
        out = fluid.layers.elementwise_add(out, identity)
        out = self.relu(out)

        return out

class Bottleneck(fluid.dygraph.Layer):
    expansion = 4

    def __init__(self,
                 inplanes,
                 planes,
                 stride=1,
                 dilation=1,
                 downsample=None,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN')):
        """Bottleneck block for ResNet.
        the stride-two layer is the 3x3 conv layer,.
        """
        super(Bottleneck, self).__init__()
        self.inplanes = inplanes
        self.planes = planes
        self.stride = stride
        self.dilation = dilation
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg

        self.conv1_stride = 1
        self.conv2_stride = stride

        self.norm1_name, norm1 = build_norm_layer(norm_cfg, planes, postfix=1)
        self.norm2_name, norm2 = build_norm_layer(norm_cfg, planes, postfix=2)
        self.norm3_name, norm3 = build_norm_layer(
            norm_cfg, planes * self.expansion, postfix=3)

        self.conv1 = build_conv_layer(
            conv_cfg,
            inplanes,
            planes,
            1,
            stride=1,
            bias_attr=False)
        self.add_sublayer(self.norm1_name, norm1)

        self.conv2 = build_conv_layer(
            conv_cfg,
            planes,
            planes,
            3,
            stride=stride,
            padding=dilation,
            dilation=dilation,
            bias_attr=False)
        self.add_sublayer(self.norm2_name, norm2)

        self.conv3 = build_conv_layer(
            conv_cfg,
            planes,
            planes * self.expansion,
            1,
            bias_attr=False)
        self.add_sublayer(self.norm3_name, norm3)

        self.relu = fluid.layers.relu
        self.downsample = downsample

    @property
    def norm1(self):
        return getattr(self, self.norm1_name)

    @property
    def norm2(self):
        return getattr(self, self.norm2_name)

    @property
    def norm3(self):
        return getattr(self, self.norm3_name)

    def forward(self, x):

        identity = x

        out = self.conv1(x)
        out = self.norm1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.norm2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.norm3(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out = fluid.layers.elementwise_add(out, identity)

        out = self.relu(out)

        return out

def make_res_layer(block,
                   inplanes,
                   planes,
                   blocks,
                   stride=1,
                   dilation=1,
                   conv_cfg=None,
                   norm_cfg=dict(type='BN')):
    downsample = None
    if stride != 1 or inplanes != planes * block.expansion:
        downsample = Sequential(
            build_conv_layer(
                conv_cfg,
                inplanes,
                planes * block.expansion,
                1,
                stride=stride,
                bias_attr=False),
            build_norm_layer(norm_cfg, planes * block.expansion)[1]
        )

    layers = []
    layers.append(
        block(
            inplanes,
            planes,
            stride,
            dilation,
            downsample,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg))
    inplanes = planes * block.expansion
    for i in range(1, blocks):
        layers.append(
            block(
                inplanes,
                planes,
                1,
                dilation,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg))

    return Sequential(*layers)

class ResNet(fluid.dygraph.Layer):
    """ResNet backbone.

    Args:
        depth (int): Depth of resnet, from {18, 34, 50, 101, 152}.
        num_stages (int): Resnet stages, normally 4.
    """

    arch_settings = {
        18: (BasicBlock, (2, 2, 2, 2)),
        34: (BasicBlock, (3, 4, 6, 3)),
        50: (Bottleneck, (3, 4, 6, 3)),
        101: (Bottleneck, (3, 4, 23, 3)),
        152: (Bottleneck, (3, 8, 36, 3))
    }

    def __init__(self,
                 depth,
                 num_stages=4,
                 strides=(1, 2, 2, 2),
                 dilations=(1, 1, 1, 1),
                 out_indices=(0, 1, 2, 3),
                 frozen_stages=-1,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 norm_eval=True,
                 zero_init_residual=True):
        super(ResNet, self).__init__()
        if depth not in self.arch_settings:
            raise KeyError('invalid depth {} for resnet'.format(depth))
        self.depth = depth
        self.num_stages = num_stages
        assert num_stages >= 1 and num_stages <= 4
        self.strides = strides
        self.dilations = dilations
        assert len(strides) == len(dilations) == num_stages
        self.out_indices = out_indices
        assert max(out_indices) < num_stages
        self.frozen_stages = frozen_stages
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.norm_eval = norm_eval
        self.zero_init_residual = zero_init_residual
        self.block, stage_blocks = self.arch_settings[depth]
        self.stage_blocks = stage_blocks[:num_stages]
        self.inplanes = 64

        self._make_stem_layer()

        self.res_layers = []
        for i, num_blocks in enumerate(self.stage_blocks):
            stride = strides[i]
            dilation = dilations[i]
            planes = 64 * 2**i
            res_layer = make_res_layer(
                self.block,
                self.inplanes,
                planes,
                num_blocks,
                stride=stride,
                dilation=dilation,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg)
            self.inplanes = planes * self.block.expansion
            layer_name = 'layer{}'.format(i + 1)
            self.add_sublayer(layer_name, res_layer)
            self.res_layers.append(layer_name)

        self._freeze_stages()

        self.feat_dim = self.block.expansion * 64 * 2**(
            len(self.stage_blocks) - 1)

    @property
    def norm1(self):
        return getattr(self, self.norm1_name)

    def _make_stem_layer(self):
        self.conv1 = build_conv_layer(
            self.conv_cfg,
            3,
            64,
            7,
            stride=2,
            padding=3,
            bias_attr=False)
        self.norm1_name, norm1 = build_norm_layer(self.norm_cfg, 64, postfix=1)
        self.add_sublayer(self.norm1_name, norm1)
        self.relu = fluid.layers.relu
        self.maxpool = nn.Pool2D(pool_size=3, pool_stride=2, pool_padding=1)

    def _freeze_stages(self):
        if self.frozen_stages >= 0:
            self.norm1.eval()
            for layer in [self.conv1, self.norm1]:
                layer.eval()
                for param in layer.parameters():
                    param.stop_gradient = True

        for i in range(1, self.frozen_stages + 1):
            layer = getattr(self, 'layer{}'.format(i))
            layer.eval()
            for param in layer.parameters():
                param.stop_gradient = True

    def init_weights(self, pretrained=None):
        logger = logging.getLogger()
        if isinstance(pretrained, str):
            logger.info('Loading pretrained model from {}'.format(pretrained))
            self.set_dict(fluid.dygraph.load_dygraph(pretrained)[0])
        elif pretrained is None:
            logger.warning('No pretrained model for Resnet')
        else:
            raise TypeError('pretrained must be a str or None')

    def forward(self, x):
        outs = []
        x = self.conv1(x)
        x = self.norm1(x)
        x = self.relu(x)
        outs.append(x)   # add for encoder
        x = self.maxpool(x)
        for i, layer_name in enumerate(self.res_layers):
            res_layer = getattr(self, layer_name)
            x = res_layer(x)
            if i in self.out_indices:
                outs.append(x)
        return tuple(outs)

    def train(self):
        super(ResNet, self).train()
        self._freeze_stages()
        if self.norm_eval:
            for layer in self.sublayers():
                # trick: eval have effect on BatchNorm only
                if isinstance(layer, nn.BatchNorm):
                    layer.eval()