#ctrgcn_stage1.py
import math
import pdb
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
from numpy import random
from graph.ntu_rgb_d import AdjMatrixGraph
from itertools import permutations
from apex import amp, optimizers
#from graph.ntu_rgb_d import AdjMatrixGraph, AdjMatrixGraph4lh, AdjMatrixGraph4rh, AdjMatrixGraph4ll, AdjMatrixGraph4rl, AdjMatrixGraph4torso

def import_class(name):
    components = name.split('.')
    mod = __import__(components[0])
    for comp in components[1:]:
        mod = getattr(mod, comp)
    return mod

def conv_branch_init(conv, branches):
    weight = conv.weight
    n = weight.size(0)
    k1 = weight.size(1)
    k2 = weight.size(2)
    nn.init.normal_(weight, 0, math.sqrt(2. / (n * k1 * k2 * branches)))
    nn.init.constant_(conv.bias, 0)

def conv_init(conv):
    if conv.weight is not None:
        nn.init.kaiming_normal_(conv.weight, mode='fan_out')
    if conv.bias is not None:
        nn.init.constant_(conv.bias, 0)

def bn_init(bn, scale):
    nn.init.constant_(bn.weight, scale)
    nn.init.constant_(bn.bias, 0)

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        if hasattr(m, 'weight'):
            nn.init.kaiming_normal_(m.weight, mode='fan_out')
        if hasattr(m, 'bias') and m.bias is not None and isinstance(m.bias, torch.Tensor):
            nn.init.constant_(m.bias, 0)
    elif classname.find('BatchNorm') != -1:
        if hasattr(m, 'weight') and m.weight is not None:
            m.weight.data.normal_(1.0, 0.02)
        if hasattr(m, 'bias') and m.bias is not None:
            m.bias.data.fill_(0)

class TemporalConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1):
        super(TemporalConv, self).__init__()
        pad = (kernel_size + (kernel_size-1) * (dilation-1) - 1) // 2
        self.conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=(kernel_size, 1),
            padding=(pad, 0),
            stride=(stride, 1),
            dilation=(dilation, 1))

        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return x

class MultiScale_TemporalConv(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 stride=1,
                 dilations=[1,2,3,4],
                 residual=True,
                 residual_kernel_size=1):

        super().__init__()
        assert out_channels % (len(dilations) + 2) == 0, '# out channels should be multiples of # branches'

        # multiple branches of temporal convolution
        self.num_branches = len(dilations) + 2
        branch_channels = out_channels // self.num_branches
        if type(kernel_size) == list:
            assert len(kernel_size) == len(dilations)
        else:
            kernel_size = [kernel_size]*len(dilations)
        # temporal convolution branches
        self.branches = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(
                    in_channels,
                    branch_channels,
                    kernel_size=1,
                    padding=0),
                nn.BatchNorm2d(branch_channels),
                nn.ReLU(inplace=True),
                TemporalConv(
                    branch_channels,
                    branch_channels,
                    kernel_size=ks,
                    stride=stride,
                    dilation=dilation),
            )
            for ks, dilation in zip(kernel_size, dilations)
        ])

        # additional max & 1x1 branch
        self.branches.append(nn.Sequential(
            nn.Conv2d(in_channels, branch_channels, kernel_size=1, padding=0),
            nn.BatchNorm2d(branch_channels),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=(3,1), stride=(stride,1), padding=(1,0)),
            nn.BatchNorm2d(branch_channels)
        ))

        self.branches.append(nn.Sequential(
            nn.Conv2d(in_channels, branch_channels, kernel_size=1, padding=0, stride=(stride,1)),
            nn.BatchNorm2d(branch_channels)
        ))

        # residual connection
        if not residual:
            self.residual = lambda x: 0
        elif (in_channels == out_channels) and (stride == 1):
            self.residual = lambda x: x
        else:
            self.residual = TemporalConv(in_channels, out_channels, kernel_size=residual_kernel_size, stride=stride)

        # initialize
        self.apply(weights_init)

    def forward(self, x):
        # input dim: (N,C,T,V)
        res = self.residual(x)
        branch_outs = []
        for tempconv in self.branches:
            out = tempconv(x)
            branch_outs.append(out)

        out = torch.cat(branch_outs, dim=1)
        out += res
        return out

class CTRGC(nn.Module):
    def __init__(self, in_channels, out_channels, rel_reduction=8, mid_reduction=1):
        super(CTRGC, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        if in_channels == 3 or in_channels == 9:
            self.rel_channels = 8
            self.mid_channels = 16
        else:
            self.rel_channels = in_channels // rel_reduction
            self.mid_channels = in_channels // mid_reduction
        self.conv1 = nn.Conv2d(self.in_channels, self.rel_channels, kernel_size=1)
        self.conv2 = nn.Conv2d(self.in_channels, self.rel_channels, kernel_size=1)
        self.conv3 = nn.Conv2d(self.in_channels, self.out_channels, kernel_size=1)
        self.conv4 = nn.Conv2d(self.rel_channels, self.out_channels, kernel_size=1)
        self.tanh = nn.Tanh()
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                conv_init(m)
            elif isinstance(m, nn.BatchNorm2d):
                bn_init(m, 1)

    def forward(self, x, A=None, alpha=1):
        x1, x2, x3 = self.conv1(x).mean(-2), self.conv2(x).mean(-2), self.conv3(x)
        x1 = self.tanh(x1.unsqueeze(-1) - x2.unsqueeze(-2))
        x1 = self.conv4(x1) * alpha + (A.unsqueeze(0).unsqueeze(0) if A is not None else 0)  # N,C,V,V
        x1 = torch.einsum('ncuv,nctv->nctu', x1, x3)
        return x1

class unit_tcn(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=9, stride=1):
        super(unit_tcn, self).__init__()
        pad = int((kernel_size - 1) / 2)
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=(kernel_size, 1), padding=(pad, 0), stride=(stride, 1))
        self.bn = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        conv_init(self.conv)
        bn_init(self.bn, 1)

    def forward(self, x):
        x = self.bn(self.conv(x))
        return x

class unit_tcn1(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=2, stride=2):
        super(unit_tcn1, self).__init__()
        pad = int((kernel_size - 1) / 2)
        self.conv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=(kernel_size, 1), padding=(0, 0), stride=(stride, 1))
        self.bn = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        conv_init(self.conv)
        bn_init(self.bn, 1)

    def forward(self, x):
        x = self.bn(self.conv(x))
        return x

class unit_gcn(nn.Module):
    def __init__(self, in_channels, out_channels, A, coff_embedding=4, adaptive=True, residual=True):
        super(unit_gcn, self).__init__()
        inter_channels = out_channels // coff_embedding
        self.inter_c = inter_channels
        self.out_c = out_channels
        self.in_c = in_channels
        self.adaptive = adaptive
        self.num_subset = A.shape[0]
        self.convs = nn.ModuleList()
        for i in range(self.num_subset):
            self.convs.append(CTRGC(in_channels, out_channels))

        if residual:
            if in_channels != out_channels:
                self.down = nn.Sequential(
                    nn.Conv2d(in_channels, out_channels, 1),
                    nn.BatchNorm2d(out_channels)
                )
            else:
                self.down = lambda x: x
        else:
            self.down = lambda x: 0
        if self.adaptive:
            self.PA = nn.Parameter(torch.from_numpy(A.astype(np.float32)))
        else:
            self.A = Variable(torch.from_numpy(A.astype(np.float32)), requires_grad=False)
        self.alpha = nn.Parameter(torch.zeros(1))
        self.bn = nn.BatchNorm2d(out_channels)
        self.soft = nn.Softmax(-2)
        self.relu = nn.ReLU(inplace=True)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                conv_init(m)
            elif isinstance(m, nn.BatchNorm2d):
                bn_init(m, 1)
        bn_init(self.bn, 1e-6)

    def forward(self, x):
        y = None
        if self.adaptive:
            A = self.PA
        else:
            A = self.A.cuda(x.get_device())
        for i in range(self.num_subset):
            z = self.convs[i](x, A[i], self.alpha)
            y = z + y if y is not None else z
        y = self.bn(y)
        y += self.down(x)
        y = self.relu(y)

        return y

# basic action feature learning and revised cross-view action feature learning
class cv_TCN_GCN_unit(nn.Module):
    def __init__(self, in_channels, out_channels, A, stride=1, residual=True, adaptive=True, kernel_size=5, dilations=[1,2]):
        super(cv_TCN_GCN_unit, self).__init__()
        self.gcn1 = unit_gcn(in_channels, out_channels, A, adaptive=adaptive)
        self.gcn2 = unit_gcn(out_channels, out_channels, A, adaptive=adaptive)
        self.tcn1 = MultiScale_TemporalConv(out_channels, out_channels, kernel_size=kernel_size, stride=stride, dilations=dilations, residual=False)
        self.tcn2 = MultiScale_TemporalConv(out_channels, out_channels, kernel_size=kernel_size, stride=1, dilations=dilations, residual=False)
        self.relu = nn.ReLU(inplace=True)
        if not residual:
            self.residual = lambda x: 0
        elif (in_channels == out_channels) and (stride == 1):
            self.residual = lambda x: x
        else:
            self.residual = unit_tcn(in_channels, out_channels, kernel_size=1, stride=stride)

    def forward(self, x, case, vlist): # or x, case, vlist, xr
        # case=0 means basic, case=1 means cross-view, vlist means view list！！！！！！！！！！！！
        if case == 0:
            y = self.relu(self.tcn1(self.gcn1(x)) + self.residual(x))
        else:
            # '''
            # CV-FCL (CV-STGCN) or GL-CVFD (only CV-GCN), the first V-FCL or C-CVFD training stage     第一阶段！！！
            temp = x # x.shape = ([20, 3, 100, 25])
            x = self.gcn1(x)
            # set batch size / gpus
            z = torch.zeros(16,2,x.size(1),x.size(2),x.size(3)).cuda().half()
            x = x.view(16,2,x.size(1),x.size(2),x.size(3))
            # cross-view spatial information interaction (25 joints, 50% probability)
            for i,view in enumerate(vlist):
                for j in range(25):
                    if random.choice([0,1]) == 0:
                        z[i,:,:,:,j] = x[i,:,:,:,j]
                    else:
                        z[i,:,:,:,j] = x[view,:,:,:,j]
            z = z.contiguous().view(16*2,x.size(2),x.size(3),x.size(4))
            # gcn2, tcn1
            z = self.tcn1(self.gcn2(z))
            # set batch size as 16
            m = torch.zeros(16,2,z.size(1),z.size(2),z.size(3)).cuda().half()
            z = z.view(16,2,z.size(1),z.size(2),z.size(3))
            # cross-view temporal information interaction (T frames, 50% probability)
            for i,view in enumerate(vlist):
                for j in range(z.size(3)):
                    if random.choice([0,1]) == 0:
                        m[i,:,:,j,:] = z[i,:,:,j,:]
                    else:
                        m[i,:,:,j,:] = z[view,:,:,j,:]
            m = m.contiguous().view(16*2,z.size(2),z.size(3),z.size(4))
            # tcn2
            y = self.relu(self.tcn2(m)+self.residual(temp))
            #y = self.relu(self.tcn1(self.gcn2(z)) + self.residual(temp))

        return y

class TCN_GCN_unit1(nn.Module):
    def __init__(self, in_channels, out_channels, A, residual=True, adaptive=True, kernel_size=1, stride=1):
        super(TCN_GCN_unit1, self).__init__()
        self.gcn1 = unit_gcn(in_channels, out_channels, A, adaptive=adaptive)
        self.tcn1 = unit_tcn1(out_channels, out_channels, kernel_size=kernel_size, stride=stride)
        self.relu = nn.ReLU(inplace=True)
        if not residual:
            self.residual = lambda x: 0
        elif (in_channels == out_channels) and (stride == 1):
            self.residual = lambda x: x
        else:
            self.residual = unit_tcn(in_channels, out_channels, kernel_size=1, stride=stride)

    def forward(self, x):
        y = self.relu(self.tcn1(self.gcn1(x)) + self.residual(x))
        return y

