import os
import sys
import time
import random
import string
import argparse

import torch
import torch.backends.cudnn as cudnn
import torch.nn.init as init
import torch.optim as optim
import torch.utils.data
import numpy as np
from tqdm import tqdm
from utils import CTCLabelConverter, AttnLabelConverter, TransLabelConverter, Averager
from dataset import hierarchical_dataset, AlignCollate, Batch_Balanced_Dataset
from model import Model
from test import validation
from torch.utils.tensorboard import SummaryWriter
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def train(opt):
    os.makedirs(opt.log, exist_ok=True)
    writer = SummaryWriter(opt.log)
    """ dataset preparation """
    if not opt.data_filtering_off:
        print('Filtering the images containing characters which are not in opt.character')
        print('Filtering the images whose label is longer than opt.batch_max_length')

    opt.select_data = opt.select_data.split('-')
    opt.batch_ratio = opt.batch_ratio.split('-')
    train_dataset = Batch_Balanced_Dataset(opt)

    log = open(f'./saved_models/{opt.exp_name}/log_dataset.txt', 'a')
    AlignCollate_valid = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD)
    valid_dataset, valid_dataset_log = hierarchical_dataset(root=opt.valid_data, opt=opt)
    valid_loader = torch.utils.data.DataLoader(
        valid_dataset, batch_size=opt.batch_size,
        shuffle=True,  # 'True' to check training progress with validation function.
        num_workers=int(opt.workers),
        collate_fn=AlignCollate_valid, pin_memory=True)

    log.write(valid_dataset_log)
    print('-' * 80)
    log.write('-' * 80 + '\n')
    log.close()

    """ model configuration """

    ctc_converter = CTCLabelConverter(opt.character)
    attn_converter = AttnLabelConverter(opt.character)
    opt.num_class = len(attn_converter.character)

    if opt.rgb:
        opt.input_channel = 3
    model = Model(opt)

    print('model input parameters', opt.imgH, opt.imgW, opt.num_fiducial, opt.input_channel, opt.output_channel,
          opt.hidden_size, opt.num_class, opt.batch_max_length)

    # weight initialization
    for name, param in model.named_parameters():
        if 'localization_fc2' in name:
            print(f'Skip {name} as it is already initialized')
            continue
        try:
            if 'bias' in name:
                init.constant_(param, 0.0)
            elif 'weight' in name:
                init.kaiming_normal_(param)
        except Exception as e:  # for batchnorm.
            if 'weight' in name:
                param.data.fill_(1)
            continue

    # data parallel for multi-GPU
    model = torch.nn.DataParallel(model).to(device)
    model.train()
    if opt.saved_model != '':
        print(f'loading pretrained model from {opt.saved_model}')
        if opt.FT:
            model.load_state_dict(torch.load(opt.saved_model), strict=False)
        else:
            model.load_state_dict(torch.load(opt.saved_model))

    """ setup loss """
    loss_avg = Averager()
    ctc_loss = torch.nn.CTCLoss(zero_infinity=True).to(device)
    attn_loss = torch.nn.CrossEntropyLoss(ignore_index=0).to(device)

    # filter that only require gradient decent
    filtered_parameters = []
    params_num = []
    for p in filter(lambda p: p.requires_grad, model.parameters()):
        filtered_parameters.append(p)
        params_num.append(np.prod(p.size()))
    print('Trainable params num : ', sum(params_num))

    # setup optimizer
    if opt.adam:
        optimizer = optim.Adam(filtered_parameters, lr=opt.lr, betas=(opt.beta1, 0.999))
    else:
        optimizer = optim.Adadelta(filtered_parameters, lr=opt.lr, rho=opt.rho, eps=opt.eps)
    print("Optimizer:")

    """ final options """
    # print(opt)
    with open(f'./saved_models/{opt.exp_name}/opt.txt', 'a') as opt_file:
        opt_log = '------------ Options -------------\n'
        args = vars(opt)
        for k, v in args.items():
            opt_log += f'{str(k)}: {str(v)}\n'
        opt_log += '---------------------------------------\n'
        print(opt_log)
        opt_file.write(opt_log)

    """ start training """
    start_iter = 0
    if opt.saved_model != '':
        try:
            start_iter = int(opt.saved_model.split('_')[-1].split('.')[0])
            print(f'continue to train, start_iter: {start_iter}')
        except:
            pass

    start_time = time.time()
    best_accuracy = -1
    best_norm_ED = -1
    iteration = start_iter
    pbar = tqdm(range(opt.num_iter))

    for iteration in pbar:

        # train part
        image_tensors, labels = train_dataset.get_batch()
        image = image_tensors.to(device)
        ctc_text, ctc_length = ctc_converter.encode(labels, batch_max_length=opt.batch_max_length)
        attn_text, attn_length = attn_converter.encode(labels, batch_max_length=opt.batch_max_length)

        batch_size = image.size(0)

        preds, refiner = model(image, attn_text[:, :-1])

        refiner_size = torch.IntTensor([refiner.size(1)] * batch_size)
        refiner = refiner.log_softmax(2).permute(1, 0, 2)
        refiner_loss = ctc_loss(refiner, ctc_text, refiner_size, ctc_length)

        total_loss = opt.lambda_ctc*refiner_loss
        target = attn_text[:, 1:]  # without [GO] Symbol
        for pred in preds:
            total_loss += opt.lambda_attn* attn_loss(pred.view(-1, pred.shape[-1]), target.contiguous().view(-1))

        model.zero_grad()
        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), opt.grad_clip)  # gradient clipping with 5 (Default)
        optimizer.step()
        loss_avg.add(total_loss)
        if loss_avg.val() <=0.6:
            opt.grad_clip = 2
        if loss_avg.val() <=0.3:
            opt.grad_clip = 1

        preds = (p.cpu() for p in preds)
        refiner = refiner.cpu()
        image = image.cpu()
        torch.cuda.empty_cache()

        writer.add_scalar('train_loss', loss_avg.val(), iteration)
        pbar.set_description('Iteration {0}/{1}, AvgLoss {2}'.format(iteration, opt.num_iter, loss_avg.val()))

        # validation part
        if (iteration + 1) % opt.valInterval == 0 or iteration==0:
            elapsed_time = time.time() - start_time
            # for log
            with open(f'./saved_models/{opt.exp_name}/log_train.txt', 'a') as log:
                model.eval()
                with torch.no_grad():
                    valid_loss, current_accuracy, current_norm_ED, preds, confidence_score, labels, infer_time, length_of_data = validation(
                        model, attn_loss, valid_loader, attn_converter, opt)
                model.train()

                # training loss and validation loss
                loss_log = f'[{iteration+1}/{opt.num_iter}] Train loss: {loss_avg.val():0.5f}, Valid loss: {valid_loss:0.5f}, Elapsed_time: {elapsed_time:0.5f}'
                writer.add_scalar('Val_loss', valid_loss)
                pbar.set_description(loss_log)
                loss_avg.reset()

                current_model_log = f'{"Current_accuracy":17s}: {current_accuracy:0.3f}, {"Current_norm_ED":17s}: {current_norm_ED:0.2f}'

                # keep best accuracy model (on valid dataset)
                if current_accuracy > best_accuracy:
                    best_accuracy = current_accuracy
                    torch.save(model.state_dict(), f'./saved_models/{opt.exp_name}/best_accuracy.pth')
                if current_norm_ED > best_norm_ED:
                    best_norm_ED = current_norm_ED
                    torch.save(model.state_dict(), f'./saved_models/{opt.exp_name}/best_norm_ED.pth')
                best_model_log = f'{"Best_accuracy":17s}: {best_accuracy:0.3f}, {"Best_norm_ED":17s}: {best_norm_ED:0.2f}'

                loss_model_log = f'{loss_log}\n{current_model_log}\n{best_model_log}'
                # print(loss_model_log)

                log.write(loss_model_log + '\n')

                # show some predicted results
                dashed_line = '-' * 80
                head = f'{"Ground Truth":25s} | {"Prediction":25s} | Confidence Score & T/F'
                predicted_result_log = f'{dashed_line}\n{head}\n{dashed_line}\n'
                for gt, pred, confidence in zip(labels[:5], preds[:5], confidence_score[:5]):
                    if 'Attn' or 'Transformer' in opt.Prediction:
                        gt = gt[:gt.find('[s]')]
                        pred = pred[:pred.find('[s]')]

                    predicted_result_log += f'{gt:25s} | {pred:25s} | {confidence:0.4f}\t{str(pred == gt)}\n'
                predicted_result_log += f'{dashed_line}'
                # print(predicted_result_log)
                log.write(predicted_result_log + '\n')

        # save model per 1e+4 iter.
        if (iteration + 1) % 1e+3 == 0:
            torch.save(
                model.state_dict(), f'./saved_models/{opt.exp_name}/SCATTER_STR.pth')

        if (iteration + 1) == opt.num_iter:
            print('end the training')
            sys.exit()
        # iteration += 1

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--exp_name', help='Where to store logs and models')
    parser.add_argument('--train_data', default="../data_lmdb_release/training/", help='path to training dataset')
    parser.add_argument('--valid_data', default="../data_lmdb_release/validation/", help='path to validation dataset')
    parser.add_argument('--manualSeed', type=int, default=1111, help='for random seed setting')
    parser.add_argument('--workers', type=int, help='number of data loading workers', default=4)
    parser.add_argument('--batch_size', type=int, default=256, help='input batch size')
    parser.add_argument('--num_iter', type=int, default=1000000, help='number of iterations to train for')
    parser.add_argument('--valInterval', type=int, default=500, help='Interval between each validation')
    parser.add_argument('--saved_model', default='saved_models/MJ-ST-20-Seed261504/SCATTER_STR.pth', help="path to model to continue training")
    parser.add_argument('--FT', action='store_true', help='whether to do fine-tuning')
    parser.add_argument('--adam', action='store_true', help='Whether to use adam (default is Adadelta)')
    parser.add_argument('--lr', type=float, default=1, help='learning rate, default=1.0 for Adadelta')
    parser.add_argument('--beta1', type=float, default=0.9, help='beta1 for adam. default=0.9')
    parser.add_argument('--rho', type=float, default=0.95, help='decay rate rho for Adadelta. default=0.95')
    parser.add_argument('--eps', type=float, default=1e-8, help='eps for Adadelta. default=1e-8')
    parser.add_argument('--grad_clip', type=float, default=5, help='gradient clipping value. default=5')
    """ Data processing """
    parser.add_argument('--select_data', type=str, default='MJ-ST',
                        help='select training data (default is MJ-ST, which means MJ and ST used as training data)')
    parser.add_argument('--batch_ratio', type=str, default='0.5-0.5',
                        help='assign ratio for each selected data in the batch')
    parser.add_argument('--total_data_usage_ratio', type=str, default='1.0',
                        help='total data usage ratio, this ratio is multiplied to total number of data.')
    parser.add_argument('--batch_max_length', type=int, default=35, help='maximum-label-length')
    parser.add_argument('--imgH', type=int, default=32, help='the height of the input image')
    parser.add_argument('--imgW', type=int, default=100, help='the width of the input image')
    parser.add_argument('--rgb', action='store_true', default=False, help='use rgb input')
    parser.add_argument('--character', type=str,
                        default='0123456789abcdefghijklmnopqrstuvwxyz', help='character label')
    parser.add_argument('--sensitive', default=True, action='store_true', help='for sensitive character mode')
    parser.add_argument('--PAD', action='store_true', help='whether to keep ratio then pad for image resize')
    parser.add_argument('--data_filtering_off', default=True, action='store_true', help='for data_filtering_off mode')
    """ Model Architecture """
    parser.add_argument('--num_fiducial', type=int, default=20, help='number of fiducial points of TPS-STN')
    parser.add_argument('--input_channel', type=int, default=1,
                        help='the number of input channel of Feature extractor')
    parser.add_argument('--output_channel', type=int, default=512,
                        help='the number of output channel of Feature extractor')
    parser.add_argument('--hidden_size', type=int, default=512, help='the size of the LSTM hidden state')
    parser.add_argument('--log', type=str, default='log', help='Save training loss')
    parser.add_argument('--lambda_ctc', type=float, default=0.1, help='refinement loss coeff')
    parser.add_argument('--lambda_attn', type=float, default=1.0, help='prediction loss coeff')

    opt = parser.parse_args()

    if not opt.exp_name:
        opt.exp_name = f'{opt.select_data}-{opt.num_fiducial}'
        opt.exp_name += f'-Seed{opt.manualSeed}'
        # print(opt.exp_name)

    os.makedirs(f'./saved_models/{opt.exp_name}', exist_ok=True)
    opt.log = f'./saved_models/{opt.exp_name}'
    """ vocab / character number configuration """
    if opt.sensitive:
        # opt.character += 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
        opt.character = string.printable[:-6]  # same with ASTER setting (use 94 char).

    """ Seed and GPU setting """
    # print("Random Seed: ", opt.manualSeed)
    random.seed(opt.manualSeed)
    np.random.seed(opt.manualSeed)
    torch.manual_seed(opt.manualSeed)
    torch.cuda.manual_seed(opt.manualSeed)

    cudnn.benchmark = True
    cudnn.deterministic = True
    opt.num_gpu = torch.cuda.device_count()
    # print('device count', opt.num_gpu)
    if opt.num_gpu > 1:
        print('------ Use multi-GPU setting ------')
        print('if you stuck too long time with multi-GPU setting, try to set --workers 0')
        # check multi-GPU issue https://github.com/clovaai/deep-text-recognition-benchmark/issues/1
        opt.workers = opt.workers * opt.num_gpu
        opt.batch_size = opt.batch_size * opt.num_gpu

        """ previous version
        print('To equilize batch stats to 1-GPU setting, the batch_size is multiplied with num_gpu and multiplied batch_size is ', opt.batch_size)
        opt.batch_size = opt.batch_size * opt.num_gpu
        print('To equalize the number of epochs to 1-GPU setting, num_iter is divided with num_gpu by default.')
        If you dont care about it, just comment out these line.)
        opt.num_iter = int(opt.num_iter / opt.num_gpu)
        """

    train(opt)


