# Licensed under a 3-clause BSD style license - see LICENSE
'''
This module provides the framework for the creation of the photometric likelihoods
used in the cross-matching of the two catalogues.
'''

import os
import sys
import numpy as np

from .misc_functions import map_large_index_to_small_index, _load_single_sky_slice
from .misc_functions_fortran import misc_functions_fortran as mff
from .photometric_likelihood_fortran import photometric_likelihood_fortran as plf

__all__ = ['compute_photometric_likelihoods']

def compute_photometric_likelihoods(joint_folder_path, a_cat_folder_path, b_cat_folder_path,
                                    afilts, bfilts, mem_chunk_num, cf_points, cf_areas,
                                    include_phot_like, use_phot_priors, bright_frac=None,
                                    field_frac=None):
    '''
    Derives the photometric likelihoods and priors for use in the catalogue
    cross-match process.

    Parameters
    ----------
    joint_folder_path : string
        The folder where all folders and files created during the cross-match
        process are stored.
    a_cat_folder_path : string
        The folder where the input data for catalogue "a" are located.
    b_cat_folder_path : string
        The location of catalogue "b"'s input data.
    afilts : list of string
        A list of the filters in catalogue "a"'s photometric data file.
    bfilts : list of string
        List of catalogue "b"'s filters.
    mem_chunk_num : integer
        Fraction of input datasets to load at once, in the case of data larger
        than the memory of the system.
    cf_points : numpy.ndarray
        The on-sky coordinates that define the locations of each small
        set of sources to be used to derive the relative match and non-match
        photometric likelihoods.
    cf_areas : numpy.ndarray
        The areas of closest on-sky separation surrounding each point in
        ``cf_points``, used to normalise numbers of sources to sky densities.
    include_phot_like : boolean
        Flag to indicate whether to derive astrophysical likelihoods ``c`` and
        ``f``, based on the common coevality of sources of given magnitudes.
    use_phot_priors : boolean
        Indicator as to whether to use astrophysical priors, based on the common
        number of likely matches and non-matches in each ``cf_points`` area, or
        use naive, asymmetric priors solely based on number density of sources.
    bright_frac : float, optional
        Expected fraction of sources inside the "bright" error circles used to
        construct the counterpart distribution, to correct for missing numbers.
        If ``include_phot_like`` or ``use_phot_prior`` is True then this must
        be supplied, otherwise it can be omitted.
    field_frac : float, optional
        Expected fraction of sources inside the "field" error circles used to
        construct the counterpart distribution, to correct for missing numbers.
        If ``include_phot_like`` or ``use_phot_prior`` is True then this must
        be supplied, otherwise it can be omitted.
    '''

    if bright_frac is None and (include_phot_like or use_phot_priors):
        raise ValueError("bright_frac must be supplied if include_phot_like or use_phot_priors "
                         "is set to True. Please supply an appropriate fraction.")
    if field_frac is None and (include_phot_like or use_phot_priors):
        raise ValueError("field_frac must be supplied if include_phot_like or use_phot_priors "
                         "is set to True. Please supply an appropriate fraction.")

    print("Creating c(m, m) and f(m)...")

    len_a = len(np.load('{}/con_cat_astro.npy'.format(a_cat_folder_path), mmap_mode='r'))
    len_b = len(np.load('{}/con_cat_astro.npy'.format(b_cat_folder_path), mmap_mode='r'))

    print("Distributing sources into sky slices...")
    sys.stdout.flush()

    a_sky_inds = distribute_sky_indices(joint_folder_path, a_cat_folder_path, 'a', mem_chunk_num,
                                        cf_points)
    b_sky_inds = distribute_sky_indices(joint_folder_path, b_cat_folder_path, 'b', mem_chunk_num,
                                        cf_points)

    print("Making bins...")
    sys.stdout.flush()

    abinlengths, abinsarray, longabinlen = create_magnitude_bins(
        cf_points, afilts, mem_chunk_num, joint_folder_path, a_cat_folder_path, 'a', a_sky_inds,
        include_phot_like or use_phot_priors)
    bbinlengths, bbinsarray, longbbinlen = create_magnitude_bins(
        cf_points, bfilts, mem_chunk_num, joint_folder_path, b_cat_folder_path, 'b', b_sky_inds,
        include_phot_like or use_phot_priors)

    print("Calculating PDFs...")
    sys.stdout.flush()

    c_priors = np.lib.format.open_memmap(
        '{}/phot_like/c_priors.npy'.format(joint_folder_path), mode='w+', dtype=float,
        shape=(len(bfilts), len(afilts), len(cf_points)), fortran_order=True)
    fa_priors = np.lib.format.open_memmap(
        '{}/phot_like/fa_priors.npy'.format(joint_folder_path), mode='w+', dtype=float,
        shape=(len(bfilts), len(afilts), len(cf_points)), fortran_order=True)
    fb_priors = np.lib.format.open_memmap(
        '{}/phot_like/fb_priors.npy'.format(joint_folder_path), mode='w+', dtype=float,
        shape=(len(bfilts), len(afilts), len(cf_points)), fortran_order=True)
    c_array = np.lib.format.open_memmap(
        '{}/phot_like/c_array.npy'.format(joint_folder_path), mode='w+', dtype=float,
        shape=(longbbinlen-1, longabinlen-1, len(bfilts), len(afilts),
               len(cf_points)), fortran_order=True)
    fa_array = np.lib.format.open_memmap(
        '{}/phot_like/fa_array.npy'.format(joint_folder_path), mode='w+', dtype=float,
        shape=(longabinlen-1, len(bfilts), len(afilts), len(cf_points)), fortran_order=True)
    fb_array = np.lib.format.open_memmap(
        '{}/phot_like/fb_array.npy'.format(joint_folder_path), mode='w+', dtype=float,
        shape=(longbbinlen-1, len(bfilts), len(afilts), len(cf_points)), fortran_order=True)

    # Within each loop, since we've already assigned all sources to have a closest
    # c/f sky position pointing to be assigned to, we simply slice within the _load
    # functions on ID, and for the initial mem_chunk_num stepping load in the range
    # [m_, m_+mem_chunk_num).
    for m_ in range(0, len(cf_points), mem_chunk_num):
        a_multi_return = _load_multiple_sky_slice(joint_folder_path, 'a', m_, m_+mem_chunk_num,
                                                  a_cat_folder_path, a_sky_inds,
                                                  include_phot_like or use_phot_priors)
        if include_phot_like or use_phot_priors:
            (a_small_photo, a_sky_ind_small, a_small_astro, a_blen_small, a_inds_small,
             a_size_small) = a_multi_return
        else:
            a_small_photo, a_sky_ind_small = a_multi_return
        del a_multi_return

        b_multi_return = _load_multiple_sky_slice(joint_folder_path, 'b', m_, m_+mem_chunk_num,
                                                  b_cat_folder_path, b_sky_inds,
                                                  include_phot_like or use_phot_priors)
        if include_phot_like or use_phot_priors:
            (b_small_photo, b_sky_ind_small, b_small_astro, b_blen_small, b_inds_small,
             b_size_small) = b_multi_return
            del b_blen_small
        else:
            b_small_photo, b_sky_ind_small = b_multi_return
        del b_multi_return

        for m in range(m_, min(len(cf_points), m_+mem_chunk_num)):
            area = cf_areas[m]
            a_sky_cut = _load_single_sky_slice(
                joint_folder_path, 'a', m, a_sky_ind_small)
            a_photo_cut = a_small_photo[a_sky_cut]
            if include_phot_like or use_phot_priors:
                a_astro_cut, a_blen_cut, a_inds_cut, a_size_cut = (
                    a_small_astro[a_sky_cut], a_blen_small[a_sky_cut], a_inds_small[:, a_sky_cut],
                    a_size_small[a_sky_cut])

            b_sky_cut = _load_single_sky_slice(
                joint_folder_path, 'b', m, b_sky_ind_small)
            b_photo_cut = b_small_photo[b_sky_cut]
            if include_phot_like or use_phot_priors:
                b_astro_cut, b_inds_cut, b_size_cut = (
                    b_small_astro[b_sky_cut], b_inds_small[:, b_sky_cut], b_size_small[b_sky_cut])

                # Return the overall to subarray mapping of *_inds_cut, as well
                # as the unique values of *_inds_cut, for use in creating the
                # opposing view slices into the other catalogue, for each
                # catalogue. Note that the lengths are swapped in the two calls,
                # as the a_inds map into length of catalogue "b" and vice versa.
                a_inds_map, a_inds_cut_unique = map_large_index_to_small_index(
                    a_inds_cut, len_b, '{}/phot_like'.format(joint_folder_path))
                b_inds_map, b_inds_cut_unique = map_large_index_to_small_index(
                    b_inds_cut, len_a, '{}/phot_like'.format(joint_folder_path))

                # Combined with b_inds_map, this subarray of the catalogue gives
                # an array that has every "a" source that overlaps the "b"
                # subarray objects. Note we use b_inds_cut_unique for catalogue
                # "a" sources.
                a_astro_ind = np.load('{}/con_cat_astro.npy'.format(a_cat_folder_path),
                                      mmap_mode='r')[b_inds_cut_unique]
                b_astro_ind = np.load('{}/con_cat_astro.npy'.format(b_cat_folder_path),
                                      mmap_mode='r')[a_inds_cut_unique]
                a_photo_ind = np.load('{}/con_cat_photo.npy'.format(a_cat_folder_path),
                                      mmap_mode='r')[b_inds_cut_unique]
                b_photo_ind = np.load('{}/con_cat_photo.npy'.format(b_cat_folder_path),
                                      mmap_mode='r')[a_inds_cut_unique]
                a_flen_ind = np.load('{}/group/aflen.npy'.format(joint_folder_path),
                                     mmap_mode='r')[b_inds_cut_unique]
                b_flen_ind = np.load('{}/group/bflen.npy'.format(joint_folder_path),
                                     mmap_mode='r')[a_inds_cut_unique]

            for i in range(0, len(afilts)):
                if not include_phot_like and not use_phot_priors:
                    a_num_photo_cut = np.sum(~np.isnan(a_photo_cut[:, i]))
                    Na = a_num_photo_cut / area
                else:
                    a_bins = abinsarray[:abinlengths[i, m], i, m]
                    a_mag = a_photo_cut[:, i]
                    a_flags = ~np.isnan(a_mag)
                    a_flags_ind = ~np.isnan(a_photo_ind[:, i])
                for j in range(0, len(bfilts)):
                    if not include_phot_like and not use_phot_priors:
                        b_num_photo_cut = np.sum(~np.isnan(b_photo_cut[:, j]))
                        Nb = b_num_photo_cut / area
                        # Without using photometric-based priors, all we can
                        # do is set the prior on one catalogue to 0.5 -- that
                        # is, equal chance of match or non-match; for this we
                        # use the less dense of the two catalogues as our
                        # "one-sided" match. Then, accordingly, we update the
                        # "field" source density of the more dense catalogue
                        # with its corresponding density, based on the input
                        # density and the counterpart density calculated.
                        c_prior = min(Na, Nb) / 2
                        fa_prior = Na - c_prior
                        fb_prior = Nb - c_prior
                        # To fake no photometric likelihoods, simply set all
                        # values to one, to cancel in the ratio later.
                        c_like, fa_like, fb_like = 1-1e-10, 1-1e-10, 1-1e-10
                    else:
                        b_bins = bbinsarray[:bbinlengths[j, m], j, m]
                        b_mag = b_photo_cut[:, j]
                        b_flags = ~np.isnan(b_mag)
                        b_mag_ind = b_photo_ind[:, j]
                        b_flags_ind = ~np.isnan(b_mag_ind)

                        c_prior, c_like, fa_prior, fa_like, fb_prior, fb_like = create_c_and_f(
                            a_astro_cut, b_astro_cut, a_mag, b_mag, a_inds_map, a_size_cut,
                            b_inds_map, b_size_cut, a_blen_cut, a_bins, b_bins, bright_frac,
                            field_frac, a_flags, b_flags, a_astro_ind, b_astro_ind, a_flen_ind,
                            b_flen_ind, a_flags_ind, b_flags_ind, b_mag_ind, area)
                    if use_phot_priors and not include_phot_like:
                        # If we only used the create_c_and_f routine to derive
                        # priors, then quickly update likelihoods here.
                        c_like, fa_like, fb_like = 1-1e-10, 1-1e-10, 1-1e-10

                    # Have to add a very small "fire extinguisher" value to all
                    # likelihoods and priors, to avoid ever having exactly zero
                    # value in either, which would mean all island permutations
                    # were rejected.
                    c_priors[j, i, m] = c_prior + 1e-10
                    fa_priors[j, i, m] = fa_prior + 1e-10
                    fb_priors[j, i, m] = fb_prior + 1e-10
                    c_array[:bbinlengths[j, m]-1,
                            :abinlengths[i, m]-1, j, i, m] = c_like + 1e-10
                    fa_array[:abinlengths[i, m]-1, j, i, m] = fa_like + 1e-10
                    fb_array[:bbinlengths[j, m]-1, j, i, m] = fb_like + 1e-10

    os.system('rm {}/a_small_sky_slice.npy'.format(joint_folder_path))
    os.system('rm {}/b_small_sky_slice.npy'.format(joint_folder_path))

    # *binsarray is passed back from create_magnitude_bins as a memmapped array,
    # but *binlengths is just a numpy array, so quickly save these before returning.
    np.save('{}/phot_like/abinlengths.npy'.format(joint_folder_path), abinlengths)
    np.save('{}/phot_like/bbinlengths.npy'.format(joint_folder_path), bbinlengths)

    return

