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

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

class OnlyTextModel(nn.Module):
    # Model which only uses text part
    def __init__(self, device, BioBert, BioBertConfig):
        super(OnlyTextModel, self).__init__()
        text_embed_size =  BioBertConfig.hidden_size

        self.FinalFC = nn.Linear(text_embed_size, 1, bias=False)

        for p in self.FinalFC.parameters():
            nn.init.xavier_uniform_(p)

        self.final_act = torch.sigmoid
        self.criterion= nn.BCEWithLogitsLoss()
        self.device = device

        self.BioBert = BioBert

    def forward(self, text, attn):
        emb = self.BioBert(text, attn)
        emb = emb[0][:,0,:]
        logit_X = emb
        logits = self.FinalFC(logit_X)
        logits = logits.squeeze()
        probs = self.final_act(logits)
        return logits, probs

    def get_l2(self):
        l2_reg = torch.tensor(0.).to(self.device)
        for param in self.FinalFC.parameters():
            l2_reg += param.norm(2)
        return l2_reg

class TimeAttn(nn.Module):
    def __init__(self, AttnType="All"):
        super(TimeAttn, self).__init__()
        self.a = torch.nn.Parameter(torch.tensor(1.0))
        self.b = torch.nn.Parameter(torch.tensor(1.0))
        self.w1 = torch.nn.Parameter(torch.tensor(1/3))

        self.e = torch.nn.Parameter(torch.tensor(0.02))
        self.f = torch.nn.Parameter(torch.tensor(1.0))
        self.w2 = torch.nn.Parameter(torch.tensor(1/3))

        self.d = torch.nn.Parameter(torch.tensor(40.0))
        self.n = torch.nn.Parameter(torch.tensor(4.5))
        self.w3 = torch.nn.Parameter(torch.tensor(1/3))

        self.AttnType = AttnType
    def GetAgg(self, Vals, Times):
        # We are cutting times values at 1 hour so we don't run into NaNs below
        Times = torch.clamp(Times, 1)
        A_conv = 1/(self.a * (Times ** self.b))
        A_lin = torch.clamp(self.f - self.e * Times, 0)
        A_conc = 1/(1+(Times/self.d) ** self.n)
        if self.AttnType == "All":
            attns = self.w1 * A_conv + self.w2 * A_lin + self.w3 * A_conc
        attn_sum = torch.sum(attns)
        if attn_sum < 1e-5:
            attns = torch.ones_like(attns) / len(attns)
        else:
            attns = attns / attn_sum
        results = attns[None, :].mm(Vals).squeeze()
        return results
