###############################################################################
# Project: ApRES
# Purpose: Classes to encapsulate an ApRES file
# Author:  Paul M. Breen
# Date:    2018-09-24
###############################################################################

__version__ = '0.1.1'

import re
import os
import sys
import warnings

import numpy as np
from netCDF4 import Dataset

class ApRESBurst(object):
    """
    Class for reading and writing ApRES bursts
    """

    DEFAULTS = {
        'autodetect_file_format_version': True,
        'forgive': True,
        'header_markers': {
            'start': '\r\n*** Burst Header ***',
            'end': '\r\n*** End Header ***'
        },
        'end_of_header_re': '\*\*\* End Header',
        'header_line_delim': '=',
        'header_line_eol': '\r\n',
        'data_type_key': 'Average',
        'data_types': ['<u2','<f4','<u4'],
        'data_dim_keys': ['NSubBursts', 'N_ADC_SAMPLES'],
        'data_dim_order': 'C'
    }

    ALTERNATIVES = [{
            'header_line_delim': '=',
            'data_dim_keys': ['NSubBursts', 'N_ADC_SAMPLES']
        }, {
            'header_line_delim': ':',
            'data_dim_keys': ['SubBursts in burst', 'Samples']
        }
    ]

    def __init__(self, fp=None):
        """
        Constructor

        * fp is not None: Then header_start will be set to the file's current
        position.  This allows reading through a file one burst block at a
        time.  This is the normal behaviour.
        * fp is None: Allows the burst object to be populated by other means,
        rather than reading from a file

        :param fp: The open file object of the input file
        :type fp: file object
        """

        self.header_start = 0
        self.data_start = -1
        self.header_lines = []
        self.header = {}
        self.data_shape = ()
        self.data_type = '<u2'
        self.data = None

        self.fp = fp

        if self.fp:
            self.header_start = self.fp.tell()

    def reset_init_defaults(self):
        """
        Reset the initial DEFAULTS parsing tokens from the ALTERNATIVES list
        """

        for key in self.ALTERNATIVES[0]:
            self.DEFAULTS[key] = self.ALTERNATIVES[0][key]

    def read_header_lines(self):
        """
        Read the raw header lines from the file

        :returns: The raw header lines
        :rtype: list
        """

        self.data_start = -1
        self.header_lines = []

        self.fp.seek(self.header_start, 0)
        line = self.fp.readline()
        self.header_lines.append(line.rstrip())

        while(line):
            # The data section follows the end of header marker
            if re.match(self.DEFAULTS['end_of_header_re'], line):
                self.data_start = self.fp.tell()
                break

            line = self.fp.readline()
            self.header_lines.append(line.rstrip())

        return self.header_lines

    def determine_file_format_version(self):
        """
        Determine the file format version from the read header lines

        The header lines are inspected to determine the file format version,
        and the DEFAULTS parsing tokens are setup accordingly
        """

        # Create a list of the first data dimension key from the alternative
        # file formats
        dim_keys = [item['data_dim_keys'][0] for item in self.ALTERNATIVES]

        # Search the header lines for one of the keys, to determine version
        for line in self.header_lines:
            for i, key in enumerate(dim_keys):
                if re.match(key, line):
                    for key in self.ALTERNATIVES[i]:
                        self.DEFAULTS[key] = self.ALTERNATIVES[i][key]
                    return

    def store_header(self):
        """
        Store the read header lines as parsed key/value pairs

        :returns: The parsed header
        :rtype: dict
        """

        self.header = {}

        for line in self.header_lines:
            item = line.split(self.DEFAULTS['header_line_delim'], maxsplit=1)

            if len(item) > 1 and item[0]:
                self.header[item[0].strip()] = item[1].strip()

        return self.header

    def define_data_shape(self):
        """
        Parse the data dimensions from the header to define the data shape

        If the Average header item is non-zero, then that indicates that the
        subbursts have been aggregated to a single value, so we rewrite that
        dimension as 1

        :returns: The data shape
        :rtype: tuple
        """

        self.data_shape = ()
        data_shape = []

        for key in self.DEFAULTS['data_dim_keys']:
            data_shape.append(int(self.header[key]))

        if int(self.header[self.DEFAULTS['data_type_key']]) > 0:
            data_shape[0] = 1

        self.data_shape = tuple(data_shape)

        return self.data_shape

    def define_data_type(self):
        """
        The Average item from the header is used to define the data type

        Average = 0: Record contains all subbursts as 16-bit ints
        Average = 1: Record contains the averaged subbursts as a 16-bit int
        Average = 2: Record contains the stacked subbursts as a 32-bit int

        :returns: The data type
        :rtype: str
        """

        try:
            self.data_type = self.DEFAULTS['data_types'][int(self.header[self.DEFAULTS['data_type_key']])]
        except IndexError as e:
            raise IndexError('Unsupported Averaging/Stacking configuration option found in header: {}'.format(int(self.header[self.DEFAULTS['data_type_key']])))

        return self.data_type

    def configure_from_header(self):
        """
        Configure the object from the raw header lines

        :returns: The parsed header
        :rtype: dict
        """

        if self.DEFAULTS['autodetect_file_format_version']:
            self.determine_file_format_version()

        self.store_header()
        self.define_data_shape()
        self.define_data_type()

        return self.header

    def read_header(self):
        """
        Read the header from the file

        * The raw header lines are available in the header_lines list
        * The parsed header is available in the header dict
        * The file offset to the start of the data is available in data_start
        * The data shape is available in the data_shape tuple
        * The data type is available in data_type

        :returns: The parsed header
        :rtype: dict
        """

        # The header lines are used to configure this object
        self.read_header_lines()
        self.configure_from_header()

        return self.header

    def reshape_data(self):
        """
        Reshape the data according to the data_shape tuple

        If the data read from the file don't conform to the expected
        data_shape as determined from the header, then when reshaping the data,
        an exception will occur:

        * If the reshaping fails because the data array is shorter than
        expected, then we will reraise the exception.
        * If forgive mode is True, then we will emit a warning and then
        truncate the data to conform to the data_shape, otherwise we will
        reraise the exception.

        :returns: The data
        :rtype: numpy.array
        """

        try:
            self.data = np.reshape(self.data, self.data_shape, order=self.DEFAULTS['data_dim_order'])
        except ValueError as e:
            expected_len = self.data_shape[0] * self.data_shape[1]

            if self.data.size < expected_len:
                warnings.warn("Data array read from file doesn't match data_shape as read from the file header: {}. It is shorter than expected. Cannot continue.")
                raise
            elif self.data.size > expected_len:
                warnings.warn("Data array read from file doesn't match data_shape as read from the file header. It is longer than expected.")

                if self.DEFAULTS['forgive']:
                    warnings.warn("{}. Will truncate data to fit to data_shape.".format(e))
                    self.data = np.reshape(self.data[:expected_len], self.data_shape, order=self.DEFAULTS['data_dim_order'])
                else:
                    raise

        return self.data

    def read_data(self):
        """
        Read the data from the file

        The data are available in the data array, shaped according to the
        data_shape tuple

        :returns: The data
        :rtype: numpy.array
        """

        if self.data_start == -1:
            self.read_header()

        count = self.data_shape[0] * self.data_shape[1]
        self.fp.seek(self.data_start, 0)
        self.data = np.fromfile(self.fp, dtype=np.dtype(self.data_type), count=count)
        self.reshape_data()

        return self.data

    def format_header_line(self, key, value):
        """
        Format a raw header line from the given key and value

        :param key: The header item key
        :type key: str
        :param value: The header item value
        :type value: str
        :returns: The formatted raw header line
        :rtype: str
        """

        return '{}{}{}'.format(key, self.DEFAULTS['header_line_delim'], value)

    def reconstruct_header_lines(self):
        """
        Reconstruct the raw header lines from the parsed header

        :returns: The reconstructed raw header lines
        :rtype: list
        """

        self.header_lines = [self.format_header_line(k,v) for k,v in self.header.items()]
        self.header_lines.insert(0, self.DEFAULTS['header_markers']['start'])
        self.header_lines.append(self.DEFAULTS['header_markers']['end'])

        return self.header_lines

    def write_header(self, fp, subbursts=None, samples=None):
        """
        Write the header to the given file object

        :param fp: The open file object of the output file
        :type fp: file object
        :param subbursts: A range specifying the subbursts to be written
        :type subbursts: range object
        :param samples: A range specifying the samples to be written
        :type samples: range object
        """

        if self.data_start == -1:
            self.read_header()

        eol = self.DEFAULTS['header_line_eol']

        for line in self.header_lines:
            # Ensure requested subbursts & samples are reflected in the header
            if subbursts and re.match(self.DEFAULTS['data_dim_keys'][0], line):
                line = self.format_header_line(self.DEFAULTS['data_dim_keys'][0], len(subbursts))
            if samples and re.match(self.DEFAULTS['data_dim_keys'][1], line):
                line = self.format_header_line(self.DEFAULTS['data_dim_keys'][1], len(samples))

            fp.write(line + eol)

    def write_data(self, fp, subbursts=None, samples=None):
        """
        Write the data to the given file object

        :param fp: The open file object of the output file
        :type fp: file object
        :param subbursts: A range specifying the subbursts to be written
        :type subbursts: range object
        :param samples: A range specifying the samples to be written
        :type samples: range object
        """

        if self.data_start == -1:
            self.read_data()

        if not subbursts:
            subbursts = range(self.data_shape[0])
        if not samples:
            samples = range(self.data_shape[1])

        fp.write(np.asarray(self.data[subbursts.start:subbursts.stop:subbursts.step, samples.start:samples.stop:samples.step], order=self.DEFAULTS['data_dim_order']))

