import os
import sys
import glob
import time
import copy
import random
import numpy as np
import torch
import utils
import logging
import argparse
import torch.nn as nn
import torch.utils
import torch.nn.functional as F
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torch.backends.cudnn as cudnn
from model_ws import NASNet
from nao_controller import NAO

parser = argparse.ArgumentParser(description='SemiNAS Search')

# Basic model parameters.
parser.add_argument('--mode', type=str, default='train', choices=['train', 'test'])
parser.add_argument('--data', type=str, default='data')
parser.add_argument('--zip_file', action='store_true', default=False)
parser.add_argument('--lazy_load', action='store_true', default=False)
parser.add_argument('--output_dir', type=str, default='models')
parser.add_argument('--seed', type=int, default=1)
parser.add_argument('--batch_size', type=int, default=512)
parser.add_argument('--eval_batch_size', type=int, default=512)
parser.add_argument('--layers', type=int, default=21)
parser.add_argument('--dropout', type=float, default=0)
parser.add_argument('--model_init', type=str, default='he_fout', choices=['he_fin', 'he_fout'])
parser.add_argument('--no_decay_keys', type=str, default='bn', choices=['None', 'bn', 'bn#bias'])
parser.add_argument('--max_num_updates', type=int, default=1000)
parser.add_argument('--grad_clip', type=float, default=0)
parser.add_argument('--weight_decay', type=float, default=4e-5)
parser.add_argument('--arch_pool', type=str, default=None)
parser.add_argument('--lr', type=float, default=0.4)
parser.add_argument('--lr_scheduler', type=str, default='cosine', choices=['cosine', 'linear'])
parser.add_argument('--label_smooth', type=float, default=0.1, help='label smoothing')
parser.add_argument('--gamma', type=float, default=0.97, help='learning rate decay')
parser.add_argument('--decay_period', type=int, default=1, help='epochs between two learning rate decays')
parser.add_argument('--num_workers', type=int, default=32)
parser.add_argument('--log_interval', type=int, default=100)
parser.add_argument('--select_narchs', type=int, default=2000)
parser.add_argument('--saved_number', type=int, default=0)

parser.add_argument('--width_stages', type=str, default='24,40,80,96,192,320')
parser.add_argument('--n_cell_stages', type=str, default='4,4,4,4,4,1')
parser.add_argument('--stride_stages', type=str, default='2,2,2,1,2,1')

parser.add_argument('--controller_iterations', type=int, default=3)
parser.add_argument('--controller_seed_arch', type=int, default=100)
parser.add_argument('--controller_new_arch', type=int, default=100)
parser.add_argument('--controller_random_arch', type=int, default=4000)
parser.add_argument('--controller_up_sample_ratio', type=int, default=None)
parser.add_argument('--controller_encoder_layers', type=int, default=1)
parser.add_argument('--controller_decoder_layers', type=int, default=1)
parser.add_argument('--controller_mlp_layers', type=int, default=2)
parser.add_argument('--controller_hidden_size', type=int, default=16)
parser.add_argument('--controller_mlp_hidden_size', type=int, default=64)
parser.add_argument('--controller_dropout', type=float, default=0.1)
parser.add_argument('--controller_l2_reg', type=float, default=1e-4)
parser.add_argument('--controller_trade_off', type=float, default=0.8)
parser.add_argument('--controller_pretrain_epochs', type=int, default=10000)
parser.add_argument('--controller_epochs', type=int, default=1000)
parser.add_argument('--controller_batch_size', type=int, default=100)
parser.add_argument('--controller_lr', type=float, default=0.001)
parser.add_argument('--controller_optimizer', type=str, default='adam')
parser.add_argument('--controller_grad_clip', type=float, default=5.0)
args = parser.parse_args()

utils.create_exp_dir(args.output_dir, scripts_to_save=glob.glob('*.py'))

log_format = '%(asctime)s %(message)s'
logging.basicConfig(stream=sys.stdout, level=logging.INFO,
    format=log_format, datefmt='%m/%d %I:%M:%S %p')

def build_imagenet(model_state_dict, optimizer_state_dict, **kwargs):
    valid_ratio = kwargs.pop('valid_ratio', None)
    valid_num = kwargs.pop('valid_num', None)
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ColorJitter(
            brightness=0.4,
            contrast=0.4,
            saturation=0.4,
            hue=0.1),
        transforms.ToTensor(),
        normalize,
    ])
    valid_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        normalize,
    ])
    if args.zip_file:
        logging.info('Loading data from zip file')
        traindir = os.path.join(args.data, 'train.zip')
        if args.lazy_load:
            data = utils.ZipDataset(traindir)
        else:
            logging.info('Loading data into memory')
            data = utils.InMemoryZipDataset(traindir, num_workers=args.num_workers)
    else:
        logging.info('Loading data from directory')
        traindir = os.path.join(args.data, 'train')
        if args.lazy_load:
            data = dset.ImageFolder(traindir)
        else:
            logging.info('Loading data into memory')
            data = utils.InMemoryDataset(traindir, num_workers=args.num_workers)

    num_data = len(data)
    indices = list(range(num_data))
    np.random.shuffle(indices)
    if valid_ratio is not None:
        split = int(np.floor(1 - valid_ratio * num_data))
        train_indices = sorted(indices[:split])
        valid_indices = sorted(indices[split:])
    else:
        assert valid_num is not None
        train_indices = sorted(indices[valid_num:])
        valid_indices = sorted(indices[:valid_num])

    train_data = utils.WrappedDataset(data, train_indices, train_transform)
    valid_data = utils.WrappedDataset(data, valid_indices, valid_transform)
    logging.info('train set = %d', len(train_data))
    logging.info('valid set = %d', len(valid_data))

    train_queue = torch.utils.data.DataLoader(
        train_data, batch_size=args.batch_size,
        sampler=torch.utils.data.sampler.SubsetRandomSampler(train_indices),
        pin_memory=True, num_workers=args.num_workers, drop_last=False)
    valid_queue = torch.utils.data.DataLoader(
        valid_data, batch_size=args.eval_batch_size,
        sampler=torch.utils.data.sampler.SubsetRandomSampler(valid_indices),
        pin_memory=True, num_workers=args.num_workers, drop_last=False)

    model = NASNet(args.width_stages, args.n_cell_stages, args.stride_stages, args.dropout)
    model.init_model(args.model_init)
    model.set_bn_param(0.1, 0.001)
    logging.info("param size = %d", utils.count_parameters(model))

    if model_state_dict is not None:
        model.load_state_dict(model_state_dict)

    if torch.cuda.device_count() > 1:
        logging.info("Use %d %s", torch.cuda.device_count(), "GPUs !")
        model = nn.DataParallel(model)
    model = model.cuda()

    if args.no_decay_keys:
        keys = args.no_decay_keys.split('#')
        net_params=[model.module.get_parameters(keys, mode='exclude'),
                    model.module.get_parameters(keys, mode='include')]
        optimizer = torch.optim.SGD([
            {'params': net_params[0], 'weight_decay': args.weight_decay},
            {'params': net_params[1], 'weight_decay': 0},],
            args.lr,
            momentum=0.9,
            nesterov=True,
        )
    else:
        optimizer = torch.optim.SGD(
            model.parameters(),
            args.lr,
            momentum=0.9,
            weight_decay=args.weight_decay,
            nesterov=True,
        )
    if optimizer_state_dict is not None:
        optimizer.load_state_dict(optimizer_state_dict)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.decay_period, gamma=args.gamma)

    train_criterion = utils.CrossEntropyLabelSmooth(1000, args.label_smooth).cuda()
    eval_criterion = nn.CrossEntropyLoss().cuda()
    return train_queue, valid_queue, model, train_criterion, eval_criterion, optimizer, scheduler

def child_train(train_queue, model, optimizer, global_step, arch_pool, arch_pool_prob, criterion, log_interval=100):
    objs = utils.AvgrageMeter()
    top1 = utils.AvgrageMeter()
    top5 = utils.AvgrageMeter()
    model.train()
    for step, (input, target) in enumerate(train_queue):
        input = utils.move_to_cuda(input)
        target = utils.move_to_cuda(target)

        optimizer.zero_grad()
        # sample an arch to train
        arch = utils.sample_arch(arch_pool, arch_pool_prob)
        logits = model(input, arch)
        loss = criterion(logits, target)
        loss.backward()
        optimizer.step()

        prec1, prec5 = utils.accuracy(logits, target, topk=(1, 5))
        n = input.size(0)
        objs.update(loss.data, n)
        top1.update(prec1.data, n)
        top5.update(prec5.data, n)

        global_step += 1

        if global_step % log_interval == 0:
            logging.info('Train %03d loss %e top1 %f top5 %f', global_step, objs.avg, top1.avg, top5.avg)
            logging.info('Arch: %s', ' '.join(map(str, arch)))

        if global_step >= args.max_num_updates:
            break

    return top1.avg, objs.avg, global_step

def child_valid(valid_queue, model, arch_pool, criterion, log_interval=1):
    valid_acc_list = []
    with torch.no_grad():
        model.eval()
        for i, arch in enumerate(arch_pool):
            # for step, (input, target) in enumerate(valid_queue):
            inputs, targets = next(iter(valid_queue))
            inputs = utils.move_to_cuda(inputs)
            targets = utils.move_to_cuda(targets)

            logits = model(inputs, arch, bn_train=True)
            loss = criterion(logits, targets)

            prec1, prec5 = utils.accuracy(logits, targets, topk=(1, 5))
            valid_acc_list.append(prec1.data/100)

            if (i+1) % log_interval == 0:
                logging.info('Valid arch %s\n loss %.2f top1 %f top5 %f', ' '.join(map(str, arch)), loss, prec1, prec5)

    return valid_acc_list

