import os
import numpy as np # noqa
import matplotlib as mpl
mpl.use('svg')  # noqa

from matplotlib import pyplot as plt
from sklearn.metrics import confusion_matrix

import sys
import json
import pickle
import torch
from torch import nn
# from torch import optim
# from torch.optim import lr_scheduler

from opts import parse_opts
from model import generate_model
from mean import get_mean, get_std
from spatial_transforms import (
    Compose, Normalize, Scale, CenterCrop, MultiScaleCornerCrop,
    MultiScaleRandomCrop, RandomHorizontalFlip, ToTensor, RGB2Gray, LowResolution)
from temporal_transforms import LoopPadding, TemporalCenterCrop
from spatio_temporal_transforms import Coded, Averaged, OneFrame, ToTemporal  # , ToRepeat
from target_transforms import ClassLabel  # , VideoID
from target_transforms import Compose as TargetCompos # noqa
from dataset import get_test_set
from utils import Logger
from utils import AverageMeter, accuracy
from glob import glob

def test_eval(data_loader, model, criterion, opt, logger, device):
    print('eval')

    model.eval()

    losses = AverageMeter()
    top1 = AverageMeter()
    top3 = AverageMeter()
    top5 = AverageMeter()

    inst_results = []
    inst_targets = []
    input_mean = []
    N = 100
    for i, (inputs, targets) in enumerate(data_loader):
        samples = inputs.to('cpu').detach().numpy()
        for j in range(opt.batch_size):
            n = i * opt.batch_size + j
            if n > N:
                break
            if j == len(samples):
                break
            if opt.compress in ["one", "avg"]:
                from PIL import Image
                save_file_path = os.path.join(opt.result_path, 'sample_%05d.png' % n)
                Image.fromarray(samples[j].reshape((112, 112)).astype(np.uint8)).resize((1120, 1120)).save(save_file_path)
            elif opt.compress in ["mask"]:
                from PIL import Image
                save_file_path = os.path.join(opt.result_path, 'sample_%05d.png' % n)
                Image.fromarray(np.clip(samples[j] * 3, 0, 255).reshape((112, 112)).astype(np.uint8)).resize((1120, 1120)).save(save_file_path)
            else:
                from utils import save_gif
                save_gif_path = os.path.join(opt.result_path, 'sample_%005d.gif' % n)
                save_gif(samples[j].reshape(16, 112, 112, 1).astype(np.uint8), save_gif_path, vmax=255, vmin=0, interval=2000/16)

        input_mean.extend([i.mean() for i in inputs.detach().cpu().numpy()])
        inputs = inputs.to(device)
        targets = targets.to(device)
        outputs = model(inputs)

        inst_results.append(outputs.detach().cpu().numpy())
        inst_targets.append(targets.detach().cpu().numpy())

        loss = criterion(outputs, targets)
        losses.update(loss.item(), inputs.size(0))

        prec1, prec3, prec5 = accuracy(outputs.data, targets, topk=(1, 3, 5))
        top1.update(prec1, inputs.size(0))
        top3.update(prec3, inputs.size(0))
        top5.update(prec5, inputs.size(0))

        sys.stdout.flush()
        sys.stdout.write('\rEVAL: [{:>6}/{:>6}] '
                         'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                         'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                         'Acc@3 {top3.val:.3f} ({top3.avg:.3f})\t'
                         'Acc@5 {top5.val:.3f} ({top5.avg:.3f})\t\t'
                         'len {len_mean},'
                         'mean {mean:.4f},'
                         'std {std:.4f},'
                         'min {min:.4f},'
                         'max {max:.4f}'
                         '\t\t'.format(
                             i + 1,
                             len(data_loader),
                             loss=losses,
                             top1=top1,
                             top3=top3,
                             top5=top5,
                             len_mean=len(input_mean),
                             mean=np.mean(input_mean),
                             std=np.std(input_mean),
                             min=np.min(input_mean),
                             max=np.max(input_mean),
                         ))
    print()
    logger.log({
        'top1': top1.avg,
        'top3': top3.avg,
        'top5': top5.avg,
    })

    res = {
        'inst_results': inst_results,
        'inst_targets': inst_targets,
        'losses': losses,
        'top1': top1,
        'top3': top3,
        'top5': top5,
    }
    with open(os.path.join(opt.result_path, 'res.pkl'), 'wb') as f:
        pickle.dump(res, f, protocol=pickle.HIGHEST_PROTOCOL)
    # with open(os.path.join(opt.result_path, 'res.pkl'), 'rb') as f:
    #     res = pickle.load(f)

    # inst_results = res['inst_results']
    # inst_targets = res['inst_targets']
    # losses = res['losses']
    # top1 = res['top1']
    # top3 = res['top3']
    # top5 = res['top5']

    inst_targets = np.concatenate(inst_targets, axis=0)
    print(inst_targets.shape)
    inst_results = np.concatenate(inst_results, axis=0)
    print(inst_results.shape)

    with open(opt.annotation_path, 'r') as f:
        label_to_id = json.load(f)
    labels = list(label_to_id.keys())

    res = []
    res3 = []
    res5 = []
    if os.path.basename(opt.video_path) in ["mask", "mask2", "mask3", "mask_old"]:
        npys = sorted(glob(os.path.join(opt.video_path, "*/clip[0-9]/npy/*.npy")),key=lambda x: x.split('/'))
        y_pred = inst_results.argmax(axis=1)
        y_true = inst_targets
        print(y_pred)
        print(y_true)

        prev_cid = None
        d = None
        d3 = None
        for i, cid in enumerate(y_true):
            if prev_cid != cid:
                prev_cid = cid
                if d is not None:
                    for k,v in d.items():
                        print(k, np.max(v), np.mean(v), np.array(v).argmax())
                        res.append(np.max(v))
                    for k,v in d3.items():
                        print(k, np.max(v), np.mean(v), np.array(v).argmax())
                        res3.append(np.max(v))
                    for k,v in d5.items():
                        print(k, np.max(v), np.mean(v), np.array(v).argmax())
                        res5.append(np.max(v))
                print("="*30)
                print(cid, labels[cid])
                d = {"clip%d" % j: [] for j in range(1, 5)}
                d3 = {"clip%d" % j: [] for j in range(1, 5)}
                d5 = {"clip%d" % j: [] for j in range(1, 5)}
            pred = y_pred[i]
            if "clip1" in npys[i]:
                d["clip1"].append(pred == cid)
                d3["clip1"].append(cid in inst_results[i].argsort()[-3:][::-1])
                d5["clip1"].append(cid in inst_results[i].argsort()[-5:][::-1])
            elif "clip2" in npys[i]:
                d["clip2"].append(pred == cid)
                d3["clip2"].append(cid in inst_results[i].argsort()[-3:][::-1])
                d5["clip2"].append(cid in inst_results[i].argsort()[-5:][::-1])
            elif "clip3" in npys[i]:
                d["clip3"].append(pred == cid)
                d3["clip3"].append(cid in inst_results[i].argsort()[-3:][::-1])
                d5["clip3"].append(cid in inst_results[i].argsort()[-5:][::-1])
            elif "clip4" in npys[i]:
                d["clip4"].append(pred == cid)
                d3["clip4"].append(cid in inst_results[i].argsort()[-3:][::-1])
                d5["clip4"].append(cid in inst_results[i].argsort()[-5:][::-1])
            y_pred_top5 = inst_results[i].argsort()[-5:][::-1]
            for k, pred_k in enumerate(y_pred_top5):
                print("%s: %s" % (k+1, labels[pred_k]))
            # if pred == cid:
            #     print(
            #         "    %s: ok" % npys[i].replace('../datasets/REAL/', ''),
            #         inst_results[i].argsort()[-10:][::-1]
            #     )
            # else:
            #     print(
            #         "    %s: %s" % (npys[i].replace('../datasets/REAL/', ''), labels[pred]),
            #         inst_results[i].argsort()[-10:][::-1]
            #     )
        else:
            for k,v in d.items():
                print(k, np.max(v), np.mean(v))
                res.append(np.max(v))
            for k,v in d3.items():
                print(k, np.max(v), np.mean(v))
                res3.append(np.max(v))
            for k,v in d5.items():
                print(k, np.max(v), np.mean(v))
                res5.append(np.max(v))
            print("="*30)
            print(len(res), np.mean(res))
            print(len(res3), np.mean(res3))
            print(len(res5), np.mean(res5))
    else:
        y_pred = np.array([[inst_results[i*4+j].argmax() for j in range(4)] for i in range(25)])
        y_pred_top5 = np.array([[inst_results[i*4+j].argsort()[-5:][::-1] for j in range(4)] for i in range(25)])
        y_true = np.array([[inst_targets[i*4+j] for j in range(4)] for i in range(25)])
        acc_class = (y_pred == y_true).mean(axis=1)
        acc_all = acc_class.mean(axis=0)
        print(y_pred)
        print(y_true)

        for i, cid in enumerate(y_true[:, 0]):
            print("%s: \t%.4f" % (labels[cid], acc_class[i]))
            for clip in range(4):
                print("clip:", clip)
                for k, pred_k in enumerate(y_pred_top5[i][clip]):
                    print("%s: %s" % (k+1, labels[pred_k]))
            # for pred in np.unique(y_pred[i]):
            #     if pred != cid:
            #         print("    *", labels[pred])
            print()

        print(acc_all)

    print('Loss {loss.avg:.4f}\t'
          'Acc@1 {top1.avg:.4f}\t'
          'Acc@3 {top3.avg:.4f}\t'
          'Acc@5 {top5.avg:.4f}\t'
          ''.format(
              loss=losses,
              top1=top1,
              top3=top3,
              top5=top5,
          ))

if __name__ == '__main__':
    opt = parse_opts()
    if opt.root_path != '':
        opt.video_path = os.path.join(opt.root_path, opt.video_path)
        opt.annotation_path = os.path.join(opt.root_path, opt.annotation_path)
        opt.result_path = os.path.join(opt.root_path, opt.result_path)
        if opt.resume_path:
            opt.resume_path = os.path.join(opt.root_path, opt.resume_path)
        if opt.pretrain_path:
            opt.pretrain_path = os.path.join(opt.root_path, opt.pretrain_path)
    opt.scales = [opt.initial_scale]
    for i in range(1, opt.n_scales):
        opt.scales.append(opt.scales[-1] * opt.scale_step)
    opt.arch = '{}-{}'.format(opt.model, opt.model_depth)
    opt.mean = get_mean(opt.norm_value, dataset=opt.mean_dataset)
    opt.std = get_std(opt.norm_value)
    print(opt)
    if opt.dataset != 'real':
        raise
    # with open(os.path.join(opt.result_path, 'opts.json'), 'w') as opt_file:
    #     json.dump(vars(opt), opt_file, indent=2)
    if not opt.resume_path and not opt.init_path:
        raise

    torch.manual_seed(opt.manual_seed)

    model, parameters = generate_model(opt)
    print(model)
    from torch.backends import cudnn
    cudnn.benchmark = True
    criterion = nn.CrossEntropyLoss()
    device = torch.device("cpu" if opt.no_cuda else "cuda")
    if not opt.no_cuda:
        criterion = criterion.to(device)

    if opt.no_mean_norm and not opt.std_norm:
        norm_method = Normalize([0, 0, 0], [1, 1, 1])
    elif not opt.std_norm:
        norm_method = Normalize(opt.mean, [1, 1, 1])
    else:
        norm_method = Normalize(opt.mean, opt.std)

    if opt.train_crop == 'random':
        crop_method = MultiScaleRandomCrop(opt.scales, opt.sample_size)
    elif opt.train_crop == 'corner':
        crop_method = MultiScaleCornerCrop(opt.scales, opt.sample_size)
    elif opt.train_crop == 'center':
        crop_method = MultiScaleCornerCrop(
            opt.scales, opt.sample_size, crop_positions=['c'])

    if not opt.no_val:
        spatial_transform = Compose([
            ToTensor(opt.norm_value), norm_method,
        ])
        temporal_transform = LoopPadding(opt.sample_duration)
        target_transform = ClassLabel()
        if opt.compress == 'mask':
            spatio_temporal_transform = None
            temporal_transform = None
        elif opt.compress == 'avg':
            spatio_temporal_transform = Averaged()
        elif opt.compress == 'one':
            spatio_temporal_transform = OneFrame()
        elif opt.compress == 'spatial':
            spatial_transform = Compose([
                LowResolution(opt.spatial_compress_size, use_cv2=opt.use_cv2),
                ToTensor(opt.norm_value), norm_method,
            ])
            spatio_temporal_transform = None
            temporal_transform = TemporalCenterCrop(opt.sample_duration)
        else:
            spatio_temporal_transform = None
            temporal_transform = TemporalCenterCrop(opt.sample_duration)
        test_data = get_test_set(
            opt, spatial_transform, temporal_transform, target_transform,
            spatio_temporal_transform)
        test_loader = torch.utils.data.DataLoader(
            test_data,
            batch_size=opt.batch_size,
            shuffle=False,
            num_workers=opt.n_threads,
            pin_memory=True)
        test_logger = Logger(
            os.path.join(opt.result_path, 'test_.log'),
            ['top1', 'top3', 'top5'])

    if opt.resume_path:
        print('loading checkpoint {}'.format(opt.resume_path))
        checkpoint = torch.load(opt.resume_path)
        assert opt.arch == checkpoint['arch']

        opt.begin_epoch = checkpoint['epoch']
        model.load_state_dict(checkpoint['state_dict'])
        # if not opt.no_train:
        #     optimizer.load_state_dict(checkpoint['optimizer'])
    if opt.init_path:
        check = False
        print('initilize checkpoint {}'.format(opt.init_path))
        checkpoint = torch.load(opt.init_path)

        for i, (name, param) in enumerate(model.named_parameters()):
            for (prev_name, prev_param) in checkpoint['state_dict'].items():
                if name == prev_name:
                    param.data = prev_param.data
                    print('{}: {}({}) -> {}({})'.format(
                        i, prev_name, prev_param.shape, name, param.shape))
                    check = True
        if not check:
            raise

    model.to(device)
    test_eval(test_loader, model, criterion, opt, test_logger, device)


