import os
import time
import yaml
import cv2
import numpy as np
import paddle
from benchmark_utils import PaddleInferBenchmark
from preprocess import preprocess, NormalizeImage, Permute
from mot_preprocess import LetterBoxResize

from tracker import JDETracker
from ppdet.modeling.mot import visualization as mot_vis
from ppdet.modeling.mot.utils import Timer as MOTTimer

from paddle.inference import Config
from paddle.inference import create_predictor
from utils import argsparser, Timer, get_current_memory_mb
from infer import get_test_images, print_arguments

# Global dictionary
MOT_SUPPORT_MODELS = {
    'JDE',
    'FairMOT',
}

class MOT_Detector(object):
    """
    Args:
        pred_config (object): config of model, defined by `Config(model_dir)`
        model_dir (str): root path of model.pdiparams, model.pdmodel and infer_cfg.yml
        use_gpu (bool): whether use gpu
        run_mode (str): mode of running(fluid/trt_fp32/trt_fp16)
        batch_size (int): size of pre batch in inference
        trt_min_shape (int): min shape for dynamic shape in trt
        trt_max_shape (int): max shape for dynamic shape in trt
        trt_opt_shape (int): opt shape for dynamic shape in trt
        trt_calib_mode (bool): If the model is produced by TRT offline quantitative
            calibration, trt_calib_mode need to set True
        cpu_threads (int): cpu threads
        enable_mkldnn (bool): whether to open MKLDNN
    """

    def __init__(self,
                 pred_config,
                 model_dir,
                 use_gpu=False,
                 run_mode='fluid',
                 batch_size=1,
                 trt_min_shape=1,
                 trt_max_shape=1088,
                 trt_opt_shape=608,
                 trt_calib_mode=False,
                 cpu_threads=1,
                 enable_mkldnn=False):
        self.pred_config = pred_config
        self.predictor, self.config = load_predictor(
            model_dir,
            run_mode=run_mode,
            batch_size=batch_size,
            min_subgraph_size=self.pred_config.min_subgraph_size,
            use_gpu=use_gpu,
            use_dynamic_shape=self.pred_config.use_dynamic_shape,
            trt_min_shape=trt_min_shape,
            trt_max_shape=trt_max_shape,
            trt_opt_shape=trt_opt_shape,
            trt_calib_mode=trt_calib_mode,
            cpu_threads=cpu_threads,
            enable_mkldnn=enable_mkldnn)
        self.det_times = Timer()
        self.cpu_mem, self.gpu_mem, self.gpu_util = 0, 0, 0
        self.tracker = JDETracker()

    def preprocess(self, im):
        preprocess_ops = []
        for op_info in self.pred_config.preprocess_infos:
            new_op_info = op_info.copy()
            op_type = new_op_info.pop('type')
            preprocess_ops.append(eval(op_type)(**new_op_info))
        im, im_info = preprocess(im, preprocess_ops)
        inputs = create_inputs(im, im_info)
        return inputs

    def postprocess(self, pred_dets, pred_embs):
        online_targets = self.tracker.update(pred_dets, pred_embs)
        online_tlwhs, online_ids = [], []
        online_scores = []
        for t in online_targets:
            tlwh = t.tlwh
            tid = t.track_id
            tscore = t.score
            vertical = tlwh[2] / tlwh[3] > 1.6
            if tlwh[2] * tlwh[3] > self.tracker.min_box_area and not vertical:
                online_tlwhs.append(tlwh)
                online_ids.append(tid)
                online_scores.append(tscore)
        return online_tlwhs, online_scores, online_ids

    def predict(self, image, threshold=0.5, repeats=1):
        '''
        Args:
            image (dict): dict(['image', 'im_shape', 'scale_factor'])
            threshold (float): threshold of predicted box' score
        Returns:
            online_tlwhs, online_ids (np.ndarray)
        '''
        self.det_times.preprocess_time_s.start()
        inputs = self.preprocess(image)
        self.det_times.preprocess_time_s.end()
        pred_dets, pred_embs = None, None
        input_names = self.predictor.get_input_names()
        for i in range(len(input_names)):
            input_tensor = self.predictor.get_input_handle(input_names[i])
            input_tensor.copy_from_cpu(inputs[input_names[i]])

        self.det_times.inference_time_s.start()
        for i in range(repeats):
            self.predictor.run()
            output_names = self.predictor.get_output_names()
            boxes_tensor = self.predictor.get_output_handle(output_names[0])
            pred_dets = boxes_tensor.copy_to_cpu()
            embs_tensor = self.predictor.get_output_handle(output_names[1])
            pred_embs = embs_tensor.copy_to_cpu()

        self.det_times.inference_time_s.end(repeats=repeats)

        self.det_times.postprocess_time_s.start()
        online_tlwhs, online_scores, online_ids = self.postprocess(pred_dets,
                                                                   pred_embs)
        self.det_times.postprocess_time_s.end()
        self.det_times.img_num += 1
        return online_tlwhs, online_scores, online_ids

def create_inputs(im, im_info):
    """generate input for different model type
    Args:
        im (np.ndarray): image (np.ndarray)
        im_info (dict): info of image
        model_arch (str): model type
    Returns:
        inputs (dict): input of model
    """
    inputs = {}
    inputs['image'] = np.array((im, )).astype('float32')
    inputs['im_shape'] = np.array((im_info['im_shape'], )).astype('float32')
    inputs['scale_factor'] = np.array(
        (im_info['scale_factor'], )).astype('float32')
    return inputs

class PredictConfig_MOT():
    """set config of preprocess, postprocess and visualize
    Args:
        model_dir (str): root path of model.yml
    """

    def __init__(self, model_dir):
        # parsing Yaml config for Preprocess
        deploy_file = os.path.join(model_dir, 'infer_cfg.yml')
        with open(deploy_file) as f:
            yml_conf = yaml.safe_load(f)
        self.check_model(yml_conf)
        self.arch = yml_conf['arch']
        self.preprocess_infos = yml_conf['Preprocess']
        self.min_subgraph_size = yml_conf['min_subgraph_size']
        self.labels = yml_conf['label_list']
        self.mask = False
        self.use_dynamic_shape = yml_conf['use_dynamic_shape']
        if 'mask' in yml_conf:
            self.mask = yml_conf['mask']
        self.print_config()

    def check_model(self, yml_conf):
        """
        Raises:
            ValueError: loaded model not in supported model type
        """
        for support_model in MOT_SUPPORT_MODELS:
            if support_model in yml_conf['arch']:
                return True
        raise ValueError("Unsupported arch: {}, expect {}".format(yml_conf[
            'arch'], MOT_SUPPORT_MODELS))

    def print_config(self):
        print('-----------  Model Configuration -----------')
        print('%s: %s' % ('Model Arch', self.arch))
        print('%s: ' % ('Transform Order'))
        for op_info in self.preprocess_infos:
            print('--%s: %s' % ('transform op', op_info['type']))
        print('--------------------------------------------')

def load_predictor(model_dir,
                   run_mode='fluid',
                   batch_size=1,
                   use_gpu=False,
                   min_subgraph_size=3,
                   use_dynamic_shape=False,
                   trt_min_shape=1,
                   trt_max_shape=1088,
                   trt_opt_shape=608,
                   trt_calib_mode=False,
                   cpu_threads=1,
                   enable_mkldnn=False):
    """set AnalysisConfig, generate AnalysisPredictor
    Args:
        model_dir (str): root path of __model__ and __params__
        run_mode (str): mode of running(fluid/trt_fp32/trt_fp16/trt_int8)
        batch_size (int): size of pre batch in inference
        use_gpu (bool): whether use gpu
        use_dynamic_shape (bool): use dynamic shape or not
        trt_min_shape (int): min shape for dynamic shape in trt
        trt_max_shape (int): max shape for dynamic shape in trt
        trt_opt_shape (int): opt shape for dynamic shape in trt
        trt_calib_mode (bool): If the model is produced by TRT offline quantitative
            calibration, trt_calib_mode need to set True
        cpu_threads (int): cpu threads
        enable_mkldnn (bool): whether to open MKLDNN
    Returns:
        predictor (PaddlePredictor): AnalysisPredictor
    Raises:
        ValueError: predict by TensorRT need use_gpu == True.
    """
    if not use_gpu and not run_mode == 'fluid':
        raise ValueError(
            "Predict by TensorRT mode: {}, expect use_gpu==True, but use_gpu == {}"
            .format(run_mode, use_gpu))
    config = Config(
        os.path.join(model_dir, 'model.pdmodel'),
        os.path.join(model_dir, 'model.pdiparams'))
    precision_map = {
        'trt_int8': Config.Precision.Int8,
        'trt_fp32': Config.Precision.Float32,
        'trt_fp16': Config.Precision.Half
    }
    if use_gpu:
        # initial GPU memory(M), device ID
        config.enable_use_gpu(200, 0)
        # optimize graph and fuse op
        config.switch_ir_optim(True)
    else:
        config.disable_gpu()
        config.set_cpu_math_library_num_threads(cpu_threads)
        if enable_mkldnn:
            try:
                # cache 10 different shapes for mkldnn to avoid memory leak
                config.set_mkldnn_cache_capacity(10)
                config.enable_mkldnn()
            except Exception as e:
                print(
                    "The current environment does not support `mkldnn`, so disable mkldnn."
                )
                pass

    if run_mode in precision_map.keys():
        config.enable_tensorrt_engine(
            workspace_size=1 << 10,
            max_batch_size=batch_size,
            min_subgraph_size=min_subgraph_size,
            precision_mode=precision_map[run_mode],
            use_static=False,
            use_calib_mode=trt_calib_mode)

        if use_dynamic_shape:
            min_input_shape = {'image': [1, 3, trt_min_shape, trt_min_shape]}
            max_input_shape = {'image': [1, 3, trt_max_shape, trt_max_shape]}
            opt_input_shape = {'image': [1, 3, trt_opt_shape, trt_opt_shape]}
            config.set_trt_dynamic_shape_info(min_input_shape, max_input_shape,
                                              opt_input_shape)
            print('trt set dynamic shape done!')

    # disable print log when predict
    config.disable_glog_info()
    # enable shared memory
    config.enable_memory_optim()
    # disable feed, fetch OP, needed by zero_copy_run
    config.switch_use_feed_fetch_ops(False)
    predictor = create_predictor(config)
    return predictor, config

def write_mot_results(filename, results, data_type='mot'):
    if data_type in ['mot', 'mcmot', 'lab']:
        save_format = '{frame},{id},{x1},{y1},{w},{h},{score},-1,-1,-1\n'
    elif data_type == 'kitti':
        save_format = '{frame} {id} pedestrian 0 0 -10 {x1} {y1} {x2} {y2} -10 -10 -10 -1000 -1000 -1000 -10\n'
    else:
        raise ValueError(data_type)

    with open(filename, 'w') as f:
        for frame_id, tlwhs, tscores, track_ids in results:
            if data_type == 'kitti':
                frame_id -= 1
            for tlwh, score, track_id in zip(tlwhs, tscores, track_ids):
                if track_id < 0:
                    continue
                x1, y1, w, h = tlwh
                x2, y2 = x1 + w, y1 + h
                line = save_format.format(
                    frame=frame_id,
                    id=track_id,
                    x1=x1,
                    y1=y1,
                    x2=x2,
                    y2=y2,
                    w=w,
                    h=h,
                    score=score)
                f.write(line)

class FLAGS(object):
	def __init__(self):
		# self.model_dir = None
		self.model_dir = ".\\models"
		self.image_file = None
		self.image_dir = None
		self.batch_size = 1
		# self.video_file = None
		self.video_file = ".\\videos\\ttt.mp4"
		self.camera_id = -1
		self.threshold = 0.5
		self.output_dir = "output"
		self.run_mode = "fluid"
		self.device = "cpu"
		self.use_gpu = False
		self.run_benchmark = False
		self.enable_mkldnn = False
		self.cpu_threads = 1
		self.trt_min_shape = 1
		self.trt_max_shape = 1280
		self.trt_opt_shape = 640
		self.trt_calib_mode = False
		self.save_images = True
		self.save_results = True