class ApRESFile(object):
    """
    Context manager for reading and writing ApRES files
    """

    DEFAULTS = {
        'file_encoding': 'latin-1',
        'apres_suffix': '.dat',
        'netcdf_suffix': '.nc',
        'netcdf_add_history': True,
        'netcdf_group_name': 'burst{:d}',
        'netcdf_var_name': 'data',
        'netcdf_attrs': {
            'units': '1',
            'long_name': 'Sub Burst ADC Samples'
        }
    }

    def __init__(self, path=None):
        """
        Constructor

        :param path: Path to the file
        :type path: str
        """

        self.path = path
        self.fp = None
        self.file_size = -1
        self.bursts = []

    def __enter__(self):
        """
        Enter the runtime context for this object

        The file is opened

        :returns: This object
        :rtype: ApRESFile
        """

        return self.open(self.path)

    def __exit__(self, exc_type, exc_value, exc_traceback):
        """
        Exit the runtime context for this object

        The file is closed

        :returns: False
        :rtype: bool
        """

        self.close()

        return False         # This ensures any exception is re-raised

    def open(self, path=None, mode='r'):
        """
        Open the given file

        :param path: Path to the file
        :type path: str
        :param mode: Mode in which to open the file
        :type mode: str
        :returns: This object
        :rtype: ApRESFile
        """

        if path:
            self.path = path

        self.fp = open(self.path, mode, encoding=self.DEFAULTS['file_encoding'])
        self.file_size = os.fstat(self.fp.fileno()).st_size

        return self

    def close(self):
        """
        Close the file

        :returns: This object
        :rtype: ApRESFile
        """

        self.fp.close()
        self.fp = None

        return self

    def eof(self):
        """
        Report whether the end-of-file has been reached

        :returns: True if EOF has been reached, False otherwise
        :rtype: bool
        """

        return self.fp.tell() >= self.file_size

    def read_burst(self):
        """
        Read a burst block from the file

        :returns: The burst
        :rtype: ApRESBurst
        """

        burst = ApRESBurst(fp=self.fp)
        burst.read_data()

        return burst

    def read(self):
        """
        Reads all burst blocks from the file

        The bursts are available in the bursts list

        :returns: The bursts
        :rtype: list
        """

        while not self.eof():
            burst = self.read_burst()
            self.bursts.append(burst)

        return self.bursts

    def to_apres_dat(self, path, bursts=None, subbursts=None, samples=None):
        """
        Write the bursts to the given file path as an ApRES .dat file

        When rewriting the ApRES data to a new file, the bursts, subbursts,
        and the ADC samples for those selected subbursts, can be subsetted.
        The bursts, subbursts, and samples keyword arguments must specify a
        range object

        :param path: The path of the output file
        :type path: str
        :param bursts: A range specifying the bursts to be written
        :type bursts: range object
        :param subbursts: A range specifying the subbursts to be written
        :type subbursts: range object
        :param samples: A range specifying the samples to be written
        :type samples: range object
        """

        if not self.bursts:
            self.read()

        if not bursts:
            bursts = range(len(self.bursts))

        # We append each burst, so ensure file is empty if it already exists
        with open(path, 'w') as fout:
            pass

        # The ApRES .dat file format is a mixed mode file.  The header is
        # text, and the data section is binary
        for burst in self.bursts[bursts.start:bursts.stop:bursts.step]:
            with open(path, 'a') as fout:
                burst.write_header(fout, subbursts=subbursts, samples=samples)

            with open(path, 'ab') as fout:
                burst.write_data(fout, subbursts=subbursts, samples=samples)

    def burst_to_nc_object(self, burst, nco):
        """
        Map the given burst to the given netCDF object

        The netCDF object can either be a Dataset or a Group

        :param burst: The burst
        :type burst: ApRESBurst
        :param nco: A netCDF object (one of Dataset or Group)
        :type nco: netCDF4._netCDF4.Dataset or netCDF4._netCDF4.Group
        :returns: The netCDF object
        :rtype: netCDF4._netCDF4.Dataset or netCDF4._netCDF4.Group
        """

        # Write the burst header items as netCDF object attributes
        for key in burst.header:
            nco.setncattr(key, burst.header[key])

        for j, key in enumerate(burst.DEFAULTS['data_dim_keys']):
            nco.createDimension(key, burst.data_shape[j])

        data = nco.createVariable(self.DEFAULTS['netcdf_var_name'], burst.data_type, tuple(burst.DEFAULTS['data_dim_keys']))
        data.setncatts(self.DEFAULTS['netcdf_attrs'])
        data[:] = burst.data

        return nco

    def nc_object_to_burst(self, nco):
        """
        Map the given netCDF object to an ApRESBurst object

        The netCDF object can either be a Dataset or a Group

        :param nco: A netCDF object (one of Dataset or Group)
        :type nco: netCDF4._netCDF4.Dataset or netCDF4._netCDF4.Group
        :returns: The burst object
        :rtype: ApRESBurst
        """

        # We make a copy, otherwise data is invalid after file is closed
        data = nco.variables[self.DEFAULTS['netcdf_var_name']][:]
        attrs = vars(nco).copy()

        burst = ApRESBurst()
        burst.data_start = 0              # Stop data being read from file

        # Remove any attributes that weren't part of the original burst's header
        try:
            del attrs['history']
        except KeyError:
            pass

        # Reconstruct the original burst from the parsed header and data.  We
        # initially set the header_lines to be the parsed header which allows us
        # to determine the burst format version, and thus setup the DEFAULTS
        # parsing tokens accordingly.  This then allows us to correctly
        # reconstruct the raw header lines with the appropriate header line
        # delimiter for the burst format version
        burst.data = data
        burst.header = attrs
        burst.header_lines = burst.header
        burst.determine_file_format_version()
        burst.reconstruct_header_lines()
        burst.configure_from_header()

        return burst

    def to_netcdf(self, path=None, mode='w'):
        """
        Write the bursts to the given file path as a netCDF4 file

        The default netCDF file path is the same as the input file, but with
        a .nc suffix

        :param path: The path of the output file
        :type path: str
        :param mode: Mode in which to open the file
        :type mode: str
        """

        if not self.bursts:
            self.read()

        path = path or os.path.splitext(self.path)[0] + self.DEFAULTS['netcdf_suffix']

        ncfile = Dataset(path, mode)

        # Add the command line invocation to global history attribute
        if self.DEFAULTS['netcdf_add_history']:
            ncfile.history = ' '.join(sys.argv)

        # Write each burst as a netCDF4/HDF5 group
        for i, burst in enumerate(self.bursts):
            group = ncfile.createGroup(self.DEFAULTS['netcdf_group_name'].format(i))
            self.burst_to_nc_object(burst, group)

        ncfile.close()


#!/usr/bin/python3
from nis import cat
from ssl import HAS_ECDH, _create_unverified_context
import sys
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import cv2
import torch
import torch.backends.cudnn as cudnn
from absl import app
from absl.flags import FLAGS

import rospy
import message_filters
from sensor_msgs.msg import Image
from rospy.numpy_msg import numpy_msg
from geometry_msgs.msg import Vector3
from std_msgs.msg import Header, Int8

# ADD YOLOv5 LIBS #
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH

from utils.augmentations import letterbox
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.plots import Annotator, colors

coco_id_mapping = {
    1: 'person'
}  # pyformat: disable

coco_ths_mapping = {
     1: 0.7
}  # pyformat: disable

#################################################################################
# Implementation of a CV_Bridge function for compatibility with ROS Melodic
## this is not a complete implementation
## some code lines dealing with exceptions have been removed for convenience

numpy_type_to_cvtype = {'uint8': '8U', 'int8': '8S', 'uint16': '16U',
                        'int16': '16S', 'int32': '32S', 'float32': '32F', 'float64': '64F'}

def cv2_to_imgmsg(cvim, encoding="passthrough", header=None):
    if not isinstance(cvim, (np.ndarray, np.generic)):
        raise TypeError('Your input type is not a numpy array')

    img_msg = Image()
    img_msg.height = cvim.shape[0]
    img_msg.width = cvim.shape[1]

    if header is not None:
        img_msg.header = header

    if cvim.ndim < 3:
        n_channels = 1
    else:
        n_channels = cvim.shape[2]

    cv_type = '%sC%d' % (numpy_type_to_cvtype[cvim.dtype.name], n_channels)

    if encoding == "passthrough":
        img_msg.encoding = cv_type
    else:
        img_msg.encoding = encoding

    if cvim.dtype.byteorder == '>':
        img_msg.is_bigendian = True

    #img_msg.data = cvim.tostring() # Deprecated since numpy version 1.19.0.
    img_msg.data = cvim.tobytes()
    img_msg.step = len(img_msg.data) // img_msg.height

    return img_msg

#################################################################################