class ST_Joint_Att(nn.Module):
    def __init__(self, channel, reduct_ratio, bias):
        super(ST_Joint_Att, self).__init__()

        inner_channel = channel // reduct_ratio

        self.fcn = nn.Sequential(
            nn.Conv2d(channel, inner_channel, kernel_size=1, bias=bias),
            nn.BatchNorm2d(inner_channel),
            nn.Hardswish(),
        )
        self.conv_t = nn.Conv2d(inner_channel, channel, kernel_size=1)
        self.conv_v = nn.Conv2d(inner_channel, channel, kernel_size=1)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                conv_init(m)
            elif isinstance(m, nn.BatchNorm2d):
                bn_init(m, 1)

    def forward(self, x):
        N, C, T, V = x.size()
        x_t = x.mean(3, keepdims=True)
        x_v = x.mean(2, keepdims=True).transpose(2, 3)
        x_att = self.fcn(torch.cat([x_t, x_v], dim=2))
        x_t, x_v = torch.split(x_att, [T, V], dim=2)
        x_t_att = self.conv_t(x_t).sigmoid()
        x_v_att = self.conv_v(x_v.transpose(2, 3)).sigmoid()
        x_att = x_t_att * x_v_att
        x = x * x_att
        return x

class Model(nn.Module):
    def __init__(self, num_class=40, num_point=25, num_person=2, graph=None, in_channels=3, drop_out=0, adaptive=True):
        super(Model, self).__init__()

        Graph = AdjMatrixGraph()
        A = Graph.A_binary

        self.num_class = num_class
        self.num_point = num_point
        self.data_bn = nn.BatchNorm1d(num_person * in_channels * num_point)

        base_channel = 64
        self.l1 = cv_TCN_GCN_unit(in_channels, base_channel, A, residual=False, adaptive=adaptive) #
        self.att1 = ST_Joint_Att(base_channel, 4, True) # ##
        self.l2 = cv_TCN_GCN_unit(base_channel, base_channel, A, adaptive=adaptive)
        self.l3 = cv_TCN_GCN_unit(base_channel, base_channel, A, adaptive=adaptive) # ##
        self.l4 = cv_TCN_GCN_unit(base_channel, base_channel, A, adaptive=adaptive) # ##
        self.l5 = cv_TCN_GCN_unit(base_channel, base_channel*2, A, stride=2, adaptive=adaptive)
        self.att2 = ST_Joint_Att(base_channel*2, 4, True) # ##
        self.l6 = cv_TCN_GCN_unit(base_channel*2, base_channel*2, A, adaptive=adaptive)
        self.l7 = cv_TCN_GCN_unit(base_channel*2, base_channel*2, A, adaptive=adaptive) # ##
        self.l8 = cv_TCN_GCN_unit(base_channel*2, base_channel*4, A, stride=2, adaptive=adaptive)
        self.att3 = ST_Joint_Att(base_channel*4, 4, True) # ##
        self.l9 = cv_TCN_GCN_unit(base_channel*4, base_channel*4, A, adaptive=adaptive) # ##
        self.l10 = cv_TCN_GCN_unit(base_channel*4, base_channel*4, A, adaptive=adaptive)

        # original classifier
        #self.fc = nn.Linear(base_channel*4, num_class)
        # transform feature-dim from 256 to 128 for view-semantic Fisher contrastive learning
        self.fc = nn.Linear(base_channel*4, 128)
        self.bn = nn.BatchNorm1d(128)
        self.relu = nn.ReLU(True)
        self.bn1 = nn.BatchNorm1d(256)
        self.relu1 = nn.ReLU(True)
        # transform feature-dim from 128 to 40/8, 60/3, 120/3 for semantic/view discrimination
        self.prototypes = nn.Linear(128, 40, bias=False)
        self.prototypev = nn.Linear(128, 8, bias=False)

        nn.init.normal_(self.fc.weight, 0, math.sqrt(2. / 128))
        nn.init.normal_(self.prototypev.weight, 0, math.sqrt(2. / 8))
        nn.init.normal_(self.prototypes.weight, 0, math.sqrt(2. / 40))

        bn_init(self.data_bn, 1)
        bn_init(self.bn, 1)
        bn_init(self.bn1, 1)

        if drop_out:
            self.drop_out = nn.Dropout(drop_out)
        else:
            self.drop_out = lambda x: x

    def forward(self, x): # x (V-FCL) or x, xr (S-FCL)
        N, C, T, V, M = x.size()

        x = x.permute(0, 4, 3, 1, 2).contiguous().view(N, M * V * C, T)
        x = self.data_bn(x)
        x = x.view(N, M, V, C, T).permute(0, 1, 3, 4, 2).contiguous().view(N * M, C, T, V)

        # # assistant view
        # xr = xr.permute(0, 4, 3, 1, 2).contiguous().view(N, M * V * C, T)
        # xr = self.data_bn(xr)
        # xr = xr.view(N, M, V, C, T).permute(0, 1, 3, 4, 2).contiguous().view(N * M, C, T, V)

        # corrsponding to batch size / gpus in cv_TCN_GCN_unit
        vlist = [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]
        random.shuffle(vlist)
        xcv = x
        # xcv1 = xr  #

        x = self.l1(x,0,vlist)
        xcv = self.l1(xcv,1,vlist)
        x = self.att1(x)
        x = self.l2(x,0,vlist)
        xcv = self.l2(xcv,1,vlist)
        x = self.l3(x,0,vlist)
        xcv = self.l3(xcv,1,vlist)
        x = self.l4(x,0,vlist)
        xcv = self.l4(xcv,1,vlist)
        x = self.l5(x,0,vlist)
        xcv = self.l5(xcv,1,vlist)
        x = self.att2(x)
        x = self.l6(x,0,vlist)
        xcv = self.l6(xcv,1,vlist)
        x = self.l7(x,0,vlist)
        xcv = self.l7(xcv,1,vlist)
        x = self.l8(x,0,vlist)
        xcv = self.l8(xcv,1,vlist)
        x = self.att3(x)
        x = self.l9(x,0,vlist)
        xcv = self.l9(xcv,1,vlist)
        x = self.l10(x,0,vlist)
        xcv = self.l10(xcv,1,vlist)

        # original version, N*M,C,T,V
        c_new = x.size(1)
        x = x.view(N, M, c_new, -1)
        x = x.mean(3).mean(1)
        x = self.drop_out(x)

        # the first view-term training stage
        c_new1 = xcv.size(1)
        xcv = xcv.view(N, M, c_new1, -1)
        xcv = xcv.mean(3).mean(1)
        xcv = self.drop_out(xcv)
        tempcv = xcv # view-common feature
        tempv = x-xcv # view-specific feature

        outv = self.relu1(self.bn1(tempv))
        outv = self.fc(outv) # same-action different-view view-specific feature
        outp = self.relu1(self.bn1(tempcv))
        outp = self.fc(outp) # same-action different-view view-common feature or different-action view-common feature
        fv = self.relu(self.bn(outv))
        fv = self.prototypev(fv) # view discrimination vector based on outv

        return outv,outp,fv

class CTRGCN(nn.Module):
    def __init__(self, num_class=40, num_point=25, num_person=2, graph=None, in_channels=3, drop_out=0, adaptive=True):
        super(CTRGCN, self).__init__()

        self.num_class = num_class
        self.num_point = num_point
        self.data_bn = nn.BatchNorm1d(150)

        self.encoderp_s = Model(num_class=40, num_point=25, num_person=2, graph=None, in_channels=3, drop_out=0, adaptive=True)

        bn_init(self.data_bn, 1)
        #self.data_bn = nn.BatchNorm1d(num_person * in_channels * num_point)

    def forward(self, x, viewt):

        # xr = _transform_rotation(x) # assistant view
        x = _transform_shape(x)
        N, C, T, V, M = x.size()
        xjmo = torch.zeros(N, C, T, V, M).cuda().half() # joint motion
        # xjmor = torch.zeros(N, C, T, V, M).cuda().half() # assistant-view joint motion
        for i in range(T-1):
            xjmo[:,:,i,:,:] = x[:,:,i+1,:,:] - x[:,:,i,:,:]
            # xjmor[:,:,i,:,:] = xr[:,:,i+1,:,:] - xr[:,:,i,:,:]

        outv, outp, fv = self.encoderp_s(xjmo) #
        # outp, outp1, fp, fp1 = self.encoderp_s(xjmo, xjmor) #

        return outv, outp, fv

#
class MyCTRGCN(nn.Module):
    def __init__(self, num_class=40, num_point=25, num_person=2, graph=None, in_channels=3, drop_out=0, adaptive=True):
        super(MyCTRGCN, self).__init__()
        self.net = CTRGCN(num_class=40, num_point=25, num_person=2, graph=None, in_channels=3, drop_out=0, adaptive=True)

        self.net = torch.nn.DataParallel(self.net)
        #self.net = torch.nn.parallel.DistributedDataParallel(self.net) # model distributed training setting

    def forward(self, x, viewt):
        return self.net(x, viewt)

def _transform_shape(inputs1):
    a1, a2 = inputs1.split([75, 75], dim=2)
    b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, b11, b12, b13, b14, b15, b16, b17, b18, b19, b20, b21, b22, b23, b24, b25 = a1.chunk(25, dim=2)
    inputs1 = torch.stack([b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, b11, b12, b13, b14, b15, b16, b17, b18, b19, b20, b21, b22, b23, b24, b25], dim=3)
    inputs1 = inputs1.permute(0, 2, 1, 3)
    inputs1 = inputs1.contiguous().view(-1, 3, 3014, 25)
    inputs1 = torch.split(inputs1, [300,2714], dim=2)[0]
    #inputs1 = inputs1.contiguous().view(-1, 3, 300, 25)
    inputs1 = torch.index_select(inputs1, 2, torch.arange(1,200,2).cuda())
    b31, b32, b33, b34, b35, b36, b37, b38, b39, b40, b41, b42, b43, b44, b45, b46, b47, b48, b49, b50, b51, b52, b53, b54, b55 = a2.chunk(25, dim=2)
    inputs2 = torch.stack([b31, b32, b33, b34, b35, b36, b37, b38, b39, b40, b41, b42, b43, b44, b45, b46, b47, b48, b49, b50, b51, b52, b53, b54, b55], dim=3)
    inputs2 = inputs2.permute(0, 2, 1, 3)
    inputs2 = inputs2.contiguous().view(-1, 3, 3014, 25)
    inputs2 = torch.split(inputs2, [300,2714], dim=2)[0]
    #inputs2 = inputs2.contiguous().view(-1, 3, 300, 25)
    inputs2 = torch.index_select(inputs2, 2, torch.arange(1,200,2).cuda())
    inputs = torch.stack([inputs1, inputs2], dim=4)

    return inputs


import numpy as np
import matplotlib.pyplot as plt

def print_cloud(points, l_x, l_y, delta):
    plt.figure()
    plt.scatter(points[:, 0], points[:, 1], s=l_x, alpha=1)

    plt.plot((0, 0), (0, l_y), 'k-', color='r')
    plt.plot((0, l_x), (0, 0), 'k-', color='r')
    plt.plot((0, l_x), (l_y, l_y), 'k-', color='r')
    plt.plot((l_x, l_x), (0, l_y), 'k-', color='r')

    delta_x = l_x * delta
    delta_y = l_y * delta
    plt.plot((delta_x, delta_x), (0, l_y), 'k-', color='r')
    plt.plot((0, l_x), (delta_y, delta_y), 'k-', color='r')
    plt.plot((0, l_x), (l_y - delta_y, l_y - delta_y), 'k-', color='r')
    plt.plot((l_x - delta_x, l_x - delta_x), (0, l_y), 'k-', color='r')

    plt.show()

def uniform_node_clouds(l_x, l_y, delta, n_points):
    xy_min = [0, 0]
    xy_max = [l_x, l_y]

    dot_dist = np.sqrt(l_x * l_y / n_points)

    x_line_qty = int(l_x/dot_dist)
    y_line_qty = int(l_y/dot_dist)

    points = np.random.uniform(low=xy_min, high=xy_max, size=(n_points, 2))
    x_line_points = np.linspace(0, l_x, x_line_qty)
    y_line_points = np.linspace(0, l_y, y_line_qty)

    x_lines = (0, delta * l_x, (1 - delta) * l_x, l_x)
    y_lines = (0, delta * l_y, (1 - delta) * l_y, l_y)

    for i in range(0, 4):
        x_line = np.concatenate((x_line_points[:, np.newaxis], np.repeat(x_lines[i], x_line_qty)[:, np.newaxis]), axis=1)
        y_line = np.concatenate((np.repeat(y_lines[i], y_line_qty)[:, np.newaxis], y_line_points[:, np.newaxis]), axis=1)

        points = np.concatenate((points, x_line))
        points = np.concatenate((points, y_line))

    return points


import importlib

from easydict import EasyDict as edict

__C = edict()
cfg = __C

__C.MODEL_NAME = ""  # this name would be the result file name
__C.MODULE = ""  # sample: dgl.network  b_global.network

