# -*- coding: utf8 -*-
import numpy as np
import pickle
import Tree

class RAE(object):
    """Class to use Recursive Neural Network on Tree

    Usage
    -----

    Methods
    -------

    """
    def __init__(self, vocab={}, dim=30, r=0.0001, reg=1):
        self.dim = dim

        #Initiate V, the tensor operator
        self.V = np.random.uniform(-r, r, size=(dim+1, 2*dim+2, 2*dim+2))
        self.V = (self.V+np.transpose(self.V, axes=[0, 2, 1]))/2

        #Initiate W, the linear operator
        self.W = np.random.uniform(-r, r, size=(dim+1, 2*dim+2))

        #Initiate Ws, the linear operator
        self.Ws = np.random.uniform(-r, r, size=(2, dim+1))

        #Initiate VE, the encoder tensor
        self.Ve = np.random.uniform(-r, r, size=(2*dim+2, dim+1, dim+1))
        self.Ve = (self.Ve+np.transpose(self.Ve, axes=[0, 2, 1]))/2

        #Initiate VE, the encoder tensor
        self.We = np.random.uniform(-r, r, size=(2*dim+2, dim+1))

        #Initiate L, the Lexicon representation
        self.L = np.random.uniform(-r, r, size=(len(vocab), dim))

        #Parameter holder
        self.params = {}
        self.params['V'] = self.V
        self.params['W'] = self.W
        self.params['Ws'] = self.Ws
        self.params['L'] = self.L
        self.params['Ve'] = self.Ve
        self.params['We'] = self.We

        #Regularisation
        self.reg = {'V': 0.001*reg, 'W': 0.001*reg, 'Ws': 0.0001*reg, 'L': 0.0001*reg,
                    'We': reg*0.0001, 'Ve': reg*0.0001}

        self.vocab = {}
        self.index = {}
        for i, w in enumerate(vocab):
            self.vocab[w] = i
            self.index[i] = w

        self.f = lambda X: np.tanh(X.T.dot(self.V).dot(X) + self.W.dot(X))
        self.grad = lambda f: 1-f**2
        self.dec = lambda X: np.tanh(X.T.dot(self.Ve).dot(X) + self.We.dot(X))

        self.y = lambda X: np.exp(self.Ws.dot(X).clip(-500, 700)) \
            / sum(np.exp(self.Ws.dot(X).clip(-500, 700)))

        self.norm = lambda X: np.sum(X*X)

    def save(self, saveFile):
        with open(saveFile, 'wb') as output:
            pickle.dump(self.dim, output, -1)
            pickle.dump(self.vocab, output, -1)
            for k in self.params.keys():
                pickle.dump(self.params[k], output, -1)
            for k in self.params.keys():
                pickle.dump(self.reg[k], output, -1)

    def load(self, loadFile):
        with open(loadFile, 'rb') as input:
            self.dim = pickle.load(input)
            self.vocab = pickle.load(input)
            for k in self.params.keys():
                self.params[k] = pickle.load(input)
            for k in self.params.keys():
                self.reg[k] = pickle.load(input)

    def error(self, val_set):
        '''
        Calcule l'erreur moyenne sur le set val_set avec les parametres actuel
        '''
        errorVal = 0.0
        for X_tree in val_set:
            errorVal += self.forward_pass(X_tree)
        return errorVal/len(val_set)

    def forward_pass(self, X_tree):
        '''
        Effectue la passe forward et calcule l'erreur actuelle pour l'arbre X_tree
        '''
        errorVal = 0.0
        for n in X_tree.leaf:
            # Met a jour le mot avec le Lexicon courant
            n.X = np.append(self.L[self.vocab[n.word]], 0)
            n.ypred = self.y(n.X)  # Mise a jour du label predit
            n.ypred += 1e-300*(n.ypred == 0)
            assert (n.ypred != 0).all()
            errorVal += -np.sum(n.y*np.log(n.ypred))

        for p, [a, b] in X_tree.parcours:
            aT = X_tree.nodes[a]  #
            bT = X_tree.nodes[b]  # Recupere pour chaque triplet parent,enfant1/2 les noeuds
            pT = X_tree.nodes[p]
            if aT.order < bT.order:
                X = np.append(aT.X, bT.X)
            else:
                X = np.append(bT.X, aT.X)
            pT.X = self.f(X)  # Mise a jour du decripteur du parent
            pT.X[-1] = 1

            pT.ypred = self.y(pT.X)  # Mise a jour du label predit
            pT.ypred += 1e-300*(pT.ypred == 0)
            errorVal += -np.sum(pT.y*np.log(pT.ypred))
        pT.c = pT.X
        for p, [a, b] in X_tree.parcours[::-1]:
            aT = X_tree.nodes[a]  #
            bT = X_tree.nodes[b]  # Recupere pour chaque triplet parent,enfant1/2 les noeuds
            pT = X_tree.nodes[p]
            C = self.dec(pT.c)
            #Propagation des erreurs vers le bas
            if aT.order < bT.order:  # Si aT est le noeud de gauche
                aT.c = C[:self.dim+1]
                bT.c = C[self.dim+1:]
            else:  # aT est a droite
                bT.c = C[:self.dim+1]
                aT.c = C[self.dim+1:]

        nn = len(X_tree.nodes)
        for n in X_tree.nodes:
            errorVal += np.sum((n.c-n.X)**2) / nn
        #E = sum([(self.y(n.X) - n.y) for n in X_tree.nodes])
        #print E
        #return self.Ws.dot(pT.X) -> Pas besoin de retourner le lbel, il faut le maj aussi?
        return errorVal

    def backward_pass(self, X_tree, w_root=1):
        '''
        Retourne le gradient du a l'erreur commise sur l'arbre X_tree
        Attention: suppose la forward_pass faite
        '''
        #Initialise les gradients
        grad = {}
        for k in self.params.keys():
            grad[k] = np.zeros(self.params[k].shape)

        #Initialise les deltas
        for n in X_tree.nodes:
            gX = self.grad(n.X)
            gC = self.grad(n.c)
            n.d = self.Ws.T.dot(n.ypred-n.y)*gX*w_root - 2*(n.c-n.X)*gX
            n.dr = 2*(n.c-n.X)*gC

        n.d += self.Ws.T.dot(n.ypred-n.y)*gX*(1-w_root)

        err_rec = 0

        for p, [a, b] in X_tree.parcours:
            aT = X_tree.nodes[a]  #
            bT = X_tree.nodes[b]  # Recupere pour chaque triplet parent,enfant1/2 les noeuds
            pT = X_tree.nodes[p]
            if aT.order < bT.order:
                dR = np.append(aT.dr, bT.dr)
            else:
                dR = np.append(bT.dr, aT.dr)
            pT.dr += dR.dot(self.We) + 2*dR.dot(self.Ve.dot(pT.X))
            grad['We'] += np.outer(dR, pT.X)
            grad['Ve'] += np.tensordot(dR, np.outer(pT.X, pT.X), axes=0)

        #Descend dans l arbre
        for p, [a, b] in X_tree.parcours[::-1]:
            aT = X_tree.nodes[a]  #
            bT = X_tree.nodes[b]  # Recupere pour chaque triplet parent,enfant1/2 les noeuds
            pT = X_tree.nodes[p]
            #Propagation des erreurs vers le bas
            if aT.order < bT.order:  # Si aT est le noeud de gauche
                X = np.append(aT.X, bT.X)
                ddown = (self.W.T.dot(pT.d+pT.dr)+2*(pT.d+pT.dr).dot(self.V.dot(X)))
                aT.d += ddown[:self.dim+1]
                bT.d += ddown[self.dim+1:]
            else:  # aT est a droite
                X = np.append(bT.X, aT.X)
                ddown = (self.W.T.dot(pT.d+pT.dr)+2*(pT.d+pT.dr).dot(self.V.dot(X)))
                aT.d += ddown[self.dim+1:]
                bT.d += ddown[:self.dim+1]
            err_rec += np.sum((aT.c-aT.X)**2)
            err_rec += np.sum((bT.c-bT.X)**2)
            #Contribution aux gradients du pT
            grad['Ws'] += np.outer(pT.ypred-pT.y, pT.X)
            grad['V'] += np.tensordot(pT.d + pT.dr, np.outer(X, X), axes=0)
            grad['W'] += np.outer(pT.d + pT.dr, X)

        #Contribution des feuilles
        for n in X_tree.leaf:
            grad['L'][self.vocab[n.word]] += n.d[:-1] + n.dr[:-1]
            grad['Ws'] += np.outer(n.ypred-n.y, n.X)

        return grad

    def train(self, X_trees, learning_rate=0.01, mini_batch_size=27,
              warm_start=True, r=0.0001, max_iter=1000, val_set=[],
              n_check=100, strat='AdaGrad', w_root=1,
              bin=False, reset_freq=-1, save_tmp='tmp.pkl', n_stop=4):
        '''
        Training avec AdaGrad (Dutchi et al.), prends en entrée une liste d'arbres X_trees
        '''
        #Remise à zero du modele
        if not warm_start:
            dim = self.dim
            #Initiate V, the tensor operator
            self.params['V'] = np.random.uniform(-r, r, size=(dim+1, 2*dim+2, 2*dim+2))
            self.params['V'] = (self.V+np.transpose(self.V, axes=[0, 2, 1]))/2

            #Initiate W, the linear operator
            self.params['W'] = np.random.uniform(-r, r, size=(dim+1, 2*dim+2))

            #Initiate Ws, the linear operator
            self.params['Ws'] = np.random.uniform(-r, r, size=(2, dim+1))

            #Initiate VE, the encoder tensor
            self.Ve = np.random.uniform(-r, r, size=(2*dim+2, dim+1, dim+1))
            self.Ve = (self.Ve+np.transpose(self.Ve, axes=[0, 2, 1]))/2

            #Initiate VE, the encoder tensor
            self.We = np.random.uniform(-r, r, size=(2*dim+2, dim+1))

            #Initiate L, the Lexicon representation
            self.params['L'] = np.random.uniform(-r, r, size=(len(self.vocab), dim))

        #Liste pour erreurs
        self.errMB = []
        self.errVal = []
        errMB = self.errMB
        errVal = self.errVal

        #Condition d'arret
        n_iter = 1
        gradNorm = 1.0
        if val_set != []:
            prevError = self.error(val_set) / len(val_set)
            for k in self.params.keys():
                prevError += self.reg[k]*self.norm(self.params[k])/2.0
            iniError = prevError
            minError = prevError  # optimal error so far
            glError = 0
            upStop = 0
            errVal.append(prevError)

        # Normalisation pour AdaGrad/Rms-prop
        eta = learning_rate
        dHist = {}
        dMask = {}
        dPrev = {}
        for k in self.params.keys():
            sh = self.params[k].shape
            dHist[k] = np.zeros(sh)
            dMask[k] = np.ones(sh)
            dPrev[k] = np.zeros(sh)

        early_stop = False
        while (not early_stop) and n_iter < max_iter:  # Critere moins random
            if n_iter % reset_freq == 0 and reset_freq > 0:  # Remise a zero des rates cf Socher
                for k in self.params.keys():
                    dHist[k] = np.zeros(self.params[k].shape)

            #Choose mini batch randomly
            mini_batch_samples = np.random.choice(X_trees, size=mini_batch_size)

            #Initialize gradients to 0
            dCurrent = {}
            for k in self.params.keys():
                    dCurrent[k] = np.zeros(self.params[k].shape)

            #Mini batch pour gradient
            currentMbe = 0.0
            for X_tree in mini_batch_samples:
                currentMbe += self.forward_pass(X_tree)
                grad = self.backward_pass(X_tree, w_root=w_root)
                for k in self.params.keys():
                    dCurrent[k] += grad[k]

            currentMbe /= mini_batch_size
            for k in self.params.keys():
                currentMbe += self.reg[k]*self.norm(self.params[k])/2.0

            #Division par le nombre de sample + regularisation
            for k in self.params.keys():
                dCurrent[k] = dCurrent[k]/mini_batch_size + self.reg[k]*self.params[k]

            #Mise a jour des poids et calcul des pas
            if strat == 'AdaGrad':
                eps = 0.001  # Adagrad >0 ? cf Socher
                for k in self.params.keys():
                    dHist[k] += dCurrent[k]*dCurrent[k]
                    dCurrent[k] = eta*dCurrent[k]/np.sqrt(dHist[k]+eps)
            else:
                eps = 0.001
                for k in self.params.keys():
                    dHist[k] = 0.9*dHist[k] + 0.1*dCurrent[k]*dCurrent[k]
                    dMask[k] *= .7*(dPrev[k]*dCurrent[k] >= 0) + .5
                    dCurrent[k] = eta*dMask[k].clip(1e-6, 50, out=dMask[k]) *\
                        dCurrent[k]/np.sqrt(dHist[k]+eps)

            #Calcul de la norme du gradient (critere d'arret)
            gradNorm = 0
            for k in self.params.keys():
                gradNorm += np.sum(np.abs(dCurrent[k]))

            #Keep previous gradient
            for k in self.params.keys():
                dPrev[k] = dCurrent[k]

            #Descente
            for k in self.params.keys():
                self.params[k] -= dCurrent[k]

            #Maj de la condition d'arret
            if val_set != [] and (n_iter % n_check) == 0:
                currentError = self.error(val_set)
                for k in self.params.keys():
                    currentError += self.reg[k]*self.norm(self.params[k])/2.0

                errVal.append(currentError)
                errMB.append(currentMbe)
                print('Error on validation set at iter {0} : {1} '
                      '(previous : {2})'.format(n_iter, currentError, prevError))
                print('Error on mini batch at iter {0} : {1} '
                      '(Gradient norm : {2})'.format(n_iter, currentMbe, gradNorm))

                with open(save_tmp, 'wb') as output:
                    pickle.dump(errVal, output, -1)
                    pickle.dump(errMB, output, -1)

                #Early stopping
                minError = min(minError, currentError)
                glError = 100*((currentError/minError)-1.0)
                if currentError > prevError:
                    upStop += 1
                else:
                    upStop = 0
                early_stop = (upStop >= n_stop) and (n_stop > 0)  # UP criterion
                prevError = currentError
            else:
                print('Error on mini batch at iter {0} : {1} '
                      '(Gradient norm : {2})'.format(n_iter, currentMbe, gradNorm))
                errMB.append(currentMbe)

            #Maj iter
            n_iter += 1

        if val_set != []:
            print('Error on training set before and after training'
                  '({2} iter) : {0}->{1}\n'.format(iniError, currentError, n_iter))
            print('Generalization error : {0}'.format(glError))
        return errMB, errVal

    def score_fine(self, X_trees):
        '''
        Score sur les predictions MAP avec 5 label
        '''
        countAll = 0
        countRoot = 0
        scAll = 0.0
        scRoot = 0.0
        for X_tree in X_trees:
            self.forward_pass(X_tree)
            for n in X_tree.nodes:
                countAll += 1
                scAll += (Tree.Tree.getSoftLabel(n.ypred[1]) == Tree.Tree.getSoftLabel(n.y[1]))
            countRoot += 1
            n = X_tree.nodes[-1]
            scRoot += (Tree.Tree.getSoftLabel(n.ypred[1]) == Tree.Tree.getSoftLabel(n.y[1]))
        return scAll/countAll, scRoot/countRoot

    def score_eps(self, X_trees, eps):
        '''
        Score sur les predictions MAP avec 5 label
        '''
        countAll = 0
        countRoot = 0
        scAll = 0.0
        scRoot = 0.0
        for X_tree in X_trees:
            self.forward_pass(X_tree)
            for n in X_tree.nodes:
                countAll += 1
                scAll += (abs(n.ypred[1] - n.y[1]) <= eps)
            countRoot += 1
            n = X_tree.nodes[-1]
            scRoot += (abs(n.ypred[1] - n.y[1]) <= eps)
        return scAll/countAll, scRoot/countRoot

    def score_binary(self, X_trees, inc_neut=False):
        '''
        Score sur les prediction MAP pos/neg
        '''
        countAll = 0
        countRoot = 0
        scAll = 0.0
        scRoot = 0.0
        for X_tree in X_trees:
            self.forward_pass(X_tree)
            for n in X_tree.nodes:
                countAll += 1 * (inc_neut or not (0.4 < n.y[1] <= 0.6))
                scAll += ((n.ypred[1] <= 0.5 and n.y[1] <= 0.5) or
                          (n.ypred[1] > 0.5 and n.y[1] > 0.5)) *\
                         (inc_neut or not (0.4 < n.y[1] <= 0.6))
            n = X_tree.nodes[-1]
            countRoot += 1 * (inc_neut or not (0.4 < n.y[1] <= 0.6))
            scRoot += ((n.ypred[1] <= 0.5 and n.y[1] <= 0.5) or
                      (n.ypred[1] > 0.5 and n.y[1] > 0.5)) *\
                      (inc_neut or not (0.4 < n.y[1] <= 0.6))
        return scAll/countAll, scRoot/countRoot

    def check_derivative(self, X_tree, eps=1e-6):
        '''
        Fait une compaaison dérivee / differences finies
        '''
        error1 = self.forward_pass(X_tree)
        dWs, dV, dW, dL = self.backward_pass(X_tree)
        dirW = np.random.uniform(size=self.W.shape)
        dirWs = np.random.uniform(size=self.Ws.shape)
        dirL = np.random.uniform(size=self.L.shape)
        dirV = np.random.uniform(size=self.V.shape)

        self.Ws += eps*dirWs
        self.W += eps*dirW
        self.L += eps*dirL
        self.V += eps*dirV

        error2 = self.forward_pass(X_tree)
        diff = (error2-error1)/eps
        diff2 = np.sum(dW*dirW)+np.sum(dWs*dirWs)+np.sum(dL*dirL)+np.sum(dV*dirV)

        return np.abs(diff-diff2)

    def confusion_matrix(self, X_trees):
        confAll = np.zeros((5, 5))
        confRoot = np.zeros((5, 5))
        for tree in X_trees:
            for n in tree.nodes:
                lp = Tree.Tree.getSoftLabel(n.ypred[1])
                l = Tree.Tree.getSoftLabel(n.y[1])
                confAll[l, lp] += 1
            lp = Tree.Tree.getSoftLabel(tree.nodes[-1].ypred[1])
            l = Tree.Tree.getSoftLabel(tree.nodes[-1].y[1])
            confRoot[l, lp] += 1
        for lp in range(5):
            confAll[lp, :] /= np.sum(confAll[lp, :])
            confRoot[lp, :] /= np.sum(confRoot[lp, :])
        return confAll, confRoot

    def plot_words_2D(self, labels, N=-1):
        from sklearn.decomposition import PCA
        import matplotlib.pyplot as plt
        import matplotlib.cm as cm

        if N == -1:
            N = self.L.shape[0]

        pca = PCA(n_components=2)
        X = pca.fit_transform(self.L[:N])
        revert_vocab = {i: (w, labels[w]) for (w, i) in self.vocab.iteritems()}
        y = np.zeros(N)
        for i in range(N):
            _, y[i] = revert_vocab[i]
        plt.scatter(X[:, 0], X[:, 1], c=y, cmap=cm.jet)

    def sample_w(self, p):
        dist = ((self.L-p)**2).sum(axis=1)
        i0 = dist.argmin()
        return self.index[i0]

    def sample(self, p):
        queue = [p]
        while len(queue) != 0:
            curr = queue.pop()
            self.dec(curr)
            C = self.dec(curr)
            ac = C[:self.dim+1:]
            bc = C[self.dim+1:]
            if ac[-1] > 0.:
                queue.append(ac)
            else:
                print(self.sample_w(ac[:-1]))
            if bc[-1] > 0.:
                queue.append(bc)
            else:
                print(self.sample_w(bc[:-1]))

    def plot_eps_curve(self, X, n_pts):
        import matplotlib.pyplot as plt

        eps = np.logspace(-3, 0, n_pts)
        curve_r = []
        curve_a = []
        for e in eps:
            a, r = self.score_eps(X, e)
            curve_r.append(r)
            curve_a.append(a)

        plt.semilogx(eps, curve_a)
        plt.semilogx(eps, curve_r)
        plt.show()

        return curve_a, curve_r


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

