"""Main driver of the fitting routine."""
import pickle
import random
import time
import warnings
from multiprocessing import Pool, Process
from tqdm import tqdm

import extinction
import astropy.units as u
import pandas as pd
import numpy as np
import scipy.stats as st
from numpy.random import choice
from astropy.constants import sigma_sb
from isochrones.interp import DFInterpolator
from termcolor import colored

from .config import filesdir, gridsdir, priorsdir
from .error import *
from .isochrone import estimate
from .phot_utils import *
from .sed_library import *
from .utils import *

try:
    import dynesty
    from dynesty.utils import resample_equal

    bma_flag = True
except ModuleNotFoundError:
    wrn = 'Dynesty package not found.\n'
    wrn += 'Install dynesty with `pip install dynesty`'
    warnings.warn(wrn)
    bma_flag = False
try:
    import pymultinest
except ModuleNotFoundError:
    warnings.warn(
        '(py)MultiNest installation (or libmultinest.dylib) not detected.'
    )

class Fitter:
    """The Fitter class handles the fitting routines and parameter estimation.

    Examples
    --------
    The fitter isn't instantiaded with any arguments, rather you instantiate a
    Fitter object, then you set up the configurations and finally you
    initialize the object by running the initialize method.
    >>> f = Fitter()
    >>> f.star = s  # s must be a valid Star object.
    >>> f.initialize()

    Attributes
    ----------
    out_folder : type
        Description of attribute `out_folder`.
    verbose : type
        Description of attribute `verbose`.
    star : type
        Description of attribute `star`.
    setup : type
        Description of attribute `setup`.
    norm : type
        Description of attribute `norm`.
    grid : type
        Description of attribute `grid`.
    estimate_logg : type
        Description of attribute `estimate_logg`.
    priorfile : type
        Description of attribute `priorfile`.
    av_law : type
        Description of attribute `av_law`.
    n_samples : type
        Description of attribute `n_samples`.
    bma : type
        Description of attribute `bma`.
    prior_setup : type
        Description of attribute `prior_setup`.
    sequential : type
        Description of attribute `sequential`.

    """

    colors = [
        'red', 'green', 'blue', 'yellow',
        'grey', 'magenta', 'cyan', 'white'
    ]

    def __init__(self):

        # Default values for attributes
        self._interpolators = []
        self._grids = []
        self.out_folder = None
        self.verbose = True
        self.star = None
        self.setup = ['dynesty']
        self.norm = False
        self.grid = 'phoenix'
        self.estimate_logg = False
        self.av_law = 'fitzpatrick'
        self.n_samples = None
        self.bma = False
        self.prior_setup = None
        self.sequential = True
        self.experimental = False

    @property
    def star(self):
        """Star to fit for."""
        return self._star

    @star.setter
    def star(self, star):
        # if not isinstance(star, Star) and star is not None:
        #     InstanceError(star, Star).__raise__()
        self._star = star

    @property
    def setup(self):
        """Set up options."""
        return self._setup

    @setup.setter
    def setup(self, setup):
        err_msg = 'The setup has to contain at least the fitting engine'
        err_msg += f', multinest or dynesty.\nThe setup was {setup}'
        if len(setup) < 1:
            InputError(err_msg).__raise__()
        self._setup = setup
        self._engine = setup[0]
        defaults = False
        if len(setup) == 1:
            defaults = True
        if self._engine == 'multinest':
            if defaults:
                self._nlive = 500
                self._dlogz = 0.5
            else:
                self._nlive = setup[1]
                self._dlogz = setup[2]
        if self._engine == 'dynesty':
            if defaults:
                self._nlive = 500
                self._dlogz = 0.5
                self._bound = 'multi'
                self._sample = 'rwalk'
                self._threads = 1
                self._dynamic = False
            else:
                self._nlive = setup[1]
                self._dlogz = setup[2]
                self._bound = setup[3]
                self._sample = setup[4]
                self._threads = setup[5]
                self._dynamic = setup[6]

    @property
    def norm(self):
        """Bool to decide if a normalization constant will be fitted.

        Set as True to not fit for radius and distance and fit for a
        normalization constant. After the fit a radius is calculated using
        the Gaia parallax.
        """
        return self._norm

    @norm.setter
    def norm(self, norm):
        if type(norm) is not bool:
            InputError('norm must be True or False.').__raise__()
        self._norm = norm

    @property
    def grid(self):
        """Model grid selected."""
        return self._grid

    @grid.setter
    def grid(self, grid):
        assert type(grid) == str
        self._grid = grid
        if grid.lower() == 'phoenix':
            with open(gridsdir + '/Phoenixv2_DF.pkl', 'rb') as intp:
                self._interpolator = DFInterpolator(pd.read_pickle(intp))
        if grid.lower() == 'btsettl':
            with open(gridsdir + '/BTSettl_DF.pkl', 'rb') as intp:
                self._interpolator = DFInterpolator(pd.read_pickle(intp))
        if grid.lower() == 'btnextgen':
            with open(gridsdir + '/BTNextGen_DF.pkl', 'rb') as intp:
                self._interpolator = DFInterpolator(pd.read_pickle(intp))
        if grid.lower() == 'btcond':
            with open(gridsdir + '/BTCond_DF.pkl', 'rb') as intp:
                self._interpolator = DFInterpolator(pd.read_pickle(intp))
        if grid.lower() == 'ck04':
            with open(gridsdir + '/CK04_DF.pkl', 'rb') as intp:
                self._interpolator = DFInterpolator(pd.read_pickle(intp))
        if grid.lower() == 'kurucz':
            with open(gridsdir + '/Kurucz_DF.pkl', 'rb') as intp:
                self._interpolator = DFInterpolator(pd.read_pickle(intp))
        if grid.lower() == 'coelho':
            with open(gridsdir + '/Coelho_DF.pkl', 'rb') as intp:
                self._interpolator = DFInterpolator(pd.read_pickle(intp))

    @property
    def bma(self):
        """Bayesian Model Averaging (BMA).

        Set to True if BMA is wanted. This loads every model grid interpolator
        and fits an SED to all of them, so the runtime will be slower!
        """
        return self._bma

    @bma.setter
    def bma(self, bma):
        self._bma = bma if bma_flag else False

    @property
    def models(self):
        """Models to be used in BMA."""
        return self._bma_models

    @models.setter
    def models(self, mods):
        self._bma_models = mods

    @property
    def sequential(self):
        """Set to True to make BMA sequentially instead of parallel."""
        return self._sequential

    @sequential.setter
    def sequential(self, sequential):
        self._sequential = sequential

    @property
    def n_samples(self):
        """Set number of samples for BMA."""
        return self._nsamp

    @n_samples.setter
    def n_samples(self, nsamp):
        self._nsamp = nsamp

    @property
    def verbose(self):
        """Program verbosity. Default is True."""
        return self._verbose

    @verbose.setter
    def verbose(self, verbose):
        if type(verbose) is not bool:
            InputError('Verbose must be True or False.').__raise__()
        self._verbose = verbose

    @property
    def out_folder(self):
        """Output folder.

        If none is provided the default will be the starname.
        """
        return self._out_folder

    @out_folder.setter
    def out_folder(self, out_folder):
        if type(out_folder) is not str and out_folder is not None:
            err_msg = 'Output folder must be an address or None.'
            InputError(err_msg).__raise__()
        self._out_folder = out_folder

    @property
    def av_law(self):
        """Select extinction law."""
        return self._av_law

    @av_law.setter
    def av_law(self, law):
        laws = [
            'cardelli',
            'odonnell',
            'calzetti',
            'fitzpatrick'
        ]
        law = law.lower()
        if law not in laws:
            err_msg = f'Extinction law {law} not recognized. Available extinction'
            err_msg += ' laws are: `cardelli`, `odonnell`'
            err_msg += ', `calzetti`, and `fitzpatrick`'
            InputError(err_msg).__raise__()
        law_f = None
        if law == laws[0]:
            law_f = extinction.ccm89
        if law == laws[1]:
            law_f = extinction.odonnell94
        if law == laws[2]:
            law_f = extinction.calzetti00
        if law == laws[3]:
            law_f = extinction.fitzpatrick99
        self._av_law = law_f

    def initialize(self):
        """Initialize the fitter.

        To be run only after every input is added.
        This function calculates the number of dimensions, runs the prior
        creation, creates output directory, initializes coordinators and sets
        up global variables.
        """
        global prior_dict, coordinator, fixed, order, star, use_norm, av_law
        self.start = time.time()
        err_msg = 'No star is detected. Please create an instance of Star.'
        if self.star is None:
            er = InputError(err_msg)
            er.log(self.out + '/output.log')
            er.__raise__()
        star = self.star
        if not self._bma:
            global interpolator
            interpolator = self._interpolator
        use_norm = self.norm
        # Extinction law
        av_law = self._av_law

        # Declare order of parameters.
        if not self.norm:
            order = np.array(
                [
                    'teff', 'logg', 'z',
                    'dist', 'rad', 'Av'
                ]
            )
        else:
            order = np.array(['teff', 'logg', 'z', 'norm', 'Av'])

        # Create output directory
        if self.out_folder is None:
            self.out_folder = self.star.starname + '/'
        create_dir(self.out_folder)

        self.star.save_mags(self.out_folder + '/')

        # Parameter coordination.
        # Order for the parameters are:
        # teff, logg, z, dist, rad, Av, noise
        # or
        # teff, logg, z, norm, Av, noise
        npars = 6 if not self.norm else 5
        npars += self.star.used_filters.sum()
        npars = int(npars)
        self.coordinator = np.zeros(npars)  # 1 for fixed params
        self.fixed = np.zeros(npars)
        coordinator = self.coordinator
        fixed = self.fixed

        # Setup priors.
        self.default_priors = self._default_priors()
        self.create_priors_from_setup()
        prior_dict = self.priors

        # Get dimensions.
        self.ndim = self.get_ndim()

        # warnings
        if len(self._setup) == 1:
            print('USING DEFAULT SETUP VALUES.')

        # BMA settings
        # if BMA is used, load all interpolators requested.
        if self._bma:
            if self.n_samples is None:
                self.n_samples = 'max'

            if self.star.offline:
                if self.star.temp is None:
                    self.star.temp = 4001
                off_msg = 'Offline mode assumes that the stellar'
                off_msg += ' temperature is greater than 4000 K'
                off_msg += '. If you believe this is not the case then please '
                off_msg += 'add a temperature to the Star constructor. '
                off_msg += f'The input temperature is {self.star.temp}'
                print(colored(off_msg, 'yellow'))

            for mod in self._bma_models:
                # We'll assume that if ARIADNE is running in offline mode
                # Then the star will have > 4000 K
                if mod.lower() == 'phoenix':
                    with open(gridsdir + '/Phoenixv2_DF.pkl', 'rb') as intp:
                        df = DFInterpolator(pd.read_pickle(intp))
                if mod.lower() == 'btsettl':
                    with open(gridsdir + '/BTSettl_DF.pkl', 'rb') as intp:
                        df = DFInterpolator(pd.read_pickle(intp))
                if mod.lower() == 'btnextgen':
                    if self.star.temp > 4000:
                        continue
                    else:
                        with open(gridsdir + '/BTNextGen_DF.pkl', 'rb') as inp:
                            df = DFInterpolator(pd.read_pickle(inp))
                if mod.lower() == 'btcond':
                    if self.star.temp > 4000:
                        continue
                    else:
                        with open(gridsdir + '/BTCond_DF.pkl', 'rb') as intp:
                            df = DFInterpolator(pd.read_pickle(intp))
                if mod.lower() == 'ck04':
                    if self.star.temp > 4000:
                        with open(gridsdir + '/CK04_DF.pkl', 'rb') as intp:
                            df = DFInterpolator(pd.read_pickle(intp))
                    else:
                        # Warning temp too low for model
                        continue
                if mod.lower() == 'kurucz':
                    if self.star.temp > 4000:
                        with open(gridsdir + '/Kurucz_DF.pkl', 'rb') as intp:
                            df = DFInterpolator(pd.read_pickle(intp))
                    else:
                        # Warning temp too low for model.
                        continue
                # if mod.lower() == 'coelho':
                #     if self.star.temp > 3500:
                #         with open(gridsdir + '/Coelho_DF.pkl', 'rb') as intp:
                #             df = DFInterpolator(pd.read_pickle(intp))
                #     else:
                #         # Warning
                #         continue
                self._interpolators.append(df)
                self._grids.append(mod)
            thr = self._threads if self._sequential else len(
                self._interpolators)
        else:
            thr = self._threads
        en = 'Bayesian Model Averaging' if self._bma else self._engine
        display_routine(en, self._nlive, self._dlogz, self.ndim, self._bound,
                        self._sample, thr, self._dynamic)

    def get_ndim(self):
        """Calculate number of dimensions."""
        ndim = 6 if not self.norm else 5
        ndim += self.star.used_filters.sum()
        ndim -= self.coordinator.sum()
        return int(ndim)

    def _default_priors(self):
        global order
        defaults = dict()
        # Logg prior setup.
        if self.star.get_logg:
            defaults['logg'] = st.norm(
                loc=self.star.logg, scale=self.star.logg_e)
        else:
            with open(priorsdir + '/logg_ppf.pkl', 'rb') as jar:
                defaults['logg'] = pickle.load(jar)
        # Teff prior from RAVE
        with open(priorsdir + '/teff_ppf.pkl', 'rb') as jar:
            defaults['teff'] = pickle.load(jar)
        # [Fe/H] prior setup.
        defaults['z'] = st.norm(loc=-0.125, scale=0.234)
        # Distance prior setup.
        if not self._norm:
            if self.star.dist != -1:
                defaults['dist'] = st.norm(
                    loc=self.star.dist, scale=5 * self.star.dist_e)
            else:
                defaults['dist'] = st.uniform(loc=1, scale=3000)
            # Radius prior setup.
            defaults['rad'] = st.uniform(loc=0.05, scale=100)
        # Normalization prior setup.
        else:
            up = 1 / 1e-20
            defaults['norm'] = st.truncnorm(a=0, b=up, loc=0, scale=1e-15)
        # Extinction prior setup.
        if self.star.Av == 0.:
            av_idx = 4 if self._norm else 5
            self.coordinator[av_idx] = 1
            self.fixed[av_idx] = 0
            defaults['Av'] = None
        else:
            defaults['Av'] = st.uniform(loc=0, scale=self.star.Av)
        # Noise model prior setup.
        mask = self.star.filter_mask
        flxs = self.star.flux[mask]
        errs = self.star.flux_er[mask]
        for filt, flx, flx_e in zip(self.star.filter_names[mask], flxs, errs):
            p_ = get_noise_name(filt) + '_noise'
            mu = 0
            sigma = flx_e * 10
            b = (1 - flx) / flx_e
            defaults[p_] = st.truncnorm(loc=mu, scale=sigma, a=0, b=b)
            # defaults[p_] = st.uniform(loc=0, scale=5)
            order = np.append(order, p_)
        return defaults

    def create_priors_from_setup(self):
        """Create priors from the manual setup."""
        prior_dict = dict()
        keys = self.prior_setup.keys()
        noise = []
        mask = self.star.filter_mask
        flxs = self.star.flux[mask]
        errs = self.star.flux_er[mask]
        for filt, flx, flx_e in zip(self.star.filter_names[mask], flxs, errs):
            p_ = get_noise_name(filt) + '_noise'
            noise.append(p_)
        prior_out = 'Parameter\tPrior\tValues\n'
        if 'norm' in keys and ('rad' in keys or 'dist' in keys):
            er = PriorError('rad or dist', 1)
            er.log(self.out_folder + '/output.log')
            er.__raise__()
        for k in keys:
            if type(self.prior_setup[k]) == str:
                if self.prior_setup[k] == 'default':
                    prior_dict[k] = self.default_priors[k]
                    prior_out += k + '\tdefault\n'
                if self.prior_setup[k].lower() == 'rave':
                    # RAVE prior only available for teff and logg. It's already
                    # the default for [Fe/H]
                    if k == 'logg' or k == 'teff':
                        if k == 'teff':
                            PriorError('teff', 2).warn()
                        with open(priorsdir + '/teff_ppf.pkl', 'rb') as jar:
                            prior_dict[k] = pickle.load(jar)
                        prior_out += k + '\tRAVE\n'

            else:
                prior = self.prior_setup[k][0]
                if prior == 'fixed':
                    value = self.prior_setup[k][1]
                    idx = np.where(k == order)[0]
                    self.coordinator[idx] = 1
                    self.fixed[idx] = value
                    prior_out += k + '\tfixed\t{}\n'.format(value)
                if prior == 'normal':
                    mu = self.prior_setup[k][1]
                    sig = self.prior_setup[k][2]
                    prior_dict[k] = st.norm(loc=mu, scale=sig)
                    prior_out += k + '\tnormal\t{}\t{}\n'.format(mu, sig)
                if prior == 'truncnorm':
                    mu = self.prior_setup[k][1]
                    sig = self.prior_setup[k][2]
                    low = self.prior_setup[k][3]
                    up = self.prior_setup[k][4]
                    b, a = (up - mu) / sig, (low - mu) / sig
                    prior_dict[k] = st.truncnorm(a=a, b=b, loc=mu, scale=sig)
                    prior_out += k
                    prior_out += '\ttruncatednormal\t{}\t{}\t{}\t{}\n'.format(
                        mu, sig, low, up)
                if prior == 'uniform':
                    low = self.prior_setup[k][1]
                    up = self.prior_setup[k][2]
                    prior_dict[k] = st.uniform(loc=low, scale=up - low)
                    prior_out += k + '\tuniform\t{}\t{}\n'.format(low, up)
        for par in noise:
            prior_dict[par] = self.default_priors[par]
        ff = open(self.out_folder + '/prior.dat', 'w')
        ff.write(prior_out)
        ff.close()
        del ff
        self.priors = prior_dict
        pass

    def fit(self):
        """Run fitting routine."""
        if self._engine == 'multinest':
            self.fit_multinest()
        else:
            self.fit_dynesty()
        elapsed_time = execution_time(self.start)
        end(self.coordinator, elapsed_time,
            self.out_folder, self._engine, self.norm, )
        pass

    def fit_bma(self):
        """Perform the fit with different models and the average the output.

        Only works with dynesty.
        """
        if len(self.star.filter_names[self.star.filter_mask]) <= 5:
            print(colored('\t\t\tNOT ENOUGH POINTS TO MAKE THE FIT! !', 'red'))
            return
        thr = self._threads if self._sequential else len(self._interpolators)
        # display('Bayesian Model Averaging', self.star, self._nlive,
        #         self._dlogz, self.ndim, self._bound, self._sample,
        #         thr, self._dynamic)
        if not self._sequential:
            jobs = []
            n_threads = len(self._interpolators)
            for intp, gr in zip(self._interpolators, self._grids):
                p = Process(target=self._bma_dynesty, args=([intp, gr]))
                jobs.append(p)
                p.start()
            for p in jobs:
                p.join()
        else:
            global interpolator
            for intp, gr in zip(self._interpolators, self._grids):
                interpolator = intp
                self.grid = gr
                out_file = self.out_folder + '/' + gr + '_out.pkl'
                print('\t\t\tFITTING MODEL : ' + gr)
                try:
                    self.fit_dynesty(out_file=out_file)
                except ValueError as e:
                    dump_out = self.out_folder + '/' + gr + '_DUMP.pkl'
                    pickle.dump(self.sampler.results, open(dump_out, 'wb'))
                    DynestyError(dump_out, gr, e).__raise__()
                    continue

        # Now that the fitting finished, read the outputs and average
        # the posteriors
        outs = []
        for g in self._grids:
            in_folder = f'{self.out_folder}/{g}_out.pkl'
            outs.append(in_folder)
            # with open(in_folder, 'rb') as out:
            #     outs.append(pickle.load(out))
        c = np.random.choice(self.colors)
        avgd = self.bayesian_model_average(outs, self._grids, self._norm,
                                           self.n_samples, c)
        self.save_bma(avgd)

        elapsed_time = execution_time(self.start)
        end(self.coordinator, elapsed_time, self.out_folder,
            'Bayesian Model Averaging', self.norm)
        pass

    def _bma_dynesty(self, intp, grid):
        global interpolator
        interpolator = intp

        # Parallel parallelized routine experiment
        if self.experimental:
            if self._dynamic:
                with Pool(self._threads) as executor:
                    sampler = dynesty.DynamicNestedSampler(
                        dynesty_loglike_bma, pt_dynesty, self.ndim,
                        bound=self._bound, sample=self._sample, pool=executor,
                        queue_size=self._threads, logl_args=([intp])
                    )
                    sampler.run_nested(dlogz_init=self._dlogz,
                                       nlive_batch=self._nlive,
                                       wt_kwargs={'pfrac': .95})
            else:
                with Pool(self._threads) as executor:
                    sampler = dynesty.NestedSampler(
                        dynesty_loglike_bma, pt_dynesty, self.ndim,
                        nlive=self._nlive, bound=self._bound,
                        sample=self._sample, pool=executor,
                        queue_size=self._threads, logl_args=([intp])
                    )
                    sampler.run_nested(dlogz=self._dlogz)

        elif self._dynamic:
            sampler = dynesty.DynamicNestedSampler(
                dynesty_loglike_bma, pt_dynesty, self.ndim,
                bound=self._bound, sample=self._sample, logl_args=([intp])

            )
            sampler.run_nested(dlogz_init=self._dlogz,
                               nlive_init=self._nlive,
                               wt_kwargs={'pfrac': .95})
        else:
            try:
                self.sampler = dynesty.NestedSampler(
                    dynesty_loglike_bma, pt_dynesty, self.ndim,
                    nlive=self._nlive, bound=self._bound,
                    sample=self._sample,
                    logl_args=([intp])
                )
                self.sampler.run_nested(dlogz=self._dlogz)
            except Error:
                dump_out = self.out_folder + '/' + grid + '_DUMP.pkl'
                pickle.dump(self.sampler.results, open(dump_out, 'wb'))
                er = DynestyError(dump_out, grid)
                er.log(self.out + '/output.log')
                er.__raise__()

        results = self.sampler.results
        out_file = self.out_folder + '/' + grid + '_out.pkl'
        self.save(out_file, results=results)
        pass

    def fit_multinest(self, out_file=None):
        """Run MultiNest."""
        # Set up some globals
        global mask, flux, flux_er, filts, wave
        mask = star.filter_mask
        flux = star.flux[mask]
        flux_er = star.flux_er[mask]
        filts = star.filter_names[mask]
        wave = star.wave[mask]
        path = self.out_folder + '/mnest/'
        create_dir(path)  # Create multinest path.
        pymultinest.run(
            multinest_log_like, pt_multinest, self.ndim,
            n_params=self.ndim,
            sampling_efficiency=0.8,
            evidence_tolerance=self._dlogz,
            n_live_points=self._nlive,
            outputfiles_basename=path + 'chains',
            max_modes=100,
            verbose=self.verbose,
            resume=False
        )
        if out_file is None:
            out_file = f'{self.out_folder}/{self._grid}_out.pkl'
        self.save(out_file=out_file)
        pass

    def fit_dynesty(self, out_file=None):
        """Run dynesty."""
        # Set up some globals
        global mask, flux, flux_er, filts, wave
        mask = star.filter_mask
        flux = star.flux[mask]
        flux_er = star.flux_er[mask]
        filts = star.filter_names[mask]
        wave = star.wave[mask]
        if self._dynamic:
            if self._threads > 1:
                with Pool(self._threads) as executor:
                    self.sampler = dynesty.DynamicNestedSampler(
                        dynesty_log_like, pt_dynesty, self.ndim,
                        bound=self._bound, sample=self._sample,
                        pool=executor, walks=25,
                        queue_size=self._threads - 1
                    )
                    self.sampler.run_nested(dlogz_init=self._dlogz,
                                            nlive_init=self._nlive,
                                            wt_kwargs={'pfrac': 1})
            else:
                self.sampler = dynesty.DynamicNestedSampler(
                    dynesty_log_like, pt_dynesty, self.ndim, walks=25,
                    bound=self._bound, sample=self._sample

                )
                self.sampler.run_nested(dlogz_init=self._dlogz,
                                        nlive_init=self._nlive,
                                        wt_kwargs={'pfrac': 1})
        else:
            if self._threads > 1:
                with Pool(self._threads) as executor:
                    self.sampler = dynesty.NestedSampler(
                        dynesty_log_like, pt_dynesty, self.ndim,
                        nlive=self._nlive, bound=self._bound,
                        sample=self._sample,
                        pool=executor, walks=25,
                        queue_size=self._threads - 1,
                    )
                    self.sampler.run_nested(dlogz=self._dlogz)
            else:
                self.sampler = dynesty.NestedSampler(
                    dynesty_log_like, pt_dynesty, self.ndim, walks=25,
                    nlive=self._nlive, bound=self._bound,
                    sample=self._sample
                )
                self.sampler.run_nested(dlogz=self._dlogz)
        results = self.sampler.results
        if out_file is None:
            out_file = f'{self.out_folder}/{self._grid}_out.pkl'
        self.save(out_file, results=results)
        pass

    def save(self, out_file, results=None):
        """Save multinest/dynesty output and relevant information.

        Saves a dictionary as a pickle file. The dictionary contains the
        following:

        lnZ : The global evidence.
        lnZerr : The global evidence error.
        posterior_samples : A dictionary containing the samples of each
                            parameter (even if it's fixed), the evidence,
                            log likelihood, the prior, and the posterior
                            for each set of sampled parameters.
        fixed : An array with the fixed parameter values.
        coordinator : An array with the status of each parameter (1 for fixed
                      0 for free)
        best_fit : The best fit is chosen to be the median of each sample.
                   It also includes the log likelihood of the best fit.
        star : The Star object containing the information of the star (name,
               magnitudes, fluxes, coordinates, etc)
        engine : The fitting engine used (i.e. MultiNest or Dynesty)

        Also creates a log file with the best fit parameters and 1 sigma
        error bars.

        """
        out = dict()
        logdat = '#Parameter\tmedian\tupper\tlower\t3sig_CI\n'
        log_out = self.out_folder + '/' + 'best_fit.dat'
        if self._engine == 'multinest':
            lnz, lnzer, posterior_samples = self.multinest_results(
                self.out_folder,
                self.ndim
            )
        else:
            lnz, lnzer, posterior_samples = self.dynesty_results(results)

        n = int(self.star.used_filters.sum())
        mask = self.star.filter_mask

        # Save global evidence

        if self._engine == 'dynesty':
            out['dynesty'] = results
        out['global_lnZ'] = lnz
        out['global_lnZerr'] = lnzer

        # Create raw samples holder

        out['posterior_samples'] = dict()
        j = 0
        k = 0  # filter counter
        for i, param in enumerate(order):
            if not self.coordinator[i]:
                samples = posterior_samples[:, j]
                if 'noise' in param:
                    filt = self.star.filter_names[mask][k]
                    flx = self.star.flux[mask][k]
                    _, samples = flux_to_mag(flx, samples, filt)
                    k += 1
                out['posterior_samples'][param] = samples
                j += 1
            else:
                out['posterior_samples'][param] = self.fixed[i]

        # Save loglike, priors and posteriors.

        out['posterior_samples']['loglike'] = np.zeros(
            posterior_samples.shape[0]
        )

        # If normalization constant was fitted, create a distribution of radii
        # only if there's a distance available.

        if use_norm and star.dist != -1:
            rad = self._get_rad(
                out['posterior_samples']['norm'], star.dist, star.dist_e
            )
            out['posterior_samples']['rad'] = rad

        # Create a distribution of masses.

        logg_samp = out['posterior_samples']['logg']
        rad_samp = out['posterior_samples']['rad']
        mass_samp = self._get_mass(logg_samp, rad_samp)
        out['posterior_samples']['grav_mass'] = mass_samp

        # Create a distribution of luminosities.

        teff_samp = out['posterior_samples']['teff']
        lum_samp = self._get_lum(teff_samp, rad_samp)
        out['posterior_samples']['lum'] = lum_samp

        # Create a distribution of angular diameters.

        if not use_norm:
            dist_samp = out['posterior_samples']['dist']
            ad_samp = self._get_angular_diameter(rad_samp, dist_samp)
            out['posterior_samples']['AD'] = ad_samp

        for i in range(posterior_samples.shape[0]):
            theta = build_params(
                posterior_samples[i, :], flux, flux_er, filts, coordinator,
                fixed, self.norm)
            out['posterior_samples']['loglike'][i] = log_likelihood(
                theta, flux, flux_er, wave, filts, interpolator, self.norm,
                av_law)
        lnlike = out['posterior_samples']['loglike']

        # Best fit
        # The logic is as follows:
        # Calculate KDE for each marginalized posterior distributions
        # Find peak
        # peak is best fit.
        # do only if not bma

        if not self.bma:
            out['best_fit'] = dict()
            out['uncertainties'] = dict()
            out['confidence_interval'] = dict()
            best_theta = np.zeros(order.shape[0])

            for i, param in enumerate(order):
                if not self.coordinator[i]:
                    if 'noise' in param:
                        continue
                    samp = out['posterior_samples'][param]

                    if param == 'z':
                        logdat = out_filler(samp, logdat, param, '[Fe/H]', out)
                    elif param == 'norm':
                        logdat = out_filler(samp, logdat, param, '(R/D)^2',
                                            out, fmt='e')
                        if star.dist != 1:
                            logdat = out_filler(
                                out['posterior_samples']['rad'], logdat, 'rad',
                                'R', out
                            )
                    else:
                        logdat = out_filler(samp, logdat, param, param, out)
                else:
                    logdat = out_filler(
                        0, logdat, param, out, fixed=self.fixed[i]
                    )
                best_theta[i] = out['best_fit'][param]

            # Add derived mass to best fit dictionary.

            samp = out['posterior_samples']['grav_mass']
            logdat = out_filler(
                samp, logdat, 'grav_mass', 'grav_mass', out
            )

            # Add derived luminosity to best fit dictionary.

            samp = out['posterior_samples']['lum']
            logdat = out_filler(samp, logdat, 'lum', 'lum', out)

            # Add derived angular diameter to best fit dictionary.

            if not use_norm:
                samp = out['posterior_samples']['AD']
                logdat = out_filler(samp, logdat, 'AD', 'AD', out)

            for i, param in enumerate(order):
                if not self.coordinator[i]:
                    if 'noise' not in param:
                        continue
                    samp = out['posterior_samples'][param]
                    logdat = out_filler(samp, logdat, param,
                                        param, out, fmt='f')

            # Fill in best loglike, prior and posterior.

            out['best_fit']['loglike'] = log_likelihood(
                best_theta, flux, flux_er, wave,
                filts, interpolator, self.norm, av_law
            )

            # Spectral type
            # Load Mamajek spt table
            mamajek_spt = np.loadtxt(
                filesdir + '/mamajek_spt.dat', dtype=str, usecols=[0])
            mamajek_temp = np.loadtxt(
                filesdir + '/mamajek_spt.dat', usecols=[1])

            # Find spt
            spt_idx = np.argmin(abs(mamajek_temp - out['best_fit']['teff']))
            spt = mamajek_spt[spt_idx]
            out['spectral_type'] = spt

        # Utilities for plotting.

        out['fixed'] = self.fixed
        out['coordinator'] = self.coordinator
        out['star'] = self.star
        out['engine'] = self._engine
        out['norm'] = self.norm
        out['model_grid'] = self.grid
        out['av_law'] = av_law
        if not self.bma:
            with open(log_out, 'w') as logfile:
                logfile.write(logdat)
        pickle.dump(out, open(out_file, 'wb'))
        pass

    def save_bma(self, avgd):
        """Save BMA output and relevant information.

        Saves a dictionary as a pickle file. The dictionary contains the
        following:

        lnZ : The global evidences.
        posterior_samples : A dictionary containing the samples of each
                            parameter (even if it's fixed), the evidence,
                            log likelihood, the prior, and the posterior
                            for each set of sampled parameters.
        fixed : An array with the fixed parameter values.
        coordinator : An array with the status of each parameter (1 for fixed
                      0 for free)
        best_fit : The best fit is chosen to be the median of each sample.
                   It also includes the log likelihood of the best fit.
        star : The Star object containing the information of the star (name,
               magnitudes, fluxes, coordinates, etc)

        Also creates a log file with the best fit parameters, 1 sigma
        error bars and 3 sigma CIs.

        """
        out = dict()
        logdat_samples = '#Parameter\tmedian\tupper\tlower\t3sig_low\t3sig_up\n'
        logdat_average = '#Parameter\tmedian\tupper\tlower\t3sig_low\t3sig_up\n'
        log_out_samples = f'{self.out_folder}/best_fit_sample.dat'
        log_out_average = f'{self.out_folder}/best_fit_average.dat'
        prob_out = f'{self.out_folder}/model_probabilities.dat'

        # Save global evidence of each model.
        out['lnZ'] = avgd['evidences']

        # Save original samples.
        out['originals'] = avgd['originals']

        # Save weights.
        out['weights'] = avgd['weights']

        # Create raw samples holder
        out['weighted_samples'] = dict()
        out['weighted_average'] = dict()
        j = 0
        for i, par in enumerate(order):
            if not self.coordinator[i]:
                out['weighted_samples'][par] = avgd['weighted_samples'][par]
                out['weighted_average'][par] = avgd['weighted_average'][par]
                j += 1
            else:
                out['weighted_samples'][par] = self.fixed[i]
                out['weighted_average'][par] = self.fixed[i]

        # If normalization constant was fitted, create a distribution of radii.

        if use_norm and star.dist != -1:
            rad_sampled = self._get_rad(
                out['weighted_samples']['norm'], star.dist, star.dist_e
            )
            rad_averageed = self._get_rad(
                out['weighted_average']['norm'], star.dist, star.dist_e
            )
            out['weighted_samples']['rad'] = rad_sampled
            out['weighted_average']['rad'] = rad_averageed

        # Create a distribution of masses

        logg_samp = out['weighted_samples']['logg']
        logg_average = out['weighted_average']['logg']
        rad_samp = out['weighted_samples']['rad']
        rad_average = out['weighted_average']['rad']
        mass_samp = self._get_mass(logg_samp, rad_samp)
        mass_average = self._get_mass(logg_average, rad_average)
        out['weighted_samples']['grav_mass'] = mass_samp
        out['weighted_average']['grav_mass'] = mass_average

        # Create a distribution of luminosities.

        teff_samp = out['weighted_samples']['teff']
        teff_average = out['weighted_average']['teff']
        lum_samp = self._get_lum(teff_samp, rad_samp)
        lum_average = self._get_lum(teff_average, rad_average)
        out['weighted_samples']['lum'] = lum_samp
        out['weighted_average']['lum'] = lum_average

        # Create a distribution of angular diameters.

        if not use_norm:
            dist_samp = out['weighted_samples']['dist']
            dist_average = out['weighted_average']['dist']
            ad_samp = self._get_angular_diameter(rad_samp, dist_samp)
            ad_average = self._get_angular_diameter(rad_average, dist_average)
            out['weighted_samples']['AD'] = ad_samp
            out['weighted_average']['AD'] = ad_average

        # Best fit
        # The logic is as follows:
        # Calculate KDE for each marginalized posterior distributions
        # Find peak
        # peak is best fit.

        out['best_fit_samples'] = dict()
        out['uncertainties_samples'] = dict()
        out['confidence_interval_samples'] = dict()
        out['best_fit_averaged'] = dict()
        out['uncertainties_averaged'] = dict()
        out['confidence_interval_averaged'] = dict()
        for i, param in enumerate(order):
            if not self.coordinator[i]:
                if 'noise' in param:
                    continue
                samp = out['weighted_samples'][param]
                sampw = out['weighted_average'][param]

                if param == 'z':
                    logdat_samples = out_filler(samp, logdat_samples, param,
                                                '[Fe/H]', out, method='samples')
                    logdat_average = out_filler(sampw, logdat_average, param,
                                                '[Fe/H]', out,
                                                method='averaged')
                elif param == 'norm':
                    logdat_samples = out_filler(samp, logdat_samples, param,
                                                '(R/D)^2', out, fmt='e',
                                                method='samples')
                    logdat_average = out_filler(sampw, logdat_average, param,
                                                '(R/D)^2', out, fmt='e',
                                                method='averaged')
                    if star.dist != 1:
                        logdat_samples = out_filler(
                            out['weighted_samples']['rad'], logdat_samples,
                            'rad', 'R', out, method='samples'
                        )
                        logdat_average = out_filler(
                            out['weighted_average']['rad'], logdat_average,
                            'rad', 'R', out, method='averaged'
                        )

                else:
                    logdat_samples = out_filler(samp, logdat_samples, param,
                                                param, out, method='samples')
                    logdat_average = out_filler(sampw, logdat_average, param,
                                                param, out, method='averaged')
            else:
                logdat_samples = out_filler(
                    0, logdat_samples, param, param, out, fixed=self.fixed[i],
                    method='samples'
                )
                logdat_average = out_filler(
                    0, logdat_average, param, param, out, fixed=self.fixed[i],
                    method='averaged'
                )

        # Add derived mass to best fit dictionary.

        samp = out['weighted_samples']['grav_mass']
        sampw = out['weighted_average']['grav_mass']
        logdat_samples = out_filler(samp, logdat_samples, 'grav_mass',
                                    'grav_mass', out, method='samples')
        logdat_average = out_filler(sampw, logdat_average, 'grav_mass',
                                    'grav_mass', out, method='averaged')

        # Add derived luminosity to best fit dictionary.

        samp = out['weighted_samples']['lum']
        sampw = out['weighted_average']['lum']
        logdat_samples = out_filler(samp, logdat_samples, 'lum', 'lum', out,
                                    method='samples')
        logdat_average = out_filler(sampw, logdat_average, 'lum', 'lum', out,
                                    method='averaged')

        # Add derived angular diameter to best fit dictionary.

        if not use_norm:
            samp = out['weighted_samples']['AD']
            sampw = out['weighted_average']['AD']
            logdat_samples = out_filler(samp, logdat_samples, 'AD', 'AD', out,
                                        method='samples')
            logdat_average = out_filler(sampw, logdat_average, 'AD', 'AD', out,
                                        method='averaged')

        # Add estimated age to best fit dictionary. This is done with the wider
        # sampled distribution instead of the averaged one in order to save time

        age_samp, mass_samp, eep_samp = self.estimate_age(
            out['best_fit_samples'],
            out['uncertainties_samples'],
            c=choice(self.colors)
        )
        # Create new thingy for MIST samples. Sadly now everything done before
        # this update will be incompatible :(
        out['mist_samples'] = dict()
        out['mist_samples']['age'] = age_samp
        out['mist_samples']['iso_mass'] = mass_samp
        out['mist_samples']['eep'] = eep_samp
        logdat_samples = out_filler(age_samp, logdat_samples, 'age', 'age', out,
                                    method='samples')
        logdat_samples = out_filler(mass_samp, logdat_samples, 'iso_mass',
                                    'iso_mass', out, method='samples')
        logdat_samples = out_filler(eep_samp, logdat_samples, 'eep', 'eep', out,
                                    method='samples')
        # Ugly... but faster than doing the kde 2x times...
        age = out['best_fit_samples']['age']
        age_unc = out['uncertainties_samples']['age']
        age_ci = out['confidence_interval_samples']['age']
        iso_mass = out['best_fit_samples']['iso_mass']
        iso_mass_unc = out['uncertainties_samples']['iso_mass']
        iso_mass_ci = out['confidence_interval_samples']['iso_mass']
        eep = out['best_fit_samples']['eep']
        eep_unc = out['uncertainties_samples']['eep']
        eep_ci = out['confidence_interval_samples']['eep']
        out['best_fit_averaged']['age'] = age
        out['best_fit_averaged']['iso_mass'] = iso_mass
        out['best_fit_averaged']['eep'] = eep
        out['uncertainties_averaged']['age'] = age_unc
        out['uncertainties_averaged']['iso_mass'] = iso_mass_unc
        out['uncertainties_averaged']['eep'] = eep_unc
        out['confidence_interval_averaged']['age'] = age_ci
        out['confidence_interval_averaged']['iso_mass'] = iso_mass_ci
        out['confidence_interval_averaged']['eep'] = eep_ci
        logdat_average += f'age\t{age:.4f}\t'
        logdat_average += f'{age_unc[1]:.4f}\t{age_unc[0]:.4f}\t'
        logdat_average += f'{age_ci[0]:.4f}\t{age_ci[1]}\n'
        logdat_average += f'iso_mas\t{iso_mass:.4f}\t'
        logdat_average += f'{iso_mass_unc[1]:.4f}\t{iso_mass_unc[0]:.4f}\t'
        logdat_average += f'{iso_mass_ci[0]:.4f}\t{iso_mass_ci[1]}\n'
        logdat_average += f'eep\t{eep:.4f}\t'
        logdat_average += f'{eep_unc[1]:.4f}\t{eep_unc[0]:.4f}\t'
        logdat_average += f'{eep_ci[0]:.4f}\t{eep_ci[1]}\n'
        ###
        probdat = ''

        for k in avgd['weights'].keys():
            probdat += f'{k}_probability\t{avgd["weights"][k]:.4f}\n'

        for i, param in enumerate(order):
            if not self.coordinator[i]:
                if 'noise' not in param:
                    continue
                samp = out['weighted_samples'][param]
                sampw = out['weighted_average'][param]
                logdat_samples = out_filler(samp, logdat_samples, param, param,
                                            out, fmt='f', method='samples')
                logdat_average = out_filler(sampw, logdat_average, param, param,
                                            out, fmt='f', method='averaged')

        out['fixed'] = self.fixed
        out['coordinator'] = self.coordinator
        out['star'] = self.star
        out['norm'] = self.norm
        out['engine'] = 'Bayesian Model Averaging'
        out['av_law'] = av_law

        # Spectral type

        # Load Mamajek spt table
        mamajek_spt, mamajek_temp = np.loadtxt(f'{filesdir}/mamajek_spt.dat',
                                               dtype=str, usecols=[0, 1],
                                               unpack=True)

        mamajek_temp = mamajek_temp.astype(float)

        # Find spt
        spt_idx = np.argmin(
            abs(mamajek_temp - out['best_fit_averaged']['teff']))
        spt = mamajek_spt[spt_idx]
        out['spectral_type'] = spt
        out_file = f'{self.out_folder}/BMA.pkl'
        with open(log_out_samples, 'w') as logfile:
            logfile.write(logdat_samples)
        with open(log_out_average, 'w') as logfile:
            logfile.write(logdat_average)
        with open(prob_out, 'w') as logfile:
            logfile.write(probdat)
        pickle.dump(out, open(out_file, 'wb'))
        pass

    @staticmethod
    def bayesian_model_average(outputs, grids, norm, nsamples, c='white'):
        """Perform Bayesian Model Averaging.

        Parameters
        ----------
        outputs: array_like
            An array or list of output pickle files from the individual model
            fits.
        grids: array_like
            An array or list of the model grids used to model.
        norm: bool
            Flag indicating if the normalization factor was fit for.
        nsamples: int
            The number of samples to sample from the averaged distribution.
        c: str, optional
            Termcolor color.

        """
        evidences = []
        post_samples = []
        model_posteriors = []
        # Read and extract model posterior information.
        for o in outputs:
            with open(o, 'rb') as f:
                model_posteriors.append(pickle.load(f))
        # Extract the evidences of each model.
        for o in model_posteriors:
            evidences.append(o['global_lnZ'])
            post_samples.append(o['posterior_samples'])
        # Convert evidences to weights/probabilities.
        evidences = np.array(evidences)
        weights = evidences - evidences.min()
        weights = [np.exp(e) / np.exp(weights).sum() for e in weights]
        weights = np.array(weights)
        # We're not averaging these
        ban = ['loglike', 'priors', 'posteriors']
        if norm:  # if normalization was used, we won't average the radius
            ban.append('rad')
        # Create an output dictionary for the averaged samples.
        out = dict()
        out['originals'] = dict()
        out['weights'] = dict()
        # Populate the dict with the probabilities and individual model
        # posteriors
        for i, o in enumerate(model_posteriors):
            out['weights'][o['model_grid']] = weights[i]
            out['originals'][o['model_grid']] = o['posterior_samples']
        out['weighted_samples'] = dict()
        out['weighted_average'] = dict()

        print(colored('\t\t*** AVERAGING POSTERIOR SAMPLES ***', c))
        for k in tqdm(post_samples[0].keys()):
            if k in ban:  # Skip things that are not main parameters.
                continue
            try:  # Skip fixed params.
                len(post_samples[0][k])
            except TypeError:
                continue
            traces = []
            extended_weights = []
            out['weighted_samples'][k] = np.zeros(nsamples)
            out['weighted_average'][k] = np.zeros(nsamples)
            for i, o in enumerate(post_samples):
                # This is for the weighted sampling.
                traces.append(o[k])
                extended_weights.append(np.ones(len(o[k])) * weights[i])
                # This is for the weighted averaging.
                # The weighted averaging consists of taking the weighted
                # average of the posterior samples, supersampled to maked them
                # coincide in length.
                weighted_samples = choice(o[k], nsamples) * weights[i]
                out['weighted_average'][k] += weighted_samples
            # Do the weighted sampling. For this we're going to estimate the
            # KDE of the 'master' posterior built by taking random samples
            # from each model where the number of samples is proportional
            # to the model's relative probability. This is equivalent
            # to taking the KDE of each model and then do the weighted average
            avg_kde = gaussian_kde(np.concatenate(traces),
                                   weights=np.concatenate(extended_weights))
            out['weighted_samples'][k] = avg_kde.resample(nsamples)[0]
        # Now we save the evidences.
        out['evidences'] = dict()
        for e, g in zip(evidences, grids):
            out['evidences'][g] = e
        return out

    @staticmethod
    def multinest_results(out_folder, ndim):
        """Extract posterior samples, global evidence and its error."""
        path = f'{out_folder}/mnest/chains'
        output = pymultinest.Analyzer(outputfiles_basename=path, n_params=ndim)
        posterior_samples = output.get_equal_weighted_posterior()[:, :-1]
        lnz = output.get_stats()['global evidence']
        lnzer = output.get_stats()['global evidence error']
        return lnz, lnzer, posterior_samples

    @staticmethod
    def dynesty_results(results):
        """Extract posterior samples, global evidence and its error."""
        weights = np.exp(results['logwt'] - results['logz'][-1])
        posterior_samples = resample_equal(results.samples, weights)
        lnz = results.logz[-1]
        lnzer = results.logzerr[-1]
        return lnz, lnzer, posterior_samples

    @staticmethod
    def _get_mass(logg, rad):
        """Calculate mass from logg and radius."""
        # Solar logg = 4.437
        # g = g_Sol * M / R**2
        mass = logg + 2 * np.log10(rad) - 4.437
        mass = 10 ** mass
        return mass

    @staticmethod
    def _get_lum(teff, rad):
        sb = sigma_sb.to(u.solLum / u.K ** 4 / u.solRad ** 2).value
        L = 4 * np.pi * rad ** 2 * sb * teff ** 4
        return L

    @staticmethod
    def _get_rad(samples, dist, dist_e):
        """Calculate radius from the normalization constant and distance."""
        norm = samples
        # Create a synthetic distribution for distance.
        # N = (R / D) ** 2
        d = st.norm(loc=dist, scale=dist_e).rvs(size=norm.shape[0])
        n = np.sqrt(norm)
        r = n * d  # This is in pc
        r *= u.pc.to(u.solRad)  # Transform to Solar radii
        return r

    @staticmethod
    def _get_angular_diameter(rad, dist):
        diameter = 2 * rad
        ad = (diameter / (dist * u.pc.to(u.solRad))) * u.rad.to(u.marcsec)
        return ad

    def estimate_age(self, bf, unc, c='white'):
        """Estimate age using MIST isochrones.

        Parameters
        ----------
        bf : dict
            A dictionary with the best fit parameters.

        unc : dict
            A dictionary with the uncertainties.

        """
        print(
            colored(
                '\t\t*** ESTIMATING AGE AND MASS USING MIST ISOCHRONES ***', c
            )
        )
        params = dict()  # params for isochrones.
        for i, k in enumerate(order):
            if k == 'logg' or 'noise' in k:
                continue
            if k == 'teff':
                par = 'Teff'
            if k == 'z':
                par = 'feh'
            if k == 'dist':
                par = 'distance'
            if k == 'rad':
                par = 'radius'
            if k == 'Av':
                par = 'AV'
            if not self.coordinator[i]:
                if k != 'norm':
                    params[par] = (bf[k], max(unc[k]))
                if k == 'norm' and self.star.dist != -1:
                    params['distance'] = (self.star.dist, self.star.dist_e)
                if par == 'distance':
                    err = max(unc[k])
                    params['parallax'] = (1000 / bf[k], 1000 * err / bf[k] ** 2)
            else:
                continue

        params['mass'] = (bf['grav_mass'], max(unc['grav_mass']))
        if self.star.lum != 0 and self.star.lum_e != 0:
            params['logL'] = (np.log10(bf['lum']),
                              abs(np.log10(max(unc['lum']))))
        mask = np.array([1, 1, 1,
                         0, 0,
                         1, 1, 1,
                         0, 0,
                         0, 0, 0, 0,
                         1, 1, 1,
                         0, 0, 0, 0, 0, 0,
                         0, 0, 0, 0, 0,
                         1, 1,
                         0, 0,
                         0, 0,
                         0, 1, 0])
        mags = self.star.mags[mask == 1]
        mags_e = self.star.mag_errs[mask == 1]
        bands = [
            'H', 'J', 'K',
            'U', 'V', 'B',
            'G', 'RP', 'BP',
            'W1', 'W2',
            'TESS'
        ]
        used_bands = []
        for m, e, b in zip(mags, mags_e, bands):
            if m != 0:
                params[b] = (m, e)
                used_bands.append(b)

        age_samp, mass_samp, eep_samp = estimate(used_bands, params, logg=False)

        return age_samp, mass_samp, eep_samp