#!/usr/bin/python3
#
# Copyright 2020 Dustin Kleckner
#
# 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 numpy as np
import numba
import struct
from threading import Lock

FRACTION_MASK = (2**32-1)
MAX_INT = 2**32

uint8_t = 'B'
uint16_t = 'H'
uint32_t = 'I'
uint64_t = 'Q'
int16_t = 'h'
int32_t = 'l'
bool32_t = 'i'
float_t = 'f'
double_t = 'd'
TIME64 = 'Q'
RECT = '4i'
WBGAIN = '2f'
IMFILTER = '28i'

MAXLENDESCRIPTION = 4096
MAXLENDESCRIPTION_OLD = 121
OLDMAXFILENAME = 65

T64_F = lambda x: int(x) / 2.**32
T64_F_ms = lambda x: '%.3f' % (float(x.rstrip('L')) / 2.**32)
T64_S = lambda s: lambda t: time.strftime(s, time.localtime(float(t.rstrip('L'))/2.**32))

@numba.jit(nopython=True, parallel=True, cache=True)
def unpack_10b(input, tone_map, output, offset=0, stride=1):
    # 10 bits data => 4 points packed into 5 bytes
    for block in numba.prange(len(input) // 5):
        i = block * 5
        val0 = (numba.uint16(input[i+0] & 0xFF) << 2) + (numba.uint16(input[i+1]) >> 6)
        val1 = (numba.uint16(input[i+1] & 0x3F) << 4) + (numba.uint16(input[i+2]) >> 4)
        val2 = (numba.uint16(input[i+2] & 0x0F) << 6) + (numba.uint16(input[i+3]) >> 2)
        val3 = (numba.uint16(input[i+3] & 0x03) << 8) + (numba.uint16(input[i+4]) >> 0)

        j = offset + (block * 4 * stride)
        output[j           ] = tone_map[val0]
        output[j +   stride] = tone_map[val1]
        output[j + 2*stride] = tone_map[val2]
        output[j + 3*stride] = tone_map[val3]

# #Processing the data in chunks keeps it in the L2 catch of the processor, increasing speed for large arrays by ~50%
# CHUNK_SIZE = 6 * 10**5 #Should be divisible by 3, 4 and 5!  This seems to be near-optimal.
#
# def ten2sixteen(a):
#     b = np.zeros(a.size//5*4, dtype='u2')
#
#     for j in range(0, len(a), CHUNK_SIZE):
#         (a0, a1, a2, a3, a4) = [a[j+i:j+CHUNK_SIZE:5].astype('u2') for i in range(5)]
#
#         k = j//5 * 4
#         k2 = k + CHUNK_SIZE//5 * 4
#
#         b[k+0:k2:4] = ((a0 & 0b11111111) << 2) + ((a1 & 0b11000000) >> 6)
#         b[k+1:k2:4] = ((a1 & 0b00111111) << 4) + ((a2 & 0b11110000) >> 4)
#         b[k+2:k2:4] = ((a2 & 0b00001111) << 6) + ((a3 & 0b11111100) >> 2)
#         b[k+3:k2:4] = ((a3 & 0b00000011) << 8) + ((a4 & 0b11111111) >> 0)
#
#     return b

# LUT used to convert 10 bit packed values to 16 bit.
LinLUT = np.array([
      2,    5,    6,    7,    8,    9,   10,   11,
     12,   13,   14,   15,   16,   17,   17,   18,
     19,   20,   21,   22,   23,   24,   25,   26,
     27,   28,   29,   30,   31,   32,   33,   33,
     34,   35,   36,   37,   38,   39,   40,   41,
     42,   43,   44,   45,   46,   47,   48,   48,
     49,   50,   51,   52,   53,   54,   55,   56,
     57,   58,   59,   60,   61,   62,   63,   63,
     64,   65,   66,   67,   68,   69,   70,   71,
     72,   73,   74,   75,   76,   77,   78,   79,
     79,   80,   81,   82,   83,   84,   85,   86,
     87,   88,   89,   90,   91,   92,   93,   94,
     94,   95,   96,   97,   98,   99,  100,  101,
    102,  103,  104,  105,  106,  107,  108,  109,
    110,  110,  111,  112,  113,  114,  115,  116,
    117,  118,  119,  120,  121,  122,  123,  124,
    125,  125,  126,  127,  128,  129,  130,  131,
    132,  133,  134,  135,  136,  137,  137,  138,
    139,  140,  141,  142,  143,  144,  145,  146,
    147,  148,  149,  150,  151,  152,  153,  154,
    156,  157,  158,  159,  160,  161,  162,  163,
    164,  165,  167,  168,  169,  170,  171,  172,
    173,  175,  176,  177,  178,  179,  181,  182,
    183,  184,  186,  187,  188,  189,  191,  192,
    193,  194,  196,  197,  198,  200,  201,  202,
    204,  205,  206,  208,  209,  210,  212,  213,
    215,  216,  217,  219,  220,  222,  223,  225,
    226,  227,  229,  230,  232,  233,  235,  236,
    238,  239,  241,  242,  244,  245,  247,  249,
    250,  252,  253,  255,  257,  258,  260,  261,
    263,  265,  266,  268,  270,  271,  273,  275,
    276,  278,  280,  281,  283,  285,  287,  288,
    290,  292,  294,  295,  297,  299,  301,  302,
    304,  306,  308,  310,  312,  313,  315,  317,
    319,  321,  323,  325,  327,  328,  330,  332,
    334,  336,  338,  340,  342,  344,  346,  348,
    350,  352,  354,  356,  358,  360,  362,  364,
    366,  368,  370,  372,  374,  377,  379,  381,
    383,  385,  387,  389,  391,  394,  396,  398,
    400,  402,  404,  407,  409,  411,  413,  416,
    418,  420,  422,  425,  427,  429,  431,  434,
    436,  438,  441,  443,  445,  448,  450,  452,
    455,  457,  459,  462,  464,  467,  469,  472,
    474,  476,  479,  481,  484,  486,  489,  491,
    494,  496,  499,  501,  504,  506,  509,  511,
    514,  517,  519,  522,  524,  527,  529,  532,
    535,  537,  540,  543,  545,  548,  551,  553,
    556,  559,  561,  564,  567,  570,  572,  575,
    578,  581,  583,  586,  589,  592,  594,  597,
    600,  603,  606,  609,  611,  614,  617,  620,
    623,  626,  629,  632,  635,  637,  640,  643,
    646,  649,  652,  655,  658,  661,  664,  667,
    670,  673,  676,  679,  682,  685,  688,  691,
    694,  698,  701,  704,  707,  710,  713,  716,
    719,  722,  726,  729,  732,  735,  738,  742,
    745,  748,  751,  754,  758,  761,  764,  767,
    771,  774,  777,  781,  784,  787,  790,  794,
    797,  800,  804,  807,  811,  814,  817,  821,
    824,  828,  831,  834,  838,  841,  845,  848,
    852,  855,  859,  862,  866,  869,  873,  876,
    880,  883,  887,  890,  894,  898,  901,  905,
    908,  912,  916,  919,  923,  927,  930,  934,
    938,  941,  945,  949,  952,  956,  960,  964,
    967,  971,  975,  979,  982,  986,  990,  994,
    998, 1001, 1005, 1009, 1013, 1017, 1021, 1025,
   1028, 1032, 1036, 1040, 1044, 1048, 1052, 1056,
   1060, 1064, 1068, 1072, 1076, 1080, 1084, 1088,
   1092, 1096, 1100, 1104, 1108, 1112, 1116, 1120,
   1124, 1128, 1132, 1137, 1141, 1145, 1149, 1153,
   1157, 1162, 1166, 1170, 1174, 1178, 1183, 1187,
   1191, 1195, 1200, 1204, 1208, 1212, 1217, 1221,
   1225, 1230, 1234, 1238, 1243, 1247, 1251, 1256,
   1260, 1264, 1269, 1273, 1278, 1282, 1287, 1291,
   1295, 1300, 1304, 1309, 1313, 1318, 1322, 1327,
   1331, 1336, 1340, 1345, 1350, 1354, 1359, 1363,
   1368, 1372, 1377, 1382, 1386, 1391, 1396, 1400,
   1405, 1410, 1414, 1419, 1424, 1428, 1433, 1438,
   1443, 1447, 1452, 1457, 1462, 1466, 1471, 1476,
   1481, 1486, 1490, 1495, 1500, 1505, 1510, 1515,
   1520, 1524, 1529, 1534, 1539, 1544, 1549, 1554,
   1559, 1564, 1569, 1574, 1579, 1584, 1589, 1594,
   1599, 1604, 1609, 1614, 1619, 1624, 1629, 1634,
   1639, 1644, 1649, 1655, 1660, 1665, 1670, 1675,
   1680, 1685, 1691, 1696, 1701, 1706, 1711, 1717,
   1722, 1727, 1732, 1738, 1743, 1748, 1753, 1759,
   1764, 1769, 1775, 1780, 1785, 1791, 1796, 1801,
   1807, 1812, 1818, 1823, 1828, 1834, 1839, 1845,
   1850, 1856, 1861, 1867, 1872, 1878, 1883, 1889,
   1894, 1900, 1905, 1911, 1916, 1922, 1927, 1933,
   1939, 1944, 1950, 1956, 1961, 1967, 1972, 1978,
   1984, 1989, 1995, 2001, 2007, 2012, 2018, 2024,
   2030, 2035, 2041, 2047, 2053, 2058, 2064, 2070,
   2076, 2082, 2087, 2093, 2099, 2105, 2111, 2117,
   2123, 2129, 2135, 2140, 2146, 2152, 2158, 2164,
   2170, 2176, 2182, 2188, 2194, 2200, 2206, 2212,
   2218, 2224, 2231, 2237, 2243, 2249, 2255, 2261,
   2267, 2273, 2279, 2286, 2292, 2298, 2304, 2310,
   2317, 2323, 2329, 2335, 2341, 2348, 2354, 2360,
   2366, 2373, 2379, 2385, 2392, 2398, 2404, 2411,
   2417, 2423, 2430, 2436, 2443, 2449, 2455, 2462,
   2468, 2475, 2481, 2488, 2494, 2501, 2507, 2514,
   2520, 2527, 2533, 2540, 2546, 2553, 2559, 2566,
   2572, 2579, 2586, 2592, 2599, 2605, 2612, 2619,
   2625, 2632, 2639, 2645, 2652, 2659, 2666, 2672,
   2679, 2686, 2693, 2699, 2706, 2713, 2720, 2726,
   2733, 2740, 2747, 2754, 2761, 2767, 2774, 2781,
   2788, 2795, 2802, 2809, 2816, 2823, 2830, 2837,
   2844, 2850, 2857, 2864, 2871, 2878, 2885, 2893,
   2900, 2907, 2914, 2921, 2928, 2935, 2942, 2949,
   2956, 2963, 2970, 2978, 2985, 2992, 2999, 3006,
   3013, 3021, 3028, 3035, 3042, 3049, 3057, 3064,
   3071, 3078, 3086, 3093, 3100, 3108, 3115, 3122,
   3130, 3137, 3144, 3152, 3159, 3166, 3174, 3181,
   3189, 3196, 3204, 3211, 3218, 3226, 3233, 3241,
   3248, 3256, 3263, 3271, 3278, 3286, 3294, 3301,
   3309, 3316, 3324, 3331, 3339, 3347, 3354, 3362,
   3370, 3377, 3385, 3393, 3400, 3408, 3416, 3423,
   3431, 3439, 3447, 3454, 3462, 3470, 3478, 3486,
   3493, 3501, 3509, 3517, 3525, 3533, 3540, 3548,
   3556, 3564, 3572, 3580, 3588, 3596, 3604, 3612,
   3620, 3628, 3636, 3644, 3652, 3660, 3668, 3676,
   3684, 3692, 3700, 3708, 3716, 3724, 3732, 3740,
   3749, 3757, 3765, 3773, 3781, 3789, 3798, 3806,
   3814, 3822, 3830, 3839, 3847, 3855, 3863, 3872,
   3880, 3888, 3897, 3905, 3913, 3922, 3930, 3938,
   3947, 3955, 3963, 3972, 3980, 3989, 3997, 4006,
   4014, 4022, 4031, 4039, 4048, 4056, 4064, 4095,
   4095, 4095, 4095, 4095, 4095, 4095, 4095, 4095
], dtype='u2')

class Cine:
    _MUVI_SUPPORTS_TONE_MAP = True

    _HEADER_FIELDS = (
        ('Type', '2s'),
        ('HeaderSize', uint16_t),
        ('Compression', uint16_t),
        ('Version', uint16_t),
        ('FirstMovieImage', int32_t),
        ('TotalImageCount', uint32_t),
        ('FirstImageNo', int32_t),
        ('ImageCount', uint32_t),
        ('OffImageHeader', uint32_t),
        ('OffSetup', uint32_t),
        ('OffImageOffsets', uint32_t),
        ('TriggerTime', TIME64),
    )

    _BITMAP_INFO_FIELDS = (
        ('Size', uint32_t),
        ('Width', int32_t),
        ('Height', int32_t),
        ('Planes', uint16_t),
        ('BitCount', uint16_t),
        ('Compression', uint32_t),
        ('SizeImage', uint32_t),
        ('XPelsPerMeter', int32_t),
        ('YPelsPerMeter', int32_t),
        ('ClrUsed', uint32_t),
        ('ClrImportant', uint32_t),
    )

    _SETUP_FIELDS = (
        ('FrameRate16', uint16_t),
        ('Shutter16', uint16_t),
        ('PostTrigger16', uint16_t),
        ('FrameDelay16', uint16_t),
        ('AspectRatio', uint16_t),
        ('Res7', uint16_t),
        ('Res8', uint16_t),
        ('Res9', uint8_t),
        ('Res10', uint8_t),
        ('Res11', uint8_t),
        ('TrigFrame', uint8_t),
        ('Res12', uint8_t),
        ('DescriptionOld', str(MAXLENDESCRIPTION_OLD) + 's'),
        ('Mark', uint16_t),
        ('Length', uint16_t),
        ('Res13', uint16_t),
        ('SigOption', uint16_t),
        ('BinChannels', int16_t),
        ('SamplesPerImage', uint8_t),
        ('BinName', 8 * '11s'),
        ('AnaOption', uint16_t),
        ('AnaChannels', int16_t),
        ('Res6', uint8_t),
        ('AnaBoard', uint8_t),
        ('ChOption', 8 * int16_t),
        ('AnaGain', 8 *  float_t),
        ('AnaUnit', 8 * '6s'),
        ('AnaName', 8 * '11s'),
        ('lFirstImage', int32_t),
        ('dwImageCount', uint32_t),
        ('nQFactor', int16_t),
        ('wCineFileType', uint16_t),
        ('szCinePath', 4 * (str(OLDMAXFILENAME) + 's')),
        ('Res14', uint16_t),
        ('Res15', uint8_t),
        ('Res16', uint8_t),
        ('Res17', uint16_t),
        ('Res18', double_t),
        ('Res19', double_t),
        ('Res20', uint16_t),
        ('Res1', int32_t),
        ('Res2', int32_t),
        ('Res3', int32_t),
        ('ImWidth', uint16_t),
        ('ImHeight', uint16_t),
        ('EDRShutter16', uint16_t),
        ('Serial', uint32_t),
        ('Saturation', int32_t),
        ('Res5', uint8_t),
        ('AutoExposure', uint32_t),
        ('bFlipH', bool32_t),
        ('bFlipV', bool32_t),
        ('Grid', uint32_t),
        ('FrameRate', uint32_t),
        ('Shutter', uint32_t),
        ('EDRShutter', uint32_t),
        ('PostTrigger', uint32_t),
        ('FrameDelay', uint32_t),
        ('bEnableColor', bool32_t),
        ('CameraVersion', uint32_t),
        ('FirmwareVersion', uint32_t),
        ('SoftwareVersion', uint32_t),
        ('RecordingTimeZone', int32_t),
        ('CFA', uint32_t),
        ('Bright', int32_t),
        ('Contrast', int32_t),
        ('Gamma', int32_t),
        ('Res21', uint32_t),
        ('AutoExpLevel', uint32_t),
        ('AutoExpSpeed', uint32_t),
        ('AutoExpRect', RECT),
        ('WBGgain', 4 * WBGAIN),
        ('Rotate', int32_t),
        ('WBView', WBGAIN),
        ('RealBPP', uint32_t),
        ('Conv8Min', uint32_t),
        ('Conv8Max', uint32_t),
        ('FilterCode', int32_t),
        ('FilterParam', int32_t),
        ('UF', IMFILTER),
        ('BlackCalSver', uint32_t),
        ('WhiteCalSver', uint32_t),
        ('GrayCalSver', uint32_t),
        ('bStampTime', bool32_t),
        ('SoundDest', uint32_t),
        ('FRPSteps', uint32_t),
        ('FRPImgNr', 16 * int32_t),
        ('FRPRate', 16 * uint32_t),
        ('FRPExp', 16 * uint32_t),
        ('FRPImageNr', 16 * int32_t),
        ('FRPRate', 16 * uint32_t),
        ('FRPExp', 16 * uint32_t),
        ('MCCnt', int32_t),
        ('MCPercent', 64 * float_t),
        ('CICalib', uint32_t),
        ('CalibWidth', uint32_t),
        ('CalibHeight', uint32_t),
        ('CalibRate', uint32_t),
        ('CalibExp', uint32_t),
        ('CalibEDR', uint32_t),
        ('CalibTemp', uint32_t),
        ('HeadSerial', 4 * uint32_t),
        ('RangeCode', uint32_t),
        ('RangeSize', uint32_t),
        ('Decimation', uint32_t),
        ('MasterSerial', uint32_t),
        ('Sensor', uint32_t),
        ('ShutterNs', uint32_t),
        ('EDRShutterNs', uint32_t),
        ('FrameDelayNs', uint32_t),
        ('ImPosXacq', uint32_t),
        ('ImPosYacq', uint32_t),
        ('ImWidthAcq', uint32_t),
        ('ImHeightAcq', uint32_t),
        ('Description', str(MAXLENDESCRIPTION) + 's'),
        ('RisingEdge', bool32_t),
        ('FilterTime', uint32_t),
        ('LongReady', bool32_t),
        ('ShutterOff', bool32_t),
        ('Res4', '16s'),
        ('bMetaWB', bool32_t),
        ('Hue', int32_t),
        ('BlackLevel', int32_t),
        ('WhiteLevel', int32_t),
        ('LensDescription', '256s'),
        ('LensAperture', float_t),
        ('LensFocusDistance', float_t),
        ('LensFocalLength', float_t),
        ('fOffset', float_t),
        ('fGain', float_t),
        ('fSaturation', float_t),
        ('fHue', float_t),
        ('fGamma', float_t),
        ('fGammaR', float_t),
        ('fGammaB', float_t),
        ('fFlat', float_t),
        ('fPedestalR', float_t),
        ('fPedestalG', float_t),
        ('fPedestalB', float_t),
        ('fChroma', float_t),
        ('ToneLabel', '256s'),
        ('TonePoints', int32_t),
        ('fTone', 4 * float_t),
    )

    def __init__(self, filename, output_bits=8, black_level=64,
                 white_level=4064, gamma=1.0, dark_clip=0.0):
        '''Open a Vision Research Cine file.

        This reader does not return the raw values in the file, but rather
        converts it to normalized values in a user-defined way.

        Parameters
        ----------
        filename : str

        Keywords
        --------
        output_bits : int (default: 8)
            The number of bits in the output array.  Should be 8 or 16
        black_level : int (default: 64)
            The level in the decoded cine file which is considered "black";
            this value will be converted to 0 in the output.  This value is
            standard for Phantom cameras, and should not need to be changed.
        white_level : int (default: 4064)
            The level in the decoded cine file which is considered "white";
            this value will be converted to 2^bits - 1 in the output.  This
            value is standard for Phantom cameras, and should not need to be
            changed.
        gamma : float (default: 1.0)
            The gamma value of the extracted images.  The actual intensity of
            the pixel is proportional to (output value)^(gamma)
        dark_clip : float (default: 0.0)
            Relative brightnesses below this value are converted to 0 in the
            output.  In other words, the clip value in the raw file is given
            by (black_level * (1 - dark_clip) + white_level * dark_clip).
            Note that this clip is applied *before* gamma correction!
        '''
        self._file = open(filename, 'rb')
        self.filename = filename
        self.black_level = black_level
        self.white_level = white_level
        self.gamma = gamma
        self.dark_clip = dark_clip

        # Read the headers
        self.header = self._read_header(self._HEADER_FIELDS)
        if self.header['Version'] != 1:
            raise ValueError('version of CINE file should be 1 (found %d)' % self.header['Version'])

        self._file.seek(self.header['OffImageHeader'])
        self.bi = self._read_header(self._BITMAP_INFO_FIELDS)

        self._file.seek(self.header['OffSetup'])
        self.setup = self._read_header(self._SETUP_FIELDS)

        # Find the image locations
        self._file.seek(self.header['OffImageOffsets'])
        self.image_locations = self._unpack(self.header['ImageCount'] * uint64_t)

        # For convenience
        self.width = self.bi['Width']
        self.height = self.bi['Height']
        self.internal_bit_depth = self.setup['RealBPP']
        self.image_bytes = self.bi['SizeImage']

        # This reader supports 10 bit packed... so lets see if thats what we
        #   have here.
        if self.bi['Compression'] != 256:
            raise ValueError('this library only supports 10 bit packed CINE files')

        # Allows Cine object to be accessed from multiple threads
        self.file_lock = Lock()

        # Set up tone map and output info
        if output_bits == 8:
            self.output_type = 'u1'
        elif output_bits == 16:
            self.output_type = 'u2'
        else:
            raise ValueError('output_bits shoud be 8 or 16 (found %s)' % output_bits)

        output_max = 2**output_bits - 1

        rel_map = (LinLUT.astype('d') - black_level) / (white_level - black_level)
        rel_map[np.where(rel_map < dark_clip)] = 0
        self.tone_map = (np.clip(rel_map, 0.0, 1.0)**(1./gamma) * output_max + 0.5).astype(self.output_type)

    def _unpack(self, format):
        s = struct.Struct('<' + format)
        vals = s.unpack(self._file.read(s.size))
        return vals[0] if len(vals) == 1 else vals

    def _read_header(self, fields, offset=None):
        return {name: self._unpack(format) for name, format in fields}

    def __getitem__(self, key):
        if type(key) == slice:
            return [self.get_frame(i) for i in range(len(self))[key] ]
        else:
            if key < 0:
                key = len(self) + key
            return self.get_frame(key)

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

    def get_frame(self, i, raw=False):
        '''Get a single frame from the CINE

        Parameters
        ----------
        i : int
            Frame number

        Keywords
        --------
        raw : bool (default: False)
            If true, don't do the tone mapping, just return the raw data

        Returns
        -------
        frame : numpy array
            The tone mapped frame
        '''

        if self.internal_bit_depth == 10:
            self.file_lock.acquire()

            self._file.seek(self.image_locations[i])

            annotation_size = self._unpack(uint32_t)
            annotation = self._unpack('%ds' % (annotation_size - 8))
            image_size = self._unpack(uint32_t)

            raw_frame = np.frombuffer(self._file.read(self.image_bytes), dtype='u1')
            self.file_lock.release()

            if raw:
                return raw_frame
            else:
                frame = np.empty((self.height, self.width), dtype=self.output_type)
                unpack_10b(raw_frame, self.tone_map, frame.reshape(-1))
                return frame

        else:
            raise ValueError('this library only supports 10 bit packed files')

    def pack_into_array(self, i, arr, offset=0, stride=1):
        '''Get a single frame from the Cine, and pack into an existing array.

        Parameters
        ----------
        i : int
            Frame number
        arr : numpy array
            Array in which to pack

        Keywords
        --------
        offset : int (default: 0)
            The integer offset to the start of the packing.  The array is
            flattened before accessing, so this is always a single number
        stride : int (default: 1)
            The spacing between values in the array.  Usually 1, but may be
            different if, for example, you are packing data into a single
            color channel of a multichannel image.
        '''
        if self.internal_bit_depth == 10:
            self.file_lock.acquire()

            self._file.seek(self.image_locations[i])

            annotation_size = self._unpack(uint32_t)
            annotation = self._unpack('%ds' % (annotation_size - 8))
            image_size = self._unpack(uint32_t)

            raw_frame = np.frombuffer(self._file.read(self.image_bytes), dtype='u1')

            self.file_lock.release()

            unpack_10b(raw_frame, self.tone_map, arr.reshape(-1), offset=offset, stride=stride)

        else:
            raise ValueError('this library only supports 10 bit packed files')

    def close(self):
        if hasattr(self, 'f'):
            self._file.close()
        del self._file

    def get_frame_times(self):
        '''Read the image time block from the Cine file.  Will raise an error
        if the timestamp block isn't found.  (This shouldn't be a problem
        with files saved from recent versions of the software.)

        Returns
        -------
        timestamps : array of 64 bit unsigned integers
            The timestamp of each frame.  First 32 bits are the unix time stamp
            in seconds; final 32 bits are the fractional seconds bit.  (I.e.
            the whole timestamp is a fixed point number with 32 decimal bits.)
        '''
        self.file_lock.acquire()

        self._file.seek(self.header['OffSetup'] + self.setup['Length'])

        while self._file.tell() < self.header['OffImageOffsets']:
            block_size, block_type, res = struct.unpack('<LHH', self._file.read(8))
            # print(block_size, type)

            dat = self._file.read(block_size-8)
            if block_type == 1002: # Time stamp block!
                self.file_lock.release()
                return np.fromstring(dat, dtype='u8')
        else:
            self.file_lock.release()
            raise ValueError("Cine file does not contain timestamp block!")

    def __del__(self):
        self._file.close()

if __name__ == "__main__":
    import sys
    cine = Cine(sys.argv[1], dark_clip=10E-3, gamma=2.0, output_bits=8)

    N = 128
    frame_offset = 560*10
    i0 = 75200

    channels = 3
    vol = np.empty((N, cine.height, cine.width, channels), dtype='u1')
    frame_size = np.prod(vol.shape[1:])
    print(frame_size)

    for channel in range(3):
        start = i0 + channel * frame_offset
        for z in range(N):
            # vol[i] = cine[i + start]
            cine.pack_into_array(start + z, vol, offset=z*frame_size+channel, stride=channels)

    from muvi import VolumetricMovie
    from muvi.view.qtview import view_volume

    vm = VolumetricMovie([vol])
    vm.save('test.vti')

    view_volume(vm)

    # frames = np.array(cine[32000:33000])
    # print(frames.shape, frames.dtype)
    # mi = frames.mean((1, 2))
    # frame = frames[np.argmax(mi)]
    #
    # import matplotlib.pyplot as plt
    # import time
    #
    # print('Fraction empty: %.2f%%' % ((frame==0).mean()*100))
    #
    #
    # plt.subplot(121)
    # plt.imshow(frame)
    # plt.colorbar()
    #
    # plt.subplot(122)
    # plt.hist(frame.reshape(-1), bins=np.arange(50))
    # plt.yscale('log')
    # plt.show()

    # for header, name in (('header', 'Header'), ('bi', 'Bitmap Info'), ('setup', 'Camera Setup')):
    #     print('-'*79)
    #     print(name)
    #     print('-'*79)
    #
    #     for k, v in getattr(cine, header).items():
    #         v = str(v)
    #         if len(v) > 59:
    #             v = v[:56] + '...'
    #         print('%18s: %s' % (k, v))
    #
    # print(cine.image_locations[:20])


import os
import glob
import numpy as np
import librosa
import soundfile as sf

def generate_mixtures(data_dir, sampling_rate, n):
    speakers_dirs = [os.path.join(data_dir, i) for i in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, i))]
    speakers_dirs.sort()
    print("# of mixes: %d" % n)
    print("\t with sampling rate: %d" % sampling_rate)
    print("\t out of %d speakers" % len(speakers_dirs))

    mixes = []
    reference = []
    labels = []

    def augment(speech):
        fac = np.random.uniform(-3, 3)
        return 10. ** (fac / 20) * speech

    def get_speaker_audios(speaker_dir):
        audio_files = glob.glob(os.path.join(
            speaker_dir, "**/*.wav"), recursive=True)
        return sorted(audio_files)

    for j in range(n):
        speaker_dir = np.random.choice(speakers_dirs)
        other_speaker_dir = np.random.choice(speakers_dirs)

        while(speaker_dir == other_speaker_dir):
            other_speaker_dir = np.random.choice(speakers_dirs)

        audio_file = np.random.choice(get_speaker_audios(speaker_dir))
        other_audio_file = np.random.choice(get_speaker_audios(other_speaker_dir))

        y1, sampling_rate1 = sf.read(audio_file, dtype='float32')
        y2, sampling_rate2 = sf.read(other_audio_file, dtype='float32')

        if sampling_rate1 != sampling_rate or sampling_rate2 != sampling_rate:
            raise Exception('Invalid sampling rate')

        length = min(len(y1), len(y2))
        #y1 = augment(y1[:length])
        #y2 = augment(y2[:length])
        y1 = y1[:length]
        y2 = y2[:length]
        mix = y1 + y2

        speaker_names = (os.path.basename(speaker_dir), os.path.basename(other_speaker_dir))

        reference.append([y1, y2])
        mixes.append(mix)

        labels.append(speaker_names)
    return mixes, reference, labels


