import os

from utils.argparse_utils import manage_required_args
from exp.experimenter import *
from data.hico.hico_constants import HicoConstants
from data.vcoco.vcoco_constants import VcocoConstants
from utils.constants import Constants, ExpConstants
import exp.hoi_classifier.data.hoi_candidates as hoi_candidates
import exp.hoi_classifier.data.label_hoi_candidates as label_hoi_candidates
from exp.hoi_classifier.models.hoi_classifier_model import HoiClassifierConstants
import exp.hoi_classifier.train as train
import exp.hoi_classifier.eval as evaluate
from exp.hoi_classifier.data.features_dataset import FeatureConstantsHico, FeatureConstantsVcoco
import exp.hoi_classifier.data.cache_box_features as cache_box_features
import exp.hoi_classifier.data.cache_pose_features as cache_pose_features
import exp.hoi_classifier.data.assign_pose_to_human_candidates as \
    assign_pose_to_human_candidates
import exp.hoi_classifier.vis.top_boxes_per_hoi as \
    vis_top_boxes_per_hoi


parser.add_argument(
    '--gen_hoi_cand',
    default=False,
    action='store_true',
    help='Apply this flag to generate hoi candidates')
parser.add_argument(
    '--label_hoi_cand',
    default=False,
    action='store_true',
    help='Apply this flag to label hoi candidates')
parser.add_argument(
    '--subset',
    type=str,
    choices=['train','train_val','val','test'],
    help='Apply this flag to specify subset of data')
parser.add_argument(
    '--imgs_per_batch',
    type=int,
    default=1,
    help='Number of images per batch')
parser.add_argument(
    '--fp_to_tp_ratio',
    type=int,
    default=1000,
    help='Number of images per batch')
parser.add_argument(
    '--model_num',
    type=int,
    help='Specify model number to evaluate')
parser.add_argument(
    '--verb_given_appearance',
    default=False,
    action='store_true',
    help='Use verb_given_human/object_appearance factor')
parser.add_argument(
    '--verb_given_human_appearance',
    default=False,
    action='store_true',
    help='Set verb_given_human_appearance factor')
parser.add_argument(
    '--verb_given_object_appearance',
    default=False,
    action='store_true',
    help='Set verb_given_object_appearance factor')
parser.add_argument(
    '--verb_given_boxes_and_object_label',
    default=False,
    action='store_true',
    help='Use verb_given_boxes_and_object_label factor')
parser.add_argument(
    '--verb_given_human_pose',
    default=False,
    action='store_true',
    help='Use verb_given_human_pose factor')
parser.add_argument(
    '--rcnn_det_prob',
    default=False,
    action='store_true',
    help='Use detection prob from Faster-RCNN')


def exp_gen_and_label_hoi_cand_hico():
    exp_name = 'hoi_candidates'
    exp_const = ExpConstants(exp_name=exp_name)
    data_const = HicoConstants()
    data_const.selected_dets_hdf5 = os.path.join(
        os.getcwd(),
        'data_symlinks/hico_exp/select_confident_boxes_in_hico/' + \
        'selected_coco_cls_dets.hdf5')
    _exp_gen_and_label_hoi_cand(exp_const, data_const, 'hico')


def exp_gen_and_label_hoi_cand_vcoco():
    exp_name = 'hoi_candidates'
    exp_const = ExpConstants(exp_name=exp_name, out_base_dir=os.path.join(os.getcwd(), 'data_symlinks/vcoco_exp'))
    data_const = VcocoConstants()
    data_const.selected_dets_hdf5 = os.path.join(
        os.getcwd(),
        'data_symlinks/vcoco_exp/select_confident_boxes_in_vcoco/' + \
        'selected_coco_cls_dets.hdf5')
    _exp_gen_and_label_hoi_cand(exp_const, data_const, 'vcoco')


