import xlrd
import sys
import xlwt
from xlwt import easyxf
from pprint import pprint
import logging
logging.basicConfig(filename='firewall_2.log', level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

def getSubnetIPs(book = "", name = ""):
    xl_sheet = book.sheet_by_name(name)
    logging.info('Processing BDA sheet - {0}'.format(name))
    req_col_name_01 = "INTERFACE CATEGORY"
    req_col_name_02 = "Subnet Details"
    req_columns_found = False
    req_column_index_01 = 0
    req_column_index_02 = 0
    sheet_ips_list = []

    num_cols = xl_sheet.ncols   # Number of columns
    num_rows = xl_sheet.nrows   # Number of rows
    for row_idx in range(0, num_rows):    # Iterate through rows
        row_data = []
        for col_idx in range(0, num_cols):  # Iterate through columns
            cell_obj = xl_sheet.cell(row_idx, col_idx)  # Get cell object by row, col
            if cell_obj:
                row_data.append(str(cell_obj.value))
#                 pprint(row_data)
            else:
                row_data.append(str("NA"))
        if (req_col_name_01 in row_data) and (req_col_name_02 in row_data):     # Checking for the required value in the
            #print("Found.. in row - {0}".format(row_idx))
            req_column_index_01 = row_data.index(req_col_name_01)
            req_column_index_02 = row_data.index(req_col_name_02)
            req_columns_found = True
            continue
            #process further
        else:
            pass
            #print("Not Found.. in row - {0}".format(row_idx))

        if req_columns_found and req_column_index_01 and req_column_index_02:
            if row_data[req_column_index_02] and row_data[req_column_index_01]:
                if row_data[req_column_index_01].lower().startswith('customer'):
                    sheet_ips_list.append(row_data[req_column_index_02])

    return sheet_ips_list

def get_fe_be_mgmt_IPs(book = "", name = ""):
    xl_sheet = book.sheet_by_name(name)
    logging.info('Processing exdata sheet - {0}'.format(name))
    req_col_name_01 = "INTERFACE CATEGORY"
    req_col_name_02 = "Subnet Details"
    req_columns_found = False
    req_column_index_01 = 0
    req_column_index_02 = 0
    fe_ips_list = []
    be_ips_list = []
    mgmt_ips_list = []
    bdcs_mgmt_ip = []

    num_cols = xl_sheet.ncols   # Number of columns
    for row_idx in range(0, xl_sheet.nrows):    # Iterate through rows
        #print ('Row: %s' % row_idx)
        row_data = []
        for col_idx in range(0, num_cols):  # Iterate through columns
            cell_obj = xl_sheet.cell(row_idx, col_idx)  # Get cell object by row, col
#             pprint(cell_obj)
            if cell_obj:
                row_data.append(str(cell_obj.value))
#                 pprint(row_data)
            else:
                row_data.append(str("NA"))
        #print ('Row: [%s] data: [%s]' % (row_idx, ", ".join(row_data)))
        if (req_col_name_01 in row_data) and (req_col_name_02 in row_data):
            #print("Found.. in row - {0}".format(row_idx))
            req_column_index_01 = row_data.index(req_col_name_01)
            req_column_index_02 = row_data.index(req_col_name_02)
            req_columns_found = True
            continue
            #process further
        else:
            pass
            #print("Not Found.. in row - {0}".format(row_idx))

        if req_columns_found and req_column_index_01 and req_column_index_02:
            if row_data[req_column_index_02] and row_data[req_column_index_01]:
                if row_data[req_column_index_01].lower().startswith('customer') and row_data[req_column_index_01].lower().endswith('fe'):
                    fe_ips_list.append(row_data[req_column_index_02])
                elif row_data[req_column_index_01].lower().startswith('customer') and row_data[req_column_index_01].lower().endswith('be'):
                    be_ips_list.append(row_data[req_column_index_02])
                elif row_data[req_column_index_01].lower().startswith('customer') and row_data[req_column_index_01].lower().endswith('mgmt'):
                    mgmt_ips_list.append(row_data[req_column_index_02])
                elif row_data[req_column_index_01] == "BDCS Management":
                    bdcs_mgmt_ip.append(row_data[req_column_index_02])

    return fe_ips_list, be_ips_list, mgmt_ips_list, bdcs_mgmt_ip

def process_all_sheets():
    all_info = []
    for sheet in sheet_name:
        if sheet.endswith("BDA"):
            try:
                sheet_data = getSubnetIPs(book = book, name = sheet)
                all_info += sheet_data
            except Exception as e:
                logging.critical("Exception > process_all_sheets() for BDA -> {0}".format(e))
    fe_info = []
    be_info = []
    mgmt_info = []
    bdcs_mgmt_info = []
    for sheet in sheet_name:
        if sheet.lower().endswith("exdata"):
            #sheet_data = getSubnetIPs(book = book, name = sheet)
            #print(sheet)
            try:
                fe_data, be_data, mgmt_data, bdcs_mgmt = get_fe_be_mgmt_IPs(book = book, name = sheet)
                fe_info += fe_data
                be_info += be_data
                mgmt_info += mgmt_data
                bdcs_mgmt_info += bdcs_mgmt
            except Exception as e:
                logging.critical("Exception > process_all_sheets() for exdata -> {0}".format(e))
    #     else:
    #         logging.info("There is no sheets with the suffix exdata")

    EM_SM = ["nlcl423ru05.nldc1.oraclecloud.com","nlcl423ru06.nldc1.oraclecloud.com","nlcl423ru07.nldc1.oraclecloud.com","nlcl423ru11.nldc1.oraclecloud.com"]
    US_SM = ["chr302ru27.usdc2.oraclecloud.com","chr302ru26.usdc2.oraclecloud.com","chr302ru25.usdc2.oraclecloud.com","chr302ru22.usdc2.oraclecloud.com"]

    if all_info and bdcs_mgmt_info and mgmt_info and be_info and fe_info:
        logging.info("Generating output xlsx ...")
        style1=xlwt.easyxf('font: name Calibri, color-index black, bold on, height 280; pattern: pattern solid, fore_colour yellow')
        style2=xlwt.easyxf('align: wrap yes; font: name Calibri, height 220')
        workbook=xlwt.Workbook(encoding="utf-8")
        sheet1=workbook.add_sheet("firewall Sheet")
        sheet1.col(0).width = 8000
        sheet1.col(1).width = 7000
        sheet1.col(2).width = 3000
        sheet1.col(3).width = 7000
        sheet1.col(4).width = 7000
        sheet1.col(5).width = 9000
        sheet1.write(0,0, "Source IP", style1)
        sheet1.write(0,1, "IP Type", style1)
        sheet1.write(0,2, "Protocol", style1)
        sheet1.write(0,3, "Destination IP", style1)
        sheet1.write(0,4, "Port", style1)
        sheet1.write(0,5, "Description", style1)
        sheet1.write(1,0, "Internet/Customer", style2)
        sheet1.write(1,1, "Internernet source IP", style2)
        sheet1.write(1,2, "SSH", style2)
        sheet1.write(1,3, "\n".join(all_info), style2)
        sheet1.write(1,4, "22", style2)
        sheet1.write(1,5, "Customer ssh to bda vm", style2)
        sheet1.write(2,0, "Internet/Customer", style2)
        sheet1.write(2,1, "Internernet source IP", style2)
        sheet1.write(2,2, "SSH", style2)
        sheet1.write(2,3, "\n".join(fe_info), style2)
        sheet1.write(2,4, "22", style2)
        sheet1.write(2,5, "customer ssh access to Exadata VMs (client Interface)", style2)
        sheet1.write(3,0, "Internet/Customer", style2)
        sheet1.write(3,1, "Internernet source IP", style2)
        sheet1.write(3,2, "SSH", style2)
        sheet1.write(3,3, "\n".join(be_info), style2)
        sheet1.write(3,4, "22", style2)
        sheet1.write(3,5, "customer ssh access to Exadata VM (backup Interface)", style2)
        sheet1.write(4,0, "Internet/Customer", style2)
        sheet1.write(4,1, "Internernet source IP", style2)
        sheet1.write(4,2, "SSH", style2)
        sheet1.write(4,3, "\n".join(mgmt_info), style2)
        sheet1.write(4,4, "22", style2)
        sheet1.write(4,5, "customer ssh access to Exadata VM (mgmt Interface)", style2)
        sheet1.write(5,0, "Internet/Customer", style2)
        sheet1.write(5,1, "Internernet source IP", style2)
        sheet1.write(5,2, "HTTPS", style2)
        sheet1.write(5,3, "\n".join(all_info), style2)
        sheet1.write(5,4, "7183", style2)
        sheet1.write(5,5, "Customer access to Cloudera Manager UI Inside BDA Vms", style2)
        sheet1.write(6,0, "Internet/Customer", style2)
        sheet1.write(6,1, "Internernet source IP", style2)
        sheet1.write(6,2, "HTTPS", style2)
        sheet1.write(6,3, "\n".join(all_info), style2)
        sheet1.write(6,4, "8888", style2)
        sheet1.write(6,5, "customer access to Hue UI inside BDA VMs", style2)
        sheet1.write(7,0, "\n".join(all_info), style2)
        sheet1.write(7,1, "---", style2)
        sheet1.write(7,2, "---", style2)
        sheet1.write(7,3, "ALL", style2)
        sheet1.write(7,4, "ALL", style2)
        sheet1.write(7,5, "Unlimted access from Customer VM to Internet", style2)
        sheet1.write(8,0, "\n".join(EM_SM), style2)
        sheet1.write(8,1, "---", style2)
        sheet1.write(8,2, "SSH", style2)
        sheet1.write(8,3, "\n".join(all_info), style2)
        sheet1.write(8,4, "22", style2)
        sheet1.write(8,5, "SM ssh access to bdcs mgmt network", style2)
        sheet1.write(9,0, "\n".join(EM_SM), style2)
        sheet1.write(9,1, "---", style2)
        sheet1.write(9,2, "SSH", style2)
        sheet1.write(9,3, "\n".join(fe_info), style2)
        sheet1.write(9,4, "22", style2)
        sheet1.write(9,5, "SM MT1/MT2 ssh access to Exadata Customer VM", style2)

        workbook.save("Firewall.xls")
        logging.info("Fire wall sheet Generated Successfully...")

if len(sys.argv) == 2 and ( str(sys.argv[1]).endswith("xlsx") or str(sys.argv[1]).endswith("xls") ):
    try:
        book = xlrd.open_workbook(sys.argv[1])
        sheet_name = book.sheet_names()
        logging.info("Processing  started with the given input file" + "---->" + "'" + str(sys.argv[1]) + "'")
        logging.info("List of input sheets - {0}".format(sheet_name))
        logging.info("Number of sheets found in the file" + ":" + str(book.nsheets))

    except Exception as k:
        logging.critical("exception occurred - {0}".format(k))
        sys.exit()
else:
    logging.info("Execution format -> # python parse_xlsx.py <xlsx / xls file>")
    logging.info("Exiting the script input file format is not xlsx")
    sys.exit()

if __name__ == '__main__':
    try:
        process_all_sheets()
    except Exception as emsg:
        logging.critical("exception occurred - {0}".format(emsg))
# else:
#     logging.info("Execution format -> # python parse_xlsx.py <xlsx / xls file>")


class Solution:
    # @param board, a 9x9 2D array
    # Solve the Sudoku by modifying the input board in-place.
    # Do not return any value.
    digits = set(['1','2','3','4','5','6','7','8','9'])
    minedge  = {
        0: 0,
        1: 0,
        2: 0,
        3: 3,
        4: 3,
        5: 3,
        6: 6,
        7: 6,
        8: 6
    }

    maxedge = {
        0:2,
        1:2,
        2:2,
        3:5,
        4:5,
        5:5,
        6:8,
        7:8,
        8:8
    }

    board = []

    def getCandidates(self, pos):
        # return candidate values
        a = self.getCandidatesX(pos)
        # #print a
        b = self.getCandidatesY(pos)
        # #print b
        c = self.getCandidatesSquare(pos)
        # #print c
        sol = a.intersection(b)
        sol = sol.intersection(c)
        return sol

    def getCandidatesX(self, pos):
        #return candidates
        tagged = [w for w in self.board[pos['y']] if w != '.']
        # #print tagged
        return self.digits.difference(set(tagged))

    def getCandidatesY(self, pos):
        #return candidates down line
        tagged = [w[pos['x']] for w in self.board if w[pos['x']] != '.']
        return self.digits.difference(set(tagged))

    def getCandidatesSquare(self, pos):
        #return candidates in square
        mylist = []
        for x in range(self.minedge[pos['x']], self.maxedge[pos['x']] + 1):
            for y in range(self.minedge[pos['y']], self.maxedge[pos['y']] + 1):
                if self.board[y][x] != '.':
                    mylist.append(self.board[y][x])
        return self.digits.difference(set(mylist))

    def getNextCandidate(self, pos):
        for y in range(pos['y'], 9):
            if (y == pos['y']):
                for x in range(pos['x'] + 1, 9):
                    if self.board[y][x] == '.':
                        return {'x':x, 'y':y}
            else:
                for x in range(0, 9):
                    if self.board[y][x] == '.':
                        return {'x':x, 'y':y}
        return None

    def solveSudoku(self, board_):
        self.board = board_;
        positions = [] #stack for positions
        if (self.board[0][0]) == '.':
            pos = {'x' : 0, 'y': 0}
        else:
            pos = self.getNextCandidate({'x':0, 'y':0})
        self.solveSudoku_(pos)

    def solveSudoku_(self, pos):
        candidates = self.getCandidates(pos)
        if (len(candidates) == 0):
            return
        nextpos = self.getNextCandidate(pos)

        for candidate in candidates:
            temp = list(self.board[pos['y']])
            temp[pos['x']] = candidate
            self.board[pos['y']]  =  "".join(temp)
            if (nextpos):
                a = self.solveSudoku_(nextpos)
                if (a is True):
                    return True
                else:
                    temp = list(self.board[pos['y']])
                    temp[pos['x']] = '.'
                    self.board[pos['y']]  =  "".join(temp)
        if (nextpos is None):
            return True


#! /usr/bin/env python3
# Base Server - Chapter 3 - basicserver.py
# 测试：telnet localhost 51423

import socket

host = ''    # Bind to all interfaces
port = 51423

s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR,1)
# setsockopt(level, opt, value); level通常是SOL_SOCKET.
s.bind((host,port))
print("Waiting for connections...")
s.listen(1)

