import json
import math
import os
import subprocess
import re
import time
from pathlib import Path

import click
import h5py
import numpy as np
from sklearn.model_selection import ParameterGrid, ParameterSampler
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data
from tqdm import tqdm

from src import qsub
from src import utils, data_utils, metrics
from src.eval_retrieval import eval_datasets
from src.modeling import batch_norm, models

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + '/'

params = {
    'ex_name': __file__.replace('.py', ''),
    'seed': 123456789,
    'lr': 1e-3,
    'batch_size': 32,
    'test_batch_size': 64,
    'optimizer': 'momentum',
    'epochs': 5,
    'scaleup_epochs': 0,
    'wd': 1e-5,
    'model_name': 'fishnet150',
    'pooling': 'G,G,G,G',
    'use_fc': True,
    'loss': 'arcface',
    'margin': 0.3,
    's': 30,
    # 'theta_zero': 1.25,
    'fc_dim': 512,
    'scale_limit': 0.2,
    'shear_limit': 0,
    'brightness_limit': 0.0,
    'contrast_limit': 0.0,
    'augmentation': 'soft',
    'data': 'train',
    'freeze_bn': True,
    'verifythresh': 50,
    'freqthresh': 5,
    'base_ckpt_path': ROOT + 'experiments/v7/ep4_batch_size-32_epochs-5_pooling-G,G,G,G.pth',
}


@click.group()
def cli():
    if not Path(ROOT + f'experiments/{params["ex_name"]}/train').exists():
        Path(ROOT + f'experiments/{params["ex_name"]}/train').mkdir(parents=True)
    if not Path(ROOT + f'experiments/{params["ex_name"]}/tuning').exists():
        Path(ROOT + f'experiments/{params["ex_name"]}/tuning').mkdir(parents=True)


