import os
import torch
from tqdm import tqdm
from torchvision import transforms, datasets
from PIL import ImageOps
import numpy as np

class DatasetLoader():
    def __init__(self):
        self.dataset_loader = None

    def get_loader(self):
        return self.dataset_loader

class ImageNetDatasetLoader(DatasetLoader):
    def __init__(self, root_dir, height, width, batch_size):
        super(ImageNetDatasetLoader, self).__init__()
        data_transform = transforms.Compose([
            transforms.RandomResizedCrop((height, width)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                std=[0.229, 0.224, 0.225])
        ])
        imagenet_dataset = datasets.ImageFolder(root=root_dir, transform=data_transform)
        self.dataset_loader = torch.utils.data.DataLoader(imagenet_dataset, batch_size=batch_size, shuffle=True, num_workers=10)

class COCODataset(datasets.CocoDetection):
    def __init__(self, root, annFile, height, width, number_classes, anchors, scaling_factors, transform):
        super(COCODatasetLoader, self).__init__(root, annFile)
        self.height = height
        self.width = width
        self.number_classes = number_classes
        self.anchors = anchors
        self.scaling_factors = scaling_factors
        self.number_classes = number_classes
        self.transform = transform

    def __len__(self):
        return super(COCODatasetLoader, self).__len__()

    def resize_image(self, image, bboxes):
        curr_width, curr_height = image.size
        bboxes = np.array(bboxes)

        if curr_width > curr_height:
            scale = self.width / curr_width
            new_width = self.width
            new_height = int(scale * curr_height)
            padding = (self.height - new_height) // 2
            padding_tuple = (0, padding, 0, padding)

        else:
            scale = self.height / curr_height
            new_height = self.height
            new_width = int(scale * curr_width)
            padding = (self.width - new_width) // 2
            padding_tuple = (padding, 0, padding, 0)

        image = image.resize((new_width, new_height))
        image = ImageOps.expand(image, border=padding_tuple, fill=0)

        padding_tuple = np.array(padding_tuple[:2])
        bboxes = bboxes * scale
        bboxes[..., :2] += padding_tuple

        return image, bboxes

    def get_closest_anchor(self, bbox):
        def map_fn(anchor):
            return np.absolute(anchor[0] - bbox[2]) + np.absolute(anchor[1] - bbox[3])

        output = np.apply_along_axis(map_fn, 1, self.anchors)
        return np.argmin(output)

    def create_output_matrices(self, bboxes, classes):
        output_sizes = [np.array([self.width, self.height]) // i for i in self.scaling_factors]
        outputs = [np.zeros(np.append(size, (4 + 1 + self.number_classes) * 3)) for size in output_sizes]
        print(bboxes.shape[0])
        for i in range(bboxes.shape[0]):
            bbox = bboxes[i]
            anchor_idx = self.get_closest_anchor(bboxes[i])
            factor_idx = anchor_idx // 3
            factor = self.scaling_factors[factor_idx]
            bbox[:2] = bbox[:2] / (bbox[2:] / 2)
            softmax_classes = np.zeros(self.number_classes)
            softmax_classes[classes[i] - 1] = 1.0
            scaled = bbox[:2] / factor
            indices = np.array(np.floor_divide(bbox[:2], factor), dtype="int")
            offset = (anchor_idx % 3) * (4 + 1 + self.number_classes)
            print("Writing to ", indices[0], indices[1], offset, factor_idx)
            outputs[factor_idx][indices[0], indices[1], offset:(offset + 2)] = scaled - indices
            outputs[factor_idx][indices[0], indices[1], (offset + 2):(offset + 4)] = np.log(bbox[2:4] / self.anchors[anchor_idx])
            outputs[factor_idx][indices[0], indices[1], (offset + 4)] = 1.0
            outputs[factor_idx][indices[0], indices[1], (offset + 5):(offset + 5 + self.number_classes)] = softmax_classes

        return outputs

    def __getitem__(self, idx):
        image, metadata = super(COCODatasetLoader, self).__getitem__(idx)
        metadata = [(i["bbox"], i["category_id"]) for i in metadata if i["category_id"] <= self.number_classes]
        bboxes, classes = zip(*metadata)
        image, bboxes = self.resize_image(image, bboxes)
        outputs = self.create_output_matrices(bboxes, classes)

        return self.transform(image), transforms.ToTensor()(outputs)

class COCODatasetLoader(DatasetLoader):
    def __init__(self, root, annFile, height, width, number_classes, anchors, scaling_factors):
        super(DatasetLoader).__init__()
        data_transform

if __name__ == "__main__":
    root = "/data/nishantb/coco/val2017"
    annFile = "/data/nishantb/coco/annotations/instances_val2017.json"
    anchors = np.array([[ 20.1142666 ,  23.5920818 ],
       [ 61.54272606,  64.4243818 ],
       [ 78.43208441, 147.16067311],
       [173.79737463, 100.61049946],
       [130.99313932, 262.08766581],
       [419.42090078, 136.95277798],
       [267.46770012, 219.12108881],
       [258.48081984, 397.93091738],
       [461.50743771, 331.21772574]])
    factors = [8, 16, 32]
    obj = COCODatasetLoader(root, annFile, 512, 512, 80, anchors, factors)
    outputs = obj[0][1]
    for y in outputs:
        for j in range(3):
            offset = (4 + 1 + 80) * j
            idx, idy = np.where(y[..., offset + 4] > 0)
            print(idx.shape, idy.shape)
            print(np.squeeze(np.dstack((idx, idy))).shape)
            output = y[idx, idy, offset: offset + 4]
            print(output.shape)


# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# 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.
"""
CUDA ImageNet Demo
"""

import sys
import numpy as np
import time
import logging

import paddle
import paddle.fluid as fluid
import paddle.fluid.profiler as profiler
import paddle_fl.mpc as pfl_mpc
from paddle_fl.mpc.data_utils.data_utils import get_datautils

import resnet

#np.set_printoptions(threshold=np.inf)

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger("fluid")
logger.setLevel(logging.INFO)

mpc_protocol_name = 'aby3'
mpc_du = get_datautils(mpc_protocol_name)

# choose party 0 as nccl "server", which receive nccl id of other parties
# party 0 listens on ports listed in endpoints
# other parties connect to those ports and send their nccl id
role, server = sys.argv[1], sys.argv[2]
pfl_mpc.init(mpc_protocol_name, int(role),
        net_server_addr=server, endpoints="33784,45888", network_mode="nccl")
role = int(role)

# data preprocessing
BATCH_SIZE = 32
epoch_num = 1

x = pfl_mpc.data(name='x', shape=[BATCH_SIZE, 3, 224, 224], dtype='int64')
y = pfl_mpc.data(name='y', shape=[BATCH_SIZE, 1000], dtype='int64')

class Model(object):
    """
   lenet model: alexnet, vgg-16
    """

    def __int__(self):
        """
        init
        """
        pass

    def alexnet(self):
        """
        alexnet
        """
        conv_1 = pfl_mpc.layers.conv2d(input=x, num_filters=96, filter_size=11, stride=4, act='relu')
        pool_1 = pfl_mpc.layers.pool2d(input=conv_1, pool_size=3, pool_stride=2, pool_type='max')
        conv_2 = pfl_mpc.layers.conv2d(input=pool_1, num_filters=256, filter_size=5, padding=2, act='relu')
        pool_2 = pfl_mpc.layers.pool2d(input=conv_2, pool_size=3, pool_stride=2, pool_type='avg')

        conv_3 = pfl_mpc.layers.conv2d(input=pool_2, num_filters=384, filter_size=3, padding=1, act='relu')
        conv_4 = pfl_mpc.layers.conv2d(input=conv_3, num_filters=384, filter_size=3, padding=1, act='relu')
        conv_5 = pfl_mpc.layers.conv2d(input=conv_4, num_filters=256, filter_size=3, padding=1, act='relu')
        pool_3 = pfl_mpc.layers.pool2d(input=conv_5, pool_size=3, pool_stride=2, pool_type='max')

        reshape1 = pfl_mpc.layers.reshape(pool_3, [2, BATCH_SIZE, -1])
        fc_1 = pfl_mpc.layers.fc(input=reshape1, size=4096, act='relu')
        fc_2 = pfl_mpc.layers.fc(input=fc_1, size=4096, act='relu')
        fc_out = pfl_mpc.layers.fc(input=fc_2, size=1000)
        return fc_out

    def vgg_16(self):
        conv_1 = pfl_mpc.layers.conv2d(input=x, num_filters=64, filter_size=3, stride=1, padding=1, act='relu')
        conv_2 = pfl_mpc.layers.conv2d(input=conv_1, num_filters=64, filter_size=3, stride=1, padding=1, act='relu')
        pool_1 = pfl_mpc.layers.pool2d(input=conv_2, pool_size=2, pool_stride=2, pool_type='max')
        conv_3 = pfl_mpc.layers.conv2d(input=pool_1, num_filters=128, filter_size=3, stride=1, padding=1, act='relu')
        conv_4 = pfl_mpc.layers.conv2d(input=conv_3, num_filters=128, filter_size=3, stride=1, padding=1, act='relu')
        pool_2 = pfl_mpc.layers.pool2d(input=conv_4, pool_size=2, pool_stride=2, pool_type='max')
        conv_5 = pfl_mpc.layers.conv2d(input=pool_2, num_filters=256, filter_size=3, stride=1, padding=1, act='relu')
        conv_6 = pfl_mpc.layers.conv2d(input=conv_5, num_filters=256, filter_size=3, stride=1, padding=1, act='relu')
        conv_7 = pfl_mpc.layers.conv2d(input=conv_6, num_filters=256, filter_size=3, stride=1, padding=1, act='relu')
        pool_3 = pfl_mpc.layers.pool2d(input=conv_7, pool_size=2, pool_stride=2, pool_type='max')
        conv_8 = pfl_mpc.layers.conv2d(input=pool_3, num_filters=512, filter_size=3, stride=1, padding=1, act='relu')
        conv_9 = pfl_mpc.layers.conv2d(input=conv_8, num_filters=512, filter_size=3, stride=1, padding=1, act='relu')
        conv_10 = pfl_mpc.layers.conv2d(input=conv_9, num_filters=512, filter_size=3, stride=1, padding=1, act='relu')
        pool_4 = pfl_mpc.layers.pool2d(input=conv_10, pool_size=2, pool_stride=2, pool_type='max')
        conv_11 = pfl_mpc.layers.conv2d(input=pool_4, num_filters=512, filter_size=3, stride=1, padding=1, act='relu')
        conv_12 = pfl_mpc.layers.conv2d(input=conv_11, num_filters=512, filter_size=3, stride=1, padding=1, act='relu')
        conv_13 = pfl_mpc.layers.conv2d(input=conv_12, num_filters=512, filter_size=3, stride=1, padding=1, act='relu')
        pool_5 = pfl_mpc.layers.pool2d(input=conv_13, pool_size=2, pool_stride=2, pool_type='max')

        reshape1 = pfl_mpc.layers.reshape(pool_5, [2, BATCH_SIZE, -1])
        fc_1 = pfl_mpc.layers.fc(input=reshape1, size=4096, act='relu')
        fc_2 = pfl_mpc.layers.fc(input=fc_1, size=4096, act='relu')
        fc_out = pfl_mpc.layers.fc(input=fc_2, size=1000)
        return fc_out

model = Model()
out = model.alexnet()
#out = model.vgg_16()

#model = resnet.ResNet(layers=18)
#out = model.net(input=x, class_dim=1000)
cost, softmax = pfl_mpc.layers.softmax_with_cross_entropy(out, y, return_softmax=True, soft_label=True)

infer_program = fluid.default_main_program().clone(for_test=False)

avg_loss = pfl_mpc.layers.mean(cost)

lr = 0.01

optimizer = pfl_mpc.optimizer.SGD(learning_rate=lr)
optimizer.minimize(avg_loss)

## write main program into file for check
#with open('mpc_prog.txt', 'w') as f:
#    f.write(fluid.default_main_program().to_string(throw_on_error=False))

# dummy reader for profiling
# otherwise prepare encrypted data. please refer to MNIST demo
def dummy_reader(shape, rep):
    shape = list(shape)
    shape.insert(0, 2)
    shape = tuple(shape)
    def reader():
        for i in range(rep):
            yield np.zeros(shape, dtype=np.int64)
    return reader

feature_reader = dummy_reader((3, 224, 224), 160)
label_reader = dummy_reader((1000, ), 160)

batch_feature = mpc_du.batch(feature_reader, BATCH_SIZE, drop_last=True)
batch_label = mpc_du.batch(label_reader, BATCH_SIZE, drop_last=True)

test_batch_feature = mpc_du.batch(
    feature_reader, BATCH_SIZE, drop_last=True)
test_batch_label = mpc_du.batch(label_reader, BATCH_SIZE, drop_last=True)

place = fluid.CUDAPlace(0)

# async data loader
loader = fluid.io.DataLoader.from_generator(
    feed_list=[x, y], capacity=BATCH_SIZE)
batch_sample = paddle.reader.compose(batch_feature, batch_label)
loader.set_batch_generator(batch_sample, places=place)

test_loader = fluid.io.DataLoader.from_generator(
    feed_list=[x, y], capacity=BATCH_SIZE)
test_batch_sample = paddle.reader.compose(test_batch_feature, test_batch_label)
test_loader.set_batch_generator(test_batch_sample, places=place)

def decrypt_online(shares, shape):
    main_program = fluid.Program()
    startup_program = fluid.Program()
    with fluid.program_guard(main_program, startup_program):
        input = pfl_mpc.data(name='input', shape=shape[1:], dtype='int64')
        out = pfl_mpc.layers.reveal(input)

        exe = fluid.Executor(place)
        exe.run(fluid.default_startup_program())
        plain = exe.run(feed={'input': np.array(shares).reshape(shape)}, fetch_list=[out])
        return plain

# infer
def infer():
    """
    MPC infer
    """

    step = 0
    acc = 0
    start_time = time.time()
    for sample in test_loader():
        step += 1
        res = exe.run(program=infer_program,
                             feed=sample,
                             fetch_list=[softmax, y])
        pred = decrypt_online(res[0], (2, BATCH_SIZE, 1000))
        label = decrypt_online(res[1], (2, BATCH_SIZE, 1000))
        p = np.argmax(pred, axis=-1)
        l = np.argmax(label, axis=-1)
        match = BATCH_SIZE - np.count_nonzero(p - l)
        acc += match
        if step % 10 == 0:
            end_time = time.time()
            logger.info('MPC infer of step={}, cost time in seconds:{}'.format(
                step, (end_time - start_time)))
    end_time = time.time()
    logger.info('MPC infer acc:{} / {}'.format(acc, BATCH_SIZE * step))
    logger.info('MPC infer time in seconds:{}'.format((end_time - start_time)))

# train
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())

