# -*- coding: utf-8 -*-

from collections import namedtuple

from numba import types
import numpy as np

from africanus.averaging.time_and_channel_mapping import (row_mapper,
                                                          channel_mapper)
from africanus.averaging.shared import (chan_corrs,
                                        merge_flags,
                                        vis_output_arrays)

from africanus.util.docs import DocstringTemplate
from africanus.util.numba import (is_numba_type_none, generated_jit,
                                  njit, overload, intrinsic)

TUPLE_TYPE = 0
ARRAY_TYPE = 1
NONE_TYPE = 2

def matching_flag_factory(present):
    if present:
        def impl(flag_row, ri, out_flag_row, ro):
            return flag_row[ri] == out_flag_row[ro]
    else:
        def impl(flag_row, ri, out_flag_row, ro):
            return True

    return njit(nogil=True, cache=True, inline='always')(impl)

def is_chan_flagged(flag, r, f, c):
    pass

@overload(is_chan_flagged, inline='always')
def _is_chan_flagged(flag, r, f, c):
    if is_numba_type_none(flag):
        def impl(flag, r, f, c):
            return True
    else:
        def impl(flag, r, f, c):
            return flag[r, f, c]

    return impl

@njit(nogil=True, inline='always')
def chan_add(output, input, orow, ochan, irow, ichan, corr):
    if input is not None:
        output[orow, ochan, corr] += input[irow, ichan, corr]

_row_output_fields = ["antenna1", "antenna2", "time_centroid", "exposure",
                      "uvw", "weight", "sigma"]
RowAverageOutput = namedtuple("RowAverageOutput", _row_output_fields)

@generated_jit(nopython=True, nogil=True, cache=True)
def row_average(meta, ant1, ant2, flag_row=None,
                time_centroid=None, exposure=None, uvw=None,
                weight=None, sigma=None):

    have_flag_row = not is_numba_type_none(flag_row)
    flags_match = matching_flag_factory(have_flag_row)

    def impl(meta, ant1, ant2, flag_row=None,
             time_centroid=None, exposure=None, uvw=None,
             weight=None, sigma=None):

        out_rows = meta.time.shape[0]

        counts = np.zeros(out_rows, dtype=np.uint32)

        # These outputs are always present
        ant1_avg = np.empty(out_rows, ant1.dtype)
        ant2_avg = np.empty(out_rows, ant2.dtype)

        # Possibly present outputs for possibly present inputs
        uvw_avg = (
            None if uvw is None else
            np.zeros((out_rows,) + uvw.shape[1:],
                     dtype=uvw.dtype))

        time_centroid_avg = (
            None if time_centroid is None else
            np.zeros((out_rows,) + time_centroid.shape[1:],
                     dtype=time_centroid.dtype))

        exposure_avg = (
            None if exposure is None else
            np.zeros((out_rows,) + exposure.shape[1:],
                     dtype=exposure.dtype))

        weight_avg = (
            None if weight is None else
            np.zeros((out_rows,) + weight.shape[1:],
                     dtype=weight.dtype))

        sigma_avg = (
            None if sigma is None else
            np.zeros((out_rows,) + sigma.shape[1:],
                     dtype=sigma.dtype))

        sigma_weight_sum = (
            None if sigma is None else
            np.zeros((out_rows,) + sigma.shape[1:],
                     dtype=sigma.dtype))

        # Iterate over input rows, accumulating into output rows
        for in_row, out_row in enumerate(meta.map):
            # Input and output flags must match in order for the
            # current row to contribute to these columns
            if flags_match(flag_row, in_row, meta.flag_row, out_row):
                if uvw is not None:
                    uvw_avg[out_row, 0] += uvw[in_row, 0]
                    uvw_avg[out_row, 1] += uvw[in_row, 1]
                    uvw_avg[out_row, 2] += uvw[in_row, 2]

                if time_centroid is not None:
                    time_centroid_avg[out_row] += time_centroid[in_row]

                if exposure is not None:
                    exposure_avg[out_row] += exposure[in_row]

                if weight is not None:
                    for co in range(weight.shape[1]):
                        weight_avg[out_row, co] += weight[in_row, co]

                if sigma is not None:
                    for co in range(sigma.shape[1]):
                        sva = sigma[in_row, co]**2

                        # Use provided weights
                        if weight is not None:
                            wt = weight[in_row, co]
                            sva *= wt ** 2
                            sigma_weight_sum[out_row, co] += wt
                        # Natural weights
                        else:
                            sigma_weight_sum[out_row, co] += 1.0

                        # Assign
                        sigma_avg[out_row, co] += sva

                counts[out_row] += 1

            # Here we can simply assign because input_row baselines
            # should always match output row baselines
            ant1_avg[out_row] = ant1[in_row]
            ant2_avg[out_row] = ant2[in_row]

        # Normalise
        for out_row in range(out_rows):
            count = counts[out_row]

            if count > 0:
                # Normalise uvw
                if uvw is not None:
                    uvw_avg[out_row, 0] /= count
                    uvw_avg[out_row, 1] /= count
                    uvw_avg[out_row, 2] /= count

                # Normalise time centroid
                if time_centroid is not None:
                    time_centroid_avg[out_row] /= count

                # Normalise sigma
                if sigma is not None:
                    for co in range(sigma.shape[1]):
                        ssva = sigma_avg[out_row, co]
                        wt = sigma_weight_sum[out_row, co]

                        if wt != 0.0:
                            ssva /= (wt**2)

                        sigma_avg[out_row, co] = np.sqrt(ssva)

        return RowAverageOutput(ant1_avg, ant2_avg,
                                time_centroid_avg,
                                exposure_avg, uvw_avg,
                                weight_avg, sigma_avg)

    return impl

_rowchan_output_fields = [
    "visibilities",
    "flag",
    "weight_spectrum",
    "sigma_spectrum"]
RowChanAverageOutput = namedtuple("RowChanAverageOutput",
                                  _rowchan_output_fields)

class RowChannelAverageException(Exception):
    pass

@intrinsic
def average_visibilities(typingctx, vis, vis_avg, vis_weight_sum,
                         weight, ri, fi, ro, fo, co):

    import numba.core.types as nbtypes

    have_array = isinstance(vis, nbtypes.Array)
    have_tuple = isinstance(vis, (nbtypes.Tuple, nbtypes.UniTuple))

    def avg_fn(vis, vis_avg, vis_ws, wt, ri, fi, ro, fo, co):
        vis_avg[ro, fo, co] += vis[ri, fi, co] * wt
        vis_ws[ro, fo, co] += wt

    return_type = nbtypes.NoneType("none")

    sig = return_type(vis, vis_avg, vis_weight_sum,
                      weight, ri, fi, ro, fo, co)

    def codegen(context, builder, signature, args):
        vis, vis_type = args[0], signature.args[0]
        vis_avg, vis_avg_type = args[1], signature.args[1]
        vis_weight_sum, vis_weight_sum_type = args[2], signature.args[2]
        weight, weight_type = args[3], signature.args[3]
        ri, ri_type = args[4], signature.args[4]
        fi, fi_type = args[5], signature.args[5]
        ro, ro_type = args[6], signature.args[6]
        fo, fo_type = args[7], signature.args[7]
        co, co_type = args[8], signature.args[8]
        return_type = signature.return_type

        if have_array:
            avg_sig = return_type(vis_type,
                                  vis_avg_type,
                                  vis_weight_sum_type,
                                  weight_type,
                                  ri_type, fi_type,
                                  ro_type, fo_type, co_type)
            avg_args = [vis, vis_avg, vis_weight_sum,
                        weight, ri, fi, ro, fo, co]

            # Compile function and get handle to output
            context.compile_internal(builder, avg_fn,
                                     avg_sig, avg_args)
        elif have_tuple:
            for i in range(len(vis_type)):
                avg_sig = return_type(vis_type.types[i],
                                      vis_avg_type.types[i],
                                      vis_weight_sum_type.types[i],
                                      weight_type,
                                      ri_type, fi_type,
                                      ro_type, fo_type, co_type)
                avg_args = [builder.extract_value(vis, i),
                            builder.extract_value(vis_avg, i),
                            builder.extract_value(vis_weight_sum, i),
                            weight, ri, fi, ro, fo, co]

                # Compile function and get handle to output
                context.compile_internal(builder, avg_fn,
                                         avg_sig, avg_args)
        else:
            raise TypeError("Unhandled visibility array type")

    return sig, codegen

@intrinsic
def normalise_visibilities(typingctx, vis_avg, vis_weight_sum, ro, fo, co):
    import numba.core.types as nbtypes

    have_array = isinstance(vis_avg, nbtypes.Array)
    have_tuple = isinstance(vis_avg, (nbtypes.Tuple, nbtypes.UniTuple))

    def normalise_fn(vis_avg, vis_ws, ro, fo, co):
        weight_sum = vis_ws[ro, fo, co]

        if weight_sum != 0.0:
            vis_avg[ro, fo, co] /= weight_sum

    return_type = nbtypes.NoneType("none")
    sig = return_type(vis_avg, vis_weight_sum, ro, fo, co)

    def codegen(context, builder, signature, args):
        vis_avg, vis_avg_type = args[0], signature.args[0]
        vis_weight_sum, vis_weight_sum_type = args[1], signature.args[1]
        ro, ro_type = args[2], signature.args[2]
        fo, fo_type = args[3], signature.args[3]
        co, co_type = args[4], signature.args[4]
        return_type = signature.return_type

        if have_array:
            # Normalise single array
            norm_sig = return_type(vis_avg_type,
                                   vis_weight_sum_type,
                                   ro_type, fo_type, co_type)
            norm_args = [vis_avg, vis_weight_sum, ro, fo, co]

            context.compile_internal(builder, normalise_fn,
                                     norm_sig, norm_args)
        elif have_tuple:
            # Normalise each array in the tuple
            for i in range(len(vis_avg_type)):
                norm_sig = return_type(vis_avg_type.types[i],
                                       vis_weight_sum_type.types[i],
                                       ro_type, fo_type, co_type)
                norm_args = [builder.extract_value(vis_avg, i),
                             builder.extract_value(vis_weight_sum, i),
                             ro, fo, co]

                # Compile function and get handle to output
                context.compile_internal(builder, normalise_fn,
                                         norm_sig, norm_args)
        else:
            raise TypeError("Unhandled visibility array type")

    return sig, codegen

