class Price(object):
    @staticmethod
    def encode(_price_list: list, _price_bit: list):
        if len(_price_bit) != len(_price_list):
            return None

        _total_bit = 0
        _encoded = 0

        for i in range(len(_price_list)):
            _encoded |= _price_list[i] << _total_bit
            print(f"{i} : {_encoded}")
            _total_bit += _price_bit[i]

        return _encoded

    @staticmethod
    def decode(price_bit: list, encoded_value):
        _total_bit = 0
        _decoded_vals = []
        for i in range(len(price_bit)):
            _decoded_val = (encoded_value >> _total_bit) & ((1 << price_bit[i]) - 1)
            print(f"{i}: {_decoded_val}")
            _decoded_vals.append(_decoded_val)
            _total_bit += price_bit[i]

        return _decoded_vals

vetho = int(1.000 * 10 ** 12)
vet = int(1.000 * 10 ** 12)
vebank = int(1.000 * 10 ** 12)
veusd = int(1.000 * 10 ** 6)
print(vetho)
print(vet)
print(vebank)
print(veusd)

_price_list = [vetho, vet, vebank, veusd]
_price_bit = [57 , 57, 57, 37]
encoded = Price.encode(_price_list, _price_bit)
print('encoded price = ', encoded)
decoded = Price.decode(price_bit=_price_bit, encoded_value=encoded)
print("Decoded list : ", decoded)

# decoded = Price.decode(encoded)
# print(decoded)


data = open("C:\\Users\\djdeb\\Desktop\\Random Stuff\\Advent of Code 2020\\day5input.txt")
str = data.read()
puzzle_array = str.splitlines()

highest_ID = 0
IDs = []
for ID in puzzle_array:
	real_ID1 = 0
	real_ID2 = 0
	for i in range(len(ID)):
		#print(len(ID))
		if i < 7:
			if ID[i] == "B":
				real_ID1 += (2 ** (6 - i))
		else:
			if ID[i] == "R":
				real_ID2 += (2 ** (9 - i))
	IDs.append((real_ID1 * 8) + real_ID2)
print(max(IDs))
for i in IDs:
	i = int(i)
for i in range(0, 1000):
	if i not in IDs:
		print(i)


# utils for working with 3d-protein structures
import os
import re
import numpy as np
import torch
import contextlib
from functools import wraps
from einops import rearrange, repeat
# import torch_sparse # only needed for sparse nth_deg adj calculation

# bio
from Bio import SeqIO
import itertools
import string

# sidechainnet

from sidechainnet.utils.sequence import ProteinVocabulary, ONE_TO_THREE_LETTER_MAP
from sidechainnet.utils.measure import GLOBAL_PAD_CHAR
from sidechainnet.structure.build_info import NUM_COORDS_PER_RES, BB_BUILD_INFO, SC_BUILD_INFO
from sidechainnet.structure.StructureBuilder import _get_residue_build_iter

# custom
import mp_nerf

# build vocabulary

VOCAB = ProteinVocabulary()

# constants

import alphafold2_pytorch.constants as constants

# helpers

def exists(val):
    return val is not None

# constants: same as in alphafold2.py

DISTANCE_THRESHOLDS = torch.linspace(2, 20, steps = constants.DISTOGRAM_BUCKETS)

# distance binning function

def get_bucketed_distance_matrix(coords, mask, num_buckets = constants.DISTOGRAM_BUCKETS, ignore_index = -100):
    distances = torch.cdist(coords, coords, p=2)
    boundaries = torch.linspace(2, 20, steps = num_buckets, device = coords.device)
    discretized_distances = torch.bucketize(distances, boundaries[:-1])
    discretized_distances.masked_fill_(~(mask[..., None] & mask[..., None, :]), ignore_index)
    return discretized_distances

# decorators

def set_backend_kwarg(fn):
    @wraps(fn)
    def inner(*args, backend = 'auto', **kwargs):
        if backend == 'auto':
            backend = 'torch' if isinstance(args[0], torch.Tensor) else 'numpy'
        kwargs.update(backend = backend)
        return fn(*args, **kwargs)
    return inner

def expand_dims_to(t, length = 3):
    if length == 0:
        return t
    return t.reshape(*((1,) * length), *t.shape) # will work with both torch and numpy

def expand_arg_dims(dim_len = 3):
    """ pack here for reuse.
        turns input into (B x D x N)
    """
    def outer(fn):
        @wraps(fn)
        def inner(x, y, **kwargs):
            assert len(x.shape) == len(y.shape), "Shapes of A and B must match."
            remaining_len = dim_len - len(x.shape)
            x = expand_dims_to(x, length = remaining_len)
            y = expand_dims_to(y, length = remaining_len)
            return fn(x, y, **kwargs)
        return inner
    return outer

def invoke_torch_or_numpy(torch_fn, numpy_fn):
    def outer(fn):
        @wraps(fn)
        def inner(*args, **kwargs):
            backend = kwargs.pop('backend')
            passed_args = fn(*args, **kwargs)
            passed_args = list(passed_args)
            if isinstance(passed_args[-1], dict):
                passed_kwargs = passed_args.pop()
            else:
                passed_kwargs = {}
            backend_fn = torch_fn if backend == 'torch' else numpy_fn
            return backend_fn(*passed_args, **passed_kwargs)
        return inner
    return outer

@contextlib.contextmanager
def torch_default_dtype(dtype):
    prev_dtype = torch.get_default_dtype()
    torch.set_default_dtype(dtype)
    yield
    torch.set_default_dtype(prev_dtype)

# preprocess data

def get_atom_ids_dict():
    """ Get's a dict mapping each atom to a token. """
    ids = set(["", "N", "CA", "C", "O"])

    for k,v in SC_BUILD_INFO.items():
        for name in v["atom-names"]:
            ids.add(name)

    return {k: i for i,k in enumerate(sorted(ids))}

def make_cloud_mask(aa):
    """ relevent points will be 1. paddings will be 0. """
    mask = np.zeros(constants.NUM_COORDS_PER_RES)
    # early stop if padding token
    if aa == "_":
        return mask
    # get num of atoms in aa
    n_atoms = 4+len( SC_BUILD_INFO[ ONE_TO_THREE_LETTER_MAP[aa] ]["atom-names"] )
    mask[:n_atoms] = 1
    return mask

def make_atom_id_embedds(aa, atom_ids):
    """ Return the tokens for each atom in the aa. """
    mask = np.zeros(constants.NUM_COORDS_PER_RES)
    # early stop if padding token
    if aa == "_":
        return mask
    # get atom id
    atom_list = ["N", "CA", "C", "O"] + SC_BUILD_INFO[ ONE_TO_THREE_LETTER_MAP[aa] ]["atom-names"]
    for i,atom in enumerate(atom_list):
        mask[i] = ATOM_IDS[atom]
    return mask

ATOM_IDS = get_atom_ids_dict()
CUSTOM_INFO = {k: {"cloud_mask": make_cloud_mask(k),
                   "atom_id_embedd": make_atom_id_embedds(k, atom_ids=ATOM_IDS),
                  } for k in "ARNDCQEGHILKMFPSTWYV_"}

# common utils

# parsing to pdb for easier visualization - other example from sidechainnet is:
# https://github.com/jonathanking/sidechainnet/tree/master/sidechainnet/structure

def download_pdb(name, route):
    """ Downloads a PDB entry from the RCSB PDB.
        Inputs:
        * name: str. the PDB entry id. 4 characters, capitalized.
        * route: str. route of the destin file. usually ".pdb" extension
        Output: route of destin file
    """
    os.system(f"curl https://files.rcsb.org/download/{name}.pdb > {route}")
    return route

def clean_pdb(name, route=None, chain_num=None):
    """ Cleans the structure to only leave the important part.
        Inputs:
        * name: str. route of the input .pdb file
        * route: str. route of the output. will overwrite input if not provided
        * chain_num: int. index of chain to select (1-indexed as pdb files)
        Output: route of destin file.
    """
    import mdtraj
    destin = route if route is not None else name
    # read input
    raw_prot = mdtraj.load_pdb(name)
    # iterate over prot and select the specified chains
    idxs = []
    for chain in raw_prot.topology.chains:
        # if arg passed, only select that chain
        if chain_num is not None:
            if chain_num != chain.index:
                continue
        # select indexes of chain
        chain_idxs = raw_prot.topology.select(f"chainid == {str(chain.index)}")
        idxs.extend( chain_idxs.tolist() )
    # sort: topology and xyz selection are ordered
    idxs = sorted(idxs)
    # get new trajectory from the sleected subset of indexes and save
    prot = mdtraj.Trajectory(xyz=raw_prot.xyz[:, idxs],
                             topology=raw_prot.topology.subset(idxs))
    prot.save(destin)
    return destin

def custom2pdb(coords, proteinnet_id, route):
    """ Takes a custom representation and turns into a .pdb file.
        Inputs:
        * coords: array/tensor of shape (3 x N) or (N x 3). in Angstroms.
                  same order as in the proteinnnet is assumed (same as raw pdb file)
        * proteinnet_id: str. proteinnet id format (<class>#<pdb_id>_<chain_number>_<chain_id>)
                         see: https://github.com/aqlaboratory/proteinnet/
        * route: str. destin route.
        Output: tuple of routes: (original, generated) for the structures.
    """
    import mdtraj
    # convert to numpy
    if isinstance(coords, torch.Tensor):
        coords = coords.detach().cpu().numpy()
    # ensure (1, N, 3)
    if coords.shape[1] == 3:
        coords = coords.T
    coords = np.newaxis(coords, axis=0)
    # get pdb id and chain num
    pdb_name, chain_num = proteinnet_id.split("#")[-1].split("_")[:-1]
    pdb_destin = "/".join(route.split("/")[:-1])+"/"+pdb_name+".pdb"
    # download pdb file and select appropiate
    download_pdb(pdb_name, pdb_destin)
    clean_pdb(pdb_destin, chain_num=chain_num)
    # load trajectory scaffold and replace coordinates - assumes same order
    scaffold = mdtraj.load_pdb(pdb_destin)
    scaffold.xyz = coords
    scaffold.save(route)
    return pdb_destin, route

def coords2pdb(seq, coords, cloud_mask, prefix="", name="af2_struct.pdb"):
    """ Turns coordinates into PDB files ready to be visualized.
        Inputs:
        * seq: (L,) tensor of ints (sidechainnet aa-key pairs)
        * coords: (3, N) coords of atoms
        * cloud_mask: (L, C) boolean mask of occupied spaces in scn format
        * prefix: str. directory to save files.
        * name: str. name of destin file (ex: pred1.pdb)
    """
    scaffold = torch.zeros( cloud_mask.shape, 3 )
    scaffold[cloud_mask] = coords.cpu().float()
    # build structures and save
    pred = scn.StructureBuilder( seq, crd=scaffold )
    pred.to_pdb(prefix+name)

# adapted from https://github.com/facebookresearch/esm

def remove_insertions(sequence: str) -> str:
    """ Removes any insertions into the sequence. Needed to load aligned sequences in an MSA. """
    deletekeys = dict.fromkeys(string.ascii_lowercase)
    deletekeys["."] = None
    deletekeys["*"] = None
    translation = str.maketrans(deletekeys)
    return sequence.translate(translation)

def read_msa(filename: str, nseq: int):
    """ Reads the first nseq sequences from an MSA file, automatically removes insertions."""
    return [(record.description, remove_insertions(str(record.seq)))
            for record in itertools.islice(SeqIO.parse(filename, "fasta"), nseq)]

# sidechainnet / MSA / other data utils

def ids_to_embed_input(x):
    """ Returns the amino acid string input for calculating the ESM and MSA transformer embeddings
        Inputs:
        * x: any deeply nested list of integers that correspond with amino acid id
    """
    assert isinstance(x, list), 'input must be a list'
    id2aa = VOCAB._int2char
    out = []

    for el in x:
        if isinstance(el, list):
            out.append(ids_to_embed_input(el))
        elif isinstance(el, int):
            out.append(id2aa[el])
        else:
            raise TypeError('type must be either list or character')

    if all(map(lambda c: isinstance(c, str), out)):
        return (None, ''.join(out))

    return out

