import collections
from typing import Callable, Iterable, List

import numpy as np
import torch
from numpy import ceil
from ppq.core import (CHECKPOINT_TOLERANCE, OPTIM_ADVOPT_GRAPH_MAXSIZE,
                      OPTIM_ADVOPT_INITIAL_THRES, OPTIM_ADVOPT_PASSIVE_BOOST,
                      OPTIM_ADVOPT_PATIENT, OPTIM_ADVOPT_STEP_PER_EPOCH,
                      OPTIM_ADVOPT_THRESHOLD_STEP, QuantizationProperty,
                      QuantizationStates, empty_ppq_cache)
from ppq.executor import BaseGraphExecutor, TorchExecutor
from ppq.IR import (BaseGraph, GraphCommandProcesser, Operation,
                    QuantableOperation)
from ppq.quantization.algorithm.training import (AdaroundRegTerm,
                                                 FinetuneCheckPoint, Lp_norm,
                                                 RandomMemDataset,
                                                 TrainableDelegate)
from ppq.quantization.analyise.util import MeasurePrinter
from ppq.quantization.measure import torch_mean_square_error, torch_snr_error
from torch.cuda import empty_cache
from tqdm import tqdm

from .base import QuantizationOptimizationPass

def has_bias(op: Operation):
    if op.type in {'Conv', 'ConvTranspose', 'Gemm'}:
        return op.meta_data.num_of_input == 3
    else: return False

class TrainingBasedPass(QuantizationOptimizationPass):
    """
    Training Based Pass is a basic class that provides necessary function for
        all training optimizition passes. Optimization will be more stable and
        accurate with functions provided by this pass. (Might be a little slower).

    This pass will collect result of interested outputs after optimization and
        check if the optimized result has a lower SNR. If so, the optimization will be
        accepted, layer weight will be updated, otherwise optimization will be rejected and
        takes no effects.

    Choose interested_outputs carefully, cause we compare loss only with those output variables.
        If interested_outputs is None, all graph output variables will be choosen.

    YOUR SHOULD NOTICE THAT SNR REFERS TO: POWER OF NOISE / POWER OF SIGNAL IN PPQ.

    Args:
        QuantizationOptimizationPass ([type]): [description]
    """
    def __init__(self, name: str = 'Default Quanzation Optim',
                 interested_outputs: List[str] = None, verbose: bool = True) -> None:
        self._loss_fn = torch_snr_error
        self._interested_outputs = interested_outputs
        self._checkpoints = {}
        self._verbose = verbose
        self._quant_state_recorder = {}
        super().__init__(name=name)

    @ empty_ppq_cache
    def initialize_checkpoints(
        self, graph: BaseGraph, executor: BaseGraphExecutor,
        dataloader: Iterable, collate_fn: Callable):
        """
        Establish a series of network checkpoints with your network.
            Checkpoint is a data structure that helps us compare quant results and fp32 results.
        Args:
            graph (BaseGraph): [description]
            executor (BaseGraphExecutor): [description]
            dataloader (Iterable): [description]
            collate_fn (Callable): [description]

        Raises:
            PermissionError: [description]
        """
        for operation in graph.operations.values():
            if isinstance(operation, QuantableOperation):
                for cfg, var in operation.config_with_variable:
                    if cfg.state in {QuantizationStates.BAKED, QuantizationStates.PASSIVE_BAKED}:
                        raise PermissionError('Can not initialize checkpoints when weight value is baked. '
                                              f'Variable {var.name} has a baked value.')

        if self._interested_outputs is None or len(self._interested_outputs) == 0:
            self._interested_outputs = [name for name in graph.outputs]

        for name in self._interested_outputs:
            self._checkpoints[name] = FinetuneCheckPoint(variable=name)

        # dequantize graph, collect references
        for op in graph.operations.values():
            if isinstance(op, QuantableOperation):
                op.dequantize()

        for data in dataloader:
            if collate_fn is not None: data = collate_fn(data)
            outputs = executor.forward(inputs=data, output_names=self._interested_outputs)
            for name, output in zip(self._interested_outputs, outputs):
                ckpt = self._checkpoints[name]
                assert isinstance(ckpt, FinetuneCheckPoint)
                ckpt.push(tensor=output, is_reference=True)

        # restore quantization state:
        for op in graph.operations.values():
            if isinstance(op, QuantableOperation):
                op.restore_quantize_state()

        # update state
        verbose, self._verbose = self._verbose, False
        self.check(executor=executor, dataloader=dataloader, collate_fn=collate_fn)
        self._verbose = verbose

    def check(self, executor: BaseGraphExecutor,
        dataloader: Iterable, collate_fn: Callable):
        """
        Check quantization error with a given dataloader with current checkpoints.
            Return whether quantization error is lower than before.

        Args:
            executor (BaseGraphExecutor): [description]
            dataloader (Iterable): [description]
            collate_fn (Callable): [description]

        Returns:
            [type]: [description]
        """

        # step - 1, collecting data
        for data in dataloader:
            if collate_fn is not None: data = collate_fn(data)
            outputs = executor.forward(inputs=data, output_names=self._interested_outputs)
            for name, output in zip(self._interested_outputs, outputs):
                self._checkpoints[name].push(tensor=output, is_reference=False)

        # step - 2, calculating loss
        losses = []
        for name in self._interested_outputs:
            ckpt = self._checkpoints[name]
            assert isinstance(ckpt, FinetuneCheckPoint)
            qt_out, fp_out = ckpt.pop()
            qt_out = torch.cat([tensor for tensor in qt_out])
            fp_out = torch.cat([tensor for tensor in fp_out])
            losses.append(self._loss_fn(y_pred=qt_out, y_real=fp_out).item())
            ckpt.clear()

        # step - 3, comparing loss
        loss_now, loss_old = sum(losses), sum([ckpt.best_loss for ckpt in self._checkpoints.values()])
        loss_now, loss_old = loss_now / len(losses), loss_old / len(losses)
        if self._verbose: print(f'SNR after optimization: {loss_old * 100 :.4f}% -> {loss_now * 100:.4f}%.')

        # if there is a loss drop, update all losses.
        if loss_old > (loss_now * CHECKPOINT_TOLERANCE):
            for idx, name in enumerate(self._interested_outputs):
                ckpt = self._checkpoints[name]
                assert isinstance(ckpt, FinetuneCheckPoint)
                ckpt.best_loss = losses[idx]
            return True

        if self._verbose: print(f'Not a perfect loss drop, skip this optimization.')
        return False

    def optimize(
        self, processer: GraphCommandProcesser,
        dataloader: Iterable, executor: BaseGraphExecutor, **kwargs) -> None:
        raise NotImplementedError('Can not invoke this function. '
                                  'Please inherit this class and give an implmenetation to override this function.')

    def dequantize_immediately(self, operation: Operation):
        """
        Dequantize an operation inplace, use this function carefully.
            if parameter value has been changed during your optimization procedure,
            then it is not safe to dequantize an operation via this function,
            use operation.dequantize to load stored fp32 value instead.

        This function will change quantization state to dequantize an operation,
            Only quantization state will be changed by this function so that it is
            extremely fast.

        If your parameter value has already been baked, an exception will be thrown.
        Args:
            operation (Operation): [description]
        """
        if isinstance(operation, QuantableOperation):
            for cfg, _ in operation.config_with_variable:
                assert cfg.state not in {QuantizationStates.BAKED, QuantizationStates.PASSIVE_BAKED}, (
                    'Value has already been baked, can not dequantize it via this function.')

                if cfg not in self._quant_state_recorder:
                    self._quant_state_recorder[cfg] = cfg.state
                    cfg.state = QuantizationStates.DEQUANTIZED

    def quantize_immediately(self, operation: Operation):
        """
        Restore quantization state of an operation, use this function carefully.
            if parameter value has been changed during your optimization procedure,
            then it is not safe to restore state via this function,
            use operation.restore_quantize_state to load stored quant value instead.

        This function will change quantization state to quantize an operation,
            Only quantization state will be changed by this function so that it is
            extremely fast.

        If your parameter value has already been baked, an exception will be thrown.
        Args:
            operation (Operation): [description]
        """
        if isinstance(operation, QuantableOperation):
            for cfg, _ in operation.config_with_variable:
                if cfg in self._quant_state_recorder:
                    stored_state = self._quant_state_recorder[cfg]
                    cfg.state = stored_state
                    self._quant_state_recorder.pop(cfg)

