import cv2
import numpy as np
from collections import defaultdict

'''
Variables that are important for
this program.
'''

times=0
numfan=0
#Video stream used for processing
vs=cv2.VideoCapture("/Users/aidanobrien/Documents/GitHub.nosync/BeeFanningDetector/Assets/test_img&videos/test_vid4.mp4")
#Background image used for initial background subtraction and binary and operations.
bk=cv2.imread('/Users/aidanobrien/Documents/GitHub.nosync/BeeFanningDetector/Assets/test_img&videos/testbkgrd1.jpg')
#Background image used for secont background subtraction and binary and operations. This is used to detect the wings.
bk2=cv2.imread('/Users/aidanobrien/Documents/GitHub.nosync/BeeFanningDetector/Assets/test_img&videos/black.png')
frames=defaultdict(dict) #Dict for holding the video frames of potentially fanning bees
foundbee=defaultdict(dict) #Dict that holds flags cooresponding to wether or not a fanning bee was found at a particular spot
fanframe=defaultdict(dict) #Dict that holds the most recent frame number when a particular bee was detected.
found=False
sframes=0
#Only use these cropping bounds in 30 fps vids. This eliminates
#areas at top/bottom of frame that never contain fanning.
bk=bk[100:100+240,0:0+640]
bk2=bk2[100:100+240,0:0+640]

'''
This function is given a wing contour and
frame image, and determines wether or not
the cooresponding bee is fanning.
'''
def checkWings(c,img):
    global numfan
    global frames
    global foundbee
    global sframes
    global fanframe
    i=0
    (x,y),(ma,Ma),angle=cv2.fitEllipse(c)
    x2,y2,w,h=cv2.boundingRect(c)
    ell=cv2.fitEllipse(c)
    found=False
    mom=cv2.moments(c)
    hy=0
    xw=0
    cx = int(mom["m10"] / mom["m00"])
    cy = int(mom["m01"] / mom["m00"])
    if(cy>100):
        hy=100
    else:
        hy=50
    if(cx>100):
        xw=100
    else:
        xw=50
    #height, width, layers = img.shape
    if(frames.get(tuple([cx,cy])) is not None and i in frames.get(tuple([cx,cy]))):
        framediff=0
        if(fanframe.get(tuple([cx,cy])) is not None):
            while framediff<100:
                if(i in fanframe.get(tuple([cx,cy]))):
                    framediff=sframes-fanframe.get(tuple([cx,cy]))[i]
                    if(cx == 428 and cy == 140):
                        print("Diff: {}".format(framediff))
                else:
                    break
                if framediff>=100:

                    i+=1
                else:
                    break
        if(i in fanframe.get(tuple([cx,cy]))):
            #hy, xw, layers = frames[cx,cy][i][0].shape
            fanframe[cx,cy][i]=sframes
            frames[cx,cy][i].append(img[cy-hy:cy+hy,cx-xw:cx+xw])
    else:
        for cX in range (cx-55,cx+55):
            for cY in range (cy-10,cy+10):

                framediff=0
                if(fanframe.get(tuple([cX,cY])) is not None):
                    while framediff<100:
                        if(i in fanframe.get(tuple([cX,cY]))):
                            if(cX == 428 and cY == 140):
                                print("Diff: {}".format(framediff))
                            framediff=sframes-fanframe.get(tuple([cX,cY]))[i]
                        else:
                            break
                        if framediff>=100:

                            i+=1
                        else:
                            break

                if(frames.get(tuple([cX,cY])) is not None and i in frames.get(tuple([cX,cY]))):
                    #print("{},{}".format(cx,cy))
                    #hy, xw, layers = frames[cX,cY][i][0].shape
                    if(cY>100):
                        hy=100
                    else:
                        hy=50
                    if(cX>100):
                        xw=100
                    else:
                        xw=50
                    frames[cX,cY][i].append(img[cY-hy:cY+hy,cX-xw:cX+xw])
                    fanframe[cX,cY][i]=sframes
                    found=True
                    if(len(frames.get(tuple([cX,cY]))[i])>=20 and foundbee.get(tuple([cX,cY]))[i]==False):
                        print("Fanning Detected")
                        #print("{}, {}".format(cX,cY))
                        cv2.ellipse(img,ell,(0,255,0),2)
                        foundbee[cX,cY][i]=True
                        numfan+=1
                    return

        if(found==False and cy < 189):
            fanframe[cx,cy][i]=sframes
            frames[cx,cy][i]=[img[cy-hy:cy+hy,cx-xw:cx+xw]]
            foundbee[cx,cy][i]=False

'''
This function iterates through every entry
in the frames dictionary and exports
videos frames for entries with atleast
20 frames. Videos can be found in the fanning_exports
directory.
'''
def make_vids():
    i = 0
    global frames
    for key in frames:
        for key2 in frames[key]:
            f=frames[key][key2]
            height, width, layers = f[0].shape
            print(len(f))
            if(len(f)>=20 and (width is not 0 and height is not 0)):

                size = (width,height)
                out = cv2.VideoWriter()
                out.open('/Users/aidanobrien/Documents/GitHub.nosync/BeeFanningDetector/Assets/fanning_exports/wings_'+str(key)+", "+str(fanframe[key][key2])+'.mov',cv2.VideoWriter_fourcc(*'mp4v'), 10, (size),True)
                for fr in f:
                    out.write(fr)
                out.release()
                i=i+1