class MOT_Predict(object):
    def __init__(self, FLAGS):
        paddle.enable_static()
        self.flag = FLAGS
        print_arguments(self.flag)

        self.flag = FLAGS
        pred_config = PredictConfig_MOT(self.flag.model_dir)
        self.detector = MOT_Detector(
            pred_config,
            self.flag.model_dir,
            use_gpu=self.flag.use_gpu,
            run_mode=self.flag.run_mode,
            trt_min_shape=self.flag.trt_min_shape,
            trt_max_shape=self.flag.trt_max_shape,
            trt_opt_shape=self.flag.trt_opt_shape,
            trt_calib_mode=self.flag.trt_calib_mode,
            cpu_threads=self.flag.cpu_threads,
            enable_mkldnn=self.flag.enable_mkldnn)
        self.fps = 30
        self.frame_id = 0

    def predict_via_frame(self, frame):
        online_tlwhs, online_scores, online_ids = self.detector.predict(
        frame, self.flag.threshold)

        online_im = mot_vis.plot_tracking(
            frame,
            online_tlwhs,
            online_ids,
            online_scores,
            frame_id=self.frame_id,
            fps=self.fps)
        count = len(online_ids)

        return online_im,count

    def test(self):
        capture = cv2.VideoCapture(self.flag.video_file)
        frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
        print('frame_count', frame_count)

        cv2.namedWindow("Tracking Detection",0);
        cv2.resizeWindow("Tracking Detection", 640, 480);
        while True:
            ret, frame = capture.read()
            if not ret:
                break
            online_im = self.predict_via_frame(frame)
            im = np.array(online_im)

            self.frame_id += 1
            print('detect frame:%d' % (self.frame_id))

            cv2.imshow('Tracking Detection', im)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

def main():
    my_mot_predict = MOT_Predict(FLAGS)
    my_mot_predict.test()

if __name__ == '__main__':
    # paddle.enable_static()
    # # parser = argsparser()
    # # FLAGS = parser.parse_args()
    FLAGS = FLAGS()
    # print_arguments(FLAGS)
    main()

def predict_video(detector, camera_id, FLAGS):
    if camera_id != -1:
        capture = cv2.VideoCapture(camera_id)
        video_name = 'mot_output.mp4'
    else:
        capture = cv2.VideoCapture(FLAGS.video_file)
        video_name = os.path.split(FLAGS.video_file)[-1]
    fps = 30
    frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    print('frame_count', frame_count)
    width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
    # yapf: disable
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    # yapf: enable
    if not os.path.exists(FLAGS.output_dir):
        os.makedirs(FLAGS.output_dir)
    out_path = os.path.join(FLAGS.output_dir, video_name)
    writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
    frame_id = 0
    timer = MOTTimer()
    results = []
    while (1):
        ret, frame = capture.read()
        if not ret:
            break
        timer.tic()
        online_tlwhs, online_scores, online_ids = detector.predict(
            frame, FLAGS.threshold)
        timer.toc()

        results.append((frame_id + 1, online_tlwhs, online_scores, online_ids))
        fps = 1. / timer.average_time
        online_im = mot_vis.plot_tracking(
            frame,
            online_tlwhs,
            online_ids,
            online_scores,
            frame_id=frame_id,
            fps=fps)
        if FLAGS.save_images:
            save_dir = os.path.join(FLAGS.output_dir, video_name.split('.')[-2])
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            cv2.imwrite(
                os.path.join(save_dir, '{:05d}.jpg'.format(frame_id)),
                online_im)
        frame_id += 1
        print('detect frame:%d' % (frame_id))
        im = np.array(online_im)
        writer.write(im)
        if camera_id != -1:
            cv2.imshow('Tracking Detection', im)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    if FLAGS.save_results:
        result_filename = os.path.join(FLAGS.output_dir,
                                       video_name.split('.')[-2] + '.txt')
        write_mot_results(result_filename, results)
    writer.release()


import os
import shutil
import zipfile
import requests
import pathlib
from tkinter import Tk
from tkinter.filedialog import askopenfilename, askdirectory

mandatory_mods = [
    "https://thunderstore.io/package/download/Lakatrazz/Fusion/1.3.2/",
    "https://thunderstore.io/package/download/gnonme/BoneLib/2.2.1/",
]

def download_file(url, save_path):
    response = requests.get(url, stream=True)
    content_disposition = response.headers.get('Content-Disposition')

    if content_disposition:
        file_name = content_disposition.split('filename=')[-1].strip('"')
    else:
        file_name = url.split('/')[-2] + '.zip'

    save_path = os.path.join(os.path.dirname(save_path), file_name)
    total_size = int(response.headers.get('content-length', 0))

    with open(save_path, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)
            bytes_downloaded = f.tell()
            percentage = int(100 * bytes_downloaded / total_size)
            print(f"Downloading {url}: {percentage}%", end='\r')

    return save_path

def extract_zip(zip_path, extract_path):
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        for member in zip_ref.infolist():
            if not member.is_dir():
                print(f"Extracting {member.filename}...")
                try:
                    zip_ref.extract(member, extract_path)
                    print(f"Successfully extracted {member.filename}")
                except Exception as e:
                    print(f"Failed to extract {member.filename}: {e}")

def extract_and_move_files(extract_path, dll_destination, appdata_destination):
    for root, dirs, files in os.walk(extract_path):
        for file in files:
            file_path = os.path.join(root, file)

            # Move files from Mods folder to the specified folder
            if "Mods" in root:
                destination_path = os.path.join(dll_destination, file)
                print(f"Moving {file_path} to {destination_path}")
                shutil.move(file_path, destination_path)

            # Move files from Plugins folder to the specified folder
            if "Plugins" in root:
                plugins_path = os.path.join(os.path.dirname(dll_destination), "Plugins")
                destination_path = os.path.join(plugins_path, file)
                print(f"Moving {file_path} to {destination_path}")
                shutil.move(file_path, destination_path)

        # Move other contents to AppData LocalLow directory
        for directory in dirs:
            if directory not in ["Mods", "Plugins"]:
                source_path = os.path.join(root, directory)
                destination_path = os.path.join(appdata_destination, directory)
                print(f"Moving {source_path} to {destination_path}")
                shutil.move(source_path, destination_path)

def get_saved_dll_folder(file_path):
    if os.path.exists(file_path):
        with open(file_path, 'r') as f:
            return f.read().strip()
    return None

def save_dll_folder(file_path, folder_path):
    with open(file_path, 'w') as f:
        f.write(folder_path)

def clear_directory(directory):
    for item in os.listdir(directory):
        item_path = os.path.join(directory, item)
        if os.path.isfile(item_path):
            os.remove(item_path)
        else:
            shutil.rmtree(item_path)

def main():
    Tk().withdraw()
    profile_file = askopenfilename(title="Select Profile Text File")

    script_dir = os.path.dirname(os.path.realpath(__file__))
    config_file = os.path.join(script_dir, 'config.txt')
    dll_folder = get_saved_dll_folder(config_file)
    if not dll_folder:
        dll_folder = askdirectory(title="Select Folder for DLL Files")
        save_dll_folder(config_file, dll_folder)

    with open(profile_file, 'r') as file:
        download_links = file.read().splitlines()

    profile_dir = os.path.dirname(profile_file)

    download_links = list(set(download_links + mandatory_mods))

    clear_directory(dll_folder)
    app_data_folder = os.path.join(os.getenv('APPDATA'), '..\\LocalLow\\Stress Level Zero\\BONELAB\\Mods')
    clear_directory(app_data_folder)

    for link in download_links:
        file_name = link.split('/')[-1]
        save_path = os.path.join(profile_dir, file_name)
        save_path = download_file(link, save_path)

        file_name = os.path.basename(save_path)
        extract_folder = os.path.join(profile_dir, os.path.splitext(file_name)[0])
        extract_zip(save_path, extract_folder)

        extract_and_move_files(extract_folder, dll_folder, app_data_folder)

        os.makedirs(app_data_folder, exist_ok=True)

        for item in os.listdir(extract_folder):
            item_path = os.path.join(extract_folder, item)
            destination = os.path.join(app_data_folder, item)
            if os.path.isfile(item_path):
                shutil.move(item_path, destination)
            else:
                if os.path.exists(destination):
                    shutil.rmtree(destination)
                shutil.move(item_path, destination)

        os.remove(save_path)
        shutil.rmtree(extract_folder)

if __name__ == "__main__":
    main()


from .compute_cluster_penalties_cy import do_compute_penalties
from numpy import zeros, amax

__all__ = ['compute_penalties']

def compute_penalties(kk, clusters):
        if clusters is None:
            clusters = kk.clusters
        num_clusters = amax(clusters)+1
        cluster_penalty = zeros(num_clusters)
        ustart = kk.data.unmasked_start
        uend = kk.data.unmasked_end
        penalty_k = kk.penalty_k
        penalty_k_log_n = kk.penalty_k_log_n
        float_num_unmasked = kk.data.float_num_unmasked
        num_spikes = zeros(num_clusters, dtype=int)

        do_compute_penalties(cluster_penalty, num_spikes, clusters,
                             penalty_k, penalty_k_log_n, float_num_unmasked)

        return cluster_penalty


from skimage.feature import peak_local_max
import cv2
import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage as ndi
import sys
import torch
# from support.layer.nms import nms
import torchvision

from skimage import filters
from skimage.measure import regionprops

rel_peak_thr = .3
rel_rel_thr = .3
ioa_thr = .6
topk_boxes = 3

def heat2bbox(heat_map, original_image_shape):
    h, w = heat_map.shape

    bounding_boxes = []

    heat_map = heat_map - np.min(heat_map)
    heat_map = heat_map / np.max(heat_map)

    bboxes = []
    box_scores = []

    peak_coords = peak_local_max(heat_map, exclude_border=False,
                                 threshold_rel=rel_peak_thr)  # find local peaks of heat map

    heat_resized = cv2.resize(heat_map, (
    original_image_shape[1], original_image_shape[0]))  ## resize heat map to original image shape
    peak_coords_resized = ((peak_coords + 0.5) *
                           np.asarray([original_image_shape]) /
                           np.asarray([[h, w]])
                           ).astype('int32')

    for pk_coord in peak_coords_resized:
        pk_value = heat_resized[tuple(pk_coord)]
        mask = heat_resized > pk_value * rel_rel_thr
        labeled, n = ndi.label(mask)
        l = labeled[tuple(pk_coord)]
        yy, xx = np.where(labeled == l)
        min_x = np.min(xx)
        min_y = np.min(yy)
        max_x = np.max(xx)
        max_y = np.max(yy)
        bboxes.append((min_x, min_y, max_x, max_y))
        box_scores.append(pk_value)  # you can change to pk_value * probability of sentence matching image or etc.

    ## Merging boxes that overlap too much
    box_idx = np.argsort(-np.asarray(box_scores))
    box_idx = box_idx[:min(topk_boxes, len(box_scores))]
    bboxes = [bboxes[i] for i in box_idx]
    box_scores = [box_scores[i] for i in box_idx]

    to_remove = []
    for iii in range(len(bboxes)):
        for iiii in range(iii):
            if iiii in to_remove:
                continue
            b1 = bboxes[iii]
            b2 = bboxes[iiii]
            isec = max(min(b1[2], b2[2]) - max(b1[0], b2[0]), 0) * max(min(b1[3], b2[3]) - max(b1[1], b2[1]), 0)
            ioa1 = isec / ((b1[2] - b1[0]) * (b1[3] - b1[1]))
            ioa2 = isec / ((b2[2] - b2[0]) * (b2[3] - b2[1]))
            if ioa1 > ioa_thr and ioa1 == ioa2:
                to_remove.append(iii)
            elif ioa1 > ioa_thr and ioa1 >= ioa2:
                to_remove.append(iii)
            elif ioa2 > ioa_thr and ioa2 >= ioa1:
                to_remove.append(iiii)

    for i in range(len(bboxes)):
        if i not in to_remove:
            bounding_boxes.append({
                'score': box_scores[i],
                'bbox': bboxes[i],
                'bbox_normalized': np.asarray([
                    bboxes[i][0] / heat_resized.shape[1],
                    bboxes[i][1] / heat_resized.shape[0],
                    bboxes[i][2] / heat_resized.shape[1],
                    bboxes[i][3] / heat_resized.shape[0],
                ]),
            })

    return bounding_boxes

