import numpy as np
import tensorflow as tf
from .bbox import *

class BBoxesLayer(object):

    def __init__(self, mod_tra=True, box_cls_num=None, img_shp=None):

        self.img_shp     = img_shp
        self.mod_tra     = mod_tra
        self.box_cls_num = box_cls_num
        self.box_pol_num = 300
        self.box_siz_min = 16
        self.box_prb_min = 0.5
        self.box_nms_pre = None
        self.box_nms_pst = 100   #200
        self.box_nms_max = 0.3   #0.2
        self.box_msk_min = 0.4
        self.box_msk_siz = [21, 21]
        self.meg_ovp_min = 0.5

    def generate_boxs_img(self, elems=None):
        #奇葩!和以前一点不一样!!! 这里看成是对ROI的分类、打分、mask预测，而不是对由ROI产生的BOX的，因此不对ROI做剔除!!!
        #按道理说，这么做比之前的做法合理!!!
        #(M, 4) #(M, C) #(M, 4) #(M, C, 21, 21, 2)
        rois, roi_prbs_pst, roi_prds_pst, roi_msks_pst, roi_num = elems
        rois         = rois        [0:roi_num]
        roi_prbs_pst = roi_prbs_pst[0:roi_num]
        roi_prds_pst = roi_prds_pst[0:roi_num]
        roi_msks_pst = roi_msks_pst[0:roi_num]

        roi_clss = tf.argmax(roi_prbs_pst, axis=-1)
        roi_clss = tf.cast(roi_clss, tf.int32)
        roi_prbs = tf.reduce_max(roi_prbs_pst, axis=-1)

        idxs     = tf.range(roi_num)
        idxs     = tf.stack([idxs, roi_clss], axis=-1)
        roi_msks = tf.gather_nd(roi_msks_pst, idxs) #(M, 21, 21, 2)
        roi_msks = roi_msks[..., 1]                 #(M, 21, 21)

        #roi_prds_pst = tf.gather_nd(roi_prds_pst, idxs) #如果roi的预测是定类的话要换上这段
        boxs     = bbox_transform_inv(rois, roi_prds_pst)
        boxs     = bbox_clip(boxs, [0.0, 0.0, self.img_shp[0]-1.0, self.img_shp[1]-1.0])
        #剔除过小的box
        idxs     = bbox_filter(boxs, self.box_siz_min)
        boxs     = tf.gather_nd(boxs, idxs)
        box_num  = tf.shape(boxs)[0]
        #padding
        paddings = [[0, self.box_pol_num-box_num], [0, 0]]
        boxs     = tf.pad(boxs,     paddings, "CONSTANT")
        paddings = [[0, self.box_pol_num-roi_num]]
        roi_clss = tf.pad(roi_clss, paddings, "CONSTANT")
        roi_prbs = tf.pad(roi_prbs, paddings, "CONSTANT")
        paddings = [[0, self.box_pol_num-roi_num], [0, 0], [0, 0]]
        roi_msks = tf.pad(roi_msks, paddings, "CONSTANT")
        return boxs, box_num, roi_clss, roi_prbs, roi_msks

    def generate_boxs(self, rois=None, roi_prbs_pst=None, roi_prds_pst=None, roi_msks_pst=None, roi_nums=None):

        elems = [rois, roi_prbs_pst, roi_prds_pst, roi_msks_pst, roi_nums]
        boxs, box_nums, roi_clss, roi_prbs, roi_msks = \
            tf.map_fn(self.generate_boxs_img, elems, dtype=(tf.float32, tf.int32, tf.int32, tf.float32, tf.float32),
                      parallel_iterations=10, back_prop=False, swap_memory=True, infer_shape=True)
        return boxs, box_nums, roi_clss, roi_prbs, roi_msks

    def concat_boxs_img(self, elems=None):

        rois, roi_clss, roi_prbs, roi_msks, roi_num, boxs, box_clss, box_prbs, box_msks, box_num = elems
        rois     = rois    [0:roi_num]                     #(M0, 4)
        roi_clss = roi_clss[0:roi_num]                     #(M0)
        roi_prbs = roi_prbs[0:roi_num]                     #(M0)
        roi_msks = roi_msks[0:roi_num]                     #(M0, 21, 21)
        boxs     = boxs    [0:box_num]                     #(M1, 4)
        box_clss = box_clss[0:box_num]                     #(M1)
        box_prbs = box_prbs[0:box_num]                     #(M1)
        box_msks = box_msks[0:box_num]                     #(M1, 21, 21)
        box_num  = roi_num + box_num
        boxs     = tf.concat([boxs, rois], axis=0)         #(M, 4)
        box_clss = tf.concat([box_clss, roi_clss], axis=0) #(M)
        box_prbs = tf.concat([box_prbs, roi_prbs], axis=0) #(M)
        box_msks = tf.concat([box_msks, roi_msks], axis=0) #(M, 21, 21)
        paddings = [[0, self.box_pol_num*2-box_num], [0, 0]]
        boxs     = tf.pad(boxs,     paddings, "CONSTANT")
        paddings = [[0, self.box_pol_num*2-box_num]]
        box_clss = tf.pad(box_clss, paddings, "CONSTANT")
        box_prbs = tf.pad(box_prbs, paddings, "CONSTANT")
        paddings = [[0, self.box_pol_num*2-box_num], [0, 0], [0, 0]]
        box_msks = tf.pad(box_msks, paddings, "CONSTANT")
        return boxs, box_clss, box_prbs, box_msks, box_num

    def concat_boxs(self, rois, roi_clss, roi_prbs, roi_msks, roi_nums, boxs, box_clss, box_prbs, box_msks, box_nums):

        elems = [rois, roi_clss, roi_prbs, roi_msks, roi_nums, boxs, box_clss, box_prbs, box_msks, box_nums]
        boxs, box_clss, box_prbs, box_msks, box_nums = \
            tf.map_fn(self.concat_boxs_img, elems, dtype=(tf.float32, tf.int32, tf.float32, tf.float32, tf.int32),
                      parallel_iterations=10, back_prop=False, swap_memory=True, infer_shape=True)
        return boxs, box_clss, box_prbs, box_msks, box_nums

    def merge_boxs_img(self, elems=None):

        boxs, box_clss, box_prbs, box_msks, box_num = elems
        boxs     = boxs    [0:box_num] #(M, 4)
        box_clss = box_clss[0:box_num] #(M)
        box_prbs = box_prbs[0:box_num] #(M)
        box_msks = box_msks[0:box_num] #(M, 21, 21)

        #设置一个box索引，避免大量的gather操作(prds、msks)，节省内存，提升速度
        kep_idxs = tf.range(box_num)
        #剔除背景box
        idxs     = tf.where(box_clss>0)
        kep_idxs = tf.gather_nd(kep_idxs, idxs)
        box_clss = tf.gather_nd(box_clss, idxs)
        box_prbs = tf.gather_nd(box_prbs, idxs)
        #剔除得分较低的box
        if self.box_prb_min is not None:
            idxs     = tf.where(box_prbs>=self.box_prb_min)
            kep_idxs = tf.gather_nd(kep_idxs, idxs)
            box_clss = tf.gather_nd(box_clss, idxs)
            box_prbs = tf.gather_nd(box_prbs, idxs)
        #进一步剔除过多的box
        if self.box_nms_pre is not None:
            box_nms_pre = tf.minimum(self.box_nms_pre, tf.shape(kep_idxs)[0])
            box_prbs, idxs = tf.nn.top_k(box_prbs, k=box_nms_pre, sorted=True)
            kep_idxs = tf.gather(kep_idxs, idxs)
            box_clss = tf.gather(box_clss, idxs)
        #根据kep_idxs进行剩余的gather操作
        boxs         = tf.gather(boxs,     kep_idxs)
        box_msks     = tf.gather(box_msks, kep_idxs)
        #做逐类的nms
        boxs_tmp     = tf.zeros(dtype=tf.float32, shape=[0, 4])
        box_clss_kep = tf.zeros(dtype=tf.int32,   shape=[0])
        box_prbs_kep = tf.zeros(dtype=tf.float32, shape=[0])
        def cond0(i, boxs, box_clss, box_prbs, boxs_tmp, box_clss_kep, box_prbs_kep):
            c = tf.less(i, self.box_cls_num)
            return c

        def body0(i, boxs, box_clss, box_prbs, boxs_tmp, box_clss_kep, box_prbs_kep):
            #选出对应类的rois
            idxs         = tf.where(tf.equal(box_clss, i))
            boxs_cls     = tf.gather_nd(boxs,     idxs)
            box_clss_cls = tf.gather_nd(box_clss, idxs)
            box_prbs_cls = tf.gather_nd(box_prbs, idxs)
            #进行非极大值抑制操作
            idxs         = tf.image.non_max_suppression(boxs_cls, box_prbs_cls, self.box_nms_pst, self.box_nms_max)
            boxs_cls     = tf.gather(boxs_cls,     idxs)
            box_clss_cls = tf.gather(box_clss_cls, idxs)
            box_prbs_cls = tf.gather(box_prbs_cls, idxs)
            #保存结果
            boxs_tmp     = tf.concat([boxs_tmp,     boxs_cls    ], axis=0)
            box_clss_kep = tf.concat([box_clss_kep, box_clss_cls], axis=0)
            box_prbs_kep = tf.concat([box_prbs_kep, box_prbs_cls], axis=0)
            return [i+1, boxs, box_clss, box_prbs, boxs_tmp, box_clss_kep, box_prbs_kep]

        i = tf.constant(1) #要剔除背景类
        [i, boxs, box_clss, box_prbs, boxs_tmp, box_clss_kep, box_prbs_kep] = \
            tf.while_loop(cond0, body0, \
                          loop_vars=[i, boxs, box_clss, box_prbs, boxs_tmp, box_clss_kep, box_prbs_kep], \
                          shape_invariants=[i.get_shape(), boxs.get_shape(), box_clss.get_shape(), box_prbs.get_shape(), \
                                            tf.TensorShape([None, 4]), tf.TensorShape([None]), tf.TensorShape([None])], \
                          parallel_iterations=10, back_prop=False, swap_memory=True)

        box_num      = tf.minimum(self.box_nms_pst, tf.shape(boxs_tmp)[0])
        box_prbs_kep, idxs = tf.nn.top_k(box_prbs_kep, k=box_num, sorted=False)
        boxs_tmp     = tf.gather(boxs_tmp,     idxs)
        box_clss_kep = tf.gather(box_clss_kep, idxs)

        #融合mask
        boxs_kep     = tf.zeros(dtype=tf.float32, shape=[0, 4])
        box_msks_kep = tf.zeros(dtype=tf.float32, shape=[0]+self.box_msk_siz)
        box_idxs_kep = tf.zeros(dtype=tf.int64,   shape=[0, 1])  #并行计算要保证对应关系
        def cond1(i, boxs, box_clss, box_prbs, box_msks, boxs_tmp, box_clss_kep, boxs_kep, box_msks_kep, box_idxs_kep):
            c = tf.less(i, self.box_cls_num)
            return c

        def body1(i, boxs, box_clss, box_prbs, box_msks, boxs_tmp, box_clss_kep, boxs_kep, box_msks_kep, box_idxs_kep):
            #选出对应类的boxs
            idxs         = tf.where(tf.equal(box_clss, i))
            boxs_cls     = tf.gather_nd(boxs,     idxs)
            box_prbs_cls = tf.gather_nd(box_prbs, idxs)
            box_msks_cls = tf.gather_nd(box_msks, idxs)

            box_idxs_cls = tf.where(tf.equal(box_clss_kep, i))   #选出对应类别的box和其相应的idxs
            boxs_run     = tf.gather_nd(boxs_tmp, box_idxs_cls)

            boxs_hld     = tf.zeros(dtype=tf.float32, shape=[0, 4])
            box_msks_hld = tf.zeros(dtype=tf.float32, shape=[0]+self.box_msk_siz)
            box_idxs_hld = tf.zeros(dtype=tf.int32,   shape=[0]) #并行计算要保证对应关系
            def cond(j, boxs_run, boxs_cls, box_prbs_cls, box_msks_cls, boxs_hld, box_msks_hld, box_idxs_hld):
                c = tf.less(j, tf.shape(boxs_run)[0])
                return c

            def body(j, boxs_run, boxs_cls, box_prbs_cls, box_msks_cls, boxs_hld, box_msks_hld, box_idxs_hld):
                box_run      = boxs_run[j][tf.newaxis, :]
                box_ovps     = bbox_overlaps(boxs_cls, box_run)     #(M0, 1)
                box_ovps     = tf.squeeze(box_ovps, axis=[1])       #(M0)
                idxs         = tf.where(box_ovps>=self.meg_ovp_min)
                boxs_meg     = tf.gather_nd(boxs_cls,     idxs)     #(M1, 4)
                box_prbs_meg = tf.gather_nd(box_prbs_cls, idxs)     #(M1)
                box_msks_meg = tf.gather_nd(box_msks_cls, idxs)     #(M1, 21, 21)
                '''
                assert_op = tf.Assert(tf.size(box_msks_meg)>0, [j, tf.shape(boxs_run)[0], box_ovps, box_run, boxs_cls, \
                                                                idxs, boxs_meg, box_prbs_meg, box_msks_meg], summarize=100)
                with tf.control_dependencies([assert_op]):
                    box_msks_meg = tf.identity(box_msks_meg)
                '''
                box_wgts     = box_prbs_meg / tf.reduce_sum(box_prbs_meg)
                #boxs_meg    = bbox_clip(boxs_meg, [0, 0, self.img_shp[0]-1.0, self.img_shp[1]-1.0]) #之前还原出boxs时已经clip了
                boxs_meg     = tf.cast(tf.round(boxs_meg), dtype=tf.int32)
                box_ymns, box_xmns, box_ymxs, box_xmxs = tf.split(boxs_meg, 4, axis=-1)
                box_hgts     = box_ymxs - box_ymns + 1
                box_wdhs     = box_xmxs - box_xmns + 1
                box_lehs     = tf.concat([box_hgts, box_wdhs], axis=-1)
                box_ymn      = tf.reduce_min(box_ymns)
                box_xmn      = tf.reduce_min(box_xmns)
                box_ymx      = tf.reduce_max(box_ymxs)
                box_xmx      = tf.reduce_max(box_xmxs)
                pads_hgt_fnt = box_ymns - box_ymn
                pads_wdh_fnt = box_xmns - box_xmn
                pads_hgt_bak = box_ymx - box_ymxs
                pads_wdh_bak = box_xmx - box_xmxs
                pads_hgt     = tf.concat([pads_hgt_fnt, pads_hgt_bak], axis=-1) #(M1, 2)
                pads_wdh     = tf.concat([pads_wdh_fnt, pads_wdh_bak], axis=-1) #(M1, 2)
                box_pads     = tf.stack([pads_hgt, pads_wdh], axis=1)           #(M1, 2, 2)

                def resize_mask(elems=None):
                    box_msk_meg, box_wgt, box_leh, box_pad = elems
                    box_msk_meg = tf.expand_dims(box_msk_meg, axis=-1)
                    box_msk_meg = tf.image.resize_images(box_msk_meg, box_leh, \
                                                         method=tf.image.ResizeMethod.BILINEAR, align_corners=False)
                    box_msk_meg = tf.squeeze(box_msk_meg, axis=[-1])
                    box_msk_meg = tf.cast(box_msk_meg>=self.box_msk_min, dtype=tf.float32)
                    box_msk_meg = box_msk_meg * box_wgt
                    box_msk_meg = tf.pad(box_msk_meg, box_pad, "CONSTANT")
                    return box_msk_meg
                #当没有box时，tf.map_fn不会拆开，所以tf.image.resize_images是安全的
                elems        = [box_msks_meg, box_wgts, box_lehs, box_pads]
                box_msks_meg = tf.map_fn(resize_mask, elems, dtype=tf.float32,
                                         parallel_iterations=10, back_prop=False, swap_memory=True, infer_shape=True)

                box_msk_meg  = tf.reduce_sum(box_msks_meg, axis=0)
                box_msk_idxs = tf.where(box_msk_meg>=self.box_msk_min)
                box_crd_min  = tf.maximum(tf.cast(tf.reduce_min(box_msk_idxs, axis=0), dtype=tf.int32), 0)
                box_crd_max  = tf.cast(tf.reduce_max(box_msk_idxs, axis=0), dtype=tf.int32)
                box_msk_meg  = box_msk_meg[box_crd_min[0]:box_crd_max[0]+1, box_crd_min[1]:box_crd_max[1]+1]
                box_msk_meg  = tf.expand_dims(box_msk_meg, axis=-1)
                box_msk_meg  = tf.image.resize_images(box_msk_meg, self.box_msk_siz, \
                                                      method=tf.image.ResizeMethod.BILINEAR, align_corners=False)
                box_msk_meg  = tf.squeeze(box_msk_meg, axis=[-1])
                box_msk_meg  = tf.expand_dims(box_msk_meg, axis=0)
                box_meg      = tf.concat([box_crd_min, box_crd_max], axis=0) + \
                               tf.stack([box_ymn, box_xmn, box_ymn, box_xmn], axis=0)
                box_meg      = tf.cast(box_meg, dtype=tf.float32)
                box_meg      = tf.expand_dims(box_meg, axis=0)
                box_idx_meg  = tf.expand_dims(j,       axis=0)
                #保存结果
                boxs_hld     = tf.concat([boxs_hld,     box_meg    ], axis=0)
                box_msks_hld = tf.concat([box_msks_hld, box_msk_meg], axis=0)
                box_idxs_hld = tf.concat([box_idxs_hld, box_idx_meg], axis=0)
                return [j+1, boxs_run, boxs_cls, box_prbs_cls, box_msks_cls, boxs_hld, box_msks_hld, box_idxs_hld]

            j = tf.constant(0)
            [j, boxs_run, boxs_cls, box_prbs_cls, box_msks_cls, boxs_hld, box_msks_hld, box_idxs_hld] = \
                tf.while_loop(cond, body, \
                              loop_vars=[j, boxs_run, boxs_cls, box_prbs_cls, box_msks_cls, boxs_hld, box_msks_hld, box_idxs_hld], \
                              shape_invariants=[j.get_shape(), boxs_run.get_shape(), boxs_cls.get_shape(), \
                                                box_prbs_cls.get_shape(), box_msks_cls.get_shape(), tf.TensorShape([None, 4]), \
                                                tf.TensorShape([None]+self.box_msk_siz), tf.TensorShape([None])], \
                              parallel_iterations=10, back_prop=False, swap_memory=True)
            box_idxs_hld = tf.gather(box_idxs_cls, box_idxs_hld)
            #保存结果
            boxs_kep     = tf.concat([boxs_kep,     boxs_hld    ], axis=0)
            box_msks_kep = tf.concat([box_msks_kep, box_msks_hld], axis=0)
            box_idxs_kep = tf.concat([box_idxs_kep, box_idxs_hld], axis=0)
            return [i+1, boxs, box_clss, box_prbs, box_msks, boxs_tmp, box_clss_kep, boxs_kep, box_msks_kep, box_idxs_kep]

        i = tf.constant(1) #要剔除背景类
        [i, boxs, box_clss, box_prbs, box_msks, boxs_tmp, box_clss_kep, boxs_kep, box_msks_kep, box_idxs_kep] = \
            tf.while_loop(cond1, body1, loop_vars=\
                          [i, boxs, box_clss, box_prbs, box_msks, boxs_tmp, box_clss_kep, boxs_kep, box_msks_kep, box_idxs_kep], \
                          shape_invariants=[i.get_shape(), boxs.get_shape(), box_clss.get_shape(), \
                                            box_prbs.get_shape(), box_msks.get_shape(), boxs_tmp.get_shape(), \
                                            box_clss_kep.get_shape(), tf.TensorShape([None, 4]), \
                                            tf.TensorShape([None]+self.box_msk_siz), tf.TensorShape([None, 1])], \
                          parallel_iterations=10, back_prop=False, swap_memory=True)
        #boxs_kep, box_msks_kep, box_idxs_kep
        box_clss_kep = tf.gather_nd(box_clss_kep, box_idxs_kep)
        box_prbs_kep = tf.gather_nd(box_prbs_kep, box_idxs_kep)
        box_num      = tf.shape(boxs_kep)[0]
        paddings     = [[0, self.box_nms_pst-box_num], [0, 0]]
        boxs_kep     = tf.pad(boxs_kep, paddings,     "CONSTANT")
        paddings     = [[0, self.box_nms_pst-box_num]]
        box_clss_kep = tf.pad(box_clss_kep, paddings, "CONSTANT")
        box_prbs_kep = tf.pad(box_prbs_kep, paddings, "CONSTANT")
        paddings     = [[0, self.box_nms_pst-box_num], [0, 0], [0, 0]]
        box_msks_kep = tf.pad(box_msks_kep, paddings, "CONSTANT")
        return boxs_kep, box_clss_kep, box_prbs_kep, box_msks_kep, box_num

    def merge_boxs(self, boxs=None, box_clss=None, box_prbs=None, box_msks=None, box_nums=None):
        elems = [boxs, box_clss, box_prbs, box_msks, box_nums]
        boxs, box_clss, box_prbs, box_msks, box_nums = \
            tf.map_fn(self.merge_boxs_img, elems, dtype=(tf.float32, tf.int32, tf.float32, tf.float32, tf.int32),
                      parallel_iterations=10, back_prop=False, swap_memory=True, infer_shape=True)
        return boxs, box_clss, box_prbs, box_msks, box_nums