__C.BATCH_SIZE = 32
__C.NUM_WORKERS = 8

__C.LOG_FILE_NAME = ""  # the suffix of log file

__C.MODEL_SAVE_PATH = ""  # auto generated

#
# Dataset
#
__C.DATASET = ""

# Other dataset specific configs should be imported from dataset_config.py

# wandb project name
__C.PROJECT = ""

#
# Training options
#

__C.TRAIN = edict()

# Total epochs
__C.TRAIN.NUM_EPOCHS = 200

# Optimizer type
__C.TRAIN.OPTIMIZER = "SGD"

# Start learning rate
__C.TRAIN.LR = 0.001

# LR Scheduler
__C.TRAIN.LR_SCHEDULER = "cosine"

# Learning rate decay
__C.TRAIN.LR_DECAY = 100.0

# Learning rate decay step (in epochs)
__C.TRAIN.LR_STEP = [10, 20]

# warmup_ratio for Adam Cosine
__C.TRAIN.WARMUP_RATIO = 0.0

# clip_grad
__C.TRAIN.CLIP_GRAD = None

# beta1, beta2 for Adam Optimizer
__C.TRAIN.beta1 = 0
__C.TRAIN.beta2 = 0.9

# weight decay for Adam or SGD
__C.TRAIN.WEIGHT_DECAY = 0.0

# SGD momentum
__C.TRAIN.MOMENTUM = 0.9

# Check val every n epoch
__C.TRAIN.VAL_EVERY = 5

# Visualization during training
__C.TRAIN.VIS = True
__C.TRAIN.VAL_SAMPLE_VIS = 5

# Loss function.
__C.TRAIN.LOSS = ""

#
# Callback
#
__C.CALLBACK = edict()
__C.CALLBACK.MATCHING_TASK = ["trans"]
__C.CALLBACK.CHECKPOINT_MONITOR = "val/loss"
__C.CALLBACK.CHECKPOINT_MODE = "min"

#
# Loss config
#
__C.LOSS = edict()

#
# Evaluation options
#

__C.EVAL = edict()

#
# MISC
#

# Parallel GPU indices ([0] for single GPU)
__C.GPUS = [0]
# Parallel strategy for multiple gpus
__C.PARALLEL_STRATEGY = "ddp"

# Float Precision, 32 for False, 16 for True
__C.FP16 = False

# CUDNN benchmark
__C.CUDNN = False

__C.WEIGHT_FILE = ""

# Output path (for checkpoints, running logs)
__C.OUTPUT_PATH = ""

# The step of iteration to print running statistics.
# The real step value will be the least common multiple of this value and batch_size
__C.STATISTIC_STEP = 100

# random seed used for data loading
__C.RANDOM_SEED = 42

# directory for collecting statistics of results
__C.STATS = ""

def _merge_a_into_b(a, b):
    """Merge config dictionary A into config dictionary B, clobbering the
    options in B whenever they are also specified in A.
    """
    if type(a) is not edict:
        return

    for k, v in a.items():
        # a must specify keys that are in b
        if k not in b:
            raise KeyError("{} is not a valid config key".format(k))

        # the types must match, too
        if type(b[k]) is not type(v):
            if type(b[k]) is float and type(v) is int:
                v = float(v)
            else:
                if k not in ["CLASS"]:
                    raise ValueError(
                        "Type mismatch ({} vs. {}) for config key: {}".format(
                            type(b[k]), type(v), k
                        )
                    )

        # recursively merge dicts
        if type(v) is edict:
            try:
                _merge_a_into_b(a[k], b[k])
            except:
                print("Error under config key: {}".format(k))
                raise
        else:
            b[k] = v

def cfg_from_file(filename):
    """Load a config file and merge it into the default options."""
    import yaml

    with open(filename, "r") as f:
        yaml_cfg = edict(yaml.full_load(f))

    # DONE: for model specific or data specific configs, please define them in the model package
    #  and then load them here
    if "MODULE" in yaml_cfg and "MODEL" not in __C:
        model_cfg_module = ".".join(
            ["model"] + [yaml_cfg.MODULE.split(".")[0]] + ["model_config"]
        )
        # model_cfg_module = 'model.model_config'
        mod = importlib.import_module(model_cfg_module)
        __C["MODEL"] = mod.get_model_cfg()

    if "DATASET" in yaml_cfg and yaml_cfg.DATASET is not None:
        dataset = importlib.import_module("dataset")
        __C["DATA"] = dataset.dataset_cfg[yaml_cfg.DATASET.split(".")[0].upper()]

    _merge_a_into_b(yaml_cfg, __C)

def cfg_from_list(cfg_list):
    """Set config keys via list (e.g., from command line)."""
    from ast import literal_eval

    assert len(cfg_list) % 2 == 0
    for k, v in zip(cfg_list[0::2], cfg_list[1::2]):
        key_list = k.split(".")
        d = __C
        for sub_key in key_list[:-1]:
            assert sub_key in d.keys()
            d = d[sub_key]
        sub_key = key_list[-1]
        assert sub_key in d.keys()
        try:
            value = literal_eval(v)
        except:
            # handle the case when v is a string literal
            value = v
        assert type(value) == type(
            d[sub_key]
        ), "type {} does not match original type {}".format(
            type(value), type(d[sub_key])
        )
        d[sub_key] = value


""" Wrappers for result objects and iterators """
import six

from .constants import NONE, MAX_GET_BATCH

def add_dicts(d1, d2):
    """ Merge two dicts of addable values """
    if d1 is None:
        return d2
    if d2 is None:
        return d1
    keys = set(d1)
    keys.update(set(d2))
    ret = {}
    for key in keys:
        v1 = d1.get(key)
        v2 = d2.get(key)
        if v1 is None:
            ret[key] = v2
        elif v2 is None:
            ret[key] = v1
        else:
            ret[key] = v1 + v2
    return ret

class Count(int):

    """ Wrapper for response to query with Select=COUNT """

    def __new__(cls, count, scanned_count, consumed_capacity=None):
        ret = super(Count, cls).__new__(cls, count)
        ret.count = count
        ret.scanned_count = scanned_count
        ret.consumed_capacity = consumed_capacity
        return ret

    @classmethod
    def from_response(cls, response):
        """ Factory method """
        return cls(response['Count'], response['ScannedCount'],
                   response.get('consumed_capacity'))

    def __add__(self, other):
        if other is None:
            return self
        if not isinstance(other, Count):
            return self.count + other
        if self.consumed_capacity is None:
            capacity = other.consumed_capacity
        else:
            capacity = self.consumed_capacity + other.consumed_capacity
        return Count(self.count + other.count, self.scanned_count +
                     other.scanned_count, capacity)

    def __radd__(self, other):
        return self.__add__(other)

    def __repr__(self):
        return "Count(%d)" % self

@six.python_2_unicode_compatible
class Capacity(object):
    """ Wrapper for the capacity of a table or index """

    def __init__(self, read, write):
        self._read = read
        self._write = write

    @property
    def read(self):
        """ The read capacity """
        return self._read

    @property
    def write(self):
        """ The write capacity """
        return self._write

    def __getitem__(self, key):
        return getattr(self, key)

    def __contains__(self, key):
        return key in ['read', 'write']

    @classmethod
    def create_read(cls, value):
        """ Create a new read capacity from a response """
        return cls(value['CapacityUnits'], 0)

    @classmethod
    def create_write(cls, value):
        """ Create a new write capacity from a response """
        return cls(0, value['CapacityUnits'])

    def __hash__(self):
        return self._read + self._write

    def __eq__(self, other):
        return (self.read == getattr(other, 'read', None) and
                self.write == getattr(other, 'write', None))

    def __ne__(self, other):
        return not self.__eq__(other)

    def __add__(self, other):
        return Capacity(self.read + other.read, self.write + other.write)

    def __str__(self):
        pieces = []
        if self.read:
            pieces.append('R:{0:.1f}'.format(self.read))
        if self.write:
            pieces.append('W:{0:.1f}'.format(self.write))
        if not pieces:
            return '0'
        return ' '.join(pieces)

@six.python_2_unicode_compatible
class ConsumedCapacity(object):
    """ Record of the consumed capacity of a request """

    def __init__(self, tablename, total, table_capacity=None,
                 local_index_capacity=None, global_index_capacity=None):
        self.tablename = tablename
        self.total = total
        self.table_capacity = table_capacity
        self.local_index_capacity = local_index_capacity
        self.global_index_capacity = global_index_capacity

    @classmethod
    def build_indexes(cls, response, key, capacity_factory):
        """ Construct index capacity map from a request fragment """
        if key not in response:
            return None
        indexes = {}
        for key, val in six.iteritems(response[key]):
            indexes[key] = capacity_factory(val)
        return indexes

    @classmethod
    def from_response(cls, response, is_read):
        """ Factory method for ConsumedCapacity from a response object """
        if is_read:
            cap = Capacity.create_read
        else:
            cap = Capacity.create_write
        kwargs = {
            'tablename': response['TableName'],
            'total': cap(response),
        }
        local = cls.build_indexes(response, 'LocalSecondaryIndexes', cap)
        kwargs['local_index_capacity'] = local
        gindex = cls.build_indexes(response, 'GlobalSecondaryIndexes', cap)
        kwargs['global_index_capacity'] = gindex
        if 'Table' in response:
            kwargs['table_capacity'] = cap(response['Table'])
        return cls(**kwargs)

    def __hash__(self):
        return hash(self.tablename) + hash(self.total)

    def __eq__(self, other):
        properties = ['tablename', 'total', 'table_capacity',
                      'local_index_capacity', 'global_index_capacity']
        for prop in properties:
            if getattr(self, prop) != getattr(other, prop, None):
                return False
        return True

    def __ne__(self, other):
        return not self.__eq__(other)

    def __radd__(self, other):
        return self.__add__(other)

    def __add__(self, other):
        # Handle identity cases when added to empty values
        if other is None:
            return self
        if self.tablename != other.tablename:
            raise TypeError("Cannot add capacities from different tables")
        kwargs = {
            'total': self.total + other.total,
        }
        if self.table_capacity is not None:
            kwargs['table_capacity'] = (self.table_capacity +
                                        other.table_capacity)
        kwargs['local_index_capacity'] = \
            add_dicts(self.local_index_capacity, other.local_index_capacity)
        kwargs['global_index_capacity'] = \
            add_dicts(self.global_index_capacity, other.global_index_capacity)

        return ConsumedCapacity(self.tablename, **kwargs)

    def __str__(self):
        lines = []
        if self.table_capacity:
            lines.append("Table: %s" % self.table_capacity)
        if self.local_index_capacity:
            for name, cap in six.iteritems(self.local_index_capacity):
                lines.append("Local index '%s': %s" % (name, cap))
        if self.global_index_capacity:
            for name, cap in six.iteritems(self.global_index_capacity):
                lines.append("Global index '%s': %s" % (name, cap))
        lines.append("Total: %s" % self.total)
        return '\n'.join(lines)

class PagedIterator(six.Iterator):

    """ An iterator that iterates over paged results from Dynamo """

    def __init__(self):
        self.iterator = None
        self.capacity = 0
        self.table_capacity = 0
        self.indexes = {}
        self.global_indexes = {}

    @property
    def can_fetch_more(self):  # pragma: no cover
        """ Return True if more results can be fetched from the server """
        raise NotImplementedError

    def fetch(self):  # pragma: no cover
        """ Fetch additional results from the server and return an iterator """
        raise NotImplementedError

    def _update_capacity(self, data):
        """ Update the consumed capacity metrics """
        if 'ConsumedCapacity' in data:
            # This is all for backwards compatibility
            consumed = data['ConsumedCapacity']
            if not isinstance(consumed, list):
                consumed = [consumed]
            for cap in consumed:
                self.capacity += cap.get('CapacityUnits', 0)
                self.table_capacity += cap.get('Table',
                                               {}).get('CapacityUnits', 0)
                local_indexes = cap.get('LocalSecondaryIndexes', {})
                for k, v in six.iteritems(local_indexes):
                    self.indexes.setdefault(k, 0)
                    self.indexes[k] += v['CapacityUnits']
                global_indexes = cap.get('GlobalSecondaryIndexes', {})
                for k, v in six.iteritems(global_indexes):
                    self.global_indexes.setdefault(k, 0)
                    self.global_indexes[k] += v['CapacityUnits']

    def __iter__(self):
        return self

    def __next__(self):
        if self.iterator is None:
            self.iterator = self.fetch()
        while True:
            try:
                return six.next(self.iterator)
            except StopIteration:
                if self.can_fetch_more:
                    self.iterator = self.fetch()
                else:
                    raise