logger.info('MPC training start...')
for epoch_id in range(epoch_num):
    step = 0
    epoch_start_time = time.time()
    for sample in loader():
        step += 1
        step_start_time = time.time()
        #profiler.start_profiler("All")
        results = exe.run(feed=sample)
        #profiler.stop_profiler("total", "./profile")

        if step % 10 == 0:
            step_end_time = time.time()
            logger.info(
                'MPC training of epoch_id={} step={},  cost time in seconds:{}'
                .format(epoch_id, step, (step_end_time - step_start_time)))

    ### For each epoch, save infer program
    #mpc_model_basedir = "./mpc_model/"
    #mpc_model_dir = mpc_model_basedir + "epoch{}/party{}".format(epoch_id,
    #                                                             role)
    #fluid.io.save_inference_model(
    #    dirname=mpc_model_dir,
    #    feeded_var_names=["x", "y"],
    #    target_vars=[softmax],
    #    executor=exe,
    #    main_program=infer_program,
    #    model_filename="__model__")

    epoch_end_time = time.time()
    logger.info(
        'MPC training of epoch_num={} batch_size={}, cost time in seconds:{}'
        .format(epoch_id, BATCH_SIZE, (epoch_end_time - epoch_start_time)))

    # infer on test set
    infer()


#!/usr/bin/env python
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import time
import ctypes
import tempfile
import numpy
import h5py
import pyscf.lib as lib
from pyscf.lib import logger
import pyscf.ao2mo
from pyscf.cc import _ccsd

# t2,l2 as ijab

def kernel(mycc, eris, t1=None, t2=None, l1=None, l2=None,
           max_cycle=50, tol=1e-8, verbose=logger.INFO):
    cput0 = (time.clock(), time.time())
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(mycc.stdout, verbose)

    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2
    if l1 is None: l1 = t1
    if l2 is None: l2 = t2

    nocc, nvir = t1.shape
    saved = make_intermediates(mycc, t1, t2, eris)

    if mycc.diis:
        adiis = lib.diis.DIIS(mycc, mycc.diis_file)
        adiis.space = mycc.diis_space
    else:
        adiis = lambda t1,t2,*args: (t1, t2)
    cput0 = log.timer('CCSD lambda initialization', *cput0)

    conv = False
    for istep in range(max_cycle):
        l1new, l2new = update_amps(mycc, t1, t2, l1, l2, eris, saved)
        normt = numpy.linalg.norm(l1new-l1) + numpy.linalg.norm(l2new-l2)
        l1, l2 = l1new, l2new
        l1new = l2new = None
        if mycc.diis:
            l1, l2 = mycc.diis(l1, l2, istep, normt, 0, adiis)
        log.info('cycle = %d  norm(lambda1,lambda2) = %.6g', istep+1, normt)
        cput0 = log.timer('CCSD iter', *cput0)
        if normt < tol:
            conv = True
            break
    return conv, l1, l2

# l2, t2 as ijab
def make_intermediates(mycc, t1, t2, eris):
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    nov = nocc * nvir
    foo = eris.fock[:nocc,:nocc]
    fov = eris.fock[:nocc,nocc:]
    fvv = eris.fock[nocc:,nocc:]

    class _Saved(object):
        pass
    saved = _Saved()

