# stardist / tensorflow env variables setup
import os
from os.path import join

os.environ["OMP_NUM_THREADS"] = "10"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

from pathlib import Path
import pickle
import napari
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import argparse
import dask.array as da
from tqdm import tqdm
from skimage.transform import rescale

from tifffile import imread
import matplotlib.pyplot as plt
from PIL import Image
import os

# from dask_image.imread import imread

def show_area_hist(params_file_path):
    params_df = pd.read_csv(params_file_path, index_col=0)
    params_df["area"].plot(kind="hist", bins=100, title="Area histogram")
    plt.show()

def initialize_viewer(imgs):
    viewer = napari.Viewer()
    viewer.window.resize(1800, 1000)
    layers = viewer.add_image(imgs, channel_axis=3, name="raw")

    return viewer

def add_detection(viewer, detection_file_path):
    full_detection = np.load(detection_file_path)
    detection_name = f"{Path(detection_file_path).name}_detection"
    edges_name = f"{Path(detection_file_path).name}_edges"
    viewer.add_image(full_detection['detection'], visible=False, name=detection_name)
    viewer.add_image(full_detection['edges'], blending="additive", colormap="magma", name=edges_name)

    return detection_name, edges_name

def add_tracks(viewer, tracks_file_path, track_label_file_path, graph_file_path):
    tracks_df = pd.read_pickle(tracks_file_path)
    track_name = Path(tracks_file_path).name
    with open(graph_file_path, 'rb') as f:
        graph = pickle.load(f)
    labels = np.load(track_label_file_path)
    track_label_name = Path(track_label_file_path).name
    # layer = viewer.add_tracks(tracks_df[["track_id", "t", "y", "x"]].values, graph=graph, name=track_name)
    layer = viewer.add_tracks(tracks_df[["track_id", "t", "y", "x"]].values, graph=graph, name=track_name)

    layer.visible = False
    viewer.add_labels(labels, name=track_label_name)

    return track_name, track_label_name

def display_folder(folder_path, img_path, config_id, n_frames=None):
    data_dir = Path(folder_path)
    normalized_path = data_dir / "normalized.npy"
    cellpose_path = data_dir / "cellpose_labels.npy"
    wscp_path = data_dir / "wscp_labels.npy"
    stardist_path = data_dir / "stardist_labels.npy"
    wssd_path = data_dir / "wssd_labels.npy"
    detection_path = data_dir / config_id / "detections.npz"
    tracks_path = data_dir / config_id / "tracks.pkl"

    track_label_path = data_dir / config_id / "track_labels.npy"
    graph_path = data_dir / config_id / "graph.pkl"

    tracks_ppc_path = data_dir / config_id / "tracks_ppc.pkl"
    graph_ppc_path = data_dir / config_id / "graph_ppc.pkl"

    tracks_ppc_pruned_path = data_dir / config_id / "tracks_ppc_pruned.pkl"
    graph_ppc_pruned_path = data_dir / config_id / "graph_ppc_pruned.pkl"
    track_label_ppc_pruned_path = data_dir / config_id / "track_labels_ppc_pruned.npy"

    imgs = imread(img_path)

    imgs = imgs[:, 1:, :, :]

    if n_frames is not None:
        imgs = imgs[:n_frames]
    imgs = np.swapaxes(imgs, 1, 3)

    chunks = (1, imgs.shape[1], imgs.shape[2], 1)  # chunk size used to compress input

    viewer = initialize_viewer(imgs)

    if normalized_path.exists():
        # image size is (t, y, x, c)
        normalized = da.from_array(np.load(normalized_path), chunks=chunks)
        # if normalized.shape != imgs.shape:
        #     rescale_factor = imgs.shape[1] / normalized.shape[1]
        #     upscaled = np.zeros(imgs.shape, dtype=normalized.dtype)
        #     for i in range(imgs.shape[0]):
        #         upscaled[i, ...] = rescale(normalized[i, ...].compute(), rescale_factor, channel_axis=2,
        #                                    anti_aliasing=True, preserve_range=True)
        #     normalized = da.from_array(upscaled, chunks=chunks)

        layers = viewer.add_image(normalized, rgb=False, channel_axis=3, name="normalized")
        for layer in layers:
            layer.visible = False

    if cellpose_path.exists():
        cellpose_labels = da.from_array(np.load(cellpose_path), chunks=chunks)
        # if cellpose_labels.shape != imgs.shape:
        #     rescale_factor = imgs.shape[1] / cellpose_labels.shape[1]
        #     upscaled = np.zeros(imgs.shape, dtype=cellpose_labels.dtype)
        #     for i in range(imgs.shape[0]):
        #         upscaled[i, ...] = rescale(cellpose_labels[i, ...].compute(), rescale_factor, channel_axis=2,
        #                                    anti_aliasing=False, preserve_range=True)
        #     cellpose_labels = da.from_array(upscaled, chunks=chunks)
        for i in range(cellpose_labels.shape[-1]):
            layer = viewer.add_labels(cellpose_labels[:, :, :, i], name=f"cellpose labels {i}")
            layer.visible = False

    if wscp_path.exists():
        wscp_labels = da.from_array(np.load(wscp_path), chunks=chunks)
        # if wscp_labels.shape != imgs.shape:
        #     rescale_factor = imgs.shape[1] / wscp_labels.shape[1]
        #     upscaled = np.zeros(imgs.shape, dtype=wscp_labels.dtype)
        #     for i in range(imgs.shape[0]):
        #         upscaled[i, ...] = rescale(wscp_labels[i, ...].compute(), rescale_factor, channel_axis=2,
        #                                    anti_aliasing=False, preserve_range=True)
        #     wscp_labels = da.from_array(upscaled, chunks=chunks)
        for i in range(wscp_labels.shape[-1]):
            layer = viewer.add_labels(wscp_labels[:, :, :, i], name=f"watershed with cellpose labels {i}")
            layer.visible = False

    if stardist_path.exists():
        stardist_labels = da.from_array(np.load(stardist_path), chunks=chunks)
        # if stardist_labels.shape != imgs.shape:
        #     rescale_factor = imgs.shape[1] / stardist_labels.shape[1]
        #     upscaled = np.zeros(imgs.shape, dtype=stardist_labels.dtype)
        #     for i in range(imgs.shape[0]):
        #         upscaled[i, ...] = rescale(stardist_labels[i, ...].compute(), rescale_factor, channel_axis=2,
        #                                    anti_aliasing=False, preserve_range=True)
        #     stardist_labels = da.from_array(upscaled, chunks=chunks)
        for i in range(stardist_labels.shape[-1]):
            layer = viewer.add_labels(stardist_labels[:, :, :, i], name=f"stardist labels {i}")
            layer.visible = False

    if wssd_path.exists():
        wssd_labels = da.from_array(np.load(wssd_path), chunks=chunks)
        # if wssd_labels.shape != imgs.shape:
        #     rescale_factor = imgs.shape[1] / wssd_labels.shape[1]
        #     upscaled = np.zeros(imgs.shape, dtype=wssd_labels.dtype)
        #     for i in range(imgs.shape[0]):
        #         upscaled[i, ...] = rescale(wssd_labels[i, ...].compute(), rescale_factor, channel_axis=2,
        #                                    anti_aliasing=False, preserve_range=True)
        #     wssd_labels = da.from_array(upscaled, chunks=chunks)
        for i in range(wssd_labels.shape[-1]):
            layer = viewer.add_labels(wssd_labels[:, :, :, i], name=f"watershed with stardist labels {i}")
            layer.visible = False

    if detection_path.exists():
        detection_name, edges_name = add_detection(viewer, detection_path)
        viewer.layers[edges_name].visible = False

    if tracks_path.exists() and graph_path.exists():
        track_name, track_label_name = add_tracks(viewer, tracks_path, track_label_path, graph_path)
    if tracks_ppc_path.exists() and graph_ppc_path.exists():
        track_name, track_label_name = add_tracks(viewer, tracks_ppc_path, track_label_path, graph_ppc_path)
    if tracks_ppc_pruned_path.exists() and graph_ppc_pruned_path.exists():
        track_name, track_label_name = add_tracks(viewer, tracks_ppc_pruned_path, track_label_ppc_pruned_path, graph_ppc_pruned_path)
    # screenshot = viewer.screenshot()
    # viewer.close()
    # return screenshot

    napari.run()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Display images and optional overlays.')
    parser.add_argument('--file', type=str, default="4T1 p27 trial period.HTD - Well D02 Field #3.tif",
                        help='Path to the image file')
    parser.add_argument('--config_id', type=str, default="100", required=False, help='Name of config file')

    parser.add_argument('--n_frames', type=int, default=100, help='Number of frames (optional)')
    args = parser.parse_args()
    # args.file = 'demo.tif'

    experiment = Path(args.file).stem
    output_dir = join(Path(__file__).parent, "output", experiment)
    print(f'rsync -avz --progress jorisg@192.168.1.203:/home/jorisg/projects/DSL/output/ "{join(Path(__file__).parent, "output")}/"')
    os.system(
        f'rsync -avz --progress jorisg@192.168.1.203:/home/jorisg/projects/DSL/output/ {join(Path(__file__).parent, "output")}/')
    input_file = join(Path(__file__).parent, "input", args.file)

    display_folder(output_dir, input_file, str(args.config_id), args.n_frames)

    # # List of config IDs
    # config_ids = [1, 2, 3, 4, 5, 6, 11, 12, 13, 14, 15, 16, 21, 22, 23, 24, 25, 26, 31, 32, 33, 34, 35, 36, 41, 42, 43, 44, 45, 46, 51, 52, 53, 54, 55, 56]
    # # config_ids = [1, 2, 3]
    #
    # # Determine the number of rows and columns for the grid
    #
    # n_cols = 6
    # n_rows = len(config_ids) // n_cols + (len(config_ids) % n_cols > 0)
    #
    # # Create a figure with specified dimensions
    # fig, axes = plt.subplots(n_rows, n_cols, figsize=(20, 20))  # Adjust figsize as needed
    #
    # # Remove the space between plots
    # plt.subplots_adjust(wspace=0, hspace=0)
    #
    # # Iterate over the grid and plot
    # for i, config_id in tqdm(enumerate(config_ids)):
    #     row = i // n_cols
    #     col = i % n_cols
    #
    #     args.config_id = str(config_id)
    #
    #     screenshot = display_folder(output_dir, input_file, str(args.config_id), args.n_frames)
    #     axes[row, col].imshow(screenshot)
    #     axes[row, col].set_title(str(config_id), fontsize=10)
    #
    #     # Remove axes
    #     axes[row, col].axis('off')
    #
    # # Hide any unused subplots
    # for i in range(len(config_ids), n_rows * n_cols):
    #     row = i // n_cols
    #     col = i % n_cols
    #     axes[row, col].axis('off')
    #
    # # Save the figure
    # # plt.show()
    # plt.savefig('combined_screenshots.png', dpi=600)  # Adjust the filename and dpi as needed
    # #
    # # # Close the plot
    # # plt.close(fig)


# Syntax error: It occurs when the parser detects a syntatically incorrect statement
# a = 5 print(a)

#a = 5
#print(a))

# Even if a statement is syntatically correct it might cause an error when executed
#and it is called exception error

#for instance

#import somemodule
#this will cause module error

#a = 5
#b = c

#this will cause name error

# value error: happend if the function or operation recieves an argument
# that has the right type but an inappropriate value

a = [1,2,3]
a.remove(1)

#this works but if we put 4 instead we will get value error

a[6]
#this will occur index error

my_dict = {
    'name': 'Jack'
}
my_dict['age']
# this will give key error

# If you want to force exception to occur
# x = -5
# if x < 0:
#     raise Exception('x should be positive')

# x = -5
# # the assert statement  will throught an assertion error when your condition is not True
# assert(x >=0 ), 'x is not positive'

try:
    a = 5 / 0
except Exception as e:
    print(e)


from .widgets import AdminIconChooser
try:
    from wagtail.admin.panels import FieldPanel
except ImportError:
    from wagtail.admin.edit_handlers import FieldPanel