def img_heat_bbox_disp(image, heat_map, title='', en_name='', alpha=0.6, cmap='viridis', cbar='False', dot_max=False,
                       bboxes=[], order=None, show=True):
    thr_hit = 1  # a bbox is acceptable if hit point is in middle 85% of bbox area
    thr_fit = .60  # the biggest acceptable bbox should not exceed 60% of the image
    H, W = image.shape[0:2]
    # resize heat map
    heat_map_resized = cv2.resize(heat_map, (H, W))

    # display
    fig = plt.figure(figsize=(15, 5))
    fig.suptitle(title, size=15)
    ax = plt.subplot(1, 3, 1)
    plt.imshow(image)
    if dot_max:
        max_loc = np.unravel_index(np.argmax(heat_map_resized, axis=None), heat_map_resized.shape)
        plt.scatter(x=max_loc[1], y=max_loc[0], edgecolor='w', linewidth=3)

    if len(bboxes) > 0:  # it gets normalized bbox
        if order == None:
            order = 'xxyy'

        for i in range(len(bboxes)):
            bbox_norm = bboxes[i]
            if order == 'xxyy':
                x_min, x_max, y_min, y_max = int(bbox_norm[0] * W), int(bbox_norm[1] * W), int(bbox_norm[2] * H), int(
                    bbox_norm[3] * H)
            elif order == 'xyxy':
                x_min, x_max, y_min, y_max = int(bbox_norm[0] * W), int(bbox_norm[2] * W), int(bbox_norm[1] * H), int(
                    bbox_norm[3] * H)
            x_length, y_length = x_max - x_min, y_max - y_min
            box = plt.Rectangle((x_min, y_min), x_length, y_length, edgecolor='w', linewidth=3, fill=False)
            plt.gca().add_patch(box)
            if en_name != '':
                ax.text(x_min + .5 * x_length, y_min + 10, en_name,
                        verticalalignment='center', horizontalalignment='center',
                        # transform=ax.transAxes,
                        color='white', fontsize=15)
                # an = ax.annotate(en_name, xy=(x_min,y_min), xycoords="data", va="center", ha="center", bbox=dict(boxstyle="round", fc="w"))
                # plt.gca().add_patch(an)

    plt.imshow(heat_map_resized, alpha=alpha, cmap=cmap)

    # plt.figure(2, figsize=(6, 6))
    plt.subplot(1, 3, 2)
    plt.imshow(image)
    # plt.figure(3, figsize=(6, 6))
    plt.subplot(1, 3, 3)
    plt.imshow(heat_map_resized)
    fig.tight_layout()
    fig.subplots_adjust(top=.85)

    if show:
        plt.show()
    else:
        plt.close()

    return fig

def filter_bbox(bbox_dict, order=None):
    thr_fit = .90  # the biggest acceptable bbox should not exceed 80% of the image
    if order == None:
        order = 'xxyy'

    filtered_bbox = []
    filtered_bbox_norm = []
    filtered_score = []
    if len(bbox_dict) > 0:  # it gets normalized bbox
        for i in range(len(bbox_dict)):
            bbox = bbox_dict[i]['bbox']
            bbox_norm = bbox_dict[i]['bbox_normalized']
            bbox_score = bbox_dict[i]['score']
            if order == 'xxyy':
                x_min, x_max, y_min, y_max = bbox_norm[0], bbox_norm[1], bbox_norm[2], bbox_norm[3]
            elif order == 'xyxy':
                x_min, x_max, y_min, y_max = bbox_norm[0], bbox_norm[2], bbox_norm[1], bbox_norm[3]
            if bbox_score > 0:
                x_length, y_length = x_max - x_min, y_max - y_min
                if x_length * y_length < thr_fit:
                    filtered_score.append(bbox_score)
                    filtered_bbox.append(bbox)
                    filtered_bbox_norm.append(bbox_norm)
    return filtered_bbox, filtered_bbox_norm, filtered_score

def crop_resize_im(image, bbox, size, order='xxyy'):
    H, W, _ = image.shape
    if order == 'xxyy':
        roi = image[int(bbox[2] * H):int(bbox[3] * H), int(bbox[0] * W):int(bbox[1] * W), :]
    elif order == 'xyxy':
        roi = image[int(bbox[1] * H):int(bbox[3] * H), int(bbox[0] * W):int(bbox[2] * W), :]
    roi = cv2.resize(roi, size)
    return roi

def im2double(im):
    return cv2.normalize(im.astype('float'), None, 0.0, 1.0, cv2.NORM_MINMAX)

def IoU(boxA, boxB):
    # order = xyxy
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])

    # compute the area of intersection rectangle
    interArea = max(0, xB - xA) * max(0, yB - yA)

    # compute the area of both the prediction and ground-truth
    # rectangles
    boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
    boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])

    # compute the intersection over union by taking the intersection
    # area and dividing it by the sum of prediction + ground-truth
    # areas - the interesection area
    iou = interArea / float(boxAArea + boxBArea - interArea)

    # return the intersection over union value
    return iou

def img_heat_bbox_disp(image, heat_map, size, title='', en_name='', alpha=0.6, cmap='viridis', cbar='False', dot_max=True,
                       gt=[], bbox=[], order=None, show=False, filename=''):

    heat_map = (256 * heat_map).astype(np.uint8)
    thr_hit = 1  # a bbox is acceptable if hit point is in middle 85% of bbox area
    thr_fit = .60  # the biggest acceptable bbox should not exceed 60% of the image
    H, W = size
    # resize heat map
    heat_map_resized = cv2.resize(heat_map, (W, H))
    image = cv2.resize(image, (W, H))
    # heat_map_resized = heat_map
    # display
    fig = plt.figure(figsize=(15, 5))
    fig.suptitle(title, size=15)
    ax = plt.subplot(1, 3, 1)
    plt.imshow(image)
    if dot_max:
        max_loc = np.unravel_index(np.argmax(heat_map_resized, axis=None), heat_map_resized.shape)
        plt.scatter(x=max_loc[1], y=max_loc[0], edgecolor='w', linewidth=3)

    # if len(bboxes) > 0:  # it gets normalized bbox
    #     if order == None:
    #         order = 'xxyy'
    #
    #     for i in range(len(bboxes)):
    #         bbox_norm = bboxes[i]
    #         if order == 'xxyy':
    #             x_min, x_max, y_min, y_max = int(bbox_norm[0] * W), int(bbox_norm[1] * W), int(bbox_norm[2] * H), int(
    #                 bbox_norm[3] * H)
    #         elif order == 'xyxy':
    #             x_min, x_max, y_min, y_max = int(bbox_norm[0] * W), int(bbox_norm[2] * W), int(bbox_norm[1] * H), int(
    #                 bbox_norm[3] * H)
    # x_length, y_length = x_max - x_min, y_max - y_min

    # if en_name != '':
    #     ax.text(x_min + .5 * x_length, y_min + 10, en_name,
    #             verticalalignment='center', horizontalalignment='center',
    #             # transform=ax.transAxes,
    #             color='white', fontsize=15)
        # an = ax.annotate(en_name, xy=(x_min,y_min), xycoords="data", va="center", ha="center", bbox=dict(boxstyle="round", fc="w"))
        # plt.gca().add_patch(an)

    plt.imshow(heat_map_resized, alpha=alpha, cmap=cmap)

    # plt.figure(2, figsize=(6, 6))
    plt.subplot(1, 3, 2)
    plt.imshow(image)
    # box = plt.Rectangle((gt[0, 0], gt[0, 1]), gt[0, 2] - gt[0, 0], gt[0, 3] - gt[0, 1],
    #                     edgecolor='w', linewidth=3, fill=False, color='green')
    # plt.gca().add_patch(box)
    # if len(bbox) > 0:
    #     box = plt.Rectangle((bbox[0][0], bbox[0][1]), bbox[0][2] - bbox[0][0], bbox[0][3] - bbox[0][1],
    #                         edgecolor='w', linewidth=3, fill=False, color='red')
    #     plt.gca().add_patch(box)
    # plt.figure(3, figsize=(6, 6))
    plt.subplot(1, 3, 3)
    plt.imshow(heat_map_resized, cmap='jet')
    fig.tight_layout()
    fig.subplots_adjust(top=.85)

    if show:
        plt.show()
    else:
        plt.savefig(filename, bbox_inches='tight')
    plt.close()

def isCorrect(bbox_annot, bbox_pred, iou_thr=.5, size_h=224):
    for bbox_p in bbox_pred:
        bbox_p = (np.array(bbox_p) / size_h).tolist()
        for bbox_a in bbox_annot:
            if IoU(bbox_p, bbox_a) >= iou_thr:
                return 1
    return 0

def isCorrectHit(bbox_annot, heatmap, orig_img_shape):
    H, W = orig_img_shape
    heatmap_resized = cv2.resize(heatmap, (W, H))
    max_loc = np.unravel_index(np.argmax(heatmap_resized, axis=None), heatmap_resized.shape)

    # try:
    #     threshold_value = filters.threshold_minimum(heatmap_resized)
    #     labeled_foreground = (heatmap_resized > threshold_value).astype(int)
    #     properties = regionprops(labeled_foreground, heatmap_resized)
    #     center_of_mass = properties[0].centroid
    #     weighted_center_of_mass = properties[0].weighted_centroid
    #     max_loc = weighted_center_of_mass
    # except:
    #     max_loc = np.unravel_index(np.argmax(heatmap_resized, axis=None), heatmap_resized.shape)

    for bbox in bbox_annot:
        if bbox[0] <= max_loc[1] <= bbox[2] and bbox[1] <= max_loc[0] <= bbox[3]:
            return 1
    return 0

def check_percent(bboxes):
    for bbox in bboxes:
        x_length = bbox[2] - bbox[0]
        y_length = bbox[3] - bbox[1]
        if x_length * y_length < .05:
            return False
    return True

def union(bbox):
    if len(bbox) == 0:
        return []
    if type(bbox[0]) == type(0.0) or type(bbox[0]) == type(0):
        bbox = [bbox]
    maxes = np.max(bbox, axis=0)
    mins = np.min(bbox, axis=0)
    return [[mins[0], mins[1], maxes[2], maxes[3]]]

def attCorrectness(bbox_annot, heatmap, orig_img_shape):
    H, W = orig_img_shape
    heatmap_resized = cv2.resize(heatmap, (W, H))
    h_s = np.sum(heatmap_resized)
    if h_s == 0:
        return 0
    else:
        heatmap_resized /= h_s
    att_correctness = 0
    for bbox in bbox_annot:
        x0, y0, x1, y1 = bbox
        att_correctness += np.sum(heatmap_resized[y0:y1, x0:x1])
    return att_correctness

def calc_correctness(annot, heatmap, orig_img_shape):
    # bbox_dict = heat2bbox(heatmap, orig_img_shape)
    size_h = heatmap.shape[-1]
    bbox_dict = generate_bbox(heatmap)
    # bbox, bbox_norm, bbox_score = filter_bbox(bbox_dict=bbox_dict, order='xyxy')
    annot = process_gt_bbox(annot, orig_img_shape)
    bbox_norm_annot = union(annot['bbox_norm'])
    bbox_annot = annot['bbox']
    bbox_dict = union(np.array(bbox_dict)[:, :4])
    bbox_correctness = isCorrect(bbox_norm_annot, bbox_dict, iou_thr=.5, size_h=size_h)
    hit_correctness = isCorrectHit(bbox_annot, heatmap, orig_img_shape)
    # att_correctness = attCorrectness(bbox_annot, heatmap, orig_img_shape)
    # return bbox_correctness, hit_correctness, att_correctness, bbox
    # return bbox_correctness, hit_correctness, att_correctness
    return bbox_correctness, hit_correctness, 0

def process_gt_bbox(annot, orig_img_shape):
    out = {}
    h, w = orig_img_shape
    bbox = torch.tensor(annot).numpy()
    out['bbox'] = bbox.copy()
    bbox = bbox.astype(np.float)
    bbox[:, 0] = bbox[:, 0] / w
    bbox[:, 1] = bbox[:, 1] / h
    bbox[:, 2] = bbox[:, 2] / w
    bbox[:, 3] = bbox[:, 3] / h
    out['bbox_norm'] = bbox.copy()
    return out

def no_tuple(a):
    out = []
    for item in a:
        out.append(item[0])
    return out