def distribute_sky_indices(joint_folder_path, cat_folder, name, mem_chunk_num, cf_points):
    '''
    Function to calculate the nearest on-sky photometric likelihood point for
    each catalogue source.

    Parameters
    ----------
    joint_folder_path : string
        Top-level folder path for the common files created during the cross-match
        process.
    cat_folder : string
        The location of a given catalogue's input data files.
    name : string
        Representation of whether we are calculating sky indices for catalogue
        "a", or catalogue "b".
    mem_chunk_num : integer
        Number of sub-sets to break larger data files down to, to preserve memory.
    cf_points : numpy.ndarray
        The two-point sky coordinates for each point to be used as a central
        point of a small sky area, for calculating "counterpart" and "field"
        photometric likelihoods.

    Returns
    -------
    sky_inds : numpy.ndarray
        The indices, matching ``cf_points``, of the closest sky position for each
        source in this catalogue.
    '''
    n_sources = len(np.load('{}/con_cat_astro.npy'.format(cat_folder), mmap_mode='r'))
    sky_inds = np.lib.format.open_memmap('{}/phot_like/{}_sky_inds.npy'.format(
                                         joint_folder_path, name), mode='w+', dtype=int,
                                         shape=(n_sources,), fortran_order=True)
    for cnum in range(0, mem_chunk_num):
        lowind = np.floor(n_sources*cnum/mem_chunk_num).astype(int)
        highind = np.floor(n_sources*(cnum+1)/mem_chunk_num).astype(int)
        a = np.load('{}/con_cat_astro.npy'.format(cat_folder), mmap_mode='r')[lowind:highind]
        sky_inds[lowind:highind] = mff.find_nearest_point(a[:, 0], a[:, 1],
                                                          cf_points[:, 0], cf_points[:, 1])

    return sky_inds

def create_magnitude_bins(cf_points, filts, mem_chunk_num, joint_folder_path,
                          cat_folder_path, cat_type, sky_inds, load_extra_arrays):
    '''
    Creates the N-dimensional arrays of single-band photometric bins, and
    corresponding array lengths.

    Parameters
    ----------
    cf_points : numpy.ndarray
        List of the two-dimensional on-sky coordinates defining the centers
        of each cutout for which the photometric likelihoods should be
        calculated.
    filts : list of strings
        List of the filters to create magnitude bins for in this catalogue.
    mem_chunk_num : integer
        Number of sub-sets to break larger catalogues down into, for memory
        saving purposes.
    joint_folder_path : string
        Location of top-level folder into which all intermediate files are
        saved for the cross-match process.
    cat_folder_path : string
        Location of the input data for this catalogue.
    cat_type : string
        String to indicate which catalogue we are creating bins for, either
        "a", or "b".
    sky_inds : numpy.ndarray
        Array of indices, showing which on-sky photometric point, from
        ``cf_points``, each source in the catalogue is closest to.
    load_extra_arrays : boolean
        Flag to indicate whether the photometric information is being used in
        the cross-match process, and whether to load additional arrays accordingly.

    Returns
    -------
    binlengths : numpy.ndarray
        Two-dimensional array, indicating the length of the magnitude bins for
        each filter-sky coordinate combination.
    binsarray : numpy.ndarray
        Three-dimensional array, containing the values of the magnitude bins for
        the filter-sky combinations.
    longbinlen : integer
        Value of the largest of all filter-sky combinations of ``binlengths``.
    '''
    binlengths = np.empty((len(filts), len(cf_points)), int)

    for m_ in range(0, len(cf_points), mem_chunk_num):
        a_multi_return = _load_multiple_sky_slice(
            joint_folder_path, cat_type, m_, m_+mem_chunk_num, cat_folder_path, sky_inds,
            load_extra_arrays)
        if load_extra_arrays:
            (a_phot_, sky_inds_, _, _, _, _) = a_multi_return
        else:
            a_phot_, sky_inds_ = a_multi_return
        del a_multi_return
        for m in range(m_, min(len(cf_points), m_+mem_chunk_num)):
            sky_cut = _load_single_sky_slice(
                joint_folder_path, cat_type, m, sky_inds_)
            for i in range(0, len(filts)):
                a = a_phot_[sky_cut, i]
                if np.sum(~np.isnan(a)) > 0:
                    f = make_bins(a[~np.isnan(a)])
                else:
                    f = np.array([0])
                del a
                binlengths[i, m] = len(f)

    longbinlen = np.amax(binlengths)
    binsarray = np.lib.format.open_memmap(
        '{}/phot_like/{}binsarray.npy'.format(joint_folder_path, cat_type), mode='w+', dtype=float,
        shape=(longbinlen, len(filts), len(cf_points)), fortran_order=True)
    binsarray[:, :, :] = -1
    for m_ in range(0, len(cf_points), mem_chunk_num):
        a_multi_return = _load_multiple_sky_slice(
            joint_folder_path, cat_type, m_, m_+mem_chunk_num, cat_folder_path, sky_inds,
            load_extra_arrays)
        if load_extra_arrays:
            (a_phot_, sky_inds_, _, _, _, _) = a_multi_return
        else:
            a_phot_, sky_inds_ = a_multi_return
        for m in range(m_, min(len(cf_points), m_+mem_chunk_num)):
            sky_cut = _load_single_sky_slice(
                joint_folder_path, cat_type, m, sky_inds_)
            for i in range(0, len(filts)):
                a = a_phot_[sky_cut, i]
                if np.sum(~np.isnan(a)) > 0:
                    f = make_bins(a[~np.isnan(a)])
                else:
                    f = np.array([0])
                del a
                binsarray[:binlengths[i, m], i, m] = f

    return binlengths, binsarray, longbinlen