class IconsChooserPanel(FieldPanel):
    def __init__(self, field_name, group=None, *args, **kwargs):
        super().__init__(field_name, *args, **kwargs)
        self.group = group

    def clone_kwargs(self):
        kwargs = super().clone_kwargs()
        kwargs.update(
            group=self.group,
        )
        return kwargs

    def widget_overrides(self):
        return {
            self.field_name: AdminIconChooser(group=self.group),
        }

    def get_form_options(self):
        options = super().get_form_options()
        options['widgets'] = {
            self.field_name: AdminIconChooser(group=self.group),
        }
        return options


class Node:
    def __init__(self,value) -> None:
        self.value=value
        self.next=None
    def set_next(self,next):
        self.next=next
    def get_next(self):
        return self.next

class SingleLinked:

    def __init__(self):
        self.head=None
        self.size=0

    def add(self,value):
        if self.size==0:
            self.head=Node(value)
            self.size+=1
        else:
            next=self.head
            while(1):
                if next.get_next()==None:
                    next.set_next(Node(value))
                    self.size+=1
                    break
                else:
                    next=next.get_next()

    def add_to_index(self,value,index):
        next=self.head
        if index==0:
            if self.size==0:
                print("Error")
                return
            elif self.size!=0:
                new_value,new_next=self.head.value,self.head.get_next()
                self.head.value=value
                self.head.set_next(new_next)
                self.add_to_index(new_value,1)
            return
        elif index>self.size or index<0:
            print("Error")
            return
        else:
            for i in range(self.size):
                if i!=index-1:
                    next=next.get_next()
                else:
                    new=Node(value)
                    new.set_next(next.get_next())
                    next.set_next(new)
                    self.size+=1

    def get(self,index):
        next=self.head
        if index>self.size-1 or index<0:
            return "Error"
        else:
            for i in range(self.size):
                if i!=index:
                    next=next.get_next()
                else:
                    return next.value

    def set(self,value,index):
        next=self.head
        if index>self.size-1 or index<0:
            print("Error")
            return
        else:
            for i in range(self.size):
                if i!=index:
                    next=next.get_next()
                else:
                    next.value=value
        self.print_linked()

    def clear(self):
        if self.size==0:
            return
        self.size=0
        self.head.set_next(None)

    def is_empty(self):
        return not bool(self.size)

    def remove(self,index):
        next=self.head
        if index==0:
            if self.size==0:
                print("Error")
                return
            elif self.size!=0:
                self.head=self.head.get_next()
                self.size-=1

        elif index>self.size-1 or index<0:
            print("Error")
            return
        else:
            for i in range(self.size):
                if i!=index-1:
                    next=next.get_next()
                else:
                    next.set_next(next.get_next().get_next())
                    self.size-=1
        #self.print_linked()
    def sublist(self,first,last):
        if first <0 or first >self.size-1 or last<0 or last > self.size-1 or last<first:
            print("Error")
            return
        else:
            next=self.head
            my_sub_list=[]
            for i in range(self.size):
                if i!=first:
                    next=next.get_next()
                else:
                    my_sub_list.append(next.value)
                    for j in range(i,self.size):
                        next=next.get_next()
                        if j == last:
                            print(my_sub_list)
                            return
                        my_sub_list.append(next.value)

    def contains(self,value):
        next=self.head
        for i in range(self.size):
            if value==next.value:
                print("True")
                return
            else:
                next=next.get_next()
        print("False")

    def linked_size(self):
        return self.size

    def print_linked(self):
        my_list=[]
        if self.size==0:
            print([])
            return
        else:
            next=self.head
            value=self.head.value

            while(next.get_next()!=None):
                my_list.append(value)
                next=next.get_next()
                value=next.value
            my_list.append(next.value)
        print(my_list)

class Stack:

    def __init__(self,stack_array):
        self.linked= SingleLinked()
        for i in range(len(stack_array)):
            try:
                self.linked.add(int(stack_array[i]))

            except (ValueError):
                    del stack_array[i]
    def push(self,value):
        if  my_stack.size()==0:
            self.linked.add(value)
        else:
            self.linked.add_to_index(value,0)

    def pop(self):
        self.linked.remove(0)

    def peek(self):
        print(self.linked.head.value)

    def isEmpty(self):
        print(self.linked.is_empty())

    def size(self):
        return self.linked.linked_size()
    def print_stack(self):
        self.linked.print_linked()

text_array=input()
text_array=text_array[1:len(text_array)-1]
stack_array=text_array.split(",")
my_stack=Stack(stack_array)

operation=input().strip()

if operation=="push":
    #print(my_stack.size())
    value=int(input().strip())
    my_stack.push(value)
    my_stack.print_stack()
elif operation=="pop":
    if my_stack.size()==0:
        print("Error")
    else:
        my_stack.pop()
        my_stack.print_stack()
elif operation=="peek":
    if my_stack.size()==0:
        print("Error")
    else:
        my_stack.peek()
elif operation=="isEmpty":
    my_stack.isEmpty()
elif operation=="size":
    print(my_stack.size())
else:
    print("Error")


import itertools
def check1(target):
    count = 0;
    for idx in range(2):
        flag = True
        tg = target[0][idx]
        for other in range(1, len(target)):
            now = target[other]
            if tg not in now:
                flag = False;
                break
        if flag:
            count += 1
    if count >= 1:
        return True
    else:
        return False

def check2(target):
    flag = True
    for item in target:
        if item[0] == item[1]:
            flag = False
            break
    return flag

def find(subList):
    subs = []
    for idx in range(0, len(subList), +2):
        subs.append(subList[idx:idx+2])
    if check1(subs) & check2(subs):
        return True
    else:
        return False

def solution(a):
    answer = 0
    n = len(a)
    if n == 1:
        return 0
    if n % 2 != 0:
        n -= 1
    for i in range(n, 0, -2):
        flag = False
        temps = itertools.combinations(a, i)
        for temp in temps:
            if find(temp):
                flag = True
                break
        if flag:
            return i

    return answer

print(solution([0]))
print(solution([5,2,3,3,5,3]))
print(solution([0,3,3,0,7,2,0,2,2,0]))


import numpy as np
from constants import SMALL_NON_ZERO, SMALL_LOG_NON_ZERO

def one_hot(value, classes):
	res = np.zeros(classes, dtype=np.int)
	res[value] = 1
	return res

def to_log_probability(pot, min_val=SMALL_NON_ZERO, max_val=1):
    a = np.clip(pot, min_val, max_val)
    idx = np.where(a == min_val)
    a = np.log(a)
    a[idx] = SMALL_LOG_NON_ZERO
    return a


import unittest
import pytest
import numpy.testing as npt
import yaml
from pathlib import Path

from rearrangement.evaluation.state_eval import StateEvaluation
from rearrangement.evaluation.scene_prediction import ScenePrediction

SAVE = False

class TestEvaluation(unittest.TestCase):
    def create_model(self):
        log_dir = Path('./data/models')
        run_id = 'state-prediction-71538'
        self.run_id = run_id
        return ScenePrediction.from_run(log_dir, run_id)

    def compare_ref_outputs(self, ref_path, outputs, datasets, save=False):
        if save:
            Path(ref_path).write_text(yaml.dump({self.run_id: outputs}))

        ref_data = yaml.load(Path(ref_path).read_text())[self.run_id]
        for ds in datasets:
            for metric, metric_value in ref_data[ds].items():
                if isinstance(metric_value, dict):
                    for k, v in metric_value.items():
                        npt.assert_almost_equal(v, outputs[ds][metric][k], decimal=5)
                else:
                    npt.assert_almost_equal(metric_value, outputs[ds][metric], decimal=5)

    @pytest.mark.eval
    def test_eval_real(self):
        model = self.create_model()
        datasets = ('real-cubes-1to12', 'real-cubes-1to6')
        outputs = dict()
        for ds in datasets:
            pose_eval = StateEvaluation(ds, model.config)
            outputs[ds], _,_ = pose_eval.evaluate(model)

        self.compare_ref_outputs('rearrangement/test/test_data_real.yaml', outputs, datasets, save=SAVE)


"""A basic chat server example."""

import sys
from logging import basicConfig, info

from earwax_server import ConnectionContext, Server

server: Server = Server()

@server.event
def on_connect(ctx: ConnectionContext) -> None:
    """Let everyone else know someone has arrived."""
    other: ConnectionContext
    for other in server.connections:
        if other is ctx:
            ctx.send_string('Welcome to the server.')
            ctx.send_string(
                f'Other connections: {len(server.connections) - 1}.'
            )
        else:
            other.send_string(f'New connection from {ctx.logger.name}.')

@server.event
def on_disconnect(ctx: ConnectionContext) -> None:
    """Tell everyone who has just left."""
    other: ConnectionContext
    for other in server.connections:
        other.send_string(f'{ctx.logger.name} has disconnected.')

@server.event
def on_data(ctx: ConnectionContext, data: bytes) -> None:
    """Send a message to everyone."""
    if data.decode().strip() == 'quit':
        ctx.send_string('Goodbye.')
        ctx.disconnect()
        return
    other: ConnectionContext
    name: bytes = ctx.logger.name.encode(sys.getdefaultencoding())
    sep: bytes = ': '.encode(sys.getdefaultencoding())
    message: bytes = name + sep + data
    for other in server.connections:
        other.send_bytes(message)

if __name__ == '__main__':
    basicConfig(stream=sys.stdout, level='INFO')
    info('Server starting...')
    try:
        server.run(1234)
    except KeyboardInterrupt:
        info('Server done.')


import math

rebuildArray = []
printTreeStrings = [""]

class Node:

	def __init__(self, key, leftchild, rightchild, parent):

		self.key = key
		self.leftchild = leftchild
		self.rightchild = rightchild
		self.parent = parent