@generated_jit(nopython=True, nogil=True, cache=True)
def row_chan_average(row_meta, chan_meta,
                     flag_row=None, weight=None,
                     visibilities=None, flag=None,
                     weight_spectrum=None, sigma_spectrum=None):

    dummy_chan_freq = None
    dummy_chan_width = None

    have_vis = not is_numba_type_none(visibilities)
    have_flag = not is_numba_type_none(flag)
    have_flag_row = not is_numba_type_none(flag_row)
    have_flags = have_flag_row or have_flag

    have_weight = not is_numba_type_none(weight)
    have_weight_spectrum = not is_numba_type_none(weight_spectrum)
    have_sigma_spectrum = not is_numba_type_none(sigma_spectrum)

    def impl(row_meta, chan_meta, flag_row=None, weight=None,
             visibilities=None, flag=None,
             weight_spectrum=None, sigma_spectrum=None):

        out_rows = row_meta.time.shape[0]
        nchan, ncorrs = chan_corrs(visibilities, flag,
                                   weight_spectrum, sigma_spectrum,
                                   dummy_chan_freq, dummy_chan_width,
                                   dummy_chan_width, dummy_chan_width)

        chan_map, out_chans = chan_meta

        in_shape = (row_meta.map.shape[0], nchan, ncorrs)
        out_shape = (out_rows, out_chans, ncorrs)

        if not have_flag:
            flag_avg = None
        else:
            flag_avg = np.full(out_shape, False, dtype=np.bool_)

        # If either flag_row or flag is present, we need to ensure that
        # effective averaging takes place.
        if have_flags:
            flags_match = np.full(in_shape, False, dtype=np.bool_)
            flag_counts = np.zeros(out_shape, dtype=np.uint32)
        else:
            flags_match = None
            flag_counts = None

        counts = np.zeros(out_shape, dtype=np.uint32)

        # Determine output bin counts both unflagged and flagged
        for ri, ro in enumerate(row_meta.map):
            row_flagged = have_flag_row and flag_row[ri] != 0
            for fi, fo in enumerate(chan_map):
                for co in range(ncorrs):
                    flagged = (row_flagged or
                               (have_flag and flag[ri, fi, co] != 0))

                    if have_flags and flagged:
                        flag_counts[ro, fo, co] += 1
                    else:
                        counts[ro, fo, co] += 1

        # ------
        # Flags
        # ------

        # Determine whether input samples should contribute to an output bin
        # and, if flags are parent, whether the output bin is flagged

        # This follows from the definition of an effective average:
        #
        # * bad or flagged values should be excluded
        #   when calculating the average
        #
        # Note that if a bin is completely flagged we still compute an average,
        # to which all relevant input samples contribute.
        for ri, ro in enumerate(row_meta.map):
            row_flagged = have_flag_row and flag_row[ri] != 0
            for fi, fo in enumerate(chan_map):
                for co in range(ncorrs):
                    if counts[ro, fo, co] > 0:
                        # Output bin should only contain unflagged samples
                        out_flag = False

                        if have_flag:
                            # Set output flags
                            flag_avg[ro, fo, co] = False

                    elif have_flags and flag_counts[ro, fo, co] > 0:
                        # Output bin is completely flagged
                        out_flag = True

                        if have_flag:
                            # Set output flags
                            flag_avg[ro, fo, co] = True
                    else:
                        raise RowChannelAverageException("Zero-filled bin")

                    # We should only add a sample to an output bin
                    # if the input flag matches the output flag.
                    # This is because flagged samples don't contribute
                    # to a bin with some unflagged samples while
                    # unflagged samples never contribute to a
                    # completely flagged bin
                    if have_flags:
                        in_flag = (row_flagged or
                                   (have_flag and flag[ri, fi, co] != 0))
                        flags_match[ri, fi, co] = in_flag == out_flag

        # -------------
        # Visibilities
        # -------------
        if not have_vis:
            vis_avg = None
        else:
            vis_avg, vis_weight_sum = vis_output_arrays(
                visibilities, out_shape)

            # # Aggregate
            for ri, ro in enumerate(row_meta.map):
                for fi, fo in enumerate(chan_map):
                    for co in range(ncorrs):
                        if have_flags and not flags_match[ri, fi, co]:
                            continue

                        wt = (weight_spectrum[ri, fi, co]
                              if have_weight_spectrum else
                              weight[ri, co] if have_weight else 1.0)

                        average_visibilities(visibilities,
                                             vis_avg,
                                             vis_weight_sum,
                                             wt, ri, fi, ro, fo, co)

            # Normalise
            for ro in range(out_rows):
                for fo in range(out_chans):
                    for co in range(ncorrs):
                        normalise_visibilities(
                            vis_avg, vis_weight_sum, ro, fo, co)

        # ----------------
        # Weight Spectrum
        # ----------------

        if not have_weight_spectrum:
            weight_spectrum_avg = None
        else:
            weight_spectrum_avg = np.zeros(out_shape, weight_spectrum.dtype)

            # Aggregate
            for ri, ro in enumerate(row_meta.map):
                for fi, fo in enumerate(chan_map):
                    for co in range(ncorrs):
                        if have_flags and not flags_match[ri, fi, co]:
                            continue

                        weight_spectrum_avg[ro, fo, co] += (
                            weight_spectrum[ri, fi, co])

        # ---------------
        # Sigma Spectrum
        # ---------------
        if not have_sigma_spectrum:
            sigma_spectrum_avg = None
        else:
            sigma_spectrum_avg = np.zeros(out_shape, sigma_spectrum.dtype)
            sigma_spectrum_weight_sum = np.zeros_like(sigma_spectrum_avg)

            # Aggregate
            for ri, ro in enumerate(row_meta.map):
                for fi, fo in enumerate(chan_map):
                    for co in range(ncorrs):
                        if have_flags and not flags_match[ri, fi, co]:
                            continue

                        wt = (weight_spectrum[ri, fi, co]
                              if have_weight_spectrum else
                              weight[ri, co] if have_weight else 1.0)

                        ssv = sigma_spectrum[ri, fi, co]**2 * wt**2
                        sigma_spectrum_avg[ro, fo, co] += ssv
                        sigma_spectrum_weight_sum[ro, fo, co] += wt

            # Normalise
            for ro in range(out_rows):
                for fo in range(out_chans):
                    for co in range(ncorrs):
                        sswsum = sigma_spectrum_weight_sum[ro, fo, co]
                        if sswsum != 0.0:
                            ssv = sigma_spectrum_avg[ro, fo, co]
                            sigma_spectrum_avg[ro, fo, co] = np.sqrt(
                                ssv / sswsum**2)

        return RowChanAverageOutput(vis_avg, flag_avg,
                                    weight_spectrum_avg,
                                    sigma_spectrum_avg)

    return impl

_chan_output_fields = ["chan_freq", "chan_width", "effective_bw", "resolution"]
ChannelAverageOutput = namedtuple("ChannelAverageOutput", _chan_output_fields)

@generated_jit(nopython=True, nogil=True, cache=True)
def chan_average(chan_meta, chan_freq=None, chan_width=None,
                 effective_bw=None, resolution=None):

    def impl(chan_meta, chan_freq=None, chan_width=None,
             effective_bw=None, resolution=None):
        chan_map, out_chans = chan_meta

        chan_freq_avg = (
            None if chan_freq is None else
            np.zeros(out_chans, dtype=chan_freq.dtype))

        chan_width_avg = (
            None if chan_width is None else
            np.zeros(out_chans, dtype=chan_width.dtype))

        effective_bw_avg = (
            None if effective_bw is None else
            np.zeros(out_chans, dtype=effective_bw.dtype))

        resolution_avg = (
            None if resolution is None else
            np.zeros(out_chans, dtype=resolution.dtype))

        counts = np.zeros(out_chans, dtype=np.uint32)

        for in_chan, out_chan in enumerate(chan_map):
            counts[out_chan] += 1

            if chan_freq is not None:
                chan_freq_avg[out_chan] += chan_freq[in_chan]

            if chan_width is not None:
                chan_width_avg[out_chan] += chan_width[in_chan]

            if effective_bw is not None:
                effective_bw_avg[out_chan] += effective_bw[in_chan]

            if resolution is not None:
                resolution_avg[out_chan] += resolution[in_chan]

        for out_chan in range(out_chans):
            if chan_freq is not None:
                chan_freq_avg[out_chan] /= counts[out_chan]

        return ChannelAverageOutput(chan_freq_avg, chan_width_avg,
                                    effective_bw_avg, resolution_avg)

    return impl

AverageOutput = namedtuple("AverageOutput",
                           ["time", "interval", "flag_row"] +
                           _row_output_fields +
                           _chan_output_fields +
                           _rowchan_output_fields)

@generated_jit(nopython=True, nogil=True, cache=True)
def time_and_channel(time, interval, antenna1, antenna2,
                     time_centroid=None, exposure=None, flag_row=None,
                     uvw=None, weight=None, sigma=None,
                     chan_freq=None, chan_width=None,
                     effective_bw=None, resolution=None,
                     visibilities=None, flag=None,
                     weight_spectrum=None, sigma_spectrum=None,
                     time_bin_secs=1.0, chan_bin_size=1):

    valid_types = (types.misc.Omitted, types.scalars.Float,
                   types.scalars.Integer)

    if not isinstance(time_bin_secs, valid_types):
        raise TypeError("time_bin_secs must be a scalar float")

    valid_types = (types.misc.Omitted, types.scalars.Integer)

    if not isinstance(chan_bin_size, valid_types):
        raise TypeError("chan_bin_size must be a scalar integer")

    def impl(time, interval, antenna1, antenna2,
             time_centroid=None, exposure=None, flag_row=None,
             uvw=None, weight=None, sigma=None,
             chan_freq=None, chan_width=None,
             effective_bw=None, resolution=None,
             visibilities=None, flag=None,
             weight_spectrum=None, sigma_spectrum=None,
             time_bin_secs=1.0, chan_bin_size=1):

        nchan, ncorrs = chan_corrs(visibilities, flag,
                                   weight_spectrum, sigma_spectrum,
                                   chan_freq, chan_width,
                                   effective_bw, resolution)

        # Merge flag_row and flag arrays
        flag_row = merge_flags(flag_row, flag)

        # Generate row mapping metadata
        row_meta = row_mapper(time, interval, antenna1, antenna2,
                              flag_row=flag_row, time_bin_secs=time_bin_secs)

        # Generate channel mapping metadata
        chan_meta = channel_mapper(nchan, chan_bin_size)

        # Average row data
        row_data = row_average(row_meta, antenna1, antenna2, flag_row=flag_row,
                               time_centroid=time_centroid, exposure=exposure,
                               uvw=uvw, weight=weight, sigma=sigma)

        # Average channel data
        chan_data = chan_average(chan_meta, chan_freq=chan_freq,
                                 chan_width=chan_width,
                                 effective_bw=effective_bw,
                                 resolution=resolution)

        # Average row and channel data
        row_chan_data = row_chan_average(row_meta, chan_meta,
                                         flag_row=flag_row, weight=weight,
                                         visibilities=visibilities, flag=flag,
                                         weight_spectrum=weight_spectrum,
                                         sigma_spectrum=sigma_spectrum)

        # Have to explicitly write it out because numba tuples
        # are highly constrained types
        return AverageOutput(row_meta.time,
                             row_meta.interval,
                             row_meta.flag_row,
                             row_data.antenna1,
                             row_data.antenna2,
                             row_data.time_centroid,
                             row_data.exposure,
                             row_data.uvw,
                             row_data.weight,
                             row_data.sigma,
                             chan_data.chan_freq,
                             chan_data.chan_width,
                             chan_data.effective_bw,
                             chan_data.resolution,
                             row_chan_data.visibilities,
                             row_chan_data.flag,
                             row_chan_data.weight_spectrum,
                             row_chan_data.sigma_spectrum)

    return impl