# 通常情况下，无限循环是不好的，因为会耗尽系统的CPU资源。
# 这里不同，调用accept()的时候，只在有一个客户端连接后才返回。
# 如果没有客户端连接，程序会停止。
while 1:
    clientsock, clientaddr = s.accept()
    print("Got connection from ", clientsock.getpeername())
    clientsock.close()


import tensorflow as tf
from corner_pooling import TopPool,BottomPool, LeftPool, RightPool

class Corner:

    def __init__(self,is_train):

        self.is_train = is_train
        self.n_deep = 5
        self.n_dims = [256, 256, 384, 384, 384, 512]
        self.n_ksize = [2, 2, 2, 2, 2, 4]

    def head(self,img):

        with tf.variable_scope('head'):

            net = tf.layers.conv2d(img,128,7,2,padding='same',name='conv1')
            net = tf.layers.batch_normalization(net,trainable=self.is_train,name='bn1')
            net = self.residual(net,256,strides=2,name='residual_start')  #128,128,256

            return net

    def conv_bn_relu(self,x,filters,k_size = 3,strides = 1,is_relu = True,is_bn = True,name = 'conv_bn_re'):

        with tf.variable_scope(name):

            net = tf.layers.conv2d(x,filters,k_size,strides=strides,padding='same',name='res_conv')
            if is_bn:
                net = tf.layers.batch_normalization(net,trainable=self.is_train,name='res_bn')
            if is_relu:
                net = tf.nn.relu(net)

            return net

    def residual(self,x,filters,k_size = 3,strides = 1,name = 'residual'):
        with tf.variable_scope(name):
            net = self.conv_bn_relu(x,filters,k_size=k_size,strides=strides,name = 'up_1')
            net = self.conv_bn_relu(net,filters,k_size=k_size,is_relu=False,name = 'up_2')

            #skip  #1x1 conv
            skip = self.conv_bn_relu(x,filters = filters,k_size= 1,strides=strides,is_relu=False,name='low')

            out = tf.nn.relu(tf.add(skip,net))

            return out

    def res_block(self,x,filters,k_size = 3,iter = 1,name = 'res_block'):

        with tf.variable_scope(name):

            net = self.residual(x,filters,k_size,name='residual_0')
            for i in range(1,iter):
                net = self.residual(net,filters,k_size,name='residual_%d' % i)
            return net

    def hourglass_net(self,x,deep,k_size,filters,name = 'hourglass_iter'):

        with tf.variable_scope(name):
            k1 = k_size[0]
            k2 = k_size[1]
            curr_filter = filters[0]
            next_filter = filters[1]

            up1 = self.res_block(x,curr_filter,k_size=k1,name='up_1')
            net = tf.layers.max_pooling2d(x,(2,2),(2,2),padding='same')

            down_1 = self.res_block(net,next_filter,k1,name='down_1')
            if deep > 1:
                down_2 = self.hourglass_net(down_1,deep - 1,k_size=k_size[1:],filters=filters[1:],name='hourglass_%d' % (deep - 1))
            else:
                down_2 = self.res_block(down_1,filters=next_filter,k_size=k2,name='down_2')
            down_3 = self.res_block(down_2,curr_filter,k_size=k1,name='down_3')

            up2 = tf.image.resize_nearest_neighbor(down_3,tf.shape(down_3)[1:3] * 2,name='up_2')

            out = tf.add(up1,up2)

            return out

    def corner_pool(self,x,filters,k_size = 3,name = 'corner_pooling'):
        with tf.variable_scope(name):

            with tf.variable_scope('top_left'):

                top_ponits = self.conv_bn_relu(x,128,name = 'top')
                top_pool = TopPool(top_ponits)

                left_points = self.conv_bn_relu(x,128,name='left')
                left_pool = LeftPool(left_points)

                top_left = tf.add(top_pool,left_pool)
                top_left = self.conv_bn_relu(top_left,filters=filters,is_relu=False,name='_top_left')
                _top_left = self.conv_bn_relu(x,filters=filters,is_relu=False,k_size=1,name='skip')
                top_left = tf.add(top_left,_top_left)
                top_left = tf.nn.relu(top_left)
                top_left = self.conv_bn_relu(top_left,filters = filters,name='top_left')

            with tf.variable_scope('bottom_right'):

                bottom_poinst = self.conv_bn_relu(x,128,name='bottom')
                bottom_pool = BottomPool(bottom_poinst)

                right_points = self.conv_bn_relu(x,128,name='right')
                right_pool = RightPool(right_points)

                bottom_right = tf.add(bottom_pool,right_pool)
                bottom_right = self.conv_bn_relu(bottom_right,filters = filters,is_relu=False,name='_bottom_right')
                _bottom_right = self.conv_bn_relu(x,filters=filters,is_relu=False,k_size=1,name='skip')

                bottom_right = tf.add(bottom_right,_bottom_right)
                bottom_right = tf.nn.relu(bottom_right)
                bottom_right = self.conv_bn_relu(bottom_right,filters=filters,name='bottom_right')

                return top_left,bottom_right

    def heatmap(self,x,filters,name = 'heatmaps'):
        #channel:80
        with tf.variable_scope(name):
            net = self.conv_bn_relu(x,filters=filters[0],is_bn=False,name = 'conv1')
            net = tf.layers.conv2d(net,filters=filters[1],kernel_size=1,name='conv2')

            return net

    def embedding(self,x,filters,name = 'embedding'):

        #channel:1
        with tf.variable_scope(name):
            net = self.conv_bn_relu(x,filters[0],is_bn=False,name='conv1')
            net = tf.layers.conv2d(net,filters[1],1,name='conv2')

            return net

    def offset(self,x,filters,name = 'offset'):
        #channel : 2
        with tf.variable_scope(name):
            net = self.conv_bn_relu(x, filters[0], is_bn=False, name='conv1')
            net = tf.layers.conv2d(net, filters[1], 1, name='conv2')

            return net

    def inter(self,x1,x2,filters,name = 'inter'):
        with tf.variable_scope(name):
            net1 = self.conv_bn_relu(x1,filters=filters,is_relu=False,k_size=1,name='branch_start')
            net2 = self.conv_bn_relu(x2,filters=filters,is_relu=False,k_size=1,name='branch_hourglass2')
            net = tf.nn.relu(tf.add(net1,net2))
            net = self.residual(net,filters=filters)
            return net

    def main_net(self,img,
                 gt_tl_embedding = None,
                 gt_br_embedding = None):
        with tf.variable_scope('corner_net'):

            start_net = self.head(img)  #[,128,128,256]
            print(start_net)
            with tf.variable_scope('hourglass_1'):

                #[,128,128,256]
                hg1 = self.hourglass_net(start_net,deep=self.n_deep,k_size=self.n_ksize,filters=self.n_dims,name='hg_1')

                hinge1 = self.conv_bn_relu(hg1,filters=256,name='hinge1')
                top_left,bottom_right = self.corner_pool(hinge1,256)

                #top_left的heatmap和embedding
                tl_heatmap = self.heatmap(top_left,[256,80],name='top_left_heatmap1')
                tl_embedding = self.embedding(top_left,[256,1],name='top_left_embedding1')
                tl_offset = self.offset(top_left,[256,2],name='top_left_offset1')

                #bottom_right
                br_heatmap = self.heatmap(bottom_right, [256, 80], name='bottom_right_heatmap1')
                br_embedding = self.embedding(bottom_right, [256, 1], name='bottom_right_embedding1')
                br_offset = self.offset(bottom_right, [256, 2], name='bottom_right_offset1')

            with tf.variable_scope('hourglass_2'):
                inter = self.inter(start_net,hinge1,256)
                hg2 = self.hourglass_net(inter,deep=self.n_deep,k_size=self.n_ksize,filters=self.n_dims,name='hg_2')
                hinge2 = self.conv_bn_relu(hg2,filters=256,name='hinge2')
                top_left2,bottom_right2 = self.corner_pool(hinge2,256)

                tl_heatmap2 = self.heatmap(top_left2,[256,80],name='top_left_heatmap2')
                tl_embedding2 = self.embedding(top_left2, [256, 1], name='top_left_embedding2')
                tl_offset2 = self.offset(top_left2, [256, 2], name='top_left_offset2')
                br_heatmap2 = self.heatmap(bottom_right2, [256, 80], name='bottom_right_heatmap2')
                br_embedding2 = self.embedding(bottom_right2, [256, 1], name='bottom_right_embedding2')
                br_offset2 = self.offset(bottom_right2, [256, 2], name='bottom_right_offset2')

        if self.is_train:
            return [tl_heatmap,br_heatmap,tl_embedding,br_embedding,tl_offset,br_offset,tl_heatmap2,br_heatmap2,tl_embedding2,br_embedding2,tl_offset2,br_offset2]
        return [tl_heatmap2,br_heatmap2,tl_embedding2,br_embedding2,tl_offset2,br_offset2]

