# -*- coding: utf-8 -*-
"""
Tencent is pleased to support the open source community by making BK-BASE 蓝鲸基础平台 available.
Copyright (C) 2021 THL A29 Limited, a Tencent company.  All rights reserved.
BK-BASE 蓝鲸基础平台 is licensed under the MIT License.
License for BK-BASE 蓝鲸基础平台:
--------------------------------------------------------------------
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
documentation files (the "Software"), to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software,
and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial
portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""

import gc
import logging
from collections import OrderedDict
from enum import Enum
from itertools import groupby
from typing import Any, List, Set, Union

import attr
import cattr
import jieba
import jsonext
import numpy as np
import requests
from cached_property import threaded_cached_property
from jinja2 import Template
from networkx import DiGraph
from networkx import faster_could_be_isomorphic as is_isomorphic
from scipy.sparse import vstack
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn.preprocessing import normalize

from metadata.exc import CompareError
from metadata.runtime import rt_context, rt_g
from metadata.util.i18n import lazy_selfish as _
from metadata_pro.analyse import AnalyseComponent, AnalyseFlow, AnalyseProcedure

logger = logging.getLogger(__name__)

class VectorizerType(Enum):
    COUNT = 'count'
    TF_IDF = 'tf_idf'

class SegmentMode(Enum):
    CHINESE = 'chinese'
    UNDER_SCORE = 'under_score'

class AnalyseMode(Enum):
    CACHED = 'cached'
    UP_TO_DATE = 'up_to_date'

@attr.s(frozen=True)
class RTField(object):
    field_name = attr.ib(type=str)
    field_type = attr.ib(type=str)

@attr.s
class RT(object):
    result_table_id = attr.ib(type=str)
    fields = attr.ib(type=Set[RTField])
    result_table_name = attr.ib(type=str)
    description = attr.ib(type=str)

def compare_content_len(instance, attribute, value):
    if not value:
        raise CompareError(_('The compare content is blank.'))

@attr.s
class CompareContext(object):
    target = attr.ib(type=Union[Set[Any], List[Any]], validator=compare_content_len)
    reference = attr.ib(type=Union[Set[Any], List[Any]], validator=compare_content_len)

def word_segmentation(sentence, segment_mode):
    if not sentence:
        return ['']
    if segment_mode is SegmentMode.CHINESE:
        return jieba.cut(sentence)
    elif segment_mode is SegmentMode.UNDER_SCORE:
        return sentence.split('_')

class BagOfWord(AnalyseComponent):
    """
    词袋。支持count，tf_idf两种词集合向量生成方式。
    """

    def __init__(self, mode=VectorizerType.COUNT):
        super(BagOfWord, self).__init__()
        self.mode = mode

    @threaded_cached_property
    def vectorizer(self):
        return CountVectorizer() if self.mode is VectorizerType.COUNT else TfidfVectorizer()

    def build(self, documents):
        return self.vectorizer.fit_transform(documents)

    def transform(self, documents):
        return self.vectorizer.transform(documents)

    @staticmethod
    def load_corpus():
        documents = []
        ret_rt = rt_context.layers.asset.complex_search(
            statement="select result_table_id,result_table_name,description from result_table", backend_type='mysql'
        )
        ret_rt_field = rt_context.layers.asset.complex_search(
            statement="select result_table_id,field_name,field_alias,field_type from result_table_field",
            backend_type='mysql',
        )
        for item in ret_rt:
            documents.append(
                ' '.join(word_segmentation(item['result_table_name'], segment_mode=SegmentMode.UNDER_SCORE))
            )
            if item['description']:
                documents.append(' '.join(word_segmentation(item['description'], segment_mode=SegmentMode.CHINESE)))
        for item in ret_rt_field:
            documents.append(item['field_name'])
        return documents, ret_rt, ret_rt_field

def renew_context():
    # 生成后台词袋
    bow = BagOfWord(VectorizerType.TF_IDF)
    documents, ret_rt, ret_rt_field = bow.load_corpus()
    bow.build(documents)
    rt_g.bow_now = bow
    logger.info('Bow is generated.')

    # 生成既有RT向量
    rt_info = OrderedDict()
    for item in ret_rt:
        rt_info[item.pop('result_table_id')] = item
    for k, g in groupby(ret_rt_field, key=lambda item: item['result_table_id']):
        if k in rt_info:
            rt_info[k]['fields'] = {item['field_name']: {'field_alias': item['field_alias']} for item in g}
    ss = SchemaSimilarity(bow)
    schema_vectors = ss.transform_input(item.get('fields', []) for item in rt_info.values())
    result_table_name_vectors = DescriptionSimilarity(bow).transform_input(
        item.pop('result_table_name') for item in rt_info.values()
    )
    description_vectors = DescriptionSimilarity(bow, SegmentMode.CHINESE).transform_input(
        item.pop('description') for item in rt_info.values()
    )
    rt_vectors = OrderedDict()
    for n, k in enumerate(rt_info):
        rt_vectors[k] = {
            'schema': schema_vectors[n],
            'result_table_name': result_table_name_vectors[n],
            'description': description_vectors[n],
        }
    rt_g.rt_vectors_now = rt_vectors
    rt_g.rt_ids_now = np.char.asarray(list(rt_vectors.keys()))
    rt_g.rt_info_now = rt_info
    logger.info('Existed RT vectors are generated.')

    gc.collect()

class WordsSetSimilarity(AnalyseProcedure):
    def __init__(self, bow):
        super(WordsSetSimilarity, self).__init__()
        self.bow = bow

    def calc_cos_distance(self, matrix_target, matrix_reference):
        matrix_x = normalize(matrix_target) if self.bow.mode is VectorizerType.COUNT else matrix_target
        matrix_y = matrix_reference.transpose()
        return np.dot(matrix_x, matrix_y)

    def transform_input(self, _input):
        return self.bow.transform(' '.join(item) for item in _input)

    def execute(self, _input=None, *args, **kwargs):
        super(WordsSetSimilarity, self).execute(_input)
        matrix_target = self.transform_input(self.input.target)
        matrix_reference = self.transform_input(self.input.reference)
        self.output = self.calc_cos_distance(matrix_target, matrix_reference)
        return self.output

class SchemaSimilarity(WordsSetSimilarity):
    pass

class DescriptionSimilarity(WordsSetSimilarity):
    def __init__(self, bow, segment_mode=SegmentMode.UNDER_SCORE):
        super(DescriptionSimilarity, self).__init__(bow)
        self.segment_mode = segment_mode

    def transform_input(self, _input):
        words_sets = [word_segmentation(sentence, self.segment_mode) for sentence in _input]
        return super(DescriptionSimilarity, self).transform_input(words_sets)

class ResultTableSimilarity(AnalyseProcedure):
    def __init__(self, bow=None):
        super(ResultTableSimilarity, self).__init__()
        if not bow:
            self.bow = rt_context.bow_now
        self.schema_matrix, self.result_table_name_matrix, self.description_matrix = None, None, None

    def calc_content_similarity(self, **kwargs):
        ws = WordsSetSimilarity(self.bow)
        ws.input = CompareContext(
            [[item.field_name for item in rt.fields] for rt in self.input.target],
            [[item.field_name for item in rt.fields] for rt in self.input.reference],
        )
        self.schema_matrix = ws.execute()
        ds = DescriptionSimilarity(
            self.bow,
        )
        ds.input = CompareContext(
            [rt.result_table_name for rt in self.input.target], [rt.result_table_name for rt in self.input.reference]
        )
        self.result_table_name_matrix = ds.execute()
        ds_2 = DescriptionSimilarity(
            self.bow,
            segment_mode=SegmentMode.CHINESE,
        )
        ds_2.input = CompareContext(
            {rt.description for rt in self.input.target}, {rt.description for rt in self.input.reference}
        )
        self.description_matrix = ds_2.execute()
        return self.schema_matrix * 0.8 + self.result_table_name_matrix * 0.1 + self.description_matrix * 0.1

    def execute(self, *args, **kwargs):
        args, kwargs = super(ResultTableSimilarity, self).execute(*args, **kwargs)
        self.output = self.calc_content_similarity(**kwargs)
        return self.output

class CachedResultTableSimilarity(ResultTableSimilarity):
    def __init__(self, *args, **kwargs):
        super(CachedResultTableSimilarity, self).__init__(*args, **kwargs)
        self.ws = WordsSetSimilarity(self.bow)

    def calc_content_similarity(self, two_way_exchange=False):
        target, reference = self.input.target, self.input.reference
        if two_way_exchange:
            target, reference = reference, target
        self.schema_matrix = self.ws.calc_cos_distance(
            vstack(item['schema'] for item in target), vstack(item['schema'] for item in reference)
        )
        self.result_table_name_matrix = self.ws.calc_cos_distance(
            vstack(item['result_table_name'] for item in target),
            vstack(item['result_table_name'] for item in reference),
        )
        self.description_matrix = self.ws.calc_cos_distance(
            vstack(item['description'] for item in target), vstack(item['description'] for item in reference)
        )
        if two_way_exchange:
            self.schema_matrix = self.schema_matrix.transpose()
            self.result_table_name_matrix = self.result_table_name_matrix.transpose()
            self.description_matrix = self.description_matrix.transpose()
        return self.schema_matrix * 0.8 + self.result_table_name_matrix * 0.1 + self.description_matrix * 0.1

class ResultTableSimilarityAnalyseFlow(AnalyseFlow):
    content_access_graphql = Template(
        "{rt_metadata(func:eq(ResultTable.result_table_id,{{ rt_ids }}))"
        "{result_table_id:ResultTable.result_table_id result_table_name:ResultTable.result_table_name  "
        "description: ResultTable.description "
        "fields:~ResultTableField.result_table "
        "{field_name:ResultTableField.field_name field_type:ResultTableField.field_type}}}"
    )

    def __init__(self):
        super(ResultTableSimilarityAnalyseFlow, self).__init__()

    def load_rt_metadata(self):
        all_rt_ids = set(self.input.target) | set(self.input.reference)
        query_statement = self.content_access_graphql.render(rt_ids=jsonext.dumps(all_rt_ids))
        ret_rt = rt_context.layers.asset.complex_search(statement=query_statement, backend_type='dgraph')
        ret_rt_dct = {item['result_table_id']: item for item in ret_rt['data']['rt_metadata']}
        not_existed_rt_id = set(all_rt_ids) - set(ret_rt_dct)
        if not_existed_rt_id:
            raise CompareError(_('Some rt id {} is not existed.').format(not_existed_rt_id))
        target_rt_metadata = [cattr.structure(ret_rt_dct[rt_name], RT) for rt_name in self.input.target]
        reference_rt_metadata = [cattr.structure(ret_rt_dct[rt_name], RT) for rt_name in self.input.reference]
        return {'target': target_rt_metadata, 'reference': reference_rt_metadata}

    def execute(self, *args, **kwargs):
        super(ResultTableSimilarityAnalyseFlow, self).execute(*args, **kwargs)
        rt_metadata = self.load_rt_metadata()
        a = ResultTableSimilarity()
        a.execute(CompareContext(**rt_metadata))
        self.output = {
            'final_goal': a.output.toarray(),
            'detail': {
                'schema': a.schema_matrix.toarray(),
                'name': a.result_table_name_matrix.toarray(),
                'description': a.description_matrix.toarray(),
            },
        }
        return self.output

class CachedResultTableSimilarityAnalyseFlow(AnalyseFlow):
    def __init__(self):
        super(CachedResultTableSimilarityAnalyseFlow, self).__init__()

    def load_rt_metadata(self):
        if set(self.input.target) == set(self.input.reference) == {'*'}:
            raise CompareError(_("Can't compare all items together."))
        for rt_id in set(self.input.target) | set(self.input.reference):
            if rt_id not in rt_context.rt_vectors_now and rt_id != "*":
                raise CompareError('RT id {} is not in cache.'.format(rt_id))
        target = (
            [rt_context.rt_vectors_now[rt_id] for rt_id in self.input.target]
            if set(self.input.target) != {'*'}
            else list(rt_context.rt_vectors_now.values())
        )
        reference = (
            [rt_context.rt_vectors_now[rt_id] for rt_id in self.input.reference]
            if set(self.input.reference) != {'*'}
            else list(rt_context.rt_vectors_now.values())
        )
        two_way_exchange = True if set(self.input.reference) == {'*'} else False
        return target, reference, two_way_exchange

    def execute(self, _input=None, reference_max_n=None, *args, **kwargs):
        super(CachedResultTableSimilarityAnalyseFlow, self).execute(_input)
        target, reference, two_way_exchange = self.load_rt_metadata()
        a = CachedResultTableSimilarity()
        a.execute(CompareContext(target, reference), two_way_exchange=two_way_exchange)
        ret = {
            'final_goal': a.output.toarray(),
            'detail': {
                'schema': a.schema_matrix.toarray(),
                'name': a.result_table_name_matrix.toarray(),
                'description': a.description_matrix.toarray(),
            },
        }
        if reference_max_n:
            self.gather_reference_top(reference_max_n, ret)
        self.output = ret
        return self.output

    @staticmethod
    def gather_reference_top(reference_max_n, ret):
        indexes = np.unravel_index(np.argsort(ret['final_goal'], axis=1)[:, -reference_max_n:], ret['final_goal'].shape)
        ret['final_goal'] = ret['final_goal'][indexes]
        for k, v in ret['detail'].items():
            ret['detail'][k] = v[indexes]

        ret['reference_result_table_ids'] = [rt_context.rt_ids_now[i] for i in indexes[1]]

class FieldsSuggestionAnalyseFlow(AnalyseFlow):
    def __init__(self, bow=None):
        super(FieldsSuggestionAnalyseFlow, self).__init__()
        if not bow:
            self.bow = rt_context.bow_now
        self.ss = SchemaSimilarity(self.bow)

    def execute(self, _input=None, match_percent=0.5, *args, **kwargs):
        super(FieldsSuggestionAnalyseFlow, self).execute(_input)
        matrix_target = self.ss.transform_input(self.input)
        schema_matrix = self.ss.calc_cos_distance(
            vstack(item['schema'] for item in rt_context.rt_vectors_now.values()),
            matrix_target,
        )
        schema_matrix = schema_matrix.transpose()
        schema_similarity, reference_result_table_ids = self.gather_reference_top(
            match_percent, schema_matrix.toarray()
        )
        rts_fields_info = self.query_rt_info(reference_result_table_ids)
        self.output = {'schema_similarity': schema_similarity, 'rts_fields_info': rts_fields_info}
        return self.output

    def query_rt_info(self, reference_result_table_ids):
        info_lst = []
        for lst in reference_result_table_ids:
            info_lst.append(OrderedDict((rt_id, rt_g.rt_info_now[rt_id]) for rt_id in lst))
        return info_lst

    @staticmethod
    def gather_reference_top(match_percent, ret):
        r, c = np.where(ret >= match_percent)
        indexes = np.split(c, np.searchsorted(r, list(range(1, ret.shape[0]))))
        indexes = np.unravel_index(indexes, ret.shape)
        ret = ret[indexes]
        reference_result_table_ids = [rt_context.rt_ids_now[i] for i in indexes[1]]
        return ret, reference_result_table_ids

def is_chinese(word):
    for ch in word:
        if '\u4e00' <= ch <= '\u9fff':
            return True
    return False

class FieldAliasSuggestionAnalyseFlow(FieldsSuggestionAnalyseFlow):
    def execute(self, _input=None, *args, **kwargs):
        ret = super(FieldAliasSuggestionAnalyseFlow, self).execute(_input, *args, **kwargs)
        suggest_aliases = {}
        for dct in ret['rts_fields_info']:
            for rt_info in dct.values():
                for field_name, info in rt_info['fields'].items():
                    if is_chinese(info['field_alias']) and field_name not in suggest_aliases:
                        suggest_aliases[field_name] = info['field_alias']
        suggested_aliases_lst = [
            {field_name: suggest_aliases.get(field_name, None) for field_name in field_name_lst}
            for field_name_lst in self.input
        ]
        return suggested_aliases_lst

class ResultTableLineageSimilarityAnalyseFlow(AnalyseFlow):
    def __init__(self):
        super(ResultTableLineageSimilarityAnalyseFlow, self).__init__()
        self.normal_conf = rt_context.config_collection.normal_config

    def query_lineage(self, qualified_name):
        url = self.normal_conf.meta_api_url + 'lineage/'
        r = requests.get(url, params={'type': 'result_table', 'qualified_name': qualified_name})
        r.raise_for_status()
        info = r.json()
        if info['result'] and info['data'] and info['data']['relations']:
            return info['data']
        else:
            raise CompareError(_('Fail to get lineage.'))

    def build_lineage_graph(self, qualified_name):
        edges = [(relation['from'], relation['to']) for relation in self.query_lineage(qualified_name)['relations']]
        g = DiGraph()
        g.add_edges_from(edges)
        return g

    def execute(self, _input=None, *args, **kwargs):
        super(ResultTableLineageSimilarityAnalyseFlow, self).execute(_input)
        target, reference = self.input.target[0], self.input.reference[0]
        target_g = self.build_lineage_graph(target)
        reference_g = self.build_lineage_graph(reference)
        self.output = is_isomorphic(target_g, reference_g)
        return self.output


import numpy as np
from torch.nn import functional as F
from torch import nn
from torchvision.ops import nms
import torch

class Region_Proposal_Network(nn.Module):
    def __init__(self, in_channels, out_channels, subsampling, scales, ratios, stride):
        super(Region_Proposal_Network, self).__init__()
        # input and output channels
        self.in_channels = in_channels
        self.out_channels = out_channels
        # anchor scales and ratios, expected a list like [8,16,32] and [0.5,1,2] respectively
        self.scales = scales
        self.ratios = ratios
        # k anchor boxes
        self.k = len(scales) * len(ratios)
        # subsampling factor, VGG16: 16, 800->50
        self.subsampling = subsampling
        # stride
        self.stride = stride
        # generate k anchor boxes on the feature map (subsampled)
        # EX: [22500, 4]
        self.anchor_box = self.get_anchor_box()
        # intermediate layer
        self.intermediate = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        # reg layer outputs 4k, 4 box coordinates of k boxes
        # deault: stride=1, padding=0
        self.reg = nn.Conv2d(out_channels, 4*self.k, kernel_size=1)
        # cls layer outputs 2k scores that estimate probability of object or not object
        # deault: stride=1, padding=0
        self.cls = nn.Conv2d(out_channels, 2*self.k, kernel_size=1)
        # initialize all new layers by drawing weights from a zero-mean Gaussian distribution with standard deviation 0.01
        self.intermediate.weight.data.normal_(0, 0.01)
        self.intermediate.bias.data.zero_()
        self.reg.weight.data.normal_(0, 0.01)
        self.reg.bias.data.zero_()
        self.cls.weight.data.normal_(0, 0.01)
        self.cls.bias.data.zero_()

    def forward(self, x, image_size, scale=1.):
        # N: batch size, C: channel, H,W: height and width of input feature map
        N, C, H, W = x.shape
        # generate all anchor boxes on the original image space
        anchor = self.get_all_anchor_box(H,W)
        x = F.relu(self.intermediate(x))
        pred_anchor_locs  = self.reg(x) # [1,4*9,50,50]
        pred_cls_scores = self.cls(x) # [1,2*9,50,50]
        # reshape to fit anchor targets, [1,36,50,50] -> [1,50,50,36] -> [1,22500,4]
        pred_anchor_locs = pred_anchor_locs.permute(0, 2, 3, 1).contiguous().view(1, -1, 4)
        # reshape to fit anchor targets, [1,18,50,50] -> [1,50,50,18]
        pred_cls_scores = pred_cls_scores.permute(0, 2, 3, 1).contiguous()
        # dim=4: normalize output score
        # cls layer as a two-class softmax layer
        pred_softmax_scores = F.softmax(pred_cls_scores.view(N, H, W, self.k, 2), dim=4)
        # EX: [1, 22500]
        objectness_score = pred_softmax_scores[:, :, :, :, 1].contiguous().view(N, -1)

        # Proposal Layer
        # generate a set of proposals
        # loop through batch, feed one image at a time
        pred_boxes = []
        pred_boxes_idx = []
        for i in range(N):
            pred_box = self.generate_proposals(pred_anchor_locs[i].cpu().data.numpy(), anchor,objectness_score[i].cpu().data.numpy(), image_size)
            batch_idx = i * np.ones((len(pred_box), ), dtype=np.int32)
            pred_boxes.append(pred_box)
            pred_boxes_idx.append(batch_idx)

        pred_boxes = np.concatenate(pred_boxes, axis=0)
        pred_boxes_idx = np.concatenate(pred_boxes_idx, axis=0)

        # pred_cls_scores: [1,50,50,18] -> [1,22500,2]
        return pred_anchor_locs, pred_cls_scores.view(N, -1, 2), pred_boxes, pred_boxes_idx, anchor

    # 针对原image 800x800 全部生成的anchor box的坐标点
    # EX: [ -37.254833  -82.50967    53.254833   98.50967 ]
    # [22500,4] = 800x800/(16x16)=2500,2500*9=22500
    def get_all_anchor_box(self, H, W):
        shift_x = np.arange(0, W*self.stride, self.stride)
        shift_y = np.arange(0, H*self.stride, self.stride)
        shift_x, shift_y = np.meshgrid(shift_x, shift_y)
        # shift.shape (2500,4)
        shift = np.stack((shift_y.ravel(), shift_x.ravel(),
                          shift_y.ravel(), shift_x.ravel()), axis=1)

        A = self.anchor_box.shape[0]
        K = shift.shape[0]
        anchor_box = self.anchor_box
        anchor = anchor_box.reshape((1, A, 4)) + \
                 shift.reshape((1, K, 4)).transpose((1, 0, 2))
        anchor = anchor.reshape((K * A, 4)).astype(np.float32)

        return anchor

    # 针对原image 16X16 region生成的9个anchor box
    # EX: [ -37.254833  -82.50967    53.254833   98.50967 ]
    def get_anchor_box(self):
        anchor_box = np.zeros((len(self.scales)*len(self.ratios), 4), dtype=np.float32)
        x_center  = self.subsampling / 2.
        y_center  = self.subsampling / 2.
        for i in range(len(self.ratios)):
            for j in range(len(self.scales)):
                h = self.subsampling * self.scales[j] * np.sqrt(self.ratios[i])
                w = self.subsampling * self.scales[j] * np.sqrt(1./ self.ratios[i])
                index = i * len(self.scales) + j
                # 0,1 -> top left coordinates, 2,3->bottom right coordinates
                # [y1,x1,y2,x2]
                anchor_box[index, 0] = y_center - h / 2.
                anchor_box[index, 1] = x_center - w / 2.
                anchor_box[index, 2] = y_center + h / 2.
                anchor_box[index, 3] = x_center + w / 2.

        return anchor_box

    def generate_proposals(self, pred_anchor_locs, anchor, objectness_score, image_size):
        # IoU threshold for NMS
        nms_thresh = 0.7
        if(self.training):  # During training
            num_pre_nms = 12000 # number of anchor box before nms
            num_after_nms = 2000 # number of anchor box after nms
        else:               # During testing
            num_pre_nms = 6000 # number of anchor box before nms
            num_after_nms = 300 # number of anchor box after nms

        # minimum height of the object required to create a proposal
        min_size = 16
        # convert pred_anchor_locs to anchor_box coordinates [y1, x1, y2, x2] format
        # [1,22500,4] ->
            # compute center of the anchor box on 800x800 image space
        # and height and width of the anchor box
        anchor = anchor.astype(anchor.dtype, copy=False)
        h = anchor[:, 2] - anchor[:, 0]
        w = anchor[:, 3] - anchor[:, 1]
        y_center = anchor[:, 0] + 0.5 * h
        x_center = anchor[:, 1] + 0.5 * w

        # outputs of the reg layer
        # [:, 0::4] to preserve 2D array
        dy = pred_anchor_locs[:, 0::4]
        dx = pred_anchor_locs[:, 1::4]
        dh = pred_anchor_locs[:, 2::4]
        dw = pred_anchor_locs[:, 3::4]
        # compute new center of the predicted anchor box by reg layer
        new_y_center = dy * h[:, np.newaxis] + y_center[:, np.newaxis]
        new_x_center = dx * w[:, np.newaxis] + x_center[:, np.newaxis]
        new_h = np.exp(dh) * h[:, np.newaxis]
        new_w = np.exp(dw) * w[:, np.newaxis]
        pred_box = np.zeros(pred_anchor_locs.shape, dtype=pred_anchor_locs.dtype)
        pred_box[:, 0::4] = new_y_center - 0.5 * new_h
        pred_box[:, 1::4] = new_x_center - 0.5 * new_w
        pred_box[:, 2::4] = new_y_center + 0.5 * new_h
        pred_box[:, 3::4] = new_x_center + 0.5 * new_w
        # make sure the predicted anchor box stay inside the image
        pred_box[:, 0] = np.clip(pred_box[:, 0], 0, image_size[0]) # [0,H]
        pred_box[:, 2] = np.clip(pred_box[:, 2], 0, image_size[0])
        pred_box[:, 1] = np.clip(pred_box[:, 1], 0, image_size[1]) # [0,W]
        pred_box[:, 3] = np.clip(pred_box[:, 3], 0, image_size[1])
        # Remove predicted boxes with either height or width less than minimum size
        height_diff = pred_box[:, 2] - pred_box[:, 0]
        width_diff = pred_box[:, 3] - pred_box[:, 1]
        idx = np.where((height_diff >= min_size) & (width_diff >= min_size))[0]
        pred_box = pred_box[idx, :]
        objectness_score = objectness_score[idx] # [22500, 1]
        # select num_pre_nms number of anchor box
        # EX: 22500 -> 12000
        idx = np.argsort(objectness_score.ravel())[::-1]
        idx = idx[0:num_pre_nms]
        pred_box = pred_box[idx, :]
        objectness_score = objectness_score[idx]

        # perform Non-maximal suppression
        # Cuda or CPU?
        # TODO: move to cpu?
        nms_result = nms(torch.from_numpy(pred_box).cuda(), torch.from_numpy(objectness_score).cuda(), nms_thresh)
        nms_result[:num_after_nms]
        pred_box = pred_box[nms_result.cpu().numpy()]
        return pred_box


# Test whether cross section are correctly created.

import numpy as np

import prince_cr.config as config
from prince_cr import cross_sections, photonfields, core

config.x_cut = 1e-4
config.x_cut_proton = 1e-2
config.tau_dec_threshold = np.inf
config.max_mass = 14
config.debug_level = 0 # suppress info statements

pf = photonfields.CombinedPhotonField(
                [photonfields.CMBPhotonSpectrum,
                photonfields.CIBGilmore2D])
cs = cross_sections.CompositeCrossSection([(0., cross_sections.TabulatedCrossSection, ('CRP2_TALYS',)),
                                        (0.14, cross_sections.SophiaSuperposition, ())])

import unittest
class TestCsec(unittest.TestCase):

    def test_kernel_1(self):
        prince_run_talys = core.PriNCeRun(max_mass = 4, photon_field=pf, cross_sections=cs)
        self.assertEqual(prince_run_talys.int_rates._batch_matrix.shape, (88344, 72))
        self.assertEqual(prince_run_talys.int_rates._batch_rows.shape, (88344,))
        self.assertEqual(prince_run_talys.int_rates._batch_cols.shape, (88344,))
        self.assertEqual(prince_run_talys.int_rates._batch_vec.shape, (88344,))

    def test_kernel_4(self):
        prince_run_talys = core.PriNCeRun(max_mass = 1, photon_field=pf, cross_sections=cs)
        self.assertEqual(prince_run_talys.int_rates._batch_matrix.shape, (17528, 72))
        self.assertEqual(prince_run_talys.int_rates._batch_rows.shape, (17528,))
        self.assertEqual(prince_run_talys.int_rates._batch_cols.shape, (17528,))
        self.assertEqual(prince_run_talys.int_rates._batch_vec.shape, (17528,))

    def test_kernel_14(self):
        prince_run_talys = core.PriNCeRun(max_mass = 14, photon_field=pf, cross_sections=cs)
        self.assertEqual(prince_run_talys.int_rates._batch_matrix.shape, (287664, 72))
        self.assertEqual(prince_run_talys.int_rates._batch_rows.shape, (287664,))
        self.assertEqual(prince_run_talys.int_rates._batch_cols.shape, (287664,))
        self.assertEqual(prince_run_talys.int_rates._batch_vec.shape, (287664,))

if __name__ == '__main__':
    unittest.main()


#!/usr/bin/env python
"""
 Copyright 2020 Johns Hopkins University  (Author: Jesus Villalba)
 Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