# As we don't have l2 in memory, hold tau temporarily in memory
    w1 = fvv - numpy.einsum('ja,jb->ba', fov, t1)
    w2 = foo + numpy.einsum('ib,jb->ij', fov, t1)
    w3 = _cp(numpy.einsum('kc,jkbc->bj', fov, t2) * 2 + fov.T)
    w3 -= numpy.einsum('kc,kjbc->bj', fov, t2)
    w3 += reduce(numpy.dot, (t1.T, fov, t1.T))
    w4 = fov.copy()

    eris_ovvv = _cp(eris.ovvv)
    eris_ovvv = lib.unpack_tril(eris_ovvv.reshape(nov,-1))
    eris_ovvv = eris_ovvv.reshape(nocc,nvir,nvir,nvir)

    wovvv = numpy.empty((nocc,nvir,nvir,nvir))
    t2tmp = numpy.empty((nocc,nvir,nocc,nvir))
    for i in range(nocc):
        wovvv[i] = eris_ovvv[i].transpose(1,0,2) * 2
        t2tmp[i] = t2[i].transpose(2,0,1)
    #:wovvv += numpy.einsum('jabd,kjdc->kabc', eris_ovvv, t2) * -1.5
    tmp = lib.dot(t2tmp.reshape(nov,-1), wovvv.reshape(-1,nvir**2),
                  -1.5/2).reshape(-1,nvir,nvir,nvir)
    g2ovvv = tmp
    for i in range(nocc):
        wovvv[i] -= eris_ovvv[i].transpose(1,2,0)
        wovvv[i] += tmp[i].transpose(1,2,0)
        g2ovvv[i] = eris_ovvv[i]*2
        g2ovvv[i] -= eris_ovvv[i].transpose(1,2,0)
    tmp = t2tmp = None

    w1 += numpy.einsum('jcba,jc->ba', eris_ovvv, t1*2)
    w1 -= numpy.einsum('jabc,jc->ba', eris_ovvv, t1)
    #:w3 += numpy.einsum('kdcb,kjdc->bj', eris_ovvv, theta)
    theta = numpy.empty(t2.shape)
    for i in range(nocc):
        theta[i] = t2[i] * 2
        theta[i] -= t2[i].transpose(0,2,1)
        lib.dot(eris_ovvv[i].reshape(-1,nvir).T,
                _cp(theta[i].reshape(nocc,-1)).T, 1, w3, 1)

    theta = _cp(theta.transpose(0,2,1,3))
    #:vkbca = numpy.einsum('jdca,kbjd->kbca', g2ovvv, theta)
    vkbca = lib.dot(_cp(theta.reshape(nov,-1)),
                    g2ovvv.reshape(-1,nvir*nvir)).reshape(-1,nvir,nvir,nvir)
    for i in range(nocc):
        wovvv[i] += vkbca[i].transpose(2,0,1)
        wovvv[i] -= vkbca[i].transpose(2,1,0) * .5
    vkabc = None

    #:wOVov = numpy.einsum('jbcd,kd->jbkc', eris_ovvv, t1)
    #:wOvOv = numpy.einsum('jdcb,kd->jbkc', eris_ovvv, -t1)
    wOVov = lib.dot(eris_ovvv.reshape(-1,nvir),
                    t1.T).reshape(-1,nvir,nvir,nocc).transpose(0,1,3,2).copy()
    for i in range(nocc):
        g2ovvv[i] = eris_ovvv[i].transpose(1,2,0) * 2
    wOvOv = lib.dot(g2ovvv.reshape(-1,nvir),
                    -t1.T, .5).reshape(-1,nvir,nvir,nocc).transpose(0,1,3,2).copy()
    for i in range(nocc):
        g2ovvv[i] -= eris_ovvv[i].transpose(1,0,2)
    eris_ovov = _cp(_cp(eris.ovov).transpose(0,2,1,3))
    tau = _ccsd.make_tau(t2, t1, t1)
    #:wooov[:,j0:j1] = numpy.einsum('icbd,jkbd->ijkc', g2ovvv, tau)
    #:woooo[:,:,j0:j1] = numpy.einsum('icjd,klcd->ijkl', eris_ovov, tau)
    tmp = lib.dot(g2ovvv.reshape(-1,nvir**2), tau.reshape(-1,nvir**2).T)
    wooov = _cp(tmp.reshape(-1,nvir,nocc,nocc).transpose(0,2,3,1))
    woooo = lib.dot(eris_ovov.reshape(-1,nvir**2),
                    tau.reshape(-1,nvir**2).T).reshape(-1,nocc,nocc,nocc)
    eris_ovov = eris_ovvv = g2ovvv = tau = tmp = None

    eris_ooov = _cp(eris.ooov)
    eris_ovoo = _cp(eris.ovoo)
    #:woooo += numpy.einsum('icjl,kc->ijkl', eris_ovoo, t1)
    #:wOVov += numpy.einsum('jblk,lc->jbkc', eris_ovoo, -t1)
    for i in range(nocc):
        woooo[i] += lib.dot(t1, eris_ovoo[i].reshape(nvir,-1)).reshape((nocc,)*3).transpose(1,0,2)
    lib.dot(eris_ovoo.reshape(-1,nocc), t1, -1, wOVov.reshape(-1,nvir), 1)
    #:wooov -= numpy.einsum('ibjl,lkcb->ijkc', eris_ovoo*1.5, t2)
    t2tmp = numpy.empty((nocc,nvir,nocc,nvir))
    for i in range(nocc):
        t2tmp[i] = t2[i].transpose(2,0,1)
    tmp_ooov = _cp(-eris_ooov.transpose(2,0,1,3)).reshape(-1,nov)
    lib.dot(tmp_ooov, t2tmp.reshape(nov,-1), 1.5, wooov.reshape(-1,nov), 1)
    t2tmp = None

    g2ooov, tmp_ooov = tmp_ooov.reshape(nocc,nocc,nocc,nvir), None
    g2ooov += eris_ooov * 2
    #:vikjc = numpy.einsum('iklb,jlcb->ikjc', g2ooov, theta)
    vikjc = lib.dot(g2ooov.reshape(-1,nov), theta.reshape(-1,nov).T)
    vikjc = vikjc.reshape(nocc,nocc,nocc,nvir)
    wooov += vikjc.transpose(0,2,1,3)
    wooov -= vikjc*.5
    g2ooov = vikjc = eris_ovoo = None

    w2 += numpy.einsum('ijkb,kb->ij', eris_ooov, t1) * 2
    w2 -= numpy.einsum('kjib,kb->ij', eris_ooov, t1)
    #:w3 -= numpy.einsum('kjlc,klbc->bj', eris_ooov, theta)
    for i in range(nocc):
        lib.dot(_cp(theta[i].transpose(1,2,0)).reshape(-1,nvir).T,
                eris_ooov[i].reshape(nocc,-1).T, -1, w3, 1)
    #:woooo += numpy.einsum('ikjc,lc->ijkl', eris_ooov, t1)
    #:wOvOv += numpy.einsum('jklb,lc->jbkc', eris_ooov, t1)
    woooo += lib.dot(eris_ooov.reshape(-1,nvir),
                     t1.T).reshape((-1,nocc,nocc,nocc)).transpose(0,2,1,3)
    for i in range(nocc):
        lib.dot(_cp(eris_ooov[i].transpose(2,0,1)).reshape(-1,nocc),
                t1, 1, wOvOv[i].reshape(-1,nvir), 1)
        wooov[i] += eris_ooov[i].transpose(1,0,2)*2
        wooov[i] -= eris_ooov[i]
    eris_ooov = theta = None

    eris_ovov = _cp(eris.ovov)
    g2ovov = numpy.empty((nocc,nocc,nvir,nvir))
    for i in range(nocc):
        g2ovov[i] = eris_ovov[i].transpose(1,0,2)*2
        g2ovov[i] -= eris_ovov[i].transpose(1,2,0)
    tmpw4 = numpy.einsum('klcd,ld->kc', g2ovov, t1)
    #:w1 -= numpy.einsum('kcja,kjcb->ba', g2ovov, t2)
    w1 -= lib.dot(t2.reshape(-1,nvir).T, g2ovov.reshape(-1,nvir))
    w1 -= numpy.einsum('ja,jb->ba', tmpw4, t1)
    #:w2 += numpy.einsum('ibkc,jkbc->ij', g2ovov, t2)
    w2 += lib.dot(g2ovov.reshape(nocc,-1), t2.reshape(nocc,-1).T)
    w2 += numpy.einsum('ib,jb->ij', tmpw4, t1)
    w3 += reduce(numpy.dot, (t1.T, tmpw4, t1.T))
    w4 += tmpw4
    vOVov  = eris_ovov.copy()
    #:vOVov += numpy.einsum('jbld,klcd->jbkc', g2ovov, t2)
    #:vOVov -= numpy.einsum('jbld,kldc->jbkc', eris_ovov, t2)
    lib.dot(_cp(g2ovov.transpose(0,2,1,3)).reshape(-1,nov),
            _cp(t2.transpose(0,2,1,3).reshape(nov,-1).T), 1,
            vOVov.reshape(nov,-1), 1)
    lib.dot(eris_ovov.reshape(-1,nov),
            _cp(t2.transpose(0,3,1,2).reshape(nov,-1).T), -1,
            vOVov.reshape(nov,-1), 1)
    g2ovov = None

    #:tmp = numpy.einsum('jbld,kd->ljbk', eris_ovov, t1)
    #:wOVov -= numpy.einsum('ljbk,lc->jbkc', tmp, t1)
    #:tmp = numpy.einsum('jdlb,kd->ljbk', eris_ovov, t1)
    #:wOvOv += numpy.einsum('ljbk,lc->jbkc', tmp, t1)
    tmp = numpy.empty((nocc,nvir,nocc))
    for j in range(nocc):
        lib.dot(_cp(eris_ovov[j].transpose(1,0,2)).reshape(-1,nvir),
                t1.T, 1, tmp.reshape(-1,nocc))
        lib.dot(tmp.reshape(nocc,-1).T, t1, -1, wOVov[j].reshape(-1,nvir), 1)
        lib.dot(eris_ovov[j].reshape(nvir,-1).T, t1.T, 1,
                tmp.reshape(-1,nocc))
        lib.dot(tmp.reshape(nocc,-1).T, t1, 1, wOvOv[j].reshape(-1,nvir), 1)
    tmp = None

    #:vOvOv = numpy.einsum('jdlb,kldc->jbkc', eris_ovov, t2)
    ovovtmp = _cp(eris_ovov.transpose(0,3,2,1).reshape(-1,nov))
    vOvOv = numpy.empty((nocc,nvir,nocc,nvir))
    for j in range(nocc):
        lib.dot(t2[j].reshape(-1,nvir).T, ovovtmp.T, 1,
                vOvOv[j].reshape(nvir,-1))
        vOvOv[j] -= eris.oovv[j].transpose(2,0,1)
    ovovtmp = eris_ovov = None
    vOvOv = lib.transpose(vOvOv.reshape(nov,-1)).reshape(nocc,nvir,nocc,nvir)
    wOVov += vOVov
    wOvOv += vOvOv
    saved.wOVov = wOVov
    saved.wOvOv = wOvOv
    ovovtmp = wOVov = wOvOv = eris_ovov = None

    ov2 = vOVov*2 + vOvOv
    w3 += numpy.einsum('kcjb,kc->bj', ov2, t1)
    #:wooov += numpy.einsum('ibjc,kb->ijkc', ov2, t1)
    #:wovvv -= numpy.einsum('jakb,jc->kabc', ov2, t1)
    for i in range(nocc):
        wooov[i] += lib.dot(t1, ov2[i].reshape(nvir,-1)).reshape(nocc,nocc,nvir).transpose(1,0,2)
    lib.dot(_cp(ov2.transpose(0,2,1,3).reshape(nocc,-1)).T,
            t1, -1, wovvv.reshape(-1,nvir), 1)
    ov2 = None
    ov1 = vOvOv*2 + vOVov
    #:wooov -= numpy.einsum('ibkc,jb->ijkc', ov1, t1)
    #:wovvv += numpy.einsum('jakc,jb->kabc', ov1, t1)
    for i in range(nocc):
        lib.dot(t1, ov1[i].reshape(nvir,-1), -1, wooov[i].reshape(nocc,-1), 1)
    wovvv += lib.dot(_cp(ov1.reshape(nocc,-1)).T,
                     t1).reshape(nvir,-1,nvir,nvir).transpose(1,0,3,2)
    ov1 = None

    woooo += _cp(eris.oooo).transpose(0,2,1,3)
    saved.woooo = woooo
    saved.wooov = wooov
    woooo = wooov = None

    w3 += numpy.einsum('bc,jc->bj', w1, t1)
    w3 -= numpy.einsum('kj,kb->bj', w2, t1)

    eris_ooov = _cp(eris.ooov)
    g2ooov = eris_ooov * 2
    g2ooov -= eris_ooov.transpose(2,0,1,3)
    #:tmp = numpy.einsum('kjla,jb->kabl', g2ooov, t1)
    #:wovvv = numpy.einsum('kabl,lc->kabc', tmp, t1)
    #:wovvv += numpy.einsum('kjla,jlbc->kabc', g2ooov, t2)
    tmp = lib.dot(g2ooov.reshape(nocc,-1).T, t1).reshape(-1,nocc,nvir,nvir).transpose(0,2,3,1)
    lib.dot(_cp(tmp.reshape(-1,nocc)), t1, 1, wovvv.reshape(-1,nvir), 1)
    tmp = None
    lib.dot(_cp(g2ooov.transpose(0,2,1,3).reshape(nocc**2,-1)).T,
            t2.reshape(nocc**2,-1), 1, wovvv.reshape(nov,-1), 1)
    g2ooov = eris_ooov = vOVov = vOvOv = None

    saved.wovvv = wovvv
    saved.w1 = w1
    saved.w2 = w2
    saved.w3 = w3
    saved.w4 = w4
    return saved