if __name__ == '__main__':
    x = tf.placeholder(dtype=tf.float32,shape=[None,511,511,3])
    test1 = Corner(True)
    test_res = test1.main_net(x)
    print(test_res)


from django.test import TestCase
from django.urls import reverse

from ..urls import app_name

USERNAME = "username"
SLUG = "slug"
ID = 1

ROUTES = (
    ("index", "/", []),
    ("post_edit", f"/posts/{ID}/edit/", [ID]),
    ("post_create", "/create/", []),
    ("post_detail", f"/posts/{ID}/", [ID]),
    ("profile", f"/profile/{USERNAME}/", [USERNAME]),
    ("add_comment", f"/posts/{ID}/comment/", [ID]),
    ("group_list", f"/group/{SLUG}/", [SLUG]),
    ("follow_index", "/follow/", []),
    ("profile_follow", f"/profile/{USERNAME}/follow/", [USERNAME]),
    ("profile_unfollow", f"/profile/{USERNAME}/unfollow/", [USERNAME]),
)

class RoutesTests(TestCase):
    def test_urls_correctly_resolved(self):
        for key, url, args in ROUTES:
            with self.subTest(key=key):
                self.assertEqual(reverse(f"{app_name}:{key}", args=args), url)


#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import datetime

import mock
from oslo_utils import timeutils
import six

from heat.common import exception
from heat.common import template_format
from heat.engine import resource
from heat.engine import scheduler
from heat.tests.autoscaling import inline_templates
from heat.tests import common
from heat.tests import utils

as_template = inline_templates.as_template
as_params = inline_templates.as_params

class TestAutoScalingPolicy(common.HeatTestCase):
    def setUp(self):
        super(TestAutoScalingPolicy, self).setUp()

    def create_scaling_policy(self, t, stack, resource_name):
        rsrc = stack[resource_name]
        self.assertIsNone(rsrc.validate())
        scheduler.TaskRunner(rsrc.create)()
        self.assertEqual((rsrc.CREATE, rsrc.COMPLETE), rsrc.state)
        return rsrc

    def test_validate_scaling_policy_ok(self):
        t = template_format.parse(inline_templates.as_template)
        t['Resources']['WebServerScaleUpPolicy']['Properties'][
            'ScalingAdjustment'] = 33
        t['Resources']['WebServerScaleUpPolicy']['Properties'][
            'AdjustmentType'] = 'PercentChangeInCapacity'
        t['Resources']['WebServerScaleUpPolicy']['Properties'][
            'MinAdjustmentStep'] = 2
        stack = utils.parse_stack(t, params=as_params)
        self.policy = stack['WebServerScaleUpPolicy']
        self.assertIsNone(self.policy.validate())

    def test_validate_scaling_policy_error(self):
        t = template_format.parse(inline_templates.as_template)
        t['Resources']['WebServerScaleUpPolicy']['Properties'][
            'ScalingAdjustment'] = 1
        t['Resources']['WebServerScaleUpPolicy']['Properties'][
            'AdjustmentType'] = 'ChangeInCapacity'
        t['Resources']['WebServerScaleUpPolicy']['Properties'][
            'MinAdjustmentStep'] = 2
        stack = utils.parse_stack(t, params=as_params)
        self.policy = stack['WebServerScaleUpPolicy']
        ex = self.assertRaises(exception.ResourcePropertyValueDependency,
                               self.policy.validate)
        self.assertIn('MinAdjustmentStep property should only '
                      'be specified for AdjustmentType with '
                      'value PercentChangeInCapacity.', six.text_type(ex))

    def test_scaling_policy_bad_group(self):
        t = template_format.parse(inline_templates.as_template_bad_group)
        stack = utils.parse_stack(t, params=as_params)

        up_policy = self.create_scaling_policy(t, stack,
                                               'WebServerScaleUpPolicy')

        ex = self.assertRaises(exception.ResourceFailure, up_policy.signal)
        self.assertIn('Alarm WebServerScaleUpPolicy could '
                      'not find scaling group', six.text_type(ex))

    def test_scaling_policy_not_alarm_state(self):
        """If the details don't have 'alarm' then don't progress."""
        t = template_format.parse(as_template)
        stack = utils.parse_stack(t, params=as_params)
        pol = self.create_scaling_policy(t, stack, 'WebServerScaleUpPolicy')

        test = {'current': 'not_an_alarm'}
        with mock.patch.object(pol, '_cooldown_inprogress',
                               side_effect=AssertionError()) as dont_call:
            self.assertRaises(resource.NoActionRequired,
                              pol.handle_signal, details=test)
            self.assertEqual([], dont_call.call_args_list)

    def test_scaling_policy_cooldown_toosoon(self):
        """If _cooldown_inprogress() returns True don't progress."""
        t = template_format.parse(as_template)
        stack = utils.parse_stack(t, params=as_params)
        pol = self.create_scaling_policy(t, stack, 'WebServerScaleUpPolicy')
        test = {'current': 'alarm'}

        with mock.patch.object(pol.stack, 'resource_by_refid',
                               side_effect=AssertionError) as dont_call:
            with mock.patch.object(pol, '_cooldown_inprogress',
                                   return_value=True) as mock_cip:
                self.assertRaises(resource.NoActionRequired,
                                  pol.handle_signal, details=test)
                mock_cip.assert_called_once_with()
            self.assertEqual([], dont_call.call_args_list)

    def test_scaling_policy_cooldown_ok(self):
        t = template_format.parse(as_template)
        stack = utils.parse_stack(t, params=as_params)
        pol = self.create_scaling_policy(t, stack, 'WebServerScaleUpPolicy')
        test = {'current': 'alarm'}

        group = self.patchobject(pol.stack, 'resource_by_refid').return_value
        group.name = 'fluffy'
        with mock.patch.object(pol, '_cooldown_inprogress',
                               return_value=False) as mock_cip:
            pol.handle_signal(details=test)
            mock_cip.assert_called_once_with()
        group.adjust.assert_called_once_with(1, 'ChangeInCapacity', None)

class TestCooldownMixin(common.HeatTestCase):
    def setUp(self):
        super(TestCooldownMixin, self).setUp()

    def create_scaling_policy(self, t, stack, resource_name):
        rsrc = stack[resource_name]
        self.assertIsNone(rsrc.validate())
        scheduler.TaskRunner(rsrc.create)()
        self.assertEqual((rsrc.CREATE, rsrc.COMPLETE), rsrc.state)
        return rsrc

    def test_cooldown_is_in_progress_toosoon(self):
        t = template_format.parse(as_template)
        stack = utils.parse_stack(t, params=as_params)
        pol = self.create_scaling_policy(t, stack, 'WebServerScaleUpPolicy')

        now = timeutils.utcnow()
        previous_meta = {'cooldown': {
            now.isoformat(): 'ChangeInCapacity : 1'}}
        self.patchobject(pol, 'metadata_get', return_value=previous_meta)
        self.assertTrue(pol._cooldown_inprogress())

    def test_cooldown_is_in_progress_scaling_unfinished(self):
        t = template_format.parse(as_template)
        stack = utils.parse_stack(t, params=as_params)
        pol = self.create_scaling_policy(t, stack, 'WebServerScaleUpPolicy')

        previous_meta = {'scaling_in_progress': True}
        self.patchobject(pol, 'metadata_get', return_value=previous_meta)
        self.assertTrue(pol._cooldown_inprogress())

    def test_cooldown_not_in_progress(self):
        t = template_format.parse(as_template)
        stack = utils.parse_stack(t, params=as_params)
        pol = self.create_scaling_policy(t, stack, 'WebServerScaleUpPolicy')

        awhile_ago = timeutils.utcnow() - datetime.timedelta(seconds=100)
        previous_meta = {
            'cooldown': {
                awhile_ago.isoformat(): 'ChangeInCapacity : 1'
            },
            'scaling_in_progress': False
        }
        self.patchobject(pol, 'metadata_get', return_value=previous_meta)
        self.assertFalse(pol._cooldown_inprogress())

    def test_scaling_policy_cooldown_zero(self):
        t = template_format.parse(as_template)

        # Create the scaling policy (with Cooldown=0) and scale up one
        properties = t['Resources']['WebServerScaleUpPolicy']['Properties']
        properties['Cooldown'] = '0'

        stack = utils.parse_stack(t, params=as_params)
        pol = self.create_scaling_policy(t, stack, 'WebServerScaleUpPolicy')

        now = timeutils.utcnow()
        previous_meta = {now.isoformat(): 'ChangeInCapacity : 1'}
        self.patchobject(pol, 'metadata_get', return_value=previous_meta)
        self.assertFalse(pol._cooldown_inprogress())

    def test_scaling_policy_cooldown_none(self):
        t = template_format.parse(as_template)

        # Create the scaling policy no Cooldown property, should behave the
        # same as when Cooldown==0
        properties = t['Resources']['WebServerScaleUpPolicy']['Properties']
        del properties['Cooldown']

        stack = utils.parse_stack(t, params=as_params)
        pol = self.create_scaling_policy(t, stack, 'WebServerScaleUpPolicy')

        now = timeutils.utcnow()
        previous_meta = {now.isoformat(): 'ChangeInCapacity : 1'}
        self.patchobject(pol, 'metadata_get', return_value=previous_meta)
        self.assertFalse(pol._cooldown_inprogress())

    def test_metadata_is_written(self):
        t = template_format.parse(as_template)
        stack = utils.parse_stack(t, params=as_params)
        pol = self.create_scaling_policy(t, stack, 'WebServerScaleUpPolicy')

        nowish = timeutils.utcnow()
        reason = 'cool as'
        meta_set = self.patchobject(pol, 'metadata_set')
        self.patchobject(timeutils, 'utcnow', return_value=nowish)
        pol._cooldown_timestamp(reason)
        meta_set.assert_called_once_with(
            {'cooldown': {nowish.isoformat(): reason},
             'scaling_in_progress': False})

