# -*- coding: utf-8 -*-
"""
author: Diogo Silva
notes: Boruvka implementation based on Sousa's "A Generic and Highly Efficient Parallel Variant of Boruvka ’s Algorithm"
       connected components from Boruvka
"""

import numpy as np
from MyML.helper.scan import exprefixsumNumbaSingle as ex_prefix_sum_cpu, exprefixsumNumba as ex_prefix_sum_cpu2,\
                             scan_gpu as ex_prefix_sum_gpu
from MyML.graph.mst import findMinEdgeNumba, removeMirroredNumba, initColorsNumba, propagateColorsNumba,\
                           buildFlag, countNewEdgesNumba, assignInsertNumba,\
                           compute_cuda_grid_dim, findMinEdge_CUDA, removeMirroredEdges_CUDA, memSet,\
                           initializeColors_CUDA, propagateColors_CUDA, buildFlag_CUDA, countNewEdges_CUDA, assignInsert_CUDA
from numba import jit, cuda, void, boolean, int8, int32, float32

def connected_comps_mst_seq(mst, dest, fe):
    pass

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
#
#                      NUMBA CPU
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

def connected_comps_seq(dest_in, weight_in, firstEdge_in, outDegree_in):
    dest = dest_in
    weight = weight_in
    firstEdge = firstEdge_in
    outDegree = outDegree_in

    n_vertices = firstEdge.size
    n_edges = dest.size

    n_components = n_vertices

    # still need edge_id for conflict resolution in find_minedge
    edge_id = np.arange(n_edges, dtype = dest.dtype)

    #labels = np.empty(n_vertices, dtype = dest.dtype)
    first_iter = True

    # initialize with name top_edge so we can recycle an array between iterations
    top_edge = np.empty(n_components, dtype = dest.dtype)

    final_converged = False
    while(not final_converged):
        vertex_minedge = top_edge
        findMinEdgeNumba(vertex_minedge, weight, firstEdge, outDegree, dest)
        removeMirroredNumba(vertex_minedge, dest)

        # intialize colors of current graph
        colors = np.empty(n_components, dtype = dest.dtype)
        initColorsNumba(vertex_minedge, dest, colors)

        # propagate colors until convergence
        converged = False
        while(not converged):
            converged = propagateColorsNumba(colors)

        # flag marks the vertices that are the representatives of the new supervertices
        # new_vertex will be initialized with he flags
        new_vertex = vertex_minedge
        buildFlag(colors, new_vertex)

        new_n_vertices = ex_prefix_sum_cpu(new_vertex, init = 0)

        if first_iter:
            # first iteration defines labels as the initial colors and updates
            labels = colors.copy()
            first_iter = False

        # update the labels with the new colors
        update_labels_single_pass(labels, colors, new_vertex)

        if new_n_vertices == 1:
            final_converged = True
            break

        # count number of edges for new supervertices and write in new outDegree
        newOutDegree = np.zeros(new_n_vertices, dtype = dest.dtype)
        countNewEdgesNumba(colors, firstEdge, outDegree, dest, new_vertex, newOutDegree)

        # new first edge array for contracted graph
        newFirstEdge = np.empty(newOutDegree.size, dtype = dest.dtype)
        new_n_edges = ex_prefix_sum_cpu2(newOutDegree, newFirstEdge, init = 0)

        # if no edges remain, then MST has converged
        if new_n_edges == 0:
            final_converged = True
            break

        # create arrays for new edges
        new_dest = np.empty(new_n_edges, dtype = dest.dtype)
        new_edge_id = np.empty(new_n_edges, dtype = dest.dtype)
        new_weight = np.empty(new_n_edges, dtype = weight.dtype)
        top_edge = newFirstEdge.copy()

        # assign and insert new edges
        assignInsertNumba(edge_id, dest, weight, firstEdge, outDegree, colors,
                          new_vertex, new_dest, new_edge_id, new_weight, top_edge)

        # delete old graph
        del new_vertex, edge_id, dest, weight, firstEdge, outDegree, colors

        # write new graph
        n_components = newFirstEdge.size
        edge_id = new_edge_id
        dest = new_dest
        weight = new_weight
        firstEdge = newFirstEdge
        outDegree = newOutDegree

    return labels

@jit
def update_labels_numba(labels, update_array):
    """
    This kernel is dual purpose.

    This kernel updates the color of each vertex in the graph.
    The current colors of the vertices are indices for the current components
    in the graph. Each current component in the graph will have a new color.
    The new color of vertex v will be the new color colors[curr_color] of its
    color labels[v] (index of a component). This step happens because the new
    color propagation is representative of the contracted graph. In this phase
    the update_array is the colors array.

    Old components are merged into new ones with new IDs. The labels need to
    have the new IDs. In this phase the update_array is the new_vertex array.
    """
    n_components = labels.size

    for v in range(n_components):
        curr_color = labels[v]
        labels[v] = update_array[curr_color]

@jit
def update_labels_single_pass(labels, colors, new_vertex):
    """
    Does all the updates on a single pass
    """
    n_components = labels.size

    for v in range(n_components):
        curr_color = labels[v]
        new_color = colors[curr_color]
        new_color_id = new_vertex[new_color]
        labels[v] = new_color_id