class Tree:

	def __init__(self, alpha, key):
		self.root = Node(key, None, None, None)
		self.size = 1
		self.maxsize = 1
		self.alpha = alpha

	def insert(self, key, node = 0):
		if node == 0:
			node = self.root
		if key < node.key:
			if node.leftchild:
				self.insert(key, node.leftchild)
			else:
				node.leftchild = Node(key, None, None, node)
				self.size+=1
				self.maxsize = max(self.size, self.maxsize)
				depth = self.depthN(node.leftchild)
				#print("inserting {}...".format(key))
				#print("{} < {}?".format(depth, math.log(self.size, 1/self.alpha)))
				if depth > (math.log(self.size, 1/self.alpha)):
					self.rebalance(node.leftchild)
		else:
			if node.rightchild:
				return self.insert(key, node.rightchild)
			else:
				node.rightchild = Node(key, None, None, node)
				self.size+=1
				self.maxsize = max(self.size, self.maxsize)
				depth = self.depthN(node.rightchild)
				#print("inserting {}...".format(key))
				#print("{} < {}?".format(depth, math.log(self.size, 1/self.alpha)))
				if depth > (math.floor(math.log(self.size, 1/self.alpha))):
					self.rebalance(node.rightchild)
		return

	def rebalance(self, node = 0):
		if node == 0:
			node = self.root
		print("rebalance key: {}".format(node.key))
		scapegoat = node
		while(node.parent):
			node = node.parent
			if self.sizeT(node.leftchild) > self.alpha * self.sizeT(node) or self.sizeT(node.rightchild) > self.alpha * self.sizeT(node):
				scapegoat = node
		print("Scapegoat key: {}".format(scapegoat.key))
		print("children of scapegoat: {}, {}".format(scapegoat.leftchild, scapegoat.rightchild.key))
		self.buildArray(scapegoat)
		print(rebuildArray)
		print("midpoint is {}".format(int(math.floor(float(len(rebuildArray))/2))))
		scapegoat.key = rebuildArray[int(math.floor(float(len(rebuildArray))/2))]
		print(0, int (math.floor(len(rebuildArray)/2))-1)
		print(int (math.floor(float(len(rebuildArray))/2)+1), len(rebuildArray))
		scapegoat.leftchild = self.rebuild(0, int (math.floor(float(len(rebuildArray))/2))-1)
		scapegoat.rightchild = self.rebuild(int (math.floor(float(len(rebuildArray))/2)+1), len(rebuildArray))
		self.reParent(scapegoat)
		for x in range(0, len(rebuildArray)):
			rebuildArray.pop()
		return

	def buildArray(self, node):
		rebuildArray.append(node.key)
		if node.leftchild:
			self.buildArray(node.leftchild)
		if node.rightchild:
			self.buildArray(node.rightchild)
		rebuildArray.sort()
		return

	def reParent(self, node):
		if node.leftchild:
			node.leftchild.parent = node
			self.reParent(node.leftchild)
		if node.rightchild:
			node.rightchild.parent = node
			self.reParent(node.rightchild)
		return

	def rebuild(self, start, end):
		if end >= len(rebuildArray):
			end = len(rebuildArray) - 1
		print(start, end)
		if start == end:
			print("midpoint is {}, value is {}".format(start, rebuildArray[start]))
			return Node(rebuildArray[start], None, None, None)
		if start > end:
			print("done")
			return None
		midpoint = int (math.floor(start + (end - start)/2))
		print("midpoint is {}, value is {}".format(midpoint, rebuildArray[midpoint]))
		temp = Node(rebuildArray[midpoint], None, None, None)
		temp.leftchild = self.rebuild(start, midpoint-1)
		if end - start == 1:
			pass
		else:
			temp.rightchild = self.rebuild(midpoint+1, end)
		return temp

	def sizeT(self, node):
		if node == None:
			return 0
		else:
		 	return self.sizeT(node.leftchild) + self.sizeT(node.rightchild) + 1

	def depthN(self, node):
		depth = 0
		while node.parent:
			node = node.parent
			depth += 1
		return depth

	def delete(self, key):
		self.size-=1
		delNode = self.search(key)
		if not delNode.rightchild and not delNode.leftchild:
			if delNode == delNode.parent.leftchild:
				delNode.parent.leftchild = None
				if self.size <= self.alpha * self.maxsize:
					self.rebalance()
				return
			else:
				delNode.parent.rightchild = None
				if self.size <= self.alpha * self.maxsize:
					self.rebalance()
				return
		if delNode.rightchild:
			successor = delNode.rightchild
			while successor.leftchild:
				successor = successor.leftchild
			if delNode.rightchild == successor:
				#print(successor.key)
				if successor.rightchild:
					delNode.rightchild = successor.rightchild
				else:
					delNode.rightchild = None
				delNode.key = successor.key
				if self.size <= self.alpha * self.maxsize:
					self.rebalance()
				return
			if successor.rightchild:
				successor.parent.leftchild = successor.rightchild
				successor.rightchild.parent = sucessor.parent
				if self.size <= self.alpha * self.maxsize:
					self.rebalance()
				return
			else:
				successor.parent.leftchild = None
				delNode.key = successor.key
				if self.size <= self.alpha * self.maxsize:
					self.rebalance()
				return
		else:
			if delNode == self.root:
				self.root = delNode.leftchild
				self.root.parent = None
				if self.size <= self.alpha * self.maxsize:
					self.rebalance()
				return
			if delNode == delNode.parent.leftchild:
				delNode.parent.leftchild = delNode.leftchild
				if self.size <= self.alpha * self.maxsize:
					self.rebalance()
				return
			else:
				delNode.parent.rightchild = delNode.leftchild
				if self.size <= self.alpha * self.maxsize:
					self.rebalance()
				return

	def search(self, key, node = 0):
		if node == 0:
			node = self.root
		if key == node.key:
			return node
		elif key > node.key:
			return self.search(key, node.rightchild)
		elif key < node.key:
			return self.search(key, node.leftchild)

	"""

	def printTree(self, node = 0):
		if node == 0:
			node = self.root
		print("-({})".format(node.key))
		if node.leftchild:
			self.pTree(1, node.leftchild, 'L')
		if node.rightchild:
			self.pTree(1, node.rightchild, 'R')

	def pTree(self, dashes, node, LR):
		pSt = "-"
		for x in range(0, dashes):
			pSt = pSt + '-'
		pSt += "({}{})".format(LR, node.key)
		print(pSt)
		if node.leftchild:
			self.pTree(dashes+1, node.leftchild, 'L')
		if node.rightchild:
			self.pTree(dashes+1, node.rightchild, 'R')
	"""

	def printTree(self, node, x):
		depth = ""
		if node == 0:
			node = self.root
		if node.rightchild:
			self.printTree(node.rightchild, x+1)
		for i in range(0, x):
			depth += '\t'
		if node.key:
			if node.leftchild and node.rightchild:
				print(depth + str(node.key) + "\t\t\t(children: {}, {})".format(node.leftchild.key, node.rightchild.key))
			elif node.leftchild:
				print(depth + str(node.key) + "\t\t\t(children: {})".format(node.leftchild.key))
			elif node.rightchild:
				print(depth + str(node.key) + "\t\t\t(children: {})".format(node.rightchild.key))
			else:
				print(depth + str(node.key) + "\t\t\t(children: None)")
		if node.leftchild:
			self.printTree(node.leftchild, x+1)

	"""
	def printTree(self, node, depth):
		#print(node.key)
		dashes = ""
		try:
			printTreeStrings[depth] += dashes + " {}".format(node.key)
		except IndexError:
			printTreeStrings.append(str(node.key))
		if node.leftchild:
			self.printTree(node.leftchild, depth+1)
		if node.rightchild:
			self.printTree(node.rightchild, depth+1)
	"""

if __name__ == "__main__":

	while(1):
		command = raw_input('Command : ')
		command = command.split()
		print("")
		if command[0] == "Done":
			exit()
		elif command[0] == "BuildTree" and len(command) == 3:
			SCT = Tree(float (command[1]), int (command[2]))
		elif command[0] == "Insert" and len(command) == 2:
			SCT.insert(int (command[1]))
		elif command[0] == "Search" and len(command) == 2:
			SCT.search(int (command[1]))
		elif command[0] == "Delete" and len(command) == 2:
			SCT.delete(int (command[1]))
		elif command[0] == "Print":
			SCT.printTree(SCT.root, 0)
			"""
			for x in printTreeStrings:
				print(x)
			"""
		else:
			print("{}: incorrect format or command not recognized".format(command[0]))


from django.db import models
from django.contrib.auth.models import Permission
from django.contrib.auth.models import User
from django.contrib.contenttypes.models import ContentType
from django.db.models.signals import post_save
from guardian.shortcuts import get_perms_for_model, assign_perm

from onadata.apps.api.models.team import Team
from onadata.apps.main.models import UserProfile

def create_owner_team_and_permissions(sender, instance, created, **kwargs):
    if created:
        team = Team.objects.create(
            name=Team.OWNER_TEAM_NAME, organization=instance.user,
            created_by=instance.created_by)
        content_type = ContentType.objects.get(
            app_label='api', model='organizationprofile')
        permission, created = Permission.objects.get_or_create(
            codename="is_org_owner", name="Organization Owner",
            content_type=content_type)
        team.permissions.add(permission)
        instance.creator.groups.add(team)

        for perm in get_perms_for_model(instance.__class__):
            assign_perm(perm.codename, instance.user, instance)

            if instance.creator:
                assign_perm(perm.codename, instance.creator, instance)

            if instance.created_by:
                assign_perm(perm.codename, instance.created_by, instance)

class OrganizationProfile(UserProfile):
    """Organization: Extends the user profile for organization specific info

        * What does this do?
            - it has a createor
            - it has owner(s), through permissions/group
            - has members, through permissions/group
            - no login access, no password? no registration like a normal user?
            - created by a user who becomes the organization owner
        * What relationships?
    """

    class Meta:
        app_label = 'api'
        permissions = (
            ('can_add_xform', "Can add/upload an xform to organization"),
            ('view_organizationprofile', "Can view organization profile"),
        )

    is_organization = models.BooleanField(default=True)
    # Other fields here
    creator = models.ForeignKey(User)

    def save(self, *args, **kwargs):
        super(OrganizationProfile, self).save(*args, **kwargs)

    def remove_user_from_organization(self, user):
        """Removes a user from all teams/groups in the organization.

        :param user: The user to remove from this organization.
        """
        for group in user.groups.filter('%s#' % self.user.username):
            user.groups.remove(group)

    def is_organization_owner(self, user):
        """Checks if user is in the organization owners team.

        :param user: User to check.

        :returns: Boolean whether user has organization level permissions.
        """
        has_owner_group = user.groups.filter(
            name='%s#%s' % (self.user.username, Team.OWNER_TEAM_NAME))
        return True if has_owner_group else False

post_save.connect(
    create_owner_team_and_permissions, sender=OrganizationProfile,
    dispatch_uid='create_owner_team_and_permissions')


from discord.ext import commands
from discord import Forbidden, Embed
import discord
import sqlite3

from itertools import islice
from contextlib import redirect_stdout
from io import StringIO
from textwrap import indent
from traceback import format_exception, print_exc
import subprocess

class Owner(commands.Cog):
    def __init__(self, bot):
        self.bot = bot
        self.save = None

    async def cog_check(self, ctx):
        if ctx.author.id == self.bot.owner_id:
            return True
        return False

    @commands.command(hidden=True)
    async def owner(self, ctx):
        return await ctx.send("Samir!")

    @commands.command(aliases=['r'], hidden=True)
    async def reload(self, ctx, arg):
        if arg.isdecimal():
            try:
                self.bot.reload_extension(self.bot.bot_extensions[int(arg)]);
                return await ctx.send("Success!")
            except Exception as e:
                print(e);
                return await ctx.send("Error!")

        try:
            self.bot.reload_extension(str(arg));
            return await ctx.send("Success!")
        except Exception as e:
            print(e);
            return await ctx.send("Error!")

    @commands.command(aliases=['l'], hidden=True)
    async def load(self, ctx, arg):
        try:
            self.bot.load_extension(str(arg));
            return await ctx.send("Success!")
        except Exception as e:
            print(e);
            return await ctx.send("Error!")

    @commands.command(aliases=['u'], hidden=True)
    async def unload(self, ctx, arg):
        try:
            self.bot.unload_extension(str(arg));
            return await ctx.send("Success!")
        except Exception as e:
            print(e);
            return await ctx.send("Error!")

    @commands.command(aliases=['eval','e'], hidden=True)
    async def _evaluate(self, ctx, *, arg):
        def clean_code(content):
            if content.startswith("```") and content.endswith("```"):
                return "\n".join(content.split("\n")[1:])[:-3]
            else:
                return content

        code = clean_code(arg)
        local_variables = {
            'ctx': ctx,
            'game': self.bot.games.get(ctx.channel.id, None),
            'disocrd': discord,
            'bot': self.bot,
            'author': ctx.author
        }
        stdout = StringIO()
        try:
            with redirect_stdout(stdout):
                exec(
                    f"async def function():\n{indent(code, '    ')}",
                    local_variables
                )
                obj = await local_variables["function"]()
                result = (f"--*Output*```python\n{stdout.getvalue()}\n---```\n" +\
                         f"--*RETURN VALUE* ```python\n{obj}```")
        except Exception as e:
            result = "```python\n"
            result += "".join(format_exception(e, e, e.__traceback__)) + "```"
            result = result.replace(r"c:...Path...\ownercog.py", "YOU MUST NOT SEE THIS PATH KEK")

        return await ctx.send(result)

    @commands.command(aliases=['exec','ex'], hidden=True)
    async def _execute(self, ctx, *, arg):
        try:
            sb = subprocess.Popen(args=arg, shell=True, stdout=subprocess.PIPE)
            return await ctx.send(f"```{sb.stdout.read()}```")
        except Exception as e:
            return await ctx.send("```"+"".join(format_exception(e, e, e.__traceback__))+"```")

    @commands.command(hidden=True)
    async def test(self, ctx, ):
        print("Hello World!", ctx.guild_prefix)

def setup(bot):
    bot.add_cog(Owner(bot))


from time import sleep

from p2p_network.node import Node