import numpy as np

from . import preliminary_computations as precomp
from .local_conditions import sky_countrate
from . import optics as opt
from . import photometry as phot
from . import utils

def etc_computation(info_dict):
    """
    Compute either the SNR, the total exposure time or
    the magnitude in function of the 2 others

    Parameters
    ----------
    info_dict: dictionary
        contains all relevant information

    wavelength : array
        wavelengths in angstrom

    Returns
    ---------
    SNR: float
        Signal to noise ratio

    mag: float
        magnitude reached

    tot_exp_time: float
        total exposure time in seconds

    """
    # display result
    verbose = info_dict["verbose"]

    etc_type = info_dict["etc_type"]
    SNR = info_dict["SNR"]
    tot_exp_time = info_dict["Nexp"] * info_dict["exptime"]
    info_dict["total_exposure_time"] = tot_exp_time
    Nexp = info_dict["Nexp"]
    # Detector Integration Time in seconds
    if etc_type == "snr" or etc_type == "mag":
        DIT = info_dict["exptime"] - info_dict["T_dithering"]

    # Display
    if verbose:
        print("\nInformation about Passband:")
        print("----------------------------")
        print("Cut_on: %.f angstroms" % info_dict["Passband_cuton"])
        print("Effective wavelength: %.f angstroms" %
              info_dict["effWavelength"])
        print("Cut_off: %.f angstroms" % info_dict["Passband_cutoff"])

        print("\nAirmass: %.2f" % info_dict["airmass"])
        print("\nSeeing: %.2f" % info_dict["seeing_los_arcsec"])

    if info_dict["detailed_trans"] == 1:
        # Computes mean transmission of each components for the given passband
        mean_trans_tel = utils.mean_efficiency_passband(
            info_dict, opt.telescope_efficiency(info_dict)
        )
        mean_trans_inst = utils.mean_efficiency_passband(
            info_dict,
            opt.instrument_channel_efficiency(info_dict)
            * info_dict["Trans_filter"]
            * info_dict["camera_efficiency"],
        )
        mean_trans_filter = utils.mean_efficiency_passband(
            info_dict, info_dict["Trans_filter"]
        )
        mean_eta_cam = utils.mean_efficiency_passband(
            info_dict, info_dict["camera_efficiency"]
        )
        mean_trans_optics = mean_trans_tel * mean_trans_inst
        mean_trans_atm = utils.mean_efficiency_passband(
            info_dict, info_dict["Trans_atmosphere"]
        )
        mean_trans_system = mean_trans_optics
        info_dict["trans_mean_tel"] = mean_trans_tel
        info_dict["trans_mean_inst"] = mean_trans_inst
        info_dict["trans_mean_optics"] = mean_trans_optics
        info_dict["trans_mean_filter"] = mean_trans_filter
        info_dict["trans_mean_atm"] = mean_trans_atm
        info_dict["trans_mean_cam"] = mean_eta_cam
        info_dict["trans_mean_system"] = mean_trans_system

        if verbose:
            print("\nMEAN EFFICENCIES:")
            print("------------------")
            print("Obscuration: %.3f" % (1.0 - info_dict["obstruction"]))
            print(
                "Telescope: %.3f (+obs: %.3f)"
                % (mean_trans_tel,
                   mean_trans_tel * (1.0 - info_dict["obstruction"]))
            )
            print("Instrument: %.3f" % mean_trans_inst)
            print(
                "Optics (tel+inst): %.3f  (+obs: %.3f)"
                % (
                    mean_trans_optics,
                    mean_trans_optics * (1.0 - info_dict["obstruction"]),
                )
            )
            print("Filter: %.3f" % mean_trans_filter)
            print("Atmosphere: %.3f" % mean_trans_atm)
            print("Camera: %.3f" % mean_eta_cam)
            print(
                "System: %.3f (+obs: %.3f)\n"
                % (
                    mean_trans_system,
                    mean_trans_system * (1 - info_dict["obstruction"]),
                )
            )

    elif info_dict["detailed_trans"] == 0:
        mean_eta_cam = utils.mean_efficiency_passband(
            info_dict, info_dict["camera_efficiency"]
        )
        mean_eta_optics = utils.mean_efficiency_passband(
            info_dict, phot.set_filter(info_dict)
        )
        mean_trans_system = mean_eta_cam * mean_eta_optics
        info_dict["trans_mean_system"] = mean_trans_system
        if verbose:
            print("\nMEAN EFFICENCIES:")
            print("------------------")
            print("Obscuration: %.3f" % (1.0 - info_dict["obstruction"]))
            print(
                "System: %.2f (+obs: %.3f)\n"
                % (
                    mean_trans_system,
                    mean_trans_system * (1 - info_dict["obstruction"]),
                )
            )

    # Number of pixels covering 1.35*FWHM of the PSF
    npix = info_dict["npix"]
    # Factor when estimating the Noise from other images
    factor_ima = precomp.factor_images_averaged(info_dict)
    # Fraction of light in the brightest pixel
    f_pix = precomp.Normalisation_factor(info_dict, True)
    # Fraction of light in the PSF
    f_PSF = precomp.Normalisation_factor(info_dict, False)

    info_dict["factor_ima"] = factor_ima
    info_dict["f_pix"] = f_pix
    info_dict["f_PSF"] = f_PSF

    # Background Noise countrate in e-/s/px
    # ---------------------------------------
    info_dict = sky_countrate(info_dict)  # e-/s/px
    BN = info_dict["Sky_CountRate"]
    # print ('Sky countrate: %.2f (e-/px/s)' % BN)
    # Thermic signal (electrons/s/pixel)  <--> Dark current
    # ------------------------------------------------------
    DC = info_dict["cameras"][info_dict["channel"]]["DC"]
    # Digitization noise   (e-/pixel)
    # ----------------------------------
    # Converter analog to digital noise of 1/2 ADU (electrons/pixel)
    DigN = info_dict["dig_noise"]

    # Readout noise (e-/pixel)
    # --------------------------
    RN = info_dict["cameras"][info_dict["channel"]]["RN"]

    # Instrument background (e-/s/pix)
    inst_bg = info_dict["Instrument_bg"]

    # Object
    # ---------
    # Count rate of the object in e-/s
    if etc_type == "snr" or etc_type == "time":
        CR, fph = info_dict["Object_fes"], info_dict["Object_fph"]  # e-/s

    # Zeropoint
    # -----------
    ZP = info_dict["zeropoint"]

    # Add some info in info_cit
    if verbose:
        print("Zeropoint: %.2f (%s mag)" % (ZP,
                                            info_dict["photometry_system"]))

    # Compute the SNR
    # ------------------
    if etc_type == "snr":
        # In the case of a given magnitude to reach, the fraction of flux
        # we kept should be included so that the given mag corresponds to the
        # measured count rate.
        if info_dict["object_type"] == "magnitude":
            CR = CR  # / f_PSF
            CR_pix = CR  # / f_pix
        else:
            CR_pix = CR

        # Peak SNR (Object signal/noise at the brightest pixel)
        SNR_pix = (
            np.sqrt(Nexp)
            * CR_pix
            * f_pix
            * DIT
            / np.sqrt(
                CR_pix * f_pix * DIT
                + factor_ima * ((RN ** 2.0 + DigN ** 2.0)
                                + DIT * (DC + BN + inst_bg))
            )
        )

        # Total integrated noise over npix (electrons/area)
        SNR = (
            np.sqrt(Nexp)
            * CR
            * f_PSF
            * DIT
            / np.sqrt(
                CR * f_PSF * DIT
                + factor_ima
                * npix
                * ((RN ** 2.0 + DigN ** 2.0) + DIT * (DC + BN + inst_bg))
            )
        )

        if info_dict["object_type"] == "magnitude":
            # mag = ZP - 2.5*np.log10(CR*f_PSF)
            mag = info_dict["object_magnitude"]
            mag_pix = ZP - 2.5 * np.log10(CR_pix * f_pix)
            mag_pix = info_dict["object_magnitude"]
        else:
            mag = ZP - 2.5 * np.log10(CR)
            mag_pix = ZP - 2.5 * np.log10(CR_pix)

        Ftot_el = CR * f_PSF * DIT  # *np.sqrt(Nexp)
        Ftot_el_pix = CR_pix * f_pix * DIT  # *np.sqrt(Nexp)
        DIT_pix = DIT

        info_dict["SNR"] = SNR
        info_dict["SNR_pix"] = SNR_pix
        info_dict["mag_pix"] = mag_pix
        info_dict["Ftot_el_pix"] = Ftot_el_pix
        info_dict["Ftot_el"] = Ftot_el
        info_dict["DIT_pix"] = DIT_pix

        if verbose:
            print(
                "\n\nA magnitude (%s system) of %.2f in %s band within a total exposure time of %.2f seconds splited in %d exposure(s), implies a total SNR of :\n"
                % (
                    info_dict["photometry_system"],
                    mag,
                    info_dict["filter_band"],
                    DIT * Nexp,
                    Nexp,
                )
            )
            # print ('\t - Peak SNR at the brightest pixel: %.2f \n' % SNR_pix)
            print("\t - Integrated SNR over %d pixels: %.2f" % (npix, SNR))
            print(
                "\n\nA magnitude (%s system) of %.2f in %s band within a total exposure time of %.2f seconds splited in %d exposure(s), implies a SNR for the central pixel of of :\n\n"
                % (
                    info_dict["photometry_system"],
                    mag_pix,
                    info_dict["filter_band"],
                    DIT_pix * Nexp,
                    Nexp,
                )
            )
            print("\t - SNR of the central pixel: %.2f \n\n" % SNR_pix)

    # Compute the total exposure time
    # --------------------------------
    elif etc_type == "time":

        # In the case of a given magnitude to reach, the fraction of flux
        # we kept should be included so that the given mag corresponds to the
        # measured count rate.
        if info_dict["object_type"] == "magnitude":
            CR = CR  # / f_PSF

        # Integrated over Npixels     (solve 2nd degree equation)
        if Nexp > 1:
            SNR_1 = SNR / np.sqrt(Nexp)
        else:
            SNR_1 = SNR
        A_sys = -((CR * f_PSF) ** 2.0)
        B_sys = SNR_1 ** 2.0 * (CR * f_PSF + factor_ima * npix
                                * (DC + BN + inst_bg))
        C_sys = SNR_1 ** 2.0 * factor_ima * npix * (RN ** 2.0 + DigN ** 2.0)

        delta = B_sys * B_sys - 4.0 * A_sys * C_sys

        DIT = (-B_sys - np.sqrt(delta)) / (2.0 * A_sys)
        # mag = -2.5*np.log10(obj_janskys/3631)

        if info_dict["object_type"] == "magnitude":
            mag = ZP - 2.5 * np.log10(CR)  # *f_PSF)
        else:
            mag = ZP - 2.5 * np.log10(CR)

        Ftot_el = CR * f_PSF * DIT  # *np.sqrt(Nexp)

        # Brightest pixel
        if info_dict["object_type"] == "magnitude":
            CR_pix = CR  # / f_pix
        else:
            CR_pix = CR

        A_sys = -((CR_pix * f_pix) ** 2.0)
        B_sys = SNR_1 ** 2.0 * (CR_pix * f_pix + factor_ima
                                * (DC + BN + inst_bg))
        C_sys = SNR_1 ** 2.0 * factor_ima * (RN ** 2.0 + DigN ** 2.0)

        delta = B_sys * B_sys - 4.0 * A_sys * C_sys

        DIT_pix = (-B_sys - np.sqrt(delta)) / (2.0 * A_sys)

        if info_dict["object_type"] == "magnitude":
            mag_pix = ZP - 2.5 * np.log10(CR_pix)  # *f_pix)
        else:
            mag_pix = ZP - 2.5 * np.log10(CR)

        Ftot_el_pix = CR_pix * f_pix * DIT_pix  # *np.sqrt(Nexp)

        info_dict["DIT"] = DIT
        info_dict["DIT_pix"] = DIT_pix
        info_dict["Ftot_el_pix"] = Ftot_el_pix
        info_dict["Ftot_el"] = Ftot_el
        info_dict["mag_pix"] = mag_pix

        if verbose:
            print(
                "\n\nReaching a magnitude (%s system) of %.2f in %s band with a SNR of %.2f requires:\n"
                % (info_dict["photometry_system"], mag, info_dict["filter_band"], SNR)
            )
            print("\t - a Total exposure time of: %.2f s\n" % (DIT * Nexp))
            print(
                "\n\nReaching a magnitude (%s system) of %.2f in %s band with a SNR of %.2f for the central pixel requires:\n\n"
                % (
                    info_dict["photometry_system"],
                    mag_pix,
                    info_dict["filter_band"],
                    SNR,
                )
            )
            print("\t - a Total exposure time of: %.2f s\n\n" % (DIT_pix * Nexp))

    # Compute the magnitude
    # ------------------------
    elif etc_type == "mag":
        # f_PSF=1
        # f_pix=1
        # Integrated over Npixels    (solve 2nd degree equation)
        A_sys = -((f_PSF * DIT * np.sqrt(Nexp)) ** 2.0)
        B_sys = SNR ** 2.0 * f_PSF * DIT
        C_sys = SNR ** 2.0 * (
            factor_ima * npix * (RN ** 2.0 + DigN ** 2.0 + DIT
                                 * (DC + BN + inst_bg))
        )

        delta = B_sys * B_sys - 4.0 * A_sys * C_sys

        CR = (-B_sys - np.sqrt(delta)) / (2.0 * A_sys)
        Ftot_el = CR * f_PSF * DIT  # *np.sqrt(Nexp)

        mag = ZP - 2.5 * np.log10(CR)

        # Central pixel
        A_sys = -((f_pix * DIT * np.sqrt(Nexp)) ** 2.0)
        B_sys = SNR ** 2.0 * f_pix * DIT
        C_sys = SNR ** 2.0 * (
            factor_ima * (RN ** 2.0 + DigN ** 2.0 + DIT * (DC + BN + inst_bg))
        )

        delta = B_sys * B_sys - 4.0 * A_sys * C_sys

        CR_pix = (-B_sys - np.sqrt(delta)) / (2.0 * A_sys)
        Ftot_el_pix = CR_pix * f_pix * DIT  # *np.sqrt(Nexp)
        mag_pix = ZP - 2.5 * np.log10(CR_pix)
        DIT_pix = DIT

        object_mag = mag * np.ones(len(info_dict["wavelength_ang"]))
        fJy = phot.mag2Jy(info_dict, object_mag)  # Jy
        # erg/s/cm2/A
        flam = utils.fJy_to_flambda(info_dict["wavelength_ang"], fJy)
        # ph/s/cm2/A
        fph = utils.flambda_to_fph(info_dict["wavelength_ang"], flam)

        info_dict["Object_mag"] = object_mag
        info_dict["object_magnitude"] = mag
        info_dict["mag_pix"] = mag_pix
        info_dict["DIT_pix"] = DIT
        info_dict["Ftot_el_pix"] = Ftot_el_pix
        info_dict["Ftot_el"] = Ftot_el

        if verbose:
            print(
                "\n\nFor a total SNR=%.2f in a total exposure time of %.2f (sec) in %d exposure(s) we reach:\n"
                % (SNR, DIT * Nexp, Nexp)
            )
            print(
                "\t - a magnitude (%s system) of: %.2f in %s band\n"
                % (info_dict["photometry_system"],
                   mag,
                   info_dict["filter_band"])
            )
            print(
                "\n\nFor the central pixel a SNR=%.2f in a total exposure time of %.2f (sec) in %d exposure(s) we reach:\n\n"
                % (SNR, DIT_pix * Nexp, Nexp)
            )
            print(
                "\t - a magnitude (%s system) of: %.2f in %s band\n\n"
                % (info_dict["photometry_system"], mag_pix, info_dict["filter_band"])
            )

    info_dict["DIT"] = DIT

    # sigma_shot_noise = np.sqrt(Ftot_el * DIT)
    # sigma_dark_current = np.sqrt(DC * npix * DIT)
    # sigma_sky = np.sqrt(BN * npix * DIT)
    # sigma_digitization = np.sqrt(npix * DigN ** 2.0)
    # sigma_readout_noise = np.sqrt(npix * RN ** 2.0)

    # Total number of electrons in the brightest pixel for 1 exposure
    N_el_tot_pix1 = Ftot_el_pix + (BN + DC + inst_bg) * DIT_pix + RN + DigN
    N_el_tot_pix2 = (Ftot_el * f_pix / f_PSF
                     + (BN + DC + inst_bg) * DIT + RN + DigN)

    info_dict["N_el_tot_pix1"] = N_el_tot_pix1
    info_dict["N_el_tot_pix2"] = N_el_tot_pix2

    if N_el_tot_pix1 > info_dict["cameras"][info_dict["channel"]]["FWC"]:
        info_dict["saturation"] = "Yes"
    else:
        info_dict["saturation"] = "No"

    if verbose:

        print(
            "\nFull well capacity of 1 pixel: %.2f (electrons)"
            % (info_dict["cameras"][info_dict["channel"]]["FWC"])
        )
        print("\n\n--------- One pixel only------------------")
        print(
            "\nPhoto-electrons created: central pix for %d exposure(s) of %.2f sec "
            % (Nexp, DIT_pix)
        )
        print("\tby:")
        print("\t- Object:         %10.2f   (electrons)" % Ftot_el_pix)
        print("\t- Sky:            %10.2f   (electrons)" % (BN * DIT_pix))
        print("\t- Readout:        %10.2f   (electrons)" % RN)
        print("\t- Dark current:   %10.2f   (electrons)" % (DC * DIT_pix))
        print("\t- Digitization:   %10.2f   (electrons)" % DigN)
        print("\t- Instrument bg:  %10.2f   (electrons)" % (inst_bg * DIT_pix))

        print(
            "\nSNR: -central pixel: %.2f"
            % (
                np.sqrt(Nexp)
                * Ftot_el_pix
                / np.sqrt(
                    Ftot_el_pix
                    + factor_ima
                    * ((RN ** 2.0 + DigN ** 2.0) + DIT_pix
                       * (DC + BN + inst_bg))
                )
            )
        )

        print(
            "\nTotal of electrons collected in the central pixel during an exposure time of %d seconds: %.2f "
            % (DIT_pix, N_el_tot_pix1)
        )
        if N_el_tot_pix1 > info_dict["cameras"][info_dict["channel"]]["FWC"]:
            print(
                "--> Central pixel saturated: number of electrons > Full well Capacity"
            )
        elif N_el_tot_pix1 > info_dict["cameras"][info_dict["channel"]]["gain"] * (
            2.0 ** (info_dict["cameras"][info_dict["channel"]]["bits"]) - 1
        ):
            print(
                "--> Central pixel saturated: number of electrons > number of digitizations"
            )
        elif (
            N_el_tot_pix1 > 1.0 / 2 * info_dict["cameras"][info_dict["channel"]]["FWC"]
        ):
            print(
                "--> Number of electrons in central pixel > 1/2 of Full well Capacity. Risk of non-linear response."
            )

        else:
            print("--> No saturation")

        print("\n\n\n--------- Integrated over %d pixels------------------" % npix)
        print(
            "\nPhoto-electrons created: brightest pix |  total of %d pixels, %d exposure(s) of %.2f sec "
            % (npix, Nexp, DIT)
        )
        print("\tby:")
        print(
            "\t- Object:         %10.2f   |   %10.2f   (electrons)"
            % (Ftot_el * f_pix / f_PSF, Ftot_el)
        )
        print(
            "\t- Sky:            %10.2f   |   %10.2f   (electrons)"
            % (BN * DIT, (BN * npix * DIT * Nexp))
        )
        print(
            "\t- Readout:        %10.2f   |   %10.2f   (electrons)"
            % (RN, (RN * npix * Nexp))
        )
        print(
            "\t- Dark current:   %10.2f   |   %10.2f   (electrons)"
            % (DC * DIT, (DC * DIT * npix * Nexp))
        )
        print(
            "\t- Digitization:   %10.2f   |   %10.2f   (electrons)"
            % (DigN, (DigN * npix * Nexp))
        )
        print(
            "\t- Instrument bg:  %10.2f   |   %10.2f   (electrons)"
            % (inst_bg * DIT, (inst_bg * DIT * npix * Nexp))
        )

        print(
            "\nSNR: -Brightest pixel: %.2f"
            % (
                np.sqrt(Nexp)
                * Ftot_el
                * f_pix
                / f_PSF
                / np.sqrt(
                    Ftot_el * f_pix / f_PSF
                    + factor_ima
                    * ((RN ** 2.0 + DigN ** 2.0) + DIT * (DC + BN + inst_bg))
                )
            )
        )
        print(
            "     -integrated over %d pixels: %.2f"
            % (
                npix,
                np.sqrt(Nexp)
                * Ftot_el
                / np.sqrt(
                    Ftot_el
                    + factor_ima
                    * npix
                    * ((RN ** 2.0 + DigN ** 2.0) + DIT * (DC + BN + inst_bg))
                ),
            )
        )
        print(
            "\nTotal of electrons collected in the brightest pixel during an exposure time of %d seconds: %.2f "
            % (DIT_pix, N_el_tot_pix2)
        )
        if N_el_tot_pix2 > info_dict["cameras"][info_dict["channel"]]["FWC"]:
            print(
                "--> Brightest pixel saturated: number of electrons > Full well Capacity"
            )
        elif N_el_tot_pix2 > info_dict["cameras"][info_dict["channel"]]["gain"] * (
            2.0 ** (info_dict["cameras"][info_dict["channel"]]["bits"]) - 1
        ):
            print(
                "--> Brightest pixel saturated: number of electrons > number of digitizations"
            )
        elif (
            N_el_tot_pix2 > 1.0 / 2 * info_dict["cameras"][info_dict["channel"]]["FWC"]
        ):
            print(
                "--> Number of electrons in brightest pixel > 1/2 of Full well Capacity. Risk of non-linear response."
            )

        else:
            print("--> No saturation")

        print(
            "\nDead time: %.2f sec \n(%.2f sec for dithering, the %.2f sec for the readout are not taken into account)"
            % (
                info_dict["deadtime_tot"],
                info_dict["T_dithering"],
                info_dict["cameras"][info_dict["channel"]]["ReadoutTime"],
            )
        )
    info_dict["SNR"] = SNR
    info_dict["mag"] = mag
    info_dict["total_exposure_time"] = DIT * Nexp
    info_dict["fph"] = fph

    return info_dict