# update L1, L2
def update_amps(mycc, t1, t2, l1, l2, eris=None, saved=None):
    if saved is None:
        saved = make_intermediates(mycc, t1, t2, eris)
    time1 = time0 = time.clock(), time.time()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    nov = nocc * nvir
    foo = eris.fock[:nocc,:nocc]
    fov = eris.fock[:nocc,nocc:]
    fvv = eris.fock[:nocc,:nocc]

    #:mba = numpy.einsum('klca,klcb->ba', l2, t2*2-t2.transpose(0,1,3,2))
    #:mij = numpy.einsum('ikcd,jkcd->ij', l2, t2*2-t2.transpose(0,1,3,2))
    #:theta = t2*2 - t2.transpose(0,1,3,2)
    theta = _ccsd.make_0132(t2, t2, 2, -1)
    mba = lib.dot(theta.reshape(-1,nvir).T, l2.reshape(-1,nvir))
    mij = lib.dot(l2.reshape(nocc,-1), theta.reshape(nocc,-1).T)
    theta = None
    mba1 = numpy.einsum('jc,jb->bc', l1, t1) + mba
    mij1 = numpy.einsum('kb,jb->kj', l1, t1) + mij
    mia1 =(t1 + numpy.einsum('kc,jkbc->jb', l1, t2) * 2
         - numpy.einsum('kc,jkcb->jb', l1, t2)
         - reduce(numpy.dot, (t1, l1.T, t1))
         - numpy.einsum('bd,jd->jb', mba, t1)
         - numpy.einsum('lj,lb->jb', mij, t1))

    tmp = mycc.add_wvvVV(numpy.zeros_like(l1), l2, eris)
    l2new = numpy.empty((nocc,nocc,nvir,nvir))
    ij = 0
    for i in range(nocc):
        for j in range(i):
            tmp1 = tmp[ij] * .5  # *.5 because of l2+l2.transpose(1,0,3,2) later
            l2new[i,j] = tmp1
            l2new[j,i] = tmp1.T
            ij += 1
        l2new[i,i] = tmp[ij] * .5
        ij += 1
    l1new =(numpy.einsum('ijab,jb->ia', l2new, t1) * 4
          - numpy.einsum('jiab,jb->ia', l2new, t1) * 2)
    tmp = tmp1 = None

    l1new += eris.fock[:nocc,nocc:]
    l1new += numpy.einsum('ib,ba->ia', l1, saved.w1)
    l1new -= numpy.einsum('ja,ij->ia', l1, saved.w2)
    l1new -= numpy.einsum('ik,ka->ia', mij, saved.w4)
    l1new -= numpy.einsum('ca,ic->ia', mba, saved.w4)
    l1new += numpy.einsum('ijab,bj->ia', l2, saved.w3) * 2
    l1new -= numpy.einsum('ijba,bj->ia', l2, saved.w3)

    l2new += numpy.einsum('ia,jb->ijab', l1, saved.w4)
    #:l2new += numpy.einsum('jibc,ca->jiba', l2, saved.w1)
    #:l2new -= numpy.einsum('kiba,jk->jiba', l2, saved.w2)
    lib.dot(l2.reshape(-1,nvir), saved.w1, 1, l2new.reshape(-1,nvir), 1)
    lib.dot(saved.w2, l2.reshape(nocc,-1),-1, l2new.reshape(nocc,-1), 1)

    eris_ooov = _cp(eris.ooov)
    l1new -= numpy.einsum('jkia,kj->ia', eris_ooov, mij1) * 2
    l1new += numpy.einsum('ikja,kj->ia', eris_ooov, mij1)
    #:l2new -= numpy.einsum('ka,kijb->jiba', l1, eris_ooov)
    lib.dot(_cp(eris_ooov.transpose(0,2,1,3).reshape(nocc,-1)).T,
            l1, -1, l2new.reshape(-1,nvir), 1)
    eris_ooov = None

    tau = _ccsd.make_tau(t2, t1, t1)
    #:l2tau = numpy.einsum('ijcd,klcd->ijkl', l2, tau)
    l2tau = lib.dot(l2.reshape(nocc**2,-1),
                    tau.reshape(nocc**2,-1).T).reshape((nocc,)*4)
    tau = None
    l2t1 = numpy.einsum('ijcd,kc->ijkd', l2, t1)

    eris_ovvv = _cp(eris.ovvv)
    eris_ovvv = lib.unpack_tril(eris_ovvv.reshape(nov,-1))
    eris_ovvv = eris_ovvv.reshape(nocc,nvir,nvir,nvir)

    l1new += numpy.einsum('iabc,bc->ia', eris_ovvv, mba1) * 2
    l1new -= numpy.einsum('ibca,bc->ia', eris_ovvv, mba1)
    #:l2new += numpy.einsum('ic,jbac->jiba', l1, eris_ovvv)
    tmp = lib.dot(l1, eris_ovvv.reshape(-1,nvir).T).reshape(nocc,-1,nvir,nvir)
    for i in range(nocc):
        l2new[i] += tmp[i].transpose(0,2,1)
    #:m4 = numpy.einsum('ijkd,kadb->ijab', l2t1, eris_ovvv)
    m4 = tmp
    lib.dot(_cp(l2t1.reshape(nocc*nocc,-1)),
            _cp(eris_ovvv.transpose(0,2,1,3).reshape(-1,nvir**2)),
            1, m4.reshape(nocc*nocc,-1))
    l2new -= m4
    l1new -= numpy.einsum('ijab,jb->ia', m4, t1) * 2
    l1new -= numpy.einsum('ijab,ia->jb', m4, t1) * 2
    l1new += numpy.einsum('jiab,jb->ia', m4, t1)
    l1new += numpy.einsum('jiab,ia->jb', m4, t1)
    eris_ovvv = tmp = None

    eris_ovov = _cp(eris.ovov)
    l1new += numpy.einsum('jb,iajb->ia', l1, eris_ovov) * 2
    #:l2new -= numpy.einsum('jbic,ca->jiba', eris_ovov, mba1)
    #:l2new -= numpy.einsum('kajb,ik->ijab', eris_ovov, mij1)
    tmp = lib.dot(eris_ovov.reshape(-1,nvir), mba1).reshape(nocc,nvir,nocc,nvir)
    lib.dot(mij1, eris_ovov.reshape(nocc,-1), 1, tmp.reshape(nocc,-1), 1)
    tmp_oovv = numpy.empty((nocc,nocc,nvir,nvir))
    for i in range(nocc):
        tmp_oovv[i] = eris_ovov[i].transpose(1,0,2) * .5
        l2new[i] += tmp_oovv[i]
        l2new[i] -= tmp[i].transpose(1,0,2)
    tmp = None
    l1new += numpy.einsum('iajb,jb->ia', eris_ovov, mia1) * 2
    l1new -= numpy.einsum('ibja,jb->ia', eris_ovov, mia1)
    #:m4 = numpy.einsum('kalb,ijkl->ijab', eris_ovov, l2tau)
    lib.dot(l2tau.reshape(nocc*nocc,-1), tmp_oovv.reshape(-1,nvir**2),
            1, m4.reshape(nocc**2,-1))
    l2new += m4
    l1new += numpy.einsum('ijab,jb->ia', m4, t1) * 4
    l1new -= numpy.einsum('ijba,jb->ia', m4, t1) * 2
    eris_ovov = m4 = tmp_oovv = None

    eris_oovv = _cp(eris.oovv)
    l1new -= numpy.einsum('jb,ijba->ia', l1, eris_oovv)
    eris_oovv = None

    saved_wooov = _cp(saved.wooov)
    #:l1new -= numpy.einsum('jkca,ijkc->ia', l2, saved_wooov)
    l1new -= lib.dot(saved_wooov.reshape(nocc,-1), l2.reshape(-1,nvir))
    saved_wovvv = _cp(saved.wovvv)
    #:l1new += numpy.einsum('kibc,kabc->ia', l2, saved_wovvv)
    for j in range(nocc):
        l1new += lib.dot(l2[j].reshape(nocc,-1),
                         saved_wovvv[j].reshape(nvir,-1).T)
    saved_wooov = saved_wovvv = None

    saved_wOvOv = _cp(saved.wOvOv)
    tmp_ovov = _cp(saved.wOVov) * 2
    tmp_ovov += saved_wOvOv
    #:tmp = l2.transpose(0,2,1,3) - l2.transpose(0,3,1,2)*.5
    #:l2new += numpy.einsum('kcia,kcjb->jiba', tmp, tmp_ovov)
    tmp = numpy.empty((nocc,nvir,nocc,nvir))
    for i in range(nocc):
        tmp[i] = l2[i].transpose(2,0,1)*-.5
        tmp[i] += l2[i].transpose(1,0,2)
    tmp = lib.dot(tmp_ovov.reshape(-1,nov),
                  tmp.reshape(nov,-1)).reshape(-1,nvir,nocc,nvir)
    #:tmp = numpy.einsum('jkca,ibkc->ijab', l2, saved_wOvOv)
    for i in range(nocc):
        l2new[i] += tmp[i].transpose(1,0,2)
        tmp_ovov[i] = l2[i].transpose(2,0,1)
    lib.dot(saved_wOvOv.reshape(-1,nov), tmp_ovov.reshape(nov,-1),
            1, tmp.reshape(nov,-1))
    for i in range(nocc):
        l2new[i] += tmp[i].transpose(1,2,0)
        l2new[i] += tmp[i].transpose(1,0,2) * .5
    saved_wOvOv = tmp = tmp_ovov = None

    saved_woooo = _cp(saved.woooo)
    #:m3 = numpy.einsum('klab,ijkl->ijab', l2, saved_woooo)
    m3 = lib.dot(saved_woooo.reshape(-1,nocc**2),
                 l2.reshape(nocc**2,-1), .5).reshape(-1,nocc,nvir,nvir)
    l2new += m3
    l1new += numpy.einsum('ijab,jb->ia', m3, t1) * 4
    l1new -= numpy.einsum('ijba,jb->ia', m3, t1) * 2
    saved_woooo = m3 = None

    mo_e = eris.fock.diagonal()
    eia = lib.direct_sum('i-j->ij', mo_e[:nocc], mo_e[nocc:])
    l1new /= eia
    l1new += l1

#    l2new = l2new + l2new.transpose(1,0,3,2)
#    l2new /= lib.direct_sum('ia+jb->ijab', eia, eia)
#    l2new += l2
    ij = 0
    for i in range(nocc):
        for j in range(i):
            dab = lib.direct_sum('a+b->ab', eia[i], eia[j])
            tmp = (l2new[i,j]+l2new[j,i].T) / dab + l2[i,j]
            l2new[i,j] = tmp
            l2new[j,i] = tmp.T
            ij += 1
        dab = lib.direct_sum('a+b->ab', eia[i], eia[i])
        l2new[i,i] = (l2new[i,i]+l2new[i,i].T)/dab + l2[i,i]
        ij += 1

    time0 = log.timer_debug1('update l1 l2', *time0)
    return l1new, l2new

def _cp(a):
    return numpy.array(a, copy=False, order='C')

if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf.cc import ccsd
    from pyscf import ao2mo

    mol = gto.M()
    mf = scf.RHF(mol)

    mcc = ccsd.CCSD(mf)

    numpy.random.seed(12)
    nocc = 5
    nmo = 12
    nvir = nmo - nocc
    eri0 = numpy.random.random((nmo,nmo,nmo,nmo))
    eri0 = ao2mo.restore(1, ao2mo.restore(8, eri0, nmo), nmo)
    fock0 = numpy.random.random((nmo,nmo))
    fock0 = fock0 + fock0.T + numpy.diag(range(nmo))*2
    t1 = numpy.random.random((nocc,nvir))
    t2 = numpy.random.random((nocc,nocc,nvir,nvir))
    t2 = t2 + t2.transpose(1,0,3,2)
    l1 = numpy.random.random((nocc,nvir))
    l2 = numpy.random.random((nocc,nocc,nvir,nvir))
    l2 = l2 + l2.transpose(1,0,3,2)

    eris = lambda:None
    eris.oooo = eri0[:nocc,:nocc,:nocc,:nocc].copy()
    eris.ooov = eri0[:nocc,:nocc,:nocc,nocc:].copy()
    eris.ovoo = eri0[:nocc,nocc:,:nocc,:nocc].copy()
    eris.oovv = eri0[:nocc,:nocc,nocc:,nocc:].copy()
    eris.ovov = eri0[:nocc,nocc:,:nocc,nocc:].copy()
    idx = numpy.tril_indices(nvir)
    eris.ovvv = eri0[:nocc,nocc:,nocc:,nocc:][:,:,idx[0],idx[1]].copy()
    eris.vvvv = pyscf.ao2mo.restore(4,eri0[nocc:,nocc:,nocc:,nocc:],nvir)
    eris.fock = fock0

    saved = make_intermediates(mcc, t1, t2, eris)
    l1new, l2new = update_amps(mcc, t1, t2, l1, l2, eris, saved)
    print(abs(l1new).sum()-38172.7896467303)
    print(numpy.dot(l1new.flatten(), numpy.arange(35)) - 739312.005491083)
    print(numpy.dot(l1new.flatten(), numpy.sin(numpy.arange(35)))-7019.50937051188)
    print(numpy.dot(numpy.sin(l1new.flatten()), numpy.arange(35))-69.6652346635955)

    print(abs(l2new).sum()-72035.4931071527)
    print(abs(l2new-l2new.transpose(1,0,3,2)).sum())
    print(numpy.dot(l2new.flatten(), numpy.arange(35**2)) - 48427109.5409886)
    print(numpy.dot(l2new.flatten(), numpy.sin(numpy.arange(35**2)))-137.758016736487)
    print(numpy.dot(numpy.sin(l2new.flatten()), numpy.arange(35**2))-507.656936701192)

    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]

    mol.basis = 'cc-pvdz'
    mol.build()
    rhf = scf.RHF(mol)
    rhf.conv_tol = 1e-16
    rhf.scf()

    mcc = ccsd.CCSD(rhf)
    mcc.conv_tol = 1e-12
    ecc, t1, t2 = mcc.kernel()

    nmo = rhf.mo_energy.size
    fock0 = numpy.diag(rhf.mo_energy)
    nocc = mol.nelectron // 2
    nvir = nmo - nocc

    eris = mcc.ao2mo()
    conv, l1, l2 = kernel(mcc, eris, t1, t2, tol=1e-8)
    print(numpy.linalg.norm(l1)-0.0132626841292)
    print(numpy.linalg.norm(l2)-0.212575609057)

    import ccsd_rdm
    dm1 = ccsd_rdm.make_rdm1(mcc, t1, t2, l1, l2)
    dm2 = ccsd_rdm.make_rdm2(mcc, t1, t2, l1, l2)
    h1 = reduce(numpy.dot, (rhf.mo_coeff.T, rhf.get_hcore(), rhf.mo_coeff))
    eri = pyscf.ao2mo.full(rhf._eri, rhf.mo_coeff)
    eri = pyscf.ao2mo.restore(1, eri, nmo).reshape((nmo,)*4)
    e1 = numpy.einsum('pq,pq', h1, dm1)
    e2 = numpy.einsum('pqrs,pqrs', eri, dm2) * .5
    print e1+e2+mol.energy_nuc() - rhf.e_tot - ecc