class BiasCorrectionPass(TrainingBasedPass):
    def __init__(self, auto_check: bool=False, interested_output: List[str] = None,
                 verbose: bool = True, max_steps:int = 8) -> None:
        """
        Quantization can introduce a biased error in the activations.
            Bias correction serves as a useful prosedure to eliminate those introduced bias error.

        let: Y = WX + b
             Quant(Y) = Qunat(W) Quant(X) + b

             bias_error = reduce_mean(Y - Quant(Y))

        Correct bias by: b = b + bias_error

        Args:
            quantize_function (BaseQuantFunction): [description]
            auto_check (bool, optional): [description]. Defaults to False.
        """
        super().__init__(name='PPQ Bias Correction Pass',
                         interested_outputs=interested_output, verbose=verbose)
        self._auto_check = auto_check
        self._max_steps = max_steps

    @ empty_ppq_cache
    def optimize(
        self,
        processer: GraphCommandProcesser,
        dataloader: Iterable,
        executor: BaseGraphExecutor,
        collate_fn: Callable,
        **kwargs
    ) -> None:
        def collect_bias(output: torch.Tensor, collector: list, op_type: str):
            if op_type in {'Conv', 'ConvTranspose'}:
                collector.append(torch.mean(output, dim=(0, 2, 3)).unsqueeze(0))
            elif op_type in {'Gemm'}:
                collector.append(torch.mean(output, dim=(0, )).unsqueeze(0))
            else: raise TypeError(f'Unsupported Operation type: {op_type}')

        assert isinstance(executor, TorchExecutor), (
            'PPQ Training-based optimization algorithm needs a TorchExecutor.')

        if self._auto_check:
            self.initialize_checkpoints(graph=processer.graph, executor=executor,
                                        dataloader=dataloader, collate_fn=collate_fn)

        for idx, operation in tqdm(enumerate(executor._executing_order),
                                   desc='Bias Correction Procedure ...',
                                   total=len(executor._executing_order)):
            assert isinstance(operation, Operation)
            if not has_bias(operation): continue

            bias, output_var = operation.inputs[-1].value, operation.outputs[0]
            qt_collector, fp_collector = [], []

            for idx, data in enumerate(dataloader):
                if collate_fn is not None: data = collate_fn(data)
                [output] = executor.forward(inputs=data, output_names=[output_var.name])
                collect_bias(output, qt_collector, op_type=operation.type)
                if idx >= self._max_steps: break
            self.dequantize_immediately(operation)

            for idx, data in enumerate(dataloader):
                if collate_fn is not None: data = collate_fn(data)
                [output] = executor.forward(inputs=data, output_names=[output_var.name])
                collect_bias(output, fp_collector, op_type=operation.type)
                if idx >= self._max_steps: break
            self.quantize_immediately(operation)

            bias_error = (torch.mean(torch.cat(fp_collector), dim=0) - torch.mean(torch.cat(qt_collector), dim=0))
            if self._auto_check:
                backup = bias.clone()
                operation.inputs[-1].value = bias + bias_error
                if not self.check(executor=executor, dataloader=dataloader, collate_fn=collate_fn):
                    operation.inputs[-1].value = backup
            else: operation.inputs[-1].value = bias + bias_error