class ClinicalNotesModel(nn.Module):
    # Main model with which combines time-series and text part
    def __init__(self, dropout_keep_prob, W_embed, model_name, model_type, BioBert, notes_aggeregate, TSModel, TS_aggeregate, Attn_Type, device):
        super(ClinicalNotesModel, self).__init__()
        self.dropout_keep_prob = dropout_keep_prob
        if model_name != "BioBert":
            # Load word2vec weights for word embedding
            self.W_embed = nn.Embedding.from_pretrained(torch.tensor(W_embed, dtype=torch.float))
            # CNN model for time series the window sizes defined below
            self.sizes = range(2, 5)
            self.CNNs = nn.ModuleList([nn.Conv1d(200, 256, ngram_size, stride=1, padding=ngram_size//2) for ngram_size in self.sizes])
            # Dropout over the text embedding
            self.EmbedDropout = nn.Dropout(1 - dropout_keep_prob)
        self.TSModel = TSModel
        self.TS_aggeregate = TS_aggeregate

        time_series_size = 256
        # LSTM model for time-series
        if TSModel == 'BiLSTM':
            self.lstm = nn.LSTM(76, time_series_size, 1, batch_first=True, bidirectional=True)
        else:
            self.lstm = nn.LSTM(76, time_series_size, 1, batch_first=True)

        if TSModel == 'BiLSTM':
            time_series_size *= 2

        # Using attention
        if TS_aggeregate == 'Attention':
            self.TS_attn = Attention(time_series_size, Attn_Type)

        if model_name == 'avg_we':
            print('Using Average Embedding Vector')
            text_embed_size =  W_embed.shape[1]
        elif model_name == 'cnn':
            print('Using CNN for text part')
            text_embed_size =  len(self.sizes) * 256
        elif model_name == 'BioBert':
            print('Using BioBert model')
            text_embed_size =  BioBert.config.hidden_size

        if model_type == "baseline":
            print('Using only time-series')
            self.FinalFC = nn.Linear(time_series_size, 1, bias=False)
        elif model_type == "text_only":
            print('Using only text part')
            self.FinalFC = nn.Linear(text_embed_size, 1, bias=False)
        else:
            print('Using both text and time series')
            self.FinalFC = nn.Linear(time_series_size + text_embed_size, 1, bias=False)

        for p in self.FinalFC.parameters():
            # Using Xavier Uniform initialization for final fully connected layer
            nn.init.xavier_uniform_(p)
        self.final_act = torch.sigmoid
        self.model_name = model_name
        self.model_type = model_type
        self.criterion= nn.BCEWithLogitsLoss()
        self.device = device
        self.notes_aggeregate = notes_aggeregate

        if self.notes_aggeregate == 'TimeAttn':
            self.TA = TimeAttn(AttnType="All")
        elif self.notes_aggeregate == 'Attn':
            self.NotesAttn = Attention(text_embed_size, Attn_Type)

        self.BioBert = BioBert

    def forward(self, X, text, attns, times):
        # if model_type is baseline we only use the time-series part
        if self.model_type != 'baseline':
            if self.model_name != 'BioBert':
                # When not using bert based models, use word2vec embedding layer to get the embedding vectors
                embeds = self.W_embed(text)
            if self.model_name == 'avg_we':
                text_embeddings = torch.mean(embeds, axis=1)
            elif self.model_name == 'BioBert':
                if self.notes_aggeregate == 'Mean' or  self.notes_aggeregate == 'TimeAttn' or self.notes_aggeregate == 'Attn':
                    # For the models where we take the mean of embeddings generated for several notes we loop through the notes and take their mean
                    txt_arr = []
                    for txts, attn, ts in zip(text, attns, times):
                        if len(txts.shape) == 1:
                            # If there is only one note for a patient we just add a dimension
                            txts = txts[None, :]
                            attn = attn[None, :]
                        txtemb = self.BioBert(txts, attn)
                        emb = txtemb[0][:,0,:]
                        if self.notes_aggeregate == 'TimeAttn':
                            txt_arr.append(self.TA.GetAgg(emb, ts))
                        elif self.notes_aggeregate == 'Attn':
                            Attn_out, _ = self.NotesAttn(emb[-1,:].unsqueeze(0).unsqueeze(0), emb.unsqueeze(0))
                            txt_arr.append(Attn_out.squeeze())
                        else:
                            txt_arr.append(torch.mean(emb, axis=0))
                    text_embeddings = torch.stack(txt_arr)
                    # deleting some tensors to free up some space
                    del txt_arr
                else:
                    # If we are only using one note we just run bert based model and get its embedding
                    text_embeddings = self.BioBert(text, attns)
                    text_embeddings = text_embeddings[0][:,0,:]
            else:
                # If using CNN for texts
                result_tensors = []
                for i in range(len(self.sizes)):
                    # running CNN with different kernel sizes and concatenating the embeddings
                    # 256 -> 2,3 best yet.
                    text_conv1d = self.CNNs[i](embeds.transpose(1,2))
                    text_conv1d = text_conv1d.transpose(1, 2)
                    text_conv1d = F.relu(text_conv1d)
                    text_conv1d, _ = torch.max(text_conv1d, axis=1)
                    result_tensors.append(text_conv1d)
                text_embeddings = torch.cat(result_tensors, axis=1)
                text_embeddings = self.EmbedDropout(text_embeddings)
        # Getting the time-series part embedding
        rnn_outputs, _ = self.lstm(X)
        if self.TS_aggeregate == 'Last':
            mean_rnn_outputs = rnn_outputs[:, -1, :]
        elif self.TS_aggeregate == 'Mean':
            mean_rnn_outputs = torch.mean(rnn_outputs, 1)
        elif self.TS_aggeregate == 'Attention':
            mean_rnn_outputs, _ = self.TS_attn(rnn_outputs[:, -1, :].unsqueeze(1), rnn_outputs)
            mean_rnn_outputs = mean_rnn_outputs.squeeze(1)

        if self.model_type == 'baseline':
            logit_X = mean_rnn_outputs
        elif self.model_type == 'text_only':
            logit_X = text_embeddings
        else:
            # Concatenating time-series and text embedding
            logit_X = torch.cat((text_embeddings.float(), mean_rnn_outputs.float()), 1)

        #Final FC layer
        logits = self.FinalFC(logit_X)
        logits = logits.squeeze(dim=-1)
        probs = self.final_act(logits)
        return logits, probs

    def get_l2(self):
        # get l2 regularization weight of the cnn and final layer
        l2_reg = torch.tensor(0.).to(self.device)
        if self.model_name != 'BioBert':
            for cnn in self.CNNs:
                for param in cnn.parameters():
                    l2_reg += param.norm(2)
        for param in self.FinalFC.parameters():
            l2_reg += param.norm(2)
        return l2_reg

class Attention(nn.Module):
    # Attention from torchnlp
    """ Applies attention mechanism on the `context` using the `query`.

    **Thank you** to IBM for their initial implementation of :class:`Attention`. Here is
    their `License
    <https://github.com/IBM/pytorch-seq2seq/blob/master/LICENSE>`__.

    Args:
        dimensions (int): Dimensionality of the query and context.
        attention_type (str, optional): How to compute the attention score:

            * dot: :math:`score(H_j,q) = H_j^T q`
            * general: :math:`score(H_j, q) = H_j^T W_a q`

    Example:

         >>> attention = Attention(256)
         >>> query = torch.randn(5, 1, 256)
         >>> context = torch.randn(5, 5, 256)
         >>> output, weights = attention(query, context)
         >>> output.size()
         torch.Size([5, 1, 256])
         >>> weights.size()
         torch.Size([5, 1, 5])
    """

    def __init__(self, dimensions, attention_type='general'):
        super(Attention, self).__init__()

        if attention_type not in ['dot', 'general']:
            raise ValueError('Invalid attention type selected.')

        self.attention_type = attention_type
        if self.attention_type == 'general':
            self.linear_in = nn.Linear(dimensions, dimensions, bias=False)

        self.linear_out = nn.Linear(dimensions * 2, dimensions, bias=False)
        self.softmax = nn.Softmax(dim=-1)
        self.tanh = nn.Tanh()

    def forward(self, query, context):
        """
        Args:
            query (:class:`torch.FloatTensor` [batch size, output length, dimensions]): Sequence of
                queries to query the context.
            context (:class:`torch.FloatTensor` [batch size, query length, dimensions]): Data
                overwhich to apply the attention mechanism.

        Returns:
            :class:`tuple` with `output` and `weights`:
            * **output** (:class:`torch.LongTensor` [batch size, output length, dimensions]):
              Tensor containing the attended features.
            * **weights** (:class:`torch.FloatTensor` [batch size, output length, query length]):
              Tensor containing attention weights.
        """
        batch_size, output_len, dimensions = query.size()
        query_len = context.size(1)

        if self.attention_type == "general":
            query = query.reshape(batch_size * output_len, dimensions)
            query = self.linear_in(query)
            query = query.reshape(batch_size, output_len, dimensions)

        # TODO: Include mask on PADDING_INDEX?

        # (batch_size, output_len, dimensions) * (batch_size, query_len, dimensions) ->
        # (batch_size, output_len, query_len)
        attention_scores = torch.bmm(query, context.transpose(1, 2).contiguous())

        # Compute weights across every context sequence
        attention_scores = attention_scores.view(batch_size * output_len, query_len)
        attention_weights = self.softmax(attention_scores)
        attention_weights = attention_weights.view(batch_size, output_len, query_len)

        # (batch_size, output_len, query_len) * (batch_size, query_len, dimensions) ->
        # (batch_size, output_len, dimensions)
        mix = torch.bmm(attention_weights, context)

        # concat -> (batch_size * output_len, 2*dimensions)
        combined = torch.cat((mix, query), dim=2)
        combined = combined.view(batch_size * output_len, 2 * dimensions)

        # Apply linear_out on every 2nd dimension of concat
        # output -> (batch_size, output_len, dimensions)
        output = self.linear_out(combined).view(batch_size, output_len, dimensions)
        output = self.tanh(output)

        return output, attention_weights


from gensim.models.doc2vec import TaggedDocument
from gensim.models import Doc2Vec
import json
from pathlib import Path

from word_parser import Parser

import time
from multiprocessing import Pool

import logging
from logging import getLogger, StreamHandler, Formatter

logger = getLogger('Log')
logger.setLevel(logging.DEBUG)

stream_handler = StreamHandler()
stream_handler.setLevel(logging.DEBUG)

handler_format = Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
stream_handler.setFormatter(handler_format)

logger.addHandler(stream_handler)

class Doc2VecModel:
    def __init__(self, path: str = ''):
        self._model = Doc2Vec.load(path) if path != '' else None
        self._corpus = []

    # 学習
    def train(self):
        # DBoW-STS
        model = Doc2Vec(documents=self._corpus,
                        dm=0, vector_size=300, window=15, alpha=0.025,
                        min_alpha=0.025, min_count=1, sample=1e-4, negative=5)

        '''
        model = Doc2Vec(documents=self._corpus,
                        dm=1, vector_size=300, window=15, alpha=0.025, min_alpha=0.025, min_count=1, sample=1e-6)
        model.train(self._corpus,
                    total_examples=model.corpus_count, epochs=1000)
        '''

        model.train(self._corpus,
                    total_examples=model.corpus_count, epochs=400)
        self._model = model
        self._corpus = []

    def add(self, name: str, file_name, words_list):
        for i, words in enumerate(words_list):
            self._corpus.append(TaggedDocument(
                words=words, tags=['{0}-{1}-{2}'.format(name, file_name, i)]))

    # モデルの保存
    def save(self, path: str):
        self._model.save(path)

    # モデルの読み込み
    def load(self, path: str):
        self._model = Doc2Vec.load(path)

    # 近傍探索
    def most_similar(self, vec, topn=10):
        return self._model.docvecs.most_similar(vec, topn=topn)

    def infer_vector(self, s: str):
        return [self._model.infer_vector(Parser(s).parse_noun())]

def add(file_name, dic, lis, i):
    logger.debug("{0} / {1}, {2}({3})".format(i, len(lis), dic['user'], file_name))
    words_list = []
    for s in dic['message']:
        words = []
        ps = Parser(s['text'])
        # for word in ps.parse():
        for word in ps.parse_noun():
            words.append(word)

        if len(words) == 0:
            continue

        words_list.append(words)

    logger.debug('len(words_list) = {0}...'.format(len(words_list)))
    if len(words_list) == 0:
        return []

    return words_list

def main():
    g = input('>> ')

    st = time.time()

    if g == '1':
        dm = Doc2VecModel()

        p = Path('../data/')

        lis = list(p.glob('*'))

        pool = Pool(10)
        multi_res = []
        for i, file_name in enumerate(lis):
            if file_name == '../data/all.json':
                continue
            dic = {}
            with open(file_name, 'r') as f:
                dic = json.load(f)

            multi_res.append((pool.apply_async(add, (file_name, dic, lis, i)), file_name))

        for res, file_name in multi_res:
            dic = {}
            with open(file_name, 'r') as f:
                dic = json.load(f)

            words_list = res.get()
            if len(words_list) == 0:
                continue
            dm.add(dic['user'], file_name, words_list)

        logger.debug('End {0} s(parse)'.format(time.time() - st))
        dm.train()
        logger.debug('End {0} s'.format(time.time() - st))
        dm.save('../model/doc2vec_v4.model')
    else:
        s = input('>> ')
        dm = Doc2VecModel('../model/doc2vec_v4.model')
        vec = dm.infer_vector(s)
        for key, value in dm.most_similar(vec):
            dic = {}
            with open(key.split('-')[1], 'r') as f:
                dic = json.load(f)

            print(dic['user'], dic['message']
                  [int(key.split('-')[2])]['text'], value)

if __name__ == '__main__':
    main()


import math
import sys

from advent_of_code.basesolver import BaseSolver

class Map:
    def __init__(self, w, h):
        self.w = w
        self.h = h
        self.asteroids = []

    def is_between(self, asteroid1, asteroid2, asteroid):
        # is asteroid between asteroid1 and asteroid2
        crossproduct = (asteroid.y - asteroid1.y) * (asteroid2.x - asteroid1.x) - (
            asteroid.x - asteroid1.x
        ) * (asteroid2.y - asteroid1.y)

        # compare versus epsilon for floating point values, or != 0 if using integers
        if abs(crossproduct) > sys.float_info.epsilon:
            return False

        dotproduct = (asteroid.x - asteroid1.x) * (asteroid2.x - asteroid1.x) + (
            asteroid.y - asteroid1.y
        ) * (asteroid2.y - asteroid1.y)
        if dotproduct < 0:
            return False

        squaredlengthba = (asteroid2.x - asteroid1.x) * (asteroid2.x - asteroid1.x) + (
            asteroid2.y - asteroid1.y
        ) * (asteroid2.y - asteroid1.y)
        if dotproduct > squaredlengthba:
            return False

        return True

    def does_see(self, asteroid1, asteroid2):
        for asteroid in self.asteroids:
            if asteroid != asteroid1 and asteroid != asteroid2:
                is_between = self.is_between(asteroid1, asteroid2, asteroid)
                if is_between:
                    return False
        return True

    def sees(self, asteroid):
        seen_asteroids = []
        for other_asteroid in self.asteroids:
            if other_asteroid != asteroid:
                if self.does_see(asteroid, other_asteroid):
                    seen_asteroids.append(other_asteroid)
        return seen_asteroids

    def calculate_asteroids_sees(self):
        for i, asteroid in enumerate(self.asteroids, 1):
            asteroid_sees = self.sees(asteroid)
            asteroid.sees = asteroid_sees
            # print('Calculated sees: {}/{}'.format(i, len(self.asteroids)))

    def get_asteroid_by_coords(self, x, y):
        for asteroid in self.asteroids:
            if asteroid.x == x and asteroid.y == y:
                return asteroid
        return None

    def get_all_asteroids_between(self, asteroid1, asteroid2):
        all_asteroids_between = []
        for asteroid in self.asteroids:
            if asteroid != asteroid1 and asteroid != asteroid2:
                if self.is_between(asteroid1, asteroid2, asteroid):
                    all_asteroids_between.append(asteroid)
        return all_asteroids_between

    def distance(self, asteroid1, asteroid2):
        return math.sqrt(
            ((asteroid1.x - asteroid2.x) ** 2) + ((asteroid1.y - asteroid2.y) ** 2)
        )

    def destroy_asteroid(self, asteroid):
        self.asteroids.remove(asteroid)

    def draw(self):
        result = ""
        for y in range(self.h):
            for x in range(self.w):
                map_asteroid = self.get_asteroid_by_coords(x, y)
                if map_asteroid:
                    result += "#"
                else:
                    result += "."
            result += "\n"
        print(result.strip())

    def draw_asteroid(self, asteroid):
        result = ""
        for y in range(self.h):
            for x in range(self.w):
                map_asteroid = self.get_asteroid_by_coords(x, y)
                if map_asteroid == asteroid:
                    result += "X"
                elif map_asteroid in asteroid.sees:
                    result += "O"
                elif map_asteroid:
                    result += "#"
                else:
                    result += "."
            result += "\n"
        print(result.strip())

    def __str__(self):
        return str(self.asteroids)

    def __repr__(self):
        return str(self)

class Asteroid:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.sees = None

    def __str__(self):
        return "({}, {})".format(self.x, self.y)

    def __repr__(self):
        return str(self)

def load_input(lines):
    return [x.strip() for x in lines]

def create_map(content):
    map = Map(w=len(content[0]), h=len(content))

    for y, line in enumerate(content):
        for x, char in enumerate(line):
            if char == "#":
                map.asteroids.append(Asteroid(x, y))
    return map

def clockwiseangle_and_distance(point):
    global origin
    global refvec
    # Vector between point and the origin: v = p - o
    vector = [point.x - origin.x, point.y - origin.y]
    # Length of vector: ||v||
    lenvector = math.hypot(vector[0], vector[1])
    # If length is zero there is no angle
    if lenvector == 0:
        return -math.pi, 0
    # Normalize vector: v/||v||
    normalized = [vector[0] / lenvector, vector[1] / lenvector]
    dotprod = normalized[0] * refvec[0] + normalized[1] * refvec[1]  # x1*x2 + y1*y2
    diffprod = refvec[1] * normalized[0] - refvec[0] * normalized[1]  # x1*y2 - y1*x2
    angle = math.atan2(diffprod, dotprod)
    # Negative angles represent counter-clockwise angles so we need to subtract them
    # from 2*pi (360 degrees)
    if angle < 0:
        return 2 * math.pi + angle, lenvector
    # I return first the angle because that's the primary sorting criterium
    # but if two vectors have the same angle then the shorter distance should come first.
    return angle, lenvector

def order_asteroids_clockwise(asteroids):
    global max_sees_asteroid
    sorted_asteroids = sorted(asteroids, key=clockwiseangle_and_distance)
    if sorted_asteroids[0].x == max_sees_asteroid.x:
        return [sorted_asteroids[0]] + list(reversed(sorted_asteroids[1:]))
    else:
        return list(reversed(sorted_asteroids))

max_sees_asteroid = None
origin = None
refvec = None

class Y2019D10Solver(BaseSolver):
    def solve_part_a(self):
        input_content = load_input(self.lines)

        map = create_map(input_content)

        map.calculate_asteroids_sees()

        max_sees_asteroid = map.asteroids[0]
        for asteroid in map.asteroids:
            if len(asteroid.sees) > len(max_sees_asteroid.sees):
                max_sees_asteroid = asteroid

        print("{}: sees {}".format(max_sees_asteroid, len(max_sees_asteroid.sees)))

        return len(max_sees_asteroid.sees)

    def solve_part_b(self):
        global max_sees_asteroid
        global origin
        global refvec
        input_content = load_input(self.lines)

        map = create_map(input_content)

        map.calculate_asteroids_sees()

        max_sees_asteroid = map.asteroids[0]
        for asteroid in map.asteroids:
            if len(asteroid.sees) > len(max_sees_asteroid.sees):
                max_sees_asteroid = asteroid

        print("Station at: {}".format(max_sees_asteroid))

        origin = max_sees_asteroid
        refvec = [0, -1]

        rotation = 0
        destroyed_asteroids = []
        while True:
            # print('Rotation: {}'.format(rotation))
            # map.draw_asteroid(max_sees_asteroid)
            if len(map.asteroids) == 1:
                break
            clockwise_asteroids = order_asteroids_clockwise(max_sees_asteroid.sees)
            # print('Destroying:', clockwise_asteroids)
            for asteroid in clockwise_asteroids:
                map.destroy_asteroid(asteroid)
                destroyed_asteroids.append(asteroid)

                if len(destroyed_asteroids) == 200:
                    return destroyed_asteroids[-1].x * 100 + destroyed_asteroids[-1].y

            map.calculate_asteroids_sees()
            rotation += 1


"""
This script creates the matrices:
the extracted geomaddataset based on the grids setting and mosaik method.

"""

from pathlib import Path
import pickle
import dill
import numpy as np
from mosaiks import config as c
from mosaiks.featurization import featurize, featurize_and_save
from mosaiks.utils import io, spatial

if __name__ == "__main__":
    for mpg in c.meters_per_grid:
        for city in ['Eldoret','Embu','Garissa','Kakamega','Kericho','Kisumu','Kitui','Machakos','Malindi','Mombasa','Nairobi','Naivasha','Nakuru','Nyeri','Thika']:
            subgrid_files = Path(c.grid_dir) / city
            base_image_dir = Path(c.data_dir) / "raw" / "imager-crop-mad" / city
            for f in subgrid_files.glob("grid*.npz"):
                if f.is_file() and str(f).split("_")[-2] == str(mpg):
                    print(f.name)
                    grid_name = f.name
                    grid_name_lst = grid_name.split("_")
                    area = grid_name_lst[1]
                    sample = grid_name_lst[4]
                    image_folder = base_image_dir  #/ f"{area}_{sample}"
                    fsettings = c.features["random"]
                    out_fpath = Path(c.features_dir) / f"{image_folder.name}_{mpg}_{fsettings['num_filters']}_{fsettings['pool_size']}.pkl"

                    assert (
                        image_folder.is_dir()
                    ), f"You have not downloaded images to {image_folder}"

                    featurize_and_save(image_folder, out_fpath, c)


# PROBLEM: Diameter of binary tree
# https://leetcode.com/problems/diameter-of-binary-tree/

class Node:
    def __init__(self, data=0, left=None, right=None):
        self.data = data
        self.left = left
        self.right = right

def NBT (elements, root, i, n) -> Node:
    if n==0 : return None

    if i<n:
        temp = Node(elements[i])
        root = temp
        root.left = NBT(elements, root.left, 2*i+1, n)
        root.right = NBT(elements, root.right, 2*i+2, n)

    return root

# SOLUTION
d = 0
def traverse (root: Node) -> int:
    if root==None : return 0
    ld = traverse(root.left)
    rd = traverse(root.right)
    global d
    d = max(d, ld+rd)
    return max(ld, rd)+1

def diameter (root: Node) -> int:
    traverse(root)
    return d

if __name__ == "__main__":
    # INPUT :
    elements = [1,2,3,4,5]
    root = NBT(elements, Node(elements[0]), 0, len(elements))

    # OUTPUT :
    print(diameter(root))


from tkinter import *
from tkinter import ttk
from PIL import Image, ImageTk

import location_weather_crawling as lwc
import location_gui as lg

def update_weather_animation(wcnt,weather_images,weather_image_label):
	weather_image = weather_images[wcnt]
	weather_image_label.configure(image=weather_image)

	wcnt += 1
	if wcnt == 7:
		wcnt = 0

	weather_info_frame.after(400, update_weather_animation, wcnt, weather_images,weather_image_label)

def weather_info(root,location):
	#날씨정보 프레임
	global weather_info_frame
	weather_info_frame = Frame(root)
	weather_info_frame.pack(side = "top", fill = "both", pady=40)

	#location_gui에서 위치 정보 받아오기
	location_text = location
	print("location text : ", location_text)

	# 받아온 위치정보를 이용해서 날씨 정보값 받아오기
	weather_infos = lwc.location_weather(location_text)

	current_tmp_text = weather_infos["온도"]
	weather_text = weather_infos["날씨"]

	print(current_tmp_text)
	print(weather_text)

	# 날씨 정보 이미지로 받아오기
	weather_dict = {"맑음" : "sunny", "흐림" : "cloudy","구름많음":"cloudy", "구름조금":"partly_sunny", \
		"비":"rainy", "눈":"snowy","천둥번개":"stormy","바람":"windy"}
	global weather_images
	weatherPhoto_file_name = "./images/weather_" + weather_dict[weather_text] + ".gif"

	weather_images = [PhotoImage(file = weatherPhoto_file_name, format="gif -index %i" %(i)) for i in range(7)]

	weather_image_label = Label(weather_info_frame)

	update_weather_animation(0, weather_images, weather_image_label)
	weather_image_label.pack()

	location_label = Label(weather_info_frame, text=location_text)
	location_label.config(font = ("Courier",12,"bold"))
	location_label.pack(anchor="s")

	weather_label = Label(weather_info_frame, text=weather_text)
	weather_label.config(font = ("Courier", 10))
	weather_label.pack(fill="both")

	current_tmp_label = Label(weather_info_frame, text=current_tmp_text)
	current_tmp_label.config(font = ("Courier", 20, "bold"))
	current_tmp_label.pack(fill="both")

	tree = ttk.Treeview(weather_info_frame, columns=["one","two"])
	tree.column("#0",width=120, anchor="center")
	tree.heading("#0",text="날씨 정보",anchor="center")

	tree.column("#1",width=195, anchor="center")
	tree.heading("#1",text="수치",anchor="center")

	for key, value in weather_infos.items():
		tree.insert("","end",text=key,values=value)

	tree.pack(padx=5, pady=5)


import heapq

def solution(n, works):
    # 합보다 크면 0
    if sum(works) <= n:
        return 0

    # 최대 힙 구현
    works = [-i for i in works]
    heapq.heapify(works)

    # 제일 큰수를 최대한 줄여주는 방식
    while n > 0:
        max_num = heapq.heappop(works)
        heapq.heappush(works, max_num+1)
        n -= 1

    return sum([i**2 for i in works])

print(solution(4, [4, 3, 3]))
print(solution(1, [2, 1, 2]))
print(solution(3, [1,1]))


import pandas as pd
import numpy as np

print("Loading Train func")
from train_dirichlet_interface import train_gan
print("Finished Loading Train func")

train_ds = pd.read_csv('data/experiment_march_21/train_set.csv', header=None).values
print(train_ds.shape)
ds_size = train_ds.shape[0]

# train_ds = train_ds[0:ds_size,:]

repetitions = 5
batch_size = 256

#partiton_n = [300, 400, 500, 700]

partiton_n = pd.read_csv('random_amounts.csv', header=None).values.flatten()

print(partiton_n)
print(len(partiton_n))

for m in range(len(partiton_n)):

    for k_rep in range(repetitions):

        # partiton_n = np.random.randint(batch_size,ds_size) #At least it has to fit one batch size
        index_list = np.random.randint(0,ds_size, size=partiton_n[m]).tolist()

        print("Calling training function")
        ratio = str(k_rep+1)+"/"+str(repetitions)
        ratio2 = str(m+1)+"/"+str(len(partiton_n))
        telegram_message = " >> Repetition ("+ratio+") of Partiton ("+ratio2+")"

        train_gan(train_ds,  index_list, partiton_n[m], k_rep,telegram_message)
        print("====> Finished repetition "+str(k_rep)+' of partition # '+str(m))

# Snippet to write a partition table
"""
partiton_n = [np.random.randint(batch_size,ds_size) for k in range(10) ]

df = pd.DataFrame(np.array(partiton_n))

save_name = 'random_amounts.csv'
with open(save_name, 'a') as f:
    df.to_csv(f, header=False, index=False)
"""


def propagar(vector):
    # encender fósforos hacia la derecha
    for i, z in enumerate(vector, start=0): # recorremos la lista
        if z == 1:   # si encontramos a e
            try:
                if vector[i+1] != -1:
                    vector[i+1] = 1
                    # pos += 1
            # da error de index por lo que agrego except
            except:
                continue

    # encender fósforos hacia la izquierda
    index = len(vector) - 1
    while index >0:
        if vector[index] == 1 and vector[index-1] != -1:
                vector[index-1] = 1
        index -= 1

    return vector

#%% otra vesion más simple

def propagar(vector):
    '''Función que recibe un vector con 0's, 1's y -1's y devuelve un vector
    en el que los 1's se propagaron a sus vecinos con 0.'''

    # loop para encender fósforos hacia la derecha
    for i, z in enumerate(vector, start=0): # recorremos la lista
        if z == 1:   # si encontramos un fosforo encendido...
            try:
                if vector[i+1] == 0: # ...chequeamos que el fosforo a la derecha este apagado
                    vector[i+1] = 1  # ...si esta apagado lo encendemos
            # da error de index cuando miramos vecino de la ultima posicion
            except IndexError:
                continue

    # loop para encender fósforos hacia la izquierda
    # aprovecho que i está en la última posición luego del loop anterior y lo uso acá
    while i >0:
        if vector[i] == 1 and vector[i-1] == 0: # si encuentra un fosforo encendido y su vecino de la izq está apagado
                vector[i-1] = 1 # enciende el fosforo apagado
        i -= 1 # nos movemos hacia la izquiera

    return vector

#%% Probar funcion y chequeo los resultados con los de la clase
if __name__ == '__main__':
    prueba_1 = propagar([ 0, 0, 0,-1, 1, 0, 0, 0,-1, 0, 1, 0, 0])
    resultado_clase1 = [ 0, 0, 0,-1, 1, 1, 1, 1,-1, 1, 1, 1, 1]
    print(f'{prueba_1}')
    print(prueba_1==resultado_clase1)
    prueba_2 = propagar([ 0, 0, 0, 1, 0, 0])
    resultado_clase2 = [1, 1, 1, 1, 1, 1]
    print(f'{prueba_2}')
    print(prueba_2 == resultado_clase2)


import smtplib
from email.mime.text import MIMEText
class Email(object):
    """docstring for Email"""
    def __init__(self, body='(empty)', subject='(empty)', recipients=None, test=False):
        super(Email, self).__init__()
        self.test = test
        self.subject = subject
        self.body = body
        if recipients is None:
            self.recipients = ['thomasklijnsma@gmail.com']
        else:
            self.recipients = recipients

    def parse_msg(self, recipient):
        msg = MIMEText(self.body, 'html')
        msg['Subject'] = '[ws] ' + self.subject
        msg['From']    = 'tklijnsm.helper@gmail.com'
        msg['To']      = recipient
        return msg

    def send(self):
        for recipient in self.recipients:
            msg = self.parse_msg(recipient)
            if not self.test:
                server = smtplib.SMTP('smtp.gmail.com:587')
                server.ehlo()
                server.starttls()
                server.login(
                    'tklijnsm.helper@gmail.com',
                    'iwanttohelp'
                    )
                server.sendmail( msg['From'], msg['To'], msg.as_string() )
                server.quit()
            else:
                logging.info('TEST MODE: Sending the following email:')
                logging.info(msg)


import numpy as np
import tensorflow as tf
from sklearn.metrics import accuracy_score

(X_train, y_train), (X_test, y_test) = tf.keras.datasets.mnist.load_data()
X_train = X_train[:, ..., np.newaxis] / 255.0
X_test = X_test[:, ..., np.newaxis] / 255.0

model = tf.keras.models.Sequential()
model.add(tf.keras.layers.Input(shape=(28, 28, 1)))
model.add(tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation="relu"))
model.add(tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3), activation="relu"))
model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
model.add(tf.keras.layers.Flatten())
model.add(tf.keras.layers.Dense(units=128, activation="relu"))
model.add(tf.keras.layers.Dense(units=10, activation="softmax"))