# This file is part of sequencing.
#
#    Copyright (c) 2021, The Sequencing Authors.
#    All rights reserved.
#
#    This source code is licensed under the BSD-style license found in the
#    LICENSE file in the root directory of this source tree.

import re
import json
from collections import defaultdict
from contextlib import contextmanager

import numpy as np
import qutip
import attr

from .parameters import Parameterized, ListParameter, DictParameter
from .modes import Mode, sort_modes

class CouplingTerm(object):
    """An object representing a coupling between two ``Modes``,
    given by a Hamiltonian term of the form ``strength * op1 * op2``.
    If the keyword argument ``add_hc`` is provided and is True,
    then the Hamiltonian term takes the form
    ``strength * ((op1 * op2) + (op1 * op2).dag())``.

    Args:
        mode1 (Mode): First mode to be coupled.
        op1_expr (str): String representation of mode1's operator in the coupling term.
        mode2 (Mode): Second mode to tbe coupled (can the same object as mode1).
        op2_expr (str): String representation of mode2's operator in the coupling term.
        strength (optional, float): Coefficient parameterizing the
            strength of the coupling. Strength should be given in
            units of 2 * pi * GHz. Default: 1.
        add_hc (optional, bool): Whether to add the Hermitian conjugate
            of the product of op1 and op2. Default: False.
    """

    def __init__(self, mode1, op1_expr, mode2, op2_expr, strength=1, add_hc=False):
        for mode in [mode1, mode2]:
            if not isinstance(mode, Mode):
                raise TypeError(f"Expected instance of Mode, but got {type(mode)}.")
        for expr in [op1_expr, op2_expr]:
            if not isinstance(expr, str):
                raise TypeError(f"Expected instance of str, but got {type(expr)}.")
        self.mode1 = mode1
        self.mode2 = mode2
        self.op1_expr = op1_expr
        self.op2_expr = op2_expr
        self.strength = float(strength)
        self.add_hc = bool(add_hc)

    @property
    def op1(self):
        """Operator for ``mode1``."""
        op = self.mode1.operator_expr(self.op1_expr)
        if not isinstance(op, qutip.Qobj):
            raise TypeError(f"Expected op1 to be a qutip.Qobj, not {type(op)}")
        if not self.add_hc and not op.isherm:
            raise ValueError("Expected op1 to be Hermitian since add_hc is False.")
        return op

    @property
    def op2(self):
        """Operator for ``mode2``."""
        op = self.mode2.operator_expr(self.op2_expr)
        if not isinstance(op, qutip.Qobj):
            raise TypeError(f"Expected op2 to be a qutip.Qobj, not {type(op)}")
        if not self.add_hc and not op.isherm:
            raise ValueError("Expected op2 to be Hermitian since add_hc is False.")
        return op

    def H(self, strength=None, add_hc=None):
        """Returns the operator representing the coupling term.

        Args:
            strength (optional, float): Coefficient parameterizing the
                strength of the coupling. Strength should be given in units
                of 2 * pi * GHz. Defaults to self.strength.
            add_hc (optional, bool): Whether to add the Hermitian conjugate
                of product of op1 and op2. Defaults to self.add_hc.

        Returns:
            ``qutip.Qobj``: Operator representing the coupling term.
        """
        if strength is None:
            strength = self.strength
        if add_hc is None:
            add_hc = self.add_hc
        op = self.op1 * self.op2
        if add_hc:
            op = op + op.dag()
        return strength * op

    def __repr__(self):
        return (
            f"{type(self).__name__}("
            f"{self.mode1.name}.{self.op1_expr}, "
            f"{self.mode2.name}.{self.op2_expr}, "
            f"strength={self.strength:.3e}, "
            f"add_hc={self.add_hc}"
            ")"
        )