def ids_to_prottran_input(x):
    """ Returns the amino acid string input for calculating the ESM and MSA transformer embeddings
        Inputs:
        * x: any deeply nested list of integers that correspond with amino acid id
    """
    assert isinstance(x, list), 'input must be a list'
    id2aa = VOCAB._int2char
    out = []

    for ids in x:
        chars = ' '.join([id2aa[i] for i in ids])
        chars = re.sub(r"[UZOB]", "X", chars)
        out.append(chars)

    return out

def get_prottran_embedd(seq, model, tokenizer, device = None):
    from transformers import pipeline

    fe = pipeline('feature-extraction', model = model, tokenizer = tokenizer, device = (-1 if not exists(device) else device.index))

    max_seq_len = seq.shape[1]
    embedd_inputs = ids_to_prottran_input(seq.cpu().tolist())

    embedding = fe(embedd_inputs)
    embedding = torch.tensor(embedding, device = device)

    return embedding[:, 1:(max_seq_len + 1)]

def get_msa_embedd(msa, embedd_model, batch_converter, device = None):
    """ Returns the MSA_tr embeddings for a protein.
        Inputs:
        * seq: ( (b,) L,) tensor of ints (in sidechainnet int-char convention)
        * embedd_model: MSA_tr model (see train_end2end.py for an example)
        * batch_converter: MSA_tr batch converter (see train_end2end.py for an example)
        Outputs: tensor of (batch, n_seqs, L, embedd_dim)
            * n_seqs: number of sequences in the MSA
            * embedd_dim: number of embedding dimensions. 768 for MSA_Transformer
    """
    # use MSA transformer
    REPR_LAYER_NUM = 12
    device = seq.device
    max_seq_len = msa.shape[-1]
    embedd_inputs = ids_to_embed_input(msa.cpu().tolist())

    msa_batch_labels, msa_batch_strs, msa_batch_tokens = batch_converter(embedd_inputs)
    with torch.no_grad():
        results = embedd_model(msa_batch_tokens.to(device), repr_layers=[REPR_LAYER_NUM], return_contacts=False)
    # index 0 is for start token. so take from 1 one
    token_reps = results["representations"][REPR_LAYER_NUM][..., 1:max_seq_len+1, :]
    return token_reps

def get_esm_embedd(seq, embedd_model, batch_converter, msa_data=None):
    """ Returns the ESM embeddings for a protein.
        Inputs:
        * seq: ( (b,) L,) tensor of ints (in sidechainnet int-char convention)
        * embedd_model: ESM model (see train_end2end.py for an example)
        * batch_converter: ESM batch converter (see train_end2end.py for an example)
        Outputs: tensor of (batch, n_seqs, L, embedd_dim)
            * n_seqs: number of sequences in the MSA. 1 for ESM-1b
            * embedd_dim: number of embedding dimensions. 1280 for ESM-1b
    """
    # use ESM transformer
    device = seq.device
    REPR_LAYER_NUM = 33
    max_seq_len = seq.shape[-1]
    embedd_inputs = ids_to_embed_input(seq.cpu().tolist())

    batch_labels, batch_strs, batch_tokens = batch_converter(embedd_inputs)
    with torch.no_grad():
        results = embedd_model(batch_tokens.to(device), repr_layers=[REPR_LAYER_NUM], return_contacts=False)
    # index 0 is for start token. so take from 1 one
    token_reps = results["representations"][REPR_LAYER_NUM][..., 1:max_seq_len+1, :].unsqueeze(dim=1)
    return token_reps

def get_t5_embedd(seq, tokenizer, encoder, msa_data=None, device=None):
    """ Returns the ProtT5-XL-U50 embeddings for a protein.
        Inputs:
        * seq: ( (b,) L,) tensor of ints (in sidechainnet int-char convention)
        * tokenizer:  tokenizer model: T5Tokenizer
        * encoder: encoder model: T5EncoderModel
                 ex: from transformers import T5EncoderModel, T5Tokenizer
                     model_name = "Rostlab/prot_t5_xl_uniref50"
                     tokenizer = T5Tokenizer.from_pretrained(model_name, do_lower_case=False )
                     model = T5EncoderModel.from_pretrained(model_name)
                     # prepare model
                     model = model.to(device)
                     model = model.eval()
                     if torch.cuda.is_available():
                         model = model.half()
        Outputs: tensor of (batch, n_seqs, L, embedd_dim)
            * n_seqs: number of sequences in the MSA. 1 for T5 models
            * embedd_dim: number of embedding dimensions. 1024 for T5 models
    """
    # get params and prepare
    device = seq.device if device is None else device
    embedd_inputs = ids_to_prottran_input(seq.cpu().tolist())

    # embedd - https://huggingface.co/Rostlab/prot_t5_xl_uniref50
    inputs_embedding = []
    shift_left, shift_right = 0, -1
    ids = tokenizer.batch_encode_plus(embedd_inputs, add_special_tokens=True,
                                                     padding=True,
                                                     return_tensors="pt")
    with torch.no_grad():
        embedding = encoder(input_ids=torch.tensor(ids['input_ids']).to(device),
                            attention_mask=torch.tensor(ids["attention_mask"]).to(device))
    # return (batch, seq_len, embedd_dim)
    token_reps = embedding.last_hidden_state[:, shift_left:shift_right].to(device)
    token_reps = expand_dims_to(token_reps, 4-len(token_reps.shape))
    return token_reps.float()

def get_all_protein_ids(dataloader, verbose=False):
    """ Given a sidechainnet dataloader for a CASP version,
        Returns all the ids belonging to proteins.
        Inputs:
        * dataloader: a sidechainnet dataloader for a CASP version
        Outputs: a set containing the ids for all protein entries.
    """
    # store ids here
    ids = set([])
    # iterate for all batches
    for i,batch in tqdm(enumerate(dataloaders['train'])):
        # for breaking from 2 loops at once
        try:
            for i in range(batch.int_seqs.shape[0]):
                # check if all fragments are : 4_LETTER_PDB + NUM + CHAIN
                max_len_10 = len(batch.pids[i]) < 10
                fragments  = [len(x) <= 4 for x in batch.pids[i].split("_")]
                fragments_under_4 = sum(fragments) == len(fragments) # AND CONDITION
                # record id
                if max_len_10 and fragments_under_4:
                    ids.add(batch.pids[i])
                else:
                    if verbose:
                        print("skip:", batch.pids[i], "under 4", fragments)
        except StopIteration:
            break
    # returns set of ids
    return ids

def scn_cloud_mask(scn_seq, boolean=True, coords=None):
    """ Gets the boolean mask atom positions (not all aas have same atoms).
        Inputs:
        * scn_seq: (batch, length) sequence as provided by Sidechainnet package
        * boolean: whether to return as array of idxs or boolean values
        * coords: optional .(batch, lc, 3). sidechainnet coords.
                  returns the true mask (solves potential atoms that might not be provided)
        Outputs: (batch, length, NUM_COORDS_PER_RES) boolean mask
    """

    scn_seq = expand_dims_to(scn_seq, 2 - len(scn_seq.shape))
    # early check for coords mask
    if coords is not None:
        batch_mask = ( rearrange(coords, '... (l c) d -> ... l c d', c=constants.NUM_COORDS_PER_RES) == 0 ).sum(dim=-1) < coords.shape[-1]
        if boolean:
            return batch_mask.bool()
        else:
            return batch_mask.nonzero()

    # do loop in cpu
    device = scn_seq.device
    batch_mask = []
    scn_seq = scn_seq.cpu().tolist()
    for i, seq in enumerate(scn_seq):
        # get masks for each prot (points for each aa)
        batch_mask.append( torch.tensor([CUSTOM_INFO[VOCAB._int2char[aa]]['cloud_mask'] \
                                         for aa in seq]).bool().to(device) )
    # concat in last dim
    batch_mask = torch.stack(batch_mask, dim=0)
    # return mask (boolean or indexes)
    if boolean:
        return batch_mask.bool()
    else:
        return batch_mask.nonzero()

def scn_backbone_mask(scn_seq, boolean=True, n_aa=3):
    """ Gets the boolean mask for N and CA positions.
        Inputs:
        * scn_seq: sequence(s) as provided by Sidechainnet package (int tensor/s)
        * n_aa: number of atoms in a backbone. (may include cbeta as 4th pos)
        * bool: whether to return as array of idxs or boolean values
        Outputs: (N_mask, CA_mask, C_mask)
    """
    wrapper = torch.zeros(*scn_seq.shape, n_aa).to(scn_seq.device)
    # N is the first atom in every AA. CA is the 2nd.
    wrapper[..., 0] = 1
    wrapper[..., 1] = 2
    wrapper[..., 2] = 3
    wrapper = rearrange(wrapper, '... l c -> ... (l c)')
    # find idxs
    N_mask  = wrapper == 1
    CA_mask = wrapper == 2
    C_mask  = wrapper == 3
    if boolean:
        return N_mask, CA_mask, C_mask
    return torch.nonzero(N_mask), torch.nonzero(CA_mask), torch.nonzero(C_mask)

def scn_atom_embedd(scn_seq):
    """ Returns the token for each atom in the aa.
        Inputs:
        * scn_seq: sequence(s) as provided by Sidechainnet package (int tensor/s)
    """
    device = scn_seq.device
    batch_tokens = []
    # do loop in cpu
    scn_seq = scn_seq.cpu().tolist()
    for i,seq in enumerate(scn_seq):
        batch_tokens.append( torch.tensor([CUSTOM_INFO[VOCAB.int2char(aa)]["atom_id_embedd"] \
                                           for aa in seq]) )
    batch_tokens = torch.stack(batch_tokens, dim=0).long().to(device)
    return batch_tokens