import pandas as pd
import numpy as np
import scipy as sp
import os
import exa
import exatomic
from exatomic import qe
from neighbors_input import *
from exatomic.algorithms import neighbors
from exatomic.algorithms.neighbors import periodic_nearest_neighbors_by_atom    # Only valid for simple cubic periodic cells
from nuc import *

def My_fn(x,analyte,symbols):
    solv = [s for s in list(set(symbols)) if s is not analyte]
    sym = x.split()[0]
    if sym == analyte:
        return 0
    else:
        for i,s in enumerate(solv):
            if sym == s:
                return i+1

def gen_inputs(traj_dir=traj_dir):
    nuc_df=pd.DataFrame.from_dict(nuc)
    trajs = [t.name for t in os.scandir(traj_dir) if t.name.isnumeric()]
    print("Parsing "+ str(len(trajs)) + " trajectories...")
    for traj in trajs:
        print("Trajectory "+traj+'...')
        pos_dir = traj_dir + traj + '/'
        xyzs = pos_dir + "xyzs/"
        pos = list(filter(lambda x: ".pos" in x, os.listdir(pos_dir)))[0]
        atom = qe.parse_xyz(pos_dir+'/'+pos,symbols=symbols)
        atom['frame'] = atom['frame'].astype(int)
        atom = atom[(atom['frame']>=atom.frame[0] + start_prod) & (atom['frame']%sample_freq==0)]

        u = exatomic.Universe()
        u.atom = atom

        # Add the unit cell dimensions to the frame table of the universe
        for i, q in enumerate(("x", "y", "z")):
            for j, r in enumerate(("i", "j", "k")):
                if i == j:
                    u.frame[q+r] = celldm
                else:
                    u.frame[q+r] = 0.0
            u.frame["o"+q] = 0.0
        u.frame['periodic'] = True

        grouped = u.atom.groupby('frame')

        if write_ADF:
            ADF_dir = pos_dir+"ADF/"
            if not os.path.isdir(ADF_dir):
                os.mkdir(ADF_dir)
            solute = "".join([n for n in nuc_symbol if n.isalpha()])
            if not skip_compute_neighbors:
                print("Computing nearest neighbors...")
                dct = neighbors.periodic_nearest_neighbors_by_atom(u,    # Universe with all frames from which we want to extract clusters
                                        solute,       # Source atom from which we will search for neighbors
                                        celldm,       # Cubic cell dimension
                                        [nn],         # Cluster sizes we want
                                        #take_prj=14,
                                        dmax=celldm/2)

                grouped_nn = dct[nn].atom.groupby('frame')
                if write_xyzs:
                    print("Writing .xyz files and ADF inputs...")
                    if not os.path.isdir(xyzs):
                        os.mkdir(xyzs)
                else:
                    print("Writing ADF inputs...")
            else:
                print("Writing ADF inputs with coordinates from previously computed clusters...")

        if write_GIPAW:
            print("Writing QE-GIPAW inputs...")
            paw_dir = pos_dir + "GIPAW/"
            if not os.path.isdir(paw_dir):
                os.mkdir(paw_dir)

        i = 0
        dt = timestep*2.418884e-5
        for frame, group in grouped:
            i+=1
            fname = str(i).zfill(4)
            time = frame*dt
            comment = "frame: {}, time: {}".format(frame, time)

            if write_ADF:
                if not skip_compute_neighbors:
                    cluster = grouped_nn.get_group(frame).copy()
                    cluster_angs = cluster.copy()
                    cluster_angs.loc[:,'x'] = cluster.loc[:,'x']*0.529177
                    cluster_angs.loc[:,'y'] = cluster.loc[:,'y']*0.529177
                    cluster_angs.loc[:,'z'] = cluster.loc[:,'z']*0.529177

                    shit = cluster_angs.loc[:,['symbol','x','y','z']].values
                    a = str(shit).replace('[','')
                    b = a.replace(']','')
                    c = b.replace('\n ', '\n')
                    xyz = c.replace("\'",'')
                    xyz_lines = xyz.split('\n')
                    sorted_xyz = sorted(xyz_lines,key=lambda x: My_fn(x,solute,symbols))
                    coord = "\n".join(sorted_xyz)
                    nat = len(cluster)

                    if write_xyzs:
                        with open(xyzs + str(i).zfill(4)+'.xyz', 'w') as f:
                            f.write(str(nat)+"\n{0}\n".format(comment)+coord)
                else:
                    with open(xyzs + str(i).zfill(4)+'.xyz', 'r') as f:
                        lines = f.readlines()
                    coord = "".join(lines[2:])
                    nat = len(lines[2:])

                if not os.path.isdir(ADF_dir + fname):
                    os.mkdir(ADF_dir +"/"+ fname)

                ch = 0
                clines = coord.split('\n')
                #print(clines)
                for line in clines:
                    ll=line.split()
                    if ll[0] == solute:
                        ch += solute_charge
                    if ll[0] == 'O':
                        ch += -2
                    if ll[0] == 'H':
                        ch += 1

                with open(ADF_dir + fname + "/" + fname + "-scf.inp", 'w') as g:
                    g.write(ADF_in.format(coord,comment,ch,solute))

                with open(ADF_dir + fname + "/" + fname + ".slm", 'w') as g:
                    g.write(ADF_slm.format(traj, fname, CCR_scratch))

            if write_GIPAW:
                #print(atom.head())
                #print(nuc_df)
                atom_angs = atom.copy()
                atom_angs.loc[:,'x'] = atom.loc[:,'x']*0.529177
                atom_angs.loc[:,'y'] = atom.loc[:,'y']*0.529177
                atom_angs.loc[:,'z'] = atom.loc[:,'z']*0.529177
                shit = atom_angs[atom_angs['frame']==frame].loc[:,['symbol','x','y','z']].values
                a = str(shit).replace('[','')
                b = a.replace(']','')
                c = b.replace('\n ', '\n')
                full_coord = c.replace("\'",'')
                nat = len(shit)

                if not os.path.isdir(paw_dir + fname):
                    os.mkdir(paw_dir +"/"+ fname)
                with open(paw_dir + fname + "/" + fname + "-scf.inp", 'w') as g:
                    g.write(PAW_scf_in.format(comment,fname,celldm,nat,full_coord))

                with open(paw_dir + fname + "/" + fname + "-efg.inp", 'w') as g:
                    g.write(PAW_efg_in.format(comment,fname,str(nuc_df.loc['Q',nuc_symbol]*1e30)))
        #Your slurm template here!
                with open(paw_dir + fname + "/" + fname + ".slm", 'w') as g:
                    g.write(PAW_slm.format(fname,traj,CCR_scratch))
    print("Done")
    return(dct[nn])


