    x0, y0, x1, y1 = p2bb[p]
    return read_cropped_image(expand_path(p), x0, y0, x1, y1, False, img_shape=(BOX_SIZE, BOX_SIZE, 3))

__author__ = 'ZFTurbo: http://example.com'

if __name__ == '__main__':
    import os

    gpu_use = 0
    print('GPU use: {}'.format(gpu_use))
    os.environ["KERAS_BACKEND"] = "tensorflow"
    os.environ["CUDA_VISIBLE_DEVICES"] = "{}".format(gpu_use)

from a00_common_functions import *
from siamese_net_v5_densenet121.a01_seamese_net_models_v5 import read_cropped_image, build_model, preprocess_image

from pandas import read_csv
from os.path import isfile
import pickle
import numpy as np
import os
import sys
from keras import regularizers
from keras.optimizers import Adam
from keras.utils import Sequence
from lap import lapjv
import time
import math
import pandas as pd
from sklearn.model_selection import KFold
import keras
import random
from keras import backend as K
from scipy.ndimage import affine_transform
from multiprocessing.pool import ThreadPool
import operator
from contextlib import contextmanager

@contextmanager
def timer(name):
    t0 = time.time()
    yield
    print('[{}] done in {:.1f} s'.format(name, time.time() - t0))

BOX_SIZE = 384
TOTAL_FOLD = 4
RUN_FOLD = int(sys.argv[1])
print("Using {} fold CV. Current fold: {} Box size: {}".format(TOTAL_FOLD, RUN_FOLD, BOX_SIZE))

THREADS_NUM = 6
print("Threads: {}".format(THREADS_NUM))

LINEAR_ASSIGNMENT_SEGMENT_SIZE = 4
segment = True
if segment:
    print("Using LINEAR_ASSIGNMENT_SEGMENT_SIZE: {}".format(LINEAR_ASSIGNMENT_SEGMENT_SIZE))

def get_image_sizes(join):
    # photo -> image shape
    cache_path = CACHE_PATH + 'p2size_{}.pklz'.format(os.path.basename(__file__))
    if not os.path.isfile(cache_path):
        p2size = {}
        for p in join:
            size = get_shape(expand_path(p))
            p2size[p] = size
        print(len(p2size), list(p2size.items())[:5])
        save_in_file(p2size, cache_path)
    else:
        p2size = load_from_file(cache_path)
        if len(p2size) != len(join):
            print('Len p2size diff!')
            exit()
    return p2size

def get_boxes():
    ### using zft bbox
    temp_p2bb = pickle.load(open(OUTPUT_PATH + 'p2bb_averaged_v1.pkl', 'rb'))
    p2bb1 = {}
    for k in temp_p2bb:
        p2bb1[k + '.jpg'] = temp_p2bb[k]
    temp_p2bb = pickle.load(open(OUTPUT_PATH + 'p2bb_averaged_playground_v1.pkl', 'rb'))
    p2bb2 = {}
    for k in temp_p2bb:
        p2bb2[k + '.jpg'] = temp_p2bb[k]
    p2bb = {**p2bb1, **p2bb2}
    return p2bb

def get_tagged_data():
    # load tagged data
    tagged = dict([(p[:-4] + '.jpg', w) for _, p, w in read_csv(INPUT_PATH + 'train.csv').to_records()])

    if 0:
        newdata = dict(
            [(k[:-4] + '.jpg', v) for k, v in load_from_file_fast(OUTPUT_PATH + 'new_train_data.pkl').items()])
        newdata_bootstrap_Dec17 = dict([(k[:-4] + '.jpg', v) for k, v in
                                        load_from_file_fast(OUTPUT_PATH + 'pseudo_label_test_data_LB953_model.pkl').items()])
        tagged = dict(**tagged, **newdata)
        for ndk in newdata_bootstrap_Dec17:
            if ndk not in tagged:
                tagged[ndk] = newdata_bootstrap_Dec17[ndk]
            elif tagged[ndk] == newdata_bootstrap_Dec17[ndk]:
                if 0:
                    print("Ignored duplicate keys with same values: {} -> {}".format(ndk, tagged[ndk]))
            else:
                print("Error. Found duplicate keys with different values: {} -> {}, {}".format(ndk, tagged[ndk],
                                                                                newdata_bootstrap_Dec17[
                                                                                ndk]))
                assert 0
    print("Total tagged images: ", len(tagged))
    return tagged

# whale ID -> photo
def get_whale_lists_for_class_dict(tagged):
    w2p = {}
    for img, lb in tagged.items():
        if img in p2bb:  # remove bad images!
            if lb not in w2p:
                w2p[lb] = [img]
            else:
                w2p[lb].append(img)
    return w2p

def print_stats(w2p):
    dfs = pd.DataFrame({'Id': list(w2p.keys())})
    dfs['counts'] = 0
    for i in range(dfs.shape[0]):
        dfs.iat[i, 1] = len(w2p[dfs.iat[i, 0]])
    print("\nStats: ")
    print('Num in class, Count')
    print(dfs['counts'].value_counts())

p2bb = get_boxes()
tagged = get_tagged_data()
w2p = get_whale_lists_for_class_dict(tagged)

# test
submit = [p for _, p, _ in read_csv(INPUT_PATH + 'sample_submission.csv').to_records()]
submit = [f[:-4] + '.jpg' for f in submit]
join = set(list(tagged.keys()) + submit)

# photo -> image shape
p2size = get_image_sizes(join)

print('Tagged items: {}'.format(len(tagged)))
print('Test items: {}'.format(len(submit)))
print('Overall items: {}'.format(len(join)))
print('Example: {} {}'.format(list(tagged.items())[:5], submit[:5]))

def get_kfold_split_weimin(fold_number):
    fold_path = OUTPUT_PATH + 'kfold/new_4_folds_split_train_val_v1.pkl'
    print('Load fold split from: {}'.format(fold_path))
    data = load_from_file_fast(fold_path)
    data = data['fold_{}'.format(fold_number)]
    train_part = []
    for el in sorted(data['train']):
        train_part.append(el[:-4] + '.jpg')
    valid_part = []
    for el in sorted(data['val']):
        valid_part.append(el[:-4] + '.jpg')
    return train_part, valid_part

def gen_train_labels_to_fotos(w2p, train):
    w2ts = {}  # valid whale label to photos
    for k in w2p:
        if len(w2p[k]) > 1 and k != 'new_whale':
            w2ts[k] = []
            for el in w2p[k]:
                if el in train:
                    w2ts[k].append(el)
            w2ts[k] = np.array(w2ts[k])
    return w2ts

train, val = get_kfold_split_weimin(RUN_FOLD)
w2ts = gen_train_labels_to_fotos(w2p, train)
print('Train size: {}'.format(len(train)))
print('Valid size: {}'.format(len(val)))
print('W2TS size: {}'.format(len(w2ts)))

def read_for_training(p):
    """
    Read and preprocess an image with data augmentation (random transform).
    """
    x0, y0, x1, y1 = p2bb[p]
    return read_cropped_image(expand_path(p), x0, y0, x1, y1, True, img_shape=(BOX_SIZE, BOX_SIZE, 3))

def read_for_validation(p):
    """
    Read and preprocess an image without data augmentation (use for testing).
    """

def prepare_val_res(score, threshold):
    """
    Generate a Kaggle submission file.
    @param threshold the score given to 'new_whale'
    @param filename the submission file name
    """
    new_whale = 'new_whale'
    res = {}
    train_arr = np.array(train)

    for i, p in enumerate(val):
        t = []
        s = set()
        a = score[i, :]

        top_label_probs = {}
        cond = a > threshold
        cond_index = np.where(cond)[0]
        cond_images = train_arr[cond_index]
        for j, img in enumerate(cond_images):
            if tagged[img] in top_label_probs:
                top_label_probs[tagged[img]] += a[cond_index[j]]
            else:
                top_label_probs[tagged[img]] = a[cond_index[j]]

        sorted_top_label_probs = sort_dict_by_values(top_label_probs)

        t = []
        for lb, _ in sorted_top_label_probs:
            t.append(lb)

        if len(t) < 5:
            t.append(new_whale)

        for index in np.argsort(a)[::-1]:
            if tagged[train_arr[index]] not in t:
                t.append(tagged[train_arr[index]])
            if len(t) >= 5:
                break

        assert len(t) >= 5

        res[p[:-4] + '.jpg'] = t[:5]

    return res

class TrainingDataFast(Sequence):
    def __init__(self, score_list, steps=1000, batch_size=32):
        """
        @param score the cost matrix for the picture matching
        @param steps the number of epoch we are planning with this score matrix
        """
        super(TrainingDataFast, self).__init__()
        #self.score      = -score # Maximizing the self.score is the same as minimuzing -score.
        self.score = np.zeros((len(train), len(train)))
        start = 0
        for s in score_list:
            the_size = s.shape[0]
            self.score[start:start+the_size, start:start+the_size] = -s
            start += the_size

        self.steps      = steps
        self.batch_size = batch_size
        for ts in w2ts.values():
            idxs = [t2i[t] for t in ts]
            for i in idxs:
                for j in idxs:
                    self.score[i,j] = 10000.0 # Set a large value for matching whales -- eliminates this potential pairing

        self.on_epoch_end()
        self.p = ThreadPool(THREADS_NUM)

    def __getitem__(self, index):
        start = self.batch_size*index
        end   = min(start + self.batch_size, len(self.match) + len(self.unmatch))
        size  = end - start
        assert size > 0

        def _read_for_training_expand_dim(p):
            return np.expand_dims(read_for_training(p), 0).astype(K.floatx())

        i_start, i_end = start//2, int(start//2+np.ceil(size/2))
        unroll_match = [[m[0], m[1], u[0], u[1]] for m, u in zip(self.match[i_start:i_end], self.unmatch[i_start:i_end])]
        unroll_match = [i for item in unroll_match for i in item]
        unroll_train_images = np.concatenate(self.p.map(_read_for_training_expand_dim, unroll_match), 0)
        unroll_train_images = preprocess_image(unroll_train_images)

        a_indices = list(range(0, len(unroll_train_images), 2))
        a = unroll_train_images[a_indices]
        b_indices = list(range(1, len(unroll_train_images), 2))
        b = unroll_train_images[b_indices]
        c = np.zeros((size,1), dtype=K.floatx())
        for i in range(0, size, 2):
            c[i,   0] = 1 # This is a match
            c[i+1, 0] = 0 # Different whales

        return [a, b], c

    def on_epoch_end(self):
        if self.steps <= 0: return # Skip this on the last epoch.
        self.steps     -= 1
        self.match      = [] # len(self.match) == len(train), each element is a pair of images that match the same label
        self.unmatch    = []
        if segment:
            # Using slow scipy. Make small batches.
            # Because algorithm is O(n^3), small batches are much faster.
            # However, this does not find the real optimum, just an approximation.
            tmp   = []
            batch = math.ceil(self.score.shape[0] / LINEAR_ASSIGNMENT_SEGMENT_SIZE)
            print("\nScore matrix size: {} and lapjv batch size: {}".format(self.score.shape[0], batch))
            for start in range(0, self.score.shape[0], batch):
                end = min(self.score.shape[0], start + batch)
                with timer("lapjv ops 1 round"):
                    _,_,x = lapjv(self.score[start:end, start:end])
                tmp.append(x + start)
            x = np.concatenate(tmp)
        else:
            print("Running lapjv on entire score matrix! size: ".format)
            with timer("lapjv ops (entire data) "):
                _,_,x = lapjv(self.score) # Solve the linear assignment problem

        y = np.arange(len(x),dtype=np.int32)

        # Compute a derangement for matching whales
        for ts in w2ts.values():
            d = ts.copy()
            while True:
                random.shuffle(d)
                if not np.any(ts == d): break
            for ab in zip(ts,d): self.match.append(ab)

        # Construct unmatched whale pairs from the LAP solution.
        for i,j in zip(x,y):
            if i == j:
                print(self.score)
                print(x)
                print(y)
                print(i,j)
            assert i != j
            self.unmatch.append((train[i],train[j]))

        # Force a different choice for an eventual next epoch.
        self.score[x,y] = 10000.0
        self.score[y,x] = 10000.0
        random.shuffle(self.match)
        random.shuffle(self.unmatch)
        assert len(self.match) == len(train) and len(self.unmatch) == len(train)

    def __len__(self):
        return (len(self.match) + len(self.unmatch) + self.batch_size - 1)//self.batch_size