class ResultSet(PagedIterator):

    """ Iterator that pages results from Dynamo """

    def __init__(self, connection, limit, *args, **kwargs):
        super(ResultSet, self).__init__()
        self.connection = connection
        # The limit will be mutated, so copy it and leave the original intact
        self.limit = limit.copy()
        self.args = args
        self.kwargs = kwargs
        self.last_evaluated_key = None
        self.consumed_capacity = None

    @property
    def can_fetch_more(self):
        """ True if there are more results on the server """
        return (self.last_evaluated_key is not None and not self.limit.complete)

    def fetch(self):
        """ Fetch more results from Dynamo """
        self.limit.set_request_args(self.kwargs)
        data = self.connection.call(*self.args, **self.kwargs)
        self.limit.post_fetch(data)
        self.last_evaluated_key = data.get('LastEvaluatedKey')
        if self.last_evaluated_key is None:
            self.kwargs.pop('ExclusiveStartKey', None)
        else:
            self.kwargs['ExclusiveStartKey'] = self.last_evaluated_key
        self._update_capacity(data)
        if 'consumed_capacity' in data:
            self.consumed_capacity += data['consumed_capacity']
        for raw_item in data['Items']:
            item = self.connection.dynamizer.decode_keys(raw_item)
            if self.limit.accept(item):
                yield item

class GetResultSet(PagedIterator):

    """ Iterator that pages the results of a BatchGetItem """

    def __init__(self, connection, tablename, keys, consistent=False,
                 attributes=None, alias=None, return_capacity=NONE):
        super(GetResultSet, self).__init__()
        self.connection = connection
        self.tablename = tablename
        self.keys = keys
        self.consistent = consistent
        if attributes is not None:
            if not isinstance(attributes, six.string_types):
                attributes = ', '.join(attributes)
        self.attributes = attributes
        self.alias = alias
        self.return_capacity = return_capacity
        self._attempt = 0
        self.consumed_capacity = None

    @property
    def can_fetch_more(self):
        return bool(self.keys)

    def build_kwargs(self):
        """ Construct the kwargs to pass to batch_get_item """
        keys, self.keys = self.keys[:MAX_GET_BATCH], self.keys[MAX_GET_BATCH:]
        query = {'ConsistentRead': self.consistent}
        if self.attributes is not None:
            query['ProjectionExpression'] = self.attributes
        if self.alias:
            query['ExpressionAttributeNames'] = self.alias
        query['Keys'] = keys
        return {
            'RequestItems': {
                self.tablename: query,
            },
            'ReturnConsumedCapacity': self.return_capacity,
        }

    def fetch(self):
        """ Fetch a set of items from their keys """
        kwargs = self.build_kwargs()
        data = self.connection.call('batch_get_item', **kwargs)
        if 'UnprocessedKeys' in data:
            for items in six.itervalues(data['UnprocessedKeys']):
                self.keys.extend(items['Keys'])
            # Getting UnprocessedKeys indicates that we are exceeding our
            # throughput. So sleep for a bit.
            self._attempt += 1
            self.connection.exponential_sleep(self._attempt)
        else:
            # No UnprocessedKeys means our request rate is fine, so we can
            # reset the attempt number.
            self._attempt = 0
        self._update_capacity(data)
        if 'consumed_capacity' in data:
            # Comes back as a list from BatchWriteItem
            self.consumed_capacity = \
                sum(data['consumed_capacity'], self.consumed_capacity)
        return iter(data['Responses'][self.tablename])

    def __next__(self):
        result = super(GetResultSet, self).__next__()
        return self.connection.dynamizer.decode_keys(result)

class TableResultSet(PagedIterator):

    """ Iterator that pages table names from ListTables """

    def __init__(self, connection, limit=None):
        super(TableResultSet, self).__init__()
        self.connection = connection
        self.limit = limit
        self.last_evaluated_table_name = None

    @property
    def can_fetch_more(self):
        if self.last_evaluated_table_name is None:
            return False
        return self.limit is None or self.limit > 0

    def fetch(self):
        kwargs = {}
        if self.limit is None:
            kwargs['Limit'] = 100
        else:
            kwargs['Limit'] = min(self.limit, 100)
        if self.last_evaluated_table_name is not None:
            kwargs['ExclusiveStartTableName'] = self.last_evaluated_table_name
        data = self.connection.call('list_tables', **kwargs)
        self.last_evaluated_table_name = data.get('LastEvaluatedTableName')
        tables = data['TableNames']
        if self.limit is not None:
            self.limit -= len(tables)
        return iter(tables)

class Result(dict):

    """
    A wrapper for an item returned from Dynamo

    Attributes
    ----------
    capacity : float
        Total consumed capacity
    table_capacity : float
        Consumed capacity on the table
    indexes : dict
        Mapping of local index name to the consumed capacity on that index
    global_indexes : dict
        Mapping of global index name to the consumed capacity on that index
    exists : bool
        False if the result is empty (i.e. no result was returned from dynamo)

    """

    def __init__(self, dynamizer, response, item_key):
        super(Result, self).__init__()
        self.exists = item_key in response
        for k, v in six.iteritems(response.get(item_key, {})):
            self[k] = dynamizer.decode(v)

        self.consumed_capacity = response.get('consumed_capacity')

        # All this shit is for backwards compatibility
        cap = response.get('ConsumedCapacity', {})
        self.capacity = cap.get('CapacityUnits', 0)
        self.table_capacity = cap.get('Table', {}).get('CapacityUnits', 0)
        self.indexes = dict((
            (k, v['CapacityUnits']) for k, v in
            six.iteritems(cap.get('LocalSecondaryIndexes', {}))
        ))
        self.global_indexes = dict((
            (k, v['CapacityUnits']) for k, v in
            six.iteritems(cap.get('GlobalSecondaryIndexes', {}))
        ))

    def __repr__(self):
        return 'Result({0})'.format(super(Result, self).__repr__())

class Limit(object):

    """
    Class that defines query/scan limit behavior

    Parameters
    ----------
    scan_limit : int, optional
        The maximum number of items for DynamoDB to scan. This will not
        necessarily be the number of items returned.
    item_limit : int, optional
        The maximum number of items to return. Fetches will continue until this
        number is reached or there are no results left. See also: ``strict``
    min_scan_limit : int, optional
        This only matters when ``item_limit`` is set and ``scan_limit`` is not.
        After doing multiple fetches, the ``item_limit`` may drop to a low
        value. The ``item_limit`` will be passed up as the query ``Limit``, but
        if your ``item_limit`` is down to 1 you may want to fetch more than 1
        item at a time. ``min_scan_limit`` determines the minimum ``Limit`` to
        send up when ``scan_limit`` is None. (default 20)
    strict : bool, optional
        This modifies the behavior of ``item_limit``. If True, the query will
        never return more items than ``item_limit``. If False, the query will
        fetch until it hits the ``item_limit``, and then return the rest of the
        page as well. (default False)
    filter : callable, optional
        Function that takes a single item dict and returns a boolean. If True,
        the item will be counted towards the ``item_limit`` and returned from
        the iterator. If False, it will be skipped.

    """

    def __init__(self, scan_limit=None, item_limit=None, min_scan_limit=20,
                 strict=False, filter=lambda x: True):
        self.scan_limit = scan_limit
        if item_limit is None:
            self.item_limit = scan_limit
        else:
            self.item_limit = item_limit
        self.min_scan_limit = min_scan_limit
        self.strict = strict
        self.filter = filter

    def copy(self):
        """ Return a copy of the limit """
        return Limit(self.scan_limit, self.item_limit, self.min_scan_limit,
                     self.strict, self.filter)

    def set_request_args(self, args):
        """ Set the Limit parameter into the request args """
        if self.scan_limit is not None:
            args['Limit'] = self.scan_limit
        elif self.item_limit is not None:
            args['Limit'] = max(self.item_limit, self.min_scan_limit)
        else:
            args.pop('Limit', None)

    @property
    def complete(self):
        """ Return True if the limit has been reached """
        if self.scan_limit is not None and self.scan_limit == 0:
            return True
        if self.item_limit is not None and self.item_limit == 0:
            return True
        return False

    def post_fetch(self, response):
        """ Called after a fetch. Updates the ScannedCount """
        if self.scan_limit is not None:
            self.scan_limit -= response['ScannedCount']

    def accept(self, item):
        """ Apply the filter and item_limit, and return True to accept """
        accept = self.filter(item)
        if accept and self.item_limit is not None:
            if self.item_limit > 0:
                self.item_limit -= 1
            elif self.strict:
                return False
        return accept


import setuptools

with open("README.md", "r") as fh:
    long_description = fh.read()

setuptools.setup(
    name="eigensheep",
    version="1.0.0",
    author="Kevin Kwok",
    author_email="antimatter15@gmail.com",
    description="Run Jupyter cells in AWS Lambda for massively parallel experimentation",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/antimatter15/eigensheep",
    packages=setuptools.find_packages(),
    classifiers=[
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 2",
        "License :: OSI Approved :: MIT License",
        "Operating System :: OS Independent",
        "Development Status :: 3 - Alpha",
        "Framework :: Jupyter",
        "Topic :: Scientific/Engineering",
        "Intended Audience :: Science/Research",
        "Intended Audience :: Developers"
    ],
    install_requires=[
        'tqdm',
        'boto3',
        'ipython',
        'ipywidgets'
    ],
    extras_require={
        ':python_version == "2.7"': ['futures']
    }
)


import pickle
import argparse
import os
from omni.isaac.kit import SimulationApp

# add argparse arguments
parser = argparse.ArgumentParser("Welcome to Orbit: Omniverse Robotics Environments!")
parser.add_argument("--headless", action="store_true", default=False, help="Force display off at all times.")
parser.add_argument("--num_envs", type=int, default=3, help="Number of environments to spawn.")
args_cli = parser.parse_args()

# launch omniverse app
config = {"headless": args_cli.headless}
simulation_app = SimulationApp(config)

"""Rest everything follows."""
from PIL import Image as im
# import os
import torch
import cv2
import numpy as np
import open3d as o3d
import random
from scipy import signal
from omni.isaac.core.prims import RigidPrim,GeometryPrim
from omni.isaac.orbit.utils.array import convert_to_torch
import omni.isaac.core.utils.prims as prim_utils
from omni.isaac.core.simulation_context import SimulationContext
from omni.isaac.core.utils.carb import set_carb_setting
from omni.isaac.core.utils.viewports import set_camera_view
from omni.isaac.orbit.utils.assets import ISAAC_NUCLEUS_DIR
import omni.isaac.orbit.utils.kit as kit_utils
from omni.isaac.orbit.robots.config.franka import FRANKA_PANDA_ARM_WITH_PANDA_HAND_CFG
from omni.isaac.orbit.robots.single_arm import SingleArmManipulator
from omni.isaac.orbit.sensors.camera import Camera, PinholeCameraCfg
from omni.isaac.orbit.sensors.camera.utils import create_pointcloud_from_rgbd
from omni.isaac.core.objects import DynamicCuboid, FixedCuboid
from omni.isaac.orbit.objects import RigidObjectCfg,RigidObject
from PIL import Image
import matplotlib.pyplot as plt
from omni.isaac.orbit.controllers.differential_inverse_kinematics import (
    DifferentialInverseKinematics,
    DifferentialInverseKinematicsCfg,
)
import asyncio
from omni.isaac.orbit.objects.rigid import RigidObject, RigidObjectCfg
from omni.isaac.orbit.utils.math import convert_quat
import scipy.spatial.transform as tf
from shapely import Polygon, STRtree, area, contains
from omni.isaac.orbit_envs.isaac_env_cfg import EnvCfg, IsaacEnvCfg, PhysxCfg, SimCfg, ViewerCfg

import carb
#################
from place_new_obj import place_new_obj_fun
"""
Main
"""
class sim_cfg:
    sim: SimCfg = SimCfg(
            dt=0.01,
            substeps=4,
            physx=PhysxCfg(
                # num_position_iterations=8,
                gpu_max_rigid_contact_count=1024**2, #1024**2*2,
                gpu_max_rigid_patch_count=160*2048*10, #160*2048*10, #160*2048*10,
                gpu_found_lost_pairs_capacity = 1024 * 1024 * 2 * 2,#1024 * 1024 * 2 * 1, #1024 * 1024 * 2 * 8,
                gpu_found_lost_aggregate_pairs_capacity=1024 * 1024 * 32 * 1,#1024 * 1024 * 32 * 1, #1024 * 1024 * 32,
                gpu_total_aggregate_pairs_capacity=1024 * 1024 * 2 *1, #1024 * 1024 * 2 *1, #1024 * 1024 * 2 * 8
                friction_correlation_distance=0.0025,
                friction_offset_threshold=0.04,
                bounce_threshold_velocity=0.5,
                gpu_max_num_partitions=8,

            ),

        )