def cluster_gt(bbox, heatmap):
    out = {}
    bbox = torch.tensor(bbox).numpy().squeeze()
    for i in range(len(bbox)):
        curr_heatmap = heatmap[i:i+1].clone()
        curr_bbox = bbox[i].copy()
        if len(out.keys()) == 0:
            out[str(len(out.keys()))] = (curr_bbox, curr_heatmap)
            continue
        flag = False
        for j, key in enumerate(out.keys()):
            old_bbox, old_heatmap = out[key]
            if IoU(curr_bbox, old_bbox) == 1:
                out[key] = (curr_bbox, torch.cat((old_heatmap, curr_heatmap), dim=0))
                flag = True
                break
        if not flag:
            out[str(len(out.keys()))] = (curr_bbox, curr_heatmap)

    for ix, key in enumerate(out.keys()):
        curr_bbox, old_heatmap = out[key]
        old_heatmap = old_heatmap.mean(dim=0).squeeze().detach().clone().cpu().numpy()
        out[key] = (np.expand_dims(curr_bbox, axis=0),
                    old_heatmap)
    return out

def intensity_to_rgb(intensity, cmap='cubehelix', normalize=False):
    assert intensity.ndim == 2, intensity.shape
    intensity = intensity.astype("float")

    if normalize:
        intensity -= intensity.min()
        intensity /= intensity.max()

    cmap = plt.get_cmap(cmap)
    intensity = cmap(intensity)[..., :3]
    return intensity.astype('float32') * 255.0

def generate_bbox(cam, threshold=0.5, nms_threshold=0.05, max_drop_th=0.5):
    heatmap = intensity_to_rgb(cam, normalize=True).astype('uint8')
    gray_heatmap = cv2.cvtColor(heatmap, cv2.COLOR_RGB2GRAY)

    thr_val = threshold * np.max(gray_heatmap)

    _, thr_gray_heatmap = cv2.threshold(gray_heatmap,
                                        int(thr_val), 255,
                                        cv2.THRESH_TOZERO)
    try:
        _, contours, _ = cv2.findContours(thr_gray_heatmap,
                                          cv2.RETR_TREE,
                                          cv2.CHAIN_APPROX_SIMPLE)
    except:
        contours, _ = cv2.findContours(thr_gray_heatmap,
                                       cv2.RETR_TREE,
                                       cv2.CHAIN_APPROX_SIMPLE)
    if len(contours) != 0:
        proposals = [cv2.boundingRect(c) for c in contours]
        # proposals = [(x, y, w, h) for (x, y, w, h) in proposals if h * w > 0.05 * 224 * 224]
        if len(proposals) > 0:
            proposals_with_conf = [thr_gray_heatmap[y:y + h, x:x + w].mean()/255 for (x, y, w, h) in proposals]
            inx = torchvision.ops.nms(torch.tensor(proposals).float(),
                                      torch.tensor(proposals_with_conf).float(),
                                      nms_threshold)
            estimated_bbox = torch.cat((torch.tensor(proposals).float()[inx],
                                        torch.tensor(proposals_with_conf)[inx].unsqueeze(dim=1)),
                                       dim=1).tolist()
            estimated_bbox = [(x, y, x+w, y+h, conf) for (x, y, w, h, conf) in estimated_bbox
                              if conf > max_drop_th * np.max(proposals_with_conf)]
        else:
            estimated_bbox = [[0, 0, 1, 1, 0], [0, 0, 1, 1, 0]]
    else:
        estimated_bbox = [[0, 0, 1, 1, 0], [0, 0, 1, 1, 0]]
    return estimated_bbox

def generate_proposals(cam, threshold=0.5):
    heatmap = intensity_to_rgb(cam, normalize=True).astype('uint8')
    gray_heatmap = cv2.cvtColor(heatmap, cv2.COLOR_RGB2GRAY)
    thr_val = threshold * np.max(gray_heatmap)
    _, gray_heatmap = cv2.threshold(gray_heatmap,
                                    int(thr_val), 255,
                                    cv2.THRESH_TOZERO)
    try:
        _, contours, _ = cv2.findContours(gray_heatmap,
                                          cv2.RETR_TREE,
                                          cv2.CHAIN_APPROX_SIMPLE)
    except:
        contours, _ = cv2.findContours(gray_heatmap,
                                       cv2.RETR_TREE,
                                       cv2.CHAIN_APPROX_SIMPLE)
    proposals = []
    for c in contours:
        if c.shape[0]<5:
            continue
        x, y, w, h = cv2.boundingRect(c)
        energy = gray_heatmap[y:y + h, x:x + w].sum() / gray_heatmap.sum()
        if energy < 1:
            proposals.append([x, y, x + w, y + h, energy])
    return proposals, gray_heatmap

def get_scores(proposals, mask):
    scores = []
    bboxes = []
    for p in proposals:
        x, y, w, h = p
        if w*h < 0.05*224*224:
            continue
        energy = (mask[y:y + h, x:x + w, :]/255/3).sum()/(w*h)
        scores.append(energy)
        bboxes.append([x, y, x+w, h+y])
    return bboxes, scores


# -*- coding: utf-8 -*-
import telebot
from telebot import types
import sqlite3
import io #для работы с битовыми данными в памяти
import docx # для работы с Word
import datetime # для работы с временем
import SQLfunctions # подключаем sqlite функции
from PIL import Image, ImageDraw, ImageFont # библиотека по работе с изображениями, добавление текста, шрифты

conn = sqlite3.connect('running.db')# подключаем sqlite

# курсор для работы с таблицами
cursor = conn.cursor()

try:
    # sql запрос на создание таблицы
    query = "CREATE TABLE \"runningData\" (\"ID\" INTEGER UNIQUE, \"user_id\" INTEGER, \"run\" TEXT, PRIMARY KEY (\"ID\"))"
    cursor.execute(query)
except:
    pass

try:
    query1 = "CREATE TABLE \"rewardData\" (\"ID\" INTEGER UNIQUE, \"user_id\" INTEGER, \"reward\" TEXT, PRIMARY KEY (\"ID\"))"
    cursor.execute(query1)
except:
    pass

# подключим токен нашего бота
bot = telebot.TeleBot(SQLfunctions.token)

# напишем, что делать нашему боту при команде старт
@bot.message_handler(commands=['start'])
def send_keyboard(message, text="Привет, я HSE_Running Bot! Твой персональный помощник для учета пробежек! \n"
                                "Записывай пробежки в формате:\n"
                                'Дистанция (км), Время (мин), Локация, Дата (дд.мм.гггг)\n\n'
                                'Пример:\n2 км, 10 мин, Москва, 12.10.2021 \n\n'
                                'Ты также можешь записать в пробежку только дистанцию и время пробежки.\n\n'
                                'Пример:\n2 км, 10 мин\n'
                                'Но в таком случае я не смогу учесть такие пробежки при показе активности за последний месяц или '
                                'при показе пробежек за определенный период.\n\n'
                                "Если у тебя уже есть записи пробежек в файлах, можешь загрузить их сюда выбрав 'Импорт записей'"):
    keyboard = types.ReplyKeyboardMarkup(row_width=2)
    itembtn1 = types.KeyboardButton('Добавить пробежку')
    itembtn2 = types.KeyboardButton('Удалить пробежку')
    itembtn3 = types.KeyboardButton('Показать все пробежки')
    itembtn4 = types.KeyboardButton('Показать пробежки за период')
    itembtn5 = types.KeyboardButton('Показать активность за месяц')
    itembtn6 = types.KeyboardButton('Импорт записей')  # загрузить записи из файла
    itembtn7 = types.KeyboardButton('Награды')
    itembtn8 = types.KeyboardButton('Пока все!')
    itembtn9 = types.KeyboardButton('Очистить runningData')
    itembtn10 = types.KeyboardButton('Очистить rewardData')

    # keyboard.add(itembtn1, itembtn2, itembtn3, itembtn4, itembtn5, itembtn6, itembtn7, itembtn8, itembtn9, itembtn10)
    keyboard.add(itembtn1, itembtn2, itembtn6, itembtn3, itembtn4, itembtn5, itembtn7, itembtn8, itembtn9, itembtn10)
    # пришлем это все сообщением и запишем выбранный вариант
    msg = bot.send_message(message.from_user.id, text=text, reply_markup=keyboard)

    # отправим этот вариант в функцию, которая его обработает
    bot.register_next_step_handler(msg, callback_worker)

#--------------------------------
# Теперь напишем функции для нашего бота

# функция "Добавить пробежку"
def add_run(msg):
    flag_marathon = 0      # переменная, которая нужна для отслеживания, пробежал ли пользователь марафон
    flag_half_marathon = 0 # переменная, которая нужна для отслеживания, пробежал ли пользователь полумарафон
    if (msg.text != 'Назад'):
        try:
            curDist = float(str(msg.text).split(',')[0].split()[0].lstrip('(').lstrip("'"))
            if curDist >= 0:
                if ',' in str(msg.text).rstrip(')').rstrip(','):
                    if curDist == 42:
                        flag_marathon = 1
                    elif curDist == 21:
                        flag_half_marathon = 1
                SQLfunctions.insertRunIntoBD(msg, msg.text)
                runText = msg.text # сохраним текущую пробежку в переменную, потом передадим ее в функцию memoryPhoto
                #-----
                bot.send_message(msg.chat.id, 'Вау, это надо зафиксировать :)')
                reward(msg, flag_marathon, flag_half_marathon) # вызываем функцию учета наград
                markup = types.ReplyKeyboardMarkup()
                markup.add(types.KeyboardButton('Да'))
                markup.add(types.KeyboardButton('Нет'))
                msg = bot.send_message(msg.from_user.id, text='Сделаем памятную фотку?', reply_markup=markup)
                bot.register_next_step_handler(msg, memoryPhoto, runText)  # переходим к функции memory photo
            else:
                bot.send_message(msg.chat.id, "Ваша дистанция - отрицательное число. Извините, я не могу добавить эту пробежку в список.")
                send_keyboard(msg, "Чем еще могу помочь?")
        except:
            bot.send_message(msg.chat.id, 'Вы ввели пробежку в неправильном формате! Извините, я не могу её добавить в список')
            send_keyboard(msg, "Чем еще могу помочь?")
    else:
        send_keyboard(msg, "Хорошо, отменяем. Чем еще могу помочь?")

# функция для выбора сделать фото или нет
def memoryPhoto(msg, runText):
    # bot.send_message(msg.chat.id, 'Попали в функцию memoryPhoto')
    if (msg.text == 'Да'):
        bot.send_message(msg.chat.id, 'Хорошо, скинь фотку... ')
        bot.register_next_step_handler(msg, image_open, runText)
        #send_keyboard(msg, 'Могу помочь чем-то еще?')
    else:
        bot.send_message(msg.chat.id, 'Нет так нет')
        send_keyboard(msg, 'Могу помочь чем-то еще?')

# реализуем функцию "Обработки изображений"
@bot.message_handler(content_types=['photo'])
def image_open(msg, runText):
    if (msg.text != 'Назад'):
        try:
            file_info = bot.get_file(msg.photo[len(msg.photo) - 1].file_id)
            downloaded_file = bot.download_file(file_info.file_path) #сохраняем фото с сервера телеграмм
            fp = io.BytesIO(downloaded_file) # декодируем изображение из байт
            img = Image.open(fp) # откроем изображение с библиотекой Pillow
            font = ImageFont.load_default() # стандартный шрифт
            text = str(runText)[-10:].encode('cp1251')#.encode('UTF-8') #добавляемый шрифт
            draw_text = ImageDraw.Draw(img) # добавим текст на изображение
            draw_text.text((10,10), text, font=font, fill=('#58F9F6')) # параметры размещения текста на изображении
            #img.save('test.jpg')
            bot.send_photo(msg.chat.id, img)
            bot.send_message(msg.chat.id, "Ваше фото готово!")
            send_keyboard(msg, 'Могу помочь чем-то еще?')
        except Exception as e:
            bot.reply_to(msg, e)

    else:
        bot.send_message(msg, "Хорошо, отменяем. Чем еще могу помочь?")

# функция "Импорт из файла"
# теперь напишем функцию "Импорт записей" для загрузки записей из файлов txt и docx
# Это функция красивых строк для текстовых файлов(docx)
def get_document_string(documentList):
    documentList_str1 = [line.rstrip() for line in documentList]
    documentList_str2 = [str(val) for val in
                    documentList_str1]
    return '\n'.join(documentList_str2)

# Это функция красивых строк для текстовых файлов(txt)
def get_document_string_txt(documentList):
    documentList_str = [str(val) for val in
                    documentList]
    return '\n'.join(documentList_str)