AVERAGING_DOCS = DocstringTemplate("""
Averages in time and channel.

Parameters
----------
time : $(array_type)
    Time values of shape :code:`(row,)`.
interval : $(array_type)
    Interval values of shape :code:`(row,)`.
antenna1 : $(array_type)
    First antenna indices of shape :code:`(row,)`
antenna2 : $(array_type)
    Second antenna indices of shape :code:`(row,)`
time_centroid : $(array_type), optional
    Time centroid values of shape :code:`(row,)`
exposure : $(array_type), optional
    Exposure values of shape :code:`(row,)`
flag_row : $(array_type), optional
    Flagged rows of shape :code:`(row,)`.
uvw : $(array_type), optional
    UVW coordinates of shape :code:`(row, 3)`.
weight : $(array_type), optional
    Weight values of shape :code:`(row, corr)`.
sigma : $(array_type), optional
    Sigma values of shape :code:`(row, corr)`.
chan_freq : $(array_type), optional
    Channel frequencies of shape :code:`(chan,)`.
chan_width : $(array_type), optional
    Channel widths of shape :code:`(chan,)`.
effective_bw : $(array_type), optional
    Effective channel bandwidth of shape :code:`(chan,)`.
resolution : $(array_type), optional
    Effective channel resolution of shape :code:`(chan,)`.
visibilities : $(array_type) or tuple of $(array_type), optional
    Visibility data of shape :code:`(row, chan, corr)`.
    Tuples of visibilities arrays may be supplied,
    in which case tuples will be output.
flag : $(array_type), optional
    Flag data of shape :code:`(row, chan, corr)`.
weight_spectrum : $(array_type), optional
    Weight spectrum of shape :code:`(row, chan, corr)`.
sigma_spectrum : $(array_type), optional
    Sigma spectrum of shape :code:`(row, chan, corr)`.
time_bin_secs : float, optional
    Maximum summed interval in seconds to include within a bin.
    Defaults to 1.0.
chan_bin_size : int, optional
    Number of bins to average together.
    Defaults to 1.

Notes
-----

The implementation currently requires unique lexicographical
combinations of (TIME, ANTENNA1, ANTENNA2). This can usually
be achieved by suitably partitioning input data on indexing rows,
DATA_DESC_ID and SCAN_NUMBER in particular.

Returns
-------
namedtuple
    A namedtuple whose entries correspond to the input arrays.
    Output arrays will be ``None`` if the inputs were ``None``.
""")

try:
    time_and_channel.__doc__ = AVERAGING_DOCS.substitute(
                                    array_type=":class:`numpy.ndarray`")
except AttributeError:
    pass


from genericpath import getsize
#from statistics import covariance
import sys
from turtle import setundobuffer
import numpy as np
#sys.path=["/home/creichardt/.local/lib/python3.7/site-packages/","/home/creichardt/spt3g_software/build","/home/creichardt/.local/lib/python3.7/site-packages/healpy-1.15.0-py3.7-linux-x86_64.egg"]+sys.path
import healpy

import os
from spt3g import core,maps, calibration
from spectra_port import utils
import pickle as pkl
import pdb
import time
AlmType = np.dtype(np.complex64)

def name_tempdir(basedir):
    while True:
        rand = np.random.randint(0,999999)
        path = "{}/workdir_{:6d}".format(basedir, rand)
        if not os.path.exists(path):
            return path

def printinplace(myString):
    digits = len(myString)
    delete = "\b" * (digits)
    print("{0}{1:{2}}".format(delete, myString, digits), end="")
    sys.stdout.flush()

def load_spt3g_healpix_ring_map(file,require_order = 'Ring',require_nside=8192,map_key='T'):
    # only taking 1st map
    frames = core.G3File(file)
    for frame in frames:
        if frame.type == core.G3FrameType.Map:
            if require_nside is not None:
                assert(require_nside == frame[map_key].nside)
            if require_order is not None:
                assert(frame[map_key].nested == (require_order == 'Nest'))
            if type(frame[map_key]) is np.ndarray:
                ind = frame[map_key].nonzero()
                map = frame[map_key][frame[map_key] != 0]
            else:
                ind, map = frame[map_key].nonzero_pixels()
            return( np.asarray(ind).astype(np.int64,casting='same_kind'), np.asarray(map).astype(np.float32,casting='same_kind') )
    raise Exception("No Map found in file: {}".format(file))

def reformat_shts(shtfilelist, processedshtfile,
                           lmax,
                           cmbweighting = True,
                           mask  = None,
                           kmask = None,
                           ell_reordering=None,
                           no_reorder=False,
                           ram_limit = None,
                          ) -> 'May be done in Fortran - output is a file':
    '''
    Output is expected to be CL (Dl if cmbweighting=Trure) * mask normalization factor * kweights
    Output ordering is expected to
    '''
    if ram_limit is None:
        ram_limit = 16 * 2**30 # 16 GB

    # number of bytes in a Dcomplex: 16
    # number of arrays we need to make to do this efficiently: 6 or less
    # number of pixels in an fft: winsize^2
    ram_required=16*6*lmax**2
    parallelism = int(np.ceil(ram_limit/ram_required))

    inv_mask_factor = 1.
    if mask is not None:
        inv_mask_factor = np.sqrt(1./np.mean(mask**2))

    #ie do parallelism SHTs at once...
    size = healpy.sphtfunc.Alm.getsize(lmax)
    if kmask is not None:
        if kmask.shape[0] != size:
            raise Exception("kmask provided is wrong size ({} vs {}), exiting".format(size,kmask.shape[0]))
        local_kmask = kmask.astype(np.float32)
    else:
        local_kmask = np.ones(size,dtype=np.float32)

    if cmbweighting:
        dummy_vec = np.arange(lmax+1,dtype=np.float32)
        dummy_vec = np.sqrt((dummy_vec*(dummy_vec+1.))/(2*np.pi)) # This will be squared since Cl =a*a
        j=0
        for i in range(lmax+1):
            nm = lmax+1-i
            local_kmask[j:j+nm]=dummy_vec[i:]
            j=j+nm

    if ell_reordering is None:  # need to make it
        #have lmax+1 m=0's, followed by lmax m=1's.... (if does do l=0,m=0)
        # healpy has indexing routines, but they only take 1 at a time...
        #make dummy vec for use
        dummy_vec = np.zeros(lmax+1,dtype=np.int)
        k=0
        for i in np.arange(lmax+1):
            dummy_vec[i] = k
            k=k+lmax-i
        ell_reordering = np.zeros(size,dtype=np.int)
        k=0
        for i in range(lmax+1):
            ell_reordering[k:k+i+1] = dummy_vec[0:i+1] + i
            k += i+1

    print('Warning: not using pixel weights in SHT')
    with open(processedshtfile,'wb') as fp:

        oldtime = time.time()
        count = 0
        for file in shtfilelist:
            newtime=time.time()
            timeinminutes = (newtime - oldtime)/60.0
            oldtime=newtime
            printinplace('SHT map: {}  Last one took: {:.1f} minutes'.format(count,timeinminutes))
            count += 1

            #TBD get SHT
            with np.load(file) as obs_alms:
                alms = obs_alms['alm']
            assert lmax ==  healpy.sphtfunc.Alm.getlmax(alms.shape[0])

            #possibly downsample alms to save later CPU cycles
            # TBD if worthwhile

            #apply weighting (ie cl-dl) and kmask
            alms *= local_kmask

            #TBD, possibly adjust for mask factor here
            alms *= inv_mask_factor

            # Get reindexing

            #reorder and write to disk
            #need to check sizing
            #32 bit floats/64b complex should be fine for this. will need to bump up by one for aggregation
            if no_reorder:
                (alms.astype(AlmType)).tofile(fp)
            else:
                (alms[ell_reordering].astype(AlmType)).tofile(fp)

def take_and_reformat_shts(mapfilelist, processedshtfile,
                           nside,lmax,
                           cmbweighting = True,
                           mask  = None,
                           kmask = None,
                           ell_reordering=None,
                           no_reorder=False,
                           ram_limit = None,
                           npmapformat=False,
                           map_key='T'
                          ) -> 'May be done in Fortran - output is a file':
    '''
    Output is expected to be CL (Dl if cmbweighting=Trure) * mask normalization factor * kweights
    Output ordering is expected to
    '''
    if ram_limit is None:
        ram_limit = 16 * 2**30 # 16 GB

    # number of bytes in a Dcomplex: 16
    # number of arrays we need to make to do this efficiently: 6 or less
    # number of pixels in an fft: winsize^2
    ram_required=16*6*lmax**2
    parallelism = int(np.ceil(ram_limit/ram_required))

    map_scratch = np.zeros(12*nside**2,dtype=np.float32)

    inv_mask_factor = 1.
    if mask is not None:
        inv_mask_factor = np.sqrt(1./np.mean(mask**2))

    if mask is not None:
        if type(mask) is np.ndarray:
            map_inds = mask.nonzero()
            cut_mask = mask[mask != 0]
        else:
            map_inds, cut_mask = mask.nonzero_pixels()
        npix = len(map_inds)

    #ie do parallelism SHTs at once...
    size = healpy.sphtfunc.Alm.getsize(lmax)
    if kmask is not None:
        if kmask.shape[0] != size:
            raise Exception("kmask provided is wrong size ({} vs {}), exiting".format(size,kmask.shape[0]))
        local_kmask = kmask.astype(np.float32)
    else:
        local_kmask = np.ones(size,dtype=np.float32)

    if cmbweighting:
        dummy_vec = np.arange(lmax+1,dtype=np.float32)
        dummy_vec = np.sqrt((dummy_vec*(dummy_vec+1.))/(2*np.pi)) #this will be squared later as Cl=a*a
        j=0
        for i in range(lmax+1):
            nm = lmax+1-i
            local_kmask[j:j+nm]=dummy_vec[i:]
            j=j+nm

    if ell_reordering is None:  # need to make it
        #have lmax+1 m=0's, followed by lmax m=1's.... (if does do l=0,m=0)
        # healpy has indexing routines, but they only take 1 at a time...
        #make dummy vec for use
        dummy_vec = np.zeros(lmax+1,dtype=np.int)
        k=0
        for i in np.arange(lmax+1):
            dummy_vec[i] = k
            k=k+lmax-i
        ell_reordering = np.zeros(size,dtype=np.int)
        k=0
        for i in range(lmax+1):
            ell_reordering[k:k+i+1] = dummy_vec[0:i+1] + i
            k += i+1

    print('Warning: not using pixel weights in SHT')
    with open(processedshtfile,'wb') as fp:

        oldtime = time.time()
        count = 0
        for file in mapfilelist:
            newtime=time.time()
            timeinminutes = (newtime - oldtime)/60.0
            oldtime=newtime
            printinplace('SHT map: {}  Last one took: {:.1f} minutes'.format(count,timeinminutes))
            count += 1

            #TBD get a map
            if not npmapformat:
                ring_indices, map_tmp = load_spt3g_healpix_ring_map(file,map_key=map_key)

                map_scratch[:]=0 #reset
                map_scratch[ring_indices]=map_tmp #fill in the temperature map

                # if not already masked, apply mask
                #may need to change the next line based on formatting
                if mask is not None:
                    map_scratch  = mask*map_scratch
            else:
                map_tmp = utils.load_spt3g_cutsky_healpix_ring_map(file,npix)
                map_scratch[map_inds]=map_tmp * cut_mask

            #gets alms
            alms = healpy.sphtfunc.map2alm(map_scratch,lmax = lmax, pol=False, use_pixel_weights=False, iter = 1,datapath='/sptlocal/user/creichardt/healpy-data/')

            #possibly downsample alms to save later CPU cycles
            # TBD if worthwhile

            #apply weighting (ie cl-dl) and kmask
            alms *= local_kmask

            #TBD, possibly adjust for mask factor here
            alms *= inv_mask_factor

            # Get reindexing

            #reorder and write to disk
            #need to check sizing
            #32 bit floats/64b complex should be fine for this. will need to bump up by one for aggregation
            if no_reorder:
                (alms.astype(AlmType)).tofile(fp)
            else:
                (alms[ell_reordering].astype(AlmType)).tofile(fp)

