import os
import numpy as np
from copy import deepcopy

from common import Common, co
from bird import Bird
from nonlinear import NonLinear
from resum import Resum
from projection import Projection
from angular import Angular
from greenfunction import GreenFunction

# import importlib, sys
# importlib.reload(sys.modules['common'])
# importlib.reload(sys.modules['bird'])
# importlib.reload(sys.modules['nonlinear'])
# importlib.reload(sys.modules['resum'])
# importlib.reload(sys.modules['projection'])
# importlib.reload(sys.modules['angular'])
# importlib.reload(sys.modules['greenfunction'])

# from common import Common, co
# from bird import Bird
# from nonlinear import NonLinear
# from resum import Resum
# from projection import Projection
# from angular import Angular
# from greenfunction import GreenFunction

class Correlator(object):

    def __init__(self, config_dict=None, load_engines=True):

        self.cosmo_catalog = {
            "P11": Option("P11", (list, np.ndarray),
                description="Linear matter power spectrum in [Mpc/h]^3",
                default=None) ,
            "k11": Option("k11", (list, np.ndarray),
                description="k-array in [h/Mpc] on which P11 is evaluated",
                default=None) ,
            "D": Option("D", (float, list, np.ndarray),
                description="Scale independent growth function. To specify if \'skycut\' > 1.",
                default=None) ,
            "f": Option("f", (float, list, np.ndarray),
                description="Scale independent growth rate (for RSD). Automatically set to 0 for \'output\': \'m__\'.",
                default=None) ,
            "bias": Option("bias", (dict, list, np.ndarray),
                description="EFT parameters in dict = \{ \'b1\', \'b2\', \'b3\', \'b4\', \'cct\', \'cr1\', \'cr2\', \'ce0\', \'ce1\', \'ce2\' \}.",
                default=None) ,
            "Omega0_m": Option("Omega0_m", float,
                description="Fractional matter abundance at present time. To specify for exact time dependence.",
                default=None) ,
            "w0_fld": Option("w0_fld", float,
                description="Dark energy equation of state parameter. To specify for exact time dependence if varied (otherwise w0 = -1).",
                default=None) ,
            "z": Option("z", (float, list, np.ndarray),
                description="Effective redshift(s). Should match the number of skycuts. To specify for exact time dependence.",
                default=None) ,
            "DA": Option("DA", (float, list, np.ndarray),
                description="Angular distance times H_0. To specify if \'with_AP\' is True.",
                default=None) ,
            "H": Option("H", (float, list, np.ndarray),
                description="Hubble parameter by H_0. To specify if \'with_AP\' is True.",
                default=None) ,
            "Dz": Option("Dz", (list, np.ndarray),
                description="Scale independent growth function over redshift bin. To specify if \'with_redshift_bin\' is True.",
                default=None) ,
            "fz": Option("fz", (list, np.ndarray),
                description="Scale independent growth rate over redshift bin. To specify if \'with_redshift_bin\' is True.",
                default=None) ,
            "DAz": Option("DAz", (list, np.ndarray),
                description="Angular distance times H_0 over redshift bin. To specify if \'with_redshift_bin\' and \'with_AP\' are True.",
                default=None) ,
            "Hz": Option("Hz", (list, np.ndarray),
                description="Hubble parameter by H_0 over redshift bin. To specify if \'with_redshift_bin\' and \'with_AP\' are True.",
                default=None) ,
            "rz": Option("rz", (list, np.ndarray),
                description="Comoving distance in [Mpc/h] over redshift bin. To specify if \'output\':\'w\'.",
                default=None) ,
        }

        self.config_catalog = {
            "output": Option("output", str, ["bPk", "bCf", "mPk", "mCf", "bmPk", "bmCf", "w"],
                description="Correlator: biased tracers / matter / biased tracers-matter -- power spectrum / correlation function ; \'w\': angular correlation function. ",
                default="bPk") ,
            "multipole": Option("multipole", int, [0, 2, 3],
                description="Number of multipoles. 0: real space. 2: monopole + quadrupole. 3: monopole + quadrupole + hexadecapole.",
                default=2) ,
            "wedge": Option("wedge", int,
                description="Number of wedges. 0: compute multipole instead. ",
                default=0) ,
            "wedges_bounds": Option("wedges_bounds", (list, np.ndarray),
                description="Wedges bounds: [0, a_1, ..., a_{n-1}, 1], n: number of wedges, such that 0 < mu < a_1, ..., a_{n-1} < mu < 1. Default: equi-spaced between 0 and 1.",
                default=None) ,
            "skycut": Option("skycut", int,
                description="Number of skycuts.",
                default=1) ,
            "with_time": Option("with_time", bool,
                description="Time (in)dependent evaluation (for multi skycuts / redshift bin). For \'with_redshift_bin\': True, \'skycut\' > 1 or \'output\': \'w\', automatically set to False.",
                default=True) ,
            "z": Option("z", (float, list, np.ndarray),
                description="Effective redshift(s). Should match the number of skycuts.",
                default=None) ,
            "km": Option("km", float,
                description="Inverse galaxy spatial extension scale in [h/Mpc].",
                default=1.) ,
            "nd": Option("nd", float,
                description="Mean galaxy density",
                default=3e-4) ,
            "with_stoch": Option("with_stoch", bool,
                description="With stochastic terms.",
                   default=False) ,
            "with_bias": Option("with_bias", bool,
                description="Bias (in)dependent evalution. Automatically set to False for \'with_time\': False.",
                   default=False) ,
            "with_exact_time": Option("with_exact_time", bool,
                description="Exact time dependence or EdS approximation.",
                default=False) ,
            "with_redshift_bin": Option("with_redshift_bin", bool,
                description="Account for the galaxy count distribution over a redshift bin.",
                default=False) ,
            "zz": Option("zz", (list, np.ndarray),
                description="Array of redshift points inside a redshift bin. For multi skycuts, a list of arrays should be provided.",
                default=None) ,
            "nz": Option("nz", (list, np.ndarray),
                description="Galaxy counts distribution over a redshift bin. For multi skycuts, a list of arrays should be provided.",
                default=None) ,
            "kmax": Option("kmax", float,
                description="kmax in [h/Mpc] for \'output\': \'_Pk\'",
                default=0.25) ,
#             "smin": Option("smin", float,
#                 description="smin in [Mpc/h] for \'output\': \'_Cf\'",
#                 default=1.) ,
            "accboost": Option("accboost", int, [1, 2, 3],
                description="Sampling accuracy boost.",
                default=1) ,
            "optiresum": Option("optiresum", bool,
                description="True: Resumming only with the BAO peak. False: Resummation on the full correlation function.",
                default=False) ,
            "xdata": Option("xdata", (np.ndarray, list),
                description="Array of data points.",
                default=None) ,
            "with_resum": Option("with_resum", bool,
                description="Apply IR-resummation.",
                default=True) ,
            "with_AP": Option("wity_AP", bool,
                description="Apply Alcock Paczynski effect. Automatically set to False for \'output\': \'w\'.",
                default=False) ,
            "z_AP": Option("z_AP", (float, list, np.ndarray),
                description="Fiducial redshift used to convert coordinates to distances. A list can be provided for multi skycuts. If only one value is passed, use it for all skycuts.",
                default=None) ,
            "Omega_m_AP": Option("Omega_m_AP", (float, list, np.ndarray),
                description="Fiducial matter abundance used to convert coordinates to distances. A list can be provided for multi skycuts. If only one value is passed, use it for all skycuts.",
                default=None) ,
            "with_window": Option("window", bool,
                description="Apply mask. Automatically set to False for \'output\': \'w\' or \'_Cf\'.",
                default=False) ,
            "windowPk": Option("windowPk", (str, list),
                description="Path to Fourier convolution window file for \'output\': \'_Pk\'. If not provided, read \'windowCf\', precompute the Fourier one and save it here.",
                default=None) ,
            "windowCf": Option("windowCf", (str, list),
                description="Path to configuration space window file with columns: s [Mpc/h], Q0, Q2, Q4. A list can be provided for multi skycuts. Put \'None\' for each skycut without window.",
                default=None) ,
            "with_binning": Option("with_binning", bool,
                description="Apply binning for linear-spaced data bins.",
                default=False) ,
            "with_fibercol": Option("with_fibercol", bool,
                description="Apply fiber collision effective window corrections.",
                default=False) ,
            "with_nlo_bias": Option("with_nlo_bias", bool,
                description="With next-to-leading bias.",
                default=False) ,
            "with_tidal_alignments": Option("with_tidal_alignments", bool,
                description="With tidal alignements: bq * (\mu^2 - 1/3) \delta_m ",
                default=False) ,
            "with_quintessence": Option("with_quintessence", bool,
                description="Clustering quintessence. Automatically set \'with_exact_time\' to True.",
                default=False) ,
            "with_MG": Option("with_MG", bool,
                description="General scale-independent modified gravity. Automatically set \'with_exact_time\' to True.",
                default=False) ,
            "with_nDGP": Option("with_nDGP", bool,
                description="nDGP modified gravity with LCDM background. Automatically set \'with_exact_time\' and \'with_MG\' to True.",
                default=False) ,
            "Omega_rc": Option("Omega_rc", float,
                description="nDGP parameter",
                default=None) , #set default?
            "w_integrator": Option("w_integrator", str,
                description="Type of integration for \'output\': \'w\' : \'trapz\', \'cuba\' or \'fast\'. \'fast\' is tuned for DESY1 precision.",
                default='cuba') ,
            "w_theta_cut": Option("w_theta_cut", (float, list, np.ndarray),
                description="Angle cut for \'output\': \'w\'. ",
                default=0) ,
            "nonmarg": Option("nonmarg", bool,
                description="General scale-independent modified gravity. Automatically set \'with_exact_time\' to True.",
                default=False) ,
        }

        if config_dict is not None: self.set(config_dict, load_engines=load_engines)

    def info(self, description=True):

        for on in ['config', 'cosmo']:

            print ("\n")
            if on is 'config':
                print ("Configuration commands [.set(config_dict)]")
                print ("----------------------")
                catalog = self.config_catalog
            elif on is 'cosmo':
                print ("Cosmology commands [.compute(cosmo_dict)]")
                print ("------------------")
                catalog = self.cosmo_catalog

            for (name, config) in zip(catalog, catalog.values()):
                if config.list is None: print("\'%s\': %s" % (name, typename(config.type)))
                else: print("\'%s\': %s ; options: %s" % (name, typename(config.type), config.list))
                if description:
                    print ('    - %s' % config.description)
                    print ('    * default: %s' % config.default)

    def set(self, config_dict, load_engines=True):

        # Reading config provided by user
        self.__read_config(config_dict)

        # Setting no-optional config
        self.config["smin"] = 1.
        self.config["smax"] = 1000.
        self.config["kmin"] = 0.001

        # Checking for config conflict
        self.__is_config_conflict()

        # Loading PyBird engines
        #print('nDGP pybird set:', self.config["with_nDGP"])
        if load_engines: self.__load_engines()
        else: self.__load_engines(only_common=True)

    def compute(self, cosmo_dict, module=None):

        cosmo_dict_local = cosmo_dict.copy()

        if module is 'class':
            cosmo_dict_class = self.setcosmo(cosmo_dict, module='class')
            cosmo_dict_local.update(cosmo_dict_class)

        self.__read_cosmo(cosmo_dict_local)
        self.__is_cosmo_conflict()

        if self.config["skycut"] == 1:
            self.bird = Bird(self.cosmo, with_bias=self.config["with_bias"], with_stoch=self.config["with_stoch"], with_nlo_bias=self.config["with_nlo_bias"], co=self.co)
            self.nonlinear.PsCf(self.bird)
            if self.config["with_bias"]: self.bird.setPsCf(self.bias)
            else: self.bird.setPsCfl()
            if self.config["with_resum"]:
                if self.config["with_cf"]: self.resum.PsCf(self.bird)
                else: self.resum.Ps(self.bird)
            if "w" in self.config["output"]: self.angular.w(self.bird, self.cosmo["Dz"], self.cosmo["fz"], self.cosmo["rz"], self.config["zz"], self.config["nz"])
            else:
                if self.config["with_redshift_bin"]: self.projection.redshift(self.bird, self.cosmo["Dz"], self.cosmo["fz"], self.cosmo["DAz"], self.cosmo["Hz"])
                elif self.config["with_AP"]: self.projection.AP(self.bird)
                if self.config["with_window"]: self.projection.Window(self.bird)
                if self.config["with_fibercol"]: self.projection.fibcolWindow(self.bird)
                if self.config["wedge"] is not 0: self.projection.Wedges(self.bird)
                if self.config["with_binning"]: self.projection.xbinning(self.bird)
                else: self.projection.xdata(self.bird)

        elif self.config["skycut"] > 1:
            if self.config["with_time"]: # if all skycuts have same redshift
                cosmoi = deepcopy(self.cosmo)
                cosmoi["f"] = self.cosmo["f"][0]
                cosmoi["D"] = self.cosmo["D"][0]
                cosmoi["z"] = self.config["z"][0]
                if self.config["with_AP"] and not self.config["with_redshift_bin"]:
                    cosmoi["DA"] = self.cosmo["DA"][0]
                    cosmoi["H"] = self.cosmo["H"][0]
                self.bird = Bird(cosmoi, with_bias=False, with_stoch=self.config["with_stoch"], with_nlo_bias=self.config["with_nlo_bias"], co=self.co)
                self.nonlinear.PsCf(self.bird)
                self.bird.setPsCfl()
                if self.config["with_resum"]:
                    if self.config["with_cf"]: self.resum.PsCf(self.bird)
                    else: self.resum.Ps(self.bird)
                self.birds = [deepcopy(self.bird) for i in range(self.config["skycut"])]

            else:
                self.birds = []
                cosmoi = deepcopy(self.cosmo)

                def mycycle(skycut, first=0, L=None):
                    if L is None: L = [i for i in range(skycut)]
                    if first == 0: first = (skycut+1)//2
                    else: first = skycut//2
                    return [item for i, item in enumerate(L+L) if i < skycut+first and first <= i]

                zbins = mycycle(self.config["skycut"], first=2) # cycle to get the middle redshift
                for i in zbins:
                    cosmoi["f"] = self.cosmo["f"][i]
                    cosmoi["D"] = self.cosmo["D"][i]
                    cosmoi["z"] = self.config["z"][i]

                    if self.config["with_AP"] and not self.config["with_redshift_bin"]:
                        cosmoi["DA"] = self.cosmo["DA"][i]
                        cosmoi["H"] = self.cosmo["H"][i]

                    if i == zbins[0]:
                        self.bird = Bird(cosmoi, with_bias=False, with_stoch=self.config["with_stoch"], with_nlo_bias=self.config["with_nlo_bias"], co=self.co)
                        self.nonlinear.PsCf(self.bird)
                        self.bird.setPsCfl()
                        if self.config["with_resum"]: self.resum.Ps(self.bird, makeIR=True, makeQ=True, setPs=False)
                        self.birds.append(self.bird)
                    else:
                        birdi = deepcopy(self.bird)
                        birdi.settime(cosmoi) # rescale by (Dnew/Dold)**(2p) # f if not given yet at this stage?
                        if self.config["with_resum"]:
                            if self.config["with_cf"]: self.resum.PsCf(birdi, makeIR=False, makeQ=True, setPs=False, setCf=True)
                            else: self.resum.Ps(birdi, makeIR=False, makeQ=True, setPs=True)
                        self.birds.append(birdi)

                if self.config["with_resum"]:
                    if self.config["with_cf"]: self.resum.PsCf(self.birds[0], makeIR=False, makeQ=False, setPs=False, setCf=True)
                    else: self.resum.Ps(self.birds[0], makeIR=False, makeQ=False, setPs=True)

                self.birds = mycycle(self.config["skycut"], first=0, L=self.birds) # cycle back the birds

            for i in range(self.config["skycut"]):
                if "w" in self.config["output"]:
                    self.angular[i].w(self.birds[i], self.cosmo["Dz"][i], self.cosmo["fz"][i], self.cosmo["rz"][i], self.config["zz"][i], self.config["nz"][i], which=self.config["w_integrator"], theta_cut=self.config["w_theta_cut"][i])
                else:
                    if self.config["with_redshift_bin"] and self.config["nz"][i] is not None: self.projection[i].redshift(self.birds[i], self.cosmo["Dz"][i], self.cosmo["fz"][i], self.cosmo["DAz"][i], self.cosmo["Hz"][i])
                    elif self.config["with_AP"]: self.projection[i].AP(self.birds[i])
                    if self.config["with_window"]: self.projection[i].Window(self.birds[i])
                    if self.config["with_fibercol"]: self.projection[i].fibcolWindow(self.birds[i])
                    if self.config["wedge"] is not 0: self.projection[i].Wedges(self.birds[i])
                    if self.config["with_binning"]: self.projection[i].xbinning(self.birds[i])
                    else: self.projection[i].xdata(self.birds[i])

    def cache(self, as_dict=True):

        if self.config["skycut"] == 1:

            if as_dict:
                correlator_cache = {}
                correlator_cache = {"f": self.cosmo["f"]}
                if "Pk" in self.config["output"]:
                    correlator_cache["lin"] = self.bird.P11l
                    correlator_cache["loop"] = self.bird.Ploopl
                    correlator_cache["ct"] = self.bird.Pctl
                    if self.config["with_stoch"]: correlator_cache["st"] = self.bird.Pstl
                    if self.config["with_nlo_bias"]: correlator_cache["nlo"] = self.bird.Pnlol
                elif "Cf" in self.config["output"]:
                    correlator_cache["lin"] = self.bird.C11l
                    correlator_cache["loop"] = self.bird.Cloopl
                    correlator_cache["ct"] = self.bird.Cctl
                    if self.config["with_nlo_bias"]: correlator_cache["nlo"] = self.bird.Cnlol
                elif "w" in self.config["output"]:
                    correlator_cache["lin"] = self.bird.wlin
                    correlator_cache["loop"] = self.bird.wloop
                    correlator_cache["ct"] = self.bird.wct
                    if self.config["with_nlo_bias"]: correlator_cache["nlo"] = self.bird.wnlo
                return correlator_cache

            else: return self.bird

        elif self.config["skycut"] > 1:

            if as_dict:
                correlator_cache = {}
                correlator_cache = {"f": self.cosmo["f"]}
                if "Pk" in self.config["output"]:
                    correlator_cache["lin"] = [self.birds[i].P11l for i in range(self.config["skycut"])]
                    correlator_cache["loop"] = [self.birds[i].Ploopl for i in range(self.config["skycut"])]
                    correlator_cache["ct"] = [self.birds[i].Pctl for i in range(self.config["skycut"])]
                    if self.config["with_stoch"]: correlator_cache["st"] = [self.birds[i].Pstl for i in range(self.config["skycut"])]
                    if self.config["with_nlo_bias"]: correlator_cache["nlo"] = [self.birds[i].Pnlol for i in range(self.config["skycut"])]
                elif "Cf" in self.config["output"]:
                    correlator_cache["lin"] = [self.birds[i].C11l for i in range(self.config["skycut"])]
                    correlator_cache["loop"] = [self.birds[i].Cloopl for i in range(self.config["skycut"])]
                    correlator_cache["ct"] = [self.birds[i].Cctl for i in range(self.config["skycut"])]
                    if self.config["with_nlo_bias"]: correlator_cache["nlo"] = [self.birds[i].Cnlol for i in range(self.config["skycut"])]
                elif "w" in self.config["output"]:
                    correlator_cache["lin"] = [self.birds[i].wlin for i in range(self.config["skycut"])]
                    correlator_cache["loop"] = [self.birds[i].wloop for i in range(self.config["skycut"])]
                    correlator_cache["ct"] = [self.birds[i].wct for i in range(self.config["skycut"])]
                    if self.config["with_nlo_bias"]: correlator_cache["nlo"] = [self.birds[i].wnlo for i in range(self.config["skycut"])]
                return correlator_cache

            else: return self.birds

    def setcache(self, correlator_cache, as_dict=True):

        self.__read_cosmo({})

        if self.config["skycut"] == 1:

            if as_dict:
                self.cosmo["f"] = correlator_cache["f"]
                self.bird = Bird(self.cosmo, with_bias=False, with_stoch=self.config["with_stoch"], with_nlo_bias=self.config["with_nlo_bias"], co=self.co)
                if "Pk" in self.config["output"]:
                    self.bird.P11l = correlator_cache["lin"]
                    self.bird.Ploopl = correlator_cache["loop"]
                    self.bird.Pctl = correlator_cache["ct"]
                    if self.config["with_stoch"]: self.bird.Pstl = correlator_cache["st"]
                    if self.config["with_nlo_bias"]: self.bird.Pnlol = correlator_cache["nlo"]
                elif "Cf" in self.config["output"]:
                    self.bird.C11l = correlator_cache["lin"]
                    self.bird.Cloopl = correlator_cache["loop"]
                    self.bird.Cctl = correlator_cache["ct"]
                    if self.config["with_nlo_bias"]: self.bird.Cnlol = correlator_cache["nlo"]
                elif "w" in self.config["output"]:
                    self.bird.wlin = correlator_cache["lin"]
                    self.bird.wloop = correlator_cache["loop"]
                    self.bird.wct = correlator_cache["ct"]
                    if self.config["with_nlo_bias"]: self.bird.wnlo = correlator_cache["nlo"]

            else: self.bird = correlator_cache

        elif self.config["skycut"] > 1:

            if as_dict:
                self.birds = []
                for i in range(self.config["skycut"]):
                    self.cosmo["f"] = correlator_cache["f"][i]
                    self.bird = Bird(self.cosmo, with_bias=False, with_stoch=self.config["with_stoch"], with_nlo_bias=self.config["with_nlo_bias"], co=self.co)
                    if "Pk" in self.config["output"]:
                        self.bird.P11l = correlator_cache["lin"][i]
                        self.bird.Ploopl = correlator_cache["loop"][i]
                        self.bird.Pctl = correlator_cache["ct"][i]
                        if self.config["with_stoch"]: self.bird.Pstl = correlator_cache["st"][i]
                        if self.config["with_nlo_bias"]: self.bird.Pnlol = correlator_cache["nlo"][i]
                    elif "Cf" in self.config["output"]:
                        self.bird.C11l = correlator_cache["lin"][i]
                        self.bird.Cloopl = correlator_cache["loop"][i]
                        self.bird.Cctl = correlator_cache["ct"][i]
                        if self.config["with_nlo_bias"]: self.bird.Cnlol = correlator_cache["nlo"][i]
                    elif "w" in self.config["output"]:
                        self.bird.wlin = correlator_cache["lin"][i]
                        self.bird.wloop = correlator_cache["loop"][i]
                        self.bird.wct = correlator_cache["ct"][i]
                        if self.config["with_nlo_bias"]: self.bird.wnlo = correlator_cache["nlo"][i]
                    self.birds.append(self.bird)

            else: self.birds = correlator_cache

    def get(self, bias=None):

        if self.config["skycut"] == 1:

            if not self.config["with_bias"]:
                self.__is_bias_conflict(bias)
                if "Pk" in self.config["output"]: self.bird.setreducePslb(self.bias)
                elif "Cf" in self.config["output"]: self.bird.setreduceCflb(self.bias)
                elif "w" in self.config["output"]: self.bird.setw(self.bias)

            if "Pk" in self.config["output"]: return self.bird.fullPs
            elif "Cf" in self.config["output"]: return self.bird.fullCf
            elif "w" in self.config["output"]: return self.bird.fullw

        elif self.config["skycut"] > 1:

            if not isinstance(bias, (list, np.ndarray)): raise Exception("Please specify bias (in a list of dicts) for each corresponding skycuts. ")
            if len(bias) is not self.config["skycut"]: raise Exception("Please specify bias (in a list of dicts) for each corresponding skycuts. ")

            for i in range(self.config["skycut"]):
                self.__is_bias_conflict(bias[i])
                if "w" in self.config["output"]: self.birds[i].setw(self.bias)
                elif "Cf" in self.config["output"]: self.birds[i].setreduceCflb(self.bias)
                elif "Pk" in self.config["output"]: self.birds[i].setreducePslb(self.bias)
            if "Pk" in self.config["output"]: return [self.birds[i].fullPs for i in range(self.config["skycut"])]
            elif "Cf" in self.config["output"]: return [self.birds[i].fullCf for i in range(self.config["skycut"])]
            elif "w" in self.config["output"]: return [self.birds[i].fullw for i in range(self.config["skycut"])]

    def getmarg(self, b1, model=1, bq=0):

        def marg(loop, ct, b1, f1, Pst=None, model=model, bq=0):

            # if self.config["with_redshift_bin"]: f = 1. #### This is wrong: the redshift integration is with f(z)/f(z_fid) ; f(z_fid) is given in setBias(), not before.
            # else: f = f1
            f = f1

            if loop.ndim is 3:
                loop = np.swapaxes(loop, axis1=0, axis2=1)
                ct = np.swapaxes(ct, axis1=0, axis2=1)
                if Pst is not None: Pst = np.swapaxes(Pst, axis1=0, axis2=1)

            if self.co.Nloop is 12: Pb3 = loop[3] + b1 * loop[7]                            # config["with_time"] = True
            elif self.co.Nloop is 18: Pb3 = loop[3] + b1 * loop[7] + bq * loop[16]          # config["with_time"] = True, config["with_tidal_alignments"] = True
            elif self.co.Nloop is 22: Pb3 = f * loop[8] + b1 * loop[16]                     # config["with_time"] = False, config["with_exact_time"] = False
            elif self.co.Nloop is 35: Pb3 = f * loop[18] + b1 * loop[29]                    # config["with_time"] = False, config["with_exact_time"] = True

            m = np.array([ Pb3.reshape(-1), 2 * (f * ct[0+3] + b1 * ct[0]).reshape(-1) / self.config["km"]**2 ])

            if "w" not in self.config["output"]:
                if self.config["multipole"] >= 2: m = np.vstack([m, 2 * (f * ct[1+3] + b1 * ct[1]).reshape(-1) / self.config["km"]**2])
                if self.config["multipole"] >= 3: m = np.vstack([m, 2 * (f * ct[2+3] + b1 * ct[2]).reshape(-1) / self.config["km"]**2])

                if self.config["with_stoch"]:
                    if model <= 4: m = np.vstack([m, Pst[2].reshape(-1) ]) # k^2 quad
                    if model == 1: m = np.vstack([m, Pst[0].reshape(-1) ]) # k^0 mono
                    if model == 3: m = np.vstack([m, Pst[1].reshape(-1) ]) # k^2 mono
                    if model == 4: m = np.vstack([m, Pst[1].reshape(-1) , Pst[0].reshape(-1) ]) # k^2 mono, k^0 mono

            return m

        if self.config["skycut"] == 1:
            if "Pk" in self.config["output"]: return marg(self.bird.Ploopl, self.bird.Pctl, b1=b1, f1=self.bird.f, Pst=self.bird.Pstl, bq=bq)
            elif "Cf" in self.config["output"]: return marg(self.bird.Cloopl, self.bird.Cctl, b1=b1, f1=self.bird.f, Pst=self.bird.Cstl, bq=bq)
            elif "w" in self.config["output"]: return marg(self.bird.wloop, self.bird.wct, b1=b1, f1=self.bird.f, bq=bq)
        elif self.config["skycut"] > 1:
            if "Pk" in self.config["output"]: return [ marg(self.birds[i].Ploopl, self.birds[i].Pctl, b1=b1[i], f1=self.birds[i].f, Pst=self.birds[i].Pstl, bq=bq[i]) for i in range(self.config["skycut"]) ]
            elif "Cf" in self.config["output"]: return [ marg(self.birds[i].Cloopl, self.birds[i].Cctl, b1=b1[i], f1=self.birds[i].f, Pst=self.birds[i].Cstl, bq=bq[i]) for i in range(self.config["skycut"]) ]
            elif "w" in self.config["output"]: return [ marg(self.birds[i].wloop, self.birds[i].wct, b1=b1[i], f1=self.birds[i].f, bq=bq[i]) for i in range(self.config["skycut"]) ]

    def getnlo(self, model=1):
        if self.config["skycut"] == 1:
            if "Pk" in self.config["output"]: return self.bird.Ps[1]**2/self.bird.Ps[0]
            elif "Cf" in self.config["output"]:
                lowpass = 1.-1./(1.+np.exp(-self.projection.xout+50)) # avoid 0-crossing by damping the nlo term for s > scut
                return self.bird.Cf[1]**2/self.bird.Cf[0] * lowpass
            elif "w" in self.config["output"]: return self.bird.w[1]**2/self.bird.w[0]
        elif self.config["skycut"] > 1:
            if "Pk" in self.config["output"]: return [ self.birds[i].Ps[1]**2/self.birds[i].Ps[0] for i in range(self.config["skycut"]) ]
            elif "Cf" in self.config["output"]: return [ self.birds[i].Cf[1]**2/self.birds[i].Cf[0] * (1.-1./(1.+np.exp(-self.projection[i].xout+50))) for i in range(self.config["skycut"]) ]
            elif "w" in self.config["output"]: return [ self.birds[i].w[1]**2/self.birds[i].w[0] for i in range(self.config["skycut"]) ]

    def __load_engines(self, only_common=False):

        #print('nDGP pybird load:', self.config["with_nDGP"])
        self.co = Common(Nl=self.config["multipole"], kmax=self.config["kmax"], km=self.config["km"], nd=self.config["nd"], halohalo=self.config["halohalo"],
            with_cf=self.config["with_cf"], with_time=self.config["with_time"], optiresum=self.config["optiresum"],
            exact_time=self.config["with_exact_time"], quintessence=self.config["with_quintessence"], MG=self.config["with_MG"], Omega_rc=self.config["Omega_rc"], nDGP=self.config["with_nDGP"], with_tidal_alignments=self.config["with_tidal_alignments"])

        if not only_common:
            self.nonlinear = NonLinear(load=True, save=True, co=self.co)
            self.resum = Resum(co=self.co)

            if "w" in self.config["output"]:
                if self.config["skycut"] == 1:
                    self.angular = Angular(self.config["xdata"], co=self.co)
                elif self.config["skycut"] > 1:
                    self.angular = []
                    for i in range(self.config["skycut"]):
                        if len(self.config["xdata"]) is self.config["skycut"]: xdata = self.config["xdata"][i]
                        else: xdata = self.config["xdata"]
                        self.angular.append(Angular(xdata, co=self.co))

            else:
                if self.config["skycut"] == 1:
                    self.projection = Projection(self.config["xdata"], Om_AP=self.config["Omega_m_AP"], z_AP=self.config["z_AP"],
                        window_fourier_name=self.config["windowPk"], path_to_window='', window_configspace_file=self.config["windowCf"],
                        binning=self.config["with_binning"], fibcol=self.config["with_fibercol"], Nwedges=self.config["wedge"], wedges_bounds=self.config["wedges_bounds"],
                        zz=self.config["zz"], nz=self.config["nz"], co=self.co)

                elif self.config["skycut"] > 1:

                    self.projection = []

                    for i in range(self.config["skycut"]):

                        if len(self.config["xdata"]) == 1: xdata = self.config["xdata"][i]
                        elif len(self.config["xdata"]) is self.config["skycut"]: xdata = self.config["xdata"][i]
                        else: xdata = self.config["xdata"]

                        if self.config["with_window"]:
                            windowPk = self.config["windowPk"][i]
                            windowCf = self.config["windowCf"][i]
                        else:
                            windowPk = None
                            windowCf = None

                        if self.config["with_AP"]:
                            if isinstance(self.config["Omega_m_AP"], float): Om_AP = self.config["Omega_m_AP"]
                            elif len(self.config["Omega_m_AP"]) is self.config["skycut"]: Om_AP = self.config["Omega_m_AP"][i]

                            if isinstance(self.config["z_AP"], float): z_AP = self.config["z_AP"]
                            elif len(self.config["z_AP"]) is self.config["skycut"]: z_AP = self.config["z_AP"][i]
                        else:
                            Om_AP = None
                            z_AP = None

                        if self.config["with_redshift_bin"]:
                            zz = self.config["zz"][i]
                            nz = self.config["nz"][i]
                        else:
                            zz = None
                            nz = None

                        self.projection.append( Projection(xdata, Om_AP=Om_AP, z_AP=z_AP,
                            window_fourier_name=windowPk, path_to_window='', window_configspace_file=windowCf,
                            binning=self.config["with_binning"], fibcol=self.config["with_fibercol"], Nwedges=self.config["wedge"], wedges_bounds=self.config["wedges_bounds"],
                            zz=zz, nz=nz,
                            co=self.co) )

    def __read_cosmo(self, cosmo_dict):

        # Checking if the inputs are consistent with the options
        for (name, cosmo) in zip(self.cosmo_catalog, self.cosmo_catalog.values()):
                for cosmo_key in cosmo_dict:
                    if cosmo_key is name:
                        cosmo.check(cosmo_key, cosmo_dict[cosmo_key])

        # Setting unspecified configs to default value
        for (name, cosmo) in zip(self.cosmo_catalog, self.cosmo_catalog.values()):
            if cosmo.value is None: cosmo.value = cosmo.default

        # Translating the catalog to a dict
        self.cosmo = translate_catalog_to_dict(self.cosmo_catalog)

    def __is_cosmo_conflict(self):

        if self.cosmo["k11"] is None or self.cosmo["P11"] is None:
            raise Exception("Please provide a linear matter power spectrum \'P11\' and the corresponding \'k11\'. ")

        if len(self.cosmo["k11"]) != len(self.cosmo["P11"]):
            raise Exception("Please provide a linear matter power spectrum \'P11\' and the corresponding \'k11\' of same length.")

        if self.cosmo["k11"][0] > 1e-4 or self.cosmo["k11"][-1] < 1.:
            raise Exception("Please provide a linear matter spectrum \'P11\' and the corresponding \'k11\' with min(k11) < 1e-4 and max(k11) > 1.")

        if self.config["skycut"] > 1:
            if self.cosmo["D"] is None:
                raise Exception("You asked multi skycuts. Please specify the growth function \'D\'. ")
            elif len(self.cosmo["D"]) is not self.config["skycut"]:
                raise Exception("Please specify (in a list) as many growth functions \'D\' as the corresponding skycuts.")

        if self.config["multipole"] == 0: self.cosmo["f"] = 0.
        elif not self.config["with_redshift_bin"]:
            if self.cosmo["f"] is None:
                raise Exception("Please specify the growth rate \'f\'.")
            if self.config["skycut"] == 1:
                if not isinstance(self.cosmo["f"], float):
                    raise Exception("Please provide a single growth rate \'f\'.")
            elif len(self.cosmo["f"]) != self.config["skycut"]:
                raise Exception("Please specify (in a list) as many \'f\' as the corresponding skycuts.")

        if self.config["wedge"] > 0:
            if self.config["wedges_bounds"] is not None:
                if len(self.config["wedges_bounds"]) != self.config["wedge"]+1 or self.config["wedges_bounds"][0] != 0 or self.config["wedges_bounds"][-1] != 1:
                    raise Exception("If specifying \'wedges_bounds\', specify them in a list as: [0, a_1, ..., a_{n-1}, 1], where n: number of wedges")

        if self.config["with_bias"]:
            self.__is_bias_conflict()

        if self.config["with_AP"] and not self.config["with_redshift_bin"]:
            if self.cosmo["DA"] is None or self.cosmo["H"] is None:
                raise Exception("You asked to apply the AP effect. Please specify \'DA\' and \'H\'. ")

            if self.config["skycut"] == 1:
                if not isinstance(self.cosmo["DA"], float) and not isinstance(self.cosmo["H"], float):
                    raise Exception("Please provide a single pair of \'DA\' and \'H\'.")
            elif len(self.cosmo["DA"]) is not self.config["skycut"] or len(self.cosmo["H"]) is not self.config["skycut"]:
                raise Exception("Please specify (in lists) as many \'DA\' and \'H\' as the corresponding skycuts.")

        if self.config["with_redshift_bin"]:
            if self.cosmo["Dz"] is None or self.cosmo["fz"] is None:
                raise Exception("You asked to account the galaxy counts distribution. Please specify \'Dz\' and \'fz\'. ")

            if self.config["skycut"] == 1:
                if len(self.cosmo["Dz"]) is not len(self.config["zz"]) or len(self.cosmo["fz"]) is not len(self.config["zz"]):
                    raise Exception("Please specify \'Dz\' and \'fz\' with same length as \'zz\'. ")
            elif len(self.cosmo["Dz"]) is not self.config["skycut"] or len(self.cosmo["fz"]) is not self.config["skycut"]:
                raise Exception("Please specify (in lists) as many \'Dz\' and \'fz\' as the corresponding skycuts.")

            if self.config["with_AP"]:
                if self.cosmo["DAz"] is None or self.cosmo["Hz"] is None:
                    raise Exception("You asked to account the galaxy counts distribution and apply the AP effect. Please specify \'DAz\' and \'Hz\'. ")

                if self.config["skycut"] == 1:
                    if len(self.cosmo["DAz"]) is not len(self.config["zz"]) or len(self.cosmo["Hz"]) is not len(self.config["zz"]):
                        raise Exception("Please specify \'DAz\' and \'Hz\' with same length as \'zz\'. ")
                elif len(self.cosmo["DAz"]) is not self.config["skycut"] or len(self.cosmo["Hz"]) is not self.config["skycut"]:
                    raise Exception("Please specify (in lists) as many \'DAz\' and \'Hz\' as the corresponding skycuts.")

            if "w" in self.config["output"]:
                if self.cosmo["rz"] is None:
                    raise Exception("You asked angular statistics. Please specify \'rz\'. ")

                if self.config["skycut"] == 1:
                    if len(self.cosmo["rz"]) is not len(self.config["zz"]):
                        raise Exception("Please specify \'rz\' with same length as \'zz\'. ")
                elif len(self.cosmo["rz"]) is not self.config["skycut"]:
                    raise Exception("Please specify (in a list) as many \'rz\'as the corresponding skycuts.")

    def __is_bias_conflict(self, bias=None):

        ###raise Exception("Input error in \'%s\'; input configs: %s. Check Correlator.info() in any doubt." % ())

        if bias is not None: self.cosmo["bias"] = bias

        if self.cosmo["bias"] is None: raise Exception("Please specify \'bias\'. ")
        if isinstance(self.cosmo["bias"], (list, np.ndarray)): self.cosmo["bias"] = self.cosmo["bias"][0]
        if not isinstance(self.cosmo["bias"], dict): raise Exception("Please specify bias in a dict. ")

        if "bm" in self.config["output"]:
            if not self.config["with_stoch"]:
                if self.config["multipole"] == 0:
                    if len(self.cosmo["bias"]) is not 6: raise Exception("Please specify a dict of 6 biases: \{ \'b1\', \'b2\', \'b3\', \'b4\', \'cct\' + matter counterterm: \'dct\' \}. ")
                    else: self.bias = { "b1": self.cosmo["bias"]["b1"], "b2": self.cosmo["bias"]["b2"], "b3": self.cosmo["bias"]["b3"], "b4": self.cosmo["bias"]["b4"], "cct": self.cosmo["bias"]["cct"], "cr1": 0., "cr2": 0., "ce0": 0., "ce1": 0., "ce2": 0., "dct": self.cosmo["bias"]["dct"], "dr1": 0., "dr2": 0. }
                elif self.config["multipole"] == 2:
                    if len(self.cosmo["bias"]) is not 8: raise Exception("Please specify a dict of 8 biases: \{ \'b1\', \'b2\', \'b3\', \'b4\', \'cct\', \'cr1\' + matter counterterms: \'dct\', \'dr1\' \}. ")
                    else: self.bias = { "b1": self.cosmo["bias"]["b1"], "b2": self.cosmo["bias"]["b2"], "b3": self.cosmo["bias"]["b3"], "b4": self.cosmo["bias"]["b4"], "cct": self.cosmo["bias"]["cct"], "cr1": self.cosmo["bias"]["cr1"], "cr2": 0., "ce0": 0., "ce1": 0., "ce2": 0., "dct": self.cosmo["bias"]["dct"], "dr1": self.cosmo["bias"]["dr1"], "dr2": 0. }
                elif self.config["multipole"] == 3:
                    if len(self.cosmo["bias"]) is not 10: raise Exception("Please specify a dict of 10 biases: \{ \'b1\', \'b2\', \'b3\', \'b4\', \'cct\', \'cr1\', \'cr2\' + matter counterterms: \'dct\', \'dr1\', \'dr2\' \}. ")
                    else: self.bias = { "b1": self.cosmo["bias"]["b1"], "b2": self.cosmo["bias"]["b2"], "b3": self.cosmo["bias"]["b3"], "b4": self.cosmo["bias"]["b4"], "cct": self.cosmo["bias"]["cct"], "cr1": self.cosmo["bias"]["cr1"], "cr2": self.cosmo["bias"]["cr2"], "ce0": 0., "ce1": 0., "ce2": 0., "dct": self.cosmo["bias"]["dct"], "dr1": self.cosmo["bias"]["dr1"], "dr2": self.cosmo["bias"]["dr2"] }
            else:
                pass # to code up

        elif "m" in self.config["output"]:
            if self.config["multipole"] == 0:
                if len(self.cosmo["bias"]) is not 1: raise Exception("Please specify a dict of 1 bias: \{ \'cct\' \}. ")
                else: self.bias = { "b1": 1., "b2": 1., "b3": 1., "b4": 0., "cct": self.cosmo["bias"]["cct"], "cr1": 0., "cr2": 0., "ce0": 0., "ce1": 0., "ce2": 0. }
            elif self.config["multipole"] == 2:
                if len(self.cosmo["bias"]) is not 2: raise Exception("Please specify a dict of 2 biases: \{ \'cct\', \'cr1\' \}. ")
                else: self.bias = { "b1": 1., "b2": 1., "b3": 1., "b4": 0., "cct": self.cosmo["bias"]["cct"], "cr1": self.cosmo["bias"]["cr1"], "cr2": 0., "ce0": 0., "ce1": 0., "ce2": 0. }
            elif self.config["multipole"] == 3:
                if len(self.cosmo["bias"]) is not 3: raise Exception("Please specify a dict of 3 biases: \{ \'cct\', \'cr1\', \'cr2\' \}. ")
                else: self.bias = { "b1": 1., "b2": 1., "b3": 1., "b4": 0., "cct": self.cosmo["bias"]["cct"], "cr1": self.cosmo["bias"]["cr1"], "cr2": self.cosmo["bias"]["cr2"], "ce0": 0., "ce1": 0., "ce2": 0. }

        else:

            Nextra = 0
            if self.config["with_nlo_bias"]: Nextra += 1
            if self.config["with_tidal_alignments"]: Nextra += 1

            if not self.config["with_stoch"]:
                if self.config["multipole"] == 0 or "w" in self.config["output"]:
                    if len(self.cosmo["bias"]) is not 5+Nextra: raise Exception("Please specify a dict of 5 biases: \{ \'b1\', \'b2\', \'b3\', \'b4\', \'cct\' \}. ")
                    else: self.bias = { "b1": self.cosmo["bias"]["b1"], "b2": self.cosmo["bias"]["b2"], "b3": self.cosmo["bias"]["b3"], "b4": self.cosmo["bias"]["b4"], "cct": self.cosmo["bias"]["cct"], "cr1": 0., "cr2": 0., "ce0": 0., "ce1": 0., "ce2": 0. }
                elif self.config["multipole"] == 2:
                    if len(self.cosmo["bias"]) is not 6+Nextra: raise Exception("Please specify a dict of 6 biases: \{ \'b1\', \'b2\', \'b3\', \'b4\', \'cct\', \'cr1\' \}. ")
                    else: self.bias = { "b1": self.cosmo["bias"]["b1"], "b2": self.cosmo["bias"]["b2"], "b3": self.cosmo["bias"]["b3"], "b4": self.cosmo["bias"]["b4"], "cct": self.cosmo["bias"]["cct"], "cr1": self.cosmo["bias"]["cr1"], "cr2": 0., "ce0": 0., "ce1": 0., "ce2": 0. }
                elif self.config["multipole"] == 3:
                    if len(self.cosmo["bias"]) is not 7+Nextra: raise Exception("Please specify a dict of 7 biases: \{ \'b1\', \'b2\', \'b3\', \'b4\', \'cct\', \'cr1\', \'cr2\' \}. ")
                    else: self.bias = { "b1": self.cosmo["bias"]["b1"], "b2": self.cosmo["bias"]["b2"], "b3": self.cosmo["bias"]["b3"], "b4": self.cosmo["bias"]["b4"], "cct": self.cosmo["bias"]["cct"], "cr1": self.cosmo["bias"]["cr1"], "cr2": self.cosmo["bias"]["cr2"], "ce0": 0., "ce1": 0., "ce2": 0. }
            else:
                if self.config["multipole"] == 0:
                    if len(self.cosmo["bias"]) is not 6+Nextra: raise Exception("Please specify a dict of 6 biases: \{ \'b1\', \'b2\', \'b3\', \'b4\', \'cct\', \'ce0\' \}. ")
                    else: self.bias = { "b1": self.cosmo["bias"]["b1"], "b2": self.cosmo["bias"]["b2"], "b3": self.cosmo["bias"]["b3"], "b4": self.cosmo["bias"]["b4"], "cct": self.cosmo["bias"]["cct"], "cr1": 0., "cr2": 0., "ce0": self.cosmo["bias"]["ce0"], "ce1": 0., "ce2": 0. }
                elif self.config["multipole"] == 2:
                    if len(self.cosmo["bias"]) is not 9+Nextra: raise Exception("Please specify a dict of 9 biases: \{ \'b1\', \'b2\', \'b3\', \'b4\', \'cct\', \'cr1\', \'ce0\', \'ce1\', \'ce2\'  \}. ")
                    else: self.bias = { "b1": self.cosmo["bias"]["b1"], "b2": self.cosmo["bias"]["b2"], "b3": self.cosmo["bias"]["b3"], "b4": self.cosmo["bias"]["b4"], "cct": self.cosmo["bias"]["cct"], "cr1": self.cosmo["bias"]["cr1"], "cr2": 0., "ce0": self.cosmo["bias"]["ce0"], "ce1": self.cosmo["bias"]["ce1"], "ce2": self.cosmo["bias"]["ce2"] }
                elif self.config["multipole"] == 3:
                    if len(self.cosmo["bias"]) is not 10+Nextra: raise Exception("Please specify a dict of 10 biases: \{ \'b1\', \'b2\', \'b3\', \'b4\', \'cct\', \'cr1\', \'cr2\', \'ce0\', \'ce1\', \'ce2\' \}. ")
                    else: self.bias = { "b1": self.cosmo["bias"]["b1"], "b2": self.cosmo["bias"]["b2"], "b3": self.cosmo["bias"]["b3"], "b4": self.cosmo["bias"]["b4"], "cct": self.cosmo["bias"]["cct"], "cr1": self.cosmo["bias"]["cr1"], "cr2": self.cosmo["bias"]["cr2"], "ce0": self.cosmo["bias"]["ce0"], "ce1": self.cosmo["bias"]["ce1"], "ce2": self.cosmo["bias"]["ce2"] }

            if self.config["with_nlo_bias"]:
                try: self.bias["bnlo"] = self.cosmo["bias"]["bnlo"]
                except: raise Exception ("Please specify the next-to-leading bias \'bnlo\'.  ")

            if self.config["with_tidal_alignments"]:
                try: self.bias["bq"] = self.cosmo["bias"]["bq"]
                except: raise Exception ("Please specify the tidal alignments bias \'bq\'.  ")

    def __read_config(self, config_dict):

        # Checking if the inputs are consistent with the options
        for (name, config) in zip(self.config_catalog, self.config_catalog.values()):
                for config_key in config_dict:
                    if config_key == name:
                        config.check(config_key, config_dict[config_key])

        # Setting unspecified configs to default value
        for (name, config) in zip(self.config_catalog, self.config_catalog.values()):
            if config.value is None: config.value = config.default

        # Translating the catalog to a dict
        self.config = translate_catalog_to_dict(self.config_catalog)

        self.config["accboost"] = float(self.config["accboost"])

    def __is_config_conflict(self):

        if "Cf" in self.config["output"]:
            self.config["with_window"] = False
            self.config["with_cf"] = True
        elif "w" in self.config["output"]:
            self.config["multipole"] = 3
            self.config["wedge"] = 0
            self.config["with_cf"] = True
            self.config["with_bias"] = False
            self.config["with_redshift_bin"] = True
            self.config["with_AP"] = False
            self.config["with_window"] = False
            self.config["with_stoch"] = False
        else:
            self.config["with_cf"] = False

        # if self.config["with_cf"]: self.config["with_stoch"] = False
        # if self.config["wedge"] is not 0: self.config["multipole"] = 3 # enforced

        if "bm" in self.config["output"]: self.config["halohalo"] = False
        else: self.config["halohalo"] = True

        if self.config["skycut"] > 1:
            self.config["with_bias"] = False
            if len(self.config["z"]) != self.config["skycut"]:
                raise Exception("Please specify as many redshifts \'z\' as the number of skycuts.")
                self.config["z"] = np.asarray(self.config["z"])
            if "w" not in self.config["output"]: ### TO CHANGE
                def checkEqual(lst): return lst[1:] == lst[:-1]
                if checkEqual(self.config["z"]): # if same redshift
                    self.config["with_time"] = True
                    #self.config["z"] = self.config["z"][0]
                else: self.config["with_time"] = False
            elif "w" in self.config["output"]:
                if isinstance(self.config["w_theta_cut"], (float, int)): self.config["w_theta_cut"] = self.config["w_theta_cut"] * np.ones(shape=(self.config["skycut"]))

        if self.config["xdata"] is None:
            raise Exception("Please specify a data point array \'xdata\'.")
        if len(self.config["xdata"]) == 1 and isinstance(self.config["xdata"][0], (list, np.ndarray)):
            self.config["xdata"] = self.config["xdata"][0]
        # else:
        #     self.config["xdata"] = np.asarray(self.config["xdata"])

        #     def is_conflict_xdata(xdata):
        #         if "w" in self.config["output"]:
        #             pass
        #         elif "Cf" in self.config["output"]:
        #             if xdata[0] < self.config["smin"] or xdata[-1] > self.config["smax"]:
        #                 raise Exception("Please specify a data point array \'xdata\' in: (%s, %s)." % (self.config["smin"], self.config["smax"]))
        #         else:
        #             if xdata[0] < self.config["kmin"] or xdata[-1] > self.config["kmax"]:
        #                 raise Exception("Please specify a data point array \'xdata\' in: (%s, %s) or increase the kmax." % (self.config["kmin"], self.config["kmax"]))

        #     if self.config["skycut"] == 1:
        #         is_conflict_xdata(self.config["xdata"])
        #     elif self.config["skycut"] > 1:
        #         if len(self.config["xdata"]) == 1: is_conflict_xdata(self.config["xdata"])
        #         if len(self.config["xdata"]) is self.config["skycut"]:
        #             for xi in self.config["xdata"]: is_conflict_xdata(xi)
        #         else:
        #             raise Exception("Please provide a commmon data point array \'xdata\' or as many arrays (in a list) as the corresponding skycuts.")

        if self.config["with_AP"]:
            if self.config["Omega_m_AP"] is None: raise Exception("You asked to apply the AP effect. Please specify \'Omega_m_AP\'.")
            if self.config["z_AP"] is None: self.config["z_AP"] = self.config["z"]
            if self.config["skycut"] == 1:
                if isinstance(self.config["z_AP"], list): self.config["z_AP"] = self.config["z_AP"][0]

        if self.config["with_window"]:

            def is_conflict_window(windowCf):
                try:
                    test = np.loadtxt(windowCf)
                    if self.config["with_cf"]: self.windowPk = None
                except IOError:
                    print("You asked to apply a mask. Please specify a correct path to the configuration space window file.")
                    raise

            if self.config["skycut"] == 1:
                if isinstance(self.config["windowCf"], list): self.config["windowCf"] = self.config["windowCf"][0]
                if isinstance(self.config["windowPk"], list): self.config["windowPk"] = self.config["windowPk"][0]
                is_conflict_window(self.config["windowCf"])
            elif self.config["skycut"] > 1:
                for windowCf in self.config["windowCf"]:
                    if windowCf is not None:
                        is_conflict_window(windowCf)
        else:
            self.config["windowPk"] = None
            self.config["windowCf"] = None

        if self.config["with_redshift_bin"]:

            self.config["with_bias"] = False
            self.config["with_time"] = False

            def is_conflict_zz(zz, nz):
                if zz is None or nz is None:
                    pass#raise Exception("You asked to account for the galaxy counts distribution over a redshift bins. Please provide a distribution \'nz\' and corresponding \'zz\'. ")
                elif len(zz) != len(nz):
                    raise Exception("Please provide \'nz\' and corresponding \'zz\' of the same length. ")

            if self.config["skycut"] == 1: is_conflict_zz(self.config["zz"], self.config["nz"])
            elif self.config["skycut"] > 1:
                self.config["zz"] = np.asarray(self.config["zz"])
                self.config["nz"] = np.asarray(self.config["nz"])
                # print (len(self.config["zz"]), len(self.config["nz"]))
                if len(self.config["zz"]) == self.config["skycut"] and len(self.config["nz"]) == self.config["skycut"]:
                    for zz, nz in zip(self.config["zz"], self.config["nz"]): is_conflict_zz(zz, nz)
                else:
                    raise Exception("Please provide as many \'nz\' with corresponding \'zz\' (in a list) as the corresponding skycuts. ")
        else:
            self.config["zz"] = None
            self.config["nz"] = None

        if self.config["with_quintessence"]: self.config["with_exact_time"] = True
        if self.config["with_MG"]: self.config["with_exact_time"] = True
        if self.config["with_nDGP"]:
            self.config["with_exact_time"] = True #needed?
            self.config["with_MG"] = True

    def setcosmo(self, cosmo_dict, module='class'):

        if module is 'class':

            from classy import Class

            # Not sure this is useful: does class read z_max_pk?
            if self.config["skycut"] == 1:
                if self.config["with_redshift_bin"]: zmax = max(self.config["zz"])
                else: zmax = self.config["z"]
            elif self.config["skycut"] > 1:
                if self.config["with_redshift_bin"]:
                    maxbin = np.argmax(self.config["z"])
                    zmax = max(self.config["zz"][maxbin])
                else: zmax = max(self.config["z"])

            cosmo_dict_local = cosmo_dict.copy()
            if self.config["with_bias"]: del cosmo_dict_local["bias"] # Class does not like dictionary with keys other than the ones it reads...

            M = Class()
            M.set(cosmo_dict_local)
            M.set({'output': 'mPk', 'P_k_max_h/Mpc': 1.0, 'z_max_pk': zmax })
            M.compute()

            cosmo = {}

            if self.config["with_bias"]:
                try: cosmo["bias"] = cosmo_dict["bias"]
                except:
                    print ("Please specify \'bias\'.")
                    raise

            if self.config["skycut"] == 1: zfid = self.config["z"]
            elif self.config["skycut"] > 1: zfid = self.config["z"][0]

            cosmo["k11"] = np.logspace(-5, 0, 200) # k in h/Mpc
            cosmo["P11"] = np.array([M.pk(k*M.h(), zfid)*M.h()**3 for k in cosmo["k11"]]) # P(k) in (Mpc/h)**3

            if self.config["skycut"] == 1:
                if self.config["multipole"] is not 0: cosmo["f"] = M.scale_independent_growth_factor_f(self.config["z"])
                if self.config["with_exact_time"]:
                    cosmo["z"] = self.config["z"]
                    cosmo["Omega0_m"] = M.Omega0_m()
                    try: cosmo["w0_fld"] = cosmo_dict["w0_fld"]
                    except: pass
                if self.config["with_AP"]:
                    cosmo["DA"] = M.angular_distance(self.config["z"]) * M.Hubble(0.)
                    cosmo["H"] = M.Hubble(self.config["z"]) / M.Hubble(0.)

            elif self.config["skycut"] > 1:
                if self.config["multipole"] is not 0: cosmo["f"] = np.array([M.scale_independent_growth_factor_f(z) for z in self.config["z"]])
                cosmo["D"] = np.array([M.scale_independent_growth_factor(z) for z in self.config["z"]])
                if self.config["with_AP"] and not self.config["with_redshift_bin"]:
                    cosmo["DA"] = np.array([M.angular_distance(z) * M.Hubble(0.) for z in self.config["z"]])
                    cosmo["H"] = np.array([M.Hubble(z) / M.Hubble(0.) for z in self.config["z"]])

            if self.config["with_redshift_bin"]:
                if self.config["skycut"] == 1:
                    cosmo["D"] = M.scale_independent_growth_factor(self.config["z"])

                    cosmo["Dz"] = np.array([M.scale_independent_growth_factor(z) for z in self.config["zz"]])
                    cosmo["fz"] = np.array([M.scale_independent_growth_factor_f(z) for z in self.config["zz"]])

                    cosmo["DAz"] = np.array([M.angular_distance(z) * M.Hubble(0.) for z in self.config["zz"]])
                    cosmo["Hz"] = np.array([M.Hubble(z) / M.Hubble(0.) for z in self.config["zz"]])

                elif self.config["skycut"] > 1:
                    cosmo["Dz"] = np.array([ [M.scale_independent_growth_factor(z) for z in zz] for zz in self.config["zz"] ])
                    cosmo["fz"] = np.array([ [M.scale_independent_growth_factor_f(z) for z in zz] for zz in self.config["zz"] ])

                    cosmo["DAz"] = np.array([ [M.angular_distance(z) * M.Hubble(0.) for z in zz] for zz in self.config["zz"] ])
                    cosmo["Hz"] = np.array([ [M.Hubble(z) / M.Hubble(0.)  for z in zz] for zz in self.config["zz"] ])

            if "w" in self.config["output"]:
                def comoving_distance(z): return M.angular_distance(z)*(1+z)*M.h()
                if self.config["skycut"] == 1: cosmo["rz"] = np.array([comoving_distance(z) for z in self.config["zz"]])
                elif self.config["skycut"] > 1: cosmo["rz"] = np.array([ [comoving_distance(z) for z in zz] for zz in self.config["zz"] ])

            if self.config["with_nDGP"]:
                #k_MGCAMB = np.loadtxt('Ben/mgcamb/matterpower_z0.dat')[:,0]
                #P_lin_MGCAMB = np.loadtxt('Ben/mgcamb/matterpower_z0.dat')[:,1]
                #PlinB =interpolate.interp1d(k_MGCAMB,P_lin_MGCAMB,kind = 'cubic')
                zm = 5 # z in matter domination
                def scale_factor(z): return 1/(1.+z)
                Omega0_m = cosmo["Omega0_m"]
                Om_rc = self.config["Omega_rc"]
                GF = GreenFunction(Omega0_m, MG = True, nDGP=True, Omega_rc = Om_rc)
                Dp = GF.D(scale_factor(zfid))/GF.D(scale_factor(zm))
                Dm = GF.Dminus(scale_factor(zfid))/GF.Dminus(scale_factor(zm)) #needed?
                D_class = M.scale_independent_growth_factor(zfid) / M.scale_independent_growth_factor(zm)
                D_lcdm = GF.D_LCDM(scale_factor(zfid))/GF.D(scale_factor(zm))
                cosmo["P11"] *= Dp**2 / D_lcdm**2 #GF.D(scale_factor(zfid))**2/M.scale_independent_growth_factor(zfid)**2 #
                #cosmo["k11"] = k_MGCAMB
                #cosmo["P11"] = PlinB(cosmo["k11"]) * (Dp**2 / D_class**2)
                cosmo["D"] = GF.D(scale_factor(zfid))/GF.D(scale_factor(0))
                #if self.config["multipole"] is not 0: cosmo["f"] = GF.fplus(scale_factor(zfid))
                if self.config["skycut"] == 1:
                    if self.config["multipole"] != 0: cosmo["f"] = GF.fplus(scale_factor(self.config["z"]))

                elif self.config["skycut"] > 1:
                    if self.config["multipole"] != 0: cosmo["f"] = np.array([GF.fplus(scale_factor(z)) for z in self.config["z"]])

            if self.config["with_quintessence"]:
                # starting deep inside matter domination and evolving to the total adiabatic linear power spectrum.
                # This does not work in the general case, e.g. with massive neutrinos (okish for minimal mass though)
                # This does not work for multi skycuts nor for redshift bins.
                zm = 5. # z in matter domination
                def scale_factor(z): return 1/(1.+z)
                Omega0_m = cosmo["Omega0_m"]
                w = cosmo["w0_fld"]
                GF = GreenFunction(Omega0_m, w=w, quintessence=True)
                Dq = GF.D(scale_factor(zfid)) / GF.D(scale_factor(zm))
                Dm = M.scale_independent_growth_factor(zfid) / M.scale_independent_growth_factor(zm)
                cosmo["P11"] *= Dq**2 / Dm**2 * ( 1 + (1+w)/(1.-3*w) * (1-Omega0_m)/Omega0_m * (1+zm)**(3*w) ) # 1611.07966 eq. (4.15)

            return cosmo