# Copyright 2019 Amazon.com, Inc. or its affiliates. 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.
# A copy of the License is located at
#     http://www.apache.org/licenses/LICENSE-2.0
# or in the "license" file accompanying this file. This file 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 mxnet as mx
import json
import os
import numpy as np
from collections import namedtuple
import logging

class MXNetVisionServiceBatching(object):
    def __init__(self):
        """
        Initialization for MXNet Vision Service supporting batch inference
        """
        self.mxnet_ctx = None
        self.mx_model = None
        self.labels = None
        self.signature = None
        self.epoch = 0
        self._context = None
        self._batch_size = 0
        self.initialized = False
        self.erroneous_reqs = set()

    def top_probability(self, data, labels, top=5):
        """
        Get top probability prediction from NDArray.

        :param data: NDArray
            Data to be predicted
        :param labels: List
            List of class labels
        :param top:
        :return: List
            List of probability: class pairs in sorted order
        """
        dim = len(data.shape)
        if dim > 2:
            data = mx.nd.array(
                np.squeeze(data.asnumpy(), axis=tuple(range(dim)[2:])))
        sorted_prob = mx.nd.argsort(data[0], is_ascend=False)
        top_prob = map(lambda x: int(x.asscalar()), sorted_prob[0:top])
        return [{'probability': float(data[0, i].asscalar()), 'class': labels[i]}
                for i in top_prob]

    def initialize(self, context):
        """
        Initialize model. This will be called during model loading time

        :param context: Initial context contains model server system properties.
        :return:
        """
        self._context = context
        self._batch_size = context.system_properties["batch_size"]
        self.initialized = True

        properties = context.system_properties
        model_dir = properties.get("model_dir")
        gpu_id = properties.get("gpu_id")

        signature_file_path = os.path.join(model_dir, "signature.json")
        if not os.path.isfile(signature_file_path):
            raise RuntimeError("Missing signature.json file.")

        with open(signature_file_path) as f:
            self.signature = json.load(f)

        model_files_prefix = context.manifest["model"]["modelName"]
        archive_synset = os.path.join(model_dir, "synset.txt")
        if os.path.isfile(archive_synset):
            synset = archive_synset
            self.labels = [line.strip() for line in open(synset).readlines()]

        data_names = []
        data_shapes = []
        for input_data in self.signature["inputs"]:
            data_name = input_data["data_name"]
            data_shape = input_data["data_shape"]

            # Set batch size
            data_shape[0] = self._batch_size

            # Replace 0 entry in data shape with 1 for binding executor.
            for idx in range(len(data_shape)):
                if data_shape[idx] == 0:
                    data_shape[idx] = 1

            data_names.append(data_name)
            data_shapes.append((data_name, tuple(data_shape)))

        checkpoint_prefix = "{}/{}".format(model_dir, model_files_prefix)

        # Load MXNet module
        self.mxnet_ctx = mx.cpu() if gpu_id is None else mx.gpu(gpu_id)
        sym, arg_params, aux_params = mx.model.load_checkpoint(checkpoint_prefix, self.epoch)

        self.mx_model = mx.mod.Module(symbol=sym, context=self.mxnet_ctx,
                                      data_names=data_names, label_names=None)
        self.mx_model.bind(for_training=False, data_shapes=data_shapes)
        self.mx_model.set_params(arg_params, aux_params, allow_missing=True, allow_extra=True)

    def inference(self, model_input):
        """
        Internal inference methods for MXNet. Run forward computation and
        return output.

        :param model_input: list of NDArray
            Preprocessed inputs in NDArray format.
        :return: list of NDArray
            Inference output.
        """
        batch = namedtuple('Batch', ['data'])

        self.mx_model.forward(batch([model_input]), is_train=False)
        outputs = self.mx_model.get_outputs()
        res = mx.ndarray.split(outputs[0], axis=0, num_outputs=outputs[0].shape[0])
        res = [res] if not isinstance(res, list) else res
        return res

    def preprocess(self, request):
        """
        Decode all input images into ndarray.

        Note: This implementation doesn't properly handle error cases in batch mode,
        If one of the input images is corrupted, all requests in the batch will fail.

        :param request:
        :return:
        """
        img_list = []
        param_name = self.signature['inputs'][0]['data_name']
        input_shape = self.signature['inputs'][0]['data_shape']
        # We are assuming input shape is NCHW
        [c, h, w] = input_shape[1:]

        # Clear error requests set.
        self.erroneous_reqs.clear()

        for idx, data in enumerate(request):
            img = data.get(param_name)
            if img is None:
                img = data.get("body")

            if img is None:
                img = data.get("data")

            if img is None or len(img) == 0:
                logging.error("Error processing request")
                self.erroneous_reqs.add(idx)
                continue

            try:
                img_arr = mx.image.imdecode(img, 1, True, None)
            except Exception as e:
                logging.error(e, exc_info=True)
                self.erroneous_reqs.add(idx)
                continue

            img_arr = mx.image.imresize(img_arr, w, h, 2)
            img_arr = mx.nd.transpose(img_arr, (2, 0, 1))
            self._num_requests = idx + 1
            img_list.append(img_arr)

        logging.debug("Worker :{} received {} requests".format(os.getpid(), self._num_requests))
        reqs = mx.nd.stack(*img_list)
        reqs = reqs.as_in_context(self.mxnet_ctx)

        if (self._batch_size - self._num_requests) != 0:
            padding = mx.nd.zeros((self._batch_size - self._num_requests, c, h, w), self.mxnet_ctx, 'uint8')
            reqs = mx.nd.concat(reqs, padding, dim=0)

        return reqs

    def postprocess(self, data):
        res = []
        for idx, resp in enumerate(data[:self._num_requests]):
            if idx not in self.erroneous_reqs:
                res.append(self.top_probability(resp, self.labels, top=5))
            else:
                res.append("This request was not processed successfully. Refer to mms.log for additional information")
        return res

_service = MXNetVisionServiceBatching()

def handle(data, context):
    if not _service.initialized:
        _service.initialize(context)

    if data is None:
        return None

    try:
        data = _service.preprocess(data)
        data = _service.inference(data)
        data = _service.postprocess(data)

        return data
    except Exception as e:
        logging.error(e, exc_info=True)
        request_processor = context.request_processor
        request_processor.report_status(500, "Unknown inference error")
        return [str(e)] * _service._batch_size


from .baseconfig import camblib, CAMB_Structure, CAMBError, dll_import
from ctypes import c_bool, c_int, c_double, c_float, byref, POINTER
from . import reionization as ion
from . import recombination as recomb
from . import initialpower as ipow
from . import constants
import numpy as np
from . import bbn
import logging

# ---Parameters

max_nu = 5
max_transfer_redshifts = 150
nthermo_derived = 13
Transfer_kh = 1
Transfer_cdm = 2
Transfer_b = 3
Transfer_g = 4
Transfer_r = 5
Transfer_nu = 6
Transfer_tot = 7
Transfer_nonu = 8
Transfer_tot_de = 9
Transfer_Weyl = 10
Transfer_Newt_vel_cdm = 11
Transfer_Newt_vel_baryon = 12
Transfer_vel_baryon_cdm = 13
Transfer_max = Transfer_vel_baryon_cdm

NonLinear_none = 0
NonLinear_pk = 1
NonLinear_lens = 2
NonLinear_both = 3

derived_names = ['age', 'zstar', 'rstar', 'thetastar', 'DAstar', 'zdrag',
                 'rdrag', 'kd', 'thetad', 'zeq', 'keq', 'thetaeq', 'thetarseq']

transfer_names = ['k/h', 'delta_cdm', 'delta_baryon', 'delta_photon', 'delta_neutrino', 'delta_nu', 'delta_tot',
                  'delta_nonu', 'delta_tot_de', 'Weyl', 'v_newtonian_cdm', 'v_newtonian_baryon', 'v_baryon_cdm']

evolve_names = transfer_names + ['a', 'etak', 'H', 'growth', 'v_photon', 'pi_photon', 'E_2', 'v_neutrino']

background_names = ['x_e', 'opacity', 'visibility', 'cs2b']

# ---Variables in modules.f90
# To set the value please just put
# variable_name.value = new_value

# logical
_HighAccuracyDefault = dll_import(POINTER(c_bool), "modelparams", "highaccuracydefault")
_HighAccuracyDefault.value = True

_lSampleBoost = dll_import(c_double, "modelparams", "lsampleboost")
_AccuracyBoost = dll_import(c_double, "modelparams", "accuracyboost")
_lAccuracyBoost = dll_import(c_float, "modelparams", "laccuracyboost")
_DoLateRadTruncation = dll_import(c_bool, "gaugeinterface", "dolateradtruncation")

DebugParam = dll_import(c_double, "modelparams", "debugparam")
# DebugParam.value = 1000000*2

# logical
do_bispectrum = dll_import(c_int, "modelparams", "do_bispectrum")
# do_bispectrum.value = False