def get_first_index_ell(l):
    # l=0 - 0
    # l = 1 -> 1 (0+1)
    # l = 2 ->  3 (1+2)
    # l = 3 -> 6 (3+3)
    if type(l) is int:
        return int(l*(l+1)/2)
    elif type(l) is np.ndarray:
        return (l*(l+1)/2).astype(np.int)
    else:
        pdb.set_trace()
        return -1

def generate_jackknife_shts( processed_shtfile, jackknife_shtfile,  lmax,
                             setdef) -> 'Does differencing to make SHT equiv file for nulls, returns new setdef':
    buffer_size = healpy.sphtfunc.Alm.getsize(lmax)
    buffer_bytes= buffer_size * np.zeros(1,dtype=AlmType).nbytes
    buffera = np.zeros(buffer_size,dtype=AlmType)
    bufferb = np.zeros(buffer_size,dtype=AlmType)
    setsize = setdef.shape[0]
    nsets   = setdef.shape[1]

    oldtime=time.time()
    with open(processed_shtfile,'rb') as fin, open(jackknife_shtfile,'wb') as fout:
        for i in range(setsize):
            newtime=time.time()
            timeinminutes = (newtime - oldtime)/60.0
            oldtime=newtime
            printinplace('Creating null SHT : {} of {}  Last one took: {:.1f} minutes'.format(i,setsize,timeinminutes))
            #need to do stuff here

            fin.seek( setdef[i,0] * buffer_bytes )
            buffera  = np.fromfile(fin,dtype=AlmType,count=buffer_size)
            fin.seek( setdef[i,1] * buffer_bytes )
            bufferb  = np.fromfile(fin,dtype=AlmType,count=buffer_size)

            buffera -= bufferb
            buffera *= 0.5
            #fout.seek( i * buffer_bytes )
            (buffera.astype(AlmType)).tofile(fout)

    return(np.reshape(np.arange(setsize,dtype=np.int32),[setsize,1]))

def generate_coadd_shts( processed_shtfile, coadd_shtfile,  lmax,
                             setdef) -> 'Does differencing to make SHT equiv file for nulls, returns new setdef':
    '''
    Setdef in: Nbundles_out (dim0) x Nmaps_to_coadd (dim1)
    Setdef out: Nbundles_out (same as input dim0)
    '''

    buffer_size = healpy.sphtfunc.Alm.getsize(lmax)
    buffer_bytes= buffer_size * np.zeros(1,dtype=AlmType).nbytes
    buffera = np.zeros(buffer_size,dtype=AlmType)
    bufferb = np.zeros(buffer_size,dtype=AlmType)
    setsize = setdef.shape[0]
    nsets   = setdef.shape[1]

    with open(processed_shtfile,'rb') as fin, open(coadd_shtfile,'wb') as fout:
        for i in range(setsize):
            #need to do stuff here

            fin.seek( setdef[i,0] * buffer_bytes )
            buffera  = np.fromfile(fin,dtype=AlmType,count=buffer_size)
            for j in range(1,nsets):
                fin.seek( setdef[i,1] * buffer_bytes )
                bufferb  = np.fromfile(fin,dtype=AlmType,count=buffer_size)
                buffera += bufferb
            buffera *= (1./nsets)
            fout.seek( i * buffer_bytes )
            buffera.tofile(fout)

    return(np.reshape(np.arange(setsize,dtype=np.int32),[setsize,1]))

def load_cross_spectra_data_from_disk(shtfile, nshts, npersht, start, stop):
    nelems = stop - start + 1
    buffer_bytes = np.zeros(1,dtype=AlmType).nbytes
    data = np.zeros([nshts,nelems],dtype=AlmType)
    print(nshts,nelems)
    with open(shtfile,'r') as fp:
        for i in range(nshts):
            fp.seek((i*npersht+start) * buffer_bytes)
            data[i,:] = np.fromfile(fp,count=nelems,dtype=AlmType)
    return data

def take_all_cross_spectra( processedshtfile, lmax,
                            setdef, banddef, ram_limit=None, auto = False,nshts=None) -> 'Returns set of all x-spectra, binned':
    '''
    ;; Step 1, copy all of the fft files and apply scalings masks etc

    ;; Step 2 (this function):  average all the bands to create binned x-spectra
    '''
    if ram_limit is None:
        ram_limit = 64 * 2**30 # default limit is 64 GB

    # Simplifying assumption axb == (a^c b + b^c a)
    # assume do *not* do x-spectra between same observation
    nsets   = setdef.shape[1] #nfreq
    setsize = setdef.shape[0] #nbundles
    nspectra=np.int((nsets*(nsets+1))/2 + 0.001)
    print(nsets,setsize,nspectra)
    if auto:
        nrealizations=setsize
    else:
        nrealizations=np.int( (setsize*(setsize-1))/2 + 0.001)

    nbands = banddef.shape[0]-1
    if nshts is  None:
        nshts  = np.int(np.max(setdef)+1.001)

    npersht = healpy.sphtfunc.Alm.getsize(lmax)
    #pdb.set_trace()
    allspectra_out = np.zeros([nbands,nspectra,nrealizations],dtype=np.float32)
    nmodes_out     = np.zeros(nbands, dtype = np.int32)

    tmpresult = np.zeros([setsize,setsize],dtype=np.float64)

    # number of bytes in a Dcomplex: 16
    # number of arrays we need to make to do this efficiently: 6 or less
    # number of pixels in an fft: winsize^2
    ram_required=16*6*lmax**2
    max_nmodes=ram_limit/nshts/32 #64 b complex

    assert(banddef[0] == 0 and banddef[-1] < lmax)
    #assumes banddef[0]=0
    #so first bin goes 1 - banddef[1]
    # second bin goes banddef[1]+1 - banddef[2], etc
    band_start_idx = get_first_index_ell(banddef+1)

    #code=reverse_linefeed_code()

    i=0 # i is the last bin to have finished. initially 0
    while (i < nbands):
        istop = np.where((band_start_idx - band_start_idx[i]) < max_nmodes)[0][-1] # get out of tuple, then take last elem of array

        if istop <= i:
            raise Exception("Insufficient ram for processing even a single bin")

        print('take_all_cross_spectra: loading bands {} {}'.format(i,istop-1))
        # technical: delete the last iteration of banddata_big first
        banddata_big=0
        # get data for as many bins as will fit in our ramlimit

        banddata_big=load_cross_spectra_data_from_disk(processedshtfile,
                                                       nshts, npersht,
                                                       band_start_idx[i],
                                                       band_start_idx[istop]-1 )
        #process this data
        for iprime in range(i, istop):
            printinplace('processing band {}    '.format(iprime))

            nmodes=(band_start_idx[iprime+1]-band_start_idx[iprime])
            nmodes_out[iprime]=nmodes
            aidx=band_start_idx[iprime]-band_start_idx[i]
            banddata=banddata_big[:,aidx:(aidx+nmodes-1)] # first index SHT; second index alm

            spectrum_idx=0
            for j in range(nsets):
                for k in range(j, nsets):
                    if not auto:

                        tmpresult  = np.real(np.matmul(banddata[setdef[:,j],:],np.conj(banddata[setdef[:,k],:]).T)) #need to check dims -- intended to end up for 3 freqs with 3x3 matrix

                        tmpresult += tmpresult.T # imposing the ab + ba condition
                        tmpresult /= (2*nmodes)
                        #it had a factor of 1/(reso**2 winsize**2)
                        # leaving this out for curved sky
                        a=0
                        for l in range(setsize-1):
                            rowlength=setsize-l-1
                            allspectra_out[iprime, spectrum_idx, a:(a+rowlength)]=tmpresult[l, l+1:setsize]
                            a+=rowlength
                    else:
                        idx=np.arange(setsize,dtype=np.int)
                        tmpresult=np.sum(np.real(banddata[setdef[:, j],:]*np.conj(banddata[setdef[:, k],:])), 1,dtype=np.float64) / (nmodes) # had been in flatsky: *reso^2*winsize^2)
                        allspectra_out[iprime, spectrum_idx, :]=tmpresult.astype(np.float32)
                    spectrum_idx+=1
                    #           pdb.set_trace()
        i=istop
    return(allspectra_out, nmodes_out)