def controller_train(train_queue, model, optimizer):
    objs = utils.AvgrageMeter()
    mse = utils.AvgrageMeter()
    nll = utils.AvgrageMeter()
    model.train()
    for step, sample in enumerate(train_queue):
        encoder_input = utils.move_to_cuda(sample['encoder_input'])
        encoder_target = utils.move_to_cuda(sample['encoder_target'])
        decoder_input = utils.move_to_cuda(sample['decoder_input'])
        decoder_target = utils.move_to_cuda(sample['decoder_target'])

        optimizer.zero_grad()
        predict_value, log_prob, arch = model(encoder_input, decoder_input)
        loss_1 = F.mse_loss(predict_value.squeeze(), encoder_target.squeeze())
        loss_2 = F.nll_loss(log_prob.contiguous().view(-1, log_prob.size(-1)), decoder_target.view(-1))
        loss = args.controller_trade_off * loss_1 + (1 - args.controller_trade_off) * loss_2
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), args.controller_grad_clip)
        optimizer.step()

        n = encoder_input.size(0)
        objs.update(loss.data, n)
        mse.update(loss_1.data, n)
        nll.update(loss_2.data, n)

    return objs.avg, mse.avg, nll.avg

def controller_infer(queue, model, step, direction='+'):
    new_arch_list = []
    model.eval()
    for i, sample in enumerate(queue):
        encoder_input = utils.move_to_cuda(sample['encoder_input'])
        model.zero_grad()
        new_arch = model.generate_new_arch(encoder_input, step, direction=direction)
        new_arch_list.extend(new_arch.data.squeeze().tolist())
    return new_arch_list

def train_controller(model, train_input, train_target, epochs):
    logging.info('Train data: {}'.format(len(train_input)))
    dataset = utils.NAODataset(train_input, train_target, True)
    queue = torch.utils.data.DataLoader(
        dataset, batch_size=args.controller_batch_size, shuffle=True, pin_memory=True)
    optimizer = torch.optim.Adam(model.parameters(), lr=args.controller_lr, weight_decay=args.controller_l2_reg)
    for epoch in range(1, epochs+1):
        loss, mse, ce = controller_train(queue, model, optimizer)
        logging.info("epoch %04d train loss %.6f mse %.6f ce %.6f", epoch, loss, mse, ce)

def generate_synthetic_controller_data(model, exclude=[], maxn=1000, archs_selected=200):
    synthetic_input = []
    synthetic_target = []
    while len(synthetic_input) < maxn:
        synthetic_arch = utils.generate_arch(1, args.layers, args.num_ops)[0]
        if synthetic_arch not in exclude and synthetic_arch not in synthetic_input:
            synthetic_input.append(synthetic_arch)

    synthetic_dataset = utils.ControllerDataset(synthetic_input, None, False)
    synthetic_queue = torch.utils.data.DataLoader(synthetic_dataset, batch_size=len(synthetic_dataset), shuffle=False, pin_memory=True)

    # with torch.no_grad():
    model.eval()
    grads = list()
    for sample in synthetic_queue:
        input = utils.move_to_cuda(sample['encoder_input'])
        encoder_outputs, _, _, predict_value = model.encoder(input)
        gradients = torch.autograd.grad(predict_value, encoder_outputs, torch.ones_like(predict_value))[0]
        synthetic_target += predict_value.data.squeeze().tolist()
        grads += [torch.norm(g).tolist() for g in gradients]

    assert len(synthetic_input) == len(synthetic_target)
    assert len(synthetic_input) == len(grads)

    synthetic_indices = np.argsort(grads)
    synthetic_input = [synthetic_input[i] for i in synthetic_indices][:archs_selected]
    synthetic_target = [synthetic_target[i] for i in synthetic_indices][:archs_selected]

    return synthetic_input, synthetic_target

def main():
    if not torch.cuda.is_available():
        logging.info('no gpu device available')
        sys.exit(1)

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    cudnn.enabled = True
    cudnn.benchmark = True
    cudnn.deterministic = True

    args.device_count = torch.cuda.device_count() if torch.cuda.is_available() else 1
    args.lr = args.lr
    args.batch_size = args.batch_size
    args.eval_batch_size = args.eval_batch_size
    args.width_stages = [int(val) for val in args.width_stages.split(',')]
    args.n_cell_stages = [int(val) for val in args.n_cell_stages.split(',')]
    args.stride_stages = [int(val) for val in args.stride_stages.split(',')]
    args.num_class = 1000
    args.num_ops = len(utils.OPERATIONS)
    args.controller_vocab_size = 1 + args.num_ops
    args.controller_source_length = args.layers
    args.controller_encoder_length = args.layers
    args.controller_decoder_length = args.controller_source_length

    logging.info("args = %s", args)

    if args.arch_pool is not None:
        logging.info('Architecture pool is provided, loading')
        with open(args.arch_pool) as f:
            archs = f.read().splitlines()
            archs = list(map(lambda x:list(map(int, x.strip().split())), archs))
            child_arch_pool = archs
    else:
        child_arch_pool = None

    train_queue, valid_queue, model, train_criterion, eval_criterion, optimizer, scheduler = build_imagenet(None, None, valid_num=5000, epoch=-1)

    controller = NAO(
        args.controller_encoder_layers,
        args.controller_mlp_layers,
        args.controller_decoder_layers,
        args.controller_vocab_size,
        args.controller_hidden_size,
        args.controller_mlp_hidden_size,
        args.controller_dropout,
        args.controller_encoder_length,
        args.controller_source_length,
        args.controller_decoder_length,
    )
    controller = controller.cuda()
    logging.info("Encoder-Predictor-Decoder param size = %d", utils.count_parameters(controller))

    if child_arch_pool is None:
        logging.info('Architecture pool is not provided, randomly generating now')
        child_arch_pool = utils.generate_arch(args.controller_seed_arch, args.layers, args.num_ops)

    arch_pool = []
    arch_pool_valid_acc = []
    archs_selected = args.select_narchs
    for controller_iteration in range(args.controller_iterations+1):
        logging.info('Iteration %d', controller_iteration+1)

        child_arch_pool_prob = None
        num_updates = 0
        max_num_updates = args.max_num_updates
        epoch = 1
        while True:
            lr = scheduler.get_lr()[0]
            logging.info('epoch %d lr %e', epoch, lr)
            # sample an arch to train
            train_acc, train_obj, num_updates = child_train(train_queue, model, optimizer, num_updates, child_arch_pool, child_arch_pool_prob, train_criterion)
            epoch += 1
            scheduler.step()
            if num_updates >= max_num_updates:
                break

        logging.info("Evaluate seed archs")
        arch_pool += child_arch_pool
        arch_pool_valid_acc = child_valid(valid_queue, model, arch_pool, eval_criterion)

        arch_pool_valid_acc_sorted_indices = np.argsort(arch_pool_valid_acc)[::-1]
        arch_pool = list(map(lambda x:arch_pool[x], arch_pool_valid_acc_sorted_indices))
        arch_pool_valid_acc = list(map(lambda x:arch_pool_valid_acc[x], arch_pool_valid_acc_sorted_indices))
        with open(os.path.join(args.output_dir, 'arch_pool.{}'.format(controller_iteration)), 'w') as fa:
            with open(os.path.join(args.output_dir, 'arch_pool.perf.{}'.format(controller_iteration)), 'w') as fp:
                for arch, perf in zip(arch_pool, arch_pool_valid_acc):
                    arch = ' '.join(map(str, arch))
                    fa.write('{}\n'.format(arch))
                    fp.write('{}\n'.format(perf))
        if controller_iteration == args.controller_iterations:
            break

        # Train Encoder-Predictor-Decoder
        logging.info('Train Encoder-Predictor-Decoder')
        inputs = arch_pool
        min_val = min(arch_pool_valid_acc)
        max_val = max(arch_pool_valid_acc)
        targets = list(map(lambda x: (x - min_val) / (max_val - min_val), arch_pool_valid_acc))

        # Pre-train
        logging.info('Pre-train EPD')
        train_controller(controller, inputs, targets, args.controller_pretrain_epochs)
        logging.info('Finish pre-training EPD')
        # Generate synthetic data
        logging.info('Generate synthetic data for EPD')
        archs_selected *= 2
        synthetic_inputs, \
        synthetic_targets = generate_synthetic_controller_data(controller,
                                                               inputs,
                                                               args.controller_random_arch,
                                                               archs_selected)
        if args.controller_up_sample_ratio:
            all_inputs = inputs * args.controller_up_sample_ratio + synthetic_inputs
            all_targets = targets * args.controller_up_sample_ratio + synthetic_targets
        else:
            all_inputs = inputs + synthetic_inputs
            all_targets = targets + synthetic_targets
        # Train
        logging.info('Train EPD')
        train_controller(controller, all_inputs, all_targets, args.controller_epochs)
        logging.info('Finish training EPD')

        # Generate new archs
        new_archs = []
        max_step_size = 100
        predict_step_size = 0.0
        topk_indices = np.argsort(all_targets)[:100]
        topk_archs = list(map(lambda x:all_inputs[x], topk_indices))
        infer_dataset = utils.ControllerDataset(topk_archs, None, False)
        infer_queue = torch.utils.data.DataLoader(
            infer_dataset, batch_size=len(infer_dataset), shuffle=False, pin_memory=True)
        while len(new_archs) < args.controller_new_arch:
            predict_step_size += 0.1
            logging.info('Generate new architectures with step size %.2f', predict_step_size)
            new_arch = controller_infer(infer_queue, controller, predict_step_size, direction='+')
            for arch in new_arch:
                if arch not in inputs and arch not in new_archs:
                    new_archs.append(arch)
                if len(new_archs) >= args.controller_new_arch:
                    break
            logging.info('%d new archs generated now', len(new_archs))
            if predict_step_size > max_step_size:
                break

        child_arch_pool = new_archs
        logging.info("Generate %d new archs", len(child_arch_pool))

    logging.info('Finish Searching')

    top_archs = arch_pool[:5]
    top_archs_perf = arch_pool_valid_acc[:5]
    with open(os.path.join(args.output_dir, 'arch_pool.final'), 'w') as fa:
        with open(os.path.join(args.output_dir, 'arch_pool.perf.final'), 'w') as fp:
            for arch, perf in zip(top_archs, top_archs_perf):
                arch = ' '.join(map(str, arch))
                fa.write('{}\n'.format(arch))
                fp.write('{}\n'.format(perf))