class Detector(object):
    def __init__(self):

        # define class variables and set initial values
        ## variables for runnig yolov5 model
        self.image_size = (int(720), int(1280))
        self.conf_thres = 0.15
        self.iou_thres = 0.45
        self.agnostic_nms= False
        self.max_det = 100
        self.max_dist = 3
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

        ## variables for topic management
        self.topic_name = {}
        self.topic_name["rgb"] = "/zed2/zed_node/left/image_rect_color"
        self.topic_name["depth"] = "/zed2/zed_node/depth/depth_registered"
        self.raw_img = {}
        self.raw_img["rgb"] = np.zeros( self.image_size + (3,) )
        self.raw_img["depth"] = np.zeros( self.image_size )
        self.is_received = {}
        self.is_received["rgb"] = False
        self.is_received["depth"] = False
        self.rcv_time = {}
        self.rcv_time["rgb"] = None
        self.rcv_time["depth"] = None
        self.proc_time = rospy.get_time()

        ## variables for controlling module status
        self.model = None
        self.is_initialized = False
        ## variables for camera to world transformation
        # self.color_intrinsic = np.array([[337.2084410968044, 0.0,               320.5],
        #                                  [0.0,               337.2084410968044, 240.5],
        #                                  [0.0,               0.0,               1.0  ]])
        # self.inverse_intrinsic = np.linalg.inv(self.color_intrinsic)
        # self.cam_to_body = np.array([[ 0,  0,  1],
        #                              [-1,  0,  0],
        #                              [ 0, -1,  0]])

        # SUBSCRIBERS
        self.sub_rgb = message_filters.Subscriber(self.topic_name["rgb"], numpy_msg(Image))
        self.sub_depth = message_filters.Subscriber(self.topic_name["depth"], numpy_msg(Image))
        self.sub_ts = message_filters.TimeSynchronizer([self.sub_rgb, self.sub_depth], queue_size=10)
        self.sub_ts.registerCallback(self.sync_callback)
        # PUBLISHERS
        self.pub_img = rospy.Publisher("/detection/image_rect", Image, queue_size=1)
        self.pub_crowd = rospy.Publisher("/detection/num_people", Int8, queue_size=1)

        # initialize the yolov5 model
        self._init_model()

    @torch.no_grad()
    def _init_model(self):
        # initialize model and load pretrained weight
        print('Initializing model...')
        self.model = attempt_load('yolov5m.pt', map_location=self.device)  # load FP32 model
        self.stride = int(self.model.stride.max())  # model stride
        self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names  # get class names
        self.model.eval()
        print('model loaded')
        if self.is_initialized == False:
            self.is_initialized = True

    def _preprocess_data(self, input_img):
        ## padding and resizing
        tmp_img = letterbox(input_img, self.image_size, stride=self.stride, auto=True)[0]
        ## switching axes
        tmp_img = tmp_img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        tmp_img = np.ascontiguousarray(tmp_img)
        ## noramlization
        tmp_img = tmp_img / 255.0
        ## convert data format from numpy to tensor
        input_tensor = torch.from_numpy(tmp_img).unsqueeze(dim=0).to(self.device).float()
        return input_tensor

    def _make_prediction(self, input_tensor):
        # make prediction for the current input
        pred, _ = self.model(input_tensor)
        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, \
                None, self.agnostic_nms, max_det=self.max_det)
        ## get detection result of the first sample
        ## since only one sample is given as input
        result = pred[0]

        return result

    def _measure_distance(self, box, depth_img):
        H, W = self.image_size
        # find 3d position of the center point of detected box
        box = [int(val) for val in box]
        x_0, y_0, x_1, y_1 = box

        ## slice the center box
        img_slice = depth_img[y_0:y_1, x_0:x_1]
        ## if over 50% of values in the center box is NaN, skip the current box
        nan_idx = np.isnan(img_slice)
        if np.mean(nan_idx) < 0.5:
            distance = np.mean(img_slice[np.invert(nan_idx)])
        else:
            distance = np.nan

        return distance

    def _process_detections(self, annotator, detections, tensor_size, depth_img):
        # initilaize a variable to store number of people in the image
        num_people = 0
        ## scale detection box cooridnates to fit the input image size
        detections[:, :4] = scale_coords(tensor_size, detections[:, :4], self.image_size).round()
        ## draw detection boxes
        for *xyxy, conf, obj in reversed(detections):
            # get the class index
            c = int(obj)
            c_coco = c+1

            # check if the detected class is in the coco_id_mapping
            if not (c_coco in coco_id_mapping.keys()):
                continue
            # get the name for the detected class
            name = coco_id_mapping[c_coco]
            # get the confidence threshold
            conf_ths = coco_ths_mapping[c_coco]
            # check if the confidence score is over the threshold
            if conf.data.cpu() < conf_ths:
                continue

            # measure distance values in the box
            distance = self._measure_distance(xyxy, depth_img)
            if distance > self.max_dist:
                continue

            # draw a detection box
            color = colors(c, bgr=True)
            annotator.box_label(xyxy, name, color=color)
            num_people += 1

        return num_people

    def sync_callback(self, rgb, depth):
        try:
            # store the rgb message
            self.raw_img["rgb"] = np.frombuffer(rgb.data, dtype=np.uint8).reshape(self.image_size+(3,))
            self.is_received["rgb"] = True
            self.rcv_time["rgb"] = rgb.header.stamp
            # store the depth message
            self.raw_img["depth"] = np.frombuffer(depth.data, dtype=np.float32).reshape(self.image_size)
            self.is_received["depth"] = True
            self.rcv_time["depth"] = depth.header.stamp

        except Exception as e:
            print(e)

    def run_model(self):
        # if no image has been received, skip the loop
        if (self.is_received["rgb"] and self.is_received["depth"]) is False:
            return

        # measure the start time of current loop
        t0 = rospy.get_time()

        # if the last process time has not been outdated than a second, skip the loop
        # if (t0 - self.proc_time) < 1.0:
        #     return

        # copy the recent images
        tmp_img = np.copy(self.raw_img["rgb"][:, :, :3])
        tmp_depth = np.copy(self.raw_img["depth"])
        tmp_stamp = self.rcv_time["rgb"]
        ## store the time stamp into a header
        header = Header()
        header.stamp = tmp_stamp

        # prepare tensor input data
        tmp_img_ = self._preprocess_data(tmp_img)

        # make prediction
        result = self._make_prediction(tmp_img_)

        # process the detection result
        ## define a annotator object for drawing detection boxes
        annotator = Annotator(tmp_img, line_width=1, pil=False)
        num_people = self._process_detections(annotator, result, tmp_img_.shape[2:], tmp_depth)
        ## get the annotation result
        output_img = annotator.result()

        # measure the time when the detection is finished
        t1 = rospy.get_time()

        # print("inference time: %4.3f"%(t1-t0))

        # publish the detection result
        self.pub_img.publish(cv2_to_imgmsg(output_img.astype(np.uint8), "rgb8", header))
        self.pub_crowd.publish(Int8(data=num_people))

        # update the last process time
        self.proc_time = t1

def main(_argv):
    rospy.init_node('simple_crowd_detector')

    detector = Detector()
    rate = rospy.Rate(10)
    while not rospy.is_shutdown():
        if detector.is_initialized is True:
            detector.run_model()
            rate.sleep()

if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit as e:
        print(e)


from task_two.day_in_month import days_in_month

# a help function to divide the date currectly
# the function accepts either dd.mm.year or dd-mm-year format
# the function returns either a string invalid input or the date as a tuple (dd,mm,year)
def split_helper(date, seperator):
    try:
        date_list = date.split(seperator)
        if len(date_list) == 3:
            dd = date_list[0]
            mm = date_list[1]
            yy = date_list[2]
            if dd.isdigit() and mm.isdigit() and yy.isdigit():
                dd, mm, yy = int(dd), int(mm), int(yy)
                # ^ made below code cleaner ^
                # check month and day validation
                if mm < 13 and dd <= days_in_month(mm, yy):
                    date_tuple = (True, dd, mm, yy)
                    return date_tuple

        date_tuple = (False, "Invalid date input")
        return date_tuple
    except AttributeError as err:
        return err


# Cleaning up sections of the camp - each section has a unique ID and each elve has a range
# Noticed many of the section assignments overlap
# Sections are assigned in comma seprated lines as ranges
# We are looking for fully contained overlap (6-6, 4-6)

def main():
    # pure overlap
    naughty_pairs = []
    # partial overlap
    extra_naughty_pairs = []
    with open("cleaning_pairs.txt") as f:
        while True:
            line = f.readline()
            if not line:
                break
            clean_line = line.strip()
            elf_1_assign, elf_2_assign = clean_line.split(",")
            e1 = [int(i) for i in elf_1_assign.split("-")]
            e2 = [int(i) for i in elf_2_assign.split("-")]

            if (e1[0] >= e2[0] and e1[1] <= e2[1]) or (e1[0] <= e2[0] and e1[1] >= e2[1]):
                naughty_pairs.append([e1, e2])

            # part 2: which have any amount of partial overlap

            if (e1[0] <= e2[0] and e1[1] >= e2[0]) or (e2[0] <= e1[0] and e2[1] >= e1[0]):
                extra_naughty_pairs.append([e1, e2])

    print("Total overlap:")
    print(len(naughty_pairs))

    print(extra_naughty_pairs)
    print(len(extra_naughty_pairs))

if __name__ == "__main__":
    main()


import re

# txt = "The rain in Spain"

# #Find all lower case characters alphabetically between "a" and "m":
'''
Write a Python program that matches a word containing 'z',
not the start or end of the word.
'''
# x = re.findall("[a-m]", txt)
# print(sorted(x))

text = "Zebra"

pattern = r'\b\w*z\w*\b'

matches = re.findall(pattern, text, re.IGNORECASE )
for i in matches:
    print(i)
else:
    print(None)

    '''
    dates
    '''


# This code is supporting material for the book
# Building Machine Learning Systems with Python
# by Willi Richert and Luis Pedro Coelho
# published by PACKT Publishing
#
# It is made available under the MIT License

import os
import sys
import collections
import csv
import json

from matplotlib import pylab
import numpy as np

DATA_DIR = "data"
CHART_DIR = "charts"

if not os.path.exists(DATA_DIR):
    raise RuntimeError("Expecting directory 'data' in current path")