def connected_comps_gpu(dest_in, weight_in, firstEdge_in, outDegree_in,
                        MAX_TPB = 512, stream = None):
    if stream is None:
        myStream = cuda.stream()
    else:
        myStream = stream

    dest = cuda.to_device(dest_in, stream = myStream)
    weight = cuda.to_device(weight_in, stream = myStream)
    firstEdge = cuda.to_device(firstEdge_in, stream = myStream)
    outDegree = cuda.to_device(outDegree_in, stream = myStream)

    n_vertices = firstEdge.size
    n_edges = dest.size

    n_components = n_vertices

    # still need edge_id for conflict resolution in find_minedge
    edge_id = cuda.to_device(np.arange(n_edges, dtype = dest.dtype), stream = myStream)

    #labels = np.empty(n_vertices, dtype = dest.dtype)
    first_iter = True

    # initialize with name top_edge so we can recycle an array between iterations
    top_edge = cuda.device_array(n_components, dtype = dest.dtype, stream = myStream)
    labels = cuda.device_array(n_components, dtype = dest.dtype, stream = myStream)

    converged = cuda.device_array(1, dtype = np.int8, stream = myStream)

    gridDimLabels = compute_cuda_grid_dim(n_components, MAX_TPB)
    gridDim = compute_cuda_grid_dim(n_components, MAX_TPB)

    final_converged = False
    while(not final_converged):
        vertex_minedge = top_edge

        findMinEdge_CUDA[gridDim, MAX_TPB, myStream](weight, firstEdge, outDegree, vertex_minedge, dest)

        removeMirroredEdges_CUDA[gridDim, MAX_TPB, myStream](dest, vertex_minedge)

        colors = cuda.device_array(shape = n_components, dtype = np.int32, stream = myStream)
        initializeColors_CUDA[gridDim, MAX_TPB, myStream](dest, vertex_minedge, colors)

        # propagate colors until convergence
        propagateConverged = False
        while(not propagateConverged):
            propagateColors_CUDA[gridDim, MAX_TPB, myStream](colors, converged)
            converged_num = converged.getitem(0, stream = myStream)
            propagateConverged = True if converged_num == 1 else False

        # first we build the flags in the new_vertex array
        new_vertex = vertex_minedge # reuse the vertex_minedge array as the new new_vertex
        buildFlag_CUDA[gridDim, MAX_TPB, myStream](colors, new_vertex)

        # new_n_vertices is the number of vertices of the new contracted graph
        new_n_vertices = ex_prefix_sum_gpu(new_vertex, MAX_TPB = MAX_TPB, stream = myStream).getitem(0, stream = myStream)
        new_n_vertices = int(new_n_vertices)

        if first_iter:
            # first iteration defines labels as the initial colors and updates
            labels.copy_to_device(colors, stream = myStream)
            first_iter = False

        # other iterations update the labels with the new colors
        update_labels_single_pass_cuda[gridDimLabels, MAX_TPB, myStream](labels, colors, new_vertex)

        if new_n_vertices == 1:
            final_converged = True
            del new_vertex
            break

        newGridDim = compute_cuda_grid_dim(n_components, MAX_TPB)

        # count number of edges for new supervertices and write in new outDegree
        newOutDegree = cuda.device_array(shape = new_n_vertices, dtype = np.int32, stream = myStream)
        memSet[newGridDim, MAX_TPB, myStream](newOutDegree, 0) # zero the newOutDegree array
        countNewEdges_CUDA[gridDim, MAX_TPB, myStream](colors, firstEdge, outDegree, dest, new_vertex, newOutDegree)

        # new first edge array for contracted graph
        newFirstEdge = cuda.device_array_like(newOutDegree, stream = myStream)
        newFirstEdge.copy_to_device(newOutDegree, stream = myStream) # copy newOutDegree to newFirstEdge
        new_n_edges = ex_prefix_sum_gpu(newFirstEdge, MAX_TPB = MAX_TPB, stream = myStream)

        new_n_edges = new_n_edges.getitem(0, stream = myStream)
        new_n_edges = int(new_n_edges)

        # if no edges remain, then MST has converged
        if new_n_edges == 0:
            final_converged = True
            del newOutDegree, newFirstEdge, new_vertex
            break

        # create arrays for new edges
        new_dest = cuda.device_array(new_n_edges, dtype = np.int32, stream = myStream)
        new_edge_id = cuda.device_array(new_n_edges, dtype = np.int32, stream = myStream)
        new_weight = cuda.device_array(new_n_edges, dtype = weight.dtype, stream = myStream)

        top_edge = cuda.device_array_like(newFirstEdge, stream = myStream)
        top_edge.copy_to_device(newFirstEdge, stream = myStream)

        # assign and insert new edges
        assignInsert_CUDA[gridDim, MAX_TPB, myStream](edge_id, dest, weight, firstEdge,
                                            outDegree, colors, new_vertex,
                                            new_dest, new_edge_id, new_weight, top_edge)

        # delete old graph
        del new_vertex, edge_id, dest, weight, firstEdge, outDegree, colors

        # write new graph
        n_components = newFirstEdge.size
        edge_id = new_edge_id
        dest = new_dest
        weight = new_weight
        firstEdge = newFirstEdge
        outDegree = newOutDegree
        gridDim = newGridDim

    returnLabels = labels.copy_to_host()

    del dest, weight, edge_id, firstEdge, outDegree, converged, labels

    return returnLabels

@cuda.jit
def update_labels_cuda(labels, update_array):
    """
    CUDA version of update_labels.
    """

    v = cuda.grid(1)
    n_components = labels.size

    if v >= n_components:
        return

    curr_color = labels[v]
    labels[v] = update_array[curr_color]

@cuda.jit
def update_labels_single_pass_cuda(labels, colors, new_vertex):
    """
    Does all the updates on a single pass
    """
    v = cuda.grid(1)
    n_components = labels.size

    if v >= n_components:
        return

    curr_color = labels[v]
    new_color = colors[curr_color]
    new_color_id = new_vertex[new_color]
    labels[v] = new_color_id


from django.views.generic import View, ListView, TemplateView, DetailView
from django.shortcuts import render
from django.http import JsonResponse, HttpResponse, QueryDict
from django.urls import reverse

from .models import students
from .form import StuModelForm

class AddStu(TemplateView):
    template_name = 'stuadd_jq.html'

class ListStu(ListView):
    template_name = 'stulist_jq.html'
    model = students
    context_object_name = "all_students"
    '''通过上面3行，就可以将数据从数据库中全部拿到，然后渲染到stulist_jq.html中'''

    def post(self, request):
        print(request.POST)
        StuFormdata = StuModelForm(request.POST)
        if StuFormdata.is_valid():
            StuFormdata.save() #将通过表单验证的数据保存到数据库中
            res = {'code':0, 'msg':'用户添加成功', 'next_url':reverse('students:jlist')}
        else:
            #print(StuFormdata.errors.as_json())
            print(StuFormdata.errors.as_text())
            res = {'code':1, 'errmsg':StuFormdata.errors.as_text(), 'next_url':reverse('students:jadd')}
        #return render(request, 'stulist.html')
        return JsonResponse(res)

    '''删除用户'''
    def delete(self, request):
        data = QueryDict(request.body).dict() #通过前端ajax将ID传递到后端
        print(data)
        pk = data.get('id')
        try:
            user = self.model.objects.filter(pk=pk)
            user.delete()
            res = {'code':0, 'result':'删除用户成功'}
        except:
            res = {'code':1, 'errmsg':'删除用户失败'}
        return JsonResponse(res, safe=True)