if __name__ == '__main__':
    main()


from collections import OrderedDict, namedtuple
import datetime as dt

import numpy as np

from .. import setup_logging
from ..utils.utils import dist, geo_dist, pairwise, find_extrema

log = setup_logging.get_logger('st.tracking')

NUM_ENSEMBLE_MEMBERS = 56

class VortmaxNearestNeighbourTracker(object):
    '''Simple nearest neighbour tracker

    Assumes that the two nearest points from one timestep to another belong to the same track.
    '''
    def __init__(self):
        self.use_geo_dist = True

        if self.use_geo_dist:
            self.dist = geo_dist
            self.dist_cutoff = geo_dist((0, 0), (2, 0)) * 8
        else:
            self.dist = dist
            self.dist_cutoff = 5

    def _construct_vortmax_tracks_by_date(self, all_vortmax_time_series):
        self.all_vort_tracks_by_date = []
        for ensemble_member in range(NUM_ENSEMBLE_MEMBERS):
            vortmax_time_series = all_vortmax_time_series[ensemble_member]
            vortmax_tracks = []
            vort_tracks_by_date = OrderedDict()
            # Find all start vortmaxes (those that have no previous vortmaxes)
            # and use these to generate a track.
            for vortmaxes in vortmax_time_series.values():
                for vortmax in vortmaxes:
                    if len(vortmax.prev_vortmax) == 0:
                        vortmax_track = VortMaxTrack(vortmax, ensemble_member)
                        if len(vortmax_track.vortmaxes) >= 6:
                            vortmax_tracks.append(vortmax_track)

                            for date in vortmax_track.vortmax_by_date.keys():
                                if date not in vort_tracks_by_date:
                                    vort_tracks_by_date[date] = []
                                vort_tracks_by_date[date].append(vortmax_track)

                    # Allows vort_tracks_by_date to be serialized.
                    vortmax.next_vortmax = []
                    vortmax.prev_vortmax = []

            self.all_vort_tracks_by_date.append(vort_tracks_by_date)
        return self.all_vort_tracks_by_date

    def track(self, df):
        all_vortmax_time_series = []
        for ensemble_member in range(NUM_ENSEMBLE_MEMBERS):
            vortmax_time_series = OrderedDict()
            all_vortmax_time_series.append(vortmax_time_series)
            for i in df[df.em == ensemble_member].index:
                row = df.loc[i]
                if row.date not in vortmax_time_series:
                    vortmax_time_series[row.date] = []
                vortmax = VortMax(row.date, (row.lon, row.lat), row.vort)
                vortmax_time_series[vortmax.date].append(vortmax)
        print('Made objects')
        return self.track_vort_maxima(all_vortmax_time_series)

    def track_vort_maxima(self, all_vortmax_time_series):
        '''Uses a generated list of vortmaxes to track them from one timestep to another

        :param vortmax_time_series: dict of vortmaxes
        :returns: self.vort_tracks_by_date (OrderedDict)
        '''
        for ensemble_member in range(NUM_ENSEMBLE_MEMBERS):
            log.info('Tracking vortmaxima: {0}'.format(ensemble_member))
            vortmax_time_series = all_vortmax_time_series[ensemble_member]
            # Loops over 2 lists of vortmaxes
            for vs1, vs2 in pairwise(vortmax_time_series.values()):
                for v1 in vs1:
                    min_dist = self.dist_cutoff
                    v2next = None

                    # Find the nearest vortmax in the next timestep.
                    for v2 in vs2:
                        d = self.dist(v1.pos, v2.pos)
                        if d < min_dist:
                            min_dist = d
                            v2next = v2

                    # Add the nearest vortmax in the next timestep.
                    if v2next:
                        v1.add_next(v2next)
                        if len(v1.next_vortmax) != 1:
                            raise Exception('There should only ever be one next_vortmax')

            # Some vortmaxes may have more than one previous vortmax.
            # Find these and choose the nearest one as the actual previous.
            for vs in vortmax_time_series.values():
                for v in vs:
                    if len(v.prev_vortmax) > 1:
                        min_dist = self.dist_cutoff
                        vprev = None
                        for pv in v.prev_vortmax:
                            d = self.dist(pv.pos, v.pos)
                            if d < min_dist:
                                min_dist = d
                                vprev = pv

                        for pv in v.prev_vortmax:
                            if pv != vprev:
                                pv.next_vortmax.remove(v)

                        v.prev_vortmax = [vprev]

        return self._construct_vortmax_tracks_by_date(all_vortmax_time_series)

class FieldFinder(object):
    def __init__(self, c20data, all_vort_tracks_by_date):
        self.c20data = c20data
        self.all_vort_tracks_by_date = all_vort_tracks_by_date
        self.all_cyclone_tracks = []
        for ensemble_member in range(NUM_ENSEMBLE_MEMBERS):
            self.all_cyclone_tracks.append({})

    def collect_fields(self, start_date, end_date):
        # import ipdb; ipdb.set_trace()
        index = np.where(self.c20data.dates == start_date)[0][0]
        end_index = np.where(self.c20data.dates == end_date)[0][0]

        while index <= end_index:
            date = self.c20data.dates[index]
            log.info('Collecting fields for: {0}'.format(date))
            self.c20data.set_date(date)
            for ensemble_member in range(NUM_ENSEMBLE_MEMBERS):
                cyclone_tracks = self.all_cyclone_tracks[ensemble_member]

                if date not in self.all_vort_tracks_by_date[ensemble_member]:
                    continue

                vort_tracks = self.all_vort_tracks_by_date[ensemble_member][date]
                for vort_track in vort_tracks:
                    cyclone_track = None
                    if vort_track not in cyclone_tracks:
                        if len(vort_track.vortmaxes) >= 6:
                            cyclone_track = CycloneTrack(vort_track, ensemble_member)
                            cyclone_tracks[vort_track] = cyclone_track
                    else:
                        cyclone_track = cyclone_tracks[vort_track]

                    if cyclone_track:
                        self.add_fields_to_track(date, cyclone_track, ensemble_member)
            index += 1

    def add_fields_to_track(self, date, cyclone_track, ensemble_member):
        actual_vmax_pos = cyclone_track.get_vmax_pos(date)
        # Round values to the nearest multiple of 2 (vmax_pos can come from an interpolated field)
        vmax_pos = tuple([int(round(p / 2.)) * 2 for p in actual_vmax_pos])
        lon_index = np.where(self.c20data.lons == vmax_pos[0])[0][0]
        lat_index = np.where(self.c20data.lats == vmax_pos[1])[0][0]

        min_lon, max_lon = lon_index - 5, lon_index + 6
        min_lat, max_lat = lat_index - 5, lat_index + 6
        local_slice = (slice(min_lat, max_lat), slice(min_lon, max_lon))

        local_prmsl = self.c20data.prmsl[ensemble_member][local_slice].copy()

	# 9950 pressure level.
        local_windspeed = np.sqrt(self.c20data.u9950[ensemble_member][local_slice] ** 2 +
                                  self.c20data.v9950[ensemble_member][local_slice] ** 2)
        max_windspeed_pos = np.unravel_index(np.argmax(local_windspeed), (11, 11))
        max_windspeed = local_windspeed[max_windspeed_pos]

        lon = self.c20data.lons[min_lon + max_windspeed_pos[1]]
        lat = self.c20data.lats[min_lat + max_windspeed_pos[0]]

        cyclone_track.max_windspeeds[date] = max_windspeed
        cyclone_track.max_windspeed_positions[date] = (lon, lat)

        e, index_pmaxs, index_pmins = find_extrema(local_prmsl)
        min_dist = 1000
        pmin = None
        pmin_pos = None
        for index_pmin in index_pmins:
            lon = self.c20data.lons[min_lon + index_pmin[1]]
            lat = self.c20data.lats[min_lat + index_pmin[0]]

            local_pmin = local_prmsl[index_pmin[0], index_pmin[1]]
            local_pmin_pos = (lon, lat)
            dist = geo_dist(actual_vmax_pos, local_pmin_pos)
            if dist < min_dist:
                min_dist = dist
                pmin = local_pmin
                pmin_pos = local_pmin_pos

        cyclone_track.min_dists[date] = min_dist
        cyclone_track.pmins[date] = pmin
        cyclone_track.pmin_positions[date] = pmin_pos

        if pmin:
            cyclone_track.p_ambient_diffs[date] = local_prmsl.mean() - pmin
        else:
            cyclone_track.p_ambient_diffs[date] = None

        cyclone_track.t850s[date] = self.c20data.t850[ensemble_member][lat_index, lon_index]
        cyclone_track.t995s[date] = self.c20data.t995[ensemble_member][lat_index, lon_index]
        cyclone_track.capes[date] = self.c20data.cape[ensemble_member][lat_index, lon_index]
        cyclone_track.pwats[date] = self.c20data.pwat[ensemble_member][lat_index, lon_index]
        # No longer using due to it not having much
        # discriminatory power and space constraints.
        cyclone_track.rh995s[date] = 0


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.shortcuts import render
from django.http import HttpResponse
from forms import SpeakerForm, AuthForm
from models import SpeakerModel, AuthModel
from scripts import ModuleML
import numpy as np
# Create your views here.