max_bessels_l_index = dll_import(c_int, "modelparams", "max_bessels_l_index")
# max_bessels_l_index.value = 1000000

max_bessels_etak = dll_import(c_double, "modelparams", "max_bessels_etak")
# max_bessels_etak.value = 1000000*2

# logical
call_again = dll_import(c_int, "modelparams", "call_again")
# call_again.value = False

grhom = dll_import(c_double, "modelparams", "grhom")
grhog = dll_import(c_double, "modelparams", "grhog")
grhor = dll_import(c_double, "modelparams", "grhor")
grhob = dll_import(c_double, "modelparams", "grhob")
grhoc = dll_import(c_double, "modelparams", "grhoc")
grhov = dll_import(c_double, "modelparams", "grhov")
grhornomass = dll_import(c_double, "modelparams", "grhornomass")
grhok = dll_import(c_double, "modelparams", "grhok")

taurst = dll_import(c_double, "modelparams", "taurst")
dtaurec = dll_import(c_double, "modelparams", "dtaurec")
taurend = dll_import(c_double, "modelparams", "taurend")
tau_maxvis = dll_import(c_double, "modelparams", "tau_maxvis")
adotrad = dll_import(c_double, "modelparams", "adotrad")

grhormass = dll_import(c_double * max_nu, "modelparams", "grhormass")
nu_masses = dll_import(c_double * max_nu, "modelparams", "nu_masses")

akthom = dll_import(c_double, "modelparams", "akthom")
fHe = dll_import(c_double, "modelparams", "fhe")
Nnow = dll_import(c_double, "modelparams", "nnow")

limber_phiphi = dll_import(c_int, "modelparams", "limber_phiphi")
# limber_phiphi.value = 0

num_extra_redshiftwindows = dll_import(c_int, "modelparams", "num_extra_redshiftwindows")
# num_extra_redshiftwindows.value = 0

num_redshiftwindows = dll_import(c_int, "modelparams", "num_redshiftwindows")

# logical
use_spline_template = dll_import(c_bool, "modelparams", "use_spline_template")
# use_spline_template.value = True

ThermoDerivedParams = dll_import(c_double * nthermo_derived, "modelparams", "thermoderivedparams")
# ThermoDerivedParams.value = 1.

# logical
Log_lvalues = dll_import(c_bool, "lvalues", "log_lvalues")
# Log_lvalues.value = False

# Variables from module ModelData

# logical
has_cl_2D_array = dll_import(c_bool, "modeldata", "has_cl_2d_array")
# has_cl_2D_array.value = False

lmax_lensed = dll_import(c_int, "modeldata", "lmax_lensed")

# Variable from module Transfer
# logical
transfer_interp_matterpower = dll_import(c_bool, "transfer", "transfer_interp_matterpower")
# transfer_interp_matterpower.value = False

transfer_power_var = dll_import(c_int, "transfer", "transfer_power_var")
# transfer_power_var.value = Transfer_tot

# logical
get_growth_sigma8 = dll_import(c_bool, "transfer", "get_growth_sigma8")
# get_growth_sigma8.value = True

CAMB_validateparams = camblib.__camb_MOD_camb_validateparams
CAMB_validateparams.restype = c_bool

CAMB_setinitialpower = camblib.__handles_MOD_camb_setinitialpower

numpy_1d = np.ctypeslib.ndpointer(c_double, flags='C_CONTIGUOUS')
CAMB_primordialpower = camblib.__handles_MOD_camb_primordialpower
CAMB_primordialpower.restype = c_bool

class TransferParams(CAMB_Structure):
    """
    Object storing parameters for the matter power spectrum calculation. PK variables are for setting main outputs.
    Other entries are used internally, e.g. for sampling to get correct non-linear corrections and lensing.

    :ivar high_precision: True for more accuracy
    :ivar kmax: k_max to output
    :ivar k_per_logint: number of points per log k interval. If zero, set an irregular optimized spacing.
    :ivar PK_num_redshifts: number of redshifts to calculate
    :ivar PK_redshifts: redshifts to output for the matter transfer and power

    """
    _fields_ = [
        ("high_precision", c_int),  # logical
        ("num_redshifts", c_int),
        ("kmax", c_double),
        ("k_per_logint", c_int),
        ("redshifts", c_double * max_transfer_redshifts),
        ("PK_redshifts", c_double * max_transfer_redshifts),
        ("NLL_redshifts", c_double * max_transfer_redshifts),
        ("PK_redshifts_index", c_int * max_transfer_redshifts),
        ("NLL_redshifts_index", c_int * max_transfer_redshifts),
        ("PK_num_redshifts", c_int),
        ("NLL_num_redshifts", c_int)
    ]