class ScalingPolicyAttrTest(common.HeatTestCase):
    def setUp(self):
        super(ScalingPolicyAttrTest, self).setUp()
        t = template_format.parse(as_template)
        self.stack = utils.parse_stack(t, params=as_params)
        self.policy = self.stack['WebServerScaleUpPolicy']
        self.assertIsNone(self.policy.validate())
        scheduler.TaskRunner(self.policy.create)()
        self.assertEqual((self.policy.CREATE, self.policy.COMPLETE),
                         self.policy.state)

    def test_alarm_attribute(self):
        self.assertIn("WebServerScaleUpPolicy",
                      self.policy.FnGetAtt('AlarmUrl'))


import os
import random

from skimage import io, transform
import scipy.io
from scipy.spatial.transform import Rotation
from scipy.spatial.transform import Slerp
import numpy as np
from torch.utils.data import Dataset, DataLoader
#import pptk

import util
import torch
# the main data class which has iterators to all datasets
# add more datasets accordingly

def get_root_dir(db):
    if db == "cad120":
        return os.path.join('..','data','rgbd-human','cad-120')
    elif db == 'biwi':
        return os.path.join('..','data','rgbd-human','kinect_head_pose_db','hpdb')
    elif db == 'biwiid':
        return os.path.join('..','data','rgbd-human','BIWI-I')
    elif db == 'human36':
        return os.path.join('..','data','human3.6m_downloader','training','subject')
    else:
        return ''

class Data():
    def __init__(self,db='all',face3dmm='lm'):
        if face3dmm == 'lm':
            self.dataloader = SyntheticLoader()
        else:
            self.dataloader = SyntheticLoader2()
        self.M = self.dataloader.M
        self.N = self.dataloader.N
        self.batchsize = 4
        self.shuffle = True
        self.transform = True
        self.batchloader = DataLoader(self.dataloader,
                batch_size=self.batchsize,
                shuffle=self.shuffle,
                num_workers=4)

    def __len__(self):
        return len(self.dataloader)

    def printinfo(self):
        print(f"RANDOM TRANSFORMS: TRUE")
        print(f"SHUFFLE: TRUE")
        print(f"BATCH SIZE: {self.batchsize}")
        print()

class TestData():
    def __init__(self):
        self.batchsize = 4

    def createLoader(self,f):
        self.dataloader = TestLoader(f)
        self.batchloader = DataLoader(self.dataloader,
                batch_size = self.batchsize,
                shuffle=False,
                num_workers=0)

        return self.batchloader

    def __len__(self):
        return len(self.dataloader)

class ImageLoader():

    def __init__(self,db='biwiid',addnoise=False):
        root_dir = get_root_dir(db)
        get_img_loader(db)

class TestLoader(Dataset):
    def __init__(self,f):
        root_dir = os.path.join("../data/synthetic_principal",f"sequencef{f:04d}")

        self.root_dir = root_dir
        self.files = [os.path.join(root_dir,f) for f in os.listdir(root_dir)]
        self.files = sorted(self.files,key=str.lower)
        self.M = 100
        self.N = 68

        self.w = 640
        self.h = 480
        self.px = self.w / 2
        self.py = self.h / 2

    def __len__(self):

        return len(self.files)

    def __getitem__(self,idx):
        fname = self.files[idx]
        data = scipy.io.loadmat(fname)
        self.M = 100
        self.N = 68

        tmp = data['sequence'][0,0]
        x_w = tmp['x_w']
        x_img_gt = tmp['x_img_true']
        x_cam = tmp['x_cam']
        R = tmp['R']
        T = tmp['T']
        f = torch.Tensor(tmp['f'].astype(np.float)[0]).float()
        d = np.mean(T[:,2])

        x_img = x_img_gt

        sample = {}
        sample['fname'] = fname
        sample['x_w_gt'] = torch.from_numpy(x_w - x_w.mean(0)).float()
        sample['x_cam_gt'] = torch.from_numpy(x_cam).float()
        sample['x_img'] = torch.from_numpy(x_img).float()
        sample['x_img_gt'] = torch.from_numpy(x_img_gt).float()
        #sample['x_img_gt'] = torch.from_numpy(x_img_gt).float().permute(0,2,1)
        sample['f_gt'] = torch.Tensor(f).float()
        sample['d_gt'] = torch.Tensor([d]).float()
        sample['T_gt'] = T

        if 'K' in data.keys():
            sample['K'] = torch.from_numpy(data['K']).float()
        else:
            sample['K'] = torch.tensor([[f,0,self.w/2],[0,f,self.h/2],[0,0,1]],dtype=torch.float32)

        return sample

class Cad60Loader(Dataset):
    def __init__(self):
        if os.path.isdir("../data"):
            self.root_dir = os.path.join("../data/tmp/cad60/processed")
        else:
            self.root_dir = os.path.join("../data0/tmp/cad60/processed")
        files = os.listdir(self.root_dir)
        sorted(files,key=str.lower)

        self.all_paths = [os.path.join(self.root_dir,f) for f in files]

    def __len__(self):

        return len(self.all_paths)

    def __getitem__(self,idx):
        path = self.all_paths[idx]
        #print(f"load: {path}")
        data = scipy.io.loadmat(path)

        x2d = data['x2d']
        xcam = data['xcam']
        f = data['fgt'][0,0]
        M = x2d.shape[0]
        N = x2d.shape[1]

        pts = x2d.reshape((M*N,2))
        x2d = np.transpose(x2d,(0,2,1))
        xcam = np.transpose(xcam,(0,2,1))

        sample = {}
        sample['x_img_gt'] = torch.from_numpy(x2d).float()
        sample['x_img'] = torch.from_numpy(pts).float()
        sample['x_cam_gt'] = torch.from_numpy(xcam).float()
        sample['f_gt'] = torch.from_numpy(np.array([f]).astype(np.float)).float()

        return sample

class Cad120Loader(Dataset):
    def __init__(self):
        if os.path.isdir("../data"):
            #self.root_dir = os.path.join("../data/tmp/cad120/processed")
            #shape_dir = "../data/face_alignment/300W_LP/Code/ModelGeneration/shape_simple.mat"
            self.root_dir = os.path.join("../data/tmp/cad120/processed_orig")
            shape_dir = "../data/face_alignment/300W_LP/Code/ModelGeneration/shape_simple.mat"
        files = os.listdir(self.root_dir)
        files = sorted(files,key=str.lower)

        self.all_paths = [os.path.join(self.root_dir,f) for f in files]
        print(files)

        self.w = 640
        self.h = 480
        self.px = 320
        self.py = 240

    def __len__(self):

        return len(self.all_paths)

    def __getitem__(self,idx):
        path = self.all_paths[idx]
        data = scipy.io.loadmat(path)

        x2d = data['x2d_raw']
        f = data['fgt'][0,0]
        M = x2d.shape[0]
        N = x2d.shape[1]

        x2d = np.transpose(x2d,(0,2,1))

        K_gt = torch.tensor([[f, 0, self.px],[0,f,self.py],[0,0,1]],dtype=torch.float32)

        sample = {}
        if 'imgs' in data.keys(): sample['imgs'] = data['imgs']
        sample['x_w_gt'] = None
        sample['x_img_gt'] = None
        sample['K_gt'] = K_gt
        sample['x_img'] = torch.from_numpy(x2d).float()
        sample['d'] = torch.from_numpy(data['depth'].astype(np.float32)).float().mean(1)
        sample['f_gt'] = torch.from_numpy(np.array([f]).astype(np.float)).float()

        return sample

class Human36Loader(Dataset):
    def __init__(self):
        if os.path.isdir("../data"):
            self.root_dir = os.path.join("../data/tmp/human36/processed")
        else:
            self.root_dir = os.path.join("../data0/tmp/human36/processed")
        files = os.listdir(self.root_dir)
        #files = sorted(files,key=str.lower)
        files.sort()
        self.all_paths = [os.path.join(self.root_dir,f) for f in files]

        self.w = 1000
        self.h = 1000
        self.px = 500
        self.py = 500

    def __len__(self):
        return len(self.all_paths)

    def __getitem__(self,idx):
        path = self.all_paths[idx]
        print(path)
        #path = '../data/tmp/human36/processed/S1_Eating 2.55011271_sequence.mat'
        #print(f"load: {path}")
        data = scipy.io.loadmat(path)

        x2d = data['x2d_raw']
        xcam = data['xcam']
        f = data['fgt'][0,0]

        x2d = np.transpose(x2d,(0,2,1))
        xcam = np.transpose(xcam,(0,2,1))

        K_gt = torch.tensor([[f, 0, self.px],[0,f,self.py],[0,0,1]],dtype=torch.float32)

        sample = {}
        sample['x_w_gt'] = None
        sample['x_img_gt'] = None
        sample['K_gt'] = K_gt
        sample['x_img'] = torch.from_numpy(x2d).float()
        sample['x_cam'] = torch.from_numpy(xcam).float()
        sample['d'] = torch.from_numpy(data['depth'].astype(np.float32)).float().mean(1)
        sample['f_gt'] = torch.from_numpy(np.array([f]).astype(np.float)).float()

        return sample