def make_bins(input_mags):
    '''
    Calculate bins for a catalogue's magnitude distribution, ensuring all stars
    are in histogram bins of sufficient number statistics.

    Parameters
    ----------
    input_mags : numpy.ndarray
        Array of magnitudes of given filter from the specific catalogue, to be
        placed in a histogram.

    Returns
    -------
    output_bins : numpy.ndarray
        Bins for the given catalogue-filter combination that produce robust
        numbers of sources within each magnitude interval.
    '''
    minamag = np.amin(input_mags)
    maxamag = np.amax(input_mags)
    da = 0.1
    maxa = da*np.ceil(maxamag/da)
    mina = da*np.floor(minamag/da)
    na = int(np.ceil((maxa - mina)/da) + 1)
    output_bins = np.linspace(mina, maxa, na)
    # If min/max magnitudes that define magnitude bins happen to lie exactly
    # on a bin edge (i.e., maxamag % da == 0), then just pad bin edge slightly.
    if np.abs(mina - minamag) < 1e-5:
        output_bins[0] -= 1e-4
    if np.abs(maxa - maxamag) < 1e-5:
        output_bins[-1] += 1e-4

    hist, output_bins = np.histogram(input_mags, bins=output_bins)
    smalllist = []
    # Minimum number statistics in each 1-D bin.
    minnum = 250

    for i in range(0, len(output_bins)-1):
        if hist[i] < minnum:
            smalllist.extend([i])
    smalllist = np.array(smalllist)
    dellist = []
    if len(smalllist) > 0:
        for i in smalllist:
            if i not in dellist:
                flag = 0
                for j in range(i+1, len(output_bins)-1):
                    if np.sum(hist[i:j+1]) > minnum:
                        dellist.extend([k for k in range(i+1, j+1)])
                        flag = 1
                        break
                if flag == 0:
                    dellist.extend([k for k in range(i+1, len(output_bins)-1)])
    output_bins = np.delete(output_bins, dellist)

    return output_bins

