import os
#os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'
import numpy as np
import scipy.misc
from keras.layers import Input, Dropout, Conv2D, MaxPooling2D, UpSampling2D, concatenate
from keras.models import load_model
from keras import backend as K
from nets import *

DATA_LOAD_PATH = 'address to load data'
WEIGHT_LOAD_PATH = 'address to load weight'
SAVE_PATH = 'address to save'

# load data
load = np.load(LOAD_PATH)
x_test = load['arr_0']
x_test = x_test / 255
mean = x_test.mean(axis=0)
x = x_test - mean
print(np.shape(x))
# model
inputs = Input((512, 512, 1))
conv1 = Conv2D(filters=64, kernel_size=3,
               activation='selu', padding='same', name='conv1_1', kernel_initializer='he_normal')(inputs)
conv1 = Conv2D(filters=64, kernel_size=3,
               activation='selu', padding='same', name='conv1_2', kernel_initializer='he_normal')(conv1)
#conv1 = Dropout(0.5)(conv1)
pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

conv2 = Conv2D(filters=128, kernel_size=3,
               activation='selu', padding='same', name='conv2_1', kernel_initializer='he_normal')(pool1)
conv2 = Conv2D(filters=128, kernel_size=3,
               activation='selu', padding='same', name='conv2_2', kernel_initializer='he_normal')(conv2)
#conv2 = Dropout(0.5)(conv2)
pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)

conv3 = Conv2D(filters=256, kernel_size=3,
               activation='selu', padding='same', name='conv3_1', kernel_initializer='he_normal')(pool2)
conv3 = Conv2D(filters=256, kernel_size=3,
               activation='selu', padding='same', name='conv3_2', kernel_initializer='he_normal')(conv3)
#conv3 = Dropout(0.5)(conv3)
pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)

conv4 = Conv2D(filters=512, kernel_size=3,  activation='selu',
               padding='same', name='conv4_1', kernel_initializer='he_normal')(pool3)
conv4 = Conv2D(filters=512, kernel_size=3,  activation='selu',
               padding='same', name='conv4_2', kernel_initializer='he_normal')(conv4)
#conv4 = Dropout(0.5)(conv4)
pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)

conv5 = Conv2D(filters=1024, kernel_size=3,
               activation='selu', padding='same', name='conv5_1', kernel_initializer='he_normal')(pool4)
conv5 = Conv2D(filters=1024,  kernel_size=3,
               activation='selu', padding='same', name='conv5_2', kernel_initializer='he_normal')(conv5)
#conv5 = Dropout(0.5)(conv5)

up6 = concatenate([Conv2D(filters=512, kernel_size=3, activation='selu', padding='same',
                          name='deconv6', kernel_initializer='he_normal')(UpSampling2D(size=(2, 2))(conv5)), conv4])
conv6 = Conv2D(filters=512, kernel_size=3,
               activation='selu', padding='same', name='conv6_1', kernel_initializer='he_normal')(up6)
#conv6 = Dropout(0.5)(conv6)
conv6 = Conv2D(filters=512, kernel_size=3,
               activation='selu', padding='same', name='conv6_2', kernel_initializer='he_normal')(conv6)

up7 = concatenate([Conv2D(filters=256, kernel_size=3, activation='selu', padding='same',
                          name='deconv7', kernel_initializer='he_normal')(UpSampling2D(size=(2, 2))(conv6)), conv3])
conv7 = Conv2D(filters=256, kernel_size=3,
               activation='selu', padding='same', name='conv7_1', kernel_initializer='he_normal')(up7)
#conv7 = Dropout(0.5)(conv7)
conv7 = Conv2D(filters=256, kernel_size=3,
               activation='selu', padding='same', name='conv7_2', kernel_initializer='he_normal')(conv7)

up8 = concatenate([Conv2D(filters=128, kernel_size=3, activation='selu', padding='same',
                          name='deconv8', kernel_initializer='he_normal')(UpSampling2D(size=(2, 2))(conv7)), conv2])
conv8 = Conv2D(filters=128,  kernel_size=3,
               activation='selu', padding='same', name='conv8_1', kernel_initializer='he_normal')(up8)
#conv8 = Dropout(0.5)(conv8)
conv8 = Conv2D(filters=128, kernel_size=3,
               activation='selu', padding='same', name='conv8_2', kernel_initializer='he_normal')(conv8)

up9 = concatenate([Conv2D(filters=64, kernel_size=3, activation='selu', padding='same',
                          name='deconv9', kernel_initializer='he_normal')(UpSampling2D(size=(2, 2))(conv8)), conv1])
conv9 = Conv2D(filters=64, kernel_size=3,
               activation='selu', padding='same', name='conv9_1', kernel_initializer='he_normal')(up9)
#conv9 = Dropout(0.5)(conv9)
conv9 = Conv2D(filters=64, kernel_size=3,
               activation='selu', padding='same', name='conv9_2', kernel_initializer='he_normal')(conv9)
#conv9 = Dropout(0.5)(conv9)
conv9 = Conv2D(filters=2, kernel_size=3, activation='selu',
               padding='same', name='conv9_3', kernel_initializer='he_normal')(conv9)

conv10 = Conv2D(filters=1, kernel_size=1,
                activation='sigmoid', name='conv10_1')(conv9)

model = Model(inputs=inputs, outputs=conv10)
model.load_weights(LOAD_PATH, by_name=True)
pred = model.predict(x, batch_size=8, verbose=1)
print('predicted')
pred = pred * 255
if not os.path.exists(SAVE_PATH):
    os.mkdir(SAVE_PATH)
for i in range(pred.shape[0]):
    scipy.misc.imsave(SAVE_PATH + os.sep +
                      str(i) + '.jpg', pred[i, :, :, 0])
print('predicted image saved')
#np.savez_compressed('pred_test.npz', pred)
print('predicted data saved')


import glob
import netharn as nh
import numpy as np
import ubelt as ub
from fishnet import coco_api
from netharn.models.yolo2 import light_yolo
from fishnet.data import load_coco_datasets, YoloCocoDataset

def predict():
    """
    Currently hacked in due to limited harness support.

    srun -c 4 -p priority --gres=gpu:1 \
            python ~/code/baseline-viame-2018/yolo_viame.py predict \
            --gpu=0
    """

    # HACK: Load the training dataset to extract the categories
    # INSTEAD: Should read the categories from a deployed model
    coco_dsets = load_coco_datasets()
    categories = coco_dsets['train'].dataset['categories']

    # Create a dataset to iterate through the images to predict on
    test_gpaths = glob.glob(ub.truepath('~/data/noaa/test_data/*/*.png'))
    predict_coco_dataset = {
        'licenses': [],
        'info': [],
        'categories': categories,
        'images': [
            {
                'id': idx,
                'file_name': fpath,
            }
            for idx, fpath in enumerate(test_gpaths)
        ],
        'annotations': [],
    }
    predict_coco_dset = coco_api.CocoDataset(predict_coco_dataset, tag='predict')
    predict_dset = YoloCocoDataset(predict_coco_dset, train=False)

    # HACK: Define the path to the model weights
    # INSTEAD: best weights should be packaged in a model deployment
    load_path = ub.truepath(
        '~/work/viame/yolo/fit/nice/baseline1/best_snapshot.pt')
    # load_path = ub.truepath(
    #     '~/work/viame/yolo/fit/nice/baseline1/torch_snapshots/_epoch_00000080.pt')

    # HACK: Define the model topology (because we know what we trained with)
    # INSTEAD: model deployment should store and abstract away the topology
    model = light_yolo.Yolo(**{
        'num_classes': predict_dset.num_classes,
        'anchors': np.asarray([(1.08, 1.19), (3.42, 4.41),
                               (6.63, 11.38), (9.42, 5.11),
                               (16.62, 10.52)],
                              dtype=np.float),
        'conf_thresh': 0.001,
        'nms_thresh': 0.5,
    })

    # Boilerplate code that could be abstracted away in a prediction harness
    xpu = nh.XPU.cast('gpu')
    print('xpu = {!r}'.format(xpu))
    model = xpu.mount(model)
    snapshot_state = xpu.load(load_path)
    model.load_state_dict(snapshot_state['model_state_dict'])

    batch_size = 16
    workers = 4
    predict_loader = predict_dset.make_loader(batch_size=batch_size,
                                              num_workers=workers,
                                              shuffle=False, pin_memory=False)

    letterbox = predict_dset.letterbox

    # HACK: Main prediction loop
    # INSTEAD: Use a prediction harness to abstract these in a similar way to
    # the fit harness.
    predictions = []

    with nh.util.grad_context(False):
        _iter = ub.ProgIter(predict_loader, desc='predicting')
        for bx, raw_batch in enumerate(_iter):
            batch_inputs, batch_labels = raw_batch

            inputs = xpu.variable(batch_inputs)
            labels = {k: xpu.variable(d) for k, d in batch_labels.items()}

            outputs = model(inputs)

            # Transform yolo outputs into the coco format
            postout = model.module.postprocess(outputs)

            indices = labels['indices']
            orig_sizes = labels['orig_sizes']
            inp_size = np.array(inputs.shape[-2:][::-1])
            bsize = len(inputs)
            for ix in range(bsize):
                postitem = postout[ix].data.cpu().numpy()

                orig_size = orig_sizes[ix].data.cpu().numpy()
                gx = int(indices[ix].data.cpu().numpy())
                gid = predict_dset.dset.dataset['images'][gx]['id']

                # Unpack postprocessed predictions
                sboxes = postitem.reshape(-1, 6)
                pred_cxywh = sboxes[:, 0:4]
                pred_scores = sboxes[:, 4]
                pred_cxs = sboxes[:, 5].astype(np.int)

                sortx = pred_scores.argsort()
                pred_scores = pred_scores.take(sortx)
                pred_cxs = pred_cxs.take(sortx)
                pred_cxywh = pred_cxywh.take(sortx, axis=0)

                norm_boxes = nh.util.Boxes(pred_cxywh, 'cxywh')
                boxes = norm_boxes.scale(inp_size)
                pred_box = letterbox._boxes_letterbox_invert(boxes, orig_size,
                                                             inp_size)
                pred_box = pred_box.clip(0, 0, orig_size[0], orig_size[1])

                pred_xywh = pred_box.toformat('xywh').data

                # print(ub.repr2(pred_cxywh.tolist(), precision=2))
                # print(ub.repr2(pred_xywh.tolist(), precision=2))

                for xywh, cx, score in zip(pred_xywh, pred_cxs, pred_scores):
                    if score > 0.1:
                        cid = predict_dset.dset.dataset['categories'][cx]['id']
                        pred = {
                            'id': len(predictions) + 1,
                            'image_id': gid,
                            'category_id': cid,
                            'bbox': list(xywh),
                            'score': score,
                        }
                        predictions.append(pred)
            # if bx > 1:
            #     break

    predict_coco_dset.dataset['annotations'] = predictions
    predict_coco_dset._build_index()

    with open('./viame_pred_dump.mscoco.json') as file:
        predict_coco_dset.dump(file)

    if False:
        import utool as ut
        from matplotlib import pyplot as plt
        gids = set([a['image_id'] for a in predict_coco_dset.anns.values()])
        for gid in ut.InteractiveIter(list(gids)):

            try:
                fig = plt.figure(1)
                fig.clf()
                predict_coco_dset.show_annotation(gid=gid)
                fig.canvas.draw()
            except Exception:
                print('cannot draw')

        z = inputs[0].cpu().numpy().transpose(1, 2, 0)
        nh.util.imshow(z, fnum=2, colorspace='rgb')

