from __future__ import print_function
import argparse
from math import log10, ceil
import random, shutil, json
from os.path import join, exists, isfile, realpath, dirname
from os import makedirs, remove, chdir, environ

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import DataLoader, SubsetRandomSampler
from torch.utils.data.dataset import Subset
import torchvision.transforms as transforms
from PIL import Image
from datetime import datetime
import torchvision.datasets as datasets
import torchvision.models as models
import h5py
import faiss

from tensorboardX import SummaryWriter
import numpy as np
import netvlad

from ipdb import set_trace as bp
import sys;sys.path.insert(0,'netvlad/ccsmmutils');import img_utils as myiu

parser = argparse.ArgumentParser(description='pytorch-NetVlad')
parser.add_argument('--mode', type=str, default='train', help='Mode', choices=['train', 'test', 'cluster'])
parser.add_argument('--batchSize', type=int, default=4,
        help='Number of triplets (query, pos, negs). Each triplet consists of 12 images.')
parser.add_argument('--cacheBatchSize', type=int, default=24, help='Batch size for caching and testing')
parser.add_argument('--cacheRefreshRate', type=int, default=1000,
        help='How often to refresh cache, in number of queries. 0 for off')
parser.add_argument('--nEpochs', type=int, default=30, help='number of epochs to train for')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
        help='manual epoch number (useful on restarts)')
parser.add_argument('--nGPU', type=int, default=1, help='number of GPU to use.')
parser.add_argument('--optim', type=str, default='SGD', help='optimizer to use', choices=['SGD', 'ADAM'])
parser.add_argument('--lr', type=float, default=0.0001, help='Learning Rate.')
parser.add_argument('--lrStep', type=float, default=5, help='Decay LR ever N steps.')
parser.add_argument('--lrGamma', type=float, default=0.5, help='Multiply LR by Gamma for decaying.')
parser.add_argument('--weightDecay', type=float, default=0.001, help='Weight decay for SGD.')
parser.add_argument('--momentum', type=float, default=0.9, help='Momentum for SGD.')
parser.add_argument('--nocuda', action='store_true', help='Dont use cuda')
parser.add_argument('--threads', type=int, default=8, help='Number of threads for each data loader to use')
parser.add_argument('--seed', type=int, default=123, help='Random seed to use.')
parser.add_argument('--dataPath', type=str, default='netvlad_v100_datasets/', help='Path for centroid data.')
parser.add_argument('--runsPath', type=str, default='checkpoints/runs/', help='Path to save runs to.')
parser.add_argument('--savePath', type=str, default='checkpoints',
        help='Path to save checkpoints to in logdir. Default=checkpoints/')
parser.add_argument('--cachePath', type=str, default=environ['TMPDIR'], help='Path to save cache to.')
parser.add_argument('--resume', type=str, default='', help='Path to load checkpoint from, for resuming training or testing.')
parser.add_argument('--ckpt', type=str, default='latest',
        help='Resume from latest or best checkpoint.', choices=['latest', 'best'])
parser.add_argument('--evalEvery', type=int, default=1,
        help='Do a validation set run, and save, every N epochs.')
parser.add_argument('--patience', type=int, default=10, help='Patience for early stopping. 0 is off.')
parser.add_argument('--dataset', type=str, default='pittsburgh',
        help='Dataset to use', choices=['pittsburgh'])
parser.add_argument('--arch', type=str, default='vgg16',
        help='basenetwork to use', choices=['vgg16', 'alexnet'])
parser.add_argument('--vladv2', action='store_true', help='Use VLAD v2')
parser.add_argument('--pooling', type=str, default='netvlad', help='type of pooling to use',
        choices=['netvlad', 'max', 'avg'])
parser.add_argument('--num_clusters', type=int, default=64, help='Number of NetVlad clusters. Default=64')
parser.add_argument('--margin', type=float, default=0.1, help='Margin for triplet loss. Default=0.1')
parser.add_argument('--split', type=str, default='val', help='Data split to use for testing. Default is val',
        choices=['test', 'test250k', 'train', 'val'])
parser.add_argument('--fromscratch', action='store_true', help='Train from scratch rather than using pretrained models')

def train(epoch):
    epoch_loss = 0
    startIter = 1 # keep track of batch iter across subsets for logging

    if opt.cacheRefreshRate > 0:
        subsetN = ceil(len(train_set) / opt.cacheRefreshRate)
        #TODO randomise the arange before splitting?
        subsetIdx = np.array_split(np.arange(len(train_set)), subsetN)
    else:
        subsetN = 1
        subsetIdx = [np.arange(len(train_set))]

    nBatches = (len(train_set) + opt.batchSize - 1) // opt.batchSize

    for subIter in range(subsetN):
        print('====> Building Cache')
        model.eval()
        train_set.cache = join(opt.cachePath, train_set.whichSet + '_feat_cache.hdf5')
        with h5py.File(train_set.cache, mode='w') as h5:
            pool_size = encoder_dim
            if opt.pooling.lower() == 'netvlad': pool_size *= opt.num_clusters
            h5feat = h5.create_dataset("features",
                    [len(whole_train_set), pool_size],
                    dtype=np.float32)
            with torch.no_grad():
                for iteration, (input, indices) in enumerate(whole_training_data_loader, 1):
                    input = input.to(device)
                    image_encoding = model.encoder(input)
                    vlad_encoding = model.pool(image_encoding)
                    h5feat[indices.detach().numpy(), :] = vlad_encoding.detach().cpu().numpy()
                    del input, image_encoding, vlad_encoding

        sub_train_set = Subset(dataset=train_set, indices=subsetIdx[subIter])

        training_data_loader = DataLoader(dataset=sub_train_set, num_workers=opt.threads,
                    batch_size=opt.batchSize, shuffle=True,
                    collate_fn=dataset.collate_fn, pin_memory=cuda)

        print('Allocated:', torch.cuda.memory_allocated())
        print('Cached:', torch.cuda.memory_cached())

        model.train()
        for iteration, (query, positives, negatives,
                negCounts, indices) in enumerate(training_data_loader, startIter):
            # some reshaping to put query, pos, negs in a single (N, 3, H, W) tensor
            # where N = batchSize * (nQuery + nPos + nNeg)
            if query is None: continue # in case we get an empty batch

            B, C, H, W = query.shape
            nNeg = torch.sum(negCounts)
            input = torch.cat([query, positives, negatives])

            input = input.to(device)
            image_encoding = model.encoder(input)
            vlad_encoding = model.pool(image_encoding)

            vladQ, vladP, vladN = torch.split(vlad_encoding, [B, B, nNeg])

            optimizer.zero_grad()

            # calculate loss for each Query, Positive, Negative triplet
            # due to potential difference in number of negatives have to
            # do it per query, per negative
            loss = 0
            for i, negCount in enumerate(negCounts):
                for n in range(negCount):
                    negIx = (torch.sum(negCounts[:i]) + n).item()
                    loss += criterion(vladQ[i:i+1], vladP[i:i+1], vladN[negIx:negIx+1])

            loss /= nNeg.float().to(device) # normalise by actual number of negatives
            loss.backward()
            optimizer.step()
            del input, image_encoding, vlad_encoding, vladQ, vladP, vladN
            del query, positives, negatives

            batch_loss = loss.item()
            epoch_loss += batch_loss

            if iteration % 50 == 0 or nBatches <= 10:
                print("==> Epoch[{}]({}/{}): Loss: {:.4f}".format(epoch, iteration,
                    nBatches, batch_loss), flush=True)
                writer.add_scalar('Train/Loss', batch_loss,
                        ((epoch-1) * nBatches) + iteration)
                writer.add_scalar('Train/nNeg', nNeg,
                        ((epoch-1) * nBatches) + iteration)
                print('Allocated:', torch.cuda.memory_allocated())
                print('Cached:', torch.cuda.memory_cached())

        startIter += len(training_data_loader)
        del training_data_loader, loss
        optimizer.zero_grad()
        torch.cuda.empty_cache()
        remove(train_set.cache) # delete HDF5 cache

    avg_loss = epoch_loss / nBatches

    print("===> Epoch {} Complete: Avg. Loss: {:.4f}".format(epoch, avg_loss),
            flush=True)
    writer.add_scalar('Train/AvgLoss', avg_loss, epoch)