def _load_multiple_sky_slice(joint_folder_path, cat_name, ind1, ind2, cat_folder_path, sky_inds,
                             load_extra_arrays):
    '''
    Function to, in a memmap-friendly way, return a sub-set of the photometry
    of a given catalogue.

    Parameters
    ----------
    joint_folder_path : string
        Folder in which common cross-match intermediate data files are stored.
    cat_name : string
        String defining whether this function was called on catalogue "a" or "b".
    ind1 : float
        The lower of the two sky indices, as defined in ``distribute_sky_indices``,
        to return a sub-set of the larger catalogue between. This value represents
        the index of a given on-sky position, used to construct the "counterpart"
        and "field" likelihoods.
    ind2 : float
        The upper of the sky indices, defining the sub-set of the photometric
        array to return.
    cat_folder_path : string
        The folder defining where this particular catalogue is stored.
    sky_inds : numpy.ndarray
        The given catalogue's ``distribute_sky_indices`` values, to compare
        with ``ind1`` and ``ind2``.
    load_extra_arrays : boolean
        Flag to indicate whether the photometric information is being used in
        the cross-match process, and whether to load additional arrays accordingly.

    Returns
    -------
    photo_cutout : numpy.ndarray
        A sub-set of the photometry of the given catalogue, those points which are
        astrometrically closest to the sky indices between ``ind1`` and ``ind2``.
    sky_ind_cutout : numpy.ndarray
        The reduced ``sky_inds`` array, containing only those between ``ind1`` and
        ``ind2``.
    list_of_arrays : list of numpy.ndarrays
        Depending on whether ``load_extra_arrays`` is ``True`` or not, this list
        contains either just a cutout of the photometric data for this catalogue
        and the corresponding subset of the sky index array, or it also contains
        subsets of the astrometry array, and "bright" source error circle length,
        and overlap index and size arrays.
    '''
    sky_cut = np.lib.format.open_memmap('{}/{}_temporary_sky_slice_combined.npy'.format(
        joint_folder_path, cat_name), mode='w+', dtype=bool, shape=(len(sky_inds),))

    di = max(1, len(sky_inds) // 20)

    for i in range(0, len(sky_inds), di):
        sky_cut[i:i+di] = (sky_inds[i:i+di] >= ind1) & (sky_inds[i:i+di] < ind2)

    if load_extra_arrays:
        astro_cutout = np.load('{}/con_cat_astro.npy'.format(cat_folder_path),
                               mmap_mode='r')[sky_cut]
        a_blen_cutout = np.load('{}/group/{}blen.npy'.format(joint_folder_path, cat_name),
                                mmap_mode='r')[sky_cut]
        a_inds_cutout = np.load('{}/group/{}inds.npy'.format(joint_folder_path, cat_name),
                                mmap_mode='r')[:, sky_cut]
        a_size_cutout = np.load('{}/group/{}size.npy'.format(joint_folder_path, cat_name),
                                mmap_mode='r')[sky_cut]

    photo_cutout = np.load('{}/con_cat_photo.npy'.format(cat_folder_path), mmap_mode='r')[sky_cut]
    sky_ind_cutout = np.load('{}/phot_like/{}_sky_inds.npy'.format(joint_folder_path, cat_name),
                             mmap_mode='r')[sky_cut]

    if load_extra_arrays:
        list_of_arrays = (photo_cutout, sky_ind_cutout, astro_cutout, a_blen_cutout, a_inds_cutout,
                          a_size_cutout)
    else:
        list_of_arrays = (photo_cutout, sky_ind_cutout)

    os.system('rm {}/{}_temporary_sky_slice_combined.npy'.format(joint_folder_path, cat_name))

    return list_of_arrays

def create_c_and_f(a_astro, b_astro, a_mag, b_mag, a_inds, a_size, b_inds, b_size, a_blen,
                   a_bins, b_bins, bright_frac, field_frac, a_flags, b_flags, a_astro_ind,
                   b_astro_ind, a_flen_ind, b_flen_ind, a_flags_ind, b_flags_ind, b_mag_ind, area):
    '''
    Functionality to create the photometric likelihood and priors from a set
    of photometric data in a given pair of filters.

    Parameters
    ----------
    a_astro : numpy.ndarray
        Array of astrometric parameters for all catalogue "a" sources in this
        given sky slice.
    b_astro : numpy.ndarray
        Astrometric parameters for small sky region catalogue "b" objects.
    a_mag : numpy.ndarray
        Catalogue "a" magnitudes for sky area.
    b_mag : numpy.ndarray
        Catalogue "b" magnitudes for sources in sky region.
    a_inds : numpy.ndarray
        Indices into catalogue "b" for each "a" object, indicating potential
        overlaps in counterparts.
    a_size : numpy.ndarray
        The number of potential overlapping sources in catalogue "b", for each
        catalogue "a" object.
    b_inds : numpy.ndarray
        Overlap indices into catalogue "a" for each catalogue "b" object.
    b_size : numpy.ndarray
        Number of overlaps from catalogue "b" into catalogue "a".
    a_blen : numpy.ndarray
        The "bright" error circle radius, integrating the joint AUF convolution
        out to ``expected_frac`` for largest "a"-"b" potential pairing, for each
        catalogue "a" object.
    a_bins : numpy.ndarray
        Array containing the magnitude bins into which to place catalogue "a"
        sources.
    b_bins : numpy.ndarray
        Array containing the magnitude bins into which to place catalogue "b"
        sources.
    bright_frac : float
        Fraction of total probability integral to consider potential counterparts
        out to, when considering potential overlaps between catalogue-catalogue
        pairings.
    field_frac : float
        Fraction of total probability integral out to which sources are removed
        from consideration as "field" sources (i.e., they are not assumed to be
        ruled out as potential counterparts), when considering potential overlaps
        in pairs of objects between the two catalogues.
    a_flags : numpy.ndarray
        Boolean flags for whether a source in catalogue "a" has a detected
        magnitude in ``a_mag``.
    b_flags : numpy.ndarray
        Detection flags for catalogue "b" sources in ``b_mag``.
    a_astro_ind : numpy.ndarray
        Astrometric information for all catalogue "a" sources with at least one
        overlap with catalogue "b" sources from ``b_astro``.
    b_astro_ind : numpy.ndarray
        Astrometric information for all catalogue "b" sources with an overlap
        with at least one catalogue "a" source in ``a_astro``.
    a_flen_ind : numpy.ndarray
        Largest joint AUF integral error circle radius for the "field" source
        integral fraction, for each catalogue "a" object in ``a_astro_ind``.
    b_flen_ind : numpy.ndarray
        Maximum AUF integral distance for each source in ``b_astro_ind``.
    a_flags_ind : numpy.ndarray
        Boolean flags for catalogue "a" sources which have an overlap with any
        catalogue "b" objects that are in this sky region.
    b_flags_ind : numpy.ndarray
        Boolean detection flags for all catalogue "b" sources for which any
        sources in catalogue "a" have a potential overlap in this sky region.
    b_mag_ind : numpy.ndarray
        All source magnitudes for which the catalogue "a" subset of sources
        have an overlap in catalogue "b".
    area : float
        Area of sky region for which photometric likelihood and prior are
        being calculated, in square degrees.

    Returns
    -------
    Nc : float
        The prior density of counterpart sources between the catalogues.
    cdmdm : numpy.ndarray
        Two-dimensional array of the photometric likelihood of counterpart between
        the two catalogues.
    Nfa : float
        So-called "field" source density in catalogue "a".
    fa : numpy.ndarray
        Probability density array of field sources for catalogue "a".
    Nfb : float
        Field source density prior for catalogue "b".
    fb : numpy.ndarray
        Field source PDF for catalogue "b".
    '''
    a_hist, a_bins = np.histogram(a_mag[a_flags], bins=a_bins)
    pa = a_hist/(np.sum(a_hist)*np.diff(a_bins))

    a_cuts = plf.find_mag_bin_inds(a_mag, a_flags, a_bins)

    # get_field_dists allows for magnitude slicing, to get f(m | m) instead of f(m),
    # but when we do want f(m) we just pass two impossible magnitudes as the limits.
    a_mask, a_area = plf.get_field_dists(a_astro[:, 0], a_astro[:, 1], b_astro_ind[:, 0],
                                         b_astro_ind[:, 1], a_inds, a_size, b_flen_ind, a_flags,
                                         b_flags_ind, b_mag, -999, 999)
    b_mask, b_area = plf.get_field_dists(b_astro[:, 0], b_astro[:, 1], a_astro_ind[:, 0],
                                         a_astro_ind[:, 1], b_inds, b_size, a_flen_ind, b_flags,
                                         a_flags_ind, a_mag, -999, 999)
    a_mask = a_mask.astype(bool)
    b_mask = b_mask.astype(bool)
    a_left = a_mag[a_mask]
    b_left = b_mag[b_mask]
    hist, a_bins = np.histogram(a_left, bins=a_bins)
    Num_fa = np.sum(a_mask)

    fa = hist / (np.sum(hist)*np.diff(a_bins))

    hist, b_bins = np.histogram(b_left, bins=b_bins)
    Num_fb = np.sum(b_mask)

    fb = hist / (np.sum(hist)*np.diff(b_bins))

    Nfa = Num_fa/(area - a_area)
    Nfb = Num_fb/(area - b_area)

    bm = np.empty((len(b_bins)-1, len(a_bins)-1), float, order='F')
    z = np.empty(len(a_bins)-1, float)

    mag_mask, aa = plf.brightest_mag(a_astro[:, 0], a_astro[:, 1], b_astro_ind[:, 0],
                                     b_astro_ind[:, 1], a_mag, b_mag_ind, a_inds, a_size, a_blen,
                                     a_flags, b_flags_ind, a_bins)
    mag_mask = mag_mask.astype(bool)
    for i in range(0, len(a_bins)-1):
        hist, b_bins = np.histogram(b_mag_ind[mag_mask[:, i]], bins=b_bins)
        q = np.sum(mag_mask[:, i])
        if q > 0:
            bm[:, i] = hist/(np.diff(b_bins)*q)
        else:
            bm[:, i] = 0
        z[i] = np.sum(hist)/np.sum(a_cuts[i])
    cdmdm = np.empty((len(b_bins)-1, len(a_bins)-1), float, order='F')
    for i in range(0, len(a_bins)-1):
        bmask, barea = plf.get_field_dists(
            b_astro[:, 0], b_astro[:, 1], a_astro[:, 0], a_astro[:, 1], b_inds, b_size, a_flen_ind,
            b_flags, a_flags_ind, a_mag, a_bins[i], a_bins[i+1])
        bmask = bmask.astype(bool)
        b_left = b_mag[bmask]
        hist, b_bins = np.histogram(b_left, bins=b_bins)
        _Num_fb = np.sum(b_mask)

        _fb = hist / (np.sum(hist)*np.diff(b_bins))
        _Nfb = _Num_fb/(area - barea)
        Fm = np.append(0, np.cumsum(_fb[:-1] * np.diff(b_bins[:-1])))
        for j in range(0, len(b_bins)-1):
            Cm = np.sum(cdmdm[:j, i]*np.diff(b_bins[:j+1]))
            cdmdm[j, i] = max(0, z[i]*bm[j, i]*np.exp(aa[i]*_Nfb*Fm[j]) - (1-Cm)*aa[i]*_Nfb*_fb[j])

    zc = np.sum(cdmdm*np.diff(b_bins).reshape(-1, 1), axis=0)
    frac = zc/bright_frac
    density_of_inputs = np.sum(a_cuts, axis=1)/area
    Nc = np.sum(frac*density_of_inputs)

    integral = 0
    for i in range(0, len(a_bins)-1):
        cdmdm[:, i] *= pa[i] / (a_bins[i+1] - a_bins[i])
        integral = integral + np.sum(cdmdm[:, i]*np.diff(b_bins))*(a_bins[i+1] - a_bins[i])
    if integral > 0:
        cdmdm /= integral

    # Correct the field priors for the fraction of counterparts that get left
    # in their "cutout" circle, by the fact that we don't use the entire integral:
    Nfa = Nfa - (1 - field_frac)*Nc
    Nfb = Nfb - (1 - field_frac)*Nc

    return Nc, cdmdm, Nfa, fa, Nfb, fb


# Lint as python3
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""A tree search ideal observer for alchemy."""

import collections
import copy
import math
from typing import Any, Counter, List, Mapping, MutableMapping, Sequence, Tuple

from dm_alchemy.ideal_observer import helpers
from dm_alchemy.ideal_observer import precomputed_maps
from dm_alchemy.types import graphs
from dm_alchemy.types import helpers as types_helpers
from dm_alchemy.types import stones_and_potions
import numpy as np

# Alias these for readability
AlignedStone = stones_and_potions.AlignedStone
PerceivedStone = stones_and_potions.PerceivedStone
PerceivedPotion = stones_and_potions.PerceivedPotion
LatentStone = stones_and_potions.LatentStone
LatentPotion = stones_and_potions.LatentPotion
StoneMap = stones_and_potions.StoneMap
PotionMap = stones_and_potions.PotionMap
PartialStoneMap = stones_and_potions.PartialStoneMap
PartialPotionMap = stones_and_potions.PartialPotionMap
PartialGraph = graphs.PartialGraph

PrecomputedMaps = precomputed_maps.PrecomputedMaps

# We use indices in place of actual types for speed
AlignedStoneIndex = stones_and_potions.AlignedStoneIndex
PerceivedPotionIndex = stones_and_potions.PerceivedPotionIndex
LatentStoneIndex = stones_and_potions.LatentStoneIndex
LatentPotionIndex = stones_and_potions.LatentPotionIndex
StoneMapIndex = stones_and_potions.StoneMapIndex
PotionMapIndex = stones_and_potions.PotionMapIndex

Action = Tuple[AlignedStoneIndex, PerceivedPotionIndex]
ActionObjective = Tuple[Action, Any]
SearchResults = MutableMapping[int, ActionObjective]
ActionObjectiveAndSearchResults = Tuple[Action, Any, SearchResults]

END_TRIAL = types_helpers.END_TRIAL

def get_possible_stone_maps(
    stone_map_indices: Sequence[StoneMapIndex],
    aligned_stones: Sequence[AlignedStone]
) -> Tuple[List[int], List[StoneMapIndex]]:
  """Gets possible stone maps (and their indices) consistent with stones passed."""
  # This doesn't need to be especially fast as it just happens once per trial.
  still_in_s = []
  poss_sms = []
  for i, sm in enumerate(stone_map_indices):
    stone_map = stones_and_potions.stone_map_from_index(sm)
    if stone_map.consistent_with_stones(aligned_stones):
      still_in_s.append(i)
      poss_sms.append(sm)
  return still_in_s, poss_sms

class WorldStateDistribution:
  """Distribution over world states."""

  def __init__(
      self, stone_map_distr: Mapping[StoneMapIndex, float],
      potion_map_distr: Mapping[PotionMapIndex, float],
      precomputed: PrecomputedMaps
  ):

    self.stone_map_possible = sorted(stone_map_distr.keys())
    self.potion_map_possible = sorted(potion_map_distr.keys())
    self.partial_potion_map_index = (
        stones_and_potions.partial_potion_map_from_possibles(
            self.potion_map_possible, precomputed.index_to_perm_index).index(
                precomputed.perm_index_to_index))
    self.partial_stone_map_index = (
        stones_and_potions.partial_stone_map_from_possibles(
            self.stone_map_possible).index())

    # Use a bitfield like structure otherwise it takes forever to copy
    self.partial_graph_possible = helpers.list_to_bitfield(
        range(precomputed.graphs_list.shape[0]))
    self.partial_graph_index = (
        precomputed.partial_graph_index_to_possible_index[
            graphs.partial_graph_from_possibles(
                precomputed.graphs_list[self.get_possible_graphs()]).index()])

    self.poss_world_states = np.zeros((
        len(potion_map_distr), len(stone_map_distr),
        precomputed.graph_index_distr.shape[0]), dtype=np.float)
    for potion_map_index, p1 in enumerate(potion_map_distr.values()):
      for stone_map_index, p2 in enumerate(stone_map_distr.values()):
        for graph_index, p3 in enumerate(precomputed.graph_index_distr):
          p = p1 * p2 * p3
          self.poss_world_states[
              potion_map_index, stone_map_index, graph_index] = p

    self.observed_no_effect_bits = 0

  def new_trial(
      self, aligned_stone_indices: Counter[AlignedStoneIndex]
  ) -> None:
    """Updates the world state distribution given the stones perceived.

    The reward indicator on the stones allows us to limit the possible maps from
    stone space to latent space. For example if we see a stone with reward 3
    then it must be at [1, 1, 1] in latent space.

    Args:
      aligned_stone_indices: The stones seen in this trial.
    """
    # The stones seen when we start a new trial could eliminate stone map
    # possibilities if there are multiple.
    aligned_stones = [
        stones_and_potions.aligned_stone_from_index(aligned_stone_index)
        for aligned_stone_index in aligned_stone_indices]
    still_in_s, poss_sms = get_possible_stone_maps(
        self.stone_map_possible, aligned_stones)
    assert still_in_s, 'Stones seen in trial are impossible.'

    self.stone_map_possible = poss_sms
    self.partial_stone_map_index = (
        stones_and_potions.partial_stone_map_from_possibles(
            self.stone_map_possible).index())
    self.poss_world_states = self.poss_world_states[:, still_in_s, :]
    total_prob = self.poss_world_states.sum()
    self.poss_world_states /= total_prob

  def get_possible_graphs(self) -> List[int]:
    return helpers.bitfield_to_list(self.partial_graph_possible)

  def potions_equivalent(
      self, p1: PerceivedPotionIndex, p2: PerceivedPotionIndex,
      s: AlignedStoneIndex, precomputed: PrecomputedMaps
  ) -> bool:
    """If the potions effect on the stone are equivalent in this belief state.

    This is the case if we have the same knowledge about the potions effect and
    the same number of these potions and their counterparts on the same
    dimension. In this case the calculation of the expected reward must be
    exactly the same. This does not imply that the actual effect of the potions
    will be the same.

    For example, if we have no knowledge of the perceptual mapping and graph and
    we have one red potion and one green potion then the calculation will
    include terms for the probability that the red potion acts on each of the
    directed edges in latent space. The calculation for the green potion would
    have all of the same terms.

    Args:
      p1: The first potion.
      p2: The second potion.
      s: The stone they will be applied to.
      precomputed: Precomputed maps used for speed.

    Returns:
      True if they are equivalent.
    """
    latent_dims = precomputed.possible_latent_dims[
        p1, self.partial_potion_map_index[0]]
    latent_dims2 = precomputed.possible_latent_dims[
        p2, self.partial_potion_map_index[0]]
    if latent_dims != latent_dims2:
      return False
    for latent_dim in latent_dims:
      this_could_stay_still1, latent_dirs1 = precomputed.possible_latent_dirs[
          self.partial_potion_map_index[1], self.partial_stone_map_index,
          latent_dim, p1, s]
      this_could_stay_still2, latent_dirs2 = precomputed.possible_latent_dirs[
          self.partial_potion_map_index[1], self.partial_stone_map_index,
          latent_dim, p2, s]
      if this_could_stay_still1 != this_could_stay_still2:
        return False
      if latent_dirs1 != latent_dirs2:
        return False
    return True

  def update_possible(
      self, stone_index: AlignedStoneIndex,
      potion_index: PerceivedPotionIndex,
      result_index: AlignedStoneIndex,
      precomputed: PrecomputedMaps
  ) -> None:
    """Updates which possibilities are consistent with this observation.

    Args:
      stone_index: The initial stone.
      potion_index: The potion applied to the stone.
      result_index: The resulting stone.
      precomputed: Precomputed maps used for speed.
    """
    # Update which potion maps are possible
    poss_p = precomputed.poss_p_maps[stone_index, potion_index, result_index]
    if poss_p is not None:
      self.partial_potion_map_index = precomputed.partial_potion_map_update[
          stone_index, potion_index, result_index,
          self.partial_potion_map_index[0], self.partial_potion_map_index[1]]
      self.potion_map_possible, still_in_p = helpers.sorted_intersection(
          self.potion_map_possible, poss_p)
      self.poss_world_states = self.poss_world_states[still_in_p, :, :]

    # Update which stone maps are possible
    poss_s = precomputed.poss_s_maps[stone_index, potion_index, result_index]
    if poss_s is not None:
      self.partial_stone_map_index = precomputed.partial_stone_map_update[
          stone_index, result_index, self.partial_stone_map_index]
      self.stone_map_possible, still_in_s = helpers.sorted_intersection(
          self.stone_map_possible, poss_s)
      self.poss_world_states = self.poss_world_states[:, still_in_s, :]

    # Update which graphs are possible
    update_graphs_possible = False
    if stone_index != result_index:
      update_graphs_possible = True
      self.partial_graph_index = precomputed.partial_graph_update[
          precomputed.drop_reward[stone_index],
          precomputed.drop_reward[result_index]][self.partial_graph_index]

    if not update_graphs_possible:
      missing_edge = precomputed.missing_edge_no_change[
          self.partial_stone_map_index, self.partial_potion_map_index[0],
          self.partial_potion_map_index[1], potion_index,
          precomputed.drop_reward[stone_index]]
      if missing_edge != -1:
        update_graphs_possible = True
        self.partial_graph_index = precomputed.update_partial_graph_no_change[
            self.partial_graph_index, missing_edge]

    if update_graphs_possible:
      new_graphs_possible = precomputed.partial_graph_to_matching_graphs[
          self.partial_graph_index]
      remaining_graphs_possible = (self.partial_graph_possible &
                                   new_graphs_possible)
      # Work out the position of the eliminated slices in poss_world_states.
      still_in_g = []
      ind = 0
      for i in range(precomputed.graphs_list.shape[0]):
        check = 1 << i
        poss_check = self.partial_graph_possible & check
        if remaining_graphs_possible & check and poss_check:
          still_in_g.append(ind)
        if poss_check:
          ind += 1
      self.partial_graph_possible = remaining_graphs_possible

      self.poss_world_states = self.poss_world_states[:, :, still_in_g]

    # If stone map is known then get info about which actions will have no
    # effect because they take the stone out of the latent cube
    stone_map_index = precomputed.partial_stone_map_to_stone_map[
        self.partial_stone_map_index]
    if stone_map_index != -1:
      self.observed_no_effect_bits |= precomputed.no_effect_from_partial_chem[
          stone_map_index, self.partial_potion_map_index[0],
          self.partial_potion_map_index[1]]

  def possible_outcomes(
      self, perceived_potion_index: PerceivedPotionIndex,
      aligned_stone_index: AlignedStoneIndex, precomputed: PrecomputedMaps
  ) -> Tuple[List[AlignedStoneIndex], bool]:
    """Gets a list of outcomes we could see applying this potion to this stone.

    Args:
      perceived_potion_index: The potion we apply.
      aligned_stone_index: The stone we apply it to.
      precomputed: Precomputed maps used for speed.

    Returns:
      A list of possible outcomes and a boolean saying whether one of them is
      the stone remaining the same.
    """
    outcomes = []
    could_stay_still = False
    for latent_dim in precomputed.possible_latent_dims[
        perceived_potion_index, self.partial_potion_map_index[0]]:
      # latent_dir may not be possible if the reward for the stone is already at
      # max or min. If you know stone position in latent space on latent_dim
      # then latent_dir can only be the opposite so you don't need to consider
      # the reward going the other way even if you don't know what direction the
      # potion acts in.
      this_could_stay_still, latent_dirs = precomputed.possible_latent_dirs[
          self.partial_potion_map_index[1], self.partial_stone_map_index,
          latent_dim, perceived_potion_index, aligned_stone_index]
      # Could stay still due to going outside the cube.
      could_stay_still |= this_could_stay_still
      for latent_dir in latent_dirs:
        result = precomputed.react_result[
            aligned_stone_index, latent_dim, (latent_dir + 1) // 2,
            self.partial_graph_index]
        if result != helpers.IMPOSSIBLE:
          outcomes.append(result)
      if latent_dirs:
        # Could stay still due to edge not existing
        this_edge_exists = precomputed.edge_exists[
            self.partial_graph_index, precomputed.drop_reward[
                aligned_stone_index], latent_dim]
        # If either we know the edge isn't there or we are not sure if the edge
        # is there then could stay still.
        if this_edge_exists != graphs.KNOWN_EDGE:
          could_stay_still = True

    if could_stay_still:
      outcomes.append(aligned_stone_index)
    return outcomes, could_stay_still

  def action_and_outcome(
      self, stone_index: AlignedStoneIndex,
      potion_index: PerceivedPotionIndex, result_index: AlignedStoneIndex,
      precomputed: PrecomputedMaps, bit_mask: int
  ) -> float:
    """Updates the world state distribution given we saw this observation."""

    # Eliminate whole slices of the world state distribution if possible given
    # the new information.
    self.update_possible(stone_index, potion_index, result_index, precomputed)

    # If the stone changed as a result of applying the potion then all
    # information gained removes whole slices otherwise we must remove
    # combinations of potion map, stone map and graph.
    if stone_index == result_index:
      for potion_i, potion_map_index in enumerate(self.potion_map_possible):
        for stone_i, stone_map_index in enumerate(self.stone_map_possible):
          graphs_bitfield = precomputed.graphs_with_edge[
              stone_map_index, potion_map_index, stone_index, potion_index]

          # If there are no graphs with an edge between the stone and result
          # then continue.
          if graphs_bitfield == 0:
            continue

          # Graphs in the list are not possible because they contain the edge so
          # the stone should have changed but didn't.
          not_possible = []
          ind = 0
          for i in range(precomputed.graphs_list.shape[0]):
            check = 1 << i
            poss_check = self.partial_graph_possible & check
            if graphs_bitfield & check and poss_check:
              not_possible.append(ind)
            if poss_check:
              ind += 1

          self.poss_world_states[potion_i, stone_i, not_possible] = 0.0
      self.observed_no_effect_bits |= bit_mask

    # Re-normalise
    total_prob = self.poss_world_states.sum()
    if total_prob > 0.0:
      self.poss_world_states /= total_prob
    return total_prob

  def __len__(self):
    return len(self.poss_world_states)

  def update_stone_map(self, new_to_old: StoneMap) -> None:
    """If we assumed the wrong rotation we may need to swap stone map dims."""
    # Change partial stone map.
    partial_stone_map = stones_and_potions.partial_stone_map_from_index(
        self.partial_stone_map_index)
    # If the partial stone map is not completely known at this point then it is
    # completely unknown since any 2 bits of information would be enough to
    # completely determine the rotation.
    if any(c == types_helpers.UNKNOWN
           for c in partial_stone_map.latent_pos_dir):
      assert all(c == types_helpers.UNKNOWN
                 for c in partial_stone_map.latent_pos_dir)
    else:
      partial_stone_map.chain(new_to_old)
    self.partial_stone_map_index = partial_stone_map.index()
    # Change poss stone maps.
    old_stone_map_possible = copy.deepcopy(self.stone_map_possible)
    old_stone_map_to_new_stone_map = {}
    for stone_map_index in self.stone_map_possible:
      stone_map = stones_and_potions.stone_map_from_index(stone_map_index)
      stone_map.chain(new_to_old)
      old_stone_map_to_new_stone_map[stone_map_index] = stone_map.index()
    self.stone_map_possible = sorted(
        old_stone_map_to_new_stone_map[stone_map_index]
        for stone_map_index in self.stone_map_possible)
    new_stone_map_to_index = {
        stone_map_index: i
        for i, stone_map_index in enumerate(self.stone_map_possible)}
    old_index_to_new_index = [
        new_stone_map_to_index[old_stone_map_to_new_stone_map[stone_map]]
        for stone_map in old_stone_map_possible]
    # Change poss world states.
    old_poss_world_states = copy.deepcopy(self.poss_world_states)
    for old_index, new_index in enumerate(old_index_to_new_index):
      self.poss_world_states[new_index] = old_poss_world_states[old_index]
    # Change observed no effect.
    old_observed_no_effect_bits = copy.deepcopy(self.observed_no_effect_bits)
    self.observed_no_effect_bits = 0
    for old_index in range(stones_and_potions.AlignedStone.num_dir_assignments):
      aligned_stone = stones_and_potions.aligned_stone_from_index(
          AlignedStoneIndex(old_index))
      new_index = new_to_old.apply(aligned_stone).index()
      for potion_index in range(stones_and_potions.PerceivedPotion.num_types):
        old_mask = 1 << (old_index * PerceivedPotion.num_types) + potion_index
        masked = old_observed_no_effect_bits & old_mask
        if masked:
          new_mask = 1 << (new_index * PerceivedPotion.num_types) + potion_index
          self.observed_no_effect_bits |= new_mask

def init_world_state_distribution(
    precomputed: PrecomputedMaps
) -> WorldStateDistribution:
  """Creates an initial world state distribution from observed stones."""
  # Initialise the ideal observer based on the stones and potions you can see
  return WorldStateDistribution(
      stones_and_potions.stone_map_distr(precomputed.stone_maps),
      stones_and_potions.potion_map_distr(precomputed.potion_maps),
      precomputed)

def stone_potion_bit_mask(
    stone_index: AlignedStoneIndex, potion_index: PerceivedPotionIndex,
    precomputed: PrecomputedMaps
) -> int:
  """Returns a mask for the bit representing a stone potion pair."""
  stone_part = precomputed.drop_reward[stone_index] * PerceivedPotion.num_types
  return 1 << (stone_part + potion_index)

class BeliefState:
  """Belief the ideal observer has about stones, potions, world and reward.

  The belief state consists of a set of perceived stones, a set of perceived
  potions, a distribution over world states, and a reward so far.
  """

  possible_partial_graph_num_bits = None

  def __init__(self, precomputed: PrecomputedMaps):

    # These should be set by calling new_trial
    self.aligned_stones: Counter[AlignedStoneIndex] = collections.Counter()
    self.perceived_potions: Counter[PerceivedPotionIndex] = (
        collections.Counter())
    self.world_state_distribution = init_world_state_distribution(precomputed)

  def representative_potions(
      self, stone_index: AlignedStoneIndex, precomputed: PrecomputedMaps
  ) -> List[PerceivedPotionIndex]:
    """Gets a representative set of potions for this stone and belief state.

    Some potions will be equivalent if we don't know what they do and we have
    the same number of them and their counterparts on the same perceptual
    dimension. For each equivalence set we return one potion as a representative
    of the set.

    Args:
      stone_index: The stone to apply potions to.
      precomputed: Precomputed maps used for speed.

    Returns:
      A representative set of potions.
    """
    potion_to_count = [0 for _ in range(PerceivedPotion.num_types)]
    for p1 in self.perceived_potions:
      p2 = precomputed.potion_to_pair[p1]
      c1 = self.perceived_potions[p1]
      c2 = self.perceived_potions[p2]
      potion_to_count[p1] = (c1, c2)
    representative_potions = []
    for p2 in self.perceived_potions:
      equiv = False
      for p1 in representative_potions:
        b1 = (self.world_state_distribution.observed_no_effect_bits &
              precomputed.potion_masks[p1]) >> p1
        b2 = (self.world_state_distribution.observed_no_effect_bits &
              precomputed.potion_masks[p2]) >> p2
        if (potion_to_count[p1] == potion_to_count[p2] and
            (b1 == b2) and
            self.world_state_distribution.potions_equivalent(
                p1, p2, stone_index, precomputed)):
          equiv = True
          break
      if not equiv:
        representative_potions.append(p2)
    return representative_potions

  def _remove_stone(self, stone_index: AlignedStoneIndex) -> None:
    if stone_index in self.aligned_stones:
      self.aligned_stones[stone_index] -= 1
    if self.aligned_stones[stone_index] == 0:
      del self.aligned_stones[stone_index]

  def _add_stone(self, stone_index: AlignedStoneIndex) -> None:
    self.aligned_stones.update([stone_index])

  def _remove_potion(self, potion_index: PerceivedPotionIndex) -> None:
    self.perceived_potions.subtract([potion_index])
    if self.perceived_potions[potion_index] == 0:
      del self.perceived_potions[potion_index]

  def possible_actions(
      self, precomputed: PrecomputedMaps
  ) -> List[Tuple[AlignedStoneIndex, PerceivedPotionIndex]]:
    """Gets representative list of possible actions which have an effect."""
    # Use -1, -1 to represent ending and putting all stones in the cauldron or
    # throwing them away. If we consider this action first then in the event
    # that it has the same expected reward as using a potion we will take this
    # action instead. This gives more intuitive behaviour, for example if we
    # have the best stone we won't transform it to something less good and then
    # transform it back.
    poss_actions = [(AlignedStoneIndex(END_TRIAL),
                     PerceivedPotionIndex(END_TRIAL))]
    for s in self.aligned_stones:
      # Don't consider potions if we have observed that they have no effect.
      potions = [p for p in self.representative_potions(s, precomputed)
                 if not (self.world_state_distribution.observed_no_effect_bits &
                         stone_potion_bit_mask(s, p, precomputed))]
      poss_actions.extend([(s, p) for p in potions])

    return poss_actions

  def use_potion(
      self, stone_index: AlignedStoneIndex,
      potion_index: PerceivedPotionIndex,
      result_index: AlignedStoneIndex
  ) -> None:
    """Uses the potion on the current stone.

    Args:
      stone_index: The stone used in the potion.
      potion_index: The potion applied to the stone.
      result_index: The result observed.
    """
    # Remove the used potion
    self._remove_potion(potion_index)
    # If the stone has not changed then we don't need to do anything else.
    if stone_index == result_index:
      return
    # Remove the initial stone and replace with the result
    self._remove_stone(stone_index)
    self._add_stone(result_index)

  def action_and_outcome(
      self, stone_index: AlignedStoneIndex,
      potion_index: PerceivedPotionIndex, result_index: AlignedStoneIndex,
      precomputed: PrecomputedMaps, bit_mask: int
  ) -> float:
    """Updates the belief state given the action and observation.

    Args:
      stone_index: The stone used in the potion.
      potion_index: The potion in which the stone is used.
      result_index: The resulting stone.
      precomputed: Precomputed maps used for speed.
      bit_mask: Mask on observed no effect for the stone and potion passed.

    Returns:
      The probability given our prior belief state of this observation.
    """
    self.use_potion(stone_index, potion_index, result_index)

    # Update the world state distribution
    total_prob = self.world_state_distribution.action_and_outcome(
        stone_index, potion_index, result_index, precomputed, bit_mask)

    return total_prob

  def new_trial(
      self, aligned_stones: Counter[AlignedStoneIndex],
      perceived_potions: Counter[PerceivedPotionIndex]
  ) -> None:
    self.aligned_stones = copy.deepcopy(aligned_stones)
    self.perceived_potions = copy.deepcopy(perceived_potions)
    self.world_state_distribution.new_trial(aligned_stones)

  def to_bitfield(self) -> int:
    """Converts to a bitfield to cache results."""

    def perceived_potions_to_bits(
        perceived_potions: Mapping[PerceivedPotionIndex, int]
    ) -> Tuple[int, int]:
      """Converts the set of perceived potions to a bitfield."""
      local_int_rep = 0
      for potion_type, count in perceived_potions.items():
        local_int_rep |= (count << (
            PerceivedPotion.count_num_bits * potion_type))
        if count > PerceivedPotion.max_present:
          raise ValueError('Too many potions present.')
        if potion_type >= PerceivedPotion.num_types:
          raise ValueError('Invalid potion type.')
      return local_int_rep, (
          PerceivedPotion.num_types * PerceivedPotion.count_num_bits)

    def aligned_stones_to_bits(
        aligned_stones: Mapping[AlignedStoneIndex, int]
    ) -> Tuple[int, int]:
      """Converts the set of perceived stones to a bitfield."""
      local_int_rep = 0
      stone_number = 0
      for stone_type, count in sorted(aligned_stones.items()):
        for _ in range(count):
          local_int_rep |= (stone_type << (
              AlignedStone.num_bits * stone_number))
          stone_number += 1
        if stone_type >= stones_and_potions.AlignedStone.num_types:
          raise ValueError('Invalid stone type')
      if stone_number > AlignedStone.max_present:
        raise ValueError('Too many stones present.')
      return local_int_rep, AlignedStone.max_present * AlignedStone.num_bits

    all_things = [
        perceived_potions_to_bits(self.perceived_potions),
        aligned_stones_to_bits(self.aligned_stones),
        (self.world_state_distribution.observed_no_effect_bits,
         LatentStone.num_types * PerceivedPotion.num_types),
        (self.world_state_distribution.partial_potion_map_index[0],
         PartialPotionMap.num_bits_axis),
        (self.world_state_distribution.partial_potion_map_index[1],
         PartialPotionMap.num_bits_dir),
        (self.world_state_distribution.partial_stone_map_index,
         PartialStoneMap.num_bits),
        (self.world_state_distribution.partial_graph_index,
         BeliefState.possible_partial_graph_num_bits)
    ]

    return helpers.pack_to_bitfield(all_things)

  def __repr__(self) -> str:
    # Convert the observed_no_effect bitfield to a matrix before printing.
    observed_no_effect = np.zeros(
        (stones_and_potions.LatentStone.num_types,
         stones_and_potions.PerceivedPotion.num_types))
    for pe_st in range(stones_and_potions.LatentStone.num_types):
      for pe_po in range(stones_and_potions.PerceivedPotion.num_types):
        bit_num = (pe_st * stones_and_potions.PerceivedPotion.num_types) + pe_po
        observed_no_effect[pe_st, pe_po] = (
            self.world_state_distribution.observed_no_effect_bits &
            (1 << bit_num))

    return (
        'BeliefState(aligned_stones={aligned_stones}, '
        'perceived_potions={perceived_potions}, '
        'observed_no_effect={observed_no_effect}, '
        'poss_world_states={poss_world_states}, '
        'partial_potion_map_index={partial_potion_map_index}, '
        'partial_stone_map_index={partial_stone_map_index}, '
        'partial_graph_index={partial_graph_index}, '
        'partial_graph_possible={partial_graph_possible}, '
        'stone_map_possible={stone_map_possible}, '
        'potion_map_possible={potion_map_possible}, '
        'num_world_states = {num_world_states}, '.format(
            aligned_stones=self.aligned_stones,
            perceived_potions=self.perceived_potions,
            observed_no_effect=types_helpers.str_np_array_construct(
                observed_no_effect),
            poss_world_states=types_helpers.str_np_array_construct(
                self.world_state_distribution.poss_world_states),
            partial_potion_map_index=(
                self.world_state_distribution.partial_potion_map_index),
            partial_stone_map_index=(
                self.world_state_distribution.partial_stone_map_index),
            partial_graph_index=(
                self.world_state_distribution.partial_graph_index),
            partial_graph_possible=(
                self.world_state_distribution.partial_graph_possible),
            stone_map_possible=self.world_state_distribution.stone_map_possible,
            potion_map_possible=(
                self.world_state_distribution.potion_map_possible),
            num_world_states=len(self.world_state_distribution)))

  def update_stone_map(
      self, new_to_old: StoneMap
  ) -> None:
    """If we assumed the wrong rotation we may need to swap stone map dims."""
    # Change poss stone maps.
    old_aligned_stones = {
        stones_and_potions.aligned_stone_from_index(stone): count
        for stone, count in self.aligned_stones.items()}
    new_aligned_stones = collections.Counter({
        AlignedStone(stone.reward, new_to_old.apply(
            stone).latent_coords).index(): count
        for stone, count in old_aligned_stones.items()})
    self.aligned_stones = new_aligned_stones
    self.world_state_distribution.update_stone_map(new_to_old)

  @property
  def num_world_states(self) -> int:
    return np.where(self.world_state_distribution.poss_world_states)[0].size

  @property
  def num_potion_maps(self) -> int:
    return len(self.world_state_distribution.potion_map_possible)

  @property
  def num_stone_maps(self) -> int:
    return len(self.world_state_distribution.stone_map_possible)

  @property
  def num_graphs(self) -> int:
    return len(self.world_state_distribution.get_possible_graphs())

  def partial_potion_map(
      self, index_to_perm_index: np.ndarray
  ) -> PartialPotionMap:
    return stones_and_potions.partial_potion_map_from_index(
        self.world_state_distribution.partial_potion_map_index,
        index_to_perm_index)

  def partial_stone_map(self) -> PartialStoneMap:
    return stones_and_potions.partial_stone_map_from_index(
        self.world_state_distribution.partial_stone_map_index)

  def partial_graph(
      self, possible_partial_graph_indices: np.ndarray
  ) -> PartialGraph:
    return graphs.partial_graph_from_index(
        possible_partial_graph_indices[
            self.world_state_distribution.partial_graph_index])

class BeliefStateWithRotation:
  """Belief state over chem including rotations."""

  def __init__(self, precomputed: precomputed_maps.PrecomputedMaps):
    self.belief_state = BeliefState(precomputed)
    self.possible_rotations = stones_and_potions.possible_rotations()
    self.rotation = None
    # We need to know 1 stone which is consistent with the selected rotation.
    self._observed_stone = None
    self._rotation_to_angles = (
        lambda rotation: tuple(stones_and_potions.rotation_to_angles(rotation)))
    stone_map_indices = [
        sm.index() for sm in stones_and_potions.possible_stone_maps()]
    self._stone_maps_for_rotation = {
        self._rotation_to_angles(rotation): copy.deepcopy(stone_map_indices)
        for rotation in stones_and_potions.possible_rotations()}

  def _update_given_stones(
      self, perceived_stones: Sequence[PerceivedStone]
  ) -> None:
    """Updates the possible rotations and belief state given observed stones."""
    if not perceived_stones:
      raise ValueError(
          'Must pass perceived stones to update possible rotations.')

    # Given the stones we see can we eliminate some possible rotations.
    valid_rotations = []
    for rotation in self.possible_rotations:
      # For a rotation to be possible all stones have to go to corners of the
      # cube and the change in latent variables has to be consistent with the
      # change in reward (i.e. at least one stone map gives the observed
      # rewards).
      aligned_stones = []
      rotation_valid = True
      for stone in perceived_stones:
        valid, coords = stones_and_potions.aligns(stone, rotation)
        if valid:
          aligned_stones.append(stones_and_potions.aligned_stone_from_coords(
              coords, stone.reward))
        else:
          rotation_valid = False
          break

      if rotation_valid:
        stone_maps = self._stone_maps_for_rotation[self._rotation_to_angles(
            rotation)]
        _, possible_stone_maps = get_possible_stone_maps(
            stone_maps, aligned_stones)
        self._stone_maps_for_rotation[self._rotation_to_angles(
            rotation)] = possible_stone_maps
        if possible_stone_maps:
          valid_rotations.append(rotation)
    assert valid_rotations, 'No rotation is valid.'
    self.possible_rotations = valid_rotations
    if self.rotation is None:
      self.rotation = self.possible_rotations[0]
      self._observed_stone = stones_and_potions.align(
          perceived_stones[0], self.rotation)
    elif not stones_and_potions.rotations_equal(
        self.rotation, self.possible_rotations[0]):
      new_to_old = stones_and_potions.get_new_mapping_to_old_mapping(
          self.rotation, self.possible_rotations[0], self._observed_stone)
      self.belief_state.update_stone_map(new_to_old)
      self.rotation = self.possible_rotations[0]
      self._observed_stone = stones_and_potions.align(
          perceived_stones[0], self.rotation)

  def new_trial(
      self, perceived_stones: Counter[PerceivedStone],
      perceived_potions: Counter[PerceivedPotion]
  ) -> None:
    """Updates belief state given that new trial has started."""
    self._update_given_stones(list(perceived_stones.keys()))
    aligned_stones = collections.Counter(
        {stones_and_potions.align(stone, self.rotation).index(): count
         for stone, count in perceived_stones.items()})
    perceived_potion_indices = collections.Counter(
        {potion.index(): count for potion, count in perceived_potions.items()})
    self.belief_state.new_trial(aligned_stones, perceived_potion_indices)

  def action_and_outcome(
      self, stone: PerceivedStone, potion: PerceivedPotion,
      result: PerceivedStone, precomputed: PrecomputedMaps
  ) -> float:
    self._update_given_stones([result])
    stone_index = stones_and_potions.align(stone, self.rotation).index()
    result_index = stones_and_potions.align(result, self.rotation).index()
    bit_mask = stone_potion_bit_mask(stone_index, potion.index(), precomputed)
    return self.belief_state.action_and_outcome(
        stone_index, potion.index(), result_index, precomputed, bit_mask)

  @property
  def num_world_states(self) -> int:
    return self.belief_state.num_world_states

  @property
  def num_potion_maps(self) -> int:
    return self.belief_state.num_potion_maps

  @property
  def num_stone_maps(self) -> int:
    return self.belief_state.num_stone_maps

  @property
  def num_graphs(self) -> int:
    return self.belief_state.num_graphs

  def partial_potion_map(
      self, index_to_perm_index: np.ndarray
  ) -> PartialPotionMap:
    return self.belief_state.partial_potion_map(index_to_perm_index)

  def partial_stone_map(self) -> PartialStoneMap:
    return self.belief_state.partial_stone_map()

  def partial_graph(
      self, possible_partial_graph_indices: np.ndarray
  ) -> PartialGraph:
    return self.belief_state.partial_graph(possible_partial_graph_indices)

def search(
    belief_state: BeliefState,
    search_results: SearchResults,
    bonus: int, precomputed: PrecomputedMaps,
    depth: int = 0,
    minimise_world_states: bool = False
) -> ActionObjective:
  """Searches iteratively over actions and outcomes to find expected reward.

  Conducts a depth first search over the DAG of available actions and the
  possible outcomes. The reward for a latent stone is assumed to be the sum of
  the latent values plus the passed bonus if all latent values are positive. We
  do not deal with reward functions with arbitrary coefficient vectors and
  offsets.

  Args:
    belief_state: The current belief state.
    search_results: A cache of previously computed results mapping the belief
      state as a bitfield to the best action and expected reward.
    bonus: The extra reward we get by reaching the stone of reward 3.
    precomputed: Precomputed maps used for speed.
    depth: Number of actions taken in this search to reach this belief state.
    minimise_world_states: Let the objective be to minimise the number of world
      states at the end of the trial instead of to maximise the accumulated
      reward. Actions selected will not necessarily produce reward but will
      narrow down the possible chemistries, e.g. given a maximum value stone on
      the first trial and a potion we would use the potion to find out the
      effect even though it could reduce the value of the stone.

  Returns:
    The best action and maximum expected reward achievable from this state.
  """

  # If we have searched from this belief state before return the cached result.
  belief_state_bitfield = belief_state.to_bitfield()
  if belief_state_bitfield in search_results:
    return search_results[belief_state_bitfield]

  # For all possible actions consider all possible outcomes and then search from
  # the outcome.
  action_rewards = {}
  for stone_index, potion_index in belief_state.possible_actions(precomputed):
    # This action means use or discard all stones.
    if potion_index == END_TRIAL:
      # Ending the trial so get the actual number of world states
      if minimise_world_states:
        expected_num_world_states = np.where(
            belief_state.world_state_distribution.poss_world_states)[0].size
      else:
        raw_reward_count = [(precomputed.stone_to_reward[s], c) for s, c in
                            belief_state.aligned_stones.items()]
        reward_per_stone_type = [
            0.0 if reward < 0 else c * (reward + bonus)
            if reward == stones_and_potions.max_reward() else c * reward
            for reward, c in raw_reward_count]
        action_reward = sum(reward_per_stone_type)
        best_action_depth = depth
    else:
      bit_mask = stone_potion_bit_mask(stone_index, potion_index, precomputed)
      if belief_state.world_state_distribution.observed_no_effect_bits & bit_mask:
        continue
      # Using a potion on a stone could lead to a number of possible outcomes
      # with various probabilities. The ideal observer must calculate what
      # reward it can expect to obtain for each of these.
      if minimise_world_states:
        # The expected number of world states when the trial ends
        expected_num_world_states = 0
      else:
        action_reward = 0.0
        best_action_depth = 0.0
      poss_outcomes, could_stay_still = (
          belief_state.world_state_distribution.possible_outcomes(
              potion_index, stone_index, precomputed))
      if len(poss_outcomes) == 1:
        # If there is only one possibility and it is staying still then there is
        # no need to search this action as it will have no effect.
        if could_stay_still:
          continue
        new_game_state = copy.deepcopy(belief_state)
        new_game_state.use_potion(stone_index, potion_index, poss_outcomes[0])
        _, objective = search(
            new_game_state, search_results, bonus, precomputed, depth + 1,
            minimise_world_states)
        if minimise_world_states:
          expected_num_world_states -= objective
        else:
          search_reward, neg_search_action_depth = objective
          action_reward += search_reward
          best_action_depth -= neg_search_action_depth
      else:
        for outcome in poss_outcomes:
          new_game_state = copy.deepcopy(belief_state)
          prob = new_game_state.action_and_outcome(
              stone_index, potion_index, outcome, precomputed, bit_mask)
          if prob > 0:
            _, objective = search(
                new_game_state, search_results, bonus, precomputed, depth + 1,
                minimise_world_states)
            if minimise_world_states:
              expected_num_world_states -= prob * objective
            else:
              search_reward, neg_search_action_depth = objective
              action_reward += prob * search_reward
              best_action_depth -= prob * neg_search_action_depth
    # Store the expected reward and the negative search depth so that when we
    # maximise reward, if 2 actions have the same expected reward then we will
    # take the action with the minimum search depth. This prevents us taking
    # actions which do not harm our expected reward but do not help us.
    if minimise_world_states:
      action_rewards[
          (stone_index, potion_index)] = -expected_num_world_states
    else:
      action_rewards[(stone_index, potion_index)] = (
          action_reward, -best_action_depth)

  result = max(action_rewards.items(), key=lambda a: a[1])

  search_results[belief_state_bitfield] = result
  return result

def ideal_observer(
    init_game_state: BeliefState,
    search_results: SearchResults,
    bonus: int, precomputed: PrecomputedMaps, minimise_world_states: bool
) -> ActionObjectiveAndSearchResults:
  """Runs the ideal observer given a set of stones and potions.

  This runs an exhaustive search from the initial state over possible actions
  and possible outcomes of those actions. It returns which action to take, the
  expected reward and a set of search results for belief states encountered.

  Args:
    init_game_state: The initial belief state of the system.
    search_results: Previously computed action and reward for belief states.
    bonus: The additional reward for getting the best stone.
    precomputed: Precomputed maps used for speed.
    minimise_world_states: Let the objective be to minimise the number of world
      states at the end of the trial instead of to maximise the accumulated
      reward.

  Returns:
    The best action to take, the expected reward and a set of search results for
    belief states encountered.
  """
  # Set the number of bits required to fit all possible partial graph indices.
  BeliefState.possible_partial_graph_num_bits = math.ceil(math.log2(
      len(precomputed.partial_graph_index_to_possible_index)))

  # Run the search over all possible next actions
  action, objective = search(
      init_game_state, search_results, bonus, precomputed,
      minimise_world_states=minimise_world_states)
  return action, objective, search_results