class DetailStu(DetailView):
    template_name = 'stuupdate_jq.html'
    model = students
    context_object_name = "stu"

    def get_object(self, queryset=None):
        obj = super(DetailStu, self).get_object()
        return obj

    def post(self, request, **kwargs):
        print(request.POST)
        print(kwargs) #例子 {'pk': '3'}
        user = students.objects.get(pk=kwargs['pk'])
        StuFormdata = StuModelForm(request.POST, instance=user)  #modelform对象instance是什么作用？
        if StuFormdata.is_valid():
            try:
                print(StuFormdata)
                StuFormdata.save()#将数据保存到数据库
                res = {'code':0, 'result': '用户更新成功', 'next_url':reverse('students:jlist')}
            except:
                res = {'code':1, 'errmsg': '用户更新失败', 'next_url':reverse('students:jlist')}
        else:
            print(StuFormdata.errors)
            res = {'code': 1, 'errmsg': StuFormdata.errors, 'next_url': reverse('students:jlist')}
        return render(request,'jump_jq.html',res)

# class AddStu(View):
#     def get(self, request):
#         return render(request, 'stuadd_jq.html')
#
#     def post(self, request):
#         print(request.POST)
#         return render(request, 'stuadd_jq.html')


import time

start_time = time.time()

f = open('names_1.txt', 'r')
names_1 = f.read().split("\n")  # List containing 10000 names
f.close()

f = open('names_2.txt', 'r')
names_2 = f.read().split("\n")  # List containing 10000 names
f.close()

duplicates = []  # Return the list of duplicates in this data structure

# Replace the nested for loops below with your improvements
for name_1 in names_1:
    for name_2 in names_2:
        if name_1 == name_2:
            duplicates.append(name_1)

# Sol1 - 1.46s - O(n^2)
# for name_1 in names_1:
    # if name_1 in names_2:
    # duplicates.append(name_1)

# Sol2 - 1.6s - O(n^2)
# duplicates = [name for name in names_1 if name in names_2]

# Sol3 - sort and then binary search - O(nlog(n)) - 0.5s
# Binary Search - log(n)
def myBinarySearch(items, item):
    if len(items) is 1:
        if items[0] is item:
            return True
        else:
            return False
    elif item == items[int(len(items)/2)]:
        return True
    elif item < items[int(len(items)/2)]:
        newItems = items[0:int(len(items)/2)]
        return myBinarySearch(newItems, item)
    else:
        newItems = items[int(len(items)/2):]
        return myBinarySearch(newItems, item)

# Built in Sorting - nlog(n)
names_1.sort()

# Iterating over array and searching - nlog(n)
for name in names_2:
    if myBinarySearch(names_1, name):
        duplicates.append(name)

end_time = time.time()
print (f"{len(duplicates)} duplicates:\n\n{', '.join(duplicates)}\n\n")
print (f"runtime: {end_time - start_time} seconds")

# ---------- Stretch Goal -----------
# Python has built-in tools that allow for a very efficient approach to this problem
# What's the best time you can accomplish?  Thare are no restrictions on techniques or data
# structures, but you may not import any additional libraries that you did not write yourself.


def check(aline):
    pools = []
    leftedge = []

    for i, c in enumerate(aline):
        if c == "\\":
            leftedge.append(i)
        elif c == "_":
            pass
        else:
            if leftedge:
                left = leftedge.pop()
                w = i - left

                while pools and left < pools[-1][0]:
                    w += pools.pop()[1]

                pools.append((left, w))

    return pools

if __name__ == '__main__':

    al = input().strip()
    ps = check(al)

    print(sum([x[1] for x in ps]))
    print(len(ps), *[x[1] for x in ps])


# api.py

class ConnectorKey:
    """
    A ConnectorKey uniquely identifies a connector. The connector key
    comprises of (bundleName, bundleVersion, connectorName).
    """

    def __init__(self, bundleName, bundleVersion, connectorName):

        if bundleName is None:
            raise IllegalArgumentException("bundleName cannot be null")

        if bundleVersion is None:
            raise IllegalArgumentException("bundleVersion cannot be null")

        if connectorName is None:
            raise IllegalArgumentException("connectorName cannot be null")

        self.bundleName = bundleName
        self.bundleVersion = bundleVersion
        self.connectorName = connectorName

    def __eq__(self, o):
        if isinstance(o, ConnectorKey):
            if self.bundleName != o.bundleName \
                or self.bundleVersion != o.bundleVersion \
                or self.connectorName != o.connectorName:
                return False
            return True
        return False

    def __hash__(self):
        return hash(self)

    def __str__(self):
        return \
        "ConnectorKey(bundleName=%s,bundleVersion=%s,connectorName=%s" \
        % (self.bundleName, self.bundleVersion, self.connectorName)


import pika

connection = pika.BlockingConnection(pika.ConnectionParameters('localhost'))
channel = connection.channel()
channel.queue_declare(queue='userlog')

def callback(ch, method, properties, body):
    print(" [x] %s" % body.decode('utf-8'))

channel.basic_consume(callback, queue='userlog', no_ack=True)

print(' [*] Waiting for logging messages. To exit press CTRL+C')
channel.start_consuming()


# -*- coding: <utf-8> -*-
from bottle import *
import urllib.request
import json
from threading import Timer
import codecs
def uberNameToID(uName):
    stats_request=urllib.request.urlopen('http://pastats.com/report/getplayerid?ubername='+uName)
    return stats_request.read().decode('utf-8')
def update_everything():
    global l
    try:
        l.updateall()
    except:
        pass
    Timer(600,update_everything).start()
def update_leaderboards():
    global l
    l.update()
    Timer(600,update_leaderboards).start()
class Leaderboard:
    def __init__(self):
        self.UNametoDName={}
        self.UIDtoName={}
        self.UNametoPAS={}
        self.ladder_url='https://cdn.palobby.com/api/ladder/'
        self.nconvert_url='http://4.uberent.com/GameClient/UserNames?TitleId=4&UberIds='
        self.update()

    def update(self):
        self.data=json.loads(urllib.request.urlopen(self.ladder_url).read().decode('utf-8'))
        self.totals={'classic':self.data['classic-ranks'],'titans':self.data['titans-ranks']}
        self.player_numbers={'classic':self.data['classic-total'],'titans':self.data['titans-total']}
        self.ladders={'classic':self.data['classic-players'],'titans':self.data['titans-players']}

        for z in self.ladders:
            for x in self.ladders[z]:
                player_place=1
                for y in self.ladders[z][x]:
                    y['place']=player_place
                    player_place+=1
                    if y['UberId'] in self.UIDtoName:
                        y['UberName']=self.UIDtoName[y['UberId']]
                    else:
                        uberResponse=urllib.request.urlopen(self.nconvert_url+y['UberId']).read().decode('utf-8')
                        uberResponse=json.loads(uberResponse)
                        uberName=uberResponse['Users'][y['UberId']]['UberName']
                        displayName=uberResponse['Users'][y['UberId']]['TitleDisplayName']
                        y['DisplayName']=displayName
                        self.UNametoDName=displayName
                        y['UberName']=uberName
                        self.UIDtoName[y['UberId']]=uberName

                    if y['UberName'] in self.UNametoPAS:
                        y['PAStatsId']=self.UNametoPAS[y['UberName']]
                    else:
                        PASId=uberNameToID(y['UberName'])
                        y['PAStatsId']=PASId
                        self.UNametoPAS[y['UberName']]=PASId
        self.ladders_final=self.ladders.copy()

    def updateall(self):
        self.__init__()