def _exp_gen_and_label_hoi_cand(exp_const, data_const, data_sign='hico'):
    args = parser.parse_args()
    not_specified_args = manage_required_args(
        args,
        parser,
        required_args=['subset'],
        optional_args=['gen_hoi_cand', 'label_hoi_cand'])
    if len(not_specified_args) > 0:
        return

    exp_const.subset = args.subset

    if args.gen_hoi_cand:
        print('Generating HOI candidates from Faster-RCNN dets...')
        hoi_candidates.generate(exp_const, data_const, data_sign)
    
    if args.label_hoi_cand:
        print('Labelling HOI candidates from Faster-RCNN dets...')
        data_const.hoi_cand_hdf5 = os.path.join(
            exp_const.exp_dir,
            f'hoi_candidates_{exp_const.subset}.hdf5')
        label_hoi_candidates.assign(exp_const, data_const)


def exp_cache_box_feats_hico():
    exp_name = 'hoi_candidates'
    exp_const = ExpConstants(exp_name=exp_name)
    data_const = HicoConstants()
    _exp_cache_box_feats(exp_const, data_const)

def exp_cache_box_feats_vcoco():
    exp_name = 'hoi_candidates'
    exp_const = ExpConstants(exp_name=exp_name, out_base_dir=os.path.join(os.getcwd(), 'data_symlinks/vcoco_exp'))
    data_const = VcocoConstants()
    _exp_cache_box_feats(exp_const, data_const)


def _exp_cache_box_feats(exp_const, data_const):
    args = parser.parse_args()
    
    not_specified_args = manage_required_args(
        args,
        parser,
        required_args=['subset'])

    exp_const.subset = args.subset
    data_const.hoi_cand_hdf5 = os.path.join(
        exp_const.exp_dir,
        f'hoi_candidates_{exp_const.subset}.hdf5')

    cache_box_features.main(exp_const, data_const)


def exp_assign_pose_to_human_cand_hico():
    exp_name = 'hoi_candidates'
    exp_const = ExpConstants(exp_name=exp_name)
    data_const = HicoConstants()
    _exp_assign_pose_to_human_cand(exp_const, data_const)

def exp_assign_pose_to_human_cand_vcoco():
    exp_name = 'hoi_candidates'
    exp_const = ExpConstants(exp_name=exp_name, out_base_dir=os.path.join(os.getcwd(), 'data_symlinks/vcoco_exp'))
    data_const = VcocoConstants()
    _exp_assign_pose_to_human_cand(exp_const, data_const)


def _exp_assign_pose_to_human_cand(exp_const, data_const):
    args = parser.parse_args()
    
    not_specified_args = manage_required_args(
        args,
        parser,
        required_args=['subset'])

    exp_const.subset = args.subset
    data_const.hoi_cand_hdf5 = os.path.join(
        exp_const.exp_dir,
        f'hoi_candidates_{exp_const.subset}.hdf5')
    data_const.human_pose_dir = os.path.join(
        data_const.proc_dir,
        'human_pose')
    data_const.num_keypoints = 18

    assign_pose_to_human_candidates.main(exp_const, data_const)


def exp_cache_pose_feats_hico():
    exp_name = 'hoi_candidates'
    exp_const = ExpConstants(exp_name=exp_name)
    data_const = HicoConstants()
    _exp_cache_pose_feats(exp_const, data_const)

def exp_cache_pose_feats_vcoco():
    exp_name = 'hoi_candidates'
    exp_const = ExpConstants(exp_name=exp_name, out_base_dir=os.path.join(os.getcwd(), 'data_symlinks/vcoco_exp'))
    data_const = VcocoConstants()
    _exp_cache_pose_feats(exp_const, data_const)


def _exp_cache_pose_feats(exp_const, data_const):
    args = parser.parse_args()
    
    not_specified_args = manage_required_args(
        args,
        parser,
        required_args=['subset'])

    exp_const.subset = args.subset
    data_const.hoi_cand_hdf5 = os.path.join(
        exp_const.exp_dir,
        f'hoi_candidates_{exp_const.subset}.hdf5')
    data_const.human_cands_pose_hdf5 = os.path.join(
        exp_const.exp_dir,
        f'human_candidates_pose_{exp_const.subset}.hdf5')
    data_const.num_keypoints = 18

    cache_pose_features.main(exp_const, data_const)


def exp_train_hico():
    out_base_dir = os.path.join(
        os.getcwd(), 'data_symlinks/hico_exp/hoi_classifier')
    data_const_train = FeatureConstantsHico(subset='train')
    data_const_val = FeatureConstantsHico(subset='val')
    _exp_train(out_base_dir, data_const_train, data_const_val, 'hico')