class BIWIID3ddfaLoader(Dataset):

    def __init__(self):
        self.root_dir = os.path.join("/home/huynshen/projects/git/3DDFA_V2/biwiid_faces/processed")
        files = os.listdir(self.root_dir)
        #files = sorted(files,key=str.lower)

        self.all_paths = [os.path.join(self.root_dir,f) for f in files]
        #files.sort()

    def __len__(self):

        return len(self.all_paths)

    def __getitem__(self,idx):
        path = self.all_paths[idx]
        #print(f"load: {path}")
        data = scipy.io.loadmat(path)

        x2d = data['x2d']
        f = data['fgt'][0,0]
        x2d = np.transpose(x2d,(0,2,1))
        #xcam = np.transpose(xcam,(0,2,1))

        sample = {}
        sample['x_w_gt'] = None
        sample['x_img_gt'] = None
        sample['x_img'] = torch.from_numpy(x2d).float()
        sample['x_cam'] = None
        sample['f_gt'] = torch.from_numpy(np.array([f]).astype(np.float)).float()

        return sample

class BIWIIDLoader(Dataset):

    def __init__(self):
        if os.path.isdir("../data"):
            #self.root_dir = os.path.join("../data/tmp/biwi-i/processed")
            self.root_dir = os.path.join("../data/tmp/biwi-i/processed_orig")
        files = os.listdir(self.root_dir)
        #files = sorted(files,key=str.lower)

        files.sort()
        self.all_paths = [os.path.join(self.root_dir,f) for f in files]

        self.w = 640
        self.h = 480
        self.px = 320
        self.py = 240

    def __len__(self):

        return len(self.all_paths)

    def __getitem__(self,idx):
        path = self.all_paths[idx]
        print(f"load: {path}")
        data = scipy.io.loadmat(path,squeeze_me=True)

        x2d = data['x2d_raw']
        #x2d = data['x2d']
        #x2d[:,:,0] = x2d[:,:,0] + self.w
        #x2d[:,:,1] = x2d[:,:,1] + self.h
        x2d[:,:,0] = x2d[:,:,0] / self.w * self.px
        x2d[:,:,1] = x2d[:,:,1] / self.h * self.py
        #xcam = data['xcam']
        f = data['fgt']
        x2d = np.transpose(x2d,(0,2,1))
        #xcam = np.transpose(xcam,(0,2,1))

        K_gt = torch.tensor([[f,0,self.px],[0,f,self.py],[0,0,1]],dtype=torch.float32)

        sample = {}
        sample['fname'] = data['fname']
        sample['x_img'] = torch.from_numpy(x2d).float()
        sample['K_gt'] = K_gt
        #sample['x_cam'] = torch.from_numpy(xcam).float()
        sample['d'] = torch.from_numpy(data['depth'].astype(np.float32)).float().mean(1)
        sample['f_gt'] = torch.from_numpy(np.array([f]).astype(np.float)).float()

        return sample

class CheckerboardLoader(Dataset):
    def __init__(self):

        self.root_dir = '/home/huynshen/projects/git/A-flexible-new-technique-for-camera-calibration/zhang'
        self.file = 'checkerboard.mat'

    def __len__(self):

        return 1

    def __getitem__(self,idx):
        path = os.path.join(self.root_dir,self.file)
        data = scipy.io.loadmat(path,struct_as_record=False,squeeze_me=True)
        Fs = data['data'].Fs
        x_w_gt = data['data'].M
        x_img = data['data'].x
        weights = data['data'].ws
        K_gt = data['data'].K

        sample = {}
        sample['x_img'] = torch.from_numpy(x_img).float().permute(2,0,1)
        sample['x_w_gt'] = torch.from_numpy(x_w_gt).float()
        sample['K_gt'] = torch.from_numpy(K_gt).float()
        sample['Fs'] = torch.from_numpy(Fs).float()
        sample['weights'] = torch.from_numpy(weights).float()
        return sample

class MeritonLoader(Dataset):
    def __init__(self):
        self.root_dir = './extras/meriton_college.mat'
        data_dict = scipy.io.loadmat(self.root_dir,squeeze_me=True,struct_as_record=False)
        self.A = data_dict['data'].A_gt
        self.F = data_dict['data'].F
        pts = data_dict['data'].pts
        obs_file = '/home/huynshen/Downloads/meriton_college_1/2D/nview-corners_masa'
        p_w_file = '/home/huynshen/Downloads/meriton_college_1/3D/p3d'
        observation = np.genfromtxt(obs_file)
        self.p_w = np.genfromtxt(p_w_file)

        m1 = observation[:,0] > 0
        m2 = observation[:,1] > 0
        m3 = observation[:,2] > 0
        m = m1 & m2 & m3

        p1 = pts[0].x2d[observation[:,0][m].astype(np.int),:]
        p2 = pts[1].x2d[observation[:,1][m].astype(np.int),:]
        p3 = pts[2].x2d[observation[:,2][m].astype(np.int),:]
        self.x_img = np.stack((p1,p2,p3))
        self.p_w = self.p_w[m]

    def __len__(self):
        return 1

    def __getitem__(self,idx):

        sample = {}
        sample['x_img'] = torch.from_numpy(self.x_img).float().permute(0,2,1)
        sample['x_w_gt'] = torch.from_numpy(self.p_w).float()
        sample['Fs'] = torch.from_numpy(self.F).float()
        sample['K_gt'] = torch.from_numpy(self.A).float()

        return sample

class BIWILoader(Dataset):
    def __init__(self):
        if os.path.isdir("../data"):
            self.root_dir = os.path.join("../data/tmp/biwi/processed")
            self.shape_dir = os.path.join("../data/tmp/biwi/shape")
        files = os.listdir(self.root_dir)
        files = sorted(files,key=str.lower)

        self.w = 640
        self.h = 480
        self.px = 320
        self.py = 240
        self.res = 300

        self.all_paths = [os.path.join(self.root_dir,f) for f in files]
        self.maxangle = 20

    def __len__(self):

        return len(self.all_paths)

    def __getitem__(self,idx):
        path = self.all_paths[idx]
        #print(f"load: {path}")
        data = scipy.io.loadmat(path)
        print(data.keys())
        quit()

        shape_path = os.path.join(self.shape_dir,f"{idx+1:02d}_shape.mat")
        shape_data = scipy.io.loadmat(shape_path)

        x2d = data['x2d']
        x2d[:,:,0] = x2d[:,:,0] + self.px
        x2d[:,:,1] = x2d[:,:,1] + self.py
        xcam = data['xcam']
        f = data['fgt'][0,0]

        K_gt = torch.tensor([[f,0,self.px],[0,f,self.py],[0,0,1]],dtype=torch.float32)

        x2d = np.transpose(x2d,(0,2,1))
        xcam = np.transpose(xcam,(0,2,1))

        sample = {}
        sample['x_w_gt'] = None
        sample['x_img'] = torch.from_numpy(x2d).float()
        sample['K_gt'] = K_gt
        sample['x_cam'] = torch.from_numpy(xcam).float()
        sample['d'] = torch.norm(sample['x_cam'].mean(2),dim=1)
        sample['f_gt'] = torch.from_numpy(np.array([f]).astype(np.float)).float()

        S = torch.from_numpy(shape_data['mean_lm']).float()
        S = S - S.mean(0).unsqueeze(0)
        S[:,2] = S[:,2] * -1
        sample['shape'] = S

        return sample

class AnalysisLoader(Dataset):
    def __init__(self):
        if os.path.isdir("../data"):
            self.root_dir = os.path.join("../data/tmp/biwi/processed")
            shape_dir = "../data/face_alignment/300W_LP/Code/ModelGeneration/shape_simple.mat"
            self.shape_dir = os.path.join("../data/tmp/biwi/shape")
        else:
            self.root_dir = os.path.join("../data0/tmp/biwi/processed")
            shape_dir = "../data0/face_alignment/300W_LP/Code/ModelGeneration/shape_simple.mat"
            self.shape_dir = os.path.join("../data0/tmp/biwi/shape")
        subjects = [f"{sub:02d}" for sub in range(1,25)]
        files = os.listdir(self.root_dir)
        files = sorted(files,key=str.lower)
        self.all_paths = [os.path.join(self.root_dir,f) for f in files]

        shape_data = scipy.io.loadmat(shape_dir)
        mu_lm = shape_data['mu_lm']
        mu_exp = shape_data['mu_exp']
        lm_eigenvec = shape_data['lm_eigenvec']
        exp_eigenvec = shape_data['exp_eigenvec']

        self.mu_lm = mu_lm.T
        self.mu_lm = self.mu_lm - np.mean(self.mu_lm,0)
        self.mu_exp = mu_exp.T
        self.mu_exp = self.mu_exp - np.mean(self.mu_exp,0)

        self.lm_eigenvec = lm_eigenvec
        self.exp_eigenvec = exp_eigenvec

        self.minangle = 20
        self.maxangle = 20

    def __len__(self):
        return 24

    def __getitem__(self,idx):
        path = self.all_paths[idx]
        #print(f"load: {path}")
        data = scipy.io.loadmat(path)

        shape_path = os.path.join(self.shape_dir,f"{idx+1:02d}_shape.mat")
        shape_data = scipy.io.loadmat(shape_path)

        x_w = shape_data['mean_lm']
        x_w[:,2] = x_w[:,2] * -1
        x2d = data['x2d']
        xcam = data['xcam']
        f = data['fgt'][0,0]
        R = data['R']
        M = x2d.shape[0]
        N = x2d.shape[1]

        # find valid views
        validview = []
        for i in range(M):
            r = Rotation.from_matrix(R[i])
            angles = r.as_euler('zyx',degrees=False)
            angles = np.arcsin(np.sin(angles)) * 180 / np.pi
            if np.any(np.abs(angles) > self.maxangle): continue
            else: validview.append(i)

        x2d = x2d[validview]
        xcam = xcam[validview]
        M = x2d.shape[0]
        N = x2d.shape[1]

        pts = x2d.reshape((M*N,2))
        x2d = np.transpose(x2d,(0,2,1))
        xcam = np.transpose(xcam,(0,2,1))

        sample = {}
        sample['x_w_gt'] = torch.from_numpy(x_w).float()
        sample['x_img_gt'] = torch.from_numpy(x2d).float()
        sample['x_img'] = torch.from_numpy(pts).float()
        sample['x_cam_gt'] = torch.from_numpy(xcam).float()
        sample['f_gt'] = torch.from_numpy(np.array([f]).astype(np.float)).float()

        return sample