def generate_table(ladder):
    table=''
    ranks={'1':'Uber','2':'Platinum','3':'Gold','4':'Silver','5':'Bronze'}
    for x in ladder:
        table+='
\n
\n	User	Last Match At	Last League Change	Replay
\n'
        for y in ladder[x]:
            table+='
'+str(y['place'])+'	\n'
            table+=''+y['DisplayName']+'	\n'
            table+=''+y['LastMatchAt'].replace('Z','')+'	\n'
            table+=''+y['LastRankChange'].replace('Z','')+'	\n'
            table+='PAStats Replay	\n
\n'
        table+='
\n\n'
    return table

l=Leaderboard()
lserver=Bottle()
Timer(300,update_everything).start()
#Timer(600,update_leaderboards).start()
@lserver.route('/api')
def return_api():
    out_dict={'totals':l.totals,'player_numbers':l.player_numbers,'ladders':l.ladders}
    return json.dumps(out_dict)
@lserver.route('/static/<file>')
def return_static(file):
    return static_file(file,'./static')
@lserver.route('/')
def index():
    return static_file('index.html','./')
@lserver.route('/titans')
def titans():
    ladder=l.ladders_final['titans']
    table=generate_table(ladder)
    template_out=SimpleTemplate(name='main.html')
    return template_out.render(tables=table)
@lserver.route('/classic')
def classic():
    ladder=l.ladders_final['classic']
    table=generate_table(ladder)
    template_out=SimpleTemplate(name='main.html')
    return template_out.render(tables=table)

lserver.run(port=80,server='twisted',host='YOURHOSTHERE.com')


import cv2
import numpy

video = cv2.VideoCapture('s_job.mp4')       # read image

while(video.isOpened()):                    # till video remains open
    ret, frame = video.read()               # read frames
    cv2.imshow('video', frame)              # show video
    if cv2.waitKey(1) & 0xFF == ord('q'):   # wait for response from keyboard
        break
video.release()                             # close opened vidoe
cv2.destroyAllWindows()


"""
# 493. Reverse Pairs

# Given an array nums, we call (i, j) an important reverse pair if i < j and nums[i] > 2*nums[j].

# You need to return the number of important reverse pairs in the given array.

# Example1:
# Input: [1,3,2,3,1]
# Output: 2

# Example2:
# Input: [2,4,3,5,1]
# Output: 3

# https://leetcode.com/problems/reverse-pairs/discuss/97268/General-principles-behind-problems-similar-to-%22Reverse-Pairs%22
"""
import bisect