import os, sys, math
import numpy as np
import torch.nn as nn, torch.nn.functional as F
import torchvision.models as models

class ResNet50(nn.Module):
    def __init__(self, num_classes=1000, pretrained=True):
        super().__init__()
        print("creating model " + self.__class__.__name__)
        self.resnet50 = models.resnet50(pretrained=pretrained)

        # modify the initial conv layer
        conv1_layer = self.resnet50.conv1
        conv1_weights = conv1_layer.weight.data.clone()

        self.resnet50.conv1 = nn.Conv2d(
            2,
            out_channels=conv1_layer.out_channels,
            kernel_size=conv1_layer.kernel_size,
            padding=conv1_layer.padding,
            stride=conv1_layer.stride,
            bias=not (conv1_layer is None),
        )

        mean_conv1_weights = conv1_weights.mean(dim=1, keepdim=True).repeat(1, 2, 1, 1)
        self.resnet50.conv1.weight.data.copy_(mean_conv1_weights)

        self.feature_dim = 2048
        self.fc = nn.Linear(self.feature_dim, num_classes)

    def forward(self, x):
        x = self.resnet50.conv1(x)
        x = self.resnet50.bn1(x)
        x = self.resnet50.relu(x)
        x = self.resnet50.maxpool(x)

        x = self.resnet50.layer1(x)
        x = self.resnet50.layer2(x)
        x = self.resnet50.layer3(x)
        x = self.resnet50.layer4(x)

        x = F.adaptive_avg_pool2d(x, output_size=1)
        features = x.view(x.size(0), -1)
        classifier_output = self.fc(features)

        return features, classifier_output


