# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# 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.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import sys
# add python path of PadleDetection to sys.path
parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 2)))
sys.path.insert(0, parent_path)

from ppdet.utils.logger import setup_logger
logger = setup_logger('ppdet.anchor_cluster')

from scipy.cluster.vq import kmeans
import numpy as np
from tqdm import tqdm

from ppdet.utils.cli import ArgsParser
from ppdet.utils.check import check_gpu, check_version, check_config
from ppdet.core.workspace import load_config, merge_config


class BaseAnchorCluster(object):
    def __init__(self, n, cache_path, cache, verbose=True):
        """
        Base Anchor Cluster

        Args:
            n (int): number of clusters
            cache_path (str): cache directory path
            cache (bool): whether using cache
            verbose (bool): whether print results
        """
        super(BaseAnchorCluster, self).__init__()
        self.n = n
        self.cache_path = cache_path
        self.cache = cache
        self.verbose = verbose

    def print_result(self, centers):
        raise NotImplementedError('%s.print_result is not available' %
                                  self.__class__.__name__)

    def get_whs(self):
        whs_cache_path = os.path.join(self.cache_path, 'whs.npy')
        shapes_cache_path = os.path.join(self.cache_path, 'shapes.npy')
        if self.cache and os.path.exists(whs_cache_path) and os.path.exists(
                shapes_cache_path):
            self.whs = np.load(whs_cache_path)
            self.shapes = np.load(shapes_cache_path)
            return self.whs, self.shapes
        whs = np.zeros((0, 2))
        shapes = np.zeros((0, 2))
        self.dataset.parse_dataset()
        roidbs = self.dataset.roidbs
        for rec in tqdm(roidbs):
            h, w = rec['h'], rec['w']
            bbox = rec['gt_bbox']
            wh = bbox[:, 2:4] - bbox[:, 0:2] + 1
            wh = wh / np.array([[w, h]])
            shape = np.ones_like(wh) * np.array([[w, h]])
            whs = np.vstack((whs, wh))
            shapes = np.vstack((shapes, shape))

        if self.cache:
            os.makedirs(self.cache_path, exist_ok=True)
            np.save(whs_cache_path, whs)
            np.save(shapes_cache_path, shapes)

        self.whs = whs
        self.shapes = shapes
        return self.whs, self.shapes

    def calc_anchors(self):
        raise NotImplementedError('%s.calc_anchors is not available' %
                                  self.__class__.__name__)

    def __call__(self):
        self.get_whs()
        centers = self.calc_anchors()
        if self.verbose:
            self.print_result(centers)
        return centers


