#  Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
#
#Licensed under the Apache License, Version 2.0 (the "License");
#you may not use this file except in compliance with the License.
#You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
#Unless required by applicable law or agreed to in writing, software
#distributed under the License is distributed on an "AS IS" BASIS,
#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#See the License for the specific language governing permissions and
#limitations under the License.

import os
import platform
import random
import numpy as np
import pandas as pd
import multiprocessing
import json
import logging
import functools
import paddle
logger = logging.getLogger(__name__)

from .reader_utils import DataReader
from models.bsn.bsn_utils import iou_with_anchors, ioa_with_anchors


class BSNVideoReader(DataReader):
    """
    Data reader for BsnTem model, which was stored as features extracted by prior networks
    dataset cfg: anno_file, annotation file path,
                 feat_path, feature path,
                 file_list, file list for infer,
                 tscale, temporal length of input,
                 anchor_xmin, anchor_xmax, the range of each point in the feature sequence,
                 batch_size, batch size of input data,
                 num_threads, number of threads of data processing

    """

    def __init__(self, name, mode, cfg):
        self.name = name
        self.mode = mode
        self.tscale = cfg.MODEL.tscale  # 100
        self.gt_boundary_ratio = cfg.MODEL.gt_boundary_ratio
        self.anno_file = cfg.MODEL.anno_file
        self.file_list = cfg.INFER.filelist
        self.subset = cfg[mode.upper()]['subset']
        self.tgap = 1. / self.tscale
        self.feat_path = cfg.MODEL.feat_path
        self.anchor_xmin = [self.tgap * i for i in range(self.tscale)]
        self.anchor_xmax = [self.tgap * i for i in range(1, self.tscale + 1)]
        self.get_dataset_dict()

        self.batch_size = cfg[mode.upper()]['batch_size']
        self.num_threads = cfg[mode.upper()]['num_threads']
        if (mode == 'test') or (mode == 'infer'):
            self.num_threads = 1  # set num_threads as 1 for test and infer

    def get_dataset_dict(self):
        self.video_dict = {}
        if self.mode == "infer":
            annos = json.load(open(self.file_list))
            for video_name in annos.keys():
                self.video_dict[video_name] = annos[video_name]
        else:
            annos = json.load(open(self.anno_file))
            for video_name in annos.keys():
                video_subset = annos[video_name]["subset"]
                if self.subset == "train_val":
                    if "train" in video_subset or "validation" in video_subset:
                        self.video_dict[video_name] = annos[video_name]
                else:
                    if self.subset in video_subset:
                        self.video_dict[video_name] = annos[video_name]
        self.video_list = list(self.video_dict.keys())
        self.video_list.sort()
        print("%s subset video numbers: %d" %
              (self.subset, len(self.video_list)))

    def get_video_label(self, video_name):
        video_info = self.video_dict[video_name]
        video_second = video_info['duration_second']
        video_labels = video_info['annotations']

        gt_bbox = []
        for gt in video_labels:
            tmp_start = max(min(1, gt["segment"][0] / video_second), 0)
            tmp_end = max(min(1, gt["segment"][1] / video_second), 0)
            gt_bbox.append([tmp_start, tmp_end])

        gt_bbox = np.array(gt_bbox)
        gt_xmins = gt_bbox[:, 0]
        gt_xmaxs = gt_bbox[:, 1]
        gt_lens = gt_xmaxs - gt_xmins
        gt_len_small = np.maximum(self.tgap, self.gt_boundary_ratio * gt_lens)
        gt_start_bboxs = np.stack(
            (gt_xmins - gt_len_small / 2, gt_xmins + gt_len_small / 2), axis=1)
        gt_end_bboxs = np.stack(
            (gt_xmaxs - gt_len_small / 2, gt_xmaxs + gt_len_small / 2), axis=1)

        match_score_action = []
        for jdx in range(len(self.anchor_xmin)):
            match_score_action.append(
                np.max(
                    ioa_with_anchors(self.anchor_xmin[jdx], self.anchor_xmax[
                        jdx], gt_xmins, gt_xmaxs)))
        match_score_start = []
        for jdx in range(len(self.anchor_xmin)):
            match_score_start.append(
                np.max(
                    ioa_with_anchors(self.anchor_xmin[jdx], self.anchor_xmax[
                        jdx], gt_start_bboxs[:, 0], gt_start_bboxs[:, 1])))
        match_score_end = []
        for jdx in range(len(self.anchor_xmin)):
            match_score_end.append(
                np.max(
                    ioa_with_anchors(self.anchor_xmin[jdx], self.anchor_xmax[
                        jdx], gt_end_bboxs[:, 0], gt_end_bboxs[:, 1])))

        gt_start = np.array(match_score_start)
        gt_end = np.array(match_score_end)
        gt_action = np.array(match_score_action)
        return gt_start, gt_end, gt_action

    def load_file(self, video_name):
        video_feat = np.load(self.feat_path + "/" + video_name + ".npy")
        video_feat = video_feat.T
        video_feat = video_feat.astype("float32")
        return video_feat

    def create_reader(self):
        """reader creator for ctcn model"""
        if self.mode == 'infer':
            return self.make_infer_reader()
        if self.num_threads == 1:
            return self.make_reader()
        else:
            sysstr = platform.system()
            if sysstr == 'Windows':
                return self.make_multithread_reader()
            else:
                return self.make_multiprocess_reader()

    def make_infer_reader(self):
        """reader for inference"""

        def reader():
            batch_out = []
            for video_name in self.video_list:
                video_idx = self.video_list.index(video_name)
                video_feat = self.load_file(video_name)
                batch_out.append((video_feat, video_idx))

                if len(batch_out) == self.batch_size:
                    yield batch_out
                    batch_out = []

        return reader

    def make_reader(self):
        """single process reader"""

        def reader():
            video_list = self.video_list
            if self.mode == 'train':
                random.shuffle(video_list)

            batch_out = []
            for video_name in video_list:
                video_idx = video_list.index(video_name)
                video_feat = self.load_file(video_name)
                gt_start, gt_end, gt_action = self.get_video_label(video_name)

                if self.mode == 'train' or self.mode == 'valid':
                    batch_out.append((video_feat, gt_start, gt_end, gt_action))
                elif self.mode == 'test':
                    batch_out.append(
                        (video_feat, gt_start, gt_end, gt_action, video_idx))
                else:
                    raise NotImplementedError('mode {} not implemented'.format(
                        self.mode))
                if len(batch_out) == self.batch_size:
                    yield batch_out
                    batch_out = []

        return reader

    def make_multithread_reader(self):
        def reader():
            if self.mode == 'train':
                random.shuffle(self.video_list)

            for video_name in self.video_list:
                video_idx = self.video_list.index(video_name)
                yield [video_name, video_idx]

        def process_data(sample, mode):
            video_name = sample[0]
            video_idx = sample[1]
            video_feat = self.load_file(video_name)
            gt_start, gt_end, gt_action = self.get_video_label(video_name)
            if mode == 'train' or mode == 'valid':
                return (video_feat, gt_start, gt_end, gt_action)
            elif mode == 'test':
                return (video_feat, gt_start, gt_end, gt_action, video_idx)
            else:
                raise NotImplementedError('mode {} not implemented'.format(
                    self.mode))

        mapper = functools.partial(process_data, mode=self.mode)

        def batch_reader():
            xreader = paddle.reader.xmap_readers(mapper, reader,
                                                 self.num_threads, 1024)
            batch = []
            for item in xreader():
                batch.append(item)
                if len(batch) == self.batch_size:
                    yield batch
                    batch = []

        return batch_reader

    def make_multiprocess_reader(self):
        """multiprocess reader"""

        def read_into_queue(video_list, queue):

            batch_out = []
            for video_name in video_list:
                video_idx = video_list.index(video_name)
                video_feat = self.load_file(video_name)
                gt_start, gt_end, gt_action = self.get_video_label(video_name)

                if self.mode == 'train' or self.mode == 'valid':
                    batch_out.append((video_feat, gt_start, gt_end, gt_action))
                elif self.mode == 'test':
                    batch_out.append(
                        (video_feat, gt_start, gt_end, gt_action, video_idx))
                else:
                    raise NotImplementedError('mode {} not implemented'.format(
                        self.mode))

                if len(batch_out) == self.batch_size:
                    queue.put(batch_out)
                    batch_out = []
            queue.put(None)

        def queue_reader():
            video_list = self.video_list
            if self.mode == 'train':
                random.shuffle(video_list)

            n = self.num_threads
            queue_size = 20
            reader_lists = [None] * n
            file_num = int(len(video_list) // n)
            for i in range(n):
                if i < len(reader_lists) - 1:
                    tmp_list = video_list[i * file_num:(i + 1) * file_num]
                else:
                    tmp_list = video_list[i * file_num:]
                reader_lists[i] = tmp_list

            queue = multiprocessing.Queue(queue_size)
            p_list = [None] * len(reader_lists)
            # for reader_list in reader_lists:
            for i in range(len(reader_lists)):
                reader_list = reader_lists[i]
                p_list[i] = multiprocessing.Process(
                    target=read_into_queue, args=(reader_list, queue))
                p_list[i].start()
            reader_num = len(reader_lists)
            finish_num = 0
            while finish_num < reader_num:
                sample = queue.get()
                if sample is None:
                    finish_num += 1
                else:
                    yield sample
            for i in range(len(p_list)):
                if p_list[i].is_alive():
                    p_list[i].join()

        return queue_reader


class BSNProposalReader(DataReader):
    """
    Data reader for BsnPem model, which was stored as features extracted by prior networks
    dataset cfg: anno_file, annotation file path,
                 file_list, file list for infer,
                 top_K, number of proposals during training/test,
                 feat_path, feature path generated by PGM,
                 prop_path, proposal path generated by PGM,
                 batch_size, batch size of input data,
                 num_threads, number of threads of data processing.
    """

    def __init__(self, name, mode, cfg):
        self.name = name
        self.mode = mode

        self.top_K = cfg[mode.upper()]['top_K']
        self.anno_file = cfg.MODEL.anno_file
        self.file_list = cfg.INFER.filelist
        self.subset = cfg[mode.upper()]['subset']

        if mode == 'infer':
            self.feat_path = cfg[mode.upper()]['feat_path']
            self.prop_path = cfg[mode.upper()]['prop_path']
        else:
            self.feat_path = cfg.MODEL.feat_path
            self.prop_path = cfg.MODEL.prop_path
        self.get_dataset_dict()

        self.batch_size = cfg[mode.upper()]['batch_size']
        self.num_threads = cfg[mode.upper()]['num_threads']
        if (mode == 'test') or (mode == 'infer'):
            self.num_threads = 1  # set num_threads as 1 for test and infer

    def get_dataset_dict(self):
        self.video_dict = {}
        if self.mode == "infer":
            annos = json.load(open(self.file_list))
            for video_name in annos.keys():
                self.video_dict[video_name] = annos[video_name]
        else:
            annos = json.load(open(self.anno_file))
            for video_name in annos.keys():
                video_subset = annos[video_name]["subset"]
                if self.subset in video_subset:
                    self.video_dict[video_name] = annos[video_name]
        self.video_list = list(self.video_dict.keys())
        self.video_list.sort()
        print("%s subset video numbers: %d" %
              (self.subset, len(self.video_list)))

    def get_props(self, video_name):
        pdf = pd.read_csv(self.prop_path + video_name + ".csv")
        pdf = pdf[:self.top_K]
        props_start = pdf.xmin.values[:]
        props_end = pdf.xmax.values[:]
        props_start_score = pdf.xmin_score.values[:]
        props_end_score = pdf.xmax_score.values[:]
        props_info = np.hstack([
            props_start[:, np.newaxis], props_end[:, np.newaxis],
            props_start_score[:, np.newaxis], props_end_score[:, np.newaxis]
        ])
        if self.mode == "infer":
            return props_info
        else:
            props_iou = pdf.match_iou.values[:]
            return props_iou, props_info

    def load_file(self, video_name):
        video_feat = np.load(self.feat_path + video_name + ".npy")
        video_feat = video_feat[:self.top_K, :]
        video_feat = video_feat.astype("float32")
        return video_feat

    def create_reader(self):
        """reader creator for ctcn model"""
        if self.mode == 'infer':
            return self.make_infer_reader()
        if self.num_threads == 1:
            return self.make_reader()
        else:
            sysstr = platform.system()
            if sysstr == 'Windows':
                return self.make_multithread_reader()
            else:
                return self.make_multiprocess_reader()

    def make_infer_reader(self):
        """reader for inference"""

        def reader():
            batch_out = []
            for video_name in self.video_list:
                video_idx = self.video_list.index(video_name)
                props_feat = self.load_file(video_name)
                props_info = self.get_props(video_name)
                batch_out.append((props_feat, props_info, video_idx))
                if len(batch_out) == self.batch_size:
                    yield batch_out
                    batch_out = []

        return reader

    def make_reader(self):
        """single process reader"""

        def reader():
            video_list = self.video_list
            if self.mode == 'train':
                random.shuffle(video_list)

            batch_out = []
            for video_name in video_list:
                video_idx = video_list.index(video_name)
                props_feat = self.load_file(video_name)
                props_iou, props_info = self.get_props(video_name)

                if self.mode == 'train' or self.mode == 'valid':
                    batch_out.append((props_feat, props_iou))
                elif self.mode == 'test':
                    batch_out.append(
                        (props_feat, props_iou, props_info, video_idx))
                else:
                    raise NotImplementedError('mode {} not implemented'.format(
                        self.mode))
                if len(batch_out) == self.batch_size:
                    yield batch_out
                    batch_out = []

        return reader

    def make_multithread_reader(self):
        def reader():
            if self.mode == 'train':
                random.shuffle(self.video_list)
            for video_name in self.video_list:
                video_idx = self.video_list.index(video_name)
                yield [video_name, video_idx]

        def process_data(sample, mode):
            video_name = sample[0]
            video_idx = sample[1]
            props_feat = self.load_file(video_name)
            props_iou, props_info = self.get_props(video_name)
            if mode == 'train' or mode == 'valid':
                return (props_feat, props_iou)
            elif mode == 'test':
                return (props_feat, props_iou, props_info, video_idx)
            else:
                raise NotImplementedError('mode {} not implemented'.format(
                    mode))

        mapper = functools.partial(process_data, mode=self.mode)

        def batch_reader():
            xreader = paddle.reader.xmap_readers(mapper, reader,
                                                 self.num_threads, 1024)
            batch = []
            for item in xreader():
                batch.append(item)
                if len(batch) == self.batch_size:
                    yield batch
                    batch = []

        return batch_reader

    def make_multiprocess_reader(self):
        """multiprocess reader"""

        def read_into_queue(video_list, queue):

            batch_out = []
            for video_name in video_list:
                video_idx = video_list.index(video_name)
                props_feat = self.load_file(video_name)
                props_iou, props_info = self.get_props(video_name)

                if self.mode == 'train' or self.mode == 'valid':
                    batch_out.append((props_feat, props_iou))
                elif self.mode == 'test':
                    batch_out.append(
                        (props_feat, props_iou, props_info, video_idx))
                else:
                    raise NotImplementedError('mode {} not implemented'.format(
                        self.mode))

                if len(batch_out) == self.batch_size:
                    queue.put(batch_out)
                    batch_out = []
            queue.put(None)

        def queue_reader():
            video_list = self.video_list
            if self.mode == 'train':
                random.shuffle(video_list)

            n = self.num_threads
            queue_size = 20
            reader_lists = [None] * n
            file_num = int(len(video_list) // n)
            for i in range(n):
                if i < len(reader_lists) - 1:
                    tmp_list = video_list[i * file_num:(i + 1) * file_num]
                else:
                    tmp_list = video_list[i * file_num:]
                reader_lists[i] = tmp_list

            queue = multiprocessing.Queue(queue_size)
            p_list = [None] * len(reader_lists)
            # for reader_list in reader_lists:
            for i in range(len(reader_lists)):
                reader_list = reader_lists[i]
                p_list[i] = multiprocessing.Process(
                    target=read_into_queue, args=(reader_list, queue))
                p_list[i].start()
            reader_num = len(reader_lists)
            finish_num = 0
            while finish_num < reader_num:
                sample = queue.get()
                if sample is None:
                    finish_num += 1
                else:
                    yield sample
            for i in range(len(p_list)):
                if p_list[i].is_alive():
                    p_list[i].join()

        return queue_reader