import os, json, base64, cv2, glob
import numpy as np
import matplotlib.pyplot as plt
from coco import CocoConfig
from Mask.config import Config
import Mask.utils as utils
import Mask.model as modellib
import Mask.visualize as visualize
from convert_file import load_image

def init():
    np.set_printoptions(threshold=np.inf)

def run(input_df):
    data = json.loads(input_df)
    im = load_image(image64=data)

    config = CocoConfig()

    model = modellib.MaskRCNN(mode="inference", model_dir="./models/", config=config)
    model.load_weights(filepath="./models/mask_rcnn_moles_0090.h5", by_name=True)

    class_names = ["BG", "malignant", "benign"]

    # predict the mask, bounding box and class of the image
    r = model.detect([im])[0]
    prediction = None
    for idx, val in enumerate(class_names):
        if idx == r["class_ids"]:
            prediction = val
            print(val)
        else:
            continue
    return prediction


"""Kubo thermal conductivity base class."""
# Copyright (C) 2022 Atsushi Togo
# All rights reserved.
#
# This file is part of phono3py.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in
#   the documentation and/or other materials provided with the
#   distribution.
#
# * Neither the name of the phonopy project nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import numpy as np
from phonopy.units import Kb, THzToEv

from phono3py.phonon.group_velocity_matrix import GroupVelocityMatrix
from phono3py.phonon.heat_capacity_matrix import mode_cv_matrix