#!/usr/bin/env python
# encoding: utf-8
"""
Sampler for Multi Pixel Gibbs modelling, using a library-based SP
estimates rather than MCMC.
"""

from collections import OrderedDict

import numpy as np

from astropy.table import Table, hstack, Column
from astropy.utils.console import ProgressBar

from sedbot.chain import MultiPixelChain

from sedbot.utils.timer import Timer

class MultiPixelLibraryGibbsBgSampler(object):
    """Library-in-Gibbs sampler for two-level hierarchical models of pixel
    parameters and a linearly sampled background parameter.

    Parameters
    ----------
    model : :class:`sedbot.multipix.libmodels.LibraryModel`
        A LibraryModel instance that encapsulates the library SEDs and the
        observed SEDs.
    """
    def __init__(self, model, ncpu=1):
        super(MultiPixelLibraryGibbsBgSampler, self).__init__()
        self.model = model
        self._n_cpu_estimator = ncpu

    def sample(self, n_iter,
               B0=None,
               chain=None):
        """Sample for `n_iter` Gibbs steps.

        Parameters
        ----------
        n_iter : int
            Number of Gibbs steps (number of full iterations over all
            parameters)
        B0 : ndarray
            Initial values for the background, a ``(n_band)`` array.
            Background is in units of flux per arcsec^2.
        chain : :class:`sedbot.chain.MultiPixelChain`
            A previously-built chain, whose last sample will be used
            as the starting points for this sampling (in liue of setting
            `theta0`, `phi0` and `B0` manually).
        """
        if chain is not None:
            B0 = np.empty(self.model.n_bands, dtype=np.float)
            for i, (n, instr) in enumerate(zip(self.model.observed_bands,
                                               self.model.instruments)):
                name = "B__{0}__{1}".format(instr, n)
                B0[i] = chain[name][-1]
        else:
            assert B0 is not None

        # Initialize chains
        self._init_chains(n_iter)

        with ProgressBar(n_iter) as bar:
            for i in xrange(n_iter):
                with Timer() as timer:  # NOQA for timing
                    # Sample stellar populations
                    self._estimate_theta(i, B=B0)
                    B0 = None

                    # Estimate background given current SP state
                    self._estimate_background(i)
                # print "Step {0:d} took {1:.2f} sec".format(i, timer.interval)

                bar.update()

    def _init_chains(self, n_iter):
        """Initialize memory for the chain."""
        # FSPS SP Parameter estimates
        self.theta_chain = np.empty((n_iter + 1,
                                     self.model.n_pix,
                                     self.model.n_theta),
                                    dtype=np.float)
        self.theta_chain.fill(np.nan)

        # Background estimates
        self.B_chain = np.empty((n_iter + 1,
                                 self.model.n_bands),
                                dtype=np.float)
        self.B_chain.fill(np.nan)

        # FSPS computed metadata values
        self.blob_chain = np.empty((n_iter + 1,
                                    self.model.n_pix,
                                    len(self.model.meta_params)),
                                   dtype=np.float)

        # M/L ratios in all library bands
        self.ml_chain = np.empty((n_iter + 1,
                                  self.model.n_pix,
                                  len(self.model.library_bands)),
                                 dtype=np.float)

        # Model SED
        self.sed_chain = np.empty((n_iter + 1,
                                   self.model.n_pix,
                                   len(self.model.library_bands)),
                                  dtype=np.float)

    def _estimate_theta(self, k, B=None):
        """Estimate the stellar population parameters.

        Parameters
        ----------
        B : ndarray
            Optional values for the background, a ``(n_band)`` array.
            Background is in units of flux per arcsec^2.
            If `None` then the background is read from the previous value of
            the background chain.
        """
        n_pixels = self.model.n_pix

        if B is None:
            B = self.B_chain[k - 1, :]

        # Compute SP parameters for each SED, with background subtracted
        for i in xrange(n_pixels):
            # print "Pixel {0:d}".format(i)
            le = self.model.estimator_for_pixel(i, B,
                                                ncpu=self._n_cpu_estimator)
            # marginal estimate of model parameters
            for j, name in enumerate(self.model.theta_params):
                self.theta_chain[k, i, j] = le.estimate(name, p=(0.5,))[0]
                print i, name, self.theta_chain[k, i, j]
            # marginal estimate of metadata parameters
            for j, name in enumerate(self.model.meta_params):
                self.blob_chain[k, i, j] = le.estimate_meta(name, p=(0.5,))[0]
            # marginal estimate of M/L
            for j, band in enumerate(self.model.library_bands):
                self.ml_chain[k, i, j] = le.estimate_ml(band, p=(0.5,))[0]
            # marginal estimate of SED
            for j, band in enumerate(self.model.library_bands):
                self.sed_chain[k, i, j] = le.estimate_flux(band, p=(0.5,))[0]

    def _estimate_background(self, k):
        """Estimate the background given the current stellar population
        parameter state.

        The background update is *always* done after the theta update.

        Parameters
        ----------
        k : int
            Index of the current Gibbs step.
        """
        # Get the model SED matching the observed bands
        model_seds = self.sed_chain[k, :, self.model.band_indices]
        # transpose to to make a (n_sed, n_band) array)... I know right?
        B_new = self.model.estimate_background(model_seds.T)
        self.B_chain[k, :] = B_new

    @property
    def table(self):
        """A :class:`MultiPixelChain` representing the Gibbs chain.

        The intention is for the library sampler to have a data output
        similar to the MH-in-Gibbs sampler.
        """
        meta = OrderedDict((
            ('observed_bands', self.model.observed_bands),
            ('instruments', self.model.instruments),
            ('computed_bands', self.model.library_bands),
            ('msun_ab', self.model.msun_ab),
            ('band_indices', self.model.band_indices),
            ('theta_params', self.model.theta_params),
            ('meta_params', self.model.meta_params),
            ('sed', self.model._seds),
            ('sed_err', self.model._errs),
            ('pixels', self.model.pixel_metadata),
            ('area', self.model._areas),
            ('d', self.model.d)))

        # Make tables for theta chain, B chain, and blob chain (that includes
        # SED and M/L datasets)
        # TODO we should probably just change the chain dtypes instead
        # original table - (step, pixel, parameter)
        # output table - (step, parameter, pixel)
        theta_table = Table(np.swapaxes(self.theta_chain, 1, 2),
                            names=self.model.theta_params,
                            meta=meta)

        background_names = ["B__{0}__{1}".format(n, b)
                            for n, b in zip(self.model.instruments,
                                            self.model.observed_bands)]
        B_table = Table(self.B_chain, names=background_names)

        meta_table = Table(np.swapaxes(self.blob_chain, 1, 2),
                           names=self.model.meta_params)

        # model_sed column is of shape (step, pixel, band)
        # conforms to MultiPixelChain standard
        meta_table.add_column(Column(name='model_sed',
                                     data=self.sed_chain))

        # output table (step, band, pixel)
        # name columns to be different from flux
        ml_names = ["logML_{0}".format(b) for b in self.model.library_bands]
        ml_table = Table(np.swapaxes(self.ml_chain, 1, 2),
                         names=ml_names)

        tbl = MultiPixelChain(hstack((theta_table,
                                      B_table,
                                      meta_table,
                                      ml_table)))

        return tbl


import argparse
import os
import numpy as np
import tensorflow as tf

from models_clevr_snmn.model import Model
from models_clevr_snmn.config import (
    cfg, merge_cfg_from_file, merge_cfg_from_list)
from util.clevr_train.data_reader import DataReader
from util import boxes
from util.losses import SharpenLossScaler

parser = argparse.ArgumentParser()
parser.add_argument('--cfg', required=True)
parser.add_argument('opts', default=None, nargs=argparse.REMAINDER)
args = parser.parse_args()
merge_cfg_from_file(args.cfg)
assert cfg.EXP_NAME == os.path.basename(args.cfg).replace('.yaml', '')
if args.opts:
    merge_cfg_from_list(args.opts)

# Start session
os.environ["CUDA_VISIBLE_DEVICES"] = str(cfg.GPU_ID)
sess = tf.Session(config=tf.ConfigProto(
    gpu_options=tf.GPUOptions(allow_growth=cfg.GPU_MEM_GROWTH)))

# Data files
imdb_file = cfg.IMDB_FILE % cfg.TRAIN.SPLIT_LOC
data_reader = DataReader(
    imdb_file, shuffle=True, one_pass=False, batch_size=cfg.TRAIN.BATCH_SIZE,
    vocab_question_file=cfg.VOCAB_QUESTION_FILE, T_encoder=cfg.MODEL.T_ENCODER,
    vocab_answer_file=cfg.VOCAB_ANSWER_FILE, load_gt_layout=True,
    vocab_layout_file=cfg.VOCAB_LAYOUT_FILE, T_decoder=cfg.MODEL.T_CTRL,
    img_H=cfg.MODEL.H_IMG, img_W=cfg.MODEL.W_IMG)
num_vocab = data_reader.batch_loader.vocab_dict.num_vocab
num_choices = data_reader.batch_loader.answer_dict.num_vocab
module_names = data_reader.batch_loader.layout_dict.word_list

# Inputs and model
input_seq_batch = tf.placeholder(tf.int32, [None, None])
seq_length_batch = tf.placeholder(tf.int32, [None])
image_feat_batch = tf.placeholder(
    tf.float32, [None, cfg.MODEL.H_FEAT, cfg.MODEL.W_FEAT, cfg.MODEL.FEAT_DIM])
model = Model(
    input_seq_batch, seq_length_batch, image_feat_batch, num_vocab=num_vocab,
    num_choices=num_choices, module_names=module_names, is_training=True)

# Loss function
bbox_ind_batch = tf.placeholder(tf.int32, [None])
bbox_offset_batch = tf.placeholder(tf.float32, [None, 4])
loss_bbox_ind = tf.reduce_mean(
    tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=model.loc_scores, labels=bbox_ind_batch))
loss_bbox_offset = model.bbox_offset_loss(bbox_ind_batch, bbox_offset_batch)
if cfg.TRAIN.USE_GT_LAYOUT:
    gt_layout_batch = tf.placeholder(tf.int32, [None, None])
    loss_layout = tf.reduce_mean(
        tf.nn.sparse_softmax_cross_entropy_with_logits(
            logits=model.module_logits, labels=gt_layout_batch))
else:
    loss_layout = tf.convert_to_tensor(0.)
loss_rec = model.rec_loss
sharpen_scale_ph = tf.placeholder(tf.float32, [])
if cfg.TRAIN.USE_SHARPEN_LOSS:
    loss_sharpen = model.sharpen_loss()
else:
    loss_sharpen = tf.convert_to_tensor(0.)
loss_train = (loss_bbox_ind * cfg.TRAIN.BBOX_IND_LOSS_WEIGHT +
              loss_bbox_offset * cfg.TRAIN.BBOX_OFFSET_LOSS_WEIGHT +
              loss_layout * cfg.TRAIN.LAYOUT_LOSS_WEIGHT +
              loss_rec * cfg.TRAIN.REC_LOSS_WEIGHT +
              loss_sharpen * cfg.TRAIN.SHARPEN_LOSS_WEIGHT * sharpen_scale_ph)