import time
import os
import torch
import numpy as np
import cv2
import openpifpaf
import pycuda.driver as cuda
import PIL

import tensorrt as trt
from libs.detectors.x86.openpifpaf_tensorrt.decoder import CifCafDecoder
from libs.detectors.utils.fps_calculator import convert_infr_time_to_fps

def allocate_buffers(engine):
    host_inputs = []
    cuda_inputs = []
    host_outputs = []
    cuda_outputs = []
    bindings = []
    for i in range(engine.num_bindings):
        binding = engine[i]
        size = trt.volume(engine.get_binding_shape(binding)) * \
               engine.max_batch_size
        host_mem = cuda.pagelocked_empty(size, np.float32)
        cuda_mem = cuda.mem_alloc(host_mem.nbytes)
        bindings.append(int(cuda_mem))
        if engine.binding_is_input(binding):
            host_inputs.append(host_mem)
            cuda_inputs.append(cuda_mem)
        else:
            host_outputs.append(host_mem)
            cuda_outputs.append(cuda_mem)
    stream = cuda.Stream()  # create a CUDA stream to run inference
    return bindings, host_inputs, cuda_inputs, host_outputs, cuda_outputs, stream

class Detector:
    """
    TODO: UPDATE for TensorRT
    Perform pose estimation with Openpifpaf model. extract pedestrian's bounding boxes from key-points.
    :param config: Is a ConfigEngine instance which provides necessary parameters.
    """

    def _load_engine(self):
        precision=int(self.model_variables['TensorrtPrecision'])
        TRTbinPath='/repo/data/tensorrt/openpifpaf_resnet50_{}_{}_d{}.trt'.format(self.w,self.h,precision)
        if not os.path.exists(TRTbinPath):
            os.system('bash /repo/generate_tensorrt.bash config-x86-gpu-tensorrt.ini')
        with open(TRTbinPath, 'rb') as f, trt.Runtime(self.trt_logger) as runtime:
            return runtime.deserialize_cuda_engine(f.read())

    def __init__(self, config, model_name, variables):
        self.config = config
        self.model_name = model_name
        self.model_variables = variables
        self.fps = None
        self.w, self.h, _ = [int(i) for i in self.model_variables['ImageSize'].split(',')]
        self.trt_logger = trt.Logger(trt.Logger.INFO)
        self.model_input_size = (self.w, self.h)
        self.device = None  # enter your Gpu id here
        self.cuda_context = None
        self._init_cuda_stuff()

    def _init_cuda_stuff(self):
        cuda.init()
        self.engine = self._load_engine()
        self.device = cuda.Device(0)  # enter your Gpu id here
        self.cuda_context = self.device.make_context()
        self.engine_context = self.engine.create_execution_context()
        bindings, host_inputs, cuda_inputs, host_outputs, cuda_outputs, stream = allocate_buffers(self.engine)
        self.bindings = bindings
        self.host_inputs = host_inputs
        self.host_outputs = host_outputs
        self.cuda_inputs = cuda_inputs
        self.cuda_outputs = cuda_outputs
        self.stream = stream

    def __del__(self):
        """ Free CUDA memory. """

        self.cuda_context.pop()
        del self.cuda_context
        del self.engine_context
        del self.engine

    def inference(self, resized_rgb_image):
        """
        This method will perform inference and return the detected bounding boxes
        Args:
            resized_rgb_image: uint8 numpy array with shape (img_height, img_width, channels)

        Returns:
            result: a dictionary contains of [{"id": 0, "bbox": [x1, y1, x2, y2], "score":s%}, {...}, {...}, ...]

        """
        image = resized_rgb_image
        image = cv2.resize(image, self.model_input_size)
        pil_im = PIL.Image.fromarray(image)
        preprocess = None

        data = openpifpaf.datasets.PilImageList([pil_im], preprocess=preprocess)
        loader = torch.utils.data.DataLoader(
            data, batch_size=1, shuffle=False, pin_memory=True,
            collate_fn=openpifpaf.datasets.collate_images_anns_meta)

        for images_batch, _, __ in loader:
            np_img = images_batch.numpy()

        bindings = self.bindings
        host_inputs = self.host_inputs
        host_outputs = self.host_outputs
        cuda_inputs = self.cuda_inputs
        cuda_outputs = self.cuda_outputs
        stream = self.stream

        host_inputs[0] = np.ravel(np.zeros_like(np_img))

        self.cuda_context.push()
        t_begin = time.perf_counter()

        np.copyto(host_inputs[0], np.ravel(np_img))
        cuda.memcpy_htod_async(
            cuda_inputs[0], host_inputs[0], stream)

        self.engine_context.execute_async(
            batch_size=1,
            bindings=bindings,
            stream_handle=stream.handle)

        cif=[None] * 1
        caf=[None] * 1
        cif_names=['cif']
        caf_names=['caf']
        for i in range(1, self.engine.num_bindings):
            cuda.memcpy_dtoh_async(
                host_outputs[i - 1], cuda_outputs[i - 1], stream)

        stream.synchronize()

        for i in range(1, self.engine.num_bindings):
            shape = self.engine.get_binding_shape(i)
            name = self.engine.get_binding_name(i)
            total_shape = np.prod(shape)
            output = host_outputs[i - 1][0: total_shape]
            output = np.reshape(output, tuple(shape))
            if name in cif_names:
                index_n = cif_names.index(name)
                tmp = torch.from_numpy(output[0])
                cif = tmp.cpu().numpy()
            elif name in caf_names:
                index_n = caf_names.index(name)
                tmp = torch.from_numpy(output[0])
                caf = tmp.cpu().numpy()

        heads = [cif, caf]
        self.cuda_context.pop()

        inference_time = time.perf_counter() - t_begin

        fields = heads

        decoder_time=time.perf_counter()
        decoder = CifCafDecoder()
        predictions = decoder.decode(fields)
        decoder_time = time.perf_counter() - t_begin
        self.fps = convert_infr_time_to_fps(inference_time+decoder_time)

        result = []

        for i, pred_object in enumerate(predictions):
            pred = pred_object.data
            pred_visible = pred[pred[:, 2] > .2]
            xs = pred_visible[:, 0]
            ys = pred_visible[:, 1]

            if len(xs) == 0 or len(ys) == 0:
                continue

            x, y, w, h = pred_object.bbox()

            x_min = int(x)
            x_max = int(x + w)
            y_min = int(y)
            y_max = int(y + h)
            xmin = int(max(x_min - .15 * w, 0))
            xmax = int(min(x_max + .15 * w, self.w))
            ymin = int(max(y_min - .2 * h, 0))
            ymax = int(min(y_max + .05 * h, self.h))
            bbox_dict = {"id": "1-" + str(i), "bbox": [ymin / self.h, xmin / self.w, ymax / self.h, xmax / self.w],
                    "score": 0.9, "face": None}

            # extract face bounding box
            if np.all(pred[[0, 1, 2, 5, 6], -1] > 0.15):
                x_min_face = int(pred[6, 0])
                x_max_face = int(pred[5, 0])
                y_max_face = int((pred[5, 1] + pred[6, 1]) / 2)
                y_eyes = int((pred[1, 1] + pred[2, 1]) / 2)
                y_min_face = 2 * y_eyes - y_max_face
                if (y_max_face - y_min_face > 0) and (x_max_face - x_min_face > 0):
                    h_crop = y_max_face - y_min_face
                    x_min_face = int(max(0, x_min_face - 0.1 * h_crop))
                    y_min_face = int(max(0, y_min_face - 0.1 * h_crop))
                    x_max_face = int(min(self.w, x_min_face + 1.1 * h_crop))
                    y_max_face = int(min(self.h, y_min_face + 1.1 * h_crop))
                    bbox_dict["face"] = [y_min_face / self.h, x_min_face / self.w, y_max_face / self.h, x_max_face / self.w]

            result.append(bbox_dict)

        return result