class ConductivityKuboMixIn:
    """Thermal conductivity mix-in for group velocity matrix."""

    @property
    def kappa(self):
        """Return kappa."""
        return self._kappa

    @property
    def mode_kappa_mat(self):
        """Return mode_kappa_mat."""
        return self._mode_kappa_mat

    def _init_velocity(self, gv_delta_q):
        self._velocity_obj = GroupVelocityMatrix(
            self._pp.dynamical_matrix,
            q_length=gv_delta_q,
            symmetry=self._pp.primitive_symmetry,
            frequency_factor_to_THz=self._pp.frequency_factor_to_THz,
        )

    def _set_cv(self, i_gp, i_data):
        """Set heat capacity matrix.

        x=freq/T has to be small enough to avoid overflow of exp(x).
        x < 100 is the hard-corded criterion.
        Otherwise just set 0.

        """
        irgp = self._grid_points[i_gp]
        freqs = self._frequencies[irgp] * THzToEv
        cutoff = self._pp.cutoff_frequency * THzToEv

        for i_temp, temp in enumerate(self._temperatures):
            if (freqs / (temp * Kb) > 100).any():
                continue
            cvm = mode_cv_matrix(temp, freqs, cutoff=cutoff)
            self._cv_mat[i_temp, i_data] = cvm[self._pp.band_indices, :]

    def _set_velocities(self, i_gp, i_data):
        gvm, gv = self._get_gv_matrix(i_gp)
        gvm_sum2, kstar_order = self._get_gvm_by_gvm(i_gp)
        self._num_sampling_grid_points += kstar_order
        self._gv_mat[i_data] = gvm
        self._gv_mat_sum2[i_data] = gvm_sum2
        self._gv[i_data] = gv

    def _get_gv_matrix(self, i_gp):
        """Get group velocity matrix.

        Returns
        -------
        gv_mat : ndarray
            Group velocity matrix at grid point of `i_gp`.
            shape=(num_band0, num_band, 3), dtype=complex
        gv : ndarray
            Group velocity at grid point of `i_gp`.
            shape=(num_band0, 3), dtype=double

        """
        irgp = self._grid_points[i_gp]
        self._velocity_obj.run([self._get_qpoint_from_gp_index(irgp)])
        gvm = np.zeros(self._gv_mat.shape[1:], dtype=self._complex_dtype, order="C")
        gv = np.zeros(self._gv.shape[1:], dtype="double", order="C")
        for i in range(3):
            gvm[:, :, i] = self._velocity_obj.group_velocity_matrices[
                0, i, self._pp.band_indices, :
            ]
            gv[:, i] = (
                self._velocity_obj.group_velocity_matrices[0, i]
                .diagonal()[self._pp.band_indices]
                .real
            )
        return gvm, gv

    def _get_gvm_by_gvm(self, i_gp):
        r"""Compute sum of gvm over k-star.

        q is an irreducible q-point. k_q is an arm of k-star of q.
        R are the elements of the reciprocal point group.
        m is the multiplicity of q, i.e., order of site-point symmetry.

        \sum_{k_q} v^\alpha_{kjj'} v^\beta_{kj'j}
        = 1/m \sum_{R} v^\alpha_{Rq jj'} v^\beta_{Rq j'j}

        Returns
        -------
        gvm_sum2 : ndarray
            sum of gvm x gvm over kstar arms.
            shape=(num_band0, num_band, 6).
            The last 6 gives 3x3 tensor element positions, xx, yy, zz, yz, xz, xy.
        order_kstar : int
            Number of kstar arms.

        """
        multi = self._get_multiplicity_at_q(i_gp)
        q = self._get_qpoint_from_gp_index(self._grid_points[i_gp])
        qpoints = [np.dot(r, q) for r in self._point_operations]
        self._velocity_obj.run(qpoints)

        gvm_sum2 = np.zeros(
            self._gv_mat.shape[1:-1] + (6,), dtype=self._complex_dtype, order="C"
        )
        for i_pair, (a, b) in enumerate(
            ([0, 0], [1, 1], [2, 2], [1, 2], [0, 2], [0, 1])
        ):
            for gvm in self._velocity_obj.group_velocity_matrices:
                gvm_by_gvm = np.multiply(gvm[a], gvm[b].T)
                gvm_sum2[:, :, i_pair] += gvm_by_gvm[self._pp.band_indices, :]
        gvm_sum2 /= multi

        return gvm_sum2, self._get_kstar_order(i_gp, multi)


#!/usr/bin/env python
import torch
import models
import utils
from torch.utils.data import DataLoader, TensorDataset
from torch import nn
import time
import numpy as np

class PGLearner(object):
    """ Reinforcement learning framework with policy gradient. This class is the base structure for all
        policy gradient-based  deep reinforcement learning models.

    Arguments:

        agent (models.Generator): The agent which generates the desired molecules

        env (utils.Env): The environment which provides the reward and judge
                                 if the genrated molecule is valid and desired.

        prior: The auxiliary model which is defined differently in each methods.
    """
    def __init__(self, agent, env, prior=None, memory=None, mean_func='geometric'):
        self.replay = 10
        self.agent = agent
        self.prior = prior
        self.batch_size = 64  # * 4
        self.n_samples = 128  # * 8
        self.env = env
        self.epsilon = 1e-3
        self.penalty = 0
        self.scheme = 'PR'
        self.out = None
        self.memory = memory
        # mean_func: which function to use for averaging: 'arithmetic' or 'geometric'
        self.mean_func = mean_func

    def policy_gradient(self):
        pass

    def fit(self):
        best = 0
        last_save = 0
        log = open(self.out + '.log', 'w')
        for epoch in range(1000):
            print('\n----------\nEPOCH %d\n----------' % epoch)
            self.policy_gradient()
            seqs = self.agent.sample(self.n_samples)
            ix = utils.unique(seqs)
            smiles = [self.agent.voc.decode(s) for s in seqs[ix]]
            scores = self.env(smiles, is_smiles=True)

            desire = (scores.DESIRE).sum() / self.n_samples
            score = scores[self.env.keys].values.mean()
            valid = scores.VALID.mean()

            if best <= score:
                torch.save(self.agent.state_dict(), self.out + '.pkg')
                best = score
                last_save = epoch

            print("Epoch: %d average: %.4f valid: %.4f unique: %.4f" %
                  (epoch, score, valid, desire), file=log)
            for i, smile in enumerate(smiles):
                score = "\t".join(['%0.3f' % s for s in scores.values[i]])
                print('%s\t%s' % (score, smile), file=log)
            if epoch - last_save > 100:
                break
        for param_group in self.agent.optim.param_groups:
            param_group['lr'] *= (1 - 0.01)
        log.close()

class Reinvent(PGLearner):
    """ REINVENT algorithm

    Reference: Olivecrona, M., Blaschke, T., Engkvist, O. et al. Molecular de-novo design
               through deep reinforcement learning. J Cheminform 9, 48 (2017).
               https://doi.org/10.1186/s13321-017-0235-x

    Arguments:

        agent (models.Generator): The agent network which is constructed by deep learning model
                                   and generates the desired molecules.

        env (utils.Env): The environment which provides the reward and judge
                                 if the genrated molecule is valid and desired.

        prior (models.Generator): The prior network which is constructed by deep learning model
                                   and ensure the agent to generate molecules with correct grammar.
    """
    def __init__(self, agent, env, prior, epsilon=60):
        super(Reinvent, self).__init__(agent, env, prior)
        for param in self.prior.parameters():
            param.requires_grad = False
        # self.agent.optim.lr = 0.0005
        self.epsilon = epsilon

    def policy_gradient(self):
        seqs = []
        for _ in range(self.replay):
            seq = self.agent.sample(self.batch_size)
            seqs.append(seq)
        seqs = torch.cat(seqs, dim=0)
        ix = utils.unique(seqs)
        seqs = seqs[ix]
        smiles = [self.agent.voc.decode(s) for s in seqs]

        scores = self.env.calc_reward(smiles, self.scheme)[:, 0]
        ds = TensorDataset(seqs, torch.Tensor(scores-self.beta).to(utils.dev))
        loader = DataLoader(ds, batch_size=self.n_samples, shuffle=True)

        for seq, score in loader:
            # Calculate gradients and make an update to the network weights
            self.agent.optim.zero_grad()
            prior_likelihood = self.prior.likelihood(seq).sum(dim=1)
            agent_likelihood = self.agent.likelihood(seq).sum(dim=1)
            augmented_likelihood = prior_likelihood + self.epsilon * score
            loss = torch.pow((augmented_likelihood - agent_likelihood), 2)
            # Calculate loss
            loss = loss.mean()

            # Add regularizer that penalizes high likelihood for the entire sequence
            loss_p = - (1 / agent_likelihood).mean()
            loss += 5 * 1e3 * loss_p
            loss.backward()
            self.agent.optim.step()

class DrugEx(PGLearner):
    """ DrugEx algorithm (version 1.0)

    Reference: Liu, X., Ye, K., van Vlijmen, H.W.T. et al. An exploration strategy improves the diversity
               of de novo ligands using deep reinforcement learning: a case for the adenosine A2A receptor.
               J Cheminform 11, 35 (2019).
               https://doi.org/10.1186/s13321-019-0355-6

    Arguments:

        agent (models.Generator): The agent network which is constructed by deep learning model
                                   and generates the desired molecules.

        env (utils.Env): The environment which provides the reward and judge
                                 if the genrated molecule is valid and desired.

        prior (models.Generator): The pre-trained network which is constructed by deep learning model
                                   and ensure the agent to explore the approriate chemical space.
    """
    def __init__(self, agent, env, prior=None, memory=None, beta=0.1):
        super(DrugEx, self).__init__(agent, env, prior, memory=memory)
        self.beta = beta

    def policy_gradient(self):
        seqs = []
        for _ in range(self.replay):
            seq = self.agent.evolve1(self.batch_size, epsilon=self.epsilon, mutate=self.prior)
            seqs.append(seq)
        seqs = torch.cat(seqs, dim=0)
        ix = utils.unique(seqs)
        seqs = seqs[ix]
        smiles = [self.agent.voc.decode(s) for s in seqs]

        scores = self.env.calc_reward(smiles, self.scheme)
        ds = TensorDataset(seqs, torch.Tensor(scores-self.beta).to(utils.dev))
        loader = DataLoader(ds, batch_size=self.n_samples, shuffle=True)
        self.agent.PGLoss(loader)