def input(request):
    s = SpeakerForm()
    return render(request, 'train.html', {'form': s})

def train(request):
    if request.method == 'POST':
        form = SpeakerForm(request.POST, request.FILES)
        if form.is_valid():
            name_val = form.cleaned_data['name']
            model = ModuleML()
            train_x, train_y, label = model.train_model(request.FILES['audio'])
            model.GBM_run()
            s = SpeakerModel()
            s.audio = 0
            s.name = name_val
            s.label = label
            s.save()
            return HttpResponse('Training successful')
        else:
            s = SpeakerForm()
            render(request, 'train.html', {'form': s})

    else:
        s = SpeakerForm()
        render(request, 'train.html', {'form': s})

def authenticate(request):
    list = SpeakerModel.objects.all()
    auth = AuthForm()
    return render(request, 'authenticate.html', {'form': auth , 'items': list})

def result(request):
    if request.method == 'POST':
        form = AuthForm(request.POST, request.FILES)
        if form.is_valid():
            label = int(request.POST['pick'])
            model = ModuleML()
            val, pred, conf = model.predict(request.FILES['audio'], label)
            if pred == label:
                st = "Match"
            else:
                st = "No Match"

            person_true = SpeakerModel.objects.get(label= pred)
        else:
            list = SpeakerModel.objects.all()
            auth = AuthForm()
            return render(request, 'authenticate.html', {'form': auth, 'items': list})
    else:
        list = SpeakerModel.objects.all()
        auth = AuthForm()
        return render(request, 'authenticate.html', {'form': auth, 'items': list})
    return HttpResponse(st + '\n Confidence Score:\t'+ str(conf))


"""
Perform detection using models created with the Tensorflow Object Detection API.
https://github.com/tensorflow/models/tree/master/research/object_detection
"""

import cv2
import numpy as np
import os

with open(os.getenv('TFODA_CLASSES_PATH'), 'r') as classes_file:
    classes = dict(enumerate([line.strip() for line in classes_file.readlines()]))
with open(os.getenv('TFODA_CLASSES_OF_INTEREST_PATH'), 'r') as coi_file:
    classes_of_interest = tuple([line.strip() for line in coi_file.readlines()])

def get_bounding_boxes(image):
    # create model using weights and config
    net = cv2.dnn.readNetFromTensorflow(os.getenv('TFODA_WEIGHTS_PATH'), os.getenv('TFODA_CONFIG_PATH'))

    # create image blob
    blob = cv2.dnn.blobFromImage(image, size=(300, 300), swapRB=True, crop=False)

    # detect objects
    net.setInput(blob)
    detections = net.forward()

    # get bounding boxes
    confidence_threshold = float(os.getenv('TFODA_CONFIDENCE_THRESHOLD'))
    boxes = []
    _classes = []
    _confidences = []
    rows, cols, _ = image.shape
    for detection in detections[0, 0]:
        confidence = float(detection[2])
        class_id = int(detection[1])
        if confidence > confidence_threshold and class_id in classes and classes[class_id] in classes_of_interest:
            left = int(detection[3] * cols)
            top = int(detection[4] * rows)
            right = int(detection[5] * cols)
            bottom = int(detection[6] * rows)

            x, y, w, h = left, top, right - left, bottom - top
            boxes.append([x, y, w, h])
            _classes.append(classes[class_id])
            _confidences.append(confidence)

    # remove overlapping bounding boxes
    nms_threshold = 0.4
    indices = cv2.dnn.NMSBoxes(boxes, _confidences, confidence_threshold, nms_threshold)

    _bounding_boxes = []
    for i in indices:
        i = i[0]
        _bounding_boxes.append(boxes[i])

    return _bounding_boxes, _classes, _confidences


#! /usr/bin/env python
"""
source.py
A source contains one or more Objects that are spatially coincident, each
containing a single color lightcurve of a ZTF object.
"""
import os
import pickle
import numpy as np
from astropy.coordinates import SkyCoord
import astropy.units as u

from zort.object import Object
from zort.plot import plot_objects
from zort.utils import return_filename, return_objects_map_filename, \
    return_radec_map_filename

################################
#                              #
#  Source Class                #
#                              #
################################

class Source:
    """
    A source contains one or more Objects that are spatially coincident, each
    containing a single color lightcurve of a ZTF object. All objects within
    source must be from the same lightcurve file.d
    """

    def __init__(self, filename, object_id_g=None,
                 object_id_r=None, object_id_i=None,
                 lightcurve_position_g=None,
                 lightcurve_position_r=None,
                 lightcurve_position_i=None,
                 apply_catmask=True, PS_g_minus_r=0,
                 objects_map=None,
                 radec_map=None,
                 lightcurve_file_pointer=None,
                 check_initialization=True):
        # Load filenames and check for existence
        self.filename = return_filename(filename)
        self._objects_map_filename = None
        self._radec_map_filename = None

        if check_initialization:
            self._check_initialization(object_id_g, object_id_r, object_id_i,
                                       lightcurve_position_g,
                                       lightcurve_position_r,
                                       lightcurve_position_i)

        self.apply_catmask = apply_catmask
        self.PS_g_minus_r = PS_g_minus_r
        if objects_map:
            self.objects_map = objects_map
        else:
            self.objects_map = None
        if radec_map:
            self.radec_map = radec_map
        else:
            self.radec_map = None

        self.lightcurve_file_pointer = lightcurve_file_pointer
        objects = self._load_objects(object_id_g, object_id_r, object_id_i,
                                     lightcurve_position_g,
                                     lightcurve_position_r,
                                     lightcurve_position_i)
        self.object_g = objects[0]
        self.object_r = objects[1]
        self.object_i = objects[2]
        self.objects = [o for o in objects if o is not None]

        self._radec = None
        self._glonlat = None

    @property
    def objects_map_filename(self):
        if self._objects_map_filename is None:
            self._objects_map_filename = return_objects_map_filename(self.filename)
        return self._objects_map_filename

    @property
    def radec_map_filename(self):
        if self._radec_map_filename is None:
            self._radec_map_filename = return_radec_map_filename(self.filename)
        return self._radec_map_filename

    @property
    def radec(self):
        if self._radec is None:
            ra = np.mean([obj.ra for obj in self.objects if obj is not None])
            dec = np.mean([obj.dec for obj in self.objects if obj is not None])
            self._radec = [ra, dec]
        return self._radec

    @property
    def ra(self):
        return self.radec[0]

    @property
    def dec(self):
        return self.radec[1]

    @property
    def glonlat(self):
        if self._glonlat is None:
            coord = SkyCoord(self.ra, self.dec, unit=u.degree, frame='icrs')
            glon, glat = coord.galactic.l.value, coord.galactic.b.value
            if glon > 180:
                glon -= 360
            self._glonlat = (glon, glat)
        return self._glonlat

    @property
    def glon(self):
        return self.glonlat[0]

    @property
    def glat(self):
        return self.glonlat[1]

    def _return_object_print_info(self, object):
        if object is None:
            object_id, nepochs = None, None
        else:
            object_id, nepochs = object.object_id, object.nepochs
        return {'id': object_id, 'nepochs': nepochs}

    def __repr__(self):
        title = 'Filename: %s\n' % self.filename.split('/')[-1]
        title += 'Object-g ID: {id} | nepochs: {nepochs}\n'.format(
            **self._return_object_print_info(self.object_g))
        title += 'Object-r ID: {id} | nepochs: {nepochs}\n'.format(
            **self._return_object_print_info(self.object_r))
        title += 'Object-i ID: {id} | nepochs: {nepochs}\n'.format(
            **self._return_object_print_info(self.object_i))
        title += 'Ra/Dec: (%.5f, %.5f)\n' % (self.ra, self.dec)

        return title

    def _check_initialization(self,
                              object_id_g, object_id_r, object_id_i,
                              lightcurve_position_g,
                              lightcurve_position_r,
                              lightcurve_position_i):
        arr = np.array([object_id_g, object_id_r, object_id_i,
                        lightcurve_position_g, lightcurve_position_r,
                        lightcurve_position_i])
        if np.all(arr == None):
            raise Exception('Source must be instantiated '
                            'with at least one Object ID '
                            'or lightcurve position.')

        if object_id_g is not None and lightcurve_position_g is not None:
            raise Exception('Only initialize g object with Object ID '
                            'or lightcurve position, but not both.')

        if object_id_r is not None and lightcurve_position_r is not None:
            raise Exception('Only initialize r object with Object ID '
                            'or lightcurve position, but not both.')

        if object_id_i is not None and lightcurve_position_i is not None:
            raise Exception('Only initialize i object with Object ID '
                            'or lightcurve position, but not both.')

    def _load_object(self, object_id, lightcurve_position, color):
        if object_id is None and lightcurve_position is None:
            return None

        if lightcurve_position is not None:
            obj = Object(self.filename,
                         lightcurve_position=lightcurve_position,
                         objects_map=self.objects_map,
                         apply_catmask=self.apply_catmask,
                         lightcurve_file_pointer=self.lightcurve_file_pointer)
        else:
            if self.objects_map is None:
                self.objects_map = self.load_objects_map()
            obj = Object(self.filename,
                         object_id=object_id,
                         objects_map=self.objects_map,
                         lightcurve_file_pointer=self.lightcurve_file_pointer)
        if obj.color != color:
            str = "Color of 'object_id_{color}' is {color_obj}. " \
                  "Must be {color}.".format(color=color,
                                            color_obj=obj.color)
            raise Exception(str)
        return obj

    def _load_objects(self, object_id_g, object_id_r, object_id_i,
                      lightcurve_position_g, lightcurve_position_r, lightcurve_position_i):
        object_g = self._load_object(object_id=object_id_g,
                                     lightcurve_position=lightcurve_position_g,
                                     color='g')
        object_r = self._load_object(object_id=object_id_r,
                                     lightcurve_position=lightcurve_position_r,
                                     color='r')
        object_i = self._load_object(object_id=object_id_i,
                                     lightcurve_position=lightcurve_position_i,
                                     color='i')
        return object_g, object_r, object_i

    def load_objects_map(self):
        objects_map_filename = self.objects_map_filename
        objects_map = pickle.load(open(objects_map_filename, 'rb'))
        return objects_map

    def load_radec_map(self):
        radec_map_filename = self.radec_map_filename
        radec_map = pickle.load(open(radec_map_filename, 'rb'))
        return radec_map

    def plot_lightcurves(self, filename=None, insert_radius=30,
                         model_params=None):
        if filename is None:
            object_ids = '_'.join([str(obj.object_id) for obj in self.objects
                                   if obj is not None])
            filename = '%s_%s.png' % (self.filename.replace('.txt', ''),
                                      object_ids)
        print('plot_lightcurves from source')
        plot_objects(filename=filename,
                     object_g=self.object_g,
                     object_r=self.object_r,
                     object_i=self.object_i,
                     insert_radius=insert_radius,
                     model_params=model_params)

