import os
import argparse
from pickle import FALSE
import csv

from douzero.evaluation.simulation import evaluate


def make_evaluate(args, vs_config):
    landlord_config = vs_config['landlord']
    farmer_config = vs_config['farmer']
    if landlord_config['prefix'] == 'random':
        args.landlord = 'random'
    else:
        args.landlord = '%s/%s_landlord_%i.ckpt' % (
            landlord_config['folder'], landlord_config['prefix'], landlord_config['frame']
        )
    if farmer_config['prefix'] == 'random':
        args.landlord_up = 'random'
        args.landlord_front = 'random'
        args.landlord_down = 'random'
    else:
        args.landlord_up = '%s/%s_landlord_up_%i.ckpt' % (
            farmer_config['folder'], farmer_config['prefix'], farmer_config['frame']
        )
        args.landlord_front = '%s/%s_landlord_front_%i.ckpt' % (
            farmer_config['folder'], farmer_config['prefix'], farmer_config['frame']
        )
        args.landlord_down = '%s/%s_landlord_down_%i.ckpt' % (
            farmer_config['folder'], farmer_config['prefix'], farmer_config['frame']
        )

    print('%s %i\tvs\t%s %i' % (
        landlord_config['prefix'], landlord_config['frame'], farmer_config['prefix'], farmer_config['frame']))


    landlord_wp, farmer_wp, landlord_adp, farmer_adp =  evaluate(args.landlord,
             args.landlord_up,
             args.landlord_front,
             args.landlord_down,
             args.eval_data,
             args.num_workers,
             args.output,
             args.title)

    with open('result.csv', 'w') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow([
            landlord_wp,
            farmer_wp,
            landlord_adp,
            farmer_adp
        ])
    


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        'Dou Dizhu Evaluation')
    parser.add_argument('--landlord', type=str,
                        default='baselines/lite_uni_5001600.ckpt')
    parser.add_argument('--landlord_up', type=str,
                        default='baselines/lite_uni_5001600.ckpt')
    parser.add_argument('--landlord_front', type=str,
                        default='baselines/lite_uni_5001600.ckpt')
    parser.add_argument('--landlord_down', type=str,
                        default='baselines/lite_uni_5001600.ckpt')
    parser.add_argument('--eval_data', type=str,
                        default='eval_data_200_r.pkl')
    parser.add_argument('--num_workers', type=int, default=3)
    parser.add_argument('--gpu_device', type=str, default='0')
    parser.add_argument('--output', type=bool, default=True)
    parser.add_argument('--title', type=str, default='New')
    args = parser.parse_args()
    # args.output = True
    args.output = False
    if args.output:
        args.num_workers = 1

    os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_device

    landlord_wp, farmer_wp, landlord_adp, farmer_adp =  evaluate(args.landlord,
             args.landlord_up,
             args.landlord_front,
             args.landlord_down,
             args.eval_data,
             args.num_workers,
             args.output,
             args.title)
    print(landlord_wp, farmer_wp, landlord_adp, farmer_adp)
    os._exit(0)

    baselines = [
        {'folder': 'baselines', 'prefix': 'legacy_general', 'frame': 736107200},
        {'folder': 'baselines', 'prefix': 'legacy_general', 'frame': 479412800},
        {'folder': 'baselines', 'prefix': 'legacy_general', 'frame': 262529600},
        {'folder': 'baselines', 'prefix': 'legacy_general', 'frame': 143539200},
    ]
    evaluations = [
            {'folder': 'baselines', 'prefix': 'lite_general', 'frame': 299457600},
            {'folder': 'baselines', 'prefix': 'lite_resnet', 'frame': 124064400},
    ]
    
    eval_list = [
        {
            'farmer': baseline,
            'landlord': evaluation,
            'two_way': True
        }
        for evaluation in evaluations for baseline in baselines
    ]

    eval_list.extend([
        {
            'farmer': evaluation,
            'landlord': evaluation,
            'two_way': False
        }
        for evaluation in evaluations
    ])

    # eval_list = [
    #     {
    #         'farmer': { 'folder': 'baselines', 'prefix': 'lite_general', 'frame': 221553600},
    #         'landlord': { 'folder': 'baselines', 'prefix': 'lite_resnet', 'frame': 102494400},
    #         'two_way': False
    #     },
    # ]

    for vs in eval_list:
        make_evaluate(args, vs)
        if vs['two_way']:
            vs['landlord_bak'] = vs['landlord']
            vs['landlord'] = vs['farmer']
            vs['farmer'] = vs['landlord_bak']
            make_evaluate(args, vs)