"""Photometric redshift analysis. Includes a wrapper to LEPHARE and BPZ.

- LEPHARE: http://www.cfht.hawaii.edu/~arnouts/LEPHARE/lephare.html
- BPZ: http://www.stsci.edu/~dcoe/BPZ
"""

from __future__ import print_function
import os
import sys
import subprocess
import numpy as N
import pylab as P
from scipy.optimize import curve_fit
from astropy.io import ascii
from astropy.table import Table, hstack
from astropy.coordinates import SkyCoord
from . import utils as cutils

class LEPHARE(object):

    """Wrapper to the LEPHARE photometric redshift code.

    http://www.cfht.hawaii.edu/~arnouts/LEPHARE/lephare.html
    """

    def __init__(self, magnitudes, errors, zpara=None, spectro_file=None, **kwargs):
        """
        Run the LEPHARE progam (zphota).

        :param list magnitudes: Magnitudes. A list of list
        :param list errors: Error on magnitudes. Same shape as the magnitude list.
        :param string zpara: default is $LEPHAREDIR/config/zphot_megacam.para
        :param dictionnary files: A dictioannry conating the names of the

        Kwargs include the following list of possible parameters

        :param string basename: Base name for all created files
        :param string cname: Name of the studied cluster. If basename if not given,
         it will be used as base name for all created file
        :param string filters: filter list
        :param list ra: list of ra
        :param list dec: list of dec
        :param list id: list of ID
        """
        self.data = {'mag': magnitudes, 'err': errors}
        self.kwargs = kwargs
        self.config = os.environ["LEPHAREDIR"] + \
            "/config/zphot_megacam.para" if zpara is None else zpara

        self.spectro_file = spectro_file
        # Name of created files?
        prefix = ""
        if 'basename' in kwargs:
            prefix = kwargs['basename'] + "_"
        elif 'cname' in kwargs:
            prefix = kwargs['cname'] + "_"
        self.files = {}
        self.files['input'] = prefix + "zphot.in"
        self.files['output'] = prefix + "zphot.out"
        self.files['pdz_output'] = prefix + "zphot"
        self.files['all_input'] = self.files['input'].replace('.in', '.all')

        # Initialize lephare output variables
        self.lephare_out = None
        self.data_out = None

        self.write_input()

    def write_input(self):
        """
        Create and write files needed to run LEPHARE.

        - the input data file for LEPHARE
        - a similar file containing the sources ID along with their RA DEC.
        """
        f = open(self.files['input'], 'w')
        if self.spectro_file is None:
            # No spectroscopic redshift file provided in config.yaml
            # --> only need the SHORT format for LePhare input file
            if 'filters' in self.kwargs:
                f.write("# id " + " ".join(["mag_%s" % filt for filt in self.kwargs['filters']]) +
                        " " + " ".join(["err_mag_%s" % filt for filt in self.kwargs['filters']]) +
                        "\n")
                for i, mags in enumerate(N.concatenate([self.data['mag'], self.data['err']]).T):
                    f.write("%i %s\n" %
                            (i, " ".join(["%.3f" % m for m in mags])))

                f.close()
                print("INFO: Input data saved in", self.files['input'])
        else:
            # Spectroscopic redshift file provided in config.yaml
            # --> Need to write LePhare input file in the LONG format,
            # i.e with 'context' and 'spectroz' of matching galaxies
            zspec = ZSPEC(self.spectro_file, names=[
                          'object', 'ra', 'dec', 'zspec'])
            zspec.skycoords = SkyCoord(
                zspec.data['ra'], zspec.data['dec'], unit='deg')
            skycoords_cat = SkyCoord(
                self.kwargs['ra'], self.kwargs['dec'], unit='deg')
            idx, d2d, d3d = skycoords_cat.match_to_catalog_sky(zspec.skycoords)
            zp = zspec.data['zspec'][idx]
            # identify galaxies with bad match, i.e. dist > 300 mas
            bad = N.where(d2d.mas > 300)
            zp[bad] = -99
            print("INFO: Using " + str(len(idx) - N.size(bad)) +
                  " galaxies for spectroz training")
            if 'filters' in self.kwargs:
                f.write("# id " + " ".join(["mag_%s" % filt for filt in self.kwargs['filters']]) +
                        " " + " ".join(["err_mag_%s" % filt for filt in self.kwargs['filters']]) +
                        " context" + " zspec" + "\n")
                # tells LePhare to run using the u, g, r, i and z bands.
                context = 31
                for i, mags in enumerate(N.concatenate([self.data['mag'], self.data['err']]).T):
                    f.write("%i %s %s\n" % (i, " ".join(["%.3f" % m for m in mags]),
                                            " ".join(("%i" % context, "%.3f" % zp[i]))))
                f.close()
        if 'ra' in self.kwargs:
            f = open(self.files['all_input'], 'w')
            if 'filters' in self.kwargs is not None:
                f.write("# id ID RA DEC " +
                        " ".join(["mag_%s" % filt for filt in self.kwargs['filters']]) + " " +
                        " ".join(["err_mag_%s" % filt for filt in self.kwargs['filters']]) +
                        "\n")
            for i, mags in enumerate(N.concatenate([self.data['mag'], self.data['err']]).T):
                f.write("%i %i %f %f %s\n" % (i, self.kwargs['id'][i],
                                              self.kwargs['ra'][i], self.kwargs['dec'][i],
                                              " ".join(["%.3f" % m for m in mags])))
            f.close()
            print("INFO: All data saved in", self.files['all_input'])

    def check_config(self, config=None):
        """
        Check that the SED and filters requested for the LePhare run do exist.

        If not: explains where the problem is and aborts.
        """
        if config is not None:
            if os.path.exists(config):
                self.config = config
            else:
                raise ValueError("%s does not exist" % config)
        with open(self.config) as f:
            for line in f:
                if "ZPHOTLIB" in line:
                    libs = line.split()[1].split(",")
        path_to_lib = os.environ["LEPHAREWORK"] + "/lib_mag/"

        for lib in libs:
            lib_tmp = path_to_lib + lib + ".bin"
            if not os.path.exists(lib_tmp):
                print("\nERROR: Requested library %s does not exist " % lib_tmp)
                print("INFO: Available SED libraries are:\n ")
                for f in os.listdir(path_to_lib):
                    if f.endswith(".bin"):
                        print(f)
                        print("--> Correct the ZPHOTLIB variable in %s" % self.config +
                              "or generate the missing LePhare SED libraries")
                sys.exit()

        for counter, lib in enumerate(libs):
            lib_tmp = path_to_lib + lib + ".doc"
            with open(lib_tmp) as f:
                for line in f:
                    if "FILTER_FILE" in line:
                        filt_tmp = line.split()[1]
                        if counter == 0:
                            filt_ref = filt_tmp
                        if filt_tmp != filt_ref:
                            print("\nERROR: Requested SED libraries %s" % libs +
                                  " have been built using different filters. Either change the " +
                                  "requested libraries or re-generate them accordingly.")
                            sys.exit()

        path_to_filt = os.environ["LEPHAREWORK"] + "/filt/"
        if not os.path.exists(path_to_filt + filt_ref):
            print("\nERROR: The FILTER_FILE %s used by the SED libraries %s does not exists." %
                  (path_to_filt + filt_ref, libs))
            sys.exit()

    def run(self, config=None):
        """
        Run LEPHARE.

        Default config file is $LEPHAREDIR/config/zphot_megacam.para.
        Can be overwritten with the config argument
        """
        if config is not None:
            if os.path.exists(config):
                self.config = config
            else:
                raise ValueError("%s does not exist" % config)

        # build command line
        cmd = "zphota"
        cmd += " -c " + self.config
        cmd += " -CAT_IN " + self.files['input']
        cmd += " -CAT_OUT " + self.files['output']
        cmd += " -PDZ_OUT " + self.files['pdz_output']
        print("INFO: Will run '%s'" % cmd)

        self.lephare_out = subprocess.check_output(
            cmd, stderr=subprocess.STDOUT, shell=True)
        print("INFO: LEPHARE output summary (full output in self.lephare_out)")
        print(
            "\n".join(["   " + zo for zo in self.lephare_out.split("\n")[-6:]]))

        self.data_out = ZPHOTO(self.files['output'], self.files['pdz_output'], zcode_name='lephare',
                               all_input=self.files['all_input'], **self.kwargs)

class BPZ(object):

    """
    Wrapper to the BPZ photometric redshift code.

    http://www.stsci.edu/~dcoe/BPZ
    """

    def __init__(self, magnitudes, errors, zpara=None, spectro_file=None, **kwargs):
        """
        Run the BPZ progam (zphota).

        :param list magnitudes: Magnitudes. A list of list
        :param list errors: Error on magnitudes. Same shape as the magnitude list.

        Kwargs include the following list of possible parameters

        :param string basename: Base name for all created files
        :param string cname: Name of the studied cluster. If basename if not given,
                             it will be used as base name for all created file
        :param string filters: filter list
        :param list ra: list of ra
        :param list dec: list of dec
        :param list id: list of ID
        """
        self.data = {'mag': magnitudes, 'err': errors}
        self.kwargs = kwargs
        self.config = zpara
        self.spectro_file = spectro_file

        # Name of created files?
        prefix = ""
        if 'basename' in kwargs:
            prefix = kwargs['basename'] + "_"
        elif 'cname' in kwargs:
            prefix = kwargs['cname'] + "_"
        self.files = {}
        self.files['input'] = prefix + "bpz.in"
        self.files['flux_comparison'] = prefix + "bpz.flux_comparison"
        self.files['output'] = prefix + "bpz.bpz"
        self.files['pdz_output'] = prefix + "bpz.probs"
        self.files['columns'] = prefix + "bpz.columns"
        self.files['all_input'] = self.files['input'].replace('.in', '.all')

        # Initialize BPZ output variables
        self.bpz_out = None
        self.data_out = None

        self.write_input()
        self.build_columns_file()

    def write_input(self):
        """
        Create and write files needed to run BPZ.

        - the input data file for BPZ
        - a similar file containing the sources ID along with their RA DEC.
        """
        f = open(self.files['input'], 'w')
        if 'filters' in self.kwargs:
            f.write("# id " + " ".join(["mag_%s" % filt for filt in self.kwargs['filters']]) +
                    " " + " ".join(["err_mag_%s" % filt for filt in self.kwargs['filters']]) +
                    "\n")
            for i, mags in enumerate(N.concatenate([self.data['mag'], self.data['err']]).T):
                f.write("%i %s\n" % (i, " ".join(["%.3f" % m for m in mags])))

            f.close()
            print("INFO: Input data saved in", self.files['input'])
        if 'ra' in self.kwargs:
            f = open(self.files['all_input'], 'w')
            if 'filters' in self.kwargs is not None:
                f.write("# id ID RA DEC " +
                        " ".join(["mag_%s" % filt for filt in self.kwargs['filters']]) + " " +
                        " ".join(["err_mag_%s" % filt for filt in self.kwargs['filters']]) +
                        "\n")
            for i, mags in enumerate(N.concatenate([self.data['mag'], self.data['err']]).T):
                f.write("%i %i %f %f %s\n" % (i, self.kwargs['id'][i],
                                              self.kwargs['ra'][i], self.kwargs['dec'][i],
                                              " ".join(["%.3f" % m for m in mags])))
            f.close()
            print("INFO: All data saved in", self.files['all_input'])

#    def build_columns_file(self, prefix='HSC-', sufix='',
#                           filters=None, ref='i', z_s=False):
    def build_columns_file(self, prefix='CFHT_megacam_', sufix='p',
                           filters=None, ref='i', z_s=False):
        """Build and write the 'columns' file.

        Hardcoded for test purpose.
        """
        if filters is None:  # set the default values as the CFHT ones
            filters = ['u', 'g', 'r', 'i', 'z']
#            filters = ['g', 'r', 'i', 'z', 'Y']
        f = open(self.files['columns'], 'w')
        f.write("# Filter  columns  AB/Vega  zp_error  zp_offset\n")
        for i, filt in enumerate(filters):
            f.write("%s%s%s     %i, %i   AB        0.01      0.00\n" %
                    (prefix, filt, sufix, i + 2, i + len(filters) + 2))
        if z_s:
            f.write("Z_S                  %i\n" % (len(filters) * 2 + 2))
        f.write("M_0                 %i\n"
                    % [j + 2 for j, filt in enumerate(filters) if filt == ref][0])
        f.write("ID                    1\n")
        f.close()

    def run(self):
        """
        Run BPZ.

        Configuration file must exist in the current directory.

        .. todo:: Build the configuration file on the fly (the .columns)
        """
        if not os.getenv('BPZPATH'):
            mess = "BPZPATH if not defined. Please install BPZ "
            mess += "(see http://www.stsci.edu/~dcoe/BPZ/install.html)."
            raise IOError(mess)

        if not os.path.exists(self.files['columns']):
            raise IOError("%s does not exist" % self.files['columns'])

        # build command line options from param file
        if self.config is not None:
            opt_arr = N.genfromtxt(self.config, dtype=None)
            option = ['-' + opt_arr[i, 0].decode() + ' ' + opt_arr[i, 1].decode()
                      for i in N.arange(len(opt_arr))]
            options = ' '.join(e for e in option)
        else:
            options = ''

        # build command line