def create_source_from_object(object, locate_siblings=True, radius_as=2, skip_filterids=None):
    if locate_siblings and object.siblings is None:
        object.locate_siblings(radius_as=radius_as,
                               skip_filterids=skip_filterids)
    source_dict = {'g': None, 'r': None, 'i': None,
                   object.color: object.object_id}
    for sibling in object.siblings:
        source_dict[sibling.color] = sibling.object_id

    return Source(filename=object.filename,
                  object_id_g=source_dict['g'],
                  object_id_r=source_dict['r'],
                  object_id_i=source_dict['i'])

def save_sources(filename, sources, overwrite=False):
    if os.path.exists(filename) and not overwrite:
        print('%s already exists, exiting without saving objects. '
              'Set overwrite=True to enable writing over existing '
              'object lists.' % filename)
        return None

    with open(filename, 'w') as f:
        for source in sources:
            f.write('%s,%s,%s,%s\n' % (source.filename,
                                       source.object_id_g,
                                       source.object_id_r,
                                       source.object_id_i))

def load_sources(filename):
    sources = []
    for line in open(filename, 'r'):
        filename, object_id_g, object_id_r, object_id_i = line.replace('\n', '').split(',')
        sources.append(Source(filename,
                              object_id_g=object_id_g,
                              object_id_r=object_id_r,
                              object_id_i=object_id_i))

    return sources

import numpy,import astropy,from astropy


# import the necessary packages
from PIL import Image
import pickle
import cv2
import numpy as np
import imutils
from keras.models import load_model
import base64
import io
from sklearn.neighbors import KNeighborsClassifier
import urllib
import certifi
from keras import backend as K

# MODEL_FILENAME = "model/captcha_model2.hdf5"
# MODEL_FILENAME2 = 'model/NN_model.sav'
MODEL_FILENAME = "model/captcha_model_final.hdf5"
MODEL_FILENAME2 = 'model/NN_model2.sav'

MODEL_LABELS_FILENAME = "model/model_labels2.dat"

def predict_captcha_text(image_file,img_type, mod='keras'):
    # Load up the model labels (so we can translate model predictions to actual letters)
    with open(MODEL_LABELS_FILENAME, "rb") as f:
        lb = pickle.load(f)

    # Load the trained neural network
    if mod=='keras':
        model = load_model(MODEL_FILENAME)
    else:
        print('scikit')
        model = pickle.load(open(MODEL_FILENAME2, 'rb'))

    image, last, orig_img = pre_process(image_file,img_type)

    # find the contours (continuous blobs of pixels) the image
    contours = cv2.findContours(last.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Hack for compatibility with different OpenCV versions
    contours = contours[0] if imutils.is_cv2() else contours[1]

    letter_image_regions = []

    # Now we can loop through each of the four contours and extract the letter
    # inside of each one
    for contour in contours:
        if len(contour) < 9:
            continue
        else:
            # Get the rectangle that contains the contour
            (x, y, w, h) = cv2.boundingRect(contour)

            # Compare the width and height of the contour to detect letters that
            # are conjoined into one chunk
            #         print (w,h,w/h)
            if w / h > 2:
                triple_width = int(w / 3)
                letter_image_regions.append((x, y, triple_width, h))
                letter_image_regions.append((x + triple_width, y, triple_width, h))
                letter_image_regions.append((x + 2 * triple_width, y, triple_width, h))

            elif w / h > 1.25:
                # This contour is too wide to be a single letter!
                # Split it in half into two letter regions!
                half_width = int(w / 2)
                letter_image_regions.append((x, y, half_width, h))
                letter_image_regions.append((x + half_width, y, half_width, h))
            else:
                # This is a normal letter by itself
                letter_image_regions.append((x, y, w, h))

    letter_image_regions = sorted(letter_image_regions, key=lambda x: x[0])

    # Create an output image and a list to hold our predicted letters
    output = cv2.merge([image] * 3)
    predictions = []

    if len(letter_image_regions) > 5:
        prev = None
        letter_image_regions2 = []
        for lir in letter_image_regions:

            if prev:
                if lir[2] < 15:
                    letter_image_regions2.append((prev[0], prev[1], prev[2] + lir[2], 25))
                if lir[2] > 15:
                    letter_image_regions2.append(prev)
                    letter_image_regions2.append(lir)
                prev = None
            else:
                if lir[2] < 15:
                    prev = lir
                else:
                    prev = None
                    letter_image_regions2.append(lir)
    else:
        letter_image_regions2 = letter_image_regions

    # loop over the lektters
    for i, letter_bounding_box in enumerate(letter_image_regions2):
        # Grab the coordinates of the letter in the image
        if i > 4:
            break
        x, y, w, h = letter_bounding_box
        if x < 2:
            x = 2
        if y < 2:
            y = 2
        # Extract the letter from the original image with a 2-pixel margin around the edge
        letter_image = image[y - 2:y + h + 2, x - 2:x + w + 2]
        #         print('shape',letter_image.shape)
        # Re-size the letter image to 20x20 pixels to match training data
        letter_image = resize_to_fit(letter_image, 20, 20)

        # Turn the single image into a 4d list of images to make Keras happy
        letter_image = np.expand_dims(letter_image, axis=2)
        letter_image = np.expand_dims(letter_image, axis=0)

        # Ask the neural network to make a prediction
        if mod=='keras':
            prediction = model.predict(letter_image)
        else:

            prediction = model.predict(letter_image.reshape(1, 400) / 255.0)

        # Convert the one-hot-encoded prediction back to a normal letter
        letter = lb.inverse_transform(prediction)[0]
        predictions.append(letter)
    K.clear_session()
    captcha_text = "".join(predictions)
#    cv2.imwrite('static/uploads/'+captcha_text+'.jpg',orig_img)
    return captcha_text

def base64toImg(b64_string):

    # reconstruct image as an numpy array
    img = io.BytesIO(base64.b64decode(b64_string))
    img = Image.open(img)
    # finally convert RGB image to BGR for opencv
    # and save result
    cv2_img = cv2.cvtColor(np.array(img), cv2.COLOR_BGR2RGB) #np.array(img)

    return cv2_img

def resize_to_fit(image, width, height):
    """
    A helper function to resize an image to fit within a given size
    :param image: image to resize
    :param width: desired width in pixels
    :param height: desired height in pixels
    :return: the resized image
    """

    # grab the dimensions of the image, then initialize
    # the padding values
    (h, w) = image.shape[:2]

    # if the width is greater than the height then resize along
    # the width
    if w > h:
        image = imutils.resize(image, width=width)

    # otherwise, the height is greater than the width so resize
    # along the height
    else:
        image = imutils.resize(image, height=height)

    # determine the padding values for the width and height to
    # obtain the target dimensions
    padW = int((width - image.shape[1]) / 2.0)
    padH = int((height - image.shape[0]) / 2.0)

    # pad the image then apply one more resizing to handle any
    # rounding issues
    image = cv2.copyMakeBorder(image, padH, padH, padW, padW,
                               cv2.BORDER_REPLICATE)
    image = cv2.resize(image, (width, height))

    # return the pre-processed image
    return image

def url_to_image(url):
    # download the image, convert it to a NumPy array, and then read
    # it into OpenCV format

    resp = urllib.request.urlopen(url, cafile=certifi.where())#urllib.urlopen(url)
    image = np.asarray(bytearray(resp.read()), dtype="uint8")
    image = cv2.imdecode(image, cv2.IMREAD_COLOR)

    # return the image
    return image

def pre_process(image_file,img_type):
    # image1 = cv2.imread(image_file)
    # image1 = url_to_image(image_file)
    if img_type=='file':
        image1 = cv2.imread(image_file)
    elif img_type=='url':
        image1 = url_to_image(image_file)
    else:
        image1 = base64toImg(image_file)

    # load the example image and convert it to grayscale
    image = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)
    image[0:10, :] = 255
    image[40:, :] = 255
    image[:, 136:] = 255

    # threshold the image (convert it to pure black and white)
    thresh = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)[1]
    kernel = np.ones((3, 3), np.uint8)

    new = thresh.copy()
    new[new == 255] = 1
    hist = np.sum(new, axis=0)
    hist[hist < 7] = 0
    k = np.where(hist == 0)
    last2 = thresh.copy()
    last2[:, k] = 0

    s = 255 - cv2.erode(last2, kernel, iterations=1)
    s = cv2.medianBlur(s, 3)

    return s, last2, image1