if __name__ == '__main__':
    """
    CommandLine:
        python ~/code/baseline-viame-2018/fishnet/predict.py all
    """
    predict()


from __future__ import print_function
import pandas as pd
import numpy as np
import argparse
import subprocess
import os
import gzip
import contextlib
from datetime import datetime
import tempfile
import shutil
import glob
from sklearn.decomposition import PCA

@contextlib.contextmanager
def cd(cd_path):
    saved_path = os.getcwd()
    os.chdir(cd_path)
    yield
    os.chdir(saved_path)

def gtf_to_bed(annotation_gtf, feature='gene'):
    """
    Parse genes from GTF, create placeholder DataFrame for BED output
    """
    chrom = []
    start = []
    end = []
    gene_id = []
    with open(annotation_gtf, 'r') as gtf:
        for row in gtf:
            row = row.strip().split('\t')
            if row[0][0]=='#' or row[2]!=feature: continue # skip header
            chrom.append(row[0])

            # TSS: gene start (0-based coordinates for BED)
            if row[6]=='+':
                start.append(np.int64(row[3])-1)
                end.append(np.int64(row[3]))
            elif row[6]=='-':
                start.append(np.int64(row[4])-1)  # last base of gene
                end.append(np.int64(row[4]))
            else:
                raise ValueError('Strand not specified.')

            gene_id.append(row[8].split(';',1)[0].split(' ')[1].replace('"',''))

    bed_df = pd.DataFrame(
        data={'chr':chrom, 'start':start, 'end':end, 'gene_id':gene_id},
        columns=['chr', 'start', 'end', 'gene_id'],
        index=gene_id)
    return bed_df

def write_bed(bed_df, output_name):
    """Write DataFrame to BED"""
    bgzip = subprocess.Popen('bgzip -c > '+output_name,
        stdin=subprocess.PIPE, shell=True)
    bed_df.to_csv(bgzip.stdin, sep='\t', index=False)
    stdout, stderr = bgzip.communicate()
    subprocess.check_call('tabix -f '+output_name, shell=True)

if __name__=='__main__':

    parser = argparse.ArgumentParser(description='Run leafcutter clustering, prepare for FastQTL')
    parser.add_argument('junc_files_list', help='File with paths to ${sample_id}.junc files')
    parser.add_argument('exons', help='Exon definitions file, with columns: chr, start, end, strand, gene_id, gene_name')
    parser.add_argument('genes_gtf', help='Collapsed gene annotation in GTF format')
    parser.add_argument('prefix', help='Prefix for output files (sample set ID)')
    parser.add_argument('--min_clu_reads', default='50', type=str, help='Minimum number of reads supporting each cluster')
    parser.add_argument('--min_clu_ratio', default='0.001', type=str, help='Minimum fraction of reads in a cluster that support a junction')
    parser.add_argument('--max_intron_len', default='500000', type=str, help='Maximum intron length')
    parser.add_argument('--num_pcs', default=10, type=int, help='Number of principal components to calculate')
    parser.add_argument('--leafcutter_dir', default='/opt/leafcutter',
                        help="leafcutter directory, containing 'clustering' directory")
    parser.add_argument('-o', '--output_dir', default='.', help='Output directory')
    args = parser.parse_args()

    print('['+datetime.now().strftime("%b %d %H:%M:%S")+'] leafcutter clustering')
    if not os.path.exists(args.output_dir):
        os.mkdir(args.output_dir)

    with cd(args.output_dir):

        print('  * decompressing and renaming junc files')
        with open(args.junc_files_list) as f:
            junc_files = f.read().strip().split('\n')

        junc_dir = os.path.join(args.output_dir, 'junc_files')
        if not os.path.exists(junc_dir):
            os.mkdir(junc_dir)
        sample_ids = []
        for f in junc_files:
            sample_id = os.path.split(f)[1].split('.')[0]
            sample_ids.append(sample_id)
            # shutil.move(f, os.path.join(junc_dir, sample_id+'.junc.gz'))
            shutil.copy2(f, os.path.join(junc_dir, sample_id+'.junc.gz'))

        subprocess.check_call('gunzip -f '+os.path.join(junc_dir, '*.junc.gz'), shell=True)
        junc_files = sorted([os.path.join(junc_dir, i+'.junc') for i in sample_ids])

        print('  * running leafcutter clustering')
        # generates ${prefix}_perind_numers.counts.gz and ${prefix}_perind.counts.gz
        with tempfile.NamedTemporaryFile(dir=args.output_dir) as tmp:
            with open(tmp.name, 'w') as f:
                f.write('\n'.join(junc_files)+'\n')
            subprocess.check_call(
                'python '+os.path.join(args.leafcutter_dir, 'clustering', 'leafcutter_cluster.py' \
                    +' --juncfiles '+tmp.name \
                    +' --outprefix '+args.prefix \
                    +' --minclureads '+args.min_clu_reads \
                    +' --mincluratio '+args.min_clu_ratio \
                    +' --maxintronlen '+args.max_intron_len), shell=True)

        print('  * compressing outputs')
        subprocess.check_call('gzip {}_pooled'.format(args.prefix), shell=True)
        subprocess.check_call('gzip {}_refined'.format(args.prefix), shell=True)

        print('  * mapping clusters to genes')
        subprocess.check_call(
            'Rscript' \
                +' '+os.path.abspath(os.path.join(os.path.dirname(__file__), 'map_clusters_to_genes.R')) \
                +' '+os.path.join(args.output_dir, args.prefix+'_perind.counts.gz') \
                +' '+args.exons \
                +' '+args.prefix + '.leafcutter.clusters_to_genes.txt', shell=True)

        print('  * filtering counts')
        counts_df = pd.read_csv(os.path.join(args.output_dir, args.prefix+'_perind.counts.gz'), sep='\s+').set_index('chrom')
        calculate_frac = lambda x: float(x[0])/float(x[1]) if x[1]>0 else 0
        frac_df = counts_df.applymap(lambda x: calculate_frac([int(i) for i in x.split('/')]))
        pct_zero = (frac_df==0).sum(axis=1) / frac_df.shape[1]  # for zero counts, frac is zero
        n_unique = frac_df.apply(lambda x: len(x.unique()), axis=1)
        zscore_df = ((frac_df.T-frac_df.mean(1)) / frac_df.std(1)).T

        # filter out introns with low counts or low complexity
        n = np.floor(frac_df.shape[1]*0.1)
        if n<10:
            n = 10
        mask = (pct_zero <= 0.5) & (n_unique >= n)
        # additional filter for low complexity
        ns = zscore_df.shape[1]
        mask2 = ((zscore_df.abs()<0.25).sum(1)>=ns-3) & ((zscore_df.abs()>6).sum(1)<=3)
        if np.any(mask & mask2):
            print('    ** dropping {} introns with low variation'.format(np.sum(mask & mask2)))
        mask = mask & ~mask2

        filtered_counts_df = counts_df.loc[mask]
        cluster_ids = np.unique(counts_df.index.map(lambda x: x.split(':')[-1]))
        filtered_cluster_ids = np.unique(filtered_counts_df.index.map(lambda x: x.split(':')[-1]))
        print('    ** dropping {} introns with counts in fewer than 50% of samples\n'
              '       {}/{} introns remain ({}/{} clusters)'.format(
                   counts_df.shape[0]-filtered_counts_df.shape[0], filtered_counts_df.shape[0], counts_df.shape[0], len(filtered_cluster_ids), len(cluster_ids)
            ))
        filtered_counts_file = os.path.join(args.output_dir, args.prefix+'_perind.counts.filtered.gz')
        with gzip.open(filtered_counts_file, 'wt') as f:
            filtered_counts_df.to_csv(f, sep=' ')

        print('  * preparing phenotype table')
        subprocess.check_call(
            'python '+os.path.join(args.leafcutter_dir, 'scripts', 'prepare_phenotype_table.py') \
            +' '+filtered_counts_file \
            +' -p '+str(args.num_pcs), shell=True)

        print('  * concatenating BED files')
        bed_files = sorted(glob.glob(os.path.join(args.output_dir, '*_perind.counts.filtered.gz.qqnorm_*')))
        bed_df = []
        for f in bed_files:
            bed_df.append(pd.read_csv(f, sep='\t', dtype=str))
        bed_df = pd.concat(bed_df, axis=0)
        print('    ** sorting')
        # leafcutter doesn't produce output for chrX --> numeric sort
        for c in ['#Chr', 'start', 'end']:
            bed_df[c] = bed_df[c].astype(int)
        bed_df = bed_df.sort_values(['#Chr', 'start', 'end'])
        print('    ** writing BED')
        bed_file = os.path.join(args.output_dir, args.prefix+'.perind.counts.filtered.qqnorm.bed.gz')
        bgzip = subprocess.Popen('bgzip -c > '+bed_file, stdin=subprocess.PIPE, shell=True, universal_newlines=True)
        stdout, stderr = bgzip.communicate(bed_df.to_csv(sep='\t', index=False))
        print('    ** indexing')
        subprocess.check_call('tabix '+bed_file, shell=True)

        print('  * converting cluster coordinates to gene coordinates')
        tss_df = gtf_to_bed(args.genes_gtf)
        cluster2gene_dict = pd.read_csv(os.path.join(args.output_dir, args.prefix + '.leafcutter.clusters_to_genes.txt'),
            sep='\t', index_col=0, squeeze=True).to_dict()

        # add 'chr' prefix
        bed_df['#Chr'] = 'chr'+bed_df['#Chr'].astype(str)
        bed_df['ID'] = 'chr'+bed_df['ID']

        print('    ** assigning introns to gene mapping(s)')
        n = 0
        gene_bed_df = []
        group_s = {}
        for _,r in bed_df.iterrows():
            s = r['ID'].split(':')
            cluster_id = s[0]+':'+s[-1]
            if cluster_id in cluster2gene_dict:
                gene_ids = cluster2gene_dict[cluster_id].split(',')
                for g in gene_ids:
                    gi = r['ID']+':'+g
                    gene_bed_df.append(tss_df.loc[g, ['chr', 'start', 'end']].tolist() + [gi] + r.iloc[4:].tolist())
                    group_s[gi] = g
            else:
                n += 1
        if n>0:
            print('    ** discarded {} introns without gene mapping'.format(n))
        print('  * writing FastQTL inputs')
        gene_bed_df = pd.DataFrame(gene_bed_df, columns=bed_df.columns)
        gene_bed_df = gene_bed_df.groupby('#Chr', sort=False, group_keys=False).apply(lambda x: x.sort_values('start'))
        # change sample IDs to participant IDs
        gene_bed_df.rename(columns={i:'-'.join(i.split('-')[:2]) for i in gene_bed_df.columns[4:]}, inplace=True)
        write_bed(gene_bed_df, os.path.join(args.output_dir, args.prefix+'.leafcutter.bed.gz'))
        pd.Series(group_s).sort_values().to_csv(os.path.join(args.output_dir, args.prefix+'.leafcutter.phenotype_groups.txt'), sep='\t')

        print('  * calculating PCs')
        pca = PCA(n_components=args.num_pcs)
        pca.fit(bed_df[bed_df.columns[4:]])
        pc_df = pd.DataFrame(pca.components_, index=['PC{}'.format(i) for i in range(1,11)],
            columns=['-'.join(i.split('-')[:2]) for i in bed_df.columns[4:]])
        pc_df.index.name = 'ID'
        pc_df.to_csv(args.prefix+'.leafcutter.PCs.txt', sep='\t')