@cli.command()
@click.option('--tuning', is_flag=True)
@click.option('--params-path', type=click.Path(), default=None, help='json file path for setting parameters')
@click.option('--devices', '-d', type=str, help='comma delimited gpu device list (e.g. "0,1")')
@click.option('--resume', type=str, default=None, help='checkpoint path')
@click.option('--save-interval', '-s', type=int, default=1, help='if 0 or negative value, not saving')
def job(tuning, params_path, devices, resume, save_interval):
    global params
    if tuning:
        with open(params_path, 'r') as f:
            params = json.load(f)
        mode_str = 'tuning'
        setting = '_'.join(f'{tp}-{params[tp]}' for tp in params['tuning_params'])
    else:
        mode_str = 'train'
        setting = ''

    # パラメーターを変えるときにseedも変えたい(seed averagingの効果を期待)
    seed = sum(ord(_) for _ in str(params.values()))
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False

    params['clean_path'] = ROOT + f'input/clean/train19_cleaned_verifythresh{params["verifythresh"]}_freqthresh{params["freqthresh"]}.csv'
    exp_path = ROOT + f'experiments/{params["ex_name"]}/'
    os.environ['CUDA_VISIBLE_DEVICES'] = devices

    logger, writer = utils.get_logger(log_dir=exp_path + f'{mode_str}/log/{setting}',
                                      tensorboard_dir=exp_path + f'{mode_str}/tf_board/{setting}')

    if params['augmentation'] == 'soft':
        params['scale_limit'] = 0.2
        params['brightness_limit'] = 0.1
    elif params['augmentation'] == 'middle':
        params['scale_limit'] = 0.3
        params['shear_limit'] = 4
        params['brightness_limit'] = 0.1
        params['contrast_limit'] = 0.1
    else:
        raise ValueError

    if resume is not None:
        params['clean_path'] = utils.load_checkpoint(path=resume, params=True)['params']['clean_path']

    train_transform, eval_transform = data_utils.build_transforms(
        scale_limit=params['scale_limit'],
        shear_limit=params['shear_limit'],
        brightness_limit=params['brightness_limit'],
        contrast_limit=params['contrast_limit'],
    )

    data_loaders = data_utils.make_verified_train_loaders(
        params=params,
        data_root=ROOT + 'input/' + params['data'],
        train_transform=train_transform,
        eval_transform=eval_transform,
        scale='S',
        test_size=0.1,
        num_workers=8)

    model = models.LandmarkFishNet(n_classes=params['class_topk'],
                                   model_name=params['model_name'],
                                   pooling_strings=params['pooling'].split(','),
                                   loss_module=params['loss'],
                                   s=params['s'],
                                   margin=params['margin'],
                                   use_fc=params['use_fc'],
                                   fc_dim=params['fc_dim'],
                                   pretrained=False,
                                   ).cuda()

    sdict = torch.load(params['base_ckpt_path'])['state_dict']
    del sdict['final.weight']  # remove fully-connected layer
    model.load_state_dict(sdict, strict=False)

    optimizer = utils.get_optim(params, model)
    criterion = nn.CrossEntropyLoss()
    scheduler = optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=params['epochs'] * len(data_loaders['train']), eta_min=3e-6)
    start_epoch = 0

    if len(devices.split(',')) > 1:
        model = nn.DataParallel(model)

    if resume is not None:
        ckpt = utils.load_checkpoint(path=resume, model=model, optimizer=optimizer, epoch=True, params=True)
        model, optimizer, start_epoch, params = ckpt['model'], ckpt['optimizer'], ckpt['epoch'] + 1, ckpt['params']

        scheduler = optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=params['epochs'] * len(data_loaders['train']), eta_min=3e-6,
            last_epoch=start_epoch * len(data_loaders['train']))

        # パラメーターを変えるときにseedも変えたい(seed averagingの効果を期待)
        seed = sum(ord(_) for _ in str(params.values())) + 12345
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.benchmark = False

        train_transform, eval_transform = data_utils.build_transforms(
            scale_limit=params['scale_limit'],
            shear_limit=params['shear_limit'],
            brightness_limit=params['brightness_limit'],
            contrast_limit=params['contrast_limit'],
        )
        data_loaders = data_utils.make_verified_train_loaders(
            params=params,
            data_root=ROOT + 'input/' + params['data'],
            train_transform=train_transform,
            eval_transform=eval_transform,
            scale='M2',
            test_size=0.1,
            num_workers=8)
        batch_norm.freeze_bn(model)
        setting += 'scaleup_' + resume.split('/')[-1].replace('.pth', '')

    for epoch in range(start_epoch, params['epochs']):
        logger.info(f'Epoch {epoch}/{params["epochs"]} | lr: {optimizer.param_groups[0]["lr"]}')

        # ============================== train ============================== #
        model.train(True)

        losses = utils.AverageMeter()
        prec1 = utils.AverageMeter()

        for i, (_, x, y) in tqdm(enumerate(data_loaders['train']),
                                 total=len(data_loaders['train']),
                                 miniters=None, ncols=55):
            x = x.to('cuda')
            y = y.to('cuda')

            outputs = model(x, y)
            loss = criterion(outputs, y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step()

            acc = metrics.accuracy(outputs, y)
            losses.update(loss.item(), x.size(0))
            prec1.update(acc, x.size(0))

            if i % 100 == 99:
                logger.info(f'{epoch+i/len(data_loaders["train"]):.2f}epoch | {setting} acc: {prec1.avg}')

        train_loss = losses.avg
        train_acc = prec1.avg

        writer.add_scalars('Loss', {'train': train_loss}, epoch)
        writer.add_scalars('Acc', {'train': train_acc}, epoch)
        writer.add_scalar('LR', optimizer.param_groups[0]['lr'], epoch)

        if (epoch + 1) == params['epochs'] or (epoch + 1) % save_interval == 0:
            output_file_name = exp_path + f'ep{epoch}_' + setting + '.pth'
            utils.save_checkpoint(path=output_file_name,
                                  model=model,
                                  epoch=epoch,
                                  optimizer=optimizer,
                                  params=params)

    model = model.module
    datasets = ('roxford5k', 'rparis6k')
    results = eval_datasets(model, datasets=datasets, ms=False, tta_gem_p=1.0, logger=logger)

    if tuning:
        tuning_result = {}
        for d in datasets:
            for key in ['mapE', 'mapM', 'mapH']:
                mapE, mapM, mapH, mpE, mpM, mpH, kappas = results[d]
                tuning_result[d + '-' + key] = [eval(key)]
        utils.write_tuning_result(params, tuning_result, exp_path + 'tuning/results.csv')


@cli.command()
@click.option('--mode', type=str, default='grid', help='Search method (tuning)')
@click.option('--n-iter', type=int, default=10, help='n of iteration for random parameter search (tuning)')
@click.option('--n-gpu', type=int, default=-1, help='n of used gpu at once')
@click.option('--devices', '-d', type=str, help='comma delimited gpu device list (e.g. "0,1")')
@click.option('--save-interval', '-s', type=int, default=-1, help='if 0 or negative value, not saving')
@click.option('--n-blocks', '-n', type=int, default=1)
@click.option('--block-id', '-i', type=int, default=0)
def tuning(mode, n_iter, n_gpu, devices, save_interval, n_blocks, block_id):
    """
    Example:
        python v12.py tuning --devices 0,1,2,3 --n-gpu 2
    """

    if n_gpu == -1:
        n_gpu = len(devices.split(','))

    space = [
        {
            'batch_size': [32],
            # 'epochs': [1],
            # 'scaleup_epochs': [0],
            'loss': ['arcface'],
            'verifythresh': [30],
            'freqthresh': [3],
            'margin': [0.3, 0.2],
            's': [30],
            # 'augmentation': ['soft'],
        },

        # for reproduce
        # {
        #     'batch_size': [32],
        #     'loss': ['arcface'],
        #     'verifythresh': [30, 40, 50],
        #     'freqthresh': [2, 3, 5],
        #     'margin': [0.3],
        #     's': [30],
        # },
        # {
        #     'batch_size': [32],
        #     'loss': ['arcface'],
        #     'verifythresh': [30],
        #     'freqthresh': [3],
        #     'margin': [0.2],
        #     's': [30],
        # },
    ]

    if mode == 'grid':
        candidate_list = list(ParameterGrid(space))
    elif mode == 'random':
        candidate_list = list(ParameterSampler(space, n_iter, random_state=params['seed']))
    else:
        raise ValueError

    n_per_block = math.ceil(len(candidate_list) / n_blocks)
    candidate_chunk = candidate_list[block_id * n_per_block: (block_id + 1) * n_per_block]

    utils.launch_tuning(mode=mode, n_iter=n_iter, n_gpu=n_gpu, devices=devices,
                        params=params, root=ROOT, save_interval=save_interval,
                        candidate_list=candidate_chunk)


@cli.command()
@click.option('--model-path', '-m', type=str)
@click.option('--devices', '-d', type=str, default="0", help='comma delimited gpu device list (e.g. "0,1")')
@click.option('--ms', is_flag=True)
@click.option('--scale', type=str, default='S')
@click.option('--batch-size', '-b', type=int, default=64)
@click.option('--splits', type=str, default='index,test')
@click.option('--n-blocks', '-n', type=int, default=1)
@click.option('--block-id', '-i', type=int, default=0)
def predict(model_path, devices, ms, scale, batch_size, splits, n_blocks, block_id):
    """
    python v7.py predict -m v7/ep4_batch_size-32_epochs-5_pooling-G,G,G,G.pth -d 3 &
    python v7.py predict -m v7/ep4_batch_size-32_epochs-5_pooling-G,G,G,G.pth --ms --scale M -b 28 -d 0 &
    """
    os.environ['CUDA_VISIBLE_DEVICES'] = devices

    ckpt = torch.load(model_path)
    params, state_dict = ckpt['params'], ckpt['state_dict']
    params['test_batch_size'] = batch_size

    splits = splits.split(',')

    model = models.LandmarkFishNet(n_classes=params['class_topk'],
                                   model_name=params['model_name'],
                                   pooling_strings=params['pooling'].split(','),
                                   loss_module=params['loss'],
                                   s=params['s'],
                                   margin=params['margin'],
                                   use_fc=params['use_fc'],
                                   fc_dim=params['fc_dim'],
                                   )
    if params['freeze_bn']:
        batch_norm.freeze_bn(model)
    model.load_state_dict(state_dict)
    model = model.to('cuda').eval()

    train_transform, eval_transform = data_utils.build_transforms()
    data_loaders = data_utils.make_predict_loaders(params,
                                                   eval_transform=eval_transform,
                                                   scale=scale,
                                                   splits=splits,
                                                   num_workers=8,
                                                   n_blocks=n_blocks,
                                                   block_id=block_id)

    exp_path = ROOT + f'experiments/{params["ex_name"]}/'

    file_suffix = model_path.split('/')[-1].replace('.pth', '')
    file_suffix = scale + '_' + file_suffix
    file_suffix = 'ms_' + file_suffix if ms else file_suffix

    scales = [0.75, 1.0, 1.25] if ms else [1.0]

    for split in splits:
        ids, feats = [], []
        for i, (img_id, x) in tqdm(enumerate(data_loaders[split]),
                                   total=len(data_loaders[split]),
                                   miniters=None, ncols=55):

            batch_size, _, h, w = x.shape
            feat_blend = np.zeros((batch_size, params['fc_dim']), dtype=np.float32)

            with torch.no_grad():
                x = x.to('cuda')

                for s in scales:
                    size = int(h * s // model.DIVIDABLE_BY * model.DIVIDABLE_BY), \
                           int(w * s // model.DIVIDABLE_BY * model.DIVIDABLE_BY)  # round off
                    scaled_x = F.interpolate(x, size=size, mode='bilinear', align_corners=True)
                    feat = model.extract_feat(scaled_x)
                    feat = feat.cpu().numpy()
                    feat_blend += feat

            feats.append(feat_blend)
            ids.extend(img_id)

        feats = np.concatenate(feats) / len(scales)

        output_path = Path(f'{exp_path}feats_{split}_{file_suffix}')
        output_path.mkdir(parents=True, exist_ok=True)
        with h5py.File(output_path / f'block{block_id}.h5', 'a') as f:
            f.create_dataset('ids', data=np.array(ids, dtype=f'S{len(ids[0])}'))
            f.create_dataset('feats', data=feats)


@cli.command()
@click.argument('job-type')
@click.option('--mode', type=str, default='grid', help='Search method (tuning)')
@click.option('--n-iter', type=int, default=10, help='n of iteration for random parameter search (tuning)')
@click.option('--n-gpu', type=int, default=-1, help='n of used gpu at once')
@click.option('--devices', '-d', type=str, help='comma delimited gpu device list (e.g. "0,1")')
@click.option('--save-interval', '-s', type=int, default=-1, help='if 0 or negative value, not saving')
@click.option('--model-path', '-m', default=None, type=str)
@click.option('--ms', is_flag=True)
@click.option('--scale', type=str, default='S')
@click.option('--batch-size', '-b', type=int, default=64)
@click.option('--splits', type=str, default='index,test')
@click.option('--n-blocks', '-n', type=int, default=1)
@click.option('--instance-type', type=str, default='rt_G.small')
def launch_qsub(job_type,
                mode, n_iter, n_gpu, devices, save_interval,  # tuning args
                model_path, ms, scale, batch_size, splits,  # predict args
                n_blocks, instance_type
                ):
    """
    python v17.py launch-qsub predict -m v17/ep --ms --scale L --batch-size 26 --splits index,test --n-blocks 32
    python v7c.py launch-qsub tuning -d 0,1,2,3 --n-gpu 4 --n-blocks 2 -s 1 --instance-type rt_F
    """
    exp_path = ROOT + f'experiments/{params["ex_name"]}/'
    logger = utils.get_logger(log_dir=exp_path)
    job_ids = []
    for block_id in range(n_blocks):
        if job_type == 'tuning':
            cmd_with_args = [
                "python", "-W", "ignore", __file__, "tuning",
                "--mode", mode,
                "--n-iter", str(n_iter),
                "--n-gpu", str(n_gpu),
                "--devices", devices,
                "--save-interval", str(save_interval),
                "--n-blocks", str(n_blocks),
                "--block-id", str(block_id),
            ]
        elif job_type == 'predict':
            cmd_with_args = [
                "python", "-W", "ignore", __file__, "predict",
                "-m", model_path,
                "--splits", splits,
                "--scale", scale,
                "--batch-size", str(batch_size),
                "--n-blocks", str(n_blocks),
                "--block-id", str(block_id),
            ]
            if ms:
                cmd_with_args.append("--ms")
        else:
            raise ValueError('job-type should be one of "tuning" or "predict"')
        proc = qsub.qsub(cmd_with_args,
                         n_hours=48,
                         instance_type=instance_type,
                         logger=logger)
        logger.info(f'Response from qsub: {proc.returncode}')

        m = re.match(r'Your job (\d+) \(', proc.stdout.decode('utf8'))
        job_id = int(m.group(1)) if m is not None else None
        logger.info(f'Job id: {job_id}')
        assert job_id is not None
        job_ids.append(job_id)
        time.sleep(1)

    qsub.monitor_jobs(job_ids, logger)


@cli.command()
@click.option('--devices', '-d', type=str, help='comma delimited gpu device list (e.g. "0,1")')
@click.option('--model-path', '-m', default=None, type=str)
@click.option('--ms', is_flag=True)
@click.option('--scale', type=str, default='S')
@click.option('--batch-size', '-b', type=int, default=64)
@click.option('--splits', type=str, default='index,test')
def multigpu_predict(devices, model_path, ms, scale, batch_size, splits):
    devices = devices.split(',')

    procs = []
    for block_id, d in enumerate(devices):
        cmd_with_args = [
            "python", __file__, "predict",
            "-d", d,
            "-m", model_path,
            "--splits", splits,
            "--scale", scale,
            "--batch-size", str(batch_size),
            "--n-blocks", str(len(devices)),
            "--block-id", str(block_id),
        ]
        if ms:
            cmd_with_args.append("--ms")
        procs.append(subprocess.Popen(cmd_with_args))

    while True:
        time.sleep(1)
        if all(proc.poll() is not None for proc in procs):
            print('All jobs have finished.')
            break


if __name__ == '__main__':
    cli()