class YOLOv2AnchorCluster(BaseAnchorCluster):
    def __init__(self,
                 n,
                 dataset,
                 size,
                 cache_path,
                 cache,
                 iters=1000,
                 verbose=True):
        super(YOLOv2AnchorCluster, self).__init__(
            n, cache_path, cache, verbose=verbose)
        """
        YOLOv2 Anchor Cluster

        Reference:
            https://github.com/AlexeyAB/darknet/blob/master/scripts/gen_anchors.py

        Args:
            n (int): number of clusters
            dataset (DataSet): DataSet instance, VOC or COCO
            size (list): [w, h]
            cache_path (str): cache directory path
            cache (bool): whether using cache
            iters (int): kmeans algorithm iters
            verbose (bool): whether print results
        """
        self.dataset = dataset
        self.size = size
        self.iters = iters

    def print_result(self, centers):
        logger.info('%d anchor cluster result: [w, h]' % self.n)
        for w, h in centers:
            logger.info('[%d, %d]' % (round(w), round(h)))

    def metric(self, whs, centers):
        wh1 = whs[:, None]
        wh2 = centers[None]
        inter = np.minimum(wh1, wh2).prod(2)
        return inter / (wh1.prod(2) + wh2.prod(2) - inter)

    def kmeans_expectation(self, whs, centers, assignments):
        dist = self.metric(whs, centers)
        new_assignments = dist.argmax(1)
        converged = (new_assignments == assignments).all()
        return converged, new_assignments

    def kmeans_maximizations(self, whs, centers, assignments):
        new_centers = np.zeros_like(centers)
        for i in range(centers.shape[0]):
            mask = (assignments == i)
            if mask.sum():
                new_centers[i, :] = whs[mask].mean(0)
        return new_centers

    def calc_anchors(self):
        self.whs = self.whs * np.array([self.size])
        # random select k centers
        whs, n, iters = self.whs, self.n, self.iters
        logger.info('Running kmeans for %d anchors on %d points...' %
                    (n, len(whs)))
        idx = np.random.choice(whs.shape[0], size=n, replace=False)
        centers = whs[idx]
        assignments = np.zeros(whs.shape[0:1]) * -1
        # kmeans
        if n == 1:
            return self.kmeans_maximizations(whs, centers, assignments)

        pbar = tqdm(range(iters), desc='Cluster anchors with k-means algorithm')
        for _ in pbar:
            # E step
            converged, assignments = self.kmeans_expectation(whs, centers,
                                                             assignments)
            if converged:
                logger.info('kmeans algorithm has converged')
                break
            # M step
            centers = self.kmeans_maximizations(whs, centers, assignments)
            ious = self.metric(whs, centers)
            pbar.desc = 'avg_iou: %.4f' % (ious.max(1).mean())

        centers = sorted(centers, key=lambda x: x[0] * x[1])
        return centers


class YOLOv5AnchorCluster(BaseAnchorCluster):
    def __init__(self,
                 n,
                 dataset,
                 size,
                 cache_path,
                 cache,
                 iters=300,
                 gen_iters=1000,
                 thresh=0.25,
                 verbose=True):
        super(YOLOv5AnchorCluster, self).__init__(
            n, cache_path, cache, verbose=verbose)
        """
        YOLOv5 Anchor Cluster

        Reference:
            https://github.com/ultralytics/yolov5/blob/master/utils/general.py

        Args:
            n (int): number of clusters
            dataset (DataSet): DataSet instance, VOC or COCO
            size (list): [w, h]
            cache_path (str): cache directory path
            cache (bool): whether using cache
            iters (int): iters of kmeans algorithm
            gen_iters (int): iters of genetic algorithm
            threshold (float): anchor scale threshold
            verbose (bool): whether print results
        """
        self.dataset = dataset
        self.size = size
        self.iters = iters
        self.gen_iters = gen_iters
        self.thresh = thresh

    def print_result(self, centers):
        whs = self.whs
        centers = centers[np.argsort(centers.prod(1))]
        x, best = self.metric(whs, centers)
        bpr, aat = (
            best > self.thresh).mean(), (x > self.thresh).mean() * self.n
        logger.info(
            'thresh=%.2f: %.4f best possible recall, %.2f anchors past thr' %
            (self.thresh, bpr, aat))
        logger.info(
            'n=%g, img_size=%s, metric_all=%.3f/%.3f-mean/best, past_thresh=%.3f-mean: '
            % (self.n, self.size, x.mean(), best.mean(),
               x[x > self.thresh].mean()))
        logger.info('%d anchor cluster result: [w, h]' % self.n)
        for w, h in centers:
            logger.info('[%d, %d]' % (round(w), round(h)))

    def metric(self, whs, centers):
        r = whs[:, None] / centers[None]
        x = np.minimum(r, 1. / r).min(2)
        return x, x.max(1)

    def fitness(self, whs, centers):
        _, best = self.metric(whs, centers)
        return (best * (best > self.thresh)).mean()

    def calc_anchors(self):
        self.whs = self.whs * self.shapes / self.shapes.max(
            1, keepdims=True) * np.array([self.size])
        wh0 = self.whs
        i = (wh0 < 3.0).any(1).sum()
        if i:
            logger.warning('Extremely small objects found. %d of %d'
                           'labels are < 3 pixels in width or height' %
                           (i, len(wh0)))

        wh = wh0[(wh0 >= 2.0).any(1)]
        logger.info('Running kmeans for %g anchors on %g points...' %
                    (self.n, len(wh)))
        s = wh.std(0)
        centers, dist = kmeans(wh / s, self.n, iter=self.iters)
        centers *= s

        f, sh, mp, s = self.fitness(wh, centers), centers.shape, 0.9, 0.1
        pbar = tqdm(
            range(self.gen_iters),
            desc='Evolving anchors with Genetic Algorithm')
        for _ in pbar:
            v = np.ones(sh)
            while (v == 1).all():
                v = ((np.random.random(sh) < mp) * np.random.random() *
                     np.random.randn(*sh) * s + 1).clip(0.3, 3.0)
            new_centers = (centers.copy() * v).clip(min=2.0)
            new_f = self.fitness(wh, new_centers)
            if new_f > f:
                f, centers = new_f, new_centers.copy()
                pbar.desc = 'Evolving anchors with Genetic Algorithm: fitness = %.4f' % f

        return centers