class Organic(PGLearner):
    """ ORGANIC algorithm

    Reference: Sanchez-Lengeling B, Outeiral C, Guimaraes GL, Aspuru-Guzik A (2017)
               Optimizing distributions over molecular space. An Objective-Reinforced
               Generative Adversarial Network for Inverse-design Chemistry (ORGANIC)
               https://doi.org/10.26434/chemrxiv.5309668.v3

    Arguments:

        agent (models.Generator): The agent network which is constructed by deep learning model
                                   and generates the desired molecules.

        env (utils.Env): The environment which provides the reward and judge
                                 if the genrated molecule is valid and desired.

        prior (models.Generator): The discriminator which is constrcuted by deep learning model and
                                   judge if the generated molecule is similar to the real molecule.
    """
    def __init__(self, agent, env, prior, real=None, memory=None):
        super(Organic, self).__init__(agent, env, prior, memory=memory)
        self.epsilon = 0.5
        self.loader = DataLoader(real, batch_size=self.batch_size, shuffle=True, drop_last=True,
                                 collate_fn=real.collate_fn)
        for param_group in self.prior.optim.param_groups:
            param_group['lr'] = 1e-3

    def policy_gradient(self):
        seqs = []
        for _ in range(self.replay):
            seq = self.agent.sample(self.batch_size)
            seqs.append(seq)
        seqs = torch.cat(seqs, dim=0)
        ix = utils.unique(seqs)
        seqs = seqs[ix]
        smiles = [self.agent.voc.decode(s) for s in seqs]

        scores = self.env.calc_reward(smiles, self.scheme)
        scores = torch.Tensor(scores).to(utils.dev)
        scores = self.epsilon * scores + (1 - self.epsilon) * self.prior(seqs).data

        ds = TensorDataset(seqs, scores)
        loader = DataLoader(ds, batch_size=self.n_samples, shuffle=True)
        self.agent.PGLoss(loader)
        self.Train_dis_BCE(epochs=1)

    def Train_dis_BCE(self, epochs=1, out=None):
        best_loss = np.Inf
        criterion = nn.BCELoss()
        log = open(out + '.log', 'w') if out is not None else None
        last_save = 0
        for epoch in range(epochs):
            t0 = time.time()
            for i, real in enumerate(self.loader):
                size = len(real)
                fake = self.agent.sample(size)
                data = torch.cat([fake, real.to(utils.dev)], dim=0)
                label = torch.cat([torch.zeros(size, 1), torch.ones(size, 1)]).to(utils.dev)
                self.prior.optim.zero_grad()
                label_ = self.prior(data)
                loss = criterion(label_, label)
                loss.backward()
                self.prior.optim.step()
                if i % 10 == 0 and i != 0:
                    for param_group in self.prior.optim.param_groups:
                        param_group['lr'] *= (1 - 0.03)
            if out is None: continue
            print('[Epoch: %d/%d] %.1fs loss: %f' % (
                epoch, epochs, time.time() - t0, loss.item()), file=log)
            if loss.item() < best_loss:
                print('[Performance] loss is improved from %f to %f, Save model to %s' %
                      (best_loss, loss, out + '.pkg'), file=log)
                torch.save(self.prior.state_dict(), out + ".pkg")
                best_loss = loss.item()
                last_save = epoch
            else:
                print('[Performance] loss  is not improved.', file=log)
            if epoch - last_save > 100:
                break
        return loss.item()

class Evolve(PGLearner):
    """ DrugEx algorithm (version 2.0)

    Reference: Liu, X., Ye, K., van Vlijmen, H.W.T. et al. DrugEx v2: De Novo Design of Drug Molecule by
               Pareto-based Multi-Objective Reinforcement Learning in Polypharmacology.
               J Cheminform (2021). https://doi.org/10.1186/s13321-019-0355-6

    Arguments:

        agent (models.Generator): The agent network which is constructed by deep learning model
                                   and generates the desired molecules.

        env (utils.Env): The environment which provides the reward and judge
                                 if the genrated molecule is valid and desired.

        prior (models.Generator): The pre-trained network which is constructed by deep learning model
                                   and ensure the agent to explore the approriate chemical space.
    """
    def __init__(self, agent, env, prior=None, crover=None, mean_func='geometric', memory=None):
        super(Evolve, self).__init__(agent, env, prior, mean_func=mean_func, memory=memory)
        self.crover = crover

    def policy_gradient(self, crover=None, memory=None, epsilon=None):
        seqs = []
        start = time.time()
        for _ in range(self.replay):
            seq = self.agent.evolve1(self.batch_size, epsilon=epsilon, crover=crover, mutate=self.prior)
            seqs.append(seq)
        t1 = time.time()
        seqs = torch.cat(seqs, dim=0)
        if memory is not None:
            mems = [memory, seqs]
            seqs = torch.cat(mems)
        smiles = np.array([self.agent.voc.decode(s) for s in seqs])
        # smiles = np.array(utils.canonicalize_list(smiles))
        ix = utils.unique(np.array([[s] for s in smiles]))
        smiles = smiles[ix]
        seqs = seqs[torch.LongTensor(ix).to(utils.dev)]

        scores = self.env.calc_reward(smiles, self.scheme)
        if memory is not None:
            scores[:len(memory), 0] = 1
            ix = scores[:, 0].argsort()[-self.batch_size * 4:]
            seqs, scores = seqs[ix, :], scores[ix, :]
        t2 = time.time()
        ds = TensorDataset(seqs, torch.Tensor(scores).to(utils.dev))
        loader = DataLoader(ds, batch_size=self.n_samples, shuffle=True)

        self.agent.PGLoss(loader)
        t3 = time.time()
        print(t1 - start, t2-t1, t3-t2)

    def fit(self):
        best = 0
        log = open(self.out + '.log', 'w')
        last_smiles = []
        last_scores = []
        interval = 250
        last_save = -1

        for epoch in range(10000):
            print('\n----------\nEPOCH %d\n----------' % epoch)
            if epoch < interval and self.memory is not None:
                self.policy_gradient(crover=None, memory=self.memory, epsilon=1e-1)
            else:
                self.policy_gradient(crover=self.crover, epsilon=self.epsilon)
            seqs = self.agent.sample(self.n_samples)
            smiles = [self.agent.voc.decode(s) for s in seqs]
            smiles = np.array(utils.canonicalize_list(smiles))
            ix = utils.unique(np.array([[s] for s in smiles]))
            smiles = smiles[ix]
            scores = self.env(smiles, is_smiles=True)

            desire = (scores.DESIRE).sum() / self.n_samples
            if self.mean_func == 'arithmetic':
                score = scores[self.env.keys].values.sum() / self.n_samples / len(self.env.keys)
            else:
                score = scores[self.env.keys].values.prod(axis=1) ** (1.0 / len(self.env.keys))
                score = score.sum() / self.n_samples
            valid = scores.VALID.sum() / self.n_samples

            print("Epoch: %d average: %.4f valid: %.4f unique: %.4f" %
                  (epoch, score, valid, desire), file=log)
            if best < score:
                torch.save(self.agent.state_dict(), self.out + '.pkg')
                best = score
                last_smiles = smiles
                last_scores = scores
                last_save = epoch

            if epoch % interval == 0 and epoch != 0:
                for i, smile in enumerate(last_smiles):
                    score = "\t".join(['%.3f' % s for s in last_scores.values[i]])
                    print('%s\t%s' % (score, smile), file=log)
                self.agent.load_state_dict(torch.load(self.out + '.pkg'))
                self.crover.load_state_dict(torch.load(self.out + '.pkg'))
            if epoch - last_save > interval: break
        log.close()


# Authors: Stephane Gaiffas <stephane.gaiffas@gmail.com>
# License: BSD 3 clause

"""
This module contains all required container types for nodes (including the
node numpy dtypes and several functions operating locally in nodes.
"""

from math import log
import numpy as np
from numba import from_dtype, jit, boolean, uint8, intp, uintp, float32, void
from numba.experimental import jitclass

from ._utils import get_type, sample_without_replacement
from ._tree_context import TreeClassifierContextType, TreeRegressorContextType