# Copyright 2020 The DDSP Authors.
#
# 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.

# Lint as: python3
r"""Apache Beam pipeline for computing TFRecord dataset of synthetic examples.

Example usage:
==============
ddsp_generate_synthetic_dataset \
--output_tfrecord_path=/tmp/synthetic_data.tfrecord \
--num_shards=1 \
--gin_param="generate_examples.generate_fn = @generate_notes_v2" \
--num_examples=100 \
--alsologtostderr

For the ICML workshop paper, we created 10,000,000 examples in 1000 shards.
"""

from absl import app
from absl import flags
import apache_beam as beam
from ddsp.training.data_preparation import synthetic_data  # pylint:disable=unused-import
import gin
import numpy as np
import pkg_resources
import tensorflow.compat.v2 as tf

GIN_PATH = pkg_resources.resource_filename(__name__, '../gin')

FLAGS = flags.FLAGS

flags.DEFINE_string(
    'output_tfrecord_path', None,
    'The prefix path to the output TFRecord. Shard numbers will be added to '
    'actual path(s).')
flags.DEFINE_integer(
    'num_shards', None,
    'The number of shards to use for the TFRecord. If None, this number will '
    'be determined automatically.')
flags.DEFINE_integer(
    'num_examples', 1000000,
    'The total number of synthetic examples to generate.')
flags.DEFINE_integer(
    'random_seed', 42,
    'Random seed to use for deterministic generation.')
flags.DEFINE_list(
    'pipeline_options', '--runner=DirectRunner',
    'A comma-separated list of command line arguments to be used as options '
    'for the Beam Pipeline.')

# Gin config flags.
flags.DEFINE_multi_string('gin_search_path', [],
                          'Additional gin file search paths.')
flags.DEFINE_multi_string('gin_file', [], 'List of paths to the config files.')
flags.DEFINE_multi_string('gin_param', [],
                          'Newline separated list of Gin parameter bindings.')

class GenerateExampleFn(beam.DoFn):
  """Gin-configurable wrapper to generate synthetic examples."""

  def __init__(self, gin_str, **kwargs):
    super().__init__(**kwargs)
    self._gin_str = gin_str

  def start_bundle(self):
    with gin.unlock_config():
      gin.parse_config(self._gin_str)

  @gin.configurable('generate_examples')
  def process(self, seed, generate_fn=gin.REQUIRED):
    np.random.seed(seed)
    batch = generate_fn(n_batch=1)
    beam.metrics.Metrics.counter('GenerateExampleFn', 'generated').inc()
    yield {k: v[0].numpy() for k, v in batch.items()}

def _float_dict_to_tfexample(float_dict):
  """Convert dictionary of float arrays to tf.train.Example proto."""
  return tf.train.Example(
      features=tf.train.Features(
          feature={
              k: tf.train.Feature(
                  float_list=tf.train.FloatList(value=v.flatten()))
              for k, v in float_dict.items()
          }
      ))

def run():
  """Run the beam pipeline to create synthetic dataset."""
  pipeline_options = beam.options.pipeline_options.PipelineOptions(
      FLAGS.pipeline_options)
  with beam.Pipeline(options=pipeline_options) as pipeline:
    for gin_search_path in [GIN_PATH] + FLAGS.gin_search_path:
      gin.add_config_file_search_path(gin_search_path)
    gin.parse_config_files_and_bindings(
        FLAGS.gin_file, FLAGS.gin_param, skip_unknown=True)

    np.random.seed(FLAGS.random_seed)
    _ = (
        pipeline
        | beam.Create(np.random.randint(2**32, size=FLAGS.num_examples))
        | beam.ParDo(GenerateExampleFn(gin.config_str()))
        | beam.Reshuffle()
        | beam.Map(_float_dict_to_tfexample)
        | beam.io.tfrecordio.WriteToTFRecord(
            FLAGS.output_tfrecord_path,
            num_shards=FLAGS.num_shards,
            coder=beam.coders.ProtoCoder(tf.train.Example))
    )

def main(unused_argv):
  """From command line."""
  run()

def console_entry_point():
  """From pip installed script."""
  app.run(main)

if __name__ == '__main__':
  console_entry_point()


import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' #supress tensorflow messages

import gym #OpenAI Gym
import retro #Gym Retro
import numpy as np #NumPy
from time import time #calculate runtime
from collections import deque #needed for replay memory
from random import sample #used to get random minibacth

#TensorFlow 2.0
from tensorflow.keras.models import Model, load_model
from tensorflow.keras.layers import Input, Dense, Conv2D, Flatten, Lambda, Add
from tensorflow.keras.optimizers import Adam
from tensorflow.keras import backend as K
from tensorflow.python.framework.ops import disable_eager_execution
disable_eager_execution() #Disable Eager, IMPORTANT!

from wrappers import * #My wrappers for the Gym Retro Environment
from utils import convert_frames, Now, get_latest_file #My utilities, useful re-usable functions