# A Keras generator to evaluate only the BRANCH MODEL
class FeatureGen(Sequence):
    def __init__(self, data, batch_size=64, verbose=1):
        super(FeatureGen, self).__init__()
        self.data       = data
        self.batch_size = batch_size
        self.verbose    = verbose
        self.p = ThreadPool(THREADS_NUM)

    def __getitem__(self, index):
        start = self.batch_size*index
        size  = min(len(self.data) - start, self.batch_size)
        def _read_for_validation_expand_dim(p):
            return np.expand_dims(read_for_validation(p), 0).astype(K.floatx())
        res = self.p.map(_read_for_validation_expand_dim, self.data[start:start+size])
        res = np.concatenate(res, 0)
        res = preprocess_image(res)

        return res

    def __len__(self):
        return (len(self.data) + self.batch_size - 1)//self.batch_size

# A Keras generator to evaluate on the HEAD MODEL on features already pre-computed.
# It computes only the upper triangular matrix of the cost matrix if y is None.
class ScoreGen(Sequence):
    def __init__(self, x, y=None, batch_size=2048, verbose=1):
        super(ScoreGen, self).__init__()
        self.x          = x
        self.y          = y
        self.batch_size = batch_size
        self.verbose    = verbose
        if y is None:
            self.y           = self.x
            self.ix, self.iy = np.triu_indices(x.shape[0],1)
        else:
            self.iy, self.ix = np.indices((y.shape[0],x.shape[0]))
            self.ix          = self.ix.reshape((self.ix.size,))
            self.iy          = self.iy.reshape((self.iy.size,))
        self.subbatch = (len(self.x) + self.batch_size - 1)//self.batch_size

    def __getitem__(self, index):
        start = index*self.batch_size
        end   = min(start + self.batch_size, len(self.ix))
        a     = self.y[self.iy[start:end],:]
        b     = self.x[self.ix[start:end],:]
        return [a,b]

    def __len__(self):
        return (len(self.ix) + self.batch_size - 1)//self.batch_size

def set_lr(model, lr):
    K.set_value(model.optimizer.lr, float(lr))

def get_lr(model):
    return K.get_value(model.optimizer.lr)

def score_reshape(score, x, y=None):
    """
    Tranformed the packed matrix 'score' into a square matrix.
    @param score the packed matrix
    @param x the first image feature tensor
    @param y the second image feature tensor if different from x
    @result the square matrix
    """
    if y is None:
        # When y is None, score is a packed upper triangular matrix.
        # Unpack, and transpose to form the symmetrical lower triangular matrix.
        m = np.zeros((x.shape[0],x.shape[0]), dtype=K.floatx())
        m[np.triu_indices(x.shape[0],1)] = score.squeeze()
        m += m.transpose()
    else:
        m        = np.zeros((y.shape[0],x.shape[0]), dtype=K.floatx())
        iy,ix    = np.indices((y.shape[0],x.shape[0]))
        ix       = ix.reshape((ix.size,))
        iy       = iy.reshape((iy.size,))
        m[iy,ix] = score.squeeze()
    return m

def compute_score_fast(verbose=1):
    """
    Compute the score matrix by scoring every pictures from the training set against every other picture O(n^2).
    """
    res = []
    batch = int(math.ceil(len(train) / LINEAR_ASSIGNMENT_SEGMENT_SIZE))
    for start in range(0, len(train), batch):
        end = min(len(train), start + batch)
        train_batch = train[start:end]

        features = branch_model.predict_generator(FeatureGen(train_batch, verbose=verbose), max_queue_size=12, verbose=0, use_multiprocessing=False)
        score    = head_model.predict_generator(ScoreGen(features, verbose=verbose), max_queue_size=12, verbose=0, use_multiprocessing=False)
        score    = score_reshape(score, features)
        res.append(score)

    return res

def make_steps_fast(step, ampl, verbose=2, batch_size=6):
    """
    Perform training epochs
    @param step Number of epochs to perform`
    @param ampl the K, the randomized component of the score matrix.
    """
    global w2ts, t2i, steps, features, score, histories

    # shuffle the training pictures
    random.shuffle(train)

    # Map training picture hash value to index in 'train' array
    t2i  = {}
    for i, t in enumerate(train): t2i[t] = i

    # Compute the match score for each picture pair of train data - square matrix score[i,j] is the score between image i and j, 0~1
    with timer("Compute score fast. "):
        score_list = compute_score_fast()

    for i in range(len(score_list)):
        score_list[i] += ampl*np.random.random_sample(size=score_list[i].shape)

    # Train the model for 'step' epochs
    history = model.fit_generator(
        TrainingDataFast(score_list, steps=step, batch_size=batch_size),
        initial_epoch=steps, epochs=steps + step, max_queue_size=12, workers=6, verbose=verbose,
        ).history
    steps += step

    # Collect history data
    history['epochs'] = steps
    history['ms'] = np.mean(score)
    history['lr'] = get_lr(model)
    print(history['epochs'], history['lr'], history['ms'])
    histories.append(history)

def run_validation(threshold):
    print("Running validation ...")
    print('Predict train with branch model...')
    ftrain = branch_model.predict_generator(FeatureGen(train), max_queue_size=20, verbose=0, use_multiprocessing=False)
    print('Predict validation with branch model...')
    fval = branch_model.predict_generator(FeatureGen(val), max_queue_size=20, verbose=0, use_multiprocessing=False)
    print('Score table calc...')
    score = head_model.predict_generator(ScoreGen(ftrain, fval), max_queue_size=20, verbose=0, use_multiprocessing=False)
    score = score_reshape(score, ftrain, fval)
    res = prepare_val_res(score, threshold)

    ## calculate score
    scores = []
    missing = []
    for img in res.keys():
        if img in tagged:
            score = apk([tagged[img]], res[img], k=5)
            scores.append(score)
        else:
            print("Warning! val image: {} not in tagged. ".format(img))
            missing.append(img)

    score = np.mean(np.array(scores))
    print("\nValid score: {:.6f}".format(score))
    return score

def get_model_to_finetune(saved_folder, fold):
    model, branch_model, head_model = build_model(64e-5, 0, img_shape=(BOX_SIZE, BOX_SIZE, 3))

    # Get best model from previous run
    best_models = glob.glob(MODELS_PATH + '{}/densenet121_384px_*.model'.format(saved_folder))
    best_score = -1.0
    best_weights = ''
    for m in best_models:
        score = float(m[:-6].split('_')[-1])
        if score > best_score:
            best_score = score
            best_weights = m

    print('Best score on previous run: {}. Use weights: {}'.format(best_score, best_weights))
    weights = best_weights
    model.load_weights(weights)
    return model, branch_model, head_model

if __name__ == '__main__':
    histories = []
    thr = 0.99
    score = 0
    steps = 0
    saved_folder = 'Res_v5_densenet121'
    if not os.path.isdir(MODELS_PATH + '{}'.format(saved_folder)):
        os.makedirs(MODELS_PATH + '{}'.format(saved_folder))

    version = 'v5_384px_finetune_' + str(RUN_FOLD)
    l2_rate = 0.0002
    model_str = 'ft'
    model, branch_model, head_model = get_model_to_finetune(saved_folder, RUN_FOLD)

    run_validation(thr)
    score = np.random.random_sample(size=(len(train), len(train)))

    # 15 epochs - freeze layers
    branch_model.trainable = False
    set_lr(model, 8e-5)
    optim = Adam(lr=8e-5)
    model.compile(optim, loss='binary_crossentropy', metrics=['binary_crossentropy', 'acc'])

    print("\n{} training block 1 {}".format('*'*17, '*'*17))
    make_steps_fast(15, 1, batch_size=16)
    val_score = run_validation(thr)
    model.save_weights(MODELS_PATH + '{}/{}_{}_temp0_{:.6f}.model'.format(saved_folder, model_str, version, val_score))

    # 4 x 5 epochs - unfreeze layers
    branch_model.trainable = True
    set_lr(model, 8e-5)
    optim  = Adam(lr=8e-5)
    model.compile(optim, loss='binary_crossentropy', metrics=['binary_crossentropy', 'acc'])

    print("\n{} training block 2 {}".format('*'*17, '*'*17))
    set_lr(model, 8e-5)
    for _ in range(4):
        make_steps_fast(5, 0.5)
    val_score = run_validation(thr)
    model.save_weights(MODELS_PATH + '{}/{}_{}_temp1_{:.6f}.model'.format(saved_folder, model_str, version, val_score))

    print("\n{} training block 3 {}".format('*'*17, '*'*17))
    set_lr(model, 4e-5)
    for _ in range(4):
        make_steps_fast(5, 0.25)
    val_score = run_validation(thr)
    model.save_weights(MODELS_PATH + '{}/{}_{}_temp2_{:.6f}.model'.format(saved_folder, model_str, version, val_score))

    print("\n{} training block 4 {}".format('*'*17, '*'*17))
    print("\n\nStarted new round: lr: 2e-5, rounds: 5, epochs: 5, ampl: 0.25\n")
    set_lr(model, 2e-5)
    for _ in range(4):
        make_steps_fast(5, 0.25)
    val_score = run_validation(thr)
    model_path = MODELS_PATH + '{}/{}_{}_temp3_{:.6f}.model'.format(saved_folder, model_str, version, val_score)
    model.save_weights(model_path)
    print("Saved model: ", model_path)

    # final fine tuning - 4 x 5 epochs
    print("\n{} training block 5 {}".format('*'*17, '*'*17))
    print("\nStarted new round: lr: 1e-5, rounds: 2, epochs: 5, ampl: 0.25\n")
    set_lr(model, 1e-5)
    for _ in range(4):
        make_steps_fast(5, 0.25)
    val_score = run_validation(thr)
    model_path = MODELS_PATH + '{}/{}_{}_final_v1_{:.6f}.model'.format(saved_folder, model_str, version, val_score)
    model.save_weights(model_path)
    print("Saved model:", model_path)

    # final fine tuning - 30 x 5 epochs (can stop at any time)
    print("\n{} training block 6 {}".format('*' * 17, '*' * 17))
    print("\nStarted new round: lr: 1e-5, rounds: 2, epochs: 5, ampl: 0.25\n")
    iters_to_run = 20
    set_lr(model, 8e-6)
    for iter1 in range(iters_to_run):
        ampl = 0.25 * ((iters_to_run - iter1) / iters_to_run)
        print('Iter: {} ampl: {}'.format(iters_to_run, ampl))
        make_steps_fast(5, ampl)
        val_score = run_validation(thr)
        model_path = MODELS_PATH + '{}/{}_{}_final_v2_{:.6f}.model'.format(saved_folder, model_str, version, val_score)
        model.save_weights(model_path)
        print("Saved model:", model_path)

    print("all done. ")