model.compile(
    optimizer="Adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"]
)

callbacks = [tf.keras.callbacks.ModelCheckpoint("model_best.h5")]

model.fit(
    X_train, y_train, validation_split=0.2, callbacks=callbacks, epochs=1, batch_size=32
)
model.load_weights("model_best.h5")
y_test_hat = model.predict(X_test)
print(accuracy_score(np.argmax(y_test_hat, axis=1), y_test))


import boto3

BUCKET_NAME = "parwizforogh12"

s3_resource = boto3.resource('s3')

s3_object = s3_resource.Object(BUCKET_NAME, 'file.pdf')

s3_object.download_file('downloaded.pdf')

print("File has been downloaded")


import itertools
import os
from collections import OrderedDict
import torchvision.utils as vutils
import torch
import torch.optim as optim
from torch import nn
from misc.utils import get_inf_iterator, mkdir
from misc import evaluate
from torch.nn import DataParallel
import random
import torch.autograd as autograd
from copy import deepcopy
from itertools import permutations, combinations
import models

from pdb import set_trace as st

def Train(args, FeatExtor, DepthEstor, FeatEmbder,
        data_loader1_real, data_loader1_fake,
        data_loader2_real, data_loader2_fake,
        data_loader3_real, data_loader3_fake,
        data_loader_target,
        summary_writer, Saver, savefilename):

    ####################
    # 1. setup network #
    ####################
    # set train state for Dropout and BN layers
    FeatExtor.train()
    DepthEstor.train()
    FeatEmbder.train()

    FeatExtor = DataParallel(FeatExtor)
    DepthEstor = DataParallel(DepthEstor)

    # setup criterion and optimizer
    criterionCls = nn.BCEWithLogitsLoss()
    criterionDepth = torch.nn.MSELoss()

    if args.optimizer_meta is 'adam':

        optimizer_all = optim.Adam(itertools.chain(FeatExtor.parameters(), DepthEstor.parameters(), FeatEmbder.parameters()),
                                   lr=args.lr_meta,
                                   betas=(args.beta1, args.beta2))

    else:
        raise NotImplementedError('Not a suitable optimizer')

    iternum = max(len(data_loader1_real),len(data_loader1_fake),
                  len(data_loader2_real),len(data_loader2_fake),
                  len(data_loader3_real),len(data_loader3_fake))

    print('iternum={}'.format(iternum))

    ####################
    # 2. train network #
    ####################
    global_step = 0

    for epoch in range(args.epochs):

        data1_real = get_inf_iterator(data_loader1_real)
        data1_fake = get_inf_iterator(data_loader1_fake)

        data2_real = get_inf_iterator(data_loader2_real)
        data2_fake = get_inf_iterator(data_loader2_fake)

        data3_real = get_inf_iterator(data_loader3_real)
        data3_fake = get_inf_iterator(data_loader3_fake)

        for step in range(iternum):

            #============ one batch extraction ============#

            cat_img1_real, depth_img1_real, lab1_real = next(data1_real)
            cat_img1_fake, depth_img1_fake, lab1_fake = next(data1_fake)

            cat_img2_real, depth_img2_real, lab2_real = next(data2_real)
            cat_img2_fake, depth_img2_fake, lab2_fake = next(data2_fake)

            cat_img3_real, depth_img3_real, lab3_real = next(data3_real)
            cat_img3_fake, depth_img3_fake, lab3_fake = next(data3_fake)

            #============ one batch collection ============#

            catimg1 = torch.cat([cat_img1_real,cat_img1_fake],0).cuda()
            depth_img1 = torch.cat([depth_img1_real,depth_img1_fake],0).cuda()
            lab1 = torch.cat([lab1_real,lab1_fake],0).float().cuda()

            catimg2 = torch.cat([cat_img2_real,cat_img2_fake],0).cuda()
            depth_img2 = torch.cat([depth_img2_real,depth_img2_fake],0).cuda()
            lab2 = torch.cat([lab2_real,lab2_fake],0).float().cuda()

            catimg3 = torch.cat([cat_img3_real,cat_img3_fake],0).cuda()
            depth_img3 = torch.cat([depth_img3_real,depth_img3_fake],0).cuda()
            lab3 = torch.cat([lab3_real,lab3_fake],0).float().cuda()

            catimg = torch.cat([catimg1,catimg2,catimg3],0)
            depth_GT = torch.cat([depth_img1,depth_img2,depth_img3],0)
            label = torch.cat([lab1,lab2,lab3],0)

           #============ doamin list augmentation ============#
            catimglist = [catimg1,catimg2,catimg3]
            lablist = [lab1,lab2,lab3]
            deplist = [depth_img1,depth_img2,depth_img3]

            domain_list = list(range(len(catimglist)))
            random.shuffle(domain_list)

            meta_train_list = domain_list[:args.metatrainsize]
            meta_test_list = domain_list[args.metatrainsize:]
            print('metatrn={}, metatst={}'.format(meta_train_list, meta_test_list[0]))

            #============ meta training ============#

            Loss_dep_train = 0.0
            Loss_cls_train = 0.0

            adapted_state_dicts = []

            for index in meta_train_list:

                catimg_meta = catimglist[index]
                lab_meta = lablist[index]
                depGT_meta = deplist[index]

                batchidx = list(range(len(catimg_meta)))
                random.shuffle(batchidx)

                img_rand = catimg_meta[batchidx,:]
                lab_rand = lab_meta[batchidx]
                depGT_rand = depGT_meta[batchidx,:]

                feat_ext_all,feat = FeatExtor(img_rand)
                pred = FeatEmbder(feat)
                depth_Pre = DepthEstor(feat_ext_all)

                Loss_cls = criterionCls(pred.squeeze(), lab_rand)
                Loss_dep = criterionDepth(depth_Pre, depGT_rand)

                Loss_dep_train+=Loss_dep
                Loss_cls_train+=Loss_cls

                zero_param_grad(FeatEmbder.parameters())
                grads_FeatEmbder = torch.autograd.grad(Loss_cls, FeatEmbder.parameters(), create_graph=True)
                fast_weights_FeatEmbder = FeatEmbder.cloned_state_dict()

                adapted_params = OrderedDict()
                for (key, val), grad in zip(FeatEmbder.named_parameters(), grads_FeatEmbder):
                    adapted_params[key] = val - args.meta_step_size * grad
                    fast_weights_FeatEmbder[key] = adapted_params[key]

                adapted_state_dicts.append(fast_weights_FeatEmbder)

            #============ meta testing ============#
            Loss_dep_test = 0.0
            Loss_cls_test = 0.0

            index = meta_test_list[0]

            catimg_meta = catimglist[index]
            lab_meta = lablist[index]
            depGT_meta = deplist[index]

            batchidx = list(range(len(catimg_meta)))
            random.shuffle(batchidx)

            img_rand = catimg_meta[batchidx,:]
            lab_rand = lab_meta[batchidx]
            depGT_rand = depGT_meta[batchidx,:]

            feat_ext_all,feat = FeatExtor(img_rand)
            depth_Pre = DepthEstor(feat_ext_all)
            Loss_dep = criterionDepth(depth_Pre, depGT_rand)

            for n_scr in range(len(meta_train_list)):
                a_dict = adapted_state_dicts[n_scr]

                pred = FeatEmbder(feat, a_dict)
                Loss_cls = criterionCls(pred.squeeze(), lab_rand)

                Loss_cls_test+=Loss_cls

            Loss_dep_test = Loss_dep

            Loss_dep_train_ave = Loss_dep_train/len(meta_train_list)
            Loss_dep_test = Loss_dep_test

            Loss_meta_train =  Loss_cls_train+args.W_depth*Loss_dep_train
            Loss_meta_test =  Loss_cls_test+args.W_depth*Loss_dep_test

            Loss_all = Loss_meta_train + args.W_metatest*Loss_meta_test

            optimizer_all.zero_grad()
            Loss_all.backward()
            optimizer_all.step()

            if (step+1) % args.log_step == 0:
                errors = OrderedDict([
                                    ('Loss_meta_train', Loss_meta_train.item()),
                                    ('Loss_meta_test', Loss_meta_test.item()),
                                    ('Loss_cls_train', Loss_cls_train.item()),
                                    ('Loss_cls_test', Loss_cls_test.item()),
                                    ('Loss_dep_train_ave', Loss_dep_train_ave.item()),
                                    ('Loss_dep_test', Loss_dep_test.item()),
                                    ])
                Saver.print_current_errors((epoch+1), (step+1), errors)

            #============ tensorboard the log info ============#
            info = {
                'Loss_meta_train': Loss_meta_train.item(),
                'Loss_meta_test': Loss_meta_test.item(),
                'Loss_cls_train': Loss_cls_train.item(),
                'Loss_cls_test': Loss_cls_test.item(),
                'Loss_dep_train_ave': Loss_dep_train_ave.item(),
                'Loss_dep_test': Loss_dep_test.item(),
                    }
            for tag, value in info.items():
                summary_writer.add_scalar(tag, value, global_step)

            global_step+=1

            #############################
            # 2.4 save model parameters #
            #############################
            if ((step + 1) % args.model_save_step == 0):
                model_save_path = os.path.join(args.results_path, 'snapshots', savefilename)
                mkdir(model_save_path)

                torch.save(FeatExtor.state_dict(), os.path.join(model_save_path,
                    "FeatExtor-{}-{}.pt".format(epoch+1, step+1)))
                torch.save(FeatEmbder.state_dict(), os.path.join(model_save_path,
                    "FeatEmbder-{}-{}.pt".format(epoch+1, step+1)))
                torch.save(DepthEstor.state_dict(), os.path.join(model_save_path,
                    "DepthEstor-{}-{}.pt".format(epoch+1, step+1)))

        if ((epoch + 1) % args.model_save_epoch == 0):
            model_save_path = os.path.join(args.results_path, 'snapshots', savefilename)
            mkdir(model_save_path)

            torch.save(FeatExtor.state_dict(), os.path.join(model_save_path,
                "FeatExtor-{}.pt".format(epoch+1)))
            torch.save(FeatEmbder.state_dict(), os.path.join(model_save_path,
                "FeatEmbder-{}.pt".format(epoch+1)))
            torch.save(DepthEstor.state_dict(), os.path.join(model_save_path,
                "DepthEstor-{}.pt".format(epoch+1)))

    torch.save(FeatExtor.state_dict(), os.path.join(model_save_path,
        "FeatExtor-final.pt"))
    torch.save(FeatEmbder.state_dict(), os.path.join(model_save_path,
        "FeatEmbder-final.pt"))
    torch.save(DepthEstor.state_dict(), os.path.join(model_save_path,
        "DepthEstor-final.pt"))