def new_message_callback(sender_id, msg, msg_type):
    if msg_type == 'string':
        print('< RECEIVED \'{0}\' FROM'.format(msg.decode('utf-8')), sender_id)
    if msg_type == 'file':
        print('< RECEIVED FILE')
        with open('received', 'bw') as write_out:
            write_out.write(msg)

def sending_callback(receiver_id, msg, msg_type):
    if msg_type == 'string':
        print('> SENDING  \'{0}\' TO'.format(msg.decode('utf-8')), receiver_id)
    elif msg_type == 'file':
        print('> SENDING  FILE TO', receiver_id)

if __name__ == '__main__':
    client_id_number = int(input('CLIENT  ID  : ') or '1')
    CLIENT_ID = bytes([client_id_number])
    CLIENT_PORT = 13370 + client_id_number
    print(CLIENT_PORT)
    PARTNER_PORT = 13370 + int(input('PARTNER ID  : ') or '1')
    print(PARTNER_PORT)

    CLIENT_IP = input('CLIENT  IP  : ') or '127.0.0.1'
    print(CLIENT_IP)
    PARTNER_IP = input('PARTNER IP  : ') or '127.0.0.1'
    print(PARTNER_IP)

    node = Node(CLIENT_IP, CLIENT_PORT, CLIENT_ID)
    node.connect_to_node(PARTNER_IP, PARTNER_PORT)

    node.conn_mgr.CALLBACKS['message'] = new_message_callback
    node.conn_mgr.CALLBACKS['sending'] = sending_callback

    sleep(0.25)

    print('UP AND RUNNING')
    while True:
        send_data = input()
        if send_data[0].upper() == 'B':
            target_id = node.conn_mgr.BROADCAST_ID
        else:
            target_id = bytes([int(send_data[0])])
        msg_str = send_data[2:]
        if msg_str[:4].upper() == 'FILE':
            path = msg_str[5:]
            with open(path, 'br') as f:
                file_data = f.read()
            node.send_msg(target_id, file_data, 'file')
        else:
            if target_id == node.conn_mgr.BROADCAST_ID:
                node.send_broadcast(msg_str.encode('utf-8'))
            else:
                node.send_msg(target_id, msg_str.encode('utf-8'))


import argparse
import itertools
import os
import pickle
import time
from collections import defaultdict
from dataclasses import dataclass
from os.path import dirname, join
from pathlib import PurePath

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch import nn
from tqdm import tqdm

from src.atlesconfig import config, arg_parse
from src.atlespredict import dbsearch, pepdataset, postprocess, specdataset, specollate_model
from src.atlestrain import model

@dataclass
class PepInfo:
    pep_list: list
    prot_list: list
    pep_mass_list: list

def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = str(config.get_config(key="master_port", section="input"))
    torch.cuda.set_device(rank)
    dist.init_process_group(backend="nccl", world_size=world_size, rank=rank)

# check if preprocessed folder exisits.
# if not do step: 1 - 2
# do step 3 - 4
# Step 1
def classify_peptides(index_path):
    pep_dir = config.get_config(key="pep_dir", section="search")
    min_pep_len = config.get_config(key="min_pep_len", section="ml")
    max_pep_len = config.get_config(key="max_pep_len", section="ml")
    max_clvs = config.get_config(key="max_clvs", section="ml")

    pep_dataset = pepdataset.PeptideDataset(pep_dir, decoy=False)
    pep_classes_path = join(index_path, "peptide_classes")
    os.makedirs(index_path)

    os.mkdir(pep_classes_path)
    os.mkdir(join(index_path, "peptide_embeddings"))
    os.mkdir(join(index_path, "decoy_embeddings"))

    # 1 - classify peptides and write to 144 separate files.
    print("Opening files")
    open_files = {}
    class_offsets = {}
    for length, clv, mod in itertools.product(range(min_pep_len, max_pep_len + 1), range(max_clvs + 1), range(2)):
        file_name = "{}-{}-{}".format(length, clv, mod)
        open_files[file_name] = open(join(pep_classes_path, file_name), "a")
        class_offsets[file_name] = 0

    print("Classifying peptides and writing to files")
    for pep, clv, mod, prot in zip(
        pep_dataset.pep_list,
        pep_dataset.missed_cleavs,
        pep_dataset.pep_modified_list,
        pep_dataset.prot_list,
    ):
        pep_len = sum(map(str.isupper, pep))
        if min_pep_len <= pep_len <= max_pep_len and 0 <= clv <= max_clvs:
            file_name = "{}-{}-{}".format(int(pep_len), int(clv), int(mod))
            if file_name in open_files:
                f = open_files[file_name]
                f.write(">" + prot + "\n")
                f.write(pep + "\n")
                class_offsets[file_name] += 1

    print("Closing files")
    for _, f in open_files.items():
        f.close()

    cum = 0
    for length, clv, mod in itertools.product(range(min_pep_len, max_pep_len + 1), range(max_clvs + 1), range(2)):
        file_name = "{}-{}-{}".format(length, clv, mod)
        offset = class_offsets[file_name]
        class_offsets[file_name] = cum
        cum += offset

    # need to get decoy offset later
    with open(join(index_path, "{}".format("peptide_class_offsets.pkl")), "wb") as f:
        pickle.dump(class_offsets, f)

def get_snap_model(rank):
    model_name = config.get_config(key="model_name", section="search")
    print("Using model: {}".format(model_name))
    snap_model = specollate_model.Net(vocab_size=30, embedding_dim=512, hidden_lstm_dim=512, lstm_layers=2).to(rank)
    snap_model = nn.parallel.DistributedDataParallel(snap_model, device_ids=[rank])
    # snap_model.load_state_dict(torch.load("models/32-embed-2-lstm-SnapLoss2-noch-3k-1k-152.pt")["model_state_dict"])
    # below one has 26975 identified peptides.
    # snap_model.load_state_dict(
    #     torch.load("models/512-embed-2-lstm-SnapLoss-noch-80k-nist-massive-52.pt")["model_state_dict"]
    # )
    # below one has 27.5k peps
    # snap_model.load_state_dict(
    #     torch.load("models/hcd/512-embed-2-lstm-SnapLoss2D-inputCharge-80k-nist-massive-116.pt")["model_state_dict"]
    # )

    snap_model.load_state_dict(torch.load("specollate-model/{}".format(model_name))["model_state_dict"])
    snap_model = snap_model.module
    snap_model.eval()
    print(snap_model)
    return snap_model

# 2 - load each class and process it using specollate, save embeddings for each class separately
def process_peptides(rank, snap_model, index_path):
    pep_dir = config.get_config(key="pep_dir", section="search")
    pep_batch_size = config.get_config(key="pep_batch_size", section="search")
    min_pep_len = config.get_config(key="min_pep_len", section="ml")
    max_pep_len = config.get_config(key="max_pep_len", section="ml")
    max_clvs = config.get_config(key="max_clvs", section="ml")
    class_offsets = {}
    cum = 0
    for length, clv, mod in itertools.product(range(min_pep_len, max_pep_len + 1), range(max_clvs + 1), range(2)):
        file_name = "{}-{}-{}".format(length, clv, mod)
        class_offsets[file_name] = cum
        pep_classes_path = join(index_path, "peptide_classes")
        pep_file_path = join(pep_classes_path, file_name)
        if os.path.exists(pep_file_path):
            if not os.path.getsize(pep_file_path):
                os.remove(pep_file_path)
                continue
            print("Processing file: {}".format(file_name))
            # process peptides
            pep_dataset = pepdataset.PeptideDataset(pep_dir, pep_file_path, decoy=rank == 1)
            cum += len(pep_dataset)
            pep_loader = torch.utils.data.DataLoader(
                dataset=pep_dataset,
                batch_size=pep_batch_size,
                collate_fn=dbsearch.pep_collate,
            )

            print("Processing {}...".format("Peptides" if rank == 0 else "Decoys"))
            e_peps = dbsearch.runSpeCollateModel(pep_loader, snap_model, "peps", rank)
            print("Peptides done!")

            # save embeddings
            embedding_type = "peptide_embeddings" if rank == 0 else "decoy_embeddings"
            embedding_path = join(index_path, embedding_type, file_name)
            print("Saving embeddings at {}".format(embedding_path))
            torch.save(e_peps, embedding_path)
            print("Done \n")

    offset_file_name = "peptide_class_offsets.pkl" if rank == 0 else "decoy_class_offsets.pkl"
    with open(join(index_path, "{}".format(offset_file_name)), "wb") as f:
        pickle.dump(class_offsets, f)

def run_atles(rank, spec_loader):
    model_ = model.Net().to(rank)
    model_ = nn.parallel.DistributedDataParallel(model_, device_ids=[rank])
    # model_.load_state_dict(
    #     torch.load("atles-out/16403437/models/pt-mass-ch-16403437-1toz70vi-472.pt")["model_state_dict"]
    # )
    # model_.load_state_dict(torch.load(
    #     '/lclhome/mtari008/DeepAtles/atles-out/123/models/pt-mass-ch-123-2zgb2ei9-385.pt'
    #     )['model_state_dict'])
    model_.load_state_dict(
        torch.load(
            "/lclhome/mtari008/DeepAtles/atles-out/1382/models/nist-massive-deepnovo-mass-ch-1382-c8mlqbq7-157.pt"
        )["model_state_dict"]
    )
    model_ = model_.module
    model_.eval()
    print(model_)

    lens, cleavs, mods = dbsearch.runAtlesModel(spec_loader, model_, rank)
    pred_cleavs_softmax = torch.log_softmax(cleavs, dim=1)
    _, pred_cleavs = torch.max(pred_cleavs_softmax, dim=1)
    pred_mods_softmax = torch.log_softmax(mods, dim=1)
    _, pred_mods = torch.max(pred_mods_softmax, dim=1)

    return (
        torch.round(lens).type(torch.IntTensor).squeeze().tolist(),
        pred_cleavs.squeeze().tolist(),
        pred_mods.squeeze().tolist(),
    )

def process_spectra(rank, snap_model):
    prep_path = config.get_config(section="search", key="prep_path")
    spec_batch_size = config.get_config(key="spec_batch_size", section="search")
    spec_dataset = specdataset.SpectraDataset(join(prep_path, "specs.pkl"))
    spec_loader = torch.utils.data.DataLoader(
        dataset=spec_dataset,
        batch_size=spec_batch_size,
        collate_fn=dbsearch.spec_collate,
    )

    print("Processing spectra...")
    e_specs = dbsearch.runSpeCollateModel(spec_loader, snap_model, "specs", rank)
    print("Spectra done!")

    atles_start_time = time.time()
    lens, cleavs, mods = run_atles(rank, spec_loader)
    atles_end_time = time.time()
    atles_time = atles_end_time - atles_start_time
    print("Atles time: {}".format(atles_time))
    return e_specs, lens, cleavs, mods, spec_dataset.masses, spec_dataset.charges

# 3 - Loop over spectra classes, load embeddings for peptides, peform db search
def create_spectra_dict(lens, cleavs, mods, e_specs, spec_masses):
    min_pep_len = config.get_config(key="min_pep_len", section="ml")
    max_pep_len = config.get_config(key="max_pep_len", section="ml")
    max_clvs = config.get_config(key="max_clvs", section="ml")
    print("Creating spectra filtered dictionary.")
    spec_filt_dict = defaultdict(list)
    for idx, (l, clv, mod) in enumerate(zip(lens, cleavs, mods)):
        if min_pep_len <= l <= max_pep_len and 0 <= clv <= max_clvs:
            key = "{}-{}-{}".format(int(l), int(clv), int(mod))
            spec_filt_dict[key].append([idx, e_specs[idx], spec_masses[idx]])

    return spec_filt_dict