# функция записи пробежек из файла
@bot.message_handler(content_types=['document'])
def import_run(msg):
    flag_marathon = 0
    flag_half_marathon = 0
    text = []
    if (msg.text != 'Назад'):
        bot.send_message(msg.chat.id, "Вы загрузили документ")
        try:
            chat_id = msg.chat.id
            file_info = bot.get_file(msg.document.file_id) #получаем информацию о файле на сервере telegram
            downloaded_file = bot.download_file(file_info.file_path) #записываем в память битовое представление файла
            end_file_txt = msg.document.file_name.endswith('.txt') #расширение файла txt
            end_file_docx = msg.document.file_name.endswith('.docx') #расширение файла docx
            if (end_file_txt == True) or (end_file_docx == True):

                if end_file_txt == True:
                    text = downloaded_file.decode('utf-8').split('\n')
                    doc_str = get_document_string_txt(text)
                    bot.send_message(msg.chat.id, 'Я сохраню эти записи:')
                    bot.send_message(msg.chat.id, doc_str)

                elif end_file_docx == True:
                     doc = docx.Document(io.BytesIO(downloaded_file))
                     text = [paragraph.text for paragraph in doc.paragraphs]
                     doc_str = get_document_string(text)
                     bot.send_message(msg.chat.id, 'Я сохраню эти записи:')
                     bot.send_message(msg.chat.id, doc_str)

            else: #если расширение другое, бот отправит стикер sorry
                bot.send_sticker(msg.chat.id, 'CAACAgIAAxkBAAIKJGFrNvGpUtjbltZpXQNpetpU1FC6AAI_AAOtZbwUhPhHiSSyKyAhBA')
        except Exception as e:
            bot.reply_to(msg, e)

        record_run = [line.rstrip() for line in text] #запишем построчно данные из файла в БД
        for i in range(len(record_run)):
            try:
                curDist = float(str(record_run[i]).split(',')[0].split()[0].lstrip('(').lstrip("'"))
                if curDist >= 0:
                    if ',' in str(record_run[i]):
                        if curDist == 42:
                            flag_marathon = 1
                        elif curDist == 21:
                            flag_half_marathon = 1
                    SQLfunctions.insertRunIntoBD(msg, record_run[i])
                else:
                    bot.send_message(msg.chat.id,
                                     "Ваша дистанция - отрицательное число. Извините, я не могу добавить следующую пробежку в список:")
                    bot.send_message(msg.chat.id, record_run[i])
            except:
                bot.send_message(msg.chat.id,
                                 'Вы ввели пробежку в неправильном формате! Извините, я не могу добавить следующую пробежку в список:')
                bot.send_message(msg.chat.id, record_run[i])

        reward(msg, flag_marathon, flag_half_marathon)
        bot.send_message(msg.chat.id, 'Я всё зафиксировал :)')
        send_keyboard(msg, 'Могу помочь чем-то еще?')
    else:
        send_keyboard(msg, "Хорошо, отменяем. Чем еще могу помочь?")

# реализуем теперь функцию "Показать все пробежки"
# для этого сначала напишем функцию get_runs_string, которая делает красивые строки и отправляет их пользователю
def get_runs_string(runsList):
    runsList_str = []
    for val in list(enumerate(runsList)):
        runsList_str.append(str(val[0] + 1) + ') ' + val[1][0] + '\n')
    return ''.join(runsList_str)

# теперь отправляем пользователю список пробежек
def showAllRuns(msg):
    runsList = get_runs_string(SQLfunctions.selectRunsfromBD(msg))
    bot.send_message(msg.chat.id, runsList)
    send_keyboard(msg, "Чем еще могу помочь?")

# реализуем функцию "Удалить пробежку"
# для этого сначала пользователь должен выбрать пробежку, которую он хочет удалить
def chooseToDelete(msg):
    markup = types.ReplyKeyboardMarkup()
    runsList = SQLfunctions.selectRunsfromBD(msg)
    if len(runsList) != 0:
        for val in runsList:
            markup.add(types.KeyboardButton(val[0]))
        msg = bot.send_message(msg.from_user.id, text="Выбери пробежку для удаления", reply_markup=markup)
        bot.register_next_step_handler(msg, confirmDelete)
    else:
        bot.send_message(msg.from_user.id,'Здесь пусто. Нечего удалить.')
        send_keyboard(msg, "Чем еще могу помочь?")

# спрашиваем, точно ли пользователь хочет удалить пробежку
def confirmDelete(msg):
    bot.send_message(msg.from_user.id, 'Вы собираетесь удалить следующую пробежку:')
    bot.send_message(msg.from_user.id, msg.text)
    markup = types.ReplyKeyboardMarkup()
    markup.add(types.KeyboardButton('Да'))
    markup.add(types.KeyboardButton('Нет'))
    msg1 = bot.send_message(msg.from_user.id, text='Удаляем?', reply_markup=markup)
    bot.register_next_step_handler(msg1, deleteRun, msg)

# теперь напишем саму функцию deleteRun
def deleteRun(msg1, msg):
    if msg1.text == 'Да':
        SQLfunctions.deleteRunFromBD(msg)
        bot.send_message(msg.chat.id, 'Пробежка удалена')
        send_keyboard(msg, "Чем еще могу помочь?")
    else:
        send_keyboard(msg, "Окей, ничего не удаляем! Чем еще могу помочь?")

# перейдем к функции Получения, просмотра и учета наград
# сначала добавим функцию для красивого вывода строк с наградой
def get_reward_string(rewardList):
    rewList_str = []
    for val in list(rewardList):
        rewList_str.append(str('"' + val[0] + '"' + '\n'))
    return ''.join(rewList_str)

# теперь переходим к самой функции
def reward(msg, flag_marathon, flag_half_marathon):
    runsList = SQLfunctions.selectRunsfromBD(msg)
    rewardList = get_reward_string(SQLfunctions.selectRewardsfromBD(msg))

    summdistance = 0  # переменная, в которую будем записывать суммарное расстояние, которое уже пробежал пользователь
    for val in runsList:
        # try:
        if ',' in str(val).rstrip(')').rstrip(','):
            summdistance += float(str(val).split(',')[0].split()[0].lstrip('(').lstrip("'"))
        else:
            bot.send_message(msg.chat.id, "================\n"
                                              "УВЕДОМЛЕНИЕ\n"
                                              "* Вы вашем списке пробежек есть пробежка в нерекомендуемом формате!\n"
                                              "* Пробежка в нерекомендуемом формате не влияет на работу бота, однако данные из данной"
                                              " пробежки никак не будут учтены\n"
                                              "* Рекомендуется удалить данную пробежку:\n")
            bot.send_message(msg.chat.id, val)
            bot.send_message(msg.chat.id, "================")
        # except:
        #     bot.send_message(msg.chat.id, "================\n"
        #                                   "УВЕДОМЛЕНИЕ\n"
        #                                   "* Вы вашем списке пробежек есть пробежка в нерекомендуемом формате!\n"
        #                                   "* Пробежка в нерекомендуемом формате не влияет на работу бота, однако данные из данной"
        #                                   " пробежки никак не будут учтены\n"
        #                                   "* Рекомендуется удалить данную пробежку:\n")
        #     bot.send_message(msg.chat.id, val)
        #     bot.send_message(msg.chat.id, "================")

    if (flag_marathon == 1) and ("Пробежать марафон!" not in rewardList):
        bot.send_message(msg.chat.id, 'Да вы марафонец!')
        bot.send_sticker(msg.chat.id, 'CAACAgIAAxkBAAEDG4JhbVH5kd2eiIFKEUWJ8g6jxBehlgACtQADwZxgDNJKKKDyOCUEIQQ')
        SQLfunctions.insertRewardIntoBD(msg, "Пробежать марафон!")

    if (flag_half_marathon == 1) and ("Пробежать полумарафон!" not in rewardList):
        bot.send_message(msg.chat.id, 'Ого, вы пробежали полумарафон!')
        bot.send_sticker(msg.chat.id, 'CAACAgIAAxkBAAEDG4RhbVUKixP0SEOvFQQ8nIb2QDvkgAACtwADwZxgDPilirtWD6kDIQQ')
        SQLfunctions.insertRewardIntoBD(msg, "Пробежать полумарафон!")

    if (summdistance > 0 and "Начало положено!" not in rewardList):
        bot.send_message(msg.chat.id, 'Ваша первая пробежка, круто!')
        bot.send_sticker(msg.chat.id, 'CAACAgIAAxkBAAEDEl1hZzjsZ1hhdMvFXy5uI4Twg_THfwACxAADMNSdEcjFvLwK6xVKIQQ')
        SQLfunctions.insertRewardIntoBD(msg, "Начало положено!")

    if (summdistance >= 100 and "Преодолеть отметку в 100 км!" not in rewardList):
        bot.send_message(msg.chat.id, f'Вау, вы преодолели отметку в 100 км! Так держать!')
        bot.send_sticker(msg.chat.id, 'CAACAgIAAxkBAAEDFKBhaFqay25yDk89nZ53k2qPmwLMOAACywADMNSdEXWK5xlwcXbPIQQ')
        SQLfunctions.insertRewardIntoBD(msg, "Преодолеть отметку в 100 км!")

    if (summdistance >= 200 and "Преодолеть отметку в 200 км!" not in rewardList):
        bot.send_message(msg.chat.id, 'Вау, вы преодолели отметку в 200 км! Очень круто!')
        bot.send_sticker(msg.chat.id, 'CAACAgIAAxkBAAEDG5xhbWhoD5PCuVHwuXIfomPf1-7XrgACzAADMNSdEbg0CDIOCTHMIQQ')
        SQLfunctions.insertRewardIntoBD(msg, "Преодолеть отметку в 200 км!")

# функция, которая выводит список всех полученных наград
def showRewards(msg):
    rewardList = get_reward_string(SQLfunctions.selectRewardsfromBD(msg))
    bot.send_message(msg.chat.id, rewardList)
    send_keyboard(msg, "Чем еще могу помочь?")

# напишем теперь функцию, которая удаляет ВСЕ пробежки из БД
# для этого сначала добавим функцию, которая будет спрашивать, уверен ли пользователь,
# что он хочет удалить все пробежки
def confirmDeleteAllRunningData(msg):
    markup = types.ReplyKeyboardMarkup()
    markup.add(types.KeyboardButton('Да'))
    markup.add(types.KeyboardButton('Нет'))
    msg = bot.send_message(msg.from_user.id, text='Вы уверены, что хотите удалить ВСЕ пробежки?', reply_markup=markup)
    bot.register_next_step_handler(msg, deleteAllRunningData)

# сама функция удаления пробежек
def deleteAllRunningData(msg):
    if (msg.text == 'Да'):
        SQLfunctions.deleteAllRunsFromBD(msg)
        send_keyboard(msg, "Все пробежки удалены. Чем еще могу помочь?")
    else:
        send_keyboard(msg, "Ок, ничего не удаляем. Чем еще могу помочь?")

# напишем теперь функцию, которая удаляет ВСЕ награды из БД
# для этого сначала добавим функцию, которая будет спрашивать, уверен ли пользователь,
# что он хочет удалить все награды
def confirmDeleteAllRewardData(msg):
    markup = types.ReplyKeyboardMarkup()
    markup.add(types.KeyboardButton('Да'))
    markup.add(types.KeyboardButton('Нет'))
    msg = bot.send_message(msg.from_user.id, text='Вы уверены, что хотите удалить ВСЕ достижения?', reply_markup=markup)
    bot.register_next_step_handler(msg, deleteAllRewardData)

# сама функция удаления наград
def deleteAllRewardData(msg):
    if (msg.text == 'Да'):
        SQLfunctions.deleteAllRewardsFromBD(msg)
        send_keyboard(msg, "Все награды удалены. Чем еще могу помочь?")
    else:
        send_keyboard(msg, "Ок, ничего не удаляем. Чем еще могу помочь?")