#        cmd = "python $BPZPATH/bpz.py %s " % self.files['input'] + options
        cmd = "bpz_run.py %s " % self.files['input'] + options
        print("INFO: Will run '%s'" % cmd)

        self.bpz_out = subprocess.check_output(
            cmd, stderr=subprocess.STDOUT, shell=True)
        print("INFO: BPZ output summary (full output in self.bpz_out)")
#        print("\n".join(["   " + zo for zo in self.bpz_out.split("\n")[:20]]))

        self.data_out = ZPHOTO(self.files['output'], self.files['pdz_output'],
                               zcode_name='bpz', all_input=self.files['all_input'],
                               **self.kwargs)

class ZPHOTO(object):

    """Read photoz code (LePhare, BPZ) output file and creates/saves astropy tables."""

    def __init__(self, zphot_output, zphot_pdz_output, zcode_name=None, all_input=None, **kwargs):
        """Read the photoz progam (LePhare, BPZ, ...) output."""
        self.files = {}
        self.files['output'] = zphot_output
        self.files['pdz_output'] = zphot_pdz_output
        self.kwargs = kwargs
        if all_input is not None:
            self.files['input'] = all_input
            self.read_input()
        self.code = zcode_name
        self.read()

    def read(self):
        """Read the output."""
        f = open(self.files['output'], 'r')
        self.data_array = N.loadtxt(self.files['output'], unpack=True)

        if self.code == 'lephare':
            self.header = [l for l in f if l.startswith('#')]
            f.close()
            self.variables = N.loadtxt(os.getenv('LEPHAREDIR') +
                                       "/config/zphot_output.para", dtype='string')
            self.data_dict = {v: a for v, a in zip(
                self.variables, self.data_array)}
            self.nsources = len(self.data_dict['Z_BEST'])
            self.pdz_zbins = N.loadtxt(
                self.files['pdz_output'] + '.zph', unpack=True)
            self.pdz_val = N.loadtxt(
                self.files['pdz_output'] + '.pdz', unpack=True)

        elif self.code == 'bpz':
            self.header = [l for l in f if l.startswith('##')]
            f.seek(0)
            self.variables = [l[4:].replace(' ', '').split(
                '\n')[0] for l in f if l.startswith('# ')]
            f.seek(0)
            # BPZ does not provide a zbins file.
            # Needs to create it from zmin, zmax and dz specified in output file
            zmin = float([line.split('=')[1]
                          for line in f if 'ZMIN' in line][0])
            f.seek(0)
            zmax = float([line.split('=')[1]
                          for line in f if 'ZMAX' in line][0])
            f.seek(0)
            dz = float([line.split('=')[1] for line in f if 'DZ' in line][0])
            f.close()

            self.data_dict = {v: a for v, a in zip(
                self.variables, self.data_array)}
            self.nsources = len(self.data_dict['Z_B'])
            self.pdz_zbins = N.arange(zmin, zmax + dz, dz)
            self.pdz_val = N.loadtxt(self.files['pdz_output'], unpack=True,
                                     usecols=N.arange(1, len(self.pdz_zbins) + 1))

    def save_zphot(self, file_out, path_output, overwrite=False):
        """Save the output of photoz code (z_best, chi^2, pdz) into astropy table."""
        # Duplicates the zbins vector for each object.
        # It is redundant information but astropy tables need each field to have the
        # same size. Or maybe I'm missing something.
        zbins = N.tile(self.pdz_zbins, (len(self.kwargs['id']), 1))

        # Converts LePhare or BPZ likelihood to actual probability density
        for i in N.arange(len(self.pdz_val.T)):
            norm = N.trapz(self.pdz_val[:, i], self.pdz_zbins)
            new_pdz_val = self.pdz_val[:, i] / norm
            self.pdz_val[:, i] = new_pdz_val

        # Creates astropy table to be saved in path_output of file_out
        new_tab = hstack([Table([self.kwargs['id']]), Table([self.kwargs['ra']]),
                          Table([self.kwargs['dec']]), Table(self.data_dict),
                          Table([zbins], names=['zbins']),
                          Table([self.pdz_val.T], names=['pdz'])],
                         join_type='inner')

        # Rename BPZ Z_B to Z_BEST to match LePhare
        if 'Z_B' in new_tab.keys():
            new_tab.rename_column('Z_B', 'Z_BEST')

        # overwrite keyword of data.write(file,path) does not only overwrites
        # the data in path, but the whole file, i.e. (we lose all other
        # paths in the process) --> overwrite_or_append (see above)

        cutils.overwrite_or_append(file_out, path_output, new_tab, overwrite=overwrite)

        print("INFO: ", self.code, "data saved in", file_out, "as", path_output)

    def read_input(self):
        """Read the input."""
        data = N.loadtxt(self.files['input'], unpack=True)
        f = open(self.files['input'], 'r')
        l = f.readlines()[0]
        self.input_data = {k: d for k, d in zip(l[2:-1].split(), data)}
        f.close()

    def hist(self, param, **kwargs):
        """Plot histograms.

        Possible kwargs

        :params float minv: Lower value of the histogram
        :params float maxv: Upper value of the histogram
        :params int nbins: Number of bins. Default is 10.
        :params string xlabel: An xlbal for the figure
        :params string title: A title for the figure
        :params float zclust: Redshift of the studies cluster
        """
        # Get the data and apply some cut if asked
        pval = self.data_dict[param]
        filt = N.array([1] * len(pval), dtype='bool')
        if 'minv' in kwargs:
            filt &= (pval >= kwargs['minv'])
        if 'maxv' in kwargs:
            filt &= (pval <= kwargs['maxv'])
        pval = pval[filt]

        # Plot the histogram
        fig = P.figure()
        ax = fig.add_subplot(111, ylabel='#')
        ax.hist(pval, bins=kwargs['nbins'] if 'nbins' in kwargs else 10)
        xlabel = kwargs['xlabel'] if 'xlabel' in kwargs else param
        ax.set_xlabel(xlabel)
        if 'title' in kwargs:
            ax.set_title(kwargs['title'])
        if 'zclust' in kwargs:
            ax.axvline(kwargs['zclust'], color='r',
                       label='Cluster redshift (%.4f)' % kwargs['zclust'])
            ax.legend(loc='best')

        # Save the figure
        fig.savefig(self.files['output'].replace(
            '.out', '') + "_" + xlabel + "_zphot_hist.png")

    def plot(self, px, py, **kwargs):
        """
        Plot x vs. y.

        Possible kwargs are:
        :params float minx: lower limit of the x axis
        :params float maxx: upper limit of the x axis
        :params float miny: lower limit of the y axis
        :params float maxy: upper limit of the y axis
        :params string xlabel: label of the x axis
        :params string ylabel: label of the y axis
        :params string title: title of the figure
        """
        pvalx = self.data_dict[px]
        pvaly = self.data_dict[py]
        filt = N.array([1] * len(pvalx), dtype='bool')
        if 'minx' in kwargs and kwargs['minx'] is not None:
            filt &= (pvalx >= kwargs['minx'])
        if 'maxx' in kwargs and kwargs['maxx'] is not None:
            filt &= (pvalx <= kwargs['maxx'])
        if 'miny' in kwargs and kwargs['miny'] is not None:
            filt &= (pvaly >= kwargs['miny'])
        if 'maxy' in kwargs and kwargs['maxy'] is not None:
            filt &= (pvaly <= kwargs['maxy'])
        pvalx, pvaly = pvalx[filt], pvaly[filt]
        fig = P.figure()
        ax = fig.add_subplot(111)
        ax.scatter(pvalx, pvaly)
        if 'xlabel' in kwargs and kwargs['xlabel'] is not None:
            px = kwargs['xlabel']
        if 'ylabel' in kwargs and kwargs['ylabel'] is not None:
            py = kwargs['ylabel']
        ax.set_xlabel(px)
        ax.set_ylabel(py)
        if 'title' in kwargs and kwargs['title'] is not None:
            ax.set_title(kwargs['title'])

        if 'figname' in kwargs and kwargs['figname'] is not None:
            fig.savefig(self.files['output'].replace(
                '.out', '') + "_%s_vs_%s_zphot.png" % (py, px))
        else:
            fig.savefig("%s_vs_%s_zphot.png" % (py, px))

    def plot_map(self, title=None, zmin=0, zmax=999):
        """Plot the redshift sky-map."""
        if not hasattr(self, 'input_data'):
            print("WARNING: No input data given. Cannot plot the redshift map.")
            return

        ra, dec, redshift = self.input_data['RA'], self.input_data['DEC'], self.data_dict['Z_BEST']

        # redshift has to be >= 0
        filt = (redshift >= zmin) & (redshift < zmax)
        ra, dec, redshift = ra[filt], dec[filt], redshift[filt]

        fig = P.figure()
        ax = fig.add_subplot(111, xlabel='RA (deg)', ylabel='DEC (deg)')
        scat = ax.scatter(ra, dec, c=redshift, cmap=(P.cm.jet))
        cb = fig.colorbar(scat)
        cb.set_label('Photometric redshift')
        if title is not None:
            ax.set_title(title)
        ax.set_xlim(xmin=min(ra) - 0.001, xmax=max(ra) + 0.001)
        ax.set_ylim(ymin=min(dec) - 0.001, ymax=max(dec) + 0.001)
        fig.savefig(self.files['output'].replace(
            '.out', '') + "_redshift_map.png")

def dict_to_array(d, filters='ugriz'):
    """Transform a dictionnary into a list of arrays."""
    return N.array([N.array(d[f]) for f in filters])