if not os.path.exists(CHART_DIR):
    os.mkdir(CHART_DIR)

def tweak_labels(Y, pos_sent_list):
    pos = Y == pos_sent_list[0]
    for sent_label in pos_sent_list[1:]:
        pos |= Y == sent_label

    Y = np.zeros(Y.shape[0])
    Y[pos] = 1
    Y = Y.astype(int)

    return Y

def load_sanders_data(dirname=".", line_count=-1):
    count = 0

    topics = []
    labels = []
    tweets = []

    with open(os.path.join(DATA_DIR, dirname, "corpus.csv"), "r") as csvfile:
        metareader = csv.reader(csvfile, delimiter=',', quotechar='"')
        for line in metareader:
            count += 1
            if line_count > 0 and count > line_count:
                break

            topic, label, tweet_id = line

            tweet_fn = os.path.join(
                DATA_DIR, dirname, 'rawdata', '%s.json' % tweet_id)
            try:
                tweet = json.load(open(tweet_fn, "r"))
            except IOError:
                print(("Tweet '%s' not found. Skip." % tweet_fn))
                continue

            if 'text' in tweet and tweet['user']['lang'] == "en":
                topics.append(topic)
                labels.append(label)
                tweets.append(tweet['text'])

    tweets = np.asarray(tweets)
    labels = np.asarray(labels)

    return tweets, labels

def plot_pr(auc_score, name, phase, precision, recall, label=None):
    pylab.clf()
    pylab.figure(num=None, figsize=(5, 4))
    pylab.grid(True)
    pylab.fill_between(recall, precision, alpha=0.5)
    pylab.plot(recall, precision, lw=1)
    pylab.xlim([0.0, 1.0])
    pylab.ylim([0.0, 1.0])
    pylab.xlabel('Recall')
    pylab.ylabel('Precision')
    pylab.title('P/R curve (AUC=%0.2f) / %s' % (auc_score, label))
    filename = name.replace(" ", "_")
    pylab.savefig(os.path.join(CHART_DIR, "pr_%s_%s.png" %
                               (filename, phase)), bbox_inches="tight")

def show_most_informative_features(vectorizer, clf, n=20):
    c_f = sorted(zip(clf.coef_[0], vectorizer.get_feature_names()))
    top = list(zip(c_f[:n], c_f[:-(n + 1):-1]))
    for (c1, f1), (c2, f2) in top:
        print("\t%.4f\t%-15s\t\t%.4f\t%-15s" % (c1, f1, c2, f2))

def plot_log():
    pylab.clf()
    pylab.figure(num=None, figsize=(6, 5))

    x = np.arange(0.001, 1, 0.001)
    y = np.log(x)

    pylab.title('Relationship between probabilities and their logarithm')
    pylab.plot(x, y)
    pylab.grid(True)
    pylab.xlabel('P')
    pylab.ylabel('log(P)')
    filename = 'log_probs.png'
    pylab.savefig(os.path.join(CHART_DIR, filename), bbox_inches="tight")

def plot_feat_importance(feature_names, clf, name):
    pylab.clf()
    coef_ = clf.coef_
    important = np.argsort(np.absolute(coef_.ravel()))
    f_imp = feature_names[important]
    coef = coef_.ravel()[important]
    inds = np.argsort(coef)
    f_imp = f_imp[inds]
    coef = coef[inds]
    xpos = np.array(list(range(len(coef))))
    pylab.bar(xpos, coef, width=1)

    pylab.title('Feature importance for %s' % (name))
    ax = pylab.gca()
    ax.set_xticks(np.arange(len(coef)))
    labels = ax.set_xticklabels(f_imp)
    for label in labels:
        label.set_rotation(90)
    filename = name.replace(" ", "_")
    pylab.savefig(os.path.join(
        CHART_DIR, "feat_imp_%s.png" % filename), bbox_inches="tight")

def plot_feat_hist(data_name_list, filename=None):
    pylab.clf()
    num_rows = 1 + (len(data_name_list) - 1) / 2
    num_cols = 1 if len(data_name_list) == 1 else 2
    pylab.figure(figsize=(5 * num_cols, 4 * num_rows))

    for i in range(num_rows):
        for j in range(num_cols):
            pylab.subplot(num_rows, num_cols, 1 + i * num_cols + j)
            x, name = data_name_list[i * num_cols + j]
            pylab.title(name)
            pylab.xlabel('Value')
            pylab.ylabel('Density')
            # the histogram of the data
            max_val = np.max(x)
            if max_val <= 1.0:
                bins = 50
            elif max_val > 50:
                bins = 50
            else:
                bins = max_val
            n, bins, patches = pylab.hist(
                x, bins=bins, normed=1, facecolor='green', alpha=0.75)

            pylab.grid(True)

    if not filename:
        filename = "feat_hist_%s.png" % name

    pylab.savefig(os.path.join(CHART_DIR, filename), bbox_inches="tight")

def plot_bias_variance(data_sizes, train_errors, test_errors, name):
    pylab.clf()
    pylab.ylim([0.0, 1.0])
    pylab.xlabel('Data set size')
    pylab.ylabel('Error')
    pylab.title("Bias-Variance for '%s'" % name)
    pylab.plot(
        data_sizes, train_errors, "-", data_sizes, test_errors, "--", lw=1)
    pylab.legend(["train error", "test error"], loc="upper right")
    pylab.grid()
    pylab.savefig(os.path.join(CHART_DIR, "bv_" + name + ".png"))

def load_sent_word_net():

    sent_scores = collections.defaultdict(list)
    sentiwordnet_path = os.path.join(DATA_DIR, "SentiWordNet_3.0.0_20130122.txt")

    if not os.path.exists(sentiwordnet_path):
        print("Please download SentiWordNet_3.0.0 from http://sentiwordnet.isti.cnr.it/download.php, extract it and put it into the data directory")
        sys.exit(1)

    with open(sentiwordnet_path, 'r') as csvfile:
        reader = csv.reader(csvfile, delimiter='\t', quotechar='"')
        for line in reader:
            if line[0].startswith("#"):
                continue
            if len(line) == 1:
                continue

            POS, ID, PosScore, NegScore, SynsetTerms, Gloss = line
            if len(POS) == 0 or len(ID) == 0:
                continue
            # print POS,PosScore,NegScore,SynsetTerms
            for term in SynsetTerms.split(" "):
                # drop #number at the end of every term
                term = term.split("#")[0]
                term = term.replace("-", " ").replace("_", " ")
                key = "%s/%s" % (POS, term.split("#")[0])
                sent_scores[key].append((float(PosScore), float(NegScore)))
    for key, value in sent_scores.items():
        sent_scores[key] = np.mean(value, axis=0)

    return sent_scores

def log_false_positives(clf, X, y, name):
    with open("FP_" + name.replace(" ", "_") + ".tsv", "w") as f:
        false_positive = clf.predict(X) != y
        for tweet, false_class in zip(X[false_positive], y[false_positive]):
            f.write("%s\t%s\n" %
                    (false_class, tweet.encode("ascii", "ignore")))

if __name__ == '__main__':
    plot_log()


#!/usr/bin/env python3.9
import time
import youtube_dl
from youtube_dl.utils import DownloadError, ExtractorError
from pathlib import Path
import requests
import pdftotext
import get_wayback_machine
from bs4 import BeautifulSoup
from moviepy.editor import VideoFileClip
from newspaper import Article
from ebooklib import epub

from user_settings import average_word_length, headers, wpm
from src.backend.log import log_

def extract_epub(url):
    "extracts book duration and title from epub file"
    try:
        book = epub.read_epub(url)
        items = [x for x in book.get_items()]
        text_content = ""
        for item in items:
            if hasattr(item, "get_body_content"):
                 html = item.get_body_content()
                 soup = BeautifulSoup(html, "html.parser")
                 text_content += soup.text.replace("\n", " ")

        total_words = len(text_content) / average_word_length
        estimatedReadingTime = str(round(total_words / wpm, 1))

        res = {"type": "epub",
               "length": estimatedReadingTime,
               "title": book.title.replace("\n", ""),
               "url": url}
    except Exception as e:
        log_(f"Error: {e}", False)
        res = {"type": "epub (exception)",
               "url": url}
    return res

def extract_youtube(url):
    "extracts video duration in minutes from youtube link, title etc"
    res = {}
    with youtube_dl.YoutubeDL({"quiet": True}) as ydl:
        try:
            video = ydl.extract_info(url, download=False)
        except (KeyError, DownloadError, ExtractorError) as e:
            log_(f"ERROR: Video link skipped because : error during information \
extraction from {url} : {e}", False)
            res.update({"type": "video not found", "url": url})
            return res

        title = video['title'].strip()
        res = {"type": "video",
               "title": title.replace("\n", ""),
               "url": url}

        if "duration" in video:
            length = str(round(video['duration'] / 60, 1))
            res.update({"length": length})
            res.update({"channel": video["uploader"]})
        else:
            log_("Youtube playlist detected.:", False)
            res["type"] = "video (youtube playlist)"
            if "playlist" not in res["title"].lower():
                res["title"] += " - Playlist"
            length = 0
            channels = []
            for ent in video["entries"]:
                length += ent["duration"]
                channels.append(ent["uploader"])
            res["length"] = str(round(length / 60, 1))
            if len(list(set(channels))) == 1:
                res["channel"] = channels[0]
            else:
                if "channel" in res:
                    del res["channel"]
    return res

