import pdb

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal
from scipy.stats import entropy, gaussian_kde, normaltest
import nflows
from nflows import distributions, transforms, utils, flows
from nflows.transforms.normalization import BatchNorm
from nflows.nn import nets
from nflows.transforms.base import (
    CompositeTransform,
    InputOutsideDomain,
    InverseTransform,
    Transform,
)
from nflows.utils import torchutils

def build_nflows(num_layers=2, hids=20, dims=2, context_dims=2,
        batch_norm=False, activation=torch.nn.functional.relu, bins = 15, tail=8.0,
        device = 'cuda', rqs=True, bimodal=False):
    context_net = Linear_2L(context_dims, 2*dims, hids, 0.5, 0,
         mc_drop = False, fixed_masks = False,
         different_heads = False, device = device)
    base_dist = nflows.distributions.ConditionalDiagonalNormal(
        shape=[dims], context_encoder= context_net)

    transforms = []

    def create_net(in_features, out_features):
        return Linear_2L(in_features, out_features, hids, 0.5,
                context_dims, fixed_masks = False,
                different_heads = False, device=device)
    for _ in range(num_layers):
        if dims > 1:
            transforms.append(nflows.transforms.RandomPermutation(features=dims))
            mask = nflows.utils.torchutils.create_mid_split_binary_mask(dims)
            transforms.append(
                    nflows.transforms.PiecewiseCubicCouplingTransform(mask, create_net,
                        tails='linear', num_bins=bins, tail_bound=tail,
                ))
        if dims == 1:
            transforms.append(
                nflows.transforms.MaskedPiecewiseQuadraticAutoregressiveTransform(
                 features=dims,
                 hidden_features=hids,
                 context_features=context_dims,
                 num_blocks = 2,
                 use_batch_norm=batch_norm,
                 num_bins=bins,
                 tails='linear',
                 tail_bound = tail,
                 activation = activation,
                 use_residual_blocks = False,))
    transform = nflows.transforms.CompositeTransform(transforms)

    flow = nflows.flows.Flow(transform, base_dist)
    return flow

def build_nflows_ensemble(num_layers=2, hids=20, dims=2, context_dims=2,
        batch_norm=False, activation=torch.nn.functional.relu, bins = 15, tail=8.0,
        device = 'cuda', rqs=True, base = True, flows = True, multihead=False,
        fixed_masks=False, ensemble_size=15, bimodal=False):
    if base:
        context_net = Linear_2L(context_dims, 2*dims, hids*2, 0.5, 0,
            fixed_masks = fixed_masks, num_masks = ensemble_size,
            different_heads = multihead, device = device)
    else:
        context_net = Linear_2L(context_dims, 2*dims, hids*2, 0.5, 0,
             fixed_masks = False, num_masks = ensemble_size,
             different_heads = False, device = device)
    base_dist = nflows.distributions.ConditionalDiagonalNormal(
        shape=[dims], context_encoder= context_net)

    transforms = []

    if flows:
        def create_net(in_features, out_features):
            return Linear_2L(in_features, out_features, hids, 0.5,
                    context_dims, fixed_masks=fixed_masks,
                    different_heads = multihead, num_masks=ensemble_size, device=device)
    else:
        def create_net(in_features, out_features):
            return Linear_2L(in_features, out_features, hids, 0.5,
                    context_dims, fixed_masks = False,
                    different_heads = False, device=device)
    for _ in range(num_layers):
        if dims > 1:
            transforms.append(nflows.transforms.RandomPermutation(features=dims))
            mask = nflows.utils.torchutils.create_mid_split_binary_mask(dims)
            transforms.append(
                    nflows.transforms.PiecewiseCubicCouplingTransform(mask, create_net,
                        tails='linear', num_bins=bins, tail_bound=tail,
                ))
        if dims == 1:
            transforms.append(
                nflows.transforms.MaskedPiecewiseQuadraticAutoregressiveTransform(
                 features=dims,
                 hidden_features=hids,
                 context_features=context_dims,
                 num_blocks = 1,
                 use_batch_norm=batch_norm,
                 num_bins=bins,
                 tails='linear',
                 tail_bound = tail,
                 activation = activation,
                 use_residual_blocks = False,
                 ensemble = flows))
                 #create_context_net = create_net))
    transform = nflows.transforms.CompositeTransform(transforms)

    flow = nflows.flows.Flow(transform, base_dist)
    return flow

class Linear_2L(nn.Module):
    def __init__(self, input_dim, output_dim, n_hid, pdrop, context_dim,
                 fixed_masks = False, num_masks = 10, different_heads = False,
                 device='cpu'):
        super(Linear_2L, self).__init__()

        self.pdrop = pdrop

        self.input_dim = input_dim
        self.output_dim = output_dim
        self.n_hid = n_hid

        self.fc1 = nn.Linear(input_dim+context_dim, n_hid)
        self.fc2 = nn.Linear(n_hid, n_hid)
        if different_heads:
            self.heads = []
            for i in range(num_masks):
                exec(f'self.head{i} = nn.Linear(n_hid, output_dim)')
                exec(f'self.heads.append(self.head{i})')
        else:
            self.fc3 = nn.Linear(n_hid, output_dim)

        self.different_heads = different_heads
        # choose your non linearity
        # self.act = nn.Tanh()
        # self.act = nn.Sigmoid()
        self.act = nn.ReLU(inplace=True)
        # self.act = nn.ELU(inplace=True)
        # self.act = nn.SELU(inplace=True)
        self.fixed_masks = fixed_masks
        if fixed_masks:
            self.create_masks(num_masks, device)
        self.num_masks = num_masks

    def forward(self, x, context=None, rand_mask=True, mask_index = 0):
        if self.fixed_masks:
            if rand_mask:
                mask = self.masks[np.random.choice(self.num_masks)]
            else:
                mask = self.masks[mask_index]

        if self.different_heads:
            if rand_mask:
                head_idx = np.random.choice(self.num_masks)
            else:
                head_idx = mask_index

        x = x.view(-1, self.input_dim)  # view(batch_size, input_dim)
        if context is None:
            pass
        else:
            x = torch.cat((x, context), dim=1)
        # -----------------
        x = self.fc1(x)
        if self.fixed_masks:
             x = mask[0].repeat(x.shape[0],1)*x
        # -----------------
        x = self.act(x)
        # -----------------
        x = self.fc2(x)
        if self.fixed_masks:
            x = mask[1].repeat(x.shape[0],1)*x
        # -----------------
        x = self.act(x)
        # -----------------
        if self.different_heads:
            y = self.heads[head_idx](x)
        else:
            y = self.fc3(x)

        return y

    def create_masks(self, num_masks, device):
        masks = []
        for i in range(num_masks):
            mask_l1 = torch.bernoulli(torch.full_like(torch.ones(self.n_hid), self.pdrop))\
                .to(device)
            mask_l2 = torch.bernoulli(torch.full_like(torch.ones(self.n_hid), self.pdrop))\
                .to(device)
            masks.append([mask_l1, mask_l2])
        self.masks = masks


import subprocess
import time
import os
import stat
import threading
import uuid

class Iperf3(object):

    def __init__(self, _ssh_machine1,
                 _ssh_key1,
                 _ssh_machine2,
                 _ssh_key2):
        self.ssh_machine1 = _ssh_machine1
        self.ssh_machine2 = _ssh_machine2
        self.ssh_key1     = _ssh_key1
        self.ssh_key2     = _ssh_key2

    def generate_test_file(self,
                           command_list,
                           filename):

        with open(filename, 'w') as f:
            f.write("#!/bin/bash\n")
            for command in command_list:
                f.write(" ".join(command) + "\n")
                f.write("sleep 5\n")
        os.chmod(filename, os.stat(filename).st_mode | stat.S_IEXEC)

    def get_result_value_from_client_iperf_file(self,client_file):
        print(client_file)
        proc = subprocess.Popen(['./get_value.sh',client_file],stdout=subprocess.PIPE)
        proc.wait()
        value_bytes = proc.communicate()[0].decode('utf-8')
        value=''.join(str(v) for v in value_bytes)
        # May return \n only
        if not value or ('\n' in value and len(value)==1):
            return None
        print(value)
        proc = subprocess.Popen(['./get_metric.sh',client_file],stdout=subprocess.PIPE)
        proc.wait()
        metric_bytes = proc.communicate()[0].decode('utf-8')
        metric=''.join(str(v) for v in metric_bytes)
        if 'M' in metric:
            return float(value)
        if 'G' in metric:
            return (float(value) * 1000.0)
        return(float(value) * 0.001)

    def get_results(self,
                    client_key,
                    client_addr,
                    flow_num=20):
        sum = 0.0
        filepath='./' + client_addr + '_'
        filepath += str(uuid.uuid4())
        filepath += '/'
        os.mkdir(filepath)
        scp = subprocess.Popen(['scp','-i',client_key,client_addr + ':~/iperf3_output.*',filepath])
        scp.wait()
        failed_flows = 0
        for i in range(0,flow_num):
            outfile = filepath + 'iperf3_output.' + str(i)
            res = self.get_result_value_from_client_iperf_file(outfile)
            if res == None:
                failed_flows += 1
            else:
                sum += res

        print('Total is: {} Mbps'.format(sum))
        print('Mean is: {} Mbps'.format(sum/float(flow_num)))

    def run_performance_tests(self,
                              use_udp=False, # protocol to be used
                              bw='500M',       # bandwidth
                              duration='300',
                              flow_num=20,
                              server_addr=None,
                              server_port=5201,
                              server_file='server_file.sh',
                              client_file='client_file.sh'):
        sleep_between_serv_clients = 30
        s_cmd_base = 'iperf3 -s -1'
        c_cmd_base = 'iperf3 -c ' + self.ssh_machine2 + ' -b ' + bw + ' -t ' + duration
        if use_udp:
            c_cmd_base += ' -u'
        port=server_port
        s_cmd_list = []
        for i in range(0,flow_num):
            outfile = 'iperf3_output.' + str(i)
            #s_cmd = ['ssh','-i',self.ssh_key2,self.ssh_machine2,
            #         'nohup',s_cmd_base,'-p',str(port+i),'&>',outfile]
            s_cmd = ['nohup',s_cmd_base,'-p',str(port+i),'&>',outfile,'&']
            s_cmd_list.append(s_cmd)

        self.generate_test_file(s_cmd_list,server_file)
        s_scp = subprocess.Popen(['scp','-i',self.ssh_key2,server_file,self.ssh_machine2 + ':~/']);
        s_scp.wait()
        #print("Running: {} as server".format(s_cmd))
        subprocess.Popen(['ssh','-i',self.ssh_key2,self.ssh_machine2,'./' + server_file])

        time.sleep(sleep_between_serv_clients)

        c_cmd_list = []
        for i in range(0,flow_num):
            outfile = 'iperf3_output.' + str(i)
            #c_cmd = ['ssh','-i',self.ssh_key1,self.ssh_machine1,
            #         'nohup',c_cmd_base,'-p',str(port+i),'&>',outfile]
            c_cmd = ['nohup',c_cmd_base,'-p',str(port+i),'&>',outfile,'&']
            c_cmd_list.append(c_cmd)

        self.generate_test_file(c_cmd_list,client_file)
        c_scp = subprocess.Popen(['scp','-i',self.ssh_key1,client_file,self.ssh_machine1 + ':~/']);
        c_scp.wait()
        #print("Running: {} as server".format(c_cmd))
        subprocess.Popen(['ssh','-i',self.ssh_key1,self.ssh_machine1,'./' + client_file])

        print("Waiting for test to finish........")
        time.sleep(int(duration) + sleep_between_serv_clients)
        print("DONE")
        #subprocess.Popen(['ssh','-i',self.ssh_key2,self.ssh_machine2,
        #                  "kill -9 $(ps aux | grep iperf | awk \'{print $2}\')"])
        self.get_results(client_key=self.ssh_key1,
                         client_addr=self.ssh_machine1,
                         flow_num=flow_num)