'''
Fold 0:
TB 0: Valid score: 0.958814
TB 1: Valid score: 0.910119
TV 2: Valid score: 0.939753
TB 3: Valid score: 0.963057
TB 4: Valid score: 0.963430
TB 5: Valid score: 0.966701
Valid score: 0.967093

Fold 1:
TB 0: Valid score: 0.955803
TB 1: Valid score: 0.889472
TB 2: Valid score: 0.926793
TB 3: Valid score: 0.957300
TB 4: Valid score: 0.963647
TB 5: Valid score: 0.966626
Valid score: 0.967650

Fold 2:
TB 0: Valid score: 0.956908
TB 1: Valid score: 0.926379
TB 2: Valid score: 0.934138
TB 3: Valid score: 0.950624
TB 4: Valid score: 0.955206
TB 5: Valid score: 0.960862
Valid score: 0.961138
Valid score: 0.961669

Fold 3:
TB 0: Valid score: 0.964243
TB 1: Valid score: 0.915626
TB 2: Valid score: 0.947841
TB 3: Valid score: 0.965813
TB 4: Valid score: 0.963862
TB 5: Valid score: 0.968446
Valid score: 0.973923
Valid score: 0.974151
'''


while True:
                random.shuffle(d)
                if not np.any(ts == d): break

__author__ = 'ZFTurbo: http://example.com'

if __name__ == '__main__':
    import os

    gpu_use = 0
    print('GPU use: {}'.format(gpu_use))
    os.environ["KERAS_BACKEND"] = "tensorflow"
    os.environ["CUDA_VISIBLE_DEVICES"] = "{}".format(gpu_use)

from a00_common_functions import *
from siamese_net_v5_densenet121.a01_seamese_net_models_v5 import read_cropped_image, build_model, preprocess_image

from pandas import read_csv
from os.path import isfile
import pickle
import numpy as np
import os
import sys
from keras import regularizers
from keras.optimizers import Adam
from keras.utils import Sequence
from lap import lapjv
import time
import math
import pandas as pd
from sklearn.model_selection import KFold
import keras
import random
from keras import backend as K
from scipy.ndimage import affine_transform
from multiprocessing.pool import ThreadPool
import operator
from contextlib import contextmanager

@contextmanager
def timer(name):
    t0 = time.time()
    yield
    print('[{}] done in {:.1f} s'.format(name, time.time() - t0))

BOX_SIZE = 384
TOTAL_FOLD = 4
RUN_FOLD = int(sys.argv[1])
print("Using {} fold CV. Current fold: {} Box size: {}".format(TOTAL_FOLD, RUN_FOLD, BOX_SIZE))

THREADS_NUM = 6
print("Threads: {}".format(THREADS_NUM))

LINEAR_ASSIGNMENT_SEGMENT_SIZE = 4
segment = True
if segment:
    print("Using LINEAR_ASSIGNMENT_SEGMENT_SIZE: {}".format(LINEAR_ASSIGNMENT_SEGMENT_SIZE))

def get_image_sizes(join):
    # photo -> image shape
    cache_path = CACHE_PATH + 'p2size_{}.pklz'.format(os.path.basename(__file__))
    if not os.path.isfile(cache_path):
        p2size = {}
        for p in join:
            size = get_shape(expand_path(p))
            p2size[p] = size
        print(len(p2size), list(p2size.items())[:5])
        save_in_file(p2size, cache_path)
    else:
        p2size = load_from_file(cache_path)
        if len(p2size) != len(join):
            print('Len p2size diff!')
            exit()
    return p2size

def get_boxes():
    ### using zft bbox
    temp_p2bb = pickle.load(open(OUTPUT_PATH + 'p2bb_averaged_v1.pkl', 'rb'))
    p2bb1 = {}
    for k in temp_p2bb:
        p2bb1[k + '.jpg'] = temp_p2bb[k]
    temp_p2bb = pickle.load(open(OUTPUT_PATH + 'p2bb_averaged_playground_v1.pkl', 'rb'))
    p2bb2 = {}
    for k in temp_p2bb:
        p2bb2[k + '.jpg'] = temp_p2bb[k]
    p2bb = {**p2bb1, **p2bb2}
    return p2bb

def get_tagged_data():
    # load tagged data
    tagged = dict([(p[:-4] + '.jpg', w) for _, p, w in read_csv(INPUT_PATH + 'train.csv').to_records()])

    if 0:
        newdata = dict(
            [(k[:-4] + '.jpg', v) for k, v in load_from_file_fast(OUTPUT_PATH + 'new_train_data.pkl').items()])
        newdata_bootstrap_Dec17 = dict([(k[:-4] + '.jpg', v) for k, v in
                                        load_from_file_fast(OUTPUT_PATH + 'pseudo_label_test_data_LB953_model.pkl').items()])
        tagged = dict(**tagged, **newdata)
        for ndk in newdata_bootstrap_Dec17:
            if ndk not in tagged:
                tagged[ndk] = newdata_bootstrap_Dec17[ndk]
            elif tagged[ndk] == newdata_bootstrap_Dec17[ndk]:
                if 0:
                    print("Ignored duplicate keys with same values: {} -> {}".format(ndk, tagged[ndk]))
            else:
                print("Error. Found duplicate keys with different values: {} -> {}, {}".format(ndk, tagged[ndk],
                                                                                newdata_bootstrap_Dec17[
                                                                                ndk]))
                assert 0
    print("Total tagged images: ", len(tagged))
    return tagged

# whale ID -> photo
def get_whale_lists_for_class_dict(tagged):
    w2p = {}
    for img, lb in tagged.items():
        if img in p2bb:  # remove bad images!
            if lb not in w2p:
                w2p[lb] = [img]
            else:
                w2p[lb].append(img)
    return w2p

def print_stats(w2p):
    dfs = pd.DataFrame({'Id': list(w2p.keys())})
    dfs['counts'] = 0
    for i in range(dfs.shape[0]):
        dfs.iat[i, 1] = len(w2p[dfs.iat[i, 0]])
    print("\nStats: ")
    print('Num in class, Count')
    print(dfs['counts'].value_counts())

p2bb = get_boxes()
tagged = get_tagged_data()
w2p = get_whale_lists_for_class_dict(tagged)

# test
submit = [p for _, p, _ in read_csv(INPUT_PATH + 'sample_submission.csv').to_records()]
submit = [f[:-4] + '.jpg' for f in submit]
join = set(list(tagged.keys()) + submit)

# photo -> image shape
p2size = get_image_sizes(join)

print('Tagged items: {}'.format(len(tagged)))
print('Test items: {}'.format(len(submit)))
print('Overall items: {}'.format(len(join)))
print('Example: {} {}'.format(list(tagged.items())[:5], submit[:5]))

def get_kfold_split_weimin(fold_number):
    fold_path = OUTPUT_PATH + 'kfold/new_4_folds_split_train_val_v1.pkl'
    print('Load fold split from: {}'.format(fold_path))
    data = load_from_file_fast(fold_path)
    data = data['fold_{}'.format(fold_number)]
    train_part = []
    for el in sorted(data['train']):
        train_part.append(el[:-4] + '.jpg')
    valid_part = []
    for el in sorted(data['val']):
        valid_part.append(el[:-4] + '.jpg')
    return train_part, valid_part

def gen_train_labels_to_fotos(w2p, train):
    w2ts = {}  # valid whale label to photos
    for k in w2p:
        if len(w2p[k]) > 1 and k != 'new_whale':
            w2ts[k] = []
            for el in w2p[k]:
                if el in train:
                    w2ts[k].append(el)
            w2ts[k] = np.array(w2ts[k])
    return w2ts

train, val = get_kfold_split_weimin(RUN_FOLD)
w2ts = gen_train_labels_to_fotos(w2p, train)
print('Train size: {}'.format(len(train)))
print('Valid size: {}'.format(len(val)))
print('W2TS size: {}'.format(len(w2ts)))

def read_for_training(p):
    """
    Read and preprocess an image with data augmentation (random transform).
    """
    x0, y0, x1, y1 = p2bb[p]
    return read_cropped_image(expand_path(p), x0, y0, x1, y1, True, img_shape=(BOX_SIZE, BOX_SIZE, 3))

def read_for_validation(p):
    """
    Read and preprocess an image without data augmentation (use for testing).
    """
    x0, y0, x1, y1 = p2bb[p]
    return read_cropped_image(expand_path(p), x0, y0, x1, y1, False, img_shape=(BOX_SIZE, BOX_SIZE, 3))

def prepare_val_res(score, threshold):
    """
    Generate a Kaggle submission file.
    @param threshold the score given to 'new_whale'
    @param filename the submission file name
    """
    new_whale = 'new_whale'
    res = {}
    train_arr = np.array(train)

    for i, p in enumerate(val):
        t = []
        s = set()
        a = score[i, :]

        top_label_probs = {}
        cond = a > threshold
        cond_index = np.where(cond)[0]
        cond_images = train_arr[cond_index]
        for j, img in enumerate(cond_images):
            if tagged[img] in top_label_probs:
                top_label_probs[tagged[img]] += a[cond_index[j]]
            else:
                top_label_probs[tagged[img]] = a[cond_index[j]]

        sorted_top_label_probs = sort_dict_by_values(top_label_probs)

        t = []
        for lb, _ in sorted_top_label_probs:
            t.append(lb)

        if len(t) < 5:
            t.append(new_whale)

        for index in np.argsort(a)[::-1]:
            if tagged[train_arr[index]] not in t:
                t.append(tagged[train_arr[index]])
            if len(t) >= 5:
                break

        assert len(t) >= 5

        res[p[:-4] + '.jpg'] = t[:5]

    return res

class TrainingDataFast(Sequence):
    def __init__(self, score_list, steps=1000, batch_size=32):
        """
        @param score the cost matrix for the picture matching
        @param steps the number of epoch we are planning with this score matrix
        """
        super(TrainingDataFast, self).__init__()
        #self.score      = -score # Maximizing the self.score is the same as minimuzing -score.
        self.score = np.zeros((len(train), len(train)))
        start = 0
        for s in score_list:
            the_size = s.shape[0]
            self.score[start:start+the_size, start:start+the_size] = -s
            start += the_size

        self.steps      = steps
        self.batch_size = batch_size
        for ts in w2ts.values():
            idxs = [t2i[t] for t in ts]
            for i in idxs:
                for j in idxs:
                    self.score[i,j] = 10000.0 # Set a large value for matching whales -- eliminates this potential pairing

        self.on_epoch_end()
        self.p = ThreadPool(THREADS_NUM)

    def __getitem__(self, index):
        start = self.batch_size*index
        end   = min(start + self.batch_size, len(self.match) + len(self.unmatch))
        size  = end - start
        assert size > 0

        def _read_for_training_expand_dim(p):
            return np.expand_dims(read_for_training(p), 0).astype(K.floatx())

        i_start, i_end = start//2, int(start//2+np.ceil(size/2))
        unroll_match = [[m[0], m[1], u[0], u[1]] for m, u in zip(self.match[i_start:i_end], self.unmatch[i_start:i_end])]
        unroll_match = [i for item in unroll_match for i in item]
        unroll_train_images = np.concatenate(self.p.map(_read_for_training_expand_dim, unroll_match), 0)
        unroll_train_images = preprocess_image(unroll_train_images)

        a_indices = list(range(0, len(unroll_train_images), 2))
        a = unroll_train_images[a_indices]
        b_indices = list(range(1, len(unroll_train_images), 2))
        b = unroll_train_images[b_indices]
        c = np.zeros((size,1), dtype=K.floatx())
        for i in range(0, size, 2):
            c[i,   0] = 1 # This is a match
            c[i+1, 0] = 0 # Different whales

        return [a, b], c

    def on_epoch_end(self):
        if self.steps <= 0: return # Skip this on the last epoch.
        self.steps     -= 1
        self.match      = [] # len(self.match) == len(train), each element is a pair of images that match the same label
        self.unmatch    = []
        if segment:
            # Using slow scipy. Make small batches.
            # Because algorithm is O(n^3), small batches are much faster.
            # However, this does not find the real optimum, just an approximation.
            tmp   = []
            batch = math.ceil(self.score.shape[0] / LINEAR_ASSIGNMENT_SEGMENT_SIZE)
            print("\nScore matrix size: {} and lapjv batch size: {}".format(self.score.shape[0], batch))
            for start in range(0, self.score.shape[0], batch):
                end = min(self.score.shape[0], start + batch)
                with timer("lapjv ops 1 round"):
                    _,_,x = lapjv(self.score[start:end, start:end])
                tmp.append(x + start)
            x = np.concatenate(tmp)
        else:
            print("Running lapjv on entire score matrix! size: ".format)
            with timer("lapjv ops (entire data) "):
                _,_,x = lapjv(self.score) # Solve the linear assignment problem

        y = np.arange(len(x),dtype=np.int32)

        # Compute a derangement for matching whales
        for ts in w2ts.values():
            d = ts.copy()

            for ab in zip(ts,d): self.match.append(ab)

        # Construct unmatched whale pairs from the LAP solution.
        for i,j in zip(x,y):
            if i == j:
                print(self.score)
                print(x)
                print(y)
                print(i,j)
            assert i != j
            self.unmatch.append((train[i],train[j]))

        # Force a different choice for an eventual next epoch.
        self.score[x,y] = 10000.0
        self.score[y,x] = 10000.0
        random.shuffle(self.match)
        random.shuffle(self.unmatch)
        assert len(self.match) == len(train) and len(self.unmatch) == len(train)

    def __len__(self):
        return (len(self.match) + len(self.unmatch) + self.batch_size - 1)//self.batch_size