print('['+datetime.now().strftime("%b %d %H:%M:%S")+'] done')


import skimage.io as io
import skimage.transform as skt
import numpy as np
from PIL import Image
from src.models.class_patcher import patcher
from src.utils.imgproc import *

class patcher(patcher):
    def __init__(self, body='./body/body_chronos.png', **options):
        super().__init__('クロノス', body=body, pantie_position=[0, 2], **options)
        self.mask = io.imread('./mask/mask_chronos.png')

    def convert(self, image):
        pantie = np.array(image)
        patch = np.copy(pantie[-100:-5, 546:, :])
        pantie[-100:, 546:, :] = 0
        patch = skt.resize(patch[::-1, ::-1, :], (patch.shape[0] + 30, patch.shape[1]), anti_aliasing=True, mode='reflect')
        [pr, pc, d] = patch.shape
        pantie[127 - 5:127 - 5 + pr, :pc, :] = np.uint8(patch * 255)

        front = pantie[:350, :300]
        arrx = (np.linspace(0, 1, 25)**2) * 43
        arrx[5:20] += np.sin(np.linspace(0, np.pi, 15)) * 4
        arrx[5:10] += np.sin(np.linspace(0, np.pi, 5)) * 13
        arry = np.zeros(25)
        arrx -= 30
        front = affine_transform_by_arr(front, arrx, arry)
        front = np.uint8(front[::-1, 7:][88:] * 255)

        back = pantie[:350, 300:-10][:, ::-1]
        arrx = (np.linspace(0, 1, 25)**2) * 115
        arrx[2:14] += np.sin(np.linspace(0, np.pi, 12)) * 7
        arry = np.zeros(25)
        arrx -= 70
        back = affine_transform_by_arr(back, arrx, arry)
        back = np.uint8(back[3:, 10:10 + front.shape[1]] * 255)

        pantie = np.concatenate((back, front), axis=0)
        pantie = np.uint8(resize(pantie, [1.55, 1.745]) * 255)
        pantie = np.bitwise_and(pantie, self.mask)
        pantie = np.concatenate((pantie[:, ::-1], pantie), axis=1)
        return Image.fromarray(pantie)


import sys
import cv2 as cv
import operator
import numpy as np
import math
import json

use_mask = False
img = None
templ = None
match_method = 3
IMG_PADDING_TOP = 150

def exceedsBlackRatio(img, ratio):
    return cv.countNonZero(cv.cvtColor(img, cv.COLOR_BGR2GRAY)) < (1.0-ratio)*img.shape[0]*img.shape[1]

def main(argv):
    if (len(argv) < 1):
        print('Not enough parameters')
        print('Usage:\nocv.py <data file path>')
        return -1

    with open(argv[0]) as json_file:
        data = json.load(json_file)

    for sheetData in data:
        image_path = sheetData['path']
        img = cv.imread(image_path, cv.IMREAD_COLOR)

        try:
            # overlay sprite sheet over an equally sized black material: https://stackoverflow.com/a/14102014/5920409
            # I failed to come up with a more elegant way to get rid of the alpha channel while avoiding the underlying color artifacts.
            alpha = cv.imread(image_path, cv.IMREAD_UNCHANGED)[:,:,3] / 255.0
            # black = np.zeros((img.shape[0],img.shape[1],3), np.uint8)
            for c in range(0, 3):
                #img[:, :, c] = (alpha * img[:, :, c] + (1.0 - alpha) * black[:, :, c])
                img[:, :, c] = alpha * img[:, :, c]     # for a black background the commented out calculation gets reduced to this, as we're multiplying by 0
            img_black_ratio_limit = 0.95
            if(exceedsBlackRatio(img, img_black_ratio_limit)):   # if the source image is almost entirely black, invert it before overlaying.
                print('Image is over {}% black. Inverting colors.'.format(img_black_ratio_limit*100))
                img = cv.imread(image_path, cv.IMREAD_COLOR)
                alpha = cv.imread(image_path, cv.IMREAD_UNCHANGED)[:,:,3] / 255.0
                img = cv.bitwise_not(img)
                for c in range(0, 3):
                    img[:, :, c] = alpha * img[:, :, c]
        except IndexError:  # occurs if image has no alpha channel. In this case we can just take the image as is.
            pass

        if (img is None):
            print('Can\'t read the image')
            return -1

        img = cv.copyMakeBorder(img,IMG_PADDING_TOP,0,0,0,cv.BORDER_CONSTANT,value=0) # used to find faces in negative y area
        sheetData = calcFaces(img, sheetData)

    with open(argv[0], 'w') as json_file:
        json.dump(data, json_file)

    return 0

def calcFaces(img, sheetData):
    # finds most likely face dimensions (currently checks for [32 / 64 / ...])
    results = {}

    min_face_dim = 256   # minimum detectable face expression sprite size
    max_black_ratio = 0.75   # 25% of a face segment must be non-black to be used for matching
    min_match_face = 0.9 # face has to match another face by 90% or it unlikely to actually be a face
    min_match_face_body = 0.50

    SPRITE_FORMAT_NORMAL = 0            # a body sprite and at least one expression sprite below
    SPRITE_FORMAT_FULLBODY = 1          # full body sprites

    row_count = img.shape[0]    # = height
    max_y = row_count-1
    col_count = img.shape[1]    # = width
    max_x = col_count-1
    face_reference = None

    # recursive function to find a proper (=non-black) face sample for the given dimensions [dim x dim]
    def findFaceLeftBound(img, dim, _x1=0, _y1=None):
        # bottom left, using previous face as offset
        height = img.shape[0]
        width = img.shape[1]
        y1 = _y1 if _y1 is not None else height - dim
        x1 = _x1

        # if face exceeds right edge bounds, restart at upper row
        if(x1 >= width):
            x1 = 0
            y1 -= dim
        # if face exceeds bottom or top bounds, there is no face for these dimensions (y1 < 0 returns empty array even if y2 is in bounds)
        if(y1 >= height or y1 < 0):
            return None, None, None

        face = img[y1:y1+dim, x1:x1+dim]
        if(exceedsBlackRatio(face, max_black_ratio)):
            print("{} {} {} is too black.".format(dim, x1, y1))
            face, x1, y1 = findFaceLeftBound(img, dim, x1+dim, y1)
        return face, x1, y1

    d = min_face_dim
    while d*2 <= row_count and d*2 <= col_count:
        face1, f1x1, f1y1 = findFaceLeftBound(img, d)
        if(face1 is not None):
            face2, f2x1, f2y1 = findFaceLeftBound(img, d, f1x1+d, f1y1)
        if(face1 is not None and face2 is not None):
            # cv.namedWindow("{} face1 {} {}".format(d, f1x1, f1y1), cv.WINDOW_AUTOSIZE)
            # cv.imshow("{} face1 {} {}".format(d, f1x1, f1y1), face1)
            # cv.waitKey(0)
            # cv.namedWindow("{} face2 {} {}".format(d, f2x1, f2y1), cv.WINDOW_AUTOSIZE)
            # cv.imshow("{} face2 {} {}".format(d, f2x1, f2y1), face2)
            # cv.waitKey(0)
            matched = cv.matchTemplate(face1, face2, match_method)
            if(matched[0][0] > min_match_face):
                results[(d, "left")] = (matched[0][0], face1)
        else:
            print("Dimensions of {} skipped due to bad face samples.".format(d))
        d *= 2

    while(len(results) > 0):
        print(results)

        # Switched highest probability to smallest dimension. (Wrong) large dimensions tend to be rated too highly
        # when compared to their neighbours and the results are already filtered for good probability.

        # highestProbabilityKey = max(results, key = lambda k : results.get(k)[0])   # get key of highest probability (large dimensions tend to rate too high though)
        smallestDimensionKey = min(results.keys(), key = lambda k : [k[0] for k in results.keys()])   # get key of smallest dimension from results
        dim = smallestDimensionKey[0]
        dirStr = smallestDimensionKey[1]
        certainty = results.get(smallestDimensionKey)[0]
        face_reference = results.get(smallestDimensionKey)[1]
        print("Most probable dimensions for this sprite sheet's faces:", dim, "aligned to the", dirStr, "side.")

        # find body height bottom to top. The first two rows missing a face are declared the bottom
        # of the body sprite. This may causes extremely rare false positives on sprites missing a face
        # on their first face slot of two following rows, but that's fine with me.
        y1 = row_count-dim
        y2 = row_count
        x1 = 0
        x2 = dim
        body_y = None
        y_offset = 0
        while(y1-y_offset > 0):
            face = img[ y1-y_offset   :   y2-y_offset , x1   :   x2]
            matched = cv.matchTemplate(face_reference, face, match_method)
            if(matched[0][0] < min_match_face):
                print("{} -> ({}, {}) => ({}, {}) is not a face. (Might be part of the body sprite.)"
                    .format(matched[0][0], x1, y1-y_offset, x2, y2-y_offset))
                if(body_y):     # stop scanning if we found missing faces on two rows in a row. (hah)
                    break
                body_y = y2-y_offset
                print("New Body sprite prediction: ({},{}) => ({},{})".format(0, 0, col_count, body_y))
            elif(body_y):
                body_y = None     # reset body prediction if we found a face on the following row.
                # win_name = "Face found on row"
                # cv.namedWindow( win_name, cv.WINDOW_AUTOSIZE )
                # cv.imshow(win_name, face)
                # cv.waitKey(0)
            y_offset += dim
        body = img[:body_y, :col_count]
        # bwin_name = "Body Prediction"
        # cv.namedWindow(bwin_name, cv.WINDOW_AUTOSIZE)
        # cv.imshow(bwin_name, body)
        # cv.waitKey(0)
        # face_ref_win_name = "Chosen Face Reference"
        # cv.namedWindow(face_ref_win_name, cv.WINDOW_AUTOSIZE)
        # cv.imshow(face_ref_win_name, face_reference)
        # cv.waitKey(0)

        # Do a final search for our reference face on our found body sprite. This sorts out some false positives.
        if(face_reference.shape[0] > 0 and face_reference.shape[1] > 0
                and body.shape[0] >= face_reference.shape[0] and body.shape[1] >= face_reference.shape[1]):
            face_body_match = cv.matchTemplate(body, face_reference, match_method)
            _minVal, maxVal, _minLoc, maxLoc = cv.minMaxLoc(face_body_match, None)
            if(maxVal < min_match_face_body):
                print("Failed to find face on body. Assuming false positive. ({})".format(maxVal))
                results.pop(smallestDimensionKey) # remove false positive from results
            else:
                certainty = (certainty+maxVal)/2
                matchLoc = maxLoc

                # final data for face.js. Everything seems to have gone well.
                sheetData["bodyHeight"] = body_y - IMG_PADDING_TOP  # subtracts image padding used to find faces in negative y
                sheetData["bodyWidth"] = col_count
                sheetData["eWidth"] = sheetData["eHeight"] = dim
                sheetData["headX"] = matchLoc[0]
                sheetData["headY"] = matchLoc[1] - IMG_PADDING_TOP # subtracts image padding used to find faces in negative y
                sheetData["dialogOffsetX"] = 0
                # No idea how to even estimate values for dialogOffsetY.
                # Would probably need some big data set and some big maths. Or machine learning.
                # Right now I just adjust values by hand as needed. Manual entries are mostly done by granting 10px of space
                # above (physical) head. Could probably do something similiar programmically.
                # sheetData["dialogOffsetY"] = 0
                sheetData["specialFormat"] = SPRITE_FORMAT_NORMAL
                sheetData["certainty"] = certainty
                break # we found our desired result, so break the result loop. Not pretty, I know.
        else:
            # This is just a fail safe from the past and might never actually happen in this project's current state,
            # but I can't be bothered checking right now.
            print("Invalid face and/or body dimensions.")
            sheetData["bodyHeight"] = row_count - IMG_PADDING_TOP # subtract image padding used to find faces in negative y
            sheetData["bodyWidth"] = col_count
            sheetData["specialFormat"] = SPRITE_FORMAT_FULLBODY
            sheetData["certainty"] = 0.0
    else:   # results became empty
        dim = None
        face_reference = None
        print("No probable faces detected. This is likely a full body sprite.")

        # final data for face.js if sheet is a full body sprite (= we failed to find face sprites)
        sheetData["bodyHeight"] = row_count - IMG_PADDING_TOP # subtracts image padding used to find faces in negative y
        sheetData["bodyWidth"] = col_count
        # sheetData["dialogOffsetX"] = 0
        # sheetData["dialogOffsetY"] = 0
        sheetData["specialFormat"] = SPRITE_FORMAT_FULLBODY   # full body sprite
        sheetData["certainty"] = 0.0
    print(sheetData)
    return sheetData