def mat_input_to_masked(x, x_mask=None, edges_mat=None, edges=None,
                          edge_mask=None, edge_attr_mat=None,
                          edge_attr=None):
    """ Turns the padded input and edges + mask into the
        non-padded inputs and edges.
        At least one of (edges_mat, edges) must be provided.
        The same format for edges and edge_attr must be provided
        (either adj matrix form or flattened form).
        Inputs:
        * x: ((batch), N, D) a tensor of N nodes and D dims for each one
        * x_mask: ((batch), N,) boolean mask for x
        * edges: (2, E) optional. indices of the corresponding adjancecy matrix.
        * edges_mat: ((batch), N, N) optional. adjacency matrix for x
        * edge_mask: optional. boolean mask of the same shape of either "edge_mat" or "edges".
        * edge_attr: (E, D_edge) optional. edge attributes of D_edge dims.
        * edge_attr_mat: ((batch), N, N) optional. adjacency matrix with features
        Outputs:
        * x: (N_, D) the masked node features
        * edge_index: (2, E_) the masked x-indices for the edges
        * edge_attr: (E_, D_edge) the masked edge attributes
        * batch: (N_,) the corresponding index in the batch for each node
    """
    # collapse batch dimension
    if len(x.shape) == 3:
        batch_dim = x.shape[1]
        # collapse for x and its mask
        x = rearrange(x, 'b n d ... -> (b n) d ...')
        if x_mask is not None:
            x_mask = rearrange(x_mask, 'b n ... -> (b n) ...')
        else:
            x_mask = torch.ones_like(x[..., 0]).bool()

        # collapse for edge indexes and attributes if needed
        if edges_mat is not None and edges is None:
            edges = torch.nonzero(edges_mat, as_tuple=False).t()
            edges = edges[1:] + edges[:1]*batch_dim
        # get the batch identifier for each node
        batch = (torch.arange(x.shape[0], device=x.device) // batch_dim)[x_mask]
    else:
        # edges to indices format
        if edges_mat is not None and edges is None:
            edges = torch.nonzero(edges_mat, as_tuple=False).t()
        # get the batch identifier for each node
        batch = torch.zeros(x.shape[0], device=x.device).to(x.device)

    # adapt edge attrs if provided
    if edge_attr_mat is not None and edge_attr is None:
            edge_attr = edge_attr[edges_mat.bool()]
    # gen edge_mask if not provided
    if edge_mask is None:
        edge_mask = torch.ones_like(edges[-1]).bool()

    # begin applying masks
    x = x[x_mask]
    # process edge indexes: get square mat and remove all non-coding atoms
    max_num = edges.max().item()+1
    wrapper = torch.zeros(max_num, max_num).to(x.device)
    wrapper[edges[0][edge_mask], edges[1][edge_mask]] = 1
    wrapper = wrapper[x_mask, :][:, x_mask]
    edge_index = torch.nonzero(wrapper, as_tuple=False).t()
    # process edge attr
    edge_attr = edge_attr[edge_mask] if edge_attr is not None else None

    return x, edge_index, edge_attr, batch

def nth_deg_adjacency(adj_mat, n=1, sparse=False):
    """ Calculates the n-th degree adjacency matrix.
        Performs mm of adj_mat and adds the newly added.
        Default is dense. Mods for sparse version are done when needed.
        Inputs:
        * adj_mat: (N, N) adjacency tensor
        * n: int. degree of the output adjacency
        * sparse: bool. whether to use torch-sparse module
        Outputs:
        * edge_idxs: ij positions of the adjacency matrix
        * edge_attrs: degree of connectivity (1 for neighs, 2 for neighs^2, ... )
    """
    adj_mat = adj_mat.float()
    attr_mat = torch.zeros_like(adj_mat)
    new_adj_mat = adj_mat.clone()

    for i in range(n):
        if i == 0:
            attr_mat += adj_mat
            continue

        if i == 1 and sparse:
            idxs = adj_mat.nonzero().t()
            vals = adj_mat[idxs[0], idxs[1]]
            new_idxs = idxs.clone()
            new_vals = vals.clone()
            m, k, n = 3 * [adj_mat.shape[0]] # (m, n) * (n, k) , but adj_mats are squared: m=n=k

        if sparse:
            new_idxs, new_vals = torch_sparse.spspmm(new_idxs, new_vals, idxs, vals, m=m, k=k, n=n)
            new_vals = new_vals.bool().float()
            # fill by indexes bc it's faster in sparse mode - will need an intersection function
            previous = attr_mat[new_idxs[0], new_idxs[1]].bool().float()
            attr_mat[new_idxs[0], new_idxs[1]] = (1 - previous)*(i+1)
        else:
            new_adj_mat = (new_adj_mat @ adj_mat).bool().float()
            attr_mat.masked_fill( (new_adj_mat - attr_mat.bool().float()).bool(), i+1 )

    return new_adj_mat, attr_mat

def prot_covalent_bond(seqs, adj_degree=1, cloud_mask=None, mat=True, sparse=False):
    """ Returns the idxs of covalent bonds for a protein.
        Inputs
        * seq: (b, n) torch long.
        * adj_degree: int. adjacency degree
        * cloud_mask: mask selecting the present atoms.
        * mat: whether to return as indexes of only atoms (PyG version)
               or matrices of masked atoms (for batched training).
               for indexes, only 1 seq is supported.
        * sparse: bool. whether to use torch_sparse for adj_mat calc
        Outputs: edge_idxs, edge_types (degree of adjacency).
    """
    device = seqs.device
    # set up container adj_mat (will get trimmed - less than 14)
    next_aa = NUM_COORDS_PER_RES
    adj_mat = torch.zeros(seqs.shape[0], *[seqs.shape[1]*NUM_COORDS_PER_RES]*2)
    # not needed to device since it's only for indices
    seq_list = seqs.cpu().tolist()
    for s,seq in enumerate(seq_list):
        next_idx = 0
        for i,idx in enumerate(seq):
            aa_bonds = constants.AA_DATA[VOCAB._int2char[idx]]['bonds']
            # if no edges -> padding token -> finish bond creation for this seq
            if len(aa_bonds) == 0:
                break
            # correct next position. for indexes functionality
            next_aa = max(aa_bonds, key=lambda x: max(x))[-1]
            # offset by pos in chain ( intra-aa bonds + with next aa )
            bonds = next_idx + torch.tensor( aa_bonds + [[2, next_aa]] ).t()
            next_idx += next_aa
            # delete link with next if final AA in seq
            if i == seqs.shape[1] - 1:
                bonds = bonds[:, :-1]
            # modify adj mat
            adj_mat[s, bonds[0], bonds[1]] = 1
        # convert to undirected
        adj_mat[s] = adj_mat[s] + adj_mat[s].t()
        # do N_th degree adjacency
        adj_mat, attr_mat = nth_deg_adjacency(adj_mat, n=adj_degree, sparse=sparse)

    if mat:
        # return the full matrix/tensor
        return attr_mat.bool().to(seqs.device), attr_mat.to(device)
    else:
        edge_idxs = attr_mat[0].nonzero().t().long()
        edge_types = attr_mat[0, edge_idxs[0], edge_idxs[1]]
        return edge_idxs.to(seqs.device), edge_types.to(seqs.device)

def sidechain_container(seqs, backbones, atom_mask, cloud_mask=None, padding_tok=20):
    """ Gets a backbone of the protein, returns the whole coordinates
        with sidechains (same format as sidechainnet). Keeps differentiability.
        Inputs:
        * seqs: (batch, L) either tensor or list
        * backbones: (batch, L*n_aa, 3): assume batch=1 (could be extended (?not tested)).
                     Coords for (N-term, C-alpha, C-term, (c_beta)) of every aa.
        * atom_mask: (14,). int or bool tensor specifying which atoms are passed.
        * cloud_mask: (batch, l, c). optional. cloud mask from scn_cloud_mask`.
                      sets point outside of mask to 0. if passed, else c_alpha
        * padding: int. padding token. same as in sidechainnet: 20
        Outputs: whole coordinates of shape (batch, L, 14, 3)
    """
    atom_mask = atom_mask.bool().cpu().detach()
    cum_atom_mask = atom_mask.cumsum(dim=-1).tolist()

    device = backbones.device
    batch, length = backbones.shape[0], backbones.shape[1] // cum_atom_mask[-1]
    predicted  = rearrange(backbones, 'b (l back) d -> b l back d', l=length)

    # early check if whole chain is already pred
    if cum_atom_mask[-1] == 14:
        return predicted

    # build scaffold from (N, CA, C, CB) - do in cpu
    new_coords = torch.zeros(batch, length, constants.NUM_COORDS_PER_RES, 3)
    predicted  = predicted.cpu() if predicted.is_cuda else predicted

    # fill atoms if they have been passed
    for i,atom in enumerate(atom_mask.tolist()):
        if atom:
            new_coords[:, :, i] = predicted[:, :, cum_atom_mask[i]-1]

    # generate sidechain if not passed
    for s,seq in enumerate(seqs):
        # format seq accordingly
        if isinstance(seq, torch.Tensor):
            padding = (seq == padding_tok).sum().item()
            seq_str = ''.join([VOCAB._int2char[aa] for aa in seq.cpu().numpy()[:-padding or None]])
        elif isinstance(seq, str):
            padding = 0
            seq_str = seq
        # get scaffolds - will overwrite oxygen since its position is fully determined by N-C-CA
        scaffolds = mp_nerf.proteins.build_scaffolds_from_scn_angles(seq_str, angles=None, device="cpu")
        coords, _ = mp_nerf.proteins.sidechain_fold(wrapper = new_coords[s, :-padding or None].detach(),
                                                    **scaffolds, c_beta = cum_atom_mask[4]==5)
        # add detached scn
        for i,atom in enumerate(atom_mask.tolist()):
            if not atom:
                new_coords[:, :-padding or None, i] = coords[:, i]

    new_coords = new_coords.to(device)
    if cloud_mask is not None:
        new_coords[torch.logical_not(cloud_mask)] = 0.

    # replace any nan-s with previous point location (or N if pos is 13th of AA)
    nan_mask = list(torch.nonzero(new_coords!=new_coords, as_tuple=True))
    new_coords[nan_mask[0], nan_mask[1], nan_mask[2]] = new_coords[nan_mask[0],
                                                                   nan_mask[1],
                                                                   (nan_mask[-2]+1) % new_coords.shape[-1]]
    return new_coords.to(device)

# distance utils (distogram to dist mat + masking)

def center_distogram_torch(distogram, bins=DISTANCE_THRESHOLDS, min_t=1., center="mean", wide="std"):
    """ Returns the central estimate of a distogram. Median for now.
        Inputs:
        * distogram: (batch, N, N, B) where B is the number of buckets.
        * bins: (B,) containing the cutoffs for the different buckets
        * min_t: float. lower bound for distances.
        Outputs:
        * central: (batch, N, N)
        * dispersion: (batch, N, N)
        * weights: (batch, N, N)
    """
    shape, device = distogram.shape, distogram.device
    # threshold to weights and find mean value of each bin
    n_bins = ( bins - 0.5 * (bins[2] - bins[1]) ).to(device)
    n_bins[0]  = 1.5
    n_bins[-1] = 1.33*bins[-1] # above last threshold is ignored
    max_bin_allowed = torch.tensor(n_bins.shape[0]-1).to(device).long()
    # calculate measures of centrality and dispersion -
    magnitudes = distogram.sum(dim=-1)
    if center == "median":
        cum_dist = torch.cumsum(distogram, dim=-1)
        medium   = 0.5 * cum_dist[..., -1:]
        central  = torch.searchsorted(cum_dist, medium).squeeze()
        central  = n_bins[ torch.min(central, max_bin_allowed) ]
    elif center == "mean":
        central  = (distogram * n_bins).sum(dim=-1) / magnitudes
    # create mask for last class - (IGNORE_INDEX)
    mask = (central <= bins[-2].item()).float()
    # mask diagonal to 0 dist - don't do masked filling to avoid inplace errors
    diag_idxs = np.arange(shape[-2])
    central   = expand_dims_to(central, 3 - len(central.shape))
    central[:, diag_idxs, diag_idxs]  *= 0.
    # provide weights
    if wide == "var":
        dispersion = (distogram * (n_bins - central.unsqueeze(-1))**2).sum(dim=-1) / magnitudes
    elif wide == "std":
        dispersion = ((distogram * (n_bins - central.unsqueeze(-1))**2).sum(dim=-1) / magnitudes).sqrt()
    else:
        dispersion = torch.zeros_like(central, device=device)
    # rescale to 0-1. lower std / var  --> weight=1. set potential nan's to 0
    weights = mask / (1+dispersion)
    weights[weights != weights] *= 0.
    weights[:, diag_idxs, diag_idxs] *= 0.
    return central, weights

# distance matrix to 3d coords: https://github.com/scikit-learn/scikit-learn/blob/42aff4e2e/sklearn/manifold/_mds.py#L279

def mds_torch(pre_dist_mat, weights=None, iters=10, tol=1e-5, eigen=False, verbose=2):
    """ Gets distance matrix. Outputs 3d. See below for wrapper.
        Assumes (for now) distogram is (N x N) and symmetric
        Outs:
        * best_3d_coords: (batch x 3 x N)
        * historic_stresses: (batch x steps)
    """
    device, dtype = pre_dist_mat.device, pre_dist_mat.type()
    # ensure batched MDS
    pre_dist_mat = expand_dims_to(pre_dist_mat, length = ( 3 - len(pre_dist_mat.shape) ))
    # start
    batch, N, _ = pre_dist_mat.shape
    diag_idxs = np.arange(N)
    his = [torch.tensor([np.inf]*batch, device=device)]

    # initialize by eigendecomposition: https://www.lptmc.jussieu.fr/user/lesne/bioinformatics.pdf
    # follow : https://www.biorxiv.org/content/10.1101/2020.11.27.401232v1.full.pdf
    D = pre_dist_mat**2
    M =  0.5 * (D[:, :1, :] + D[:, :, :1] - D)
    # do loop svd bc it's faster: (2-3x in CPU and 1-2x in GPU)
    # https://discuss.pytorch.org/t/batched-svd-lowrank-being-much-slower-than-loop-implementation-both-cpu-and-gpu/119336
    svds = [torch.svd_lowrank(mi) for mi in M]
    u = torch.stack([svd[0] for svd in svds], dim=0)
    s = torch.stack([svd[1] for svd in svds], dim=0)
    v = torch.stack([svd[2] for svd in svds], dim=0)
    best_3d_coords = torch.bmm(u, torch.diag_embed(s).abs().sqrt())[..., :3]

    # only eigen - way faster but not weights
    if weights is None and eigen==True:
        return torch.transpose( best_3d_coords, -1, -2), torch.zeros_like(torch.stack(his, dim=0))
    elif eigen==True:
        if verbose:
            print("Can't use eigen flag if weights are active. Fallback to iterative")

    # continue the iterative way
    if weights is None:
        weights = torch.ones_like(pre_dist_mat)

    # iterative updates:
    for i in range(iters):
        # compute distance matrix of coords and stress
        best_3d_coords = best_3d_coords.contiguous()
        dist_mat = torch.cdist(best_3d_coords, best_3d_coords, p=2).clone()

        stress   = ( weights * (dist_mat - pre_dist_mat)**2 ).sum(dim=(-1,-2)) * 0.5
        # perturb - update X using the Guttman transform - sklearn-like
        dist_mat[ dist_mat <= 0 ] += 1e-7
        ratio = weights * (pre_dist_mat / dist_mat)
        B = -ratio
        B[:, diag_idxs, diag_idxs] += ratio.sum(dim=-1)

        # update
        coords = (1. / N * torch.matmul(B, best_3d_coords))
        dis = torch.norm(coords, dim=(-1, -2))

        if verbose >= 2:
            print('it: %d, stress %s' % (i, stress))
        # update metrics if relative improvement above tolerance
        if (his[-1] - stress / dis).mean() <= tol:
            if verbose:
                print('breaking at iteration %d with stress %s' % (i,
                                                                   stress / dis))
            break

        best_3d_coords = coords
        his.append( stress / dis )

    return torch.transpose(best_3d_coords, -1,-2), torch.stack(his, dim=0)

def mds_numpy(pre_dist_mat, weights=None, iters=10, tol=1e-5, eigen=False, verbose=2):
    """ Gets distance matrix. Outputs 3d. See below for wrapper.
        Assumes (for now) distrogram is (N x N) and symmetric
        Out:
        * best_3d_coords: (3 x N)
        * historic_stress
    """
    if weights is None:
        weights = np.ones_like(pre_dist_mat)

    # ensure batched MDS
    pre_dist_mat = expand_dims_to(pre_dist_mat, length = ( 3 - len(pre_dist_mat.shape) ))
    # start
    batch, N, _ = pre_dist_mat.shape
    his = [np.inf]
    # init random coords
    best_stress = np.inf * np.ones(batch)
    best_3d_coords = 2*np.random.rand(batch, 3, N) - 1
    # iterative updates:
    for i in range(iters):
        # compute distance matrix of coords and stress
        dist_mat = np.linalg.norm(best_3d_coords[:, :, :, None] - best_3d_coords[:, :, None, :], axis=-3)
        stress   = (( weights * (dist_mat - pre_dist_mat) )**2).sum(axis=(-1, -2)) * 0.5
        # perturb - update X using the Guttman transform - sklearn-like
        dist_mat[dist_mat == 0] = 1e-7
        ratio = weights * (pre_dist_mat / dist_mat)
        B = -ratio
        B[:, np.arange(N), np.arange(N)] += ratio.sum(axis=-1)
        # update - double transpose. TODO: consider fix
        coords = (1. / N * np.matmul(best_3d_coords, B))
        dis = np.linalg.norm(coords, axis=(-1, -2))
        if verbose >= 2:
            print('it: %d, stress %s' % (i, stress))
        # update metrics if relative improvement above tolerance
        if (best_stress - stress / dis).mean() <= tol:
            if verbose:
                print('breaking at iteration %d with stress %s' % (i,
                                                                   stress / dis))
            break

        best_3d_coords = coords
        best_stress = stress / dis
        his.append(best_stress)

    return best_3d_coords, np.array(his)

def get_dihedral_torch(c1, c2, c3, c4):
    """ Returns the dihedral angle in radians.
        Will use atan2 formula from:
        https://en.wikipedia.org/wiki/Dihedral_angle#In_polymer_physics
        Can't use torch.dot bc it does not broadcast
        Inputs:
        * c1: (batch, 3) or (3,)
        * c1: (batch, 3) or (3,)
        * c1: (batch, 3) or (3,)
        * c1: (batch, 3) or (3,)
    """
    u1 = c2 - c1
    u2 = c3 - c2
    u3 = c4 - c3

    return torch.atan2( ( (torch.norm(u2, dim=-1, keepdim=True) * u1) * torch.cross(u2,u3, dim=-1) ).sum(dim=-1) ,
                        (  torch.cross(u1,u2, dim=-1) * torch.cross(u2, u3, dim=-1) ).sum(dim=-1) )

def get_dihedral_numpy(c1, c2, c3, c4):
    """ Returns the dihedral angle in radians.
        Will use atan2 formula from:
        https://en.wikipedia.org/wiki/Dihedral_angle#In_polymer_physics
        Inputs:
        * c1: (batch, 3) or (3,)
        * c1: (batch, 3) or (3,)
        * c1: (batch, 3) or (3,)
        * c1: (batch, 3) or (3,)
    """
    u1 = c2 - c1
    u2 = c3 - c2
    u3 = c4 - c3

    return np.arctan2( ( (np.linalg.norm(u2, axis=-1, keepdims=True) * u1) * np.cross(u2,u3, axis=-1)).sum(axis=-1),
                       ( np.cross(u1,u2, axis=-1) * np.cross(u2, u3, axis=-1) ).sum(axis=-1) )

def calc_phis_torch(pred_coords, N_mask, CA_mask, C_mask=None,
                    prop=True, verbose=0):
    """ Filters mirrors selecting the 1 with most N of negative phis.
        Used as part of the MDScaling wrapper if arg is passed. See below.
        Angle Phi between planes: (Cterm{-1}, N, Ca{0}) and (N{0}, Ca{+1}, Cterm{+1})
        Inputs:
        * pred_coords: (batch, 3, N) predicted coordinates
        * N_mask: (batch, N) boolean mask for N-term positions
        * CA_mask: (batch, N) boolean mask for C-alpha positions
        * C_mask: (batch, N) or None. boolean mask for C-alpha positions or
                    automatically calculate from N_mask and CA_mask if None.
        * prop: bool. whether to return as a proportion of negative phis.
        * verbose: bool. verbosity level
        Output: (batch, N) containing the phi angles or (batch,) containing
                the proportions.
        Note: use [0] since all prots in batch have same backbone
    """
    # detach gradients for angle calculation - mirror selection
    pred_coords_ = torch.transpose(pred_coords.detach(), -1 , -2).cpu()
    # ensure dims
    N_mask = expand_dims_to( N_mask, 2-len(N_mask.shape) )
    CA_mask = expand_dims_to( CA_mask, 2-len(CA_mask.shape) )
    if C_mask is not None:
        C_mask = expand_dims_to( C_mask, 2-len(C_mask.shape) )
    else:
        C_mask = torch.logical_not(torch.logical_or(N_mask,CA_mask))
    # select points
    n_terms  = pred_coords_[:, N_mask[0].squeeze()]
    c_alphas = pred_coords_[:, CA_mask[0].squeeze()]
    c_terms  = pred_coords_[:, C_mask[0].squeeze()]
    # compute phis for every pritein in the batch
    phis = [get_dihedral_torch(c_terms[i, :-1],
                               n_terms[i,  1:],
                               c_alphas[i, 1:],
                               c_terms[i,  1:]) for i in range(pred_coords.shape[0])]

    # return percentage of lower than 0
    if prop:
        return torch.stack([(x<0).float().mean() for x in phis], dim=0 )
    return phis

def calc_phis_numpy(pred_coords, N_mask, CA_mask, C_mask=None,
                    prop=True, verbose=0):
    """ Filters mirrors selecting the 1 with most N of negative phis.
        Used as part of the MDScaling wrapper if arg is passed. See below.
        Angle Phi between planes: (Cterm{-1}, N, Ca{0}) and (N{0}, Ca{+1}, Cterm{+1})
        Inputs:
        * pred_coords: (batch, 3, N) predicted coordinates
        * N_mask: (N, ) boolean mask for N-term positions
        * CA_mask: (N, ) boolean mask for C-alpha positions
        * C_mask: (N, ) or None. boolean mask for C-alpha positions or
                    automatically calculate from N_mask and CA_mask if None.
        * prop: bool. whether to return as a proportion of negative phis.
        * verbose: bool. verbosity level
        Output: (batch, N) containing the phi angles or (batch,) containing
                the proportions.
    """
    # detach gradients for angle calculation - mirror selection
    pred_coords_ = np.transpose(pred_coords, (0, 2, 1))
    n_terms  = pred_coords_[:, N_mask.squeeze()]
    c_alphas = pred_coords_[:, CA_mask.squeeze()]
    # select c_term auto if not passed
    if C_mask is not None:
        c_terms = pred_coords_[:, C_mask]
    else:
        c_terms  = pred_coords_[:, (np.ones_like(N_mask)-N_mask-CA_mask).squeeze().astype(bool) ]
    # compute phis for every pritein in the batch
    phis = [get_dihedral_numpy(c_terms[i, :-1],
                               n_terms[i,  1:],
                               c_alphas[i, 1:],
                               c_terms[i,  1:]) for i in range(pred_coords.shape[0])]

    # return percentage of lower than 0
    if prop:
        return np.array( [(x<0).mean() for x in phis] )
    return phis

# alignment by centering + rotation to compute optimal RMSD
# adapted from : https://github.com/charnley/rmsd/

def kabsch_torch(X, Y, cpu=True):
    """ Kabsch alignment of X into Y.
        Assumes X,Y are both (Dims x N_points). See below for wrapper.
    """
    device = X.device
    #  center X and Y to the origin
    X_ = X - X.mean(dim=-1, keepdim=True)
    Y_ = Y - Y.mean(dim=-1, keepdim=True)
    # calculate convariance matrix (for each prot in the batch)
    C = torch.matmul(X_, Y_.t()).detach()
    if cpu:
        C = C.cpu()
    # Optimal rotation matrix via SVD
    if int(torch.__version__.split(".")[1]) < 8:
        # warning! int torch 1.<8 : W must be transposed
        V, S, W = torch.svd(C)
        W = W.t()
    else:
        V, S, W = torch.linalg.svd(C)

    # determinant sign for direction correction
    d = (torch.det(V) * torch.det(W)) < 0.0
    if d:
        S[-1]    = S[-1] * (-1)
        V[:, -1] = V[:, -1] * (-1)
    # Create Rotation matrix U
    U = torch.matmul(V, W).to(device)
    # calculate rotations
    X_ = torch.matmul(X_.t(), U).t()
    # return centered and aligned
    return X_, Y_

def kabsch_numpy(X, Y):
    """ Kabsch alignment of X into Y.
        Assumes X,Y are both (Dims x N_points). See below for wrapper.
    """
    # center X and Y to the origin
    X_ = X - X.mean(axis=-1, keepdims=True)
    Y_ = Y - Y.mean(axis=-1, keepdims=True)
    # calculate convariance matrix (for each prot in the batch)
    C = np.dot(X_, Y_.transpose())
    # Optimal rotation matrix via SVD
    V, S, W = np.linalg.svd(C)
    # determinant sign for direction correction
    d = (np.linalg.det(V) * np.linalg.det(W)) < 0.0
    if d:
        S[-1]    = S[-1] * (-1)
        V[:, -1] = V[:, -1] * (-1)
    # Create Rotation matrix U
    U = np.dot(V, W)
    # calculate rotations
    X_ = np.dot(X_.T, U).T
    # return centered and aligned
    return X_, Y_

# metrics - more formulas here: http://predictioncenter.org/casp12/doc/help.html

def distmat_loss_torch(X=None, Y=None, X_mat=None, Y_mat=None, p=2, q=2,
                       custom=None, distmat_mask=None, clamp=None):
    """ Calculates a loss on the distance matrix - no need to align structs.
        Inputs:
        * X: (N, d) tensor. the predicted structure. One of (X, X_mat) is needed.
        * X_mat: (N, N) tensor. the predicted distance matrix. Optional ()
        * Y: (N, d) tensor. the true structure. One of (Y, Y_mat) is needed.
        * Y_mat: (N, N) tensor. the predicted distance matrix. Optional ()
        * p: int. power for the distance calculation (2 for euclidean)
        * q: float. power for the scaling of the loss (2 for MSE, 1 for MAE, etc)
        * custom: func or None. custom loss over distance matrices.
                  ex: lambda x,y: 1 - 1/ (1 + ((x-y))**2) (1 is very bad. 0 is good)
        * distmat_mask: (N, N) mask (boolean or weights for each ij pos). optional.
        * clamp: tuple of (min,max) values for clipping distance matrices. ex: (0,150)
    """
    assert (X is not None or X_mat is not None) and \
           (Y is not None or Y_mat is not None), "The true and predicted coords or dist mats must be provided"
    # calculate distance matrices
    if X_mat is None:
        X = X.squeeze()
        if clamp is not None:
            X = torch.clamp(X, *clamp)
        X_mat = torch.cdist(X, X, p=p)
    if Y_mat is None:
        Y = Y.squeeze()
        if clamp is not None:
            Y = torch.clamp(Y, *clamp)
        Y_mat = torch.cdist(Y, Y, p=p)
    if distmat_mask is None:
        distmat_mask = torch.ones_like(Y_mat).bool()

    # do custom expression if passed
    if custom is not None:
        return custom(X_mat.squeeze(), Y_mat.squeeze()).mean()
    # **2 ensures always positive. Later scale back to desired power
    else:
        loss = ( X_mat - Y_mat )**2
        if q != 2:
            loss = loss**(q/2)
        return loss[distmat_mask].mean()

def rmsd_torch(X, Y):
    """ Assumes x,y are both (B x D x N). See below for wrapper. """
    return torch.sqrt( torch.mean((X - Y)**2, axis=(-1, -2)) )

def rmsd_numpy(X, Y):
    """ Assumes x,y are both (B x D x N). See below for wrapper. """
    return np.sqrt( np.mean((X - Y)**2, axis=(-1, -2)) )

def gdt_torch(X, Y, cutoffs, weights=None):
    """ Assumes x,y are both (B x D x N). see below for wrapper.
        * cutoffs is a list of `K` thresholds
        * weights is a list of `K` weights (1 x each threshold)
    """
    device = X.device
    if weights is None:
        weights = torch.ones(1,len(cutoffs))
    else:
        weights = torch.tensor([weights]).to(device)
    # set zeros and fill with values
    GDT = torch.zeros(X.shape[0], len(cutoffs), device=device)
    dist = ((X - Y)**2).sum(dim=1).sqrt()
    # iterate over thresholds
    for i,cutoff in enumerate(cutoffs):
        GDT[:, i] = (dist <= cutoff).float().mean(dim=-1)
    # weighted mean
    return (GDT*weights).mean(-1)

def gdt_numpy(X, Y, cutoffs, weights=None):
    """ Assumes x,y are both (B x D x N). see below for wrapper.
        * cutoffs is a list of `K` thresholds
        * weights is a list of `K` weights (1 x each threshold)
    """
    if weights is None:
        weights = np.ones( (1,len(cutoffs)) )
    else:
        weights = np.array([weights])
    # set zeros and fill with values
    GDT = np.zeros( (X.shape[0], len(cutoffs)) )
    dist = np.sqrt( ((X - Y)**2).sum(axis=1) )
    # iterate over thresholds
    for i,cutoff in enumerate(cutoffs):
        GDT[:, i] = (dist <= cutoff).mean(axis=-1)
    # weighted mean
    return (GDT*weights).mean(-1)

def tmscore_torch(X, Y):
    """ Assumes x,y are both (B x D x N). see below for wrapper. """
    L = max(15, X.shape[-1])
    d0 = 1.24 * (L - 15)**(1/3) - 1.8
    # get distance
    dist = ((X - Y)**2).sum(dim=1).sqrt()
    # formula (see wrapper for source):
    return (1 / (1 + (dist/d0)**2)).mean(dim=-1)

def tmscore_numpy(X, Y):
    """ Assumes x,y are both (B x D x N). see below for wrapper. """
    L = max(15, X.shape[-1])
    d0 = 1.24 * np.cbrt(L - 15) - 1.8
    # get distance
    dist = np.sqrt( ((X - Y)**2).sum(axis=1) )
    # formula (see wrapper for source):
    return (1 / (1 + (dist/d0)**2)).mean(axis=-1)

def mdscaling_torch(pre_dist_mat, weights=None, iters=10, tol=1e-5,
                    fix_mirror=True, N_mask=None, CA_mask=None, C_mask=None,
                    eigen=False, verbose=2):
    """ Handles the specifics of MDS for proteins (mirrors, ...) """
    # batched mds for full parallel
    preds, stresses = mds_torch(pre_dist_mat, weights=weights,iters=iters,
                                              tol=tol, eigen=eigen, verbose=verbose)
    if not fix_mirror:
        return preds, stresses

    # no need to caculate multiple mirrors - just correct Z axis
    phi_ratios = calc_phis_torch(preds, N_mask, CA_mask, C_mask, prop=True)
    to_correct = torch.nonzero( (phi_ratios < 0.5)).view(-1)
    # fix mirrors by (-1)*Z if more (+) than (-) phi angles
    preds[to_correct, -1] = (-1)*preds[to_correct, -1]
    if verbose == 2:
        print("Corrected mirror idxs:", to_correct)

    return preds, stresses

def mdscaling_numpy(pre_dist_mat, weights=None, iters=10, tol=1e-5,
                    fix_mirror=True, N_mask=None, CA_mask=None, C_mask=None, verbose=2):
    """ Handles the specifics of MDS for proteins (mirrors, ...) """
    # batched mds for full parallel
    preds, stresses = mds_numpy(pre_dist_mat, weights=weights,iters=iters,
                                              tol=tol, verbose=verbose)
    if not fix_mirror:
        return preds, stresses

    # no need to caculate multiple mirrors - just correct Z axis
    phi_ratios = calc_phis_numpy(preds, N_mask, CA_mask, C_mask, prop=True)
    for i,pred in enumerate(preds):
        # fix mirrors by (-1)*Z if more (+) than (-) phi angles
        if phi_ratios < 0.5:
            preds[i, -1] = (-1)*preds[i, -1]
            if verbose == 2:
                print("Corrected mirror in struct no.", i)

    return preds, stresses

def lddt_ca_torch(true_coords, pred_coords, cloud_mask, r_0=15.):
    """ Computes the lddt score for each C_alpha.
        https://academic.oup.com/bioinformatics/article/29/21/2722/195896
        Inputs:
        * true_coords: (b, l, c, d) in sidechainnet format.
        * pred_coords: (b, l, c, d) in sidechainnet format.
        * cloud_mask : (b, l, c) adapted for scn format.
        * r_0: float. maximum inclusion radius in reference struct.
        Outputs:
        * (b, l) lddt for c_alpha scores (ranging between 0 and 1)
        See wrapper below.
    """
    device, dtype = true_coords.device, true_coords.type()
    thresholds = torch.tensor([0.5, 1, 2, 4], device=device).type(dtype)
    # adapt masks
    cloud_mask = cloud_mask.bool().cpu()
    c_alpha_mask  = torch.zeros(cloud_mask.shape[1:], device=device).bool() # doesn't have batch dim
    c_alpha_mask[..., 1] = True
    # container for c_alpha scores (between 0,1)
    wrapper = torch.zeros(true_coords.shape[:2], device=device).type(dtype)

    for bi, seq in enumerate(true_coords):
        # select atoms for study
        c_alphas = cloud_mask[bi]*c_alpha_mask # only pick c_alpha positions
        selected_pred = pred_coords[bi, c_alphas, :]
        selected_target = true_coords[bi, c_alphas, :]
        # get number under distance
        dist_mat_pred = torch.cdist(selected_pred, selected_pred, p=2)
        dist_mat_target = torch.cdist(selected_target, selected_target, p=2)
        under_r0_target = dist_mat_target < r_0
        compare_dists = torch.abs(dist_mat_pred - dist_mat_target)[under_r0_target]
        # measure diff below threshold
        score = torch.zeros_like(under_r0_target).float()
        max_score = torch.zeros_like(under_r0_target).float()
        max_score[under_r0_target] = 4.
        # measure under how many thresholds
        score[under_r0_target] = thresholds.shape[0] - \
                                 torch.bucketize( compare_dists, boundaries=thresholds ).float()
        # dont include diagonal
        l_mask = c_alphas.float().sum(dim=-1).bool()
        wrapper[bi, l_mask] = ( score.sum(dim=-1) - thresholds.shape[0] ) / \
                              ( max_score.sum(dim=-1) - thresholds.shape[0] )

    return wrapper

################
### WRAPPERS ###
################

@set_backend_kwarg
@invoke_torch_or_numpy(mdscaling_torch, mdscaling_numpy)
def MDScaling(pre_dist_mat, **kwargs):
    """ Gets distance matrix (-ces). Outputs 3d.
        Assumes (for now) distrogram is (N x N) and symmetric.
        For support of ditograms: see `center_distogram_torch()`
        Inputs:
        * pre_dist_mat: (1, N, N) distance matrix.
        * weights: optional. (N x N) pairwise relative weights .
        * iters: number of iterations to run the algorithm on
        * tol: relative tolerance at which to stop the algorithm if no better
               improvement is achieved
        * backend: one of ["numpy", "torch", "auto"] for backend choice
        * fix_mirror: int. number of iterations to run the 3d generation and
                      pick the best mirror (highest number of negative phis)
        * N_mask: indexing array/tensor for indices of backbone N.
                  Only used if fix_mirror > 0.
        * CA_mask: indexing array/tensor for indices of backbone C_alpha.
                   Only used if fix_mirror > 0.
        * verbose: whether to print logs
        Outputs:
        * best_3d_coords: (3 x N)
        * historic_stress: (timesteps, )
    """
    pre_dist_mat = expand_dims_to(pre_dist_mat, 3 - len(pre_dist_mat.shape))
    return pre_dist_mat, kwargs

@expand_arg_dims(dim_len = 2)
@set_backend_kwarg
@invoke_torch_or_numpy(kabsch_torch, kabsch_numpy)
def Kabsch(A, B):
    """ Returns Kabsch-rotated matrices resulting
        from aligning A into B.
        Adapted from: https://github.com/charnley/rmsd/
        * Inputs:
            * A,B are (3 x N)
            * backend: one of ["numpy", "torch", "auto"] for backend choice
        * Outputs: tensor/array of shape (3 x N)
    """
    # run calcs - pick the 0th bc an additional dim was created
    return A, B

@expand_arg_dims()
@set_backend_kwarg
@invoke_torch_or_numpy(rmsd_torch, rmsd_numpy)
def RMSD(A, B):
    """ Returns RMSD score as defined here (lower is better):
        https://en.wikipedia.org/wiki/
        Root-mean-square_deviation_of_atomic_positions
        * Inputs:
            * A,B are (B x 3 x N) or (3 x N)
            * backend: one of ["numpy", "torch", "auto"] for backend choice
        * Outputs: tensor/array of size (B,)
    """
    return A, B

@expand_arg_dims()
@set_backend_kwarg
@invoke_torch_or_numpy(gdt_torch, gdt_numpy)
def GDT(A, B, *, mode="TS", cutoffs=[1,2,4,8], weights=None):
    """ Returns GDT score as defined here (highre is better):
        Supports both TS and HA
        http://predictioncenter.org/casp12/doc/help.html
        * Inputs:
            * A,B are (B x 3 x N) (np.array or torch.tensor)
            * cutoffs: defines thresholds for gdt
            * weights: list containing the weights
            * mode: one of ["numpy", "torch", "auto"] for backend
        * Outputs: tensor/array of size (B,)
    """
    # define cutoffs for each type of gdt and weights
    cutoffs = [0.5,1,2,4] if mode in ["HA", "ha"] else [1,2,4,8]
    # calculate GDT
    return A, B, cutoffs, {'weights': weights}

@expand_arg_dims()
@set_backend_kwarg
@invoke_torch_or_numpy(tmscore_torch, tmscore_numpy)
def TMscore(A, B):
    """ Returns TMscore as defined here (higher is better):
        >0.5 (likely) >0.6 (highly likely) same folding.
        = 0.2. https://en.wikipedia.org/wiki/Template_modeling_score
        Warning! It's not exactly the code in:
        https://zhanglab.ccmb.med.umich.edu/TM-score/TMscore.cpp
        but will suffice for now.
        Inputs:
            * A,B are (B x 3 x N) (np.array or torch.tensor)
            * mode: one of ["numpy", "torch", "auto"] for backend
        Outputs: tensor/array of size (B,)
    """
    return A, B


#Dossier de configuration des variables globales, liens d'API, ou autres services général
SYMBOL = "ETHUSDT"
PAIR = "USDT"

#Path
PATH_FEATURE = "data/usdt_data_features/"
PATH_DATA = "data/usdt_data/"
PATH_SYMBOL = "data/symbols.csv"
PATH_BAD_SYMBOL = "data/badSymbols.csv"

#DEBUG FLAG FOR DL
DEBUG_DL = True
TIMER_DL = True

#Numver of value needed for update indicator
NBR_VALUE_FEATURE = 7000


from django.shortcuts import render, redirect
from case.models import Case
from case.task import (repeat_case)
# Create your views here.
def home(request):
    all_case = Case.objects.all()
    return render(request, 'case/case.html', {'cases': all_case})

def make_repeat(request):
    if request.method == "POST":
        for i in dict(request.POST)["cases"]:
            repeat_case.delay(int(i), int(request.POST.get("repeat")))
    return redirect("home")


import pandas as pd
to=int(input("Enter the total number of elements you want in the series:"))
lst,le = [],0
for i in range(0, to):
	ele = str(input("Enter a word:"))
	lst.append(ele)
le=len(lst)
series=pd.Series(lst)
print("After converting it to uppercase:",end=" ")
for i in series:
    i=i.capitalize()
    print(i ,end=" ")


import random
#  To testn
import math

class MyHashableKey:

    def __init__(self, int_value, string_value):
        self.int_value = int_value
        self.string_value = string_value

    def __eq__(self, other):
        return (self.int_value == other.int_value) and (self.string_value ==
                other.string_value)

    def __hash__(self):

        string_int = 0
        last_char = chr(1)
        offset = 31

         #  Use offset to help make string int more unique for each
         #  string since 'alfred' and 'falder' would generate the same integer
         #  value
        for char in self.string_value:
            string_int += ord(char) * pow(offset, ord(last_char))
            last_char = char
            offset += 1

        return string_int * self.int_value

if __name__ == "__main__":
    arr = [0] * 16

    for i in range(200):
        test_str = ""
        for i in range(14):
            test_str += chr(random.randint(65, 122))
        test_test = random.randint(1, 100)
        key = MyHashableKey(test_test, test_str)
        index = hash(key)%16
        arr[index] += 1
    print(arr)


from __future__ import absolute_import
from datetime import timedelta

from st2common import log as logging
from st2common.garbage_collection.workflows import purge_workflow_executions
from st2common.models.db.workflow import WorkflowExecutionDB
from st2common.persistence.workflow import WorkflowExecution
from st2common.util import date as date_utils
from st2tests.base import CleanDbTestCase

LOG = logging.getLogger(__name__)

class TestPurgeWorkflowExecutionInstances(CleanDbTestCase):
    @classmethod
    def setUpClass(cls):
        CleanDbTestCase.setUpClass()
        super(TestPurgeWorkflowExecutionInstances, cls).setUpClass()

    def setUp(self):
        super(TestPurgeWorkflowExecutionInstances, self).setUp()

    def test_no_timestamp_doesnt_delete(self):
        now = date_utils.get_datetime_utc_now()

        instance_db = WorkflowExecutionDB(
            start_timestamp=now - timedelta(days=20),
            end_timestamp=now - timedelta(days=20),
            status="running",
        )
        WorkflowExecution.add_or_update(instance_db)

        self.assertEqual(len(WorkflowExecution.get_all()), 1)
        expected_msg = "Specify a valid timestamp"
        self.assertRaisesRegexp(
            ValueError,
            expected_msg,
            purge_workflow_executions,
            logger=LOG,
            timestamp=None,
        )
        self.assertEqual(len(WorkflowExecution.get_all()), 1)

    def test_purge(self):
        now = date_utils.get_datetime_utc_now()

        instance_db = WorkflowExecutionDB(
            start_timestamp=now - timedelta(days=20),
            end_timestamp=now - timedelta(days=20),
            status="failed",
        )
        WorkflowExecution.add_or_update(instance_db)

        instance_db = WorkflowExecutionDB(
            start_timestamp=now - timedelta(days=20),
            status="running",
        )
        WorkflowExecution.add_or_update(instance_db)

        instance_db = WorkflowExecutionDB(
            start_timestamp=now - timedelta(days=5),
            end_timestamp=now - timedelta(days=5),
            status="succeeded",
        )
        WorkflowExecution.add_or_update(instance_db)

        self.assertEqual(len(WorkflowExecution.get_all()), 3)
        purge_workflow_executions(logger=LOG, timestamp=now - timedelta(days=10))
        self.assertEqual(len(WorkflowExecution.get_all()), 2)

    def test_purge_incomplete(self):
        now = date_utils.get_datetime_utc_now()

        instance_db = WorkflowExecutionDB(
            start_timestamp=now - timedelta(days=20),
            end_timestamp=now - timedelta(days=20),
            status="cancelled",
        )
        WorkflowExecution.add_or_update(instance_db)

        instance_db = WorkflowExecutionDB(
            start_timestamp=now - timedelta(days=20),
            status="running",
        )
        WorkflowExecution.add_or_update(instance_db)

        instance_db = WorkflowExecutionDB(
            start_timestamp=now - timedelta(days=5),
            end_timestamp=now - timedelta(days=5),
            status="succeeded",
        )
        WorkflowExecution.add_or_update(instance_db)

        self.assertEqual(len(WorkflowExecution.get_all()), 3)
        purge_workflow_executions(
            logger=LOG, timestamp=now - timedelta(days=10), purge_incomplete=True
        )
        self.assertEqual(len(WorkflowExecution.get_all()), 1)


# Conversion from kilos to pounds
# Author: Brian Hamburg

# set kg counter
kg = 1

# print table labels
print("Kilograms    Pounds")

# print table
while kg < 200:
    print(format(kg, "<3d"), format(kg * 2.2, "15.1f"))
    kg += 2


##operation on image

import numpy as np
import cv2

img = cv2.imread('C:\\Users\\YASHU\\Desktop\\image_repo\\flowers.jpeg',cv2.IMREAD_COLOR)  #reading an images

px=img[55,55]

img[100:150, 100:150] = [0,0,0] # convert the colour of the part specifies by the points into the specified color
cv2.imshow('image',img)

# cut and paste
img_part = img[40:100, 100:150]
img[0:60,0:50]=img_part
cv2.imshow('imge',img)

cv2.waitKey()
cv2.destroyAllWindows()


# -*- coding: utf-8 -*-
"""
Bias
Created on Wed Feb 12 07:17:20 2020
Return distance between layers
@author: Juan G Diosa

This function calculates the distance between layers in a multilayer
structure, given the number of layers, the thickness of the structure,
and a bias factor.

"""
def BiasV(layers,thickness,bias):
    import numpy as np

    #-------------------------------------------------------------------------#
    # Create an array of zeros to store the distances between layers.
    dz=np.zeros((layers,1))

    # Bias factor is equal to 1.
    if bias==1:
        # If it is, the distance between each layer is the same, and the function calculates
        # the distance using the formula (thickness/(layers-1)).
        for m in range(2,layers+1):
            dz[m-1,0]=dz[m-2,0]+(thickness/(layers-1))

    else:
        # If the bias factor is not equal to 1, the function uses a different
        # formula to calculate the distance between layers.
        # The formula takes into account the bias factor and ensures
        # that the distance between layers changes exponentially from one layer to the next.

        # Calculate the sum of the geometric series (bias^0 + bias^1 + ... + bias^(layers-2)).
        bfac=0
        for a in range(1,layers):
            bfac=bfac+bias**(a-1)

        # Calculate the distance between the layers based on the bias
        # factor and the total thickness of the structure.
        zbias=thickness/bfac
        dz[1,0]=dz[0,0]+zbias
        for e in range(3,layers):
            dz[e-1,0]=dz[e-2,0]+(zbias*bias**(e-2))
        dz[layers-1,0]=thickness

    # Return the array of distances between layers.
    return(dz)


import pylab as pl
import numpy as np
from leastsq import polynomial_fit_normal_equations, polynomial_fit
import matplotlibpreferences

matplotlibpreferences.load_preferences()

# Temperature, resistance
# T (degC), R (Ohm)
# 28,30,...,60
temp = range(28, 62, 2)
copper = [
    106.0,
    106.8,
    107.0,
    108.0,
    108.6,
    109.3,
    110.1,
    110.8,
    111.5,
    111.9,
    112.6,
    113.3,
    113.7,
    114.3,
    114.9,
    115.6,
    116.1,
]
platinum = [
    111.0,
    111.6,
    112.4,
    113.2,
    114.0,
    114.8,
    115.4,
    115.9,
    117.1,
    117.6,
    118.6,
    119.2,
    119.9,
    120.9,
    121.5,
    122.2,
    122.9,
]
# Print data
print(np.array([temp, copper, platinum]).T)

print(u"Cuivre:")
betCopper1 = polynomial_fit_normal_equations(temp, copper, 1)
print(u"(Normal) bet=", betCopper1)
print(u"    Resistivité=", betCopper1[0] / betCopper1[1], "Ohm/Deg")
betCopper2 = polynomial_fit(temp, copper, 1)
print(u"(QR) bet=", betCopper2)
print(u"    Resistivité=", betCopper2[0] / betCopper2[1], "Ohm/Deg")

print(u"Platine:")
betPlatine1 = polynomial_fit_normal_equations(temp, platinum, 1)
print(u"(Normal) bet=", betPlatine1)
print(u"    Resistivité=", betPlatine1[0] / betPlatine1[1], "Ohm/Deg")
betPlatine2 = polynomial_fit(temp, platinum, 1)
print(u"(QR) bet=", betPlatine2)
print(u"    Resistivité=", betPlatine2[0] / betPlatine2[1], "Ohm/Deg")

#
t = np.linspace(28, 60)
pl.figure(figsize=(3.5, 2.0))
pl.plot(temp, copper, "o", label=u"Cu (données)", color="tab:blue")
modeleCuivre = betCopper2[0] * t + betCopper2[1]
pl.plot(t, modeleCuivre, "-", label=u"Cu (modèle)", color="tab:blue")
pl.plot(temp, platinum, "x", label=u"Pt (données)", color="tab:orange")
modelePlatine = betPlatine2[0] * t + betPlatine2[1]
pl.plot(t, modelePlatine, "--", label=u"Pt (modèle)", color="tab:orange")
pl.legend(loc="best")
pl.xlim(left=20.0)
pl.ylim(top=130.0)
pl.xlabel(u"Température (°C)")
pl.ylabel(u"Résistance (Ohm)")
pl.title(u"Résistance en fonction de la température.")
pl.savefig("resistivite.pdf", bbox_inches="tight")

# Resolution
# Design matrix
n = 2
X = np.vander(temp, n)
print("X=")
print(X)
y = copper
print("y=")
print(y)
Q, R = np.linalg.qr(X)
z = Q.T @ y
print("z=")
print(z)
beta = np.linalg.solve(R, z)
print("beta=")
print(beta)
alpha = beta[0] / beta[1]
print(u"Resistivité= %.e3 Ohm/Deg" % (alpha))


from random import choice

card_types = list(range(1,14))
cards = []
print (card_types)
suits = ('H', 'D', 'C', 'S')
print(suits)
card_deck = ()
deck_cards = []
for i in suits:
	for n in card_types:
		acard = str(n) + str(i)
		deck_cards.append(acard)
print(deck_cards)


from django.shortcuts import render
from .serializers import *
from .models import *
from rest_framework.decorators import api_view
from rest_framework.response import Response

# Create your views here.

@api_view(['GET'])
def book_list(request):
    book_obj = BooksModel.objects.all()
    serializer = BookSerializer(book_obj, many=True)

    return Response(serializer.data)

@api_view(['POST'])
def post_book(request):
    book_obj = BooksModel.objects.all()
    serializer = BookSerializer(data=request.data)
    if serializer.is_valid():
        serializer.save()
    return Response(serializer.data)

@api_view(['POST'])
def update_book(request, id):
    book_obj = BooksModel.objects.get(id=id)
    serializer = BookSerializer(instance=book_obj, data=request.data)
    if serializer.is_valid():
        serializer.save()
    return Response(serializer.data)

@api_view(['DELETE'])
def delete_book(request, id):
    book_obj = BooksModel.objects.get(id=id)
    book_obj.delete()
    return Response('Book is deleted')


import os
from pathlib import Path
from typing import Callable, Dict

import matplotlib.pyplot as plt
import torch
import torch.distributions
import torch.utils.data
import torchvision
from torch.optim.lr_scheduler import CyclicLR
from torch.utils.tensorboard import SummaryWriter
from torchsummary import summary
from tqdm import tqdm
from yaml import SafeLoader, load

import texture_vae.utils.texture_dataset as dataset
from texture_vae.models.texture_vae import Autoencoder
from texture_vae.utils.mmsim import MSSIM
from texture_vae.utils.utils import plot_image_list

def kl_loss(mu, log_var):
    # see Appendix B from VAE paper:
    # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
    # https://arxiv.org/abs/1312.6114
    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    kld = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp(), dim=1)
    return kld