class ZSPEC(object):

    """Compare spectroscopic and photometric redshifts."""

    def __init__(self, sfile, names, unit='deg'):
        """Read the input data file and organize the data.

        :param str sfile: File containing the spectroscopic redshift
        :param list names: A list of names for the columns. You must use 'ra' and 'dec'
         for the coordinates.
        :param str unit: Unit of the given coordinates ('deg', 'rad'). Should be understandable by
         astropy.coordinates.SkyCoord

        The spectroscopic redshift must be named 'zspec'.
        """
        self.sfile = sfile
        self.data = ascii.read(sfile, names=names)
        if 'ra' not in self.data.keys():
            raise IOError("RA coordinate must be called 'ra'")
        elif 'dec' not in self.data.keys():
            raise IOError("DEC coordinate must be called 'dec'")
        elif 'zspec' not in self.data.keys():
            raise IOError("Spectroscopic reshift must be 'zspec'")

        # Check whether duplicate galaxies exist in the spectroz sample
        # and remove them. Ideally would average out the various occurances.
        # (To be done later)
        ra = ['{:.20}'.format(x) for x in self.data['ra']]
        dec = ['{:.20}'.format(x) for x in self.data['dec']]
        radec = N.core.defchararray.add(ra, dec)
        unique_radec, good = N.unique(radec, return_index=True)
        if len(unique_radec) < len(radec):
            print("INFO: There are " + str(len(radec) - len(unique_radec)) +
                  " duplicate galaxies in spectroz sample. They are removed.")
        bad = N.delete(N.arange(len(self.data)), good)
        self.data.remove_rows(bad)

        self.skycoords = SkyCoord(self.data['ra'], self.data['dec'], unit=unit)
        self.zphot = self.skycoords_phot = self.match = None

    def load_zphot(self, ra, dec, zphot, unit='deg'):
        """Load the photometric informations and match them to the spectro ones.

        :param list ra: List of RA coordinates
        :param list dec: List of DEC coordinates
        :param list zphot: List of photometric redshift
        :param list unit: List of RA coordinates

        All lists must have the same length.
        """
        assert len(ra) == len(dec) == len(zphot)
        self.zphot = Table([ra, dec, zphot], names=['ra', 'dec', 'zphot'])
        self.skycoords_phot = SkyCoord(ra, dec, unit=unit)
        idx, d2d, d3d = self.skycoords.match_to_catalog_sky(
            self.skycoords_phot)
        self.match = Table([idx, d2d.marcsec, d3d], names=['idx', 'd2d', 'd3d'],
                           meta={'description': ['Indices into first catalog.',
                                                 'On-sky separation between the '
                                                 'closest match for each element',
                                                 '3D distance between the closest'
                                                 'match for each element'],
                                 'unit': ['', 'marcsec', '']})

    def plot(self, cut=300, path_to_png=None):
        """Plot a sky-map of the matches."""
        if self.match is None:
            raise IOError(
                "ERROR: You must load the photometric data first (load_zphot).")

        zspec = self.data['zspec']
        zphot = self.zphot['zphot'][self.match['idx']]
        sdist = N.array(self.match['d2d'])
        filt = sdist < cut
        zspec, zphot, sdist = [x[filt] for x in [zspec, zphot, sdist]]

        fig = P.figure(figsize=(15, 8))

        # z-phot as a function of z-spec
        ax = fig.add_subplot(121, xlabel='Z-spec', ylabel='Z-phot')
        ax.set_xlim([0, 1.5])
        ax.set_ylim([0, 3.5])
        scat = ax.scatter(zspec, zphot, c=sdist, cmap=(P.cm.copper))
        ax.plot(N.arange(int(max(zspec + 1))),
                N.arange(int(max(zspec + 1))), c='red')
        cb = fig.colorbar(scat)
        cb.set_label('On-sky distance (marcsec)')
        ax.set_title("%i galaxies" % len(self.match[filt]))

        # z-phot - zspec as a function of sources on-sky distances
        ax = fig.add_subplot(122, xlabel='On-sky distance',
                             ylabel='(Z-phot - Z-spec)')
        scat = ax.scatter(sdist, zphot - zspec, color='k')
        ax.set_title("%i galaxies" % len(self.match[filt]))
        ax.set_xscale('log')
        ax.set_xlim([1., 1.e6])
        ax.set_ylim([-1., 3.5])

        if path_to_png is not None:
            fig.savefig(path_to_png)

        fig = P.figure()

        # radec scatter plot of all catalogue and zspec galaxies
        ax = fig.add_subplot(121, xlabel='ra', ylabel='dec')
        ax.scatter(self.skycoords_phot.ra, self.skycoords_phot.dec,
                   color='k', label='Photo-z', s=15)
        ax.scatter(self.skycoords.ra, self.skycoords.dec,
                   color='r', label='Spectro-z', s=12)

        # radec scatter plot of matched catalogue and zspec galaxies, within the cut criterion
        ax = fig.add_subplot(122, xlabel='ra', ylabel='dec')
        ax.scatter(self.skycoords_phot.ra[self.match['idx'][filt]],
                   self.skycoords_phot.dec[self.match['idx'][filt]],
                   color='k', label='Photo-z', s=5)
        ax.scatter(self.skycoords.ra[filt], self.skycoords.dec[filt],
                   color='r', label='Spectro-z', s=5)

        if path_to_png is not None:
            fig.savefig(path_to_png.replace('.png', '_map.png'))

        P.show()

    def scatter(self, zclust, cluster=None, cut=0.1, stability=False):
        """Redshift scatter in the cluster.

        Plot the spectroscopic redshift distribution and apply a gaussian fit.
        """
        # Apply selection
        zspec = self.data['zspec'][N.abs(self.data['zspec'] - zclust) < cut]

        # The gaussian fit
        hist, bin_edges = N.histogram(zspec, bins=len(zspec) / 4)
        bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2

        # Use initial guess for the fitting coefficients (A, mu and sigma above)
        coeff = curve_fit(gauss, bin_centers, hist, p0=[
                          N.max(hist[0]), zclust, 0.3])[0]

        # Plot
        fig = P.figure()
        ax = fig.add_subplot(111, xlabel='Z-spec', ylabel='#',
                             title='' if cluster is None else cluster +
                             ', %i object included (cut=%.2f)' % (len(zspec), cut))
        ax.hist(zspec, bins=len(zspec) / 4, histtype='stepfilled')
        ax.axvline(zclust, color='r', label='Z-cluster (%.3f)' % zclust, lw=2)
        ax.plot(bin_centers, gauss(bin_centers, *coeff),
                label='Gaussian fit (mean, std)=(%.3f, %.3f)' %
                (coeff[1], N.abs(coeff[2])), lw=2, color='k')
        ax.legend(loc='best')

        if stability:
            cuts = [0.5, 0.1, 0.05, 0.03]
            z, ze = zip(*[self.scatter(zclust, cluster=cluster, cut=c)
                          for c in cuts])
            fig = P.figure()
            ax = fig.add_subplot(111, xlabel='Selection cut around cluste rredshift',
                                 ylabel='Estimated redshift',
                                 title='' if cluster is None else cluster +
                                 ', %i object included (cut=%.2f)' % (len(zspec), cut))
            ax.axhline(N.mean(z), label='Average value', color='r', lw=2)
            ax.errorbar(cuts, z, yerr=ze, label='Individual fits', color='k',
                        capsize=20, elinewidth=3)
            ax.legend(loc='best')
            print("INFO: Stability over the followed range of selection cuts:")
            print("       ", cuts)
            print("INFO: Input redshift: %.4f" % zclust)
            print("INFO: Average redshift: %.4f =/- %.4f" %
                  (N.mean(z), N.sqrt(N.std(z)**2 + N.mean(ze)**2)))
            P.show()
            return N.mean(z), N.sqrt(N.std(z)**2 + N.mean(ze)**2)
        P.show()
        return coeff[1], N.abs(coeff[2])

def gauss(x, *p):
    """Model function to be used to fit a gaussian distribution."""
    A, mu, sigma = p
    return A * N.exp(- (x - mu) ** 2 / (2. * sigma ** 2))


#!/usr/bin/env python

""" A set of tools for modelling the antenna reponse and for callibrating the amplitude vs frequency of the antennas
based on pyCRtools. see Schellart et al. Detecting cosmic rays with the LOFAR radio telescope,  and Nelles et al. Calibrating the absolute amplitude scale for air showers measured at LOFAR

Note: LBA_ant_calibrator still needs some work.

author: Brian hare
"""

##internal
import glob
from pickle import load
import datetime
from os.path import  dirname, abspath

##external
import numpy as np
from matplotlib import pyplot as plt
from scipy.interpolate import interp1d
from scipy.interpolate import RegularGridInterpolator, pchip_interpolate, PchipInterpolator
from scipy.io import loadmat

##mine
from LoLIM.utilities import processed_data_dir, MetaData_directory, RTD, SId_to_Sname, v_air
import LoLIM.transmogrify as tmf
import LoLIM.IO.metadata as md

# def get_absolute_calibration(frequencies, freq_mode='30-90'):
#     """this function returns the absolute calibration given a set of frequencies (in Hz). Note that this is the INVERSE of the factor in pyCRtools. You should multiply this by the jones matrix to set the absolute correction
#     Note this is spefically derived for 80m cable lengths, which is consistant with relative calibraiton below.

#     If frequency is outside 30-80 MHz. A is assumed to be constant outside 30-80. Cable lengths and RCU gain are correct between 10-90 MHz. Assumed to be constant outside 10-90.

#     Returns 0 where frequencies<0

#     freq-mode: can be '30-90' or '10-90'. This sets which RCU setting is used. This should probably be the same as the data set being used, but probably doesn't matter.
#           The calibration was derived with 30-80, which is default."""

#     A = [

#         ]