# A Keras generator to evaluate only the BRANCH MODEL
class FeatureGen(Sequence):
    def __init__(self, data, batch_size=64, verbose=1):
        super(FeatureGen, self).__init__()
        self.data       = data
        self.batch_size = batch_size
        self.verbose    = verbose
        self.p = ThreadPool(THREADS_NUM)

    def __getitem__(self, index):
        start = self.batch_size*index
        size  = min(len(self.data) - start, self.batch_size)
        def _read_for_validation_expand_dim(p):
            return np.expand_dims(read_for_validation(p), 0).astype(K.floatx())
        res = self.p.map(_read_for_validation_expand_dim, self.data[start:start+size])
        res = np.concatenate(res, 0)
        res = preprocess_image(res)

        return res

    def __len__(self):
        return (len(self.data) + self.batch_size - 1)//self.batch_size

# A Keras generator to evaluate on the HEAD MODEL on features already pre-computed.
# It computes only the upper triangular matrix of the cost matrix if y is None.
class ScoreGen(Sequence):
    def __init__(self, x, y=None, batch_size=2048, verbose=1):
        super(ScoreGen, self).__init__()
        self.x          = x
        self.y          = y
        self.batch_size = batch_size
        self.verbose    = verbose
        if y is None:
            self.y           = self.x
            self.ix, self.iy = np.triu_indices(x.shape[0],1)
        else:
            self.iy, self.ix = np.indices((y.shape[0],x.shape[0]))
            self.ix          = self.ix.reshape((self.ix.size,))
            self.iy          = self.iy.reshape((self.iy.size,))
        self.subbatch = (len(self.x) + self.batch_size - 1)//self.batch_size

    def __getitem__(self, index):
        start = index*self.batch_size
        end   = min(start + self.batch_size, len(self.ix))
        a     = self.y[self.iy[start:end],:]
        b     = self.x[self.ix[start:end],:]
        return [a,b]

    def __len__(self):
        return (len(self.ix) + self.batch_size - 1)//self.batch_size

def set_lr(model, lr):
    K.set_value(model.optimizer.lr, float(lr))

def get_lr(model):
    return K.get_value(model.optimizer.lr)

def score_reshape(score, x, y=None):
    """
    Tranformed the packed matrix 'score' into a square matrix.
    @param score the packed matrix
    @param x the first image feature tensor
    @param y the second image feature tensor if different from x
    @result the square matrix
    """
    if y is None:
        # When y is None, score is a packed upper triangular matrix.
        # Unpack, and transpose to form the symmetrical lower triangular matrix.
        m = np.zeros((x.shape[0],x.shape[0]), dtype=K.floatx())
        m[np.triu_indices(x.shape[0],1)] = score.squeeze()
        m += m.transpose()
    else:
        m        = np.zeros((y.shape[0],x.shape[0]), dtype=K.floatx())
        iy,ix    = np.indices((y.shape[0],x.shape[0]))
        ix       = ix.reshape((ix.size,))
        iy       = iy.reshape((iy.size,))
        m[iy,ix] = score.squeeze()
    return m

def compute_score_fast(verbose=1):
    """
    Compute the score matrix by scoring every pictures from the training set against every other picture O(n^2).
    """
    res = []
    batch = int(math.ceil(len(train) / LINEAR_ASSIGNMENT_SEGMENT_SIZE))
    for start in range(0, len(train), batch):
        end = min(len(train), start + batch)
        train_batch = train[start:end]

        features = branch_model.predict_generator(FeatureGen(train_batch, verbose=verbose), max_queue_size=12, verbose=0, use_multiprocessing=False)
        score    = head_model.predict_generator(ScoreGen(features, verbose=verbose), max_queue_size=12, verbose=0, use_multiprocessing=False)
        score    = score_reshape(score, features)
        res.append(score)

    return res

def make_steps_fast(step, ampl, verbose=2, batch_size=6):
    """
    Perform training epochs
    @param step Number of epochs to perform`
    @param ampl the K, the randomized component of the score matrix.
    """
    global w2ts, t2i, steps, features, score, histories

    # shuffle the training pictures
    random.shuffle(train)

    # Map training picture hash value to index in 'train' array
    t2i  = {}
    for i, t in enumerate(train): t2i[t] = i

    # Compute the match score for each picture pair of train data - square matrix score[i,j] is the score between image i and j, 0~1
    with timer("Compute score fast. "):
        score_list = compute_score_fast()

    for i in range(len(score_list)):
        score_list[i] += ampl*np.random.random_sample(size=score_list[i].shape)

    # Train the model for 'step' epochs
    history = model.fit_generator(
        TrainingDataFast(score_list, steps=step, batch_size=batch_size),
        initial_epoch=steps, epochs=steps + step, max_queue_size=12, workers=6, verbose=verbose,
        ).history
    steps += step

    # Collect history data
    history['epochs'] = steps
    history['ms'] = np.mean(score)
    history['lr'] = get_lr(model)
    print(history['epochs'], history['lr'], history['ms'])
    histories.append(history)

def run_validation(threshold):
    print("Running validation ...")
    print('Predict train with branch model...')
    ftrain = branch_model.predict_generator(FeatureGen(train), max_queue_size=20, verbose=0, use_multiprocessing=False)
    print('Predict validation with branch model...')
    fval = branch_model.predict_generator(FeatureGen(val), max_queue_size=20, verbose=0, use_multiprocessing=False)
    print('Score table calc...')
    score = head_model.predict_generator(ScoreGen(ftrain, fval), max_queue_size=20, verbose=0, use_multiprocessing=False)
    score = score_reshape(score, ftrain, fval)
    res = prepare_val_res(score, threshold)

    ## calculate score
    scores = []
    missing = []
    for img in res.keys():
        if img in tagged:
            score = apk([tagged[img]], res[img], k=5)
            scores.append(score)
        else:
            print("Warning! val image: {} not in tagged. ".format(img))
            missing.append(img)

    score = np.mean(np.array(scores))
    print("\nValid score: {:.6f}".format(score))
    return score

def get_model_to_finetune(saved_folder, fold):
    model, branch_model, head_model = build_model(64e-5, 0, img_shape=(BOX_SIZE, BOX_SIZE, 3))

    # Get best model from previous run
    best_models = glob.glob(MODELS_PATH + '{}/densenet121_384px_*.model'.format(saved_folder))
    best_score = -1.0
    best_weights = ''
    for m in best_models:
        score = float(m[:-6].split('_')[-1])
        if score > best_score:
            best_score = score
            best_weights = m

    print('Best score on previous run: {}. Use weights: {}'.format(best_score, best_weights))
    weights = best_weights
    model.load_weights(weights)
    return model, branch_model, head_model

if __name__ == '__main__':
    histories = []
    thr = 0.99
    score = 0
    steps = 0
    saved_folder = 'Res_v5_densenet121'
    if not os.path.isdir(MODELS_PATH + '{}'.format(saved_folder)):
        os.makedirs(MODELS_PATH + '{}'.format(saved_folder))

    version = 'v5_384px_finetune_' + str(RUN_FOLD)
    l2_rate = 0.0002
    model_str = 'ft'
    model, branch_model, head_model = get_model_to_finetune(saved_folder, RUN_FOLD)

    run_validation(thr)
    score = np.random.random_sample(size=(len(train), len(train)))

    # 15 epochs - freeze layers
    branch_model.trainable = False
    set_lr(model, 8e-5)
    optim = Adam(lr=8e-5)
    model.compile(optim, loss='binary_crossentropy', metrics=['binary_crossentropy', 'acc'])

    print("\n{} training block 1 {}".format('*'*17, '*'*17))
    make_steps_fast(15, 1, batch_size=16)
    val_score = run_validation(thr)
    model.save_weights(MODELS_PATH + '{}/{}_{}_temp0_{:.6f}.model'.format(saved_folder, model_str, version, val_score))

    # 4 x 5 epochs - unfreeze layers
    branch_model.trainable = True
    set_lr(model, 8e-5)
    optim  = Adam(lr=8e-5)
    model.compile(optim, loss='binary_crossentropy', metrics=['binary_crossentropy', 'acc'])

    print("\n{} training block 2 {}".format('*'*17, '*'*17))
    set_lr(model, 8e-5)
    for _ in range(4):
        make_steps_fast(5, 0.5)
    val_score = run_validation(thr)
    model.save_weights(MODELS_PATH + '{}/{}_{}_temp1_{:.6f}.model'.format(saved_folder, model_str, version, val_score))

    print("\n{} training block 3 {}".format('*'*17, '*'*17))
    set_lr(model, 4e-5)
    for _ in range(4):
        make_steps_fast(5, 0.25)
    val_score = run_validation(thr)
    model.save_weights(MODELS_PATH + '{}/{}_{}_temp2_{:.6f}.model'.format(saved_folder, model_str, version, val_score))

    print("\n{} training block 4 {}".format('*'*17, '*'*17))
    print("\n\nStarted new round: lr: 2e-5, rounds: 5, epochs: 5, ampl: 0.25\n")
    set_lr(model, 2e-5)
    for _ in range(4):
        make_steps_fast(5, 0.25)
    val_score = run_validation(thr)
    model_path = MODELS_PATH + '{}/{}_{}_temp3_{:.6f}.model'.format(saved_folder, model_str, version, val_score)
    model.save_weights(model_path)
    print("Saved model: ", model_path)

    # final fine tuning - 4 x 5 epochs
    print("\n{} training block 5 {}".format('*'*17, '*'*17))
    print("\nStarted new round: lr: 1e-5, rounds: 2, epochs: 5, ampl: 0.25\n")
    set_lr(model, 1e-5)
    for _ in range(4):
        make_steps_fast(5, 0.25)
    val_score = run_validation(thr)
    model_path = MODELS_PATH + '{}/{}_{}_final_v1_{:.6f}.model'.format(saved_folder, model_str, version, val_score)
    model.save_weights(model_path)
    print("Saved model:", model_path)

    # final fine tuning - 30 x 5 epochs (can stop at any time)
    print("\n{} training block 6 {}".format('*' * 17, '*' * 17))
    print("\nStarted new round: lr: 1e-5, rounds: 2, epochs: 5, ampl: 0.25\n")
    iters_to_run = 20
    set_lr(model, 8e-6)
    for iter1 in range(iters_to_run):
        ampl = 0.25 * ((iters_to_run - iter1) / iters_to_run)
        print('Iter: {} ampl: {}'.format(iters_to_run, ampl))
        make_steps_fast(5, ampl)
        val_score = run_validation(thr)
        model_path = MODELS_PATH + '{}/{}_{}_final_v2_{:.6f}.model'.format(saved_folder, model_str, version, val_score)
        model.save_weights(model_path)
        print("Saved model:", model_path)

    print("all done. ")