from texture_vae.utils.mmsim import MSSIM

def train(
    autoencoder: Autoencoder,
    train_data,
    epochs: int,
    kl_weight: float,
    lr: float,
    post_epoch: Callable,
):
    sim = MSSIM()
    opt = torch.optim.Adam(params=autoencoder.parameters(), lr=lr)
    scheduler = CyclicLR(
        opt, base_lr=lr / 2, max_lr=lr * 2, step_size_up=100, cycle_momentum=False
    )
    for epoch in range(epochs):
        losses = []
        with tqdm(train_data, unit="batch") as tepoch:
            for x, l in tepoch:
                x = x.to(autoencoder.device)
                tepoch.set_description(f"epoch {epoch}")
                opt.zero_grad()
                x_hat, mu, logvar, z = autoencoder(x)
                # mse = torch.sum((x - x_hat) ** 2, dim=(1, 2, 3))
                # mse = mse.mean()
                sim_loss = sim(x, x_hat)
                kld = kl_loss(mu, logvar)
                kld = kld.mean()
                loss = (
                    sim_loss + kl_weight * kld
                )  # higher weight gives better generalization result in sampling
                loss.backward()
                opt.step()
                scheduler.step()
                losses.append(loss)
                tepoch.set_postfix(loss=loss.item())
        post_epoch(epoch, x_hat, sim_loss, kld, sum(losses) / len(losses))