if __name__ == "__main__":
    main(sys.argv[1:])


import argparse
import copy
import glob
import json
import os
import sys
from itertools import groupby, count, chain
from typing import List

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

from imapper.logic.joints import Joint
from imapper.logic.scenelet import Scenelet
from imapper.pose.filter_full_fit import Span2
from imapper.pose.span import Span

import numpy as np
from scipy.ndimage.filters import gaussian_filter

_lower_body = [Joint.LKNE, Joint.RKNE,
               Joint.LANK, Joint.RANK,
               Joint.LHIP, Joint.RHIP]

def _check_exists(l):
    if os.path.exists(l):
        return l
    else:
        l = os.path.join(os.path.dirname(__file__), l)
        if os.path.exists(l):
            return l

    raise RuntimeError("Path does not exist: \"%s\"" % l)

def get_pose_diff(pose_q, pose_m):
    """Estimates distance with aligned pelvises in 2D.

    Args:
        pose (np.ndarray): (3, 16) Pose to compare.
        q_pose (np.ndarray): (3, 16) Pose to compare to.

    Returns:
        (float): Pose distance
    """

    # TODO: Project onto ground instead of assume y-down
    pose_m[[0, 2]] -= (pose_m[:, Joint.PELV:Joint.PELV + 1]
                       - pose_q[:, Joint.PELV:Joint.PELV + 1])[[0, 2]]
    diff = np.mean(np.linalg.norm(pose_q[:, _lower_body]
                                  - pose_m[:, _lower_body], axis=0))

    return diff

def group_spans(above):
    """Groups consecutive entries to spans.

    Args:
        above (List[Span2]): Sorted spans with length == 1, having
        distance above the distance threshold.

    Returns:
         (List[Span2]): Spans with length >=1, the min distance between
         two consecutive spans is now > 1.
         Distance := second.start - first.end.

    """
    if not len(above):
        return []

    assert all(a.end <= b.start for a, b in zip(above[:-1], above[1:])), \
        'Assumed sorted spans: {}'.format(above)
    spans = [copy.deepcopy(above[0])]
    it = iter(above)
    next(it)
    prev_frame_id = above[0].start
    for span2 in it:
        frame_id = span2.start
        if prev_frame_id + 1 < frame_id:
            spans[-1].end = prev_frame_id
            spans.append(Span2(start=frame_id, end=frame_id,
                               value=span2.value))
        prev_frame_id = frame_id
    spans[-1].end = prev_frame_id

    return spans

def get_delta_frame_id(span0, span1):
    """Computes the minimum distance between two non-overlapping spans.

    Args:
         span0 (Span): First span.
         span1 (Span): Second span.
    """
    if span0.overlaps(span1):
        assert False, (span0, span1)
        return 0

    if span0.end < span1.start:
        return span1.start - span0.end
    else:
        return span0.start - span1.end

def filter_spans_to_replace(spans, chosen, c_threshold):
    """Estimates remaining spans that are not covered in chosen.

    Args:
        spans (List[Span2]):
            Spans to replace because of distance (will be used without
            objects).
        chosen (List[Span]):
            Spans chosen for candidate generation (will be used with
            objects).
        c_threshold (float):
            Charness threshold.

    Returns:
        (List[Span2]): Spans to add to gap_command.sh with the
        --remove-objects flag.
    """
    # keep track of activated actor ids (we'll keep one each)
    actor_ids = set()

    # sort in decreasing characteristicness order
    chosen_srtd = sorted(chosen, key=lambda _c: _c.charness, reverse=True)

    # filter through all spans
    for c in chosen_srtd:

        # if not characteristic enough and actor already covered
        if c.actor_id in actor_ids and c.charness < c_threshold:
            c.active = False
            continue

        # we'll need at least one scenelet for each actor
        actor_ids.add(c.actor_id)

        # find closest active Span
        try:
            closest = min(
                (c2 for c2 in chosen if c2 != c and c2.active),
                key=lambda _c: get_delta_frame_id(c, _c))

            # time difference to closest active Span
            dt = get_delta_frame_id(c, closest)

            # remove, if has a more characteristic neighbour that's active
            # TODO: this might be unnecessary, could be swapped for Span2
            # immediately
            if dt == 1 and closest.active and closest.charness > c.charness:
                c.active = False
        except ValueError as error:
            pass

    # re-activate inactive ones above threshold without objects (Span2)
    # for i in range(len(chosen_srtd)):
    #     c = chosen_srtd[i]
    #     if c.active or c.charness < c_threshold:
    #         continue
    #
    #     # Span2 means without objects
    #     chosen_srtd[i] = Span2(start=c.start, end=c.end, value=99.,
    #                            actor_id=c.actor_id)

    # find uncovered parts
    chosen_set = frozenset(
        f for s in chosen for f in range(s.start, s.end + 1) if s.active)
    # gather output spans
    o_spans = []
    for span in spans:
        frames = frozenset(f for f in range(span.start, span.end + 1))
        remainder = sorted(frames.difference(chosen_set))
        if not len(remainder):
            continue

        for _, it in groupby(remainder, lambda n, k=count(): n - next(k)):
            l_it = list(it)
            start, end = l_it[0], l_it[-1]
            if start == end:
                end = end + 1
            o_spans.append(Span2(start=start, end=end, value=span.value,
                                 actor_id=span.actor_id))
    print('o_spans: {}'.format(o_spans))

    return chosen, o_spans

def get_pose_distance(query3d, path_match, gap):
    skel_q = query3d.skeleton
    mid_frame = gap[0] + (gap[1] - gap[0]) // 2
    frame_q = min((frame_id for frame_id in skel_q.get_frames()),
                  key=lambda frame_id: abs(frame_id - mid_frame))
    time = skel_q.get_time(frame_id=frame_q)
    print('Closest to {} is {} with time {}'.format(mid_frame, frame_q, time))
    match = Scenelet.load(path_match)
    skel_m = match.skeleton
    frame_m = skel_m.find_time(time=time)
    time_m = skel_m.get_time(frame_id=frame_m)
    print('Closest match is {} with time {}'.format(frame_m, time_m))

    diffs = []
    for frame_id in range(frame_q - 1, frame_q + 2):
        if not skel_q.has_pose(frame_id):
            print('Skipping frame_id {} in query because missing'
                  .format(frame_id))
            continue
        _time_q = skel_q.get_time(frame_id=frame_id)
        _frame_m = skel_m.find_time(time=_time_q)
        _time_m = skel_m.get_time(frame_id=_frame_m)
        if abs(_time_m - _time_q) > 1.:
            print('Skipping matched time of {} because too far from {}'
                  .format(_time_m, _time_q))
            continue
        pose_q = skel_q.get_pose(frame_id=frame_id)
        pose_m = skel_m.get_pose(frame_id=_frame_m)
        diff = get_pose_diff(pose_q=pose_q, pose_m=pose_m)
        print('Diff: {}'.format(diff))
        diffs.append(diff)

    # return np.mean(diffs), len(diffs) > 0
    return np.max(diffs), len(diffs) > 0