def exp_train_vcoco():
    out_base_dir = os.path.join(
        os.getcwd(), 'data_symlinks/vcoco_exp/hoi_classifier')
    data_const_train = FeatureConstantsVcoco(subset='train')
    data_const_val = FeatureConstantsVcoco(subset='val')
    _exp_train(out_base_dir, data_const_train, data_const_val, 'vcoco')


# data_sign指明当前处理的数据集
def _exp_train(out_base_dir, data_const_train, data_const_val, data_sign='hico'):
    args = parser.parse_args()
    not_specified_args = manage_required_args(
        args,
        parser,
        required_args=['imgs_per_batch', 'fp_to_tp_ratio'],
        optional_args=[
            'verb_given_appearance',
            'verb_given_human_appearance',
            'verb_given_object_appearance',
            'verb_given_boxes_and_object_label',
            'verb_given_human_pose',
            'rcnn_det_prob'])

    exp_name = 'factors'
    if args.rcnn_det_prob:
        exp_name += '_rcnn_det_prob'
    if args.verb_given_appearance:
        exp_name += '_appearance'
    if args.verb_given_human_appearance:
        exp_name += '_human_appearance'
    if args.verb_given_object_appearance:
        exp_name += '_object_appearance'
    if args.verb_given_boxes_and_object_label:
        exp_name += '_boxes_and_object_label'
    if args.verb_given_human_pose:
        pose = '_human_pose'
        exp_name += pose

    exp_const = ExpConstants(
        exp_name=exp_name,
        out_base_dir=out_base_dir)
    exp_const.log_dir = os.path.join(exp_const.exp_dir, 'log')
    exp_const.model_dir = os.path.join(exp_const.exp_dir, 'models')
    exp_const.num_epochs = 10
    exp_const.imgs_per_batch = args.imgs_per_batch
    exp_const.lr = 1e-3

    model_const = Constants()
    model_const.hoi_classifier = HoiClassifierConstants(data_sign)
    model_const.hoi_classifier.verb_given_appearance = args.verb_given_appearance
    model_const.hoi_classifier.verb_given_human_appearance = args.verb_given_human_appearance
    model_const.hoi_classifier.verb_given_object_appearance = args.verb_given_object_appearance
    model_const.hoi_classifier.verb_given_boxes_and_object_label = args.verb_given_boxes_and_object_label
    model_const.hoi_classifier.verb_given_human_pose = args.verb_given_human_pose
    model_const.hoi_classifier.rcnn_det_prob = args.rcnn_det_prob

    train.main(exp_const, data_const_train, data_const_val, model_const, data_sign)


def exp_eval_hico():
    out_base_dir = os.path.join(
        os.getcwd(), 'data_symlinks/hico_exp/hoi_classifier')
    data_const = FeatureConstantsHico(subset='test')
    _exp_eval(out_base_dir, data_const)


def exp_eval_vcoco():
    out_base_dir = os.path.join(
        os.getcwd(), 'data_symlinks/vcoco_exp/hoi_classifier')
    data_const = FeatureConstantsVcoco(subset='test')
    _exp_eval(out_base_dir, data_const)