def main(config: Dict):
    if not os.path.exists(config["snapshot"]):
        os.makedirs(config["snapshot"])

    LATENT_DIMS = config["model"]["latent_dims"]
    TEXTURE_SIZE = config["model"]["texture_size"]
    KL_WEIGHT = float(config["model"]["kl_weight"])
    LEARNING_RATE = float(config["lr"])
    SNAPSHOT_FN = Path(config["snapshot"]).as_posix()

    autoencoder = Autoencoder(
        latent_dims=LATENT_DIMS, image_size=TEXTURE_SIZE, device="cuda"
    )
    autoencoder.weight_init(0, 0.02)
    try:
        autoencoder.load_state_dict(torch.load(SNAPSHOT_FN))
    except FileNotFoundError as err:
        pass

    data = torch.utils.data.DataLoader(
        dataset.TextureDataset(
            config["train_images"],
            transform=torchvision.transforms.Resize(TEXTURE_SIZE),
        ),
        batch_size=25,
        shuffle=True,
    )
    summary(autoencoder, (3, TEXTURE_SIZE, TEXTURE_SIZE), device=autoencoder.device)
    input_images, labels = next(iter(data))
    images_plt = [
        (input_images[b].permute(1, 2, 0).numpy() * 255).astype("uint8")
        for b in range(input_images.shape[0])
    ]
    plot_image_list(images_plt)
    plt.show()
    plt.rcParams["figure.dpi"] = 300
    tb_writer = SummaryWriter(
        f"{config['logs']['logdir']}/{config['logs']['experiment_name']}"
    )

    def _log_epoch(epoch: int, x_hat, mse_loss, kl_loss, total_loss):
        torch.save(autoencoder.state_dict(), SNAPSHOT_FN)
        tb_writer.add_scalar("loss/mse", mse_loss.item(), epoch)
        tb_writer.add_scalar("loss/kl", kl_loss.item(), epoch)
        tb_writer.add_scalar("loss/total", total_loss.item(), epoch)
        tb_writer.add_images("reconstruction", x_hat, global_step=epoch)
        samples = torch.randn(16, LATENT_DIMS).to("cuda")
        sampled_images = autoencoder.sample(samples)
        tb_writer.add_images("sampled", sampled_images, global_step=epoch)

    train(
        autoencoder,
        data,
        epochs=config["epochs"],
        post_epoch=_log_epoch,
        kl_weight=KL_WEIGHT,
        lr=LEARNING_RATE,
    )
    tb_writer.close()

