"""
file structure:
flip_labels_and_scans.py
	 scan_directrory - raw scans folder
	 label_directrory - labels folder
	 save_dir_scan - flipped scans folder (where they will be saved)
	 save_dir_labels - flipped labels folder (where they will be saved)

This script flips nii (nifti) labels and scans along the sagittal plane. The plane flipping occurs on can be modified by changing the transformation matrix in the flip3dlabel and flip3dscan functions.

This script assumes the following file naming conventions:
scans: "scanIdentifier_somesuffix.nii"
labels: "scanIdentifier_50um_segmentation_IE-label.nii", the suffix can be modified by altering "label_name" in the "scan_flip_iterator" function.
note that scanIdentifier should be unique.
"""

#imports
import numpy as np
import SimpleITK as sitk
import os

#specify directory of scans you would like to flip
scan_directrory = 'scan_154um'

#specify directory of where labelmaps are
label_directrory = 'lab'

#specify directory where you want to save flipped scans
save_dir_scan = 'scan_save'

#specify directory where you want to save flipped labels
save_dir_labels = 'lab_save'

def get_center(img):
	"""
	This function returns the physical center point of a 3d sitk image
	:param img: The sitk image we are trying to find the center of
	:return: The physical center point of the image
	"""

	width, height, depth = img.GetSize()
	return img.TransformIndexToPhysicalPoint((int(np.ceil(width/2)),
		                          int(np.ceil(height/2)),
		                          int(np.ceil(depth/2))))

def flip3dlabel(img):
	"""
	This function flips the sitk label passeed to it with NN interpolation
	:param img: An sitk labelmap
	:return: The flipped label
	"""

	affineTrans = sitk.AffineTransform(3)
	image_center = get_center(img)
	affineTrans.SetMatrix([-1,0,0,0,1,0,0,0,1])
	affineTrans.SetCenter(image_center)
	flipped = sitk.Resample(img, affineTrans,sitk.sitkNearestNeighbor)
	return flipped

def flip3dscan(img,lab):
	"""
	This function flips the sitk image passeed to it with BSpline interpolation
	:param img: An sitk image
	:param lab: An sitk label associated with the given image - used to maintain alignment
	:return: The flipped image
	"""

	affineTrans = sitk.AffineTransform(3)
	image_center = get_center(lab)
	affineTrans.SetMatrix([-1,0,0,0,1,0,0,0,1])
	affineTrans.SetCenter(image_center)
	interpolator = sitk.sitkBSpline
	flipped = sitk.Resample(img, img, affineTrans,
		     interpolator, -2000)
	return flipped

def label_flip_iterator(file):
	"""
	This function is called each time a label is flipped. Naming and saving is done here.
	:param file: filename of label
	"""

	prefix = file.split("_")[0] #get the sample prefix IE '1932L'
	name_without_filetype = file.split(".nii")[0] #file name before the extension (.nii)
	newname = name_without_filetype+"_flipped.nii"

	lab = sitk.ReadImage(label_directrory+'/'+file)
	flipped_lab = flip3dlabel(lab)
	sitk.WriteImage(flipped_lab,save_dir_labels+"/"+newname)#labels are saved with _flipped appended to their original names

def scan_flip_iterator(file):
	"""
	This function is called each time a scan is flipped. Naming and saving is done here.
	:param file: filename of scan
	"""

	prefix = file.split("_")[0] #get the scan prefix IE '1932L'
	name_without_filetype = file.split(".nii")[0] #everything before the extension (.nii)
	newname = name_without_filetype+"_flipped.nii"
	label_name = prefix+"_50um_segmentation_IE-label_flipped.nii" #labels corresponding to scans need this naming convention following prefix

	im = sitk.ReadImage(scan_directrory+"/"+file)
	lab = sitk.ReadImage(save_dir_labels+'/'+label_name)
	flipped_im = flip3dscan(im,lab) #flip the image with respect to its already flipped label
	sitk.WriteImage(flipped_im,save_dir_scan+"/"+newname) #scans are saved with _flipped appended to their original names

dir=os.listdir(label_directrory)
for i in range(0,len(dir)): #iterate through the directory of labels
	label_flip_iterator(dir[i])

dir=os.listdir(scan_directrory)
for i in range(0,len(dir)): #iterate through the directory of raw scans
	scan_flip_iterator(dir[i])


import numpy as np
from scipy.spatial.distance import cdist
import matplotlib.pyplot as plt

def has_converged(centers, new_centers):
    return set([tuple(a) for a in centers]) == set([tuple(a) for a in new_centers])

def kmeans(X, K):
    # centroids = X[np.random.choice(X.shape[0], K, replace=False)]
    centroids = np.array([[1.0, 1.0], [5.0, 7.0]])
    it = 0
    while True:
        it += 1
        D = cdist(X, centroids)
        labels = np.argmin(D, axis=1)
        new_centroids = np.zeros((K, X.shape[1]))
        for k in range(K):
            new_centroids[k, :] = np.mean(X[labels == k, :], axis=0)
        display(X, K, labels)
        plt.show()
        if has_converged(centroids, new_centroids):
            break

        centroids = new_centroids

    return labels, centroids

def display(X, K, labels):
    for i in range(K):
        X0 = X[labels == i, :]
        plt.plot(X0[:, 0], X0[:, 1], '.')

def error(X, K, labels):
    sum = 0
    for i in range(K):
        X0 = X[labels == i, :]
        sum += np.std(X0)
    print(sum / K)

def random_data():
    for i in range(6):
        mean = 200 * np.random.random_sample((1, 2))
        X0 = np.random.multivariate_normal(mean[0], [[10, 0], [0, 10]], np.random.randint(20, 50))
        if i == 0:
            X = X0
        else:
            X = np.concatenate((X, X0))
    return X

from sklearn.cluster import KMeans

A = np.array([[1.0, 1.5, 3.0, 5.0, 3.5, 4.5, 3.5]])
B = np.array([[1.0, 2.0, 4.0, 7.0, 5.0, 5.0, 4.5]])
X = np.append(A.T, B.T, axis=1)
# X = random_data()
for K in range(2, 10):
    (labels, centroids) = kmeans(X, K)
    display(X, K, labels)
    plt.show()
    error(X, K, labels)
    cls = KMeans(n_clusters=K, random_state=0)
    cls.fit(X)
    lbl = cls.labels_
    display(X, K, lbl)
    plt.show()


import numpy as np
import tensorflow as tf
import cv2

def colormap_jet(img):
    color_image = cv2.applyColorMap(np.uint8(img), cv2.COLORMAP_JET)
    return color_image

def color_disparity(disparity):
    with tf.variable_scope('color_disparity'):
        batch_size = disparity.shape[0]
        color_maps = []
        for i in range(batch_size):
            color_disp = tf.py_func(colormap_jet, [-disparity[i]], tf.uint8)
            color_maps.append(color_disp)
        color_batch = tf.stack(color_maps, axis=0)
        return color_batch

def count_text_lines(file_path):
    f = open(file_path, 'r')
    lines = f.readlines()
    f.close()
    return len(lines)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
===================================
Timer --- Create a timer decorator.
===================================

Largely this module was simply practice on writing decorators.

Might need to review logging best practices. I don't want the logger from
this module to emit anything, but it seems tedious to place that burden
on any module that imports from here.

.. seealso::

    :mod:`cProfile`
    :mod:`pstats`
    :mod:`timeit`
    :magic:`timeit`