import torch.nn.functional as F
from torch.autograd import Variable
import torch.nn as nn
import numpy as np
import torch
import lib
import random

class EncoderRNN(nn.Module):
    def __init__(self, opt, vocab):
        super(EncoderRNN, self).__init__()
        self.vocab = vocab
        self.opt = opt
        self.vocab_size = len(self.vocab)
        self.num_directions = 2 if self.opt.brnn else 1
        self.embedding = nn.Embedding(self.vocab_size, opt.emb_size, padding_idx=lib.constants.PAD)
        self.rnn = getattr(nn, self.opt.rnn_type)(
            input_size=self.opt.emb_size,
            hidden_size=opt.rnn_size // self.num_directions,
            num_layers=self.opt.layers,
            dropout=self.opt.dropout,
            bidirectional=self.opt.brnn)

    def forward(self, src, src_lens, hidden=None):
        emb = self.embedding(src)
        packed_emb = nn.utils.rnn.pack_padded_sequence(emb, src_lens)
        packed_outputs, self.hidden = self.rnn(packed_emb, hidden)
        outputs, output_lens =  nn.utils.rnn.pad_packed_sequence(packed_outputs)
        if self.opt.brnn: self.hidden = self._cat_directions(self.hidden)
        return outputs, self.hidden

    def _cat_directions(self, hidden):
        def _cat(h):
            return torch.cat([h[0:h.size(0):2], h[1:h.size(0):2]], 2)
        if isinstance(hidden, tuple): #LSTM
            hidden = tuple([_cat(h) for h in hidden])
        else: #GRU
            hidden = _cat(hidden)
        return hidden

class LuongAttnDecoderRNN(nn.Module):
    def __init__(self, opt, vocab):
        super(LuongAttnDecoderRNN, self).__init__()
        self.opt =  opt
        self.vocab = vocab
        self.vocab_size = len(self.vocab)
        self.tanh = nn.Tanh()
        self.embedding =  nn.Embedding(self.vocab_size, opt.emb_size, padding_idx=lib.constants.PAD)
        self.rnn = getattr(nn, self.opt.rnn_type)(
            input_size=self.opt.emb_size,
            hidden_size=self.opt.rnn_size,
            num_layers=self.opt.layers,
            dropout=self.opt.dropout)

        if self.opt.attention:
            self.W_a = nn.Linear(self.opt.rnn_size, self.opt.rnn_size, bias=opt.bias)
            self.W_c = nn.Linear(self.opt.rnn_size + self.opt.rnn_size, self.opt.rnn_size, bias=opt.bias)

        if self.opt.tie_decoder_embeddings and self.vocab_size!=1:
            self.W_proj = nn.Linear(self.opt.rnn_size, self.opt.emb_size, bias=opt.bias)
            self.W_s = nn.Linear(self.opt.emb_size, self.vocab_size, bias=opt.bias)
            self.W_s.weight = self.embedding.weight
        else:
            self.W_s = nn.Linear(self.opt.rnn_size, self.vocab_size, bias=opt.bias)

    def forward(self, src, src_lens, encoder_outputs, decoder_hidden):
        emb = self.embedding(src.unsqueeze(0))
        decoder_output, self.decoder_hidden = self.rnn(emb, decoder_hidden)
        decoder_output = decoder_output.transpose(0,1)
        if self.opt.attention:
            attention_scores = torch.bmm(decoder_output, self.W_a(encoder_outputs).transpose(0,1).transpose(1,2))
            attention_mask = lib.metric.sequence_mask(src_lens).unsqueeze(1)
            attention_scores.data.masked_fill_(1 - attention_mask.data, -float('inf'))
            attention_weights = F.softmax(attention_scores.squeeze(1), dim=1).unsqueeze(1)
            context_vector = torch.bmm(attention_weights, encoder_outputs.transpose(0,1))
            concat_input = torch.cat([context_vector, decoder_output], -1)
            concat_output = self.tanh(self.W_c(concat_input))
            attention_weights = attention_weights.squeeze(1)
        else:
            attention_weights = None
            concat_output = decoder_output
        if self.opt.tie_decoder_embeddings and self.vocab_size!=1:
            output = self.W_s(self.W_proj(concat_output))
        else:
            output = self.W_s(concat_output)
        output = output.squeeze(1)
        del src_lens
        return output, self.decoder_hidden, attention_weights