def extract_local_video(link):
    """
    extracts video duration in minutes from local files
    https://stackoverflow.com/questions/3844430/how-to-get-the-duration-of-a-video-in-python
    """
    vid = Path(link)
    if not vid.exists():
        link = link.replace("\\", "")
        vid = Path(link)
    if not vid.exists():
        log_(f"ERROR : Thought this was a local video but file was not found! \
: {link}", False)
        return {"type": "local video not found",
                "url": link}
    clip = VideoFileClip(link)
    duration = round(clip.duration / 60, 1)
    title = clip.filename.strip()
    dic = {"type": "local video",
           "title": title.replace("\n", ""),
           "length": duration,
           "url": link}
    return dic

def extract_pdf_url(url):
    "extracts reading time from an online pdf"
    try:
        downloaded = requests.get(url, headers=headers, timeout=5)
    except requests.exceptions.ConnectTimeout as e:
        log_(f"Connection timed out, skipping url: {url}", False)
        temp_dic = {}
        temp_dic["type"] = "connection timed out"
        temp_dic["url"] = url
        return temp_dic
    open("./.temporary.pdf", "wb").write(downloaded.content)
    temp_dic = extract_pdf_local("./.temporary.pdf")
    temp_dic["type"] = "online pdf"
    temp_dic["url"] = url
    del temp_dic["title"]  # can't yet scrap pdf title reliably
    Path("./.temporary.pdf").unlink()
    return temp_dic

def extract_pdf_local(path):
    "extracts reading time from a local pdf file"
    if not Path(path).exists():
        path = path.replace(" ", "\\ ")
    try:
        with open(path, "rb") as f:
            text_content = pdftotext.PDF(f)
            text_content = " ".join(text_content).replace("\n", " ")
    except FileNotFoundError:
        log_(f"ERROR: Cannot find {path}, I thought it was a PDF", False)
        return {"type": "pdf not found",
                "url": path}

    total_words = len(text_content) / average_word_length
    estimatedReadingTime = str(round(total_words / wpm, 1))

    title = path.split(sep="/")[-1].strip()
    res = {"type": "local pdf",
           "length": estimatedReadingTime,
           "title": title.replace("\n", ""),
           "url": path}
    return res

def extract_txt(path):
    "extracts reading time from a text file"
    try:
        txt_file = Path(path)
        with txt_file.open() as f:
            lines = f.readlines()
        text_content = ' '.join(lines).replace("\n", " ")

        total_words = len(text_content) / average_word_length
        estimatedReadingTime = str(round(total_words / wpm, 1))

        title = path.split(sep="/")[-1].strip()
        res = {"type": "text",
               "length": estimatedReadingTime,
               "url": path,
               "title": title.replace("\n", "")}
        return res

    except ValueError as e:
        log_(f"ERROR: Cannot find {path} : {e}", False)
        res = {"type": "txt file not found",
               "url": path}
        return res

def extract_webpage(url, simple_method=True):
    """
    extracts reading time from a webpage, output is a dictionnary containing
    estimation of the reading time ; title of the page ; if the wayback
    machine was used
    """
    wayback_used = 0
    if simple_method:
        article = Article(url)
        try:
            article.download()
            article.parse()
            title = article.title
            if title == "WWW Error Blocked Diagnostic":
                raise Exception("Exception, found title: 'WWW Error Blocked Diagnostic'")
            text_content = " ".join(article.text.replace("\n", " ").split())
        except Exception as e:
            log_(f"Error: {e}")
            simple_method = False
    if simple_method is False:
        log_("Using fallack extractor")
        try:
            req = requests.get(url, headers=headers, timeout=5)
        except requests.exceptions.ConnectTimeout as e:  # if timed out
            log_("Connection timed out, retrying once in 5 seconds...", False)
            time.sleep(5)
            try:
                req = requests.get(url, headers=headers, timeout=10)
            except requests.exceptions.ConnectTimeout as e:
                log_(f"Connection timed out again, skipping url: {url}", False)
                res = {"type": "connection timed out",
                       "url": url}
                return res
        except requests.exceptions.ConnectionError: # if url is dead use wayback machine
            log_(f"Using the wayback machine for {url}", False)
            wayback_used = 1
            wb = get_wayback_machine.get(url)
            try:
                url = wb.links['last memento']['url']
                req = requests.get(url, headers=headers, timeout=5)
            except (requests.exceptions.ConnectionError, AttributeError) as e:
                log_(f"ERROR: url could not be found even using wayback machine : \
{url} : {e}", False)
                res = {"type": "web page not found",
                       "url": url,
                       "used_wayback_machine": "wayback url not found"}
                return res
            except requests.exceptions.ConnectTimeout as e:
                log_(f"Connection timed out on wayback machine, skipping url: {url}", False)
                res = {"type": "connection timed out",
                       "url": url}
                return res

        # extracting divs and p elements and keeping the the smallest text
        html_page = req.content
        soup = BeautifulSoup(html_page, 'html.parser')
        parsed_text_trial = []
        parsed_text_trial.append(' '.join([x.text.replace("\n", " ")
                                           for x in soup.find_all("div")]))
        parsed_text_trial.append(soup.get_text().replace("\n", " "))
        parsed_text_trial.append(" ".join([x.text.replace("\n", " ")
                                           for x in soup.find_all('p')]))
        parsed_text_trial = [" ".join(x.split()) for x in parsed_text_trial]
        parsed_text_trial.sort(key=lambda x: len(x))
        text_content = parsed_text_trial[0]

        titles = soup.find_all('title')
        if len(titles) != 0:
            title = soup.find_all('title')[0].get_text()
        else:
            title = "No title found"

    total_words = len(text_content) / average_word_length
    estimatedReadingTime = str(round(total_words / wpm, 1))
    title = title.replace("\n", "").replace("\\n", "").strip()
    log_(f"Found title: {title}")
    res = {"title": title,
           "type": "webpage",
           "length": estimatedReadingTime,
           "used_wayback_machine": wayback_used,
           "url": url}
    return res


import numpy as np
import os
from pathlib import Path
import scipy
import torch
from torch.nn import functional as F
from typing import *

from src.enums import *