class AdaRoundPass(QuantizationOptimizationPass):
    def __init__(self,
                 collecting_device: str = 'cpu',
                 epoch: int = 512,
                 batch_size: int = 32) -> None:
        super().__init__(name='PPQ AdaRound Pass')
        self._collecting_device = collecting_device
        self.epoch = epoch
        self.batch_size = batch_size

    @ empty_ppq_cache
    def optimize(
        self,
        processer: GraphCommandProcesser,
        dataloader: Iterable,
        executor: BaseGraphExecutor,
        collate_fn: Callable,
        **kwargs
    ) -> None:
        assert isinstance(executor, TorchExecutor), ('PPQ Training-based optimization algorithm needs a TorchExecutor.')
        graph = processer.graph
        sorted_ops = graph.topological_sort()
        for idx, target_op in tqdm(enumerate(sorted_ops), desc='AdaRound...', total=len(graph.operations)):
            if not isinstance(target_op, QuantableOperation): continue
            if not target_op.type in {'Conv', 'ConvTranspose', 'Gemm'}: continue

            fp_outputs, quant_inputs = [], []
            interested_var = (target_op.inputs[0].name, target_op.outputs[0].name)

            for op in sorted_ops[: idx + 1]:
                if isinstance(op, QuantableOperation): op.dequantize()
            for data in tqdm(dataloader, desc='AdaRound Procedure 1', total=len(dataloader)):
                if collate_fn is not None: data = collate_fn(data)
                fp_input, fp_output = executor.forward(inputs=data, output_names=interested_var)
                fp_outputs.append(fp_output)
            fp_weight = target_op.parameters[0].value.clone()

            for op in sorted_ops[: idx + 1]:
                if isinstance(op, QuantableOperation): op.restore_quantize_state()
            for data in tqdm(dataloader, desc='AdaRound Procedure 2', total=len(dataloader)):
                if collate_fn is not None: data = collate_fn(data)
                quant_input, _ = executor.forward(inputs=data, output_names=interested_var)
                quant_inputs.append(quant_input)

            fp_outputs_concat = torch.cat(fp_outputs)
            quant_inputs_concat = torch.cat(quant_inputs)
            weight, bias = target_op.parameters[0].value, None
            if target_op.num_of_input == 3:
                bias = target_op.parameters[1].value
                bias = bias.clone()
            weight = weight.clone()
            params = [weight, bias] if bias is not None else [weight]
            for param in params: param.requires_grad = True

            print ('Adaround optimize {}'.format(target_op.name))
            weight_quantization_config = target_op.config.input_quantization_config[1].dominated_by
            weight_scale = weight_quantization_config.scale
            weight_offset = weight_quantization_config.offset

            max_iter = self.epoch * fp_outputs_concat.shape[0] / self.batch_size
            reg = AdaroundRegTerm(max_iter)

            # per-channel scale preprocess
            if weight_quantization_config.policy.has_property(QuantizationProperty.PER_CHANNEL):
                view_shape = [
                    1 if axis != weight_quantization_config.channel_axis else -1
                    for axis in range(fp_weight.ndim)]
                weight_scale = weight_scale.view(view_shape)
                weight_offset = weight_offset.view(view_shape)

            # init continuous_v, make sure h(v) = round_diff
            round_diff = (fp_weight / weight_scale) - (fp_weight / weight_scale).floor()
            v_init = -torch.log((reg.zeta - reg.gamma) / (round_diff - reg.gamma) - 1)
            continuous_v = torch.nn.Parameter(v_init.to(executor._device), True)
            optimizer = torch.optim.Adam([continuous_v])

            cur_iter = 0
            data_len = quant_inputs_concat.shape[0]
            for ep_idx in range(self.epoch):
                batch_num = int(ceil(data_len / self.batch_size))
                # shuffle data
                index = np.arange(data_len)
                np.random.shuffle(index)
                for idx in range(batch_num):
                    st = idx * self.batch_size
                    ed = min(st + self.batch_size, data_len)

                    # soft AdaRound quant weight
                    params[0] = self.adaround_quant_weight(fp_weight, weight_scale, weight_offset, weight_quantization_config, continuous_v)
                    in_snap = [ quant_inputs_concat[index[st:ed,]] ]
                    [quant_output] = executor.operation_forward(target_op, inputs=in_snap + params)
                    fp32_output = fp_outputs_concat[index[st:ed],]

                    loss = Lp_norm(fp32_output, quant_output) + reg(continuous_v, cur_iter)
                    optimizer.zero_grad()
                    loss.backward(retain_graph=True)
                    optimizer.step()
                    cur_iter += 1

                if ep_idx % 100 == 0:
                    print("Epoch: {:<4} L2 Loss: {:>10.3f} Beta: {:>3.3f}".format(ep_idx, loss, reg.beta))
            h_v = AdaroundRegTerm().rectified_sigmoid(continuous_v)
            print("Loss: {:>5.3f} Ceil: {:>5} Floor: {:>5} Total: {:>5} Ratio: {:>.3f}".format(
                loss,
                h_v[h_v + 1e-4 >= 1.0].numel(), h_v[h_v <= 1e-4].numel(), torch.numel(h_v),
                (h_v[h_v + 1e-4 >= 1.0].numel() + h_v[h_v <= 1e-4].numel()) / torch.numel(h_v)))

            # update weight
            rounded_weight = self.adaround_quant_weight(fp_weight, weight_scale, weight_offset, weight_quantization_config, continuous_v, soft=False)
            target_op.parameters[0].value.copy_(rounded_weight)
            del fp_outputs_concat
            del quant_inputs_concat
            target_op.config.input_quantization_config[1].state = QuantizationStates.ACTIVATED
            if bias is not None:
                target_op.parameters[1].value.copy_(bias)
                target_op.config.input_quantization_config[-1].state = QuantizationStates.PASSIVE

    def adaround_quant_weight(self, weight, scale, offset, weight_quantization_config, round_var, soft=True):
        quant_max = weight_quantization_config.quant_max
        quant_min = weight_quantization_config.quant_min
        if soft:
            weight = (weight / scale).floor() + AdaroundRegTerm().rectified_sigmoid(round_var)
        else:
            weight = (weight / scale).floor() + (round_var >= 0).float()
        weight = torch.clamp(weight + offset, quant_min, quant_max)
        weight = (weight - offset) * scale
        return weight

class LearningStepSizeOptimization(TrainingBasedPass):
    def __init__(self, name: str = 'PPQ LSQ Optimization') -> None:
        super().__init__(name=name)

    def optimize(self, processer: GraphCommandProcesser,
                 dataloader: Iterable, executor: BaseGraphExecutor,
                 **kwargs) -> None:

        return super().optimize(processer, dataloader, executor, **kwargs)