class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, opt):
        super(Seq2Seq, self).__init__()
        self.torch = torch.cuda if opt.cuda else torch
        self.encoder = encoder
        self.decoder = decoder
        self.opt = opt
        if opt.share_embeddings:
            self.encoder.embedding.weight = self.decoder.embedding.weight

    def forward(self, batch, eval=False):
        tgt, tgt_lens = batch['tgt']
        src, src_lens = batch['src']
        batch_size = src.size(1)
        # print(batch_size)
        # print('Day la tgt.size()[0]: ', tgt.size()[0])
        assert(batch_size == tgt.size(1))
        input_seq = Variable(torch.LongTensor([lib.constants.BOS] * batch_size))
        decoder_outputs = Variable(torch.zeros(self.opt.max_train_decode_len, batch_size, self.decoder.vocab_size))
        # print('len cua decoder: ', len(decoder_outputs))
        if self.opt.cuda: input_seq, decoder_outputs = input_seq.cuda(), decoder_outputs.cuda()
        max_tgt_len = tgt.size()[0]
        encoder_outputs, encoder_hidden = self.encoder(src, src_lens.data.tolist())
        decoder_hidden = encoder_hidden
        use_teacher_forcing = False if eval else random.random() < self.opt.teacher_forcing_ratio
        for t in range(max_tgt_len):
            decoder_output, decoder_hidden, attention_weights = self.decoder(input_seq, src_lens, encoder_outputs, decoder_hidden)
            decoder_outputs[t] = decoder_output
            if use_teacher_forcing:
                input_seq = tgt[t]
            else:
                topv, topi = decoder_output.topk(1)
                input_seq = topi.squeeze()
        return decoder_outputs

    def backward(self, outputs, tgt_seqs, mask, criterion, eval=False, normalize=True):
        max_tgt_len = tgt_seqs.size()[0]
        logits = outputs[:max_tgt_len]
        loss, num_corrects = criterion(logits, tgt_seqs, mask, normalize=normalize)
        if(not eval): loss.backward()
        return loss.item(), num_corrects

    def translate(self, batch):
        tgt, tgt_lens = batch['tgt']
        src, src_lens = batch['src']
        batch_size = src.size(1)
        assert (batch_size == tgt.size(1))
        input_seq = Variable(torch.LongTensor([lib.constants.BOS] * batch_size))
        decoder_outputs = Variable(torch.zeros(self.opt.max_train_decode_len, batch_size, self.decoder.vocab_size))
        if self.opt.cuda: input_seq, decoder_outputs = input_seq.cuda(), decoder_outputs.cuda()
        encoder_outputs, encoder_hidden = self.encoder(src, src_lens.data.tolist())
        decoder_hidden = encoder_hidden
        if self.opt.attention: all_attention_weights = torch.zeros(self.opt.max_train_decode_len, src.size(1), len(src))
        end_of_batch_pred = np.array([lib.constants.EOS] * len(src_lens))
        preds = np.ones((self.opt.max_train_decode_len, len(src_lens))) * 2
        probs = np.ones((self.opt.max_train_decode_len, len(src_lens))) * 2
        for t in range(self.opt.max_train_decode_len):
            decoder_output, decoder_hidden, attention_weights = self.decoder(input_seq, src_lens, encoder_outputs, decoder_hidden)
            if self.opt.attention:
                all_attention_weights[t] = attention_weights.cpu().data
            prob, token_ids = decoder_output.data.topk(1)
            token_ids = token_ids.squeeze()
            prob = prob.squeeze()
            preds[t,:] = token_ids
            probs[t,:] = prob
            input_seq = Variable(token_ids)
            if np.sum(np.equal(token_ids.cpu().numpy(),end_of_batch_pred)) == len(src):
                break
        preds = torch.LongTensor(preds)
        return probs, preds


import time

import numpy as np
from numpy.linalg import inv
from scipy.optimize import fsolve

from Trajectory import Trajectory
from cosraKinematics import cosraKinematics
from dynamixel import dynamixel

class cosraCtrl:
    def __init__(self):
        pass

    """ Dynamixel Test: Basic """
    def dynamixelTest(self):
        # checking dynamixel
        # 기구부 구동시키기 위한 모터 출력 체크하기
        for i in range(4):
            dxls.set_pos(dxl_ids[i], 0)
            time.sleep(0.5)

        time.sleep(1)

        for i in range(4):
            dxls.set_pos(dxl_ids[i], 4095)
            time.sleep(0.5)
        dxls.set_pos_sync(dxl_ids, [2048, 2048, 2048, 2048])

    def cableDrivingTest(self):
        while True:
            dxls.set_pos(dxl_ids[3], 1000)
            # dxls.set_pos(dxl_ids[2], 0)
            time.sleep(2)
            dxls.set_pos(dxl_ids[3], 1500)
            # dxls.set_pos(dxl_ids[2], 4095)
            time.sleep(2)

    """ Pick N Place Algorithm """
    def pickNplace(self, start=[], goal=[]):
        t = 0.0
        tf = 3.0        # 시간 조정해야할 수도? 다이나믹셀 팍팍 움직임.
        gripAngle = 333  # 그리퍼가 물체 잡기 위해 필요한 각도 (0 ~ 4095 값 -> tuning)

        # manipulator는 home position에서 수평 방향으로 이동
        # vision 처리를 통해 start object와 goal object 인식하고 좌표에 대한 정보를 보내주면,
        # 해당 좌표의 수직 방향으로 manipulator가 움직이는 평면에서의 via point를 2개 형성.
        # 따라서, manipulator는 home -> via1 -> start -> via1 -> via2 -> goal -> via2 -> home 의 trajectory를 갖는다.

        ''' Home Position '''
        self.homePosition()
        q_home = self.dxlPos2rad(dxls.get_pos_sync(dxl_ids[:3]))  # 각 id별 dxl position을 받아 rad으로 변환
        # print(q_home)
        cart_home = cosraKinematics.fk(joints=q_home)[:3, -1]  # home position EE cartesian
        # print(cart_home)

        ''' Via Points '''
        cart_pick = np.array(start)  # pick 해야 하는 물체의 horizontal cartesian
        # print(cart_pick)
        cart_via1 = np.array(start)
        cart_via1[2] += 0.04        # manipulator gripper 초기 높이만큼 offset
        # print(cart_via1)
        cart_place = np.array(goal)  # place 해야 하는 곳의 horizontal cartesian
        # print(cart_place)
        cart_via2 = np.array(goal)
        cart_via2[2] += 0.04        # manipulator gripper 초기 높이만큼 offset
        # print(cart_via2)

        """ Cartesian Trajectories """
        # Trajectory 1: Move to Start point horizontally
        traj_s1, t1 = Trajectory.LSPB(q0=cart_home, qf=cart_via1, tf=tf, tb=tf / 3)

        # Trajectory 2: Go down for Gripping
        traj_s2, t2 = Trajectory.LSPB(q0=cart_via1, qf=cart_pick, tf=tf, tb=tf / 3)

        # Trajectory 3: Pick up
        traj_s3, t3 = Trajectory.LSPB(q0=cart_pick, qf=cart_via1, tf=tf, tb=tf / 3)

        # Trajectory 4: Move to Goal point horizontally
        traj_s4, t4 = Trajectory.LSPB(q0=cart_via1, qf=cart_via2, tf=tf, tb=tf / 3)

        # Trajectory 5: Place
        traj_s5, t5 = Trajectory.LSPB(q0=cart_via2, qf=cart_place, tf=tf, tb=tf / 3)

        # Trajectory 6: Return
        traj_s6, t6 = Trajectory.LSPB(q0=cart_place, qf=cart_via2, tf=tf, tb=tf / 3)

        # Trajectory 7: Home coming
        traj_s7, t7 = Trajectory.LSPB(q0=cart_via2, qf=cart_home, tf=tf, tb=tf / 3)

        """ Joint Trajectories using IK """
        inHome = False

        for i in range(len(t1)):
            # q_cur = self.dxlPos2rad(dxls.get_pos_sync(dxl_ids[:3]))
            q_cur = np.zeros(3)
            for j in range(3):
                q_cur[j] = dxls.get_pos(dxl_ids[j])
            q_cur = dxlPos2rad(q_cur)
            # print(q_cur)
            joints = fsolve(ik, x0=q_cur, args=tuple(traj_s1[i]))
            # joints1 = cosraKinematics.ik(traj_s1[i])  # radians
            # print(joints1)
            dxl_pos = self.rad2dxlPos(joints)
            # dxls.set_pos_sync(dxl_ids, dxl_pos)
            for k in range(3):
                dxls.set_pos(dxl_ids[k], int(dxl_pos[k]))
            # time.sleep(0.005)

        for i in range(len(t2)):
            q_cur = np.zeros(3)

            for j in range(3):
                q_cur[j] = dxls.get_pos(dxl_ids[j])

            q_cur = dxlPos2rad(q_cur)
            joints = fsolve(ik, x0=q_cur, args=tuple(traj_s2[i]))
            dxl_pos = self.rad2dxlPos(joints)

            for k in range(3):
                dxls.set_pos(dxl_ids[k], int(dxl_pos[k]))

            cart_cur = cosraKinematics.fk(joints=q_cur)[:3, -1]  # home position EE cartesian

            ''' Grip '''
            if np.linalg.norm(cart_cur - cart_pick) < 10e-4:
                dxls.set_pos(dxl_ids[3], gripAngle)
                time.sleep(2)

        for i in range(len(t3)):
            q_cur = np.zeros(3)

            for j in range(3):
                q_cur[j] = dxls.get_pos(dxl_ids[j])

            q_cur = dxlPos2rad(q_cur)
            joints = fsolve(ik, x0=q_cur, args=tuple(traj_s3[i]))
            dxl_pos = self.rad2dxlPos(joints)

            for k in range(3):
                dxls.set_pos(dxl_ids[k], int(dxl_pos[k]))

        for i in range(len(t4)):
            q_cur = np.zeros(3)

            for j in range(3):
                q_cur[j] = dxls.get_pos(dxl_ids[j])

            q_cur = dxlPos2rad(q_cur)
            joints = fsolve(ik, x0=q_cur, args=tuple(traj_s4[i]))
            dxl_pos = self.rad2dxlPos(joints)

            for k in range(3):
                dxls.set_pos(dxl_ids[k], int(dxl_pos[k]))

        for i in range(len(t5)):
            q_cur = np.zeros(3)

            for j in range(3):
                q_cur[j] = dxls.get_pos(dxl_ids[j])

            q_cur = dxlPos2rad(q_cur)
            joints = fsolve(ik, x0=q_cur, args=tuple(traj_s5[i]))
            dxl_pos = self.rad2dxlPos(joints)

            for k in range(3):
                dxls.set_pos(dxl_ids[k], int(dxl_pos[k]))

            cart_cur = cosraKinematics.fk(joints=q_cur)[:3, -1]  # home position EE cartesian

            ''' Place '''
            if np.linalg.norm(cart_cur - cart_pick) < 10e-4:
                dxls.set_pos(dxl_ids[3], -1 * gripAngle)
                time.sleep(2)

        for i in range(len(t6)):
            q_cur = np.zeros(3)

            for j in range(3):
                q_cur[j] = dxls.get_pos(dxl_ids[j])

            q_cur = dxlPos2rad(q_cur)
            joints = fsolve(ik, x0=q_cur, args=tuple(traj_s6[i]))
            dxl_pos = self.rad2dxlPos(joints)

            for k in range(3):
                dxls.set_pos(dxl_ids[k], int(dxl_pos[k]))

        for i in range(len(t7)):
            q_cur = np.zeros(3)

            for j in range(3):
                q_cur[j] = dxls.get_pos(dxl_ids[j])

            q_cur = dxlPos2rad(q_cur)
            joints = fsolve(ik, x0=q_cur, args=tuple(traj_s7[i]))
            dxl_pos = self.rad2dxlPos(joints)

            for k in range(3):
                dxls.set_pos(dxl_ids[k], int(dxl_pos[k]))

            cart_cur = cosraKinematics.fk(joints=q_cur)[:3, -1]  # home position EE cartesian

            ''' Place '''
            if np.linalg.norm(cart_cur - cart_home) < 10e-4:
                inHome = True
                time.sleep(2)
                return inHome

    def dxlPos2rad(self, pos=[]):
        # dxl position: 0 ~ 4095 (0 degree at 2048)
        joints = np.zeros(3)
        for i in range(3):
            joints[i] = 2 * pos[i] * np.pi / 4096
        return joints

    def rad2dxlPos(self, joints=[]):
        pos = np.zeros(3)
        for i in range(3):
            pos[i] = 4096 * joints[i] / (2 * np.pi)
        return pos

    def homePosition(self):
        homePos = [2048, 2048, 2048, 2048]
        dxls.set_pos_sync(dxl_ids, homePos)