# loads M views and N points depending on the setting
# ablation for synthetic dataset
class MNLoader(Dataset):
    def __init__(self,f,M=100,N=68,addnoise=True,seed=0):
        if os.path.isdir("../data"):
            root_dir = os.path.join("../data/synthetic_3dface",f"sequencef{f:04d}")
        else:
            root_dir = os.path.join("../data0/synthetic_3dface",f"sequencef{f:04d}")
        random.seed(seed)
        self.root_dir = root_dir
        self.files = [os.path.join(root_dir,f) for f in os.listdir(root_dir)]
        self.files = sorted(self.files,key=str.lower)
        self.addnoise = addnoise

        #self.transform = transforms.Compose([ToTensor()])
        if os.path.isdir("../data"):
            root_dir = "../data/face_alignment/300W_LP/Code/ModelGeneration/shape_simple.mat"
        else:
            root_dir = "../data0/face_alignment/300w_lp/code/modelgeneration/shape_simple.mat"

        # load shape data
        shape_data = scipy.io.loadmat(root_dir)
        self.mu_s = shape_data['mu_lm'].T
        self.mu_exp = shape_data['mu_exp'].T
        lm_eigenvec = shape_data['lm_eigenvec']
        exp_eigenvec = shape_data['exp_eigenvec']
        sigma = np.diag(shape_data['sigma'][:,0] / 100)

        if N > 68: N=68

        #lm_eigenvec = lm_eigenvec.reshape(53215,3,199)
        #exp_eigenvec = exp_eigenvec.reshape(53215,3,29)
        lm_eigenvec = lm_eigenvec.reshape(68,3,199)
        exp_eigenvec = exp_eigenvec.reshape(68,3,29)
        self.lm_eigenvec = torch.from_numpy(lm_eigenvec[:N].reshape(-1,199)).float()
        self.exp_eigenvec = torch.from_numpy(exp_eigenvec[:N].reshape(-1,29)).float()
        self.mu_s = torch.from_numpy(self.mu_s[:N]).float()
        self.mu_exp = torch.from_numpy(self.mu_exp[:N]).float()
        self.sigma = torch.from_numpy(sigma).float()

        # bideo sequence length
        self.M = M
        self.N = N

        # get indices
        indices = []
        le = list(range(36,42)) + list(range(17,22))
        re = list(range(42,48)) + list(range(22,27))
        nose = list(range(27,36))
        mouth = list(range(48,68))
        jaw = list(range(0,17))
        random.shuffle(le)
        random.shuffle(re)
        random.shuffle(nose)
        random.shuffle(mouth)
        random.shuffle(jaw)
        lm = [le,re,nose,mouth,jaw]
        for i in range(68):
            idx = lm[i % len(lm)].pop()
            indices.append(idx)
            if len(lm[i % len(lm)]) == 0:
                lm.pop(i % len(lm))
        self.indices = indices

    def __len__(self):

        return len(self.files)

    def __getitem__(self,idx):
        fname = self.files[idx]
        data = scipy.io.loadmat(fname)

        tmp = data['sequence'][0,0]
        x_w = tmp['x_w']
        x_img_gt = tmp['x_img_true']
        x_cam = tmp['x_cam']
        R = tmp['R']
        T = tmp['T']
        f = torch.Tensor(tmp['f'].astype(np.float)[0]).float()
        d = np.mean(T[:,2])

        if self.addnoise:
            le = x_img_gt[:,36,:]
            re = x_img_gt[:,45,:]
            std = np.max(np.linalg.norm(le - re,axis=1)*0.05)
            noise = np.random.rand(self.M,self.N,2) * std
        else:
            noise = 0

        x_img_gt = x_img_gt[:self.M,:self.N,:]
        x_w = x_w[:self.N,:]
        x_cam = x_cam[:self.M,:self.N]
        x_img_gt[:,:,0] = x_img_gt[:,:,0] - 320
        x_img_gt[:,:,1] = x_img_gt[:,:,1] - 240
        x_img = x_img_gt + noise
        x_img = x_img.reshape((self.M*self.N,2))

        sample = {}
        sample['x_w_gt'] = torch.from_numpy(x_w).float()
        sample['x_cam_gt'] = torch.from_numpy(x_cam).float().permute(0,2,1)
        sample['x_img'] = torch.from_numpy(x_img).float()
        sample['x_img_gt'] = torch.from_numpy(x_img_gt).float().permute(0,2,1)
        sample['f_gt'] = torch.Tensor(f).float()

        return sample

    def get3DMM_(self,a):
        a = a.unsqueeze(-1)
        delta = torch.mm(torch.mm(self.lm_eigenvec,self.sigma),a).view(self.N,3)
        shape = self.mu_s + delta
        shape = shape - shape.mean(0).unsqueeze(0)
        return shape

class SyntheticLoader2(Dataset):

    def __init__(self):

        #self.transform = transforms.Compose([ToTensor()])
        if os.path.isdir("../data"):
            root_dir = '/home/huynshen/data/face_alignment/300W_LP/Code/ModelGeneration/Model_Shape.mat'

        # load shape data
        shape_data = scipy.io.loadmat(root_dir)
        self.sigma = shape_data['sigma'][:,0]
        self.s_eig = shape_data['shape_eigenvec']
        self.s_mu = shape_data['mu_shape']

        # video sequence length
        self.M = 100
        self.N = self.s_mu.shape[0] // 3

        # extra boundaries on camera coordinates
        self.maxn = 500
        self.minn = 20
        self.w = 640
        self.h = 480
        self.minf = 300; self.maxf = 1500
        self.minz = 800; self.maxz = 8000;
        self.max_rx = 20;
        self.max_ry = 20; self.max_rz = 20;
        self.std = 0.05

    def __len__(self):
        return 10000

    def __getitem__(self,idx):

        # data holders
        M = self.M
        x_w = np.zeros((self.N,3))

        # define intrinsics
        while True:
            f = 500 + np.random.randn() * (self.maxf - self.minf);
            if f >= self.minf and f < self.maxf: break
        K = np.array([[f,0,self.w/2],[ 0,f,self.h/2], [0,0,1]])

        # create random 3dmm shape
        s,alpha = self.generateRandomFace()

        #import pptk
        #v = pptk.viewer(s)
        #v.set(point_size=1.1)

        # create random 3dmm expression
        #beta = np.random.randn(29) * 0.1
        #e = np.sum(np.expand_dims(beta,0)*self.exp_eigenvec,1)
        #exp = e.reshape(N,3)

        # define depth
        tz = np.random.random() * (self.maxz-self.minz) + self.minz
        minz = np.maximum(tz - 500,self.minz)
        maxz = np.minimum(tz + 500,self.maxz)

        # get initial and final rotation
        while True:
            r_init, q_init = self.generateRandomRotation()
            r_final, q_final = self.generateRandomRotation()
            t_init = self.generateRandomTranslation(K,minz,maxz)
            t_final = self.generateRandomTranslation(K,minz,maxz)

            ximg_init = self.project2d(r_init,t_init,K,s)
            ximg_final = self.project2d(r_final,t_final,K,s)
            if np.any(np.amin(ximg_init,axis=0) < -320): continue
            if np.any(np.amin(ximg_final,axis=0) < -320): continue
            if np.any(np.amin(ximg_init,axis=1) < -240): continue
            if np.any(np.amin(ximg_final,axis=1) < -240): continue
            init = np.amax(ximg_init,axis=0)
            final = np.amax(ximg_final,axis=0)
            if init[0] > 320: continue
            if final[0] > 320: continue
            if init[1] > 240: continue
            if final[1] > 240: continue
            break
        d = (t_init[2] + t_final[2]) / 2

        # interpolate quaternion using spherical linear interpolation
        qs = np.stack((q_init,q_final))
        Rs = Rotation.from_quat(qs)
        times = np.linspace(0,1,M)
        slerper = Slerp([0,1],Rs)
        rotations = slerper(times)
        matrices = rotations.as_matrix()

        T = np.stack((np.linspace(t_init[0],t_final[0],M),
                np.linspace(t_init[1],t_final[1],M),
                np.linspace(t_init[2],t_final[2],M))).T

        x_cam = np.matmul(matrices,np.stack(M*[x_w.T])) + np.expand_dims(T,-1)
        proj = np.matmul(np.stack(M*[K]),x_cam)
        proj = proj / np.expand_dims(proj[:,2,:],1)
        proj = proj.transpose(0,2,1)
        x_img_true = proj[:,:,:2]

        #le = x_img_true[:,36,:]
        #re = x_img_true[:,45,:]
        #std = np.mean(np.linalg.norm(le - re,axis=1)*0.05)

        noise = np.random.randn(x_img_true.shape[0],x_img_true.shape[1],x_img_true.shape[2]) * self.std
        x_img = x_img_true + noise

        # create dictionary for results
        sample = {}
        sample['alpha_gt'] = torch.from_numpy(alpha).float()
        sample['x_w_gt'] = torch.from_numpy(x_w).float()
        sample['x_cam_gt'] = torch.from_numpy(x_cam).float()
        sample['x_img'] = x_img
        sample['x_img_gt'] = x_img_true
        sample['f_gt'] = torch.Tensor([f]).float()
        #sample['x_img_norm'] = torch.from_numpy(x_img_norm).float()
        #sample['K_gt'] = torch.from_numpy(K).float()
        #sample['R_gt'] = torch.from_numpy(R).float()
        #sample['Q_gt'] = torch.from_numpy(Q).float()
        #sample['T_gt'] = torch.from_numpy(T).float()
        return sample

    def generateRandomFace(self):
        N = self.s_mu.shape[0] // 3
        alphas = np.random.randn(199)[:,np.newaxis]
        d = np.matmul(self.s_eig,np.diag(self.sigma))
        d = np.matmul(d,alphas)
        s = self.s_mu + d
        s = s / 1000
        s = s.reshape((N,3))
        s[2,:] = s[2,:] * -1

        return s, alphas

    def generateRandomRotation(self):

        ax = self.max_rx;
        ay = self.max_ry;
        az = self.max_rz;
        rx = np.random.random()*2*ax - ax;
        ry = np.random.random()*2*ay - ay;
        rz = np.random.random()*2*az - az;

        r = Rotation.from_euler('zyx',[rz,ry,rx],degrees=True)
        q = r.as_quat()

        return r,q

    def generateRandomTranslation(self,K,minz,maxz,w=640,h=480):

        xvec = np.array([[w],[w/2],[1]])
        yvec = np.array([[h/2],[h],[1]])
        vz = np.array([[0],[0],[1]]);
        vx = np.matmul(np.linalg.inv(K),xvec)
        vy = np.matmul(np.linalg.inv(K),yvec)
        vx = np.squeeze(vx)
        vy = np.squeeze(vy)
        vz = np.array([0,0,1])
        thetax = np.arctan2(np.linalg.norm(np.cross(vz,vy)),np.dot(vz,vy));
        thetay = np.arctan2(np.linalg.norm(np.cross(vz,vx)),np.dot(vz,vx));

        tz = np.random.random()*(maxz-minz) + minz;
        maxx = tz * np.tan(thetax);
        maxy = tz * np.tan(thetay);
        tx = np.random.random()*maxx*2 - maxx;
        ty = np.random.random()*maxy*2 - maxy;
        t = [tx,ty,tz];

        return np.array(t)

    def project2d(self,r,t,K,pw):

        R = r.as_matrix()
        xc = np.matmul(R,pw.T) + np.expand_dims(t,1);

        proj = np.matmul(K,xc)
        proj = proj / proj[2,:]
        ximg = proj.T

        return ximg