def take_all_sim_cross_spectra( processedshtfile, lmax,
                            setdef1,  banddef, setdef2=None, ram_limit=None, auto=False) -> 'Returns set of all x-spectra, binned':
    '''
    ;; Step 1, copy all of the fft files and apply scalings masks etc

    ;; Step 2 (this function):  average all the bands to create binned x-spectra
    ;; this assumes sims are created with two bundles
    '''
    if ram_limit is None:
        ram_limit = 32 * 2**30 # default limit is 32 GB

    # Simplifying assumption axb == (a^c b + b^c a)
    # assume do *not* do x-spectra between same observation
    nsets   = setdef1.shape[1]
    setsize = setdef1.shape[0]
    nspectra=np.int((nsets*(nsets+1))/2 + 0.001)
    print(nsets,setsize,nspectra)

    if auto is False:
        assert setdef2 is not None

    nrealizations=setsize

    nbands = banddef.shape[0]-1

    nshts  = np.int(np.max([setdef1,setdef2])+1.001)
    npersht = healpy.sphtfunc.Alm.getsize(lmax)
    #pdb.set_trace()
    allspectra_out = np.zeros([nbands,nspectra,nrealizations],dtype=np.float32)
    nmodes_out     = np.zeros(nbands, dtype = np.int32)

    tmpresult = np.zeros([setsize,setsize],dtype=np.float64)

    # number of bytes in a Dcomplex: 16
    # number of arrays we need to make to do this efficiently: 6 or less
    # number of pixels in an fft: winsize^2
    ram_required=16*6*lmax**2
    max_nmodes=ram_limit/nshts/32 #64 b complex

    assert(banddef[0] == 0 and banddef[-1] <= lmax)
    #assumes banddef[0]=0
    #so first bin goes 1 - banddef[1]
    # second bin goes banddef[1]+1 - banddef[2], etc
    band_start_idx = get_first_index_ell(banddef+1)

    #code=reverse_linefeed_code()

    i=0 # i is the last bin to have finished. initially 0
    while (i < nbands):
        istop = np.where((band_start_idx - band_start_idx[i]) < max_nmodes)[0][-1] # get out of tuple, then take last elem of array

        if istop <= i:
            print('ram hit:',max_modes, band_start_idx[i],band_start_idx[i+1])
            raise Exception("Insufficient ram for processing even a single bin")

        print('take_all_cross_spectra: loading bands {} {}'.format(i,istop-1))
        # technical: delete the last iteration of banddata_big first
        banddata_big=0
        # get data for as many bins as will fit in our ramlimit

        banddata_big=load_cross_spectra_data_from_disk(processedshtfile,
                                                       nshts, npersht,
                                                       band_start_idx[i],
                                                       band_start_idx[istop]-1 )
        #process this data
        for iprime in range(i, istop):
            printinplace('processing band {}    '.format(iprime))

            nmodes=(band_start_idx[iprime+1]-band_start_idx[iprime])
            nmodes_out[iprime]=nmodes
            aidx=band_start_idx[iprime]-band_start_idx[i]
            banddata=banddata_big[:,aidx:(aidx+nmodes-1)] # first index SHT; second index alm

            spectrum_idx=0
            for j in range(nsets):
                for k in range(j, nsets):
                    if not auto:

                        #hypothetically, have 150a, 150b, 220a,220b
                        #want to end with:
                        # 150a * 220b + 150b * 220a
                        #iew 1j* 2k + 2j* 1k
                        tmpresult  =np.sum(np.real(banddata[setdef1[:, j],:]*np.conj(banddata[setdef2[:, k],:])), 1,dtype=np.float64)
                        tmpresult +=np.sum(np.real(banddata[setdef2[:, j],:]*np.conj(banddata[setdef1[:, k],:])), 1,dtype=np.float64)
                        tmpresult /= (2*nmodes)

                        allspectra_out[iprime, spectrum_idx, :]=tmpresult.astype(np.float32)

                    else:
                        #j/k are freqs
                        #first index is nrealizations
                        tmpresult=np.sum(np.real(banddata[setdef1[:, j],:]*np.conj(banddata[setdef1[:, k],:])), 1,dtype=np.float64) / (nmodes) # had been in flatsky: *reso^2*winsize^2)
                        #tmpresult is nrealizations long
                        allspectra_out[iprime, spectrum_idx, :]=tmpresult.astype(np.float32)
                    spectrum_idx+=1
                    #           pdb.set_trace()
        i=istop
    return(allspectra_out, nmodes_out)

def process_all_cross_spectra(allspectra, nbands, nsets,setsize,
                              auto=False,
                              skipcov=False ) -> 'Returns mean and covarariance estimates':

    print("Correlating Cross Spectra")
    nspectra = int( (nsets * (nsets+1))/2 + 0.001)

    if auto:
        nrealizations = setsize
    else:
        nrealizations=int( (setsize*(setsize-1))/2 + 0.001)

    allspectra = np.reshape(allspectra, [nbands*nspectra, nrealizations])
    #cov  = np.zeros([nbands*nspectra, nbands*nspectra],dtype=np.float64)

    spectrum = np.sum(allspectra,-1,dtype=np.float64)/nrealizations

    spectrum = np.reshape(spectrum,[nbands,nspectra])
    if skipcov:
        return spectrum,None,None,None
    # [nbands*nspectra, nrealizations])
    spectrum_2d = np.tile(np.reshape(spectrum,[nbands*nspectra,1]), [1,nrealizations])

    cov1 = np.matmul((allspectra-spectrum_2d) , (allspectra-spectrum_2d).T)
    cov1/= (nrealizations*(nrealizations-1))

    cov2 = None
    if not auto:
        realization_to_complement=np.zeros([nrealizations, setsize],dtype=np.float64)

        for i in range(setsize):
            realization_idx = 0
            for j in range(setsize):
                for k in range(j+1,setsize):
                    if (i == j) or (i == k):
                        realization_to_complement[realization_idx, i]=1./(setsize-1)
                    realization_idx += 1

        allcomplementspectra=np.matmul(allspectra,realization_to_complement)
        #nbands*nspectra, setsize)
        spectrum_2d=np.tile(np.reshape(spectrum,[nbands*nspectra,1]), [1,setsize])

        cov2=np.matmul( (allcomplementspectra-spectrum_2d), (allcomplementspectra-spectrum_2d).T )
        cov2/=(setsize**2 / 2)
        cov=2*cov2-cov1

    else:
        cov=cov1*(nrealizations)

    return spectrum,cov,cov1,cov2

'''
Create a class instance to simplify storing all the arguments along with the output
'''
class unbiased_multispec:
    def __init__(self,
                 # Maps/SHT flags ################################################
                 mapfile, #required -array of map filenames, g3 format
                 window, # required -- mask to apply for SHT
                 banddef, # required. [0,lmax_bin1, lmax_bin2, ...]
                 nside, #required. eg 8192
                 lmax=None, #optional, but should be set. Defaults to 2*nside
                 cmbweighting=True, # True ==> return Dl. False ==> Return Cl
                 kmask = None, #If not none, must be the right size for the Alms. A numpy array/vector
                 setdef=None, # optional -- will take from mapfile array dimensions if not provided
                 setdef2 = None, #optional -- if provided will assume doing sim cross-spectra
                 jackknife = False, #If true, will difference SHTs to do null spectrum
                 auto=False, #If true will do autospectra instead of cross-spectra
                 apply_windowfactor = True, #if true, calculate and apply normalization correction for partial sky mask.
                 map_key = 'T', #where to fetch maps from
                 skipcov=False, #don't calculate covariances
                 # Run time processing flags ################################################
                 ramlimit=64 * 2**30, # optional -- set to change default RAM limit from 64gb
                 resume=True, #optional -- will use existing files if true
                 basedir=None, # strongly suggested. defaults to current directory and can use a lot of disk space
                 persistdir=None, # optional - can be unset. will create a temp directory within basedir
                 remove_temporary_files= False, # optional. Defaults to off (user has to do cleanup, but can restart runs later)
                 verbose = False ): #extra print statements
                #maybe sometime I'll put in more input file arguments...
        '''
                 # Outputs ################################################
                 allspectra -- array of all cross-spectra (binned according to banddef)
                 cov -- array estimated covariance
                 est1_cov -- array estimated covariance from estimator 1
                 est2_cov -- array estimated covariance from estimator 2
                 nmodes -- array of number of alms per bandpower bin (form banddef)
                 windowfactor -- value used to normalize spectrum for apodization window. May be 1 (ie not corrected)
        '''
        self.mapfile = mapfile
        self.window = window
        self.banddef = banddef
        self.nside = nside
        self.lmax = lmax
        if self.lmax is None:
            self.lmax = 2*self.nside
        self.cmbweighting = cmbweighting
        self.kmask = kmask
        self.setdef = setdef
        self.jackknife = jackknife
        self.auto = auto
        self.apply_windowfactor = apply_windowfactor
        self.ramlimit = ramlimit
        self.resume = resume
        self.basedir = basedir
        self.persistdir = persistdir
        self.remove_temporary_files = remove_temporary_files
        self.verbose = verbose
        self.allspectra = None
        self.spectrum = None
        self.cov = None
        self.est1_cov = None
        self.est2_cov = None
        self.nmodes = None
        self.windowfactor = 1.0

        #################
        # figure out scratch directories
        #################
        try:
            if not os.path.isdir(self.basedir):
                raise TypeError  # to be caught below
        except TypeError:
            self.basedir = os.getcwd()
        try:
            if os.path.exists(self.persistdir) and not os.path.isdir(self.persistdir):
                print("WARNING -- Requested scratch exists, but is not a directory: {}".format(self.persistdir))
                raise TypeError
        except TypeError:
            self.persistdir = name_tempdir(self.basedir)
            print("WARNING -- using {} for scratch".format(self.persistdir))
            if not os.path.isdir(self.persistdir):
                os.makedirs(self.persistdir)

        #maybe at some point, we'll use status. right now nothing is done. Resume will only affect the full step level - no partial steps yet.
        status_file = self.persistdir + '/status.pkl'

        processed_sht_file = self.persistdir + '/shts_processed.bin'
        if not self.resume:
            try:
                os.remove(processed_sht_file)
            except FileNotFoundError:
                pass

        #################
        # Figure out set def based on structure of map file names, if not provided
        #################
        if self.setdef is None:
            #may need to change this -- unsure if it's right or transpose
            #may also need to make it 2d
            #remove warning printout when checked
            self.setdef = self.mapfile.shape
            print('Warning - check set def: inferred {}'.format(self.setdef))

        #get SHTs done
        sht_size = os.path.getsize(processed_sht_file)
        desired_size = healpy.sphtfunc.Alm.getsize(lmax) * np.zeros(1,dtype=AlmType).nbytes
        if (sht_size < desired_size):  #this will be false if resume==False since deleted file above.
            print("Dont expect to be here")
            pdb.set_trace()
            take_and_reformat_shts(self.fuile, processed_sht_file,
                   self.nside,self.lmax,
                   cmbweighting = self.cmbweighting,
                   mask  = self.window,
                   kmask = self.kmask,
                   ell_reordering=None,
                   no_reorder=False,
                   ram_limit = self.ramlimit,
                   map_key=map_key
                  )

        use_setdef  = setdef
        use_shtfile = processed_sht_file
        if self.jackknife:
            jackknife_sht_file = self.persistdir + '/null_shts_processed.bin'
            use_setdef = generate_jackknife_shts( processed_sht_file, jackknife_sht_file,  self.lmax, self.setdef)
            use_shtfile = jackknife_sht_file

        self.use_setdef = use_setdef

        #figure out cross-spectra (or autospectra)
        if setdef2 is None:
            allspectra, nmodes= take_all_cross_spectra( use_shtfile, self.lmax,
                                                        self.use_setdef, self.banddef,  ram_limit=self.ramlimit, auto = self.auto) #-> 'Returns set of all x-spectra, binned':
        else:
            allspectra, nmodes= take_all_sim_cross_spectra( use_shtfile, self.lmax,
                                                        self.use_setdef,self.banddef, setdef2=setdef2, ram_limit=self.ramlimit, auto = self.auto) #-> 'Returns set of all x-spectra, binned':

        self.allspectra = allspectra
        self.nmodes = nmodes

        #bring it all together
        nbands = banddef.shape[0]-1
        nsets   = use_setdef.shape[1]
        setsize = use_setdef.shape[0]

        process_auto = self.auto or (setdef2 is not None) # only get 1 per set for the sim crosses too
        spectrum,cov,cov1,cov2 = process_all_cross_spectra(self.allspectra, nbands, nsets,setsize,
                                                            auto=process_auto)
        self.spectrum = spectrum
        self.cov      = cov
        self.est1_cov = cov1
        self.est2_cov = cov2