class aartfaac_LBA_model:

    def __init__(self, model_loc=None, R=700, C=15e-12, mode="LBA_OUTER"):
        if model_loc is None:
            model_loc = dirname(abspath(__file__)) + '/AARTFAAC_LBA_MODEL/'

        self.model_loc = model_loc
        self.antenna_mode = mode

        self.fine_model_loc = self.model_loc + "LBA_core_fine"
        self.course_model_loc = self.model_loc + "LBA_core"

        ## get base data, assume is same between models
        data_fine = loadmat( self.fine_model_loc , variable_names=['Positions', 'Theta', 'Phi', 'Freq', 'Zant', 'g_z'] ) ## Zant is antenna impedences, g_z is angular dependence
        self.AART_ant_positions = data_fine['Positions']
        self.AART_thetas = data_fine['Theta'][0]  ## zenithal angle in degrees
        self.AART_Phis = data_fine['Phi'][0]      ## azimuthal angle in degrees

        self.AART_fineFreqs = data_fine['Freq'][0]
        fine_gz = data_fine['g_z']
        fine_Zant =  data_fine['Zant']

        data_course = loadmat( self.course_model_loc , variable_names=['Freq', 'Zant', 'g_z'] )
        self.AART_courseFreqs = data_course['Freq'][0]
        course_gz = data_course['g_z']
        course_Zant =  data_course['Zant']

        ### figure out how to fold fine and course models together
        fineFreq_low = self.AART_fineFreqs[0]
        # fineFreq_high = self.AART_fineFreqs[-1]

        freqs = []
        course_range_1 = [0]
        fine_range = []
        course_range_2 = []

        ## first we search AART_courseFreqs to find highest frequency lower than fine-range, adding them to the frequency list
        for i in range(len(self.AART_courseFreqs)):
            Fc = self.AART_courseFreqs[i]
            if Fc < fineFreq_low:
                freqs.append(Fc)
            else:
                break
        course_range_1.append(i)
        fine_range.append(i)

        ## add all fine frequencies
        freqs += list(self.AART_fineFreqs)
        fine_range.append(len(freqs))
        course_range_2.append(len(freqs))

        ## find first course frequency greater than fine frequencies
        for i in range(i, len(self.AART_courseFreqs)):
            if self.AART_courseFreqs[i] > freqs[-1]:
                break

        ## add rest of course frequencies to list
        freqs += list(self.AART_courseFreqs[i:])
        course_range_2.append(len(freqs))

        self.all_frequencies = np.array(freqs, dtype=np.double)
        self.course_frequencyRange_low = np.array(course_range_1)
        self.fine_frequencyRange = np.array(fine_range)
        self.course_frequencyRange_high = np.array(course_range_2)

        #### combine antenna impedences
        num_ants = len(fine_Zant)
        num_thetas = len(self.AART_thetas)
        num_phis = len(self.AART_Phis)
        num_freqs = len(self.all_frequencies)

        self.antenna_Z = np.empty([num_ants, num_ants, num_freqs ], dtype=np.complex)

        self.antenna_Z[:,:, :self.course_frequencyRange_low[1]] = course_Zant[:,:,  :self.course_frequencyRange_low[1]]
        self.antenna_Z[:,:, self.fine_frequencyRange[0]:self.fine_frequencyRange[1]] = fine_Zant[:,:,  :]
        self.antenna_Z[:,:, self.fine_frequencyRange[1]:] = course_Zant[:,:, self.fine_frequencyRange[1]-len(self.all_frequencies):]

        self.total_impedence = np.empty((num_ants, num_ants, num_freqs), dtype=np.complex)
        self.set_RC( R, C ) ## this sets self.total_impedence

        ### now we combine total G_z, which is voltage on antenna
        self.total_Gz = np.empty( (num_ants,2,num_thetas,num_phis,num_freqs), dtype=np.complex )

        # for ant_i in range(num_ants):
            # for pol_i in range(2):
        self.total_Gz[:,:, :,:,                             : self.course_frequencyRange_low[1]] = course_gz[:,:, :,:,    : self.course_frequencyRange_low[1]]
        self.total_Gz[:,:, :,:, self.fine_frequencyRange[0] : self.fine_frequencyRange[1]]       = fine_gz  [:,:, :,:,    : ]
        self.total_Gz[:,:, :,:, self.fine_frequencyRange[1] :]                                   = course_gz[:,:, :,:, self.fine_frequencyRange[1]-len(self.all_frequencies): ]

    def get_LNA_filter(self):
        """given the antenna mode, return two things. First is np.array of indecies of LNAs that are on, second is np.array of LNAs that are off"""

        station_ant_positions = []
        for station in ['CS002', 'CS003', 'CS004', 'CS005', 'CS006', 'CS007']:
            station_ant_positions.append(  md.convertITRFToLocal( md.getItrfAntennaPosition(station, self.antenna_mode) )[::2, :] )

        found_indeces = []
        unfound_indeces = []

        AARFAACC_X_locations = self.AART_ant_positions[0]
        AARFAACC_Y_locations = self.AART_ant_positions[1]

        err2 = 0.1*0.1
        for AAR_i in range(len(AARFAACC_X_locations)):

            found = False
            for stat_ant_XYZ in station_ant_positions:
                for XYZ in stat_ant_XYZ:
                    KX, KY, KZ = XYZ
                    AX = AARFAACC_X_locations[AAR_i]
                    AY = AARFAACC_Y_locations[AAR_i]
                    R2 = (AX-KX)**2 + (AY-KY)**2
                    if R2<err2:
                        found = True
                        break
                if found:
                    break

            if found:
                found_indeces.append(AAR_i)
            else:
                unfound_indeces.append(AAR_i)

        return np.array(found_indeces), np.array(unfound_indeces)

    def set_RC(self, R, C ):
        self.R = R
        self.C = C

        num_ants = len(self.antenna_Z)
        num_freqs = len(self.all_frequencies)

        # functioning_indeces, nonfunctioning_indeces = self.get_LNA_filter()
        # LNA_filter = np.zeros(num_ants, dtype=np.int)
        # LNA_filter[ functioning_indeces ] = 1
        # LNA_Off_filter = np.zeros(num_ants, dtype=np.int)
        # LNA_Off_filter[ nonfunctioning_indeces ] = 1

        tmp_mat = np.empty((num_ants, num_ants), dtype=np.complex)
        ZLMA_matrix =  np.zeros((num_ants, num_ants), dtype=np.complex)

        for Fi in range(num_freqs):
            # set LNA matrix
            LNA = R/(1+self.all_frequencies[Fi]*(2*np.pi*1j*R*C)) # LNA_filter*(R/(1+self.all_frequencies[Fi]*(2*np.pi*1j*R*C)))  +  LNA_Off_filter*( 1/(self.all_frequencies[Fi]*2*np.pi*1j*C) )
            np.fill_diagonal(ZLMA_matrix,  LNA)

            ## caclulate denominator bit
            tmp_mat[:,:] = self.antenna_Z[:,:,Fi]
            tmp_mat += ZLMA_matrix

            ## try sign difference
            # tmp_mat *= -1
            # np.fill_diagonal(tmp_mat,  np.diagonal(tmp_mat)*-1)

            ## invert and multiple
            inv = np.linalg.inv( tmp_mat )
            np.matmul(ZLMA_matrix, inv, out=tmp_mat)

            # finally
            self.total_impedence[:,:,Fi] = tmp_mat

    def loc_to_anti(self, antenna_XY):
        """given an antenna XY location, return internal index of X-dipole of that antenna"""

        KX, KY = antenna_XY

        AARFAACC_X_locations = self.AART_ant_positions[0, 0::2]
        AARFAACC_Y_locations = self.AART_ant_positions[1, 0::2]

        found_index = None
        err2 = 0.1*0.1
        for AAR_i in range(len(AARFAACC_X_locations)):
            AX = AARFAACC_X_locations[AAR_i]
            AY = AARFAACC_Y_locations[AAR_i]
            R2 = (AX-KX)**2 + (AY-KY)**2
            if R2<err2:
                found_index = AAR_i*2
                break
        if found_index is None:
            print("ERROR! AARTFAAC model cannot find your antenna loc:", antenna_XY)
            quit()

        X_ant_i = found_index
        return X_ant_i

    def RCUname_to_anti(self, antenna_name):
        """given RCUid (and known mode), return internal index of the X-dipole of that antenna"""

        station = SId_to_Sname[ int(antenna_name[:3]) ]
        known_positions = md.convertITRFToLocal( md.getItrfAntennaPosition(station, self.antenna_mode) )
        this_antenna_location = known_positions[ int(antenna_name[-3:]) ]

        return self.loc_to_anti( this_antenna_location[0:2] )

    class single_LBA_model:
        def __init__(self, jones_functions, freq_bounds):
            self.jones_functions = jones_functions
            self.freq_bounds = freq_bounds

        def Jones_Matrices(self, frequencies, zenith, azimuth, freq_fill=1.0):
            """ if frequencies is numpy array in Hz, zenith and azimuth in degrees, than return numpy array of jones matrices,
            that when doted with [zenithal,azimuthal] component of incidident E-field, then will give [X,Y] voltages on dipoles"""

            return_matrices = np.empty( (len(frequencies), 2,2), dtype=np.complex )

            if zenith<0:
                zenith = 0
            elif zenith > 90:
                zenith = 90

            while azimuth<0:
                azimuth += 360
            while azimuth>360:
                azimuth -= 360

            # good_frequency_filter = np.logical_and(frequencies>=self.freq_bounds[0], frequencies<self.freq_bounds[-1])

            # points = np.empty( (np.sum(good_frequency_filter), 3), dtype=np.double )

            J00,J01 = self.jones_functions[0]
            J10,J11 = self.jones_functions[1]

            for fi,f in enumerate(frequencies):
                if f<=self.freq_bounds[0] or f>=self.freq_bounds[-1]:
                    return_matrices[fi, 0,0] = freq_fill
                    return_matrices[fi, 1,1] = freq_fill
                    return_matrices[fi, 0,1] = 0.0
                    return_matrices[fi, 1,0] = 0.0
                else:
                    return_matrices[fi, 0,0] = J00( (zenith,azimuth,f) )
                    return_matrices[fi, 0,1] = J01( (zenith,azimuth,f) )
                    return_matrices[fi, 1,0] = J10( (zenith,azimuth,f) )
                    return_matrices[fi, 1,1] = J11( (zenith,azimuth,f) )

                    # print(fi)
                    # M = return_matrices[fi]
                    # print("J:", M)
                    # print( M[ 0,0]*M[ 1,1] - M[  0,1]*M[  1,0] )
                    # print()

            return return_matrices

    def get_antenna_model(self, ant_i):
        """note: ant_i should be internal index!"""

        X_ant_i = int(ant_i/2)*2

        KX, KY = self.AART_ant_positions[:, X_ant_i]

        num_thetas = len(self.AART_thetas)
        num_phis = len(self.AART_Phis)
        num_frequencies = len( self.all_frequencies)
        num_antennas = len( self.total_Gz )

    #### some setup
        ## phase shifter to remove geometric delay
        shifter = np.empty([num_thetas, num_phis, num_frequencies ], dtype=np.complex)
        for zi in range( num_thetas ):
            for ai in range( num_phis ):
                Zenith = self.AART_thetas[zi]/RTD
                Azimuth = self.AART_Phis[ai]/RTD
                dt = ( KX*np.sin(Zenith)*np.cos(Azimuth) + KY*np.sin(Zenith)*np.sin(Azimuth))/v_air
                np.exp( self.all_frequencies*(1j*2*np.pi*(-dt)), out=shifter[zi,ai] )

        ## frequency interpolation
        frequency_bin = 0.5e6
        minF = self.all_frequencies[0]
        maxF = self.all_frequencies[-1]
        num_interpolant_freqs = int((maxF+minF)/frequency_bin )
        interpolant_frequencies = np.linspace(minF, maxF, num_interpolant_freqs)

        ## memory
        tmp1 = np.empty( (num_antennas,num_frequencies), dtype=np.complex )
        tmp2 = np.empty(num_frequencies, dtype=np.complex )
        tmp3 = np.empty(num_frequencies, dtype=np.double )

        def make_interpolant(antenna, polarization):
            """ antenna is 0 or 1 for X or Y, pol is 0 or 1 for zenith or azimithal component"""
            nonlocal tmp1, tmp2, tmp3

            antenna = antenna + X_ant_i

            grid = np.empty([num_thetas, num_phis, len(interpolant_frequencies) ], dtype=np.complex)

            for theta_i in range(num_thetas):
                for phi_i in range(num_phis):

                    ## dot product between voltages and impedences
                    tmp1[:,:]  = self.total_Gz[:, polarization,theta_i,phi_i, :]
                    tmp1 *= self.total_impedence[antenna, :, :] # this shouldn't matter??
                    # tmp1 *= self.total_impedence[:, antenna, :]

                    np.sum( tmp1, axis=0, out = tmp2 )

                    ## shift phase due to arival direction
                    tmp2 *= shifter[theta_i,phi_i]

                    ## interpolate amplitude and phase
                    np.abs(tmp2, out=tmp3)
                    interp_ampltude = pchip_interpolate(self.all_frequencies,tmp3, interpolant_frequencies)

                    # if theta_i==0 and phi_i==0:
                        # print('amp')
                        # plt.plot( interpolant_frequencies, interp_ampltude, 'o' )
                        # plt.plot( self.all_frequencies, tmp3, 'o' )
                        # plt.show()

                    angles = np.angle(tmp2)
                    angles = np.unwrap(angles)

                    interp_angle = pchip_interpolate(self.all_frequencies,angles, interpolant_frequencies)

                    # if theta_i==0 and phi_i==0:
                    #     print('angle')
                    #     plt.plot( interpolant_frequencies, interp_angle, 'o' )
                    #     plt.plot( self.all_frequencies, angles, 'o' )
                    #     plt.show()

                    ## now convert back to real and imag
                    interp_angle = interp_angle*1j
                    np.exp( interp_angle, out=grid[theta_i,phi_i] )
                    grid[theta_i,phi_i] *= interp_ampltude

            ## correct for different in definition
            if polarization==0: ## zenith points in oppisite directions in two models??
                grid *= -1
            ## and final angle-frequency interpolant
            interpolant = RegularGridInterpolator((self.AART_thetas, self.AART_Phis, interpolant_frequencies),  grid,  bounds_error=False,fill_value=0.0)
            return interpolant

        J00_interpolant = make_interpolant(0, 0)
        J01_interpolant = make_interpolant(0, 1)
        J10_interpolant = make_interpolant(1, 0)
        J11_interpolant = make_interpolant(1, 1)

        return self.single_LBA_model([[J00_interpolant,J01_interpolant],[J10_interpolant,J11_interpolant]],  [interpolant_frequencies[0],interpolant_frequencies[-1]])

    def get_average_antenna_model(self):
        """return model averaged over antenna set"""

        if self.antenna_mode == "LBA_OUTER":
            start_i = 576
            end_i = 576*2
        elif self.antenna_mode == "LBA_INNER":
            start_i = 0
            end_i = 576
        else:
            print("unknown mode:", self.antenna_mode )

        total_num_antennas = len( self.total_Gz )
        num_frequencies = len(self.all_frequencies)
        num_zeniths = len(self.AART_thetas)
        num_azimuths = len(self.AART_Phis)

        ### define a utility function for extracting grid data
        temp_matrix = np.empty( (num_zeniths, num_azimuths, num_frequencies), dtype=np.complex )
        shifterTMP = np.empty( [ num_frequencies ], dtype=np.complex)
        tmp1 = np.empty( (total_num_antennas,num_frequencies), dtype=np.complex )
        def get_ant_i(ant_i, pol_i):
            """given internal antenna index, and polarization, fill temp_matrix with correct response"""
            nonlocal  shifterTMP, tmp1, temp_matrix

            KX, KY = self.AART_ant_positions[:, ant_i]
            for zi, ze in enumerate(self.AART_thetas):
                for ai, az in enumerate(self.AART_Phis):

                    ## dot product between voltages and impedences
                    tmp1[:,:]  = self.total_Gz[:, pol_i,zi,ai, :]
                    tmp1 *= self.total_impedence[ant_i, :, :] # this shouldn't matter??
                    # tmp1 *= self.total_impedence[:, antenna, :]

                    np.sum( tmp1, axis=0, out = temp_matrix[zi,ai, : ] )

                    ## caculate phase shifts
                    Zenith = ze/RTD
                    Azimuth = az/RTD
                    dt = ( KX*np.sin(Zenith)*np.cos(Azimuth) + KY*np.sin(Zenith)*np.sin(Azimuth))/v_air
                    shifterTMP[:] = self.all_frequencies
                    shifterTMP *= -1j*2*np.pi*dt
                    np.exp( shifterTMP, out=shifterTMP )

                    ## apply the shits
                    temp_matrix[zi,ai, : ] *= shifterTMP

            if pol_i==0: ## zenith points in oppisite directions in two models??
                temp_matrix *= -1

        ### now define a utility function for calculating the average, and interpolating

        ## frequency interpolation
        frequency_bin = 0.5e6
        minF = self.all_frequencies[0]
        maxF = self.all_frequencies[-1]
        num_interpolant_freqs = int((maxF+minF)/frequency_bin )
        interpolant_frequencies = np.linspace(minF, maxF, num_interpolant_freqs)

        temp_AVE_matrix = np.empty( (num_zeniths, num_azimuths, num_frequencies), dtype=np.complex )
        tmp3 = np.empty(num_frequencies, dtype=np.double )
        def calc_func(antenna_pol, field_pol):
            """antenna_pol should be 0 or 1 for X or Y antenna, and field_pol is 0 or 1 for zenithal or azimuthal field"""
            nonlocal temp_AVE_matrix, tmp3

            ## first average
            Npairs = int( (end_i-start_i)/2 )
            for pair_i in range(Npairs):
                ant_i = 2*pair_i + antenna_pol + start_i
                get_ant_i(ant_i, field_pol) ## this fills temp_matrix
                temp_AVE_matrix += temp_matrix ## error prone, but should be okay

                if not np.all( np.isfinite(temp_AVE_matrix) ):
                    print('YO problem!', pair_i, np.all( np.isfinite(temp_matrix) ) )
                    quit()

            temp_AVE_matrix /= Npairs

            upsample_grid = np.empty([num_zeniths, num_azimuths, len(interpolant_frequencies) ], dtype=np.complex)

            ## now interpolate frequencies
            for zi, ze in enumerate(self.AART_thetas):
                for ai, az in enumerate(self.AART_Phis):

                    np.abs(temp_AVE_matrix[zi,ai, : ], out=tmp3)
                    interp_ampltude = pchip_interpolate(self.all_frequencies, tmp3, interpolant_frequencies)

                    angles = np.angle(temp_AVE_matrix[zi,ai, : ])
                    angles = np.unwrap(angles)

                    interp_angle = pchip_interpolate(self.all_frequencies, angles, interpolant_frequencies)

                    ## now convert back to real and imag
                    interp_angle = interp_angle*1j
                    np.exp( interp_angle, out = upsample_grid[zi,ai, :] )
                    upsample_grid[zi,ai, :] *= interp_ampltude

            ## and final angle-frequency interpolant
            return RegularGridInterpolator((self.AART_thetas, self.AART_Phis, interpolant_frequencies),  upsample_grid,  bounds_error=False, fill_value=0.0)

        ### now we actually make the responses
        J00 = calc_func(0, 0)
        J01 = calc_func(0, 1)
        J10 = calc_func(1, 0)
        J11 = calc_func(1, 1)
        return self.single_LBA_model([[J00,J01],[J10,J11]],  [interpolant_frequencies[0],interpolant_frequencies[-1]])

    def get_antenna_locs(self):
        """return two arrays. First is X locations of all antennas. Second is Y locations. Each location is in pairs, the first is the X dipole, second is Y dipole"""
        return self.AART_ant_positions[0], self.AART_ant_positions[1]

    def get_grid_info(self):
        """return info about the internal grid. returns three real-valued numpy arrays: frequencies [Hz], zenithal [degrees], and azimuthal [degrees]"""
        return self.all_frequencies, self.AART_thetas, self.AART_Phis

    def get_response_grid(self, antenna_i, frequency_i, ant_pol_i, field_pol_i):
        """given an antenna_i, and frequency_i, return response of ant_pol_i (0 for X 1 for Y dipole) to field_pol_i (0 is zenithal, 1 is azimuthal). Response is a 2D matrix, first index is zenith angle, second is azimuthal"""

        total_antenna_i = int(antenna_i/2)*2 + ant_pol_i

        #### some setup
        KX, KY = self.AART_ant_positions[:, total_antenna_i]

        num_thetas = len(self.AART_thetas)
        num_phis = len(self.AART_Phis)

        frequency = self.all_frequencies[ frequency_i ]

        ret = np.empty([num_thetas, num_phis], dtype=np.complex)

        ## the calculation
        for zenith_i in range( num_thetas ):
            for azimuth_i in range( num_phis ):

                ## phase shifter to remove geometric delay
                Zenith = self.AART_thetas[zenith_i]/RTD
                Azimuth = self.AART_Phis[azimuth_i]/RTD
                dt = ( KX*np.sin(Zenith)*np.cos(Azimuth) + KY*np.sin(Zenith)*np.sin(Azimuth))/v_air
                phase_shift = np.exp( frequency*(1j*2*np.pi*(-dt)) )

                ## the antenna response
                response = np.dot( self.total_impedence[total_antenna_i, :, frequency_i] , self.total_Gz[:, field_pol_i, zenith_i, azimuth_i, frequency_i] )

                ret[ zenith_i,  azimuth_i] = response*phase_shift

        if field_pol_i==0: ## zenith points in oppisite directions in two models??
            ret *= -1

        return ret