'''
This is the main driver loop
that iterates through the provided
video and calls the appropriate functions
to detect fanning bees.
'''
while True:
    sframes+=1
    hasframes,img=vs.read()
    if(hasframes == False):
        break
    #Again, in 30 fps top and bottom areas  of the frame are removed.
    img=img[100:100+240,0:0+640]
    subImage=(bk.astype('int32')-img.astype('int32')).clip(0).astype('uint8')
    grey=cv2.cvtColor(subImage,cv2.COLOR_BGR2GRAY)
    retval,thresh=cv2.threshold(grey,35,255,cv2.THRESH_BINARY)
    kernel=np.ones((5,5),np.uint8)
    thresh=cv2.morphologyEx(thresh,cv2.MORPH_OPEN,kernel)
    noback = cv2.bitwise_and(img, img, mask= thresh)
    #Color Bounds for 60 fps vids
    #upper = np.array([220,220,220])
    #lower = np.array([160,160,160])
    cv2.imshow('noback',noback)
    #Color Bounds for 30 fps vids
    upper = np.array([255,255,255])
    lower = np.array([128,128,128])
    mask = cv2.inRange(noback, lower, upper)

    wings = cv2.bitwise_and(noback, noback, mask=mask)

    cv2.imshow('Just_Wings/Shadows',wings)

    subImage2=(wings.astype('int32')-bk2.astype('int32')).clip(0).astype('uint8')
    grey2=cv2.cvtColor(subImage2,cv2.COLOR_BGR2GRAY)
    retval2,thresh2=cv2.threshold(grey2,35,255,cv2.THRESH_BINARY)
    kernel2=np.ones((5,5),np.uint8)
    thresh2=cv2.morphologyEx(thresh2,cv2.MORPH_OPEN,kernel)
    im2, contours1, hierarchy1 = cv2.findContours(thresh2, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    cnt2=[]
    for c in contours1:
        x,y,w,h=cv2.boundingRect(c)
        r=w/h
        #wing ellipse bounds for 30 fps video
        if(w*h>150 and w*h < 200 and w > h):
        #wing ellipse bounds for 60 fps video (still refining these)
        #if(w*h>300 and w > h and w > 25 and w < 53 and h > 10 and h < 30 and r > 1.44 and r < 3.9):
            ell=cv2.fitEllipse(c)
            checkWings(c,img)
            cv2.ellipse(img,ell,(0,255,0),2)

            #print(w*h,w,h)

        else:
            cnt2.append(c)

    cv2.putText(img, "Fanning Bees: {}".format(numfan), (0, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2, cv2.LINE_4)

    cv2.imshow('Result',img)

    cv2.imshow('Thresh',thresh2)
    if times > 0:
        key=cv2.waitKey(1) & 0xFF
        #if q is pressed, stop loop
        if key == ord("c"):
            continue
        if key == ord("q"):
            break
    times = times + 1
vs.release()
cv2.destroyAllWindows()
make_vids()


import cv2 as cv
import os
import tensorflow as tf
import numpy as np

class TensoflowMobilNetSSDFaceDector():
    def __init__(self,
                 det_threshold=0.3,
                 model_path=os.path.join(os.path.dirname(__file__), 'models/ssd/frozen_inference_graph_face.pb')):
        self.det_threshold = det_threshold
        self.detection_graph = tf.Graph()
        with self.detection_graph.as_default():
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(model_path, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')

        with self.detection_graph.as_default():
            config = tf.ConfigProto()
            config.gpu_options.allow_growth = True
            self.sess = tf.Session(graph=self.detection_graph, config=config)

    def detect_face(self, image_path):
        # Load image
        image = cv.imread(image_path)

        h, w, c = image.shape

        image_np = cv.cvtColor(image, cv.COLOR_BGR2RGB)

        image_np_expanded = np.expand_dims(image_np, axis=0)
        image_tensor = self.detection_graph.get_tensor_by_name(
            'image_tensor:0')

        boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')

        scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
        classes = self.detection_graph.get_tensor_by_name(
            'detection_classes:0')
        num_detections = self.detection_graph.get_tensor_by_name(
            'num_detections:0')

        (boxes, scores, classes, num_detections) = self.sess.run(
            [boxes, scores, classes, num_detections],
            feed_dict={image_tensor: image_np_expanded})

        boxes = np.squeeze(boxes)
        scores = np.squeeze(scores)

        filtered_score_index = np.argwhere(
            scores >= self.det_threshold).flatten()
        selected_boxes = boxes[filtered_score_index]

        faces = np.array([[
            int(x1 * w),
            int(y1 * h),
            int(x2 * w),
            int(y2 * h),
        ] for y1, x1, y2, x2 in selected_boxes])

        return faces


import tempfile
import os
import cv2
import argparse
import numpy as np
from zipfile import ZipFile
from shutil import copyfile, rmtree
from scipy.spatial import distance

def parse_arguments():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '-s', '--source', type=str,
        default='/home/markus/KITTI/semantic/training/image_2',
        help='Source directory for images you are looking matches for'
    )

    parser.add_argument(
        '-i', '--input', type=str,
        default='/home/markus/KITTI/raw_data_synced',
        help='Input directory containing KITTI raw data zips'
    )

    parser.add_argument(
        '-b', '--batch', type=int,
        default=500,
        help='Batch size of images to read at once for distance calculations'
    )

    parser.add_argument(
        '-o', '--output', type=str,
        default='/home/markus/KITTI/semantic/training/',
        help='Output directory that should contain the lidar and calibration data directories after running this'
    )

    return parser.parse_args()

def read_images(images_path, nr_images=None, offset=0):
    image_filenames = sorted(os.listdir(images_path))
    images = []

    if nr_images:
        start = offset
        end = offset + nr_images
        image_filenames = image_filenames[start:end]

    for image_filename in image_filenames:
        image_path = os.path.join(images_path, image_filename)
        image = cv2.imread(image_path).flatten()[:1358640]
        images.append(image)

    return image_filenames, np.array(images)

def get_target_zips(input_path):
    zip_files_path = os.path.join(input_path, '*.zip')
    return sorted(os.listdir(zip_files_path))

def unpack_zip(zip_path, target_path):
    with ZipFile(zip_path, 'r') as zip_object:
        zip_object.extractall(target_path)

def copy_velodyne(source_image_name, zip_image_name, zip_velodyne_path, output_dir):
    zip_velodyne_filename = zip_image_name[:-4] + '.bin'
    velodyne_path = os.path.join(
        zip_velodyne_path, zip_velodyne_filename)

    target_velodyne_filename = source_image_name[:-4] + '.bin'
    target_velodyne_path = os.path.join(
        output_dir, 'velodyne', target_velodyne_filename)

    copyfile(velodyne_path, target_velodyne_path)

def calculate_distances_and_find_matches(source_image_names, source_images, zip_image_names, zip_images, zip_velodyne_path, extraction_name, output_dir):

    distances = distance.cdist(
        source_images, zip_images, 'sqeuclidean')
    current_matches = np.where(distances == 0)

    matches = []
    for source_frame_id, zip_frame_id in zip(current_matches[0], current_matches[1]):
        source_image_name = source_image_names[source_frame_id]
        zip_image_name = zip_image_names[zip_frame_id]

        copy_velodyne(source_image_name, zip_image_name,
                      zip_velodyne_path, output_dir)

        match = (source_image_name, extraction_name, zip_image_name)
        matches.append(match)

    return matches

# TODO: Add copying for calibration matrices as well
def match_images(source_images_dir, target_images_dir, batch_size, output_dir):

    source_image_names, source_images = read_images(source_images_dir)
    target_zips = get_target_zips(target_images_dir)

    total_nr_source_images = len(source_image_names)

    matches_file_path = os.path.join(output_dir, 'matches.txt')
    if os.path.isfile(matches_file_path):
        all_matches = []
        with open(matches_file_path, 'r') as f:
            for line in f:
                line = line.rstrip().split(' ')
                all_matches.append((line[0], line[1], line[2]))

    else:
        all_matches = []

    if len(all_matches) > 0:
        processed_zips = ['{name}.zip'.format(name=x[1]) for x in all_matches]
        target_zips = [
            name for name in target_zips if name not in processed_zips]

        processed_source_image_names = [match[0] for match in all_matches]

        # Filter source
        def filter_sources(x): return x[0] not in processed_source_image_names
        filtered_sources = list(
            filter(filter_sources, zip(source_image_names, source_images)))

        source_image_names = [x[0] for x in filtered_sources]
        source_images = np.array([x[1] for x in filtered_sources])

    with tempfile.TemporaryDirectory() as tmp_dir:

        for i, target_zip in enumerate(target_zips):
            if len(all_matches) == total_nr_source_images:
                break

            print('Working on archive {index:d}/{number_archives:d}: {name:s}'.format(
                index=(i+1), number_archives=len(target_zips), name=target_zip))

            target_zip_path = os.path.join(target_images_dir, target_zip)
            unpack_zip(target_zip_path, tmp_dir)

            zip_date = target_zip[:10]
            extraction_name = target_zip[:-4]
            zip_images_path = os.path.join(
                tmp_dir, zip_date, extraction_name, 'image_02', 'data')

            zip_velodyne_path = os.path.join(
                tmp_dir, zip_date, extraction_name, 'velodyne_points', 'data')

            if batch_size > 0:

                batch_offset = 0
                while 1:
                    zip_image_names, zip_images = read_images(
                        zip_images_path, batch_size, batch_offset)

                    batch_matches = calculate_distances_and_find_matches(
                        source_image_names, source_images, zip_image_names, zip_images, zip_velodyne_path, extraction_name, output_dir)
                    all_matches.extend(batch_matches)

                    if zip_images.shape[0] < batch_size:
                        break

                    batch_offset += batch_size
            else:
                zip_image_names, zip_images = read_images(zip_images_path)
                matches = calculate_distances_and_find_matches(
                    source_image_names, source_images, zip_image_names, zip_images, zip_velodyne_path, extraction_name, output_dir)

                all_matches.extend(matches)

            # Remove unpacked folder
            unzipped_folder_path = os.path.join(
                tmp_dir, zip_date, extraction_name)
            rmtree(unzipped_folder_path)

    print('Matches contained {:d} elements'.format(len(all_matches)))

    all_matches = sorted(all_matches, key=lambda x: x[0])
    with open(matches_file_path, 'w') as f:
        for match in all_matches:
            f.write('{semantic_image} {extract} {extract_image}\n'.format(
                semantic_image=match[0], extract=match[1], extract_image=match[2]))

    matched_image_names = [x[0] for x in all_matches]
    sources_not_matched = [
        name for name in source_image_names if name not in matched_image_names]
    print('Source images not matched: ', sources_not_matched)

if __name__ == '__main__':
    args = parse_arguments()
    SOURCE_DIRECTORY = args.source
    INPUT_DIRECTORY = args.input
    BATCH_SIZE = args.batch
    OUTPUT_DIRECTORY = args.output

    if not os.path.isdir(OUTPUT_DIRECTORY):
        os.mkdir(OUTPUT_DIRECTORY)

    velodyne_path = os.path.join(OUTPUT_DIRECTORY, 'velodyne')
    if not os.path.isdir(velodyne_path):
        os.mkdir(velodyne_path)

    match_images(SOURCE_DIRECTORY, INPUT_DIRECTORY,
                 BATCH_SIZE, OUTPUT_DIRECTORY)


import collections
import string

import numpy
import six

import cupy
from cupy import carray
from cupy import elementwise
from cupy import util

six_range = six.moves.range
six_zip = six.moves.zip

_broadcast = elementwise._broadcast
_check_args = elementwise._check_args
_decide_params_type = elementwise._decide_params_type
_get_kernel_params = elementwise._get_kernel_params
_get_args_info = elementwise._get_args_info
_get_out_args = elementwise._get_out_args
_get_out_args_with_params = elementwise._get_out_args_with_params
_get_param_info = elementwise._get_param_info
_get_typename = elementwise._get_typename
_guess_routine = elementwise._guess_routine
_reduce_dims = elementwise._reduce_dims

def _get_simple_reduction_kernel(
        name, block_size, reduce_type, params, identity,
        pre_map_expr, reduce_expr, post_map_expr,
        type_preamble, input_expr, output_expr, preamble, options):
    if identity is None:
        identity = ''
    module_code = string.Template('''
    ${type_preamble}
    ${preamble}
    #define REDUCE(a, b) (${reduce_expr})
    #define POST_MAP(a) (${post_map_expr})

    typedef ${reduce_type} _type_reduce;
    extern "C" __global__ void ${name}(${params}) {
      if (_out_clp2_size > 256) {
        CUPY_FOR(_i, _out_ind.size()) {
          _type_reduce _s = _type_reduce(${identity});
          for (int _j = _i, _J = 0;
               _j < _in_ind.size();
               _j += _out_ind.size(), _J++) {
            _in_ind.set(_j);
            ${input_expr}
            _type_reduce _a = ${pre_map_expr};
            _s = REDUCE(_s, _a);
          }
          _out_ind.set(_i);
          ${output_expr}
          POST_MAP(_s);
        }
      } else {
        extern __shared__ _type_reduce _sdata_raw[];
        _type_reduce *_sdata = _sdata_raw;
        int _tid = threadIdx.x;
        _sdata[_tid] = _type_reduce(${identity});
        unsigned int _i = _tid % _out_clp2_size;
        if (_i >= _out_ind.size()) return;
        _type_reduce _s = _type_reduce(${identity});
        int _J_offset = _tid / _out_clp2_size;
        int _j_offset = _J_offset * _out_ind.size();
        int _J_stride = ${block_size} / _out_clp2_size;
        int _j_stride = _J_stride * _out_ind.size();
        for (int _j = _i + _j_offset, _J = _J_offset;
             _j < _in_ind.size();
             _j += _j_stride, _J += _J_stride) {
          _in_ind.set(_j);
          ${input_expr}
          _type_reduce _a = ${pre_map_expr};
          _s = REDUCE(_s, _a);
        }
        _sdata[_tid] = _s;
        __syncthreads();
        if (_tid >= 256) return;
        _sdata[_tid] = REDUCE(_sdata[_tid], _sdata[_tid + 256]);
        __syncthreads();
        if (_out_clp2_size <= 128) {
          _sdata[_tid] = REDUCE(_sdata[_tid], _sdata[_tid + 128]);
          __syncthreads();
          if (_out_clp2_size <= 64) {
            _sdata[_tid] = REDUCE(_sdata[_tid], _sdata[_tid + 64]);
            __syncthreads();
            if (_out_clp2_size <= 32) {
              _sdata[_tid] = REDUCE(_sdata[_tid], _sdata[_tid + 32]);
              if (_out_clp2_size <= 16) {
                _sdata[_tid] = REDUCE(_sdata[_tid], _sdata[_tid + 16]);
                if (_out_clp2_size <= 8) {
                  _sdata[_tid] = REDUCE(_sdata[_tid], _sdata[_tid + 8]);
                  if (_out_clp2_size <= 4) {
                    _sdata[_tid] = REDUCE(_sdata[_tid], _sdata[_tid + 4]);
                    if (_out_clp2_size <= 2) {
                      _sdata[_tid] = REDUCE(_sdata[_tid], _sdata[_tid + 2]);
                      if (_out_clp2_size <= 1) {
                        _sdata[_tid] = REDUCE(_sdata[_tid], _sdata[_tid + 1]);
                      }
                    }
                  }
                }
              }
            }
          }
        }
        _s = _sdata[_tid];
        if (_tid >= _out_ind.size()) return;
        _out_ind.set(_i);
         ${output_expr}
        POST_MAP(_s);
      }
    }''').substitute(
        name=name,
        block_size=block_size,
        reduce_type=reduce_type,
        params=params,
        identity=identity,
        reduce_expr=reduce_expr,
        pre_map_expr=pre_map_expr,
        post_map_expr=post_map_expr,
        type_preamble=type_preamble,
        input_expr=input_expr,
        output_expr=output_expr,
        preamble=preamble)
    module = carray.compile_with_cache(module_code, options)
    return module.get_function(name)

def _get_axis(axis, ndim):
    if axis is None:
        axis = tuple(six_range(ndim))
    elif isinstance(axis, collections.Sequence):
        axis = tuple(axis)
    else:
        axis = axis,

    for dim in axis:
        if dim < -ndim or dim >= ndim:
            raise ValueError('Axis overrun')
    axis = tuple(sorted([dim % ndim for dim in axis]))
    raxis = tuple([dim for dim in six_range(ndim) if dim not in axis])
    return axis, raxis

def _get_out_shape(shape, axis, raxis, keepdims):
    if keepdims:
        out_shape = list(shape)
        for i in axis:
            out_shape[i] = 1
        return tuple(out_shape)
    return tuple([shape[i] for i in raxis])

def _get_trans_args(args, trans, shape, params=None):
    if trans == tuple(six_range(len(shape))):
        return args, shape
    if params is not None and any(p.raw for p in params):
        raise NotImplementedError('Illegal conditions')
    args = [cupy.transpose(a, trans) if isinstance(a, cupy.ndarray) else a
            for a in args]
    shape = tuple([shape[i] for i in trans])
    return args, shape

def _get_inout_args(in_args, out_args, in_indexer, out_indexer, out_clp2_size,
                    params, reduce_dims):
    if reduce_dims:
        in_args, in_shape = _reduce_dims(
            in_args, params, in_indexer.shape)
        out_args, out_shape = _reduce_dims(
            out_args, params[len(in_args):], out_indexer.shape)
        in_indexer.shape = in_shape
        out_indexer.shape = out_shape
    args = in_args + out_args + [in_indexer, out_indexer,
                                 numpy.int32(out_clp2_size)]
    return args

@util.memoize(for_each_device=True)
def _get_simple_reduction_function(
        routine, params, args_info, in_arg_dtype, out_arg_dtype, out_types,
        name, block_size, identity, input_expr, output_expr, _preamble,
        options):
    reduce_type = routine[3]
    if reduce_type is None:
        reduce_type = _get_typename(out_types[0])

    t = (_get_typename(in_arg_dtype), _get_typename(out_arg_dtype))
    type_preamble = 'typedef %s type_in0_raw; typedef %s type_out0_raw;' % t

    params = _get_kernel_params(params, args_info)
    return _get_simple_reduction_kernel(
        name, block_size, reduce_type, params, identity,
        routine[0], routine[1], routine[2],
        type_preamble, input_expr, output_expr, _preamble, options)

class simple_reduction_function(object):

    def __init__(self, name, ops, identity, preamble):
        self.name = name
        self._ops = ops
        self.identity = identity
        self._preamble = preamble
        self.nin = 1
        self.nout = 1
        in_params = _get_param_info('T in0', True)
        out_params = _get_param_info('T out0', False)
        self._params = (
            in_params + out_params +
            _get_param_info(
                'CIndexer _in_ind, CIndexer _out_ind', False) +
            _get_param_info('int32 _out_clp2_size', True))
        self._input_expr = 'const type_in0_raw in0 = _raw_in0[_in_ind.get()];'
        self._output_expr = 'type_out0_raw &out0 = _raw_out0[_out_ind.get()];'
        self._routine_cache = {}

    def __call__(self, a, axis=None, dtype=None, out=None, keepdims=False):
        if not isinstance(a, cupy.ndarray):
            raise TypeError('Input type must be cupy.ndarray')
        if self.identity is None:
            assert a.size != 0
        if dtype is not None:
            dtype = numpy.dtype(dtype).type

        in_args = [a]
        if out is None:
            _check_args((a,))
            out_args = []
        else:
            _check_args((a, out))
            out_args = [out]

        in_types, out_types, routine = _guess_routine(
            self.name, self._routine_cache, self._ops, in_args, dtype)

        axis, raxis = _get_axis(axis, a.ndim)
        out_shape = _get_out_shape(a.shape, axis, raxis, keepdims)
        out_args = _get_out_args(out_args, out_types, out_shape)
        in_args, in_shape = _get_trans_args(
            in_args, axis + raxis, in_args[0].shape)

        in_indexer = carray.Indexer(in_shape)
        out_indexer = carray.Indexer(out_shape)
        out_clp2_size = 2 ** int.bit_length(int(out_indexer.size - 1))

        inout_args = _get_inout_args(
            in_args, out_args, in_indexer, out_indexer, out_clp2_size,
            self._params, True)
        args_info = _get_args_info(inout_args)

        block_size = 512
        kern = _get_simple_reduction_function(
            routine, self._params, args_info,
            in_args[0].dtype.type, out_args[0].dtype.type, out_types,
            self.name, block_size, self.identity,
            self._input_expr, self._output_expr, self._preamble, ())

        shared_mem = 32 * block_size
        if out_clp2_size > 256:
            shared_mem = 0
        # TODO(okuta) set actual size
        kern.linear_launch(max(out_indexer.size, block_size), inout_args,
                           shared_mem, block_size)

        if len(out_args) == 1:
            return out_args[0]
        return tuple(out_args)

@util.memoize(for_each_device=True)
def _get_reduction_kernel(
        params, args_info, types,
        name, block_size, reduce_type, identity, map_expr, reduce_expr,
        post_map_expr, preamble, options):
    kernel_params = _get_kernel_params(params, args_info)
    arrays = [p for p, a in six_zip(params, args_info)
              if not p.raw and a[0] is cupy.ndarray]
    type_preamble = '\n'.join(
        'typedef %s %s;' % (_get_typename(v), k)
        for k, v in types)
    input_expr = '\n'.join(
        ['const {0} {1} = _raw_{1}[_j];'.format(p.ctype, p.name)
         for p in arrays if p.is_const])
    output_expr = '\n'.join(
        ['{0} &{1} = _raw_{1}[_i];'.format(p.ctype, p.name)
         for p in arrays if not p.is_const])

    return _get_simple_reduction_kernel(
        name, block_size, reduce_type, kernel_params, identity,
        map_expr, reduce_expr, post_map_expr,
        type_preamble, input_expr, output_expr, preamble, options)

class ReductionKernel(object):

    """User-defined reduction kernel.

    This class can be used to define a reduction kernel with or without
    broadcasting.

    The kernel is compiled at an invocation of the
    :meth:`~ReductionKernel.__call__` method, which is cached for each device.
    The compiled binary is also cached into a file under the
    ``$HOME/.cupy/kernel_cache/`` directory with a hashed file name. The cached
    binary is resued by other processes.

    Args:
        in_params (str): Input argument list.
        out_params (str): Output argument list.
        map_expr (str): Mapping expression for input values.
        reduce_expr (str): Reduction expression.
        post_map_expr (str): Mapping expression for reduced values.
        identity (str): Identity value for starting the reduction.
        name (str): Name of the kernel function. It should be set for
            readability of the performance profiling.
        reduce_type (str): Type of values to be used for reduction. This type
            is used to store the special variables ``a``.
        reduce_dims (bool): If True, input arrays are reshaped without copy to
            smaller dimensions for efficiency.
        preamble (str): Fragment of the CUDA-C/C++ code that is inserted at the
            top of the cu file.
        options (tuple of str): Additional compilation options.

    """
    def __init__(self, in_params, out_params,
                 map_expr, reduce_expr, post_map_expr,
                 identity, name='reduce_kernel', reduce_type=None,
                 reduce_dims=True, preamble='', options=()):
        self.in_params = _get_param_info(in_params, True)
        self.out_params = _get_param_info(out_params, False)
        self.nin = len(self.in_params)
        self.nout = len(self.out_params)
        self.nargs = self.nin + self.nout
        self.params = (
            self.in_params + self.out_params +
            _get_param_info('CIndexer _in_ind, CIndexer _out_ind', False) +
            _get_param_info('int32 _out_clp2_size', True))
        self.identity = identity
        self.reduce_expr = reduce_expr
        self.map_expr = map_expr
        self.name = name
        self.options = options
        self.reduce_dims = reduce_dims
        self.post_map_expr = post_map_expr
        if reduce_type is None:
            self.reduce_type = self.out_params[0].ctype
        else:
            self.reduce_type = reduce_type
        self.preamble = preamble

    def __call__(self, *args, **kwargs):
        """Compiles and invokes the reduction kernel.

        The compilation runs only if the kernel is not cached. Note that the
        kernels with different argument dtypes, ndims, or axis are not
        compatible. It means that single ReductionKernel object may be compiled
        into multiple kernel binaries.

        Args:
            args: Arguments of the kernel.

        Returns:
            Arrays are returned according to the ``out_params`` argument of the
            ``__init__`` method.

        """

        out = kwargs.pop('out', None)
        axis = kwargs.pop('axis', None)
        keepdims = kwargs.pop('keepdims', False)
        if kwargs:
            raise TypeError('Wrong arguments %s' % kwargs)

        n_args = len(args)
        if n_args != self.nin and n_args != self.nargs:
            raise TypeError('Wrong number of arguments for %s' % self.name)

        out_args = list(args[self.nin:])
        if out is not None:
            if self.nout != 1:
                raise NotImplementedError('')
            if len(out_args) != 0:
                raise ValueError("cannot specify 'out' as both "
                                 "a positional and keyword argument")
            out_args = [out]

        in_args, broad_shape = _broadcast(args, self.in_params, False)
        _check_args(in_args + out_args)

        if self.identity is None:
            assert 0 in broad_shape

        cp_array = cupy.ndarray
        in_ndarray_types = tuple(
            [a.dtype.type if isinstance(a, cp_array) else None
             for a in in_args])
        out_ndarray_types = tuple(
            [a.dtype.type if isinstance(a, cp_array) else None
             for a in out_args])
        in_types, out_types, types = _decide_params_type(
            self.in_params, self.out_params,
            in_ndarray_types, out_ndarray_types)

        axis, raxis = _get_axis(axis, len(broad_shape))
        out_shape = _get_out_shape(broad_shape, axis, raxis, keepdims)
        in_args = [x if isinstance(x, cp_array) else t(x)
                   for x, t in six_zip(in_args, in_types)]
        in_args, in_shape = _get_trans_args(
            in_args, axis + raxis, broad_shape, self.in_params)
        out_args = _get_out_args_with_params(
            out_args, out_types, out_shape, self.out_params)

        in_indexer = carray.Indexer(in_shape)
        out_indexer = carray.Indexer(out_shape)
        out_clp2_size = 2 ** int.bit_length(int(out_indexer.size - 1))

        inout_args = _get_inout_args(
            in_args, out_args, in_indexer, out_indexer, out_clp2_size,
            self.params, self.reduce_dims)
        args_info = _get_args_info(inout_args)

        block_size = 512
        kern = _get_reduction_kernel(
            self.params, args_info, types,
            self.name, block_size, self.reduce_type, self.identity,
            self.map_expr, self.reduce_expr, self.post_map_expr,
            self.preamble, self.options)

        shared_mem = 32 * block_size
        if out_clp2_size > 256:
            shared_mem = 0
        # TODO(okuta) set actual size
        kern.linear_launch(max(out_indexer.size, block_size), inout_args,
                           shared_mem, block_size)
        return out_args[0]

def create_reduction_func(name, ops, routine=None, identity=None,
                          preamble=''):
    _ops = []
    for t in ops:
        if not isinstance(t, tuple):
            typ = t
            rt = routine
        else:
            typ, rt = t
            rt = tuple(i or j for i, j in six_zip(rt, routine))

        types = typ.split('->')
        if len(types) == 1:
            in_types = out_types = tuple(types)
        else:
            in_types, out_types = map(tuple, types)
        in_types = tuple([numpy.dtype(t).type for t in in_types])
        out_types = tuple([numpy.dtype(t).type for t in out_types])
        _ops.append((in_types, out_types, rt))

    return simple_reduction_function(name, _ops, identity, preamble)

_min_max_preamble = '''
struct min_max_st{
    type_in0_raw value;
    int index;
    __device__ min_max_st() : index(-1) { }
    __device__ min_max_st(type_in0_raw v) : value(v), index(0) { }
    __device__ min_max_st(type_in0_raw v, int i) : value(v), index(i) { }
};
__device__ min_max_st my_min(const min_max_st& a, const min_max_st& b) {
    if (a.index == -1) return b;
    if (b.index == -1) return a;
    return min_max_st(min(a.value, b.value));
}
__device__ min_max_st my_max(const min_max_st& a, const min_max_st& b) {
    if (a.index == -1) return b;
    if (b.index == -1) return a;
    return min_max_st(max(a.value, b.value));
}
__device__ min_max_st my_argmin(const min_max_st& a, const min_max_st& b) {
    if (a.index == -1) return b;
    if (b.index == -1) return a;
    return (a.value <= b.value) ? a : b;
}
__device__ min_max_st my_argmax(const min_max_st& a, const min_max_st& b) {
    if (a.index == -1) return b;
    if (b.index == -1) return a;
    return (a.value >= b.value) ? a : b;
}'''

amin = create_reduction_func(
    'cupy_min',
    ('?->?', 'B->B', 'h->h', 'H->H', 'i->i', 'I->I', 'l->l', 'L->L',
     'q->q', 'Q->Q', 'e->e', 'f->f', 'd->d'),
    ('min_max_st(in0)', 'my_min(a, b)', 'out0 = a.value', 'min_max_st'),
    None, _min_max_preamble)

amax = create_reduction_func(
    'cupy_max',
    ('?->?', 'B->B', 'h->h', 'H->H', 'i->i', 'I->I', 'l->l', 'L->L',
     'q->q', 'Q->Q', 'e->e', 'f->f', 'd->d'),
    ('min_max_st(in0)', 'my_max(a, b)', 'out0 = a.value', 'min_max_st'),
    None, _min_max_preamble)

argmin = create_reduction_func(
    'cupy_argmin',
    ('?->l', 'B->l', 'h->l', 'H->l', 'i->l', 'I->l', 'l->l', 'L->l',
     'q->l', 'Q->l', 'e->l', 'f->l', 'd->l'),
    ('min_max_st(in0, _J)', 'my_argmin(a, b)', 'out0 = a.index', 'min_max_st'),
    None, _min_max_preamble)

argmax = create_reduction_func(
    'cupy_argmax',
    ('?->l', 'B->l', 'h->l', 'H->l', 'i->l', 'I->l', 'l->l', 'L->l',
     'q->l', 'Q->l', 'e->l', 'f->l', 'd->l'),
    ('min_max_st(in0, _J)', 'my_argmax(a, b)', 'out0 = a.index', 'min_max_st'),
    None, _min_max_preamble)


"""
Module description:

"""

__version__ = '0.1'
__author__ = 'Vito Walter Anelli, Claudio Pomo'
__email__ = 'vitowalter.anelli@poliba.it, claudio.pomo@poliba.it'

import numpy as np
from ast import literal_eval as make_tuple
from tqdm import tqdm

from elliot.dataset.samplers import pointwise_pos_neg_sampler as pws
from elliot.recommender.neural.NeuMF.neural_matrix_factorization_model import NeuralMatrixFactorizationModel
from elliot.recommender.recommender_utils_mixin import RecMixin
from elliot.utils.write import store_recommendation

from elliot.recommender.base_recommender_model import BaseRecommenderModel
from elliot.recommender.base_recommender_model import init_charger

np.random.seed(42)

class NeuMF(RecMixin, BaseRecommenderModel):
    r"""
    Neural Collaborative Filtering

    For further details, please refer to the `paper <https://arxiv.org/abs/1708.05031>`_

    Args:
        mf_factors: Number of MF latent factors
        mlp_factors: Number of MLP latent factors
        mlp_hidden_size: List of units for each layer
        lr: Learning rate
        dropout: Dropout rate
        is_mf_train: Whether to train the MF embeddings
        is_mlp_train: Whether to train the MLP layers

    To include the recommendation model, add it to the config file adopting the following pattern:

    .. code:: yaml

      models:
        NeuMF:
          meta:
            save_recs: True
          epochs: 10
          batch_size: 512
          mf_factors: 10
          mlp_factors: 10
          mlp_hidden_size: (64,32)
          lr: 0.001
          dropout: 0.0
          is_mf_train: True
          is_mlp_train: True
    """
    @init_charger
    def __init__(self, data, config, params, *args, **kwargs):

        self._random = np.random

        self._sampler = pws.Sampler(self._data.i_train_dict)

        self._params_list = [
            ("_learning_rate", "lr", "lr", 0.001, None, None),
            ("_mf_factors", "mf_factors", "mffactors", 10, int, None),
            ("_mlp_factors", "mlp_factors", "mlpfactors", 10, int, None),
            ("_mlp_hidden_size", "mlp_hidden_size", "mlpunits", "(64,32)", lambda x: list(make_tuple(str(x))), lambda x: self._batch_remove(str(x), " []").replace(",", "-")),
            ("_dropout", "dropout", "drop", 0, None, None),
            ("_is_mf_train", "is_mf_train", "mftrain", True, None, None),
            ("_is_mlp_train", "is_mlp_train", "mlptrain", True, None, None),
        ]
        self.autoset_params()

        if self._batch_size < 1:
            self._batch_size = self._data.transactions

        self._ratings = self._data.train_dict
        self._sp_i_train = self._data.sp_i_train
        self._i_items_set = list(range(self._num_items))

        self._model = NeuralMatrixFactorizationModel(self._num_users, self._num_items, self._mf_factors,
                                                     self._mlp_factors, self._mlp_hidden_size,
                                                     self._dropout, self._is_mf_train, self._is_mlp_train,
                                                     self._learning_rate)

    @property
    def name(self):
        return "NeuMF"\
               + "_e:" + str(self._epochs) \
               + "_bs:" + str(self._batch_size) \
               + f"_{self.get_params_shortcut()}"

    def train(self):
        if self._restore:
            return self.restore_weights()

        best_metric_value = 0
        early_stopping = 5

        for it in range(self._epochs):
            loss = 0
            steps = 0
            with tqdm(total=int(self._data.transactions // self._batch_size), disable=not self._verbose) as t:
                for batch in self._sampler.step(self._data.transactions, self._batch_size):
                    steps += 1
                    loss += self._model.train_step(batch)
                    t.set_postfix({'loss': f'{loss.numpy() / steps:.5f}'})
                    t.update()

            if not (it + 1) % self._validation_rate:
                recs = self.get_recommendations(self.evaluator.get_needed_recommendations())
                result_dict = self.evaluator.eval(recs)
                self._results.append(result_dict)

                print(f'Epoch {(it + 1)}/{self._epochs} loss {loss/steps:.5f}')

                if self._results[-1][self._validation_k]["val_results"][self._validation_metric] > best_metric_value:
                    print("******************************************")
                    best_metric_value = self._results[-1][self._validation_k]["val_results"][self._validation_metric]
                    early_stopping = 5
                    if self._save_weights:
                        self._model.save_weights(self._saving_filepath)
                    if self._save_recs:
                        store_recommendation(recs, self._config.path_output_rec_result + f"{self.name}-it:{it + 1}.tsv")
                else:
                    early_stopping -= 1
                    if early_stopping == 0:
                        print('Reached Early Stopping Condition at Epoch {0}\n\tEXIT'.format(it + 1))
                        break

    def get_recommendations(self, k: int = 100):
        predictions_top_k = {}
        for index, offset in enumerate(range(0, self._num_users, self._batch_size)):
            offset_stop = min(offset + self._batch_size, self._num_users)
            predictions = self._model.get_recs(
                (
                    np.repeat(np.array(list(range(offset, offset_stop)))[:, None], repeats=self._num_items, axis=1),
                    np.array([self._i_items_set for _ in range(offset, offset_stop)])
                 )
            )
            v, i = self._model.get_top_k(predictions, self.get_train_mask(offset, offset_stop), k=k)
            items_ratings_pair = [list(zip(map(self._data.private_items.get, u_list[0]), u_list[1]))
                                  for u_list in list(zip(i.numpy(), v.numpy()))]
            predictions_top_k.update(dict(zip(map(self._data.private_users.get,
                                                  range(offset, offset_stop)), items_ratings_pair)))
        return predictions_top_k


from functools import wraps
import importlib
import logging

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

from core.config import cfg
from model.roi_pooling.functions.roi_pool import RoIPoolFunction
#from model.roi_crop.functions.roi_crop import RoICropFunction
#from modeling.roi_xfrom.roi_align.functions.roi_align import RoIAlignFunction
import modeling.rpn_heads as rpn_heads
import modeling.fast_rcnn_heads as fast_rcnn_heads
import modeling.mask_rcnn_heads as mask_rcnn_heads
import modeling.keypoint_rcnn_heads as keypoint_rcnn_heads
import modeling.dispSeg_heads as dispSeg_heads
import modeling.semseg_heads as semseg_heads
import utils.blob as blob_utils
import utils.net as net_utils
import utils.Lossfuction as Lossfuction
import utils.resnet_weights_helper as resnet_utils
from lib.nn import SynchronizedBatchNorm2d
import pynvml
import cv2
import modeling.CRL as CRL
from torchvision.utils import make_grid
from tensorboardX import SummaryWriter
from torch.utils.checkpoint import checkpoint
import numpy as np
logger = logging.getLogger(__name__)

pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
def get_func(func_name):
    """Helper to return a function object by name. func_name must identify a
    function in this module or the path to a function relative to the base
    'modeling' module.
    """
    if func_name == '':
        return None
    try:
        parts = func_name.split('.')
        # Refers to a function in this module
        if len(parts) == 1:
            return globals()[parts[0]]
        # Otherwise, assume we're referencing a module under modeling
        module_name = 'modeling.' + '.'.join(parts[:-1])
        module = importlib.import_module(module_name)
        return getattr(module, parts[-1])
    except Exception:
        logger.error('Failed to find function: %s', func_name)
        raise

def compare_state_dict(sa, sb):
    if sa.keys() != sb.keys():
        return False
    for k, va in sa.items():
        if not torch.equal(va, sb[k]):
            return False
    return True

def check_inference(net_func):
    @wraps(net_func)
    def wrapper(self, *args, **kwargs):
        if not self.training:
            if cfg.PYTORCH_VERSION_LESS_THAN_040:
                return net_func(self, *args, **kwargs)
            else:
                with torch.no_grad():
                    return net_func(self, *args, **kwargs)
        else:
            raise ValueError('You should call this function only on inference.'
                              'Set the network in inference mode by net.eval().')

    return wrapper

class SegmentationModuleBase(nn.Module):
    def __init__(self):
        super(SegmentationModuleBase, self).__init__()
        self.iter = 0

    def pixel_acc(self, preds, label):
        #label = label.view(-1).contiguous()
        valid = (label != 255).long()
        acc_sum = torch.sum(valid*(preds == label).long())
        pixel_sum = torch.sum(valid)
        acc = acc_sum.float() / (pixel_sum.float() + 1e-10)
        return acc

class disparityregression_bat(nn.Module):
    def __init__(self):
        super(disparityregression, self).__init__()
        maxdisp = cfg.DISP.MAX_DISPLACEMENT
        self.disp = Variable(torch.Tensor(np.reshape(np.array(range(maxdisp)),[1,maxdisp])).cuda(), requires_grad=False)
        self.zeros = Variable(torch.Tensor(np.zeros(np.array(range(maxdisp)),[1,maxdisp])).cuda(), requires_grad=False)
        def forward(self, x, semseg):
            zeros = self.zeros
            for i in range(x.size()[0]):
                zeros = torch.cat([zeros, self.zeros.repeat(1,semseg[i])],dim=0)
            disp = self.disp.repeat(x.size()[0],1)
            scans = torch.cat([zeros, disp], dim=1)
            for i in range(x.size()[0]):
                zeros = torch.cat([zeros, self.zeros.repeat(1,(cfg.MODEL.NUM_CLASSES-semseg[i]))],dim=0)
            scans = torch.cat([disp, zeros], dim=1)
            out = torch.sum(x*disp,1)
            return out

class disparityregression(nn.Module):
    def __init__(self):
        super(disparityregression, self).__init__()
        maxdisp = cfg.DISP.MAX_DISPLACEMENT
        self.disp= Variable(torch.Tensor(np.reshape(np.array(range(maxdisp)),[1,maxdisp])).cuda(), requires_grad=False)
    def forward(self, x):
        disp = self.disp.repeat(x.size()[0],cfg.MODEL.NUM_CLASSES)
        out = torch.sum(x*disp,1)
        return out

class Generalized_3DSD(SegmentationModuleBase):
    def __init__(self):
        super(SegmentationModuleBase, self).__init__()

        # For cache
        self.mapping_to_detectron = None
        self.orphans_in_detectron = None
        self.iter=0
        self.draw=False
        builder = semseg_heads.ModelBuilder()

        #define encoder
        if cfg.SEM.USE_RESNET:
            self.encoder = get_func(cfg.MODEL.CONV_BODY)()
        else:
            builder = semseg_heads.ModelBuilder()
            self.encoder = builder.build_encoder(
                arch=cfg.SEM.ARCH_ENCODER,
                fc_dim=cfg.SEM.FC_DIM,
                weights=cfg.RESNETS.IMAGENET_PRETRAINED_WEIGHTS
                )

        #define shape weights
        self.decoder = builder.build_decoder(
                arch=cfg.SEM.DECODER_TYPE,
                fc_dim=cfg.SEM.FC_DIM,
                num_class=cfg.MODEL.NUM_CLASSES,
                use_softmax=not self.training,
                weights='')

        ##define 3DSD loss
        self.crit = nn.NLLLoss(ignore_index=255*cfg.DISP.MAX_DISPLACEMENT+1)
        self.SmoothL1Loss = nn.SmoothL1Loss(reduction='none') #elementwise_mean
        self.disp_gain = 1.0*cfg.DISP.EXPECT_MAXDISP / cfg.DISP.MAX_DISPLACEMENT
        print ('epe of each pixel is set to %.2f'%self.disp_gain)
        self.BCELoss = nn.BCELoss(reduction='elementwise_mean')
    def segdisp_loss(self, pred, label_info, lbl_i):
        semseg_label = label_info['{}_{}'.format(cfg.SEM.OUTPUT_PREFIX, lbl_i)]
        '''
        disp_label   = label_info['{}_{}'.format(cfg.DISP.OUTPUT_PREFIX, lbl_i)]
        h, w = semseg_label.shape[-2:]
        disp_label = F.adaptive_max_pool2d(disp_label, (h, w))
        disp_label_match = torch.round(disp_label / self.disp_gain).long()
        assert torch.sum(disp_label_match>cfg.DISP.MAX_DISPLACEMENT)==0, 'errors in disp match'
        hybrid_label = semseg_label*cfg.DISP.MAX_DISPLACEMENT + disp_label_match
        ####### ignore wrong label
        hybrid_label[semseg_label==255] = 255 * cfg.DISP.MAX_DISPLACEMENT+1
        hybrid_label[disp_label_match<0] = 255 * cfg.DISP.MAX_DISPLACEMENT+1
        '''
        pred_semseg, pred_ids  = torch.max(pred, dim=1)
        pred_ids = pred_ids // cfg.DISP.MAX_DISPLACEMENT
        mask_expect = (semseg_label == pred_ids).float()
        mask_expect.requires_grad = False
        loss = self.BCELoss(pred_semseg, mask_expect)
        return loss
    def segdisp_pred(self, pred, label_info, lbl_i=0):
        #using 3D softmax for predict
        #The shape of pred is (b, num_classes * max_disp, h, w)
        max_scans, max_ids = torch.max(pred, dim=1)
        #semseg pred
        semseg_pred = max_ids // cfg.DISP.MAX_DISPLACEMENT
        #disp pred
        disp_pred = (max_ids % cfg.DISP.MAX_DISPLACEMENT + self.disp_gain).float() * max_scans
        #disp_scans = label_info['disp_scans'].repeat(pred.shape[0]//cfg.TRAIN.IMS_PER_BATCH, 1)
        #disp_zeros = label_info['disp_zeros'].repeat(pred.shape

        #disp_pred = torch.sum(disp_scans * pred, dim=1)
        #scans = torch.stack([semseg_pred*cfg.DISP.MAX_DISPLACEMENT + i for i in range(1, cfg.DISP.MAX_DISPLACEMENT+1)], dim=1)
        #disp_pred = pred[scans]
        return semseg_pred, disp_pred

    def disp_loss(self, pred, label_info, label_i=0, sparse=True, shape=(30,30)):
        targe = label_info['disp_label_0']
        if sparse:
            targe = F.adaptive_max_pool2d(targe, shape)
        #calucate disp output
        disp_scans = label_info['disp_scans'].repeat(pred.shape[0]//cfg.TRAIN.IMS_PER_BATCH, 1,1,1)
        disp_pred = torch.sum(pred*disp_scans, dim=1)
        #loss and pred
        EPE_map = self.SmoothL1Loss(disp_pred, targe)
        epe_pred = torch.abs(disp_pred - targe)
        #ignore false disp values
        positive = targe.ge(0)
        EPE_map = torch.masked_select(EPE_map, positive)
        epe_pred = torch.masked_select(epe_pred, positive)
        #normlization
        loss = EPE_map.mean()
        epe_pred = epe_pred.mean()
        return loss, epe_pred

    def forward(self,data,**label_info):
        return_dict = {}
        pred = self.decoder(self.encoder(data, return_feature_maps=True))
        if self.training: # training
            return_dict['losses'] = {}
            return_dict['metrics'] = {}
            if not isinstance(pred, list):
                pred = [pred]
            for label_i, pred_i in enumerate(pred):
                # use 3D interp or not
                stride = 2**cfg.SEM.DOWNSAMPLE[label_i]
                pred_i = nn.functional.interpolate(pred_i, (cfg.DISP.MAX_DISPLACEMENT,
                    cfg.SEM.INPUT_SIZE[0]//stride, cfg.SEM.INPUT_SIZE[1]//stride),
                    mode='trilinear',align_corners=False)
                _, num_class, max_d, hei, wid = pred_i.shape
                #transfer 3D to 2D maps
                pred_i = pred_i.view(-1, num_class * max_d, hei, wid).contiguous()
                pred_i = nn.functional.softmax(pred_i, dim=1)
                #3D sematic disparity loss
                loss = self.segdisp_loss(pred_i, label_info, label_i)
                return_dict['losses']['loss_segdisp'] = loss
                #stat for predict
                pred_semseg, pred_disp = self.segdisp_pred(pred_i, label_info, label_i)
                return_dict['disp_image']=pred_disp
                return_dict['semseg_image']=pred_semseg
                return_dict['metrics']['pixel_accuary'] = \
                        self.pixel_acc(pred_semseg, label_info['{}_{}'.format(cfg.SEM.OUTPUT_PREFIX, label_i)])
                return_dict['metrics']['epe'] = \
                        self.disp_loss(pred_disp, label_info['{}_{}'.format(cfg.DISP.OUTPUT_PREFIX, label_i)],
                                shape=(hei, wid))
            # pytorch0.4 bug on gathering scalar(0-dim) tensors
            for k, v in return_dict['losses'].items():
                return_dict['losses'][k] = v.unsqueeze(0)
            for k, v in return_dict['metrics'].items():
                return_dict['metrics'][k] = v.unsqueeze(0)
        else: # inference
            pred_semseg, pred_disp = self.segdisp_pred(pred)
            return_dict['pred_semseg'] = pred_semseg
            return_dict['pred_disp'] = pred_disp

        return return_dict

    @check_inference
    def convbody_net(self, data):
        """For inference. Run Conv Body only"""
        blob_conv = self.Conv_Body(data)
        return blob_conv

    def _add_loss(self, return_dict, key, value):
        """Add loss tensor to returned dictionary"""
        return_dict['losses'][key] = value


from itertools import chain

import numpy as np
import xarray as xr
import geopandas as gpd
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt

RED = "#A01914"
BLUE = "#4F6DB8"
RED_TO_BLUE = [ # from https://gka.github.io using lightness correction
    '#002d6e', '#375aa2', '#6f8ad1', '#a7bffa',
    '#f5f5f5', '#fdad97', '#e36b55', '#b23125', '#720000'
]
CMAP = matplotlib.colors.LinearSegmentedColormap.from_list("signature-BlRd", RED_TO_BLUE)
NORM = matplotlib.colors.Normalize(vmin=-1, vmax=3)
LAYER_FONT_WEIGHT = "medium"
EDGE_WIDTH = 0.06
EDGE_COLOR = "white"

EPSG_3035_PROJ4 = "+proj=laea +lat_0=52 +lon_0=10 +x_0=4321000 +y_0=3210000 +ellps=GRS80 +units=m +no_defs "

MAP_MIN_X = 2500000
MAP_MIN_Y = 1400000
MAP_MAX_X = 6600000
MAP_MAX_Y = 5500000

CONTINENTAL_SCENARIO = "continental-autarky-100-continental-grid"
NATIONAL_SCENARIO = "national-autarky-100-national-grid"
REGIONAL_SCENARIO = "regional-autarky-100-regional-grid"
DEMAND_TECH = "demand_elec"

def plot_map(path_to_national_shapes, path_to_regional_shapes, connected_regions,
             path_to_aggregated_results, path_to_shapes, path_to_plot):
    """Plot maps of results."""
    fig = plt.figure(figsize=(4.41, 2.2), constrained_layout=True)
    axes = fig.subplots(1, 3, gridspec_kw={'width_ratios': [30, 30, 1]}).flatten()
    shapes = (gpd.read_file(path_to_shapes)
                 .to_crs(EPSG_3035_PROJ4)
                 .rename(columns={"id": "locs"})
                 .set_index("locs")
                 .rename(index=lambda idx: idx.replace(".", "-")))

    national = gpd.read_file(path_to_national_shapes).to_crs(EPSG_3035_PROJ4).set_index("country_code")
    regional = (gpd.read_file(path_to_regional_shapes)
                   .to_crs(EPSG_3035_PROJ4)
                   .set_index("id")
                   .rename(index=lambda idx: idx.replace(".", "-")))

    cost_reader = CostReader(
        path_to_aggregated_results=path_to_aggregated_results,
        connected_regions=connected_regions,
        country_codes=shapes.country_code.to_xarray()
    )
    national["cost"] = cost_reader.read(scenario=NATIONAL_SCENARIO, level="national")
    regional["cost"] = cost_reader.read(scenario=REGIONAL_SCENARIO, level="regional")

    base_costs = cost_reader.read(scenario=CONTINENTAL_SCENARIO, level="continental")
    national["cost"] = national["cost"] / base_costs
    regional["cost"] = regional["cost"] / base_costs

    total_costs_national = cost_reader.read(scenario=NATIONAL_SCENARIO, level="continental") / base_costs
    total_costs_regional = cost_reader.read(scenario=REGIONAL_SCENARIO, level="continental") / base_costs

    _plot_layer(national, total_costs_national, "National supply and balancing", "A", axes[0])
    _plot_layer(regional, total_costs_regional, "Regional supply and balancing", "B", axes[1])

    _plot_colorbar(fig, axes[2].inset_axes([0, 0.175, 1, 0.65]))
    axes[2].axis("off")

    fig.savefig(path_to_plot, pil_kwargs={"compression": "tiff_lzw"})

def _plot_layer(units, total_cost, layer_name, panel_id, ax):
    ax.set_aspect('equal')
    units.plot(
        linewidth=EDGE_WIDTH,
        edgecolor=EDGE_COLOR,
        column="cost",
        vmin=NORM.vmin,
        vmax=NORM.vmax,
        cmap=CMAP,
        ax=ax
    )
    ax.set_xlim(MAP_MIN_X, MAP_MAX_X)
    ax.set_ylim(MAP_MIN_Y, MAP_MAX_Y)
    ax.set_xticks([])
    ax.set_yticks([])
    sns.despine(ax=ax, top=True, bottom=True, left=True, right=True)

    ax.annotate(
        layer_name,
        xy=[0.5, -0.05],
        xycoords='axes fraction',
        ha='center',
        va='center',
        weight=LAYER_FONT_WEIGHT
    )
    ax.set_title(panel_id, loc="left")

def _plot_colorbar(fig, axis):
    s_m = matplotlib.cm.ScalarMappable(cmap=CMAP, norm=NORM)
    cmap = s_m.get_cmap()
    colors = cmap(np.linspace(1 / 4, 1, cmap.N))
    cmap = matplotlib.colors.LinearSegmentedColormap.from_list('cut_jet', colors)
    s_m = matplotlib.cm.ScalarMappable(cmap=cmap, norm=matplotlib.colors.Normalize(vmin=0, vmax=3))
    s_m.set_array([])
    cbar = fig.colorbar(s_m, cax=axis)
    cbar.set_ticks([0, 1, 2, 3])
    cbar.set_ticklabels(["0", "1", "2", "≥3"])
    cbar.outline.set_linewidth(0)
    cbar.ax.set_ylabel('Relative system cost', rotation=90)

class CostReader:

    def __init__(self, path_to_aggregated_results, country_codes, connected_regions):
        self.__results = xr.open_dataset(path_to_aggregated_results)
        self.__country_codes = country_codes
        self.__connected_regions = connected_regions

    def read(self, scenario, level):
        assert level in ["continental", "national", "regional"]
        cost = self.__results.cost.sel(scenario=scenario).sum("techs")
        demand = -self.__results.energy_cap.sel(techs=DEMAND_TECH, scenario=scenario)
        if level == "continental":
            cost = cost.sum(dim="locs").item()
            demand = demand.sum(dim="locs").item()
        elif level == "national":
            cost = (
                cost
                .groupby(self.__country_codes.sortby(cost.locs))
                .sum("locs")
                .to_series()
            )
            demand = (
                demand
                .groupby(self.__country_codes.sortby(demand.locs))
                .sum("locs")
                .to_series()
            )
        elif level == "regional":
            cost = self._postprocess_combined_regions(cost).to_series()
            demand = self._postprocess_combined_regions(demand).to_series()
        return (cost / demand)

    def _postprocess_combined_regions(self, da):
        region_pairs = list(self.__connected_regions.items())
        regions = list(chain(*region_pairs))
        if all([region in da.locs for region in regions]):
            # config is valid, and resolution is regional. apply config
            for region1, region2 in region_pairs:
                sum_data = da.sel(locs=region1) + da.sel(locs=region2)
                da.loc[dict(locs=region1)] = sum_data
                da.loc[dict(locs=region2)] = sum_data
            return da
        elif all([region not in da.locs for region in regions]):
            # config is valid, but resolution is not regional. do nothing
            return da
        else:
            raise ValueError("Config of connected regions is invalid.")

if __name__ == "__main__":
    plot_map(
        path_to_shapes=snakemake.input.shapes,
        path_to_national_shapes=snakemake.input.national_shapes,
        path_to_regional_shapes=snakemake.input.regional_shapes,
        path_to_aggregated_results=snakemake.input.results,
        connected_regions=snakemake.params.connected_regions,
        path_to_plot=snakemake.output[0]
    )


import threading
#from multiprocessing import Pool
import time
import tkinter as Tk
import json
import cv2
import numpy as np
from datetime import datetime
import os
import matplotlib.pyplot as plt
import pandas as pd
import glob
import logging
import myutill
from sms_utill import sms

class SnowDepthReader(object):

    def __init__(self, camera):

        logging.basicConfig(level=logging.INFO, format='[%(asctime)s %(levelname)s] (%(threadName)-9s) %(message)s',)
        self.log = logging.getLogger(__name__)
        self.drawing = False
        self.time = 0
        self.image_file = ''
        self.img = np.empty(0)
        self.imgt = np.empty(0)
        self.hh = 0
        self.hl = 0
        self.wh = 0
        self.wl = 0
        self.ix = 0
        self.iy = 0
        for n, c in camera.items():
            self.name = n
            self.camera = c
            self.rotation = c['rotation']
        with open('notification_list.json', 'r') as f:
            self.notification_list = json.load(f)

        self.read_snow_stake()

    def download_new_image(self):

        def local_address(url):
            local = url.find('http') < 0
            return local

        if local_address(str(self.camera['url'])):
            img_list = glob.glob(str(self.camera['url']+'*.jpg'))
            last_img = max(img_list, key=os.path.getctime)
            self.img = cv2.imread(last_img)
        else:
            self.img = myutill.download_last_image(str(self.camera['url']))

        return

    # this functions rotate the image so the snow stake will be vertical and convert to a threshold image.
    def arrange_img(self):

        rows, cols = self.img.shape[:2]
        img = cv2.cvtColor(self.img, cv2.COLOR_BGR2GRAY)
        M = cv2.getRotationMatrix2D((cols/2, rows/2), self.camera['rotation'], 1)
        img = cv2.warpAffine(img, M, (cols, rows))
        img = cv2.GaussianBlur(img, (5, 5), 0)
        self.img = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 5, 2)
        img = cv2.GaussianBlur(img, (7, 7), 0)

    # This function return the roi (where is the snow stake in the image)
    def get_roi(self):

        # read the rois from the station
        if 'roi' in self.camera:
            return self.camera['roi']

    def read_snow_stake(self):

        """

        :rtype:
        """
        self.download_new_image()
        self.arrange_img()
        roi_bound = self.get_roi()
        roi = self.img[int(roi_bound['hl']):int(roi_bound['hh']), int(roi_bound['wl']):int(roi_bound['wh'])]
        try: # Mostly white stick
            snow_line = np.where(roi<roi.mean()-roi.std())[0][-1]
        except IndexError: # Mostly dark stick
            snow_line = np.where(roi<roi.mean()+roi.std())[0][-1]
        rows, cols = roi.shape
        depth = self.camera['max depth'] * (rows-snow_line) / rows
        #Round depth to the nearest 1/2 inch:
        depth = round(depth * 2)/2
        if os.path.isfile(self.camera['output file']):
            old_sh = pd.read_csv(self.camera['output file'])
        else:
            old_sh = pd.DataFrame(columns=['Time', 'HS'])
        new_sh = pd.DataFrame([{'Time': str(datetime.now())[:-10], 'HS': depth}])
        sh = pd.concat([new_sh, old_sh], ignore_index=True)[['Time', 'HS']]
        sh.to_csv(self.camera['output file'], index=False)
        depths = sh.iloc[:48, :][::-1]
        fig, _ = plt.subplots(figsize=(12, 7))
        depths.set_index('Time').plot(grid=True, rot=30)
        plt.title('{0} snow depth'.format(self.name))
        plt.ylim(0, depths['HS'].max()*1.5)#self.camera['max depth']+2)
        plt.ylabel('Snow depth in in.')
        plt.tight_layout()
        plt.savefig('../data/{0}/{0}_snow_depth.png'.format(self.name))
        if key in self.notification_list:
            if (time.time()-self.notification_list[key]['delay'] > 6*3600):
                if key == 'Loveland':
                    if depth > 4 and (time.time()-self.notification_list[key]['delay'] > 6*3600):
                        message = sms()
                        if datetime.now().hour > 16 or datetime.now().hour < 7:
                            for name, number in self.notification_list[key].items():
                                if name == 'delay':
                                    continue
                                message.send(to=number, text='Hi {0}, The {1} snow stake is at {2}'.format(name, key, depth))
                            self.notification_list[key]['delay'] = time.time()
                            with open('notification_list.json', 'w') as f:
                                json.dump(self.notification_list, f)
                if key == 'Monument':
                    if depth - sh.iloc[-24, 1] > 6:
                        message = sms()
                        for name, number in self.notification_list[key].items():
                            if name == 'delay':
                                continue
                            message.send(to=number, text='Hi {0}, The {1} snow stake is at {2}'.format(name, key, depth))
                        self.notification_list[key]['delay'] = time.time()
                        with open('notification_list.json', 'w') as f:
                            json.dump(self.notification_list, f)

'''
Callback function for image angle setting slide bar
'''
def on_change():

    pass

class NewCamera(object):

    def __init__(self, url=''):

        self.drawing = False
        self.hh = 0
        self.hl = 0
        self.wh = 0
        self.wl = 0
        self.ix = 0
        self.iy = 0
        self.rotation = 0
        self.zoom = 1
        self.name = ''
        self.camera = {}
        self.cam = {}
        self.root = Tk.Tk()
        self.root.wm_title('Enter new camera')

        self.label1 = Tk.Label(self.root, text='Name:')
        self.E1 = Tk.Entry(self.root, bd=5)

        defalult_url = Tk.StringVar(self.root, url)
        self.label2 = Tk.Label(self.root, text='URL:')
        self.E2 = Tk.Entry(self.root, textextvariable=defalult_url, bd=5)

        self.label3 = Tk.Label(self.root, text='Max height:')
        self.E3 = Tk.Entry(self.root, bd=5)

        self.submit = Tk.Button(self.root, text="Submit", command=self.save_new_camera)

        self.label1.pack()
        self.E1.pack()
        self.label2.pack()
        self.E2.pack()
        self.label3.pack()
        self.E3.pack()
        self.submit.pack(side=Tk.BOTTOM)

        self.root.mainloop()

    def download_image(self):

        def local_address(url):
            return url.find('http') != 0

        url = self.cam['url']
        if local_address(str(self.cam['url'])):
            if os.path.splitext(self.cam['url'])[1] not in ['.jpg', '.png']:
                img_list = glob.glob(str(self.cam['url']+'*.jpg'))
                url = max(img_list, key=os.path.getctime)
            self.img = cv2.imread(url)
        else:
            self.img = myutill.download_image(self.cam['url'])
        return self.img

    def draw_roi(self, event, x, y, flags, param):

        if self.drawing:
            self.imgt = cv2.resize(self.img.copy(), None, fx=self.zoom, fy=self.zoom, interpolation=cv2.INTER_AREA)

        if event == cv2.EVENT_LBUTTONDOWN:
            self.drawing = True
            self.ix = x
            self.iy = y

        if event == cv2.EVENT_MOUSEMOVE:
            if self.drawing:
                print(x, y)
                self.imgt = cv2.rectangle(self.imgt, (self.ix, self.iy), (x, y), (0, 0, 255), 2)

        if event == cv2.EVENT_LBUTTONUP:
            self.drawing = False
            cv2.rectangle(self.imgt, (self.ix, self.iy), (x, y), (0, 0, 255), 2)
            self.wh = max(self.ix, x)/self.zoom
            self.wl = min(self.ix, x)/self.zoom
            self.hh = max(self.iy, y)/self.zoom
            self.hl = min(self.iy, y)/self.zoom

    def set_roi(self):

        """

        :rtype: object
        """
        self.download_image()

        cv2.namedWindow('set camera images')
        cv2.createTrackbar('img_rotation', 'set camera images', 0, 8, on_change)
        cv2.setMouseCallback('set camera images', self.draw_roi)
        self.imgt = self.img.copy()
        rows, cols = self.img.shape[:2]
        rotate = 0
        roi = {}
        # Set the image rotation
        while True:
            cv2.imshow('set camera images', self.imgt)
            k = cv2.waitKey(30) & 0xFF
            if k == 27:
                break
            rotate = (cv2.getTrackbarPos('img_rotation', 'set camera images') - 4)
            M = cv2.getRotationMatrix2D((cols/2, rows/2), rotate, 1)
            self.imgt = cv2.warpAffine(self.img, M, (cols, rows))

        # Set anchore
        self.imf = self.imgt.copy()
        self.zoom = 1
        while True:
            cv2.imshow('set camera images', self.imgt)
            k = cv2.waitKey(30) & 0xFF
            if k == 27:
                break
            if k == ord('o'):
                self.zoom = self.zoom/1.5
                self.imgt = cv2.resize(self.img.copy(), None, fx=self.zoom, fy=self.zoom, interpolation=cv2.INTER_AREA)
            anchore_ul_x, anchore_ul_y = self.wl, self.hl
            cv2.imwrite(self.name + '_anc.jpg', self.imgt[self.hl:self.hh, self.wl:self.wh,:])
            anc_h, anc_w = self.hh-self.hl, self.wh-self.wl

        # Set ROI
        self.img = self.imgt.copy()
        self.zoom = 1
        while True:
            cv2.imshow('set camera images', self.imgt)
            k = cv2.waitKey(30) & 0xFF
            if k == 27:
                break
            if k == ord('o'):
                self.zoom = self.zoom/1.5
                self.imgt = cv2.resize(self.img.copy(), None, fx=self.zoom, fy=self.zoom, interpolation=cv2.INTER_AREA)
            roi = {'upper_left_y': self.hh, 'hl': self.hl, 'wh': self.wh, 'wl': self.wl}

        cv2.destroyAllWindows()
        self.cam['rotation'] = rotate
        self.cam['anchore'] = self.name + '_anc.jpg'
        self.cam['roi'] = roi

    def save_new_camera(self):

        if os.path.isfile('cameras.json'):
            with open('cameras.json', 'r') as f:
                cameras = json.load(f)
        else:
            cameras = {}
        url = self.E2.get()
        self.cam['url'] = url#[:url.rfind('/')+1]
        self.cam['max depth'] = float(self.E3.get())
        self.cam['ext'] = 'jpg'
        self.name = self.E1.get()
        self.cam['output file'] = os.path.join('../data/{0}/{0}.dat'.format(self.name))
        cameras[self.name] = self.cam
        self.root.destroy()
        self.set_roi()
        with open('cameras.json', 'w') as f:
            json.dump(cameras, f)
        return

    def get_new_camera(self):

        return self.camera

# A worker thread class that listen to new camera entry
class NewCameraEntryListener(object):

    def __init__(self, interval=1, cameras = {}):
        """ Constructor

        :type interval: int, list of cameras
        :param interval: Check interval, in seconds, cameras: dic with the cameras to add the new one
        """
        self.interval = interval
        self.cameras = cameras

        thread = threading.Thread(target=self.run, args=())
        thread.daemon = True                            # Daemonize thread
        thread.start()                                  # Start the execution

    def run(self):

        while True:
            time.sleep(self.interval)
            enc = input('\nClick n and Enter to add a new camera \n')
            if enc == 'n':
                new_cam = NewCamera().get_new_camera()
                for name in new_cam:
                    self.cameras[name] = new_cam[name]
                    with open('cameras.json', 'w') as fp:
                        json.dump(cameras, fp)

def send_to_snow_readers(camera):

    _ = SnowDepthReader(camera)

if __name__ =='__main__':

    cam = NewCamera()

    '''
    '''

#cam = NewCamera()
#while True:
    with open('cameras.json', 'r') as fp:
        cameras = json.load(fp)
    for key, cam in cameras.items():
        camera = key[0].upper() + key[1:].lower()
        os.makedirs('../data/{}'.format(camera), exist_ok=True)
        #os.makedirs('/home/www/html/snow_stake_reader/data/{0}/'.format(camera), exist_ok=True)
        reader = SnowDepthReader({key: cam})

#    cameras_list = [{name: cam} for name, cam in cameras.items()]
#    print(cameras_list)
#    pool = Pool(4)
#    pool.map_async(SnowDepthReader, cameras_list)
#    pool.close()
#    pool.join()
#
#    time.sleep(120) # wait 15 min between reads

'''
read_dispatcher = ReadDispatcher()
read_dispatcher.daemon = True
read_dispatcher.start()

# enter new camera
while True:
  time.sleep(1)
  i = raw_input('\nPress n + Enter to enter a new camera\n')
  if i == 'n':
    with open('cameras.json', 'r') as fp:
          cameras = json.load(fp)

    new_cam = NewCamera().get_new_camera()
    for name in new_cam:
        cameras[name] = new_cam[name]

    with open('cameras.json', 'w') as fp:
        json.dump(cameras, fp)
'''