"""
The ExpErrorgenOp class and supporting functionality.
"""
#***************************************************************************************************
# Copyright 2015, 2019 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
# Under the terms of Contract DE-NA0003525 with NTESS, the U.S. Government retains certain rights
# in this software.
# 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 or in the LICENSE file in the root pyGSTi directory.
#***************************************************************************************************

import warnings as _warnings

import numpy as _np
import scipy.linalg as _spl
import scipy.sparse as _sps
import scipy.sparse.linalg as _spsl

from pygsti.modelmembers.operations.linearop import LinearOperator as _LinearOperator
from pygsti.modelmembers.operations.lindbladerrorgen import LindbladParameterization as _LindbladParameterization
from pygsti.modelmembers import modelmember as _modelmember, term as _term
from pygsti.modelmembers.errorgencontainer import ErrorGeneratorContainer as _ErrorGeneratorContainer
from pygsti.baseobjs.polynomial import Polynomial as _Polynomial
from pygsti.tools import matrixtools as _mt

IMAG_TOL = 1e-7  # tolerance for imaginary part being considered zero
MAX_EXPONENT = _np.log(_np.finfo('d').max) - 10.0  # so that exp(.) doesn't overflow
TODENSE_TRUNCATE = 1e-11

class ExpErrorgenOp(_LinearOperator, _ErrorGeneratorContainer):
    """
    An operation parameterized by the coefficients of an exponentiated sum of Lindblad-like terms.
    TODO: update docstring!

    The exponentiated terms give the operation's action.

    Parameters
    ----------
    errorgen : LinearOperator
        The error generator for this operator.  That is, the `L` if this
        operator is `exp(L)`.
    """

    def __init__(self, errorgen):
        # Extract superop dimension from 'errorgen'
        state_space = errorgen.state_space
        self.errorgen = errorgen  # don't copy (allow object reuse)

        evotype = self.errorgen._evotype

        #Create representation object
        rep_type_order = ('dense', 'experrgen') if evotype.prefer_dense_reps else ('experrgen', 'dense')
        rep = None
        for rep_type in rep_type_order:
            try:
                if rep_type == 'experrgen':
                    # "sparse mode" => don't ever compute matrix-exponential explicitly
                    rep = evotype.create_experrorgen_rep(self.errorgen._rep)
                elif rep_type == 'dense':
                    rep = evotype.create_dense_superop_rep(None, state_space)

                    # Cache values - for later work with dense rep
                    self.exp_err_gen = None   # used for dense_rep=True mode to cache qty needed in deriv_wrt_params
                    self.base_deriv = None
                    self.base_hessian = None
                else:
                    assert(False), "Logic error!"

                self._rep_type = rep_type
                break

            except AttributeError:
                pass  # just go to the next rep_type

        if rep is None:
            raise ValueError("Unable to construct representation with evotype: %s" % str(evotype))

        # Caches in case terms are used
        self.terms = {}
        self.exp_terms_cache = {}  # used for repeated calls to the exponentiate_terms function
        self.local_term_poly_coeffs = {}

        _LinearOperator.__init__(self, rep, evotype)
        _ErrorGeneratorContainer.__init__(self, self.errorgen)
        self.init_gpindices()  # initialize our gpindices based on sub-members
        self._update_rep()  # updates self._rep
        #Done with __init__(...)

    #Note: no to_memoized_dict needed, as ModelMember version does all we need.

    @classmethod
    def _from_memoized_dict(cls, mm_dict, serial_memo):
        errorgen = serial_memo[mm_dict['submembers'][0]]
        return cls(errorgen)

    def submembers(self):
        """
        Get the ModelMember-derived objects contained in this one.

        Returns
        -------
        list
        """
        return [self.errorgen]

    def _update_rep(self, close=False):
        """
        Updates self._rep as needed after parameters have changed.
        """
        if self._rep_type == 'dense':
            # compute matrix-exponential explicitly
            self.exp_err_gen = _spl.expm(self.errorgen.to_dense(on_space='HilbertSchmidt'))  # used in deriv_wrt_params

            dense = self.exp_err_gen
            self._rep.base.flags.writeable = True
            self._rep.base[:, :] = dense
            self._rep.base.flags.writeable = False
            self.base_deriv = None
            self.base_hessian = None
        else:  # if not close:
            self._rep.errgenrep_has_changed(self.errorgen.onenorm_upperbound())

            #CHECK that sparsemx action is correct (DEBUG CHECK)
            #from pygsti.modelmembers.states import StaticState
            #Mdense = _spl.expm(self.errorgen.to_dense())
            #if Mdense.shape == (4,4):
            #    for i in range(4):
            #        v = _np.zeros(4); v[i] = 1.0
            #
            #        staterep = StaticState(v)._rep
            #        check_acton = self._rep.acton(staterep).data
            #
            #        #check_sparse_scipy = _spsl.expm_multiply(self.errorgen.to_sparse(), v.copy())
            #        prep = _mt.expm_multiply_prep(self.errorgen.to_sparse())
            #        check_sparse = _mt.expm_multiply_fast(prep, v)
            #        check_dense = _np.dot(Mdense, v)
            #
            #        diff = _np.linalg.norm(check_dense - check_acton)
            #        #diff2 = _np.linalg.norm(check_sparse_scipy - check_sparse)
            #        if diff > 1e-6: # or diff2 > 1e-3:
            #            print("PROBLEM (%d)!!" % i, " Expop diff = ", diff)

    def set_gpindices(self, gpindices, parent, memo=None):
        """
        Set the parent and indices into the parent's parameter vector that are used by this ModelMember object.

        Parameters
        ----------
        gpindices : slice or integer ndarray
            The indices of this objects parameters in its parent's array.

        parent : Model or ModelMember
            The parent whose parameter array gpindices references.

        memo : dict, optional
            A memo dict used to avoid circular references.

        Returns
        -------
        None
        """
        _modelmember.ModelMember.set_gpindices(self, gpindices, parent, memo)
        self.terms = {}  # clear terms cache since param indices have changed now
        self.exp_terms_cache = {}
        self.local_term_poly_coeffs = {}

    def to_dense(self, on_space='minimal'):
        """
        Return this operation as a dense matrix.

        Returns
        -------
        numpy.ndarray
        """
        if self._rep_type == 'dense':
            # Then self._rep contains a dense version already
            return self._rep.base  # copy() unnecessary since we set to readonly

        else:
            # Construct a dense version from scratch (more time consuming)
            return _spl.expm(self.errorgen.to_dense(on_space))

    #FUTURE: maybe remove this function altogether, as it really shouldn't be called
    def to_sparse(self, on_space='minimal'):
        """
        Return the operation as a sparse matrix.

        Parameters
        ----------
        on_space : {'minimal', 'Hilbert', 'HilbertSchmidt'}
            The space that the returned dense operation acts upon.  For unitary matrices and bra/ket vectors,
            use `'Hilbert'`.  For superoperator matrices and super-bra/super-ket vectors use `'HilbertSchmidt'`.
            `'minimal'` means that `'Hilbert'` is used if possible given this operator's evolution type, and
            otherwise `'HilbertSchmidt'` is used.

        Returns
        -------
        scipy.sparse.csr_matrix
        """
        if self._rep_type == 'dense':
            return _sps.csr_matrix(self.to_dense(on_space))
        else:
            return _spsl.expm(self.errorgen.to_sparse(on_space).tocsc()).tocsr()

    def deriv_wrt_params(self, wrt_filter=None):
        """
        The element-wise derivative this operation.

        Construct a matrix whose columns are the vectorized
        derivatives of the flattened operation matrix with respect to a
        single operation parameter.  Thus, each column is of length
        op_dim^2 and there is one column per operation parameter.

        Parameters
        ----------
        wrt_filter : list or numpy.ndarray
            List of parameter indices to take derivative with respect to.
            (None means to use all the this operation's parameters.)

        Returns
        -------
        numpy array
            Array of derivatives, shape == (dimension^2, num_params)
        """
        if not self._rep_type == 'dense':
            #raise NotImplementedError("deriv_wrt_params(...) can only be used when a dense representation is used!")
            #_warnings.warn("Using finite differencing to compute ExpErrogenOp derivative!")
            return super(ExpErrorgenOp, self).deriv_wrt_params(wrt_filter)

        if self.base_deriv is None:
            d2 = self.dim

            #Deriv wrt hamiltonian params
            derrgen = self.errorgen.deriv_wrt_params(None)  # apply filter below; cache *full* deriv
            derrgen.shape = (d2, d2, -1)  # separate 1st d2**2 dim to (d2,d2)
            dexpL = _d_exp_x(self.errorgen.to_dense(on_space='minimal'), derrgen, self.exp_err_gen)
            derivMx = dexpL.reshape(d2**2, self.num_params)  # [iFlattenedOp,iParam]

            assert(_np.linalg.norm(_np.imag(derivMx)) < IMAG_TOL), \
                ("Deriv matrix has imaginary part = %s.  This can result from "
                 "evaluating a Model derivative at a 'bad' point where the "
                 "error generator is large.  This often occurs when GST's "
                 "starting Model has *no* stochastic error and all such "
                 "parameters affect error rates at 2nd order.  Try "
                 "depolarizing the seed Model.") % str(_np.linalg.norm(_np.imag(derivMx)))
            # if this fails, uncomment around "DB COMMUTANT NORM" for further debugging.
            derivMx = _np.real(derivMx)
            self.base_deriv = derivMx

            #check_deriv_wrt_params(self, derivMx, eps=1e-7)
            #fd_deriv = finite_difference_deriv_wrt_params(self, wrt_filter, eps=1e-7)
            #derivMx = fd_deriv

        if wrt_filter is None:
            return self.base_deriv.view()
            #view because later setting of .shape by caller can mess with self.base_deriv!
        else:
            return _np.take(self.base_deriv, wrt_filter, axis=1)

    def has_nonzero_hessian(self):
        """
        Whether this operation has a non-zero Hessian with respect to its parameters.

        (i.e. whether it only depends linearly on its parameters or not)

        Returns
        -------
        bool
        """
        return True

    def hessian_wrt_params(self, wrt_filter1=None, wrt_filter2=None):
        """
        Construct the Hessian of this operation with respect to its parameters.

        This function returns a tensor whose first axis corresponds to the
        flattened operation matrix and whose 2nd and 3rd axes correspond to the
        parameters that are differentiated with respect to.

        Parameters
        ----------
        wrt_filter1 : list or numpy.ndarray
            List of parameter indices to take 1st derivatives with respect to.
            (None means to use all the this operation's parameters.)

        wrt_filter2 : list or numpy.ndarray
            List of parameter indices to take 2nd derivatives with respect to.
            (None means to use all the this operation's parameters.)

        Returns
        -------
        numpy array
            Hessian with shape (dimension^2, num_params1, num_params2)
        """
        if not self._rep_type == 'dense':
            #raise NotImplementedError("hessian_wrt_params is only implemented for *dense-rep* LindbladOps")
            #_warnings.warn("Using finite differencing to compute ExpErrogenOp Hessian!")
            return super(ExpErrorgenOp, self).hessian_wrt_params(wrt_filter1, wrt_filter2)

        if self.base_hessian is None:
            d2 = self.dim
            nP = self.num_params
            hessianMx = _np.zeros((d2**2, nP, nP), 'd')

            #Deriv wrt other params
            dEdp = self.errorgen.deriv_wrt_params(None)  # filter later, cache *full*
            d2Edp2 = self.errorgen.hessian_wrt_params(None, None)  # hessian
            dEdp.shape = (d2, d2, nP)  # separate 1st d2**2 dim to (d2,d2)
            d2Edp2.shape = (d2, d2, nP, nP)  # ditto

            series, series2 = _d2_exp_series(self.errorgen.to_dense(on_space='minimal'), dEdp, d2Edp2)
            term1 = series2
            term2 = _np.einsum("ija,jkq->ikaq", series, series)
            d2expL = _np.einsum("ikaq,kj->ijaq", term1 + term2,
                                self.exp_err_gen)
            hessianMx = d2expL.reshape((d2**2, nP, nP))

            #hessian has been made so index as [iFlattenedOp,iDeriv1,iDeriv2]
            assert(_np.linalg.norm(_np.imag(hessianMx)) < IMAG_TOL)
            hessianMx = _np.real(hessianMx)  # d2O block of hessian

            self.base_hessian = hessianMx

            #TODO: check hessian with finite difference here?

        if wrt_filter1 is None:
            if wrt_filter2 is None:
                return self.base_hessian.view()
                #view because later setting of .shape by caller can mess with self.base_hessian!
            else:
                return _np.take(self.base_hessian, wrt_filter2, axis=2)
        else:
            if wrt_filter2 is None:
                return _np.take(self.base_hessian, wrt_filter1, axis=1)
            else:
                return _np.take(_np.take(self.base_hessian, wrt_filter1, axis=1),
                                wrt_filter2, axis=2)

    @property
    def parameter_labels(self):
        """
        An array of labels (usually strings) describing this model member's parameters.
        """
        return self.errorgen.parameter_labels

    @property
    def num_params(self):
        """
        Get the number of independent parameters which specify this operation.

        Returns
        -------
        int
            the number of independent parameters.
        """
        return self.errorgen.num_params

    def to_vector(self):
        """
        Extract a vector of the underlying operation parameters from this operation.

        Returns
        -------
        numpy array
            a 1D numpy array with length == num_params().
        """
        return self.errorgen.to_vector()

    def from_vector(self, v, close=False, dirty_value=True):
        """
        Initialize the operation using a vector of parameters.

        Parameters
        ----------
        v : numpy array
            The 1D vector of operation parameters.  Length
            must == num_params()

        close : bool, optional
            Whether `v` is close to this operation's current
            set of parameters.  Under some circumstances, when this
            is true this call can be completed more quickly.

        dirty_value : bool, optional
            The value to set this object's "dirty flag" to before exiting this
            call.  This is passed as an argument so it can be updated *recursively*.
            Leave this set to `True` unless you know what you're doing.

        Returns
        -------
        None
        """
        self.errorgen.from_vector(v, close, dirty_value)
        self._update_rep(close)
        self.dirty = dirty_value

    def taylor_order_terms(self, order, max_polynomial_vars=100, return_coeff_polys=False):
        """
        Get the `order`-th order Taylor-expansion terms of this operation.

        This function either constructs or returns a cached list of the terms at
        the given order.  Each term is "rank-1", meaning that its action on a
        density matrix `rho` can be written:

        `rho -> A rho B`

        The coefficients of these terms are typically polynomials of the operation's
        parameters, where the polynomial's variable indices index the *global*
        parameters of the operation's parent (usually a :class:`Model`), not the
        operation's local parameter array (i.e. that returned from `to_vector`).

        Parameters
        ----------
        order : int
            Which order terms (in a Taylor expansion of this :class:`LindbladOp`)
            to retrieve.

        max_polynomial_vars : int, optional
            maximum number of variables the created polynomials can have.

        return_coeff_polys : bool
            Whether a parallel list of locally-indexed (using variable indices
            corresponding to *this* object's parameters rather than its parent's)
            polynomial coefficients should be returned as well.

        Returns
        -------
        terms : list
            A list of :class:`RankOneTerm` objects.
        coefficients : list
            Only present when `return_coeff_polys == True`.
            A list of *compact* polynomial objects, meaning that each element
            is a `(vtape,ctape)` 2-tuple formed by concatenating together the
            output of :method:`Polynomial.compact`.
        """
        if order not in self.terms:
            self._compute_taylor_order_terms(order, max_polynomial_vars)

        if return_coeff_polys:
            return self.terms[order], self.local_term_poly_coeffs[order]
        else:
            return self.terms[order]

    def _compute_taylor_order_terms(self, order, max_polynomial_vars):  # separated for profiling

        mapvec = _np.ascontiguousarray(_np.zeros(max_polynomial_vars, _np.int64))
        for ii, i in enumerate(self.gpindices_as_array()):
            mapvec[ii] = i

        def _compose_poly_indices(terms):
            for term in terms:
                #term.map_indices_inplace(lambda x: tuple(_modelmember._compose_gpindices(
                #    self.gpindices, _np.array(x, _np.int64))))
                term.mapvec_indices_inplace(mapvec)
            return terms

        assert(self.gpindices is not None), "LindbladOp must be added to a Model before use!"
        mpv = max_polynomial_vars

        #Note: for now, *all* of an error generator's terms are considered 0-th order,
        # so the below call to taylor_order_terms just gets all of them.  In the FUTURE
        # we might want to allow a distinction among the error generator terms, in which
        # case this term-exponentiation step will need to become more complicated...
        postTerm = _term.RankOnePolynomialOpTerm.create_from(_Polynomial({(): 1.0}, mpv),
                                                             None, None, self._evotype, self.state_space)  # identity
        loc_terms = _term.exponentiate_terms(self.errorgen.taylor_order_terms(0, max_polynomial_vars),
                                             order, postTerm, self.exp_terms_cache)
        #OLD: loc_terms = [ t.collapse() for t in loc_terms ] # collapse terms for speed

        poly_coeffs = [t.coeff for t in loc_terms]
        tapes = [poly.compact(complex_coeff_tape=True) for poly in poly_coeffs]
        if len(tapes) > 0:
            vtape = _np.concatenate([t[0] for t in tapes])
            ctape = _np.concatenate([t[1] for t in tapes])
        else:
            vtape = _np.empty(0, _np.int64)
            ctape = _np.empty(0, complex)
        coeffs_as_compact_polys = (vtape, ctape)
        self.local_term_poly_coeffs[order] = coeffs_as_compact_polys

        # only cache terms with *global* indices to avoid confusion...
        self.terms[order] = _compose_poly_indices(loc_terms)

    def taylor_order_terms_above_mag(self, order, max_polynomial_vars, min_term_mag):
        """
        Get the `order`-th order Taylor-expansion terms of this operation that have magnitude above `min_term_mag`.

        This function constructs the terms at the given order which have a magnitude (given by
        the absolute value of their coefficient) that is greater than or equal to `min_term_mag`.
        It calls :method:`taylor_order_terms` internally, so that all the terms at order `order`
        are typically cached for future calls.

        The coefficients of these terms are typically polynomials of the operation's
        parameters, where the polynomial's variable indices index the *global*
        parameters of the operation's parent (usually a :class:`Model`), not the
        operation's local parameter array (i.e. that returned from `to_vector`).

        Parameters
        ----------
        order : int
            The order of terms to get (and filter).

        max_polynomial_vars : int, optional
            maximum number of variables the created polynomials can have.

        min_term_mag : float
            the minimum term magnitude.

        Returns
        -------
        list
            A list of :class:`Rank1Term` objects.
        """
        mapvec = _np.ascontiguousarray(_np.zeros(max_polynomial_vars, _np.int64))
        for ii, i in enumerate(self.gpindices_as_array()):
            mapvec[ii] = i

        assert(self.gpindices is not None), "LindbladOp must be added to a Model before use!"
        mpv = max_polynomial_vars

        postTerm = _term.RankOnePolynomialOpTerm.create_from(_Polynomial({(): 1.0}, mpv), None, None,
                                                             self._evotype, self.state_space)  # identity term
        postTerm = postTerm.copy_with_magnitude(1.0)
        #Note: for now, *all* of an error generator's terms are considered 0-th order,
        # so the below call to taylor_order_terms just gets all of them.  In the FUTURE
        # we might want to allow a distinction among the error generator terms, in which
        # case this term-exponentiation step will need to become more complicated...
        errgen_terms = self.errorgen.taylor_order_terms(0, max_polynomial_vars)

        #DEBUG CHECK MAGS OF ERRGEN COEFFS
        #poly_coeffs = [t.coeff for t in errgen_terms]
        #tapes = [poly.compact(complex_coeff_tape=True) for poly in poly_coeffs]
        #if len(tapes) > 0:
        #    vtape = _np.concatenate([t[0] for t in tapes])
        #    ctape = _np.concatenate([t[1] for t in tapes])
        #else:
        #    vtape = _np.empty(0, _np.int64)
        #    ctape = _np.empty(0, complex)
        #v = self.to_vector()
        #errgen_coeffs = _bulk_eval_compact_polynomials_complex(
        #    vtape, ctape, v, (len(errgen_terms),))  # an array of coeffs
        #for coeff, t in zip(errgen_coeffs, errgen_terms):
        #    coeff2 = t.coeff.evaluate(v)
        #    if not _np.isclose(coeff,coeff2):
        #        assert(False), "STOP"
        #    t.set_magnitude(abs(coeff))

        #evaluate errgen_terms' coefficients using their local vector of parameters
        # (which happends to be the same as our paramvec in this case)
        egvec = self.errorgen.to_vector()   # we need errorgen's vector (usually not in rep) to perform evaluation
        errgen_terms = [egt.copy_with_magnitude(abs(egt.coeff.evaluate(egvec))) for egt in errgen_terms]

        terms = []
        for term in _term.exponentiate_terms_above_mag(errgen_terms, order,
                                                       postTerm, min_term_mag=min_term_mag):
            #poly_coeff = term.coeff
            #compact_poly_coeff = poly_coeff.compact(complex_coeff_tape=True)
            term.mapvec_indices_inplace(mapvec)  # local -> global indices

            # DEBUG CHECK - to ensure term magnitudes are being set correctly (i.e. are in sync with evaluated coeffs)
            # t = term
            # vt, ct = t._rep.coeff.compact_complex()
            # coeff_array = _bulk_eval_compact_polynomials_complex(vt, ct, self.parent.to_vector(), (1,))
            # if not _np.isclose(abs(coeff_array[0]), t._rep.magnitude):  # DEBUG!!!
            #     print(coeff_array[0], "vs.", t._rep.magnitude)
            #     import bpdb; bpdb.set_trace()
            #     c1 = _Polynomial.from_rep(t._rep.coeff)

            terms.append(term)
        return terms

    @property
    def total_term_magnitude(self):
        """
        Get the total (sum) of the magnitudes of all this operator's terms.

        The magnitude of a term is the absolute value of its coefficient, so
        this function returns the number you'd get from summing up the
        absolute-coefficients of all the Taylor terms (at all orders!) you
        get from expanding this operator in a Taylor series.

        Returns
        -------
        float
        """
        # return exp( mag of errorgen ) = exp( sum of absvals of errgen term coeffs )
        # (unitary postfactor has weight == 1.0 so doesn't enter)
        return _np.exp(min(self.errorgen.total_term_magnitude, MAX_EXPONENT))
        #return _np.exp(self.errorgen.total_term_magnitude)  # overflows sometimes

    @property
    def total_term_magnitude_deriv(self):
        """
        The derivative of the sum of *all* this operator's terms.

        Computes the derivative of the total (sum) of the magnitudes of all this
        operator's terms with respect to the operators (local) parameters.

        Returns
        -------
        numpy array
            An array of length self.num_params
        """
        return _np.exp(self.errorgen.total_term_magnitude) * self.errorgen.total_term_magnitude_deriv

    def set_dense(self, m):
        """
        Set the dense-matrix value of this operation.

        Attempts to modify operation parameters so that the specified raw
        operation matrix becomes mx.  Will raise ValueError if this operation
        is not possible.

        Parameters
        ----------
        m : array_like or LinearOperator
            An array of shape (dim, dim) or LinearOperator representing the operation action.

        Returns
        -------
        None
        """
        mx = _LinearOperator.convert_to_matrix(m)
        errgen_cls = self.errorgen.__class__

        #Note: this only really works for LindbladErrorGen objects now... make more general in FUTURE?
        truncate = TODENSE_TRUNCATE  # can't just be 'True' since we need to throw errors when appropriate
        new_errgen = errgen_cls.from_operation_matrix_and_blocks(
            mx, self.errorgen.coefficient_blocks, 'auto', self.errorgen.matrix_basis,
            truncate, self.errorgen.evotype, self.errorgen.state_space)
        self.errorgen.from_vector(new_errgen.to_vector())
        self._update_rep()  # needed to rebuild exponentiated error gen
        self.dirty = True

    def transform_inplace(self, s):
        """
        Update operation matrix `O` with `inv(s) * O * s`.

        Generally, the transform function updates the *parameters* of
        the operation such that the resulting operation matrix is altered as
        described above.  If such an update cannot be done (because
        the operation parameters do not allow for it), ValueError is raised.

        Parameters
        ----------
        s : GaugeGroupElement
            A gauge group element which specifies the "s" matrix
            (and it's inverse) used in the above similarity transform.

        Returns
        -------
        None
        """
        #assert(_np.allclose(U, _np.linalg.inv(Uinv)))
        #just conjugate postfactor and Lindbladian exponent by U:
        self.errorgen.transform_inplace(s)
        self._update_rep()  # needed to rebuild exponentiated error gen
        self.dirty = True

    def spam_transform_inplace(self, s, typ):
        """
        Update operation matrix `O` with `inv(s) * O` OR `O * s`, depending on the value of `typ`.

        This functions as `transform_inplace(...)` but is used when this
        operation is used as a part of a SPAM vector.  When `typ == "prep"`,
        the spam vector is assumed to be `rho = dot(self, <spamvec>)`,
        which transforms as `rho -> inv(s) * rho`, so `self -> inv(s) * self`.
        When `typ == "effect"`, `e.dag = dot(e.dag, self)` (note that
        `self` is NOT `self.dag` here), and `e.dag -> e.dag * s`
        so that `self -> self * s`.

        Parameters
        ----------
        s : GaugeGroupElement
            A gauge group element which specifies the "s" matrix
            (and it's inverse) used in the above similarity transform.

        typ : { 'prep', 'effect' }
            Which type of SPAM vector is being transformed (see above).

        Returns
        -------
        None
        """
        assert(typ in ('prep', 'effect')), "Invalid `typ` argument: %s" % typ
        from pygsti.models import gaugegroup as _gaugegroup

        if isinstance(s, _gaugegroup.UnitaryGaugeGroupElement) \
           or isinstance(s, _gaugegroup.TPSpamGaugeGroupElement):
            U = s.transform_matrix
            Uinv = s.transform_matrix_inverse
            mx = self.to_dense(on_space='minimal') if self._rep_type == 'dense' else self.to_sparse(on_space='minimal')

            #just act on postfactor and Lindbladian exponent:
            if typ == "prep":
                mx = _mt.safe_dot(Uinv, mx)
            else:
                mx = _mt.safe_dot(mx, U)
            self.set_dense(mx)  # calls _update_rep() and sets dirty flag

    def __str__(self):
        s = "Exponentiated operation map with dim = %d, num params = %d\n" % \
            (self.dim, self.num_params)
        return s

    def _oneline_contents(self):
        """ Summarizes the contents of this object in a single line.  Does not summarize submembers. """
        return "exponentiates"