def test(eval_set, epoch=0, write_tboard=False):
    # TODO what if features dont fit in memory?
    test_data_loader = DataLoader(dataset=eval_set,
                num_workers=opt.threads, batch_size=opt.cacheBatchSize, shuffle=False,
                pin_memory=cuda)

    model.eval()
    with torch.no_grad():
        print('====> Extracting Features')
        pool_size = encoder_dim
        if opt.pooling.lower() == 'netvlad': pool_size *= opt.num_clusters
        dbFeat = np.empty((len(eval_set), pool_size))

        for iteration, (input, indices) in enumerate(test_data_loader, 1):
            input = input.to(device) #[24, 3, 480, 640]
            image_encoding = model.encoder(input) #[24, 512, 30, 40]
            vlad_encoding = model.pool(image_encoding) #[24,32768]

            #dbFeat : [17608, 32768]
            dbFeat[indices.detach().numpy(), :] = vlad_encoding.detach().cpu().numpy() #[24,32768]
            if iteration % 50 == 0 or len(test_data_loader) <= 10:
                print("==> Batch ({}/{})".format(iteration,len(test_data_loader)), flush=True)
                myiu.clf()
                myiu.imshow(input,221,'input')
                myiu.imshow(image_encoding[:,:3,:,:],222,'encoding')
                myiu.plot(vlad_encoding,223,'vlad')

            del input, image_encoding, vlad_encoding
    del test_data_loader

    # extracted for both db and query, now split in own sets
    qFeat = dbFeat[eval_set.dbStruct.numDb:].astype('float32') #[7608,32768]
    dbFeat = dbFeat[:eval_set.dbStruct.numDb].astype('float32') #[10000,32768]

    print('====> Building faiss index')
    #qFeat  : [7608,32768], pool_size = 32768 as dimension of feature
    #dbFeat : [10000,32768]
    faiss_index = faiss.IndexFlatL2(pool_size) #32768
    faiss_index.add(dbFeat)

    print('====> Calculating recall @ N')
    n_values = [1,5,10,20] #n nearest neighbors

    #we want to see 20 nearnest neighbors using following search command.
    _, predictions = faiss_index.search(qFeat, max(n_values)) #predictions : [7608,20]

    # for each query get those within threshold distance
    gt = eval_set.getPositives() #Ground Truth

    correct_at_n = np.zeros(len(n_values))
    #TODO can we do this on the matrix in one go?
    for qIx, pred in enumerate(predictions):
        for i,n in enumerate(n_values):
            # if in top N then also in top NN, where NN > N
            if np.any(np.in1d(pred[:n], gt[qIx])):
                correct_at_n[i:] += 1
                break
    recall_at_n = correct_at_n / eval_set.dbStruct.numQ

    recalls = {} #make dict for output
    for i,n in enumerate(n_values):
        recalls[n] = recall_at_n[i]
        print("====> Recall@{}: {:.4f}".format(n, recall_at_n[i]))
        if write_tboard: writer.add_scalar('Val/Recall@' + str(n), recall_at_n[i], epoch)

    return recalls

def get_clusters(cluster_set):
    nDescriptors = 50000
    nPerImage = 100
    nIm = ceil(nDescriptors/nPerImage)

    sampler = SubsetRandomSampler(np.random.choice(len(cluster_set), nIm, replace=False))
    data_loader = DataLoader(dataset=cluster_set,
                num_workers=opt.threads, batch_size=opt.cacheBatchSize, shuffle=False,
                pin_memory=cuda,
                sampler=sampler)

    if not exists(join(opt.dataPath, 'centroids')):
        makedirs(join(opt.dataPath, 'centroids'))

    initcache = join(opt.dataPath, 'centroids', opt.arch + '_' + cluster_set.dataset + '_' + str(opt.num_clusters) + '_desc_cen.hdf5')
    with h5py.File(initcache, mode='w') as h5:
        with torch.no_grad():
            model.eval()
            print('====> Extracting Descriptors')
            dbFeat = h5.create_dataset("descriptors",
                        [nDescriptors, encoder_dim],
                        dtype=np.float32)

            for iteration, (input, indices) in enumerate(data_loader, 1):
                input = input.to(device)
                image_descriptors = model.encoder(input).view(input.size(0), encoder_dim, -1).permute(0, 2, 1)

                batchix = (iteration-1)*opt.cacheBatchSize*nPerImage
                for ix in range(image_descriptors.size(0)):
                    # sample different location for each image in batch
                    sample = np.random.choice(image_descriptors.size(1), nPerImage, replace=False)
                    startix = batchix + ix*nPerImage
                    dbFeat[startix:startix+nPerImage, :] = image_descriptors[ix, sample, :].detach().cpu().numpy()

                if iteration % 50 == 0 or len(data_loader) <= 10:
                    print("==> Batch ({}/{})".format(iteration,
                        ceil(nIm/opt.cacheBatchSize)), flush=True)
                del input, image_descriptors

        print('====> Clustering..')
        niter = 100
        kmeans = faiss.Kmeans(encoder_dim, opt.num_clusters, niter=niter, verbose=False)
        kmeans.train(dbFeat[...])

        print('====> Storing centroids', kmeans.centroids.shape)
        h5.create_dataset('centroids', data=kmeans.centroids)
        print('====> Done!')