class DDDQNAgent(object):
    def __init__(self, game, combos, time_limit=None, batch_size=32, learn_every=10, update_every=10000,
                 alpha=1e-4, gamma=0.99, epsilon=1.0, epsilon_min=0.01, epsilon_decay=0.9999, memory_size=100000):
        #retro environment
        self.game = game #game rom name
        self.combos = combos #valid discrete button combinations
        self.env = self.build_env(time_limit=time_limit) #retro environment
        self.num_actions = len(combos) #number of possible actions for env
        self.state_shape = self.env.observation_space.shape #env state dims
        self.state = self.reset() #initialize state
        #training
        self.batch_size=batch_size #batch size
        self.steps = 0 #number of steps ran
        self.learn_every = learn_every #interval of steps to fit model
        self.update_every = update_every #interval of steps to update target model
        self.alpha = alpha #learning rate
        self.gamma = gamma #discount factor
        self.epsilon = epsilon #exploration factor
        self.epsilon_min = epsilon_min #minimum exploration probability
        self.epsilon_decay = epsilon_decay #exponential decay rate of epsilon
        #memory
        self.memory_size = memory_size #replay memory size
        self.memory = deque(maxlen=memory_size) #replay memory
        self.log = [] #stores information from training
        #models
        self.q_eval = self.build_network() #Q eval model
        self.q_target = self.build_network() #Q target model

    def build_env(self, time_limit=None, downsampleRatio=2, numStack=4):
    #Build the gym retro environment.
        env = retro.make(game=self.game, state=retro.State.DEFAULT, scenario='scenario',
                         record=False, obs_type=retro.Observations.IMAGE)
        env = Discretizer(env, combos=self.combos)
        if time_limit is not None: env = TimeLimit(env, time_limit)
        env = SkipFrames(env)
        env = Rgb2Gray(env)
        env = Downsample(env, downsampleRatio)
        env = FrameStack(env, numStack)
        env = ScaledFloatFrame(env)
        return env

    def build_network(self):
    #Build the Dueling DQN Network
        X_input = Input(self.state_shape, name='input')
        X = Conv2D(32, 8, 4, activation='relu')(X_input)
        X = Conv2D(64, 4, 2, activation='relu')(X)
        X = Conv2D(64, 3, 1, activation='relu')(X)
        X = Flatten()(X)
        X = Dense(1024, activation='relu', kernel_initializer='he_uniform')(X)
        X = Dense(512, activation='relu', kernel_initializer='he_uniform')(X)
        #value layer
        V = Dense(1, activation='linear', name='V')(X) #V(S)
        #advantage layer
        A = Dense(self.num_actions, activation='linear', name='Ai')(X) #A(s,a)
        A = Lambda(lambda a: a[:, :] - K.mean(a[:, :], keepdims=True), name='Ao')(A) #A(s,a)
        #Q layer (V + A)
        Q = Add(name='Q')([V, A]) #Q(s,a)
        Q_model = Model(inputs=[X_input], outputs=[Q], name='qvalue')
        Q_model.compile(loss='mse', optimizer=Adam(lr=self.alpha))
        return Q_model

    def update_target(self):
    #Update the q_target model to weights of q_eval model.
        #don't learn until we have at least one batch and update interval steps reached
        if len(self.memory) < self.batch_size or self.steps % self.update_every != 0: return
        self.q_target.set_weights(self.q_eval.get_weights())

    def remember(self, action, next_state, reward, done):
    #Store data in memory and update current state.
        self.memory.append([self.state, action, next_state, reward, done])
        self.state = next_state

    def choose_action(self, training):
    #Predict next action based on current state and decay epsilon.
        if training: #when training allow random exploration
            if np.random.random() < self.epsilon: #get random action
                action = np.random.randint(self.num_actions)
            else: #predict best actions
                action = np.argmax(self.q_eval.predict(self.state)[0])
            #decay epsilon, if epsilon falls below min then set to min
            if self.epsilon > self.epsilon_min:
                self.epsilon *= self.epsilon_decay
            elif self.epsilon < self.epsilon_min:
                self.epsilon = self.epsilon_min
                print('Epsilon mininum of {} reached'.format(self.epsilon_min))
        else: #if not training then always get best action
            action = np.argmax(self.q_eval.predict(self.state)[0])
        return action

    def learn(self):
        #don't learn until we have at least one batch and learn interval steps reached
        if len(self.memory) < self.batch_size or self.steps % self.learn_every != 0: return
        #sample memory for a minibatch
        mini_batch = sample(self.memory, self.batch_size)
        #separate minibatch into elements
        state, action, next_state, reward, done = [np.squeeze(i) for i in zip(*mini_batch)]
        Q = self.q_eval.predict(state) #get Q values for starting states
        Q_next = self.q_eval.predict(next_state) #get Q values for ending states
        Q_target = self.q_target.predict(next_state) #get Q values from target model
        #update q values
        for i in range(self.batch_size):
            if done[i]:
                Q[i][action[i]] = 0.0 #terminal state
            else:
                a = np.argmax(Q_next[i]) ## a'_max = argmax(Q(s',a'))
                Q[i][action[i]] = reward[i] + self.gamma * Q_target[i][a] #Q_max = Q_target(s',a'_max)
        #fit network on batch_size = minibatch_size
        self.q_eval.fit(state, Q, batch_size=self.batch_size, verbose=0, shuffle=False)

    def load(self, directory, qeval_name=None, qtarget_name=None):
    #Load the actor and critic weights.
        print('Loading models ...', end=' ')
        #if no names supplied try to load most recent
        if qeval_name is not None and qtarget_name is not None:
            qeval_path = os.path.join(directory, qeval_name)
            qtarget_path = os.path.join(directory, qtarget_name)
        elif qeval_name is None and qtarget_name is None:
            qeval_path = get_latest_file(directory + '/*QEval.h5')
            qtarget_path = get_latest_file(directory + '/*QTarget.h5')
        self.q_eval.load_weights(qeval_path)
        self.q_target.load_weights(qtarget_path)
        print('Done. Models loaded from {}'.format(directory))
        print('Loaded Q_Eval model {}'.format(qeval_path))
        print('Loaded Q_Target model {}'.format(qtarget_path))

    def save(self, directory, fileName):
    #Save the actor and critic weights.
        print('Saving models ...', end=' ')
        if not os.path.exists(directory): os.makedirs(directory)
        qeval_name = fileName + '_QEval.h5'
        qtarget_name = fileName + '_QTarget.h5'
        self.q_eval.save_weights(os.path.join(directory, qeval_name))
        self.q_target.save_weights(os.path.join(directory, qtarget_name))
        print('Done. Saved to {}'.format(os.path.abspath(directory)))

    def save_log(self, directory, fileName, clear=False):
    #Save the information currently stored in log list.
        print('Saving log ...', end=' ')
        if not os.path.exists(directory): os.makedirs(directory)
        f = open(os.path.join(directory, fileName + '.csv'), 'w')
        for line in self.log:
            f.write(str(line)[1:-1].replace('None', '') + '\n')
        f.close()
        if clear: self.log = []
        print('Done. Saved to {}'.format(os.path.abspath(directory)))

    def reset(self):
    #Reset environment and return expanded state.
        self.state = np.expand_dims(self.env.reset(), axis=0)

    def close(self):
    #Close the environment.
        self.env.close()

    def step(self, action):
    #Run one step for given action and return data.
        observation, reward, done, info = self.env.step(action)
        observation = np.expand_dims(observation, axis=0)
        return observation, reward, done, info

    def run(self, num_episodes=100, render=False, checkpoint=False, cp_render=False, cp_interval=None, otype='AVI'):
    #Run num_episodes number of episodes and train the q model after each learn_every number of steps. The target
    #model is updated every update_every number of steps. If render is true then render each episode to monitor.
    #If checkpoint is true then save model weights and log and evaluate the model and convert and save the frames
    #every cp_interval number of of episodes. The evaluation is rendered if cp_render is true.
        printSTR = 'Episode: {}/{} | Score: {:.4f} | AVG 50: {:.4f} | Elapsed Time: {} mins'
        start_time = time()
        scores = []
        self.reset()
        for e in range(1, num_episodes + 1):
            score = 0
            LIVES = None #will store starting lives
            while True:
                action = self.choose_action(training=True) #predict action
                next_state, reward, done, info = self.step(action) #perform action
                score += reward #cumulative score for episode
                reward = np.clip(reward, -1.0, 1.0).item() #clip reward to range [-1.0, 1.0]
                if LIVES is None: LIVES = info['lives'] #get starting lives
                if info['lives'] < LIVES: done = True #flag for reset when dead
                self.remember(action, next_state, reward, done) #store results
                self.steps += 1 #increment steps
                self.update_target() #update target network (update_every)
                self.learn() #fit q model (learn_every)
                if render: self.env.render()
                if done:
                    scores.append(score) #store scores for all epsisodes
                    self.reset()
                    break
            elapsed_time = round((time() - start_time)/60, 2)
            print(printSTR.format(e, num_episodes, round(score, 4), np.average(scores[-50:]), elapsed_time))
            if checkpoint and (e % cp_interval) == 0:
                eval_score, frames = self.evaluate(render=cp_render)
                print('EVALUATION: {}'.format(round(eval_score, 4)))
                self.log.append([e, score, np.average(scores[-50:]), elapsed_time, eval_score])
                fileName = 'DDDQN_{}_{}_{}'.format(e, self.game, Now(separate=False))
                self.save('models', fileName)
                self.save_log('logs', fileName)
                convert_frames(frames, 'renders', fileName, otype=otype)
            elif checkpoint:
                self.log.append([e, score, np.average(scores[-50:]), elapsed_time, None])
        elapsed_time = round((time() - start_time)/60, 2)
        print('Finished training {} episodes in {} minutes.'.format(num_episodes, (time() - start_time)/60))

    def evaluate(self, render=False):
    #Run an episode and return the score and frames.
        frames = []
        score = 0
        self.reset()
        while True:
            action = self.choose_action(training=False) #get best action
            observation, reward, done, info = self.step(action) #perform action
            self.state = observation #update current state
            score += reward #cumulative score for episode
            if render: self.env.render()
            frames.append(self.env.render(mode='rgb_array'))
            if done:
                self.reset()
                break
        return score, frames

    def play_episode(self, render=False, render_and_save=False, otype='AVI'):
    #Run one episode. If render is true then render each episode to monitor.
    #If render_and_save is true then save frames and convert to GIF image or AVI movie.
    #The reward for the epsiode is returned.
        frames = []
        score = 0
        self.reset()
        while True:
            action = self.choose_action(training=False) #get best action
            observation, reward, done, info = self.step(action) #perform action
            self.state = observation #update current state
            score += reward #cumulative score for episode
            if render: self.env.render()
            if render_and_save: frames.append(self.env.render(mode='rgb_array'))
            if done:
                print('Finished! Score: {}'.format(score))
                self.reset()
                break
        if render_and_save:
            fileName = 'DDDQN_PLAY_{}_{}'.format(self.game, Now(separate=False))
            convert_frames(frames, 'renders', fileName, otype=otype)
        return score


__copyright__ = "Copyright (c) 2021 Jina AI Limited. All rights reserved."
__license__ = "Apache-2.0"

import os
import pytest

import numpy as np
from jina.executors.indexers import BaseIndexer
from jina.executors.metas import get_default_metas

from .. import NGTIndexer

# fix the seed here
np.random.seed(500)
vec_idx = np.random.randint(0, high=100, size=[10]).astype(str)
vec = np.random.random([10, 10]).astype('float32')
query = np.array(np.random.random([10, 10]), dtype=np.float32)

@pytest.fixture(scope='function', autouse=True)
def metas(tmpdir):
    os.environ['TEST_WORKSPACE'] = str(tmpdir)
    metas = get_default_metas()
    metas['workspace'] = os.environ['TEST_WORKSPACE']
    yield metas
    del os.environ['TEST_WORKSPACE']