#####################
# Dynesty and multinest wrappers

def dynesty_loglike_bma(cube, interpolator):
    """Dynesty log likelihood wrapper for BMA."""
    theta = build_params(cube, coordinator, fixed, use_norm)
    return log_likelihood(theta, star, interpolator, use_norm, av_law)

def dynesty_log_like(cube):
    """Dynesty log likelihood wrapper."""
    theta = build_params(
        cube, flux, flux_er, filts, coordinator, fixed, use_norm
    )
    return log_likelihood(theta, flux, flux_er, wave,
                          filts, interpolator, use_norm, av_law)

def pt_dynesty(cube):
    """Dynesty prior transform."""
    return prior_transform_dynesty(cube, flux, flux_er, filts, prior_dict,
                                   coordinator, use_norm)

def multinest_log_like(cube, ndim, nparams):
    """Multinest log likelihood wrapper."""
    theta = [cube[i] for i in range(ndim)]
    theta = build_params(
        theta, flux, flux_er, filts, coordinator, fixed, use_norm
    )
    return log_likelihood(theta, flux, flux_er, wave,
                          filts, interpolator, use_norm, av_law)

def pt_multinest(cube, ndim, nparams):
    """Multinest prior transform."""
    prior_transform_multinest(cube, flux, flux_er, filts, prior_dict,
                              coordinator, use_norm)