class BiasCorrelator(Correlator):
    '''
    Class to load pre-computed correlator
    '''
    def __init__(self, config_dict=None, load_engines=False):
        Correlator.__init__(self, config_dict, load_engines=load_engines)

def translate_catalog_to_dict(catalog):
    newdict = dict.fromkeys(catalog)
    for key, option in zip(catalog, catalog.values()):
        newdict[key] = option.value
    return newdict

def typename(onetype):
    if isinstance(onetype, tuple): return [t.__name__ for t in onetype]
    else: return [onetype.__name__]

class Option(object):

    def __init__(self, config_name, config_type, config_list=None, description='', default=None, verbose=False):

        self.verbose = verbose
        self.name = config_name
        self.type = config_type
        self.list = config_list
        self.description = description
        self.default = default
        self.value = None

    def check(self, config_key, config_value):
        is_config = False
        if self.verbose: print("\'%s\': \'%s\'" % (config_key, config_value))
        if isinstance(config_value, self.type):
            if self.list is None: is_config = True
            elif isinstance(config_value, str):
                if any(config_value in o for o in self.list): is_config = True
            elif isinstance(config_value, (int, float, bool)):
                if any(config_value == o for o in self.list): is_config = True
        if is_config:
            self.value = config_value
        else: self.error()
        return is_config

    def error(self):
        if self.list is None:
            try: raise Exception("Input error in \'%s\'; input configs: %s. Check Correlator.info() in any doubt." % (self.name, typename(self.type)))
            except Exception as e: print(e)
        else:
            try: raise Exception("Input error in \'%s\'; input configs: %s. Check Correlator.info() in any doubt." % (self.name, self.list))
            except Exception as e: print(e)