if __name__ == "__main__":
    from optparse import OptionParser

    parser = OptionParser()
    parser.add_option(
        "-c",
        "--config",
        dest="config_fn",
        default="config.yml",
        help="Config file",
        metavar="FILE",
    )
    (options, args) = parser.parse_args()
    with open(options.config_fn, "r") as fp:
        config = load(fp, SafeLoader)
    main(config)


# Export public PythonMonkey APIs
from .pythonmonkey import *
from .helpers import *
from .require import *

# Expose the package version
import importlib.metadata
__version__= importlib.metadata.version(__name__)
del importlib

# Load the module by default to expose global APIs
## builtin_modules
require("console")
require("base64")
require("timers")
require("url")
require("XMLHttpRequest")

# Add the `.keys()` method on `Object.prototype` to get JSObjectProxy dict() conversion working
# Conversion from a dict-subclass to a strict dict by `dict(subclass)` internally calls the .keys() method to read the dictionary keys,
# but .keys on a JSObjectProxy can only come from the JS side
pythonmonkey.eval("""
(makeList) => {
  const keysMethod = {
    get() {
      return () => makeList(...Object.keys(this))
    }
  }
  Object.defineProperty(Object.prototype, "keys", keysMethod)
  Object.defineProperty(Array.prototype, "keys", keysMethod)
}
""", { 'filename': __file__, 'fromPythonFrame': True })(lambda *args: list(args))