if __name__=="__main__":
    print("*************************************")
    print("** Make sure SSH keys for servers  **")
    print("** SSH address should of form:     **")
    print("** name@IP                         **")
    print("** or                              **")
    print("** name@hostname                   **")
    print("** Key should be a filepath        **")
    print("**                                 **")
    print("** Make sure iperf3 is installed   **")
    print("*************************************")

    ##### test STARTUP parameters:
    use_udp=False
    bw='500M'
    duration='300'
    flow_num=20
    server_addr=None
    server_port=5201

    ####
    # test_list syntax:
    # ( IP MACHINE 1, KEY MACHINE 1, IP MACHINE 2, KEY MACHINE 2)
    test_list = [('10.5.0.3','./id_iperf_test','10.5.0.30','./id_iperf_test')]
                 #('10.5.0.3','./id_iperf_test','10.5.0.30','./id_iperf_test')]
    thread_list = []
    for tup in test_list:
        test = Iperf3(tup[0],tup[1],tup[2],tup[3])
        thread = threading.Thread(test.run_performance_tests(use_udp=use_udp,
                                                             bw=bw,
                                                             duration=duration,
                                                             flow_num=flow_num,
                                                             server_port=server_port))
        thread_list.append(thread)
        thread.start()

    #waiting threads to finish:
    for t in thread_list:
        t.join()


#Perulangan For

#for nilai in sequence:
#   blok code

#Contoh 1
for i in "mizard":
    print(i)

#Contoh 2 Menggunakan fungsi range()
for i in range(10): #Start(dimulai dari 0 dan berhenti diangka sebelum angka terakhir)
    print(i)
for i in range(2,11): #Stop(berhenti diangka sebelumnya angka terakhir)
    print(i)
for i in range(5,20,2): #Step(bertambah sejumlah angka yang diisikan)
    print(i)

#Contoh 3 perulangan menggunakan continue
for i in range(10):
    print(i)
    if i == 5:
        print(i)
        continue

#Contoh 4 perulangan menggunakan break
for i in range(20):
   if i == 15:
        print("Ini perulangan ke-",i)
        break

#Contoh 5 perulangan menggunakan list
data = ["mizard","Jamal","Udin","Ngab"]
for i in data:
    if i == "Udin":
        print(i)
        break

#Perulangan while
#while nilai operator sequence:
#   blok kode

#Contoh 1
i = 2
while i < 11:
    print(i)
    i += 1

#Contoh 2
data = [10,20,30,40]
for i in data:
    if i == 30:
        print(i)
        break
i = 1
while i < len(data):
    if data[i] == 30:
      print(data[i])
      break
    i += 1

#Challenge
for i in range(2,41):
    if i == 10:
        print("ini adalah nilai",i)
        continue
    elif i == 20:
        print("ini adalah nilai",i)
        continue
    elif i == 30:
        print("ini adalah nilai",i)
        break


"""
create model
Creator: Xiaoshui Huang
Date: 2020-06-19
"""
from se_math.so3 import inverse, transform
import torch
import numpy as np
from random import sample

import se_math.se3 as se3
import se_math.invmat as invmat

import igl
import os
import sys

sys.path.append('./../')
sys.path.append('./../../')
from loss import cal_loss_intersection_batch_whole_median_pts_lines, Reconstruction_point, Random_uniform_distribution_lines_batch_efficient_resample, chamfer_dist, Sample_neighs

from utils import npmat2euler

# we also make chamfer_loss for data!

def dict_all_to_device(tensor_dict, device):
    """Sends everything into a certain device """
    for k in tensor_dict:
        if isinstance(tensor_dict[k], torch.Tensor):
            tensor_dict[k] = tensor_dict[k].to(device)

def save_pred_gt_obj(V_src, V_pred, V_gt, V_tgt_trans, paths_src, paths_pred,
                     paths_gt, paths_gt_pred):
    Face = np.zeros(3).reshape(1, 3).astype(np.int32)
    for i in range(V_pred.shape[0]):
        igl.write_triangle_mesh(paths_src[i], V_src[i].numpy(), Face)
        igl.write_triangle_mesh(paths_pred[i], V_pred[i].numpy(), Face)
        igl.write_triangle_mesh(paths_gt[i], V_gt[i].numpy(), Face)
        igl.write_triangle_mesh(paths_gt_pred[i], V_tgt_trans[i].numpy(), Face)

# a global function to flatten a feature
def flatten(x):
    return x.view(x.size(0), -1)

# a global function to calculate max-pooling
def symfn_max(x):
    # [B, K, N] -> [B, K, 1]
    a = torch.nn.functional.max_pool1d(x, x.size(-1))
    return a

# a global function to generate mlp layers
def _mlp_layers(nch_input,
                nch_layers,
                b_shared=True,
                bn_momentum=0.1,
                dropout=0.0):
    """ [B, Cin, N] -> [B, Cout, N] or
        [B, Cin] -> [B, Cout]
    """
    layers = []
    last = nch_input
    for i, outp in enumerate(nch_layers):
        if b_shared:
            weights = torch.nn.Conv1d(last, outp, 1)
        else:
            weights = torch.nn.Linear(last, outp)
        layers.append(weights)
        # layers.append(torch.nn.BatchNorm1d(outp, momentum=bn_momentum))
        layers.append(torch.nn.GroupNorm(8, outp))
        layers.append(torch.nn.ReLU())
        if b_shared == False and dropout > 0.0:
            layers.append(torch.nn.Dropout(dropout))
        last = outp
    return layers

# a class to generate MLP network
class MLPNet(torch.nn.Module):
    """ Multi-layer perception.
        [B, Cin, N] -> [B, Cout, N] or
        [B, Cin] -> [B, Cout]
    """
    def __init__(self,
                 nch_input,
                 nch_layers,
                 b_shared=True,
                 bn_momentum=0.1,
                 dropout=0.0):
        super().__init__()
        list_layers = _mlp_layers(nch_input, nch_layers, b_shared, bn_momentum,
                                  dropout)
        self.layers = torch.nn.Sequential(*list_layers)

    def forward(self, inp):
        out = self.layers(inp)
        return out

# encoder network
class PointNet(torch.nn.Module):
    def __init__(self, dim_k=1024):
        super().__init__()
        scale = 1
        mlp_h1 = [int(64 / scale), int(64 / scale)]
        mlp_h2 = [int(64 / scale), int(128 / scale), int(dim_k / scale)]

        self.h1 = MLPNet(3, mlp_h1, b_shared=True).layers
        self.h2 = MLPNet(mlp_h1[-1], mlp_h2, b_shared=True).layers
        self.sy = symfn_max

    def forward(self, points):
        """ points -> features
            [B, N, 3] -> [B, K]
        """
        # for pointnet feature extraction
        x = points.transpose(1, 2)  # [B, 3, N]
        x = self.h1(x)
        x = self.h2(x)  # [B, K, N]
        x = flatten(self.sy(x))

        return x