class AdvancedQuantOptimization(TrainingBasedPass):
    """
    PPQ Advanced Quantization Optimization

    This optimization pass minimize the quantization errors of each subgraph separately
        by optimizing its parameters over the calibration set.

    Where:
        qout = quant( quant(W + W_offset) * quant(X) + quant(bias + bias_offset) )

        fout = W * B + bias

        error = Mean((qout - fout)^2)

    This training procedure trys to solve best W_offest and bias_offset to minimize error
        Based on your setting and network size, the training procedure will takes 5~120 minutes.

    This function will treat your network as a series of subgraph, you should notice that
        ONLY THE OUTPUT VALUE OF A SUBGRAPH IS OPTIMIZED IN THIS PASS,
        ACTIVATIONS THAT INSIDE YOUR SUBGRAPH MIGHT BE GREATLY CHANGED!
        DO NOT ATTEMPT TO COMPARE THOSE QUANTIZED ACTIVATION WITH ITS FP32 VERSION.

    We use graph search engine to build subgraph from your network with pattern below:

    while len(graph.get_downstream_operations(start_op)) == 1:
        end_op = graph.get_downstream_operations(start_op)[0]
        if len(graph.get_upstream_operations(end_op)) == 1:
            path.append(end_op)
            start_op = end_op
        else: break

    Args:
        TrainingBasedPass ([type]): [description]
    """
    def __init__(self, collecting_device: str, limit: float = 3.0, lr: float = 1e-3,
                 interested_outputs: List[str] = None, interested_layers: List[str] = None,
                 verbose: bool = True, check: bool = True) -> None:

        super().__init__(
            name='PPQ Advanced Optimization Procedure(Blockwise)',
            interested_outputs=interested_outputs, verbose=verbose)

        self.lr                = lr
        self.collecting_device = collecting_device
        self.check_flag        = check
        self.offset_limit      = limit
        self.interested_layers = interested_layers
        self.t_step            = OPTIM_ADVOPT_THRESHOLD_STEP
        self.steps_per_epoch   = OPTIM_ADVOPT_STEP_PER_EPOCH
        self.patient           = OPTIM_ADVOPT_PATIENT
        self.passive_boost     = OPTIM_ADVOPT_PASSIVE_BOOST
        self.max_iter          = 10000

        if isinstance(self.interested_layers, list) and len(self.interested_layers) == 0:
            self.interested_layers = None

    def collect_training_data(
        self, output_name: str,
        dataloader: Iterable,
        executor: BaseGraphExecutor,
        collate_fn: Callable) -> List[List[torch.Tensor]]:

        output_collector = []
        for data in dataloader:
            if collate_fn is not None: data = collate_fn(data)
            [output] = executor.forward(data, output_names=[output_name])
            output_collector.append(output.to(self.collecting_device))
        return output_collector

    @ empty_ppq_cache
    def finetune(
        self, quant_inputs: List[torch.Tensor], fp32_outputs: List[torch.Tensor],
        executor: TorchExecutor, block: List[Operation],
        dataloader: Iterable, collate_fn:Callable) -> None:

        # initialize training environment.
        losses     = []
        last_loss  = 1e9
        threshold  = OPTIM_ADVOPT_INITIAL_THRES
        trys_count = 0
        cur_iter   = 0
        delegates  = []
        device     = executor._executing_contenxt.executing_device
        loss_recorder = {}
        output_var = block[-1].outputs[0]
        input_var  = block[0].inputs[0]

        dataset = RandomMemDataset(data=[[qt, fp] for qt, fp in zip(quant_inputs, fp32_outputs)])

        # create trainable delegates for each parameter.
        for operation in block:
            if operation.is_computing_op and isinstance(operation, QuantableOperation):
                for cfg, var in operation.config_with_variable:
                    if not var.is_parameter: continue
                    boost = 1 if cfg.state == QuantizationStates.PASSIVE else self.passive_boost
                    delegates.append(TrainableDelegate(
                        value=var.value, config=cfg,
                        limit=self.offset_limit, boost=boost, binding=var
                        )
                    )

        # set up optimizer, ready for training.
        optimizer = torch.optim.Adam(params=[d.offset for d in delegates], lr=self.lr)
        while cur_iter < self.max_iter:
            for _ in range(self.steps_per_epoch):
                qt_input, fp_output = dataset.pop()
                # update weights:
                for parameter in delegates:
                    assert isinstance(parameter, TrainableDelegate)
                    parameter.quantize(threshold=threshold)

                qt_input, fp_output = qt_input.to(device), fp_output.to(device)
                qt_output = executor.partial_graph_forward(
                    operations=block, feed_dict={input_var.name: qt_input},
                    output_names=[output_var.name])[0]

                # compute loss
                optimizer.zero_grad()
                loss = torch_mean_square_error(qt_output, fp_output)
                loss.backward()
                optimizer.step()

                cur_iter += 1
                losses.append(loss.detach().item())

            # pleatu interval schedule.
            cur_loss = sum(losses) / len(losses)
            if cur_loss < last_loss * .99:
                last_loss, trys_count = cur_loss, 0
                # record loss
                loss_recorder[threshold] = cur_loss
            else:
                trys_count += 1
                if trys_count > self.patient:
                    # rebuild optimizer, clear all state.
                    optimizer.state = collections.defaultdict(dict)
                    trys_count, last_loss, threshold = 0, 1e9, threshold - self.t_step
                    if threshold <= 0.5: break

            # clear loss state
            losses.clear()

        # DEBUG INFO, JUST IN CASE.
        '''
        for offset in [d.offset for d in delegates]:
            print(offset.shape)
            print(' ------ GARD ------')
            print(offset._grad.flatten().max())
            print(' ------ VALUE ------')
            print(offset.flatten().max())
        '''

        # clear all delegates
        for delegate in delegates:
            assert isinstance(delegate, TrainableDelegate)
            delegate.clear()

        # display loss
        if self._verbose:
            loss_recorder = {f'{threshold * 100:.1f}%': loss for threshold, loss in loss_recorder.items()}
            print(f'Optimize Result For Block: ', end='')
            # display your block with following ugly code.
            print(block[0].name, end='')
            for operation in block[1:]: print('->' + operation.name, end='')
            print('')

            MeasurePrinter(
                data=loss_recorder, measure='MSE',
                label='Threshold', order=None).print()

        # Check
        if self.check_flag:
            if not self.check(executor=executor, dataloader=dataloader, collate_fn=collate_fn):
                for delegate in delegates:
                    assert isinstance(delegate, TrainableDelegate)
                    delegate.withdraw()

        # detach weight
        for delegate in delegates:
            assert isinstance(delegate, TrainableDelegate)
            delegate.binding.value = delegate.binding.value.detach()

    def build_block_from_start(self, graph: BaseGraph, start_op: QuantableOperation) -> List[Operation]:
        path = [start_op]
        while len(graph.get_downstream_operations(start_op)) == 1:
            end_op = graph.get_downstream_operations(start_op)[0]
            if len(graph.get_upstream_operations(end_op)) == 1:
                path.append(end_op)
                start_op = end_op
            else: break

        num_of_computing_ops = sum([1 for op in path if op.is_computing_op])
        while num_of_computing_ops > OPTIM_ADVOPT_GRAPH_MAXSIZE:
            if path[-1].is_computing_op: num_of_computing_ops -= 1
            path.pop(-1)
        return path

    def optimize(
        self, processer: GraphCommandProcesser, dataloader: Iterable,
        executor: TorchExecutor, collate_fn: Callable, **kwargs) -> None:

        if self._interested_outputs is None:
            self._interested_outputs = [name for name in processer.graph.outputs]

        if self.collecting_device == 'executor':
            self.collecting_device = executor._device

        graph = processer.graph
        visited = set()

        # check if there is any baked value inside your graph
        for operation in graph.operations.values():
            if isinstance(operation, QuantableOperation):
                for cfg, var in operation.config_with_variable:
                    if cfg.state in {QuantizationStates.BAKED, QuantizationStates.PASSIVE_BAKED}:
                        raise PermissionError('Can not apply advanced optimization pass when weight value is baked. '
                                              f'Variable {var.name} has a baked value.')

        # find all operations that need to be finetuned.
        interested_ops = []
        for target_op in graph.topological_sort():
            if isinstance(target_op, QuantableOperation) and target_op.is_computing_op:
                if self.interested_layers is None: interested_ops.append(target_op)
                elif self.interested_layers is not None and target_op.name in self.interested_layers:
                    interested_ops.append(target_op)

        # set up checkpoints
        if self.check_flag:
            self.initialize_checkpoints(
                graph=graph, executor=executor,
                dataloader=dataloader, collate_fn=collate_fn)

        for start_op in tqdm(interested_ops, total=len(interested_ops), desc='Advanced Optim Procedure...'):
            assert isinstance(start_op, QuantableOperation)

            if start_op in visited: continue
            block = self.build_block_from_start(graph=graph, start_op=start_op)

            end_op       = block[-1]
            block_input  = start_op.inputs[0]
            block_output = end_op.outputs[0]

            # dequantize prefix operations and block operations
            for op in graph.operations.values():
                if isinstance(op, QuantableOperation):
                    op.dequantize()
                    # can not use dequantize_immediately cause weight has been changed.
                    # self.dequantize_immediately(op)

            fp32_outputs = self.collect_training_data(
                output_name=block_output.name, dataloader=dataloader,
                executor=executor, collate_fn=collate_fn)

            # quantize prefix operations and block operations
            for op in graph.operations.values():
                if isinstance(op, QuantableOperation):
                    op.restore_quantize_state()

            quant_inputs = self.collect_training_data(
                output_name= block_input.name, dataloader=dataloader,
                executor=executor, collate_fn=collate_fn)

            # start training, solve the best parameters
            self.finetune(
                quant_inputs=quant_inputs, fp32_outputs=fp32_outputs,
                executor=executor, block=block,
                dataloader=dataloader, collate_fn=collate_fn)

            for op in block: visited.add(op)

            # empty cache.
            fp32_outputs.clear()
            quant_inputs.clear()
            empty_cache()