# This data-type describes all the information saved about a node.
# It's used in _tree.py where the Tree dataclass contains an attribute called nodes
# that uses this data type.
node_dtype = np.dtype(
    [
        # Index of the node in the nodes array
        ("node_id", np.uintp),
        #
        # Index of the parent node
        ("parent", np.uintp),
        #
        # Index of the left child node (it's TREE_LEAF is node is a leaf)
        ("left_child", np.intp),
        #
        # Index of the right child node (it's TREE_LEAF is node is a leaf)
        ("right_child", np.intp),
        #
        # Is the node a leaf ?
        ("is_leaf", np.bool),
        #
        # Is the node a left child ? (it's ? for the root node)
        ("is_left", np.bool),
        #
        # Depth of the node in the tree
        ("depth", np.uintp),
        #
        # Feature used for splitting the node (it's ??? if node is a leaf)
        ("feature", np.uintp),
        #
        # Continuous threshold used for splitting the node (not used for now)
        ("threshold", np.float32),
        #
        # Index of the bin threshold used for splitting the node
        ("bin_threshold", np.uint8),
        #
        # Impurity of the node. Used to avoid to split a "pure" node (with impurity=0).
        #   Note that impurity is computed using training samples (as all tree-growing
        #   related things)
        ("impurity", np.float32),
        #
        # Validation loss of the node, computed on validation (out-of-the-bag) samples
        ("loss_valid", np.float32),
        #
        # Logarithm of the subtree aggregation weight
        ("log_weight_tree", np.float32),
        #
        # Number of training samples in the node
        ("n_samples_train", np.uintp),
        #
        # Number of validation (out-of-the-bag) samples in the node
        ("n_samples_valid", np.uintp),
        #
        # Weighted number of training samples in the node
        ("w_samples_train", np.float32),
        #
        # Weighted number of validation (out-of-the-bag) samples in the node
        ("w_samples_valid", np.float32),
        #
        # Index of the first training sample in the node. We have that
        #   partition_train[start_train:end_train] contains the indexes of the node's
        #   training samples
        ("start_train", np.uintp),
        #
        # End-index of the slice containing the node's training samples indexes
        ("end_train", np.uintp),
        #
        # Index of the first validation (out-of-the-bag) sample in the node. We have
        #   that partition_valid[start_valid:end_valid] contains the indexes of the
        #   node's validation samples
        ("start_valid", np.uintp),
        #
        # End-index of the slice containing the node's validation samples indexes
        ("end_valid", np.uintp),
    ]
)

# Convert this dtype to a numba-usable data-type
node_type = from_dtype(node_dtype)

# TODO: many attributes are not used in the node_context ?

# A node_context contains all the data required to find the best split of the node
node_context_type = [
    # This array contains the index of all the features. This will be modified
    # inplace when sampling features to be considered for splits
    ("features_pool", uintp[::1]),
    #
    # This array will contain the features sampled uniformly at random (without
    # replacement) to be considered for splits
    ("features_sampled", uintp[::1]),
    #
    # Do we need to perform features sampling ?
    ("sample_features", boolean),
    #
    # Number of training samples in the node
    ("n_samples_train", intp),
    #
    # Number of validation (out-of-the-bag) samples in the node
    ("n_samples_valid", intp),
    #
    # Weighted number of training samples in the node
    ("w_samples_train", float32),
    #
    # Weighted number of validation (out-of-the-bag) samples in the node
    ("w_samples_valid", float32),
    #
    # Index of the first training sample in the node. We have that
    #   partition_train[start_train:end_train] contains the indexes of the node's
    #   training samples
    ("start_train", intp),
    #
    # End-index of the slice containing the node's training samples indexes
    ("end_train", intp),
    #
    # Index of the first validation (out-of-the-bag) sample in the node. We have
    #   that partition_valid[start_valid:end_valid] contains the indexes of the
    #   node's validation samples
    ("start_valid", intp),
    #
    # End-index of the slice containing the node's validation samples indexes
    ("end_valid", intp),
    #
    # Validation loss of the node, computed on validation (out-of-the-bag) samples
    ("loss_valid", float32),
    #
    # Weighted number of training samples for each (feature, bin) in the node
    ("w_samples_train_in_bins", float32[:, ::1]),
    #
    # Weighted number of validation samples for each (feature, bin) in the node
    ("w_samples_valid_in_bins", float32[:, ::1]),
]

node_classifier_context_type = [
    *node_context_type,
    #
    # Weighted number of training samples for each (feature, bin, label) in the node
    ("y_sum", float32[:, :, ::1]),
    #
    # Prediction produced by the node using the training data it contains
    ("y_pred", float32[::1]),
]

node_regressor_context_type = [
    *node_context_type,
    #
    # Weighted sum of the labels for each (feature, bin) in the node
    ("y_sum", float32[:, ::1]),
    #
    # Weighted sum of the squared labels for each (feature, bin) in the node
    ("y_sq_sum", float32[:, ::1]),
    #
    # Prediction produced by the node using the training data it contains
    ("y_pred", float32),
]

@jitclass(node_classifier_context_type)
class NodeClassifierContext:
    """A node_context contains all the data required to find the best split of the
    node for a classification tree.

    Parameters
    ----------
    tree_context : TreeClassifierContext
        An object describing the context of the tree

    Attributes
    ----------
    features_pool : ndarray
        Array of shape (n_features,) of intp dtype containing all the features
        indexes, namely [0, ..., n_features-1]

    features_sampled : ndarray
        Array of shape (max_features,) of intp dtype containing sampled features to
        be considered for possible splitting

    sample_features : bool
        Do we need to perform features sampling ? (This if True if n_features !=
        max_features and False otherwise)

    features : ndarray
        Array of shape (max_features,) of intp dtype containing the candidate features
        for splitting. These features are chosen uniformly at random each time a
        split is looked for

    n_samples_train : int
        Number of training samples in the node

    n_samples_valid : int
        Number of validation (out-of-the-bag) samples in the node

    w_samples_train : float
        Weighted number of training samples in the node

    w_samples_valid : float
        Weighted number of validation (out-of-the-bag) samples in the node

    start_train : int
        Index of the first training sample in the node. We have that
        partition_train[start_train:end_train] contains the indexes of the node's
        training samples

    end_train : int
        End-index of the slice containing the node's training samples indexes

    start_valid : int
        Index of the first validation (out-of-the-bag) sample in the node. We have
        that partition_valid[start_valid:end_valid] contains the indexes of the
        node's validation samples

    end_valid : int
        End-index of the slice containing the node's validation samples indexes

    loss_valid : float
        Validation loss of the node, computed on validation (out-of-the-bag) samples

    w_samples_train_in_bins : ndarray
        Weighted number of training samples for each (feature, bin) in the node

    w_samples_valid_in_bins : ndarray
        Weighted number of validation samples for each (feature, bin) in the node

    y_sum : ndarray
        Weighted number of training samples for each (feature, bin, label) in the node

    y_pred : ndarray
        Prediction produced by the node using the training data it contains
    """

    def __init__(self, tree_context):
        init_node_context(tree_context, self)
        max_features = tree_context.max_features
        max_bins = tree_context.max_bins
        n_classes = tree_context.n_classes
        self.y_sum = np.empty((max_features, max_bins, n_classes), dtype=np.float32)
        self.y_pred = np.empty(n_classes, dtype=np.float32)

@jitclass(node_regressor_context_type)
class NodeRegressorContext:
    """A node_context contains all the data required to find the best split of the
    node for a regression tree.

    Parameters
    ----------
    tree_context : TreeRegressorContext
        An object describing the context of the tree

    Attributes
    ----------
    features_pool : ndarray
        Array of shape (n_features,) of intp dtype containing all the features
        indexes, namely [0, ..., n_features-1]

    features_sampled : ndarray
        Array of shape (max_features,) of intp dtype containing sampled features to
        be considered for possible splitting

    sample_features : bool
        Do we need to perform features sampling ? (This if True if n_features !=
        max_features and False otherwise)

    features : ndarray
        Array of shape (max_features,) of intp dtype containing the candidate features
        for splitting. These features are chosen uniformly at random each time a
        split is looked for

    n_samples_train : int
        Number of training samples in the node

    n_samples_valid : int
        Number of validation (out-of-the-bag) samples in the node

    w_samples_train : float
        Weighted number of training samples in the node

    w_samples_valid : float
        Weighted number of validation (out-of-the-bag) samples in the node

    start_train : int
        Index of the first training sample in the node. We have that
        partition_train[start_train:end_train] contains the indexes of the node's
        training samples

    end_train : int
        End-index of the slice containing the node's training samples indexes

    start_valid : int
        Index of the first validation (out-of-the-bag) sample in the node. We have
        that partition_valid[start_valid:end_valid] contains the indexes of the
        node's validation samples

    end_valid : int
        End-index of the slice containing the node's validation samples indexes

    loss_valid : float
        Validation loss of the node, computed on validation (out-of-the-bag) samples

    w_samples_train_in_bins : ndarray
        Weighted number of training samples for each (feature, bin) in the node

    w_samples_valid_in_bins : ndarray
        Weighted number of validation samples for each (feature, bin) in the node

    y_sum : ndarray
        Weighted sum of the labels for each (feature, bin) in the node

    y_sq_sum : ndarray
        Weighted sum of the squared labels for each (feature, bin) in the node

    y_pred : float
        Prediction produced by the node using the training data it contains
    """

    def __init__(self, tree_context):
        init_node_context(tree_context, self)
        max_features = tree_context.max_features
        max_bins = tree_context.max_bins
        self.y_sum = np.empty((max_features, max_bins), dtype=np.float32)
        self.y_sq_sum = np.empty((max_features, max_bins), dtype=np.float32)
        self.y_pred = 0.0

NodeClassifierContextType = get_type(NodeClassifierContext)
NodeRegressorContextType = get_type(NodeRegressorContext)