def zero_param_grad(params):
    for p in params:
        if p.grad is not None:
            p.grad.zero_()


import os

from setuptools import setup

def rel(*xs):
    return os.path.join(os.path.abspath(os.path.dirname(__file__)), *xs)

with open(rel("README.md")) as f:
    long_description = f.read()

with open(rel("molten", "__init__.py"), "r") as f:
    version_marker = "__version__ = "
    for line in f:
        if line.startswith(version_marker):
            _, version = line.split(version_marker)
            version = version.strip().strip('"')
            break
    else:
        raise RuntimeError("Version marker not found.")

dependencies = [
    "typing-extensions>=3.6,<4.0",
    "typing-inspect>=0.3.1,<0.7",
]

extra_dependencies = {}
extra_dependencies["all"] = dependencies
extra_dependencies["dev"] = extra_dependencies["all"] + [
    # Docs
    "alabaster>0.7",
    "sphinx<1.8",
    "sphinxcontrib-napoleon",

    # Linting
    "flake8",
    "flake8-bugbear",
    "flake8-quotes",
    "isort",
    "mypy",

    # Misc
    "bumpversion>0.5,<0.6",

    # Contrib
    "dramatiq[rabbitmq]>1.3,<2.0",
    "gevent",
    "gunicorn>19.8",
    "jinja2>=2.10,<3.0",
    "msgpack>0.5,<0.6",
    "prometheus-client>=0.2,<0.3",
    "sqlalchemy>1.2,<2.0",
    "toml>0.9,<0.10",
    "wsgicors>=0.7,<0.8",

    # Testing
    "pytest",
]

setup(
    name="molten",
    version=version,
    author="Bogdan Popa",
    author_email="bogdan@cleartype.io",
    description="A minimal, extensible, fast and productive API framework.",
    long_description=long_description,
    long_description_content_type="text/markdown",
    packages=[
        "molten",
        "molten.contrib",
        "molten.http",
        "molten.openapi",
        "molten.openapi.templates",
        "molten.testing",
        "molten.validation",
    ],
    include_package_data=True,
    install_requires=dependencies,
    python_requires=">=3.6",
    extras_require=extra_dependencies,
    classifiers=[
        "Environment :: Web Environment",
        "Intended Audience :: Developers",
        "License :: OSI Approved :: GNU Lesser General Public License v3 or later (LGPLv3+)",
        "Operating System :: OS Independent",
        "Programming Language :: Python :: 3 :: Only",
        "Programming Language :: Python :: 3.6",
        "Programming Language :: Python :: 3.7",
        "Programming Language :: Python",
        "Topic :: Internet :: WWW/HTTP :: Dynamic Content",
        "Topic :: Internet :: WWW/HTTP :: WSGI :: Application",
        "Topic :: Software Development :: Libraries :: Application Frameworks",
        "Topic :: Software Development :: Libraries :: Python Modules",
    ],
)