if __name__ == '__main__':
    cosra = cosraCtrl()
    dxls = dynamixel()
    dxl_ids = [0, 1, 2, 3]  # ID setting

    for i in range(4):
        dxls.enable_torque(dxl_ids[i], True)  # torque enable
        # dxls.set_pos(dxl_ids[i], 0)  # home position

    # checking current dynamixel position
    # print(dxls.get_pos(dxl_ids[0]))
    # print(cosra.dxlPos2rad(dxls.get_pos(dxl_ids[0])))
    # print(dxls.get_pos(dxl_ids[1]))
    # print(dxls.get_pos(dxl_ids[2]))
    # print(dxls.get_pos_sync(dxl_ids))
    # print(cosra.dxlPos2rad(dxls.get_pos_sync(dxl_ids[:3])))

    # checking fk results: ee cartesian
    # print(cosraKinematics.fk(cosra.dxlPos2rad(dxls.get_pos_sync(dxl_ids[:3])))[:3,-1])

    """ Dynamixel Test: Basic """
    # cosra.dynamixelTest()
    # cosra.cableDrivingTest()

    # reached = False  # whether all the blocks are arranged correctly
    # inHome = False  # whether robot is in motion

    """ Pick N Place Algorithm """
    # while not reached:
    #     # 상자 분류 완료하지 못 했을 때
    #     cosra.pickNplace()

        # 상자 분류 완성했을 때
        # if ...:
        #     reached = True
        #     break

    cosra.homePosition()
    dxls.close_port()


import numpy as np
import pandas as pd

'''
@X: Tokenized samples

__________________________________________________________________
index	Stemmed Text 	            Lemmatized Text
__________________________________________________________________
0	    husband live nearbi stop	husband live nearby stop
1	    great attent servic great	great attentive service great
2	    locat great fresh brew	    location great fresh brew
3	    nice restaur park vega	    nice restaurant park vegas
4	    star sushi joint serv size	make star sushi joint serve
__________________________________________________________________

'''

'''
@y: the corresponding target values

______________________________________________
index   Multiclass Multilabel(one-hot encoded)
______________________________________________
0	    -1	        1	0	0
1	    1	        0	0	1
2	    1	        0	0	1
3	    0	        0	1	0
4	    1	        0	0	1
______________________________________________
'''

class Dataset():
    def __init__(self, path, **kwargs):
        print('Starting initialization')

        self.df = pd.read_csv(path)
        self.X = self.df['text']
        print(self.X[:10])
        # self.X_hours = get_regular_features(self.df)
        # print(self.X_hours[:10])
        cutWords_list = [x.split(" ") for x in self.X]

        from sklearn.preprocessing import LabelEncoder
        labelEncoder = LabelEncoder()
        self.y = labelEncoder.fit_transform(self.df['category'])
        print(self.y[:10])

        from sklearn.preprocessing import OneHotEncoder
        encoder = OneHotEncoder(sparse=False)
        self.y_one_hot = encoder.fit_transform(self.y.reshape((6000, 1)))

        for arg, value in kwargs.items():
            if arg == 'glove' and value == 1:
                self._isGloVe = True
                # _ = glove2word2vec('vectors.txt', 'test_glove.txt')
                glove_model = KeyedVectors.load_word2vec_format('test_glove.txt')
                self.X_glove = [get_contentVector(cutWords, glove_model) for cutWords in cutWords_list]
            elif arg == 'sg'and value == 1:
                self._isSG = True
                # train SG and save word vectors as Binary File
                # sg_model = Word2Vec(cutWords_list, sg = 1, size=300, iter=30, min_count=5, workers=multiprocessing.cpu_count())
                # sg_model.wv.save_word2vec_format("sg_wv.bin", binary=True)
                # Load Binary SG word vectors
                wv_sg = KeyedVectors.load_word2vec_format("sg_wv.bin", binary=True)
                self.X_sg = [get_contentVector(cutWords, wv_sg) for cutWords in cutWords_list]
            elif arg == 'cbow' and value == 1:
                self._isCBOW = True
                # cbow_model = Word2Vec(cutWords_list, sg = 0, size=300, iter=30, min_count=5, workers=multiprocessing.cpu_count())
                # cbow_model.wv.save_word2vec_format("cbow_wv.bin", binary=True)
                wv_cbow = KeyedVectors.load_word2vec_format("cbow_wv.bin", binary=True)
                self.X_cbow = [get_contentVector(cutWords, wv_cbow) for cutWords in cutWords_list]

        print('Initialization finished')

    def get_X_SG(self):
        print(self.X_sg[:1])
        return self.X_sg

    def get_X_CBOW(self):
        print(self.X_cbow[:1])
        return self.X_cbow

    def get_X_GloVe(self):
        print(self.X_glove[:1])
        return self.X_glove

    def get_y(self):
        return self.y


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import time
import numpy as np
import tensorflow as tf
import model_blstm

FLAGS = tf.app.flags.FLAGS

tf.app.flags.DEFINE_string('data_dir', '../data/',
                           "data directory")
tf.app.flags.DEFINE_string('checkpoints_dir', '../checkpoints/',
                           "training checkpoints directory")
tf.app.flags.DEFINE_string('log_dir', '../train_log/',
                           "summary directory")
tf.app.flags.DEFINE_string('restore_path', None,
                           "path of saving model eg: ../checkpoints/model.ckpt-5")
tf.app.flags.DEFINE_integer('batch_size', 128,
                            "mini-batch size")
tf.app.flags.DEFINE_integer('total_epoches', 300,
                            "total training epoches")
tf.app.flags.DEFINE_integer('hidden_size', 128,
                            "size of LSTM hidden memory")
tf.app.flags.DEFINE_integer('num_layers', 2,
                            "number of stacked blstm")
tf.app.flags.DEFINE_integer("input_dims", 10,
                            "input dimensions")
tf.app.flags.DEFINE_integer("num_classes", 69,  # 68 letters + 1 blank
                            "num_labels + 1(blank)")
tf.app.flags.DEFINE_integer('save_freq', 5,
                            "ephoches of frequency of saving model")
tf.app.flags.DEFINE_float('learning_rate', 0.001,
                          "learning rate of RMSPropOptimizer")
tf.app.flags.DEFINE_float('decay_rate', 0.99,
                          "decay rate of RMSPropOptimizer")
tf.app.flags.DEFINE_float('momentum', 0.9,
                          "momentum of RMSPropOptimizer")
tf.app.flags.DEFINE_float('max_length', 1940,
                          "pad to same length")
tf.app.flags.DEFINE_integer('label_pad', 63,
                            "label pad size")
tf.app.flags.DEFINE_boolean('if_valid_vr', False,
                            "label pad size")
tf.app.flags.DEFINE_boolean('if_lowercase_only', False,
                            "if letter table only contain lowercase")