def _d_exp_series(x, dx):
    TERM_TOL = 1e-12
    tr = len(dx.shape)  # tensor rank of dx; tr-2 == # of derivative dimensions
    assert((tr - 2) in (1, 2)), "Currently, dx can only have 1 or 2 derivative dimensions"
    #assert( len( (_np.isnan(dx)).nonzero()[0] ) == 0 ) # NaN debugging
    #assert( len( (_np.isnan(x)).nonzero()[0] ) == 0 ) # NaN debugging
    series = dx.copy()  # accumulates results, so *need* a separate copy
    last_commutant = term = dx; i = 2

    #take d(matrix-exp) using series approximation
    while _np.amax(_np.abs(term)) > TERM_TOL:  # _np.linalg.norm(term)
        if tr == 3:
            #commutant = _np.einsum("ik,kja->ija",x,last_commutant) - \
            #            _np.einsum("ika,kj->ija",last_commutant,x)
            commutant = _np.tensordot(x, last_commutant, (1, 0)) - \
                _np.transpose(_np.tensordot(last_commutant, x, (1, 0)), (0, 2, 1))
        elif tr == 4:
            #commutant = _np.einsum("ik,kjab->ijab",x,last_commutant) - \
            #        _np.einsum("ikab,kj->ijab",last_commutant,x)
            commutant = _np.tensordot(x, last_commutant, (1, 0)) - \
                _np.transpose(_np.tensordot(last_commutant, x, (1, 0)), (0, 3, 1, 2))

        term = 1 / _np.math.factorial(i) * commutant

        #Uncomment some/all of this when you suspect an overflow due to x having large norm.
        #print("DB COMMUTANT NORM = ",_np.linalg.norm(commutant)) # sometimes this increases w/iter -> divergence => NaN
        #assert(not _np.isnan(_np.linalg.norm(term))), \
        #    ("Haddamard series = NaN! Probably due to trying to differentiate "
        #     "exp(x) where x has a large norm!")

        #DEBUG
        #if not _np.isfinite(_np.linalg.norm(term)): break # DEBUG high values -> overflow for nqubit operations
        #if len( (_np.isnan(term)).nonzero()[0] ) > 0: # NaN debugging
        #    #WARNING: stopping early b/c of NaNs!!! - usually caused by infs
        #    break

        series += term  # 1/_np.math.factorial(i) * commutant
        last_commutant = commutant; i += 1
    return series