class CAMBparams(CAMB_Structure):
    """
    Object storing the parameters for a CAMB calculation, including cosmological parameters and
    settings for what to calculate. When a new object is instantiated, default parameters are set automatically.
    """

    def __init__(self):
        getattr(camblib, '__camb_MOD_camb_setdefparams')(byref(self))

    _fields_ = [
        ("WantCls", c_int),  # logical
        ("WantTransfer", c_int),  # logical
        ("WantScalars", c_int),  # logical
        ("WantTensors", c_int),  # logical
        ("WantVectors", c_int),  # logical
        ("DoLensing", c_int),  # logical
        ("want_zstar", c_int),  # logical
        ("want_zdrag", c_int),  # logical
        ("PK_WantTransfer", c_int),  # logical
        ("NonLinear", c_int),
        ("Want_CMB", c_int),  # logical
        ("max_l", c_int),
        ("max_l_tensor", c_int),
        ("max_eta_k", c_double),
        ("max_eta_k_tensor", c_double),
        ("omegab", c_double),
        ("omegac", c_double),
        ("omegav", c_double),
        ("omegan", c_double),
        ("H0", c_double),
        ("TCMB", c_double),
        ("YHe", c_double),
        ("num_nu_massless", c_double),
        ("num_nu_massive", c_int),
        ("nu_mass_eigenstates", c_int),
        ("share_delta_neff", c_int),  # logical
        ("nu_mass_degeneracies", c_double * max_nu),
        ("nu_mass_fractions", c_double * max_nu),
        ("nu_mass_numbers", c_int * max_nu),
        ("scalar_initial_condition", c_int),
        ("OutputNormalization", c_int),
        ("AccuratePolarization", c_int),  # logical
        ("AccurateBB", c_int),  # logical
        ("AccurateReionization", c_int),  # logical
        ("MassiveNuMethod", c_int),
        ("InitPower", ipow.InitialPowerParams),
        ("Reion", ion.ReionizationParams),
        ("Recomb", recomb.RecombinationParams),
        ("Transfer", TransferParams),
        ("InitialConditionVector", c_double * 10),
        ("OnlyTransfers", c_int),  # logical
        ("DerivedParameters", c_int),  # logical
        ("ReionHist", ion.ReionizationHistory),
        ("flat", c_int),  # logical
        ("closed", c_int),  # logical
        ("open", c_int),  # logical
        ("omegak", c_double),
        ("curv", c_double),
        ("r", c_double),
        ("Ksign", c_double),
        ("tau0", c_double),
        ("chi0", c_double)
    ]

    def validate(self):
        """
        Do some quick tests for sanity

        :return: True if OK
        """
        return CAMB_validateparams(byref(self))

    def set_accuracy(self, AccuracyBoost=1., lSampleBoost=1., lAccuracyBoost=1.,
                     HighAccuracyDefault=True, DoLateRadTruncation=True):
        """
        Set parameters determining calculation accuracy (large values may give big slow down).
        Note curently these are set globally, not just per parameter set.

        :param AccuracyBoost: increase AccuracyBoost to decrease integration step size, increase density of k sampling, etc.
        :param lSampleBoost: increase lSampleBoost to increase density of L sampling for CMB
        :param lAccuracyBoost: increase lAccuracyBoost to increase the maximum L included in the Boltzmann hierarchies
        :param HighAccuracyDefault: True for Planck-level accuracy (False is WMAP)
        :param DoLateRadTruncation: If True, use approximation to radiation perturbation evolution at late times
        :return: self
        """
        _lSampleBoost.value = lSampleBoost
        _AccuracyBoost.value = AccuracyBoost
        _lAccuracyBoost.value = lAccuracyBoost
        _HighAccuracyDefault.value = HighAccuracyDefault
        _DoLateRadTruncation.value = DoLateRadTruncation
        logging.warning('accuracy parameters are changed globally, not yet per parameter set')
        return self

    def set_initial_power(self, initial_power_params):
        """
        Set the InitialPower primordial power spectrum parameters

        :param initial_power_params: :class:`.initialpower.InitialPowerParams` instance
        :return: self
        """
        assert (isinstance(initial_power_params, ipow.InitialPowerParams))
        CAMB_setinitialpower(byref(self), byref(initial_power_params))
        return self

    def set_bbn_helium(self, ombh2, delta_nnu, tau_neutron=bbn.tau_n):
        """
        Set the Helium abundance parameter YHe using BBN consistency (using fitting formula as Planck 2015 papers)

        :param ombh2: physical density of baryons
        :param delta_nnu: additional relativistic Delta_Neff = N_eff - 3.046
        :param tau_neutron: neutron half life in seconds
        :return: self
        """
        Yp = bbn.yhe_fit(ombh2, delta_nnu, tau_neutron)
        self.YHe = bbn.ypBBN_to_yhe(Yp)
        return self

    def set_cosmology(self, H0=67, cosmomc_theta=None, ombh2=0.022, omch2=0.12, omk=0.0, num_massive_neutrinos=1,
                      mnu=0.06, nnu=3.046,
                      YHe=None, meffsterile=0, standard_neutrino_neff=3.046, TCMB=constants.COBE_CMBTemp, tau=None,
                      tau_neutron=bbn.tau_n):
        """
        Sets cosmological parameters in terms of physical densities and parameters used in Planck 2015 analysis.
        Assumes a single distinct neutrino mass eigenstate, by default one neutrino with mnu = 0.06eV.
        If you require more fine-grained control can set the neutrino parameters directly rather than using this function.

        :param H0: Hubble parameter (in km/s/Mpc)
        :param cosmomc_theta The CosmoMC theta parameter. You must set H0=None to solve for H0 given cosmomc_theta
        :param ombh2: physical density in baryons
        :param omch2:  physical density in cold dark matter
        :param omk: Omega_K curvature parameter
        :param num_massive_neutrinos:  number of massive neutrinos
        :param mnu: sum of neutrino masses (in eV)
        :param nnu: N_eff, effective relativistic degrees of freedom
        :param YHe: Helium mass fraction. If None, set from BBN consistency.
        :param meffsterile: effective mass of sterile neutrinos
        :param standard_neutrino_neff:  default value for N_eff in standard cosmology (non-integer to allow for partial
                heating of neutrinos at electron-positron annihilation and QED effects)
        :param TCMB: CMB temperature (in Kelvin)
        :param tau: optical depth; if None, current Reion settings are not changed
        :param tau_neutron: neutron lifetime, for setting YHe using BBN consistency
        """

        if YHe is None:
            # use BBN prediction
            self.set_bbn_helium(ombh2, nnu - standard_neutrino_neff, tau_neutron)
            YHe = self.YHe
        else:
            self.YHe = YHe

        if cosmomc_theta is not None:
            if not (0.001 < cosmomc_theta < 0.1):
                raise CAMBError('cosmomc_theta looks wrong (parameter is just theta, not 100*theta)')

            kw = locals();
            [kw.pop(x) for x in ['self', 'H0', 'cosmomc_theta']]

            if H0 is not None:
                raise CAMBError('Set H0=None when setting cosmomc_theta.')

            try:
                from scipy.optimize import brentq
            except ImportError:
                raise CAMBError('You need SciPy to set cosmomc_theta.')

            from . import camb

            def f(H0):
                self.set_cosmology(H0=H0, **kw)
                return camb.get_background(self, no_thermo=True).cosmomc_theta() - cosmomc_theta

            self.H0 = brentq(f, 10, 100, rtol=1e-4)
        else:
            self.H0 = H0

        self.TCMB = TCMB
        fac = (self.H0 / 100.0) ** 2
        self.omegab = ombh2 / fac
        self.omegac = omch2 / fac

        neutrino_mass_fac = 94.07
        # conversion factor for thermal with Neff=3 TCMB=2.7255

        omnuh2 = mnu / neutrino_mass_fac * (standard_neutrino_neff / 3.0) ** 0.75
        omnuh2_sterile = meffsterile / neutrino_mass_fac
        if omnuh2_sterile > 0 and nnu < standard_neutrino_neff:
            raise CAMBError('sterile neutrino mass required Neff>3.046')
        if omnuh2 and not num_massive_neutrinos:
            raise CAMBError('non-zero mnu with zero num_massive_neutrinos')

        omnuh2 = omnuh2 + omnuh2_sterile
        self.omegan = omnuh2 / fac
        self.omegav = 1 - omk - self.omegab - self.omegac - self.omegan
        self.share_delta_neff = False
        self.nu_mass_eigenstates = 0
        self.num_nu_massless = nnu
        self.nu_mass_numbers[0] = 0
        self.num_nu_massive = 0
        if omnuh2 > omnuh2_sterile:
            neff_massive_standard = num_massive_neutrinos * standard_neutrino_neff / 3.0
            self.num_nu_massive = num_massive_neutrinos
            self.nu_mass_eigenstates = self.nu_mass_eigenstates + 1
            if nnu > neff_massive_standard:
                self.num_nu_massless = nnu - neff_massive_standard
            else:
                self.num_nu_massless = 0
                neff_massive_standard = nnu

            self.nu_mass_numbers[self.nu_mass_eigenstates - 1] = num_massive_neutrinos
            self.nu_mass_degeneracies[self.nu_mass_eigenstates - 1] = neff_massive_standard
            self.nu_mass_fractions[self.nu_mass_eigenstates - 1] = (omnuh2 - omnuh2_sterile) / omnuh2
        else:
            neff_massive_standard = 0
        if omnuh2_sterile > 0:
            if nnu < standard_neutrino_neff:
                raise CAMBError('nnu < 3.046 with massive sterile')
            self.num_nu_massless = standard_neutrino_neff - neff_massive_standard
            self.num_nu_massive = self.num_nu_massive + 1
            self.nu_mass_eigenstates = self.nu_mass_eigenstates + 1
            self.nu_mass_numbers[self.nu_mass_eigenstates - 1] = 1
            self.nu_mass_degeneracies[self.nu_mass_eigenstates - 1] = max(1e-6, nnu - standard_neutrino_neff)
            self.nu_mass_fractions[self.nu_mass_eigenstates - 1] = omnuh2_sterile / omnuh2

        if tau is not None:
            self.Reion.set_tau(tau)

        return self

    def set_dark_energy(self, w=-1.0, sound_speed=1.0, dark_energy_model='fluid'):
        """
        Set dark energy parameters. Not that in this version these are not actually stored in
        the CAMBparams variable but set globally. So be careful!

        :param w: p_de/rho_de, assumed constant
        :param sound_speed: rest-frame sound speed of dark energy fluid
        :param dark_energy_model: model to use, default is 'fluid'
        :return: self
        """
        # Variables from module LambdaGeneral
        if dark_energy_model != 'fluid':
            raise CAMBError('This version only supports the fluid energy model')
        if w != -1 or sound_speed != 1:
            logging.warning('Currently dark energy parameters are changed globally, not per parameter set')
        w_lam = dll_import(c_double, "lambdageneral", "w_lam")
        w_lam.value = w
        cs2_lam = dll_import(c_double, "lambdageneral", "cs2_lam")
        cs2_lam.value = sound_speed
        return self

    def get_omega_k(self):
        """
        Get curvature parameter Omega_k

        :return: Omega_k
        """
        return 1 - self.omegab - self.omegac - self.omegan - self.omegav

    def set_matter_power(self, redshifts=[0.], kmax=1.2, k_per_logint=None, silent=False):
        """
        Set parameters for calculating matter power spectra and transfer functions.

        :param redshifts: array of redshifts to calculate
        :param kmax: maximum k to calculate
        :param k_per_logint: number of k steps per log k. Set to zero to use default optimized spacing.
        :param silent: if True, don't give warnings about sort order
        :return:  self
        """
        self.WantTransfer = True
        self.Transfer.high_precision = True
        self.Transfer.kmax = kmax
        if not k_per_logint:
            self.Transfer.k_per_logint = 0
        else:
            self.Transfer.k_per_logint = k_per_logint
        zs = sorted(redshifts, reverse=True)
        if not silent and np.any(np.array(zs) - np.array(redshifts) != 0):
            print("Note: redshifts have been re-sorted (earliest first)")
        if len(redshifts) > max_transfer_redshifts:
            raise CAMBError('You can have at most %s redshifts' % max_transfer_redshifts)
        self.Transfer.PK_num_redshifts = len(redshifts)
        for i, z in enumerate(zs):
            self.Transfer.PK_redshifts[i] = z
        return self

    def set_for_lmax(self, lmax, max_eta_k=None, lens_potential_accuracy=0,
                     lens_margin=150, k_eta_fac=2.5, lens_k_eta_reference=18000.0):
        """
        Set parameters to get CMB power spectra accurate to specific a l_lmax.
        Note this does not fix the actual output L range, spectra may be calculated above l_max (but may not be accurate there).
        To fix the l_max for output arrays use the optional input argument to :meth:`.camb.CAMBdata.get_cmb_power_spectra` etc.

        :param lmax: l_max you want
        :param max_eta_k: maximum value of k*eta_* to use, which indirectly sets k_max. If None, sensible value set automatically.
        :param lens_potential_accuracy: Set to 1 or higher if you want to get the lensing potential accurate
        :param lens_margin: the delta l_max to use to ensure lensed C_L are correct at l_max
        :param k_eta_fac:  k_eta_fac default factor for setting max_eta_k = k_eta_fac*lmax if max_eta_k=None
        :param lens_k_eta_reference:  value of max_eta_k to use when lens_potential_accuracy>0; use k_eta_max = lens_k_eta_reference*lens_potential_accuracy
        :return: self
        """
        if self.DoLensing:
            self.max_l = lmax + lens_margin
        else:
            self.max_l = lmax
        self.max_eta_k = max_eta_k or min(self.max_l, 3000) * k_eta_fac
        if lens_potential_accuracy:
            if self.NonLinear == NonLinear_none:
                self.NonLinear = NonLinear_lens
            else:
                self.NonLinear = NonLinear_both
            self.max_eta_k = max(self.max_eta_k, lens_k_eta_reference * lens_potential_accuracy)
        return self

    def scalar_power(self, k):
        return self.primordial_power(k, 0)

    def tensor_power(self, k):
        return self.primordial_power(k, 2)

    def primordial_power(self, k, ix):
        if np.isscalar(k):
            karr = np.array([float(k)])
        else:
            karr = np.array(k)
        n = karr.shape[0]
        powers = np.empty(n)
        CAMB_primordialpower(byref(self), karr, powers, byref(c_int(n)), byref(c_int(ix)))
        if np.isscalar(k):
            return powers[0]
        else:
            return powers

def Transfer_SetForNonlinearLensing(P):
    camblib.__transfer_MOD_transfer_setfornonlinearlensing(byref(P))

def Transfer_SortAndIndexRedshifts(P):
    camblib.__transfer_MOD_transfer_sortandindexredshifts(byref(P))

CAMB_primordialpower.argtypes = [POINTER(CAMBparams), numpy_1d, numpy_1d, POINTER(c_int), POINTER(c_int)]


"""
Class definition of YOLO_v3 style detection model on image and video
"""
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import cv2
import colorsys
import os
from timeit import default_timer as timer
from keras.layers.convolutional import Conv2D
import tensorflow as tf
from keras import models
from keras import layers

import numpy as np
from keras import backend as K
from keras.models import load_model
from keras.layers import Input
from PIL import Image, ImageFont, ImageDraw
from keras.layers import Conv2D, Add, ZeroPadding2D, UpSampling2D, Concatenate, MaxPooling2D

from yolo3.model import yolo_eval, yolo_body, tiny_yolo_body
from yolo3.utils import letterbox_image
import os
from keras.utils import multi_gpu_model