def find_surrounding(span, spans):
    """Finds the span in spans that best surrounds span.

    Args:
        span (Span): Span to surround.
        spans (List[Span]): Surrounding candidates.

    Returns:
        (Span): Surrounding span.
    """
    mid = (span.start + span.end) // 2
    return max((s for s in spans if s.start < mid < s.end),
               key=lambda _s: min(_s.end - mid, mid - _s.start))

def construct_command(name_query, span, tc, surr, path_scenes_root,
                      path_scenelets_db):
    comm = '# ' if not span.active else ''

    # Replace bad walking poses in input
    remove_objects = ' --remove-objects' \
        if span.active and isinstance(span, Span2) \
        else ''

    cmd = "{comm:s}D=\"{root:s}/{scene:s}/opt2b/{t0:03d}_{t1:03d}\";\n" \
          "{comm:s}if [ ! -d ${{D}} ]; then\n" \
          "{comm:s}\tpython3 imapper/pose/opt_consistent.py -silent " \
          "--wp 1 --ws 0.05 --wi 1. --wo 1. " \
          "-w-occlusion -w-static-occlusion --maxiter 15 " \
          "--nomocap -v " \
          "{root:s}/{scene:s}/skel_{scene:s}_unannot.json " \
          "independent " \
          "-s {scenelets_db:s} " \
          "--gap {t0:03d} {t1:03d} --batch-size ${{BATCH_SIZE}} " \
          "--dest-dir opt2 " \
          "--candidates opt1/{s0:03d}_{s1:03d} --n-candidates 200 " \
          "-tc {tc:.2f}{remove_objects:s}\n" \
          "{comm:s}fi" \
        .format(scene=name_query, t0=span.start, t1=span.end,
                tc=tc, s0=surr.start, s1=surr.end,
                root=path_scenes_root, comm=comm,
                remove_objects=remove_objects,
                scenelets_db=path_scenelets_db)
    return cmd

class LoggedSequence(object):
    def __init__(self):
        self._data = []

    def add_point(self, value, time):
        assert isinstance(time, (int, float, np.float32, np.float64)), \
            print("Expected time, not {} ({})".format(time, type(time)))
        self._data.append([value, time])

    def plot(self, path_root):
        values, times = zip(*self._data)
        print('values: {}\ntimes: {}'.format(values, times))
        f = plt.figure()
        ax = f.add_subplot(111)

        ax.plot(times, values, 'x--',
                label='Distance to best Kinect fit\'s center frame')
        plt.xlabel('Time (s)')
        plt.ylabel('Sum local squared distance')
        plt.legend()
        plt.savefig(os.path.join(path_root, 'tmp.pdf'))

def main(argv):
    pjoin = os.path.join

    parser = argparse.ArgumentParser("")
    parser.add_argument(
        'd', type=_check_exists,
        help="Input directory")
    parser.add_argument(
        'db', type=_check_exists,
        help="Path to scenelets database"
    )
    parser.add_argument(
        '-o', '--opt-folder',
        help="Which optimization output to process. Default: opt1",
        default='opt1')
    parser.add_argument(
        '-limit', type=int,
        help="How many scenelets to aggregate.",
        default=3
    )
    parser.add_argument(
        '--c-threshold',
        help='Distance threshold. Default: 0.3',
        type=float, default=0.3)
    parser.add_argument(
        '--d-threshold',
        help='Distance threshold. Default: 0.4',
        type=float, default=0.4)
    args = parser.parse_args(argv)
    path_scenes_root = os.path.normpath(pjoin(args.d, os.pardir))
    print("Working with %s" % args.d)
    d = pjoin(args.d, args.opt_folder)

    name_query = d.split(os.sep)[-2]
    query3d = Scenelet.load(
        os.path.join(args.d,
                     'skel_{:s}_unannot.json'.format(name_query)))
    n_actors = query3d.skeleton.n_actors  # type: int

    log = LoggedSequence()
    elements = [[], [], []]
    spans = []
    above = []
    for p in sorted(os.listdir(d)):
        d_time = pjoin(d, p)
        if not os.path.isdir(d_time) or 'bak' in p:
            continue
        parts = p.split('_')
        start = int(parts[0])
        end = int(parts[-1])

        sum_charness = 0.
        sum_weight = 0.
        sum_charness_unw = 0.
        sum_weight_unw = 0
        diff_succ, diff = False, 100.
        for f in glob.iglob("%s/skel_%s_*.json" % (d_time, name_query)):
            rank = int(os.path.splitext(f)[0].rpartition('_')[-1])
            if rank >= args.limit:
                continue
            # print(f)
            data = json.load(open(f, 'r'))
            charness = data['charness']
            weight = max(0., .1 - data['score_fit'])
            sum_charness += weight * charness
            sum_weight += weight
            sum_charness_unw += charness
            elements[sum_weight_unw].append(charness)
            sum_weight_unw += 1

            if rank == 0:
                diff, diff_succ = get_pose_distance(
                    query3d=query3d, path_match=f, gap=(start, end))

        if sum_weight > 0.:
            sum_charness /= sum_weight
        if sum_weight_unw > 0.:
            sum_charness_unw /= sum_weight_unw
        frame_id = (start + end) // 2
        actor_id = query3d.skeleton.get_actor_id(frame_id=frame_id) \
            if n_actors > 1 else 0

        spans.append(Span(start, end, sum_charness, sum_charness_unw,
                          actor_id=actor_id))

        # check for pose replacement
        if diff_succ and diff > args.d_threshold:
            # time = query3d.skeleton.get_time(frame_id=frame_id)
            above.append(Span2(start=frame_id, end=frame_id, value=diff))

        if diff_succ:
            log.add_point(value=diff, time=frame_id)

    for actor_id in range(n_actors):
        cs, span_ids = zip(*[(span.charness, span_id)
                             for span_id, span in enumerate(spans)
                             if span.actor_id == actor_id])
        cs2 = gaussian_filter(cs, sigma=2.5).tolist()

        for smoothed, span_id in zip(cs2, span_ids):
            spans[span_id].smoothed_charness = smoothed

    plt.figure()
    plt.plot(cs, label="orig")
    plt.plot([span.smoothed_charness for span in spans], label="smoothed")
    # c0 = [span.charness_unw for span in spans]
    # plt.plot(c0, 'o-', label="unweighted")
    # for sigma in [2.5]:
    #     cs2 = gaussian_filter(cs, sigma=sigma)
    #     plt.plot(cs2, '+-', label="%f" % sigma)

    # for i, elem in enumerate(elements):
    #     plt.plot(elem, 'x-', label="cand%d" % i)
    plt.legend()
    # plt.show()
    p_graph = os.path.join(args.d, 'charness2.svg')
    plt.savefig(p_graph)

    log.plot(path_root=args.d)

    spans = sorted(spans, key=lambda s: s.smoothed_charness, reverse=True)

    chosen = [spans[0]]
    for span_id in range(1, len(spans)):
        span = spans[span_id]
        overlap = next((c for c in chosen if span.overlaps(c)), None)
        if overlap is not None:
            # print("Skipping %s because overlaps %s" % (span, overlap))
            continue
        elif span.start == query3d.skeleton.get_frames()[0] \
                or span.end == query3d.skeleton.get_frames()[-1]:
            print("Skipping %s because first/last frame." % span)
            continue
        else:
            chosen.append(span)

    # Spans to replace, because of quality, if not replaced already
    spans_r = group_spans(above=above)
    chosen, spans_r_filtered = filter_spans_to_replace(
        spans=spans_r, chosen=chosen, c_threshold=args.c_threshold)
    print('Filtered: {}\nchosen: {}'.format(spans_r_filtered, chosen))
    # keep at least one characteristic pose for each actor
    # actors_output = set()
    p_cmd = os.path.join(args.d, 'gap_command_new.sh')
    with open(p_cmd, 'w') as f:
        f.write('BATCH_SIZE=8\n\n')
        for actor_id in range(n_actors):
            f.write('######\n# Actor {:d}\n######\n\n'.format(actor_id))
            for span in chain(chosen, spans_r_filtered):
                if not span.actor_id == actor_id:
                    continue
                surr = find_surrounding(span, spans)
                cmd = construct_command(
                    name_query=name_query, span=span, tc=args.c_threshold,
                    surr=surr, path_scenes_root=path_scenes_root,
                    path_scenelets_db=args.db)

                if isinstance(span.charness, float):
                    f.write("# charness: %g\n" % span.charness)

                # Comment out if too low
                # if span.charness < args.c_threshold \
                #         and actor_id in actors_output:
                # actors_output.add(actor_id)

                f.write('{}\n\n'.format(cmd))
            f.write('\n\n')
        # with open(pjoin(d_time, 'avg_charness.json')) as fch:
        #     data = json.load(fch)
        #     charness = data['avg_charness']
        #     print(charness)

if __name__ == '__main__':
    main(sys.argv[1:])


#  Version: 2020.02.21
#
#  MIT License
#
#  Copyright (c) 2018 Jiankang Deng and Jia Guo
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to deal
#  in the Software without restriction, including without limitation the rights
#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#  copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#
#  The above copyright notice and this permission notice shall be included in all
#  copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#  SOFTWARE.
#

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from scipy import misc
import sys
import os
import argparse
#import tensorflow as tf
import numpy as np
import mxnet as mx
import random
import cv2
import sklearn
from sklearn.decomposition import PCA
from time import sleep
from easydict import EasyDict as edict
from mtcnn_detector import MtcnnDetector
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'src', 'common'))
import face_image
import face_preprocess

def do_flip(data):
  for idx in range(data.shape[0]):
    data[idx,:,:] = np.fliplr(data[idx,:,:])

def get_model(ctx, image_size, model_str, layer):
  _vec = model_str.split(',')
  assert len(_vec)==2
  prefix = _vec[0]
  epoch = int(_vec[1])
  print('loading',prefix, epoch)
  sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
  all_layers = sym.get_internals()
  sym = all_layers[layer+'_output']
  model = mx.mod.Module(symbol=sym, context=ctx, label_names = None)
  #model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))], label_shapes=[('softmax_label', (args.batch_size,))])
  model.bind(data_shapes=[('data', (1, 3, image_size[0], image_size[1]))])
  model.set_params(arg_params, aux_params)
  return model