def search_database(rank, spec_filt_dict, spec_charges, index_path, out_pin_dir):
    pep_dir = config.get_config(key="pep_dir", section="search")
    search_spec_batch_size = config.get_config(key="search_spec_batch_size", section="search")
    min_pep_len = config.get_config(key="min_pep_len", section="ml")
    max_pep_len = config.get_config(key="max_pep_len", section="ml")
    length_filter = config.get_config(key="length_filter", section="filter")
    len_tol_pos = config.get_config(key="len_tol_pos", section="filter") if length_filter else 0
    len_tol_neg = config.get_config(key="len_tol_neg", section="filter") if length_filter else 0
    # dist.barrier()
    # Run database search for each dict item
    unfiltered_time = 0

    print("Running filtered {} database search.".format("target" if rank == 0 else "decoy"))
    for key in spec_filt_dict:
        print("Searching for key {}.".format(key))
        spec_inds = []
        pep_inds = []
        psm_vals = []
        pep_info = PepInfo([], [], [])
        cum = 0
        for tol in range(len_tol_neg, len_tol_pos + 1):
            key_len, key_clv, key_mod = (
                int(key.split("-")[0]),
                int(key.split("-")[1]),
                int(key.split("-")[2]),
            )
            if key_len + tol < min_pep_len or key_len + tol > max_pep_len:
                continue
            file_name = "{}-{}-{}".format(key_len + tol, key_clv, key_mod)
            pep_classes_path = join(index_path, "peptide_classes")
            pep_file_path = join(pep_classes_path, file_name)
            if not os.path.exists(pep_file_path):
                print("Key {} not found in pep_dataset".format(pep_file_path))
                continue
            print("Processing file: {}".format(file_name))
            # process peptides
            pep_dataset = pepdataset.PeptideDataset(pep_dir, pep_file_path, decoy=rank == 1)
            # pep_loader = torch.utils.data.DataLoader(
            #     dataset=pep_dataset, batch_size=pep_batch_size,
            #     collate_fn=dbsearch.pep_collate)
            pep_info.pep_list += pep_dataset.pep_list
            pep_info.prot_list += pep_dataset.prot_list
            pep_info.pep_mass_list += pep_dataset.pep_mass_list

            # load embeddings
            pep_embeddings_path = join(index_path, "peptide_embeddings" if rank == 0 else "decoy_embeddings")
            embedding_file_path = join(pep_embeddings_path, file_name)
            e_peps = torch.load(embedding_file_path)
            # pep_data = [[idx + class_offsets[file_name], e_pep, mass]
            #             for idx, (e_pep, mass) in enumerate(zip(e_peps, pep_dataset.pep_mass_list))]
            pep_data = [
                [idx + cum, e_pep, mass] for idx, (e_pep, mass) in enumerate(zip(e_peps, pep_dataset.pep_mass_list))
            ]
            cum += len(pep_data)

            print("Searching against key {} with {} peptides.".format(file_name, len(pep_dataset.pep_mass_list)))
            spec_subset = spec_filt_dict[key]
            search_loader = torch.utils.data.DataLoader(
                dataset=spec_subset,
                num_workers=0,
                batch_size=search_spec_batch_size,
                shuffle=False,
            )
            pep_dataset = None
            unfiltered_start_time = time.time()
            l_spec_inds, l_pep_inds, l_psm_vals = dbsearch.filtered_parallel_search(search_loader, pep_data, rank)
            unfiltered_time += time.time() - unfiltered_start_time

            if not l_spec_inds:
                continue
            spec_inds.extend(l_spec_inds)
            pep_inds.append(l_pep_inds)
            psm_vals.append(l_psm_vals)

        if not spec_inds:
            continue
        # spec_inds.extend(l_spec_inds)
        # pep_inds.append(l_pep_inds)
        # psm_vals.append(l_psm_vals)

        pep_inds = torch.cat(pep_inds, 0)
        psm_vals = torch.cat(psm_vals, 0)

        print("{} PSMS: {}".format("Target" if rank == 0 else "Decoy", len(pep_inds)))

        # 4 - Write PSMs to pin file
        postprocess.write_to_pin(rank, pep_inds, psm_vals, spec_inds, pep_info, spec_charges, out_pin_dir)

def run_atles_search(rank, world_size, config_path, args_dict):
    setup(rank, world_size)
    config.init_config(config_path)
    arg_parse.process_args_dict(args_dict)
    pep_dir = config.get_config(key="pep_dir", section="search")
    pep_index_name = PurePath(pep_dir).name
    index_path = join(config.get_config(key="index_path", section="search"), pep_index_name, "filtered")

    length_filter = config.get_config(key="length_filter", section="filter")
    len_tol_pos = config.get_config(key="len_tol_pos", section="filter") if length_filter else 0

    spectra = PurePath(config.get_config(key="prep_path", section="search")).name
    filt = "filt{}".format(len_tol_pos)
    out_pin_dir = join(os.getcwd(), "percolator", pep_index_name + "-" + filt + "-" + spectra)

    print(f"Pep dir name {pep_dir}")
    print(f"length filter {length_filter}")
    print(f"len_tol_pos {len_tol_pos}")

    with torch.no_grad():
        print("Running filtered ooc search on {}.".format(pep_index_name))
        snap_model = get_snap_model(rank)
        if not os.path.exists(index_path):
            if rank == 0:
                classify_peptides(index_path)
            dist.barrier()
            process_peptides(rank, snap_model, index_path)

        t_time = time.time()
        e_specs, lens, cleavs, mods, spec_masses, spec_charges = process_spectra(rank, snap_model)
        spec_filt_dict = create_spectra_dict(lens, cleavs, mods, e_specs, spec_masses)
        l_time = time.time()
        search_database(rank, spec_filt_dict, spec_charges, index_path, out_pin_dir)
        print("Search time: {}".format(time.time() - l_time))
        postprocess.post_process_pin_files(rank, out_pin_dir)
        print("Total time: {}".format(time.time() - t_time))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    # Adding optional argument
    parser.add_argument("-c", "--config", help="Path to the config file.")
    parser.add_argument("-p", "--preprocess", help="Preprocess data?", default="True")

    # Read arguments from command line
    input_params = parser.parse_args()

    if input_params.config:
        tqdm.write("config: %s" % input_params.path)
    config.param_path = input_params.config if input_params.config else join((dirname(__file__)), "config.ini")

    num_gpus = torch.cuda.device_count()
    print("Num GPUs: {}".format(num_gpus))
    start_time = time.time()
    mp.spawn(run_atles_search, args=(2,), nprocs=2, join=True)
    # run_atles_search(0, 1)
    print("Total time: {}".format(time.time() - start_time))

    # if all filters disabled, call a different function


from utils import get_all_prime
import math
from tqdm import tqdm

def main():
    n = 800800
    target = n ** n
    log_2 = int(math.log2(target))
    primes = get_all_prime(log_2)
    # print(primes)
    l_t = math.log2(target)
    res = 0
    l = 0
    r = len(primes) - 1
    while l < r:
        print(l, r)
        while primes[r] * math.log2(primes[l]) + primes[l] * math.log2(primes[r]) > l_t:
            r -= 1
        res += (r - l)
        # print(l, r)
        l += 1

    print(res)

if __name__ == "__main__":
    main()


from sys import argv
from string import punctuation
from collections import *
from difflib import SequenceMatcher
from isText import istext

#Check if file names passed to script

if((len(argv)<2)):
    print("Usage: python scrape.py file1.txt file2.txt etc.txt")

topics = []
headers = []

#TODO fix this initialization so that the key words are read from files

#Year
year_keyWords = ['2018','2017','2016']
topics.append(year_keyWords)
headers.append("Year")
#Source
source_keyWords = ['Waters','OmniSpray','Prosolia','Custom','custom','lab-built',"source","ion"]
topics.append(source_keyWords)
headers.append("Source")
#MS
MS_keyWords = ["Xevo","xevo","LTQ","Thermo","Thermofisher","Orbitrap"]
topics.append(MS_keyWords)
headers.append("MS")
#Modality
modality_keyWords = ["linear","ion","trap"]
topics.append(modality_keyWords)
headers.append("Modality")
#Inlet-to-capillary distance
inletToCap_keyWords = ["inlet","capillary","mm","to","inlet-to-capillary"]
topics.append(inletToCap_keyWords)
headers.append("inlet-to-cap (mm)")
#Tip-to-surface distance
tipToSurface_keyWords = ["tip","surface","to","-","tip-to-surface"]
topics.append(tipToSurface_keyWords)
headers.append("tip-to-surface (mm)")
#Solvent
solvent_keyWords = ["solvent","methanol","water","1:1","acetonitrile","leucine","ACN","DMF","dimethylformamide","formic acid"]
topics.append(solvent_keyWords)
headers.append("Solvent")
#Flowrate
flowrate_keyWords = ["flowrate","min"]
topics.append(flowrate_keyWords)
headers.append("Flow uL/min")
#Pressure
pressure_keyWords = ["pressure","psi","PSI","N2"]
topics.append(pressure_keyWords)
headers.append("N2 Pressure (PSI)")
#Ion mode
ionMode_keyWords = ["ion","mode","positive","negative"]
topics.append(ionMode_keyWords)
headers.append("Ion Mode")
#injection
injection_keyWords = ["injection","ms"]
topics.append(injection_keyWords)
headers.append("Injection ms")
#Microscans
microscans_keyWords = ["microscan","microscans"]
topics.append(microscans_keyWords)
headers.append("Microscans")
#Spray voltage
sprayVoltage_keyWords = ["spray","voltage","V","v","kV","volts"]
topics.append(sprayVoltage_keyWords)
headers.append("Spray voltage kV")
#Capillary voltage
capillaryVoltage_keyWords = ["capillary","voltage","V","v","kV","volts","cap"]
topics.append(capillaryVoltage_keyWords)
headers.append("Capillary Voltage V")
#Tube lens
tubeLens_keyWords = ["tube","lens","voltage","V","kV"]
topics.append(tubeLens_keyWords)
headers.append("Tube Lens V")
#Capillary Temperature
capillaryTemperature_keyWords = ["capillary","temperature","C","°C","°"]
topics.append(capillaryTemperature_keyWords)
headers.append("Capillary Temperature °C")
#m/z range
mzRange_keyWords = ["m/z","range","to"]
topics.append(mzRange_keyWords)
headers.append("m/z range")
#Resolution
resolution_keyWords = ["pixel","resolution","um"]
topics.append(resolution_keyWords)
headers.append("Resolution um")
#Resolving Power
resolvingPower_keyWords = ["resolving","power","kV"]
topics.append(resolvingPower_keyWords)
headers.append("Resolving Power kV")
#Length
length_keyWords = ["length","m","mm","cm"]
topics.append(length_keyWords)
headers.append("Length mm")
#Line Seperation
lineSeperation_keyWords = ["line","seperation","mm","cm","step","size"]
topics.append(lineSeperation_keyWords)
headers.append("Line Seperation mm")
#Step
step_keyWords = ["step","x-axis","axis","size"]
topics.append(step_keyWords)
headers.append("Step mm")
#Scan duration
scanDuration_keyWords = ["scan","duration","time","ms"]
topics.append(scanDuration_keyWords)
headers.append("Scan duration ms")
#Velocity
velocity_keyWords = ["velocity","rate","um/sec","um/s","mm/s"]
topics.append(velocity_keyWords)
headers.append("Velocity um/s")
#Substrate
substrate_keyWords = ["substrate","slide","glass"]
topics.append(substrate_keyWords)
headers.append("Substrate")
#Acquisition Software
acquisitionSoftware_keyWords = ["software","acquisition","Xcalibur","data","MassLynx"]
topics.append(acquisitionSoftware_keyWords)
headers.append("Acquisition Software")
#Conversion Software
conversionSoftware_keyWords = ["software","conversion","convert","process","processing","raw","firefly","FireFly"]
topics.append(conversionSoftware_keyWords)
headers.append("Conversion software")
#Visualization Software
visualizationSoftware_keyWords = ["software","visualization","visualize","visual","image","images","processed"]
topics.append(visualizationSoftware_keyWords)
headers.append("Visualization Software")

output_file = "output.txt"

def append(output):
  with open(output_file,"a") as myfile:
    myfile.write(output)