# -*- coding: utf-8 -*-
import logging
from pathlib import Path
from typing import Optional
from typing import Sequence
from typing import Tuple
from typing import Union

import numpy as np

from nata.containers import GridArray
from nata.containers import ParticleArray
from nata.utils.container_tools import register_backend
from nata.utils.zdf import info
from nata.utils.zdf import read

@register_backend(GridArray)
class Osiris_zdf_GridFile:
    name = "osiris_dev_grid_zdf"
    location: Optional[Path] = None

    def __init__(self, location: Union[str, Path]) -> None:
        self.location = location if isinstance(location, Path) else Path(location)

    @staticmethod
    def is_valid_backend(file_path: Union[Path, str]) -> bool:
        if isinstance(file_path, str):
            file_path = Path(file_path)

        if not isinstance(file_path, Path):
            return False

        if not file_path.is_file():
            return False

        if not file_path.suffix == ".zdf":
            return False

        z_info = info(str(file_path))
        if hasattr(z_info, "type"):
            if z_info.type == "grid":
                return True

        return False

    @property
    def _dset_name(self) -> str:
        logging.info(f"Accessing '{self.location}' for '_dset_name'")
        z_info = info(str(self.location))
        label = z_info.grid.label
        return z_info.grid.name or self.clean(label)

    def get_data(self=None, indexing=None):
        # TODO: apply indexing here
        (z_data, z_info) = read(str(self.location))
        return z_data.transpose()

    @property
    def dataset_name(self) -> str:
        logging.info(f"Accessing '{self.location}' for 'dataset_name'")
        z_info = info(str(self.location))
        label = z_info.grid.label
        return z_info.grid.name or self.clean(label)

    @property
    def dataset_label(self) -> str:
        logging.info(f"Accessing '{self.location}' for 'dataset_label'")
        z_info = info(str(self.location))
        return z_info.grid.label

    @property
    def ndim(self):
        logging.info(f"Accessing '{self.location}' for 'ndim'")
        z_info = info(str(self.location))
        return z_info.grid.ndims

    @property
    def shape(self):
        logging.info(f"Accessing '{self.location}' for 'shape'")
        z_info = info(str(self.location))
        nx = []
        for n in z_info.grid.nx:
            nx.append(n.item())
        return tuple(nx)

    @property
    def dtype(self):
        logging.info(f"Accessing '{self.location}' for 'dtype'")
        (z_data, z_info) = read(str(self.location))
        return z_data.dtype

    @property
    def dataset_unit(self):
        logging.info(f"Accessing '{self.location}' for 'dataset_unit'")
        z_info = info(str(self.location))
        return z_info.grid.units

    @property
    def axes_min(self):
        logging.info(f"Accessing '{self.location}' for 'axes_min'")
        min_values = []
        z_info = info(str(self.location))
        for axis in z_info.grid.axis:
            min_values.append(axis.min)
        return np.array(min_values)

    @property
    def axes_max(self):
        logging.info(f"Accessing '{self.location}' for 'axes_max'")
        max_values = []
        z_info = info(str(self.location))
        for axis in z_info.grid.axis:
            max_values.append(axis.max)
        return np.array(max_values)

    @property
    def axes_names(self):
        logging.info(f"Accessing '{self.location}' for 'axes_names'")
        names = []
        z_info = info(str(self.location))
        for axis in z_info.grid.axis:
            label = axis.label
            names.append(axis.name or self.clean(label))
        return np.array(names)

    @property
    def axes_labels(self):
        logging.info(f"Accessing '{self.location}' for 'axes_labels'")
        long_names = []
        z_info = info(str(self.location))
        for axis in z_info.grid.axis:
            long_names.append(axis.label)
        return np.array(long_names)

    @property
    def axes_units(self):
        logging.info(f"Accessing '{self.location}' for 'axes_units'")
        units = []
        z_info = info(str(self.location))
        for axis in z_info.grid.axis:
            units.append(axis.units)
        return np.array(units)

    @property
    def iteration(self):
        logging.info(f"Accessing '{self.location}' for 'iteration'")
        z_info = info(str(self.location))
        return z_info.iteration.n

    @property
    def time_step(self):
        logging.info(f"Accessing '{self.location}' for 'time_step'")
        z_info = info(str(self.location))
        return z_info.iteration.t

    @property
    def time_unit(self):
        logging.info(f"Accessing '{self.location}' for 'time_unit'")
        z_info = info(str(self.location))
        return z_info.iteration.tunits

    def clean(self, name):
        return (
            name.lower()
            .replace("_", "")
            .replace("^", "")
            .replace("{", "")
            .replace("}", "")
            .replace("\\", "")
            .replace(" ", "_")
        )

@register_backend(ParticleArray)
class Osiris_zdf_ParticleFile:
    name = "osiris_dev_particles_zdf"
    location: Optional[Path] = None

    def __init__(self, location: Union[str, Path]) -> None:
        self.location = location if isinstance(location, Path) else Path(location)

    @staticmethod
    def is_valid_backend(file_path: Path) -> bool:
        if not file_path.is_file():
            return False

        if not file_path.suffix == ".zdf":
            return False

        z_info = info(str(file_path))
        if hasattr(z_info, "type"):
            if z_info.type == "particles":
                return True

        return False

    @property
    def dataset_name(self) -> str:
        logging.info(f"Accessing '{self.location}' for 'dataset_name'")
        z_info = info(str(self.location))
        return z_info.particles.name

    @property
    def dataset_label(self) -> str:
        logging.info(f"Accessing '{self.location}' for 'dataset_label'")
        z_info = info(str(self.location))
        return z_info.particles.label

    @property
    def num_particles(self) -> int:
        logging.info(f"Accessing '{self.location}' for 'num_particles'")
        z_info = info(str(self.location))
        return z_info.particles.nparts

    def get_data(self, indexing=None, fields=None) -> np.ndarray:
        logging.info(f"Reading data in '{self.location}'")
        (z_data, z_info) = read(str(self.location))
        if fields is None:
            # create a structured array
            dset = np.empty(self.num_particles, dtype=self.dtype)

            # fill the array
            for quant in self.quantity_names:
                dset[quant] = z_data[quant]
        else:
            if indexing is None:
                dset = z_data[fields][:]
            else:
                dset = z_data[fields][indexing]

        return dset

    @property
    def quantity_names(self) -> Sequence[str]:
        logging.info(f"Accessing '{self.location}' for 'quantity_names'")
        z_info = info(str(self.location))
        quantities = []
        for key in z_info.particles.quants:
            if key == "tag":
                continue
            quantities.append(key)

        return quantities

    @property
    def quantity_labels(self) -> Sequence[str]:
        logging.info(f"Accessing '{self.location}' for 'quantity_labels'")
        z_info = info(str(self.location))
        names = []
        for quant in self.quantity_names:
            names.append(z_info.particles.qlabels[quant])
        return names

    @property
    def quantity_units(self) -> Sequence[str]:
        logging.info(f"Accessing '{self.location}' for 'quantity_units'")
        z_info = info(str(self.location))
        units = []
        for quant in self.quantity_names:
            units.append(z_info.particles.qunits[quant])

        return units

    @property
    def dtype(self) -> np.dtype:
        logging.info(f"Accessing '{self.location}' for 'dtype'")
        (z_data, z_info) = read(str(self.location))
        fields = []
        for quant in self.quantity_names:
            fields.append((quant, z_data[quant].dtype))

        return np.dtype(fields)

    @property
    def ndim(self) -> int:
        raise NotImplementedError

    @property
    def shape(self) -> Tuple[int, ...]:
        raise NotImplementedError

    @property
    def iteration(self) -> int:
        logging.info(f"Accessing '{self.location}' for 'iteration'")
        z_info = info(str(self.location))
        return z_info.iteration.n

    @property
    def time_step(self) -> float:
        logging.info(f"Accessing '{self.location}' for 'time_step'")
        z_info = info(str(self.location))
        return z_info.iteration.t

    @property
    def time_unit(self) -> str:
        logging.info(f"Accessing '{self.location}' for 'time_unit'")
        z_info = info(str(self.location))
        return z_info.iteration.tunits


"""
@file hough_lines.py
@brief This program demonstrates line finding with the Hough transform
"""
import sys
import math
import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

def afstand(c0, c1, c2, maxlen, hoek):
    if c0[0] == c1[0]:
        alpha1 = np.pi * np.sign(c1[1] - c0[1])
    else:
        t = (c1[0] - c0[0]) / (c1[1] - c0[1])
        alpha1 = np.arctan(t)

    if c2[0] == c1[0]:
        alpha2 = np.pi * np.sign(c1[1] - c2[1])
    else:
        t = (c1[0] - c2[0]) / (c1[1] - c2[1])
        alpha2 = np.arctan(t)

    alpha = alpha1 - alpha2
    if abs(alpha) > hoek:
        return maxlen + 5

    else:
        a1 = c1[0] - c2[0]
        a2 = c1[1] - c2[1]
        afst = math.sqrt(a1 ** 2 + a2 ** 2)
        return afst

def main(argv):
    ## [load]
    default_file = 'paashaas.jpg'
    maxlen = 7
    maxhoek = np.pi/2  # in radialen

    offset_y = 200 + 200
    offset_x = 180

    filename = argv[0] if len(argv) > 0 else default_file

    # Loads an image
    src = cv.imread(cv.samples.findFile(filename), 0)  # cv.IMREAD_GRAYSCALE)

    # Check if image is loaded fine
    if src is None:
        print('Error opening image!')
        print('Usage: hough_lines.py [image_name -- default ' + default_file + '] \n')
        return -1
    ## [load]

    ## [edge_detection]
    # Edge detection
    dst = cv.Canny(src, 50, 200, None, 3)
    ## [edge_detection]

    # Copy edges to the images that will display the results in BGR
    cdst = cv.cvtColor(dst, cv.COLOR_GRAY2BGR)
    cdstP = np.copy(cdst)

    ## [hough_lines]
    #  Standard Hough Line Transform
    lines = cv.HoughLines(dst, 1, np.pi / 180, 150, None, 0, 0)
    ## [hough_lines]

    ## [draw_lines]
    # Draw the lines
    if lines is not None:
        for i in range(0, len(lines)):
            rho = lines[i][0][0]
            theta = lines[i][0][1]
            a = math.cos(theta)
            b = math.sin(theta)
            x0 = a * rho
            y0 = b * rho
            pt1 = (int(x0 + 1000 * (-b)), int(y0 + 1000 * (a)))
            pt2 = (int(x0 - 1000 * (-b)), int(y0 - 1000 * (a)))

            cv.line(cdst, pt1, pt2, (0, 0, 255), 3, cv.LINE_AA)
    ## [draw_lines]

    ## [hough_lines_p]
    # Probabilistic Line Transform
    linesP = cv.HoughLinesP(dst, 1, np.pi / 2, 1, None, 0, 0)
    ## [hough_lines_p]

    ## [draw_lines_p]
    # Draw the lines
    coorda = []  # array van coordinaten
    if linesP is not None:
        File1 = open("tabel2.txt", "w")
        for i in range(0, len(linesP)):
            l = linesP[i][0]
            cv.line(cdstP, (l[0], l[1]), (l[2], l[3]), (0, 0, 255), 3, cv.LINE_AA)
            coorda.append([[l[0] + offset_x, -l[1] + offset_y], [l[2] + offset_x, -l[3] + offset_y]])
            File1.write("[" + str(l[0]) + " " + str(l[1]) + "]" + "\n")
        File1.close()

    coorda = sorted(coorda)  # array of coordinates

    groepnr = 0
    coordgroepen = [coorda[0]]  # array of coordinates

    while len(coorda) > 0:
        afstanden = []
        weg = []
        for i in coorda:
            if len(coordgroepen[groepnr]) >= 3:
                afst = afstand(coordgroepen[groepnr][-3], coordgroepen[groepnr][-1], i[1], maxlen, maxhoek)
            else:
                afst = afstand(coordgroepen[groepnr][1], coordgroepen[groepnr][-1], i[1], maxlen, maxhoek)
            if afst < maxlen:
                    afstanden.append([afst, i])

        if len(afstanden) == 0:
            if len(coorda) > 0:
                coordgroepen.extend([coorda[0]])
            groepnr += 1
        else:
            afstanden.sort()
            coordgroepen[groepnr].extend(afstanden[0][1])
            coorda.remove(afstanden[0][1])

    print(groepnr)

    plt.axis('off')
    plt.axis('equal')
    for i in coordgroepen:
        xs = [x[0] for x in i]
        ys = [x[1] for x in i]
        plt.plot(xs, ys)
    # plt.savefig('dobot\\dobot_min' + minlen.__str__() + '_max' + maxlen.__str__() + '_' + default_file)
    plt.show()

    ## [draw_lines_p]

    ## [imshow]
    # Show results
    # cv.imshow("Source", src)
    # cv.imshow("Detected Lines (in red) - Standard Hough Line Transform", cdst)
    cv.imshow("Detected Lines (in red) - Probabilistic Line Transform", cdstP)

    ## [imshow]
    ## [exit]
    # Wait and Exit
    cv.waitKey()

    return 0
    ## [exit]