class FaceModel:
  def __init__(self, args):
    self.args = args
    ctx = mx.gpu(args.gpu)
    _vec = args.image_size.split(',')
    assert len(_vec)==2
    image_size = (int(_vec[0]), int(_vec[1]))
    self.model = None
    self.ga_model = None
    if len(args.model)>0:
      self.model = get_model(ctx, image_size, args.model, 'fc1')
    if len(args.ga_model)>0:
      self.ga_model = get_model(ctx, image_size, args.ga_model, 'fc1')

    self.threshold = args.threshold
    self.det_minsize = 50
    self.det_threshold = [0.6,0.7,0.8]
    #self.det_factor = 0.9
    self.image_size = image_size
    mtcnn_path = os.path.join(os.path.dirname(__file__), 'mtcnn-model')
    if args.det==0:
      detector = MtcnnDetector(model_folder=mtcnn_path, ctx=ctx, num_worker=1, accurate_landmark = True, threshold=self.det_threshold)
    else:
      detector = MtcnnDetector(model_folder=mtcnn_path, ctx=ctx, num_worker=1, accurate_landmark = True, threshold=[0.0,0.0,0.2])
    self.detector = detector

  def get_input(self, face_img):
    ret = self.detector.detect_face(face_img, det_type = self.args.det)
    if ret is None:
      return None
    bbox, points = ret
    if bbox.shape[0]==0:
      return None
    bbox = bbox[0,0:4]
    points = points[0,:].reshape((2,5)).T
    #print(bbox)
    #print(points)
    nimg = face_preprocess.preprocess(face_img, bbox, points, image_size='112,112')
    nimg = cv2.cvtColor(nimg, cv2.COLOR_BGR2RGB)
    aligned = np.transpose(nimg, (2,0,1))
    return aligned

  def get_feature(self, aligned):
    input_blob = np.expand_dims(aligned, axis=0)
    data = mx.nd.array(input_blob)
    db = mx.io.DataBatch(data=(data,))
    self.model.forward(db, is_train=False)
    embedding = self.model.get_outputs()[0].asnumpy()
    embedding = sklearn.preprocessing.normalize(embedding).flatten()
    return embedding

  def get_ga(self, aligned):
    input_blob = np.expand_dims(aligned, axis=0)
    data = mx.nd.array(input_blob)
    db = mx.io.DataBatch(data=(data,))
    self.ga_model.forward(db, is_train=False)
    ret = self.ga_model.get_outputs()[0].asnumpy()
    g = ret[:,0:2].flatten()
    gender = np.argmax(g)
    a = ret[:,2:202].reshape( (100,2) )
    a = np.argmax(a, axis=1)
    age = int(sum(a))

    return gender, age


#!/usr/bin/env python
# -*- coding: utf-8 -*-

# ___        InstaBot V 1.2.0 by nickpettican           ___
# ___        Automate your Instagram activity           ___

# ___        Copyright 2017 Nicolas Pettican            ___

# ___   This software is licensed under the Apache 2    ___
# ___   license. You may not use this file except in    ___
# ___   compliance with the License.                    ___

# ___                   DISCLAIMER                      ___

# ___ InstaBot was created for educational purposes and ___
# ___ the end-user assumes sole responsibility of any   ___
# ___ consequences of it's misuse. Please be advised of ___
# ___ Instagram's monitoring, 1000 likes a day is the   ___
# ___ maximum you will get or you risk a temporary ban. ___
# ___ Choose your tags wisely or you may risk liking    ___
# ___ and commenting on undesirable media or spam.      ___

from lxml import etree
import json, itertools, socket
import numpy as np

def internet_connection(host = '8.8.8.8', port = 53, timeout = 3):
    # check for internet connection

    try:
        socket.setdefaulttimeout(timeout)
        socket.socket(socket.AF_INET, socket.SOCK_STREAM).connect((host, port))
        return True
    except:
        return False

def return_random_sequence(data, time_in_day):
    # generates random delay values that add up to time in day

    if data == 0:
        return [0]
    random_values = np.random.random(data)
    random_values /= random_values.sum()
    return [int(i*time_in_day) for i in random_values]


import os
import warnings
from pathlib import Path

import nnabla as nn
import numpy as np
from api.separators.util import download_and_verify
from billiard.pool import Pool
from nnabla.ext_utils import get_extension_context
from spleeter.audio.adapter import AudioAdapter
from tqdm import trange
from xumx.test import separate

MODEL_URL = 'https://nnabla.org/pretrained-models/ai-research-code/x-umx/x-umx.h5'