'''
Fold 0:
TB 0: Valid score: 0.958814
TB 1: Valid score: 0.910119
TV 2: Valid score: 0.939753
TB 3: Valid score: 0.963057
TB 4: Valid score: 0.963430
TB 5: Valid score: 0.966701
Valid score: 0.967093

Fold 1:
TB 0: Valid score: 0.955803
TB 1: Valid score: 0.889472
TB 2: Valid score: 0.926793
TB 3: Valid score: 0.957300
TB 4: Valid score: 0.963647
TB 5: Valid score: 0.966626
Valid score: 0.967650

Fold 2:
TB 0: Valid score: 0.956908
TB 1: Valid score: 0.926379
TB 2: Valid score: 0.934138
TB 3: Valid score: 0.950624
TB 4: Valid score: 0.955206
TB 5: Valid score: 0.960862
Valid score: 0.961138
Valid score: 0.961669

Fold 3:
TB 0: Valid score: 0.964243
TB 1: Valid score: 0.915626
TB 2: Valid score: 0.947841
TB 3: Valid score: 0.965813
TB 4: Valid score: 0.963862
TB 5: Valid score: 0.968446
Valid score: 0.973923
Valid score: 0.974151
'''


    tally = defaultdict(list)
    for i, item in enumerate(seq):
        tally[item].append(i)
    return ((key, locs) for key, locs in tally.items() if len(locs) > 1)

# pylint: disable=too-many-branches,too-many-statements
"""General interferometer compiler for the X class of circuits."""

from collections import defaultdict
import copy

import numpy as np
from thewalrus.symplectic import expand

from strawberryfields.program_utils import CircuitError, Command, group_operations

import strawberryfields.ops as ops

from .compiler import Compiler
from .gbs import GBS
from .gaussian_unitary import GaussianUnitary

def list_duplicates(seq):
    """Returns a generator representing the duplicated values in the sequence
    mapped to the indices they appear at."""

class Xunitary(Compiler):
    r"""General interferometer or unitary compiler for the X class of circuits.

    This compilation routine is performed at the interferometer/unitary matrix level.

    This compiler accepts the following gates, decompositions, and measurements:

    * :class:`~.ops.S2gate`
    * :class:`~.ops.Rgate`
    * :class:`~.ops.BSgate`
    * :class:`~.ops.MZgate`
    * :class:`~.ops.Interferometer`
    * :class:`~.ops.BipartiteGraphEmbed`
    * :class:`~.ops.MeasureFock`

    All ``S2gate`` operations, if present, **must** be placed at the start of the program,
    and match the X-series topology. That is, for a device with :math:`2N` modes, the
    ``S2gate`` operations can only be applied to modes :math:`(m, m+N)`.

    Subsequent operations represent the interferometer, and may consist of any combination of
    :class:`~.ops.BSgate`, :class:`~.ops.MZgate`,
    :class:`~.ops.Interferometer`, :class:`~.ops.BipartiteGraphEmbed`, as long as the unitary
    on modes :math:`(0, 1,\dots, N-1)` is repeated on modes :math:`(N, N+1, \dots, 2N-1)`.
    The unitary will automatically be compiled to match the topology of the X-series.

    Finally, the circuit must complete with Fock measurements.

    **Example**

    The compiler may be used on its own:

    >>> prog.compile(compiler="Xunitary")

    Alternatively, it can be combined with an X series device specification to include additional
    information, such as allowed parameter ranges.

    >>> eng = sf.RemoteEngine("X8")
    >>> spec = eng.device_spec
    >>> prog.compile(device=spec, compiler="Xunitary")
    """
    short_name = "Xunitary"
    interactive = False

    primitives = {
        "S2gate",
        "MeasureFock",
        "Rgate",
        "BSgate",
        "MZgate",
        "Interferometer",
    }

    decompositions = {
        "BipartiteGraphEmbed": {
            "mesh": "rectangular_symmetric",
            "drop_identity": False,
        },
    }

    def compile(self, seq, registers):
        # the number of modes in the provided program
        n_modes = len(registers)

        # Number of modes must be even
        if n_modes % 2 != 0:
            raise CircuitError("The X series only supports programs with an even number of modes.")
        half_n_modes = n_modes // 2
        # Call the GBS compiler to do basic measurement validation.
        # The GBS compiler also merges multiple measurement commands
        # into a single MeasureFock command at the end of the circuit.
        seq = GBS().compile(seq, registers)

        # ensure that all modes are measured
        if len(seq[-1].reg) != n_modes:
            raise CircuitError("All modes must be measured.")

        # Check circuit begins with two-mode squeezers
        A, B, C = group_operations(seq, lambda x: isinstance(x, ops.S2gate))
        # If there are no two-mode squeezers add squeezers at the beginning with squeezing param equal to zero.
        if B == []:
            initS2 = [
                Command(ops.S2gate(0, 0), [registers[i], registers[i + half_n_modes]])
                for i in range(half_n_modes)
            ]
            seq = initS2 + seq
            A, B, C = group_operations(seq, lambda x: isinstance(x, ops.S2gate))

        if A != []:
            raise CircuitError("There can be no operations before the S2gates.")

        regrefs = set()

        if B:
            # get set of circuit registers as a tuple for each S2gate
            regrefs = {(cmd.reg[0].ind, cmd.reg[1].ind) for cmd in B}

        # the set of allowed mode-tuples the S2gates must have
        allowed_modes = set(zip(range(0, half_n_modes), range(half_n_modes, n_modes)))

        if not regrefs.issubset(allowed_modes):
            raise CircuitError("S2gates do not appear on the correct modes.")

        # determine which modes do not have input S2gates specified
        missing = allowed_modes - regrefs

        for i, j in missing:
            B.insert(0, Command(ops.S2gate(0, 0), [registers[i], registers[j]]))

        # get list of circuit registers as a tuple for each S2gate
        regrefs = [(cmd.reg[0].ind, cmd.reg[1].ind) for cmd in B]

        # merge S2gates
        if len(regrefs) > half_n_modes:
            for mode, indices in list_duplicates(regrefs):
                r = 0
                phi = 0

                for k, i in enumerate(sorted(indices, reverse=True)):
                    removed_cmd = B.pop(i)
                    r += removed_cmd.op.p[0]
                    phi_new = removed_cmd.op.p[1]

                    if k > 0 and phi_new != phi:
                        raise CircuitError("Cannot merge S2gates with different phase values.")

                    phi = phi_new

                i, j = mode
                B.insert(indices[0], Command(ops.S2gate(r, phi), [registers[i], registers[j]]))

        meas_seq = [C[-1]]
        seq = GaussianUnitary().compile(C[:-1], registers)

        # extract the compiled symplectic matrix
        if seq == []:
            S = np.identity(2 * n_modes)
            used_modes = list(range(n_modes))
        else:
            S = seq[0].op.p[0]
            # determine the modes that are acted on by the symplectic transformation
            used_modes = [x.ind for x in seq[0].reg]

        if not np.allclose(S @ S.T, np.identity(len(S))):
            raise CircuitError(
                "The operations after squeezing do not correspond to an interferometer."
            )

        if len(used_modes) != n_modes:
            # The symplectic transformation acts on a subset of
            # the programs registers. We must expand the symplectic
            # matrix to one that acts on all registers.
            # simply extract the computed symplectic matrix
            S = expand(seq[0].op.p[0], used_modes, n_modes)

        U = S[:n_modes, :n_modes] - 1j * S[:n_modes, n_modes:]
        U11 = U[:half_n_modes, :half_n_modes]
        U12 = U[:half_n_modes, half_n_modes:]
        U21 = U[half_n_modes:, :half_n_modes]
        U22 = U[half_n_modes:, half_n_modes:]
        if not np.allclose(U12, 0) or not np.allclose(U21, 0):
            # Not a bipartite graph
            raise CircuitError(
                "The applied unitary cannot mix between the modes {}-{} and modes {}-{}.".format(
                    0, half_n_modes - 1, half_n_modes, n_modes - 1
                )
            )

        if not np.allclose(U11, U22):
            # Not a symmetric bipartite graph
            raise CircuitError(
                "The applied unitary on modes {}-{} must be identical to the applied unitary on modes {}-{}.".format(
                    0, half_n_modes - 1, half_n_modes, n_modes - 1
                )
            )
        U1 = ops.Interferometer(U11, mesh="rectangular_symmetric", drop_identity=False)._decompose(
            registers[:half_n_modes]
        )
        U2 = copy.deepcopy(U1)

        for Ui in U2:
            Ui.reg = [registers[r.ind + half_n_modes] for r in Ui.reg]

        return B + U1 + U2 + meas_seq


    LOG.info("Controller process is going shut down")

    SHUTDOWN_EVENT.set()

"""Main loop of the controller process."""

import atexit
import os
import sys
import threading

from decapod_common import config
from decapod_common import log
from decapod_common.models import execution
from decapod_common.models import server
from decapod_common.models import task
from decapod_controller import taskpool

CONF = config.make_controller_config()
"""Config."""

LOG = log.getLogger(__name__)
"""Logger."""

SHUTDOWN_EVENT = threading.Event()
"""Event which should be set by signal handler."""

TASK_POOL = taskpool.TaskPool(CONF["controller"]["worker_threads"])

def main():
    """Daemon main loop."""

    atexit.register(TASK_POOL.stop)

    LOG.info("Controller process has been started. PID %s", os.getpid())

    exc_info = None
    try:
        for tsk in task.Task.watch(stop_condition=SHUTDOWN_EVENT):
            LOG.info(
                "Fetched task %s (execution_id=%s, task_type=%s)",
                tsk._id, tsk.execution_id, tsk.task_type
            )
            try:
                tsk.bounce()
            except Exception as exc:
                LOG.warning("Cannot bounce task %s. Skip. %s", tsk._id, exc)
                continue

            if not possible_to_process(tsk):
                LOG.info("Task %s is temporarily skipped.", tsk._id)
            else:
                process_task(tsk)
    except BaseException:  # SystemExit is raised by shutdown_callback
        exc_info = sys.exc_info()
        if isinstance(exc_info[0], Exception):
            LOG.exception("Uncontroller exception has been raised: %s",
                          exc_info[0])

    wait_to_finish()

    LOG.info("Controller has been shutdown")

    if exc_info:
        exc = exc_info[0](exc_info[1])
        exc.__traceback__ = exc_info[2]
        raise exc

def wait_to_finish():
    SHUTDOWN_EVENT.set()

def shutdown_callback():
    """deamonocle callback to shutdown deamon."""

def possible_to_process(tsk):
    if tsk.task_type == task.TaskType.playbook:
        servers = get_servers_for_task(tsk.execution_id)
        try:
            server.ServerModel.lock_servers(servers)
        except ValueError:
            LOG.info("Cannot execute task because servers are locked.")
            return False

    return True

def process_task(tsk):
    LOG.info("Start to process task %s", tsk._id)

    if tsk.task_type in (task.TaskType.playbook,
                         task.TaskType.server_discovery):
        TASK_POOL.submit(tsk)
    elif tsk.task_type == task.TaskType.cancel:
        tsk.start()
        executing_task = tsk.get_executing_task()
        if executing_task:
            TASK_POOL.cancel(executing_task._id)
            tsk.complete()
        else:
            LOG.error("Cannot find executing task for %s", tsk._id)
            tsk.fail("Cannot find executing task")
    else:
        LOG.error("Unknown task %s", tsk._id)
        tsk.fail("Unknown task")