def main():
    parser = ArgsParser()
    parser.add_argument(
        '--n', '-n', default=9, type=int, help='num of clusters')
    parser.add_argument(
        '--iters',
        '-i',
        default=1000,
        type=int,
        help='num of iterations for kmeans')
    parser.add_argument(
        '--gen_iters',
        '-gi',
        default=1000,
        type=int,
        help='num of iterations for genetic algorithm')
    parser.add_argument(
        '--thresh',
        '-t',
        default=0.25,
        type=float,
        help='anchor scale threshold')
    parser.add_argument(
        '--verbose', '-v', default=True, type=bool, help='whether print result')
    parser.add_argument(
        '--size',
        '-s',
        default=None,
        type=str,
        help='image size: w,h, using comma as delimiter')
    parser.add_argument(
        '--method',
        '-m',
        default='v2',
        type=str,
        help='cluster method, [v2, v5] are supported now')
    parser.add_argument(
        '--cache_path', default='cache', type=str, help='cache path')
    parser.add_argument(
        '--cache', action='store_true', help='whether use cache')
    FLAGS = parser.parse_args()

    cfg = load_config(FLAGS.config)
    merge_config(FLAGS.opt)
    check_config(cfg)
    # check if set use_gpu=True in paddlepaddle cpu version
    check_gpu(cfg.use_gpu)
    # check if paddlepaddle version is satisfied
    check_version()

    # get dataset
    dataset = cfg['TrainDataset']
    if FLAGS.size:
        if ',' in FLAGS.size:
            size = list(map(int, FLAGS.size.split(',')))
            assert len(size) == 2, "the format of size is incorrect"
        else:
            size = int(FLAGS.size)
            size = [size, size]
    elif 'inputs_def' in cfg['TrainReader'] and 'image_shape' in cfg[
            'TrainReader']['inputs_def']:
        size = cfg['TrainReader']['inputs_def']['image_shape'][1:]
    else:
        raise ValueError('size is not specified')

    if FLAGS.method == 'v2':
        cluster = YOLOv2AnchorCluster(FLAGS.n, dataset, size, FLAGS.cache_path,
                                      FLAGS.cache, FLAGS.iters, FLAGS.verbose)
    elif FLAGS.method == 'v5':
        cluster = YOLOv5AnchorCluster(FLAGS.n, dataset, size, FLAGS.cache_path,
                                      FLAGS.cache, FLAGS.iters, FLAGS.gen_iters,
                                      FLAGS.thresh, FLAGS.verbose)
    else:
        raise ValueError('cluster method: %s is not supported' % FLAGS.method)

    anchors = cluster()


if __name__ == "__main__":
    main()