"""Energy calculations and histograms."""

import abc
import collections
import datetime
import functools
import itertools
import math
import os
import pickle
import warnings

import numpy
import pandas
from matplotlib import pyplot

# Use seaborn's plotting styles.
try:
    import seaborn
    seaborn.set()
except ModuleNotFoundError:
    msg = "Seaborn wasn't found."
    warnings.warn(msg)

# Energy-z histogram z-limits for the 8 by 4 plate arrangements. The
# z-limits for other plate arrangements are adjusted in increments equal
# to the `_plate_separation`.
_z_lims_8_4 = (-65, 95)
_plate_separation = 4.15 + 3.5

# Energy-z histogram energy-limits.
_e_lims_200 = (0, 400)  # for 200 GeV
_e_lims_350 = (0, 650)  # for 350 GeV
_tube_e_lims_200 = tuple(lim / 25 for lim in _e_lims_200)
_tube_e_lims_350 = tuple(lim / 25 for lim in _e_lims_350)

_default_bin_density = 10
_default_dpi = 300
_default_length_units = 'mm'
_default_energy_units = 'MeV'

# Just used to split energy-z plots.
_PEEK_z_lims = (-43.5 / 2, 43.5 / 2)

_default_tube_z_lims = (-35 / 2, 35 / 2)

# Convert the _default_middle_z_lims to integers so that the plot
# tick labels are displayed without decimals.
_default_middle_z_lims = tuple(int(lim) for lim in (-8 / 2, 8 / 2))

_default_xy_lims = 2 * ((-3, 3),)
_default_xy_hist_z_lims = _default_middle_z_lims

_image_format = 'jpg'
_data_format = 'csv'
linewidth = 0.5

class Calc(abc.ABC):
    """
    A base for calculations and presentations of results a certain
    kind.

        Attributes:
            resultss: A collection of the results. It is a pandas DataFrame
            in this base class, but that can change.

            piece: Piece of analysis that this calculation is for.

            save_dir: Save output in this folder.

            save_name: A name for the saved files.
    """

    def __init__(self, piece):
        self._piece = piece
        self.resultss = pandas.DataFrame()

        if self._piece.out_dir:
            # Save when Piece has an output directory.
            self.save_dir = self._piece.out_dir
            self.save_name = self._piece.name
        else:
            self.save_dir = None
            self.save_name = None

    def __repr__(self):
        return repr(self.resultss)

    def get(self, i=None):
        """
        Get the ith set of results, or if i isn't given, make sure
        resultss contains only one entry and return it.

        :param i: The index of the results to get.
        :type i: int or None
        :return: Results.
        :rtype:
        """
        if i is None:
            self._assert_single_entry()
            return self.resultss.loc[0]

        return self.resultss.loc[i]

    def add_data(self, data):
        """
        Calculate results from data and keep them.

        :param data: Hits data.
        :type data: pandas.DataFrame
        :return:
        :rtype:
        """
        self.add_results(self._data2results(data))

    def add_results(self, results):
        """
        Keep some already calculated results (after making sure they are
        good results).

        :param results: Results to add.
        :type results: pandas.Series or dict or collections.OrderedDict
        :return:
        :rtype: None
        """
        self._check_results(results)
        self.resultss = _ordered_append(self.resultss, results)

    def save(self):
        """
        Save results to a csv file if save_dir is set.

        :return: Path to saved results file.
        :rtype: str
        """
        if not self.save_dir:
            return None

        filename = '.'.join((self.save_name, _data_format))
        filepath = os.path.join(self.save_dir, filename)

        return save_dataframe(self.resultss, filepath)

    def pickle_save(self, filepath=None):
        """Save self in a pickle format. Does nothing if `filepath`
        isn't given and `self.save_dir` is not set.

        :param filepath: Location to save. If not given, use
        `self.save_dir`.
        """
        if not filepath:
            if not self.save_dir:
                return
            filepath = os.path.join(self.save_dir, self.save_name + '.p')
        with open(filepath, 'wb') as file:
            pickle.dump(self, file)

    def _assert_single_entry(self):
        """Make sure only one set of results has been added."""
        assert len(self.resultss) == 1, \
            "This Calc container does not have a single entry."

    def _assert_multiple_entries(self):
        """Make sure more than one set of results has been added."""
        assert len(self.resultss) > 1, \
            "This Calc container does not have more than one set of" \
            " results."

    def _assert_nonempty(self):
        """Make sure at least one set of results has been added."""
        assert len(self.resultss) >= 1, \
            "This Calc container is empty."

    @abc.abstractmethod
    def _data2results(self, data):
        """
        Calculate results from hits data and return them.

        :param data: Hits data.
        :return: Calc results.
        """

    @abc.abstractmethod
    def _check_results(self, results):
        """Make sure some calculated results are good to keep."""