import torch
import torch.nn as nn
import torch.nn.functional as F
from .module import (DecoderPrenet, EmotionPrenet, Postnet, PositionEmbedding,
                     TransformerDecoderGuided, TransformerDecoderLayer,
                     TransformerEncoder, TransformerEncoderLayer, EmoDecoderLayer, PairedTransformerDecoderGuided, PairedDecoderLayer,
                     PairedTransformerEncoderLayer, PairedTransformerEncoder)
from g2p_en import G2p
import random
import numpy as np

class ETTSTransformer(nn.Module):
    def __init__(self, text_embed_dim, emo_embed_dim, nmels, maxlength, symbols, ngst, nlst,
                 d_model=512, hidden_size=1536, nlayers=4, nheads=4):
        super().__init__()
        self.d_model = d_model
        self.nheads = nheads
        self.text_embedding = nn.Embedding(len(symbols.keys()), text_embed_dim, padding_idx=symbols['<pad>'])
        self.global_style_embedding = nn.Embedding(ngst, d_model)
        self.local_style_embedding = nn.Embedding(nlst, d_model)
        self.encoder = PairedTransformerEncoder(
            [PairedTransformerEncoderLayer(d_model, nheads, hidden_size) for i in range(nlayers)]
        )
        self.decoder = TransformerDecoderGuided(
            [TransformerDecoderLayer(d_model, nheads, hidden_size) for i in range(nlayers)]
        )
        self.emotion_global_prenet = EmotionPrenet(emo_embed_dim, hidden_size=d_model, dropout=0.1)
        self.emotion_local_prenet = nn.LSTM(emo_embed_dim, d_model, 1, batch_first=True)
        self.emotion_local_attn = nn.MultiheadAttention(d_model, num_heads=4, batch_first=True)
        self.emotion_global_attn = nn.MultiheadAttention(d_model, num_heads=4, batch_first=True)
        self.text_norm = nn.LayerNorm(text_embed_dim)
        self.decoder_prenet = DecoderPrenet(nmels, bottleneck_size=d_model//8, hidden_size=d_model)
        self.mel_linear = nn.Linear(d_model, nmels)
        self.postnet = Postnet(nmels, hidden_size=d_model, dropout=0.1)
        self.stop_pred = nn.Linear(d_model, 1)

        self.pos_txt = PositionEmbedding(maxlength, d_model, scaled=False)
        self.pos_mel = PositionEmbedding(maxlength, d_model, scaled=False)
        self.proj = nn.Linear(emo_embed_dim, d_model)
        tgt_mask = (torch.tril(torch.ones(maxlength, maxlength), diagonal=0) == 0)
        self.register_buffer('tgt_mask', tgt_mask)

        self.maxlength = maxlength
        self.nmels = nmels
#        self.acc = []

    def length2mask(self, x, length):
        mask = torch.zeros(
            x.shape[:2], dtype=x.dtype, device=x.device
        )
        mask[
            (torch.arange(mask.shape[0], device=x.device), length - 1)
        ] = 1
        mask = mask.flip([-1]).cumsum(-1).flip([-1]).bool()
        return ~mask

    def forward(self, emotion_global, emotion_local, text, mels, length_emo_global, length_emo_local, length_text, length_mels):
        text = self.pos_txt(self.text_norm(self.text_embedding(text)))
        txt_mask = self.length2mask(text, length_text)
        text = text.transpose(0, 1) #T, B, C

        if emotion_global is not None:
            emo_mask_global = self.length2mask(emotion_global, length_emo_global)
            emotion_global = self.emotion_global_prenet(emotion_global, mask=emo_mask_global) #(B,T,C)
            keys = torch.tanh(self.global_style_embedding.weight).unsqueeze(0).expand(emotion_global.size(0), -1, -1)
            emotion_global = self.emotion_global_attn(emotion_global, keys, keys)[0].transpose(0, 1)

        emotion_local, _ = self.emotion_local_prenet(emotion_local)
        keys = torch.tanh(self.local_style_embedding.weight).unsqueeze(0).expand(emotion_local.size(0), -1, -1)
        emotion_local = self.emotion_local_attn(emotion_local, keys, keys)[0]
        stride, ksize = 4, 8
        min_crops = 15
        emotion_local = F.avg_pool1d(emotion_local.transpose(1, 2), kernel_size=ksize, stride=stride).permute(2, 0, 1)
        crop_size = random.randint(min_crops, max(emotion_local.size(0), min_crops))
        emotion_local = emotion_local[: crop_size]
        d_length_emo = torch.minimum(length_emo_local // stride + 1, torch.zeros_like(length_emo_local) + emotion_local.size(0))
        emo_mask2 = self.length2mask(emotion_local.transpose(0, 1), d_length_emo)

        text, enc_gloss, enc_attn = self.encoder(text, global_emotion=emotion_global, local_emotion=emotion_local,
                                                 src_key_padding_mask=txt_mask, emotion_key_padding_mask=emo_mask2)
        src, src_padding_mask = text, txt_mask

        mels = self.pos_mel(self.decoder_prenet(mels, dropout=0.2))
        mel_padding_mask = self.length2mask(mels, length_mels)
        mels = mels.transpose(0, 1) #T, B, C
        mels = emotion_global.expand(mels.size(0), -1, -1) + mels
        mels = torch.cat([torch.zeros((1, mels.size(1), mels.size(2)), device=mels.device), mels[: -1]], 0) #shift
        tgt_len = mels.size(0)
        tgt_mask = self.tgt_mask[: tgt_len, : tgt_len]

        output, ed_gloss, attn, dec_attn = self.decoder(mels, src,
                                                        tgt_mask=tgt_mask,
                                                        tgt_key_padding_mask=mel_padding_mask,
                                                        text_key_padding_mask=src_padding_mask) #TBC
        gloss = ed_gloss + enc_gloss
        gate_logit = self.stop_pred(output).squeeze(-1).transpose(0, 1)
        output = output.transpose(0, 1) #B, T, C
        mel_out = self.mel_linear(output)
        mel_out_post = mel_out.detach() + self.postnet(mel_out.detach())

        return mel_out, mel_out_post, gate_logit, mel_padding_mask, gloss, attn, enc_attn, dec_attn

    def inference(self, emotion_global, emotion_local, text, maxlen=1000, threshold=.5, ref_wav=True):
        text = self.pos_txt(self.text_norm(self.text_embedding(text)))
        text = text.transpose(0, 1) #T, B, C
        mels = torch.zeros((1, 1, self.nmels), device=text.device) #B, T, C, B=1
        attn_mask = torch.zeros((self.nheads, 1, text.size(0)), device=text.device)
        keys = torch.tanh(self.global_style_embedding.weight).unsqueeze(0)
        emotion_global = self.emotion_global_prenet(emotion_global)
        emotion_global = self.emotion_global_attn(emotion_global, keys, keys)
        emotion_global = emotion_global[0].transpose(0, 1)
        keys = torch.tanh(self.local_style_embedding.weight).unsqueeze(0)
        if ref_wav:
            emotion_local, _ = self.emotion_local_prenet(emotion_local)
            emotion_local = self.emotion_local_attn(emotion_local, keys, keys)[0]
        else:
            emotion_local = torch.matmul(emotion_local, keys)
        emotion_local = F.avg_pool1d(emotion_local.transpose(1, 2), kernel_size=8, stride=4).permute(2, 0, 1)
        text, _, _ = self.encoder(text, global_emotion=emotion_global, local_emotion=emotion_local)
        src = text
        start_token = torch.zeros((1, 1, self.d_model), device=text.device)
        #Decode
        for i in range(maxlen):
            if i == 0:
                mels_pos = start_token
            else:
                mels_pos = self.pos_mel(self.decoder_prenet(mels[:, 1:], dropout=0.2)).transpose(0, 1)
                mels_pos = emotion_global.expand(mels_pos.size(0), -1, -1) + mels_pos
                mels_pos = torch.cat([start_token, mels_pos], 0)
            tgt_mask = self.tgt_mask[: mels_pos.size(0), : mels_pos.size(0)]
            output, _, attns, _ = self.decoder(mels_pos, src, tgt_mask=tgt_mask)#, memory_mask=attn_mask)
            gate_logit = self.stop_pred(output[-1]).item()
            output = output[-1] #B=1, C
            mel_out = self.mel_linear(output)
            mels = torch.cat([mels, mel_out.unsqueeze(1)], 1)
            if gate_logit > threshold:
                break
        mels = mels[:, 1:]
        mels = mels + self.postnet(mels)
        return mels, attns


from __future__ import division
import smtplib
import os
import calendar
import time
import numpy as np
from astropy.io import ascii
import glob
import subprocess

def findPhone(mac_address):
    with open('../../database.csv') as dataBase:
        for line in dataBase:
	    column = line.split(', ')
	    if mac_address in line:
		return (column[2], column[1])

def writeFile():
    if os.path.isfile("./sent_list.dat"):
        print "There is a SENT file that exists..."
    else:
        ascii.write([ np.array([]), np.array([])], names=['mac','sent_time'], output='./sent_list.dat')

restricted_list = np.array(["00:26:75:BB:E9:40","F4:F2:6D:0E:D7:74","BC:6C:21:02:FC:49","A9:8E:24:68:98:80","69:94:23:A4:6F:23","8D:3A:E3:45:88:C7","88:53:2E:B3:CC:98", "48:E9:F1:2A:C1:E3", "A0:18:28:2A:0F:5C"])

def mainCommand():
    while True:
        check_sent = glob.glob("./sent_list.*")
        if check_sent!=[]:
            sent_already = ascii.read("./sent_list.dat")
        else:
            sent_already = np.array([])

        if (os.path.isfile("./trigger.dat")):
            trigger_input = ascii.read("./trigger.dat", guess=False)
            if len(trigger_input['mac'])>0:
		additionalCondition = 1
		time.sleep(1)
                if len(sent_already["mac"])!=0:
		  try:
                    trig_unique_1 = np.array([x if ((x not in sent_already['mac'])|( (x in sent_already['mac'])and(np.max(sent_already['sent_time'][sent_already['mac']==x]) < calendar.timegm(time.gmtime()) - 10*60))) else "-99" for x in trigger_input['mac']])
                    trig_unique_2 = ((trig_unique_1[trig_unique_1!="-99"]))
                    print "Success"
		    print trig_unique_2
		    additionalCondition = 0
		  except Exception,e:
		    print str(e)
		    trig_unique_2 = np.array([])
		    print trig_unique_2
		    print "Pass3-Failed"
		   # print trig_unique_2
                else:
                    trig_unique_1 = trigger_input['mac']
                    trig_unique_2 = trig_unique_1
                    print "Initializing"
		    additionalCondition = 0

                if len(trig_unique_2)>0:
                    outbound_macs = np.array([])
                    outbound_time = np.array([])

                    for x in trig_unique_2:
                        if x not in restricted_list and additionalCondition == 0:
                            outbound_macs = np.append(outbound_macs, x)
                            outbound_time = np.append(outbound_time, (calendar.timegm(time.gmtime())) )
                            print_message = "Mac:%s, Last Visited:%s, Time Spent:%s minutes \n"%(x, time.strftime('%Y-%m-%d %H:%M:%S', time.gmtime(trigger_input['last_seen'][trigger_input['mac']==x]+ 8*60*60 )) , (trigger_input['time_spent'][trigger_input['mac']==x]))
                            print print_message
			    time.sleep(1)
			    namewithNumber = findPhone(x)
			    if namewithNumber is None:
				whatsapp_message = './yowsup-cli demos -c file.config -M -s 6584983348 "%s"' % print_message
			    else:
			        whatsapp_message = './yowsup-cli demos -c file.config -M -s 65%s "%s, have you registered for Balik Kampung?"' % namewithNumber
			    subprocess.call(whatsapp_message, shell=True)

                    if len(sent_already)==0:
                        ascii.write([outbound_macs, outbound_time], names=['mac','sent_time'], output='./sent_list.dat')
                    else:
                        ascii.write([np.append(sent_already['mac'], outbound_macs), np.append(sent_already['sent_time'],outbound_time)], names=['mac','sent_time'], output='./sent_list.dat')

while True:
    try:
        writeFile()
        mainCommand()
    except ValueError:
	pass


import os
from typing import Any, List, Optional, Tuple, Union

import jax
import jax.numpy as jnp
import flax
import numpy as np
from tqdm import tqdm

from .config import Config
from .impl import VLBDiffWave

class VLBDiffWaveApp:
    """Application for vlb-diffwave.
    """
    def __init__(self, config: Config):
        """Initializer.
        Args:
            config: model configuration.
        """
        self.config = config
        self.model = VLBDiffWave(config)
        self.param = None
        # since `denoise` is stateless, pre-jit is available
        self.denoiser = jax.jit(self.model.denoise)

    def __call__(self,
                 mel: jnp.ndarray,
                 timesteps: Union[int, jnp.ndarray] = 10,
                 key: Optional[jnp.ndarray] = None,
                 noise: Optional[jnp.ndarray] = None,
                 use_tqdm: bool = False) -> \
            Tuple[jnp.ndarray, List[np.ndarray]]:
        """Generate audio from mel-spectrogram.
        Args:
            mel: [float32; [B, T // H, M]], condition mel-spectrogram.
            timesteps: [float32; [S]], time steps, from one to zero, including endpoint.
            key: jax random prng key.
            noise: [float32; [B, T]], starting noise.
                neither key nor noise should be None.
            use_tqdm: whether use tqdm on iteration or not.
        Returns:
            [float32; [B, T]], generated audio and intermediate representations.
        """
        # assertion
        assert key is not None or noise is not None
        if noise is None:
            # split
            key, sub = jax.random.split(key)
            # B, T // H, _
            bsize, mellen, _ = mel.shape
            # [B, T]
            noise = jax.random.normal(sub, shape=(bsize, mellen * self.config.hop))
        if isinstance(timesteps, int):
            # [S]
            timesteps = jnp.linspace(1., 0., timesteps + 1)
        # scanning, outputs and intermediate representations
        return self.inference(mel, timesteps, noise, key, use_tqdm)

    def inference(self,
                  mel: jnp.ndarray,
                  timesteps: jnp.ndarray,
                  signal: jnp.ndarray,
                  key: Optional[jnp.ndarray] = None,
                  use_tqdm: bool = False) -> \
            Tuple[jnp.ndarray, List[np.ndarray]]:
        """Generate audio, just-in-time compiled.
        Args:
            mel: [float32; [B, T // H, M]], condition mel-spectrogram.
            timesteps: [float32; [S + 1]], time steps.
            signal: [float32; [B, T]], starting signal.
                neither key nor signal should be None.
            key: jax random key.
                if not provided, inference with mean only.
            use_tqdm: whether use tqdm for iteration visualization or not.
        Returns:
            [float32; [B, T]], generated audio and intermdeidate representations.
        """
        iterator = zip(timesteps[:-1], timesteps[1:])
        # use tqdm iterator
        if use_tqdm:
            iterator = tqdm(iterator, total=len(timesteps) - 1, leave=False)
        # S x [B, T]
        ir = []
        # [], []
        for time_t, time_s in iterator:
            # [B, T], [B]
            mean, std = self.denoiser(self.param, signal, mel, time_t[None], time_s[None])
            if key is None:
                # mean only
                signal = mean
            else:
                # split
                key, sub = jax.random.split(key)
                # [B, T]
                signal = mean + jax.random.normal(sub, shape=mean.shape) * std[:, None]
            # write it as cpu array for preventing oom
            ir.append(np.asarray(signal))
        # [B, T], S x [B, T]
        return signal, ir

    def init(self, key: np.ndarray):
        """Initialize model parameters.
        Args:
            key: jax prng keys.
        Returns:
            model parameters.
        """
        # placeholders
        signal = jnp.zeros([1, self.config.hop])
        time = jnp.zeros([1])
        mel = jnp.zeros([1, 1, 80])
        # initialize
        self.param = self.model.init(key, signal, time, mel)

    def write(self, path: str, optim: Optional[Any] = None):
        """Write model checkpoints.
        Args:
            path: path to the checkpoint.
            optim: optimizer state.
        """
        with open(path, 'wb') as f:
            f.write(flax.serialization.to_bytes(self.param))
        if optim is not None:
            name, ext = os.path.splitext(path)
            with open(f'{name}_optim{ext}', 'wb') as f:
                f.write(flax.serialization.to_bytes(optim))

    def restore(self, path: str, optim: Optional[Any] = None) -> Optional[Any]:
        """Restore model parameters from `path` checkpoint.
        Args:
            path: path to the checkpoint.
            optim: optimizer state, if provided.
        Returns:
            loaded optimizer states if provided
        """
        with open(path, 'rb') as f:
            binary = f.read()
        if self.param is None:
            # initialize parameters with dummy key.
            self.init(jax.random.PRNGKey(0))
        # restore
        self.param = flax.serialization.from_bytes(self.param, binary)
        # auxiliary restoration
        if optim is not None:
            name, ext = os.path.splitext(path)
            with open(f'{name}_optim{ext}', 'rb') as f:
                binary = f.read()
            # restore
            return flax.serialization.from_bytes(optim, binary)
        # explicit returns
        return None


""" ± Imports: ±
• os (Used to interact with the OS)
• numpy (Used to handle arrays)
• PIL (Python Imaging Library - Used to handle generation of images, i.e. Programmatic Penguins)
• secrets (randbelow - Used for random number generation)"""
import os
import numpy as np
from PIL import Image
from secrets import randbelow

DIMENSIONS = 480, 480
DIRNAME = os.path.dirname(__file__)

# ± determineSecondColor() method ± #

def determineSecondColor():
    # ± Get randint (b). If b is greater than 250, orange; if <= 250 or > 28, grey; else, white. ± #
    b = randbelow(750)

    if b > 250:
        sk = (223, 113, 38)
        return sk
    elif b <= 250 and b > 28:
        sk = (89, 86, 82)
        return sk
    else:
        sk = (255, 255, 255)
        return sk

# ± determineEyeColor() method ± #

def determineEyeColor():
    # ± If flipper color = beak color, eyes black; if determiner is between 1359 and 1530, eyes 'crazy red'; else, black. ± #
    x = randbelow(1000)
    if x < 134:
        ey = (172, 50, 50)
        return ey
    else:
        ey = (0, 0, 0)
        return ey

# ± rollColors() method ± #
def rollColors():
    colors = []

    # ± Generate a random RGB color for the penguin's b(o)d(y) - sk determined using functions. ± #
    bd = (randbelow(256), randbelow(256), randbelow(256))
    colors.append(bd)

    sk = determineSecondColor()
    colors.append(sk)

    ey = determineEyeColor()
    colors.append(ey)

    # ± Hair, Background & Outline ± #
    hr = (255, 242, 0)
    colors.append(hr)
    bg = (238, 238, 238)
    colors.append(bg)
    ol = (0, 0, 0)
    colors.append(ol)

    return colors

def rollPenguins():
    e = randbelow(1000)
    if e > 250:
        pixels = tall_penguin
    elif 250 >= e > 150:
        pixels = emperor_penguin
    else:
        pixels = plump_penguin
    return pixels

init = []
flag = 0
# ± Generate Programmatic Penguins ± #
for i in range(0, 1500):

    """ Colors List - 0: bd, 1: sk, 2: ey, 3: hr, 4: bg, 5: ol """
    c = rollColors()
    init.append(c)

    bd = c[0]
    sk = c[1]
    ey = c[2]
    hr = c[3]
    bg = c[4]
    ol = c[5]

    """ Pixel Sets """
    tall_penguin = [
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, ol, ol,
         ol, ol, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, ol, ol, bd, bd,
         bd, bd, ol, ol, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol, bd, bd, ey, bd,
         ey, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol, bd, bd, ey, bd,
         ey, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol, bd, bd, bd, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, ol, ol, ol, ol, ol,
         ol, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, sk, sk, sk, sk, sk, sk,
         sk, ol, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, ol, ol, ol, ol, ol, ol,
         ol, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, bg, bg, bd, bd, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, bg, bg, bg, bd, bd, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, bg, bg, bg, bg, bd, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, bg, bg, bg, bg, bg, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, bg, bg, bg, bd, bd,
         bd, bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, bg, ol, ol, ol, ol,
         ol, ol, ol, ol, bd, bd, ol, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, ol,  sk,  sk, ol,  sk,
         sk,  sk,  sk,  sk, ol, bd, ol, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol,  sk,  sk, ol,  sk,
         sk,  sk,  sk,  sk, ol, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol, ol, ol, ol, ol,
         ol, ol, ol, ol, ol, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg]
    ]

    plump_penguin = [
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, ol, ol, ol,
         ol, ol, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, ol, ol, bd, bd, bd,
         bd, bd, ol, ol, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol, bd, bd, ey, bd, bd,
         ey, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol, bd, bd, bd, bd, bd,
         bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, ol, ol, ol, ol, ol, ol,
         ol, ol, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, sk, sk, sk, sk, sk, sk, sk,
         sk, ol, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, ol, ol, ol, ol, ol, ol, ol,
         ol, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, bg, bg, bd, bd, bd, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, bg, bg, bg, bg, bd, bd, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, bg, bg, bg, bg, bg, bd, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, bg, bg, bg, bg, bg, bd, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, bg, bg, bg, bd, bd, bd,
         bd, bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, ol, ol, ol, ol, ol, ol,
         ol, ol, ol, ol, ol, bd, ol, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol,  sk,  sk,  sk, ol,  sk,
         sk,  sk,  sk,  sk, ol, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol, ol, ol, ol, ol, ol,
         ol, ol, ol, ol, ol, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg]
    ]

    emperor_penguin = [
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, hr, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, hr, hr, bg, bg,
         hr, hr, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, hr, hr,
         ol, bg, bg, bg, hr, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, hr, hr, ol, ol, bg, bg,
         ol, ol, bg, ol, hr, hr, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, ol, ol, bg,
         bg, ol, ol, bg, bg, hr, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, ol, ol,
         ol, ol, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, ol, ol, bd, bd,
         bd, bd, ol, ol, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol, bd, bd, ey, bd,
         ey, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol, bd, bd, ey, bd,
         ey, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol, bd, bd, bd, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, ol, ol, ol, ol, ol,
         ol, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, sk, sk, sk, sk, sk, sk,
         sk, ol, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, ol, ol, ol, ol, ol, ol,
         ol, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, bg, bg, bd, bd, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, bg, bg, bg, bd, bd, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, bg, bg, bg, bg, bd, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, ol, bg, bg, bg, bg, bg, bd,
         bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, bg, bg, bg, bd, bd,
         bd, bd, bd, bd, bd, bd, ol, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, bg, ol, ol, ol, ol,
         ol, ol, ol, ol, bd, bd, ol, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, ol, ol,  sk,  sk, ol,  sk,
         sk,  sk,  sk,  sk, ol, bd, ol, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol,  sk,  sk, ol,  sk,
         sk,  sk,  sk,  sk, ol, ol, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, ol, ol, ol, ol, ol,
         ol, ol, ol, ol, ol, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg],
        [bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg,
         bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg, bg]
    ]

    if flag == 0:
        last = 0
        flag += 1
    else:
        last = init[i-1]
        current = init[i]
        if current == last:
            same = True
            while same:
                c = rollColors()
                init.append(c)
                if current != last:
                    same = False
                else:
                    pass
        else:
            pass

    pixels = rollPenguins()
    arr = np.array(pixels, dtype=np.uint8)

    # ± Use PIL to create an image using the array of pixels supplied. ± #
    img = Image.fromarray(arr)
    img = img.resize(DIMENSIONS, resample=0)
    filename = DIRNAME + '/output/' + (str(i+1)) + '.png'
    img.save(filename)

    print(f'Created & Saved Image: {i+1}.')
print(f'Created {i+1} Images.')