import asyncio

from aiogram import types, Dispatcher
from aiogram.dispatcher import DEFAULT_RATE_LIMIT
from aiogram.dispatcher.handler import CancelHandler, current_handler
from aiogram.dispatcher.middlewares import BaseMiddleware
from aiogram.utils.exceptions import Throttled
from loguru import logger

class ThrottlingMiddleware(BaseMiddleware):
    """
    Throtling middleware that react to last message
    """

    def __init__(self, limit=DEFAULT_RATE_LIMIT, key_prefix='antiflood_'):
        self.rate_limit = limit
        self.prefix = key_prefix
        super(ThrottlingMiddleware, self).__init__()

    async def on_process_message(self, message: types.Message, data: dict):
        handler = current_handler.get()
        dispatcher = Dispatcher.get_current()
        if handler:
            limit = getattr(handler, 'throttling_rate_limit', self.rate_limit)
            key = getattr(handler, 'throttling_key', f"{self.prefix}_{handler.__name__}")
        else:
            limit = self.rate_limit
            key = f"{self.prefix}_message"
        try:
            logger.debug(key)
            await dispatcher.throttle(key, rate=limit)
        except Throttled as t:
            r = await self.message_throttled(message, t, data)
            if not r:
                raise CancelHandler()

    async def message_throttled(self, message: types.Message, throttled: Throttled, data: dict):
        handler = current_handler.get()
        dispatcher = Dispatcher.get_current()
        if handler:
            key = getattr(handler, 'throttling_key', f"{self.prefix}_{handler.__name__}")
        else:
            key = f"{self.prefix}_message"
        delta = throttled.rate - throttled.delta

        if throttled.exceeded_count <= 2:
            logger.debug(f'{key} - Too many requests!')
        await asyncio.sleep(delta)
        thr = await dispatcher.check_key(key)

        if thr.exceeded_count == throttled.exceeded_count:
            logger.debug(f'Unlocked.{thr} {throttled}')
            return True


import numpy as np
from PIL import Image
import vec_tool as vt
import sphere_function as sphere
import time

def cam(t):
    return np.array([t, 0.0, 0.0, 0.0])

def lin_solve(a, b):
    if a == 0:
        return False
    else:
        return -b/a

# Determine where the vector intersects the plane
def intersect(point, vector, n_E, d):
    return lin_solve(np.dot(vector, n_E), np.dot(point, n_E) - d)

def inside(v, n_E, d=0):
    return np.dot(v, n_E)-d <= 0

# Create a numpy array with all values zero
size = 101
x = 10
pixels = np.full((size, size, 3), (0, 0, 0), dtype=np.uint8)
colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255)]

# Objects
n = 8
obj_scl = sphere.sphere_scl(n)
obj_vct = sphere.sphere_vct(n)

obj_scl = obj_scl + vt.translate(obj_vct, [x, 0, 0])

# Boost
speed = -0.0
obj_vct = vt.boosty(obj_vct, speed)

# Create camera and rays.
camera = cam(0)
ray = np.ones(4)

for m in range(1):
    start_time = time.time()
    camera = cam(0)
    # Iterate through all pixels
    for i in range(size):
        if i%100 == 0:
            pass
            #print(i)
        for j in range(size):
            ray[2] = (j/size - 0.5)/2
            ray[3] = (i/size - 0.5)/2
            ray[0] = - np.sqrt(ray[1]*ray[1] + ray[2]*ray[2] + ray[3]*ray[3])

            best = 100000

            for l in range(4*n*n+2):
                t = intersect(camera, ray, obj_vct[l], obj_scl[l])
                point = camera + t*ray
                if not inside(ray, obj_vct[l]):
                    pass
                else:
                    if t < best and t >= 0:
                        on_face = True
                        for r in range(4*n*n + 2):
                            if r != l:
                                if not inside(point, obj_vct[r], obj_scl[r]):
                                    on_face = False
                                    break

                        if on_face:
                            best = t
                            pixels[i][j] = vt.color_vct(obj_vct[l], np.array([0, 0, 1]))

    end_time = time.time()

    # Convert the numpy array to an image
    img = Image.fromarray(pixels)

    # Save the image to a file
    img.save(f'sphere-loc-{n}.png')

    print(end_time-start_time)


import math
from ..color_utils import *
from .effects import BaseEffect, DYNAMIC

class CropEffect(BaseEffect):
    def __init__(self, effect, size, offset=0):
        super().__init__(effect.type)
        self.effect = effect
        self.size = size
        self.offset = offset

    def tick(self, pixels, time_delta):
        colors = clone_spliced_pixels(pixels, self.size, self.offset)
        self.effect.tick(colors, time_delta)

        N = len(pixels)
        for i in range(0, self.size):
            pixels[(i + self.offset) % N] = colors[i]

    def clone(self):
        return CropEffect(self.effect.clone(), self.size, self.offset)

class ResizeEffect(BaseEffect):
    def __init__(self, effect, size):
        super().__init__(effect.type)
        self.effect = effect
        self.size = size
        self.colors = None

    def tick(self, pixels, time_delta):
        N = len(pixels)
        ratio = (N - 1) / (self.size - 1)

        if self.colors is None or self.effect.type == DYNAMIC:
            self.colors = resize_clone(pixels, self.size)
            self.effect.tick(self.colors, time_delta)

        for i in range(N):
            pixels[i] = self.colors[int(i / ratio)]

    def clone(self):
        return ResizeEffect(self.effect.clone(), self.size)


#Python code for socket practice

#importing socket and sys modules
import socket
import sys

#Setting host and port variables for server based on user input
host = input ("Type the server IP address: ")
port = int(input("Type the server port #: "))

#Create socket based on inet4 family address and TCP connection
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

try:
    #Bind socket to user specified host addr and port num
    s.bind((host, port))
# Return error message if bind fails, and exit
except socket.error as msg:
    print('Bind failed')
    sys.exit()

#Indicate successful bind of socket
print("# Socket bind complete")

#Setting the number of queued connections to the socket listening for a connection
s.listen(3)
print("# Socket now listening")

#Accepting incoming connections, returning the socket object (conect) and the client address bound to the socket (addr)
connect, addr = s.accept()

#Print the info describing how the client is connected to the server socket
print('# Connected to ' + addr[0] + ':' + str(addr[1]))

#Iterate over the data sent between the server and client, and continue to print as long as there is data being transmitted
#When no data exists, close the connection
while True:
    data = connect.recv(1024)
    line = data.decode('UTF-8')    # convert to string
    line = line.replace("\n","")   # remove newline character
    print(line)

s.close()


from numpy.random import rand, randn
from numpy import tanh, zeros, ones, eye, arctanh
from numpy import concatenate, dot, diag, sum
from numpy.linalg import norm, pinv
from copy import deepcopy