if __name__ == "__main__":
    main(sys.argv[1:])


import os
import re
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import glob
from PIL import Image
from pathlib import Path
from torch import distributions as pyd
from torch.distributions.utils import _standard_normal

__REDUCE__ = lambda b: "mean" if b else "none"

def l1(pred, target, reduce=False):
    """Computes the L1-loss between predictions and targets."""
    return F.l1_loss(pred, target, reduction=__REDUCE__(reduce))

def mse(pred, target, reduce=False):
    """Computes the MSE loss between predictions and targets."""
    return F.mse_loss(pred, target, reduction=__REDUCE__(reduce))

def _get_out_shape(in_shape, layers):
    """Utility function. Returns the output shape of a network for a given input shape."""
    x = torch.randn(*in_shape).unsqueeze(0)
    return (
        (nn.Sequential(*layers) if isinstance(layers, list) else layers)(x)
        .squeeze(0)
        .shape
    )

def orthogonal_init(m):
    """Orthogonal layer initialization."""
    if isinstance(m, nn.Linear):
        nn.init.orthogonal_(m.weight.data)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Conv2d):
        gain = nn.init.calculate_gain("relu")
        nn.init.orthogonal_(m.weight.data, gain)
        if m.bias is not None:
            nn.init.zeros_(m.bias)

def soft_update_params(m, m_target, tau):
    """Update slow-moving average of online network (target network) at rate tau."""
    with torch.no_grad():
        for p, p_target in zip(m.parameters(), m_target.parameters()):
            p_target.data.lerp_(p.data, tau)

def set_requires_grad(net, value):
    """Enable/disable gradients for a given (sub)network."""
    for param in net.parameters():
        param.requires_grad_(value)

class TruncatedNormal(pyd.Normal):
    """Utility class implementing the truncated normal distribution."""

    def __init__(self, loc, scale, low=-1.0, high=1.0, eps=1e-6):
        super().__init__(loc, scale, validate_args=False)
        self.low = low
        self.high = high
        self.eps = eps

    def _clamp(self, x):
        clamped_x = torch.clamp(x, self.low + self.eps, self.high - self.eps)
        x = x - x.detach() + clamped_x.detach()
        return x

    def sample(self, clip=None, sample_shape=torch.Size()):
        shape = self._extended_shape(sample_shape)
        eps = _standard_normal(shape, dtype=self.loc.dtype, device=self.loc.device)
        eps *= self.scale
        if clip is not None:
            eps = torch.clamp(eps, -clip, clip)
        x = self.loc + eps
        return self._clamp(x)

class NormalizeImg(nn.Module):
    """Module that divides (pixel) observations by 255."""

    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x.div(255.0)

class Flatten(nn.Module):
    """Module that flattens its input to a (batched) vector."""

    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x.view(x.size(0), -1)

def enc(cfg):
    """Returns our MoDem encoder that takes a stack of 224x224 frames as input."""
    C = int(3 * cfg.frame_stack)
    layers = [
        NormalizeImg(),
        nn.Conv2d(C, cfg.num_channels, 7, stride=2),
        nn.ReLU(),
        nn.Conv2d(cfg.num_channels, cfg.num_channels, 5, stride=2),
        nn.ReLU(),
        nn.Conv2d(cfg.num_channels, cfg.num_channels, 3, stride=2),
        nn.ReLU(),
        nn.Conv2d(cfg.num_channels, cfg.num_channels, 3, stride=2),
        nn.ReLU(),
        nn.Conv2d(cfg.num_channels, cfg.num_channels, 3, stride=2),
        nn.ReLU(),
        nn.Conv2d(cfg.num_channels, cfg.num_channels, 3, stride=2),
        nn.ReLU(),
    ]
    out_shape = _get_out_shape((C, cfg.img_size, cfg.img_size), layers)
    layers.extend([Flatten(), nn.Linear(np.prod(out_shape), cfg.latent_dim)])
    return nn.Sequential(*layers)

def state_enc(cfg):
    """Returns a proprioceptive state encoder + modality fuse."""
    return (
        nn.Sequential(
            nn.Linear(cfg.state_dim, cfg.enc_dim),
            nn.ELU(),
            nn.Linear(cfg.enc_dim, cfg.latent_dim),
        ),
        nn.Sequential(
            nn.Linear(cfg.latent_dim, cfg.enc_dim),
            nn.ELU(),
            nn.Linear(cfg.enc_dim, cfg.latent_dim),
        ),
    )

def mlp(in_dim, mlp_dim, out_dim, act_fn=nn.ELU()):
    """Returns an MLP."""
    if isinstance(mlp_dim, int):
        mlp_dim = [mlp_dim, mlp_dim]
    return nn.Sequential(
        nn.Linear(in_dim, mlp_dim[0]),
        act_fn,
        nn.Linear(mlp_dim[0], mlp_dim[1]),
        act_fn,
        nn.Linear(mlp_dim[1], out_dim),
    )

def q(cfg, act_fn=nn.ELU()):
    """Returns a Q-function that uses Layer Normalization."""
    return nn.Sequential(
        nn.Linear(cfg.latent_dim + cfg.action_dim, cfg.mlp_dim),
        nn.LayerNorm(cfg.mlp_dim),
        nn.Tanh(),
        nn.Linear(cfg.mlp_dim, cfg.mlp_dim),
        act_fn,
        nn.Linear(cfg.mlp_dim, 1),
    )

class RandomShiftsAug(nn.Module):
    """
    Random shift image augmentation.
    Adapted from https://github.com/facebookresearch/drqv2
    """

    def __init__(self, cfg):
        super().__init__()
        self.pad = int(
            cfg.img_size / 21
        )  # maintain same padding ratio as in original implementation

    def forward(self, x):
        n, _, h, w = x.size()
        assert h == w
        padding = tuple([self.pad] * 4)
        x = F.pad(x, padding, "replicate")
        eps = 1.0 / (h + 2 * self.pad)
        arange = torch.linspace(
            -1.0 + eps, 1.0 - eps, h + 2 * self.pad, device=x.device, dtype=x.dtype
        )[:h]
        arange = arange.unsqueeze(0).repeat(h, 1).unsqueeze(2)
        base_grid = torch.cat([arange, arange.transpose(1, 0)], dim=2)
        base_grid = base_grid.unsqueeze(0).repeat(n, 1, 1, 1)
        shift = torch.randint(
            0, 2 * self.pad + 1, size=(n, 1, 1, 2), device=x.device, dtype=x.dtype
        )
        shift *= 2.0 / (h + 2 * self.pad)
        grid = base_grid + shift
        return F.grid_sample(x, grid, padding_mode="zeros", align_corners=False)

class Episode(object):
    """Storage object for a single episode."""

    def __init__(self, cfg, init_obs, init_state=None):
        self.cfg = cfg
        self.device = torch.device("cpu")
        self.obs = torch.empty(
            (cfg.episode_length + 1, *init_obs.shape),
            dtype=torch.uint8,
            device=self.device,
        )
        self.obs[0] = torch.tensor(init_obs, dtype=torch.uint8, device=self.device)
        self.state = torch.empty(
            (cfg.episode_length + 1, *init_state.shape),
            dtype=torch.float32,
            device=self.device,
        )
        self.state[0] = torch.tensor(
            init_state, dtype=torch.float32, device=self.device
        )
        self.action = torch.empty(
            (cfg.episode_length, cfg.action_dim),
            dtype=torch.float32,
            device=self.device,
        )
        self.reward = torch.empty(
            (cfg.episode_length,), dtype=torch.float32, device=self.device
        )
        self.cumulative_reward = 0
        self.done = False
        self._idx = 0

    @classmethod
    def from_trajectory(cls, cfg, obs, states, action, reward, done=None):
        """Constructs an episode from a trajectory."""
        episode = cls(cfg, obs[0], states[0])
        episode.obs[1:] = torch.tensor(
            obs[1:], dtype=episode.obs.dtype, device=episode.device
        )
        episode.state[1:] = torch.tensor(
            states[1:], dtype=episode.state.dtype, device=episode.device
        )
        episode.action = torch.tensor(
            action, dtype=episode.action.dtype, device=episode.device
        )
        episode.reward = torch.tensor(
            reward, dtype=episode.reward.dtype, device=episode.device
        )
        episode.cumulative_reward = torch.sum(episode.reward)
        episode.done = True
        episode._idx = cfg.episode_length
        return episode

    def __len__(self):
        return self._idx

    @property
    def first(self):
        return len(self) == 0

    def __add__(self, transition):
        self.add(*transition)
        return self

    def add(self, obs, state, action, reward, done):
        self.obs[self._idx + 1] = torch.tensor(
            obs, dtype=self.obs.dtype, device=self.obs.device
        )
        self.state[self._idx + 1] = torch.tensor(
            state, dtype=self.state.dtype, device=self.state.device
        )
        self.action[self._idx] = action
        self.reward[self._idx] = reward
        self.cumulative_reward += reward
        self.done = done
        self._idx += 1

def get_demos(cfg):
    fps = glob.glob(str(Path(cfg.demo_dir) / "demonstrations" / f"{cfg.task}/*.pt"))
    episodes = []
    for fp in fps:
        data = torch.load(fp)
        frames_dir = Path(os.path.dirname(fp)) / "frames"
        assert frames_dir.exists(), "No frames directory found for {}".format(fp)
        frame_fps = [frames_dir / fn for fn in data["frames"]]
        obs = np.stack([np.array(Image.open(fp)) for fp in frame_fps]).transpose(
            0, 3, 1, 2
        )
        state = torch.tensor(data["states"], dtype=torch.float32)
        if cfg.task.startswith("mw-"):
            state = torch.cat((state[:, :4], state[:, 18 : 18 + 4]), dim=-1)
        elif cfg.task.startswith("adroit-"):
            if cfg.task == "adroit-door":
                state = np.concatenate([state[:, :27], state[:, 29:32]], axis=1)
            elif cfg.task == "adroit-hammer":
                state = state[:, :36]
            elif cfg.task == "adroit-pen":
                state = np.concatenate([state[:, :24], state[:, -9:-6]], axis=1)
            else:
                raise NotImplementedError()
        actions = np.array(data["actions"], dtype=np.float32).clip(-1, 1)
        if cfg.task.startswith("mw-") or cfg.task.startswith("adroit-"):
            rewards = (
                np.array(
                    [
                        _data[
                            "success" if "success" in _data.keys() else "goal_achieved"
                        ]
                        for _data in data["infos"]
                    ],
                    dtype=np.float32,
                )
                - 1.0
            )
        else:  # use dense rewards for DMControl
            rewards = np.array(data["rewards"])
        episode = Episode.from_trajectory(cfg, obs, state, actions, rewards)
        episodes.append(episode)
    return episodes