@attr.s
class System(Parameterized):
    """A collection of ``Modes`` that can be coupled together.

    Attributes:
        modes (list[Mode]): List of all Modes in the system.
        coupling_terms (dict[frozenset[str], list[CouplingTerm]]):
            A dictionary of CouplingTerm objects specifying all
            interactions in the system.
        cross_kerrs (dict[frozenset[str], float]): A dictionary
            of cross-Kerr values in units of GHz.
    """

    modes = ListParameter()
    cross_kerrs = DictParameter()

    order_modes = True

    def initialize(self):
        super().initialize()
        if self.order_modes:
            self.modes = sort_modes(self.modes)
        self.active_modes = self.modes
        self.coupling_terms = defaultdict(list)

    def __getattribute__(self, name):
        # Access modes like system.qubit.
        # System.modes can be changed at any time, so we cannot
        # just setattr(self, name, mode) in initialize().
        try:
            for mode in object.__getattribute__(self, "modes"):
                if mode.name == name:
                    return mode
        except AttributeError:
            pass
        return object.__getattribute__(self, name)

    def get_mode(self, mode):
        """Fetch a mode by name.

        Args:
            mode (str or Mode): Name of Mode to fetch, or the Mode itself.

        Returns:
            Mode: The requested ``Mode``.
        """
        if isinstance(mode, Mode):
            mode = mode.name
        if mode not in [m.name for m in self.modes]:
            raise ValueError(f"{mode} is not a mode of {self.name}.")
        return getattr(self, mode)

    @property
    def levels(self):
        """Dictionary of (mode_name, number_of_levels)"""
        return {mode.name: mode.levels for mode in self.modes}

    @property
    def active_modes(self):
        """List of the modes currently being used."""
        for mode in self._active_modes:
            if mode not in self.modes:
                raise ValueError(
                    f"{mode.name} is not in {self.name}.modes. "
                    f"This likely happened because {self.name}.modes "
                    f"was changed after {self.name} was created."
                    f"Please set {self.name}.active_modes to be "
                    f"a subset of {self.name}.modes."
                )
        return self._active_modes

    @active_modes.setter
    def active_modes(self, modes):
        if isinstance(modes[0], str):
            modes = [getattr(self, m) for m in modes]
        if self.order_modes:
            modes = sort_modes(modes)
        if not all(mode in self.modes for mode in modes):
            raise ValueError("active_modes must be a subset of the system's modes.")
        self._active_modes = modes
        for mode in self._active_modes:
            mode.space = self._active_modes

    @contextmanager
    def use_modes(self, modes):
        """A context manager that temporarily sets ``self.active_modes``
        to ``modes``, then reverts ``self.active_modes``
        to its previous value.

        Args:
            modes (list[Mode]): List of ``Modes`` to temporarily
                assign to ``self.active_modes``.
        """
        if isinstance(modes, (str, Mode)):
            modes = [modes]
        old_modes = self.active_modes
        try:
            self.active_modes = modes
            yield
        finally:
            self.active_modes = old_modes

    @staticmethod
    def tensor(*args):
        """Calculates the tensor product of input operators."""
        return qutip.tensor(*args)

    def I(self, modes=None):  # noqa: E741, E743
        """Identity operator.

        Args:
            modes (optional, list[Mode]): List of Modes to use in
                constructing H0. If None, will use self.active_modes.
                Default: None.

        Returns:
            ``qutip.Qobj``: Identity operator on the Hilbert space
            defined by ``modes``.
        """
        if modes is None:
            modes = self.active_modes
        elif not all(mode in self.modes for mode in modes):
            raise ValueError("modes must be a subset of the system's modes.")
        return self.tensor(*[qutip.qeye(mode.levels) for mode in modes])

    eye = I

    def fock(self, *args, **kwargs):
        """Returns a product state in the Fock basis. States can be
        specified either positionally or as keyword arguments.

        Args:
            *args (tuple): Fock states of Modes in the order of self.modes.
            **kwargs (dict): Fock states of Modes specified as keyword
                arguments, mode_name=n.

        Returns:
            ``qutip.Qobj``: The requested product state.
        """
        if args:
            if kwargs:
                raise ValueError(
                    "If positional arguments are provided, "
                    "no keyword arguments are allowed."
                )
            if len(args) != len(self.active_modes):
                raise ValueError(
                    "The number of positional argument must match "
                    "the number of active modes."
                )
            states = [
                qutip.fock(mode.levels, val)
                for mode, val in zip(self.active_modes, args)
            ]
        else:
            states = [
                qutip.fock(mode.levels, kwargs.get(mode.name, 0))
                for mode in self.active_modes
            ]
        return self.tensor(*states)

    def fock_dm(self, *args, **kwargs):
        """Returns a product state in the Fock basis, as a density matrix.

        States can be specified either positionally or as keyword arguments.

        Args:
            *args (tuple): Fock states of Modes in the order of self.modes.
            **kwargs (dict): Fock states of Modes specified as keyword
                arguments, mode_name=n.

        Returns:
            ``qutip.Qobj``: The requested product state, as a density matrix.
        """
        ket = self.fock(*args, **kwargs)
        return qutip.ket2dm(ket)

    basis = fock

    def ground_state(self):
        """Returns the ground state of the system.

        Returns:
            ``qutip.Qobj``: The system's ground state.
        """
        return self.fock()

    def logical_basis(self, *args, **kwargs):
        """Returns a product state in the basis spanned by the logical states
        of all modes. Logical states can be specified either positionally
        or as keyword arguments.

        Args:
            *args (tuple): Logical states of Modes in the order of self.modes.
            **kwargs (dict): Logical states of Modes specified as keyword
                arguments, mode_name=n.

        Returns:
            ``qutip.Qobj``: The requested product state.
        """
        if args:
            if kwargs:
                raise ValueError(
                    "If positional arguments are provided, "
                    "no keyword arguments are allowed."
                )
            if len(args) != len(self.active_modes):
                raise ValueError(
                    "The number of positional argument must match "
                    "the number of active modes."
                )
            states = [
                mode.logical_states(full_space=False)[val]
                for mode, val in zip(self.active_modes, args)
            ]
        else:
            states = [
                mode.logical_states(full_space=False)[kwargs.get(mode.name, 0)]
                for mode in self.active_modes
            ]
        return self.tensor(*states)

    def set_cross_kerr(self, mode1, mode2, chi=0):
        """Set the cross-Kerr (in GHz) between two modes. Note that the order
        of mode1 and mode2 doesn't matter.

        Args:
            mode1 (Mode or str): Instance of Mode or
                the name of a member of ``self.modes``.
            mode2 (Mode or str): Instance of Mode or
                the name of a member of ``self.modes``.
            chi (optional, float): Cross-Kerr between mode0 and mode1 in GHz.
                Default: 0.
        """
        if isinstance(mode1, str):
            mode1 = self.get_mode(mode1)
        if isinstance(mode2, str):
            mode2 = self.get_mode(mode2)
        if mode1 is mode2:
            raise ValueError("If mode1 is mode2, then it's not a cross-Kerr.")
        key = frozenset([mode1.name, mode2.name])
        # Replace this cross-Kerr if it already exists
        if key in self.coupling_terms:
            for i, term in enumerate(self.coupling_terms[key][:]):
                if (
                    term.mode1 is mode1
                    and term.op1_expr == "n"
                    and term.mode2 is mode2
                    and term.op2_expr == "n"
                ) or (
                    term.mode1 is mode2
                    and term.op1_expr == "n"
                    and term.mode2 is mode1
                    and term.op2_expr == "n"
                ):
                    _ = self.coupling_terms[key].pop(i)
        self.coupling_terms[key].append(
            CouplingTerm(mode1, "n", mode2, "n", strength=2 * np.pi * chi)
        )
        self.cross_kerrs[key] = chi

    def couplings(self, modes=None, clean=True):
        """Returns all of the static coupling terms in the Hamiltonian.

        Args:
            modes (optional, list[Mode]): List of Modes to use in
                constructing H0. If None, will use self.active_modes.
                Default: None.
            clean (optional, bool): Only keep operators with nonzero elements.
                Default: True.

        Returns:
            list[qutip.Qobj]: List of static coupling terms.
        """
        if modes is None:
            modes = self.active_modes
        mode_names = [mode.name for mode in modes]
        coupling_terms = []
        with self.use_modes(modes):
            for (mode0, mode1), terms in self.coupling_terms.items():
                if mode0 in mode_names and mode1 in mode_names:
                    coupling_terms.extend([term.H() for term in terms])
        if clean:
            return [term for term in coupling_terms if term.data.nnz]
        return coupling_terms

    def H0(self, modes=None, clean=True):
        """Returns the static Hamiltonian consisting of all
        self-Kerrs and cross-Kerrs.

        Args:
            modes (optional, list[Mode]): List of Modes to use in
                constructing H0. If None, will use self.active_modes.
                Default: None.
            clean (optional, bool): Only keep operators with nonzero elements.
                Default: True.

        Returns:
            list[qutip.Qobj]: Static Hamiltonian in list form.
        """
        if modes is None:
            modes = self.active_modes
        detunings = [mode.detuning for mode in modes]
        self_kerrs = [mode.self_kerr for mode in modes]
        couplings = self.couplings(modes=modes, clean=clean)
        H0 = detunings + self_kerrs + couplings
        if clean:
            return [H for H in H0 if H.data.nnz]
        return H0

    def c_ops(self, modes=None, clean=True):
        """Returns a list of collapse operators corresponding to
        loss (decay/excitation) and dephasing of all modes.

        Args:
            modes (optional, list[Mode]): List of Modes to use in
                constructing H0. If None, will use self.active_modes.
                Default: None.
            clean (optional, bool): Only keep operators with nonzero elements.
                Default: True.

        Returns:
            list[qutip.Qobj]: List of all collapse operators.
        """
        if modes is None:
            modes = self.active_modes
        decay = [mode.decay for mode in modes]
        excitation = [mode.excitation for mode in modes]
        dephasing = [mode.dephasing for mode in modes]
        c_ops = decay + excitation + dephasing
        if clean:
            return [c for c in c_ops if c.data.nnz]
        return c_ops

    def as_dict(self, json_friendly=False):
        """Overrides Parameterized.as_dict() in order to deal with cross_kerrs.

        Args:
            json_friendly (optional, bool): Whether to return
                a JSON-friendly dictionary. Default:True.

        Returns:
            dict: Dictionary representation of the System object.
        """
        d = super().as_dict(json_friendly=json_friendly)
        cross_kerrs = d.pop("cross_kerrs")
        d["cross_kerrs"] = {}
        if json_friendly:
            # turn frozenset({mode0, mode1}) into '{mode0, mode1}' for json
            for key, val in cross_kerrs.items():
                new_key = "{" + ", ".join(key) + "}"
                d["cross_kerrs"][new_key] = val
        else:
            for key, val in cross_kerrs.items():
                d["cross_kerrs"][key] = val
        return d

    @classmethod
    def from_json(cls, json_path=None, json_str=None):
        """Overrides Parameterized.from_json()
        in order to deal with cross_kerrs.

        Args:
            json_path (optional, str): Path to JSON file from which
                to load parameters. Required if ``json_str`` is ``None``.
                Default: None.
            json_str (optional, str): JSON string like that returned by
                ``self.to_json(dumps=True)``. Required if ``json_path``
                is ``None`` Default: None.

        Returns:
            System: Instance of ``System``
            whose parameters have been populated from the JSON data.
        """

        def json_decode(obj):
            # decode CrossKerr
            for key in list(obj):
                if re.match(r"\{(.*?)\}", key):
                    # turn '{mode0, mode1}' into frozenset({mode0, mode1})
                    k = key.replace("{", "").replace("}", "").split(", ")
                    new_key = frozenset(k)
                    obj[new_key] = obj[key]
                    del obj[key]
            return obj

        if json_str is not None:
            if json_path is not None:
                raise ValueError(
                    "You must provide either json_path " "or json_str, not both."
                )
            d = json.loads(json_str, object_hook=json_decode)
        else:
            if json_path is None:
                raise ValueError("You must provide either json_path " "or json_str.")
            if not json_path.endswith(".json"):
                json_path = json_path + ".json"
            with open(json_path, "r") as f:
                d = json.load(f, object_hook=json_decode)

        return cls.from_dict(d)