class Numbers(Calc):
    """
    The numbers don't lie.

    This class manages calculations of different numbers from the data
    and writes them to files.
    """

    def add_data(self, data, tags=None):
        self.add_results(self._data2results(data, tags))

    def append_mean_and_uncertainties(
            self, mean_tags=None, std_tags=None, sem_tags=None
    ):
        """
        Append mean, standard deviation and standard error entries to
        the collection of results.

        :param mean_tags: Tags to use for the new means entry.
        :type mean_tags: dict or collections.OrderedDict or None
        :param std_tags: Tags to use for the new standard deviation
            entry.
        :type std_tags: dict or collections.OrderedDict or None
        :param sem_tags: Tags to use for the new standard error in the
            mean entry.
        :type sem_tags: dict or collections.OrderedDict or None
        :return: The new mean, standard deviation, and standard error.
        :rtype:
        """
        # Pandas' std function gives NANs for DataFrames with only one
        # entry.
        self._assert_multiple_entries()

        new_rows = [self.resultss.mean(), self.resultss.std()]
        new_rows.append(  # standard error
            new_rows[1] / math.sqrt(len(self.resultss))
        )

        for i, tags in enumerate((mean_tags, std_tags, sem_tags)):
            if tags:
                new_rows[i] = pandas.Series(tags).combine_first(new_rows[i])
            self.add_results(new_rows[i])

        return new_rows

    def _data2results(self, data, tags=None):
        """

        :param data:
        :type data:
        :param tags: Extra result values for tagging an entry (e.g. with
            its event and run).
        :type tags: dict or collections.OrderedDict or None
        :return:
        :rtype:
        """
        new_results = collections.OrderedDict()

        if tags:
            new_results.update(tags)

        new_results.update(self.__split_z(data))

        return new_results

    def _check_results(self, results):
        assert len(results) <= 9, "Tried to add funny results."

    @staticmethod
    def __split_z(data):
        """
        Calculate energy deposit mean and std dev on sections of
        data split up along z.

        :param data: Hits data.
        :type data: pandas.DataFrame
        :return: Calculation results.
        :rtype: dict
        """
        tube_indices = numpy.logical_and(
            _default_tube_z_lims[0] <= data.z,
            data.z <= _default_tube_z_lims[1]
        )
        middle_indices = numpy.logical_and(
            _default_middle_z_lims[0] <= data.z,
            data.z <= _default_middle_z_lims[1]
        )
        return {
            'full_e_dep':
                data.energy_deposit.sum(),
            'tube_e_dep':
                data.energy_deposit[tube_indices].sum(),
            'middle_e_dep':
                data.energy_deposit[middle_indices].sum(),
        }

    @staticmethod
    def __tubes(data):
        """TODO: Document and organize this."""
        offset = 7.5 / 4
        # TODO: Change default to self (probably don't need to do this).
        data = data[numpy.logical_and(
            _default_middle_z_lims[0] < data.z,
            data.z < _default_middle_z_lims[1]
        )]

        bottom = data.y <= 0
        top = data.y > 0

        bottom_left = numpy.logical_and(bottom, data.x <= -offset)
        bottom_right = numpy.logical_and(bottom, data.x > -offset)
        top_left = numpy.logical_and(top, data.x <= offset)
        top_right = numpy.logical_and(top, data.x > offset)

        bottom_left_sum = data.energy_deposit[bottom_left].sum()
        bottom_right_sum = data.energy_deposit[bottom_right].sum()
        top_left_sum = data.energy_deposit[top_left].sum()
        top_right_sum = data.energy_deposit[top_right].sum()

        return {
            'top_right': top_right_sum,
            'top_left': top_left_sum,
            'bottom_left': bottom_left_sum,
            'bottom_right': bottom_right_sum
        }

class Histogram(Calc):
    """
    Base histogram.

        Attributes:
            bin_density: E.g. bins per mm.

            dpi: Dots per inch for images.

            title: The plot title.

    TODO: Add the option to close all the plot figures created.
    """

    def __init__(
            self,
            piece,
            title=None,
            bin_density=None,
            dpi=None,
            energy_units=None,
            length_units=None
    ):
        super().__init__(piece)

        # Attributes with defaults.
        self.title = title or self._default_title
        self.bin_density = bin_density or _default_bin_density
        self.dpi = dpi or _default_dpi
        self.energy_units = energy_units or _default_energy_units
        self.length_units = length_units or _default_length_units

    def save_fig(self, fig, file_suffix):
        """
        Save a figure to a new file if save_dir is set.

        :param fig: The figure to save.
        :type fig: figure
        :param file_suffix: Added to the end of the saved filename.
        :type file_suffix: str or None
        :return: File path of the saved figure.
        :rtype: str
        """
        if not self.save_dir:
            return None

        if file_suffix:
            save_name = '-'.join((self.save_name, file_suffix))
        else:
            save_name = self.save_name

        filename = '.'.join((save_name, _image_format))
        filepath = os.path.join(self.save_dir, filename)

        fig.savefig(
            filepath, dpi=self.dpi, format=_image_format, bbox_inches='tight'
        )

        return filepath

    def _to_density(self, sums):
        """
        Convert values from units of energy/bin to units of
        energy/distance.

        :param sums: Values to convert.
        :type sums: numpy.ndarray
        :return: Converted values.
        :rtype: numpy.ndarray
        """
        return sums * self.bin_density

    @property
    def _default_title(self):
        """Default plot title."""
        raise NotImplementedError

    @abc.abstractmethod
    def plot_single(self, i=None, save=True):
        """
        Plot the ith set of sums. If i isn't given, assume there is only
        set of sums and plot it.

        :param save: Save to file with the standard filename if True.
        :type save: bool
        :param i: Index of the sums.
        :type i: int
        :return: New figure and axis/axes.
        :rtype: tuple
        """

    @abc.abstractmethod
    def plot_means(self):
        """
        Plot the mean, and maybe the standard deviation, of all results.

        :return: New figure and axis/axes.
        :rtype: tuple
        """

    @abc.abstractmethod
    def _make_fig_and_axes(self):
        """
        Make a new labeled figure and axis/axes.

        :return: The figure and axis/axes.
        :rtype: tuple
        """