"""
import datetime
import functools
import logging
from os import scandir
from runpy import run_path
import time
from timeit import Timer

from IPython.core.getipython import get_ipython

# noinspection PyProtectedMember
from IPython.core.magics.execution import _format_time as format_delta

logging.basicConfig(level=logging.INFO)

def timer(func):
    """Print the runtime of the decorated function.

    Utilizes `time.perf_counter`.

    .. todo:: Begin using the :mod:`timeit` module.

        There are more specialized ways of profiling things in
        other modules; however, this works for a rough estimate.

    Parameters
    ----------
    func : function
        Function to profile

    Returns
    -------
    value : float
        Output of function :func:`time.perf_counter()`.

    """

    @functools.wraps(func)
    def wrapper_timer(*args, **kwargs):
        start_time = time.perf_counter()
        value = func(*args, **kwargs)
        end_time = time.perf_counter()
        run_time = end_time - start_time
        logging.info(f"Finished {func.__name__!r} in {run_time:.4f} secs")
        return value

    return wrapper_timer

# class ModuleTimer()
# I mean while we're practicing decorators throw this in the mix
def debug(func):
    """Print the function signature and return value"""

    @functools.wraps(func)
    def wrapper_debug(*args, **kwargs):
        args_repr = [repr(a) for a in args]  # 1
        kwargs_repr = [f"{k}={v!r}" for k, v in kwargs.items()]  # 2
        signature = ", ".join(args_repr + kwargs_repr)  # 3
        print(f"Calling {func.__name__}({signature})")
        value = func(*args, **kwargs)
        print(f"{func.__name__!r} returned {value!r}")  # 4
        return value

    return wrapper_debug

def exc_timer(statement, setup=None):
    """A non-decorator implementation that uses `timeit`."""
    t = Timer(stmt=statement, setup=setup)  # outside the try/except
    try:
        return t.timeit()
    except Exception:  # noqa E722
        t.print_exc()

class ArgReparser:
    """Class decorator that echoes out the arguments a function was called with."""

    def __init__(self, func):
        """Initialize the reparser with the function it wraps."""
        self.func = func

    def __call__(self, *args, **kwargs):
        print("entering function " + self.func.__name__)
        i = 0
        for arg in args:
            print("arg {0}: {1}".format(i, arg))
            i = i + 1

        return self.func(*args, **kwargs)

def time_dir(directory=None):
    """How long does it take to exec(compile(file)) every file in the startup dir?"""
    if directory is None:
        directory = get_ipython().startup_dir
    result = []
    for i in scandir("."):
        if i.name.endswith(".py"):
            file = i.name
            print(file)
            print(time.time())
            start_time = time.time()
            exec(compile(open(file).read(), "timer", "exec"))
            end = time.time()
            diff = end - start_time
            print(f"{diff}")
            result.append((file, diff))

    return result

class LineWatcher:
    """Class that implements a basic timer.

    Registers the `start` and `stop` methods with the IPython events API.
    """

    def __init__(self):
        """Define the classes start_time parameter."""
        self.start_time = self.start()

    def start(self):
        """Return `time.time`."""
        return time.time()

    def __repr__(self):
        return f"{self.__class__.__name__} {self.start_time}"

    def stop(self):
        """Determine the difference between start time and end time."""
        stop_time = time.time()

        diff = abs(stop_time - self.start_time)
        print("time: {}".format(format_delta(diff)))
        return diff

def load_ipython_extension(ip=None, line_watcher=None):
    """Initialize a `LineWatcher` and register start and stop with IPython."""
    if ip is None:
        ip = get_ipython()
    if ip is None:
        return

    if line_watcher is None:
        line_watcher = LineWatcher()
    ip.events.register("pre_run_cell", line_watcher.start)
    ip.events.register("post_run_cell", line_watcher.stop)

def unload_ipython_extension(ip=None, line_watcher=None):
    if ip is None:
        ip = get_ipython()
    if ip is None:
        return

    if line_watcher is None:
        line_watcher = LineWatcher()
    ip.events.unregister("pre_run_cell", line_watcher.start)
    ip.events.unregister("post_run_cell", line_watcher.stop)


from PIL import Image

imgx = 512
imgy = 512

image = Image.new("RGB",(imgx,imgy))

for x in range(imgx):
	for y in range(imgy):
		if ((x//64)%2 == 1) or ((x//64)%2 == 2) and (y//64)%2 == 1 or ((y//64)%2 == 2):
			image.putpixel ((x,y), (0,0,0) )
		else:
			image.putpixel ((x,y), (250,0,0) )

image.save("demo_image.png", "PNG")


import requests
import time
import json
from hoshino import aiorequests

apiroot = 'https://help.tencentbot.top'

async def getprofile(viewer_id: int, interval: int = 1, full: bool = False) -> dict:
    reqid = json.loads(await aiorequests.get(f'{apiroot}/enqueue?full={full}&target_viewer_id={viewer_id}').content.decode('utf8'))['reqeust_id']

    if reqid is None:
        return "id err"

    while True:
        query = json.loads(await aiorequests.get(f'{apiroot}/query?request_id={reqid}').content.decode('utf8'))
        status = query['status']
        if status == 'done':
            return query['data']
        elif status == 'queue':
            time.sleep(interval)
        else: # notfound or else
            return "queue"

async def queryarena(defs: list, page: int) -> dict:
    return json.loads(await aiorequests.get(f'{apiroot}/arena?def={",".join([str(x) for x in defs])}&page={page}').content.decode('utf8'))


import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression

df = pd.read_csv('dividedsamples/training.csv')
dfval = pd.read_csv('dividedsamples/testing.csv')

train_features = df.copy()
test_features = dfval.copy()

train_labels = train_features.pop('price')
test_labels = test_features.pop('price')

regressor = LinearRegression()
regressor.fit(train_features, train_labels)
coeff_df = pd.DataFrame(regressor.coef_, train_features.columns, columns=['Coefficient'])
print(coeff_df)

y_pred = regressor.predict(test_features)
boi = pd.DataFrame({'Actual': test_labels, 'Predicted': y_pred})
print(boi)


from django import forms

from django.forms import TextInput, SplitDateTimeWidget

class NumberInput(TextInput):
    """
    HTML5 Number input
    Left for backwards compatibility
    """
    input_type = 'number'

class AdminDateWidget(forms.DateInput):

    @property
    def media(self):
        js = ["calendar.js", "admin/DateTimeShortcuts.js"]
        return forms.Media(js=["admin/js/%s" % path for path in js])

    def __init__(self, attrs=None, format=None):
        final_attrs = {'class': 'form-control datepicker',
                       'size': '10', 'type': 'date'}

        if attrs is not None:
            final_attrs.update(attrs)
        super(AdminDateWidget, self).__init__(attrs=final_attrs, format=format)

class AdminTimeWidget(forms.TimeInput):

    @property
    def media(self):
        js = ["calendar.js", "admin/DateTimeShortcuts.js"]
        return forms.Media(js=["admin/js/%s" % path for path in js])

    def __init__(self, attrs=None, format=None):
        final_attrs = {'class': 'form-control timepicker',
                       'size': '8', 'type': 'time'}

        if attrs is not None:
            final_attrs.update(attrs)
        super(AdminTimeWidget, self).__init__(attrs=final_attrs, format=format)

class LteAdminSplitDateTime (forms.SplitDateTimeWidget):

    #template_name = 'admin/widgets/split_datetime.html'

    def __init__(self, attrs=None):
        widgets = [AdminDateWidget, AdminTimeWidget]
        # Note that we're calling MultiWidget, not SplitDateTimeWidget, because
        # we want to define widgets.
        forms.MultiWidget.__init__(self, widgets, attrs)


__author__ = 'yueli'
import numpy as np
import matplotlib.pyplot as plt
from config.config import *

mrList = np.linspace(1, 13, 13)
negativeList = [-1, -1, -10, -10, -1, -1, -1, -10 ,-1, -1, -1, -10, -10]
noMapReplyList = np.linspace(0, 0, 13)
rlocSet1 = [-10, 1, 1, 1, -10, -10, 1, 1, 1, -10, 1, 1, 1]
rlocSet2 = [-10, 2, 2, 2, 2, 2, 2, 2, 2, -10, 2, 2, 2]
rlocSet3 = [-10, -10, 3, 3, -10, -10, -10, 3, -10, -10, -10, 3, 3]

plt.xlim(0.5, 13.5)
plt.ylim(-1.5, 3.5)

plt.scatter(mrList,negativeList, color = 'blue')
plt.scatter(mrList,noMapReplyList, color = 'yellow')
plt.scatter(mrList,rlocSet1, color = 'purple')
plt.scatter(mrList,rlocSet2, color = 'red')
plt.scatter(mrList,rlocSet3, color = 'green')

plt.xlabel("13 different Map Resolvers")
plt.ylabel("Responses from MRs")
plt.title("Responses from 13 MRs for EID-153.16.49.112 at liege(by MR)")
plt.xticks(mrList, ['MR1', 'MR2', 'MR3', 'MR4', 'MR5', 'MR6', 'MR7', 'MR8', 'MR9', 'MR10', 'MR11', 'MR12', 'MR13', 'MR14', 'MR15'])
plt.yticks([-1, 0, 1, 2, 3], ['Negative Reply', 'No Map Reply', '82.121.231.67', '192.168.1.66', '132.227.85.231'])

# plt.savefig(
#     os.path.join(PLOT_DIR, 'Plot_variable_MR', 'Plot_variable_MR.eps'),
#     dpi=300,
#     transparent=True
# )
plt.show()


#!/usr/bin/env python3.8

def missing_element(arr1, arr2):
	arr1.sort()
	arr2.sort()

	for num1, num2 in zip(arr1, arr2):
		if num1 != num2:
			return num1

	return -1

def missing_element1(arr1, arr2):
	count = {}
	output = []
	for i in arr1:
		if i in count:
			count[i] += 1
		else:
			count[i] = 1

	for i in arr2:
		if i in count:
			count[i] -= 1
		else:
			count[i] = -1

	for k in count:
		if count[k] > 0:
			output.append(k)

	return output

import collections
def missing_element2(arr1, arr2):
	count = collections.defaultdict(int)
	output = []
	for i in arr2:
			count[i] += 1

	for i in arr1:
		if count[i] == 0:
			output.append(i)
		else:
			count[i] -= 1

	return output

def missing_element3(arr1, arr2):
	return sum(arr1) - sum(arr2)

def missing_element4(arr1, arr2):
	result = 0

	for num in arr1+arr2:
		result ^= num

	return result

arr1 = [5,5,7,7]
arr2 = [5,7,7]

print(missing_element4(arr1,arr2))

print( ord("A")^ord("A"))


"""Guide Eye 01 module"""

from functools import partial
from mgear.shifter.component import guide
from mgear.core import transform, pyqt
from mgear.vendor.Qt import QtWidgets, QtCore

from maya.app.general.mayaMixin import MayaQWidgetDockableMixin
from maya.app.general.mayaMixin import MayaQDockWidget

from . import settingsUI as sui

# guide info
AUTHOR = "Jeremie Passerin, Miquel Campos"
URL = ", www.miquletd.com"
EMAIL = ", "
VERSION = [1, 0, 0]
TYPE = "eye_01"
NAME = "eye"
DESCRIPTION = "eye control rig"

##########################################################
# CLASS
##########################################################

class Guide(guide.ComponentGuide):
    """Component Guide Class"""

    compType = TYPE
    compName = NAME
    description = DESCRIPTION

    author = AUTHOR
    url = URL
    email = EMAIL
    version = VERSION

    def postInit(self):
        """Initialize the position for the guide"""
        self.save_transform = ["root", "look"]

    def addObjects(self):
        """Add the Guide Root, blade and locators"""

        # eye guide
        self.root = self.addRoot()
        vTemp = transform.getOffsetPosition(self.root, [0, 0, 1])
        self.look = self.addLoc("look", self.root, vTemp)

        centers = [self.root, self.look]
        self.dispcrv = self.addDispCurve("crv", centers)

    def addParameters(self):
        """Add the configurations settings"""

        self.pUpVDir = self.addEnumParam(
            "upVectorDirection", ["X", "Y", "Z"], 1)

        self.pIkRefArray = self.addParam("ikrefarray", "string", "")
        self.pUseIndex = self.addParam("useIndex", "bool", False)

        self.pParentJointIndex = self.addParam(
            "parentJointIndex", "long", -1, None, None)

##########################################################
# Setting Page
##########################################################

class settingsTab(QtWidgets.QDialog, sui.Ui_Form):
    """The Component settings UI"""

    def __init__(self, parent=None):
        super(settingsTab, self).__init__(parent)
        self.setupUi(self)

class componentSettings(MayaQWidgetDockableMixin, guide.componentMainSettings):
    """Create the component setting window"""

    def __init__(self, parent=None):
        self.toolName = TYPE
        # Delete old instances of the componet settings window.
        pyqt.deleteInstances(self, MayaQDockWidget)

        super(self.__class__, self).__init__(parent=parent)
        self.settingsTab = settingsTab()

        self.setup_componentSettingWindow()
        self.create_componentControls()
        self.populate_componentControls()
        self.create_componentLayout()
        self.create_componentConnections()

    def setup_componentSettingWindow(self):
        self.mayaMainWindow = pyqt.maya_main_window()

        self.setObjectName(self.toolName)
        self.setWindowFlags(QtCore.Qt.Window)
        self.setWindowTitle(TYPE)
        self.resize(350, 350)

    def create_componentControls(self):
        return

    def populate_componentControls(self):
        """Populate Controls

        Populate the controls values from the custom attributes of the
        component.

        """
        # populate tab
        self.tabs.insertTab(1, self.settingsTab, "Component Settings")

        # populate component settings
        self.settingsTab.upVectorDirection_comboBox.setCurrentIndex(
            self.root.attr("upVectorDirection").get())

        ikRefArrayItems = self.root.attr("ikrefarray").get().split(",")
        for item in ikRefArrayItems:
            self.settingsTab.ikRefArray_listWidget.addItem(item)

    def create_componentLayout(self):

        self.settings_layout = QtWidgets.QVBoxLayout()
        self.settings_layout.addWidget(self.tabs)
        self.settings_layout.addWidget(self.close_button)

        self.setLayout(self.settings_layout)

    def create_componentConnections(self):

        cBox = self.settingsTab.upVectorDirection_comboBox
        cBox.currentIndexChanged.connect(
            partial(self.updateComboBox,
                    self.settingsTab.upVectorDirection_comboBox,
                    "upVectorDirection"))

        self.settingsTab.ikRefArrayAdd_pushButton.clicked.connect(
            partial(self.addItem2listWidget,
                    self.settingsTab.ikRefArray_listWidget,
                    "ikrefarray"))

        self.settingsTab.ikRefArrayRemove_pushButton.clicked.connect(
            partial(self.removeSelectedFromListWidget,
                    self.settingsTab.ikRefArray_listWidget,
                    "ikrefarray"))

        self.settingsTab.ikRefArray_listWidget.installEventFilter(self)

    def eventFilter(self, sender, event):
        if event.type() == QtCore.QEvent.ChildRemoved:
            if sender == self.settingsTab.ikRefArray_listWidget:
                self.updateListAttr(sender, "ikrefarray")
            return True
        else:
            return QtWidgets.QDialog.eventFilter(self, sender, event)

    def dockCloseEventTriggered(self):
        pyqt.deleteInstances(self, MayaQDockWidget)


import logging
import sys

from loguru import logger
from starlette.config import Config
from starlette.datastructures import Secret

from app.core.logger import InterceptHandler

config = Config(".env")

API_PREFIX = "/api"
VERSION = "0.1.0"
DEBUG: bool = config("DEBUG", cast=bool, default=False)
MAX_CONNECTIONS_COUNT: int = config("MAX_CONNECTIONS_COUNT", cast=int, default=10)
MIN_CONNECTIONS_COUNT: int = config("MIN_CONNECTIONS_COUNT", cast=int, default=10)
HOST: str = config("HOST", cast=str, default="0.0.0.0")
PORT: int = config("PORT", cast=int, default=35100)
SECRET_KEY: Secret = config("SECRET_KEY", cast=Secret, default="")

PROJECT_NAME: str = config("PROJECT_NAME", default="augmentation")

# logging configuration
LOGGING_LEVEL = logging.DEBUG if DEBUG else logging.INFO
logging.basicConfig(
    handlers=[InterceptHandler(level=LOGGING_LEVEL)], level=LOGGING_LEVEL
)
logger.configure(handlers=[{"sink": sys.stderr, "level": LOGGING_LEVEL}])

FASTTEXT_PATH = config("FASTTEXT_PATH", default="./model/cc.vi.300.vec")
PHOBERT_PATH = config("PHOBERT_PATH", default="./model/PhoBERT_base_fairseq")
STOPWORD_PATH = config("STOPWORD_PATH", default="./data/vietnamese-stopwords.txt")
IRRELEVANT_WORD_PATH = config("IRRELEVANT_WORD_PATH", default="./data/irrelevant_words.txt")
EDIT_DISTANCE_PATH = config("EDIT_DISTANCE_PATH", default="./data/edit_distance.txt")
MAX_CACHE_SIZE = config("MAX_CACHE_SIZE", cast=int, default=1000)
PHO_NLP_URL = config("PHO_NLP_URL", default="http://172.29.13.23:20217/")
VN_CORE_PATH = config("VN_CORE_PATH", default="http://172.29.13.23")
VN_CORE_PORT = config("VN_CORE_PORT", cast=int, default=20215)


from typing import List, Optional

from fastapi import APIRouter, Header
from fastapi.exceptions import HTTPException

from server.models.subscription import (
    ExchangeKlineSubscriptionRequest,
    ExchangeSubscription,
    ExchangeSubscriptionType,
)

router = APIRouter()

@router.get("/")
async def list(x_connection_id: str = Header()) -> List[ExchangeSubscription]:
    return await ExchangeSubscription.find(
        ExchangeSubscription.type == ExchangeSubscriptionType.KLINE,
        ExchangeSubscription.connection == x_connection_id,
    ).to_list()

@router.get("/{symbol:path}/interval/{interval:path}/")
async def retrieve(
    symbol: str,
    interval: str,
    x_connection_id: str = Header(),
) -> Optional[ExchangeSubscription]:
    try:
        return await ExchangeSubscription.find_one(
            ExchangeSubscription.type == ExchangeSubscriptionType.KLINE,
            ExchangeSubscription.symbol == symbol,
            ExchangeSubscription.interval == interval,
            ExchangeSubscription.connection == x_connection_id,
        )

    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@router.get("/{symbol:path}/")
async def list_symbol(
    symbol: str,
    x_connection_id: str = Header(),
) -> List[ExchangeSubscription]:
    try:
        return await ExchangeSubscription.find(
            ExchangeSubscription.type == ExchangeSubscriptionType.KLINE,
            ExchangeSubscription.symbol == symbol,
            ExchangeSubscription.connection == x_connection_id,
        ).to_list()

    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@router.post("/")
async def create(
    request: ExchangeKlineSubscriptionRequest, x_connection_id: str = Header()
) -> ExchangeSubscription:
    try:
        existing_subscription = await ExchangeSubscription.find_one(
            ExchangeSubscription.type == ExchangeSubscriptionType.KLINE,
            ExchangeSubscription.symbol == request.symbol,
            ExchangeSubscription.interval == request.interval,
            ExchangeSubscription.connection == x_connection_id,
        )

        if existing_subscription:
            return existing_subscription

        subscription = ExchangeSubscription(
            type=ExchangeSubscriptionType.KLINE,
            interval=request.interval,
            connection=x_connection_id,
            symbol=request.symbol,
        )

        return await subscription.create()

    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@router.delete("/")
async def destroy(
    request: ExchangeKlineSubscriptionRequest, x_connection_id: str = Header()
):
    try:
        item = await ExchangeSubscription.find_one(
            ExchangeSubscription.type == ExchangeSubscriptionType.KLINE,
            ExchangeSubscription.interval == request.interval,
            ExchangeSubscription.connection == x_connection_id,
            ExchangeSubscription.symbol == request.symbol,
        )

        if not item:
            raise HTTPException(status_code=400, detail="subscription not found")

        await ExchangeSubscription.delete(item)

    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


import datetime as dt
import matplotlib.pyplot as plt
import numpy as np
import os

class HobsHeader(object):
    sim_head = '"SIMULATED EQUIVALENT"'
    obs_head = '"OBSERVED VALUE"'
    obs_name = '"OBSERVATION NAME"'
    date = 'DATE'
    dyear = 'DECIMAL_YEAR'

    header = {sim_head: None,
              obs_head: None,
              obs_name: None,
              date: None,
              dyear: None}

class HobsOut(dict):
    """
    Reads output data from Hobs file and prepares it for post processing.

    Class sets observations to an ordered dictionary based on observation name

    If observation name is consistant for a site, a time series is created
    for plotting!

    Parameters
    ----------
    filename : str
        hobs filename
    strip_after : str
        flag to indicate a character to strip the hobs label after for
        grouping wells.

        Example:  OBS_1
                  OBS_2
        strip_after could be set to "_" and then all OBS observations will
        be stored under the OBS key. This is extremely useful for plotting
        and calculating statistics

    """

    def __init__(self, filename, strip_after=""):
        super(HobsOut, self).__init__()
        self.name = filename
        self._strip_after = strip_after
        self._dataframe = None
        self.__read_hobs_output()

    def __read_hobs_output(self):
        """
        Method to read a hobs output file. Dynamically sets header information
        and reads associated values.

        Sets values to HobsOut dictionary
        """

        with open(self.name) as hobout:
            for ix, line in enumerate(hobout):
                if ix == 0:
                    self.__set_header(line)

                else:
                    self.__set_dictionary_values(line)

    def __set_dictionary_values(self, line):
        """
        Method to set incoming hobs line to dictionary data values

        Args:
            line: (str)
        """
        t = line.strip().split()
        obsname = t[HobsHeader.header[HobsHeader.obs_name]]
        dict_name = obsname
        if self._strip_after:
            dict_name = obsname.split(self._strip_after)[0]
        simval = float(t[HobsHeader.header[HobsHeader.sim_head]])
        obsval = float(t[HobsHeader.header[HobsHeader.obs_head]])
        residual = simval - obsval
        date = self.__set_datetime_object(t[HobsHeader.header[HobsHeader.date]])
        decimal_date = float(t[HobsHeader.header[HobsHeader.dyear]])

        if dict_name in self:
            self[dict_name]['simval'].append(simval)
            self[dict_name]['obsval'].append(obsval)
            self[dict_name]['date'].append(date)
            self[dict_name]['decimal_date'].append(decimal_date)
            self[dict_name]['residual'].append(residual)
            self[dict_name]["obsname"].append(obsname)
        else:
            self[dict_name] = {"obsname": [obsname], "date": [date],
                               "decimal_date": [decimal_date],
                               "simval": [simval], "obsval": [obsval],
                               "residual": [residual]}

    def __set_header(self, line):
        """
        Reads header line and sets header index

        Parameters
        ----------
        line : str
            first line of the HOB file

        """
        n = 0
        s = ""

        for i in line:
            s += i
            if s in HobsHeader.header:
                HobsHeader.header[s] = n
                n += 1
                s = ""

            elif s in (" ", "\t", "\n"):
                s = ""

            else:
                pass

        for key, value in HobsHeader.header.items():
            if value is None:
                raise AssertionError("HobsHeader headings must be updated")

    def __set_datetime_object(self, s):
        """
        Reformats a string of YYYY-mm-dd to a datetime object

        Parameters
        ----------
        s : str
            string of YYYY-mm-dd

        Returns
        -------
            datetime.date
        """
        return dt.datetime.strptime(s, "%Y-%m-%d")

    def __get_date_string(self, date):
        """

        Parmaeters
        ----------
            date: datetime.datetime object

        Returns
        -------
            string

        """
        return date.strftime("%Y/%m/%d")

    @property
    def obsnames(self):
        """
        Return a list of obsnames from the HobsOut dictionary
        """
        return self.keys()

    def to_dataframe(self):
        """
        Method to get a pandas dataframe object of the
        HOBs data.

        Returns
        -------
            pd.DataFrame
        """
        import pandas as pd

        if self._dataframe is None:
            df = None
            for hobsname, d in self.items():
                t = pd.DataFrame(d)
                if df is None:
                    df = t
                else:
                    df = pd.concat([df, t], ignore_index=True)
            self._dataframe = df

        return self._dataframe

    def get_sum_squared_errors(self, obsname):
        """
        Returns the sum of squared errors from the residual

        Parameters
        ----------
        obsname : str
            observation name

        Returns
        -------
            float: sum of square error
        """
        return sum([i**2 for i in self[obsname]['residual']])

    def get_rmse(self, obsname):
        """
        Returns the RMSE from the residual

        Parameters
        ----------
        obsname : str
            observation name

        Returns
        -------
            float: rmse
        """
        return np.sqrt(np.mean([i**2 for i in self[obsname]['residual']]))

    def get_number_observations(self, obsname):
        """
        Returns the number of observations for an obsname

        Parameters
        ----------
        obsname : str
            observation name

        Returns
        -------
            int
        """
        return len(self[obsname]['simval'])

    def get_maximum_residual(self, obsname):
        """
        Returns the datetime.date and maximum residual value

        Parameters
        ----------
        obsname : str
            observation name

        Returns
        -------
            tuple: (datetime.date, residual)
        """
        data = self[obsname]['residual']
        index = data.index(max(data))
        date = self[obsname]['date'][index]
        return date, max(data)

    def get_minimum_residual(self, obsname):
        """
        Returns the datetime.date, minimum residual value

        Parameters
        ----------
        obsname : str
            observation name

        Returns
        -------
            tuple: (datetime.date, residual)
        """
        data = self[obsname]['residual']
        index = data.index(min(data))
        date = self[obsname]['date'][index]
        return date, min(data)

    def get_mean_residual(self, obsname):
        """
        Returns the datetime.date, minimum residual value

        Parameters
        ----------
        obsname : str
            observation name

        Returns
        -------
            tuple: (datetime.date, residual)
        """
        data = self[obsname]['residual']
        return np.mean(data)

    def get_median_residual(self, obsname):
        """
        Returns the datetime.date, minimum residual value

        Parameters
        ----------
        obsname : str
            observation name

        Returns
        -------
            tuple: (datetime.date, residual)
        """
        data = self[obsname]['residual']
        return np.median(data)

    def get_maximum_residual_heads(self, obsname):
        """
        Returns the datetime.date, simulated, and observed
        heads at the maximum residual value

        Parameters
        ----------
        obsname : str
            observation name

        Returns
        -------
            tuple: (datetime.date, simulated head, observed head)
        """
        resid = self[obsname]['residual']
        index = resid.index(max(resid))
        observed = self[obsname]['obsval'][index]
        simulated = self[obsname]['simval'][index]
        date = self[obsname]['date'][index]
        return date, simulated, observed

    def get_minimum_residual_heads(self, obsname):
        """
        Returns the datetime.date, simulated, and observed
        heads at the maximum residual value

        Parameters
        ----------
        obsname : str
            observation name

        Returns
        -------
            tuple: (datetime.date, simulated head, observed head)
        """
        resid = self[obsname]['residual']
        index = resid.index(min(resid))
        observed = self[obsname]['obsval'][index]
        simulated = self[obsname]['simval'][index]
        date = self[obsname]['date'][index]
        return date, simulated, observed

    def get_residual_bias(self, filter=None):
        """
        Method to determine the bias of measurements +-
        by checking the residual. Returns fraction of residuals
        > 0.

        Parameters
        ----------
        filter: (str, list, tuple, or function)
            filtering criteria for writing statistics.
            Function must return True for filter out, false to use

        Returns
        -------
            (float) fraction of residuals greater than zero
        """
        nobs = 0.
        ngreaterzero = 0.

        for obsname, meta_data in self.items():
            if self.__filter(obsname, filter):
                continue

            residual = np.array(meta_data['residual'])
            rgreaterzero = sum((residual > 0))

            nobs += residual.size
            ngreaterzero += rgreaterzero

        try:
            bias = ngreaterzero / nobs
        except ZeroDivisionError:
            raise ZeroDivisionError("No observations found!")

        return bias

    def write_dbf(self, dbfname, filter=None):
        """
        Method to write a dbf file from a the HOBS dictionary

        Parameters
        ----------
        dbfname : str
            dbf file name
        filter: (str, list, tuple, or function)
            filtering criteria for writing statistics.
            Function must return True for filter out, false for write to file

        """
        import shapefile
        data = []
        for obsname, meta_data in self.items():

            if self.__filter(obsname, filter):
                continue

            for ix, val in enumerate(meta_data['simval']):
                data.append([obsname,
                             self.__get_date_string(meta_data['date'][ix]),
                             val,
                             meta_data['obsval'][ix],
                             meta_data['residual'][ix]])

        try:
            # traps for pyshp 1 vs. pyshp 2
            w = shapefile.Writer(dbf=dbfname)
        except Exception:
            w = shapefile.Writer()

        w.field("HOBSNAME", fieldType="C")
        w.field("HobsDate", fieldType="D")
        w.field("HeadSim", fieldType='N', decimal=8)
        w.field("HeadObs", fieldType="N", decimal=8)
        w.field("Residual", fieldType="N", decimal=8)

        for rec in data:
            w.record(*rec)

        try:
            w.save(dbf=dbfname)
        except AttributeError:
            w.close()

    def write_min_max_residual_dbf(self, dbfname, filter=None):
        """
        Method to write a dbf of transient observations
        using observation statistics

        Parameters
        ----------
        dbfname : str
            dbf file name
        filter: (str, list, tuple, or function)
            filtering criteria for writing statistics.
            Function must return True for filter out, false for write to file

        """
        import shapefile
        data = []
        for obsname, meta_data in self.items():
            if self.__filter(obsname, filter):
                continue

            max_date, resid_max = self.get_maximum_residual(obsname)
            min_date, resid_min = self.get_minimum_residual(obsname)
            simval_max, obsval_max = self.get_maximum_residual_heads(obsname)[1:]
            simval_min, obsval_min = self.get_minimum_residual_heads(obsname)[1:]

            data.append([obsname,
                         self.get_number_observations(obsname),
                         self.__get_date_string(max_date), resid_max,
                         self.__get_date_string(min_date), resid_min,
                         simval_max, obsval_max, simval_min, obsval_min])

        try:
            # traps for pyshp 1 vs. pyshp 2
            w = shapefile.Writer(dbf=dbfname)
        except Exception:
            w = shapefile.Writer()

        w.field("HOBSNAME", fieldType="C")
        w.field("FREQUENCY", fieldType="N")
        w.field("MaxDate", fieldType="C")
        w.field("MaxResid", fieldType='N', decimal=8)
        w.field("MinDate", fieldType="C", decimal=8)
        w.field("MinResid", fieldType="N", decimal=8)
        w.field("MaxHeadSim", fieldType="N", decimal=8)
        w.field("MaxHeadObs", fieldType="N", decimal=8)
        w.field("MinHeadSim", fieldType="N", decimal=8)
        w.field("MinHeadObs", fieldType="N", decimal=8)

        for rec in data:
            w.record(*rec)

        try:
            w.save(dbf=dbfname)
        except AttributeError:
            w.close()

    def __filter(self, obsname, filter):
        """
        Boolean filetering method, checks if observation name
        is in the filter.

        Parameters
        ----------
        obsname : str
            observation name
        filter: (str, list, tuple, or function)
            filtering criteria for writing statistics.
            Function must return True for filter out, false for write to file

        Returns
        -------
            bool: True if obsname in filter

        """
        if filter is None:
            return False

        elif isinstance(filter, list) or isinstance(filter, tuple):
            if obsname in list:
                return True

        elif isinstance(filter, str):
            if obsname == filter:
                return True

        elif callable(filter):
            if filter(obsname):
                return True

        else:
            raise Exception("Filter is not an appropriate type")

        return False

    def write_summary_statistics_csv(self, csvname, filter=None):
        """
        Method to write summary calibration statistics to a
        CSV file for analysis and reports

        Parameters
        ----------
        csvname : str
            csv file name
        filter: (str, list, tuple, or function)
            filtering criteria for writing statistics.
            Function must return True for filter out, false for write to file

        """
        data = []
        header = ["Well name", "Average", "Median",
                  "Minimum", "Maximum", "RMSE ft", "Frequency"]

        for obsname, meta_data in sorted(self.items()):

            if self.__filter(obsname, filter):
                continue

            resid_mean = self.get_mean_residual(obsname)
            resid_median = self.get_median_residual(obsname)
            resid_max = self.get_maximum_residual(obsname)[-1]
            resid_min = self.get_minimum_residual(obsname)[-1]
            rmse = self.get_rmse(obsname)
            frequency = self.get_number_observations(obsname)
            data.append((obsname, resid_mean, resid_median,
                         resid_min, resid_max, rmse, frequency))

        data = np.array(data, dtype=[('id', 'O'), ('mean', float),
                                     ('med', float), ('min', float),
                                     ('max', float), ('rmse', float),
                                     ('num', np.int)])

        with open(csvname, "w") as foo:
            foo.write(",".join(header) + "\n")
            np.savetxt(foo, data, fmt="%15s,%.2f,%2f,%2f,%2f,%2f,%d")

    def plot(self, obsname, *args, **kwargs):
        """
        Plotting functionality from the hobs dictionary

        Parameters
        ----------
        obsname: str
            hobs package observation name
        *args: matplotlib args
        **kwargs: matplotlib kwargs

        Returns
        -------
            matplotlib.pyplot.axes object

        """
        simulated = True
        if "observed" in kwargs:
            simulated = False
            kwargs.pop('observed')

        observed = True
        if "simulated" in kwargs:
            observed = False
            kwargs.pop('simulated')

        if obsname not in self:
            raise AssertionError("Obsname {}: not valid".format(obsname))

        axes = False
        if 'ax' in kwargs:
            ax = kwargs.pop('ax')
            axes = True

        if not axes:
            ax = plt.subplot(111)

        obsval = self[obsname]['obsval']
        simval = self[obsname]['simval']
        date = self[obsname]['date']

        if observed:
            kwargs['label'] = "Observed"
            kwargs['color'] = 'r'

            ax.plot(date, obsval, *args, **kwargs)

        if simulated:
            kwargs['label'] = "Simulated"
            kwargs['color'] = 'b'

            ax.plot(date, simval, *args, **kwargs)

        return ax

    def plot_measured_vs_simulated(self, filter=None, **kwargs):
        """
        Plots measured vs. simulated data along a 1:1 profile.

        Parameters
        ----------
        filter: (str, list, tuple, or function)
            filtering criteria for writing statistics.
            Function must return True for filter out, false for write to file
        **kwargs: matplotlib.pyplot plotting kwargs

        Returns
        -------
            matplotlib.pyplot.axes object

        """
        axes = plt.subplot(111)

        for obsname, meta_data in self.items():

            if self.__filter(obsname, filter):
                continue

            simulated = meta_data['simval']
            observed = meta_data['obsval']

            axes.plot(observed, simulated, 'bo', markeredgecolor='k')

        return axes

    def plot_simulated_vs_residual(self, filter=None,
                                   histogram=False, **kwargs):
        """
        Creates a matplotlib plot of simulated heads vs residual

        Parameters
        ----------
        filter: (str, list, tuple, or function)
            filtering criteria for writing statistics.
            Function must return True for filter out, false for write to file
        histogram: (bool)
            Boolean variable that defines either a scatter plot (False)
            or a histogram (True) of residuals

        **kwargs: matplotlib.pyplot plotting kwargs

        Returns
        -------
            matplotlib.pyplot.axes object

        """
        axes = plt.subplot(111)

        if not histogram:
            for obsname, meta_data in self.items():

                if self.__filter(obsname, filter):
                    continue

                residual = meta_data['residual']
                observed = meta_data['obsval']

                axes.plot(observed, residual, 'bo', markeredgecolor="k")

        else:
            bins = np.arange(-25, 26, 5)

            d = {}
            for ix, abin in enumerate(bins):
                frequency = 0
                for obsname, meta_data in self.items():
                    if self.__filter(obsname, filter):
                        continue

                    for residual in meta_data['residual']:
                        if ix == 0:
                            if residual < abin:
                                frequency += 1

                        elif ix == (len(bins) - 1):
                            if residual > abin:
                                frequency += 1

                        else:
                            if bins[ix - 1] <= residual < abin:
                                frequency += 1

                if ix == 0:
                    name = "Less than {}".format(abin)

                elif ix == (len(bins) - 1):
                    name = "Greater than {}".format(abin)

                else:
                    name = "{} to {}".format(bins[ix - 1] + 1, abin)

                d[ix + 1] = {'name': name,
                             'frequency': frequency}

            tick_num = []
            tick_name = []
            for index, meta_data in sorted(d.items()):
                axes.bar(index, meta_data['frequency'], width=0.8,
                         **kwargs)
                tick_num.append(index)
                tick_name.append(meta_data['name'])

            plt.xticks(tick_num, tick_name, rotation=45, fontsize=10)
            plt.xlim([0.5, len(tick_num) + 1])
            plt.subplots_adjust(left=0.12, bottom=0.22,
                                right=0.90, top=0.90,
                                wspace=0.20, hspace=0.20)
            plt.ylabel("Frequency")

        return axes

if __name__ == "__main__":
    ws = r'C:\Users\jlarsen\Desktop\Lucerne\Lucerne_OWHM\V0_initial_from_MODOPTIM\output'
    hobs_name = "hobs.out"

    tmp = HobsOut(os.path.join(ws, hobs_name))
    tmp.plot("04N01W01R04S", "o-")
    plt.legend(loc=0, numpoints=1)
    plt.show()
    print('break')


import os
import sys

try:
    from dreamberd import interprete
except ModuleNotFoundError:
    sys.exit("Use -m keyword.")

from argparse import ArgumentParser

parser = ArgumentParser(
    prog="DreamBerd Interpreter (Python)",
    description="The perfect programming language.",
)
parser.add_argument("content", help="The file or code to run.")
args = parser.parse_args()

if os.path.exists(args.content):
    with open(args.content, "r", encoding="utf-8") as file:
        content: str = file.read()
else:
    content = args.content

interprete(content)


from __future__ import annotations
import re
from typing import TYPE_CHECKING

from twitchio import User, PartialUser, Chatter, PartialChatter, Channel, Clip
from .errors import BadArgument

if TYPE_CHECKING:
    from .core import Context

__all__ = (
    "convert_Chatter",
    "convert_Clip",
    "convert_Channel",
    "convert_PartialChatter",
    "convert_PartialUser",
    "convert_User",
)

async def convert_Chatter(ctx: Context, arg: str) -> Chatter:
    """
    Converts the argument into a chatter in the chat. If the chatter is not found, BadArgument is raised.
    """
    arg = arg.lstrip("@")
    resp = [x for x in filter(lambda c: c.name == arg, ctx.chatters or tuple())]
    if not resp:
        raise BadArgument(f"The user '{arg}' was not found in {ctx.channel.name}'s chat.")

    return resp[0]

async def convert_PartialChatter(ctx: Context, arg: str) -> PartialChatter:
    """
    Converts the argument into a chatter in the chat. As opposed to Chatter converter, this will return a PartialChatter regardless of the cache state.
    """
    return PartialChatter(ctx._ws, name=arg.lstrip("@"), channel=ctx.channel, message=None)

async def convert_Clip(ctx: Context, arg: str) -> Clip:
    finder = re.search(r"(https://clips.twitch.tv/)?(?P<slug>.*)", arg)
    if not finder:
        raise RuntimeError(
            "regex failed to match"
        )  # this should never ever raise, but its here to make type checkers happy

    slug = finder.group("slug")
    clips = await ctx.bot.fetch_clips([slug])
    if not clips:
        raise BadArgument(f"Clip '{slug}' was not found")

    return clips[0]

async def convert_User(ctx: Context, arg: str) -> User:
    """
    Similar to convert_Chatter, but fetches from the twitch API instead,
    returning a :class:`twitchio.User` instead of a :class:`twitchio.Chatter`.
    To use this, you most have a valid client id and API token or client secret
    """
    arg = arg.lstrip("@")
    user = await ctx.bot.fetch_users(names=[arg])
    if not user:
        raise BadArgument(f"User '{arg}' was not found.")
    return user[0]

async def convert_PartialUser(ctx: Context, arg: str) -> User:
    """
    This is simply a shorthand to :ref:`~convert_User`, as fetching from the api will return a full user model
    """
    return await convert_User(ctx, arg)

async def convert_Channel(ctx: Context, arg: str) -> Channel:
    if arg not in ctx.bot._connection._cache:
        raise BadArgument(f"Not connected to channel '{arg}'")

    return ctx.bot.get_channel(arg)

_mapping = {
    User: convert_User,
    PartialUser: convert_PartialUser,
    Channel: convert_Channel,
    Chatter: convert_Chatter,
    PartialChatter: convert_PartialChatter,
    Clip: convert_Clip,
}


from typing import Any, Callable, Dict

from torchvision import transforms as T

from rikai.types.vision import Image

"""
Adapted from https://github.com/pytorch/pytorch.github.io/blob/site/assets/hub/pytorch_vision_resnet.ipynb
"""  # noqa E501

def pre_processing(options: Dict[str, Any]) -> Callable:
    """
    All pre-trained models expect input images normalized in the same way, i.e.
    mini-batches of 3-channel RGB images of shape (3 x H x W), where H and W
    are expected to be at least 224. The images have to be loaded in to a range
    of [0, 1] and then normalized using mean = [0.485, 0.456, 0.406] and std =
    [0.229, 0.224, 0.225].
    """
    return T.Compose(
        [
            T.Resize(256),
            T.CenterCrop(224),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )

def post_processing(options: Dict[str, Any]) -> Callable:
    def post_process_func(batch):
        results = []
        for result in batch:
            results.append(result.detach().cpu().tolist())
        return results

    return post_process_func

OUTPUT_SCHEMA = "array<float>"


"""Python module for common workflows and library methods.