# функция создания сводки
def activity_while_month(msg):
    runsList = SQLfunctions.selectRunsfromBD(msg)
    distance = 0
    time = 0
    places = []
    for i in runsList:
        for j in i:
            try:
                info = str(j).split(',')
                date = info[3].replace(' ','').split('.')
                now = datetime.datetime.now().strftime('%y-%m-%d')
                bd = datetime.date(int(date[2]),int(date[1]),int(date[0])).strftime('%y-%m-%d')
                if (datetime.datetime.strptime(now,'%y-%m-%d') - datetime.datetime.strptime(bd,'%y-%m-%d')).days <= 30:
                    distance += float(str(i).split(',')[0].split()[0].lstrip('(').lstrip("'"))
                    time += int(str(i).split(',')[1].split()[0].lstrip('(').lstrip("'"))
                    places.append(str(info[2]))
            except:
                # bot.send_message(msg.chat.id, "Следующая пробежка не удовлетворяет требованиям к формату, поэтому не будет учтена")
                # bot.send_message(msg.chat.id, i)
                pass

    if len(places) != 0:
        a = max(set(places), key=places.count)
    else:
        a = 'Нет мест'
    bot.send_message(msg.chat.id, 'Расстояние: {} км'.format(distance))
    bot.send_message(msg.chat.id, 'Время: {} минут'.format(time))
    bot.send_message(msg.chat.id, 'Самое часто посещаемое место: {}'.format(a))
    send_keyboard(msg, "Чем еще могу помочь?")

# програмка на вход получает количество месяцев, а дальше смотрит на разницу по датам и подбирает все пробежки
def runs_in_period(msg):
    answer = ''
    runsList = SQLfunctions.selectRunsfromBD(msg)
    if len(runsList) != 0:
        spisok = []
        for i in runsList:
            for j in i:
                try:
                    info = str(j).split(',')
                    date = info[3].replace(' ', '').split('.')
                    now = datetime.datetime.now().strftime('%y-%m-%d')
                    bd = datetime.date(int(date[2]), int(date[1]), int(date[0])).strftime('%y-%m-%d')
                    if (datetime.datetime.strptime(now, '%y-%m-%d') - datetime.datetime.strptime(bd, '%y-%m-%d')).days <= int(msg.text) * 30:
                        spisok.append(j)
                except:
                    pass

        if len(spisok) != 0:
            for i in spisok:
                answer += str(i) + '\n'
            bot.send_message(msg.from_user.id, answer)
            send_keyboard(msg, "Чем еще могу помочь?")
        else:
            bot.send_message(msg.from_user.id, "У вас нет пробежек за указанный период")
            send_keyboard(msg, "Чем еще могу помочь?")
    else:
        bot.send_message(msg.from_user.id, "У вас еще нет пробежек")
        send_keyboard(msg, "Чем еще могу помочь?")
#--------------------------------

# привязываем функции к кнопкам на клавиатуре
def callback_worker(call):
    if call.text == "Добавить пробежку":
        markup = types.ReplyKeyboardMarkup(row_width=1)
        markup.add(types.KeyboardButton('Назад'))
        msg = bot.send_message(call.chat.id, text='Давайте добавим пробежку! Напишите ее в чат.\n'
                                       'Помните, что пробежку нужно записать в формате:\n'
                                       'Дистанция (км), Время (мин), Локация, Дата (дд.мм.гггг)\n\n'
                                       'Пример:\n2 км, 10 мин, Москва, 12.10.2021 \n\n'
                                       'Если нажали на команду по ошибке, нажмите "Назад" на клавиатуре',
                                        reply_markup=markup)
        bot.register_next_step_handler(msg, add_run)

    # кнопка импорта файлов
    elif call.text == "Импорт записей":
        markup = types.ReplyKeyboardMarkup(row_width=1)
        markup.add(types.KeyboardButton('Назад'))
        msg = bot.send_message(call.chat.id,
                               'Давайте добавим ваши записи пробежек! Нажмите на скрепку и прикрепите файл!\n'
                               'Я умею работать только с файлами .txt или .docx\n '
                               'Желательно, чтобы записи пробежек были на отдельных строках и в формате:\n'
                               'Дистанция (км), Время (мин), Локация, Дата (дд.мм.гггг)\n\n'
                               'Если нажали на команду по ошибке, нажмите "Назад" на клавиатуре',
                               reply_markup=markup)
        bot.register_next_step_handler(msg, import_run)

    elif call.text == "Показать все пробежки":
        try:
            showAllRuns(call)
        except:
            bot.send_message(call.chat.id, 'Здесь пусто. Может начнем бежать прямо сейчас? ;)')
            send_keyboard(call, "Чем еще могу помочь?")

    elif call.text == "Удалить пробежку":
        chooseToDelete(call)

    elif call.text == "Награды":
        bot.send_message(call.chat.id, 'Давайте посмотрим на ваши награды!...')
        try:
            showRewards(call)
        except:
            bot.send_message(call.chat.id, 'У вас пока нет наград :(')
            bot.send_sticker(call.chat.id, 'CAACAgIAAxkBAAEDFKRhaGPl0VAFJ14oaW1t7nnHVPEw8wACdQADwDZPE5B0WbJxOIvjIQQ')
            send_keyboard(call, "Чем еще могу помочь?")

    elif call.text == 'Показать активность за месяц':
        bot.send_message(call.chat.id, 'Вот ваша активность за последний месяц')
        activity_while_month(call)

    elif call.text == 'Показать пробежки за период':
        msg = bot.send_message(call.from_user.id, 'Введите количество месяцев')
        bot.register_next_step_handler(msg, runs_in_period)

    elif call.text == "Пока все!":
        bot.send_message(call.chat.id, 'Хорошего дня! Когда захотите продолжить нажмите на команду\n /start')

    elif call.text == 'Очистить runningData':
        confirmDeleteAllRunningData(call)

    elif call.text == 'Очистить rewardData':
        confirmDeleteAllRewardData(call)

    else:
        bot.send_message(call.chat.id, 'Извините, я вас не понимаю :(')
        send_keyboard(call, "Пожалуйста, выберите функции из предложенных")

bot.polling(none_stop=True)


import gzip
import bz2

s = ''
with open('BinaryData.py', 'r') as f:
    s = f.read()

with gzip.open('B.gz', 'w') as f:
    f.write(s.encode())

with bz2.open('B.bz2', 'w') as f:
    f.write(s.encode())


import random

from simple_yahtzee import SuperSimpleYahtzeeEnv

class BaseAgent(object):
    def act(self, state, ava_actions):
        return random.choice(ava_actions)

def play(max_episode=100, num_dice=4, num_eyes=6):
    env = SuperSimpleYahtzeeEnv(num_dice, num_eyes, 'silent')
    agent = BaseAgent()
    total_scores = []

    for _ in range(max_episode):
        state = env.reset()
        done = False
        while not done:
            ava_actions = env.available_actions()
            action = agent.act(state, ava_actions)
            state, reward, done, info = env.step(action)
        env.render()
        env.show_result()
        total_scores.append(env.get_total_score())

    print(total_scores)
    print(sum(total_scores)/len(total_scores))

if __name__ == '__main__':
    play()


import os
import sys
from sqlalchemy import Column, ForeignKey, Integer, String
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship
from sqlalchemy import create_engine
from eralchemy import render_er

Base = declarative_base()
class Characters(Base):
    __tablename__ = "characters"
    id = Column(Integer, primary_key=True)
    name = Column(String(250), nullable=False)
    gender = Column(String(250), nullable=False)
    Age = Column(Integer, nullable=False)

class Planet(Base):
    __tablename__= "planets"
    id = Column(Integer, primary_key=True)
    name = Column(String(250), nullable=False)
    location = Column(String(250), nullable=False)
class User(Base):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True)
    name = Column(String(250), nullable=False)
    favoritesid = Column(Integer, ForeignKey("favorites.id"))

class Favorite(Base):
    __tablename__ = 'favorites'
    id = Column(Integer, primary_key=True)
    planetid = Column(Integer, ForeignKey("planets.id"))
    characterid = Column(Integer, ForeignKey("characters.id"))
    user = relationship(User)
    character = relationship(Characters)
    planet = relationship(Planet)

    def to_dict(self):
        return {}

## Draw from SQLAlchemy base
render_er(Base, 'diagram.png')


import os
import boto3
import numpy as np
import tensorflow as tf
from flask import Flask
from dotenv import load_dotenv
from pymongo import MongoClient
from keras.models import load_model
from sklearn.preprocessing import LabelEncoder
from werkzeug.datastructures import FileStorage
from werkzeug.middleware.proxy_fix import ProxyFix
from flask_restplus import Api, Resource
from utils.Model import ModelManager
load_dotenv()

# Mongodb connection
client = MongoClient(os.environ['MONGO_CLIENT_URL'])
db = client.registry

# AWS S3 connection
session = boto3.Session(
    aws_access_key_id=os.environ['AWS_ACCESS_KEY_ID'],
    aws_secret_access_key=os.environ['AWS_SECRET_KEY']
)
s3 = session.resource('s3')

# App and API setup
app = Flask(__name__)
app.wsgi_app = ProxyFix(app.wsgi_app)

api = Api(app, version="1.0", title="Anomaly Detection", description="")
ns = api.namespace('api')
single_parser = api.parser()
single_parser.add_argument("files", location="files", type=FileStorage, action='append', required=True)

graph = tf.get_default_graph()
backup_model = load_model("./models/backup/model.h5")
backup_label_encoder = LabelEncoder()
backup_label_encoder.classes_ = np.load("./models/backup/classes.npy")

@ns.route("/classify")
class MultiClassification(Resource):
    @api.doc(parser=single_parser, description='Upload an image of a solar panel')
    def post(self):
        model = ModelManager(db, s3, graph, backup_model, backup_label_encoder,
            bucket_name=os.environ['AWS_BUCKET_NAME'])
        model.load_latest_model()

        args = single_parser.parse_args()
        image_files = args.files

        preds = []
        for image in image_files:
            image_array = model.preprocess(image)
            preds.append(model.predict(image_array)[0])

        return {'prediction': str(preds)}

if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0")


import pandas as pn
import numpy as np
import math
from keras.layers.core import Flatten
from keras.layers import Embedding, Multiply, Dense, Input
from keras.models import Model
from keras.layers.merge import concatenate
from keras.models import load_model
from time import time
import pickle
import argparse
import matplotlib.pyplot as plt

def parse_args() :
    parser = argparse.ArgumentParser(description="Training the model.")
    parser.add_argument('--path_save', nargs = '?', default= 'D:/',
                        help= 'Path to save the results.')
    parser.add_argument('--model', type = int, default= 1,
                        help= 'Model to train : 1)GMF  2)HybMLP  3)NHybNeuMF 4)MLP  5)NCF.')
    parser.add_argument('--pretrain', type = int, default= 0,
                        help= 'Pretrain : Enable 1, Disable 0.')
    parser.add_argument('--epochs', type=int, default=1,
                        help='Number of epochs.')
    parser.add_argument('--optimizer_subs', nargs='?', default='adam',
                        help='Specify an optimizer for the submodels : adam, sgd.')
    parser.add_argument('--optimizer_nhf', nargs='?', default='sgd',
                        help='Specify an optimizer for NHybF : adam, sgd.')
    parser.add_argument('--optimizer_ncf',  nargs='?' , default='sgd',
                        help='Specify an optimizer for NCF : adam, sgd.')
    parser.add_argument('--batch_size', type=int, default=100,
                        help ='Batch size.')
    parser.add_argument('--num_layers', type=int, default=5,
                        help ='HybMLP number of layers.')
    parser.add_argument('--emb_size_gmf', type=int, default=64,
                        help ='GMF embedding size.')
    parser.add_argument('--emb_size_hmlp', type=int, default=32,
                        help ='HybMLP embedding size.')
    parser.add_argument('--emb_size_mlp', type=int, default=32,
                        help ='MLP embedding size.')
    parser.add_argument('--num_factors', type=int, default=16,
                        help ='Number of predictive factors for HybMLP.')
    parser.add_argument('--num_neg', type=int, default=1,
                        help ='Number of negative instances to pair with a positive instance.')
    parser.add_argument('--estop', type = int, default = -1,
                        help = 'Early Stopping : Enable 1, Disable -1.')
    parser.add_argument('--path_data', nargs = '?', default= 'D:/',
                        help = 'Input trainset and testset folder path.')
    parser.add_argument('--dataset', type = int, default = 1,
                        help = 'Choose a dataset:\n 1)MovieLens 100K  2)MovieLens 1M  3)Yelp.')
    return parser.parse_args()

def save_obj(obj, path):
    with open(path+'.pkl', 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)