loss_total = loss_train + cfg.TRAIN.WEIGHT_DECAY * model.l2_reg

# Train with Adam
solver = tf.train.AdamOptimizer(learning_rate=cfg.TRAIN.SOLVER.LR)
solver_op = solver.minimize(loss_total)
# Save moving average of parameters
ema = tf.train.ExponentialMovingAverage(decay=cfg.TRAIN.EMV_DECAY)
ema_op = ema.apply(model.params)
with tf.control_dependencies([solver_op]):
    train_op = tf.group(ema_op)

# Save snapshot
snapshot_dir = cfg.TRAIN.SNAPSHOT_DIR % cfg.EXP_NAME
os.makedirs(snapshot_dir, exist_ok=True)
snapshot_saver = tf.train.Saver(max_to_keep=None)  # keep all snapshots
if cfg.TRAIN.START_ITER > 0:
    snapshot_file = os.path.join(snapshot_dir, "%08d" % cfg.TRAIN.START_ITER)
    print('resume training from %s' % snapshot_file)
    snapshot_saver.restore(sess, snapshot_file)
else:
    sess.run(tf.global_variables_initializer())
    if cfg.TRAIN.INIT_FROM_WEIGHTS:
        snapshot_saver.restore(sess, cfg.TRAIN.INIT_WEIGHTS_FILE)
        print('initialized from %s' % cfg.TRAIN.INIT_WEIGHTS_FILE)
# Save config
np.save(os.path.join(snapshot_dir, 'cfg.npy'), np.array(cfg))

# Write summary to TensorBoard
log_dir = cfg.TRAIN.LOG_DIR % cfg.EXP_NAME
os.makedirs(log_dir, exist_ok=True)
log_writer = tf.summary.FileWriter(log_dir, tf.get_default_graph())
loss_bbox_ind_ph = tf.placeholder(tf.float32, [])
loss_bbox_offset_ph = tf.placeholder(tf.float32, [])
loss_layout_ph = tf.placeholder(tf.float32, [])
loss_rec_ph = tf.placeholder(tf.float32, [])
loss_sharpen_ph = tf.placeholder(tf.float32, [])
accuracy_ph = tf.placeholder(tf.float32, [])
summary_trn = []
summary_trn.append(tf.summary.scalar("loss/bbox_ind", loss_bbox_ind_ph))
summary_trn.append(tf.summary.scalar("loss/bbox_offset", loss_bbox_offset_ph))
summary_trn.append(tf.summary.scalar("loss/layout", loss_layout_ph))
summary_trn.append(tf.summary.scalar("loss/rec", loss_rec_ph))
summary_trn.append(tf.summary.scalar("loss/sharpen", loss_sharpen_ph))
summary_trn.append(tf.summary.scalar("loss/sharpen_scale", sharpen_scale_ph))
summary_trn.append(tf.summary.scalar("eval/loc/P1", accuracy_ph))
log_step_trn = tf.summary.merge(summary_trn)

# Run training
avg_accuracy, accuracy_decay = 0., 0.99
iou_th = cfg.TRAIN.BBOX_IOU_THRESH
sharpen_loss_scaler = SharpenLossScaler(cfg)
for n_batch, batch in enumerate(data_reader.batches()):
    n_iter = n_batch + cfg.TRAIN.START_ITER
    if n_iter >= cfg.TRAIN.MAX_ITER:
        break

    sharpen_scale = sharpen_loss_scaler(n_iter)
    feed_dict = {input_seq_batch: batch['input_seq_batch'],
                 seq_length_batch: batch['seq_length_batch'],
                 image_feat_batch: batch['image_feat_batch'],
                 bbox_ind_batch: batch['bbox_ind_batch'],
                 bbox_offset_batch: batch['bbox_offset_batch'],
                 sharpen_scale_ph: sharpen_scale}
    if cfg.TRAIN.USE_GT_LAYOUT:
        feed_dict[gt_layout_batch] = batch['gt_layout_batch']
    loc_scores_val, bbox_offset_val, loss_bbox_ind_val, loss_bbox_offset_val, \
        loss_layout_val, loss_rec_val, loss_sharpen_val, _ = sess.run(
            (model.loc_scores, model.bbox_offset, loss_bbox_ind,
             loss_bbox_offset, loss_layout, loss_rec, loss_sharpen, train_op),
            feed_dict)

    # compute accuracy
    bbox_pred = boxes.batch_feat_grid2bbox(
        np.argmax(loc_scores_val, axis=1), bbox_offset_val,
        data_reader.batch_loader.stride_H, data_reader.batch_loader.stride_W,
        data_reader.batch_loader.feat_H, data_reader.batch_loader.feat_W)
    bbox_gt = batch['bbox_batch']
    accuracy = np.mean(boxes.batch_bbox_iou(bbox_pred, bbox_gt) >= iou_th)
    avg_accuracy += (1-accuracy_decay) * (accuracy-avg_accuracy)

    # Add to TensorBoard summary
    if (n_iter+1) % cfg.TRAIN.LOG_INTERVAL == 0:
        print("exp: %s, iter = %d\n\t" % (cfg.EXP_NAME, n_iter+1) +
              "loss (bbox_ind) = %f, loss (bbox_offset) = %f, "
              "loss (layout) = %f, loss (rec) = %f, "
              "loss (sharpen) = %f, sharpen_scale = %f\n\t" % (
                loss_bbox_ind_val, loss_bbox_offset_val, loss_layout_val,
                loss_rec_val, loss_sharpen_val, sharpen_scale) +
              "P1@%.2f (cur) = %f, P1@%.2f (avg) = %f" % (
                iou_th, accuracy, iou_th, avg_accuracy))
        summary = sess.run(log_step_trn, {
            loss_bbox_ind_ph: loss_bbox_ind_val,
            loss_bbox_offset_ph: loss_bbox_offset_val,
            loss_layout_ph: loss_layout_val,
            loss_rec_ph: loss_rec_val,
            loss_sharpen_ph: loss_sharpen_val, sharpen_scale_ph: sharpen_scale,
            accuracy_ph: avg_accuracy})
        log_writer.add_summary(summary, n_iter+1)

    # Save snapshot
    if ((n_iter+1) % cfg.TRAIN.SNAPSHOT_INTERVAL == 0 or
            (n_iter+1) == cfg.TRAIN.MAX_ITER):
        snapshot_file = os.path.join(snapshot_dir, "%08d" % (n_iter+1))
        snapshot_saver.save(sess, snapshot_file, write_meta_graph=False)
        print('snapshot saved to ' + snapshot_file)


#!/usr/bin/env python3

import sys
import argparse
import numpy as np
import cv2
import time
from PIL import Image

from edgetpu.basic.basic_engine import BasicEngine
from edgetpu.basic import edgetpu_utils

fps = ""
framecount = 0
time1 = 0

LABEL_CONTOURS = [(0, 0, 0),  # 0=None
                  # 1=Buildings , 2=Fences, 3=Other, 4=Pedestrians, 5=Poles
                  (70, 70, 70), (190, 153, 153), (250, 170, 160), (220, 20, 60), (153, 153, 153),
                  # 6=RoadLines, 7=Roads, 8=Sidewalks, 9=Vegetation, 10=Vehicles
                  (157, 234, 50), (128, 64, 128), (244, 35, 232), (107, 142, 35), (0, 0, 142),
                  # 11=Walls, 12=TrafficSigns
                  (102, 102, 156), (220, 220, 0)]