def get_servers_for_task(execution_id):
    execmodel = execution.ExecutionModel.find_by_model_id(execution_id)
    if not execmodel:
        return []

    return execmodel.servers


    blocks = []

    for block_string in parse_tf_file_for_block_strings(path):
        if block_string.strip().startswith("variable "):
            try:
                parsed = hcl.loads(block_string)
            except ValueError:
                log.bad(f"error parsing {path}")
                print(block_string)
                raise
            else:
                blocks.append(parsed)

    return blocks

import enum
import json
import re
import sys
from pathlib import Path
from typing import Generator, List

import hcl

from . import log

class State(enum.Enum):
    BLOCK = enum.auto()
    ESCAPE = enum.auto()
    ROOT = enum.auto()
    STRING = enum.auto()

def clean_block_string(block_string: str) -> str:

    # Remove comments.
    lines = []
    in_comment = False
    for line in block_string.splitlines():
        stripped = line.lstrip()
        if not stripped:
            continue
        if stripped.startswith("/*"):
            in_comment = True
        elif in_comment:
            if stripped.startswith("*/"):
                in_comment = False
        elif stripped.startswith("#"):
            continue
        else:
            lines.append(line)
    block_string = "\n".join(lines)

    # Add quotes around bare expressions
    # because the parser doesn't support them.

    space = r" *"
    identifier = r"[a-zA-Z0-9_-]+"
    not_number = r"(?!\d+$|\d+\.\d+$)"
    not_boolean = r"(?!true$|false$)"
    not_list = r"(?!\[)"
    not_map = r"(?!{)"
    not_string = r'(?!")'
    value = r"(\S.*$|$)"

    block_string = re.sub(
        f"^{space}({identifier}){space}={space}{not_number}{not_boolean}{not_list}{not_map}{not_string}({value})$",
        r'\1 = "\2"',
        block_string,
        flags=re.MULTILINE,
    )

    return block_string

def get_outputs_from_block(block: dict) -> Generator[dict, None, None]:

    if "output" not in block:
        return

    output = block["output"]

    if isinstance(output, dict):
        outputs = [output]
    else:
        outputs = output

    for output in outputs:
        for name, block in output.items():
            yield {"name": name, "value": block["value"]}

def parse_apply_outputs(stdout: str) -> dict:

    try:
        outputs_string = stdout.split("Outputs:\n", 1)[1]
    except IndexError:
        return {}

    cleaned = clean_block_string(outputs_string)
    if not cleaned:
        return {}

    try:
        return hcl.loads(cleaned)
    except ValueError as error:
        print(file=sys.stderr)
        log.bad("Error parsing:")
        print(stdout, file=sys.stderr)
        log.bad(f"Raising: {error.__class__.__name__}")
        raise

def parse_environment_variable_for_variables(name: str, value: str) -> dict:

    cleaned = clean_block_string(f"{name[7:]} = {value}")

    try:
        return hcl.loads(cleaned)
    except ValueError as error:
        print(file=sys.stderr)
        log.bad("Error parsing:")
        print(value, file=sys.stderr)
        log.bad(f"Raising: {error.__class__.__name__}")
        raise

def parse_tf_file_for_block_strings(path: Path) -> Generator[str, None, None]:

    states = [State.ROOT]

    buffer = []

    for char in read_chars_from_file(path):

        buffer.append(char)

        state = states[-1]

        if state is State.ROOT:

            if char == '"':
                states.append(State.STRING)
            elif char == "{":
                if buffer[-1:] != "\n":
                    buffer.append("\n")
                states.append(State.BLOCK)
            elif char == "\n":
                block = clean_block_string("".join(buffer))
                if block:
                    yield block
                buffer.clear()

        elif state is State.STRING:

            if char == "\\":
                states.append(State.ESCAPE)
            elif char == '"':
                states.pop()

        elif state is State.ESCAPE:

            states.pop()

        elif state is State.BLOCK:

            if char == '"':
                states.append(State.STRING)
            elif char == "{":
                states.append(State.BLOCK)
            elif char == "}":
                if buffer[-1:] != "\n":
                    buffer[-1] = "\n"
                    buffer.append(char)
                states.pop()
                if states[-1] is State.ROOT:
                    block = clean_block_string("".join(buffer))
                    if block:
                        yield block
                    buffer.clear()

        else:
            raise ValueError(state)

    # There shouldn't be anything left over.
    block = clean_block_string("".join(buffer))
    if block:
        raise ValueError(block)

def parse_tf_file_for_variables(path: Path) -> List[dict]:
    """
    This is a really bad parser for *.tf and *.tfvars files,
    with the only goal being to parse variable definitions and
    variable values.

    From http://example.com

        In future versions of Terraform, we will also support native tooling
        to convert HCL to JSON and JSON to HCL cleanly (including comments).

    When that happens, consider replacing this code with calls to that.

    """

def parse_json_file_for_blocks(path: Path) -> List[dict]:

    with open(path) as open_file:
        contents = json.load(open_file)

    if isinstance(contents, dict):
        blocks = [contents]
    else:
        blocks = contents

    return blocks

def parse_tfvars_file_for_variables(path: Path) -> List[dict]:
    cleaned = clean_block_string(path.read_text())
    if cleaned:
        return [hcl.loads(cleaned)]
    else:
        return []

def read_chars_from_file(path: Path) -> Generator[str, None, None]:
    with path.open() as open_file:
        while True:
            char = open_file.read(1)
            if char:
                yield char
            else:
                break


while True:
            char = open_file.read(1)
            if char:
                yield char
            else:
                break

import enum
import json
import re
import sys
from pathlib import Path
from typing import Generator, List

import hcl

from . import log

class State(enum.Enum):
    BLOCK = enum.auto()
    ESCAPE = enum.auto()
    ROOT = enum.auto()
    STRING = enum.auto()

def clean_block_string(block_string: str) -> str:

    # Remove comments.
    lines = []
    in_comment = False
    for line in block_string.splitlines():
        stripped = line.lstrip()
        if not stripped:
            continue
        if stripped.startswith("/*"):
            in_comment = True
        elif in_comment:
            if stripped.startswith("*/"):
                in_comment = False
        elif stripped.startswith("#"):
            continue
        else:
            lines.append(line)
    block_string = "\n".join(lines)

    # Add quotes around bare expressions
    # because the parser doesn't support them.

    space = r" *"
    identifier = r"[a-zA-Z0-9_-]+"
    not_number = r"(?!\d+$|\d+\.\d+$)"
    not_boolean = r"(?!true$|false$)"
    not_list = r"(?!\[)"
    not_map = r"(?!{)"
    not_string = r'(?!")'
    value = r"(\S.*$|$)"

    block_string = re.sub(
        f"^{space}({identifier}){space}={space}{not_number}{not_boolean}{not_list}{not_map}{not_string}({value})$",
        r'\1 = "\2"',
        block_string,
        flags=re.MULTILINE,
    )

    return block_string

def get_outputs_from_block(block: dict) -> Generator[dict, None, None]:

    if "output" not in block:
        return

    output = block["output"]

    if isinstance(output, dict):
        outputs = [output]
    else:
        outputs = output

    for output in outputs:
        for name, block in output.items():
            yield {"name": name, "value": block["value"]}

def parse_apply_outputs(stdout: str) -> dict:

    try:
        outputs_string = stdout.split("Outputs:\n", 1)[1]
    except IndexError:
        return {}

    cleaned = clean_block_string(outputs_string)
    if not cleaned:
        return {}

    try:
        return hcl.loads(cleaned)
    except ValueError as error:
        print(file=sys.stderr)
        log.bad("Error parsing:")
        print(stdout, file=sys.stderr)
        log.bad(f"Raising: {error.__class__.__name__}")
        raise

def parse_environment_variable_for_variables(name: str, value: str) -> dict:

    cleaned = clean_block_string(f"{name[7:]} = {value}")

    try:
        return hcl.loads(cleaned)
    except ValueError as error:
        print(file=sys.stderr)
        log.bad("Error parsing:")
        print(value, file=sys.stderr)
        log.bad(f"Raising: {error.__class__.__name__}")
        raise

def parse_tf_file_for_block_strings(path: Path) -> Generator[str, None, None]:

    states = [State.ROOT]

    buffer = []

    for char in read_chars_from_file(path):

        buffer.append(char)

        state = states[-1]

        if state is State.ROOT:

            if char == '"':
                states.append(State.STRING)
            elif char == "{":
                if buffer[-1:] != "\n":
                    buffer.append("\n")
                states.append(State.BLOCK)
            elif char == "\n":
                block = clean_block_string("".join(buffer))
                if block:
                    yield block
                buffer.clear()

        elif state is State.STRING:

            if char == "\\":
                states.append(State.ESCAPE)
            elif char == '"':
                states.pop()

        elif state is State.ESCAPE:

            states.pop()

        elif state is State.BLOCK:

            if char == '"':
                states.append(State.STRING)
            elif char == "{":
                states.append(State.BLOCK)
            elif char == "}":
                if buffer[-1:] != "\n":
                    buffer[-1] = "\n"
                    buffer.append(char)
                states.pop()
                if states[-1] is State.ROOT:
                    block = clean_block_string("".join(buffer))
                    if block:
                        yield block
                    buffer.clear()

        else:
            raise ValueError(state)

    # There shouldn't be anything left over.
    block = clean_block_string("".join(buffer))
    if block:
        raise ValueError(block)

def parse_tf_file_for_variables(path: Path) -> List[dict]:
    """
    This is a really bad parser for *.tf and *.tfvars files,
    with the only goal being to parse variable definitions and
    variable values.

    From http://example.com

        In future versions of Terraform, we will also support native tooling
        to convert HCL to JSON and JSON to HCL cleanly (including comments).

    When that happens, consider replacing this code with calls to that.

    """

    blocks = []

    for block_string in parse_tf_file_for_block_strings(path):
        if block_string.strip().startswith("variable "):
            try:
                parsed = hcl.loads(block_string)
            except ValueError:
                log.bad(f"error parsing {path}")
                print(block_string)
                raise
            else:
                blocks.append(parsed)

    return blocks

def parse_json_file_for_blocks(path: Path) -> List[dict]:

    with open(path) as open_file:
        contents = json.load(open_file)

    if isinstance(contents, dict):
        blocks = [contents]
    else:
        blocks = contents

    return blocks

def parse_tfvars_file_for_variables(path: Path) -> List[dict]:
    cleaned = clean_block_string(path.read_text())
    if cleaned:
        return [hcl.loads(cleaned)]
    else:
        return []

def read_chars_from_file(path: Path) -> Generator[str, None, None]:
    with path.open() as open_file:


while True:
        z =  mci.recvPlist()
        if not z:
            break
        completion = z.get('PercentComplete')
        if completion:
            print('Installing, %s: %s %% Complete' % (ipaPath, z['PercentComplete']))
        if z.get('Status') == 'Complete':
            print("Installation %s\n" % z['Status'])
            break

# This file is part of pymobiledevice2

from __future__ import print_function
import os
import warnings
from pymobiledevice2.lockdown import LockdownClient
from pymobiledevice2.afc import AFCClient, AFCShell
from optparse import OptionParser
warnings.warn(
"""The libraries upon which this program depends will soon be deprecated in
favor of the house_arrest.py and installation_proxy.py libraries.
See those files for example program written using the new libraries."""
)

def house_arrest(lockdown, applicationId):
    try:
        mis = lockdown.startService("com.apple.mobile.house_arrest")
    except:
        lockdown = LockdownClient()
        mis = lockdown.startService("com.apple.mobile.house_arrest")

    if mis == None:
        return
    mis.sendPlist({"Command": "VendDocuments", "Identifier": applicationId})
    res = mis.recvPlist()
    if res.get("Error"):
        print("Unable to Lookup the selected application: You probably trying to access to a system app...")
        return None
    return AFCClient(lockdown, service=mis)

def house_arrest_shell(lockdown, applicationId):
    afc =  house_arrest(lockdown, applicationId)
    if afc: AFCShell(client=afc).cmdloop()