class aartfaac_average_LBA_model:

    def __init__(self, mode="LBA_OUTER"):
        if mode == "LBA_OUTER":
            fname = MetaData_directory+"/lofar/antenna_response_model/AARTFAAC_AVE_LBAOUTER_R700_C17.npz"
        else:
            print('mode unknown:', mode)
            quit()

        data = np.load(fname)
        initial_frequencies = data['freqs']
        zeniths = data['zeniths']
        azimuths = data['azimuths']

        original_J00 = data['J00']
        original_J01 = data['J01']
        original_J10 = data['J10']
        original_J11 = data['J11']

        num_zeniths  = len(zeniths)
        num_azimuths = len(azimuths)

        ## frequency interpolation
        frequency_bin = 0.5e6
        minF = initial_frequencies[0]
        maxF = initial_frequencies[-1]
        num_interpolant_freqs = int((maxF+minF)/frequency_bin )
        interpolant_frequencies = np.linspace(minF, maxF, num_interpolant_freqs)

        # print('arg', initial_frequencies.shape, original_J00.shape)

        tmp = np.empty(len(initial_frequencies), dtype=np.double )
        def upsample_and_interpolate(GRID):
            nonlocal tmp

            upsample_grid = np.empty([num_zeniths, num_azimuths, num_interpolant_freqs ], dtype=np.complex)

            ## now interpolate frequencies
            for zi, ze in enumerate(zeniths):
                for ai, az in enumerate(azimuths):

                    np.abs(GRID[zi,ai, : ], out=tmp)
                    interp_ampltude = pchip_interpolate(initial_frequencies, tmp, interpolant_frequencies)

                    angles = np.angle(GRID[zi,ai, : ])
                    angles = np.unwrap(angles)

                    interp_angle = pchip_interpolate(initial_frequencies, angles, interpolant_frequencies)

                    ## now convert back to real and imag
                    interp_angle = interp_angle*1j
                    np.exp( interp_angle, out = upsample_grid[zi,ai, :] )
                    upsample_grid[zi,ai, :] *= interp_ampltude

            ## and final angle-frequency interpolant

            return RegularGridInterpolator((zeniths, azimuths, interpolant_frequencies),  upsample_grid,  bounds_error=False, fill_value=0.0)

        self.freq_bounds = [ interpolant_frequencies[0], interpolant_frequencies[-1] ]

        self.J00 = upsample_and_interpolate( original_J00 )
        self.J01 = upsample_and_interpolate( original_J01 )
        self.J10 = upsample_and_interpolate( original_J10 )
        self.J11 = upsample_and_interpolate( original_J11 )

    def Jones_Matrices(self, frequencies, zenith, azimuth, freq_fill=1.0):
        """ if frequencies is numpy array in Hz, zenith and azimuth in degrees, than return numpy array of jones matrices,
        that when doted with [zenithal,azimuthal] component of incidident E-field, then will give [X,Y] voltages on dipoles"""

        return_matrices = np.empty( (len(frequencies), 2,2), dtype=np.complex )

        if zenith<0:
            zenith = 0
        elif zenith > 90:
            zenith = 90

        while azimuth<0:
            azimuth += 360
        while azimuth>360:
            azimuth -= 360

        # good_frequency_filter = np.logical_and(frequencies>=self.freq_bounds[0], frequencies<self.freq_bounds[-1])

        # points = np.empty( (np.sum(good_frequency_filter), 3), dtype=np.double )

        for fi,f in enumerate(frequencies):
            if f<self.freq_bounds[0] or f>self.freq_bounds[-1]:
                return_matrices[fi, 0,0] = freq_fill
                return_matrices[fi, 1,1] = freq_fill
                return_matrices[fi, 0,1] = 0.0
                return_matrices[fi, 1,0] = 0.0
            else:
                return_matrices[fi, 0,0] = self.J00( (zenith,azimuth,f) )
                return_matrices[fi, 0,1] = self.J01( (zenith,azimuth,f) )
                return_matrices[fi, 1,0] = self.J10( (zenith,azimuth,f) )
                return_matrices[fi, 1,1] = self.J11( (zenith,azimuth,f) )

                # print(fi)
                # M = return_matrices[fi]
                # print("J:", M)
                # print( M[ 0,0]*M[ 1,1] - M[  0,1]*M[  1,0] )
                # print()

        return return_matrices

    def save_to_text(self):

        freq_grid = np.linspace(10,90, num=int((90-10)/1)+1 )
        zenith_grid = np.linspace(0,90, num=int(90/5)+1 )
        azimuth_grid = np.linspace(0,360, num=int(360/10)+1 )

        J00_out = open('./J00_text.txt', 'w')
        J01_out = open('./J01_text.txt', 'w')
        J10_out = open('./J10_text.txt', 'w')
        J11_out = open('./J11_text.txt', 'w')

        J00_out.write('voltage on X-dipole to zenithal field. Azimuthal=0 points 45 degrees south from East.\n')
        J00_out.write('f (MHz) Zenith(deg.) Azimuth(deg.) real(Vout) imag(Vout)\n')

        J01_out.write('voltage on X-dipole to azimuthal field. Azimuthal=0 points 45 degrees south from East.\n')
        J01_out.write('f (MHz) Zenith(deg.) Azimuth(deg.) real(Vout) imag(Vout)\n')

        J10_out.write('voltage on Y-dipole to zenithal field\. Azimuthal=0 points 45 degrees south from East.n')
        J10_out.write('f (MHz) Zenith(deg.) Azimuth(deg.) real(Vout) imag(Vout)\n')

        J11_out.write('voltage on Y-dipole to azimuthal field. Azimuthal=0 points 45 degrees south from East.\n')
        J11_out.write('f (MHz) Zenith(deg.) Azimuth(deg.) real(Vout) imag(Vout)\n')

        for f in freq_grid:
            for z in zenith_grid:
                for a in azimuth_grid:
                    jonesy = self.Jones_Matrices( [f*1e6],z,a+45 )
                    start = str(f)+ ' ' + str(z)+' '+str(a)+' '
                    J00_out.write(start+str(np.real(jonesy[0,0,0])) + ' ' + str(np.imag(jonesy[0,0,0])) +'\n')
                    J01_out.write(start+str(np.real(jonesy[0,0,1])) + ' ' + str(np.imag(jonesy[0,0,1])) +'\n')
                    J10_out.write(start+str(np.real(jonesy[0,1,0])) + ' ' + str(np.imag(jonesy[0,1,0])) +'\n')
                    J11_out.write(start+str(np.real(jonesy[0,1,1])) + ' ' + str(np.imag(jonesy[0,1,1])) +'\n')

class calibrated_AARTFAAC_model:
    """returns the AARTFAAC model multiplied by katies cal."""

    def __init__(self):
        self.AARTFAAC = aartfaac_average_LBA_model()

        calibration = [42484.88872879, 41519.47733373, 39694.16854372, 38435.36963118,
                       36974.13596039, 35819.34454985, 35072.53901876, 33960.74197721,
                       32944.65142405, 32112.44688046, 31590.52174516, 30433.52586868,
                       29756.92041985, 28880.95581826, 28204.37711364, 27646.83132496,
                       27058.11219529, 26693.43614747, 25952.89684011, 25517.72346825,
                       25220.5602153 , 24782.62398452, 24349.88724441, 23755.78027908,
                       23289.45439997, 22800.30039329, 22176.17569116, 21619.22598549,
                       21341.84386379, 21368.96227069, 22032.38713345, 22127.96304476,
                       22678.25184787, 24670.14420208, 25703.00236409, 24792.9736945,
                       23817.8752368 , 22912.1048524 , 22393.11267467, 21844.38535201,
                       20831.10515113, 20220.27315072, 19223.64518551, 18099.23669312,
                       17569.76228552, 16298.99230863, 14845.6897604 , 13488.79405579,
                       11966.53709451, 10895.66530218,  9703.34502309]

        cal_frequencies = np.arange(30e6, 80.5e6, 1e6)
        self.calibration_interpolator = PchipInterpolator(cal_frequencies, calibration, extrapolate=True  )

    def get_calibrator(self, frequencies):
        return self.calibration_interpolator( frequencies )

    def Jones_ONLY(self, frequencies, zenith, azimuth, freq_fill=1.0):
        return self.AARTFAAC.Jones_Matrices( frequencies, zenith, azimuth, freq_fill )

    def Jones_Matrices(self, frequencies, zenith, azimuth, freq_fill=1.0):
        """return calibrated jones matrices"""
        JM = self.AARTFAAC.Jones_Matrices( frequencies, zenith, azimuth, freq_fill )
        C = self.get_calibrator( frequencies )
        JM[:, 0,0] *= C
        JM[:, 0,1] *= C
        JM[:, 1,0] *= C
        JM[:, 1,1] *= C
        return JM