Authors: Prasad Hegde
"""

import os
import json
import pathlib
import inspect
import random
import string

class Workflows():
    """
    Common Workflows and library methods
    """
    def get_config_data(self, test_method):
        """
        This routine retuns the config data specific to the test case
        :param test_method: Name of the test method
        :return: tuple containing global_config and test_args
        """
        path = pathlib.Path(inspect.getfile(self.__class__)).parent.absolute()
        config_path = os.path.join(path, "config.json")
        with open(config_path) as f_in:
            config_data = json.load(f_in)
            return config_data["global_config"], config_data["test_args"][self.__class__.__name__]\
                                                            [test_method]

    @staticmethod
    def generate_new_email(length=16, suffix=None):
        """
        This routine generates a new email id
        :param length: Length of the email(int)
        :param suffix: domain(str)
        :return: email id (str)
        """
        retval = ''.join(random.choice(string.ascii_lowercase + string.digits) \
                         for i in range(length))
        return retval + suffix if suffix else retval

    @staticmethod
    def verify_response_header(expected_header, actual_header):
        """
        This routine is used to validate expected response header against actual
        :param expected_header: dict
        :param actual_header: dict
        :return: Boolean
        """
        if not any(item in actual_header.items() for item in expected_header.items()):
            return False
        return True

    @staticmethod
    def verify_response_time(expected_response_time, actual_response_time):
        """
        This routine is used to verify response time of api call
        :param actual_response_time: sec
        :return: Boolean
        """
        if actual_response_time <= expected_response_time:
            return True
        return False

    @staticmethod
    def update_user_details(test_args, **kwargs):
        """
        This Routine is used to build user details
        :param test_args: test args of the test method
        :param kwargs: first_name, last_name, dob, image_url, email_id
        :return: user data (dict)
        """
        first_name = kwargs.get('first_name', test_args["updated_user_details"]["first_name"])
        last_name = kwargs.get('last_name', test_args["updated_user_details"]["last_name"])
        dob = kwargs.get('dob', test_args["updated_user_details"]["dob"])
        image_url = kwargs.get('image_url', test_args["updated_user_details"]["image_url"])
        email = kwargs.get('email_id', None)
        user_data = {"first_name": first_name, "last_name": last_name, "date_of_birth": dob,
                     "image_url": image_url}
        if email:
            user_data["email"] = email

        return user_data


# First Solution: 58476KB / 200ms / 674B
def BS(array,start,end):
    while start<=end:
        mid = (start+end)//2
        if array[mid][1] == 1 and array[mid-1][1]==2: return mid
        elif array[mid][1] == 2: start = mid+1
        else: end = mid-1
    return None

def Solution(data):
    data = sorted(data.items(), key=lambda x:(-x[1],x[0]))
    midpoint = BS(data,0,N+M-1)
    if midpoint == None:
      print(0)
    else:
      print(midpoint)
      stdout.write('\n'.join(map(str,dict(data[:midpoint]).keys())))

from sys import stdin,stdout
from collections import Counter

N, M = map(int, stdin.readline().split())
data = Counter([stdin.readline().rstrip() for _ in range(N+M)])
Solution(data)

# ---------------------------------------------------------
# More Advanced Solution: 41884KB / 124ms / 272B

from sys import stdin,stdout

N, M = map(int, stdin.readline().split())
hear = set([stdin.readline().rstrip() for _ in range(N)])
see = set([stdin.readline().rstrip() for _ in range(M)])

common = sorted(list(hear & see))
print(len(common))
stdout.write('\n'.join(common))


import json

def getAdminAccount():

    with open("./Data/admins.json", "r") as file:
        JSON = file.read()

    accounts = json.loads(JSON)

    return accounts

def getAccount():

    with open("./Data/accounts.json", "r") as file:
        JSON = file.read()

    accounts = json.loads(JSON)

    return accounts


import math as ma
import multiprocessing as mp
import time

def first_test(n):
    """ test naïf de primalité
    retourne True si l'entier est premier, et inversement
    n : un entier naturel
    """
    for a in range(2, int(ma.sqrt(n) + 1)):
        if n % a == 0:
            return False
    return True

def pi(x):
    """retourne le nombre de premiers inférieurs à x
    X : un réel
    """
    cpt = 0
    for n in range(1, int(x)):
        if first_test(n):
            cpt += 1
    return cpt

def gen_carmichael(t):
    """retourne tous les nombres de Carmichael inférieurs à x
    t : un réel
    """
    res = []
    for x in range(3, int(t), 2):  # les nombres de Carmichael sont impairs
        valid = False
        for y in range(2, x):
            if ma.gcd(x, y) == 1:
                if pow(y, x-1, x) != 1:
                    valid = False
                    break
            else:
                valid = True
        if valid:
            res.append(x)

    return res

def worker_proc(x):
    valid = False
    for y in range(2, x):
        if ma.gcd(x, y) == 1:
            if pow(y, x - 1, x) != 1:
                return
        else:
            valid = True
    if valid:
        print(x)

def gen_carmichael_mp(t):
    """retourne tous les nombres de Carmichael inférieurs à t
    version multiprocess
    t : un réel
    """
    pool = mp.Pool(processes=mp.cpu_count())
    pool.map(worker_proc, range(3, int(t), 2))

def gen_carmichael_3(k):
    """ genère les nombres de Carmicheal de longueur binaire k à partir de trois diviseurs premiers
    k : un entier
    """
    # t = int(t)
    prime = []
    for n in range(3, 2 ** k, 2):
        if first_test(n):
            prime.append(n)
    res = []
    for i_a, a in enumerate(prime):
        for i_b, b in enumerate(prime[:i_a]):
            ab = a * b
            for c in prime[:i_b]:
                # on  a obtenu 3 premiers, on teste si leur produit est Carmichael
                # worker_proc(a * b * c)
                tst = ab * c - 1
                if tst.bit_length() != k:
                    continue
                if tst % 2 == 0 and tst % (a - 1) == 0 and tst % (b - 1) == 0 and tst % (c - 1) == 0 and a % (b * c) != 0:
                    res.append(tst + 1)
    return sorted(res)

def gen_carmichael_3_all(t):
    """ genère un nombre de Carmicheal inférieur t à partir de trois diviseurs premiers
    version sans contrainte de taille
    """
    t = int(t)
    prime = []
    for n in range(3, t, 2):
        if first_test(n):
            prime.append(n)
    res = []
    for i_a, a in enumerate(prime):
        for i_b, b in enumerate(prime[:i_a]):
            ab = a * b
            for c in prime[:i_b]:
                tst = ab * c - 1
                if tst % 2 == 0 and tst % (a - 1) == 0 and tst % (b - 1) == 0 and tst % (c - 1) == 0 and a % (b * c) != 0:
                    res.append(tst + 1)
    return sorted(res)

def gen_carmichael_2(p):
    """retourne tous les nombre de carmichael de la forme pqr pour un p donné"""
    prime = []
    for n in range(3, 2 * p * (p ** 2 + 1), 2):
        if n == p:
            continue
        if first_test(n):
            prime.append(n)
    res = []
    for i_r, r in enumerate(prime):
        for q in prime[:i_r]:
            tst = p * q * r - 1
            if tst % 2 == 0 and tst % (p - 1) == 0 and tst % (q - 1) == 0 and tst % (r - 1) == 0 and r % (p * q) != 0:
                res.append(tst + 1)
    return sorted(res)

if __name__ == '__main__':
    t = time.time()
    gen_carmichael_mp(10000)
    print("mt : ", str(time.time() - t))
    t = time.time()
    print(gen_carmichael(64000))
    print("naif : ", str(time.time() - t))
    t = time.time()
    print(gen_carmichael_3_all(100))
    print("3 : ", str(time.time() - t))


import os
import luigi
import json
import z5py
import numpy as np

from ..cluster_tasks import WorkflowBase
from ..relabel import RelabelWorkflow
from ..relabel import find_uniques as unique_tasks
from ..node_labels import NodeLabelWorkflow
from ..features import RegionFeaturesWorkflow
from .. import write as write_tasks

from . import size_filter_blocks as size_filter_tasks
from . import background_size_filter as bg_tasks
from . import filling_size_filter as filling_tasks
from . import filter_blocks as filter_tasks
from . import id_filter as id_tasks
from . import orphan_assignments as orphan_tasks
from . import graph_watershed_assignments as gws_tasks
from . import graph_connected_components as cc_tasks

class SizeFilterWorkflow(WorkflowBase):
    input_path = luigi.Parameter()
    input_key = luigi.Parameter()
    output_path = luigi.Parameter()
    output_key = luigi.Parameter()
    size_threshold = luigi.IntParameter()
    hmap_path = luigi.Parameter(default='')
    hmap_key = luigi.Parameter(default='')
    relabel = luigi.BoolParameter(default=True)
    preserve_zeros = luigi.BoolParameter(default=False)

    def _bg_filter(self, dep):
        filter_task = getattr(bg_tasks,
                              self._get_task_name('BackgroundSizeFilter'))
        dep = filter_task(tmp_folder=self.tmp_folder,
                          max_jobs=self.max_jobs,
                          config_dir=self.config_dir,
                          input_path=self.input_path,
                          input_key=self.input_key,
                          output_path=self.output_path,
                          output_key=self.output_key,
                          dependency=dep)
        return dep

    def _ws_filter(self, dep):
        filter_task = getattr(filling_tasks,
                              self._get_task_name('FillingSizeFilter'))
        dep = filter_task(tmp_folder=self.tmp_folder,
                          max_jobs=self.max_jobs,
                          config_dir=self.config_dir,
                          input_path=self.input_path,
                          input_key=self.input_key,
                          output_path=self.output_path,
                          output_key=self.output_key,
                          hmap_path=self.hmap_path,
                          hmap_key=self.hmap_key,
                          preserve_zeros=self.preserve_zeros,
                          dependency=dep)
        return dep

    def requires(self):
        un_task = getattr(unique_tasks,
                          self._get_task_name('FindUniques'))
        dep = un_task(tmp_folder=self.tmp_folder,
                      max_jobs=self.max_jobs,
                      config_dir=self.config_dir,
                      input_path=self.input_path,
                      input_key=self.input_key,
                      return_counts=True,
                      dependency=self.dependency,
                      prefix='size_filter')
        sf_task = getattr(size_filter_tasks,
                          self._get_task_name('SizeFilterBlocks'))
        dep = sf_task(tmp_folder=self.tmp_folder,
                      max_jobs=self.max_jobs,
                      config_dir=self.config_dir,
                      input_path=self.input_path,
                      input_key=self.input_key,
                      size_threshold=self.size_threshold,
                      dependency=dep)

        if self.hmap_path == '' or self.hmap_path is None:
            assert self.hmap_key == '' or self.hmap_key is None
            dep = self._bg_filter(dep)
        else:
            assert self.hmap_key != ''
            dep = self._ws_filter(dep)

        if self.relabel:
            dep = RelabelWorkflow(tmp_folder=self.tmp_folder,
                                  max_jobs=self.max_jobs,
                                  config_dir=self.config_dir,
                                  target=self.target,
                                  input_path=self.output_path,
                                  input_key=self.output_key,
                                  assignment_path=self.output_path,
                                  assignment_key='relabel_size_filter',
                                  prefix='size_filter',
                                  dependency=dep)
        return dep

    @staticmethod
    def get_config():
        configs = super(SizeFilterWorkflow, SizeFilterWorkflow).get_config()
        configs.update({'size_filter_blocks': size_filter_tasks.SizeFilterBlocksLocal.default_task_config(),
                        'background_size_filter': bg_tasks.BackgroundSizeFilterLocal.default_task_config(),
                        'filling_size_filter': filling_tasks.FillingSizeFilterLocal.default_task_config(),
                        **RelabelWorkflow.get_config()})
        return configs

class FilterLabelsWorkflow(WorkflowBase):
    input_path = luigi.Parameter()
    input_key = luigi.Parameter()
    label_path = luigi.Parameter()
    label_key = luigi.Parameter()
    node_label_path = luigi.Parameter()
    node_label_key = luigi.Parameter()
    output_path = luigi.Parameter()
    output_key = luigi.Parameter()
    filter_labels = luigi.ListParameter()

    def requires(self):
        dep = NodeLabelWorkflow(tmp_folder=self.tmp_folder, config_dir=self.config_dir,
                                target=self.target, max_jobs=self.max_jobs,
                                ws_path=self.input_path, ws_key=self.input_key,
                                input_path=self.label_path, input_key=self.label_key,
                                output_path=self.node_label_path,
                                output_key=self.node_label_key,
                                prefix='filter_labels', max_overlap=True,
                                dependency=self.dependency)
        id_task = getattr(id_tasks,
                          self._get_task_name('IdFilter'))
        id_filter_path = os.path.join(self.output_path, 'filtered_ids.json')
        dep = id_task(tmp_folder=self.tmp_folder, config_dir=self.config_dir,
                      dependency=dep, max_jobs=self.max_jobs,
                      node_label_path=self.node_label_path,
                      node_label_key=self.node_label_key,
                      output_path=id_filter_path,
                      filter_labels=self.filter_labels)
        filter_task = getattr(filter_tasks,
                              self._get_task_name('FilterBlocks'))
        dep = filter_task(tmp_folder=self.tmp_folder, config_dir=self.config_dir,
                          dependency=dep, max_jobs=self.max_jobs,
                          input_path=self.input_path, input_key=self.input_key,
                          filter_path=id_filter_path,
                          output_path=self.output_path, output_key=self.output_key)
        return dep

    @staticmethod
    def get_config():
        configs = super(FilterLabelsWorkflow, FilterLabelsWorkflow).get_config()
        configs.update({'id_filter':
                        id_tasks.IdFilterLocal.default_task_config(),
                        'filter_blocks':
                        filter_tasks.FilterBlocksLocal.default_task_config(),
                        **NodeLabelWorkflow.get_config()})
        return configs

class ApplyThreshold(luigi.Task):
    feature_path = luigi.Parameter()
    feature_key = luigi.Parameter()
    out_path = luigi.Parameter()
    threshold = luigi.FloatParameter()
    threshold_mode = luigi.Parameter(default='less')
    dependency = luigi.TaskParameter()

    threshold_modes = ('less', 'greater', 'equal')

    def requires(self):
        return self.dependency

    def run(self):
        f = z5py.File(self.feature_path)
        ds = f[self.feature_key]
        feats = ds[:, 0]

        assert self.threshold_mode in self.threshold_modes
        if self.threshold_mode == 'less':
            filter_ids = feats < self.threshold
        elif self.threshold_mode == 'greater':
            filter_ids = feats > self.threshold
        elif self.threshold_mode == 'equal':
            filter_ids = feats == self.threshold

        filter_ids = np.where(filter_ids)[0].tolist()
        with open(self.out_path, 'w') as f:
            json.dump(filter_ids, f)

    def output(self):
        return luigi.LocalTarget(self.out_path)

class FilterByThresholdWorkflow(WorkflowBase):
    input_path = luigi.Parameter()
    input_key = luigi.Parameter()
    seg_in_path = luigi.Parameter()
    seg_in_key = luigi.Parameter()
    seg_out_path = luigi.Parameter()
    seg_out_key = luigi.Parameter()
    threshold = luigi.FloatParameter()
    relabel = luigi.BoolParameter(default=True)

    def requires(self):
        # calculate the region features
        feat_path = os.path.join(self.tmp_folder, 'reg_feats.n5')
        feat_key = 'feats'
        dep = RegionFeaturesWorkflow(tmp_folder=self.tmp_folder, max_jobs=self.max_jobs,
                                     target=self.target, config_dir=self.config_dir,
                                     input_path=self.input_path, input_key=self.input_key,
                                     labels_path=self.seg_in_path, labels_key=self.seg_in_key,
                                     output_path=feat_path, output_key=feat_key)

        # apply threshold to get the ids to filter out
        id_filter_path = os.path.join(self.tmp_folder, 'filtered_ids.json')
        dep = ApplyThreshold(feature_path=feat_path, feature_key=feat_key,
                             out_path=id_filter_path, threshold=self.threshold,
                             dependency=dep)

        # filter all blocks
        filter_task = getattr(filter_tasks,
                              self._get_task_name('FilterBlocks'))
        dep = filter_task(tmp_folder=self.tmp_folder, config_dir=self.config_dir,
                          dependency=dep, max_jobs=self.max_jobs,
                          input_path=self.seg_in_path, input_key=self.seg_in_key,
                          filter_path=id_filter_path,
                          output_path=self.seg_out_path, output_key=self.seg_out_key)
        if self.relabel:
            dep = RelabelWorkflow(tmp_folder=self.tmp_folder,
                                  max_jobs=self.max_jobs,
                                  config_dir=self.config_dir,
                                  target=self.target,
                                  input_path=self.seg_out_path,
                                  input_key=self.seg_out_key,
                                  assignment_path=self.seg_out_path,
                                  assignment_key='assignments/relabel_filter',
                                  prefix='pp_filter_by_threshold',
                                  dependency=dep)
        return dep

    @staticmethod
    def get_config():
        configs = super(FilterByThresholdWorkflow, FilterByThresholdWorkflow).get_config()
        configs.update({'filter_blocks': filter_tasks.FilterBlocksLocal.default_task_config(),
                        **RegionFeaturesWorkflow.get_config()})
        return configs

class FilterOrphansWorkflow(WorkflowBase):

    graph_path = luigi.Parameter()
    graph_key = luigi.Parameter()

    path = luigi.Parameter()
    segmentation_key = luigi.Parameter()
    assignment_key = luigi.Parameter()

    output_path = luigi.Parameter()
    assignment_out_key = luigi.Parameter()
    output_key = luigi.Parameter(default=None)
    relabel = luigi.BoolParameter(default=False)

    def requires(self):
        assert False, "FIXME not debugged yet"
        dep = self.dependency
        orphan_task = getattr(orphan_tasks,
                              self._get_task_name('OrphanAssignments'))
        dep = orphan_task(tmp_folder=self.tmp_folder, max_jobs=self.max_jobs,
                          config_dir=self.config_dir, dependency=dep,
                          graph_path=self.graph_path, graph_key=self.graph_key,
                          assignment_path=self.path, assignment_key=self.assignment_key,
                          output_path=self.path, output_key=self.assignment_out_key,
                          relabel=self.relabel)
        if self.output_key is not None:
            write_task = getattr(write_tasks,
                                 self._get_task_name('Write'))
            dep = write_task(tmp_folder=self.tmp_folder, max_jobs=self.max_jobs,
                             config_dir=self.config_dir, dependency=dep,
                             input_path=self.path, input_key=self.segmentation_key,
                             output_path=self.path, output_key=self.output_key,
                             assignment_path=self.output_path, assignment_key=self.assignment_out_key,
                             identifier='filter-orphans')
        return dep

    @staticmethod
    def get_config():
        configs = super(FilterOrphansWorkflow, FilterOrphansWorkflow).get_config()
        configs.update({'orphan_assignments': orphan_tasks.OrphanAssignmentsLocal.default_task_config(),
                        'write': write_tasks.WriteLocal.default_task_config()})
        return configs

class ConnectedComponentsWorkflow(WorkflowBase):
    problem_path = luigi.Parameter()
    graph_key = luigi.Parameter()

    assignment_path = luigi.Parameter()
    assignment_key = luigi.Parameter()

    output_path = luigi.Parameter()
    assignment_out_key = luigi.Parameter()

    output_key = luigi.Parameter(default='')
    path = luigi.Parameter(default='')
    fragments_key = luigi.Parameter(default='')

    def requires(self):
        cc_task = getattr(cc_tasks,
                          self._get_task_name('GraphConnectedComponents'))
        dep = cc_task(max_jobs=self.max_jobs, tmp_folder=self.tmp_folder,
                      config_dir=self.config_dir,
                      problem_path=self.problem_path,
                      graph_key=self.graph_key,
                      assignment_path=self.assignment_path,
                      assignment_key=self.assignment_key,
                      output_path=self.output_path,
                      output_key=self.assignment_out_key,
                      dependency=self.dependency)
        if self.output_key != '':
            write_task = getattr(write_tasks,
                                 self._get_task_name('Write'))
            assert self.fragments_key != '' and self.path != ''
            dep = write_task(tmp_folder=self.tmp_folder, max_jobs=self.max_jobs,
                             config_dir=self.config_dir, dependency=dep,
                             input_path=self.path, input_key=self.fragments_key,
                             output_path=self.output_path, output_key=self.output_key,
                             assignment_path=self.output_path,
                             assignment_key=self.assignment_out_key,
                             identifier='graph-connected-components')
        return dep

    @staticmethod
    def get_config():
        configs = super(ConnectedComponentsWorkflow, ConnectedComponentsWorkflow).get_config()
        configs.update({'graph_connected_components': cc_tasks.GraphConnectedComponentsLocal.default_task_config(),
                        'write': write_tasks.WriteLocal.default_task_config()})
        return configs

class SizeFilterAndGraphWatershedWorkflow(WorkflowBase):

    problem_path = luigi.Parameter()
    graph_key = luigi.Parameter()
    features_key = luigi.Parameter()

    #
    path = luigi.Parameter()
    # path to the merged segmentation
    segmentation_key = luigi.Parameter()
    # path to the underlying fragments
    fragments_key = luigi.Parameter(default='')
    # path to the fragment segment assignment
    assignment_key = luigi.Parameter()

    # the size filter threshold
    size_threshold = luigi.IntParameter(default=None)
    target_number = luigi.IntParameter(default=None)
    relabel = luigi.BoolParameter(default=False)
    from_costs = luigi.BoolParameter(default=False)

    output_path = luigi.Parameter()
    assignment_out_key = luigi.Parameter()
    output_key = luigi.Parameter(default='')

    def find_sizes(self, dep):
        # find segemnts that should be merged according to the size filter
        un_task = getattr(unique_tasks,
                          self._get_task_name('FindUniques'))
        dep = un_task(tmp_folder=self.tmp_folder,
                      max_jobs=self.max_jobs,
                      config_dir=self.config_dir,
                      input_path=self.path,
                      input_key=self.segmentation_key,
                      return_counts=True,
                      dependency=dep,
                      prefix='size-filter-and-graph-watershed')
        sf_task = getattr(size_filter_tasks,
                          self._get_task_name('SizeFilterBlocks'))
        dep = sf_task(tmp_folder=self.tmp_folder,
                      max_jobs=self.max_jobs,
                      config_dir=self.config_dir,
                      input_path=self.path,
                      input_key=self.segmentation_key,
                      size_threshold=self.size_threshold,
                      target_number=self.target_number,
                      dependency=dep)
        return dep

    def requires(self):
        assert (self.size_threshold is None) != (self.target_number is None)
        dep = self.dependency

        # find the sizes for all segments
        dep = self.find_sizes(dep)

        # run graph watershed to merge in all small segments
        filter_path = os.path.join(self.tmp_folder, 'discard_ids.npy')
        gws_task = getattr(gws_tasks,
                           self._get_task_name('GraphWatershedAssignments'))
        dep = gws_task(tmp_folder=self.tmp_folder, max_jobs=self.max_jobs,
                       config_dir=self.config_dir, dependency=dep,
                       problem_path=self.problem_path, graph_key=self.graph_key,
                       features_key=self.features_key,
                       assignment_path=self.path, assignment_key=self.assignment_key,
                       output_path=self.output_path, output_key=self.assignment_out_key,
                       filter_nodes_path=filter_path,
                       relabel=self.relabel, from_costs=self.from_costs)

        if self.output_key != '':
            assert self.fragments_key != ''
            write_task = getattr(write_tasks,
                                 self._get_task_name('Write'))
            dep = write_task(tmp_folder=self.tmp_folder, max_jobs=self.max_jobs,
                             config_dir=self.config_dir, dependency=dep,
                             input_path=self.path, input_key=self.fragments_key,
                             output_path=self.output_path, output_key=self.output_key,
                             assignment_path=self.output_path, assignment_key=self.assignment_out_key,
                             identifier='size-filter-graph-ws')
        return dep

    @staticmethod
    def get_config():
        configs = super(SizeFilterAndGraphWatershedWorkflow,
                        SizeFilterAndGraphWatershedWorkflow).get_config()
        configs.update({'size_filter_blocks': size_filter_tasks.SizeFilterBlocksLocal.default_task_config(),
                        'graph_watershed_assignments':
                        gws_tasks.GraphWatershedAssignmentsLocal.default_task_config(),
                        'write': write_tasks.WriteLocal.default_task_config()})
        return configs


##
# The model uses elements from both the Transformer Encoder as introduced in
# “Attention is All You Need” (https://arxiv.org/pdf/1706.03762.pdf) and the
# Message Passing Neural Network (MPNN) as described in "Neural Message Passing
# for Quantum Chemistry" paper (https://arxiv.org/pdf/1704.01212.pdf) .
#
# The overall architecture most closely resembles the Transformer Encoder with
# stacked encoder blocks and layers connected through residual connections with
# layer norm. In this case however the encoder blocks are build up of two
# message passing layers, followed by three different types of attention layers
# with a final pointwise feed-forward network.
#
# Both message passing layers use a slightly modified version of the edge
# networks as detailed in the MPNN paper. The first layer allows message passing
# between bonded atoms, whereas the second layer does so for the atom pairs for
# which we need to predict the scalar coupling constant. Unlike the attention
# layers the message passing layers' parameters are tied across blocks.
#
# The three attention layers are:
#   1. distance based gaussian attention
#   2. graph distance based attention
#   3. scaled dot product self attention
#
# Although the final layers in the block resemble the encoder blocks of the
# Transformer model, there are several additional layers designed specifically
# to capture the structure and relationships among atoms in a molecule.
#
# Much of the code is adopted from the annotated version of the Transformer
# paper, which can be found here
# (http://nlp.seas.harvard.edu/2018/04/03/attention.html).

import math
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
from fcnet import FullyConnectedNet, hidden_layer
from scatter import scatter_mean
from layernorm import LayerNorm

def clones(module, N):
    """Produce N identical layers."""
    return torch.nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class SublayerConnection(nn.Module):
    """
    A residual connection followed by a layer norm.
    Note for code simplicity the norm is first as opposed to last.
    """
    def __init__(self, size, dropout):
        super().__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        """Apply residual connection to any sublayer with the same size."""
        return x + self.dropout(sublayer(self.norm(x)))

def _gather_nodes(x, idx, sz_last_dim):
    idx = idx.unsqueeze(-1).expand(-1, -1, sz_last_dim)
    return x.gather(1, idx)

class ENNMessage(nn.Module):
    """
    The edge network message passing function from the MPNN paper. Optionally
    adds and additional cosine angle based attention mechanism over incoming
    messages.
    """
    PAD_VAL = -999

    def __init__(self, d_model, d_edge, kernel_sz, enn_args={}, ann_args=None):
        super().__init__()
        assert kernel_sz <= d_model
        self.d_model, self.kernel_sz = d_model, kernel_sz
        self.enn = FullyConnectedNet(d_edge, d_model*kernel_sz, **enn_args)
        if ann_args: self.ann = FullyConnectedNet(1, d_model, **ann_args)
        else: self.ann = None

    def forward(self, x, edges, pairs_idx, angles=None, angles_idx=None, t=0):
        """Note that edges and pairs_idx raw inputs are for a unidirectional
        graph. They are expanded to allow bidirectional message passing."""
        if t==0:
            self.set_a_mat(edges)
            if self.ann: self.set_attn(angles)
            # concat reversed pairs_idx for bidirectional message passing
            self.pairs_idx = torch.cat([pairs_idx, pairs_idx[:,:,[1,0]]], dim=1)
        return self.add_message(torch.zeros_like(x), x, angles_idx)

    def set_a_mat(self, edges):
        n_edges = edges.size(1)
        a_vect = self.enn(edges)
        a_vect = a_vect / (self.kernel_sz ** .5) # rescale
        mask = edges[:,:,0,None].expand(a_vect.size())==self.PAD_VAL
        a_vect = a_vect.masked_fill(mask, 0.0)
        self.a_mat = a_vect.view(-1, n_edges, self.d_model, self.kernel_sz)
        # concat a_mats for bidirectional message passing
        self.a_mat = torch.cat([self.a_mat, self.a_mat], dim=1)

    def set_attn(self, angles):
        angles = angles.unsqueeze(-1)
        self.attn = self.ann(angles)
        mask = angles.expand(self.attn.size())==self.PAD_VAL
        self.attn = self.attn.masked_fill(mask, 0.0)

    def add_message(self, m, x, angles_idx=None):
        """Add message for atom_{i}: m_{i} += sum_{j}[attn_{ij} A_{ij}x_{j}]."""
        # select the 'x_{j}' feeding into the 'm_{i}'
        x_in = _gather_nodes(x, self.pairs_idx[:,:,1], self.d_model)

        # do the matrix multiplication 'A_{ij}x_{j}'
        if self.kernel_sz==self.d_model: # full matrix multiplcation
            ax = (x_in.unsqueeze(-2) @ self.a_mat).squeeze(-2)
        else: # do a convolution
            x_padded = F.pad(x_in, self.n_pad)
            x_unfolded = x_padded.unfold(-1, self.kernel_sz, 1)
            ax = (x_unfolded * self.a_mat).sum(-1)

        # apply atttention
        if self.ann:
            n_pairs = self.pairs_idx.size(1)
            # average all attn(angle_{ijk}) per edge_{ij}.
            # i.e.: attn_{ij} = sum_{k}[attn(angle_{ijk})] / n_angles_{ij}
            ave_att = scatter_mean(self.attn, angles_idx, num=n_pairs, dim=1,
                                   out=torch.ones_like(ax))
            ax = ave_att * ax

        # sum up all 'A_{ij}h_{j}' per node 'i'
        idx_0 = self.pairs_idx[:,:,0,None].expand(-1, -1, self.d_model)
        return m.scatter_add(1, idx_0, ax)

    @property
    def n_pad(self):
        k = self.kernel_sz
        return (k // 2, k // 2 - int(k % 2 == 0))

class MultiHeadedDistAttention(nn.Module):
    """Generalizes the euclidean and graph distance based attention layers."""
    def __init__(self, h, d_model):
        super().__init__()
        self.d_model, self.d_k, self.h = d_model, d_model // h, h
        self.attn = None
        self.linears = clones(nn.Linear(d_model, d_model), 2)

    def forward(self, dists, x, mask):
        batch_size = x.size(0)
        x = self.linears[0](x).view(batch_size, -1, self.h, self.d_k)
        x, self.attn = self.apply_attn(dists, x, mask)
        x = x.view(batch_size, -1, self.h * self.d_k)
        return self.linears[-1](x)

    def apply_attn(self, dists, x, mask):
        attn = self.create_raw_attn(dists, mask)
        attn = attn.transpose(-2,-1).transpose(1, 2)
        x = x.transpose(1, 2)
        x = torch.matmul(attn, x)
        x = x.transpose(1, 2).contiguous()
        return x, attn

    def create_raw_attn(self, dists, mask):
        pass

class MultiHeadedGraphDistAttention(MultiHeadedDistAttention):
    """Attention based on an embedding of the graph distance matrix."""
    MAX_GRAPH_DIST = 10
    def __init__(self, h, d_model):
        super().__init__(h, d_model)
        self.embedding = nn.Embedding(self.MAX_GRAPH_DIST+1, h)

    def create_raw_attn(self, dists, mask):
        emb_dists = self.embedding(dists)
        mask = mask.unsqueeze(-1).expand(emb_dists.size())
        emb_dists = emb_dists.masked_fill(mask==0, -1e9)
        return F.softmax(emb_dists, dim=-2).masked_fill(mask==0, 0)

class MultiHeadedEuclDistAttention(MultiHeadedDistAttention):
    """Attention based on a parameterized normal pdf taking a molecule's
    euclidean distance matrix as input."""
    def __init__(self, h, d_model):
        super().__init__(h, d_model)
        self.log_prec = nn.Parameter(torch.Tensor(1, 1, 1, h))
        self.locs = nn.Parameter(torch.Tensor(1, 1, 1, h))
        nn.init.normal_(self.log_prec, mean=0.0, std=0.1)
        nn.init.normal_(self.locs, mean=0.0, std=1.0)

    def create_raw_attn(self, dists, mask):
        dists = dists.unsqueeze(-1).expand(-1, -1, -1, self.h)
        z = torch.exp(self.log_prec) * (dists - self.locs)
        pdf = torch.exp(-0.5 * z ** 2)
        return pdf / pdf.sum(dim=-2, keepdim=True).clamp(1e-9)

def attention(query, key, value, mask=None, dropout=None):
    """Compute 'Scaled Dot Product Attention'."""
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None: scores = scores.masked_fill(mask==0, -1e9)
    p_attn = F.softmax(scores, dim=-1).masked_fill(mask==0, 0)
    if dropout is not None: p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

class MultiHeadedSelfAttention(nn.Module):
    """Applies self-attention as described in the Transformer paper."""
    def __init__(self, h, d_model, dropout=0.1):
        super().__init__()
        self.d_model, self.d_k, self.h = d_model, d_model // h, h
        self.attn = None
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.dropout = nn.Dropout(p=dropout) if dropout > 0.0 else None

    def forward(self, x, mask):
        # Same mask applied to all h heads.
        mask = mask.unsqueeze(1)
        batch_size = x.size(0)

        # 1) Do all the linear projections in batch from d_model => h x d_k
        query, key, value = [
            l(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
            for l in self.linears[:3]
        ]

        # 2) Apply attention on all the projected vectors in batch.
        x, self.attn = attention(query, key, value, mask, self.dropout)

        # 3) "Concat" using a view and apply a final linear.
        x = x.transpose(1, 2).contiguous()
        x = x.view(batch_size, -1, self.d_model)
        return self.linears[-1](x)

class AttendingLayer(nn.Module):
    """Stacks the three attention layers and the pointwise feedforward net."""
    def __init__(self, size, eucl_dist_attn, graph_dist_attn, self_attn, ff,
                 dropout):
        super().__init__()
        self.eucl_dist_attn = eucl_dist_attn
        self.graph_dist_attn = graph_dist_attn
        self.self_attn = self_attn
        self.ff = ff
        self.subconns = clones(SublayerConnection(size, dropout), 4)
        self.size = size

    def forward(self, x, eucl_dists, graph_dists, mask):
        eucl_dist_sub = lambda x: self.eucl_dist_attn(eucl_dists, x, mask)
        x = self.subconns[0](x, eucl_dist_sub)
        graph_dist_sub = lambda x: self.graph_dist_attn(graph_dists, x, mask)
        x = self.subconns[1](x, graph_dist_sub)
        self_sub = lambda x: self.self_attn(x, mask)
        x = self.subconns[2](x, self_sub)
        return self.subconns[3](x, self.ff)

class MessagePassingLayer(nn.Module):
    """Stacks the bond and scalar coupling pair message passing layers."""
    def __init__(self, size, bond_mess, sc_mess, dropout, N):
        super().__init__()
        self.bond_mess = bond_mess
        self.sc_mess = sc_mess
        self.linears = clones(nn.Linear(size, size), 2*N)
        self.subconns = clones(SublayerConnection(size, dropout), 2*N)

    def forward(self, x, bond_x, sc_pair_x, angles, mask, bond_idx, sc_idx,
                angles_idx, t=0):
        bond_sub = lambda x: self.linears[2*t](
            self.bond_mess(x, bond_x, bond_idx, angles, angles_idx, t))
        x = self.subconns[2*t](x, bond_sub)
        sc_sub = lambda x: self.linears[(2*t)+1](
            self.sc_mess(x, sc_pair_x, sc_idx, t=t))
        return self.subconns[(2*t)+1](x, sc_sub)

class Encoder(nn.Module):
    """Encoder stacks N attention layers and one message passing layer."""
    def __init__(self, mess_pass_layer, attn_layer, N):
        super().__init__()
        self.mess_pass_layer = mess_pass_layer
        self.attn_layers = clones(attn_layer, N)
        self.norm = LayerNorm(attn_layer.size)

    def forward(self, x, bond_x, sc_pair_x, eucl_dists, graph_dists, angles,
                mask, bond_idx, sc_idx, angles_idx):
        """Pass the inputs (and mask) through each block in turn. Note that for
        each block the same message passing layer is used."""
        for t, attn_layer in enumerate(self.attn_layers):
            x = self.mess_pass_layer(x, bond_x, sc_pair_x, angles, mask,
                                     bond_idx, sc_idx, angles_idx, t)
            x = attn_layer(x, eucl_dists, graph_dists, mask)
        return self.norm(x)

# After N blocks of message passing and attending, the encoded atom states are
# transferred to the head of the model: a customized feed-forward net for
# predicting the scalar coupling (sc) constant.

# First the relevant pairs of atom states for each sc constant in the batch
# are selected, concatenated and stacked. Also concatenated to the encoded
# states are a set of raw molecule and sc pair specific features. These states
# are fed into a residual block comprised of a dense layer followed by a type
# specific dense layer of dimension 'd_ff' (the same as the dimension used for
# the pointwise feed-forward net).

# The processed states are passed through to a relatively small feed-forward
# net, which predicts each sc contribution seperately plus a residual.
# Ultimately, the predictions of these contributions and the residual are summed
# to predict the sc constant.

def create_contrib_head(d_in, d_ff, act, dropout=0.0, layer_norm=True):
    layers = hidden_layer(d_in, d_ff, False, dropout, layer_norm, act)
    layers += hidden_layer(d_ff, 1, False, 0.0) # output layer
    return nn.Sequential(*layers)

class ContribsNet(nn.Module):
    """The feed-forward net used for the sc contribution and final sc constant
    predictions."""
    N_CONTRIBS = 5
    CONTIB_SCALES = [1, 250, 45, 35, 500] # scales used to make the 5 predictions of similar magnitude

    def __init__(self, d_in, d_ff, vec_in, act, dropout=0.0, layer_norm=True):
        super().__init__()
        contrib_head = create_contrib_head(d_in, d_ff, act, dropout, layer_norm)
        self.blocks = clones(contrib_head, self.N_CONTRIBS)

    def forward(self, x):
        ys = torch.cat(
            [b(x)/s for b,s in zip(self.blocks, self.CONTIB_SCALES)], dim=-1)
        return torch.cat([ys[:,:-1], ys.sum(dim=-1, keepdim=True)], dim=-1)

class MyCustomHead(nn.Module):
    """Joins the sc type specific residual block with the sc contribution
    feed-forward net."""
    PAD_VAL = -999
    N_TYPES = 8

    def __init__(self, d_input, d_ff, d_ff_contribs, pre_layers=[],
                 post_layers=[], act=nn.ReLU(True), dropout=3*[0.], norm=False):
        super().__init__()
        fc_pre = hidden_layer(d_input, d_ff, False, dropout[0], norm, act)
        self.preproc = nn.Sequential(*fc_pre)
        fc_type = hidden_layer(d_ff, d_input, False, dropout[1], norm, act)
        self.types_net = clones(nn.Sequential(*fc_type), self.N_TYPES)
        self.contribs_net = ContribsNet(
            d_input, d_ff_contribs, d_ff, act, dropout[2], layer_norm=norm)

    def forward(self, x, sc_types):
        # stack inputs with a .view for easier processing
        x, sc_types = x.view(-1, x.size(-1)), sc_types.view(-1)
        mask =  sc_types != self.PAD_VAL
        x, sc_types = x[mask], sc_types[mask]

        x_ = self.preproc(x)
        x_types = torch.zeros_like(x)
        for i in range(self.N_TYPES):
            t_idx = sc_types==i
            if torch.any(t_idx): x_types[t_idx] = self.types_net[i](x_[t_idx])
            else: x_types = x_types + 0.0 * self.types_net[i](x_) # fake call (only necessary for distributed training - to make sure all processes have gradients for all parameters)
        x = x + x_types
        return self.contribs_net(x)

class Transformer(nn.Module):
    """Molecule transformer with message passing."""
    def __init__(self, d_atom, d_bond, d_sc_pair, d_sc_mol, N=6, d_model=512,
                 d_ff=2048, d_ff_contrib=128, h=8, dropout=0.1, kernel_sz=128,
                 enn_args={}, ann_args={}):
        super().__init__()
        assert d_model % h == 0
        self.d_model = d_model
        c = copy.deepcopy
        bond_mess = ENNMessage(d_model, d_bond, kernel_sz, enn_args, ann_args)
        sc_mess = ENNMessage(d_model, d_sc_pair, kernel_sz, enn_args)
        eucl_dist_attn = MultiHeadedEuclDistAttention(h, d_model)
        graph_dist_attn = MultiHeadedGraphDistAttention(h, d_model)
        self_attn = MultiHeadedSelfAttention(h, d_model, dropout)
        ff = FullyConnectedNet(d_model, d_model, [d_ff], dropout=[dropout])

        message_passing_layer = MessagePassingLayer(
            d_model, bond_mess, sc_mess, dropout, N)
        attending_layer = AttendingLayer(
            d_model, c(eucl_dist_attn), c(graph_dist_attn), c(self_attn), c(ff),
            dropout
        )

        self.projection = nn.Linear(d_atom, d_model)
        self.encoder = Encoder(message_passing_layer, attending_layer, N)
        self.write_head = MyCustomHead(
            2 * d_model + d_sc_mol, d_ff, d_ff_contrib, norm=True)

    def forward(self, atom_x, bond_x, sc_pair_x, sc_mol_x, eucl_dists,
                graph_dists, angles, mask, bond_idx, sc_idx, angles_idx,
                sc_types):
        x = self.encoder(
            self.projection(atom_x), bond_x, sc_pair_x, eucl_dists, graph_dists,
            angles, mask, bond_idx, sc_idx, angles_idx
        )
        # for each sc constant in the batch select and concat the relevant pairs
        # of atom  states.
        x = torch.cat(
            [_gather_nodes(x, sc_idx[:,:,0], self.d_model),
             _gather_nodes(x, sc_idx[:,:,1], self.d_model),
             sc_mol_x], dim=-1
        )
        return self.write_head(x, sc_types)