class NeuralNetwork:

    def __init__(self,
                 number_of_inputs,
                 number_of_hidden_neurons,
                 number_of_outputs,
                 activation_function=lambda x: tanh(x),
                 activation_function_derivation=lambda x: 1+tanh(x)**2,
                 output_function=lambda x: x,
                 output_function_derivation=lambda x: ones(x.shape),
                 output_function_inv=lambda x: x,
                 max_epoch = 10000,
                 log=False) -> None:

        # initialize NN values
        self.InputNumber = number_of_inputs
        self.OutputNumber = number_of_outputs
        self.HiddenLayerNumber = len(number_of_hidden_neurons)
        self.HiddenLayerNeurons = number_of_hidden_neurons
        self.ParametersNumber = 0
        self.W = [randn(number_of_hidden_neurons[0], number_of_inputs+1)*0.01];
        self.ParametersNumber += self.W[-1].size
        for i in range(1,self.HiddenLayerNumber):
            self.W.append(randn(number_of_hidden_neurons[i], number_of_hidden_neurons[i-1]+1)*0.1)
            self.ParametersNumber += self.W[-1].size
        self.W.append(randn(number_of_outputs , number_of_hidden_neurons[-1]+1))
        self.ParametersNumber += self.W[-1].size
        self.W_new = deepcopy(self.W)
        self.af = activation_function
        self.daf = activation_function_derivation
        self.of = output_function
        self.dof = output_function_derivation
        self.ofi = output_function_inv
        self.max_epoch = max_epoch
        self.log = log
        pass

    def __feedforward(self, Input):
        # modify input (add a 1 as last input for neurons bias)
        Modified_Input = concatenate((Input, ones((1,Input.shape[1]))))

        # all layers inputs are saved here
        # initialize with input 2 layer value
        net = [dot(self.W_new[0] , Modified_Input)]

        # all layers outputs are saved here
        # initialize with first layer output
        y = [self.af(net[0])]

        # all layers slopes are saved here
        # initialize with first layer slope
        s = [self.daf(net[0])]

        for i in range(1 , self.HiddenLayerNumber):
            # modified last calculating layer output to inserting bias effect
            Modified_Input = concatenate((y[-1], ones((1,y[-1].shape[1]))))

            # calculate and add last calculating layer input
            net.append(dot(self.W_new[i] , Modified_Input))
            # calculate and add last calculating layer output
            y.append(self.af(net[-1]))
            # calculate and add last calculating layer slope
            s.append(self.daf(net[-1]))

        # modified last calculating layer output to inserting bias effect
        Modified_Input = concatenate((y[-1], ones((1,y[-1].shape[1]))))
        # calculate and add last layer input
        net.append(dot(self.W_new[self.HiddenLayerNumber] , Modified_Input))
        # calculate and add last layer output
        y.append(self.of(net[-1]))
        # calculate and add last layer slope
        s.append(self.dof(net[-1]))

        return y , net , s

    def __call__(self, Input):
        output, _ , _ = self.__feedforward(Input)
        return output[-1]

    def Train(self, Inputs, Outputs, learning_factor=0.1, thershold=1e-3) -> None:
        # delta list making
        delta = []

        # initialize delta
        for i in range(self.HiddenLayerNumber):
            delta.append(zeros((self.HiddenLayerNeurons[i]+1, self.HiddenLayerNeurons[i]+1)))
        delta.append(zeros((self.OutputNumber, self.OutputNumber)))

        # Fast Correction
        y , net , s = self.__feedforward(Inputs)
        Modified_Input = concatenate((y[self.HiddenLayerNumber-1], ones((1,y[self.HiddenLayerNumber-1].shape[1]))))
        self.W[self.HiddenLayerNumber] = dot(pinv(Modified_Input.T) , self.ofi(Outputs).T).T
        self.W_new = deepcopy(self.W)

        # Forward computation
        y , net , s = self.__feedforward(Inputs)
        errors = y[-1] - Outputs
        max_error = sum(errors**2)**0.5/Inputs.shape[1]

        learning_rate_decision = 0
        learning_factor_tmp = learning_factor

        if self.log:
            print('{:05} , {:.2E} , {:.20f}'.format(0, learning_factor, max_error))

        epoch = 0
        while  max_error > thershold and epoch < self.max_epoch:
            epoch += 1
            delta_wp = zeros((self.ParametersNumber,1))
            H = eye(self.ParametersNumber)*learning_factor_tmp

            # Back Propagation for Jacobian
            for p in range(Inputs.shape[1]):
                J = zeros((self.OutputNumber , self.ParametersNumber))

                param_index = 0;
                delta[self.HiddenLayerNumber] = diag(s[self.HiddenLayerNumber][:,p])
                for j in range(self.W[self.HiddenLayerNumber].shape[0]):
                    for k in range(self.W[self.HiddenLayerNumber].shape[1]-1):
                        J[:,param_index] = delta[self.HiddenLayerNumber][:,j]*y[self.HiddenLayerNumber-1][k,p]
                        param_index+=1
                    J[:,param_index] = delta[self.HiddenLayerNumber][:,j]
                    param_index+=1

                for i in range(self.HiddenLayerNumber-1,-1,-1):
                    delta[i] = dot(dot(delta[i+1], self.W[i+1][:,:-1]), diag(s[i][:,p]))
                    for j in range(self.W[i].shape[0]):
                        for k in range(self.W[i].shape[1]-1):
                            if i == 0:
                                J[:,param_index] = delta[i][:,j]*Inputs[k,p]
                            else:
                                J[:,param_index] = delta[i][:,j]*y[i-1][k,p]
                            param_index+=1
                        J[:,param_index] = delta[i][:,j]
                        param_index+=1
                H += dot(J.T , J)
                delta_wp += dot(J.T , errors[:,p].reshape((-1,1)))

            # Levenberge-Marquart Algorithm
            delta_w = dot(pinv(H) , delta_wp)

            # Gradient Descent (LMA Like)
            #delta_w = (delta_wp /norm(H))/ learning_factor_tmp

            # update parameters
            param_index = 0;
            for i in range(self.HiddenLayerNumber,-1,-1):
                for j in range(self.W[i].shape[0]):
                    for k in range(self.W[i].shape[1]):
                        self.W_new[i][j,k] -= delta_w[param_index , 0]
                        param_index+=1

            # Forward computation
            new_y , new_net , new_s = self.__feedforward(Inputs)
            new_errors = new_y[-1] - Outputs
            new_max_error = sum(errors**2)**0.5/Inputs.shape[1]

            if (new_max_error > max_error):
                learning_rate_decision += 1
                if learning_rate_decision > 10:
                    learning_rate_decision = 0
                    learning_factor_tmp *= 10

            self.W = deepcopy(self.W_new)
            y = new_y
            net = new_net
            s = new_s
            errors = new_errors
            max_error=new_max_error

            if self.log:
                print('{:05} , {:.2E} , {:.20f}'.format(epoch, learning_factor_tmp, new_max_error))

        pass

if __name__ == '__main__':
    from math import pi
    import numpy
    from matplotlib import pyplot as plt

    N = 100;
    Inputs_Delay = 4;
    func = lambda x: numpy.sin(x*pi/180)
    time = numpy.arange(N).reshape((1,-1))*9
    time_delayed = numpy.array([numpy.arange(N)-1-i for i in range(Inputs_Delay)])*9
    Outputs = func(time)
    Inputs = func(time_delayed)
    #print(Inputs)
    #print(Outputs)
    NN = NeuralNetwork(
              number_of_inputs=Inputs_Delay,
              number_of_hidden_neurons=[2],
              number_of_outputs=1,
              #activation_function=lambda x: 2/(1+numpy.exp(-2*x))-1,
              #activation_function_derivation=lambda x: 1+(2/(1+numpy.exp(-2*x))-1)**2,
              max_epoch=1000,
              log=True)

    fig1 = plt.figure("Figure 1")
    plt.plot(time[0,:] , Outputs[0,:] , time[0,:] , NN(Inputs)[0,:])

    NN.Train(Inputs=Inputs, Outputs=Outputs, learning_factor=1e+10, thershold=1e-10)

    fig2 = plt.figure("Figure 2")
    plt.plot(time[0,:] , Outputs[0,:] , time[0,:] , NN(Inputs)[0,:])

    plt.show()


def parse_input(path="input.txt"):
    with open(path, "r") as f:
        lines = f.readlines()

    rounds = [line.split() for line in lines]
    return rounds

SHAPE_SCORE = {
    "X": 1,
    "Y": 2,
    "Z": 3,
    "A": 1,
    "B": 2,
    "C": 3
}

def game_logic(rounds):
    score = 0
    for _round in rounds:
        round_score = 0
        shapes = "".join(str(shape) for shape in _round)

        # win
        if (
            shapes == "AY" or
            shapes == "BZ" or
            shapes == "CX"
        ):
            round_score += 6
        # draw
        elif (
            shapes == "AX" or
            shapes == "BY" or
            shapes == "CZ"
        ):
            round_score += 3

        round_score += SHAPE_SCORE[_round[1]]
        score += round_score

    return score

WIN_SHAPE = {
    "A": "B",
    "B": "C",
    "C": "A"
}

def game_logic2(rounds):
    score = 0

    for _round in rounds:
        round_score = 0
        strategy = _round[1]
        opponent = _round[0]

        # win
        if strategy == "Z":
            round_score += 6
            shape = WIN_SHAPE[opponent]
            round_score += SHAPE_SCORE[shape]

        # draw
        if strategy == "Y":
            round_score += 3
            round_score += SHAPE_SCORE[opponent]

        # lose
        if strategy == "X":
            shape = WIN_SHAPE[WIN_SHAPE[opponent]]
            round_score += SHAPE_SCORE[shape]

        score += round_score
    return score

if __name__ == "__main__":
    # test 1
    rounds = parse_input(path="test.txt")
    test_score = int(rounds[0][0])
    score = game_logic(rounds[1:])
    print(test_score, score)
    assert(test_score == score)

    # puzzle 1
    rounds = parse_input()
    score = game_logic(rounds)
    print("Strategy 1 Score", score)

    # test 2
    rounds = parse_input(path="test2.txt")
    test_score = int(rounds[0][0])
    score = game_logic2(rounds[1:])
    print(test_score, score)
    assert(test_score == score)

    # puzzle 2
    rounds = parse_input()
    score = game_logic2(rounds)
    print("Strategy 2 Score:", score)


import argparse
import json
import os

from itertools import product

archs = [
    "x86_64",
]
build_types = [
    "RelWithDebInfo",
]
actions = [
    "build",
    "debug"
]

def main(args):
    pipeline_steps = []
    cur_build_types = build_types
    if args.build_type is not None:
        cur_build_types = [args.build_type]

    for arch, build_type in product(archs, cur_build_types):
        pipeline_steps.append({
            "label": f"Build & test :cpp: for Windows-{arch}-{build_type} :windows:",
            "timeout_in_minutes": "240",
            "agents": {
              "provider": "gcp",
              "machineType": "c2-standard-16",
              "minCpuPlatform": "Intel Cascade Lake",
              "image": "family/ml-cpp-1-windows-2016",
            },
            "commands": [
              f'if ( "{args.action}" -eq "debug" ) {{\$Env:ML_DEBUG="1"}}',
              "& .buildkite\\scripts\\steps\\build_and_test.ps1"
            ],
            "depends_on": "check_style",
            "key": f"build_test_Windows-{arch}-{build_type}",
            "env": {
              "ML_DEBUG": "0",
              "CPP_CROSS_COMPILE": "",
              "CMAKE_FLAGS": "-DCMAKE_TOOLCHAIN_FILE=cmake/windows-x86_64.cmake",
              "RUN_TESTS": "true",
              "BOOST_TEST_OUTPUT_FORMAT_FLAGS": "--logger=JUNIT,error,boost_test_results.junit",
            },
            "artifact_paths": ["*/**/unittest/boost_test_results.junit"],
            "plugins": {
              "test-collector#v1.2.0": {
                "files": "*/*/unittest/boost_test_results.junit",
                "format": "junit"
              }
            },
            "notify": [
              {
                "github_commit_status": {
                  "context": f"Build and test on Windows {arch} {build_type}",
                },
              },
            ],
        })

    pipeline = {
        "steps": pipeline_steps,
    }
    print(json.dumps(pipeline, indent=2))

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--build-type",
                        required=False,
                        choices=build_types,
                        default=None,
                        help="Specify a specific build type to build")
    parser.add_argument("--action",
                        required=False,
                        choices=actions,
                        default="build",
                        help="Specify a build action.")

    args = parser.parse_args()

    main(args)