def invert_2X2_matrix_list( matrices ):
    """ if matrices is an array of 2x2 matrices, then return the array of inverse matrices """
    num = len(matrices)
    out = np.zeros( (num, 2,2), dtype=matrices.dtype)

    out[:, 0,0] = matrices[:, 1,1]
    out[:, 0,1] = -matrices[:, 0,1]
    out[:, 1,0] = -matrices[:, 1,0]
    out[:, 1,1] = matrices[:, 0,0]

    determinants = matrices[:, 0,0]*matrices[:, 1,1] - matrices[:, 0,1]*matrices[:, 1,0]

    out /= determinants[:, np.newaxis, np.newaxis]

    return out

def fourier_series( x, p):
    """Evaluates a partial Fourier series

        F(x) \\approx \\frac{a_{0}}{2} + \\sum_{n=1}^{\\mathrm{order}} a_{n} \\sin(nx) + b_{n} \\cos(nx)
    """

    r = p[0] / 2

    order = int( (len(p) - 1) / 2 )

    for i in range(order):

        n = i + 1

        r += p[2*i + 1] * np.sin(n * x) + p[2*i + 2] * np.cos(n * x)

    return r

def getGalaxyCalibrationData(antenna_noise_power, timestamp, antenna_type="outer"):
    """return factor to correct for amplitude shifts. Essenturally returns sqrt( P_{expected} / P_{measured} ). Where P is noise power.
    for antenna_type outer it returns factor for Y/X dipoles, for "inner" returns "X/Y". antenna_noise_power is an array of measured powers for each antenna.
    Even/odd indecies should be Y/X dipole for outer and oppisite for inner.
    timestamp should be posix timestamp"""

    longitude = 6.869837540/RTD

    ## this is in outer order:  Y,X
    # coefficients_lba = [ np.array([ 0.01489468, -0.00129305,  0.00089477, -0.00020722, -0.00046507]),   ## for Y antennas
    #                       np.array( [ 0.01347391, -0.00088765,  0.00059822,  0.00011678, -0.00039787] )  ] ## for X antennas
    coefficients_lba = [ np.array( [ 3.85712631e+01, -2.17182149e+00,  1.68114451e+00 , 4.24076969e-01,  -9.24289199e-01,  2.11372242e-01, 1.09281884e-01, -1.74674795e-01, 3.70793388e-03] ),   ## for Y antennas
                         np.array( [38.13314007, -3.02861767 , 2.11558435, -0.30123627, -0.94641864, -0.14297615,  0.05037442, -0.04133833, -0.11443689])  ] ## for X antennas

    # Convert timestamp to datetime object
    t = datetime.datetime.utcfromtimestamp(timestamp)
    # Calculate JD(UT1)
    ut = tmf.gregorian2jd(t.year, t.month, float(t.day) + ((float(t.hour) + float(t.minute) / 60. + float(t.second) / 3600.) / 24.))
    # Calculate JD(TT)
    dtt = tmf.delta_tt_utc(tmf.date2jd(t.year, t.month, float(t.day) + ((float(t.hour) + float(t.minute) / 60. + float(t.second) / 3600.) / 24.)))
    tt = tmf.gregorian2jd(t.year, t.month, float(t.day) + ((float(t.hour) + float(t.minute) / 60. + (float(t.second) + dtt / 3600.)) / 24.))
    # Calculate Local Apparant Sidereal Time
    last = tmf.rad2circle(tmf.last(ut, tt, longitude))

    galactic_noise_power =[  fourier_series(last, coefficients_lba[0]),  fourier_series(last, coefficients_lba[1])  ]

    antenna_noise_power[ antenna_noise_power==0 ] = np.nan

    if antenna_type == 'outer':
        Y_measured_powers = antenna_noise_power[0::2]
        X_measured_powers = antenna_noise_power[1::2]

        Y_expected_power = galactic_noise_power[0]
        X_expected_power = galactic_noise_power[1]

    else:
        X_measured_powers = antenna_noise_power[0::2]
        Y_measured_powers = antenna_noise_power[1::2]

        X_expected_power = galactic_noise_power[0]
        Y_expected_power = galactic_noise_power[1]

    ## note this should make new arrays
    Y_factors = Y_expected_power / Y_measured_powers
    X_factors = X_expected_power / X_measured_powers

    np.sqrt(Y_factors, out=Y_factors)
    np.sqrt(X_factors, out=X_factors)

    if antenna_type == 'outer':
        return Y_factors, X_factors
    else:
        return X_factors, Y_factors


#!/usr/bin/env python

"""Module to deal with trees

Some more documentation

For now (and for some time), this module assumes that in the Trees folder,
there is only one set of tree_bricks (i.e. TreeMaker was run for only one set
of tree_bricks).

TODO:
    - everything
    - pep8

"""
import matplotlib
import numpy as np
#from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt
import pandas as pd
from scipy.io import FortranFile as FF
import yt
from glob import glob
import os
import yt.utilities.fortran_utils as fpu
from tqdm import tqdm

from . import physics
from .utils import find_outputs