def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
    model_out_path = join(opt.savePath, filename)
    torch.save(state, model_out_path)
    if is_best:
        shutil.copyfile(model_out_path, join(opt.savePath, 'model_best.pth.tar'))

class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

class L2Norm(nn.Module):
    def __init__(self, dim=1):
        super().__init__()
        self.dim = dim

    def forward(self, input):
        return F.normalize(input, p=2, dim=self.dim)

if __name__ == "__main__":
    opt = parser.parse_args()

    restore_var = ['lr', 'lrStep', 'lrGamma', 'weightDecay', 'momentum',
            'runsPath', 'savePath', 'arch', 'num_clusters', 'pooling', 'optim',
            'margin', 'seed', 'patience']
    if opt.resume:
        flag_file = join(opt.resume, 'checkpoints', 'flags.json')
        if exists(flag_file):
            with open(flag_file, 'r') as f:
                stored_flags = {'--'+k : str(v) for k,v in json.load(f).items() if k in restore_var}
                to_del = []
                for flag, val in stored_flags.items():
                    for act in parser._actions:
                        if act.dest == flag[2:]:
                            # store_true / store_false args don't accept arguments, filter these
                            if type(act.const) == type(True):
                                if val == str(act.default):
                                    to_del.append(flag)
                                else:
                                    stored_flags[flag] = ''
                for flag in to_del: del stored_flags[flag]

                train_flags = [x for x in list(sum(stored_flags.items(), tuple())) if len(x) > 0]
                print('Restored flags:', train_flags)
                opt = parser.parse_args(train_flags, namespace=opt)

    print(opt)

    if opt.dataset.lower() == 'pittsburgh':
        import pittsburgh as dataset
    else:
        raise Exception('Unknown dataset')

    cuda = not opt.nocuda
    if cuda and not torch.cuda.is_available():
        raise Exception("No GPU found, please run with --nocuda")

    device = torch.device("cuda" if cuda else "cpu")

    random.seed(opt.seed)
    np.random.seed(opt.seed)
    torch.manual_seed(opt.seed)
    if cuda:
        torch.cuda.manual_seed(opt.seed)

    print('===> Loading dataset(s)')
    if opt.mode.lower() == 'train':
        whole_train_set = dataset.get_whole_training_set()
        whole_training_data_loader = DataLoader(dataset=whole_train_set,
                num_workers=opt.threads, batch_size=opt.cacheBatchSize, shuffle=False,
                pin_memory=cuda)

        train_set = dataset.get_training_query_set(opt.margin)

        print('====> Training query set:', len(train_set))
        whole_test_set = dataset.get_whole_val_set()
        print('===> Evaluating on val set, query count:', whole_test_set.dbStruct.numQ)
    elif opt.mode.lower() == 'test':
        if opt.split.lower() == 'test':
            whole_test_set = dataset.get_whole_test_set()
            print('===> Evaluating on test set')
        elif opt.split.lower() == 'test250k':
            whole_test_set = dataset.get_250k_test_set()
            print('===> Evaluating on test250k set')
        elif opt.split.lower() == 'train':
            whole_test_set = dataset.get_whole_training_set()
            print('===> Evaluating on train set')
        elif opt.split.lower() == 'val':
            whole_test_set = dataset.get_whole_val_set()
            print('===> Evaluating on val set')
        else:
            raise ValueError('Unknown dataset split: ' + opt.split)
        print('====> Query count:', whole_test_set.dbStruct.numQ)
    elif opt.mode.lower() == 'cluster':
        whole_train_set = dataset.get_whole_training_set(onlyDB=True)

    print('===> Building model')

    pretrained = not opt.fromscratch
    if opt.arch.lower() == 'alexnet':
        encoder_dim = 256
        encoder = models.alexnet(pretrained=pretrained)
        # capture only features and remove last relu and maxpool
        layers = list(encoder.features.children())[:-2]

        if pretrained:
            # if using pretrained only train conv5
            for l in layers[:-1]:
                for p in l.parameters():
                    p.requires_grad = False

    elif opt.arch.lower() == 'vgg16':
        encoder_dim = 512
        encoder = models.vgg16(pretrained=pretrained)
        # capture only feature part and remove last relu and maxpool

        layers = list(encoder.features.children())[:-2]

        if pretrained:
            # if using pretrained then only train conv5_1, conv5_2, and conv5_3
            for l in layers[:-5]:
                for p in l.parameters():
                    p.requires_grad = False

    if opt.mode.lower() == 'cluster' and not opt.vladv2:
        layers.append(L2Norm())

    encoder = nn.Sequential(*layers)
    model = nn.Module()
    model.add_module('encoder', encoder)

    if opt.mode.lower() != 'cluster':
        if opt.pooling.lower() == 'netvlad':
            net_vlad = netvlad.NetVLAD(num_clusters=opt.num_clusters, dim=encoder_dim, vladv2=opt.vladv2)
            if not opt.resume:
                if opt.mode.lower() == 'train':
                    initcache = join(opt.dataPath, 'centroids', opt.arch + '_' + train_set.dataset + '_' + str(opt.num_clusters) +'_desc_cen.hdf5')
                else:
                    initcache = join(opt.dataPath, 'centroids', opt.arch + '_' + whole_test_set.dataset + '_' + str(opt.num_clusters) +'_desc_cen.hdf5')

                if not exists(initcache):
                    raise FileNotFoundError('Could not find clusters, please run with --mode=cluster before proceeding')

                with h5py.File(initcache, mode='r') as h5:
                    clsts = h5.get("centroids")[...]
                    traindescs = h5.get("descriptors")[...]
                    net_vlad.init_params(clsts, traindescs)
                    del clsts, traindescs

            model.add_module('pool', net_vlad)
        elif opt.pooling.lower() == 'max':
            global_pool = nn.AdaptiveMaxPool2d((1,1))
            model.add_module('pool', nn.Sequential(*[global_pool, Flatten(), L2Norm()]))
        elif opt.pooling.lower() == 'avg':
            global_pool = nn.AdaptiveAvgPool2d((1,1))
            model.add_module('pool', nn.Sequential(*[global_pool, Flatten(), L2Norm()]))
        else:
            raise ValueError('Unknown pooling type: ' + opt.pooling)

    isParallel = False
    if opt.nGPU > 1 and torch.cuda.device_count() > 1:
        model.encoder = nn.DataParallel(model.encoder)
        if opt.mode.lower() != 'cluster':
            model.pool = nn.DataParallel(model.pool)
        isParallel = True

    if not opt.resume:
        model = model.to(device)

    if opt.mode.lower() == 'train':
        if opt.optim.upper() == 'ADAM':
            optimizer = optim.Adam(filter(lambda p: p.requires_grad,
                model.parameters()), lr=opt.lr)#, betas=(0,0.9))
        elif opt.optim.upper() == 'SGD':
            optimizer = optim.SGD(filter(lambda p: p.requires_grad,
                model.parameters()), lr=opt.lr,
                momentum=opt.momentum,
                weight_decay=opt.weightDecay)

            scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=opt.lrStep, gamma=opt.lrGamma)
        else:
            raise ValueError('Unknown optimizer: ' + opt.optim)

        # original paper/code doesn't sqrt() the distances, we do, so sqrt() the margin, I think :D
        criterion = nn.TripletMarginLoss(margin=opt.margin**0.5,
                p=2, reduction='sum').to(device)

    if opt.resume:
        if opt.ckpt.lower() == 'latest':
            resume_ckpt = join(opt.resume, 'checkpoints', 'checkpoint.pth.tar')
        elif opt.ckpt.lower() == 'best':
            resume_ckpt = join(opt.resume, 'checkpoints', 'model_best.pth.tar')

        if isfile(resume_ckpt):
            print("=> loading checkpoint '{}'".format(resume_ckpt))
            checkpoint = torch.load(resume_ckpt, map_location=lambda storage, loc: storage)
            opt.start_epoch = checkpoint['epoch']
            best_metric = checkpoint['best_score']
            model.load_state_dict(checkpoint['state_dict'])
            model = model.to(device)
            if opt.mode == 'train':
                optimizer.load_state_dict(checkpoint['optimizer'])
            print("=> loaded checkpoint '{}' (epoch {})"
                  .format(resume_ckpt, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(resume_ckpt))

    bp()
    if opt.mode.lower() == 'test':
        print('===> Running evaluation step')
        epoch = 1
        recalls = test(whole_test_set, epoch, write_tboard=False)
    elif opt.mode.lower() == 'cluster':
        print('===> Calculating descriptors and clusters')
        get_clusters(whole_train_set)
    elif opt.mode.lower() == 'train':
        print('===> Training model')
        writer = SummaryWriter(log_dir=join(opt.runsPath, datetime.now().strftime('%b%d_%H-%M-%S')+'_'+opt.arch+'_'+opt.pooling))

        # write checkpoints in logdir
        logdir = writer.file_writer.get_logdir()
        opt.savePath = join(logdir, opt.savePath)
        if not opt.resume:
            makedirs(opt.savePath)

        with open(join(opt.savePath, 'flags.json'), 'w') as f:
            f.write(json.dumps(
                {k:v for k,v in vars(opt).items()}
                ))
        print('===> Saving state to:', logdir)

        not_improved = 0
        best_score = 0
        for epoch in range(opt.start_epoch+1, opt.nEpochs + 1):
            if opt.optim.upper() == 'SGD':
                scheduler.step(epoch)
            train(epoch)
            if (epoch % opt.evalEvery) == 0:
                recalls = test(whole_test_set, epoch, write_tboard=True)
                is_best = recalls[5] > best_score
                if is_best:
                    not_improved = 0
                    best_score = recalls[5]
                else:
                    not_improved += 1

                save_checkpoint({
                        'epoch': epoch,
                        'state_dict': model.state_dict(),
                        'recalls': recalls,
                        'best_score': best_score,
                        'optimizer' : optimizer.state_dict(),
                        'parallel' : isParallel,
                }, is_best)

                if opt.patience > 0 and not_improved > (opt.patience / opt.evalEvery):
                    print('Performance did not improve for', opt.patience, 'epochs. Stopping.')
                    break

        print("=> Best Recall@5: {:.4f}".format(best_score), flush=True)
        writer.close()


from hierarc.Likelihood.transformed_cosmography import TransformedCosmography
from hierarc.Likelihood.LensLikelihood.base_lens_likelihood import LensLikelihoodBase
from hierarc.Likelihood.anisotropy_scaling import AnisotropyScalingIFU
from hierarc.Util.distribution_util import PDFSampling
import numpy as np
import copy

class LensLikelihood(TransformedCosmography, LensLikelihoodBase, AnisotropyScalingIFU):
    """
    master class containing the likelihood definitions of different analysis
    """
    def __init__(self, z_lens, z_source, name='name', likelihood_type='TDKin', anisotropy_model='NONE',
                 ani_param_array=None, ani_scaling_array_list=None, ani_scaling_array=None,
                 num_distribution_draws=50, kappa_ext_bias=False, kappa_pdf=None, kappa_bin_edges=None, mst_ifu=False,
                 lambda_scaling_property=0, **kwargs_likelihood):
        """

        :param z_lens: lens redshift
        :param z_source: source redshift
        :param name: string (optional) to name the specific lens
        :param likelihood_type: string to specify the likelihood type
        :param ani_param_array: array of anisotropy parameter values for which the kinematics are predicted
        :param ani_scaling_array: velocity dispersion sigma**2 scaling (also J scaling) of anisotropy parameter relative
         to default prediction. The scaling corresponds to the ani_param_array parameter spacing
         (to generate an interpolation function). A value =1 in ani_scaling_array results in the value stored in the
         provided J() predictions.
        :param ani_scaling_array_list: list of array with the scalings of J() for each IFU
        :param num_distribution_draws: int, number of distribution draws from the likelihood that are being averaged over
        :param kappa_ext_bias: bool, if True incorporates the global external selection function into the likelihood.
        If False, the likelihood needs to incorporate the individual selection function with sufficient accuracy.
        :param kappa_pdf: array of probability density function of the external convergence distribution
         binned according to kappa_bin_edges
        :param kappa_bin_edges: array of length (len(kappa_pdf)+1), bin edges of the kappa PDF
        :param mst_ifu: bool, if True replaces the lambda_mst parameter by the lambda_ifu parameter (and distribution)
         in sampling this lens.
        :param lambda_scaling_property: float (optional), scaling of lambda_mst = lambda_mst_global + alpha * lambda_scaling_property
        :param kwargs_likelihood: keyword arguments specifying the likelihood function,
        see individual classes for their use
        """
        TransformedCosmography.__init__(self, z_lens=z_lens, z_source=z_source)
        if ani_scaling_array_list is None and ani_scaling_array is not None:
            ani_scaling_array_list = [ani_scaling_array]
        AnisotropyScalingIFU.__init__(self, anisotropy_model=anisotropy_model, ani_param_array=ani_param_array,
                                      ani_scaling_array_list=ani_scaling_array_list)
        LensLikelihoodBase.__init__(self, z_lens=z_lens, z_source=z_source, likelihood_type=likelihood_type, name=name,
                                    **kwargs_likelihood)
        self._num_distribution_draws = int(num_distribution_draws)
        self._kappa_ext_bias = kappa_ext_bias
        self._mst_ifu = mst_ifu
        if kappa_pdf is not None and kappa_bin_edges is not None:
            self._kappa_dist = PDFSampling(bin_edges=kappa_bin_edges, pdf_array=kappa_pdf)
            self._draw_kappa = True
        else:
            self._draw_kappa = False
        self._lambda_scaling_property = lambda_scaling_property

    def lens_log_likelihood(self, cosmo, kwargs_lens=None, kwargs_kin=None, kwargs_source=None):
        """

        :param cosmo: astropy.cosmology instance
        :param kwargs_lens: keywords of the hyper parameters of the lens model
        :param kwargs_kin: keyword arguments of the kinematic model hyper parameters
        :param kwargs_source: keyword argument of the source model (such as SNe)
        :return: log likelihood of the data given the model
        """

        # here we compute the unperturbed angular diameter distances of the lens system given the cosmology
        # Note: Distances are in physical units of Mpc. Make sure the posteriors to evaluate this likelihood is in the
        # same units
        ddt, dd = self.angular_diameter_distances(cosmo)
        delta_lum_dist = self.luminosity_distance_modulus(cosmo)
        # here we effectively change the posteriors of the lens, but rather than changing the instance of the KDE we
        # displace the predicted angular diameter distances in the opposite direction
        return self.hyper_param_likelihood(ddt, dd, delta_lum_dist, kwargs_lens=kwargs_lens, kwargs_kin=kwargs_kin,
                                           kwargs_source=kwargs_source)

    def hyper_param_likelihood(self, ddt, dd, delta_lum_dist, kwargs_lens=None, kwargs_kin=None, kwargs_source=None):
        """

        :param ddt: time-delay distance
        :param dd: angular diameter distance to the deflector
        :param delta_lum_dist: relative luminosity distance to pivot redshift
        :param kwargs_lens: keywords of the hyper parameters of the lens model
        :param kwargs_kin: keyword arguments of the kinematic model hyper parameters
        :param kwargs_source: keyword argument of the source model (such as SNe)
        :return: log likelihood given the single lens analysis for the given hyper parameter
        """
        kwargs_lens = self._kwargs_init(kwargs_lens)
        kwargs_kin = self._kwargs_init(kwargs_kin)
        kwargs_source = self._kwargs_init(kwargs_source)
        kwargs_kin_copy = copy.deepcopy(kwargs_kin)
        sigma_v_sys_error = kwargs_kin_copy.pop('sigma_v_sys_error', None)

        if self.check_dist(kwargs_lens, kwargs_kin, kwargs_source):  # sharp distributions
            return self.log_likelihood_single(ddt, dd, delta_lum_dist, kwargs_lens, kwargs_kin_copy, kwargs_source,
                                              sigma_v_sys_error=sigma_v_sys_error)
        else:
            likelihood = 0
            for i in range(self._num_distribution_draws):
                logl = self.log_likelihood_single(ddt, dd, delta_lum_dist, kwargs_lens, kwargs_kin_copy, kwargs_source,
                                                  sigma_v_sys_error=sigma_v_sys_error)
                exp_logl = np.exp(logl)
                if np.isfinite(exp_logl) and exp_logl > 0:
                    likelihood += exp_logl
            if likelihood <= 0:
                return -np.inf
            return np.log(likelihood/self._num_distribution_draws)

    def log_likelihood_single(self, ddt, dd, delta_lum_dist, kwargs_lens, kwargs_kin, kwargs_source, sigma_v_sys_error=None):
        """

        :param ddt: time-delay distance
        :param dd: angular diameter distance to the deflector
        :param delta_lum_dist: relative luminosity distance to pivot redshift
        :param kwargs_lens: keywords of the hyper parameters of the lens model
        :param kwargs_kin: keyword arguments of the kinematic model hyper parameters
        :param kwargs_source: keyword arguments of source brightness
        :param sigma_v_sys_error: unaccounted uncertainty in the velocity dispersion measurement
        :return: log likelihood given the single lens analysis for a single (random) realization of the hyper parameter distribution
        """
        lambda_mst, kappa_ext, gamma_ppn = self.draw_lens(**kwargs_lens)
        # draw intrinsic source magnitude
        mag_source = self.draw_source(lum_dist=delta_lum_dist, **kwargs_source)
        ddt_, dd_, mag_source_ = self.displace_prediction(ddt, dd, gamma_ppn=gamma_ppn, lambda_mst=lambda_mst,
                                                          kappa_ext=kappa_ext, mag_source=mag_source)
        aniso_param_array = self.draw_anisotropy(**kwargs_kin)
        aniso_scaling = self.ani_scaling(aniso_param_array)

        lnlikelihood = self.log_likelihood(ddt_, dd_, aniso_scaling=aniso_scaling, sigma_v_sys_error=sigma_v_sys_error,
                                           mu_intrinsic=mag_source_)
        return lnlikelihood

    def angular_diameter_distances(self, cosmo):
        """
        time-delay distance Ddt, angular diameter distance to the lens (dd)

        :param cosmo: astropy.cosmology instance (or equivalent with interpolation)
        :return: ddt, dd, ds in units physical Mpc
        """
        dd = cosmo.angular_diameter_distance(z=self._z_lens).value
        ds = cosmo.angular_diameter_distance(z=self._z_source).value
        dds = cosmo.angular_diameter_distance_z1z2(z1=self._z_lens, z2=self._z_source).value
        ddt = (1. + self._z_lens) * dd * ds / dds
        return ddt, dd

    def luminosity_distance_modulus(self, cosmo):
        """
        the difference in luminosity distance between a pivot redshift (z=0.1) and the source redshift
        (effectively the ratio as this is the magnitude transform)

        :param cosmo: astropy.cosmology instance (or equivalent with interpolation)
        :return: lum_dist(z_source) - lum_dist(z_pivot)
        """
        angular_diameter_distances = cosmo.angular_diameter_distance(self._z_source).value
        lum_dists = (5 * np.log10((1 + self._z_source) * (1 + self._z_source) * angular_diameter_distances))

        z_anchor = 0.1
        ang_dist_anchor = cosmo.angular_diameter_distance(z_anchor).value
        lum_dist_anchor = (5 * np.log10((1 + z_anchor) * (1 + z_anchor) * ang_dist_anchor))
        delta_lum_dist = lum_dists - lum_dist_anchor
        return delta_lum_dist

    def check_dist(self, kwargs_lens, kwargs_kin, kwargs_source):
        """
        checks if the provided keyword arguments describe a distribution function of hyper parameters or are single
        values

        :param kwargs_lens: lens model hyper parameter keywords
        :param kwargs_kin: kinematic model hyper parameter keywords
        :param kwargs_source: source brightness hyper parameter keywords
        :return: bool, True if delta function, else False
        """
        lambda_mst_sigma = kwargs_lens.get('lambda_mst_sigma', 0)  # scatter in MST
        kappa_ext_sigma = kwargs_lens.get('kappa_ext_sigma', 0)
        a_ani_sigma = kwargs_kin.get('a_ani_sigma', 0)
        beta_inf_sigma = kwargs_kin.get('beta_inf_sigma', 0)
        sne_sigma = kwargs_source.get('sigma_sne', 0)
        if a_ani_sigma == 0 and lambda_mst_sigma == 0 and kappa_ext_sigma == 0 and beta_inf_sigma == 0 and sne_sigma == 0:
            if self._draw_kappa is False:
                return True
        return False

    def draw_lens(self, lambda_mst=1, lambda_mst_sigma=0, kappa_ext=0, kappa_ext_sigma=0, gamma_ppn=1, lambda_ifu=1,
                  lambda_ifu_sigma=0, alpha_lambda=0):
        """

        :param lambda_mst: MST transform
        :param lambda_mst_sigma: spread in the distribution
        :param kappa_ext: external convergence mean in distribution
        :param kappa_ext_sigma: spread in the distribution
        :param gamma_ppn: Post-Newtonian parameter
        :param lambda_ifu: secondary lambda_mst parameter for subset of lenses specified for
        :param lambda_ifu_sigma: secondary lambda_mst_sigma parameter for subset of lenses specified for
        :param alpha_lambda: float, linear slope of the lambda_int scaling relation with lens quantity self._lambda_scaling_property
        :return: draw from the distributions
        """
        if self._mst_ifu is True:
            lambda_lens = lambda_ifu + alpha_lambda * self._lambda_scaling_property
            lambda_mst_draw = np.random.normal(lambda_lens, lambda_ifu_sigma)
        else:
            lambda_lens = lambda_mst + alpha_lambda * self._lambda_scaling_property
            lambda_mst_draw = np.random.normal(lambda_lens, lambda_mst_sigma)
        if self._draw_kappa is True:
            kappa_ext_draw = self._kappa_dist.draw_one
        elif self._kappa_ext_bias is True:
            kappa_ext_draw = np.random.normal(kappa_ext, kappa_ext_sigma)
        else:
            kappa_ext_draw = 0
        return lambda_mst_draw, kappa_ext_draw, gamma_ppn

    @staticmethod
    def draw_source(mu_sne=1, sigma_sne=0, lum_dist=0):
        """

        :param mu_sne: mean brightness of SNe
        :param sigma_sne: std of brightness distribution of SNe relative to the mean brightness
        :param lum_dist: luminosity distance
        (astronomical magnitude scaling of defined brightness to the source redshift)
        :return: realization of source amplitude given distribution
        """
        # draw apparent magnitude at pivot luminosity distance (z=0.1)
        mag_draw = np.random.normal(loc=mu_sne, scale=sigma_sne/mu_sne)
        # move apparent magnitude to redshift of source with relative luminosity distance
        mag_source = mag_draw + lum_dist
        # return linear amplitude with base log 10
        return mag_source

    def sigma_v_measured_vs_predict(self, cosmo, kwargs_lens={}, kwargs_kin={}):
        """
        mean and error covariance of velocity dispersion measurement
        mean and error covariance of velocity dispersion predictions

        :param cosmo: astropy.cosmology instance
        :param kwargs_lens: keywords of the hyper parameters of the lens model
        :param kwargs_kin: keyword arguments of the kinematic model hyper parameters
        :return: sigma_v_measurement, cov_error_measurement, sigma_v_predict_mean, cov_error_predict
        """
        # if no kinematics is provided, return None's
        if not self.likelihood_type in ['DdtHistKin', 'IFUKinCov', 'DdtGaussKin']:
            return None, None, None, None
        kwargs_kin_copy = copy.deepcopy(kwargs_kin)
        sigma_v_sys_error = kwargs_kin_copy.pop('sigma_v_sys_error', None)
        ddt, dd = self.angular_diameter_distances(cosmo)
        sigma_v_measurement, cov_error_measurement = self.sigma_v_measurement(sigma_v_sys_error=sigma_v_sys_error)
        sigma_v_predict_list = []
        sigma_v_predict_mean = np.zeros_like(sigma_v_measurement)
        cov_error_predict = np.zeros_like(cov_error_measurement)
        for i in range(self._num_distribution_draws):
            lambda_mst, kappa_ext, gamma_ppn = self.draw_lens(**kwargs_lens)
            ddt_, dd_, _ = self.displace_prediction(ddt, dd, gamma_ppn=gamma_ppn, lambda_mst=lambda_mst,
                                                    kappa_ext=kappa_ext)
            aniso_param_array = self.draw_anisotropy(**kwargs_kin_copy)
            aniso_scaling = self.ani_scaling(aniso_param_array)
            sigma_v_predict_i, cov_error_predict_i = self.sigma_v_prediction(ddt_, dd_, aniso_scaling=aniso_scaling)
            sigma_v_predict_mean += sigma_v_predict_i
            cov_error_predict += cov_error_predict_i
            sigma_v_predict_list.append(sigma_v_predict_i)

        sigma_v_predict_mean /= self._num_distribution_draws
        cov_error_predict /= self._num_distribution_draws
        sigma_v_mean_std = np.std(sigma_v_predict_list, axis=0)
        cov_error_predict += np.outer(sigma_v_mean_std, sigma_v_mean_std)
        return sigma_v_measurement, cov_error_measurement, sigma_v_predict_mean, cov_error_predict

    def ddt_dd_model_prediction(self, cosmo, kwargs_lens={}):
        """
        predicts the model uncertainty corrected ddt prediction of the applied model (e.g. power-law)

        :param cosmo: astropy.cosmology instance
        :param kwargs_lens: keywords of the hyper parameters of the lens model
        :return: ddt_model mean, ddt_model sigma, dd_model mean, dd_model sigma
        """
        ddt, dd = self.angular_diameter_distances(cosmo)
        ddt_draws = []
        dd_draws = []
        for i in range(self._num_distribution_draws):
            lambda_mst, kappa_ext, gamma_ppn = self.draw_lens(**kwargs_lens)
            ddt_, dd_, _ = self.displace_prediction(ddt, dd, gamma_ppn=gamma_ppn, lambda_mst=lambda_mst,
                                                    kappa_ext=kappa_ext)
            ddt_draws.append(ddt_)
            dd_draws.append(dd_)
        return np.mean(ddt_draws), np.std(ddt_draws), np.mean(dd_draws), np.std(dd_draws)

    @staticmethod
    def _kwargs_init(kwargs=None):
        """

        :param kwargs: keyword argument or None
        :return: keyword argument
        """
        if kwargs is None:
            kwargs = {}
        return kwargs


"""
Code implement MEMR in https://arxiv.org/abs/2006.04802
"""

from typing import Callable
from typing import Dict

import gym
import numpy as np
import rlutils.algos.tf.mf.sac as sac
import rlutils.infra as rl_infra
import rlutils.np as rln
import rlutils.replay_buffers as replay_buffers
import rlutils.tf as rlu
import tensorflow as tf

class PyPrioritizedReplayBuffer(replay_buffers.DictPrioritizedReplayBuffer):
    def reset(self):
        super(PyPrioritizedReplayBuffer, self).reset()
        self.last_priority_ptr = 0

    def get_priority_uninitialized_data(self):
        idx = np.arange(self.last_priority_ptr, len(self))
        self.last_priority_ptr = len(self)
        return idx, self.__getitem__(idx)

class SegmentReplayBuffer(replay_buffers.DictReplayBuffer):
    """
    When adding, the input size must be equal to the segment size.
    When sampling, first sample a segment, then sample the data.
    """

    def __init__(self, segment_size, *args, **kwargs):
        self.segment_size = segment_size
        super(SegmentReplayBuffer, self).__init__(*args, **kwargs)

    def add(self, data: Dict[str, np.ndarray]):
        batch_size = list(data.values())[0].shape[0]
        assert batch_size == self.segment_size
        super(SegmentReplayBuffer, self).add(data=data)

    def sample(self):
        assert not self.is_empty()
        num_segment = len(self) // self.segment_size
        segment_idx = self.np_random.randint(num_segment)
        start = segment_idx * self.segment_size
        end = (segment_idx + 1) * self.segment_size
        idx_range = np.arange(start, end)
        idxs = self.np_random.choice(idx_range, size=self.batch_size)
        return self.__getitem__(idxs)

class MEMRUpdater(rl_infra.OffPolicyUpdater):
    def __init__(self, total_steps, model_update_every, model_replay_buffer, model_rollout_freq,
                 **kwargs):
        super(MEMRUpdater, self).__init__(**kwargs)
        self.total_steps = total_steps
        self.model_update_every = model_update_every
        self.model_replay_buffer = model_replay_buffer
        self.model_rollout_freq = model_rollout_freq
        self.beta_scheduler = rln.schedulers.LinearSchedule(total_steps, initial_p=0.4, final_p=1.0)
        self.update_scheduler = rln.schedulers.LinearSchedule(total_steps, initial_p=1,
                                                              final_p=self.update_per_step)

    def log_tabular(self):
        super(MEMRUpdater, self).log_tabular()
        self.logger.log_tabular('Priority', with_min_and_max=True)

    def update(self, global_step):
        if global_step % self.model_update_every == 0:
            self.agent.update_model(self.replay_buffer.get())

        if global_step % self.model_rollout_freq == 0:
            # update uninitialized priority
            idx, data = self.replay_buffer.get_priority_uninitialized_data()
            priorities = self.agent.compute_priority(obs=data['obs'])
            self.replay_buffer.update_priorities(idx, priorities)
            # sample obs and unroll transitions.
            data, idx = self.replay_buffer.sample(beta=self.beta_scheduler.value(global_step))
            transitions = self.agent.unroll_trajectory(data['obs'])
            transitions['weights'] = data['weights'] / np.mean(data['weights'])
            self.agent.behavior_policy.train_on_batch(x=transitions['obs'],
                                                      y=transitions['act'])
            # update the priority after update the behavior policy
            priorities = self.agent.compute_priority(obs=transitions['obs']).numpy()
            self.replay_buffer.update_priorities(idx, priorities)
            self.logger.store(Priority=priorities)
            # add transitions to the model_replay_buffer
            transitions = rlu.functional.to_numpy_or_python_type(transitions)
            self.model_replay_buffer.add(transitions)

        if global_step % self.update_every == 0:
            update_per_step = int(self.update_scheduler.value(global_step))
            for _ in range(self.update_every):
                for _ in range(update_per_step):
                    batch = self.model_replay_buffer.sample()
                    batch['update_target'] = ((self.policy_updates + 1) % self.policy_delay == 0)
                    self.agent.train_on_batch(data=batch)
                    self.policy_updates += 1

class SquashedGaussianMLPActor(rlu.nn.SquashedGaussianMLPActor):
    @tf.function
    def compute_log_prob_and_log_std(self, obs, raw_actions):
        pi_distribution = self((obs, tf.constant(False)))[-1]
        logp_pi = pi_distribution.log_prob(raw_actions)
        log_std = tf.reduce_mean(tf.math.log(pi_distribution.stddev()), axis=-1)
        return logp_pi / self.ac_dim, log_std

class MEMRAgent(tf.keras.Model):
    def __init__(self, obs_spec, act_spec,
                 model_mlp_hidden=512, model_lr=1e-3, model_num_ensembles=5, reward_fn=None, terminate_fn=None,
                 policy_mlp_hidden=256, policy_lr=3e-4):
        super(MEMRAgent, self).__init__()
        self.obs_spec = obs_spec
        self.act_spec = act_spec
        self.act_dim = self.act_spec.shape[0]
        if len(self.obs_spec.shape) == 1:  # 1D observation
            self.obs_dim = self.obs_spec.shape[0]
        else:
            raise NotImplementedError
        self.dynamics_model = rlu.nn.EnsembleWorldModel(obs_dim=self.obs_dim, act_dim=self.act_dim,
                                                        mlp_hidden=model_mlp_hidden, num_layers=4,
                                                        num_ensembles=model_num_ensembles, lr=model_lr,
                                                        reward_fn=reward_fn, terminate_fn=terminate_fn)
        self.agent = sac.SACAgent(obs_spec=obs_spec, act_spec=act_spec, policy_mlp_hidden=policy_mlp_hidden,
                                  policy_lr=policy_lr, q_mlp_hidden=policy_mlp_hidden, q_lr=policy_lr, alpha=1.0,
                                  alpha_lr=policy_lr, tau=5e-3, gamma=0.99, target_entropy=-(self.act_dim // 2),
                                  auto_alpha=True)
        self.behavior_policy = SquashedGaussianMLPActor(ob_dim=self.obs_dim, ac_dim=self.act_dim,
                                                        mlp_hidden=policy_mlp_hidden)
        self.behavior_policy.compile(optimizer=rlu.future.get_adam_optimizer(policy_lr))

    def set_logger(self, logger):
        self.dynamics_model.set_logger(logger)
        self.agent.set_logger(logger)

    def log_tabular(self):
        self.agent.log_tabular()
        self.dynamics_model.log_tabular()

    @tf.function
    def compute_priority(self, obs):
        raw_actions = self.agent.policy_net((obs, tf.constant(False)))[2]
        logp_pi, log_std = self.behavior_policy.compute_log_prob_and_log_std(obs, raw_actions)
        constant = 0.5 * np.log(2. * np.pi)
        priorities = -(logp_pi + log_std + constant)
        priorities = tf.clip_by_value(priorities, clip_value_min=1e-4, clip_value_max=50.)
        return priorities

    def update_model(self, data):
        self.dynamics_model.update(inputs=data, sample_weights=None, batch_size=512, num_epochs=100,
                                   patience=5, validation_split=0.1, shuffle=True)

    @tf.function
    def unroll_trajectory(self, obs):
        act = self.agent.act_batch_explore_tf(obs)
        next_obs, rew, done = self.dynamics_model.predict_on_batch_tf(obs, act)
        done = tf.cast(done, tf.float32)
        return dict(
            obs=obs,
            act=act,
            next_obs=next_obs,
            rew=rew,
            done=done
        )

    def train_on_batch(self, data, **kwargs):
        return self.agent.train_on_batch(data=data)

    def act_batch_explore(self, obs):
        return self.agent.act_batch_explore(obs)

    def act_batch_test(self, obs):
        return self.agent.act_batch_test(obs)

class Runner(rl_infra.runner.TFOffPolicyRunner):
    def on_epoch_end(self, epoch):
        self.tester.test_agent(get_action=lambda obs: self.agent.act_batch_test(obs),
                               name=self.agent.__class__.__name__,
                               num_test_episodes=self.num_test_episodes)
        # Log info about epoch
        self.logger.log_tabular('Epoch', epoch)
        self.tester.log_tabular()
        self.sampler.log_tabular()
        self.updater.log_tabular()
        self.logger.log_tabular(key='EnvDatsetSize', val=len(self.replay_buffer))
        self.logger.log_tabular(key='ModelDatasetSize', val=len(self.model_replay_buffer))
        self.timer.log_tabular()
        self.logger.dump_tabular()

    def setup_updater(self, update_after, policy_delay, update_per_step, update_every, model_rollout_freq=None):
        self.update_after = update_after
        self.updater = MEMRUpdater(total_steps=self.epochs * self.steps_per_epoch,
                                   model_update_every=self.steps_per_epoch // 4,
                                   agent=self.agent,
                                   replay_buffer=self.replay_buffer,
                                   model_replay_buffer=self.model_replay_buffer,
                                   model_rollout_freq=model_rollout_freq,
                                   policy_delay=policy_delay,
                                   update_per_step=update_per_step,
                                   update_every=update_every
                                   )

    def setup_agent(self, **kwargs):
        from rlutils.gym import static
        static_fn = static.get_static_fn(self.env_name)
        self.agent = MEMRAgent(obs_spec=self.env.single_observation_space,
                               act_spec=self.env.single_action_space,
                               reward_fn=None,
                               terminate_fn=static_fn.terminate_fn_tf_batch,
                               **kwargs)

    def setup_replay_buffer(self, replay_size, batch_size, segment_size=None):
        segment_size = segment_size
        self.seeds_info['replay_buffer'] = self.seeder.generate_seed()
        self.seeds_info['model_replay_buffer'] = self.seeder.generate_seed()
        self.replay_buffer = PyPrioritizedReplayBuffer.from_vec_env(vec_env=self.env, capacity=replay_size,
                                                                    batch_size=segment_size, alpha=0.6,
                                                                    seed=self.seeds_info['replay_buffer'])
        data_spec = {
            'obs': self.env.single_observation_space,
            'act': self.env.single_action_space,
            'next_obs': self.env.single_observation_space,
            'rew': gym.spaces.Space(shape=None, dtype=np.float32),
            'done': gym.spaces.Space(shape=None, dtype=np.float32),
            'weights': gym.spaces.Space(shape=None, dtype=np.float32)
        }
        self.model_replay_buffer = SegmentReplayBuffer(data_spec=data_spec, capacity=replay_size,
                                                       batch_size=batch_size,
                                                       seed=self.seeds_info['model_replay_buffer'],
                                                       segment_size=segment_size)

    @classmethod
    def main(cls,
             env_name,
             env_fn: Callable = None,
             exp_name: str = None,
             steps_per_epoch=1000,
             epochs=100,
             start_steps=3000,
             update_after=750,
             update_every=50,
             update_per_step=20,
             policy_delay=1,
             batch_size=4000,
             model_rollout_freq=50,
             num_model_rollouts=400,
             num_parallel_env=1,
             num_test_episodes=30,
             seed=1,
             # agent
             model_mlp_hidden=512,
             model_lr=1e-3,
             model_num_ensembles=5,
             policy_mlp_hidden=256,
             policy_lr=3e-4,
             # replay
             replay_size=int(1e6),
             logger_path='data'
             ):
        config = locals()
        runner = cls(seed=seed, steps_per_epoch=steps_per_epoch, epochs=epochs, exp_name=None,
                     logger_path=logger_path)
        runner.setup_env(env_name=env_name, env_fn=env_fn, num_parallel_env=1, asynchronous=False,
                         num_test_episodes=num_test_episodes)
        agent_kwargs = dict(

        )
        runner.setup_agent(**agent_kwargs)
        runner.setup_replay_buffer(replay_size=replay_size,
                                   batch_size=batch_size,
                                   segment_size=num_model_rollouts * model_rollout_freq
                                   )
        runner.setup_sampler(start_steps=start_steps)
        runner.setup_tester(num_test_episodes=num_test_episodes)
        runner.setup_updater(update_after=update_after,
                             policy_delay=policy_delay,
                             update_per_step=update_per_step,
                             update_every=update_every,
                             model_rollout_freq=model_rollout_freq)
        runner.setup_logger(config=config, tensorboard=False)
        runner.run()

if __name__ == '__main__':
    rl_infra.runner.run_func_as_main(Runner.main)