class SyntheticLoader(Dataset):

    def __init__(self):

        root_dir = "../extras/shape_simple.mat"

        # load shape data
        shape_data = scipy.io.loadmat(root_dir)
        mu_lm = shape_data['mu_lm']
        mu_exp = shape_data['mu_exp']
        self.lm_eigenvec = shape_data['lm_eigenvec']
        exp_eigenvec = shape_data['exp_eigenvec']
        self.sigma = shape_data['sigma'] / 100

        self.mu_lm = mu_lm.T
        self.mu_lm = self.mu_lm - np.mean(self.mu_lm,0)

        self.mu_exp= mu_exp.T
        self.mu_exp= self.mu_exp - np.mean(self.mu_exp,0)

        self.exp_eigenvec = exp_eigenvec

        # video sequence length
        self.M = 100
        self.N = 68
        self.res = 640

        # extra boundaries on camera coordinates
        self.w = 640
        self.h = 480
        self.minf = 300; self.maxf = 1500
        self.minz = 800; self.maxz = 8000;
        self.max_rx = 20;
        self.max_ry = 20; self.max_rz = 20;
        self.xstd = 1; self.ystd = 1;

    def __len__(self):
        return 10000

    # betas (199,1)
    def get_3dmm(self,betas):
        b = betas.shape[0]
        mu_s  = torch.from_numpy(self.mu_lm).float().to(betas.device).clone()
        lm_eig = torch.from_numpy(self.lm_eigenvec).float().to(betas.device)
        sigma = torch.from_numpy(self.sigma).float().to(betas.device)

        mu_s[:,2] = mu_s[:,2] * -1
        sigma = torch.diag(sigma.squeeze())

        # adjust batch size
        lm_eig = torch.mm(lm_eig, sigma).unsqueeze(0).repeat(b,1,1)

        # predict batchwise 3d shapes
        shape = mu_s.unsqueeze(0).repeat(b,1,1) + torch.bmm(lm_eig,betas).view(b,-1,3)
        shape = shape - shape.mean(1).unsqueeze(1)
        return shape

    def __getitem__(self,idx):
        # data holders
        M = self.M
        N = self.N
        x_w = np.zeros((N,3));
        x_cam = np.zeros((M,N,3));
        x_img = np.zeros((M,N,2));
        x_img_true = np.zeros((M,N,2));

        # define intrinsics
        f = self.minf + random.random() * (self.maxf - self.minf);
        K = np.array([[f,0,self.w/2],[ 0,f,self.h/2], [0,0,1]])

        # create random 3dmm shape
        alpha = np.random.randn(199,1) * 5
        lm_eigenvec = np.matmul(self.lm_eigenvec,np.diag(self.sigma[:,0]))
        s = np.squeeze(np.matmul(lm_eigenvec,alpha))
        s = s.reshape(N,3)
        lm = self.mu_lm + s
        lm = lm - np.expand_dims(np.mean(lm,axis=0),axis=0)
        x_w = lm

        # create random 3dmm expression
        beta = np.random.randn(29) * 0.1
        e = np.sum(np.expand_dims(beta,0)*self.exp_eigenvec,1)
        exp = e.reshape(N,3)

        # define depth
        tz = np.random.random() * (self.maxz-self.minz) + self.minz
        minz = np.maximum(tz - 500,self.minz)
        maxz = np.minimum(tz + 500,self.maxz)

        # get initial and final rotation
        while True:
            r_init, q_init = self.generateRandomRotation()
            r_final, q_final = self.generateRandomRotation()
            t_init = self.generateRandomTranslation(K,minz,maxz)
            t_final = self.generateRandomTranslation(K,minz,maxz)

            ximg_init = self.project2d(r_init,t_init,K,x_w)
            ximg_final = self.project2d(r_final,t_final,K,x_w)
            if np.any(np.amin(ximg_init,axis=0) < 0): continue
            if np.any(np.amin(ximg_final,axis=0) < 0): continue
            if np.any(np.amin(ximg_init,axis=1) < 0): continue
            if np.any(np.amin(ximg_final,axis=1) < 0): continue
            init = np.amax(ximg_init,axis=0)
            final = np.amax(ximg_final,axis=0)
            if init[0] > self.w: continue
            if final[0] > self.w: continue
            if init[1] > self.h: continue
            if final[1] > self.h: continue
            break
        d = (t_init[2] + t_final[2]) / 2

        # interpolate quaternion using spherical linear interpolation
        qs = np.stack((q_init,q_final))
        Rs = Rotation.from_quat(qs)
        times = np.linspace(0,1,M)
        slerper = Slerp([0,1],Rs)
        rotations = slerper(times)
        matrices = rotations.as_matrix()

        T = np.stack((np.linspace(t_init[0],t_final[0],M),
                np.linspace(t_init[1],t_final[1],M),
                np.linspace(t_init[2],t_final[2],M))).T

        x_cam = np.matmul(matrices,np.stack(M*[x_w.T])) + np.expand_dims(T,-1)
        proj = np.matmul(np.stack(M*[K]),x_cam)
        proj = proj / np.expand_dims(proj[:,2,:],1)
        proj = proj.transpose(0,2,1)
        x_img_true = proj[:,:,:2]

        #le = x_img_true[:,36,:]
        #re = x_img_true[:,45,:]
        #std = np.mean(np.linalg.norm(le - re,axis=1)*0.05)

        noise = np.random.randn(M,N,2)

        x_img = x_img_true + noise

        # create dictionary for results
        sample = {}
        sample['alpha_gt'] = torch.from_numpy(alpha).float()
        sample['x_w_gt'] = torch.from_numpy(x_w).float()
        sample['x_cam_gt'] = torch.from_numpy(x_cam).float()
        sample['x_img'] = torch.from_numpy(x_img).float().permute(0,2,1)
        sample['x_img_gt'] = torch.from_numpy(x_img_true).float().permute(0,2,1)
        sample['f_gt'] = torch.Tensor([f]).float()
        #sample['x_img_norm'] = torch.from_numpy(x_img_norm).float()
        #sample['K_gt'] = torch.from_numpy(K).float()
        #sample['R_gt'] = torch.from_numpy(R).float()
        #sample['Q_gt'] = torch.from_numpy(Q).float()
        #sample['T_gt'] = torch.from_numpy(T).float()
        return sample

    def generateRandomRotation(self):

        ax = self.max_rx;
        ay = self.max_ry;
        az = self.max_rz;
        rx = np.random.random()*2*ax - ax;
        ry = np.random.random()*2*ay - ay;
        rz = np.random.random()*2*az - az;

        r = Rotation.from_euler('zyx',[rz,ry,rx],degrees=True)
        q = r.as_quat()

        return r,q

    def generateRandomTranslation(self,K,minz,maxz,w=640,h=480):

        xvec = np.array([[w],[w/2],[1]])
        yvec = np.array([[h/2],[h],[1]])
        vz = np.array([[0],[0],[1]]);
        vx = np.matmul(np.linalg.inv(K),xvec)
        vy = np.matmul(np.linalg.inv(K),yvec)
        vx = np.squeeze(vx)
        vy = np.squeeze(vy)
        vz = np.array([0,0,1])
        thetax = np.arctan2(np.linalg.norm(np.cross(vz,vy)),np.dot(vz,vy));
        thetay = np.arctan2(np.linalg.norm(np.cross(vz,vx)),np.dot(vz,vx));

        tz = np.random.random()*(maxz-minz) + minz;
        maxx = tz * np.tan(thetax);
        maxy = tz * np.tan(thetay);
        tx = np.random.random()*maxx*2 - maxx;
        ty = np.random.random()*maxy*2 - maxy;
        t = [tx,ty,tz];

        return np.array(t)

    def project2d(self,r,t,K,pw):

        R = r.as_matrix()
        xc = np.matmul(R,pw.T) + np.expand_dims(t,1);

        proj = np.matmul(K,xc)
        proj = proj / proj[2,:]
        ximg = proj.T

        return ximg

# LOADER FOR BIWI KINECT DATASET ONLY
class Face3DMM():
    def __init__(self,
            ):

        # load shape data
        shape_data = scipy.io.loadmat(shape_path)
        self.lm = torch.from_numpy(shape_data['keypoints'].astype(np.int32)).long().squeeze()
        self.mu_shape = torch.from_numpy(mu_shape.reshape(53215,3)).float()
        self.mu_shape = self.mu_shape - self.mu_shape.mean(0).unsqueeze(0)
        self.mu_shape = self.mu_shape / torch.max(torch.abs(self.mu_shape))
        self.shape_eigenvec = torch.from_numpy(shape_eigenvec).float()

        # load expression data
        exp_data = scipy.io.loadmat(exp_path)
        mu_exp = exp_data['mu_exp']
        exp_eigenvec = exp_data['w_exp']

        self.mu_exp = torch.from_numpy(mu_exp.reshape(53215,3)).float()
        self.exp_eigenvec = torch.from_numpy(exp_eigenvec).float()