# decoder network
class Decoder(torch.nn.Module):
    def __init__(self, num_points=2048, bottleneck_size=1024):
        super(Decoder, self).__init__()
        self.num_points = num_points
        self.bottleneck_size = bottleneck_size
        # self.bn1 = torch.nn.BatchNorm1d(bottleneck_size)
        # self.bn2 = torch.nn.BatchNorm1d(bottleneck_size // 2)
        # self.bn3 = torch.nn.BatchNorm1d(bottleneck_size // 4)
        self.bn1 = torch.nn.GroupNorm(8, bottleneck_size)
        self.bn2 = torch.nn.GroupNorm(8, bottleneck_size // 2)
        self.bn3 = torch.nn.GroupNorm(8, bottleneck_size // 4)
        self.fc1 = torch.nn.Linear(self.bottleneck_size, bottleneck_size)
        self.fc2 = torch.nn.Linear(self.bottleneck_size, bottleneck_size // 2)
        self.fc3 = torch.nn.Linear(bottleneck_size // 2, bottleneck_size // 4)
        self.fc4 = torch.nn.Linear(bottleneck_size // 4, self.num_points * 3)
        self.th = torch.nn.Tanh()

    def forward(self, x):
        batchsize = x.size()[0]
        x = torch.nn.functional.relu(self.bn1(self.fc1(x)))
        x = torch.nn.functional.relu(self.bn2(self.fc2(x)))
        x = torch.nn.functional.relu(self.bn3(self.fc3(x)))
        x = self.th(self.fc4(x)) * 10
        x = x.view(batchsize, 3, self.num_points).transpose(1, 2).contiguous()
        return x

# the neural network of feature-metric registration
class SolveRegistration(torch.nn.Module):
    def __init__(self, ptnet, decoder=None):
        super().__init__()
        # network
        self.encoder = ptnet
        self.decoder = decoder
        # functions
        self.inverse = invmat.InvMatrix.apply
        self.exp = se3.Exp  # [B, 6] -> [B, 4, 4]
        self.transform = se3.transform  # [B, 1, 4, 4] x [B, N, 3] -> [B, N, 3]

        # initialization for dt: [w1, w2, w3, v1, v2, v3], 3 rotation angles and 3 translation
        delta = 1.0e-2  # step size for approx. Jacobian (default: 1.0e-2)
        dt_initial = torch.autograd.Variable(
            torch.Tensor([delta, delta, delta, delta, delta, delta]))
        self.dt = torch.nn.Parameter(dt_initial.view(1, 6), requires_grad=True)

        # results
        self.last_err = None
        self.g_series = None  # for debug purpose
        self.prev_r = None
        self.g = None  # estimated transformation T
        self.device = None
        self.g_series_gpu = None

    # estimate T
    # noly return the encoder loss, but also return intersection loss

    def estimate_t(self,
                   data,
                   maxiter=5,
                   xtol=1.0e-7,
                   p0_zero_mean=True,
                   p1_zero_mean=True,
                   mode='train'):
        """
        give two point clouds, estimate the T by using IC algorithm
        :param p0: point cloud
        :param p1: point cloud
        :param maxiter: maximum iteration
        :param xtol: a threshold for early stop of transformation estimation
        :param p0_zero_mean: True: normanize p0 before IC algorithm
        :param p1_zero_mean: True: normanize p1 before IC algorithm
        :return: feature-metric projection error (r), encoder-decoder loss (loss_ende) and intersection loss!
        """
        p1 = data['points_src_sample']
        p0 = data['points_tar_sample']
        a0 = torch.eye(4).view(1, 4, 4).expand(p0.size(0), 4,
                                               4).to(p0)  # [B, 4, 4]
        a1 = torch.eye(4).view(1, 4, 4).expand(p1.size(0), 4,
                                               4).to(p1)  # [B, 4, 4]
        self.device = p1.device
        batch_size = p1.shape[0]
        # normalization
        if p0_zero_mean:
            p0_m = p0.mean(dim=1)  # [B, N, 3] -> [B, 3]
            a0 = a0.clone()
            a0[:, 0:3, 3] = p0_m
            q0 = p0 - p0_m.unsqueeze(1)
        else:
            q0 = p0
        if p1_zero_mean:
            p1_m = p1.mean(dim=1)  # [B, N, 3] -> [B, 3]
            a1 = a1.clone()
            a1[:, 0:3, 3] = -p1_m
            q1 = p1 - p1_m.unsqueeze(1)
        else:
            q1 = p1

        # use IC algorithm to estimate the transformation
        # generate the transform!
        g0 = torch.eye(4).to(q0).view(1, 4, 4).expand(q0.size(0), 4,
                                                      4).contiguous()
        r, g, loss_ende = self.ic_algo(g0, q0, q1, maxiter, xtol)
        # the g don't backgrade the gradinent？
        self.g = g
        # re-normalization
        if p0_zero_mean or p1_zero_mean:
            est_g = self.g
            if p0_zero_mean:
                est_g = a0.to(est_g).bmm(est_g)
            if p1_zero_mean:
                est_g = est_g.bmm(a1.to(est_g))
            self.g = est_g

            est_gs = self.g_series  # [M, B, 4, 4]
            if p0_zero_mean:
                est_gs = a0.unsqueeze(0).contiguous().to(est_gs).matmul(est_gs)
            if p1_zero_mean:
                est_gs = est_gs.matmul(a1.unsqueeze(0).contiguous().to(est_gs))
            self.g_series = est_gs

            est_gs_gpu = self.g_series_gpu  # [M, B, 4, 4]
            if p0_zero_mean:
                est_gs_gpu = a0.unsqueeze(0).contiguous().to(
                    est_gs_gpu).matmul(est_gs_gpu)
            if p1_zero_mean:
                est_gs_gpu = est_gs_gpu.matmul(
                    a1.unsqueeze(0).contiguous().to(est_gs_gpu))
            self.g_series_gpu = est_gs_gpu

        loss_pp_wise = (torch.mean(
            torch.abs(
                self.transform(self.g.unsqueeze(1), data['points_src_sample'])
                - self.transform(
                    torch.inverse(data['igt']).unsqueeze(1),
                    data['points_src_sample']))))
        if mode is 'train':
            R = (torch.norm(
                data['tar_box'][:, 0, :] - data['tar_box'][:, -1, :],
                dim=-1,
                p=2) * 0.5).reshape(-1, 1)
            lines = None
            points_ref = data['points_tar_sample'].contiguous()

            tar_faces_tensor = data['points_based_neighs_tar'].reshape(
                points_ref.shape[0], -1, 9)
            # if we used the transformed, we may generate better results!

            temp_g = self.g_series_gpu[-1]

            pred_src_transformed_final_sample = self.transform(
                temp_g.unsqueeze(1),
                data['points_src_sample'].contiguous()).detach()

            # pred_src_transformed_final_sample = data['points_src_sample']
            if lines is None:
                lines = Random_uniform_distribution_lines_batch_efficient_resample(
                    R, data['centers'], 15000,
                    pred_src_transformed_final_sample.contiguous(),
                    data['points_tar_sample'].contiguous(), self.device)

            # set our loss;
            loss_intersection = torch.FloatTensor([0]).to(self.device)
            for i in range(maxiter - 3, maxiter):
                temp_g = self.g_series_gpu[i]
                pred_src_transformed_final_sample = self.transform(
                    temp_g.unsqueeze(1), data['points_src_sample'])

                pred_src_faces_tensor = self.transform(
                    temp_g.unsqueeze(1),
                    data['points_based_neighs_src']).reshape(
                        pred_src_transformed_final_sample.shape[0], -1, 9)
                tp_loss_intersection = torch.FloatTensor([0]).to(self.device)
                for j in range(pred_src_faces_tensor.shape[0]):
                    tp_loss_intersection += cal_loss_intersection_batch_whole_median_pts_lines(
                        1, 1, 5, 5, pred_src_faces_tensor[j:j + 1, :, :],
                        tar_faces_tensor[j:j + 1, :, :], lines[j:j + 1, :, :],
                        self.device) / 5.0
                loss_intersection = loss_intersection + \
                    tp_loss_intersection*0.5**(maxiter-i-1)

            loss_chamfer = chamfer_dist(pred_src_transformed_final_sample,
                                        data['points_tar_sample'])

            return r, loss_ende, loss_intersection / batch_size, loss_pp_wise, loss_chamfer
        return r, loss_ende, loss_pp_wise,

    # IC algorithm
    # encoder, we just use the chamfer!
    def ic_algo(self, g0, p0, p1, maxiter, xtol):
        """
        use IC algorithm to estimate the increment of transformation parameters
        :param g0: initial transformation
        :param p0: point cloud
        :param p1: point cloud
        :param maxiter: maxmimum iteration
        :param xtol: a threashold to check increment of transformation  for early stop
        :return: feature-metric projection error (r), updated transformation (g), encoder-decoder loss
        """
        training = self.encoder.training
        # training = self.decoder.training
        batch_size = p0.size(0)

        self.last_err = None
        g = g0
        self.g_series = torch.zeros(maxiter + 1, *g0.size(), dtype=g0.dtype)
        self.g_series[0] = g0.clone()
        self.g_series_gpu = torch.zeros(maxiter, *g0.size(),
                                        dtype=g0.dtype).to(self.device)

        # generate the features
        f0 = self.encoder(p0)
        f1 = self.encoder(p1)

        # task 1
        loss_enco_deco = 0.0
        if self.decoder is not None:
            # we generate the decoder f0?
            # make an encoder decoder!
            decoder_out_f0 = self.decoder(f0)
            decoder_out_f1 = self.decoder(f1)

            # the decoder meets AE!
            p0_dist1, p0_dist2 = self.chamfer_loss(
                p0.contiguous(), decoder_out_f0)  # loss function
            loss_net0 = (torch.mean(p0_dist1)) + (torch.mean(p0_dist2))
            p1_dist1, p1_dist2 = self.chamfer_loss(
                p1.contiguous(), decoder_out_f1)  # loss function
            loss_net1 = (torch.mean(p1_dist1)) + (torch.mean(p1_dist2))
            loss_enco_deco = loss_net0 + loss_net1

        # self.encoder.eval()  # and fix them BN.
        # if fix, ho to backward gradients？
        # task 2
        f0 = self.encoder(p0)  # [B, N, 3] -> [B, K]
        # approx. J by finite difference
        dt = self.dt.to(p0).expand(batch_size,
                                   6)  # convert to the type of p0. [B, 6]
        J = self.approx_Jac(p0, f0, dt)
        # compute pinv(J) to solve J*x = -r
        try:
            Jt = J.transpose(1, 2)  # [B, 6, K]
            H = Jt.bmm(J)  # [B, 6, 6]
            # H = H + u_lamda * iDentity
            B = self.inverse(H)
            pinv = B.bmm(Jt)  # [B, 6, K]
        except RuntimeError as err:
            self.last_err = err
            f1 = self.encoder(p1)  # [B, N, 3] -> [B, K]
            r = f1 - f0
            self.ptnet.train(training)
            return r, g, -1

        itr = 0
        r = None
        # we
        for itr in range(maxiter):
            p = self.transform(g.unsqueeze(1),
                               p1)  # [B, 1, 4, 4] x [B, N, 3] -> [B, N, 3]
            f1 = self.encoder(p)  # [B, N, 3] -> [B, K]
            r = f1 - f0  # [B,K]
            # generate the r！
            dx = -pinv.bmm(r.unsqueeze(-1)).view(batch_size, 6)

            check = dx.norm(p=2, dim=1, keepdim=True).max()
            if float(check) < xtol:
                if itr == 0:
                    self.last_err = 0  # no update.
                break

            g = self.update(g, dx)
            self.g_series_gpu[itr] = g
            self.g_series[itr + 1] = g.clone()
            self.prev_r = r

        self.encoder.train(training)
        return r, g, loss_enco_deco

    # estimate Jacobian matrix
    def approx_Jac(self, p0, f0, dt):
        # p0: [B, N, 3], Variable
        # f0: [B, K], corresponding feature vector
        # dt: [B, 6], Variable
        # Jk = (ptnet(p(-delta[k], p0)) - f0) / delta[k]
        batch_size = p0.size(0)
        num_points = p0.size(1)

        # compute transforms
        transf = torch.zeros(batch_size, 6, 4, 4).to(p0)
        for b in range(p0.size(0)):
            d = torch.diag(dt[b, :])  # [6, 6]
            D = self.exp(-d)  # [6, 4, 4]
            transf[b, :, :, :] = D[:, :, :]
        transf = transf.unsqueeze(2).contiguous()  # [B, 6, 1, 4, 4]
        p = self.transform(transf,
                           p0.unsqueeze(1))  # x [B, 1, N, 3] -> [B, 6, N, 3]

        f0 = f0.unsqueeze(-1)  # [B, K, 1]
        f1 = self.encoder(p.view(-1, num_points, 3))
        f = f1.view(batch_size, 6, -1).transpose(1, 2)  # [B, K, 6]

        df = f0 - f  # [B, K, 6]
        J = df / dt.unsqueeze(1)  # [B, K, 6]

        return J

    # update the transformation
    def update(self, g, dx):
        # [B, 4, 4] x [B, 6] -> [B, 4, 4]
        dg = self.exp(dx)
        return dg.matmul(g)

    # calculate the chamfer loss
    def chamfer_loss(self, a, b):
        x, y = a, b
        bs, num_points, points_dim = x.size()
        xx = torch.bmm(x, x.transpose(2, 1))
        yy = torch.bmm(y, y.transpose(2, 1))
        zz = torch.bmm(x, y.transpose(2, 1))
        # diag_ind = torch.arange(0, num_points).type(torch.cuda.LongTensor)
        diag_ind = torch.arange(0, num_points)
        rx = xx[:, diag_ind, diag_ind].unsqueeze(1).expand_as(xx)
        ry = yy[:, diag_ind, diag_ind].unsqueeze(1).expand_as(yy)
        P = (rx.transpose(2, 1) + ry - 2 * zz)
        return torch.min(P, 1)[0], torch.min(P, 2)[0]

    @staticmethod
    def rsq(r):
        # |r| should be 0
        z = torch.zeros_like(r)
        return torch.nn.functional.mse_loss(r, z, reduction='sum')

    @staticmethod
    def comp(g, igt):
        """ |g*igt - I| (should be 0) """
        assert g.size(0) == igt.size(0)
        assert g.size(1) == igt.size(1) and g.size(1) == 4
        assert g.size(2) == igt.size(2) and g.size(2) == 4
        A = g.matmul(igt)
        I = torch.eye(4).to(A).view(1, 4, 4).expand(A.size(0), 4, 4)
        return torch.nn.functional.mse_loss(A, I, reduction='mean') * 16

    @staticmethod
    def comp_inv(g, igt):
        """ |g*igt - I| (should be 0) """
        assert g.size(0) == igt.size(0)
        assert g.size(1) == igt.size(1) and g.size(1) == 4
        assert g.size(2) == igt.size(2) and g.size(2) == 4
        # A = g.matmul(igt)
        gt = torch.inverse(igt)
        # I = torch.eye(4).to(A).view(1, 4, 4).expand(A.size(0), 4, 4)

        return torch.nn.functional.mse_loss(g, gt, reduction='mean')

# main algorithm class
class FMRTrain:
    def __init__(self, dim_k, num_points, train_type):
        self.dim_k = dim_k
        self.num_points = num_points
        self.max_iter = 5  # max iteration time for IC algorithm
        # 0: unsupervised, 1: semi-supervised see. self.compute_loss()
        self._loss_type = train_type

        self.transform = se3.transform  # [B, 1, 4, 4] x [B, N, 3] -> [B, N, 3]

    def create_model(self):
        # Encoder network: extract feature for every point. Nx1024
        ptnet = PointNet(dim_k=self.dim_k)
        # Decoder network: decode the feature into points
        decoder = Decoder(num_points=self.num_points)
        # feature-metric ergistration (fmr) algorithm: estimate the transformation T
        fmr_solver = SolveRegistration(ptnet, decoder)
        return fmr_solver

    def compute_loss(self, solver, data, device, mode='train', maxiter=5):
        # p0, p1, igt = data
        # p0 = p0.to(device)  # template
        # p1 = p1.to(device)  # source
        # igt = igt.to(device)  # igt: p0 -> p1
        dict_all_to_device(data, device)
        p1 = data['points_src_sample']
        p0 = data['points_tar_sample']
        igt = data['igt']

        if mode is 'train':
            r, loss_ende, loss_intersection, loss_pp_wise, loss_chamfer = solver.estimate_t(
                data, self.max_iter, mode=mode)
        else:
            # test model!
            r, loss_ende, loss_pp_wise = solver.estimate_t(data,
                                                           maxiter,
                                                           mode=mode)
        loss_r = solver.rsq(r)
        est_g = solver.g
        # generate the difference between the pred and gt!
        loss_g = solver.comp_inv(est_g, igt)

        # unsupervised learning, set max_iter=0
        if self.max_iter == 0:
            return loss_ende

        # semi-supervised learning, set max_iter>0
        if self._loss_type == 0:
            loss = loss_ende
        elif self._loss_type == 1:
            loss = loss_ende + loss_g
        elif self._loss_type == 2:
            loss = loss_r + loss_g
        else:
            loss = loss_g

        # we need use the multiple indicators to measure the quality!
        np_pred_rotation = est_g[:, :3, :3].transpose(
            2, 1).detach().cpu().numpy()
        np_pred_euler = npmat2euler(np_pred_rotation, 'xyz')
        np_gt_rotation = data['R'].detach().cpu().numpy()
        np_gt_euler = npmat2euler(np_gt_rotation, 'xyz')
        loss_rotation_euler_mae = np.mean(np.abs(np_pred_euler - np_gt_euler))
        loss_rotation_euler_rmse = np.sqrt(
            np.mean((np_pred_euler - np_gt_euler)**2))

        np_loss = {
            'loss_rot_euler_mae': loss_rotation_euler_mae,
            'loss_rot_euler_rmse': loss_rotation_euler_rmse
        }

        # set the weights
        if mode is 'train':
            return 0.01 * loss_ende + 1.0 * loss_intersection + .0 * loss_g + 0.0 * loss_chamfer, loss_g.detach(
            ), loss_intersection.detach(), loss_pp_wise.detach(
            ), loss_ende.detach(), np_loss
        return loss_g, loss_g.detach(), loss_pp_wise.detach(
        ), loss_ende.detach(), np_loss

    def train(self,
              model,
              trainloader,
              optimizer,
              device,
              epoch,
              train_writer=None):
        model.train()

        Debug = True
        total_loss = 0
        total_loss_gt = 0
        total_loss_intersection = 0
        total_loss_pp_wise = 0
        total_loss_encoder = 0
        total_loss_rot_euler_mae = 0
        total_loss_rot_euler_rmse = 0
        if Debug:
            epe = 0
            count = 0
            count_mid = 9
        for i, data in enumerate(trainloader):
            loss, loss_gt, loss_intersection, loss_pp_wise, loss_ende, np_loss = self.compute_loss(
                model, data, device)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            loss_item = loss.item()
            total_loss += loss_item
            total_loss_gt += loss_gt.item()
            total_loss_pp_wise += loss_pp_wise.item()
            total_loss_intersection += loss_intersection.item()
            total_loss_encoder += loss_ende.item()
            total_loss_rot_euler_mae += np_loss['loss_rot_euler_mae']
            total_loss_rot_euler_rmse += np_loss['loss_rot_euler_rmse']
            if Debug:
                epe += loss_item
                if count % 10 == 0:
                    print('i=%d, fmr_loss=%f ' % (i, float(epe) /
                                                  (count_mid + 1)))
                    epe = 0.0
            count += 1
            print(
                "ba/ep{:0d}/{:0d},l_insec:{:4f}, l_gt{:4f},l_pp_w{:4f}, l_en{:4f}, l_rot_eu_mae{:4f}, l_rot_eu_rmse{:4f}"
                .format(i, epoch, loss_intersection.item(), loss_gt.item(),
                        loss_pp_wise.item(), loss_ende.item(),
                        np_loss['loss_rot_euler_mae'],
                        np_loss['loss_rot_euler_rmse']))
        ave_loss = float(total_loss) / count
        ave_loss_gt = float(total_loss_gt) / count
        ave_loss_intersection = float(total_loss_intersection) / count
        ave_loss_wise = float(total_loss_pp_wise) / count
        ave_loss_encoder = float(total_loss_encoder) / count
        ave_loss_rot_euler_mae = (float)(total_loss_rot_euler_mae) / count
        ave_loss_rot_euler_rmse = (float)(total_loss_rot_euler_rmse) / count
        if train_writer is not None:
            train_writer.add_scalar('./loss/loss_sum', ave_loss, epoch)
            train_writer.add_scalar('./loss/loss_gt', ave_loss_gt, epoch)
            train_writer.add_scalar('./loss/loss_intersec',
                                    ave_loss_intersection, epoch)
            train_writer.add_scalar('./loss/loss_wise_mse', ave_loss_wise,
                                    epoch)
            train_writer.add_scalar('./loss/loss_ende', ave_loss_encoder,
                                    epoch)
            train_writer.add_scalar('./lr', optimizer.param_groups[0]['lr'],
                                    epoch)
            train_writer.add_scalar('./loss/loss_rot_euler_mae',
                                    ave_loss_rot_euler_mae, epoch)
            train_writer.add_scalar('./loss/loss_rot_euler_rmse',
                                    ave_loss_rot_euler_rmse, epoch)

# \033[36m,test gt:{:4f}, pp_wise:{:4f}, rot_mae{:4f}, rot_rmse{:4f}\033[0m
        print(
            " \033[36m,train:l_gt:{:4f}, l_intersec:{:4f}, l_pp_wise{:4f}, l_encoder{:4f}, l_rot_eu_mae{:4f}, l_rot_eu_rmse{:4f} \033[0m, "
            .format(ave_loss_gt, ave_loss_intersection, ave_loss_wise,
                    ave_loss_encoder, ave_loss_rot_euler_mae,
                    ave_loss_rot_euler_rmse))
        return ave_loss

    def validate(self, model, testloader, device, epoch, save_results=None):
        # model.eval()
        vloss = 0.0
        vloss_gt = 0.0
        vloss_pp_wise = 0.0
        vloss_rot_euler_mae = 0.0
        vloss_rot_euler_rmse = 0.0
        count = 0
        count_i = 0
        with torch.no_grad():
            for i, data in enumerate(testloader):
                loss_net, loss_gt, loss_pp_wise, loss_ende, np_loss = self.compute_loss(
                    model, data, device, mode='test')
                vloss += loss_net.item()
                vloss_gt += loss_gt.item()
                vloss_pp_wise += loss_pp_wise.item()
                vloss_rot_euler_mae += np_loss['loss_rot_euler_mae']
                vloss_rot_euler_rmse += np_loss['loss_rot_euler_rmse']
                count += 1
                print("Test:sample{:0d},loss_pp_wise:{:4f}".format(
                    i, loss_pp_wise.item()))
                if epoch % 10 == 0:
                    est_g = model.g  # (1, 4, 4)

                    igt = data['igt']
                    ig_gt = igt.cpu().contiguous().view(-1, 4,
                                                        4)  # --> [1, 4, 4]
                    g_hat = est_g.cpu().contiguous().view(-1, 4,
                                                          4)  # --> [1, 4, 4]
                    p1 = data['points_src_sample']
                    p0 = data['points_tar_sample']
                    if save_results is not None:

                        paths_pred = []
                        paths_gt = []
                        paths_src = []
                        paths_gt_pred = []
                        src_transform = self.transform(est_g.unsqueeze(1), p1)
                        src_transform_sample = self.transform(
                            est_g.unsqueeze(1), data['points_src_sample'])
                        tgt_transform = self.transform(igt.unsqueeze(1), p0)
                        V_src = p0.cpu().detach()
                        V_pred = src_transform.cpu().detach()
                        V_gt = p1.cpu().detach()
                        V_tgt_trans = tgt_transform.cpu().detach()
                        for j in range(p0.shape[0]):

                            paths_pred.append(
                                os.path.join(
                                    save_results,
                                    str(epoch) + "pred_src" + str(count_i) +
                                    ".obj"))
                            paths_gt.append(
                                os.path.join(
                                    save_results,
                                    str(epoch) + "gt" + str(count_i) + ".obj"))
                            paths_src.append(
                                os.path.join(
                                    save_results,
                                    str(epoch) + "src" + str(count_i) +
                                    ".obj"))
                            paths_gt_pred.append(
                                os.path.join(
                                    save_results,
                                    str(epoch) + "pred_gt" + str(count_i) +
                                    ".obj"))

                            F = np.zeros([1, 3]).astype(np.int32)
                            igl.write_obj(
                                paths_gt_pred[j].replace(
                                    'pred_gt', 'transformed_sample', 1),
                                src_transform_sample.cpu().detach().numpy().
                                reshape(-1, 3), F)
                            igl.write_obj(
                                paths_gt_pred[j].replace(
                                    'pred_gt', 'src_sample', 1),
                                data['points_src_sample'].cpu().detach().numpy(
                                ).reshape(-1, 3), F)
                            igl.write_obj(
                                paths_gt_pred[j].replace(
                                    'pred_gt', 'tar_sample', 1),
                                data['points_tar_sample'].cpu().detach().numpy(
                                ).reshape(-1, 3), F)

                            count_i += 1
                        save_pred_gt_obj(V_src, V_pred, V_gt, V_tgt_trans,
                                         paths_src, paths_pred, paths_gt,
                                         paths_gt_pred)

        ave_vloss = float(vloss) / count
        ave_vloss_gt = float(vloss_gt) / count
        ave_vloss_pp_wise = float(vloss_pp_wise) / count
        ave_vloss_rot_euler_mae = float(vloss_rot_euler_mae) / count
        ave_vloss_rot_euler_rmse = float(vloss_rot_euler_rmse) / count
        print(
            "\033[36m,test gt:{:4f}, pp_wise:{:4f}, rot_mae{:4f}, rot_rmse{:4f}\033[0m, "
            .format(ave_vloss_gt, ave_vloss_pp_wise, ave_vloss_rot_euler_mae,
                    ave_vloss_rot_euler_rmse))
        return ave_vloss

class FMRTest:
    def __init__(self, args):
        self.filename = args.outfile
        self.dim_k = args.dim_k
        self.max_iter = 10  # max iteration time for IC algorithm
        self._loss_type = 3  # see. self.compute_loss()
        self.transform = se3.transform  # [B, 1, 4, 4] x [B, N, 3] -> [B, N, 3]

    def create_model(self):
        # Encoder network: extract feature for every point. Nx1024
        ptnet = PointNet(dim_k=self.dim_k)
        # feature-metric ergistration (fmr) algorithm: estimate the transformation T
        fmr_solver = SolveRegistration(ptnet)
        return fmr_solver

    # we save the results!
    # pay attention to final results!
    def evaluate(self,
                 solver,
                 testloader,
                 device,
                 save_results=None,
                 writer=None):
        solver.eval()
        with open(self.filename, 'w') as fout:
            self.eval_1__header(fout)
            count_i = 0
            total_loss_pp_wise = 0
            total_loss_gt = 0
            with torch.no_grad():
                for i, data in enumerate(testloader):
                    # p0, p1, igt = data  # igt: p0->p1
                    dict_all_to_device(data, device)
                    p1 = data['points_src_sample']
                    p0 = data['points_tar_sample']
                    igt = data['igt']
                    # igt =
                    # # compute trans from p1->p0
                    # g = se3.log(igt)  # --> [-1, 6]
                    # igt = se3.exp(-g)  # [-1, 4, 4]
                    # p0, p1 = self.ablation_study(p0, p1)

                    p0 = p0.to(device)  # template (1, N, 3)
                    p1 = p1.to(device)  # source (1, M, 3)
                    # When we evaluate, we ignore the chafer, ignore any loss function!
                    r, loss_ende, loss_pp_wise = solver.estimate_t(
                        data, self.max_iter, mode='test')
                    total_loss_pp_wise += loss_pp_wise
                    est_g = solver.g  # (1, 4, 4)

                    ig_gt = igt.cpu().contiguous().view(-1, 4,
                                                        4)  # --> [1, 4, 4]
                    g_hat = est_g.cpu().contiguous().view(-1, 4,
                                                          4)  # --> [1, 4, 4]

                    dg = g_hat.bmm(ig_gt)  # if correct, dg == identity matrix.
                    dx = se3.log(
                        dg)  # --> [1, 6] (if corerct, dx == zero vector)
                    dn = dx.norm(p=2, dim=1)  # --> [1]
                    dm = dn.mean()

                    self.eval_1__write(fout, ig_gt, g_hat)
                    print('test, %d/%d, %f, %f' %
                          (i, len(testloader), dm, loss_pp_wise))
                    if writer is not None:
                        writer.add_scalar('./loss/test', dm, i)
                    # p = self.transform(g.unsqueeze(1),
                    #                    p1)  # [B, 1, 4, 4] x [B, N, 3] -> [B, N, 3]
                    # est_g:p1--->p0
                    # igt: p0-->p1
                    if save_results is not None:
                        paths_pred = []
                        paths_gt = []
                        paths_src = []
                        paths_gt_pred = []
                        src_transform = self.transform(est_g.unsqueeze(1), p1)

                        tgt_transform = self.transform(igt.unsqueeze(1), p0)
                        V_src = p0.cpu().detach()
                        V_pred = src_transform.cpu().detach()
                        V_gt = p1.cpu().detach()
                        V_tgt_trans = tgt_transform.cpu().detach()
                        for i in range(p0.shape[0]):

                            paths_pred.append(
                                os.path.join(save_results,
                                             str(count_i) + "pred_src.obj"))
                            paths_gt.append(
                                os.path.join(save_results,
                                             str(count_i) + "gt.obj"))
                            paths_src.append(
                                os.path.join(save_results,
                                             str(count_i) + "src.obj"))
                            paths_gt_pred.append(
                                os.path.join(save_results,
                                             str(count_i) + "pred_gt.obj"))

                            count_i += 1
                        save_pred_gt_obj(V_src, V_pred, V_gt, V_tgt_trans,
                                         paths_src, paths_pred, paths_gt,
                                         paths_gt_pred)

    def ablation_study(self, p0, p1, add_noise=False, add_density=False):
        # ablation study
        # mesh = self.plyread("./box1Kinect1.ply")
        # p0 = torch.tensor(mesh).to(device).unsqueeze(0)
        # mesh = self.plyread("./box11.ply")
        # p1 = torch.tensor(mesh).to(device).unsqueeze(0)

        # add noise
        if add_noise:
            p1 = torch.tensor(np.float32(np.random.normal(p1, 0.01)))

        # add outliers
        if add_density:
            density_ratio = 0.5
            pts_num = p1.shape[0]
            sampleNum = int(pts_num *
                            density_ratio)  # the number of remaining points
            if pts_num > sampleNum:
                num = sample(range(1, pts_num), sampleNum)
            elif pts_num > 0:
                num = range(0, pts_num)
            else:
                print("No points in this point cloud!")
                return
            p1 = p1[num, :]
        return p0, p1

    def eval_1__header(self, fout):
        cols = [
            'h_w1', 'h_w2', 'h_w3', 'h_v1', 'h_v2', 'h_v3', 'g_w1', 'g_w2',
            'g_w3', 'g_v1', 'g_v2', 'g_v3'
        ]  # h: estimated, g: ground-truth twist vectors
        print(','.join(map(str, cols)), file=fout)
        fout.flush()

    def eval_1__write(self, fout, ig_gt, g_hat):
        x_hat = se3.log(g_hat)  # --> [-1, 6]
        mx_gt = se3.log(ig_gt)  # --> [-1, 6]
        for i in range(x_hat.size(0)):
            x_hat1 = x_hat[i]  # [6]
            mx_gt1 = mx_gt[i]  # [6]
            vals = torch.cat((x_hat1, -mx_gt1))  # [12]
            valn = vals.cpu().numpy().tolist()
            print(','.join(map(str, valn)), file=fout)
        fout.flush()


import sys
import multiprocessing
from controls import ManualControl
from cam import Camera
from server import get_command_keyboard, stream_frame, get_command
import threading

# Klavye ile hareket için mode = 1
# Sesli komut ile hareket için mode = 2
# Klavye ile hareket ve Aynı anda Raspberryden PC'ye frame aktarma için mode = 3
# Sesli komut ile hareket ve Aynı anda Raspberryden PC'ye frame aktarma için mode = 4
# default mode = 1
mode = 1

def cam(targets, isRead, phase, frm):
    # set camera object with Camera class
    camera = Camera(show=False, captureIndex=-1, camRes=(640, 480))
    camera.set_camera_settings(966.9541358947754)
    camera.set_aruco_settings(markerSize=4, totalMarkers=50, arucoWidth=6)

    while True:
        camera.set_frame()
        isRead.value = camera.isRead
        camera.detect_aruco()
        if camera.target is not None:
            camera.target.set_instant_phase_angle(phase.value)
            targets.append(camera.target)
        frm["data"] = camera.frame
        camera.break_and_release()
        if camera.out:
            break

if __name__ == '__main__':
    manager = multiprocessing.Manager()
    targets = manager.list()
    isRead = multiprocessing.Value('i', 0)
    phase = multiprocessing.Value('i', 0)
    frm = manager.dict()
    frm["command"] = "dur"

    # PC'den raspberry'yi klavye ile kontrol etmek istiyorsanız mode = 1 yapın.
    if mode == 1:
        t1 = threading.Thread(target=get_command_keyboard, args=(frm,))
        t2 = threading.Thread(target=ManualControl.get_command_keyboard_from_pc, args=(frm,))
        try:
            t1.start()
            t2.start()
        except (KeyboardInterrupt, SystemExit):
            sys.exit()

    # PC'den raspberry'yi sesli komut ile kontrol etmek istiyorsanız mode = 2 yapın.
    elif mode == 2:
        t1 = threading.Thread(target=get_command, args=(frm,))
        t2 = threading.Thread(target=ManualControl.speech_move, args=(frm,))
        try:
            t1.start()
            t2.start()
        except (KeyboardInterrupt, SystemExit):
            sys.exit()

    # Klavye ile hareket ve Aynı anda Raspberry'den PC'ye frame aktarma için mode = 3
    elif mode == 3:
        p1 = multiprocessing.Process(target=cam, args=(targets, isRead, phase, frm))
        t1 = threading.Thread(target=stream_frame, args=(frm,))
        t2 = threading.Thread(target=get_command_keyboard, args=(frm,))
        t3 = threading.Thread(target=ManualControl.get_command_keyboard_from_pc, args=(frm,))
        try:
            p1.start()
            t1.start()
            t2.start()
            t3.start()
        except (KeyboardInterrupt, SystemExit):
            sys.exit()

    # Sesli komut ile hareket ve Aynı anda Raspberry'den PC'ye frame aktarma için mode = 4
    elif mode == 4:
        p1 = multiprocessing.Process(target=cam, args=(targets, isRead, phase, frm))
        t1 = threading.Thread(target=stream_frame, args=(frm,))
        t2 = threading.Thread(target=get_command, args=(frm,))
        t3 = threading.Thread(target=ManualControl.speech_move, args=(frm,))
        try:
            p1.start()
            t1.start()
            t2.start()
            t3.start()
        except (KeyboardInterrupt, SystemExit):
            p1.kill()
            sys.exit()


from . animation import Animation

class Off(Animation):
    """A trivial animation that turns all pixels in a layout off."""

    def __init__(self, layout, timeout=1, **kwds):
        super().__init__(layout, **kwds)
        self.internal_delay = timeout

    def step(self, amt=1):
        self.layout.all_off()

from .. util import deprecated
if deprecated.allowed():  # pragma: no cover
    OffAnim = Off


k = int(input())

def mos(n):
    for i in range(len(n)):
        if n[i] == "0":
            n += "1"
        elif n[i] == "1":
            n += "0"
        if len(n) == k:
            return n[k - 1]
    return mos(n)

print(mos("0"))


from sports.nba.nba_team import NBA_Team

class PortlandTrailBlazers(NBA_Team):
    """
    NBA's Portland TrailBlazers Static Information

    """
    full_name = "Portland TrailBlazers"
    name = "TrailBlazers"
    team_id = 1610612757

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


from numpy.lib.polynomial import RankWarning
import torch as pt
import numpy as np
from dataset.GuidedBraTSDataset3D import GuidedBraTSDataset3D
from model.PFSeg import PFSeg3D
import cv2
import SimpleITK as sitk

lr=0.0001
epoch=100
batch_size=1
model_path='/path/to/Saved_models'

img_size=(64,96,96)

model=PFSeg3D().cuda()
model.load_state_dict(pt.load(model_path+'/PFSeg_3D_BraTS_patch-free_bs_best.pt',map_location = 'cpu'))

trainset=GuidedBraTSDataset3D('/path/to/BraTS20',mode='all',augment=False)
# valset=BraTSDataset3D('/path/to/BraTS20',mode='val')
# testset=GuidedBraTSDataset3D('/path/to/BraTS20',mode='test')

train_dataset=pt.utils.data.DataLoader(trainset,batch_size=batch_size,shuffle=False,drop_last=True)
# val_dataset=pt.utils.data.DataLoader(valset,batch_size=1,shuffle=True,drop_last=True)
# test_dataset=pt.utils.data.DataLoader(testset,batch_size=1,shuffle=True,drop_last=True)

def GenerateCoarseMask():
    model.eval()
    dice_sum=0
    hd_sum=0
    jc_sum=0

    for i,data in enumerate(train_dataset):
        output_list=np.zeros((1,1,2*img_size[0],2*img_size[1],2*img_size[2]))
        label_list=np.zeros((1,1,2*img_size[0],2*img_size[1],2*img_size[2]))

        (inputs,labels,raw_image,guidance,_)=data
        labels3D = pt.autograd.Variable(labels).type(pt.FloatTensor).cuda().unsqueeze(1)
        guidance = pt.autograd.Variable(guidance).type(pt.FloatTensor).cuda().unsqueeze(1)

        inputs3D = pt.autograd.Variable(inputs).type(pt.FloatTensor).cuda().unsqueeze(1)
        with pt.no_grad():
            outputs3D,_ = model(inputs3D,guidance)
        outputs3D=np.array(outputs3D.squeeze(0).squeeze(0).cpu().data.numpy())
        output_list=np.zeros((raw_image.shape[1]+64,raw_image.shape[2]+64,raw_image.shape[3]+64))
        output_list[32:-32,32:-32,32:-32]=outputs3D
        label_list=np.zeros((raw_image.shape[1]+64,raw_image.shape[2]+64,raw_image.shape[3]+64))
        label_list[32:-32,32:-32,32:-32]=np.array(labels3D.squeeze(0).squeeze(0).cpu().data.numpy())
        input_real=np.array(raw_image.squeeze(0).numpy())
        input_list=np.zeros((raw_image.shape[1]+64,raw_image.shape[2]+64,raw_image.shape[3]+64))
        input_list[32:-32,32:-32,32:-32]=input_real

        output_list[output_list<0.5]=0.
        output_list[output_list>=0.5]=1.

        results=np.where(output_list!=0)
        x_list=results[0]
        y_list=results[1]
        z_list=results[2]

        x_max=x_list.max()
        x_min=x_list.min()
        y_max=y_list.max()
        y_min=y_list.min()
        z_max=z_list.max()
        z_min=z_list.min()

        x_length=64*(1+(x_max-x_min)//64)  #确保是16的倍数
        y_length=64*(1+(y_max-y_min)//64)
        z_length=64*(1+(z_max-z_min)//64)

        x_center=(x_max-x_min)//2+x_min
        y_center=(y_max-y_min)//2+y_min
        z_center=(z_max-z_min)//2+z_min

        bbox_xmin=x_center-x_length//2
        bbox_xmax=x_center+x_length//2
        bbox_ymin=y_center-y_length//2
        bbox_ymax=y_center+y_length//2
        bbox_zmin=z_center-z_length//2
        bbox_zmax=z_center+z_length//2

        # cropped_coarse=np.zeros((x_length,y_length,z_length))
        # cropped_image=np.zeros((x_length,y_length,z_length))
        # cropped_mask=np.zeros((x_length,y_length,z_length))

        cropped_image=input_list[bbox_xmin:bbox_xmax,bbox_ymin:bbox_ymax,bbox_zmin:bbox_zmax]
        cropped_coarse=output_list[bbox_xmin:bbox_xmax,bbox_ymin:bbox_ymax,bbox_zmin:bbox_zmax]
        cropped_mask=label_list[bbox_xmin:bbox_xmax,bbox_ymin:bbox_ymax,bbox_zmin:bbox_zmax]

        if not(cropped_mask.shape==cropped_image.shape):
            raise Exception()
        if not(cropped_image.shape[0]%16==0 and cropped_image.shape[1]%16==0 and cropped_image.shape[2]%16==0):
            raise Exception()

        # save the cropped images for next round training
        np.save('/path/to/BraTS20/cropped_coarse/Case_{:3d}_64image.npy'.format(i+1),cropped_image)
        np.save('/path/to/BraTS20/cropped_coarse/Case_{:3d}_64coarse.npy'.format(i+1),cropped_coarse)
        np.save('/path/to/BraTS20/cropped_coarse/Case_{:3d}_64mask.npy'.format(i+1),cropped_mask)

        # final_img=np.zeros(shape=(2*img_size[1],2*2*img_size[2]))
        # final_img[:,:2*img_size[2]]=output_list[0,0,64,:,:]*255
        # final_img[:,2*img_size[2]:]=label_list[0,0,64,:,:]*255
        # cv2.imwrite('TestPhase_BraTS.png',final_img)

        pr_sum = output_list.sum()
        gt_sum = label_list.sum()
        pr_gt_sum = np.sum(output_list[label_list == 1])
        dice = 2 * pr_gt_sum / (pr_sum + gt_sum)
        dice_sum += dice
        print("dice:",dice)

        # hausdorff=hd95(output_list.squeeze(0).squeeze(0),label_list.squeeze(0).squeeze(0))
        # jaccard=jc(output_list.squeeze(0).squeeze(0),label_list.squeeze(0).squeeze(0))

        # hd_sum+=hausdorff
        # jc_sum+=jaccard

    print("Finished. Total dice: ",dice_sum/len(train_dataset),'\n')
    print("Finished. Avg Jaccard: ",jc_sum/len(train_dataset))
    print("Finished. Avg hausdorff: ",hd_sum/len(train_dataset))
    return dice_sum/len(train_dataset)

GenerateCoarseMask()


#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
Update map explorers
--------------------

"""
import logging
from os.path import join

from hdx.data.dataset import Dataset
from hdx.data.resource import Resource

from src.acled import update_lc_acled, update_ssd_acled
from src.cbpf import update_cbpf
from src.fts import update_fts
#from src.rowca import update_rowca

logger = logging.getLogger(__name__)

def get_valid_names(downloader, url, headers):
    rows_gen = downloader.get_tabular_rows(url, dict_rows=True, headers=headers)
    return [x['Name'] for x in rows_gen if x['Name'] != 'Name']

def update_resources(resource_updates):
    for resource_info in resource_updates.values():
        resource = Resource.read_from_hdx(resource_info['id'])
        resource.set_file_to_upload(resource_info['path'])
        resource.update_in_hdx()

def update_lc(today, downloader, folder, lc_names_url, lc_mappings_url,
              acled_base_url, fts_base_url, rowca_base_url):
    logger.info('Lake Chad Map Explorer Data')
    country_list = ['Cameroon', 'Nigeria', 'Niger', 'Chad']
    valid_names = get_valid_names(downloader, lc_names_url, headers=['ISO', 'Name'])
    replace_values = downloader.download_tabular_key_value(lc_mappings_url)
    resource_updates = dict()
    resource_updates['acled_events'] = {'id': 'fc396bf2-d204-48b2-84d2-337ada015273',
                                        'path': join(folder, 'Lake_Chad_Basin_Recent_Conflict_Events.csv')}
    resource_updates['acled_fatalities'] = {'id': '3792ee5d-ca30-4e5c-96c8-618c6b625d12',
                                            'path': join(folder, 'Lake_Chad_Basin_Recent_Conflict_Event_Total_Fatalities.csv')}
    resource_updates['fts'] = {'id': '2890c719-4fb2-4178-acdb-e0c5c91cfbce',
                               'path': join(folder, 'Lake_Chad_Basin_Appeal_Status.csv')}
    # resource_updates['rowca_population'] = {'id': '048df35c-e35f-4b1f-aa1a-2d1ce1292f22',
    #                                         'path': join(folder, 'Lake_Chad_Basin_Estimated_Population.csv')}
    # resource_updates['rowca_displaced'] = {'id': '1bdcc8f3-223c-4f7d-9bc6-48be317d50c5',
    #                                        'path': join(folder, 'Lake_Chad_Basin_Displaced.csv')}
    logger.info('Lake Chad - ACLED')
    update_lc_acled(today, acled_base_url, country_list, valid_names, replace_values, resource_updates)
    logger.info('Lake Chad - FTS')
    update_fts(fts_base_url, downloader, country_list, resource_updates)
    # logger.info('Lake Chad - ROWCA')
    # update_rowca(rowca_base_url, downloader, valid_names, replace_values, resource_updates)
    logger.info('Lake Chad - Dataset Date')
    update_resources(resource_updates)
    dataset = Dataset.read_from_hdx('lake-chad-crisis-map-explorer-data')
    dataset.set_dataset_date_from_datetime(today)
    dataset.update_in_hdx()

def update_ssd(today, downloader, folder, ssd_adm1_names_url, ssd_adm2_names_url, ssd_mappings_url,
               acled_base_url, cbpf_base_url):
    logger.info('South Sudan Map Explorer Data')
    country_list = ['South Sudan']
    valid_adm1_names = get_valid_names(downloader, ssd_adm1_names_url, headers=['Name'])
    valid_adm2_names = get_valid_names(downloader, ssd_adm2_names_url, headers=['Name'])
    replace_values = downloader.download_tabular_key_value(ssd_mappings_url)
    resource_updates = dict()
    resource_updates['acled_events'] = {'id': '3480f362-67bb-44d0-b749-9e8fc0963fc0',
                                        'path': join(folder, 'South_Sudan_Recent_Conflict_Events.csv')}
    resource_updates['acled_fatalities'] = {'id': 'a67b85ee-50b4-4345-9102-d88bf9091e95',
                                            'path': join(folder, 'South_Sudan_Recent_Conflict_Event_Total_Fatalities.csv')}
    resource_updates['cbpf'] = {'id': 'd6b18405-5982-4075-bb0a-a1a85f09d842',
                                'path': join(folder, 'South_Sudan_Country_Based_Pool_Funds.csv')}
    logger.info('South Sudan - ACLED')
    update_ssd_acled(today, acled_base_url, country_list, valid_adm2_names, replace_values, resource_updates)
    logger.info('South Sudan - CBPF')
    update_cbpf(cbpf_base_url, downloader, 'SSD19', today, valid_adm1_names, replace_values, resource_updates)
    logger.info('South_Sudan_ - Dataset Date')
    update_resources(resource_updates)
    dataset = Dataset.read_from_hdx('south-sudan-crisis-map-explorer-data')
    dataset.set_dataset_date_from_datetime(today)
    dataset.update_in_hdx()


from flask import Flask, send_file, request, abort
from pathlib import Path
import youtube_dl
import json

app = Flask(__name__)

@app.route('/queuemp3', methods=['GET', 'POST'])
def queuemp3():
    if request.method == 'POST':
        try:
            data = request.get_json()
            url = data['url']
            print(url)

            ydl = youtube_dl.YoutubeDL()
            r = None
            with ydl:
                # don't download, much faster
                r = ydl.extract_info(url, download=False)

            options = {
                'format': 'bestaudio/best',
                'extractaudio': True,  # only keep the audio
                'audioformat': "mp3",  # convert to mp3
                'outtmpl': '{}.mp3'.format(r['title']),  # name the file the ID of the video
                'noplaylist': True,  # only download single song, not playlist
            }

            ''' print some typical fields if needed
            print("%s uploaded by '%s', has %d views, %d likes, and %d dislikes" % (
                r['title'], r['uploader'], r['view_count'], r['like_count'], r['dislike_count']))'''

            with youtube_dl.YoutubeDL(options) as ydl:
                ydl.download([url])
            try:
                return json.dumps({'filename': r['title']})
            except Exception as e:
                return str(e)
        finally:
            print("A request was sent for queueing a conversion")

@app.route('/downloadmp3', methods=['GET', 'POST'])
def downloadmp3():
    if request.method == 'POST':
        filename = request.form['filename']
        print(filename)
        audio_file = Path("./{}.mp3".format(filename))
        if audio_file.is_file():
            return send_file('./{}.mp3'.format(filename),
                             attachment_filename='{}.mp3'.format(filename))
        else:
            abort(404)

if __name__ ==  "__main__":
    app.run(host="0.0.0.0", port=8080, debug=True)


import typing as t
import json
import re
from pathlib import Path
from PIL import Image
from torch.utils.data import Dataset

from .types.marked_image \
    import MarkedImage, MarkedImageTensor
from .transforms import (
    ToTensor
)
from ..utils import coord

class BdcDataSet(Dataset):

    def __init__(self, img_path: str, land_path: str, transform=None):
        super().__init__()

        if transform is None:
            self.transform = ToTensor()
        else:
            self.transform = transform

        self.image_files = [
            p for p in Path(img_path).glob("**/*")
            if re.search('/*.(jpg|png)', str(p))
        ]
        if land_path is not None:
            with open(land_path) as lm:
                landmarks = json.load(lm)
            self.landmarks = self.__normalize_landmarks(landmarks)
        else:
            self.landmarks = {}

    def __len__(self) -> int:
        return len(self.image_files)

    def __getitem__(self, idx: int) -> MarkedImageTensor:
        p = self.image_files[idx]
        with Image.open(str(p)).convert('RGB') as img:
            img.load()
            lmarks = self.landmarks.get(p.name, [])

            sample: MarkedImage = {
                'image': img,
                'landmarks': lmarks
            }

            sample = self.transform(sample)

            return sample

    def __normalize_landmarks(self, landmarks) -> t.Dict:
        norm_lands = {}

        for p in self.image_files:
            lmarks = landmarks[p.name]
            with Image.open(str(p)).convert('RGB') as img:
                img.load()
                norm_lands[p.name] = list(map(
                    lambda x: coord.to_ml_coord(x, img.size),
                    lmarks
                ))

        return norm_lands


class Solution:
    def maxCandies(self, status: List[int], candies: List[int], keys: List[List[int]], containedBoxes: List[List[int]],
                   initialBoxes: List[int]) -> int:
        n = len(status)
        can_open = [status[i] for i in range(n)]
        has_box, used = [False] * n, [False] * n

        q = collections.deque()
        ans = 0

        for box in initialBoxes:
            has_box[box] = True
            if can_open[box]:
                q.append(box)
                used[box] = True
                ans += candies[box]

        while len(q) > 0:
            big_box = q.popleft()
            for key in keys[big_box]:
                can_open[key] = True
                if not used[key] and has_box[key]:
                    q.append(key)
                    used[key] = True
                    ans += candies[key]
            for box in containedBoxes[big_box]:
                has_box[box] = True
                if not used[box] and can_open[box]:
                    q.append(box)
                    used[box] = True
                    ans += candies[box]

        return ans

temp/toy/python/1298. Maximum Candies You Can Get from Boxes.py


import os
import sys
import random
import tables as tb
import numpy  as np
import pandas  as pd

import invisible_cities.reco.paolina_functions as plf
import invisible_cities.reco.dst_functions     as dstf

from   invisible_cities.io.mcinfo_io           import load_mchits
from   invisible_cities.io.mcinfo_io           import load_mcparticles

start = int(sys.argv[1])
numb = int(sys.argv[2])
size = float(sys.argv[3])
blob_radius = float(sys.argv[4])

vox_size = np.array([size,size,size],dtype=np.float16)    # voxel size
pe2keV = 1.

loop_events = []
event, track_ID = [], []
maxR, minX, maxX, minY, maxY, minZ, maxZ = [], [], [], [], [], [], []
evt_energy, energy = [], []
length, numb_of_hits, numb_of_voxels, numb_of_tracks = [], [], [], []
v_size_x, v_size_y, v_size_z = [], [], []
extreme1_x, extreme1_y, extreme1_z = [], [], []
extreme2_x, extreme2_y, extreme2_z = [], [], []
eblob1, eblob2 = [], []
eblob1_bary, eblob2_bary = [], []
blob1_bary_x, blob1_bary_y, blob1_bary_z = [], [], []
blob2_bary_x, blob2_bary_y, blob2_bary_z = [], [], []
event_vxls, track_ID_vxls = [], []
voxel_x, voxel_y, voxel_z = [], [], []
voxel_e = []
signal = []

hits_file = ''
events_in = 0
for n in range(start,start+numb):

    for part in range(10):
        hits_file = '/home/paolafer/data/MC/Tl_upper_port/hits/Tl208_NEW_v1_03_01_nexus_v5_03_04_UPPER_PORT_10.2bar_run4_1hit_perSiPM_hits.{0}_{1}.h5'.format(n, part)

        if not os.path.isfile(hits_file):
            print('{0} not existing'.format(hits_file))
            continue

        print('Analyzing {0}'.format(hits_file))

        hits_dict = load_mchits(hits_file)
        p_dict = load_mcparticles(hits_file)

        events_in += len(hits_dict)

        for nevt, hitc in hits_dict.items():

            tot_e = sum([hh.E for hh in hitc])

            ### smear hit energy to create 1% FWHM resolution at 1592 keV
            sigma_e = 0.01/2.35 * np.sqrt(1.592/tot_e) ### remember, this is relative!
            smeared_tot_e = tot_e + tot_e*np.random.normal(0., 1.) * sigma_e
            sm_factor = smeared_tot_e / tot_e
            #print(tot_e, smeared_tot_e)
            for h in hitc:
                h.energy = h.energy * sm_factor

            voxels = plf.voxelize_hits(hitc, vox_size)
            trks = plf.make_track_graphs(voxels)

            ### Is it a e+e- events?
            positron = False
            for _, particle in p_dict[nevt].items():
                if (particle.name == 'e+') & (len(particle.hits) > 0):
                    positron = True

            for c, t in enumerate(trks, 0):

                etrk = sum([vox.E for vox in t.nodes()])
                extr1, extr2 = plf.find_extrema(t)

                ## first way to calculate blobs: using hits within a sphere from the extremes
                e_blob1 = e_blob2 = 0.
                for h in hitc:
                    dist1 = np.linalg.norm(h.pos - extr1.pos)
                    dist2 = np.linalg.norm(h.pos - extr2.pos)
                    if dist1 < blob_radius:
                        e_blob1 += h.E
                    if dist2 < blob_radius:
                        e_blob2 += h.E

                if (e_blob2 > e_blob1):
                    e_blob1, e_blob2 = e_blob2, e_blob1

                ## second way to calculate blob (a la Michel)
                positions1 = [h.pos for h in extr1.hits]
                qs1 = [h.E for h in extr1.hits]
                if sum(qs1):
                    bary_pos1 = np.average(positions1, weights=qs1, axis=0)
                else:
                    bary_pos1 = extr1.pos

                positions2 = [h.pos for h in extr2.hits]
                qs2 = [h.E for h in extr2.hits]
                if sum(qs2):
                    bary_pos2 = np.average(positions2, weights=qs2, axis=0)
                else:
                    bary_pos2 = extr2.pos

                e_blob1_bary = e_blob2_bary = 0.
                for h in hitc:
                    dist1 = np.linalg.norm(h.pos - bary_pos1)
                    dist2 = np.linalg.norm(h.pos - bary_pos2)
                    if dist1 < blob_radius:
                        e_blob1_bary += h.E
                    if dist2 < blob_radius:
                        e_blob2_bary += h.E

                if (e_blob2_bary > e_blob1_bary):
                    e_blob1_bary, e_blob2_bary = e_blob2_bary, e_blob1_bary

                ## event-related
                event += [nevt]
                signal += [positron]
                evt_energy += [tot_e/pe2keV]
                numb_of_hits += [len(hitc)]
                v_size_x += [voxels[0].size[0]]
                v_size_y += [voxels[0].size[1]]
                v_size_z += [voxels[0].size[2]]

                ## track-related
                track_ID += [c]
                length += [plf.length(t)]
                energy += [etrk/pe2keV]
                numb_of_voxels += [len(t.nodes())]
                numb_of_tracks += [len(trks)]
                extreme1_x += [extr1.X]
                extreme1_y += [extr1.Y]
                extreme1_z += [extr1.Z]
                extreme2_x += [extr2.X]
                extreme2_y += [extr2.Y]
                extreme2_z += [extr2.Z]
                eblob1 += [e_blob1/pe2keV]
                eblob2 += [e_blob2/pe2keV]
                eblob1_bary += [e_blob1_bary/pe2keV]
                eblob2_bary += [e_blob2_bary/pe2keV]
                blob1_bary_x += [bary_pos1[0]]
                blob1_bary_y += [bary_pos1[1]]
                blob1_bary_z += [bary_pos1[2]]
                blob2_bary_x += [bary_pos2[0]]
                blob2_bary_y += [bary_pos2[1]]
                blob2_bary_z += [bary_pos2[2]]

                min_x = 1e+06
                max_x = -1e+06
                min_y = 1e+06
                max_y = -1e+06
                min_z = 1e+06
                max_z = 0.
                max_r = 0

                for v in t.nodes():
                    ## voxel-related
                    event_vxls = event_vxls + [nevt]
                    track_ID_vxls = track_ID_vxls + [c]
                    voxel_x = voxel_x + [v.X]
                    voxel_y = voxel_y + [v.Y]
                    voxel_z = voxel_z + [v.Z]
                    voxel_e = voxel_e + [v.E]

                    for h in v.hits:
                        if h.X < min_x:
                            min_x = h.X
                        if h.X > max_x:
                            max_x = h.X
                        if h.Y < min_y:
                            min_y = h.Y
                        if h.Y > max_y:
                            max_y = h.Y
                        if h.Z < min_z:
                            min_z = h.Z
                        if h.Z > max_z:
                            max_z = h.Z
                        if np.sqrt(h.X*h.X + h.Y*h.Y) > max_r:
                            max_r = np.sqrt(h.X*h.X + h.Y*h.Y)

                minX += [min_x]
                maxX += [max_x]
                minY += [min_y]
                maxY += [max_y]
                minZ += [min_z]
                maxZ += [max_z]
                maxR += [max_r]

loop_events = [events_in]
blob_radius = [blob_radius]

df = pd.DataFrame({ 'event': event, 'evt_energy': evt_energy, 'signal': signal,
                    'minX': minX, 'maxX': maxX, 'minY': minY, 'maxY': maxY, 'minZ': minZ, 'maxZ': maxZ,
                    'maxR': maxR,
                    'numb_of_hits': numb_of_hits, 'energy': energy,
                    'numb_of_tracks': numb_of_tracks, 'length': length, 'track_ID': track_ID,
                    'numb_of_voxels': numb_of_voxels,
                    'voxel_size_x': v_size_x, 'voxel_size_y': v_size_y,
                    'voxel_size_z': v_size_z, 'eblob1': eblob1, 'eblob2': eblob2,
                    'extreme1_x': extreme1_x, 'extreme1_y': extreme1_y, 'extreme1_z': extreme1_z,
                    'extreme2_x': extreme2_x, 'extreme2_y': extreme2_y, 'extreme2_z': extreme2_z,
                    'eblob1_bary': eblob1_bary, 'eblob2_bary': eblob2_bary,
                    'blob1_bary_x': blob1_bary_x, 'blob1_bary_y': blob1_bary_y, 'blob1_bary_z': blob1_bary_z,
                    'blob2_bary_x': blob2_bary_x, 'blob2_bary_y': blob2_bary_y, 'blob2_bary_z': blob2_bary_z,
                     })
df_vxls = pd.DataFrame({'event': event_vxls, 'track_ID': track_ID_vxls,
                        'voxel_x': voxel_x, 'voxel_y': voxel_y, 'voxel_z': voxel_z, 'voxel_e': voxel_e
                        })

df_run_info = pd.DataFrame({'events_in': loop_events, 'blob_radius': blob_radius
                           })

out_name = '/home/paolafer/analysis/tracking_trueinfo_TlMC_run4_vxl{0}mm_R{1}mm_{2}_{3}.hdf5'.format(int(size), int(blob_radius[0]), start, numb)

store = pd.HDFStore(out_name, "w", complib=str("zlib"), complevel=4)
store.put('tracks', df, format='table', data_columns=True)
store.put('voxels', df_vxls, format='table', data_columns=True)
store.put('run_info', df_run_info, format='table', data_columns=True)
store.close()


import os
import sys
import logging

if sys.version_info >= (3, 0):
    from io import StringIO
else:
    try:
        from cStringIO import StringIO
    except ImportError:
        from StringIO import StringIO
        assert StringIO

from pylint import lint
from pylint.__pkginfo__ import numversion

class PyLinter(object):
    """PyLinter class for Anaconda
    """

    def __init__(self, filename, rcfile):
        self.filename = filename
        self.exit = sys.exit
        self.rcfile = rcfile
        self.stdout = sys.stdout
        self.output = StringIO()

        sys.exit = lambda x: None
        sys.stdout = self.output

        self.execute()

    def execute(self):
        """Execute the linting process
        """

        if numversion < (1, 0, 0):
            args = '--include-ids=y -r n'.split(' ')
        else:
            args = '--msg-template={msg_id}:{line}:{column}:{msg} -r n'.split(
                ' ')

        if self.rcfile:
            args.append('--rcfile={0}'.format(os.path.expanduser(self.rcfile)))

        args.insert(0, self.filename)
        lint.Run(args)

    def parse_errors(self):
        """Parse the output given by PyLint
        """

        errors = {'E': [], 'W': [], 'V': []}
        data = self.output.getvalue()

        sys.exit = self.exit
        sys.stdout = self.stdout

        for error in data.splitlines():
            if '************* Module ' in error:
                _, module = error.split('************* Module ')
                if not module in self.filename:
                    continue
            else:
                offset = None
                try:
                    if numversion >= (1, 0, 0):
                        code, line, offset, message = error.split(':', 3)
                    else:
                        code, line, message = error.split(':', 2)
                except ValueError as exception:
                    logging.debug(
                        'unhandled exception in PyLinter parse_errors '
                        'this is a non fatal error: {0}'.format(exception)
                    )
                    logging.debug(
                        'the error string that raised this exception was: '
                        '{0}, please, report this in the GitHub site'.format(
                            error
                        )
                    )
                    continue

                if numversion < (1, 0, 0):
                    try:
                        line, offset = line.split(',')
                    except ValueError:
                        # seems like some versions (or packagers) of pylint
                        # prior to 1.0.0 adds offset to the output but others
                        # doesn't
                        pass

                errors[self._map_code(code)[0]].append({
                    'line': int(line),
                    'offset': offset,
                    'code': self._map_code(code)[1],
                    'message': '[{0}] {1}'.format(
                        self._map_code(code)[1], message
                    )
                })

        return errors

    def _map_code(self, code):
        """Map the given code to fit Anaconda codes
        """

        mapping = {'C': 'V', 'E': 'E', 'F': 'E', 'I': 'V', 'R': 'W', 'W': 'W'}
        return (mapping[code[0]], code[1:])


from lk.utils.config_util import ConfigUtil
from lk.utils.shell_util import run_and_confirm, run, run_and_return_output

from furl import furl

bitbucket = 'bitbucket'
bitbucket_domain = 'bitbucket.org'
github = 'github'
github_domain = 'github.com'

class SourceCodeRepo(object):

    def __init__(self, url=None, service=None, user=None, repo_name=None):

        self._url = url

        self._service = service
        self._user = user
        self._repo_name = repo_name

    @property
    def url(self):

        if self._url:
            return self._url

        else:

            url = 'https://{service_domain}/{user}/{repo}'.format(
                service_domain=self.service_domain,
                user=self.user,
                repo=self.repo_name
            )
            return url

    @property
    def hosting_service_host(self):

        hosting_service_host = self._url.split('/')[2]

        return hosting_service_host

    @property
    def hosting_service(self):

        hosting_service = self.hosting_service_host.split('.')[0]

        return hosting_service

    @property
    def user(self):

        if self._user:
            return self._user

        else:
            user = self._url.split('/')[3]
            return user

    @property
    def repo_name(self):

        if self._repo_name:
            return self._repo_name

        else:
            repo_name = self._url.split('/')[4]
            return repo_name

    @property
    def clone_command(self):

        # https://github.com/lk-commands/default
        # git@github.com:lk-commands/default.git

        # git clone git@bitbucket.org:eyalev/lk-commands.git

        # clone_command = 'git clone git@{hosting_service_host}:{user}/{repo_name}.git'.format(
        # clone_command = 'git clone {repo_url}.git'.format(
        clone_command = 'git clone {git_url}'.format(
            git_url=self.git_url
        )

        return clone_command

    @property
    def git_url(self):

        url = self.url

        if 'github' in url:
            return url

        _furl = furl(url)

        git_url = 'git@{host}:{user}/{repo}.git'.format(
            host=_furl.host,
            user=str(_furl.path).split('/')[1],
            repo=str(_furl.path).split('/')[2]
        )

        return git_url

    def clone(self):

        print('# Cloning lk-repo')

        clone_command = SourceCodeRepo(self.url).clone_command

        command = '{clone_command} {local_repo_path}'.format(
            clone_command=clone_command,
            local_repo_path=self.local_repo_string_path
        )

        run_and_confirm(command)

    @property
    def commands_dir_string_path(self):
        return self.local_repo_string_path + '/commands'

    @property
    def local_repo_string_path(self):

        commands_repo_local_path = '{local_repos_dir}/{repo_service}/{repo_user}/{commands_repo_name}'.format(
            local_repos_dir=ConfigUtil().local_repos_dir,
            repo_service=self.hosting_service,
            repo_user=self.user,
            commands_repo_name=self.repo_name
        )

        return commands_repo_local_path

    @property
    def service(self):

        if self._service:
            return self._service

        if 'bitbucket.org' in self.url:
            return bitbucket

        elif 'github.com' in self.url:
            return github

        else:
            raise NotImplementedError

    @property
    def bitbucket(self):
        return self.service == bitbucket

    @property
    def github(self):
        return self.service == github

    @property
    def service_domain(self):

        if self.bitbucket:
            return bitbucket_domain

        if self.github:
            return github_domain

        else:
            raise NotImplementedError

    def remote_file_source(self, file_name):

        if self.bitbucket:
            shell_command = 'git archive --remote=git@{service_domain}:{user}/{repo}.git HEAD commands/{file_name} | tar -x -O'.format(
                service_domain=self.service_domain,
                user=self.user,
                repo=self.repo_name,
                file_name=file_name
            )
            output = run_and_return_output(shell_command)
            return output

        elif self.github:
            raise NotImplementedError

        else:
            raise NotImplementedError


# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

from collections import deque
class Solution(object):
    def widthOfBinaryTree(self, root):
        """
        :type root: TreeNode
        :rtype: int
        """
        if root == None:
            return 0
        maxWidth = 1
        q = deque([(0, root)])

        while len(q) != 0:
            cnt = len(q)
            start = q[0]
            end = q[-1]
            width = end[0] - start[0] + 1
            maxWidth = max(maxWidth, width)

            while cnt > 0:
                cnt -= 1
                idx, node = q.popleft()
                if node.left != None:
                    q.append((idx * 2, node.left))
                if node.right != None:
                    q.append((idx * 2 + 1, node.right))

        return maxWidth

leetcode/662. Maximum Width of Binary Tree.py


from django.core.management import call_command
from django.core.management.base import BaseCommand, CommandError

import requests, json

from app_comments.models import RedditPost, Comment
from annoying.functions import get_object_or_None
from app_comments.lib.comments import CommentBuilder, RedditPostBuilder
from bs4 import BeautifulSoup
from app_comments.management.commands.get_comments import PostGetter
from time import sleep

class Command(BaseCommand):
    args = ""
    help = ""

    def add_arguments(s, parser):
        parser.add_argument('--url', nargs='+', type=str)

    def process_args(s, options):
        url = options['url'][0] if options['url'] else None
        return url
        # orig_url = url[:]
        # if url:
        #     if url[-5:] != '.json':
        #         url = url[:-1] + '.json'

        # return url, orig_url

    def handle(s, *args, **options):
        #url = s.process_args(options)
        #print(url)
        url = 'https://www.reddit.com/top.json?sort=top&t=year'
        base_url = 'https://www.reddit.com'
        resp = requests.get(url)
        if resp.status_code == 200:
            text_json = resp.text
        else:
            print(resp.text)
            return

        page_json = json.loads(text_json)

        for post_info in page_json['data']['children']:
            comments_url = base_url + post_info['data']['permalink']
            comments_json_url = comments_url[:-1]+'.json'
            pg = PostGetter()
            resp = pg.get(comments_json_url, comments_url)
            print(resp, 1)
            if resp == 'bad http':
                sleep_time = 5
                print('sleeping (%s)...' % sleep_time)
                sleep(sleep_time)
                resp = pg.get(comments_json_url, comments_url)
                if resp == 'bad http':
                    print('sleeping (%s)...' % sleep_time)
                    sleep(sleep_time)
                    resp = pg.get(comments_json_url, comments_url)
                    if resp == 'bad http':
                        print('sleeping (%s)...' % sleep_time)
                        sleep(sleep_time)
            # cmd_data = {'--url': comments_url}
            # call_command('get_comments', **cmd_data)
            # break


import json
import sys
import argparse

sys.path.append("../evaluation")
from evaluate import tuple_f1, convert_opinion_to_tuple

def get_args():
    """
    Helper function to get the gold json, predictions json and negation jsons
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("gold")
    parser.add_argument("predictions")
    parser.add_argument("metadata")
    args = parser.parse_args()
    return args

def open_json(json_file):
    """
    Helper function to open the json files
    """
    with open(json_file) as o:
        file = json.load(o)
    sent_dict = {sent["sent_id"]: sent for sent in file}
    sent_keys = set(sent_dict.keys())
    return sent_keys, sent_dict

def main():
    args = get_args()
    with open(args.metadata) as o:
        metadata = json.load(o)

    test_domains = {}

    gold_keys, gold = open_json(args.gold)
    pred_keys, pred = open_json(args.predictions)

    # get the domains found in the test data
    for sent_id in gold_keys:
        domain = metadata[sent_id[:6]]["category"]
        if domain not in test_domains:
            test_domains[domain] = [sent_id]
        else:
            test_domains[domain].append(sent_id)

    # print the domains in descending order
    for key, value in sorted(test_domains.items(), key=lambda kv: len(kv[1])):
        print("{}:     \t{}".format(key, len(value)))
    print()
    print()

    # get the sentiment graph F1 for each domain
    for domain, sent_ids in sorted(test_domains.items(),
                                   key=lambda kv: len(kv[1])):
        domain_gold = dict([(sent_id, convert_opinion_to_tuple(gold[sent_id]))                for sent_id in sent_ids])
        domain_pred = dict([(sent_id, convert_opinion_to_tuple(pred[sent_id]))                for sent_id in sent_ids])
        f1 = tuple_f1(domain_gold, domain_pred)
        print("{0}: {1:.3f}".format(domain, f1))

if __name__ == "__main__":
    main()


from __future__ import division
import numpy as np
from scipy import sparse
from sklearn.metrics.pairwise import euclidean_distances
import time

# Produce grid points for a 2d grayscale image
def get_points_2d(image, res):
    rows, columns = image.shape
    grid_x, grid_y = np.mgrid[0:columns:res, 0:rows:res]
    grid = np.array((grid_x.flatten(), grid_y.flatten())).T
    return grid

# Produce grid points for a 3d grayscale image
def get_points_3d(image, res):
    rows, columns, z = image.shape
    grid_z, grid_x, grid_y = np.mgrid[0:z:res, 0:columns:res, 0:rows:res]
    grid = np.array((grid_x.flatten(), grid_y.flatten(), grid_z.flatten())).T
    return grid

# Wendland kernel as a function of r = norm(x-y)/c_sup
def dist_kernel(r):
    return max((1-r, 0))**4 * (4*r + 1)

def blowup_S(S, dim):
    (m, n) = S.shape
    if dim == 3:
        S_full = sparse.lil_matrix((3 * m, 3 * n), dtype = np.float32)
        #S_full = np.zeros((3 * m, 3 * n))
        S_full[0::3, 0::3] = S
        S_full[1::3, 1::3] = S
        S_full[2::3, 2::3] = S
    else:
        S_full = np.zeros((2 * m, 2 * n))
        S_full[0::2, 0::2] = S
        S_full[1::2, 1::2] = S
    return S_full.tocsc()

# Create evaluation matrix given kernel centers (grid points), evaluation points
# and kernel support
def evaluation_matrix(kernels, points, c_sup, dim):
    dim = kernels.shape[1]
    vect_kernel = np.vectorize(dist_kernel)
    start = time.time()
    S = euclidean_distances(points, kernels) / c_sup
    #print("VEC -- euc dist ", (time.time() - start) / 60)
    # Mark entries with 0 kernel support
    start = time.time()
    S[np.where(S > 1)] = -1
    non_zero_indices = np.where(S >= 0)
    #print("VEC -- S[np.where(S > 1)] and np.where(S>=0) ", (time.time() - start) / 60)
    # Evaluate kernel at points within support
    start = time.time()
    S[non_zero_indices] = vect_kernel(S[non_zero_indices])
    #print("VEC -- S[non_zero] = vect_kernel ", (time.time() - start) / 60)
    start = time.time()
    S[np.where(S == -1)] = 0
    #print("VEC -- S[np.where(S == -1)] = 0 ", (time.time() - start) / 60)
    start = time.time()
    #full_S = blowup_S_old(S, dim)
    #print("VEC -- blowup ", (time.time() - start) / 60)
    return sparse.csc_matrix(S)

def evaluation_matrix_blowup(kernels, points, c_sup, dim):
    dim = kernels.shape[1]
    vect_kernel = np.vectorize(dist_kernel)
    start = time.time()
    S = euclidean_distances(points, kernels) / c_sup
    #print("VEC -- euc dist ", (time.time() - start) / 60)
    # Mark entries with 0 kernel support
    start = time.time()
    S[np.where(S > 1)] = -1
    non_zero_indices = np.where(S >= 0)
    #print("VEC -- S[np.where(S > 1)] and np.where(S>=0) ", (time.time() - start) / 60)
    # Evaluate kernel at points within support
    start = time.time()
    S[non_zero_indices] = vect_kernel(S[non_zero_indices])
    #print("VEC -- S[non_zero] = vect_kernel ", (time.time() - start) / 60)
    start = time.time()
    S[np.where(S == -1)] = 0
    #print("VEC -- S[np.where(S == -1)] = 0 ", (time.time() - start) / 60)
    start = time.time()
    full_S = blowup_S(S, dim)
    #print("VEC -- blowup ", (time.time() - start) / 60)
    return full_S

# Create velocity field by weighing kernels by alphas
def make_V(S, alpha, d):
    alpha = alpha.flatten()
    if (S.shape[1] == alpha.shape[0]):
        lmda = S.dot(alpha)
        return lmda.reshape(-1, d)
    else:
        alpha = alpha.reshape(-1, d)
        return S.dot(alpha)


from vector import Vector
import turtle

scale = 40

def print_vector(vector, color):
    turtle.pencolor(color)
    turtle.penup()
    turtle.home()
    turtle.pendown()
    turtle.goto(vector.elements[0]*scale,vector.elements[1]*scale)

def print_system(x,y):
    turtle.home()
    for i in range(x):
        turtle.dot(3)
        turtle.write(i, align='right')
        turtle.setx(scale*(i+1))

    turtle.home()
    for j in range(y):
        turtle.dot(3)
        turtle.write(j, align='right')
        turtle.sety(scale*(j+1))

turtle.speed(10)
print_system(10,10)

vector1 = Vector([3, 2])

print_vector(vector1, 'red')

vector2 = Vector([1,-4])

print_vector(vector2, 'blue')

vector1.add_vector(vector2)

print_vector(vector1, 'green')

turtle.done()


import sys
sys.stdin = open('../input.txt', 'r')

N = int(input())
numbers = list(map(int, sys.stdin.readline().split()))

min_num, max_num = 1000000, -1000000

for number in numbers:
    if number < min_num:
        min_num = number
    if number > max_num:
        max_num = number

print(min_num, max_num)


from matplotlib import pyplot as plt
from numpy import loadtxt, zeros
from skimage.measure import label
from os import path

if __name__ == '__main__':
    current_dir = path.dirname(__file__)
    file_names = ['mat_p0.70.dat', 'mat_p0.72.dat']

    for file_name in file_names:
        file_path = path.join(current_dir, file_name)
        lattice = loadtxt(file_path)

        # change connectivity to 2 if you want to consider Moore neighborhood
        labelled_lattice = label(lattice, background=0, connectivity=1)
        num_clusters = labelled_lattice.max()

        cluster_sizes = []
        for cluster_id in range(1, num_clusters + 1):
            cluster_sizes.append((labelled_lattice == cluster_id).sum())

        cluster_size_distribution = zeros(max(cluster_sizes))
        for cluster_size in cluster_sizes:
            cluster_size_distribution[cluster_size - 1] += 1

        inverse_cdf = zeros(max(cluster_sizes))
        for cluster_size in range(max(cluster_sizes)):
            inverse_cdf[cluster_size] = (cluster_size_distribution[cluster_size:]).sum()
        inverse_cdf /= sum(cluster_size_distribution)

        plt.figure(figsize=(11, 5))
        plt.subplot(1, 2, 1)
        plt.title(f"Lattice from {file_name}")
        plt.imshow(lattice)

        plt.subplot(1, 2, 2)
        plt.title("Cluster Size Distribution")
        plt.xlabel("Cluster Size s")
        plt.ylabel("P(S > s)")
        plt.loglog(range(1, max(cluster_sizes) + 1), inverse_cdf, 'bo')
        plt.show()