class ReversePairs:

    """
    Complexity analysis

    Time complexity: O(n\log n)O(nlogn)

    In each step we divide the array into 2 sub-arrays, and hence, the maximum times we need to divide is equal to O(\log n)O(logn)
    Additional O(n)O(n) work needs to be done to count the inversions and to merge the 2 sub-arrays after sorting. Hence total time complexity is O(n\log n)O(nlogn)
    Space complexity: O(n)O(n). Additional space for storing LL and RR arrays

    """
    def doit_divide_conquer(self, nums: list) -> int:

        def merge(nums):
            if len(nums) <= 1:
                return nums, 0

            mid = len(nums) // 2
            left, lcount = merge(nums[:mid])
            right, rcount = merge(nums[mid:])

            count = lcount + rcount
            left_len = len(left)

            for r in right:
                add_to = left_len - bisect.bisect(left, 2 * r)
                if add_to == 0:
                    break
                count += add_to

            return sorted(left + right), count

        arr, count = merge(nums)
        return count

    def doit_binarysearch(self, nums: list) -> int:
        import bisect

        A, res = [], 0
        for i in reversed(range(len(nums))):
            index = bisect.bisect_left(A, nums[i])
            res += index
            bisect.insort(A, 2 * nums[i])
        return res

    # n*log(n) <important> <treeArray> <BIT> Binary Index Tree
    def doit(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """
        # get all values used, num and  (2 * num + 1)
        s = set(nums + [num * 2 + 1 for num in nums])

        # get a dict to map the position of each 'value' in "bit"
        d = {v: i + 1 for i, v in enumerate(sorted(s))}

        # the bit tree array, each one represent the 'cnt' of the 'value', greater than
        bit = [0] * (len(s) + 1)

        # <important> 'pre' != pre(next('pre'))   next(pre(next)) != next

        # insert
        def insert(i):
            while i:
                # add one to i and all i's lower point,
                # so when you locate to 'i', you know how many 'cnt' is more than current 'i'
                bit[i] += 1
                # contine to go "pre" node, and add cnt. then when 'pre' go his 'next', it can acculuate it
                # go down to next level
                i -= i & (-i)

        def getSum(i):
            s = 0
            # current 'i' amount bit[i] elements more than 'd[2 * num + 1]'
            while i < len(bit):
                s += bit[i]
                # go up level to larger ones
                i += i & (-i)
            return s

        res=0
        for num in nums:
            res += getSum(d[2 * num + 1])
            insert(d[num])

        return res

    # O(nlog(n))
    def doit2(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """
        def msort(list):
            L = len(list)
            if len(list) < 2:
                return list
            else:
                return merge(msort(list[:L//2]), msort(list[L//2:]))

        self.res_ = 0

        def merge(left, right):

            l, r = 0, 0
            while l < len(left) and r < len(right):
                if left[l] <= 2 * right[r]:
                    l += 1
                else:
                    self.res_ += len(left) - l
                    r += 1

            return sorted(left + right)

        msort(nums)
        return self.res_

    # O(n*log(n))
    def doit_array_1(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """
        import bisect
        A, res = [], 0
        for i in reversed(range(len(nums))):
            index = bisect.bisect_left(A, nums[i])
            res += index
            bisect.insort(A, 2 * nums[i])

        return res

    # O(n*log(n))
    def doit_array_2(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """
        import bisect
        A, res = [], 0
        for i in range(len(nums)):
            index = bisect.bisect_left(A, 2 * nums[i] + 1)
            res += len(A) - index
            bisect.insort(A, nums[i])
        return res

    # O(n^2)
    def doit1(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """
        res = 0
        for i in range(0, len(nums)):

            for j in range( i+1, len(nums)):

                if nums[i] > 2 * nums[j]:
                    res += 1

        return res

    # O(n^2)
    def doit1(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """
        res = 0
        for i in reversed(range(0, len(nums))):

            for j in range(0, i-1):

                if nums[i] > 2 * nums[j]:
                    res += 1

        return res

# It looks like a host of solutions are out there (BST-based, BIT-based, Merge-sort-based).

# Here I�d like to focus on the general principles behind these solutions and its possible application to a number of similar problems.

# The fundamental idea is very simple: break down the array and solve for the subproblems.

# A breakdown of an array naturally reminds us of subarrays. To smoothen our following discussion,
# let�s assume the input array is nums, with a total of n elements. Let nums[i, j] denote the subarray starting from index i to index j (both inclusive),
# T(i, j) as the same problem applied to this subarray (for example, for Reverse Pairs, T(i, j) will represent the total number of important reverse pairs for subarray nums[i, j]).

# With the definition above, it�s straightforward to identify our original problem as T(0, n - 1).
# Now the key point is how to construct solutions to the original problem from its subproblems.
# This is essentially equivalent to building recurrence relations for T(i, j).
# Since if we can find solutions to T(i, j) from its subproblems, we surely can build solutions to larger subarrays
# until eventually the whole array is spanned.

# While there may be many ways for establishing recurrence relations for T(i, j), here I will only introduce the following two common ones:

# T(i, j) = T(i, j - 1) + C, i.e., elements will be processed sequentially and C denotes the subproblem
# for processing the last element of subarray nums[i, j]. We will call this sequential recurrence relation.

# T(i, j) = T(i, m) + T(m + 1, j) + C where m = (i+j)/2, i.e., subarray nums[i, j] will be further partitioned into
# two parts and C denotes the subproblem for combining the two parts. We will call this partition recurrence relation.

# For either case, the nature of the subproblem C will depend on the problem under consideration,
# and it will determine the overall time complexity of the original problem. So usually it�s crucial to find efficient
# algorithm for solving this subproblem in order to have better time performance. Also pay attention to possibilities of overlapping subproblems,
# in which case a dynamic programming (DP) approach would be preferred.

# Next, I will apply these two recurrence relations to this problem �Reverse Pairs� and list some solutions for your reference.

# I -- Sequential recurrence relation

# Again we assume the input array is nums with n elements and T(i, j) denotes the total number of important reverse pairs for subarray nums[i, j].
# For sequential recurrence relation, we can set i = 0, i.e., the subarray always starts from the beginning. Therefore we end up with:

# T(0, j) = T(0, j - 1) + C

# where the subproblem C now becomes "find the number of important reverse pairs with the first element of the pair coming
# from subarray nums[0, j - 1] while the second element of the pair being nums[j]".

# Note that for a pair (p, q) to be an important reverse pair, it has to satisfy the following two conditions:

# p < q: the first element must come before the second element;
# nums[p] > 2 * nums[q]: the first element has to be greater than twice of the second element.
# For subproblem C, the first condition is met automatically; so we only need to consider the second condition,
# which is equivalent to searching for all elements within subarray nums[0, j - 1] that are greater than twice of nums[j].

# The straightforward way of searching would be a linear scan of the subarray, which runs at the order of O(j).
# From the sequential recurrence relation, this leads to the naive O(n^2) solution.

# To improve the searching efficiency, a key observation is that the order of elements in the subarray does not matter,
# since we are only interested in the total number of important reverse pairs.
# This suggests we may sort those elements and do a binary search instead of a plain linear scan.

# If the searching space (formed by elements over which the search will be done) is �static� (it does not vary from run to run),
# placing the elements into an array would be perfect for us to do the binary search. However, this is not the case here.
# After the j-th element is processed, we need to add it to the searching space so that it becomes searchable for later elements,
# which renders the searching space expanding as more and more elements are processed.

# Therefore we�d like to strike a balance between searching and insertion operations.
# This is where data structures like binary search tree (BST) or binary indexed tree (BIT) prevail,
# which offers relatively fast performance for both operations.

# 1. BST-based solution
# we will define the tree node as follows, where val is the node value and cnt is the total number of elements
# in the subtree rooted at current node that are greater than or equal to val:

class treeNode:

    def __init__(self, val):
        self.val_ = val
        self.cnt_ = 1

        self.left_, self.right_ = None, None

def search(root, val):

    if not root:
        return 0

    elif root.val_ == val:
        return root.cnt_

    elif root.val_ > val:
        return root.cnt_ + search(root.left_, val)

    elif root.val_ < val:
        return search(root.right_, val)

def insert(root, val):

    if root == None:
        root = treeNode(val)

    elif root.val_ == val:
        root.cnt_ += 1

    elif val < root.val_:
        root.left_ = insert(root.left_, val)

    elif val > root.val_:
        root.cnt_ += 1
        root.right_ = insert(root.right_, val)

    return root

# And finally the main program, in which we will search for all elements no less than twice of current element plus 1
# (converted to long type to avoid overflow) while insert the element itself into the BST.

# Note: this homemade BST is not self-balanced and the time complexity can go as bad as O(n^2)
# (in fact you will get TLE if you copy and paste the solution here). To guarantee O(nlogn) performance,
# use one of the self-balanced BST�s (e.g. Red-black tree, AVL tree, etc.).

class Solution:

   def doit(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """
        res = 0
        root = None
        for c in nums:
            res += search(root, 2 * c + 1)
            root = insert(root, c)

        return res

# 2. BIT-based solution

# For BIT, the searching and insertion operations are:

def search(bit, i):
    sum = 0
    while i < len(bit):
        sum += bit[i]
        i += i & -i

    return sum

def insert(bit, i):
    while i > 0:
        bit[i] += 1
        i -= i & -i

class solution:

    def doit(self, nums):

        def index(val):
            l, r = 0, len(sort_nums) - 1
            while l <= r:
                m = l + ((r - l) >> 1)

                if sort_nums[m] >= val:
                    r = m - 1
                else:
                    l = m + 1

            return l + 1

        bit = [0] * (len(nums) + 1)
        res = 0
        sort_nums = sorted(nums)

        for c in nums:
            res += search(bit, index(sort_nums, 2 * c + 1))
            insert(bit, index(sort_nums, c))

        return res

# More explanation for the BIT-based solution:

# We want the elements to be sorted so there is a sorted version of the input array which is copy.

# The bit is built upon this sorted array. Its length is one greater than that of the copy array to account for the root.

# Initially the bit is empty and we start doing a sequential scan of the input array. For each element being scanned, we first search the bit to find all elements greater than twice of it and add the result to res. We then insert the element itself into the bit for future search.

# Note that conventionally searching of the bit involves traversing towards the root from some index of the bit, which will yield a predefined running total of the copy array up to the corresponding index. For insertion, the traversing direction will be opposite and go from some index towards the end of the bit array.

# For each scanned element of the input array, its searching index will be given by the index of the first element in the copy array that is greater than twice of it (shifted up by 1 to account for the root), while its insertion index will be the index of the first element in the copy array that is no less than itself (again shifted up by 1). This is what the index function is for.

# For our case, the running total is simply the number of elements encountered during the traversal process. If we stick to the convention above, the running total will be the number of elements smaller than the one at the given index, since the copy array is sorted in ascending order. However, we�d actually like to find the number of elements greater than some value (i.e., twice of the element being scanned), therefore we need to flip the convention. This is what you see inside the search and insert functions: the former traversing towards the end of the bit while the latter towards the root.

# II -- Partition recurrence relation
# For partition recurrence relation, setting i = 0, j = n - 1, m = (n-1)/2, we have:

# T(0, n - 1) = T(0, m) + T(m + 1, n - 1) + C

# where the subproblem C now reads "find the number of important reverse pairs with the first element of the pair coming
# from the left subarray nums[0, m] while the second element of the pair coming from the right subarray nums[m + 1, n - 1]".

# Again for this subproblem, the first of the two aforementioned conditions is met automatically.
# As for the second condition, we have as usual this plain linear scan algorithm, applied for each element in the left (or right) subarray.
# This, to no surprise, leads to the O(n^2) naive solution.

# Fortunately the observation holds true here that the order of elements in the left or right subarray does not matter,
# which prompts sorting of elements in both subarrays. With both subarrays sorted, the number of important reverse pairs
# can be found in linear time by employing the so-called two-pointer technique: one pointing to elements in the left subarray while the other to those in the right subarray and both pointers will go only in one direction due to the ordering of the elements.

# The last question is which algorithm is best here to sort the subarrays. Since we need to partition the array into halves anyway, it is most natural to adapt it into a Merge-sort. Another point in favor of Merge-sort is that the searching process above can be embedded seamlessly into its merging stage.

# So here is the Merge-sort-based solution, where the function "reversePairsSub" will return the total number of important reverse pairs within subarray nums[l, r]. The two-pointer searching process is represented by the nested while loop involving variable p, while the rest is the standard merging algorithm.

# III -- Summary

# Many problems involving arrays can be solved by breaking down the problem into subproblems applied on subarrays and then link the solution to the original problem with those of the subproblems, to which we have sequential recurrence relation and partition recurrence relation. For either case, it�s crucial to identify the subproblem C and find efficient algorithm for approaching it.

# If the subproblem C involves searching on �dynamic searching space�, try to consider data structures that support relatively fast operations on both searching and updating (such as self-balanced BST, BIT, Segment tree, ...).

# If the subproblem C of partition recurrence relation involves sorting, Merge-sort would be a nice sorting algorithm to use. Also, the code could be made more elegant if the solution to the subproblem can be embedded into the merging process.

# If there are overlapping among the subproblems T(i, j), it�s preferable to cache the intermediate results for future lookup.

# Lastly let me name a few leetcode problems that fall into the patterns described above and thus can be solved with similar ideas.

# 315. Count of Smaller Numbers After Self
# 327. Count of Range Sum

# For leetcode 315, applying the sequential recurrence relation (with j fixed), the subproblem C reads: find the number of elements out of visited ones that are smaller than current element, which involves searching on �dynamic searching space�; applying the partition recurrence relation, we have a subproblem C: for each element in the left half, find the number of elements in the right half that are smaller than it, which can be embedded into the merging process by noting that these elements are exactly those swapped to its left during the merging process.

# For leetcode 327, applying the sequential recurrence relation (with j fixed) on the pre-sum array, the subproblem C reads: find the number of elements out of visited ones that are within the given range, which again involves searching on �dynamic searching space�; applying the partition recurrence relation, we have a subproblem C: for each element in the left half, find the number of elements in the right half that are within the given range, which can be embedded into the merging process using the two-pointer technique.

# Anyway, hope these ideas can sharpen your skills for solving array-related problems.


import re

discs = []

with open("input.txt") as f:
    for l in f.readlines():
        m = re.match(r'Disc #(\d+) has (\d+) positions; at time=0, it is at position (\d+).', l)
        if m:
            discs.append({"disc": int(m.group(1)), "size": int(m.group(2)), "pos": int(m.group(3))})

t = 0

while True:
    res = []
    for x in discs:
        # starting at time t, disc X takes x steps to reach the destination
        res.append((x["pos"] + t + x["disc"]) % x["size"])

    if sum(res) == 0:
        print(t)
        break

    t += 1


# <Source: https://github.com/rtqichen/ffjord/blob/master/lib/toy_data.py >

import numpy as np
import sklearn
import torch
import sklearn.datasets
from PIL import Image
import os

# Dataset iterator
def inf_train_gen(data, rng=None, batch_size=200):
    if rng is None:
        rng = np.random.RandomState()
        #rng = np.random.RandomState(42)
        #print(rng)

    if data == "1Normal":
        z_dist = torch.distributions.Normal(loc=2, scale=.5)
        z = z_dist.sample((batch_size, 1))
        return z

    if data == "2spirals-8gaussians":
        data1 = inf_train_gen("2spirals", rng=rng, batch_size=batch_size)
        data2 = inf_train_gen("8gaussians", rng=rng, batch_size=batch_size)
        return np.concatenate([data1, data2], axis=1)

    if data == "4-2spirals-8gaussians":
        data1 = inf_train_gen("2spirals", rng=rng, batch_size=batch_size)
        data2 = inf_train_gen("8gaussians", rng=rng, batch_size=batch_size)
        data3 = inf_train_gen("2spirals", rng=rng, batch_size=batch_size)
        data4 = inf_train_gen("8gaussians", rng=rng, batch_size=batch_size)
        return np.concatenate([data1, data2, data3, data4], axis=1)

    if data == "8-2spirals-8gaussians":
        data1 = inf_train_gen("4-2spirals-8gaussians", rng=rng, batch_size=batch_size)
        data2 = inf_train_gen("4-2spirals-8gaussians", rng=rng, batch_size=batch_size)
        return np.concatenate([data1, data2], axis=1)

    if data == "3-MIX_DEP":
        data1 = inf_train_gen("2spirals", rng=rng, batch_size=batch_size)
        data2 = inf_train_gen("8gaussians", rng=rng, batch_size=batch_size)
        data3 = inf_train_gen("circles", rng=rng, batch_size=batch_size)
        p = rng.randint(0, 2, batch_size).reshape(-1, 1)
        p = np.concatenate([p, p], axis=1)
        dist = 2*((data1[:, 0]**2 + data1[:, 1]**2)**.5).reshape(-1, 1) - .5
        data3 = np.concatenate([dist, dist], axis=1) * data2*.5 + rng.randn(batch_size, 2) * .1
        return np.concatenate([data1, data2, data3], axis=1).astype(data2.dtype)

    if data == "7-pyramid":
        data1 = inf_train_gen("circles", rng=rng, batch_size=batch_size)
        data2 = inf_train_gen("8gaussians", rng=rng, batch_size=batch_size)
        data3 = rng.randn(batch_size, 1) + data1.max(1).reshape(-1, 1)
        data4 = rng.randn(batch_size, 1) + data2.min(1).reshape(-1, 1)
        p = rng.randint(0, 2, batch_size).reshape(-1, 1)
        data5 = p * np.sin(data3 + data4) + (1 - p) * np.cos(data3 + data4) + rng.randn(batch_size, 1)
        return np.concatenate([data1, data2, data3, data4, data5], axis=1).astype(data2.dtype)

    if data == "8-MIX":
        data1 = inf_train_gen("2spirals", rng=rng, batch_size=batch_size)
        data2 = inf_train_gen("8gaussians", rng=rng, batch_size=batch_size)
        data3 = inf_train_gen("swissroll", rng=rng, batch_size=batch_size)
        data4 = inf_train_gen("circles", rng=rng, batch_size=batch_size)
        data8 = inf_train_gen("moons", rng=rng, batch_size=batch_size)
        data6 = inf_train_gen("pinwheel", rng=rng, batch_size=batch_size)
        data7 = inf_train_gen("checkerboard", rng=rng, batch_size=batch_size)
        data5 = inf_train_gen("line", rng=rng, batch_size=batch_size)
        std = np.array([1.604934 , 1.584863 , 2.0310535, 2.0305095, 1.337718 , 1.4043778,  1.6944685, 1.6935346,
                        1.7434783, 1.0092416, 1.4860426, 1.485661 , 2.3067558, 2.311637 , 1.4430547, 1.4430547], dtype=np.float32)
        data = np.concatenate([data1, data2, data3, data4, data5, data6, data7, data8], axis=1)

        return data/std
    if data == "7-MIX":
        data1 = inf_train_gen("2spirals", rng=rng, batch_size=batch_size)
        data2 = inf_train_gen("8gaussians", rng=rng, batch_size=batch_size)
        data3 = inf_train_gen("swissroll", rng=rng, batch_size=batch_size)
        data4 = inf_train_gen("circles", rng=rng, batch_size=batch_size)
        data5 = inf_train_gen("moons", rng=rng, batch_size=batch_size)
        data6 = inf_train_gen("pinwheel", rng=rng, batch_size=batch_size)
        data7 = inf_train_gen("checkerboard", rng=rng, batch_size=batch_size)
        std = np.array([1.604934 , 1.584863 , 2.0310535, 2.0305095, 1.337718 , 1.4043778,  1.6944685, 1.6935346,
                        1.7434783, 1.0092416, 1.4860426, 1.485661 , 2.3067558, 2.311637], dtype=np.float32)
        data = np.concatenate([data1, data2, data3, data4, data5, data6, data7], axis=1)

        return data/std

    if data == "swissroll":
        data = sklearn.datasets.make_swiss_roll(n_samples=batch_size, noise=1.0)[0]
        data = data.astype("float32")[:, [0, 2]]
        data /= 5
        return data

    elif data == "circles":
        data = sklearn.datasets.make_circles(n_samples=batch_size, factor=.5, noise=0.08)[0]
        data = data.astype("float32")
        data *= 3
        return data

    elif data == "moons":
        data = sklearn.datasets.make_moons(n_samples=batch_size, noise=0.1)[0]
        data = data.astype("float32")
        data = data * 2 + np.array([-1, -0.2])
        data = data.astype("float32")
        return data

    elif data == "8gaussians":
        scale = 4.
        centers = [(1, 0), (-1, 0), (0, 1), (0, -1), (1. / np.sqrt(2), 1. / np.sqrt(2)),
                   (1. / np.sqrt(2), -1. / np.sqrt(2)), (-1. / np.sqrt(2),
                                                         1. / np.sqrt(2)), (-1. / np.sqrt(2), -1. / np.sqrt(2))]
        centers = [(scale * x, scale * y) for x, y in centers]

        dataset = []
        for i in range(batch_size):
            point = rng.randn(2) * 0.5
            idx = rng.randint(8)
            center = centers[idx]
            point[0] += center[0]
            point[1] += center[1]
            dataset.append(point)
        dataset = np.array(dataset, dtype="float32")
        dataset /= 1.414
        return dataset

    elif data == "2gaussians":
        scale = 4.
        centers = [(.5, -.5), (-.5, .5)]
        centers = [(scale * x, scale * y) for x, y in centers]

        dataset = []
        for i in range(batch_size):
            point = rng.randn(2) * .75
            idx = rng.randint(2)
            center = centers[idx]
            point[0] += center[0]
            point[1] += center[1]
            dataset.append(point)
        dataset = np.array(dataset, dtype="float32")
        #dataset /= 1.414
        return dataset

    elif data == "4gaussians":
        scale = 4.
        centers = [(.5, -.5), (-.5, .5), (.5, .5), (-.5, -.5)]
        centers = [(scale * x, scale * y) for x, y in centers]

        dataset = []
        for i in range(batch_size):
            point = rng.randn(2) * .75
            idx = rng.randint(4)
            center = centers[idx]
            point[0] += center[0]
            point[1] += center[1]
            dataset.append(point)
        dataset = np.array(dataset, dtype="float32")
        # dataset /= 1.414
        return dataset

    elif data == "2igaussians":
        scale = 4.
        centers = [(.5, 0.), (-.5, .0)]
        centers = [(scale * x, scale * y) for x, y in centers]

        dataset = []
        for i in range(batch_size):
            point = rng.randn(2) * .75
            idx = rng.randint(2)
            center = centers[idx]
            point[0] += center[0]
            point[1] += center[1]
            dataset.append(point)
        dataset = np.array(dataset, dtype="float32")
        # dataset /= 1.414
        return dataset

    elif data == "conditionnal8gaussians":
        scale = 4.
        centers = [(1, 0), (-1, 0), (0, 1), (0, -1), (1. / np.sqrt(2), 1. / np.sqrt(2)),
                   (1. / np.sqrt(2), -1. / np.sqrt(2)), (-1. / np.sqrt(2),
                                                         1. / np.sqrt(2)), (-1. / np.sqrt(2), -1. / np.sqrt(2))]
        centers = [(scale * x, scale * y) for x, y in centers]

        dataset = []
        context = np.zeros((batch_size, 8))
        for i in range(batch_size):
            point = rng.randn(2) * 0.5
            idx = rng.randint(8)
            context[i, idx] = 1
            center = centers[idx]
            point[0] += center[0]
            point[1] += center[1]
            dataset.append(point)
        dataset = np.array(dataset, dtype="float32")
        dataset /= 1.414
        return dataset, context

    elif data == "pinwheel":
        radial_std = 0.3
        tangential_std = 0.1
        num_classes = 5
        num_per_class = batch_size // 5
        rate = 0.25
        rads = np.linspace(0, 2 * np.pi, num_classes, endpoint=False)

        features = rng.randn(num_classes*num_per_class, 2) \
            * np.array([radial_std, tangential_std])
        features[:, 0] += 1.
        labels = np.repeat(np.arange(num_classes), num_per_class)

        angles = rads[labels] + rate * np.exp(features[:, 0])
        rotations = np.stack([np.cos(angles), -np.sin(angles), np.sin(angles), np.cos(angles)])
        rotations = np.reshape(rotations.T, (-1, 2, 2))

        return 2 * rng.permutation(np.einsum("ti,tij->tj", features, rotations)).astype("float32")

    elif data == "2spirals":
        n = np.sqrt(rng.rand(batch_size // 2, 1)) * 540 * (2 * np.pi) / 360
        d1x = -np.cos(n) * n + rng.rand(batch_size // 2, 1) * 0.5
        d1y = np.sin(n) * n + rng.rand(batch_size // 2, 1) * 0.5
        x = np.vstack((np.hstack((d1x, d1y)), np.hstack((-d1x, -d1y)))) / 3
        x += rng.randn(*x.shape) * 0.1
        return x.astype("float32")

    elif data == "checkerboard":
        x1 = np.random.rand(batch_size) * 4 - 2
        x2_ = np.random.rand(batch_size) - np.random.randint(0, 2, batch_size) * 2
        x2 = x2_ + (np.floor(x1) % 2)
        return np.concatenate([x1[:, None], x2[:, None]], 1).astype("float32") * 2

    elif data == "line":
        x = rng.rand(batch_size)
        #x = np.arange(0., 1., 1/batch_size)
        x = x * 5 - 2.5
        y = x #- x + rng.rand(batch_size)
        return np.stack((x, y), 1).astype("float32")
    elif data == "line-noisy":
        x = rng.rand(batch_size)
        x = x * 5 - 2.5
        y = x + rng.randn(batch_size)
        return np.stack((x, y), 1).astype("float32")
    elif data == "cos":
        x = rng.rand(batch_size) * 6 - 3
        y = np.sin(x*5) * 2.5 + np.random.randn(batch_size) * .3
        return np.stack((x, y), 1).astype("float32")
    elif data == "joint_gaussian":
        x2 = torch.distributions.Normal(0., 4.).sample((batch_size, 1))
        x1 = torch.distributions.Normal(0., 1.).sample((batch_size, 1)) + (x2**2)/4

        return torch.cat((x1, x2), 1)
    elif data == "woodStructural":
        log_lik = 0.
        z0 = torch.distributions.log_normal.LogNormal(5., 1., validate_args=None).sample((batch_size, 1))
        log_lik += torch.distributions.log_normal.LogNormal(5., 1., validate_args=None).log_prob(z0)

        z1 = torch.distributions.log_normal.LogNormal(-2., 1., validate_args=None).sample((batch_size, 1))
        log_lik += torch.distributions.log_normal.LogNormal(-2., 1., validate_args=None).log_prob(z1)

        z2 = torch.distributions.Normal(loc=torch.tanh(z0 + z1 - 2.8), scale=.1).sample()
        log_lik += torch.distributions.Normal(loc=torch.tanh(z0 + z1 - 2.8), scale=.1).log_prob(z2)

        z3 = torch.distributions.Normal(loc=z0 * z1, scale=.1).sample()
        log_lik += torch.distributions.Normal(loc=z0 * z1, scale=.1).log_prob(z3)

        z4 = torch.distributions.Normal(loc=7, scale=2.).sample((batch_size, 1))
        log_lik += torch.distributions.Normal(loc=7, scale=2.).log_prob(z4)

        z5 = torch.distributions.Normal(loc=torch.tanh(z3 + z4), scale=.1).sample()
        log_lik += torch.distributions.Normal(loc=torch.tanh(z3 + z4), scale=.1).log_prob(z5)

        x0 = torch.distributions.Normal(loc=z3, scale=.1).sample()
        log_lik += torch.distributions.Normal(loc=z3, scale=.1).log_prob(x0)

        x1 = torch.distributions.Normal(loc=z5, scale=.1).sample()
        log_lik += torch.distributions.Normal(loc=z5, scale=.1).log_prob(x1)

        z = torch.cat((z0, z1, z2, z3, z4, z5, x0, x1), 1)
        #print(log_lik.mean() + torch.log(z.std(0)).sum())
        #exit()
        return z
    else:
        return inf_train_gen("8gaussians", rng, batch_size)

def getA(toy):
    if toy == "8-MIX":
        A = torch.zeros(16, 16)
        for i in range(8):
            A[i * 2, 2*i+1] = 1
    elif toy == "7-MIX":
        A = torch.zeros(14, 14)
        for i in range(7):
            A[i * 2, 2 * i + 1] = 1

    elif toy == "3-MIX_DEP":
        A = torch.zeros(6, 6)
        for i in range(2):
            A[i * 2, 2 * i + 1] = 1
        A[4, 0] = 1
        A[4, 2] = 1
        A[4, 1] = 1
        A[4, 3] = 1
        A[5, 0] = 1
        A[5, 2] = 1
        A[5, 1] = 1
        A[5, 3] = 1
    elif toy == "woodStructural":
        A = torch.zeros(8, 8)
        A[2, 0] = 1
        A[2, 1] = 1
        A[3, 0] = 1
        A[3, 1] = 1
        A[5, 3] = 1
        A[5, 4] = 1
        A[6, 3] = 1
        A[7, 5] = 1
    elif toy == "7-pyramid":
        A = torch.zeros(7, 7)
        A[0, 1] = 1
        A[2, 3] = 1
        A[4, 0] = 1
        A[4, 1] = 1
        A[5, 2] = 1
        A[5, 3] = 1
        A[6, 4] = 1
        A[6, 5] = 1
    elif toy == '1Normal':
        A = torch.zeros(1, 1)
    else:
        return None
    return A