class Forest(object):
    """
    Read the outputs from TreeMaker and gather in a dataframe
    in self.trees
    units are:
    x,y,z -> Mpccm
    vx,vy,xz -> km/s (to check)
    m,mvir -> 1e11Msun
    r,rvir -> Mpc
    """

    def __init__(self, Galaxy=False, path='.', big_run=False):
        paths = find_outputs(path+'/Outputs')
        self.prefix = path
        yt.funcs.mylog.setLevel(40)
        ds = yt.load(paths[-1])

        _sim = {}
        _sim['h'] = float(ds.cosmology.hubble_constant)
        _sim['Om'] = ds.omega_matter
        _sim['Ol'] = ds.omega_lambda
        _sim['Lbox'] = float(ds.length_unit.in_units('Mpccm'))

        treefolder = os.path.join(self.prefix, 'Trees')
        self.galaxies = False
        if Galaxy:
            treefolder = os.path.join(self.prefix, 'TreeStars')
            self.galaxies = True

        self.treefolder = treefolder
        self.sim = _sim
        self.ds = ds
        self.folder = treefolder
        self.snap = self._get_timestep_number()

        self.read_tree(big_run)

    def read_tree(self, big_run=False):
        """
        """

        if self.galaxies:
            tree_file = os.path.join(self.folder, 'tree.dat')
            self = self._read_tree(tree_file, 'd', big_run)
        else:
            tree_file = glob('{}/tree_file_???.001'.format(self.folder))[0]
            self = self._read_tree(tree_file, 'f', big_run)

        return

    def map_halo_ts_to_output(self, timestep):
        timestep = timestep.astype('int64')
        bricks = sorted(glob(os.path.join(self.treefolder, 'tree_bricks???')))
        mapping = np.empty(timestep.max(), dtype=np.int64)
        for istep in sorted(timestep.unique()):
            #ioutput = int(brick.split('bricks')[1])
            ioutput = int(bricks[istep-1].split('bricks')[1])
            mapping[istep-1] = ioutput

        return mapping[timestep-1]

    def get_all_progenitors(self, hnum, timestep=None):
        """Return the reduced tree containing ONLY the progenitors of halo hid
        ==========
        * hid: ID of the halo at output timestep given by the halo/galaxy finder
        * timestep (None): timestep at which the ID must be taken, default is last timestep
        """

        if timestep is None:
            ts = self.halo_ts.max()
        print('Care, this can be very long for galaxies. The get_main_progenitor is much faster')
        print('It is OK for Haloes')
        progenitors = self._get_progenitors(hnum, timestep=timestep)

        return progenitors

    def get_main_progenitor(self, hnum, timestep=None):
        """Return the reduced tree containing ONLY the main progenitors of halo hid
        ==========
        * hid: ID of the halo at output timestep given by the halo/galaxy finder
        * timestep (None): timestep at which the ID must be taken, default is last timestep
        """
        id_main = {}
        if timestep is None:
            ts = self.trees.halo_ts.max()
        else:
            ts = timestep

        list_progenitors = [(hnum, ts)]
        fatherMass = self.trees.loc[hnum,timestep].fatherMass
        fathersID = self.trees.loc[hnum,timestep].fathersID
        while (fathersID.size > 1):
            #update ts
            ts -= 1
            #remove background
            fatherMass = fatherMass[fathersID > 0]
            fathersID = fathersID[fathersID > 0]
            #main father is the one that contains most mass
            main_father = fathersID[fatherMass.argmax()]
            list_progenitors += [(main_father, ts)]
            fatherMass = self.trees.loc[main_father, ts].fatherMass
            fathersID = self.trees.loc[main_father, ts].fathersID
        main_progs = self.trees.loc[list_progenitors]

        return main_progs

    def get_main_children(self, hnum, timestep=None):
        """Return the reduced tree containing ONLY the main children of halo hid
        ==========
        * hid: ID of the halo at output timestep given by the halo/galaxy finder
        * timestep (None): timestep at which the ID must be taken, default is last timestep
        """
        # Get current timestep
        max_ts = self.trees.halo_ts.max()
        if timestep is None:
            ts = max_ts
        else:
            ts = timestep

        list_sons = [(hnum, ts)]
        while (ts < max_ts):
            main_son = self.trees.loc[hnum,ts].mainSon
            #update ts
            ts += 1
            list_sons += [(main_son, ts)]
        children = self.trees.loc[list_sons]

        return children

    def get_family(self, hnum, timestep=None):
        """Return the reduced tree containing ONLY the main progenitors/children of halo hid
        ==========
        * hid: ID of the halo at output timestep given by the halo/galaxy finder
        * timestep (None): timestep at which the ID must be taken, default is last timestep
        """
        # Get current timestep
        if timestep is None:
            ts = self.trees.halo_ts.max()
        else:
            ts = timestep

        child = self.get_main_children(hnum, timestep)
        fathers = self.get_main_progenitor(hnum, timestep)
        fathers = fathers.loc[fathers.halo_ts != ts]
        family = pd.concat((child, fathers))
        family = family.sort_index(level=1)

        return family

    def plot_all_trees(self, minmass=-1, maxmass=1e99, radius=1.0,
                       output=None, loc='./'):
        """
        See plot_halo_tree.
        Compute the above function for all halos with virial mass
        between `minmass` and `maxmass`. The width in Mpccm for the
        dynamics of the BH is radius. `loc` is the location to save the
        PDF.
        """

        if output is None:
            output = str(minmass) + 'To' + str(maxmass) + '_r' + str(radius)

        # Normalize min/max mass to 1e11
        minmass /= 1.e11
        maxmass /= 1.e11

        # Select halos at the last timestep of the tree, with the right mass
        # Also selects halos with at least 1 progenitor...
        mask = ((self.trees.halo_ts == self.snap) &
                (self.trees.mvir > minmass) &
                (self.trees.mvir < maxmass) &
                (self.trees.first_prog != -1))

        tid = self.trees[mask].halo_id
        print('Total: {} halos'.format(len(tid)))

        pdf = PdfPages(loc + 'trees{}.pdf'.format(output))
        try:
            fig = self.fig
        except AttributeError:
            self.fig = plt.figure(figsize=(12, 12))
        self.fig.savefig(pdf, format='pdf', dpi=200)

        for ihalo in tid:
            self.plot_halo_tree(hid=int(ihalo), radius=radius, pdffile=pdf)
            plt.close()
        pdf.close()
        plt.close('all')

        return 'OK'

    def plot_halo_tree(self, hid=None, hnum=None, hts=None, radius=1.0,
                       pdffile=None, loc='./'):
        """
        Plot Mass/Merger history of a halo
        Parameters
        ----------
        * hid (None): halo_id of the halo you want to have the merger/mass history
        * hnum (None): halo_num of the halo you want to have the
          merger history, if used then the timestep of this halo must
          be given with hts
        * radius (1): width of the window to plot the dynamics of the halo
        * loc ('./'): where to save the plot
        """

        # Define the selected halo
        if hid is None:
            hid = (self.trees.loc[(self.trees.halo_num == hnum) & (
                self.trees.halo_ts == hts)].halo_id).astype(int)
        halo = self.trees.ix[hid]
        sim = self.sim

        xh = halo['x']  # - sim['Lbox']/2 #/sim['Lbox'] + .5
        yh = halo['y']  # - sim['Lbox']/2   #/sim['Lbox'] + .5
        zh = halo['z']  # - sim['Lbox']/2   #/sim['Lbox'] + .5

        # Get progenitors
        progs = self._get_progenitors(hid)
        progs.loc[hid].descendent_id = -1
        main_progs = self.get_main_progenitor(hid)
        main_progs.loc[hid].descendent_id = -1

        # Recenter on selected halo and correct for periodic boundaries
        progs['x'] -= xh
        progs['y'] -= yh
        progs['z'] -= zh
        progs['x'].where(progs.x <= sim['Lbox'] / 2.,
                         progs.x - sim['Lbox'], inplace=True)
        progs['y'].where(progs.y <= sim['Lbox'] / 2.,
                         progs.y - sim['Lbox'], inplace=True)
        progs['z'].where(progs.z <= sim['Lbox'] / 2.,
                         progs.z - sim['Lbox'], inplace=True)
        progs['x'].where(progs.x >= -sim['Lbox'] / 2.,
                         progs.x + sim['Lbox'], inplace=True)
        progs['y'].where(progs.y >= -sim['Lbox'] / 2.,
                         progs.y + sim['Lbox'], inplace=True)
        progs['z'].where(progs.z >= -sim['Lbox'] / 2.,
                         progs.z + sim['Lbox'], inplace=True)
        main_progs['x'] -= xh
        main_progs['y'] -= yh
        main_progs['z'] -= zh
        main_progs['x'].where(main_progs.x <= sim['Lbox'] / 2.,
                              main_progs.x - sim['Lbox'], inplace=True)
        main_progs['y'].where(main_progs.y <= sim['Lbox'] / 2.,
                              main_progs.y - sim['Lbox'], inplace=True)
        main_progs['z'].where(main_progs.z <= sim['Lbox'] / 2.,
                              main_progs.z - sim['Lbox'], inplace=True)
        main_progs['x'].where(main_progs.x >= -sim['Lbox'] / 2.,
                              main_progs.x + sim['Lbox'], inplace=True)
        main_progs['y'].where(main_progs.y >= -sim['Lbox'] / 2.,
                              main_progs.y + sim['Lbox'], inplace=True)
        main_progs['z'].where(main_progs.z >= -sim['Lbox'] / 2.,
                              main_progs.z + sim['Lbox'], inplace=True)

        # Define plot range
        xc = (main_progs['x'].min() + main_progs['x'].max()) * .5
        yc = (main_progs['y'].min() + main_progs['y'].max()) * .5
        zc = (main_progs['z'].min() + main_progs['z'].max()) * .5
        xmin = ymin = zmin = -radius
        xmax = ymax = zmax = radius
        progs[['x', 'y', 'z']] -= [xc, yc, zc]
        main_progs[['x', 'y', 'z']] -= [xc, yc, zc]

        # Recenter other halos, and correct for periodic boundaries
        x = self.trees.x - (xc + xh)
        y = self.trees.y - (yc + yh)
        z = self.trees.z - (zc + zh)
        x.where(x <= sim['Lbox'] / 2., x - sim['Lbox'], inplace=True)
        y.where(y <= sim['Lbox'] / 2., y - sim['Lbox'], inplace=True)
        z.where(z <= sim['Lbox'] / 2., z - sim['Lbox'], inplace=True)
        x.where(x >= -sim['Lbox'] / 2., x + sim['Lbox'], inplace=True)
        y.where(y >= -sim['Lbox'] / 2., y + sim['Lbox'], inplace=True)
        z.where(z >= -sim['Lbox'] / 2., z + sim['Lbox'], inplace=True)

        # Select other halos within plot range
        others = ((x < radius) & (-radius < x) &
                  (y < radius) & (-radius < y) &
                  (z < radius) & (-radius < z))

        # Scalings
        scc = np.log10(progs.m)
        mainscc = np.log10(main_progs.m)
        sc = (scc - scc.min()) / (scc.max() - scc.min()) * 500.
        mainsc = (mainscc - scc.min()) / (mainscc.max() - scc.min()) * 500.

        osc = np.log10(self.trees[others].m)
        osc = (osc - scc.min()) / (scc.max() - scc.min()) * 500.
        ocol = self.trees[others].halo_ts
        ocol = .15 + .7 * (ocol - ocol.min()) / (ocol.max() - ocol.min())
        edg = np.where(self.trees[others].bush_id ==
                       halo['bush_id'], 'orange', 'k')

        # Plot halos dynamics in the comoving space
        # try:
        #    fig = self.fig
        # except AttributeError:
        self.fig = plt.figure(figsize=(12, 12))
        fig = self.fig
        # try:
        #    ax = self.ax
        # except AttributeError:
        ax10 = self.fig.add_axes([0.07, 0.05, 0.4, 0.4])
        ax00 = self.fig.add_axes([0.07, 0.5, 0.4, 0.4])
        ax01 = self.fig.add_axes([0.55, 0.5, 0.4, 0.4])
        ax11 = self.fig.add_axes([0.55, 0.05, 0.4, 0.4])
        self.ax = [ax00, ax01, ax10, ax11]
        ax = self.ax

        self.ax[0].scatter(x[others], y[others], c=ocol, s=osc, cmap='Greys',
                           vmin=0, vmax=1., edgecolor=edg,
                           rasterized=True)
        self.ax[0].scatter(progs.x, progs.y, c=progs.halo_ts, s=sc,
                           cmap='summer', rasterized=True)
        self.ax[0].scatter(main_progs.x, main_progs.y,
                           c=main_progs.halo_ts, s=mainsc,
                           cmap='magma', rasterized=True)
        self.ax[0].set_xlabel(r'$x$ (cMpc)')
        self.ax[0].set_ylabel(r'$y$ (cMpc)')
        self.ax[0].set_xlim(xmin, xmax)
        self.ax[0].set_ylim(ymin, ymax)

        self.ax[1].scatter(z[others], y[others], c=ocol, s=osc, cmap='Greys',
                           vmin=0, vmax=1., edgecolor=edg,
                           rasterized=True)
        self.ax[1].scatter(progs.z, progs.y, c=progs.halo_ts, s=sc,
                           cmap='summer', rasterized=True)
        self.ax[1].scatter(main_progs.z, main_progs.y,
                           c=main_progs.halo_ts, s=mainsc,
                           cmap='magma', rasterized=True)
        self.ax[1].set_xlabel(r'$z$ (cMpc)')
        self.ax[1].set_ylabel(r'$y$ (cMpc)')
        self.ax[1].set_xlim(zmin, zmax)
        self.ax[1].set_ylim(ymin, ymax)

        self.ax[2].scatter(x[others], z[others], c=ocol, s=osc, cmap='Greys',
                           vmin=0, vmax=1., edgecolor=edg,
                           rasterized=True)
        self.ax[2].scatter(progs.x, progs.z, c=progs.halo_ts, s=sc,
                           cmap='summer', rasterized=True)
        self.ax[2].scatter(main_progs.x, main_progs.z,
                           c=main_progs.halo_ts, s=mainsc,
                           cmap='magma', rasterized=True)
        self.ax[2].set_xlabel(r'$x$ (cMpc)')
        self.ax[2].set_ylabel(r'$z$ (cMpc)')
        self.ax[2].set_xlim(xmin, xmax)
        self.ax[2].set_ylim(zmin, zmax)

        # Draw some lines
        for progid in progs.halo_id:
            descid = progs.ix[progid].descendent_id
            if descid > 0:
                progenitor = progs.ix[progid]
                descendent = progs.ix[descid]
                px, py, pz = progenitor[['x', 'y', 'z']]
                dx, dy, dz = descendent[['x', 'y', 'z']]
                l0 = self.ax[0].plot([px, dx], [py, dy],
                                     lw=1, c='k', zorder=-1)
                l1 = self.ax[1].plot([pz, dz], [py, dy], lw=1, c='k',
                                     zorder=-1)
                l2 = self.ax[2].plot([px, dx], [pz, dz], lw=1, c='k',
                                     zorder=-1)

        title = "Halo #{t:d}\n(x, y, z) = ({x:3.3f}, {y:3.3f}, {z:3.3f})"
        title = title.format(t=int(halo.halo_num),
                             x=xh / sim['Lbox'] + .5,
                             y=yh / sim['Lbox'] + .5,
                             z=zh / sim['Lbox'] + .5)
        suptitle = self.fig.suptitle(title, fontsize=18)

        if not self.fig.texts:
            self.fig.texts.append(suptitle)
        if pdffile:
            self.fig.savefig(pdffile, format='pdf', dpi=200)
        else:
            plt.savefig(loc + 'halo_{:d}dynamics.png'.format(int(halo.halo_num)), dpi=100, format='png')

        # A bit of cleaning
        # for axx in self.ax:
        #    axx.clear()
        #self.fig.texts = []
        fig.clf()
        # try:
        #    fig = self.fig
        # except AttributeError:
        #    self.fig = plt.figure(figsize=(12, 12))
        # try:
        #    ax = self.ax
        # except AttributeError:
        ax10 = self.fig.add_axes([0.07, 0.05, 0.4, 0.4])
        ax00 = self.fig.add_axes([0.07, 0.5, 0.4, 0.4])
        ax01 = self.fig.add_axes([0.55, 0.5, 0.4, 0.4])
        ax11 = self.fig.add_axes([0.55, 0.05, 0.4, 0.4])
        self.ax = [ax00, ax01, ax10, ax11]
        ax = self.ax

        # Plot halo mass/merger history

        time = np.array([self.timestep['age'][int(hts) - 1]
                         for hts in progs.halo_ts])
        maintime = np.array([self.timestep['age'][int(hts) - 1]
                             for hts in main_progs.halo_ts])
        self.ax[0].scatter(time, progs.m * 1e11, c=progs.halo_ts,
                           s=sc, cmap='summer', rasterized=True)
        self.ax[0].semilogy(
            [self.ds.cosmology.t_from_z(z) / (3600 * 24 * 365 * 1e9)
             for z in np.logspace(-10, 2, 1000)],
            1e12 * physics.MofZ(halo.m * 1e11 / 1e12,
                                np.logspace(-10, 2, 1000),
                                1 / self.timestep['aexp'][progs.loc[hid].halo_ts - 1] - 1),
            'k')
        self.ax[0].scatter(maintime, main_progs.m * 1e11,
                           c=main_progs.halo_ts, s=mainsc,
                           cmap='magma', rasterized=True)
        self.ax[0].set_xlabel(r'$t$ (Gyr)')
        self.ax[0].set_ylabel(r'$M$ ($\mathrm{M}_{\odot}$)')
        self.ax[0].set_ylim(1e11 * progs.m.min() / 3.,
                            1e11 * progs.m.max() * 3.)
        self.ax[0].set_xlim(time.min(), time.max())
        self.ax[0].set_yscale('log')
        z = np.unique([int((1 / self.timestep['aexp'][int(hts) - 1] - 1) * 2) / 2.
                       for hts in main_progs.halo_ts])
        z = z[1:-1]
        ax_z = self.ax[0].twiny()
        ax_z.set_xlim(self.ax[0].get_xlim())
        ax_z.set_xticks([float(self.ds.cosmology.t_from_z(
            zz) / (1e9 * 365 * 24 * 3600)) for zz in z])
        ax_z.set_xticklabels(z)
        ax_z.set_xlabel("redshift")

        main_prog = self.get_main_progenitor(hid)
        minor_mergers = []
        major_mergers = []
        for current_id in main_prog.halo_id:
            if main_prog.halo_ts[current_id] != main_prog.halo_ts.min():
                halo_ts = main_prog.halo_ts[current_id]
                progs = self.get_all_progenitors(
                    current_id, timestep=halo_ts - 1)
                mask_minor = ((progs.m / main_prog.m[current_id + 1] > 1. / 20) &
                              (progs.m / main_prog.m[current_id + 1] < 1. / 4))
                minor_mergers += [len(progs.m[mask_minor])]
                mask_major = (progs.m / main_prog.m[current_id + 1] >= 1. / 4)
                major_mergers += [len(progs.m[mask_major]) - 1]

        scc = np.log10(main_prog.m)
        sc = (scc - scc.min()) / (scc.max() - scc.min()) * 500.
        time = np.array([self.timestep['age'][hts - 1]
                         for hts in main_prog[main_prog.halo_ts !=
                                              self.trees.halo_ts[hid]].halo_ts])
        self.ax[2].plot(time, # np.cumsum(major_mergers)[-1] -
                        np.cumsum(major_mergers),
                        color='b', label='1:1 > mass ratio > 1:4 (sim)')

        self.ax[2].plot(time,
                        [physics.N_merger_until_z(
                            0.25,
                            halo.m * 1e11 / 1e12,
                            self.ds.cosmology.z_from_t(t * 1e9 * 365 * 24 * 3600),
                            self.ds.cosmology.z_from_t(time.max() * 1e9 * 365 * 24 * 3600))
                         for t in time],
                        color='b', linestyle='--',
                        label='1:1 > mass ratio > 1:4 (theory)')

        self.ax[2].plot(time, # np.cumsum(minor_mergers)[-1] -
                        np.cumsum(minor_mergers), color='r',
                        label='1:4 > mass ratio > 1:20 (sim)')

        self.ax[2].plot(time,
                        [physics.N_merger_until_z(
                            1. / 20,
                            halo.m * 1e11 / 1e12,
                            self.ds.cosmology.z_from_t(t * 1e9 * 365 * 24 * 3600),
                            self.ds.cosmology.z_from_t(time.max() * 1e9 * 365 * 24 * 3600))
                         - physics.N_merger_until_z(
                             1. / 4,
                             halo.m * 1e11 / 1e12,
                             self.ds.cosmology.z_from_t(t * 1e9 * 365 * 24 * 3600),
                             self.ds.cosmology.z_from_t(time.max() * 1e9 * 365 * 24 * 3600))
                         for t in time],
                        color='r', linestyle='--',
                        label='1:4 > mass ratio > 1:20 (theory)')
        # self.ax[2].scatter(time, [minor_mergers[m] for m in major_mergers],
        #    c=main_prog[main_prog.halo_ts != main_prog.halo_ts.max()].halo_ts,
        #    s=sc, cmap='summer', rasterized=True)
        # self.ax[2].scatter(time, [major_mergers[m] for m in major_mergers],
        #    c=main_prog[main_prog.halo_ts != main_prog.halo_ts.max()].halo_ts,
        #    s=sc, cmap='magma', rasterized=True, marker='*')
        self.ax[2].legend(loc='best')
        self.ax[2].set_ylabel('#merger left before t$_{max}$')
        self.ax[2].set_xlabel(r'$t$ (Gyr)')
        self.ax[2].set_xlim(time.min(), time.max())
        self.ax[2].set_ylim(ymin=0)
        z = np.unique([int(1 / self.timestep['aexp'][int(hts) - 1] - 1)
                       for hts in main_progs.halo_ts])
        z = z[1:-1]
        ax_z = self.ax[2].twiny()
        ax_z.set_xlim(self.ax[0].get_xlim())
        ax_z.set_xticks([np.copy(self.ds.cosmology.t_from_z(
            zz) / (1e9 * 365 * 24 * 3600)) for zz in z])
        ax_z.set_xticklabels(z)

        title = "Halo #{t:d}\nFinal virial mass = {x:.2e} M$_\odot$"
        title = title.format(t=int(halo.halo_num),
                             x=halo['mvir'] * 1e11)
        suptitle = self.fig.suptitle(title, fontsize=18)

        if not self.fig.texts:
            self.fig.texts.append(suptitle)
        if pdffile:
            self.fig.savefig(pdffile, format='pdf', dpi=200)
        else:
            plt.savefig('halo_{:d}.png'.format(
                int(halo.halo_num)), dpi=100, format='png')

        # A bit of cleaning
        # for axx in self.ax:
        #    axx.clear()
        #self.fig.texts = []
        fig.clf()

        return hid

    def plot_accretion_history(self, hid, scale='mass'):
        """Plot halo accretion history.

        Blah
        """
        # Define the selected halo
        halo = self.trees.ix[hid]
        sim = self.sim

        # Get progenitors
        progs = self._get_progenitors(hid)

        # Scalings
        scc = np.log10(progs.m)
        sc = (scc - scc.min()) / (scc.max() - scc.min()) * 500.

        try:
            fig = self.fig
        except AttributeError:
            self.fig = plt.figure(figsize=(12, 12))
        try:
            ax = self.ax
        except AttributeError:
            self.ax = self.fig.add_subplot(111)

        time = np.array([self.timestep['age'][hts - 1]
                         for hts in progs.halo_ts])
        self.ax.scatter(time, progs.m * 1e11, c=progs.halo_ts, s=sc,
                        cmap='summer', rasterized=True)
        self.ax.set_xlabel(r'$t$ (Gyr)')
        self.ax.set_ylabel(r'$M$ ($\mathrm{M}_{\odot}$)')
        self.ax.set_ylim(1e11 * progs.m.min() / 3., 1e11 * progs.m.max() * 3.)
        self.ax.set_xlim(time.min(), time.max())
        self.ax.set_yscale('log')

        self.fig.tight_layout()
        self.fig.savefig('halo_{:d}_history.pdf'.format(
            int(halo.halo_num)), dpi=200, format='pdf')

        return hid

    # CONVENIENCE FUNCTIONS
    def _get_timestep_number(self):
        l = len(glob(os.path.join(self.folder, 'halos_results.???')))
        return l

    def get_ioutlist(self, inputfile='input_TreeMaker.dat'):
        """Return the number of treebricks files from an input file

        The input_TreeMaker.dat file must be a list of:
            nsnaps some_number
            '/path/to/tree_bricks{n:03d}'
            '/path/to/tree_bricks{n+1:03d}'
            ...
            '/path/to/tree_bricks{n+nsnaps:03d}'
            ...

        What should be done:
            * Read all the /path/to/tree_bricks lines
            * For each line, get the tree_bricks number
            * Return the mapping between [1, nsnaps] and the tree_bricks list.
        """
        import os.path

        with open(self.folder + inputfile, 'r') as ifile:
            lines = ifile.readlines()

        outputs = []

        # We know that the name of the tree_bricks file should be
        # "tree_bricksXXX"
        tbbasename = 'tree_bricks'
        # We can skip the first line, which should be "nsnaps 1"
        for path in lines[1:]:
            # Remove those nasty "'" and \n
            tbname = os.path.basename(path[1:-2])
            assert tbname.startswith(tbbasename)
            outputs.append(int(tbname[len(tbbasename):]))

        return outputs

    def _get_progenitors(self, hid, timestep):
        if self.galaxies:
            progenitors = self.trees.loc[(self.trees.halo_num == hid) & (self.trees.halo_ts == timestep)]
            if len(progenitors) == 0:
                return
            else:
                if len(progenitors.fathersID.item()) == 1:
                    return progenitors
                else:
                    for fatherID in progenitors.fathersID.item():
                        progenitors = pd.concat([progenitors, self._get_progenitors(fatherID, timestep-1)])
                    return progenitors

        else:
            target = self.trees.loc[(self.trees.halo_num == hid) & (self.trees.halo_ts == timestep)].index
            print(target)
            mask = ((self.trees.halo_id >= self.trees.loc[target].halo_id.item()) &
                    (self.trees.halo_id <= self.trees.loc[target].last_prog.item()))

            progenitors = self.trees.loc[mask].copy()
            return progenitors

    def _read_halo(self, F, precision, big_run=False):
        [ID] = fpu.read_vector(F, 'i')
        [BushID] = fpu.read_vector(F, 'i')
        [timestep] = fpu.read_vector(F, 'i')
        [level, hosthaloID, hostsubID, nbsub, nextsub] = fpu.read_vector(F, 'i')
        [m] = fpu.read_vector(F, precision)
        [dmacc] = fpu.read_vector(F, 'd')
        [x, y, z] = fpu.read_vector(F, precision)
        [vx, vy, vz] = fpu.read_vector(F, precision)
        [Lx, Ly, Lz] = fpu.read_vector(F, precision)
        [r, a, b, c] = fpu.read_vector(F, precision)
        [ek, ep, et] = fpu.read_vector(F, precision)
        [spin] = fpu.read_vector(F, precision)
        [nbfather] = fpu.read_vector(F, 'i')
        if nbfather == 0:
            fatherID = []
            fatherMass = []
        else:
            fatherID = fpu.read_vector(F, 'i')
            fatherMass = fpu.read_vector(F, precision)
        [nbsons] = fpu.read_vector(F, 'i')
        if nbsons == 0:
            sonsID = []
            mainSon = []
        else:
            sonsID = fpu.read_vector(F, 'i')
        mainSon = fpu.read_vector(F, 'i')[0]
        if self.galaxies:
            [rvir, mvir, tvir, cvel,Reff] = fpu.read_vector(F, precision)
        else:
            [rvir, mvir, tvir, cvel] = fpu.read_vector(F, precision)
        [rho_0, r_c] = fpu.read_vector(F, precision)
        if not big_run:
            fpu.read_vector(F, 'i')  # ncont

        return (ID, timestep, level, hosthaloID, hostsubID, m,
                           dmacc, x, y, z, vx, vy, vz, Lx, Ly, Lz, r, a, b,
                           c, ek, ep, et, spin, rvir, mvir, tvir, cvel,
                           fatherID, fatherMass, sonsID, mainSon, Reff)

    def _read_tree(self, tree_file, precision, big_run=False):
        Key_tree = [
            'halo_num', 'tree_step', 'level', 'host_halo_id', 'host_sub_id', 'm',
            'dmacc', 'x', 'y', 'z', 'vx', 'vy', 'vz', 'Lx', 'Ly', 'Lz', 'r',
            'a', 'b', 'c', 'ek', 'ep', 'et', 'spin', 'rvir', 'mvir', 'tvir',
            'cvel', 'fathersID', 'fatherMass', 'sonsID', 'mainSon']
        if self.galaxies:
            Key_tree += ['Reff']

        with open(tree_file, 'rb') as F:
            self.timestep = {}
            [self.timestep['nsteps']] = fpu.read_vector(F, 'i')
            n_halo_tot = fpu.read_vector(F, 'i')
            self.timestep['nhalos'] = n_halo_tot[:self.timestep['nsteps']]
            self.timestep['aexp'] = fpu.read_vector(F, precision)
            fpu.read_vector(F, precision)
            self.timestep['age'] = fpu.read_vector(F, precision)

            data = np.empty(shape=(n_halo_tot.sum(), len(Key_tree)), dtype=object)
            if (os.path.exists('{}/tree.hdf'.format(self.folder))):
              self.trees = pd.read_hdf('{}/tree.hdf'.format(self.folder))
            else:
              print('Reading the tree for the first time... be patient!')
              j = 0
              for istep in tqdm(range(self.timestep['nsteps'])):
                for ihalo in range(n_halo_tot[istep]+n_halo_tot[istep+self.timestep['nsteps']]):
                    data[j] = self._read_halo(F, precision, big_run)
                    j += 1

              dd = {k: data[:, i]
                  for i, k in enumerate(Key_tree)}
              self.trees = pd.DataFrame(dd)
              for k in ['halo_num', 'tree_step', 'level', 'host_halo_id', 'host_sub_id']:
                    self.trees[k] = self.trees[k].astype(np.int32)
              for k in ['m', 'dmacc', 'x', 'y', 'z', 'vx', 'vy', 'vz', 'Lx', 'Ly',
                  'Lz', 'r', 'a', 'b', 'c', 'ek', 'ep', 'et', 'spin', 'rvir',
                  'mvir', 'tvir', 'cvel']:
                self.trees[k] = self.trees[k].astype(np.float32)

              #process a bit the data
              self.timestep['aexp'] = 1/(1+self.ds.cosmology.z_from_t(self.ds.arr(self.timestep['age'],'Gyr')))
              ####### Create halo_ts from the step in the tree
              self.trees['halo_ts'] = self.map_halo_ts_to_output(self.trees.tree_step)
              iout = self.trees.tree_step.values.astype(int)-1
              self.trees['aexp'] = self.timestep['aexp'][iout]
              aexp = self.trees['aexp']
              aexp_last = self.timestep['aexp'].max()

              factor = float(self.ds.length_unit.in_units('cm') / 3.08e24) * aexp / aexp_last
              self.trees['x'] = self.trees['x'] / factor + 0.5
              self.trees['y'] = self.trees['y'] / factor + 0.5
              self.trees['z'] = self.trees['z'] / factor + 0.5

              # Find outputs given halo_ts
              paths = find_outputs(self.prefix+'/Outputs')
              self.outputs = paths[-int(self.trees.halo_ts.max()):]
              step_first_gal = len(paths) - self.trees.halo_ts.max()
              self.trees.halo_ts += step_first_gal

              self.trees.m *= 1e11
              self.trees.dmacc *= 1e11

              self.trees.set_index(['halo_num','halo_ts'], drop=False, inplace=True)
              self.trees.to_hdf('{}/tree.hdf'.format(self.folder), 'hdf5')

        return self