def point_cloud_process(pcd):

    plane_model, inliers = pcd.segment_plane(distance_threshold=0.01,
                                         ransac_n=3,
                                         num_iterations=1000)
    outlier_cloud = pcd.select_by_index(inliers, invert=True)
    plane_model, inliers = outlier_cloud.segment_plane(distance_threshold=0.01,
                                         ransac_n=3,
                                         num_iterations=1000)
    print(plane_model)
    return plane_model
def _configure_simulation_flags( sim_params: dict = None):
        """Configure simulation flags and extensions at load and run time."""
        # acquire settings interface
        carb_settings_iface = carb.settings.get_settings()
        # enable hydra scene-graph instancing
        # note: this allows rendering of instanceable assets on the GUI
        carb_settings_iface.set_bool("/persistent/omnihydra/useSceneGraphInstancing", True)
        # change dispatcher to use the default dispatcher in PhysX SDK instead of carb tasking
        # note: dispatcher handles how threads are launched for multi-threaded physics
        carb_settings_iface.set_bool("/physics/physxDispatcher", True)
        # disable contact processing in omni.physx if requested
        # note: helpful when creating contact reporting over limited number of objects in the scene
        if sim_params["disable_contact_processing"]:
            carb_settings_iface.set_bool("/physics/disableContactProcessing", True)

def main():
    """Spawns lights in the stage and sets the camera view."""
    sim_params = sim_cfg.sim.to_dict()
    if sim_params is not None:
        if "physx" in sim_params:
            physx_params = sim_params.pop("physx")
            sim_params.update(physx_params)
    # set flags for simulator
    _configure_simulation_flags(sim_params)
    # create a simulation context to control the simulator
    sim = SimulationContext(
        stage_units_in_meters=1.0,
        physics_dt=0.01,
        rendering_dt=0.1,
        backend="torch",
        sim_params=sim_params,
        # physics_prim_path="/physicsScene",
        device='cuda:0',
    )
    # Load kit helper
    # sim = SimulationContext(physics_dt=0.01, rendering_dt=1, backend="torch",device='cuda:0',)
    # Set main camera
    set_camera_view([0, 2, 3.], [0.0, 0.0, 0])
    # Enable GPU pipeline and flatcache
    if sim.get_physics_context().use_gpu_pipeline:
        sim.get_physics_context().enable_flatcache(True)
    # PhysicsContext.set_gpu_total_aggregate_pairs_capacity(4000)
    # Enable hydra scene-graph instancing
    set_carb_setting(sim._settings, "/persistent/omnihydra/useSceneGraphInstancing", True)

    ####################################################
    ####################################################
    ################# SCENE ############################
    # Ground-plane
    kit_utils.create_ground_plane("/World/defaultGroundPlane",z_position=-0.5,
        static_friction=0.5,
        dynamic_friction=0.5,
        restitution=0.8,
        improve_patch_friction=True,)
    # Lights-1
    prim_utils.create_prim(
        "/World/Light/GreySphere",
        "SphereLight",
        translation=(4.5, 3.5, 10.0),
        attributes={"radius": 2.5, "intensity": 600.0, "color": (0.75, 0.75, 0.75)},
    )
    # Lights-2
    prim_utils.create_prim(
        "/World/Light/WhiteSphere",
        "SphereLight",
        translation=(-4.5, 3.5, 10.0),
        attributes={"radius": 2.5, "intensity": 600.0, "color": (1.0, 1.0, 1.0)},
    )
    #################### create table
    table_path = f"{ISAAC_NUCLEUS_DIR}/Props/Shapes/cube.usd"
    Table = FixedCuboid(prim_path="/World/Table",position=(0,0,-0.25),scale=(0.5,0.5,0.5))
    # Table.set_mass(10000000)
    sideTable = FixedCuboid(prim_path="/World/sideTable",position=(0.35,-0.9,-0.3),scale=(0.4,0.4,0.4))
    # sideTable.set_mass(10)
    #################### robot base
    prim_utils.create_prim("/World/Robotbase", usd_path=table_path,position=(0,-0.45,-0.2),scale=(0.3,0.26,0.4))
    #################### ycb path
    ######################################### load ycb objects
    # ycb_usd_paths = {
    #     "crackerBox": f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/003_cracker_box.usd",
    #     "sugarBox": f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/004_sugar_box.usd",
    #     "tomatoSoupCan": f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/005_tomato_soup_can.usd",
    #     "mustardBottle": f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/006_mustard_bottle.usd",
    #     "mug":f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned/025_mug.usd",
    #     "largeMarker":f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned/040_large_marker.usd",
    #     "tunaFishCan":f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned/007_tuna_fish_can.usd",
    #     "banana":f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned/011_banana.usd",
    #     # "pitcherBase":f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned/019_pitcher_base.usd",
    #     "bowl":f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned/024_bowl.usd",
    #     "largeClamp":f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned/051_large_clamp.usd",
    #     "scissors":f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned/037_scissors.usd",
    # }
    # ycb_name = ['crackerBox','sugarBox','tomatoSoupCan','mustardBottle','mug','largeMarker','tunaFishCan',
    #             'banana','bowl','largeClamp','scissors']
    # ycb_usd_paths = {
    #     "crackerBox": f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/003_cracker_box.usd",
    #     "sugarBox": f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/004_sugar_box.usd",
    #     # "tomatoSoupCan": f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/005_tomato_soup_can.usd",
    #     "mustardBottle": f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/006_mustard_bottle.usd",
    # }
    # ycb_name = ['crackerBox','sugarBox','mustardBottle']
    ycb_usd_paths = {
        # "crackerBox": f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/003_cracker_box.usd",
        "sugarBox": f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/004_sugar_box.usd",
        # "tomatoSoupCan": f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/005_tomato_soup_can.usd",
        # "mustardBottle": f"{ISAAC_NUCLEUS_DIR}/Props/YCB/Axis_Aligned_Physics/006_mustard_bottle.usd",
    }
    ycb_name = ['sugarBox']
    obj1 = []
    obj_name_list = []
    for i in range(9):
        obj_cfg1 = RigidObjectCfg()
        obj_cfg1.meta_info = RigidObjectCfg.MetaInfoCfg(usd_path=ycb_usd_paths[ycb_name[0]],scale=(1.0, 1.0, 1.0),)
        obj_cfg1.init_state = RigidObjectCfg.InitialStateCfg(
        pos=(2-0.25*i, 0.8, -0.4), rot=(1.0, 0.0, 0.0, 0.0), lin_vel=(0.0, 0.0, 0.0), ang_vel=(0.0, 0.0, 0.0)
        )
        obj_cfg1.rigid_props = RigidObjectCfg.RigidBodyPropertiesCfg(
            solver_position_iteration_count=16,
            solver_velocity_iteration_count=1,
            max_angular_velocity=0.5,
            max_linear_velocity=0.5,
            max_depenetration_velocity=0.5,
            disable_gravity=False,
        )
        obj_cfg1.physics_material = RigidObjectCfg.PhysicsMaterialCfg(
            static_friction=0.5, dynamic_friction=0.5, restitution=0.0, prim_path="/World/Materials/cubeMaterial"
        )
        obj1.append(RigidObject(obj_cfg1))
        obj_name_list.append(ycb_name[0]+str(i))
    ################################ robot setting
    robot_cfg = FRANKA_PANDA_ARM_WITH_PANDA_HAND_CFG
    robot_cfg.data_info.enable_jacobian = True
    robot_cfg.rigid_props.disable_gravity = True
    robot = SingleArmManipulator(cfg=robot_cfg)
    robot.spawn("/World/Robot", translation=(0.0, -.45, 0))

    ###################################### controller
    ik_control_cfg = DifferentialInverseKinematicsCfg(
        command_type="pose_abs",
        ik_method="dls",
        position_offset=robot.cfg.ee_info.pos_offset,
        rotation_offset=robot.cfg.ee_info.rot_offset,
    )
    ik_controller = DifferentialInverseKinematics(cfg=ik_control_cfg, num_robots=1, device=sim.device)
    ###################################### spawn items
    for i,obj_t in enumerate(obj1):
        obj_t.spawn(f"/World/Objs/obj1/obj_{i}")
    ###################################### sensor extension camera

    camera_cfg = PinholeCameraCfg(
        sensor_tick=0,
        height=240,
        width=240,
        data_types=["rgb", "distance_to_image_plane", "normals", "motion_vectors"],
        usd_params=PinholeCameraCfg.UsdCameraCfg(
            focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 1.0e5)
        ),
    )
    camera = Camera(cfg=camera_cfg, device="cuda")
    # hand_camera = Camera(cfg=camera_cfg,device='cuda')
    # hand_camera.spawn("/World/Robot/panda_hand/hand_camera", translation=(0.1, 0.0, 0.0),orientation=(0,0,1,0))
    # hand_camera.spawn("/World/hand_camera")
    # Spawn camera
    camera.spawn("/World/CameraSensor")
    sim.reset()
    ##########################################
    ##########################################
    for _ in range(14):
        sim.render()
    # hand_camera.initialize()
    camera.initialize()
    robot.initialize()
    ik_controller.initialize()
    for i,obj_t in enumerate(obj1):
            obj_t.initialize(f"/World/Objs/obj1/obj_{i}")
    # Reset states
    robot.reset_buffers()
    ik_controller.reset_idx()
    position = [0, 0, 1.2]
    orientation = [0, 0, -1, 0]
    camera.set_world_pose_ros(position, orientation)
    # hand_camera.set_world_pose_ros([0.35,-0.9,0.55], orientation)
    Table.initialize()
    sideTable.initialize()
    sideTable.set_collision_enabled(True)
    Table.set_collision_enabled(True)
    for i,obj_t in enumerate(obj1):
            obj_t.update_buffers(0.01)
    for _ in range(10):
        sim.render()
    ##################################################################### get plane model of the table
    camera.update(dt=0.0)
    pcd = get_pcd(camera)
    plane_model = point_cloud_process(pcd)
    plane_model_ori = plane_model
    plane_model = np.array([plane_model[0],plane_model[1],plane_model[2]])
    ##################################################################### load ycb
    obj_dict = dict()
    table_obj_pos_rot = dict()
    for i in range(1000):
        print('i')
        print(i)
        ## initialize parameters
        for _ in ycb_name:
            obj_dict[_] = 0
        table_obj_pos_rot = dict()
        ################## random drop one item
        randi = np.random.randint(0,len(ycb_name))
        obj_name_i = ycb_name[randi]
        angle = np.random.randint(0,180)
        translation = torch.rand(2).tolist()
        translation = [0.3*translation[0]-0.15,0.3*translation[1]-0.15,0.07]
        rot = convert_quat(tf.Rotation.from_euler("XYZ", (0,0,angle), degrees=True).as_quat(), to="wxyz")
        root_state = obj1[0].get_default_root_state()
        root_state[:,:3] = torch.tensor(translation).cuda()
        root_state[:,3:7] = torch.tensor(rot).cuda()
        obj1[0].set_root_state(root_state)
        obj1[0].update_buffers(0.01)
        table_obj_pos_rot[obj_name_i] = [(translation,rot)]
        for j in range(10):
            sim.step()
        obj_dict[obj_name_i] += 1
        #################### start the loop to drop item iteratively
        for j in range(1,9):
            print(j)
            if simulation_app.is_running():
                print('running')
            if sim.step(render=not args_cli.headless):
                print('pause')
                sim.play()
                continue
            camera.update(dt=0.01)
            for _ in range(10):
                sim.render()
            camera.update(dt=0.01)
            pcd = get_pcd(camera)
            pointcloud_w = np.array(pcd.points)
            select_m = np.dot(pointcloud_w,plane_model) + float(plane_model_ori[3])
            index_inliers = np.argwhere((select_m >=-0.01)).reshape(-1).astype(int)
            inliers = pointcloud_w[index_inliers]
            select_m = np.dot(inliers,plane_model) + float(plane_model_ori[3])
            index_inliers = np.argwhere((select_m <=0.3)).reshape(-1).astype(int)
            inliers = inliers[index_inliers]
            index_inliers = np.argwhere((inliers[:,1]>=-0.25)).reshape(-1).astype(int)
            inliers = inliers[index_inliers]
            index_inliers = np.argwhere((inliers[:,1]<=0.25)).reshape(-1).astype(int)
            inliers = inliers[index_inliers]
            index_inliers = np.argwhere((inliers[:,0]>=-0.25)).reshape(-1).astype(int)
            inliers = inliers[index_inliers]
            index_inliers = np.argwhere((inliers[:,0]<=0.25)).reshape(-1).astype(int)
            inliers = inliers[index_inliers]
            select_m = np.dot(inliers,plane_model) + float(plane_model_ori[3])
            index_objects = np.argwhere((select_m>=0.001)).reshape(-1).astype(int)
            objects_point = inliers[index_objects].copy()
            objects_pcd = o3d.geometry.PointCloud()
            objects_pcd.points = o3d.utility.Vector3dVector(objects_point)
            # o3d.visualization.draw_geometries([objects_pcd])
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(inliers)
            # o3d.visualization.draw_geometries([pcd])
            pts_tab = np.array(pcd.points)
            Nx,Ny = 50,50
            x = np.linspace(np.min(pts_tab[:,0]), np.max(pts_tab[:,0]), Nx)
            y = np.linspace(np.min(pts_tab[:,1]), np.max(pts_tab[:,1]), Ny)
            xv, yv = np.meshgrid(x, y)
            pts = np.array(objects_pcd.points)
            u = (pts[:,0] - np.min(pts_tab[:,0]))/ ( np.max(pts_tab[:,0])-np.min(pts_tab[:,0]) )
            v = (pts[:,1] - np.min(pts_tab[:,1]))/ ( np.max(pts_tab[:,1])-np.min(pts_tab[:,1]) )
            u = (Nx-1)*u
            v = (Ny-1)*v
            occupancy = np.zeros( (Ny,Nx) )
            u = u.astype(int)
            v = v.astype(int)
            u_ind = np.where(u<Nx)
            u = u[u_ind]
            v = v[u_ind]
            v_ind = np.where(v<Ny)
            u = u[v_ind]
            v = v[v_ind]
            u_ind = np.where(u>=0)
            u = u[u_ind]
            v = v[u_ind]
            v_ind = np.where(v>=0)
            u = u[v_ind]
            v = v[v_ind]
            occupancy[v,u] = 1
            occupancy = np.fliplr(occupancy)
            randi = np.random.randint(0,len(ycb_name))
            obj_name_i = ycb_name[randi]
            vertices_new_obj = get_new_obj_info(obj_name_i)
            flag_found, new_poly_vetices,occu_tmp,new_obj_pos = place_new_obj_fun(occupancy,vertices_new_obj)
            if flag_found:
                obj_dict[obj_name_i] += 1
                if obj_name_i in ["mug","tomatoSoupCan","pitcherBase","tunaFishCan","bowl","banana"]:
                    rot = convert_quat(tf.Rotation.from_euler("XYZ", (-90,np.rad2deg(new_obj_pos[2]),0), degrees=True).as_quat(), to="wxyz")
                else:
                    rot = convert_quat(tf.Rotation.from_euler("XYZ", (0,0,-np.rad2deg(new_obj_pos[2])), degrees=True).as_quat(), to="wxyz")
                print(new_obj_pos)
                translation = [(Nx/2-new_obj_pos[1])*1./100.,(new_obj_pos[0]-Ny/2)*1./100.,0.05]
                print(translation)
                root_state = obj1[j].get_default_root_state()
                root_state[:,:3] = torch.tensor(translation).cuda()
                root_state[:,3:7] = torch.tensor(rot).cuda()
                obj1[j].set_root_state(root_state)
                for k in range(250):
                    sim.step
                obj1[j].update_buffers(0.01)
                print('current state')
                print(obj1[j].data.root_pos_w)
                if obj_name_i not in table_obj_pos_rot:
                    table_obj_pos_rot[obj_name_i] = [(translation,rot)]
                else:
                    table_obj_pos_rot[obj_name_i].append((translation,rot))
                for k in range(250):
                    sim.step
                for k in range(20):
                    sim.render()
                print('current state')
                print(obj1[j].data.root_pos_w)
            else:
                file_name_ori = "dict_"
                file_list = os.listdir("generated_table/")
                print('done')
                print(table_obj_pos_rot)
                for k in table_obj_pos_rot:
                    print(len(table_obj_pos_rot[k]))
                num_file = 1
                while True:
                    file_name = file_name_ori+str(num_file)+".pkl"
                    if file_name in file_list:
                        num_file +=1
                    else:
                        file_path = "generated_table/"+file_name
                        f_save = open(file_path,'wb')
                        table_obj_pos_rot = [table_obj_pos_rot,obj_name_i]
                        pickle.dump(table_obj_pos_rot,f_save)
                        f_save.close()
                        break
                # f_save = open('generated_table/dict_file.pkl','wb')
                # pickle.dump(table_obj_pos_rot,f_save)
                # f_save.close()
                break
        #################### reset
        for i,obj_t in enumerate(obj1):
            root_state = obj_t.get_default_root_state()
            obj_t.set_root_state(root_state)
        for j in range(200):
            sim.step()