if FLAGS.if_lowercase_only:
    letter_table = [' ', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'ga', 'h', 'i', 'j', 'k', 'km', 'l', 'm', 'n', 'o', 'p', 'pt',
                    'q', 'r', 's', 'sc', 'sp', 't', 'u', 'v', 'w', 'x', 'y', 'z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '<b>']
else:
    letter_table = [' ', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f',
                    'g', 'ga', 'h', 'i', 'j', 'k', 'km', 'l', 'm', 'n', 'o', 'p', 'pt', 'q', 'r', 's', 'sc', 'sp', 't', 'u', 'v', 'w', 'x', 'y', 'z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '<b>']

class ModelConfig(object):
    """
    testing config
    """

    def __init__(self):
        self.data_dir = FLAGS.data_dir
        self.checkpoints_dir = FLAGS.checkpoints_dir
        self.log_dir = FLAGS.log_dir
        self.restore_path = FLAGS.restore_path
        self.batch_size = FLAGS.batch_size
        self.total_epoches = FLAGS.total_epoches
        self.hidden_size = FLAGS.hidden_size
        self.num_layers = FLAGS.num_layers
        self.input_dims = FLAGS.input_dims
        self.num_classes = FLAGS.num_classes
        self.save_freq = FLAGS.save_freq
        self.learning_rate = FLAGS.learning_rate
        self.decay_rate = FLAGS.decay_rate
        self.momentum = FLAGS.momentum
        self.max_length = FLAGS.max_length
        self.label_pad = FLAGS.label_pad
        self.if_valid_vr = FLAGS.if_valid_vr

    def show(self):
        print("data_dir:", self.data_dir)
        print("checkpoints_dir:", self.checkpoints_dir)
        print("log_dir:", self.log_dir)
        print("restore_path:", self.restore_path)
        print("batch_size:", self.batch_size)
        print("total_epoches:", self.total_epoches)
        print("hidden_size:", self.hidden_size)
        print("num_layers:", self.num_layers)
        print("input_dims:", self.input_dims)
        print("num_classes:", self.num_classes)
        print("save_freq:", self.save_freq)
        print("learning_rate:", self.learning_rate)
        print("decay_rate:", self.decay_rate)
        print("momentum:", self.momentum)
        print("max_length:", self.max_length)
        print("label_pad:", self.label_pad)
        print("if_valid_vr:", self.if_valid_vr)

def train_model():
    with tf.get_default_graph().as_default() as graph:
        # config setting
        config = ModelConfig()
        config.show()
        # load data
        input_data = np.load(FLAGS.data_dir + 'data.npy')
        target_data = np.load(FLAGS.data_dir + 'dense.npy').item()
        label_data = target_data['dense'].astype(np.int32)
        label_data_length = target_data['length'].astype(np.int32)

        # label_seq_len = target_data['length'].astype(np.int32)
        seq_len_list = []
        for _, v in enumerate(input_data):
            seq_len_list.append(v.shape[0])
        seq_len_list = np.array(seq_len_list).astype(np.int32)
        k = np.argmax(seq_len_list)
        max_length = input_data[k].shape[0]

        if FLAGS.if_valid_vr:
            vr_valid_data_raw = np.load('VRdataValidation.npy')[
                :FLAGS.batch_size]
            vr_valid_target = np.load('VRdenseValidation.npy').item()
            vr_valid_label_raw = vr_valid_target['dense'].astype(np.int32)[
                :FLAGS.batch_size]
            vr_seq_len_list = []
            for _, v in enumerate(vr_valid_data_raw):
                vr_seq_len_list.append(v.shape[0])
            vr_seq_len_list = np.array(vr_seq_len_list).astype(np.int32)
            # padding each textline to maximum length -> max_length (1940)
            vr_valid_data = []
            for _, v in enumerate(vr_valid_data_raw):
                residual = max_length - v.shape[0]
                padding_array = np.zeros([residual, FLAGS.input_dims])
                vr_valid_data.append(
                    np.concatenate([v, padding_array], axis=0))
            vr_valid_data = np.array(vr_valid_data)
            vr_valid_label = []
            # padding each label to same dense length -> label_pad (64)
            for _, v in enumerate(vr_valid_label_raw):
                residual = FLAGS.label_pad - v.shape[0]
                padding_array = np.zeros([residual]) - 1
                vr_valid_label.append(
                    np.concatenate([v, padding_array], axis=0))
            vr_valid_label = np.array(vr_valid_label).astype(np.int32)
            print("vr_valid_data.shape", vr_valid_data.shape)
            print("vr_valid_label.shape", vr_valid_label.shape)

        # padding each textline to maximum length -> max_length (1940)
        padded_input_data = []
        for _, v in enumerate(input_data):
            residual = max_length - v.shape[0]
            padding_array = np.zeros([residual, FLAGS.input_dims])
            padded_input_data.append(
                np.concatenate([v, padding_array], axis=0))
        padded_input_data = np.array(padded_input_data)

        train_data, valid_data = np.split(
            padded_input_data, [padded_input_data.shape[0] * 9 // 10])
        train_seq_len, valid_seq_len = np.split(
            seq_len_list, [seq_len_list.shape[0] * 9 // 10])
        train_label, valid_label = np.split(
            label_data, [label_data.shape[0] * 9 // 10])
        train_label_length, valid_label_length = np.split(
            label_data_length, [label_data_length.shape[0] * 9 // 10])

        # number of batches
        train_num_batch = int(train_label.shape[0] / config.batch_size)
        valid_num_batch = int(valid_label.shape[0] / config.batch_size)
        # model
        model = model_blstm.HWRModel(config, graph)
        # Add an op to initialize the variables.
        init = tf.global_variables_initializer()
        # Add ops to save and restore all the variables.
        saver = tf.train.Saver()
        train_summary_writer = tf.summary.FileWriter(
            FLAGS.log_dir + 'ephoch_train', graph=graph)
        valid_summary_writer = tf.summary.FileWriter(
            FLAGS.log_dir + 'ephoch_valid', graph=graph)
        vr_valid_summary_writer = tf.summary.FileWriter(
            FLAGS.log_dir + 'ephoch_vr_valid', graph=graph)
        # Session
        with tf.Session() as sess:
            sess.run(init)
            # restore model if exist
            if FLAGS.restore_path is not None:
                saver.restore(sess, FLAGS.restore_path)
                print("Model restored:", FLAGS.restore_path)
            # time cost evaluation
            start_time = time.time()
            end_time = 0.0
            for _ in range(config.total_epoches):
                # Shuffle the data
                shuffled_indexes = np.random.permutation(train_data.shape[0])
                train_data = train_data[shuffled_indexes]
                train_seq_len = train_seq_len[shuffled_indexes]
                train_label = train_label[shuffled_indexes]
                loss_sum = 0.0
                for b in range(train_num_batch):
                    batch_idx = b * config.batch_size
                    # input, sequence length, label
                    input_batch = train_data[batch_idx:batch_idx +
                                             config.batch_size]
                    seq_len_batch = train_seq_len[batch_idx:batch_idx +
                                                  config.batch_size]
                    dense_batch = train_label[batch_idx:batch_idx +
                                              config.batch_size]
                    # train
                    global_step, losses = model.step(sess, input_batch,
                                                     seq_len_batch, dense_batch)
                    loss_sum += losses
                global_ephoch = int(global_step // train_num_batch)
                # logging per ephoch
                # validation
                v_loss_sum = 0.0
                for v_b in range(valid_num_batch):
                    v_batch_idx = v_b * config.batch_size
                    # input, sequence length, label
                    v_input_batch = valid_data[v_batch_idx:v_batch_idx +
                                               config.batch_size]
                    v_seq_len_batch = valid_seq_len[v_batch_idx:v_batch_idx +
                                                    config.batch_size]
                    v_dense_batch = valid_label[v_batch_idx:v_batch_idx +
                                                config.batch_size]
                    v_losses = model.compute_losses(sess, v_input_batch,
                                                    v_seq_len_batch, v_dense_batch)
                    v_loss_sum += v_losses

                # predict result
                v_batch_idx = global_ephoch % valid_num_batch
                v_input_batch = valid_data[v_batch_idx:v_batch_idx +
                                           config.batch_size]
                v_seq_len_batch = valid_seq_len[v_batch_idx:v_batch_idx +
                                                config.batch_size]
                v_dense_batch = valid_label[v_batch_idx:v_batch_idx +
                                            config.batch_size]
                v_label_length_batch = valid_label_length[v_batch_idx:v_batch_idx +
                                                          config.batch_size]
                #visualize first data in validation batch
                visual_target_length = v_label_length_batch[0]
                predict, levenshtein = model.predict(
                    sess, v_input_batch, v_seq_len_batch, v_dense_batch)
                str_decoded = ''.join(
                    [letter_table[x] for x in np.asarray(predict.values[:visual_target_length])])
                val_original = ''.join(
                    [letter_table[x] for x in v_dense_batch[0]])
                end_time = time.time()
                print('Original val: %s' % val_original)
                print('Decoded  val: %s' % str_decoded)
                print("%d epoches, %d steps, mean loss: %f, valid mean loss: %f, time cost: %f(sec/batch), levenshtein: %f" %
                      (global_ephoch,
                       global_step,
                       loss_sum / train_num_batch,
                       v_loss_sum / valid_num_batch,
                       (end_time - start_time) / train_num_batch,
                       levenshtein))
                start_time = end_time
                train_summary = tf.Summary(value=[tf.Summary.Value(
                    tag="ephoch_mean_loss", simple_value=loss_sum / train_num_batch)])
                valid_summary = tf.Summary(value=[tf.Summary.Value(
                    tag="ephoch_mean_loss", simple_value=v_loss_sum / valid_num_batch)])
                train_summary_writer.add_summary(
                    train_summary, global_step=global_ephoch)
                valid_summary_writer.add_summary(
                    valid_summary, global_step=global_ephoch)
                train_summary_writer.flush()
                valid_summary_writer.flush()

                # VR validation
                if FLAGS.if_valid_vr:
                    vr_v_loss_sum = model.compute_losses(sess, vr_valid_data,
                                                         vr_seq_len_list, vr_valid_label)
                    # predict vr result
                    vr_idx = global_step % FLAGS.batch_size
                    vr_predict, vr_levenshtein = model.predict(
                        sess, vr_valid_data[vr_idx:vr_idx + 1], vr_seq_len_list[vr_idx:vr_idx + 1], vr_valid_label[vr_idx:vr_idx + 1])
                    vr_str_decoded = ''.join(
                        [letter_table[x] for x in np.asarray(vr_predict.values)])
                    vr_val_original = ''.join(
                        [letter_table[x] for x in vr_valid_label[vr_idx]])
                    print('Original vr_val: %s' % vr_val_original)
                    print('Decoded  vr_val: %s' % vr_str_decoded)

                    vr_valid_summary = tf.Summary(value=[tf.Summary.Value(
                        tag="ephoch_mean_loss", simple_value=vr_v_loss_sum)])
                    vr_valid_summary_writer.add_summary(
                        vr_valid_summary, global_step=global_ephoch)
                    vr_valid_summary_writer.flush()

                    print("VR valid mean loss: %f, vr_levenshtein: %f" %
                          (vr_v_loss_sum, vr_levenshtein))

                if (global_ephoch % FLAGS.save_freq) == 0:
                    save_path = saver.save(
                        sess, FLAGS.checkpoints_dir + "model.ckpt",
                        global_step=global_step)
                    print("Model saved in file: %s" % save_path)

def main(_):
    train_model()

if __name__ == "__main__":
    if not os.path.exists(FLAGS.checkpoints_dir):
        os.makedirs(FLAGS.checkpoints_dir)
    tf.app.run()


import os
import sys
import numpy as np
import time
import cv2

# Deterimens if a photo is a face and returns a scaled image if it is. Returns
# False if it is not a face.
# Params:
#   path - path to the image
#   haar_map - path to the haar_cascade training data
#   pw - the picture width
#   ph - the picture height
def is_face( path, haar_map, pw, ph ):
  good_face = True
  face_cascade = cv2.CascadeClassifier( haar_map )
  #eye_cascade = cv2.CascadeClassifier('haarcascade_eye.xml')
  #smile_cascade = cv2.CascadeClassifier('haarcascade_smile.xml')

  img = cv2.imread( path )
  try:
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  except cv2.error:
    # Some images have problems being converted to grayscale
    print 'Couldn\'t convert to grayscale'
    return False

  faces = face_cascade.detectMultiScale(
    gray,
    scaleFactor=1.2,
    minNeighbors=5,
    minSize=(30, 30),
    flags = cv2.cv.CV_HAAR_SCALE_IMAGE
  )

  # Draw rectangles around detected features
  #for (x,y,w,h) in faces:
    # Draws a square around the face
    #cv2.rectangle(img,(x,y),(x+w,y+h),(255,0,0),2)
    #roi_gray = gray[y:y+h, x:x+w]
    #roi_color = img[y:y+h, x:x+w]
     # eyes = eye_cascade.detectMultiScale(roi_gray, 1.1, 5)
     # for (ex,ey,ew,eh) in eyes:
     #     cv2.rectangle(roi_color,(ex,ey),(ex+ew,ey+eh),(0,0,255),2)
     # smiles = smile_cascade.detectMultiScale(roi_gray, 1.1, 8)
     # for (sx,sy,sw,sh) in smiles:
     #     cv2.rectangle(roi_color,(sx,sy),(sx+sw,sy+sh),(0,255,0),2)

  # Check that there is only one face and that is large enough
  if( len(faces) != 1 ):
    # not a single face
    good_face = False
  else:
    # Dimensions of the face
    x,y,w,h = faces[0]
    if( w < 100 or h < 100 ):
      # too small
      good_face = False

  # If those conditions were met
  if( good_face ):
    img = center_face(img, x, y, w, h, pw, ph)
    return cv2.resize(img, ( pw, ph ) )
  else:
    return good_face

# Center a face in the image and return the new image
# Params:
#   img - the face image
#   x,y - the coords of the bottom left corner of the face
#   w,h - the dimensions of the face
#   pw,ph - the image dimensions
def center_face(img, x, y, w, h, pw, ph):
  # Margins
  m1 = 0.2*pw
  m2 = 0.8*pw
  pts1 = np.float32([[x,y],[x+w,y],[x,y+h],[x+w,y+h]])
  pts2 = np.float32([[m1,m1],[m2,m1],[m1,m2],[m2,m2]])
  M = cv2.getPerspectiveTransform(pts1,pts2)
  dst = cv2.warpPerspective(img,M,(pw,ph))
  return dst

# Shows a face and waits for user to press a key. exit the program if the
# escape key is pressed
# Params
#   img - the image to display
def show_face( img ):
  cv2.destroyAllWindows()
  cv2.imshow('img',img)
  k = cv2.waitKey(50)
  if k == 27 and accepted > 0:
    exit()

################################################################################

def main( prof_path, output_path='output', haar_map='haarcascade_frontalface_default.xml' ):

  # Access all JPG files in directory
  imlist = [os.path.join(root, name)
             for root, dirs, files in os.walk( os.path.join(os.getcwd(), prof_path) )
             for name in files
             if name.endswith((".jpeg", ".jpg"))]

  # Make output directory if it doesn't exist
  if( not os.path.isdir(output_path) ):
    os.makedirs( output_path )

  # Assuming all images are the same size, get dimensions of first image
  img = cv2.imread(imlist[0],0)
  h, w = img.shape[:2]
  # Book keeping
  N = len(imlist)
  accepted = 0.0
  rate = 0.0
  avg_face = None

  print 'Found', N, 'profiles'

  for im in range( N ):
    try:
      # Face match
      face = is_face(imlist[im], haar_map, w, h)
      # It returned a scaled face
      if( type(face) != bool ):
        if( avg_face == None ):
          avg_face = face
        else:
          alpha = 1.0/(float(accepted%101) + 1.0)
          # Do running average
          #alpha = 1.0/(float(accepted) + 1.0)
          #alpha = np.ceil(alpha*100.0)/100.0
          avg_face = cv2.addWeighted(face, alpha, avg_face, 1.0-alpha, 0.0, avg_face)
          #show_face(avg_face)

        accepted += 1.0
    except IOError:
      print 'Couldn\'t open image',im

    # Give Status report
    if ( accepted % 100 == 1 ):
      print 'Progress:', im,'/',N,'=',100.0*np.round(float(im)/float(N)*1000.0)/1000.0,'%'
      rate = accepted/float(im+1)
      print 'Accepted:', accepted,'/',N,'=',np.round(rate*100000.0)/1000.0,'%'
      if( avg_face != None ):
        #show_face( avg_face )
        print 'Saving Current Average'
        img_out = os.path.join(output_path, 'Average_Face_'+str(int(accepted))+'.png')
        cv2.imwrite( img_out, avg_face )
      # If we want to reset the average after every 100
      avg_face = None
  cv2.destroyAllWindows()

if(len(sys.argv) < 2):
  print 'Usage: python average.py profiles_dir/ <output/> <haar_cascade_default.xml>'
elif( len(sys.argv) == 2 ):
  main( sys.argv[1] )
elif( len(sys.argv) == 3 ):
  main( sys.argv[1], sys.argv[2] )
elif( len(sys.argv) == 4 ):
  main( sys.argv[1], sys.argv[2], sys.argv[3] )