@jit(
    [
        void(TreeClassifierContextType, NodeClassifierContextType),
        void(TreeRegressorContextType, NodeRegressorContextType),
    ],
    nopython=True,
    nogil=True,
)
def init_node_context(tree_context, node_context):
    """A common initializer for NodeContextClassifier and NodeContextRegressor.
    Since numba's jitclass support for inheritance is inexistant, we use this
    function both in the constructors of NodeContextClassifier and NodeContextRegressor.

    Parameters
    ----------
    tree_context : TreeContext
        An object describing the context of the tree

    node_context : NodeClassifierContext or NodeRegressorContext
        The node context that we be partly initialized by this function
    """
    max_features = tree_context.max_features
    max_bins = tree_context.max_bins
    n_features = tree_context.n_features
    # If max_features is the same as n_features, we don't need to sample features
    node_context.sample_features = max_features != n_features
    # This array contains the index of all the features. This will be modified
    # inplace when sampling features to be considered for splits
    node_context.features_pool = np.arange(0, n_features, dtype=np.uintp)
    # This array will contain the features sampled uniformly at random (without
    # replacement) to be considered for splits
    node_context.features_sampled = np.arange(0, max_features, dtype=np.uintp)
    node_context.w_samples_train_in_bins = np.empty(
        (max_features, max_bins), dtype=np.float32
    )
    node_context.w_samples_valid_in_bins = np.empty(
        (max_features, max_bins), dtype=np.float32
    )

@jit(
    void(
        TreeClassifierContextType, NodeClassifierContextType, uintp, uintp, uintp, uintp
    ),
    nopython=True,
    nogil=True,
    boundscheck=False,
    locals={
        "w_samples_train_in_bins": float32[:, ::1],
        "w_samples_valid_in_bins": float32[:, ::1],
        "y_sum": float32[:, :, ::1],
        "y_pred": float32[::1],
        "features": uintp[::1],
        "X": uint8[:, :],
        "y": float32[::1],
        "sample_weights": float32[::1],
        "partition_train": uintp[::1],
        "partition_valid": uintp[::1],
        "n_classes": uintp,
        "dirichlet": float32,
        "train_indices": uintp[::1],
        "valid_indices": uintp[::1],
        "w_samples_train": float32,
        "w_samples_valid": float32,
        "f": uintp,
        "loss_valid": float32,
        "feature": uintp,
        "sample": uintp,
        "bin": uint8,
        "label": uintp,
        "sample_weight": float32,
        "k": uintp,
    },
)
def compute_node_classifier_context(
    tree_context, node_context, start_train, end_train, start_valid, end_valid
):
    """Computes the node context from the data and from the tree context for
    classification. Computations are saved in the passed node_context.

    Parameters
    ----------
    tree_context : TreeContext
        The tree context

    node_context : NodeClassifierContext
        The node context that this function will compute

    start_train : int
        Index of the first training sample in the node. We have that
        partition_train[start_train:end_train] contains the indexes of the node's
        training samples

    end_train : int
        End-index of the slice containing the node's training samples indexes

    start_valid : int
        Index of the first validation (out-of-the-bag) sample in the node. We have
        that partition_valid[start_valid:end_valid] contains the indexes of the
        node's validation samples

    end_valid : int
        End-index of the slice containing the node's validation samples indexes
    """
    # Initialize the things from the node context
    w_samples_train_in_bins = node_context.w_samples_train_in_bins
    w_samples_valid_in_bins = node_context.w_samples_valid_in_bins
    y_sum = node_context.y_sum
    y_pred = node_context.y_pred

    # If necessary, sample the features
    if node_context.sample_features:
        sample_without_replacement(
            node_context.features_pool, node_context.features_sampled
        )

    features = node_context.features_sampled
    w_samples_train_in_bins.fill(0.0)
    w_samples_valid_in_bins.fill(0.0)
    y_sum.fill(0.0)
    y_pred.fill(0.0)

    # Get information from the tree context
    X = tree_context.X
    y = tree_context.y
    sample_weights = tree_context.sample_weights
    partition_train = tree_context.partition_train
    partition_valid = tree_context.partition_valid
    n_classes = tree_context.n_classes
    dirichlet = tree_context.dirichlet

    # The indices of the training samples contained in the node
    train_indices = partition_train[start_train:end_train]
    valid_indices = partition_valid[start_valid:end_valid]

    # Weighted number of training and validation samples
    w_samples_train = 0.0
    w_samples_valid = 0.0

    # A counter for the features
    f = 0
    # The validation loss
    loss_valid = 0.0

    # TODO: unrolling the for loop could be faster
    # For-loop on features first and then samples (X is F-major)

    for feature in features:
        # Compute statistics about training samples
        for sample in train_indices:
            bin = X[sample, feature]
            label = uintp(y[sample])
            sample_weight = sample_weights[sample]
            if f == 0:
                w_samples_train += sample_weight
                y_pred[label] += sample_weight
            # One more sample in this bin for the current feature
            w_samples_train_in_bins[f, bin] += sample_weight
            # One more sample in this bin for the current feature with this label
            y_sum[f, bin, label] += sample_weight

        # TODO: we should put this outside so that we can change the dirichlet parameter
        # without rebuilding the tree
        # The prediction is given by the formula
        #   y_k = (n_k + dirichlet) / (n_samples + dirichlet * n_classes)
        # where n_k is the number of samples with label class k
        if f == 0:
            for k in range(n_classes):
                y_pred[k] = (y_pred[k] + dirichlet) / (
                    w_samples_train + n_classes * dirichlet
                )

        # Compute sample counts about validation samples
        for sample in valid_indices:
            bin = X[sample, feature]
            sample_weight = sample_weights[sample]
            if f == 0:
                w_samples_valid += sample_weight
                label = uintp(y[sample])
                # TODO: aggregation loss is hard-coded here. Call a function instead
                #  when implementing other losses
                loss_valid += - w_samples_valid * log(y_pred[label])

            w_samples_valid_in_bins[f, bin] += sample_weight

        f += 1

    # Save remaining things in the node context
    node_context.n_samples_train = end_train - start_train
    node_context.n_samples_valid = end_valid - start_valid
    node_context.loss_valid = loss_valid
    node_context.w_samples_train = w_samples_train
    node_context.w_samples_valid = w_samples_valid

@jit(
    void(
        TreeRegressorContextType, NodeRegressorContextType, uintp, uintp, uintp, uintp
    ),
    nopython=True,
    nogil=True,
    boundscheck=False,
    locals={
        "w_samples_train_in_bins": float32[:, ::1],
        "w_samples_valid_in_bins": float32[:, ::1],
        "y_sum": float32[:, ::1],
        "y_sq_sum": float32[:, ::1],
        "y_pred": float32,
        "features": uintp[::1],
        "X": uint8[:, :],
        "y": float32[::1],
        "sample_weights": float32[::1],
        "partition_train": uintp[::1],
        "partition_valid": uintp[::1],
        "train_indices": uintp[::1],
        "valid_indices": uintp[::1],
        "w_samples_train": float32,
        "w_samples_valid": float32,
        "f": uintp,
        "loss_valid": float32,
        "feature": uintp,
        "sample": uintp,
        "bin": uint8,
        "label": float32,
        "sample_weight": float32,
        "k": uintp,
        "w_y": float32,
    },
)
def compute_node_regressor_context(
    tree_context, node_context, start_train, end_train, start_valid, end_valid
):
    """Computes the node context from the data and from the tree context for regression.
    Computations are saved in the passed node_context.

    Parameters
    ----------
    tree_context : TreeContext
        The tree context

    node_context : NodeRegressorContext
        The node context that this function will compute

    start_train : int
        Index of the first training sample in the node. We have that
        partition_train[start_train:end_train] contains the indexes of the node's
        training samples

    end_train : int
        End-index of the slice containing the node's training samples indexes

    start_valid : int
        Index of the first validation (out-of-the-bag) sample in the node. We have
        that partition_valid[start_valid:end_valid] contains the indexes of the
        node's validation samples

    end_valid : int
        End-index of the slice containing the node's validation samples indexes
    """
    # Initialize the things from the node context
    w_samples_train_in_bins = node_context.w_samples_train_in_bins
    w_samples_valid_in_bins = node_context.w_samples_valid_in_bins
    y_sum = node_context.y_sum
    y_sq_sum = node_context.y_sq_sum

    # If necessary, sample the features
    if node_context.sample_features:
        sample_without_replacement(
            node_context.features_pool, node_context.features_sampled
        )

    features = node_context.features_sampled
    w_samples_train_in_bins.fill(0.0)
    w_samples_valid_in_bins.fill(0.0)
    y_sum.fill(0.0)
    y_sq_sum.fill(0.0)
    y_pred = 0.0

    # Get information from the tree context
    X = tree_context.X
    y = tree_context.y
    sample_weights = tree_context.sample_weights
    partition_train = tree_context.partition_train
    partition_valid = tree_context.partition_valid

    # The indices of the training samples contained in the node
    train_indices = partition_train[start_train:end_train]
    valid_indices = partition_valid[start_valid:end_valid]

    # Weighted number of training and validation samples
    w_samples_train = 0.0
    w_samples_valid = 0.0

    # A counter for the features
    f = 0
    # The validation loss
    loss_valid = 0.0

    # TODO: unrolling the for loop could be faster
    # For-loop on features first and then samples (X is F-major)

    for feature in features:
        # Compute statistics about training samples
        for sample in train_indices:
            bin = X[sample, feature]
            label = y[sample]
            sample_weight = sample_weights[sample]
            w_y = sample_weight * label
            if f == 0:
                w_samples_train += sample_weight
                y_pred += w_y
            # One more sample in this bin for the current feature
            w_samples_train_in_bins[f, bin] += sample_weight
            # One more sample in this bin for the current feature with this label

            y_sum[f, bin] += w_y
            y_sq_sum[f, bin] += w_y * label

        # The prediction is simply the weighted average of labels
        if f == 0:
            y_pred /= w_samples_train

        # Compute sample counts about validation samples
        for sample in valid_indices:
            bin = X[sample, feature]
            sample_weight = sample_weights[sample]
            if f == 0:
                w_samples_valid += sample_weight
                label = y[sample]
                # TODO: aggregation loss is hard-coded here. Call a function instead
                #  when implementing other losses
                loss_valid += sample_weight * (label - y_pred) * (label - y_pred)

            w_samples_valid_in_bins[f, bin] += sample_weight

        f += 1

    # Save remaining things in the node context
    node_context.y_pred = y_pred
    node_context.n_samples_train = end_train - start_train
    node_context.n_samples_valid = end_valid - start_valid
    node_context.loss_valid = loss_valid
    node_context.w_samples_train = w_samples_train
    node_context.w_samples_valid = w_samples_valid