"""
import sys
import os
from jsonargparse import (
    ArgumentParser,
    ActionConfigFile,
    ActionParser,
    namespace_to_dict,
)
import time
import logging
import multiprocessing

import numpy as np

import torch
import torch.nn as nn

from hyperion.hyp_defs import config_logger, set_float_cpu
from hyperion.torch.utils import open_device
from hyperion.torch.utils import ddp
from hyperion.torch.trainers import XVectorTrainerFromWav as Trainer
from hyperion.torch.models import TDNNXVector as XVec
from hyperion.torch.data import AudioDataset as AD
from hyperion.torch.data import ClassWeightedSeqSampler as Sampler
from hyperion.torch.metrics import CategoricalAccuracy
from hyperion.torch.narchs import AudioFeatsMVN as AF

def init_data(
    audio_path,
    train_list,
    val_list,
    train_aug_cfg,
    val_aug_cfg,
    num_workers,
    num_gpus,
    rank,
    **kwargs
):

    ad_args = AD.filter_args(**kwargs)
    sampler_args = Sampler.filter_args(**kwargs)
    if rank == 0:
        logging.info("audio dataset args={}".format(ad_args))
        logging.info("sampler args={}".format(sampler_args))
        logging.info("init datasets")

    train_data = AD(audio_path, train_list, aug_cfg=train_aug_cfg, **ad_args)
    val_data = AD(audio_path, val_list, aug_cfg=val_aug_cfg, is_val=True, **ad_args)

    if rank == 0:
        logging.info("init samplers")
    train_sampler = Sampler(train_data, **sampler_args)
    val_sampler = Sampler(val_data, **sampler_args)

    num_workers_per_gpu = int((num_workers + num_gpus - 1) / num_gpus)
    largs = (
        {"num_workers": num_workers_per_gpu, "pin_memory": True} if num_gpus > 0 else {}
    )

    train_loader = torch.utils.data.DataLoader(
        train_data, batch_sampler=train_sampler, **largs
    )

    test_loader = torch.utils.data.DataLoader(
        val_data, batch_sampler=val_sampler, **largs
    )

    return train_loader, test_loader

def init_feats(rank, **kwargs):
    feat_args = AF.filter_args(**kwargs["feats"])
    if rank == 0:
        logging.info("feat args={}".format(feat_args))
        logging.info("initializing feature extractor")
    feat_extractor = AF(trans=True, **feat_args)
    if rank == 0:
        logging.info("feat-extractor={}".format(feat_extractor))
    return feat_extractor

def init_xvector(num_classes, rank, **kwargs):
    xvec_args = XVec.filter_args(**kwargs)
    if rank == 0:
        logging.info("xvector network args={}".format(xvec_args))
    xvec_args["num_classes"] = num_classes
    model = XVec(**xvec_args)
    if rank == 0:
        logging.info("x-vector-model={}".format(model))
    return model

def train_xvec(gpu_id, args):

    config_logger(args.verbose)
    del args.verbose
    logging.debug(args)

    kwargs = namespace_to_dict(args)
    torch.manual_seed(args.seed)
    set_float_cpu("float32")

    ddp_args = ddp.filter_ddp_args(**kwargs)
    device, rank, world_size = ddp.ddp_init(gpu_id, **ddp_args)
    kwargs["rank"] = rank
    train_loader, test_loader = init_data(**kwargs)
    feat_extractor = init_feats(**kwargs)
    model = init_xvector(train_loader.dataset.num_classes, **kwargs)

    trn_args = Trainer.filter_args(**kwargs)
    if rank == 0:
        logging.info("trainer args={}".format(trn_args))
    metrics = {"acc": CategoricalAccuracy()}
    trainer = Trainer(
        model,
        feat_extractor,
        device=device,
        metrics=metrics,
        ddp=world_size > 1,
        **trn_args
    )
    if args.resume:
        trainer.load_last_checkpoint()
    trainer.fit(train_loader, test_loader)

    ddp.ddp_cleanup()

if __name__ == "__main__":

    parser = ArgumentParser(
        description="Train XVector with ResNet encoder from audio files"
    )

    parser.add_argument("--cfg", action=ActionConfigFile)
    parser.add_argument("--audio-path", required=True)
    parser.add_argument("--train-list", required=True)
    parser.add_argument("--val-list", required=True)

    AD.add_class_args(parser)
    Sampler.add_class_args(parser)

    parser.add_argument("--train-aug-cfg", default=None)
    parser.add_argument("--val-aug-cfg", default=None)

    parser.add_argument(
        "--num-workers", type=int, default=5, help="num_workers of data loader"
    )

    AF.add_class_args(parser, prefix="feats")
    XVec.add_class_args(parser)
    Trainer.add_class_args(parser)

    ddp.add_ddp_args(parser)
    parser.add_argument("--seed", type=int, default=1123581321, help="random seed")
    parser.add_argument(
        "--resume",
        action="store_true",
        default=False,
        help="resume training from checkpoint",
    )
    parser.add_argument(
        "-v", "--verbose", dest="verbose", default=1, choices=[0, 1, 2, 3], type=int
    )

    args = parser.parse_args()
    try:
        gpu_id = int(os.environ["LOCAL_RANK"])
    except:
        gpu_id = 0

    if gpu_id == 0:
        try:
            config_file = Path(args.exp_path) / "config.yaml"
            parser.save(args, str(config_file), format="yaml", overwrite=True)
        except:
            pass

    # torch docs recommend using forkserver
    multiprocessing.set_start_method("forkserver")
    train_xvec(gpu_id, args)


#!/usr/bin/env python3

import sys, os, time
sys.dont_write_bytecode = True
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

from datetime import datetime, timedelta
import tensorflow as tf
import numpy as np
import pandas as pd
import gym
import retro
import retro_contest.local

from policy import Policy
from actor import Actor, process_state

def main():

    #train/test vars
    recover=True
    train_steps = 10000 * 12 #batches per hour * hours
    train_time = 0 #seconds
    test_render = False
    test_episodes = 3
    print_interval = 10000 #FIXME saves on this interval also
    n_actors = 1

    gym_envs = False
    #output an int action if true, else one hot action
    act_int = True if gym_envs else False

    #init env
    #def get_env(seed=42):
    #    if gym_envs:
    #        env_name = ['CartPole-v0', 'Acrobot-v1'][0]
    #        env = gym.make(env_name)
    #        env.seed(seed)
    #        return env
    #    else:
    #        env_name = 'SonicTheHedgehog-Genesis'
    #        env_state = 'GreenHillZone.Act1'
    #        env = retro_contest.local.make(game=env_name, state=env_state)
    #        env.seed(seed)
    #        return env

    #get some info for the policy, init the emulator wrapper
    env = Env_Controller('./sonic-train.csv',
            init=('SonicTheHedgehog-Genesis', 'GreenHillZone.Act1'))
    state = env.reset()
    n_actions = env.action_space.n
    state_shape = process_state(state).shape

    #init the session
    with tf.Session() as sess:

        #global learner
        learner_name = 'learner_global'
        learner_policy = Policy(state_shape, n_actions, learner_name,
                act_int=act_int, recover=True, sess=sess)

        actors = []
        for i in range(n_actors):
            #FIXME: for now just stay synchronous
                #only one emulator allowed per thread
            #env = get_env(42 + i)
            actor_policy = Policy(state_shape, n_actions,
                    'actor_%s' % i, act_int=act_int, sess=sess,
                    pull_scope=learner_name)
            actor = Actor(env, actor_policy, learner_policy)
            actors.append(actor)

        #FIXME: new thread for each actor
        end_time = datetime.now() + timedelta(seconds=train_time)
        print('[*] training for %ss with %s actors (ending at %s)' % (
                train_time, len(actors), '{:%H:%M:%S}'.format(end_time)))
        start_time = time.time()
        start_step = learner_policy.get_step()
        #while(train_time > (time.time() - start_time)):
        while(True):
            for actor in actors:
                batch = actor.run()
                learner_policy.train(batch)
            step = learner_policy.get_step()
            if step % print_interval == 0:
                learner_policy.save()
                print('[*] learner step: %s' % step)
                info = actor.test(episodes=3, render=False)
                print(pd.DataFrame.from_dict(data=info
                        ).describe().loc[['min', 'max', 'mean', 'std']])

            #FIXME: debug limit train steps
            if (step - start_step) >= train_steps:
                break

        steps_per_s = float(learner_policy.get_step() - start_step) / float(
                time.time() - start_time)
        steps_per_min = round(steps_per_s * 60, 3)
        steps_per_sec = round(steps_per_s, 3)
        print('[+] done training: %s steps/min (%s /s)' % (steps_per_min,
                steps_per_sec))
        learner_policy.save()

        #print('learner at: %s' % learner_policy.get_step())
        #print('actor at: %s' % actor_policy.get_step())

        #test with clean actor, clean environment
        print('[*] testing for %s episodes per state' % test_episodes)
        #env = get_env(429)
        env.keep_env = True
        actor_policy = Policy(state_shape, n_actions,
                'actor_test', act_int=act_int, sess=sess)
        actor = Actor(env, actor_policy, learner_policy)

        all_stats = {g: {k:[] for k in env.game_states[g]}
                for g in env.game_states.keys()}
        for game in env.game_states.keys():
            for state in env.game_states[game]:
                print('[*] testing %s (%s)' % (game, state))
                actor.env.switch_env(game, state)
                info = actor.test(episodes=3,
                        render=test_render)

                #extract some basic stats from raw results
                all_stats[game][state] = pd.DataFrame.from_dict(data=info
                        ).describe().loc[['min', 'max', 'mean', 'std']]
                #change column names to include state name
                all_stats[game][state].columns = ['%s_%s' % (state, x)
                        for x in all_stats[game][state].columns.values]

                #print(all_stats[game][state], '\n')

        #output stats FIXME: write dfs_concat, calc_* to disk
        cols = {'rewards': [], 'steps': []}
        for game in all_stats:
            print(game, (79-len(game))*'*')
            dfs = [all_stats[game][state] for state in all_stats[game]]
            dfs_concat = pd.concat(dfs, axis=1)
            print(dfs_concat, '\n')
            #for col in dfs_concat.columns.values:
            #    if 'reward' in col:
            #        cols['rewards'].append(col)
            #    elif 'step' in col:
            #        cols['steps'].append(col)
            #r = dfs_concat[cols['rewards']]
            #s = dfs_concat[cols['steps']]

            #calc_s = pd.concat([s.min(axis=1).rename('calc_min'),
            #        s.max(axis=1).rename('calc_max'),
            #        s.mean(axis=1).rename('calc_mean'),
            #        s.std(axis=1).rename('calc_std')], axis=1)
            #calc_r = pd.concat([r.min(axis=1).rename('calc_min'),
            #        r.max(axis=1).rename('calc_max'),
            #        r.mean(axis=1).rename('calc_mean'),
            #        r.std(axis=1).rename('calc_std')], axis=1)

            cols = [x for x in dfs_concat.columns.values if 'reward' in x]
            calc_r = pd.concat([dfs_concat[cols].mean(
                    axis=1).rename('calc_mean'), dfs_concat[cols].std(
                    axis=1).rename('calc_std')], axis=1)
            print(calc_r, '\n\n')

            #print('\nsteps:\n', calc_s, '\nrewards:\n', calc_r)

#wrapper that handles learning on multiple retro env states
    #changes to different state / game each episode
    #uses delta rewards as recommended
class Env_Controller():
    def __init__(self, csv, init=None):
        #games, states read from csv
        #self.game_states = self._csv_dict(csv)

        games = ['SonicTheHedgehog-Genesis',
                'SonicTheHedgehog2-Genesis',
                'SonicAndKnuckles3-Genesis']
        self.all_game_states = {k:retro.list_states(k) for k in games}
        #58 total states

        #states yet to be added
        self.game_states = {} #current states used for training
        self.available_states = []
        for game in games:
            for state in self.all_game_states[game]:
                self.available_states.append((game, state))

        #start with two env states
        add_n_games = 44 #FIXME: steps / 10k + 2
        for _ in range(add_n_games):
            self._add_state()

        if init:
            self.env = retro_contest.local.make(game=init[0], state=init[1])
        self.act_int = False
        self.action_space = self.env.action_space
        self.done = False

        #FIXME: every x steps add an env state
            #progressively add
        self.keep_env = False
        self.switch_interval = 10000 #steps before adding a state
        self.switch_steps = 0
        self.add_interval = 300000 #350000 is ~ 1h at 3 batch/sec
        self.add_steps = 0
            # ~60 hrs to full dataset

        self.episode_steps = 0
        self.episode_len = 200 #limit length of episodes while training

        self._cur_r = 0.0
        self._max_r = 0.0

    def _add_state(self, game_state=None):
        if len(self.available_states) == 0:
            return
        if game_state == None:
            game_state = self.available_states.pop(0)

        self._last_added = game_state
        g, s = game_state
        if g in self.game_states.keys():
            self.game_states[g].append(s)
        else:
            self.game_states[g] = [s]

        print('[*] added: %s %s' % (g, s))

    def _csv_dict(self, csv):
        df = pd.read_csv(csv, header=0, index_col=0)
        return {k: list(v) for k,v in df.groupby("game")["state"]}

    def create_env(self, game=None, state=None):
        assert(game or not state) #dont provide a state without a game
        self._cur_r = 0.0
        self._max_r = 0.0
        self.switch_steps = 0
        self.episode_steps = 0
        #chooses random game, state
        if game == None:
            game = np.random.choice(list(self.game_states.keys()))
        if state == None:
            state = np.random.choice(self.game_states[game])

        #priority game switching during adding process
            #chance gradually decreases, has choice to be prev chosen
        min_chance = 0.25 #real chance is min + (1 / num_envs)
        if len(self.available_states) > 0:
            arr = [self._last_added, (game, state)]
            idx = np.random.choice(len(arr), p=[min_chance, 1 - min_chance])
            game, state = arr[idx]

        self.env = retro_contest.local.make(game=game, state=state)
        #print('    curr env: %s(%s)' % (game, state))

    def switch_env(self, game=None, state=None):
        self.env.close()
        self.create_env(game, state)

    def seed(self, n):
        self.env.seed(n)

    def reset(self):
        self._cur_r = 0.0
        self._max_r = 0.0
        self.episode_steps = 0
        return self.env.reset()

    def render(self):
        self.env.render()

    def step(self, action):
        #add state on interval
        if self.add_steps > self.add_interval:
            self._add_state()
            self.add_steps = 0
        #swap env on interval
        if self.done:
            if not self.keep_env and self.switch_steps > self.switch_interval:
                self.switch_env() #when episode is finished, swap envs
            self.reset()
            obs = self.env.step(action)
        else:
            obs = self.env.step(action)

        state, reward, done, info = obs

        #scale reward
        reward = reward * 0.01

        #FIXME: enforce and increase limit on episode length on interval
        # if self.episode_steps >= self.episode_len:
            #done = False
        #if ???:
            #self.episode_len += 100

        self.done = done

        #delta rewards (currently not working, 4hr bad policy)
        #self._cur_r += reward
        #reward = max(0.0, self._cur_r - self._max_r)
        #self._max_r = max(self._max_r, self._cur_r)
        self.switch_steps += 1
        self.add_steps += 1
        self.episode_steps += 1
        return state, reward, done, info

if __name__ == '__main__':
    main()


import argparse
import sys
sys.path.extend(['../', '.', '../../', '../../../'])
from json_config import CONFIG
from training import BaseTrainer, Base, Crop, Padding, Normalize
from testing import CVBaseTester
import os
from residualizer_dataset import BIOBD_BSNIP_Residualized_Dataset, IndexDataManager
import torch, json
from collections import OrderedDict
import numpy as np
import logging

class BIOBD_BSNIP_Model(CVBaseTester, BaseTrainer):
    """
        Defines all the hyperparameters/directories by hand and instantiates the DL model.
    """

    def __init__(self, args, **kwargs):
        self.args = args
        self.logger = logging.getLogger("pynet")
        # Training/Testing split manually defined
        OUTPUT_DIR = "/neurospin/tmp/psy_sbox/analyses/202104_biobd-bsnip_cata12vbm_predict-dx"
        cv_filename = os.path.join(OUTPUT_DIR, 'mwp1-gs_all_cvlso-learningcurves_train-test-folds-by-size.json')
        with open(cv_filename) as json_file:
            cv_lso_lrncurv = json.load(json_file)
        # Important: keep the same order while iterating through the dict
        cv_lso_lrncurv = OrderedDict({k: [np.array(x) for x in v] for k, v in cv_lso_lrncurv.items()})
        k = 0
        for i, fold in enumerate(cv_lso_lrncurv):
            if int(fold.split('-')[1]) == args.N_train_max:
                self.logger.info("Site %s <--> Fold %i"%(fold, k))
                k += 1

        train_indices = tuple([split[0] for (fold, split) in cv_lso_lrncurv.items()
                               if int(fold.split('-')[1]) == args.N_train_max])
        test_indices = tuple([split[1] for (fold, split) in cv_lso_lrncurv.items()
                               if int(fold.split('-')[1]) == args.N_train_max])

        # Directories
        self.args.checkpoint_dir = "/neurospin/psy_sbox/bd261576/checkpoints/bipolar_prediction/BIOBD_BSNIP/" \
                                   "DenseNet/N_%i"%self.args.N_train_max
        self.args.input_path = os.path.join(OUTPUT_DIR, "biobd-bsnip_cat12vbm_mwp1-gs.npy")
        self.args.metadata_path = os.path.join(OUTPUT_DIR, "biobd-bsnip_cat12vbm_participants.csv")
        self.args.exp_name = self.args.exp_name or "DenseNet_BIOBD_BSNIP_N%i"%args.N_train_max
        self.args.outfile_name = None

        ## Hyperparameters
        self.args.gamma_lr = 0.2
        self.args.step_size_lr = 10
        self.args.lr = 1e-4
        self.args.num_classes = 1
        self.args.nb_epochs = 100
        self.args.nb_folds = 13
        self.args.folds = None # do not select a fold
        self.args.nb_epochs_per_saving = 50
        self.args.batch_size = 8
        self.args.sampler = "random"
        self.args.metrics = ["balanced_accuracy", "roc_auc"]
        self.args.labels = ["dx"]

        # Testing
        self.args.test_all_epochs = False
        self.args.test_best_epoch = False
        self.args.with_logit = False
        self.args.predict = False
        self.args.cuda = True

        ## Input tf
        self.args.input_transforms = [Crop((1, 121, 128, 121)), Padding([1, 128, 128, 128], mode='constant'),
                                      Normalize()]

        # Build network
        self.args.dropout, self.args.concrete_dropout, args.bayesian = 0, None, None
        self.net = BaseTrainer.build_network("densenet121", num_classes=self.args.num_classes,
                                             args=self.args, in_channels=1)

        # Build DataManager
        self.manager = IndexDataManager([self.args.input_path], [self.args.metadata_path],
                                        training_indices=train_indices,
                                        val_indices=test_indices,
                                        test_indices=[],
                                        mmap_mode=True,
                                        labels=self.args.labels,
                                        sampler=self.args.sampler,
                                        input_transforms=self.args.input_transforms,
                                        dataset_cls=BIOBD_BSNIP_Residualized_Dataset,
                                        dataset_kwargs={"residualize": self.args.residualize},
                                        batch_size=self.args.batch_size,
                                        pin_memory=True)

        # Build loss
        self.loss = BaseTrainer.build_loss("BCE", net=self.net, args=self.args)

        # Build optimizer
        self.optimizer = torch.optim.Adam(self.net.parameters(), lr=self.args.lr, **CONFIG['optimizer']['Adam'])
        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, gamma=self.args.gamma_lr,
                                                         step_size=self.args.step_size_lr)
        # Build the model
        self.model = Base(model=self.net,
                          metrics=self.args.metrics,
                          use_cuda=True,
                          loss=self.loss,
                          pretrained=self.args.pretrained_path,
                          load_optimizer=False,
                          optimizer=self.optimizer)

    def training(self):
        train_history, valid_history = self.model.training(self.manager,
                                                           nb_epochs=self.args.nb_epochs,
                                                           scheduler=self.scheduler,
                                                           with_validation=False,
                                                           checkpointdir=self.args.checkpoint_dir,
                                                           nb_epochs_per_saving=self.args.nb_epochs_per_saving,
                                                           exp_name=self.args.exp_name,
                                                           standard_optim=True,
                                                           with_visualization=False,
                                                           gpu_time_profiling=False)

        return train_history, valid_history

    def testing(self):
        super().run()

def str2bool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

if __name__=="__main__":
    logger = logging.getLogger("pynet")
    parser = argparse.ArgumentParser()
    parser.add_argument("--N_train_max", type=int, required=True, choices=[100, 300, 500, 700, 903])
    parser.add_argument("--pretrained_path", type=str)
    parser.add_argument("--residualize", type=str2bool, default=True)
    parser.add_argument("--exp_name", type=str)
    args = parser.parse_args()

    model = BIOBD_BSNIP_Model(args)

    # Train the model on every fold
    model.training()
    # Test it at the end
    model.testing()


from ._arcgis_model import ArcGISModel
import tempfile
from pathlib import Path
import json
from ._codetemplate import code
from ._arcgis_model import _EmptyData
import logging
from ._codetemplate import instance_detector_prf

try:
    import torch
    from fastai.vision.learner import cnn_learner
    from fastai.callbacks.hooks import model_sizes
    from fastai.vision.learner import create_body
    from fastai.vision.image import open_image
    from torchvision.models import resnet34
    from torchvision import models
    import numpy as np
    from .._data import prepare_data, _raise_fastai_import_error
    from fastai.callbacks import EarlyStoppingCallback
    from ._arcgis_model import SaveModelCallback, _set_multigpu_callback
    import torchvision
    from torchvision import models
    from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
    from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
    from fastai.basic_train import Learner
    from ._maskrcnn_utils import is_no_color, mask_rcnn_loss, train_callback
    from fastai.torch_core import split_model_idx
    import matplotlib.pyplot as plt
    import matplotlib.patches as patches

    HAS_FASTAI = True
except Exception as e:
    #raise Exception(e)
    HAS_FASTAI = False

class MaskRCNN(ArcGISModel):
    """
    Creates a ``MaskRCNN`` Instance segmentation object

    =====================   ===========================================
    **Argument**            **Description**
    ---------------------   -------------------------------------------
    data                    Required fastai Databunch. Returned data object from
                            ``prepare_data`` function.
    ---------------------   -------------------------------------------
    backbone                Optional function. Backbone CNN model to be used for
                            creating the base of the `MaskRCNN`, which
                            is `resnet50` by default.
                            Compatible backbones: 'resnet50'
    ---------------------   -------------------------------------------
    pretrained_path         Optional string. Path where pre-trained model is
                            saved.
    =====================   ===========================================

    :returns: ``MaskRCNN`` Object
    """
    def __init__(self, data, backbone=None, pretrained_path=None):

        super().__init__(data, backbone)

        self._backbone = models.resnet50

        #if not self._check_backbone_support(self._backbone):
        #    raise Exception (f"Enter only compatible backbones from {', '.join(self.supported_backbones)}")

        self._code = instance_detector_prf

        model = models.detection.maskrcnn_resnet50_fpn(pretrained=True, min_size = data.chip_size)
        in_features = model.roi_heads.box_predictor.cls_score.in_features
        model.roi_heads.box_predictor = FastRCNNPredictor(in_features, data.c)
        in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
        hidden_layer = 256
        model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
                                                       hidden_layer,
                                                       data.c)

        self.learn = Learner(data, model, loss_func = mask_rcnn_loss)
        self.learn.callbacks.append(train_callback(self.learn))
        self.learn.model = self.learn.model.to(self._device)

        # fixes for zero division error when slice is passed
        self.learn.layer_groups = split_model_idx(self.learn.model, [28])
        self.learn.create_opt(lr=3e-3)

        if pretrained_path is not None:
            self.load(pretrained_path)

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return '<%s>' % (type(self).__name__)

    @property
    def supported_backbones(self):
        return [models.detection.maskrcnn_resnet50_fpn.__name__]

    @classmethod
    def from_model(cls, emd_path, data=None):
        """
        Creates a ``MaskRCNN`` Instance segmentation object from an Esri Model Definition (EMD) file.

        =====================   ===========================================
        **Argument**            **Description**
        ---------------------   -------------------------------------------
        emd_path                Required string. Path to Esri Model Definition
                                file.
        ---------------------   -------------------------------------------
        data                    Required fastai Databunch or None. Returned data
                                object from ``prepare_data`` function or None for
                                inferencing.

        =====================   ===========================================

        :returns: `MaskRCNN` Object
        """

        emd_path = Path(emd_path)
        with open(emd_path) as f:
            emd = json.load(f)

        model_file = Path(emd['ModelFile'])

        if not model_file.is_absolute():
            model_file = emd_path.parent / model_file

        model_params = emd['ModelParameters']

        try:
            class_mapping = {i['Value'] : i['Name'] for i in emd['Classes']}
            color_mapping = {i['Value'] : i['Color'] for i in emd['Classes']}
        except KeyError:
            class_mapping = {i['ClassValue'] : i['ClassName'] for i in emd['Classes']}
            color_mapping = {i['ClassValue'] : i['Color'] for i in emd['Classes']}

        if data is None:
            empty_data = _EmptyData(path=tempfile.TemporaryDirectory().name, loss_func=None, c=len(class_mapping) + 1, chip_size=emd['ImageHeight'])
            empty_data.class_mapping = class_mapping
            empty_data.color_mapping = color_mapping
            return cls(empty_data, **model_params, pretrained_path=str(model_file))
        else:
            return cls(data, **model_params, pretrained_path=str(model_file))

    def _create_emd(self, path):
        import random
        super()._create_emd(path)
        self._emd_template["Framework"] = "arcgis.learn.models._inferencing"
        self._emd_template["ModelConfiguration"] = "_maskrcnn_inferencing"
        self._emd_template["InferenceFunction"] = "ArcGISInstanceDetector.py"

        self._emd_template["ExtractBands"] = [0, 1, 2]
        self._emd_template['Classes'] = []
        class_data = {}
        for i, class_name in enumerate(self._data.classes[1:]):  # 0th index is background
            inverse_class_mapping = {v: k for k, v in self._data.class_mapping.items()}
            class_data["Value"] = inverse_class_mapping[class_name]
            class_data["Name"] = class_name
            color = [random.choice(range(256)) for i in range(3)] if is_no_color(self._data.color_mapping) else \
            self._data.color_mapping[inverse_class_mapping[class_name]]
            class_data["Color"] = color
            self._emd_template['Classes'].append(class_data.copy())

        json.dump(self._emd_template, open(path.with_suffix('.emd'), 'w'), indent=4)
        return path.stem

    @property
    def _model_metrics(self):
        return {}

    def _predict_results(self, xb):

        self.learn.model.eval()
        predictions = self.learn.model(xb.cuda())
        predictionsf =[]
        for i in range(len(predictions)):
            predictionsf.append({})
            predictionsf[i]['masks'] = predictions[i]['masks'].detach().cpu().numpy()
            predictionsf[i]['boxes'] = predictions[i]['boxes'].detach().cpu().numpy()
            predictionsf[i]['labels'] = predictions[i]['labels'].detach().cpu().numpy()
            predictionsf[i]['scores'] = predictions[i]['scores'].detach().cpu().numpy()
            del predictions[i]['masks']
            del predictions[i]['boxes']
            del predictions[i]['labels']
            del predictions[i]['scores']
        del xb
        torch.cuda.empty_cache()

        return predictionsf

    def _predict_postprocess(self, predictions, threshold=0.5, box_threshold = 0.5):

        pred_mask = []
        pred_box = []

        for i in range(len(predictions)):
            out = predictions[i]['masks'].squeeze()
            pred_box.append([])

            if out.shape[0] != 0:  # handle for prediction with n masks
                if len(out.shape) == 2: # for out dimension hxw (in case of only one predicted mask)
                    out = out[None]
                ymask = np.where(out[0]> threshold, 1, 0)
                #if torch.max(out[0]) > threshold:
                if predictions[i]['scores'][0] > box_threshold:
                    pred_box[i].append(predictions[i]['boxes'][0])
                for j in range(1,out.shape[0]):
                    ym1 = np.where(out[j]> threshold, j+1, 0)
                    ymask += ym1
                    #if torch.max(out[j]) > threshold:
                    if predictions[i]['scores'][j] > box_threshold:
                        pred_box[i].append(predictions[i]['boxes'][j])
            else:
                ymask = np.zeros((self._data.chip_size, self._data.chip_size)) # handle for not predicted masks
            pred_mask.append(ymask)
        return pred_mask, pred_box

    def show_results(self, mode='mask', mask_threshold=0.5, box_threshold=0.7, nrows=None, imsize=5, index=0, alpha=0.5, cmap='tab20'):
        """
        Displays the results of a trained model on a part of the validation set.

        =====================   ===========================================
        **Argument**            **Description**
        ---------------------   -------------------------------------------
        mode                    Required arguments within ['bbox', 'mask', 'bbox_mask'].
                                    * ``bbox`` - For visualizing only boundig boxes.
                                    * ``mask`` - For visualizing only mask
                                    * ``bbox_mask`` - For visualizing both mask and bounding boxes.
        ---------------------   -------------------------------------------
        mask_threshold          Optional float. The probabilty above which
                                a pixel will be considered mask.
        ---------------------   -------------------------------------------
        box_threshold           Optional float. The pobabilty above which
                                a detection will be considered valid.
        ---------------------   -------------------------------------------
        nrows                   Optional int. Number of rows of results
                                to be displayed.
        =====================   ===========================================
        """

        if mode not in ['bbox', 'mask', 'bbox_mask']:
            raise Exception("mode can be only ['bbox', 'mask', 'bbox_mask']")

        # Get Number of items
        if nrows is None:
            nrows = self._data.batch_size
        ncols=2

        # Get Batch
        xb,yb = self._data.one_batch('DatasetType.Valid')

        predictions = self._predict_results(xb)

        pred_mask, pred_box = self._predict_postprocess(predictions, mask_threshold, box_threshold)

        fig, ax = plt.subplots(nrows=nrows, ncols=ncols, figsize=(ncols*imsize, nrows*imsize))
        fig.suptitle('Ground Truth / Predictions', fontsize=20)

        for i in range(nrows):
            ax[i][0].imshow(xb[i].numpy().transpose(1,2,0))
            ax[i][0].axis('off')
            if mode in ['mask', 'bbox_mask']:
                yb_mask = yb[i][0].numpy()
                for j in range(1, yb[i].shape[0]):
                    max_unique = np.max(np.unique(yb_mask))
                    yb_j = np.where(yb[i][j]>0, yb[i][j] + max_unique, yb[i][j])
                    yb_mask += yb_j
                ax[i][0].imshow(yb_mask, cmap = cmap, alpha = alpha)
            ax[i][0].axis('off')
            ax[i][1].imshow(xb[i].numpy().transpose(1,2,0))
            ax[i][1].axis('off')
            if mode in ['mask', 'bbox_mask']:
                ax[i][1].imshow(pred_mask[i], cmap=cmap, alpha = alpha)
            if mode in ['bbox','bbox']:
                if pred_box[i] != []:
                    for num_boxes in pred_box[i]:
                        rect = patches.Rectangle((num_boxes[0], num_boxes[1]), num_boxes[2]-num_boxes[0], num_boxes[3]-num_boxes[1], linewidth=1, edgecolor='r', facecolor='none')
                        ax[i][1].add_patch(rect)
            ax[i][1].axis('off')
        plt.subplots_adjust(top=0.95)
        torch.cuda.empty_cache()


import sys
sys.path.append('..') # Add src to path
import os
os.environ["HDF5_USE_FILE_LOCKING"]='FALSE'
import datetime
import h5py
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tensorflow as tf
from src.display import get_cmap
from src.utils import make_log_dir
# comment these out if you don't have cartopy
import cartopy.feature as cfeature
from src.display.cartopy import make_ccrs,make_animation
from make_dataset import NowcastGenerator,get_nowcast_train_generator,get_nowcast_test_generator
from unet_benchmark import create_model
from unet_benchmark import nowcast_mae, nowcast_mse
import pdb
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import time
import sys
import json
from pathlib import Path

data_path="/scratch/li.baol/SEVIR"
# Target locations of sample training & testing data
DEST_TEST_FILE= os.path.join(data_path, 'data/processed/nowcast_testing_000.h5')
DEST_TEST_META= os.path.join(data_path, 'data/processed/nowcast_testing_000_META.csv')

# Control how many samples are read.   Set to -1 to read all 5000 samples.
N_TRAIN=-1
TRAIN_VAL_FRAC=0.8
#set_trace()
N_TEST=-1
num_iters = 1000
batch_size = int(sys.argv[1])
#for i in range(num_iters):
#    print(f'{i}/{num_iters} progress', end='\r', flush=True)
#    time.sleep(0.1)
#pdb.set_trace()

with h5py.File(DEST_TEST_FILE,'r') as hf:
    Nr = N_TEST if N_TEST>=0 else hf['IN_vil'].shape[0]
    X_test = hf['IN_vil'][:Nr]
    Y_test = hf['OUT_vil'][:Nr]
    testing_meta=pd.read_csv(DEST_TEST_META).iloc[:Nr]

# Add more as needed
params={
    'start_neurons'   :16,      # Controls size of hidden layers in CNN, higher = more complexity
    'activation'      :'relu',  # Activation used throughout the U-Net,  see https://www.tensorflow.org/api_docs/python/tf/keras/activations
    'loss'            :'mae',   # Either 'mae' or 'mse', or others as https://www.tensorflow.org/api_docs/python/tf/keras/losses
    'loss_weights'    :0.021,    # Scale for loss.  Recommend squaring this if using MSE
    'opt'             :tf.keras.optimizers.Adam,  # optimizer, see https://www.tensorflow.org/api_docs/python/tf/keras/optimizers
    'learning_rate'   :0.001,   # Learning rate for optimizer
    'num_epochs'      :10,       # Number of epochs to train for
    'batch_size'      :8        # Size of batches during training
}

unet = create_model(start_neurons=params['start_neurons'],activation=params['activation'])
unet.load_weights('logs/unet.hdf5')

#batch_size,batch_num=8,0
#bs,be=batch_size*batch_num,batch_size*(batch_num+1)
#x_test,y_test,meta = X_test[bs:be],Y_test[bs:be],testing_meta.iloc[bs:be]
#
## warm up
#unet.predict(x_test)
time_record = []

for i in range(num_iters):
    indexs = np.random.choice(len(X_test), batch_size, replace=False)
    x_test = X_test[indexs]
#    if len(x_test) != 8:
#        pdb.set_trace()
    t_start = time.time()
    unet.predict(x_test)
    t_end = time.time()
    lat = round((t_end - t_start)*1000,3)
    time_record.append(lat) # ms
    print(f'{i}/{num_iters} done, latency is {lat}ms', end='\r', flush=True)

gpu = sys.argv[2]

Path('logs/time_records').mkdir(parents=True, exist_ok=True)
with open(f'logs/time_records/{gpu}.json', 'w') as fp:
    json.dump(time_record, fp, indent=4)


import tensorflow as tf
import numpy as np

from models import vae
from models import layers as nw

class JointVAE(vae.VAE):
    """
    Variational Auto-Encoder with Two Input Modalities
    """
    def __init__(self, arguments, name, tracker, init_minibatches, session=None, log_dir=None, model_dir=None):

        self.x1_init, self.x2_init, self.x1p_init, self.x2p_init = init_minibatches

        # constructor
        super(JointVAE, self).__init__(arguments=arguments, name=name, tracker=tracker, init_minibatch=None,
                                       session=session, log_dir=log_dir, model_dir=model_dir)

    def _initialize(self,):
        # additional parameters
        self.objective = self.args["objective"]

        # input placeholders
        self.x1 = tf.placeholder(tf.float32, [None, self.n_x], name='x1')
        self.x2 = tf.placeholder(tf.float32, [None, self.n_x], name='x2')
        self.x1p = tf.placeholder(tf.float32, [None, self.n_x], name='x1_paired')
        self.x2p = tf.placeholder(tf.float32, [None, self.n_x], name='x2_paired')

        # data-dependent weight initialization (Salisman, Kingma - 2016)
        x1_init = tf.constant(self.x1_init, tf.float32)
        x2_init = tf.constant(self.x2_init, tf.float32)
        x1p_init = tf.constant(self.x1p_init, tf.float32)
        x2p_init = tf.constant(self.x2p_init, tf.float32)

        # compute weight initializations
        self._model((x1_init, x2_init, x1p_init, x2p_init), init=True)

        # model specification
        self._model((self.x1, self.x2, self.x1p, self.x2p), init=False)

        # marginal bounds

        self.lx1, self.lx1rec, self.lx1pen, self.logq1, self.logp1 = self._marginal_bound(self.rx1_1, self.x1,
                                        self.z1_mu, self.z1_sigma, self.log_q1, self.z1, scope='marg_x1')

        self.lx2, self.lx2rec, self.lx2pen, self.logq2, self.logp2 = self._marginal_bound(self.rx2_2, self.x2,
                                        self.z2_mu, self.z2_sigma, self.log_q2, self.z2, scope='marg_x2')

        self.lx1p, self.lx1prec, self.lx1ppen, self.logq1p, self.logp1p = self._marginal_bound(self.rx1_1p, self.x1p,
                             self.z1p_mu, self.z1p_sigma, self.log_q1p, self.z1p, scope='marg_x1p')

        self.lx2p, self.lx2prec, self.lx2ppen, self.logq2p, self.logp2p = self._marginal_bound(self.rx2_2p, self.x2p,
                             self.z2p_mu, self.z2p_sigma, self.log_q2p, self.z2p, scope='marg_x2p')

        # joint bound

        self.lx12, self.lx12rec1, self.lx12rec2, self.lx12pen, self.logq12, self.logp12 = self._joint_bound(
            self.rx1_12, self.x1p, self.rx2_12, self.x2p,
            self.z12_mu, self.z12_sigma, self.log_q12, self.z12, scope='joint_bound')

        # translation bounds

        self.tx1 = self._translation_bound(self.rx1_2p, self.x1p, scope='translate_to_x1')
        self.tx2 = self._translation_bound(self.rx2_1p, self.x2p, scope='translate_to_x2')

        # loss function
        self.loss = self._loss(scope='loss')

        # optimizer
        self.step = self._optimizer(self.loss)

        # summary variables
        self.summary = self._summaries()

    def _model(self, xs, init):

        with tf.variable_scope('joint_autoencoder') as scope:
            if not init:
                scope.reuse_variables()

            extra = self.args['flow']

            x1, x2, x1p, x2p = xs

            # encoders
            z1p_mu, z1p_sigma, h1p, he1p = self._encoder(x1p, init=init, scope='x1_enc')
            z2p_mu, z2p_sigma, h2p, he2p = self._encoder(x2p, init=init, scope='x2_enc')

            # joint encoder
            z12_mu, z12_sigma, h12 = self._joint_encoder(z1p_mu, z1p_sigma, h1p, z2p_mu, z2p_sigma, h2p,
                                x1, x2, x1h=he1p, x2h=he2p,
                                extra=extra, init=init, scope='x1x2_enc')

            # sample network
            z12, log_q12 = self._sample(z12_mu, z12_sigma, h12, init=init, scope='sample')

            # decoders
            rx1_12, rx1_12_probs = self._decoder(z12, x1p, init=init, scope='x1_dec')
            rx2_12, rx2_12_probs = self._decoder(z12, x2p, init=init, scope='x2_dec')

        if not init:

            with tf.variable_scope('joint_autoencoder') as scope:

                scope.reuse_variables()

                # unpaired encodings
                z1_mu, z1_sigma, h1, _ = self._encoder(x1, init=init, scope='x1_enc')
                z2_mu, z2_sigma, h2, _ = self._encoder(x2, init=init, scope='x2_enc')

                # additional samples
                z1, log_q1 = self._sample(z1_mu, z1_sigma, h1, init=init, scope='sample')
                z2, log_q2 = self._sample(z2_mu, z2_sigma, h2, init=init, scope='sample')
                z1p, log_q1p = self._sample(z1p_mu, z1p_sigma, h1p, init=init, scope='sample')
                z2p, log_q2p = self._sample(z2p_mu, z2p_sigma, h2p, init=init, scope='sample')

                # reconstructions
                rx1_1, rx1_1_probs = self._decoder(z1, x1, init=init, scope='x1_dec')
                rx2_2, rx2_2_probs = self._decoder(z2, x2, init=init, scope='x2_dec')

                # translations
                rx1_2, rx1_2_probs = self._decoder(z2, x1, init=init, scope='x1_dec')
                rx2_1, rx2_1_probs = self._decoder(z1, x2, init=init, scope='x2_dec')

                # reconstructions (from paired input)
                rx1_1p, rx1_1p_probs = self._decoder(z1p, x1p, init=init, scope='x1_dec')
                rx2_2p, rx2_2p_probs = self._decoder(z2p, x2p, init=init, scope='x2_dec')

                # translations (from paired input)
                rx1_2p, rx1_2p_probs = self._decoder(z2p, x1p, init=init, scope='x1_dec')
                rx2_1p, rx2_1p_probs = self._decoder(z1p, x2p, init=init, scope='x2_dec')

                # final tensors
                self.z12_mu, self.z12_sigma = (z12_mu, z12_sigma)
                self.z1_mu, self.z1_sigma = (z1_mu, z1_sigma)
                self.z2_mu, self.z2_sigma = (z2_mu, z2_sigma)
                self.z1p_mu, self.z1p_sigma = (z1p_mu, z1p_sigma)
                self.z2p_mu, self.z2p_sigma = (z2p_mu, z2p_sigma)

                self.z12, self.log_q12 = (z12, log_q12)
                self.z1, self.log_q1 = (z1, log_q1)
                self.z2, self.log_q2 = (z2, log_q2)
                self.z1p, self.log_q1p = (z1p, log_q1p)
                self.z2p, self.log_q2p = (z2p, log_q2p)

                self.rx1_12, self.rx1_12_probs = (rx1_12, rx1_12_probs)
                self.rx2_12, self.rx2_12_probs = (rx2_12, rx2_12_probs)

                self.rx1_1, self.rx1_1_probs = (rx1_1, rx1_1_probs)
                self.rx2_2, self.rx2_2_probs = (rx2_2, rx2_2_probs)
                self.rx1_2, self.rx1_2_probs = (rx1_2, rx1_2_probs)
                self.rx2_1, self.rx2_1_probs = (rx2_1, rx2_1_probs)

                self.rx1_1p, self.rx1_1p_probs = (rx1_1p, rx1_1p_probs)
                self.rx2_2p, self.rx2_2p_probs = (rx2_2p, rx2_2p_probs)
                self.rx1_2p, self.rx1_2p_probs = (rx1_2p, rx1_2p_probs)
                self.rx2_1p, self.rx2_1p_probs = (rx2_1p, rx2_1p_probs)

    def _joint_encoder(self, mu1, sigma1, h1, mu2, sigma2, h2, x1, x2, x1h, x2h, extra, init, scope):

        with tf.variable_scope(scope):

            jtype = self.args['joint_type']

            n_units = self.args['n_units']
            n_z = self.n_z

            if jtype == 'constrain':
                z_mu, z_sigma = self._constrain(mu1, sigma1, mu2, sigma2, scope='constrain')

                h12 = None
                if h1 is not None and h2 is not None:
                    h12 = tf.nn.elu(h1 + h2)

            elif jtype == 'small':
                z_mu, z_sigma, h12 = nw.joint_fc_encode(x1h, x2h, n_units, n_z, extra, init, scope='small_enc')

            elif jtype == 'large':
                if self.dataset == 'halved_mnist':
                    raise NotImplementedError

                elif self.dataset == 'mnist':
                    raise NotImplementedError

                else:
                    raise NotImplementedError

            else:
                raise NotImplementedError

            return z_mu, z_sigma, h12

    def _constrain(self, x1_mu, x1_sigma, x2_mu, x2_sigma, scope):
        """
        Computes mean and standard deviation of the product of two Gaussians.
        """
        with tf.variable_scope(scope):
            x1_var = tf.square(x1_sigma)
            x2_var = tf.square(x2_sigma)
            x1v_inv = tf.reciprocal(x1_var)
            x2v_inv = tf.reciprocal(x2_var)
            x12_var = tf.reciprocal(x1v_inv + x2v_inv)
            xx = tf.multiply(x1v_inv, x1_mu)
            yy = tf.multiply(x2v_inv, x2_mu)
            x12_mu = tf.multiply(x12_var, xx + yy)
            x12_sigma = tf.sqrt(x12_var)

            return x12_mu, x12_sigma

    def _marginal_bound(self, logits, labels, mu, sigma, log_q, z_K, scope):

        with tf.variable_scope(scope):

            l1 = self._reconstruction(logits=logits, labels=labels, scope='reconstruction')
            l2, log_q, log_p = self._penalty(mu=mu, sigma=sigma, log_q=log_q, z_K=z_K, scope='penalty')
            bound = l1 + l2

            return bound, l1, l2, log_q, log_p

    def _joint_bound(self, x1_logits, x1_labels, x2_logits, x2_labels, mu, sigma, log_q, z_K, scope):

        with tf.variable_scope(scope):

            l1_x1 = self._reconstruction(logits=x1_logits, labels=x1_labels, scope='reconstruction_x1')
            l1_x2 = self._reconstruction(logits=x2_logits, labels=x2_labels, scope='reconstruction_x2')

            l2, log_q, log_p = self._penalty(mu=mu, sigma=sigma, log_q=log_q, z_K=z_K, scope='penalty')

            bound = l1_x1 + l1_x2 + l2

            return bound, l1_x1, l1_x2, l2, log_q, log_p

    def _translation_bound(self, logits, labels, scope):

        with tf.variable_scope(scope):
            return self._reconstruction(logits=logits, labels=labels, scope='reconstruction')

    def _freebits(self, l2, log_q, log_p, alpha):
        """
        Compute freebits penalty if alpha < 0, otherwise return original penalty l2.
        """
        if alpha < 0:
            l2 = tf.reduce_mean(-log_q + log_p, axis=0)
            l2 = tf.minimum(l2, alpha)
            l2 = tf.reduce_sum(l2)

        return l2

    def _loss(self, scope):

        with tf.variable_scope(scope):
            alpha = self.args['anneal']
            ja = self.args['joint_anneal']

            # marginal x1
            lx1pen = self._freebits(self.lx1pen, self.logq1, self.logp1, alpha)
            lx1 = self.lx1rec + lx1pen

            # marginal x2
            lx2pen = self._freebits(self.lx2pen, self.logq2, self.logp2, alpha)
            lx2 = self.lx2rec + lx2pen

            if self.objective == "joint":
                # joint x1 and x2
                lx12pen = self._freebits(self.lx12pen, self.logq12, self.logp12, alpha)
                lx12 = self.lx12rec1 + self.lx12rec2 + lx12pen

                bound = ja*lx12 + lx1 + lx2

            elif self.objective == "translate":
                # marginal x1p
                lx1ppen = self._freebits(self.lx1ppen,  self.logq1p, self.logp1p, alpha)
                lx1p = self.lx1prec + lx1ppen

                # marginal x2p
                lx2ppen = self._freebits(self.lx2ppen, self.logq2p, self.logp2p, alpha)
                lx2p = self.lx2prec + lx2ppen

                bound = (self.tx1 + lx2p) + (self.tx2 + lx1p) + lx1 + lx2

            else:
                raise NotImplementedError

            loss = -bound
            return loss

    def _summaries(self,):

        with tf.variable_scope("summaries"):
            tf.summary.scalar('loss_(ignore_test)', self.loss)

            if self.objective == "joint":
                tf.summary.scalar('lower_bound_on_log_p_x_y', self.lx12)

            elif self.objective == "translate":
                tf.summary.scalar('lower_bound_on_log_p_x_y_t1', self.tx1 + self.lx2p)
                tf.summary.scalar('lower_bound_on_log_p_x_y_t2', self.tx2 + self.lx1p)

            else:
                raise NotImplementedError

            tf.summary.scalar('marg_x1_(ignore_test)', self.lx1)
            tf.summary.scalar('marg_x2_(ignore_test)', self.lx2)
            tf.summary.scalar('marg_x1p', self.lx1p)
            tf.summary.scalar('marg_x2p', self.lx2p)

            tf.summary.scalar('joint', self.lx12)
            tf.summary.scalar('trans_to_x1', self.tx1)
            tf.summary.scalar('trans_to_x2', self.tx2)

            return tf.summary.merge_all()

    def train(self, xs):
        """
        Performs single training step.
        """
        x1, x2, x1_pairs, x2_pairs = xs

        feed = {self.x1: x1, self.x2: x2, self.x1p: x1_pairs, self.x2p: x2_pairs}
        outputs = [self.summary, self.step, self.loss, self.lx1, self.lx2, self.lx12, self.tx1, self.tx2, self.lx1p, self.lx2p]

        summary, _, loss, lx1, lx2, lx12, tx1, tx2, lx1p, lx2p = self.sess.run(outputs, feed_dict=feed)

        if self.objective == 'joint':
            bound = lx12
            terms = {'lower_bound_on_log_p_x_y': bound, 'loss': loss,
                     'lx1': lx1, 'lx2': lx2, 'lx12': lx12, 'tx1': tx1, 'tx2': tx2}
        else:  # translate
            bound_t1 = tx1 + lx2p
            bound_t2 = tx2 + lx1p
            terms = {'lower_bound_on_log_p_x_y_t1': bound_t1, 'lower_bound_on_log_p_x_y_t2': bound_t2, 'loss': loss,
                     'lx1': lx1, 'lx2': lx2, 'lx12': lx12, 'tx1': tx1, 'tx2': tx2}

        self._track(terms, prefix='train_')
        self.tr_writer.add_summary(summary, self.n_steps)

        self.n_steps = self.n_steps + 1

    def test(self, xs):
        """
        Computes lower bound on test data.
        """
        x1, x2 = xs

        feed = {self.x1: x1, self.x2: x2, self.x1p: x1, self.x2p: x2}
        outputs = [self.summary, self.loss, self.lx1p, self.lx2p, self.lx12, self.tx1, self.tx2, self.lx1p, self.lx2p]

        summary, loss, lx1, lx2, lx12, tx1, tx2, lx1p, lx2p = self.sess.run(outputs, feed_dict=feed)

        if self.objective == 'joint':
            bound = lx12
            terms = {'lower_bound_on_log_p_x_y': bound, 'loss': loss,
                     'lx1': lx1, 'lx2': lx2, 'lx12': lx12, 'tx1': tx1, 'tx2': tx2}
        else:  # translate
            bound_t1 = tx1 + lx2p
            bound_t2 = tx2 + lx1p
            terms = {'lower_bound_on_log_p_x_y_t1': bound_t1, 'lower_bound_on_log_p_x_y_t2': bound_t2, 'loss': loss,
                     'lx1': lx1, 'lx2': lx2, 'lx12': lx12, 'tx1': tx1, 'tx2': tx2}

        self._track(terms, prefix='test_')
        self.te_writer.add_summary(summary, self.n_steps)

    def reconstruct(self, xs):
        """
        Reconstruct input.
        If one input is None, then this can be interpretted as a translation from the other input.

        xs: (x1, x2) or (x1, None) or (None, x2)
        """
        z = self.encode(xs, mean=False)
        x1, x2 = self.decode(z)

        return x1, x2

    def encode(self, xs, mean=False):
        """
        Encodes x1 or x2 or both x1 and x2 to the latent space.
        """
        x1, x2 = xs
        feed = {self.x1p: x1, self.x2p: x2}
        outputs = self.z12_mu, self.z12

        if x1 is None and x2 is None:
            raise ValueError

        elif x1 is None:
            feed = {self.x2p: x2}
            outputs = self.z2p_mu, self.z2p

        elif x2 is None:
            feed = {self.x1p: x1}
            outputs = self.z1p_mu, self.z1p

        if mean:
            assert self.is_flow == False
            return self.sess.run(outputs[0], feed_dict=feed)
        else:
            return self.sess.run(outputs[1], feed_dict=feed)

    def decode(self, z):

        if self.is_autoregressive:
            x1 = np.random.rand(z.shape[0], self.n_x)
            x2 = np.random.rand(z.shape[0], self.n_x)

            x1, x2 = self._joint_autoregressive_sampling(z, x1, x2, n_pixels=0,
                    z_var=self.z12, p1_var=self.rx1_12_probs, p2_var=self.rx2_12_probs)

        else:
            feed = {self.z12: z}
            outputs = [self.rx1_12_probs, self.rx2_12_probs]
            rx1, rx2 = self.sess.run(outputs, feed_dict=feed)

            x1 = self._factorized_sampling(rx1)
            x2 = self._factorized_sampling(rx2)

        return x1, x2

    def _joint_autoregressive_sampling(self, z, x1, x2, n_pixels, z_var, p1_var, p2_var):
        """
        Synthesize images autoregressively.
        """
        def _locate_2d(idx, w):
            pos = idx + 1
            r = np.ceil(pos / w)
            c = pos - (r-1)*w

            return int(r-1), int(c-1)

        h = self.h
        w = self.w
        ch = self.n_ch
        n_x = h * w * ch

        remain = h*w - n_pixels

        x1 = x1.copy()
        x2 = x2.copy()
        x1_empty, x2_empty = self._empty_like(x1, x2)

        for i in range(remain):

            feed = {z_var: z, self.x1p: x1, self.x2p: x2, self.x1: x1_empty, self.x2: x2_empty}
            p1, p2 = self.sess.run([p1_var, p2_var], feed_dict=feed)

            hp, wp = _locate_2d(n_pixels + i, w)

            x1 = np.reshape(x1, newshape=[-1, h, w, ch])
            x2 = np.reshape(x2, newshape=[-1, h, w, ch])

            if self.n_ch == 1:
                p1 = np.reshape(p1, newshape=[-1, h, w, ch])
                p1 = p1[:, hp, wp, :]
                x1[:, hp, wp, :] = np.random.binomial(n=1, p=p1)

                p2 = np.reshape(p2, newshape=[-1, h, w, ch])
                p2 = p2[:, hp, wp, :]
                x2[:, hp, wp, :] = np.random.binomial(n=1, p=p2)

            elif self.distribution == 'discrete':
                p1 = np.reshape(p1, newshape=[-1, h, w, ch, 256])
                p1 = p1[:, hp, wp, :, :]
                x1[:, hp, wp, :] = self._categorical_sampling(p1) / 255

                p2 = np.reshape(p2, newshape=[-1, h, w, ch, 256])
                p2 = p2[:, hp, wp, :, :]
                x2[:, hp, wp, :] = self._categorical_sampling(p2) / 255

            elif self.distribution == 'continuous':
                p1 = np.reshape(p1, newshape=[-1, h, w, ch])
                x1[:, hp, wp, :] = p1[:, hp, wp, :]

                p2 = np.reshape(p2, newshape=[-1, h, w, ch])
                x2[:, hp, wp, :] = p2[:, hp, wp, :]

            else:
                raise NotImplementedError

            x1 = np.reshape(x1, newshape=[-1, n_x])
            x2 = np.reshape(x2, newshape=[-1, n_x])

        return x1, x2

    def _empty_like(self, x1, x2):

        x1_shape = list(x1.shape)
        x1_shape[0] = 0
        x2_shape = list(x2.shape)
        x2_shape[0] = 0
        x1_empty = np.zeros(shape=x1_shape)
        x2_empty = np.zeros(shape=x2_shape)

        return x1_empty, x2_empty


# encoding: utf-8
"""
@author: Jiang Wen
@contact: Wenjiang.wj@foxmail.com
This file is an middle-ware to encapsulate light-head-rcnn for academic use.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from .light_config import cfg, config