"""
"Install"
"Upgrade"
"Uninstall"
"Lookup"
"Browse"
"Archive"
"Restore"
"RemoveArchive"
"LookupArchives"
"CheckCapabilitiesMatch"

installd
if stat("/var/mobile/tdmtanf") => "TDMTANF Bypass" => SignerIdentity bypass
"""

def mobile_install(lockdown,ipaPath):
    #Start afc service & upload ipa
    afc = AFCClient(lockdown)
    afc.set_file_contents("/" + os.path.basename(ipaPath), open(ipaPath,'rb').read())
    mci = lockdown.startService("com.apple.mobile.installation_proxy")
    mci.sendPlist({"Command":"Install",
                         "PackagePath": os.path.basename(ipaPath)})

def list_apps(lockdown):
    mci = lockdown.startService("com.apple.mobile.installation_proxy")
    mci.sendPlist({"Command":"Lookup"})
    res = mci.recvPlist()
    for app in res["LookupResult"].values():
        if app.get("ApplicationType") != "System":
            print(app["CFBundleIdentifier"], "=>", app.get("Container"))
        else:
            print(app["CFBundleIdentifier"], "=> N/A")

def get_apps_BundleID(lockdown,appType="User"):
    appList = []
    mci = lockdown.startService("com.apple.mobile.installation_proxy")
    mci.sendPlist({"Command":"Lookup"})
    res = mci.recvPlist()
    for app in res["LookupResult"].values():
        if app.get("ApplicationType")  == appType:
            appList.append(app["CFBundleIdentifier"])
    mci.close()
    return appList

if __name__ == "__main__":
    parser = OptionParser(usage="%prog")
    parser.add_option("-l", "--list", dest="list", action="store_true", default=False,
                  help="List installed applications (non system apps)")
    parser.add_option("-a", "--app", dest="app", action="store", default=None,
                  metavar="APPID", help="Access application files with AFC")
    parser.add_option("-i", "--install", dest="installapp", action="store", default=None,
                  metavar="FILE", help="Install an application package")

    (options, args) = parser.parse_args()
    if options.list:
        lockdown = LockdownClient()
        list_apps(lockdown)
    elif options.app:
        lockdown = LockdownClient()
        house_arrest_shell(lockdown, options.app)
    elif options.installapp:
        lockdown = LockdownClient()
        mobile_install(lockdown, options.installapp)
    else:
        parser.print_help()


return serpent_cbc_encrypt(key, data, iv=iv or b"\x00" * 16)

from typing import Optional

from .components.pyserpent import serpent_cbc_encrypt, serpent_cbc_decrypt

__all__ = ["serpent"]

class SerpentCbc:
    def encrypt(self, key: bytes, data: bytes, iv: Optional[bytes] = None) -> bytes:
        """
        Encrypts buffer using Serpent algorithm in CBC mode.

        :param key: Cryptographic key (4-32 bytes, must be multiple of four)
        :type key: bytes
        :param data: Buffer to be encrypted
        :type data: bytes
        :param iv: Initialization vector (defaults to `b"\\\\x00" * 16`)
        :type iv: bytes, optional
        :return: Encrypted data
        :rtype: bytes
        """

    def decrypt(self, key: bytes, data: bytes, iv: Optional[bytes] = None) -> bytes:
        """
        Decrypts buffer using Serpent algorithm in CBC mode.

        :param key: Cryptographic key (4-32 bytes, must be multiple of four)
        :type key: bytes
        :param data: Buffer to be decrypted
        :type data: bytes
        :param iv: Initialization vector (defaults to `b"\\\\x00" * 16`)
        :type iv: bytes, optional
        :return: Decrypted data
        :rtype: bytes
        """
        return serpent_cbc_decrypt(key, data, iv=iv or b"\x00" * 16)

class Serpent:
    cbc = SerpentCbc()

serpent = Serpent()


    for server in cluster:
        s = manager.create_server(cluster[server])

    for server in cluster:
        cluster[server].ensure_started()

    return manager.get_servers()

def create_cluster(manager, cluster):
    """Create all servers in cluster."""

def firewall_test(manager, firewall_rules, server):
    """Run tests on firewall rules."""
    # add 1 rule and remove it
    server.add_firewall_rule(firewall_rules[0])

    fs = server.get_firewall_rules()
    assert len(fs) == 1

    fs[0].destroy()
    fs = server.get_firewall_rules()
    assert len(fs) == 0

    # add several rules and remove them
    server.configure_firewall(firewall_rules)

    fs = server.get_firewall_rules()
    assert len(fs) == 2

    for _ in fs:
        manager.delete_firewall_rule(server.uuid, 1)

    fs = server.get_firewall_rules()
    assert len(fs) == 0

def server_test(manager, server):
    """Run tests on a server instance."""
    server.populate()

    server.core_number = '3'
    server.memory_amount = '1024'
    server.save()

    server.add_ip()

    storage = manager.create_storage(size=10, tier='maxiops', zone='uk-lon1')
    server.add_storage(storage)

    server.start()

    # sync new info from API and assert the changes from above have happened
    server.populate()
    assert server.core_number == '3'
    assert server.memory_amount == '1024'
    assert len(server.storage_devices) == 3
    assert len(server.ip_addresses) == 4

    server.ensure_started()

def tag_servers_test(manager, tags, cluster):
    """Run tests on tags."""
    # create tags
    for t in tags:
        manager.create_tag(str(t))

    cluster['web1'].add_tags(['testweb'])
    cluster['web2'].add_tags(['testweb'])
    cluster['lb'].add_tags([tags[1]])  # tags[1] is 'db'
    cluster['db'].add_tags(['testlb'])

    fetched_servers = manager.get_servers(tags_has_one=['testlb'])
    assert len(fetched_servers) == 1
    assert fetched_servers[0].tags[0] == 'testlb'


"""
Context library to allow dropping tensorboardX anywhere in the codebase.
If there is no SummaryWriter in the context, function calls will be no-op.

Usage:

    writer = SummaryWriter()

    with summary_writer_context(writer):
        some_func()

    def some_func():
        SummaryWriterContext.add_scalar("foo", tensor)

"""

import contextlib
import logging
from typing import Any, Dict, List

from torch.utils.tensorboard import SummaryWriter

logger = logging.getLogger(__name__)

class SummaryWriterContextMeta(type):
    def __getattr__(cls, func):
        if func.startswith("__"):
            return super().__getattr__(func)

        if not cls._writer_stacks:

            def noop(*args, **kwargs):
                return

            return noop

        writer = cls._writer_stacks[-1]

        def call(*args, **kwargs):
            if "global_step" not in kwargs:
                kwargs["global_step"] = cls._global_step
            try:
                return getattr(writer, func)(*args, **kwargs)
            except Exception as e:
                if hasattr(writer, "exceptions_to_ignore") and isinstance(
                    e, writer.exceptions_to_ignore
                ):
                    logger.warning("Ignoring exception: {}".format(e))
                    if hasattr(writer, "exception_logging_func"):
                        writer.exception_logging_func(e)
                    return
                raise

        return call

class SummaryWriterContext(metaclass=SummaryWriterContextMeta):
    _writer_stacks: List[SummaryWriter] = []
    _global_step = 0
    _custom_scalars: Dict[str, Any] = {}

    @classmethod
    def _reset_globals(cls):
        cls._global_step = 0
        cls._custom_scalars = {}

    @classmethod
    def increase_global_step(cls):
        cls._global_step += 1

    @classmethod
    def add_histogram(cls, key, val, *args, **kwargs):
        try:
            return cls.__getattr__("add_histogram")(key, val, *args, **kwargs)
        except ValueError:
            logger.warning(f"Cannot create histogram for {key}, got values: {val}")

    @classmethod
    def add_custom_scalars(cls, writer):
        """
        Call this once you are satisfied setting up custom scalar
        """

    @classmethod
    def add_custom_scalars_multilinechart(cls, tags, category=None, title=None):
        assert category and title, "category & title must be set"
        if category not in cls._custom_scalars:
            cls._custom_scalars[category] = {}
        assert (
            title not in cls._custom_scalars[category]
        ), "Title ({}) is already in category ({})".format(title, category)
        cls._custom_scalars[category][title] = ["Multiline", tags]

    @classmethod
    def push(cls, writer):
        assert isinstance(
            writer, SummaryWriter
        ), "writer is not a SummaryWriter: {}".format(writer)
        cls._writer_stacks.append(writer)

    @classmethod
    def pop(cls):
        return cls._writer_stacks.pop()

@contextlib.contextmanager
def summary_writer_context(writer):
    if writer is not None:
        SummaryWriterContext.push(writer)
    try:
        yield
    finally:
        if writer is not None:
            SummaryWriterContext.pop()


    assert isinstance(account_name, str), "Account is not a string: {}".format(account_name)
    assert isinstance(account_types, AccountTypes), (
        "Account types has invalid type: {}".format(account_types))
    account_type = get_account_type(account_name)
    return account_type == account_types.equity

"""Definition for global account types.

This is where we keep the global account types value and definition.

Note that it's unfortunate that we're using globals and side-effect here, but
this is the best solution in the short-term, the account types are used
in too many places to pass around that state everywhere. Maybe we change
this later on.
"""
__copyright__ = "Copyright (C) 2014-2017  Martin Blais"
__license__ = "GNU GPLv2"

import re
from collections import namedtuple
from typing import Tuple

from beancount.core import account
from beancount.core.account import Account

# A tuple that contains the names of the root accounts.
#   assets: a str, the name of the prefix for the Asset subaccounts.
#   liabilities: a str, the name of the prefix for the Liabilities subaccounts.
#   equity: a str, the name of the prefix for the Equity subaccounts.
#   income: a str, the name of the prefix for the Income subaccounts.
#   expenses: a str, the name of the prefix for the Expenses subaccounts.
AccountTypes = namedtuple('AccountTypes', "assets liabilities equity income expenses")

# Default values for root accounts.
DEFAULT_ACCOUNT_TYPES = AccountTypes("Assets",
                                     "Liabilities",
                                     "Equity",
                                     "Income",
                                     "Expenses")

def get_account_type(account_name: Account):
    """Return the type of this account's name.

    Warning: No check is made on the validity of the account type. This merely
    returns the root account of the corresponding account name.

    Args:
      account_name: A string, the name of the account whose type is to return.
    Returns:
      A string, the type of the account in 'account_name'.

    """
    assert isinstance(account_name, str), "Account is not a string: {}".format(account_name)
    return account.split(account_name)[0]

def get_account_sort_key(account_types: AccountTypes,
                         account_name: Account) -> Tuple[str, Account]:
    """Return a tuple that can be used to order/sort account names.

    Args:
      account_types: An instance of AccountTypes, a tuple of account type names.
    Returns:
      An object to use as the 'key' argument to the sort function.
    """
    return (account_types.index(get_account_type(account_name)), account_name)

def is_account_type(account_type: str, account_name: Account) -> bool:
    """Return the type of this account's name.

    Warning: No check is made on the validity of the account type. This merely
    returns the root account of the corresponding account name.

    Args:
      account_type: A string, the prefix type of the account.
      account_name: A string, the name of the account whose type is to return.
    Returns:
      A boolean, true if the account is of the given type.
    """
    return bool(re.match('^{}{}'.format(account_type, account.sep), account_name))

def is_root_account(account_name: Account) -> bool:
    """Return true if the account name is a root account.

    This function does not verify whether the account root is a valid
    one, just that it is a root account or not.

    Args:
      account_name: A string, the name of the account to check for.
    Returns:
      A boolean, true if the account is root account.
    """
    assert isinstance(account_name, str), "Account is not a string: {}".format(account_name)
    return (account_name and
            bool(re.match(r'([A-Z][A-Za-z0-9\-]+)$', account_name)))