#For each file in the arguments
for one_filename in argv[1:]:
  output_string=""
  #Open the file and read its contents
  print("Text file to import and read:", one_filename)
  print("\nReading file...\n")

  text_file = open(one_filename, 'r')
  all_lines = text_file.readlines()
  text_file.close()

  print("\nFile read finished!")

  topicIndex=0
  #For each topic's keywords
  for keyWords in topics:
    print("Current key words",keyWords)
    print("Current Topic: ",headers[topicIndex])
    matches=[]
    matchness=[]
    topicIndex=topicIndex+1

    # Search through all lines within the current text file
    for line in all_lines:
      matchFrequency=0 #Keeps track of how many key words appears in a line

      # Store how many key words appear in the line
      for key in keyWords:
        # Check to see if the line contains one of the topic's keywords
        if key in line:
          if matchFrequency == 0:
            matches.append(line) # Add the current line to array containing all lines containing keywords
            matchness.append(1)
            matchFrequency = 1
          else:
            matchness[-1]+=1

    if(len(matches)>0):
      # Sort the lines based on how many keywords they contain in ascending order
      matchness,matches = zip(*sorted(zip(matchness,matches)))

      matchness,matches = (list(t) for t in zip(*sorted(zip(matchness,matches))))
      matchness = matchness[::-1]
      matches = matches[::-1]

    # USER PROCESSING SECTION

    exit = True
    i=0
    #Loop through all lines containing keywords for the current topic
    #Loop until the user enters some text
    while(exit and i<len(matches)):
      print(matchness[i],matches[i])
      user_input = input()
      #maybe make this into a function that can call itself
      if not user_input:
        pass# Do nothing. Just go to next string
      else:
        line_index = all_lines.index(matches[i])
        startIndex=line_index-1
        endIndex=line_index+1
        print(user_input)
        # Add more context if the user wants to see more
        while user_input == "COMMAND:MORE":

          #TODO Check if line_index-1 is not out of bounds
          for moreIndex in range(startIndex,endIndex):
            print(all_lines[moreIndex])

          startIndex=startIndex-1
          endIndex=endIndex+1
          user_input = input()

        exit = False
        # ENTER STRING INTO APPROPRIATE SPREADSHEET CELL
        output_string = output_string + "," + user_input
        # MOVE TO NEXT COLUMN
      i+=1
      if(i>=len(matches) and exit):
        print("No more matches. Put in best guess?")
        user_input = input()
        if not user_input:
          output_string = output_string + ",-" # Add - into cell indicating it was not found
        else:
          exit = False
          # ENTER STRING INTO APPROPRIATE SPREADSHEET CELL
          output_string = output_string + "," + user_input
          # MOVE TO NEXT COLUMN

  append(one_filename[:3]+output_string+"\n")


'''
Aum Sri Sai Ram
Implementation of ECCT: Ensemble Consensual Collaborative Training for Facial Expression Recognition with Noisy Annotations
Authors: Darshan Gera and Dr. S. Balasubramanian, SSSIHL
Date: 20-01-2021
Email: darshangera@sssihl.edu.in
'''

# -*- coding:utf-8 -*-
import torch
import torch.nn.functional as F
from torch.autograd import Variable
import torchvision.transforms as transforms
from model.cnn import resModel
import numpy as np
from common.utils import accuracy
import os
from algorithm.loss import ECCTloss