def test_simple_ngt(tmpdir):
    import ngtpy
    path = os.path.join(tmpdir, 'ngt-index')
    dimension, queries, top_k, batch_size, num_batch = 10, 3, 5, 8, 3

    ngtpy.create(path=path, dimension=dimension, distance_type='L2')
    _index = ngtpy.Index(path=path)
    for i in range(num_batch):
        _index.batch_insert(np.random.random((batch_size, dimension)), num_threads=4)
    assert os.path.exists(path)

    idx = []
    dist = []
    for key in np.random.random((queries, dimension)):
        results = _index.search(key, size=top_k, epsilon=0.1)
        index_k = []
        distance_k = []
        [(index_k.append(result[0]), distance_k.append(result[1])) for result in results]
        idx.append(index_k)
        dist.append(distance_k)

    idx = np.array(idx)
    dist = np.array(dist)

    assert idx.shape == dist.shape
    assert idx.shape == (queries, top_k)

def test_ngt_indexer(metas):
    with NGTIndexer(index_filename='ngt.test.gz', metas=metas) as indexer:
        indexer.add(vec_idx, vec)
        indexer.save()
        assert os.path.exists(indexer.index_abspath)
        save_abspath = indexer.save_abspath

    with BaseIndexer.load(save_abspath) as indexer:
        assert isinstance(indexer, NGTIndexer)
        idx, dist = indexer.query(query, top_k=4)
        assert idx.shape == dist.shape
        assert idx.shape == (10, 4)

def test_ngt_indexer_known(metas):
    vectors = np.array([[1, 1, 1],
                        [10, 10, 10],
                        [100, 100, 100],
                        [1000, 1000, 1000]], dtype=np.float32)
    keys = np.array([4, 5, 6, 7]).reshape(-1, 1).astype(str)
    with NGTIndexer(index_filename='ngt.test.gz', metas=metas) as indexer:
        indexer.add(keys, vectors)
        indexer.save()
        assert os.path.exists(indexer.index_abspath)
        save_abspath = indexer.save_abspath

    queries = np.array([[1, 1, 1],
                        [10, 10, 10],
                        [100, 100, 100],
                        [1000, 1000, 1000]], dtype=np.float32)
    with BaseIndexer.load(save_abspath) as indexer:
        assert isinstance(indexer, NGTIndexer)
        idx, dist = indexer.query(queries, top_k=2)
        np.testing.assert_equal(idx, np.array([[4, 5], [5, 4], [6, 5], [7, 6]]).astype(str))
        assert idx.shape == dist.shape
        assert idx.shape == (4, 2)
        np.testing.assert_equal(indexer.query_by_key(['7', '4']), vectors[[3, 0]])

def test_ngt_indexer_known_big(metas):
    """Let's try to have some real test. We will have an index with 10k vectors of random values between 5 and 10.
     We will change tweak some specific vectors that we expect to be retrieved at query time. We will tweak vector
     at index [0, 1000, 2000, 3000, 4000, 5000, 6000, 7000, 8000, 9000], this will also be the query vectors.
     Then the keys will be assigned shifted to test the proper usage of `_int2ext_id` and `ext2int_id`
    """
    vectors = np.random.uniform(low=5.0, high=10.0, size=(10000, 1024)).astype('float32')

    queries = np.empty((10, 1024))
    for idx in range(0, 10000, 1000):
        array = idx * np.ones((1, 1024))
        queries[int(idx / 1000)] = array
        vectors[idx] = array

    keys = np.arange(10000, 20000).reshape(-1, 1).astype(str)

    with NGTIndexer(index_filename='ngt.test.gz', num_threads=4, metas=metas) as indexer:
        indexer.add(keys, vectors)
        indexer.save()
        assert os.path.exists(indexer.index_abspath)
        save_abspath = indexer.save_abspath

    with BaseIndexer.load(save_abspath) as indexer:
        assert isinstance(indexer, NGTIndexer)
        idx, dist = indexer.query(queries, top_k=1)
        np.testing.assert_equal(idx, np.array(
            [[10000], [11000], [12000], [13000], [14000], [15000], [16000], [17000], [18000], [19000]]).astype(str))
        assert idx.shape == dist.shape
        assert idx.shape == (10, 1)
        np.testing.assert_equal(indexer.query_by_key(['10000', '15000']), vectors[[0, 5000]])


#
# Copyright (c) 2020, NVIDIA CORPORATION.
#
# 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.
#
#
# Copyright (c) 2020, NVIDIA CORPORATION.
#
# 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 cudf
import cupy as cp
import numpy as np
import pandas

import cuml.common.logger as logger
from cuml.experimental.explainer.common import get_dtype_from_model_func
from cuml.experimental.explainer.common import get_handle_from_cuml_model_func
from cuml.experimental.explainer.common import get_link_fn_from_str_or_fn
from cuml.experimental.explainer.common import get_tag_from_model_func
from cuml.experimental.explainer.common import model_func_call
from cuml.common.input_utils import input_to_cupy_array

class SHAPBase():
    """
    Base class for SHAP based explainers.

    Parameters
    ----------
    model : function
        Function that takes a matrix of samples (n_samples, n_features) and
        computes the output for those samples with shape (n_samples). Function
        must use either CuPy or NumPy arrays as input/output.
    data : Dense matrix containing floats or doubles.
        Background dataset. Dense arrays are supported.
    order : 'F', 'C' or None (default = None)
        Set to override detection of row ('C') or column ('F') major order,
        if None it will be attempted to be inferred from model.
    order_default : 'F' or 'C' (default = 'C')
        Used when `order` is None. If the order cannot be inferred from the
        model, then order is set to `order_default`.
    link : function or str (default = 'identity')
        The link function used to map between the output units of the
        model and the SHAP value units.
    random_state: int, RandomState instance or None (default = None)
        Seed for the random number generator for dataset creation.
    is_gpu_model : bool or None (default = None)
        If None Explainer will try to infer whether `model` can take GPU data
        (as CuPy arrays), otherwise it will use NumPy arrays to call `model`.
        Set to True to force the explainer to use GPU data,  set to False to
        force the Explainer to use NumPy data.
    handle : cuml.raft.common.handle
        Specifies the handle that holds internal CUDA state for
        computations in this model. Most importantly, this specifies the CUDA
        stream that will be used for the model's computations, so users can
        run different models concurrently in different streams by creating
        handles in several streams.
        If it is None, a new one is created.
    dtype : np.float32 or np.float64 (default = None)
        Parameter to specify the precision of data to generate to call the
        model. If not specified, the explainer will try to get the dtype
        of the model, if it cannot be queried, then it will default to
        np.float32.
    output_type : 'cupy' or 'numpy' (default = None)
        Parameter to specify the type of data to output.
        If not specified, the explainer will try to see if model is gpu based,
        if so it will be set to `cupy`, otherwise it will be set to `numpy`.
        For compatibility with SHAP's graphing libraries, specify `numpy`.

    """

    def __init__(self,
                 *,
                 model,
                 background,
                 order=None,
                 order_default='C',
                 link='identity',
                 verbose=False,
                 random_state=None,
                 is_gpu_model=None,
                 handle=None,
                 dtype=None,
                 output_type=None):

        if verbose is True:
            self.verbose = logger.level_debug
        elif verbose is False:
            self.verbose = logger.level_error
        else:
            self.verbose = verbose

        if handle is None:
            self.handle = get_handle_from_cuml_model_func(model,
                                                          create_new=True)
        else:
            self.handle = handle

        if order is None:
            self.order = get_tag_from_model_func(func=model,
                                                 tag='preferred_input_order',
                                                 default=order_default)
        else:
            self.order = order

        self.link = link
        self.link_fn = get_link_fn_from_str_or_fn(link)
        self.model = model
        if is_gpu_model is None:
            # todo (dgd): when sparse support is added, use this tag to see if
            # model can accept sparse data
            self.is_gpu_model = \
                get_tag_from_model_func(func=model,
                                        tag='X_types_gpu',
                                        default=None) is not None
        else:
            self.is_gpu_model = is_gpu_model

        # we are defaulting to numpy for now for compatibility
        if output_type is None:
            # self.output_type = 'cupy' if self.is_gpu_model else 'numpy'
            self.output_type = 'numpy'
        else:
            self.output_type = output_type

        # if not dtype is specified, we try to get it from the model
        if dtype is None:
            self.dtype = get_dtype_from_model_func(func=model,
                                                   default=np.float32)
        else:
            if dtype in [np.float32, np.float64]:
                self.dtype = np.dtype(dtype)
            raise ValueError("dtype must be either np.float32 or np.float64")

        self.background, self.N, self.M, _ = \
            input_to_cupy_array(background, order=self.order,
                                convert_to_dtype=self.dtype)

        self.random_state = random_state

        if isinstance(background,
                      pandas.DataFrame) or isinstance(background,
                                                      cudf.DataFrame):
            self.feature_names = background.columns.to_list()
        else:
            self.feature_names = [None for _ in range(len(background))]

        # evaluate the model in background to get the expected_value
        self.expected_value = self.link_fn(
            cp.mean(
                model_func_call(X=self.background,
                                model_func=self.model,
                                gpu_model=self.is_gpu_model),
                axis=0
            )
        )

        # D tells us the dimension of the model. For example, `predict_proba`
        # functions typically return n values for n classes as opposed to
        # 1 valued for a typical `predict`
        if len(self.expected_value.shape) == 0:
            self.D = 1
        else:
            self.D = self.expected_value.shape[0]