class YOLO(object):
    _defaults = {
        "model_path": 'model_data/yolo.h5',
        "anchors_path": 'model_data/yolo_anchors.txt',
        "classes_path": 'model_data/coco_classes.txt',
        "score" : 0.3,
        "iou" : 0.45,
        "model_image_size" : (416, 416),
        "gpu_num" : 1,
    }

    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"

    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults) # set up default values
        self.__dict__.update(kwargs) # and update with user overrides
        self.class_names = self._get_class()
        self.anchors = self._get_anchors()
        self.sess = K.get_session()
        self.boxes, self.scores, self.classes = self.generate()

    def _get_class(self):
        classes_path = os.path.expanduser(self.classes_path)
        with open(classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names

    def _get_anchors(self):
        anchors_path = os.path.expanduser(self.anchors_path)
        with open(anchors_path) as f:
            anchors = f.readline()
        anchors = [float(x) for x in anchors.split(',')]
        return np.array(anchors).reshape(-1, 2)

    def generate(self):
        model_path = os.path.expanduser(self.model_path)
        assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'

        # Load model, or construct model and load weights.
        num_anchors = len(self.anchors)
        num_classes = len(self.class_names)
        is_tiny_version = num_anchors==6 # default setting
        try:
            self.yolo_model = load_model(model_path, compile=False)
        except:
            self.yolo_model = tiny_yolo_body(Input(shape=(None,None,3)), num_anchors//2, num_classes) \
                if is_tiny_version else yolo_body(Input(shape=(None,None,3)), num_anchors//3, num_classes)
            self.yolo_model.load_weights(self.model_path) # make sure model, anchors and classes match
        else:
            assert self.yolo_model.layers[-1].output_shape[-1] == \
                num_anchors/len(self.yolo_model.output) * (num_classes + 5), \
                'Mismatch between model and given anchor and class sizes'

        print('{} model, anchors, and classes loaded.'.format(model_path))

        # Generate colors for drawing bounding boxes.
        hsv_tuples = [(x / len(self.class_names), 1., 1.)
                      for x in range(len(self.class_names))]
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(
            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
                self.colors))
        np.random.seed(10101)  # Fixed seed for consistent colors across runs.
        np.random.shuffle(self.colors)  # Shuffle colors to decorrelate adjacent classes.
        np.random.seed(None)  # Reset seed to default.

        # Generate output tensor targets for filtered bounding boxes.
        self.input_image_shape = K.placeholder(shape=(2, ))
        if self.gpu_num>=2:
            self.yolo_model = multi_gpu_model(self.yolo_model, gpus=self.gpu_num)
        boxes, scores, classes = yolo_eval(self.yolo_model.output, self.anchors,
                len(self.class_names), self.input_image_shape,
                score_threshold=self.score, iou_threshold=self.iou)
        return boxes, scores, classes

    def detect_image(self, image):
        start = timer()
        if self.model_image_size != (None, None):
            assert self.model_image_size[0]%32 == 0, 'Multiples of 32 required'
            assert self.model_image_size[1]%32 == 0, 'Multiples of 32 required'
            boxed_image = letterbox_image(image, tuple(reversed(self.model_image_size)))
        else:
            new_image_size = (image.width - (image.width % 32),
                              image.height - (image.height % 32))
            boxed_image = letterbox_image(image, new_image_size)
        image_data = np.array(boxed_image, dtype='float32')

        print(image_data.shape)
        image_data /= 255.
        image_data = np.expand_dims(image_data, 0)  # Add batch dimension.

        print(self.yolo_model.layers)
        #  my code here:
        print("Bingo! let's look into it!")
        print(self.yolo_model.layers[1])
        self.yolo_model.summary()
        print(self.yolo_model.get_layer("conv2d_1").output[0,0,0,1])
        for layer in self.yolo_model.layers:
            print(layer)
            print(layer.name)
        layer_dict = dict([(layer.name, layer) for layer in self.yolo_model.layers])

        filter_out = []
        filter_num_sum = 0
            # for i in range(75):
            # layer_name = "conv2d_%d" %(i+1)
        layer_name = "conv2d_70"
        print(layer_name)
        layer_output = layer_dict[layer_name].output
        print("this is the output")
        print(layer_output)
        print(layer_output[:,:,:,0])
        filter_shape = K.shape(layer_output)
        filter_num = filter_shape[3]
        #  loss = tf.math.divide (sum_val, tf.shape (abs_val))
        #  loss = K.mean(layer_output[:, :, :, filter_index])brooch        #  part1 of my code ends
        out_boxes, out_scores, out_classes, layer_output, filter_num = self.sess.run(
                                                           [self.boxes, self.scores, self.classes, layer_output, filter_num],
                                                           feed_dict={
                                                           self.yolo_model.input: image_data,
                                                           self.input_image_shape: [image.size[1], image.size[0]],
                                                           K.learning_phase(): 0
                                                           })
        print("filter_num",filter_num)
        print("this is the layer_output",layer_output)
        filter_num_sum += filter_num
        # filter_num is the index of the number of filter
        abs_val = K.abs(layer_output[:, :, :, 22])
        sum_val = K.sum(abs_val)
        sum_val = K.eval(sum_val)
        print("sum_val",sum_val)
        tensor_shape = K.shape(abs_val)
        sum_num = tensor_shape[1]*tensor_shape[2]
        sum_num = tf.to_float(sum_num, name='ToFloat')
        sum_num = K.eval(sum_num)
        print("sum_num", sum_num)
        loss = sum_val/sum_num
        print("loss", loss)
        filter_out.append(loss)
        print("filter_num_sum", filter_num_sum)
        print(filter_out)
#  x = np.arange(0, filter_num_sum, 1)
#  lines = plt.plot(x, filter_out, 'o')
#  plt.show()
        print(tensor_shape)
        print(tensor_shape[1])
        print(tensor_shape[2])
        print("filter_num: ", filter_num)
        # my code_2 here:
        # compute the gradient of the input picture wrt this loss
        # true_loss = K.mean(layer_output[:, :, :, filter_index])
        # grads = K.gradients(true_loss, self.yolo_model.input)[0]
        # layer_dict= dict([(layer.name, layer) for layer in self.yolo_model.layers])
        # layer_name = 'conv2d_1'
        # filter_index = 0
        # layer_output = layer_dict[layer_name].output
        # normalization trick: we normalize the gradient
        # grads /= (K.sqrt(K.mean(K.square(grads))) + 1e-5)

        # this function returns the loss and grads given the input picture
        # iterate = K.function([self.yolo_model.input], [true_loss, grads])

#   init = tf.global_variables_initializer()         # When init is run later(session.run(init)),
#   with tf.Session() as session:                    # Create a session and print the output
#            session.run(init)                            # Initializes the variables
#            print(session.run(loss))

        #input_data= mpimg.imread('/Users/chenyun/Downloads/7.png')
        #print("shape of input image: ", input_data.shape)
        #input_data = cv2.resize(input_data, (224, 224))
        #print(input_data.shape)
        #input_img_data[0,:,:,:]=input_data

# run gradient ascent for 500 steps
#        for i in range(500):
#            loss_value, grads_value = iterate([input_img_data])
#            input_img_data += grads_value * 10
#            print(loss_value)

# util function to convert a tensor into a valid image
#   def deprocess_image(x):
# normalize tensor: center on 0., ensure std is 0.1
#        x -= x.mean()
#        x /= (x.std() + 1e-5)
#        x *= 0.1

# clip to [0, 1]
#        x += 0.5
#        x = np.clip(x, 0, 1)

# convert to RGB array
#        x *= 255
#        x = x.transpose((1, 2, 0))
#        x = np.clip(x, 0, 255).astype('uint8')
#        return 0

#        img = input_data_converted[0]
#        img = deprocess_image(img)

# print(img.shape)
# img=np.transpose(img, (0, 2, 1))
# imgplot = plt.imshow(img)
#        print(img.shape)

# img = input_img_data[0]
# img_2 = deprocess_image(img)

# print(img_2.shape)
# img_2=np.transpose(img_2, (0, 2, 1))
# plt.savefig('img_2.png')
        #imgplot=plt.imshow(img_2)
# print(img_2.shape)

# part_2 code ends here

        layer_output = layer_dict[layer_name].output

        print('Found {} boxes for {}, hello'.format(len(out_boxes), 'img'))
        font = ImageFont.truetype(font='font/FiraMono-Medium.otf',
                    size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
        thickness = (image.size[0] + image.size[1]) // 300

        for i, c in reversed(list(enumerate(out_classes))):
            predicted_class = self.class_names[c]
            box = out_boxes[i]
            score = out_scores[i]

            label = '{} {:.2f}'.format(predicted_class, score)
            draw = ImageDraw.Draw(image)
            label_size = draw.textsize(label, font)

            top, left, bottom, right = box
            top = max(0, np.floor(top + 0.5).astype('int32'))
            left = max(0, np.floor(left + 0.5).astype('int32'))
            bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
            right = min(image.size[0], np.floor(right + 0.5).astype('int32'))
            print(label, (left, top), (right, bottom))

            if top - label_size[1] >= 0:
                text_origin = np.array([left, top - label_size[1]])
            else:
                text_origin = np.array([left, top + 1])

            # My kingdom for a good redistributable image drawing library.
            for i in range(thickness):
                draw.rectangle(
                    [left + i, top + i, right - i, bottom - i],
                    outline=self.colors[c])
            draw.rectangle(
                [tuple(text_origin), tuple(text_origin + label_size)],
                fill=self.colors[c])
            draw.text(text_origin, label, fill=(0, 0, 0), font=font)
            del draw

        end = timer()
        print(end - start)

        return loss

    def close_session(self):
        self.sess.close()

def detect_video(yolo, video_path, output_path=""):
    import cv2
    vid = cv2.VideoCapture(video_path)
    if not vid.isOpened():
        raise IOError("Couldn't open webcam or video")
    video_FourCC    = int(vid.get(cv2.CAP_PROP_FOURCC))
    video_fps       = vid.get(cv2.CAP_PROP_FPS)
    video_size      = (int(vid.get(cv2.CAP_PROP_FRAME_WIDTH)),
                        int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    isOutput = True if output_path != "" else False
    if isOutput:
        print("!!! TYPE:", type(output_path), type(video_FourCC), type(video_fps), type(video_size))
        out = cv2.VideoWriter(output_path, video_FourCC, video_fps, video_size)
    accum_time = 0
    curr_fps = 0
    fps = "FPS: ??"
    prev_time = timer()
    while True:
        return_value, frame = vid.read()
        image = Image.fromarray(frame)
        image = yolo.detect_image(image)
        result = np.asarray(image)
        curr_time = timer()
        exec_time = curr_time - prev_time
        prev_time = curr_time
        accum_time = accum_time + exec_time
        curr_fps = curr_fps + 1
        if accum_time > 1:
            accum_time = accum_time - 1
            fps = "FPS: " + str(curr_fps)
            curr_fps = 0
        cv2.putText(result, text=fps, org=(3, 15), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                    fontScale=0.50, color=(255, 0, 0), thickness=2)
        cv2.namedWindow("result", cv2.WINDOW_NORMAL)
        cv2.imshow("result", result)
        if isOutput:
            out.write(result)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    yolo.close_session()


'''edit.py'''

# Internal Libraries
import os
from typing import List, Tuple

import numpy as np

from auto_editor.utils.effects import Effect
from auto_editor.utils.log import Log
from auto_editor.utils.func import fnone, set_output_name, append_filename

def get_chunks(inp, segment, fps, args, log, audio_samples=None, sample_rate=None):
    from auto_editor.cutting import (combine_audio_motion, combine_segment, to_speed_list,
        set_range, chunkify, apply_mark_as, apply_margin, seconds_to_frames, cook)

    start_margin, end_margin = args.frame_margin

    start_margin = seconds_to_frames(start_margin, fps)
    end_margin = seconds_to_frames(end_margin, fps)
    min_clip = seconds_to_frames(args.min_clip_length, fps)
    min_cut = seconds_to_frames(args.min_cut_length, fps)

    def get_has_loud(inp, args, fps, audio_samples, sample_rate, log):
        # type: (...) -> np.ndarray
        from auto_editor.analyze.generic import get_np_list

        if(args.edit_based_on == 'none'):
            return get_np_list(inp, audio_samples, sample_rate, fps, np.ones)
        if(args.edit_based_on == 'all'):
            return get_np_list(inp, audio_samples, sample_rate, fps, np.zeros)

        audio_list, motion_list = None, None

        if('audio' in args.edit_based_on):
            from auto_editor.analyze.audio import audio_detection

            if(audio_samples is None):
                audio_list = get_np_list(inp, audio_samples, sample_rate, fps, np.ones)
            else:
                audio_list = audio_detection(audio_samples, sample_rate,
                    args.silent_threshold, fps, log)

        if('motion' in args.edit_based_on):
            from auto_editor.analyze.motion import motion_detection

            if(len(inp.video_streams) == 0):
                motion_list = get_np_list(inp, audio_samples, sample_rate, fps, np.ones)
            else:
                motion_list = motion_detection(inp, args.motion_threshold, log,
                    width=args.md_width, dilates=args.md_dilates, blur=args.md_blur)

        if(audio_list is not None and motion_list is not None):
            if(len(audio_list) > len(motion_list)):
                audio_list = audio_list[:len(motion_list)]
            elif(len(motion_list) > len(audio_list)):
                motion_list = motion_list[:len(audio_list)]

        return combine_audio_motion(audio_list, motion_list, args.edit_based_on, log)

    has_loud = get_has_loud(inp, args, fps, audio_samples, sample_rate, log)
    has_loud_length = len(has_loud)
    has_loud = apply_mark_as(has_loud, has_loud_length, fps, args, log)
    has_loud = cook(has_loud, min_clip, min_cut)
    has_loud = apply_margin(has_loud, has_loud_length, start_margin, end_margin)

    if(segment is not None):
        has_loud = combine_segment(has_loud, segment, fps)

    # Remove small clips/cuts created by applying other rules.
    has_loud = cook(has_loud, min_clip, min_cut)

    speed_list = to_speed_list(has_loud, args.video_speed, args.silent_speed)

    if(args.cut_out != []):
        speed_list = set_range(speed_list, args.cut_out, fps, 99999, log)

    if(args.add_in != []):
        speed_list = set_range(speed_list, args.add_in, fps, args.video_speed, log)

    if(args.set_speed_for_range != []):
        for item in args.set_speed_for_range:
            speed_list = set_range(speed_list, [item[1:]], fps, item[0], log)

    return chunkify(speed_list, has_loud_length)

def edit_media(i, inp, ffmpeg, args, progress, segment, exporting_to_editor, data_file,
    temp, log):
    from auto_editor.utils.container import get_rules

    chunks = None
    if(inp.ext == '.json'):
        from auto_editor.formats.make_json import read_json_cutlist

        input_path, chunks = read_json_cutlist(inp.path, log)
        inp = ffmpeg.file_info(input_path)

        output_path = set_output_name(inp.path, inp.ext, data_file, args)
    else:
        output_path = args.output_file[i]
        if(not os.path.isdir(inp.path) and os.path.splitext(output_path)[1] == ''):
            output_path = set_output_name(output_path, inp.ext, data_file, args)

    log.debug('{} -> {}'.format(inp.path, output_path))

    output_container = os.path.splitext(output_path)[1].replace('.', '')

    # Check if export options make sense.
    rules = get_rules(output_container)
    codec_error = "'{}' codec is not supported in '{}' container."

    if(not fnone(args.sample_rate)):
        if(rules['samplerate'] is not None and args.sample_rate not in rules['samplerate']):
            log.error("'{}' container only supports samplerates: {}".format(output_container,
                rules['samplerate']))

    vcodec = args.video_codec
    if(vcodec == 'uncompressed'):
        vcodec = 'mpeg4'
    if(vcodec == 'copy'):
        vcodec = inp.video_streams[0]['codec']

    if(vcodec != 'auto'):
        if(rules['vstrict'] and vcodec not in rules['vcodecs']):
            log.error(codec_error.format(vcodec, output_container))

        if(vcodec in rules['disallow_v']):
            log.error(codec_error.format(vcodec, output_container))

    acodec = args.audio_codec
    if(acodec == 'copy'):
        acodec = inp.audio_streams[0]['codec']
        log.debug(f'Settings acodec to {acodec}')

    if(acodec not in ['unset', 'auto']):
        if(rules['astrict'] and acodec not in rules['acodecs']):
            log.error(codec_error.format(acodec, output_container))

        if(acodec in rules['disallow_a']):
            log.error(codec_error.format(acodec, output_container))

    if(args.keep_tracks_seperate and rules['max_audio_streams'] == 1):
        log.warning(
            "'{}' container doesn't support multiple audio tracks.".format(container)
        )

    if(os.path.isfile(output_path) and inp.path != output_path and not args.preview
        and not args.timeline):
        log.debug('Removing already existing file: {}'.format(output_path))
        os.remove(output_path)

    audio_samples = None
    tracks = len(inp.audio_streams)

    fps = 30 if inp.fps is None else float(inp.fps)

    if(fps < 1):
        log.error('{}: Frame rate cannot be below 1. fps: {}'.format(inp.basename, fps))

    # Extract subtitles in their native format.
    if(len(inp.subtitle_streams) > 0):
        cmd = ['-i', inp.path, '-hide_banner']
        for s, sub in enumerate(inp.subtitle_streams):
            cmd.extend(['-map', '0:s:{}'.format(s)])
        for s, sub in enumerate(inp.subtitle_streams):
            cmd.extend([os.path.join(temp, '{}s.{}'.format(s, sub['ext']))])
        ffmpeg.run(cmd)

    sample_rate = None

    if(args.cut_by_this_track >= tracks and 'cut_by_this_track' in args._set):
        message = "You choose a track that doesn't exist.\nThere "
        if(tracks == 1):
            message += 'is only {} track.\n'.format(tracks)
        else:
            message += 'are only {} tracks.\n'.format(tracks)
        for t in range(tracks):
            message += ' Track {}\n'.format(t)
        log.error(message)

    # Split audio tracks into: 0.wav, 1.wav, etc.
    log.conwrite('Extracting audio')

    cmd = ['-i', inp.path, '-hide_banner']
    for t in range(tracks):
        cmd.extend(['-map', '0:a:{}'.format(t), '-ac', '2',
            os.path.join(temp, '{}.wav'.format(t))])
    ffmpeg.run(cmd)
    del cmd

    if(tracks != 0):
        if(args.cut_by_all_tracks):
            temp_file = os.path.join(temp, 'combined.wav')
            cmd = ['-i', inp.path, '-filter_complex',
                '[0:a]amix=inputs={}:duration=longest'.format(tracks), '-ac', '2',
                '-f', 'wav', temp_file]
            ffmpeg.run(cmd)
            del cmd
        else:
            temp_file = os.path.join(temp, '{}.wav'.format(args.cut_by_this_track))

        from auto_editor.scipy.wavfile import read
        sample_rate, audio_samples = read(temp_file)

    if(chunks is None):
        chunks = get_chunks(inp, segment, fps, args, log, audio_samples, sample_rate)

    def is_clip(chunk: Tuple[int, int, float]) -> bool:
        return chunk[2] != 99999

    def number_of_cuts(chunks):
        # type: (List[Tuple[int, int, float]]) -> int
        return len(list(filter(is_clip, chunks)))

    num_cuts = number_of_cuts(chunks)

    effects = Effect(args, log, _vars={})

    if(args.export_as_json):
        from auto_editor.formats.make_json import make_json_cutlist
        make_json_cutlist(inp.path, output_path, chunks, log)
        return num_cuts, output_path

    if(args.timeline):
        from auto_editor.formats.make_json import make_json_cutlist
        make_json_cutlist(inp.path, 0, chunks, log)
        return num_cuts, None

    if(args.preview):
        from auto_editor.preview import preview
        preview(inp, chunks, log)
        return num_cuts, None

    if(args.export_to_premiere):
        from auto_editor.formats.premiere import premiere_xml
        premiere_xml(inp, temp, output_path, chunks, sample_rate, fps, log)
        return num_cuts, output_path

    if(args.export_to_final_cut_pro):
        from auto_editor.formats.final_cut_pro import fcp_xml

        fcp_xml(inp, temp, output_path, chunks, fps, log)
        return num_cuts, output_path

    if(args.export_to_shotcut):
        from auto_editor.formats.shotcut import shotcut_xml

        shotcut_xml(inp, temp, output_path, chunks, fps, log)
        return num_cuts, output_path

    def pad_chunk(item, total_frames):
        start = None
        end = None
        if(item[0] != 0):
            start = [0, item[0], 2]
        if(item[1] != total_frames -1):
            end = [item[1], total_frames -1, 2]

        if(start is None):
            return [item] + [end]
        if(end is None):
            return [start] + [item]
        return [start] + [item] + [end]

    def make_media(inp, chunks, output_path):
        from auto_editor.utils.video import mux_quality_media

        if(rules['allow_subtitle']):
            from auto_editor.render.subtitle import cut_subtitles
            cut_subtitles(ffmpeg, inp, chunks, fps, temp, log)

        if(rules['allow_audio']):
            from auto_editor.render.audio import make_new_audio

            for t in range(tracks):
                temp_file = os.path.join(temp, '{}.wav'.format(t))
                new_file = os.path.join(temp, 'new{}.wav'.format(t))
                make_new_audio(temp_file, new_file, chunks, log, fps, progress)

                if(not os.path.isfile(new_file)):
                    log.bug('Audio file not created.')

        video_stuff = []

        if(rules['allow_video']):
            from auto_editor.render.video import render_av
            for v, vid in enumerate(inp.video_streams):
                if(vid['codec'] in ['png', 'jpeg']):
                    video_stuff.append(('image', None, None))
                else:
                    video_stuff.append(render_av(ffmpeg, v, inp, args, chunks, fps,
                        progress, effects, rules, temp, log))

        log.conwrite('Writing the output file.')

        mux_quality_media(ffmpeg, video_stuff, rules, output_path, output_container,
            args, inp, temp, log)
        if(output_path is not None and not os.path.isfile(output_path)):
            log.bug(f'The file {output_path} was not created.')

    if(args.export_as_clip_sequence):
        total_frames = chunks[-1][1]
        clip_num = 0
        for chunk in chunks:
            if(chunk[2] == 99999):
                continue
            make_media(inp, pad_chunk(chunk, total_frames),
                append_filename(output_path, f'-{clip_num}'))
            clip_num += 1
    else:
        make_media(inp, chunks, output_path)
    return num_cuts, output_path


import math
import random
import time
import os
import sys
from datetime import datetime

import collections
from collections import namedtuple
from collections import deque
from itertools import count

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision.transforms as T

import matplotlib
import matplotlib.pyplot as plt

import numpy as np
from pymss import EnvManager
from IPython import embed
from Model import *
use_cuda = torch.cuda.is_available()
FloatTensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if use_cuda else torch.LongTensor
ByteTensor = torch.cuda.ByteTensor if use_cuda else torch.ByteTensor
Tensor = FloatTensor
Episode = namedtuple('Episode',('s','a','r', 'value', 'logprob'))

class EpisodeBuffer(object):
    def __init__(self):
        self.data = []

    def Push(self, *_args):
        self.data.append(Episode(*_args))

    def Pop(self):
        self.data.pop()

    def GetData(self):
        return self.data

MuscleTransition = namedtuple('MuscleTransition',('JtA','tau_des','L','b'))

class MuscleBuffer(object):
    def __init__(self, buff_size=10000):
        super(MuscleBuffer, self).__init__()
        self.buffer = deque(maxlen=buff_size)

    def Push(self, *_args):
        self.buffer.append(MuscleTransition(*_args))

    def Clear(self):
        self.buffer.clear()

Transition = namedtuple('Transition',('s','a', 'logprob', 'TD', 'GAE'))

class ReplayBuffer(object):
    def __init__(self, buff_size=10000):
        super(ReplayBuffer, self).__init__()
        self.buffer = deque(maxlen=buff_size)

    def Push(self,*args):
        self.buffer.append(Transition(*args))

    def Clear(self):
        self.buffer.clear()

MarginalTransition = namedtuple('MargianlTransition', ('sb', 'v'))

class MargianlBuffer(object):
    def __init__(self, buff_size=10000):
        super(MargianlBuffer, self).__init__()
        self.buffer = deque(maxlen=buff_size)

    def Push(self, *_args):
        self.buffer.append(MarginalTransition(*_args))

    def Clear(self):
        self.buffer.clear()

class PPO(object):
    def __init__(self, meta_file, num_slaves=16):
        # plt.ion()
        np.random.seed(seed=int(time.time()))
        self.num_slaves = num_slaves
        self.env = EnvManager(meta_file, self.num_slaves)
        self.use_muscle = self.env.UseMuscle()
        self.num_state = self.env.GetNumState()
        self.num_action = self.env.GetNumAction()
        self.num_muscles = self.env.GetNumMuscles()

        self.num_epochs = 10
        self.num_epochs_muscle = 3
        self.num_evaluation = 0
        self.num_tuple_so_far = 0
        self.num_episode = 0
        self.num_tuple = 0
        self.num_simulation_Hz = self.env.GetSimulationHz()
        self.num_control_Hz = self.env.GetControlHz()
        self.num_simulation_per_control = self.num_simulation_Hz // self.num_control_Hz

        self.gamma = 0.99
        self.lb = 0.99

        self.buffer_size = 2048
        self.batch_size = 128
        self.muscle_batch_size = 128
        self.replay_buffer = ReplayBuffer(30000)
        self.muscle_buffer = MuscleBuffer(30000)

        self.model = SimulationNN(self.num_state,self.num_action)

        self.muscle_model = MuscleNN(self.env.GetNumTotalMuscleRelatedDofs(),self.num_action,self.num_muscles)

        if use_cuda:
            self.model.cuda()
            self.muscle_model.cuda()

        self.default_learning_rate = 1E-4
        self.default_clip_ratio = 0.2
        self.learning_rate = self.default_learning_rate
        self.clip_ratio = self.default_clip_ratio
        self.optimizer = optim.Adam(self.model.parameters(),lr=self.learning_rate)
        self.optimizer_muscle = optim.Adam(self.muscle_model.parameters(),lr=self.learning_rate)
        self.max_iteration = 50000

        self.w_entropy = -0.001

        self.loss_actor = 0.0
        self.loss_critic = 0.0
        self.loss_muscle = 0.0
        self.rewards = []
        self.sum_return = 0.0
        self.max_return = -1.0
        self.max_return_epoch = 1
        self.tic = time.time()

        # for adaptive sampling, marginal value training
        self.use_adaptive_sampling = self.env.UseAdaptiveSampling()
        self.marginal_state_num = self.env.GetMarginalStateNum()
        self.marginal_buffer = MargianlBuffer(30000)
        self.marginal_model = MarginalNN(self.marginal_state_num)
        self.marginal_value_avg = 1.
        self.marginal_learning_rate = 1e-3
        if use_cuda:
            self.marginal_model.cuda()
        self.marginal_optimizer = optim.SGD(self.marginal_model.parameters(), lr=self.marginal_learning_rate)
        self.marginal_loss = 0.0
        self.marginal_samples = []
        self.marginal_sample_num = 2000
        self.marginal_k = self.env.GetMarginalParameter()
        self.mcmc_burn_in = 1000
        self.mcmc_period = 20

        self.total_episodes = []
        self.episodes = [None]*self.num_slaves
        for j in range(self.num_slaves):
            self.episodes[j] = EpisodeBuffer()
        self.env.Resets(True)

    def SaveModel(self):
        self.model.save('../nn/current.pt')
        self.muscle_model.save('../nn/current_muscle.pt')

        if self.max_return_epoch == self.num_evaluation:
            self.model.save('../nn/max.pt')
            self.muscle_model.save('../nn/max_muscle.pt')
        if self.num_evaluation % 100 == 0:
            self.model.save('../nn/'+str(self.num_evaluation//100)+'.pt')
            self.muscle_model.save('../nn/'+str(self.num_evaluation//100)+'_muscle.pt')

    def LoadModel(self, path):
        self.model.load('../nn/'+path+'.pt')
        self.muscle_model.load('../nn/'+path+'_muscle.pt')

    def ComputeTDandGAE(self):
        self.replay_buffer.Clear()
        self.muscle_buffer.Clear()
        self.marginal_buffer.Clear()
        self.sum_return = 0.0
        for epi in self.total_episodes:
            data = epi.GetData()
            size = len(data)
            if size == 0:
                continue
            states, actions, rewards, values, logprobs = zip(*data)

            values = np.concatenate((values, np.zeros(1)), axis=0)
            advantages = np.zeros(size)
            ad_t = 0

            epi_return = 0.0
            for i in reversed(range(len(data))):
                epi_return += rewards[i]
                delta = rewards[i] + values[i+1] * self.gamma - values[i]
                ad_t = delta + self.gamma * self.lb * ad_t
                advantages[i] = ad_t
            self.sum_return += epi_return
            TD = values[:size] + advantages

            for i in range(size):
                self.replay_buffer.Push(states[i], actions[i], logprobs[i], TD[i], advantages[i])

            if self.use_adaptive_sampling:
                for i in range(size):
                    self.marginal_buffer.Push(states[i][-self.marginal_state_num:], values[i])

        self.num_episode = len(self.total_episodes)
        self.num_tuple = len(self.replay_buffer.buffer)
        # print('SIM : {}'.format(self.num_tuple))
        self.num_tuple_so_far += self.num_tuple

        muscle_tuples = self.env.GetMuscleTuples()
        for i in range(len(muscle_tuples)):
            self.muscle_buffer.Push(muscle_tuples[i][0],muscle_tuples[i][1],muscle_tuples[i][2],muscle_tuples[i][3])

    def SampleStatesForMarginal(self):
        # MCMC : Metropolitan-Hastings
        _marginal_samples = []
        marginal_sample_prob = []
        marginal_sample_cumulative_prob = []
        p_sb = 0.
        mcmc_idx = 0
        while len(_marginal_samples) < self.marginal_sample_num:
            # Generation
            state_sb_prime = self.env.SampleMarginalState()

            # Evaluation
            marginal_value = self.marginal_model(Tensor(state_sb_prime)).cpu().detach().numpy().reshape(-1)
            # print(marginal_value, state_sb_prime)
            p_sb_prime = math.exp(self.marginal_k * (1. - marginal_value/self.marginal_value_avg) )

            # Rejection
            if p_sb_prime > np.random.rand() * p_sb:
                if mcmc_idx > self.mcmc_burn_in:
                    _marginal_samples.append(state_sb_prime)
                    marginal_sample_prob.append(p_sb_prime)
                p_sb = p_sb_prime
                mcmc_idx += 1

        sorted_y_idx_list = sorted(range(len(marginal_sample_prob)), key=lambda x: marginal_sample_prob[x])
        marginal_samples = [_marginal_samples[i] for i in sorted_y_idx_list]
        marginal_sample_prob.sort()

        marginal_sample_cumulative_prob.append(marginal_sample_prob[0])

        for i in range(1, len(marginal_sample_prob)):
            marginal_sample_cumulative_prob.append(marginal_sample_prob[i] + marginal_sample_cumulative_prob[-1])

        for i in range(len(marginal_sample_cumulative_prob)):
            marginal_sample_cumulative_prob[i] = marginal_sample_cumulative_prob[i]/marginal_sample_cumulative_prob[-1]

        # print(self.marginal_value_avg, sum(marginal_sample_cumulative_prob))
        # plt.figure(0)
        # plt.clf()
        # stride_idx = len(marginal_samples[0])-2
        # speed_idx = len(marginal_samples[0])-1
        # xx = []
        # yy = []
        #
        # for marginal_sample in marginal_samples:
        #     marginal_sample_exact = marginal_sample.copy()
        #     marginal_sample_exact[stride_idx] *= math.sqrt(0.00323409929)
        #     marginal_sample_exact[speed_idx] *= math.sqrt(0.00692930964)
        #     marginal_sample_exact[stride_idx] += 1.12620703
        #     marginal_sample_exact[speed_idx] += 0.994335964
        #
        #     xx.append(marginal_sample_exact[stride_idx])
        #     yy.append(marginal_sample_exact[speed_idx])
        #
        # plt.scatter(xx, yy)
        #
        # # plt.xlim(left=-3., right=3.)
        # # plt.ylim(bottom=-3., top=3.)
        # plt.xlim(left=0., right=2.)
        # plt.ylim(bottom=0., top=2.)
        # plt.show()
        # plt.pause(0.001)

        self.env.SetMarginalSampled(np.asarray(marginal_samples), marginal_sample_cumulative_prob)

    def GenerateTransitions(self):
        self.total_episodes = []
        states = [None]*self.num_slaves
        actions = [None]*self.num_slaves
        rewards = [None]*self.num_slaves
        states_next = [None]*self.num_slaves
        states = self.env.GetStates()
        local_step = 0
        terminated = [False]*self.num_slaves
        counter = 0
        while True:
            counter += 1
            # if counter % 10 == 0:
            #     print('SIM : {}'.format(local_step),end='\r')
            a_dist,v = self.model(Tensor(states))
            actions = a_dist.sample().cpu().detach().numpy()
            # actions = a_dist.loc.cpu().detach().numpy()
            logprobs = a_dist.log_prob(Tensor(actions)).cpu().detach().numpy().reshape(-1)
            values = v.cpu().detach().numpy().reshape(-1)
            self.env.SetActions(actions)
            if self.use_muscle:
                mt = Tensor(self.env.GetMuscleTorques())
                for _ in range(self.num_simulation_per_control//2):
                    dt = Tensor(self.env.GetDesiredTorques())
                    activations = self.muscle_model(mt,dt).cpu().detach().numpy()
                    self.env.SetActivationLevels(activations)

                    self.env.Steps(2)
            else:
                self.env.StepsAtOnce()

            for j in range(self.num_slaves):
                nan_occur = False
                terminated_state = True

                if np.any(np.isnan(states[j])) or np.any(np.isnan(actions[j])) or np.any(np.isnan(states[j])) or np.any(np.isnan(values[j])) or np.any(np.isnan(logprobs[j])):
                    nan_occur = True

                elif self.env.IsEndOfEpisode(j) is False:
                    terminated_state = False
                    rewards[j] = self.env.GetReward(j)
                    self.episodes[j].Push(states[j], actions[j], rewards[j], values[j], logprobs[j])
                    local_step += 1

                if terminated_state or nan_occur:
                    if nan_occur:
                        self.episodes[j].Pop()
                    self.total_episodes.append(self.episodes[j])
                    self.episodes[j] = EpisodeBuffer()

                    self.env.Reset(True, j)

            if local_step >= self.buffer_size:
                break

            states = self.env.GetStates()

    def OptimizeSimulationNN(self):
        all_transitions = np.array(self.replay_buffer.buffer)
        for j in range(self.num_epochs):
            np.random.shuffle(all_transitions)
            for i in range(len(all_transitions)//self.batch_size):
                transitions = all_transitions[i*self.batch_size:(i+1)*self.batch_size]
                batch = Transition(*zip(*transitions))

                stack_s = np.vstack(batch.s).astype(np.float32)
                stack_a = np.vstack(batch.a).astype(np.float32)
                stack_lp = np.vstack(batch.logprob).astype(np.float32)
                stack_td = np.vstack(batch.TD).astype(np.float32)
                stack_gae = np.vstack(batch.GAE).astype(np.float32)

                a_dist,v = self.model(Tensor(stack_s))
                '''Critic Loss'''
                loss_critic = ((v-Tensor(stack_td)).pow(2)).mean()

                '''Actor Loss'''
                ratio = torch.exp(a_dist.log_prob(Tensor(stack_a))-Tensor(stack_lp))
                stack_gae = (stack_gae-stack_gae.mean())/(stack_gae.std()+ 1E-5)
                stack_gae = Tensor(stack_gae)
                surrogate1 = ratio * stack_gae
                surrogate2 = torch.clamp(ratio,min =1.0-self.clip_ratio,max=1.0+self.clip_ratio) * stack_gae
                loss_actor = - torch.min(surrogate1,surrogate2).mean()
                '''Entropy Loss'''
                loss_entropy = - self.w_entropy * a_dist.entropy().mean()

                self.loss_actor = loss_actor.cpu().detach().numpy().tolist()
                self.loss_critic = loss_critic.cpu().detach().numpy().tolist()

                loss = loss_actor + loss_entropy + loss_critic

                self.optimizer.zero_grad()
                loss.backward(retain_graph=True)
                for param in self.model.parameters():
                    if param.grad is not None:
                        param.grad.data.clamp_(-0.5,0.5)
                self.optimizer.step()
            # print('Optimizing sim nn : {}/{}'.format(j+1,self.num_epochs),end='\r')
        # print('')

    def OptimizeMuscleNN(self):
        muscle_transitions = np.array(self.muscle_buffer.buffer)
        for j in range(self.num_epochs_muscle):
            np.random.shuffle(muscle_transitions)
            for i in range(len(muscle_transitions)//self.muscle_batch_size):
                tuples = muscle_transitions[i*self.muscle_batch_size:(i+1)*self.muscle_batch_size]
                batch = MuscleTransition(*zip(*tuples))

                stack_JtA = np.vstack(batch.JtA).astype(np.float32)
                stack_tau_des = np.vstack(batch.tau_des).astype(np.float32)
                stack_L = np.vstack(batch.L).astype(np.float32)

                stack_L = stack_L.reshape(self.muscle_batch_size,self.num_action,self.num_muscles)
                stack_b = np.vstack(batch.b).astype(np.float32)

                stack_JtA = Tensor(stack_JtA)
                stack_tau_des = Tensor(stack_tau_des)
                stack_L = Tensor(stack_L)
                stack_b = Tensor(stack_b)

                activation = self.muscle_model(stack_JtA,stack_tau_des)
                tau = torch.einsum('ijk,ik->ij',(stack_L,activation)) + stack_b

                loss_reg = (activation).pow(2).mean()
                loss_target = (((tau-stack_tau_des)/100.0).pow(2)).mean()

                loss = 0.01*loss_reg + loss_target
                # loss = loss_target

                self.optimizer_muscle.zero_grad()
                loss.backward(retain_graph=True)
                for param in self.muscle_model.parameters():
                    if param.grad is not None:
                        param.grad.data.clamp_(-0.5,0.5)
                self.optimizer_muscle.step()

            # print('Optimizing muscle nn : {}/{}'.format(j+1,self.num_epochs_muscle),end='\r')
        self.loss_muscle = loss.cpu().detach().numpy().tolist()
        # print('')

    def OptimizeMarginalNN(self):
        marginal_transitions = np.array(self.marginal_buffer.buffer)
        for j in range(self.num_epochs):
            np.random.shuffle(marginal_transitions)
            for i in range(len(marginal_transitions)//self.batch_size):
                transitions = marginal_transitions[i*self.batch_size:(i+1)*self.batch_size]
                batch = MarginalTransition(*zip(*transitions))

                stack_sb = np.vstack(batch.sb).astype(np.float32)
                stack_v = np.vstack(batch.v).astype(np.float32)

                v = self.marginal_model(Tensor(stack_sb))

                # Marginal Loss
                loss_marginal = ((v-Tensor(stack_v)).pow(2)).mean()
                self.marginal_loss = loss_marginal.cpu().detach().numpy().tolist()
                self.marginal_optimizer.zero_grad()
                loss_marginal.backward(retain_graph=True)

                for param in self.marginal_model.parameters():
                    if param.grad is not None:
                        param.grad.data.clamp_(-0.5, 0.5)
                self.marginal_optimizer.step()

                # Marginal value average
                avg_marginal = Tensor(stack_v).mean().cpu().detach().numpy().tolist()
                self.marginal_value_avg -= self.marginal_learning_rate * (self.marginal_value_avg - avg_marginal)

            # print('Optimizing margin nn : {}/{}'.format(j+1, self.num_epochs), end='\r')
        # print('')

    def OptimizeModel(self):
        self.ComputeTDandGAE()
        self.OptimizeSimulationNN()
        if self.use_muscle:
            self.OptimizeMuscleNN()
        if self.use_adaptive_sampling:
            self.OptimizeMarginalNN()

    def Train(self, idx):
        if self.use_adaptive_sampling and (idx % self.mcmc_period == 0):
            self.SampleStatesForMarginal()
        self.GenerateTransitions()
        self.OptimizeModel()

    def Evaluate(self):
        self.num_evaluation = self.num_evaluation + 1
        h = int((time.time() - self.tic)//3600.0)
        m = int((time.time() - self.tic)//60.0)
        s = int((time.time() - self.tic))
        s = s - m*60
        m = m - h*60
        if self.num_episode is 0:
            self.num_episode = 1
        if self.num_tuple is 0:
            self.num_tuple = 1
        if self.max_return < self.sum_return/self.num_episode:
            self.max_return = self.sum_return/self.num_episode
            self.max_return_epoch = self.num_evaluation
        with open('../nn/log.txt', 'a') as f:
            f.write('# {} === {}h:{}m:{}s ===\n'.format(self.num_evaluation,h,m,s))
            f.write('||Loss Actor               : {:.4f}\n'.format(self.loss_actor))
            f.write('||Loss Critic              : {:.4f}\n'.format(self.loss_critic))
            if self.use_muscle:
                f.write('||Loss Muscle              : {:.4f}\n'.format(self.loss_muscle))
            if self.use_adaptive_sampling:
                f.write('||Loss Marginal            : {:.4f}\n'.format(self.marginal_loss))
            f.write('||Noise                    : {:.3f}\n'.format(self.model.log_std.exp().mean()))
            f.write('||Num Transition So far    : {}\n'.format(self.num_tuple_so_far))
            f.write('||Num Transition           : {}\n'.format(self.num_tuple))
            f.write('||Num Episode              : {}\n'.format(self.num_episode))
            f.write('||Avg Return per episode   : {:.3f}\n'.format(self.sum_return/self.num_episode))
            f.write('||Avg Reward per transition: {:.3f}\n'.format(self.sum_return/self.num_tuple))
            f.write('||Avg Step per episode     : {:.1f}\n'.format(self.num_tuple/self.num_episode))
            f.write('||Max Avg Retun So far     : {:.3f} at #{}\n'.format(self.max_return,self.max_return_epoch))
            f.write('=============================================\n')
        self.rewards.append(self.sum_return/self.num_episode)

        self.SaveModel()

        return np.array(self.rewards)

plt.ion()

def Plot(y,title,num_fig=1,ylim=True):
    temp_y = np.zeros(y.shape)
    if y.shape[0]>5:
        temp_y[0] = y[0]
        temp_y[1] = 0.5*(y[0] + y[1])
        temp_y[2] = 0.3333*(y[0] + y[1] + y[2])
        temp_y[3] = 0.25*(y[0] + y[1] + y[2] + y[3])
        for i in range(4,y.shape[0]):
            temp_y[i] = np.sum(y[i-4:i+1])*0.2

    plt.figure(num_fig)
    plt.clf()
    plt.title(title)
    plt.plot(y,'b')

    plt.plot(temp_y,'r')

    plt.show()
    if ylim:
        plt.ylim([0,1])
    plt.pause(0.001)

import argparse
import os
if __name__=="__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-m', '--model', help='model path')
    parser.add_argument('-d', '--meta', help='meta file')
    parser.add_argument('-p', '--parallel', help='num slaves')

    args = parser.parse_args()
    if args.meta is None:
        print('Provide meta file')
        exit()

    if args.parallel is None:
        ppo = PPO(args.meta)
    else:
        ppo = PPO(args.meta, int(args.parallel))
    nn_dir = '../nn'
    if not os.path.exists(nn_dir):
        os.makedirs(nn_dir)
    if args.model is not None:
        ppo.LoadModel(args.model)
    else:
        ppo.SaveModel()
    print('num states: {}, num actions: {}'.format(ppo.env.GetNumState(),ppo.env.GetNumAction()))
    for _i in range(ppo.max_iteration-5):
        ppo.Train(_i)
        rewards = ppo.Evaluate()
    # Plot(rewards,'reward',0,False)


import numpy as np
from scipy import spatial

from datetime import datetime

max_doc_len = 500

time_start = datetime.now()
print("Starting time is " + str(time_start))
glove_embeddings = {}
embeds_file = open('glove/simple.txt', 'r')
#embeds_file = open('glove/glove.840B.300d.txt', 'r')
for line in embeds_file:
    try:
        splitLines = line.split()
        word = splitLines[0]
        vector = np.array([float(value) for value in splitLines[1:]])
        glove_embeddings[word] = vector
    except:
        continue
#print(len(glove_embeddings))
time_glove = datetime.now()
print("Glove completed " + str(time_glove) + " required " + str(time_glove - time_start))

def get_embedding(text):
    words = text.split(' ')
    count_words = 0
    text_vector = np.zeros((300), dtype='float32')
    mod_text = ""
    for word in words:
        if word in glove_embeddings:
            text_vector += glove_embeddings[word]
            count_words += 1
            mod_text += (word + ' ')
    if count_words > 0:
        text_vector /= count_words
    return [text_vector, mod_text]

queries = []
#queries_file = open('collection_queries/queries.dev.small.tsv', 'r')
queries_file = open('collection_queries/small_queries_test.tsv', 'r')
q_line = queries_file.readlines()[:10]
for q in q_line:
    queries.append((q.split('\t')[1]).rstrip())
#print(queries)

docs = []
docs_file = open('corpus/small.tsv', 'r')
#docs_file = open('corpus/msmarco-docs.tsv', 'r')
docs_lines = docs_file.readlines()[:100]
for doc in docs_lines:
    docs.append((doc.split('\t')[3]).rstrip())
#print(doc)

document_embeddings = []
document_lengths = []
modified_documents = []
for (i, doc) in enumerate(docs):
    sentences = doc.split('.')
    curr_doc = []
    curr_lengths = []
    mod_doc = []
    for sentence in sentences:
        embed, length = get_embedding(sentence)
        curr_doc.append(embed)
        curr_lengths.append(len(length))
        mod_doc.append(length)
    #print(curr_lengths)
    document_embeddings.append(curr_doc)
    document_lengths.append(curr_lengths)
    modified_documents.append(mod_doc)
#print(document_lengths)
time_docs = datetime.now()
time_prev = datetime.now()
print("Document embeddings completed " + str(time_docs) + " required " + str(time_docs - time_glove))

#queries_vector = []
for (i, query) in enumerate(queries):
    query_vector, mod_query = get_embedding(query)
    #print(query_vector.shape)
    ids = i
    fw = open('docs/' + str(ids) + '.txt', 'w')

    for (j, doc) in enumerate(docs):
        sentences = doc.split('.')
        sentences_scores = [[0]*3 for _ in range(len(sentences))] # creating one for index, one for length and one for relevance score
        for (k, sentence) in enumerate(sentences):
            #sen_embed = get_embedding(sentence)
            #print(sen_embed.shape)
            sentences_scores[k][0] = k
            #sentences_scores[k][1] = len(sentence.split(' '))
            sentences_scores[k][1] = document_lengths[j][k]
            sentences_scores[k][2] = 1 - spatial.distance.cosine(query_vector, document_embeddings[j][k])
        sentences_scores = sorted(sentences_scores, key=lambda x: -x[2])
        #print("query num " + str(i) + " having sentence scores as  " + str(sentences_scores))
        #print("\n\n\n\n\n")
        final_doc = ""
        final_doc += mod_query
        new_doc_len = len(mod_query.split(' '))
        idx = 0
        while idx < len(sentences) and (new_doc_len + sentences_scores[idx][1]) < 512:
            final_doc += sentences[sentences_scores[idx][0]]
            new_doc_len += sentences_scores[idx][1]
            idx += 1
        #print(final_doc)
        fw.write(final_doc + '\n')
    time_query = datetime.now()
    print("query " + str(i) + " completed by " + str(time_query) + " took " + str(time_query - time_prev))
    time_prev = time_query

time_final = datetime.now()
print("final timing " + str(time_final))
#print(queries_vector)