# Copyright (c) 2022 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.

import logging
from tqdm import tqdm
import numpy as np
import argparse
import paddle
from paddleslim.common import load_config as load_slim_config
from paddleslim.common import get_logger
from paddleslim.auto_compression import AutoCompression
from paddleslim.common.dataloader import get_feed_vars

import sys
sys.path.append('../../../')
from ppocr.data import build_dataloader
from ppocr.postprocess import build_post_process
from ppocr.metrics import build_metric

logger = get_logger(__name__, level=logging.INFO)


def argsparser():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        '--config_path',
        type=str,
        default=None,
        help="path of compression strategy config.",
        required=True)
    parser.add_argument(
        '--save_dir',
        type=str,
        default='output',
        help="directory to save compressed model.")
    parser.add_argument(
        '--devices',
        type=str,
        default='gpu',
        help="which device used to compress.")
    return parser


def reader_wrapper(reader, input_name):
    if isinstance(input_name, list) and len(input_name) == 1:
        input_name = input_name[0]

    def gen(): # 形成一个字典输入
        for i, batch in enumerate(reader()):
            yield {input_name: batch[0]}

    return gen


def eval_function(exe, compiled_test_program, test_feed_names, test_fetch_list):
    post_process_class = build_post_process(all_config['PostProcess'],
                                            global_config)
    eval_class = build_metric(all_config['Metric'])
    model_type = global_config['model_type']

    with tqdm(
            total=len(val_loader),
            bar_format='Evaluation stage, Run batch:|{bar}| {n_fmt}/{total_fmt}',
            ncols=80) as t:
        for batch_id, batch in enumerate(val_loader):
            images = batch[0]
            
            try:
                preds, = exe.run(compiled_test_program,
                             feed={test_feed_names[0]: images},
                             fetch_list=test_fetch_list)
            except:
                preds, _ = exe.run(compiled_test_program,
                             feed={test_feed_names[0]: images},
                             fetch_list=test_fetch_list)

            batch_numpy = []
            for item in batch:
                batch_numpy.append(np.array(item))

            if model_type == 'det':
                preds_map = {'maps': preds}
                post_result = post_process_class(preds_map, batch_numpy[1])
                eval_class(post_result, batch_numpy)
            elif model_type == 'rec':
                post_result = post_process_class(preds, batch_numpy[1])
                eval_class(post_result, batch_numpy)
            t.update()
        metric = eval_class.get_metric()
    logger.info('metric eval ***************')
    for k, v in metric.items():
        logger.info('{}:{}'.format(k, v))

    if model_type == 'det':
        return metric['hmean']
    elif model_type == 'rec':
        return metric['acc']
    return metric


def main():
    rank_id = paddle.distributed.get_rank()
    if args.devices == 'gpu':
        place = paddle.CUDAPlace(rank_id)
        paddle.set_device('gpu')
    else:
        place = paddle.CPUPlace()
        paddle.set_device('cpu')

    global all_config, global_config
    all_config = load_slim_config(args.config_path)

    if "Global" not in all_config:  
        raise KeyError(f"Key 'Global' not found in config file. \n{all_config}")
    global_config = all_config["Global"]

    gpu_num = paddle.distributed.get_world_size()

    train_dataloader = build_dataloader(all_config, 'Train', args.devices,
                                        logger)

    global val_loader
    val_loader = build_dataloader(all_config, 'Eval', args.devices, logger)

    if isinstance(all_config['TrainConfig']['learning_rate'],
                  dict) and all_config['TrainConfig']['learning_rate'][
                      'type'] == 'CosineAnnealingDecay':
        steps = len(train_dataloader) * all_config['TrainConfig']['epochs']
        all_config['TrainConfig']['learning_rate']['T_max'] = steps
        print('total training steps:', steps)

    global_config['input_name'] = get_feed_vars(
        global_config['model_dir'], global_config['model_filename'],
        global_config['params_filename'])
    
    ac = AutoCompression(
        model_dir=global_config['model_dir'],
        model_filename=global_config['model_filename'],
        params_filename=global_config['params_filename'],
        save_dir=args.save_dir,
        config=all_config,
        train_dataloader=reader_wrapper(train_dataloader,
                                        global_config['input_name']),
        eval_callback=eval_function if rank_id == 0 else None,
        eval_dataloader=reader_wrapper(val_loader, global_config['input_name']))
    ac.compress()


if __name__ == '__main__':
    paddle.enable_static()
    parser = argsparser()
    args = parser.parse_args()
    main()