def _exp_eval(out_base_dir, data_const):
    args = parser.parse_args()
    not_specified_args = manage_required_args(
        args,
        parser,
        required_args=['model_num'],
        optional_args=[
            'verb_given_appearance',
            'verb_given_human_appearance',
            'verb_given_object_appearance',
            'verb_given_boxes_and_object_label',
            'verb_given_human_pose',
            'rcnn_det_prob'])

    exp_name = 'factors'
    if args.rcnn_det_prob:
        exp_name += '_rcnn_det_prob'
    if args.verb_given_appearance:
        exp_name += '_appearance'
    if args.verb_given_human_appearance:
        exp_name += '_human_appearance'
    if args.verb_given_object_appearance:
        exp_name += '_object_appearance'
    if args.verb_given_boxes_and_object_label:
        exp_name += '_boxes_and_object_label'
    if args.verb_given_human_pose:
        exp_name += '_human_pose'

    exp_const = ExpConstants(
        exp_name=exp_name,
        out_base_dir=out_base_dir)
    exp_const.model_dir = os.path.join(exp_const.exp_dir, 'models')
    data_const.balanced_sampling = False
    
    model_const = Constants()
    model_const.model_num = args.model_num
    model_const.hoi_classifier = HoiClassifierConstants()
    model_const.hoi_classifier.verb_given_appearance = args.verb_given_appearance
    model_const.hoi_classifier.verb_given_human_appearance = args.verb_given_human_appearance
    model_const.hoi_classifier.verb_given_object_appearance = args.verb_given_object_appearance
    model_const.hoi_classifier.verb_given_boxes_and_object_label = args.verb_given_boxes_and_object_label
    model_const.hoi_classifier.verb_given_human_pose = args.verb_given_human_pose
    model_const.hoi_classifier.rcnn_det_prob = args.rcnn_det_prob
    model_const.hoi_classifier.model_pth = os.path.join(
        exp_const.model_dir,
        f'hoi_classifier_{model_const.model_num}')

    if isinstance(data_const, FeatureConstantsVcoco):
        data_sign = 'vcoco'
    else:
        data_sign = 'hico'
    evaluate.main(exp_const, data_const, model_const, data_sign)


def exp_top_boxes_per_hoi_hico():
    out_base_dir = os.path.join(
        os.getcwd(), 'data_symlinks/hico_exp/hoi_classifier')
    data_const = FeatureConstantsHico(subset='test')
    _exp_top_boxes_per_hoi(out_base_dir, data_const)


def exp_top_boxes_per_hoi_vcoco():
    out_base_dir = os.path.join(
        os.getcwd(), 'data_symlinks/vcoco_exp/hoi_classifier')
    data_const = FeatureConstantsVcoco(subset='test')
    _exp_top_boxes_per_hoi(out_base_dir, data_const)


def _exp_top_boxes_per_hoi(out_base_dir, data_const):
    args = parser.parse_args()
    not_specified_args = manage_required_args(
        args,
        parser,
        required_args=['model_num'],
        optional_args=[
            'verb_given_appearance',
            'verb_given_human_appearance',
            'verb_given_object_appearance',
            'verb_given_boxes_and_object_label',
            'verb_given_human_pose',
            'rcnn_det_prob'])

    exp_name = 'factors'
    if args.rcnn_det_prob:
        exp_name += '_rcnn_det_prob'
    if args.verb_given_appearance:
        exp_name += '_appearance'
    if args.verb_given_human_appearance:
        exp_name += '_human_appearance'
    if args.verb_given_object_appearance:
        exp_name += '_object_appearance'
    if args.verb_given_boxes_and_object_label:
        exp_name += '_boxes_and_object_label'
    if args.verb_given_human_pose:
        exp_name += '_human_pose'

    exp_const = ExpConstants(
        exp_name=exp_name,
        out_base_dir=out_base_dir)
    exp_const.model_dir = os.path.join(exp_const.exp_dir, 'models')
    exp_const.num_to_vis = 10

    data_const.pred_hoi_dets_h5py = os.path.join(
        exp_const.exp_dir,
        f'pred_hoi_dets_test_{args.model_num}.hdf5')
    hoi_cand_dir = os.path.join(
        os.getcwd(),
        'data_symlinks/vcoco_exp/hoi_candidates')
    data_const.human_pose_feats_hdf5 = os.path.join(
        hoi_cand_dir,
        'human_pose_feats_test.hdf5')
    data_const.num_pose_keypoints = 18
    
    model_const = Constants()
    model_const.model_num = args.model_num
    model_const.hoi_classifier = HoiClassifierConstants()
    model_const.hoi_classifier.verb_given_appearance = args.verb_given_appearance
    model_const.hoi_classifier.verb_given_boxes_and_object_label = args.verb_given_boxes_and_object_label
    model_const.hoi_classifier.verb_given_human_pose = args.verb_given_human_pose
    model_const.hoi_classifier.rcnn_det_prob = args.rcnn_det_prob
    model_const.hoi_classifier.model_pth = os.path.join(
        exp_const.model_dir,
        f'hoi_classifier_{model_const.model_num}')

    vis_top_boxes_per_hoi.main(exp_const, data_const, model_const)


if __name__ == '__main__':
    list_exps(globals())