def _d2_exp_series(x, dx, d2x):
    TERM_TOL = 1e-12
    tr = len(dx.shape)  # tensor rank of dx; tr-2 == # of derivative dimensions
    tr2 = len(d2x.shape)  # tensor rank of dx; tr-2 == # of derivative dimensions
    assert((tr - 2, tr2 - 2) in [(1, 2), (2, 4)]), "Current support for only 1 or 2 derivative dimensions"

    series = dx.copy()  # accumulates results, so *need* a separate copy
    series2 = d2x.copy()  # accumulates results, so *need* a separate copy
    term = last_commutant = dx
    last_commutant2 = term2 = d2x
    i = 2

    #take d(matrix-exp) using series approximation
    while _np.amax(_np.abs(term)) > TERM_TOL or _np.amax(_np.abs(term2)) > TERM_TOL:
        if tr == 3:
            commutant = _np.einsum("ik,kja->ija", x, last_commutant) - \
                _np.einsum("ika,kj->ija", last_commutant, x)
            commutant2A = _np.einsum("ikq,kja->ijaq", dx, last_commutant) - \
                _np.einsum("ika,kjq->ijaq", last_commutant, dx)
            commutant2B = _np.einsum("ik,kjaq->ijaq", x, last_commutant2) - \
                _np.einsum("ikaq,kj->ijaq", last_commutant2, x)

        elif tr == 4:
            commutant = _np.einsum("ik,kjab->ijab", x, last_commutant) - \
                _np.einsum("ikab,kj->ijab", last_commutant, x)
            commutant2A = _np.einsum("ikqr,kjab->ijabqr", dx, last_commutant) - \
                _np.einsum("ikab,kjqr->ijabqr", last_commutant, dx)
            commutant2B = _np.einsum("ik,kjabqr->ijabqr", x, last_commutant2) - \
                _np.einsum("ikabqr,kj->ijabqr", last_commutant2, x)

        term = 1 / _np.math.factorial(i) * commutant
        term2 = 1 / _np.math.factorial(i) * (commutant2A + commutant2B)
        series += term
        series2 += term2
        last_commutant = commutant
        last_commutant2 = (commutant2A + commutant2B)
        i += 1
    return series, series2

def _d_exp_x(x, dx, exp_x=None):
    """
    Computes the derivative of the exponential of x(t) using
    the Haddamard lemma series expansion.

    Parameters
    ----------
    x : ndarray
        The 2-tensor being exponentiated

    dx : ndarray
        The derivative of x; can be either a 3- or 4-tensor where the
        3rd+ dimensions are for (multi-)indexing the parameters which
        are differentiated w.r.t.  For example, in the simplest case
        dx is a 3-tensor s.t. dx[i,j,p] == d(x[i,j])/dp.

    exp_x : ndarray, optional
        The value of `exp(x)`, which can be specified in order to save
        a call to `scipy.linalg.expm`.  If None, then the value is
        computed internally.

    Returns
    -------
    ndarray
        The derivative of `exp(x)` given as a tensor with the
        same shape and axes as `dx`.
    """
    tr = len(dx.shape)  # tensor rank of dx; tr-2 == # of derivative dimensions
    assert((tr - 2) in (1, 2)), "Currently, dx can only have 1 or 2 derivative dimensions"

    series = _d_exp_series(x, dx)
    if exp_x is None: exp_x = _spl.expm(x)

    if tr == 3:
        #dExpX = _np.einsum('ika,kj->ija', series, exp_x)
        dExpX = _np.transpose(_np.tensordot(series, exp_x, (1, 0)), (0, 2, 1))
    elif tr == 4:
        #dExpX = _np.einsum('ikab,kj->ijab', series, exp_x)
        dExpX = _np.transpose(_np.tensordot(series, exp_x, (1, 0)), (0, 3, 1, 2))

    return dExpX