class EnergyVsZ(Histogram):
    """
    Histograms of energy deposit vs. z.

        Attributes:
            e_lims: Overall energy/y-axis limits.

            tube_e_lims: Tube energy/y-axis limits.

            z_lims: Overall z limits.

            tube_z_lims: Tube z limits.

            middle_z_lims: Tube middle section z limits.
    """

    _default_title = 'Energy vs. z.'

    def __init__(
            self,
            piece,
            title=None,
            e_lims=None,
            tube_e_lims=None,
            z_lims=None,
            tube_z_lims=None,
            middle_z_lims=None,
            **kwargs
    ):
        super().__init__(piece, title, **kwargs)

        # Limits according to piece info.
        self.__e_lims, self.__tube_e_lims, self.__z_lims = self.__get_lims(
            e_lims, tube_e_lims, z_lims
        )

        self.__tube_z_lims = tube_z_lims or _default_tube_z_lims
        self.__middle_z_lims = middle_z_lims or _default_middle_z_lims

        # Histogram bins.
        self.__bins = _make_bins(
            self.__z_lims[0], self.__z_lims[1], self.bin_density
        )
        self.__bin_mids = _make_bin_midpoints(self.__bins)

    def plot_single(self, i=None, save=True, energy_label=None):
        """
        Plot a single event.

        :param energy_label: Energy deposit in middle tube sections,
            displayed in a text label on the plot.
        """
        fig, ax, ax_middle = self._make_fig_and_axes()
        fig.suptitle(self.title)
        self.__label_middle(ax_middle, energy_label)

        tube_kwargs = {
            'linewidth': linewidth, 'color': 'purple', 'label': 'Tube Cals'
        }
        plate_kwargs = {
            'linewidth': linewidth, 'color': 'blue', 'label': 'Plate Cals'
        }
        self.__split_plot(
            x=self.__bin_mids,
            ys=(self._to_density(self.get(i)),),
            plot_fn=ax_middle.show_plots,
            main_ax=ax,
            kwargs1=tube_kwargs,
            kwargs2=plate_kwargs
        )

        self.__add_legend(fig, ax, ax_middle)
        if save:
            self.save_fig(fig, file_suffix='z')

    def plot_means(self, energy_label=None):
        fig, ax, ax_middle = self._make_fig_and_axes()
        fig.suptitle(self.title + ' Averaged.')
        self.__label_middle(ax_middle, energy_label)

        means = self._to_density(numpy.mean(self.resultss, axis=0))
        stds = self._to_density(numpy.std(self.resultss, axis=0))

        tube_means_kwargs = {
            'linewidth': linewidth,
            'color': 'purple',
            'label': 'Tube Cals Average'
        }
        plate_means_kwargs = {
            'linewidth': linewidth,
            'color': 'blue',
            'label': 'Plate Cals Average'
        }
        tube_stds_kwargs = tube_means_kwargs.copy()
        plate_stds_kwargs = plate_means_kwargs.copy()
        tube_stds_kwargs.update({'alpha': 0.3, 'label': 'Standard Deviation'})
        plate_stds_kwargs.update({'alpha': 0.3, 'label': 'Standard Deviation'})

        self.__split_plot(  # Plot means.
            x=self.__bin_mids,
            ys=(means,),
            plot_fn=ax_middle.show_plots,
            main_ax=ax,
            kwargs1=tube_means_kwargs,
            kwargs2=plate_means_kwargs
        )
        self.__split_plot(  # Plot standard deviations.
            x=self.__bin_mids,
            ys=(means + stds, means - stds),
            plot_fn=ax_middle.fill_between,
            main_ax=ax,
            kwargs1=tube_stds_kwargs,
            kwargs2=plate_stds_kwargs
        )

        self.__add_legend(fig, ax, ax_middle)
        self.save_fig(fig, file_suffix='z')

    def plot_multi(self, energy_label=None):
        """Plot all of the events on one graph at once."""
        fig, ax, ax_middle = self._make_fig_and_axes()
        num_events = len(self.resultss)
        fig.suptitle(self.title + f' {num_events} events.')
        self.__label_middle(ax_middle, energy_label)

        kwargs = {'linewidth': linewidth, 'alpha': 0.2}
        colors = itertools.cycle((
            'red', 'orange', 'yellow', 'pink', 'purple', 'brown', 'black'
        ))
        for _, sums in self.resultss.iterrows():
            kwargs['color'] = next(colors)
            self.__split_plot(
                x=self.__bin_mids,
                ys=(self._to_density(sums),),
                plot_fn=ax_middle.show_plots,
                main_ax=ax,
                kwargs1=kwargs,
                kwargs2=kwargs
            )

        self.save_fig(fig, file_suffix=f'z-{num_events}events')

    def _data2results(self, data):
        return pandas.Series(
            numpy.histogram(
                data.z, bins=self.__bins, weights=data.energy_deposit
            )[0],
            index=self.__bin_mids
        )

    def _check_results(self, sums):
        assert len(sums) == (len(self.__bins) - 1), \
            "Tried to add funny histogram sums."

    def _make_fig_and_axes(self):
        fig, ax = pyplot.subplots()
        ax_middle = ax.twinx()

        ax.set_xlabel(f'z ({self.length_units})')
        ax.set_ylabel(
            'E Dep Density - Plate Cals'
            f' ({self.energy_units} / {self.length_units})'
        )
        ax_middle.set_ylabel(
            'E Dep Density - Tube Cals'
            f' ({self.energy_units} / {self.length_units})'
        )

        # Indicate the middle tube sections with vertical bars and axis
        # ticks.
        for lim in self.__middle_z_lims:
            ax_middle.axvline(
                lim, linestyle='--', linewidth=1, color='gray', zorder=0
            )
        ax_middle.set_xticks(self.__middle_z_lims, minor=True)
        ax_middle.set_xticklabels(self.__middle_z_lims, minor=True)

        # for axes in (ax, ax_middle):
        #     axes.spines['right'].set_visible(False)
        #     axes.spines['top'].set_visible(False)

        for axes, e_limits in zip(
                (ax, ax_middle), (self.__e_lims, self.__tube_e_lims)
        ):
            axes.set_xlim(self.__z_lims)
            axes.set_ylim(e_limits)

        return fig, ax, ax_middle

    def __split_plot(
            self, x, ys, plot_fn, main_ax, kwargs1=None, kwargs2=None
    ):
        """
        Split up data and plot it on two axes.

        The data with x values inside split_lims (inclusive) is plotted
        using `plot_fn` normally, and the rest is scaled to `main_ax`.

        :param x: x values.
        :type x: 1d array
        :param ys: All the y values.
        :type ys: 2d array
        :param plot_fn: Function that plots data to the middle axes.
            The middle axes are displayed on top, so all plotting uses
            this function to avoid overlapping axes.
        :type plot_fn: function
        :param main_ax: The main axes, which lie underneath the middle
            axes.
        :type ax2_plot_func: matplotlib.pyplot.axes
        :param kwargs1: Keyword arguments passed to the middle plot.
        :type kwargs1: dict or None
        :param kwargs2: Keyword arguments passed to the full plot.
        :type kwargs2: dict or None
        :return: The two plot results.
        :rtype: (handle, handle)
        """
        kwargs1 = kwargs1 or None
        kwargs2 = kwargs2 or None

        inside_indices = _range2indices(x, _PEEK_z_lims)
        outside_indices = numpy.logical_not(inside_indices)

        plots = plot_fn(
            x[outside_indices],
            *(y[outside_indices] for y in ys),
            transform=main_ax.transData,
            **kwargs2
        )
        plots_middle = plot_fn(
            x[inside_indices],
            *(y[inside_indices] for y in ys),
            **kwargs1
        )
        return plots, plots_middle

    def __label_middle(self, ax, energy_label):
        """Label the energy deposit in the middle tube sections."""
        if energy_label:
            ax.annotate(
                energy_label,
                bbox={
                    'boxstyle': 'round4',
                    'facecolor':
                        pyplot.style.library['seaborn']['axes.facecolor'],
                    'edgecolor': 'silver',
                    'alpha': 0.7
                },
                xy=(0, 0),
                xycoords='data',
                xytext=(0, 25),
                textcoords='offset points',
                horizontalalignment='center',
                verticalalignment='center'
            )

    def __get_lims(self, e_lims=None, tube_e_lims=None, z_lims=None):
        """
        Get the plot energy-limits for both sets of axes (full and
        tubes) and the plot z-limits from Piece info.

        The z-limits are shifted from the default 8 by 4 limits in
        increments equal to the `_plate_separation`. Arguments passed in
        will override the calculated limits.
        """
        if ('incident_energy' not in self._piece.info) \
                or (self._piece.info['incident_energy'] == '350gev'):
            if not e_lims:
                e_lims = _e_lims_350
            if not tube_e_lims:
                tube_e_lims = _tube_e_lims_350
        else:
            if not e_lims:
                e_lims = _e_lims_200
            if not tube_e_lims:
                tube_e_lims = _tube_e_lims_200

        if not z_lims:
            plates = self._piece.info.get('plates')
            if (not plates) or plates == (8, 4):
                z_lims = _z_lims_8_4
            else:
                z_lims = (
                    # limit at back plate
                    _z_lims_8_4[0] - (plates[1] - 4) * _plate_separation,
                    # limit at front plate
                    _z_lims_8_4[1] + (plates[0] - 8) * _plate_separation
                )

        return e_lims, tube_e_lims, z_lims

    @staticmethod
    def __add_legend(fig, ax, ax_middle):
        """Place a legend at the upper-right corner of the axes"""
        # Combine legend handles and labels from both axes.
        handles, labels = (val1 + val2 for val1, val2 in zip(
            ax.get_legend_handles_labels(),
            ax_middle.get_legend_handles_labels()
        ))

        legend = ax_middle.legend(
            handles,
            labels,
            loc='upper right',
            ncol=2,
            fontsize='small',

            # Place legend at the corner of the axes rather.
            # `figure.legend` places the legend at the corner of the
            # figure instead.
            # bbox_to_anchor=(1, 1),
            # bbox_transform=ax.transAxes
        )
        legend.set_zorder(1)