class Transformer:
    def __init__(self, purpose: str, transforms: list):
        self.purpose = purpose

        self.transforms = transforms

        if self.purpose == "test":
            self.deterministic = True
        else:
            self.deterministic = False

        self.rng = np.random.RandomState(seed=1)

    def __call__(self, data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        transformed_data = data

        for transform_config in self.transforms:
            transform_type = TransformEnum(transform_config["type"])

            transform_fct = getattr(self, transform_type.value)
            transformed_data = transform_fct(transform_config, transformed_data)

        return transformed_data

    def white_noise(self, transform_config: dict,
                    data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        stdev = transform_config["stdev"]

        noise = None
        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                if noise is None:
                    noise_value = self.rng.normal(loc=0, scale=stdev, size=tuple(value.size()))
                    noise = value.new_tensor(noise_value, dtype=value.dtype)

                transformed_value = value + noise

                data[channel] = transformed_value

        return data

    def range_adjusted_white_noise(self, transform_config: dict,
                                   data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        res_grid = data[ChannelEnum.RES_GRID]
        robot_position = data[ChannelEnum.REL_POSITION]

        sample_grid = data[ChannelEnum.OCC_DEM]

        # distance of every pixel from the robot
        lin_x = np.arange(start=-sample_grid.shape[0] / 2, stop=sample_grid.shape[0] / 2, step=1) * res_grid[0].numpy()
        lin_y = np.arange(start=-sample_grid.shape[1] / 2, stop=sample_grid.shape[1] / 2, step=1) * res_grid[1].numpy()
        off_y, off_x = np.meshgrid(lin_x, lin_y)

        dist_x = off_x - robot_position[0].numpy()
        dist_y = off_y - robot_position[1].numpy()
        dist_p2_norm = np.sqrt(np.square(dist_x) + np.square(dist_y))

        stdev = transform_config["stdev"]
        range = transform_config["range"]

        scale = stdev * np.square(1 / range * dist_p2_norm)

        noise = None
        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                if noise is None:
                    noise_value = rng.normal(loc=0, scale=scale, size=tuple(value.size()))
                    noise = value.new_tensor(noise_value, dtype=value.dtype)

                transformed_value = value + noise

                data[channel] = transformed_value

        return data

    def gaussian_filtered_white_noise(self, transform_config: dict,
                                      data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        probability = transform_config["probability"]
        horizontal_stdev = transform_config["horizontal_stdev"]
        vertical_stdev = transform_config["vertical_stdev"]

        filtered_noise = None
        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                if filtered_noise is None:
                    peaks = rng.choice([0, 1], p=[1 - probability, probability], size=tuple(value.size()))
                    white_noise = rng.normal(loc=0, scale=vertical_stdev, size=tuple(value.size()))
                    noise = peaks * white_noise
                    filtered_noise = scipy.ndimage.gaussian_filter(input=noise, sigma=horizontal_stdev)

                    filtered_noise = value.new_tensor(filtered_noise, dtype=value.dtype)

                transformed_value = value + filtered_noise

                data[channel] = transformed_value

        return data

    def random_vertical_scale(self, transform_config: dict,
                              data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        min, max = transform_config["min"], transform_config["max"]

        scale = rng.uniform(low=min, high=max)

        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                if channel is ChannelEnum.REL_POSITION:
                    transformed_value = value.clone()
                    transformed_value[2] = scale * transformed_value[2]
                else:
                    transformed_value = scale * value

                data[channel] = transformed_value

        return data

    def random_vertical_offset(self, transform_config: dict,
                               data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        min, max = transform_config["min"], transform_config["max"]

        offset = rng.uniform(low=min, high=max)

        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                if channel is ChannelEnum.REL_POSITION:
                    transformed_value = value.clone()
                    transformed_value[2] = offset + transformed_value[2]
                else:
                    transformed_value = offset + value

                data[channel] = transformed_value

        return data

    def random_occlusion(self, transform_config: dict,
                         data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        probability = transform_config["probability"]

        occlusion = None
        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                if occlusion is None:
                    occlusion = rng.choice([0, 1], p=[1 - probability, probability], size=tuple(value.size()))
                    occlusion = torch.tensor(occlusion)

                if channel == ChannelEnum.OCC_MASK:
                    # assert value.dtype == torch.bool
                    #
                    # transformed_value = value.to(dtype=torch.int) + occlusion
                    #
                    # transformed_value[transformed_value > 1] = 1
                    #
                    # data[channel] = transformed_value.to(dtype=torch.bool)

                    value[occlusion == 1] = True
                elif channel in [ChannelEnum.OCC_DEM, ChannelEnum.OCC_DATA_UM]:
                    value[occlusion == 1] = np.nan
                else:
                    raise NotImplementedError

                data[channel] = value

        return data

    def random_occlusion_dilation(self, transform_config: dict,
                                  data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        mask = data[ChannelEnum.OCC_MASK]

        if mask.dtype == torch.bool:
            mask = mask.to(dtype=torch.float)

        if mask.dim() == 2:
            # we need a channel composition with minibatches for 2d convolutions
            mask = mask.unsqueeze(dim=0).unsqueeze(dim=0)
        else:
            raise NotImplementedError

        kernel_size = int(transform_config["kernel_size"])
        # probability to set kernel cell to one
        probability = transform_config["probability"]
        kernel_choice = rng.choice([0, 1], p=[1 - probability, probability], size=(kernel_size, kernel_size))
        # we always set the center of the kernel to 1
        kernel_choice[int(kernel_size // 2), int(kernel_size // 2)] = 1
        kernel = torch.tensor(kernel_choice, dtype=torch.float).unsqueeze(dim=0).unsqueeze(dim=0)

        with torch.no_grad():
            dilated_mask = F.conv2d(mask, weight=kernel, padding=1) > 0

        dilated_occlusion = dilated_mask.squeeze(dim=0).squeeze(dim=0)

        data[ChannelEnum.OCC_MASK] = dilated_occlusion

        if ChannelEnum.OCC_DEM in data:
            data[ChannelEnum.OCC_DEM][dilated_occlusion == 1] = np.nan
        if ChannelEnum.OCC_DATA_UM in data:
            data[ChannelEnum.OCC_DATA_UM][dilated_occlusion == 1] = np.nan

        return data

    def random_occlusion_raycasting(self, transform_config: dict,
                               data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        import grid_map_raycasting

        occ_mask = data[ChannelEnum.OCC_MASK]
        gt_dem = data[ChannelEnum.GT_DEM]
        res_grid = data[ChannelEnum.RES_GRID]

        vantage_point_u = 0
        vantage_point_v = 0
        vantage_point_elevation = np.NaN
        # we only accept coordinates for which we have elevation information
        while np.isnan(vantage_point_elevation):
            u_choices = np.arange(start=0, stop=gt_dem.size(0))
            v_choices = np.arange(start=0, stop=gt_dem.size(1))

            vantage_point_u = rng.choice(u_choices)
            vantage_point_v = rng.choice(v_choices)

            vantage_point_elevation = gt_dem[vantage_point_u, vantage_point_v]

        if transform_config.get("occluded_area_constraint") is not None:
            constraint_config = transform_config["occluded_area_constraint"]
            max_iterations = constraint_config["max_iterations"]
            min_fraction = constraint_config["min"]
            max_fraction = constraint_config["max"]
            min_step = constraint_config.get("min_step", 0.05)
        else:
            max_iterations = 1
            min_fraction = 0
            max_fraction = 1
            min_step = 0

        low = transform_config["height_viewpoint"]["min"]
        high = transform_config["height_viewpoint"]["max"]
        found_suitable_occlusion = False
        i = 0
        while found_suitable_occlusion is False and i < max_iterations:
            i += 1

            height_viewpoint = rng.uniform(low=low, high=high)

            vantage_point_x = (-gt_dem.size(0) / 2 + vantage_point_u) * res_grid[0]
            vantage_point_y = (-gt_dem.size(1) / 2 + vantage_point_v) * res_grid[1]
            vantage_point_z = vantage_point_elevation + height_viewpoint

            vantage_point = np.array([vantage_point_x, vantage_point_y, vantage_point_z], dtype=np.double)

            np_gt_dem = gt_dem.detach().numpy().astype(np.double)

            np_raycasted_occ_mask = grid_map_raycasting.rayCastGridMap(vantage_point, np_gt_dem, res_grid)

            occ_mask_shape = np_raycasted_occ_mask.shape
            occlusion_fraction = np_raycasted_occ_mask.sum() / (occ_mask_shape[0] * occ_mask_shape[1])

            if transform_config.get("occluded_area_constraint") is not None:
                if min_fraction <= occlusion_fraction <= max_fraction:
                    found_suitable_occlusion = True
                else:
                    if occlusion_fraction < min_fraction:
                        high = height_viewpoint

                        if np.abs(high - low) < min_step:
                            low = max(0, low - min_step)

                    elif max_fraction < occlusion_fraction:
                        low = height_viewpoint

                        if np.abs(high - low) < min_step:
                            high += min_step
                    else:
                        raise ValueError
            else:
                found_suitable_occlusion = True

        data[ChannelEnum.OCC_MASK] = torch.logical_or(occ_mask, torch.tensor(np_raycasted_occ_mask, dtype=torch.bool))

        if ChannelEnum.OCC_DEM in data:
            data[ChannelEnum.OCC_DEM][data[ChannelEnum.OCC_MASK] == 1] = np.nan
        if ChannelEnum.OCC_DATA_UM in data:
            data[ChannelEnum.OCC_DATA_UM][data[ChannelEnum.OCC_MASK] == 1] = np.nan

        return data

    def range_data_uncertainty(self, transform_config: dict,
                               data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        res_grid = data[ChannelEnum.RES_GRID]
        robot_position = data[ChannelEnum.REL_POSITION]

        sample_grid = data[ChannelEnum.OCC_DATA_UM]

        # distance of every pixel from the robot
        lin_x = np.arange(start=-sample_grid.shape[0] / 2, stop=sample_grid.shape[0] / 2, step=1) * res_grid[0].numpy()
        lin_y = np.arange(start=-sample_grid.shape[1] / 2, stop=sample_grid.shape[1] / 2, step=1) * res_grid[1].numpy()
        off_y, off_x = np.meshgrid(lin_x, lin_y)

        dist_x = off_x - robot_position[0].numpy()
        dist_y = off_y - robot_position[1].numpy()
        dist_p2_norm = np.sqrt(np.square(dist_x) + np.square(dist_y))

        stdev = transform_config["stdev"]
        range = transform_config["range"]

        range_data_variance = stdev**2 * np.square(1 / range * dist_p2_norm)

        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                transformed_value = value + value.new_tensor(range_data_variance, dtype=value.dtype)

                data[channel] = transformed_value

        return data


import os
from PyQt4 import QtGui

class DocumentListController(object):

    def __init__(self):
        self.association = {}

    def add(self, document):
        if document.path is None:
            text = "New File"
        else:
            text = os.path.basename(document.path)

        item = QtGui.QStandardItem(text)
        self.association[item] = document

        return item

    def remove(self, document):
        for item in self.association.items():
            if document in item:
                self.association.pop(item[0])
                return item[0]

    def get_item_from_document(self, document):
        for item in self.association.items():
            if document in item:
                return item[0]

    def change_filename(self, document, new_filepath):
        item = self.get_item_from_document(document)
        item.setText(os.path.basename(new_filepath))


"""Test functions for inference.py."""

from fbt import inference, util
import numpy as np

def test_factorised_neighbor_step():
  """Tests Factorised Neighbors step."""

  contacts_all = [
    (0, 1, 2, 1),
    (1, 0, 2, 1),
    (3, 2, 2, 1),
    (2, 3, 2, 1),
    (4, 5, 2, 1),
    (5, 4, 2, 1),
    ]
  observations_all = np.array([
    (0, 2, 1)
  ])
  num_users = 6
  num_time_steps = 5

  user_interval = (0, num_users)

  p0, p1 = 0.01, 0.3
  g_param, h_param = 0.2, 0.2
  alpha, beta = 0.0001, 0.001

  seq_array = np.stack(list(
    util.iter_sequences(time_total=num_time_steps, start_se=False)))
  seq_array_hot = np.transpose(util.state_seq_to_hot_time_seq(
    seq_array, time_total=num_time_steps), [1, 2, 0]).astype(np.int8)

  prior = [1-p0, p0, 0., 0.]

  log_A_start = util.enumerate_log_prior_values(
    prior, [1-p0, 1-g_param, 1-h_param],
    seq_array, num_time_steps)

  # Precompute log(C) terms, relating to observations
  obs_array = util.make_inf_obs_array(int(num_time_steps), alpha, beta)
  log_c_z_u = util.calc_c_z_u(
    user_interval,
    obs_array,
    observations_all)

  q_marginal_infected = np.zeros((num_users, num_time_steps), dtype=np.float32)

  infect_counter = util.InfectiousContactCount(
    contacts=contacts_all,
    samples=None,
    num_users=num_users,
    num_time_steps=num_time_steps,
  )
  past_contacts = infect_counter.get_past_contacts_slice(
    list(range(user_interval[0], user_interval[1])))

  obs_diff = np.max(log_c_z_u) - np.min(log_c_z_u)
  assert obs_diff > 1.0, f"Observation difference is too small {obs_diff}"

  # 1 update
  post_exp, tstart, tend = inference.fn_step_wrapped(
    user_interval,
    seq_array_hot,
    log_c_z_u,
    log_A_start,
    q_marginal_infected,
    num_time_steps,
    p0,
    p1,
    past_contacts_array=past_contacts,
    start_belief=np.ones((num_users, 4), dtype=np.float32),
    quantization=-1)

  time_spent = tend - tstart

  assert time_spent < 1.0, f"FN takes way too long: {time_spent}"
  np.testing.assert_array_almost_equal(
    post_exp.shape, [num_users, num_time_steps, 4])

def test_fact_neigh_with_start_belief():

  contacts_all = np.array([
    (0, 1, 2, 1),
    ], dtype=np.int32)
  observations_all = np.array([
    (0, 2, 1)
  ], dtype=np.int32)
  num_users = 2
  num_time_steps = 5

  p0, p1 = 0.01, 0.5

  start_belief = np.array(
    [[.1, .4, .5, .0],
     [.9, .1, .0, .0]], dtype=np.float32)

  post_exp = inference.fact_neigh(
    num_users=num_users,
    num_time_steps=num_time_steps,
    observations_all=observations_all,
    contacts_all=contacts_all,
    probab_0=p0,
    probab_1=p1,
    g_param=.5,
    h_param=.5,
    start_belief=start_belief,
    alpha=0.001,
    beta=0.01,
    num_updates=5)

  text = ("Note this is a stochastic test. And may fail one in a thousand times"
          "Please rerun a few times")
  # Start belief for u0 is high in E and I states, so after the contact between
  # u0 and u1 on day 2, then u1 should be in E state and I state after

  with np.printoptions(precision=3, suppress=True):
    assert post_exp[1][3][1] > .2, text + "\n" + repr(post_exp)
    assert post_exp[1][4][2] > .1, text + "\n" + repr(post_exp)


import os, torch, argparse, sys
sys.path.append("..")
import torch.optim as optim
import torch.utils.data as utils
import pandas as pd
from model import *
from tqdm import tqdm
from sklearn.model_selection import train_test_split, StratifiedKFold
from torch.utils import data
from torch.optim.lr_scheduler import ReduceLROnPlateau
from tqdm import tqdm
from generator import *
from sklearn.metrics import precision_score, recall_score

parser = argparse.ArgumentParser()

parser.add_argument("--root", type=str,required=True)
parser.add_argument("--k_folds", type=int, default=1)
parser.add_argument("--weight_llrgb", type=bool, default=False)
parser.add_argument("--llrgb_factor", type=float, default=0.5)
parser.add_argument("--holdout_frac", type=float, default=0.15)
parser.add_argument("--obs_len", type=float, default=82)
parser.add_argument("--init_lr", type=float, default=0.001)
parser.add_argument("--num_epochs", type=int, default=250)

config = parser.parse_args()

def train(model, model_optimizer, input_image, input_label, input_dnu, batch_idx, train_dataloader, config):

    model_optimizer.zero_grad()

    # Combined forward pass
    outputs, kl = model.probforward(input_image, input_dnu)  # Bayesian

    # Calculate loss and backpropagate
    loss = elbo(outputs, input_label, kl, get_beta(math.ceil(len(train_dataloader) / 32), beta_type="Blundell", batch_idx=batch_idx))  # 32 is batch size

    if config.weight_llrgb: # upweight for Kepler (5), but downweight for K2 (0.25)?
        if loss[input_dnu.squeeze(1) > 9].size()[0] < 1:
            loss = torch.mean(loss[input_dnu.squeeze(1) <= 9])
        elif loss[input_dnu.squeeze(1) <= 9].size()[0] < 1:
            loss = config.llrgb_factor*torch.mean(loss[input_dnu.squeeze(1) > 9])
        else:
            loss = torch.mean(loss[input_dnu.squeeze(1) <= 9]) + config.llrgb_factor*torch.mean(loss[input_dnu.squeeze(1) > 9])
    else:
        loss = torch.mean(loss)
    loss.backward()

    # Update parameters
    model_optimizer.step()
    pred = torch.max(outputs, dim=1)[1]
    correct = torch.sum(pred.eq(input_label)).item()
    total = input_label.numel()
    acc = 100. * correct / total
    return loss.item(), acc, pred

def validate(model, val_dataloader, config):
    model.eval()  # set to evaluate mode
    val_loss = 0
    val_batch_acc = 0
    val_batches = 0
    val_pred = []
    val_truth = []
    val_prob = []
    val_kic_array = []
    val_dnuz = []
    val_logits = []
    for batch_idy, val_data in enumerate(val_dataloader, 0):  # indices,scaled_indices, numax, teff, fe_h, age, tams_age

        val_image = val_data[0].cuda().float()
        val_label = val_data[1].cuda().long().squeeze(1)
        val_dnu_var = val_data[2].cuda().float()
        val_flag = val_data[3].cuda().float()
        val_kic = val_data[4].cuda().float()

        val_image = val_image[val_flag != 1]
        val_label = val_label[val_flag != 1]
        val_dnu_var = val_dnu_var[val_flag != 1]
        val_kic = val_kic[val_flag != 1]

        if len(val_label) < 1:
            continue

        with torch.no_grad():

            outputs, kl = model.probforward(val_image, val_dnu_var) # Bayesian
            val_batch_loss = elbo(outputs, val_label, kl,
                                  get_beta(math.ceil(len(val_dataloader)) / 32, beta_type="Blundell", batch_idx=batch_idy))

            if config.weight_llrgb: # upweight for Kepler (5), but downweight for K2 (0.25)?
                if val_batch_loss[val_dnu_var.squeeze(1) > 9].size()[0] < 1:
                    val_batch_loss = torch.mean(loss[val_dnu_var.squeeze(1) <= 9])
                elif val_batch_loss[val_dnu_var.squeeze(1) <= 9].size()[0] < 1:
                    val_batch_loss = config.llrgb_factor*torch.mean(loss[val_dnu_var.squeeze(1) > 9])
                else:
                    val_batch_loss = torch.mean(val_batch_loss[val_dnu_var.squeeze(1) <= 9]) + config.llrgb_factor*torch.mean(loss[val_dnu_var.squeeze(1) > 9])
            else:
                val_batch_loss = torch.mean(val_batch_loss)

            val_logits.append(outputs.data.cpu().numpy())

        pred = torch.max(outputs, dim=1)[1]
        correct = torch.sum(pred.eq(val_label)).item()
        total = val_label.numel()
        val_loss += val_batch_loss.mean().item()
        val_batch_acc += 100. * correct / total
        val_batches += 1
        val_pred.append(pred.data.cpu().numpy())
        val_dnuz.append(val_dnu_var.data.cpu().numpy())
        val_truth.append(val_label.data.cpu().numpy())
        val_prob.append(F.softmax(outputs, dim=1).data.cpu().numpy())
        val_kic_array.append(val_kic.data.cpu().numpy())
    val_kic_array = np.concatenate(val_kic_array, 0)
    val_dnuz = np.concatenate(val_dnuz, 0)
    val_logits = np.concatenate(val_logits, 0)

    return (val_loss / val_batches), (val_batch_acc / val_batches), np.concatenate(val_pred, axis=0), np.concatenate(
        val_truth, axis=0), np.concatenate(val_prob, axis=0), val_kic_array, val_logits

def train_model(config):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print('Using device:', device)

    model = Bayes_Classifier()
    model.to(device)

    folder_kic = []
    root_folder = config.root
    for dirpath, dirnames, filenames in os.walk(root_folder):
        for filex in filenames:
            folder_kic.append(int(filex.split('-')[0]))

    folder_kic = np.unique(folder_kic)

    train_kics, val_kics = train_test_split(folder_kic, test_size=config.holdout_frac, random_state=137)
    print('Number of Train Stars: ', len(train_kics))
    print('Number of Validation Stars: ', len(val_kics))
    train_gen = npz_generator(root=root_folder, select_kic=train_kics, obs_len=config.obs_len)
    train_dataloader = utils.DataLoader(train_gen, shuffle=True, batch_size=32)

    val_gen = npz_generator(root=root_folder, select_kic=val_kics, obs_len=config.obs_len)
    val_dataloader = utils.DataLoader(val_gen, shuffle=False, batch_size=32)

    learning_rate = config.init_lr
    model_optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    scheduler = ReduceLROnPlateau(model_optimizer, mode='min', factor=0.5, patience=20, verbose=True, min_lr=1E-6)

    n_epochs = config.num_epochs
    best_loss = 1.e9
    model_checkpoint = True
    print(model_checkpoint)
    for epoch in range(1, n_epochs + 1):
        print('---------------------')
        print('Epoch: ', epoch)
        train_loss = 0
        train_batches = 0
        acc_cum = 0
        train_kic_array = []
        dnu_array = []
        pred_array = []
        model.train()  # set to training mode

        for batch_idx, data in tqdm(enumerate(train_dataloader, 0), total=len(train_dataloader), unit='batches'):

            image = data[0].cuda().float()
            label = data[1].cuda().long().squeeze(1)
            dnu = data[2].cuda().float()
            flag = data[3].cuda().float()
            train_kic = data[4].cuda().float()

            image = image[flag != 1] # flag 1 is bad stuff
            label = label[flag != 1]
            dnu = dnu[flag != 1]
            train_kic = train_kic[flag != 1]

            if len(label) < 2:
                print('Insufficient Batch!')
                continue

            loss, acc, predz = train(model, model_optimizer, image, label, dnu, batch_idx, train_dataloader, config)
            train_loss += loss  # Summing losses across all batches, so if you want the mean for EACH sample, divide by number of batches
            train_batches += 1
            acc_cum += acc
            train_kic_array.append(train_kic.data.cpu().numpy())
            pred_array.append(predz.data.cpu().numpy())
            dnu_array.append(dnu.data.cpu().numpy())
        train_loss = train_loss / train_batches
        train_acc = acc_cum / train_batches
        train_kic_array = np.concatenate(train_kic_array, 0)
        pred_array = np.concatenate(pred_array, axis = 0)
        dnu_array = np.concatenate(dnu_array, axis = 0)

        val_loss, val_acc, val_pred, val_truth, val_prob, val_kic, val_logits = validate(model, val_dataloader, config)
        scheduler.step(train_loss)  # reduce LR on loss plateau

        print('\n\nTrain Loss: ', train_loss)
        print('Train Acc: ', train_acc)

        val_loss = val_loss
        val_acc = val_acc
        print('Val Loss: ', val_loss)
        print('Val Acc: ', val_acc)

        for param_group in model_optimizer.param_groups:
            print('Current Learning Rate: ', param_group['lr'])

        if model_checkpoint:
            is_best = val_loss < best_loss
            best_loss = min(val_loss, best_loss)

            if is_best:
                filepath = 'checkpoint/Epoch%d-ACC:%.2f-Loss:%.3f' % (
                    epoch,val_acc, val_loss)
                print('Model saved to %s' % filepath)
                torch.save(model.state_dict(), filepath)
            else:
                print('No improvement over the best of %.4f' % best_loss)

def train_model_k_fold(config):

    folder_kic = []
    root_folder = config.root
    for dirpath, dirnames, filenames in os.walk(root_folder):
        for filex in filenames:
            folder_kic.append(int(filex.split('-')[0]))

    train_kics = np.copy(folder_kic)
    folder_kic = np.unique(folder_kic)

    label_data = pd.read_csv('Elsworth_Jie_2019_ID_Label.dat', header=0, delim_whitespace=True)
    label_kic = label_data['KIC'].values
    label_truth = label_data['Label'].values
    train_label = np.array([label_truth[np.where(label_kic == kicz)][0] for kicz in train_kics])

    cv = StratifiedKFold(n_splits=config.k_folds, shuffle=True, random_state=137)
    prob_over_folds = []
    truth_over_folds = []
    acc_over_folds = []
    rgb_precision_over_folds = []
    heb_precision_over_folds = []
    rgb_recall_over_folds = []
    heb_recall_over_folds = []
    fold_count = 0

    for train_idx, valid_idx in cv.split(train_kics, train_label):
        fold_count += 1
        print('Validation Fold: ', fold_count)

        model = Bayes_Classifier()
        model.cuda()
        torch.backends.cudnn.benchmark = True

        train_gen = npz_generator(root=root_folder, select_kic=train_kics[train_idx], obs_len=config.obs_len)
        train_dataloader = utils.DataLoader(train_gen, shuffle=True, batch_size=32)

        val_gen = npz_generator(root=root_folder, select_kic=train_kics[valid_idx], obs_len=config.obs_len)
        val_dataloader = utils.DataLoader(val_gen, shuffle=False, batch_size=32)

        learning_rate = config.init_lr
        model_optimizer = optim.Adam(model.parameters(), lr=learning_rate)
        scheduler = ReduceLROnPlateau(model_optimizer, mode='min', factor=0.5, patience=3, verbose=True, min_lr=1E-6)

        n_epochs = config.num_epochs
        for epoch in range(1, n_epochs + 1):
            print('---------------------')
            print('Epoch: ', epoch)
            print('Validation Fold: ', fold_count)
            train_loss = 0
            train_batches = 0
            acc_cum = 0

            model.train()  # set to training mode

            for batch_idx, data in tqdm(enumerate(train_dataloader, 0), total=len(train_dataloader), unit='batches'):

                image = data[0].cuda().float()
                label = data[1].cuda().long().squeeze(1)
                dnu = data[2].cuda().float()
                flag = data[3].cuda().float()

                image = image[flag != 1]
                label = label[flag != 1]
                dnu = dnu[flag != 1]

                if len(label) < 2:
                    print('Insufficient Batch!')
                    continue

                loss, acc, predz = train(model, model_optimizer, image, label, dnu, batch_idx, train_dataloader, config)
                train_loss += loss  # Summing losses across all batches, so if you want the mean for EACH sample, divide by number of batches
                train_batches += 1
                acc_cum += acc

            train_loss = train_loss / train_batches
            train_acc = acc_cum / train_batches

            val_loss, val_acc, val_pred, val_truth, val_prob, val_kic, val_logits = validate(model, val_dataloader, config)
            scheduler.step(train_loss)  # reduce LR on loss plateau

            print('\n\nTrain Loss: ', train_loss)
            print('Train Acc: ', train_acc)

            print('Val Loss: ', val_loss)
            print('Val Acc: ', val_acc)

            for param_group in model_optimizer.param_groups:
                print('Current Learning Rate: ', param_group['lr'])

        # After training, one last validation pass after training
        _, final_acc, final_pred, final_truth, final_prob, final_kic, final_logits = validate(model, val_dataloader, config)
        fold_precision_heb = precision_score(y_true=final_truth, y_pred=final_pred, pos_label=1)
        fold_precision_rgb = precision_score(y_true=final_truth, y_pred=final_pred, pos_label=0)
        fold_recall_heb = recall_score(y_true=final_truth, y_pred=final_pred, pos_label=1)
        fold_recall_rgb = recall_score(y_true=final_truth, y_pred=final_pred, pos_label=0)

        print('Accuracy for this fold: ', final_acc)
        print('HeB Precision for this fold: ', fold_precision_heb)
        print('HeB Recall for this fold: ', fold_recall_heb)
        print('RGB Precision for this fold: ', fold_precision_rgb)
        print('RGB Recall for this fold: ', fold_recall_rgb)

        prob_over_folds.append(final_prob)
        truth_over_folds.append(final_truth)
        acc_over_folds.append(final_acc)
        print('Accumulated Accuracy Thus Far: ', acc_over_folds)
        rgb_precision_over_folds.append(fold_precision_rgb)
        rgb_recall_over_folds.append(fold_recall_rgb)
        heb_precision_over_folds.append(fold_precision_heb)
        heb_recall_over_folds.append(fold_recall_heb)

    prob_over_folds = np.concatenate(prob_over_folds, axis=0)
    truth_over_folds = np.concatenate(truth_over_folds, axis=0)

    print(prob_over_folds.shape)
    print(truth_over_folds.shape)

    print('Validation Accuracy: %.3f +/- %.3f' % (np.mean(acc_over_folds), np.std(acc_over_folds)))
    print('Validation HeB Precision: %.3f +/- %.3f' % (
    np.mean(heb_precision_over_folds), np.std(heb_precision_over_folds)))
    print('Validation HeB Recall: %.3f +/- %.3f' % (np.mean(heb_recall_over_folds), np.std(heb_recall_over_folds)))
    print('Validation RGB Precision: %.3f +/- %.3f' % (
    np.mean(rgb_precision_over_folds), np.std(rgb_precision_over_folds)))
    print('Validation RGB Recall: %.3f +/- %.3f' % (np.mean(rgb_recall_over_folds), np.std(rgb_recall_over_folds)))

if __name__ == '__main__':
    if config.k_folds == 1:
        train_model(config)
    else:
        train_model_k_fold(config)


from django.contrib import admin

# Register your models here.
from . import models

@admin.register(models.ListModel)
class ListModel(admin.ModelAdmin):

    list_display =('ip','method','t_code','create_time')
    # 默认不配置的话，第一个字段会存在链接到记录编辑页面
    # list_display_links = None
    list_display_links = ('ip','method','t_code','create_time')
    # 每页显示条目数 缺省值100
    list_per_page = 10
    readonly_fields = ('id',)
    # 设置过滤选项
    list_filter = ('id', 'ip', 'method', 't_code','create_time')
    # 按发布日期降序排序
    ordering = ('-create_time',)
    # 搜索条件设置
    search_fields = ('ip','method','create_time')


import numpy as np
from skimage.feature import greycomatrix, greycoprops, hog
from sklearn.decomposition import PCA

def glcm_F(image, angles, distances, prop, d, features, symmetric=True, normed=True):
    n_features = len(angles) * len(distances)

    glcm = greycomatrix(image, distances=distances, angles=angles, levels=int(256/d), symmetric=symmetric, normed=normed)

    for idx, p in enumerate(prop):  # obtenim features de la matriu GLCM
        f = greycoprops(glcm, prop=p)[0, 0]
        features[(idx * n_features): (idx + 1) * n_features] = f.flatten()

def HOG(image, px, orientations=9, block_norm="L2"):
    features, img = hog(image, orientations=orientations, pixels_per_cell=(px, px),
                        cells_per_block=(1, 1), transform_sqrt=False,
                        block_norm=block_norm, feature_vector=True, visualise=True)

    return features, img

def pca_F(train_set, test_set, percent=0.95):
    pca = PCA(n_components=train_set.shape[1])

    pca.fit(train_set)

    cumsum = np.cumsum(pca.explained_variance_ratio_)
    r = np.where(cumsum > percent)
    pca = PCA(n_components=r[0][0])

    x_train = pca.fit_transform(train_set)
    x_test = pca.transform(test_set)

    print("Train", x_train.shape[0], "Test", x_test.shape[0])

    return x_train, x_test