def is_balance_sheet_account(account_name: Account, account_types: AccountTypes) -> bool:
    """Return true if the given account is a balance sheet account.
    Assets, liabilities and equity accounts are balance sheet accounts.

    Args:
      account_name: A string, an account name.
      account_types: An instance of AccountTypes.
    Returns:
      A boolean, true if the account is a balance sheet account.
    """
    assert isinstance(account_name, str), "Account is not a string: {}".format(account_name)
    assert isinstance(account_types, AccountTypes), (
        "Account types has invalid type: {}".format(account_types))
    account_type = get_account_type(account_name)
    return account_type in (account_types.assets,
                            account_types.liabilities,
                            account_types.equity)

def is_income_statement_account(account_name: Account, account_types: AccountTypes) -> bool:
    """Return true if the given account is an income statement account.
    Income and expense accounts are income statement accounts.

    Args:
      account_name: A string, an account name.
      account_types: An instance of AccountTypes.
    Returns:
      A boolean, true if the account is an income statement account.
    """
    assert isinstance(account_name, str), "Account is not a string: {}".format(account_name)
    assert isinstance(account_types, AccountTypes), (
        "Account types has invalid type: {}".format(account_types))
    account_type = get_account_type(account_name)
    return account_type in (account_types.income,
                            account_types.expenses)

def is_equity_account(account_name: Account, account_types: AccountTypes) -> bool:
    """Return true if the given account is an equity account.

    Args:
      account_name: A string, an account name.
      account_types: An instance of AccountTypes.
    Returns:
      A boolean, true if the account is an equity account.
    """

def is_inverted_account(account_name: Account, account_types: AccountTypes) -> bool:
    """Return true if the given account has inverted signs.

    An inverted sign is the inverse as you'd expect in an external report, i.e.,
    with all positive signs expected.

    Args:
      account_name: A string, an account name.
      account_types: An instance of AccountTypes.
    Returns:
      A boolean, true if the account has an inverted sign.
    """
    assert isinstance(account_name, str), "Account is not a string: {}".format(account_name)
    assert isinstance(account_types, AccountTypes), (
        "Account types has invalid type: {}".format(account_types))
    account_type = get_account_type(account_name)
    return account_type in (account_types.liabilities,
                            account_types.income,
                            account_types.equity)

def get_account_sign(account_name: Account, account_types: AccountTypes=None) -> int:
    """Return the sign of the normal balance of a particular account.

    Args:
      account_name: A string, the name of the account whose sign is to return.
      account_types: An optional instance of the current account_types.
    Returns:
      +1 or -1, depending on the account's type.
    """
    if account_types is None:
        account_types = DEFAULT_ACCOUNT_TYPES
    assert isinstance(account_name, str), "Account is not a string: {}".format(account_name)
    account_type = get_account_type(account_name)
    return (+1
            if account_type in (account_types.assets,
                                account_types.expenses)
            else -1)


        import json

        try:
            return json.loads(stream)
        except ValueError as e:
            cause = "invalid json format. %s" % str(e)
            raise InvalidFormatError(cause=cause)

#     Santiago Dueñas <example@nnthink.com>

import logging

from ..db.model import UniqueIdentity, Identity, Profile,\
    Enrollment, Organization, Domain, Country, MatchingBlacklist
from ..exceptions import InvalidFormatError, InvalidDateError
from ..utils import str_to_datetime

logger = logging.getLogger(__name__)

class SortingHatParser(object):
    """Parse identities and organizations using Sorting Hat format.

    The Sorting Hat data format is a JSON stream that contains
    information about unique identities and organizations.

    The unique identities are stored in an object named 'uidentities'.
    The keys of this object are the UUID of the unique identities.
    Each unique identity object stores a list of identities and
    enrollments.

    Organizations are stored in 'organizations' object. Its keys
    are the name of the organizations and each organization object is
    related to a list of domains.

    :param stream: stream to parse

    :raises InvalidFormatError: raised when the format of the stream is
        not valid.
    """

    def __init__(self, stream):
        self._blacklist = {}
        self._identities = []
        self._organizations = {}
        self.__parse(stream)

    @property
    def blacklist(self):
        bl = [b for b in self._blacklist.values()]
        bl.sort(key=lambda b: b.excluded)
        return bl

    @property
    def identities(self):
        self._identities.sort(key=lambda u: u.uuid)
        return [u for u in self._identities]

    @property
    def organizations(self):
        orgs = [o for o in self._organizations.values()]
        orgs.sort(key=lambda o: o.name)
        return orgs

    def __parse(self, stream):
        """Parse Sorting Hat stream"""

        if not stream:
            raise InvalidFormatError(cause="stream cannot be empty or None")

        json = self.__load_json(stream)

        self.__parse_organizations(json)
        self.__parse_identities(json)
        self.__parse_blacklist(json)

    def __parse_blacklist(self, json):
        """Parse blacklist entries using Sorting Hat format.

        The Sorting Hat blacklist format is a JSON stream that
        stores a list of blacklisted entries.

        Next, there is an example of a valid stream:

        {
            "blacklist": [
                "John Doe",
                "John Smith",
                "root@example.com"
            ]
        }

        :param stream: stream to parse

        :raises InvalidFormatError: raised when the format of the stream is
            not valid.
        """
        try:
            for entry in json['blacklist']:
                if not entry:
                    msg = "invalid json format. Blacklist entries cannot be null or empty"
                    raise InvalidFormatError(cause=msg)

                excluded = self.__encode(entry)

                bl = self._blacklist.get(excluded, None)

                if not bl:
                    bl = MatchingBlacklist(excluded=excluded)
                    self._blacklist[excluded] = bl
        except KeyError as e:
            msg = "invalid json format. Attribute %s not found" % e.args
            raise InvalidFormatError(cause=msg)

    def __parse_identities(self, json):
        """Parse identities using Sorting Hat format.

        The Sorting Hat identities format is a JSON stream on which its
        keys are the UUID of the unique identities. Each unique identity
        object has a list of identities and enrollments.

        When the unique identity does not have a UUID, it will be considered
        as an anonymous unique identity. This means that the UUID of these
        identities will be regenerated during the loading process.

        Next, there is an example of a valid stream:

        {
            "uidentities": {
                "johnsmith@example.net": {
                    "enrollments": [],
                    "identities": [],
                    "uuid": null
                },
                "03e12d00e37fd45593c49a5a5a1652deca4cf302": {
                    "enrollments": [
                        {
                            "end": "2100-01-01T00:00:00",
                            "start": "1900-01-01T00:00:00",
                            "organization": "Example",
                            "uuid": "03e12d00e37fd45593c49a5a5a1652deca4cf302"
                        }
                    ],
                    "identities": [
                        {
                            "email": "jsmith@example.com",
                            "id": "03e12d00e37fd45593c49a5a5a1652deca4cf302",
                            "name": "John Smith",
                            "source": "scm",
                            "username": "jsmith",
                            "uuid": "03e12d00e37fd45593c49a5a5a1652deca4cf302"
                        },
                        {
                            "email": "jsmith@example.com",
                            "id": "75d95d6c8492fd36d24a18bd45d62161e05fbc97",
                            "name": "John Smith",
                            "source": "scm",
                            "username": null,
                            "uuid": "03e12d00e37fd45593c49a5a5a1652deca4cf302"
                        }
                    ],
                    "profile": {
                        "country": {
                            "alpha3": "USA",
                            "code": "US",
                            "name": "United States of America"
                        },
                        "email": "jsmith@example.com",
                        "gender": "male",
                        "gender_acc": 100,
                        "name": null,
                        "is_bot": true,
                        "uuid": "03e12d00e37fd45593c49a5a5a1652deca4cf302"
                    },
                    "uuid": "03e12d00e37fd45593c49a5a5a1652deca4cf302"
                }
            }
        }

        :param stream: stream to parse

        :raises InvalidFormatError: raised when the format of the stream is
            not valid.
        """
        try:
            for uidentity in json['uidentities'].values():
                uuid = self.__encode(uidentity['uuid'])

                uid = UniqueIdentity(uuid=uuid)

                if uidentity['profile']:
                    profile = uidentity['profile']

                    if type(profile['is_bot']) != bool:
                        msg = "invalid json format. 'is_bot' must have a bool value"
                        raise InvalidFormatError(cause=msg)

                    is_bot = profile['is_bot']

                    gender = profile.get('gender', None)

                    if gender is not None:
                        gender = self.__encode(gender)

                    gender_acc = profile.get('gender_acc', None)

                    if gender_acc is not None:
                        if type(gender_acc) != int:
                            msg = "invalid json format. 'gender_acc' must have an integer value"
                            raise InvalidFormatError(cause=msg)
                        elif not 0 <= gender_acc <= 100:
                            msg = "invalid json format. 'gender_acc' is not in range (0,100)"
                            raise InvalidFormatError(cause=msg)

                    name = self.__encode(profile['name'])
                    email = self.__encode(profile['email'])

                    prf = Profile(uuid=uuid, name=name, email=email,
                                  gender=gender, gender_acc=gender_acc,
                                  is_bot=is_bot)

                    if profile['country']:
                        alpha3 = self.__encode(profile['country']['alpha3'])
                        code = self.__encode(profile['country']['code'])
                        name = self.__encode(profile['country']['name'])

                        c = Country(alpha3=alpha3, code=code, name=name)

                        prf.country_code = code
                        prf.country = c

                    uid.profile = prf

                for identity in uidentity['identities']:
                    identity_id = self.__encode(identity['id'])
                    name = self.__encode(identity['name'])
                    email = self.__encode(identity['email'])
                    username = self.__encode(identity['username'])
                    source = self.__encode(identity['source'])

                    sh_id = Identity(id=identity_id, name=name,
                                     email=email, username=username,
                                     source=source, uuid=uuid)

                    uid.identities.append(sh_id)

                for enrollment in uidentity['enrollments']:
                    organization = self.__encode(enrollment['organization'])

                    org = self._organizations.get(organization, None)

                    if not org:
                        org = Organization(name=organization)
                        self._organizations[organization] = org

                    try:
                        start = str_to_datetime(enrollment['start'])
                        end = str_to_datetime(enrollment['end'])
                    except InvalidDateError as e:
                        raise InvalidFormatError(cause=str(e))

                    rol = Enrollment(start=start, end=end, organization=org)

                    uid.enrollments.append(rol)

                self._identities.append(uid)
        except KeyError as e:
            msg = "invalid json format. Attribute %s not found" % e.args
            raise InvalidFormatError(cause=msg)

    def __parse_organizations(self, json):
        """Parse organizations using Sorting Hat format.

        The Sorting Hat organizations format is a JSON stream which
        its keys are the name of the organizations. Each organization
        object has a list of domains. For instance:

        {
            "organizations": {
                "Bitergia": [
                    {
                        "domain": "api.bitergia.com",
                        "is_top": false
                    },
                    {
                        "domain": "bitergia.com",
                        "is_top": true
                    }
                ],
                "Example": []
            },
            "time": "2015-01-20 20:10:56.133378"
        }

        :param json: stream to parse

        :raises InvalidFormatError: raised when the format of the stream is
            not valid.
        """
        try:
            for organization in json['organizations']:
                name = self.__encode(organization)

                org = self._organizations.get(name, None)

                if not org:
                    org = Organization(name=name)
                    self._organizations[name] = org

                domains = json['organizations'][organization]

                for domain in domains:
                    if type(domain['is_top']) != bool:
                        msg = "invalid json format. 'is_top' must have a bool value"
                        raise InvalidFormatError(cause=msg)

                    dom = Domain(domain=domain['domain'],
                                 is_top_domain=domain['is_top'])
                    org.domains.append(dom)
        except KeyError as e:
            msg = "invalid json format. Attribute %s not found" % e.args
            raise InvalidFormatError(cause=msg)

    def __load_json(self, stream):
        """Load json stream into a dict object """

    def __encode(self, s):
        return s if s else None