class EnergyVsXY(Histogram):
    """Histogram of energy deposit vs. x and y."""

    _default_title = 'Energy vs. x and y.'

    def __init__(self, piece, title=None, xy_lims=None, z_lims=None):
        super().__init__(piece, title)
        self.resultss = []  # Ust a list instead of a DataFrame.

        # Attributes with defaults.
        self.__xy_lims = xy_lims or _default_xy_lims
        self.__z_lims = z_lims or _default_xy_hist_z_lims

        self.__binss = tuple(
            _make_bins(limits[0], limits[1], self.bin_density)
            for limits in self.__xy_lims
        )
        self.__binss_mesh = numpy.meshgrid(*self.__binss)
        # self.__bin_midss = tuple(
        #     _make_bin_midpoints(bins) for bins in self.__binss
        # )

    def plot_single(self, i=None, save=True):
        fig, ax = self._make_fig_and_axes()
        fig.suptitle(self.title)

        ax.pcolormesh(*self.__binss_mesh, self._to_density(self.get(i)))

        if save:
            self.save_fig(fig, file_suffix='xy')

    def plot_means(self):
        fig, ax = self._make_fig_and_axes()
        fig.suptitle(self.title + ' Averaged.')

        ax.pcolormesh(
            *self.__binss_mesh,
            self._to_density(numpy.mean(self.resultss, axis=0))
        )

        self.save_fig(fig, file_suffix='xy')

    def get(self, i=None):
        """
        Overrides the base to get from a list instead of a DataFrame.
        """
        if i is None:
            self._assert_single_entry()
            return self.resultss[0]

        return self.resultss[i]

    def add_results(self, sums):
        """
        Overrides the base to append to a list instead of a DataFrame.
        """
        self._check_results(sums)
        self.resultss.append(sums)

    def _data2results(self, data):
        sliced = data[numpy.logical_and(
            self.__z_lims[0] <= data.z, data.z <= self.__z_lims[1]
        )]
        # Transpose the sums to match the xy/Cartesian indexing used by
        # `numpy.meshgrid` and Matplotlib's `pcolormesh`.
        return numpy.histogram2d(
            sliced.x, sliced.y, bins=self.__binss, weights=sliced.energy_deposit
        )[0].T

    def _check_results(self, sums):
        """Check that `sums` is the correct shape for the bins."""
        assert numpy.all(numpy.equal(
            sums.shape,
            tuple(len(bins) - 1 for bins in self.__binss)
        ))

    def _make_fig_and_axes(self):
        fig, ax = pyplot.subplots()
        fig.tight_layout(rect=(0, 0, 1, 0.86))

        ax.set_xlabel(f'x ({self.length_units})')
        ax.set_ylabel(f'y ({self.length_units})')

        # No need to set x limits since the axes will be equal.
        # ax.set_xlim(self.__xy_lims[0])
        ax.set_ylim(self.__xy_lims[1])

        ax.set_aspect('equal')

        return fig, ax

def save_dataframe(dataframe, path):
    """
    Write a DataFrame to a csv file with nicely spaced columns and a
    header, and also write it to a second file compatible with Office
    365 (since it seems to prohibit the use of csv's).

    :param dataframe: The DataFrame to save.
    :param path: Path to save the files at. The extension of the
        filename at `path` is ignored and replaced with 'csv' and
        'xlsx'.
    :return: The paths to the saved files, (csv_path, excel_path).
    """
    tail, ext = os.path.splitext(path)

    csv_path = tail + '.csv'
    header = (
        'FCal and SCal analysis output.'
        f' {datetime.datetime.now().ctime()}'
    )
    with open(csv_path, 'w') as file:
        file.write('\n'.join((
            header,
            dataframe.to_string(index=False),
            ''
        )))

    excel_path = tail + '.xlsx'
    dataframe.to_excel(excel_path, index=False)

    return csv_path, excel_path

def _range2indices(array, range_lims):
    """
    Find the values of an array that are inclusively within range_lims.

    :param array: Array of values that are compared to range_lims.
    :type array: numpy.ndarray
    :param range_lims: (lower limit, upper limit)
    :type range_lims: (float, float)
    :return: Boolean array locating the indices of the values within the
        range.
    :rtype: numpy.ndarray
    """
    return numpy.logical_and(range_lims[0] <= array, array <= range_lims[1])

@functools.lru_cache()
def _make_bins(start, end, bin_density):
    """
    Make histogram bins with a given bin density.

    Since the bin density is fixed, the bins start exactly at `start`
    and end at a rounded up value near `end`.

    :param start: Exact start of the first bin.
    :type start: float
    :param end: Lower limit of the end of last bin.
    :type end: float
    :param bin_density: Exact bin density.
    :type bin_density: float
    :return: Bin endpoints.
    :rtype: numpy.ndarray
    """
    assert (end > start)

    # Round the end limit so the bin size is correct.
    num_bins = math.ceil(bin_density * (end - start))
    new_end = start + num_bins / bin_density

    return numpy.linspace(start, new_end, num_bins + 1)

# @functools.lru_cache()
def _make_bin_midpoints(bins):
    """Get midpoints of histogram bins."""
    return (bins[:-1] + bins[1:]) / 2

def _ordered_append(dataframe, other):
    """
    Append `other` to `dataframe` while keeping the order of the columns
    in both.

    The `append` method that comes with pandas Dataframes alphabetizes
    the order of the columns, which may not always be desired. This
    function provides an alternative.

    :param dataframe: The DataFrame to append to.
    :type dataframe: pandas.DataFrame
    :param other: The data to append.
    :type other: pandas.DataFrame or pandas.Series/dict-like object,
        or list of these
    :return: New DataFrame with appended data.
    :rtype: pandas.DataFrame
    """
    # The new column order. New columns from `other` are inserted in
    # front of the original `dataframe` columns.
    columns = [
                  col for col in other.keys() if col not in dataframe.keys()
              ] + list(dataframe.keys())

    # pandas' append alphabetizes the columns here. Indexing the result
    # with `columns` sets our own column order.
    return dataframe.append(other, ignore_index=True)[columns]