def get_pcd(camera):
    camera.update(dt=0.0)
    pointcloud_w, pointcloud_rgb = create_pointcloud_from_rgbd(
            camera.data.intrinsic_matrix,
            depth=camera.data.output["distance_to_image_plane"],
            rgb=camera.data.output["rgb"],
            position=camera.data.position,
            orientation=camera.data.orientation,
            normalize_rgb=True,
            num_channels=4,
    )
    if not isinstance(pointcloud_w, np.ndarray):
        pointcloud_w = pointcloud_w.cpu().numpy()
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(pointcloud_w)
    rgb=camera.data.output["rgb"]
    rgb = convert_to_torch(rgb, device='cuda:0', dtype=torch.float32)
    rgb = rgb[:, :, :3].cpu().data.numpy()

    img = Image.fromarray((rgb).astype(np.uint8))
    plt.imshow(img)
    plt.show()
    # o3d.visualization.draw_geometries([pcd])
    return pcd
def get_new_obj_info(obj_type):
# def get_new_obj_info(camera,size,hand_plane_model,obj_type):
    '''
    pcd = get_pcd(camera)
    # o3d.visualization.draw_geometries([pcd])
    plane_model, inliers = pcd.segment_plane(distance_threshold=0.01,
                                         ransac_n=3,
                                         num_iterations=1000)
    outlier_cloud = pcd.select_by_index(inliers, invert=True)
    plane_model, inliers = outlier_cloud.segment_plane(distance_threshold=0.005,
                                         ransac_n=3,
                                         num_iterations=1000)
    inlier_cloud = outlier_cloud.select_by_index(inliers)
    plane_model_ori = plane_model
    plane_model = np.array([plane_model[0],plane_model[1],plane_model[2]])
    pointcloud_w = np.array(outlier_cloud.points)

    # o3d.visualization.draw_geometries([inlier_cloud])
    # o3d.visualization.draw_geometries([outlier_cloud])

    select_m = np.dot(pointcloud_w,plane_model) + float(plane_model_ori[3])
    index_objects = np.argwhere((select_m>=0.01)).reshape(-1).astype(int)
    objects_point = pointcloud_w[index_objects].copy()
    # objects_pcd = outlier_cloud.select_by_index(inliers,invert = True)
    objects_pcd = o3d.geometry.PointCloud()
    objects_pcd.points = o3d.utility.Vector3dVector(objects_point)
    # o3d.visualization.draw_geometries([objects_pcd])
    pcd = inlier_cloud
    # o3d.visualization.draw_geometries([pcd])
    pts_tab = np.array(pcd.points)
    # pts_tab[:,2] = 0
    # pcd.points = o3d.utility.Vector3dVector(pts_tab)
    aabb = objects_pcd.get_oriented_bounding_box()
    print(np.array(aabb.get_box_points()))
    aabb_points = np.array(aabb.get_box_points()).reshape((-1,3))
    aabb.color = (1, 0, 0)
    # o3d.visualization.draw_geometries([objects_pcd, aabb])
    Nx = size[1]
    Ny = size[0]
    x = np.linspace(np.min(pts_tab[:,0]), np.max(pts_tab[:,0]), Nx)
    y = np.linspace(np.min(pts_tab[:,1]), np.max(pts_tab[:,1]), Ny)
    xv, yv = np.meshgrid(x, y)
    pts = np.array(objects_pcd.points)
    u = (pts[:,0] - np.min(pts_tab[:,0]))/ ( np.max(pts_tab[:,0])-np.min(pts_tab[:,0]) )
    v = (pts[:,1] - np.min(pts_tab[:,1]))/ ( np.max(pts_tab[:,1])-np.min(pts_tab[:,1]) )
    u = (size[0]-1)*u
    v = (size[1]-1)*v
    occupancy = np.zeros( (Ny,Nx) )
    u = u.astype(int)
    v = v.astype(int)
    u_ind = np.where(u<size[0])
    u = u[u_ind]
    v = v[u_ind]
    v_ind = np.where(v<size[1])
    u = u[v_ind]
    v = v[v_ind]
    u_ind = np.where(u>=0)
    u = u[u_ind]
    v = v[u_ind]
    v_ind = np.where(v>=0)
    u = u[v_ind]
    v = v[v_ind]
    occupancy[v,u] = 1
    occupancy = np.fliplr(occupancy)
    file_list = os.listdir("obj_mask/")
    file_name = obj_type +"_mask.pkl"
    if file_name not in file_list:
        file_path = "obj_mask/"+file_name
        f_save = open(file_path,'wb')
        pickle.dump(occupancy,f_save)
        f_save.close()
    '''
    occupancy = np.zeros((40,40))
    file_list = os.listdir("obj_mask/")
    for i in range(len(file_list)):
        if obj_type in file_list[i]:
            fileObject2 = open('obj_mask/'+file_list[i], 'rb')
            occupancy=  pickle.load(fileObject2)

            fileObject2.close()
    # plt.imshow(occupancy)
    # plt.show()
    vertices_new_obj = get_new_obj_contour_bbox(occupancy)
    return vertices_new_obj