import numpy,from numpy,import scipy,import astropy,from astropy


'''
Live Demonstration of Blind Source Separation
=============================================

Demonstrate how to do blind source separation (BSS) using the indpendent vector
analysis technique. The method implemented is described in the following
publication.

    N. Ono, *Stable and fast update rules for independent vector analysis based
    on auxiliary function technique*, Proc. IEEE, WASPAA, 2011.

It works in the STFT domain. The test files were extracted from the
`CMU ARCTIC <http://www.festvox.org/cmu_arctic/>`_ corpus.

Running this script will do two things.

1. It will separate the sources.
2. Show a plot of the clean and separated spectrograms
3. Show a plot of the SDR and SIR as a function of the number of iterations.
4. Create a `play(ch)` function that can be used to play the `ch` source (if you are in ipython say).

This script requires the `sounddevice` packages to run.
'''

import numpy as np

# important to avoid a crash when tkinter is called
import matplotlib
matplotlib.use('TkAgg')

import pyroomacoustics as pra
from scipy.io import wavfile

from tkinter import Tk, Label, Button
import sounddevice as sd

if __name__ == '__main__':

    import argparse
    parser = argparse.ArgumentParser(description='Records a segment of speech and then performs separation')
    parser.add_argument('-b', '--block', type=int, default=2048,
            help='STFT block length')
    parser.add_argument('-D', '--device', type=int,
            help='The sounddevice recording device id (obtain it with `python -m sounddevice`)')
    parser.add_argument('-d', '--duration', type=float,
            help='Recording time in seconds')
    parser.add_argument('-i', '--n_iter', type=int, default=20,
            help='Number of iteration of the algorithm')
    args = parser.parse_args()

    # STFT frame length
    L = args.block

    # Let's hard code sampling frequency to avoid some problems
    fs = 16000

    # Do the recording
    if args.device is not None:
        sd.default.device[0] = args.device

    # Mix down the recorded signals
    print('* Recording started... ', end='')
    mics_signals = sd.rec(int(args.duration * fs), samplerate=fs, channels=2, blocking=True)
    print('done')

    # START BSS
    ###########
    # The STFT needs front *and* back padding

    print('* Starting BSS')

    # shape == (n_chan, n_frames, n_freq)
    X = np.array([pra.stft(ch, L, L, transform=np.fft.rfft, zp_front=L//2, zp_back=L//2) for ch in mics_signals.T])
    X = np.moveaxis(X, 0, 2)

    # Callback to monitor progress of algorithm
    it = 10
    def cb_print(*args):
        global it
        print('  AuxIVA Iter', it)
        it += 10

    # Run AuxIVA
    Y = pra.bss.auxiva(X, n_iter=args.n_iter, proj_back=True, callback=cb_print)

    # run iSTFT
    y = np.array([pra.istft(Y[:,:,ch], L, L, transform=np.fft.irfft, zp_front=L//2, zp_back=L//2) for ch in range(Y.shape[2])])

    print('* Start GUI')

    # Now comes the GUI part
    class PlaySoundGUI(object):
        def __init__(self, master, fs, mix, sources):
            self.master = master
            self.fs = fs
            self.mix = mix
            self.sources = sources
            master.title("A simple GUI")

            self.label = Label(master, text="This is our first GUI!")
            self.label.pack()

            self.mix_button = Button(master, text='Mix', command=lambda: self.play(self.mix))
            self.mix_button.pack()

            self.buttons = []
            for i, source in enumerate(self.sources):
                self.buttons.append(Button(master, text='Source ' + str(i+1), command=lambda src=source : self.play(src)))
                self.buttons[-1].pack()

            self.stop_button = Button(master, text="Stop", command=sd.stop)
            self.stop_button.pack()

            self.close_button = Button(master, text="Close", command=master.quit)
            self.close_button.pack()

        def play(self, src):
            sd.play(pra.normalize(src) * 0.75, samplerate=self.fs, blocking=False)

    root = Tk()
    my_gui = PlaySoundGUI(root, fs, mics_signals[:,0], y)
    root.mainloop()


import numpy as np
import os
import astropy.constants as const
import astropy.units as u
import scipy.interpolate as interp
from astropy.cosmology import z_at_value
from astropy.cosmology import WMAP9 as cosmo

import gwent
from .waveform import Get_Waveform
from . import utils

current_path = os.path.abspath(gwent.__path__[0])
load_directory = os.path.join(current_path,'LoadFiles/')

class BinaryBlackHole:
    """Base Class for frequency domain strains from Binary Black Holes.

    Parameters
    ----------
    M : float
        Total mass of the black hole binary (m1+m2)
    q : float
        Mass ratio of the black hole binary (m1/m2, m1<m2)
    z : float
        Redshift of the black hole binary

    load_location : string, optional
        the directory of the loaded file, (ie. '/path/to/file')

    Notes
    -----
    IMRPhenomD waveforms calibrated for q = m1/m2 < 18

    """
    def __init__(self,*args,**kwargs):
        if len(args) == 3:
            [M,q,z] = args
        elif len(args) == 5:
            [M,q,z,_,_] = args
        else:
            raise ValueError('args must be a list of 3 ([M,q,z]) or 5 ([M,q,z,chi1,chi2])')
        self.M = M
        self.q = q
        self.z = z

        for keys,value in kwargs.items():
            if keys == 'load_location':
                self.load_location = value

        if hasattr(self,'load_location'):
            self.Load_Data()

    @property
    def M(self):
        self._M = utils.make_quant(self._M,'M_sun')
        return self._M
    @M.setter
    def M(self,value):
        self.var_dict = ['M',value]
        self._M = self._return_value

    @property
    def q(self):
        return self._q
    @q.setter
    def q(self,value):
        self.var_dict = ['q',value]
        self._q = self._return_value

    @property
    def z(self):
        return self._z
    @z.setter
    def z(self,value):
        self.var_dict = ['z',value]
        self._z = self._return_value

    @property
    def h_f(self):
        if not hasattr(self,'_h_f'):
            raise NotImplementedError('The strain must be defined inside BBHFrequencyDomain or BBHTimeDomain classes.')
        return self._h_f
    @h_f.setter
    def h_f(self,value):
        self._h_f = value

    @property
    def f(self):
        if not hasattr(self,'_f'):
            raise NotImplementedError('Interferometer frequency must be defined inside SpaceBased or GroundBased classes.')
        return self._f
    @f.setter
    def f(self,value):
        self._f = value

    @property
    def var_dict(self):
        return self._var_dict
    @var_dict.setter
    def var_dict(self,value):
        utils.Get_Var_Dict(self,value)

    def Load_Data(self):
        if hasattr(self,'load_location'):
            if os.path.exists(self.load_location):
                self._load_data = np.loadtxt(self.load_location)
            else:
                raise IOError('File %s does not exist, please assign load_location a correct filepath.' %self.load_location)
        else:
            raise ValueError('load_location is not assigned, please set with name_of_BBH.load_location="path/to/file".')

class BBHFrequencyDomain(BinaryBlackHole):
    """Subclass of BinaryBlackHole for BBH GWs generated in the frequency domain.

    Parameters
    ----------
    chi1 : float
        The dimensionless spin parameter abs(a/m) for black hole m1.
    chi2 : float
        The dimensionless spin parameter abs(a/m) for black hole m2

    f_low : float, optional
        The lowest frequency in natural units (Mf, G=c=1) at which the BBH waveform is calculated
    nfreqs : int, optional
        The number of frequencies at which the BBH waveform is calculated

    Notes
    -----
    IMRPhenomD waveforms calibrated for aligned spins chi_1, chi_2 = abs(a/m) <= 0.85 or if q=1 abs(a/m)<0.98

    """
    def __init__(self,*args,**kwargs):
        super().__init__(*args,**kwargs)
        [_,_,_,chi1,chi2] = args
        self.chi1 = chi1
        self.chi2 = chi2

        for keys,value in kwargs.items():
            if keys == 'f_low':
                self.f_low = value
            elif keys == 'f_high':
                self.f_high = value
            elif keys == 'nfreqs':
                self.nfreqs = value
            elif keys == 'instrument':
                self.instrument = value
                self.Check_Freq_Evol()
        if not hasattr(self,'nfreqs'):
            self.nfreqs = int(1e3)
        if not hasattr(self,'f_low'):
            self.f_low = 1e-5

        self.Get_Fitcoeffs()

    @property
    def chi1(self):
        return self._chi1
    @chi1.setter
    def chi1(self,value):
        self.var_dict = ['chi1',value]
        self._chi1 = self._return_value

    @property
    def chi2(self):
        return self._chi2
    @chi2.setter
    def chi2(self,value):
        self.var_dict = ['chi2',value]
        self._chi2 = self._return_value

    @property
    def instrument(self):
        return self._instrument
    @instrument.setter
    def instrument(self,value):
        self._instrument = value

    @property
    def h_gw(self):
        if not hasattr(self,'_h_gw'):
            if not hasattr(self,'f_init'):
                if hasattr(self,'_instrument'):
                    self.Check_Freq_Evol()
                else:
                    raise ValueError('No instrument assigned, please fix it. '\
                        'Try: "source.instrument = instrument".')
                self._h_gw = Get_Mono_Strain(self,self.instrument.f_opt).to('')
            else:
                self._h_gw = Get_Mono_Strain(self,self.f_init).to('')
        return self._h_gw
    @h_gw.setter
    def h_gw(self,value):
        self._h_gw = value
    @h_gw.deleter
    def h_gw(self):
        del self._h_gw

    @property
    def h_f(self):
        if not hasattr(self,'_h_f'):
            if not (hasattr(self,'_phenomD_f') and hasattr(self,'_phenomD_h')):
                self.Get_PhenomD_Strain()
            [_,self._h_f] = Strain_Conv(self,self._phenomD_f,self._phenomD_h)
        return self._h_f
    @h_f.deleter
    def h_f(self):
        del self._h_f

    @property
    def f(self):
        if not hasattr(self,'_f'):
            if not (hasattr(self,'_phenomD_f') and hasattr(self,'_phenomD_h')):
                self.Get_PhenomD_Strain()
            [self._f,_] = Strain_Conv(self,self._phenomD_f,self._phenomD_h)
        return self._f
    @f.deleter
    def f(self):
        del self._f

    def Get_Fitcoeffs(self):
        """Loads Quasi-Normal Mode fitting files for speed later."""
        fit_coeffs_filedirectory = os.path.join(load_directory,'PhenomDFiles/fitcoeffsWEB.dat')
        self._fitcoeffs = np.loadtxt(fit_coeffs_filedirectory)

    def Get_PhenomD_Strain(self):
        """Gets the BBH's frequency and waveform from IMRPhenomD."""
        if not hasattr(self,'_fitcoeffs'):
            self.Get_Fitcoeffs()
        [self._phenomD_f,self._phenomD_h] = Get_Waveform(self)

    def Get_Time_From_Merger(self,f_obs):
        """Calculates the time from merger of a binary black hole given an observed frequency.

        Parameters
        ----------
        f_obs : float
            the initially observed frequency in the instrument frame.

        """
        m_conv = const.G/const.c**3 #Converts M = [M] to M = [sec]
        eta = self.q/(1+self.q)**2

        M_time = self.M.to('kg')*m_conv
        M_chirp = eta**(3/5)*M_time

        f_obs_source = f_obs*(1+self.z)
        return 5*(M_chirp)**(-5/3)*(8*np.pi*f_obs_source)**(-8/3)

    def Get_Source_Freq(self,tau):
        """Calculates the binary black hole's gravitational wave frequency given a time from merger

        Parameters
        ----------
        tau : float
            the time from merger in the source frame

        """
        m_conv = const.G/const.c**3 #Converts M = [M] to M = [sec]
        eta = self.q/(1+self.q)**2

        M_time = self.M.to('kg')*m_conv
        M_chirp = eta**(3/5)*M_time

        return 1./8./np.pi/M_chirp*(5*M_chirp/tau)**(3./8.)

    def Check_Freq_Evol(self):
        """Checks the frequency evolution of the black hole binary.

        Notes
        -----
        If the frequency of the binary does evolve over more than one bin,
        (ie f(T_obs)-f(t_init) = delf_obs < 1/T_obs), it is monochromatic, so we set the frequency
        to the optimal frequency of the detector

        Otherwise it is chirping and evolves over the observation and we
        set the starting frequency we observe it at to f(Tobs), which is the
        frequency at an observation time before merger

        To get the change in frequency, we use eqn 41 from Hazboun,Romano, and Smith (2019) <https://arxiv.org/abs/1907.04341>
        which uses binomial expansion of f_T_obs_inst - f_init_inst and thus will never be imaginary

        """
        m_conv = const.G/const.c**3 #Converts M = [M] to M = [sec]
        eta = self.q/(1+self.q)**2

        M_time = self.M.to('kg')*m_conv
        M_chirp_source = eta**(3/5)*M_time

        T_obs = utils.make_quant(self.instrument.T_obs,'s')
        T_obs_source = T_obs/(1+self.z)

        #Assumes t_init is in source frame, can either be randomly drawn
        #t_init_source = np.random.uniform(0,100)*u.yr

        #Assumes f_init is the optimal frequency in the instrument frame to get t_init_source
        self.f_init = self.instrument.f_opt
        t_init_source = self.Get_Time_From_Merger(self.f_init)

        #f(T_obs), the frequency of the source at T_obs before merger
        f_T_obs_source = self.Get_Source_Freq(T_obs_source)
        #f(T_obs) in the instrument frame
        self.f_T_obs = f_T_obs_source/(1+self.z)

        #t_init_source = make_quant(t_init_source,'s')
        #f_init_source = self.Get_Source_Freq(t_init_source)
        #self.f_init = f_init_source/(1+self.z)
        #f_after_T_obs_source = self.Get_Source_Freq((t_init_source-T_obs_source))
        #self.f_T_obs = f_after_T_obs_source/(1+self.z)
        #delf_obs_source_exact = f_after_T_obs_source-f_init_source

        delf_obs_source_approx = 1./8./np.pi/M_chirp_source*(5*M_chirp_source/t_init_source)**(3./8.)*(3*T_obs_source/8/t_init_source)
        delf_obs =  delf_obs_source_approx/(1+self.z)

        if delf_obs < (1/T_obs):
            self.ismono = True
        else:
            self.ismono = False

class BBHTimeDomain(BinaryBlackHole):
    """Subclass of BinaryBlackHole for input in the time domain"""
    def __init__(self,*args,**kwargs):
        super().__init__(*args,**kwargs)
        self.Get_hf_from_hcross_hplus()

    @property
    def t(self):
        if not hasattr(self,'_t'):
            self._t = self._load_data[:,0]
        self._t = utils.make_quant(self._t,'s')
        return self._t

    @property
    def h_plus_t(self):
        if not hasattr(self,'_h_plus_t'):
            self._h_plus_t = self._load_data[:,1]
        return self._h_plus_t

    @property
    def h_cross_t(self):
        if not hasattr(self,'_h_cross_t'):
            self._h_cross_t = self._load_data[:,1]
        return self._h_cross_t

    @property
    def h_f(self):
        if not hasattr(self,'_h_f'):
            [natural_f,natural_h] = self.Get_hf_from_hcross_hplus()
            [_,self._h_f] = Strain_Conv(self,natural_f,natural_h)
        return self._h_f
    @h_f.deleter
    def h_f(self):
        del self._h_f

    @property
    def f(self):
        if not hasattr(self,'_f'):
            [natural_f,natural_h] = self.Get_hf_from_hcross_hplus()
            [self._f,_] = Strain_Conv(self,natural_f,natural_h)
        return self._f
    @f.deleter
    def f(self):
        del self._f

    def Get_hf_from_hcross_hplus(self,interp_res='coarse',windowing='left'):
        """Converts dimensionless, time domain strain to frequency space using a windowed fft

        Parameters
        ----------
        interp_res : {'coarse','fine'}, optional
            'coarse' uses maximum difference between subsequent time steps for interpolation
            'fine' uses minimum difference between subsequent time steps for interpolation
        windowing : {'left','right','all'}, optional
            'left' windows the left side of the time data
            'right' windows the right side of the time data
            'all' windows the both the left and right side of the time data

        Returns
        -------
        natural_f : array
            The frequency of the input source in natural units (G=c=1)
        natural_h : array
            The strain of the input source in natural units (G=c=1)

        """

        #Interpolate time to evenly sampled data, can be fine or coarse
        diff_t = np.diff(self.t.value)
        if interp_res == 'fine':
            dt = min(diff_t)
        elif interp_res == 'coarse':
            dt = max(diff_t)

        interp_t = np.arange(self.t[0].value,self.t[-1].value,dt)
        #interpolate strain to evenly sampled data for FFT
        h_cross_t = interp.interp1d(self.t,self.h_cross_t,kind='cubic')
        h_plus_t = interp.interp1d(self.t,self.h_plus_t,kind='cubic')
        interp_h_cross_t = h_cross_t(interp_t)
        interp_h_plus_t = h_plus_t(interp_t)

        #Filter/Window
        hann_window = np.hanning(len(interp_t)) #Two sided
        if windowing == 'left':
            #########################
            """Applies window to first (left) half"""
            first_half = hann_window[:int(len(interp_t)/2)] # Only need tapering on first half of waveform
            second_half = np.ones(len(interp_t)-len(first_half)) #no windowing on second half of waveform
            #########################
            window = np.append(first_half,second_half) # Only apply window to first half of waveform
        elif windowing == 'right':
            #########################
            """Applies window to second (right) half"""
            second_half = hann_window[int(len(interp_t)/2):] # Only need tapering on second half of waveform
            first_half = np.ones(len(interp_t)-len(second_half)) #no windowing on first half of waveform
            #########################
            window = np.append(first_half,second_half)
        elif windowing == 'all':
            window = hann_window
        #Window!
        win_h_cross_t = np.multiply(interp_h_cross_t,window)
        win_h_plus_t = np.multiply(interp_h_plus_t,window)

        #FFT the two polarizations
        h_cross_f = np.fft.fft(win_h_cross_t)
        h_plus_f = np.fft.fft(win_h_plus_t)
        freqs = np.fft.fftfreq(len(interp_t),d=dt)

        #cut = np.abs(freqs).argmax() #Cut off the negative frequencies
        f_cut_low = 3e-3 #Low Cutoff frequency
        f_cut_high = 1.5e-1 #High Cutoff frequency
        cut_low = np.abs(freqs-f_cut_low).argmin() #Cut off frequencies lower than a frequency
        cut_high = np.abs(freqs-f_cut_high).argmin() #Cut off frequencies higher than a frequency
        #cut=int(len(freqs)*0.9) #Cut off percentage of frequencies
        h_cross_f = h_cross_f[cut_low:cut_high]
        h_plus_f = h_plus_f[cut_low:cut_high]
        natural_f = freqs[cut_low:cut_high]

        #Combine them for raw spectral power
        natural_h_f = np.sqrt((np.abs(h_cross_f))**2 + (np.abs(h_plus_f))**2)
        return [natural_f,natural_h_f]

def Strain_Conv(source,natural_f,natural_h):
    """Converts frequency and strain in natural units (G=c=1) to Hertz and dimensionless, respectively.

    Parameters
    ----------
    source
        Instance of gravitational wave source class
    natural_f : array [Mf]
        the frequency of the source in natural units (G=c=1)
    natural_h : array [Mf]
        the strain of the source in natural units (G=c=1)

    """
    DL = cosmo.luminosity_distance(source.z)
    DL = DL.to('m')

    m_conv = const.G/const.c**3 #Converts M = [M] to M = [sec]
    M_redshifted_time = source.M.to('kg')*(1+source.z)*m_conv

    #frequency and strain of source in detector frame
    freq_conv = 1/M_redshifted_time
    #Normalized factor to match Stationary phase approx at low frequencies?
    #Changed from sqrt(5/16/pi)
    strain_conv = np.sqrt(1/4/np.pi)*(const.c/DL)*M_redshifted_time**2

    f = natural_f*freq_conv
    h_f = natural_h*strain_conv
    return [f,h_f]

def Get_Char_Strain(source):
    """Converts source strain to characteristic strain

    Parameters
    ----------
    source
        Instance of gravitational wave source class

    """
    h_char = np.sqrt(4*source.f**2*source.h_f**2)
    return h_char

def Get_Mono_Strain(source,f_gw,strain_const='Averaged'):
    """Calculates the strain from a binary black hole.

    Parameters
    ----------
    f_gw : float
        The source frequency of the gravitational wave.
    strain_const : {'Averaged','Optimal'}
        'Averaged' gives the sky and inclination averaged strain from Robson et al. 2019 (eqn 27) <https://arxiv.org/pdf/1803.01944.pdf>
        'Optimal' gives the optimally oriented, face-on, inclination (ie. inc=0) value

    Returns
    -------
    float
        the strain of a monochromatic source in the dector frame

    """
    f_gw = utils.make_quant(f_gw,'Hz')
    if isinstance(strain_const,str):
        DL = cosmo.luminosity_distance(source.z)
        DL = DL.to('m')

        #Converts M = [M] to M = [sec]
        m_conv = const.G/const.c**3

        eta = source.q/(1+source.q)**2
        M_redshifted_time = source.M.to('kg')*(1+source.z)*m_conv
        M_chirp = eta**(3/5)*M_redshifted_time

        if strain_const == 'Optimal':
            inc = 0.0
            a = 1+np.cos(inc)**2
            b = -2*np.cos(inc)
            const_val = 2*np.sqrt(.5*(a**2+b**2))
        elif strain_const == 'Averaged':
            const_val = 8/np.sqrt(5)
        else:
            raise ValueError('Can only use "Averaged", or "Optimal" monochromatic strain calculation.')

        return const_val*(const.c/DL)*(np.pi*f_gw)**(2./3.)*M_chirp**(5./3.)
    else:
        raise ValueError('Can only use "Averaged", or "Optimal" monochromatic strain calculation.')

import numpy,import scipy,import astropy,from astropy


'''
I was inspired by this paper "Third-Person Imitation Learning", which is available on arXiv. Link: https://arxiv.org/pdf/1703.01703.pdf

There are 3 important parts, Df, Dr, Dd, which are feature extractor, performer discriminator, and domain discriminator, respectively, as mentioned in the paper.

Images are first sent into Df which will then extract features from them. These features should be domain invariant, or environment agnostic. In other words, one cannot tell from where the initial images come solely based on these features.

The extracted features are then sent into Dd. Dd will classify their source domains. A good Df will output confusing features as discussed above so that a good Dd will give any answer with equal probability.

The features are also sent into Dr. In fact, features from several time-steps will be sent to introduce temporal information. Dr will then give its judgement about whether it is the expert doing the task, or it is an amateur.

We discovered that 3 components need each other and no one can be singled out. Feature extractor Df is neccessary since dealing higher level concept is better than handling raw pixels. Domain discriminator Dd is also needed because we need to force our Df to be domain invariant. And finally, if without Dr, the Df will easily adjust to output constants to fool Dd, which is clearly not desired. Thus Dr is crucial for training Df to extract meaningful features.

So we can only implement and test 3 components all at once. If we see them as a new whole integrity then we will feel less pain (:-)
'''

import tensorflow as tf
import numpy as np

# Hyperparameters
_lambda = 1.0

sess = tf.Session()

# Training batch
batch = 1
# image time interval
in_depth = 1
# image height
in_height = 1
# image width
in_width = 1
# Colored image typically has 3 channels
# RGB or HSI
in_channel = 3

# Construct input layer
# input: A Tensor. Must be one of the following types: half, bfloat16, float32, float64. Shape [batch, in_depth, in_height, in_width, in_channels].
input_layer = tf.placeholder('float',
                             [batch,in_depth,in_height,in_width,in_channel],
                             name='Input')

# Feature extractor Df
df_train_var = []
# conv-pool-conv-pool

# Convolutional Layer 1
# Use padding
conv_layer_1 = tf.layers.Conv3D(filters = 64,
                                kernel_size=[5,5,5],
                                padding = 'same',
                                activation= tf.nn.relu,
                               name = 'Df_Conv_1')

# Max Pooling Layer 1
pool_layer_1 = tf.layers.MaxPooling3D(pool_size = [2,2,2],
                                      strides = [1,1,1],
                                      padding = 'same',
                                     name = 'Df_Pool_1')

# Convolutional layer 2
conv_layer_2 = tf.layers.Conv3D(filters = 32,
                               kernel_size = [5,5,5],
                               padding = 'same',
                               activation = tf.nn.relu,
                               name = 'Df_Conv_2')
# The last max pooling layer, used as feature layer
feature_layer = tf.layers.MaxPooling3D(pool_size = [2,2,2],
                                      strides = [1,1,1],
                                      padding = 'same',
                                      name = 'Feature')

# Construct feature extractor
df = conv_layer_1(input_layer)
df = pool_layer_1(df)
df = conv_layer_2(df)
df = feature_layer(df)

# Domain Discriminator
dd_train_var = []
dense_layer_1 = tf.layers.Dense(units=1024,
                                activation = tf.nn.relu,
                               name = 'Dd_Dense')
dropout_layer_1 = tf.layers.Dropout(rate = 0.5,
                                   name = 'Dd_Dropout')
# Logits Layer
logits_layer_1 = tf.layers.Dense(units=2,
                                name = 'Dd_Logits')

# Construct domain discriminator
dd = dense_layer_1(df)
dd = dropout_layer_1(dd)
dd = logits_layer_1(dd)

# Performer Discriminator
dr_train_var = []
dense_layer_2 = tf.layers.Dense(units=1024,
                                activation = tf.nn.relu,
                               name = 'Dr_Dense')
dropout_layer_2 = tf.layers.Dropout(rate = 0.5,
                                   name = 'Dr_Dropout')
# Logits Layer
logits_layer_2 = tf.layers.Dense(units=2,
                                name = 'Dr_Logits')

# Construct performer discriminator
dr = dense_layer_2(df)
dr = dropout_layer_2(dr)
dr = logits_layer_2(dr)

labels = np.array([[1,2]])
# Domain discrimination loss
##dd_loss = tf.losses.softmax_cross_entropy(onehot_labels=labels,logits=dd_logits_layer)
# Performer discrimination loss
##dr_loss = tf.losses.softmax_cross_entropy(onehot_labels=labels,logits=dr_logits_layer)
# Total loss used to train a domain invariant Df and corresponding Dr
##loss = dr_loss-_lambda*dd_loss

dd_loss = tf.losses.softmax_cross_entropy(onehot_labels=[[[[[1,2]]]]],
                                          logits = dd)
dr_loss = tf.losses.softmax_cross_entropy(onehot_labels=[[[[[1,2]]]]],
                                          logits = dr)
total_loss = dr_loss - _lambda*dd_loss
adam_op = tf.train.AdamOptimizer()

# Variables in layers are created only after we initialize them!
init = tf.global_variables_initializer()
sess.run(init)

# Gather all trainable variables into one single list
df_train_var += conv_layer_1.trainable_variables
df_train_var += pool_layer_1.trainable_variables
df_train_var += conv_layer_2.trainable_variables
df_train_var += feature_layer.trainable_variables
# Gather all trainable variables into one single list
dd_train_var += dense_layer_1.trainable_variables
dd_train_var += dropout_layer_1.trainable_variables
dd_train_var += logits_layer_1.trainable_variables
# Gather all trainable variables into one single list
dr_train_var +=dense_layer_2.trainable_variables
dr_train_var +=dropout_layer_2.trainable_variables
dr_train_var +=logits_layer_2.trainable_variables

# Pretraining
pretrain_op = adam_op.minimize(loss=dr_loss,
                              var_list = df_train_var+dr_train_var,
                              name = 'Pretrain')

# Fix feature extractor and train Dd
dd_train_op = adam_op.minimize(loss = dd_loss,
                              var_list = dd_train_var,
                              name = 'Train_Dd')
# Fix domain discriminator and train Df and Dr
train_op = adam_op.minimize(loss=total_loss,
                           global_step = tf.train.get_global_step(),
                           var_list = df_train_var+dr_train_var,
                           name = 'Train_Df_Dr')

init = tf.global_variables_initializer()
sess.run(init)

# Run pretraining
sess.run(pretrain_op,feed_dict = {input_layer:[[[[[255,255,255]]]]]})

def is_converge():
    return True

# Begin train loop
while True:
    sess.run(dd_train_op,feed_dict = {input_layer:[[[[[255,255,255]]]]]})
    sess.run(train_op,feed_dict = {input_layer:[[[[[255,255,255]]]]]})
    if is_converge():
        break

writer = tf.summary.FileWriter("./model/", sess.graph)

writer.close()