def load_data(path, dataset):
    if dataset == 3 :
        yelp_columns = ['user_id','item_id', 'rating', 'fans','average_stars','friends','vote_funny','useful','vote_cool','hot','more',
        'profile','cute','list','note','plain', 'cool','funny','writer','photos',
        'latitude','longitude','Breakfast & Brunch', 'American (Traditional)', 'Burgers', 'Fast Food',
        'American (New)', 'Chinese', 'Pizza', 'Italian', 'Sandwiches', 'Sushi Bars', 'Japanese',
        'Indian', 'Mexican', 'Vietnamese', 'Thai', 'Asian Fusion','Take-out', 'Wi-Fi',
        'dessert','latenight','lunch','dinner','breakfast','brunch', 'Caters',
        'Noise Level','Takes Reservations', 'Delivery', 'romantic','intimate','touristy',
        'hipster','divey','classy','trendy','upscale','casual',
        'Parking', 'Has TV','Outdoor Seating', 'Attire', 'Alcohol', 'Waiter Service',
        'Accepts Credit Cards', 'Good for Kids', 'Good For Groups',
        'Price Range', 'Wheelchair Accessible']

        trainset = pn.read_csv(path+'train.csv', delimiter = '::',
        names = yelp_columns, engine= 'python')

        testset = pn.read_csv(path+'test.csv', delimiter = '::',
        names = yelp_columns, engine= 'python')
    else :
        movielens_columns = ['user_id', 'item_id', 'rating','sexe', 'age','occupation', 'Action' ,
        'Adventure', 'Animation', 'Children\'s', 'Comedy', 'Crime', 'Documentary', 'Drama', 'Fantasy', 'Film-Noir',
        'Horror', 'Musical', 'Mystery', 'Romance', 'Sci-Fi', 'Thriller', 'War', 'Western']

        trainset = pn.read_csv(path+'train.csv', delimiter = '::',
        names = movielens_columns, engine = 'python')

        testset = pn.read_csv(path+'test.csv', delimiter = '::',
        names = movielens_columns, engine = 'python')

    num_users = int(max(testset.user_id.max()+1,trainset.user_id.max()+1))
    num_items = int(max(testset.item_id.max()+1,trainset.item_id.max()+1))

    return trainset, testset, num_users, num_items

def evaluate_model(model, model_name, test, k, dataset):
    if dataset == 3 :
        #Yelp
        test_userID = test['user_id']
        test_userDATA = test[['fans','average_stars','friends','vote_funny','useful','vote_cool','hot','more',
        'profile','cute','list','note','plain', 'cool','funny','writer','photos']]

        test_itemID = test['item_id']
        test_itemDATA = test.iloc[:,-50:]
    else :
        #MovieLens
        test_userID = test['user_id']
        test_userDATA = test[['sexe', 'age','occupation']]

        test_itemID = test['item_id']
        test_itemDATA = test.iloc[:,-18:]

    if(model_name in ['GMF','NCF','MLP','pretrained NCF']):
        predictions = model.predict([test_userID, test_itemID], verbose=2)
    else:
        predictions = model.predict([test_userID, test_userDATA, test_itemID, test_itemDATA], verbose=2)

    predictions = pn.DataFrame(data=predictions,columns=['predicted'])
    predictions = pn.concat([test,predictions], axis=1)
    predictions = predictions.sort_values(by=['predicted'], ascending=False)

    users=test.user_id.unique()

    hrs, ndcgs = 0,0
    for u in users:
        p= predictions[predictions['user_id']==u].loc[:,'rating'].head(k)
        hr, ndcg = evaluate_by_user(p)
        hrs = hrs + hr
        ndcgs = ndcgs + ndcg

    mean_hr = hrs/len(users)
    mean_ndcg = ndcgs/len(users)

    return mean_hr, mean_ndcg

def evaluate_by_user(test):
    hr, ndcg = 0, 0
    i=1
    for rating in test:
        if rating == 1 :
            hr = 1
            ndcg = math.log(2)/math.log(i+1)
            break
        i=i+1
    return hr, ndcg

#create model
def build_NHF_model(emb_size_hmlp, emb_size_gmf, predictive_factors, num_layers, dataset):
    input_userID = Input(shape = [1], name = 'user_ID')
    input_itemID = Input(shape = [1], name = 'item_ID')

    if dataset == 3 :
        #Yelp
        input_userDATA = Input(shape = [17], name = 'user_data')
        input_itemDATA = Input(shape = [50], name = 'item_data')
    else :
        #MovieLens
        input_userDATA = Input(shape = [3], name = 'user_data')
        input_itemDATA = Input(shape = [18], name = 'item_data')

    #GMF part
    user_latent_factors_GMF = emb_size_gmf
    item_latent_factors_GMF = emb_size_gmf

    user_emb_GMF = Embedding(num_users, user_latent_factors_GMF, name = 'user_emb_GMF')(input_userID)
    item_emb_GMF = Embedding(num_items, item_latent_factors_GMF, name = 'item_emb_GMF')(input_itemID)

    flat_u_GMF = Flatten()(user_emb_GMF)
    flat_i_GMF = Flatten()(item_emb_GMF)

    mul_layer = Multiply()([flat_u_GMF, flat_i_GMF])

    #HybMLP part
    user_latent_factors_hMLP = emb_size_hmlp
    item_latent_factors_hMLP = emb_size_hmlp

    user_emb_hMLP = Embedding(num_users, user_latent_factors_hMLP, name= 'user_emb_hMLP')(input_userID)
    item_emb_hMLP = Embedding(num_items, item_latent_factors_hMLP, name= 'item_emb_hMLP')(input_itemID)

    flat_u_hMLP = Flatten()(user_emb_hMLP)
    flat_i_hMLP = Flatten()(item_emb_hMLP)

    concat_hMLP = concatenate([flat_u_hMLP, flat_i_hMLP, input_userDATA, input_itemDATA])
    layer = concat_hMLP
    for l in range(num_layers,0,-1):
        layer = Dense(predictive_factors*(2**(l-1)), activation='relu', name= 'layer%d' %(num_layers-l+1))(layer)

    #NeuHybMF part
    concat_NeuhMF = concatenate([mul_layer,layer])

    out = Dense(1, activation='sigmoid', name='output')(concat_NeuhMF)

    model = Model([input_userID, input_userDATA, input_itemID, input_itemDATA], out)

    return model

#create GMF model
def build_GMF_model(emb_size):
    input_userID = Input(shape = [1], name = 'user_ID')
    input_itemID = Input(shape = [1], name = 'item_ID')

    user_emb_GMF = Embedding(num_users, emb_size, name = 'user_emb_GMF')(input_userID)
    item_emb_GMF = Embedding(num_items, emb_size, name = 'item_emb_GMF')(input_itemID)

    flat_u_GMF = Flatten()(user_emb_GMF)
    flat_i_GMF = Flatten()(item_emb_GMF)

    mul_layer = Multiply()([flat_u_GMF, flat_i_GMF])

    out = Dense(1, activation='sigmoid', name='output')(mul_layer)

    GMF_model = Model([input_userID, input_itemID], out)

    return GMF_model

#create HybMLP model
def build_hMLP_model(emb_size, predictive_factors, num_layers, dataset):
    input_userID = Input(shape = [1], name = 'user_ID')
    input_itemID = Input(shape = [1], name = 'item_ID')

    if dataset == 3 :
        #Yelp
        input_userDATA = Input(shape = [17], name = 'user_data')
        input_itemDATA = Input(shape = [50], name = 'item_data')
    else :
        #MovieLens
        input_userDATA = Input(shape = [3], name = 'user_data')
        input_itemDATA = Input(shape = [18], name = 'item_data')

    user_emb_hMLP = Embedding(num_users, emb_size, name= 'user_emb_hMLP')(input_userID)
    item_emb_hMLP = Embedding(num_items, emb_size, name= 'item_emb_hMLP')(input_itemID)

    flat_u_hMLP = Flatten()(user_emb_hMLP)
    flat_i_hMLP = Flatten()(item_emb_hMLP)

    concat_hMLP = concatenate([flat_u_hMLP, flat_i_hMLP, input_userDATA, input_itemDATA])
    layer = concat_hMLP
    for l in range(num_layers,0,-1):
        layer = Dense(predictive_factors*(2**(l-1)), activation='relu', name= 'layer%d' %(num_layers-l+1))(layer)

    out = Dense(1, activation='sigmoid', name='output')(layer)

    hMLP_model = Model([input_userID, input_userDATA, input_itemID, input_itemDATA], out)

    return hMLP_model

#create NCF model
def build_NCF_model(emb_size_mlp, emb_size_gmf, predictive_factors, num_layers):
    input_userID = Input(shape = [1], name = 'user_ID')
    input_itemID = Input(shape = [1], name = 'item_ID')

    #GMF part
    user_latent_factors_GMF = emb_size_gmf
    item_latent_factors_GMF = emb_size_gmf

    user_emb_GMF = Embedding(num_users, user_latent_factors_GMF, name = 'user_emb_GMF')(input_userID)
    item_emb_GMF = Embedding(num_items, item_latent_factors_GMF, name = 'item_emb_GMF')(input_itemID)

    flat_u_GMF = Flatten()(user_emb_GMF)
    flat_i_GMF = Flatten()(item_emb_GMF)

    mul_layer = Multiply()([flat_u_GMF, flat_i_GMF])

    #MLP part
    user_latent_factors_MLP = emb_size_mlp
    item_latent_factors_MLP = emb_size_mlp

    user_emb_MLP = Embedding(num_users, user_latent_factors_MLP, name= 'user_emb_MLP')(input_userID)
    item_emb_MLP = Embedding(num_items, item_latent_factors_MLP, name= 'item_emb_MLP')(input_itemID)

    flat_u_MLP = Flatten()(user_emb_MLP)
    flat_i_MLP = Flatten()(item_emb_MLP)

    concat_MLP = concatenate([flat_u_MLP, flat_i_MLP])
    layer = concat_MLP
    for l in range(num_layers,0,-1):
        layer = Dense(predictive_factors*(2**(l-1)), activation='relu', name= 'layer%d' %(num_layers-l+1))(layer)

    #NeuMF part
    concat_NeuhMF = concatenate([mul_layer,layer])

    out = Dense(1, activation='sigmoid', name='output')(concat_NeuhMF)

    model = Model([input_userID, input_itemID], out)

    return model

#create MLP model
def build_MLP_model(emb_size,predictive_factors,num_layers):
    input_userID = Input(shape = [1], name = 'user_ID')
    input_itemID = Input(shape = [1], name = 'item_ID')

    user_latent_factors_MLP = emb_size
    item_latent_factors_MLP = emb_size

    user_emb_MLP = Embedding(num_users, user_latent_factors_MLP, name= 'user_emb_MLP')(input_userID)
    item_emb_MLP = Embedding(num_items, item_latent_factors_MLP, name= 'item_emb_MLP')(input_itemID)

    flat_u_MLP = Flatten()(user_emb_MLP)
    flat_i_MLP = Flatten()(item_emb_MLP)

    concat_MLP = concatenate([flat_u_MLP, flat_i_MLP])
    layer = concat_MLP
    for l in range(num_layers,0,-1):
        layer = Dense(predictive_factors*(2**(l-1)), activation='relu', name= 'layer%d' %(num_layers-l+1))(layer)

    out = Dense(1, activation='sigmoid', name='output')(layer)

    MLP_model = Model([input_userID, input_itemID], out)

    return MLP_model

def load_pretrained_model(model, model_name, gmf_model, hmlp_model, num_layers):
    # MF embeddings
    user_emb_GMF = gmf_model.get_layer('user_emb_GMF').get_weights()
    item_emb_GMF = gmf_model.get_layer('item_emb_GMF').get_weights()
    model.get_layer('user_emb_GMF').set_weights(user_emb_GMF)
    model.get_layer('item_emb_GMF').set_weights(item_emb_GMF)
    if(model_name=='NHF'):
        # HybMLP embeddings
        user_emb_MLP = hmlp_model.get_layer('user_emb_hMLP').get_weights()
        item_emb_MLP = hmlp_model.get_layer('item_emb_hMLP').get_weights()
        model.get_layer('user_emb_hMLP').set_weights(user_emb_MLP)
        model.get_layer('item_emb_hMLP').set_weights(item_emb_MLP)
    elif(model_name=='NCF'):
        # MLP embeddings
        user_emb_MLP = hmlp_model.get_layer('user_emb_MLP').get_weights()
        item_emb_MLP = hmlp_model.get_layer('item_emb_MLP').get_weights()
        model.get_layer('user_emb_MLP').set_weights(user_emb_MLP)
        model.get_layer('item_emb_MLP').set_weights(item_emb_MLP)
    # HybMLP/MLP layers
    for i in range(num_layers):
        mlp_layer_weights = hmlp_model.get_layer('layer%d' %(i+1)).get_weights()
        model.get_layer('layer%d' %(i+1)).set_weights(mlp_layer_weights)

    # Prediction weights
    gmf_prediction = gmf_model.get_layer('output').get_weights()
    mlp_prediction = hmlp_model.get_layer('output').get_weights()
    new_weights = np.concatenate((gmf_prediction[0], mlp_prediction[0]), axis=0)
    new_b = gmf_prediction[1] + mlp_prediction[1]
    model.get_layer('output').set_weights([0.5*new_weights, 0.5*new_b])
    return model