def get_new_obj_contour_bbox(occu:np.array):
    mask = occu.copy()
    shape_occu = occu.shape
    Nx = shape_occu[1]
    Ny = shape_occu[0]
    mask = np.array((mask-np.min(mask))*255/(np.max(mask)-np.min(mask)),dtype=np.uint8)
    ret,mask = cv2.threshold(mask,50,255,0)
    contours,hierarchy = cv2.findContours(mask,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
    max_area = 0
    cnt = []
    for i in contours:
        area_tmp = cv2.contourArea(i)
        if area_tmp>max_area:
            max_area = area_tmp
            cnt = i
    # approx = cv2.minAreaRect(cnt)
    x,y,w,h = cv2.boundingRect(cnt)
    # print(approx)
    # box = cv2.boxPoints(approx)
    # approx = np.int0(box)
    if x+w >=occu.shape[1]:
        w = occu.shape[1]-x-1
    if y+h >=occu.shape[0]:
        h = occu.shape[0]-1-y
    approx = np.array([[x,y],[x+w,y],[x+w,y+h],[x,y+h]])
    vertices_new_obj = []
    mask_tmp = mask.copy()
    if len(approx) >=2:
        approx = approx.reshape((-1,2))
        for i in range(len(approx)):
            mask_tmp[approx[i][1],approx[i][0]] = 130
        vertices_new_obj = approx
        # print(vertices_new_obj)
        vertices_new_obj = vertices_new_obj - np.array([Nx/2,Ny/2])
        # print(vertices_new_obj)
        # plt.imshow(mask_tmp)
        # plt.show()

        l = []
        for i in range(2):
            l.append(np.linalg.norm(vertices_new_obj[i]-vertices_new_obj[i+1]))
        # print(l)
        return vertices_new_obj
    else:
        return None
def get_new_obj_pcd(camera,size,hand_plane_model):
    pcd = get_pcd(camera)
    # o3d.visualization.draw_geometries([pcd])
    plane_model, inliers = pcd.segment_plane(distance_threshold=0.01,
                                         ransac_n=3,
                                         num_iterations=1000)
    outlier_cloud = pcd.select_by_index(inliers, invert=True)
    plane_model, inliers = outlier_cloud.segment_plane(distance_threshold=0.005,
                                         ransac_n=3,
                                         num_iterations=1000)
    inlier_cloud = outlier_cloud.select_by_index(inliers)
    plane_model_ori = plane_model
    plane_model = np.array([plane_model[0],plane_model[1],plane_model[2]])
    pointcloud_w = np.array(outlier_cloud.points)

    # o3d.visualization.draw_geometries([inlier_cloud])
    # o3d.visualization.draw_geometries([outlier_cloud])

    select_m = np.dot(pointcloud_w,plane_model) + float(plane_model_ori[3])
    index_objects = np.argwhere((select_m>=0.01)).reshape(-1).astype(int)
    objects_point = pointcloud_w[index_objects].copy()
    # objects_pcd = outlier_cloud.select_by_index(inliers,invert = True)
    objects_pcd = o3d.geometry.PointCloud()
    objects_pcd.points = o3d.utility.Vector3dVector(objects_point)
    # o3d.visualization.draw_geometries([objects_pcd])
    pcd = inlier_cloud
    o3d.visualization.draw_geometries([pcd])
    pts_tab = np.array(pcd.points)
    # pts_tab[:,2] = 0
    # pcd.points = o3d.utility.Vector3dVector(pts_tab)
    aabb = objects_pcd.get_oriented_bounding_box()
    # print(np.array(aabb.get_box_points()))
    aabb_points = np.array(aabb.get_box_points()).reshape((-1,3))
    aabb.color = (1, 0, 0)
    o3d.visualization.draw_geometries([objects_pcd, aabb])
    return aabb_points
def place_new_object(occu,ycb_list,ycb_path,num_new,obj_dict):
    angle = 0
    while True:
        randi = np.random.randint(0,len(ycb_list))
        # angle = np.random.randint(0,180)
        key_ori = ycb_list[randi]
        usd_path = ycb_path[key_ori]
        if key_ori not in obj_dict:
            obj_dict[key_ori] = 1
        else:
            if obj_dict[key_ori]<9:
                obj_dict[key_ori] +=1
                break

    print(obj_dict)
    key = key_ori+str(obj_dict[key_ori])
    translation = 0
    translation = [0.35,-0.9,0.1]
    rot = convert_quat(tf.Rotation.from_euler("XYZ", (0,0,angle), degrees=True).as_quat(), to="wxyz")
    if key_ori in ["mug","tomatoSoupCan","pitcherBase","tunaFishCan","bowl","banana"]:
        rot = convert_quat(tf.Rotation.from_euler("XYZ", (-90,angle,0), degrees=True).as_quat(), to="wxyz")
    prim_utils.create_prim(f"/World/newObjects/{key}", usd_path=usd_path, translation=translation,orientation=rot)
    new_obj = GeometryPrim(f"/World/newObjects/{key}",collision=True)
    RigidPrim(f"/World/newObjects/{key}",mass=0.3)
    new_obj_path = f"/World/newObjects/{key}"
    return obj_dict,new_obj,key_ori,new_obj_path
    # prim_utils.delete_prim(f"/World/newObjects/{key+str(num_new)}")

def bound_detect(occu):
    occu_p = occu.copy()/np.max(occu)
    occu_p = np.array(occu*255,dtype=np.uint8)
    # data = im.fromarray(occu_p)
    # data.save('pic/occu.png')
    # img = cv2.imread('pic/occu.png')
    # gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    ret,thresh = cv2.threshold(occu_p,50,255,0)
    contours,hierarchy = cv2.findContours(thresh,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
    print("Number of contours detected:",len(contours))
    shape_dict = dict()
    i = 0
    polygons = []
    for cnt in contours:
        i += 1
        # approx = cv2.approxPolyDP(cnt,0.1*cv2.arcLength(cnt,True),True)
        # approx = cv2.convexHull(cnt)
        approx = cv2.minAreaRect(cnt)
        (x,y) = cnt[0,0]
        # print(approx)
        box = cv2.boxPoints(approx)
        approx = np.int0(box)
        # print(box)
        if len(approx) >=2:
            # img = cv2.drawContours(img,[approx],-1,(0,255,255),3)
            # print(approx)
            approx = np.array(approx).reshape((-1,2))
            shape_dict[i] = approx
            polygons.append(Polygon(approx))
    print(polygons)
    tree = STRtree(polygons)
    print("remove containing boxes")
    del_ind = []
    for i,poly in enumerate(polygons):
        indice = tree.query(poly, predicate="contains").tolist()
        print(i)
        print(tree.geometries.take(indice))
        print(tree.geometries.take(i))
        if len(indice) >0:
            for j in indice:
                if contains(poly,tree.geometries.take(j)) and area(poly)>area(tree.geometries.take(j)):
                    if j+1 in shape_dict:
                        del(shape_dict[j+1])
                        del_ind.append(int(j))
    print(del_ind)
    for i in shape_dict:
        points_tmp = shape_dict[i].copy()
        points_tmp[:,0] = points_tmp[:,1]
        points_tmp[:,1] = shape_dict[i][:,0].copy()
        for j in range(len(points_tmp)):
            p_s = points_tmp[j]
            if j < len(points_tmp)-1:
                p_e = points_tmp[j+1]
            else:
                p_e = points_tmp[0]
            line = p_e - p_s
            length = np.linalg.norm(line)
            # print(p_s,p_e,length)
            for k in range(int(np.ceil(length))):
                if np.ceil(p_s[0]+k*line[0]/length) < occu.shape[0] and np.ceil(p_s[1]+k*line[1]/length)<occu.shape[1] and np.ceil(p_s[1]+k*line[1]/length)>=0 and np.ceil(p_s[0]+k*line[0]/length)>=0:
                    occu[int(np.ceil(p_s[0]+k*line[0]/length)),int(np.ceil(p_s[1]+k*line[1]/length))] = 2
    polygons = []
    for i in shape_dict:
        polygons.append(shape_dict[i])
    print(polygons)
    plt.imshow(occu)
    plt.show()
    return shape_dict

if __name__ == "__main__":
    # Run empty stage
    # for i in range(3):
        # simulation_app = SimulationApp(config)
        main()
        # Close the simulator
        simulation_app.close()


import pyttsx3
import PyPDF2

book = open('sql.pdf', 'rb')
audio = pyttsx3.init()
pdfReader = PyPDF2.PdfFileReader(book)

total_pages = pdfReader.numPages
print(total_pages)

for page in range(total_pages):
    page = pdfReader.getPage(page)
    text = page.extractText()
    audio.say(text)
    audio.runAndWait()


import os
import logging

import click
import pickle
import pandas as pd

# LOGGING
log_format = logging.Formatter("%(asctime)s - [%(filename)s %(funcName)s] - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)  # create logger
logger.setLevel(logging.DEBUG)  # logger level
ch = logging.StreamHandler()  # console handler
ch.setFormatter(log_format)  # format
logger.addHandler(ch)

# DATA
TARGET = "condition"
# FILENAMES
DATA_FILENAME = "data.csv"
MODEL_FILENAME = "model.pkl"
PREDS_FILENAME = "predictions.csv"

@click.command("predict")
@click.option("--input-dir")
@click.option("--output-dir")
@click.option("--model-dir")
def predict(input_dir: str, output_dir: str, model_dir: str):
    logger.info("predicting")

    # is data exist?
    data_path = os.path.join(input_dir, DATA_FILENAME)
    if not os.path.exists(data_path):
        raise ValueError(f"data not found: {data_path}")
    logger.info(f"data exists: {data_path}")
    # is model exist?
    model_path = os.path.join(model_dir, MODEL_FILENAME)
    if not os.path.exists(model_path):
        raise ValueError(f"model not found: {model_path}")
    logger.info(f"model exist: {model_path}")

    # load data
    data = pd.read_csv(data_path)
    # load model
    with open(model_path, "rb") as file:
        model = pickle.load(file)
    # predict
    preds = model.predict(data)
    preds = pd.DataFrame({TARGET: preds})

    os.makedirs(output_dir, exist_ok=True)
    preds_dir = os.path.join(output_dir, PREDS_FILENAME)
    preds.to_csv(preds_dir)

    logger.info(f"predictions are ready: {preds_dir}")

if __name__ == '__main__':
    predict()


import logging
import pickle
from hashlib import sha1
from pathlib import Path
from typing import Any, Union

import typer
from rich import print
from rich.logging import RichHandler

from gil import paths
from gil.core import Blob, Commit, Sha, Tree
from gil.utils import NoRefException, get_ref, set_ref

FORMAT = "%(message)s"
logging.basicConfig(level=logging.DEBUG, format=FORMAT, datefmt="[%X]", handlers=[RichHandler()])
logger = logging.getLogger(__name__)

app = typer.Typer(help="`gil` is a git-like tool for educational purposes.")

########################################################################################
@app.command()
def init():
    """Create a gil repo."""
    pass
    path = Path.cwd() / ".gil"
    if path.exists():
        logger.warning(f"gil repo at {path} already exists")
        return

    path.mkdir()
    (path / "objects").mkdir()
    (path / "refs").mkdir()
    (path / "refs/tags").mkdir()
    (path / "refs/heads").mkdir()
    (path / "refs/remotes").mkdir()
    with (path / "HEAD").open("w") as f:
        f.write("refs/heads/main")
    logger.info(f"gil repo created in '{path}'")

########################################################################################
# BLOB
@app.command()
def hash_object(path: Path):
    """Hash the object and add it to the repo."""
    assert path.is_file()
    data = path.read_bytes()
    sha = hash_data(data)
    dump_obj(sha, Blob(data), path)
    return sha

@app.command()
def cat_file(sha: str, do_print: bool = True):
    """Read the content of the object with the given sha."""
    path = paths.OBJECTS_DIR / sha
    assert path.is_file()
    with path.open("rb") as f:
        obj = pickle.load(f)
    if do_print:
        print(obj)
    return obj

########################################################################################
# TREE
@app.command()
def hash_tree(path: Path):
    """Hash the tree (recursively) and add it to the repo."""
    assert path.is_dir()
    # collect all the data
    items = []
    for p in path.iterdir():
        if p.name == ".gil":
            continue
        elif p.is_file():
            items.append((hash_object(p), str(p.name)))
        elif p.is_dir():
            items.append((hash_tree(p), str(p.name)))
        else:
            raise Exception("WTF!")

    # create and dump tree object
    items = sorted(items)
    sha = hash_data("".join(sha for sha, _ in items))
    dump_obj(sha, Tree(items), path)
    return sha

########################################################################################
# COMMIT
@app.command()
def snapshot(commit_message: str):
    """Think: git add and git commit"""
    tree_sha = hash_tree(Path("."))

    try:
        prev_commit_sha = get_ref()
        _prev_commit = cat_file(prev_commit_sha, do_print=False)
        if _prev_commit.tree == tree_sha:
            logger.warning("Nothing changed, nothing to commit.")
            return
    except NoRefException:
        prev_commit_sha = ""

    commit = Commit(tree_sha, prev_commit_sha, commit_message)
    commit_sha = hash_data((str(commit.parent) + commit.tree).encode())
    dump_obj(commit_sha, commit, "commit")
    set_ref(commit_sha)

########################################################################################
# MISC
@app.command()
def log():
    """log of all commits"""
    commit_sha = get_ref()
    while commit_sha:
        commit: Commit = cat_file(commit_sha, do_print=False)
        print(f"# {commit_sha}")
        print(commit)
        commit_sha = commit.parent

# GRAPH
@app.command()
def graph():
    from graphviz import Digraph

    dot = Digraph()
    for path in paths.OBJECTS_DIR.iterdir():
        sha = path.name
        obj = cat_file(sha)

        if isinstance(obj, Blob):
            dot.node(sha[:8], shape="box")
        elif isinstance(obj, Tree):
            dot.node(sha[:8], shape="triangle")
            for child_sha, child_name in obj.items:
                dot.edge(sha[:8], child_sha[:8], child_name)
        elif isinstance(obj, Commit):
            dot.node(sha[:8], shape="circle")
            dot.edge(sha[:8], obj.tree[:8], obj.commit_message)
            dot.edge(sha[:8], obj.parent[:8], obj.commit_message)
        else:
            raise ValueError("WTF?")

    dot.render("/tmp/gil.pdf", view=True)

########################################################################################
# UTILS
def dump_obj(sha: Sha, obj: Any, message) -> None:
    dst = paths.OBJECTS_DIR / sha
    if dst.exists():
        logger.warning(f"{sha} --> '{message}' already hashed")
    else:
        with dst.open("wb") as f:
            pickle.dump(obj, f)
        logger.info(f"{sha} --> '{message}' hashed")

def hash_data(data: Union[str, bytes]) -> Sha:
    if isinstance(data, str):
        data = data.encode()
    return sha1(data).hexdigest()

########################################################################################
if __name__ == "__main__":
    app()


# -*- coding: utf-8 -*-
from __future__ import annotations

import smtplib

from email.mime.application import MIMEApplication
from email.mime.image import MIMEImage
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from matplotlib.backends.backend_pdf import PdfPages
from pathlib import Path
from tempfile import TemporaryDirectory

SIGNATURE = (
    '
Name
\n'
    "
Company
\n"
    "
Address
\n"
    "
City & Zip
\n"
    "
Office: (000) 000-0000
\n"
    "
Cell: (111) 111-1111
\n"
    "
user@company.com
\n"
)

class EmailMsg:
    def __init__(self, recipient_list, subject, signature=SIGNATURE):
        self.sender = "user@company.com"
        self.subject = subject
        self.recipient_list = recipient_list
        self.signature = signature
        self.attachments = []
        self.images = []
        self.text = []

    def send(self):
        msg = self.construct_msg()
        with smtplib.SMTP("smtp.company.com") as mailer:
            server_response = smtplib.SMTP.ehlo(mailer)
            mailer.sendmail(msg["From"], self.recipient_list, msg.as_string())

    def construct_msg(self):
        msg = MIMEMultipart()
        msg["Subject"] = self.subject
        msg["From"] = self.sender
        for attachment in self.attachments:
            msg.attach(attachment)
        for display in self.images:
            msg.attach(display)
        all_text = "

".join(self.text)
        all_text += "

" + self.signature
        msg.attach(MIMEText(all_text, "html"))
        return msg

    def convert_plots_to_attachment(self, figure_name, figures):
        if not isinstance(figures, (list, tuple)):
            figures = [figures]
        with TemporaryDirectory() as temp_dir:
            dir_path = Path(str(temp_dir))
            file = dir_path / "plots.pdf"
            with PdfPages(file) as pdf_file:
                for figure in figures:
                    pdf_file.savefig(figure, dpi=300, bbox_inches="tight")
            self.attach_file(figure_name, file, "pdf")

    def attach_file(self, attachment_name, file, file_type):
        with file.open("rb") as f:
            attachment = MIMEApplication(f.read(), _subtype=file_type)
        attachment.add_header(
            "Content-Disposition", "attachment", filename=attachment_name
        )
        self.attachments.append(attachment)

    def add_image(self, figure_id, figure):
        with TemporaryDirectory() as temp_dir:
            dir_path = Path(str(temp_dir))
            file = dir_path / "image.png"
            figure.savefig(file, dpi=150, bbox_inches="tight")
            fp = open(file, "rb")
            try:
                attachment = MIMEImage(fp.read())
                attachment.add_header("Content-ID", "<%s>" % figure_id)
                attachment.add_header(
                    "Content-Disposition", "inline; filename=image.png"
                )
            except:
                pass
            fp.close()
        self.add_text(f'

')
        self.images.append(attachment)

    def add_text(self, text):
        self.text.append(text)


"""
Solution for Instant Insanity puzzle - https://en.wikipedia.org/wiki/Instant_Insanity
"""

from enum import Enum, auto

class Faces(Enum):
    TOP = auto()
    LEFT = auto()
    BACK = auto()
    RIGHT = auto()
    FRONT = auto()
    BOTTOM = auto()

    def rotate_faces(self, face_list):
        rotated_faces = []
        rotated_faces.extend(face_list)
        return rotated_faces

    def face_permutations(self):
        fp = []
        orig = list(Faces)
        # Keep top and bottom same
        fp.extend(self.rotate_faces(orig))
        # Flip top and bottom
        # Move left to top
        # Move right to top
        # Move front to top
        # Move back to top
        return fp

Color = Enum('Color', ['RED', 'YELLOW', 'GREEN', 'BLUE'])

def find_orientations(cube):
    orientations = []
    fp = Faces.face_permutations()
    # for p in fp:
    #
    return orientations

print(dir(Color))
print(dir(Faces))

cubes = [dict(zip(list(Faces), [Color.RED, Color.BLUE, Color.YELLOW, Color.GREEN, Color.RED, Color.BLUE])),
         dict(zip(list(Faces), [Color.RED, Color.RED, Color.BLUE, Color.YELLOW, Color.GREEN, Color.GREEN])),
         dict(zip(list(Faces), [Color.RED, Color.BLUE, Color.YELLOW, Color.GREEN, Color.BLUE, Color.YELLOW])),
         dict(zip(list(Faces), [Color.RED, Color.BLUE, Color.YELLOW, Color.YELLOW, Color.GREEN, Color.GREEN]))]
print(cubes)

for d in cubes:
    cube_colors = [(f.name, c.name) for f, c in d.items()]
    print(cube_colors)

print(list(Faces))
print(Faces(1))


"""
Common gamedata utilities
"""

import logging

import mlbv.mlbam.common.config as config
import mlbv.mlbam.common.util as util

LOG = logging.getLogger(__name__)

def get_feedtype_keystring(feedtype_map):
    reverse_list = list()
    for longkey in feedtype_map:
        reverse_list.append("{}:{}".format(feedtype_map[longkey], longkey))
    return ", ".join(reverse_list)

def convert_feedtype_to_short(feedtype, feedtype_map):
    if feedtype in feedtype_map:
        return feedtype_map[feedtype]
    return feedtype

def convert_to_long_feedtype(feed, feedtype_map):
    if feed in feedtype_map:
        return feed
    for feedtype in feedtype_map:
        if feedtype_map[feedtype] == feed:
            return feedtype
    return feed

def is_fav(game_rec):
    if "favourite" in game_rec:
        return game_rec["favourite"]
    if config.CONFIG.parser["favs"] is None or config.CONFIG.parser["favs"] == "":
        return False
    for fav in util.get_csv_list(config.CONFIG.parser["favs"]):
        if fav in (game_rec["away"]["abbrev"], game_rec["home"]["abbrev"]):
            return True
    return False

def apply_filter(game_rec, arg_filter, filters=None):
    """Returns the game_rec if the game matches the filter, or if no filtering is active."""
    if arg_filter == "favs":
        arg_filter = config.CONFIG.parser["favs"]
    elif filters and arg_filter in filters:
        arg_filter = filters[arg_filter]
    elif not arg_filter:
        return game_rec

    # apply the filter
    for team in util.get_csv_list(arg_filter):
        if team in (game_rec["away"]["abbrev"], game_rec["home"]["abbrev"]):
            return game_rec

    # no match
    return None


#Please refer to the README File for the full formulas and explanation

import time #To Calculate the execution time
import numpy as np #To deal with Math
import pandas as pd #To handle the data's part
import scipy.optimize as opt

strt = time.time() #Calculate the time

#------------------------------------------------------------

def sigmoid(z):
    return 1/(1+np.exp(-z) )

#------------------------------------------------------------

def costfunc(thetas,x,y):

    thetas = np.matrix(thetas)
    x = np.matrix(x)
    y = np.matrix(y)

    first = np.multiply( -y ,np.log( sigmoid( np.dot(x, thetas.T) )  )   )
    second = np.multiply( (1-y)  , np.log(  1 - ( sigmoid( np.dot(x, thetas.T) )  )   )  )
    c = np.sum( first - second ) / len(x)

    return c

#------------------------------------------------------------

def Gradientdes( thetas, x, y):

    thetas = np.matrix(thetas)
    x = np.matrix(x)
    y = np.matrix(y)

    parameters = int(thetas.ravel().shape[1])

    grd =  np.zeros(parameters)
    pred = sigmoid( x.dot( thetas.T ) )

    for e in range(parameters):

         term = np.multiply( ( pred - y) ,x[: , e] )
         grd[e] =  ( 1/len(x) )     *   np.sum(term)

    return grd

#------------------------------------------------------------

def predc(thetas,x):
    classif = sigmoid( x * thetas.T )
    e = [ 1 if i >= 0.5 else 0 for i in classif ]
    return e

#------------------------------------------------------------

#Import the Dataset file, The txt file
datas = pd.read_csv(r'C:\Users\NADY\Desktop\dsc.txt', header = None , names=['Xdata0','Xdata1','Ydata'] ) #Write your path here

datas.insert(0,'Ones',1)

#We have the data, So, I want to choose specific columns
cols = datas.shape[1]
x = datas.iloc[:,0:cols-1]
y = datas.iloc[:,cols-1:cols]

x = np.array(x.values)
y = np.array(y.values)
thetas =np.zeros( x.shape[1] )

#--------------------

print('\n The Cost Function before optimization',costfunc(thetas, x, y),'\n')

results = opt.fmin_tnc(func = costfunc, x0 = thetas , fprime= Gradientdes, args=(x,y)  )
print(results)

#--------------------

endt = time.time()

costafteropti = costfunc(results[0], x, y)
print('\nThe Cost Function before optimization\n',costafteropti)

newvalues = predc( np.matrix( results[0] ) , x)
print( '\nThe Expected Values \n', newvalues,'\n' )

#--------------------

correctansrs = [ 1 if (a == 1 and b == 1) or (a == 0 and b == 0) else 0 for (a,b) in zip(y,newvalues)]
print('The Correct Answers are 1 wrong are 0\n',correctansrs)

accuracy = (sum(correctansrs) * 100 )/len(correctansrs)

print( f'\nacc =  {accuracy}% \n' )

print(f'\nThe Execution Time is : {endt-strt} Seconds' )


import matplotlib.pyplot as plt
import numpy as np
import math
import random

class Part2:
    def __init__(self, t=50, y=0.7, separation=25, n=50, chi_squared=13.09051):
        # CONST
        self.T = t
        self.Y = y
        self.SEPARATION = separation
        self.N = n
        self.CHI_SQUARED = chi_squared
        # VARS
        self.x = np.arange(0, self.T, 0.1)
        self.y = [0] * self.T * 10
        self.separateY = [-1, 1]
        self.separateX = np.arange(0, self.T, self.T / self.SEPARATION)
        # Computing vars
        self.flow = self.flows2flow(self.generate())
        [self.countInRange, self.uniqueList] = self.catch_in_interval()
        self.lambda_t = sum(self.flow) / sum(self.countInRange)
        self.theoreticalN = self.theoretical()

    # Генерим из первой части 50 потоков
    def generate(self):
        iterations = []

        for i in range(0, self.T):
            t = 0
            i = 0
            times = []
            while t < self.T:
                y = random.uniform(0, 1)
                x = math.log((1 - y), math.e) / (-self.Y)
                u = x
                t += u
                i += 1

                if t < self.T:
                    rounded_t = round(t, 1)
                    if rounded_t not in times:
                        times.append(rounded_t)

            iterations.append(times)

        return iterations

    # Картинка
    def figure(self):
        plt.plot(self.x, self.y)
        for sepX in self.separateX:
            plt.plot([sepX] * 2, self.separateY)

        plt.margins(0, 5)
        plt.ylabel('y')
        plt.xlabel('x')

        plt.plot(self.flow, [0] * len(self.flow), 'ro', markersize=2)
        plt.show()

    # Получаем количество попаданий в интервал
    def catch_in_interval(self):
        # Левая границы
        left = 0
        # Правая границы
        right = left + self.T / self.SEPARATION
        # Уникальные
        u = None
        # Число попаданий
        c = []

        for i in range(0, self.SEPARATION):
            c.append(len(list(filter(lambda x: left < x < right, self.flow))))
            u = list(range(0, 10))

            # Переход к следующему отрезку
            # левая стала правой
            left = right
            # Правая
            right = left + self.T / self.SEPARATION

            self.flow[:] = [f / 10 for f in self.flow]

        return [c, u]

    # Все в один поток
    @staticmethod
    def flows2flow(iterations):
        a = []
        for x in iterations:
            a += x
        return a

    # Получаем n теоритическое
    def theoretical(self):
        n_theories = []
        n_sum = sum(self.countInRange)

        for n in self.uniqueList:
            p = pow(self.lambda_t, n) / math.factorial(n) * pow(math.e, -self.lambda_t)
            n_theories.append(p * n_sum)

        self.uniqueList[:] = [n - n / self.SEPARATION for n in n_theories]
        return n_theories

    # Получаем хи квадрат
    def chi_squared(self):
        chi = 0
        for count, n in zip(self.uniqueList, self.theoreticalN):
            chi += pow((count - n), 2) / n
        return chi

    def hypothesis(self):
        return self.chi_squared() < self.CHI_SQUARED

laba3 = Part2()
# print(laba3.flow)
print('lambda теориритическое: ' + str(laba3.lambda_t))
print('N теоритическое: ' + str(laba3.theoreticalN))
print('Хи квадрат: ' + str(laba3.chi_squared()))
if laba3.hypothesis():
    print('Гипотеза о пуассоновском потоке принимается')
else:
    print('Гипотеза о пуассоновском потоке отвергается')
laba3.figure()