def decode_prediction_mask(mask):
    mask_shape = mask.shape
    mask_color = np.zeros(shape=[mask_shape[0], mask_shape[1], 3], dtype=np.uint8)
    unique_label_ids = [v for v in np.unique(mask) if v != 0 and v != 255]
    for label_id in unique_label_ids:
        idx = np.where(mask == label_id)
        mask_color[idx] = LABEL_CONTOURS[label_id]
    return mask_color

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="bisenetv2_carla_256x256_full_integer_quant_edgetpu.tflite", help="Path of the BiSeNetV2 model.")
    parser.add_argument("--video_file", default="cityscapes_demo.mp4", help="Input video to run segmentation on.")
    parser.add_argument("--output_name", default="output.avi", help="Name of the output file.")
    parser.add_argument('--vidfps', type=int, default=30, help='FPS of Video. (Default=30)')
    args = parser.parse_args()

    model    = args.model
    video_file    = args.video_file
    video_file    = args.video_file
    vidfps        = args.vidfps

    devices = edgetpu_utils.ListEdgeTpuPaths(edgetpu_utils.EDGE_TPU_STATE_UNASSIGNED)
    engine = BasicEngine(model_path=model, device_path=devices[0])
    model_height = engine.get_input_tensor_shape()[1]
    model_width  = engine.get_input_tensor_shape()[2]

    cap = cv2.VideoCapture(video_file)
    cap.set(cv2.CAP_PROP_FPS, vidfps)
    video_width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    video_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    out = cv2.VideoWriter("output.avi",
        cv2.VideoWriter_fourcc(*"MJPG"), 20, (video_width, video_height))

    waittime = 1
    window_name = "Segmentation"

    cv2.namedWindow(window_name, cv2.WINDOW_AUTOSIZE)

    while(cap.isOpened()):
        t1 = time.perf_counter()

        ret, color_image = cap.read()
        if not ret:
            break

        # Normalization
        prepimg_deep = cv2.resize(color_image, (model_width, model_height))
        prepimg_deep = cv2.cvtColor(prepimg_deep, cv2.COLOR_BGR2RGB)
        prepimg_deep = np.expand_dims(prepimg_deep, axis=0)

        # Run model
        prepimg_deep = prepimg_deep.astype(np.uint8)
        inference_time, predictions = engine.run_inference(prepimg_deep.flatten())

        # Segmentation
        predictions = predictions.reshape(model_height, model_width, 13)
        predictions = np.argmax(predictions, axis=-1)
        imdraw = decode_prediction_mask(predictions)
        imdraw = cv2.cvtColor(imdraw, cv2.COLOR_RGB2BGR)
        imdraw = cv2.resize(imdraw, (video_width, video_height))
        imdraw = cv2.addWeighted(color_image, 1.0, imdraw, 0.9, 0)

        cv2.putText(imdraw, fps, (video_width-170,15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (38,0,255), 1, cv2.LINE_AA)
        cv2.imshow(window_name, imdraw)
        out.write(imdraw)

        if cv2.waitKey(waittime)&0xFF == ord('q'):
            break

        # FPS calculation
        framecount += 1
        if framecount >= 10:
            fps       = "(Playback) {:.1f} FPS".format(time1/10)
            framecount = 0
            time1 = 0
        t2 = time.perf_counter()
        elapsedTime = t2-t1
        time1 += 1/elapsedTime
    out.release()
    cap.release()
    cv2.destroyAllWindows()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Apr  4 06:03:28 2019

@author: tanujsinghal
"""
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
import re
import copy
from sklearn.metrics import jaccard_similarity_score,confusion_matrix
from sklearn.externals import joblib
import os
import pathlib
import pickle

if __name__=='__main__':
        hin=pd.read_csv("Hinglish_Profanity_List.csv")
        train = pd.read_csv("train.csv")
        test = pd.read_csv("test.csv")
        columns = ['obscene','insult','toxic','severe_toxic','identity_hate','threat']
        hin_bad_words = hin.iloc[:,0].values.tolist()
        bad_words_to_english = hin.iloc[:,1].values.tolist()
        hin = hin.iloc[:,:-1].values.tolist()
        train, test = train_test_split(train, test_size=0.2)
        labels = train.iloc[:,2:]
        train_data = train.iloc[:,1]
        test_data = test.iloc[:,1]

        features = 5000
        ngram = (1,2)
        vectorizer = TfidfVectorizer(stop_words='english',\
                                token_pattern = "\w*[a-z]\w*",\
                                ngram_range=ngram,\
                                max_features=features)

        train_features = vectorizer.fit_transform(train_data)
        filename='vect'
        pickle.dump(vectorizer, open(filename, 'wb'))
        test_features = vectorizer.transform(test_data)
        logreg = LogisticRegression(C=10,solver="liblinear")
        models={}
        logistic_results = pd.DataFrame(columns=columns)
        cnt=0

        for i in columns:
                y = train[i]
                models[i]=copy.copy(logreg.fit(train_features, y))
                filename = "model_"+ str(cnt)
                pickle.dump(models[i], open(filename, 'wb'))
                ypred_X = logreg.predict(train_features)
                testy_prob = logreg.predict_proba(test_features)[:,1]
                logistic_results[i] = testy_prob
                cnt+=1
def abusive_hinglish_to_english(data):
    hin=pd.read_csv("Hinglish_Profanity_List.csv")
    hin_bad_words = hin.iloc[:,0].values.tolist()
    bad_words_to_english = hin.iloc[:,1].values.tolist()
    hin = hin.iloc[:,:-1].values.tolist()
    cnt=0
    for sentence in data:
        wordList = sentence.split()
        for word in hin_bad_words:
            if word in wordList:
                x=wordList.index(word)
                wordList[x]=bad_words_to_english[hin_bad_words.index(word)]
        sentence = ' '.join(wordList)
        data[cnt]=sentence
        cnt+=1
    return data

def check_dir(directory):
    if not os.path.exists(directory):
        os.makedirs(directory)

def myinput(vectorizer,model,val):
    sent="Thank you for understanding. I think very highly of you and would not revert without discussion."
    sent2="Yo bitch Ja Rule is more succesful then you'll ever be whats up with you and hating you sad mofuckas...i should bitch slap ur pethedic white faces and get you to kiss my ass you guys sicken me. Ja rule is about pride in da music man. dont diss that shit on him. and nothin is wrong bein like tupac he was a brother too...fuckin white boys get things right next time.,"
    sen3="Explanation Why the edits made under my username Hardcore Metallica Fan were reverted? They weren't vandalisms, just closure on some GAs after I voted at New York Dolls FAC. And please don't remove the template from the talk page since I'm retired now.89.205.38.27"
    sen4="COCKSUCKER BEFORE YOU PISS AROUND ON MY WORK"
    sen5="While booking during rush hour, it is always advisable to check this box ON. It will deduct your FULL ticket amount first. After that it check whether confirmed ticket available or not. If confirmed ticket not available it will show ticket not available. Within 2 days will get your full amount Refund."
    sen6="madarchod chutia"
    sent7='F**K YOU!good day'
    l=[sent,sent2,sen3,sen4,sen5,sen6,sent7]
    l=abusive_hinglish_to_english(l)
    df = pd.DataFrame(l)
    user_data = vectorizer.transform(l)
    results2 = pd.DataFrame(columns=columns)
    for i in columns:
        user_results = models[i].predict_proba(user_data)[:,1]
        results2[i] = user_results
    y=results2.iloc[val].values
    x = columns
    plt.ylim(0, 100)
    plt.tight_layout()
    plt.bar(x, height= y)
    plt.show()
    plt.savefig('foo.png')
    return df,results2

def myinput_network(text):
    columns = ['obscene','insult','toxic','severe_toxic','identity_hate','threat']
    sent="Thank you for understanding. I think very highly of you and would not revert without discussion."
    sent2="Yo bitch Ja Rule is more succesful then you'll ever be whats up with you and hating you sad mofuckas...i should bitch slap ur pethedic white faces and get you to kiss my ass you guys sicken me. Ja rule is about pride in da music man. dont diss that shit on him. and nothin is wrong bein like tupac he was a brother too...fuckin white boys get things right next time.,"
    sen3="Explanation Why the edits made under my username Hardcore Metallica Fan were reverted? They weren't vandalisms, just closure on some GAs after I voted at New York Dolls FAC. And please don't remove the template from the talk page since I'm retired now.89.205.38.27"
    sen4="COCKSUCKER BEFORE YOU PISS AROUND ON MY WORK"
    sen5="While booking during rush hour, it is always advisable to check this box ON. It will deduct your FULL ticket amount first. After that it check whether confirmed ticket available or not. If confirmed ticket not available it will show ticket not available. Within 2 days will get your full amount Refund."
    sen6="madarchod chutia"
    sent7='F**K YOU!good day'
    l=[sent,sent2,sen3,sen4,sen5,sen6,sent7,text]
    l=[text,sent2]
    if len(text)>1 and type(text) is list:
            l=text
            res,x=myinput_network2(l)
            return res,x
    l=abusive_hinglish_to_english(l)
    df = pd.DataFrame(l)
    f='vect'
    vect= pickle.load(open(f, 'rb'))
    user_data = vect.transform(l)
    results2 = pd.DataFrame(columns=columns)
    cnt=0
    mymodels={}
    for i in range(6):
        filename='model_'+str(i)
        mymodels[columns[i]]= pickle.load(open(filename, 'rb'))
    for i in range(6):
        user_results = mymodels[columns[i]].predict_proba(user_data)[:,1]
        results2[columns[i]] = user_results
    x = columns
    return results2.iloc[0].values,x
def myinput_network2(text):
        columns = ['obscene','insult','toxic','severe_toxic','identity_hate','threat']
        x = columns
        if len(text)>1:
                l=text
                print(l)
                print(type(l))

        l=abusive_hinglish_to_english(l)
        df = pd.DataFrame(l)
        f='vect'
        vect= pickle.load(open(f, 'rb'))
        user_data = vect.transform(l)
        results2 = pd.DataFrame(columns=columns)
        cnt=0
        mymodels={}
        filename='model_'+str(2)
        mymodels[columns[2]]= pickle.load(open(filename, 'rb'))
        user_results = mymodels[columns[2]].predict_proba(user_data)[:,1]
        results2[columns[2]] = user_results
        return results2[columns[2]].values,x


import logging, sys, os
import numpy as np

from PIL import Image
import io

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision
import torchvision.models
import torchvision.transforms as transforms
from model_def import Net

import subprocess

subprocess.call(['pip', 'install', 'sagemaker_inference'])
from sagemaker_inference import content_types, decoder

def _get_logger():
    '''
    로깅을 위해 파이썬 로거를 사용
    # https://stackoverflow.com/questions/17745914/python-logging-module-is-printing-lines-multiple-times
    '''
    loglevel = logging.DEBUG
    l = logging.getLogger(__name__)
    if not l.hasHandlers():
        l.setLevel(loglevel)
        logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
        l.handler_set = True
    return l

logger = _get_logger()

# 파이토치 서브의 디폴트 model_fn, input_fn 코드
# https://github.com/aws/sagemaker-pytorch-inference-toolkit/blob/master/src/sagemaker_pytorch_serving_container/default_pytorch_inference_handler.py

import logging
def model_fn(model_dir):
    logger.info("--> model_dir : {}".format(model_dir))
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print("device: ", device)

    model = Net()

    logger.info("--> model network is loaded")

#     if torch.cuda.device_count() > 1:
#         logger.info("Gpu count: {}".format(torch.cuda.device_count()))
#         model = nn.DataParallel(model)

    # 디폴트로 DataParallel 를 기술함.
    # 이유는 weight 이름 앞에 module 을 붙이기 위해서 임
    model = nn.DataParallel(model)

    model_file_path = os.path.join(model_dir, "model.pth")
    print("model_file_path: ", model_file_path)

    try:
        with open(model_file_path, "rb") as f:
            model.load_state_dict(torch.load(f))
        logger.info("---> ####### Model is loaded #########")
    except BaseException:
        logging.exception("An exception was thrown!")
        logger.info("---> ########## Failure loading a Model #######")
        # 디버깅에만 sleep을 사용하세요.
        # import time
        # time.sleep(600) # 디버깅을 위해서 10분 정지

    return model.to(device)

def input_fn(input_data, content_type):
    '''
    content_type == 'application/x-npy' 일시에 토치 텐서의 변환 작업 수행
    '''
    logger.info("#### input_fn starting ######")
    logger.info(f"content_type: {content_type}")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if isinstance(input_data, (np.ndarray)):
        np_array = input_data # 로컬에서 테스트시에 numpy.ndarray 로 제공 되는 경우
    else: # 토치 서브를 통해서 numpy.ndarray가 제공되면, IOByte( 확인 필요) 로 제공되기에 디코딩이 필요함.
        np_array = decoder.decode(input_data, content_type)

    logger.info(f"np_array shape: {np_array.shape}  ")

    tensor = torch.FloatTensor(
            np_array) if content_type in content_types.UTF8_TYPES else torch.from_numpy(np_array)

    return tensor.to(device)

def predict_fn(data, model):
    logger.info("#### predict_fn starting ######")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    input_data = data.to(device)
    model.eval()
    with torch.no_grad():
        output = model(input_data)
    return output


import matplotlib
matplotlib.use('TkAgg')
from keras import backend as K
from keras.models import load_model
from keras.preprocessing import image
from keras.optimizers import Adam
from imageio import imread
import numpy as np
from matplotlib import pyplot as plt

from models.keras_ssd300 import ssd_300
from keras_loss_function.keras_ssd_loss import SSDLoss
from keras_layers.keras_layer_AnchorBoxes import AnchorBoxes
from keras_layers.keras_layer_DecodeDetections import DecodeDetections
from keras_layers.keras_layer_DecodeDetectionsFast import DecodeDetectionsFast
from keras_layers.keras_layer_L2Normalization import L2Normalization

from ssd_encoder_decoder.ssd_output_decoder import decode_detections, decode_detections_fast

from data_generator.object_detection_2d_data_generator import DataGenerator
from data_generator.object_detection_2d_photometric_ops import ConvertTo3Channels
from data_generator.object_detection_2d_geometric_ops import Resize
from data_generator.object_detection_2d_misc_utils import apply_inverse_transforms

# Set the image size.
img_height = 300
img_width = 300

# 1: Build the Keras model

K.clear_session()  # Clear previous models from memory.

model = ssd_300(image_size=(img_height, img_width, 3),
                n_classes=20,
                mode='inference',
                l2_regularization=0.0005,
                scales=[0.1, 0.2, 0.37, 0.54, 0.71, 0.88, 1.05],
                # The scales for MS COCO are [0.07, 0.15, 0.33, 0.51, 0.69, 0.87, 1.05]
                aspect_ratios_per_layer=[[1.0, 2.0, 0.5],
                                         [1.0, 2.0, 0.5, 3.0, 1.0 / 3.0],
                                         [1.0, 2.0, 0.5, 3.0, 1.0 / 3.0],
                                         [1.0, 2.0, 0.5, 3.0, 1.0 / 3.0],
                                         [1.0, 2.0, 0.5],
                                         [1.0, 2.0, 0.5]],
                two_boxes_for_ar1=True,
                steps=[8, 16, 32, 64, 100, 300],
                offsets=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
                clip_boxes=False,
                variances=[0.1, 0.1, 0.2, 0.2],
                normalize_coords=True,
                subtract_mean=[123, 117, 104],
                swap_channels=[2, 1, 0],
                confidence_thresh=0.5,
                iou_threshold=0.45,
                top_k=200,
                nms_max_output_size=400)

# 2: Load the trained weights into the model.

# TODO: Set the path of the trained weights.
weights_path = 'pre_trained/VGG_VOC0712Plus_SSD_300x300_ft_iter_160000.h5'

model.load_weights(weights_path, by_name=True)

# 3: Compile the model so that Keras won't complain the next time you load it.

adam = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)

ssd_loss = SSDLoss(neg_pos_ratio=3, alpha=1.0)

model.compile(optimizer=adam, loss=ssd_loss.compute_loss)

orig_images = []  # Store the images here.
input_images = []  # Store resized versions of the images here.

# We'll only load one image in this example.
img_path = 'examples/fish_bike.jpg'

orig_images.append(imread(img_path))
img = image.load_img(img_path, target_size=(img_height, img_width))
img = image.img_to_array(img)
input_images.append(img)
input_images = np.array(input_images)

y_pred = model.predict(input_images)

confidence_threshold = 0.5

y_pred_thresh = [y_pred[k][y_pred[k, :, 1] > confidence_threshold] for k in range(y_pred.shape[0])]

np.set_printoptions(precision=2, suppress=True, linewidth=90)
print("Predicted boxes:\n")
print('   class   conf xmin   ymin   xmax   ymax')
print(y_pred_thresh[0])

# Display the image and draw the predicted boxes onto it.

# Set the colors for the bounding boxes
colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
classes = ['background',
           'aeroplane', 'bicycle', 'bird', 'boat',
           'bottle', 'bus', 'car', 'cat',
           'chair', 'cow', 'diningtable', 'dog',
           'horse', 'motorbike', 'person', 'pottedplant',
           'sheep', 'sofa', 'train', 'tvmonitor']

plt.figure(figsize=(20, 12))
plt.imshow(orig_images[0])

current_axis = plt.gca()

for box in y_pred_thresh[0]:
    # Transform the predicted bounding boxes for the 300x300 image to the original image dimensions.
    xmin = box[2] * orig_images[0].shape[1] / img_width
    ymin = box[3] * orig_images[0].shape[0] / img_height
    xmax = box[4] * orig_images[0].shape[1] / img_width
    ymax = box[5] * orig_images[0].shape[0] / img_height
    color = colors[int(box[0])]
    label = '{}: {:.2f}'.format(classes[int(box[0])], box[1])
    current_axis.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, color=color, fill=False, linewidth=2))
    current_axis.text(xmin, ymin, label, size='x-large', color='white', bbox={'facecolor': color, 'alpha': 1.0})

plt.show()


# This code is originated from https://github.com/andywu0913/OpenAI-GYM-CarRacing-DQN
# Also uses https://github.com/pythonlessons/Reinforcement_Learning/blob/master/03_CartPole-reinforcement-learning_Dueling_DDQN/Cartpole_Double_DDQN.py
import gym
import os, sys
import random
import numpy as np
from collections import deque
import cv2
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, Lambda, Flatten, Dense, Add
from tensorflow.keras.optimizers import Adam
from tensorflow import keras
import tensorflow.keras.backend as K
from model_tester import TesterAgent
from gym_multi_car_racing import MultiCarRacing

# ACTIONS  = [
#             (-1, 0.6, 0.2), (0, 1, 0.2), (1, 0.6, 0.2), #           Action Space Structure
#             (-1, 0.6,   0), (0, 1,   0), (1, 0.6,   0), #        (Steering Wheel, Gas, Break)
#             (-1, 0, 0.2), (0, 0, 0.2), (1, 0, 0.2), # Range        -1~1       0~1   0~1
#             (-1, 0,   0), (0, 0,   0), (1, 0,   0)
#            ]

ACTIONS  = [
            (-1, 0.2, 0.2), (0, 0.2, 0.2), (1, 0.2, 0.2), #           Action Space Structure
            (-1, 0.2,   0), (0, 0.6,   0), (1, 0.2,   0), #        (Steering Wheel, Gas, Break)
            (-1, 0, 0.2), (0, 0, 0.2), (1, 0, 0.2), # Range        -1~1       0~1   0~1
            (-1, 0,   0), (0, 0,   0), (1, 0,   0)
           ]

def process_state_image(state, car_id):
    state = state[car_id, ...]
    #state = np.dot(state, [0.299, 0.587, 0.144])
    state = cv2.cvtColor(state, cv2.COLOR_BGR2GRAY)
    state = state.astype(float)
    state /= 255.0
    return state

def add_frame(s_new, s_stack):
    s_stack[:,:,:-1] = s_stack[:,:,1:]
    s_stack[:,:,-1] = s_new
    return s_stack

class DDQNTesterAgent(TesterAgent):
    def __init__(self,
                 model_path='..\\algorithms\dueling\Image\save\\26_09_15_27\episode_500.h5',
                 car_id=0,
                 num_frames=4,
                 learning_rate = 0.001,
                 actions=ACTIONS,
                 **kwargs
                 ):
        super().__init__(**kwargs)
        self.learning_rate   = learning_rate
        self.car_id = car_id
        self.frame_buffer = None
        self.num_frames = num_frames
        self.actions = actions
        self.model = self.build_model()
        self.agent = self._load_model(model_path)

    def build_model(self):
        # Neural Net for Deep-Q learning Model
        input = Input(shape=(96, 96, self.num_frames), name='input')
        x = Conv2D(filters=6, kernel_size=(7,7), strides=3, activation='relu')(input)
        x = MaxPooling2D(pool_size=(2,2))(x)
        x = Conv2D(filters=12, kernel_size=(4,4), activation='relu')(x)
        x = MaxPooling2D(pool_size=(2,2))(x)
        x = Flatten()(x)
        x = Dense(256, activation='relu')(x)
        x = Dense(len(self.actions) + 1, activation='linear')(x)
        output = Lambda(lambda i: K.expand_dims(i[:,0],-1) + i[:,1:] - K.mean(i[:,1:], keepdims=True), output_shape=(len(self.actions),))(x)

        model = Model(inputs = input, outputs = output, name='Dueling_DDQN_model')
        model.compile(loss="mean_squared_error", optimizer=Adam(learning_rate=self.learning_rate), metrics=["accuracy"])
        model.summary()
        return model

    def _load_model(self, model_path):
        self.model.load_weights(model_path)

    def _update_frame_buffer(self, new_frame):
        if self.frame_buffer is None:
            self.frame_buffer = np.repeat(new_frame[:,:,np.newaxis], self.num_frames, axis=2)
        else:
            self.frame_buffer = add_frame(new_frame, self.frame_buffer)

    def state_to_action(self, state):
        """
        This function should take the most recent state and return the
        action vector in the exact same form it is needed for the environment.
        If you are using frame buffer see example in _update_frame_buffer
        how to take care of that.
        """
        state = process_state_image(state, self.car_id)
        self._update_frame_buffer(state)
        act_values = self.model.predict(np.expand_dims(self.frame_buffer, axis=0))
        action_index = np.argmax(act_values[0])

        return self.actions[action_index]

    @staticmethod
    def setup_action_space(env):
        """
        This should be the same action space setup function that you used for training.
        Make sure that the actions set here are the same as the ones used to train the model.
        """
        env.cont_action_space = ACTIONS
        env.action_space = gym.spaces.Discrete(len(env.cont_action_space))

    @staticmethod
    def get_observation_type():
        """
        Simply return 'frames' or 'features'
        """
        return 'frames'


#!/usr/bin/python3.4
# News Recommender version 2
# written by Max Reinisch

# Given a URL for a news article, recommend TV news
# clips similar to article contents

####################
##### IMPORTS ######
####################

import urllib3
from urllib.parse import urlencode
import json
from sklearn.feature_extraction.text import TfidfVectorizer
from scipy.spatial.distance import cosine
import spacy
from collections import Counter
import en_core_web_sm
import logging

####################
#### VARIABLES #####
####################

CALAIS_ENDPOINT = 'https://api.thomsonreuters.com/permid/calais'
GDELT_HEADER = {'Content-Type' : 'application/json', 'outputformat' : 'application/json'}
log = logging.getLogger(__name__)
urllib3.disable_warnings()
http = urllib3.PoolManager()
nlp = en_core_web_sm.load()

def makeRecommendations(article, calais_token = False):
    query = getSearchQuery(article, calais_token)
    GDELT_response = getGDELTv2Response(query) #GDELT_response in JSON format
    if(GDELT_response.get('clips')):
        return list(sortClipsBySimilarity(GDELT_response.get("clips"), article))
    else:
        return []

def getSearchQuery(article, calais_token = False):
    if(calais_token):
        calais_json = getOpenCalaisResponse(article, calais_token)
        if calais_json != "SpaCy":
            return parseCalais(calais_json)
    entities = extractEntities(article)
    return parseEntities(entities)

def extractEntities(article):
    doc = nlp(article["title"] + article["title"] + article["body"])
    return [x.text for x in doc.ents]

def parseEntities(entities):
    if len(entities) == 0:
        return []

    return [e[0] for e in Counter(entities).most_common(3)]

def getOpenCalaisResponse(article, calais_token):
    CALAIS_HEADER = {'X-AG-Access-Token' : calais_token, 'Content-Type' : 'text/raw', 'outputformat' : 'application/json'}
    response = http.request('POST', CALAIS_ENDPOINT, body= (article.get("title")+" "+article.get("body")).encode('utf-8'), headers=CALAIS_HEADER, timeout=80)
    if response.status >= 400:
        log.warning("OpenCalais returned status code: " + str(response.status) + ".  Falling back to SpaCy extraction protocol...")
        return "SpaCy"
    content = response.data.decode('utf-8')

    c = json.loads(content)
    return c

def parseCalais(c):
    # This function is an unsolved problem.  How to find the best search query
    # Currently returns a list of all entities' names in order of number of mentions
    entities = [values for values in c.values() if values.get("_typeGroup") == "entities"]
    if len(entities) ==0:
        return []
    return [e.get('lastname') or e['name'] for e in sorted(entities, key=lambda x: len(x['instances']), reverse=True)]

def getGDELTv2Response(query):
    good_result = False
    entities = query[:3]
    while(not good_result):
        if len(entities) ==0:
            log.warning("No search found.  Returning empty result.")
            return {}
        query = "+".join(['"' + entity + '"' for entity in entities])
        # url = 'https://api.gdeltproject.org/api/v2/tv/tv?query='+query+ 'market:%22National%22&mode=clipgallery&format=json&datanorm=perc&timelinesmooth=0&datacomb=sep&last24=yes&timezoom=yes&TIMESPAN=14days'
        encoded_args = urlencode({'query': query+' market:"National"', 'mode':'clipgallery', 'format':'json', 'datanorm':'perc',"timelinesmooth":0, "datacomb":"sep", "last24":"yes", "timezoom":"yes", "TIMESPAN":"14days"})
        url = 'https://api.gdeltproject.org/api/v2/tv/tv?' + encoded_args
        res = http.urlopen("GET", url)

        if(res.status>=400):
            log.error("GDELT returned status code: " + str(res.status)+ ".  Exiting...")
            return {}
        try:
            gdelt_json = json.loads(res.data.decode('utf-8'))
        except ValueError:
            log.warning("Bad GDELT response, Returning empty result.");
            log.warning(url)
            gdelt_json = {}
        log.info(url)
        if len(gdelt_json.keys()) ==0:
            entities = entities[0:-1]
            log.info("GDLET returned 0 results. Simplifying search to: " + str(entities))
        else:
            good_result = True
    return gdelt_json

def sortClipsBySimilarity(clips, article):
    vectorizer = TfidfVectorizer(stop_words="english", ngram_range=(1,2) )
    bow = vectorizer.fit_transform([clip.get('snippet') for clip in clips])
    article_bow = vectorizer.transform([article.get('body')])
    cosine_distances = [cosine(vec.todense(), article_bow.todense()) for vec in bow]
    ret = []
    for clip, dist in zip(clips, cosine_distances):
        clip.update({"similarity": dist})
        ret.append(clip)
    log.info(ret)
    return sorted(ret, key = lambda x: x.get('similarity'))