def train(model, model_name, train, test, num_epochs, path, batch, dataset):
    if dataset == 3 :
        #Yelp
        train_userID = train['user_id']
        train_userDATA = train[['fans','average_stars','friends','vote_funny','useful','vote_cool','hot','more',
        'profile','cute','list','note','plain','cool','funny','writer','photos']]

        train_itemID = train['item_id']
        train_itemDATA = train.iloc[:,-50:]

        train_y = train['rating']
    else :
        #MovieLens
        train_userID = train['user_id']
        train_userDATA = train[['sexe', 'age','occupation']]

        train_itemID = train['item_id']
        train_itemDATA = train.iloc[:,-18:]

        train_y = train['rating']

    print('\nTRAINING '+model_name+'...\n')

    #intitialisation
    best_hr, best_ndcg = evaluate_model(model,model_name,test, 10, dataset)
    best_iteration = 0
    all_hrs, all_ndcgs = {}, {}
    all_hrs[0], all_ndcgs[0] = best_hr, best_ndcg
    model.save_weights(path+' '+model_name+'.h5', overwrite = True)
    bad_epochs=0

    #for every epoch : train and test
    for epoch in range(1, num_epochs+1) :
        t1 = time()

        if(model_name in ['GMF','NCF','MLP','pretrained NCF']):
            history = model.fit([train_userID, train_itemID], train_y, batch_size = batch , epochs = 1, verbose = 0, shuffle = True)
        else:
            history = model.fit([train_userID, train_userDATA, train_itemID, train_itemDATA], train_y, batch_size = batch , epochs = 1, verbose = 0, shuffle = True)

        hr, ndcg = evaluate_model(model, model_name, test, 10, dataset)
        all_hrs[epoch], all_ndcgs[epoch] = hr, ndcg

        if ndcg > best_ndcg :
            best_hr, best_ndcg, best_iteration = hr, ndcg, epoch
            model.save_weights(path+' '+model_name+'.h5', overwrite = True)
            bad_epochs = 0
        else:
            bad_epochs=bad_epochs+1
            if(bad_epochs==early_stopping):
                break

        t2 = time()
        print('Iteration %d [%.2f s]: loss = %.4f, HR = %.4f, NDCG = %.4f' % (epoch, t2-t1, history.history['loss'][0], hr, ndcg))

    print("Best iteration %d, best HR = %.4f, best NDCG = %.4f" % (best_iteration, best_hr, best_ndcg))
    save_obj(all_hrs,path+' '+model_name+' HRs')
    save_obj(all_ndcgs,path+' '+model_name+' NDCGs')
    return all_hrs, all_ndcgs

if __name__ == "__main__" :

    args = parse_args()
    path_save = args.path_save
    dataset = args.dataset
    model_to_train = int(args.model)
    pretrain = int(args.pretrain)
    epochs = args.epochs
    optimizer_nhf = args.optimizer_nhf.lower()
    optimizer_subs = args.optimizer_subs.lower()
    optimizer_ncf = args.optimizer_ncf.lower()
    emb_size_gmf = int(args.emb_size_gmf)
    emb_size_hmlp = int(args.emb_size_hmlp)
    emb_size_mlp = args.emb_size_mlp #check cast
    pred_facts = int(args.num_factors)
    num_neg = int(args.num_neg)
    num_layers = int(args.num_layers)
    early_stopping = int(args.estop)
    path_data = args.path_data
    batch = int(args.batch_size)

    '''
    number of epochs (early stopping)
    '''

    #load data
    trainset, testset, num_users, num_items = load_data(path_data, dataset)

    #select model to train
    if model_to_train == 1 :
        #train GMF
        path = path_save+'EMB %d'%(emb_size_gmf)
        GMF = build_GMF_model(emb_size_gmf)
        GMF.compile(optimizer = optimizer_subs, loss = 'binary_crossentropy')
        hrs,ndcgs = train(GMF, 'GMF', trainset, testset, epochs, path, batch, dataset)

    elif model_to_train == 2 :
        #train HybMLP
        path = path_save+'HybMLP %d EMB %d PF %d'%(num_layers, emb_size_hmlp, pred_facts)
        hMLP = build_hMLP_model(emb_size_hmlp, pred_facts, num_layers, dataset)
        hMLP.compile(optimizer = optimizer_subs, loss = 'binary_crossentropy')
        hrs,ndcgs = train(hMLP, 'HybMLP', trainset, testset, epochs, path, batch, dataset)

    elif  model_to_train == 3 :
        #train NeuHybMF
        if pretrain == 1 : #with pretrain models
            #pretrain GMF
            path = path_save+'EMB %d'%(emb_size_gmf)
            GMF = build_GMF_model(emb_size_gmf)
            GMF.compile(optimizer = optimizer_subs, loss = 'binary_crossentropy')
            train(GMF, 'GMF', trainset, testset, epochs, path, batch, dataset)

            #pretrain HybMLP

            if num_layers!=0:
                path = path_save+'HybMLP %d EMB %d PF %d'%(num_layers, emb_size_hmlp, pred_facts)
            else:
                path = path_save+'HybMLP %d EMB %d'%(num_layers, emb_size_hmlp)

            hMLP = build_hMLP_model(emb_size_hmlp, pred_facts, num_layers, dataset)
            hMLP.compile(optimizer = optimizer_subs, loss = 'binary_crossentropy')
            train(hMLP, 'HybMLP', trainset, testset, epochs, path, batch, dataset)

            #load pretrained models and train NeuHybMF (NHF)
            NHF = build_NHF_model(emb_size_hmlp, emb_size_gmf, pred_facts, num_layers, dataset)
            NHF.compile(optimizer = optimizer_nhf, loss = 'binary_crossentropy')
            NHF = load_pretrained_model(NHF,'NHF', GMF, hMLP, num_layers)
            hrs,ndcgs = train(NHF, 'pretrained NHF', trainset, testset, epochs, path, batch, dataset)

        else : #without pretrain models
            NHF = build_NHF_model(emb_size_hmlp, emb_size_gmf, pred_facts, num_layers, dataset)
            NHF.compile(optimizer = optimizer_nhf, loss = 'binary_crossentropy')
            hrs,ndcgs = train(NHF,'NHF', trainset, testset, epochs, path_save, batch, dataset)

    elif model_to_train == 4 :
        #train MLP
        if num_layers!=0:
            path = path_save+'HybMLP %d EMB %d PF %d'%(num_layers, emb_size_hmlp, pred_facts)
        else:
            path = path_save+'HybMLP %d EMB %d'%(num_layers, emb_size_hmlp)

        MLP = build_MLP_model(emb_size_mlp, pred_facts, num_layers)
        MLP.compile(optimizer = optimizer_subs, loss = 'binary_crossentropy')
        hrs,ndcgs = train(MLP, 'MLP', trainset, testset, epochs, path, batch, dataset)

    elif  model_to_train == 5 :
        #train NCF
        if pretrain == 1 : #with pretrain models
            #pretrain GMF
            path = path_save+'EMB %d'%(emb_size_gmf)
            GMF = build_GMF_model(emb_size_gmf)
            GMF.compile(optimizer = optimizer_subs, loss = 'binary_crossentropy')
            train(GMF, 'GMF', trainset, testset, epochs, path, batch, dataset)

            #train MLP
            if num_layers!=0:
                path = path_save+'HybMLP %d EMB %d PF %d'%(num_layers, emb_size_hmlp, pred_facts)
            else:
                path = path_save+'HybMLP %d EMB %d'%(num_layers, emb_size_hmlp)

            MLP = build_MLP_model(emb_size_mlp, pred_facts, num_layers)
            MLP.compile(optimizer = optimizer_subs, loss = 'binary_crossentropy')
            train(MLP, 'MLP', trainset, testset, epochs, path, batch, dataset)

            #load pretrained models and train NCF
            NCF = build_NCF_model(emb_size_mlp, emb_size_gmf, pred_facts, num_layers)
            NCF.compile(optimizer = optimizer_ncf, loss = 'binary_crossentropy')
            NCF = load_pretrained_model(NCF,'NCF', GMF, MLP, num_layers)
            hrs,ndcgs = train(NCF, 'pretrained NCF', trainset, testset, epochs, path, batch, dataset)

        else : #without pretrain models
            NCF = build_NCF_model(emb_size_mlp, emb_size_gmf, pred_facts, num_layers)
            NCF.compile(optimizer = optimizer_ncf, loss = 'binary_crossentropy')
            hrs,ndcgs = train(NCF,'NCF', trainset, testset, epochs, path_save, batch, dataset)

    #PLOT

    x1, y1 = zip(*(hrs.items()))
    x2, y2 = zip(*(ndcgs.items()))
    plt.plot(x1,y1,'ro',linestyle = 'dashed')

    plt.subplot(211)
    plt.plot(x1,y1,'ro',linestyle = 'dashed')
    plt.title('The results of the training')
    plt.ylabel('HR@10')
    plt.subplot(212)
    plt.plot(x2,y2,'bd',linestyle = 'dashed')
    plt.ylabel('NDCG@10')
    plt.xlabel('Epochs')
    plt.show()


#Scripts\activate เพื่อเข้าไปใน environment
#ก่อนใช้ ให้เข้าไปใน tutorial ก่อน โดยการ cd week03
#ใช้ scrapy crawl quotes -o LMAOXD.csv โดย LMAOXD คือชื่อไฟล์ และ .csv คือนามสกุลที่อยากได้ สามารถเปลี่ยนเป็น json ได้
#ใช้ scrapy shell https://store.steampowered.com/app/551170/Onmyoji/ เพื่อ test
#Tutorial Vid : https://youtu.be/ALizgnSFTwQ

import scrapy
import time
import random
from scrapy_selenium import SeleniumRequest

from selenium.webdriver.common.by import By
from selenium.webdriver.support import expected_conditions as EC

class QuotesSpider(scrapy.Spider):
    name = 'quotes'

    def start_requests(self):
        url = 'https://store.steampowered.com/app/2089350/Hatsune_Miku_Logic_Paint_S/'

        yield SeleniumRequest(
            url=url,
            callback=self.parse,
            wait_time=1
            )

    def parse(self, response):
        for quote in response.css('div.page_content_ctn'):
            yield {
                'url' : response.request.url,
                'title': quote.css('.apphub_AppName::text').get(),
                'tag': quote.css('.app_tag::text').re('Hero Shooter|Battle Royale|Competitive|Roguelite|JRPG|Story Rich|Automobile Sim|Sports|Racing|Driving|Immersive Sim|Tactical|Management|Third Person|Simulation|Co-op|Local Multiplayer|Online Co-Op|PvE|Football|Soccer|PvP|Controller|Team-Based|Realistic|Local Co-Op|Multiplayer|Singleplayer|Anime|Strategy|Card Battler|Casual|PvP|Turn-Based Tactics|Card Game|Deckbuilding|Comic Book|Free to Play|Early Access|Superhero|Stylized|RPG|Free to Play|Turn-Based|Adventure|Fantasy|Stealth|Horror|Action|Parody|Retro|Funny|FPS|Shooter|Open World'),
                #([A-Z])\w+')
                # 'price' : quote.css('.game_purchase_price.price::text').get(),
                # 'price02' : quote.css('.game_purchase_price.price::text').getall(),
                'price' : quote.css('.game_purchase_price.price::text').re('Free To Play|Free|\d+'),
                'discounted price' : quote.css('.discount_original_price::text').get()
            }

        next_page = response.css('div.store_horizontal_autoslider_ctn a::attr(href)').get()
        next_pages = response.css('div.block_responsive_horizontal_scroll a::attr(href)').getall()
        for next_page in next_pages:
            if next_page is not None:
                #time.sleep(5)
                next_page = response.urljoin(next_page)
                yield SeleniumRequest(url=next_page, callback=self.parse,wait_time=1)