import torch
import torch.nn as nn
from torch.utils.data import Dataset
import numpy as np
import os
import cv2
import random

from torchOnVideo.denoising.vnlnet.video_patch_search import VideoPatchSearch

# download the dataset from f16-video-set

# Using the same weights for each channel
# gives optimal contrast over noise.
def rgb_to_gray(img):
    img = np.asarray(img, dtype=np.float32)
    if (len(img.shape) == 3 or img.shape[3] == 1):
        return img
    res = np.dot(img[...,:3], [0.57735, 0.57735, 0.57735])
    return np.asarray(res, dtype=np.float32)

class F16Denoising(Dataset):
    def __init__(self, data_path, color_mode=False, sigma=25, oracle_mode=0, past_frames=7, future_frames=7, search_window_width=41, nn_patch_width=41, pass_nn_value=False, patch_width=44, patch_stride=5):
        super(Dataset, self).__init__()
        self.color = color_mode
        self.sigma = sigma

        categories = os.scandir(data_path)
        categories = [c for c in categories if not c.name.startswith('.') and c.is_dir()]

        video_paths_dict = {}
        video_paths_list = []
        for c in categories:
            list_for_category = []
            paths = os.scandir(c.path)
            paths = [p for p in paths if not p.name.startswith('.') and p.is_dir()]
            video_paths_list.extend(paths)
            video_paths_dict[c.name] = [p.path for p in paths]
        categories = [c.name for c in categories]

        print ('%d categories' % len(categories))
        print ('%d videos' % len(video_paths_list))

        self.categories = categories
        self.video_paths_dict = video_paths_dict

        self.patch_width_nn = nn_patch_width
        self.patch_data_width_nn = 1
        self.past_frames = past_frames
        self.future_frames = future_frames
        self.num_neighbors = 1 + past_frames + future_frames
        self.patch_width = patch_width
        self.patch_stride = patch_stride
        self.oracle_mode = oracle_mode
        self.pass_nn_value = pass_nn_value

        self.ps = VideoPatchSearch(patch_search_width=self.patch_width_nn, patch_data_width=self.patch_data_width_nn,
                                   past_frames=self.past_frames, future_frames=self.future_frames,
                                   input_dtype=np.float32,
                                   search_width=search_window_width)

        np.random.seed(2018)
        self.prepare_epoch()

    def prepare_epoch(self):
        """
        Read random videos of the database
        """
        subdirnames = ['/01/', '/02/', '/03/', '/04/', '/05/', '/06/', '/07/', '/08/', '/09/']
        filenames= ['001.png', '002.png', '003.png', '004.png', '005.png', '006.png', '007.png', '008.png', '009.png', '010.png', '011.png', '012.png', '013.png', '014.png', '015.png']
        burst_nums = np.random.randint(len(subdirnames), size=len(self.categories*10))
        frame_nums = np.random.randint(self.past_frames, high=len(filenames)-self.future_frames, size=len(self.categories*10))
        i = 0
        self.videos = []
        self.keys = []
        for c in self.categories:
            paths = self.video_paths_dict[c]
            paths = np.random.permutation(paths)
            for p in paths[:3]:
                dir_path = p + subdirnames[burst_nums[i]]
                if not (os.path.exists(dir_path)):
                    continue
                video = []
                for f in range(frame_nums[i]-self.past_frames, frame_nums[i]+self.future_frames+1):
                    img = cv2.imread(dir_path + filenames[f])
                    img = np.asarray(img, dtype=np.float32)
                    if self.color:
                        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                    else:
                        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                        img = np.expand_dims(img, 2)
                    img = np.asarray(img, dtype=np.float32)
                    img = img/255.
                    video.append(img)
                ref_image = video[self.past_frames]
                video = np.asarray(video, dtype=np.float32)

                video_noised = video + self.sigma * np.random.randn(video.shape[0], video.shape[1], video.shape[2], video.shape[3])
                video_noised = np.asarray(video_noised, dtype=np.float32)

                video_search_gray = rgb_to_gray(video if self.oracle_mode == 1 else video_noised)

                nn = self.ps.compute(video_search_gray, self.past_frames)

                self.videos.append((video_noised[self.past_frames, ...] - ref_image, video_noised, nn))
                ys = range(2*self.patch_width, ref_image.shape[0]-2*self.patch_width, self.patch_stride)
                xs = range(2*self.patch_width, ref_image.shape[1]-2*self.patch_width, self.patch_stride)
                xx, yy = np.meshgrid(xs, ys)
                xx = np.asarray(xx.flatten(), dtype=np.uint32)
                yy = np.asarray(yy.flatten(), dtype=np.uint32)
                self.keys.append(np.stack([i*np.ones([len(xx)], dtype=np.uint32), xx, yy]).T)

                i = i + 1
        self.keys = np.concatenate(self.keys, axis=0)
        self.num_keys = self.keys.shape[0]
        self.indices = [i for i in range(self.num_keys)]

        random.shuffle(self.indices)

    def __len__(self):
        return self.num_keys

    def data_num_channels(self):
        return (3 if self.color else 1) * (self.num_neighbors if self.pass_nn_value else 1)

    def __getitem__(self, index):
        key = self.keys[self.indices[index],:]
        patch_width = self.patch_width
        i = key[0]
        x = key[1]
        y = key[2]
        anchor = self.patch_width_nn//2
        noise = self.videos[i][0][(y-patch_width):y, (x-patch_width):x,:]

        if self.pass_nn_value:
            nn_patch = self.videos[i][2][(y-anchor-patch_width):(y-anchor), (x-anchor-patch_width):(x-anchor),:]
            patch_stack = self.ps.build_neighbors_array(self.videos[i][1], nn_patch)
        else:
            patch_stack = self.videos[i][1][self.past_frames, (y-patch_width):y, (x-patch_width):x,:]

        patch_stack = patch_stack.transpose(2, 0, 1)
        noise = noise.transpose(2, 0, 1)

        return (torch.Tensor(patch_stack), torch.Tensor(noise))


import cv2
import numpy as np
import imgreco.resources
from util import cvimage as Image
inventory_mask = imgreco.resources.load_image('inventory/inventory_mask.png', imread_flags=cv2.IMREAD_GRAYSCALE)
inventory_mask16 = inventory_mask.resize((16, 16), resample=Image.NEAREST).array
inventory_mask = inventory_mask.array
filenames = imgreco.resources.get_entries('inventory/icons')[1]
templates = {filename.removesuffix('.webp').removesuffix('.png'): imgreco.resources.load_image('inventory/icons/' + filename, imread_flags=cv2.IMREAD_COLOR).resize((48, 48)) for filename in filenames if filename.endswith('.webp') or filename.endswith('.png')}

all_itemid = np.array(list(templates.keys()), dtype=object)
all_icons = np.concatenate([templates[x].array[None, ...] for x in all_itemid]).astype(np.float32)
all_icons_16 = np.concatenate([templates[x].resize((16, 16)).array[None, ...] for x in all_itemid]).astype(np.float32)
all_icons[:, inventory_mask == 0] = 0
all_icons_16[:, inventory_mask16 == 0] = 0