class Rescale(object):
    """Rescale the image in a sample to a given size.

    Args:
        output_size (tuple or int): Desired output size. If tuple, output is
            matched to output_size. If int, smaller of image edges is matched
            to output_size keeping aspect ratio the same.
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        self.output_size = output_size

    def __call__(self, image):

        h, w = image.shape[:2]
        if isinstance(self.output_size, int):
            if h > w:
                new_h, new_w = self.output_size * h / w, self.output_size
            else:
                new_h, new_w = self.output_size, self.output_size * w / h
        else:
            new_h, new_w = self.output_size

        new_h, new_w = int(new_h), int(new_w)

        # torch image: C X H X W
        img = transform.resize(image, (new_h, new_w))
        img = img.transpose((2, 0, 1))

        return img

class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""

    def __call__(self, data):
        # swap color axis because
        # numpy image: H x W x C
        return torch.from_numpy(data).float()

#Batch LOader
#BatchLoader = DataLoader(LFWLoader(),batch_size=8,shuffle=True,num_workers=4)

# UNIT TESTING
if __name__ == '__main__':

    # test dataloaders
    loader = MeritonLoader()
    print(loader.A.shape)
    print(loader.x_img.shape)
    print(loader.p_w.shape)
    quit()

    loader = CheckerboardLoader()
    data = loader[0]
    print(data['x_img'].shape)
    data['x_w_gt']
    data['K_gt']
    data['Fs']
    data['weights']
    quit()

    loader = AllLoader()
    data = loader[0]
    print(data.keys())
    quit()

    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D

    loader = MNLoader(400)
    quit()

    loader = AnalysisLoader()
    sample = loader[1]

    print(sample['x_w_gt'].shape)
    print(sample['x_cam_gt'].shape)
    print(sample['x_img'].shape)
    print(sample['x_img_gt'].shape)
    print(sample['f_gt'].shape)

    quit()

    loader = BIWILoader()
    sample = loader[1]
    print(sample['x_img'].shape)
    print(sample['x_img_gt'].shape)
    print(sample['f_gt'].shape)
    print(sample['x_cam_gt'].shape)

    import matplotlib.pyplot as plt
    for i in range(499):
        x = sample['x_img_gt'].numpy()[i].T
        print(x.shape)
        plt.scatter(x[:68,0],x[:68,1])
        plt.show()
    quit()

    loader = BIWIIDLoader()
    sample = loader[1]
    ximg = sample['x_img_gt']
    M = ximg.shape[0]
    for i in range(M):
        fig = plt.figure()
        ax = Axes3D(fig)
        x = sample['x_cam_gt'].numpy()[i].T
        ax.scatter(x[:,0],x[:,1],x[:,2])
        plt.show()
    print(ximg.shape)
    quit()

    loader = TestLoader(1000)
    sample = loader[1]

    for i in range(100):
        x = sample['x_img_gt'].numpy()[i].T
        plt.scatter(x[:,0],x[:,1])
        plt.show()
    quit()

    # test training loader
    loader = SyntheticLoader()
    sample = loader[1]
    print("sample made")
    print(sample.keys())


from ortools.sat.python import cp_model
from SolutionPrinter import SolutionPrinter

def main(num_candy, num_people, min_candy, max_candy):
	model = cp_model.CpModel()

	cVar = [model.NewIntVar(0, num_candy, 'cVar_%i' % i) for i in range(1, num_people + 1)]
	model.Add(sum(cVar) == num_candy)
	for i in range(num_people):
		model.Add(cVar[i] >= min_candy)
		model.Add(cVar[i] <= max_candy)

	solver = cp_model.CpSolver()
	solution_printer = SolutionPrinter(cVar)
	status = solver.SearchForAllSolutions(model, solution_printer)
	print()
	print('Solutions found : %i' % solution_printer.SolutionCount())

if __name__ == "__main__":
	num_candy = 10
	num_people = 3
	min_candy = 2
	max_candy = 5
	main(num_candy, num_people, min_candy, max_candy)


def append(ol,ele,**kwargs):
    '''
        from elist.elist import *
        ol = [1,2,3,4]
        ele = 5
        id(ol)
        append(ol,ele,mode="original")
        ol
        id(ol)
        ####
        ol = [1,2,3,4]
        ele = 5
        id(ol)
        new = append(ol,ele)
        new
        id(new)
    '''
    if('mode' in kwargs):
        mode = kwargs["mode"]
    else:
        mode = "new"
    if(mode == "new"):
        new = copy.deepcopy(ol)
        new.append(ele)
        return(new)
    else:
        ol.append(ele)
        return(ol)


import math     #수학 함수 (일반 숫자)
import cmath    #켤레 복소수연산 함수 모듈 (복소수 전용)

#모듈을 사용하지 않고 부동소수점 자리 변경을 하는 방법
#print = format 1) "format"%(value)     2)"객체".format(value,,,,(자리수))      3)f""
#1. 문자열 사용 방법 = 출력 서식으로    %c : 한문자, %s: string, %d: 10진, %o: 8진, %x: 16진, %f: 부동소수
                                     # %5.1f-> %전체자리수.소수이하자리수f
number  = math.pi
decimal_places =2
formatter_number = "{:.{}f}".format(number, decimal_places) #str.format
print(formatter_number)

#2. round() 이용 = 수치값 변경
#number=3.141592
number=math.pi
decimal_places = 2
round_number = round(number, decimal_places)
print(round_number)     #3.14
print(math.pi, cmath.pi)    #3.141592653589793 3.141592653589793


"""
Record a realsense video.

When done, press ctrl-c once and give it a second to close up.
"""
import os
import pyrealsense2 as rs

if __name__ == '__main__':
    FILENAME = 'realsense_video'

    number = -1

    for filename in os.listdir('.'):
        if FILENAME in filename:
            file_number = int(filename.replace(FILENAME, '').replace('.bag', ''))

            number = max(number, file_number)

    FILENAME += f'{number + 1}.bag'

    ##
    pipeline = rs.pipeline()
    config = rs.config()

    config.enable_stream(rs.stream.depth, 0, 0, rs.format.z16, 0)
    config.enable_stream(rs.stream.color, 0, 0, rs.format.rgb8, 0)

    config.enable_record_to_file(FILENAME)

    ##
    pipeline.start(config)

    try:
        i = 0
        while True:
            i += 1
            frames = pipeline.wait_for_frames()

            depth_frame = frames.get_depth_frame()
            color_frame = frames.get_color_frame()

            if not depth_frame or not color_frame:
                print("Failed to capture image.")
                continue

    except Exception as e:
        print(e)

    pipeline.stop()


import os, shutil, fire
from typing import List, Optional
from pyfiglet import Figlet

def main(path: Optional[str] = None) -> None:
    __print_logo('url - desktop')
    fire.Fire(__convert)

def __convert(path: Optional[str] = None) -> None:
    if path is None:
        path = os.getcwd()

    try:
        url_list_path = __get_url_files_path(path)
        exit_if_no_files(url_list_path)
        url_filenames = __get_url_filenames_without_extension(url_list_path)
        old_url_folder_name = __create_old_url_folder(path, url_list_path)
        __copy_url_files_to_old_url_folder(
            path, url_list_path, old_url_folder_name)
        __change_url_to_desktop(url_list_path, url_filenames)
        print('\nSuccessful conversion!\n')
    except Exception as e:
        print(f'Error: {str(e)}')

def __print_logo(text_logo: str) -> None:
    figlet = Figlet(font='slant')
    print(figlet.renderText(text_logo))

def __get_url_files_path(path: str) -> List[str]:
    url_list_path = list()
    for file in os.listdir(path):
        if file.endswith(".url") or file.endswith(".URL"):
            url_list_path.append(os.path.join(path, file))
    return url_list_path

def exit_if_no_files(url_list_path: List[str]) -> None:
    if len(url_list_path) <= 0:
        print('No URL to convert!')
        exit()

def __get_url_filenames_without_extension(
    url_list_path: List[str]
) -> List[str]:
    url_filenames = list()
    for path in url_list_path:
        url_filenames.append(os.path.splitext(os.path.basename(path))[0])
    return url_filenames

def __create_old_url_folder(
    path: str,
    url_list_path: List[str]
) -> str:
    if path is None:
        path = os.getcwd()

    old_url_folder_name = 'old urls'
    if not os.path.isdir(os.path.join(path, old_url_folder_name)):
        os.mkdir(os.path.join(path, old_url_folder_name))
    return old_url_folder_name

def __copy_url_files_to_old_url_folder(
    path: str,
    url_list_path: List[str],
    old_url_folder_name: str
) -> None:
    for file_path in url_list_path:
        shutil.copy(file_path, os.path.join(path, old_url_folder_name))

def __change_url_to_desktop(
    url_list_path: List[str],
    url_filenames: List[str]
) -> None:
    for file_path, filename in zip(url_list_path, url_filenames):
        # Take the URL from file
        url_file = open(file_path, 'r')
        for i, line in enumerate(url_file):
            if i == 1:
                url = line
        url_file.close()

        # Verify if we got the URL
        if not url.startswith('URL'):
            break

        # Convert to DESKTOP
        url_file = open(file_path, 'w')
        url_file.write('[Desktop Entry]\n')
        url_file.write('Encoding=UTF-8\n')
        url_file.write(f'Name={filename}\n')
        url_file.write('Type=Link\n')
        url_file.write(url)
        url_file.write('Icon=text-html\n')
        url_file.close()

        # Rename the file from .url to .desktop
        os.rename(
            file_path, f'{os.path.join(os.path.dirname(file_path),filename)}.desktop')


from urllib import request

url = 'http://httpbin.org/get'
headers = {
    'User-Agent':'Opera/9.80 (Windows NT 6.1; U; zh-cn) Presto/2.9.168 Version/11.50'
}

#1.创建请阿牛对象
req = request.Request(url=url, headers=headers)

#2.获取响应对象
res = request.urlopen(req)

#3.读取响应内容

html = res.read().decode('utf-8')

print(html)


import sys

def input():
    return sys.stdin.readline().rstrip()

n,m=map(int,input().split(' '))

ans=[]

def solve():
    if(len(ans)==m):
        print(*ans)
        return
    for i in range(1,n+1):
        if i not in ans:
            ans.append(i)
            solve()
            ans.pop()

solve()


#
# @lc app=leetcode.cn id=35 lang=python3
#
# [35] 搜索插入位置
#

# @lc code=start

# version I
# Your runtime beats 96.97 % of python3 submissions(32ms)
# Your memory usage beats 32.48 % of python3 submissions (14 MB)
class Solution:
    def searchInsert(self, nums:list, target):
        left, right = 0, len(nums)-1
        while left <= right:
            mid = (left + right)//2
            if nums[mid] < target:
                left = mid + 1
            elif nums[mid] > target:
                right = mid - 1
            else:
                return mid
        nums.insert(left + 1, target)
        return left

# @lc code=end

# solution = Solution()
# l = [1,3,5,6]
# print(solution.searchInsert(l, 2))
# print(l)

35.搜索插入位置.py