class ReplayBuffer(object):
    """
    Storage and sampling functionality for training MoDem.
    Uses prioritized experience replay by default.
    """

    def __init__(self, cfg):
        self.cfg = cfg
        self.device = torch.device("cpu")
        self.capacity = 2 * cfg.train_steps + 1
        obs_shape = (3, *cfg.obs_shape[-2:])
        self._state_dim = cfg.state_dim
        self._obs = torch.empty(
            (self.capacity + 1, *obs_shape), dtype=torch.uint8, device=self.device
        )
        self._last_obs = torch.empty(
            (self.capacity // cfg.episode_length, *cfg.obs_shape),
            dtype=torch.uint8,
            device=self.device,
        )
        self._action = torch.empty(
            (self.capacity, cfg.action_dim), dtype=torch.float32, device=self.device
        )
        self._reward = torch.empty(
            (self.capacity,), dtype=torch.float32, device=self.device
        )
        self._state = torch.empty(
            (self.capacity, self._state_dim), dtype=torch.float32, device=self.device
        )
        self._last_state = torch.empty(
            (self.capacity // cfg.episode_length, self._state_dim),
            dtype=torch.float32,
            device=self.device,
        )
        self._priorities = torch.ones(
            (self.capacity,), dtype=torch.float32, device=self.device
        )
        self._eps = 1e-6
        self.idx = 0

    def __len__(self):
        return self.idx

    def __add__(self, episode: Episode):
        self.add(episode)
        return self

    def add(self, episode: Episode):
        obs = episode.obs[:-1, -3:]
        if episode.obs.shape[1] == 3:
            last_obs = episode.obs[-self.cfg.frame_stack :].view(
                self.cfg.frame_stack * 3, *self.cfg.obs_shape[-2:]
            )
        else:
            last_obs = episode.obs[-1]
        self._obs[self.idx : self.idx + self.cfg.episode_length] = obs
        self._last_obs[self.idx // self.cfg.episode_length] = last_obs
        self._action[self.idx : self.idx + self.cfg.episode_length] = episode.action
        self._reward[self.idx : self.idx + self.cfg.episode_length] = episode.reward
        states = torch.tensor(episode.state, dtype=torch.float32)
        self._state[
            self.idx : self.idx + self.cfg.episode_length, : self._state_dim
        ] = states[:-1]
        self._last_state[
            self.idx // self.cfg.episode_length, : self._state_dim
        ] = states[-1]
        max_priority = (
            1.0
            if self.idx == 0
            else self._priorities[: self.idx].max().to(self.device).item()
        )
        mask = (
            torch.arange(self.cfg.episode_length)
            >= self.cfg.episode_length - self.cfg.horizon
        )
        new_priorities = torch.full(
            (self.cfg.episode_length,), max_priority, device=self.device
        )
        new_priorities[mask] = 0
        self._priorities[self.idx : self.idx + self.cfg.episode_length] = new_priorities
        self.idx = (self.idx + self.cfg.episode_length) % self.capacity

    def update_priorities(self, idxs, priorities):
        self._priorities[idxs] = priorities.squeeze(1).to(self.device) + self._eps

    def _get_obs(self, arr, idxs):
        obs = torch.empty(
            (self.cfg.batch_size, self.cfg.frame_stack * 3, *arr.shape[-2:]),
            dtype=arr.dtype,
            device=torch.device("cuda"),
        )
        obs[:, -3:] = arr[idxs].cuda(non_blocking=True)
        _idxs = idxs.clone()
        mask = torch.ones_like(_idxs, dtype=torch.bool)
        for i in range(1, self.cfg.frame_stack):
            mask[_idxs % self.cfg.episode_length == 0] = False
            _idxs[mask] -= 1
            obs[:, -(i + 1) * 3 : -i * 3] = arr[_idxs].cuda(non_blocking=True)
        return obs.float()

    def sample(self):
        probs = self._priorities[: self.idx] ** self.cfg.per_alpha
        probs /= probs.sum()
        total = len(probs)
        idxs = torch.from_numpy(
            np.random.choice(
                total, self.cfg.batch_size, p=probs.cpu().numpy(), replace=True
            )
        ).to(self.device)
        weights = (total * probs[idxs]) ** (-self.cfg.per_beta)
        weights /= weights.max()
        obs = (
            self._get_obs(self._obs, idxs)
            if self.cfg.frame_stack > 1
            else self._obs[idxs].cuda(non_blocking=True)
        )
        next_obs_shape = (3 * self.cfg.frame_stack, *self._last_obs.shape[-2:])
        next_obs = torch.empty(
            (self.cfg.horizon + 1, self.cfg.batch_size, *next_obs_shape),
            dtype=obs.dtype,
            device=obs.device,
        )
        action = torch.empty(
            (self.cfg.horizon + 1, self.cfg.batch_size, *self._action.shape[1:]),
            dtype=torch.float32,
            device=self.device,
        )
        reward = torch.empty(
            (self.cfg.horizon + 1, self.cfg.batch_size),
            dtype=torch.float32,
            device=self.device,
        )
        state = self._state[idxs, : self._state_dim]
        next_state = torch.empty(
            (self.cfg.horizon + 1, self.cfg.batch_size, *state.shape[1:]),
            dtype=state.dtype,
            device=state.device,
        )
        for t in range(self.cfg.horizon + 1):
            _idxs = idxs + t
            next_obs[t] = (
                self._get_obs(self._obs, _idxs + 1)
                if self.cfg.frame_stack > 1
                else self._obs[_idxs + 1].cuda(non_blocking=True)
            )
            action[t] = self._action[_idxs]
            reward[t] = self._reward[_idxs]
            next_state[t] = self._state[_idxs + 1, : self._state_dim]

        mask = (_idxs + 1) % self.cfg.episode_length == 0
        next_obs[-1, mask] = (
            self._last_obs[_idxs[mask] // self.cfg.episode_length]
            .to(next_obs.device, non_blocking=True)
            .float()
        )
        state = state.cuda(non_blocking=True)
        next_state[-1, mask] = (
            self._last_state[_idxs[mask] // self.cfg.episode_length, : self._state_dim]
            .to(next_state.device)
            .float()
        )
        next_state = next_state.cuda(non_blocking=True)
        next_obs = next_obs.cuda(non_blocking=True)
        action = action.cuda(non_blocking=True)
        reward = reward.unsqueeze(2).cuda(non_blocking=True)
        idxs = idxs.cuda(non_blocking=True)
        weights = weights.cuda(non_blocking=True)

        return obs, next_obs, action, reward, state, next_state, idxs, weights

def linear_schedule(schdl, step):
    """Outputs values following a linear decay schedule"""
    try:
        return float(schdl)
    except ValueError:
        match = re.match(r"linear\((.+),(.+),(.+)\)", schdl)
        if match:
            init, final, duration = [float(g) for g in match.groups()]
            mix = np.clip(step / duration, 0.0, 1.0)
            return (1.0 - mix) * init + mix * final
    raise NotImplementedError(schdl)


import cv2
from PIL import Image, ImageChops
import pytesseract
import caesar

def main():
    '''
        PARTE I: MANIPULACION DE IMAGEN Y CAPTURA DE MENSAJE SECRETO
    '''
    img = Image.open('matroska.png')
    format = img.format
    size = img.size
    mode = img.mode

    print("Image Format: ", format, "Image Size: ", size, "Image Mode: ", mode)

    img_width = size[0]
    img_height = size[1]

    img_bitmap = img.load()

    # Convierte el archivo en blanco y negro

    print("filtering image...")
    for i in range(img_width):
        for j in range(img_height):
            old_pixel = img_bitmap[i, j]
            red = old_pixel[0]
            green = old_pixel[1]
            blue = old_pixel[2]

            new_pixel = pixel_filter(red, green, blue)

            img_bitmap[i, j] = new_pixel

    img.save('matroska-filtered.png')

    print("done...")

    message_img = cv2.imread('matroska-filtered-message.png')
    ciphered_text = message_extract(message_img)

    print("Ciphered Message: ", ciphered_text)

    ciphered_text_lines = ciphered_text.split("{!@#$%}")
    ciphered_text_lines.pop()
    ciphered_text_lines[:] = [item for item in ciphered_text_lines if item]
    length = len(ciphered_text_lines)

    for i in range(length):
        ciphered_text_lines[i] = ciphered_text_lines[i].replace(" ", "")

    print("Lines in text: ", length)
    # for line in ciphered_text_lines:
    #     print(line + "\n")

    with open("ciphered-message.txt", "w") as cFile:
        for line in ciphered_text_lines:
            cFile.write(line + "\n")

    '''
        PARTE II: PROCESO DE DESENCRIPCION DEL MENSAJE

    '''

    # text = caesar.decrypt("EXXEGOexsrgi", 4)
    # print("Decripted message: ", text)

    for key in range(1, 51):
        # print("Key: ", key)
        with open("decoded-message" + str(key) + ".txt", "w") as dFile:
            for line in ciphered_text_lines:
                text = caesar.decrypt(line, key)
                dFile.write(text + "\n")

    deciphered_text = []

    with open("decoded-message31.txt", "r") as rFile:
        line = rFile.readline()
        while line:
            deciphered_text.append(line)
            line = rFile.readline()

    length = len(deciphered_text)
    for i in range(length):
        deciphered_text[i] = deciphered_text[i].replace("X", " ")
        deciphered_text[i] = deciphered_text[i].replace("=", "")
        print(deciphered_text[i] + "\n")

def pixel_filter(red, green, blue):
    new_red = 0
    new_green = 0
    new_blue = 0

    if(red % 2 == 0):
        new_red = 1
    else:
        new_red = 0

    if(green % 2 == 0):
        new_green = 1
    else:
        new_green = 0

    if(blue % 2 == 0):
        new_blue = 1
    else:
        new_blue = 0

    return (new_red, new_green, new_blue)

# returns Image object with just the text
def message_extract(image):
    text = pytesseract.image_to_string(image)

    text = text.replace("\n", "{!@#$%}")

    result = text

    return result

if __name__ == "__main__":
    main()


import sockslib
import socket
import threading
import select
import boto3
import random
import json
import argparse
import interceptors

import LaunchNewSOCKS5

from loguru import logger

ec2 = boto3.resource('ec2',region_name='us-west-1')

class ForwardServer:
    def __init__(self, fwd_host, fwd_port, proxy=None, _interceptors=[], host='127.0.0.1', port=1234) -> None:
        self.proxy_addr = (proxy['ip'], proxy['port']) if proxy else None
        self.proxy_authentication = [sockslib.NoAuth(),
                                     sockslib.UserPassAuth(proxy['username'], proxy['password'])] if proxy else None

        self.bind_addr = (host, port)
        self.fwd_addr  = (fwd_host, fwd_port)

        self.server_socket = None
        self.interceptors = _interceptors

        self.client_threads = []
        self.client_signals = []

    def __safe_shutdown(self, sock):
        try:
            sock.shutdown()
        except:
            pass

        try:
            sock.close()
        except:
            pass

    def __forward(self, sock1, sock2, stop_event, client_name, _interceptors=[]):
        sock1.setblocking(False)

        while True:
            try:
                ready = select.select([sock1], [], [], 10)
                if ready[0]:
                    data = sock1.recv(4096)
                    if not data:
                        break

                    for interceptor in _interceptors:
                        data = interceptor(data)

                    sock2.sendall(data)
                if stop_event.is_set():
                    break
            except (socket.error, ValueError) as e:
                stop_event.set()
                break
            # except Exception as e:
            #     logger.error(f"[{client_name}] Forwarding stopped: {e}")
            #     stop_event.set()
            #     break

        self.__safe_shutdown(sock1)
        self.__safe_shutdown(sock2)

        stop_event.set()

    def __handle_client(self, conn, addr):
        if self.proxy_addr is None:
            client_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        else:
            client_sock = sockslib.SocksSocket()
            client_sock.set_proxy(self.proxy_addr,
                                sockslib.Socks.SOCKS5,
                                self.proxy_authentication)
        client_name = f"{addr[0]}:{addr[1]}"

        logger.info(f"[{client_name}] Connecting to {self.fwd_addr[0]}:{self.fwd_addr[1]}")
        client_sock.settimeout(10)

        try:
            client_sock.connect(self.fwd_addr)
        except TimeoutError:
            logger.info(f"[{client_name}] Connection timed out.")

            self.__safe_shutdown(client_sock)
            self.__safe_shutdown(conn)

            return
        except Exception as e:
            logger.error(f"[{client_name}] Connection failed: {e}")

            self.__safe_shutdown(client_sock)
            self.__safe_shutdown(conn)

            return

        logger.info(f"[{client_name}] Connected to {self.fwd_addr[0]}:{self.fwd_addr[1]}")

        serverbound_interceptors = []
        clientbound_interceptors = []
        if len(self.interceptors) > 0:
            logger.info(f"[{client_name}] Starting interceptors...")

            active_interceptors = []
            for interceptor in self.interceptors:
                active_interceptor = interceptor[0](*interceptor[1])
                active_interceptor.setup(conn, client_sock, client_name)
                active_interceptors.append(active_interceptor)

            serverbound_interceptors = [interceptor.handle_serverbound_chunk for interceptor in active_interceptors]
            clientbound_interceptors = [interceptor.handle_clientbound_chunk for interceptor in active_interceptors]

        stop_event = threading.Event()
        client_thread = threading.Thread(target = self.__forward,
                                        args    = [conn, client_sock, stop_event, client_name, serverbound_interceptors],
                                        daemon  = True)
        client_thread.start()

        server_thread = threading.Thread(target = self.__forward,
                                        args    = [client_sock, conn, stop_event, client_name, clientbound_interceptors],
                                        daemon  = True)
        server_thread.start()

        self.client_threads.append(client_thread)
        self.client_threads.append(server_thread)
        self.client_signals.append(stop_event)

        client_thread.join()
        server_thread.join()

        self.__safe_shutdown(client_sock)
        self.__safe_shutdown(conn)

        logger.info(f"[{client_name}] Connection closed")

    def __run_server(self):
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True)
        self.server_socket.bind(self.bind_addr)
        self.server_socket.listen()

        logger.info(f"Listening on {self.bind_addr[0]}:{self.bind_addr[1]}")

        while True:
            conn, addr = self.server_socket.accept()
            client_thread = threading.Thread(target = self.__handle_client,
                                            args    = [conn, addr],
                                            daemon  = True)
            self.client_threads.append(client_thread)
            client_thread.start()

    def start(self):
        self.__run_server()

    def stop(self):
        for signal in self.client_signals:
            signal.set()
        for thread in self.client_threads:
            thread.join()
        self.server_socket.close()

def load_available_proxies():
    with open('proxy_servers.json') as json_file:
        data = json.load(json_file)
    return data

def regen_proxies(num_proxies):
    logger.info(f"Spawning {num_proxies} replacement proxie(s)...")
    LaunchNewSOCKS5.launch(num_proxies)
    logger.info(f"Succesfully spawned {num_proxies} proxie(s)")

def destroy_proxy(available_proxies, proxy):
    ec2.instances.filter(InstanceIds=[proxy['instance_id']]).terminate()
    available_proxies.remove(proxy)
    with open('proxy_servers.json', 'w') as outfile:
        json.dump(available_proxies, outfile)

def main():
    parser = argparse.ArgumentParser(prog="ForwardProxy",
                                     description='Acts as a forward proxy, optional (primary) SOCKS5 transport with automatic EC2 integration for increased anonymity.')

    parser.add_argument('-d', '--disable-proxy', dest='proxy', action='store_false', help='Disable use of a SOCKS5 proxy for the connection to the target server')
    parser.add_argument('remote_addr', metavar='remote_addr', type=str, help='The address of the target server')
    parser.add_argument('remote_port', metavar='remote_port', type=int, help='The port of the target server')
    parser.add_argument('--bind-addr', metavar='Bind Address', type=str, help='The address to bind the proxy server to', default='0.0.0.0')
    parser.add_argument('--bind-port', metavar='Bind Port', type=int, help='The port to bind the proxy server to', default=1234)
    parser.add_argument('-i', '--ideal', dest='ideal', type=int, help='The ideal number of proxies to have available', default=2)

    args = parser.parse_args()

    try:
        fwd_server = None
        proxy = None
        available_proxies = None

        if args.proxy:
            available_proxies = load_available_proxies()
        waitthreads = []

        if args.proxy:
            if len(available_proxies) == 0:
                logger.warning(f"Proxy deficit detected")
                spawn_thread = threading.Thread(target=regen_proxies, args=[args.ideal], daemon=True)

                spawn_thread.start()
                spawn_thread.join()

                available_proxies = load_available_proxies()

            if len(available_proxies) < args.ideal:
                spawn_thread = threading.Thread(target=regen_proxies, args=[1])
                waitthreads.append(spawn_thread)
                spawn_thread.start()

        if args.proxy:
            proxy = random.choice(available_proxies)
            logger.info(f"Using proxy {proxy['ip']}:{proxy['port']}")

        logger.info("Starting server...")
        fwd_server = ForwardServer(args.remote_addr, args.remote_port, proxy, [[interceptors.MCHandshakeInterceptor, ("play.picraft.cc ", 25565)]], args.bind_addr, args.bind_port)
        fwd_server.start()
    except KeyboardInterrupt as e:
        try:
            print("", end="\r")
            if fwd_server:
                logger.info("Server shutting down...")
                fwd_server.stop()

            if len(waitthreads) > 0:
                logger.info("Waiting for worker threads...")
            for thread in waitthreads:
                thread.join()

            if args.proxy and proxy:
                logger.info("Destroying proxy...")
                available_proxies = load_available_proxies()
                destroy_proxy(available_proxies, proxy)
        except KeyboardInterrupt as e:
            print("", end="\r")
            logger.critical("Ctrl+C pressed twice, forcefully shutting down.")

if __name__ == "__main__":
    main()


import torch.nn as nn
import torch.nn.functional as F
from ..neural_network import ResNetBlock

class ActorCritic(nn.Module):

    def __init__(self, feature_dim, action_dim, embedding_dim=64):
        super(ActorCritic, self).__init__()
        self.actor = Actor(feature_dim, action_dim, embedding_dim)
        self.critic = Critic(feature_dim, action_dim, embedding_dim)

    def act(self, x):
        return self.actor(x)

    def evaluate(self, obs):
        """Return logits of actions"""
        values = self.critic(obs)
        return values

class Actor(nn.Module):
    def __init__(self, feature_dim, action_dim, embedding_dim=64):
        super(Actor, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(feature_dim, feature_dim),
            nn.ReLU(),
            nn.Linear(feature_dim, action_dim),
            nn.ReLU(),
            nn.Linear(action_dim, action_dim)
        )
        # n_mid_channcels = action_dim * 2
        # self.net = nn.Sequential(
        #     # + channcel
        #     nn.Conv2d(action_dim, n_mid_channcels, kernel_size=[1, 1], stride=[1, 1]),
        #     ResNetBlock(n_mid_channcels),
        #     # - feature
        #     nn.Conv2d(n_mid_channcels, n_mid_channcels, kernel_size=[feature_dim, 1], stride=[1, 1]),
        #     ResNetBlock(n_mid_channcels),
        #     # - channcel
        #     nn.Conv2d(n_mid_channcels, action_dim, kernel_size=[1, 1], stride=[1, 1]),
        #     ResNetBlock(action_dim),
        #     # last
        #     nn.Conv2d(action_dim, action_dim, kernel_size=[1, 1], stride=[1, 1]),
        #     nn.ReLU(inplace=False),
        #     nn.Flatten(),
        # )
    def forward(self, obs):
        """Return logits of actions"""
        action_logits = self.net(obs)
        # action_logits = self.relu(out)
        return action_logits

    def act(self, obs):
        return self(obs)

class Critic(nn.Module):
    def __init__(self, feature_dim, action_dim, embedding_dim=64):
        super(Critic, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(feature_dim, feature_dim),
            nn.ReLU(),
            nn.Linear(feature_dim, action_dim),
            nn.ReLU(),
            nn.Linear(action_dim, 1)
        )

    def forward(self, obs):
        """Return logits of actions"""
        values = self.net(obs)
        return values


from zipline.pipeline.common import SID_FIELD_NAME
from zipline.pipeline.loaders.blaze.core import ffill_query_in_range

def load_raw_data(assets,
                  data_query_cutoff_times,
                  expr,
                  odo_kwargs,
                  checkpoints=None):
    """
    Given an expression representing data to load, perform normalization and
    forward-filling and return the data, materialized. Only accepts data with a
    `sid` field.

    Parameters
    ----------
    assets : pd.int64index
        the assets to load data for.
    data_query_cutoff_times : pd.DatetimeIndex
        The datetime when data should no longer be considered available for
        a session.
    expr : expr
        the expression representing the data to load.
    odo_kwargs : dict
        extra keyword arguments to pass to odo when executing the expression.
    checkpoints : expr, optional
        the expression representing the checkpointed data for `expr`.

    Returns
    -------
    raw : pd.dataframe
        The result of computing expr and materializing the result as a
        dataframe.
    """
    lower_dt, upper_dt = data_query_cutoff_times[[0, -1]]
    raw = ffill_query_in_range(
        expr,
        lower_dt,
        upper_dt,
        checkpoints=checkpoints,
        odo_kwargs=odo_kwargs,
    )
    sids = raw[SID_FIELD_NAME]
    raw.drop(
        sids[~sids.isin(assets)].index,
        inplace=True
    )
    return raw


import unittest

class TheTuple(unittest.TestCase):
    """
    Python 的元组与列表类似，不同之处在于元组的元素不能修改。
    元组使用小括号 ( )，列表使用方括号 [ ]。
    元组创建很简单，只需要在括号中添加元素，并使用逗号隔开即可。
    也可以不要括号
    """

    def test_a(self):
        tuple0 = ()
        print(type(tuple0))
        tuple1 = ('google', 'amd', 'vdos', 1, 2, 3)
        print(tuple1)
        # 不加括号
        tuple2 = 'niggers', 'pigs', 'dogs'
        print(type(tuple2))
        # 元组中只包含一个元素时，需要在元素后面添加逗号，否则括号会被当作运算符使用
        # x = (1)
        # print(type(x))
        tuple3 = (4396,)
        self.assertEqual(1, 1)

    # 访问、截取，同列表

    # 不允许修改、删除、新增
    def test_unmodifiable(self):
        tuple1 = ('apple', 'amd', 'DJI', 'vdos')
        print(tuple1[3])
        # 修改元组元素操作是非法的
        # tuple1[0] = 'huawei'
        # tuple1[4] = 'google'

        # del tuple1[0]
        # tuple1.remove('apple')
        # tuple1.pop()

        # 但是可以通过截取、拼接，创建新的元组
        tuple2 = ('huawei',) + tuple1[1:]
        print(tuple2)
        self.assertEqual(1, 1)

    '''
    所谓元组的不可变指的是元组所指向的内存中的内容不可变。（普通常量指针）
    重新赋值的元组，会绑定到新的对象，而非修改了原来的对象。
    '''

    def test_constPointer(self):
        tup = (1,)
        print(id(tup))
        tup = (2,)
        print(id(tup))
        self.assertEqual(1, 1)

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