"""
Copyright 2020 The OneFlow Authors. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import math

import oneflow as flow
from oneflow.nn import init
from oneflow.nn.common_types import _size_2_t
from oneflow.nn.module import Module
from oneflow.nn.modules.utils import _pair

def slice(x, begin, size):
    ndim = len(x.shape)
    if not isinstance(begin, (list, tuple)) or len(begin) != ndim:
        raise ValueError(
            "begin must be a list/tuple with the same length as input tensor's number of dimensions"
        )
    if not all((isinstance(b, int) or b is None for b in begin)):
        raise ValueError("element of begin must be a int or None")
    if not isinstance(size, (list, tuple)) or len(size) != ndim:
        raise ValueError(
            "size must be a list/tuple with the same length as input tensor's number of dimensions."
        )
    if not all((isinstance(s, int) or s is None for s in size)):
        raise ValueError("element of size must be a int or None")
    slice_tup_list = []
    for (b, s, dim_size) in zip(begin, size, x.shape):
        (start, stop, step) = (None, None, 1)
        if b is not None:
            if b < -dim_size or b >= dim_size:
                raise ValueError("element of begin is out of range")
            start = b
        if s is not None:
            if s == -1:
                stop = dim_size
            else:
                if s <= 0 or s > dim_size:
                    raise ValueError("element of size is invalid")
                if b + s < dim_size:
                    stop = b + s
        slice_tup_list.append((start, stop, step))
    return flow.slice(x, slice_tup_list)

class ConvUtil(object):
    @classmethod
    def split(cls, x, axis, split_num):
        split_len = x.shape[axis] // split_num
        result_list = []
        slice_begin = [0] * len(x.shape)
        slice_size = [-1] * len(x.shape)
        slice_size[axis] = split_len
        for i in range(split_num):
            slice_begin[axis] = i * split_len
            result = slice(x, slice_begin, slice_size)
            result_list.append(result)
        return result_list

class ConvTranspose2d(Module):
    """

    Applies a 2D transposed convolution operator over an input image composed of several input planes.

    This module can be seen as the gradient of Conv2d with respect to its input.
    It is also known as a fractionally-strided convolution or
    a deconvolution (although it is not an actual deconvolution operation).

    Args:
        in_channels (int): Number of channels in the input image
        out_channels (int): Number of channels produced by the convolution
        kernel_size (int or tuple): Size of the convolving kernel
        stride (int or tuple, optional): Stride of the convolution. Default: 1
        padding (int or tuple, optional): ``dilation * (kernel_size - 1) - padding`` zero-padding
            will be added to both sides of each dimension in the input. Default: 0
        output_padding (int or tuple, optional): Additional size added to one side
            of each dimension in the output shape. Default: 0
        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``
        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1

    Shape:
        - Input: :math:`(N, C_{in}, H_{in}, W_{in})`
        - Output: :math:`(N, C_{out}, H_{out}, W_{out})` where

        .. math::
              H_{out} = (H_{in} - 1) \\times \\text{stride}[0] - 2 \\times \\text{padding}[0] + \\text{dilation}[0]

                        \\times (\\text{kernel_size}[0] - 1) + \\text{output_padding}[0] + 1
        .. math::
              W_{out} = (W_{in} - 1) \\times \\text{stride}[1] - 2 \\times \\text{padding}[1] + \\text{dilation}[1]

                        \\times (\\text{kernel_size}[1] - 1) + \\text{output_padding}[1] + 1

    Attributes:
        ConvTranspose2d.weight (Tensor): the learnable weights of the module of shape
                         :math:`(\\text{in_channels}, \\frac{\\text{out_channels}}{\\text{groups}},`
                         :math:`\\text{kernel_size[0]}, \\text{kernel_size[1]})`.
                         The values of these weights are sampled from
                         :math:`\\mathcal{U}(-\\sqrt{k}, \\sqrt{k})` where
                         :math:`k = \\frac{groups}{C_\\text{out} * \\prod_{i=0}^{1}\\text{kernel_size}[i]}`
        ConvTranspose2d.bias (Tensor): the learnable bias of the module of shape (out_channels)
                         If :attr:`bias` is ``True``, then the values of these weights are
                         sampled from :math:`\\mathcal{U}(-\\sqrt{k}, \\sqrt{k})` where
                         :math:`k = \\frac{groups}{C_\\text{out} * \\prod_{i=0}^{1}\\text{kernel_size}[i]}`

    Examples::

        >>> import numpy as np
        >>> import oneflow as flow
        >>> import oneflow.nn as nn

        >>> m = nn.ConvTranspose2d(16, 33, 3, stride=2)
        >>> # non-square kernels and unequal stride and with padding
        >>> m = nn.ConvTranspose2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2))
        >>> m = m.to("cuda")
        >>> input = flow.Tensor(np.random.randn(20, 16, 50, 100), device=flow.device("cuda"))
        >>> output = m(input)
        >>> output.size()
        flow.Size([20, 33, 93, 100])

    .. _cross-correlation:
        https://en.wikipedia.org/wiki/Cross-correlation

    .. _link:
        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: _size_2_t,
        stride: _size_2_t = 1,
        padding: _size_2_t = 0,
        output_padding: _size_2_t = 0,
        groups: int = 1,
        bias: bool = True,
        dilation: int = 1,
        padding_mode: str = "zeros",
    ) -> None:
        super().__init__()
        assert padding_mode == "zeros"
        kernel_size = _pair(kernel_size)
        stride = _pair(stride)
        padding = _pair(padding)
        output_padding = _pair(output_padding)
        dilation = _pair(dilation)
        self.groups = groups
        assert in_channels % groups == 0
        assert out_channels % groups == 0
        self.weight = flow.nn.Parameter(
            flow.Tensor(in_channels, out_channels // groups, *kernel_size)
        )
        self.in_channel_groups = in_channels // groups
        self.bias = None
        self._bias_add_op = None
        if bias:
            self.bias = flow.nn.Parameter(flow.Tensor(out_channels))
            self._bias_add_op = (
                flow.builtin_op("bias_add")
                .Input("a")
                .Input("b")
                .Output("out")
                .Attr("axis", 1)
                .Build()
            )
        self._op = (
            flow.builtin_op("deconv2d")
            .Input("in")
            .Input("weight")
            .Attr("filters", out_channels // groups)
            .Attr("padding_before", padding)
            .Attr("data_format", "channels_first")
            .Attr("kernel_size", kernel_size)
            .Attr("strides", stride)
            .Attr("dilation_rate", dilation)
            .Attr("output_padding", output_padding)
            .Attr("groups", 1)
            .Output("out")
            .Build()
        )
        self.reset_parameters()

    def reset_parameters(self) -> None:
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            (fan_in, _) = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in)
            init.uniform_(self.bias, -bound, bound)

    def forward(self, x):
        if self.groups > 1:
            in_channel_axis = 1
            in_split_list = ConvUtil.split(
                x, axis=in_channel_axis, split_num=self.groups
            )
            out_list = []
            for i in range(len(in_split_list)):
                out_list.append(
                    self._op(
                        in_split_list[i],
                        self.weight[
                            i
                            * self.in_channel_groups : (i + 1)
                            * self.in_channel_groups,
                            :,
                            :,
                            :,
                        ],
                    )[0]
                )
            res = flow.cat(out_list, dim=in_channel_axis)
        else:
            res = self._op(x, self.weight)[0]
        if self._bias_add_op is not None:
            res = self._bias_add_op(res, self.bias)[0]
        return res

if __name__ == "__main__":
    import doctest

    doctest.testmod(raise_on_error=True)


#!/usr/bin/env python3
"""Compute a background mask for X-ray microscopy data.

Functions
---------
parse_args
    Parse command line arguments.
initialize_cloudvolume
    Create a new CloudVolume archive.
load_image
    Load an image from CloudVolume.
create_bg_mask
    Create a mask of background regions in x-ray microscopy.
write_image
    Write an image to CloudVolume.

Dependencies
------------
cloud-volume
mpi4py
numpy
scipy
scikit-image
"""

import argparse
import logging
import os
import re

from cloudvolume import CloudVolume
from mpi4py import MPI
import numpy as np
from scipy.signal import find_peaks, peak_prominences, peak_widths
from skimage.exposure import histogram
from skimage.filters import gaussian
from skimage.measure import label, regionprops
from skimage.morphology import remove_small_holes

COMM = MPI.COMM_WORLD
RANK = COMM.Get_rank()
SIZE = COMM.Get_size()

LOGGER = logging.getLogger('create_background_mask.py')
syslog = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s %(name)s Rank %(rank)s : %(message)s')
syslog.setFormatter(formatter)
LOGGER.setLevel(logging.INFO)
LOGGER.addHandler(syslog)
LOGGER = logging.LoggerAdapter(LOGGER, {'rank': str(RANK)})

def parse_args():
    """Parse command line arguments."""
    p = argparse.ArgumentParser()

    p.add_argument('--input', type=str,
                   help='path to the CloudVolume archive')
    p.add_argument('--output', type=str,
                   help='path to the bg mask CloudVolume archive')
    p.add_argument('--resolution', type=int, nargs='*', default=[10, 10, 10],
                   help='resolution of the dataset')
    p.add_argument('--mip', type=int, default=0,
                   help='number of mip levels to create')
    p.add_argument('--chunk-size', type=int, nargs='*', default=[64, 64, 64],
                   help='size of each CloudVolume block file')
    p.add_argument('--z-step', type=int, default=None)
    p.add_argument('--factor', type=int, nargs='*', default=[2, 2, 2],
                   help='factor to scale between mip levels')
    p.add_argument('--flip-xy', action='store_true',
                   help='pass to transplose the X and Y axes')
    p.add_argument('--memory-limit', type=float, default=10000,
                   help='max memory available to CloudVolume')
    p.add_argument('--offset', type=int, nargs='*', default=[0, 0, 0],
                   help='offset into the volume from the upper-left corner')
    p.add_argument('--quiet', action='store_true',
                   help='pass to deactivate logging')

    return p.parse_args()

def initialize_cloudvolume(path, resolution, offset, volume_size, chunk_size,
                           mip, factor):
    """Create a new CloudVolume archive.

    Parameters
    ----------
    path : str
        Filepath to the location to write the archive.
    resolution : tuple of int
        Imaging resolution of the images in each dimension.
    offset : tuple of int
        Offset within the volume to the start of the archive.
    volume_size : tuple of int
        The dimensions of the volume in pixels.
    chunk_size : tuple of int
        The size of each CloudVolume block in pixels.
    mip : int
        The number of mip levels to include.
    factor : tuple of int
        The factor of change in each dimension across mip levels.

    Returns
    -------
    cv_args : dict
        The parameters needed to re-access the CloudVolume archive.
    """
    # Set the parameters of the info file.
    info = CloudVolume.create_new_info(
        num_channels=1,
        layer_type='segmentation',
        data_type='uint32',
        encoding='compressed_segmentation',
        resolution=resolution,
        voxel_offset=offset,
        volume_size=volume_size[:-1],
        chunk_size=chunk_size,
        max_mip=0,
        factor=factor
    )

    # Set up and initialize the CloudVolume object
    cv_args = dict(
        bounded=True, fill_missing=True, autocrop=False,
        cache=False, compress_cache=None, cdn_cache=False,
        progress=False, info=info, provenance=None, compress=True,
        non_aligned_writes=True, parallel=1)

    # for i in range(1, mip + 1):
    #     info['scales'][i]['compressed_segmentation_block_size'] = \
    #         info['scales'][0]['compressed_segmentation_block_size']

    cv = CloudVolume(path, mip=0, **cv_args)

    # Create the info file.
    LOGGER.info('Initializing image layer with config {}'.format(cv_args))
    cv.commit_info()
    return cv_args

def load_subvolume(cv, z_start, z_end, flip_xy=False):
    """Load an image from CloudVolume.

    Parameters
    ----------
    cv : cloudvolume.CloudVolume
        CloudVolume image layer to mask.
    z_start : int
        The index of the first image in the layer.
    z_end : int
        The index of the last image in the layer.
    flip_xy : bool
        CloudVolume reorders the dimension of image volumes, and the order of
        the x and y dimensions can vary. If True, indicates that the CloudVolume
        layer is saved in (Y, X, Z) order; otherwise it is saved as (X, Y, Z).

    Returns
    -------
    subvol : numpy.ndarray
        The subvolume with the dimensions reordered as (Z, Y, X).
    """
    # Each entry in the z dimension represents one image. Extract an image.
    subvol = cv[:, :, z_start:z_end, :]
    subvol = np.squeeze(subvol)

    # Transpose the dimensions back to
    if not flip_xy:
        subvol = np.transpose(subvol, axes=[2, 1, 0])

    LOGGER.info('Loaded subvolume with shape {}.'.format(subvol.shape))
    return subvol

def find_bg_mask(img):
    """Create a mask of background regions in x-ray microscopy.

    Parameters
    ----------
    img : numpy.ndarray
        X-ray microscopy image.

    Returns
    -------
    bgmask : numpy.ndarray
        Binary mask of the background of ``img``.
    """
    if img.ndim == 2:
        img = np.expand_dims(img, axis=0)

    bgmask = np.zeros((3,) +img.shape, dtype=np.uint8)

    for d in range(img.ndim):
        for i in range(img.shape[d]):
            if d == 0:
                subimg = img[i, :, :]
            elif d == 1:
                subimg = img[:, i, :]
            elif d == 2:
                subimg = img[:, :, i]

            # Blur the image to smooth any background artifacts.
            LOGGER.info('Blurring image.')
            blur = gaussian(subimg, sigma=5, preserve_range=True)

            # Compute the image histogram and find the peaks.
            LOGGER.info('Finding histogram peaks.')
            hist, bins = histogram(blur)
            peaks, properties = find_peaks(hist)  # , height=(0.3 * img.size))
            prominences = peak_prominences(hist, peaks)
            widths = peak_widths(hist, peaks, rel_height=0.333,
                                 prominence_data=prominences)

            # Select the left-most peak (backgrounds are usually dark) and use the
            # width of the peak to select a threshold value. Create a mask of all
            # pixels less than or equal to the threshold.
            ordered = np.argsort(peaks)
            threshold = peaks[ordered[0]] + (widths[0][ordered[0]] / 2.0)
            # threshold = peaks[0] + (widths[0][0] / 2.0)
            LOGGER.info('Setting hard threshold {} for image.'.format(threshold))
            mask = np.zeros(subimg.shape, dtype=np.uint8)
            mask[np.where(subimg <= threshold)] = 1
            # Perform some clean up and find the largest connected component.
            LOGGER.info('Cleaning mask of image.')
            # remove_small_holes(mask, area_threshold=30, connectivity=2,
            #                    in_place=True)
            labels = label(mask)
            objs = regionprops(labels)
            # bg = None
            # for obj in objs:
            #     if obj.bbox_area >= 0.85 * img.size:
            #         coords = obj.coords
            #         break

            # Select the connected component with the largest bounding box as the
            # background mask.
            objs.sort(key=lambda x: x.bbox_area, reverse=True)
            # objs = [o for o in objs
            #         if np.any(np.asarray(o.bbox[:mask.ndim]) == np.asarray(mask.shape))
            #         or np.any(np.asarray(o.bbox[mask.ndim:]) == 0)]
            print(len(objs))
            if len(objs) > 0:
                coords = tuple([objs[0].coords[:, j] for j in range(subimg.ndim)])
                LOGGER.info('Setting background mask of image.')

                if d == 0:
                    bgmask[d, i, coords[0], coords[1]] = 1
                elif d == 1:
                    bgmask[d, coords[0], i, coords[1]] = 1
                elif d == 2:
                    bgmask[d, coords[0], coords[1], i] = 1
    LOGGER.info('Full background mask covers {} voxels.'.format(np.sum(bgmask)))

    consensus = bgmask[0] + bgmask[1] + bgmask[2]
    consensus[np.where(consensus == 1)] = 0
    consensus[consensus.nonzero()] = 1
    objs = sorted(regionprops(label(consensus)), key=lambda x: x.bbox_area, reverse=True)
    for obj in objs[1:]:
        coords = tuple([obj.coords[:, j] for j in range(img.ndim)])
        consensus[coords] = 0

    LOGGER.info('Full background mask covers {} voxels.'.format(np.sum(consensus)))
    return consensus.astype(np.uint32)

def write_subvolume(path, subvolume, flip_xy, z_start, mip, factor):
    """Write an image to CloudVolume.

    Parameter
    ---------
    path : str
        Filepath to the location to write the archive.
    subvolume : numpy.ndarray
        Image data to write to the archive.
    flip_xy : bool
        If True, order ``layer`` as [Y, X, Z]. Otherwise, order ``layer`` as
        [X, Y, Z].
    z_start
        The starting index of ``layer`` within the archive.
    mip
        The number of mip levels to compute.
    factor
        The factor by which to reduce each mip level along each dimension.
    """
    # Transpose the axes to match the CloudVolume order
    if subvolume.ndim == 2:
        subvolume = np.expand_dims(subvolume, 0)

    if flip_xy:
        subvolume = np.transpose(subvolume, axes=[1, 2, 0])
    else:
        subvolume = np.transpose(subvolume, axes=[2, 1, 0])

    if subvolume.ndim == 3:
        subvolume = np.expand_dims(subvolume, -1)

    cv_args = dict(
        bounded=True, fill_missing=True, autocrop=False,
        cache=False, compress_cache=None, cdn_cache=False,
        progress=False, info=None, provenance=None, compress=True,
        non_aligned_writes=True, parallel=1)

    # Set the volume for each mip level
    for m in range(1):
        # Access the CloudVolume
        LOGGER.info('Writing MIP level {}.'.format(mip))
        cv = CloudVolume(path, mip=m, **cv_args)

        # Compute the index of this layer in the CloudVolume archive
        offset = cv.mip_voxel_offset(m)
        step = np.power(np.array(factor), m)
        cv_z_start = int(z_start // step[2] + offset[2])
        cv_z_end = int(min(cv_z_start + subvolume.shape[-2], cv.shape[-2]))

        # Set the layer
        cv[:, :, cv_z_start:cv_z_end] = subvolume

        # Reduce the size of the layer to match the next mip level
        subvolume = subvolume[::factor[0], ::factor[1], ::factor[2]]

def create_background_mask(input, output, resolution=(10, 10, 10), mip=0,
                           chunk_size=(64, 64, 64), z_step=None,
                           factor=(2, 2, 2), flip_xy=False, memory_limit=10000,
                           offset=(0, 0, 0), quiet=False):
    """Create and write data to a new CloudVolume archive."""
    if quiet:
        LOGGER.logger.removeHandler(syslog)
        noop = logging.NullHandler()
        LOGGER.logger.addHandler(noop)

    if 'image' not in os.path.basename(input):
        inpath = input + '/image'
    else:
        inpath = input

    if os.path.isdir(inpath) and not re.search(r'^file://', inpath):
        inpath = 'file://' + os.path.abspath(inpath)

    if RANK == 0:
        LOGGER.info('Loading CloudVolume image layer {}.'.format(inpath))

    img_cv = CloudVolume(inpath)
    volume_shape = img_cv.shape

    outpath = os.path.abspath(output)

    if os.path.dirname(inpath) == outpath:
        outpath = outpath + 'background'

    if not re.search(r'^[\w]+://.+$', outpath):
        outpath = 'file://' + os.path.abspath(output)

    # On rank 0, initialize the CloudVolume info file, and load in the list of
    # images to insert into the archive.
    if RANK == 0:
        LOGGER.info('Initialized CloudVolume image layer at {}'.format(outpath))
        cv_args = initialize_cloudvolume(
            outpath,
            resolution,
            offset,
            volume_shape,
            chunk_size,
            mip,
            factor)

    # Block until the background CloudVolume layer is initialized.
    GOGOGO = COMM.bcast(1, root=0)

    # Iterate over layers of the volume. Each rank will load and write one
    # layer at a time. If there are fewer ranks than layers, increment to
    # n_ranks + rank and load the layer at that index.
    # offset from the volume origin.
    layer_idx = RANK * chunk_size[-1]
    while layer_idx < volume_shape[-2]:
        # Compute the index of the first image in this layer, including any
        layer_shape = int(min(layer_idx + chunk_size[-1],
                              img_cv.shape[-2]))
        LOGGER.info('Loading images {}-{}.'.format(layer_idx, layer_shape))
        image = load_subvolume(img_cv, layer_idx, layer_shape,
                               flip_xy=flip_xy)

        LOGGER.info('Creating background mask.')
        mask = find_bg_mask(image)

        # Write the layer to the archive.
        LOGGER.info('Writing mask of images {}-{}'.format(layer_idx, layer_shape))
        write_subvolume(
            outpath,
            mask,
            flip_xy,
            layer_idx,
            mip,
            factor)

        # Increment to the next known layer that does not overlap with any
        # other rank.
        layer_idx += SIZE * chunk_size[-1]

    LOGGER.info('Done.')

def main():
    args = parse_args()
    create_background_mask(
        args.input,
        args.output,
        resolution=args.resolution,
        mip=args.mip,
        chunk_size=args.chunk_size,
        z_step=args.z_step,
        factor=args.factor,
        flip_xy=args.flip_xy,
        memory_limit=args.memory_limit,
        offset=args.offset,
        quiet=args.quiet)

if __name__ == '__main__':
    main()