from . import dataset
import os.path as osp
from . import network_desp
import tensorflow as tf
import numpy as np
import cv2, os, sys, math, json, pickle

from tqdm import tqdm
from utils.py_faster_rcnn_utils.cython_nms import nms, nms_new
from utils.py_utils import misc

from multiprocessing import Queue, Process
from detection_opr.box_utils.box import DetBox
from detection_opr.utils.bbox_transform import clip_boxes, bbox_transform_inv
from functools import partial

class PersonDetector ( object ):
    """
    This is an person detcor based on light-head-rcnn, which is adapted from test.py in experiments/.../ directory
    Some thing in light_config.py is hacked to simplify usage.
    Such as directory and output form.
    """

    def __init__(self, epoch_num=26, show_image=False):
        """
        Person detector init fuction, which will call tf.restore()
        The initialize process may cause many time with many tensorflow log printed out.
        :param epoch_num: 26 is the default model checkpoint
        :param show_image: show image to debug.
        """
        self._show_image = show_image
        self.model_file = osp.join (
            config.output_dir, 'model_dump',
            'epoch_{:d}'.format ( epoch_num ) + '.ckpt' )

        self.func, self.inputs = self._load_model ( self.model_file )

    def detect(self, image, image_id):
        """
        detect on a single image.
        :param image: full path of the image, not necessarily an absolute path
        :return: an dict of detection results.
        """
        all_results = []
        data_dict = dict ( data=image, image_id=image_id )
        result_dict = self.inference ( self.func, self.inputs, data_dict )
        all_results.append ( result_dict )

        if self._show_image:
            image = result_dict['data']
            for db in result_dict['result_boxes']:
                if db.score > config.test_vis_threshold:
                    db.draw ( image )
            if 'boxes' in result_dict.keys ():
                for db in result_dict['boxes']:
                    db.draw ( image )
            import matplotlib.pyplot as plt
            im = cv2.cvtColor ( image, cv2.COLOR_BGR2RGB )
            plt.imshow ( im )
            plt.show ()

        return self._make_result_dict ( all_results )

    def _load_model(self, model_file):
        from tensorflow.python import pywrap_tensorflow
        def get_variables_in_checkpoint_file(file_name):
            try:
                reader = pywrap_tensorflow.NewCheckpointReader ( file_name )
                var_to_shape_map = reader.get_variable_to_shape_map ()
                return var_to_shape_map
            except Exception as e:  # pylint: disable=broad-except
                print ( str ( e ) )
                if "corrupted compressed block contents" in str ( e ):
                    print (
                        "It's likely that your checkpoint file has been compressed "
                        "with SNAPPY." )

        def load_model(sess, model_path):
            # TODO(global variables ?? how about _adam weights)
            variables = tf.global_variables ()
            var_keep_dic = get_variables_in_checkpoint_file ( model_path )
            if 'global_step' in var_keep_dic:
                var_keep_dic.pop ( 'global_step' )

            # vis_var_keep_dic = []
            variables_to_restore = []
            for v in variables:
                if v.name.split ( ':' )[0] in var_keep_dic:
                    # print('Varibles restored: %s' % v.name)
                    variables_to_restore.append ( v )
                    # vis_var_keep_dic.append(v.name.split(':')[0])
                else:
                    # print('Unrestored Variables: %s' % v.name)
                    pass
            # print('Extra Variables in ckpt', set(var_keep_dic) - set(vis_var_keep_dic))

            if len ( variables_to_restore ) > 0:
                restorer = tf.train.Saver ( variables_to_restore )
                restorer.restore ( sess, model_path )
            else:
                print ( 'No variables in {} fits the network'.format ( model_path ) )

        tfconfig = tf.ConfigProto ( allow_soft_placement=True )
        tfconfig.gpu_options.allow_growth = True
        sess = tf.Session ( config=tfconfig )
        net = network_desp.Network ()
        inputs = net.get_inputs ()
        net.inference ( 'TEST', inputs )
        test_collect_dict = net.get_test_collection ()
        test_collect = [it for it in test_collect_dict.values ()]
        # The way in light-head is not work in my modified version since I rename variable names
        # saver = tf.train.Saver ()
        # saver.restore ( sess, model_file )
        # Use function in tf_cpn
        load_model ( sess, model_file )
        return partial ( sess.run, test_collect ), inputs

    def inference(self, val_func, inputs, data_dict):
        image = data_dict['data']
        ori_shape = image.shape

        if config.eval_resize == False:
            resized_img, scale = image, 1
        else:
            resized_img, scale = dataset.resize_img_by_short_and_max_size (
                image, config.eval_image_short_size, config.eval_image_max_size )
        height, width = resized_img.shape[0:2]

        resized_img = resized_img.astype ( np.float32 ) - config.image_mean
        resized_img = np.ascontiguousarray ( resized_img[:, :, [2, 1, 0]] )

        im_info = np.array (
            [[height, width, scale, ori_shape[0], ori_shape[1], 0]],
            dtype=np.float32 )

        feed_dict = {inputs[0]: resized_img[None, :, :, :], inputs[1]: im_info}

        _, scores, pred_boxes, rois = val_func ( feed_dict=feed_dict )

        boxes = rois[:, 1:5] / scale

        if cfg.TEST.BBOX_REG:
            pred_boxes = bbox_transform_inv ( boxes, pred_boxes )
            pred_boxes = clip_boxes ( pred_boxes, ori_shape )

        pred_boxes = pred_boxes.reshape ( -1, config.num_classes, 4 )
        result_boxes = []
        for j in range ( 1, config.num_classes ):
            inds = np.where ( scores[:, j] > config.test_cls_threshold )[0]
            cls_scores = scores[inds, j]
            cls_bboxes = pred_boxes[inds, j, :]
            cls_dets = np.hstack ( (cls_bboxes, cls_scores[:, np.newaxis]) ).astype (
                np.float32, copy=False )

            keep = nms ( cls_dets, config.test_nms )
            cls_dets = np.array ( cls_dets[keep, :], dtype=np.float, copy=False )
            for i in range ( cls_dets.shape[0] ):
                db = cls_dets[i, :]
                dbox = DetBox (
                    db[0], db[1], db[2] - db[0], db[3] - db[1],
                    tag=config.class_names[j], score=db[-1] )
                result_boxes.append ( dbox )
        if len ( result_boxes ) > config.test_max_boxes_per_image:
            result_boxes = sorted (
                result_boxes, reverse=True, key=lambda t_res: t_res.score ) \
                [:config.test_max_boxes_per_image]

        result_dict = data_dict.copy ()
        result_dict['result_boxes'] = result_boxes
        return result_dict

    def _make_result_dict(self, all_results):
        coco_records = []

        for result in all_results:
            result_boxes = result['result_boxes']
            if config.test_save_type == 'coco':
                image_id = int ( result['image_id'] )
                for rb in result_boxes:
                    if rb.tag == 'person':
                        record = {'image_id': image_id, 'category_id': config.datadb.classes_originID[rb.tag],
                                  'score': rb.score, 'bbox': [rb.x, rb.y, rb.w, rb.h], 'data': result['data']}
                        coco_records.append ( record )
            else:
                raise Exception (
                    "Unimplemented save type: " + str ( config.test_save_type ) )

        return coco_records

if __name__ == '__main__':
    detector = PersonDetector ( show_image=True )
    img = cv2.imread ( './img_000000.png' )
    print ( detector.detect ( img, 0 ) )