class noisyfer:
    def __init__(self, args, train_dataset, device, input_channel, num_classes):

        # Hyper Parameters
        self.batch_size = args.batch_size
        learning_rate = args.lr

        self.device = device
        self.num_iter_per_epoch = args.num_iter_per_epoch
        self.print_freq = args.print_freq
        self.n_epoch = args.n_epoch
        self.train_dataset = train_dataset
        self.co_lambda_max = args.co_lambda_max
        self.beta = args.beta
        self.num_classes  = args.num_classes
        self.max_epochs = args.n_epoch

        if  args.model_type=="res":
            self.model1 = resModel(args)
            self.model2 = resModel(args)
            self.model3 = resModel(args)

        self.model1.to(device)
        self.model2.to(device)
        self.model3.to(device)

        filter_list = ['module.classifier.weight', 'module.classifier.bias']

        base_parameters_model1 = list(map(lambda x: x[1],list(filter(lambda kv: kv[0] not in filter_list, self.model1.named_parameters()))))
        base_parameters_model2 = list(map(lambda x: x[1],list(filter(lambda kv: kv[0] not in filter_list, self.model2.named_parameters()))))
        base_parameters_model3 = list(map(lambda x: x[1],list(filter(lambda kv: kv[0] not in filter_list, self.model3.named_parameters()))))

        self.optimizer = torch.optim.Adam([{'params': base_parameters_model1}, {'params': list(self.model1.module.classifier.parameters()), 'lr': learning_rate}], lr=1e-3)
        self.optimizer.add_param_group({'params': base_parameters_model2, 'lr':1e-3})
        self.optimizer.add_param_group({'params': list(self.model2.module.classifier.parameters()), 'lr':learning_rate})
        self.optimizer.add_param_group({'params': base_parameters_model3, 'lr':1e-3})
        self.optimizer.add_param_group({'params': list(self.model3.module.classifier.parameters()), 'lr':learning_rate})

        print('\n Initial learning rate is:')
        for param_group in self.optimizer.param_groups:
            print(  param_group['lr'])

        if not args.resume:

           pretrained = torch.load(args.resume)
           pretrained_state_dict1 = pretrained['model_1']
           pretrained_state_dict2 = pretrained['model_2']
           pretrained_state_dict3 = pretrained['model_3']

           model1_state_dict =  self.model1.state_dict()
           model2_state_dict =  self.model2.state_dict()
           model3_state_dict =  self.model3.state_dict()
           loaded_keys = 0
           total_keys = 0
           for key in pretrained_state_dict1:
               if  ((key=='module.fcx.weight')|(key=='module.fcx.bias')):
                   print(key)
                   pass
               else:
                   model1_state_dict[key] = pretrained_state_dict1[key]
                   model2_state_dict[key] = pretrained_state_dict2[key]
                   model3_state_dict[key] = pretrained_state_dict3[key]
                   total_keys+=1
                   if key in model1_state_dict and key in model2_state_dict and  key in model3_state_dict:
                      loaded_keys+=1
           print("Loaded params num:", loaded_keys)
           print("Total params num:", total_keys)
           self.model1.load_state_dict(model1_state_dict)
           self.model2.load_state_dict(model2_state_dict)
           self.model3.load_state_dict(model3_state_dict)

           print('All 3 models loaded from ',args.resume)

        else:
           print('\n No checkpoint found from FEC trained dataset.\n')

        self.loss_fn = ECCTloss  #Dynamic balancing of Suprevision Loss and Consistency Loss

        self.m1_statedict =  self.model1.state_dict()
        self.m2_statedict =  self.model2.state_dict()
        self.m3_statedict =  self.model3.state_dict()
        self.o_statedict = self.optimizer.state_dict()

        self.adjust_lr = args.adjust_lr

    # Evaluate the Model
    def evaluate(self, test_loader):
        print('Evaluating ...')
        self.model1.eval()
        self.model2.eval()
        self.model3.eval()
        correct1 = 0
        total1 = 0
        correct2 = 0
        total2 = 0
        correct3 = 0
        total3 = 0
        correct  = 0
        with torch.no_grad():
            for images, labels, _ in test_loader:
                images = (images).to(self.device)
                logits1 = self.model1(images)
                outputs1 = F.softmax(logits1, dim=1)
                _, pred1 = torch.max(outputs1.data, 1)
                total1 += labels.size(0)
                correct1 += (pred1.cpu() == labels).sum()

                logits2 = self.model2(images)
                outputs2 = F.softmax(logits2, dim=1)
                _, pred2 = torch.max(outputs2.data, 1)
                total2 += labels.size(0)
                correct2 += (pred2.cpu() == labels).sum()

                logits3 = self.model3(images)
                outputs3 = F.softmax(logits3, dim=1)
                _, pred3 = torch.max(outputs3.data, 1)
                total3 += labels.size(0)
                correct3 += (pred3.cpu() == labels).sum()

                avg_output = 0.33*(outputs1.data + outputs2.data+ outputs3.data)
                _, avg_pred = torch.max(avg_output, 1)
                correct += (avg_pred.cpu() == labels).sum()

            acc1 = 100 * float(correct1) / float(total1)
            acc2 = 100 * float(correct2) / float(total2)
            acc3 = 100 * float(correct3) / float(total3)
            acc = 100 * float(correct) / float(total2)
        return acc1, acc2,acc3, acc

    def save_model(self, epoch, acc, noise):
        torch.save({' epoch':  epoch,
                    'model_1': self.m1_statedict,
                    'model_2': self.m2_statedict,
                    'model_3': self.m3_statedict,
                    'optimizer':self.o_statedict,},
                     os.path.join('checkpoints/', "epoch_"+str(epoch)+'_noise_'+noise+"_acc_"+str(acc)[:5]+".pth"))
        print('Models saved '+os.path.join('checkpoints/', "epoch_"+str(epoch)+'_noise_'+noise+"_acc_"+str(acc)[:5]+".pth"))

    # Train the Model
    def train(self, train_loader, epoch):
        print('Training ...')
        self.model1.train()
        self.model2.train()
        self.model3.train()

        if epoch > 0:
           self.adjust_learning_rate(self.optimizer, epoch)

        train_total = 0
        train_correct = 0
        train_total2 = 0
        train_correct2 = 0
        train_total3 = 0
        train_correct3 = 0
        pure_ratio_1_list = []
        pure_ratio_2_list = []

        for i, (images, labels, indexes) in enumerate(train_loader):
            ind = indexes.cpu().numpy().transpose()

            if i > self.num_iter_per_epoch:
                break

            images = images.to(self.device)
            labels = labels.to(self.device)
            # Forward + Backward + Optimize

            logits1 = self.model1(images)
            prec1 = accuracy(logits1, labels, topk=(1,))
            train_total += 1
            train_correct += prec1

            logits2 = self.model2.forward(images)

            prec2 = accuracy(logits2, labels, topk=(1,))

            train_total2 += 1
            train_correct2 += prec2

            logits3 = self.model3(images)

            prec3 = accuracy(logits3, labels, topk=(1,))
            train_total3 += 1
            train_correct3 += prec3

            avg_prec = accuracy(0.33*(logits1+logits2+logits3), labels, topk=(1,))

            loss = self.loss_fn(logits1, logits2, logits3, labels, self.co_lambda_max,  self.beta, epoch, self.n_epoch)

            self.optimizer.zero_grad()
            loss.backward()

            self.optimizer.step()

            if (i + 1) % self.print_freq == 0:
                print(
                    'Epoch [%d/%d], Iter [%d/%d] Training Accuracy1: %.4F, Training Accuracy2: %.4f, Training Accuracy3: %.4f,Avg Accuracy: %.4f, Loss1: %.4f, %%'
                    % (epoch + 1, self.n_epoch, i + 1, len(self.train_dataset) // self.batch_size, prec1, prec2,prec3,avg_prec,  loss.data.item() ))

        train_acc1 = float(train_correct) / float(train_total)
        train_acc2 = float(train_correct2) / float(train_total2)
        train_acc3 = float(train_correct3) / float(train_total2)
        return train_acc1, train_acc2, train_acc3 , avg_prec

    def adjust_learning_rate(self, optimizer, epoch):
        print('\n******************************\n\tAdjusted learning rate: '+str(epoch) +'\n')
        for param_group in optimizer.param_groups:
           param_group['lr'] *= 0.95
           print(param_group['lr'])
        print('******************************')

    # Train the Model
    def train(self, train_loader, epoch):
        print('Training ...')
        self.model1.train()
        self.model2.train()
        self.model3.train()

        if epoch > 0:
           self.adjust_learning_rate(self.optimizer, epoch)

        train_total = 0
        train_correct = 0
        train_total2 = 0
        train_correct2 = 0
        train_total3 = 0
        train_correct3 = 0
        pure_ratio_1_list = []
        pure_ratio_2_list = []

        for i, (images, labels, indexes) in enumerate(train_loader):
            ind = indexes.cpu().numpy().transpose()

            if i > self.num_iter_per_epoch:
                break

            images = images.to(self.device)
            labels = labels.to(self.device)
            # Forward + Backward + Optimize

            logits1 = self.model1(images)
            prec1 = accuracy(logits1, labels, topk=(1,))
            train_total += 1
            train_correct += prec1

            logits2 = self.model2.forward(images)

            prec2 = accuracy(logits2, labels, topk=(1,))

            train_total2 += 1
            train_correct2 += prec2

            logits3 = self.model3(images)

            prec3 = accuracy(logits3, labels, topk=(1,))
            train_total3 += 1
            train_correct3 += prec3

            avg_prec = accuracy(0.33*(logits1+logits2+logits3), labels, topk=(1,))

            loss = self.loss_fn(logits1, logits2, logits3, labels, self.co_lambda_max,  self.beta, epoch, self.n_epoch)

            self.optimizer.zero_grad()
            loss.backward()

            self.optimizer.step()

            if (i + 1) % self.print_freq == 0:
                print(
                    'Epoch [%d/%d], Iter [%d/%d] Training Accuracy1: %.4F, Training Accuracy2: %.4f, Training Accuracy3: %.4f,Avg Accuracy: %.4f, Loss1: %.4f, %%'
                    % (epoch + 1, self.n_epoch, i + 1, len(self.train_dataset) // self.batch_size, prec1, prec2,prec3,avg_prec,  loss.data.item() ))

        train_acc1 = float(train_correct) / float(train_total)
        train_acc2 = float(train_correct2) / float(train_total2)
        train_acc3 = float(train_correct3) / float(train_total2)
        return train_acc1, train_acc2, train_acc3 , avg_prec

    def adjust_learning_rate(self, optimizer, epoch):
        #print('\n******************************\n\tAdjusted learning rate: '+str(epoch) +'\n')
        for param_group in optimizer.param_groups:
           param_group['lr'] *= 0.95
           #print(param_group['lr'])
        #print('******************************')


import copy
import numpy as np
import astropy.units as u
import logging
from scipy.interpolate import interp1d
from astropy.coordinates import SkyCoord
from gammapy.data import Observations
from gammapy.datasets import MapDataset
from gammapy.irf import Background2D
from gammapy.makers import MapDatasetMaker, SafeMaskMaker, FoVBackgroundMaker
from gammapy.maps import WcsNDMap, WcsGeom
from regions import CircleAnnulusSkyRegion

class RadialAcceptanceMapCreator:

    def __init__(self, energy_axis, offset_axis, oversample_map=10, exclude_regions=[],
                 min_run_per_cos_zenith_bin=3, initial_cos_zenith_binning=0.01):
        """
            Create the class for calculating radial acceptance model

            Parameters
            ----------
            energy_axis : gammapy.maps.geom.MapAxis
                The energy axis for the acceptance model
            offset_axis : gammapy.maps.geom.MapAxis
                The offset axis for the acceptance model
            oversample_map : int
                Oversample in number of pixel of the spatial axis used for the calculation
            exclude_regions : list of 'regions.SkyRegion'
                Region with known or putative gamma-ray emission, will be excluded of the calculation of the acceptance map
            min_run_per_cos_zenith_bin : int
                Minimum number of runs per zenith bins
            initial_cos_zenith_binning : float
                Initial bin size for cos zenith binning
        """

        # Store base parameter
        self.energy_axis = energy_axis
        self.offset_axis = offset_axis
        self.oversample_map = oversample_map
        self.exclude_regions = exclude_regions
        self.min_run_per_cos_zenith_bin = min_run_per_cos_zenith_bin
        self.initial_cos_zenith_binning = initial_cos_zenith_binning

        # Calculate map parameter
        self.n_bins_map = self.offset_axis.nbin * self.oversample_map * 2
        self.binsz = self.offset_axis.edges[-1] / (self.n_bins_map / 2)

    def __create_sky_map(self, obs, add_bkg=False):
        """
            Create the sky map used for computation

            Parameters
            ----------
            obs : gammapy.data.observations.Observation
                The observation used to make the sky map
            add_bkg : bool
                If true will also add the background model to the map

            Returns
            -------
            map_obs : gammapy.datasets.MapDataset
            exclusion_mask : gammapy.maps.WcsNDMap
        """

        geom_obs = WcsGeom.create(skydir=obs.pointing_radec, npix=(self.n_bins_map, self.n_bins_map),
                                  binsz=self.binsz, frame="icrs", axes=[self.energy_axis])
        maker = MapDatasetMaker(selection=["counts"])
        if add_bkg:
            maker = MapDatasetMaker(selection=["counts", "background"])

        maker_safe_mask = SafeMaskMaker(methods=["offset-max"], offset_max=np.max(self.offset_axis.edges))

        geom_image = geom_obs.to_image()
        exclusion_mask = ~geom_image.region_mask(self.exclude_regions)

        map_obs = maker.run(MapDataset.create(geom=geom_obs), obs)
        map_obs = maker_safe_mask.run(map_obs, obs)

        return map_obs, exclusion_mask

    def create_radial_acceptance_map(self, observations):
        """
            Calculate a radial acceptance map

            Parameters
            ----------
            observations : gammapy.data.observations.Observations
                The collection of observations used to make the acceptance map

            Returns
            -------
            background : gammapy.irf.background.Background2D
        """

        center_map = SkyCoord(ra=0. * u.deg, dec=0. * u.deg, frame='icrs')

        geom = WcsGeom.create(skydir=center_map, npix=(self.n_bins_map, self.n_bins_map), binsz=self.binsz,
                              frame="icrs", axes=[self.energy_axis])

        count_map_background = WcsNDMap(geom=geom)
        exp_map_background = WcsNDMap(geom=geom, unit=u.s)
        exp_map_background_total = WcsNDMap(geom=geom, unit=u.s)
        livetime = 0. * u.s

        for obs in observations:
            count_map_obs, exclusion_mask = self.__create_sky_map(obs)

            exp_map_obs = MapDataset.create(geom=count_map_obs.geoms['geom'])
            exp_map_obs_total = MapDataset.create(geom=count_map_obs.geoms['geom'])
            exp_map_obs.counts.data = obs.observation_live_time_duration.value
            exp_map_obs_total.counts.data = obs.observation_live_time_duration.value

            for i in range(count_map_obs.counts.data.shape[0]):
                count_map_obs.counts.data[i, :, :] = count_map_obs.counts.data[i, :, :] * exclusion_mask
                exp_map_obs.counts.data[i, :, :] = exp_map_obs.counts.data[i, :, :] * exclusion_mask

            count_map_background.data += count_map_obs.counts.data
            exp_map_background.data += exp_map_obs.counts.data
            exp_map_background_total.data += exp_map_obs_total.counts.data
            livetime += obs.observation_live_time_duration

        data_background = np.zeros((self.energy_axis.nbin, self.offset_axis.nbin)) * u.Unit('s-1 MeV-1 sr-1')
        for i in range(self.offset_axis.nbin):
            selection_region = CircleAnnulusSkyRegion(center=center_map, inner_radius=self.offset_axis.edges[i],
                                                      outer_radius=self.offset_axis.edges[i + 1])
            selection_map = geom.to_image().region_mask([selection_region])
            for j in range(self.energy_axis.nbin):
                value = u.dimensionless_unscaled * np.sum(count_map_background.data[j, :, :] * selection_map)
                value *= np.sum(exp_map_background_total.data[j, :, :] * selection_map) / np.sum(
                    exp_map_background.data[j, :, :] * selection_map)

                value /= (self.energy_axis.edges[j + 1] - self.energy_axis.edges[j])
                value /= 2. * np.pi * \
                         (self.offset_axis.edges[i + 1].to('radian') - self.offset_axis.edges[i].to('radian')) * \
                         self.offset_axis.center[i].to('radian')
                value /= livetime
                data_background[j, i] = value

        background = Background2D(axes=[self.energy_axis, self.offset_axis], data=data_background)

        return background

    def create_radial_acceptance_map_cos_zenith_binned(self, observations):
        """
            Calculate a radial acceptance map using cos zenith binning and interpolation

            Parameters
            ----------
            observations : gammapy.data.observations.Observations
                The collection of observations used to make the acceptance map

            Returns
            -------
            background : dict of gammapy.irf.background.Background2D
                A dict with observation number as keu and a bakground model that could be used as an acceptance model associated at each key

        """

        cos_zenith_bin = np.sort(np.arange(1.0, 0. - self.initial_cos_zenith_binning, -self.initial_cos_zenith_binning))
        cos_zenith_observations = [np.cos(obs.pointing_zen) for obs in observations]
        run_per_bin = np.histogram(cos_zenith_observations, bins=cos_zenith_bin)[0]

        i = 0
        while i < len(run_per_bin):
            if run_per_bin[i] < self.min_run_per_cos_zenith_bin and (i + 1) < len(run_per_bin):
                run_per_bin[i] += run_per_bin[i + 1]
                run_per_bin = np.delete(run_per_bin, i + 1)
                cos_zenith_bin = np.delete(cos_zenith_bin, i + 1)
            elif run_per_bin[i] < self.min_run_per_cos_zenith_bin and (i + 1) == len(run_per_bin) and i > 0:
                run_per_bin[i - 1] += run_per_bin[i]
                run_per_bin = np.delete(run_per_bin, i)
                cos_zenith_bin = np.delete(cos_zenith_bin, i)
                i -= 1
            else:
                i += 1

        binned_observations = []
        for i in range((len(cos_zenith_bin) - 1)):
            binned_observations.append(Observations())
        for obs in observations:
            binned_observations[np.digitize(np.cos(obs.pointing_zen), cos_zenith_bin)-1].append(obs)

        binned_model = [self.create_radial_acceptance_map(binned_obs) for binned_obs in binned_observations]
        bin_center = []
        for i in range(len(binned_observations)):
            bin_center.append(np.mean([np.cos(obs.pointing_zen) for obs in binned_observations[i]]))

        radial_acceptance_map = {}
        if len(binned_model) <= 1:
            logging.warning('Only one zenith bin, zenith interpolation deactivated')
            for obs in observations:
                radial_acceptance_map[obs.obs_id] = binned_model[0]
        else:
            data_cube = np.zeros(tuple([len(binned_model), ] + list(binned_model[0].data.shape))) * binned_model[0].unit
            for i in range(len(binned_model)):
                data_cube[i] = binned_model[i].data * binned_model[i].unit
            interp_func = interp1d(x=np.array(bin_center),
                                   y=np.log10(data_cube.value+np.finfo(np.float64).tiny),
                                   axis=0,
                                   fill_value='extrapolate')
            for obs in observations:
                radial_acceptance_map[obs.obs_id] = Background2D(axes=binned_model[0].axes,
                                                                         data=(10.**interp_func(np.cos(obs.pointing_zen)))*data_cube.unit)

        return radial_acceptance_map

    def create_radial_acceptance_map_per_observation(self, observations, zenith_bin=True):
        """
            Calculate a radial acceptance map with the norm adjusted for each run

            Parameters
            ----------
            observations : gammapy.data.observations.Observations
                The collection of observations used to make the acceptance map
            zenith_bin : bool,
                If true the acceptance maps will be generated using zenith binning and interpolation

            Returns
            -------
            background : dict of gammapy.irf.background.Background2D
                A dict with observation number as keu and a bakground model that could be used as an acceptance model associated at each key
        """

        base_radial_acceptance_map = {}
        if zenith_bin:
            base_radial_acceptance_map = self.create_radial_acceptance_map_cos_zenith_binned(observations)
        else:
            unique_base_radial_acceptance_map = self.create_radial_acceptance_map(observations)
            for obs in observations:
                base_radial_acceptance_map[obs.obs_id] = unique_base_radial_acceptance_map

        radial_acceptance_map = {}
        # Fit norm of the model to the observations
        for obs in observations:
            id_observation = obs.obs_id

            # replace the background model
            modified_observation = copy.deepcopy(obs)
            modified_observation.bkg = base_radial_acceptance_map[id_observation]

            # Fit the background model
            logging.info('Fit to model to run ' + str(id_observation))
            map_obs, exclusion_mask = self.__create_sky_map(modified_observation, add_bkg=True)
            maker_FoV_background = FoVBackgroundMaker(method='fit', exclusion_mask=exclusion_mask)
            map_obs = maker_FoV_background.run(map_obs)

            # Extract the normalisation
            parameters = map_obs.models.to_parameters_table()
            norm_background = parameters[parameters['name'] == 'norm']['value'][0]

            # Apply normalisation to the background model
            radial_acceptance_map[id_observation] = copy.deepcopy(base_radial_acceptance_map[id_observation])
            radial_acceptance_map[id_observation].data = radial_acceptance_map[id_observation].data * norm_background

        return radial_acceptance_map


'''
You are given an array of non-negative integers that represents a two-dimensional elevation map where each element is unit-width wall and the integer is the height. Suppose it will rain and all spots between two walls get filled up.

Compute how many units of water remain trapped on the map in O(N) time and O(1) space.

For example, given the input [2, 1, 2], we can hold 1 unit of water in the middle.

Given the input [3, 0, 1, 3, 0, 5], we can hold 3 units in the first index, 2 in the second, and 3 in the fourth index (we cannot hold 5 since it would run off to the left), so we can trap 8 units of water.
'''
import random

def rainWall(walls):
    def calculateRain(highest, ind1, ind2, step=1):
        secondHighest = 0
        secondHighestInd = 0
        sumOfWalls = 0
        sumOfRain = 0
        for ind in range(ind1, ind2, step):
            sumOfWalls += walls[ind]
            if(walls[ind] > secondHighest):
                sumOfRain += (secondHighest*(abs(ind-secondHighestInd)))
                secondHighest = walls[ind]
                secondHighestInd = ind
        sumOfRain += (secondHighest * (abs(ind2 - secondHighestInd)))
        # return (secondHighest*(abs(ind2-ind1))-sumOfWalls)
        return sumOfRain - sumOfWalls
    highest = walls[0]
    ind = 0
    for i in range(len(walls)):
        if walls[i] > highest:
            highest = walls[i]
            ind = i

    rain = 0
    rain += calculateRain(highest,0, ind)
    rain += calculateRain(highest,len(walls)-1, ind, -1)
    # print(rain)
    return(rain)

# sampleWall = [3, 0, 1, 3, 0, 5]
# print(sampleWall,"\n",rainWall(sampleWall))

# bigWall = [2,5,3,6,9,5,3,9,1,5,6,4]
# print(bigWall,"\n",rainWall(bigWall))

for i in range(0, 3):
    randomWall = [random.randint(0,15) for x in range(random.randint(5,15))]
    print(randomWall,"\n",rainWall(randomWall))

'''
.....X
.....X
XOOXOX
XOOXOX
XOXXOX
301305
SUM = 8

....XOOX....
....XOOX....
....XOOX....
...XXOOXOOX.
.XOXXXOXOXX.
.XOXXXOXOXXX
.XXXXXXXXXXX
XXXXXXXXXXXX
XXXXXXXXXXXX
253695391564
SUM = 18
'''


import json

import jwt
import pytest
import unittest

from citrine.exceptions import (
    NonRetryableException,
    WorkflowConflictException,
    WorkflowNotReadyException,
    RetryableException,
    BadRequest)

from datetime import datetime, timedelta
import pytz
import mock
import requests
import requests_mock
from citrine._session import Session
from citrine.exceptions import UnauthorizedRefreshToken, Unauthorized, NotFound

def refresh_token(expiration: datetime = None) -> dict:
    token = jwt.encode(
        payload={'exp': expiration.timestamp()},
        key='garbage'
    )
    return {'access_token': token.decode('utf-8')}

@pytest.fixture
def session():
    session = Session(
        refresh_token='12345',
        scheme='http',
        host='citrine-testing.fake'
    )
    # Default behavior is to *not* require a refresh - those tests can clear this out
    # As rule of thumb, we should be using freezegun or similar to never rely on the system clock
    # for these scenarios, but I thought this is light enough to postpone that for the time being
    session.access_token_expiration = datetime.utcnow() + timedelta(minutes=3)

    return session

def test_get_refreshes_token(session: Session):
    session.access_token_expiration = datetime.utcnow() - timedelta(minutes=1)
    token_refresh_response = refresh_token(datetime(2019, 3, 14, tzinfo=pytz.utc))

    with requests_mock.Mocker() as m:
        m.post('http://citrine-testing.fake/api/v1/tokens/refresh', json=token_refresh_response)
        m.get('http://citrine-testing.fake/api/v1/foo', json={'foo': 'bar'})

        resp = session.get_resource('/foo')

    assert {'foo': 'bar'} == resp
    assert datetime(2019, 3, 14) == session.access_token_expiration

def test_get_refresh_token_failure(session: Session):
    session.access_token_expiration = datetime.utcnow() - timedelta(minutes=1)

    with requests_mock.Mocker() as m:
        m.post('http://citrine-testing.fake/api/v1/tokens/refresh', status_code=401)

        with pytest.raises(UnauthorizedRefreshToken):
            session.get_resource('/foo')

def test_get_no_refresh(session: Session):
    with requests_mock.Mocker() as m:
        m.get('http://citrine-testing.fake/api/v1/foo', json={'foo': 'bar'})
        resp = session.get_resource('/foo')

    assert {'foo': 'bar'} == resp

def test_get_not_found(session: Session):
    with requests_mock.Mocker() as m:
        m.get('http://citrine-testing.fake/api/v1/foo', status_code=404)
        with pytest.raises(NotFound):
            session.get_resource('/foo')

class SessionTests(unittest.TestCase):
    @mock.patch.object(Session, '_refresh_access_token')
    @mock.patch.object(requests.Session, 'request')
    def test_status_code_409(self, mock_request, _):
        resp = mock.Mock()
        resp.status_code = 409
        mock_request.return_value = resp
        with pytest.raises(NonRetryableException):
            Session().checked_request('method', 'path')
        with pytest.raises(WorkflowConflictException):
            Session().checked_request('method', 'path')

    @mock.patch.object(Session, '_refresh_access_token')
    @mock.patch.object(requests.Session, 'request')
    def test_status_code_425(self, mock_request, _):
        resp = mock.Mock()
        resp.status_code = 425
        mock_request.return_value = resp
        with pytest.raises(RetryableException):
            Session().checked_request('method', 'path')
        with pytest.raises(WorkflowNotReadyException):
            Session().checked_request('method', 'path')

    @mock.patch.object(Session, '_refresh_access_token')
    @mock.patch.object(requests.Session, 'request')
    def test_status_code_400(self, mock_request, _):
        resp = mock.Mock()
        resp.status_code = 400
        resp_json = {
            'code': 400,
            'message': 'a message',
            'validation_errors': [
                {
                    'failure_message': 'you have failed',
                },
            ],
        }
        resp.json = lambda: resp_json
        resp.text = json.dumps(resp_json)
        mock_request.return_value = resp
        with pytest.raises(BadRequest) as einfo:
            Session().checked_request('method', 'path')
        assert einfo.value.api_error.validation_errors[0].failure_message \
            == resp_json['validation_errors'][0]['failure_message']

    @mock.patch.object(Session, '_refresh_access_token')
    @mock.patch.object(requests.Session, 'request')
    def test_status_code_401(self, mock_request, _):
        resp = mock.Mock()
        resp.status_code = 401
        mock_request.return_value = resp
        with pytest.raises(NonRetryableException):
            Session().checked_request('method', 'path')
        with pytest.raises(Unauthorized):
            Session().checked_request('method', 'path')

    @mock.patch.object(Session, '_refresh_access_token')
    @mock.patch.object(requests.Session, 'request')
    def test_status_code_404(self, mock_request, _):
        resp = mock.Mock()
        resp.status_code = 404
        mock_request.return_value = resp
        with pytest.raises(NonRetryableException):
            Session().checked_request('method', 'path')

def test_post_refreshes_token_when_denied(session: Session):
    token_refresh_response = refresh_token(datetime(2019, 3, 14, tzinfo=pytz.utc))

    with requests_mock.Mocker() as m:
        m.post('http://citrine-testing.fake/api/v1/tokens/refresh', json=token_refresh_response)
        m.register_uri('POST', 'http://citrine-testing.fake/api/v1/foo', [
            {'status_code': 401, 'json': {'reason': 'invalid-token'}},
            {'json': {'foo': 'bar'}}
        ])

        resp = session.post_resource('/foo', json={'data': 'hi'})

    assert {'foo': 'bar'} == resp
    assert datetime(2019, 3, 14) == session.access_token_expiration

def test_delete_unauthorized_without_json(session: Session):
    with requests_mock.Mocker() as m:
        m.delete('http://citrine-testing.fake/api/v1/bar/something', status_code=401)

        with pytest.raises(Unauthorized):
            session.delete_resource('/bar/something')

def test_failed_put_with_stacktrace(session: Session):
    with requests_mock.Mocker() as m:
        m.put(
            'http://citrine-testing.fake/api/v1/bad-endpoint',
            status_code=500,
            json={'debug_stacktrace': 'blew up!'}
        )

        with pytest.raises(Exception) as e:
            session.put_resource('/bad-endpoint', json={})

    assert '{"debug_stacktrace": "blew up!"}' == str(e.value)


__author__ = 'Igor Zagorodnov'

from ocelot.cpbd.wake3D_old import *
from ocelot.gui import *

filename='WakeXFEL0.1600.001'
Ps0, q, z0, gamref = astraBeam2particleArray(filename)

#
Ps = Ps0.particles.view()
Np = len(Ps)/6
Ps.shape = (Np, 6)

y0=0
x0=0.1e-3
Ps[:,2]=Ps[:,2]+y0
Ps[:,0]=Ps[:,0]+x0
Ns=500              # wake sampling
NF=20               # smoothing filter order
wakeFile='wake_vert_1m.txt'
THv=LoadWakeTable (wakeFile)
Px,Py,Pz,I00=AddTotalWake (Ps[:,0],Ps[:,2],Ps[:,4],q,THv,Ns,NF)
wakefile='WakeLQD.txt'
w1=np.loadtxt(wakefile,comments='%')
s=-w1[:,0]*1e3+1

P = exact_xxstg_2_xp(Ps, gamref)
sp=P[:,2]*1e6

plt.figure(1)
koef=np.min(Pz)/np.max(I00[:,1])
plt.plot(sp,Pz*1e-6,'.',s,w1[:,1]*1e12/4*1e-6,I00[:,0]*1e6,-I00[:,1]*koef*1e-6)
plt.figure(2)
koef=np.max(Px)/np.max(I00[:,1]);
plt.plot(sp,Px*1e-6,'.',s,w1[:,2]*1e15/2*1e-6*x0,I00[:,0]*1e6,I00[:,1]*koef*1e-6)
plt.show()


#Body mass index (BMI) is a value derived from the mass (weight) and height of a person.
#The BMI is defined as the body mass divided by the square of the body height, and is expressed in units of kg/m2, resulting from mass in kilograms and height in metres.

# height_in and weight_lb are available as a regular lists
# Calculate the BMI: bmi
#np_height_m = np.array(height_in) * 0.0254
#np_weight_kg = np.array(weight_lb) * 0.453592
#bmi = np_weight_kg / np_height_m ** 2
# Create the light array
#light = bmi < 21
#np_light = np.array(light)
# Print out light
#print(light)
# Print out BMIs of all baseball players whose BMI is below 21
#print(bmi[light])

#Improved version
# Import numpy
import numpy as np

# Ask user to input their height and weight
height = float(input("Enter your height in meters: "))
weight = float(input("Enter your weight in kilograms: "))

# Calculate the BMI: bmi
bmi = weight / height ** 2

# Create the light array
light = bmi < 21
heavy = bmi > 25
np_light = np.array(light)

# Print out light
print("Is the BMI below 21? " + str(light))

# Print out the BMI
print("Your BMI is: " + str(round(bmi, 2)))

# Print out BMI if it's below 21
if light:
    print("Your BMI is below 21, which is considered to be light.")
if heavy:
    print("Your BMI is over 24, which is considered to be overweighted aka being a fat ass.")