class XUMXSeparator:
    """Performs source separation using X-UMX API."""
    def __init__(
        self,
        cpu_separation: bool,
        bitrate=256,
        softmask=False,
        alpha=1.0,
        iterations=1
    ):
        """Default constructor.
        :param config: Separator config, defaults to None
        """
        self.model_file = 'x-umx.h5'
        self.model_dir = Path('pretrained_models')
        self.model_file_path = self.model_dir / self.model_file
        self.context = 'cpu' if cpu_separation else 'cudnn'
        self.softmask = softmask
        self.alpha = alpha
        self.iterations = iterations
        self.bitrate = bitrate
        self.sample_rate = 44100
        self.residual_model = False
        self.audio_adapter = AudioAdapter.default()
        self.chunk_duration = 30

    def get_estimates(self, input_path: str):
        ctx = get_extension_context(self.context)
        nn.set_default_context(ctx)
        nn.set_auto_forward(True)

        audio, _ = self.audio_adapter.load(input_path,
                                           sample_rate=self.sample_rate)

        if audio.shape[1] > 2:
            warnings.warn('Channel count > 2! '
                          'Only the first two channels will be processed!')
            audio = audio[:, :2]

        if audio.shape[1] == 1:
            print('received mono file, so duplicate channels')
            audio = np.repeat(audio, 2, axis=1)

        # Split and separate sources using moving window protocol for each chunk of audio
        # chunk duration must be lower for machines with low memory
        chunk_size = self.sample_rate * self.chunk_duration
        if (audio.shape[0] % chunk_size) == 0:
            nchunks = (audio.shape[0] // chunk_size)
        else:
            nchunks = (audio.shape[0] // chunk_size) + 1

        print('Separating...')
        estimates = {}
        for chunk_idx in trange(nchunks):
            cur_chunk = audio[chunk_idx *
                              chunk_size:min((chunk_idx + 1) *
                                             chunk_size, audio.shape[0]), :]
            cur_estimates = separate(cur_chunk,
                                     model_path=str(self.model_file_path),
                                     niter=self.iterations,
                                     alpha=self.alpha,
                                     softmask=self.softmask,
                                     residual_model=self.residual_model)
            if any(estimates) is False:
                estimates = cur_estimates
            else:
                for key in cur_estimates:
                    estimates[key] = np.concatenate(
                        (estimates[key], cur_estimates[key]), axis=0)
        return estimates

    def create_static_mix(self, parts, input_path: str, output_path: Path):
        download_and_verify(MODEL_URL, self.model_dir, self.model_file_path)
        estimates = self.get_estimates(input_path)

        final_source = None

        for name, source in estimates.items():
            if not parts[name]:
                continue
            final_source = source if final_source is None else final_source + source

        print('Writing to MP3...')
        self.audio_adapter.save(output_path, final_source, self.sample_rate, 'mp3', self.bitrate)

    def separate_into_parts(self, input_path: str, output_path: Path):
        download_and_verify(MODEL_URL, self.model_dir, self.model_file_path)
        estimates = self.get_estimates(input_path)

        # Export all source MP3s in parallel
        pool = Pool()
        tasks = []
        output_path = Path(output_path)

        for name, estimate in estimates.items():
            filename = f'{name}.mp3'
            print(f'Exporting {name} MP3...')
            task = pool.apply_async(self.audio_adapter.save, (output_path / filename, estimate, self.sample_rate, 'mp3', self.bitrate))
            tasks.append(task)

        pool.close()
        pool.join()


#!/usr/bin/env python3

import rospy
from numpy import maximum, minimum
from sensor_msgs.msg import JointState
from std_msgs.msg import Float64MultiArray
from time import time, sleep

# processes the data from the ROSTopic named "joint_states"
def joint_callback(data):
    pass
    #commented out because it throws an index out of range error which I don't understand
    '''
    print("Msg: {}".format(data.header.seq))
    print("Wheel Positions:\n\tLeft: {0:.2f}rad\n\tRight:{0:.2f}rad\n\n".format(data.position[0],data.position[1]))
    print("Joint Positions:\n\tShoulder1: {0:.2f}rad\n\tShoulder2:{0:.2f}rad\n\tElbow: {0:.2f}rad\n\tWrist:{0:.2f}rad\n\n".format(data.position[2],data.position[3],data.position[4],data.position[5]))
    print("Gripper Position:\n\tGripper:{0:.2f}rad\n".format(data.position[6]))
    print("----------")
    '''
# listens to the "joint_states" topic and sends them to "joint_callback" for processing

def read_joint_states():
    rospy.Subscriber("joint_states",JointState,joint_callback)

# Makes sure the joints do not go outside the joint limits/break the servos
def clean_joint_states(data):
    lower_limits = [0, -1.57, -1.57, -1.57, -1.57, -1]
    upper_limits = [0, 1.57, 1.57, 1.57, 1.57, 1.57]
    clean_lower = maximum(lower_limits,data)
    clean_upper = minimum(clean_lower,upper_limits)
    return list(clean_upper)

# publishes a set of joint commands to the 'joint_trajectory_point' topic
def move_arm():
    jointpub = rospy.Publisher('joint_trajectory_point',Float64MultiArray,queue_size =10)
    joint_pos = Float64MultiArray()
# Joint Position vector should contain 6 elements:
# [0, shoulder1, shoulder2, elbow, wrist, gripper]
    joint_pos.data = clean_joint_states([0, 0, 0, -0.78, 0, -1.1])
    jointpub.publish(joint_pos)
    read_joint_states()

#loops over the commands at 20Hz until shut down
if __name__ == '__main__':
    rospy.init_node('move_arm',anonymous=True)
    rate = rospy.Rate(20)
    while not rospy.is_shutdown():
        move_arm()
        rate.sleep()


#Predict
from deepforest import main
from deepforest.utilities import annotations_to_shapefile
import glob
import geopandas as gpd
import rasterio
from src.main import TreeModel
from src import data
from torch.utils.data import Dataset
import os
import numpy as np
from torchvision import transforms
from torch.nn import functional as F
import torch
from torch.utils.data.dataloader import default_collate

class on_the_fly_Dataset(Dataset):
    """A csv file with a path to image crop and label
    Args:
       crowns: geodataframe of crown locations from a single rasterio src
       image_path: .tif file location
    """
    def __init__(self, crowns, image_path, config=None):
        self.config = config
        self.crowns = crowns
        self.image_size = config["image_size"]
        self.src = rasterio.open(image_path)

    def __len__(self):
        #0th based index
        return self.crowns.shape[0]

    def __getitem__(self, index):
        geom = self.crowns.iloc[index].geometry
        individual = self.crowns.iloc[index].individual
        left, bottom, right, top = geom.bounds
        crop = self.src.read(window=rasterio.windows.from_bounds(left, bottom, right, top, transform=self.src.transform))

        if crop.size == 0:
            return individual, None

        #preprocess and batch
        image = data.preprocess_image(crop, channel_is_first=True)
        image = transforms.functional.resize(image, size=(self.config["image_size"],self.config["image_size"]), interpolation=transforms.InterpolationMode.NEAREST)

        inputs = {}
        inputs["HSI"] = image

        return individual, inputs

def my_collate(batch):
    batch = [x for x in batch if x[1] is not None]
    return default_collate(batch)

def predict_tile(PATH, model_path, config):
    #get rgb from HSI path
    HSI_basename = os.path.basename(PATH)
    if "hyperspectral" in HSI_basename:
        rgb_name = "{}.tif".format(HSI_basename.split("_hyperspectral")[0])
    else:
        rgb_name = HSI_basename
    rgb_pool = glob.glob(config["rgb_sensor_pool"], recursive=True)
    rgb_path = [x for x in rgb_pool if rgb_name in x][0]
    crowns = predict_crowns(rgb_path)
    crowns["tile"] = PATH

    #Load species model
    m = TreeModel.load_from_checkpoint(model_path)
    trees, features = predict_species(HSI_path=PATH, crowns=crowns, m=m, config=config)

    #Spatial smooth
    trees = smooth(trees=trees, features=features, size=config["neighbor_buffer_size"], alpha=config["neighborhood_strength"])
    trees["spatial_taxonID"] = trees["spatial_label"]
    trees["spatial_taxonID"] = trees["spatial_label"].apply(lambda x: m.index_to_label[x])

    return trees

def predict_crowns(PATH):
    m = main.deepforest()
    if torch.cuda.is_available():
        m.config["gpus"] = 1
    m.use_release(check_release=False)
    boxes = m.predict_tile(PATH)
    r = rasterio.open(PATH)
    transform = r.transform
    crs = r.crs
    gdf = annotations_to_shapefile(boxes, transform=transform, crs=crs)

    #Dummy variables for schema
    basename = os.path.splitext(os.path.basename(PATH))[0]
    individual = ["{}_{}".format(x, basename) for x in range(gdf.shape[0])]
    gdf["individual"] = individual
    gdf["plotID"] = None
    gdf["siteID"] = None #TODO
    gdf["box_id"] = None
    gdf["plotID"] = None
    gdf["taxonID"] = None

    return gdf

def predict_species(crowns, HSI_path, m, config):
    ds = on_the_fly_Dataset(crowns, HSI_path, config)
    data_loader = torch.utils.data.DataLoader(
        ds,
        batch_size=config["predict_batch_size"],
        shuffle=False,
        num_workers=config["workers"],
        collate_fn=my_collate
    )
    df, features = m.predict_dataloader(data_loader, train=False, return_features=True)
    crowns["bbox_score"] = crowns["score"]

    #If CHM exists
    try:
        df = df.merge(crowns[["individual","geometry","bbox_score","tile","CHM_height"]], on="individual")
    except:
        df = df.merge(crowns[["individual","geometry","bbox_score","tile"]], on="individual")

    return df, features

def smooth(trees, features, size, alpha):
    """Given the results dataframe and feature labels, spatially smooth based on alpha value"""
    trees = gpd.GeoDataFrame(trees, geometry="geometry")
    sindex = trees.sindex
    tree_buffer = trees.buffer(size)
    smoothed_features = []
    for index, geom in enumerate(tree_buffer):
        intersects = sindex.query(geom)
        focal_feature = features[index,]
        neighbor_features = np.mean(features[intersects,], axis=0)
        smoothed_feature = focal_feature + alpha * neighbor_features
        smoothed_features.append(smoothed_feature)
    smoothed_features = np.vstack(smoothed_features)
    spatial_label = np.argmax(smoothed_features, axis=1)
    spatial_score = np.max(smoothed_features, axis=1)
    trees["spatial_label"] = spatial_label
    trees["spatial_score"] = spatial_score

    return trees


# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.

import torch
import numpy as np

# from ssd_mobilenet_v1 import create_mobilenetv1_ssd
from ssd_r34 import SSD_R34

model = SSD_R34()
model.eval()

image = torch.rand(1,3,1200,1200)

results = model(image)


# python
import os

# site-packages
import numpy as np

# local imports
import utils, tasks

class simulation( object ):

    def __init__(self, dir_name):

        # intialize data structure to store simulation data
        self.data = {
          # from simulation configuration file
          'name': None,
          'friction': None,
          'rough_fault_seed': None,
          'fault_roughness': None,
          'ihypo': None,

          # calculated from simulation output
          'rupture_length'  : None,
          'rupture_area'    : None,
          'avg_stress_drop' : None,
          'avg_slip': None,
          'avg_psv' : None,
          'avg_surf_slip': None,
          'avg_surf_psv': None,
          'max_slip' : None,
          'max_psv' : None,
          'max_surf_slip' : None,
          'max_surf_psv' : None,
          'moment': None,
          'mw': None,
        }

        # populated from parsing the parameter file
        self.params = {}

        # hardcoded
        self.__config_file = 'meta.py'

        # simulation directory
        self.dir = os.path.expanduser( dir_name )

        # read simulation details and update data
        self.data.update( self.parse_simulation_details() )

        # add some parameters to top level class for easy access
        self.dt = self.params['dt']
        self.dx = self.params['dx'][0]
        self.nx = int(self.params['nn'][0])
        self.nz = int(self.params['nn'][1])
        self.ihypo = self.params['ihypo']
        self.name = self.params['name']

        # load slip, psv, and trup
        nx = int(self.params['shape']['su1'][0])
        nz = int(self.params['shape']['su1'][1])
        self.dtype = self.params['dtype']

        su1 = np.fromfile(os.path.join(self.dir, 'out/su1'), dtype=self.dtype).reshape([nz,nx])
        su2 = np.fromfile(os.path.join(self.dir, 'out/su2'), dtype=self.dtype).reshape([nz,nx])
        self.slip = np.sqrt( su1**2 + su2**2 )
        self.trup = np.fromfile(os.path.join(self.dir, 'out/trup'), dtype=self.dtype).reshape([nz,nx])
        self.psv = np.fromfile(os.path.join(self.dir, 'out/psv'), dtype=self.dtype).reshape([nz,nx])

        # load shear stress
        nx = int(self.params['shape']['tsm'][0])
        nz = int(self.params['shape']['tsm'][1])
        nt = int(self.params['shape']['tsm'][2]) # some models didn't write out all timesteps

        self.tsm = np.fromfile( os.path.join(self.dir, 'out/tsm'), dtype=self.dtype )

        # sometimes the simulations finish early
        try:
            self.tsm = self.tsm.reshape([nt,nz,nx])
        except ValueError:
            n = len( self.tsm )
            sdim = int(n / ( nx * nz ))
            self.tsm = self.tsm.reshape([sdim,nz,nx])

    def parse_simulation_details( self ):
        """
        parse parameters from meta.py file. will populate specifics for database.

        inputs
            config_file (string of file object) : configuration file for simulation

        returns
            params (dict) : params.keys() = ['name', 'friction', rough_fault_seed', 'fault_roughness']

        """
        # read meta.py file
        meta = utils.parse( os.path.join( self.dir, self.__config_file ) )

        # update simulation object with all parameters
        self.params.update( meta )

        # extract parameters
        data = {}
        data['name'] = meta['name']
        data['friction'] = meta['friction']
        data['rough_fault_seed'] = utils.extract_from_string( data['name'], ['rf','seed'] )
        data['fault_roughness'] = utils.extract_from_string( data['name'], 'a' )
        data['ihypo'] = str(meta['ihypo'])
        return data

    def get_masked_field(self, field_name, tind = None, mask='psv', tol=0.001, hypo_tol=4000):
        """
        masks a rupture field around hypocenter, velocity strengthening region, and based on
        field values.  currently, masking where psv < 0.001 seems to be the most general approach.

        notes: returns the mask that is the same shape as 'field'

        inputs
            sim : simulation object
            field : string of field in simulation object to use.
            hypo_tol : radius to mask around hypocenter.
            vel_str_tol : velocity strengthening tolerance.
            field_tol : tolerance of simulation field.

        return:
            field (ndarray.masked)

        """

        # get reference to fields
        field = getattr(self, field_name)
        mask_field = getattr(self, mask)

        nskipx_mask = self.params['indices'][mask][0][2]
        if field_name is 'slip':
            nskipx_field = self.params['indices']['su1'][0][2]
        else:
            nskipx_field = self.params['indices'][field_name][0][2]

        # get dims of field
        nt = 0
        try:
            nx,nz,nt = field.shape
        except ValueError:
            nx,nz = field.shape

        # some sanity checks
        if len(mask_field.shape) > len(field.shape):
            raise ValueError("Mask should not have higher dimensionality that field being masked.")

        if mask_field.shape != (self.nz, self.nx):
            raise ValueError("Masking field should cover entire fault plane.  Try trup, or psv, or slip.")

        if nskipx_mask > nskipx_field:
            raise ValueError("Mask has less resolution that field. Unable to mask field.")

        # calculate hypocentral mask
        hx = int(self.ihypo[0])
        hz = int(self.ihypo[1])
        r = int(hypo_tol / self.dx)
        z,x = np.ogrid[-hz:self.nz-hz, -hx:self.nx-hx]
        mask1 = x*x + z*z <= r*r

        # mask unruptured area based on mask and tol
        mask2 = np.zeros( mask_field.shape, dtype=bool )
        mask2[np.where( mask_field < tol )] = True

        # decimate mask if necessary
        ratio = 1
        if nskipx_field != nskipx_mask:
            ratio = int(nskipx_field / nskipx_mask)

        # combine masks using logical or
        mask = np.ma.mask_or(mask1, mask2)[::ratio, ::ratio]

        # return at time
        if tind is not None:
            masked = np.ma.array( np.squeeze(field[tind, :, :]), mask = mask )
        else:
            masked = np.ma.array( field, mask = mask )

        return masked

    def process( self ):
        """ process the simulation from folders in tasks. make sure to update data as we go along. """
        print(f'processing sim: {self.name}')

        self.data['rupture_length'] = tasks.get_fault_extent( self )
        self.data['rupture_area'] = tasks.get_area( self )
        self.data['rupture_area_bbox'] = tasks.get_area_bbox( self )
        self.data['avg_stress_drop'] = tasks.compute_stress_drop( self )

        avg, avg_surf, mx, mx_surf = tasks.rupture_field_stats( self , 'slip' )
        self.data['avg_slip'] = avg
        self.data['avg_surf_slip']= avg_surf
        self.data['max_slip'] = mx
        self.data['max_surf_slip'] = mx_surf

        avg, avg_surf, mx, mx_surf = tasks.rupture_field_stats( self , 'psv' )
        self.data['avg_psv'] = avg
        self.data['avg_surf_psv'] = avg_surf
        self.data['max_psv'] = mx
        self.data['max_surf_psv'] = mx_surf

        self.data['mw'] = tasks.get_mw( self, dtype=self.dtype )
        self.data['moment'] = utils.mw_to_m0( self.data['mw'] )

        return self


import numpy as np
import torch

from .assignment import BBoxAssigner

def random_choice(gallery, num):
    """Random select some elements from the gallery.

    It seems that Pytorch's implementation is slower than numpy so we use numpy
    to randperm the indices.
    """
    assert len(gallery) >= num
    if isinstance(gallery, list):
        gallery = np.array(gallery)
    cands = np.arange(len(gallery))
    np.random.shuffle(cands)
    rand_inds = cands[:num]
    if not isinstance(gallery, np.ndarray):
        rand_inds = torch.from_numpy(rand_inds).long().to(gallery.device)
    return gallery[rand_inds]

def assign_and_sample(bboxes, gt_bboxes, gt_bboxes_ignore, gt_labels, cfg):
    if isinstance(gt_bboxes, np.ndarray):
        gt_bboxes = torch.from_numpy(gt_bboxes).to(bboxes.device)
    if isinstance(gt_bboxes_ignore, np.ndarray):
        gt_bboxes_ignore = torch.from_numpy(gt_bboxes_ignore).to(bboxes.device)
    if isinstance(gt_labels, np.ndarray):
        gt_labels = torch.from_numpy(gt_labels).to(bboxes.device)

    bbox_assigner = BBoxAssigner(**cfg.assigner)
    bbox_sampler = BBoxSampler(**cfg.sampler)

    assign_result = bbox_assigner.assign(bboxes, gt_bboxes, gt_bboxes_ignore,
                                         gt_labels)
    sampling_result = bbox_sampler.sample(assign_result,
                                          bboxes,
                                          gt_bboxes,
                                          gt_labels)
    return assign_result, sampling_result

class BBoxSampler(object):
    """Sample positive and negative bboxes given assigned results.

    Args:
        pos_fraction (float): Positive sample fraction.
        neg_pos_ub (float): Negative/Positive upper bound.
        pos_balance_sampling (bool): Whether to sample positive samples around
            each gt bbox evenly.
        neg_balance_thr (float, optional): IoU threshold for simple/hard
            negative balance sampling.
        neg_hard_fraction (float, optional): Fraction of hard negative samples
            for negative balance sampling.
    """

    def __init__(self,
                 num,
                 pos_fraction,
                 neg_pos_ub=-1,
                 add_gt_as_proposals=True,
                 pos_balance_sampling=False,
                 neg_balance_thr=0,
                 neg_hard_fraction=0.5):
        self.num = num
        self.pos_fraction = pos_fraction
        self.neg_pos_ub = neg_pos_ub
        self.add_gt_as_proposals = add_gt_as_proposals
        self.pos_balance_sampling = pos_balance_sampling
        self.neg_balance_thr = neg_balance_thr
        self.neg_hard_fraction = neg_hard_fraction

    def _sample_pos(self, assign_result, num_expected):
        """Balance sampling for positive bboxes/anchors.

        1. calculate average positive num for each gt: num_per_gt
        2. sample at most num_per_gt positives for each gt
        3. random sampling from rest anchors if not enough fg
        """
        pos_inds = torch.nonzero(assign_result.gt_inds > 0)
        if pos_inds.numel() != 0:
            pos_inds = pos_inds.squeeze(1)
        if pos_inds.numel() <= num_expected:
            return pos_inds
        elif not self.pos_balance_sampling:
            return random_choice(pos_inds, num_expected)
        else:
            unique_gt_inds = torch.unique(
                assign_result.gt_inds[pos_inds].cpu())
            num_gts = len(unique_gt_inds)
            num_per_gt = int(round(num_expected / float(num_gts)) + 1)
            sampled_inds = []
            for i in unique_gt_inds:
                inds = torch.nonzero(assign_result.gt_inds == i.item())
                if inds.numel() != 0:
                    inds = inds.squeeze(1)
                else:
                    continue
                if len(inds) > num_per_gt:
                    inds = random_choice(inds, num_per_gt)
                sampled_inds.append(inds)
            sampled_inds = torch.cat(sampled_inds)
            if len(sampled_inds) < num_expected:
                num_extra = num_expected - len(sampled_inds)
                extra_inds = np.array(
                    list(set(pos_inds.cpu()) - set(sampled_inds.cpu())))
                if len(extra_inds) > num_extra:
                    extra_inds = random_choice(extra_inds, num_extra)
                extra_inds = torch.from_numpy(extra_inds).to(
                    assign_result.gt_inds.device).long()
                sampled_inds = torch.cat([sampled_inds, extra_inds])
            elif len(sampled_inds) > num_expected:
                sampled_inds = random_choice(sampled_inds, num_expected)
            return sampled_inds

    def _sample_neg(self, assign_result, num_expected):
        """Balance sampling for negative bboxes/anchors.

        Negative samples are split into 2 set: hard (balance_thr <= iou <
        neg_iou_thr) and easy (iou < balance_thr). The sampling ratio is
        controlled by `hard_fraction`.
        """
        neg_inds = torch.nonzero(assign_result.gt_inds == 0)
        if neg_inds.numel() != 0:
            neg_inds = neg_inds.squeeze(1)
        if len(neg_inds) <= num_expected:
            return neg_inds
        elif self.neg_balance_thr <= 0:
            # uniform sampling among all negative samples
            return random_choice(neg_inds, num_expected)
        else:
            max_overlaps = assign_result.max_overlaps.cpu().numpy()
            # balance sampling for negative samples
            neg_set = set(neg_inds.cpu().numpy())
            easy_set = set(
                np.where(
                    np.logical_and(max_overlaps >= 0,
                                   max_overlaps < self.neg_balance_thr))[0])
            hard_set = set(np.where(max_overlaps >= self.neg_balance_thr)[0])
            easy_neg_inds = list(easy_set & neg_set)
            hard_neg_inds = list(hard_set & neg_set)

            num_expected_hard = int(num_expected * self.neg_hard_fraction)
            if len(hard_neg_inds) > num_expected_hard:
                sampled_hard_inds = random_choice(hard_neg_inds,
                                                  num_expected_hard)
            else:
                sampled_hard_inds = np.array(hard_neg_inds, dtype=np.int)
            num_expected_easy = num_expected - len(sampled_hard_inds)
            if len(easy_neg_inds) > num_expected_easy:
                sampled_easy_inds = random_choice(easy_neg_inds,
                                                  num_expected_easy)
            else:
                sampled_easy_inds = np.array(easy_neg_inds, dtype=np.int)
            sampled_inds = np.concatenate((sampled_easy_inds,
                                           sampled_hard_inds))
            if len(sampled_inds) < num_expected:
                num_extra = num_expected - len(sampled_inds)
                extra_inds = np.array(list(neg_set - set(sampled_inds)))
                if len(extra_inds) > num_extra:
                    extra_inds = random_choice(extra_inds, num_extra)
                sampled_inds = np.concatenate((sampled_inds, extra_inds))
            sampled_inds = torch.from_numpy(sampled_inds).long().to(
                assign_result.gt_inds.device)
            return sampled_inds

    def sample(self, assign_result, bboxes, gt_bboxes, gt_labels=None):
        """Sample positive and negative bboxes.

        This is a simple implementation of bbox sampling given candidates,
        assigning results and ground truth bboxes.

        1. Assign gt to each bbox.
        2. Add gt bboxes to the sampling pool (optional).
        3. Perform positive and negative sampling.

        Args:
            assign_result (:obj:`AssignResult`): Bbox assigning results.
            bboxes (Tensor): Boxes to be sampled from.
            gt_bboxes (Tensor): Ground truth bboxes.
            gt_labels (Tensor, optional): Class labels of ground truth bboxes.

        Returns:
            :obj:`SamplingResult`: Sampling result.
        """
        bboxes = bboxes[:, :7]

        gt_flags = bboxes.new_zeros((bboxes.shape[0], ), dtype=torch.uint8)
        if self.add_gt_as_proposals:
            bboxes = torch.cat([gt_bboxes, bboxes], dim=0)
            assign_result.add_gt_(gt_labels)
            gt_flags = torch.cat([
                bboxes.new_ones((gt_bboxes.shape[0], ), dtype=torch.uint8),
                gt_flags
            ])

        num_expected_pos = int(self.num * self.pos_fraction)
        pos_inds = self._sample_pos(assign_result, num_expected_pos)
        # We found that sampled indices have duplicated items occasionally.
        # (mab be a bug of PyTorch)
        pos_inds = pos_inds.unique()
        num_sampled_pos = pos_inds.numel()
        num_expected_neg = self.num - num_sampled_pos
        if self.neg_pos_ub >= 0:
            num_neg_max = int(self.neg_pos_ub *
                              num_sampled_pos) if num_sampled_pos > 0 else int(
                                  self.neg_pos_ub)
            num_expected_neg = min(num_neg_max, num_expected_neg)
        neg_inds = self._sample_neg(assign_result, num_expected_neg)
        neg_inds = neg_inds.unique()

        return SamplingResult(pos_inds, neg_inds, bboxes, assign_result.max_overlaps, gt_bboxes,
                              assign_result, gt_flags)

class SamplingResult(object):

    def __init__(self, pos_inds, neg_inds, bboxes, max_overlaps, gt_bboxes, assign_result,
                 gt_flags):
        self.pos_inds = pos_inds
        self.neg_inds = neg_inds
        self.pos_bboxes = bboxes[pos_inds]
        self.neg_bboxes = bboxes[neg_inds]
        self.pos_is_gt = gt_flags[pos_inds]
        self.pos_overs = max_overlaps[pos_inds]
        self.neg_overs = max_overlaps[neg_inds]
        self.num_gts = gt_bboxes.shape[0]
        self.pos_assigned_gt_inds = assign_result.gt_inds[pos_inds] - 1
        self.pos_gt_bboxes = gt_bboxes[self.pos_assigned_gt_inds, :]
        if assign_result.labels is not None:
            self.pos_gt_labels = assign_result.labels[pos_inds]
        else:
            self.pos_gt_labels = None

    @property
    def bboxes(self):
        return torch.cat([self.pos_bboxes, self.neg_bboxes])

    @property
    def overs(self):
        return torch.cat([self.pos_overs, self.neg_overs])